Sequential KV Cache Compression via Probabilistic Language Tries: Beyond the Per-Vector Shannon Limit

Resumen Técnico: Compresión Secuencial de KV Cache mediante Árboles de Lenguaje Probabilísticos

1. El Problema: El Cuello de Botella de la Memoria en la Inferencia

En la inferencia de modelos de lenguaje transformer, el KV Cache (caché de claves y valores) actúa como la memoria de trabajo del modelo. Para un contexto largo, este caché consume una cantidad masiva de memoria, a menudo superando el tamaño de los propios pesos del modelo (ej. ~80 GB para un contexto de 128k tokens en un modelo de 70B parámetros).

• Estado del Arte Actual: Métodos recientes como TurboQuant han alcanzado el límite de Shannon para la compresión por vector. TurboQuant logra comprimir cada componente de un vector KV a aproximadamente 3 bits mediante rotación polar y corrección de residuos, demostrando que no se puede comprimir más sin perder precisión si se tratan los vectores como muestras independientes.
• La Limitación Fundamental: El enfoque actual trata los vectores KV como datos flotantes aislados. Sin embargo, el KV cache es una secuencia generada por un modelo que predice un lenguaje estructurado. Los vectores no son aleatorios; son funciones deterministas de una secuencia de tokens que sigue una distribución de probabilidad conocida. Ignorar la correlación secuencial entre tokens (y por tanto, entre vectores KV) deja sin explotar una redundancia masiva.

2. Metodología: Una Arquitectura de Dos Capas

Los autores proponen la Compresión KV Secuencial, una arquitectura que explota la estructura del lenguaje para superar el límite de compresión por vector. Se basa en dos capas ortogonales:

Capa 1: Deducuplicación de Prefijos Probabilísticos (Probabilistic Prefix Deduplication)

• Concepto: En lugar de buscar coincidencias exactas de tokens (como hacen vLLM o SGLang), el método utiliza la métrica de árbol de lenguaje probabilístico (PLT) definida como $d_T(s, s') = -\log_2 P_M(s \land s')$.
• Mecanismo: Agrupa sesiones que comparten prefijos "semánticamente equivalentes" (alta probabilidad de coincidencia en el espacio de probabilidad del modelo), incluso si los tokens no son idénticos lexicalmente.
• Almacenamiento: Se almacena el caché KV completo para el "centroide" (la secuencia más probable del grupo) y solo los residuos (deltas) para las otras sesiones a partir del punto de divergencia.
• Ventaja: Elimina redundancia inter-sesión que la coincidencia exacta pierde, aprovechando que prefijos semánticamente similares generan vectores KV muy cercanos.

Capa 2: Codificación Delta Predictiva (Predictive Delta Coding)

• Concepto: Dentro de una sesión, el vector KV en la posición $t$ es altamente predecible dado el contexto $t-1$.
• Mecanismo: En lugar de almacenar el vector KV $KV_t$, el sistema calcula el vector KV predicho ($\hat{KV}_t$) basado en la distribución de probabilidad del modelo sobre el siguiente token. Luego, almacena solo el residuo $R_t = KV_t - \hat{KV}_t$.
• Cuantización Adaptativa: La magnitud del residuo es pequeña cuando el modelo tiene alta certeza (baja entropía/surprisal del token). Se utiliza una cuantización adaptativa que asigna menos bits a los residuos pequeños (tokens predecibles) y más a los grandes, en lugar de usar una precisión fija por componente.

3. Contribuciones Clave

1. El Límite de Entropía Secuencial (Teorema 1):

   • Se demuestra formalmente que la entropía condicional de un vector KV dado el historial de caché está acotada por la sorpresa (surprisal) por token del modelo:
     $$H(KV_t | KV_{<t}) \leq H(token_t | token_{<t})$$
   • Esto implica que la información necesaria para codificar el KV cache no es la entropía del vector flotante (que es alta), sino la incertidumbre del siguiente token (que es baja en textos coherentes).

2. Acotación Teórica de Compresión:

   • Para un modelo con una perplejidad típica de 10-20 (inglés fluido), el límite de entropía es de 3.3 a 4.3 bits por posición de token (para todo el estado KV de todas las capas y cabezas).
   • En comparación, TurboQuant requiere ~3 bits por componente (donde un token tiene $2 \times L \times H_{head} \times d \approx 10^6$ componentes).
   • Ratio Teórico: La compresión secuencial ofrece un ratio teórico de ~914,000x superior a TurboQuant en el límite de Shannon. Incluso con sobrecargas prácticas pesimistas (1000x sobre el límite de entropía), el ratio sigue siendo ~914x.

3. Composibilidad y Ortogonalidad:

   • Las dos capas propuestas son ortogonales a los métodos de cuantización por vector. Se pueden apilar debajo de TurboQuant: primero se deduplican prefijos, luego se codifican deltas predictivos, y finalmente los residuos se cuantizan con TurboQuant.

4. Comportamiento Asintótico:

   • A diferencia de los métodos actuales donde el costo de memoria crece linealmente con la longitud del contexto, la compresión secuencial mejora a medida que crece el contexto. Cuantos más tokens procesa el modelo, más predecible es el siguiente, reduciendo la entropía marginal por token.

4. Resultados y Análisis

• Eficiencia Teórica: El artículo establece que el "suelo" de Shannon para la compresión de KV cache es mucho más bajo de lo que se creía, porque la redundancia del lenguaje (estructura predecible) no se explota en los enfoques actuales.
• Implementación Práctica:
  • Se propone el uso de aproximaciones de Top-k para calcular el vector KV predicho sin realizar una pasada completa del modelo para cada token del vocabulario.
  • La integración con el bucle de inferencia es viable, ya que la distribución de probabilidad (softmax) ya se calcula durante la inferencia estándar.
• Relación con la Codificación de Especulación: El método comparte fundamentos con la codificación de especulación (Speculative Decoding); la distribución de los tokens candidatos para especulación puede usarse simultáneamente para calcular la media del KV predicho sin costo adicional.

5. Significado e Implicaciones

• Ruptura del "Muro de Memoria": Este trabajo sugiere que el problema del crecimiento lineal de la memoria en contextos largos no es solo un problema de hardware, sino de software. La compresión secuencial convierte el costo marginal de un token adicional en una función decreciente, haciendo viables aplicaciones de contextos extremadamente largos (millones de tokens) que actualmente son prohibitivas.
• Nueva Perspectiva Teórica: Demuestra que el límite de compresión de TurboQuant es correcto para su subproblema (vectores aislados), pero que el problema real (secuencia de vectores) tiene un límite de Shannon estrictamente inferior.
• Paradoja de Jevons: Al reducir drásticamente el costo marginal de la inferencia de larga duración, esta tecnología podría impulsar la demanda de aplicaciones más ambiciosas (agentes persistentes, análisis de documentos masivos), aumentando la demanda total de memoria a pesar de la mayor eficiencia por token.

Conclusión:
El artículo propone un cambio de paradigma: dejar de ver el KV cache como un conjunto de vectores flotantes independientes y tratarlo como una secuencia de datos generada por un lenguaje formal. Mediante la deduplicación probabilística de prefijos y la codificación delta predictiva, es posible alcanzar ratios de compresión órdenes de magnitud superiores a los métodos actuales, transformando la viabilidad de la inferencia de larga duración.