Resumen Técnico: MOMENTKV

Planteamiento del Problema

La decodificación autorregresiva en los modelos de lenguaje basados en Transformers depende del caché de Clave-Valor (KV), que almacena proyecciones computadas previamente para reutilizarlas en los pasos de decodificación. A medida que las longitudes de secuencia escalan a decenas o cientos de miles de tokens, la huella de memoria del caché KV crece linealmente, convirtiéndose en el principal cuello de botella para la inferencia de contexto largo.

La evacuación del caché KV aborda esto reteniendo un subconjunto de tamaño fijo de pares clave-valor y descartando el resto. Aunque los métodos existentes han mejorado las estrategias de selección de tokens (por ejemplo, basadas en pesos de atención acumulados o puntuación consciente del valor), comparten una limitación fundamental: la inferencia post-evacuación renormaliza la atención exclusivamente sobre los tokens retenidos, descartando por completo el conjunto evacuado.

Los autores identifican que la principal fuente de degradación de la salida no es simplemente la masa de atención residual en los tokens evacuados (que los métodos existentes ya minimizan), sino una desajuste direccional entre el conjunto de tokens retenidos y el evacuado. El análisis empírico revela que, en la práctica, los tokens evacuados suelen ser casi ortogonales a los retenidos. En consecuencia, incluso una pequeña cantidad de masa de atención evacuada puede producir un error direccional sustancial que la renormalización estándar no puede recuperar. Esto crea una "brecha direccional" donde mejorar la selección de tokens por sí solo produce rendimientos decrecientes, ya que una retención más selectiva concentra el conjunto retenido en direcciones específicas alineadas con la consulta (query), ampliando aún más la brecha con el conjunto evacuado.

Metodología: MOMENTKV

Para abordar esta brecha direccional, los autores proponen MOMENTKV, un marco que mantiene estadísticas de momentos compactas y de tamaño reducido sobre el conjunto de tokens evacuados. Estas estadísticas sirven para un doble propósito: guiar las decisiones de evacuación y corregir la salida de la atención durante la inferencia.

1. Estadísticas de Momentos Compactas

En lugar de almacenar los tokens evacuados, MOMENTKV mantiene cuatro sumas acumulativas por cabeza de atención, actualizadas incrementalmente tras cada evacuación:

• Conteo ($n_e$): El número de tokens evacuados.
• Suma de Claves ($s_k$): $\sum_{i \in E} k_i$.
• Suma de Valores ($s_v$): $\sum_{i \in E} v_i$.
• Suma de Producto Exterior ($S$): $\sum_{i \in E} v_i k_i^\top$.

Estas estadísticas permiten el cálculo de la media de claves ($\bar{k}$), la media de valores ($\bar{v}$) y la covarianza valor-clave ($\tilde{S}$) sin necesidad de volver a visitar los tokens evacuados. La sobrecarga de almacenamiento total es $O(d^2)$ por cabeza, independiente de la longitud del contexto.

2. Aproximación de Primer Orden de la Salida Evacuada

Durante la inferencia, MOMENTKV deriva una aproximación de primer orden de forma cerrada de la salida de atención evacuada, $\hat{f}_E(q)$. Al recentrar los logits alrededor de la media de las claves y aplicar una expansión de Taylor de primer orden a la función softmax, los autores derivan:
$$\hat{f}_E(q) = \bar{v} + \frac{1}{n_e \sqrt{d}} \tilde{S} q$$

• El término $\bar{v}$ actúa como una base independiente de la consulta, recuperando la dirección media de la sub-salida evacuada.
• El término $\tilde{S} q$ actúa como una corrección adaptativa a la consulta, capturando la variación perpendicular (contenido direccional) que la base por sí sola no logra capturar.

Esta aproximación permite al modelo reconstruir una estimación de la contribución de la salida evacuada perdida, restaurando efectivamente la información direccional que la evacuación estándar oblitera.

3. Evacuación Informada por Momentos

Las estadísticas de momentos también guían el proceso de evacuación mismo. Para cada token retenido $j$, el algoritmo calcula un residuo de momento:
$$r_j = v_j - \bar{v} - \frac{1}{n_e \sqrt{d}} \tilde{S} k_j$$
Este residuo mide cuánta información direccional novedosa aporta el token $j$ más allá de lo que el resumen acumulado ya captura. La puntuación de evacuación se define como el producto del peso de atención y la norma del residuo: $\text{score}(j) = \alpha_j \cdot \|r_j\|$.

• Los tokens con bajos residuos son bien predichos por el modelo de momentos y son priorizados para su evacuación.
• Los tokens con altos residuos (que portan información direccional única) son retenidos incluso si sus pesos de atención son menores.

Esto crea un bucle de refuerzo mutuo: la evacuación informada por momentos mantiene el conjunto evacuado "geométricamente regular" (baja dispersión de logit $\sigma$), lo que estrecha el límite de la aproximación para la corrección de inferencia. A su vez, la corrección precisa proporciona mejores residuos para las futuras decisiones de evacuación.

4. Inferencia Corregida por Normalización

La salida de atención final se reconstruye como una combinación convexa de la salida retenida ($f_R$) y la salida evacuada aproximada ($\hat{f}_E$):
$$\hat{f}(q) de \hat{w}_R f_R(q) + (1 - \hat{w}_R) \hat{f}_E(q)$$
El peso de mezcla $\hat{w}_R$ se estima utilizando una cota inferior de la función de partición evacuada derivada mediante la desigualdad de Jensen. Esto asegura que el modelo tenga en cuenta la magnitud de la masa evacuada mientras corrige su dirección.

Contribuciones Clave

1. Identificación de la Brecha Direccional: El artículo establece que la fuente dominante de error en la evacuación del caché KV es la casi ortogonalidad entre las sub-salidas retenidas y evacuadas, no solo la magnitud de la masa de atención evacuada.
2. Estadísticas de Momentos Compactas: Introducción de un método eficiente en memoria para rastrear las propiedades geométricas (media y covarianza) de los tokens evacuados utilizando solo $O(d^2)$ de almacenamiento por cabeza.
3. Corrección de Forma Cerrada: Derivación de una aproximación de primer orden para la salida de atención evacuada que recupera tanto la dirección media como las variaciones dependientes de la consulta.
4. Mecanismo de Doble Rol: Un marco unificado donde las estadísticas de momentos optimizan simultáneamente la selección de evacuación (identificando tokens redundantes) y corrigen las salidas de inferencia (restando la información direccional perdida).

Resultados Experimentales

Los autores evaluaron MOMENTKV en los benchmarks LongBench y RULER utilizando LLaMA-3.1-8B-Instruct y Qwen3-4B-Instruct.

• Rendimiento: MOMENTKV superó consistentemente a todos los modelos base (H2O, SnapKV, PyramidKV, Ada-KV) en todos los presupuestos de caché ($L \in \{128, 256, 512, 1024\}$).
• Compresión Agresiva: Las mayores ganancias se observaron bajo compresión agresiva. En LongBench con $L=128$ y LLaMA-3.1-8B, MOMENTKV logró una mejora de puntuación promedio de +1.35 sobre el modelo base más fuerte (Ada-KV), reteniendo el 94.1% del rendimiento del caché completo.
• Tareas de Recuperación: En RULER, que pone a prueba la recuperación precisa de tokens, MOMENTKV mostró ganancias proporcionales de +3.4 y +3.5 sobre Ada-KV para LLaMA y Qwen respectivamente, validando su efectividad en entornos intensivos en recuperación.
• Estudios de Ablación:
  • La Inferencia Corregida por Normalización (NC) contribuyó significativamente más (+2.41) que la Evacuación Informada por Momentos (MI) por sí sola (+0.57), confirmando que recuperar la dirección de la evacuación es el principal motor del rendimiento.
  • La combinación de ambos mecanismos produjo un efecto sinérgico (+3.17 de ganancia total), superando la suma de las ganancias individuales.
• Eficiencia: El método añade solo 1.6 ms de latencia por token y requiere 4.1 MB de memoria auxiliar (6.4% del caché retenido), mientras reduce la latencia por token en más de un 65% en comparación con un caché completo.

Significancia

El artículo argumenta que el camino para mejorar la inferencia de contexto largo no reside únicamente en una mejor selección de tokens, sino en recuperar la información direccional que las estrategias de evacuación actuales obliteran. MOMENTKV demuestra que, al mantener estadísticas compactas sobre el conjunto evacuado, es posible cerrar la brecha direccional, transformando el "costo irreducible" de la evacuación en un error de aproximación corregible. Este enfoque ofrece una solución práctica y escalable para desplegar modelos de lenguaje de contexto largo sin sacrificar la eficiencia de los cachés KV de presupuesto fijo.