皆可为词:一种用于扩散模型的 ViT 骨干网络
计算机视觉与模式识别
2023-03-28 v4 人工智能
机器学习
摘要
视觉Transformer(ViT)已在多种视觉任务中展现出潜力,而基于卷积神经网络(CNN)的 U-Net 在扩散模型中仍占主导地位。我们设计了一种简单且通用的基于 ViT 的架构(命名为 U-ViT),用于扩散模型的图像生成。U-ViT 的特点是将所有输入(包括时间、条件和含噪图像块)视为 token,并在浅层与深层之间采用长跳跃连接。我们在无条件与类条件图像生成以及文本到图像生成任务上评估了 U-ViT,其表现与类似规模的基于 CNN 的 U-Net 相当甚至更优。特别地,采用 U-ViT 的隐扩散模型在 ImageNet 256x256 类条件图像生成中取得了 2.29 的破纪录 FID 分数,在 MS-COCO 文本到图像生成中取得了 5.48 的分数,上述成绩均来自训练生成模型时未使用大型外部数据集的方法。我们的结果表明,对于基于扩散的图像建模,长跳跃连接至关重要,而基于 CNN 的 U-Net 中的下采样与上采样算子并非总是必要。我们相信 U-ViT 能为未来扩散模型骨干网络的研究提供见解,并有益于大规模跨模态数据集上的生成建模。
引用
@article{arxiv.2209.12152,
title = {All are Worth Words: A ViT Backbone for Diffusion Models},
author = {Fan Bao and Shen Nie and Kaiwen Xue and Yue Cao and Chongxuan Li and Hang Su and Jun Zhu},
journal= {arXiv preprint arXiv:2209.12152},
year = {2023}
}
备注
Accepted to CVPR 2023