Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Soft-Masked Diffusion Language Models

Michael Hersche, Samuel Moor-Smith, Thomas Hofmann, Abbas Rahimi

TL;DR

This work addresses the information loss inherent in binary masking in masked diffusion language models (MDLMs) by introducing Soft-Masked Diffusion Language Models (SM). SM blends the mask token with a weighted mixture of the top-$k$ predictions from the previous step, guided by a confidence-based weight $\lambda(\mathbf{p}) = \omega_s \sigma(\omega_a(-H(\mathbf{p})-\omega_b))$, enabling continuous feedback and partial information propagation across decoding steps. The authors propose a training procedure that jointly optimizes SM parameters with the backbone MDLM in a two-pass fashion, and demonstrate consistent performance gains on both a $169$M-parameter language model and large-scale Dream models for code generation, including improvements in perplexity, MAUVE, and coding benchmarks under high-throughput decoding budgets. SM also complements existing efficiency techniques like unmasking schedulers, caching, and Fast-dLLM, yielding faster, higher-quality generation in high-NFE regimes. The results suggest continuous feedback is a viable direction to enhance discrete diffusion models for language and code tasks, with practical implications for faster, more reliable generation systems.

Abstract

Diffusion models have demonstrated strong potential in language modeling, offering various advantages over traditional autoregressive approaches. Their ability to generate and revise entire responses in parallel enables faster generation and built-in self-correction mechanisms. Most modern diffusion-based language models employ masked diffusion, where decoding involves iteratively processing masked tokens based on a binary decision: either retaining the mask or replacing it with the predicted token. However, this binary choice discards valuable predictive information when the mask is retained. To address this limitation, we introduce soft-masking (SM), a novel method that dynamically blends the embedding of the mask token with the embeddings of the top-$k$ predicted tokens from the previous decoding step, for each retained mask. This provides the model with a more informative prior, preserving context from earlier computations and allowing partial information about masked tokens to propagate beyond a single step. We propose a training methodology that adapts a pretrained masked diffusion language model to incorporate SM. We demonstrate that continuing pretraining a 169M parameter model with SM leads to improved perplexity and MAUVE scores. Furthermore, we finetune two state-of-the-art diffusion models, Dream-7B and Dream-Coder-7B, with SM. SM consistently improves performance across multiple coding benchmarks, particularly in high-throughput settings.

Soft-Masked Diffusion Language Models

TL;DR

This work addresses the information loss inherent in binary masking in masked diffusion language models (MDLMs) by introducing Soft-Masked Diffusion Language Models (SM). SM blends the mask token with a weighted mixture of the top- predictions from the previous step, guided by a confidence-based weight , enabling continuous feedback and partial information propagation across decoding steps. The authors propose a training procedure that jointly optimizes SM parameters with the backbone MDLM in a two-pass fashion, and demonstrate consistent performance gains on both a M-parameter language model and large-scale Dream models for code generation, including improvements in perplexity, MAUVE, and coding benchmarks under high-throughput decoding budgets. SM also complements existing efficiency techniques like unmasking schedulers, caching, and Fast-dLLM, yielding faster, higher-quality generation in high-NFE regimes. The results suggest continuous feedback is a viable direction to enhance discrete diffusion models for language and code tasks, with practical implications for faster, more reliable generation systems.

Abstract

Diffusion models have demonstrated strong potential in language modeling, offering various advantages over traditional autoregressive approaches. Their ability to generate and revise entire responses in parallel enables faster generation and built-in self-correction mechanisms. Most modern diffusion-based language models employ masked diffusion, where decoding involves iteratively processing masked tokens based on a binary decision: either retaining the mask or replacing it with the predicted token. However, this binary choice discards valuable predictive information when the mask is retained. To address this limitation, we introduce soft-masking (SM), a novel method that dynamically blends the embedding of the mask token with the embeddings of the top- predicted tokens from the previous decoding step, for each retained mask. This provides the model with a more informative prior, preserving context from earlier computations and allowing partial information about masked tokens to propagate beyond a single step. We propose a training methodology that adapts a pretrained masked diffusion language model to incorporate SM. We demonstrate that continuing pretraining a 169M parameter model with SM leads to improved perplexity and MAUVE scores. Furthermore, we finetune two state-of-the-art diffusion models, Dream-7B and Dream-Coder-7B, with SM. SM consistently improves performance across multiple coding benchmarks, particularly in high-throughput settings.
Paper Structure (54 sections, 11 equations, 6 figures, 6 tables, 1 algorithm)

This paper contains 54 sections, 11 equations, 6 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. This work: Continuous feedback in MDLMs via soft-masking
  3. Background: Masked Diffusion Language Models
  4. MDLM Modeling
  5. Forward corruption
  6. Reverse process
  7. Reverse process in natural language processing
  8. Training objective
  9. MDLMs in Practice
  10. Unmasking strategies
  11. Conditional generation
  12. Soft-Masked Diffusion Language Models
  13. Soft-Masking (SM)
  14. General formulation
  15. Confidence-based weighting
  16. ...and 39 more sections

Figures (6)

  • Figure 1: Illustrative answer generation using masked diffusion language models (MDLMs) via iterative decoding with (a) standard binary masking or (b) our proposed soft-masking. Our soft-masking enriches the feedback for the next decoding step by superposing the masked tokens with the previously predicted top-$k$ candidates, enabling more accurate and faster generation.
  • Figure 2: Iterative denoising in MDLMs using the proposed soft-masking (SM). Given a context, the aim is to predict the answer via iterative denoising of an initially fully masked response. Here, a bidirectional Transformer ($f_\theta$) performs a single denoising step. This output is passed through an unmasking function that, based on the Transformer's token probabilities, determines which tokens remain masked. Our proposed SM enriches the masked tokens by superposing them with the normalized top-$k$ tokens at each position, weighted by a confidence parameter ($\lambda$).
  • Figure 3: Continuing MDLM pretraining on OpenWebText. We show the average $\pm$ standard deviation (shaded) across 5 seeds. SM is configured with $k=3$. (a) Our SM yields better (lower) validation perplexity than binary masking. (b) The model learns to fully use SM by increasing its influence over the scaling factor $\omega_s$.
  • Figure 4: Integrating SM into Fast-dLLM. We plot Dream-Coder-7B performance vs. throughput with both binary feedback and our SM. SM again excels in high-throughput settings.
  • Figure 5: Ablation study language modeling on OWT. Default SM parameters are $p_{sm}=0.8$ and $k=3$.
  • ...and 1 more figures