Kuleshov Group | How to Build a Diffusion Language Model

Contents Masked Diffusion in a Nutshell Understanding Masked Diffusion: A Probabilistic Perspective Block Diffusion for Flexible-Length Generation Architectures: Encoder, Decoder, and Encoder-Decoder Iterative Refinement and Built-In Error Correction Accelerating Diffusion Sampling via Distillation Diffusion Enables Controllable Generation Post-Training Diffusion Language Models Biological and Scientific Domains Diffusion-Based Large Language Models Is Diffusion A Path Towards More Intelligent Models? A Scaling-Law Perspective Introduction: Autoregressive and Diffusion Language Models Two families of generative AI algorithms are widely used today. For continuous data such as images or video, the state-of-the-art approach is based on diffusion models. For discrete data such as text or code, the standard approach is instead autoregressive models. This article explores an alternative for discrete data, one built on the modern paradigm of diffusion. Mainstream language models are autoregressive: they generate tokens left-to-right, one at a time, each conditioned on the tokens before it. This approach is powerful, but it also has inherent limitations: No error correction: once a token is emitted it cannot be revised, so early mistakes compound. Generation is slow: producing a sequence takes as many steps as there are tokens, and does not naturally lend itself to fast, parallel generation. Causal attention: generation only ever looks backward, never at future context. Diffusion models take a different approach. Rather than producing text one token at a time, they generate the whole sequence at once, starting from an initial guess and iteratively refining it over a number of steps. This unlocks several advantages: generation can trade off speed and quality by using fewer or more steps, mistakes can be corrected along the way, and every step attends to bidirectional context. Autoregressive LLMs generate one token at a time, left to right, taking as many steps as there are tokens (top). Diffusion LLMs -- such as Gemma Diffusion shown here -- instead start from a rough, full-length draft and refine every position in parallel over a few rounds (bottom), rewriting the whole sequence at each step rather than emitting a single token. Figure credit: M. Grootendorst & Gemma Diffusion. Applying diffusion to language had long been an open problem. In 2024 the field reached a turning point, as diffusion models became competitive with autoregressive models on quality. By 2026, diffusion LLMs are a reality, with releases from leading industry labs — Mercury 2 (Inception Labs) , Gemma Diffusion (Google) , and Nemotron Diffusion (NVIDIA) . This article traces the ideas and papers that underlie these modern models. Background: Gaussian Diffusion Before introducing diffusion for language, we start with a brief overview of Gaussian diffusion for image generation. We will then build up discrete diffusion by analogy. Generating by iterative denoising The central concept underlying diffusion models is denoising. Instead of painting an image in one shot, a diffusion model produces images step by step, starting from pure random noise and removing a little of it at every step until a coherent image emerges. Generating an image through many small steps turns out to be far simpler than producing it all at once, and this is what makes diffusion models so effective. How does a model learn to denoise? The trick is to teach it by showing examples of noise being gradually transformed into an image. Diffusion achieves this via two complementary processes. First, a forward process takes a clean source image and turns it into pure noise, one step at a time. Second, a reverse process learns to invert this transformation, turning pure noise back into an image; it is trained on the image-to-noise trajectories produced by the forward process. Forward process The forward process takes a clean training image and produces a sequence of increasingly noisy images that trace a path from clean data to pure noise.

A Gaussian diffusion trajectory on an image. Reading left to right, the forward process gradually adds noise until a clean photo of a dog dissolves into pure static. This trajectory will serve as training data for the reverse process. Reverse process The reverse process is where the actual learning happens. We train a model to transform noise into images by following the steps produced by the forward process in reverse. Reverse path for Gaussian diffusion. Again going left to right, the generative reverse process is trained to reproduce the trajectory from the forward process in reverse, starting with noise, and reconstructing the original image. Concretely, given a noisy image, we train a machine learning model to separate the noise from the underlying image or, equivalently, to predict either the noise that was added or the clean image itself, since given the noisy input, knowing one determines the other.

The two processes that define diffusion. The forward process (top, left to right) turns clean data $x_0$ into complete noise $x_T$ by adding a little Gaussian noise at each step; the generative reverse process (bottom, right to left) is trained on this data to denoise $x_T$ back to $x_0$ using the same sequence of steps. After training on a sufficiently large set of trajectories, the model learns to generalize and generates new images starting from random samples of white noise. This forward/reverse recipe — corrupt data with noise, then learn to reverse the corruption one step at a time — is the blueprint for every diffusion model. Simple Masked Diffusion Models The main obstacle in bringing diffusion to language is deciding what "noise" should mean for discrete tokens. For example, the noise used in classical diffusion is Gaussian, and adding continuous Gaussian noise to categorical variables is not well-defined. Below we introduce one simple yet effective approach that defines noise via masking. Our group popularized this approach, and it now forms the basis of most open-source diffusion language models. Masked Diffusion in a Nutshell The easiest way to understand masked diffusion is as an unmasking transformer. We train the model by taking clean sequences, masking a random fraction of their tokens, and asking a bidirectional transformer to fill in the blanks. If you know BERT, this is essentially BERT with a randomized masking rate — but unlike BERT, the resulting model is generative. You can think of masked diffusion as a generative BERT. Training masked diffusion as an unmasking transformer. The forward masking process samples a random noise level $0 < t < 1$ and hides that fraction of the tokens in a clean datapoint $x$, producing a partially masked $z_t$; the model $x_\theta$ is then trained to reconstruct the original tokens. Figure credit: Sasha Rush. Once we trained the unmasking transformer, we can generate text by starting from a fully masked sequence and repeating two steps many times: Infilling: Ask the model to fill in every blank in the current sequence, yielding a rough guess of the clean tokens. Remasking: Randomly re-noise the infilled sequence by replacing tokens with masks, but keep a few more tokens unmasked than in the previous round. One sampling step. A denoising model fills in the masked positions of the current sequence $z_t$, then a random subset of those fresh predictions is re-masked to form $z_s$, the slightly-less-masked sequence at the next step. Iterating this fills the sequence in an arbitrary order. Figure credit: Sasha Rush. Each round leaves fewer positions masked, until the sequence converges to a clean sample from the model. Generation thus amounts to starting from a sequence full of blanks and gradually filling in words in an arbitrary order. Sampling in action on the opening of One Hundred Years of Solitude.