ReFreeKV: К сжатию KV-кэша без порогов

Аннотация

Для снижения потребления памяти при инференсе больших языковых моделей (LLM) был предложен ряд методов сокращения KV-кэша. Хотя эти техники позволяют достичь безошибочного уменьшения памяти на многих наборах данных, они часто опираются на неявно выраженное условие: для достижения оптимальной производительности необходимо заранее определить пороговое значение, зависящее от входных данных, для бюджета KV-кэша. Однако такой чувствительный к входным данным подход может быть существенно ограничен в реальных сценариях, поскольку входные данные из открытых доменов охватывают разнообразные области, длины и уровни сложности без четких границ для выбора порога. В результате зависимость от такого чувствительного к входным данным порога может стать фундаментальным ограничением, вызывающим значительное ухудшение производительности на произвольных входных данных. В данной работе мы предлагаем новую цель, которая снимает ограничения, связанные с порогами, для надежного сжатия KV-кэша, выступая за "беспороговые" методы, которые адаптивно распределяют бюджет, сохраняя производительность полного кэша. Затем мы предлагаем новый метод, ReFreeKV, служащий первой реализацией этой цели. Обширные эксперименты на 13 наборах данных с различной длиной контекста, типами задач и размерами моделей демонстрируют его эффективность и результативность. Наш код опубликован по адресу https://github.com/Patrick-Ni/ReFreeKV.

English

To reduce memory consumption during LLM inference, a handful of methods have been proposed for KV cache pruning. While these techniques can accomplish lossless memory reduction on many datasets, they often hinge on an under-emphasized condition: an input/domain-specific threshold for KV cache budget needs to be pre-determined to achieve the optimal performance. However, such input-sensitive design may be considerably limited in real-world scenarios, as open-domain inputs span diverse domains, lengths and difficulty levels, without clear boundaries for threshold selection. As a result, the dependence of such input-sensitive threshold can be a fundamental limitation that causes large degradation on arbitrary inputs. In this work, we propose a new objective that lifts the threshold constraints for robust KV compression, advocating for "threshold-free" methods that adaptively adjust budget allocation while preserving full-cache performance. We then propose a novel method, ReFreeKV, serving as the first instantiation of this objective. Extensive experiments across 13 datasets with diverse context lengths, task types, and model sizes demonstrate its efficacy and efficiency. Our code is publicly released at https://github.com/Patrick-Ni/ReFreeKV.