Context Memorization for Efficient Long Context Generation

Resumen Técnico: Memoria de Estado de Atención para la Generación Eficiente de Contextos Largos

Enunciado del Problema
Las aplicaciones modernas de Modelos de Lenguaje Grandes (LLM) dependen cada vez más de prefijos de condicionamiento largos (por ejemplo, ejemplos en contexto, documentos recuperados o instrucciones de agentes) para guiar la inferencia. Sin embargo, la inferencia aumentada con prefijos enfrenta dos limitaciones estructurales principales:

1. Decaimiento del Prefijo: A medida que avanza la generación, la atención del modelo se distribuye entre los tokens recién generados, lo que provoca que la influencia del prefijo se desvanezca, particularmente en escenarios de contexto largo.
2. Ineficiencia en la Inferencia: El cálculo de atención sobre el prefijo escala linealmente con su longitud ($O(L)$), imponiendo una latencia significativa y una sobrecarga de memoria tanto durante la fase de prellenado como en cada paso de decodificación. Aunque la caché de prefijos amortiza los costos de prellenado, aún incurre en un consumo de memoria sustancial.

Las soluciones existentes se dividen en dos categorías, ambas con inconvenientes:

• Internalización del Prefijo: Métodos que ajustan los parámetros del modelo mediante fine-tuning o utilizan hiperredes para internalizar el comportamiento del prefijo. Estos requieren entrenamiento basado en gradientes, lo que los hace intensivos en recursos, lentos para actualizar e inadecuados para cambios dinámicos en el prefijo.
• Compresión del Prefijo: Métodos que mantienen el prefijo en el mecanismo de atención pero lo comprimen (por ejemplo, poda de tokens o expulsión de la caché KV). Estos aún requieren atender al prefijo comprimido durante la inferencia, lo que significa que la latencia escala con la longitud comprimida y la influencia del prefijo sigue sujeta al decaimiento.

Metodología: Memoria de Estado de Atención (ASM)
Los autores proponen la Memoria de Estado de Atención, un enfoque libre de entrenamiento que externaliza el prefijo en una memoria ligera basada en búsqueda de estados de atención precalculados. Este método aprovecha la identidad de softmax en línea para descomponer la atención sobre un bloque de claves concatenado en atención sobre subbloques.

La idea central es que, para una consulta dada, el estado de atención $(a, Z)$ —que comprende la salida de atención y el término de normalización— es suficiente para reconstruir la contribución del prefijo a la atención sin volver a atender a las claves y valores originales.

El marco de trabajo opera en dos fases:

1. Fase de Calibración Offline (Construcción de Memoria):

   • Recolección: El modelo realiza pasadas hacia adelante sobre un conjunto de trazas representativas de consultas-respuestas concatenadas con el prefijo. Para cada token de consulta en la respuesta, el sistema registra el vector de consulta y el estado de atención correspondiente $(a, Z)$ sobre el prefijo.
   • Agrupamiento: Los vectores de consulta recolectados se agrupan (por ejemplo, mediante K-medias) en $K$ centroides.
   • Agregación: Para cada grupo, se crea una única entrada representativa. El vector de consulta del centroide es la media de las consultas del grupo. El estado de atención $(\bar{a}, \bar{Z})$ se calcula mediante agregación consciente de la atención, fusionando los estados de todos los miembros del grupo utilizando la regla de actualización de softmax en línea. Esto asegura que el centroide represente un promedio sin pérdidas de la atención del prefijo del grupo.
   • Eficiencia: Para manejar prefijos muy largos (por ejemplo, más de 16K tokens) sin exceder la memoria de la GPU, el prefijo puede dividirse en fragmentos, codificarse de forma independiente y fusionarse dentro de la memoria utilizando la propiedad de composibilidad de la identidad de softmax en línea.

2. Fase de Inferencia Online:

   • Recuperación: Para una consulta entrante, el sistema identifica el centroide de memoria más cercano utilizando similitud coseno (opcionalmente con transformaciones de blanqueamiento y manejo de RoPE).
   • Fusión: El estado de atención precalculado recuperado $(\bar{a}, \bar{Z})$ se fusiona con la salida de autoatención de la consulta (calculada sobre tokens no pertenecientes al prefijo) utilizando el operador de fusión definido por la identidad de softmax en línea. Esto recupera la salida de atención exacta que se habría producido si el prefijo hubiera estado presente, sin atender a los tokens del prefijo.
   • Complejidad: El costo de búsqueda escala logarítmicamente ($O(\log K)$) con el tamaño de la memoria mediante indexación jerárquica, independientemente de la longitud original del prefijo.

Contribuciones Clave

1. Externalización Libre de Entrenamiento: El artículo propone un método para externalizar prefijos largos en un diccionario de memoria compacto sin requerir entrenamiento basado en gradientes ni modelos auxiliares.
2. Extensión del Softmax en Línea: Los autores extienden la identidad de softmax en línea desde el cálculo eficiente de atención al dominio de la reutilización de prefijos entre consultas, permitiendo la reconstrucción sin pérdidas de la atención del prefijo mediante recuperación.
3. Rendimiento y Eficiencia: Los experimentos demuestran que ASM iguala o supera el rendimiento de la atención completa mientras reduce significativamente los costos computacionales.

Resultados Experimentales
El método se evaluó en ManyICLBench (Aprendizaje en Contexto) y RuleArena/NBA (Generación Aumentada por Recuperación) utilizando LLaMA-3.1-8B.

• Aprendizaje en Contexto (ICL): En ManyICLBench, ASM mejoró la precisión sobre el aprendizaje en contexto estándar en presupuestos de memoria de 1K a 8K entradas. Con un presupuesto de memoria de 8K, redujo la latencia de atención en 1.36× en comparación con la atención completa. ASM superó a las líneas base de compresión de caché KV (KVzip), que a menudo colapsaban información crítica como los tokens de etiqueta.
• Generación Aumentada por Recuperación (RAG): En la prueba de NBA (que involucra ~20K tokens de reglas), ASM superó el rendimiento de RAG con atención completa utilizando solo el 20% de la huella de memoria. El rendimiento óptimo se observó con un tamaño de libro de códigos intermedio (4K entradas), lo que sugiere que el número de entradas es un hiperparámetro específico de la tarea en lugar de uno que deba simplemente maximizarse.
• Latencia: Aunque la búsqueda jerárquica incurre en una sobrecarga mayor en presupuestos de memoria muy pequeños, el crecimiento de la latencia es sublineal. ASM se vuelve más rápido que la atención completa en aproximadamente 4K entradas de memoria y logra una aceleración de 1.8× en 16K entradas.

Significado y Afirmaciones
El artículo afirma que la Memoria de Estado de Atención aborda los costos estructurales de la generación de contextos largos al desacoplar el conocimiento del prefijo del cálculo de atención por consulta. Al externalizar el prefijo en una memoria precalculada basada en búsqueda, el método elimina la escala lineal de la latencia de inferencia con la longitud del prefijo y elimina el problema de decaimiento del prefijo asociado con los mecanismos de atención estándar.

Los autores posicionan este trabajo como un paso hacia una dirección de investigación más amplia: externalizar el conocimiento de los LLM en representaciones compactas y reutilizables que existen más allá de los prompts de texto y los parámetros del modelo. Enfatizan que su enfoque ofrece una compensación favorable entre precisión y eficiencia, proporcionando una alternativa libre de entrenamiento a la internalización de parámetros y una alternativa más eficiente a la compresión de prefijos.