Resumen Técnico: KV-Fold

Enunciado del Problema

Los transformadores grandes preentrenados están limitados por el costo cuadrático de memoria de la atención completa y la ventana de contexto fija de una sola pasada hacia adelante. Aunque los modelos recientes soportan ventanas de contexto de hasta cientos de miles de tokens, las cargas de trabajo del mundo real (por ejemplo, código a escala de repositorio, registros médicos longitudinales) a menudo exceden estos límites. Las soluciones existentes enfrentan un compromiso persistente:

• Métodos de transmisión (por ejemplo, StreamingLLM) mantienen una memoria acotada reteniendo solo los tokens recientes y los "sumideros de atención", pero pierden la capacidad de recuperar información una vez que esta cae fuera de la ventana deslizante.
• La compresión de la caché KV y los métodos de memoria recurrente a menudo requieren modificación del modelo, ajuste fino o introducen aproximaciones con pérdida que degradan la fidelidad de recuperación.
• La atención completa es computacionalmente inviable para secuencias extremadamente largas en hardware estándar debido al tamaño masivo de la matriz de puntuaciones de atención (por ejemplo, ~1 TB para una secuencia de 128K en una GPU de 40 GB).

El artículo pregunta si un transformador preentrenado congelado puede mantener su estado interno a través de fragmentos sucesivos para realizar inferencia de contexto largo como una recurrencia, sin cambios arquitectónicos ni entrenamiento.

Metodología: KV-Fold

Los autores proponen KV-Fold, un protocolo de inferencia sin entrenamiento que trata la caché de Clave-Valor (KV) como el acumulador en un pliegue izquierdo (foldl) sobre fragmentos de secuencia.

• Mecanismo: Una secuencia larga de longitud $T$ se divide en fragmentos de longitud $C$. Al procesar el fragmento $t$, el modelo realiza una pasada hacia adelante donde:
  1. Las consultas ($Q$) se generan a partir de los tokens del fragmento actual.
  2. Las claves ($K$) y los valores ($V$) se forman concatenando la caché KV acumulada de todos los fragmentos anteriores ($0$a$t-1$) con los estados KV del fragmento actual.
  3. Los embebidos posicionales (RoPE) se aplican continuamente a través de los límites de los fragmentos para mantener la alineación absoluta de la secuencia.
  4. La caché KV ampliada se pasa hacia adelante al fragmento $t+1$.
• Estado: Los parámetros del modelo permanecen congelados. No se introducen tokens de memoria especiales ni ajuste fino. La caché KV crece linealmente con la longitud de la secuencia, actuando como un registro estructurado y sin compresión del cómputo pasado del modelo.
• Analogía: El proceso es análogo a un pliegue izquierdo de programación funcional, donde la caché KV sirve como el estado acumulador actualizado por una función de recurrencia de un solo paso $F_\theta$ en cada límite de fragmento.

Contribuciones Clave

1. Protocolo KV-Fold: Un cambio operativo simple que convierte un transformador estándar en un modelo recurrente de contexto largo al transportar la caché KV acumulada a través de los fragmentos.
2. Estabilidad de la Recurrencia: Evidencia empírica que muestra que la recurrencia inducida no sufre de acumulación de errores ilimitada. En cambio, la "deriva" (diferencia en la Probabilidad Logarítmica Negativa en comparación con la atención completa) aumenta brevemente y luego se satura en una meseta plana que persiste a través de cadenas profundas.
3. Deriva Estructural vs. Numérica: Se demuestra que la meseta es una propiedad estructural de la dinámica de atención en lugar de un artefacto numérico. Permanece estable bajo un aumento de 10,000 veces en la precisión numérica (de bf16 a fp32) y es robusta en diferentes familias de modelos (OLMoE, Qwen, Llama) y tamaños de fragmentos.
4. Recuperación Exacta de Largo Alcance: El método preserva información factual exacta a largas distancias. En las pruebas de "aguja en un pajar", KV-Fold logra una recuperación de coincidencia exacta del 100% en contextos de hasta 128K tokens y profundidades de cadena de hasta 511, igualando el rendimiento de la atención completa.
5. Compromisos Prácticos: El artículo caracteriza la posición del método en el espectro memoria-versus-recuperación. Aunque requiere un crecimiento lineal de memoria (a diferencia de la transmisión de memoria acotada), permite la recuperación a profundidades arbitrarias donde los métodos de transmisión fallan, haciendo viable ejecutar una recuperación equivalente a la atención completa en hardware que no puede soportar una sola pasada hacia adelante monolítica.

Resultados Clave

• Saturación de la Deriva: En Qwen2.5-7B-Instruct, la deriva se satura en la profundidad 7 y permanece plana hasta la profundidad 63, con un cambio total de solo -0.0003 nats entre las profundidades 15 y 60.
• Rendimiento de Recuperación:
  • Aguja Única: 100% de precisión de recuperación en 152 ensayos que abarcan contextos de 16K a 128K tokens y profundidades de cadena de hasta 511 en Llama-3.1-8B.
  • Múltiples Agujas: 99.4% de precisión de recuperación (175/176) al recuperar múltiples hechos distintos simultáneamente.
• Viabilidad de Hardware: En una GPU A100 de 40 GB, una sola pasada hacia adelante de atención completa para una secuencia de 128K es inviable (requiriendo ~1 TB para las puntuaciones de atención). KV-Fold procesa la misma secuencia en 171 segundos con un uso máximo de memoria de 35.6 GB, manteniéndose dentro de los límites del hardware.
• Comparación con StreamingLLM: En un contexto de 128K, StreamingLLM mantiene una memoria acotada (16.6 GB) y un tiempo de ejecución más rápido (22s), pero cae al 0% de recuperación una vez que la "aguja" sale de su ventana de 1024 tokens. KV-Fold mantiene el 100% de recuperación en todas las distancias probadas, pero requiere más memoria (35.6 GB) y un tiempo de ejecución más largo.

Significado y Afirmaciones

El artículo afirma que los transformadores preentrenados congelados ya soportan una forma estable de recurrencia de caché KV. El significado principal es demostrar que la inferencia de contexto largo no requiere estrictamente cambios arquitectónicos, entrenamiento o compresión con pérdida para lograr una recuperación de alta fidelidad.

Los autores argumentan que KV-Fold revela un "régimen de atención ligeramente desplazado pero estable" donde el modelo puede procesar información recursivamente sin degradar su capacidad para acceder a tokens pasados específicos. Esto proporciona una ruta práctica para extender las ventanas de contexto en hardware existente, intercambiando el crecimiento lineal de memoria por la capacidad de realizar recuperación exacta sobre secuencias que de otro modo se truncarían o aproximarían. El artículo concluye que, aunque el método no reduce los FLOPs totales ni la memoria por debajo de la atención completa, reorganiza el cómputo en pasos manejables, haciendo viable la inferencia de contexto largo donde antes era imposible. El trabajo futuro sugerido incluye un ajuste fino modesto para reducir la deriva residual y combinar el método con técnicas de extensión posicional para ir más allá de los límites de contexto nativos.