TIPS: A Tiny Inference-time Policy for Scaling in Diffusion Language Models

Ya-Shi Zhang; Jian Tang; Pierre-André Noël; Oleksiy Ostapenko

cs.LGOct 15, 2026

TIPS: A Tiny Inference-time Policy for Scaling in Diffusion Language Models

Authors: Ya-Shi Zhang, Jian Tang, Pierre-André Noël, Oleksiy Ostapenko

Organizations: ServiceNow AI Research · Mila - Quebec AI Institute · HEC Montréal

Abstract

Masked diffusion language models expose an inference-time choice: which positions become visible after each denoising pass. Many existing decoding methods use hand-designed rules or require visibility to grow monotonically. We present a $\textbf{T}$ iny $\textbf{I}$ nference-time $\textbf{P}$ olicy for $\textbf{S}$ caling ( $\textbf{TIPS}$ ), an approximately $600$ K-parameter controller that learns both unmasking and remasking while keeping the underlying language model frozen. TIPS uses visibility-aware confidence features and a shallow transformer to select which tokens should be visible, while a fixed schedule controls the expected visible fraction. We train the controller through supervised warmup followed by reinforcement learning from terminal correctness rewards. On GSM8K and MATH-500 with LLaDA-8B-Instruct, the full-canvas and two-block TIPS variants achieve $58.46\%$ and $57.97\%$ micro-averaged accuracy at a budget of $32$ decoder evaluations, exceeding the strongest baseline in our main comparison by $12.1$ and $11.6$ percentage points, respectively. Both variants are trained at a budget of $32$ evaluations and remain competitive at higher budgets without retraining. These results demonstrate that learning when to reveal and reconsider tokens can improve the accuracy-compute trade-off without updating the decoder backbone.

Figures & tables

\qquad\qquad\qquad x_{t+1,i}=\begin{cases}\texttt{[MASK]},&a_{t,i}=0,\\ x_{t,i},&a_{t,i}=1,\ m_{t,i}=0,\\ \widehat{x}_{t,i},&a_{t,i}=1,\ m_{t,i}=1.\end{cases}

Algorithm 1 TIPS: Inference

Method	Blocks	4	8	16	32	64	128	256	AUC
Fast-dLLM	$8\!\times\!32$	17.48	17.48	14.73	37.93	65.59	67.01	67.01	40.83
WINO	$1\!\times\!256$	18.14	21.06	34.80	46.34	49.26	49.37	49.37	39.10
UNCODE	$1\!\times\!256$	12.26	14.35	21.50	29.74	53.66	63.66	67.73	37.15
Jazbec et al. (2026)	$1\!\times\!256$	5.63	7.42	9.18	12.39	47.72	64.80	64.63	29.44
TIPS full	$1\!\times\!256$	16.46 [-1.5pt] $\pm$ 1.79	29.12 [-1.5pt] $\pm$ 1.68	48.05 [-1.5pt] $\pm$ 0.73	58.46 [-1.5pt] $\pm$ 0.10	63.09 [-1.5pt] $\pm$ 0.50	65.32 [-1.5pt] $\pm$ 0.55	66.28 [-1.5pt] $\pm$ 0.71	50.90
TIPS block	$2\!\times\!128$	12.53 [-1.5pt] $\pm$ 0.45	20.95 [-1.5pt] $\pm$ 0.99	42.12 [-1.5pt] $\pm$ 2.04	57.97 [-1.5pt] $\pm$ 1.23	64.49 [-1.5pt] $\pm$ 0.80	67.12 [-1.5pt] $\pm$ 0.28	68.18 [-1.5pt] $\pm$ 0.36	48.83

Table 1: Micro-averaged test accuracy (%) on GSM8K and MATH-500 (1,819 problems) at different decoder-call budgets. For compactness, we display a subset of baselines that preserves the highest or competitive accuracy at every evaluated budget across the full comparison suite. TIPS accuracy entries show the mean with

\pm

standard error underneath, across three training-seed means; each seed mean averages three inference draws. AUC is the normalized area under the accuracy versus

\log_{2}

-budget curve over budgets 4–256, computed using trapezoidal quadrature and reported on the same 0–100 scale as accuracy. TIPS AUC values are computed from the mean accuracy curves. Bold indicates the highest accuracy among all evaluated configurations in each budget column; in the AUC column, it indicates the highest value among the displayed methods. Blocks denotes number of blocks

\times

tokens per block. The policy from Jazbec et al. (2026) corresponds to the expert-steering variant. Fast-dLLM’s unavailable 4-call accuracy is imputed as its 8-call accuracy.

Figure 1: Accuracy–compute tradeoff on mathematical reasoning. Micro-averaged test accuracy on GSM8K and MATH-500 (1,819 problems) versus total decoder-calls, shown on a logarithmic scale. We compare

\mathrm{TIPS}_{\mathrm{full}}

(

1\times 256

) and

\mathrm{TIPS}_{\mathrm{block}}

(

2\times 128

) against Fast-dLLM (

8\times 32

), WINO, UNCODE, and the learned unmasking policy of Jazbec et al. (2026) with expert-steering; the latter three use

1\times 256

full-canvas decoding. The displayed baselines include a highest or competitive accuracy baseline at every evaluated budget from the full comparison suite. TIPS curves show means with shaded bands indicating

\pm 1

standard error across three training-seed means, each averaged over three inference draws. Baselines are shown as point estimates. Other adaptive baselines may use fewer calls than the specified budget.

Table 4

Panel of Figure 2 — Figure 2: Generalization of ${\mathrm{TIPS}_{\mathrm{full}}}$ . Left: the LLaDA-trained full-canvas policy transfers to Dream-7B. Right: $\mathrm{TIPS}_{\mathrm{full}}$ is applied to different block sizes on LLaDA, with decoding budget shared across blocks.

Appendix figures & tables57 assets

Supplementary material from the paper’s appendix.

Appendix

Single-canvas configurations: 256 response tokens
Method / configuration	Blocks	4	8	16	32	64	128	256
Random / uniform	$1\!\times\!256$	11.45	16.91	25.16	33.06	38.34	41.82	42.57
High confidence (fitted)	$1\!\times\!256$	16.22	16.82	23.42	34.52	39.36	32.71	28.26
High confidence (linear)	$1\!\times\!256$	5.06	12.26	30.84	37.82	33.32	27.32	29.19
Probability margin	$1\!\times\!256$	9.84	17.59	31.61	39.53	35.51	27.54	30.40
Predictive entropy	$1\!\times\!256$	4.45	7.81	20.51	32.49	25.73	19.79	28.59

Appendix

Table 4: Complete accuracy–budget comparison on mathematical reasoning. Micro-averaged test accuracy (%) on GSM8K and MATH-500 (1,819 problems). All baseline configurations are retained, with blocks denoting number of blocks

\times

tokens per block. TIPS entries give the mean and, underneath, its standard error across three training-seed means; each seed mean averages three inference draws. Baselines are point estimates. Bold indicates the highest reported mean in each budget column among the 256-token configurations. Dashes indicate unreported cells. The non-adjusted Fast-dLLM shared-cap adaptation is retained for reference, and the updated mean-budget adaptation is added as a separate row. The final row uses only 128 response tokens and is displayed separately, without entering the matched comparison.

Configuration	$T=16$	$T=32$	$T=64$
Warmup $\rightarrow$ RLVR	43.16	56.15	58.20
Supervised warmup only	28.61	38.28	43.75
RLVR only	42.68	47.66	50.39
Remasking disabled (frozen)	36.13	44.53	45.21
Monotone policy (retrained)	42.38	47.75	51.27

Appendix

Table 5: Training stages and reversible visibility. Validation-set micro-accuracy (%) over 1,024 problems, with one training seed and one evaluation seed. RLVR continuations use the same number of updates.

Panel of Figure 3 — Figure 3: Supervised initialization and reversible visibility. Validation-set accuracy versus decoder-call budget on a logarithmic axis. Each curve is one trained policy evaluated with one inference seed.

Feature width $K$	$T=16$	$T=32$	$T=64$	$T=256$
$1$	40.12 $\pm$ 0.93	48.44 $\pm$ 1.03	51.38 $\pm$ 1.00	56.01 $\pm$ 0.96
$2$	43.68 $\pm$ 1.27	50.59 $\pm$ 0.48	54.30 $\pm$ 0.64	57.68 $\pm$ 0.48
$10$	38.90 $\pm$ 0.35	45.94 $\pm$ 0.60	49.97 $\pm$ 0.53	55.35 $\pm$ 1.08
$50$	40.13 $\pm$ 0.42	47.12 $\pm$ 1.12	49.97 $\pm$ 0.76	54.03 $\pm$ 1.10

Appendix

Table 6: Feature-width study. Test micro-accuracy (%), reported as mean

\pm

sample standard deviation across five evaluation seeds for each fixed policy.

Figure 4: Confidence-feature width. Test micro-accuracy versus log decoder-call budget for

K\in\{1,2,10,50\}

raw-probability features. Markers include intermediate budgets as well as powers of two. Shading is one standard error across five evaluation seeds.

Figure 5: Finalization deferral in RLVR-only policies. Realized visible-token fraction at

T=32

: mean

\pm

one standard deviation over training seeds

123

124

, and

125

, with evaluation seed

42

and example-weighted pooling of GSM8K test and MATH-500. The shaded final interval and dashed terminal segments denote deterministic finalization. The origin is the prescribed all-mask state.

Occupancy control	Masks before finalization	Masked (%)	Accuracy (%)
None	$146.20\pm 15.68$	$57.11$	$45.90$
Soft schedule loss	$122.91\pm 9.96$	$48.01$	$47.18$
Fixed power-law calibration	$17.65\pm 0.02$	$6.89$	$50.83$

Appendix

Table 7: Occupancy control reduces the finalization burden. At

L=256

and

T=32

, mask counts are means

\pm

one standard deviation across three training seeds, measured at

\tau=31/32

with evaluation seed

42

Frozen-policy setting	$T=16$	$T=32$	$T=64$
Fitted: $\gamma=0.14544$	43.16	56.15	58.20
$\gamma=-0.2$	30.86	47.95	54.49
$\gamma=0$ (linear)	39.55	50.49	55.66
$\gamma=0.2$	45.70	53.22	57.91
Calibration disabled	26.66	25.20	25.59

Appendix

Table 8: Inference-time schedule interventions. Validation-set accuracy (%) for frozen

\mathrm{TIPS}_{\mathrm{full}}

. All

\gamma

settings retain occupancy calibration; the final row bypasses the scalar logit shift entirely.

Panel of Figure 6 — Figure 6: Schedule shape and occupancy calibration have distinct effects. Left: validation-set accuracy at the four measured $\gamma$ values, with a separate curve for each budget. Right: the fitted schedule, calibrated linear schedule, and calibration-disabled policy versus log budget. All settings use the same frozen checkpoint.

Policy	Parameters	$T=16$	$T=32$	$T=64$	Policy time (ms)
Two-layer DiT	628,739	43.16	56.15	58.20	8.06
One-layer DiT	331,395	45.12	53.03	55.86	6.88
Per-token MLP	634,835	38.18	48.24	54.88	6.22

Appendix

Table 9: Architecture and controller cost. Validation-set accuracy (%) and total CUDA-measured policy time per generated example at

T=32

. Timings accumulate all 31 policy decisions; they were collected on an H100 80GB with float32 policy weights and batch size at most 16. The decoder and feature-extraction cost are excluded from the policy-time column.

Training configuration	$T=16$	$T=32$	$T=64$
$T_{\mathrm{train}}=32$	43.16	56.15	58.20
$T_{\mathrm{train}}=128$	39.16	49.32	53.22

Appendix

Table 10: Training-budget comparison for

\mathrm{TIPS}_{\mathrm{full}}

. Validation-set accuracy (%) at the indicated inference budgets.

Panel of Figure 7 — Figure 7: Controller architecture and training budget. Validation-set accuracy versus log decoder-call budget.

Method	Layers	Parameters	$4$	$8$	$16$	$32$	$64$	$128$
$\mathrm{TIPS}_{\mathrm{full}}$	1	331,395	17.38	25.98	45.12	53.03	55.86	59.47
$\mathrm{TIPS}_{\mathrm{full}}$	2	628,739	18.26	28.42	43.16	56.15	58.20	59.67
$\mathrm{TIPS}_{\mathrm{block}}$	1	331,395	12.70	20.90	40.43	52.64	58.79	61.33
$\mathrm{TIPS}_{\mathrm{block}}$	2	628,739	15.04	23.54	41.60	53.03	59.18	61.33
Jazbec et al. (2026)	1	331,265	7.32	14.16	32.32	43.85	43.75	43.95
Jazbec et al. (2026)	2	628,609	7.23	11.62	20.70	40.33	43.65	44.82

Appendix

Table 11: Matched policy-depth comparisons. Validation-set accuracy (%) at the requested decoder-call caps. Each TIPS decoding geometry is unchanged when depth is varied.

Method	Cap $16$	Cap $32$	Cap $64$	Cap $128$
TIPS (either geometry or depth)	16.00	32.00	64.00	128.00
Jazbec et al., one layer	16.00	29.52	35.42	35.92
Jazbec et al., two layers	16.00	31.95	50.30	51.38

Appendix

Table 12: Consumed decoder calls in the depth comparison. Mean calls per validation sample. TIPS uses every allocated call, whereas the comparison policy may stop before reaching the cap. Equal caps therefore do not always imply equal consumed compute.

Figure 8: Full-canvas comparison at matched policy depth. Validation-set accuracy for one- and two-layer TIPS and the learned unmasking policy of Jazbec et al. (2026) . The horizontal axis is the requested decoder-call cap; consumed calls for adaptive policies are reported in Table 12 .

Figure 9: Policy depth under two-block decoding. Both curves use the same

2\times 128

geometry and exact global decoder-call budget. The one-layer variant approaches the two-layer reference at larger budgets and matches it at

T=128

in this single-seed comparison.

Block	Phase	Native	Random reveal	Random remask	Random both
1	0.25	44.42	41.29	43.65	40.52
1	0.50	44.42	43.32	42.33	42.33
1	0.75	44.42	44.91	44.58	43.87
2	0.25	44.42	41.89	43.05	41.34
2	0.50	44.42	40.79	43.76	40.79
2	0.75	44.42	42.99	44.53	43.21

Appendix

Table 13: Position-selection case study:

\mathrm{TIPS}_{\mathrm{block}}

T=16

. Test micro-accuracy (%) for a count-preserving intervention at one decision. All six sites for this selected setting are included. Block indices are one-based.

Figure 10: Count-preserving position randomization for

\mathrm{TIPS}_{\mathrm{block}}

T=16

. Positive values mean that native position selection yields higher final accuracy. All six intervention sites are shown, with one-based block indices. Reveal counts and remask counts are separately preserved at the intervened decision.

Policy / active block	Accepted references	Eligible	Too few tokens
$\mathrm{TIPS}_{\mathrm{full}}$	1047/1819	1047	0
$\mathrm{TIPS}_{\mathrm{block}}$ , block 1	1095/1819	1095	0
$\mathrm{TIPS}_{\mathrm{block}}$ , block 2	1095/1819	948	147

Appendix

Table 14: Reference selection for controlled corruption. One

T=32

reference attempt is made per test problem for each policy. We retain verifier-accepted responses with at least eight eligible visible tokens at the intervention site. The same eligible pool is used for all

k

within a row.

Policy / active block	$k$	Correct / eligible	Answer (%)	Restored (%)	$\Delta$ (pp)
$\mathrm{TIPS}_{\mathrm{full}}$	0	1038/1047	99.14	—	+0.00
$\mathrm{TIPS}_{\mathrm{full}}$	1	1030/1047	98.38	42.50	-0.76
$\mathrm{TIPS}_{\mathrm{full}}$	2	1022/1047	97.61	42.50	-1.53
$\mathrm{TIPS}_{\mathrm{full}}$	4	1007/1047	96.18	43.79	-2.96
$\mathrm{TIPS}_{\mathrm{full}}$	8	979/1047	93.51	46.41	-5.64
$\mathrm{TIPS}_{\mathrm{block}}$ , block 1	0	1076/1095	98.26	—	+0.00

Appendix

Table 15: Controlled corruption and recovery. Answer correctness is measured over eligible problems; exact restoration is measured over the

k

substituted tokens per problem.

k=0

denotes the reconstructed clean continuation.

\Delta

compares answer accuracy with that matched clean control.

Panel of Figure 11 — Figure 11: Recovery from controlled corruption. The horizontal axis is the number of substituted tokens. Each curve uses its own accepted-reference population and remaining decoding horizon, as detailed in the text. Panel (a) weights corrupted tokens; panel (b) weights eligible problems.

Policy / active block	$k$	Restored	Learned reveal	Finalization	Transient only
$\mathrm{TIPS}_{\mathrm{full}}$	1	445	434	11	0
$\mathrm{TIPS}_{\mathrm{full}}$	2	890	857	33	0
$\mathrm{TIPS}_{\mathrm{full}}$	4	1834	1767	67	0
$\mathrm{TIPS}_{\mathrm{full}}$	8	3887	3747	140	4
$\mathrm{TIPS}_{\mathrm{block}}$ , block 1	1	838	736	102	0
$\mathrm{TIPS}_{\mathrm{block}}$ , block 1	2	1651	1442	209	2

Appendix

Table 16: Where restoration occurs. The learned-reveal and finalization columns partition the finally restored tokens by their last restoration event. The last column counts tokens restored temporarily but not present in their reference form at termination.

Policy / active block	$k$	Corrupt (%)	Clean (%)	$\Delta_{\rm token}$	$\Delta_{\rm example}$
$\mathrm{TIPS}_{\mathrm{full}}$	1	0.15	0.06	+0.08	+0.08
$\mathrm{TIPS}_{\mathrm{full}}$	2	0.19	0.06	+0.13	+0.13
$\mathrm{TIPS}_{\mathrm{full}}$	4	0.30	0.06	+0.24	+0.24
$\mathrm{TIPS}_{\mathrm{full}}$	8	0.54	0.07	+0.47	+0.49
$\mathrm{TIPS}_{\mathrm{block}}$ , block 1	1	0.30	0.14	+0.16	+0.16
$\mathrm{TIPS}_{\mathrm{block}}$ , block 1	2	0.37	0.14	+0.23	+0.23

Appendix

Table 17: Drift at untouched positions. We compare identical initially visible, nonspecial positions in each corrupted/clean pair, excluding substitutions. The first two rates and

\Delta_{\rm token}

weight token positions;

\Delta_{\rm example}

averages per-problem drift fractions. Differences are in percentage points. A token mismatch does not necessarily change the final answer.

Figure 12: Excess drift at initially untouched positions. Token-weighted disagreement with the reference after a corrupted continuation, minus disagreement at the identical untouched positions in its matched clean continuation. Block numbers are one-based.

Policy	$T$	Remasks	Visible (%)	All positions (%)	Steps (%)	Trajectories (%)
$\mathrm{TIPS}_{\mathrm{full}}$	16	503,803	18.04	7.21	93.29	100.00
$\mathrm{TIPS}_{\mathrm{full}}$	32	1,119,399	17.97	7.75	96.57	100.00
$\mathrm{TIPS}_{\mathrm{full}}$	64	2,457,136	18.71	8.38	97.97	100.00
$\mathrm{TIPS}_{\mathrm{block}}$	16	324,263	30.50	9.95	85.31	100.00
$\mathrm{TIPS}_{\mathrm{block}}$	32	929,911	34.23	13.31	92.91	100.00
$\mathrm{TIPS}_{\mathrm{block}}$	64	2,315,040	38.11	16.04	96.02	100.00

Appendix

Table 18: How often TIPS remasks. All 1,819 test problems are included in each row. The four rates use, respectively, visible position-opportunities, all active position-opportunities, learned steps containing a remask, and trajectories containing a remask. A position-opportunity is one active position at one learned decision. Deterministic finalization is excluded.

Policy	$T$	Changed	Changed (%)	Same (%)	Delay	Finalized	Special (%)
$\mathrm{TIPS}_{\mathrm{full}}$	16	42,714	8.48	91.52	2.245	29,553	45.36
$\mathrm{TIPS}_{\mathrm{full}}$	32	72,426	6.47	93.53	2.699	16,413	36.39
$\mathrm{TIPS}_{\mathrm{full}}$	64	118,743	4.83	95.17	3.126	8,332	31.96
$\mathrm{TIPS}_{\mathrm{block}}$	16	22,988	7.09	92.91	1.198	48,771	25.16
$\mathrm{TIPS}_{\mathrm{block}}$	32	31,264	3.36	96.64	1.294	31,638	18.69
$\mathrm{TIPS}_{\mathrm{block}}$	64	42,633	1.84	98.16	1.412	17,001	14.76

Appendix

Table 19: What happens after remasking. An episode begins when a visible token is remasked and ends at its next reveal. Changed and same-token percentages use completed episodes; every recorded episode completes. Delay is the mean number of global forward-call intervals until re-reveal. Finalized counts episodes ending at deterministic finalization. Special is the fraction of remask events involving tokenizer-defined special tokens, without opportunity adjustment.

Policy	$T$	$q=0$	$q=1/4$	$q=1/2$	$q=3/4$	$q=1$
$\mathrm{TIPS}_{\mathrm{full}}$	16	46.23	24.02	16.99	17.26	17.70
$\mathrm{TIPS}_{\mathrm{full}}$	32	57.56	31.45	20.56	19.79	19.74
$\mathrm{TIPS}_{\mathrm{full}}$	64	63.39	40.85	31.34	25.73	23.91
$\mathrm{TIPS}_{\mathrm{block}}$	16	44.42	17.98	20.45	20.62	15.94
$\mathrm{TIPS}_{\mathrm{block}}$	32	60.20	24.85	22.76	26.06	20.67
$\mathrm{TIPS}_{\mathrm{block}}$	64	64.54	38.70	38.15	40.19	34.91

Appendix

Table 20: Complete heuristic-to-TIPS handoff results. Test micro-accuracy (%) at all global switching fractions.

q=0

is pure TIPS;

q=1

is the matched fixed-horizon non-adjusted confidence-threshold heuristic. Every row uses exactly

T

real decoder calls, including no-op and finalization calls.

Panel of Figure 13 — Figure 13: Heuristic prefixes followed by TIPS. Test micro-accuracy at every evaluated switching fraction, including pure TIPS ( $q=0$ ) and the matched heuristic ( $q=1$ ). Each curve has a fixed total decoder-call budget.

Policy	$T$	$q$	Wrong $\to$ correct	Correct $\to$ wrong	$\Delta$ (pp)
$\mathrm{TIPS}_{\mathrm{full}}$	16	0.25	267	152	+6.32
$\mathrm{TIPS}_{\mathrm{full}}$	16	0.50	91	104	-0.71
$\mathrm{TIPS}_{\mathrm{full}}$	16	0.75	42	50	-0.44
$\mathrm{TIPS}_{\mathrm{full}}$	32	0.25	394	181	+11.71
$\mathrm{TIPS}_{\mathrm{full}}$	32	0.50	137	122	+0.82
$\mathrm{TIPS}_{\mathrm{full}}$	32	0.75	50	49	+0.05

Appendix

Table 21: Paired final-answer changes after handoff. Counts compare each hybrid with the matched pure-heuristic endpoint on the same 1,819 problems. Correctness refers to terminal answers. Both beneficial and adverse transitions are retained.

Dataset	$N$	Non-adjusted Fast-dLLM	$+\mathrm{TIPS}_{\mathrm{full}}$	$\Delta$ (pp)	W $\to$ C	C $\to$ W
		Accuracy (%)			Answer transitions
GSM8K	1,319	79.30	79.61	+0.30	4	0
MATH-500	500	34.60	34.80	+0.20	1	0
Micro-average	1,819	67.01	67.29	+0.27	5	0

Appendix

Table 22: Four-call repair of completed non-adjusted Fast-dLLM responses. Test accuracy before and after applying

\mathrm{TIPS}_{\mathrm{full}}

to the same saved responses, with training and evaluation seed

123

. W

\to

C and C

\to

W count wrong-to-correct and correct-to-wrong final-answer transitions, respectively. Every response is repaired, regardless of its original correctness. Non-adjusted Fast-dLLM has a

256

-call global cap but terminates naturally on every problem. Differences are computed from unrounded accuracies.

Figure 14: Accuracy changes after four additional repair calls. Points show the paired accuracy change from completed non-adjusted Fast-dLLM responses to their

\mathrm{TIPS}_{\mathrm{full}}

revisions. Whiskers are

95\%

paired bootstrap intervals from

10{,}000

resamples, sampling problems within each benchmark at its original size. The dashed line marks no accuracy change.

Dataset	Method	Accuracy (%)			Mean decoder calls
		$16$	$32$	$64$	$16$	$32$	$64$
GSM8K	$\mathrm{TIPS}_{\mathrm{full}}$	41.77	55.72	63.84	16.00	32.00	64.00
GSM8K	Fast-dLLM	25.32	30.78	34.57	15.48	32.75	66.33
MATH-500	$\mathrm{TIPS}_{\mathrm{full}}$	17.80	22.40	22.20	16.00	32.00	64.00
MATH-500	Fast-dLLM	16.40	18.60	17.80	17.50	30.70	58.21
Micro-average	$\mathrm{TIPS}_{\mathrm{full}}$	35.18	46.56	52.39	16.00	32.00	64.00

Appendix

Table 23: Dream transfer by dataset and consumed compute. Columns correspond to requested budgets 16, 32, and 64. Both methods decode a single 256-token canvas. The pooled row weights individual problems equally.

Figure 15: Zero-shot backbone transfer on a matched full canvas. Test micro-accuracy for the frozen LLaDA-trained

\mathrm{TIPS}_{\mathrm{full}}

policy and Fast-dLLM on Dream-7B. Both use

1\times 256

decoding. The horizontal axis is the decoder-call budget; actual baseline consumption is reported in Table 23 .

Method	Geometry	HumanEval (164)			MBPP (500)
		$T=16$	$T=32$	$T=64$	$T=16$	$T=32$	$T=64$
UNCODE / PC-Sampler	$1\times 256$	4.27	6.10	16.46	0.00	2.80	15.20
WINO	$1\times 256$	15.24	21.95	24.39	16.20	26.80	32.40
WINO	$2\times 128$	12.20	23.78	30.49	1.80	11.40	29.80
Non-adjusted Fast-dLLM	$1\times 256$	15.24	21.95	23.78	16.80	23.40	26.20
Non-adjusted Fast-dLLM	$8\times 32$	14.63	32.93	40.24	1.20	13.40	38.20

Appendix

Table 24: Zero-shot coding transfer. TIPS uses frozen math-trained policies. Pass@1 (%), one generation per test task, against each of the five retained baselines individually. Baseline presets/assets are inherited from mathematical reasoning. The

8\!\times\!32

baseline uses a shared global cap. Adaptive methods may consume fewer than

T

calls.

Figure 16: Zero-shot coding transfer against individual baselines. One-generation pass@1 versus log decoder-call cap. TIPS uses its frozen math-trained policies.

Method	Geometry	HumanEval (164)			MBPP (500)
		$T=16$	$T=32$	$T=64$	$T=16$	$T=32$	$T=64$
UNCODE / PC-Sampler	$1\times 256$	4.27	6.10	16.46	0.00	2.80	15.20
WINO	$1\times 256$	15.24	21.95	24.39	16.20	26.80	32.40
WINO	$2\times 128$	12.20	23.78	30.49	1.80	11.40	29.80
Non-adjusted Fast-dLLM	$1\times 256$	15.24	21.95	23.78	16.80	23.40	26.20
Non-adjusted Fast-dLLM	$8\times 32$	14.63	32.93	40.24	1.20	13.40	38.20

Appendix

Table 25: Coding-trained TIPS. TIPS policies are trained separately on coding; the decoder remains frozen. Pass@1 (%), one generation per test task, against each of the five retained baselines individually. Baseline presets/assets and TIPS settings are inherited from mathematical reasoning and untuned. The

8\!\times\!32

baseline uses a shared global cap. Adaptive methods may consume fewer than

T

calls.

Figure 17: Coding-trained TIPS against individual baselines. The decoder is unchanged; only the TIPS policies are trained for coding. Baselines are identical to those in Figure 16 . Curves show one evaluation sample.

Method	Geometry	HumanEval (164)			MBPP (500)
		$T=16$	$T=32$	$T=64$	$T=16$	$T=32$	$T=64$
UNCODE / PC-Sampler	$1\times 256$	16.00	32.00	64.00	16.00	32.00	64.00
WINO	$1\times 256$	15.84	25.26	31.24	15.94	25.20	33.71
WINO	$2\times 128$	16.00	24.84	31.19	16.00	28.73	41.09
Non-adjusted Fast-dLLM	$1\times 256$	13.21	18.83	21.01	14.95	23.45	28.65
Non-adjusted Fast-dLLM	$8\times 32$	15.93	28.51	37.82	16.00	31.76	53.07

Appendix

Table 26: Consumed compute in coding experiments. Mean decoder calls at the indicated caps, for the same conditions as Tables 24 and 25 . Both math-trained and coding-trained TIPS use the exact requested horizon.

Target geometry	Dataset	$T=16$	$T=32$	$T=64$
$1\times 256$	GSM8K	57.39	71.72	74.75
$1\times 256$	MATH-500	19.60	25.80	30.80
$1\times 256$	Micro-average	47.00	59.10	62.67
$2\times 128$	GSM8K	44.73	66.41	75.74
$2\times 128$	MATH-500	15.80	25.80	29.00
$2\times 128$	Micro-average	36.78	55.25	62.89

Appendix

Table 27: Frozen-policy block-size transfer. All rows use the full-canvas-trained checkpoint, without block-specific policy training. Values are test accuracy (%). The total 256-token response and total decoder-call budget remain fixed.

Figure 18: Frozen-policy transfer across decoding geometries. Every curve uses the same full-canvas-trained checkpoint and the same global 256-token response length. The decoder-call budget is shared across blocks.

Proposal temperature	$T=16$	$T=32$	$T=64$
0.0	47.00	59.10	62.67
0.1	48.32	58.11	63.33
0.5	46.89	56.02	62.62
1.0	37.77	51.35	60.75

Appendix

Table 28: Sensitivity to proposal temperature. Frozen

\mathrm{TIPS}_{\mathrm{full}}

on the pooled math test set. Only token-proposal temperature changes; finalization remains greedy. Values are single-seed micro-accuracy (%).

Figure 19: Sensitivity to proposal temperature. Test micro-accuracy for the same frozen

\mathrm{TIPS}_{\mathrm{full}}

policy. Only token proposals change temperature; finalization remains greedy, and visibility actions retain the same sampling rule.

Figure 20: Supervised warmup on mathematics. Recorded training soft binary cross-entropy in the initial and final warmup phases. Faint traces are raw observations; the darker curves are exponential moving averages with smoothing coefficient 0.1. The phases have separate optimizer-step axes.

Figure 21: Recomputed warmup validation curves. Recorded mathematical-reasoning training loss and validation observations recomputed on a fixed state bank, with exponential moving averages for visualization.

Figure 22: Supervised warmup on coding. Supplied training and validation soft binary cross-entropy traces for the initial and final warmup phases. The displayed averages summarize within-run loss observations.

Domain	Policy	First 50 updates	Last 50 updates
Math	$\mathrm{TIPS}_{\mathrm{full}}$	38.19	50.73
Math	$\mathrm{TIPS}_{\mathrm{block}}$	37.55	50.41
Coding	$\mathrm{TIPS}_{\mathrm{full}}$	16.28	37.55
Coding	$\mathrm{TIPS}_{\mathrm{block}}$	13.66	27.41

Appendix

Table 29: RLVR training diagnostics. Mean raw terminal-correctness reward (%) over the first and last 50 recorded updates of each run.

Panel of Figure 23 — Figure 23: RLVR training curves. Terminal correctness reward per nominal update. Faint lines show the raw observations; darker lines show exponential moving averages with smoothing coefficient 0.1.

Figure 24: Accepted full-canvas trajectory on GSM8K. Janet’s ducks lay 16 eggs daily; she uses three for breakfast and four for muffins, then sells the remainder for two dollars each. The final response correctly computes

(16-3-4)\times 2=18

dollars. The snapshots show progressive visibility, remasking, and finalization.

Figure 25: Accepted full-canvas trajectory on MATH-500. The task asks how many positive whole-number divisors 196 has. The response arrives at

196=2^{2}\times 7^{2}

and the accepted answer

(2+1)(2+1)=9

. Token-level reveal and remask annotations are retained in the displayed intermediate and final canvases.

Figure 26: Accepted blockwise trajectory on GSM8K. A robe requires two bolts of blue fiber and half as much white fiber, for a total of three bolts. The first block contains the solution and is finalized after 16 calls; the second block is then decoded under the same global budget. This example also illustrates that TIPS continues to the allocated horizon even when the answer-bearing text appears earlier.

Figure 27: Verifier-accepted blockwise answer with flawed intermediate reasoning. The task asks for the angle between the lines

2x=3y=-z

and

6x=-y=-4z

. The final answer of

90^{\circ}

is accepted, but the displayed derivation uses incorrect direction vectors and an invalid inference from their dot product. We retain the example to distinguish final-answer verification from verification of the reasoning trace.

Figure 28: Unsuccessful full-canvas revision on GSM8K. Josh buys a house for

80,000, spends

50,000 on repairs, and increases its value by 150%. Despite repeated visibility changes, the response ends with an incorrect calculation and a generated answer of $40,000. The trajectory demonstrates that remasking is an opportunity to revise, not a guarantee of correct reasoning or a correct final answer.

Setting	$\mathrm{TIPS}_{\mathrm{full}}$	$\mathrm{TIPS}_{\mathrm{block}}$
Response geometry	$1\times 256$	$2\times 128$
DiT layers / hidden width	$2/128$	$2/128$
Attention heads / head width	$2/64$	$2/64$
Feed-forward width	512	512
Confidence channels	2	2
Time-embedding width	128	128

Appendix

Table 30: Configuration of the final mathematical-reasoning policies. The architecture is shared; decoding geometry, local timing, and fitted schedule differ.

Stage	Updates	State presentations	Initialization
Initial supervised warmup	3,230	396,056	Fresh policy
Refreshed supervised warmup	4,074	495,070	Initial warmup endpoint
Full-canvas RLVR	935	—	Refreshed warmup endpoint
Two-block RLVR	935	—	Same warmup endpoint

Appendix

Table 31: Training stages per seed. State presentations count supervised examples processed, not distinct training problems.

Setting	Supervised warmup	RLVR
Optimizer	AdamW	AdamW
Peak learning rate	$3\times 10^{-5}$	$3\times 10^{-5}$
Adam moments	$(0.9,0.999)$	$(0.9,0.99)$
Adam numerical constant	$10^{-8}$	$10^{-8}$
Weight decay	0	0.1; 0 for bias/norm
Learning-rate warmup	100 updates	100 updates

Appendix

Table 32: Optimization settings.

Method	Configurations
Random / uniform	$1\times 256$
High confidence	$1\times 256$ , fitted schedule; $1\times 256$ , linear schedule
Probability margin	$1\times 256$
Predictive entropy	$1\times 256$
Fast-dLLM	$8\times 32$ , mean-adjusted as described in section 4.1
Non-adjusted Fast-dLLM	$1\times 256$ ; $8\times 32$

Appendix

Table 33: Baseline configurations. A geometry denotes number of blocks

\times

tokens per block. Separate entries within a row are distinct saved configurations. The

4\times 32

UPO configuration uses a 128-token response and is recorded only at

T=128

Population	GSM8K	MATH / MATH-500	Total
Gradient-training input	6,280	6,496	12,776
Validation and schedule fitting	512	512	1,024
Final mathematical-reasoning test	1,319	500	1,819

Appendix

Table 34: Problem populations for the main mathematical-reasoning training and evaluation. The validation population is also used, without correctness labels, to fit the schedule.

Diffusion Language Model Decoding

Diffusion language model decoding generates text through iterative parallel refinement, balancing fewer denoising steps with accuracy, coherence, and control.

See all 235 papers

Mar 9, 2026

DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention

Masked diffusion language models enable parallel token decoding, providing a promising alternative to the sequential nature of autoregressive generation. However, their iterative denoising process remains computationally expensive because it repeatedly processes the entire sequence at every step. We observe that across these diffusion steps, most token representations remain stable; only a small subset, which we term salient tokens, contributes meaningfully to the next update. Leveraging this temporal sparsity, we present DyLLM, a training-free inference framework that accelerates decoding by selectively computing only these salient tokens. DyLLM identifies saliency by measuring the cosine similarity of attention contexts between adjacent denoising steps. It recomputes feed-forward and attention operations only for salient tokens while reusing cached activations for the remainder. Across diverse reasoning and code-generation benchmarks, DyLLM achieves up to 9.6x higher throughput while largely preserving the baseline accuracy of representative open-source diffusion LLMs, LLaDA, and Dream.

Younjoo Lee, Seungkyun Dan, Junghoo Lee +2

Seoul National University, Seoul, Republic of Korea

Diffusion Language Model Inference Diffusion Language Model Decoding Diffusion Model Inference Acceleration Token Pruning

Ask PDF

May 22, 2026

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Diffusion language models intrinsically fail to capture correlations between decoded tokens, which leads to a harsh trade-off between sampling quality and throughput. To solve this issue, we propose DiLaDiff, a variant of masked diffusion language models with three components: (1) a continuous latent space with semantic capabilities, learned by an auto-encoder fine-tuned from an existing masked diffusion language model; (2) a latent diffusion model learning the prior over the encoder distribution; (3) a consistency model distilling the learned prior into a few-step latent generative model. We show that, even without distillation, our latent-guided diffusion model outperforms the masked diffusion baseline while significantly accelerating inference. Consistency distillation further lowers the computational overhead of continuous diffusion, such that the latent is generated in negligible time compared to discrete decoding.

Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3

NVIDIA

Diffusion Language Models Diffusion Language Model Distillation Diffusion Language Model Inference Masked Diffusion Models Continuous Diffusion Language Models Diffusion Model Inference Acceleration+1

Ask PDF

Jul 16, 2026

Mask-Aware Policy Gradients for Diffusion Language Models

Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term. Combining optimization of both terms leads to state-of-the-art outcomes on mathematical reasoning and coding benchmarks, with scores of 87.1% on GSM8K and 53.4% on MBPP.

Haran Raajesh, Kulin Shah, Adam Klivans +1

Department of Computer Science The University of Texas at Austin

Masked Diffusion Models RL for Language Model Reasoning Policy Gradient Methods RL for Diffusion Models

Ask PDF

Abstract

Diffusion Language Model Decoding

Explore similar work

DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Mask-Aware Policy Gradients for Diffusion Language Models