北大&华为“逆向革新”扩散模型：U-DiT以六分之一算力超越DiT，U-Net架构重回巅峰
引言： 大型语言模型（LLM）的浪潮席卷全球，而生成式AI的另一支重要力量——扩散模型，也在不断演进。近期，北京大学和华为的研究团队发表了一篇题为“U-DiTs: Downsample Tokensin U-Shaped Diffusion Transformers”的论文，提出了一种名为U-DiT的新型扩散模型架构。该架构巧妙地结合了曾被主流Transformer架构取代的U-Net结构，并通过创新的下采样自注意力机制，以极低的算力消耗实现了超越现有DiT模型的生成效果，堪称对扩散模型领域的一次“逆向革新”。
主体：
近年来，基于Transformer架构的扩散模型，例如DiT，凭借其强大的建模能力在图像生成领域取得了显著进展。DiT模型抛弃了传统的U-Net架构，采用了直筒型去噪模型，并在隐空间生成任务中展现出优异性能，引领了PixArt、SD3等后续工作的潮流。然而，U-Net架构作为先前扩散模型的基石，其在图像空间和隐空间的生成效果也同样出色，其“归纳偏置”（inductive bias）在扩散任务中的有效性早已得到广泛验证。
这种现状引发了北大与华为研究团队的思考：能否将U-Net架构与Transformer的优势相结合，从而进一步提升扩散模型的性能？ 带着这一疑问，他们深入研究，最终提出了U-DiT架构。
论文首先回顾了一个简单的实验：直接将U-Net与DiT模块结合。结果显示，在相似的算力条件下，这种简单的结合（DiT-UNet）仅比原始DiT略有提升，这表明简单的叠加并不能充分发挥两种架构的优势。
研究团队随后深入分析了U-Net在扩散模型中的作用机制。他们发现，U-Net的主干结构特征图主要包含低频信号。而Transformer中的全局自注意力机制计算量巨大，在U-Net的主干中使用全局自注意力可能存在冗余计算。
基于这一观察，研究团队提出了一个关键的创新：下采样自注意力机制。该机制在进行自注意力计算之前，先对特征图进行2倍下采样，以滤除噪声较多的高频信息，并保留信息丰富的低频信息。为了避免信息丢失，他们生成了四个维度相同的下采样图，确保下采样前后特征总维度不变。然后，在四个下采样图上分别独立进行自注意力运算（共享QKV映射），最后将结果融合成一个完整的特征图。
这种下采样自注意力机制巧妙地降低了计算复杂度。与传统的全局自注意力相比，它将计算量降低了3/4，显著提升了效率。令人惊喜的是，这种看似简单的改进不仅降低了算力消耗，反而提升了模型的生成效果。
基于下采样自注意力机制，研究团队构建了U-DiT模型。他们训练了三个不同规模的U-DiT模型（S/B/L），并与DiT模型进行了对比。实验结果表明，在ImageNet生成任务上，U-DiT-L在400K次迭代后，其FID分数比DiT-XL高约10；U-DiT-S/B模型比同级别DiT模型高约30 FID。更令人瞩目的是，U-DiT-B模型仅需DiT-XL六分之一的算力，就能达到甚至超越后者在ImageNet上的生成效果。
此外，U-DiT在有条件生成任务和大图生成任务上也展现出显著优势。 通过延长训练迭代次数，U-DiT-L在600K次迭代后，其无条件生成效果甚至超越了DiT在7M次迭代后的结果。 在1M次迭代下，U-DiT的有条件生成效果已经非常逼真。
结论：
北大&华为提出的U-DiT模型，通过巧妙地结合U-Net架构和创新的下采样自注意力机制，在降低算力消耗的同时，显著提升了扩散模型的生成效果，为扩散模型的发展提供了新的方向。这项研究不仅证明了U-Net架构在扩散模型中的持续价值，也为高效、高质量的图像生成提供了新的解决方案。U-DiT的成功，也为未来研究人员探索更有效率、更强大的扩散模型架构提供了宝贵的经验和启示。 U-DiT的开源也为广大的研究者和开发者提供了便捷的工具，有望进一步推动扩散模型技术的普及和应用。 未来，我们有理由期待基于U-DiT架构的更多创新应用，以及其在更广泛领域的应用前景。
参考文献：
- Tian,Y., et al. (2024). U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers. arXiv preprint arXiv:2405.02730. (论文链接：https://arxiv.org/pdf/2405.02730)
- (机器之心报道链接: 此处应补充机器之心报道的链接)
*(注：由于无法访问互联网，我无法提供机器之心报道的具体链接。请读者自行搜索相关报道。) *
Views: 0