PaperlayerView exact revision on arXiv

Source-linked AI summary

ViTGAN: Training GANs with Vision Transformers

Kwonjoon Lee, Huiwen Chang, Lu Jiang, Han Zhang, Zhuowen Tu, Ce Liu

arXiv:2107.04589v2cs.CVcs.LGeess.IV

TL;DR

The paper asks whether Vision Transformers can support image generation in GANs with quality comparable to CNN-based GANs, despite instability from high-variance discriminator gradients and inadequate conventional regularization. It introduces stabilization and architectural modifications for ViT-based GANs, and ViTGAN achieves comparable performance to leading CNN-based models across CIFAR-10, CelebA, and LSUN bedroom.

  • Problem

    The paper investigates whether Vision Transformers can generate images in GANs with quality comparable to CNN-based GANs, while noting instability from high-variance discriminator gradients and ineffective conventional regularization.

  • Method

    ViTGAN stabilizes ViT-based GANs through improved discriminator spectral normalization, generator layer-normalization and output-mapping modifications, and architectural studies of latent and pixel mappings.

  • Results

    ViTGAN achieves comparable performance to leading CNN-based GANs on CIFAR-10, CelebA, and LSUN bedroom, and can outperform StyleGAN2 when combining its discriminator with the ViTGAN generator.

  • Takeaways & Limitations

    The experiments establish a concrete ViT-based GAN baseline whose stabilization techniques support competitive image generation and facilitate further research on Transformer-only GAN architectures.

  • Takeaways & Limitations

    The paper does not claim superiority over highly optimized StyleGAN2 + ADA models equipped with sophisticated data augmentation and tuned configurations.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper asks whether convolution-free Vision Transformers can support GAN-based image generation at quality comparable to CNN-based GANs. It identifies instability in ViT GAN training and proposes modifications that stabilize convergence, achieving comparable benchmark performance.

  • ViTGAN investigates whether Vision Transformers can train GANs for image generation without convolution or pooling.
  • ViT-based GAN training becomes highly unstable because discriminator training develops high-variance gradients, while conventional regularization does not resolve the issue.
  • The paper modifies discriminator spectral normalization and generator layer normalization and output mapping to stabilize training and facilitate convergence.
  • Ablation experiments validate the necessity of the proposed techniques for stable and superior image generation.
  • On CIFAR-10, CelebA, and LSUN bedroom, ViTGAN achieves performance comparable to StyleGAN2 under the same training setting.
  • The paper does not claim superiority over highly optimized StyleGAN2 + ADA, instead targeting the performance gap between CNN-based GANs and vanilla ViT-based GANs.

2 RELATED WORK

The related work situates ViTGAN among GAN stabilization methods, Vision Transformers, and Transformer-based image generators. Unlike autoregressive or cross-modal approaches, ViTGAN uses Vision Transformers within adversarial training.

  • Generative Adversarial Networks: GANs model target distributions through adversarial minmax optimization using divergences or integral probability metrics such as Wasserstein distance.
  • Generative Adversarial Networks: Common GAN stabilization methods include spectral normalization, gradient penalties, consistency regularization, and data augmentation, primarily developed for CNNs.
  • Vision Transformers: Vision Transformers are convolution-free Transformers that classify images as sequences of patches and are among the first Transformer architectures applied to GAN-based image generation here.
  • Generative Transformers in Vision: Other Transformer vision-generation methods formulate image synthesis as autoregressive sequence learning or cross-modal image-text learning.
  • Generative Transformers in Vision: ViTGAN instead trains Vision Transformers in the generative adversarial paradigm and proposes stability techniques within the original ViT backbone.
  • ViTGAN: The proposed framework uses ViT-based generator and discriminator components, with discriminator scores from the classification embedding and pixels generated patch-by-patch.

3 PRELIMINARIES: VISION TRANSFORMERS (VITS)

Vision Transformers represent images as sequences of patches and process them with Transformer components. Their sequence begins with a classification token and positional embeddings, followed by repeated self-attention and feed-forward transformations.

  • A Vision Transformer flattens a 2D image into a raster-scanned sequence of image patches.
  • Each patch has spatial size P × P and channel dimension C, while the sequence length is determined by the number of patches.
  • ViT prepends a learnable classification embedding to the patch sequence and adds one-dimensional positional embeddings.
  • The displayed sequence notation includes the patch embeddings that form the Transformer input.
  • Each Transformer layer applies layer normalization, multi-headed self-attention, and a residual connection to the preceding representation.
  • Self-attention forms query, key, and value representations from the input using learnable matrices, then aggregates multiple heads through concatenation and linear projection.

4 METHOD

ViTGAN stabilizes ViT-based GAN training through discriminator regularization and a redesigned generator. The generator combines self-modulated Transformer blocks with an output mapping from token embeddings to image patches.

  • Discriminator design: ViTGAN regularizes its discriminator with improved spectral normalization to strengthen Lipschitz continuity and stabilize training.The update multiplies each normalized layer weight by its spectral norm at initialization.
  • Discriminator design: Overlapping image patches reduce sensitivity to arbitrary grid partitions while retaining the same sequence length.The design can also provide stronger locality cues between neighboring patches.
  • Discriminator design: Convolutional projections apply 3×3 convolutions when computing Q, K, and V, combining local context with Transformer global context.The original projection remains recoverable with an identity convolution kernel.
  • Generator design: The generator maps Gaussian noise z through an MLP to latent vector w, then uses positional embeddings, latent conditioning, and Transformer blocks to produce token embeddings.Baseline variants either add w to every positional embedding or prepend w to the input sequence.
  • Generator design: Self-modulated layer normalization uses the latent vector to control adaptive normalization parameters γ_l and β_l within the generator.The proposed generator includes a Transformer block and an output mapping layer.
  • Generator design: An implicit neural representation maps each patch embedding y_i to patch pixels x_i using Fourier-encoded spatial locations and a two-layer MLP.The generated patches are reshaped into a complete image; this representation is reported as particularly helpful for training.

5 EXPERIMENTS

Experiments evaluate ViTGAN across standard image-generation benchmarks, compare it with Transformer and CNN-based GANs, and test its architectural and regularization choices through ablations. ViTGAN stabilizes Transformer-based GAN training and achieves competitive or improved image-generation quality.

  • Experimental Setup: ViTGAN is evaluated on CIFAR-10, CelebA, and LSUN bedroom across resolutions from 32×32 to 256×256.The experiments use consistent settings across datasets and resolutions, varying key architectural parameters such as Transformer-block count and patch size.
  • Main Results: ViTGAN outperforms other Transformer-based GANs by a large margin and converges to lower or comparable FID across datasets.The method overcomes gradient-magnitude spikes, improving training stability and resulting image quality.
  • Main Results: ViTGAN achieves comparable performance to leading CNN-based models including BigGAN and StyleGAN2, with comparable qualitative image quality and diversity.The comparison uses a generic StyleGAN2 configuration in the main table, while more comprehensive comparisons appear in the appendix.
  • Compatibility with CNN-based GANs: Mixing ViTGAN and StyleGAN2 components shows that the proposed techniques work with both Transformer-based and CNN-based generators and discriminators.The ViTGAN discriminator improves performance with a StyleGAN2 generator, and the ViTGAN generator also works with a StyleGAN2 discriminator.
  • Ablation Studies: Ablations identify implicit neural representation, L2+ISN regularization, overlapping patches, adaptive token connection weights, and convolutional projection as beneficial design choices.L2+ISN improves stability and scores over R1 and standard spectral normalization; convolutional projection combines local and global context and outperforms the alternatives tested.

6 CONCLUSION AND LIMITATIONS

ViTGAN applies Vision Transformers throughout a GAN and adds techniques for stable training and convergence. Experiments establish a baseline up to 256×256, while higher-resolution performance remains unverified.

  • 6 CONCLUSION AND LIMITATIONS: ViTGAN leverages Vision Transformers in GANs and introduces techniques intended to stabilize training and improve convergence.The architecture consists solely of Transformer layers.
  • 6 CONCLUSION AND LIMITATIONS: Experiments on standard benchmarks show performance comparable to leading CNN-based GANs.
  • 6 CONCLUSION AND LIMITATIONS: ViTGAN provides a concrete Vision Transformer baseline for image synthesis at resolutions up to 256×256.
  • 6 CONCLUSION AND LIMITATIONS: The framework could incorporate advanced training methods or coarse-to-fine architectures.
  • 6 CONCLUSION AND LIMITATIONS: The paper does not verify ViTGAN on high-resolution images beyond its baseline scope.

ETHICS STATEMENT

The authors acknowledge misuse risks for image-generation techniques and unresolved ethical and fairness questions, especially for identifiable-face generation.

  • ETHICS STATEMENT: Image-generation techniques can be misused to produce misleading information.
  • ETHICS STATEMENT: Future work should consider ethical and fairness implications when generating identifiable faces.
  • ETHICS STATEMENT: The paper leaves unclear whether the architecture affects groups differently by appearance, age, or physical disability.
  • ETHICS STATEMENT: The authors report conducting no subject studies or data collection for this project.

REPRODUCIBILITY STATEMENT

The reproducibility materials describe public-benchmark experiments, shared implementation details, and controlled comparisons using consistent augmentation and training settings.

  • REPRODUCIBILITY STATEMENT: All experiments use public image-generation benchmarks, with implementation details provided in the experiment section and appendices.
  • REPRODUCIBILITY STATEMENT: The authors plan to release code and models to support reproducibility.
  • REPRODUCIBILITY STATEMENT: Table 5 compares convolution-based BigGAN and StyleGAN2 with Transformer-based ViTGAN, with and without data augmentation.
  • REPRODUCIBILITY STATEMENT: DiffAug mitigates Transformer-discriminator overfitting by increasing training-data diversity, enabling ViTGAN to perform comparably to BigGAN and StyleGAN2.
  • REPRODUCIBILITY STATEMENT: The comparisons use StyleGAN2 with DiffAug to isolate architecture differences under the same training setting.
  • REPRODUCIBILITY STATEMENT: Ablations indicate adaptive generator weights matter more in the generator than discriminator, while locality benefits both components.

A.4 HIGH-RESOLUTION SAMPLES

ViTGAN is evaluated on uncurated LSUN Bedroom samples at 128×128 and 256×256 resolutions. At 256×256, it achieves FID 4.67 and outperforms the cited GANformer result.

  • A.4 HIGH-RESOLUTION SAMPLES: ViTGAN produces uncurated LSUN Bedroom samples at 128×128 and 256×256 resolutions.
  • A.4 HIGH-RESOLUTION SAMPLES: At 256×256 resolution, the generator uses sequence length 1024 while retaining the 128×128 architecture.
  • A.4 HIGH-RESOLUTION SAMPLES: 4.67 FID is achieved by ViTGAN on LSUN Bedroom 256×256.
  • A.4 HIGH-RESOLUTION SAMPLES: ViTGAN outperforms GANformer on LSUN Bedroom 256×256, with FID 4.67 versus 6.51.
  • A.4 HIGH-RESOLUTION SAMPLES: 2.48 FID is achieved by ViTGAN on LSUN Bedroom 128×128.

B COMPUTATION COST ANALYSIS

The paper argues that Transformer computation is manageable at standard image resolutions and focuses on GAN-specific challenges rather than inventing a new general-purpose Transformer block. Its Swin experiments remain limited by inferior FID and computational constraints.

  • Computation cost: A few hundred tokens typically keep Transformer sequence lengths within a region where quadratic attention complexity is manageable for standard image resolutions.The paper attributes this to tokenization strategies such as non-overlapping patches and strided convolutions.
  • Scope: The proposed techniques target challenges unique to combining Transformers and GANs rather than designing a new Transformer block for general computer vision.The authors state that the techniques remain valid for several ViT variants.
  • Swin limitation: Swin Transformer experiments produced inferior FID, while more extensive architectural tuning was left for future work because of limited computational resources.The authors identify numbers of heads, channels, and layers as possible tuning dimensions.
  • Comparison: Table 7 compares generator computation costs among methods.The supplied passage provides the table caption but no numerical entries.

C EXPERIMENT DETAILS

Experiments cover CIFAR-10, CelebA, and LSUN bedroom at multiple resolutions, using ViT-based generators and discriminators with specified block configurations and regularization. Additional design choices address discriminator overfitting and patch generation.

  • Datasets and resolutions: The experiments use 32×32 CIFAR-10, 64×64 CelebA and LSUN bedroom, and 128×128 LSUN bedroom settings.CIFAR-10 uses 50K training and 10K test images; LSUN bedroom has approximately 3 million training images and 300 validation images.
  • Model configurations: The 32×32 setting uses four-block ViT generators and discriminators, while the 64×64 setting increases both models to six blocks.Transformer blocks use feature dimension 384 and MLP hidden dimension 1,536 in the described settings.
  • Model configurations: At 128×128, the generator uses 8×8 patches and eight Transformer blocks, while the discriminator maintains sequence length 64 with four blocks.The discriminator reaches the desired sequence length using 3 × 3 convolutions with stride 2.
  • Regularization: Training applies Translation, Color, Cutout, and Scaling augmentations with probability 0.8, alongside balanced consistency regularization with λreal = λfake = 10.0.Dropout, weight decay, and Stochastic Depth are not used beyond bCR among the listed ViT regularization methods.
  • Patch extraction: Patch overlap extends each patch edge by o pixels, using effective patch size (P + 2o)×(P + 2o) with o = P/2 to mitigate discriminator overfitting.The authors distinguish this implementation from using a convolution operator.
  • Embeddings and patch generation: Positional embeddings use sine-activated linear projections of normalized patch positions, while patch generation uses Fourier-encoded pixel coordinates and a modulated two-layer MLP.The MLP is conditioned on patch embedding yi via weight modulation.
Loading 2107.04589v2…