Technische Samenvatting: Data-efficiënte Autoregressieve-naar-Diffusie Taalmodellen via On-Policy Distillatie

1. Probleemstelling

Het artikel behandelt de uitdagingen die inherent zijn aan het transformeren van vooraf getrainde Autoregressieve Taalmodellen (ARLMs) naar Diffusie Taalmodellen (DLMs). Hoewel recent werk heeft geprobeerd ARLMs naar DLMs te converteren door causale aandacht te vervangen door bidirectionele aandacht en opnieuw te trainen met een diffusie-doelstelling, identificeren de auteurs twee kritieke "distributieverschuivingen" die de trainingsefficiëntie en prestaties beperken:

1. Kennisbehoud-mismatch: De overgang van een doelstelling voor het voorspellen van het volgende token (ARLM) naar een diffusie-doelstelling (DLM) kan de substantiële kennis die tijdens de oorspronkelijke ARLM-pretraining is verworven, verzwakken of weggooien.
2. Train-Inference State Mismatch: Standaard DLMs worden getraind op sequenties die gecorrumpeerd zijn door uniform willekeurige masking (forward proces). Echter, tijdens inferentie genereert het model sequenties via reverse unmasking trajecten geleid door confidence-gebaseerde decoding-heuristieken. De toestanden die tijdens inferentie worden aangetroffen, verschillen aanzienlijk van de willekeurige toestanden die tijdens training worden gezien, wat leidt tot suboptimale leerprocessen.

Bestaande conversiemethoden behandelen dit vaak als een voortzetting van de pretraining, waarbij ze er niet expliciet in slagen deze mismatches aan te pakken of vereisen prohibitief grote datasets (biljoenen tokens) om DLMs vanaf nul te pretrainen.

2. Methodologie: On-Policy Diffusie Taalmodellen (OPDLM)

De auteurs stellen OPDLM voor, een framework dat ARLMs converteert naar DLMs met behulp van On-Policy Distillatie (OPD). Deze benadering herkadert de conversie als een post-training probleem in plaats van een pretraining probleem.

Kernmechanisme

OPDLM maakt gebruik van een "self-OPD" strategie waarbij het studentmodel zijn eigen trainingsdata genereert, en het oorspronkelijke ARLM fungeert als een bevroren docent (teacher).

• Studentmodel: Een ARLM-checkpoint die is aangepast om blokwijze causale (bidirectionele) aandacht te gebruiken, geïnitialiseerd vanuit de oorspronkelijke ARLM-gewichten.
• Docentmodel (Teacher): Het oorspronkelijke, bevroren ARLM.
• Trainings-traject (On-Policy): In plaats van willekeurige masks te samplen uit een schone dataset, samplet het student DLM zijn eigen reverse decoding trajecten ($\tau \sim \Gamma^{rev}_{\theta}$). Deze trajecten vertegenwoordigen de werkelijke toestanden die het model tegenkomt tijdens inferentie.
• Distillatie-doelstelling:
  • Vanuit een gesampled reverse traject wordt een gedeeltelijk gedenoised staat $x_t$ geselecteerd.
  • Voor elke gemaskeerde tokenpositie $i$ in een blok, wordt de voorspellende distributie van de student $p_{\theta}$ getraind om de distributie van de docent te matchen.
  • Cruciaal is dat de docent-distributie geconditioneerd is op de unmasked causale prefix van de terminale sequentie $x_0$ (de uiteindelijke gegenereerde output van dat traject). Specifiek is het doel $p_{ARLM}(\cdot | x_{<b}^0, x_{<i}^0)$.
  • De loss is een Kullback-Leibler (KL) divergentie tussen de token-niveau distributie van de docent en de distributie van de student op de gemaskeerde posities.

Curriculum Learning

Om instabiliteit direct na de conversie aan te pakken (waarbij de student kwalitatief minder goede trajecten kan genereren door de structurele verschuiving van causale naar blokwijze aandacht), introduceren de auteurs een Rollout-Length Curriculum. De lengte van de on-policy generatie wordt progressief verhoogd van een minimum ($L_{min}$) naar een maximum ($L_{max}$) over de warmup-stappen, wat ervoor zorgt dat de student leert de docent te matchen op coherente contexten voordat hij langere sequenties aanpakt.

3. Belangrijkste Bijdragen

1. Formulering van ARLM-naar-DLM als Post-Training: Het artikel toont aan dat DLMs verkregen kunnen worden zonder de computationeel intensieve fase van pretraining vanaf nul, waardoor de conversie wordt gepositioneerd als een efficiënte post-training procedure.
2. Resolutie van Distributieverschuivingen: Door te trainen op de eigen reverse trajecten van het model (on-policy), elimineert OPDLM de train-inference state mismatch. Door het oorspronkelijke ARLM als docent te gebruiken, waarborgt het de kennisbehoud.
3. Extreme Data-efficiëntie: De methode bereikt concurrerende prestaties met 15× tot 7.000× minder trainings-tokens vergeleken met gevestigde DLM-baselines (bijv. LLaDA, Dream, SDAR).
4. Generale en Taakspecifieke Toepasbaarheid: Het framework laat zien zowel algemene doelgerichte DLMs als taakspecifieke modellen (bijv. voor wiskunde) te produceren zonder dat daar een aparte DLM-pretraining fase voor nodig is.

4. Experimentele Resultaten

De auteurs evalueerden OPDLM (4B en 8B schalen) tegenover baselines zoals LLaDA, Dream, SDAR en Fast-dLLM-v2 over algemene kennis, wiskunde en coding benchmarks.

• Prestaties vs. Efficiëntie: OPDLM-8B behaalde sterke resultaten op benchmarks zoals AIME-24 (14,7% nauwkeurigheid) en GPQA-Diamond (36,1%) terwijl het slechts 0,066 miljard trainings-tokens en 4,2 × 10¹⁸ FLOPs vereiste. In contrast vereisten baselines honderden miljarden tokens (bijv. LLaDA-8B gebruikte 1.500B tokens).
• Ablatie-studies:
  • On-Policy vs. Off-Policy: Het vervangen van de on-policy traject-sampling door standaard willekeurige masking (Off-Policy) resulteerde in consistente prestatieverliezen, wat bevestigt dat het trainen op de eigen inferentie-toestanden van het model de primaire drijfveer voor succes is.
  • Kennisbehoud: OPDLM behield zero-shot capaciteiten die niet expliciet tijdens de conversie werden beoogd, inclusief extended thinking (redeneren binnen <thought> tags) en meertalige generatie, ondanks dat het getraind werd op een gecureerde Engelse corpus.
• Inference Efficiëntie: De modellen ondersteunen multi-token decoding. Door de decoding confidence threshold en blokgrootte aan te passen, demonstreerden de auteurs een trade-off waarbij de throughput meer dan 2–3 tokens per stap kan overschrijden met gematigde kosten voor nauwkeurigheid.
• Taakspecifieke Adaptatie: Wanneer strikt getraind op wiskundige data (MATH dataset), presteerde OPDLM-MATH beter dan standaard DLM-baselines en kwam het overeen met gespecialiseerde RL-getrainde modellen (TraDo) zonder dat een vooraf getrainde DLM of reinforcement learning met verifieerbare beloningen nodig was.

5. Betekenis en Claims

Het artikel claimt dat OPDLM een nieuwe Pareto-frontier vestigt voor de efficiëntie van DLM-training. Door gebruik te maken van on-policy distillatie, stellen de auteurs dat de conversie van ARLMs naar DLMs kan worden bereikt met minimale computationele middelen, waardoor diffusie-gebaseerde tekstgeneratie toegankelijk wordt zonder de "prohibitieve kosten" van pretraining vanaf nul.

De auteurs benadrukken dat deze benadering de gemeenschap in staat stelt om toekomstige ARLMs efficiënt aan te passen naar diffusiemodellen, waarbij de rijke kennis van de oorspronkelijke modellen behouden blijft terwijl de voordelen van parallelle decoding worden verkregen. Het werk wordt gepresenteerd als een praktisch pad om de kloof te overbruggen tussen de volwassenheid van autoregressieve modellen en de potentiële voordelen van diffusiemodellen.