Tensor Cache: Eviction-conditioned Associative Memory for Transformers

Resumen Técnico: Tensor Cache

Enunciado del Problema
La inferencia de Transformadores autoregresivos depende del almacenamiento en caché de pares Clave-Valor (KV) para evitar recalcular prefijos. Sin embargo, el estado KV retenido crece linealmente con la longitud del contexto, la profundidad y la concurrencia, creando un cuello de botella de memoria. El almacenamiento en caché de ventana deslizante limita el uso de memoria reteniendo solo los $W$ tokens más recientes, pero este enfoque sufre de "olvido forzado": una vez que un token es expulsado de la ventana, su información se descarta por completo, haciendo inaccesible la evidencia relevante fuera de la ventana. Los métodos de estado acotado existentes o bien descartan por completo los tokens expulsados o intentan comprimir cada token en un promedio móvil, lo que puede introducir ruido o fallar en distinguir eficazmente entre el contexto reciente y el distante.

Metodología: Tensor Cache (TC)
Tensor Cache introduce una arquitectura de estado acotado de dos niveles que transforma la expulsión de KV de un evento de eliminación en una operación de escritura de memoria asociativa.

1. Arquitectura de Dos Niveles:

   • L1 (Caché de Primer Nivel): Un búfer circular de tamaño fijo de KV que contiene los $W$ tokens más recientes. Esto soporta atención softmax exacta para el contexto local.
   • L2 (Caché de Segundo Nivel): Una matriz de producto externo de tamaño fijo $A$ (memoria de pesos rápidos) que almacena información exclusivamente de los tokens expulsados de la ventana L1.
   • Mecanismo de Escritura: Cuando el búfer circular sobrescribe una entrada, el par $(k, v)$ desplazado no se descarta. En su lugar, se escribe en la matriz L2 $A$ mediante una actualización de producto externo: $A \leftarrow \lambda A + \eta (k \otimes v)$. Aquí, $\lambda$ (decaimiento) y $\eta$ (tasa de escritura) son escalares aprendidos por cabeza.
   • Mecanismo de Lectura: Las consultas futuras $q_t$ acceden al pasado comprimido calculando $r_t = q_t A$. Aprovechando la identidad de atención lineal $q(k \otimes v) = \langle q, k \rangle v$, esta operación recupera efectivamente información de todos los tokens expulsados sin almacenarlos individualmente.
   • Fusión: Una puerta escalar aprendida $g$ fusiona las salidas de la atención local (L1) y la lectura de memoria (L2): $y_t = y_{local} + \sigma(g) m_t$.

2. Optimización del Entrenamiento (Escaneo Paralelo):
   El artículo identifica un problema crítico con los atajos de entrenamiento por bloques estándar. Un enfoque ingenuo resume un bloque por su media ($\bar{k} \otimes \bar{v}$), lo que introduce $C^2 - C$ productos externos espurios entre tokens ($k_i \otimes v_j$ donde $i \neq j$) que no existen durante la inferencia en flujo.

   • Solución: Los autores implementan un escaneo ponderado paralelo que es matemáticamente equivalente a las escrituras secuenciales por token. Esto cierra la brecha entre el entrenamiento y la inferencia dentro de la epsilon de punto flotante 32 bits (error relativo < $10^{-7}$), asegurando que el modelo aprenda la dinámica asociativa correcta sin el ruido de la interferencia entre tokens.

Contribuciones Clave

• Memoria Asociativa Condicionada a la Expulsión: A diferencia de las memorias de producto externo anteriores (por ejemplo, Infini-attention, mLSTM) que ingieren cada token, TC utiliza la matriz de producto externo estrictamente como una caché L2 alimentada solo por expulsiones de ventana deslizante. Esto particiona limpiamente los tokens en "residentes en L1" y "comprimidos en L2".
• Corrección de Artefactos de Entrenamiento: El artículo identifica y resuelve un atajo de entrenamiento previamente pasado por alto (media por bloques) que introduce productos espurios entre tokens. El escaneo ponderado paralelo propuesto elimina esta discrepancia al mismo costo computacional.
• Eficiencia de Estado Acotado: TC proporciona un mecanismo donde el contexto más antiguo influye en las predicciones sin retener entradas expulsadas individuales, ofreciendo una alternativa práctica a la retención completa de KV.

Resultados Experimentales
Los autores evaluaron TC en sistemas de escalado, recuperación asociativa sintética y modelado de lenguaje de contexto largo con texto real (OpenWebText y Shakespeare).

• Recuperación Asociativa Sintética: En tareas que requieren recuperar información más allá de la ventana deslizante, TC logró una precisión del 100%, igualando el rendimiento de KV completo. Redujo el estado de inferencia retenido en un 72–84% en comparación con KV completo, superando significativamente a Infini-attention (que se degradaba a medida que aumentaban las brechas) y a las líneas base de ventana deslizante (que funcionaban a niveles de azar).
• Modelado de Lenguaje de Contexto Largo: En OpenWebText, TC logró la menor Negativa Log-Verosimilitud (NLL) media en cada longitud de contexto de evaluación desde 1,024 hasta 32,768 tokens, con la única excepción de $L=2,048$, donde fue estadísticamente indistinguible de Infini-attention. En $L=32,768$, TC alcanzó una NLL de 5.14, superando a KV completo (6.00) e Infini-attention (5.42), mientras utilizaba 2.4$\times$ menos memoria pico de GPU que KV completo.
• Rendimiento: TC incurrió en una sobrecarga de rendimiento de ~30% en comparación con los métodos de ventana deslizante puros (debido a la lectura/actualización de L2), pero permaneció consistentemente más rápido que Infini-attention en todas las longitudes de contexto probadas.
• Diagnósticos de Capacidad: Los experimentos sobre la capacidad bruta de memoria asociativa mostraron que el rendimiento de recuperación se degrada a medida que se comprimen más asociaciones no relacionadas en el estado fijo, confirmando la naturaleza acotada de la memoria.

Significado y Afirmaciones
El artículo afirma que Tensor Cache cierra exitosamente la brecha entre la eficiencia de memoria de la atención de ventana deslizante y las capacidades de recuperación de la atención de contexto completo. Al tratar la expulsión de KV como una operación de memoria estructurada en lugar de una eliminación destructiva, TC permite que los Transformadores accedan a evidencia relevante fuera de la ventana inmediata sin el costo de memoria lineal de las cachés KV completas. Los autores enfatizan que su contribución principal no es la invención de la memoria de producto externo en sí misma, sino su aplicación específica como una caché L2 condicionada a la expulsión y la corrección rigurosa de los artefactos de entrenamiento que anteriormente obstaculizaban tales métodos. TC se presenta como una alternativa práctica de estado acotado que mejora la frontera de calidad-memoria para la inferencia de contexto largo.