Technische Samenvatting: Read the Trace, Steer the Path (CAPR)

Probleemstelling

Diffusion Large Language Models (dLLMs) genereren tekst door iteratief een volledig gemaskeerde sequentie te denoisen over parallelle stappen. In tegen tegenstelling tot autoregressieve modellen, legt de dLLM voor elke gemaskeerde positie bij elke stap een voorspellende distributie bloot, wat een rijke "denoising trace" (denoising-spoor) creëert die onthult welke tokens zelfverzekerd en stabiel worden versus welke onstabiel blijven of oscilleren.

Huidige Reinforcement Learning (RL) post-training methoden voor dLLMs maken geen effectief gebruik van dit spoor en vallen uiteen in twee extremen:

1. Flat Rollouts: Kennen een enkel resultaat-beloning toe aan de gehele traject. Hoewel computationeel goedkoop, bieden ze geen signaal over welke specifieke denoising-stappen of tokenblokken cruciaal waren voor het uiteindelijke resultaat.
2. Tree Rollouts: Vertakken alternatieve voortzettingen om beloningen te vergelijken en segment-niveau advantages te genereren. Hoewel ze fijnere granulariteit bieden, groeit de computationele kosten snel met de vertakkingsfactor en diepte, wat ze vaak onbetaalbaar duur maakt.

De centrale vraag die dit werk adresseert is: Kan de denoising trace zelf tree-achtige supervisiesignalen bieden zonder de computationele kosten van een volledige tree-expansie?

Methodologie: CAPR

De auteurs introduceren CAPR (Cached-Amortized Path Refinement), een blok-gebaseerd RL-algoritme ontworpen om per-blok advantages te extraheren uit de denoising trace. CAPR werkt via drie primaire mechanismen:

1. Cache & Steer

CAPR registreert de vroege denoising-stappen om een compact path state ($h_j$) voor elke positie $j$ op te bouwen. Deze staat vat samen:

• Zelfvertrouwen in de huidige top token.
• Entropie van de waarschijnlijkheidsdistributie.
• Stabiliteit van de top token over de stappen heen.

Deze staat wordt bijgewerkt via een GRU-achtige mechanisme waarbij hoge entropie het geheugen reset (het vergeten van verouderde bewijslast) en stabiliteit het behoudt. De path state dient twee functies:

• Steering: Het stuurt de logits voor de volgende reverse stap bij door een positieve push te geven aan stabiele, zelfverzekerde posities, terwijl oscillerende posities worden gedempt.
• Unmasking Control: Het beïnvloedt welke gemaskeerde posities worden onthuld (unmasked) op basis van hun zekerheid, zodat het model zich focust op stabiele commitments.

2. Branch & Prune

Om trainingssignalen te genereren zonder volledige tree-expansie, voert CAPR één vertakkingsgebeurtenis uit op het middelpunt van de denoising-traject ($f_b = 0.5$).

• Caching: Het model cachet de prefix (logits, voorspelde tokens en path state) tot aan het vertakkingspunt.
• Forking: Het genereert een klein aantal sibling continuaties ($G_b$, bijv. 2) door Gumbel-perturbaties toe te passen op de vocabulaire-waarschijnlijkheden. Deze siblings hergebruiken de gecachte prefix, waardoor redundante forward passes worden vermeden.
• Pruning: Op blokgrenzen worden branches gepruned op basis van hun path state kwaliteit (gemiddelde stabiliteit). Alleen de meest veelbelovende branches gaan door, wat voorkomt dat de zoekruimte uitmondt in een volledige tree.

3. Block Critic

Een lichtgewicht value head (Block Critic) leest de path state bij elk blokgrenspunt.

• Reward Redistribution: Nadat de rollout is voltooid, wordt de uiteindelijke resultaat-beloning verdeeld over de blokken op basis van het aantal onthulde tokens in elk blok en de path state features.
• Advantage Calculation: Met behulp van TD/GAE (Generalized Advantage Estimation) zet de critic de schaarse uiteindelijke beloning om in een per-blok advantage ($\hat{A}_b$).
• PPO Update: Deze blok-advantages herwegen de standaard rollout-niveau advantage gebruikt in PPO. Dit stelt het model in staat om te leren welke specifieke blokken het meest hebben bijgedragen aan het succes, waardoor de granulariteit van tree search effectief wordt hersteld zonder de kosten van een tree.

Trainingsdoelstelling: De totale loss combineert de clipped PPO loss (met blok-gewogen advantages), KL-divergentie (om dicht bij het referentiemodel te blijven), een $\chi^2$ regularisator (om zeldzame log-ratio pieken te onderdrukken), self-distillation loss (met behulp van opgeloste rollouts als supervised targets) en een value head loss.

Belangrijkste Bijdragen

1. Denoising Trace-Based RL: Het artikel stelt een nieuw perspectief voor waarbij zelfvertrouwen en stabiliteit langs het denoising-pad worden omgezet in per-blok advantages, wat de noodzaak voor volledige tree-expansie elimineert.
2. CAPR Algoritme: Een praktische implementatie die vroege outputs cachet, sibling rollouts amortiseert, denoising-stappen stuurt via path states, en een block critic gebruikt om granulaire PPO-gewichten te generen.
3. State-of-the-Art Prestaties: CAPR bereikt nieuwe state-of-the-art resultaten voor RL-getunede dLLMs op vier redeneer-benchmarks (4×4 Sudoku, Countdown, GSM8K, Math500) op zowel dense als Mixture-of-Experts (MoE) LLaDA backbones bij 256- en 512-token budgetten.

Experimentele Resultaten

De auteurs evalueerden CAPR tegen baselines inclusief Diffu-GRPO, VRPO, wd1, SAPO, d2-stepMerge, TraceRL, en de boomstructuur-gebaseerde d-TreeRPO.

• Nauwkeurigheid: Op het 256-token budget presteerde CAPR beter dan of gelijk aan alle baselines op alle taken voor zowel LLaDA-8B-Instruct als LLaDA-MoE-7B-A1B-Instruct.
  • Sudoku: CAPR behaalde 93.0% (LLaDA-8B) en 84.8% (LLaDA-MoE), wat overeenkomt met of de sterkste tree-gestructureerde baseline overtreft (d-TreeRPO: 92.9% en 78.3%).
  • Countdown: Significante winsten werden waargenomen, waarbij CAPR 72.7% (8B) en 68.8% (MoE) bereikte, waarmee het d-TreeRPO verslaat.
• Efficiëntie: CAPR vermindert de computationele kosten aanzienlijk vergeleken met tree rollouts.
  • Generatiekosten: CAPR vereist ongeveer 0.75× de kosten van flat rollouts en 0.6× de kosten van tree rollouts (bij standaard instellingen).
  • Wall-Clock Tijd: Op Sudoku nam CAPR 165 seconden per optimalisatiestap in vergelijking met 596 seconden voor d-TreeRPO, waarbij vergelijkbare nauwkeurigheid werd bereikt met minder dan één derde van de compute per stap.
• Ablatie Studies:
  • De Block Critic leverde de grootste tussenliggende winst op, wat de leercurve substantieel veranderde.
  • Branch & Prune voegde nuttige diversiteit toe maar was minder effectief als standalone methode.
  • Self-Distillation was cruciaal voor het stabiliseren van de training op groepen met nul variantie (opgeloste rollouts) en het verminderen van gradiëntpieken.

Betekenis en Claims

Het artikel claimt dat CAPR een praktisch "middengemiddelde" biedt voor dLLM-RL. Het slaagt erin de granulariteit van tree search (lokale vergelijking van alternatieven) te herstellen terwijl de computationele explosie van volledige tree-expansie wordt vermeden.

• Efficiëntie vs. Granulariteit: CAPR laat zien dat de denoising trace zelf voldoende informatie bevat om beloningen op blokniveau toe te wijzen, wat de noodzaak voor dure tree-gebaseerde rollouts wegneemt.
• Generalisatie: De methode wordt effectief getransfereerd van dense naar sparse (MoE) backbones zonder taakspecifieke hertuning.
• Practicaliteit: Door de per-stap compute te reduceren tot ongeveer 0.6× van tree rollouts terwijl de nauwkeurigheid gelijk blijft of zelfs hoger is, maakt CAPR training met hoge granulariteit haalbaar voor grotere modellen en langere contexten.

De auteurs merken beperkingen op en erkennen dat de methode momenteel wordt geëvalueerd op automatisch verifieerbare redeneertaken en mogelijk aanpassing vereist voor open eind instructie-opvolging of domeinen met ruisende, niet-lokale beloningen. Daarnaast wordt de block critic online geleerd, wat betrouwbaarheidsproblemen kan opleveren voor zeer lange redeneerketens.