Метод

1️⃣ Первое наблюдение полностью идентично KIVI - ключи обладают ярко выраженными выбросами, а значения - нет. Потому ключи квантизуются поканально, значения - потокенно.

2️⃣ Далее замечают, что RoPE перемешивает соседние каналы и делает выбросы и channel-specific поведение не столько ярко выраженным, потому предлагают квантизовать ключи до применения RoPE и каждый раз применять его уже к квантизованным кэшам.

3️⃣ Однородная квантизация приятна и удобна в использовании, но обычно далека от оптимальной в плане ошибки. И не мудрствуя лукаво предлагают использовать 1-мерный k-means, но используя в качестве метрики не стандарную Евклидову, а диагональ матрицы Фишера, которую оценивают по маленькой выборке.

4️⃣ Некоторые значения сильно выбиваются среди остальных, потому их предлагается не квантизовать, а держать в виде разреженной матрицы (CSR, CSC). На практике берут 1% самых больших по модулю значений.

5️⃣ Модели чувствительны к изменениям attention sinks, отличающихся большими нормами активаций, потому предлагается их не квантизовать - все равно их мало, и на общей битности не скажется. Не квантизуют только первый токен.

6️⃣ Динамическое определение диапазона значений на лету дорого, и для поканального квантования требует пересчета значений всех токенов. По этой причине предлагают калибровку констант квантизации для ключей проводить offline, на небольшой выборке, а потокенную квантизацию кэшей на лету.

И для всей этой навороченной красоты пишут специальный кернел, что работало все с адекватной скоростью.

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

Метод валидируют на перплексии на стандартной длине контекста на 🦙-1.

- В 4 бита просадка порядка 0.01 ppl
- В 3 бита просадка порядка 0.1 ppl
- В 2 бита - десятые

Затем метод проверяют на версиях Llama-2, адаптированных под длинный контекст, где KVQuant почти не просаживает метрики в 4 битах, и умеренно в 3 бита. На Passkey 🗝 retrieval метод не просаживается по качеству, в отличие от KIVI.

Далее KVQuant прогоняют на LongBench/RULER (где одна из задач - “иголка в сене”, вытащить нужный факт из далекого прошлого) - популярных бенчах на длинный контекст.

KVQuant снова лучше KIVI, на RULER KIVI (эффективно 3-битный) заметно страдает (56%->40%), но KVQuant теряет 3% качества (56->53%).

Метод совместим с квантизацией весов. Так как авторы KVQuant - авторы еще и SqueezeLLM, то совмещают именно с этим подходом. Квантизация кэшей почти не меняет метрики квантизованной модели.

В ablation показывают, что все компоненты важны для успеха:
1️⃣ Выбор размерностей для квантизации
2️⃣ Pre-RoPE квантизация
3️⃣ Неоднородная квантизация
4️⃣ Учет выбросов

Итоговое ускорение операций умножения в KV - 1.2x-1.7x (не end-to-end latency). При 2-битной квантизации удается впихнуть на одну A100 1M токенов для Llama-2-7b, и на 8 - 10M токенов.

Вывод

Метод композитный, солидный, заточенный под выбивание хорошего качества. Однако сложность имеет свою цену - и в данном случае это достаточно дорогая деквантизация, из-за которой ускорение инференса меньше, чем у конкуретных подходов (того же KIVI).