Make U-Nets Great Again!北大&华为提出扩散架构U-DiT,六分之一算力即可超越DiT

Sora 的发布让广大研究者及开发者深刻认识到基于 Transformer 架构扩散模型的巨大潜力。作为这一类的代表性工作,DiT 模型抛弃了传统的 U-Net 扩散架构,转而使用直筒型去噪模型。鉴于直筒型 DiT 在隐空间生成任务上效果出众,后续的一些工作如 PixArt、SD3 等等也都不约而同地使用了直筒型架构。

然而令人感到不解的是,U-Net 结构是之前最常用的扩散架构,在图像空间和隐空间的生成效果均表现不俗;可以说 U-Net 的 inductive bias 在扩散任务上已被广泛证实是有效的。因此,北大和华为的研究者们产生了一个疑问:能否重新拾起 U-Net,将 U-Net 架构和 Transformer 有机结合,使扩散模型效果更上一层楼?带着这个问题,他们提出了基于 U-Net 的 DiT 架构 U-DiT。

从一个小实验谈开去

首先,研究者开展了一个小实验,在实验中尝试着将 U-Net 和 DiT 模块简单结合。然而,如表 1 所示,在相似的算力比较下,U-Net 的 DiT(DiT-UNet)仅仅比原始的 DiT 有略微的提升。

在图 3 中,作者们展示了从原始的直筒 DiT 模型一步步演化到 U-DiT 模型的过程。

根据先前的工作,在扩散中 U-Net 的主干结构特征图主要为低频信号。由于全局自注意力运算机制需要消耗大量算力,在 U-Net 的主干自注意力架构中可能存在冗余。这时作者注意到,简单的下采样可以自然地滤除噪声较多的高频,强调信息充沛的低频。既然如此,是否可以通过下采样来消除对特征图自注意力中的冗余?

Token 下采样后的自注意力

由此,作者提出了下采样自注意力机制。在自注意力之前,首先需将特征图进行 2 倍下采样。为避免重要信息的损失,生成了四个维度完全相同的下采样图,以确保下采样前后的特征总维度相同。随后,在四个特征图上使用共用的 QKV 映射,并分别独立进行自注意力运算。最后,将四个 2 倍下采样的特征图重新融为一个完整特征图。和传统的全局自注意力相比,下采样自注意力可以使得自注意力所需算力降低 3/4。

令人惊讶的是,尽管加入下采样操作之后能够显著模型降低所需算力,但是却反而能获得比原来更好的效果(表 1)。

U-DiT:全面超越 DiT

根据此发现,作者提出了基于下采样自注意力机制的 U 型扩散模型 U-DiT。对标 DiT 系列模型的算力,作者提出了三个 U-DiT 模型版本(S/B/L)。在完全相同的训练超参设定下,U-DiT 在 ImageNet 生成任务上取得了令人惊讶的生成效果。其中,U-DiT-L 在 400K 训练迭代下的表现比直筒型 DiT-XL 模型高约 10 FID,U-DiT-S/B 模型比同级直筒型 DiT 模型高约 30 FID;U-DiT-B 模型只需 DiT-XL/2 六分之一的算力便可达到更好的效果(表 2、图 1)。

在有条件生成任务(表 3)和大图(512*512)生成任务(表 5)上,U-DiT 模型相比于 DiT 模型的优势同样非常明显。

研究者们还进一步延长了训练的迭代次数,发现 U-DiT-L 在 600K 迭代时便能优于 DiT 在 7M 迭代时的无条件生成效果(表 4、图 2)。

U-DiT 模型的生成效果非常出众,在 1M 次迭代下的有条件生成效果已经非常真实。

论文已被 NeurIPS 2024 接收,更多内容,请参考原论文。

相关推荐
算力魔方AIPC13 分钟前
Meta重磅发布Llama 3.3 70B:开源AI模型的新里程碑
人工智能·llama
CSBLOG22 分钟前
深度学习试题及答案解析(一)
人工智能·深度学习
四口鲸鱼爱吃盐24 分钟前
Pytorch | 利用VMI-FGSM针对CIFAR10上的ResNet分类器进行对抗攻击
人工智能·pytorch·python
hanglove_lucky25 分钟前
本地摄像头视频流在html中打开
前端·后端·html
四口鲸鱼爱吃盐32 分钟前
Pytorch | 利用PI-FGSM针对CIFAR10上的ResNet分类器进行对抗攻击
人工智能·pytorch·python
边缘计算社区1 小时前
吉快科技荣膺“金边奖·最佳大模型一体机”,引领AI边缘新时代
人工智能·科技
新智元1 小时前
LeCun 八年前神预言,大模型路线再颠覆?OpenAI 宣告:强化学习取得稳定性突破
人工智能·openai
电子海鸥1 小时前
迁移学习--fasttext概述
人工智能·机器学习·迁移学习
因_果_律1 小时前
亚马逊云科技 re:Invent 2024重磅发布!Amazon Bedrock Data Automation 预览版震撼登场
大数据·人工智能·科技·亚马逊云科技·re invent
新智元1 小时前
李飞飞谢赛宁:多模态 LLM「空间大脑」觉醒,惊现世界模型雏形!
人工智能·llm