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

## Abstract

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.

## 1 Introduction

Discrete diffusion models have emerged as a promising alternative to autoregressive language models for text generation. Unlike traditional left-to-right decoding, diffusion language models (LMs) learn to iteratively denoise corrupted text by predicting clean tokens from noisy versions (Gong et al., 2023). During training, these models corrupt response tokens with a relaxed categorical process and learn to recover the original text; during inference, they start from pure noise and progressively refine it into coherent text through multiple denoising steps.

Recent advances have explored semi-autoregressive decoding as a practical compromise between fully parallel generation and strict sequential processing. Systems like SSD-LM generate text in fixed-size blocks, where tokens within each block are produced in parallel while maintaining causal dependencies across blocks (Han et al., 2023). Modern discrete diffusion LMs, including the LLaDA family, have adopted blockwise inference as their primary generation strategy (Nie et al., 2025). This widespread adoption of blockwise decoding raises a fundamental question: *can a training objective based on bidirectional attention and full-sequence random masking effectively serve a model that will be decoded causally, block by block?*

The prevailing supervised fine-tuning (SFT) approach for diffusion LMs is fundamentally misaligned with blockwise inference. Classical SFT randomly masks tokens across the entire response and reconstructs them using bidirectional context in a single pass. However, during inference, the model operates differently: it receives a clean, deterministic prefix and must produce exactly one block while all future tokens remain completely hidden. As illustrated in figure [1](https://arxiv.org/html/2508.19529v2#S1.F1), this training-inference mismatch creates three critical problems: (i) noisy prefixes: the model trains on corrupted contexts that never appear during generation; (ii) dependency leakage: training may reveal tokens within or beyond the target block, violating the causal constraints enforced at inference; and (iii) granularity mismatch: training optimizes token-level decisions while inference requires block-level coordination.

Figure: 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.
Refer to caption: x1.png

To address this misalignment, we introduce Blockwise SFT (Figure [2](https://arxiv.org/html/2508.19529v2#S1.F2)), a training objective that mirrors deployment: at each step we supervise a single active block, keep the prefix clean and frozen, hide all future tokens, and compute loss only within the active block. This aligns gradients with blockwise decoding, preserves cross-block causality, and concentrates supervision where decisions are actually made. The recipe is architecture-agnostic and applies to any diffusion LM that supports blockwise inference.

Figure: 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.
Refer to caption: x2.png

Contributions.
(1) We diagnose the core mismatch between classical full-response SFT and blockwise semi-autoregressive decoding, showing how global random masking induces noisy prefixes and suffix leakage.
(2) We propose Blockwise SFT, a drop-in objective that supervises one active block at a time under a clean prefix and hidden future, requiring no architectural or sampling changes.
(3) We provide theory: a variational upper bound on blockwise likelihoods, unbiased timestep-sampled gradients, and a characterization of the gradient bias in classical SFT under mismatch.

Aligning training with inference yields concrete gains. Fine-tuning on MetaMathQA and evaluating on GSM8K and MATH, Blockwise SFT consistently improves Pass@1 under matched FLOPs and matched supervised tokens, with faster early-training progress. Performance peaks when the training block size matches the inference block size and drops with injected prefix noise or suffix visibility, corroborating our alignment analysis.

## 2 Related Work

Discrete Diffusion Models for Text Generation
Diffusion has been extended to discrete text via Argmax Flows and Multinomial Diffusion, which map tokens to continuous space and add categorical noise during training (Hoogeboom et al., 2021). D3PMs generalize this with structured corruption processes to better capture token relationships (Austin et al., 2023). DiffuSeq applies diffusion to conditional sequence-to-sequence tasks, matching or surpassing autoregressive baselines (Gong et al., 2023). Recent evaluations highlight trade-offs between modeling quality and generation speed (Weligalle, 2025).

Supervised Fine-Tuning Methods
For autoregressive models, SFT trains on cross-entropy loss over response tokens (Ouyang et al., 2022), with variants such as instruction tuning improving generalization (Wei et al., 2023; Chung et al., 2022). Diffusion-based models use iterative refinement, with SFT approaches reconstructing randomly masked tokens in one step (Austin et al., 2023; Li et al., 2022). This departs from the blockwise, semi-autoregressive decoding common in diffusion LLMs, creating a training–inference mismatch.

Autoregressive Decoding in Diffusion-Based Language Models
Early discrete diffusion generators were slow and sometimes incoherent. SSD-LM introduced semi-autoregressive decoding: generating fixed-size blocks with diffusion denoising while freezing prior blocks (Han et al., 2023). Block Diffusion interpolates between pure diffusion and fully autoregressive sampling via block decomposition and KV-caching, improving flexibility and throughput (Arriola et al., 2025). Adaptive Parallel Decoding (APD) adjusts block size by uncertainty, combining diffusion marginals with an autoregressive verifier to approach AR speeds with minimal quality loss (Israel et al., 2025).

Contemporary diffusion–SFT variants.
Recent variants emphasize efficiency or stability rather than decoding alignment: MDLM mixes masked-LM–style losses to narrow the gap to autoregressive models (Sahoo et al., 2024); Soft-Masked Diffusion LM uses linguistically informed soft masking with per-step cross-entropy to reduce cost (Chen et al., 2023); RDM reparameterizes the discrete process to improve training and sampling (Zheng et al., 2024); and Two-Step Loss & Scheduling adopts a two-step diffusion with gradually increased self-conditioning to mitigate mismatch (Asada & Miwa, 2025).
These methods chiefly modify the objective, schedule, or parameterization while supervising across the full response.
In contrast, Blockwise SFT aligns supervision with semi-autoregressive blockwise decoding (clean prefix, hidden future, loss on the active block), yielding consistent gains on GSM8K and MATH under equal compute (Table [1](https://arxiv.org/html/2508.19529v2#S4.T1)).

## 3 Methodology

### 3.1 Overview: From Training-Inference Mismatch to Blockwise Alignment

Semi-autoregressive diffusion models face a fundamental challenge: they are trained to denoise tokens anywhere in a sequence, but deployed to generate specific blocks conditioned on clean prefixes. This section develops Blockwise SFT to bridge this gap.

Consider generating a simple expression like “$5-3=\,$?”. During training, classical methods might mask the “5” or reveal the answer “2” in the suffix, but scenarios that never occur during inference, which always starts from a pristine prefix and generates the next block without future information. This mismatch creates three specific issues we quantify in Section [3.2](https://arxiv.org/html/2508.19529v2#S3.SS2): (i) training corrupts prefixes that inference keeps clean, (ii) training leaks future tokens that inference never sees, and (iii) training spreads supervision across all positions while inference focuses on the next block.

Our solution: blockwise alignment.
We make training mirror deployment. At each step, we identify the active block, which is the next chunk the model would generate during inference. We keep its prefix exactly as it appears at deployment (clean and fixed), completely hide the suffix (no future leakage), and concentrate the training signal solely on this active block. This alignment ensures that gradient updates directly optimize the objective used during generation.

Roadmap.
Section [3.2](https://arxiv.org/html/2508.19529v2#S3.SS2) analyzes classical SFT’s limitations under blockwise decoding and quantifies the resulting gradient bias. Section [3.3](https://arxiv.org/html/2508.19529v2#S3.SS3) derives our blockwise objective from first principles, proving it provides an unbiased variational bound on the true decoding risk. Section [3.4](https://arxiv.org/html/2508.19529v2#S3.SS4) translates the theory into a practical algorithm. Throughout, we interleave formal results with intuitive explanations to maintain clarity.

### 3.2 Classical SFT: Analysis of Training-Inference Mismatch

We first recall the classical approach, then characterize how it diverges from blockwise decoding.

Classical SFT recap.
Given an instruction–response sequence $\mathbf{x}_{1:L}=[\mathbf{c};\mathbf{r}]$, classical SFT samples independent masks on response tokens with rate $\pi$ (i.e., $m_{i}\!\sim\!\mathrm{Bernoulli}(\pi)$ for $i\!>\!L_{c}$) and trains the model to reconstruct only the masked positions at each diffusion step $t$:

$$ $\mathcal{L}_{\mathrm{SFT}}(\theta)=\sum_{t=1}^{T}\omega_{t}\,\mathbb{E}_{\mathbf{x},\mathbf{m},\mathbf{z}_{t}}\Bigl[-\!\!\sum_{i>L_{c}}\!\mathbf{1}[m_{i}=1]\,\log p_{\theta}\bigl(x_{i}\mid\mathbf{z}_{t},t\bigr)\Bigr].$ (1) $$

This objective treats masked positions uniformly and leverages bidirectional context—a capability unavailable during blockwise generation.

Quantifying the mismatch.
At inference, the model generates block $a$ (indices $\mathcal{I}_{a}$) conditioned on a clean prefix $\mathcal{I}_{\mathrm{prefix}}^{(a)}$ and without access to the future $\mathcal{I}_{\mathrm{suffix}}^{(a)}$. Classical training violates all three conditions:

- •
*Prefix corruption.* Random masking almost surely corrupts long prefixes (e.g., with $\pi{=}0.3$ and 100 prefix tokens, corruption occurs with probability $>1-3.3\times 10^{-16}$), whereas inference always conditions on a pristine prefix.
- •
*Suffix leakage.* A nontrivial fraction of future tokens remains visible during training even for modest suffix lengths, enabling shortcuts that are disallowed at inference.
- •
*Diluted supervision.* The loss in Eq. equation [1](https://arxiv.org/html/2508.19529v2#S3.E1) spreads gradients over all response positions. The active block obtains only a $|\mathcal{I}_{a}|/L_{r}$ share of the total signal, which is identical to any other block despite being the sole target at inference.

These mismatches induce a systematic bias in the learning signal:

###### Theorem 3.1 (Gradient bias under training-inference mismatch) .

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 prefix perturbations and $L_{\text{suf}}$-Lipschitz with respect to suffix visibility, then:

$$ $\bigl\|\,\mathbb{E}\bigl[\nabla_{\theta}\ell^{\text{cls}}(\theta;\mathbf{x},a)\bigr]-\nabla_{\theta}\ell^{\star}(\theta;\mathbf{x},a)\,\bigr\|\;\leq\;L_{\text{pre}}\!\cdot\!\bigl(1-(1-\pi)^{|\mathcal{I}_{\mathrm{prefix}}^{(a)}|}\bigr)\;+\;L_{\text{suf}}\!\cdot\!\bigl(1-\pi^{|\mathcal{I}_{\mathrm{suffix}}^{(a)}|}\bigr).$ (2) $$

*Intuition.* The bound quantifies how prefix corruption and suffix visibility pull gradients away from the true objective; the bias grows with sequence length and is not removed by tuning $\pi$, motivating a structurally aligned training procedure. (Proof in Appendix [A.6](https://arxiv.org/html/2508.19529v2#A1.SS6).)

###### Theorem 3.1 (Gradient bias under training-inference mismatch) .

### 3.3 Deriving the Blockwise Objective

We now develop an objective that aligns training with inference, ensuring unbiased gradient estimates for blockwise generation.

Target: blockwise generation risk.
Semi-autoregressive decoding generates a response as $M$ consecutive blocks $\mathbf{b}^{(1)},\dots,\mathbf{b}^{(M)}$ following instruction $\mathbf{c}$. The likelihood factorizes as:

$$ $p_{\theta}(\mathbf{x})=\prod_{a=1}^{M}p_{\theta}\!\bigl(\mathbf{b}^{(a)}\mid\mathbf{context}^{(a)},t{=}0\bigr),$ (3) $$

where $\mathbf{context}^{(a)}=\mathbf{x}_{1:L_{c}+B(a-1)}$ represents the clean prefix through block $a{-}1$. The risk we minimize is:

$$ $\mathcal{R}_{\text{block}}(\theta)=\mathbb{E}_{\mathbf{x}}\!\left[-\sum_{a=1}^{M}\log p_{\theta}\!\bigl(\mathbf{b}^{(a)}\mid\mathbf{context}^{(a)},t{=}0\bigr)\right].$ (4) $$

This is the exact objective optimized during inference.

Constructing a faithful training surrogate.
Direct optimization of Eq. equation [4](https://arxiv.org/html/2508.19529v2#S3.E4) is intractable due to the discrete diffusion process. We derive a surrogate that maintains three critical properties: it upper-bounds the true risk, admits efficient stochastic gradient estimation, and preserves the blockwise structure.

For active block $a$ with indices $\mathcal{I}_{a}$, we apply diffusion training only within this block while keeping the prefix clean and suffix hidden:

$$ $\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)=-\!\!\sum_{i\in\mathcal{I}_{a}}\log p_{\theta}\!\bigl(x_{i}\mid\mathbf{z}_{t},t\bigr),$ (5) $$

where $\mathbf{z}_{t}\!\sim\!q_{t}(\cdot\mid\mathbf{x})$ incorporates noise at diffusion step $t$. The full objective averages over blocks and diffusion steps:

$$ $\mathcal{L}_{\text{BW--SFT}}(\theta)=\mathbb{E}_{\mathbf{x},\,a}\!\Bigl[\sum_{t=1}^{T}\omega_{t}\,\mathbb{E}_{\mathbf{z}_{t}}\bigl[\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)\bigr]\Bigr].$ (6) $$

This construction provides rigorous guarantees connecting training to inference:

###### Theorem 3.2 (Variational bound on blockwise generation) .

For appropriate weights $\{\omega_{t}\}_{t=1}^{T}$, the blockwise diffusion objective upper-bounds the generation risk:

$$ $-\log p_{\theta}\!\bigl(\mathbf{b}^{(a)}\mid\mathbf{context}^{(a)},t{=}0\bigr)\;\leq\;\sum_{t=1}^{T}\omega_{t}\;\mathbb{E}_{\mathbf{z}_{t}\sim q_{t}(\cdot\mid\mathbf{x})}\bigl[\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)\bigr]\;+\;C,$ (7) $$

where $C$ is independent of $\theta$. Therefore, $\mathcal{R}_{\text{block}}(\theta)\leq\mathcal{L}_{\text{BW--SFT}}(\theta)+C^{\prime}$.

*Intuition.* By focusing diffusion training on individual blocks with proper conditioning, we obtain a tractable upper bound on the exact inference objective. Minimizing this bound improves generation quality. (Proof in Appendix [A.6](https://arxiv.org/html/2508.19529v2#A1.SS6).)

###### Theorem 3.3 (Unbiased gradient estimation) .

Sampling block $a\sim\rho$ and diffusion step $t\sim\tilde{\omega}$ (where $\tilde{\omega}_{t}\propto\omega_{t}$) yields an unbiased gradient estimator:

$$ $\widehat{g}(\theta)=\frac{1}{\rho(a)}\!\left(\sum_{s=1}^{T}\omega_{s}\right)\nabla_{\theta}\,\mathbb{E}_{\mathbf{z}_{t}\sim q_{t}(\cdot\mid\mathbf{x})}\bigl[\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)\bigr].$ (8) $$

Then $\mathbb{E}[\widehat{g}(\theta)]=\nabla_{\theta}\,\mathcal{L}_{\text{BW--SFT}}(\theta)$ for any sampling distribution $\rho$.

*Intuition.* Despite sampling only one block and one diffusion step per update, our gradients remain unbiased estimates of the full objective. This enables efficient training while maintaining theoretical soundness. (Proof in Appendix [A.6](https://arxiv.org/html/2508.19529v2#A1.SS6).)

Summary.
The blockwise objective in Eq. equation [6](https://arxiv.org/html/2508.19529v2#S3.E6) bridges the gap between tractable training and faithful inference by: (i) providing a variational bound on the true generation risk, (ii) enabling unbiased stochastic optimization, and (iii) eliminating the systematic bias inherent in classical approaches.

###### Theorem 3.2 (Variational bound on blockwise generation) .

###### Theorem 3.3 (Unbiased gradient estimation) .

### 3.4 Practical Implementation

We now translate the theoretical framework into a practical training algorithm. Algorithm [1](https://arxiv.org/html/2508.19529v2#alg1) implements one training step of Blockwise SFT. The key insight is simplicity: we modify only the masking strategy, leaving the model architecture and inference procedure unchanged.

Figure: Algorithm 1 Blockwise SFT — One Training Step

Implementation details.
Blockwise SFT is a drop-in replacement for classical SFT: the only change is a structured mask that preserves the prefix and hides all future tokens. We keep the standard diffusion training recipe (Appendix [A.1](https://arxiv.org/html/2508.19529v2#A1.SS1)); each step uses the stochastic gradient estimator in Theorem [3.3](https://arxiv.org/html/2508.19529v2#S3.Thmtheorem3), with uniform block sampling $\rho(a)=1/M$ and importance-weighted diffusion steps to optimize—without bias—the variational bound in Theorem [3.2](https://arxiv.org/html/2508.19529v2#S3.Thmtheorem2). In practice we use sequence length $L{=}128$ to balance speed and context; uniform block sampling works well, though difficulty-aware reweighting could further improve efficiency. The method requires no architectural changes and integrates cleanly with existing diffusion LM codebases. By maintaining clean prefixes, hiding suffixes, and focusing loss on the active block, Blockwise SFT aligns training with blockwise inference and removes the systematic bias of classical SFT at comparable compute.

## 4 Experiments

### 4.1 Experiment Setup

We train on meta-math/MetaMathQA (Yu et al., 2024) and evaluate on the test splits of openai/gsm8k (Cobbe et al., 2021) and HuggingFaceH4/MATH-500 (Hendrycks et al., 2021), reporting Pass@1 via exact string match. These mathematical reasoning sets align with our blockwise paradigm: each step depends causally on preceding derivations without access to future tokens, mirroring clean-prefix, hidden-future supervision.

Configurations.
We compare (1) base (no fine-tuning), (2) Classical SFT with random masking over the full response, and (3) Blockwise SFT with active-block masking. All experiments use GSAI-ML/LLaDA-8B-Instruct with LoRA (Hu et al., 2021), following the recipe of Biderman et al. (2024).

Evaluation Protocols.
Recall that Classical SFT aggregates loss over the entire response each step, whereas Blockwise SFT aggregates loss only on the active block (clean prefix, fully hidden suffix). To isolate supervision granularity from raw compute, we report two comparisons: (i) hold model-side compute fixed while changing which tokens receive loss (Equal-FLOPS); and (ii) hold the total number of supervised tokens fixed while allowing the number of optimization steps to differ (Equal-Tokens).

Equal-FLOPS.
We fix the FLOP-determining knobs: parameter count, sequence length $L$, batch size, gradient accumulation, and number of optimizer updates. Thus, both methods consume identical forward and backward FLOPs; only the loss-bearing positions differ. For example, with $L_{c}{=}32$, $L_{r}{=}96$, and block size $B{=}32$, both run the same updates; Classical applies loss over masked positions among $96$ response tokens per step, while Blockwise applies loss only within the sampled active block of $32$ tokens.

Equal-Tokens.
We equalize the total number of supervised tokens over the run and index progress by $\tau$ (traversals): for Classical SFT, $\tau$ equals epochs; for Blockwise, $\tau{=}1$ means every sample’s blocks are each supervised once (requiring $\approx L_{r}/B$ epochs). In the same example, Classical supervises $96$ tokens/step while Blockwise supervises $32$; accordingly, Blockwise runs $3{\times}$ more epochs to reach the same $\tau$.

Implementation Details.
Bernoulli masking rate $\pi\!\sim\!\mathrm{Uniform}(10^{-3},1)$ is sampled per sequence (Classical) or per active block (Blockwise); if nothing is masked, we mask the last token.
Training uses AdamW (Loshchilov & Hutter, 2019) with linear warmup and cosine decay (Smith, 2017) in bfloat16 (Micikevicius et al., 2018), runs on a single A100 (80GB) at $\sim$5.1k tokens/s, and peaks at $\approx$23.5GB memory (deployable on RTX 3090) when the batch size is set to 1. Full hyperparameters are in Appendix [A.1](https://arxiv.org/html/2508.19529v2#A1.SS1).

### 4.2 Main Results

Table [1](https://arxiv.org/html/2508.19529v2#S4.T1) presents our primary results comparing Blockwise SFT against the original model and five baseline methods under matched compute. Blockwise SFT achieves the highest Pass@1 accuracy on both benchmarks, with a substantial +5.2 point gain over the strongest baseline (Two-Step Loss & Scheduling) on GSM8K and +1.6 points on MATH. Notably, Classical SFT underperforms the base model on MATH (29.6 vs. 31.7), suggesting that naive supervision misalignment can be detrimental. Recent variants offer only marginal improvements, while Blockwise SFT’s aligned supervision delivers consistent and significant gains.

**Table 1: Head-to-head comparison under Equal-FLOPS: Pass@1 (%) on GSM8K and MATH.**
| Method | GSM8K (Pass@1) | MATH (Pass@1) |
| --- | --- | --- |
| Base (no FT) | $62.1\pm 1.0$ | $31.7\pm 0.4$ |
| Classical SFT (Austin et al., 2023) | $67.7\pm 1.4$ | $29.6\pm 0.7$ |
| MDLM (Sahoo et al., 2024) | $68.4\pm 1.9$ | $31.9\pm 0.5$ |
| Soft-Masked Diffusion LM (Chen et al., 2023) | $67.7\pm 0.8$ | $29.9\pm 0.6$ |
| RDM (Zheng et al., 2024) | $65.5\pm 1.5$ | $32.3\pm 0.6$ |
| Two-Step Loss & Scheduling (Asada & Miwa, 2025) | $70.8\pm 1.3$ | $32.6\pm 0.7$ |
| Blockwise SFT (Ours) | $76.0\pm 1.6$ | $34.2\pm 0.5$ |

Training Dynamics under Equal-FLOPS and Equal-Tokens.
Figs. [3](https://arxiv.org/html/2508.19529v2#S4.F3)a–b show that Blockwise SFT converges to a lower loss (0.24) than Classical SFT (0.36), indicating more effective learning under aligned supervision.
Figs. [4](https://arxiv.org/html/2508.19529v2#S4.F4)a–b track test Pass@1 throughout training: on GSM8K, Blockwise SFT maintains a $\sim$10-point margin; on MATH, Classical SFT degrades below the base model while Blockwise SFT improves steadily. These dynamics highlight that training–inference alignment is crucial to preserve and enhance reasoning ability.

Figure: (a) Classical SFT
Refer to caption: x3.png

Figure: (a) GSM8K
Refer to caption: x5.png

Fig. [5](https://arxiv.org/html/2508.19529v2#S4.F5) evaluates accuracy under a fixed supervised-token budget.
Blockwise SFT delivers large gains at $\tau{=}1$ and remains stable through $\tau{=}3$.
By contrast, Classical SFT behaves differently across datasets: on GSM8K it shows a small early uptick at $\tau{=}1$ before degrading as $\tau$ increases, whereas on MATH it declines monotonically across traversals.
The late-epoch drops further indicate clear overfitting for Classical SFT when the same tokens are revisited over multiple epochs, while Blockwise SFT remains robust under the same budget. Thus, Blockwise SFT’s advantage stems from supervision quality, not quantity.

Figure: (a) GSM8K
Refer to caption: x7.png

### 4.3 Block Size Consistency Study

Figure [6](https://arxiv.org/html/2508.19529v2#S4.F6) investigates robustness to block size misalignment between training and inference. For each dataset, we train four checkpoints with $B_{\text{train}}\!\in\!\{8,16,32,64\}$ under the Equal-FLOPS protocol (identical sequence length, batch size, and number of optimizer updates). Each checkpoint is then evaluated with semi-autoregressive decoding at $B_{\text{infer}}\!\in\!\{8,16,32,64\}$, producing a $4{\times}4$ grid per dataset. Diagonal cells capture matched granularity; off-diagonal cells impose a controlled mismatch at inference.

In both experimental groups, performance peaks along the diagonal ($B_{\text{train}}{=}B_{\text{infer}}$) and degrades progressively as the mismatch grows. This diagonal dominance indicates that gains stem from aligning supervision granularity with the decoding procedure, rather than from any single block size. Near-diagonal cells degrade gracefully, suggesting robustness to minor misalignment. The factorial design removes compute confounds, attributing the observed pattern to training–inference alignment.

Figure: (a) GSM8K
Refer to caption: x9.png

### 4.4 Ablation Studies

Prefix Noise Analysis.
Figs. [7](https://arxiv.org/html/2508.19529v2#S4.F7)a–b examine the effect of corrupting prefix blocks during training.
Performance degrades progressively as $\pi_{\text{prefix}}$ increases, with the steepest drop at $\pi_{\text{prefix}}=1.0$.
This degradation pattern has an important nuance: when block sizes are large, noisy prefixes often have no effect because the first block (which has no prefix) is selected more frequently.
This explains why Blockwise SFT maintains reasonable performance even with substantial prefix corruption, because the fast convergence compensates for the effectively reduced training data.
The key insight is that maintaining clean prefix context during training mirrors the deterministic prefix available at inference, reinforcing proper conditioning behavior. Implementation details are provided in Appendix [A.7](https://arxiv.org/html/2508.19529v2#A1.SS7).

Figure: (a) Prefix noise on GSM8K
Refer to caption: x11.png

Suffix Leakage Analysis.
Figs. [8](https://arxiv.org/html/2508.19529v2#S4.F8)a–b reveal that suffix visibility is critically detrimental.
As $\pi_{\text{suffix}}$ decreases (more future tokens become visible), performance drops sharply.
At $\pi_{\text{suffix}}=0$ (fully visible suffix), Blockwise SFT degenerates to Classical SFT performance.
This dramatic degradation occurs because visible suffixes fundamentally corrupt the learning objective: instead of learning causal reasoning from prefix to active block, the model learns to extract answers from future context, which is a capability unavailable at inference.
This finding underscores that strict future masking is the primary mechanism behind Blockwise SFT’s success, ensuring the model develops genuine reasoning capabilities rather than spurious correlations with leaked information.

Figure: (a) Suffix leakage on GSM8K
Refer to caption: x13.png

## 5 Conclusion

We introduce Blockwise SFT, a simple recipe that aligns supervision with semi-autoregressive, blockwise decoding in discrete diffusion LMs. By freezing clean prefixes, masking only the active block, and strictly hiding future tokens, it removes the contextual shift and dependency leakage of classical SFT. We ground this method with a variational upper bound on blockwise likelihoods, unbiased block/time-sampled gradients, and an analysis of classical SFT’s mismatch bias (Theorems [3.1](https://arxiv.org/html/2508.19529v2#S3.Thmtheorem1)–[3.2](https://arxiv.org/html/2508.19529v2#S3.Thmtheorem2)). Empirically, fine-tuning on MetaMathQA with evaluation on GSM8K and MATH yields consistent gains under matched compute and token budgets, state-of-the-art head-to-head results, and diagnostics that track alignment (block-size consistency in §[4.3](https://arxiv.org/html/2508.19529v2#S4.SS3), prefix/suffix ablations in §[4.4](https://arxiv.org/html/2508.19529v2#S4.SS4)). The recipe is easy to adopt: no architectural changes, a drop-in loss, and practical single-step estimators (§[3.4](https://arxiv.org/html/2508.19529v2#S3.SS4)). Limitations and next steps include adaptive or uncertainty-aware block sizing, combining with preference or RL-style objectives, and extending the alignment principle to other semi- or non-autoregressive generators. Overall, honoring the decoder’s structural constraints in SFT is a core driver of performance in diffusion-based language models.

## References

- Arriola et al. (2025)
Marianne Arriola, Aaron Gokaslan, Justin T. Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, and Volodymyr Kuleshov.
Block diffusion: Interpolating between autoregressive and diffusion language models, 2025.
URL [https://arxiv.org/abs/2503.09573](https://arxiv.org/abs/2503.09573).
- Asada & Miwa (2025)
Masaki Asada and Makoto Miwa.
Addressing the training-inference discrepancy in discrete diffusion for text generation.
In Owen Rambow, Leo Wanner, Marianna Apidianaki, Hend Al-Khalifa, Barbara Di Eugenio, and Steven Schockaert (eds.), *Proceedings of the 31st International Conference on Computational Linguistics*, pp. 7156–7164, Abu Dhabi, UAE, January 2025. Association for Computational Linguistics.
URL [https://aclanthology.org/2025.coling-main.477/](https://aclanthology.org/2025.coling-main.477/).
- Austin et al. (2023)
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg.
Structured denoising diffusion models in discrete state-spaces, 2023.
URL [https://arxiv.org/abs/2107.03006](https://arxiv.org/abs/2107.03006).
- Biderman et al. (2024)
Dan Biderman, Jacob Portes, Jose Javier Gonzalez Ortiz, Mansheej Paul, Philip Greengard, Connor Jennings, Daniel King, Sam Havens, Vitaliy Chiley, Jonathan Frankle, Cody Blakeney, and John P. Cunningham.
Lora learns less and forgets less, 2024.
URL [https://arxiv.org/abs/2405.09673](https://arxiv.org/abs/2405.09673).
- Chen et al. (2023)
Jiaao Chen, Aston Zhang, Mu Li, Alex Smola, and Diyi Yang.
A cheaper and better diffusion language model with soft-masked noise, 2023.
URL [https://arxiv.org/abs/2304.04746](https://arxiv.org/abs/2304.04746).
- Chung et al. (2022)
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pellat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei.
Scaling instruction-finetuned language models, 2022.
URL [https://arxiv.org/abs/2210.11416](https://arxiv.org/abs/2210.11416).
- Cobbe et al. (2021)
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman.
Training verifiers to solve math word problems, 2021.
URL [https://arxiv.org/abs/2110.14168](https://arxiv.org/abs/2110.14168).
- Gong et al. (2023)
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong.
Diffuseq: Sequence to sequence text generation with diffusion models, 2023.
URL [https://arxiv.org/abs/2210.08933](https://arxiv.org/abs/2210.08933).
- Han et al. (2023)
Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov.
Ssd-lm: Semi-autoregressive simplex-based diffusion language model for text generation and modular control, 2023.
URL [https://arxiv.org/abs/2210.17432](https://arxiv.org/abs/2210.17432).
- Hendrycks et al. (2021)
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt.
Measuring mathematical problem solving with the math dataset, 2021.
URL [https://arxiv.org/abs/2103.03874](https://arxiv.org/abs/2103.03874).
- Hoogeboom et al. (2021)
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling.
Argmax flows and multinomial diffusion: Learning categorical distributions, 2021.
URL [https://arxiv.org/abs/2102.05379](https://arxiv.org/abs/2102.05379).
- Hu et al. (2021)
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.
Lora: Low-rank adaptation of large language models, 2021.
URL [https://arxiv.org/abs/2106.09685](https://arxiv.org/abs/2106.09685).
- Israel et al. (2025)
Daniel Israel, Guy Van den Broeck, and Aditya Grover.
Accelerating diffusion llms via adaptive parallel decoding, 2025.
URL [https://arxiv.org/abs/2506.00413](https://arxiv.org/abs/2506.00413).
- Li et al. (2022)
Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B. Hashimoto.
Diffusion-lm improves controllable text generation, 2022.
URL [https://arxiv.org/abs/2205.14217](https://arxiv.org/abs/2205.14217).
- Loshchilov & Hutter (2019)
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization, 2019.
URL [https://arxiv.org/abs/1711.05101](https://arxiv.org/abs/1711.05101).
- Micikevicius et al. (2018)
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu.
Mixed precision training, 2018.
URL [https://arxiv.org/abs/1710.03740](https://arxiv.org/abs/1710.03740).
- Nie et al. (2025)
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li.
Large language diffusion models, 2025.
URL [https://arxiv.org/abs/2502.09992](https://arxiv.org/abs/2502.09992).
- Ouyang et al. (2022)
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe.
Training language models to follow instructions with human feedback, 2022.
URL [https://arxiv.org/abs/2203.02155](https://arxiv.org/abs/2203.02155).
- Sahoo et al. (2024)
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov.
Simple and effective masked diffusion language models, 2024.
URL [https://arxiv.org/abs/2406.07524](https://arxiv.org/abs/2406.07524).
- Smith (2017)
Leslie N. Smith.
Cyclical learning rates for training neural networks, 2017.
URL [https://arxiv.org/abs/1506.01186](https://arxiv.org/abs/1506.01186).
- Wei et al. (2023)
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou.
Chain-of-thought prompting elicits reasoning in large language models, 2023.
URL [https://arxiv.org/abs/2201.11903](https://arxiv.org/abs/2201.11903).
- Weligalle (2025)
Ashen Weligalle.
Discrete diffusion models for language generation, 2025.
URL [https://arxiv.org/abs/2507.07050](https://arxiv.org/abs/2507.07050).
- Yu et al. (2024)
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T. Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu.
Metamath: Bootstrap your own mathematical questions for large language models, 2024.
URL [https://arxiv.org/abs/2309.12284](https://arxiv.org/abs/2309.12284).
- Zheng et al. (2024)
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong.
A reparameterized discrete diffusion model for text generation, 2024.
URL [https://arxiv.org/abs/2302.05737](https://arxiv.org/abs/2302.05737).

## Appendix A Appendix

### A.1 Full Hyperparameter Settings

Table [2](https://arxiv.org/html/2508.19529v2#A1.T2) lists the full hyperparameter configuration used in all experiments reported in the main paper.
Unless otherwise specified, parameters not listed (e.g., weight decay, gradient clipping) are set to their default values.
The configuration is kept identical across all experimental protocols (Equal-FLOPS, Equal-Tokens) and other studies to ensure fair comparison.

**Table 2: hyperparameter configuration for all experiments.**
| Category | Parameter | Value / Description |
| --- | --- | --- |
| Batch Size | Number of GPUs | $1$ |
| Per-device batch size | $32$ |  |
| Gradient accumulation steps | $1$ |  |
| Global batch size | $32$ |  |
| Learning Rate Schedule | Learning rate | $1\times 10^{-5}$ |
| Scheduler type | Warmup + Cosine decay |  |
| Warmup steps | $10\%$ of total steps |  |
| Cosine decay minimum ratio | $0.1$ |  |
| Masking | Mask ratio $\pi$ | $\mathrm{Uniform}(10^{-3},1)$ sampled per sequence / active block |
| Optimizer (AdamW) | $\beta_{1}$ | $0.95$ |
| $\beta_{2}$ | $0.99$ |  |
| Weight decay | 0 |  |
| LoRA | Rank $r$ | $256$ |
| $\alpha$ | $512$ |  |
| Dropout | $0.05$ |  |
| Target modules | Attention and MLP layers |  |
| Bias | None |  |

### A.2 Dataset Preprocessing and Evaluation Protocols

#### Training Set.

We follow the preprocessing instructions provided by the MetaMathQA dataset.
For each sample, the query field is inserted into the following instruction template:

The placeholder {instruction} is replaced with the actual query text.
We then concatenate this instruction with the response field using the model’s chat template to form the final training sequence.
To reduce compute and facilitate reproducibility, the maximum sequence length is set to $256$ tokens; samples exceeding this limit (fewer than $10\%$ of the dataset) are discarded.
The MetaMathQA dataset contains no samples from the GSM8K or MATH test sets.

#### Test Sets.

For evaluation, we apply the same instruction template to the prompts from the test splits of GSM8K and MATH(MATH-500), ensuring identical formatting between training and evaluation inputs.

#### Scoring Rules.

We report Pass@1 accuracy for both datasets using exact match:

- •
GSM8K: We extract the last numeric value from the model output (e.g., using a regex that matches the final signed/decimal number) and perform exact string match against the reference answer.
- •
MATH: No extraction or normalization is applied because of the complexity of the answer format; we compare the full model output to the reference answer by exact string match.

### A.3 Detailed Inference Settings

Table [3](https://arxiv.org/html/2508.19529v2#A1.T3) summarises the inference-time configurations used throughout our experiments, including both the default settings and the variations for the Block Size Consistency Study.
Unless otherwise noted, all other parameters follow the default configuration of the LLaDA-8B-Instruct model.

**Table 3: Inference-time settings for all experiments.**
| Parameter | Value / Range | Description |
| --- | --- | --- |
| Maximum new tokens | $128$ | Maximum number of tokens generated per example. |
| Block size | $\{8,16,32,64\}$ | Block Size Consistency Study. |
|  | $8$ | Ablation experiments. |
|  | $32$ (default) | Default for all other experiments. |
| Diffusion steps | $128$ (default) | Number of denoising steps during inference. |
| Other parameters | Default (LLaDA-8B-Instruct) | All remaining inference parameters follow the model’s official default configuration. |

### A.4 Additional Preliminaries and Notation

#### Tokenisation and sequences.

Let $\mathcal{V}$ be a finite vocabulary of size $V$. A tokenised sequence is $\mathbf{x}_{1:L}\in\mathcal{V}^{L}$ with $x_{i}\in\mathcal{V}$. Each training instance is

$$ $\mathbf{x}_{1:L}=\bigl[\mathbf{c}_{1:L_{c}};\mathbf{r}_{1:L_{r}}\bigr],\qquad L=L_{c}+L_{r}.$ $$

#### Masking operators.

Let $\mathbf{m}_{1:L}\in\{0,1\}^{L}$ with $m_{i}{=}0$ for observed and $m_{i}{=}1$ for masked. For Classical SFT we sample i.i.d. masks on all response positions:
$m_{i}\sim\mathrm{Bernoulli}(\pi)$ for $i\in\{L_{c}{+}1,\dots,L\}$.
For Blockwise SFT, with active block index $a$ and block size $B$,

$$ $m_{i}=\begin{cases}0,&i\leq L_{c}+B(a-1)\quad\text{(prefix)},\[2.0pt] \mathrm{Bernoulli}(\pi),&L_{c}+B(a-1)<i\leq L_{c}+Ba\quad\text{(active block)},\[2.0pt] 1,&i>L_{c}+Ba\quad\text{(suffix)}.\end{cases}$ $$

We use $\tilde{\mathbf{x}}=\mathbf{x}\odot(1-\mathbf{m})$ to denote the masked sequence.
(For variants used in ablations, see Eq. equation [10](https://arxiv.org/html/2508.19529v2#A1.E10) in Appendix [A.7](https://arxiv.org/html/2508.19529v2#A1.SS7).)

#### Discrete diffusion objective.

Let $t\in\{0,\dots,T\}$ index diffusion steps and $q_{t}(\mathbf{z}_{t}\mid\mathbf{x})$ be the forward noising kernel (relaxed categorical diffusion). The reverse model $\mathbf{p}_{\theta}(\mathbf{x}\mid\mathbf{z}_{t},t)$ predicts the original tokens at masked positions. The per–step cross-entropy and the weighted objective are

$$ $\mathcal{L}_{t}(\theta)=\!\mathbb{E}_{\mathbf{x},\mathbf{m},\mathbf{z}_{t}}\!\left[\mathrm{CE}\!\big(\mathbf{p}_{\theta}(\cdot\mid\mathbf{z}_{t},t),\mathbf{x}\big)\,\middle|\,\mathbf{m}\right],\qquad\sum_{t=1}^{T}\omega_{t}\,\mathcal{L}_{t}(\theta),$ $$

where $\{\omega_{t}\}_{t=1}^{T}$ are nonnegative weights.

#### Notation table.

See Table [4](https://arxiv.org/html/2508.19529v2#A1.T4).

**Table 4: Complete notation used in the paper.**
| Symbol | Meaning |
| --- | --- |
| $\mathcal{V},V$ | Vocabulary and its size |
| $\mathbf{x}_{1:L}$ | Full instruction–response sequence |
| $L_{c},L_{r}$ | Instruction and response lengths |
| $B,M$ | Block size and number of response blocks |
| $a$ | Index of the active block |
| $\mathbf{m}_{1:L}$ | Binary mask (1 = latent, 0 = observed) |
| $T,t$ | Total diffusion steps and current step |
| $q_{t},\mathbf{p}_{\theta}$ | Forward kernel and reverse model |
| $\pi$ | Masking rate (sequence or active-block level) |

### A.5 Practical Estimators and Weighting

#### Single- t t sampling.

To reduce cost, we may sample a single diffusion step $t$ with probability $\tilde{\omega}_{t}\propto\omega_{t}$ and multiply the per-step gradient by the normalizer $Z=\sum_{s=1}^{T}\omega_{s}$. Formally, letting

$$ $\widehat{\mathcal{L}}(\theta;\mathbf{x},a,t)=Z\;\mathbb{E}_{\mathbf{z}_{t}\sim q_{t}(\cdot\mid\mathbf{x})}\bigl[\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)\bigr]\quad\text{with}\quad t\sim\tilde{\omega},$ $$

we have $\nabla_{\theta}\,\mathbb{E}_{t\sim\tilde{\omega}}[\widehat{\mathcal{L}}(\theta;\mathbf{x},a,t)]=\nabla_{\theta}\,\sum_{t=1}^{T}\omega_{t}\,\mathbb{E}_{\mathbf{z}_{t}}[\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)]$.

#### Block sampling and importance weights.

If the active block is drawn from a non-uniform distribution $\rho(a)$, multiply the per-example gradient by $1/\rho(a)$ to obtain an unbiased estimator of the block-averaged objective (cf. Theorem [3.3](https://arxiv.org/html/2508.19529v2#S3.Thmtheorem3)).

### A.6 Proofs for Section 3.3

###### Proof of Theorem 3.1 .

Express the classical-SFT gradient as an expectation over prefix/suffix masking patterns with mask rate $\pi$. If the prefix is fully clean and the suffix fully hidden, the two gradients coincide. Otherwise they differ by at most $L_{\text{pre}}$ for any masked prefix token and by at most $L_{\text{suf}}$ for any unmasked suffix token by Lipschitzness. The probability that at least one prefix token is masked equals $1-(1-\pi)^{|\mathcal{I}_{\mathrm{prefix}}^{(a)}|}$; the probability that at least one suffix token remains unmasked equals $1-\pi^{|\mathcal{I}_{\mathrm{suffix}}^{(a)}|}$. Taking expectation and applying the triangle inequality yields the bound in equation [2](https://arxiv.org/html/2508.19529v2#S3.E2); averaging over $a$ extends the result to the population gradient bias.
∎

###### Proof of Theorem 3.2 .

Consider the joint over $(\mathbf{b}^{(a)},\{\mathbf{z}_{t}\}_{t=1}^{T})$ under the forward noising process $q_{t}(\cdot\mid\mathbf{x})$ restricted to indices $\mathcal{I}_{a}$, conditioning on $\mathbf{c\!ontext}^{(a)}$ and deterministically masking tokens outside $\mathcal{I}_{a}$ (suffix). The discrete-diffusion ELBO applied to this restricted chain gives

$$ $\displaystyle-\log p_{\theta}\!\bigl(\mathbf{b}^{(a)}\mid\mathbf{c\!ontext}^{(a)}\bigr)$ $\displaystyle\leq\sum_{t=1}^{T}\mathbb{E}_{q_{t}}\Bigl[\mathrm{KL}\!\bigl(q(\mathbf{z}_{t-1}\mid\mathbf{z}_{t},\mathbf{b}^{(a)})\,\|\,p_{\theta}(\mathbf{z}_{t-1}\mid\mathbf{z}_{t},\mathbf{c\!ontext}^{(a)})\bigr)$ $\displaystyle\qquad\quad-\log p_{\theta}\!\bigl(\mathbf{b}^{(a)}\mid\mathbf{z}_{1},\mathbf{c\!ontext}^{(a)}\bigr)\Bigr]+C.$ (9) $$

Standard manipulations for discrete diffusion convert the KL reconstruction terms into tokenwise cross-entropies on $\mathcal{I}_{a}$, producing a weighted sum $\sum_{t}\omega_{t}\,\mathbb{E}_{q_{t}}[\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)]$; $C$ and the weights depend only on the forward process and are independent of $\theta$. Taking expectation over data and summing across $a$ yields $\mathcal{R}_{\text{block}}(\theta)\leq\mathcal{L}_{\text{BW--SFT}}(\theta)+C^{\prime}$ with $C^{\prime}$ independent of $\theta$.
∎

###### Proof of Theorem 3.3 .

Taking total expectation over $a\sim\rho$ and $t\sim\tilde{\omega}$ and using linearity of $\nabla_{\theta}$,

$$ $\mathbb{E}\bigl[\widehat{g}(\theta)\bigr]=\sum_{a=1}^{M}\rho(a)\frac{1}{\rho(a)}\sum_{t=1}^{T}\tilde{\omega}_{t}\Bigl(\sum_{s=1}^{T}\omega_{s}\Bigr)\nabla_{\theta}\mathbb{E}_{\mathbf{z}_{t}}[\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)]=\nabla_{\theta}\mathbb{E}_{\mathbf{x},a}\sum_{t=1}^{T}\omega_{t}\mathbb{E}_{\mathbf{z}_{t}}[\tilde{\mathcal{L}}_{t}(\theta;\mathbf{x},a)],$ $$

which equals $\nabla_{\theta}\,\mathcal{L}_{\text{BW--SFT}}(\theta)$. Finally, using the standard diffusion–ELBO correspondence—that the weighted per-step reconstruction terms act as a faithful surrogate for the $t{=}0$ blockwise negative log-likelihood (see Theorem [3.2](https://arxiv.org/html/2508.19529v2#S3.Thmtheorem2))—we may identify the inner gradient with that of the $t{=}0$ blockwise NLL; hence the estimator targets $\nabla_{\theta}\,\mathcal{R}_{\text{block}}(\theta)$ up to an overall positive scalar.
∎

###### Proof of Theorem 3.1 .

###### Proof of Theorem 3.2 .

###### Proof of Theorem 3.3 .

### A.7 Ablation Implementation Details

We formalize the two ablation settings (Noisy Prefix and Leaky Suffix) as direct modifications to the masking rule in Blockwise SFT (§[3.4](https://arxiv.org/html/2508.19529v2#S3.SS4)).
Recall that in standard Blockwise SFT, the binary mask $m_{i}$ for token $i$ is defined as:

$$ $m_{i}=\begin{cases}0,&i\leq L_{c}+B(a-1)\quad\text{(prefix)},\[2.0pt] \mathrm{Bernoulli}(\pi),&L_{c}+B(a-1)<i\leq L_{c}+Ba\quad\text{(active block)},\[2.0pt] 1,&i>L_{c}+Ba\quad\text{(suffix)},\end{cases}$ (10) $$

where $\pi\in(0,1)$ is the active-block masking rate, fixed for a given sample.
Let $\mathcal{I}_{\mathrm{prefix}}=\{L_{c}+1,\dots,L_{c}+B(a-1)\}$ denote all prefix tokens (after the prompt and before the active block),
$\mathcal{I}_{a}$ the indices of the active block,
and $\mathcal{I}_{\mathrm{suffix}}=\{L_{c}+Ba+1,\dots,L\}$ the suffix tokens.

#### 1. Noisy Prefix.

Given a fixed $\pi_{\mathrm{prefix}}\in[0,1]$ for the entire experiment, the mask $m_{i}$ is modified from equation [10](https://arxiv.org/html/2508.19529v2#A1.E10) as:

$$ $m_{i}=\begin{cases}\mathrm{Bernoulli}(\pi_{\mathrm{prefix}}),&i\in\mathcal{I}_{\mathrm{prefix}},\[3.0pt] \mathrm{Bernoulli}(\pi),&i\in\mathcal{I}_{a},\[3.0pt] 1,&i\in\mathcal{I}_{\mathrm{suffix}},\end{cases}$ $$

injecting random noise into the prefix tokens and simulating contextual distribution shift.

#### 2. Leaky Suffix.

Given a fixed $\pi_{\mathrm{suffix}}\in[0,1]$ for the entire experiment, the mask $m_{i}$ is modified from equation [10](https://arxiv.org/html/2508.19529v2#A1.E10) as:

$$ $m_{i}=\begin{cases}0,&i\in\mathcal{I}_{\mathrm{prefix}},\[3.0pt] \mathrm{Bernoulli}(\pi),&i\in\mathcal{I}_{a},\[3.0pt] \mathrm{Bernoulli}(\pi_{\mathrm{suffix}}),&i\in\mathcal{I}_{\mathrm{suffix}},\end{cases}$ $$

allowing the model to observe a random fraction of future tokens beyond the active block, thus introducing block-level dependency leakage.

In both cases, $\pi_{\mathrm{prefix}}$ and $\pi_{\mathrm{suffix}}$ are fixed for the duration of an experiment and applied identically to all training samples.

### A.8 Disclosure of Large Language Model Usage

Large Language Models were used exclusively for grammatical refinement and language polishing of this manuscript. All research design, methodology, experiments, analysis, and scientific conclusions are original contributions. The LLMs provided no substantive input on technical content, experimental design, or theoretical derivations.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 4
presentation: 3
questions: - What about baselines like the following?
    - (i) Causal SFT: classical SFT but with a causal mask (clean prefix, hidden suffix) over the *entire* response
    - (ii) Blockwise-but-bidirectional: blockwise loss but allowing future visibility
    - (iii) Span-corruption within a causal window (T5-style) matched to block size

-  The $(B_{\text{train}})$ vs $(B_{\text{infer}})$ grid is promising; what are the results for $B=1$ (degenerating to token-wise) and very large $B$ (near fully parallel)?

- Are there any failure modes for Blockwise SFT? When does Blockwise SFT hurt (e.g., very short responses, high-entropy next blocks)?
soundness: 2
strengths: - The paper nails a real but often-overlooked problem: standard SFT trains on full sequences with bidirectional context, but blockwise inference sees only a clean prefix and hides the future. The proposed Blockwise SFT fixes this by supervising only the active block with a clean prefix and hidden suffix. The theoretical contributions are solid—gradient-bias analysis (Theorem 3.1), a variational upper bound matching the inference factorization (Theorem 3.2), and an unbiased single-block estimator (Theorem 3.3).

- The experiments are well-designed with proper controls: EQUAL-FLOPS and EQUAL-TOKENS setups isolate the effect of the objective itself. Results are consistently better (GSM8K: 76.0 vs. 70.8, MATH: 34.2 vs. 32.6) with smoother training. The diagnostics are thorough—the block-size consistency grid shows performance peaks when training and inference block sizes match, and ablations confirm that prefix cleanliness and strict future masking are both critical. Algorithm 1 is simple, requires no architecture changes.

- The writing is clear and efficient. Figures communicate the core idea instantly, theorems come with intuition, and experimental protocols include concrete examples that make FLOPs/token accounting transparent.

- This work matters because blockwise decoding is now standard in discrete diffusion LMs. A drop-in training fix that mirrors deployment behavior has immediate practical value and points toward a broader principle: match your supervision structure to your decoding structure. The simplicity of the method plus strong empirical results make it likely to be adopted.
summary: This paper identifies a training–inference mismatch for discrete diffusion language models (DLMs) that are decoded semi‑autoregressively in fixed‑size blocks but are trained with bidirectional, full‑sequence random masking. The authors propose Blockwise SFT, a drop‑in training objective that (i) partitions responses into blocks, (ii) samples one active block per step, (iii) keeps the prefix clean (no noise), (iv) fully hides the future, and (v) computes loss only within the active block, thereby aligning supervision with the deployed decoding regime. They analyze classical SFT’s bias under this mismatch (Theorem 3.1), derive a variational upper bound on the blockwise negative log‑likelihood (Theorem 3.2), and show an unbiased stochastic gradient estimator when sampling a block and diffusion step (Theorem 3.3). Experiments fine‑tuning on MetaMathQA and evaluating on GSM8K / MATH report consistent Pass@1 gains over classical SFT and several diffusion‑SFT variants under matched FLOPs and matched supervised‑token budgets; block‑size consistency studies find peak accuracy when training and inference block sizes match, and prefix/suffix ablations indicate that preserving a clean prefix and strictly hiding future tokens are key. The work is simple, architecture‑agnostic, and argues that matching supervision granularity to blockwise decoding materially improves diffusion LMs.
weaknesses: 1. Loss inside the active block (Eq. (5))

   The paper defines

   $$
   \tilde{\mathcal{L}}_t(\theta; \mathbf{x}, a) = -\sum_{i\in\mathcal{I}_a} \log p_\theta(x_i \mid \mathbf{z}_t, t),
   $$

   yet Algorithm 1 samples an intra-block mask ($m_i\sim\mathrm{Bernoulli}(\pi)$) and states that loss is "only on the active block." In discrete diffusion / masked-LM style training, loss is normally computed only on masked positions; otherwise, unmasked tokens create an identity path and can yield degenerate gradients. The indicator ($\mathbf{1}[m_i=1]$) that appears in Eq. (1) is missing in Eq. (5), and Eq. (6) omits the expectation over the mask distribution.

2. Conditioning on context

   Eq. (5) conditions $p_\theta(\cdot \mid \mathbf{z}_t,t)$ but omits the explicit context (clean prefix, hidden suffix), whereas Eq. (9) does condition on $\text{context}^{(a)}$.

3. In Thm. 3.1, the bound scales with the probability that at least one prefix token is corrupted, $1-(1-\pi)^{|\mathcal{I}_{\text{prefix}}|}$, and the probability that at least one suffix token leaks, $1-\pi^{|\mathcal{I}_{\text{suffix}}|}$. This is a very loose surrogate for the actual perturbation magnitude, which typically scales with the expected number of corrupted tokens, e.g., $\pi\cdot|\mathcal{I}_{\text{prefix}}|$ and $(1-\pi)\cdot|\mathcal{I}_{\text{suffix}}|$. Moreover, the theorem does not quantify the third mismatch ("diluted supervision"), even though the text highlights it as critical. Finally, gradients being "$L_{\text{pre}}$- and $L_{\text{suf}}$-Lipschitz" with respect to discrete masking patterns is a strong, nonstandard assumption that needs justification.

4. In Thm. 3.2, it is correct that a diffusion ELBO can upper-bound the conditional block NLL. However, the proof text for Thm. 3.3 asserts that the inner gradient can be "identified with that of the $t{=}0$ blockwise NLL … hence the estimator targets $\nabla_\theta \mathcal{R}_{\text{block}}$ up to a scalar." This is not implied by the ELBO inequality: minimizing a surrogate bound does not make its gradient equal (or proportional) to the true NLL gradient.

5. In Thm. 3.3, the estimator multiplies by $\sum_s \\omega_s$ when sampling $t\sim\\tilde{\\omega}\propto\\omega$, which is fine, but the theorem statement should be explicit that unbiasedness is for the bound in Eq. (6) (not for the true NLL) and that the expectation is also over $\\mathbf{m}$ when masks are stochastic.

6. §3.4 states "In practice we use sequence length $L=128$," while Appendix A.2 says the max length is 256, and the running example uses $L_c=32$, $L_r=96$ ($L=128$). It is unclear which length was actually used in reported results.

7. Claims of generality are supported only on math datasets. Blockwise alignment may behave differently on open-ended generation (coding, dialogue).

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How does Blockwise SFT perform when combined with RLHF or preference-based fine-tuning ?
2. Is the performance gain mostly due to reduced prefix noise, or does strict hidden-suffix causality contribute more?  
3. Can adaptive block sizes during training (e.g., uncertainty-driven scheduling) outperform fixed ones?  
4. What is the actual training-time overhead of sampling one active block per step versus full-sequence masking?  
5. Does the variational bound (Theorem 3.2) still hold under non-uniform or variable-length block sampling?
soundness: 3
strengths: - **Clear motivation & strong alignment insight:** The paper precisely diagnoses the training–inference mismatch in diffusion LMs and proposes a clean fix grounded in the decoding procedure.  
- **Theory–practice coherence:** Provides formal bias analysis, variational bound, and unbiased gradient theorem; yet remains simple to implement (mask-only change).  
- **Empirical rigor:** Demonstrates consistent improvements across reasoning datasets, with ablations (block size, prefix noise, suffix leakage) directly supporting the central hypothesis.
summary: This paper introduces **Blockwise SFT**, a training-time objective designed to align **diffusion language models (DLMs)** with their **semi-autoregressive, blockwise decoding** behavior.  
Classical SFT randomly masks tokens across full responses, leading to a mismatch: noisy prefixes, suffix leakage, and token-level supervision that misaligns with block-level generation.  
Blockwise SFT resolves this by training only on a **single active block per step** — freezing all preceding tokens (clean prefix), fully hiding future ones (no leakage), and computing loss only within the active block.  
The authors theoretically derive it as a **variational upper bound on blockwise likelihood**, prove unbiased gradient estimation, and implement it without architectural change.  
Empirical results on **GSM8K**, **MATH**, and **MetaMathQA** show consistent gains: +5.2 pts on GSM8K and +1.6 pts on MATH over the strongest baselines, validating that aligning supervision granularity with decoding yields measurable improvements.
weaknesses: ## Major
- **Limited architectural coverage.** All experiments use only **LLaDA-8B-Instruct** with LoRA fine-tuning. The claim of being “architecture-agnostic” is unsupported without testing on other diffusion backbones such as **BlockDiffusion**, **APD**, or **RDM**.  
- **Scope restricted to reasoning tasks.** Evaluation is limited to GSM8K, MATH, and MetaMathQA, leaving uncertainty about performance on open-ended or dialogue data.  
- **Lack of scaling analysis.** The paper does not study behavior under larger models (e.g., 30B+) or longer block sizes (e.g., 512+), limiting insight into scalability.  
- **Theoretical framing stops short of efficiency guarantees.** While unbiasedness is proven, practical convergence or compute-efficiency trade-offs are not quantified.

## Minor
- **Compute fairness.** EQUAL-FLOPS and EQUAL-TOKENS protocols are well explained but omit wall-clock runtime comparisons.  
- **Notation heaviness.** Some derivations (Eq. 7–8) obscure intuition for non-diffusion audiences.  
- **Ablation breadth.** Although prefix/suffix studies are thorough, broader tests on other modalities or instruction-following data would strengthen the empirical story.

---

confidence: 3
contribution: 1
initial_rating: 2
presentation: 1
questions: I am not fully convinced about the necessity of Blockwise SFT. If blockwise generation is already established, wouldn’t directly training with a block diffusion objective be more natural and efficient—while also enabling the use of KV caching at inference?
soundness: 1
strengths: The paper is clearly written and well-organized. The methodology and theoretical sections are presented with both formal rigor and intuitive explanations. The algorithm and the diagrams make the proposed training recipe easy to follow, which enhances readability.
summary: This paper proposes Blockwise SFT, a supervised fine-tuning method for discrete diffusion language models that aligns training with semi-autoregressive, blockwise decoding used at inference. Classical SFT applies random full-sequence masking and bidirectional attention,  Blockwise SFT resolves this by freezing clean prefixes, masking only one active block per step, and fully hiding future tokens, optimizing loss solely within that block. The authors provide theoretical guarantees—showing that the new objective gives a variational upper bound on the true blockwise likelihood and yields unbiased gradient estimates—and demonstrate consistent empirical gains on GSM8K, MATH, and MetaMathQA benchmarks.
weaknesses: - Limited novelty: The proposed Blockwise SFT can be viewed as a straightforward adaptation of existing ideas for aligning training objectives with blockwise decoding. While the paper frames the problem clearly, the methodological innovation appears incremental.

- Training inefficiency: Since each training step only supervises one active block, the number of effective supervised tokens per update is significantly lower than in standard SFT under the same FLOP budget.

- Missing baselines: The experimental comparisons lack a direct baseline using block diffusion.

---

confidence: 5
contribution: 2
initial_rating: 2
presentation: 3
questions: 1. What specific advantages does blockwise SFT offer compared to block diffusion models? 
2. When employing alternative sampling strategies—such as reducing EOS generation probability without blockwise inference—does blockwise SFT still demonstrate gains over classical SFT?
soundness: 2
strengths: 1. The writing is clear and easy to follow. 
2. The ablation study is comprehensive and addresses several of my concerns.
summary: This paper proposes a simple yet effective strategy for supervised fine-tuning of diffusion-based large language models such as LLaDA and Dream. Specifically, the authors introduce a blockwise SFT approach where, during training, a single block is selected while masking information from future blocks and preserving information from preceding blocks without masking. The authors argue that this training strategy better aligns with the blockwise inference process inherent to these models. Experimental results demonstrate that the proposed blockwise SFT yields notable improvements over classical SFT methods.
weaknesses: 1. I believe this method offers limited advantages compared to block diffusion models. Block diffusion models, while requiring twice the number of training tokens, enable training across all blocks and kv cache in the inference process. In contrast, blockwise SFT only trains one block at a time, which restricts its training coverage. Consequently, I find the contribution of this work to be limited. 
2. Blockwise inference is not strictly necessary for LLaDA. During LLaDA's training, EOS tokens are padded to enable variable-length generation, which can lead to premature EOS generation during inference, resulting in truncated responses. However, several strategies can mitigate this issue without requiring blockwise inference. For instance, suppressing EOS token probability in early inference stages (as demonstrated in https://arxiv.org/pdf/2509.23924) can effectively address this problem. This undermines the necessity of the paper's core motivation.


Decision: Reject