Riassunto Tecnico: Multi-Resolution Residual Routing per Deep Decoder-Only Transformers

Problematica

I Transformer standard decoder-only si affidano a connessioni residue PreNorm in cui gli aggiornamenti dei sottostrati vengono aggregati con pesi unitari fissi. Sebbene questa struttura consenta un addestramento stabile, impone una regola di accumulo rigida che può diluire il contributo dei singoli strati all'aumentare della profondità della rete. Recenti progressi, come gli Attention Residuals e i Block Attention Residuals (Block AttnRes), affrontano questo problema sostituendo l'accumulo fisso con un routing basato sul contenuto e appreso sugli output precedenti. Block AttnRes migliora specificamente l'efficienza comprimendo più strati in riassunti residui a livello di blocco ($C_b$).

Tuttavia, gli autori identificano un limite critico in Block AttnRes: un singolo riassunto di blocco rappresenta solo la componente a bassa frequenza (DC) della traiettoria residua all'interno di un blocco. Memorizzando solo la somma degli aggiornamenti ($\sum u_{b,i}$), il meccanismo scarta la preziosa struttura direzionale, come:

1. Squilibrio tra i sottostrati: Se il blocco è dominato dagli aggiornamenti di Attention o MLP.
2. Dinamiche temporali: Se gli aggiornamenti iniziali all'interno del blocco puntano in una direzione diversa rispetto a quelli tardivi.

Gli autori ipotizzano che, man mano che la lunghezza della traiettoria residua cresce, questa informazione direzionale perduta diventi sempre più vitale per un efficace routing depth-wise.

Metodologia: WAV v1

Gli autori propongono WAV v1 (Waveform Attention v1), un'estensione multi-risoluzione leggera di Block AttnRes. Invece di rappresentare un blocco esclusivamente tramite la sua somma totale, WAV v1 aumenta ogni rappresentazione di blocco con due basi di dettaglio direzionale a somma zero, preservando al contempo i moduli originali di Attention e MLP.

1. Rappresentazione Multi-Risoluzione

Per un blocco $b$ contenente $m$ aggiornamenti di sottostrato $\{u_{b,i}\}_{i=1}^m$, WAV v1 calcola:

• State Basis ($C_b$): La somma standard del blocco, che rappresenta lo stato a bassa frequenza.
  $$C_b = \sum_{i=1}^m u_{b,i}$$
• Phase Basis ($D^{phase}_b$): Una sorgente di dettaglio che contrasta gli aggiornamenti di Attention vs. MLP. Assegna un segno $+1$ agli aggiornamenti di Attention e $-1$ agli aggiornamenti di MLP.
  $$D^{phase}_b = \sum_{i=1}^m a_i u_{b,i}, \quad \text{dove } a_i \in \{+1, -1\}$$
• Split Basis ($D^{split}_b$): Una sorgente di dettaglio che contrasta la prima metà rispetto alla seconda metà del blocco. Assegna $+1$ alla prima metà e $-1$ alla seconda metà.
  $$D^{split}_b = \sum_{i=1}^m r_i u_{b,i}, \quad \text{dove } r_i \in \{+1, -1\}$$

Queste basi vengono accumulate online man mano che gli aggiornamenti vengono prodotti, comportando un overhead computazionale trascurabile.

2. Meccanismo di Routing

Il mixer softmax depth-wise riceve un pool di sorgenti espanso. Per Block AttnRes, il pool è $\{e, C_0, \dots, C_{b-1}, PC\}$. Per WAV v1, esso si espande includendo le basi di dettaglio per ogni blocco completato e il blocco parziale corrente:
$$S_{WAV} = \{e, C_0, \tilde{D}^{phase}_0, \tilde{D}^{split}_0, \dots, C_{b-1}, \tilde{D}^{phase}_{b-1}, \tilde{D}^{split}_{b-1}, PC, \tilde{P}^{phase}, \tilde{P}^{split}\}$$

3. Meccanismi di Stabilità

Per evitare la destabilizzazione durante le prime fasi dell'addestramento, WAV v1 impiega due strategie conservative:

• Negative Detail Bias: Le sorgenti di dettaglio sono inizializzate con un forte bias negativo ($\beta_D = -2.0$), garantendo che il modello si comporti come lo standard Block AttnRes all'inizializzazione e utilizzi i dettagli solo quando benefici.
• Detached RMS Matching: I tensori di dettaglio sono scalati rispetto al loro corrispondente riassunto di blocco ($C$) utilizzando un rapporto RMS staccato (detached), limitato da un fattore di scala massimo $\rho$. Ciò evita che i dettagli dominino semplicemente a causa di magnitudo grezze più grandi.
• Readout Conservativo: Per impostazione predefinita, il mixer di predizione finale legge solo le sorgenti di embedding e $C$, non le sorgenti di dettaglio, garantendo ulteriormente la stabilità.

Contributi Chiave

1. Formulazione: Il paper inquadra il block residual routing come un problema di rappresentazione multi-risoluzione, distinguendo tra l'informazione di stato a bassa frequenza ($C_b$) e i dettagli direzionali intra-blocco.
2. Algoritmo: Introduzione di WAV v1, un'estensione "drop-in" che aumenta i riassunti di blocco con basi di dettaglio di fase e di split senza modificare le funzioni core di Attention o MLP.
3. Evidenza Empirica: Esperimenti di scaling preliminari che dimostrano come l'utilità del routing multi-risoluzione sia dipendente dalla profondità, diventando significativamente vantaggiosa solo in architetture più profonde.

Risultati Sperimentali

Gli autori hanno valutato modelli GPT decoder-only da 12, 24 e 48 strati su dataset a livello di carattere TinyStories e Text8.

• Trend Dipendente dalla Profondità:
  • 12 Strati: WAV v1 è inferiore a Block AttnRes, suggerendo che i dettagli direzionali siano ridondanti o rumorosi in reti poco profonde.
  • 24 Strati: WAV v1 diventa competitivo, mostrando lievi miglioramenti su TinyStories e prestazioni comparabili su Text8.
  • 48 Strati: WAV v1 raggiunge la migliore loss di validazione su entrambi i dataset, superando sensibilmente Block AttnRes, ReZero e LayerScale.
• Guadagni Quantitativi (48 Strati):
  • TinyStories: Riduzione della loss di validazione da 0.4960 (Block AttnRes) a 0.4738 (WAV v1), ovvero una riduzione del 2.19% della PPL.
  • Text8: Riduzione della loss di validazione da 0.9363 a 0.9305, ovvero una riduzione dello 0.57% della PPL.
• Dinamiche di Addestramento: A 48 strati, WAV v1 si separa da Block AttnRes precocemente durante l'addestramento e mantiene un vantaggio persistente, indicando che il miglioramento non è un artefatto di fine addestramento ma un cambiamento fondamentale nella traiettoria di ottimizzazione.

Significato e Rivendicazioni

Il paper sostiene che il routing dei residui profondi beneficia non solo della selezione di quale blocco leggere, ma anche dal routing della struttura direzionale interna di ogni blocco. I risultati suggeriscono che, man mano che la traiettoria residua si allunga, la "forma" della traiettoria di aggiornamento (ad esempio, dominanza dell'attenzione o dinamiche early-vs-late) diventa un segnale informativo per il mixer depth-wise.

Gli autori posizionano WAV v1 come un approccio complementare ai metodi esistenti:

• A differenza di ReZero o LayerScale, che modulano la magnitudo degli aggiornamenti, WAV v1 cambia la rappresentazione disponibile al router.
• A differenza di Delta Attention Residuals, che effettuano il routing su singoli delta (costo elevato), WAV v1 offre un'approssimazione a livello di blocco e multi-risoluzione che bilancia efficienza e ricchezza strutturale.

Gli autori concludono che le valutazioni future dei metodi di residual-routing dovrebbero considerare la profondità, poiché i benefici dell'informazione multi-risoluzione possono essere legati alla lunghezza e alla struttura della traiettoria residua. Riconoscono le attuali limitazioni, inclusa la mancanza di deviazione standard tra i seed, la restrizione a piccoli modelli a livello di carattere e l'uso di basi di dettaglio progettate manualmente, suggerendo che il lavoro futuro dovrebbe esplorare basi apprendibili e modelli a livello di token su scala maggiore.