DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs

Technische Samenvatting: DSL-LLaDA: Schalen van Continue Denoising naar 8B Masked Diffusion LMs

Probleemstelling
Discrete Masked Diffusion Language Models (DLMs), zoals LLaDA, genereren tekst via iteratieve parallelle decodering. Hoewel deze modellen zijn opgeschaald naar miljarden parameters, verslechtert de inferentiekwaliteit sterk bij lage stap-budgetten (generatie met weinig stappen). Het kernproblek is een afweging tussen outputlengte en kwaliteit: met een vast aantal forward passes produceert standaard iteratieve unmasking ofwel korte, hoogwaardige outputs, ofwel lange, repetitieve en incoherente tekst. Deze mislukking komt voort uit de factorisatie van de reverse transitie, waarbij elke tokenpositie onafhankelijk wordt gedecodeerd van zijn context-geconditioneerde marginale distributie, waardoor inter-token correlaties worden genegeerd. Wanneer de stap-budgetten laag zijn, moet elke stap veel posities simultaan vastleggen, wat deze fouten versterkt en leidt tot voortijdige terminatie of repetitie-loops.

Hoewel continue denoising in de embedding-ruimte een theoretische oplossing biedt door alle posities gezamenlijk te evolueren om de factorisatie in de token-ruimte te omzeilen, blijft het bouwen van een dergelijk model vanaf nul op grote schaal een openstaand ontwerpprobleem vanwege de complexiteit van embedding-geometrie, ruis-scheduling en trainingsdoelstellingen.

Methodologie
De auteurs stellen DSL-LLaDA voor, een methode om een bestaande pretrained masked DLM (specifiek LLaDA-8B-Instruct) licht aan te passen om continue embedding-ruimte denoising te ondersteunen, in plaats van een continu model vanaf nul te trainen.

1. Discrete Stochastic Localization (DSL) Adaptatie:

   • De methode vervangt de binaire clean-of-mask input van standaard DLMs door continue ruisige staten. Voor een clean token $x_i$ wordt een Signal-to-Noise Ratio (SNR) $\gamma_i$ gesampled, en een ruisige staat $z_i$ wordt geconstrueerd als een mengeling van de token-embedding en Gaussische ruis.
   • Een softmax converter brengt deze continue DSL-staten terug naar de input embedding-ruimte van de backbone. Deze converter produceert een "soft mask" waarbij een lage SNR overeenkomt met een mask-achtige staat, een gemiddelde SNR met onbetrouwbare willekeurige tokens, en een hoge SNR met duidelijke token-embeddings.
   • De backbone wordt slechts voor 1.000 stappen continue-pretrained op FineWeb-Edu. De trainingsdoelstelling blijft standaard token cross-entropy, maar de input-staten zijn nu continue mengelingen in plaats van binaire masks. Cruciaal is dat de converter de model mogelijk maakt deze inputs te verwerken zonder architecturale wijzigingen, aangezien de SNR-invariantie van de optimale denoiser de noodzaak voor expliciete timestep-inputs wegneemt.

2. Continue Inferentie (SDE Sampler):

   • Tijdens inferentie behoudt het model een continue staat voor alle posities. Het gebruikt een op Stochastic Differential Equations (SDE) gebaseerde sampler (Heun stappen) om de staat te evolueren van een lage SNR (hoge onzekerheid) naar een hoge SNR.
   • Het model voorspelt een waarschijnlijkheidsverdeling over tokens, die wordt omgezet in een gewogen gemiddelde van token-embeddings (gebruikmakend van de top-$M$ meest waarschijnlijke tokens). Deze embedding-schatting werkt de continue staat bij.
   • Hard token commitment wordt uitgesteld tot de laatste stap. Dit stelt het model in staat om alle posities gezamenlijk te verfijnen, waardoor inter-token correlaties gedurende het hele generatieproces behouden blijven.

3. Conditionele Generatie:

   • Voor taken zoals samenvatting worden context-tokens vastgezet (clamped) op hun clean embeddings, terwijl de respons-regio een continue denoising ondergaat. De backbone attention-ert bidirectioneel, waarbij de denoising van de respons wordt geconditioneerd op de clean context.

Belangrijkste Bijdragen

1. Lichtgewicht Conversie: Gedemonstreerd dat een pretrained 8B masked DLM kan worden geconverteerd om continue embedding-ruimte denoising te ondersteunen via een minimale 1.000-staps adaptatie via DSL, waardoor het trainen van een continu model vanaf nul wordt vermeden.
2. Vermijden van Low-NFE Failure Modes: Getoond dat continue SDE-inferentie effectief de voortijdige terminatie en repetitie-loops voorkomt die kenmerkend zijn voor iteratieve unmasking bij lage stap-budgetten (NFE $\le$ 16).
3. Selectieve Noisy-State Robustheid: Geïdentificeerd dat de adaptatie een robuustheid induceert waarbij het model gecorrumpeerde tokens kan corrigeren terwijl clean tokens worden behouden (selectieve correctie), een gedrag dat afwezig is bij discrete baselines met gelijke rekenkracht.

Resultaten

• Zero-Shot Samenvatting: Op vier benchmarks (XSum, CNN/DailyMail, arXiv, PubMed) bereikte DSL-LLaDA-SDE de beste ROUGE-1 scores bij lage stap-budgetten ($\le$ 16 forward passes) zonder lengtecontrole-heuristieken. In contrast hiermee ofwelmed de discrete baselines de outputs voortijdig afgekapt of leden ze aan hoge repetitiegraat.
• Repetitie en Lengte: Bij open-ended generatie hield DSL-LLaDA-SDE de repetitiegraad onder de 9% over alle stap-budgetten, terwijl discrete baselines de repetitie overschreden met meer dan 60% bij lage NFE of afkortten naar zeer korte lengtes.
• Foutcorrectie: Het aangepaste model vertoonde "selectieve correctie", waarbij gecorrumpeerde tokens werden hersteld terwijl meer dan 98% van de clean tokens behouden bleef. Ter vergelijking: discrete baselines met willekeurige token-corruptie (XDLM) beschadigden een significant deel van de clean tokens, en standaard masked diffusion (MDM-CPT) vertoonde verwaarloosbare correctiecapaciteit.
• Controle-experimenten: Controle-experimenten met standaard masked diffusion training met dezelfde rekenkracht (MDM-CPT) en willekeurige token-corruptie (XDLM) slaagden er niet in de voordelen van continue denoising of de selectieve robuustheid te repliceren, wat bevestigt dat de verbeteringen voortkomen uit continue-ruis blootstelling in plaats van enkel extra rekenkracht.

Betekenis en Claims
Het artikel claimt dat DSL-LLaDA een algemeen bruikbare conversieprocedure biedt voor masked DLMs, waardoor ze de voordelen van continue denoising (gezamenlijke evolutie van tokens, uitgestelde commitment) kunnen benutten met minimale trainingsoverhead. De resultaten suggereren dat continue-ruis adaptatie het probleem van de lengte-kwaliteit afweging, inherent aan iteratieve unmasking bij lage stap-budgetten, kan oplossen.

De auteurs blijven bescheiden over de bredere implicaties:

• De evaluatie is primair gericht op anchored generation (samenvatting en instructie-opvolging). Open-ended creatief schrijven vertoont nog steeds hogere graden van degeneratie (24% vs. $\le$4% bij samenvatting).
• De methode is momenteel alleen gevalideerd op de LLaDA-8B schaal; prestaties op andere architecturen of grotere schalen zijn nog niet getest.
• Hoewel SDE beter presteert dan discrete methoden bij lage NFE, evenaart of overtreft vanilla LLaDA de prestaties van SDE bij long-reference taken bij hoge NFE (waar voortijdige terminatie minder een probleem is).
• De auteurs claimen geen gelijkwaardigheid met autoregressieve (AR) modellen van dezelfde schaal, en merken op dat AR-modellen nog steeds leiden op long-reference taken bij hoge stap-budgetten.
• Menselijke evaluatie wordt geïdentificeerd als toekomstig werk.