Source-linked AI summary
Continuous diffusion for categorical data
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H. Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, Curtis Hawthorne, Rémi Leblond, Will Grathwohl, Jonas Adler
TL;DR
Language is categorical, making standard continuous diffusion unsuitable despite its benefits for generative modelling. CDCD embeds tokens in Euclidean space and combines continuous-time diffusion with score interpolation and time warping; it supports non-autoregressive language modelling and flexible generation, while retaining important efficiency trade-offs relative to autoregressive models.
Problem
Standard diffusion models are poorly suited to discrete categorical language, while diffusion-inspired alternatives depart from capabilities such as continuous input processing and classifier-free guidance.
Method
CDCD embeds categorical tokens in Euclidean space and jointly learns embeddings with a continuous diffusion model using score interpolation, cross-entropy training, and time warping.
Results
CDCD provides a framework for continuous diffusion of categorical data that enables non-autoregressive language modelling and flexible sampling, including arbitrary infilling.
Takeaways & Limitations
Continuous diffusion can represent token-level uncertainty during sampling and retain diffusion capabilities such as flexible sampling and classifier-free guidance.
Takeaways & Limitations
Compared with autoregressive models, diffusion can be less data-efficient and more costly because training uses one noise level per example and sampling requires full-sequence passes.
Abstract
from arXiv · showhide
Diffusion models have quickly become the go-to paradigm for generative modelling of perceptual signals (such as images and sound) through iterative refinement. Their success hinges on the fact that the underlying physical phenomena are continuous. For inherently discrete and categorical data such as language, various diffusion-inspired alternatives have been proposed. However, the continuous nature of diffusion models conveys many benefits, and in this work we endeavour to preserve it. We propose CDCD, a framework for modelling categorical data with diffusion models that are continuous both in time and input space. We demonstrate its efficacy on several language modelling tasks.
1. Introduction
CDCD preserves continuous diffusion for categorical language data by embedding tokens in Euclidean space and combining score interpolation, cross-entropy training, and time warping.
- Diffusion-based language models have had limited success because standard diffusion is poorly suited to discrete categorical text representations.
- CDCD models categorical data with diffusion that is continuous in both time and input space by embedding discrete tokens in Euclidean space.
- Score interpolation replaces score matching, enabling cross-entropy training and joint end-to-end learning of the diffusion model and Euclidean embeddings.
- Time warping adapts the distribution of noise levels sampled during training to improve training efficiency.
- The framework applies continuous diffusion to language modelling and machine translation while following a procedure familiar from masked language modelling.
2. Diffusion models
Diffusion generates data through iterative denoising, but applying its continuous formalism to categorical inputs requires embeddings and introduces trade-offs against autoregressive language models.
- Diffusion models learn to reverse gradual corruption step by step, decomposing generation into smaller denoising problems.
- 2.1. Formalism: The diffusion formalism describes forward and reverse corruption with stochastic differential equations and can also use a probability-flow ODE with matching marginals.
- 2.2. Diffusion for discrete data: For discrete inputs, the score function is undefined, so CDCD embeds inputs in continuous space to retain diffusion capabilities such as classifier-free guidance.
- 2.2. Diffusion for discrete data: Continuous embeddings can represent superpositions of possible token outcomes at intermediate sampling times, delaying commitment until the end.
- 2.3. Diffusion and autoregression: Diffusion supports flexible iterative refinement and arbitrary connectivity, whereas autoregressive models offer greater training and sampling efficiency through sequential factorization and activation caching.
- 2.3. Diffusion and autoregression: Diffusion models may be less data-efficient and more costly per sampling step because each training example uses one noise level and each step requires a full sequence pass.
3. The CDCD framework
CDCD adapts continuous diffusion to categorical data by jointly learning Euclidean token embeddings and a diffusion model trained with cross-entropy through score interpolation. Time warping reallocates training toward useful noise levels, while normalization prevents embedding growth.
- Overview: CDCD jointly learns Euclidean token embeddings and the diffusion model, enabling continuous diffusion for categorical data.The framework trains both components end-to-end with a shared loss.
- Score interpolation: Score interpolation estimates diffusion scores by linearly combining the vocabulary’s possible conditional scores using predicted token probabilities.Because the conditional score has one possible value per vocabulary token, probability-weighted interpolation yields a score estimate.
- Score interpolation: Cross-entropy training provides stable optimization and allows end-to-end learning, but changes the relative weighting assigned to noise levels.This motivates adapting timestep sampling rather than relying on the unmodified weighting induced by the loss.
- Diffusion on embeddings: L2-normalizing embedding vectors prevents the uncontrollable parameter growth caused by jointly training embeddings with cross-entropy.The underlying embedding parameters remain trainable, with gradients propagated through normalization.
- Time warping: Time warping adaptively changes the sampled noise-level distribution to focus training on levels where it is most useful.The procedure uses inverse-CDF sampling and is equivalent to time reweighting in expectation while being more statistically efficient.
- Time warping: Time warping typically concentrates weight on intermediate noise levels because very low noise is trivial and very high noise favors marginal predictions.The stated heuristic seeks entropy that increases linearly with the warped time variable.
4. Diffusion language models
CDCD supports conditional language generation by denoising selected token positions with a mask-conditioned Transformer. Its setup handles prefix completion, arbitrary infilling, self-conditioning, and translation without decoder-side causal masking.
- Mask-conditional Transformer: A boolean conditioning mask marks sequence positions as clean inputs or noisy tokens to generate.The model receives noisy embeddings, conditioning embeddings, and the mask as stacked inputs.
- Mask-conditional Transformer: CDCD reduces language modeling to denoising classification, allowing a Transformer without attention masking or causal connectivity restrictions.This supports non-autoregressive refinement across sequence positions.
- Masking strategies: Prefix masks support completion, while fully random masks support conditioning on arbitrary subsets of sequence positions.Training on an equal mixture of prefix and fully random masks slightly improves prefix-completion performance.
- Noise level conditioning: The Transformer receives the timestep through conditional normalization so its computations adapt to the noise level.Each layer normalization is followed by timestep-dependent shifting and scaling.
- Self-conditioning: Self-conditioning adds embeddings interpolated from previous token predictions, increasing training time by 10–15% while yielding significant performance gains.The previous predictions are supplied as an additional stacked input sequence.
- Machine translation: The translation model uses separate encoder and decoder Transformer stacks and predicts complete fixed-length sequences including BOS, EOS, and padding.At sampling time, tokens after the first EOS are discarded.
5. Related work
Prior work explores discrete corruption, embedding-based continuous diffusion, continuous-time models, and iterative refinement for language and translation. CDCD distinguishes itself by combining continuity in both input space and time while targeting language modeling.
- Discrete diffusion: Discrete diffusion and diffusion-inspired language models provide non-autoregressive iterative refinement but forgo some capabilities associated with continuous diffusion.Replacing continuous diffusion with discrete corruption offers flexibility while changing the modeling framework.
- Continuous diffusion for discrete data: Embedding-based approaches apply continuous diffusion to discrete inputs, but several use discrete-time diffusion rather than continuity in both input space and time.CDCD is described as continuous in both dimensions.
- Machine translation: Prior machine-translation work addresses non-autoregressive multimodality through distillation, latent transformers, repeated decoding, editing, and verification strategies.These methods target incoherent parallel sampling decisions such as repeated tokens.
- Time warping: Parameterizing the noise schedule has been proposed previously, but time warping instead aims to linearize prediction entropy rather than minimize diffusion-loss variance.The comparison distinguishes the objectives of the two approaches.
6. Experiments
CDCD experiments examine architectural choices, sampling strategies, prompt completion, infilling, and translation. The results show several useful design effects, favourable language-generation metrics under selected settings, but weaker translation performance and recurring token defects.
- Evaluation: CDCD models are evaluated on prompt completion, infilling, and machine translation against autoregressive counterparts.Mask-conditional models support prompt completion and infilling, while encoder-decoder models are used for translation.
- Design decisions: Renormalisation and clamping significantly reduce AR-NLL but also substantially reduce entropy.Models without renormalisation produce better samples after sampling hyperparameters are tuned.
- Design decisions: AR-NLL initially improves with embedding dimensionality but stabilises beyond w = 64 through time warping’s focus on relevant noise levels.The experiments use w = 256 and t_min = 1.0 elsewhere.
- Design decisions: Fully random masking improves results even when evaluation uses prefix masks, so other experiments use a 50-50 masking strategy.The authors suspect bidirectional context makes embedding learning easier.
- Design decisions: Self-conditioning significantly improves AR-NLL while leaving unigram entropy roughly unchanged.The method reuses computation from preceding sampling steps.
- Sampling: Score-temperature and initial-noise-scale changes provide a better trade-off than manipulating p(x_0|x,t), whereas high guidance scales tend to produce repeated phrases.Score temperature and classifier-free guidance have complementary effects to a degree.
- Sampling: Time warping at sampling improves over the heuristic alone, while the effects of combining them do not compound.The authors interpret this as evidence that time-warped step spacing is close to optimal because uncertainty decreases approximately constantly.
- Prompt completion and infilling: CDCD produces favourable MAUVE scores for several score-temperature and guidance-scale settings, providing evidence of compelling samples.The reported metrics also include AR-NLL and unigram entropy for generated sequences, but MAUVE should be interpreted cautiously because repetition may be penalised unevenly.
7. Discussion
CDCD enables non-autoregressive diffusion language models with arbitrary infilling and flexible sampling, while retaining continuous time and input space. The discussion highlights broader applicability, architectural flexibility, and unresolved limitations around variable-length outputs and sampling efficiency.
- CDCD enables non-autoregressive language modelling with arbitrary infilling and a procedure reminiscent of BERT.
- Its flexible sampling procedure trades off sample quality and computational cost, while iterative refinement permits changing the number of sampling steps without retraining.
- Variable-length output remains an unresolved research question because diffusion models require a fixed-size canvas during sampling.Random padding during training might allow some length variation, but the discussion does not establish a complete solution.
- The authors identify improved diffusion sampling algorithms as an open direction for determining whether CDCD’s sampling efficiency can be improved.
- CDCD’s components are not language-specific and are expected to apply to generative modelling problems involving structured categorical data.The authors hypothesise that jointly fitting embeddings with the diffusion model could improve latent diffusion models, especially with time warping.
A.1. Parameterisation
The time-warping parameterisation represents a monotone CDF with piecewise-linear bins, making it easy to normalise, invert, and fit to expected training loss. Importance weighting compensates for the feedback introduced by adapting timestep sampling.
- Time warping fits a monotonically increasing function to expected loss across timesteps so noise levels can be sampled efficiently during training.The function must be easy to normalise and invert for inverse transform sampling.
- The timestep range is shifted and scaled to [0, 1], allowing the normalised CDF to be parameterised over a unit interval.
- Softmax-normalised input and output logits define partitions of the unit interval whose bin edges are obtained from cumulative sums.
- Within each bin, linear interpolation produces a monotone piecewise-linear CDF, while nonzero minimum bin sizes improve numerical stability.
- Using exponentials instead of softmax yields an unnormalised CDF whose bin edges can fit positive loss values directly.The normalised CDF is recovered afterward by applying softmax to the output logits.
- The learned CDF is straightforward to invert by swapping input and output logits, and its derivative is piecewise constant within each bin.
- Importance weights correct the bias created when adaptive timestep sampling oversamples certain regions during training.Gradients through the warping parameters are stopped to limit changes in the training distribution, although this was not empirically necessary.
A.5. Warping sampling timesteps
Warping timestep spacing improves sampling when applied during both training and sampling. The piecewise-linear CDF also supports temperature and uniform-mixture manipulations of the sampling distribution.
- Warping both the training noise-level distribution and sampling timestep spacing yields the best results.Linearly spaced uniform-time steps are warped afterward and perform better than uniformly spaced timesteps.
- The authors hypothesise that decreasing entropy at a constant rate from step to step is a useful sampling heuristic.
- Changing output bin sizes changes the temperature of the distribution represented by the CDF.This follows from the piecewise-constant probability density determined by the ratio of output to input bin sizes.
- The parameterisation can also derive the CDF of a mixture between the learned distribution and a uniform distribution.
- Equal input and output bin sizes produce the identity CDF corresponding to a uniform distribution on the unit interval.The effects of these distribution manipulations are visualised in Figure 8.
B. Negative results
The paper reports failed attempts to improve results as context-dependent evidence intended to clarify design choices rather than discourage further investigation.
- The authors present failed improvement attempts to clarify design choices and support future investigation, while stressing that outcomes depend strongly on experimental context and details.
B.1. Constraining the embedding parameters
The authors compared normalization with alternative regularization strategies for controlling embedding parameters and found normalization most effective.
- Embedding parameters can be controlled either by normalization or by adding a regularization loss term.
- The experiments included L2 regularization, margin penalties, and automatic adaptation of the L2 penalty weight.
- Normalization worked best among the tested approaches.
B.2. Removing the time dependency
Removing the time-embedding MLP to make the architecture more BERT-like significantly reduced performance, despite the hypothesis that noisy inputs could reveal the timestep.
- The authors removed the time-embedding MLP to simplify the architecture and make it more closely resemble BERT.
- They hypothesized that the classifier could infer the noise level directly from noisy input vectors.
- Removing explicit timestep information significantly hurt performance.
B.3. Simplex diffusion
The authors explored diffusion on the simplex of categorical distributions, but its uneven corruption process impeded language modelling. Score interpolation and time warping proved more effective with Gaussian diffusion in Euclidean embedding space.
- Simplex diffusion: Discrete token sequences were lifted into categorical distributions, whose nonnegative components sum to 1 and therefore lie on the simplex.
- Simplex diffusion: The simplex diffusion formulation was based on the tractable Cox-Ingersoll-Ross process.
- Simplex diffusion: High vocabulary dimensionality made the corruption process uneven, producing heavy-tailed noise and frequent outliers.
- Simplex diffusion: These outliers could make corrupted vectors appear to represent incorrect tokens even at very low noise levels.
- Simplex diffusion: Score interpolation and time warping were originally developed for simplex diffusion.
- Simplex diffusion: Both ideas were more effective when combined with Gaussian diffusion in a Euclidean embedding space.