Новый метод сжатия KV-кэша SelKV ускоряет LLM в 3 раза при 100k токенов
Учёные опубликовали на arXiv новый метод сжатия ключ-значение (KV) кэша для больших языковых моделей (LLM). Разработка под названием SelKV (Selective KV Cache Merging) позволяет значительно снизить потребление памяти и ускорить инференс, не требуя дообучения.
Проблема KV-кэша заключается в том, что его объём линейно растёт с длиной контекста, что становится узким местом для LLM. Существующие методы сжатия, такие как слияние токенов, часто теряют качество из-за грубой агрегации и эффекта «провисания внимания» (attention sag), когда объединённые токены получают ту же массу softmax, что и отдельные.
SelKV решает эти проблемы двумя компонентами. Первый — мягкое косинусное стробирование, которое адаптивно регулирует решения о слиянии на основе схожести векторов значений. Оно подавляет или отбрасывает несхожие токены, сохраняя семантическую точность. Второй — механизм компенсации отношения внимания, который применяет смещение логитов на этапе декодирования, полученное из статистики префильного внимания, исправляя дисбаланс softmax, вызванный слиянием.
Эксперименты проводились на наборе LongBench (16 англоязычных наборов данных). При сохранении лишь 25% KV-кэша SelKV демонстрирует сильную производительность сжатия по сравнению с базовыми методами. Особенно эффективен метод на моделях с групповым запросным вниманием (GQA), где достигается почти безошибочное качество генерации.
Более того, на задачах множественного ответа по нескольким документам SelKV даже превосходит полный кэш, а при длине контекста в 100 000 токенов обеспечивает ускорение декодирования в 3,3 раза. Результаты подтверждают, что предложенный фреймворк может быть эффективной альтернативой для развёртывания LLM на устройствах с ограниченной памятью.
Разработчики подчёркивают, что SelKV не требует дополнительного обучения и может быть легко интегрирован в существующие модели. Это делает его перспективным инструментом для ускорения работы чат-ботов, ассистентов и других приложений на базе LLM.
