Riepilogo Tecnico: Full-Bandwidth Transformer

Definizione del Problema

I transformer autoregressivi computano lungo due assi: orizzontalmente attraverso i token generati e verticalmente attraverso la profondità del modello. Mentre l'attenzione densa fornisce un ampio accesso orizzontale ai token passati, il canale di feedback verticale tra i passi di decodifica rimane stretto. Nella decodifica standard, il modello comprime l'intero stato nascosto $D$-dimensionale dell'ultimo strato in un singolo token discreto (che trasporta al massimo $\log_2 |V|$ bit) per servirlo come input per il passo successivo. Di conseguenza, la computazione non verbalizzata (attivazioni intermedie, incertezza, piani parziali) rimane "congelata nella profondità" nel KV cache; essa rimane accessibile agli strati sopra il suo punto di produzione, ma non può tornare alla base dello stack per un'ulteriore elaborazione. Ciò costringe il modello a dover o verbalizzare gli stati intermedi token per token o a ricomputarli da zero, limitando l'efficacia della profondità di ragionamento per ogni token.

Metodologia: Decodifica del Feedback Latente

Gli autori propongono il Full-Bandwidth Transformer, che allarga il canale di feedback inter-step iniettando lo stato nascosto dell'ultimo strato precedente nell'input insieme all'embedding del token campionato.

1. Meccanismo di Feedback Latente

Ad ogni passo di decodifica $t$, l'input $u_t$ è formato fondendo l'embedding del token campionato $e_t$ con l'ultimo stato nascosto del passo precedente $h^L_{t-1}$:
$$u_t = e_t \otimes h^L_{t-1}$$
L'operatore di fusione $\otimes$ è una Gated Linear Unit (GLU):
$$e_t \otimes h^L_{t-1} = (W_U h^L_{t-1}) \odot \sigma(W_G e_t)$$

• Asimmetria: Lo stato nascosto occupa il percorso del valore, mentre il token agisce come un gate moltiplicativo. Questo design impedisce al modello di apprendere a ignorare il percorso dello stato (un potenziale scorciatoia nella fusione simmetrica), forzandolo a utilizzare l'intero stato nascosto.
• Architettura: Questo aggiunge solo due moltiplicazioni di matrici $D \times D$ per ogni token. L'architettura standard del transformer, il layout del KV cache e lo stack di serving rimangono invariati.
• Effetto: Ciò consente agli stati non verbalizzati di rientrare alla base dello stack con un budget di profondità rinnovato. Ogni strato, inclusi quelli più superficiali, può ora accedere alla storia completa come elaborata dall'intero stack, anziché solo dagli strati sottostanti.

2. Strategia di Addestramento Parallelo

L'addestramento di questa ricorrenza direttamente farebbe perdere il parallelismo del teacher forcing, poiché l'input alla posizione $t$ dipende dall'output al tempo $t-1$. Gli autori introducono un obiettivo multi-passo per risolvere questo problema:

• Parallelismo Temporale: Invece di srotolare sequenzialmente sui token, il modello esegue $K$ passaggi forward paralleli sulla sequenza.
  • Pass 1: Passaggio forward standard con plain token embeddings.
  • Pass $k$ ($k > 1$): Gli stati nascosti dal Pass $k-1$ sono traslati a destra di una posizione, fusi con gli embedding dei token e reinseriti nello stack in parallelo.
• Funzione di Perdita: L'obiettivo combina la standard perdita di Next-Token Prediction (NTP) dal primo passaggio con perdite NTP ausiliarie dai passaggi successivi. I gradienti dai passaggi successivi backpropagano negli stati latenti precedenti, incoraggiandoli a essere input riutilizzabili.
• Scheduling: Per garantire la stabilità e gestire i costi computazionali, gli autori utilizzano uno scheduling progressivo:
  • L'addestramento inizia con obiettivi a singolo passaggio standard.
  • Il feedback latente viene introdotto a metà addestramento (ad esempio, batch a 2 passaggi).
  • Una piccola frazione (ad esempio, il 3%) di batch con un numero maggiore di passaggi (ad esempio, 3 passaggi) viene miscelata. Empiricamente, questa piccola frazione stabilizza la mappa di feedback appresa, rendendola una contrazione verso un punto fisso che rimane stabile anche quando estrapolata a orizzonti lunghi (centinaia di passaggi) durante l'inferenza.
• Prefix Mix-in: Per gestire il mismatch della distribuzione tra l'addestramento (input completamente fusi) e l'inferenza (prompt semplice + generazione fusa), il processo di addestramento randomizza la lunghezza del prefisso dove iniziano gli input fusi, rispecchiando la struttura prompt-poi-genera dell'inferenza.

3. Tecniche di Stabilità

Per prevenire la divergenza durante l'inferenza a lungo orizzonte, gli autori impiegano:

• Scaling della Profondità: Assicurare che la norma dello stato dell'ultimo strato rimanga $O(1)$ anziché crescere con la profondità.
• Weight Tying: Legare i pesi del layer di embedding e del layer di readout per garantire la compatibilità tra input di token semplici e stati fusi.
• Regolarizzazione del Rumore: Aggiungere un piccolo rumore di jitter allo stato nascosto trasportato durante l'addestramento per migliorare la robustezza.

Risultati Chiave

Gli autori hanno addestrato modelli da 1 miliardo di parametri fino a 400 miliardi di token e li hanno valutati contro transformer standard addestrati su 100 miliardi - 1 trilione di token.

• Efficienza dei Dati: I Full-Bandwidth Transformer eguagliano o si avvicinano alle prestazioni dei transformer standard addestrati con circa 1,5x - 2x più token. Ad esempio, un modello full-bandwidth da 200 miliardi di token si avvicina alle prestazioni di un baseline standard da 400 miliardi di token.
• Guadagni di Performance:
  • Validation Loss: Migliorata a tutte le scale.
  • Ragionamento e Coding: Nei compiti di generazione libera (GSM8K, Math500, HumanEval, MBPP), la decodifica con feedback latente (regimi SOFT e FUSED) supera costantemente la decodifica standard con gli stessi pesi.
    • SOFT (Feedback Latente): Migliore per i compiti matematici, suggerendo che lo stato nascosto trasporta il ragionamento meglio dei token verbalizzati.
    • FUSED (Multi-pass Prefill): Migliore per i compiti di coding, suggerendo che rappresentazioni raffinate del prompt sono cruciali.
  • Instruction Tuning: I guadagni persistono dopo l'estensione del contesto lungo e l'instruction tuning.
• Concisione: Sui modelli base (pre-instruction tuning), il feedback latente produce tracce di ragionamento marcatamente più brevi a parità di accuratezza. Il modello esegue la computazione "silenziosamente" nello stato nascosto invece di narrarla token per token. Questo effetto diminuisce dopo l'instruction tuning, probabilmente perché i dati di tuning imitano un ragionamento verboso standard.
• Accessibilità dello Stato: Esperimenti di probing confermano che il prefilling ricorrente espone le informazioni di stato globale agli strati superficiali (Layer 0) molto più efficacemente del prefilling standard, validando la tesi del "full-bandwidth".

Significato e Rivendicazioni

Il paper sostiene che il Full-Bandwidth Transformer offre un nuovo asse di scaling per i Large Language Models: allocare più computazione per token anziché semplicemente più token o parametri.

• Rivisitare lo Scaling: Mentre i dati unici di alta qualità diventano un vincolo, gli autori sostengono la necessità di estrarre un segnale di apprendimento maggiore da ogni token. Il feedback latente raggiunge questo obiettivo permettendo al modello di iterare sul proprio stato interno senza il collo di bottiglia della verbalizzazione.
• Efficienza di Inferenza: L'approccio fornisce guadagni di performance con un overhead di decodifica per token trascurabile (solo due moltiplicazioni di matrici). La computazione aggiuntiva è sostenuta principalmente durante l'eventuale multi-pass prefilling o l'addestramento, non durante la generazione sequenziale di ogni token.
• Minimalismo Architetturale: Il metodo preserva l'architettura standard del transformer e il KV cache, richiedendo solo un cambiamento nella costruzione dell'input e uno specifico schema di addestramento.
• Stabilità: Il lavoro dimostra che i loop di feedback ricorrenti possono essere addestrati efficacemente nei transformer utilizzando obiettivi multi-pass paralleli, superando l'instabilità spesso associata all'addestramento di reti ricorrenti.

Gli autori concludono che i Full-Bandwidth Transformer rappresentano un passaggio dal dipendere esclusivamente dalla scala dei dati al leva di una computazione basata su feedback e ricorrenza per migliorare la qualità della rappresentazione e l'efficienza del ragionamento.