Riassunto Tecnico: Predizione Latente Gerarchica per Modelli di Linguaggio

Definizione del Problema

La standard Next-Token Prediction (NTP) con addestramento teacher-forced ha permesso lo sviluppo di modelli linguistici transformer altamente capaci. Tuttavia, questo paradigma introduce un disallineamento tra addestramento e inferenza noto come exposure bias (bias di esposizione). Durante l'inferenza, la generazione auto-regressiva si affida agli output precedenti del modello piuttosto che alla verità di base (ground truth), portando a errori cumulativi che degradano la qualità della generazione a lungo raggio.

Recenti tentativi per mitigare questo problema, come la Multi-Token Prediction (MTP) e la Next-Latent Prediction (NextLat), introducono obiettivi ausiliari per predire token futuri o stati latenti. Tuttavia, questi approcci presentano dei limiti:

• MTP soffre spesso di errori cumulativi durante i rollout multi-step.
• NextLat impone l'auto-predizione in uno spazio latente ma tipicamente predice solo un passo in avanti. Sebbene sia efficace per la coerenza delle transizioni locali, fornisce solo una pressione di gradiente indiretta per catturare la struttura su scale temporali più lunghe, lasciando le dipendenze a lungo orizzonte attenuate dalla dinamica sequenziale che si dispiega.

La sfida centrale è come modellare la rappresentazione di un modello di linguaggio per codificare strutture predittive multi-scala durante il pretraining senza incorrere in overhead architettonico o latenza di inferenza.

Metodologia: Hierarchical Latent Prediction (HiLP)

Gli autori propongono la Hierarchical Latent Prediction (HiLP), un metodo di addestramento che introduce una gerarchia temporale nell'obiettivo di auto-predizione nello spazio latente. L'architettura consiste in quattro mappe parametrizzate apprese congiuntamente: una testa di output ($p_\theta$), dinamiche di basso livello ($p_\psi$), dinamiche di alto livello ($p_\phi$) e una testa NTP combinata ($p_\rho$).

Componenti dell'Architettura

1. Latenti di Basso Livello ($h_t$): Un transformer backbone $G_\theta$ produce latenti di basso livello $h_t$ dal prefisso di token $X_{1:t}$.
2. Latenti Astratti di Alto Livello ($u_t$): Un modulo di sliding-window attention (SWA) con dimensione della finestra $k$ aggrega gli ultimi $k$ stati di basso livello in un latente astratto di alto livello: $u_t = \text{SWA}(h_{t-k+1:t})$.
3. Dinamiche di Alto Livello: Un modello $p_\phi$ predice $k$ passi in avanti nello spazio astratto ($\hat{u}_{t+k}$) senza input di token, incoraggiando la pianificazione lookahead oltre le transizioni a singolo passo.
4. Testa NTP Combinata: La predizione del prossimo token è condizionata sia dal latente di basso livello $h_t$ che dal latente di alto livello $\tilde{u}_t$ (dove $\tilde{u}_t$ include un stop-gradient per evitare che la testa combinata aggiorni i latenti di basso livello attraverso la finestra SWA).

Obiettivi di Addestramento

La perdita totale è una somma pesata di cinque termini:

1. Standard NTP Loss ($L_{ntp}$): Cross-entropy standard per la predizione del prossimo token.
2. Lower-Level Transition Consistency ($L_h$): Perdita SmoothL1 che assicura che il modello di dinamica di basso livello $p_\psi$ predica accuratamente il successivo stato latente di basso livello.
3. KL Divergence ($L_{KL}$): Assicura che la distribuzione dei latenti predetti corrisponda alla distribuzione target.
4. Higher-Level Transition Consistency ($L_u$): Perdita SmoothL1 che assicura che il modello di dinamica di alto livello $p_\phi$ predica accuratamente il latente astratto $k$ passi in avanti.
5. Combined NTP Loss ($L_{cntp}$): Cross-entropy per la testa combinata, sfruttando sia $h_t$ che $\tilde{u}_t$.

Strategia di Inferenza: Una caratteristica di design critica di HiLP è che l'intero apparato gerarchico (modulo SWA, modelli di dinamica, testa combinata) viene scartato al momento dell'inferenza. Viene utilizzata solo la testa NTP standard, garantendo zero overhead architettonico e nessun aumento della latenza di inferenza.

Contributi Chiave

1. Addestramento della Rappresentazione Gerarchica: L'introduzione di HiLP, che integra l'apprendimento auto-predittivo multi-scala nel pretraining del transformer tramite sliding-window attention sui latenti, un modello di dinamica di alto livello e una testa NTP combinata.
2. Valutazione: Benchmarking completo dei modelli risultanti sulla precisione nei task downstream (coding e ragionamento multi-step) e sull'efficienza della speculative decoding, confrontando HiLP con MTP e NextLat.
3. Intervento Puro al Tempo di Addestramento: Dimostrazione che la complessa struttura gerarchica serve puramente come segnale di addestramento ausiliario, permettendo al modello di mantenere l'efficienza del decoding auto-regressivo standard durante il deployment.

Risultati Sperimentali

Gli autori hanno addestrato modelli da 1 miliardo di parametri su 100 miliardi di token utilizzando 8 GPU NVIDIA B200.

Performance nei Benchmark

• Coding (HumanEval): HiLP ha raggiunto un punteggio pass@1 di 11.33, superando NTP standard (8.77), MTP (9.21) e NextLat (10.58).
• Ragionamento Multi-Step: Sui benchmark simbolici e multi-step (DataComp per LLM), HiLP ha dimostrato prestazioni migliorate, attribuite alla predizione latente multi-scala.

Speculative Decoding

HiLP ha mostrato una superiore efficienza nei compiti di speculative decoding:

• Dati Code: HiLP ha ottenuto una media di 3.41 token accettati per ogni step di drafting, rispetto a 2.57 per MTP e 3.26 per NextLat.
• Dati Nemotron-ClimbMix: HiLP ha ottenuto 2.59 token accettati, superando leggermente NextLat (2.57) e MTP (2.11).
• Stabilità del Rollout Latente: L'analisi delle curve di cross-entropy latente ha rivelato che, mentre HiLP mantiene una perdita di predizione a breve termine comparabile a NextLat, riduce significativamente l'errore ad orizzonti più lunghi (fino a 8 step), indicando rollout latenti multi-step più stabili.

Efficienza

Sebbene HiLP aggiunga parametri ausiliari (totale 1.27B rispetto a 1.06B per il backbone), il percorso di inferenza rimane identico al baseline. Il throughput di addestramento per HiLP è stato circa il 65% del baseline NTP, che è superiore a MTP (49%) ma inferiore a NextLat (~83%), riflettendo il costo del grafo di perdita gerarchica aggiuntivo.

Significato e Rivendicazioni

Il paper sostiene che HiLP affronta con successo il "gap di miopia" nella next-token prediction introducendo la gerarchia temporale direttamente nell'obiettivo di auto-predizione nello spazio latente. A differenza dei precedenti modelli gerarchici che consumano complessità al tempo di inferenza (come i Large Concept Models o i Byte Latent Transformers), HiLP modella la rappresentazione durante il pretraining per codificare una struttura predittiva multi-scala, che viene poi rimossa interamente per l'inferenza.

Gli autori pongono l'ipotesi che questo approccio permetta al modello di apprendere dinamiche di transizione latente più ricche e di catturare meglio le dipendenze a lungo orizzonte senza sacrificare la velocità o la semplicità del decoding auto-regressivo standard a livello di token. Il lavoro si posiziona all'intersezione tra obiettivi di pretraining oltre la NTP, la modellazione del linguaggio nello spazio latente e i modelli gerarchici, offrendo un metodo innovativo per migliorare le capacità di ragionamento e pianificazione a lungo orizzonte nei modelli di linguaggio.