JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation

Technische Samenvatting: JEPA-Reasoner: Ontkoppeling van Latente Redenering van Tokengeneratie

Probleemstelling
Huidige Large Language Models (LLMs) opereren onder een verenigd autoregressief paradigma waarbij hoogwaardige redenering en laagwaardige tokengeneratie onlosmakelijk aan elkaar gekoppeld zijn. In dit kader voorspelt het model het volgende token op basis van een geschiedenis van eerder gesamplede discrete tokens. Deze koppeling creëert een fundamentele bottleneck: elke lokale fout in de tokenselectie (bijv. een grammaticale misstap of een suboptimale woordkeuze) vervuilt het contextvenster, wat direct de trajectorie van alle daaropvolgende redeneringsstappen corrumpeert. Dit leidt tot cumulatieve fouten in taken met een lange horizon. Bovendien blijven bestaande pogingen tot redeneren in de latente ruimte (bijv. COCONUT, Recurrent Depth-modellen) beperkt door deze koppeling, omdat ze latente vectoren met tokens vermengen of gebruikmaken van causale masking die voorkomt dat vroege tokens worden geïnformeerd door toekomstige redeneringstoestanden. Deze modellen hebben ook moeite met het representeren van onzekerheid, omdat ze vaak gedwongen worden tot een ineenstorting naar een enkel redeneerpad vanwege de discrete aard van de token-niveau koppeling.

Methodologie
De auteurs stellen JEPA-Reasoner voor, een nieuw architecturaal paradigma dat latente-ruimte redenering volledig ontkoppelt van tokengeneratie. De architectuur bestaat uit twee afzonderlijke modules:

1. JEPA-Reasoner (De Redeneerder): Deze module opereert volledig binnen een continue, genormaliseerde latente ruimte. Het maakt gebruik van een Joint-Embedding Predictive Architecture (JEPA) die is aangepast voor autoregressieve latente-ruimte generatie.

   • Mechanisme: Het neemt een input, codeert deze, en gebruikt aangepaste Transformer-blokken om de volgende latente matrix te voorspellen die de daaropvolgende redeneringsstappen vertegenwoordigt.
   • Restrictie: Cruciaal is dat de Redeneerder de outputs niet projecteert naar vocabulaire-waarschijnlijkheden. In plaats daarvan wordt de gegenereerde latente matrix genormaliseerd (via een hybride RMS- en L2-normalisatielaag) tot een eenheids-hypersfeer en teruggekoppeld als input voor de volgende stap.
   • Training: Het wordt getraind via Self-Supervised Training (SST) om latente representaties te voorspellen die worden geleverd door een target encoder (wiens gewichten worden bijgewerkt via Exponential Moving Average). De doelstelling is een geschaalde cosinus-gelijkenis-loss, wat het model aanmoedigt om hoekverschillen en consistente feature-dynamiek te leren in plaats van discrete tokenvoorspelling.

2. Talker (De Generator): Dit is een aparte, lichtgewicht Transformer-gebaseerde module die onafhankelijk is getraind om menselijk leesbare tekst te reconstrueren.

   • Mechanisme: Het ontvangt de volledige latente redeneerketen gegenereerd door de JEPA-Reasoner en reconstrueert de tokensequentie.
   • Varianten: Het artikel introduceert Mono-Talker (alleen decoders, voor context-vrije taken) en Dual-Talker (encoders en decoders, voor context-bewuste taken).
   • Afhankelijkheid: De Talker is uitsluitend getraind voor reconstructie op basis van de output van de Reasoner; het kan niet onafhankelijk van de Reasoner voorspellingen doen.

Belangrijkste Bijdragen
Het artikel identificeert drie emergente eigenschappen van deze ontkoppelde architectuur:

1. Wiskundige Foutbeheersing: De auteurs bieden een probabilistische factorisatie aan waaruit blijkt dat, omdat de redeneerketen ($R$) onafhankelijk wordt gegenereerd van de token-sampling ($X$), token-niveau fouten geen wiskundige weg hebben om de logica van de Reasoner te beïnvloeden ($\frac{\partial r_{t+1}}{\partial \hat{x}_t} = 0$). Dit isoleert de redeneermotor van discrete sampling-fouten.
2. Continue Begeleiding: De generator (Talker) heeft toegang tot de volledige verliesloze redeneertrajectorie voordat er enige tokens worden gegenereerd, in tegenstelling tot gekoppelde modellen waarbij generatie wordt afgewisseld met redenering.
3. Representatie van Onzekerheid: De Reasoner kan "gemengde latente vectoren" produceren die meerdere mogelijke redeneerpaden simultaan coderen. Deze vectoren fungeren als lineaire combinaties van vocabulaire-embeddings, waardoor het model meerdere hypothesen kan behouden zonder de overhead van beam search of meerdere decodes.

Experimentele Resultaten
De auteurs hebben gecontroleerde experimenten uitgevoerd op synthetische en natuurlijke taaltaken:

• Synthetische Taken (Tree Search & CFG):
  • In een binaire boomzoektaak bereikte het model een nauwkeurigheid van 99,87%. PCA-analyse bevestigde dat de voorspelde latente vectoren een continue wolk vormen binnen de ruimte gespannen door discrete vocabulaire-embeddings, wat de hypothese van de "gemengde latente staat" valideert.
  • In Context-Free Grammar (CFG) generatietaken toonde de JEPA-Reasoner een superieure robuustheid tegen zowel input-tokenruis (0–30% vervanging) als latente-ruimte Gaussische ruis vergeleken met een standaard Transformer-baseline en het gekoppelde COCONUT-model.
• Natuurlijke Taal Taken (GSM8K):
  • Een 0.9B parameter JEPA-Reasoner (gecombineerd met een 198M Talker) werd getraind op identieke data en hyperparameters als een 0.9B gekoppelde Transformer-baseline.
  • Prestaties: De JEPA-Reasoner behaalde een nauwkeurigheid van 51,9% op de 8-shot GSM8K benchmark, een verbetering van 149,5% ten opzichte van de baseline (20,8%).
  • Opschaling: Terwijl de baseline een "redeneerplateau" vertoonde (0,1% verbetering van 5-shot naar 8-shot), maakte de JEPA-Reasoner gebruik van extra voorbeelden om een prestatieverbetering van 131,7% te realiseren, wat wijst op effectieve in-context opschaling.
  • Vergelijking: Het 0.9B model presteerde beter dan een 1B parameter CoT-model (Llama 3.2) en een 3.5B parameter Recurrent Depth-model (Huginn-0125) op dezelfde benchmark.

Betekenis en Claims
Het artikel positioneert JEPA-Reasoner niet louter als een prestatieverbetering door schaling, maar als een architecturale onderzoek naar de voordelen van ontkoppeling. De auteurs beweren dat de overgang van een gekoppeld naar een ontkoppeld paradigma significante winsten oplevert in robuustheid en redeneerdiepte zonder dat geavanceerde reinforcement learning of industriële trainingsrecepten vereist zijn.

Het werk suggereert dat het "redeneerplateau" dat wordt waargenomen bij kleinschalige token-gebaseerde modellen, kan worden overwonnen door het redeneerproces te verplaatsen naar een continue, abstracte latente ruimte. Door de redeneermotor te isoleren van het discrete token-samplingproces, biedt de architectuur een robuustere fundering voor complexe redenering, wat foutbeheersing en de representatie van onzekerheid mogelijk maakt die moeilijk te bereiken zijn in standaard autoregressieve modellen. De auteurs concluderen dat deze ontkoppelde aanpak een nieuwe weg voorwaarts biedt voor het ontwerp van foundation models.