PaperlayerView exact revision on arXiv

Source-linked AI summary

Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

Marianne Arriola, Aaron Gokaslan, Justin T. Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, Volodymyr Kuleshov

arXiv:2503.09573v3cs.LGcs.AI

TL;DR

Existing diffusion language models offer potential gains in parallel generation and controllability but are limited in likelihood modeling and fixed-length generation. The paper introduces BD3-LMs, combining autoregressive block structure with discrete diffusion, and reports state-of-the-art perplexity among discrete diffusion models while enabling arbitrary-length generation.

  • Problem

    Diffusion language models are limited by fixed-length generation and lag autoregressive models in likelihood modeling, despite potential benefits in parallelized generation and controllability.

  • Method

    BD3-LMs model blocks autoregressively, apply discrete diffusion within each block, and use efficient training algorithms plus data-driven noise schedules that reduce gradient variance.

  • Results

    BD3-LMs achieve new state-of-the-art perplexity among discrete diffusion models and generate arbitrary-length sequences, including sequences longer than the training context.

  • Takeaways & Limitations

    Block diffusion combines variable-length generation, KV caching, and parallel within-block sampling while narrowing the gap between diffusion and autoregressive language modeling.

  • Takeaways & Limitations

    Training BD3-LMs is more expensive than regular diffusion training, and sequential block generation may face autoregressive-like speed and controllability constraints when blocks are small.

Abstract

from arXiv · show

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

The paper introduces BD3-LMs, which combine autoregressive block structure with discrete diffusion to address fixed-length generation, inefficient inference, and diffusion quality gaps. It develops efficient training and variance-reducing noise schedules, achieving strong benchmark results.

  • Motivation: Existing discrete diffusion models are limited by fixed-length generation, unavailable KV caching, and lower quality than autoregressive models.These constraints are especially relevant for applications requiring arbitrary-length responses.
  • Core approach: BD3-LMs define an autoregressive distribution over token blocks while modeling each block’s conditional distribution with discrete denoising diffusion.This interpolates between autoregressive and diffusion language models.
  • Training: The paper introduces specialized training algorithms that efficiently leverage the entire batch of tokens.The training objective cannot be computed efficiently with one standard neural-network forward pass.
  • Training: Gradient variance is identified as a contributor to the perplexity gap, motivating data-driven noise schedules that reduce variance.The proposed schedules target the variance of diffusion-model training gradients.
  • Results: BD3-LMs achieve new state-of-the-art perplexity among discrete diffusion models and generate arbitrary-length sequences.They support variable-length generation and improve inference efficiency through KV caching and parallel sampling.

2 BACKGROUND: LANGUAGE MODELING PARADIGMS

The background contrasts autoregressive language modeling with discrete diffusion. Autoregressive models train efficiently but generate sequentially, while diffusion models reverse progressive corruption and can support parallel generation but introduce different modeling constraints.

  • Autoregressive models: Autoregressive models factorize token probabilities conditionally on preceding tokens and can be trained efficiently through next-token prediction.Their sequential dependencies require L generation steps for L tokens.
  • Autoregressive models: Autoregressive generation takes L steps to produce L tokens because of sequential dependencies.This contrasts with diffusion’s potential for parallelized generation.
  • Diffusion models: Diffusion models learn to reverse a forward corruption process whose latent variables represent progressively noisier versions of clean data.The process is discretized into T steps.
  • Discrete diffusion: D3PM applies a Markov forward process independently to tokens, with diffusion matrices supporting masking, random token changes, and related substitutions.The reverse model predicts clean tokens from noisy sequences.
  • Discrete diffusion: Discrete diffusion models are trained with variational inference using the Negative ELBO objective.The formalism also extends to continuous-time Markov-chain formulations and score-based generalizations.

3 BLOCK DIFFUSION LANGUAGE MODELING

Block diffusion models autoregressively generate token blocks while applying discrete diffusion within each block. The paper supplies likelihood-based training and sampling procedures that use vectorization, parallel within-block generation, and KV caching.

  • Block diffusion distributions: Block diffusion factorizes the likelihood over B blocks of length L′ and models each block conditional with discrete diffusion.The number of blocks satisfies B = L/L′ when the sequence length is divisible by the block length.
  • Training objective: The block diffusion objective sums per-block NELBO terms, yielding a valid NELBO for the full model.Each term models log pθ(xb | x<b).
  • Model architecture: A single transformer with block-causal attention parameterizes the denoisers for all blocks.Tokens in block b attend to tokens in blocks 1 through b.
  • Training: Vectorized training computes the block objective efficiently using concatenated noisy and clean sequences with a tailored attention mask.The implementation avoids looping over blocks and keeps training overhead tractable.
  • Sampling: Sampling proceeds one block at a time, using previously sampled blocks as context and caching their keys and values.The cache avoids recalculating context representations for each new block.
  • Sampling: Block diffusion supports arbitrary-length sequences and parallel generation within each block.This differs from fixed-length diffusion sampling and token-by-token autoregressive sampling.

4 UNDERSTANDING LIKELIHOOD GAPS BETWEEN DIFFUSION & AR MODELS

The single-token case shows that block diffusion and autoregression have equivalent objectives in expectation, yet diffusion training exhibits a perplexity gap because it estimates gradients from fewer masked tokens. Full masking removes this gap and lowers training-loss variance.

  • 4.1 MASKED BD3-LMS: Masked BD3-LMs use a diffusion formulation in which tokens are gradually replaced by a special mask token.The masking probability is 1 −α_t under a schedule with α_0 = 1 and α_1 = 0.
  • 4.1 MASKED BD3-LMS: The NELBO is tight for L′ = 1 but becomes a looser approximation of the true negative log-likelihood as L′ approaches L.Its continuous-time formulation uses the instantaneous rate of change of α_t.
  • 4.2 CASE STUDY: SINGLE TOKEN GENERATION: For L′ = 1, the block diffusion parameterization is equivalent in expectation to the autoregressive NLL, but shows a two point perplexity gap on LM1B.The comparison uses single-token generation on LM1B.
  • 4.2 CASE STUDY: SINGLE TOKEN GENERATION: The gap arises because diffusion computes cross-entropy only for masked tokens, using 2x fewer tokens than autoregressive training.This reduced token count produces higher-variance loss-gradient estimates.
  • 4.2 CASE STUDY: SINGLE TOKEN GENERATION: Full masking makes the block diffusion objective equivalent to the autoregressive objective and yields the same perplexity as AR training.With q(x_t = m|x_ℓ) = 1, the measured NELBO variance falls from 1.52 to 0.11 after training on 328M tokens.
  • 4.3 DIFFUSION GAP FROM HIGH VARIANCE TRAINING: The training estimator’s variance depends on the noise schedule even though the NELBO itself is schedule-invariant.The paper therefore estimates gradient variance across random batches and sampled diffusion times.

5 LOW-VARIANCE NOISE SCHEDULES FOR BD3-LMS

The paper reduces diffusion-training variance with clipped noise schedules that avoid uninformative extreme masking rates. It adaptively selects schedule parameters because the variance-minimizing mask distribution depends on block size.

  • 5.2 CLIPPED SCHEDULES FOR LOW-VARIANCE GRADIENTS: Very low or very high masking rates produce weak learning signals and high-variance gradients.Few masked tokens are easy to reconstruct, while masking everything reduces the target to token marginals.
  • 5.2 CLIPPED SCHEDULES FOR LOW-VARIANCE GRADIENTS: Clipped schedules sample mask rates 1 −α_t uniformly between β and ω.The schedule approximates near-zero and near-one masking probabilities outside the selected interval.
  • 5.2 CLIPPED SCHEDULES FOR LOW-VARIANCE GRADIENTS: The optimal mask-rate distribution can vary with block size, so the schedule is learned adaptively during training.The method directly optimizes β and ω using diffusion-ELBO variance as a proxy for gradient-estimator variance.
  • 5.3 DATA-DRIVEN CLIPPED SCHEDULES ACROSS BLOCK SIZES: For block sizes L′ ∈ {4, 16, 128}, each block size has a unique clipped distribution minimizing both NELBO variance and test perplexity.The results show a correlation between diffusion-NELBO variance and test perplexity.
  • 5.3 DATA-DRIVEN CLIPPED SCHEDULES ACROSS BLOCK SIZES: Table 2 compares test perplexity with Var_X,t [L_BD(X; θ)] after LM1B training and finetuning.Models use a linear schedule for 65B tokens and are finetuned for 10B tokens.

6 EXPERIMENTS

Experiments evaluate BD3-LMs on likelihood, zero-shot generalization, sample quality, variable-length generation, noise schedules, and training efficiency. Across these settings, BD3-LMs improve diffusion-model perplexity, generate longer sequences, and accelerate training with vectorization.

  • Likelihood evaluation: BD3-LMs outperform prior diffusion methods on LM1B and improve perplexity over MDLM by up to 13%.A similar trend is observed on OpenWebText.
  • Likelihood evaluation: BD3-LMs achieve the best zero-shot perplexity on Pubmed and the best diffusion-model perplexity on Wikitext, LM1B, and AG News.Models are trained on OpenWebText and evaluated on unseen datasets.
  • Sample quality and variable-length sequence generation: BD3-LMs generate sequences up to ≈10× longer than SEDD, which is restricted to its training context size.The comparison uses 500 variable-length samples from models trained on OpenWebText.
  • Sample quality and variable-length sequence generation: BD3-LMs achieve the best generative perplexities among prior diffusion methods and approach autoregressive sample quality in qualitative comparisons.Relative to SSD-LM, improved generative perplexity uses an order of magnitude fewer generation steps.
  • Noise schedules: Clipped noise schedules are most effective for reducing training variance, with optimal masking rates depending on block size.Heavier masking is effective for L′ = 4, whereas lighter masking is optimal for L′ = 16.
  • Training efficiency: A vectorized single-forward-pass training approach provides a 20-25% speed-up over two forward passes.The approach reduces memory-bandwidth bottlenecks by using efficient attention kernels.

7 DISCUSSION AND PRIOR WORK

The paper positions BD3-LMs as a discrete block-diffusion alternative to D3PM, MDLM, Gaussian diffusion, and semi-autoregressive diffusion. It also notes practical limitations involving training cost, sequential block generation, block-size selection, and general generative-model risks.

  • Comparison to D3PM: BD3-LMs extend D3PM to variable-length generation, address gradient variance, and improve D3PM perplexity.The method applies D3PM-style diffusion to autoregressive conditionals over blocks.
  • Comparison to MDLM: BD3-LMs push perplexity beyond MDLM while showing that noise schedules significantly affect gradient variance despite schedule-invariant NELBO.The reported perplexity improvements stem from more than block diffusion alone.
  • Comparison to Gaussian and semi-autoregressive diffusion: Compared with Gaussian diffusion over embeddings, discrete diffusion provides tractable likelihood estimates and avoids the worse perplexity reported for continuous alternatives.The comparison concerns diffusion over continuous embeddings of discrete tokens.
  • Comparison to semi-autoregressive diffusion: Compared with SSD-LM, BD3-LMs use fewer generation calls and produce improved sample quality while retaining likelihood estimation.BD3-LM model calls are bounded by the number of generated tokens.
  • Limitations: Training BD3-LMs is more expensive than regular diffusion, while sequential block generation can impose autoregressive-like speed and controllability constraints.The optimal block size is task specific and may be larger when greater control is desired.

8 CONCLUSION

The paper introduces BD3-LMs to address arbitrary-length generation and the perplexity gap between discrete diffusion and autoregressive models. Its block-wise formulation, specialized training algorithm, and custom noise schedules improve long-form generation and diffusion-model perplexity.

  • Conclusion: BD3-LMs represent a block-wise extension of D3PM using specialized training algorithms and custom noise schedules.The models are designed around the paper’s two motivating problems: arbitrary-length generation and the perplexity gap to autoregressive models.
  • Conclusion: BD3-LMs generate long-form documents and establish a new state of the art among discrete diffusion models in perplexity.

C Experimental Details

The experimental details define masked block diffusion, its NELBO and single-token autoregressive limit, and the efficient transformer implementation. They also describe attention masks that process noised tokens and conditional context together.

  • Block diffusion objective: The sequence likelihood is factorized over blocks, with discrete diffusion applied within each block over T discretization steps.The derivation uses a block-level NELBO.
  • Masked diffusion process: Masked diffusion uses a decreasing noise schedule α_t from 1 to 0, with a dedicated mask state and independently applied token transitions.The forward process is defined through diffusion matrices Q_t.
  • NELBO: The masked-diffusion NELBO simplifies to a weighted average of cross-entropy terms and is invariant to the noise schedule.
  • Autoregressive limit: For block size one, the block-diffusion NELBO recovers the autoregressive NLL, although its estimator can exhibit a two-point perplexity gap in practice.The equivalence is stated for single-token generation, while the empirical gap is reported when training on LM1B.
  • Efficient training: An efficient training algorithm concatenates the noised sequence and conditional context, then updates them with a custom 2L × 2L attention mask.The mask combines block-diagonal, offset block-causal, and block-causal components.
  • Efficient training: The attention-mask example uses L = 6 and block size L′ = 2 to visualize the block structure.

B.7 OPTIMIZED ATTENTION KERNEL WITH FLEXATTENTION

The paper uses FlexAttention with a custom sparse block-diffusion mask to reduce memory and accelerate attention computation. This optimization combines structured sparsity, precomputed masking, and fused GPU kernels.

  • Sparse attention design: The custom mask exploits the extreme sparsity of the block-diffusion attention matrix to improve computational efficiency.The mask is implemented through FlexAttention, which supports structured sparse attention mechanisms.
  • Sparse attention design: FlexAttention precomputes block-level masking and skips fully masked blocks, reducing FLOPs, memory accesses, and attention-memory overhead.The implementation uses compile-time mask construction and avoids full materialization of large attention matrices.
  • Kernel optimization: torch.compile fuses masked-attention operations into optimized Triton-based GPU kernels that parallelize execution.The custom kernel is designed around the block-wise sparsity pattern.
  • Measured performance: Up to ≈5X speedup is achieved over native PyTorch scaled_dot_product_attention on an A5000 GPU with L = 1024 and batch size B = 16.The comparison uses the naive native implementation in PyTorch ≥2.5.
  • Measured performance: Replacing FlashAttention with FlexAttention produces an ≈15% end-to-end model-forward-pass speedup under the same L = 1024 and B = 16 setting.This reports model forward-pass speed rather than only the attention-kernel speed.

C.1 DATASETS

Experiments evaluate the models on LM1B and OpenWebText using wrapped fixed-length training sequences and specified sampling and likelihood protocols. The setup also addresses variable-length generation, diffusion-step costs, and sampling efficiency.

  • Datasets and models: Experiments use LM1B with context length 128 and OpenWebText with context length 1024, reporting LM1B test perplexity and reserving 100k OWT documents for validation.LM1B uses bert-base-uncased tokenization, while OWT uses the GPT2 tokenizer.
  • Data preparation: LM1B and OWT examples are concatenated and wrapped to the training context length rather than relying on short-example padding or sequence truncation.LM1B baselines are retrained with wrapping to make larger block sizes useful during training.
  • Datasets and models: The evaluated models use a 12-layer, 768-dimensional, 12-head transformer with approximately 110M parameters and rotary positional embeddings.The parameterization covers autoregressive baselines, SEDD, MDLM, and BD3-LMs.
  • Training protocol: BD3-LMs train for 850K steps at maximum context length, then fine-tune for 150K steps under varying block sizes using noise-schedule optimization.Training uses LM1B and OWT across 3090, A5000, A6000, and A100 GPUs.
  • Inference protocol: The first-hitting sampler reduces sampling-variable complexity from O(TLV) to O(LV) while remaining theoretically equivalent to the MDLM sampler.It analytically samples the first timestep at which a masked token becomes unmasked, especially benefiting settings where T ≫ L.

D SAMPLES

The samples compare outputs from MDLM, BD3-LM, and an autoregressive model using generative perplexity under GPT2-Large and sample entropy. BD3-LM generates a longer sample than MDLM with lower generative perplexity, while AR has the lowest perplexity among the shown examples.

  • MDLM sample: MDLM produces a length-1024 sample with generative perplexity 69.26 and entropy 5.6 under GPT2-Large.The sample uses T = 5K diffusion steps.
  • BD3-LM sample: BD3-LM produces a length-2031 sample with block size L′ = 16, generative perplexity 24.3, and entropy 5.5 under GPT2-Large.The model was trained with context length L = 1024 and uses T = 5K diffusion steps.
  • Cross-model comparison: Among these examples, BD3-LM has lower generative perplexity than MDLM at similar entropy while generating a comparable-length sample to the autoregressive model.The reported values are 24.3 versus 69.26 perplexity, entropy 5.5 versus 5.6, and lengths 2031 versus 2003.
  • Autoregressive sample: The autoregressive model produces a length-2003 sample with generative perplexity 10.6 and entropy 5.5 under GPT2-Large.It was trained with context length L = 1024.
Loading 2503.09573v3…