Diffusion In Diffusion: Breaking the Autoregressive Bottleneck in Block Diffusion Models

TL;DR

This work tackles the irreversibility and myopia of block diffusion language models by introducing Diffusion in Diffusion, a draft-then-refine framework that progressively expands the block receptive field across multiple stages. It adds a Snapshot Confidence Remask mechanism and a mixed-scale training objective to enable both rapid drafting with small blocks and global revision with large blocks, yielding strong gains on OpenWebText with only a fraction of the fine-tuning budget. Empirically, the approach establishes new state-of-the-art Gen PPL for discrete diffusion on the target scale, with Stage 2 revision bringing perplexities from the draft level down to 21.9 (and down to 20.6 on longer sequences), approaching autoregressive performance while retaining diffusion-based efficiency. The key contributions—structural diffusion, a robust remasking strategy, and mixed-scale training—offer a practical path to combining local coherence and global planning in diffusion-based text generation.

Abstract

Block diffusion language models, operating as semi-autoregressive paradigms, combine the strengths of both autoregressive and diffusion paradigms. However, their strict unidirectional block dependencies introduce irreversibility and sacrifice the global planning capabilities for which diffusion models are renowned. In order to address these issues, 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 utilise 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 just 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.

Figures (2)

• Figure 1: Overview of Diffusion in Diffusion method
• Figure 2: Ablation on Revision Scope. Generative Perplexity (Gen PPL) as a function of the Stage 2 revision ratio $\gamma$ across varying block sizes $\mathcal{B}^{(2)}$. The gray dashed line represents the Stage 1 baseline (BS=4).