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

Although ViT has shown promise in vision, CNN-based U-Nets remain dominant in diffusion models, leaving open whether they are necessary. U-ViT treats time, condition, and noisy image patches as tokens and connects shallow and deep layers with long skip connections. Across unconditional, class-conditional, and text-to-image generation, it is comparable to or better than similar-size CNN-based U-Nets.

  • Problem

    Whether diffusion models need CNN-based U-Nets remains an open question, despite ViT's promise in vision and the central role of backbones in diffusion models.

    Source: §1 · pp. 1–2

  • Method

    U-ViT treats time, condition, and noisy image patches as tokens and uses long skip connections between shallow and deep layers.

    Source: §1 · p. 2

  • Results

    Across unconditional, class-conditional, and text-to-image generation, U-ViT is comparable to or better than similar-size CNN-based U-Nets, with FID 2.29 on class-conditional ImageNet 256×256.

    Source: §1 · p. 2

  • Takeaways & Limitations

    For diffusion image modeling, long skip connections are crucial, while CNN-based U-Net down/up-sampling operators are not always necessary.

    Source: §1 · p. 2

    U-ViT can also parameterize other prediction types, such as x0 prediction.

    Source: §3 · p. 2

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 adapts a ViT for diffusion by treating time, condition, and noisy image patches as tokens and adding long skip connections. It matches or exceeds similar-sized CNN-based U-Nets across three image-generation tasks, with especially strong latent-diffusion FID scores.

1. Introduction

  • CNN-based U-Nets dominate diffusion image generation, motivating the question of whether their convolutional architecture is necessary despite ViTs’ promise in vision tasks.The paper frames backbone design as central to diffusion models, alongside the development of algorithms.Source: §1 · pp. 1–2
  • U-ViT treats time, condition, and noisy image patches as tokens, and uses long skip connections between shallow and deep layers.The authors motivate the skips as preserving low-level features useful for pixel-level prediction and easing training; an optional 3×3 convolutional block before output may improve visual quality.Source: Figure 1 · p. 1, §1 · p. 2
  • FID scores of 2.29 on ImageNet 256×256 and 5.48 on MS-COCO make U-ViT latent diffusion record-breaking among methods trained without large external datasets.Across unconditional, class-conditional, and text-to-image generation, U-ViT is comparable to or better than a similarly sized CNN-based U-Net.Source: §1 · p. 2
  • The results suggest long skip connections are crucial, while CNN-based U-Net down- and up-sampling operators are not always necessary for image diffusion models.The authors propose U-ViT as a source of insights for diffusion backbones and generative modeling on large-scale cross-modality datasets.Source: §1 · p. 2

2. Background

Diffusion models generate data by reversing a gradual noise-injection process, with learning framed as noise prediction and extended to conditional generation. The background contrasts the CNN-based U-Net commonly used for image diffusion with ViT’s token-based image representation.

2. Background

  • Diffusion models [25] add noise through a Markov-chain forward process and approximate its reverse transitions with Gaussian distributions to generate data from noise.The forward transition uses a noise schedule, while the reverse model approximates the ground-truth reverse transition.Source: §2 · p. 2
  • Learning is equivalent to predicting injected Gaussian noise, and conditional models feed class or text information into the noise-prediction objective.The condition c may be the condition itself or its continuous embedding.Source: §2 · p. 2
  • Prior image-diffusion work relies heavily on CNN-based U-Net, which combines down-sampling and up-sampling blocks with long skip connections.Conditions are incorporated through mechanisms such as adaptive group normalization [13] and cross attention.Source: §2 · p. 2
  • ViT [16] represents images as flattened-patch tokens with learnable 1D position embeddings, but its suitability for diffusion-based image modeling remains unclear.The patch embeddings are fed into a transformer encoder.Source: §2 · p. 2

3. Method

U-ViT treats time, condition, and noisy-image patches as tokens and uses long skip connections to carry low-level features. The section evaluates implementation choices and studies how depth, width, and patch size affect performance on CIFAR10.

3. Method

  • U-ViT tokenizes time, condition, and noisy-image patches to predict injected noise, and can also parameterize x0 prediction.The image is split into patches, following ViT design.Source: §3 · p. 2

3.1. Implementation Details

  • Concatenating the main and long-skip embeddings performs best; direct addition offers no benefit, while other projected skip combinations outperform removing the long skip.The study ablates these choices on CIFAR10, evaluating FID every 50K iterations on 10K generated samples for efficiency.Source: §3.1 · p. 3, Figure 2 · p. 3
  • Treating time as a token performs better than injecting it through adaptive layer normalization (AdaLN).AdaLN applies a time-conditioned scale and bias after transformer-block layer normalization.Source: §3.1 · p. 3
  • A 3×3 convolution after the patch-output projection performs slightly better than placing it before that projection or omitting it, and linear patch embedding beats convolutional embedding.The optional convolution is intended to prevent potential transformer image artifacts and improves sample visual quality in the reported experiments.Source: §3 · p. 3, §3.1 · p. 4
  • One-dimensional learnable position embeddings outperform two-dimensional sinusoidal embeddings; without position embeddings, the model fails to generate meaningful images.The authors conclude that positional information is critical for image generation.Source: §3.1 · p. 4

3.2. Effect of Depth, Width and Patch Size

  • On CIFAR10, performance improves as depth increases from 9 to 13 and width from 256 to 512, but gains stop at depth 17 and width 768; patch size 2 is needed for good performance.Reducing patch size from 8 to 2 improves performance, but further reduction to 1 brings no gain; the authors hypothesize that low-level noise prediction needs small patches, so high-resolution images are modeled in low-dimensional latent space.Source: §3.2 · p. 4

4. Related Work

The section contrasts U-ViT with transformer- and CNN-based diffusion backbones, and distinguishes backbone design from improvements to sampling, training, and controllable generation.

Transformers in diffusion models

  • U-ViT performs much better than GenViT, which lacks long skip connections and a 3×3 convolutional block and incorporates time before normalization layers.The passage attributes the improvement to careful implementation design.Source: §4 · p. 4
  • U-ViT achieves a better FID than VQ-Diffusion [20] (see Table 1 and Table 4), treating all inputs as tokens and using long skip connections between shallow and deep layers.Source: §4 · p. 4
  • Transformers in diffusion models also encode texts [48], decode texts [7] [29] [39] [46], and generate CLIP embeddings.Source: §4 · p. 4

U-Net in diffusion models

  • U-ViT has comparable or better performance than a similarly sized CNN-based U-Net, while using a conceptually simple ViT-based design.CNN-based U-Nets have been improved with group normalization [25], multi-head attention, improved residual blocks [13], and cross attention.Source: §4 · pp. 4–5

Improvements of diffusion models

5. Experiments

U-ViT matches U-Net on unconditional image generation, improves on prior diffusion models for class-conditional ImageNet, and achieves a 5.48 FID on text-to-image MS-COCO.

5.1. Experimental Setup

  • Experiments cover unconditional CIFAR10 and CelebA, class-conditional ImageNet, and MS-COCO text-to-image; high-resolution images are modeled in autoencoder latents, with CLIP text embeddings supplied as tokens.The setup uses ImageNet at 64×64, 256×256, and 512×512, and MS-COCO at 256×256.Source: §5.1 · p. 5

5.2. Unconditional and Class-Conditional Image Generation

  • U-ViT is comparable to U-Net on unconditional CIFAR10 and CelebA 64×64, while performing much better than GenViT.The comparison uses FID on 50K generated samples.Source: Table 1 · p. 6, §5.2 · p. 7
  • On ImageNet 64×64, U-ViT-M scores 5.85 FID versus IDDPM’s 6.92, and U-ViT-L improves the score to 4.26.The configurations have 131M and 287M parameters, respectively.Source: Table 1 · p. 6, §5.2 · p. 7
  • On class-conditional ImageNet 256×256, U-ViT achieves a state-of-the-art FID of 2.29, outperforming all prior diffusion models; it also outperforms LDM across sampling-step counts using the same sampler.U-ViT also outperforms VQ-Diffusion, a discrete diffusion model that uses a transformer backbone.Source: §5.2 · p. 7, Table 3 · p. 7
  • On ImageNet 512×512, U-ViT outperforms ADM-G, which directly models image pixels.The paper also reports selected samples from ImageNet 256×256 and 512×512 with good quality and clear semantics.Source: §5.2 · p. 7
  • Long skip connections are critical on ImageNet, extending the paper’s earlier finding on CIFAR10.Figure 5 ablates the long skip connection on ImageNet 256×256 without classifier-free guidance.Source: §5.2 · p. 7, Figure 5 · p. 7
  • U-ViT is comparable to state-of-the-art diffusion models on other reported metrics and has comparable or lower GFLOPs.The paper refers to additional metrics, including sFID, inception score, precision, and recall, in Appendix C.Source: §5.2 · p. 7

5.3. Text-to-Image Generation on MS-COCO

  • 5.48 FID on MS-COCO is achieved by U-ViT-S (Deep), a state-of-the-art result among methods trained without large external datasets.With the same random seed, U-ViT generated higher-quality samples that better matched text semantics than U-Net; the authors hypothesize that text and image tokens interact at every U-ViT layer.Source: §5.3 · p. 7

6. Conclusion

U-ViT is a ViT-based diffusion architecture that tokenizes time, condition, and noisy image patches and uses long skip connections. Across image-generation tasks, it is comparable to or better than a similar-size CNN-based U-Net, suggesting long skips matter while CNN U-Net sampling operators may not be necessary.

6. Conclusion

  • U-ViT treats time, condition, and noisy image patches as tokens and uses long skip connections between shallow and deep layers.The paper presents it as a simple, general ViT-based architecture for image generation with diffusion models.Source: §6 · pp. 7–8
  • Across unconditional, class-conditional, and text-to-image generation, U-ViT is comparable to or better than a similar-size CNN-based U-Net.The evaluation covers all three image-generation task types.Source: §6 · p. 8
  • The results suggest long skip connections are crucial for diffusion-based image modeling, while CNN U-Net down- and up-sampling operators are not always necessary.The authors propose that U-ViT may inform future diffusion backbones and benefit generative modeling on large-scale cross-modality datasets.Source: §6 · p. 8

A. Experimental Setup · B. Details of the U-Net Baseline on MS-COCO

The appendix summarizes U-ViT’s experimental setup and efficiency, then specifies an LDM U-Net baseline on MS-COCO with a comparable parameter count and matched training iterations and batch size.

A. Experimental Setup

  • U-ViT’s experimental setup is tabulated, with definitions for its decoder weights, noise schedules, samplers, classifier-free guidance probability, and output convolution option.The table identifies “ft-EMA” and “original” decoder weights, VP and SD schedules, and EM and DPM-Solver samplers.Source: Table 5 · p. 13, A. Experimental Setup · p. 13
  • Across datasets, U-ViT uses a learning rate of 2e-4, weight decay of 0.03, and AdamW coefficients (β1, β2) = (0.99, 0.99).Exceptions that could further improve performance are a 3e-4 learning rate on ImageNet 64×64, β2 = 0.999 on CIFAR10, and (β1, β2) = (0.9, 0.9) on MS-COCO.Source: A. Experimental Setup · p. 13
  • Gradient checkpointing reduces U-ViT-L/2 training memory on one A100 from 53GB to 10GB for batch size 128 on ImageNet 256×256.Training uses mixed precision for efficiency; training devices and durations are listed in Table 6.Source: A. Experimental Setup · p. 13, Table 6 · p. 14
  • 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 these times.Source: A. Experimental Setup · p. 13

B. Details of the U-Net Baseline on MS-COCO

  • The MS-COCO baseline is the LDM cross-attention U-Net, operating on 32×32 latents and downsampling to 16×16, 8×8, and 4×4 resolutions.It has 128 channels at 32×32 and 256 at other resolutions, two residual blocks per resolution, and self- and cross-attention at 16×16 and 8×8.Source: B. Details of the U-Net Baseline on… · p. 14
  • The baseline has 53M parameters, comparable to U-ViT-Small’s 45M, and uses the same training iterations and batch size as U-ViT for a fair comparison.It uses AdamW with weight decay 0.01 and coefficients (β1, β2) = (0.9, 0.999), matching LDM settings; its tuned learning rate is 2e-4.Source: B. Details of the U-Net Baseline on… · p. 14

C. Results of Other Metrics and Configurations on ImageNet

On ImageNet, U-ViT remains comparable to state-of-the-art U-Net-based diffusion models across FID, sFID, inception score, precision, and recall, with comparable or lower GFLOPs.

C. Results of Other Metrics and Configurations on ImageNet

  • U-ViT is comparable to state-of-the-art U-Net-based diffusion models on ImageNet FID, sFID, inception score, precision, and recall, with comparable or lower GFLOPs.Table 7 also reports parameter counts and GFLOPs; the supplied passages do not provide specific table values.Source: Table 7 · p. 14, C. Results of Other Metrics and… · p. 14

D. CKA Analysis

CKA analysis finds that concatenating long-skip branches changes transformer representations, while addition does not; addition and no-skip networks show similar representations across neighboring blocks.

D. CKA Analysis

  • The analysis compares concatenation, addition, and no long skip connection using CKA on hidden representations after each transformer block, with input time fixed at t = 0.5 on CIFAR10.The three skip-branch combination settings are Linear(Concat(hm, hs)), hm + hs, and no long skip connection.Source: D. CKA Analysis · p. 15
  • CKA is used to analyze similarity between hidden representations in deep neural networks [11] [35].Source: D. CKA Analysis · p. 15
  • Concatenation significantly changes transformer representations, whereas addition does not; addition and no-skip settings show similar representations across neighboring blocks.Figure 7 illustrates low similarity between neighboring blocks under concatenation and similar representations in neighboring blocks under addition and no long skip connection.Source: Figure 7 · p. 15, D. CKA Analysis · p. 15

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

On ImageNet 256×256, U-ViT outperforms a similarly computationally costly U-Net at 500K training iterations, both with and without classifier-free guidance.

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

  • At 500K iterations, U-ViT achieves FID 6.58 versus U-Net’s 10.69 without guidance, and FID 2.29 versus 2.66 with guidance scale 0.4.The models have 501M and 646M parameters, respectively, and use 133 and 135 GFLOPs; without guidance, U-ViT also consistently outperforms U-Net at different training iterations.Source: Figure 8 · p. 15, E. Compare with U-Net Under Similar… · p. 15

F. Additional Samples

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

F. Additional Samples

  • ImageNet 512×512 samples are shown conditioned on twelve classes, including goldfish, arctic fox, monarch butterfly, and cheeseburger.The figures also include class-conditioned samples of fountain, balloon, tabby cat, lorikeet, and agaric, alongside african elephant, flamingo, and tennis ball.Source: Figure 9 · p. 16, Figure 10 · p. 17
  • Random samples are also shown for ImageNet 512×512 and 256×256.The 512×512 random samples appear separately from the class-conditioned samples.Source: Figure 11 · p. 18, Figure 13 · p. 20
  • ImageNet 256×256 samples include twelve class-conditioned categories, from goldfish and arctic fox to lorikeet and agaric.A separate figure presents random samples at this resolution.Source: Figure 12 · p. 19, Figure 13 · p. 20
  • MS-COCO samples use prompts randomly drawn from the validation set.The figure presents random samples on MS-COCO.Source: Figure 14 · p. 21

Related papers

Loading 2209.12152v4…