Riepilogo Tecnico: Residui di Attenzione con Ruoli Disaccoppiati (Role-Decoupled Attention Residuals)

Definizione del Problema

Nelle architetture Transformer standard, lo stream residuo agisce come una memoria fissa dove ogni strato eredita lo stato immediatamente precedente tramite addizione seriale. Le recenti architetture di routing della profondità, come i Block Attention Residuals (Block AttnRes), migliorano questo approccio permettendo ai sottostrati di recuperare un mix dipendente dal contenuto di rappresentazioni precedenti, anziché solo dal predecessore immediato. Ciò trasforma la profondità da un ordine computazionale fisso in un asse di recupero appreso.

Tuttavia, Block AttnRes introduce un vincolo di progettazione specifico: utilizza una singola distribuzione di profondità per costruire l'input condiviso per tutte e tre le proiezioni di attenzione (Query, Key e Value). Questo accoppia due ruoli funzionalmente distinti:

1. Geometria di Corrispondenza (Matching Geometry): Determinata da Query e Key, che decidono dove guardare (quali posizioni di token corrispondono).
2. Recupero del Contenuto (Content Retrieval): Determinato dai Value, che decidono cosa recuperare.

Il paper sostiene che forzare la rappresentazione migliore per decidere "dove guardare" a essere anche la migliore per decidere "cosa leggere" sia un accoppiamento non necessario. La domanda centrale è se il disaccoppiamento di questi ruoli migliori il pretraining dei modelli linguistici a parità di altri fattori.

Metodologia: Residui di Attenzione con Ruoli Disaccoppiati (RD-AttnRes)

L'autore introduce i Role-Decoupled Attention Residuals (RD-AttnRes), un rilassamento architettonico minimo di Block AttnRes.

• Architettura: RD-AttnRes mantiene l'insieme di sorgenti dell'architettura genitore, la topologia dei blocchi, la normalizzazione, i percorsi MLP e i meccanismi di writeback. L'unica modifica avviene nella lettura della profondità sul lato attenzione.
• Meccanismo di Disaccoppiamento:
  • Query e Key (Q/K): Condividono una singola query di percorso appresa ($q^{QK}_\ell$) per calcolare un mix ($h^{QK}_\ell$) dal set di memoria della profondità residua. Ciò preserva una geometria di corrispondenza coerente.
  • Value (V): Imparano una query di percorso indipendente ($q^V_\ell$) per calcolare un mix separato ($h^V_\ell$) sopra lo stesso identico set di sorgenti residue.
• Efficienza dei Parametri: La modifica aggiunge solo un vettore di dimensione $d$ (larghezza del modello) per ogni strato. Per le scale testate (120M e 343M di parametri), ciò rappresenta un aumento trascurabile dei parametri (<0,01%).
• Ripristino Esatto: Se la query di percorso del valore è legata alla query Q/K ($q^V_\ell = q^{QK}_\ell$), RD-AttnRes si riduce esattamente all'architettura Block AttnRes genitrice. Ciò garantisce che la proposta sia un'estensione rigorosa e verificabile piuttosto che un sistema con molteplici cambiamenti simultanei.
• Costo Computazionale: Il metodo aggiunge un lavoro di $O(BTK_\ell d)$ per ogni strato di attenzione (dove $B$ è la dimensione del batch, $T$ la lunghezza della sequenza e $K_\ell$ il numero di sorgenti), ma non introduce nuovi termini di attenzione token-to-token $O(T^2d)$.

Disegno Sperimentale

Lo studio impiega un protocollo sperimentale accoppiato e rigoroso per isolare l'effetto del disaccoppiamento dei ruoli:

• Dataset: FineWeb-Edu (campione da 10BT), con una divisione deterministica per training, validazione e test.
• Modelli: Due modelli Transformer Pre-Norm decoder-only di scala:
  • S: 120M di parametri (16 strati, larghezza 576).
  • M: 343M di parametri (24 strati, larghezza 832).
• Protocollo: Cinque seed corrispondenti per scala. Sia Block AttnRes che RD-AttnRes sono addestrati con identica inizializzazione, ordine dei dati, programmi di ottimizzazione e un budget di 2.0B di token.
• Controlli: Controlli direzionali a 0.5B di token includono:
  • Tied (Legato): Forza la query di percorso del valore a corrispondere alla query Q/K (ripristinando il genitore).
  • Average (Medio): Esegue due percorsi ma media i loro mix (testa se la doppia esecuzione da sola aiuti).
  • Frozen-V (V Congelato): Mantiene la query di percorso del valore uniforme (testa se una query fissa aiuti).
• Diagnostica: La divergenza di Jensen–Shannon (JS) è misurata tra le distribuzioni di profondità di Q/K e V per verificare se i percorsi imparano strategie distinte.

Risultati Chiave

1. Miglioramento della Loss di Validazione: RD-AttnRes migliora la negative log-likelihood (NLL) di validazione in tutti i 10 confronti accoppiati (5 seed × 2 scale).
   • Scala 120M: Riduzione media della NLL di 0,0301, corrispondente a una riduzione della perplexity del 2,97%.
   • Scala 343M: Riduzione media della NLL di 0,0247, corrispondente a una riduzione della perplexity del 2,43%.
2. Persistenza del Guadagno: Il miglioramento è visibile in tutte le pietre miliari di valutazione (0,25B, 0,5B, 1,0B e 2,0B di token), sebbene l'entità assoluta del divario si restringa leggermente con l'avanzare del training.
3. Validazione dei Controlli:
   • I controlli Tied e Average performano in modo simile al genitore Block AttnRes, indicando che il guadagno non è dovuto al numero di parametri o alla duplicazione del calcolo.
   • Il controllo Frozen-V performa male, confermando che la query di percorso del valore deve essere appresa, non fissa.
4. Divergenza dei Percorsi: Le distribuzioni apprese di Q/K e V non collassano alla stessa soluzione. La divergenza JS media tra gli strati è di circa 0,218 (120M) e 0,215 (343M), confermando che il modello impara letture di profondità distinte per la corrispondenza rispetto al recupero del contenuto.

Significato e Rivendicazioni

Il paper formula una rivendicazione stretta e limitata piuttosto che una universale. L'autore non sostiene che RD-AttnRes sia superiore a tutte le architetture di routing dei residui o che stabilisca un comportamento di scaling universale.

Inveve, i risultati supportano la conclusione specifica che nell'ambito del regime di addestramento testato (FineWeb-Edu, 2B di token, scale specifiche del modello), la condivisione della lettura di profondità QKV in Block AttnRes è un vincolo consequenziale. L'evidenza suggerisce che la corrispondenza (Q/K) e il recupero del contenuto (V) beneficiano di letture di profondità differenti.

Limitazioni Riconosciute:

• Lo studio è limitato a un corpus web in inglese e a modelli con meno di un miliardo di parametri.
• Il divario di performance assoluto si restringe su orizzonti di addestramento più lunghi; ricerche future sono necessarie per determinare se il divario persista o cresca a 4B–10B di token.
• Esiste una nota sistemistica: l'implementazione attuale comporta una penalità di throughput (~9–10%) e un aumento dell'uso della memoria rispetto al riferimento Block AttnRes, sebbene kernel fusi possano mitigare ciò.
• I confronti con altre architetture recenti (es. MUDDFormer, Delta AttnRes) sono concettuali piuttosto che basati su baseline empiriche corrispondenti.

Conclusione:
Il paper stabilisce che quando la profondità diventa un asse di recupero appreso, la condivisione dovrebbe seguire la funzione a valle (ruolo) piuttosto che la convenienza di implementazione. Disaccoppiando il percorso del valore con un overhead di parametri minimo, RD-AttnRes ottiene miglioramenti costanti nella loss di pretraining, convalidando l'ipotesi che la geometria di corrispondenza e il recupero del contenuto siano compiti distinti che richiedono rappresentazioni di profondità diverse.