Riepilogo Tecnico: Hidden Decoding su Larga Scala

Definizione del Problema

La scalabilità dei Large Language Models (LLM) si è tradizionalmente basata sull'ampliamento del backbone Transformer, il che richiede cicli costosi di pre-addestramento e aumenta i costi di servizio. Per i modelli di fondazione già robusti, l'ulteriore espansione del backbone è spesso impraticabile. Sebbene studi recenti suggeriscano che allocare più computazione per token (anche con parametri fissi) possa migliorare il ragionamento e l'accuratezza, gli approcci esistenti affrontano significativi ostacoli di scalabilità:

• Transformer a Profondità Ricorrente (Looped): Questi riutilizzano gli stessi blocchi Transformer più volte per ogni token per aumentare la computazione. Tuttavia, sono difficili da scalare verso modelli frontier (es. 100B+ parametri) perché la natura iterativa della computazione a loop interrompe le assunzioni di parallelismo di pipeline utilizzate per addestrare i modelli più grandi, causando tempi di inattività delle GPU.
• Espansione della Lunghezza della Sequenza: Alimentare input più lunghi è compatibile con le pipeline di addestramento standard, ma richiede un metodo per espandere la computazione senza limitarsi ad aumentare la finestra di contesto per lo stesso token.

Il documento affronta la sfida di migliorare un esistente backbone Transformer aumentando la computazione per token senza alterare la profondità o la larghezza del backbone, mantenendo al contempo la compatibilità con il parallelismo di pipeline e le infrastrutture di addestramento su larga scala.

Metodologia: Hidden Decoding

Gli autori propongono l'Hidden Decoding, un metodo di scalabilità della lunghezza della sequenza applicato durante il Continued Pretraining (CPT). Il meccanismo centrale espande ogni token di input in $n$ flussi paralleli elaborati all'interno di un singolo passaggio in avanti (forward pass).

1. Espansione dei Token Multi-Stream

• Espansione dell'Embedding: Invece di una singola tabella di embedding del vocabolario, il modello utilizza $n$ tabelle di embedding indipendenti ($E_1, \dots, E_n$). Un token di input $x_i$ viene espanso in una sequenza di $n$ rappresentazioni: $(E_1(x_i), E_2(x_i), \dots, E_n(x_i))$.
• Costruzione della Sequenza: Queste rappresentazioni sono intercalate per formare una sequenza espansa di lunghezza $nL$ (dove $L$ è la lunghezza originale della sequenza).
• Obiettivo di Addestramento: Solo l'ultimo stream ($E_n$) è supervisionato. Il modello predice il token successivo $x_{i+1}$ basandosi sullo stato nascosto della posizione dell'ultimo stream. I primi $n-1$ stream agiscono come "stati di computazione latente" senza perdita diretta, permettendo loro di raffinare internamente la rappresentazione prima della predizione finale.

2. Attenzione Fattorizzata per Stream (Stream-Factorized Attention)

Per evitare che il costo dell'attenzione cresca quadraticamente ($O(n^2 L^2)$) con la lunghezza della sequenza espansa, gli autori introducono la Stream-Factorized Attention:

• Layer Intra-Stream: La maggior parte dei layer limita l'attenzione alle posizioni all'interno dello stesso stream. Ciò riduce la complessità a circa $O(nL^2)$.
• Layer Cross-Stream: Un piccolo sottoinsieme di layer (ad esempio, full attention o sliding-window layers dal modello base) consente il mixing tra gli stream.
• Efficienza: Questo design mantiene il costo aggiuntivo dell'attenzione vicino alla linearità rispetto a $n$, anziché quadratico, rendendo il metodo addestrabile per modelli con oltre 100B parametri.

3. Espansione Progressiva e Inizializzazione

Per gestire l'addestramento di grandi fattori di espansione, gli autori impiegano l'Espansione Progressiva (es. $n=1 \to 2 \to 4 \to 8$). Utilizzano la Cyclic Replication Initialization, in cui le nuove tabelle di embedding vengono inizializzate copiando quelle esistenti (es. $E_{n+k} \leftarrow E_k$). Ciò consente al modello di riprodurre inizialmente il comportamento della configurazione più piccola e poi di adattarsi fluidamente.

4. Ottimizzazioni Architetturali (Specifiche per WeLM)

Per il backbone WeLM, gli autori sfruttano un design KV-Mirror. Poiché gli stream intermedi non portano alcuna perdita e contribuiscono solo come contesto tramite Chiavi e Valori (KV), i layer specchiati (che calcolano KV dai layer precedenti) possono saltare completamente l'elaborazione degli stream intermedi. Ciò riduce ulteriormente il tempo di addestramento di circa il 20% nella loro configurazione di modello da 80B.

Contributi Chiave

1. Scalabilità del Backbone Fisso su Scala Frontier: Il documento dimostra la prima applicazione di successo della scalabilità della lunghezza della sequenza alla scala MoE di 100B+. Hanno addestrato WeLM-HD4-80B e WeLM-HD4-617B (con $n=4$), ottenendo miglioramenti delle prestazioni rispetto alle controparti non-HD senza ingrandire il backbone Transformer.
2. Espansione Efficiente della Lunghezza della Sequenza: L'introduzione della Stream-Factorized Attention consente l'addestramento con sequenze espanse a un costo quasi lineare in $n$ (ad esempio, un'espansione di 4x costa circa 4.4x–5.1x per batch, ben al di sotto del costo di 16x di un'attenzione densa).
3. Scalabilità del Fattore di Espansione: Gli esperimenti mostrano che aumentare il fattore di espansione $n$ (da 2 a 8) migliora costantemente la perdita di linguaggio (language modeling loss) e l'accuratezza downstream, validando $n$ come un pratico parametro di scalabilità per backbone fissi.

Risultati Sperimentali

Gli autori hanno valutato l'Hidden Decoding su modelli MoE di scala frontier e su baseline dense:

• Prestazioni su Scala Frontier (WeLM-HD4-80B & 617B):

  • WeLM-HD4-80B ha superato WeLM-80B in tutti i nove benchmark condivisi, con guadagni significativi in SciCode (+4.2), PHYBench (+4.0) e FrontierMath (+3.2).
  • WeLM-HD4-617B ha migliorato le prestazioni rispetto a WeLM-617B in tutti i nove benchmark, inclusi GPQA Diamond (+2.1), HLE (+1.8) e FrontierMath (+2.0).
  • Questi modelli sono stati addestrati utilizzando una ricetta di post-addestramento precoce basata solo su SFT per garantire un confronto controllato, isolando l'effetto dell'Hidden Decoding.

• Scalabilità con il Fattore di Espansione ($n$):

  • Sul modello denso Qwen3-8B-Base, aumentare $n$ da 1 a 8 ha portato a miglioramenti monotoni nell'accuratezza media (es. MMLU è migliorato da 78.2 a 81.6).
  • Sul MoE da 80B, aumentare $n$ da 1 a 8 ha elevato l'MMLU da 85.0 a 87.5 e ridotto il Pile-test BPB da 0.386 a 0.378.

• Studi di Ablazione:

  • Design della Supervisione: Supervisionare solo l'ultimo stream ha prodotto risultati migliori rispetto all'applicazione della perdita a tutti gli stream o alla somma degli output degli stream.
  • Ritenzione KV: Mantenere cache KV separate per ogni stream (invece di condividerle) ha migliorato l'accuratezza, confermando che gli stream intermedi forniscono un contesto distinto e utile.
  • Composizione dell'Attenzione: Modelli con solo pochi layer cross-stream completi (es. 4 su 27) hanno performato quasi quanto i modelli con full cross-stream mixing in tutti i layer, validando l'efficienza della Stream-Factorized Attention.

• Throughput di Inferenza:

  • L'Hidden Decoding mantiene un alto throughput in scenari con bassa latenza e piccoli batch (mantenendo l'83–88% del throughput della baseline con batch size 1).
  • Il throughput diminuisce in scenari con contesti lunghi e batch elevati (scendendo a circa il 27% con batch size 32 e 32k di input), riflettendo l'aumento della computazione per token.

Significato e Rivendicazioni

Il documento afferma che l'Hidden Decoding fornisce una via pratica e ingegneristicamente realistica per scalare i LLM frontier quando l'ulteriore espansione del backbone è troppo costosa.

• Compatibilità: A differenza dei Transformer a loop, l'Hidden Decoding si integra naturalmente con il parallelismo di pipeline e gli standard di addestramento per grandi modelli perché la computazione aggiuntiva si manifesta come una sequenza più lunga piuttosto che come profondità iterativa.
• Efficacia: È il primo metodo dimostrato che riesce a scalare con successo l'espansione della lunghezza della sequenza nel regime MoE da 100B+, producendo miglioramenti misurabili nei benchmark di ragionamento e conoscenza.
• Scalabilità: Il metodo dimostra che il fattore di espansione $n$ è un asse di scalabilità valido, offrendo guadagni di prestazione che crescono con $n$ senza richiedere nuove architetture di modello o un massiccio ri-addestramento del backbone.

Gli autori pongono questo lavoro come una soluzione al "muro di scalabilità" affrontato dai forti modelli di fondazione, offrendo un modo per estrarre maggiori capacità dai parametri esistenti attraverso una computazione latente efficiente.