CT-OPD: Counterfactual Trace On-Policy Distillation for Diffusion Vision-Language Models
Authors: Long Qian, Bingke Zhu, Jiaqi Wei, Yu Li, Yingying Chen, Jinqiao Wang
Organizations: Foundation Model Research Center, Institute of Automation, Chinese Academy of Sciences · School of Future Technology, University of Chinese Academy of Sciences · School of Engineering, Cardiff University · School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences · Wuhan AI Research
Diffusion vision-language models generate answers by gradually resolving masked tokens, making accurate conditional prediction in partially resolved states central to post-training. Masking completed answers yields coherent contexts and targets, but prescribed masks do not reflect the model's reveal decisions. Its trajectories capture these decisions, yet their provisional visible tokens can conflict with the target response. Outcome-based reinforcement learning follows these trajectories but provides only response-level feedback, which loses contrast when sampled rewards tie. To align coherent token-level supervision with the model's reveal decisions, we introduce Counterfactual Trace On-Policy Distillation (CT-OPD), which combines completed teacher responses with trajectory masks from the current student. CT-OPD retokenizes each teacher response in the student's vocabulary and extracts unresolved-position masks at successive stages of the student's reverse process. For each mask, it discards provisional rollout values and reconstructs the partial state from the teacher endpoint, so the supervised positions follow the current trajectory while the visible context and targets remain consistent with the same response. The student is trained on these reconstructed states with its native categorical loss, and trajectories are refreshed as the model evolves. Across dense and sparse diffusion architectures, CT-OPD consistently enhances multimodal understanding and reasoning capabilities, with gains of up to 9.80 points on the nine-benchmark average. On the unified understanding-and-generation architecture, it also improves both visual understanding and image generation, showing that the same principle transfers across architectures and modalities. Ablations further attribute these gains to coherent reconstruction and current-model trajectory masks.
Figures & tables
Figure 1: Motivation of CT-OPD. Completed responses provide coherent content but do not specify the student’s unresolved positions. Student trajectories reveal these positions, yet their provisional tokens may conflict with endpoint targets. CT-OPD combines trajectory-selected positions with context and targets from one endpoint.
Figure 2: Overview of CT-OPD. A completed teacher response is retokenized as the student endpoint, while a current-student rollout supplies trajectory masks at four denoising stages. CT-OPD overlays each mask on the endpoint to form coherent partial states and supervises unresolved endpoint tokens with the native categorical loss. The updated student generates a fresh trajectory for the next cycle.
Method
Perception and knowledge
Scientific reasoning
Avg.
MMMU
MMMU-Pro
MMStar
MMBench
SEEDBench
MathVerse
MathVista
GPQA
GPQA-D
Dense diffusion VLMs
LLaDA-V
41.78
19.02
50.07
69.45
69.70
22.72
53.60
16.96
16.67
40.00
+ MaskGRPO
42.33 (+0.55)
25.66 (+6.64)
52.20 (+2.13)
72.74 (+3.29)
69.79 (+0.09)
14.72 ( − 8.00)
44.40 ( − 9.20)
28.57 (+11.61)
29.29 (+12.62)
42.19 (+2.19)
+ MDPO
41.56 ( − 0.22)
19.25 (+0.23)
49.27 ( − 0.80)
71.83 (+2.38)
69.87 (+0.17)
13.45 ( − 9.27)
40.70 ( − 12.90)
22.77 (+5.81)
24.24 (+7.57)
39.22 ( − 0.78)
+ Endpoint-only
44.33 (+2.55)
23.87 (+4.85)
52.93 (+2.86)
75.87 (+6.42)
70.37 (+0.67)
19.54 ( − 3.18)
48.90 ( − 4.70)
25.45 (+8.49)
26.77 (+10.10)
43.12 (+3.12)
Table 1: CT-OPD improves dense, sparse, and unified diffusion-VLMs. Parentheses give changes from Base (red: gain; blue: decrease); bold marks the best result per backbone. Endpoint-only uses full masks; Random matches mask counts and budget; OPDLM+BPM aligns teacher distributions.
Terminal contrast and zero-contrast CT coverage
Frozen terminal groups
Zero contrast: 60.74%
With contrast: 39.26%
CT coverage (zero contrast)
Prompt: 100.00% State: 91.48% Token: 60.79%
Common-state learning on 292 all-wrong groups
876 states; 35,935 scored tokens
Method
NLL ↓
Acc. ↑
Base (frozen)
4.6073
( − 1.2412)
20.56
(+6.87)
95% CI [ − 1.3331, − 1.1545]
95% CI [6.11, 7.66]
Table 2: Learning under zero terminal reward contrast. Coverage spans four stages on 311 zero-contrast prompts; prediction uses their 292 all-wrong prompts. Parentheses give CT-OPD differences with paired 95% CIs below; NLL is in nats/token and accuracy in percentage points.
Mask setting
MMMU
MMMU-Pro
MMStar
MMBench
SEEDBench
MathVerse
MathVista
GPQA
GPQA-D
Avg.
Online matched ( CT-OPD )
46.22
33.12
55.87
83.16
74.26
27.79
61.30
32.59
33.84
49.79
Online shuffled
44.00
31.97
53.93
79.69
72.11
25.89
58.90
31.03
29.29
47.42
Frozen-Base matched
45.33
32.60
55.27
82.82
73.57
26.02
59.20
31.70
30.81
48.59
Frozen-Base shuffled
42.67
31.79
53.40
76.90
71.24
24.75
58.40
29.91
29.80
46.54
Random
42.00
28.55
56.00
74.91
69.26
23.10
56.60
28.35
33.84
45.85
Table 3: Online, prompt-matched trajectory masks give the strongest transfer. The controls vary mask source and prompt pairing; Random preserves mask counts while randomizing positions. The lower block measures positions retained by the trajectory on the same reconstructed parent.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Setting
LLaDA-V
LaViDa-LLaDA
DiMoE-VL
Training data
Vision-SR1-47K
Vision-SR1-47K
Vision-SR1-47K
Trajectory-mask ratios
1,.75,.5,.25
1,.75,.5,.25
1,.75,.5,.25
Endpoint-only mask slots
1,1,1,1
1,1,1,1
1,1,1,1
Global batch
16
16
16
Learning rate
3×10−7
3×10−7
3×10−7
Response canvas
128
128
128
Appendix
Table 4: Shared CT-OPD configuration across diffusion-VLM backbones.
Component
Total
Active
Language model
7.357B
1.720B
Projector and image newline
0.007B
0.007B
Routing auxiliaries
0.000173B
0.000173B
SigLIP2 visual tower
1.107B
1.107B
DiMoE-VL
8.471B
2.834B
Appendix
Table 5: DiMoE-VL parameter accounting. Active parameters include always-active modules, the runtime visual tower, and the top-8 experts in each routed layer.
Configuration
Stage I
Stage II
Stage III
Training objective
Vision-language alignment
Full multimodal SFT
Counterfactual trace on-policy distillation
Optimization target
Multimodal projector
Projector and LLaDA-MoE decoder
Endpoint-token likelihood on masks from student traces
Trainable modules
Projector only
Projector and LLaDA-MoE decoder
LLaDA-MoE decoder, sparse routers/experts, and LM head
Training data
LLaVA-Pretrain
MAmmoTH-VL SI-10M + OV-2M
Vision-SR1-47K endpoints
Endpoint teacher
-
-
Qwen3-VL-32B-Instruct
Vision encoder
SigLIP2-Giant-Patch16-384 (frozen)
Appendix
Table 6: Three-stage training recipe for DiMoE-VL . Stage II produces DiMoE-VL-Base ; Stage III applies CT-OPD to the resulting native masked-diffusion policy. A dash marks a component unused in that stage.
Setting
MDPO
MaskGRPO
UniGRPO
Optimizer
AdamW
AdamW
AdamW
Learning rate
3×10−7
3×10−7
3×10−7
Weight decay
0
0
0
LR schedule
constant
constant
cosine
Gradient clipping
1.0
1.0
1.0
Candidates per prompt
4
4
4
Appendix
Table 7: RL-specific training settings. MDPO and MaskGRPO use 128-token understanding canvases. UniGRPO jointly trains MMaDA understanding and image generation; paired entries denote text and image settings.
Method
GenEval
CLIP score
ImageReward
Overall
Single
Two
Count
Colors
Position
Color attr.
MMaDA-8B-MixCoT
55.45
91.56
67.68
42.19
80.05
19.25
32.00
28.43
0.585
+ Endpoint-only Distillation
60.70 (+5.25)
96.25 (+4.69)
74.49 (+6.81)
42.81 (+0.62)
81.65 (+1.60)
23.50 (+4.25)
45.50 (+13.50)
28.99 (+0.56)
0.732 (+0.147)
+ UniGRPO
59.90 (+4.44)
95.31 (+3.75)
76.77 (+9.09)
43.44 (+1.25)
80.85 (+0.80)
24.50 (+5.25)
38.50 (+6.50)
28.81 (+0.38)
0.752 (+0.167)
+ Random
60.06 (+4.60)
95.00 (+3.44)
76.01 (+8.33)
40.00 ( − 2.19)
80.59 (+0.53)
22.75 (+3.50)
46.00 (+14.00)
29.01 (+0.58)
0.765 (+0.180)
+ CT-OPD
61.69 (+6.23)
96.88 (+5.31)
77.53 (+9.85)
42.81 (+0.63)
81.65 (+1.60)
24.25 (+5.00)
47.00 (+15.00)
29.16 (+0.73)
0.774 (+0.189)
Appendix
Table 8: CT-OPD improves every reported image-generation metric over the frozen MMaDA-8B base. GenEval reports the overall score and its six compositional categories. Parenthesized values are absolute changes over MMaDA-8B-MixCoT. Bold marks the best result for each metric.
Figure 3: Matched qualitative comparison on compositional image generation. Each column uses the same GenEval prompt, sample index, and sampling seed for MMaDA-8B-MixCoT, UniGRPO, Endpoint-only Distillation, and CT-OPD . The examples display differences in attribute binding, object separation, cardinality, and relative placement.
Mask
Δ NLL (nats/token) ↓
Δ accuracy (pp) ↑
100%
+0.0384[0.0299,0.0469]
−0.24[−0.44,−0.03]
75%
−0.5086[−0.5531,−0.4657]
+7.78[6.90,8.72]
50%
−0.9615[−1.0334,−0.8925]
+16.41[14.81,18.14]
25%
−1.6434[−1.7449,−1.5460]
+30.79[28.31,33.35]
Partial aggregate
−0.8272[−0.8717,−0.7847]
+14.06[13.13,15.07]
Appendix
Table 9: Canonical reconstruction and partial-state prediction. Direct Trace-Target preserves raw student-visible values while matching CT-OPD’s endpoints, online mask-selection protocol, optimizer, and exposure. Entries report CT-OPD minus Direct Trace-Target with prompt-paired 95% CIs.
Backbone / branch
Endpoint adapter
Training trace
Support
Results
LLaDA-V
Native text tokenizer
128-token canvas; 32 steps
VL
Table 1
LaViDa-LLaDA
Native text tokenizer
128-token canvas; 32 steps
VL
Table 1
DiMoE-VL
Native text tokenizer; answer-preserving fitting
128-token canvas; 32 steps
VD
Table 1
MMaDA / understanding
Native text tokenizer
512-token canvas; 128-token blocks; 256 steps
VM
Table 1
MMaDA / generation
Frozen MAGVITv2
1,024-token image canvas; 16 steps
VM
Table 8
Appendix
Table 10: Backbone interfaces and native training traces. Prediction supports specify the normalization domain of the training loss; trace settings describe the student rollouts used to select intermediate states.
Family
Endpoint adapter
Trajectory variable ξt
Transported state
Conditional score
Absorbing masked
Detokenize-retokenize
Unresolved-position mask gt
C(YT,gt)
Unresolved-token NLL
Categorical
Domain-specific endpoint map
Base noise and scheduler randomness
Ft(YT,ξt)
Denoising cross-entropy
Gaussian/latent
Frozen endpoint encoder
Trace residual ϵtS
αtYT+σtϵtS
x0 , ϵ , or v regression
Block diffusion
Domain-specific endpoint map
Block/reveal/remasking variables
Ft(YT,ξt)
Active-block NLL
Appendix
Table 11: CT-OPD specializations under the unified corruption map.
GRPO-style reinforcement learning has advanced vision-language reasoning, but its sequence-level rewards provide limited credit assignment for individual reasoning tokens. On-policy distillation (OPD) addresses this issue by providing dense supervision on student-generated trajectories, yet standard OPD assigns similar importance to all teacher-preferred tokens, including template phrases and visually irrelevant continuations. We introduce Visual Counterfactual On-Policy Distillation (VC-OPD), a grounding-aware distillation method that prioritizes tokens whose correctness depends on image-specific evidence. After a teacher-supervised warm start that stabilizes response formats and reduces noisy rollouts, VC-OPD performs on-policy distillation with two teacher evaluations for each student trajectory: one under the original image and one under a counterfactual image with corrupted instance-specific evidence. The original-image teacher distribution determines the distillation target, while the original--counterfactual teacher gap estimates token-level visual dependence and softly reweights the loss. To avoid unstable response-level update scales, VC-OPD further applies mass-preserving normalization to the token weights. This design preserves the dense supervision of OPD while shifting learning capacity toward visually grounded reasoning tokens. Experiments on multimodal reasoning benchmarks show that VC-OPD improves visual reasoning performance and produces more effective token-level distillation.
Jiahe Feng
College of Intelligence and Computing, Tianjin University, Tianjin, China
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× fewer rollout rounds, corresponding to an estimated 96.0× to-accuracy model-compute speedup.
Haolin Ren, Ziyang Huang, Chenhao Yuan +2
Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China
Improving the reasoning capabilities of vision-language models (VLMs) often relies on post-training with supervision from stronger models. Two common strategies are supervised fine-tuning (SFT), which trains on teacher-generated demonstrations, and on-policy distillation (OPD), which provides teacher feedback on trajectories generated by the student itself. While OPD can be more data-efficient by supervising student-generated trajectories, it is unclear how the teacher's supervision at each state affects the student's learned behavior. To study this question, we conduct a controlled comparison of OPD and filtered SFT across data budgets and tasks where the teacher exhibits different levels of reliability, and analyze how these differences shape student accuracy and completion behavior. Across the two tasks, we find that OPD is effective only when the teacher remains reliable on student-generated trajectories. On ChartQA, where the teacher completes 99% of its solutions, OPD trained on only 10% of the data comes within 1.4 accuracy points of full-data SFT. In contrast, on Geometry3K, where the teacher fails to complete 22% of its solutions, OPD underperforms SFT by 7.9 points, with most of the gap explained by incomplete responses. More strikingly, starting OPD from an SFT student reduces its completion rate from 76% to 62%, effectively erasing the completion behavior learned through SFT. Filtering teacher demonstrations for completeness alone recovers much of SFT's advantage. These results show that on-policy supervision can transfer not only teacher knowledge but also persistent failure modes on student-visited states. Code is available at https://github.com/ffy208/opd-inherits-teacher-vlm.