Technische Samenvatting: GDSD: Versterkend Leren als Geleide Denoiser Zelfdistillatie voor Diffusie Taalmodellen

1. Probleemstelling

Diffusie Large Language Models (dLLMs) bieden aanzienlijke winsten in inferentie-efficiëntie ten opzichte van autoregressieve modellen (ARMs) door gelijktijdig meerdere tokens te genereren en de strikte links-naar-rechts generatievolgorde te doorbreken. Echter, dLLMs lopen momenteel achter op state-of-the-art ARMs wat betreft generatiekwaliteit, wat effectieve fine-tuning-methoden zoals Versterkend Leren (RL) noodzakelijk maakt.

De centrale belemmering voor het toepassen van RL op dLLMs is de onbehandelbaarheid van de policy-likelihood. Bestaande oplossingen vallen over het algemeen in twee categorieën:

1. Traject-Likelihood Methodes: Schatten exacte likelihoods door overgangskansen op te tellen. Hoewel ze onbevooroordeeld zijn, zijn ze rekenkundig duur en niet afgestemd op pre-training doelen.
2. Sequentie-Likelihood (ELBO) Methodes: Gebruiken de Evidence Lower Bound (ELBO), geschat vanuit willekeurig gemaskerde sequenties, als vervanger voor de ware likelihood. Deze zijn rekenkundig efficiënt en afgestemd op pre-training.

Het Kernprobleem: Het paper betoogt dat ELBO-gebaseerde RL-methoden lijden onder een Training-Inferentie Mismatch (TIM)-bias. Deze bias ontstaat uit twee bronnen:

• Likelihood Gap: De ELBO is een ondergrens met een niet-verwaarloosbare kloof ten opzichte van de ware likelihood, wat bias en variantie introduceert in de importance sampling-ratio's die worden gebruikt in Policy Gradient (PG) of PPO-doelstellingen.
• Steekproefverdelingsmismatch: dLLM-inferentie gebruikt iteratief her-masking en blokgewijze decoding, wat een steekproefverdeling ($\pi^{rm}_{old}$) induceert die verschilt van de trainingspolicy-verdeling ($\hat{\pi}_{old}$) die via ELBO wordt geschat.

Deze mismatch zorgt ervoor dat de correctie via importance sampling ineffectief is, wat leidt tot verslechterde prestaties en, in ernstige gevallen, tot catastrofale trainingstorting.

2. Methodologie: Geleide Denoiser Zelfdistillatie (GDSD)

De auteurs stellen Guided Denoiser Self-Distillation (GDSD) voor, een kader dat RL voor dLLMs herformuleert als een off-policy zelfdistillatieprobleem, waardoor de behoefte aan likelihood-schatting wordt omzeild en TIM-bias wordt vermeden.

2.1 Theoretische Grondslag

Onder reverse-KL-geregulariseerd RL heeft de optimale policy $\pi^*$ een gesloten vorm-oplossing:
$$\pi^*(x_0) \propto \pi_{old}(x_0)^{1-\beta} \cdot \pi_{ref}(x_0)^\beta \cdot \exp(\psi A(x_0))$$
waarbij $A(x_0)$ de voordeelfunctie is, $\beta$ een regularisatiecoëfficiënt is, en $\psi$ een geleidingscoëfficiënt is.

De auteurs tonen aan dat deze optimale policy een geleide denoising-verdeling induceert (een "leraar"-denoiser $p^*$) voor het gemaskerde diffusiemodel:
$$p^*(x_0|x_t) = p^{ref}_{old}(x_0|x_t) \cdot \exp(\psi A(x_0) - A_t(x_t))$$
waarbij $p^{ref}_{old}$ een geometrische mengeling is van de oude en referentie-policies, en $A_t(x_t)$ een normalisatieconstante is.

2.2 Het GDSD-doel

In plaats van een op likelihood gebaseerd doel te optimaliseren, distilleert GDSD de huidige denoiser $p_\theta$ direct vanuit de leraar $p^*$ door hun logits te laten overeenkomen. Het doel minimaliseert de gekwadrateerde $L_2$-afstand tussen de student- en leraar-logits:
$$\mathcal{L}_{GDSD} = \mathbb{E} \left[ \left( \log p_\theta(x_0|x_t) - \log p^*(x_0|x_t) \right)^2 \right]$$

2.3 Normalisatie-vrije Herformulering

Een grote uitdaging bij het implementeren hiervan is de onbehandelbare normalisatieconstante $Z_t$ (of $A_t$) in de leraarverdeling. De auteurs stellen twee oplossingen voor om dit te elimineren:

1. Directe Matching: Door gebruik te maken van de translatie-invariantie van de Softmax-operator, kan de constante worden genegeerd als de student-logits overeenkomen met de genormaliseerde leraar-logits tot op een additieve constante.
2. Token-niveau Logit Centralisatie (TLC): Om ongecontroleerde verschuivingen in logit-schaal tijdens iteratieve updates te voorkomen, stellen de auteurs voor om de logits te centraliseren. Door het gemiddelde logit over het vocabulaire voor elke tokenpositie af te trekken, wordt de normalisatieconstante volledig uit het doel verwijderd.

Het uiteindelijke praktische doel (met TLC) is:
$$\mathcal{L}_{GDSD} = \mathbb{E} \left[ \left( \log \bar{p}_\theta(x_0|x_t) - \log \bar{p}_{old}(x_0|x_t) - \psi A(x_0) \right)^2 + \beta \left( \log \bar{p}_\theta(x_0|x_t) - \log \bar{p}_{ref}(x_0|x_t) \right)^2 \right]$$
waarbij $\bar{p}$ het logit-gecentraliseerde model aanduidt.

3. Belangrijkste Bijdragen

1. Reductie tot Zelfdistillatie: Het paper stelt vast dat reverse-KL-geregulariseerd RL voor dLLMs equivalent is aan het distilleren van een geleide denoiser vanuit een zelf-leraar die is afgeleid van de gesloten vorm optimale policy.
2. GDSD Kader: Een nieuwe, likelihood-vrije zelfdistillatiemethode die TIM-bias per ontwerp elimineert, omdat deze niet afhankelijk is van importance sampling-ratio's of ELBO-vervangers.
3. Normalisatie-vrije Optimalisatie: De introductie van Token-niveau Logit Centralisatie (TLC) maakt efficiënte online RL mogelijk zonder onbehandelbare partitiefuncties te berekenen.
4. Theoretische Unificatie: Het paper demonstreert dat bestaande ELBO-gebaseerde methodes (zoals wd1, DMPO, SPG, UniGRPO) kunnen worden gezien als specifieke voorbeelden van distillatie met verschillende divergentiemaatstaven (Forward-KL of Reverse-KL), maar dat ze lijden onder data-inefficiëntie of TIM-bias die GDSD vermijdt.

4. Experimentele Resultaten

De auteurs evalueerden GDSD op LLaDA-8B en Dream-7B over zes benchmarks:

• Wiskundig Redeneren: GSM8K, MATH500
• Planning: Countdown, Sudoku
• Coderen: HumanEval, MBPP

Belangrijkste Bevindingen:

• Prestatiewinsten: GDSD presteert consequent beter dan state-of-the-art ELBO-gebaseerde methodes. Op het Dream-7B model bereikte GDSD verbeteringen in testnauwkeurigheid van maximaal +19,6% op planningstaken (Sudoku en Countdown). Op LLaDA-8B varieerden de verbeteringen van +0,6% tot +5% over planning-, wiskunde- en coderingsbenchmarks.
• Trainingsstabiliteit: GDSD toonde aanzienlijk stabielere trainingsbeloningsdynamiek in vergelijking met ELBO-gebaseerde baselines (zoals SPG, ESPO), die oscillaties of instorting vertoonden.
• Ablatiestudies: Het Token-niveau Logit Centralisatie (TLC)-component bleek training verder te stabiliseren en hielp in sommige gevallen (zoals Sudoku) het model om lokale minima te ontvluchten om hogere uiteindelijke beloningen te bereiken.

5. Betekenis en Beweringen

Het paper beweert dat directe denoiser zelfdistillatie een stabielere en effectievere alternatief biedt voor ELBO-gebaseerde RL voor dLLMs. Door de afhankelijkheid van likelihood-vervangers te verwijderen, vermijdt GDSD de structurele biases die inherent zijn aan importance-sampling-benaderingen. De resultaten suggereren dat de waargenomen prestatiewinsten niet louter het gevolg zijn van hyperparameter-tuning, maar voortkomen uit een meer principieel optimalisatieroute die het trainingsdoel afstemt op de daadwerkelijke inferentie-verdeling van dLLMs. De auteurs concluderen dat ELBO niet noodzakelijk is voor RL met dLLMs en dat het verwijderen ervan leidt tot robuustere afstemming en redeneervermogens.