Resumen Técnico: COMPILERKV

1. Planteamiento del Problema

Los Modelos de Lenguaje Grande (LLM) enfrentan un cuello de botella de memoria debido al crecimiento lineal de la caché de Clave-Valor (KV) con la longitud del contexto. La compresión de KV solo en prellenado ha surgido como una solución amigable para la implementación, donde un subconjunto de tokens se selecciona al final de la etapa de prellenado y se congela durante todo el paso de decodificación. Este enfoque convierte el uso de memoria en tiempo de ejecución en un costo fijo con cero intervención en tiempo de decodificación.

Sin embargo, el artículo identifica una discrepancia estadística fundamental en los métodos existentes de solo prellenado (por ejemplo, SnapKV, PyramidKV, HeadKV). Estos métodos intentan estimar señales correctivas —específicamente la fiabilidad por cabeza (qué cabezas de atención son críticas para la recuperación) y la sensibilidad de compresión a nivel de prompt (cuánto puede comprimirse un prompt específico)— en línea, utilizando estadísticas de un único prompt ruidoso. Los autores argumentan que esta es la unidad estadística incorrecta:

• E1 (Ruido de prompt único): Las puntuaciones de atención por prompt tienen colas pesadas y están sesgadas por la escala; grandes magnitudes no garantizan importancia semántica, y pequeños errores de estimación causan errores de selección irreversibles.
• E2 (Heterogeneidad de cabezas): Las cabezas de atención están especializadas funcionalmente. Algunas son "críticas para la recuperación", mientras que otras son consistentemente ruidosas. Una agregación ingenua permite que las cabezas ruidosas dominen las puntuaciones.
• E3 (Varianza de riesgo a nivel de prompt): Diferentes prompts tienen diferentes límites de compresión. Una regla de retención fija falla en prompts de alta entropía o alta perplejidad, lo que lleva a una pérdida catastrófica de información.

Los métodos existentes reestiman estas señales en línea para cada prompt, fallando en aprovechar el hecho de que estas señales exhiben una alta regularidad entre prompts en lugar de estabilidad dentro del prompt.

2. Metodología: COMPILERKV

Los autores proponen COMPILERKV, una política de retención de KV que compila tablas correctivas fuera de línea a partir de un corpus de calibración, en lugar de estimarlas en línea. La idea central es que las clasificaciones de fiabilidad de las cabezas y la asignación de características de riesgo a umbrales seguros son priors arquitectónicos estables que pueden aprenderse una vez y reutilizarse.

El marco consta de tres componentes acoplados que operan durante la etapa de prellenado:

A. Utilidad de Token Estabilizada (Abordando E1)

Para eliminar picos de atención transitorios y el sesgo de escala entre cabezas sin aprendizaje, el método calcula una puntuación de utilidad libre de parámetros:

• Atención agregada por ventana: Suaviza los picos transitorios agregando la atención sobre una ventana de observación $\Omega$ en la cola del prompt.
• Normalización relativa a la cabeza: Normaliza las puntuaciones de los tokens por la norma del valor medio de la cabeza, eliminando el sesgo de escala.
• Fórmula: $u^{(l,h)}_t = \alpha_t \cdot \rho^{(l,h)}_t$, donde $\alpha_t$ es la masa de la ventana y $\rho$ es la norma del valor normalizado.

B. Tabla de Heterogeneidad de Cabezas (Abordando E2)

En lugar de tratar todas las cabezas por igual o reestimar la importancia en línea, COMPILERKV compila una tabla congelada $W_{head}[l, h]$ de pesos continuos de fiabilidad.

• Compilación fuera de línea: Aprendida mediante Aprendizaje por Refuerzo Q Conservador (CQL) en un corpus de calibración. El estado es $(l, h, B_l)$, y la acción es un peso de fiabilidad discretizado.
• Mecanismo: Durante la inferencia, la utilidad de un token se reescala por la fiabilidad de su cabeza. Crucialmente, la agregación utiliza max-pooling ponderado ($\hat{u}^{(l)}_t = \max_h (u^{(l,h)}_t \cdot W_{head}[l, h])$). Esto otorga "poder de veto" a las cabezas críticas para la recuperación: si una sola cabeza crítica atiende fuertemente a un token, este se retiene, independientemente del consenso de las cabezas ruidosas.

C. Puerta de Umbral Adaptativa al Riesgo (Abordando E3)

Para manejar prompts con límites de compresión variables, una Puerta de Umbral Adaptativa al Riesgo mapea características de riesgo a nivel de prompt a un umbral de retención.

• Señales de Riesgo: Utiliza Entropía de Atención (riesgo estructural) y Perplejidad Local (riesgo semántico) calculadas sobre la ventana de observación.
• Compilación fuera de línea: Una tabla de búsqueda (LUT) mapea los bins de riesgo discretizados $(b_{ent}, b_{ppl})$ a un umbral de retención $\tau^{(l)}$. Esto también se compila mediante CQL.
• Sesgo Conservador: El objetivo de CQL penaliza los umbrales agresivos en estados raros y de alto riesgo (alta entropía/perplejidad), sesgando la tabla hacia una retención más segura donde las heurísticas en línea suelen fallar.
• Selección Elástica: La selección final retiene todos los tokens con utilidad por encima del umbral $\tau^{(l)}$. Si el conjunto candidato excede el presupuesto $B_l$, se limita a los tokens superiores-$B_l$. Si está por debajo, se mantienen todos los candidatos (sub-retención elástica).

D. Fundamentación Teórica

El artículo proporciona un análisis formal que muestra que el error de aproximación está controlado por la masa de la cola descartada (Proposición 4.1). Argumenta que el error de calibración de tablas finitas se controla a nivel de prompt (a través de prompts independientes) en lugar de asumir independencia dentro de las filas de atención de un único prompt, reduciendo así la varianza.

3. Contribuciones Clave

1. Reencuadre del Problema: Identifica que la unidad estadística para la corrección en la compresión de solo prellenado debe ser la escala entre prompts, no el prompt único.
2. Compilación Fuera de Línea: Introduce un método para compilar la fiabilidad de las cabezas y los umbrales adaptativos al riesgo fuera de línea, convirtiéndolos en priors arquitectónicos portátiles.
3. Política Acoplada: Propone una optimización conjunta donde los pesos de fiabilidad de las cabezas y los umbrales de riesgo se calibran contra la misma recompensa aguas abajo, en lugar de ser módulos independientes.
4. Portabilidad: Demuestra que las tablas compiladas se transfieren entre corpus disjuntos e incluso entre diferentes familias de modelos (por ejemplo, de LLaMA a Mistral) con una degradación mínima del rendimiento.

4. Resultados Experimentales

El método fue evaluado en cuatro backbones de 7B–8B (LLaMA-3, Mistral, Qwen2, InternLM) en los benchmarks LongBench, RULER y Needle-in-a-Haystack (NIAH).

• LongBench (presupuesto de 512 tokens): COMPILERKV logra un SOTA comprimido en los cuatro backbones, superando a la línea base más fuerte (por ejemplo, ChunkKV, HeadKV) en un promedio de +1.67 puntos (IC 95% [+1.08, +2.37]). En algunos backbones, supera ligeramente el rendimiento de FullKV, lo que los autores atribuyen a la eliminación de ruido en lugar de una ganancia de información.
• Régimen de Presión:
  • RULER (contexto de 128k): Bajo una relación fija de caché de 512/32k, COMPILERKV mantiene una puntuación de 73, superando significativamente a SnapKV (38) y KVZip (66.5), mientras se mantiene por debajo de FullKV (79).
  • NIAH (32k): Logra una precisión de 0.89 frente a 0.42 de SnapKV.
  • Agrupamiento (Batching): Con una entrada de 32k y un tamaño de lote de 16, COMPILERKV sigue siendo viable, mientras que FullKV se queda sin memoria (OOM).
• Transferibilidad:
  • Entre Corpus: Las clasificaciones de fiabilidad de las cabezas muestran una alta correlación de Spearman ($\bar{\rho} \approx 0.90$) entre corpus de calibración disjuntos.
  • Entre Modelos: La transferencia directa de tablas entre modelos (por ejemplo, de LLaMA-3 a Mistral) incurre en una pérdida de solo 0.4–0.8 puntos en LongBench en comparación con la compilación específica del objetivo.
• Eficiencia: La sobrecarga en línea es mínima: $O(1)$ búsquedas en tabla por capa y una única ordenación sobre el conjunto candidato. Logra hasta 2.45× el rendimiento de decodificación en comparación con FullKV en un contexto de 32k.

5. Significado y Afirmaciones

El artículo afirma que COMPILERKV cierra la brecha con FullKV en presupuestos fijos al corregir la discrepancia estadística en cómo se calculan las políticas de retención. Al desplazar el cálculo de señales correctivas de una estimación en línea, ruidosa y de prompt único, a un prior fuera de línea, compilado y entre prompts, el método estabiliza el proceso de selección frente al ruido y la heterogeneidad.

Los autores enfatizan que la novedad no es meramente la "conciencia de las cabezas" (que poseen métodos anteriores como HeadKV), sino la compilación conjunta fuera de línea de la fiabilidad continua de las cabezas y los umbrales de riesgo del prompt. Esto transforma la fiabilidad de las cabezas de un truco de ajuste en un objeto científico reutilizable que puede compilarse una vez y desplegarse en modelos y contextos con una sobrecarga en línea de $O(1)$. El trabajo sugiere que para la compresión de solo prellenado, las señales "correctas" son propiedades arquitectónicas del modelo, no características transitorias del prompt de entrada.