RSQ: Aprender com Tokens Importantes Leva a LLMs Quantizados Melhores

RSQ: Learning from Important Tokens Leads to Better Quantized LLMs

March 3, 2025

Autores: Yi-Lin Sung, Prateek Yadav, Jialu Li, Jaehong Yoon, Mohit Bansal

cs.AI

Resumo

A quantização camada por camada é uma técnica fundamental para comprimir eficientemente modelos grandes sem a necessidade de retreinamento dispendioso. Métodos anteriores normalmente quantizam os pesos de cada camada otimizando "uniformemente" a perda de reconstrução da camada em todos os tokens de saída. No entanto, neste artigo, demonstramos que modelos melhor quantizados podem ser obtidos priorizando o aprendizado a partir de tokens importantes (por exemplo, aqueles que possuem grandes pontuações de atenção). Com base nessa descoberta, propomos o RSQ (Rotate, Scale, then Quantize), que (1) aplica rotações (transformação ortogonal) ao modelo para mitigar outliers (aqueles com magnitude excepcionalmente grande), (2) dimensiona o recurso do token com base em sua importância, e (3) quantiza o modelo usando o framework GPTQ com as estatísticas de segunda ordem calculadas por tokens dimensionados. Para calcular a importância do token, exploramos tanto estratégias heurísticas quanto dinâmicas. Com base em uma análise detalhada de todas as abordagens, adotamos a concentração de atenção, que usa as pontuações de atenção de cada token como sua importância, como a melhor abordagem. Demonstramos que o RSQ supera consistentemente os métodos de base em várias tarefas subsequentes e três famílias de modelos: LLaMA3, Mistral e Qwen2.5. Além disso, modelos quantizados com RSQ alcançam desempenho superior em tarefas de contexto longo, destacando ainda mais sua eficácia. Por fim, o RSQ demonstra generalização em várias configurações, incluindo diferentes tamanhos de modelos, conjuntos de dados de calibração, precisões de bits e métodos de quantização.

English

Layer-wise quantization is a key technique for efficiently compressing large models without expensive retraining. Previous methods typically quantize the weights of each layer by "uniformly" optimizing the layer reconstruction loss across all output tokens. However, in this paper, we demonstrate that better-quantized models can be obtained by prioritizing learning from important tokens (e.g. which have large attention scores). Building on this finding, we propose RSQ (Rotate, Scale, then Quantize), which (1) applies rotations (orthogonal transformation) to the model to mitigate outliers (those with exceptionally large magnitude), (2) scales the token feature based on its importance, and (3) quantizes the model using the GPTQ framework with the second-order statistics computed by scaled tokens. To compute token importance, we explore both heuristic and dynamic strategies. Based on a thorough analysis of all approaches, we adopt attention concentration, which uses attention scores of each token as its importance, as the best approach. We demonstrate that RSQ consistently outperforms baseline methods across multiple downstream tasks and three model families: LLaMA3, Mistral, and Qwen2.5. Additionally, models quantized with RSQ achieve superior performance on long-context tasks, further highlighting its effectiveness. Lastly, RSQ demonstrates generalizability across various setups, including different model sizes, calibration datasets, bit precisions, and quantization methods.