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

TL;DR

This paper introduces a principled method for adapting Auto-Regressive models into Diffusion Language Models by gradually transitioning block sizes, resulting in improved long-context generation and state-of-the-art performance.

Contribution

It proposes a novel adaptation pathway from AR to DLM using a block-diffusion framework with gradual block size increase, enhancing long-context capabilities.

Findings

Achieves state-of-the-art results among 7B-class DLMs.

Demonstrates effective long-context modeling and reasoning.

Proves the adaptation method is competitive across various model scales.

Abstract

Diffusion Language Models (DLMs) enable fast generation, yet training large DLMs from scratch is costly. As a practical shortcut, adapting off-the-shelf Auto-Regressive (AR) model weights into a DLM could quickly equip the DLM with strong long-context generation capabilies. 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 two key questions unaddressed: where is the final destination of adaptation, and how to adapt better? For manifold benefits, we reframe the whole AR-to-DLM adaptation under the Block-Diffusion paradigm, transitioning from block size 1 to the final Block-Diffusion state. Concretely, the principled pathway of adaptation is designed as follows: we keep a context-causal path where causal attention is kept in the prefix, an efficient parallel adaptation procedure where an AR guidance is maintained, and gradual increment of the generation block size for a smoother transition. Built on these components, the adaptation is proved competitive on various models at different scales. With better adaptation, we propose NBDiff-7B that could inherit the long-context modeling and reasoning capabilities, and achieve state-of-the-art performance among the 7B-class DLMs. Codes: https://github.com/YuchuanTian/NBDiff.