PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffiT: Diffusion Vision Transformers for Image Generation

Ali Hatamizadeh, Jiaming Song, Guilin Liu, Jan Kautz, Arash Vahdat

arXiv:2312.02139v3cs.CVcs.AIcs.LG

TL;DR

The paper studies how Vision Transformers can serve as effective denoising networks for diffusion-based image generation. It introduces DiffiT with time-dependent self-attention and reports SOTA generation quality across latent and image spaces, including a 1.73 ImageNet-256 FID with fewer parameters than competing Transformer diffusion models.

  • Problem

    The paper studies whether Vision Transformers can effectively support diffusion-based image generation while providing fine-grained control over the denoising process.

  • Method

    DiffiT uses a ViT-based diffusion model with Time-dependent Multihead Self-Attention that adapts query, key, and value projections across denoising time.

  • Results

    DiffiT achieves SOTA performance across latent- and image-space generation tasks, including a 1.73 FID on ImageNet-256 and fewer parameters than MDT and DiT.

  • Takeaways & Limitations

    Time-dependent self-attention provides a common DiffiT design for high-quality generation with improved parameter efficiency across multiple synthesis settings.

Abstract

from arXiv · show

Diffusion models with their powerful expressivity and high sample quality have achieved State-Of-The-Art (SOTA) performance in the generative domain. The pioneering Vision Transformer (ViT) has also demonstrated strong modeling capabilities and scalability, especially for recognition tasks. In this paper, we study the effectiveness of ViTs in diffusion-based generative learning and propose a new model denoted as Diffusion Vision Transformers (DiffiT). Specifically, we propose a methodology for finegrained control of the denoising process and introduce the Time-dependant Multihead Self Attention (TMSA) mechanism. DiffiT is surprisingly effective in generating high-fidelity images with significantly better parameter efficiency. We also propose latent and image space DiffiT models and show SOTA performance on a variety of class-conditional and unconditional synthesis tasks at different resolutions. The Latent DiffiT model achieves a new SOTA FID score of 1.73 on ImageNet256 dataset while having 19.85%, 16.88% less parameters than other Transformer-based diffusion models such as MDT and DiT,respectively. Code: https://github.com/NVlabs/DiffiT

1 Introduction

DiffiT applies Vision Transformers to diffusion image generation by dynamically adapting attention across denoising time while modeling spatial and temporal dependencies. It reports strong generation quality, parameter efficiency, and SOTA results across latent- and image-space tasks.

  • 1 Introduction: TMSA adapts self-attention across denoising stages by integrating time-dependent query, key, and value weights with spatial tokens.This design models spatial and temporal dependencies and their interaction during sampling.
  • 1 Introduction: Attention maps with TMSA progressively localize toward salient details during denoising, whereas models without TMSA do not recover those details.The comparison is qualitative and uses CIFAR10 sampling trajectories.
  • 1 Introduction: Compared with AdaLN, TMSA improves parameter efficiency by learning three temporal components per block rather than six conditioning components.The comparison concerns query, key, and value temporal components versus AdaLN shift, scale, and gate parameters for self-attention and MLP.
  • 1 Introduction: Window-based TMSA reduces self-attention computation by shortening the token sequence through local regions without cross-communication.The paper describes this windowed design as effective for reducing computational cost.
  • 1 Introduction: DiffiT achieves SOTA performance across conditional and unconditional generation tasks in latent and image space, including a new ImageNet-256 FID of 1.73.The latent model uses 19.85% and 16.88% fewer parameters than MDT and DiT, respectively.
  • 1 Introduction: DiffiT uses a ViT-based denoising model for both latent- and image-space image generation.The contribution unifies the model design across these generation settings.

2 Related Work

Related work places DiffiT within two lines of research: diffusion models for broad generative tasks and Transformers for visual generation. Prior work includes convolutional diffusion networks, Transformer-based generators, and architectures combining latent and image features.

  • 2 Related Work: Diffusion models have achieved strong results in generative learning, motivating architectures that address their denoising-network design.The paper notes that diffusion denoising networks have multiple proposed architecture variants.
  • 2 Related Work: Diffusion models generate samples through iterative denoising and have been applied across image, language, speech, 3D, time-series, molecular, and security domains.The related work describes their use across multiple generative learning settings.
  • 2 Related Work: Transformer-based visual generators include pure Transformer GANs, bipartite latent-image architectures, and Linformer-based models for higher-resolution synthesis.These examples illustrate several Transformer designs for visual generation.

3 Methodology

Diffusion models learn to reverse progressive noising, while DiffiT adapts attention to denoising time and spatial structure through TMSA. The architecture supports latent- and image-space generation using Transformer blocks, U-shaped processing, and local attention.

  • Diffusion Model Preliminaries: Diffusion models train a denoising network to predict noise added to corrupted samples, approximating the reverse process from an approximately Gaussian distribution toward the data distribution.Sampling can use ODE solvers when stochastic noise is absent or SDE solvers otherwise.
  • Time-dependent Self-Attention: DiffiT’s time-dependent attention forms queries, keys, and values from both spatial tokens and time embeddings, allowing attention behavior to adapt across denoising steps.The resulting attention uses time-dependent Q, K, and V within the Transformer block.
  • DiffiT Transformer Block: DiffiT’s Transformer block combines time-dependent multi-head self-attention with LayerNorm, an MLP, residual connections, and a time-embedding token.The block receives spatial tokens arranged on a two-dimensional grid and a token representing the sampling time.
  • Latent Space: The latent DiffiT model encodes images with a pretrained variational autoencoder, patchifies latent feature maps, and applies a ViT denoiser without upsampling or downsampling layers.It also uses three-channel classifier-free guidance and a final linear layer to decode the output.
  • DiffiT Architecture: The image-space model uses a symmetrical U-shaped encoder-decoder with skip connections, convolutional resolution changes, and DiffiT blocks at each resolution.Local window attention reduces the cost for large token grids, while bottleneck layers permit information sharing across windows.

4 Results

Latent DiffiT achieves strong image-generation performance across ImageNet resolutions, with a new ImageNet-256 FID result and diverse generated images. Image-space DiffiT also reaches state-of-the-art FID performance on CIFAR-10.

  • 1.73 FID: latent DiffiT sets a new ImageNet-256 state-of-the-art result against competing approaches.
  • Latent DiffiT significantly outperforms DiT on ImageNet-512 in both FID and Inception Score.
  • Latent DiffiT generates diverse, high-quality images across classes on ImageNet-256 and ImageNet-512.
  • DiffiT achieves a state-of-the-art 1.95 FID on CIFAR-10, outperforming EDM, LSGM, U-ViT, and GenViT.

5 Ablation

The ablations show that DiffiT benefits from balanced token dimensions, its multi-resolution encoder-decoder design, direct time-dependent attention, and appropriately sized attention windows. These choices improve FID and computational efficiency across image-generation settings.

  • 5.1 Time and Feature Token Dimensions: Increasing feature-token dimensions from 256 to 512 and decreasing time-embedding dimensions from 512 to 256 both degrade performance.
  • 5.2 Effect of Architecture Design: FID improves from 4.64 to 3.71, 2.27, and 1.95 as multi-resolution decoding and DiffiT Transformer blocks are added.The progression evaluates configurations B through E on CIFAR-10.
  • 5.3 Effectiveness of TMSA: TMSA lowers FID by 0.28 in VE and 0.25 in VP DDPM++ settings on CIFAR-10.
  • 5.4 Impact of Self-Attention Components: Directly integrating time tokens into query, key, and value formation outperforms adding them to positional bias or the MLP.Removing the temporal component raises FID to 3.97 from 1.95, while the MLP alternative reaches 3.81.
  • 5.8 Comparison of TMSA and DiT Modulation: TMSA performs better than DiT modulation across training iterations on ImageNet-256.
  • 5.10 Effect of Window Size: Increasing the TMSA window from 2 to 4 decreases FID by 23.23% on CIFAR-10 and 12.17% on FFHQ-64.Further increasing the window from 4 to 8 yields only marginal improvements of 1.53% and 3.60%.

6 Conclusion

DiffiT is a ViT-based diffusion model for latent- and image-space generation. Its TMSA mechanism dynamically adapts self-attention during denoising while improving parameter efficiency and achieving state-of-the-art ImageNet-256 performance.

  • DiffiT uses TMSA to dynamically adapt self-attention across denoising stages while modeling spatial and temporal dependencies.
  • DiffiT achieves state-of-the-art ImageNet-256 performance with fewer parameters than competitive Transformer-based diffusion models.

G.1 Comparison to DiT and LDM

DiffiT incorporates time through time-dependent self-attention rather than AdaLN-style shift and scale or time-embedding concatenation. It reports state-of-the-art results across resolutions and outperforms prior models on ImageNet benchmarks.

  • DiffiT jointly learns spatial and temporal dependencies through TMSA instead of AdaLN shift-scale modulation or time-embedding concatenation.
  • DiffiT improves ImageNet-256 FID over LDM and DiT by 31.26% and 51.94%, respectively, and ImageNet-512 FID over DiT by 13.85%.

H Architecture

DiffiT uses encoder–decoder components whose architecture is specified by resolution and dataset. The latent model uses a deep ViT-style configuration with 30 layers and 1152-dimensional hidden states.

  • H Architecture: For FFHQ-64, the architecture uses four resolution stages with four blocks at each stage.For CIFAR-10, it instead uses three resolution levels: 32, 16, and 8.
  • H Architecture: The encoder and decoder component layouts are documented for DiffiT models trained at 64×64 resolution.These specifications appear in Tables S.2 and S.3.
  • H Architecture: The latent DiffiT configuration uses 30 layers, hidden size 1152, 16 attention heads, and an MLP ratio of 4.The model is described as comparable to DiT-XL/2-G, which uses patch size 2.

I Implementation Details

Experiments follow established EDM and ADM configurations across unconditional image-space and ImageNet latent-space settings. Sampling uses dataset-specific EDM steps or a 250-step DDPM procedure for ImageNet evaluation.

  • I Implementation Details: CIFAR-10 and FFHQ-64 experiments follow EDM training configurations and data augmentation strategies in an unconditional setting.Both use 200,000 Adam iterations, PyTorch, and eight NVIDIA A100 GPUs.
  • I Implementation Details: CIFAR-10 uses 18 deterministic EDM sampling steps, whereas FFHQ-64 uses 40 steps.Their training images are 32×32 and 64×64, respectively.
  • I Implementation Details: ImageNet-256 and ImageNet-512 use learning rates of 3 × 10^-4 and 1 × 10^-4, respectively, with batch sizes of 256 and 512.Both use EMA decay 0.9999, ADM diffusion hyperparameters, and a 250-step DDPM sampler for FID-50K evaluation.

J Qualitative Results

Qualitative visualizations cover CIFAR-10, FFHQ-64, and ImageNet at 256×256 and 512×512 resolutions. The authors report that DiffiT captures fine-grained details and produces high-fidelity images across these datasets.

  • J Qualitative Results: DiffiT visualizations span CIFAR-10 and FFHQ-64 image-space generation tasks.These samples are shown in Figures S.2 and S.3.
  • J Qualitative Results: Latent DiffiT visualizations include randomly sampled 512×512 ImageNet images.The samples are presented in Figures S.4–S.7.
  • J Qualitative Results: The authors report that DiffiT captures fine-grained details and produces high-fidelity images across the evaluated datasets.ImageNet-256 samples are shown in Figures S.8–S.10.
Loading 2312.02139v3…