研究团队为突破XR边缘计算瓶颈，提出高效DiT模型压缩方案
设计参数和计算效率高的Dit
（映维网Nweon 2025年06月13日）具有数十亿个模型参数的Diffusion Transformer（DiT）构成了DALL.E，Stable-Diffusion和SORA等流行图像和视频生成模型的主干。尽管相关模型在增强现实/虚拟现实等低延迟应用中十分必要，但由于其巨大的计算复杂性，它们无法部署在资源受限的边缘设备，如Apple Vision Pro或Meta Ray-Ban眼镜。
为了克服这个问题，伊利诺伊大学厄巴纳-香槟分校团队转向Knowledge Distillation，并执行彻底的设计空间探索，以实现给定参数大小的最佳DiT。特别是，团队提供了如何选择design knob的原则，如深度，宽度，attention head。
在模型性能、大小和速度之间出现了三方面的权衡，而这对边缘实现扩散至关重要。研究人员同时提出了两种蒸馏方法Teaching Assistant (TA)和Multi-In-One (MI1)，以在DiT上下文中执行特征蒸馏。
由于其高保真度、可泛化性、易于训练和可扩展性，DiT已经成为生成图像和视频的重要方法。DiT构成了各种实际部署的图像和视频生成模型的支柱，例如DALL.E，Stable-Diffusion和SORA。
出于模型的大参数大小和计算复杂性，必须使用云服务来远程运行它们。从云到边缘的数据传输相关的显著延迟无法为需要在资源受限的边缘设备实现的高帧率应用提供支持，例如增强现实/虚拟现实。
在边缘设备直接实现神经网络推理的主要挑战来自边缘硬件有限的内存和能量容量。为了解决这个问题，我们需要设计参数和计算效率高的Dit。
边缘设备通常拥有几兆字节的片上存储器，现有的实用模型则需要模型尺寸百万参数的数量级，涉及数十亿个参数。
先前关注高效Dit的研究只对特定层进行优化，或者只关注精度，或者没有推动实现预期性能所需的参数限制。伊利诺伊大学厄巴纳-香槟分校团队的重点不是通过新颖的算法方法提供SOTA DiT模型。相反，目标是使用有原则的设计选择，在给定的参数大小下提供最佳的DiT模型（在性能和速度方面）。
对于设计空间探索，在提取DiT模型的几个design knob中，研究人员选择了以下最相关的几个——深度、宽度、attention head，以及用于提取的设置（损失函数和teacher模型）。前两个knob影响效率和性能，后两个knob只影响性能。
尽管时间步长十分重要，但团队不认为它是一个design knob，因为业界已经对其进行了广泛的研究。
下面团队提出了两种方法来探索DiT的新蒸馏设置：
-
Teaching Assistant (TA)：用于提取卷积网络，而团队探索了使用具有LPIPS损失的teacher和TA组合特征蒸馏的可能性。
Multi-In-One (MI1)：通过将扩散样本映射到student的特定层，在单个步骤中执行多个扩散时间步。采用前向扩散概率流ODE，利用多步扩散的teacher模型的噪点-图像对计算中间噪点图像。
表3表明，使用TA进行特征蒸馏没有帮助。只有直接使用TA的蒸馏提供了边际效益。MI1的性能比基线差，如表4所示。
需要注意的一个关键方面是，中间层的约束并不是导致性能较差的原因，因为（2,4,6）的性能优于（3,6）。最后，团队将基线方法与唯一基于SOTA扩散的transfomer模型进行了模型参数蒸馏比较，并在所有指标上都打败了它们，包括模型大小、FID和延迟。
相关论文：Designing Parameter and Compute Efficient Diffusion Transformers using Distillation
总的来说，团队对DiT蒸馏进行了彻底的设计空间探索，并提供了获得给定模型尺寸的SOTA DiT的设计原则。当DiT在NVIDIA Jetson Orin Nano实现时，团队确定了模型性能-尺寸-速度之间的关键权衡，并可以指导未来研究人员在实际领域进行创新。
对于这个研究，研究人员同时希望强调在将新方法与先前的作品进行比较之前，创建强大而明显的基线的实践。尽管TA方法略优于基线，但结论是，与TA和MI1方法相比，根据原则设计的student模型是一个更好的选择，因为它的训练成本更低。
未来的方向可以包括分析地证明上述指导方针，或将设计空间扩展到MLP比率和扩散时间步等knob，或为每个层定制attention head，特别是因为改变attention head大小会对延迟产生影响。
