Проблема памяти и скорости

В процессе генерации текста трансформеры сохраняют данные о ранее обработанных токенах в специальном хранилище — KV-кэше. Этот кэш растет линейно вместе с длиной контекста и размером пакета данных. При работе с длинными текстами (например, 128 тысяч токенов) объем кэша может в разы превышать размер весов самой модели.

Проблема не только в объеме памяти. Для генерации каждого нового токена GPU должен считывать весь KV-кэш из памяти устройства. Пока идет этот процесс, вычислительные ядра простаивают. Сжатие кэша позволяет не только освободить место, но и ускорить работу, сокращая время ожидания передачи данных.

Ловушка квантования

Обычно квантование (сжатие данных до меньшего количества бит) воспринимается как простой регулятор: чем меньше бит, тем меньше точность. Однако внутри KV-кэша всё сложнее. Здесь важно не то, сколько бит вы используете, а то, по какой оси вы группируете данные для вычисления масштаба сжатия.

При использовании всего 2 бит у нас есть всего четыре уровня для представления значений. Если в группе оказывается одно аномально большое число, оно «раздувает» шаг масштабирования для всех остальных элементов в этой группе. В итоге обычные значения становятся слишком грубыми и неточными. Выбор оси группировки — это решение о том, какие элементы будут страдать от этой ошибки вместе.

Разная природа ключей и значений

Анализ показывает, что ключи (Keys) и значения (Values) в кэше имеют разную структуру. В ключах есть «каналы-аномалии» — специфические параметры, которые имеют огромные значения во всех токенах последовательности. Если группировать ключи по токенам, эти аномалии испортят точность всех остальных каналов. Если же группировать их по каналам, аномалии окажутся в своих собственных группах, не мешая остальным данным.

С значениями (Values) ситуация иная: они распределены более равномерно. Но здесь возникает парадокс. Если измерять ошибку просто по точности восстановления данных, то группировка по каналам кажется лучше. Однако на практике такая стратегия полностью разрушает работу модели.

Почему стандартные метрики врут

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

При группировке по токенам ошибка одного токена остается внутри него и почти не влияет на результат, так как веса внимания для неважных токенов близки к нулю. При группировке по каналам ошибка «размазывается» по всем токенам, и искажения от неважных токенов начинают портить важные данные. Таким образом, метрика ошибки реконструкции данных не является надежным показателем качества работы модели.

Что это значит

Для эффективного сжатия LLM необходимо измерять ошибку не в месте хранения данных, а в месте их непосредственного использования. При работе с ключами важно изолировать аномальные каналы, а при работе со значениями — использовать группировку по токенам, чтобы ошибки не распространялись на важные элементы контекста.