- AI, But Simple
- Posts
- (PRO) Diffusion Language Models (DLMs), Simply Explained
(PRO) Diffusion Language Models (DLMs), Simply Explained
AI, But Simple Issue #109

Diffusion Language Models (DLMs), Simply Explained (PRO)
AI, But Simple Issue #109
This week’s PRO issue includes a Colab notebook with practice exercises and solutions for diffusion generation techniques. We’ll dive deep into diffusion language models, and see how they work with interactive graphs.
Please do read the full article before attempting the code tutorial.
As always, the code in this issue can be executed on a normal home PC and does not require a large amount of VRAM.
The monthly interview question bank can be found here.
These are industry standard questions that are tested across major tech companies, AI research firms, and more. The questions are sorted by difficulty and topic.
Check back frequently for more questions, and let us know what type of questions you would like to see next!
Many language models you use today work autoregressively, generating tokens left to right given the context of what came before.
LLMs like GPT, LLaMA, and Claude all mostly follow this structure.
This idea is so deeply baked into how we build these models that most researchers also treat it as a given.
A team from Renmin University and Ant Group decided to ask a simple question: Does it always have to be this way?
Can language models instead be guided by the stochastic process of diffusion for generating text, much like for generating high-fidelity images and video?
Their answer is in the form of a model called LLaDA (Nie et al., 2025), and it is a new interesting direction.

What You’ll Learn
The reversal curse (the consequence of autoregressive reasoning)
How the LLaDA masked diffusion model works
Why LLaDA opts for bidirectional attention
LLaDA results and benchmarks and what they show
DiffusionGemma, Google’s new diffusion language model
Flow matching and embedding diffusion
What's Helpful to Know
Autoregressive models (ARM)
Models that generate tokens one at a time (left to right), where each token is conditioned on all previous tokens. This is how most current LLMs work.
More precisely, they factorize a joint distribution into a product of conditional distributions.


Originally popularized for images, they are generative models that learn to reverse a noising process. Data is iteratively corrupted, then a diffusion model learns to reverse this corruption.
For images, this means adding Gaussian noise and learning to remove it. For text, they use a discrete corruption process over tokens.

Masked diffusion model (MDM)
A diffusion model that corrupts text by randomly replacing tokens with [MASK] tokens, then trains a neural network to predict what the masked tokens originally were.
When Autoregressive Models Break Down
Before getting into LLaDA, it helps to understand with a simple example the one specific place where the left-to-right approach in autoregressive models can break down.
When you ask GPT-4 to complete the next line of a poem, it does fine, but if you ask it to complete the previous line, performance drops sharply.
This is not because the model lacks knowledge of poetry, but instead because the model was trained to predict what comes after and not what came before.
Essentially, it has no mechanism to sharply reason in the other direction.
This failure mode has a term in the AI research community called the reversal curse.

Fine-Tuning Test for Reversal Curse
This is considered a structural consequence of the left-to-right design. The model is great at "A comes before B" and worse at "B comes after A," even when those are logically the same fact.
LLaDA addresses this directly with striking results. On a reversal poem completion task, it outperformed GPT-4o by a significant margin, LLaDA scoring 45.6 while GPT-4o scored 34.3.
This all comes from a model with a fraction of the data and compute.
Note: Although today's frontier models benefit from more pretraining data, instruction tuning, retrieval, and reasoning, which make the reversal curse less visible, it’s still an interesting area for research and theory.
How Does LLaDA Actually Work?
The elegant core idea is that, instead of generating tokens one at a time from left to right, LLaDA starts with a completely masked sequence and gradually works on unmasking that full sequence.
Think of it as starting with a page of blacked-out text and progressively revealing words until the full response appears. You'll realize that this is very similar to diffusion tasks such as image generation.
There are two main processes:
Forward process
Reverse process
The forward process corrupts a clean piece of text. For a sequence of tokens x0, at time t between [0,1], each token is independently masked with probability t.
At t = 0, the text is untouched. At t = 1, everything is masked. At any point in between, roughly a fraction t of tokens are masked.
The reverse process starts from t = 1 (fully masked) and works backwards to t = 0 (fully revealed), gradually predicting what the masked tokens should be.

The model being trained is called the Masked Diffusion Model (MDM) (Sahoo et al., 2024), which uses the forward data masking process and trains a mask predictor to approximate its reverse process.
It takes a partially masked sequence xt and predicts all masked tokens simultaneously. This masked diffusion objective is shown below:

Here, 1[.] is the 0-1 indicator function. For each masked position, the model predicts the original token.
The 1/t term upweights positions that were harder to mask (low t means fewer masks, so each masked position was more selectively chosen).
This loss is provably an upper bound on the negative log-likelihood (log[pθ(x)]) of the model, which gives it the same theoretical training objective consistent with diffusion models.
It is important to note that unlike BERT, which uses a fixed masking ratio (0.15), LLaDA draws t uniformly from [0, 1] as a continuous random variable at every training step.
This means the model can mask from 0 to 100% of the tokens, seeing every possible masking level during training, which turns out to matter a lot when working at scale.
LLaDA Training and Inference
LLaDA is pre-trained on large text corpora (2.3T tokens) like most modern LLMs, but unlike autoregressive LLMs, it uses bidirectional attention instead of causal attention during training.

This is because the model has to iteratively refine full sequences at once during inference, not diffuse over tokens one-by-one.
For supervised fine-tuning (SFT), the same idea is adapted to instruction data: the prompt is kept fully visible, while only the response tokens are partially masked.
The model is then trained to reconstruct the missing response tokens conditioned on the prompt, using 4.5 million pairs.
At inference time, LLaDA will start from a fully masked response. Just like images, it begins its generation with full noise, or full masks.
The model then runs iterative denoising for N steps, repeatedly predicting and refining the missing tokens conditioned on the prompt until an output sequence is obtained.
The hyperparameter N can be tuned to balance output quality and inference time or computational cost. The higher the N, the higher the quality, the slower the inference speed.
What allows LLaDa to achieve better results at inference is the fact that the paper uses low-confidence remasking at each step where the tokens the model was least confident about get masked again, while certain ones don't.
What The Results Show
LLaDA 8B was pre-trained on 2.3 trillion tokens using a Transformer base that mirrors LLaMA 3 8B (which was trained on a much larger corpus of 15T tokens).
It shows competitive performance against LLaMA, both having identical MMLU scores (65.9 vs 65.4), being stronger on math (GSM8K: 70.3 vs 48.7) and Chinese benchmarks but weaker on tasks like Hellaswag.
After SFT, LLaDA improves at handling multi-turn dialogue, code, and constrained writing. It is the first non-autoregressive model to demonstrate chat-level capability at this scale.

LLaDA Benchmark Results From Paper
The low-confidence remasking design choice turned out to matter enormously and shows greatly improved results, achieving 70.0 on GSM8K, while random remasking achieves only 21.3.
The reason bidirectionality helps is straightforward. When LLaDA predicts a masked token, it sees every other unmasked token in the sequence (before and after), so there is no left-to-right constraint.
As mentioned, this turns out to be mathematically equivalent to training on all possible token orderings simultaneously.
An autoregressive model trained to predict "A then B" learns a specific direction, but LLaDA has no preferred direction at all, so the reversal curse simply does not apply.
DeepMind’s DiffusionGemma
Around the same time LLaDA was published, Google Deepmind unveiled Gemini Diffusion at I/O 2025, using the same diffusive generation approach in the form of a product.
Building on this direction, Google recently released DiffusionGemma (June 2026) to make this method more broadly accessible to the research community.
DiffusionGemma is Google DeepMind's new experimental open-weight text diffusion model.

Using diffusion for generation, the model generates at 1000+ tokens per second compared to 100 to 300 tokens per second for most frontier models today.
That’s around 4x faster generation than autoregressive equivalents.
The speed advantage most importantly comes from the structural property that LLaDA uses. The model generates entire blocks of tokens at once, and every denoising pass updates all positions at the same time.
DiffusionGemma also has a very accessible hardware footprint, with 26B parameters and a Mixture-of-Experts (MoE) architecture with only 3.8B active parameters.
It takes about 18GB of VRAM when 4-bit quantized, which fits on consumer GPUs (like the 4090).

The model scored 77.6% on MMLU Pro, almost matching Gemma 4’s 82.6% performance, which isn’t a bad performance tradeoff for a massive inference speedup.
But beyond all the important improvements and benchmarks, what DiffusionGemma also signals is emerging industrial commitment.
Google seems to be investing deeply into this research, offering fine-tuning tutorials, support, and NVIDIA hardware optimizations for both consumer and enterprise GPUs.
Embedding Diffusion and Flow Matching
Instead of token masking like LLaDA, there is another family of diffusion-based models called continuous embedding diffusion (e.g., Diffusion-LM, Plaid), which adds Gaussian noise directly to token embeddings.
The models corrupt the embeddings, denoise them, and round back to the nearest discrete token at the end.
The appeal with this is the access to the full continuous diffusion toolkit, such as score matching, ODE-based deterministic sampling, and classifier-free guidance that works exactly as it does in diffusion for image generation.
But the main problem is the rounding error, where the embedding space boundary between tokens is not clean, and small errors can flip a token to a semantically adjacent incorrect word.

On the other hand, Flow Matching offers a different route that sidesteps the issue mentioned above by specifying a fixed noise schedule and learning to undo it.
It specifies a direct transport path (like a straight line) between a source distribution and the data distribution where the model learns a velocity field vθ(xt, t) that moves samples along this path:

This helps as the paths are more linear than vanilla diffusion and training is more stable as fewer sampling steps are needed to get good results.

Discrete flow matching (Gat et al., 2024) applies flow matching to language by defining flows over probability distributions on the vocabulary rather than the tokens themselves.
At each timestep, the state becomes a probability distribution over tokens at each position, flowing from uniform uncertainty toward the dataset distribution. Since we’re working with a discrete categorical distribution, the rounding problem never arises.
Embedded Language Flows (ELF), Hu et al., 2026) uses the continuous route with flow matching by using more linear interpolation paths in embedding space.

Conceptual illustration of ELF
This significantly drops the accumulated rounding error and makes high-quality generation possible with much fewer steps than earlier embedding diffusion approaches.
Conclusion
The key takeaway from LLaDA, DiffusionGemma, discrete flow matching, and embedded language flows is that there is no single winning approach.
They are multiple research threads that converge to the same idea that autoregressive generation is not a law of nature but a design choice.
LLaDA established the research case, DiffusionGemma showed the industrial case, and flow matching further improved efficiency.
Here’s a special thanks to our biggest supporters:
If you enjoy our content, consider supporting us so we can keep doing what we do. Please share this with a friend!
Want to reach 8000+ ML engineers? Advertising Inquiries:
Feedback or inquiries? Send us an email at [email protected].
That’s it for this week’s issue of AI, but simple. See you next week!
—AI, but simple team