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

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

Problem Statement
Discrete Masked Diffusion Language Models (DLMs), such as LLaDA, generate text via iterative parallel decoding. While these models have scaled to billions of parameters, their inference quality degrades sharply under low step budgets (few-step generation). The core issue is a trade-off between output length and quality: with a fixed number of forward passes, standard iterative unmasking either produces short, high-quality outputs or long, repetitive, and incoherent text. This failure stems from the factorization of the reverse transition, where each token position is decoded independently from its context-conditioned marginal distribution, neglecting inter-token correlations. When step budgets are low, each step must commit many positions simultaneously, amplifying these errors and leading to premature termination or repetition loops.

While continuous denoising in embedding space offers a theoretical solution by evolving all positions jointly to sidestep token-space factorization, building such a model from scratch at scale remains an open design problem due to complexities in embedding geometry, noise scheduling, and training objectives.

Methodology
The authors propose DSL-LLaDA, a method to lightly adapt an existing pretrained masked DLM (specifically LLaDA-8B-Instruct) to support continuous embedding-space denoising, rather than training a continuous model from scratch.

1. Discrete Stochastic Localization (DSL) Adaptation:

   • The method replaces the binary clean-or-mask input of standard DLMs with continuous noisy states. For a clean token $x_i$, a Signal-to-Noise Ratio (SNR) $\gamma_i$ is sampled, and a noisy state $z_i$ is constructed as a mixture of the token embedding and Gaussian noise.
   • A softmax converter maps these continuous DSL states back to the backbone's input embedding space. This converter produces a "soft mask" where low SNR corresponds to a mask-like state, mid SNR to unreliable random tokens, and high SNR to clear token embeddings.
   • The backbone is continue-pretrained for only 1,000 steps on FineWeb-Edu. The training objective remains standard token cross-entropy, but the input states are now continuous mixtures rather than binary masks. Crucially, the converter allows the model to process these inputs without architectural changes, as the SNR-invariance of the optimal denoiser removes the need for explicit timestep inputs.

2. Continuous Inference (SDE Sampler):

   • During inference, the model maintains a continuous state for all positions. It uses a Stochastic Differential Equation (SDE) based sampler (Heun steps) to evolve the state from low SNR (high uncertainty) to high SNR.
   • The model predicts a probability distribution over tokens, which is converted into a weighted average of token embeddings (using the top-$M$ likely tokens). This embedding estimate updates the continuous state.
   • Hard token commitment is deferred until the final step. This allows the model to jointly refine all positions, preserving inter-token correlations throughout the generation process.

3. Conditional Generation:

   • For tasks like summarization, context tokens are clamped to their clean embeddings, while the response region undergoes continuous denoising. The backbone attends bidirectionally, conditioning the denoising of the response on the clean context.

Key Contributions

1. Lightweight Conversion: Demonstrated that a pretrained 8B masked DLM can be converted to support continuous-space denoising via a minimal 1,000-step adaptation using DSL, avoiding the need to train a continuous model from scratch.
2. Avoidance of Low-NFE Failure Modes: Showed that continuous SDE inference effectively avoids the premature termination and repetition loops characteristic of iterative unmasking at low step budgets (NFE $\le$ 16).
3. Selective Noisy-State Robustness: Identified that the adaptation induces a robustness where the model can correct corrupted tokens while preserving clean ones (selective correction), a behavior absent in same-compute discrete baselines.

Results

• Zero-Shot Summarization: On four benchmarks (XSum, CNN/DailyMail, arXiv, PubMed), DSL-LLaDA-SDE achieved the best ROUGE-1 scores at low step budgets ($\le$ 16 forward passes) without length-control heuristics. In contrast, discrete baselines either truncated outputs early or suffered from high repetition rates.
• Repetition and Length: In open-ended generation, DSL-LLaDA-SDE maintained repetition rates below 9% across all step budgets, whereas discrete baselines exceeded 60% repetition at low NFE or truncated to very short lengths.
• Error Correction: The adapted model demonstrated "selective correction," recovering corrupted tokens while preserving over 98% of clean tokens. In comparison, discrete baselines with random-token corruption (XDLM) damaged a significant portion of clean tokens, and standard masked diffusion (MDM-CPT) showed negligible correction ability.
• Control Experiments: Control experiments using standard masked diffusion training with the same compute (MDM-CPT) and random-token corruption (XDLM) failed to replicate the continuous denoising benefits or the selective robustness, confirming that the improvements stem from continuous-noise exposure rather than additional compute alone.

Significance and Claims
The paper claims that DSL-LLaDA provides a general-purpose conversion procedure for masked DLMs, enabling them to leverage the benefits of continuous denoising (joint evolution of tokens, delayed commitment) with minimal training overhead. The results suggest that continuous-noise adaptation can resolve the length-quality trade-off inherent in iterative unmasking at low step budgets.

The authors remain modest regarding broader implications:

• The evaluation is scoped primarily to anchored generation (summarization and instruction following). Open-ended creative writing still shows higher rates of degeneration (24% vs. $\le$4% on summarization).
• The method is currently validated only on the LLaDA-8B scale; performance on other architectures or larger scales is untested.
• While SDE outperforms discrete methods at low NFE, vanilla LLaDA matches or exceeds SDE performance on long-reference tasks at high NFE (where premature termination is less of an issue).
• The authors do not claim parity with same-scale autoregressive (AR) models, noting that AR models still lead on long-reference tasks at high step budgets.
• Human evaluation is identified as future work.