Resumen Técnico: OSDN: Mejora de la Regla Delta con Precondicionamiento Online Demostrable en Atención Lineal

1. Enunciado del Problema

La atención lineal y los modelos de espacio de estados (SSM) ofrecen una complejidad de inferencia de $O(N)$ y alternativas de memoria constante frente a la atención softmax, lo que las hace adecuadas para el procesamiento de secuencias largas. Sin embargo, a menudo tienen dificultades con el recuerdo asociativo en contexto, especialmente en comparación con los transformadores basados en softmax.

La Regla Delta (tal como se implementa en DeltaNet) intenta cerrar esta brecha interpretando la actualización del estado recurrente como un solo paso de descenso de gradiente online (OGD) sobre una pérdida de regresión por token: $f_t(S) = \frac{1}{2}\|S k_t - v_t\|_F^2$. La regla de actualización es $S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)k_t^\top$, donde $\beta_t$ es una puerta escalar.

La Limitación Central: La puerta escalar $\beta_t$ actúa como una tasa de aprendizaje uniforme aplicada a todas las dimensiones de la clave. Esto ignora la curvatura específica de cada característica del objetivo interno. En tareas de recuerdo asociativo, diferentes claves (por ejemplo, identificadores frecuentes frente a raros, valores estables frente a de alta entropía) exhiben curvaturas empíricas distintas. Un único tamaño de paso escalar fuerza un compromiso: aumentarlo para ayudar a sobrescribir memoria obsoleta en una dirección puede corregir en exceso otra dirección que ya está bien calibrada. Los optimizadores adaptativos estándar (como Adam) utilizan precondicionamiento diagonal para abordar esto, pero aplicar tales mecanismos a capas recurrentes típicamente requiere estados de segundo orden densos o rompe las pipelines de paralelización por fragmentos amigables con el hardware (como la transformada WY) que hacen eficiente a DeltaNet.

2. Metodología: DeltaNet Escalada Online (OSDN)

Los autores proponen OSDN, que aumenta la puerta escalar con un precondicionador diagonal $D_t = \text{diag}(d_t)$, aprendido online mediante retroalimentación de hipergradiente.

2.1. Equivalencia Algebraica y Preservación del Hardware

Una idea clave de OSDN es que precondicionar a la derecha el paso del gradiente es algebraicamente equivalente a una escalado por característica de la clave del lado de escritura.
La actualización estándar:
$$S_t = S_{t-1} - \beta_t \nabla f_t(S_{t-1}) D_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)(d_t \odot k_t)^\top$$
Definiendo una clave precondicionada $\tilde{k}_t = d_t \odot k_t$, la actualización se convierte en:
$$S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)\tilde{k}_t^\top$$
Esta transformación permite que OSDN preserve estrictamente la pipeline paralela WY por fragmentos de DeltaNet. La implementación solo requiere sustituir la clave de almacenamiento $K$ por $\tilde{K}$ en la matriz Gram de fragmentos y en las puntuaciones intra-fragmento, incurriendo en ninguna sobrecarga de estado de alta dimensión más allá del vector precondicionador de $O(K)$.

2.2. Retroalimentación de Hipergradiente Desacoplada

El precondicionador $d_t$ se actualiza mediante Descenso de Gradiente Online (OGD) sobre un surrogado de hipergradiente derivado del marco del Método de Gradiente Escalado Online (OSGM) [18].
La pérdida surrogada para el precondicionador es:
$$h_t(d_t) = \frac{f_t(S_t) - f_t(S_{t-1})}{\|\nabla f_t(S_{t-1})\|_F^2}$$
Críticamente, debido a la estructura exactamente cuadrática de la pérdida de regresión interna $f_t$, este surrogado admite una expresión en forma cerrada que depende únicamente de la clave $k_t$ y la puerta escalar $\beta_t$, desacoplando así del estado de alta dimensión $S_t$, el valor $v_t$ o el residuo $u_t$.
$$h_t(d_t) = \frac{(1 - \beta_t \langle d_t, k_t^2 \rangle)^2 - 1}{2 \|k_t\|_2^2}$$
Este desacoplamiento permite una implementación de dos fases:

1. Fase 1: Un escaneo ligero de $O(K)$ sobre el flujo de claves para calcular la secuencia de claves precondicionadas $\{\tilde{k}_t\}$ y actualizar $d_t$.
2. Fase 2: El paso estándar de DeltaNet por fragmentos utilizando $\tilde{K}$ en lugar de $K$.

2.3. Olvido Adaptativo del Precondicionador (APF)

Para manejar contextos lingüísticos no estacionarios donde las distribuciones de claves cambian, los autores introducen APF. Esto añade una puerta de retención aprendida, específica por token y por cabeza, $r_{t,h}$ aplicada solo al estado del precondicionador $d_t$ (no al estado de memoria $S_t$):
$$\bar{d}_{t+1} = r_{t,h} d_t + \eta \nabla_d h_t(d_t)$$
Esto mantiene la estructura de recurrencia afín requerida para el escaneo de dos fases, permitiendo al meta-optimizador "olvidar" la calibración obsoleta.

3. Garantías Teóricas

El artículo establece dos garantías de convergencia basadas en la naturaleza exactamente cuadrática de la pérdida interna:

1. Convergencia Super-Geométrica en el Límite de Población (Teorema 5.1): Frente a un comparador de Newton derecho ideal ($D^\star = \Sigma_k^\dagger$), si el aprendiz online logra un arrepentimiento sublineal, la suboptimalidad del objetivo global se contrae a una tasa super-geométrica de $O((C/T)^T)$. Esto es significativamente más rápido que la convergencia geométrica (lineal) estándar.
2. Contracción de Residuo Local al Algoritmo (Teorema 5.2): Para la actualización diagonal implementada, el producto de las relaciones de residuo locales por token $\prod_{t=1}^T \frac{f_t(S_t)}{f_t(S_{t-1})}$ está acotado. Si un comparador diagonal satisface la condición de Newton con puerta ($\beta_t \langle d^\star, k_t^2 \rangle = 1$), la tasa de contracción también es super-geométrica, específicamente $O((1/\sqrt{T})^T)$ bajo límites de arrepentimiento OGD estándar.

4. Resultados Experimentales

Los experimentos se realizaron en las arquitecturas base DeltaNet, DeltaNet con Puerta (GDN) y Kimi Delta Attention (KDA) a escalas de 340M y 1.3B parámetros.

4.1. Recuerdo en Contexto (Escala 340M)

• Rendimiento: OSDN vanilla mejoró el recuerdo en contexto estilo JRT en un 32% sobre la línea base DeltaNet (0.150 $\to$ 0.198). La ganancia se concentró en tareas de contexto repetido (por ejemplo, variantes "JRT-dos veces"), donde las claves se repiten, permitiendo que el precondicionador se calibre.
• Impacto de APF: OSDN-APF mantuvo una ganancia del 17% (0.176) y proporcionó la mejor perplejidad en WikiText/LAMBADA entre las variantes de DeltaNet (30.99 vs. 32.00).
• Generalización: En benchmarks más amplios (Commonsense, LongBench), las variantes de OSDN permanecieron a la par con sus líneas base, confirmando que el mecanismo es dirigido y no un potenciador universal de rendimiento.

4.2. Escalamiento a 1.3B Parámetros

• Contracción de Residuo: La señal a nivel de mecanismo se amplificó a escala. La relación geométrica media de residuo ($q_{geo}$) disminuyó de 0.432 a 0.265 (una reducción del 39%), casi duplicando la mejora observada a 340M.
• Paridad en Tareas Secundarias: A pesar de la masiva mejora en la contracción de residuo interna, las métricas secundarias (perplejidad en WikiText/LAMBADA, Commonsense, LongBench) permanecieron a la par con la línea base DeltaNet de 1.3B, demostrando que el mecanismo se transfiere efectivamente sin degradar las capacidades generales.

4.3. Eficiencia

• Rendimiento: Las variantes de OSDN incurrieron en una sobrecarga de inferencia mínima. A 340M, los cambios en el rendimiento estuvieron dentro de $\pm 2.2\%$. A 1.3B, OSDN-APF fue un 6.8% más lento que la línea base, atribuido al escaneo secuencial adicional para el precondicionador.
• Sobrecarga de Estado: El estado recurrente persistente aumentó solo en $\le 0.05\%$ debido al vector precondicionador de $O(K)$.

5. Significado y Afirmaciones

El artículo posiciona a OSDN como una extensión mínima y dirigida a la Regla Delta que cierra la brecha entre las actualizaciones escalares de primer orden y la optimización adaptativa.

• Mejora a Nivel de Mecanismo: La contribución principal es una mejora demostrable a nivel de mecanismo en el recuerdo asociativo mediante precondicionamiento online, validada por un diagnóstico directo de la relación de contracción de residuo.
• Compatibilidad con Hardware: A diferencia de otros métodos de segundo orden o precondicionamiento denso, OSDN preserva la pipeline paralela por fragmentos y la complejidad de estado $O(K)$, haciéndolo práctico para despliegue a gran escala.
• Escalabilidad: Los resultados demuestran que los beneficios del precondicionamiento online se transfieren y amplifican a la escala de miles de millones de parámetros, sugiriendo que la "desajuste de curvatura" en la atención lineal es un cuello de botella fundamental que puede abordarse sin una reestructuración arquitectónica.
• Limitaciones: Los autores notan explícitamente que OSDN no es un mejorador universal de benchmarks; las ganancias se concentran en tareas de recuperación que involucran claves repetidas. Las garantías teóricas dependen de la suposición de pérdida interna exactamente cuadrática y descenso monótono, que son condiciones idealizadas. El método no afirma resolver todas las limitaciones del aprendizaje en contexto, sino que aborda específicamente la geometría del paso de escritura de la Regla Delta.