Continuous diffusion for categorical data

Introduction

Generative models have seen a rapid increase in scale and capabilities over the past few years, across many modalities, including images, audio signals, video and text (Brown et al., 2020; Dhariwal et al., 2020; Borsos et al., 2022; Ramesh et al., 2022; Saharia et al., 2022b; Ho et al., 2022a). In language modelling, the focus has been on scaling up and expanding the capabilities of autoregressive models, instigated by the development of the Transformer architecture (Vaswani et al., 2017). This has resulted in general-purpose language models that are suitable for practical use.

Until recently, work on visual modalities lagged behind in terms of scale and practicability, but the development of diffusion models (Sohl-Dickstein et al., 2015; Song and Ermon, 2019; Ho et al., 2020) has resulted in a noticeable step change in capabilities. Whereas previous generative models of images were relatively inflexible and tended to produce low-resolution outputs, modern text-conditional image generators such as DALL-E 2 (Ramesh et al., 2022) and Imagen (Saharia et al., 2022b) are able to produce high-resolution outputs for any conceivable textual prompt. While this trend cannot be attributed exclusively to the advent of diffusion models (models with similar capabilities that are not based on diffusion do exist, e.g. Parti (Yu et al., 2022)), this new paradigm for generative modelling through iterative refinement has indisputably played a key role in the ‘mainstreaming’ of generative models of images.

Diffusion-based language models have seen relatively little success so far. This is in part due to the discrete categorical nature of textual representations of language, which standard diffusion models are ill-equipped to deal with. As a result, several diffusion-inspired approaches to language modelling have recently been proposed (Austin et al., 2021; Hoogeboom et al., 2021b, a; Savinov et al., 2021; Reid et al., 2022), but these depart from the diffusion modelling framework used for perceptual data in several important ways (with a few exceptions, e.g. Li et al. (2022); Strudel et al. (2022)). This usually implies having to give up some of the unique capabilities of this model class, such as the ability to use classifier-free guidance to enhance conditional generation (Ho and Salimans, 2022), which has been instrumental to the success of diffusion-based text-conditional image generators.

In this paper, we study the suitability of continuous diffusion as a generative modelling paradigm for discrete categorical data, and for textual representations of language in particular. We develop a framework, Continuous diffusion for categorical data (CDCD), based on the diffusion framework proposed by Karras et al. (2022), which enables efficient and straightforward training of diffusion-based language models that are continuous in both time and input space, by embedding discrete tokens in Euclidean space.

Our approach very closely mirrors the training procedure for masked language models such as BERT (Devlin et al., 2018, which are non-autoregressive and non-generativeAlthough several works have explored generative approaches based on masked language models (Wang and Cho, 2019; Ghazvininejad et al., 2019; Goyal et al., 2021; Shih et al., 2022), they were originally introduced and are still mainly used for representation learning.), and hence should appear familiar to language modelling practitioners. We hope that this will help lower the barrier to entry, and encourage researchers to explore continuous diffusion models for other domains for which categorical representations are best suited.

We propose score interpolation as an alternative to score matching for diffusion model training. This allows us to use the familiar cross-entropy loss function for training, which in turn enables end-to-end training of the diffusion model and the Euclidean embeddings with a single loss function;

We introduce time warping, an active learning strategy which automatically adapts the distribution of noise levels sampled during training to maximise efficiency;

We describe CDCD, a framework for continuous diffusion models of categorical data (see Figure 1), and explore its application to language modelling and machine translation.

Diffusion models

Diffusion models enable generative modelling via iterative denoising. Given a gradual corruption process which turns the data distribution into a simple distribution that is easy to sample from (usually an isotropic Gaussian distribution), we can train a model that learns to revert this process step by step. Each step in the reverse direction attempts to reconstruct a small amount of information that the corruption process removed. This is a much easier task than learning to generate data in a single forward pass through a model, as variational autoencoders (VAEs, Kingma and Welling, 2013; Rezende et al., 2014) and generative adversarial networks (GANs, Goodfellow et al., 2014) do. Autoregressive models similarly enable decomposition of the generative modelling problem into smaller subproblems that are easier to solve. Both of these approaches to iterative refinement are compared in §2.3.

Many different formalisms have been proposed for diffusion models, e.g. based on score matching (Song and Ermon, 2019) or latent variable models (Ho et al., 2020). In this work, we will follow Song et al. (2020) and use differential equations to describe the corruption process, as well as the reverse process. We believe that all these different perspectives are largely interchangeable and complementary to some degree.

Song et al. (2020) suggest modelling a diffusion process with the following stochastic differential equation (SDE):

where $\mathbf{w}$ is the standard Wiener process, $\mathbf{f}$ is the (vector-valued) drift coefficient, $g$ is the diffusion coefficient and time $t$ ranges from (clean data) to $T$ (fully corrupted). The reverse process can then be described by the following SDE:

where $\mathbf{\bar{w}}$ is the standard Wiener process in reversed time. $\mathbf{s}(\mathbf{x},t):=\nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})$ is the so-called score function, i.e. the gradient of the density of $\mathbf{x}$ at time $t$. We can train a model to predict this quantity given $\mathbf{x}$ and $t$ using score matching (Hyvärinen and Dayan, 2005):

The estimate $\hat{\mathbf{s}}(\mathbf{x},t)$ can then be plugged into this SDE to produce samplesWe use denoising score matching in practice (Vincent, 2011)..

It turns out that we can instead describe the evolution of $\mathbf{x}$ over time deterministically with an ordinary differential equation (ODE):

This is the probability flow ODE, which has the same marginals $p_{t}(\mathbf{x})$ as the forward SDE at all timesteps $t$. This equivalence is quite powerful, because it enables us to deterministically map data examples $\mathbf{x}_{0}$ to latent representations $\mathbf{x}_{T}$, and vice versa (with $\mathbf{x}_{T}$ approximately following a Gaussian distribution).

Karras et al. (2022) thoroughly explored the design space of diffusion models based on the probability flow ODE formulation, and we will largely follow their recommendations here. Concretely, we will choose $\mathbf{f}(\mathbf{x},t)=0$ and $g(t)=\sqrt{2t}$, which yields:

In this formulation, $t$ corresponds directly to the standard deviation of the Gaussian noise that is added to $\mathbf{x}_{0}$ to simulate samples from $p_{t}(\mathbf{x})$ (and therefore they refer to $t$ as $\sigma$ insteadSee Appendix B.1 of Karras et al. (2022).).

2 Diffusion for discrete data

When $\mathbf{x}$ is discrete, the score function is undefined. This can be worked around in two ways: we can try to define a similar iterative refinement procedure through denoising for discrete data, or we can embed $\mathbf{x}$ into a continuous space and apply continuous diffusion to the embeddings. While most of the literature has focused on the former approach (see §5.1 for an overview), in this work we will explore the latter – abandoning continuity of the input usually means that we have to forgo a lot of useful capabilities, such as classifier-free guidance, which we would like to keep.

Another potential advantage of lifting the discrete input into a continuous space is that it becomes possible to represent superpositions of possible outcomes at intermediate timesteps of the sampling process. In language modelling, this means that we can represent uncertainty at the individual token level: the sampling procedure only commits to specific tokens at the very end. Denoising models that operate directly in the discrete input space do not have this ability: they are only able to represent specific tokens, or the absence of a decision (through use of a ‘mask’ token). This requirement to commit early to some subset of tokens can lead to inconsistencies in the resulting samples, which are difficult to correct retroactivelySome approaches attempt to mitigate this by allowing some proportion of tokens to be resampled multiple times (Savinov et al., 2021)..

Changing the input representation to be continuous enables us to use the standard diffusion framework that has been exceptionally successful for perceptual modalities, but we should not necessarily expect it to work as well for language modelling out of the box. In fact, the most commonly used modelling setup for images implicitly reduces the loss weighting of high frequency content relative to likelihood-based models, allowing for a more efficient use of model capacity which is well aligned with human perception (Song et al., 2021a). Furthermore, the underlying physical phenomena that are being modelled (e.g. light intensity, air pressure) are inherently continuous. The relative ease with which diffusion models of images have been scaled to high resolution inputs can at least partially be attributed to these facts. We cannot expect to benefit from this for language modelling, as the notion of ‘high frequency content’ is not meaningful in this settingAt least, not in the traditional sense; Tamkin et al. (2020) suggest an approach to obtain and analyse multi-scale representations of language.. Finally, we also need to consider the impact of the choice of embedding procedure on generative modelling performance.

3 Diffusion and autoregression

Autoregressive (AR) models currently dominate language modelling research at scale. They factorise the joint distribution over a token sequence $p(x_{1},x_{2},...,x_{N})$ into sequential conditionals $p(x_{k}|x_{1},...,x_{k-1})$ and model each of them separately (with shared parameters). This means sampling always proceeds along the direction of the sequence (i.e. from left to right, when modelling text in English), and in this case, sampling an additional token constitutes an ‘iterative refinement’ step. Autoregression is a very natural fit for language, because it is best represented as a one-dimensional sequence of tokens. That said, the way humans tend to produce language, especially in written form, is far from linear. For many tasks, the ability to go back and refine earlier parts of the sequence, or to construct it hierarchically, is useful.

Changing the modelling paradigm to a more flexible form of iterative refinement (such as diffusion) is desirable, because the increased flexibility would facilitate new applications and could potentially reduce the computational cost of sampling. However, this is a challenging prospect because of the statistical efficiency of AR model training. Because the same parameters can be used to model all sequential conditionals, each training example provides a useful gradient signal for every step of the iterative refinement procedure. This is not the case for diffusion models, where we can only train on a single noise level for each training example. As a result, diffusion models are likely to be less data-efficient, and will converge more slowly. AR models are also able to benefit from caching of previous model activations at sampling time, which significantly reduces the computational cost of each step. Diffusion models require a full forward pass across the entire sequence at each step, which can be much more costly.

Nonetheless, this apparent efficiency benefit of AR models does impose a rather strict constraint on the connectivity patterns within these models – specifically, causality with respect to the input sequence. This constraint is usually implemented using some form of masking, which implies that a significant amount of computation is wasted during training. It also complicates the use of multiresolution architectures, which is very common in other domains of machine learning (e.g. computer vision). Diffusion models on the other hand are completely architecturally unrestricted, so the use of multiresolution architectures (or more exotic variants) is straightforward.

This architectural flexibility compounds with the adaptivity of the denoising procedure, which enables trading off the computational cost and sample quality at sampling time by choosing the appropriate number of iterative refinement steps, without requiring retraining or finetuning. Conversely, for AR models, the number of steps is necessarily the same as the length of the sequence to be generatedStrictly speaking, it is possible to decouple the cost of sampling from the sequence length even for autoregressive models, using probability density distillation (van den Oord et al., 2018) or alternative sampling algorithms (Song et al., 2021b; Jayaram and Thickstun, 2021), but these approaches have not been used for language models, to the best of our knowledge.. More sophisticated sampling procedures for diffusion models are also being developed on a regular basis, which can be applied to existing models without any changes. Therefore, we believe that diffusion models for language are a worthwhile pursuit, despite their relatively reduced data efficiency.

The CDCD framework

We will first describe how we can train diffusion models using the familiar categorical cross-entropy loss with score interpolation. We then show how to map categorical inputs to continuous embeddings in a way that is amenable to diffusion, which we achieve by jointly learning the embeddings and the diffusion model itself, allowing them to co-adapt. Finally, we will discuss time warping, an active learning strategy which automatically adapts the distribution of noise levels sampled during training. Together, these components constitute a framework for continuous diffusion of categorical data, or CDCD, which is summarised in a diagram in Figure 1.

Diffusion models are typically trained by minimising the score matching objective (Equation 3), where the model learns to approximate the score function $\mathbf{s}(\mathbf{x},t)$ in the least-squares sense. The model predictions can then be substituted directly into Equations 2 or 4 for sampling.

We observe that when the data is discrete and categorical, with tokens taken from a vocabulary of size $V$, the conditional score function $\mathbf{s}(\mathbf{x},t|\mathbf{x}_{0})$ can only assume $V$ possible values. Therefore, if we have a probabilistic prediction of $\mathbf{x}_{0}$, we can use it to linearly interpolate the $V$ possible values to obtain a score function estimate:

To obtain an estimate of $p(\mathbf{x}_{0}|\mathbf{x},t)$, we can make our model predict $V$ logits and apply a softmax nonlinearity, and minimise the categorical cross-entropy loss. This is also the standard setup used to train autoregressive language models (as well as classifiers in general), so it is well-studied and understood, and it ensures stability during training. Compared to the score matching loss, the cross-entropy loss will of course weight errors in the score function estimates differently relative to each other. This difference is important in practice, because we are only able to optimise the loss approximately (i.e. the global optimum is unlikely to be reached), and the relative weighting of the noise levels will also be different.

The conditional score function corresponding to the ODE in Equation 5 is given by:

2 Diffusion on embeddings

To embed the input in a continuous space, we could arbitrarily assign embeddings to different tokens, or use a representation learning technique to obtain embeddings (Strudel et al., 2022). However, since we are able to backpropagate gradients from the diffusion model into the embeddings using the reparameterisation trick (Kingma and Welling, 2013; Rezende et al., 2014), we explore learning the embeddings and the diffusion model jointly. This yields a simpler setup, with a single shared loss function for all model parameters.

If we were to train our diffusion model with score matching, joint training would result in collapse of the embedding space. Since the model is effectively predicting the noise that is added to the embeddings, this task becomes trivial when all embeddings correspond to the same vector. This minimises the loss function, but it does not yield a useful model. Therefore, additional loss terms are necessary to prevent collapse (Li et al., 2022).

Using score interpolation, we can train the diffusion model with the cross-entropy loss instead. Since the objective is now to distinguish the true embedding from all other embeddings, given a noisy embedding as input, the model is encouraged to push the embeddings as far apart as possible (as this minimises the confounding impact of the noise). We now have the opposite problem, where joint training leads to uncontrollable growth of the embedding parameters, unless they are constrained in some way.

We could again implement such a constraint using additional loss terms, but a simpler alternative is to explicitly normalise the embedding vectors. We find that this approach is very effective, and it yields a model that is trainable end-to-end with the cross-entropy loss, without requiring any additional terms. Concretely, we always L2-normalise the embedding vectors before they are used, but we allow the underlying parameters to vary freely. We backpropagate through the normalisation operation as needed.

3 Time warping

Diffusion models are essentially denoisers that can operate at many different noise levels with a single set of shared parameters. Therefore, the degree to which the model dedicates capacity to different noise levels has a significant impact on the perceived quality of the resulting samples. We can control this by appropriately weighting the noise levels during training. The impact of different weightings has been studied extensively for diffusion models of images (Nichol and Dhariwal, 2021; Song et al., 2021a; Kingma et al., 2021; Karras et al., 2022).

We note that our use of the cross-entropy loss, while conveying many benefits (such as stability and end-to-end training), also changes the relative weighting of the noise levels corresponding to different timesteps $t$. Because the effects of this change are difficult to quantify, we seek to determine a time reweighting strategy that maximises sample quality. In practice, this is best implemented through importance sampling: rather than explicitly multiplying loss contributions from different noise levels with a weighting function $\lambda(t)$, we will instead sample $t$ from a non-uniform distribution, whose density directly corresponds to the desired relative weighting. This way, we avoid introducing significant variance in the loss estimates for timesteps $t$ for which $\lambda(t)$ is particularly large.

To sample $t$ non-uniformly in practice, we can use inverse transform sampling: we first sample uniform samples $u\in$ and then warp them using the inverse cumulative distribution function (CDF) of the distribution which corresponds to the desired weighting: $t=F^{-1}(u)$. This time warping procedure is equivalent to time reweighting in expectation, but more statistically efficient.

To estimate the CDF $F(t)$ in question, we propose to use the following heuristic:

This heuristic has an intuitive interpretation: it implies that the uncertainty of the model predictions (measured in bits or nats) increases at a constant rate as a function of ‘uniform time’ $u$. Therefore, the capacity of the model should be evenly distributed across the information content of the input sequences. Furthermore, when sampling from the model, using equally spaced timesteps in uniform time will result in a gradual decrease of uncertainty with an approximately constant rate. This ensures that all sampling steps do an equal amount of work in terms of resolving uncertainty.

This yields an active learning strategy where noise levels are initially sampled uniformly, but as training progresses, the noise level distribution changes to focus attention towards those levels for which training is the most useful.

Typically, we find that time warping puts most of the weight on intermediate noise levels. At very low noise levels, the denoising classification problem becomes trivial, because the embeddings corresponding to each token are easy to identify. At very high noise levels, the optimal strategy is to predict the marginal distribution of tokens (given the conditioning), which is also relatively straightforward to learn.

Diffusion language models

Using the CDCD framework described in the previous section, we can now construct language models for different tasks. We will describe a general Transformer (Vaswani et al., 2017) architecture which can be used for prompt completion and infilling, using a boolean conditioning mask to indicate which tokens in the sequence are to be sampled (‘noisy’), and which tokens are given as conditioning (‘clean’). The full model setup is visualised in a diagram in Figure 3. We will also describe an encoder-decoder model architecture for machine translation.

Since CDCD enables us to reduce the language modelling problem to a denoising classification task, without imposing any restrictions on the model architecture (such as causality, see §2.3), we are able to use the Transformer architecture without any form of attention masking.

However, practical applications of language modelling require the ability to fix some subset of the tokens in a sequence, while generating the rest, conditioned on these given tokens. One way to achieve this would be to ‘clamp’ certain token positions throughout the sampling procedure, by reinjecting these tokens at each step, corrupted by the appropriate level of noise. This method was proposed by Song et al. (2020) and referred to as ‘the replacement method’ by Ho et al. (2022b). It seems attractive, because it would allow us to treat the model as fully unconditional during training. Nonetheless, it is not as effective as training the model to specifically support conditional sampling out of the box. Indeed, diffusion models for image inpainting are also more effective when trained specifically for that task (Saharia et al., 2022a).

Therefore, we construct the input to the model by stacking three sequences:

$\mathbf{x}$: the embeddings corresponding to the noisy input sequence, with embeddings for conditioning tokens set to the zero vector;

$\mathbf{c}$: the embeddings corresponding to the conditioning tokens, with embeddings for the tokens to be sampled set to the zero vector;

$m$: the boolean conditioning mask, indicating which tokens are given (‘clean’, $m_{i}=0$) and which are to be generated (‘noisy’, $m_{i}=1$).

As discussed before in §3.1, the output of the model consists of a sequence of logit vectors, which correspond to the predicted probabilities of each token in the vocabulary occurring at each sequence position. When calculating the training loss, we zero out the positions for given tokens.

Autoregressive language models naturally allow prefix conditioning, where the start of a sequence is given and the model generates a completion. This is not as straightforward with diffusion models. Since prefix conditioning is a very general procedure for interacting with language models, we would like our models to support it. To achieve this, we randomly sample mask sequences $m$ during training which correspond to prefix conditioning. Such masks consist of a sequence of zeros of a certain length, followed by a sequence of ones, where the length of the prefix is sampled uniformly at random.

Fully random masking

Since diffusion models are able to iteratively refine all tokens in a sequence in parallel, we are not restricted to prefix conditioning. To enable conditioning on an arbitrary subset of sequence positions, we can sample masks $m$ fully randomly during training. Rather than sampling a mask value independently for each sequence position, we first sample a clean position count uniformly at random, and then randomly select a subset of that size from the sequence. This ensures that the model is able to support conditioning on any number of tokens.

Mixed masking

While fully random masking yields the most general model, supporting conditioning on any arbitrary subset of tokens in a sequence, prefix masking is sufficient to support the most common use cases for language models. Somewhat surprisingly, we find that training on an equal mixture of prefix masks and fully random masks actually slightly improves prefix completion performance (see §6.3).

2 Noise level conditioning

Diffusion models operate on inputs corrupted with varying levels of noise. These levels correspond directly to timesteps in the diffusion process. We provide the timestep as an additional input, which is incorporated into the model using conditional normalisation: each layer normalisation operation in the model is followed by shifting and scaling the activations, with the shift and scale parameters depending on the timestep (Perez et al., 2018).

3 Self-conditioning

Chen et al. (2022) introduced self-conditioning, which significantly improves the performance of diffusion models in certain contexts. They noted that the predictions produced by diffusion models are only used to determine the direction in which to update the noisy input, and are then discarded, which is wasteful. Giving the model direct access to the predictions it produced at the previous sampling step enables a more efficient use of model capacity, by allowing it to refine previous predictions, rather than constructing them from scratch at each step. This approach bears a strong resemblance to the unrolled denoising strategy proposed by Savinov et al. (2021).

To enable the model to make use of this additional input without requiring unrolling across multiple sampling steps during training, they propose a training procedure which only requires an additional forward pass on half of the batch at each training step. Following Strudel et al. (2022), we use this procedure and find that it only increases training time by 10-15% for our models in practice, while yielding significant performance gains.

To use self-conditioning with CDCD, the input to the model now consists of four stacked sequences. In addition to $\mathbf{x}$, $\mathbf{c}$ and $m$ (see §4.1), we add $\mathbf{p}$, which is a sequence of embeddings found by interpolating the vocabulary embeddings using the token probabilities predicted in the previous sampling step (see §3.1; embeddings for conditioning tokens are set to the zero vector).

4 Machine translation model

For machine translation, we use an encoder-decoder architecture with two separate Transformer stacks. Since the conditioning (source) sequence and target sequence are separate, there is no need for masking. We adopt an architecture very similar to the original Transformer (Vaswani et al., 2017), with absolute positional embeddings (sinusoidal on the source side, learned on the target side) instead of RoPE, and ReLU as activation function.

Generating samples of the correct length is an important concern for translation. As we do not use any causal mask on the decoder side, we cannot simply disregard the contribution to the loss of the positions corresponding to padding tokens during training. Instead, we simply predict the whole sequence of tokens, including beginning-of-sentence (BOS), end-of-sentence (EOS), and padding tokens, up to a constant length. At sampling time, tokens past the first EOS are discarded.

In order to provide a strong conditioning signal to the decoder, we extend the conditioning described in §4.2 to additionally use the length of the source sequence.

5 Comparison to BERT

The model architecture and training procedure we have described so far is very similar to BERT (Devlin et al., 2018). Given the widespread use of that model, and its popularity among language modelling practitioners, we provide a side-by-side comparison.

Both models are sequence denoisers, but the nature of the noise differs. While BERT is trained on sequences corrupted by masking noise, which randomly removes a subset of tokens altogether, inputs to our model are corrupted by Gaussian noise which is added directly to the token embeddings. The input consists of a stack of multiple sequences (see §4.1 and §4.3), rather than a single sequence where some of the tokens have been masked. Because the intensity of the noise varies according to the timesteps of the diffusion process, we also provide the timestep as an additional input, which is not required for BERT. Finally, to avoid uncontrollable growth of the embedding parameters, we force the embeddings to be normalised. Other than that, the model architectures are essentially identical during training, and the loss functions are the same.

Related work

Several diffusion-based and diffusion-inspired approaches have been proposed for non-autoregressive iterative refinement of discrete data, and especially for language in particular (Ghazvininejad et al., 2019; Chang et al., 2022; Austin et al., 2021; Hoogeboom et al., 2021a, b; Savinov et al., 2021; Reid et al., 2022). Replacing continuous diffusion with a discrete corruption process affords some flexibility, but it also requires forgoing several capabilities associated with the continuous paradigm, such as efficient sampling algorithms based on advanced ODE solvers, or classifier-free guidance.

More recently, several papers have proposed approaches to apply the continuous diffusion framework to discrete data. It is important to distinguish continuity in the input space from continuity of the time variable of the corruption process; for CDCD, both are continuous. Li et al. (2022), Strudel et al. (2022) and Han et al. (2022) all target language modelling as the primary application and use an embedding-based strategy in combination with discrete-time diffusion. Campbell et al. (2022) and Sun et al. (2022) propose continuous-time models for discrete input, though they do not explore the application to language modelling. Meng et al. (2022) propose concrete score matching, which can be applied to both discrete and continuous inputs. Chen et al. (2022) use continuous-time diffusion applied to continuous relaxations of binary representations of the input. Chen et al. (2022); Li et al. (2022); Strudel et al. (2022) also suggested using the cross-entropy loss (in combination with other loss terms).

2 Iterative refinement for machine translation

There have been considerable efforts to apply non-autoregressive iterative refinement models to the task of machine translation. The first attempts by Gu et al. (2017) already uncovered the issue of ‘multi-modality’: since non-AR models usually predict all tokens in parallel and independently of each other, uncoordinated sampling decisions might lead to incoherencies like repeated tokens. Earlier advances in sequence-level distillation (Kim and Rush, 2016) allowed to alleviate those issues, albeit at the cost of expensive training dataset creation. Another line of work introduced a latent transformer (Kaiser et al., 2018): discrete latents are first sampled autoregressively, and then decoded non-autoregressively.

LVM-DAE (Lee et al., 2018) applied a non-autoregressive decoder multiple times to alleviate multi-modality. Insertion (Stern et al., 2019) and Levenshtein (Gu et al., 2019) transformers demonstrated good parallel decoding results on machine translation along with editing capabilities. The LVM-DAE line of work was later significantly improved upon by the CMLM (Ghazvininejad et al., 2019) and DisCo (Kasai et al., 2020a) methods via novel training and decoding procedures. A later work by Kong et al. (2020) combined CMLM with local autoregression. Promising advances were achieved by another follow-up of CMLM called SMART (Ghazvininejad et al., 2020), further closing the gap between AR and non-AR methods. More recently, Imputer (Chan et al., 2020; Saharia et al., 2020) achieved good results by optimizing alignment between source and target.

Kasai et al. (2020b) questioned the speed advantage of non-AR models in machine translation by comparing them to a shallow-decoder AR baseline. SUNDAE (Savinov et al., 2021) introduced step-unrolls and obtained excellent results both in machine translation and unconditional generation without relying on sequence-level distillation. Huang et al. (2022) also eliminated distillation with an unroll-like technique. Aggressive decoding (Xia et al., 2022) generated tokens in parallel while using AR models to verify generation and re-generate after the first deviation. DiffusER (Reid et al., 2022) used edit-based reconstructions and 2D beam search to almost completely close the gap between AR and non-AR models in machine translation.

3 Other related work

Kingma et al. (2021) suggest parameterising and optimising the noise schedule during training in a similar fashion to time warping (§3.3), though the objective is different: the goal is to minimise the variance of the diffusion loss, whereas our goal is to linearise the entropy of the model predictions.

Experiments

We study the effect of various design choices, and compare language models based on CDCD with the standard autoregressive approach. We train mask-conditional models for tasks such as prompt completion and infilling, and encoder-decoder models for machine translation.

We set the embedding dimensionality $d=256$. We L2-normalise the embeddings and scale them by $\sqrt{d}$, so that each component has a standard deviation of 1. We choose $t_{\min}=1.0$ and $t_{\max}=300.0$. While these values are quite different from the ones suggested by Karras et al. (2022) for image diffusion, we note that the discrete underlying nature of the input data makes it possible to predict the original tokens with 100% accuracy even when noise with $\sigma=1.0$ is added (recall that $\sigma=t$, see §2.1). Therefore, it is not useful to consider lower noise levels, except when the embedding dimensionality $d$ is reduced. We scale the noisy embeddings by $\frac{1}{\sqrt{\sigma^{2}+\sigma_{data}^{2}}}=\frac{1}{\sqrt{t^{2}+1}}$ before passing them into the model, so that the components again have a standard deviation of 1.

We drop out the conditioning by zeroing out the corresponding embeddings for 10% of training examples, in order to be able to support sampling with classifier-free guidance (Ho and Salimans, 2022). We use low-discrepancy sampling (Kingma et al., 2021) to sample uniform timesteps $u$ before applying time warping, to reduce the variance of the loss estimates. We use a mixed masking strategy (see §4.1): 50% of the masks sampled during training are prefix masks (with the length of the prefix uniformly sampled), and the other 50% are fully random masks (with the number of clean token positions again uniformly sampled).

We use the Adam optimiser (Kingma and Ba, 2015) with a learning rate of $10^{-4}$, $\beta_{1}=0.9$ and $\beta_{2}=0.99$. We use 200 Euler steps for sampling, and do not tune any sampling parameters (such as temperatures or guidance scales). While this is a large number of steps relative to the sequence length, we wanted to ensure that our measurements would not be negatively affected by discretisation errors introduced by the ODE solver.

2 Evaluation

We train mask-conditional models on the MassiveText dataset (Rae et al., 2021), except for the larger model used in §6.5, which is trained on the publicly available C4 dataset (Raffel et al., 2019). Both datasets contain web-crawled content, and are very diverse as a result. For machine translation, we train models on the WMT2014 German-English / English-German and WMT2020 Chinese-English datasets.

To evaluate our machine translation models, we follow the literature and use the BLEU score (Papineni et al., 2002). For mask-conditional models, quantitative evaluation is more challenging. Following Savinov et al. (2021); Strudel et al. (2022), we measure the likelihood of generated samples under a 1.3B parameter autoregressive language model (AR-NLL), as well as the unigram (per-token) entropy of the samples. As long as the entropy remains high enough, we found this negative log-likelihood to be strongly correlated with subjectively assessed model quality, and we made extensive use of this metric for hyperparameter exploration. To ensure a fair comparison, we always calculate these metrics using a fixed prefix mask with a prefix length of half the sequence length, regardless of the masking strategy used during training.

For experiments with a larger mask-conditional model, we also report the MAUVE metric (Pillutla et al., 2021), which is specifically designed for open-ended text generation and has been shown to correlate with human judgement (see §6.5).

3 Design decisions

We conduct various ablation experiments to justify our architecture and hyperparameter choices, reporting both the autoregressive negative log likelihood (AR-NLL) and the unigram entropy at the token level (H) in each case. The results are summarised in Table 1 and visualised in Figure 4.

Both manipulations of the score estimate (see §3.2) significantly reduce the AR-NLL, but also reduce the entropy quite a lot. Anecdotally, we also find that models using renormalisation are less amenable to improvements from sampling hyperparameter tuning, and models without renormalisation produce better samples when tuned.

Because we also make use of self-conditioning (see below), it is important to treat renormalisation and clamping as training-time hyperparameters, because they will affect the previous predictions $\mathbf{p}$ that the model receives as input (see §4.3). When not using self-conditioning, renormalisation and clamping could instead be treated as sampling hyperparameters (see §6.4).

Embedding dimensionality

We reduced $t_{\min}$ from $1.0$ to $0.1$ for this experiment, because at lower embedding dimensionalities, the same amount of noise erases more information. We verified that the cross-entropy is nearly zero at $t=0.1$ for all values of $w$ we considerThe cross-entropy is significantly higher than zero at $t=1.0$ for $w=16$, for example.. AR-NLLs initially improve with increasing dimensionality, but stabilise beyond $w=64$, thanks to time warping automatically shifting focus to the relevant noise levels (see Figure 2). We used $w=256$ and $t_{\min}=1.0$ for all other experiments, including the base model, but the result obtained with $t_{\min}=0.1$ is very similar.

Embedding initialisation scale

Results degrade when the scale of the initial embedding parameters is too high. We use $\sigma=0.001$ throughout to avoid this pitfall. Given its impact on performance, tuning this parameter is important, which is worth noting because weight initialisation scales are usually chosen heuristically, and generally are not treated as important hyperparameters to tune. We normalise the embeddings whenever they are used (see §3.2), so the scale of the underlying parameters is not relevant for inference, but clearly it significantly affects optimisation.

Masking strategy

We find that training with fully random masks improves results, even when evaluating using a prefix mask. We suspect that learning the embeddings becomes easier when bidirectional context is available. To ensure that the model uses enough capacity for prefix completion tasks, we use a 50-50 masking strategy for all other experiments.

Self-conditioning

Like Strudel et al. (2022), we find that self-conditioning (Chen et al., 2022) has a significant positive impact on model performance, by enabling reuse of computation from preceding sampling steps. The AR-NLL improves significantly, while the entropy stays roughly at the same level.

Time warping

During training, focusing on the right noise levels is clearly important, but we also find that spacing the sampling steps accordingly is essential to benefit from this improvement (see §3.3). We experimented with several manipulations of the warping function to verify the quality of our proposed entropy linearisation heuristic. We find that changing the temperature or the uniformity of the weighting (see Appendix A and Figure 8) can sometimes yield improvements, but they are relatively minor.

4 Sampling

Using the base model, we investigate the impact of various sampling hyperparameters when they are varied individually. Note that there are some significant interactions between sampling hyperparameters, which are not reflected in per-parameter experiments. We will also look at the interaction between the score temperature and the classifier-free guidance scale as an example of this.

We compare the Euler and Heun samplers suggested by Karras et al. (2022), reducing the number of sampling steps by a factor of two for the Heun sampler, so that the total number of required function evaluations does not change. For self-conditioning, we always use the most recent model prediction.

The Heun sampler seems to offer little benefit when using a sufficient number of stepsNote that the suitability of higher-order samplers is known to be highly dependent on other sampling hyperparameters such as the guidance scale (Lu et al., 2022b)., and entropy degrades when the number of steps is decreased, which is not the case for the Euler sampler. We use the Euler sampler with 200 steps for most experiments, but reasonable quality can be achieved with as few as 50 function evaluations. Based on results in the image domain, it may be possible to reduce this further with stochastic samplersDeterministic samplers have some important benefits over stochastic ones, such as support for manipulations in latent space., or more advanced sampling algorithms (Lu et al., 2022a, b).

Scaling hyperparameters

Many classes of generative models offer some notion of ‘temperature tuning’. The CDCD framework is particularly versatile in this regard:

scaling the score function estimate by a factor corresponds to changing the temperature of $p_{t}(\mathbf{x})$;

the standard deviation of the initial noise can be reduced below $1.0$, which is often done for flow-based models (Kingma and Dhariwal, 2018);

since the model produces a categorical probability distribution $p(\mathbf{x}_{0}|\mathbf{x},t)$, we can also manipulate this using various truncation strategies, such as temperature tuning and nucleus samplingNucleus sampling is a misnomer in this case, because we never actually sample from the categorical distribution, but we can still use the same strategy to shape the logits. (Holtzman et al., 2019);

classifier-free guidance (Ho and Salimans, 2022) can be used to amplify the influence of the conditioning.

Out of these, we find that scaling the score temperature or changing the initial noise scale offer a strictly better trade-off than manipulating $p(\mathbf{x}_{0}|\mathbf{x},t)$. The effectiveness of changing the score temperature is surprising, because this tends to work very poorly for diffusion-based models in the visual domain. While the trade-off offered by classifier-free guidance seems even better at first glance, in practice we find that samples obtained with high guidance scales tend to contain a lot of repeated phrases. We also study the interaction between the score temperature and the guidance scale (see Figure 6), and find that their effects are complementary to a degree.

Step spacing

Karras et al. (2022) suggest that spacing the sampling steps non-uniformly significantly improves sample quality. For CDCD, our use of time warping at sampling time already results in non-uniform step spacing. We compare their heuristic for $\rho=1,2,4$ or $8$ with time warping, and we also investigate whether they compound, by applying time warping to non-uniformly spaced steps obtained using their heuristic.

When not using time warping at sampling time, the heuristic is clearly helpful, but it does not reach the same performance. We also find that the effects do not compound, strengthening our intuition that spacing the steps using time warping is close to optimal, because it makes the rate of decrease of uncertainty approximately constant during sampling.

Time warping

In §6.3, we established that manipulating the warping function at training time is not particularly helpful. We can also choose to manipulate it only at sampling time, where it can still affect step spacing (but not model training). We again find that these manipulations do not have any meaningful positive effect on sample quality.

Final prediction

We run the model one additional time after all sampling steps are completed, and take the argmax of the predicted distributions at each sequence position to determine the sampled tokens. Instead, we could use the tokens whose embeddings are nearest to the predicted embeddings in the Euclidean sense. This works equally well, so it can save some computation, though this only becomes significant if the number of sampling steps is greatly reduced.

5 Prompt completion and infilling

We compare a 1.3B parameter model based on the CDCD framework with a pre-trained autoregressive model with the same architecture (24 Transformer blocks, 2048 units, 16 attention heads). Both models are trained on the C4 dataset for 600,000 steps with 524,288 tokens per batch. Due to the reduced data efficiency of diffusion model training (see §2.3), it is likely that the CDCD model would benefit more from further training. Relative to the autoregressive model, the CDCD model has some extra learnable parameters in the MLP for timestep embedding, and in an initial linear layer which maps the token embeddings to the Transformer hidden state. The token embeddings themselves on the other hand account for $8\times$ fewer parameters, because we use an embedding dimensionality of $256$ instead of $2048$.

While the autoregressive model was trained with a batch size of 256 and a sequence length of 2048, we trained the CDCD model with a batch size of 2048 and a sequence length of 256 instead. This is partly because at this point, we are not focusing our evaluation on long-range coherence, but also because diffusion model training benefits from larger batch sizes: since noise levels are sampled on a per-sequence basis, a larger batch size yields a lower variance loss estimate. Note that the number of sampling steps is still 200, which is now smaller than the sequence length.

We take 5,000 token sequences of length 256 from the C4 validation set (rejecting shorter sequences and cropping longer ones). For each sequence, we select a random-length prefix to use as the prompt, with prompt lengths varying between 0 and 128 (half the sequence length). We sample from the autoregressive model using these prompts with nucleus sampling, for various values of $p$. We prevent the end-of-sentence (EOS) token from being sampled, to ensure that a full-length sequence is produced every time. This is necessary because the MAUVE metric is calculated on the full sequences, including the prompts, so the presence of shorter sequences would bias the results. We also sample from the CDCD model using the same set of prompts, with different score temperatures ($T$) and classifier-free guidance scales ($\gamma$).

The resulting completions are evaluated by comparing them with the original sequences from the dataset using MAUVE, which we report in Table 3, alongside the AR-NLL (measured with the autoregressive model itselfThese numbers are not directly comparable to those reported in Tables 1 and 2, which were obtained with a model trained on a different dataset.) and the unigram entropy for the generated sequences. We get favourable MAUVE scores with the CDCD model for several settings of the score temperature $T$ and guidance scale $\gamma$. Note however that the MAUVE scores for the autoregressive samples do not seem to be convex in $p$, and it is unclear to what extent any excessive repetition introduced by increasing the guidance scale is penalised, so these numbers should be interpreted with care. Nevertheless, they provide some evidence that the CDCD model is able to produce compelling samples. Selected prompt completion and infilling samples are shown in Figure 7.

6 Machine translation

We compare a Transformer autoregressive machine translation model to a CDCD model with an encoder-decoder of the same size, on three translation tasks – WMT 2014 English-German, WMT 2014 German-English, and WMT 2020 Chinese-English – and two model sizes – Transformer base and big. For each task, all models are trained on the same training set of the corresponding year, and evaluated on newstest2014 and newstest2020 accordingly. The standard validation sets are used for early stopping and for tuning the hyperparameters of CDCD. We train two shared SentencePiece (Kudo and Richardson, 2018) tokenizers (one for English/German and one for Chinese/English) of size 32,768, using the unigram method (Kudo, 2018), with byte fallback.

For the autoregressive baseline, we use the hyperparameters described by Vaswani et al. (2017), except for the following: we tie the source/target/output embeddings following Press and Wolf (2017), we use the fixed vocabulary size indicated above, and, for Chinese-English only, we use a beam size of 6, and a length penalty $\alpha=0.6$. The CDCD models are trained for 2M steps, with a batch size of 512 for base and 1024 for big, and a sequence length of 160. We use the same dropout values as for the autoregressive baseline. At sampling time, swe set the score temperature $T=0.8$ and classifier-free guidance scales $\gamma=4.0$ for English-German and German-English, and $\gamma=8.0$ for Chinese-English, found using a manual hyperparameter search on the validation sets.

Using a diffusion framework provides a natural way of approximating Minimum Bayes-Risk decoding (Kumar and Byrne, 2004; Eikema and Aziz, 2021) by sampling a number of hypotheses, each from a different initial noise vector, and selecting the one minimizing in expectation a metric of interest, in our case BLEU.

We report BLEU scores computed on the test set for each model, year and language pair, using sacreBLEUThe signature is nrefs:1|case:mixed|eff:no|tok:13a|smooth:exp|version:2.3.1 when English is the target language, and nrefs:1|case:mixed|eff:no|tok:intl|smooth:exp|version:2.3.1 for German. (Post, 2018), in Table 4. Overall, CDCD models perform worse (between 3 and 7 BLEU points) than autoregressive counterparts of the same size. The difference is largest on English-German, where manual examination of the samples reveals that hypotheses in English sometimes appear, even though German output is expected. This certainly contributes to the poor score that the CDCD model obtains on this pair. It might be explained by the noisy character of the training data itself (English appearing on the German side). CDCD does comparatively better on Chinese-English, especially at larger scale, which may enable the model to make use of the larger training set. Finally, sampling-based MBR decoding improves translation quality, and monotonically so as a function of the number of samples, providing a gain of 0.7 to 1.8 BLEU points when using 100 samples compared to a single one.

Nonetheless, hypotheses produced by CDCD in our experiments tend to suffer from systematic defects, in particular repeated or missing tokens. We hypothesise that the model is not trained to recover from these errors, and that data augmentation, in the form of extra padding, or corruption, may prove beneficial. We leave exploring this direction to future work.

Discussion

We have proposed CDCD, a framework for diffusion of categorical data that is continuous both in time and input space, which enables training of non-autoregressive language models with a procedure reminiscent of BERT (see §4.5).

Our proposed approach has some advantages over autoregressive models, such as the ability to perform arbitrary infillingInfilling is also possible with autoregressive models if they are specifically trained for this task (Bavarian et al., 2022). and a flexible sampling procedure which allows trading off sample quality and compute requirements. It also stands to benefit from diffusion model enhancements such as classifier-free guidance and improved sampling algorithms, and from the ability to deterministically map inputs to latents and vice versa (which we have not explored so far). Some limitations compared to autoregressive models were discussed in §2.3. An important research question which we have not yet investigated, is how best to handle variable-length output. With autoregression, this is very naturally handled by introducing an end-of-sentence (EOS) token, which the model can predict to indicate that sampling should be halted. Diffusion models on the other hand require a fixed-size ‘canvas’ which is iteratively refined during sampling, although inserting padding tokens at random during training (and collapsing them after sampling) could allow for some length variation (Strudel et al., 2022), and it may also reduce the incidence of repeated tokens in the samples.

We intentionally limited the degree of architectural exploration in this work, because the choice of architecture is largely orthogonal to the framework we have proposed. We focused on the Transformer as it is currently the canonical language model architecture. Nevertheless, we pointed out the absence of any architectural restrictions in this modelling paradigm (see §2.3), which considerably simplifies the use of more intricate patterns, such as multi-resolution or Perceiver-based architectures (Jaegle et al., 2021).

We would also like to explore recent innovations for sampling from diffusion models (e.g. Lu et al. (2022a, b)), which have so far been demonstrated chiefly in the image domain, and determine to what extent they improve sampling efficiency in the CDCD framework.

Application domains

While we concentrated our empirical evaluation of the proposed CDCD framework on language tasks, none of its components are specific to language. We expect it to be suitable for any generative modelling problem involving structured categorical data. As an example, we note that latent diffusion models such as Stable Diffusion (Rombach et al., 2022) first use VQ-VAE (Van Den Oord et al., 2017) or VQ-GAN (Esser et al., 2021) to learn a latent space in which the diffusion process is then applied, by mapping the categorical latents to the continuous embeddings from the vector quantisation bottleneck. While this works well, we hypothesise that fitting new embeddings jointly with the diffusion model could bring further improvements, especially in combination with time warping. Finally, our proposed time warping heuristic may also be useful beyond diffusion models of categorical data.

Acknowledgements

We would like to thank Andy Brock, Bart Chrzaszcz, Noah Constant, Jeff Donahue, Douglas Eck, Dominik Grewe, Jordan Hoffmann, Patrick Kidger, Skanda Koppula, Lena Martens, Katie Millican, Ben Moran, Simon Osindero, Evan Shelhamer, Miloš Stanojević, Federico Vaggi, Björn Winckler, and the wider DeepMind team for their assistance and input.

We are also thankful to the creators and maintainers of the open source software used in this work, including Python (Van Rossum and Drake, 2009), NumPy (Harris et al., 2020), SciPy (Virtanen et al., 2020), JAX (Bradbury et al., 2018), TensorFlow (Abadi et al., 2016), the DeepMind JAX Ecosystem (Babuschkin et al., 2020), Diffrax (Kidger, 2021) and Matplotlib (Hunter, 2007).

References

Appendix A Fitting the unnormalised CDF

To implement time warping (see §3.3), we fit a monotonically increasing function to the expected loss at each timestep. It is essential to parameterise this function in a way that is easy to normalise and invert, so we can use it for inverse transform sampling.

We take inspiration from Müller et al. (2019); Durkan et al. (2019) and parameterise this function by dividing both the input and output range into bins. For convenience, we will assume that the normalised CDF maps the interval $$to$$. In practice, timesteps $t$ range from $t_{\min}$ to $t_{\max}$ (with $t_{\max}>t_{\min}>0$), but we can simply shift and scale them:

The normalised CDF $u=F(t^{\prime})$ is then parameterised using two sets of logits $l_{n}^{t}$ and $l_{n}^{u}$, $n=1,\ldots,N$. Applying the softmax nonlinearity to both sets yields two partitions of the unit interval, which we use to define $N$ input and output regions. The sizes of the input and output bins are:

The left edges of the input and output bins are found as the cumulative sum of the sizes of all preceding bins:

and similar for $e_{n}^{u}$. The right edges are found by adding the bin sizes to the left edges.

A.2 Fitting without normalisation

We now wish to fit $u=F(t^{\prime})$ so that it has the same shape as the expected loss as a function of time $t^{\prime}$. However, $F(t^{\prime})$ is normalised so that $u\in$, whereas the expected loss will vary between and $H$, the unigram entropy of the data. For an unconditional diffusion model, we can estimate said entropy remarkably accurately from a single batch of training data, simply by counting the frequencies of all tokens in the vocabulary to estimate the marginal token distribution. We could then use this entropy estimate to scale the loss values to the unit interval. Unfortunately, this does not work for conditional models, because in that case, the unigram entropy varies with the conditioning.

A.3 Inverse and derivative

Apart from being trivial to normalise, it is also very easy to invert the learnt CDF, simply by switching the roles of the input and output logits $l_{n}^{t}$ and $l_{n}^{u}$. We can also easily evaluate the derivative of the CDF (i.e. the probability density function), which is piecewise constant: within each bin, it is equal to the ratio of the output and input bin size, $w_{n}^{u}/w_{n}^{t}$ (Figure 2, middle).

A.4 Importance weighting

Since we are using this fitting mechanism to change the distribution of sampled timesteps over the course of training, we have created a feedback loop, because the training data for $F(t^{\prime})$ will itself become biased towards oversampled timesteps over the course of training. To compensate for this, we can use importance weights, which correspond to the reciprocal of the derivative of the CDF (or equivalently, the derivative of the inverse CDF). We take care not to backpropagate gradients through the time warping operation and the importance weights. For additional stability, we use an exponential moving average of the parameters $l_{n}^{t}$ and $l_{n}^{u}$ when performing the warping, to limit the rate of change of the training distribution, though empirically we have found that this is not strictly necessary.

A.5 Warping sampling timesteps

For sampling, we find that using steps that are linearly spaced in uniform time, and subsequently warped, works considerably better than using uniformly spaced timesteps. Using warping for both training (to change the distribution of noise levels on the fly) and sampling (to space the timesteps nonlinearly) yields the best results. We hypothesise that decreasing the entropy at a constant rate from step to step is also a useful heuristic for sampling.

A.6 Temperature and uniformity

The piecewise linear parameterisation also enables some useful distribution manipulations: we can easily change the temperature $T$ of the distribution represented by the CDF, by changing the output bin sizes:

To see why this corresponds to a temperature change, recall that the PDF is piecewise constant, and the values it assumes (up to a normalisation constant) are given by:

Similarly, we can also derive the CDF of a mixture of the distribution with a uniform distribution, with mixture weight $\mu$:

Note that when the corresponding input and output bin sizes are all equal to each other (i.e. $w_{n}^{u}=w_{n}^{t},\forall n$), we always obtain the identity function, which corresponds to the CDF of the uniform distribution on the unit interval.

The effect of these manipulations is visualised in Figure 8.

Appendix B Negative results

We informally discuss some failed attempts to improve our results. This information is provided to help understand some of our design choices, and to aid researchers and practitioners who are interested in using these methods or investigating ways to improve them. Bearing in mind that context and details significantly impact experimental results, especially in machine learning research, we expressly do not wish to discourage anyone from pursuing the ideas described here.

Instead of normalisation, the embedding parameters can also be prevented from growing uncontrollably by adding a regularisation loss term. We experimented with L2 regularisation, margin penalties (Donahue et al., 2019; Dieleman et al., 2021), as well as automatic adaptation of the L2 penalty weight to target unit variance for the embedding vector components (Rezende and Viola, 2018; Dieleman et al., 2021), but found that normalisation works best.

B.2 Removing the time dependency

In an effort to simplify the model architecture and make it resemble BERT even more strongly, we investigated models without time embedding (i.e. we removed the time embedding MLP in Figure 3). Since the model is solving a classification task, our hypothesis was that it could simply infer the noise level from the noisy input vectors, instead of requiring explicit knowledge of the timestep $t$. Unfortunately this significantly hurt performance. We suspect that this is a result of the relative scale of the embeddings and the noise (which varies greatly across noise levels), as well as the rescaling we apply to the noisy embeddings to ensure unit variance (see §6.1).

B.3 Simplex diffusion

We started out by lifting discrete token sequences into the space of categorical distributions over tokens, and performing diffusion in that space. Categorical distributions are nonnegative real-valued vectors whose components sum to 1, so they live on the simplex. We explored a tractable formulation of diffusion on the simplex, based on the Cox-Ingersoll-Ross process (Cox et al., 1985), as described by Richemond et al. (2022).

We found that language modelling with this process is impeded by the uneven nature of the corruption process, which is a consequence of the high dimensionality of the simplex (corresponding to the number of tokens in the vocabulary $V$). In practice, the noise distribution is heavy-tailed and introduces frequent outliers, which make the corrupted vectors look like they correspond to the wrong tokens, even at very low noise levels. We were not able to circumvent this issue even with powerful Transformer models. We hypothesise that mitigating this issue requires modifying the corruption process so that it does not operate independently on all components. Formulating such a correlated process with a tractable transition density is non-trivial however, and this would complicate the model to a certain degree.

Both score interpolation (§3.1) and time warping (§3.3) were originally developed in the context of simplex diffusion, but we found these ideas to be more effective in combination with Gaussian diffusion in a Euclidean embedding space.