Source-linked AI summary
DiffusionBERT: Improving Generative Masked Language Models with Diffusion Models
Zhengfu He, Tianxiang Sun, Kuanning Wang, Xuanjing Huang, Xipeng Qiu
TL;DR
Text generation with diffusion models is difficult in the discrete domain, and prior text diffusion approaches had not combined diffusion with pre-trained language models. DiffusionBERT uses BERT within an absorbing-state discrete diffusion process, adding a token-informed spindle schedule and exploring time-step designs; experiments report improved generation quality over prior diffusion and generative masked-language models.
Problem
Text diffusion is challenging because text is discrete, and prior approaches had not leveraged pre-trained language models.
Method
DiffusionBERT combines a pre-trained language model with absorbing-state discrete diffusion, using a token-informed spindle noise schedule and investigating how to incorporate time steps.
Results
DiffusionBERT significantly improves unconditional text-generation quality over existing text diffusion models and previous generative masked language models in perplexity and BLEU score.
Takeaways & Limitations
Combining diffusion models with pre-trained language models yields a successful non-autoregressive text-generation approach with improved perplexity and generation-quality tradeoffs.
Takeaways & Limitations
The comparison omits SUNDAE because no public implementation was available, with comparison deferred to later work.
Abstract
from arXiv · showhide
We present DiffusionBERT, a new generative masked language model based on discrete diffusion models. Diffusion models and many pre-trained language models have a shared training objective, i.e., denoising, making it possible to combine the two powerful models and enjoy the best of both worlds. On the one hand, diffusion models offer a promising training strategy that helps improve the generation quality. On the other hand, pre-trained denoising language models (e.g., BERT) can be used as a good initialization that accelerates convergence. We explore training BERT to learn the reverse process of a discrete diffusion process with an absorbing state and elucidate several designs to improve it. First, we propose a new noise schedule for the forward diffusion process that controls the degree of noise added at each step based on the information of each token. Second, we investigate several designs of incorporating the time step into BERT. Experiments on unconditional text generation demonstrate that DiffusionBERT achieves significant improvement over existing diffusion models for text (e.g., D3PM and Diffusion-LM) and previous generative masked language models in terms of perplexity and BLEU score.
1 Introduction
DiffusionBERT combines pre-trained denoising language models with discrete diffusion to address challenging text generation. It introduces a token-informed spindle schedule and time-agnostic decoding, achieving improved generation quality over prior approaches.
- Text diffusion remains challenging because text has a discrete nature, and prior approaches had not leveraged pre-trained language models.
- DiffusionBERT combines diffusion models and pre-trained language models through their shared denoising objective.Pre-trained denoising models provide initialization for reverse diffusion, while diffusion supplies an alternative generative training strategy.
- The spindle schedule distributes corrupted information across the forward process using token information and makes it non-Markovian without changing the training objective.Noise for x_t depends on both x_t−1 and the original text x_0.
- Time-agnostic decoding discards time-step information after comparing several ways to incorporate it into the language model.The paper reports that throwing away time information achieved the best result.
- Experiments on unconditional text generation show significantly better generation quality than existing text diffusion models and previous generative masked language models.The reported improvements use perplexity and BLEU score, with ablations confirming the effectiveness of the spindle schedule and time-agnostic decoding.
2 Background
Diffusion models learn generation by progressively corrupting data and then reversing that process. In discrete text diffusion, transition matrices corrupt tokens independently, and a parameterized model learns the reverse denoising process.
- A diffusion model uses a forward process that progressively adds noise and a reverse process that learns to denoise the corrupted sample.The reverse transition can be modeled parametrically when the noise step is sufficiently small.
- The variational objective combines a terminal distribution-matching term, intermediate reverse-process terms, and reconstruction of the original sample.The objective uses expectations over the joint forward-process distribution.
- Given the tractable posterior q(x_t−1|x_t, x_0), a parameterized model p_θ(x_t−1|x_t, t) learns the reverse diffusion process.
- In discrete diffusion, each token is a categorical variable over the vocabulary, and noise is applied through transition matrices independently across sequence positions.The transition matrix specifies q(x_t = j|x_t−1 = i), while cumulative transitions multiply across steps.
3 DiffusionBERT
DiffusionBERT integrates discrete diffusion with pretrained language models through an absorbing [MASK] state, a token-information-based spindle schedule, and alternative time-step designs. The section defines how these choices support reverse diffusion learning and iterative decoding.
- 3.1 Diffusion Models with a Discrete Absorbing State: Discrete diffusion uses an absorbing [MASK] state, with tokens either unchanged or transitioning to [MASK] during the forward process.The process converges to an all-[MASK] sequence, providing a form compatible with BERT.
- 3.1 Diffusion Models with a Discrete Absorbing State: The reverse model learns pθ(x_t−1|x_t, t) and generates text by iteratively reversing the diffusion process.The training objective is derived from the forward and reverse transition formulation.
- 3.2 Spindle Noise Schedule: The spindle schedule assigns noise according to token information, distributing corrupted information across diffusion steps rather than treating all tokens identically.It orders tokens by information and divides them into T buckets, so lower-information tokens emerge earlier during reverse generation.
- 3.2 Spindle Noise Schedule: The spindle schedule depends on the original text, making the forward process non-Markovian without changing the original training objective.Its sinusoidal control term satisfies S(0) = S(T) = 0, and λ = 0 recovers the schedule β_t = (T − t + 1)^−1.
- 3.3 The Design Space of Feeding Time Steps: The section explores layer-wise, prefix, and time-agnostic approaches for supplying or inferring the diffusion time step in pretrained language models.Time-agnostic decoding infers the step from the number of [MASK] tokens and predicts pθ(x_t−1|x_t).
4 Experiments
Experiments evaluate DiffusionBERT for unconditional LM1B generation against discrete and continuous diffusion baselines, BERT-Mouth, and autoregressive references. Results cover generation quality, diversity, convergence, time-step designs, noise schedules, and efficiency.
- Main Results: DiffusionBERT achieves better perplexity and BLEU scores than BERT-Mouth, while one-time sampling of all tokens supports more coherent long-range text.Later diffusion steps can mask and resample low-probability tokens, unlike BERT-Mouth’s fixed sampled tokens.
- Main Results: DiffusionBERT generates higher-quality text than Diffusion-LM, although the comparison also reflects its discrete diffusion process and pretrained initialization.Initializing Diffusion-LM with BERT performs worse than training it from scratch, whereas DiffusionBERT benefits from BERT initialization.
- Main Results: Time-agnostic decoding performs best for DiffusionBERT, while removing time-step embeddings significantly harms D3PM.The results indicate that DiffusionBERT can generate effectively without explicitly providing time information.
- Main Results: The spindle noise schedule improves perplexity by 18% for D3PM and 19% for DiffusionBERT.D3PM with the spindle schedule also outperforms D3PM with the standard schedule.
- Quality-Diversity Trade-off: DiffusionBERT matches a Transformer decoder’s generation ability, improves the non-autoregressive quality-diversity Pareto front, and remains behind pretrained autoregressive models of the same size.Its time-agnostic decoding uses pretrained parameters without adapting new parameters, reducing model and decoding overhead.
- Efficiency of Training and Generation: DiffusionBERT converges faster than D3PM and matches reported performance with only 30% of the training budget.That reduced budget corresponds to 0.5 million training steps.
- Efficiency of Training and Generation: DiffusionBERT offers competitive performance at speeds comparable to GPT and a better efficiency-performance trade-off than BERT-Mouth.Inference cost can be adjusted by controlling reverse-process step size.
5 Related Work
Diffusion models have achieved strong results in continuous domains but remain challenging for text because text is discrete. Prior text approaches use discrete-state diffusion, continuous diffusion with embedding and rounding, or masked-language-model-like iterative refinement.
- Diffusion models have achieved state-of-the-art sample quality in continuous domains such as image and audio generation.
- Text diffusion remains difficult because Gaussian noise cannot be directly applied to discrete data.
- Prior approaches either extend diffusion models to discrete state spaces or perform diffusion continuously before bridging back through embedding and rounding.
- Absorbing discrete diffusion models resemble conditional masked language models by predicting whole sequences and iteratively refining generated text.
- Their main difference is that DiffusionBERT models a stochastic recovery process, whereas conditional masked language models use deterministic prediction.
6 Conclusion
DiffusionBERT combines pretrained language models with absorbing-state discrete diffusion for unconditional non-autoregressive text generation. Its spindle noise schedule and time-step designs improve perplexity and the quality-variance tradeoff.
- DiffusionBERT combines pretrained language models with absorbing-state discrete diffusion models for unconditional non-autoregressive text generation.
- The method introduces new time-step designs and a spindle noise schedule as deviations from current discrete diffusion models.
- The spindle noise schedule assigns each token a schedule according to its frequency in the training corpus.
- DiffusionBERT improves perplexity and substantially advances the Pareto front of quality-variance tradeoffs for non-autoregressive methods.