Match the Canvas, Not Just the Answer: Predictability-Aware SFT for Diffusion Language Models

Miao Li; Hanyang Jiang; Hengyu Fu; Baihe Huang; Yating Liu; Yuhang Cai; Changkun Guan; Zhongzhu Chen

cs.LGOct 13, 2026

Match the Canvas, Not Just the Answer: Predictability-Aware SFT for Diffusion Language Models

Authors: Miao Li, Hanyang Jiang, Hengyu Fu, Baihe Huang, Yating Liu, Yuhang Cai, Changkun Guan, Zhongzhu Chen,

Organizations: Georgia Institute of Technology · University of California, Berkeley · University of Chicago · University of Michigan

Abstract

Diffusion language models (DLMs) enable flexible parallel generation by resolving tokens according to the model's evolving predictability rather than a fixed left-to-right order. This flexibility, however, introduces training choices absent from standard autoregressive teacher forcing: which reference tokens are visible when a target is supervised and how much generation space is allocated. Standard supervised fine-tuning (SFT) typically makes these choices through random masking, independently of inference. Our analysis shows how this training-inference mismatch can misdirect supervision. Random masking can conceal prediction errors or penalize uncertainty that decoding would otherwise resolve, while the attention-visible input length can change which response formulations the model favors, even when the underlying solution is unchanged. Motivated by these findings, we propose Predictability-Aware Canvas Selection and Supervision (PACS), an SFT method. PACS uses the model's own predictability to select each example's attention-visible input length and supervises reference tokens immediately before their confidence-guided revelation. Across reasoning and code-generation benchmarks, PACS consistently achieves the highest mean accuracy among the compared SFT methods, with gains of up to 12.84 percentage points over vanilla SFT on individual benchmarks. Code will be released upon publication.

Figures & tables

Figure 1: The training canvas changes reference supervision. Top: A canvas includes prompt tokens (green), visible answer tokens (blue), and masked positions (red). A: Visible reference content can conceal or expose a commitment-time disagreement. B: Canvas extent can alter the prediction problem, inducing an extra stopping signal on a mismatched canvas but not on the matched canvas.

Figure 2: Longer canvases induce longer responses. Realized response length increases with allocated generation length.

Figure 3: Changing context around commitment can create or conceal reference disagreement. LLaDA-8B-Instruct on GSM8K. A: Insufficient context. Hiding reference tokens revealed in the preceding step makes 37.3% of initially correct predictions wrong. B: Future context leakage. Revealing one reference token scheduled for the following reveal step early makes 51.2% of initially wrong predictions correct.

Figure 4: Canvas-extent mismatch increases reference cross-entropy. Target responses are generated at source generation length

K_{0}=384

, and

\Delta\mathrm{CE}

is the change of CE compared to that evaluated at

K_{0}

. A (center): Mean

\Delta\mathrm{CE}

across generation lengths, evaluated along each candidate’s own reference-conditioned trajectory. B (left) and C (right): Detailed blockwise

\Delta\mathrm{CE}

across response positions for

K=256

(blue) and

K=512

(orange).

Method	GSM8K	MATH500	MBPP	HumanEval
Initial	75.44	36.40	40.86	37.80
Vanilla SFT	76.27	37.00	49.03	38.41
Semi-AR block	78.20	38.40	48.60	43.90
LIFT	77.63	37.00	46.69	39.63
TRIMS	77.63	38.00	52.92	40.24
PACS	77.94	39.00	49.42	47.56

Table 1: LLaDA-8B-Instruct results (%) at

K_{\mathrm{eval}}=256

Method	GSM8K	MATH500	MBPP	HumanEval
Initial	80.52	43.80	45.53	51.83
Vanilla SFT	80.06	44.20	60.31	54.27
Semi-AR block	80.21	43.60	61.09	58.53
LIFT	80.36	44.80	50.58	52.44
TRIMS	80.44	45.60	52.14	53.66
PACS	80.52	49.20	61.48	60.98

Table 2: Dream-7B-Instruct results (%) at

K_{\mathrm{eval}}=256

Canvas rule	Reveal order	GSM8K	MATH500	MBPP	HumanEval
Batch-max	Random	76.27	37.00	45.14	39.63
Batch-max	Confidence-guided (ours)	77.48	38.00	47.08	40.90
Trajectory-selected (ours)	Random	77.03	37.40	47.86	44.51
Trajectory-selected (ours)	Confidence-guided (ours)	77.94	39.00	49.42	47.56

Table 3: Contributions of canvas selection and reveal ordering. LLaDA at

K_{\mathrm{eval}}=256

Figure 5: Canvas choice and reference predictability. Target CE (blue bars, left) and post-SFT test accuracy (orange diamonds, right) for four canvas-length rules on MATH500 and HumanEval with LLaDA-8B-Instruct at

K_{\mathrm{eval}}=256

Figure 6: Canvas search extent at

K_{\mathrm{eval}}=256

. Bars show average reference CE (left axis); orange diamonds show held-out validation accuracy (right axis).

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Learning rate	$5\times 10^{-6}$
Learning rate scheduler type	Cosine
Adam $\beta$ parameters	$\beta_{1}=0.9,\ \beta_{2}=0.999$
Gradient accumulation steps	4
Per-device train batch size	2
Epochs	1

Appendix

Table A.1: Hyperparameters used for PACS.

Figure A.1: Selected-bin proportions across candidate families. Each horizontal bar is normalized independently to

100\%

and shows the fraction of examples for which PACS selects each candidate in the nested family

\Lambda_{K}(P,A)

. We use

L_{\mathrm{ref}}=|P|+|A|

L_{\min}=B\left\lceil L_{\mathrm{ref}}/B\right\rceil

, and

B=32

, so the bins correspond to

L_{\mathrm{ref}}

L_{\min}

L_{\min}+B

L_{\min}+2B

, and

L_{\min}+3B

. Segment labels give percentages rounded to one decimal place. Bins not contained in a given

\Lambda_{K}

have no segment, and colors identify the same candidate lengths across datasets and families.

\Lambda_{4}

is the main configuration.

	GSM8K			MATH500
Model / Seq Len	128	256	512	128	256	512
LLaDA-Instruct	76.12	75.44	79.08	35.60	36.40	38.80
+ Vanilla	77.03	76.27	78.77	35.60	37.00	39.40
+ Semi-AR Block	76.12	78.20	79.23	34.00	38.40	39.20
+ LIFT	75.89	77.63	78.85	35.40	37.00	38.00
+ TRIMS	76.88	77.63	79.53	34.80	38.00	39.60

Appendix

Table A.2: Complete LLaDA-8B-Instruct results (%) across

K_{\mathrm{eval}}\in\{128,256,512\}

	GSM8K			MATH500
Model / Seq Len	128	256	512	128	256	512
Dream-Instruct-7B	70.66	80.52	80.74	37.20	43.80	44.20
+ Vanilla	70.35	80.06	80.36	37.20	44.20	44.80
+ Semi-AR Block	70.43	80.21	80.06	37.00	43.60	43.40
+ LIFT	71.41	80.36	80.81	36.00	44.80	45.20
+ TRIMS	71.26	80.44	80.59	37.40	45.60	45.40

Appendix

Table A.3: Complete Dream-7B-Instruct results (%) across

K_{\mathrm{eval}}\in\{128,256,512\}

Method	GSM8K	MATH500	MBPP	HumanEval
Vanilla SFT	75.82	36.20	42.41	36.59
Semi-AR block	77.56	37.60	41.63	37.80
LIFT	77.10	36.40	40.08	38.41
TRIMS	77.03	37.20	43.58	35.98
PACS	77.41	38.20	42.80	40.24

Appendix

Table B.1: Performance on the alternative training corpus. All methods use Math14k and KodCode and are evaluated using LLaDA-8B-Instruct at

K_{\mathrm{eval}}=256

. Bold marks the highest score in each benchmark column.

Method	Benchmark	Reported	Reproduction
LIFT 3	GSM8K	79.40	79.30
LIFT 3	MATH500	38.40	38.20
TRIMS	MBPP	41.80	41.63
TRIMS	HumanEval	37.20	36.59

Appendix

Table B.2: Sanity reproduction of prior SFT baselines on LLaDA-8B-Instruct. Reported values are taken from the corresponding papers; each reproduction follows the training and evaluation configuration of the original result.

Method	Runtime
Vanilla SFT	$\approx$ 1 h 13 min
Semi-AR block	$\approx$ 1 h 16 min
LIFT	$\approx$ 8 h 24 min
TRIMS	$\approx$ 1 h 21 min
PACS	$\approx$ 14 h 55 min

Appendix

Table B.3: Runtime on the OpenCoder workload.

	4 GPU-hours		8 GPU-hours		12 GPU-hours
Method	MATH500	HumanEval	MATH500	HumanEval	MATH500	HumanEval
Vanilla SFT	37.20	39.02	37.40	39.63	37.20	39.02
Semi-AR block	38.40	43.90	38.60	44.51	38.60	43.90
LIFT	36.40	38.41	37.00	39.63	37.20	40.85
TRIMS	37.60	40.24	38.20	40.85	38.20	39.63
PACS	37.40	40.24	38.80	43.90	39.00	46.34

Appendix

Table B.4: Compute-matched comparison on LLaDA-8B-Instruct at

K_{\mathrm{eval}}=256

. Budgets include method-specific preprocessing and optimization, excluding shared corpus construction. Accuracy is reported on MATH500 (

n=500

) and HumanEval (

n=164

). Bold marks the highest score for each benchmark and budget.

Supervision	Length rule	GSM8K	MATH500	MBPP	HumanEval
Random-mask SFT	Batch-max	76.27	37.00	49.03	38.41
Random-mask SFT	Trajectory-selected	76.50	37.20	49.42	39.02
Semi-AR block	Batch-max	78.20	38.40	48.60	43.90
Semi-AR block	Trajectory-selected	78.39	38.80	49.03	45.73
Confidence trajectory	Batch-max	77.48	38.00	47.08	40.90
Confidence trajectory	Trajectory-selected	77.94	39.00	49.42	47.56

Appendix

Table B.5: Canvas selection under different supervision rules. LLaDA-8B-Instruct at

K_{\mathrm{eval}}=256

. For trajectory-selected rows, canvas lengths are chosen from the same four-candidate family used in the main experiments.

Figure B.1: Semi-AR block canvas-choice test accuracy. Post-SFT test accuracy (orange diamonds) on MATH500 and HumanEval at

K_{\mathrm{eval}}=256

for four canvas-length rules.

Diffusion Language Models

Diffusion language models generate and refine text through iterative denoising, enabling bidirectional context and parallel token updates.

See all 473 papers

Learnability-Informed Fine-Tuning of Diffusion Language Models
May 21, 2026Shubham Parashar, Atharv Chagi, Jacob Helwig +5Supervised Fine-Tuning Diffusion Language Models
Department of Computer Science and Engineering, Texas A&M University, College Station, TX, USA · Department of Electrical and Computer Engineering, Texas A&M University, College Station, TX, USA
CARVE: Verified Expansion for Variable-Length Generation in Diffusion Language Models
Aug 31, 2026Wail Bouhedja, Amr Mohamed, Guokan ShangMasked Diffusion Models LLM Inference Acceleration
1MBZUAI · 2Sorbonne Université · 3Ecole Polytechnique
Unlocking Prompt Infilling Capability for Diffusion Language Models
Apr 4, 2026Yoshinari Fujinuma, Keisuke SakaguchiDiffusion Language Models Masked Diffusion Models
Patronus AI · Tohoku University, RIKEN

May 21, 2026J/K move · Enter open · S save

Learnability-Informed Fine-Tuning of Diffusion Language Models

Shubham Parashar, Atharv Chagi, Jacob Helwig, Lakshmi Jotsna+4

Department of Computer Science and Engineering, Texas A&M University, College Station, TX, USA · Department of Electrical and Computer Engineering, Texas A&M University, College Station, TX, USA

Supervised Fine-Tuning Diffusion Language Models Diffusion Language Model Training

Open paper AskPDF arXiv

We aim to improve the reasoning capabilities of diffusion language models (DLMs). While SFT is a popular post-training recipe for autoregressive models, its use in DLMs faces challenges and can even hurt performance, though the underlying causes remain understudied. Our analysis reveals that vanilla SFT overlooks learnability, namely what and when tokens are learned. Specifically, rare tokens are difficult to learn when most of the input is masked, whereas it is straightforward and thus of little value to learn common tokens when most of the input is unmasked. Motivated by our analysis, we propose LIFT, an efficient SFT-based post-training algorithm for DLMs. LIFT learns easy tokens when most of the input is masked and hard tokens when more context is available, thus aligning the training with the information available at different diffusion time steps. Our results show that LIFT outperforms existing SFT baselines across six reasoning benchmarks, achieving up to a 3x relative gain on AIME'24 and AIME'25. Our code is publicly available at https://github.com/divelab/LIFT.

Abstract

Diffusion Language Models

Explore similar work

Learnability-Informed Fine-Tuning of Diffusion Language Models

CARVE: Verified Expansion for Variable-Length Generation in Diffusion Language Models

Unlocking Prompt Infilling Capability for Diffusion Language Models