DAMI: inizializzazione della memoria a doppio asse per migliorare il recupero su contesti lunghi nell’attenzione lineare con gating

Abstract degli autori

L’attenzione lineare scala linearmente con la lunghezza della sequenza, ma la memoria a stati finiti può limitare il recupero su contesti lunghi. Nelle ricorrenze lineari con gating, il mantenimento della memoria dipende dal decadimento dello stato, determinato congiuntamente dai parametri di gating A_log e dt_bias e dalle attivazioni in ingresso. Negli esperimenti controllati su Gated DeltaNet (GDN), entrambi i parametri restano fortemente correlati ai rispettivi valori iniziali anche dopo l’addestramento, evidenziando l’importanza dell’inizializzazione. Introduciamo quindi Dual-Axis Memory Initialization (DAMI), che configura congiuntamente questi parametri prima dell’addestramento. DAMI assegna scale temporali geometriche alle diverse teste, per gestire distanze di recupero variate, e un profilo di sensibilità all’input basso-alto-basso tra i livelli. Un piccolo campione di attivazioni calibra l’implementazione della famiglia GDN per farle raggiungere entrambi gli obiettivi comportamentali senza aggiungere parametri addestrabili. Su GDN, KDA, GDN2, Mamba-2 e RWKV-7, DAMI migliora l’accuratezza media nel recupero in stile RULER di 5,29-16,71 punti percentuali rispetto alle corrispondenti baseline con inizializzazione casuale. Le prestazioni nella modellazione del linguaggio, nel ragionamento di buon senso e nel recupero su dati reali da 2K restano nel complesso paragonabili a quelle delle baseline. Il nostro codice è disponibile all’indirizzo https://speedai.com/research/dami.