Forget, Then Recall: Learnable Compression and Selective Unfolding via Gist Sparse Attention

1. El Problema

El escalado de los Modelos de Lenguaje Grandes (LLM) a contextos largos se ve limitado por el costo computacional cuadrático ($O(n^2)$) de los mecanismos de atención estándar, donde cada token debe interactuar con todos los tokens anteriores.

Existen dos enfoques principales para mitigar esto, pero ambos tienen desventajas:

• Selección/Compresión de KV-cache (Inferencia): Métodos como H2O o StreamingLLM eliminan tokens antiguos o seleccionan claves basándose en heurísticas. Sin embargo, operan sobre modelos preentrenados con atención completa, lo que limita su capacidad para aprender patrones óptimos de dispersión desde cero.
• Atención Escasa (Entrenamiento): Métodos como NSA, DSA o MoBA introducen dispersión durante el preentrenamiento. No obstante, suelen requerir modificaciones arquitectónicas (ramas paralelas), módulos externos no diferenciables (indexadores) o operaciones no diferenciables (promedios de agrupación) que impiden la optimización basada en gradientes del proceso de selección.

Además, en escenarios como la Generación Aumentada por Recuperación (RAG), la atención completa sobre múltiples documentos irrelevantes introduce ruido y un sesgo inductivo pobre, diluyendo la atención sobre la información relevante.

2. Metodología: Gist Sparse Attention (GSA)

Los autores proponen GSA, un marco unificado que conecta la compresión basada en "gist" (esencia) con la atención escasa mediante un mecanismo de despliegue selectivo (selective unfolding). La idea central es tratar los tokens comprimidos no solo como resúmenes, sino como señales de enrutamiento aprendibles.

Componentes Clave:

1. Tokens Gist Entrelazados:

   • El contexto se divide en fragmentos (chunks) de longitud $L$.
   • Se inserta un token "gist" ($g_m$) después de cada fragmento.
   • Se aplica una máscara causal especial: los tokens dentro de un fragmento pueden atenderse entre sí y a los gists anteriores, pero los tokens posteriores solo pueden atender al gist actual y a los gists anteriores. Esto fuerza al token gist a aprender una representación comprimida y densa de su fragmento.

2. Puntuación de Relevancia y Selección:

   • En cada paso de decodificación, la consulta actual ($q_t$) calcula una puntuación de relevancia con cada token gist mediante un producto punto ($s_{t,m} = q_t^T k_{g_m}$).
   • Se seleccionan los top-$k$ fragmentos más relevantes basándose en estas puntuaciones. A diferencia de métodos anteriores que usan promedios fijos (no diferenciables), aquí la selección es totalmente diferenciable y aprendida.

3. Despliegue Selectivo (Selective Unfolding):

   • Una vez seleccionados los $k$ fragmentos relevantes, el modelo "despliega" o restaura sus tokens crudos (raw tokens) en el contexto de atención.
   • La atención final se realiza sobre una mezcla híbrida: los tokens gist de todos los fragmentos (como resúmenes globales) más los tokens crudos de los fragmentos seleccionados.
   • Esto permite un mecanismo de búsqueda de lo grueso a lo fino: el modelo mantiene una visión global compacta y recupera detalles granulares solo donde es necesario.

4. Extensión Jerárquica (H-GSA):

   • Para lograr una complejidad logarítmica, los autores proponen una estructura recursiva "gist de gist".
   • Los tokens gist se comprimen a su vez en "meta-gist" tokens.
   • La selección se realiza en múltiples niveles: primero se seleccionan los segmentos de meta-gist más relevantes, luego los gists dentro de esos segmentos, y finalmente los tokens crudos. Esto reduce la complejidad de enrutamiento de $O(n)$ a $O(\log n)$.

Entrenamiento:

El marco se entrena en dos etapas (opcionalmente):

1. Preentrenamiento Continuo: El modelo aprende a comprimir el contexto en tokens gist.
2. Ajuste Fino Selectivo (Opcional): Se introduce el mecanismo de despliegue selectivo en el entrenamiento, permitiendo que el modelo aprenda a activar la atención sobre tokens crudos específicos basándose en la consulta.

3. Contribuciones Principales

• Puente Aprendible: Conecta la compresión de contexto y la atención escasa en un marco unificado y end-to-end, sin modificar la arquitectura del Transformer ni usar módulos externos.
• Tokens Gist como Enrutadores: Demuestran que los tokens gist entrelazados sirven eficazmente como señales de enrutamiento para seleccionar sub-contextos relevantes, superando a los promedios de agrupación no diferenciables.
• Mecanismo de Despliegue Selectivo: Propone una estrategia de "búsqueda de lo grueso a lo fino" que recupera detalles granulares solo cuando es necesario, mejorando la precisión bajo presupuestos de tokens estrictos.
• Escalabilidad Logarítmica: La extensión jerárquica (H-GSA) permite manejar contextos ilimitados con complejidad de decodificación por paso logarítmica.

4. Resultados Experimentales

Los autores evaluaron GSA en LongBench (compresión de contexto largo) y en 5 benchmarks de RAG (Generación Aumentada por Recuperación), utilizando modelos Qwen2-7B y Llama3.2-1B.

• Rendimiento Superior: GSA superó consistentemente a las líneas base de compresión (ActivationBeacon, UniGist, KVLink) y métodos de atención escasa de inferencia (H2O, StreamingLLM) en ratios de compresión de 8x a 32x.
  • En RAG con preentrenamiento continuo, GSA superó a los modelos de atención completa (Full-PT) en más de 11 puntos en promedio, demostrando su capacidad para filtrar documentos distractores.
  • Con ajuste fino, GSA alcanzó un rendimiento cercano al de la atención completa, superando a las líneas base de compresión en 8-12 puntos.
• Ventaja en RAG: En escenarios de múltiples documentos, GSA eliminó la necesidad de atención cruzada sobre tokens crudos entre documentos, reduciendo el ruido y mejorando la eficiencia.
• H-GSA: La variante jerárquica mostró ventajas significativas en ratios de compresión altos (16x y 32x), validando la eficacia de la selección multinivel.
• Estabilidad: La selección basada en top-k adaptativo demostró ser más estable y precisa que las variantes basadas en top-p (umbral de probabilidad).

5. Significado e Impacto

Este trabajo representa un avance significativo en la eficiencia de los LLMs para contextos largos al demostrar que la compresión no tiene que ser irreversible ni estática.

• Eficiencia sin Sacrificio de Calidad: Permite reducir drásticamente la memoria y el costo computacional (hasta 32x) manteniendo o incluso mejorando la precisión, algo que los métodos de compresión estática no logran.
• Integración Práctica: Al no requerir cambios arquitectónicos ni módulos externos, GSA es fácilmente integrable en infraestructuras existentes y modelos preentrenados.
• Nueva Paradigma de Atención: Introduce la idea de que la compresión puede servir como un mecanismo de enrutamiento dinámico y diferenciable, alineando mejor el comportamiento del modelo con la forma en que los humanos procesan información (resúmenes generales con acceso selectivo a detalles).

En resumen, GSA ofrece una solución elegante y efectiva para el cuello de botella de la atención cuadrática, permitiendo que los modelos de lenguaje manejen contextos de millones de tokens de manera eficiente y precisa.