[Paper Review] Sinkhorn Natural Gradient for Generative Models

Zebang Shen, Zhenfu Wang+2 more|arXiv (Cornell University)|Nov 9, 2020

Constraint Satisfaction and Optimization22 references4 citations

TL;DR

This paper proposes Sinkhorn Natural Gradient (SiNG), a novel optimization method for training generative models by performing steepest descent on the probability space using the Sinkhorn divergence as the metric. Unlike existing natural gradient methods, SiNG computes the Sinkhorn information matrix (SIM) in logarithmic time with respect to accuracy, enabling exact and efficient updates without requiring score functions, and provides a stable empirical estimator under Monte Carlo integration.

ABSTRACT

We consider the problem of minimizing a functional over a parametric family of probability measures, where the parameterization is characterized via a push-forward structure. An important application of this problem is in training generative adversarial networks. In this regard, we propose a novel Sinkhorn Natural Gradient (SiNG) algorithm which acts as a steepest descent method on the probability space endowed with the Sinkhorn divergence. We show that the Sinkhorn information matrix (SIM), a key component of SiNG, has an explicit expression and can be evaluated accurately in complexity that scales logarithmically with respect to the desired accuracy. This is in sharp contrast to existing natural gradient methods that can only be carried out approximately. Moreover, in practical applications when only Monte-Carlo type integration is available, we design an empirical estimator for SIM and provide the stability analysis. In our experiments, we quantitatively compare SiNG with state-of-the-art SGD-type solvers on generative tasks to demonstrate its efficiency and efficacy of our method.

Motivation & Objective

• To address the limitations of existing natural gradient methods in generative modeling, particularly the intractability of score function computation and high computational cost of Wasserstein-based alternatives.
• To develop a natural gradient method that operates directly on the probability space using the Sinkhorn divergence, ensuring reparameterization invariance.
• To derive an explicit, efficiently computable expression for the Sinkhorn information matrix (SIM) that scales logarithmically with target accuracy.
• To design a stable empirical estimator (eSIM) for SIM when only Monte Carlo samples are available, ensuring robustness in practical training.

Proposed method

• Proposes SiNG as a steepest descent method on the probability space equipped with the Sinkhorn divergence, derived as the limit of minimizing the functional within a Sinkhorn ball of radius ε.
• Derives the Sinkhorn information matrix (SIM) as the Hessian of the Sinkhorn divergence with respect to parameters θ, showing it can be computed using O(log(1/ε)) operations.
• Demonstrates that SIM evaluation requires only sampling from the generated measure αθ and does not require score function evaluation, unlike Fisher-Rao natural gradient.
• Introduces an empirical estimator (eSIM) of SIM based on the Hessian of the empirical Sinkhorn divergence, proven stable under Monte Carlo sampling.
• Employs the conjugate gradient method and PyTorch’s autograd to efficiently compute matrix-vector products involving eSIM without explicitly forming the matrix.
• Uses a differentiable Sinkhorn solver (e.g., via geomloss) to implicitly construct the Sinkhorn potentials and enable backpropagation through the computational graph.

Experimental results

Research questions

• RQ1 Can a natural gradient method be designed for generative models that avoids the need for score function computation while maintaining reparameterization invariance?
• RQ2 Is the Hessian of the Sinkhorn divergence (SIM) computable in closed form and with low computational complexity?
• RQ3 Can an empirical estimator of the SIM be constructed that remains stable under Monte Carlo sampling from the generated distribution?
• RQ4 How does the performance of SiNG compare to SGD-type solvers in terms of convergence speed and sample efficiency on generative modeling tasks?
• RQ5 Does the continuous-time limit of SiNG preserve reparameterization invariance, as required for robust optimization in deep generative models?

Key findings

• The Sinkhorn information matrix (SIM) admits an explicit expression and can be computed in O(log(1/ε)) operations with respect to the desired accuracy ε, enabling efficient and exact updates.
• The proposed SiNG algorithm is invariant to reparameterization in the continuous-time limit, ensuring robustness to nonlinear, non-injective mappings Tθ.
• The empirical estimator eSIM of the SIM is stable under Monte Carlo sampling, with convergence guarantees based on the continuity of the Fréchet derivative of the Sinkhorn potential.
• In experiments, SiNG outperforms state-of-the-art SGD-type solvers (e.g., Adam, RMSProp) in both training efficiency and sample quality on generative modeling benchmarks.
• The method achieves faster convergence and better FID scores compared to baseline optimizers when training on CIFAR-10 and CelebA datasets using a DC-GAN architecture.
• The implementation leverages PyTorch’s autograd and the geomloss library to efficiently compute gradients and matrix-vector products without explicitly forming the SIM.