Source-linked AI summary
A Cheaper and Better Diffusion Language Model with Soft-Masked Noise
Jiaao Chen, Aston Zhang, Mu Li, Alex Smola, Diyi Yang
TL;DR
Diffusion models face challenges on discrete text, where existing approaches may neglect linguistic structure and struggle with high-dimensional continuous-to-discrete conversion. Masked-Diffuse LM uses linguistically informed soft-masking and direct categorical prediction; the authors report better performance and efficiency than prior diffusion language models across their evaluated tasks.
Problem
Applying diffusion models to text is challenging because text is discrete, while prior methods may neglect word-level linguistic differences and rely on unstable, inefficient high-dimensional rounding.
Source: §1 · pp. 1–2Method
Masked-Diffuse LM corrupts text through linguistically informed soft-masking and predicts discrete tokens directly with cross-entropy during iterative denoising.
Source: §1 · pp. 1–2Results
Across five controllable generation tasks, Masked-Diffuse LM consistently outperforms previous models; it reports a 1.7% improvement on the POS task over Diffusion-LM.
Source: §6.4 · p. 9Takeaways & Limitations
The experiments report better generation performance and more efficient training and inference than prior diffusion language models on the evaluated tasks.
Source: §1 · p. 2Controllable generation follows a conditional independence assumption when decomposing the joint probability into per-time-step control tasks.
Source: §5 · p. 6
Abstract
from arXiv · showhide
Diffusion models that are based on iterative denoising have been recently proposed and leveraged in various generation tasks like image generation. Whereas, as a way inherently built for continuous data, existing diffusion models still have some limitations in modeling discrete data, e.g., languages. For example, the generally used Gaussian noise can not handle the discrete corruption well, and the objectives in continuous spaces fail to be stable for textual data in the diffusion process especially when the dimension is high. To alleviate these issues, we introduce a novel diffusion model for language modeling, Masked-Diffuse LM, with lower training cost and better performances, inspired by linguistic features in languages. Specifically, we design a linguistic-informed forward process which adds corruptions to the text through strategically soft-masking to better noise the textual data. Also, we directly predict the categorical distribution with cross-entropy loss function in every diffusion step to connect the continuous space and discrete space in a more efficient and straightforward way. Through experiments on 5 controlled generation tasks, we demonstrate that our Masked-Diffuse LM can achieve better generation quality than the state-of-the-art diffusion models with better efficiency.
1 Introduction
Text diffusion remains challenging because language is discrete and prior methods often neglect linguistic differences among words or rely on unstable, inefficient high-dimensional rounding. Masked-Diffuse LM uses linguistically informed soft-masking and a cross-entropy objective, and the paper reports stronger performance and efficiency than prior diffusion language models.
1 Introduction
- The paper positions its approach within iterative-denoising diffusion models, previously successful in visual and audio generation.It motivates adapting that generative approach to language despite text’s discrete nature.Source: §1 · p. 1
- Prior text-diffusion methods often ignore differences in word importance, while high-dimensional k-nearest-neighbor rounding can be unstable and inefficient.The introduction identifies a need for language-tailored corruption and a more direct continuous-to-discrete objective.Source: §1 · pp. 1–2
- On E2E and five controllable generation tasks, Masked-Diffuse LM reports state-of-the-art performance and more efficient training and inference than Diffusion-LM.The tasks are Semantic Content, Parts-of-speech, Syntax Tree, Syntax Spans, and Length.Source: §1 · p. 2
- The method soft-masks text according to linguistic features and uses linear layers with cross-entropy objectives to bridge continuous and discrete spaces.The soft-masking order runs from harder or more informative words toward simpler or less informative words.Source: §1 · pp. 1–2
2 Related Work
The related work situates Masked-Diffuse LM among language diffusion, non-autoregressive generation, and plug-and-play control methods. The paper distinguishes its language-tailored corruption and iterative recovery from prior approaches.
2 Related Work
- The paper frames its work as related to diffusion models, non-autoregressive text generation, and plug-and-play controllable generation.These three areas provide the broader context for the method’s design and use.Source: §2 · p. 2
2.1 Diffusion Models for Language
- Language-diffusion approaches use discrete corruption or continuous embeddings with uniform Gaussian noise; the paper instead proposes structured soft-masking and direct token prediction.The cited prior methods include discrete transition kernels and continuous embedding-space models.Source: §2.1 · p. 2
2.2 Non-Autoregressive Text Generation
- Left-to-right autoregressive generation fixes token order and limits editing based on later results, especially for global control; Masked-Diffuse LM recovers latent variables in a planned sequence.Prior non-autoregressive approaches include masked language models, sequence alignment, insertion and deletion, and unrolled generation paths.Source: §2.2 · pp. 2–3, Figure 1 · p. 3
2.3 Plug-and-Play Controllable Generation
- Plug-and-play methods guide generation with classifiers or other signals without changing or fine-tuning the pretrained language model; this work applies classifier-guided gradients to intermediate latent variables.Earlier methods edit hidden representations, reweight predicted tokens, or fine-tune a smaller language model for token reweighting.Source: §2.3 · p. 3
3 Background: Diffusion Models
Diffusion models gradually corrupt data through a forward process, then learn a reverse process to denoise it step by step. The background also describes variational-likelihood training and a reweighted mean-squared-error objective.
3 Background: Diffusion Models
- A forward Markov chain progressively adds Gaussian noise to sampled input data, with a schedule controlling the amount added at each step and no trainable forward-process parameters.The resulting terminal latent is described as an isotropic Gaussian distribution.Source: §3 · pp. 3–4
- A parameterized reverse transition denoises the terminal latent step by step, with each transition modeled as a Gaussian whose mean and covariance are learned.The learned transition can be implemented with a U-Net or Transformer.Source: §3 · p. 4
- Training targets the marginal likelihood through a variational lower bound, but the objective is described as unstable, motivating a reweighted mean-squared-error loss.The passage identifies this reweighting as following Ho et al. [2020].Source: §3 · p. 4
4 Method: the Masked-Diffuse LM
Masked-Diffuse LM combines linguistically informed soft-masking with stepwise mapping from continuous states to discrete predictions. Its design also supports integration of pretrained language models.
4.1 Embedding
- Input tokens are first mapped into a continuous initial state through a learnable embedding layer or encoder, bridging discrete and continuous spaces.Source: §4.1 · p. 4
4.2 Forward Process with Soft-Masking
- The forward process ranks words using normalized tf–idf relevance and entropy, then soft-masks higher-importance words earlier to create structured corruption.Importance-ranked buckets guide the masking order.Source: §4.2 · pp. 4–5
- A square-root schedule gradually increases β_t, while s sets the starting noise level and less noise is added to harder words to stabilize training.Source: §4.2 · p. 5
4.3 Diffusion Process
- At each step, the model maps continuous states to discrete predictions and minimizes weighted cross-entropy against original and masked sentences, prioritizing currently masked tokens.Source: §4.3 · p. 5
4.4 Adapting Pre-trained Language Models
- The model can replace its embedding layer with a pretrained language model; experiments freeze that model and train only the transition parameters.Source: §4.4 · p. 6
5 Controllable Text Generation with Masked-Diffuse LM
Masked-Diffuse LM applies control guidance to latent variables at every diffusion step and combines it with the reverse transition. Gradient updates balance fluency and control, while MBR decoding selects among generated samples.
5 Controllable Text Generation with Masked-Diffuse LM
- External controls are applied directly to latent variables at every intermediate diffusion step, with the joint transition decomposed into the reverse transition and a control likelihood.The decomposition follows a conditional-independence assumption.Source: §5 · p. 6
- Gradient updates use differentiable transition and classifier terms, while λ balances fluency from Masked-Diffuse LM against control from classifiers.The extra classifiers take diffusion latent variables as input, allowing direct gradient updates in latent space.Source: §5 · pp. 6–7
- Minimum Bayes Risk decoding selects the sample with the lowest expected loss under the specified loss function.MBR aggregates candidate samples from Masked-Diffuse LM.Source: §5 · p. 7
6 Experiments
Across five controllable generation tasks, Masked-Diffuse LM outperforms prior models, including when combined with BERT; it also uses less training and inference time than Diffusion-LM. Human evaluation and ablations further support the reported quality and design choices.
6.1 Datasets
- The E2E dataset contains 50K restaurant reviews with food-type, price, and customer-rating labels; controls cover semantic content, POS, syntax trees, syntax spans, and length.Evaluation uses exact-match rates for content, POS, spans, and length, and F1 for syntax-tree matching.Source: §6.1 · p. 7
- Each task uses 200 validation control targets and 50 generated samples per target; the first four tasks use classifier guidance, while length control is classifier-free.A fine-tuned GPT-2 teacher LM evaluates generated-text perplexity as a measure of quality and fluency.Source: §6.1 · pp. 7–8
6.2 Baselines
- Baselines include PPLM and FUDGE, which guide pretrained language models, and Diffusion-LM, which uses Gaussian diffusion over embedded text followed by rounding.Diffusion-LM also incorporates control signals at every diffusion step.Source: §6.2 · p. 8
6.3 Experimental Setting
- Experiments use an 80M-parameter Transformer for 20,000 training steps, with 4 A100 GPUs for training and one GPU for sampling.The setup uses sequence length 64, 500 diffusion steps, and BERT-base with frozen parameters when incorporating a pretrained model.Source: §6.3 · p. 8
6.4 Results
- Masked-Diffuse LM outperforms previous models across all five controllable generation tasks, including with BERT; human evaluation also gives it and its BERT variation the best average ranks.The paper reports a 1.7% improvement on the POS task and significantly less training and inference time than Diffusion-LM.Source: §6.4 · p. 9
6.5 Ablation Studies
- On Semantic Content, the entropy-and-Reluency noise strategy performs significantly better than the tested Gaussian, random-mask, and POS-based strategies.The authors interpret the result as evidence that considering linguistic features provides more appropriate perturbation for textual diffusion.Source: §6.5 · p. 9
- Cross-entropy performs slightly better than L2 without a pretrained model, while CE-BERT significantly outperforms L2-BERT; Table 5 reports the objective comparisons.The authors attribute the BERT result to the effectiveness of their objectives for incorporating large language models.Source: Table 5 · p. 8, §6.5 · p. 9
6.6 Case Studies
- Intermediate examples show easy words generated first, followed by more informative words such as “Mill” and “Indian” as denoising proceeds.The examples cover Length and Semantic Content tasks.Source: Table 6 · p. 8, §6.6 · p. 9
7 Conclusion
Masked-Diffuse LM combines linguistic-informed soft-masking with iterative denoising through direct discrete-token prediction and cross-entropy. On E2E and five controllable generation tasks, it achieves state-of-the-art performance and more efficient training and inference than Diffusion-LM.
7 Conclusion
- Masked-Diffuse LM corrupts discrete text through linguistic-informed soft-masking and iteratively denoises it by directly predicting text.The forward process gradually soft-masks more informative words before less informative words.Source: §7 · p. 10
- Cross-entropy loss directly predicts discrete tokens during diffusion, which the authors say stabilizes intermediate steps and makes the approach orthogonal to large pretrained models.The denoising order encourages easy-first generation.Source: §7 · p. 10
- On E2E and five controllable generation tasks, Masked-Diffuse LM achieves state-of-the-art performance and more efficient training and inference than Diffusion-LM.The tasks are Semantic Content, Parts-of-speech, Syntax Tree, Syntax Spans, and Length.Source: §7 · p. 10