Technische Samenvatting: CForce: Het verbeteren van parallelle decodering voor dLLM's via Consistency Forcing

1. Probleemstelling

Diffusion Large Language Models (dLLM's) versnellen tekstgeneratie door meerdere gemaskeerde tokens in een enkele forward pass te voorspellen, wat een hoge parallelliteit bereikt. Echter, onder agressieve parallelle decoderingsstrategieën (waarbij veel tokens gelijktijdig worden vastgelegd op basis van betrouwbaarheidsdrempels), lijden de vroege denoising-stadia vaak aan onbetrouwbare voorspellingen door beperkte context. Deze vroege fouten propageren naar opeenvolgende stadia, wat de algehele generatiekwaliteit verslechtert.

Bestaande acceleratiemethoden kampen met twee primaire beperkingen:

1. Inference-time benaderingen (bijv. caching, speculatieve decodering) wijzigen geen modelparameters om de inherente voorspellingsbetrouwbaarheid te verbeteren.
2. Training-gebaseerde benaderingen vertrouwen vaak op geconstrueerde, door docenten gegenereerde of geprivilegieerde trajecten. Deze doelen verschillen van de werkelijke inferentiestaten van het studentmodel (lage context, door drempelwaarden besliste paden), wat leidt tot een mismatch tussen training en inferentie.

De kernuitdaging is om de betrouwbaarheid van vroege voorspellingsstadia onder hoge parallelliteit te verbeteren zonder externe docenten te gebruiken of het fundamentele decoderingsalgoritme te wijzigen.

2. Methodologie: Consistency Forcing (CForce)

CForce is een distillatiemethode die is ontworpen om de voorspellingen van vroege, laag-contextuele denoising-stadia af te stemmen op die van latere, meer informatieve stadia binnen het eigen decoderingspad van het model.

2.1 Constructie van gestaffelde trajecten

In plaats van te trainen op elk inheems decoderingsstap, construeert CForce "gestaffelde" trajecten uit zelf-rollouts van een vooraf getraind dLLM:

• Zelf-rollout: Het model genereert een traject met behulp van confidence-gebaseerde parallelle decodering (het onthullen van tokens boven een drempelwaarde $\tau$).
• Stadia-partitionering: Het traject wordt gedownsampled in stadia op basis van het cumulatieve aantal nieuw onthulde tokens (stadiumgrootte $S$), in plaats van het aantal natuurlijke stappen. Dit zorgt ervoor dat aangrenzende stadia significante contextverschuivingen vertegenwoordigen.
• Aangrenzende paren: Trainingsparen bestaan uit een vroeger stadium ($\bar{s}_i$) en een later stadium ($\bar{s}_{i+1}$). Het latere stadium dient als een stop-gradient doel voor het eerdere stadium.

2.2 Trainingsdoelstelling

De methode maakt gebruik van een samengestelde verliesfunctie om consistentie af te dwingen:

1. Confidence Adaptive KL Divergence (CAD):

   • Toegepast op posities die in beide stadia gemaskeerd blijven ($U_i$).
   • Interpoleert dynamisch tussen Forward KL en Reverse KL op basis van de betrouwbaarheid ($c_j$) van de voorspelling in het latere stadium.
   • Forward KL: Minimaliseert distributionele drift (stabiel).
   • Reverse KL: Stimuleert mode-seeking (het verscherpen van voorspellingen), maar wordt gewogen door betrouwbaarheid om instorting op onzekere doelen te voorkomen.
   • Formule: $L_{CAD} = \frac{1}{|U_i|} \sum_{j \in U_i} [D_{KL}(p^{tar}_j \| p^{stu}_j) + c_j D_{KL}(p^{stu}_j \| p^{tar}_j)]$

2. Cross-Entropy (CE) Anchor:

   • Toegepast op nieuw onthulde posities ($\Delta_i$) tussen stadia.
   • Verankert de voorspelling van de student aan het schone doel-token direct vóór de commitment, waardoor de overgang van gemaskeerde naar onthulde staten wordt gestabiliseerd.

3. Curriculum Learning:

   • Om de moeilijkheid van het voorspellen van tokens die ver vooruit worden onthuld aan te pakken, wordt het model geleidelijk blootgesteld aan de volledige kloof tussen de stadia. Het deel van de onthulde tokens dat in het doelstadium wordt gebruikt, neemt toe van 10% naar 100% over de trainingsstappen heen.

2.3 Extensie naar Edit-Capable Modellen

Voor modellen die Token-to-Token (T2T) verfijning ondersteunen (bijv. LLaDA 2.1), maakt CForce gebruik van correcties in latere stadia:

• Edit-signalen: Posities waar tokens worden herzien in latere stadia ($\bar{E}_i$) worden opgenomen in het CAD-domein.
• Supervisie: Deze T2T-correcties bieden rijkere supervisie voor eerdere Mask-to-Token (M2T) voorspellingen, waardoor het model kan leren van zijn eigen vermogen tot zelfcorrectie.

2.4 Theoretische Rechtvaardiging

Het artikel biedt een theoretische bovengrens die aantoont dat de voorspellingsfout in het vroege stadium begrensd wordt door de som van de aangrenzende stadium-distributionele drift en de tokenfout bij de onthulingsgrens. De CAD-term minimaliseert de drift, terwijl de CE-anchor de grensfout minimaliseert, waardoor de voorspellingsfout in het vroege stadium benaderend wordt geminimaliseerd.

3. Belangrijkste Bijdragen

1. Herformuleerde Acceleratie: Verschuift de focus van dLLM-acceleratie van enkel het reduceren van sampling-stappen naar het verbeteren van de betrouwbaarheid van vroege voorspellingsstadia onder agressieve drempelwaardedecodering.
2. Consistency Forcing: Introduceert een distillatiemethode die aangrenzende stadia afstemt op het eigen decoderingspad van het model met behulp van een stop-gradient doel van hetzelfde model, waardoor de noodzaak voor bevroren docenten vervalt.
3. Nieuwe Doelstelling: Stelt Confidence Adaptive KL Divergence (CAD) voor om een balans te vinden tussen distributionele stabiliteit en mode-seeking, gecombineerd met een CE-anchor voor token commitment.
4. Edit-Capable Integratie: Demonstreert hoe T2T-verfijningssignalen in edit-capable modellen de eerdere M2T-drafting kunnen superviseren, wat de prestaties verbetert in zowel edit- als non-edit scenario's.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op LLaDA 2.0 (non-edit) en LLaDA 2.1 (edit-capable) varianten over vier benchmarks: GSM8K, MATH500, MBPP en HumanEval.

• Edit-Capable Modellen (LLaDA 2.1-mini):

  • CForce verhoogde de gemiddelde Tokens Per Forward (TPF) van 6.94 naar 9.08.
  • De gemiddelde nauwkeurigheid verbeterde van 85.57% naar 86.41%.
  • De winst was bijzonder sterk op benchmarks voor codegeneratie (MBPP, HumanEval).

• Non-Edit Modellen (LLaDA 2.0-mini):

  • CForce verschoof het model naar een punt met hogere parallelliteit, waarbij de gemiddelde TPF steeg van 3.60 naar 6.42.
  • Hoewel de gemiddelde nauwkeurigheid iets lager was dan die van het basismodel (door de inherente trade-off tussen snelheid en kwaliteit), presteerde het aanzienlijk beter dan de CAP-baseline bij hoge parallelliteit.
  • Onder een vast, agressief TPF-budget (bijv. TPF=8) behaalde CForce de hoogste nauwkeurigheid van alle varianten.

• Doorvoersnelheid (Throughput):

  • CForce verbeterde de wall-clock throughput met 42.82% ten opzichte van de basis LLaDA 2.0-mini en met 22.84% ten opzichte van de basis LLaDA 2.1-mini gemiddeld.

• Ablatie-studies:

  • Stadiumgrootte: $S=8$ bood de beste balans tussen parallelliteit en nauwkeurigheid.
  • Doeltype: Same-model stop-gradient doelen presteerden beter dan frozen-teacher baselines.
  • KL-vorm: CAD presteerde beter dan pure Forward KL (hogere parallelliteit) en pure Reverse KL (hogere nauwkeurigheid), waarbij een superieur compromis werd bereikt.

5. Betekenis en Claims

Het artikel claimt dat Consistency Forcing effectief de betrouwbaarheidsbottleneck in hoog-parallele dLLM-decodering aanpakt. Door het afdwingen van vroege voorspellingen om af te stemmen op latere, meer geïnformeerde staten afgeleid van het eigen traject van het model, maakt CForce het volgende mogelijk:

• Hogere Parallelliteit: Modellen kunnen meer tokens per forward pass vastleggen zonder catastrofale foutpropagatie.
• Verbeterde Snelheid-Kwaliteit Trade-off: De methode duwt de operationele frontier, waardoor snellere generatie mogelijk is met behoud of verbetering van de kwaliteit, vooral in edit-capable scenario's waar zelfcorrectie-signalen beschikbaar zijn.
• Algemene Toepasbaarheid: De aanpak is effectief voor zowel standaard gemaskeerde diffusie als edit-capable modellen, waarbij gebruik wordt gemaakt van de specifieke sterktes van elk (M2T drafting en T2T verfijning).

De auteurs merken op dat de methode momenteel afhankelijk is van vooraf verzamelde trajecten van een vaste initialisatie, wat kan afwijken van de distributie van de student naarmate de training vordert. Toekomstig werk wordt gesuggereerd om online trajectverzameling te verkennen om de forcing-doelstelling aan te passen aan de evoluerende het model.