PaperlayerView exact revision on arXiv

Source-linked AI summary

All are Worth Words: A ViT Backbone for Diffusion Models

Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, Jun Zhu

arXiv:2209.12152v4cs.CVcs.AIcs.LG

TL;DR

Diffusion models commonly rely on CNN-based U-Nets, motivating whether their design is necessary for diffusion image generation. The paper introduces U-ViT, which processes time, conditions, and noisy image patches as tokens while using long skip connections. U-ViT matches or exceeds similarly sized CNN-based U-Nets across several generation tasks, including reported FID scores of 2.29 on ImageNet and 5.48 on MS-COCO.

  • Problem

    Diffusion models remain dominated by CNN-based U-Nets, leaving it unclear whether their reliance on down/up-sampling architecture is necessary for diffusion image modeling.

  • Method

    U-ViT parameterizes the diffusion noise-prediction network by treating time, condition, and noisy image patches as tokens and connecting shallow and deep layers with long skips.

  • Results

    U-ViT is comparable if not superior to similarly sized CNN-based U-Nets across unconditional, class-conditional, and text-to-image generation, with FID scores of 2.29 on ImageNet 256×256 and 5.48 on MS-COCO.

  • Takeaways & Limitations

    The results suggest long skip connections are crucial for diffusion-based image modeling, while CNN-based down/up-sampling operators are not always necessary.

  • Takeaways & Limitations

    U-ViT can also parameterize other prediction types, such as x0 prediction, but the supplied text does not establish comparative results for them.

Abstract

from arXiv · show

Vision transformers (ViT) have shown promise in various vision tasks while the U-Net based on a convolutional neural network (CNN) remains dominant in diffusion models. We design a simple and general ViT-based architecture (named U-ViT) for image generation with diffusion models. U-ViT is characterized by treating all inputs including the time, condition and noisy image patches as tokens and employing long skip connections between shallow and deep layers. We evaluate U-ViT in unconditional and class-conditional image generation, as well as text-to-image generation tasks, where U-ViT is comparable if not superior to a CNN-based U-Net of a similar size. In particular, latent diffusion models with U-ViT achieve record-breaking FID scores of 2.29 in class-conditional image generation on ImageNet 256x256, and 5.48 in text-to-image generation on MS-COCO, among methods without accessing large external datasets during the training of generative models. Our results suggest that, for diffusion-based image modeling, the long skip connection is crucial while the down-sampling and up-sampling operators in CNN-based U-Net are not always necessary. We believe that U-ViT can provide insights for future research on backbones in diffusion models and benefit generative modeling on large scale cross-modality datasets.

1. Introduction

U-ViT asks whether diffusion models need the CNN-based U-Net and proposes a ViT-based alternative that treats all inputs as tokens while retaining long skip connections. Across unconditional, class-conditional, and text-to-image generation, U-ViT is comparable or superior to similarly sized CNN-based U-Nets, with strong reported FID scores.

  • Diffusion models widely use CNN-based U-Nets with down-sampling blocks, up-sampling blocks, and long skip connections.
  • U-ViT treats time, condition, and noisy image patches as tokens and adds long skip connections between shallow and deep layers.An optional 3×3 convolutional block before output is included for visual quality.
  • U-ViT is comparable if not superior to a similarly sized CNN-based U-Net across unconditional, class-conditional, and text-to-image generation.
  • 2.29 FID was achieved on ImageNet 256×256 for class-conditional generation without large external training datasets.
  • 5.48 FID was achieved on MS-COCO for text-to-image generation without large external training datasets.
  • The results suggest long skip connections are crucial, whereas CNN U-Net down/up-sampling operators are not always necessary.

2. Background

Diffusion models learn to reverse a gradual noise-injection process, typically through CNN-based U-Nets. The background frames the remaining question of whether ViT-style token processing can serve diffusion-based image modeling.

  • Diffusion models gradually inject noise into data and reverse that process to generate data from noise.The forward noise-injection process is formalized as a Markov chain.
  • The forward transition is Gaussian, with α_t and β_t defining the noise schedule and satisfying α_t + β_t = 1.
  • A Gaussian reverse model approximates the ground-truth reverse transition, with its mean learned through noise prediction.
  • Conditional diffusion models feed class or text condition information into the noise-prediction objective.The condition may be represented directly or through a continuous embedding.
  • Prior image-modeling work relies heavily on CNN-based U-Nets, which combine down-sampling, up-sampling, long skips, and conditioning mechanisms.
  • ViT represents images as sequences of flattened patches with learnable one-dimensional position embeddings, but its suitability for diffusion image modeling remains unclear.

3. Method

U-ViT is a ViT-based diffusion backbone that tokenizes time, conditions, and noisy image patches, while adding long skip connections and selected convolutional components. Ablations and scaling studies identify effective implementation choices and show that performance depends on depth, width, and especially small patch sizes.

  • Architecture: U-ViT treats time, condition, and noisy image patches as tokens while parameterizing the diffusion model’s noise prediction network.The noisy image is split into patches before tokenization.
  • Architecture: Long skip connections link shallow and deep layers to provide shortcuts for low-level features in the pixel-level prediction task.The paper states that these shortcuts ease training of the noise prediction network.
  • Implementation ablations: Concatenating the main and long-skip branches followed by linear projection performs best among the tested branch-combination choices.Direct addition provides no benefit, while other projected combinations improve performance over dropping the long skip.
  • Implementation ablations: Treating time as a token outperforms adaptive layer normalization, while a 3×3 convolution after output projection performs slightly better than the alternatives.The convolutional comparison also includes placing the block before projection and removing it.
  • Implementation ablations: The original linear patch embedding and 1-dimensional learnable position embedding outperform their tested alternatives; removing position embeddings prevents meaningful image generation.The position study indicates that positional information is critical for image generation.
  • Scaling properties: Performance improves from 9 to 13 layers, from width 256 to 512, and from patch size 8 to 2, but gains stop at depth 17, width 768, or patch size 1 in the reported studies.Because small patches are costly at high resolution, the paper models low-dimensional latent representations with U-ViT.

4. Related Work

Prior diffusion systems largely rely on CNN-based U-Nets, while ViTs had shown promise elsewhere without an established role in diffusion image modeling. U-ViT is positioned against these backbones and broader diffusion-model improvements, with comparable or superior performance to similarly sized U-Nets.

  • Transformers in diffusion models: ViTs represent images as sequences of flattened patches with learnable one-dimensional position embeddings, but their suitability for diffusion image modeling was previously unclear.U-ViT addresses this unresolved application of ViTs.
  • U-Net in diffusion models: CNN-based U-Nets dominate diffusion image generation through down-sampling blocks, up-sampling blocks, and long skip connections.Prior work also added normalization, attention, residual blocks, and cross attention.
  • Transformers in diffusion models: U-ViT has comparable or superior performance to a CNN-based U-Net of similar size across the paper’s reported comparisons.This comparison is cited to Tables 1 and 4.
  • Other diffusion improvements: Related diffusion research also improves sampling speed, training methodology, and controllable generation in addition to changing the backbone.These directions are presented as complementary improvements to diffusion models.

5. Experiments

The experiments evaluate U-ViT across unconditional, class-conditional, and text-to-image generation, including latent diffusion settings and comparisons with CNN-based U-Nets. U-ViT is competitive or better, with particularly strong FID results on ImageNet 256×256 and MS-COCO.

  • Experimental setup: U-ViT is evaluated on unconditional, class-conditional, and text-to-image generation tasks using datasets including CIFAR10, CelebA, ImageNet, and MS-COCO.High-resolution ImageNet and MS-COCO experiments use latent representations.
  • Unconditional generation: U-ViT is comparable to U-Net on unconditional CIFAR10 and CelebA 64×64 and performs much better than GenViT.The comparison reports FID on 50K generated samples.
  • Class-conditional generation: FID 5.85 with U-ViT-M on ImageNet 64×64 improves over IDDPM’s 6.92, while U-ViT-L further improves FID to 4.26.U-ViT-M has 131M parameters and U-ViT-L has 287M parameters.
  • Class-conditional generation: FID 2.29 is achieved on class-conditional ImageNet 256×256, outperforming prior diffusion models and a similarly sized U-Net in latent diffusion.The comparison also reports better performance than LDM under different sampling-step counts.
  • Ablations and broader evaluation: Long skip connections are critical on ImageNet, while U-ViT remains comparable to state-of-the-art diffusion models on other metrics with comparable or smaller GFLOPs.Reported metrics include sFID, inception score, precision, and recall.
  • Text-to-image generation: FID 5.48 is achieved on MS-COCO by increasing U-ViT-S depth from 13 to 17 layers, with generated samples showing improved quality and text-semantic matching over U-Net.The example describes U-ViT generating a ball and, in the deeper model, a bat where U-Net generated neither.

6. Conclusion

The conclusion presents U-ViT as a general ViT backbone for diffusion-based image generation. Across multiple generation tasks, it is comparable or superior to similarly sized CNN-based U-Nets, while long skip connections remain important and down/up-sampling is not always necessary.

  • 6. Conclusion: U-ViT treats time, condition, and noisy image patches as tokens and connects shallow and deep layers with long skip connections.This is the architecture’s central design.
  • 6. Conclusion: Across unconditional, class-conditional, and text-to-image generation, U-ViT is comparable if not superior to a similarly sized CNN-based U-Net.The conclusion summarizes the paper’s experimental scope and comparison.
  • 6. Conclusion: The results suggest that long skip connections are crucial, whereas CNN-based U-Net down-sampling and up-sampling operators are not always necessary.The conclusion frames this as a supported implication for diffusion-based image modeling.

A. Experimental Setup

The appendix records experimental configurations, optimization settings, inference timing, and the CNN-based U-Net baseline used for comparison. These details cover schedules, training durations, sampling, and model scale.

  • Experimental configuration: The experimental setup table summarizes decoder weights, noise schedules, samplers, and classifier-free-guidance settings.It defines abbreviations including VP, SD, EM, DPM-Solver, and p_uncond.
  • Training: Learning rates, weight decay, AdamW coefficients, iteration counts, batch sizes, and classifier-free guidance are specified for each dataset.Most datasets use a learning rate of 2e-4, with ImageNet 64×64 using 3e-4.
  • Inference: Generating 500 samples with DPM-Solver on one A100 takes about 19, 34, 59, and 89 seconds for U-ViT-S, M, L, and H, respectively.Classifier-free guidance doubles the reported time.
  • Baseline: The baseline U-Net operates on 32×32 latent representations, downsamples to 16×16, 8×8, and 4×4, and has 53M parameters versus 45M for U-ViT-S.Self- and cross-attention are applied at 16×16 and 8×8 resolutions.

C. Results of Other Metrics and Configurations on ImageNet

Additional ImageNet evaluations report several quality metrics and computational cost across U-ViT configurations. U-ViT remains comparable to state-of-the-art diffusion models while using comparable or fewer GFLOPs.

  • Additional metrics: ImageNet evaluation reports FID, sFID, inception score, precision, recall, parameter counts, and GFLOPs.The results include multiple U-ViT configurations.
  • Comparison: U-ViT is comparable to state-of-the-art U-Net-based diffusion models on these metrics and has comparable or smaller GFLOPs.This conclusion covers both quality metrics and computational cost.

D. CKA Analysis

The CKA analysis compares hidden-representation similarity across transformer blocks for concatenation, addition, and no long skip connections. Concatenation produces substantially less similarity between neighboring blocks than the other settings.

  • CKA compares hidden representations after each transformer block across three long-skip combination strategies.The strategies are concatenation, addition, and no long skip connection.
  • Concatenation yields low similarity between neighboring transformer blocks, unlike addition and no long skip connection.The addition and no-skip settings show similar neighboring blocks, including blocks 0–3 and 6–11 or blocks 0–5 and 6–11.

E. Compare with U-Net Under Similar Amount of Parameters and Computational Cost

Under similar parameter counts and computational costs, U-ViT consistently outperforms the ADM-based U-Net on ImageNet 256×256. The advantage holds both without classifier-free guidance and with guidance scale 0.4.

  • U-ViT consistently outperforms a similarly sized U-Net across training iterations without classifier-free guidance.The comparison uses ImageNet 256×256 and models with 501M versus 646M parameters and 133 versus 135 GFLOPs.
  • 6.58 versus 10.69 FID: U-ViT outperforms U-Net at 500K iterations without classifier-free guidance.FID is evaluated with 50K samples.
  • 2.29 versus 2.66 FID: U-ViT outperforms U-Net with a classifier-free guidance scale of 0.4.Both models are evaluated at 500K training iterations on ImageNet 256×256.

F. Additional Samples

The paper presents additional generated samples for ImageNet at 512×512 and 256×256, including class-conditioned and random samples, plus random-prompt samples from MS-COCO.

  • ImageNet 512×512 samples are shown for class conditions including goldfish, arctic fox, monarch butterfly, elephant, flamingo, and tennis ball.
  • Additional ImageNet 512×512 class-conditioned samples include cheeseburger, fountain, balloon, tabby cat, lorikeet, and agaric.
  • ImageNet 256×256 includes both class-conditioned samples spanning twelve labels and random samples.
  • The displayed MS-COCO prompts describe people, objects, indoor scenes, animals, and outdoor settings.
  • MS-COCO examples are random samples generated from prompts drawn from the validation set.
Loading 2209.12152v4…