RSQ : Apprendre à partir des tokens importants conduit à de meilleurs LLM quantifiés

RSQ: Learning from Important Tokens Leads to Better Quantized LLMs

March 3, 2025

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

cs.AI

Résumé

La quantification couche par couche est une technique clé pour compresser efficacement les grands modèles sans recourir à un réentraînement coûteux. Les méthodes précédentes quantifient généralement les poids de chaque couche en optimisant "uniformément" la perte de reconstruction de la couche sur tous les tokens de sortie. Cependant, dans cet article, nous démontrons que des modèles mieux quantifiés peuvent être obtenus en priorisant l'apprentissage à partir des tokens importants (par exemple, ceux ayant des scores d'attention élevés). Sur la base de cette observation, nous proposons RSQ (Rotate, Scale, then Quantize), qui (1) applique des rotations (transformation orthogonale) au modèle pour atténuer les valeurs aberrantes (celles ayant une magnitude exceptionnellement grande), (2) met à l'échelle les caractéristiques des tokens en fonction de leur importance, et (3) quantifie le modèle en utilisant le cadre GPTQ avec les statistiques du second ordre calculées par les tokens mis à l'échelle. Pour calculer l'importance des tokens, nous explorons à la fois des stratégies heuristiques et dynamiques. Sur la base d'une analyse approfondie de toutes les approches, nous adoptons la concentration d'attention, qui utilise les scores d'attention de chaque token comme mesure de son importance, comme la meilleure approche. Nous démontrons que RSQ surpasse systématiquement les méthodes de référence sur plusieurs tâches en aval et trois familles de modèles : LLaMA3, Mistral et Qwen2.5. De plus, les modèles quantifiés avec RSQ obtiennent des performances supérieures sur les tâches à contexte long, soulignant davantage son efficacité. Enfin, RSQ démontre une généralisabilité dans divers contextes, incluant différentes tailles de modèles, jeux de données d'étalonnage, précisions en bits et méthodes de quantification.

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.