Source-linked AI summary
On the Convergence and Robustness of Training GANs with Regularized Optimal Transport
Maziar Sanjabi, Jimmy Ba, Meisam Razaviyayn, Jason D. Lee
TL;DR
Wasserstein GAN optimization is difficult because its objective is non-convex, non-smooth, and computationally challenging. The paper uses regularized optimal transport to obtain a smooth objective and controlled approximate gradients, proving convergence to stationarity with approximate discriminator solutions. Experiments on MNIST and CIFAR-10 report computational efficiency and image quality comparable to state-of-the-art methods under matched architectures and computational power.
Problem
Wasserstein GAN optimization involves a non-convex, non-smooth, and difficult-to-compute objective, complicating first-order optimization and convergence analysis.
Method
The paper smooths the Wasserstein GAN objective with regularized optimal transport and uses approximate discriminator solutions to estimate generator gradients.
Results
The analysis proves convergence to a stationary solution, while experiments on MNIST and CIFAR-10 report computational efficiency and image quality comparable to state-of-the-art methods under matched settings.
Takeaways & Limitations
Regularized optimal transport provides a practical GAN objective that supports first-order optimization without requiring exact discriminator solutions.
Takeaways & Limitations
The analysis assumes a one-to-one generator mapping and access to a discriminator solver achieving predefined accuracy, while practical implementations use neural-network parameterizations.
Abstract
from arXiv · showhide
Generative Adversarial Networks (GANs) are one of the most practical methods for learning data distributions. A popular GAN formulation is based on the use of Wasserstein distance as a metric between probability distributions. Unfortunately, minimizing the Wasserstein distance between the data distribution and the generative model distribution is a computationally challenging problem as its objective is non-convex, non-smooth, and even hard to compute. In this work, we show that obtaining gradient information of the smoothed Wasserstein GAN formulation, which is based on regularized Optimal Transport (OT), is computationally effortless and hence one can apply first order optimization methods to minimize this objective. Consequently, we establish theoretical convergence guarantee to stationarity for a proposed class of GAN optimization algorithms. Unlike the original non-smooth formulation, our algorithm only requires solving the discriminator to approximate optimality. We apply our method to learning MNIST digits as well as CIFAR-10images. Our experiments show that our method is computationally efficient and generates images comparable to the state of the art algorithms given the same architecture and computational power.
1 Introduction
The paper addresses the computational and convergence difficulties of Wasserstein GAN optimization by smoothing the objective through regularized optimal transport. It establishes smoothness, controls gradient error from approximate discriminator solutions, and derives convergence to stationarity for SGD-based methods.
- The analysis allows any discriminator-solving method, provided it reaches a specified accuracy that determines closeness to stationarity.
- Regularized Wasserstein distance is smooth in the generator parameters, enabling first-order optimization of the smoothed GAN objective.
- Approximate solutions of the regularized discriminator problem provide controlled error in the generator’s stochastic gradients.
- Smoothness and approximate first-order information yield convergence of vanilla stochastic gradient descent to a stationary solution.
- Sinkhorn loss is used to avoid regularization bias while improving robustness across regularization strengths, and the resulting methods are intended to be computationally efficient and competitive.
2 Background
The background formulates generative modeling through optimal transport and its dual discriminator problem, then introduces regularization to relax hard constraints and enable parametric solution methods. The paper also states technical assumptions and the use of an approximate dual-solving oracle.
- Generative modeling seeks a parameterized mapping whose generated distribution minimizes optimal transport cost to the target distribution.
- Wasserstein GANs use the dual optimal-transport formulation, where the discriminator is constrained to be 1-Lipschitz when the cost is a distance.
- The analysis assumes the generator mapping is one-to-one and that a dual oracle solves each regularized discriminator problem to predefined accuracy, although practical implementations use neural networks.
- Regularized optimal transport adds a strongly convex regularizer to the transport problem and yields an equivalent generator gradient after subtracting the regularization term at the optimum.
- Regularization relaxes the hard dual constraint into a soft objective penalty, while the dual remains suitable for neural-network parameterization.
3 Smoothness of the generative objective
The regularized generator objective is smooth under stated cost and mapping conditions, unlike the original non-regularized objective. This smoothness also stabilizes the optimal transport plan as generator parameters change.
- Under Lipschitz cost variation and regularized-plan stability assumptions, hλ(θ)=d_c,λ(G_θ(q),p) is smooth in θ.
- Small changes in generator parameters produce small changes in the optimal transport plan, making subsequent discriminator optimization easier.
4 Solving the generator problem to stationarity using first order methods
Regularization makes the generator objective smooth and enables approximate discriminator solutions to provide controlled stochastic gradients for first-order optimization. The resulting SGD procedure converges to approximate stationarity, while Sinkhorn loss avoids bias across regularization choices but introduces computational trade-offs as λ decreases.
- 4 Solving the generator problem to stationarity using first order methods: The original non-regularized objective lacks smoothness and reliable sub-gradient approximations, whereas regularization supplies the smoothness needed for gradient-based analysis.The paper contrasts the non-regularized objective with hλ and derives gradient estimates by approximately solving the regularized dual problem.
- 4 Solving the generator problem to stationarity using first order methods: Approximate discriminator solutions yield controlled stochastic gradients for the smooth regularized objective, enabling convergence analysis for first-order GAN optimization.The dual solver provides approximate stochastic gradients, and the generator algorithm uses vanilla mini-batch SGD with a stationarity guarantee.
- 4 Solving the generator problem to stationarity using first order methods: The convergence guarantee depends on both generator updates and the quality of solving the discriminator problem, while average inner-problem error can remain sufficiently small.The analysis permits step-dependent approximation errors through an averaged error term.
- 4.1 Sinkhorn loss: a more robust generative objective: Reducing λ to avoid regularization bias destabilizes optimization by increasing discriminator sensitivity, tightening inner-solver accuracy, increasing smoothness constants, and slowing convergence.The paper reports that these effects worsen with problem complexity and scale.
- 4.1 Sinkhorn loss: a more robust generative objective: Sinkhorn loss remains zero at the true generator and makes that generator stationary for every λ > 0, while interpolating between 2d_c and MMD limits.For symmetric cost c, the true generator parameter θ0 is stationary and Lλ(Gθ0(p), q) = 0 for any λ > 0.
- 4.1 Sinkhorn loss: a more robust generative objective: For the Sinkhorn objective, only two terms depend on generator parameters, requiring two discriminator-oracle calls and producing overall gradient error bounded by 3δ.The corresponding SGD convergence guarantee follows by replacing the variance and gradient-error terms with the Sinkhorn-specific quantities.
5 Experiments
SWGAN experiments evaluate regularized-OT objectives and cost functions on MNIST and CIFAR-10. Latent representations improve MNIST results, while CIFAR-10 experiments show faster convergence under matched computational settings but lower quality for regularized-OT objectives.
- 5.1 Learning handwritten digits: SWGAN methods with latent representations perform better than methods using pixel-domain costs on MNIST.The authors attribute this to easier optimization in a lower-dimensional representation and occasional updates that reduce overfitting to a specific representation.
- 5.1 Learning handwritten digits: Sinkhorn loss outperforms the regularized-OT objective in the MNIST comparison, where the latter produces many faint images.The authors suggest the change of objective helps find a better stationary solution.
- 5.1 Learning handwritten digits: Compared with benchmark methods, SWGAN generates higher-quality MNIST images and qualitatively converges faster.The comparison includes methods from and [18], including a Sinkhorn-based regularized-OT method.
- 5.2 Generating tiny color images: On CIFAR-10, latent Sinkhorn-loss SWGAN methods converge faster than WGAN-GP, OT-GAN, and Sinkhorn-solver baselines under the same architecture and computational budget.OT-GAN is slower because it uses two batches for each label and requires more computation; the experiments use L1 and Cosine costs.
- 5.2 Generating tiny color images: Regularized-OT SWGAN objectives do not generate high-quality CIFAR-10 samples even with carefully tuned hyperparameters.Because of computational cost, the CIFAR-10 evaluation focuses on latent Sinkhorn-loss SWGAN with L1 and Cosine costs.
C Proof of Theorem 3.1
The proof establishes differentiability and smoothness of the regularized OT objective through the unique optimal transport plan and its stability under generator-parameter perturbations. It also supports approximate discriminator optimization through gradient-based and sample-based termination criteria.
- C Proof of Theorem 3.1: Strong convexity of the regularized transport problem yields a unique optimum, enabling Danskin-based differentiability of the objective.The proof compares optimal plans at two generator parameters and uses strong convexity to control their difference.
- C Proof of Theorem 3.1: The optimal regularized transport plan is stable under small changes in generator parameters, with a bound proportional to L0/λ.The proof concludes the stability inequality after applying the stated assumptions.
- C Proof of Theorem 3.1: The unconstrained dual formulation permits first-order optimization to obtain an ε-accurate discriminator solution.The paper notes that gradient descent can solve the finite-dimensional dual problem, while parametric neural networks can exploit relationships between empirical samples.
- C Proof of Theorem 3.1: Warm-starting the discriminator from the previous dual solution is justified because optimal dual variables change stably after small generator updates.This can reduce the difficulty of solving each successive dual transport problem.
- C Proof of Theorem 3.1: A sample-based measure of transport-plan marginal consistency provides a high-probability dynamic stopping criterion for discriminator optimization.The criterion is reliable when the sample size is sufficiently large and the transport weights are bounded.
D.1 Proof of Lemma 2.1
This section proves smoothness and transport-plan stability for regularized OT, including the 2-norm case, and discusses how regularization shapes the resulting GAN objective. The analysis also identifies a trade-off between smoothness and bias.
- D.1 Proof of Lemma 2.1: For 2-norm regularized OT, the objective hλ(θ) is Lipschitz smooth under the theorem’s boundedness assumptions.The discrete-uniform specialization sets pmax = 1/(MN).
- D.1 Proof of Lemma 2.1: The optimal transport plan varies continuously with parameters, with 2-norm stability bounded by ℓ0/λ.The paper states this bound for any two parameter values.
- D.1 Proof of Lemma 2.1: The 2-norm regularized surrogate is smoother than the KL-regularized surrogate when both use the same λ, because typically ℓ0 ≪ L0.This comparison follows the stability constants stated for the two regularizers.
- D.1 Proof of Lemma 2.1: Adding a strongly convex regularizer makes the otherwise non-convex and non-smooth linear-generator objective differentiable and smooth.The unregularized objective is non-convex because of a minimum over linear functions and non-smooth when the minimizer is non-unique.
- D.1 Proof of Lemma 2.1: Increasing λ can bias the obtained solution toward zero, and sufficiently large regularization can make the solution obsolete.The simple-case analysis describes convex behavior far from the optimum and non-convexity near optimal solutions.
G Proof of Theorem 4.1
The proof reformulates the 2-norm and KL-regularized OT dual problems with auxiliary transport variables. Feasibility and approximate optimality of the resulting triples yield bounds on transport-plan error.
- G Proof of Theorem 4.1: The proof augments dual variables with transport weights chosen optimally for each dual solution.This constructs feasible triples for both an ε-accurate solution and the exact optimum.
- G Proof of Theorem 4.1: Convexity of the feasible set allows comparison between the approximate and optimal triples through their connecting line segment.The argument uses optimality of the exact solution to control the directional derivative.
- G Proof of Theorem 4.1: The same construction is applied to KL-regularized OT after reparameterizing the transport plan with aij = πij/(MN).The proof then compares the resulting approximate and optimal variables using the corresponding regularized dual formulation.
- G Proof of Theorem 4.1: The proof obtains an L1 transport-plan error bound, written as ∥π − π∗∥1 ≤ δ.The bound follows from the preceding approximate-optimality argument and the regularization parameter.
H Proof of Theorem 4.2
The proof analyzes stochastic gradient updates using smoothness, approximate discriminator gradients, and bounded variance, yielding convergence behavior across iteration regimes.
- H Proof of Theorem 4.2: The proof combines smoothness, gradient-estimation error bounded by δ, and conditional variance bounded by σ2 to analyze expected gradient norms.The conditional expectation of the stochastic error is zero, while ∥∇t − Gt∥ ≤ δ and E[∥Gt − gt∥2|ξt] ≤ σ2.
- H Proof of Theorem 4.2: For sufficiently large T, choosing a constant step size yields the asymptotic convergence regime as T approaches infinity.This is the regime highlighted in the theorem because it gives an asymptotic rate for the expected gradient norm.
- H Proof of Theorem 4.2: When T is too small, the proof uses a separate constant-step-size choice and obtains a lower bound involving 2σ2 + δ2.The finite-iteration regime is treated separately from the asymptotic case.
I Stability of dual variables
The paper studies how optimal dual variables change with generator parameters and reports stability under regularization, while identifying deterioration as transport probabilities become small.
- I Stability of dual variables: Under KL regularization and finite supports, optimal dual variables can be related across generator parameters when transport-plan probabilities remain bounded away from zero.The theorem assumes π, π′ ≥ πmin and establishes a corresponding dual solution for the perturbed parameter.
- I Stability of dual variables: The stability problem becomes harder as πmin decreases, including when the distributions get closer or the regularization weight λ becomes very small.The paper explicitly identifies these conditions as worsening dual-variable stability.
- I Stability of dual variables: The proof obtains a nearby dual solution by treating the perturbed variables as an approximate linear system and projecting onto its exact solution space.Accurately solving the system gives an optimal dual solution for the perturbed transport plan, with the projection minimizing distance from the old variables.
- I Stability of dual variables: For uniform distributions, the stability bound worsens linearly with the numbers of support points, unlike the primal stability result.The dual formulation therefore becomes less stable as M and N increase, even when vector sizes are controlled.
- I Stability of dual variables: In the synthetic Gaussian-mixture experiment, SWGAN recovered all modes with d = 4 and remained robust with d = 2, where WGAN-GP quality deteriorated substantially.Figure 3 compares SWGAN after 10,000 iterations with WGAN-GP after 30,000 iterations; the reported runtimes were approximately 13 and 14 minutes, respectively.
L.1 Learning data-dependent cost function
The paper describes adversarially learned data-dependent cost functions in a latent representation and specifies training configurations for the Gaussian-mixture and MNIST experiments.
- L.1 Learning data-dependent cost function: A data-dependent cost function can be applied in a learned latent space whose representation is updated adversarially while held fixed for several iterations.The convergence analysis applies when the representation is fixed during the relevant optimization period.
- L.1 Learning data-dependent cost function: The adversarial cost-function training is intended to make the representation more informative for distinguishing real and generated samples.The paper describes this as promoting more meaningful representations.
- L.1 Learning data-dependent cost function: The Gaussian-mixture setup used two fully connected 128-neuron generator layers with tanh activations and a discriminator with 128-neuron hidden layers.The WGAN-GP discriminator used one fully connected discriminator with hidden layers of size 128.
- L.1 Learning data-dependent cost function: The implementation performed at most 20 discriminator iterations per generator iteration and stopped dynamically when the optimality condition was met.The hyper-parameter tables summarize the training configurations for the Gaussian-mixture and MNIST experiments.
M Training details for CIFAR
The CIFAR-10 training configuration uses specified generator and discriminator architectures, while regularized-OT training exhibits blur, instability, and eventual divergence as regularization is annealed.
- M Training details for CIFAR: Table 3 summarizes the hyper-parameters used for training CIFAR-10 images.The supplied table caption identifies its scope but does not provide individual hyper-parameter values.
- M Training details for CIFAR: Regularized-OT SWGAN produced blurry CIFAR-10 samples, became unstable during regularization annealing, and eventually diverged as discriminator accuracy was lost.Figure 4 also includes direct pixel-space learning with Sinkhorn loss and fixed regularization.
- M Training details for CIFAR: The CIFAR-10 generator uses a fully connected layer followed by batch normalization, GLU activations, upsampling, convolutional layers, and tanh output.The listed generator begins with FC(128, 16384) and ends with CONV(128, 3,stride=1)-Tanh.
- M Training details for CIFAR: The discriminator uses convolutional layers with batch normalization and CReLU, followed by normalization, a fully connected layer, and a scalar output layer.The architecture includes progressively wider convolutional layers before the final fully connected stages.