Thinking into the Future: Latent Lookahead Training for Transformers

Titel: Thinking into the Future: Latent Lookahead Training for Transformers

Auteurs: Lorenzo Noci, Gregor Bachmann, Seyed-Mohsen Moosavi-Dezfooli, Moin Nabi (Apple)

1. Het Probleem

Moderne Large Language Models (LLM's) worden doorgaans getraind met een autoregressieve benadering op basis van next-token prediction (NTP) en teacher-forcing. Hoewel deze methode schaalbaar is, introduceert ze twee fundamentele beperkingen:

1. Korte horizon en gebrek aan reflectie: Het model leert slechts één stap vooruit te kijken. Het moet bij elke stap direct een discrete token kiezen zonder de downstream-consequenties van die keuze verderop in de reeks te evalueren. Dit leidt tot "myopisch" gedrag, waarbij modellen in complexere taken (zoals plannen of redeneren) vastlopen in doodlopende wegen of fouten maken in chain-of-thought redeneringen.
2. Uniforme compute-toewijzing: Elke token wordt gegenereerd via één enkele forward-pass. Dit betekent dat het model niet meer rekenkracht kan investeren in moeilijke beslissingen die meer nadenken vereisen, terwijl eenvoudige tokens evenveel compute krijgen.

Bestaande oplossingen zoals Chain-of-Thought (CoT) of Reinforcement Learning breiden de context uit met zichtbare tokens, maar dit vereist nog steeds het vastleggen van discrete tokens. Andere methoden gebruiken recursieve berekening (zoals looped transformers), maar vaak zonder de context expliciet uit te breiden met "gedachte"-tokens.

2. Methodologie: Latent Lookahead

De auteurs introduceren Latent Lookahead, een trainingsstrategie die het model toestaat om "na te denken" voordat het een zichtbaar token genereert. In plaats van direct een token te samplen, voert het model een meervoudige stap-voorwaartse simulatie uit in de latente ruimte.

Kernconcepten:

• Zichtbare vs. Latente Tokens:
  • Zichtbare tokens ($x_i$): De standaard tokens uit de trainingdata.
  • Latente tokens ($z_{i,j}$): Auxiliaire voorspellingen in de verborgen ruimte die het model genereert om $\tau$ stappen vooruit te kijken zonder een keuze te maken.
• Het Proces:
  1. Op geselecteerde posities in de sequentie (bijv. na een vraag) stapt het model over in de "latent lookahead"-modus.
  2. Het model unrollt zijn verborgen staat $\tau$ stappen vooruit door de voorspelde latente states recursief terug te voeren in de context.
  3. Er worden $\tau$ latente tokens gegenereerd: $\{z_{i,1}, ..., z_{i,\tau}\}$.
  4. Pas na deze $\tau$ stappen wordt het volgende zichtbare token gesampled, gebaseerd op de geüpdatete verborgen staat.
• Training en Supervisie:
  • De latente tokens worden supervised tegen de ground-truth tokens die $\tau$ stappen verderop in de sequentie liggen. Dit dwingt het model om alternatieve toekomstscenario's te evalueren en de juiste paden te selecteren.
  • De totale loss functie is de som van de standaard NTP-loss en de loss op de latente tokens: $L = L_{NTP} + L_{latent}$.
• Aandachtsmasker (Attention Mask):
  • Om parallelle training mogelijk te maken en consistentie te garanderen, gebruiken de auteurs een speciaal masker.
  • Binnen één "latente gedachte" (latent thought) is de aandacht niet-causaal (bidirectioneel). Dit stelt het model in staat om vroege latente tokens te verfijnen op basis van latere voorspellingen binnen dezelfde gedachte.
  • Verschillende latente gedachten (op verschillende posities in de tekst) kunnen niet naar elkaar kijken, wat de complexiteit van $\mathcal{O}(n\tau)$ terugbrengt naar $\mathcal{O}(\tau)$ forward passes per gradient stap.

3. Belangrijkste Bijdragen

1. Het Latent Lookahead Framework: Een nieuw trainingsparadigma dat context uitbreidt met latente "gedachte"-tokens en deze recursief berekent zonder sampling.
2. Efficiënte Parallelle Training: Ontwerp van een attention mask dat parallelle generatie van meerdere latente gedachten mogelijk maakt, wat de trainingskosten drastisch verlaagt ten opzichte van sequentiële benaderingen.
3. Prestatieverbetering in Plannings Taken: Het aantonen dat deze methode aanzienlijk beter presteert dan autoregressieve baselines en methoden met pauze-tokens (pause tokens) op taken die foresight vereisen.
4. Interpretatie van Latente Tokens: Visualisatie van de latente tokens toont aan dat ze een "superpositie van staten" vertegenwoordigen. Het model verkent meerdere opties en verhardt de distributie naar het juiste antwoord tijdens het denkproces, in plaats van direct te commiten.

4. Resultaten

De auteurs testten het model op drie planningsgerichte datasets: Sudoku (4x4 en 9x9), ProsQA (planning in DAG's) en Maze (doolhoven).

• Sudoku:
  • Op de 9x9 Sudoku steeg de nauwkeurigheid van 12,5% (standaard NTP) naar 35,5% met Latent Lookahead.
  • De standaard "pause token" baseline (die extra context toevoegt zonder latente supervisie) presteerde slecht (12,5%) en gaf geen verbetering, wat aantoont dat de extra compute alleen niet genoeg is; de richting van de compute (via lookahead-supervisie) is cruciaal.
• ProsQA & Maze:
  • Significant verbeteringen ten opzichte van baselines (bijv. ProsQA: 91,8% vs 80,5% voor NTP).
• Schaalbaarheid:
  • Het verhogen van het aantal latente tokens ($\tau$) leidt tot een monotoon stijgende prestatie, terwijl andere methoden verzadigen.
  • Een model met Latent Lookahead (2 lagen) presteerde beter (93,5%) dan een standaard transformer met 32 lagen (80,0%) op kleine Sudoku's, wat aantoont dat de verbetering niet puur door diepte komt, maar door de kwaliteit van de berekening.
• Ablatie-studies:
  • Niet-causale aandacht binnen latente tokens is essentieel voor het verfijnen van voorspellingen.
  • Het gebruik van het model als multi-token predictor (meerdere tokens direct decoderen) werkt minder goed dan het uitstellen van de predictie na het denkproces.

5. Betekenis en Conclusie

Het paper toont aan dat het vermogen om vooruit te kijken en alternatieve toekomstscenario's intern te simuleren essentieel is voor complexe redeneertaken.

• Verschil met bestaande werken: In tegenstelling tot Chain-of-Thought (die zichtbare tokens genereert) of Looped Transformers (die vaak recursief zijn zonder contextuitbreiding), combineert Latent Lookahead contextuitbreiding met recursieve, niet-causale berekening in de latente ruimte.
• Toekomstperspectief: De auteurs suggereren dat toekomstig werk zich moet richten op het leren van waar en hoe lang het model moet nadenken (bijv. via onzekerheidsmetingen) en het schaalbaar maken van deze trainingsmethode voor pre-training van grote modellen.

Kortom, Latent Lookahead biedt een mechanisme om compute dynamisch toe te wijzen aan moeilijke beslissingen door "nadenken" in de verborgen ruimte te laten plaatsvinden voordat een definitief antwoord wordt gegeven.