TamedFlow: Trajectory-Aware Distillation\\for Teacher-Assisted Diffusion Sampling

Shuaiqi Cheng; Jiahao Zhang; Boyuan Gu; Siyu You; Yuxi Chen; Haiying Xu; Zexin Lu; Huiyu Zhou

cs.CVOct 19, 2026

TamedFlow: Trajectory-Aware Distillation\for Teacher-Assisted Diffusion Sampling

Authors: Shuaiqi Cheng, Jiahao Zhang, Boyuan Gu, Siyu You, Yuxi Chen, Haiying Xu, Zexin Lu, Huiyu Zhou,

Organizations: The Hong Kong University of Science and Technology (Guangzhou) · University of Electronic Science and Technology of China · The Hong Kong University of Science and Technology · Huawei Technologies Ltd. ·

Abstract

Diffusion Transformers (DiTs) deliver high-fidelity image and video generation, but inference requires many sequential model evaluations. Two acceleration strategies are increasingly important: model compression makes each evaluation cheaper, while adaptive inference avoids selected teacher computation through reuse, forecasting, or speculative execution. Their combination, however, is brittle: directly applying the evaluated adaptive methods to a compressed student can degrade quality, an observed failure mode we call the Combination Gap. We link this gap to a mismatch between pointwise distillation and short-horizon execution: a student trained at teacher states need not remain reliable when its own predictions are chained. We propose TamedFlow, which trains a compressed DiT for this deployed role through cross-architecture feature alignment, a near-identity Velocity Alignment Head, and randomized rollout supervision. At inference, a lightweight Learnable Trajectory Estimator routes completed drafts: accepted drafts bypass the teacher, while rejected drafts roll back and receive one teacher rescue step. Distilling DiT-XL/2 into DiT-B/2, TamedFlow reaches FID $2.58$ at $2.25\times$ measured speedup over DDIM-50 on ImageNet $256\times256$ . Ablations support trajectory supervision, and routing controls favor LTE over random decisions at nearly matched latency. Evaluations on ImageNet $512\times512$ , MS-COCO, and Latte/FaceForensics extend the empirical scope.

Figures & tables

Figure 1: The Combination Gap and TamedFlow. Directly applying adaptive reuse to a parameter-compressed student severely degrades quality in our diagnostic. TamedFlow instead trains the compressed student for short-horizon trajectory agreement. A 4M-parameter LTE routes drafts: accepted drafts bypass the teacher; rejected drafts roll back for one teacher rescue step.

Figure 2: TamedFlow overview. Left: AlignNet, VAH, and LTE attached to the student backbone. Middle: two-stage curriculum. Stage 1 is feature alignment and Stage 2 is rollout supervision. Right: inference routes drafts via the LTE; accepted drafts bypass the teacher, rejected drafts roll back and take one teacher step.

Full-Scale Few-Step Distillation
Method	Latency (s) $\downarrow$	Speedup $\uparrow$	FLOPs (T) $\downarrow$	FID $\downarrow$	IS $\uparrow$
DDIM-50 Teacher	1.044	1.00 $\times$	23.74	2.35	237.62
Pure Student	0.180	5.80 $\times$	2.30	5.56 (+3.21)	195.07 (-42.55)
Progressive Distill. (step=4) ( Salimans and Ho, 2022 )	0.068	15.35 $\times$	$0.474$	10.25 (+7.90)	181.58 (-56.04)
LCM (step=4) ( Luo et al., 2023 )	0.072	14.50 $\times$	$0.476$	6.81 (+4.46)	202.73 (-34.89)
ARD (step=4) ( Kim et al., 2025 ) ( $N=6$ )	0.074	14.11 $\times$	$0.480$	4.32 (+1.97)	209.03 (-28.59)

Table 1: ImageNet

256\times 256

class-conditional generation with DiT-XL/2.

Figure 3: Qualitative comparison on ImageNet $256\times 256$ .

Study	Variant	Accept % $\uparrow$	Speedup $\uparrow$	FID $\downarrow$	IS $\uparrow$
Component	Full TamedFlow	63.6	2.25 $\times$	2.58	233.43
	w/o AlignNet	46.2	1.87 $\times$	2.95	211.34
	w/o VAH	39.3	1.63 $\times$	3.23	205.71
	w/o LTE	–	5.80 $\times$	5.56	195.07
Training	Full TamedFlow	63.6	2.25 $\times$	2.58	233.43
	w/o $\mathcal{L}_{\mathrm{roll}}$	54.2	2.02 $\times$	2.71	230.15

Table 2: Ablation studies on ImageNet

256\times 256

K=3

and

\tau=0.80

. Speedup is relative to the 50-step teacher baseline.

Figure 4: LTE routing behavior across generation. Left: confidence across sampling steps; the upper axis uses a noise-level coordinate, distinct from forward flow time. Right: score distributions of accepted and rejected drafts under

\tau=0.80

Method	Latency (s) $\downarrow$	Speedup $\uparrow$	FID $\downarrow$	CLIP-S $\uparrow$
Teacher DDIM-50	1.184	1.00 $\times$	13.54	30.2
Teacher DDIM-25	0.622	1.90 $\times$	16.84	28.9
Pure Student	0.203	5.83 $\times$	19.30	27.8
TaylorSeer ( $N=3$ )	0.580	2.04 $\times$	15.47	29.1
SpeCa	0.561	2.11 $\times$	16.63	28.4
FoCa ( $N=3$ )	0.565	2.09 $\times$	16.27	28.5

Table 3: MS-COCO 30K

256\times 256

text-to-image generation using fixed-step DDIM. TaylorSeer, SpeCa, and FoCa are re-implemented within our DiT-XL/2 pipeline;

N

denotes each method’s acceleration setting rather than TamedFlow’s draft horizon

K

Method	Speedup	FVD $\downarrow$	IS $\uparrow$
Teacher (Latte-XL)	1.00 $\times$	208.50	71.20
FoCa	2.06 $\times$	248.30	66.80
TaylorSeer	2.10 $\times$	253.10	65.90
TamedFlow	2.22 $\times$	225.10	68.90

Table 4: Latte video-DiT transfer on FaceForensics (

256\times 256

, 16 frames).

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Stage 1	Stage 2
Optimizer	AdamW ( $\beta_{1}=0.9$ , $\beta_{2}=0.999$ )
Base learning rate	$5\times 10^{-5}$	$1\times 10^{-4}$
Weight decay	0.0
Learning rate schedule	cosine decay with warmup
Warmup epochs	5
Max epochs	10	80

Appendix

Table 5: Training hyperparameters for TamedFlow.

Method	Latency (s) $\downarrow$	Speedup $\uparrow$	FID $\downarrow$	IS $\uparrow$
Teacher (DiT-XL/2, 675M)	3.842	1.00 $\times$	2.87	235.10
TaylorSeer ( $N=3$ )	1.715	2.24 $\times$	3.76	226.75
FoCa ( $N=3$ )	1.628	2.36 $\times$	3.47	228.10
TamedFlow ( $\tau=0.80$ )	1.594	2.41 $\times$	3.39	230.85

Appendix

Table 6: ImageNet

512\times 512

evaluation with DiT-XL/2 using fixed-step DDIM.

Method	Speedup	FID	IS
DDIM-50 Teacher	1.00 $\times$	$2.36\pm 0.02$	$237.24\pm 0.48$
TaylorSeer ( $N=3$ )	2.19 $\times$	$2.71\pm 0.04$	$229.13\pm 0.62$
FoCa ( $N=3$ )	2.21 $\times$	$2.64\pm 0.03$	$231.24\pm 0.58$
TamedFlow, $\tau=0.80$	2.26 $\times$	$2.56\pm 0.03$	$233.89\pm 0.51$

Appendix

Table 7: Three-seed ImageNet evaluation (mean

\pm

standard deviation over seeds).

Figure 5: Compact routing and execution accounting at

K=3

\tau=0.80

. The panels use distinct denominators: (a) 14 accepted and 8 rejected routing events; (b) 42 student-covered and 8 teacher-rescue nominal steps; and (c) the separately collected component profile. The component profile totals

474.85

ms and is not a decomposition of the separate

464

ms headline average.

Figure 6: Cross-setting transfer shown as six compact native-unit panels. Columns are COCO, Latte/FaceForensics, and ImageNet $512\times 512$ ; the top row reports quality metrics and the bottom row reports semantic or recognition scores. Row labels include the reported speedup relative to the teacher for that setting.

Support (Reference)	Accept	Reject
Count	11,984	4,683
Metric	Value
False Accept Rate (FAR)	0.063
False Reject Rate (FRR)	0.051
AUC	0.94
Continuous-target ECE (10 bins)	0.034

Appendix

Table 8: LTE classification and calibration metrics on the ImageNet

256\times 256

validation set at

\tau=0.80

K=3

. Support counts reference labels obtained by thresholding

c^{\star}

at 0.80. FAR and FRR use reference-reject and reference-accept events as their respective denominators. AUC uses continuous predictions; continuous-target ECE uses 10 confidence bins.

Routing policy	Teacher updates	Latency (s) $\downarrow$	FID $\downarrow$
Pure Student	0	0.180	5.56
Static: Random Steps	8	0.348	3.72
Static: Fixed Periodic	8	0.335	3.45
Dynamic: Random Accept	8	0.465	4.05
Dynamic: LTE	Adaptive	0.464	2.58

Appendix

Table 9: Routing controls on ImageNet

256\times 256

. Latency and FID-50K are measured. Teacher updates are fixed by the control policy or adaptive for LTE; this is not an equal-teacher-budget comparison with LTE.

Variant	Accept Rate % $\uparrow$	Speedup $\uparrow$	FID $\downarrow$	IS $\uparrow$
Full TamedFlow (default)	63.6	2.25 $\times$	2.58	233.43
Module and loss ablations
w/o AlignNet	46.2	1.87 $\times$	2.95	211.34
w/o VAH	39.3	1.63 $\times$	3.23	205.71
w/o $\mathcal{L}_{\mathrm{roll}}$	54.2	2.02 $\times$	2.71	230.15
w/o $\mathcal{L}_{z}$ warmup	46.2	1.81 $\times$	3.29	203.50

Appendix

Table 10: Protocol-consistent extended ablation on ImageNet

256\times 256

. Inference uses

K=3

\tau=0.80

; the randomized horizons refer to training.

Strategy	FID	Strategy	FID
FoCa	10.40	TaylorSeer	11.92
SpeCa	13.18	Dense Student	5.56

Appendix

Table 11: Same-student validation under the common DDIM-50 protocol. Only FID is reported because these direct adaptive baselines were evaluated as quality diagnostics, not as the headline latency comparison.

Figure 7: Module and training ablations from Tables 2 and 10 . The panels relate FID to acceptance or speedup at

K=3

\tau=0.80

. The no-routing student has no acceptance statistic and is omitted from the acceptance panels. Points are reported operating points; no fitted trend or uncertainty interval is implied.

Solver	Method	Latency	Speedup	Accept	FID	IS
Euler	Teacher	1.044	1.00 $\times$	–	2.35	237.62
Euler	TamedFlow	0.464	2.25 $\times$	63.6%	2.58	233.43
DPM-Solver++(2M)	Teacher	1.089	1.00 $\times$	–	2.25	244.12
DPM-Solver++(2M)	TamedFlow	0.455	2.39 $\times$	75.0%	2.44	237.55

Appendix

Table 12: Solver robustness on ImageNet

256\times 256

. DPM-Solver++(2M) is evaluated as a handoff-safe higher-order solver diagnostic; the formal bound remains DDIM-specific.

Figure 8: Paired teacher and TamedFlow measurements under Euler and DPM-Solver++(2M), reproducing Table 12 . Speedups are relative to the corresponding solver teacher. The DPM++ result uses the stated history-reset policy; the theoretical bound remains DDIM-specific.

Figure 9: Held-out velocity mismatch and three-step endpoint drift by model variant. The panels replace the redundant numerical table and use separate zero-based axes. These means do not measure Jacobian mismatch or establish that training minimizes the theoretical bound constants; no uncertainty estimates are available.

Figure 10: Impact of draft length

K

and routing threshold

\tau

on generation quality and inference efficiency. Columns: FID (left) and speedup (right). Rows: threshold

\tau

(top) and draft length

K

(bottom). The sweeps motivate the default

K=3

\tau=0.80

configuration.

Diffusion Sampling

Diffusion sampling studies how to turn a trained diffusion model into high-quality samples efficiently and reliably.

See all 260 papers

Consistent Distribution Matching for Data-Free Diffusion Distillation
Oct 6, 2026Yuxiang Fu, Qi Yan, Zike Wu +4Diffusion Model Distillation Diffusion Model Inference Acceleration
University of British Columbia · Vector Institute for AI · Canada CIFAR AI Chair
Stabilizing, Scaling & Enhancing MeanFlow for Large-scale Diffusion Distillation
May 18, 2026Xiao He, Yang Li, Peizhen Zhang +3Flow Matching Diffusion Model Distillation
1State Key Laboratory of Integrated Services Networks, Xidian University · 2Tencent Hunyuan
Predict to Skip: Linear Multistep Feature Forecasting for Efficient Diffusion Transformers
Feb 20, 2026Hanshuai Cui, Zhiqing Tang, Qianli Ma +3Diffusion Model Acceleration Diffusion Sampling
School of Artificial Intelligence, Beijing Normal University, Beijing, China · Institute of Artificial Intelligence and Future Networks, Beijing Normal University, Zhuhai, China · Shenzhen University of Advanced Technology, Shenzhen, China

Oct 6, 2026J/K move · Enter open · S save

Consistent Distribution Matching for Data-Free Diffusion Distillation

Yuxiang Fu, Qi Yan, Zike Wu, Yongxing Zhang+3

University of British Columbia · Vector Institute for AI · Canada CIFAR AI Chair

Diffusion Model Distillation Diffusion Model Inference Acceleration Distribution Matching Distillation Data-Free Knowledge Distillation

Open paper AskPDF arXiv

Flow and diffusion models suffer from slow inference due to computationally expensive numerical integration. Distillation provides a promising way for a student model to learn from a teacher's dynamics, enabling one-step or few-step generation. However, existing methods often depend on curated distillation datasets, costly teacher rollouts, or auxiliary proxy networks, which complicate model training and scaling. In this work, we propose Consistent Distribution Matching, a simulation-free and data-free distillation method for accelerating diffusion and flow models while preserving strong generative capacity. Our key insight is to unify sample generation and score estimation with one student network. Thus, our framework uses only two models, a frozen teacher and a trainable student, and optimizes one objective. We prove that minimizing our objective indicates Wasserstein convergence of the student flow-map pushforwards to the teacher marginals. On ImageNet 256

\times

256, our method attains an FID of 2.04 with a single function evaluation (1-NFE) and a 4-NFE FID of 1.37 within 40 epochs of training, surpassing the state-of-the-art distillation baselines without data. Our code code and model are available at https://consistentdmd.github.io/.

Abstract

Diffusion Sampling

Explore similar work

Consistent Distribution Matching for Data-Free Diffusion Distillation

Stabilizing, Scaling & Enhancing MeanFlow for Large-scale Diffusion Distillation

Predict to Skip: Linear Multistep Feature Forecasting for Efficient Diffusion Transformers