Source-linked AI summary
Diffusion In Diffusion: Reclaiming Global Coherence in Semi-Autoregressive Diffusion
Linrui Ma, Yufei Cui, Kai Han, Yunhe Wang
TL;DR
Global diffusion models offer bidirectional global context, but block diffusion trades some of that perspective for faster semi-autoregressive inference. DIFFUSION IN DIFFUSION uses draft-then-revise generation, confidence-based remasking, and mix-scale training to recover global capabilities. On OpenWebText, it reduced generative perplexity from 25.7 to 21.9 using 26% of baseline fine-tuning budget.
Problem
Block diffusion improves inference efficiency through autoregressive structure but can lose the global coherence and receptive field of global diffusion models.
Method
DIFFUSION IN DIFFUSION generates fast small-block drafts, then revises them with larger-block bidirectional diffusion using snapshot confidence remasking and mix-scale training.
Results
25.7 to 21.9: generative perplexity on OpenWebText decreased using only 26% of the baseline fine-tuning budget.
Takeaways & Limitations
The draft-then-revise paradigm achieves state-of-the-art generation perplexity for diffusion-based models while narrowing the performance gap with autoregressive models.
Abstract
from arXiv · showhide
One of the most compelling features of global discrete diffusion language models is their global bidirectional contextual capability. However, existing block-based diffusion studies tend to introduce autoregressive priors, which, while offering benefits, can cause models to lose this global coherence at the macro level. To regain global contextual understanding while preserving the advantages of the semi-autoregressive paradigm, we propose Diffusion in Diffusion, a 'draft-then-refine' framework designed to overcome the irreversibility and myopia problems inherent in block diffusion models. Our approach first employs block diffusion to generate rapid drafts using small blocks, then refines these drafts through global bidirectional diffusion with a larger bidirectional receptive field. We utilize snapshot confidence remasking to identify the most critical tokens that require modification, and apply mix-scale training to expand the block diffusion model's global capabilities. Empirical results demonstrate that our approach sets a new benchmark for discrete diffusion models on the OpenWebText dataset. Using only 26% of the fine-tuning budget of baseline models, we reduce generative perplexity from 25.7 to 21.9, significantly narrowing the performance gap with autoregressive models.
1 Introduction
Discrete diffusion research balances global coherence against efficient inference: global models preserve bidirectional context but are slow, while block diffusion accelerates generation yet sacrifices global perspective. DIFFUSION IN DIFFUSION addresses this tension through staged drafting and global revision, achieving strong OpenWebText results with limited tuning.
- Motivation: Global diffusion models provide global planning and complete bidirectional context, producing text with high global consistency but incurring O(L^2) inference complexity.Their lack of Key-Value Cache prevents the standard inference acceleration used by modern language models.
- Motivation: Block diffusion enables parallel execution within blocks and inter-block autoregression, providing rapid inference with Key-Value Cache and autoregressive priors.Its macro-level autoregressive structure sacrifices diffusion’s global receptive field and creates shortsightedness.
- Approach: DIFFUSION IN DIFFUSION uses small-block drafting followed by larger-block revision to restore global contextual understanding while retaining semi-autoregressive efficiency.The revision stage re-examines generated sequences with a larger bidirectional receptive field.
- Contributions: Snapshot confidence remasking selects tokens for modification, while mix-scale training combines minimal and maximal block masks to support generation and revision.The remasking strategy was reported as more effective than post-hoc likelihood evaluation.
- Results: 25.7 to 21.9: generative perplexity fell on OpenWebText using only one-quarter of comparable models’ fine-tuning budget.The result was reported as a new benchmark for discrete diffusion models and narrowed the gap with autoregressive benchmarks.
2 Preliminaries
Autoregressive models generate sequentially from prefix context, whereas block diffusion treats blocks as autoregressive-level units and denoises tokens within each block in parallel. Its training objective learns conditional infilling of the current noisy block given previously denoised blocks.
- Autoregressive Models: Autoregressive models decompose sequence probability into conditional probabilities over each token given its preceding prefix.The conditional distribution is typically modeled with causal attention.
- Autoregressive Models: Autoregressive generation supports Key-Value Caching but is inherently sequential, preventing straightforward parallel generation.This creates a direct trade-off between caching-based acceleration and parallel decoding.
- Block Diffusion Models: Block Diffusion Language Models partition a sequence into fixed-size blocks and treat each block as a metatoken at the autoregressive level.The sequence is divided into B blocks, with each block conditioned on preceding blocks.
- Block Diffusion Models: Within each block, discrete diffusion models the conditional distribution pθ(xb|x<b) instead of using autoregressive token-level conditioning.This allows parallel denoising inside blocks while preserving inter-block autoregressive structure.
- Diffusion Process: The forward process independently converts block tokens to [MASK] according to a noisy scheduling parameter, with t = 1 representing full noise.Larger t denotes a noisier timestep.
- Diffusion Process: The reverse process predicts the clean current block from its noisy state and previously denoised blocks, while training minimizes the negative evidence lower bound.The objective trains conditional generation or infilling for the current block at any noise level given the history.
3 Structural Block Diffusion
Structural Block Diffusion uses progressively larger diffusion blocks to draft quickly, then revise uncertain regions with broader bidirectional context. Snapshot confidence remasking and mixed-scale training support this transition from local drafting to global refinement.
- Multi-Stage Generation: Structural Block Diffusion progressively increases block size across stages, combining fast small-block drafting with later bidirectional refinement.The initial stage prioritizes local coherence and speed, while larger later blocks expand the receptive field.
- Inter-Stage Remasking: Each refinement stage estimates token confidence, retains the highest-confidence tokens, and remasks the remainder for targeted infilling.The remasking ratio γ(k) controls how many tokens are revised between stages.
- Global Refinement: When the final block size equals the sequence length, refinement becomes full-sequence masked diffusion initialized by the structured draft.The final stage combines partial initialization with the structural prior produced during drafting.
- Snapshot Confidence: Snapshot confidence records each token’s probability when it is unmasked, preserving generation-time uncertainty for identifying potential errors.This avoids relying solely on static post-generation evaluation, which can be overconfident about hallucinated tokens.
- Mixed-Scale Training: Mixed-scale training samples block sizes from a distribution, using small drafting blocks and large global blocks to train both capabilities.The bimodal design uses B_draft and B_global; the stated example assigns 90% of samples to small blocks and 10% to large blocks.
4 Experiments
Experiments evaluate generation quality, efficiency, revision scope, remasking, and training-mix choices on OpenWebText. Structural refinement improves perplexity over block-diffusion baselines while requiring substantially less fine-tuning.
- Generation Quality: 40K fine-tuning steps, or 26% of the 150K baseline budget, produce the reported performance.The experiments describe this as data-efficient structural refinement.
- Generation Quality: The evaluation uses GPT-2-Large-based generative perplexity for unconditional generations at lengths L = 1024 and L = 2048.Generative perplexity measures coherence and fluency of open-ended generations rather than static ground-truth likelihood.
- Generation Quality: 21.9 Gen PPL is achieved after global revision, down from 27.4 for the Stage 1-only drafting baseline.The revision stage delivers an approximately 20% relative improvement.
- Generation Quality: 21.9 Gen PPL at L = 1024 and 20.6 at L = 2048 outperform BD3-LM’s 25.0 and 22.8, respectively, at the same NFEs.The method narrows the gap with the autoregressive benchmark at 14.1.
- Quality-Efficiency Trade-off: 24.6 PPL at 1.1K NFEs surpasses standard BD3-LM’s 25.7 with a 10% compute increase, while 21.9 at 1.5K NFEs beats the baseline’s 25.0.At 3.0K NFEs, the method reaches 20.6 PPL.
- Impact of Revision Scope: Revision quality improves with larger Stage 2 blocks: significant gains begin at B(2) ≥64, with best performance at B(2) = 1024.B(2) values of 4 or 16 can fail to beat or can degrade the Stage 1 baseline.
- Impact of Revision Scope: Revision ratios show a convex U-shaped trend, with lowest perplexity between 0.25 and 0.5.Ratios γ ≤0.1 constrain revision, whereas γ ≥0.75 discard too much of the draft structure.
- Effectiveness of Remasking Strategy: Snapshot confidence remasking reduces PPL to 21.85, whereas random masking raises it from 27.36 to 30.26 and post-hoc selection reaches 29.85.Snapshot confidence records uncertainty during diffusion rather than relying only on final-sequence confidence.
5 Related Works
Related work spans discrete diffusion, semi-autoregressive block generation, and iterative draft-and-revise methods. Diffusion in Diffusion distinguishes itself by using revision to restore global planning and coherence rather than primarily accelerating generation.
- Discrete Diffusion Language Models: Discrete diffusion models generate in parallel and offer controllability and flexibility compared with autoregressive models.The related-work discussion traces discrete diffusion from D3PM’s transition-matrix formulation.
- Semi-Autoregressive Block Generation: Semi-autoregressive block generation combines autoregressive macro-level behavior with diffusion-based micro-level parallelism and uses KV Cache for efficiency.This paradigm addresses inference efficiency while retaining local bidirectional generation.
- Iterative Refinement and Multi-Stage Generation: Speculative Decoding uses rapid drafts and verification mainly to speed generation, whereas Diffusion in Diffusion uses revision to improve quality through global planning and coherence.The comparison separates the paper’s quality objective from speculative decoding’s acceleration objective.
- Iterative Refinement and Multi-Stage Generation: Snapshot confidence identifies structural defects for repair, and nested iterative diffusion restores long-range modeling within an efficient block-based framework.The method does not simply reject tokens based on likelihood.
6 Conclusion
The paper presents Diffusion in Diffusion as a structural framework for combining semi-autoregressive efficiency with global coherence. On OpenWebText, it reports state-of-the-art generation perplexity and high data efficiency while narrowing the gap with autoregressive models.
- Conclusion: Diffusion in Diffusion targets the trade-off between global coherence and inference efficiency in semi-autoregressive diffusion.Its multi-stage refinement uses snapshot confidence and mix-scale training to reintroduce long-range planning.
- Conclusion: The Draft-then-Revise paradigm achieves state-of-the-art generation perplexity on OpenWebText with exceptional data efficiency.The conclusion reports a narrowed performance gap with traditional autoregressive models.