BLOCK DIFFUSION: INTERPOLATING BETWEEN AU-TOREGRESSIVE AND DIFFUSION LANGUAGE MODELS

Abstract

Diffusion language models offer unique benefits over autoregressive models due to their potential for parallelized generation and controllability, yet they lag in likelihood modeling and are limited to fixed-length generation. In this work, we introduce a class of block diffusion language models that interpolate between discrete denoising diffusion and autoregressive models. Block diffusion overcomes key limitations of both approaches by supporting flexible-length generation and improving inference efficiency with KV caching and parallel token sampling. We propose a recipe for building effective block diffusion models that includes an efficient training algorithm, estimators of gradient variance, and data-driven noise schedules to minimize the variance. Block diffusion sets a new state-of-the-art performance among diffusion models on language modeling benchmarks and enables generation of arbitrary-length sequences. We provide the code¹, along with the model weights and blog post on the project page: https://m-arriola.com/bd3lms

1. Introduction

Diffusion models are widely used to generate images (Ho et al., 2020; Dhariwal & Nichol, 2021; Sahoo et al., 2024b) and videos (Ho et al., 2022; Gupta et al., 2023), and are becoming increasingly effective at generating discrete data such as text (Lou et al., 2024; Sahoo et al., 2024a) or biological sequences (Avdeyev et al., 2023; Goel et al., 2024). Compared to autoregressive models, diffusion models have the potential to accelerate generation and improve the controllability of model outputs (Schiff et al., 2024; Nisonoff et al., 2024; Li et al., 2024; Sahoo et al., 2024c).

Discrete diffusion models currently face at least three limitations. First, in applications such as chat systems, models must generate output sequences of arbitrary length (e.g., a response to a user’s question). However, most recent diffusion architectures only generate fixed-length vectors (Austin et al., 2021; Lou et al., 2024). Second, discrete diffusion uses bidirectional context during generation and therefore cannot reuse previous computations with KV caching, which makes inference less efficient (Israel et al., 2025). Third, the quality of discrete diffusion models, as measured by standard metrics such as perplexity, lags behind autoregressive approaches and further limits their applicability (Gulrajani & Hashimoto, 2024; Sahoo et al., 2024a).

This paper makes progress towards addressing these limitations by introducing Block Discrete Denoising Diffusion Language Models (BD3-LMs), which interpolate between discrete diffusion and autoregressive models. Specifically, block diffusion models (also known as semi-autoregressive models) define an autoregressive probability distribution over blocks of discrete random variables (Si et al., 2022; 2023); the conditional probability of a block given previous blocks is specified by a discrete denoising diffusion model (Austin et al., 2021; Sahoo et al., 2024a).

Developing effective BD3-LMs involves two challenges. First, efficiently computing the training objective for a block diffusion model is not possible using one standard forward pass of a neural network and requires developing specialized algorithms. Second, training is hampered by the high variance of the gradients of the diffusion objective, causing BD3-LMs to under-perform autoregression even with a block size of one (when both models should be equivalent). We derive estimators of gradient variance, and demonstrate that it is a key contributor to the gap in perplexity between autoregression and diffusion. We then propose custom noise processes that minimize gradient variance and make progress towards closing the perplexity gap.

We evaluate BD3-LMs on language modeling benchmarks, and demonstrate that they are able to generate sequences of arbitrary length, including lengths that exceed their training context. In addition, BD3-LMs achieve new state-of-the-art perplexities among discrete diffusion models. Compared to alternative semi-autoregressive formulations that perform Gaussian diffusion over embeddings (Han et al., 2022; 2023), our discrete approach features tractable likelihood estimates and yields samples with improved generative perplexity using an order of magnitude fewer generation steps. In summary, our work makes the following contributions:

• We introduce block discrete diffusion language models, which are autoregressive over blocks of tokens; conditionals over each block are based on discrete diffusion. Unlike prior diffusion models, block diffusion supports variable-length generation and KV caching.

• We introduce custom training algorithms for block diffusion models that enable efficiently leveraging the entire batch of tokens provided to the model.

• We identify gradient variance as a limiting factor of the performance of diffusion models, and we propose custom data-driven noise schedules that reduce gradient variance.

• Our results establish a new state-of-the-art perplexity for discrete diffusion and make progress toward closing the gap to autoregressive models.

2. Background: Language Modeling Paradigms

Notation

We consider scalar discrete random variables with $V$ categories as ‘one-hot’ column vectors in the space $\mathcal{V}=\left\{\mathbf{x}\in {\{0,1\}}^V:{\sum }_i{\mathbf{x}}_i=1\right\}\subset {\Delta }^V$ for the simplex ${\Delta }^V$. Let the $V$ -th category denote a special [MASK] token, where $\mathbf{m}\in \mathcal{V}$ is its one-hot vector. We define ${\mathbf{x}}^{1:L}$ as a sequence of $L$ tokens, where ${\text{x}}^{\mathcal{l}}\in \mathcal{V}$ for all tokens $\mathcal{l}\in \{1,\dots ,L\}$, and use ${\mathcal{V}}^L$ to denote the set of all such sequences. Throughout the work, we simplify notation and refer to the token sequence as $\mathbf{x}$ and an individual token as ${\mathbf{x}}^{\mathcal{l}}$. Finally, let Cat (·;$p$) be a categorical distribution with probability $p\in {\Delta }^V$.

2.1. Autoregressive Models

Consider a sequence of $L$ tokens $\mathbf{x}=\left[{\mathbf{x}}^1,\dots ,{\mathbf{x}}^L\right]$ drawn from the data distribution $q(\mathbf{x})$. Autoregressive (AR) models define a factorized distribution of the form

$$\log p_{\theta }(\mathbf{x})=\sum _{\mathcal{l}=1}^L\log p_{\theta }\left({\mathbf{x}}^{\mathcal{l}}\mid {\mathbf{x}}^{<\mathcal{l}}\right),$$ (1)

where each $p_{\theta }\left({\mathbf{x}}^{\mathcal{l}}\mid {\mathbf{x}}^{<\mathcal{l}}\right)$ is parameterized directly with a neural network. As a result, AR models may be trained efficiently via next token prediction. However, AR models take $L$ steps to generate $L$ tokens due to the sequential dependencies.

2.2. Discrete Denoising Diffusion Probabilistic Models

Diffusion models fit a model $p_{\theta }(\mathbf{x})$ to reverse a forward corruption process $q$ (Sohl-Dickstein et al., 2015; Ho et al., 2020; Sahoo et al., 2024b). This process starts with clean data $\mathbf{x}$ and defines latent variables ${\mathbf{x}}_t=\left[{\mathbf{x}}_t^1,\dots ,{\mathbf{x}}_t^L\right]$ for $t\in [0,1]$, which represent progressively noisier versions of $\mathbf{x}$. Given a discretization into $T$ steps, we define $s(j)=(j-1)/T$ and $t(j)=j/T$. For brevity, we drop j from $t(j)$ and $s(j)$ below; in general, $s$ denotes the time step preceding $t$.

The D3PM framework (Austin et al., 2021) defines $q$ as a Markov forward process acting independently on each token ${\mathbf{x}}^{\mathcal{l}}:q\left({\mathbf{x}}_t^{\mathcal{l}}\mid {\mathbf{x}}_s^{\mathcal{l}}\right)=\text{Cat}\left({\mathbf{x}}_t^{\mathcal{l}};Q_t{\mathbf{x}}_s^{\mathcal{l}}\right)$ where $Q_t\in {ℝ}^{V\times V}$ is the diffusion matrix. The matrix $Q_t$ can model various transformations, including masking, random token changes, and related word substitutions.

An ideal diffusion model $p_{\theta }$ is the reverse of the process $q$. The D3PM framework defines $p_{\theta }$ as

$$p_{\theta }\left({\mathbf{x}}_s\mid {\mathbf{x}}_t\right)=\prod _{\mathcal{l}=1}^L{p}_{\theta }\left({\mathbf{x}}_s^{\mathcal{l}}\mid {\mathbf{x}}_t\right)=\sum _{\mathbf{x}}\left[\prod _{\mathcal{l}=1}^{L}q\left({\mathbf{x}}_s^{\mathcal{l}}\mid {\mathbf{x}}_t^{\mathcal{l}},{\mathbf{x}}^{\mathcal{l}}\right)p_{\theta }\left({\mathbf{x}}^{\mathcal{l}}\mid {\mathbf{x}}_t\right)\right],$$ (2)

where the denoising base model $p_{\theta }\left({\mathbf{x}}^{\mathcal{l}}\mid {\mathbf{x}}_t\right)$ predicts clean token ${\mathbf{x}}^{\mathcal{l}}$ given the noisy sequence ${\mathbf{x}}_t$, and the reverse posterior $q\left({\mathbf{x}}_s^{\mathcal{l}}\mid {\mathbf{x}}_t^{\mathcal{l}},\mathbf{x}\right)$ is defined following Austin et al. (2021) in Suppl. B.3.

The diffusion model $p_{\theta }$ is trained using variational inference. Let KL[·] denote the Kullback-Leibler divergence. Then, the Negative ELBO (NELBO) is given by (Sohl-Dickstein et al., 2015):

$$\mathcal{L}(\mathbf{x};\theta )={\mathbb{E}}_q\left[-\log p_{\theta }\left(\mathbf{x}\mid {\mathbf{x}}_{t(1)}\right)+\sum _{j=1}^T{D}_{\text{KL}}\left[q\left({\mathbf{x}}_{s(j)}\mid {\mathbf{x}}_{t(j)},\mathbf{x}\right)\Vert p_{\theta }\left({\mathbf{x}}_{s(j)}\mid {\mathbf{x}}_{t(j)}\right)\right]+D_{\text{KL}}\left[q\left({\mathbf{x}}_{t(T)}\mid \mathbf{x}\right)\Vert p_{\theta }\left({\mathbf{x}}_{t(T)}\right)\right]\right]$$ (3)

This formalism extends to continuous time via Markov chain (CTMC) theory and admits score-based generalizations (Song & Ermon, 2019; Lou et al., 2024; Sun et al., 2022). Further simplifications (Sahoo et al., 2024a; Shi et al., 2024; Ou et al., 2025) tighten the ELBO and enhance performance.

3. Block Diffusion Language Modeling

We explore a class of Block Discrete Denoising Diffusion Language Models (BD3-LMs) that interpolate between autoregressive and diffusion models by defining an autoregressive distribution over blocks of tokens and performing diffusion within each block. We provide a block diffusion objective for maximum likelihood estimation and efficient training and sampling algorithms. We show that for a block size of one, the diffusion objective suffers from high variance despite being equivalent to the autoregressive likelihood in expectation. We identify high training variance as a limitation of diffusion models and propose data-driven noise schedules that reduce the variance of the gradient updates during training.

3.1. Block Diffusion Distributions and Model Architectures

We propose to combine the language modeling paradigms in Sec. 2 by autoregressively modeling blocks of tokens and performing diffusion within each block. We group tokens in $\mathbf{x}$ into $B$ blocks of length $L$ with $B=L/L^{\prime }$ (we assume that B is an integer). We denote each block ${\mathbf{x}}^{(b-1)L^{\prime }:b{L}^{\prime }}$ from token at positions $(b-1)L^{\prime }$ to $b{L}^{\prime }$ for blocks $b\in \{1,\dots ,B\}$ as ${\mathbf{x}}^b$ for simplicity. Our likelihood factorizes over blocks as

$$\log p_{\theta }(\mathbf{x})=\sum _{b=1}^B\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}^{<b}\right),$$ (4)

and each $p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}^{<b}\right)$ is modeled using discrete diffusion over a block of $L^{\prime }$ tokens. Specifically, we define a reverse diffusion process as in (2), but restricted to block $b$:

$$p_{\theta }\left({\mathbf{x}}_s^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)=\sum _{{\mathbf{x}}^b}q\left({\mathbf{x}}_s^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^b\right)p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)$$ (5)

We obtain a principled learning objective by applying the NELBO in (3) to each term in (4) to obtain

$$-\log p_{\theta }(\mathbf{x})\le {\mathcal{L}}_{\text{BD}}(\mathbf{x};\theta ):=\sum _{b=1}^B\mathcal{L}\left({\mathbf{x}}^b,{\mathbf{x}}^{<b};\theta \right),$$ (6)

where each $\mathcal{L}\left({\mathbf{x}}^b,{\mathbf{x}}^{<b};\theta \right)$ is an instance of (3) applied to $\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}^{<b}\right)$. Since the model is conditioned on ${\mathbf{x}}^{<b}$, we make the dependence on ${\mathbf{x}}^{<b}$,$\theta$ explicit in $\mathcal{L}$. We denote the sum of these terms ${\mathcal{L}}_{\text{BD}}(\mathbf{x};\theta )$ (itself a valid NELBO).

Model Architecture

Crucially, we parameterize the $B$ base denoiser models $p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)$ using a single neural network ${\mathbf{x}}_{\theta }$. The neural network ${\mathbf{x}}_{\theta }$ outputs not only the probabilities $p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)$, but also computational artifacts for efficient training. This will enable us to compute the loss ${\mathcal{L}}_{\text{BD}}(\mathbf{x};\theta )$ in parallel for all $B$ blocks in a memory-efficient manner. Specifically, we parameterize ${\mathbf{x}}_{\theta }$ using a transformer (Vaswani et al., 2017) with a block-causal attention mask. The transformer ${\mathbf{x}}_{\theta }$ is applied to L tokens, and tokens in block b attend to tokens in blocks 1 to $b$. When ${\mathbf{x}}_{\theta }$ is trained, ${\mathbf{x}}_{\theta }^b\left({\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)$ yields $L^{\prime }$ predictions for denoised tokens in block $b$ based on noised ${\mathbf{x}}_t^b$ and clean ${\mathbf{x}}^{<b}$.

In autoregressive generation, it is normal to cache keys and values for previously generated tokens to avoid recomputing them at each step. Similarly, we use ${\mathbf{K}}^b$, ${\mathbf{V}}^b$ to denote the keys and values at block $b$, and we define ${\mathbf{x}}_{\theta }$ to support these as input and output. The full signature of ${\mathbf{x}}_{\theta }$ is

$${\mathbf{x}}_{\text{logits}}^b,{\mathbf{K}}^b,{\mathbf{V}}^b\leftarrow {\mathbf{x}}_{\theta }^b\left({\mathbf{x}}_t^b,{\mathbf{K}}^{1:b-1},{\mathbf{V}}^{1:b-1}\right):={\mathbf{x}}_{\theta }^b\left({\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right),$$ (7)

where ${\mathbf{x}}_{\text{logits}}^b$ are the predictions for the clean ${\mathbf{x}}^b$, and ${\mathbf{K}}^b$,${\mathbf{V}}^b$ is the key-value cache in the forward pass of ${\mathbf{x}}_{\theta }$, and ${\mathbf{K}}^{1:b-1}$,${\mathbf{V}}^{1:b-1}$ are keys and values cached on a forward pass of ${\mathbf{x}}_{\theta }$ over ${\mathbf{x}}^{<b}$ (hence the inputs ${\mathbf{x}}^{<b}$ and ${\mathbf{K}}^{1:b-1}$,${\mathbf{V}}^{1:b-1}$ are equivalent).

3.2. Efficient Training and Sampling Algorithms

Ideally, we wish to compute the loss ${\mathcal{L}}_{\text{BD}}(\mathbf{x};\theta )$ in one forward pass of ${\mathbf{x}}_{\theta }$. However, observe that denoising ${\mathbf{x}}_t^b$ requires a forward pass on this noisy input, while denoising the next blocks requires running ${\mathbf{x}}_{\theta }$ on the clean version ${\mathbf{x}}^b$. Thus every block has to go through the model at least twice.

Training

Based on this observation, we propose a training algorithm with these minimal computational requirements (Alg. 1). Specifically, we precompute keys and values ${\mathbf{K}}^{1:B}$,${\mathbf{V}}^{1:B}$ for the full sequence $\mathbf{x}$ in a first forward pass $\left(\varnothing ,{\mathbf{K}}^{1:B},{\mathbf{V}}^{1:B}\right)\leftarrow {\mathbf{x}}_{\theta }(\mathbf{x})$. We then compute denoised predictions for all blocks using ${\mathbf{x}}_{\theta }^b\left({\mathbf{x}}_t^b,{\mathbf{K}}^{1:b-1},{\mathbf{V}}^{1:b-1}\right)$. Each token passes through ${\mathbf{x}}_{\theta }$ twice.

Vectorized Training

Naively, we would compute the logits by applying ${\mathbf{x}}_{\theta }^b\left({\mathbf{x}}_t^b,{\mathbf{K}}^{1:b-1},{\mathbf{V}}^{1:b-1}\right)$ in a loop $B$ times. We propose a vectorized implementation that computes ${\mathcal{L}}_{\text{BD}}(\mathbf{x};\theta )$ in one forward pass on the concatenation ${\mathbf{x}}_{\text{noisy}}\oplus \mathbf{x}$ of clean data $\mathbf{x}$ with noisy data ${\mathbf{x}}_{\text{noisy}}={\mathbf{x}}_{t_1}^1\oplus \cdots \oplus {\mathbf{x}}_{t_B}^B$ obtained by applying a noise level $t_b$ to each block ${\mathbf{x}}^b$. We design an attention mask for ${\mathbf{x}}_{\text{noisy}}\oplus \mathbf{x}$ such that noisy tokens attend to other noisy tokens in their block and to all clean tokens in preceding blocks (see Suppl. B.6). Our method keeps the overhead of training BD3-LMs tractable and combines with pretraining to further reduce costs.

Sampling

We sample one block at a time, conditioned on previously sampled blocks (Alg 2). We may use any sampling procedure SAMPLE $\left({\mathbf{x}}_{\theta }^b,{\mathbf{K}}^{1:b-1},{\mathbf{V}}^{1:b-1}\right)$ to sample from the conditional distribution $p_{\theta }\left({\mathbf{x}}_s^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)$, where the context conditioning is generated using cross-attention with pre-computed keys and values ${\mathbf{K}}^{1:b-1}$,${\mathbf{V}}^{1:b-1}$. Similar to AR models, caching the keys and values saves computation instead of recalculating them when sampling a new block.

Notably, our block diffusion decoding algorithm enables us to sample sequences of arbitrary length, whereas diffusion models are restricted to fixed-length generation. Further, our sampler admits parallel generation within each block, whereas AR samplers are constrained to generate token-by-token.

Algorithm 1.

Block Diffusion Training

Input: datapoint $\mathbf{x}$, # of blocks $B$, forward

noise process $q_t(\cdot \mid \mathbf{x})$, model ${\mathbf{x}}_{\theta }$, loss ${\mathcal{L}}_{\text{BD}}$

repeat

Sample $t_1,\dots ,t_B~\mathcal{U}[0,1]$

$\forall b\in \{1,\dots ,B\}:{\mathbf{x}}_{t_b}^b~q_{t_b}\left(\cdot \mid {\mathbf{x}}^b\right)$

$\varnothing$, ${\mathbf{K}}^{1:B}$, ${\mathbf{V}}^{1:B}\leftarrow {\mathbf{x}}_{\theta }(\mathbf{x})$  ⊳ KV cache

$\forall b:{\mathbf{x}}_{\mathrm{logit}}^b$, $\varnothing$, $\varnothing \leftarrow {\mathbf{x}}_{\theta }^b\left({\mathbf{x}}_{t_b}^b,{\mathbf{K}}^{1:b-1},{\mathbf{V}}^{1:b-1}\right)$

Let ${\mathbf{x}}_{\text{logit}}\leftarrow {\mathbf{x}}_{\text{logit}}^1\oplus \cdots \oplus {\mathbf{x}}_{\text{logit}}^B$

Take gradient step on ${\nabla }_{\theta }{\mathcal{L}}_{\text{BD}}\left({\text{x}}_{\text{logit}};\theta \right)$

until converged

Algorithm 2.

Block Diffusion Sampling

Input: # blocks $B$, model ${\mathbf{x}}_{\theta }$, diffusion sampling algorithm SAMPLE

$\mathbf{x}$, $\mathbf{K}$, $\mathbf{V}\leftarrow \varnothing$    ⊳ output & KV cache

for $b=1$ to $B$ do

${\mathbf{x}}^b\leftarrow \mathrm{Sample}\left({\mathbf{x}}_{\theta }^b,{\mathbf{K}}^{1:b-1},{\mathbf{V}}^{1:b-1}\right)$

$\varnothing$, ${\mathbf{K}}^b$, ${\mathbf{V}}^b\leftarrow {\mathbf{x}}_{\theta }^b\left({\mathbf{x}}^b\right)$

$\mathbf{x}\leftarrow {\mathbf{x}}^{1:b-1}\oplus {\mathbf{x}}^b$

$(\mathbf{K},\mathbf{V})\leftarrow \left({\mathbf{K}}^{1:b-1}\oplus {\mathbf{K}}^b,{\mathbf{V}}^{1:b-1}\oplus {\mathbf{V}}^b\right)$

end for

return $\mathbf{x}$

4. Understanding Likelihood Gaps Between Diffusion & AR Models

4.1. Masked BD3-LMS

The most effective diffusion language models leverage a masking noise process (Austin et al., 2021; Lou et al., 2024; Sahoo et al., 2024a), where tokens are gradually replaced with a special mask token. Here, we introduce masked BD3-LMs, a special class of block diffusion models based on the masked diffusion language modeling framework (Sahoo et al., 2024a; Shi et al., 2024; Ou et al., 2025).

More formally, we adopt a per-token noise process $q\left({\mathbf{x}}_t^{\mathcal{l}}\mid {\mathbf{x}}^{\mathcal{l}}\right)=\mathrm{Cat}\left({\mathbf{x}}_t^{\mathcal{l}};{\alpha }_t{\mathbf{x}}^{\mathcal{l}}+\left(1-{\alpha }_t\right)\mathbf{m}\right)$ for tokens $\mathcal{l}\in \{1,\dots ,L\}$ where $\mathbf{m}$ is a one-hot encoding of the mask token, and ${\alpha }_t\in [0,1]$ is a strictly decreasing function in $t$, with ${\alpha }_0=1$ and ${\alpha }_1=0$. We employ the linear schedule where the probability of masking a token at time $t$ is $1-{\alpha }_t$. We adopt the simplified objective from Sahoo et al. (2024a); Shi et al. (2024); Ou et al. (2025) (the full derivation is provided in Suppl. B.3):

$$-\log p_{\theta }(\mathbf{x})\le {\mathcal{L}}_{\text{BD}}(\mathbf{x};\theta ):=\sum _{b=1}^B{\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\frac{{\alpha }_t^{\prime }}{1-{\alpha }_t}\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)$$ (8)

where ${\alpha }_t^{\prime }$ is the instantaneous rate of change of ${\alpha }_t$ under the continuous-time extension of (3) that takes $T\to \infty$. The NELBO is tight for $L^{\prime }=1$ but becomes a looser approximation of the true negative log-likelihood for $L^{\prime }\to L$ (see Suppl. B.5).

4.2. Case Study: Single Token Generation

Our block diffusion parameterization (8) is equivalent in expectation to the autoregressive NLL (1) in the limiting case where $L^{\prime }=1$ (see Suppl. B.4). Surprisingly, we find a two point perplexity gap between our block diffusion model for $L^{\prime }=1$ and AR when training both models on the LM1B dataset.

Although the objectives are equivalent in expectation, we show that the remaining perplexity gap is a result of high training variance. Whereas AR is trained using the cross-entropy of $L$ tokens, our block diffusion model for $L^{\prime }=1$ only computes the cross-entropy for masked tokens ${\mathbf{x}}_t^{\mathcal{l}}=\mathbf{m}\forall \mathcal{l}\in \{1,\dots L\}$ so that ${\mathbb{E}}_{t~\mathcal{U}[0,1]}q\left({\mathbf{x}}_t^{\mathcal{l}}=\mathbf{m}\mid {\mathbf{x}}^{\mathcal{l}}\right)=0.5$. Thus, training on the diffusion objective involves estimating loss gradients with 2x fewer tokens and is responsible for higher training variance compared to AR.

To close the likelihood gap, we train a BD3-LM for $L^{\prime }=1$ by designing the forward process to fully mask tokens, i.e. $q\left({\mathbf{x}}_t^{\mathcal{l}}=\mathbf{m}\mid {\mathbf{x}}^{\mathcal{l}}\right)=1$. Under this schedule, the diffusion objective becomes equivalent to the AR objective (Suppl. B.4). In Table 1, we show that training under the block diffusion objective yields the same perplexity as AR training. Empirically, we see that this reduces the variance of the training loss in Figure 2. We verify that tuning the noise schedule reduces the variance of the objective by measuring ${\mathrm{Var}}_{\mathbf{x},t}\left[{\mathcal{L}}_{\text{BD}}(\mathbf{x};\theta )\right]$ after training on 328M tokens: while training on the NELBO results in a variance of 1.52, training under full masking reduces the variance to 0.11.

Table 1:

Test perplexities for single-token generation (PPL; ↓) across 16B tokens on LM1B.

	PPL (↓)
AR	22.88
+ random batch size	24.37
BD3-LM $L^{\prime }=1$	≤ 25.56
+ tuned schedule	22.88

Figure 2:

Train NLLs for modeling the per-token likelihood on LM1B. Models are trained on 16B tokens. Training under the discrete diffusion NELBO, where half of the tokens in a batch are masked on average, has similar training variance to an AR model with a random batch size.

4.3. Diffusion Gap from High Variance Training

Next, we formally describe the issue of gradient variance in training diffusion models. Given our empirical observations for single-token generation, we propose an estimator for gradient variance that we use to minimize the variance of diffusion model training for $L^{\prime }\ge 1$. While the NELBO is invariant to the choice of noise schedule (Suppl. B.3), this invariance does not hold for our Monte Carlo estimator of the loss used during training. As a result, the variance of the estimator and its gradients are dependent on the schedule. First, we express the estimator of the NELBO with a batch size $K$. We denote a batch of sequences as $\mathbf{X}=\left[{\mathbf{x}}^{(1)},{\mathbf{x}}^{(2)},\dots ,{\mathbf{x}}^{(K)}\right]$, with each ${\mathbf{x}}^{(k)}\stackrel{\text{iid}}{~}q(\mathbf{x})$. We obtain the batch NELBO estimator below, where $t(k,b)$ is sampled in sequence $k$ and block $b$:

$${\mathcal{L}}_{\text{BD}}(\mathbf{X};\theta ):=l(\mathbf{X};\theta )=\frac{1}{K}\sum _{k=1}^K\sum _{b=1}^B\frac{{\alpha }_{t(k,b)}^{\prime }}{1-{\alpha }_{t(k,b)}}\log p_{\theta }\left({\mathbf{x}}^{(k),b}\mid {\mathbf{x}}_{t(k,b)}^{(k),b},{\mathbf{x}}^{(k),<b}\right)$$ (9)

The variance of the gradient estimator over $M$ batches for each batch ${\mathbf{X}}^m\forall m\in \{1,\dots ,M\}$ is:

$${\mathrm{Var}}_{\mathbf{X},t}\left[{\nabla }_{\theta }l(\mathbf{X};\theta )\right]\approx \frac{1}{M-1}\sum _{m=1}^M{‖{\nabla }_{\theta }l\left({\mathbf{X}}^m;\theta \right)-\frac{1}{M}\sum _{m=1}^M{\nabla }_{\theta }l\left({\mathbf{X}}^m;\theta \right)‖}_2^2$$ (10)

5. Low-Variance Noise Schedules for BD3-LMS

5.1. Intuition: Avoid Extreme Mask Rates

We aim to identify schedules that minimize the variance of the gradient estimator and make training most efficient. In a masked setting, we want to mask random numbers of tokens, so that the model learns to undo varying levels of noise, which is important during sampling. However, if we mask very few tokens, reconstructing them is easy and does not provide useful learning signal. If we mask everything, the optimal reconstruction are the marginals of each token in the data distribution, which is easy to learn, and again is not useful. These extreme masking rates lead to poor high-variance gradients: we want to learn how to clip them via a simple and effective new class of schedules.

5.2. Clipped Schedules for Low-Variance Gradients

We propose a class of “clipped” noise schedules that sample mask rates $1-{\alpha }_t~\mathcal{U}[\beta ,\omega ]$ for $0\le \beta$,$\omega \le 1$. We argue that from the perspective of deriving Monte Carlo gradient estimates, these schedules are equivalent to a continuous schedule where the mask probability is approximately 0 before the specified range such that $1-{\alpha }_{<\beta }\approx ϵ$ and approximately 1 after the specified range $1-{\alpha }_{>\omega }\approx 1-ϵ$. Consequently, ${\alpha }_t^{\prime }$ is linear within the range: ${\alpha }_t^{\prime }\approx 1/(\beta -\omega )$.

5.3. Data-Driven Clipped Schedules Across Block Sizes

As the optimal mask rates may differ depending on the block size $L^{ʹ}$, we adaptively learn the schedule during training. While Kingma et al. (2021) perform variance minimization by isolating a variance term using their squared diffusion loss, this strategy is not directly applicable to our variance estimator in Equation 10 since we seek to reduce variance across random batches in addition to random $t_b$.

Instead, we optimize parameters $\beta$,$\omega$ to directly minimize training variance. To limit the computational burden of the optimization, we use the variance of the estimator of the diffusion ELBO as a proxy for the gradient estimator to optimize $\beta$,$\omega :{\min }_{\beta ,\omega }{\mathrm{Var}}_{\mathbf{X},t}[\mathcal{L}(\mathbf{X};\theta ,\beta ,\omega )]$. We perform a grid search at regular intervals during training to find the optimal $\beta$,$\omega$ (experimental details in Sec. 6).

In Table 2, we show that variance of the diffusion NELBO is correlated with test perplexity. Under a range of “clipped” noise rate distributions, we find that there exists a unique distribution for each block size $L^{\prime }\in \{4,16,128\}$ that minimizes both the variance of the NELBO and the test perplexity.

Table 2:

Perplexities (PPLs; ↓) and variances of the NELBO $\text{Va}{\text{r}}_{\mathbf{X},t}\left[{\mathcal{L}}_{\text{BD}}(\mathbf{X};\theta )\right]$ (Var. NELBO; ↓). Models are trained on LM1B using a linear schedule for 65B tokens, then finetuned for 10B tokens.

$L^{\prime }$	$\mathcal{U}[0,.5]$		$\mathcal{U}[.3,.8]$		$\mathcal{U}[.5,1]$		$\mathcal{U}[0,1]$
	PPL	Var. NELBO	PPL	Var. NELBO	PPL	Var. NELBO	PPL	Var. NELBO
128	31.72	1.03	31.78	1.35	31.92	1.83	31.78	3.80
16	31.27	7.90	31.19	3.62	31.29	3.63	31.33	7.39
4	29.23	32.68	29.37	10.39	29.16	8.28	29.23	23.65

6. Experiments

We evaluate BD3-LMs across standard language modeling benchmarks and demonstrate their ability to generate arbitrary-length sequences unconditionally. We pre-train a base BD3-LM using the maximum block size $L^{\prime }=L$ for 850K gradient steps and fine-tune under varying $L^{\prime }$ for 150K gradient steps on the One Billion Words dataset (LM1B; Chelba et al. (2014)) and OpenWebText (OWT; Gokaslan et al. (2019)). Details on training and inference are provided in Suppl C.

To reduce the variance of training on the diffusion NELBO, we adaptively learn the range of masking rates by optimizing parameters $\beta$,$\omega$ as described in Section 5.3. In practice, we do so using a grid search during every validation epoch (after ∼5K gradient updates) to identify $\beta$,$\omega :{\min }_{\beta ,\omega }{\mathrm{Var}}_{\mathbf{X},t}[\mathcal{L}(\mathbf{X};\theta ,\beta ,\omega )]$. During evaluation, we report likelihood under uniformly sampled mask rates (8) as in Austin et al. (2021); Sahoo et al. (2024a).

6.1. Likelihood Evaluation

On LM1B, BD3-LMs outperform all prior diffusion methods in Table 3. Compared to MDLM (Sahoo et al., 2024a), BD3-LMs achieve up to 13% improvement in perplexity. We observe a similar trend on OpenWebText in Table 4.

Table 3:

Test perplexities (PPL; ↓) of models trained for 65B tokens on LM1B. Best diffusion value is bolded.

	PPL (↓)
Autoregressive
Transformer-X Base (Dai et al., 2019)	23.5
Transformer (Sahoo et al., 2024a)	22.83
Diffusion
D3PM (absorb) (Austin et al., 2021)	≤ 82.34
SEDD (Lou et al., 2024)	≤ 32.68
MDLM (Sahoo et al., 2024a)	≤ 31.78
Block diffusion (Ours)
BD3-LMs $L^{\prime }=16$	≤ 30.60
$L^{\prime }=8$	≤ 29.83
$L^{\prime }=4$	≤ 28.23

Table 4:

Test perplexities (PPL; ↓) on OWT for models trained for 524B tokens. Best diffusion value is bolded.

	PPL (↓)
AR (Sahoo et al., 2024a)	17.54
SEDD (Lou et al., 2024)	≤ 24.10
MDLM (Sahoo et al., 2024a)	≤ 22.98
BD3-LMs $L^{\prime }=16$	≤ 22.27
$L^{\prime }=8$	≤ 21.68
$L^{\prime }=4$	≤ 20.73

We also evaluate the ability of BD3-LMs to generalize to unseen datasets in a zero-shot setting, following the benchmark from Radford et al. (2019). We evaluate the likelihood of models trained with OWT on datasets Penn Tree Bank (PTB; (Marcus et al., 1993)), Wikitext (Merity et al., 2016), LM1B, Lambada (Paperno et al., 2016), AG News (Zhang et al., 2015), and Scientific Papers (Pubmed and Arxiv subsets; (Cohan et al., 2018)). In Table 5, BD3-LM achieves the best zero-shot perplexity on Pubmed, surpassing AR, and the best perplexity among diffusion models on Wikitext, LM1B, and AG News.

Table 5:

Zero-shot validation perplexities (↓) of models trained for 524B tokens on OWT. All perplexities for diffusion models are upper bounds.

	PTB	Wikitext	LM1B	Lambada	AG News	Pubmed	Arxiv
AR	81.07	25.32	51.14	52.13	52.11	48.59	41.22
SEDD	96.33	35.98	68.14	48.93	67.82	45.39	40.03
mdlm	90.96	33.22	64.94	48.29	62.78	43.13	37.89
BD3-LM $L^{\prime }=4$	96.81	31.31	60.88	50.03	61.67	42.52	39.20

6.2. Sample Quality and Variable-Length Sequence Generation

One key drawback of many existing diffusion language models (e.g,. Austin et al. (2021); Lou et al. (2024)) is that they cannot generate full-length sequences that are longer than the length of the output context chosen at training time. The OWT dataset is useful for examining this limitation, as it contains many documents that are longer than the training context length of 1024 tokens.

We record generation length statistics of 500 variable-length samples in Table 6. We continue sampling tokens until an end-of-sequence token [EOS] is generated or sample quality significantly degrades (as measured by sample entropy). BD3-LMs generate sequences up to ≈10× longer than those of SEDD (Lou et al., 2024), which is restricted to the training context size.

Table 6:

Generation length statistics from sampling 500 documents from models trained on OWT.

	Median # tokens	Max # tokens
OWT train set	717	131K
AR	4008	131K
SEDD	1021	1024
BD3-LM $L^{\prime }=16$	798	9982

We also examine the sample quality of BD3-LMs through quantitative and qualitative analyses. In Table 7, we generate sequences of lengths $L=1024$, 2048 and measure their generative perplexity under GPT2-Large. To sample $L=2048$ tokens from MDLM, we use their block-wise decoding technique (which does not feature block diffusion training as in BD3-LMs).

Table 7:

Generative perplexity (Gen. PPL; ↓) and number of function evaluations (NFEs; ↓) of 300 samples of lengths $L=1024$,2048. All models are trained on OWT. AR, SEDD, MDLM, BD3-LMs use 110M parameters and are trained on 524B tokens, while SSD-LM uses 400M parameters and is pre-trained on 122B tokens. Best diffusion value is bolded. We provide further details in Suppl. C.5.

	$L=1024$		$L=2048$
Model	Gen. PPL	NFEs	Gen. PPL	NFEs
AR	14.1	1K	13.2	2K
Diffusion SEDD	52.0	1K	–	–
MDLM	46.8	1K	41.3	2K
Block Diffusion SSD-LM L′ = 25	37.2	40K	35.3	80K
	281.3	1K	281.9	2K
BD3-LMs $L^{\prime }=16$	33.4	1K	31.5	2K
$L^{\prime }=8$	30.4	1K	28.2	2K
$L^{\prime }=4$	25.7	1K	23.6	2K

We also compare to SSD-LM (Han et al., 2022), an alternative block diffusion formulation. Unlike our discrete diffusion framework, SSD-LM uses Gaussian diffusion and does not support likelihood estimation. Further, BD3-LM adopts an efficient sampler from masked diffusion, where the number of generation steps (NFEs) is upper-bounded by L since tokens are never remasked (Sahoo et al., 2024a; Ou et al., 2025). For SSD-LM, we compare sample quality using $T=1\text{K}$ diffusion steps per block, matching their experimental setting (yielding ≥40K NFEs), and $T=25$ where NFEs are comparable across methods.

BD3-LMs achieve the best generative perplexities compared to previous diffusion methods. Relative to SSD-LM, our discrete approach yields samples with improved generative perplexity using an order of magnitude fewer generation steps. We also qualitatively examine samples taken from BD3-LM and baselines (AR, MDLM) trained on the OWT dataset; we report samples in Suppl. D. We observe that BD3-LM samples have higher coherence than MDLM samples and approach the quality of AR.

6.3. Ablations

We assess the impact of the design choices in our proposed block diffusion recipes, namely 1) selection of the noise schedule and 2) the efficiency improvement of the proposed training algorithm relative to a naive implementation.

Selecting noise schedules to reduce training variance

Compared to the linear schedule used in Lou et al. (2024); Sahoo et al. (2024a), training under

“clipped” noise schedules is the most effective for reducing the training variance which correlates with test perplexity. In Table 8, the ideal “clipped” masking rates, which are optimized during training, are specific to the block size and further motivate our optimization.

Table 8:

Effect of the noise schedule on likelihood estimation. We finetune BD3-LMs on 3B tokens from LM1B and evaluate on a linear schedule. For clipped schedules, we compare optimal clipping for $L^{\prime }=4$, 16.

Noise schedule	PPL	Var. NELBO
L’ = 4 Clipped  $\mathcal{U}[0.45,0.95]$	29.21	6.24
$\mathcal{U}[0.3,0.8]$	29.38	10.33
Linear $\mathcal{U}[0,1]$	30.18	23.45
Logarithmic	30.36	23.53
Square root	31.41	26.43
L’ = 16 Clipped  $\mathcal{U}[0.45,0.95]$	31.42	3.60
$\mathcal{U}[0.3,0.8]$	31.12	3.58
Linear $\mathcal{U}[0,1]$	31.72	7.62
Square	31.43	13.03
Cosine	31.41	13.00

Relative to other standard noise schedules (Chang et al., 2022), “clipped” masking achieves the best performance. As heavier masking is effective for the smaller block size $L^{\prime }=4$, we compare with logarithmic and square root schedules that also encourage heavy masking. As lighter masking is optimal for $L^{\prime }=16$, we compare with square and cosine schedules.

Efficiency of Training Algorithm

In the BD3-LM training algorithm (Sec. 3.2), we compute ${\mathbf{x}}_{\text{logit}}$ using two options. We may perform two forward passes through the network (precomputing keys and values for the full sequence $\mathbf{x}$, then computing denoised predictions), or combine these passes by concatenating the two inputs into the same attention kernel.

We find that a single forward pass is more efficient as we reduce memory bandwidth bottlenecks by leveraging efficient attention kernels (Dao et al., 2022; Dong et al., 2024), see Suppl. B.7. Instead of paying the cost of two passes through the network, we only pay the cost of a more expensive attention operation. Our vectorized approach has 20–25% speed-up during training relative to performing two forward passes.

7. Discussion and Prior Work

Comparison to D3PM

Block diffusion builds off D3PM (Austin et al., 2021) and applies it to each autoregressive conditional. We improve over D3PM in three ways: (1) we extend D3PM beyond fixed sequence lengths; (2) we study the perplexity gap of D3PM and AR models, identify gradient variance as a contributor, and design variance-minimizing schedules; (3) we improve over the perplexity of D3PM models. Our work applies to extensions of D3PM (He et al., 2022; Lou et al., 2024) including ones in continuous time (Campbell et al., 2022; Sun et al., 2022).

Comparison to MDLM

BD3-LMs further make use of the perplexity-enhancing improvements in MDLM (Sahoo et al., 2024a; Shi et al., 2024; Ou et al., 2025). We also build upon MDLM: (1) while Sahoo et al. (2024a) point out that their NELBO is invariant to the noise schedule, we show that the noise schedule has a significant effect on gradient variance; (2) we push the state-of-the-art in perplexity beyond MDLM. Note that our perplexity improvements stem not only from block diffusion, but also from optimized schedules, and could enhance standard MDLM and D3PM models.

Comparison to Gaussian

Diffusion Alternatively, one may perform diffusion over continuous embeddings of discrete tokens (Li et al., 2022; Dieleman et al., 2022; Chen et al., 2022). This allows using algorithms for continuous data (Song et al., 2020; Ho & Salimans, 2022), but yields worse perplexity (Graves et al., 2023; Gulrajani & Hashimoto, 2024).

Comparison to Semi-Autoregressive Diffusion

Han et al. (2022; 2023) introduced a block formulation of Gaussian diffusion. BD3-LMs instead extend Austin et al. (2021), and feature: (1) tractable likelihood estimates for principled evaluation; (2) faster generation, as our number of model calls is bounded by the number of generated tokens, while SSD-LM performs orders of magnitude more calls; (3) improved sample quality. AR-Diffusion (Wu et al., 2023) extends SSD-LM with a left-to-right noise schedule; Chen et al. (2025); Ye et al. (2024) apply to decision traces and videos; Hao et al. (2024); Kong et al. (2025) extend to latent reasoning. PARD (Zhao et al., 2024) applies discrete block diffusion to graphs. In contrast, we (1) interpolate between AR/diffusion performance; (2) support KV caching; (3) perform attention within noised blocks, whereas PARD injects new empty blocks.

Autoregressive diffusion models (Hoogeboom et al., 2021b;a) extend any-order AR models (AO-ARMs; Uria et al. (2014)) to support parallel sampling. Zheng et al. (2024) prove equivalence between MDLM and AO-ARM training. Further extensions of ARMs that compete with diffusion include iterative editing (Gu et al., 2019), parallel and speculative decoding (Gu et al., 2017; Santilli et al., 2023; Cai et al., 2024; Gloeckle et al., 2024), consistency training (Kou et al., 2024), guidance (Sanchez et al., 2023), and cross-modal extensions (Liu et al., 2023; Tian et al., 2025).

Limitations

Training BD3-LMs is more expensive than regular diffusion training. We propose a vectorized algorithm that keeps training speed within <2x of diffusion training speed; in our experiments, we also pre-train with a standard diffusion loss to further reduce the speed gap. Additionally, BD3-LMs generate blocks sequentially, and hence may face the same speed and controllability constraints as AR especially when blocks are small. Their optimal block size is task specific (e.g., larger for greater control). BD3-LMs are subject to inherent limitations of generative models, including hallucinations (Achiam et al., 2023), copyright infringement (Gokaslan et al., 2024), controllability (Schiff et al., 2024; Wang et al., 2023) and harmful outputs (Bai et al., 2022).

8. Conclusion

This work explores block diffusion and is motivated by two problems with existing discrete diffusion: the need to generate arbitrary-length sequences and the perplexity gap to autoregressive models. We introduce BD3-LMs, which represent a block-wise extension of the D3PM framework (Austin et al., 2021), and leverage a specialized training algorithm and custom noise schedules that further improve performance. We observe that in addition to being able to generate long-form documents, these models also improve perplexity, setting a new state-of-the-art among discrete diffusion models.

A. Block Diffusion NELBO

Below, we provide the Negative ELBO (NELBO) for the block diffusion parameterization. Recall that the sequence ${\mathbf{x}}^{1:L}=\left[{\mathbf{x}}^1,\dots ,{\mathbf{x}}^L\right]$ is factorized over $B$ blocks, which we refer to as $\mathbf{x}$ for simplicity, drawn from the data distribution $q(\mathbf{x})$. Specifically, we will factorize the likelihood over $B$ blocks of length $L^{ʹ}$, then perform diffusion in each block over $T$ discretization steps. Let $D_{\text{KL}}[\cdot ]$ to denote the Kullback-Leibler divergence, $t$,$s$ be shorthand for $t(i)=i/T$ and $s(i)=(i-1)/T\forall i\in [1,T]$. We derive the NELBO as follows:

$$\begin{gathered} -\log p_{\theta }(\mathbf{x})=-\sum _{b=1}^B\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}^{<b}\right) \\ =-\sum _{b=1}^B\log {\mathbb{E}}_q\frac{p_{\theta }\left({\mathbf{x}}_{t(1):t(T)}^b\mid {\mathbf{x}}^{<b}\right)}{q\left({\mathbf{x}}_{t(1):t(T)}^b\mid {\mathbf{x}}^b\right)} \\ =-\sum _{b=1}^B\log {\mathbb{E}}_q\frac{p_{\theta }\left({\mathbf{x}}_{t(T)}^b\mid {\mathbf{x}}^{<b}\right){\prod }_{i=1}^T{p}_{\theta }\left({\mathbf{x}}_{s(i)}^b\mid {\mathbf{x}}_{t(i)}^b,{\mathbf{x}}^{<b}\right)}{{\prod }_{i=1}^{T}q\left({\mathbf{x}}_{t(i)}^b\mid {\mathbf{x}}_{s(i)}^b\right)} \\ \le \sum _{b=1}^B\underset{{\mathcal{L}}_{\text{recons}}}{\underbrace{\left[-{\mathbb{E}}_q\right.\log p_{\theta }\left({\mathbf{x}}^b\left|{\mathbf{x}}_{t=\frac{1}{T}}^b\right.,{\mathbf{x}}^{<b}\right)}} \\ \underset{{\mathcal{L}}_{\text{diffusion}}}{\underbrace{+{\mathbb{E}}_{t\in \left\{\frac{2}{T},\dots ,\frac{T-1}{T},1\right\}}{\mathbb{E}}_{q}T{\text{D}}_{KL}\left(q\left({\mathbf{x}}_s^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^b\right)\Vert p_{\theta }\left({\mathbf{x}}_s^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right)}} \\ \left.\underset{{\mathcal{L}}_{\text{prior}}}{\underbrace{+{\text{D}}_{KL}\left(q\left({\mathbf{x}}_{t=1}^b\mid {\mathbf{x}}^b\right)\Vert p_{\theta }\left({\mathbf{x}}_{t=1}^b\right)\right)}}\right] \end{gathered}$$ (11)

B. Masked BD3-LMS

We explore a specific class of block diffusion models that builds upon the masked diffusion language modeling framework. In particular, we focus on masking diffusion processes introduced by Austin et al. (2021) and derive a simplified NELBO under this framework as proposed by Sahoo et al. (2024a); Shi et al. (2024); Ou et al. (2025).

First, we define the diffusion matrix $Q_t$ for states $i\in \{1,\dots ,V\}$. Consider the noise schedule function ${\alpha }_t\in [0,1]$, which is a strictly decreasing function in $t$ satisfying ${\alpha }_0=1$ and ${\alpha }_1=0$. Denote the mask index as $m=V$. The diffusion matrix is defined by Austin et al. (2021) as:

$${\left[Q_t\right]}_{ij}=\left\{\begin{array}{ll}1\hfill & \text{if}i=j=m\hfill \\ {\alpha }_t\hfill & \text{if}i=j\ne m\hfill \\ 1-{\alpha }_t\hfill & \text{if}j=m,i\ne m\hfill \end{array}\right.$$ (12)

The diffusion matrix for the forward marginal $Q_{t\mid s}$ is:

$${\left[Q_{t\mid s}\right]}_{ij}=\left\{\begin{array}{ll}1\hfill & \text{if}i=j=m\hfill \\ {\alpha }_{t\mid s}\hfill & \text{if}i=j\ne m\hfill \\ 1-{\alpha }_{t\mid s}\hfill & \text{if}j=m,i\ne m\hfill \end{array}\right.$$ (13)

where ${\alpha }_{t\mid s}={\alpha }_t/{\alpha }_s$

B.1. Forward Process

Under the D3PM framework (Austin et al., 2021), the forward noise process applied independently for each token $\mathcal{l}\in \{1,\dots L\}$ is defined using diffusion matrices $Q_t\in {ℝ}^{V\times V}$ as

$$q\left({\mathbf{x}}_t^{\mathcal{l}}\mid {\mathbf{x}}^{\mathcal{l}}\right)=\mathrm{Cat}\left({\mathbf{x}}_t^{\mathcal{l}};{\overline{Q}}_t{\mathbf{x}}^{\mathcal{l}}\right),\text{with}{\overline{Q}}_{t(i)}=Q_{t(1)}{Q}_{t(2)}\dots Q_{t(i)}$$ (14)

B.2. Reverse Process

Let $Q_{t\mid s}$ denote the diffusion matrix for the forward marginal. We obtain the reverse posterior $q\left({\mathbf{x}}_s^{\mathcal{l}}\mid {\mathbf{x}}_t^{\mathcal{l}},{\mathbf{x}}^{\mathcal{l}}\right)$ using the diffusion matrices:

$$q\left({\mathbf{x}}_s^{\mathcal{l}}\mid {\mathbf{x}}_t^{\mathcal{l}},{\mathbf{x}}^{\mathcal{l}}\right)=\frac{q\left({\mathbf{x}}_t^{\mathcal{l}}\mid {\mathbf{x}}_s^{\mathcal{l}},{\mathbf{x}}^{\mathcal{l}}\right)q\left({\mathbf{x}}_s^{\mathcal{l}}\mid {\mathbf{x}}^{\mathcal{l}}\right)}{q\left({\mathbf{x}}_t^{\mathcal{l}}\mid {\mathbf{x}}^{\mathcal{l}}\right)}=\mathrm{Cat}\left({\mathbf{x}}_s^{\mathcal{l}};\frac{Q_{t\mid s}{\mathbf{x}}_t^{\mathcal{l}}\odot Q_s^{\top }{\mathbf{x}}^{\mathcal{l}}}{{\left({\mathbf{x}}_t^{\mathcal{l}}\right)}^{\top }{Q}_t^{\top }{\mathbf{x}}^{\mathcal{l}}}\right)$$ (15)

where $\odot$ denotes the Hadmard product between two vectors.

B.3. Simplified NELBO for Masked Diffusion Processes

Following Sahoo et al. (2024a); Shi et al. (2024); Ou et al. (2025), we simplify the NELBO in the case of masked diffusion processes. Below, we provide the outline of the NELBO derivation; see the full derivation in Sahoo et al. (2024a); Shi et al. (2024); Ou et al. (2025).

We will first focus on simplifying the diffusion loss term ${\mathcal{L}}_{\text{diffusion}}$ in Eq. 11. We employ the SUBS-parameterization proposed in Sahoo et al. (2024b) which simplifies the denoising model $p_{\theta }$ for masked diffusion. In particular, we enforce the following constraints on the design of $p_{\theta }$ by leveraging the fact that there only exists two possible states in the diffusion process ${\mathbf{x}}_t^{\mathcal{l}}\in \left\{{\mathbf{x}}^{\mathcal{l}},\mathbf{m}\right\}\forall \mathcal{l}\in \{1,\dots ,L\}$.

1. Zero Masking Probabilities. We set $p_{\theta }\left({\mathbf{x}}^{\mathcal{l}}=\mathbf{m}\mid {\mathbf{x}}_t^{\mathcal{l}}\right)=0$ (as the clean sequence $\mathbf{x}$ doesn’t contain masks).

2. Carry-Over Unmasking. The true posterior for the case where ${\mathbf{x}}_t^{\mathcal{l}}\ne \mathbf{m}$ is $q\left.\left({\mathbf{x}}_s^{\mathcal{l}}={\mathbf{x}}_t^{\mathcal{l}}\mid {\mathbf{x}}_t^{\mathcal{l}}\ne \right.\mathbf{m}\right)=1$ (if a token is unmasked in the reverse process, it is never remasked). Thus, we simplify the denoising model by setting $p_{\theta }\left({\mathbf{x}}_s^{\mathcal{l}}={\mathbf{x}}_t^{\mathcal{l}}\mid {\mathbf{x}}_t^{\mathcal{l}}\ne \mathbf{m}\right)=1$.

As a result, we will only approximate the posterior $p_{\theta }\left({\mathbf{x}}_s^{\mathcal{l}}={\mathbf{x}}^{\mathcal{l}}\mid {\mathbf{x}}_t^{\mathcal{l}}=\mathbf{m}\right)$. Let ${\mathbf{x}}^{b,\mathcal{l}}$ denote a token in the $\mathcal{l}$-th position in block $b\in \{1,\dots ,B\}$. The diffusion loss term becomes:

$$\begin{gathered} {\mathcal{L}}_{\text{diffusion}}=\sum _{b=1}^B{\mathbb{E}}_t{\mathbb{E}}_{q}T\left[{\text{D}}_{\text{KL}}\left[q\left({\mathbf{x}}_s^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^b\right)\Vert p_{\theta }\left({\mathbf{x}}_s^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right]\right] \\ =\sum _{b=1}^B{\mathbb{E}}_t{\mathbb{E}}_{q}T\left[\sum _{\mathcal{l}=1}^{L^{\prime }}{\text{D}}_{\text{KL}}\left[q\left({\mathbf{x}}_s^{b,\mathcal{l}}\mid {\mathbf{x}}_t^{b,\mathcal{l}},{\mathbf{x}}^{b,\mathcal{l}}\right)\Vert p_{\theta }\left({\mathbf{x}}_s^{b,\mathcal{l}}\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right]\right] \end{gathered}$$

${\text{D}}_{\text{KL}}$ is simply the discrete-time diffusion loss for the block $b$; hence, from Sahoo et al. (2024a) (Suppl. B.1), we get:

$$\begin{gathered} =\sum _{b=1}^B{\mathbb{E}}_t{\mathbb{E}}_{q}T\left[\sum _{\mathcal{l}=1}^{L^{\prime }}\frac{{\alpha }_t-{\alpha }_s}{1-{\alpha }_t}\log p_{\theta }\left({\mathbf{x}}^{b,\mathcal{l}}\mid {\mathbf{x}}_t^{b,\mathcal{l}},{\mathbf{x}}^{<b}\right)\right] \\ =\sum _{b=1}^B{\mathbb{E}}_t{\mathbb{E}}_{q}T\left[\frac{{\alpha }_t-{\alpha }_s}{1-{\alpha }_t}\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right] \end{gathered}$$ (16)

Lastly, we obtain a tighter approximation of the likelihood by taking the diffusion steps $T\to \infty$ (Sahoo et al., 2024a), for which $T\left({\alpha }_t-{\alpha }_s\right)={\alpha }_t^{\prime }$:

$${\mathcal{L}}_{\text{diffusion}}=\sum _{b=1}^B{\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\left[\frac{{\alpha }_t^{\prime }}{1-{\alpha }_t}\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right]$$ (17)

For the continuous time case, Sahoo et al. (2024a) (Suppl. A.2.4) show the reconstruction loss reduces to 0 as ${\mathbf{x}}_{t(1)}^b~{\lim }_{T\to \infty }\mathrm{Cat}\left(.;{\mathbf{x}}_{t=\frac{1}{T}}^b\right)=\mathrm{Cat}\left(.;{\mathbf{x}}^b\right)$. Using this, we obtain:

$$\begin{gathered} {\mathcal{L}}_{\text{recons}}=-{\mathbb{E}}_q\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_{t(1)}^b,{\mathbf{x}}^{<b}\right) \\ =-\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_{t(1)}^b={\mathbf{x}}^b,{\mathbf{x}}^{<b}\right) \\ =0 \end{gathered}$$ (18)

The prior loss ${\mathcal{L}}_{\text{prior}}={\text{D}}_{KL}\left(q\left({\mathbf{x}}_{t=1}^b\mid {\mathbf{x}}^b\right)\Vert p_{\theta }\left({\mathbf{x}}_{t=1}^b\right)\right)$ also reduces to 0 because ${\alpha }_{t=1}=0$ which ensures $q\left({\mathbf{x}}_{t=1}^b\mid {\mathbf{x}}^b\right)=\text{Cat}(.;\mathbf{m})$ and $p_{\theta }\left({\mathbf{x}}_{t=1}^b\right)=\text{Cat}(.;\mathbf{m})$; see Sahoo et al. (2024a) (Suppl. A.2.4).

Finally, we obtain a simple objective that is a weighted average of cross-entropy terms:

$${\mathcal{L}}_{\text{BD}}(\mathbf{x};\theta )=\sum _{b=1}^B{\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\left[\frac{{\alpha }_t^{\prime }}{1-{\alpha }_t}\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right]$$ (19)

The above NELBO is invariant to the choice of noise schedule ${\alpha }_t$; see Sahoo et al. (2024a) (Suppl.

E.1.1).

B.4. Recovering the NLL from the NELBO for Single Token Generation

Consider the block diffuson NELBO for a block size of 1 where $L^{\prime }=1$,$B=L$. The block diffusion NELBO is equivalent to the AR NLL when modeling a single token:

$$\begin{gathered} -\log p(\mathbf{x})\le \sum _{b=1}^L{\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\left[\frac{{\alpha }_t^{\prime }}{1-{\alpha }_t}\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right] \\ \because {\alpha }_t^{\prime }=-1\text{and}{\alpha }_t=1-t\text{,} \\ =-\sum _{b=1}^L{\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\left[\frac{1}{t}\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right] \\ =-\sum _{b=1}^L{\mathbb{E}}_{t~[0,1]}\frac{1}{t}{\mathbb{E}}_q\left[\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)\right] \\ \text{Expanding}{\mathbb{E}}_q[.]\text{,} \\ =-\sum _{b=1}^L{\mathbb{E}}_{t~[0,1]}\frac{1}{t}\left[q\left({\mathbf{x}}_t^b=\mathbf{m}\mid {\mathbf{x}}^b\right)\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b=\mathbf{m},{\mathbf{x}}^{<b}\right)\right. \\ \left.+q\left({\mathbf{x}}_t^b={\mathbf{x}}^b\mid {\mathbf{x}}^b\right)\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b={\mathbf{x}}^b,{\mathbf{x}}^{<b}\right)\right] \end{gathered}$$ (20)

Recall that our denoising model employs the SUBS-parameterization proposed in Sahoo et al. (2024b). The “carry-over unmasking” property ensures that $\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b={\mathbf{x}}^b,{\mathbf{x}}^{<b}\right)=0$, as an unmasked token is simply copied over from from the input of the denoising model to the output. Hence, (20) reduces to following:

$$\begin{gathered} -\log p_{\theta }(\mathbf{x})\le -\sum _{b=1}^L{\mathbb{E}}_{t~[0,1]}\frac{1}{t}q\left({\mathbf{x}}_t^b=\mathbf{m}\mid {\mathbf{x}}^b\right)\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b=\mathbf{m},{\mathbf{x}}^{<b}\right) \\ \because q\left({\mathbf{x}}_t^b=\mathbf{m}\mid {\mathbf{x}}^b\right)=t\text{, we get:} \\ =-\sum _{b=1}^L{\mathbb{E}}_{t~[0,1]}\log p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b=\mathbf{m},{\mathbf{x}}^{<b}\right) \\ =-\sum _{b=1}^L\log p_{\theta }\left({\mathbf{x}}^b\mid \mathbf{m},{\mathbf{x}}^{<b}\right) \end{gathered}$$ (21)

For single-token generation $\left(L^{\prime }=1\right)$ we recover the autoregressive NLL.

B.5. Tightness of the NELBO

For block sizes $1\le K\le L$, we show that $-\log p(\mathbf{x})\le {\mathcal{L}}_K\le {\mathcal{L}}_{K+1}$. Consider $K=1$, where we recover the autoregressive NLL (see Suppl B.4):

$$\begin{gathered} {\mathcal{L}}_1=\sum _{b=1}^L\log {\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\frac{{\alpha }_t^{\prime }}{1-{\alpha }_t}{p}_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right) \\ =-\sum _{b=1}^L\log p_{\theta }\left({\mathbf{x}}^b\mid \mathbf{m},{\mathbf{x}}^{<b}\right) \end{gathered}$$ (22)

Consider the ELBO for block size $K=2$:

$${\mathcal{L}}_2=\sum _{b=1}^{L/2}\log {\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\frac{{\alpha }_t^{\prime }}{1-{\alpha }_t}{p}_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)$$ (23)

We show that ${\mathcal{L}}_1\le {\mathcal{L}}_2$, and this holds for all $1\le K\le L$ by induction. Let ${\mathbf{x}}^{b,\mathcal{l}}$ correspond to the token in position $\mathcal{l}\in \left[1,L^{\prime }\right]$ of block $b$. We derive the below inequality:

$$\begin{gathered} -\sum _{b=1}^L\log p_{\theta }\left({\mathbf{x}}^b\mid \mathbf{m},{\mathbf{x}}^{<b}\right)=-\sum _{b=1}^{L/2}\log {\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\frac{1}{1-{\alpha }_t}{p}_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right) \\ =-\sum _{b=1}^{L/2}\log {\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\prod _{i=1}^2\frac{1}{1-{\alpha }_t}{p}_{\theta }\left({\mathbf{x}}^{b,\mathcal{l}}\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right) \\ =-\sum _{b=1}^{L/2}\log \prod _{i=1}^2{\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\frac{1}{1-{\alpha }_t}{p}_{\theta }\left({\mathbf{x}}^{b,\mathcal{l}}\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right) \\ \le -\sum _{b=1}^{L/2}\sum _{i=1}^2\log {\mathbb{E}}_{t~[0,1]}{\mathbb{E}}_q\frac{1}{1-{\alpha }_t}{p}_{\theta }\left({\mathbf{x}}^{b,\mathcal{l}}\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right) \end{gathered}$$ (24)

B.6. Specialized Attention Masks

We aim to model conditional probabilities $p_{\theta }\left({\mathbf{x}}^b\mid {\mathbf{x}}_t^b,{\mathbf{x}}^{<b}\right)$ for all blocks $b\in [1,B]$ simultaneously by designing an efficient training algorithm with our transformer backbone. However, modeling all B conditonal terms requires processing both the noised sequence ${\mathbf{x}}_t^b$ and the conditional context ${\mathbf{x}}^{<b}$ for all $b$.

Rather than calling the denoising network $B$ times, we process both sequences simultaneously by concatenating them ${\mathbf{x}}_{\text{full}}\leftarrow {\mathbf{x}}_t\oplus \mathbf{x}$ as input to a transformer. We update this sequence ${\mathbf{x}}_{\text{full}}$ of length $2L$ tokens using a custom attention mask ${\mathcal{M}}_{\text{full}}\in {\{0,1\}}^{2L\times 2L}$ for efficient training.

The full attention mask is comprised of four $L\times L$ smaller attention masks:

$${\mathcal{M}}_{\text{full}}=\left[\begin{array}{cc}{\mathcal{M}}_{BD}& {\mathcal{M}}_{OBC}\\ \mathbf{0}& {\mathcal{M}}_{BC}\end{array}\right]$$

where ${\mathcal{M}}_{BD}$ and ${\mathcal{M}}_{OBC}$ are used to update the representation of ${\text{x}}_t$ and ${\mathcal{M}}_{BC}$ is used to update the representation of $\mathbf{x}$. We define these masks as follows:

• _{${\mathcal{M}}_{BD}$} (Block-diagonal mask): Self-attention mask within noised blocks ${\mathbf{x}}_t^b$

$${\left[{\mathcal{M}}_{BD}\right]}_{ij}=\left\{\begin{array}{ll}1\hfill & \text{if}i,j\text{are in the same block}\hfill \\ 0\hfill & \text{otherwise}\hfill \end{array}\right.$$

• _{${\mathcal{M}}_{OBC}$} (Offset block-causal mask): Cross-attention to conditional context ${\mathbf{x}}^{<b}$

$${\left[{\mathcal{M}}_{OBC}\right]}_{ij}=\left\{\begin{array}{ll}1\hfill & \text{if}j\text{belongs in a block before}i\hfill \\ 0\hfill & \text{otherwise}\hfill \end{array}\right.$$

• _{${\mathcal{M}}_{BC}$} (Block-causal mask): Attention mask for updating ${\mathbf{x}}^b$

$${\left[{\mathcal{M}}_{BC}\right]}_{ij}=\left\{\begin{array}{ll}1\hfill & \text{if}j\text{belongs in the same block as}i,\text{or a block before}i\hfill \\ 0\hfill & \text{otherwise}\hfill \end{array}\right.$$

We visualize an example attention mask for $L=6$ and block size $L^{\prime }=2$ in Figure 3.

Figure 3:

Example of Specialized Attention Mask

B.7. Optimized Attention Kernel with Flexattention

As Figure 3 demonstrates, our attention matrix is extremely sparse. We can exploit this sparsity to massively improve the efficiency of BD3-LMs.

FlexAttention (Dong et al., 2024) is a compiler-driven programming model that enables efficient implementation of attention mechanisms with structured sparsity in PyTorch. It provides a flexible interface for defining custom attention masks while maintaining high performance comparable to manually optimized attention kernels.

Below in Fig. 4 we define a block-wise attention mask, block_diff_mask, based on its definition as ${\mathcal{M}}_{\text{full}}\in {\{0,1\}}^{2L\times 2L}$ in Suppl. B.6. We fuse the attention operations into a single FlexAttention kernel designed to exploit the sparsity in our attention matrix to increase computational efficiency. By doing so, we perform the following optimizations:

Figure 4:

We can adapt the masking strategy from Fig. 3 to a FlexAttention compatible sparse masking function as above. This enables the creation of a customized JIT attention operation that uses significantly less memory with up to ≈5X speedup over the naive native scaled_dot_product_attention implementation in PyTorch (≥ 2.5) on a A5000 GPU with L = 1024 and batch size B = 16.

• Precomputed Block Masking: The create_block_mask utility generates a sparse attention mask at compile-time, avoiding per-step computation of invalid attention entries. Through sparsity-aware execution, FlexAttention kernels reduce the number of FLOPs in the attention computation.

• Reduced Memory Footprint: By leveraging block-level sparsity, the attention mechanism avoids full materialization of large-scale attention matrices, significantly reducing memory overhead. FlexAttention minimizes memory accesses by skipping fully masked blocks.

• Optimized Computation via torch.compile: The integration of torch.compile enables kernel fusion and efficient execution on GPUs by generating optimized Triton-based kernels. This efficiently parallelizes masked attention computations using optimized GPU execution paths.

This implementation exploits FlexAttention’s ability to dynamically optimize execution based on the provided sparsity pattern. By precomputing block-level sparsity and leveraging efficient kernel fusion, it enables scalable attention computation for long sequences.

Overall, this approach provides a principled method to accelerate attention computations while preserving structured dependency constraints. End-to-end, replacing FlashAttention kernels using a custom mask with FlexAttention kernels leads to ≈ 15% speedup in a model forward pass. We use a single A5000 for $L=1024$ and batch size $B=16$.

C. Experimental Details

We closely follow the same training and evaluation setup as used by Sahoo et al. (2024a).

C.1. Datasets

We conduct experiments on two datasets: The One Billion Word Dataset (LM1B; Chelba et al. (2014)) and OpenWebText (OWT; Gokaslan et al. (2019)). Models trained on LM1B use the bert-base-uncased tokenizer and a context length of 128. We report perplexities on the test split of LM1B. Models trained on OWT use the GPT2 tokenizer Radford et al. (2019) and a context length of 1024. Since OWT does not have a validation split, we leave the last 100k documents for validation.

In preparing LM1B examples, Sahoo et al. (2024a) pad each example to fit in the context length of $L=128$ tokens. Since most examples consist of only a single sentence, block diffusion modeling for larger block sizes $L^{\prime }>4$ would not be useful for training. Instead, we concatenate and wrap sequences to a length of 128. As a result, we retrain our autoregressive baseline, SEDD, and MDLM on LM1B with wrapping.

Similarly for OWT, we do not pad or truncate sequences, but concatenate them and wrap them to a length of 1024 similar to LM1B. For unconditional generation experiments in Section 6.2, we wish to generate sequences longer than the context length seen during training. However, Sahoo et al. (2024a) inject beginning-of-sequence and end-of-sequence tokens ([BOS], [EOS] respectively) at the beginning and end of the training context. Thus, baselines from Sahoo et al. (2024a) will generate sequences that match the training context size. To examine model generations across varying lengths in Section 6.2, we retrain our AR, SEDD, and MDLM baselines without injecting [BOS] and [EOS] tokens in the examples. We also adopt this preprocessing convention for training all BD3-LMs on OWT.

C.2. Architecture

The model architecture augments the diffusion transformer (Peebles & Xie, 2023) with rotary positional embeddings (Su et al., 2021). We parameterize our autoregressive baselines, SEDD, MDLM, and BD3-LMs with a transformer architecture from Sahoo et al. (2024a) that uses 12 layers, a hidden dimension of 768, and 12 attention heads. This corresponds to 110M parameters. We do not include timestep conditioning as Sahoo et al. (2024a) show it does not affect performance. We use the AdamW optimizer with a batch size of 512 and constant learning rate warmup from 0 to 3e-4 for 2.5K gradient updates.

C.3. Training

We train a base BD3-LM using the maximum context length $L^{\prime }=L$ for 850K gradient steps. Then, we fine-tune under varying $L^{\prime }$ using the noise schedule optimization for 150K gradient steps on the One Billion Words dataset (LM1B) and OpenWebText (OWT). This translates to 65B tokens and 73 epochs on LM1B, 524B tokens and 60 epochs on OWT. We use 3090, A5000, A6000, and A100 GPUs.

C.4. Likelihood Evaluation

We use a single Monte Carlo estimate for sampling t to evaluate the likelihood of a token block. We adopt a low-discrepancy sampler proposed in Kingma et al. (2021) that reduces the variance of this estimate by ensuring the time steps are more evenly spaced across the interval [0,1] following Sahoo et al. (2024a). In particular, we sample the time step for each block $b\in \{1,\dots ,B\}$ and sequence $k\in \{1,\dots ,K\}$ from a different partition of the uniform interval $t(k,b)\sim \mathcal{U}\left.\left[\frac{(k-1)B+b-1}{KB},\frac{(k-1)B+b}{KB}\right.\right]$.

This low-discrepancy sampler is used for evaluation. For training, each masking probability may be sampled from a “clipped” range $1-{\alpha }_t\sim \mathcal{U}[\beta ,\omega ]$. During training, we uniformly sample $t\in [0,1]$ under the low-discrepancy sampler. We then apply a linear interpolation to ensure that the masking probability is linear within the desired range: $1-{\alpha }_t=\beta +(\omega -\beta )t$.

When reporting zero-shot likelihoods on benchmark datasets from Radford et al. (2019) using models trained on OWT, we wrap all sequences to 1024 tokens and do not add [EOS] between sequences following Sahoo et al. (2024a).

C.5. Inference

Generative Perplexity

We report generative perplexity under GPT2-Large from models trained on OWT using a context length of 1024 tokens. Since GPT2-Large uses a context size of 1024, we compute the generative perplexity for samples longer than 1024 tokens using a sliding window with a stride length of 512 tokens.

Nucleus Sampling

Following SSD-LM (Han et al., 2022), we employ nucleus sampling for BD3-LMs and our baselines. For SSD-LM, we use their default hyperparameters p = 0.95 for block size $L^{\prime }=25$. For BD3-LMs, AR and MDLM, we use p = 0.9. For SEDD, we find that p = 0.99 works best.

Number of Diffusion Steps

In Table 7, BD3-LMs and MDLM use $T=5$ K diffusion steps. BD3-LMs and MDLM use efficient sampling by caching the output of the denoising network as proposed by Sahoo et al. (2024a); Ou et al. (2025), which ensures that the number of generation steps does not exceed the sample length $L$. Put simply, once a token is unmasked, it is never remasked as a result of the simplified denoising model (Suppl. B.3). We use MDLM’s block-wise decoding algorithm for generating variable-length sequences, however these models are not trained with block diffusion. We adopt their default stride length of 512 tokens.

SSD-LM (first row in Table 7) and SEDD use $T=5$ K diffusion steps. Since block diffusion performs $T$ diffusion steps for each block $b\in \{1,\dots ,B\}$, SSD-LM undergoes $BT$ generation steps. Thus to fairly compare with SSD-LM, we also report generative perplexity for $T=25$ diffusion steps so that the number of generation steps does not exceed the sequence length (second row in Table 7).

Improved Categorical Sampling of Diffusion Models

We employ two improvements to Gumbel-based categorical sampling of diffusion models as proposed by Zheng et al. (2024).

First, we use the corrected Gumbel-based categorical sampling from Zheng et al. (2024) by sampling 64-bit Gumbel variables. Reducing the precision to 32-bit has been shown to significantly truncate the Gumbel variables, lowering the temperature and decreasing the sentence entropy.

Second, Zheng et al. (2024) show that the MDLM sampling time scales with the diffusion steps $T$, even though the number of generation steps is bounded by the sequence length. For sample length $L$ and vocabulary size $V$, the sampler requires sampling $\mathcal{O}(TLV)$ uniform variables and performing logarithmic operations on them.

We adopt the first-hitting sampler proposed by Zheng et al. (2024) that requires sampling $\mathcal{O}(LV)$ uniform variables, and thus greatly improves sampling speed especially when $T\gg L$. The first-hitting sampler is theoretically equivalent to the MDLM sampler and leverages two observations: (1) the transition probability is independent of the denoising network, (2) the transition probability is the same for all masked tokens for a given $t$. Thus, the first timestep where a token is unmasked can be analytically sampled as follows (assuming a linear schedule where ${\alpha }_t=1-t$):

$$t_{n-1}=t_n{u}^{1/n},$$ (25)

where $n\in \{L,\dots ,1\}$ denotes the number of masked tokens, $u_n\sim \mathcal{U}[0,1]$ and $t_{n-1}$ corresponds to the first timestep where $n-1$ tokens are masked.

Variable-Length Sequence Generation

For arbitrary-length sequence generation using BD3-LMs and AR in Table 6, we continue to sample tokens until the following stopping criteria are met:

1. an [EOS] token is sampled

2. the average entropy of the the last 256-token chunk is below 4

where criterion 2 are necessary to prevent run-on samples from compounding errors (for example, a sequence of repeating tokens). We find that degenerate samples with low entropy result in significantly low perplexities under GPT2 and lower the reported generative perplexity. Thus, when a sample meets criterion 2, we regenerate the sample when reporting generative perplexity in Table 7.

D SAMPLES

Figure 1:

Block diffusion sequentially generates blocks of tokens by performing diffusion within each block and conditioning on previous blocks. By combining strength from autoregressive and diffusion models, block diffusion overcomes the limitations of both approaches by supporting variable-length, higher-quality generation and improving inference efficiency with KV caching and parallel sampling.

Figure 5:

Attention computation using FlexAttention with our proposed custom mask.

Figure 6:

Sample from MDLM (Sahoo et al., 2024a) of length L = 1024 and T = 5K diffusion steps. The generative perplexity of this sample under GPT2-Large is 69.26 and its entropy is 5.6.

Figure 7:

Sample from BD3-LM for block size L^′ = 16 of length L = 2031 under T = 5K diffusion steps (trained with a context length of L = 1024). The generative perplexity of this sample under GPT2-Large is 24.3, and its entropy is 5.5.

Figure 8:

Sample from an AR model (Sahoo et al., 2024a) with length L = 2003 (trained with a context length of L = 1024). The generative perplexity of this sample under GPT2-Large is 10.6 and its entropy is 5.5.