Riepilogo Tecnico: OSDN: Miglioramento della Regola Delta con Precondizionamento Online Provabile nell'Attenzione Lineare

1. Enunciato del Problema

L'attenzione lineare e i modelli a spazio di stato (SSM) offrono una complessità di inferenza $O(N)$ e alternative a memoria costante rispetto all'attenzione softmax, rendendoli adatti all'elaborazione di sequenze lunghe. Tuttavia, spesso faticano con il richiamo associativo in contesto, specialmente se confrontati con i transformer basati su softmax.

La Regola Delta (come implementata in DeltaNet) tenta di colmare questo divario interpretando l'aggiornamento dello stato ricorrente come un singolo passo di discesa del gradiente online (OGD) su una perdita di regressione per token: $f_t(S) = \frac{1}{2}\|S k_t - v_t\|_F^2$. La regola di aggiornamento è $S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)k_t^\top$, dove $\beta_t$ è un gate scalare.

Il Limite Fondamentale: Il gate scalare $\beta_t$ agisce come un tasso di apprendimento uniforme applicato a tutte le dimensioni delle chiavi. Ciò ignora la curvatura specifica per caratteristica dell'obiettivo interno. Nei compiti di richiamo associativo, chiavi diverse (ad esempio, identificatori frequenti vs. rari, valori stabili vs. ad alta entropia) esibiscono curvature empiriche diverse. Un singolo passo scalare forza un compromesso: aumentarlo per aiutare a sovrascrivere la memoria obsoleta per una direzione può correggere eccessivamente un'altra direzione che è già ben calibrata. Gli ottimizzatori adattivi standard (come Adam) utilizzano un precondizionamento diagonale per affrontare questo problema, ma applicare tali meccanismi ai livelli ricorrenti richiede tipicamente stati di secondo ordine densi o interrompe le pipeline di parallelizzazione a blocchi favorevoli all'hardware (come la trasformazione WY) che rendono DeltaNet efficiente.

2. Metodologia: DeltaNet Scalata Online (OSDN)

Gli autori propongono OSDN, che potenzia il gate scalare con un precondizionatore diagonale $D_t = \text{diag}(d_t)$, appreso online tramite feedback di ipergradiente.

2.1. Equivalenza Algebrica e Preservazione dell'Hardware

Un'idea chiave di OSDN è che il precondizionamento a destra del passo del gradiente è algebricamente equivalente a una scala per caratteristica della chiave lato scrittura.
L'aggiornamento standard:
$$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$$
Definendo una chiave precondizionata $\tilde{k}_t = d_t \odot k_t$, l'aggiornamento diventa:
$$S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)\tilde{k}_t^\top$$
Questa trasformazione permette a OSDN di preservare rigorosamente la pipeline parallela a blocchi WY di DeltaNet. L'implementazione richiede solo la sostituzione della chiave di archiviazione $K$ con $\tilde{K}$ nella matrice Gram a blocchi e nei punteggi intra-blocco, senza incorrere in overhead di stato ad alta dimensionalità oltre al vettore precondizionatore $O(K)$.

2.2. Feedback di Ipergradiente Disaccoppiato

Il precondizionatore $d_t$ viene aggiornato tramite Discesa del Gradiente Online (OGD) su un surrogato di ipergradiente derivato dal framework del Metodo di Gradiente Scalato Online (OSGM) [18].
La perdita surrogata per il precondizionatore è:
$$h_t(d_t) = \frac{f_t(S_t) - f_t(S_{t-1})}{\|\nabla f_t(S_{t-1})\|_F^2}$$
Crucialmente, a causa della struttura esattamente quadratica della perdita di regressione interna $f_t$, questo surrogato ammette un'espressione in forma chiusa che dipende solo dalla chiave $k_t$ e dal gate scalare $\beta_t$, disaccoppiandolo dallo stato ad alta dimensionalità $S_t$, dal valore $v_t$ o dal 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}$$
Questo disaccoppiamento abilita un'implementazione in due fasi:

1. Fase 1: Una scansione leggera $O(K)$ sul flusso di chiavi per calcolare la sequenza di chiavi precondizionate $\{\tilde{k}_t\}$ e aggiornare $d_t$.
2. Fase 2: Il passaggio standard di DeltaNet a blocchi utilizzando $\tilde{K}$ invece di $K$.

2.3. Dimenticanza del Precondizionatore Adattivo (APF)

Per gestire contesti linguistici non stazionari dove le distribuzioni delle chiavi cambiano, gli autori introducono l'APF. Questo aggiunge un gate di ritenzione appreso, specifico per token e per testa, $r_{t,h}$ applicato solo allo stato del precondizionatore $d_t$ (non allo stato di memoria $S_t$):
$$\bar{d}_{t+1} = r_{t,h} d_t + \eta \nabla_d h_t(d_t)$$
Ciò mantiene la struttura di ricorrenza affine richiesta per la scansione in due fasi, consentendo all'ottimizzatore meta di "dimenticare" la calibrazione obsoleta.

3. Garanzie Teoriche

Il documento stabilisce due garanzie di convergenza basate sulla natura esattamente quadratica della perdita interna:

1. Convergenza Super-Geometrica nel Limite della Popolazione (Teorema 5.1): Contro un comparatore di Newton a destra ideale ($D^\star = \Sigma_k^\dagger$), se l'apprendista online raggiunge un rimpianto sublineare, la subottimalità dell'obiettivo globale si contrae a un tasso super-geometrico di $O((C/T)^T)$. Questo è significativamente più veloce della convergenza geometrica standard (lineare).
2. Contrazione del Residuo Locale all'Algoritmo (Teorema 5.2): Per l'aggiornamento diagonale implementato, il prodotto dei rapporti di residuo locali ai token $\prod_{t=1}^T \frac{f_t(S_t)}{f_t(S_{t-1})}$ è limitato. Se un comparatore diagonale soddisfa la condizione di Newton gateata ($\beta_t \langle d^\star, k_t^2 \rangle = 1$), anche il tasso di contrazione è super-geometrico, specificamente $O((1/\sqrt{T})^T)$ sotto i limiti di rimpianto OGD standard.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su backbones DeltaNet, Gated DeltaNet (GDN) e Kimi Delta Attention (KDA) alle scale di 340M e 1.3B parametri.

4.1. Richiamo in Contesto (Scala 340M)

• Prestazioni: OSDN vanilla ha migliorato il richiamo in contesto stile JRT del 32% rispetto alla baseline DeltaNet (0.150 $\to$ 0.198). Il guadagno si è concentrato su compiti a contesto ripetuto (ad esempio, varianti "JRT-due"), dove le chiavi ricorrono, permettendo al precondizionatore di calibrarsi.
• Impatto APF: OSDN-APF ha mantenuto un guadagno del 17% (0.176) e ha fornito la migliore perplessità WikiText/LAMBADA tra le varianti DeltaNet (30.99 vs 32.00).
• Generalizzazione: Su benchmark più ampi (Commonsense, LongBench), le varianti OSDN sono rimaste alla pari con le rispettive baseline, confermando che il meccanismo è mirato e non un potenziatore universale delle prestazioni.

4.2. Scalabilità a 1.3B Parametri

• Contrazione del Residuo: Il segnale a livello di meccanismo si è amplificato alla scala. Il rapporto geometrico medio dei residui ($q_{geo}$) è sceso da 0.432 a 0.265 (una riduzione del 39%), quasi raddoppiando il miglioramento osservato a 340M.
• Parità a valle: Nonostante il massiccio miglioramento nella contrazione del residuo interno, le metriche a valle (perplessità WikiText/LAMBADA, Commonsense, LongBench) sono rimaste alla pari con la baseline DeltaNet da 1.3B, dimostrando che il meccanismo si trasferisce efficacemente senza degradare le capacità generali.

4.3. Efficienza

• Throughput: Le varianti OSDN hanno comportato un overhead di inferenza minimo. A 340M, le variazioni di throughput sono state entro $\pm 2.2\%$. A 1.3B, OSDN-APF è stato 6.8% più lento della baseline, attribuito alla scansione sequenziale aggiuntiva per il precondizionatore.
• Overhead di Stato: Lo stato ricorrente persistente è aumentato solo di $\le 0.05\%$ a causa del vettore precondizionatore $O(K)$.

5. Significato e Affermazioni

Il documento posiziona OSDN come un'estensione minima e mirata alla Regola Delta che colma il divario tra aggiornamenti scalari del primo ordine e ottimizzazione adattiva.

• Miglioramento a Livello di Meccanismo: Il contributo principale è un miglioramento provabile a livello di meccanismo nel richiamo associativo tramite precondizionamento online, validato da una diagnosi diretta del rapporto di contrazione del residuo.
• Compatibilità Hardware: A differenza di altri metodi di precondizionamento di secondo ordine o densi, OSDN preserva la pipeline parallela a blocchi e la complessità di stato $O(K)$, rendendolo pratico per il deployment su larga scala.
• Scalabilità: I risultati dimostrano che i benefici del precondizionamento online si trasferiscono e si amplificano alla scala dei miliardi di parametri, suggerendo che il "disadattamento di curvatura" nell'attenzione lineare è un collo di bottiglia fondamentale che può essere affrontato senza un ripensamento architetturale.
• Limitazioni: Gli autori notano esplicitamente che OSDN non è un miglioratore universale di benchmark; i guadagni sono concentrati su compiti di recupero che coinvolgono chiavi ripetute. Le garanzie teoriche si basano sull'assunzione di perdita interna esattamente quadratica e discesa monotona, che sono condizioni idealizzate. Il metodo non afferma di risolvere tutte le limitazioni dell'apprendimento in contesto, ma affronta specificamente la geometria del passo di scrittura della Regola Delta.