Discrete diffusion for text
All articles

Discrete diffusion for text: the math behind masked token refinement

Diffusion models for images work in continuous pixel space. You start from Gaussian noise and learn to iteratively denoise. Text is discrete: tokens are category indices, not real-valued coordinates. You cannot add Gaussian noise to a token index and get a meaningful gradient signal back. The formulation that makes diffusion-style generation work for text is different enough from the image case that it deserves its own careful treatment. This post walks through the core math.

The forward process: corruption by masking

In continuous diffusion, the forward process adds noise progressively until the distribution converges to a known prior, typically an isotropic Gaussian. In discrete diffusion, the forward process corrupts tokens progressively by replacing them with a special mask token. At each step in the forward chain, each token has some probability of being replaced by the mask. By the final step of the forward process, all tokens are masked.

Formally, for a sequence x of N tokens, the forward process defines a series of latent sequences x_1 through x_T, where x_0 is the original uncorrupted sequence and x_T is fully masked. At each step t, each token in x_{t-1} is independently replaced with the mask token [M] with probability given by the masking schedule, or left unchanged. The masking schedule determines how quickly corruption accumulates over steps, analogous to the noise schedule in continuous diffusion.

The conditional distribution q(x_t | x_{t-1}) is simple: for each position i, x_t[i] equals [M] with probability proportional to the mask rate at step t, or x_{t-1}[i] otherwise. This makes the forward process easy to compute and, importantly, it makes q(x_t | x_0) tractable in closed form. You can sample a corrupted version at any step t directly from the original sequence without running the forward chain sequentially, which is critical for efficient training.

The reverse process: predicting clean tokens from masked context

The generative process reverses the forward corruption. Starting from x_T, which is fully masked, the model iteratively predicts and unmasks tokens to recover a high-quality sequence. At each reverse step, the model receives a partially masked sequence and produces a distribution over what the clean tokens at masked positions should be.

The key model p_theta(x_{t-1} | x_t) predicts, for each masked position, which token from the vocabulary should occupy that position given all non-masked positions in x_t as context. Crucially, non-masked positions already hold their correct values from prior reverse steps. So the model at each step is doing a form of masked language modeling: predict masked tokens given surrounding context, but the surrounding context includes both original unmasked tokens and tokens that have already been unmasked in earlier reverse steps.

This is substantially different from standard BERT-style masked language modeling, which masks a fixed fraction of tokens and predicts them all at once from a single forward pass. In diffusion decoding, the masking fraction changes across steps and the model must handle the entire range, from fully masked at t equals T down to nearly unmasked at t equals 1. The model is trained to be accurate across this entire range, not just at a single fixed masking rate.

Training objective: predicting x_0 from x_t

The standard training objective in discrete masked diffusion is a denoising objective: given a corrupted sequence x_t sampled from the forward process, predict the original x_0. The loss is the cross-entropy between the model's predicted token distribution and the true token at each masked position.

An elegant property of the masking corruption structure is that you can train with a single forward pass per training example. Given x_0 from the training data, you sample a step t uniformly, apply the forward process to obtain x_t, and compute the loss on the masked positions. The model learns to predict clean tokens from corrupted context across all corruption levels simultaneously. This is efficient and does not require unrolling the reverse chain during training.

There is an important subtlety here: the loss should be weighted appropriately across steps. Steps near t equals T have mostly masked tokens and require the model to predict with less context, which is harder. Steps near t equals 1 have mostly unmasked context and require predicting only a few remaining positions. Uniform step weighting can lead to the model over-optimizing for easy steps. Various weighting schemes exist to address this, and the choice of weighting interacts with the masking schedule design.

The masking schedule and its effect on quality

The masking schedule, the function that maps step t to mask probability alpha(t), determines the shape of the forward process and consequently the difficulty profile of the reverse process. A linear schedule masks tokens at a constant rate per step. Cosine and exponential schedules concentrate the masking either at the beginning or end of the forward chain.

The choice of schedule matters because it determines when in the reverse chain the model faces the hardest predictions. With a schedule that starts masking slowly, early reverse steps see mostly masked context and must make nearly unconditional predictions. With a schedule that starts masking quickly, early reverse steps see richer context but must still make difficult local decisions. In practice, a cosine schedule that concentrates the transition in the middle of the chain tends to produce well-distributed difficulty across steps.

The total number of reverse steps K is a hyperparameter of inference, not training. The model is trained to reverse corruption from any level t, so at inference time you can choose K. Smaller K means coarser steps and potentially lower quality. Larger K means finer steps and higher quality at the cost of more forward passes. Finding the right K for a given quality target on a given task is empirical work. This is one of the active areas of refinement we spend time on at Inception: understanding how the K-quality tradeoff curve looks for different task types, particularly code generation versus natural language.

Decoding strategy during inference

At each reverse step, the model produces a probability distribution over vocabulary tokens for each masked position. You must decide how to convert that distribution into a concrete token to unmask. Several strategies exist, and they produce different tradeoffs.

Greedy unmasking picks the highest-probability token at each position at each step. This is fast and deterministic but tends to produce outputs that are locally overconfident. The model may commit to a token at step t that turns out to be suboptimal given what gets resolved in neighboring positions at step t-1.

Temperature sampling applies a temperature parameter to the per-position distributions before sampling, controlling how concentrated the draws are. Higher temperature produces more diverse outputs, lower temperature produces more deterministic outputs, similar to temperature sampling in autoregressive generation.

A more interesting strategy is to decide, at each step, not just which tokens to sample but also which positions to unmask. Rather than unmasking all positions at each step, you can unmask only the positions where the model is most confident, measured by the entropy of the predicted distribution. This schedule-within-a-schedule approach allows high-confidence positions to resolve first and provide context for the harder positions that resolve later. The order of resolution becomes a learned or heuristic function of the model's own confidence, which tends to produce better outputs for the same number of steps.

What discrete diffusion gives up compared to continuous

The mask-and-predict formulation is elegant but it is not the only possible discrete diffusion formulation. An alternative replaces tokens not with a mask but with uniformly random vocabulary tokens, creating an absorbing-state variant or a uniform corruption variant. The uniform corruption formulation has different properties: it forces the model to identify corrupted tokens rather than having a distinct mask token identity signal them. Training is harder but the model learns more general noise robustness.

The absorbing-state formulation we have described, where masking is the corruption type, is simpler to train and has been more reliable in our experience for code generation tasks. The mask token provides a clean signal about which positions need prediction and eliminates the need for the model to detect corruption, reducing one source of training difficulty. We are not claiming the masking approach is universally superior. For some task types and some model scales, the uniform corruption approach may be preferable. Treating this as a settled question would be a mistake.

Connection to the inference cost model

Everything described above feeds into the practical serving question: how does discrete masked diffusion compare to autoregressive generation in terms of the number of forward passes required?

Autoregressive generation of an N-token sequence requires N forward passes. Discrete diffusion generation requires K forward passes, where K is the number of reverse steps. Each forward pass in the diffusion case touches all N positions rather than one, but the total number of passes is K rather than N. When K is substantially less than N, the advantage is real. When N is small, the advantage narrows. The break-even point depends on the specific model, the masking schedule, and the acceptable quality level, which is why characterizing it requires empirical measurement on the task distribution you care about rather than theoretical calculation alone.

The formulation described here is the mathematical foundation underlying the inference engine we are building at Inception. The connection between training objective design and inference cost is direct: how you train the model for masked token prediction across the full masking schedule determines how efficiently the reverse process converges, which determines how small K can be while preserving output quality. That connection is where the interesting research work lives.