Resumen Técnico: PolyKV: Retención y Asignación Heterogénea para la Compresión del Caché KV

Declaración del Problema

La inferencia de contexto largo en los Modelos de Lenguaje de Gran Escala (LLMs) enfrenta un cuello de botella significativo en la memoria debido al crecimiento del caché de Clave-Valor (KV). Aunque los métodos existentes de evacuación de caché KV reducen eficazmente el uso de memoria reteniendo subconjuntos específicos de tokens, típicamente aplican una política de compresión uniforme y un presupuesto de caché uniforme en todas las capas del transformer.

Este diseño uniforme pasa por alto la heterogeneidad inherente de las capas del transformer. Investigaciones previas indican que diferentes capas varían en sus patrones de atención, abstracciones representadas, roles contextuales y sensibilidad a la poda. En consecuencia, una única regla de evacuación global es una abstracción tosca que no tiene en cuenta que:

1. Diferentes capas pueden preferir diferentes patrones de evacuación (por ejemplo, heavy hitters frente a attention sinks).
2. Diferentes capas exhiben variaciones en su sensibilidad a la reducción del caché KV, lo que implica que requieren diferentes capacidades de caché.

Los métodos adaptativos actuales han comenzado a explorar presupuestos no uniformes, pero a menudo fallan en tratar el patrón de evacuación y el presupuesto de caché como decisiones de diseño acopladas a nivel de capa.

Metodología: PolyKV

Los autores proponen PolyKV, un marco que reformula la evacuación del caché KV como un espacio de diseño por capa. PolyKV acopla explícitamente dos decisiones para cada capa y fase de inferencia (prefill y decoding):

1. Qué evacuar: Seleccionar un patrón de evacuación específico (por ejemplo, H2O, TOVA, PyramidKV) adecuado para esa capa.
2. Cuánto retener: Asignar un presupuesto de caché específico a esa capa.

Componentes Principales

• Formulación del Espacio de Diseño: PolyKV define una estrategia $\Pi$ donde a cada capa $l$ se le asigna un patrón de evacuación $r_{l,\phi}$ y un presupuesto $b_{l,\phi}$ para cada fase $\phi$. El presupuesto total está restringido por un límite global fijo $B$.
• Calibración Offline: Para evitar una búsqueda exhaustiva en línea, PolyKV utiliza un pequeño conjunto de calibración para generar una estrategia estática fija. Recopila dos categorías de señales específicas de la capa:
  • Señales de Preferencia de Evacuación: Mide qué patrón de evacuación preserva mejor la computación FullKV de una capa. Las métricas incluyen la similitud de coseno entre las activaciones perturbadas y FullKV, y la sensibilidad de la Perplejidad (PPL) (degradación causada por un patrón específico).
  • Señales de Sensibilidad del Presupuesto de Caché: Estima con qué fuerza se ve afectada una capa por la reducción de KV. Las métricas incluyen la entropía de atención (la atención difusa requiere presupuestos más grandes) y la sensibilidad de PPL (degradación agregada del peor caso).
• Construcción de la Estrategia:
  • Selección de Patrones: Para cada par capa-fase, se selecciona el patrón de evacuación con la puntuación de preferencia más alta (por ejemplo, la mayor similitud de coseno o la menor degradación de PPL).
  • Asignación de Presupuesto: Los presupuestos se asignan proporcionalmente según las señales de sensibilidad. Las capas con mayor sensibilidad (o entropía) reciben presupuestos más grandes, sujetos a una restricción de presupuesto mínimo ($b_{min}$).
• Implementación: PolyKV se implementa como un marco basado en hooks. Inyecta lógica en puntos comunes del flujo de generación (pre-kernel, post-kernel, actualización de caché) para permitir la integración modular de varios métodos de evacuación sin modificar el código subyacente del modelo.

Contribuciones Clave

1. Reformulación de la Evacuación KV: El artículo redefine la evacuación del caché KV como un espacio de diseño por capa que acopla la elección del patrón de evacuación con la asignación de capacidad de caché bajo un presupuesto de memoria compartido.
2. El Marco PolyKV: Un sistema que operacionaliza este espacio de diseño componiendo patrones de evacuación y presupuestos a nivel de capa. Integra múltiples primitivas de evacuación mediante una interfaz unificada y utiliza señales de calibración offline para guiar la construcción de la estrategia.
3. Validación Empírica: Los autores demuestran que las estrategias heterogéneas y por capa superan a las líneas base de política única y uniforme bajo presupuestos de memoria fijos.

Resultados Experimentales

Los experimentos se realizaron en LLaMA-3.1-8B y Qwen3-8B utilizando los benchmarks LongBench y RULER.

• Desempeño en LongBench (presupuesto promedio de 512 tokens):
  • LLaMA-3.1-8B: PolyKV recuperó el 54.5% de la brecha de rendimiento entre la línea base de política única más fuerte y FullKV.
  • Qwen3-8B: PolyKV recuperó el 25.7% de la brecha.
  • PolyKV superó consistentemente a la línea base más fuerte en un barrido de presupuesto (128–1024 tokens), con mejoras que oscilaron entre el 1.7% y el 6.4%.
• Desempeño en RULER:
  • Los resultados fueron mixtos. PolyKV superó ligeramente a la línea base en LLaMA, pero tuvo un desempeño inferior en Qwen (por 3.98 puntos). Los autores atribuyen esto a la naturaleza de las tareas de RULER, que dependen de tokens de evidencia dispersos que pueden ser descartados por las políticas heterogéneas si las señales de calibración no los identifican perfectamente.
• Estudios de Ablación:
  • Patrón de Evacuación vs. Presupuesto: La principal ganancia de rendimiento proviene de asignar diferentes patrones de evacuación a diferentes capas, particularmente durante la fase de prefill.
  • Asignación de Presupuesto: La asignación de presupuesto no uniforme proporciona ganancias adicionales pero es más dependiente del modelo y menos robusta que la selección de patrones. Es más efectiva cuando el presupuesto total es suficiente para evitar la inanición de las capas sensibles.
• Análisis de Tareas: Las ganancias fueron más pronunciadas en tareas que requieren una comprensión contextual amplia (por ejemplo, QA, aprendizaje de pocos disparos/few-shot). El desempeño se degradó en tareas que requieren la recuperación exacta de tokens dispersos (por ejemplo, NIAH en RULER) o la adherencia estricta a la forma local (por ejemplo, generación de código).

Significancia y Reivindicaciones

El artículo sostiene que tratar la evacuación del caché KV como un problema de optimización composicional y consciente de las capas, en lugar de una única política de retención global, produce resultados superiores.

• La Heterogeneidad es la Clave: El estudio valida que las capas del transformer no son homogéneas; una estrategia de evacuación de "talla única" es suboptimal.
• Eficiencia Estática: Al utilizar la calibración offline, PolyKV logra estos beneficios heterogéneos sin la sobrecarga de tiempo de ejecución de una búsqueda de política en línea, lo que lo hace práctico para el despliegue.
• Limitaciones: Los autores señalan modestamente que las señales de calibración actuales no siempre preservan los tokens de evidencia "dispersos pero decisivos" requeridos para tareas de recuperación exacta, lo que indica una dirección para el refinamiento futuro en el diseño de señales.

En resumen, PolyKV demuestra que al alinear las estrategias de evacuación y las capacidades de caché con las características específicas de cada capa del transformer, es posible reducir significamente la compensación entre memoria y calidad en la inferencia de LLMs de contexto largo.