How Transformers Learn to Plan via Multi-Token Prediction

Dit paper toont aan dat multi-token-predictie (MTP) Transformers beter leert plannen dan de standaard next-token-predictie door een schoner trainingssignaal te bieden dat een reversibel redeneringsproces met twee fasen induceert.

De kern

Stel je voor dat je een grote stad moet verkennen om een schat te vinden. Je hebt een kaart (de grafiek) en je weet waar je begint en waar de schat ligt. Maar je mag alleen één stap tegelijk zetten en moet elke stap hardop zeggen voordat je de volgende zet.

Dit is precies wat de meeste kunstmatige intelligenties (LLMs) tot nu toe deden. Ze gebruiken een methode genaamd "Next-Token Prediction" (NTP). Het is alsof ze blindelings de volgende straatnaam raden op basis van wat ze net hebben gezegd. Het probleem? Ze raken vaak verward, lopen in cirkels of volgen een vals spoor omdat ze niet kunnen "vooruitkijken" naar het einddoel. Ze denken te kort door de bocht.

In dit nieuwe onderzoek ontdekken de auteurs een slimme truc: "Multi-Token Prediction" (MTP). Laten we uitleggen hoe dit werkt met een paar creatieve vergelijkingen.

1. Het oude probleem: De "Clever Hans" valstrik

Stel je voor dat je een leerling hebt die een test doet. De leraar (de computer) geeft de leerling de eerste helft van het antwoord en vraagt: "Wat is het volgende woord?"

• Bij de oude methode (NTP): De leerling kijkt niet echt naar de kaart. Hij ziet dat de leraar net "Straat A" heeft gezegd en weet dat daar "Straat B" achteraan komt. Hij raadt dus "Straat B" zonder echt na te denken over de route. Dit noemen de auteurs de "Clever Hans" truc (vernoemd naar een paard dat wiskundige sommen leek op te lossen, maar eigenlijk alleen naar onbewuste signalen van de trainer keek). De AI leert de patronen van de test, maar niet hoe te plannen.

2. De nieuwe methode: De "Time-Traveler"

De nieuwe methode (MTP) is als een tijdreiziger. In plaats van alleen te vragen "Wat is het volgende woord?", vraagt de trainer: "Wat is het woord over 2 stappen, en wat is het woord over 3 stappen?"

Dit klinkt als een klein detail, maar het verandert alles voor de hersenen van de AI:

• De "Rugzak" van de toekomst: Omdat de AI nu moet weten wat er over 2 of 3 stappen gebeurt, kan hij niet meer blindelings de eerste stap volgen. Hij moet de hele route in zijn hoofd hebben voordat hij de eerste stap zet.
• Het "Terugwaarts" Effect: Dit is het meest fascinerende deel. De onderzoekers ontdekten dat de AI door deze methode leert om terug te redeneren.
  • Vergelijking: Stel je voor dat je een labyrint moet doorlopen. In plaats van te proberen elke doodlopende weg van voren af aan te verkennen, kijkt de AI eerst naar de uitgang (de schat) en werkt dan terug naar de ingang.
  • Door te kijken naar het einddoel (de "end node"), weet de AI precies welke weg hij moet kiezen bij de eerste splitsing. Het is alsof je een touw van de uitgang naar de ingang trekt; de weg wordt ineens heel duidelijk.

3. Waarom werkt dit? (De "Gescheiden Signalen")

Waarom faalt de oude methode en slaagt de nieuwe?

• Oude methode (NTP): De signalen voor de eerste stap en de laatste stap zijn door elkaar gehusseld. Het is alsof je probeert te leren autorijden terwijl iemand constant op je rem trapt en tegelijkertijd gas geeft. De hersenen van de AI weten niet welke richting ze op moeten.
• Nieuwe methode (MTP): De onderzoekers tonen wiskundig aan dat deze methode de signalen ontkoppelt.
  • De eerste laag van de AI krijgt een heel duidelijk signaal: "Kijk naar het einddoel!"
  • De tweede laag krijgt een ander signaal: "Vind nu de weg terug."
  • Het is alsof je een team hebt waar de ene persoon de bestemming op de GPS invoert (en dat is heel duidelijk), en de andere persoon de route volgt. Ze werken niet meer tegen elkaar, maar in een perfecte keten.

4. De resultaten in de echte wereld

De auteurs testten dit niet alleen op simpele kaarten, maar ook op moeilijke puzzels zoals:

• Countdown: Een spelletje waarbij je met getallen en rekenkundige bewerkingen een doelgetal moet halen.
• SAT: Een logische puzzel waarbij je moet bepalen of een reeks voorwaarden wel of niet waar kunnen zijn.

In al deze gevallen presteerde de "Time-Traveler" (MTP) veel beter dan de "Blindganger" (NTP). De AI kon complexe plannen maken en oplossingen vinden die de oude methode niet kon begrijpen.

Conclusie: Van "Gokken" naar "Plannen"

Kortom, deze paper laat zien dat we AI's niet alleen hoeven te trainen om de volgende woord te raden, maar dat we ze moeten leren om na te denken over de toekomst.

Door hen te dwingen om meerdere stappen vooruit te kijken, veranderen we hun manier van denken. Ze stoppen met het raden van lokale patronen en beginnen met het bouwen van een globaal plan. Ze leren niet alleen wat er gebeurt, maar waarom het gebeurt en waar het naartoe leidt. Het is de verschuiving van een dromer die in de mist loopt, naar een strateeg die de hele kaart in één oogopslag ziet.

Technische samenvatting

Probleemstelling

Grote Taalmodellen (LLM's) tonen steeds meer redeneervermogen naarmate ze schalen, maar de onderliggende mechanismen hiervan zijn slecht begrepen. De standaard trainingsdoelstelling voor taalmodellen is Next-Token Prediction (NTP) met teacher forcing. Hoewel NTP effectief is voor kennisretrieval, heeft het moeite met het vastleggen van globale structuren in complexe redeneertaken, zoals plannen (planning).

• Beperkingen van NTP: NTP is beperkt in het modelleren van lange-termijn afhankelijkheden en neigt tot overfitting op lokale patronen. In planningsopdrachten (zoals het vinden van een pad in een graaf) kan NTP "trucs" gebruiken (zoals de "Clever Hans" cheat), waarbij het model zich baseert op reeds onthulde antwoorden in de prefix in plaats van een echte oplossing te plannen.
• De Oplossing: Multi-Token Prediction (MTP) is een recente alternatieve doelstelling waarbij het model meerdere toekomstige tokens parallel voorspelt in plaats van slechts één. Hoewel MTP empirisch succesvol is gebleken (bijv. in DeepSeek-V3), is het theoretisch mechanisme waarom dit leidt tot beter planningsvermogen nog onbekend.

Methodologie

De auteurs combineren empirische experimenten met een strikte theoretische analyse om te begrijpen hoe MTP planningsvermogen faciliteert.

1. Empirische Evaluatie:

   • Synthetische Taken: Ze testen modellen op graf-padvindingsopdrachten, specifiek een Ster-Graph (waarbij NTP faalt door de "Clever Hans" cheat) en een Binaire Boom (waarbij de cheat niet mogelijk is omdat beslissingen op elke stap nodig zijn).
   • Realistische Taken: Ze evalueren op complexere redeneertaken zoals Countdown (wiskundige puzzels) en Boolean Satisfiability (SAT) (NP-complete problemen).
   • Vergelijking: Ze vergelijken NTP (k=1) met MTP (k≥2) over verschillende schalen van data en modelparameters.

2. Theoretische Analyse:

   • Ze analyseren een vereenvoudigd tweelaags "disentangled" Transformer-model op de Ster-Graph taak.
   • Ze gebruiken een architectuur met een gedeelde backbone en meerdere onafhankelijke output-heads.
   • Ze analyseren de convergentiedynamiek van de gradiënten onder NTP versus MTP om te begrijpen waarom het ene model wel en het andere niet een planningsmechanisme leert.

Belangrijkste Bijdragen en Resultaten

1. Empirische Resultaten: MTP overtreft NTP consistent

• Ster-Graph: Modellen getraind met NTP blijven steken op 50% nauwkeurigheid (willekeurig gokken) omdat ze de "Clever Hans" cheat gebruiken. MTP-modellen bereiken daarentegen 100% nauwkeurigheid, zelfs met beperkte data.
• Binaire Boom: Zelfs wanneer de cheat wordt uitgesloten (door de complexiteit van de boom), presteert MTP aanzienlijk beter dan NTP. Dit bewijst dat het voordeel van MTP niet alleen komt door het elimineren van trucs, maar door dieper liggende mechanismen.
• Countdown en SAT: Op deze realistische taken behalen MTP-modellen consistent hogere testnauwkeurigheid dan NTP-baselines.

2. Theoretisch Mechanisme: Omgekeerd Redeneren (Reverse Reasoning)

De kernbijdrage is het aantonen dat MTP een tweefasig omgekeerd redeneerproces induceert:

• Fase 1 (Aandacht voor het einddoel): Het model leert eerst aandacht te richten op het eindnode (het doel) van de graaf.
• Fase 2 (Terugwaartse reconstructie): Vervolgens reconstrueert het model het pad door de tussenliggende knopen achterwaarts te traceren vanuit het doel.
• Waarom werkt dit? In een ster-graaf is het vinden van het pad makkelijker als je begint bij het einde (het doel) en terugwerkt naar het begin, in plaats van een grote zoekruimte vooruit te plannen.

3. De Rol van Gradiënt-Decoupling

Het paper legt uit waarom MTP dit mechanisme activeert terwijl NTP dit niet doet:

• NTP (Verstrengelde signalen): Bij NTP zijn de trainingsignalen van verschillende lagen verstrengeld. De gradiënt voor de eerste laag moet door de nog niet-geoptimaliseerde tweede laag gaan, wat leidt tot "misdirected gradients" (verkeerd gerichte gradiënten). Het model wordt actief weggestuurd van het vereiste omgekeerde redeneerpatroon.
• MTP (Gedecoupleerde signalen): MTP heeft een gradiënt-decoupling eigenschap. De "shallow" loss (het voorspellen van het tweede token, het einddoel) stuurt gradiënten uitsluitend naar de eerste laag, volledig om de tweede laag heen.
  • Dit zorgt voor een schoon trainingsignaal dat de eerste laag direct leert om naar het eindnode te kijken.
  • Zodra de eerste laag dit heeft geleerd, kan de tweede laag het pad reconstrueren via eenvoudige rand-matching.
• Conclusie: MTP verandert de optimalisatiedynamiek zodanig dat het model een interpreteerbaar, robuust planningscircuit ontdekt, terwijl NTP vastloopt in lokale minima.

Significantie

Dit paper is significant omdat het de eerste is die de convergentiedynamiek van MTP formeel analyseert en het onderscheid met NTP wiskundig bewijst.

• Het onthult dat het trainingsdoel (loss function) een fundamentele invloed heeft op het type redeneercircuit dat een Transformer leert.
• Het toont aan dat MTP de optimalisatie "bias" (vooroordeelt) naar het vinden van interpreteerbare algoritmen (zoals omgekeerd redeneren) in plaats van louter statistische correlaties.
• Het biedt een theoretische basis voor het ontwerp van toekomstige trainingsparadigma's die gericht zijn op geavanceerd redeneer- en planningsvermogen in AI-systemen.

Kortom, het paper bewijst dat Multi-Token Prediction niet alleen een efficiëntere inferentiemethode is, maar een fundamentele verschuiving in het leerproces die Transformers in staat stelt om echte planningsstrategieën te ontwikkelen door het gebruik van een schoner, gedecoupleerd trainingsignaal.