Resumen Técnico: Compresión de Contexto de Extremo a Extremo a Escala

Planteamiento del Problema

La inferencia de contexto largo en los Modelos de Lenguaje de Gran Escala (LLMs) está cada vez más limitada por la memoria y la latencia debido al crecimiento lineal del caché de Clave-Valor (KV) con la longitud del contexto. Aunque las técnicas recientes de compresión de caché KV intentan mitigar esto expulsando entradas del caché o entrenando cachés compactos, enfrentan limitaciones significativas:

1. Compromiso entre Calidad y Eficiencia: Muchos métodos degradan sustancialmente el rendimiento del modelo o requieren un entrenamiento específico para el dominio que no se generaliza.
2. Compatibilidad con Motores de Inferencia: Las técnicas que expulsan entradas de forma no uniforme a través de las cabezas de atención o capas son difíciles de integrar con los motores de inferencia de producción modernos (por ejemplo, vLLM, SGLang), los cuales asumen longitudes de secuencia compartidas.
3. Sobrecarga de Prefill: Muchos enfoques requieren que el contexto completo sea pre-llenado (prefilled) antes de la compresión, lo que anula las ganancias potenciales de latencia, o requieren un tiempo de cómputo adicional sustancial.
4. Limitaciones de los Soft-Tokens: Los métodos existentes de "soft-token" de codificador-decodificador suelen depender de un preprocesamiento offline, fallan en preservar el comportamiento general del modelo base, o se evalúan solo en tareas de contexto largo y dominio específico, en lugar de en benchmarks de información densa.

La pregunta central abordada es si un compresor general y agnóstico a la tarea puede ser entrenado para preservar las capacidades originales de un modelo mientras reduce significativamente el uso de memoria y el tiempo de inferencia.

Metodología

Los autores introducen los Modelos de Lenguaje de Contexto Latente (LCLMs), una familia de compresores codificador-decodificador entrenados de extremo a extremo. La arquitectura central mapea una secuencia de tokens de entrada $x_{1:T}$ en una secuencia más corta de embeddings latentes $z_{1:M}$ a través de un codificador, que luego son proyectados por un adaptador y consumidos por un decodificador en lugar del contexto original.

Diseño de la Arquitectura

A través de una búsqueda de arquitectura a gran escala desde cero, los autores identificaron elecciones de diseño óptimas:

• Agrupación (Pooling): Se encontró que el mean pooling (agrupación por media) es superior al pooling basado en tokens (EOS/CLS) y competitivo con la concatenación, particularmente en altas tasas de compresión (16×).
• Granularidad de Codificación: Se adoptó un tamaño de ventana del codificador ($W$) de 1024 tokens, lo que permite al codificador atender sobre un contexto local más grande que el bloque de compresión ($N$) sin una sobrecarga de memoria prohibitiva.
• Enmascaramiento de Atención: El enmascaramiento causal dentro del codificador superó consistentemente al enmascaramiento bidireccional.
• Adaptador: Un adaptador MLP ligero (sin auto-atención) resultó ser más efectivo y eficiente que los adaptadores basados en atención para proyectar los latentes del codificador a la dimensión del decodificador.
• Inicialización: Inicializar el codificador a partir de un modelo de embedding pre-entrenado (Qwen3-Embedding-0.6B) en lugar de un LLM estándar proporcionó mejores representaciones iniciales, aunque la brecha se redujo con un entrenamiento extensivo.

Receta de Entrenamiento

Los autores desarrollaron una receta de entrenamiento de múltiples etapas para prevenir el olvido catastrófico y asegurar la estabilidad:

1. Etapa 0 (Calentamiento del Adaptador): Congelar el codificador y el decodificador; entrenar solo el adaptador.
2. Etapa 1 (Entrenamiento del Codificador): Descongelar el codificador; mantener el decodificador congelado.
3. Etapa 2 (Pre-entrenamiento Continuo): Descongelar el decodificador con una tasa de aprendizaje pequeña; entrenar de extremo a extremo con una mezcla de predicción de siguiente token y datos de reconstrucción.
4. Etapa 3 (Ajuste Fino Supervisado - SFT): Entrenar con datos de SFT para mejorar el seguimiento de instrucciones y el razonamiento con entradas comprimidas.

Los datos de entrenamiento incluyen una mezcla cuidadosamente curada de código, texto, razonamiento y documentos de contexto largo. Crucialmente, los datos de pre-entrenamiento utilizan un formato intercalado donde los segmentos comprimidos (envueltos en tokens especiales) alternan con segmentos no comprimidos, forzando al modelo a condicionarse sobre el contexto latente en múltiples posiciones en lugar de solo al principio.

Integración de Agentes

Los autores proponen un sistema agéntico donde el agente LCLM recibe el contexto comprimido completo. El agente puede invocar selectivamente una herramienta EXPAND para recuperar el texto original no comprimido para segmentos específicos, permitiendo una estrategia de "escanear y hacer zoom" (skim-and-zoom) para tareas de largo horizonte.

Resultados Clave

La familia LCLM (específicamente un codificador de 0.6B y un decodificador de 4B) fue evaluada en tasas de compresión de 4×, 8× y 16× en benchmarks que incluyen RULER, LongBench, LongHealth y GSM8K.

• Frontera de Pareto: Los LCLMs establecen una nueva frontera de Pareto en el compromiso entre precisión y eficiencia. Logran un Tiempo al Primer Token (TTFT) significativamente más rápido y un menor uso de memoria GPU pico en comparación con las bases de comparación de compresión de caché KV (por ejemplo, SnapKV, KVzip, Attention Matching), especialmente en altas tasas de compresión.
• Preservación de la Precisión: Con una compresión de 16×, los LCLMs mantienen una alta precisión en benchmarks de contexto largo (por ejemplo, ~65% en RULER 16k frente al ~93% de la KV completa, superando significativamente a las bases de comparación que a menudo caen por debajo del 20-40%).
• Eficiencia de Memoria: A diferencia de los métodos de caché KV donde la memoria escala con la longitud total del contexto, el uso de memoria de LCLM se estabiliza en contextos largos porque el codificador procesa ventanas de tamaño fijo. Esto permite que los LCLMs manejen contextos de hasta 1M de tokens en una sola GPU H200, mientras que las bases de comparación se quedan sin memoria en 512K o 1M de tokens.
• Rendimiento de Grano Fino: En GSM8K (contextos cortos y densos), los LCLMs alcanzan la mayor precisión en todas las tasas de compresión, demostrando que no están limitados solo a QA de contexto largo, sino que pueden comprimir eficazmente entradas de información densa.
• Rendimiento de Agentes: El sistema agéntico con expansión adaptativa mejora significativamente la precisión de coincidencia exacta de cadenas en tareas de "aguja en un pajar" (needle-in-a-haystack), a menudo igualando el rendimiento del modelo no comprimido mientras retiene la eficiencia del contexto comprimido.

Significado y Reivindicaciones

El artículo afirma que los LCLMs ofrecen una solución práctica y escalable para la inferencia de contexto largo que es compatible con los motores de inferencia de producción estándar. Al desplazar la carga de la compresión hacia un codificador más pequeño y paralelizable, los LCLMs desacoplan los costos de memoria y latencia del contexto de la longitud del proceso de generación del decodificador.

Los autores enfatizan que su enfoque:

1. Preserva las Capacidades Generales: A diferencia de los métodos previos de soft-tokens que degradan el rendimiento o requieren ajuste específico de la tarea, los LCLMs mantienen un fuerte rendimiento en tareas generales, razonamiento y seguimiento de instrucciones.
2. Permite la Recuperación Adaptativa: La integración con marcos agénticos permite la expansión dinámica del contexto relevante, cerrando la brecha entre la visibilidad global del corpus y la recuperación de información de grano fino.
3. Escalabilidad: El método escala efectivamente a contextos de un millón de tokens en hardware de consumo (GPUs H200), superando los cuellos de botella de memoria que actualmente limitan el despliegue de sistemas de LLM de largo horizonte.

El trabajo sugiere que la compresión aprendida de extremo a extremo es una alternativa viable a la expulsión heurística de caché KV, proporcionando un sustrato robusto para futuros agentes de largo horizonte con memorias de trabajo persistentes.