PrefixQuant: Static Quantization Beats Dynamic through Prefixed Outliers in LLMs
[Статья] [Код]

Введение

Для получения максимального ускорения, особенно при инференсе большими глотками батчами, нужно квантизовать не только веса, но и активации. Однако, богомерзкие выбросы сильно затрудняют квантизацию, и в прошлых работах (QuaRot, SpinQuant, DuQuant) было предложено применять Адамаровы матрицы (или некоторые обучаемые), которые размазывают 🧈 выбросы по каналам.

Однако, выбросы встречают не только в отдельных каналах, но есть и токены, например <BOS> и иные разделители, отличающиеся большими значениями активаций, и их квантизация большими группами (per-tensor) с максимальным ускорением приводит к большим ошибкам. Поэтому в упонятых выше методах использовали per-token квантизацию с некоторым замедлением инференса.

В этой статье авторы предлагают вынести аутлаерные токены в префикс последовательности дабы упростить квантизацию.

Метод

Сам подход мотивирован идеями про attention sinks и Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing. Иногда трансформерному блоку не нужно ничего делать, и дабы реализовать почти тождественную операцию, приходится выбрасывать весь attention в токены с низкой семантикой.

Потому авторы предлагают просто найти все эти токены и положить в префикс последовательности.

Токены находят на основе некоторой калибровочной выборки смотря на выход down_proj в MLP. Отбирают токены, чья L_{\infty} норма сильно больше медианной по последовательности.

Сам процесс достаточно быстрый - занимает всего несколько минут. Таких токенов оказывается от 1 до 4 в рассмотренных моделях (Llama-2,3, Qwen, Mistral). Обычно это <BOS>, и еще символы переноса строки \n, артикли и знаки препинания.

Данная модификация позволяет существенно уменьшить норму последующих токенов и упростить процесс квантизации.
Для максимизации качества в конце еще предлагается делать поблочный QAT, следуя EfficientQAT.

Эксперименты

Метод валидидируют на 🦙-2,3 в режимах W8A8KV8, W4A8KV4, and W4A4KV4 квантизации (W - точность весов, A - активаций, KV - KV-кэшей).

Веса квантизуются поканально, кэши поголовно, активации потензорно.

Предложенный подход оказывается примерно на одном уровне с per-token квантизациями, чуть уступая лишь SpinQuant с дообучением матриц вращения каналов. При этом скорость инференса на 20-30% быстрее.

Поблочный файнтьюн немного накидывает.

Вывод

Простой и полезный трюк для квантизации. Тем не менее, кажется, что возможность делать ничего в блоке модели стоит закладывать изначально, и скорее всего, ввиду потребности в низко-битных моделях, сообщество до этого дойдет (через softmax+1 или иным образом).