Technische Samenvatting: Next-Latent Prediction Transformers Leren Compacte Wereldmodellen

Probleemstelling

Standaard Transformer-architecturen vervangen recurrentie door self-attention en een geheugenmechanisme dat schaalt met de sequentielengte. Hoewel dit uitzonderlijke parallelisatie en voorspellende kracht mogelijk maakt, verwijdert het de inherente inductieve bias om geschiedenis te comprimeren tot compacte latente toestanden met consistente overgangsregels. Bijgevolg leren Transformers vaak complexe, taakspecifieke shortcuts die goed passen bij trainingsdata maar slecht generaliseren, waardoor ze geen coherente "wereldmodellen" of overtuigingsstaten vormen die noodzakelijk zijn voor planning en redeneren op lange termijn. Het standaarddoel van next-token predictie garandeert niet dat de interne verborgen toestanden van het model evolueren als toereikende statistieken voor de vorige geschiedenis, noch dat deze toestanden toekomstige observaties recursief kunnen voorspellen.

Methodologie: Next-Latent Prediction (NextLat)

De auteurs introduceren Next-Latent Prediction (NextLat), een raamwerk dat standaard next-token training aanvult met zelftoezicht in de latente ruimte. Het kernidee is om de Transformer te trainen om latente representaties (specifiek de verborgen toestanden van de laatste laag $h_t$) te leren die voorspellend zijn voor de volgende latente toestand gegeven de volgende token.

Theoretische Grondslag

Het artikel stelt vast dat als een model optimaliseert voor zowel Next-Token Consistentie (het voorspellen van $X_{t+1}$ vanuit $h_t$) als Overgangsconsistentie (het voorspellen van $h_{t+1}$ vanuit $h_t$ en $X_{t+1}$), de verborgen toestanden $h_t$ theoretisch gegarandeerd convergeren naar overtuigingsstaten. Een overtuigingsstaat wordt gedefinieerd als een toereikende statistiek van de geschiedenis $X_{1:t}$ voor het voorspellen van de toekomst $X_{t+1:T}$. In tegenstelling tot standaard autoregressieve training, die ad-hoc opzoeken van vorige tokens toestaat zonder informatie te comprimeren, dwingt NextLat een recurrente inductieve bias af waarbij de latente toestand alle noodzakelijke informatie moet bevatten om over te gaan naar de volgende toestand en de volgende token te voorspellen.

Implementatie

NextLat werkt door een lichtgewicht hulpdoel toe te voegen aan de standaard Transformer-traininglus:

1. Latente Dynamiekmodel ($p_\psi$): Een eenvoudige Multi-Layer Perceptron (MLP) wordt getraind om de volgende verborgen toestand $\hat{h}_{t+1}$ te voorspellen gegeven de huidige verborgen toestand $h_t$ en de volgende token $X_{t+1}$.
2. Verliesfunctie: Het totale verlies combineert:
   • Next-Token Verlies ($L_{next-token}$): Standaard cross-entropy voor tokenpredictie.
   • Next-Verborgen Toestand Verlies ($L_{next-h}$): Een Smooth L1-verlies dat het voorspelde staat $\hat{h}_{t+1}$ regresseert tegen de ware volgende toestand $h_{t+1}$ (met stop-gradient toegepast op het doel om representatieve instorting te voorkomen).
   • KL-Divergentie Verlies ($L_{KL}$): Een complementair doel dat ervoor zorgt dat de tokenverdeling voorspeld vanuit de voorspelde latente toestand overeenkomt met de verdeling vanuit de ware latente toestand, fungerend als een zacht toezichtsignaal.
3. Meerstaps Toezicht: Het model kan worden getraind om de latente dynamiek recursief te ontrollen voor $d$ stappen, waardoor rijkere gradiëntsignalen voor langetermijnafhankelijkheden worden geboden.

Cruciaal is dat NextLat de architectuur van de Transformer, de inferentieprocedure en de efficiëntie van parallelle training onveranderd laat. Het latente dynamiekmodel wordt uitsluitend tijdens training gebruikt om representaties te vormen; tijdens inferentie decodeert het model zoals gebruikelijk autoregressief.

Belangrijkste Bijdragen

1. Theoretische Garantie: Het artikel bewijst dat optimalisatie voor next-token en next-latent consistentie bewezen transformer-representaties vormt tot overtuigingsstaten, een eigenschap die niet gegarandeerd wordt door next-token predictie alleen.
2. Praktisch Raamwerk: NextLat introduceert een computatie-efficiënt hulpdoel dat recurrente inductieve bias injecteert zonder de parallelle training of inferentiestructuur van de Transformer te veranderen.
3. Variabele-Lengte Zelf-Speculatieve Decoding: Door coherente latente dynamiek te leren, maakt NextLat recursieve predictie in de latente ruimte mogelijk. Dit staat variabele-lengte zelf-speculatieve decoding toe, waarbij het model tokens kan opstellen buiten de vaste trainingshorizon, wat de inferentie aanzienlijk versnelt.

Experimentele Resultaten

De auteurs evalueren NextLat over vier domeinen: wereldmodelleren, redeneren, plannen en taalmodelleren.

• Wereldmodelleren (Manhattan Taxi Ritten): NextLat leert interne kaarten die aanzienlijk consistenter zijn met het ware wereldmodel (straattopologie) in vergelijking met standaard GPT-, MTP- en JTP-baselines. Het bereikt hogere sequentiecompressie (0,71 versus 0,65 voor GPT) en een veel lagere effectieve latente rang (52,7 versus 160,1 voor GPT), wat wijst op meer compacte en efficiënte representaties.
• Redeneren (Countdown): Op de taak voor wiskundig redeneren presteert NextLat aanzienlijk beter dan baselines (54,8% nauwkeurigheid bij $d=1$ versus 33,1% voor GPT). Het toont superieur vooruitkijkend plannen, wat "berouwvolle compromissen" vermindert waarbij modellen ongeldige uiteindelijke vergelijkingen forceren om een doel te matchen.
• Plannen (Path-Star Grafen): NextLat behoudt bijna 100% oplossingspercentages over verschillende grafentopologieën, terwijl baselines zoals BST en JTP falen op grotere grafen ($G_{7,7}$). Dit suggereert dat NextLat de "Clever Hans" shortcut-learning vermijdt die gebruikelijk is in token-ruimte predictie.
• Taalmodelleren (FineWeb-Edu):
  • Downstream Nauwkeurigheid: NextLat behoudt next-token perplexiteit beter dan MTP en JTP, terwijl het bescheiden winsten boekt in downstream zero-shot nauwkeurigheid (59,21% gemiddeld versus 58,82% voor GPT).
  • Inferentiesnelheid: Door gebruik te maken van variabele-lengte zelf-speculatieve decoding, bereikt NextLat tot 3,3× snellere inferentie in vergelijking met standaard autoregressieve sampling. Deze snelheidswinst blijft bestaan zelfs bij het opstellen buiten de trainingshorizon, een capaciteit die MTP en JTP missen vanwege hun vaste token-ruimte horizons.

Betekenis en Claims

Het artikel positioneert NextLat als een eenvoudig maar effectief paradigma voor het leren van compacte, voorspellende en generaliseerbare representaties in Transformers. De primaire betekenis ligt in het verenigen van de recurrente inductieve bias van compacte wereldmodellen met de schaalbaarheid en parallelisme van Transformers.

De auteurs claimen dat NextLat een oplossing biedt voor de "myopische" aard van next-token predictie, waardoor modellen worden aangemoedigd gestructureerde, voorspellende interne toestanden te vormen in plaats van te vertrouwen op oppervlakkige token-correlaties. Door variabele-lengte speculatieve decoding mogelijk te maken zonder architecturale wijzigingen, biedt NextLat een praktische weg naar snellere inferentie en beter redeneren op lange termijn. Het artikel concludeert dat NextLat een stap is naar trainingsdoelen die autoregressieve modellen toerusten met eenvoudigere, compacter en daarom generaliseerbaardere representaties van complexe systemen.

De auteurs blijven bescheiden wat betreft toekomstige richtingen, met de opmerking dat ze geen complexe latente dynamiekarchitecturen hebben onderzocht (slechts eenvoudige MLP's gebruikend) en dat het optimale ontwerp van het doel (bijvoorbeeld de noodzaak van KL-verlies op grotere schaal) verder systematisch onderzoek vereist. Ze erkennen ook dat hoewel NextLat representaties verbetert, de onderliggende circuitcomplexiteit van de vaste diepte Transformer begrensd blijft door de TC0-klasse, hoewel het meegetrainde latente dynamiekmodel (RNN) blijkbaar in staat is om deze grenzen te overstijgen.