Score-based Generative Neural Networks for Large-Scale Optimal Transport
Authors: Mara Daniels
Tyler Maunu
PAul HAnd
Venue: NeurIPS 2021 Poster (2021)
Keywords: optimal transport
regularized optimal transport
langevin sampling
generative modelling
score-based generative model

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. 

Score-based Generative Neural Networks for Large-Scale Optimal Transport
Max Daniels Northeastern University daniels, www.sss.ooooeeeteeeenn.ooo edu
Tyler Maunu Brandeis University maunu@brandeis. edu
Paul Hand Northeastern University p hand@northeastemn. edu
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 one-to-one mapping from the source support to the target support, but learning even approximating such a map is computationally challenging for large and high-dimensional datasets due to the high cost 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 novel framework for learning the Sinkhorn coupling between two distributions the form of a score-based generative model. Conditioned on source data, procedure iterates Langevin Dynamics to sample target data according to the regularized optimal coupling. Key to this approach is 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.
Introduction
It is often useful to compare two data distributions by computing distance between them in some appropriate metric. For instance, statistical distances can be used to fit the parameters of distribution match some given data. Comparison of statistical distances 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 Wasserstein Distance. Intuitively, optimal transport plan from source distribution E M4 (X) target distribution ())) is a blueprint for transporting the mass of a to match that of 7 as cheaply as possible with respect to some ground cost. Here, X' and y are compact metric spaces and M (2)) denotes the set of positive Radon measures over X, and it assumed that σ, 7 are supported over all of X, respectively. The Wasserstein Distance between distributions is defined be the 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 and y. A famous example given by Brenier's Theorem, which states that, when X, y C [Rd and a, have finite variance, the optimal transport plan under a squared-l2 ground cost is realized by map T X + ) [26 Theorem 2.1 However, is often computationally challenging exactly compute optimal transport plans, as one must exactly solve a linear program requiring time which is super-quadrad in the size of input datasets [5].
Work done while Instructor in Applied Mathematics MIT.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Source
SCONES Samples
Figure 1: We SCONES sample the mean- wean-squured-L' cost, entropy regularized optimal transport mapping between 2x downsampled CelebA images (Source) and unmodified CelebA images (Target) 0.005 regularization.
Instead, we opt to study a regularized form of the optimal transport problem whose solution takes the form of joint density T [2, with marginals (x) (x)) and xy (v) =(v). A correspondence between points given by the conditional distribution Ty ( (v), which relates each input point to distribution over output points.
In recent work [22] the authors propose large-scale stochastic dual approach in which π (z, 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 (xx y). they approximate an optimal transport map using a barycentric projection of the form T' + arg ming E [d(y, Y)J, where d: y x y + R is convex cost on V. Their method is extended by (15) to the problem of learning regularized Wasserstein barycenters. both cases, the Barycentric projection observed to induce averaging artifacts such as those shown in Figure 2
Instead, we propose a direct sampling strategy to generate samples from Ry (y) using score- based generative model. Score-based generative models are trained to sample generic probability density by iterating a stochastic dynamical system knows as Langevin dynamics [24]. 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 results in sharper samples because eliminates averaging artifacts incurred by such 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 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 stability analysis of the dual problem.
4. We demonstrate the empirical success of our method on synthetic optimal transport task and on optimal transport of high dimensional image data.
Background and Related Work
We will briefly review some key facts about optimal transport and generative modeling. For more expansive background optimal transport, recommend the references (26) and [25).

---

Source
Target samples
KL
x=0.005
KL
A=
Figure 2: Samples generated by SCONES for entropy regularized optimal transport including the samples shown in Figure At regularization À = 0.005, optimal transportation with L cost has 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 a E (M') and t E M+ be probability measures supported compact sets X, y. Let X x x + be convex, lower semi-continuous function representing cost of transporting a point x E X' to E y. The regularized optimal transport distance OTx 7) is given by
OT 7) = min E [c(,, 8)1 (H
(1)
subject to
xx a, x < v)
ny =
where M 3) + a convex regularizer and 2 is regularization parameter.
We are mainly concerned with optimal transport of empirical distributions, where X' and Y are finite and o. are empirical probability vectors. In most of the following theorems, we will work in the empirical setling of Definition 2.1 so that X' and ) are finite subsets of IR and a, 7 are vectors in the probability simplices dimension | and [/]. respectively.
We refer to the objective (T) = Eg [c(x, y)] + (H)) as the primal objective, and we will use Jx [x(s,, 사b) to refer to the associated dual objective, with dual variables p. v. Two common regularizers are (T) KL(T||o x) and H (T)) ( XT), sometimes called entropy and l2 regularization respectively:
dr di((:, y) (x )
(
dr( 9) do
KL(T| x xT) Ex log
( x x = Exxr
do(x))r(()) is the Radon-Nikodym derivative of with respect the product measure These regularizers contribute useful optimization properties to the primal and dual problems.
where dr((n.)
For example, KT.(T|/a x x) is exactly the mutual information I, (X; ; Y) of the coupling (X, Y) ~ n, so intuitively speaking, entropy regularization explicitl prevents Tf | from concentrating 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.7 the entropy regularized primal prob- lem Kx(T) is X-strongly convex in l norm. The dual problem Jx(p, v)) is concave, uncon- strained, and 1 -strongly smooth in loo norm. Additionally, these objectives witness strong duality: inf (Xxx)) Kx(T)) ) sup Jx(y, 1) and the extrema of each objective are attained over their respective domains.
In addition these optimization properties, regularizing the OT problem induces 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
Jx(e; 16) EE [x(x]] + Er(((v)]
L Eexp (1 +)())))(())+ (()) -cx))c((.))))))
-B
The optimal solutions x²,, 10= = arg max p, Ja(p, and x² = ((g Rx(T) satisfy
1 x² (x, exp (H60)))+ -(x, d(x)T(v).
These propositions are specializations of Proposition 2.4 and they well-known to the literature entropy regularized optimal transport 2]. The solution x >"(x, (r, y) of the entropy regularized problem often called the Sinkhorn coupling between and in reference Sinkhorn's Algorithm [23], popular approach to efficiently solving the discrete entropy regularized OT problem. For arbitrary choices of regularization, we call . (2z, y) Sinkhorn coupling.
Propositions 2.2 and 2.3 illustrate the main desiderata when choosing the regularizer: that ( and hence Xxx (Tr). be strongly convex in norm and that (T) induce nice analytic form of x² in terms of 40" r6... In regards to the former, H (T) is akin to barrier functions used by interior point methods [S]. Prior work (2) is example of the latter, in which it is shown that for discrete optimal transport, x" regularization yields an analytic form of æ having a thresholding operation that promotes sparsity.
Conveniently, the KL and x2 regularizers both belong to the class of f-Divergences, which are statistical divergences of the form
p(z) g(x))
Dy(pll9) Eg Eg f
where IR -
R is convex, f(1) = 0, p.q q are probability measures, and p is absolutely continuous with respect to q. For example, the KI. regularizer has (t) log(t) and the x" regularizer has The f-Divergences are good choices for regularizing optimal transport: strong convexity of is sufficient condition for strong convexity of ( == DT(rllo x T) in ( norm, and the form of f the aspect which determines the form of x in terms yp".. ".. This relationship
(t= = 82
captured by the following generalization of Propositions 2.2 and 2.3 of the Supplemental Materials.
which we prove in Section B
Proposition 2.4. Consider the empirical setting of Definition 2.1 Let f(v) : + R be differen- riable a- strongly convex function with convex conjugate f (v). Set (") (v) Define the violation function (zx, yi ip, 1)) ((x) + (1)) c(x, y). Then,
1. The Dj regularized primal problem K2 is Aa-strongly convex in li norm. With respect to dual variables ip IR X and vi E DH the dual problem Jx(p,, b) is concave, uncon- strained, and Xo -strongly smooth in loo norm. Strong duality holds: Kx(T) 2 Jxly, f all ip, v. with equality for some triple πt",, ,"",,"""""" "²"
2. Jx(p, 1b) takes the form Jxlp, 1) [x(((()) E. [+(() E. IH) p. 4))) where Hj(u) = (A tt).
3. The optimal solutions (πt .4", .b") (b) satisfy
n <' = My(((,,,r
where
For this reason, we focus this work on o J-Divergence-based regularizers. Where clear, we will drop subscripts on regularizer (x)) and the so-called compatibility function M (()) and we will omit the dual variable arguments V(x,, v). The specific form of these terms for KT. regularization, x2 regularization, and a variety of other regularizers may be found in Section A of the Appendix.

---

2.2 Langevin Sampling and Score Based Generative Modeling
Given access optimal dual variables x²((x) 1 (v). 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 generic (possibly unnormalized) probability measure p(x) is its score function, given by Va log p(x) for x E X. The algorithm is an iterative Monte Carlo method which generates approximate samples it by iterating the map
i = f1 + cV2 log p(0)-1) + 2ez
where 0 is step size parameter and where ,t (0, 1) I at each time step /. 2 O. In the limit - 0 and 7 + oc, the samples रल converge weakly in distribution to p(z). Song and Ermon 24 introduce method to estimate the score with neural network sq (3r), trained on samples from p(z), so that it approximates sx)x) & V p(z) for given Z E 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(z) the noise is gradually reduced. One first samples noised distribution p(z) * (0, 71 at noise level T1. 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 72 > > TN > O. 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.
3 Conditional Sampling of Regularized Optimal Transport Plans
Our approach can be split into two main steps. First, we approximate the density of the optimal Sinkhorn coupling x² (zz y) which minimizes Kx (T) 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 (p : X' + R and the y + RR that serve as parametrized dual variables. We then maximize (x v9)) with respect to 0 via gradient descent and take the resulting parameters 0" and the associated transport plan (xx,)) y) M(V (z. y/: ; 1/0- (z) This procedure is shown in Algorithm I Note that when the dual problem is only approximately maximized, क need not be normalized density. We therefore call the psendo-coupling which approximates the true Sinkhorn coupling स.
After optimizing 0", we sample the conditional #t|x== (v) using Langevin dynamics. The score estimator for the conditional distribution is,
V log 7 -y|x==xv) (v)
V (V de ) ())(()())) - log(a(x))]
V, log (V Po ()g + sv)).
We therefore approximate V =x (v) by directly differentiating log (V((z, y)) using standard automatic differentiation tools and adding the result to unconditional score estimate sa(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 fy|x=r 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 known result that for the entropy regularized problem given in Proposition 2.3 the compatibility M(V y sp" ,)"))
e exp (1 (v0 (z) + (v(()) c(z, y))) is log-concave with respect y. For X - 0, this localizes around the optimal transport of x. T(x), and so heuristically should lead to faster mixing.
4 Theoretical Analysis
In principle, Definition 2.1 poses unconstrained optimization problem over R2d, which could be optimized directly by gradient descent on s,, 1/ as vectors. The point of a more expensive neural network parametrization ype, th is to learn continuous distribution that agrees with the empirical
5

---

Algorithm Density Estimatiin..
Algorithm 2 SCONES Sampling Procedure
Input: Step size 7, batch size m
Input: Noise levels T1 > > TN.
Input: Nets 40 4'0g-
Input: Dual vars. (f(x), (y)). Source x E X.
Input: Datasets a, r. Time steps > 0.
Input: Time steps T > 0. Step size 0.
Output: Trained yea Vag. fo do Sample X I,.... Xm O,
Output: Data sample y ~ Ty|x (V).
Initialize V1.0 ~ N (0, I). for T., i .....N... do
and ,,...., Ym ~ T.
for t 1.... T. do
Stochastic gradient update p0,, 470z:
Sample ~ (0, m.). Compute score update:
A
Σ
Ve, Ipe, (X,) (V X,)
A, + so (
m
A f V, log (V 1:,
2 <
Σ
Vog [R00, (V;) (V
LL -i,--1 + (c/2) + An) + Vez.
01 f 01 + A
end for Initialize Ti+1.0 Ti,T
82 + 02 + 742.
end for
end for Output sample ēN.T-
Output parameters {01 02].
optimal transport plan between discrete a, T, and that generalizes a continuous space containing X' V. By training Pe, the, 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| D to the global maximizer of Jx(p; 1)), under assumptions of large network width, along with quantitative bound on optimization error. In Section B. I. of the Appendix, we provide cursory analysis of rates of statistical estimation of entropy regularized Sinkhorn couplings. We make the following main assumptions on the neural networks 6og and thg.
Assumption 4.1 (Approximate Linearity). Let fo(x) be a neural network and set Ko( (J (x)][= (=)]² where Já (ar) is the Jacobian of fo(x)) with respect to 0. Let θ be of fea- sible weights, for example those reachable by gradient descent. Then fo(x) must satisfy,
There exists R > so that θ C B(0, R), where B(0, R) is the Euclidean ball of radius R.
2. There exist PM > Pm > 0 such that for and for all data points {X:]
PM 2 Amax (Ko 2 Amin (Ke((X))) 2 Pm > O.
3. For 0 E θ and for all data points [X,) the Hessian matrix (o.)) is bounded in spectral norm: [o)x(())) < PM C, where C₁ >> 0 depends only on R, N, and the
regularization X.
The dependencies of C₁, are made clear in the Supplemental Materials, Section B The quantity K3g (z) is called the neural tangent kernel (NTK) associated with the network fo(z) (z). 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 instance, fully-connected networks with smooth and Lipschitz-contimune activations fall into this class and hence satisfy Assumption when the width of all layers sufficiently large (16)
convex objective function [16] 6 7 For
First, we show in Theorem 4.2 that when sp, v 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 (x( v6) is --strongly smooth in L.0 norm. Let yPa. vh be neural networks satisfying Assumption 4.1 for the dataset ((r,, W.) N |X| |
Then gradient descent of Jx(sp, vho) with respect to O at learning rate 1J Zpms
converges to an
e-upproximate global maximizer of Jx in at most Rx 6 -1 iterations, where K. PM

---

Source
BP
SCONES Samples
Source
BP
SCONES Samples
2
S
1
L
I
4
9
4
9
9
4
3
3
3
8
3
6
6
5
6
+
4
6
9
8
7
8
8
8
7
7
7
2
4
9
Figure 3: Comparison of Barycentric Projection 22 to SCONES for optimal transport between USPS
and MNIST datasets of handwritten digits. (Left) Transporting MNIST USPS. (Right) Transporting USPS to MNIST. Here, we show transportation of the x² regularized problem at À = 0.001.
Given outputs & v of Algorithm 1 may assume by Theorem 4.2 that the networks approximate global maximizers of (x v)). Due to Aa-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 (1p, +) is sufficient to produce close approximation of the true empirical Sinkhorn coupling.
Theorem 4.3 (Stability of the OT Problem). Suppose (x( is s-strongly convex in l norm and let L(sp, v. v) be the Lagrangian of the regularized optimal transport problem. For &, tb which are c-approximate maximizers of Jx((p, 1)) the pseudo-plan क M3 (2 v; &, (z))(()) satisfies
2e
S
L
L v #)
8
Theorem 4.3 guarrantts that if one can approximately optimize the dual objective using Algorithm 1. then the corresponding coupling á is close in lı norm to the true optimal transport coupling. This approximation guarantee justifies the choice to draw samples x as an approximation to
sampling Tyjx= ma instead. Both Theorems 4.2 and 4.3 are proven in Section B of the Appendix.
Experiments
Our main point of comparison is the barycentric projection method proposed by Seguy al. 22]. which trains neural network Te X + y to map source data to target data by optimizing the Y [3. 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 ReL.U networks. We train score estimators for MNIST and USPS at 16px resolution using the method and architecture of (24] For transportation experiments using CelebA [I7]. dual variables y,, 3 are parametrized as ReLU FCNs with 8. -048-dimensi hidden layers. Both the barycentric projection and the score estimators use the U-Net based image- to-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.
objective arg ming E, Exyix_,, (T)
5.1 Optimal Transportation of Image Data
We show in Figure 3 qualitative plot of SCONES samples on transportation between MNIST and qualitative plot of transportation of CelebA
USPS digits. We also show in Section [ Figure D images. Because barycentric projection averages Ty 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 becomes more expensive and can become numerically unstable. As shown in Figures 3 and [ intermediate regularization regimes, where optimal transport has a nontrivial effect despite Ty [x=x concentrating on single image.
SCONES can be used sample the Sinkhorn coupling in
7

---

KL regularization,
x² regularization,
= 0.1
À 0.01
X = 0.005
0.1
À 0.01
À 0.001
SCONES,
35.59
35.77
43.80
25.84
25.64
25.59
Super-res.
Bary. Proj..
193.92
230.85
228.78
190.10
216.54
212.72
Super-res.
SCONES,
Identity 36.62
34.84
43.99
25.51
25.65
27.88
Bary. Proj..
195.64
217.24
217.67
188.29
219.96
214.90
Identity
Table 1: FID metric of samples generated by barycentric projection and SCONES, computed on 5000 samples from each model. For comparison to unregularized OT methods, we also trained Wasserstein-2 GAN 2W₂ GAN) 14] and a Wasserstein-2 Generative Network (W₂ Gen) [T]. W GAN achieves FIDs 55.77 on the super-res. task and 32.617 on the identity task. W2 Gen achieves FIDs 32.80 on the super-res, task and 20.57 on the identity task.
To quantitatively assess the quality of images generated by SCONES, compute the FID scores of generated CelebA images on two op- timal transport problems: transporting 2x down- sampled CelebA images to CelebA (the 'Super- res.' task) and transporting CelebA to CelebA (the 'Identity' task) for variety of regulariza- tion parameters. The FID score is popular measurement of sample quality for image gen- erative models and it is proxy for agreement between the distribution of SCONES samples of the marginal Ry (y) and the true distribution T3)). In both cases, we partition CelebA into two datasets of equal size and optimize Algo- rithm Dusing separated partitions as source and target data, resizing the source data in the super- resolution task. As shown in Table SCONES has a significantly lower FID score than samples generated by barycentric projection. However, under ideal tuning, the unconditional score net- work generates CelebA samples with FID score 10.23 [24], so there is some cost in sample qual- ity incurred when using SCONES.
SCONES Samples
-11.0
14.5
Barycentric Projection
5.2 Sampling Synthetic Data
To compare SCONES to a ground truth Sinkhorn coupling in continuous setting, sider entropy regularized optimal transport be- tween Gaussian measures on IRd Given N(,, Σ1) and - Nuu,, Mg), the Sinkhorn coupling of a, is itself a Gaussian mea- sure and it can be written in closed form in terms of the regularization À and the means and covariances of a, []. In dimensions (2, 16, 54, 128, 256). we consider Σ,, E2 whose eigenvectors are uniform random (i.e. drawn from the Haar measure on SO(d) and whose eigenvalues sampled uniform i.i.d.
-110
14.5
Figure Entropy regularized, À 2, I cost SCONES and BP samples transportation from unit Gaussian source distribution to the Swiss Roll target distribution. For many samples, the Barycen- tric average lies off the manifold high target density, whereas SCONES can separate multiple modes of the conditional coupling and correctly recover the target distribution.

---

d d=
d = 16
= 64
d 128
d 256
SCONES
0.025 + 0.0014
0.52 + 0.0086
1.2 + 0.014
1.4 + 0.066
2.0 + 0.047
BP
7.1 0.13
35 = 0.23
12:0.14
41 + 0.088
41 0.098
Table 2: Comparison of SCONES to BP on KL-regularized optimal transport between random high-dimensional gaussians. In each cell, we report the average BW-UVP between sample empirical covariance and the analytical solution. We report the average over n 10 independent random source, target Gaussians and the standard error of the mean.
from [1, 10j. In all cases, we set means μ1, M2 equal to zero and choose regularization À = 2d. In the Gaussian setting, B[[]== - =|اي] is of order d. so this choice of scaling ensures fair comparison across problem dimensions by fixing the relative magnitudes of the cost and regularization terms.
We evaluate performance on this task using the Wasseseasnn Unexplained Variance Percentage 44, BW-UV(fi, ~²), where T is the closed form solution given by Janati et al. [8] and where is the joint empirical covariance of k 10000 samples (x, y) SCONES or Barycentric Projection. We train SCONES according to Algorithm and generate samples according Algorithm In place of a score estimate,
generated using either
the ground truth target score Vy y log T rly) = Σg (y- (12) and omit annealing. We compare SCONES samples to the true solution in the BW-UVP metric [4] which is measured on a scale from to 100, lower is better. We report BW-UVP(#1 (ft, >) where ग is 2d-by-2d joint empirical covariance of SCONES samples (ar, v) i or of BP samples (x))),)), æ and T is the closed-form covariance.
6 Discussion and Future Work
We introduce and analyze the SCONES method for learning and sampling large-scale scale optimal trans- port plans. Our method takes the form of a conditional sampling problem for which the conditional score decomposes naturally into prior, unconditional score V y log (y)) and a "compatibility term" v, log (( (x) 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 mrwwardmooee--eee
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 hours using one NVIDIA 2080 Ti GPU. The cost sample score-based models may fall with future engineering advances, but iterative sampling intrinsically require multiple forward pass evaluations of the score estimator opposed to single evaluation of learned transportation mapping.
There much future work to done. First, we study only simple fully connected ReLU networks as parametrizations of the dual variables. Interestingly, we observe that under L transportation cost, parametrization by multi-layer convolutional networks perform equally or worse than their FCN counterparts when optimizing Al Algorithm 1 One explanation may be the permutation invariance of 1?2 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 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 potential synergistic effect between Langevin sampling and optimal transport. Heuristically, as X - 0 the conditional plan Ty|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 known result, that the entropy regularized 2 cost compatibility term ((V(x,,)) = = V(..)//> is log concave function of y for fixed x. It the target distribution itself log-concave, the conditional coupling Tyjx=a 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-concavinn 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.
References
[1] Jean-David Benamou and Mélanie Martinet. Capacity constrained entropic optimal transport, sinkhorn saturated domain out summation and vanishing temperature. 2020.
[2] Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In Amos Storkey and Fernando Perez-Cruz Cruz, editors, Proceedings of the Twenty-Fii International Conference on Artificial Intelligence and Statisties, volume 84 of Proceedings of Machine Learning Research, page 880-889, PMLR, Apr 2018. URL http::/pprcceeding mlr
press/v84 hw8///////eelle0
[3] S. Boyd, S.P. Boyd, L.. Vandenberghe, and Cambridge University Press. Convex Optimization. Berichte über verteilte messysteme. Cambridge University Press, 2004, ISBN 978-0-521-83378-
URL URL httpp::/books..ooogttt/// wtt.oaoooooooooooooooo/o http:/oolsk7iddmmmooodff.f..o
[4] Yongxin Chen, Jiaojiao Fan, and Amirhossein Taghvaei. Scalable computations of wasser- stein barycenter via input convex neural networks. Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research, pages 1571-1581. PMLR, 18-24 Jul 2021. URL .mm.r.g.........mmmmmm..........
httpp:///rooeediings mlr pprsss/v19/ccen////// html
[5] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, ed- itors, Advances in Neural Information Processing Systems, volume 26. Curran As- sociates, Inc., 2013. URL https::/prooeeednngg www.oddiinnnnnn.ns....e. pdf
neurips. cc htt/pappr/20000 /
[6] Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, page 1655-1664. PMLR, Jun 2019. URL hattp::/proceeeeee html
mlr, wwt.prrrss/vv/usee dul1a
[7] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, page 1675-1685 PMLR, Jun 2019. URL
http: http://prooceddings mlr Ttl..ppees/vv// v97/du19c..hmm
[8] Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic optimal transport between unbalanced gaussian measures has a closed form. In Neurips 2020, 2020.
[9] Alexia Jolicoeur-Martines Rémi Piché-Taillefer, Ioannis Mitliagkas, and Remi Tachet des Combes. Adversarial score matching and improved sampling for image generation. In Inter- national Conference on Learning Representations, 2021. URL htttp::/ net/ forum? 1iruumiiddeLLfM1133311
[10] Sham M Kakade, Shai Shalev-Shwards and Ambuj Tewari. On the duality of strong convexity and strong smoothness: Learning applications and matrix regularization. page 10.
[11] Alexander Korotin, Vage Egiazarian, Arip Asadulaev. Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations,
2021. URL http:p///pppaerrrreeeeeer/r/raap .nwt//forumm/dd-Eoxxll
__XXa
[12] Wouter M. Kouw and Marco Loog. An introduction to domain adaptation and transfer learning. arXiv: 18/2. 1806 [cs, stat), Jan 2019. 1812. 1806.
11806, arXiv:
10

---

[13] Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecum lecun.com/exdb/nnisst 2, 2010.
[14] Jacob Leygonie, Jennifer She, Amjad Almahairi, Sai Rajeswar, and Aaron Courville. Adversarial computation of optimal transport maps, 2019.
[15] Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin M Solomon. Continuous regularized wasserstein barycenters. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, page 17755-17766 Curran Associates, Inc., 2020, URL http:/t::ptooceeeeiiiiiiie www.nnnn...uir.....e. wttp:pppccpppppee/000/
file/cdf
pdf
[16] Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, page 15954-15964, Curran Associates, Inc., 2020. URL https:
neurips cc/ Paper. pdf..
[17] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
[18] Giulia Luise, Saverio Salzo, Massimiliano Pontil, and Carlo Ciliberto. Sinkhorn barycenters with free support via frank-wolfe algorithm. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Gamett, editors, Advances in Neural Information Processing Systems, volume 32, page 9322-9333, Curran Associates, Inc., 2019. URL http:s:: pdf
[19] O. Matan. R. Kiang, C. E. Stenard, B. E. Boser, Denker, Denker, D. Henderson, R. Howard, W. Hubbard, L. Jackel, and et al. Handwrit- ten character recognition using neural network architectures. 1990. URL Kiang/
[20] James Melbourne, Strongly convex divergences. Entropy, 22(1111):1327, Nov 2020. doi:
10.3390/e2211 13377.
[21] Sehastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. Advances in Neural Information Pro- cessing Systems, 29, 2016. URL/atpps://rrocedinng
[22] Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Conference on Learning Representations, 2018. URL
wtw.ppenrrrreie.eeettrrrnnn net/forum?
id=Biz1p1bRW
[23] Richard Sinkhorn. relationship between arbitrary positive matrices and stochastic matrices. Canadian Journal of Mathematics, 18:303-306, 1966. ISSN 0008-414X, 1496-4279. doi: 1AA1533/CMM-196-0033
[24] Yang Song and Stefano Ermon. Improved techniques for training score- generative models. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[25] Matthew Thorpe. Introduction to optimal transport. page 56.
[26] C. Villani. Topics in Optimal Transportation. Graduate studies in mathematics. American Mathematical Society, 2003. ISBN 978-0-8218-3311--4 URL http/tttts/:/ookte http:/bbookkggggle..coooo com/ books??d=R_nie hidddd___njjgooo


Reviews:
[{"review_id": "JVekULsX2Yc", "reviewer": "Reviewer_zJ3u", "summary": "The authors propose a method to sample from the regularized OT plan (Sinkhorn plan) via a score-based method and Langevin dynamics.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality**\n\nThe proposed method is original. Nevertheless, the main topic of this paper is OT but the only related paper discussed/compared to in this paper is Seguy et al. (2018). There are already a lot of other methods (primarily based on the unregularized OT), which outperform the BP method of Seguy et al. (2018), some even consider applications to images. I will list some of them in the next sections. Therefore, a comparison with some of them should probably be included. Most of them solve a slightly different problem (recover pure \"unregularized\" OT map), but it seems ok since the generative modeling applications the authors consider require roughly the same.\n\n**Quality**\n\nThere are some questions, comments which probably should be discussed in detail.\n\n(1) The authors use dense networks (potentials) which is a contrast to most works computing OT for images. Why do the authors not use Conv Nets?\n\n(2.1) To train OT plan, the authors use entropy-regularized OT, the optimization of which is known to be practically unstable in the continuous case. For example, Seguy et al. (2018) in Section 5.3 considered exactly L2 regularization, not entropy. Most later works Leygonie et al. (2019), Li et al. (2020), Makkuva et al. (2020), Korotin et al. (2020) which either use regularized OT or compare their method to it, all primarily use L2 reg. as it is claimed to be more stable than entropy. Therefore, I have doubts about the practical stability of the continuous entropy-OT training procedure, especially in the dims which are considered by the authors (64x64x3). My concern is that exp(lambda^-1*(phi+psi-c)) term appearing in training explodes since, e.g., for lambda=0.01, we already have exp(100). Can the authors comment on this?\n\n(2.2)  The main baseline Seguy et al. (2020) considered in this paper is rather outdated (a lot of newer related work exists), i.e. comparison just with it might seem unfair.  Comparison with barycentric projection (BP) baseline by Seguy et al. (2018) might seem unreasonable. BP uses conditional expectation E\\pi(y|x) while this paper considers the entire conditional distribution \\pi(y|x). One may say that it is obvious that the method outperforms BP. Besides, BP is also known to suffer from bias. In README.MD of the official repo of Seguy et al. (2018) there is an image demonstrating the performance of BP in 2D. Even in this toy example, it is seen how different the BP is from the true samples (RED vs. GREEN points). \n\n(3) In the synthetic data experiment, it is not transparent how the ground truth optimal plan is computed. Besides, the values of Frechet distance themself are not interpretable. Could the authors please provide values of its normalized variant BW-UVP, see Fan et. al. (2020)? In general, probably a more thorough evaluation is needed on Gaussian synthetic data (e.g. in various dimensions 2,4,...), to understand the scalability of the method.\n\n(4) How are the samples shown in Figures 1,2 chosen? Are they sequential iterations of Langevin dynamics? Please comment.\n\n(5) As far as I understand, in the CelebA experiment, for each face its downsampled copy is available on the other dataset. Such pairing might simplify the problem. Probably, half of the CelebA dataset should be coupled with the downsampled other half to avoid any inductive bias in data.\n\n(6) Experimental details are given in Appendix, some code is present in the supplementary. However, reproducing the results might be hard due to the high required computational power. It seems to me that the paper lacks \"easy\" experiments, e.g. those which a reader can briefly reproduce on a single GPU or even CPU. Why do I think this is essential? The authors clearly point out that, e.g., Celeba Experiment, requires 4 days on 8xP100 GPUs. Even not all research groups who might be interested in improving/using current work are capable of spending so much computational time on reproducing the experiments.\n\n**Clarity**\n\nThe paper is well written in general. Some places, e.g., Figures 1, 2, could be enriched. It would be nice to see a toy qualitative performance on 2D examples for better transparency, e.g. how well do the samples match the true distribution? This seems to be a necessary attribute of all OT papers nowadays.\n\n**Significance**\n\nThis is probably the first paper demonstrating the application of a continuous regularized OT plan at a truly large scale. The performance of the method is promising. Besides, the method is supported by the convergence/stability theoretical results.\n\nI admit the contribution of the authors, but I think a lot of details should be clarified and a more thorough evaluation to be performed. I set the score based on the current state of the paper but will be happy to reconsider it after the rebuttal phase.\n\n**Additional feedback**\n\nPlease do not use \\mu and \\sigma simultaneously to denote probability distributions. This is misleading for obvious reasons. (!)\n\n**Post rebuttal**\nI raise the score to 5.\n\n**Post rebuttal v2**\nI raise the score to 7, see my last post.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Social impact is ok. About the limitations:\n\nThe proposed pipeline consists of several successive elements (OT plan, score function, sampling) which might be practically complex and hard to set up. This is a limitation and should be emphasized and discussed. Besides, it is unclear how each of these elements performs separately (it is not clearly discussed in the paper). In particular, the quantitative evidence that the procedure actually samples from the true Sinkhorn transport plan requires more discussion.\n\nThe computational complexity is large (4 days of computing on 16 NVIDIA p100 GPUs for CelebA experiment), but this is discussed fairly well.\n\n", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10 hours"}, "extra_text": {}}, {"review_id": "HU3Jag3B1m", "reviewer": "Reviewer_R3r7", "summary": "The authors propose methodology to sample from high dimensional regularized transport maps using Langevin diffusion with neural network parameterized score (via score-matching), and neural network parameterized potentials. \n\nThe authors derive novel quantitative convergence results regarding approximating dual potentials with wide neural nets and the subsequent quality of approximated transport map.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper uses NN parameterized dual potentials and NN parameterized score, both of which are not original, to sample from a transport map, which is the novel application. Arguably more interestingly, the authors derive quantitative convergence rates relating the NN parameterized dual potentials to quality of transport map approximation.\n\nThe experimental results are not very strong e.g. the upsampled CelebA results start from down sampled CelebA rather than noise and are of worse quality than using score-matching to sample images starting from noise. \n\nIt is not clear what the advantages of using score-matching + NN parameterized potentials are over just score-matching for generative models and conditional generation. However, the theoretical results and proof using NN parameterized potentials is very interesting and may be used in other OT related work.\n\nStrengths:\nLeverages existing pre-trained score networks for sampling high dimensional OT transport maps e.g. for conditional generation.\nTheoretical contributions regarding using NN parameterized potentials for OT gives legitimacy to this approach which appears novel and significant\n\nWeaknesses:\nSCONES requires pre-trained score network, yet the method’s performance (for generative modeling) appears to be worse than simply using the pre-trained score network in terms of FID score. \nSimilar to previous score based methods, it requires large neural networks, and many diffusion steps resulting in slow sampling.\n\nMissing references: Parameterizing potential with Neural networks, ICNN [1] or RKHS [2], flows [3]\n\nThe statement of “breaking curse of dimensionality” is tangential to the main work and is only true under strong sub-Gaussian assumption, line 544 in appendices.\n\nThe method appears expensive e.g. CelebA upsampling uses 16 GPUs for 4 days. Is this just for training the potential networks or both potential and score networks?\n\nNo comparison with other transport samplers rather than barycentric average is provided. See e.g. [3], I imagine there are others sampling the transport map rather than averaging.\n\nMinor:\nNotation: \n$\\lambda_\\max$ not formally defined, I assume eigenvalues, line 193\nDo you mean M or H in proof, appendices line 533-542\nTypo using $S_{\\theta_1}- S_{\\theta_2}$ in line 485 of appendices but I assume means $\\theta_1-\\theta_2$ as in proof line 497\n\nCareful with splitting main/supplementary,  I assume these will be corrected when files are combined:\nLatex errors, links/refs e.g.  “??” missing in Appendices in main file\nLink [15] in the supplementary should be [14], line 527\nComment with initials “MD” left in appendices on main file, line 425\n\n[1] Optimal transport mapping via input convex neural networks, \n[2] Stochastic Optimization for Large-scale Optimal Transport, \n[3] Learning normalizing flows from Entropy-Kantorovich potentials, \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "They mention that FID is not as good as regular score matching generative modeling\nThere is limited societal impact compared to existing methods.\n", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "15"}, "extra_text": {}}, {"review_id": "H2WLUubtv_G", "reviewer": "Reviewer_jbX8", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Solving large scale optimal transport has important practical applications, and is of high relevance to the NeurIPS community. This paper proposes a new method to sample from the optimal coupling distribution.\n\nAdvantages of this submission:\n1. Very clear writing. \n2. I am not an expert on optimal transport, but the following theoretical contributions seem to be new: (i) generalizing regularized optimal transport to f-divergence regularizers; (ii) convergence rate of learning neural network based $\\phi_\\theta$ and $\\psi_\\theta$ models with gradient descent (under some rather strong assumptions). \n3. The proposed SCONES method is very general, and can be applied to regularized optimal transport problems with any f-divergence regularizer. When a pre-trained score-model is available, SCONES is also quite efficient.\n\nDisadvantages of the submission:\n1. It is unclear whether samples from Langevin dynamics really capture the optimal coupling distribution in practice. Authors provide some sanity check on synthetic Gaussian distributions, but it is important to check on image problems as well. It is not hard anyways: why not just estimate the transportation cost using samples from SCONES, and compare it with the value of the dual problems obtained after learning $\\phi_\\theta$ and $\\psi_\\theta$ models. Given that authors only used 8 pages out of a 9-page limit, there should be enough space to include such an experiment.\n\n2. Authors need to demonstrate more practical applications of sampling from the optimal coupling distribution. Why would we want to convert MNIST to USPS images and vice versa? Why should we use optimal transport to upsample images? It's not clear to me that optimal transport will provide correct upsampled images and can have any practical advantage over existing methods for superresolution.\n\nSome minor issues:\n1. \"by providing\" repeated twice on line 184.\n2. Some broken references in the appendix.\n\nPost rebuttal update:\nAuthors addressed some of my concerns in the response. I am happy to increase my score to 6.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No obvious limitations or negative societal impact to me.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "AFRX-o2PPr", "reviewer": "Reviewer_6aG5", "summary": "The authors consider a regularised optimal transport problem. They \n- obtained a dual objective in case of a general regularisation based on f-divergence,\n- used \na) the large-scale stochastic dual approach, introduced by Seguy et al., to solve the optimal transport problem in their generalised setting,\nb)  Langevin Sampling to sample from the conditional distribution of Y given X,\n- proved some convergence guarantees.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- There are a number of recent papers with large-scale approaches to OT, e.g.\n1) Jacob Leygonie, Jennifer She, Amjad Almahairi, Sai Rajeswar, and Aaron Courville. Adversarial\ncomputation of optimal transport maps. arXiv preprint arXiv:1906.09691, 2019.\n2) Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev.\nWasserstein-2 generative networks. In International Conference on Learning Representations, 2021.\n3) Ashok Vardhan Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason D Lee. Optimal\ntransport mapping via input convex neural networks. arXiv preprint arXiv:1908.10962, 2019, etc.\n\nThese methods solve un-regularized OT and recover an OT map; they significantly outperform the approach of Seguy et al. (2018). Any comments? Any comparisons? Although the problem statement is slightly different, the comparison can be done using some generative modeling application. Moreover, it is easy to outperform the approach of Seguy et al. (2018), since it uses the barycentric projection (BP) baseline which is weak compared to using the entire conditional distribution pi(y|x).\n\n- In general, the text of the paper is rather well written. However, some details are missing, see comments below.\n\n- Can the authors prove their theoretical result about the convergence for CNNs? It seems not optimal to use FCNN for images, so CNNs are needed. The authors mentioned that the assumptions of the theorem are valid for FCNN. What is about CNN? Moreover, could the authors provide some more visual results of the experiment with Celeba? FCNN are not efficient for such high dimensionality (64x64x3), so it could be good to analyse failure cases.\n\n- In case of experiments with synthetic Gaussian data I would like to see how the method scales with the dimension. Moreover, the authors provide the value of the Frechet distance. However, to understand whether they are big or not we need some characteristic value or normalisation. So the current representation of the results does not allow to comprehensively analyse them.\n\n- The authors used the Langevin dynamics to generate images. Could the authors provide some intermediate results on how the dynamics converges?\n\n- Protocols of experiments are not fully detailed, see comments above. Also, in case of the experiment with the Celeba dataset, how the samples X and Y were organised? For each image from X there is a high-res image of the same face in Y? If yes, this simplifies the problem. Also, I would prefer to see more  results of the corresponding super-resolution task, not only values of the FID score and some faces in Figure 1. Also, why not to include more toy examples and illustrate the convergence?\n\n- The results on Celeba are promising, they demonstrates that we can construct OT rather efficiently even for rather high-res images. \n\n- Overall, the results are promising and some theoretical guarantees are proved, but more details about the experiments are needed.  Also, comparisons with recent method are missing.\n\nI need more details/comments to upgrade/downgrade the current rating.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The approach consists of several steps, and it is important to discuss limitations of different steps separately, it is unclear how these different steps perform separately.", "review_text": "", "extra_scores": {"ethical_concerns": "not applicable", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.0"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a method to sample from high dimensional regularized optimal transport maps using Langevin diffusion on neural network parameterized potentials. The reviewers compliment the authors on being some of the first to tackle this problem at a truly large scale. The problem is important and the proposed solution seems technically sound and practically promising.

After the first round of reviews some concerns surfaced about the correctness of some of the presented results. In addition to general concerns about experimental validation from multiple reviewers, reviewer zJ3u pointed out some irregularities in the reported BW-UVP values which turned out to be caused by a scaling error in the regularization used in some of the experiments. The authors responded by providing updated results, in addition to acknowledging and fixing the error. The reviewers are satisfied with the author response and are now unanimous in their recommendation to accept the paper.

Authors, please carefully revise your paper for the camera ready version, and make sure to incorporate all new results and all changes that were promised in the discussion with the reviewers.
Decision: Accept (Poster)