RSQ: Learning from Important Tokens Leads to Better Quantized LLMs
This work addresses the challenge of efficiently compressing large models for deployment, offering a novel approach to quantization that enhances performance, though it is incremental in building on existing quantization frameworks.
The paper tackles the problem of improving layer-wise quantization for large language models by prioritizing learning from important tokens, such as those with large attention scores, and proposes RSQ (Rotate, Scale, then Quantize) to achieve this. The result is that RSQ consistently outperforms baseline methods across multiple downstream tasks and model families, with superior performance on long-context tasks and generalizability across various setups.
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.