Resumen Técnico: PReM (Preserve and Refresh Memory)

1. Planteamiento del Problema

La inferencia eficiente de contexto largo en los Modelos de Lenguaje Grande (LLMs) enfrenta una tensión fundamental entre la reducción de los costes computacionales y el mantenimiento del acceso a la evidencia contextual necesaria. A medida que aumenta la longitud del contexto, el caché de Clave-Valor (KV) crece, incrementando significamente la latencia de decodificación y el uso de memoria.

Los enfoques existentes para mitigar esto, como la compresión del caché KV y la compresión de contexto, sufren de dos limitaciones primarias:

1. Retención de Evidencia Estática: La mayoría de los métodos toman una decisión única y temprana sobre qué contexto retener. Esta compresión estática no logra adaptarse a las necesidades evolutivas del razonamiento de múltiples pasos, donde diferentes pasos de razonamiento pueden requerir diferentes piezas de evidencia.
2. Selección de Evidencia Desacoplada: Los métodos que dependen de compresores externos separan el proceso de compresión de la generación. Esto introduce computación adicional y evita que el modelo personalice la selección de evidencia dinámicamente para pasos de razonamiento específicos.

El artículo argumenta que la compresión efectiva debe ir más allá de la simple reducción de tamaño; debe preservar dinámicamente la evidencia útil y refrescarla cuando los pasos de razonamiento posteriores requieran información diferente.

2. Metodología: PReM

PReM (Preserve and Refresh Memory) es un marco de compresión de contexto que mantiene los contextos largos como una memoria KV interna por capa dentro del LLM. En lugar de depender de compresores externos, PReM aprende a decidir qué preservar y cuándo refrescar basándose en el estado interno del modelo.

Mecanismos Principales

• Representación de Memoria KV por Capas: El contexto largo se codifica como estados KV por capa ($M_\ell$) almacenados dentro del modelo. Estos estados se particionan en fragmentos contiguos.
• Capa de Memoria Dedicada: Se designa una capa específica ($\ell_m$) para tomar decisiones de selección de memoria. Esta capa utiliza una matriz de proyección de consulta de selección dedicada ($W^c_\ell$) para puntuar los fragmentos de contexto.
  • Selección de Fragmentos: El modelo puntúa los fragmentos basándose en la consulta actual (de la pregunta o de un token de refresco).
  • Preservar-y-Agrupar (Preserve-and-Pool): Los fragmentos seleccionados se preservan a nivel de token. Los fragmentos no seleccionados se comprimen en un único representante promedio de KV.
• El Token de Memoria $\langle m \rangle$: Un token especial sirve como una señal interna explícita para refrescar la memoria. Cuando el modelo genera $\langle m \rangle$, la capa de memoria reevalúa los fragmentos de contexto, selecciona nueva evidencia relevante y actualiza la memoria KV comprimida para las capas subsiguientes. Esto permite al modelo controlar cuándo refrescar basándose en el paso de razonamiento actual.

Entrenamiento: Entrenamiento de Refresco con Separación de Fases

Para entrenar al modelo para alinear la selección de memoria con la generación, PReM introduce una estrategia de Entrenamiento de Refresco con Separación de Fases. Cada paso de razonamiento se divide en dos pasadas hacia adelante distintas:

1. Fase de Selección de Memoria: El modelo procesa la entrada (pregunta o $\langle m \rangle$) para calcular las puntuaciones de los fragmentos y seleccionar los $k$ fragmentos de evidencia superiores. Esta fase es supervisada con una Pérdida de Selección de Memoria ($L_{ms}$) que incentiva al modelo a clasificar los fragmentos de evidencia por encima de los fragmentos que no son evidencia.
2. Fase de Generación Condicionada por Memoria: El modelo genera el segmento de la respuesta ($y_s$) condicionado al estado de memoria recién refrescado resultante de la fase de selección.

Para asegurar la consistencia entre el entrenamiento y la inferencia, el objetivo de entrenamiento incluye términos de entropía cruzada de frontera ($L_{bd}$). Estos términos conectan la fase de selección con el primer token de la respuesta y el final de un segmento de respuesta con el siguiente token $\langle m \rangle$, asegurando que el modelo aprenda los vínculos autorregresivos entre las actualizaciones de memoria y la generación.

3. Contribuciones Clave

1. Formulación: El artículo reformula la compresión de contexto como un problema aprendido de preservación y refresco de memoria, eliminando la necesidad de compresores externos.
2. Arquitectura: PReM introduce una capa de memoria dedicada y un token especial $\langle m \rangle$ para permitir actualizaciones de memoria dinámicas y específicas de cada paso dentro de los estados KV internos del LLM.
3. Estrategia de Entrenamiento: La propuesta de Entrenamiento de Refresco con Separación de Fases, que desacopla la selección de memoria de la generación durante el entrenamiento para asegurar que el modelo aprenda a condicionar la generación a la memoria específica requerida para ese paso.
4. Validación Empírica: Experimentos exhaustivos que demuestran que las actualizaciones de memoria aprendidas y dinámicas superan a la compresión estática y al prompting de contexto completo.

4. Resultados Experimentales

Los experimentos se realizaron en variantes de 3B y 7B de Qwen2.5-Instruct utilizando contextos de 32K tokens a través de seis benchmarks (TriviaQA, SQuAD, NaturalQuestions, 2WikiMQA, HotpotQA, MuSiQue).

• Rendimiento vs. Baselines: Bajo relaciones de compresión de 16× y 32×, PReM superó a ocho baselines fuertes, incluyendo métodos de caché KV (SnapKV, StreamingLLM, CAKE) y métodos de compresión de contexto (LongLLMLingua, EXIT, ICAE).
  • En el backbone de 3B con compresión de 16×, PReM mejoró el promedio de Exact Match (EM) y F1 en 5.30 y 5.15 puntos respectivamente sobre el baseline más fuerte.
  • En el backbone de 7B con compresión de 32×, las mejoras alcanzaron 8.20 (EM) y 9.89 (F1) puntos.
• Comparación con Contexto Completo: PReM superó al prompting directo de 32K tokens (sin compresión) tanto en modelos de 3B como de 7B, sugiriendo que la selección de memoria aprendida es más efectiva que exponer el modelo al contexto completo sin filtrar.
• Superioridad de Modelos Pequeños: El modelo PReM de 3B superó a los baselines de compresión de contexto suave de 7B liberados (ej. Activation Beacon), indicando que la preservación y el refresco aprendidos pueden compensar las escalas menores del backbone.
• Eficiencia: PReM redujo significativamente la memoria GPU pico (ej. ~19.57GB vs. 30.0GB para memoria completa) manteniendo una latencia de inferencia competitiva. El uso de tokens $\langle m \rangle$ permite un compromiso controlable entre la computación de decodificación (frecuencia de refresco) y el ahorro de memoria.

5. Significado y Reivindicaciones

El artículo afirma que PReM demuestra que la inferencia de contexto largo efectiva requiere actualizaciones de memoria acopladas a la generación en lugar de un acortamiento de contexto estático. Al mantener el contexto como memoria KV interna y aprender a refrescarlo dinámicamente, el modelo puede acceder adaptativamente a la evidencia específica necesaria para cada paso de razonamiento.

Los autores enfatizan que su enfoque aborda las limitaciones de la retención estática y la selección desacoplada, mostrando que un modelo puede aprender a "saber qué preservar y cuándo refrescar", lo que conduce a un rendimiento superior en tareas de razonamiento de múltiples saltos (multi-hop) y un mejor equilibrio entre eficiencia y calidad en comparación con las técnicas de compresión existentes. El trabajo sugiere que la gestión de memoria interna, entrenada conjuntamente con la generación, es un camino viable para escalar los LLMs a contextos más largos sin costes computacionales prohibitivos.