ReFreeKV : Vers une compression du cache KV sans seuil

Résumé

Pour réduire la consommation mémoire lors de l'inférence des grands modèles de langage, plusieurs méthodes ont été proposées pour l'élagage du cache KV. Bien que ces techniques puissent réaliser une réduction mémoire sans perte sur de nombreux jeux de données, elles reposent souvent sur une condition sous-estimée : un seuil spécifique à l’entrée/au domaine pour le budget du cache KV doit être prédéfini afin d’atteindre la performance optimale. Cependant, une telle conception sensible aux entrées peut être considérablement limitée dans des scénarios réels, car les entrées de domaines ouverts couvrent des domaines, des longueurs et des niveaux de difficulté variés, sans limites claires pour la sélection des seuils. Par conséquent, la dépendance à ce seuil sensible aux entrées peut constituer une limitation fondamentale entraînant une dégradation importante sur des entrées arbitraires. Dans ce travail, nous proposons un nouvel objectif qui lève les contraintes de seuil pour une compression robuste du cache KV, prônant des méthodes « sans seuil » qui ajustent de manière adaptative l’allocation du budget tout en préservant la performance du cache complet. Nous proposons ensuite une nouvelle méthode, ReFreeKV, comme première instanciation de cet objectif. Des expériences approfondies sur 13 jeux de données avec des longueurs de contexte, des types de tâches et des tailles de modèle variés démontrent son efficacité et son efficience. Notre code est publié publiquement à l’adresse 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.