Score-based Generative Neural Networks for Large-Scale Optimal Transport

1 Introduction . It is often useful to compare two data distributions by computing a distance between them in some appropriate metric . For instance , statistical distances can be used to fit the parameters of a distribution to match some given data . Comparison of statistical distances can also enable distribution testing , quantification of distribution shifts , and provide methods to correct for distribution shift through domain adaptation [ 12 ] . Optimal transport theory provides a rich set of tools for comparing distributions in Wasserstein Distance . Intuitively , an optimal transport plan from a source distribution σ ∈ M+ ( X ) to a target distribution τ ∈M+ ( Y ) is a blueprint for transporting the mass of σ to match that of τ as cheaply as possible with respect to some ground cost . Here , X and Y are compact metric spaces andM+ ( X ) denotes the set of positive Radon measures over X , and it is assumed that σ , τ are supported over all of X , Y respectively . The Wasserstein Distance between two distributions is defined to be the cost of an optimal transport plan . Because the ground cost can incorporate underlying geometry of the data space , optimal transport plans often provide a meaningful correspondence between points in X and Y . A famous example is given by Brenier ’ s Theorem , which states that , when X , Y ⊆ Rd and σ , τ have finite variance , the optimal transport plan under a squared-l2 ground cost is realized by a map T : X → Y [ 26 , Theorem 2.12 ] . However , it is often computationally challenging to exactly compute optimal transport plans , as one must exactly solve a linear program requiring time which is super-quadratic in the size of input datasets [ 5 ] . ∗Work done while an Instructor in Applied Mathematics at MIT . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Instead , we opt to study a regularized form of the optimal transport problem whose solution takes the form of a joint density π ( x , y ) with marginals πX ( x ) = σ ( x ) and πY ( y ) = τ ( y ) . A correspondence between points is given by the conditional distribution πY |X=x ( y ) , which relates each input point to a distribution over output points . In recent work [ 22 ] , the authors propose a large-scale stochastic dual approach in which π ( x , y ) is parametrized by two continuous dual variables that may be represented by neural networks and trained at large-scale via stochastic gradient ascent . Then , with access to π ( x , y ) , they approximate an optimal transport map using a barycentric projection of the form T : x 7→ arg miny EπY |X=x [ d ( y , Y ) ] , where d : Y × Y → R is a convex cost on Y . Their method is extended by [ 15 ] to the problem of learning regularized Wasserstein barycenters . In both cases , the Barycentric projection is observed to induce averaging artifacts such as those shown in Figure 2 . Instead , we propose a direct sampling strategy to generate samples from πY |X=x ( y ) using a scorebased generative model . Score-based generative models are trained to sample a generic probability density by iterating a stochastic dynamical system knows as Langevin dynamics [ 24 ] . In contrast to projection methods for large-scale optimal transport , we demonstrate that pre-trained score based generative models can be naturally applied to the problem of large-scale regularized optimal transport . Our main contributions are as follows : 1 . We show that pretrained score based generative models can be easily adapted for the purpose of sampling high dimensional regularized optimal transport plans . Our method eliminates the need to estimate a barycentric projection and it results in sharper samples because it eliminates averaging artifacts incurred by such a projection . 2 . Score based generative models have been used for unconditional data generation and for conditional data generation in settings such as inpainting . We demonstrate how to adapt pretrained score based generative models for the more challenging conditional sampling problem of regularized optimal transport . 3 . Our method relies on a neural network parametrization of the dual regularized optimal transport problem . Under assumptions of large network width , we prove that gradient descent w.r.t . neural network parameters converges to a global maximizer of the dual problem . We also prove optimization error bounds based on a stability analysis of the dual problem . 4 . We demonstrate the empirical success of our method on a synthetic optimal transport task and on optimal transport of high dimensional image data . 2 Background and Related Work . We will briefly review some key facts about optimal transport and generative modeling . For a more expansive background on optimal transport , we recommend the references [ 26 ] and [ 25 ] . Source Target samples KL λ=0.005 KL λ=0.1 Figure 2 : Samples generated by SCONES for entropy regularized optimal transport including the samples shown in Figure 1 . At regularization λ = 0.005 , optimal transportation with L2 cost has a visible effect on generated images . This effect diminishes at increased regularization λ = 0.1 . 2.1 Regularized Optimal Transport . We begin by reviewing the formulation of the regularized OT problem . Definition 2.1 ( Regularized OT ) . Let σ ∈ M+ ( X ) and τ ∈ M+ ( Y ) be probability measures supported on compact sets X , Y . Let c : X × Y → R be a convex , lower semi-continuous function representing cost of transporting a point x ∈ X to y ∈ Y . The regularized optimal transport distance OTλ ( σ , τ ) is given by OTλ ( σ , τ ) = min π Eπ [ c ( x , y ) ] + λH ( π ) ( 1 ) subject to πX = σ , πY = τ π ( x , y ) ≥ 0 where H : M+ ( X × Y ) → R is a convex regularizer and λ ≥ 0 is a regularization parameter . We are mainly concerned with optimal transport of empirical distributions , where X and Y are finite and σ , τ are empirical probability vectors . In most of the following theorems , we will work in the empirical setting of Definition 2.1 , so that X and Y are finite subsets of Rd and σ , τ are vectors in the probability simplices of dimension |X | and |Y| , respectively . We refer to the objective Kλ ( π ) = Eπ [ c ( x , y ) ] + λH ( π ) as the primal objective , and we will use Jλ ( ϕ , ψ ) to refer to the associated dual objective , with dual variables ϕ , ψ . Two common regularizers areH ( π ) = KL ( π||σ×τ ) andH ( π ) = χ2 ( π||σ×τ ) , sometimes called entropy and l2 regularization respectively : KL ( π||σ × τ ) = Eπ [ log ( dπ ( x , y ) dσ ( x ) dτ ( y ) ) ] , χ2 ( π||σ × τ ) = Eσ×τ [ ( dπ ( x , y ) dσ ( x ) dτ ( y ) ) 2 ] where dπ ( x , y ) dσ ( x ) dτ ( y ) is the Radon-Nikodym derivative of π with respect to the product measure σ × τ . These regularizers contribute useful optimization properties to the primal and dual problems . For example , KL ( π||σ× τ ) is exactly the mutual information Iπ ( X ; Y ) of the coupling ( X , Y ) ∼ π , so intuitively speaking , entropy regularization explicitly prevents πY |X=x from concentrating on a point by stipulating that the conditional measure retain some bits of uncertainty after conditioning . The effects of this regularization are described by Propositions 2.2 and 2.3 . First , regularization induces convexity properties which are useful from an optimization perspective . Proposition 2.2 . In the empirical setting of Definition 2.1 , the entropy regularized primal problem Kλ ( π ) is λ-strongly convex in l1 norm . The dual problem Jλ ( ϕ , ψ ) is concave , unconstrained , and 1λ -strongly smooth in l∞ norm . Additionally , these objectives witness strong duality : infπ∈M+ ( X×Y ) Kλ ( π ) = supϕ , ψ∈R2d Jλ ( ϕ , ψ ) , and the extrema of each objective are attained over their respective domains . In addition to these optimization properties , regularizing the OT problem induces a specific form of the dual objective and resulting optimal solutions . Proposition 2.3 . In the setting of Proposition 2.2 , the KL-regularized dual objective takes the form Jλ ( ϕ , ψ ) : = Eσ [ ϕ ( x ) ] + Eτ [ ψ ( y ) ] − λEσ×τ [ 1 e exp ( 1 λ ( ϕ ( x ) + ψ ( y ) − c ( x , y ) ) ) ] . The optimal solutions ϕ∗ , ψ∗ = arg maxϕ , ψ∈R2d Jλ ( ϕ , ψ ) and π ∗ = arg minπ∈M+ ( X×Y ) Kλ ( π ) satisfy π∗ ( x , y ) = 1 e exp ( 1 λ ( ϕ∗ ( x ) + ψ∗ ( y ) − c ( x , y ) ) ) σ ( x ) τ ( y ) . These propositions are specializations of Proposition 2.4 and they are well-known to the literature on entropy regularized optimal transport [ 5 , 2 ] . The solution π∗ ( x , y ) of the entropy regularized problem is often called the Sinkhorn coupling between σ and τ in reference to Sinkhorn ’ s Algorithm [ 23 ] , a popular approach to efficiently solving the discrete entropy regularized OT problem . For arbitrary choices of regularization , we call π∗ ( x , y ) a Sinkhorn coupling . Propositions 2.2 and 2.3 illustrate the main desiderata when choosing the regularizer : that H ( π ) , and hence Kλ ( π ) , be strongly convex in l1 norm and that H ( π ) induce a nice analytic form of π∗ in terms of ϕ∗ , ψ∗ . In regards to the former , H ( π ) is akin to barrier functions used by interior point methods [ 5 ] . Prior work [ 2 ] is an example of the latter , in which it is shown that for discrete optimal transport , χ2 regularization yields an analytic form of π∗ having a thresholding operation that promotes sparsity . Conveniently , the KL and χ2 regularizers both belong to the class of f -Divergences , which are statistical divergences of the form Df ( p||q ) = Eq [ f ( p ( x ) q ( x ) ) ] . where f : R→ R is convex , f ( 1 ) = 0 , p , q are probability measures , and p is absolutely continuous with respect to q . For example , the KL regularizer has fKL ( t ) = t log ( t ) and the χ2 regularizer has fχ2 ( t ) = t 2 − 1 . The f -Divergences are good choices for regularizing optimal transport : strong convexity of f is a sufficient condition for strong convexity of Hf ( π ) : = Df ( π||σ × τ ) in l1 norm , and the form of f is the aspect which determines the form of π∗ in terms of ϕ∗ , ψ∗ . This relationship is captured by the following generalization of Propositions 2.2 and 2.3 , which we prove in Section B of the Supplemental Materials . Proposition 2.4 . Consider the empirical setting of Definition 2.1 . Let f ( v ) : R→ R be a differentiable α-strongly convex function with convex conjugate f∗ ( v ) . Set f∗′ ( v ) = ∂vf∗ ( v ) . Define the violation function V ( x , y ; ϕ , ψ ) = ϕ ( x ) + ψ ( y ) − c ( x , y ) . Then , 1 . The Df regularized primal problem Kλ ( π ) is λα-strongly convex in l1 norm . With respect to dual variables ϕ ∈ R|X | and ψ ∈ R|Y| , the dual problem Jλ ( ϕ , ψ ) is concave , unconstrained , and 1λα -strongly smooth in l∞ norm . Strong duality holds : Kλ ( π ) ≥ Jλ ( ϕ , ψ ) for all π , ϕ , ψ , with equality for some triple π∗ , ϕ∗ , ψ∗ . 2 . Jλ ( ϕ , ψ ) takes the form Jλ ( ϕ , ψ ) = Eσ [ ϕ ( x ) ] + Eτ [ ψ ( y ) ] − Eσ×τ [ H∗f ( V ( x , y ; ϕ , ψ ) ) ] , where H∗f ( v ) = λf ∗ ( λ−1v ) . 3 . The optimal solutions ( π∗ , ϕ∗ , ψ∗ ) satisfy π∗ ( x , y ) = Mf ( V ( x , y ; ϕ , ψ ) ) σ ( x ) τ ( y ) where Mf ( x , y ) = f∗′ ( λ−1v ) . For this reason , we focus in this work on f -Divergence-based regularizers . Where it is clear , we will drop subscripts on regularizer H ( π ) and the so-called compatibility function M ( v ) and we will omit the dual variable arguments of V ( x , y ) . The specific form of these terms for KL regularization , χ2 regularization , and a variety of other regularizers may be found in Section A of the Appendix .

Summary: Using score-based generative modeling and Langevin dynamics, this paper proposes a new method to sample from the Sinkhorn coupling distribution. The method first solves the dual problem of regularized optimal transport by parameterizing the dual variables with neural networks. After learning the dual variables, the coupling distribution is known up to a normalizing constant, and can therefore be sampled from using Langevin dynamics and a score-based model as the prior. Experiments demonstrate advantages over barycenter methods.