Score-based Generative Neural Networks for Large-Scale Optimal Transport
Year: 2022

Abstract:
We consider the fundamental problem of sampling the optimal transport coupling between given source and target distributions. In certain cases, the optimal transport plan takes the form of a one-to-one mapping from the source support to the target support, but learning or even approximating such a map is computationally challenging for large and high-dimensional datasets due to the high cost of linear programming routines and an intrinsic curse of dimensionality. We study instead the Sinkhorn problem, a regularized form of optimal transport whose solutions are couplings between the source and the target distribution. We introduce a novel framework for learning the Sinkhorn coupling between two distributions in the form of a score-based generative model. Conditioned on source data, our procedure iterates Langevin Dynamics to sample target data according to the regularized optimal coupling. Key to this approach is a neural network parametrization of the Sinkhorn problem, and we prove convergence of gradient descent with respect to network parameters in this formulation. We demonstrate its empirical success on a variety of large scale optimal transport tasks.

heading: 1 Introduction
text: 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.

heading: 2 Background and Related Work
text: 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.

heading: 2.1 Regularized Optimal Transport
text: 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.

heading: 2.2 Langevin Sampling and Score Based Generative Modeling
text: Given access to optimal dual variables ϕ∗(x), ψ∗(y), it is easy to evaluate the density of the corresponding optimal coupling according to Proposition 2.4. To generate samples distributed according to this coupling, we apply Langevin Sampling. The key quantity used in Langevin sampling of a generic (possibly unnormalized) probability measure p(x) is its score function, given by ∇x log p(x) for x ∈ X . The algorithm is an iterative Monte Carlo method which generates approximate samples x̃t by iterating the map
x̃t = x̃t−1 + ∇x log p(x̃t−1) + √ 2 zt
where > 0 is a step size parameter and where zt ∼ N (0, I) independently at each time step t ≥ 0. In the limit → 0 and T →∞, the samples x̃T converge weakly in distribution to p(x). Song and Ermon [24] introduce a method to estimate the score with a neural network sϑ(x), trained on samples from p(x), so that it approximates sϑ(x) ≈ ∇x log p(x) for a given x ∈ X . To generate samples, one may iterate Langevin dynamics with the score estimate in place of the true score.
To scale this method to high dimensional image datasets, Song and Ermon [24] propose an annealing scheme which samples noised versions of p(x) as the noise is gradually reduced. One first samples a noised distribution p(x) ∗ N (0, τ1), at noise level τ1. The noisy samples, which are presumed to lie near high density regions of p(x), are used to initialize additional rounds of Langevin dynamics at diminishing noise levels τ2 > . . . > τN > 0. At the final round, Annealed Langevin Sampling outputs approximate samples according to the noiseless distribution. Song and Ermon [24] demonstrate that Annealed Langevin Sampling (ALS) with score estimatation can be used to generate sample images that rival the quality of popular generative modeling tools like GANs or VAEs.

heading: 3 Conditional Sampling of Regularized Optimal Transport Plans
text: Our approach can be split into two main steps. First, we approximate the density of the optimal Sinkhorn coupling π∗(x, y) which minimizes Kλ(π) over the data. To do so, we apply the large-scale stochastic dual approach introduced by Seguy et al. [22], which involves instantiating neural networks ϕθ : X → R and ψθ : Y → R that serve as parametrized dual variables. We then maximize Jλ(ϕθ, ψθ) with respect to θ via gradient descent and take the resulting parameters θ∗ and the associated transport plan π̂(x, y) = M(V (x, y;ϕθ∗ , ψθ∗))σ(x)τ(y). This procedure is shown in Algorithm 1. Note that when the dual problem is only approximately maximized, π̂ need not be a normalized density. We therefore call π̃ the pseudo-coupling which approximates the true Sinkhorn coupling π̂.
After optimizing θ∗, we sample the conditional π̂Y |X=x(y) using Langevin dynamics. The score estimator for the conditional distribution is,
∇y log π̂Y |X=x(y) = ∇y[log(M(V (x, y;ϕθ∗ , ψθ∗))σ(x)τ(y))− log(σ(x))] ≈ ∇y log(M(V (x, y;ϕθ∗ , ψθ∗))) + sϑ(y).
We therefore approximate∇yπ̂Y |X=x(y) by directly differentiating logM(V (x, y)) using standard automatic differentiation tools and adding the result to an unconditional score estimate sϑ(y). The full Langevin sampling algorithm for general regularized optimal transport is shown in Algorithm 2.
We note that our method has the effect of biasing the Langevin iterates towards the region where π̂Y |X=x is localized. This may be beneficial for Langevin sampling, which enjoys exponentially fast mixing when sampling log-concave distributions. In the supplementary material, we prove a known result that for the entropy regularized problem given in Proposition 2.3: the compatibility M(V (x, y;ϕ∗, ψ∗)) = 1e exp ( 1 λ (ϕ ∗(x) + ψ∗(y)− c(x, y)) )
is log-concave with respect to y. For λ → 0, this localizes around the optimal transport of x, T (x), and so heuristically should lead to faster mixing.

heading: 4 Theoretical Analysis
text: In principle, Definition 2.1 poses an unconstrained optimization problem over R2d, which could be optimized directly by gradient descent on ϕ, ψ as vectors. The point of a more expensive neural network parametrization ϕθ, ψθ is to learn a continuous distribution that agrees with the empirical
Algorithm 1 Density Estimation. Input: Step size γ, batch size m Input: Nets ϕθ1 , ψθ2 . Input: Datasets σ, τ . Time steps T > 0. Output: Trained ϕθ∗1 , ψθ∗2 . for t = 1 . . . T . do
Sample X1, . . . , Xm ∼ σ, and Y1, . . . , Ym ∼ τ . Stochastic gradient update ϕθ1 , ψθ2 :
∆1 ← m∑
i,j=1
∇θ1 [ϕθ1(Xi)−H∗(V (Xi, Yj))].
∆2 ← m∑
i,j=1
∇θ2 [ψθ2(Yj)−H∗(V (Xi, Yj))].
θ1 ← θ1 + γ∆1. θ2 ← θ2 + γ∆2.
end for Output parameters {θ1, θ2}.
Algorithm 2 SCONES Sampling Procedure Input: Noise levels τ1 > . . . > τN . Input: Dual vars. ϕ̃(x), ψ̃(y). Source x ∈ X . Input: Time steps T > 0. Step size > 0. Output: Data sample ỹ ∼ πY |X=x(y). Initialize ỹ1,0 ∼ N (0, I). for τi, i = 1 . . . N do
for t = 1 . . . T . do Sample z ∼ N (0, τi). Compute score update: ∆s ← sϑ(ỹi,t−1). ∆π ← ∇y logM(V (x, ỹi,t−1; ϕ̃, ψ̃)). x̃i,t = x̃i,t−1 + ( /2)(∆s + ∆π) + √ z. end for Initialize x̃i+1,0 = x̃i,T .
end for Output sample x̃N,T .
optimal transport plan between discrete σ, τ , and that generalizes to a continuous space containing X × Y . By training ϕθ, ψθ, we approximate the underlying continuous data distribution up to optimization error and up to statistical estimation error between the empirical coupling and the population coupling. In the present section, we justify this approach by proving convergence of Algorithm 1 to the global maximizer of Jλ(ϕ,ψ), under assumptions of large network width, along with a quantitative bound on optimization error. In Section B.1 of the Appendix, we provide a cursory analysis of rates of statistical estimation of entropy regularized Sinkhorn couplings. We make the following main assumptions on the neural networks ϕθ and ψθ.
Assumption 4.1 (Approximate Linearity). Let fθ(x) be a neural network and set Kθ(x) = [Jfθ (x)][J f θ (x)]
T where Jfθ (x) is the Jacobian of fθ(x) with respect to θ. Let Θ be a set of feasible weights, for example those reachable by gradient descent. Then fθ(x) must satisfy,
1. There exists R 0 so that Θ ⊆ B(0, R), where B(0, R) is the Euclidean ball of radius R.
2. There exist ρM > ρm > 0 such that for θ ∈ Θ and for all data points {Xi}Ni=1,
ρM ≥ λmax(Kθ(Xi)) ≥ λmin(Kθ(Xi)) ≥ ρm > 0.
3. For θ ∈ Θ and for all data points {Xi}Ni=1, the Hessian matrix D2θfθ(xi) is bounded in spectral norm: ‖D2θfθ(xi)‖ ≤ ρM Ch
, where Ch 0 depends only on R, N , and the regularization λ.
The dependencies of Ch are made clear in the Supplemental Materials, Section B. The quantityKθ(x) is called the neural tangent kernel (NTK) associated with the network fθ(x). It has been shown for a variety of nets that, at sufficiently large width, the NTK is well conditioned and nearly constant on the set of weights reachable by gradient descent on a convex objective function [16, 6, 7]. For instance, fully-connected networks with smooth and Lipschitz-continuous activations fall into this class and hence satisfy Assumption 4.1 when the width of all layers is sufficiently large [16].
First, we show in Theorem 4.2 that when ϕ, ψ are parametrized by neural networks satisfying Assumption 4.1, gradient descent converges to global maximizers of the dual objective. This provides additional justification for the large-scale approach of Seguy et al. [22]. Theorem 4.2 (Optimizing Neural Nets). Suppose Jλ(ϕ,ψ) is 1s -strongly smooth in l∞ norm. Let ϕθ, ψθ be neural networks satisfying Assumption 4.1 for the dataset {(xi, yi)}Ni=1, N = |X | · |Y|.
Then gradient descent of Jλ(ϕθ, ψθ) with respect to θ at learning rate η = λ2ρM converges to an -approximate global maximizer of Jλ in at most ( 2κR2
s ) −1 iterations, where κ = ρMρm .
Given outputs ϕ̂, ψ̂ of Algorithm 1, we may assume by Theorem 4.2 that the networks are - approximate global maximizers of Jλ(ϕ,ψ). Due to λα-strong convexity of the primal objective, the optimization error bounds the distance of the underlying pseudo-plan π̂ from the true global extrema. We make this bound concrete in Theorem 4.3, which guarantees that approximately maximizing Jλ(ϕ,ψ) is sufficient to produce a close approximation of the true empirical Sinkhorn coupling. Theorem 4.3 (Stability of the OT Problem). Suppose Kλ(π) is s-strongly convex in l1 norm and let L(ϕ,ψ, π) be the Lagrangian of the regularized optimal transport problem. For ϕ̂, ψ̂ which are -approximate maximizers of Jλ(ϕ,ψ), the pseudo-plan π̂ = Mf (V (x, y; ϕ̂, ψ̂))σ(x)τ(y) satisfies
|π̂ − π∗|1 ≤ √ 2
s ≤ 1 s ∣∣∣∇π̂L(ϕ̂, ψ̂, π̂)∣∣∣ 1 .
Theorem 4.3 guarantees that if one can approximately optimize the dual objective using Algorithm 1, then the corresponding coupling π̂ is close in l1 norm to the true optimal transport coupling. This approximation guarantee justifies the choice to draw samples π̃Y |X=x as an approximation to sampling πY |X=x instead. Both Theorems 4.2 and 4.3 are proven in Section B of the Appendix.

heading: 5 Experiments
text: Our main point of comparison is the barycentric projection method proposed by Seguy et al. [22], which trains a neural network Tθ : X → Y to map source data to target data by optimizing the objective θ := arg minθ EπY |X=x [|Tθ(x)− Y |2]. For transportation experiments between USPS [19] and MNIST [13] datasets, we scale both datasets to 16px and parametrize the dual variables and barycentric projections by fully connected ReLU networks. We train score estimators for MNIST and USPS at 16px resolution using the method and architecture of [24]. For transportation experiments using CelebA [17], dual variables ϕ, ψ are parametrized as ReLU FCNs with 8, 2048-dimensional hidden layers. Both the barycentric projection and the score estimators use the U-Net based imageto-image architecture introduced in [24]. Numerical hyperparameters like learning rates, optimizer parameters, and annealing schedules, along with additional details of our neural network architectures, are tabulated in Section C of the Appendix.

heading: 5.1 Optimal Transportation of Image Data
text: We show in Figure 3 a qualitative plot of SCONES samples on transportation between MNIST and USPS digits. We also show in Section 1, Figure 1 a qualitative plot of transportation of CelebA images. Because barycentric projection averages πY |X=x, output images are blurred and show visible mixing of multiple digits. By directly sampling the optimal transport plan, SCONES can separate these modes and generate more realistic images.
At low regularization levels, Algorithm 1 becomes more expensive and can become numerically unstable. As shown in Figures 3 and 1, SCONES can be used to sample the Sinkhorn coupling in intermediate regularization regimes, where optimal transport has a nontrivial effect despite πY |X=x not concentrating on a single image.
To quantitatively assess the quality of images generated by SCONES, we compute the FID scores of generated CelebA images on two optimal transport problems: transporting 2x downsampled CelebA images to CelebA (the ‘Superres.’ task) and transporting CelebA to CelebA (the ’Identity’ task) for a variety of regularization parameters. The FID score is a popular measurement of sample quality for image generative models and it is a proxy for agreement between the distribution of SCONES samples of the marginal πY (y) and the true distribution τ(y). In both cases, we partition CelebA into two datasets of equal size and optimize Algorithm 1 using separated partitions as source and target data, resizing the source data in the superresolution task. As shown in Table 1, SCONES has a significantly lower FID score than samples generated by barycentric projection. However, under ideal tuning, the unconditional score network generates CelebA samples with FID score 10.23 [24], so there is some cost in sample quality incurred when using SCONES.

heading: 5.2 Sampling Synthetic Data
text: To compare SCONES to a ground truth Sinkhorn coupling in a continuous setting, we consider entropy regularized optimal transport between Gaussian measures on Rd. Given σ = N (µ1,Σ1) and τ = N (µ2,Σ2), the Sinkhorn coupling of σ, τ is itself a Gaussian measure and it can be written in closed form in terms of the regularization λ and the means and covariances of σ, τ [8]. In dimensions d ∈ {2, 16, 54, 128, 256}, we consider Σ1,Σ2 whose eigenvectors are uniform random (i.e. drawn from the Haar measure on SO(d)) and whose eigenvalues are sampled uniform i.i.d.
from [1, 10]. In all cases, we set means µ1, µ2 equal to zero and choose regularization λ = 2d. In the Gaussian setting, E[‖x− y‖22] is of order d, so this choice of scaling ensures a fair comparison across problem dimensions by fixing the relative magnitudes of the cost and regularization terms.
We evaluate performance on this task using the Bures-Wasserstein Unexplained Variance Percentage [4], BW-UV(π̂, πλ), where πλ is the closed form solution given by Janati et al. [8] and where π̂ is the joint empirical covariance of k = 10000 samples (x, y) ∼ π generated using either SCONES or Barycentric Projection. We train SCONES according to Algorithm 1 and generate samples according to Algorithm 2. In place of a score estimate, we use the ground truth target score ∇y log τ(y) = Σ−12 (y − µ2) and omit annealing. We compare SCONES samples to the true solution in the BW-UVP metric [4] which is measured on a scale from 0 to 100, lower is better. We report BW-UVP(π̂, πλ) where π̂ is a 2d-by-2d joint empirical covariance of SCONES samples (x, y) ∼ π̂ or of BP samples (x, Tθ(x)), x ∼ σ and πλ is the closed-form covariance.

heading: 6 Discussion and Future Work
text: We introduce and analyze the SCONES method for learning and sampling large-scale optimal transport plans. Our method takes the form of a conditional sampling problem for which the conditional score decomposes naturally into a prior, unconditional score ∇y log τ(y) and a “compatibility term” ∇y logM(V (x, y)). This decomposition illustrates a key benefit of SCONES: one score network may re-used to cheaply transport many source distributions to the same target. In contrast, learned forward-model-based transportation maps require an expensive training procedure for each distinct pair of source and target distribution. This benefit comes in exchange of increased computational cost of iterative sampling. For example, generating 1000 samples requires roughly 3 hours using one NVIDIA 2080 Ti GPU. The cost to sample score-based models may fall with future engineering advances, but iterative sampling intrinsically require multiple forward pass evaluations of the score estimator as opposed to a single evaluation of a learned transportation mapping.
There is much future work to be done. First, we study only simple fully connected ReLU networks as parametrizations of the dual variables. Interestingly, we observe that under L2 transportation cost, parametrization by multi-layer convolutional networks perform equally or worse than their FCN counterparts when optimizing Algorithm 1. One explanation may be the permutation invariance of L2 cost: applying a permutation of coordinates to the source and target distribution does not change the optimal objective value and the optimal coupling is simply conjugated by a coordinate permutation. As a consequence, the optimal coupling may depend non-locally on input data coordinates, violating the inductive biases of localized convolutional filters. Understanding which network parametrizations or inductive biases are best for a particular choice of transportation cost, source distribution, and target distribution, is one direction for future investigation.
Second, it remains to explore whether there is a potential synergistic effect between Langevin sampling and optimal transport. Heuristically, as λ→ 0 the conditional plan πY |X=x concentrates around the transport image of x, which should improve the mixing time required by Langevin dynamics to explore high density regions of space. In Section B of the Appendix, we prove a known result, that the entropy regularized L2 cost compatibility term M(V (x, y)) = eV (x,y)/λ is a log-concave function of y for fixed x. It the target distribution is itself log-concave, the conditional coupling πY |X=x is also log-concanve and hence Langevin sampling enjoys exponentially fast mixing time. However, more work is required to understand the impacts of non-log-concavity of the target and of optimization errors when learning the compatibility and score functions in practice. We look forward to future developments on these and other aspects of large-scale regularized optimal transport.
Acknowledgments and Disclosure of Funding
M.D. acknowledges funding from Northeastern University’s Undergraduate Research & Fellowships office and the Goldwater Award. P.H. was supported in part by NSF awards 2053448, 2022205, and 1848087.

Review summaries:
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.