ReFreeKV: Hacia la Compresión de Caché KV sin Umbrales

Resumen

Para reducir el consumo de memoria durante la inferencia de LLM, se han propuesto diversos métodos para la poda de caché KV. Si bien estas técnicas pueden lograr una reducción de memoria sin pérdidas en muchos conjuntos de datos, a menudo dependen de una condición subestimada: es necesario predeterminar un umbral específico para la entrada o el dominio del presupuesto de caché KV con el fin de alcanzar un rendimiento óptimo. Sin embargo, este diseño sensible a la entrada puede verse considerablemente limitado en escenarios del mundo real, ya que las entradas de dominio abierto abarcan diversos dominios, longitudes y niveles de dificultad, sin límites claros para la selección del umbral. Como resultado, la dependencia de dicho umbral sensible a la entrada puede ser una limitación fundamental que provoca una gran degradación en entradas arbitrarias. En este trabajo, proponemos un nuevo objetivo que elimina las restricciones de umbral para una compresión KV robusta, abogando por métodos "sin umbral" que ajusten adaptativamente la asignación del presupuesto mientras preservan el rendimiento completo de la caché. A continuación, proponemos un método novedoso, ReFreeKV, como la primera instanciación de este objetivo. Experimentos exhaustivos en 13 conjuntos de datos con diversas longitudes de contexto, tipos de tareas y tamaños de modelo demuestran su eficacia y eficiencia. Nuestro código se publica públicamente en 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.