Designing Parameter and Compute Efficient Diffusion Transformers using Distillation

📄 arXiv: 2502.14226v1 📥 PDF

作者: Vignesh Sundaresha

分类: cs.CV, eess.IV

发布日期: 2025-02-20

备注: 4 pages

💡 一句话要点

利用知识蒸馏设计参数和计算高效的Diffusion Transformer，适用于边缘设备。

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)

关键词: Diffusion Transformer 知识蒸馏 边缘计算 模型压缩 特征蒸馏

📋 核心要点

  1. 现有Diffusion Transformer模型参数巨大，计算复杂度高，难以部署在资源受限的边缘设备上。
  2. 论文提出利用知识蒸馏，探索DiT模型深度、宽度等设计空间，优化模型在边缘设备上的性能。
  3. 论文提出了教学助手（TA）和多合一（MI1）两种特征蒸馏方法，并在NVIDIA Jetson Orin Nano上验证了有效性。

📝 摘要（中文）

Diffusion Transformer (DiT) 拥有数十亿的模型参数，是 DALL.E、Stable-Diffusion 和 SORA 等流行的图像和视频生成模型的基础。虽然这些模型在增强现实/虚拟现实等许多低延迟应用中是必需的，但由于其巨大的计算复杂度，它们无法部署在资源受限的边缘设备（如 Apple Vision Pro 或 Meta Ray-Ban 眼镜）上。为了克服这个问题，本文转向知识蒸馏，并进行了彻底的设计空间探索，以针对给定的参数大小实现最佳的 DiT。特别地，本文为如何选择 DiT 的深度、宽度、注意力头和蒸馏设置等设计旋钮提供了原则。在此过程中，模型性能、大小和速度之间出现了三向权衡，这对于扩散的边缘实现至关重要。本文还提出了两种蒸馏方法——教学助手（TA）方法和多合一（MI1）方法——以在 DiT 上下文中执行特征蒸馏。与现有解决方案不同，本文在 NVIDIA Jetson Orin Nano 等实际边缘设备上展示并基准测试了本文方法的有效性。

🔬 方法详解

问题定义：论文旨在解决Diffusion Transformer模型参数量过大，计算复杂度过高，难以在资源受限的边缘设备上部署的问题。现有方法通常直接对大型模型进行压缩，但忽略了针对边缘设备特点进行模型结构和训练策略的优化，导致性能下降明显。

核心思路：论文的核心思路是利用知识蒸馏技术，将大型教师模型的知识迁移到小型学生模型，同时进行设计空间探索，针对边缘设备的计算资源和延迟要求，优化学生模型的结构和训练策略。通过在模型性能、大小和速度之间进行权衡，找到最适合边缘设备的DiT模型。

技术框架：整体框架包括一个大型的预训练Diffusion Transformer教师模型和一个小型化的学生模型。首先，使用教师模型生成图像或视频。然后，利用提出的教学助手（TA）或多合一（MI1）方法，将教师模型的特征信息传递给学生模型。学生模型通过最小化与教师模型特征的差异进行训练。最后，将训练好的学生模型部署到边缘设备上。

关键创新：论文的关键创新在于提出了两种新的特征蒸馏方法：教学助手（TA）方法和多合一（MI1）方法。TA方法使用一个中间层作为教学助手，帮助学生模型更好地学习教师模型的特征。MI1方法将多个教师模型的知识融合到一个学生模型中，提高学生模型的泛化能力。此外，论文还系统地探索了DiT模型的设计空间，为边缘设备的模型设计提供了指导。

关键设计：在模型设计方面，论文探索了深度、宽度、注意力头数等参数对模型性能的影响，并提出了相应的选择原则。在损失函数方面，论文使用了特征匹配损失，鼓励学生模型的特征与教师模型的特征尽可能接近。在训练策略方面，论文使用了多阶段训练，首先训练学生模型的浅层网络，然后逐步训练深层网络，以提高训练效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在NVIDIA Jetson Orin Nano边缘设备上进行了实验，结果表明，使用提出的蒸馏方法可以显著减小DiT模型的参数量和计算复杂度，同时保持较高的生成质量。具体性能数据未知，但论文强调了在实际边缘设备上的有效性验证，优于以往仅在模拟环境下的实验。

🎯 应用场景

该研究成果可广泛应用于增强现实（AR）、虚拟现实（VR）、智能监控、自动驾驶等需要在边缘设备上进行图像和视频生成的领域。例如，在AR/VR应用中，可以利用该方法在眼镜等设备上实时生成高质量的图像，提升用户体验。在智能监控中，可以在摄像头等设备上进行实时的图像分析和处理，提高监控效率。在自动驾驶中，可以利用该方法在车辆上进行实时的环境感知和决策，提高驾驶安全性。

📄 摘要（原文）

Diffusion Transformers (DiTs) with billions of model parameters form the backbone of popular image and video generation models like DALL.E, Stable-Diffusion and SORA. Though these models are necessary in many low-latency applications like Augmented/Virtual Reality, they cannot be deployed on resource-constrained Edge devices (like Apple Vision Pro or Meta Ray-Ban glasses) due to their huge computational complexity. To overcome this, we turn to knowledge distillation and perform a thorough design-space exploration to achieve the best DiT for a given parameter size. In particular, we provide principles for how to choose design knobs such as depth, width, attention heads and distillation setup for a DiT. During the process, a three-way trade-off emerges between model performance, size and speed that is crucial for Edge implementation of diffusion. We also propose two distillation approaches - Teaching Assistant (TA) method and Multi-In-One (MI1) method - to perform feature distillation in the DiT context. Unlike existing solutions, we demonstrate and benchmark the efficacy of our approaches on practical Edge devices such as NVIDIA Jetson Orin Nano.