Technische Samenvatting: LESS: Mutual-Stability Sampling voor Diffusion Language Models

Probleemstelling

Diffusion Large Language Models (dLLMs) bieden een veelbelovend alternatief voor autoregressieve (AR) decodering door parallelle token-updates en bidirectionele conditionering mogelijk te maken via iteratieve denoising van gemaskeerde sequenties. De praktische efficiëntie wordt echter momenteel beperkt door sampling-procedures die vertrouwen op vaste reverse denoising-schema's. Standaard samplers voeren een vooraf bepaald aantal stappen uit, wat vaak rekenkracht verspilt aan posities waar voorspellingen al gestabiliseerd zijn, terwijl het tegelijkertijd het risico loopt op voortijdige toezegging aan instabiele tokens. Zodra een token in een diffusieproces is ontmaskerd, wordt deze doorgaans als vast beschouwd; als er door een rigide schema te vroeg een foutief token wordt vastgelegd, kunnen opeenvolgende reverse stappen dit niet meer herzien, wat leidt tot kwaliteitsdegradatie. Bestaande training-vrije adaptieve samplers pakken dit gedeeltelijk aan door gebruik te maken van confidence- of margin-gebaseerde criteria, maar zij slagen er vaak niet in om de volledige temporele en distributionele stabiliteit van het voorspellende traject van een token te vatten.

Methodologie

De auteurs stellen LESS (Less is More) voor, een training-vrije, model-agnostische adaptieve sampler die token-toezegging (commitment) frameert als een per-positie online stop-probleem. LESS implementeert mutual-stability sampling, waarbij een gemaskeerde positie pas in aanmerking komt voor ontmaskering wanneer drie specifieke stabiliteitssignalen gezamenlijk aangeven dat de voorspelling is gestabiliseerd.

De Joint Stability Rule

Voor elke gemaskeerde positie $i$ bij reverse stap $t$, evalueert LESS drie signalen afgeleid van de voorspellende distributie $p_{t,i}$:

1. Top-1 Predictive Confidence: De waarschijnlijkheid die wordt toegewezen aan de huidige top-1 token ($w_{t,i}$). Een positie moet een confidence-drempel $c$ overschrijden.
2. Top-1 Token Persistence: De huidige top-1 token moet ongewijzigd blijven over een recente window van $P$ reverse stappen. Dit voorkomt dat er een toezegging wordt gedaan aan een token dat net is geflipt, zelfs als de huidige confidence hoog is.
3. Top-K Inter-Step Jensen–Shannon Divergence (JSD): Dit is de innovatieve bijdrage van LESS. Het meet de stabiliteit van de voorspellende distributie zelf tussen opeenvolgende reverse stappen.
   • In plaats van de asymmetrische Kullback-Leibler (KL) divergentie, gebruikt LESS de symmetrische en begrensde JSD.
   • Om computationele efficiëntie te garanderen en de focus op relevante kandidaten te houden, wordt de JSD berekend op een verruimde (coarsened) vocabulaire bestaande uit de unie van de top-$K$ tokens van de huidige en vorige stappen, waarbij de resterende massa is samengevoegd in een residu-symbool.
   • De JSD moet onder een drempel $d$ vallen.

Een positie wordt geaccepteerd voor ontmaskering indien en slechts indien:
$$\text{accept}_{t,i} = (\text{conf}_{t,i} \ge c) \land \text{persist}_{t,i} \land (\text{JSD}_{t,i} \le d)$$

Unmasking Policy

LESS hanteert een frontier-first policy voor volledige sequentie-diffusiemodellen. Het probeert de meest linkse gemaskeerde positie te ontmaskeren als deze in aanmerking komt. Indien dit niet het geval is, scant het een look-ahead window om de eerste in aanmerking komende positie te vinden. Als er geen posities in aanmerking komen, wordt de stap overgeslagen. Om beëindiging te garanderen, wordt een forced-unmasking fallback getriggerd na $\Pi$ opeenvolgende overgeslagen stappen, hoewel deze fallback niet wordt geteld als een "stabiele" toezegging.

Belangrijkste Bijdragen

1. Formulering van Token Commitment: Het artikel frameert adaptieve ontmaskering in dLLMs als een per-positie online stop-probleem, geïnduceerd door de onomkeerbaarheid van ontmaskering.
2. Top-K Inter-Step JSD: De auteurs introduceren de top-$K$ inter-step Jensen–Shannon divergentie als een begrensde, symmetrische stabiliteitssignaal. Dit adresseert de beperkingen van eerder werk waarbij confidence alleen niet de voortdurende distributionele beweging kan detecteren, en KL-divergentie directioneel en onbegrensd is.
3. LESS Sampler: Een training-vrije sampler die confidence, token persistence en inter-step JSD combineert in een gezamenlijke stabiliteitsregel. Het vereist geen hertraining van het model en is compatibel met bestaande dLLM checkpoints.
4. Empirische Validatie: Uitgebreide evaluatie over drie dLLM-families (Dream-7B, LLaDA-8B, LLaDA-1.5-8B) en zeven benchmarks die betrekking hebben op wiskunde, code en algemene kennis.

Resultaten

De evaluatie toont aan dat LESS de nauwkeurigheid-efficiëntie trade-off significant verbetert in vergelijking met zowel fixed-budget decoding als andere sterke training-vrije adaptieve samplers (Prophet en KLASS).

• Efficiëntie: Over drie model-families heen gebruikt LESS gemiddeld 72,1% minder reverse stappen vergeleken met fixed-budget decoding (Base).
• Nauwkeurigheid: LESS bereikt een gemiddelde winst in nauwkeurigheid van +0,91 punten ten opzichte van de fixed-budget baseline terwijl het minder stappen gebruikt. Het presteert ook beter dan Prophet en KLASS in de gemiddelde nauwkeurigheid over alle model-families.
• Taakprestaties: De winst is het meest uitgesproken in taken voor wiskunde en codegeneratie. Bijvoorbeeld, op Dream-7B bereikt LESS een gemiddelde nauwkeurigheid van 65,18% (vs. 63,30% voor Base) met slechts 64,3 stappen (vs. 256 voor Base).
• Compute-besparingen: Aangezien elke reverse stap een Transformer forward pass vereist, vertaalt de reductie in stappen zich direct in lagere inference compute. Op Dream-7B GSM8K vermindert LESS de wall-clock latency van 19,45s naar 5,16s en vermindert het het aantal forward evaluaties met ongeveer 75% bij hoge step-budgets.
• Ablatie-studies:
  • Signaal Complementariteit: Geen enkel signaal (confidence, persistence of JSD) presteert zo goed als de gezamenlijke regel. Confidence is de primaire drijfveer, maar persistence en JSD bieden essentiële waarborgen tegen voortijdige toezegging en distributionele drift.
  • Robuustheid: De methode presteert consistent over verschillende modellen en decoding-regimes (full-sequence vs. semi-autoregressive blockwise) met behulp van een enkele globale configuratie $(c=0,75, d=0,040)$.

Betekenis en Claims

Het artikel claimt dat LESS een praktische, training-vrije mechanisme biedt om het decoding-proces van diffusion language models te versnellen door tokens alleen toe te zeggen wanneer hun voorspellingen confident, temporeel persistent en distributioneel stabiel zijn.

De auteurs benadrukken dat stabiliteit meer is dan alleen confidence. Een voorspelling met een hoge confidence kan nog steeds onbetrouwbaar zijn als de top-1 token recentelijk is veranderd of als de waarschijnlijkheidsmassa nog steeds verschuift tussen alternatieven. Door stabiliteit te behandelen als een gezamenlijke eigenschap van de huidige voorspelling en het recente traject, vermijdt LESS de valkuilen van voortijdige toezegging.

Het werk positioneert LESS als orthogonaal aan en complementair met andere acceleratiemethoden:

• Het complementeert training-time methoden (bijv. distillatie) die het reverse proces verkorten.
• Het complementeert systemen-methoden (bijv. caching, speculative decoding) die de per-stap latentie verminderen.
• Het fungeert als een drop-in inference-time policy die gecombineerd kan worden met deze methoden zonder het onderliggende model te wijzigen.

Het artikel concludeert dat mutual stability een levensvatbaar criterium is voor robuuste vroege token-toezegging, waardoor dLLMs een hogere efficiëntie kunnen bereiken zonder de generatiekwaliteit op te offeren, en deze zelfs te verbeteren. De auteurs erkennen dat de methode een kwaliteit-efficiëntie trade-off blootlegt die wordt gecontroleerd door de stabiliteitsdrempels, wat tuning voor specifieke modellen of taken kan vereisen, hoewel een vaste globale configuratie effectief bleek in hun diverse evaluatie-suite.