Outcome-Routed Distillation and Reinforcement Learning for Diffusion Language Models

Wentao Lu; Tianyu Zhu; Jesse Clark

cs.LGOct 13, 2026

Outcome-Routed Distillation and Reinforcement Learning for Diffusion Language Models

Authors: Wentao Lu, Tianyu Zhu, Jesse Clark

Organizations: Celeris San Francisco, CA

Abstract

On-policy distillation provides dense supervision on the student's own samples, but its targets still come from the teacher. Verifier reinforcement learning (RL) instead uses an outcome signal, but it provides no direct constraint on capabilities outside the verifier. Combining the two is difficult for diffusion language models (dLLMs), because bidirectional denoising makes the policy likelihoods required by ratio-based RL unreliable. We develop a likelihood-free pipeline for a mixture-of-experts dLLM. Each rollout group is routed by verifier outcome. Mixed groups use ratio-free, group-centred REINFORCE. If every student rollout fails but a teacher succeeds, the model receives dense KL supervision on the student's denoising canvases. All other groups are skipped. A KL term anchored to the distilled checkpoint limits changes outside the verifier reward. The tool-calling teacher succeeds on only 18.6% of the training prompts. At 24 denoising steps, distillation exceeds the teacher by 6.77 percentage points, and routing improves the distilled checkpoint by a further 18.31 points. Routing lowers reasoning-trace accuracy by 0.71 points but raises direct-answer accuracy on the same items by 13.21 points. The anchored arm restores trace accuracy. Its tool-calling score differs by -0.16 points, with a confidence interval that includes zero, although this comparison also removes a small fallback distillation route. A domain-conditional recipe retains a 16.37-point tool-calling gain; a second seed gives 18.23 points and reproduces both preservation results, but the math result is seed-dependent. We pre-registered every arm and evaluated each comparison on paired items in a single sweep on one machine.

Figures & tables

Figure 1: Likelihood-free training pipeline. Mixed groups use ratio-free, group-centred REINFORCE. All-fail groups use dense canvas KL only when a teacher succeeds. All-pass groups and all-fail groups without a successful teacher are skipped. The KL anchor to m-dist is scaled each step to a measured

0.1

gradient-norm ratio using separate backward passes.

arm	objective	change from previous
m-dist	dense multi-teacher KL on the student’s denoising canvases, with one same-architecture teacher per domain	baseline
m-route	route by verifier outcome: mixed groups use ratio-free REINFORCE; all-fail groups use dense KL when a teacher succeeds; all other groups are skipped	$+$ routing
m-anchor	add KL to m-dist at a measured $0.1$ gradient-norm ratio; math and tool-calling teachers are not loaded, so the fallback route is unavailable	$+$ anchor, $-$ fallback route
m-comp	math uses dense KL to m-dist as its primary objective; tool-calling and workflow follow m-anchor	domain-conditional objective

Table 1: Training arms. Each row changes one component relative to the row above.

teacher	method	training data	steps	benchmark gain vs base
math	verifier RL: per-position GDSD variant [ 14 ] , frozen experts, init from base	$\sim$ 200–320 competition-math prompts consumed of 5,833 sampled (GSM8K [ 3 ] +MATH [ 7 ] ), $G{=}16$ , mixed-outcome groups only	25–40	$+6.9$ @24 ( ci-sep at 8/24/48)
workflow	SFT: frozen-experts full fine-tune, seq-2048, thought-channel format	30.5M tokens; native long-CoT + pass@ $k$ -harvested completions	199	$+3.8$ macro (classification $+12$ )
tool-calling	verifier RL: the same GDSD variant, frozen experts, init from base; strict+graded AST-multiset verifier	350 single-turn prompts $\to$ 232 mixed groups ( $G{=}8$ , mixed denoising steps $\{8,24,48\}$ ), each visited once	29	$+8.6$ @24 ( ci-sep at 8/24/48)

Table 2: Teacher provenance. Each teacher starts from base and is post-trained in an earlier single-domain experiment using the best method for that domain. Mixed groups satisfy

0<\hat{r}<1

. Each RL teacher consumes every mixed group once. Benchmark gains and inheritance ratios use the same distillation-era scoring sweep; Table 3 reports a later sweep. No table combines scores from different sweeps.

domain	steps	base	m-dist	m-route	m-anchor	m-comp
math	8	57.89	60.20	62.34	62.34	58.72
	24	65.62	72.53	69.74	70.07	70.39
	48	66.28	73.36	70.72	69.90	71.71
workflow	8	43.98	49.88	50.60	51.12	51.48
	24	45.47	50.33	51.09	51.60	51.60
	48	45.28	49.61	50.97	50.79	51.02

Table 3: Scores on the three training domains, measured in one paired, single-machine sweep. Math is pass@1 accuracy; workflow is a continuous macro score scaled by

100

; tool-calling is AST-match accuracy. Bold marks m-comp . For seed 1, none of its paired accuracy cells is CI-separated below m-dist , and all workflow point estimates are at least as high.

cell (difference from m-dist )	ratio $0$	ratio $0.1$
tool-calling @24	$+18.31$ ci-sep	$+18.15$ ci-sep (difference from ratio 0: $-0.16$ , n.s. )
math @8	$+2.14$ n.s.	$+2.14$ n.s.
math @24	$-2.80$ ci-sep	$-2.47$ n.s. (difference: $+0.33$ )
math @48	$-2.63$ ci-sep	$-3.45$ ci-sep (difference: $-0.82$ )
preserve CoT vs. base	$-0.71$ ci-sep	$+0.27$ n.s. (difference: $+0.98$ , ci-sep )
multi-turn set vs. base	$+7.00$ ci-sep	$+9.50$ ci-sep

Table 4: Key differences from m-dist for m-route (anchor ratio

0

) and m-anchor (ratio

0.1

). Deltas use full-precision paired scores and can differ in the last digit from the rounded values in Table 3 . m-anchor also lacks the fallback distillation route available to m-route , so the between-arm differences do not isolate the anchor.

Figure 2: (a) Scores at 24 denoising steps. (b, c) Scores across step counts. Tool-calling gains persist for every routed arm. Math distinguishes the methods: m-anchor narrows the deficit at 24 steps and widens it at 48, while m-comp remains close to m-dist at 24 and 48 steps but gives up the routed arms’

2.14

-point advantage at 8 steps, ending

1.48

points below m-dist ( n.s. ).

cell	seed 1	seed 2	gap
tool-calling @24 vs m-dist	$+16.37$ ci-sep	$+18.23$ ci-sep	$1.86$
math @8 vs m-dist	$-1.48$ n.s.	$-3.95$ ci-sep below	$2.47$
math @24 vs m-dist	$-2.14$ n.s.	$+0.00$ n.s.	$2.14$
math @48 vs m-dist	$-1.64$ n.s.	$-1.48$ n.s.	$0.16$
preserve CoT vs base	$+0.48$ n.s.	$+0.29$ n.s.	$0.19$
preserve answer vs base	$+12.47$ ci-sep	$+12.33$ ci-sep	$0.14$

Table 5: Replication of m-comp . The prompt slices and recipe are fixed, and only the run seed changes. All arms are re-scored in one sweep, which first reproduces every prior cell exactly. Tool-calling and preservation replicate; math does not.

Figure 3: Capability scores on a common 0–100 scale. (a) Relative to base and m-dist , routed RL improves tool-calling while reducing math and chain-of-thought accuracy. (b) The anchored and composed arms retain most of the tool-calling gain while recovering part of the math and chain-of-thought losses. Each panel contains three polygons to keep every pair distinguishable under simulated colour-vision deficiency. Table 3 and Figure 2 show all five arms.

Diffusion Language Model Distillation

Diffusion language model distillation transfers knowledge or decoding behavior to make diffusion-based text generation faster and more capable.

See all 37 papers

Trace-Based On-Policy Distillation for Masked Diffusion Language Models
Jul 18, 2026Haolin Ren, Ziyang Huang, Chenhao Yuan +2Masked Diffusion Models Diffusion Model Distillation
Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China
GDSD: Reinforcement Learning as Guided Denoiser Self-Distillation for Diffusion Language Models
May 28, 2026Xiaohang Tang, Keyue Jiang, Che Liu +4Diffusion Model Distillation RL for Diffusion Models
UNIST · University of Basel
Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages
Mar 13, 2026Vishnu Teja Kunde, Fatemeh Doudi, Mahdi Farahbakhsh +3RL for Diffusion Models Diffusion Language Model Training
Department of Electrical & Computer Engineering, Texas A&M University, College Station, TX 77843

Jul 18, 2026J/K move · Enter open · S save

Trace-Based On-Policy Distillation for Masked Diffusion Language Models

Haolin Ren, Ziyang Huang, Chenhao Yuan, Jun Zhao+1

Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China

Masked Diffusion Models Diffusion Model Distillation Language Model Distillation Diffusion Language Models LLM Mathematical Reasoning Diffusion Language Model Distillation

Open paper AskPDF arXiv

Diffusion large language models (dLLMs) are a promising alternative to autoregressive generation. However, reasoning-oriented post-training for dLLMs remains challenging. Supervised fine-tuning (SFT) for dLLMs requires dense but often off-policy masked states, while reinforcement learning (RL) relies on sparse rewards or value modeling. This paper proposes trace-based on-policy distillation (TOPD), a teacher-supervised framework that transfers reasoning ability to a target dLLM without reward estimation. The key idea is to supervise a dLLM on its own denoising trajectory, focusing on the trace-aligned token decisions that form the final response. Specifically, TOPD samples on-policy diffusion trajectories from the target dLLM, obtains teacher token distributions from a teacher model on the corresponding partially denoised states, and updates the target dLLM with a token-level Reverse Kullback-Leibler (Reverse-KL) objective. This design preserves dense teacher supervision while aligning training with the model's own denoising states. On mathematical reasoning benchmarks, TOPD enables SDAR-4B-Chat to match the MATH500 accuracy of its RL-trained counterpart TraDo-4B-Instruct, with gains of +5.7 under static evaluation and +4.5 under dynamic evaluation. Compared with the RL-trained counterpart, TOPD achieves this with 4

\times

fewer rollout rounds, corresponding to an estimated 96.0

\times

to-accuracy model-compute speedup.

Abstract

Diffusion Language Model Distillation

Explore similar work

Trace-Based On-Policy Distillation for Masked Diffusion Language Models

GDSD: Reinforcement Learning as Guided Denoiser Self-Distillation for Diffusion Language Models

Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages