Riepilogo Tecnico: AMOR (Adaptive Metacognitive Output Router)

Enunciato del Problema

I modelli ibridi a ricorrenza-attenzione mirano a combinare l'efficienza inferenziale dei modelli a spazio di stato ricorrenti lineari (SSM) con il potere espressivo dei meccanismi di attenzione. Tuttavia, le architetture ibride esistenti si basano tipicamente su strategie a programma fisso, in cui l'attenzione viene applicata uniformemente su tutti i token o a livelli predeterminati (ad esempio, ogni $N$-esimo livello). Questo approccio comporta l'intero costo quadratico dell'attenzione ($O(N^2)$) anche quando lo stato ricorrente da solo è sufficiente a prevedere il token successivo con alta confidenza. Di conseguenza, questi modelli non riescono a sfruttare appieno i guadagni di efficienza degli SSM, in particolare in scenari a contesto lungo, dove il sovraccarico computazionale dei blocchi di attenzione non necessari diventa un collo di bottiglia. Inoltre, gli ibridi fissi spesso si degradano sotto spostamenti di distribuzione (ad esempio, estrapolazione della lunghezza) perché impongono meccanismi di attenzione che potrebbero non essere calibrati per la difficoltà specifica dell'input.

Metodologia: Architettura AMOR

Gli autori propongono AMOR (Adaptive Metacognitive Output Router), un'architettura ibrida post-hoc che alloca dinamicamente l'attenzione basandosi sull'incertezza predittiva piuttosto che su posizioni fisse.

Architettura Principale

AMOR è costruita su una spina dorsale ricorrente standard (testata con Mamba2 e Gated DeltaNet) e aggiunge $K$ (tipicamente 3) blocchi di attenzione controllati dall'entropia dopo i livelli ricorrenti ma prima della testa LM finale.

• Spina dorsale: Il mixer ricorrente elabora l'intera sequenza, mantenendo uno stato di dimensione fissa.
• Blocchi Post-Hoc: A differenza degli ibridi seriali che sostituiscono i livelli ricorrenti, i blocchi AMOR sono additivi. Operano sull'output normalizzato della spina dorsale.
• Asimmetria: Il primo blocco AMOR (Blocco 0) opera direttamente sul residuo normalizzato finale della spina dorsale (post-norm), mentre i blocchi successivi (1 e 2) utilizzano una formulazione standard pre-norm. Questo design preserva la spina dorsale come modello linguistico valido e autonomo, trattando i blocchi AMOR come raffinamenti condizionali.

Il Meccanismo del Gate di Entropia

L'innovazione centrale è un meccanismo di gate rigido, privo di gradienti, ispirato alla distinzione tra Sistema 1 e Sistema 2 nelle scienze cognitive (Kahneman, 2011).

1. Stima dell'Incertezza: Per ogni token, la testa LM condivisa calcola i logit. L'entropia normalizzata dell'output ($H_t$) viene calcolata dalla distribuzione softmax.
2. Soglia Adattiva: Una soglia dinamica $\tau$ viene calcolata come media mobile esponenziale (EMA) della mediana dell'entropia del batch ($\mu$) più una deviazione standard scalata ($\kappa\sigma$):
   $$\tau = \mu + \kappa\sigma$$
   Ciò consente alla soglia di tracciare gli spostamenti nella distribuzione dell'incertezza del modello durante l'addestramento senza richiedere parametri di instradamento appresi.
3. Esecuzione Condizionale:
   • Se $H_t > \tau$, il gate si attiva ($g_t=1$): Il token subisce un'attenzione causale (SDPA) con RoPE, e l'output viene aggiunto al residuo con un fattore di scalatura appreso per canale $\alpha$.
   • Se $H_t \le \tau$, il gate rimane inattivo ($g_t=0$): Il calcolo dell'attenzione viene saltato e il token passa invariato.
4. Flusso dei Gradienti: Il gate è non differenziabile. I gradienti fluiscono attraverso i pesi dell'attenzione ($W_Q, W_K, W_V, W_O$) solo nelle posizioni di attivazione. Crucialmente, $W_K$ e $W_V$ ricevono gradienti dalle future posizioni di attivazione per garantire che la cache KV rimanga valida per il recupero, anche se il token corrente ha saltato l'attenzione.

Addestramento e Inferenza

• Addestramento: AMOR viene addestrato in modalità "dense with mask" utilizzando FlashAttention-2 per l'efficienza, matematicamente equivalente a una vera esecuzione sparsa. La soglia $\tau$ viene congelata durante l'inferenza.
• Inferenza: Il modello esegue un salto condizionale. Quando il gate è spento, la proiezione $Q$ costosa, il calcolo dell'attenzione e la proiezione $O$ vengono aggirati, risparmiando risorse computazionali. Tuttavia, le proiezioni $K$ e $V$ e gli aggiornamenti della cache avvengono sempre per mantenere la continuità del contesto.

Contributi Chiave

1. Instradamento Guidato dall'Incertezza: AMOR introduce un meccanismo semplice e post-hoc che invoca selettivamente l'attenzione solo quando la spina dorsale ricorrente mostra alta incertezza predittiva (alta entropia), invece di affidarsi a programmi di livelli fissi.
2. Ibridazione Efficiente in Parametri: Il metodo aggiunge un sovraccarico di parametri minimo (circa <5% per i blocchi di gate e attenzione) raggiungendo prestazioni comparabili o superiori agli ibridi a programma fisso.
3. Robustezza agli Spostamenti di Distribuzione: Affidandosi alla spina dorsale ricorrente per i token a bassa incertezza, AMOR mitiga il degrado delle prestazioni osservato negli ibridi fissi durante l'estrapolazione della lunghezza, dove i meccanismi di attenzione spesso soffrono di deriva posizionale.
4. Gate Senza Gradienti: L'approccio evita la complessità di router appresi o obiettivi di rinforzo, utilizzando una soglia statistica derivata dalla distribuzione di output del modello stesso.

Risultati Sperimentali

Gli autori hanno valutato AMOR su tre scale (180M, 440M, 1.5B parametri) utilizzando le spine dorsali Mamba2 e Gated DeltaNet, confrontandole con modelli ricorrenti puri, Transformer puri e ibridi seriali/fusi a programma fisso.

• Ragionamento di Senso Comune: AMOR ha costantemente eguagliato o superato sia le spine dorsali ricorrenti pure che gli ibridi a programma fisso su tutte le scale e i benchmark (LAMBADA, HellaSwag, PIQA, ARC, ecc.).
• Recupero in Contesto: Mentre i modelli ricorrenti puri faticano nel recupero di copie esatte, AMOR ha migliorato sostanzialmente la propria spina dorsale, ottenendo risultati competitivi con gli ibridi fissi attivando l'attenzione su solo ~22% dei token.
• Prestazioni a Contesto Lungo: Sul benchmark LongBench, AMOR ha mantenuto prestazioni stabili all'aumentare della lunghezza del contesto. Al contrario, gli ibridi a programma fisso e i Transformer puri si sono degradati significativamente a causa degli spostamenti di distribuzione RoPE e dei vincoli di memoria. La capacità di AMOR di saltare l'attenzione sui token prevedibili gli ha permesso di scalare in modo più fluido.
• Efficienza Inferenziale: AMOR ha dimostrato una migliore scalabilità del throughput di decodifica nei regimi a contesto lungo. L'analisi teorica e i benchmark empirici (fino a 448K di contesto) hanno mostrato che i risparmi derivanti dal salto dell'attenzione superano il sovraccarico del calcolo dell'entropia una volta che la lunghezza del contesto supera un certo punto di intersezione (circa 256K per i modelli da 440M).

Significato e Affermazioni

Il paper sostiene che quando viene applicata l'attenzione è tanto critico quanto quanto viene applicata. Gli autori affermano che:

1. Efficienza ed Espressività non sono un compromesso: Trattando l'allocazione delle risorse computazionali come un problema di instradamento basato sulla difficoltà dell'input, i modelli possono raggiungere un'alta espressività senza sostenere l'intero costo dell'attenzione.
2. La semplicità è efficace: Router complessi appresi o obiettivi ausiliari non sono necessari; un semplice gate di entropia guidato dall'output può coordinare efficacemente le architetture ibride.
3. Computazione Adattiva: AMOR indica una classe di modelli che adattano dinamicamente il calcolo alla difficoltà dell'input, preservando i punti di forza delle leggere spine dorsali a spazio di stato mentre dispiegano selettivamente pesanti meccanismi di attenzione solo quando il modello è incerto.

I risultati suggeriscono che l'allocazione selettiva dell'attenzione basata sull'incertezza predittiva offre un'alternativa robusta alle strategie di instradamento uniformi o fisse, in particolare per compiti di generazione e recupero a contesto lungo.