A Reparameterized Discrete Diffusion Model for Text Generation

Lin Zheng; Jianbo Yuan; Lei Yu; Lingpeng Kong

Source-linked AI summary

A Reparameterized Discrete Diffusion Model for Text Generation

Lin Zheng, Jianbo Yuan, Lei Yu, Lingpeng Kong

arXiv:2302.05737v3Feb 11, 2023cs.CLcs.LG

TL;DR

Text generation with discrete diffusion remains challenging because prior approaches often underperform, while continuous alternatives require rounding and many iterations. The paper derives an equivalent routing-based sampling formulation and builds RDMs with simplified training and decoding. Across text-generation benchmarks, RDMs improve generation quality and can achieve higher-quality samples with substantially fewer iterations, while fixed-length generation and explicit length prediction remain limitations.

Problem
Discrete diffusion models are relatively under-explored and often inferior for text generation, while continuous approaches require rounding to discrete tokens and many sampling iterations.
Method
RDMs reparameterize discrete diffusion sampling as an explicit routing-and-denoising process and parameterize both routing and denoising.
Results
RDMs significantly improve text-generation quality across benchmarks, including about 3~20 BLEU improvements across translation datasets and iteration steps.
Takeaways & Limitations
The reparameterized framework supports effective discrete diffusion generation with flexible decoding and substantially reduced sampling runtime in evaluated tasks.
Takeaways & Limitations
RDMs are currently confined to fixed-length sentences and rely on an explicit length prediction module.

Abstract

from arXiv · show

This work studies discrete diffusion probabilistic models with applications to natural language generation. We derive an alternative yet equivalent formulation of the sampling from discrete diffusion processes and leverage this insight to develop a family of reparameterized discrete diffusion models. The derived generic framework is highly flexible, offers a fresh perspective of the generation process in discrete diffusion models, and features more effective training and decoding techniques. We conduct extensive experiments to evaluate the text generation capability of our model, demonstrating significant improvements over existing diffusion models.

1 Introduction

Discrete diffusion models have been less successful for large-scale text generation because language is discrete, while many diffusion approaches operate continuously and require rounding or many sampling iterations. The paper reparameterizes discrete diffusion sampling into an explicit routing-and-denoising process and introduces RDMs with simplified training and flexible decoding.

Continuous approaches diffuse token embeddings or logits but require a rounding scheme to recover discrete tokens and often many sampling iterations.
Discrete diffusion models remain relatively under-explored and often perform poorly on text generation.
RDMs reinterpret discrete diffusion sampling as routing each token between denoising and noisy states.
The framework reduces training to a re-weighted cross-entropy objective and makes the loss invariant to routing probabilities.
The paper derives a compact formulation, parameterizes routing and denoising explicitly, and evaluates RDMs across several text-generation benchmarks.

2 Background

Discrete diffusion models transform data into noise through a forward process and generate samples by approximating the intractable reverse process. Training uses an ELBO whose timestep terms compare analytical conditional reverse transitions with neural parameterizations that predict the original data.

The forward process gradually transforms one-hot data into a noise distribution through intermediate discrete latent variables.
Multinomial diffusion uses uniform vocabulary noise, whereas absorbing diffusion uses a point mass on an absorbing state.
Generation starts from noise and samples ancestral reverse transitions, but the exact backward transition is generally intractable.
The generative model approximates each reverse transition with pθ(xt−1|xt) and trains it by maximizing an evidence lower bound.
The ELBO decomposes into timestep-wise KL divergences, with the conditional reverse distribution analytically computable for most discrete diffusion models.
For text generation, token positions are typically diffused independently, while a neural network predicts x0 through a softmax probability vector.

3 Reparameterizing the Backward Processes

The paper shows that discrete diffusion's conditional backward transition can be expressed as a route-and-denoise process. Tokens are routed according to whether they are noisy, then either denoised toward x0 or reset to noise, yielding an equivalent but more explicit sampling path.

3.1 An Alternative Backward Formulation: The backward transition admits a compact equivalent formulation for the forward process q(xt|xt−1) = βtxt−1 + (1 − βt)qnoise.
3.1 An Alternative Backward Formulation: The interpolated noise distribution qnoise(xt) combines the current state xt with the underlying noise distribution.
3.1 An Alternative Backward Formulation: When xt = x0, the process either copies the current token or resets it to noise; when xt ≠ x0, it either denoises to x0 or remains noisy.
3.1 An Alternative Backward Formulation: The probabilities of denoising noisy tokens and reverting denoised tokens to noise are controlled by routing weights.
3.2 Reparameterized Sampling: Sampling the mixture can be reparameterized by first selecting a component and then sampling from that component distribution.
3.2 Reparameterized Sampling: The resulting model explicitly represents a routing variable that distinguishes noisy tokens and controls transitions between noisy and ground-truth states.

4 Reparameterized Discrete Diffusion Models

Reparameterized discrete diffusion models make routing explicit, preserving an equivalent sampling process while simplifying training and enabling flexible decoding. The framework supports adaptive token routing and recursive sampling over sequences.

4.1 Joint Diffusion Modeling: RDMs augment discrete diffusion with step-wise routing indicators and retain an equivalent sampling process by marginalizing the routing variables.The joint model separates routing from token transitions, making the latent routing mechanism explicit.
4.2 Training: The sequence-level training objective reduces to a weighted multi-class cross-entropy loss over noisy tokens.This replaces the original objective's KL divergence between complicated categorical distributions and permits loss reweighting through λ.
4.2 Training: The training loss can remain invariant to the routing distribution up to reweighting, allowing one trained model to support a broad family of sampling routes.A model trained with one convenient routing distribution can therefore be reused with other routing distributions indexed by q(v_t−1).
4.3 Sampling: RDM sampling starts from noisy tokens, predicts denoised token states, generates routing indicators, and iteratively updates the sequence until returning x_0,1:N.The sampling algorithm uses a trained network and temperature τ, with all token positions processed in parallel where applicable.
4.3 Sampling: Recursive computation of b_t,n tracks which tokens have already been denoised, avoiding access to the unavailable ground-truth x_0,n during decoding.The resulting frontier-set update adds no extra computational cost.
4.3 Sampling: Instead of assigning every token the same routing probability, RDMs denoise only tokens whose model-confidence scores rank among the k largest.The scores are obtained from the maximum value of the network's token-probability output.

5 Experiments

Experiments across translation, question generation, and paraphrasing show that RDMs improve discrete diffusion quality while substantially reducing decoding cost. Ablations and speed analyses attribute these gains to improved training, discriminative decoding, and the reparameterized iterative process.

Machine Translation: RDMs improve translation quality by about 3–20 BLEU across datasets and iteration steps, outperforming prior non-autoregressive baselines and approaching autoregressive results.The gains apply to both absorbing and multinomial diffusion, with larger improvements for multinomial diffusion.
Question Generation and Paraphrasing: RDMs outperform DiffuSeq and autoregressive baselines on question generation and paraphrasing while producing higher-quality samples with 10 iterations.DiffuSeq requires 2000 decoding steps, whereas RDMs reduce runtime by over 200×.
Analysis: Improved training or decoding each yields over 10–20 BLEU points for multinomial diffusion, with their combination providing further gains.The ablation identifies inadequate training and ineffective decoding as contributors to vanilla multinomial diffusion’s inferior performance.
Analysis: RDMs usually run up to 10× faster than similarly sized autoregressive models while maintaining satisfactory quality with 2 or 5 decoding steps.The comparison uses models with roughly 90–110M parameters, batch size 32, one RTX 3090 GPU, and averages over 30 runs.
Examples: Vanilla multinomial diffusion generates only at its first iteration, whereas reparameterization fixes this degenerated behavior and enables iterative refinement.The resulting behavior leads to better generation quality across the evaluated text-generation tasks.
Examples: Continuous diffusion decodes meaningful tokens slowly, while discrete diffusion produces relevant tokens within only a few steps.DiffuSeq generally does not emit meaningful tokens until around the 1000-th iteration or later.

6 Conclusion

The work analyzes discrete diffusion models and proposes reparameterized diffusion models (RDMs) that improve training and decoding, with boosted text-generation quality across benchmarks. It also identifies fixed-length generation as a current scope limitation.

RDMs are proposed as a family of reparameterized discrete diffusion models based on an extensive analysis of discrete diffusion.
RDMs significantly improve previous work in both training and decoding.
The model family is evaluated on various text-generation benchmarks and demonstrates boosted generation quality.
RDMs are currently confined to fixed-length sentences and rely on an explicit length prediction module.

Appendices

The appendices cover discrete diffusion’s connections to language modeling, its forward and backward formulations, and applications to text generation. They also describe alternative parameterizations and absorbing and multinomial diffusion processes.

Related work: Discrete diffusion connects to autoregressive language models and masked language models through deterministic and absorbing diffusion formulations.
Related work: Research improves discrete-diffusion text generation through unrolling, pretrained models, editing-based refinement, and more effective decoding strategies.
Related work: Continuous-diffusion text generation typically applies Gaussian diffusion over embedding space, while discrete diffusion uses categorical corruption processes.
Model formulations: The standard discrete-diffusion objective approximates q(x_t−1|x_t,x_0) with a parameterized reverse distribution, commonly using a Transformer prediction f(x_t;θ).
Diffusion processes: Absorbing diffusion sends tokens to a persistent mask state, whereas multinomial diffusion corrupts tokens toward a uniform categorical noise distribution.

C Derivation for Proposition 3.1

This appendix derives the discrete-diffusion transition formulas using transition matrices, Bayes’ rule, and case distinctions based on the relationship between noisy and original states.

Backward transition probabilities are derived with Bayes’ rule for cases where the current state matches or differs from the original state.
The derivation represents each forward transition with a probability matrix and composes these matrices across diffusion steps.
The derivation distinguishes the element-wise product of x_t and x_0 from the probability that the noise distribution generates x_t.
The resulting formulation also supports transitions with a time gap by considering q(x_s|x_t,x_0) for s=t−Δt.

D Derivations for the ELBO of RDMs

The appendix derives the RDM evidence lower-bound objective by decomposing its time-step losses into reconstruction and conditional-KL terms. For sequences, the objective can be reduced to a simple token-factorized expression.

The RDM loss uses a reconstruction term at t=1 and a conditional KL-divergence term for later time steps.
The time-step loss is defined over the joint distribution of noisy states and routing variables conditioned on the observed sequence.

E Derivation for Equation 6 and Discussions

The derivation decomposes the training objective into cases based on whether a token is clean or noisy. Under teacher forcing, tokenwise losses combine into a whole-sequence objective, with connections to absorbing diffusion objectives.

Case analysis: The derivation separates each token’s loss into cases defined by the binary state indicator b_t,n.The noisy case uses the noise distribution, while the clean case makes the forward and model backward distributions identical.
Case analysis: When b_t,n = 1, the true and model backward transitions coincide, yielding zero KL divergence.This removes the corresponding contribution from the objective.
Case analysis: When b_t,n = 0, the backward transition is expressed using q_noise(x_t,n), and the possible outcomes are enumerated to obtain the loss contribution.The derivation expands this noisy-token case before combining it with the clean-token case.
Sequence objective: Because tokens are modeled conditionally independently, the per-token losses are summed to produce the objective for the entire sequence.The sequence-level expression follows directly after computing the tokenwise cases.
Connection to prior objectives: The resulting simplified loss resembles absorbing-diffusion cross-entropy over masked positions but is derived from the reparameterization perspective and applies more broadly.The stated scope includes multinomial and absorbing noise processes.

F.1 Tasks

The experiments cover machine translation, question generation, paraphrasing, and related decoding settings using Transformer-based architectures and multiple evaluation protocols. The implementation includes configurable decoding, conditioning, and routing strategies across these tasks.

Machine translation: The machine-translation evaluation uses IWSLT14 DE-EN, WMT14 EN-DE, and WMT16 EN-RO benchmarks without knowledge distillation.The datasets range from approximately 160K to 4.0M training sentence pairs.
Machine translation: Evaluation reports tokenized BLEU, sacreBLEU, and COMET, with sacreBLEU and COMET showing trends similar to tokenized BLEU on WMT14 EN-DE.The additional metrics use the stated sacreBLEU signature and COMET version 2.0.0 configuration.
Other tasks: Question generation and paraphrasing use the Quasar-T and QQP datasets with WordPiece tokenization and a 30522-token vocabulary.The tasks follow preprocessing and data splits described for these benchmarks.
Architectures: The models use bidirectional Transformer self-attention, a length-prediction module, and task-dependent Transformer-base or smaller Transformer architectures.The architecture and hyper-parameter settings are specified for the different benchmark groups.
Training and decoding: Conditioned training consistently improves over vanilla training, while several additional components have smaller effects than reweighted training and adaptive decoding.The stated conditioned-training improvement includes nearly 2 BLEU on the reported setting.
Decoding: RDM decoding uses argmax token selection by default, while translation uses five length candidates and other tasks use MBR decoding with ten samples.The framework also supports stochastic top-k routing and alternative constraints on already denoised tokens.

G Extension: Improved Training with Conditioning

Conditioned training reduces randomness by coupling samples from two diffusion times while preserving an unbiased loss estimate. The paper evaluates this training approach across text-generation tasks and reports comparisons involving sample count and decoding runtime.

Training procedure: Conditioned training samples two time steps s and t, draws x_t normally, and samples x_s conditionally from x_t and x_0.When s = t, the method drops conditioning and samples directly from the marginal process.
Training procedure: The losses at the two time steps are estimated separately and averaged to form the training objective.The procedure remains unbiased while using multiple samples to estimate the objective.
Benefits: Conditioned training applies Rao-blackwellization-like coupling, constrains randomness, and better aligns simulated backward transitions with generation.The method requires only doubling the batch size within one model forward pass and applies to most existing diffusion processes.
Candidate sampling: DiffuSeq benefits slightly more than RDMs from increasing the MBR sample set, suggesting room for RDMs to exploit multiple decoding samples more effectively.The comparison concerns text diffusion models evaluated with different numbers of candidate samples.
Runtime evaluation: Runtime comparisons evaluate CMLM and RDMs on WMT14 EN-DE using the same FairSeq codebase, GPU, batch size, and sixteen iterations.The reported runtime statistics average ten runs over the full test set.
Open-ended generation: Open-ended generation experiments use Wikitext-103 and assess diversity, MAUVE, and coherence across autoregressive, vanilla discrete-diffusion, and RDM models.The compared models have approximately 430M parameters and are trained for 100K steps.

H.5 Extended Qualitative Analysis

The qualitative analysis identifies distinct decoding behaviors across diffusion variants. Reparameterization permits token rollback and avoids degeneracies that make vanilla absorbing or multinomial diffusion stop refining, while continuous diffusion may waste many steps without changing outputs.

Multinomial diffusion: Multinomial diffusion completes most sentences in the first iteration and remains unchanged despite additional allocated iterations.Its later behavior is attributed to the backward transition becoming dominated by copying previous states.
Multinomial diffusion: For large vocabularies, scaling by 1/K suppresses non-copying terms in multinomial diffusion’s backward transition.The Transformer prediction can influence decoding initially, but later transitions become dominated by the first term.
Reparameterization: Reparameterization normalizes separate transition cases independently, avoiding relative-scale suppression and producing more expected decoding behavior.The paper connects this change with better generation quality.
Continuous diffusion: Continuous diffusion often spends many consecutive iterations without modifying the decode, even before convergence, creating potential computation waste.The analysis highlights unchanged outputs across iterations 1000–1009 and contrasts this with discrete diffusion generation in ten steps or fewer.
Absorbing diffusion: Vanilla absorbing diffusion cannot fix previously made errors because decoded tokens remain fixed after their first prediction.Independent token predictions can therefore create unrecoverable duplicate-token errors.
Reparameterized decoding: RDMs can roll decoded tokens back to noisy states when their model scores fall, enabling re-prediction of low-confidence outputs.The example re-masks repetitive “months months” tokens at a later iteration.