RSQ: Aprender de los tokens importantes conduce a mejores LLMs cuantizados

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

Resumen

La cuantificación por capas es una técnica clave para comprimir eficientemente modelos grandes sin un costoso reentrenamiento. Los métodos anteriores suelen cuantificar los pesos de cada capa optimizando "uniformemente" la pérdida de reconstrucción de la capa en todos los tokens de salida. Sin embargo, en este artículo demostramos que se pueden obtener modelos mejor cuantificados priorizando el aprendizaje de tokens importantes (por ejemplo, aquellos con puntuaciones de atención altas). Basándonos en este hallazgo, proponemos RSQ (Rotar, Escalar y luego Cuantificar), que (1) aplica rotaciones (transformación ortogonal) al modelo para mitigar los valores atípicos (aquellos con una magnitud excepcionalmente grande), (2) escala las características del token según su importancia, y (3) cuantifica el modelo utilizando el marco GPTQ con las estadísticas de segundo orden calculadas por tokens escalados. Para calcular la importancia de los tokens, exploramos tanto estrategias heurísticas como dinámicas. Basándonos en un análisis exhaustivo de todos los enfoques, adoptamos la concentración de atención, que utiliza las puntuaciones de atención de cada token como su importancia, como el mejor enfoque. Demostramos que RSQ supera consistentemente a los métodos de referencia en múltiples tareas posteriores y en tres familias de modelos: LLaMA3, Mistral y Qwen2.5. Además, los modelos cuantificados con RSQ logran un rendimiento superior en tareas de contexto largo, lo que destaca aún más su efectividad. Por último, RSQ demuestra generalizabilidad en varias configuraciones, incluyendo diferentes tamaños de modelos, conjuntos de datos de calibración, precisiones de bits y métodos de cuantificación.

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.