Muir: attenzione lineare con gating allocato in base alla sensibilità e lettura con filtro adattato

Abstract degli autori

L’attenzione lineare comprime il contesto in uno stato chiave-valore di dimensioni fisse, per un addestramento a tempo lineare e un’inferenza con memoria costante. I modelli gated delta-rule fanno decadere questo stato e cancellano la sua previsione per una chiave in ingresso prima di scrivere un nuovo valore; Gated DeltaNet-2 (GDN-2) generalizza questo approccio con un decadimento canale per canale e gate di cancellazione e scrittura disaccoppiati e a rango pieno. Tuttavia, mostriamo che gli aggiornamenti dello stato sono molto meno sensibili alla cancellazione che alle scritture: l’attivazione delle chiavi e la normalizzazione a norma unitaria concentrano l’energia delle chiavi su un sottoinsieme sparso di canali, mentre le scritture dei valori scalano tutti i canali; il gate di cancellazione può quindi essere parametrizzato a un rango inferiore. Inoltre, le righe dello stato che decadono lentamente accumulano norme molto più grandi di quelle che decadono rapidamente, dominando la lettura non ponderata e imponendo una cancellazione più forte al momento della scrittura per sopprimere le interferenze. Sulla base di questa analisi, introduciamo Muir, che fonde tutti i gate in una proiezione a basso rango allocata in base alla sensibilità, ripesa le righe dello stato con un gate di lettura a filtro adattato e reinveste i parametri risparmiati in una rete feed-forward più ampia. Poiché gli algoritmi standard a blocchi diventano numericamente instabili alle grandi dimensioni dei blocchi necessarie per saturare le ampie unità di calcolo matriciale, progettiamo un aggiornamento parallelo per un addestramento stabile ed efficiente. Con un numero di parametri compreso tra 166M e 1,34B, Muir riduce la perplexity dell’1,8-2,0% rispetto a GDN-2 e del 2,5-3,6% rispetto a Transformer, aumenta fino a 1,9 punti l’accuratezza nei test di conoscenza comune e migliora il richiamo estrattivo rispetto a GDN-2. A 1,34B, Muir supera entrambi i modelli di riferimento nel richiamo multi-hop e multi-needle, estrapola di 2.560x dalla lunghezza di addestramento di 4K a flussi di 10M di token senza variazioni della perplexity e mantiene un recupero delle passkey prossimo al 100% fino a 1M di token con una lettura statica. Muir ottiene questi risultati mantenendo la stessa velocità di elaborazione di GDN-2, raggiunge 1,2x più rapidamente la sua loss di validazione a 100B di token e si addestra 2,7x più velocemente di Transformer a 32K.