Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

We thank the reviewer for their constructive feedback! We address concerns and questions below.

Concern 1: Identifying Advantages of SAD3-LM relative to AR

Relative to AR, SAD3-LMs are amenable to faster & controllable sampling and better generalize to out-of-domain data distributions.

Advantage 1: Support for faster generation

SAD3-LMs have the potential to be faster than AR models because they can output multiple tokens in parallel, which yields an algorithm with improved arithmetic efficiency that better utilizes underlying hardware. We use the number of calls to the transformer (NFEs) as a proxy for inference cost; we see below that the NFEs and sample quality are comparable to AR, and SAD3-LM is much faster compared to previous semi-autoregressive model SSD-LM [1].

The reviewer is concerned that SAD3-LMs have significantly more NFEs than AR: we clarify that sampling from SAD3-LMs use fewer or equal NFEs as AR. We use an efficient sampler [2] that caches the predictions of the denoising model until a token is sampled, using at most $L$ NFEs. We clarify that SAD3-LM uses 1015 NFEs, and that the previously reported 10K NFEs is a typo (updated in line 431).

A SAD3-LM NFE can oftentimes be run in the same time as an AR NFE by better utilizing FLOPs. While an AR NFE is memory-bound, SAD3-LMs may perform more computation to generate multiple tokens using the same memory transfer. Hence, SAD3-LMs do not incur higher compute costs than AR.

Although we find that SAD3-LMs require the $\text{NFES} \approx L$ to achieve competitive quality, distillation for discrete diffusion enables sampling 8x faster than AR with improved quality (concurrent work presented in [3]). We expect SAD3-LMs will be faster than AR after distillation.

Advantage 2: Better generalization to new domains

SAD3-LMs better generalize to out-of-domain datasets in a zero-shot setting, whereas AR overfits to the training data. SAD3-LMs reach better zero-shot likelihoods than AR on out-of-domain datasets Lambada and Scientific Papers (Table 5). SAD3-LM’s discrete diffusion objective prevents overfitting [2]: since the loss gradient is updated only using masked tokens, the noise schedule effectively applies dropout.

Advantage 3: Support for improved controllability

SAD3-LM has potential for modular control of samples similar to SSD-LM by incorporating classifier guidance. Autoregressive models are inherently less amenable to control: they sample a token left-to-right, and a token that has been generated can never change. On the other hand, diffusion models have the potential to edit output tokens multiple times output using a bidirectional context, making it easier to enforce output constraints. While guidance for discrete diffusion is nontrivial and out of the scope of this manuscript (concurrent work presented in [4, 5]), we discuss its potential below (Concern 3).

---

Concern 2: “Complex training & sampling algorithms”

The reviewer is concerned about the complexity of the proposed training and sampling algorithms.

We highlight that our sampler has a simple structure: to sample each block, we may run any diffusion algorithm as a black-box (Alg. 2). To sample subsequent blocks, we simply feed in previously generated blocks as conditional context similar to AR. Thus, our sampling algorithm is as complicated as the chosen diffusion algorithm.

SAD3-LM training is also flexible to the choice in training algorithm: we may use any diffusion backbone to model each block $\mathbf{x}_\theta^b(\mathbf{x}_t^b)$. However, we find that conditioning the prediction on $\mathbf{x}^{<b}$ in Alg 1 has a significant improvement on likelihoods.

By using a diffusion objective on $\mathbf{x}_t$, SAD3-LMs enjoy benefits of diffusion models such as better generalization and long-term planning without being restricted to purely sequential generation.

---

Concern 3: Experiments on controlled text generation

We agree with the reviewer that controllability is one advantage of diffusion over AR. However, standard guidance techniques are not directly applicable because they require taking gradients with respect to the discrete latents $\mathbf{x}_t$ which are undefined. Although discrete guidance is out of the scope of this manuscript, recent work in this space which applies guidance over transition rate matrices [4] is promising to incorporate in future work. Compared to SSD-LM which performs guidance without principled likelihoods [1], we believe classifier guidance will be better suited for SAD3-LMs because they produce likelihoods that are necessary to compute classifier gradients.

We thank the reviewer for their thoughtful feedback! We address specific concerns and questions below.

Concern 1: Motivation to use diffusion models for language modeling

Relative to AR, SAD3-LMs are amenable to faster & controllable sampling and better generalize to out-of-domain data distributions.

Advantage 1: Support for faster generation

SAD3-LMs have the potential to be faster than AR models because they can output multiple tokens in parallel, which yields an algorithm with improved arithmetic efficiency that better utilizes underlying hardware. We use the number of calls to the transformer (NFEs) as a proxy for inference cost; we see below that the NFEs and sample quality are comparable to AR, and SAD3-LM is much faster compared to previous semi-autoregressive model SSD-LM [1].

The reviewer is concerned that SAD3-LMs have significantly more NFEs than AR: we clarify that sampling from SAD3-LMs use fewer or equal NFEs as AR. We use an efficient sampler [2] that caches the predictions of the denoising model until a token is sampled, using at most $L$ NFEs. We clarify that SAD3-LM uses 1015 NFEs, and that the previously reported 10K NFEs is a typo (updated in line 431).

A SAD3-LM NFE can oftentimes be run in the same time as an AR NFE by better utilizing FLOPs. While an AR NFE is memory-bound, SAD3-LMs may perform more computation to generate multiple tokens using the same memory transfer. Hence, SAD3-LMs do not incur higher compute costs than AR.

Although we find that SAD3-LMs require the $\text{NFES} \approx L$ to achieve competitive quality, distillation for discrete diffusion enables sampling 8x faster than AR with improved quality (concurrent work presented in [3]). We expect SAD3-LMs will be faster than AR after distillation.

Advantage 2: Better generalization to new domains

SAD3-LMs better generalize to out-of-domain datasets in a zero-shot setting, whereas AR overfits to the training data. SAD3-LMs reach better zero-shot likelihoods than AR on out-of-domain datasets Lambada and Scientific Papers (Table 5). SAD3-LM’s discrete diffusion objective prevents overfitting [2]: since the loss gradient is updated only using masked tokens, the noise schedule effectively applies dropout.

Advantage 3: Support for improved controllability

SAD3-LM has potential for modular control of samples by incorporating classifier guidance. Autoregressive models are inherently less amenable to control: they sample a token left-to-right, and a token that has been generated can never change. On the other hand, diffusion models have the potential to edit output tokens multiple times output using a bidirectional context, making it easier to enforce output constraints.

However, standard guidance techniques are not directly applicable because they require taking gradients with respect to the discrete latents $\mathbf{x}_t$ which are undefined. Although discrete guidance is out of the scope of this manuscript, concurrent work in this space [4,5] are promising to incorporate in future work, such as [4] which applies guidance over transition rate matrices.

---

Concern 2: Speed measurements for training, evaluation, and sampling

At the request of the reviewer, we provide speed measurements across 12K examples (batch size 32) from LM1B on a single A5000. SAD3-LM uses block size $L'=16$ and SSD-LM [1] uses block size $L'=25$.

Compared to SSD-LM, SAD3-LM training is ~4x more compute-efficient and 5x more token-efficient by training over the entire sequence, whereas SSD-LM only models one conditional term $p(\mathbf{x}_0^b | \mathbf{x}_t^{<b}, \mathbf{x}_t^b)$ per forward pass. A naive implementation that models each block sequentially (forward pass per block) is infeasible.

During sampling, SAD3-LMs use fewer or equal NFEs as AR. For $L$ tokens and $B$ blocks, $B \leq \text{SAD3-LM NFES} \leq L$ and $\text{AR NFEs} = L$. Although we find that SAD3-LMs require the $\text{NFES} \approx L$ to achieve competitive quality, we believe distillation is promising for up to 8x speedup relative to AR (noted in Concern 1: Advantage #1).

Does the variance of the gradient use the L2 Norm? Yes, we use the L2 norm for the variance of the gradient in Eq 9. The variance of the gradient estimator derived in Eq. 9 is computed by summing the component-wise variances of the gradient vectors across batches.

What about other methods to reduce the gradient variance? Besides our noise schedule optimization, we agree with the reviewer that other strategies may be useful. Our training algorithm already implements gradient clipping following [7], yet we still find that training variance still affects performance. Drawing multiple samples from $q(\mathbf{x}_t | \mathbf{x})$ may also be effective yet is costly to perform during training.

What is meant by “computational artifacts for efficient training”? We clarify that “computational artifacts for efficient training” in line 143 refers to the precomputed keys and values used for denoising a block conditioned on previous blocks (Line 163).

Training details? We have updated training details in the bottom of Suppl F.

Missing sign in the negative ELBO. We thank the reviewer for pointing out this detail, which is updated in Eq 5.

---

References:

[1] Han, X, et al. “SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control.” ACL 2023.

[2] Sahoo, S. S.., et al. “Simple and Effective Masked Diffusion Language Models.” NeurIPS 2024.

[3] Deschenaux, J., et al. “Beyond Autoregression: Fast LLMs via Self-Distillation Through Time.” arXiv preprint 2024.

[4] Nisonoff, H., et al. “Unlocking Guidance for Discrete State-Space Diffusion and Flow Models.” arXiv preprint 2024.

[5] Li, X., et al. “Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding.” arXiv preprint 2024.

[6] Austin, J, et al. “Structured Denoising Diffusion Models in Discrete State-Spaces.” NeurIPS 2021.

Additional questions

Experimental results for $L’=1$. At the request of the reviewer we provide the test perplexities of SAD3-LM for $L’=1$, finetuning from a base model where $L’=L$ that is pre-trained for 850K gradient steps following our other experiments. SAD3-LM for $L’=1$ is within a point of AR after only 150k finetuning steps on both OWT and LM1B.

We have compared the effect of the schedule on PPL and variance for SAD3-LM $L’=1$ in Table 1 and Line 261.

“What happens when you do one forward pass on noisy inputs and directly compute the loss that way?” We find that a single forward pass on the noisy sequence is a harder task and achieves significantly worse likelihoods compared to training on both $\mathbf{x}_t^b$ and conditional context $\mathbf{x}^{<b}$ (Alg 1). We train a model on $\mathbf{x}_t$ using a single block-causal forward pass. We report the following test perplexities from performing 50K gradient steps on LM1B for block size $L'=4$:

Likelihood calculation: forward process

What is T?

We define $T$ in Lines 82 & 93 as the number of time steps in the forward diffusion process. We tighten the likelihood bound by taking $T \rightarrow \infty$.

How is $x_{t(0)}$ sampled?

We clarify that $\mathbf{x}_{t(0)}$ is the first latent from the forward noise process: since $T \rightarrow \infty$, the first diffusion step is $t(0)$. This is exactly equal to the sample from the data distribution $\mathbf{x} \sim q$ so the reconstruction term in Eq. 3 evaluates to 0.

How are other $x_t$ sampled?

We show in line 207 that other $\mathbf{x}_t$ is sampled by randomly masking tokens with probability $1 - \alpha_t$ [2].

What $Q_t$ is used?

Since we can sample $\mathbf{x}_t$ by randomly masking tokens, materializing a transition matrix $Q_t$ is unnecessary and inefficient. We include the equivalent forward process using this $Q_t$ in Line 82 to show how our diffusion process relates with D3PM [6]. The form of $Q_t$ for absorbing state diffusion is:

[Q_t]_{ij} = \\begin{cases} 1 & \\text{if } i = j = m, \\\\ \\alpha_t & \\text{if } i = j \\neq m, \\\\ 1 - \\alpha_t & \\text{if } j = m, i \\neq m. \\end{cases}

Connecting the generic discrete diffusion ELBO to the SAD3-LM ELBO. We clarify that SAD3-LM is a meta-algorithm that uses a diffusion algorithm as a black-box. It leverages that algorithm's noising process $q$ and applies it in each block $\mathbf{x}^b$ to define a new probabilistic model. It also leverages the diffusion algorithm's learning and sampling processes. In SAD3-LMs, we use the discrete diffusion model from [2] to model a block of tokens. We have provided the full derivation of the ELBO from Eq 3 to Eq 7 under this diffusion model in Suppl A.

How do you compute and sample from the posterior? We clarify that the reverse posterior (Eq 2) is intractable (without assuming tokens are modeled independently) because we must marginalize over all possible sequences $\mathbf{x}$. This summation involves an exponential number of terms, making it computationally infeasible in high dimensions.

Since noise is applied to tokens independently and we model clean tokens independently, we make the posterior tractable by modeling it as:

$p\_{\theta} (\mathbf{x}\_s \mid \mathbf{x}\_t) = \prod_i q(x^i\_s \mid x^i\_t, x^i) p\_{\theta} (x^i \mid \mathbf{x}\_t)$

During training, we may compute the posterior probabilities in the ELBO as in Eq 7 which does not require sampling. At inference, we sample $\mathbf{x}\_s \sim p\_{\theta} (\mathbf{x}\_s | \mathbf{x}_t)$ using Gumbel-Max sampling as in standard language modeling.

Show that the AR and SAD3-LM $L’=1$ objectives are equivalent? We show their equivalence updated in Suppl B.

How is perplexity calculated? We calculate the PPL for SAD3-LM as follows for a sequence $\mathbf{x} = (x^1, \dots, x^L)$ and $B$ blocks using the ELBO in Eq 5: \text{PPL}(\mathbf{x}) = \text{exp} \Bigg\\{ - \frac{1}{L} \sum_{b=1}^B \mathcal{L} (\mathbf{x}^b, \mathbf{x}^{<b}, \theta) \Bigg\\} \

We thank the reviewer for their feedback! We answer specific questions below.

Concern 1: Statistical significance of results

The reviewer asks for comment on the statistical significance of the results. We show that our reported likelihoods are significantly significant by evaluating over three random seeds (for sampling $t$, $\mathbf{x}_t$) on OWT below.

SAD3-LMs achieve 16% improvement in likelihood from [1], and the best-performing SAD3-LM is within 17% of the AR perplexity.

We expect that the significance of the likelihood gap to AR will continue to improve with further training, as AR training overfits unlike diffusion models (since token masking acts as a form of dropout).

Additional questions

The reviewer asks for clarification on the notation in the training and sampling algorithms (Algs 1, 2). The reviewer is correct that in Alg 1, $\mathbf{x}_\theta (\mathbf{x})$ is a forward pass on the clean input to generate keys and values $\mathbf{K}^{1:B}$ and $\mathbf{V}^{1:B}$.

We clarify that during sampling, we may use a diffusion sampling algorithm $\text{Sample}$ as a black-box to generate $\mathbf{x}\_T^b, \mathbf{x}\_{T-1}^b, \dots, \mathbf{x}^b$. Thus, the noisy sequence $\mathbf{x}_t$ is handled internally in $\text{Sample}$ and not included in the rest of the sampling algorithm.

We thank the reviewer for pointing out two typos: we have updated Eq 5 and the surrounding text to show the negative ELBO, and fixed the grammatical error in line 202.

References

[1] Sahoo, S. S.., et al. “Simple and Effective Masked Diffusion Language Models.” NeurIPS 2024.

Concern 4: Comparison with Jacobi Decoding & Autoregressive Diffusion

We thank the reviewer for additional references on Jacobi decoding and autoregressive diffusion. Below we provide a comparison which we have included in Section 7 (line 431).

Jacobi Decoding

Jacobi decoding [5] is an AR inference technique that iteratively refines a random sequence and supports parallel generation of token blocks. Consistency LLMs [6] extend [5] to include a fine-tuning objective. We highlight key differences:

• Attention algorithm: [5,6] preserve causal masking from AR whereas SAD3-LMs may leverage more context (attending to tokens within a block and in previous blocks). SAD3-LMs use bidirectional context in the limiting case for $L’=L$
• Noising process: [5,6] use uniform noise whereas SAD3-LMs use masking. Masking noise has been shown to be superior in language modeling [7]
• Block-wise generation: [6] denoises the entire sequence, whereas our semi-autoregressive decoding may use clean conditional context $\mathbf{x}^{<b}$ to enhance predictions.
• Task: SAD3-LMs support unconditional generation whereas [5,6] assume a translation task

Autoregressive diffusion

Autoregressive Diffusion Models [8] generalize order-agnostic autoregressive models and discrete diffusion. In contrast, our approach is not order-agnostic as blocks are modeled autoregressively. AR-Diffusion [9] is a variant of SSD-LM [1] that uses a noise schedule that encourages sampling in a left-to-right manner. However, they sacrifice parallelism by assigning unique, per-token timesteps. [10] performs time series forecasting under a diffusion objective conditioned on observations at the previous timestep. However, their framework is autoregressive, so it is not amenable to parallel sampling or controllability.

---

Additional questions

“Why is AR worse on out-of-distribution settings? Why are the results bad on AG News?”

SAD3-LMs better generalize to out-of-domain datasets in a zero-shot setting, whereas AR overfits to the training data. SAD3-LM’s discrete diffusion objective prevents overfitting [1]: since the loss gradient is updated only using masked tokens, the noise schedule can be seen as applying dropout. We believe that diffusion models perform worse on AG News because they contain shorter sequences (~200 tokens on average), where a purely autoregressive model benefits from capturing short-range dependencies.

“The authors should define $\alpha’$. Why is comparison to SSD-LM under ablations?”

We thank the reviewer for pointing out these details. We included the definition of $\alpha’$ (line 212). We rearranged comparison to SSD-LM to be under Experiments as Section 6.3 (line 409).

---

References:

[1] Han, X, et al. “SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control.” ACL 2023.

[2] Sahoo, S. S., et al. “Simple and Effective Masked Diffusion Language Models.” NeurIPS 2024.

[3] Deschenaux, J, et al. “Beyond Autoregression: Fast LLMs via Self-Distillation Through Time.” arXiv preprint 2024.

[4] Dieleman, S, et al. “Continuous diffusion for categorical data.” arXiv preprint 2022.

[5] Santilli, A, et al. “Accelerating Transformer Inference for Translation via Parallel Decoding.” ACL 2023.

[6] Kou, S, et al. “CLLMs: Consistency Large Language Models.” ICML 2024.

[7] Lou, A. “Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution.” ICML 2024.

[8] Hoogeboom, E, et al. “Autoregressive Diffusion Models.” ICLR 2022.

[9] Wu, T, et al. “AR-Diffusion: Auto-Regressive Diffusion Model for Text Generation.” NeurIPS 2023.

[10] Rasul, K, et al. “Autoregressive Denoising Diffusion Models for Multivariate Probabilistic Time Series Forecasting.” ICML 2021.

We thank the reviewer for their thorough feedback! We address specific concerns below.

Concern 1: Efficiency improvement over AR should be discussed in future work

SAD3-LMs have the potential to be faster than AR models because they can output multiple tokens in parallel, which yields an algorithm with improved arithmetic efficiency that better utilizes underlying hardware. We use the number of calls to the transformer (NFEs) as a proxy for inference cost; we see below that the NFEs and sample quality are comparable to AR, and SAD3-LM is much faster compared to previous semi-autoregressive model SSD-LM [1].

The reviewer is concerned that SAD3-LMs have significantly more NFEs than AR: we clarify that sampling from SAD3-LMs use fewer or equal NFEs as AR. We use an efficient sampler [2] that caches the predictions of the denoising model until a token is sampled, using at most $L$ NFEs. We clarify that SAD3-LM uses 1015 NFEs, and that the previously reported 10K NFEs is a typo (updated in line 431).

A SAD3-LM NFE can oftentimes be run in the same time as an AR NFE by better utilizing FLOPs. While an AR NFE is memory-bound, SAD3-LMs may perform more computation to generate multiple tokens using the same memory transfer. Hence, SAD3-LMs do not incur higher compute costs than AR.

Although we find that SAD3-LMs require the $\text{NFES} \approx L$ to achieve competitive quality, distillation for discrete diffusion enables sampling 8x faster than AR with improved quality (concurrent work presented in [3]). We expect SAD3-LMs will be faster than AR after distillation.

---

Concern 2: “Experiments should have been performed without pre-training”

Pre-training is not strictly necessary. At the request of the reviewer, we train a SAD3-LM from scratch on LM1B without the pre-training step, and find that it further closes the gap to AR.

We performed the pretraining step in our experiments to train multiple SAD3-LMs for varying block sizes under a limited compute budget.

The reviewer asks if it is fair to train SAD3-LMs from an autoregressive checkpoint: we clarify that we do not initialize from an autoregressive checkpoint, and instead initialize from a base SAD3-LM with block size $L’=L$ (Line 338). We aim to close the gap from diffusion (equivalent to our base SAD3-LM with $L’=L$) to AR by fine-tuning SAD3-LMs with block size $L’ < L$. We compare all models for the same number of gradient steps.

---

Concern 3: Some aspects of Sections 4 & 5 are unclear

The reviewer notes that aspects of Sections 4 and 5 (the problem of gradient variance in diffusion models) are unclear. We clarify these aspects below.

“It is not very clear why the variance of the gradient estimator has the formula in Eq. 9”

We clarify that the variance of the gradient estimator derived in Eq. 9 is computed by summing the component-wise variances of the gradient vectors across batches. This simple derivation provides an interpretable and computationally efficient way to quantify training variance.

“It is not clear why extreme masking leads to poor variance”

We provide clearer intuition on the effect of extreme masking on variance. There are two extreme cases:

1. If we only mask a single token, then the model doesn’t train on most tokens in the sequence. This is equivalent to performing stochastic gradient descent with a smaller batch size, which increases training variance.
2. If we mask almost all tokens, then recovering the clean sequence is near impossible, so the optimal solution is to simply predict the marginals of tokens in the data distribution. Since the sequence doesn’t provide a useful training signal, it is also equivalent to reducing the batch size.

CDCD [4], a continuous diffusion framework, draws a similar conclusion. Whereas [4] proposes to maintain a linear relationship between the noise schedule and the uncertainty of model predictions, our problem formulation is more intuitive.

Hello, thank you for bringing this to our attention! We have included a citation to your work in the final camera-ready version, we will also update our arXiv paper accordingly

We acknowledge SSD-LM [1] as one of the earlier works studying block-autoregressive diffusion.

We note that while our approaches share conceptual similarity applied to different domains (text vs graphs), BD3-LMs focus on interpolating between AR/diffusion performance by tuning the block size. BD3-LMs also feature KV caching to further improve inference efficiency. Additionally, our training algorithm employs bidirectional attention within noised blocks containing both masked and unmasked tokens.

Thank you again, we appreciate your engagement with our paper.

[1] Han, X, et al. “SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control.” ACL 2023.