Resumen Técnico: CoMem (Comprehension Memory)

Planteamiento del Problema

El modelado de lenguaje de contexto largo enfrenta dos cuellos de botella primarios: el costo cuadrático de pre-llenado (prefill) de la auto-atención y el crecimiento lineal de la memoria del caché de Clave-Valor (KV). Los métodos eficientes existentes suelen abordar esto mediante la expulsión o compresión de estados KV a lo largo del eje de tokens (reteniendo un subconjunto de tokens) o mediante la recuperación de unidades de texto externas. Sin embargo, estos enfoques a menudo luchan por equilibrar la eficiencia de la memoria con la fidelidad requerida para el razonamiento complejo sobre contextos largos. Además, muchos métodos dependen del entrenamiento de redes de memoria especializadas o de la compresión de todo el contexto, lo que puede introducir una sobrecarga significativa o una pérdida de fidelidad.

Los autores postulan que las capas del Transformer exhiben una división del trabajo: la información semántica se vuelve accesible en las capas inferiores y medias, mientras que las capas superiores se especializan cada vez más en las representaciones para la consulta (query) y la predicción actual. Los métodos actuales no explotan plenamente esta estructura basada en la profundidad para organizar la memoria.

Metodología: CoMem

CoMem (Comprehension Memory) reorganiza la memoria de contexto largo a lo largo del eje de profundidad en lugar de centrarse únicamente en el eje de tokens. Combina la selección de tokens acotada con la reutilización de computación particionada por profundidad.

Mecanismo Central

1. Fase de Escritura (Partición de Profundidad):

   • El contexto de entrada se transmite en fragmentos (chunks) disjuntos (por ejemplo, 512 tokens).
   • Cada fragmento se procesa solo a través de una capa intermedia específica $j$ (la "profundidad de división" o split depth).
   • En lugar de almacenar los estados KV de profundidad completa, CoMem almacena el tensor residual $h_j$ (el estado oculto) en la capa $j$ para cada token en el fragmento.
   • Estos tensores se almacenan en un índice externo indexado por el texto del fragmento.
   • Eficiencia de Almacenamiento: Para Qwen3-8B, almacenar $h_j$ requiere solo 1/18 parte de los bytes de un par KV completo en bf16 por token.

2. Fase de Lectura (Recuperación Acotada y Recomputación):

   • Al recibir una consulta, un selector (usando BM25) recupera un conjunto acotado de los $k$ fragmentos más relevantes.
   • Los tensores residuales $h_j$ almacenados para estos fragmentos se empaquetan junto con un token "sumidero" (sink token) y la consulta actual.
   • El modelo recomputa solo las capas superiores $[j:L]$ sobre este paquete de tamaño fijo utilizando atención de entre-fragmentos (cross-chunk attention) completa.
   • Complejidad: Para un presupuesto de recuperación ($k$) y un tamaño de fragmento fijos, el cómputo de lectura en línea y el conjunto de trabajo de memoria del modelo son independientes de la longitud total del contexto almacenado.

Entrenamiento y Adaptación

• Columna Vertebral (Backbone): El modelo base (Qwen3-8B) permanece congelado.
• Autodestilación: Para mitigar la pérdida de fidelidad causada por saltarse las capas inferiores durante la fase de lectura, se entrena un LoRA ligero (Rango-32) mediante autodestilación. El modelo estudiante (en la división $j=12$) aprende a igualar los logits de un modelo profesor congelado (en $j=0$) sobre el soporte de los top-64 logits.
• Datos de Entrenamiento: El LoRA se entrena con texto plano PG-19 sin supervisión de tareas específicas o datos de recuperación, lo que permite al modelo generalizar la capacidad de "división de profundidad".

Contribuciones Clave

1. La Profundidad como Eje de Reutilización: El artículo establece que la accesibilidad semántica alcanza su pico cerca de la profundidad media, mientras que el "límite de legibilidad zero-shot" (donde los estados almacenados pueden usarse directamente sin recomputación) depende de la escala del modelo. CoMem trata esta división de profundidad ($j$) como un control ajustable entre el trabajo precomputado y el costo en el momento de la consulta.
2. Memoria Particionada por Profundidad Acotada: CoMem almacena un estado residual por token y retoma exactamente la parte superior del Transformer. Combina la recuperación acotada con la atención de entre-fragmentos completa, logrando una huella de almacenamiento de 1/18 de los cachés KV completos por token en Qwen3-8B.
3. Compromiso Controlado entre Calidad y Eficiencia: Los autores proporcionan evidencia de que un almacenamiento en caché más profundo reduce el cómputo por consulta pero incurre en una pérdida de fidelidad, la cual la autodestilación repara sustancialmente. El método aísla los beneficios de la recuperación de los costos de la reutilización de profundidad.

Resultados Experimentales

Evaluado mediante un protocolo unificado, sin plantillas de chat, utilizando Qwen3-8B:

• Rendimiento en Benchmarks:
  • RULER: CoMem alcanza 97.05, superando al baseline KV-Direct de contexto completo (78.80).
  • LoCoMo (Memoria de Diálogo): CoMem obtiene 38.27 frente a 34.59 de KV-Direct. Esta ventaja persiste bajo el remuestreo de clústeres de conversación y auditorías de jueces independientes.
  • LongEval & LongBench: CoMem iguala o supera a los baselines (69.0 en LongEval; 12.15 macro F1 en LongBench).
• Eficiencia (NVIDIA H20 a 128k de contexto):
  • Memoria: CoMem utiliza 18.26 GB de memoria pico en comparación con los 89.36 GB de KV-Direct de contexto completo.
  • Velocidad: CoMem logra una aceleración de 7.83× en el tiempo de pre-llenado que incluye la escritura completa.
  • Decodificación: La decodificación de banda reanudada se mantiene dentro del 11% de la inferencia densa hasta 64k y se vuelve 1.07× más rápida a 128k.
• Ablaciones:
  • Recuperación vs. Profundidad: La selección acotada por sí sola (sin reutilización de profundidad) mejora la memoria de diálogo, pero la reutilización de profundidad es crítica para la eficiencia.
  • Destilación: La autodestilación recupera aproximadamente 13.75 puntos de LoCoMo perdidos por la división congelada en $j=12$.
  • Atención de Entre-Fragmentos (Cross-Chunk Attention): La atención completa entre fragmentos es esencial para la desambiguación de múltiples hechos; el uso de bloques diagonales causa caídas significativas de rendimiento.

Significado y Afirmaciones

El artículo afirma que la memoria de contexto largo puede organizarse eficazmente a lo largo del eje de capas, no solo del eje de tokens. CoMem demuestra que:

• La información semántica es accesible en profundidades intermedias, lo que permite el almacenamiento en caché de estados ciegos a la consulta que luego son especializados mediante la recomputación de las capas superiores.
• Una estrategia de recuperación acotada combinada con la computación particionada por profundidad puede desacoplar el costo de la inferencia en línea de la longitud del contexto almacenado.
• La autodestilación ligera puede reparar la pérdida de fidelidad inherente a saltarse las capas inferiores, permitiendo una alta precisión sin reentrenar la columna vertebral.

Los autores enfatizan que CoMem es un método de eficiencia que ofrece un compromiso calidad-memoria-cómputo, más que una solución de máxima precisión para cada tarea dentro de la ventana. Declaran explícitamente que sus resultados en contextos que exceden la ventana nativa (40,960 tokens) son "referencias de estrés no extendidas" y no reclaman superioridad sobre modelos que utilizan técnicas de extensión de posición como YaRN. El método se presenta como un enfoque complementario a los métodos existentes de compresión de tokens y recuperación.