Technische Samenvatting: Hiërarchische Latente Predictie voor Taalmodellen

Probleemstelling

Standaard Next-Token Prediction (NTP) met teacher-forced training heeft de ontwikkeling van zeer capabele transformer-taalmodellen mogelijk gemaakt. Dit paradigma introduceert echter een mismatch tussen training en inferentie, bekend als exposure bias. Tijdens de inferentie vertrouwt autoregressieve generatie op de eigen eerdere outputs van het model in plaats van op de grondwaarheid (ground truth), wat leidt tot cumulatieve fouten die de kwaliteit van langetermijngeneratie verslechteren.

Recente pogingen om dit te verzachten, zoals Multi-Token Prediction (MTP) en Next-Latent Prediction (NextLat), introduceren hulpdoelen om toekomstige tokens of latente toestanden te voorspellen. Deze benaderingen kampen echter met beperkingen:

• MTP lijdt vaak aan cumulatieve fouten tijdens multi-step rollouts.
• NextLat dwingt zelf-predictie af in een latente ruimte, maar voorspelt doorgaans slechts één stap vooruit. Hoewel dit effectief is voor lokale transitieconsistentie, biedt het slechts indirecte gradiëntdruk voor het vastleggen van structuren op langere temporele schalen, waardoor langetermijnafhankelijkheden die door sequentiële dynamiek worden afgezwakt, onvoldoende worden gevangen.

De kernuitdaging is hoe men de representatie van een taalmodel kan vormgeven om multi-schaal predictieve structuren te coderen tijdens de pretraining, zonder architecturale overhead of inferentielatentie te veroorzaken.

Methodologie: Hiërarchische Latente Predictie (HiLP)

De auteurs stellen Hierarchical Latent Prediction (HiLP) voor, een trainingsmethode die temporele hiërarchie introduceert in de latent-ruimte zelf-predictie doelstelling. De architectuur bestaat uit vier gezamenlijk geleerde geparametriseerde kaarten: een output head ($p_\theta$), lower-level dynamica ($p_\psi$), upper-level dynamica ($p_\phi$) en een gecombineerde NTP-head ($p_\rho$).

Architectuurcomponenten

1. Lower-Level Latents ($h_t$): Een backbone transformer $G_\theta$ produceert lower-level latents $h_t$ vanuit de token-prefix $X_{1:t}$.
2. Higher-Level Abstract Latents ($u_t$): Een sliding-window attention (SWA) module met venstergrootte $k$ aggregeert de laatste $k$ lower-level toestanden tot een higher-level abstracte latent: $u_t = \text{SWA}(h_{t-k+1:t})$.
3. Higher-Level Dynamics: Een model $p_\phi$ voorspelt $k$ stappen vooruit in de abstracte ruimte ($\hat{u}_{t+k}$) zonder token-input, wat stimuleert tot lookahead planning voorbij single-step transities.
4. Gecombineerde NTP Head: De next-token prediction is geconditioneerd op zowel de lower-level latent $h_t$ als de higher-level latent $\tilde{u}_t$ (waarbij $\tilde{u}_t$ een stop-gradient bevat om te voorkomen dat de gecombineerde head de lower latents via het SWA-venster bijwerkt).

Trainingsdoelstellingen

De totale loss is een gewogen som van vijf termen:

1. Standaard NTP Loss ($L_{ntp}$): Standaard cross-entropy voor next-token prediction.
2. Lower-Level Transition Consistency ($L_h$): SmoothL1 loss die ervoor zorgt dat het lower dynamics model $p_\psi$ de volgende lower-level latent toestand accuraat voorspelt.
3. KL Divergentie ($L_{KL}$): Zorgt ervoor dat de distributie van voorspelde lower latents overeenkomt met de doeldistributie.
4. Higher-Level Transition Consistency ($L_u$): SmoothL1 loss die ervoor zorgt dat het higher-level dynamics model $p_\phi$ de abstracte latent $k$ stappen vooruit accuraat voorspelt.
5. Gecombineerde NTP Loss ($L_{cntp}$): Cross-entropy voor de gecombineerde head, gebruikmakend van zowel $h_t$ als $\tilde{u}_t$.

Inferentiestrategie: Een cruciaal ontwerpkenmerk van HiLP is dat de gehele hiërarchische apparatuur (SWA-module, dynamics modellen en gecombineerde head) bij inferentie wordt verworpen. Alleen de standaard NTP head wordt gebruikt, wat zorgt voor nul architecturale overhead en geen toename in inferentielatentie.

Belangrijkste Bijdragen

1. Hiërarchische Representatie Training: De introductie van HiLP, die multi-schaal zelf-predictieve leerprocessen integreert in de transformer pretraining via sliding-window attention over latents, een higher-level dynamics model en een gecombineerde NTP-head.
2. Evaluatie: Uitgebreide benchmarking van de resulterende modellen op downstream nauwkeurigheid (codering en multi-step redeneren) en de efficiëntie van speculative decoding, waarbij HiLP wordt vergeleken met MTP en NextLat.
3. Pure Training-Time Interventie: Demonstratie dat de complexe hiërarchische structuur puur dient als een auxiliary trainingssignaal, waardoor het model de efficiëntie van standaard autoregressieve decoding behoudt tijdens deployment.

Experimentele Resultaten

De auteurs trainden 1B-parameter modellen op 100B tokens met behulp van 8× NVIDIA B200 GPU's.

Benchmark Prestaties

• Codering (HumanEval): HiLP behaalde een pass@1 score van 11.33, waarmee het standaard NTP (8.77), MTP (9.21) en NextLat (10.58) overtreft.
• Multi-Step Redeneren: Op symbolische en multi-step benchmarks (DataComp voor LLM's) vertoonde HiLP verbeterde prestaties, toegeschreven aan de multi-schaal latent predictie.

Speculative Decoding

HiLP toonde superieure efficiëntie in speculative decoding taken:

• Code Data: HiLP behaalde gemiddeld 3.41 geaccepteerde tokens per drafting stap, vergeleken met 2.57 voor MTP en 3.26 voor NextLat.
• Nemotron-ClimbMix Data: HiLP behaalde 2.59 geaccepteerde tokens, waarmee het NextLat (2.57) en MTP (2.11) licht overtrof.
• Latent Rollout Stabiliteit: Analyse van de latent cross-entropy curves toonde aan dat hoewel HiLP een nabij-termijn predictie loss vergelijkbaar met NextLat behoudt, het de fout bij langere horizonten (tot 8 stappen) aanzienlijk vermindert, wat wijst op stabielere multi-step latent rollouts.

Efficiëntie

Hoewel HiLP auxiliary parameters toevoegt (totaal 1.27B vs. 1.06B voor de backbone), blijft het inferentiepad identiek aan de baseline. De training throughput voor HiLP was ongeveer 65% van de NTP baseline, wat hoger is dan MTP (49%) maar lager dan NextLat (~83%), wat de kosten reflecteert van de extra hiërarchische loss graph.

Betekenis en Claims

Het artikel claimt dat HiLP succesvol de "myope kloof" (myopic gap) in next-token prediction aanpakt door een temporele hiërarchie direct in de latent-ruimte zelf-predictie doelstelling te introduceren. In tegenstelling tot eerdere hiërarchische modellen die complexiteit consumeren tijdens de inferentie (zoals Large Concept Models of Byte Latent Transformers), vormt HiLP de representatie tijdens de pretraining om een multi-schaal predictieve structuur te coderen, die vervolgens volledig wordt verwijderd tijdens de inferentie.

De auteurs stellen dat deze aanpak het model in staat stelt om rijkere latente transitiedynamica te leren en beter langetermijn-afhankelijkheden te vatten, zonder de snelheid of eenvoud van standaard token-level autoregressieve decoding op te offeren. Het werk positioneert zich op het snijvlak van pretraining doelstellingen buiten NTP, latent-ruimte taalmodellering en hiërarchische modellen, en biedt een nieuwe methode om de langetermijnredeneer- en planningcapaciteiten van taalmodellen te verbeteren.