Latent Reasoning in TRMs is Secretly a Policy Improvement Operator

Technische Samenvatting: Latente Redenering in TRM's is Stiekem een Policy Improvement Operator

1. Probleemstelling

Recente kleine modellen die gebruikmaken van latente recursie, zoals Tiny Recursive Models (TRM's) en Hierarchical Reasoning Models (HRM's), hebben veelbelovende resultaten laten zien op complexe redeneertaken zoals ARC-AGI. De heersende theorie suggereert dat deze modellen prestatiewinst boeken door de effectieve netwerkdiepte te vergroten, waardoor ze grotere modellen met minder parameters kunnen emuleren. Empirische observaties onthullen echter een discrepantie: recursief toegevoegde lagen presteren vaak slechter dan one-pass modellen met een equivalente feed-forward diepte. Dit geeft aan dat in geloopte architecturen niet elke recursieve stap betekenisvol bijdraagt aan de oplossing, wat resulteert in "dead compute" (dode berekening).

De kernvragen die dit werk adresseert zijn:

• Wanneer en waarom verbetert latente redenering de prestaties?
• Wanneer leidt het tot redundante berekeningen?
• Hoe kan de training van deze modellen worden verbeterd om ervoor te zorgen dat elke stap effectief is?

Huidige trainingsmethoden voor TRM's passen supervisie primair toe op de uiteindelijke output. Dit laat vroege verfijningsstappen met een zwakke credit assignment achter, wat optimalisatie broos maakt, vooral in scenario's met kleine hoeveelheden data.

2. Methodologie

De auteurs stellen een theoretische herformulering van latente redenering voor en een bijbehorende trainingsmethodologie.

Theoretische Formalisering: Latente Redenering als Policy Improvement
Het artikel formaliseert latente recursieve redenering als een policy improvement operator.

• Referentiebeleid ($\hat{\pi}_t$): De outputdistributie van het model vóór een latente redeneringsstap (de "eerste gok").
• Verbeterd Beleid ($\pi^+_t$): De outputdistributie na één latente redeneringsupdate.
• Advantage Signaal: Het verschil tussen de log-waarschijnlijkheden van het verbeterde beleid en het referentiebeleid fungeert als een impliciet advantage signaal ($A_t(s, a) = \log \pi^+_t(a|s) - \log \hat{\pi}_t(a|s)$).
• Mechanisme: De auteurs demonstreren dat een enkele TRM-stap een policy update implementeert die het referentiebeleid herweegt op basis van een optimaliteitsconditie. Specifiek gedraagt de latente redeneringsstap zich als een beleid geconditioneerd op het evenement dat het antwoord wordt verbeterd. Dit verbindt het algoritme in de latente ruimte met het bekende framework van policy improvement in Reinforcement Learning (RL) en Classifier-Free Guidance (CFG) in diffusiemodellen.

Voorgestelde Trainingsschema: Deep Improvement Supervision (DIS)
Op basis van het inzicht in policy improvement introduceren de auteurs Deep Improvement Supervision (DIS) om stappen met dode berekening te elimineren.

• Stapsgewijze Doelstellingen: In plaats van alleen de uiteindelijke output te superviseren, biedt DIS tussenliggende targets $\{y^\dagger_s\}$ voor elke recursiestap $s$. Deze targets worden gegenereerd via een monotone discrete corruptie-planning van de grondwaarheid output $y^*$.
• Trainingsdoel: Het model wordt getraind om de vorige target $y^\dagger_{s-1}$ (via het referentiebeleid) te mappen naar de volgende, minder gecorrumpeerde target $y^\dagger_s$ (via het verbeterde beleid).
• Verliesfunctie: De training minimaliseert een cross-entropy loss tussen de post-redenering logits en de stap-specifieke target $y^\dagger_s$. Dit dwingt de log-ratio advantage op elke stap positief te zijn voor het correcte token, wat ervoor zorgt dat elke recursieve stap een gerichte verbetering uitvoert in plaats van redundante berekening.

Architecturale Vereenvoudigingen
De voorgestelde methode vereenvoudigt de TRM-architectuur door:

• De noodzaak voor een geleerd halting mechanisme (Adaptive Computation Time) te verwijderen.
• Het aantal supervisiestappen en latente redeneringsstappen te verminderen.
• Een vast aantal recursiestappen ($N_{sup}$) te gebruiken, bepaald door de corruptie-planning.

3. Belangrijkste Bijdragen

1. Theoretisch Inzicht: Het artikel stelt vast dat latente redenering in TRM's wiskundig equivalent is aan een policy improvement operator. Het overbrugt de kloof tussen algoritmen voor redenering in de latente ruimte en RL-frameworks, en verklaart modelgedrag door de lens van advantage maximalisatie.
2. Deep Improvement Supervision (DIS): Een nieuw trainingsschema dat stapsgewijze tussenliggende targets aanbiedt. Dit transformeert het moeilijke probleem van langetermijn credit assignment naar een hanteerbare supervised learning taak waarbij elke stap expliciet getraind wordt om te verbeteren ten opzichte van de vorige staat.
3. Efficiëntiewinst: De methode bereikt concurrerende prestaties terwijl de computationele overhead aanzienlijk wordt verminderd. Het vermijdt het trainen van halting modules en vermindert het totaal aantal forward passes met een factor 18× vergeleken met standaard TRM-opstellingen, terwijl de nauwkeurigheid gelijk blijft of verbetert.

4. Experimentele Resultaten

De auteurs evalueerden hun methode op de N-Queens problematiek en de ARC-AGI benchmarks (ARC-AGI-1 en ARC-AGI-2).

• N-Queens: Gebruikmakend van een compact model (0.8M parameters), bereikte de DIS-methode een nauwkeurigheid van 0.69, wat overeenkomt met de prestaties van een standaard TRM met een halting head en 16 supervisiestappen, maar met aanzienlijk minder inference stappen ($T=1, n=2$ vs. $T=3, n=6$).
• ARC-AGI:
  • Compact Model (0.8M parameters): DIS bereikte 24.0% nauwkeurigheid op ARC-AGI-1, waarmee het de originele TRM-compact (12.0%) en de meeste bestaande open-source LLM's zonder externe kennis overtreft.
  • Medium Model (7M parameters): DIS bereikte 41.3% op ARC-AGI-1 en 6.0% op ARC-AGI-2, waarmee het de standaard TRM-medium (40.4% / 3.3%) overtreft en de originele TRM (40.4% / 3.3%) evenaart, terwijl het minder latente redeneringsstappen gebruikt.
• Reductie in Berekening: De methode verminderde het totaal aantal forward passes met een factor 18 vergeleken met de baseline TRM-configuratie, wat aantoont dat expliciete supervisie van tussenliggende stappen de noodzaak voor excessieve latente redeneringslussen wegneemt.

5. Betekenis en Claims

Het artikel beweert dat men, door latent redeneren te beschouwen als een policy improvement proces, kan verklaren waarom standaard recursieve modellen soms falen (door "dead compute" waarbij stappen de advantage margin niet verbeteren) en hoe dit te herstellen.

De primaire betekenis ligt in het Deep Improvement Supervision framework, dat:

• Een constructieve trainingsregel biedt om ervoor te zorgen dat elke recursieve stap een niet-triviale verbetering uitvoert.
• Het trainen van recursieve modellen vereenvoudigt door de noodzaak voor complexe halting mechanismen en broze credit assignment te verwijderen.
• Demonstreert dat kleine, parameterefficiënte modellen state-of-the-art resultaten kunnen behalen op complexe redeneertaken (zoals ARC-AGI) wanneer het trainingsdoel is afgestemd op het theoretische mechanisme van policy improvement.

De auteurs concluderen dat hun aanpak een pad biedt naar efficiëntere en effectievere redeneermodellen door het trainingssignaal af te stemmen op de intrinsieke algoritmische aard van latente recursie.