HiCI: Hierarchical Construction-Integration for Long-Context Attention

Resumen Técnico: HiCI (Construcción-Integración Jerárquica)

1. El Problema

Los Grandes Modelos de Lenguaje (LLMs) actuales enfrentan dos desafíos fundamentales al procesar secuencias largas:

• Eficiencia: La complejidad cuadrática ($O(n^2)$) de la atención auto-atención estándar hace que el costo computacional y de memoria sea prohibitivo a medida que aumenta la longitud de la secuencia.
• Efectividad: Simplemente extender la ventana de contexto no garantiza una modelización fiable de las dependencias de largo alcance. Las soluciones existentes (como atención dispersa o arquitecturas recurrentes) a menudo carecen de un sesgo inductivo explícito para organizar la información de un contexto largo en una jerarquía de "local a global" que guíe la atención.

Métodos actuales como la interpolación posicional (PI) o la atención dispersa (LongLoRA) mejoran la longitud o la eficiencia, pero no estructuran semánticamente la información, lo que lleva a una degradación en tareas de recuperación y comprensión.

2. Metodología: HiCI

Los autores proponen HiCI (Hierarchical Construction–Integration), un módulo de atención jerárquico inspirado en teorías cognitivas de la comprensión del discurso (Modelo de Construcción-Integración de Kintsch y la Teoría del Espacio de Trabajo Global). HiCI organiza el cálculo de la atención en tres etapas dentro de la arquitectura Transformer:

1. Construcción Local (Local Construction):

   • La secuencia de entrada se divide en segmentos de longitud fija.
   • Para cada segmento, se utiliza una atención cruzada con un conjunto compartido de $M$ "ranuras" (slots) de consulta aprendibles para extraer una representación local compacta ($L_i$).
   • Esto actúa como un cuello de botella que distila la información del segmento, priorizando la estructura saliente sobre los detalles de tokens individuales (análogo a la memoria de trabajo humana).

2. Integración Global (Global Integration):

   • Las representaciones locales de todos los segmentos se agregan en un contexto global compartido ($G$).
   • Este proceso utiliza estadísticas de múltiples vistas (media, máximo, mínimo, desviación estándar y media normalizada) sobre las representaciones locales.
   • Estas estadísticas se comprimen mediante un mecanismo compartido y se ponderan mediante atención para formar un contexto global de tamaño fijo ($K$), independiente de la longitud total de la secuencia.

3. Difusión Top-Down (Top-down Broadcast):

   • Tanto el contexto global ($G$) como la representación local específica del segmento ($L_i$) se preponen a la secuencia de claves y valores (KV) de cada segmento.
   • Esto permite que la atención a nivel de token se condicione sobre el contexto jerárquico (global y local) mientras se mantiene la paralelización entre segmentos.
   • Cada token puede así acceder indirectamente a la información de toda la secuencia a través de los slots globales.

Eficiencia: HiCI se implementa como un módulo ligero que se añade a modelos preentrenados (LLaMA-2) utilizando solo LoRA y capas de normalización, añadiendo aproximadamente un 5.5% de parámetros adicionales.

3. Contribuciones Clave

• Nuevo Módulo de Atención: Introducción de HiCI, que institucionaliza el sesgo inductivo de "construcción-integración-difusión" explícitamente en los Transformers.
• Extensión de Contexto Eficiente: Demostración de que HiCI permite extender el contexto de modelos LLaMA-2 preentrenados de 4K a 100K tokens (para 7B) y 64K tokens (para 13B) con un costo de parámetros mínimo.
• Rendimiento Superior: HiCI supera a los baselines fuertes (incluyendo LongLoRA) en modelado de lenguaje, recuperación y seguimiento de instrucciones.
• Análisis de Escalado: Los estudios de ablación revelan que la agregación jerárquica se beneficia de una granularidad de segmentos más fina, y que el rendimiento es casi invariante a la longitud bajo evaluación consistente con el entrenamiento.

4. Resultados Experimentales

Los experimentos se realizaron en modelos LLaMA-2 (7B y 13B) con entrenamiento en RedPajama y ajuste fino en LongAlpaca.

• Modelado de Lenguaje (Perplejidad): HiCI supera consistentemente a LongLoRA en conjuntos de prueba como PG-19 y Proof-pile. La mejora es más notable en contextos de evaluación más cortos, lo que sugiere una mejor preservación de la coherencia local.
• Recuperación (Retrieval):
  • Passkey Retrieval: HiCI logra 100% de precisión dentro del régimen de entrenamiento (32K) y muestra una degradación más suave al extrapolar a 56K en comparación con LongLoRA.
  • Topic Retrieval: HiCI-13B-16K iguala el rendimiento de modelos propietarios (como GPT-3.5-Turbo-16K) hasta 13K tokens y supera a todos los modelos de código abierto en 16K tokens (0.94 de precisión).
• Tareas de Instrucción (LongBench):
  • HiCI supera a LongLoRA en la mayoría de las categorías, con mejoras significativas en QA de documentos únicos (+7.4%) y tareas en chino (+11.8%).
  • Destacado: En la categoría de Código, HiCI supera a GPT-3.5-Turbo-16K con un +9.7% de precisión.
• Análisis de Atención: Se observa que las capas más profundas del modelo asignan progresivamente más atención a los slots globales (hasta un 40% en la última capa), indicando un enrutamiento emergente de información jerárquica.

5. Significado e Impacto

El trabajo demuestra que la estructuración jerárquica explícita de la información es un sesgo inductivo poderoso para el modelado de contextos largos. A diferencia de los métodos que solo ajustan la posición o la dispersión de la atención, HiCI imita procesos cognitivos humanos para construir representaciones locales, integrarlas globalmente y difundirlas.

Esto permite:

1. Escalabilidad Real: Extender ventanas de contexto a 100K tokens sin reentrenar el modelo base desde cero y con un sobrecosto de parámetros mínimo.
2. Robustez: Mejorar la estabilidad ante la extrapolación de longitud y la recuperación de información crítica en documentos largos.
3. Generalización: La arquitectura propuesta parece aplicable más allá de la familia LLaMA, sugiriendo un principio general para mejorar la comprensión de contextos largos en LLMs.

En resumen, HiCI ofrece una solución eficiente y efectiva al problema de la atención de largo alcance, superando a los modelos propietarios en tareas específicas de comprensión de código y superando a los baselines de código abierto en recuperación y modelado de lenguaje.