Make U-Nets Great Again！北大&华为强强联合

科技 2024-11-19 14:00 广东

来源：机器之心

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-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers
论文地址：https://arxiv.org/pdf/2405.02730
GitHub 地址：https://github.com/YuchuanTian/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 接收，更多内容，请参考原论文。

http://mp.weixin.qq.com/s?__biz=MzU0MDQ1NjAzNg==&mid=2247587221&idx=2&sn=dfe93e1d82791ec260bb5519bdd315b1

机器学习算法那些事

号主是大厂人工智能专家，专注于机器学习，深度学习以及计算机视觉等研究方向，每天会更新人工智能最前沿知识和分享自己的论文总结和学习笔记，让你系统化的学习每个知识点，每天进步一点点。

最新文章

深入浅出Batch Size，对模型训练的影响分析

TPAMI 2024 | 扩散模型中的测量指导：来自医学图像合成的见解

特朗普上台，第一刀再次扎在了留学生身上。。。

Make U-Nets Great Again！北大&华为强强联合

电子科大提出注意力机制创新，改进 MobileViT变体在早期降采样阶段的注意力 QKV 操作研究！

MIT教授，何恺明大神最新课程资料下载！

基于有效样本的类别不平衡损失CB-Loss

又一985，落地深圳！

清华大学提出 TANet 用于综合恶劣天气图像复原的三重注意网络！

首次明确！国务院: 博士学位论文，不再非写不可

TPAMI 2024 | 烘焙神经辐射场以实现实时视图合成

今年顶会这情况。。。大家提前做准备吧！

使用 Transformers 进行概率时间序列预测实战

深度学习常用损失函数总览：基本形式、原理、特点

李飞飞团队新作：空间智能版ImageNet来了！

TPAMI 2024 | 延迟感知的统一动态网络用于高效图像识别

50K*16薪，进字节了！

三种Transformer模型中的注意力机制介绍及Pytorch实现：从自注意力到因果自注意力

浅谈频率学派和贝叶斯学派

教授何恺明在MIT的第二门课！

谷歌&MIT何恺明团队：视觉大模型像LLM一样高效Scaling，指路连续token+随机生成顺序

TPAMI 2024 | 不确定性增强的鲁棒视频活动预测

transformer的细节到底是怎么样的？Transformer 连环18问！

1700 页，6 大主题！AI 学习路线手册发布！

降维和特征选择的对比介绍

北大韦东奕上课照片走红，板书潇洒，新发型吸睛，网友：这要打脸多少“只会翻PPT”的老师？

TPAMI 2024 | 什么样的地方会成为异常之地？

姜萍造假实锤！阿里数赛出结果。。。

如何理解attention中的Q,K,V？

TPAMI 2024 | 迁移标注者和实例依赖的转移矩阵以实现从群体学习

网传北师大教授与女下属建立“永久亲密关系”，承诺对方读博和自由列支经费？北师大回应

TPAMI 2024 | 通过寻找相关子空间对神经网络预测进行解耦解释

算法岗和开发岗有什么区别？

新增9所985！这份名单，再扩容！

TPAMI 2024 | 多模态图像融合的通用空间-频率学习框架

热议！博士一毕业，直接回老家县城大专任教，事业编、副教授待遇，外加几十万安家费…

三种Transformer模型中的注意力机制介绍及Pytorch实现：从自注意力到因果自注意力

导师：自己每天科研工作近10小时，都觉得不够。研究生们，每天工作不够5小时，拿什么去竞争？

TPAMI 2024 | 风格化学习：跨任务和领域的持续语义分割

仅32所高校！教育部公布重磅名单

梯度累积的隐藏陷阱：Transformer库中梯度累积机制的缺陷与修正

甘露寺2024年招聘公告！环境清幽，包吃包住，午休3小时。。。

TPAMI 2024 | 桥接动作：在照片之间生成3D姿态和形状

浙大高飞老师：读博士有多苦？

TPAMI 2024 | 边缘感知自监督全景深度估计Transformer与球面几何

比亚迪开奖，只要双211以上？！

TPAMI 2024 | 局部非刚性运动恢复结构的封闭形式、成对解决方案

又一院士被“除名”！院士终生制，没了？

欧式距离、曼哈顿距离、切比雪夫距离三种距离的可视化展示

NeurIPS'24｜大模型LoRA超级变体！仅需参数0.4%达微调效率巅峰！

分类

时事

民生

政务

教育

文化

科技

财富

体娱

健康

情感

旅行

百科

职场

楼市

企业

乐活

学术

汽车

时尚

创业

美食

幽默

美体

文摘

原创标签

时事社会财经军事教育体育科技汽车科学房产搞笑综艺明星音乐动漫游戏时尚健康旅游美食生活摄影宠物职场育儿情感小说曲艺文化历史三农文学娱乐电影视频图片新闻宗教电视剧纪录片广告创意壁纸头像心灵鸡汤星座命理教育培训艺术文化金融财经健康医疗美妆时尚餐饮美食母婴育儿社会新闻工业农业时事政治星座占卜幽默笑话独立短篇连载作品文化历史科技互联网

发布位置

广东北京山东江苏河南浙江山西福建河北上海四川陕西湖南安徽湖北内蒙古江西云南广西甘肃辽宁黑龙江贵州新疆重庆吉林天津海南青海宁夏西藏香港澳门台湾美国加拿大澳大利亚日本新加坡英国西班牙新西兰韩国泰国法国德国意大利缅甸菲律宾马来西亚越南荷兰柬埔寨俄罗斯巴西智利卢森堡芬兰瑞典比利时瑞士土耳其斐济挪威朝鲜尼日利亚阿根廷匈牙利爱尔兰印度老挝葡萄牙乌克兰印度尼西亚哈萨克斯坦塔吉克斯坦希腊南非蒙古奥地利肯尼亚加纳丹麦津巴布韦埃及坦桑尼亚捷克阿联酋安哥拉