Resumen Técnico: Un JOLT para el KV Cache

Declaración del Problema

En la inferencia de transformadores autorregresivos, el caché de Clave-Valor (KV cache) representa el costo de memoria dominante, escalando como $O(\text{batch} \times T \times L \times n_h \times d_h)$. Para contextos largos o batches grandes, el KV cache, en lugar de los pesos del modelo, se convierte en el principal cuello de botella de memoria y limita el rendimiento (throughput). Los métodos de compresión existentes se dividen en dos familias, ambas las cuales fallan al no explotar plenamente la redundancia estructural del caché:

1. Métodos de bajo rango (Low-rank): Estos factorizan secciones bidimensionales (por ejemplo, matrices por cabeza o bloques entre capas), pero se comprometen con un único eje de redundancia, ignorando los demás.
2. Métodos de cuantización: Estos reducen el ancho de bits de cada entrada, pero están limitados por tasas de bits fijas (por ejemplo, 4 bits), lo que crea un "suelo de compresión" que impide alcanzar ratios de compresión intermedios (2–3×) donde la calidad suele preservarse.

La observación central que motiva este trabajo es que el KV cache en una capa es un tensor de tercer orden ($n_h \times T \times d_h$) donde los ejes poseen cantidades de redundancia muy distintas. Específicamente, los ejes de cabeza y de capa son mayormente incompresibles, mientras que los ejes de tokens y de características contienen la mayor parte de la estructura compresible. Además, empíricamente los valores son 2–3× más difíciles de comprimir que las claves debido a diferencias en sus propiedades espectrales.

Metodología: JoLT

El método propuesto, JoLT (Joint Tucker and JL-residual allocation), aborda estos desafíos tratando el caché como un tensor de tercer orden y aplicando una estrategia de compresión dirigida.

1. Descomposición de Tucker Parcial

A diferencia de las factorizaciones tensoriales completas que comprimen todos los ejes, JoLT aplica una descomposición de Tucker parcial. Trunca únicamente los ejes de tokens y de características (que portan la redundancia), manteniendo intactos los ejes de cabeza y de capa (que son incompresibles). Esto se logra mediante una SVD de orden superior truncada secuencialmente (ST-HOSVD).

• Motivación: Comprimir los ejes de cabeza o de capa incurre en un error significativo, mientras que truncar los ejes de tokens y de características captura la mayor parte de la redundancia.

2. Residuo Rotado por JL

Un núcleo de bajo rango puro es insuficiente para lograr una fidelidad casi sin pérdidas, particularmente para los valores, que tienen un espectro plano donde la truncación descarta energía significativa. JoLT restaura esta energía descartada utilizando un residuo de pocos bits rotado por Johnson-Lindenstrauss (JL):

• Se calcula el residuo de truncación $R = X - \tilde{X}$.
• Se aplica una rotación aleatoria de JL al residuo.
• El residuo rotado se cuantiza uniformemente a $b \in \{0, 2, 4, 8\}$ bits.
• La decodificación implica invertir la rotación. Este residuo actúa como una palanca para convertir un núcleo con pérdida en uno casi sin pérdidas, lo cual es especialmente valioso para los valores.

3. Asignación Conjunta Lagrangiana

El método emplea un único Lagrangiano dual para asignar recursos de forma conjunta a través de las dos etapas (rangos de Tucker y anchos de bits del residuo) y a través de diferentes objetivos (Claves vs. Valores, Grupos de Capas).

• Objetivo: Minimizar el error de reconstrucción total sujeto a un presupuesto global de bytes.
• Mecanismo: La variable dual permite al asignador desplazar bits entre rangos y residuos, y entre claves y valores, explotando dinámicamente la asimetría en la compresibilidad (por ejemplo, financiando fuertemente el residuo de los valores mientras se reducen los bits del residuo de las claves).
• Granularidad: La asignación se realiza por grupo de capas y de forma separada para Claves y Valores.

4. FlashJoLT (Variante Rápida)

Para abordar el costo computacional de la SVD exacta en contextos largos, FlashJoLT reemplaza la SVD exacta del modo de tokens con una SVD de bajo rango aleatorizada.

• Calcula únicamente las $q_{cap}$ direcciones principales del modo de tokens.
• Emplea un conteo de masa de la cola (tail-mass accounting), añadiendo la cola espectral omitida como un valor singular de residuo sintético para evitar que el asignador subestime la energía capturada por un rango determinado.
• Esta variante logra una aceleración de 5–13× en el tiempo de compresión con una pérdida de calidad insignificante en la "zona libre".

Contribuciones Clave

1. Análisis Espectral: Un estudio empírico que identifica que los ejes de cabeza y capa son incompresibles, mientras que los ejes de tokens y características poseen redundancia, siendo los valores significativamente más difíciles de comprimir que las claves.
2. Algoritmo JoLT: Un compresor novedoso que combina la descomposición de Tucker parcial con un residuo rotado por JL, acoplados mediante un único dual lagrangiano para la asignación conjunta de rangos y bits.
3. FlashJoLT: Una variante de SVD aleatorizada que reduce drásticamente el tiempo de compresión preservando la calidad.
4. Validación Empírica: Un estudio exhaustivo en modelos GQA (Mistral-7B) y MHA (LLaMA-2-13B) que establece una "zona libre" de compresión casi sin pérdidas.

Resultados

La "Zona Libre" (Compresión 2–3×)

El artículo identifica un ratio de compresión de 2–3× como una "zona libre" donde la calidad del modelo se preserva dentro del ruido estadístico de la línea base sin comprimir:

• Perplejidad: Tanto en Mistral-7B (GQA) como en LLaMA-2-13B (MHA), la deriva de la perplejidad es insignificante a 2–3× para longitudes de contexto de hasta 1024.
• Error de Reconstrucción: A una compresión de 2×, JoLT logra un error de Frobenius relativo de 0.009 (Claves) y 0.006 (Valores). Esto es aproximadamente un orden de magnitud inferior que la cuantización de 4 bits y las líneas base de SVD entre capas.
• Tareas de Desempeño (Downstream Tasks):
  • GSM8K: Los modelos comprimidos (2–3×) permanecen dentro del intervalo de confianza del 95% de la línea base.
  • RULER (Needle-in-a-Haystack): Mistral mantiene una precisión de recuperación del 100% a 2–3×; LLaMA iguala su techo de la línea base.

Divergencia de Arquitectura

Más allá de la zona libre de 2–3×, las dos arquitecturas divergen:

• Mistral (GQA): Degrada de forma gradual, perdiendo aproximadamente un 4% de perplejidad por cada paso de ratio entero más allá de 4×.
• LLaMA (MHA): Degrada abruptamente entre 4× y 5×, con la perplejidad aumentando un ~27% a 4× y un ~68% a 5×. Las tareas de razonamiento (GSM8K) para LLaMA colapsan significativamente en ratios altos.

Eficiencia

• Memoria: La huella persistente del KV disminuye linealmente con el ratio de compresión (por ejemplo, de 268.4 MB a 134.6 MB a 2× para Mistral en ctx-2048). El pico de memoria de proceso permanece acotado por los pesos y activaciones.
• Velocidad: FlashJoLT ofrece una aceleración de 5–13× en el tiempo de compresión en comparación con el núcleo (backbone) exacto.

Significancia y Reivindicaciones

El artículo afirma que JoLT proporciona una compresión de 2–3× casi sin pérdidas del KV cache, eliminando efectivamente el cuello de botella de memoria para la inferencia de contexto largo sin sacrificar la calidad del modelo. Al tratar el caché como un tensor y explotar la asimetría espectral específica entre claves y valores, JoLT supera a los métodos existentes de bajo rango y cuantización en el régimen de compresión intermedia.

Los autores posicionan el método como una palanca directa para reducir el costo de servir contextos largos. Señalan dos límites primarios:

1. Límites de Arquitectura: El método es robusto para modelos GQA, pero enfrenta una degradación aguda para modelos MHA más allá de 4×, lo que sugiere que la banda de 2–3× es la recomendación segura y agnóstica al modelo.
2. Costo de Despliegue: Aunque el almacenamiento se reduce, la decodificación requiere actualmente reconstruir el caché desde los factores en cada paso. Los autores identifican el desarrollo de un kernel fusionado de reconstrucción y atención como el siguiente paso crítico para lograr la paridad de inferencia y realizar plenamente el potencial del método.