PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving GANs Using Optimal Transport

Tim Salimans, Han Zhang, Alec Radford, Dimitris Metaxas

arXiv:1803.05573v1cs.LGstat.ML

TL;DR

Generative modeling requires matching generated and data distributions, while primal optimal transport can suffer from biased mini-batch gradients and fixed costs can be insufficiently discriminative. OT-GAN addresses this with Mini-batch Energy Distance and an adversarially learned transport cost, achieving stable large-batch training and state-of-the-art benchmark results.

  • Problem

    Generative models must match generated and data distributions, but primal optimal transport can yield biased mini-batch gradients and fixed transport costs may be insufficiently discriminative.

  • Method

    OT-GAN combines primal optimal transport with energy distance over mini-batches and learns the transport cost adversarially in a neural latent space.

  • Results

    OT-GAN is uniquely stable with large mini-batches and achieves state-of-the-art results on several common benchmarks, including an 8.47 ± .12 score on CIFAR-10.

  • Takeaways & Limitations

    The mini-batch energy distance provides a highly discriminative training objective with unbiased mini-batch gradients, while learned cost functions can capture image similarity and data geometry.

  • Takeaways & Limitations

    OT-GAN requires substantial computation and memory, with best results requiring very large mini-batches and experiments generally using 8 GPUs for several days.

Abstract

from arXiv · show

We present Optimal Transport GAN (OT-GAN), a variant of generative adversarial nets minimizing a new metric measuring the distance between the generator distribution and the data distribution. This metric, which we call mini-batch energy distance, combines optimal transport in primal form with an energy distance defined in an adversarially learned feature space, resulting in a highly discriminative distance function with unbiased mini-batch gradients. Experimentally we show OT-GAN to be highly stable when trained with large mini-batches, and we present state-of-the-art results on several popular benchmark problems for image generation.

1 INTRODUCTION

Generative modeling seeks to match generated and training-data distributions, while GAN critics provide an optimizable distributional distance. OT-GAN combines primal optimal transport with adversarially learned energy distance to obtain discriminative, unbiased mini-batch training.

  • Generative modeling learns models that generate images, audio, video, text, and other data, with applications including compression and representation learning.
  • GANs use a neural-network discriminator or critic to distinguish generated data from training data and define an optimizable distance between their distributions.
  • Primal optimal transport offers closed-form solutions and tractable training objectives, but mini-batch use can produce biased gradients.
  • OT-GAN incorporates primal optimal transport into the GAN critic through Mini-batch Energy Distance, combining transport with energy distance in an adversarially learned feature space.
  • The paper develops the metric theoretically, applies it to generative modeling, evaluates it on benchmarks, and discusses strengths, weaknesses, and future work.

2 GANS AND OPTIMAL TRANSPORT

Optimal transport provides principled distributional objectives for generative modeling, but its primal and dual forms pose different practical challenges. Prior work addresses these challenges with neural critics, Sinkhorn approximations, and energy distance, motivating OT-GAN’s synthesis.

  • Earth-Mover distance measures the minimum transport cost needed to transform the generator distribution into the data distribution and is a metric for suitable costs.
  • The primal optimal-transport minimization is generally intractable, motivating a dual formulation over 1-Lipschitz functions approximated by constrained neural critics.
  • GAN critics only approximately solve the dual transport problem because critic optimization is imperfect and obtainable critics roughly match 1-Lipschitz functions.
  • Sinkhorn AutoDiff makes entropically smoothed primal transport tractable on GPUs by minimizing over soft matchings in mini-batches.
  • Mini-batch Sinkhorn distance has biased fixed-size mini-batch gradients and is not a valid distributional metric, motivating Energy Distance and OT-GAN’s synthesis of these ideas.

3 MINI-BATCH ENERGY DISTANCE

The Mini-batch Energy Distance extends energy distance to mini-batch distributions and uses Sinkhorn distance as its within-batch metric. Its repulsive generator term preserves unbiased gradients while retaining primal transport’s discriminative and stabilizing benefits.

  • Mini-batch Energy Distance measures distributions over mini-batches rather than individual images, pursuing the greater expressive power of mini-batch distributions.
  • Generalized energy distance remains valid for random objects such as mini-batches and is a metric when its underlying distance is a metric.
  • The method chooses entropy-regularized Wasserstein, or Sinkhorn, distance as the mini-batch metric because it is valid between individual mini-batches.
  • The final metric uses independently sampled data and generator mini-batches and learns its transport cost function adversarially.
  • Adding a repulsive generator term to the attractive transport term makes mini-batch gradients unbiased and the objective statistically consistent.
  • Using primal Sinkhorn transport leads to stronger discriminative power and more stable generative modeling.

4 OPTIMAL TRANSPORT GAN (OT-GAN)

OT-GAN learns an adversarial transport cost in a feature space to make mini-batch energy distance more discriminative. Training alternates generator and critic updates, uses Sinkhorn soft matchings, and extends to conditional image generation.

  • Adversarial transport cost: OT-GAN learns the transport cost adversarially because fixed costs can be statistically inefficient in high dimensions.With simple fixed costs, many poor generator distributions may appear close to the data distribution without enormous sample sizes.
  • Adversarial transport cost: The learned cost uses cosine distance between deep-network embeddings of images in an adversarially learned latent space.The feature network maps mini-batch images into vectors, and its parameters maximize the resulting mini-batch energy distance.
  • Optimization: Training alternates gradient updates, with the generator updated more often than the critic to prevent the learned cost from becoming degenerate.This update schedule discourages zero transport costs between non-identical image regions that the generator could exploit.
  • Optimization: Sinkhorn computes the matching matrix for each transport term, while OT-GAN does not backpropagate through the matching algorithm.The paper justifies ignoring matching gradients using the envelope theorem because the matching minimizes the transport objective.
  • Conditional generation: The training procedure generalizes to conditional image generation by conditioning the generator on side information sampled from the training data.The remaining algorithm is unchanged apart from replacing batches with data-side-information pairs.
  • Optimization: Algorithm 1 combines four real-generated transport terms with within-real and within-generated correction terms across independent mini-batches.The procedure samples X, X′ from real data and Y, Y′ from generated samples, then alternates critic and generator updates.

5 EXPERIMENTS

Experiments evaluate OT-GAN’s stability, consistency, and image quality across toy, CIFAR-10, ImageNet Dog, and conditional bird-generation tasks. OT-GAN remains consistent with fixed critics, benefits from larger mini-batches, and achieves strong inception-score results against baselines.

  • Mixture of Gaussian Dataset: With a fixed critic on 8 Gaussian mixtures, OT-GAN remains consistent, whereas the original GAN loss collapses to 3 of 8 modes.The experiment tests whether the generator diverges after critic updates stop.
  • CIFAR-10: On CIFAR-10, larger mini-batches improve training stability and increase the inception score.The mini-batch energy distance becomes more discriminative as larger batches cover more modes and provide lower-variance estimates.
  • CIFAR-10: 8.47 ± .12 is OT-GAN’s CIFAR-10 inception score, outperforming all baseline models.The model is trained with a batch size of 8000 and without labels.
  • ImageNet Dogs: On 128×128 ImageNet Dog generation, OT-GAN achieves an inception score of 8.97±0.09 versus DCGAN’s 8.19±0.11.The reported samples contain less nonsensical images, while the tuned DCGAN variant still suffers from mode collapse.
  • Conditional Generation of Birds: On conditional text-to-image generation for the CUB test set, OT-GAN achieves the best inception score among the compared state-of-the-art methods.The experiment uses conditional OT-GAN with batch size 2048.

6 DISCUSSION

OT-GAN combines a novel mini-batch energy distance with an adversarially learned transport cost. It is stable with large mini-batches and achieves state-of-the-art benchmark results, but requires substantial computation and memory.

  • 6 DISCUSSION: OT-GAN trains its generator by minimizing mini-batch energy distance, combining primal optimal transport with adversarially learned energy distance.The resulting metric is highly discriminative and has unbiased mini-batch gradients.
  • 6 DISCUSSION: Large mini-batches make OT-GAN uniquely stable during training and support state-of-the-art results on several common benchmarks.
  • 6 DISCUSSION: OT-GAN requires large amounts of computation and memory, with best results requiring very large mini-batches and lengthy multi-GPU training.Except for the mixture-of-Gaussians toy example, experiments used 8 GPUs for several days.
  • 6 DISCUSSION: The mini-batch energy distance remains a valid objective after critic training stops, enabling more frequent generator updates than critic updates.This procedure learns a relatively stable transport cost function and an image-embedding function capturing image geometry.

A CONDITIONAL GENERATION

The conditional OT-GAN algorithm alternates generator and critic updates using independent real and generated mini-batches with shared side information. Its loss combines cross-distribution and within-distribution optimal transport costs.

  • A CONDITIONAL GENERATION: Conditional OT-GAN training uses minibatch SGD with a configurable number of generator iterations per critic iteration.
  • A CONDITIONAL GENERATION: Each iteration samples two independent real mini-batches and two independent generator mini-batches, reusing the same side information.
  • A CONDITIONAL GENERATION: The loss adds cross real-generated transport costs and subtracts real-real and generated-generated transport costs.The algorithm updates critic parameters on critic iterations and generator parameters otherwise.

B CIFAR-10 ARCHITECTURE AND TRAINING DETAILS

For CIFAR-10, OT-GAN uses convolutional generator and critic networks loosely based on DCGAN, with normalization and data-dependent initialization. The generator maps 100-dimensional uniform latent codes to 32 × 32 color images.

  • B CIFAR-10 ARCHITECTURE AND TRAINING DETAILS: The generator and critic are convolutional networks whose architectures are loosely based on DCGAN with various modifications.
  • B CIFAR-10 ARCHITECTURE AND TRAINING DETAILS: Weight normalization and data-dependent initialization are used for both the generator and critic.
  • B CIFAR-10 ARCHITECTURE AND TRAINING DETAILS: The generator maps 100-dimensional uniform latent codes sampled between -1 and 1 to 32 × 32 color images.
  • B CIFAR-10 ARCHITECTURE AND TRAINING DETAILS: Tables 3 and 4 specify the CIFAR-10 generator and critic architectures.

C ADVERSARIALLY LEARNING THE TRANSPORT COST FUNCTION

The CIFAR-10 ablation tests a fixed cosine transport cost in the original image-pixel space instead of adversarially learning the cost function. This substantially lowers the reported inception score.

  • C ADVERSARIALLY LEARNING THE TRANSPORT COST FUNCTION: The ablation uses cosine distance in the original feature space, with the transport cost fixed while other OT-GAN settings remain unchanged.The distance is defined using the original image pixel values.
  • C ADVERSARIALLY LEARNING THE TRANSPORT COST FUNCTION: 4.93 was the highest inception score with a fixed cosine cost, versus 8.47 when the cost function was learned adversarially.
  • C ADVERSARIALLY LEARNING THE TRANSPORT COST FUNCTION: Figure 8 presents CIFAR-10 samples generated without adversarially learning the cost function.

D MODEL COLLAPSE AND SAMPLE DIVERSITY

On ImageNet dog data, DCGAN develops persistent mode collapse during prolonged training, whereas OT-GAN preserves sample diversity across substantially more epochs.

  • DCGAN begins showing mode collapse after about 900 epochs, and the model does not recover with continued training.
  • OT-GAN shows no mode collapse or reduction in sample diversity after 13,000 training epochs.
  • Figure 9 contrasts ImageNet dog samples from DCGAN after 900 epochs with OT-GAN samples after 13,000 epochs.
Loading 1803.05573v1…