让大模型在手机眼镜上跑起来，靠蒸馏压缩扩散模型
Designing Parameter and Compute Efficient Diffusion Transformers using Distillation
- 用知识蒸馏压缩百亿参数扩散模型，适配边缘设备
- 提出两类新蒸馏方法，在Jetson设备上实现高效推理
- 平衡性能、体积与速度，适合可穿戴设备部署
百亿参数的扩散Transformer（DiT）是DALL·E、Stable Diffusion和SORA等图像与视频生成模型的核心。尽管这些模型在增强现实/虚拟现实等低延迟场景中至关重要，但其巨大的计算复杂度使其无法部署于资源受限的边缘设备（如Apple Vision Pro或Meta Ray-Ban眼镜）。为此，本文采用知识蒸馏并开展全面的设计空间探索，以在给定参数量下优化DiT性能。我们提出选择深度、宽度、注意力头数及蒸馏配置的设计原则，过程中发现性能、大小与速度之间的三重权衡对边缘部署至关重要。此外，我们提出两种新蒸馏方法：教学助教（TA）与多合一（MI1），用于在DiT框架内进行特征蒸馏。不同于现有方案，我们在NVIDIA Jetson Orin Nano等实际边缘设备上验证并基准测试了所提方法的有效性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。