Pretraining with Token-Level Adaptive Latent Chain-of-Thought

Dit paper introduceert een methode voor pretraining waarbij modellen een variabele lengte aan latente Chain-of-Thought-trajecten genereren per token om de prestaties te verbeteren zonder de modelgrootte te vergroten, wat leidt tot lagere perplexiteit en betere downstream-accuraatheid met minder rekencapaciteit.

De kern

Hier is een uitleg van het onderzoek van LUMIA Lab in simpel, alledaags Nederlands, met behulp van creatieve vergelijkingen.

Het Grote Probleem: De "Grote Brein"-Dilemma

Stel je voor dat we proberen een superintelligent computerbrein (een AI) te bouwen. Tot nu toe hebben we dit gedaan door twee dingen te doen:

1. Het brein groter maken (meer neuronen/parameters).
2. Het brein meer boeken te laten lezen (meer trainingsdata).

Maar er is een probleem: we hebben bijna alle goede boeken op de aarde al opgebruikt, en het maken van nog grotere breinen kost onbetaalbare hoeveelheden energie en geld. Het is alsof je probeert een auto sneller te maken door hem steeds zwaarder te maken; op een gegeven moment is de motor te zwaar om nog vooruit te komen.

De Oplossing: "Denken voor je spreekt"

De onderzoekers van LUMIA Lab hebben een slimme nieuwe manier bedacht. In plaats van het brein groter te maken, laten ze het brein dieper nadenken voordat het een woord uitspreekt.

Ze noemen dit "Token-Level Adaptive Latent Chain-of-Thought". Dat klinkt ingewikkeld, maar het is eigenlijk heel simpel:

• Normaal: Een AI denkt één seconde na en zegt dan een woord.
• Nieuwe methode: De AI denkt intern na (in een onzichtbare "gedachtenruimte") voordat hij een woord zegt.

De Creatieve Vergelijking: De Slimme Chef-kok

Stel je een zeer drukke restaurantkeuken voor (de AI) met één chef-kok (het model).

1. De oude manier (Standaard AI): De chef kookt elke maaltijd precies even lang, of het nu een simpele salade is of een ingewikkelde soufflé. Hij gebruikt evenveel tijd en energie voor alles. Dat is inefficiënt.
2. De nieuwe manier (Adaptieve Latente CoT): De chef heeft nu een magisch vermogen om te beslissen hoeveel tijd hij aan een gerecht besteedt voordat hij het serveert.
   • Een simpele salade? Hij denkt 0,1 seconde na en serveert het direct. (Korte "latente keten").
   • Een complexe soufflé? Hij denkt 5 minuten lang stil in zijn hoofd na over de temperatuur en ingrediënten, voordat hij het serveert. (Lange "latente keten").

Het mooie is: de chef hoeft hiervoor geen extra mensen in dienst te nemen (geen groter model) en hij hoeft niet meer boeken te lezen. Hij wordt gewoon slimmer in het verdelen van zijn tijd.

Hoe werkt het technisch? (De Magische Regels)

De onderzoekers hebben drie slimme regels bedacht om dit te laten werken:

1. De Parallelle Keuken (Parallel Masking):
   In de oude methoden moest de chef wachten tot de soufflé klaar was voordat hij aan de salade kon beginnen. Dat was traag.
   In deze nieuwe methode kan de chef alle gerechten tegelijk in zijn hoofd voorbereiden. Hij denkt na over de salade, de soep en de soufflé allemaal op hetzelfde moment, maar in verschillende "tijdslagen" in zijn hoofd. Dit maakt het veel sneller.

2. De Stopknop (Probabilistic Halting):
   De chef heeft een knop die zegt: "Is dit gerecht klaar?"

   • Als hij denkt: "Ja, dit is perfect," stopt hij direct.
   • Als hij denkt: "Nog niet zeker," blijft hij doorgaan met nadenken.
     Dit bespaart enorm veel energie, want hij hoeft niet te nadenken over dingen die al makkelijk zijn.

3. De "Geen Onnodig Werk"-Regel (Correctness-Aware Loss):
   Dit is de slimste regel. De chef leert dat als hij al 99% zeker is dat de salade goed is, het niet slim is om er nog langer over na te denken. Soms maakt extra nadenken het zelfs slechter (je verpest de salade door er te veel aan te zitten). De AI leert dus: "Stop als je het al weet."

Wat is het resultaat?

De onderzoekers hebben dit getest met hun eigen AI-modellen (op basis van LLaMA). De resultaten zijn indrukwekkend:

• Beter presteren: De AI maakt minder fouten en begrijpt taal beter.
• Minder werk: Ze hebben minder rekenkracht nodig dan andere geavanceerde methoden om hetzelfde resultaat te bereiken.
• Slimmer verdelen: De AI leert zelf om meer tijd te steken in moeilijke vragen (zoals wiskunde of logica) en minder tijd in simpele vragen (zoals "wat is de kleur van de lucht?").

Samenvatting

In plaats van een AI te bouwen die groter is en meer energie verbruikt, hebben deze onderzoekers een AI gebouwd die slimmer omgaat met zijn tijd. Het is alsof je een student niet dwingt om langer te studeren, maar hem leert om precies te weten wanneer hij iets al begrijpt en wanneer hij nog even moet nadenken.

Dit is een grote stap naar efficiëntere, goedkopere en slimmere kunstmatige intelligentie in de toekomst.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Pretraining with Token-Level Adaptive Latent Chain-of-Thought" in het Nederlands.

Titel: Pretraining met Token-Level Adaptieve Latente Chain-of-Thought (Adaptive Latent CoT)

Auteurs: Boyi Zeng, Yiqin Hao, et al. (LUMIA Lab, Shanghai Jiao Tong University)
Datum: Februari 2026

---

1. Het Probleem

De huidige schaalvergroting van Large Language Models (LLMs) door het verhogen van het aantal parameters en trainingsdata stuit op kritieke beperkingen:

• Data-exhaustie: Er is een tekort aan hoogwaardige publieke trainingsdata.
• Communicatiekosten: Het vergroten van modelgrootte leidt tot enorme communicatie-overhead in gedistribueerde training.
• Uniforme Berekening: Bestaande methoden om per-token berekening te verhogen (zoals recursieve parameter-sharing of expliciete Chain-of-Thought) distribueren vaak rekenkracht uniform over alle tokens, ongeacht de moeilijkheidsgraad. Dit is suboptimaal vergeleken met menselijke cognitie, waarbij denken adaptief is (kort voor eenvoudige concepten, lang voor complexe redenering).
• Efficiëntieproblemen: Bestaande methoden voor latente redenering vereisen vaak multi-stap training, extra supervisie, of leiden tot een enorme toename van de trainings-FLOPs (bijvoorbeeld door Jacobi-iteraties).

2. Methodologie

De auteurs stellen Adaptive Latent CoT voor: een framework dat het model in staat stelt om een variabele lengte aan "latente" denkstappen (CoT) te genereren in een continue ruimte voordat een token wordt uitgegeven, zonder de modelparameters te vergroten.

De kerncomponenten zijn:

A. Parallel Masking (Oplossing voor Sequential Dependency)

In traditionele latente CoT-modellen is er een strikte sequentiële afhankelijkheid: de staat van token $t$ hangt af van de volledige keten van latente stappen van token $t-1$. Dit maakt parallelle berekening onmogelijk.

• Oplossing: De auteurs definiëren een Parallel Attention Mask over een 2D-index $(t, k)$, waarbij $t$ de tokenpositie is en $k$ de latente stap.
• Mechanisme: Dit masker zorgt voor causaliteit in beide dimensies, maar stelt het model in staat om voor een vaste latente stap $k$ alle tokens $t$ in het sequence parallel te verwerken. Dit reduceert de sequentiële afhankelijkheid van $O(L \times K)$ naar $O(K)$, wat enorme GPU-parallelisatie mogelijk maakt.

B. Adaptieve Berekening via Probabilistische Halting

Niet alle tokens vereisen evenveel denkstappen.

• Router: Een lichtgewicht module voorspelt bij elke stap $k$ de kans om door te gaan naar de volgende stap ($g^{(k)}_t$).
• Reach Probability: De kans dat een token een bepaalde stap bereikt, wordt recursief berekend op basis van de opeenvolgende "doorgang"-kansen.
• Threshold Pruning: Als de kans om een volgende stap te bereiken onder een drempelwaarde $\tau$ daalt, wordt het token uit de batch verwijderd voor verdere stappen. Dit bespaart FLOPs direct tijdens training en inferentie.
• Expectation-based Mixing: De finale representatie is een gewogen som van de uitgevoerde latente stappen, waarbij de resterende waarschijnlijkheidsmassa (van de afgekapte stappen) wordt herschapen naar de laatste uitgevoerde stap om de totale massa te behouden.

C. Correctness-Aware Adaptive Loss

Om te voorkomen dat het model onnodig blijft rekenen, introduceren ze een extra verliesfunctie.

• Principe: Als het model al een hoge waarschijnlijkheid heeft voor het ground-truth token ($p_{target}$), moet het stoppen met rekenen.
• Straf: De loss straft het "doorgaan" (continue) af, evenredig met de huidige $p_{target}$. Dit stimuleert het model om vroeg te stoppen bij makkelijke tokens en door te gaan bij moeilijke, onzekere tokens.

3. Belangrijkste Bijdragen

1. Eén-staps Pretraining: Het framework leert adaptieve berekening end-to-end tijdens één pretraining-fase op algemene tekst, zonder extra supervisie (zoals annotaties voor CoT) of multi-stap training.
2. Efficiëntie: Het reduceert zowel trainings- als inferentie-FLOPs door token-gebaseerde adaptieve halting en pruning, terwijl het toch meer rekenkracht per token toewijst waar nodig.
3. Parallellisatie: Door het 2D attention-mechanisme wordt het sequentiële bottleneck van latente redenering opgelost, wat training op lange contexten mogelijk maakt.
4. Natuurlijke Emergentie: Het gedrag van "eenvoudige tokens = weinig stappen, complexe tokens = veel stappen" ontstaat natuurlijk uit de data, zonder handmatige regels.

4. Resultaten

Experimenten zijn uitgevoerd op LLaMA-architecturen (410M en 1.4B parameters) getraind op The Pile dataset.

• Taalmodellering (Perplexity): Het model behaalt consistent de laagste perplexiteit op alle datasets (The Pile, WikiText, LAMBADA) vergeleken met sterke baselines zoals PonderLM2, LoopedLM en PausedLM.
• Efficiëntie: Het 1.4B model met Adaptive Latent CoT presteert beter dan de sterkste baseline (PonderLM-2) met minder dan de helft van de trainings-FLOPs (7.47 vs 17.47 $\times 10^{20}$ FLOPs).
• Downstream Taken: Het model behaalt de beste gemiddelde nauwkeurigheid op diverse benchmarks (ARC, HellaSwag, PIQA, etc.) in zowel 0-shot als 5-shot settings.
• Vergelijking: Een 410M model met deze techniek presteert beter dan een standaard 1.4B model (vanaf scratch getraind) binnen hetzelfde rekenbudget, wat aantoont dat adaptieve berekening effectiever is dan puur het vergroten van het aantal parameters.
• Analyse:
  • Tokens met een lage $p_{target}$ (moeilijk) krijgen automatisch meer latente stappen.
  • Tokens met een hoge $p_{target}$ (makkelijk) stoppen vroeg (soms 0 extra stappen).
  • Case studies tonen aan dat het model extra stappen toewijst aan entiteiten, cijfers en inhoudswaarden, en minder aan voegwoorden.

5. Betekenis en Impact

Dit werk biedt een fundamenteel nieuwe richting voor het schalen van LLMs. In plaats van te vertrouwen op het vinden van meer data of het vergroten van het model, toont het aan dat het interneren van adaptieve redenering in de pretraining-fase een krachtige manier is om modelcapaciteit te verhogen.

• Duurzaamheid: Het verlaagt de energie- en rekenkosten voor training en inferentie door alleen te rekenen waar het nodig is.
• Schalbaarheid: Het biedt een oplossing voor de "data wall" door meer intelligentie uit dezelfde hoeveelheid data te halen.
• Toekomst: Het bewijst dat adaptieve computationele strategieën, die eerder beperkt waren tot inferentie of gesuperviseerde settingen, succesvol en efficiënt kunnen worden geleerd tijdens de basispretraining.

Kortom, Adaptive Latent CoT stelt modellen in staat om "te denken" in een continue, verborgen ruimte, waarbij de hoeveelheid denken dynamisch wordt afgestemd op de complexiteit van het probleem, wat leidt tot superieure prestaties met minder rekenkracht.