All are Worth Words: a ViT Backbone for Score-based Diffusion Models

Fan Bao, Chongxuan Li, Yue Cao, Jun Zhu

Keywords: score-based model, diffusion model, vision transformer

TL;DR: Propose a vision transformer backbone for score-based diffusion models

Abstract: Vision transformers (ViT) have shown promise in various vision tasks including low-level ones while the U-Net remains dominant in score-based diffusion models. In this paper, we perform a systematical empirical study on the ViT-based architectures in diffusion models. Our results suggest that adding extra long skip connections (like the U-Net) to ViT is crucial to diffusion models. The new ViT architecture, together with other improvements, is referred to as U-ViT. On several popular visual datasets, U-ViT achieves competitive generation results to SOTA U-Net while requiring comparable amount of parameters and computation if not less.

Student Paper: Yes

Community Implementations: [ 2 code implementations](https://www.catalyzex.com/paper/arxiv:2209.12152/code)