[Paper Review] ViTGAN: Training GANs with Vision Transformers

Kwonjoon Lee, Hui‐Wen Chang+4 more|arXiv (Cornell University)|Jul 9, 2021

Advanced Neural Network ApplicationsComputer Science60 references79 citations

TL;DR

ViTGAN integrates Vision Transformers into GANs, introducing stabilization techniques for ViT-based discriminators and generators, achieving competitive image generation quality to StyleGAN2 on standard benchmarks.

ABSTRACT

Recently, Vision Transformers (ViTs) have shown competitive performance on image recognition while requiring less vision-specific inductive biases. In this paper, we investigate if such performance can be extended to image generation. To this end, we integrate the ViT architecture into generative adversarial networks (GANs). For ViT discriminators, we observe that existing regularization methods for GANs interact poorly with self-attention, causing serious instability during training. To resolve this issue, we introduce several novel regularization techniques for training GANs with ViTs. For ViT generators, we examine architectural choices for latent and pixel mapping layers to facilitate convergence. Empirically, our approach, named ViTGAN, achieves comparable performance to the leading CNN-based GAN models on three datasets: CIFAR-10, CelebA, and LSUN bedroom.

Motivation & Objective

• Explore whether Vision Transformers can generate high-quality images without convolutional layers.
• Develop stabilization techniques for GAN training with ViT-based discriminators and generators.
• Demonstrate ViTGAN's performance on standard benchmarks and compare to CNN-based GANs.

Proposed method

• Adopt a ViT-based discriminator with enhanced Lipschitz regularization using improved L2 attention and augmented spectral normalization (ISN).
• Stabilize training with overlapping image patches and tailored regularization to address ViT-specific instability.
• Design a ViT-based generator featuring a two-part structure: transformer blocks with self-modulated layernorm and an implicit neural representation mapping (patch embedding to pixels).
• Incorporate self-modulated layernorm (SLN) controlled by latent z through an MLP for dynamic normalization.
• Use implicit neural representation to map patch embeddings to patch pixels, enhanced with Fourier features for smoothness.

Experimental results

Research questions

• RQ1 Can Vision Transformers be effectively used in GANs for image synthesis with stable training dynamics?
• RQ2 What regularization and architectural adjustments are necessary to stabilize ViT-based GAN training?
• RQ3 How does a ViT-based generator compare against CNN-based generators in image quality and diversity on standard benchmarks?

Key findings

• ViTGAN achieves FID scores of 6.66 (CIFAR-10), 3.74 (CelebA), and 2.65 (LSUN Bedroom) with IS scores of 9.30, 3.21, and 2.36 respectively, outperforming vanilla ViT baselines.
• ViTGAN surpasses a Transformer-based baseline by a large margin and is competitive with StyleGAN2 on all three benchmarks.
• Regularization combining L2 attention with improved spectral normalization (ISN) plus overlapping patches stabilizes training and reduces gradient spikes.
• Generator architecture with implicit neural representation for patch decoding improves GAN stability and image fidelity over linear mappings.
• ViTGAN is compatible with CNN-based discriminators/generators, maintaining strong performance when mixed with StyleGAN2 components.