RSQ: Learning from Important Tokens Leads to Better Quantized LLMs

2025-03-03

AI Summary

To address performance degradation in large language model (LLM) layer-wise quantization caused by neglecting token-level importance, this paper proposes a token-importance-aware hierarchical quantization method. Methodologically, it introduces token importance—specifically attention scores—into the layer-wise quantization objective for the first time, and designs a Rotation-Scaling-Quantization (RSQ) framework: (1) orthogonal rotation to decouple outliers; (2) dynamic importance-weighted scaling guided by attention concentration to preserve critical features; and (3) precise quantization via joint GPTQ optimization and second-order statistical calibration. Evaluated on LLaMA3, Mistral, and Qwen2.5 families, the method consistently outperforms state-of-the-art quantization approaches—especially in long-context tasks—while maintaining compatibility across bit-widths (e.g., 4-bit/6-bit), diverse calibration datasets, and model scales (1B–70B).

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Improves quantization by focusing on important tokens
• Proposes RSQ method for better model compression
• Enhances performance on long-context tasks

Innovation

• Prioritizes learning from important tokens
• Applies rotations to mitigate outlier effects
• Uses attention scores for token importance scaling