From Next-Token to Next-Block: A Principled Adaptation Path for Diffusion LLMs

TL;DR

<3-5 sentence high-level summary> The paper tackles the throughput bottleneck of autoregressive decoding by reframing diffusion language model adaptation as a principled AR-to-Block-Diffusion transition. It introduces a context-causal attention mask, parallel training with an auxiliary AR loss, and a block-size growth curriculum to smoothly migrate from blocksize=1 to larger blocks, thereby reusing pretrained AR knowledge for long-context generation. The proposed NBDiff-7B achieves state-of-the-art performance among 7B diffusion models on general knowledge, math, and code benchmarks, demonstrating a compute-efficient alternative to training DLMs from scratch. This work provides a practical pathway to leverage mature AR checkpoints for scalable, parallel-generation diffusion LLMs with potential for scaling and multimodal extensions.

Abstract

Large language models (LLMs) excel at generation but dominant autoregressive (AR) decoding is inherently sequential, creating a throughput bottleneck. Diffusion Language Models (DLMs)--especially block-wise variants--enable parallel generation and intra-block bidirectional reasoning, yet training large DLMs from scratch is costly and wastes the knowledge in mature AR checkpoints. Prior "adaptation" attempts either modify logits or randomly grow attention masks to full-sequence diffusion, or simply transplant AR weights into a block-diffusion recipe, leaving a fundamental mismatch between AR causality and block-wise bidirectionality unaddressed. We reframe adaptation as a intra-paradigm path from AR to Block-Diffusion by viewing AR as Block-Diffusion with blocksize=1. Concretely, we design the pathway of adaptation as follows: we use a context-causal attention mask (causal in context, bidirectional only within the active block), an efficient parallel adaptation procedure, an auxiliary AR loss to maximize data utilization and retain pretrained knowledge, and gradual increment of the generation block size. The recipe integrates cleanly with masked block-diffusion and maintains train-inference consistency. Built on these components, NBDiff-7B (Base and Instruct) could inherit the long-context modeling and reasoning capabilities, and achieve state-of-the-art performance among the 7B-class DLMs, delivering strong gains on general-knowledge, math, and code benchmarks over strong baselines. These results demonstrate that principled AR-to-block-diffusion adaptation is an effective and compute-efficient alternative to training DLMs from scratch. Codes: https://github.com/YuchuanTian/NBDiff.

Figures (4)

• Figure 1: Comparison of our model with baselines. After adaptation from an open-sourced AR LLM, our model has good long-sequence and reasoning capabilities and shows outstanding performance in various benchmarks.
• Figure 2: The diffusion paradigm of our NBDiff-7B-Instruct model. We compare popular language generation paradigms. Diffusion LLMs adapted from AR adopt logit shift and attention mask growth; Block-Diffusion uses block-wise autoregressive and maintains an intra-block bidirectional mask; Our model adopts Block-Diffusion where bidirectional attention is used intra-block, but features a causal context.
• Figure 3: Our Parallel Training Diagram. The diagram shows the parallel training form of our Context-Causal setting (we use $blocksize=4$ as an example; the actual $blocksize$ is 32). We concatenate a clean, unmasked token sequence to the noised sequence. The attention mask $\mathbf{M}_{\mathrm{all}}$ is designed (shown in the right) such that strictly-causal attention is applied in the unmasked input; for the masked input, each token has bidirectional attention intra-block, but causal attention to past inter-block tokens that are unmasked. AR loss $\mathcal{L}_{\mathrm{AR}}$ is introduced in addition to the canonical masked loss $\mathcal{L}_{\mathrm{MDM}}$ for faster adaptation.
• Figure 4: Our Parallel Training Diagram. The diagram shows the parallel training form of our Context-Causal setting (we use $blocksize=4$ as an example; the actual $blocksize$ is 32). We concatenate a clean, unmasked token sequence to the noised sequence. An attention mask is designed such that strictly-causal attention is applied in the unmasked input; for the masked input, each token has bidirectional attention intra-block, but causal attention to past inter-block tokens that are unmasked. AR loss $\mathcal{L}_{\mathrm{AR}}$ is introduced in addition to the canonical masked loss $\mathcal{L}_{\mathrm{MDM}}$ for faster adaptation.