Blockwise SFT for Diffusion Language Models: Reconciling Bidirectional Attention and Autoregressive Decoding
This addresses a training-inference alignment problem for researchers and practitioners using diffusion-based language models, offering an incremental improvement to enhance text generation performance.
The paper tackled the mismatch between standard supervised fine-tuning and semi-autoregressive inference in discrete diffusion language models, which causes noisy gradients, and proposed Blockwise SFT to align training with blockwise decoding, resulting in consistent gains on benchmarks like GSM8K, MATH, and MetaMathQA under equal compute or token budgets.
Discrete diffusion language models have shown strong potential for text generation, yet standard supervised fine-tuning (SFT) misaligns with their semi-autoregressive inference: training randomly masks tokens across the entire response, while inference generates fixed-size blocks sequentially. This mismatch introduces noisy prefixes and leaky suffixes, biasing gradients away from the desired blockwise likelihood. We propose Blockwise SFT, which partitions responses into fixed-size blocks, selects one active block per step for stochastic masking, freezes all preceding tokens, and fully hides future ones. Loss is computed only over the active block, directly mirroring the blockwise decoding process. Experiments on GSM8K, MATH, and MetaMathQA show consistent gains over classical SFT under equal compute or token budgets. Block size consistency studies and ablations confirm that improvements stem from faithful training-inference alignment rather than incidental masking effects. Our results highlight the importance of matching supervision granularity to the decoding procedure in diffusion-based language models.