FG$^2$-GDN: Enhancing Long-Context Gated Delta Networks with Doubly Fine-Grained Control

1. El Problema

Los mecanismos de atención lineal han surgido como alternativas prometedoras a la atención softmax tradicional, ofreciendo complejidad temporal lineal ($O(L)$) en lugar de cuadrática ($O(L^2)$) durante la inferencia, lo cual es crucial para aplicaciones de contexto largo.

A pesar de los avances recientes como Gated DeltaNet (GDN) y Kimi Delta Attention (KDA), que han mejorado la recuperación asociativa mediante la regla delta (una actualización de descenso de gradiente en línea), persiste una limitación fundamental:

• En KDA, aunque se refinó la puerta de decaimiento ($\alpha_t$) de un escalar a un vector por canal (permitiendo tasas de olvido independientes por dimensión), la tasa de aprendizaje ($\beta_t$) en la actualización delta sigue siendo un escalar.
• Esto significa que todas las dimensiones de la característica comparten la misma "fuerza de escritura" o tasa de adaptación, lo que limita la capacidad del modelo para realizar adaptaciones específicas por dimensión, análogo a la diferencia entre el optimizador SGD y los optimizadores adaptativos por coordenada (como AdaGrad o Adam).

2. Metodología

Los autores proponen FG2-GDN (Gated Delta Network con Control de Beta Fino-Granularidad), que extiende el marco de KDA introduciendo tasas de aprendizaje adaptativas por coordenada.

A. FG2-GDN: Tasa de Aprendizaje Vectorial

El núcleo de la propuesta es reemplazar el escalar $\beta_t$ por un vector $\beta_t \in \mathbb{R}^{d_k}$.

• Mecanismo: En lugar de aplicar una matriz diagonal asimétrica que rompería la estructura de rango-1 necesaria para la paralelización eficiente, los autores absorben simétricamente el vector $\beta_t$ en los vectores de clave ($k_t$) y valor ($v_t$):
  $$\tilde{k}_t = \sqrt{\beta_t} \odot k_t, \quad \tilde{v}_t = \sqrt{\beta_t} \odot v_t$$
• Regla de Actualización: La actualización del estado oculto $S_t$ se transforma en:
  $$S_t = (I - \tilde{k}_t \tilde{k}_t^\top) \text{Diag}(\alpha_t) S_{t-1} + \tilde{k}_t \tilde{v}_t^\top$$
• Interpretación: Desde la perspectiva del aprendizaje en línea, esto equivale a usar tasas de aprendizaje por coordenada en el descenso de gradiente estocástico. Esto permite que cada dimensión del estado mantenga dinámicas de adaptación independientes, concentrando la "energía de escritura" en las dimensiones más discriminativas.

B. FG2-GDN+: Desacoplamiento de Escala

Se introduce una variante avanzada, FG2-GDN+, que desacopla la escala para claves y valores mediante dos vectores independientes ($\beta^k_t$ y $\beta^v_t$):

• Esto permite controlar independientemente la fuerza de borrado (a través de $\beta^k_t$) y la fuerza de escritura (a través de $\beta^v_t$).
• Esta flexibilidad ayuda a gestionar múltiples asociaciones competitivas en contextos largos, reduciendo la interferencia entre pares clave-valor almacenados.

C. Eficiencia Computacional y Arquitectura Híbrida

• Paralelización por Bloques (Chunkwise): La actualización propuesta mantiene la estructura Diagonal-Plus-Low-Rank (DPLR). Esto permite reutilizar algoritmos de entrenamiento paralelos existentes (como la representación WY) sin aumentar significativamente la complejidad computacional.
• Arquitectura: Se utiliza una arquitectura híbrida que intercala capas de FG2-GDN con capas de Atención Latente Multi-cabeza (MLA) en una proporción de 3:1, combinando la eficiencia de la recurrencia lineal con la capacidad de interacción completa de la atención estándar.

3. Contribuciones Clave

1. FG2-GDN: Introducción de tasas de aprendizaje vectoriales por canal en la regla delta, cerrando la brecha entre el control de olvido (ya existente en KDA) y el control de escritura.
2. FG2-GDN+: Propuesta de un mecanismo desacoplado para claves y valores, permitiendo un control independiente sobre la erasure (borrado) y la escritura de información.
3. Eficiencia: Demostración de que este mayor grado de granularidad se logra con una sobrecarga computacional mínima (<5% en el prefill respecto a KDA) y sin sacrificar la escalabilidad lineal.
4. Validación Empírica: Resultados exhaustivos que muestran mejoras consistentes en recuperación asociativa y comprensión de contexto largo.

4. Resultados Experimentales

Los experimentos se realizaron en escalas de 340M y 1.3B parámetros, utilizando datasets como SlimPajama y evaluando en benchmarks estándar.

• Modelado de Lenguaje (Perplejidad y Precisión):

  • FG2-GDN logró la mejor precisión promedio en tareas de razonamiento de sentido común y la perplejidad más baja en LAMBADA (13.09 en 1.3B), superando a KDA y GDN.
  • Esto indica que las tasas de escritura por dimensión benefician particularmente la comprensión de lenguaje en contextos cortos.

• Evaluación de Contexto Largo (RULER y LongBench):

  • RULER: FG2-GDN y FG2-GDN+ mostraron mejoras sustanciales en tareas de recuperación (Single-key y Multi-key NIAH) y extracción de palabras. FG2-GDN+ lideró en tareas de recuperación multi-llave, sugiriendo que el desacoplamiento ayuda a manejar interferencias.
  • LongBench: En tareas de QA de documentos largos, resumen y completado de código, FG2-GDN+ obtuvo el puntaje promedio más alto (18.3%), superando significativamente a los baselines.
  • Extrapolación: Los modelos entrenados con secuencias de 8k mantuvieron un rendimiento robusto al evaluar en 16k, demostrando una mejor capacidad de extrapolación gracias al control de memoria más estructurado.

• Eficiencia:

  • La velocidad de prefill (throughput) de FG2-GDN es comparable a KDA, con menos del 5% de sobrecarga.
  • Mantiene la escalabilidad lineal, superando a la atención completa (MLA) en velocidad a secuencias largas (hasta 32k tokens).

5. Significado

Este trabajo representa un avance significativo en la evolución de los modelos de atención lineal y recurrentes:

• Paradigma de Optimización: Establece un puente teórico y práctico entre los optimizadores adaptativos por coordenada (como Adam) y los mecanismos de atención recurrente, demostrando que la granularidad fina en la actualización de estados es crucial para el rendimiento.
• Gestión de Memoria a Largo Plazo: Proporciona un mecanismo más sofisticado para la gestión de la memoria asociativa, permitiendo que el modelo decida no solo qué olvidar (vía $\alpha_t$), sino también cómo y con qué fuerza escribir nueva información en cada dimensión específica.
• Viabilidad Práctica: Al mantener la estructura DPLR, demuestra que las mejoras de rendimiento no requieren sacrificar la eficiencia computacional, haciendo que estos modelos sean viables para aplicaciones de contexto largo en producción (agentes, multimodalidad) donde la latencia y el costo son críticos.

En resumen, FG2-GDN mejora la capacidad de los modelos lineales para retener y recuperar información precisa en contextos muy largos, superando las limitaciones de los enfoques anteriores mediante un control de doble granularidad (olvido y escritura) sin comprometer la eficiencia.