Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Blockwise SFT for Diffusion Language Models: Reconciling Bidirectional Attention and Autoregressive Decoding

Bowen Sun, Yujun Cai, Ming-Hsuan Yang, Yiwei Wang

TL;DR

Blockwise SFT addresses a core mismatch between training and blockwise inference in diffusion language models by supervising exactly one active block at a time with a clean prefix and fully hidden suffix, instead of full-sequence random masking. The authors derive a variational upper bound and an unbiased gradient estimator for the blockwise objective, providing theoretical guarantees that align training with the actual decoding process. Empirically, Blockwise SFT yields consistent gains over classical SFT on GSM8K, MATH, and MetaMathQA under matched compute and token budgets, with ablations showing the gains stem from improved supervision granularity rather than incidental masking. The approach is architecture-agnostic, requires no changes to the decoder, and highlights the importance of aligning supervision granularity with decoding procedure in diffusion-based language models.

Abstract

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.

Blockwise SFT for Diffusion Language Models: Reconciling Bidirectional Attention and Autoregressive Decoding

TL;DR

Blockwise SFT addresses a core mismatch between training and blockwise inference in diffusion language models by supervising exactly one active block at a time with a clean prefix and fully hidden suffix, instead of full-sequence random masking. The authors derive a variational upper bound and an unbiased gradient estimator for the blockwise objective, providing theoretical guarantees that align training with the actual decoding process. Empirically, Blockwise SFT yields consistent gains over classical SFT on GSM8K, MATH, and MetaMathQA under matched compute and token budgets, with ablations showing the gains stem from improved supervision granularity rather than incidental masking. The approach is architecture-agnostic, requires no changes to the decoder, and highlights the importance of aligning supervision granularity with decoding procedure in diffusion-based language models.

Abstract

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.

Paper Structure

This paper contains 33 sections, 3 theorems, 17 equations, 8 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Overview: From Training-Inference Mismatch to Blockwise Alignment
  5. Classical SFT: Analysis of Training-Inference Mismatch
  6. Deriving the Blockwise Objective
  7. Practical Implementation
  8. Experiments
  9. Experiment Setup
  10. Main Results
  11. Block Size Consistency Study
  12. Ablation Studies
  13. Conclusion
  14. Appendix
  15. Full Hyperparameter Settings
  16. ...and 18 more sections

Key Result

Theorem 3.1

Let $\nabla_\theta \ell^{\star}(\theta;\mathbf{x},a)$ denote the ideal gradient for generating block $a$ from a clean prefix (as in inference), and $\nabla_\theta \ell^{\text{cls}}(\theta;\mathbf{x},a)$ the gradient from classical SFT. If gradients are $L_{\text{pre}}$-Lipschitz with respect to pref

Figures (8)

  • Figure 1: Mismatch between training and inference: Classical SFT uses bidirectional attention and full-sequence masking, exposing future tokens and corrupting parts of the prefix; semi-autoregressive decoding is one-directional, conditioning only on a clean prefix while future blocks are strictly hidden.
  • Figure 2: Blockwise SFT step: split the completion, sample an active block, keep a clean prefix and a hidden future, and train only on masked tokens inside that block.
  • Figure 3: Training loss on MetaMathQA under Equal-FLOPS.
  • Figure 4: Test Pass@1 during training under Equal-FLOPS.
  • Figure 5: Test Pass@1 during training under Equal-Tokens.
  • ...and 3 more figures

Theorems & Definitions (6)

  • Theorem 3.1: Gradient bias under training-inference mismatch
  • Theorem 3.2: Variational bound on blockwise generation
  • Theorem 3.3: Unbiased gradient estimation
  • proof : Proof of Theorem \ref{['thm:bias_bound']}
  • proof : Proof of Theorem \ref{['thm:upperbound']}
  • proof : Proof of Theorem \ref{['thm:unbiased_general']}