在图像生成领域，扩散模型因其卓越的性能和可扩展性而备受瞩目。然而，尽管Diffusion Transformers（DiTs）等模型在生成任务中表现出色，但它们却摒弃了U-Net架构，这一在先前工作中被广泛采用的架构。对此，北京大学和华为的研究人员提出了U-shaped Diffusion Transformers（U-DiTs），旨在重新审视U-Net架构在扩散模型中的潜力。
U-Net架构在图像分割和生成任务中一直扮演着重要角色，其独特的编码器-解码器结构以及跳跃连接设计，使得它在处理图像数据时具有天然的优势。然而，在扩散模型的领域中，U-Net架构却逐渐被边缘化，取而代之的是更为简洁的同构架构。
北大与华为的研究人员对此现象进行了深入的思考和探索。他们首先进行了一个简单的实验，将DiT块嵌入到一个典型的U-Net架构中，形成了DiT-UNet。通过与同构DiT的比较，他们发现DiT-UNet在相似的计算成本下，性能仅略优于DiT。这表明，简单地将U-Net与Transformer块结合，并不能充分发挥U-Net的归纳偏置优势。
为了更好地将Transformer的注意力机制与U-Net架构相结合，研究人员对U-Net作为扩散去噪器的角色进行了重新审视。他们发现，U-Net骨干网络的特征在低频域中占主导地位，这暗示了骨干特征中可能存在冗余。基于这一发现，他们提出了一种名为Token Downsampling的创新方法。
Token Downsampling的核心思想是通过下采样操作来过滤高频噪声，从而突出低频域的特征。与之前仅对键值对进行下采样的方法不同，研究人员采取了更为激进的策略，同时对查询、键和值进行下采样。这种做法不仅减少了计算成本，还避免了信息损失。
通过将Token Downsampling与DiT-UNet相结合，研究人员取得了令人惊喜的结果。在ImageNet 256x256的生成任务中，他们提出的U-DiT模型在FID指标上取得了显著的改进，同时计算成本也大幅降低。
基于Token Downsampling的发现，研究人员进一步扩展了U-Net架构，提出了一系列的U-DiT模型。这些模型在ImageNet 256x256和512x512的生成任务中，都展现出了卓越的性能。
在ImageNet 256x256的生成任务中，U-DiT-B模型仅用六分之一的计算成本，就超越了DiT-XL/2模型的性能。而在ImageNet 512x512的生成任务中，U-DiT-B模型也以更低的计算成本，取得了更好的性能。
这些结果表明，U-DiT模型不仅在性能上超越了之前的DiT模型，而且在计算效率上也具有显著的优势。这对于实际应用中的图像生成任务来说，具有重要的意义。
U-DiT模型的提出，为扩散模型的发展注入了新的活力。它不仅重新审视了U-Net架构在扩散模型中的潜力，还通过Token Downsampling等创新方法，实现了性能和计算效率的双重提升。
然而，U-DiT模型也存在一些潜在的挑战和限制。首先，由于计算资源的限制，研究人员无法进一步扩展训练迭代次数和模型规模，以充分探索U-DiT的潜力。其次，由于训练数据集的偏见，生成的内容可能包含色情、种族主义、仇恨和暴力等信息。虽然研究人员强调了通过谨慎应用来减轻滥用风险，但这仍然是一个需要关注的问题。
