KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
[Статья] [Код]

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

Группа исследователей (с забавным примечанием The order of authors is determined by flipping a coin 🪙) из разных мест на послендем ICML представила статью со звучным названием KIVI (терминал оплаты, фрукт и птица пишутся по-другому).

Метод

Предыдущие методы квантизации кэшей обычно квантизовали кэши и ключи и потокенно, так как статистики (scale и zero) при таком подходе можно пересчитывать на лету. Для 4-битной квантизации такой подход работает хорошо, но при 2-х битных кэшах модель становится полностью негодной.

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

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

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

Метод валидируют на бенчах из LM-Eval, LongBench (наборе задач на длинный контекст) и Needle-in-a-Haystack (где нужно достать вытащить нужный факт из длинного промпта) на моделях Llama-2, Mistral.

Наивная квантизация в 2 бита сильно страдает по качеству, а предложенный подход работает с совсем небольшой просадкой. Есть, правда, маленький нюанс, что 2-битная квантизация использует размер группы 32 с fp zero-point, что означает на практике 3 бита 😅. Needle-in-a-Haystack решается почти идеально.

Поддержание неквантизованного буфера свежим токенов важно для лучше качества (свежие токены влияют сильнее, чем далекие, в большинстве задач).

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

Вывод

Сжатие кэшей - богоугодное дело. А дабы эффективно квантизовать кэши полезно смотреть на распределение того - что сжимаешь.