Masked Diffusion Language Models: Denoising Loss as a Likelihood Surrogate

Welcome back. Last time, you saw that the denoiser trains on masked clean text but generates from states shaped by its own earlier predictions. Today we’ll look at what its training loss actually measures. The central distinction is between predicting the next token of a sequence and predicting missing tokens at many noise levels. That distinction explains both why your mini diffusion model can generate text and why its loss is not an ordinary next-token likelihood.

A likelihood you can calculate left to right

For an autoregressive model, the probability assigned to a sequence $x=(x_1,\dots ,x_L)$ is built from next-token probabilities:

Taking the logarithm turns that product into a sum. During training, you can therefore score each token given the true preceding tokens and add the scores to obtain the model’s log-likelihood for the whole sequence.

The opening of Large Language Diffusion Models states this formulation explicitly. Read only the short passage around its Equation (2):

A masked diffusion model is different as a model, not because the chain rule has stopped being true. Any distribution over sequences can be written as a left-to-right product by the chain rule. But your denoiser does not directly provide those next-token conditionals: it takes a partially masked sequence, may use visible context on both sides of a blank, and predicts tokens for masked positions. Generation then uses several such predictions to turn a fully masked sequence into text.

You can picture the diffusion model’s probability for a completed sentence as combining the probabilities of possible intermediate masking and unmasking paths. Those paths are hidden when you look only at the final sentence. Consequently, a call to your denoiser does not give you a term to insert into the autoregressive product above.

What the denoiser learns instead

Your training loop already has the ingredients. Take a clean sequence $x$, choose a noise level $t$, and independently mask each eligible position with probability $t$. Call the resulting sequence $z_t$. The model sees $z_t$ and predicts the original token $x_i$ wherever position $i$ was masked.

Sasha Rush’s Simple Diffusion Language Models connects this reconstruction task to the training objective:

For the simple schedule in which the masking probability equals $t$, a standard masked-diffusion objective can be written as

Here the indicator selects only masked positions. The expression inside the sum is cross-entropy for the correct original token, not a score for the token that happens to be visible at position $i$. In your padded batches, the indicator must also exclude padding. Each new training pass can draw a different $t$ and mask pattern from the same clean example, so the denoiser practices filling blanks with little context, much context, and amounts in between.

The paper gives both this objective and its connection to sequence likelihood. Read Section 2.1 in two short stretches:

Under the paper’s specified reverse-process construction, its masked denoising objective is an upper bound on expected negative log-likelihood. Equivalently, its negative gives a lower bound on expected log-likelihood. This is why masked cross-entropy is a principled training surrogate: reducing it trains a model whose reverse process defines a distribution over completed sequences, without requiring you to calculate the exact probability of every sequence by summing over its possible hidden paths.

“Surrogate” does not mean “identical to likelihood.” A lower denoising loss need not translate directly into a particular improvement in the samples from your greedy, confidence-based notebook sampler. That sampler makes reveal choices differently from the stochastic reverse process used to establish the bound.

Why weight the loss by noise level?

Suppose there are four valid positions, and the average cross-entropy at a masked position is $2$. At $t=0.25$, you expect one masked position, so the expected sum of masked losses is $2$. At $t=0.75$, you expect three, so the expected sum is $6$. Dividing each sum by its respective $t$ makes both expected weighted sums $8$.

The point is not to force every actual mask draw to contain exactly its expected number of blanks. It is to account for the fact that higher masking rates expose more positions to the loss on average. For this linear schedule, the factor is $1/t$. With a different masking schedule, the weighting associated with a likelihood-bound derivation generally changes too; it is not safe to assume that $1/t$ always applies.

This also clarifies how to interpret the safe loss you implemented earlier. Dividing the batch’s total masked cross-entropy by its actual number of masked tokens gives a convenient, stable average loss per masked prediction. It handles batches with different mask counts, provided the zero-mask case is handled safely. But it is not, in general, the same numerical objective as first taking each example’s sum of masked losses, weighting it by $1/t$, and then averaging examples. Both can train a denoiser; only the latter, under the stated model and schedule assumptions, is the particular objective discussed in the paper.

Takeaways

An autoregressive model directly multiplies next-token probabilities to score text. A masked diffusion model instead learns to reconstruct hidden tokens across randomly chosen noise levels; its sequence distribution comes from an iterative reverse process. Masked-position cross-entropy is the practical prediction task, and schedule weighting connects a particular version of that task to a bound on negative log-likelihood.

Next, you’ll use the clean sequence, corrupted state, noise level, denoiser prediction, and reverse update to translate the linked paper’s notation into the tensors and operations in your Colab notebook.