Generative Modeling by Minimizing the Wasserstein-2 Loss (2406.13619v2)

Abstract: This paper approaches the unsupervised learning problem by minimizing the second-order Wasserstein loss (the $W_2$ loss) through a distribution-dependent ordinary differential equation (ODE), whose dynamics involves the Kantorovich potential associated with the true data distribution and a current estimate of it. A main result shows that the time-marginal laws of the ODE form a gradient flow for the $W_2$ loss, which converges exponentially to the true data distribution. An Euler scheme for the ODE is proposed and it is shown to recover the gradient flow for the $W_2$ loss in the limit. An algorithm is designed by following the scheme and applying persistent training, which naturally fits our gradient-flow approach. In both low- and high-dimensional experiments, our algorithm outperforms Wasserstein generative adversarial networks by increasing the level of persistent training appropriately.

• The paper presents a novel generative modeling framework by formulating a distribution-dependent ODE whose time-marginal law converges exponentially to the true data distribution.
• It proposes an Euler-discretized algorithm (W2-FE) that outperforms traditional Wasserstein GANs in low- and high-dimensional experiments.
• The study leverages optimal transport theory and persistent training techniques to establish a robust foundation for advanced unsupervised learning research.

Generative Modeling by Minimizing the Wasserstein-2 Loss: An In-depth Analysis

Introduction

The paper explores a novel approach to the unsupervised learning problem by focusing on minimizing the second-order Wasserstein loss ($W_2$ loss). This method is posited as a more efficient alternative to existing generative modeling techniques, particularly those leveraging different forms of generative adversarial networks (GANs). The main contribution of the paper is the introduction of a distribution-dependent ordinary differential equation (ODE) whose dynamics involve the Kantorovich potential, enabling the time-marginal law of the ODE to converge exponentially to the true data distribution.

Main Results

Distribution-dependent ODE and Fokker-Planck Equation

The cornerstone of the paper's approach is the formulation of a distribution-dependent ODE:

$dY_t = -\nabla \phi_{\mu^{Y_t}}(Y_t) \, dt,$

where $\phi_{\mu^{Y_t}}$ is the Kantorovich potential between the current estimated distribution $\mu^{Y_t}$ and the true data distribution $\mu^*$. A significant theoretical result is the proof that the time-marginal law of this ODE converges exponentially to $\mu^*$. This is achieved by establishing that the ODE has a unique solution, constructed from the associated nonlinear Fokker-Planck equation.

Euler Scheme and Algorithm Design

An Euler scheme is proposed to discretize the ODE, and it is shown to correctly approximate the gradient flow for the $W_2$ loss in the limit. An algorithm, termed W2-FE (Wasserstein-2 Forward Euler), is designed based on this scheme, leveraging persistent training to enhance performance. The algorithm demonstrates superior convergence speed and performance compared to Wasserstein GANs (WGANs) in both low- and high-dimensional experiments.

Theoretical Framework

The paper builds upon the rich theoretical foundation of optimal transport, particularly the properties of the $W_2$ distance. Here are the essential theoretical components utilized:

1. Subdifferential Calculus in $\mathcal{P}_2(\mathbb{R}^d)$:
   • The subdifferential of the function $J(\mu) = \frac{1}{2} W_2^2(\mu, \mu^*)$ is always well-defined and can be explicitly characterized for measures absolutely continuous with respect to the Lebesgue measure.
2. Gradient Flows and Exponential Convergence:
   • The gradient flow of $J$ exists and converges exponentially to $\mu^*$ under the $W_2$ distance. The paper constructs a gradient flow as a time-rescaled geodesic between the initial estimate and $\mu^*$.
3. Fokker-Planck Equation and Solution Construction:
   • The solution to the nonlinear Fokker-Planck equation associated with the gradient-descent ODE is constructed, ensuring that the time-marginal laws of the solution coincide with the true gradient flow.

Numerical Results

Low-Dimensional Experiments

In a two-dimensional setting, the W2-FE algorithm is evaluated on a task of learning a ring-shaped mixture of Gaussians from an initial Gaussian distribution. The algorithm outperforms the refined WGAN algorithm (W1-LP), demonstrating faster convergence and superior performance when suitable levels of persistent training are employed.

High-Dimensional Experiments

For domain adaptation from the USPS dataset to the MNIST dataset, the W2-FE algorithm is assessed using a 1-nearest neighbor (1-NN) classifier accuracy metric. The results showcase that the W2-FE algorithm (with higher levels of persistent training) converges significantly faster and achieves higher accuracy compared to the W1-LP algorithm.

Implications and Future Directions

The practical implications of this research are manifold. By leveraging the Wasserstein-2 distance and persistent training techniques, the proposed algorithm can achieve state-of-the-art performance in generative modeling tasks, both in low and high dimensions. This opens avenues for further exploration, including:

1. Extension to Other Optimal Transport Distances:
   • Investigating the use of higher-order Wasserstein distances or alternative optimal transport metrics to potentially enhance generative modeling frameworks.
2. Scaling to Even Higher Dimensions:
   • Testing the robustness and scalability of the W2-FE algorithm in more complex datasets and higher-dimensional spaces.
3. Persistent Training Paradigms:
   • Further examining the effects of varying persistent training levels across different neural network architectures and datasets to optimize training efficiency and model performance.

Conclusion

The paper presents a robust theoretical and practical framework for generative modeling by minimizing the Wasserstein-2 loss. Through detailed analytical results and comprehensive numerical experiments, it demonstrates the efficacy of the proposed approach over existing methods, particularly in terms of convergence speed and modeling accuracy. This work not only advances the state of generative modeling but also provides a solid foundation for future research and development in the field of unsupervised learning.