RSQ: Обучение на важных токенах приводит к созданию более качественных квантованных языковых моделей

RSQ: Learning from Important Tokens Leads to Better Quantized LLMs

March 3, 2025

Авторы: Yi-Lin Sung, Prateek Yadav, Jialu Li, Jaehong Yoon, Mohit Bansal

cs.AI

Аннотация

Послойное квантование является ключевой техникой для эффективного сжатия больших моделей без дорогостоящего переобучения. Традиционные методы обычно квантуют веса каждого слоя, "равномерно" оптимизируя потерю реконструкции слоя для всех выходных токенов. Однако в данной работе мы показываем, что более качественно квантованные модели могут быть получены путем приоритизации обучения на важных токенах (например, тех, которые имеют высокие оценки внимания). Основываясь на этом наблюдении, мы предлагаем метод RSQ (Rotate, Scale, then Quantize), который (1) применяет вращения (ортогональные преобразования) к модели для смягчения выбросов (элементов с исключительно большой величиной), (2) масштабирует признаки токенов в зависимости от их важности и (3) квантует модель с использованием фреймворка GPTQ с вычислением статистики второго порядка на основе масштабированных токенов. Для вычисления важности токенов мы исследуем как эвристические, так и динамические стратегии. На основе тщательного анализа всех подходов мы выбираем концентрацию внимания, которая использует оценки внимания каждого токена в качестве его важности, как наилучший метод. Мы демонстрируем, что RSQ стабильно превосходит базовые методы в различных задачах и для трех семейств моделей: LLaMA3, Mistral и Qwen2.5. Кроме того, модели, квантованные с помощью RSQ, показывают превосходную производительность в задачах с длинным контекстом, что дополнительно подчеркивает его эффективность. Наконец, RSQ демонстрирует универсальность в различных настройках, включая разные размеры моделей, калибровочные наборы данных, битовые точности и методы квантования.

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.