RSQ: Learning from Important Tokens Leads to Better Quantized LLMs

TL;DR

RSQ introduces a token-aware approach to post-training quantization by rotating weights to mitigate outliers, scaling token features by token importance (favoring influential tokens), and applying GPTQ with a Hessian adapted for scaled tokens. The method demonstrates that prioritizing learning from important tokens yields consistent gains across model families and tasks, including long-context benchmarks, and proves robust across model sizes, calibration datasets, bit-precisions, and even vector quantization. Key innovations include a formal objective modification using a diagonal importance matrix $\mathbf{R}$, per-layer token importance strategies (most notably Attention Concentration, AttnCon), and dataset expansion to reduce positional bias. Overall, RSQ advances practical, data-efficient quantization of large transformers by aligning the learning objective with information-critical tokens, delivering meaningful improvements in both short- and long-context scenarios.

Abstract

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.

Figures (14)

• Figure 1: Illustration of layer-wise quantization (left), three-step process of RSQ (middle) and the dataset expansion (right). On the middle, circle size and red color intensity represent weight magnitude, with larger circles and deeper red colors indicating greater magnitudes.
• Figure 2: Evaluation of two heuristic approaches with varying numbers of used tokens.
• Figure 3: Evaluation of five dynamic approaches with varying $r_{min}$.
• Figure 4: The effect of expanding the dataset on different methods.
• Figure 5: Ablation on model sizes.