Technische Samenvatting: Trajectbalans-Na-training voor Diffusie-taalmodellen

Probleem: Trajectvergrendeling in Diffusie-RL

Diffusie-taalmodellen (dLLMs) bieden een veelbelovend alternatief voor autoregressieve modellen voor redenering en codegeneratie. Echter, bestaande na-trainingsmethoden voor dLLMs passen grotendeels beloningsmaximerende Reinforcement Learning (RL)-doelstellingen (zoals PPO, GRPO) toe die zijn ontworpen voor autoregressieve modellen. De auteurs identificeren een centraal faalmechanisme in deze aanpassing, genaamd trajectvergrendeling.

Bij diffusiemodellen wordt de beloning alleen toegewezen aan de uiteindelijke gedenoiseerde voltooiing ($y$), terwijl de modelverdeling wordt opgelegd over het volledige gedenoiseerde traject ($\tau$). Omdat terminale beloningen "padonverschillend" zijn—wat betekent dat elk traject dat dezelfde correcte $y$ bereikt, dezelfde beloning ontvangt—neigen bemonsterde policy-gradient-updates ertoe om de specifieke trajecten te versterken die toevallig vroeg in de training worden bemonsterd. Dit creëert een zichzelf versterkende feedbacklus waarbij waarschijnlijkheidsmassa instort op een smal scala aan gedenoiseerde paden. Bijgevolg verliest een model, zelfs als het leert een correct antwoord te produceren, het vermogen om een divers scala aan geldige oplossingen (modi) te genereren bij herhaalde bemonstering. De auteurs formaliseren dit via Stelling 1, waarin wordt bewezen dat de dekking van terminale modi een bovengrens heeft die wordt bepaald door trajectdiversiteit; trajectvergrendeling beperkt derhalve direct de oplossingsdekking van het model.

Methodologie: TraFL (Traject Flow Balancing)

Om trajectvergrendeling aan te pakken, stellen de auteurs TraFL voor, een na-trainingsdoelstelling die is gebaseerd op het traject-balansprincipe uit Generative Flow Networks (GFlowNets). In plaats van de beloning direct te maximaliseren, traint TraFL het beleid om te matchen met een beloningsgeoriënteerde doelverdeling die is verankerd aan een bevroren referentiemodel.

De Doelstelling

De doelverdeling $p^*(\tau | x)$ wordt gedefinieerd als:
$$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$$
waarbij:

• $p_{\text{ref}}$ een bevroren referentiemodel is.
• $r(x, y)$ de terminale beloning is.
• $\beta$ de scherpte van de beloning regelt.
• $Z(x)$ een prompt-afhankelijke normalisatieterm is (partitiefunctie).

De trainingsdoelstelling minimaliseert het kwadratische residu tussen de log-kans van het huidige model en de doelstelling:
$$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$$
$$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$$

Praktische Implementatie voor Diffusie

Omdat dLLMs geen exacte traject-log-kansen blootleggen, introduceren de auteurs twee belangrijke componenten om TraFL praktisch toepasbaar te maken:

1. Diffusie-Compatibele Sequentie-niveau Surrogaat: Zij leiden een hanteerbare log-kans surrogaat ($\log \hat{p}_\theta(y|x)$) af gebaseerd op een ondergrens voor gemaskerde reconstructie. Dit omvat het bemonsteren van een maskerniveau $l$, het corrumperen van $l$ tokens, en het middelen van de log-kansen van de gereconstrueerde tokens. Een factor $1/l$ wordt gebruikt om de score te normaliseren, onafhankelijk van het aantal gemaskerde tokens.
2. Gestudeerde Prompt-afhankelijke Normalisatie: In tegenstelling tot methoden die $Z(x)$ schatten via softmax over een batch (wat concentratie kan veroorzaken), leert TraFL een partitiefunctie $Z_\phi(x)$ via een kleine MLP. Deze term ontvangt gradiëntsignalen bij elke stap, ongeacht welke voltooiingen worden bemonsterd, en voorkomt zo het "batch-lokale" normalisatiefoutpatroon dat wordt gezien bij andere benaderingen voor verdelingsmatching.

Belangrijkste Bijdragen

1. Identificatie van Trajectvergrendeling: Het artikel identificeert en bewijst formeel dat beloningsmaximerende doelstellingen in dLLMs leiden tot een instorting van trajectdiversiteit, wat op zijn beurt de dekking van verschillende correcte oplossingsmodi beperkt.
2. TraFL-doelstelling: Het voorstellen van een traject-balansdoelstelling die het beleid verankert aan een referentiemodel, waarbij de rol van beloning (het kantelen van terminale voltooiingen) wordt gescheiden van die van het referentiemodel (het toewijzen van waarschijnlijkheidsmassa over plausibele trajecten).
3. Praktisch Surrogaat en Normalisatie: De ontwikkeling van een diffusie-compatibele sequentie-niveau surrogaat en een gestudeerde normalisatieterm die stabiele training mogelijk maakt zonder exacte trajectkansen.
4. Empirische Superioriteit: TraFL is de enige geëvalueerde na-trainingsmethode die consistent presteert boven het basismodel over alle benchmark-lengte-instellingen (GSM8K, MATH-500, HumanEval, MBPP) en bemonsteringsbudgetten.

Resultaten

De auteurs evalueren TraFL op benchmarks voor wiskundige redenering en codegeneratie tegen het basismodel (LLaDA-8B-Instruct) en sterke baselines (ESPO, JUSTGRPO).

• In-distributie Prestaties: TraFL verbetert Pass@k over alle bemonsteringsbudgetten ($k \in \{1, 2, 4, 8, 16\}$) en temperaturen. Opmerkelijk is dat het Pass@1 (nauwkeurigheid bij enkele bemonstering) verbetert, terwijl het ook beter schaalt met verhoogde bemonsteringsbudgetten, in tegenstelling tot baselines die vaak de nauwkeurigheid bij enkele bemonstering opofferen voor winst bij meerdere bemonsteringen.
• Houd-out Generalisatie: Winsten transfereren naar out-of-distribution benchmarks. Op Minerva Math blijft TraFL boven het basismodel. Op LiveCodeBench is TraFL de sterkste methode over alle moeilijkheidsgraad-splits (Gemakkelijk, Middel, Moeilijk), en presteert het aanzienlijk beter dan baselines op middelzware en moeilijke problemen.
• Diversiteitsanalyse: Met behulp van een LLM-as-a-judge-protocol tonen de auteurs aan dat de winsten van TraFL gepaard gaan met een bredere dekking van verschillende oplossingsstrategieën (bijvoorbeeld verschillende bewijsaanpakken of algoritmische decomposities) in plaats van alleen het scherper maken van de score van een enkele modus. Baselines zoals ESPO, hoewel ze de nauwkeurigheid verbeteren, vertonen geen significante diversiteitswinsten ten opzichte van het basismodel, wat consistent is met de trajectvergrendelingshypothese.

Betekenis en Claims

Het artikel stelt dat TraFL een principiële oplossing biedt voor het "modus-zoekende" faalmechanisme van standaard RL na-training voor diffusiemodellen. Door over te stappen van beloningsmaximalisatie naar traject-balans, behoudt de methode de diversiteit aan gedenoiseerde paden die nodig is om meerdere geldige oplossingen te dekken.

De auteurs stellen dat TraFL de enige geëvalueerde methode is die in elke geteste benchmark-lengte-instelling presteert boven het basismodel. Zij benadrukken dat de verbeteringen robuust zijn tegen verhogingen van het bemonsteringsbudget en generaliseren naar houd-out datasets, wat suggereert dat de traject-balansdoelstelling het beloningssignaal effectief ontkoppelt van het pad-toewijzingsmechanisme. Het werk pleit voor verdelingsmatching-doelstellingen boven pure beloningsmaximalisatie voor dLLM na-training, en biedt een praktisch raamwerk dat concentratiefalenmodi vermijdt die inherent zijn aan standaard policy-gradient-benaderingen.

Het artikel blijft bescheiden wat beperkingen betreft, en merkt op dat de methode vertrouwen stelt op een benaderende sequentie-niveau surrogaat in plaats van exacte trajectkansen, en gebruik maakt van grove binaire beloningen (exacte overeenkomst of testgeslaagd), wat de leerefficiëntie kan beperken in scenario's met zeldzame successen.