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

📄 arXiv: 2610.26025v1 📥 PDF

作者: Shuaiqi Cheng, Jiahao Zhang, Boyuan Gu, Siyu You, Yuxi Chen, Haiying Xu, Zexin Lu, Huiyu Zhou, Xuming Hu

分类: cs.CV, cs.AI

发布日期: 2026-10-19

💡 一句话要点

提出TamedFlow以解决教师辅助扩散采样中的组合差距问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散变换器 模型压缩 自适应推理 轨迹估计 图像生成 视频生成 蒸馏训练

📋 核心要点

  1. 现有的扩散变换器在推理时需要多次模型评估，导致效率低下，且模型压缩与自适应推理的结合存在质量下降的风险。
  2. TamedFlow通过跨架构特征对齐和速度对齐头，结合随机回滚监督，训练压缩的DiT以适应推理阶段的需求，解决了组合差距问题。
  3. 在ImageNet数据集上，TamedFlow将FID降低至2.58，并实现了2.25倍的速度提升，验证了其有效性和优越性。

📝 摘要（中文）

扩散变换器（DiTs）在图像和视频生成中表现出色，但推理过程需要多次模型评估。现有的加速策略包括模型压缩和自适应推理，但将这两者结合时可能导致质量下降，称为组合差距。TamedFlow通过跨架构特征对齐、近似身份的速度对齐头和随机回滚监督来训练压缩的DiT，以解决这一问题。在推理阶段，轻量级的可学习轨迹估计器负责路由生成草稿，提升了推理效率和生成质量。

🔬 方法详解

问题定义：本论文旨在解决教师辅助扩散采样中的组合差距问题，即在模型压缩与自适应推理结合时，生成质量可能下降的现象。现有方法在将自适应推理应用于压缩学生模型时，往往无法保持可靠性。

核心思路：TamedFlow的核心思路是通过跨架构特征对齐和速度对齐头来训练压缩的DiT，使其在推理阶段能够有效利用教师模型的知识，同时避免组合差距带来的影响。

技术框架：整体架构包括三个主要模块：跨架构特征对齐模块、速度对齐头和随机回滚监督机制。在推理阶段，使用可学习的轨迹估计器来路由生成的草稿，优化推理效率。

关键创新：TamedFlow的关键创新在于引入了速度对齐头和随机回滚监督，这些设计使得压缩模型在推理时能够更好地利用教师模型的知识，显著提高了生成质量与效率。

关键设计：在模型训练中，采用了特征对齐损失和轨迹估计损失，确保模型在生成草稿时能够保持高质量的输出。同时，设计了轻量级的轨迹估计器，以便在推理时快速决策，减少延迟。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在ImageNet数据集上，TamedFlow将FID降低至2.58，并实现了2.25倍的速度提升，相较于DDIM-50表现出显著的性能优势。实验结果表明，轨迹监督和路由控制的设计优于随机决策，且在延迟几乎相当的情况下，提升了生成质量。

🎯 应用场景

TamedFlow的研究成果在图像和视频生成领域具有广泛的应用潜力，尤其是在需要高效推理的场景中，如实时视频处理、虚拟现实和增强现实等。其高效的生成能力和较低的计算需求使其在工业界和学术界都具备重要的实际价值，未来可能推动更多智能生成应用的发展。

📄 摘要（原文）

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 \textbf{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 \textbf{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.