ReFreeKV: Rumo à Compressão de Cache KV Sem Limiar

Resumo

Para reduzir o consumo de memória durante a inferência de LLMs, alguns métodos foram propostos para a poda de cache KV. Embora essas técnicas possam alcançar uma redução de memória sem perdas em muitos conjuntos de dados, elas frequentemente dependem de uma condição pouco enfatizada: um limiar específico de entrada/domínio para o orçamento do cache KV precisa ser pré-determinado para alcançar o desempenho ideal. No entanto, tal projeto sensível à entrada pode ser consideravelmente limitado em cenários do mundo real, pois entradas de domínio aberto abrangem diversos domínios, comprimentos e níveis de dificuldade, sem limites claros para a seleção do limiar. Como resultado, a dependência de um limiar tão sensível à entrada pode ser uma limitação fundamental que causa grande degradação em entradas arbitrárias. Neste trabalho, propomos um novo objetivo que elimina as restrições de limiar para uma compressão KV robusta, defendendo métodos "sem limiar" que ajustam adaptativamente a alocação de orçamento enquanto preservam o desempenho do cache completo. Em seguida, propomos um método inovador, ReFreeKV, servindo como a primeira instanciação desse objetivo. Experimentos extensivos em 13 conjuntos de dados com comprimentos de contexto, tipos de tarefa e tamanhos de modelo diversos demonstram sua eficácia e eficiência. Nosso código está disponível publicamente em 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.