Resumen Técnico: Atención Lineal Kaczmarz (KLA)

Declaración del Problema

El modelado de lenguaje de contexto largo se ve obstaculizado por los costos computacionales y de memoria cuadráticos de la atención Softmax estándar en los Transformers. Si bien los modelos recurrentes de tiempo lineal abordan esto comprimiendo el contexto en un estado de tamaño fijo, el diseño de la regla de actualización de estado se convierte en el cuello de botella crítico. Los enfoques existentes, como Gated DeltaNet (GDN), utilizan una "regla delta" para escribir el residuo entre un valor objetivo y la predicción actual del estado, editando así las asociaciones existentes en lugar de simplemente acumularlas.

Sin embargo, GDN depende de un coeficiente escalar aprendido para equilibrar el olvido y la actualización. Los autores argumentan que este coeficiente es una elección de diseño empírica en lugar de un tamaño de paso derivado teóricamente. Un coeficiente inadecuado puede conducir a magnitudes de actualización subóptimas: sobre-corregir para claves de gran norma y sub-corregir para claves de pequeña norma, lo que resulta en corrupción del estado y un rendimiento de recuperación subóptimo.

Metodología: Atención Lineal Kaczmarz

El artículo propone la Atención Lineal Kaczmarz (KLA), una modificación de GDN que reemplaza el escalar aprendido con un tamaño de paso derivado teóricamente basado en el método de proyección de Kaczmarz.

Derivación Teórica

Los autores plantean la operación de escritura clave-valor como un problema de optimización con restricciones. Para un token $t$, el objetivo es encontrar el estado $S_t$ más cercano al estado anterior decaído $\tilde{S}_{t-1}$ que satisfaga la restricción lineal $S_t^T k_t = v_t$ (donde $k_t$ es la clave y $v_t$ es el valor).

• Objetivo de Optimización: $\min_S \frac{1}{2}\|S - \tilde{S}_{t-1}\|_F^2$ sujeto a $S^T k_t = v_t$.
• Solución: La solución óptima es la proyección ortogonal de $\tilde{S}_{t-1}$ sobre el hiperplano de restricción. Esto produce un tamaño de paso (coeficiente) de $\beta_t = 1/\|k_t\|_2^2$.

La Regla de Actualización KLA

KLA generaliza esta proyección exacta en una regla de actualización práctica que incluye una puerta de relajación aprendible $\eta_t$ y un estabilizador numérico $\epsilon$:
$$\beta_t = \frac{\eta_t}{\|k_t\|_2^2 + \epsilon}$$
La actualización del estado se convierte en:
$$S_t = \tilde{S}_{t-1} + \beta_t k_t e_t^T$$
donde $e_t = v_t - \tilde{S}_{t-1}^T k_t$ es el error residual.

Propiedades Clave

El artículo establece que este coeficiente satisface tres propiedades deseables simultáneamente (en el caso no relajado donde $\eta_t=1, \epsilon=0$):

1. Proyección Ortogonal: Realiza el cambio de norma mínima al estado para satisfacer la restricción.
2. Búsqueda de Línea Exacta: Es el minimizador exacto de la pérdida por token a lo largo de la dirección de actualización.
3. Paso de Gradiente Normalizado: Actúa como un paso de Descenso de Gradiente Estocástico (SGD) normalizado.

Implementación

KLA está diseñado como un reemplazo directo para GDN. Preserva:

• La forma del estado ($S \in \mathbb{R}^{d_k \times d_v}$).
• Los mecanismos de puerta (puertas de decaimiento y actualización).
• La estructura de recurrencia lineal.
• El algoritmo de entrenamiento paralelo por bloques.

El único cambio estructural es el cálculo de la matriz de coeficientes diagonal $B = \text{Diag}(\beta_1, \dots, \beta_C)$ dentro del solucionador por bloques. No se requieren nuevos núcleos de hardware ni diseños de estado.

Contribuciones Clave

1. Regla Delta Auto-normalizada: Derivación del coeficiente KLA a partir de la proyección de Kaczmarz, demostrando que es una proyección ortogonal, un minimizador de búsqueda de línea exacta y un paso de gradiente normalizado para la pérdida por token.
2. Algoritmo Recurrente Directo: Demostración de que KLA conserva la recurrencia de GDN y el solucionador por bloques, requiriendo solo un cambio en la matriz de coeficientes diagonal.
3. Ganancias Empíricas: Validación de que KLA logra una menor perplejidad y un mejor seguimiento de estado que GDN y otras líneas base de tiempo lineal sin expansión arquitectónica.
4. Eficiencia Preservada: Confirmación de que KLA mantiene una complejidad de tiempo/memoria de $O(n)$ y iguala la latencia de prellenado de GDN mientras mejora el rendimiento de decodificación.

Resultados Experimentales

Los experimentos se realizaron a la escala de 0.4B de parámetros con un presupuesto de 1B de tokens en el conjunto de datos SlimPajama-672B.

Modelado de Lenguaje

• Perplejidad de Validación (PPL): KLA logró 8.09 PPL, superando a GDN (8.50), DeltaNet (8.26), GLA (8.42), Longhorn (8.79) y Mamba2 (9.41).
• Extrapolación de Contexto Largo: Sin ajuste fino de contexto largo, KLA mantuvo la menor perplejidad en longitudes de contexto de 1K a 65K tokens, mostrando una ventaja consistente de 3.7–4.1 PPL sobre GDN.

Tareas Sintéticas

• Aguja en un Pajare (S-NIAH) Simple: KLA alcanzó 100% de precisión en todas las longitudes de contexto probadas (1K–8K), mientras que GDN disminuyó ligeramente y Mamba2 falló significativamente.
• Recordatorio Asociativo de Múltiples Consultas (MQAR): KLA mejoró la precisión de extrapolación de longitud 8× en 7.03 puntos sobre GDN (73.84% vs 66.81%).
• Pila y Palíndromo: KLA mostró un rendimiento competitivo o superior en tareas de seguimiento de estado y reversión.

Eficiencia

• Latencia de Prellenado: KLA igualó la latencia de prellenado de GDN (por ejemplo, ~1220 ms en contexto de 131K), confirmando que el cálculo del coeficiente de Kaczmarz no impacta el costo dominante del núcleo por bloques.
• Rendimiento de Decodificación: En contexto de 32K, KLA logró 3.54 tokens/s, lo cual es 2.1× mayor que GDN (1.65 tokens/s) y 1.7× mayor que Mamba2.

Estudios de Ablación

Los autores ablandaron la estrategia de normalización, el mecanismo de puerta y la expansión del estado. Los resultados indicaron que la normalización por norma de clave es crítica para la estabilidad; eliminarla o usar solo un escalar aprendido condujo a una degradación significativa en la estabilidad del preentrenamiento (PPL) y en el rendimiento de las tareas sintéticas.

Significado y Afirmaciones

El artículo afirma que el coeficiente de Kaczmarz normalizado por norma de clave es un "eje de diseño de primer orden" para los modelos de secuencia de regla delta. Al tratar la operación de escritura como una proyección con restricciones, KLA resuelve el problema de las magnitudes de actualización subóptimas inherentes a los escalares aprendidos empíricamente.

El significado radica en el hecho de que estas mejoras se logran sin cambiar el estado recurrente, el núcleo de hardware ni la complejidad arquitectónica. Los autores postulan que el coeficiente en sí mismo, en lugar de la capacidad del estado o la estructura de puertas, es el factor principal que determina si una escritura de regla delta resuelve con éxito su problema de regresión local. Esto sugiere que los futuros modelos de tiempo lineal pueden lograr ganancias sustanciales en precisión, extrapolación y eficiencia de decodificación simplemente adoptando este tamaño de paso fundamentado teóricamente.