Masked Diffusion: Generative Models for Text

Masked Diffusion: Generative Models for Text#

Extra Credit

This page is extra credit for the machine learning unit. It builds directly on Generative Diffusion Models, so finish that page first. You don’t need to have done the restricted Boltzmann machine page.

On the continuous diffusion page we built a diffusion model for a continuous variable: we took a point \(x_0\), slowly added gaussian noise to it until nothing of the original was left, and then learned how to run that process backwards. This is how image generators like stable diffusion work, because a pixel is (essentially) a continuous number.

Our goal in this page is to do the same thing for text. By the end you will

have a conversation with a real language model that writes by undiffusing instead of by writing one word at a time left to right (like chatGPT does). We will give you the neural network and its weights. You will write the diffusion.
train your own (small) text diffusion model from scratch.

Everything you learned on the continuous diffusion page is going to carry over: there will be a forward process which destroys information, a fast way to jump to any step of it, an undiffusion step which needs a guess for \(x_0\), and a neural network that is trained to make that guess. Only one thing is going to have to change, and it is the first thing we need to deal with.

Getting Set Up#

You are going to need a few things installed and (for the second half of the page) an internet connection, because we will download a pre-trained network and some text to train on.

Packages. You need torch, transformers (version 4.48 or newer), requests, numpy and matplotlib.

On Google Colab all of these are already installed. This is the easiest way to do this page. Select Runtime > Change runtime type > T4 GPU before you start.
On your own machine run
```
pip install torch "transformers>=4.48" requests numpy matplotlib
```
(or put a ! in front of it and run it in a notebook cell). If you already have transformers and loading the network fails with an error about not recognizing modernbert, your version is too old: pip install -U transformers.

Downloads. You don’t have to download anything by hand. The code we give you fetches the following the first time you run it:

what	where it comes from	size	used in
the pre-trained chat network	dllm-hub/ModernBERT-large-chat-v0.1 on Hugging Face	0.9 GB	A Real Language Model
a piece of the TinyStories dataset	roneneldan/TinyStories on Hugging Face	4 MB	A tiny corpus
(optional) a bigger chat network	dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1 on Hugging Face	1.5 GB	Talking to it

The networks are saved in ~/.cache/huggingface so they are only downloaded once (on Colab they are downloaded again each time you start a new session). Neither needs an account or a login.

Hardware. Everything up to A Real Language Model runs instantly on any laptop. Chatting with the pre-trained network and training your own on stories are much faster with a GPU (Colab’s free one is plenty) but will run on a laptop CPU if you are patient.

Noise for Words#

To a computer, text is a list of integers. We chop the text into pieces called tokens (for now a token is just a word) and replace each token by its index in a list of all the allowed tokens called the vocabulary. Here is the vocabulary we will use for the first half of this page and a sentence written in it:

vocab = ["the", "cat", "box", "is", "alive", "empty", "nobody", "knows", "if", "or", "until", "somebody", "looks", "_"]
MASK = 13   # the index of "_"

def Encode(sentence):
    return torch.tensor([vocab.index(word) for word in sentence.split()])

def Decode(x):
    return " ".join(vocab[i] for i in x.tolist())

x_0 = Encode("nobody knows if the cat is alive or if the box is empty until somebody looks")
print(x_0)
print(Decode(x_0))

So our data \(x_0\) is no longer a single number. It is a list of \(L\) integers \(x_0 = (x_0^1, x_0^2, \ldots, x_0^L)\). (Here \(L=16\).)

On the continuous diffusion page our forward step was \(x_t = \sqrt{1-\beta_t}\, x_{t-1} + \sqrt{\beta_t} z\). That isn’t going to work here. Token 1 is cat and token 2 is box, but there isn’t any sense in which token 1.3 is “mostly a cat and a little bit of a box.” We can’t add a small gaussian to an integer label.

We need a different way of slowly destroying information. The simplest one you could imagine is to erase things. We add one extra token to our vocabulary, the mask token (we will print it as _), and our noise is going to be:

At step \(t\), each token that hasn’t already been erased gets replaced by the mask with probability \(\beta_t\). Once a token has been masked it stays masked forever.

This is called masked diffusion (or absorbing diffusion, because the mask state absorbs everything and never gives it back). Notice that this is still a Markov chain: \(x_t\) only depends on \(x_{t-1}\). On the continuous diffusion page, if you ran the forward process for long enough, every starting point ended up in the same place (a unit gaussian). Here, if you run long enough, every sentence ends up in the same place too: a sentence of all masks.

There is a nice physical picture here. Think of each token as an unstable nucleus. At every step each surviving nucleus has some probability \(\beta_t\) of decaying, the decays are independent of each other, and nothing ever un-decays.

The forward process#

Just like last time we need to pick the \(\beta_t\) for each of our \(T\) steps. As before define

\[\alpha_t = 1-\beta_t\]

and

\[\overline{\alpha_t}= \prod_{i=1}^t \alpha_i\]

Last time \(\overline{\alpha_t}\) told you how much of the original signal was left at step \(t\). It means the same thing here but now it is even easier to interpret: \(\alpha_t\) is the probability that a token survives step \(t\), so \(\overline{\alpha_t}\) is the probability that a token is still unmasked after \(t\) steps.

We are going to choose

\[\beta_t = \frac{1}{T-t+1} \qquad t = 1,2,\ldots,T\]

so that we take small steps early (at \(t=1\) each token is masked with probability \(1/T\)) and large steps late (at \(t=T\) everything that is left gets masked for sure).

To keep ourselves from getting confused by off-by-one errors, let’s store beta and alpha_bar as arrays of length \(T+1\) so that beta[t] really is \(\beta_t\). We will set beta[0]=0 (nothing happens at step 0) so that alpha_bar[0]=1.

Write a function

def MakeSchedule(T):
  # do stuff
  return beta, alpha_bar

(torch.cumprod will be useful.)

Grading

Plot \(\beta_t\) and \(\overline{\alpha_t}\) as a function of \(t\) for \(T=50\).

You should find that \(\overline{\alpha_t}\) is a straight line. Convince yourself (with pencil and paper - the product telescopes) that for our choice of \(\beta_t\)

\[\overline{\alpha_t} = 1-\frac{t}{T}\]

and plot this on top of your \(\overline{\alpha_t}\) to check. In other words, at step \(t\) a fraction \(t/T\) of the sentence is masked (on average).

Now write a function

def ForwardMask(x_0, t, beta):
  # do stuff
  return xs

which takes a sentence x_0 (a tensor of integers), a number of steps t to run for and the list beta, and runs the forward process one step at a time. It should return the list of sentences [x_0, x_1, ..., x_t] that it went through.

Some things that will be helpful:

torch.rand(len(x)) < p gives you a list of True/False which are each True with probability p.
If hit is such a list then x[hit] = MASK masks those tokens.
Make sure you do x = x.clone() before you change x. Otherwise you will also be changing every earlier sentence you’ve stored in your list (and x_0!)

Grading

Run your forward process on the sentence x_0 above with \(T=16\) and print out (using Decode) the sentence at every step. You should see it slowly dissolve into masks.
Now use \(T=50\). Run your forward process 1000 times and plot the average fraction of the sentence that is masked as a function of the step \(t\). Plot the theory curve \(1-\overline{\alpha_t}\) on top of it.

Faster masking#

Just like last time, it is going to be painful (when we get to training) to have to take every step of the forward process just to get a sample at step \(t\). Last time we could jump straight to step \(t\) because the sum of a bunch of gaussians is a gaussian.

Here it is even easier. Each token is doing its own thing independent of all the other tokens, and all that can happen to a token is that it is either still there or it has been masked. We already know the probability that a token is still there after \(t\) steps: it’s \(\overline{\alpha_t}\). So to jump to step \(t\):

Mask each token of \(x_0\) independently with probability \(1-\overline{\alpha_t}\).

Compare this to \(x_t = \sqrt{\overline{\alpha_t}}\, x_0 + \sqrt{1-\overline{\alpha_t}}\, z\) from last time. It’s the same idea: keep the signal with weight \(\overline{\alpha_t}\) and replace it with noise with weight \(1-\overline{\alpha_t}\).

Grading

Go ahead and write a

def ForwardMaskFast(x_0, t, alpha_bar):

which quickly generates a sample of \(x_t\) (just the one sentence, not the whole history).

To check that it gives the same distribution as ForwardMask, use \(T=50\) and \(t=20\). Run each function 5000 times and each time count how many of the 16 tokens are masked. Plot the two histograms of this count on top of each other.

You can also work out what this histogram should be: each of the \(L=16\) tokens is independently masked with probability \(p=1-\overline{\alpha_t}\), so the number masked follows the binomial distribution

\[P(n) = \binom{L}{n} p^n (1-p)^{L-n}\]

Plot this theory curve on top of your histograms. (math.comb(L,n) is the binomial coefficient.)

Undiffusing#

Our next step is to figure out how to undiffuse - i.e. to unmask.

As before, we have samples \(x_0\) from an initial distribution \(p_{init}(x_0)\) which we forward diffuse

x_0 = SamplePInit()
x_T = ForwardMaskFast(x_0, T, alpha_bar)

Last time this gave us a (largely gaussian) final distribution. This time we know exactly what the final distribution is: \(x_T\) is a sentence that is nothing but masks.

Now we’d like to start with \(x_T\) (a blank page) and step backwards \(x_T \rightarrow x_{T-1} \rightarrow \cdots \rightarrow x_0\) so that we end up with a sample from \(p_{init}\):

x = all masks
for t in range(T, 0, -1):
   x = Undiffuse(x, t, ...)
# x doesn't have to be any particular sentence, but if we histogram the x we get
# and the sentences from SamplePInit() those histograms should be the same

If \(p_{init}\) is “sentences that people write in English” then we will have built a machine that writes.

The plan is the same as last time. We can’t get \(p(x_{t-1}|x_t)\) directly, but we can get \(p(x_{t-1}|x_t,x_0)\). So to run backwards we take \(x_t\), guess the \(x_0\) that it came from, and then use our rule to get \(x_{t-1}\). The rule turns out to be really simple:

To go from \(x_t\) to \(x_{t-1}\):

a token that is not masked is left alone

a token that is masked is unmasked (i.e. replaced with our guess for \(x_0\) at that position) with probability

\[\frac{\overline{\alpha}_{t-1}-\overline{\alpha}_t}{1-\overline{\alpha}_t}\]

and otherwise stays masked.

For our choice of schedule, \(\overline{\alpha}_t = 1-t/T\), this probability is just

\[\frac{1}{t}\]

This makes some sense. If we are at step \(t\), we have \(t\) steps left to unmask everything and so (on average) we should do a fraction \(1/t\) of what is left right now. At the very last step (\(t=1\)) the probability is one and everything that is left gets unmasked.

Derivation

Let’s go ahead and derive the rule for \(p(x_{t-1}|x_t,x_0)\). Because every token is noised independently we can do this one token at a time.

The first steps are exactly the same as last time. Expanding \(p(x_t, x_{t-1}, x_0)\) in two ways using Bayes theorem and doing some algebra gives

\[p(x_{t-1}|x_t,x_0) = p(x_t| x_{t-1})\frac{p(x_{t-1}|x_0)}{p(x_t|x_0)}\]

(where we’ve used the fact that it’s a Markov chain to drop the \(x_0\) in the first term). All three pieces on the right hand side come from our forward process.

If the token is not masked at step \(t\), then it wasn’t masked at step \(t-1\) either (masks never go away) and so \(x_{t-1}=x_t=x_0\) with probability one.

If the token is masked at step \(t\), there are two things it could have been at step \(t-1\):

It was still \(x_0\) and it got masked on this very step.

\[p(x_{t-1}=x_0|x_t=\textrm{mask},x_0) = \beta_t \frac{\overline{\alpha}_{t-1}}{1-\overline{\alpha}_t}\]

It was already masked.

\[p(x_{t-1}=\textrm{mask}|x_t=\textrm{mask},x_0) = 1 \cdot \frac{1-\overline{\alpha}_{t-1}}{1-\overline{\alpha}_t}\]

Using \(\beta_t \overline{\alpha}_{t-1} = (1-\alpha_t)\overline{\alpha}_{t-1} = \overline{\alpha}_{t-1} - \overline{\alpha}_t\) the first of these becomes the probability in the rule above. You should check that the two probabilities add up to one.

Compare this to the continuous diffusion page, where \(p(x_{t-1}|x_t,x_0)\) was a gaussian with mean \(\tilde{\mu}_t\) and variance \(\tilde{\beta_t}\). Here it is just a (biased) coin flip.

Finally, we don’t know \(x_0\), so what do we do? If the token gets unmasked we need to say what it turns into, and the right thing to do is to sample it from our best knowledge of what \(x_0\) is at that position given everything we can see - i.e. from \(p(x_0^i | x_t)\). This is the job of the guessing function.

Let’s go ahead and write some functions now to get this working.

Write a function def Undiffuse(x_t, t, GuessX0, mask) which takes \(x_t\) and (stochastically) returns \(x_{t-1}\). It gets handed

the sentence x_t and the step t
a function GuessX0(x_t) which you will also write. Notice that last time your guess was a single number (the noise). This time, for every position in the sentence, the guess is a probability for each token in the vocabulary: GuessX0(x_t) should return a tensor p of shape (L, len(vocab)) where p[i,v] is the probability that the token at position i was originally token v.
the integer mask which is used for the mask token (for now this will be MASK; later, when we use a real model, it will be a different number)

Your Undiffuse should then

call p = GuessX0(x_t)
sample a guess for every position: guess = torch.multinomial(p, 1).squeeze(1) gives you a list of \(L\) tokens where the \(i\)’th one has been sampled from the probabilities p[i]
decide which positions get unmasked: those which are currently masked and pass a coin flip with probability \(1/t\)
copy the guess into those positions (and only those positions)

Undiffuse is generic. All of the knowledge about what sentences are likely is hiding in GuessX0. Over the course of this page we are going to hand your Undiffuse four different guessers and you will never have to change it.

For our first example, we are going to have a SamplePInit() that always returns the sentence the cat is alive. In other words our probability distribution is \(p(\texttt{the cat is alive})=1\).

When you write GuessX0, since you (secretly) know what \(x_0\) always is, you can probably do a really good job of guessing :). It should return probability one for the right word at every position and zero for everything else.

Grading

Once these functions are written, let’s work on the undiffusion. Use \(T=10\).

Run your undiffusion 5 times starting from a sentence of four masks (torch.full((4,), MASK)). For each run print the sentence at every step. You should see the words appear in a different (random) order each time.
Run your undiffusion 1000 times all the way to \(t=0\) and verify that you always get the cat is alive.
Run your ForwardMaskFast from \(t=0\) to \(t=5\) and your undiffusion from \(t=10\) to \(t=5\) (1000 times each) and plot histograms of the number of masked tokens in each case showing they are the same.

Two sentences#

We now want to make life a little harder. We are going to have a new probability distribution.

\[p(\texttt{the cat is alive}) = 0.8\]

\[p(\texttt{the box is empty}) = 0.2\]

\[p(\textrm{anything else}) = 0\]

Write the new SamplePInit function.

Now we also need a new GuessX0(x_t) function. Notice that (just like last time) this is much trickier and doesn’t always have an obvious answer. If I show you _ _ is _ you don’t really know which sentence it came from. It turns out the right thing to report is how probable each sentence is given what you can see. This is Bayes theorem:

Give each sentence a weight: its prior probability (0.8 or 0.2) if it agrees with every token of \(x_t\) that isn’t masked, and zero if it disagrees anywhere.
Divide the weights by their sum so they add to one. Call them \(w_A\) and \(w_B\).
Then at each position, the probability of sentence \(A\)’s word is \(w_A\) and the probability of sentence \(B\)’s word is \(w_B\). (At a position where both sentences have the same word, like the, it gets \(w_A+w_B=1\).)

So for example

\(x_t\)	\(w_A\)	\(w_B\)	guess for position 2
`_ _ is _`	0.8	0.2	`cat` 80%, `box` 20%
`_ _ is alive`	1	0	`cat` 100%
`the _ _ empty`	0	1	`box` 100%

This is the analog of the weighted average of \(z_a\) and \(z_b\) that you reported on the continuous diffusion page.

One technical note: it is possible that your GuessX0 gets handed something that agrees with neither sentence (like _ cat is empty), in which case both weights are zero and you’d be dividing by zero. If that happens just fall back to using the weights 0.8 and 0.2. (Whether this should ever happen is something you are about to find out.)

Grading

Do the undiffusion from this new probability distribution.

Using \(T=50\), generate 2000 sentences and make a bar chart of how often you got each one. You should get the cat is alive 80% of the time and the box is empty 20% of the time (and hopefully almost nothing else).
At each step of your undiffusion you are making a guess for \(x_0\). For ten runs (again at \(T=50\)) plot the guessed probability that position 2 is cat as a function of \(t\).
Now be greedy and try to do it faster. Repeat your bar chart for \(T=1, 2, 4, 10\). What happens?
Plot the fraction of your sentences which are not one of the two real sentences as a function of \(T\) (include \(T=50\) and use a log-log plot).

You should have found that with \(T=1\) almost a third of your sentences are things like the cat is empty and the box is alive which have probability zero in \(p_{init}\)! Your GuessX0 is exactly right - it’s impossible to do any better. So what went wrong?

The problem is that GuessX0 tells you about each position separately. From _ _ _ _ it (correctly) says that position 2 is cat 80% of the time and that position 4 is alive 80% of the time. What it doesn’t tell you is that those two things always happen together. If you unmask both positions in the same step, you are flipping two separate coins and 2 × 0.8 × 0.2 = 32% of the time they disagree.

If instead you only unmask one of them, the next call to GuessX0 sees it and becomes certain about the other. Taking many small steps is what turns a guesser that only knows about single positions into samples in which all the positions fit together. This is the entire reason for doing diffusion: at \(T=1\) you have a model that fills in every blank independently; at large \(T\) you have a model of whole sentences.

This is the same thing that happened on the continuous diffusion page in a slightly different guise. There, the step you took (a gaussian around \(\tilde{\mu}_t\)) was only the right answer if the step was small. The price for taking a small number of big steps is that you get the wrong distribution.

Why \(0.32/T\)?

Things only go wrong if the two positions that matter (2 and 4) get unmasked on the same step. If they are both still masked at step \(t\), the probability that they both get unmasked right now is \((1/t)^2\) and the probability neither does is \((1-1/t)^2\). The probability that both are still masked when we get to step \(t\) is \((t/T)^2\) (each one is still masked with probability \(t/T\)). So the probability that they come out together at step \(t\) is \((t/T)^2 (1/t)^2 = 1/T^2\) and summing over the \(T\) steps gives \(1/T\). When they do come out together they disagree with probability \(2\times 0.8\times 0.2=0.32\).

Being greedy#

Our Undiffuse makes two random choices: which positions to unmask (a coin flip for each) and what token to put there (a sample from the guess). It is the honest thing to do in the sense that (with enough steps) it gives you samples from exactly \(p_{init}\).

But there’s another thing you might want. You might not care about getting a fair sample of everything that could be said. You might just want one good sentence. Then there is an obvious thing to try:

What token? Don’t sample. At each position take the most probable token.
Which positions? Don’t flip coins. Unmask the positions where your guess is most confident - i.e. where that most probable token has the highest probability. We still want to be done in \(t\) more steps so unmask the top \(k\) where \(k\) is (number still masked)/\(t\) rounded up.

Write a def UndiffuseGreedy(x_t, t, GuessX0, mask) which does this. It takes exactly the same arguments as Undiffuse so you can swap one for the other.

Some things that will be helpful:

confidence, guess = p.max(dim=-1) gives you (for each position) the largest probability and which token has it
you don’t want to unmask things that aren’t masked, so set their confidence to -1
confidence.topk(k).indices gives you the positions of the k largest entries
math.ceil rounds up

Grading

Using the two-sentence GuessX0 and \(T=4\), run your greedy undiffusion a few times printing the sentence at each step.
Make the bar chart of outcomes for the greedy undiffusion with \(T=1\), \(T=2\) and \(T=4\).

You should find that greedy never produces a nonsense sentence, even with \(T=1\). But it has paid a price for this. What is it?

Keep this tradeoff in mind: the stochastic undiffusion is faithful (it samples everything in \(p_{init}\) with the right probabilities) but it needs lots of steps. The greedy one is safe but it always says the same thing.

A Real Language Model#

So far we’ve managed to do undiffusion where we know what the right guess for \(x_0\) is. For a distribution with two sentences in it we could just write GuessX0 down. The distribution we really care about - all the text that people write - has rather more sentences in it and we can’t.

So, just like last time, the guessing function is going to be replaced by a neural network. On the continuous diffusion page the network took in \((x_t, t)\) and returned its guess for the noise. Here the network takes in a sentence \(x_t\) that has some masks in it and returns, for every position, a probability for every token in the vocabulary. That is exactly the thing that your GuessX0 returned.

Later in this page you will train one of these yourself. But training a network that knows enough about the world to answer questions takes a lot of text and a lot of computer time, so first we are going to hand you one that has already been trained and get it to talk.

The network we are going to use is called ModernBERT. It is a transformer with about 400 million parameters (which makes it about three times bigger than GPT-2 and far smaller than the models behind chatGPT or Claude). The version we are using was trained with (essentially) the recipe you will use at the end of this page on a large collection of questions and answers.

We aren’t going to worry about what’s inside a transformer here. For us it is a black box that makes guesses. All the code that turns it into something that can write is going to be yours - in fact you have already written it.

The following code downloads the network and its weights (about 0.9 GB so it will take a minute the first time) and wraps it into a guessing function. If you are on Google Colab, select Runtime > Change runtime type > T4 GPU before you start as it will make everything faster. If the import fails, see Getting Set Up.

# GIVEN: load the pre-trained network and wrap it as a guessing function
from transformers import AutoTokenizer, AutoModelForMaskedLM

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
name = "dllm-hub/ModernBERT-large-chat-v0.1"
tokenizer = AutoTokenizer.from_pretrained(name)
bert = AutoModelForMaskedLM.from_pretrained(name).to(device).eval()
BERT_MASK = tokenizer.mask_token_id

def Tokens(text):
    """text -> tensor of token ids"""
    return torch.tensor(tokenizer.encode(text, add_special_tokens=False))

def Text(x):
    """tensor of token ids -> text"""
    return tokenizer.decode(x)

def GuessBERT(x_t):
    """The network's guess for x_0. Returns p of shape (L, vocabulary size):
    p[i, v] is the probability that position i is token v."""
    with torch.no_grad():
        logits = bert(x_t[None].to(device)).logits[0]
    return logits.float().softmax(dim=-1).cpu()

print(f"{sum(p.numel() for p in bert.parameters()) / 1e6:.0f} million parameters, running on {device}")
print(f"{len(tokenizer)} tokens in the vocabulary; the mask is token {BERT_MASK} which prints as {Text(torch.tensor([BERT_MASK]))}")

Tokens and blanks#

Let’s first get a feeling for what we’ve been given.

Tokens(text) and Text(x) play the role of your Encode and Decode. The vocabulary now has about 50,000 tokens in it rather than 14. Common words get their own token and less common words get built out of multiple pieces. You can put a mask into your text by typing [MASK].

GuessBERT(x_t) plays the role of your GuessX0. You give it a sentence with some masks in it and it gives you back the probability of every one of the 50,000 tokens at every position.

Here is a function that will show you the network’s top guesses at each masked position:

# GIVEN
def ShowGuesses(text, top=5):
    x = Tokens(text)
    p = GuessBERT(x)
    for i in torch.where(x == BERT_MASK)[0].tolist():
        best = p[i].topk(top)
        print(f"position {i}: " + ",  ".join(f"{Text(v)!r} {100 * q:.0f}%" for q, v in zip(best.values, best.indices)))

ShowGuesses("The capital of France is [MASK].")

Grading

Print out Tokens("Masked diffusion is surprisingly simple."). Then use Text on each of the tokens separately (i.e. Text(x[i:i+1])) to see how the sentence was chopped up.
Try ShowGuesses on four or five sentences of your own with one [MASK] in them. Find one where the network is very sure of the answer and one where it isn’t.
Now try this one, which has two masks in it:

ShowGuesses("[Question]\nName a country and its capital city.\n[/Question]\n\n[Answer]\nThe capital of [MASK] is [MASK].\n[/Answer]")

(The funny [Question] and [Answer] markers are how this network was shown questions and answers during its training. We will come back to this in a moment.)

Look at the top guess at each of the two positions. What sentence do you get if you take them both?
Fill in the first mask yourself with one of the countries the network suggested and run it again. What happens to the guess for the capital? Then try putting in a capital and masking the country.

You should have found that with both blanks open you get something like The capital of India is Bangkok. The network isn’t confused about geography: once you tell it the country, it is nearly certain about the capital.

This is exactly the the cat is empty problem. The network has in mind a bunch of perfectly good sentences (...of India is Delhi, ...of Thailand is Bangkok, ...of France is Paris) and when it reports its guesses one position at a time you can’t tell which country goes with which capital.

And you already know the fix. Don’t fill in everything at once. Undiffuse.

Grading

Take the two-mask sentence from above and turn it into tokens with Tokens. Then call your Undiffuse on it (you shouldn’t need to change your function - just hand it GuessBERT and BERT_MASK instead of GuessX0 and MASK).

Use \(T=1\) (so everything is filled in at once) and do this 10 times printing out the sentences you get. How many of the capitals are right?
Now use \(T=20\) and do it 10 times again.

Getting it to answer questions#

On the continuous diffusion page we ended by giving our model a prompt (left or right). To do this we had to feed the prompt into the network as an extra input. How do we give our text model a prompt, like a question we want it to answer?

It turns out that this is the place where masked diffusion is much nicer than what we did before. The prompt is just text. The answer is just text. So we write down a sentence where the first part is the question and the rest is masks

[Question]
Why is the sky blue?
[/Question]

[Answer]
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

and undiffuse. Your Undiffuse never touches a token that isn’t a mask, so the question stays put. And the network sees the whole thing - question included - when it makes its guess for the masks. The prompt is just the part of \(x_t\) that happens to already be unmasked.

Here is a function that writes a question in the format this network was trained on:

# GIVEN
def MakePrompt(question):
    return f"[SYS]\nYou are a helpful AI assistant.\n[/SYS]\n\n[Question]\n{question}\n[/Question]\n\n[Answer]\n"

print(MakePrompt("Why is the sky blue?"))

Write a function

def Chat(question, answer_length=64, T=64, step=Undiffuse, show=False):
  # do stuff
  return answer

which

Turns MakePrompt(question) into tokens
Sticks answer_length masks on the end of it. (torch.cat([a, b]) sticks two tensors together and torch.full((n,), BERT_MASK) gives you n masks.)
Undiffuses from \(t=T\) to \(t=0\) using step(x, t, GuessBERT, BERT_MASK)
Returns the text of the answer (i.e. only the tokens after the prompt)

If show=True it should also print out the answer-so-far every few steps so that you can watch it being written. When you are printing things out it looks nicer if you do Text(x).replace("[MASK]", "_").replace("[PAD]", "·").

There is one thing which will look odd. We have to decide how long the answer is before we start (we need to know how many masks to put down). But the network doesn’t have to use all of the space. When it is done it writes [/Answer] and then fills whatever is left over with a special padding token [PAD]. You should cut these off of what you return. This will do it: answer.replace("[PAD]", "").split("[/Answer")[0].strip()

Grading

Ask your model Why is the sky blue? with show=True and watch the answer appear. Is it writing the answer from left to right?
Ask it three more questions of your own.
Ask it the same question three times. Do you get the same answer?

Don’t expect too much. This is a small model, and with this first (completely faithful) way of undiffusing its answers will be sloppy: odd word choices, sentences that wander off, and plenty of things which aren’t true. But it should be writing English sentences which are about the thing you asked it. We will clean things up in the next two sections.

Being greedy with a real model#

You probably found that the answers were a bit rough around the edges: a repeated word here, a sentence that wanders off there. Our network is far from a perfect guesser. And the stochastic undiffusion (faithfully!) samples from the network’s guesses, including the 2% of the time that the guess is something silly. One silly token in an answer is enough to make it look bad and (worse) once it has been written it never gets erased so all the later guesses have to work around it.

This is where your greedy undiffusion earns its keep.

Grading

Your Chat takes a step argument. Use it to answer questions with UndiffuseGreedy instead.

Ask the same questions that you asked above. How do the answers compare?
Use show=True on one of them. What gets written first now?
Ask the same question three times. Do you get the same answer? Why?
Fewer steps. The appeal of diffusion for text is that you could write 64 tokens in fewer than 64 steps (a model that writes left to right has no choice but to take 64). For one question, get the answer from both of your undiffusers using \(T=64, 32, 16, 4,\) and \(1\). How few steps can you get away with? What do the mistakes look like when \(T\) is too small, and how does that connect to the cat is empty?

Temperature#

We now have two extremes. The stochastic undiffusion gives varied answers that are a bit sloppy and the greedy one gives a clean answer but always the same one. It would be nice to have a knob that goes between them, and a physicist knows exactly what that knob should be called.

Think of the guess \(p(v)\) for a token \(v\) as a Boltzmann distribution \(p(v) \propto \exp[-E(v)]\) at temperature one, where the “energy” of a token is \(E(v) = -\log p(v)\). Then the same system at temperature \(\tau\) has the distribution

\[p_\tau(v) \propto \exp[-E(v)/\tau] = p(v)^{1/\tau}\]

At \(\tau=1\) nothing changes. As \(\tau \rightarrow 0\) the system freezes into its ground state - the most probable token. For \(\tau > 1\) things get more random than the network intended.

Add an extra argument temperature=1.0 to the end of your Undiffuse. After you get p from the guesser, replace it by p**(1/temperature) (normalized so that each row sums to one again) before you sample from it.

To use this with your Chat you can do step=lambda x, t, guess, mask: Undiffuse(x, t, guess, mask, temperature=0.5).

Grading

Ask the same question at temperatures of 0.2, 0.5, 1.0 and 1.5 (a few times each). What do you see? What temperature would you use?

Talking to it#

You now have everything you need to sit down and have a conversation with your diffusion model.

# GIVEN: type a question, get an answer. Enter an empty line to stop.
def ChatLoop(**settings):
    while True:
        question = input("You: ")
        if question.strip() == "":
            break
        print("Model:", Chat(question, **settings), "\n")

# ChatLoop(step=UndiffuseGreedy)        # <-- uncomment me (once you've written Chat)

Grading

Have a conversation with your model using your favorite settings of step, T and answer_length. Include in what you hand in

a question it answers well
a question where it confidently tells you something false
a question whose answer needs more than 64 tokens (what should you change?)

A bigger model (optional)

Nothing about your code is specific to this network. Any network that takes in a sentence with masks and returns guesses for \(x_0\) will work with your Undiffuse, UndiffuseGreedy and Chat. If you want better answers, here is a somewhat larger network (600 million parameters, another 1.5 GB download) which you can swap in. It needs a recent version of transformers (pip install -U transformers if the import fails). You will need to re-run your Chat after running this.

from transformers import Qwen3ForCausalLM
name = "dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1"
tokenizer = AutoTokenizer.from_pretrained(name)
qwen = Qwen3ForCausalLM.from_pretrained(name).to(device).eval()
BERT_MASK = tokenizer.mask_token_id

def GuessBERT(x_t):
    L = len(x_t)
    see_everything = torch.zeros(1, 1, L, L, device=device)   # let every position look at every other position
    with torch.no_grad():
        logits = qwen(x_t[None].to(device), attention_mask=see_everything).logits[0]
    return logits.float().softmax(dim=-1).cpu()

def MakePrompt(question):
    return ("<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>\n<|im_start|>user\n" + question +
            "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n")

This network marks the end of its answer with <|im_end|> instead of [/Answer] so you will have to change that in your Chat.

If you have a big GPU, the same code runs LLaDA (GSAI-ML/LLaDA-8B-Instruct), an 8 billion parameter masked diffusion model whose answers are competitive with the left-to-right models of the same size.

Training#

So far someone else has trained the network for us. Now, we’d like to train our own.

Last time the training loop was

Pick a random \(t\)
Get the noisy data from your ForwardDiffusionFast function
Have your network guess the noise. Using a loss function, modify your network to match the noise more closely.

It is going to be exactly the same here:

Pick a random \(t\)
Mask each token of \(x_0\) with probability \(t/T\) (this is your ForwardMaskFast)
Have your network guess \(x_0\). Using a loss function, modify your network to make the correct tokens more probable.

The only thing we need to figure out is the loss function. Last time the network’s guess was a number and we used the mean squared error. This time the network’s guess is a list of probabilities and what we want is for the probability \(p_\theta(x_0^i | x_t)\) that it assigns to the correct token \(x_0^i\) to be large. The standard way to say “make this probability large” is to minimize \(-\log p_\theta(x_0^i | x_t)\) (this is called the cross entropy). It is zero if the network gave the right answer a probability of one and gets large as that probability gets small.

Putting it together, the loss for one sentence is

\[\mathcal{L} = \frac{T}{t} \sum_{i \textrm{ masked}} -\log p_\theta(x_0^i | x_t)\]

There are two things to notice. The sum is only over the positions which are masked (there is nothing to learn from a token that you can see). And there is a weight of \(T/t\) out front. At small \(t\) very few tokens are masked so there are very few terms in the sum. The weight makes up for it so that every \(t\) ends up counting about the same.

Where does the loss come from?

Here is a sketch. We’d like our undiffusion to produce the real sentences with high probability, which means we’d like each step of our undiffusion (which uses the network’s guess) to look as much as possible like the true reverse step \(p(x_{t-1}|x_t,x_0)\) (which knows \(x_0\)).

Consider one masked token at step \(t\). The true reverse step unmasks it into \(x_0^i\) with probability \(1/t\). Our undiffusion step unmasks it with probability \(1/t\) into a token sampled from the network’s guess. The mismatch between these two (measured by the KL divergence, the standard way to compare two probability distributions) works out to

\[\frac{1}{t}\left(-\log p_\theta(x_0^i|x_t)\right)\]

Adding this up over all the masked tokens and over all \(T\) steps gives \(\sum_{t=1}^T \frac{1}{t}\sum_{i \textrm{ masked}} -\log p_\theta(x_0^i | x_t)\). Rather than doing the sum over \(t\), we pick one \(t\) at random and multiply by \(T\) to make up for it. That’s our loss.

It is worth comparing with the continuous diffusion page where the same argument (compare the true reverse step to the one the network would take) gave a gaussian mismatch, which is why you ended up with a mean squared error.

The total has a nice meaning: averaged over sentences, \(\mathcal{L}\) is an upper bound on \(-\log p(x_0)\), i.e. on how surprised your model is by real sentences. The best any model can possibly do is the entropy of \(p_{init}\).

Here is the network we are going to train. It is a small transformer - the same kind of network as the one you were just chatting with but about 300 times smaller. As before we aren’t going to worry about what’s inside. (If you’d like to know, the extra credit page Building a Transformer from Scratch has you build this network yourself one piece at a time.) What you need to know is:

net = TinyDenoiser(vocab_size) makes a new (untrained) network
net(x) takes a tensor of integers of shape (batch, L) - i.e. a whole batch of sentences at once - and returns a tensor of shape (batch, L, vocab_size)
What it returns are not probabilities but logits: you get the probabilities by doing a softmax (i.e. \(p(v) \propto \exp[\textrm{logit}(v)]\)) over the last index. So a logit is just minus an energy.

Notice that (unlike last time) we never tell the network what \(t\) is. It doesn’t need it: it can see how much of the sentence is masked.

# GIVEN: a small transformer which takes tokens (batch, L) and returns logits (batch, L, vocab_size)
import torch.nn as nn
import torch.nn.functional as F

def Rotate(x):
    """Tell attention where each token sits by rotating pairs of features by an angle proportional to position."""
    L, d = x.shape[-2], x.shape[-1]
    frequency = 1.0 / 10000 ** (torch.arange(0, d, 2, device=x.device) / d)
    angle = torch.arange(L, device=x.device)[:, None] * frequency
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * angle.cos() - x2 * angle.sin(), x1 * angle.sin() + x2 * angle.cos()], dim=-1).flatten(-2)

class Block(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.heads = heads
        self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
        self.qkv, self.proj = nn.Linear(dim, 3 * dim), nn.Linear(dim, dim)
        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

    def forward(self, x):
        B, L, D = x.shape
        q, k, v = self.qkv(self.norm1(x)).view(B, L, 3, self.heads, D // self.heads).permute(2, 0, 3, 1, 4)
        # every position looks at every other position (a left-to-right model would only look left)
        attention = F.scaled_dot_product_attention(Rotate(q), Rotate(k), v)
        x = x + self.proj(attention.transpose(1, 2).reshape(B, L, D))
        return x + self.mlp(self.norm2(x))

class TinyDenoiser(nn.Module):
    def __init__(self, vocab_size, dim=128, heads=4, layers=4, max_length=64):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim)
        self.position = nn.Embedding(max_length, dim)
        self.blocks = nn.ModuleList([Block(dim, heads) for _ in range(layers)])
        self.norm = nn.LayerNorm(dim)
        self.out = nn.Linear(dim, vocab_size)
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.Embedding)):
                nn.init.normal_(m.weight, std=0.02)

    def forward(self, x):
        h = self.embed(x) + self.position(torch.arange(x.shape[1], device=x.device))
        for block in self.blocks:
            h = block(h)
        return self.out(self.norm(h))

def MakeGuesser(net, forbidden=()):
    """Wrap a network as a guessing function for Undiffuse: x_t of shape (L,) -> probabilities (L, vocab_size).
    `forbidden` is a list of tokens the guess is never allowed to be."""
    def Guess(x_t):
        with torch.no_grad():
            logits = net(x_t[None].to(next(net.parameters()).device))[0].cpu()
        logits[:, list(forbidden)] = -math.inf
        return logits.softmax(dim=-1)
    return Guess

Now write the loss function.

def DiffusionLoss(net, x_0, mask, T=100):
  # do stuff
  return loss

Here x_0 is a batch of sentences of shape (batch, L). You should

Pick a random \(t\) between 1 and \(T\) for each sentence: t = torch.randint(1, T + 1, (batch, 1), device=x_0.device)
Decide what to mask: masked = torch.rand(batch, L, device=x_0.device) < t / T (this is ForwardMaskFast for the whole batch at once: row b gets masked with probability t[b]/T)
Make \(x_t\): x_t = torch.where(masked, mask, x_0)
Run the network: logits = net(x_t)
Get \(-\log p_\theta(x_0^i | x_t)\) at every position. Pytorch will do the softmax, the log, and the picking out of the correct token for you: F.cross_entropy(logits.transpose(1, 2), x_0, reduction="none") returns a tensor of shape (batch, L).
Zero out the positions that weren’t masked, multiply by the weight \(T/t\), sum over the positions in each sentence, and average over the batch.

Two sentences (again)#

Before we go after real text, let’s go back to the place where we know what the right answer is: the distribution with the cat is alive 80% of the time and the box is empty 20% of the time.

Here is the general framework for the optimization (it is the same as last time).

net = TinyDenoiser(len(vocab), dim=32, heads=2, layers=2)
opt = torch.optim.Adam(net.parameters(), lr=1e-3)

for step in range(0, 1500):
  opt.zero_grad()
  x_0 = torch.stack([SamplePInit() for b in range(64)])   # a batch of 64 sentences
  loss = DiffusionLoss(net, x_0, MASK)
  loss.backward()
  opt.step()

Note that loss.item() gives you the loss.

Grading

Train your network.

Plot the loss as a function of training step. You may have to do a window averaging over approximately 100 steps to generate this plot. The best that any guesser can possibly do is the entropy of our distribution, \(-0.8\log 0.8 - 0.2 \log 0.2 = 0.50\). Draw this as a horizontal line on your plot. How close do you get?
Make your network into a guesser with GuessNet = MakeGuesser(net). Compare GuessNet(x_t) with your exact GuessX0(x_t) for _ _ _ _, for _ _ is alive and for the _ _ empty by printing out the probability each assigns to cat and to box at position 2. Has the network learned Bayes theorem? (Only compare at positions which are masked: the network never gets trained on what to say at a position it can already see, so its guesses there are meaningless. Your Undiffuse never uses them.)
Now you can use GuessNet in your Undiffuse in place of GuessX0. Go ahead and make this replacement and then generate 1000 sentences (with \(T=20\)) and make the bar chart of what you get.

A tiny corpus#

Now for some real text. We are going to train on a small piece of TinyStories, which is a collection of very short stories written using only words that a three year old would know. This is about the simplest English there is which means a small network has some chance of learning it in a few minutes.

The following code downloads the first 4 MB of it (about 5000 stories), chops it into words, and builds a vocabulary out of the 2000 most common ones. It gives you

story_vocab: the vocabulary. As before the last entry is the mask _. Any word not among the 2000 most common is replaced by <unk>.
STORY_MASK and STORY_UNK: the integers for those two tokens
GetBatch(batch, L): returns batch random chunks of text each L tokens long as a tensor of shape (batch, L). This is our SamplePInit.
StoryTokens(text) and StoryText(x): go back and forth between text and tokens

# GIVEN: a tiny corpus
import re
import requests
from collections import Counter

url = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt"
raw = requests.get(url, headers={"Range": "bytes=0-3999999"}).content[:4_000_000].decode("utf-8", errors="ignore")
words = re.findall(r"[a-z']+|[.,!?\"]", raw.lower().replace("<|endoftext|>", " "))

story_vocab = [w for w, n in Counter(words).most_common(2000)] + ["<unk>", "_"]
STORY_UNK, STORY_MASK = len(story_vocab) - 2, len(story_vocab) - 1
index = {w: i for i, w in enumerate(story_vocab)}

def StoryTokens(text):
    return torch.tensor([index.get(w, STORY_UNK) for w in re.findall(r"[a-z']+|[.,!?\"_]", text.lower())])

def StoryText(x):
    return " ".join(story_vocab[i] for i in x.tolist())

corpus = torch.tensor([index.get(w, STORY_UNK) for w in words])

def GetBatch(batch=64, L=32):
    start = torch.randint(len(corpus) - L, (batch,))
    return torch.stack([corpus[s:s + L] for s in start])

print(f"{len(corpus)} tokens of text, {len(story_vocab)} tokens in the vocabulary")
print(StoryText(GetBatch(1)[0]))

Now train a network on it. This is the same loop that you just wrote with three changes:

The network should be the default size: net = TinyDenoiser(len(story_vocab)).to(device)
Your data is x_0 = GetBatch().to(device) and your mask is STORY_MASK
Use 5000 steps. This takes a few minutes on a GPU. (On a laptop CPU it is something like ten times slower. Use the free GPU on Google Colab if you can. Otherwise you can get away with 1500 steps, which takes about ten minutes, but your stories will be rougher.)

It is worth printing out the loss every few hundred steps so that you know it’s working.

Since each sentence has \(L=32\) tokens in it, a more useful number to look at is the loss per token, \(\mathcal{L}/L\). A network which has no idea and guesses all 2002 tokens with equal probability would get \(\log 2002 = 7.6\). A network which has learned which words are common, but nothing about what order they go in, gets about 5.4.

Grading

Before you train, make a guesser for your untrained network and generate a story with it so that you can see where we are starting from:

GuessStory = MakeGuesser(net, forbidden=[STORY_MASK, STORY_UNK])
x = torch.full((32,), STORY_MASK)
for t in range(32, 0, -1):
    x = Undiffuse(x, t, GuessStory, STORY_MASK)
print(StoryText(x))

Then train your network.

Plot the loss per token as a function of training step (again with a window average). Draw horizontal lines at 7.6 and 5.4.
Generate three stories each using
- the stochastic undiffusion with \(T=32\)
- the stochastic undiffusion with \(T=2\)
- the stochastic undiffusion with \(T=32\) and a temperature of 0.7
- the greedy undiffusion with \(T=32\)

Which settings give the best stories? What goes wrong with the others? (You have seen each of these failures before.)

Prompts#

Finally, let’s give our own model a prompt. As you saw when you were chatting, there isn’t anything to it: the prompt is just the part of the sentence that starts out unmasked.

There is something here that a model that writes left to right can’t do. Nothing says that the unmasked part has to be at the beginning. We can give our model the beginning and the end of a story and ask it to fill in the middle, or give it a few words scattered around and have it write a story around them.

StoryTokens turns a _ into a mask so you can write things like

StoryTokens("once upon a time , there was a little dog named _ . _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ and they were all happy .")

Grading

Using your trained network (and whichever of your undiffusers you like the best), generate three stories for each of the following.

A story which starts once upon a time , there was a little dog named (with the rest of the 32 tokens masked)
A story which starts one day , lily went to the park . and ends and they were all happy .
A story which has the words dragon and cake somewhere in the middle of it (and everything else masked)

Where the Chat Model Came From#

You have now built every piece of a masked diffusion language model: the forward process, the undiffusion (two ways), and the training. The model you chatted with in the middle of this page was built out of exactly these pieces. There are only three differences between it and the model you just trained.

It is bigger and was trained for longer on more text. It has 400 million parameters instead of 1 million and its vocabulary has 50,000 tokens instead of 2000.
It didn’t start from scratch. It started from a network (BERT) which had already been trained on a large amount of text to fill in blanks when about a third of the tokens were masked. You now know what that is: it’s your training loop with \(t/T\) stuck at one value. Training it at every masking fraction is what turns something which can fill in a few blanks into something which can write on a blank page.
It was trained on questions and answers, and only the answers were masked. Each training example was a question followed by its answer, written with the [Question] and [Answer] markers you saw. In step 2 of the training loop, the tokens of the question were never masked. That’s it. The network learned to guess masked answers given visible questions, which is exactly the situation you put it in when you wrote Chat.

It is worth putting the two kinds of diffusion side by side.

	Continuous diffusion (the earlier page)	Masked diffusion (this page)
data \(x_0\)	a number	a list of tokens
noise	add a gaussian	replace tokens with a mask
\(\overline{\alpha_t}\)	how much of \(x_0\) is left in \(x_t\)	probability a token is still unmasked
jump to step \(t\)	\(\sqrt{\overline{\alpha_t}} x_0 + \sqrt{1-\overline{\alpha_t}} z\)	mask each token with probability \(1-\overline{\alpha_t}\)
where the forward process ends	a unit gaussian	all masks
what the network guesses	the noise \(z\) (equivalently \(x_0\))	probabilities for \(x_0\) at every position
undiffusion step	gaussian around \(\tilde{\mu}_t\)	unmask each masked token with probability \(1/t\)
loss	mean squared error on the noise	(weighted) \(-\log p\) of the masked tokens
prompt	an extra input to the network	the part of \(x_t\) that starts out unmasked
too few steps	wrong distribution	tokens that don’t fit together

Masked Diffusion: Generative Models for Text

Contents

Masked Diffusion: Generative Models for Text#

Getting Set Up#

Noise for Words#

The forward process#

Faster masking#

Undiffusing#

Two sentences#

Being greedy#

A Real Language Model#

Tokens and blanks#

Getting it to answer questions#

Being greedy with a real model#

Temperature#

Talking to it#

Training#

Two sentences (again)#

A tiny corpus#

Prompts#

Where the Chat Model Came From#