SpinQuant: LLM quantization with learned rotations
[Статья][Код]

Если в названии статьи встречаются вместе Spin и Quant - это наверняка что-то из области физики. Но нет, это снова про квантизацию LLM…

Ранее в ряде работ (QuIP, QuIP#, QuaRot) было показано, что поворот на случайную ортогональную матрицу сильно облегчает квантизацию.

И в этой работе ребята из Meta предложили развитие идеи, дающее лучшее качество на Weight + Activation (+cache) квантизации.

Метод

Наличие выбросов в отдельных каналах активаций сильно затрудняет квантизацию (так как per-channel квантизация обычно не поддерживается в железе). Ортогональные матрицы “размазывают” выбросы в весах и активациях между каналами, так что все активации имеют примерно один и тот же порядок. Но конкретный выбор ортогональной матрицы имеет значение.

Авторы замечают, что качество сильно варьируется от конкретного выбора матриц, если сэмплировать их произвольным образом (и редко бывает достаточно хороших).

Случайные Адамаровы матрицы из QuIP гораздо лучше, но все разброс существенен и нет гарантий на оптимальность.

Потому предлагается искать “хорошие” матрицы поворота с помощью оптимизации на Stiefel многообразиях, параметризующих ортогональные матрицы. Оптимизация представляет собой градиентный спуск на многообразии. А именно Кэли SGD.

Есть 4 места куда можно вкрутить поворот:
R1 - В residual перед Attention и FFN
R2 - Между V и O матрицами в Attention
R3 - В Query/Key
R4 - Между активацией в FFN

Первые два поворота можно влить в параметры модели без изменения выходов модели, последние же придется считать на каждом прогоне. Потому предложенная процедура оптимизации применяется для поиска оптимальных R1, R2 матриц. А в случаях R3, R4 используют дешевые Адамаровы матрицы, кои практически не замедляют работу LLMки и почти ничего не весят.

Для оптимизации используют небольшой набор из Wikitext-2.

Результаты

Процедура оптимизации матриц занимает час с небольшим на меньших представителях семейства 🦙 на 8 A100 и полдня для 70B моделей.

Замеряют по классике качество на перплексии и 0-шотах.

Рассматривают 3 сценария:
1️⃣ Weight only quantization
2️⃣ Weight + activation quantization
3️⃣ Weight + activation + KV cache quantization

По качеству операжают бейзлайны (ближайший QuaRot примерно на 1% в среднем). В Weight-only quantization где-то бага при замерах AQLM и QuIP#, ибо эти методы сжимают модели почти без просадки (если верить бенчам, а не здравому смыслу) почти без просадки.

Метод работает как с RTN (Round-to-nearest), так и c GPTQ.

Важны все матрицы, но большой прирост дает будто бы добавление R4 - вращение внутри FFN.

Вывод

Квантизация весов/активаций и кэша в 4 бита на текущий момент, наверное, то что дает самое большое ускорение инференса на современных GPU, потому и представляет значительный практический интерес. И улучшение качества - это хорошо. При квантизации активаций (и кэша) просадки все еще значительные и есть куда улучшаться.