PACE сжимает состояние линейных LLM до 4 бит и экономит треть памяти

Команда исследователей предложила PACE — посттренировочный квантизатор, который хранит рекуррентное состояние линейно-attention моделей в четырёх битах и при этом не требует дообучения. На четырёх гибридных моделях размером от 4B до 35B параметров вариант со смешанными 4- и 8-битными головами сопоставим с 8-битным DAMP на трёх крупных моделях, занимая около двух третей его памяти. Версия, где все головы четырёхбитные, при равном объёме памяти набирает больше, чем Q-Mamba.

Почему состояние становится узким местом

У моделей с линейным вниманием контекст живёт не в растущем KV-кэше, а в рекуррентном состоянии фиксированного размера. Казалось бы, это удобно. Но инференс-системе приходится держать отдельную копию такого состояния для каждой последовательности, которую она декодирует одновременно. Чем больше копий помещается в память ускорителя, тем больше запросов обслуживается параллельно, а значит, дешевле каждый токен.

Отсюда интерес к квантизации самого состояния. Существующие методы округляют его после каждого сгенерированного токена. Ошибка округления при этом накапливается шаг за шагом, и на четырёх битах точность заметно проседает. Авторы прежних работ приходили к одному выводу: либо поднимай разрядность, либо дообучай модель под квантизацию. PACE показывает, что есть третий путь.

Два наблюдения, на которых держится PACE

Первое касается частоты округления. PACE округляет состояние только раз в несколько токенов, а промежуточные токены держит в небольшом буфере. Похожим образом уже работают chunked-ядра на этапе префилла: они обновляют состояние раз на чанк, а не на каждый токен. Меньше округлений означает меньше накопленной ошибки.

Второе наблюдение касается структуры самого состояния в gated-delta слоях. Крупные значения там не разбросаны равномерно, а собираются в столбцах. Причина в том, что некоторые каналы обновления, записываемого на каждом токене, остаются большими на протяжении многих токенов. Поэтому PACE использует общий масштаб квантизации на целый столбец. Одна крупная запись больше не портит точность всех остальных значений, которые вынуждены делить с ней масштаб.

Третья деталь — смешанная точность. Некоторые головы дольше других удерживают ошибку округления. PACE оставляет именно их в восьми битах, а остальные переводит в четыре.

• Метод: посттренировочный, без дообучения модели
• Целевой формат: четыре бита на элемент состояния, восемь для самых «долгоживущих по ошибке» голов
• Округление: раз в несколько токенов, промежуточные токены хранятся в буфере
• Масштаб квантизации: общий на столбец состояния
• Проверка: четыре гибридные модели от 4B до 35B параметров

Что показали замеры

Смешанный PACE с 4- и 8-битными головами оказался на уровне 8-битного DAMP на трёх более крупных моделях из четырёх. Память при этом составила примерно две трети от DAMP. Чисто четырёхбитный вариант выигрывает у Q-Mamba в сравнении при равном объёме памяти. Отдельно авторы отмечают, что округление раз в несколько токенов сохраняет ускорение декодирования, которое даёт буфер: экономия памяти не оплачивается потерей скорости.

Последний пункт важен для практики. Квантизация состояния, которая каждый шаг добавляет операции округления, может съесть выигрыш от уплотнения батча. Здесь округление вынесено из горячего цикла, и этот риск уменьшается.

Где это пригодится на инфраструктуре

Гибридные архитектуры, в которых линейные слои чередуются с обычным вниманием, всё чаще появляются в открытых и коммерческих моделях. Для тех, кто их обслуживает, PACE даёт понятный рычаг: уменьшить память на состояние примерно на треть относительно 8-битной схемы без повторного обучения и дообучения. Освободившееся место можно отдать под больший батч или под более длинную очередь одновременных сессий.

Есть и ограничения. Результат подтверждён на четырёх гибридных моделях, и для самой маленькой из них смешанный PACE не достиг уровня DAMP. Командам, которые планируют внедрение, стоит проверить, как метод ведёт себя на их модели и нагрузке, прежде чем перестраивать продакшен под четырёхбитное состояние. Но сама идея — округлять реже и подстраивать масштаб под реальную структуру данных — выглядит переносимой и недорогой в реализации.