OTCOP: Learning optimal transport maps via constraint optimizations
Authors: Xiaokai Huo
Hailiang Liu
Venue: Submitted to ICLR 2023 (2023)
Keywords: optimal transport
constraint optimization
Monge problem

Abstract: The approximation power of neural networks makes it an ideal tool to learn optimal transport maps. However, existing methods are mostly based on the Kantorovich duality and require regularization and/or special network structures such as Input Convex Neural Networks (ICNN). In this paper, we propose a direct constraint optimization algorithm for the computation of optimal transport maps based on the  Monge formulation. We solve this constraint optimization problem by using three different methods: the penalty method, the augmented Lagrangian method, and the alternating direction method of multipliers method (AMDD). We demonstrate a significant improvement in the accuracy of learned optimal transport maps on benchmarks. Moreover, we show that our methods reduce the regularization effects and accurately learn the target distributions at lower transport cost.


Under review as a conference paper at ICLR 2023
ОТСОР: LEARNING OPTIMAL TRANSPORT MAP VIA CONSTRAINT OPTIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
The approximation power of the neural network makes it an ideal tool to learn optimal transport maps. However, existing methods are mostly based on the Kan- torovich duality and require regularizations and/or special network structures. In this paper, we propose direct constraint optimization algorithm for the com- putation of optimal transport maps based on the Monge formulation. We solve this constraint optimization problem by using three different methods: the Lan- grangian multiplier method, the augmented Lagrangian method, and the alternat- ing direction method of multipliers (ADMM). We demonstrate a significant accu- racy of learned optimal transport maps on high dimensional benchmarks. More- over, we show that our methods reduce the regularization effects and accurately learn the target distributions at lower transport cost.
INTRODUCTION
There has been great interest in applying modern machine learning techniques for finding optimal transport maps between two distributions. Different from traditional computational methods that solve PDEs for optimal transport maps (Benamou & Brenier (2000); Angenent et al. (2003): : et al. (018 ). modern machine learning techniques aim solve the problem directly by optimizations. The Sinkhorn Distance method Cuturi (213);; Peyré et al. (2019), the regularized OT dual Seguy al. (2017) have been used to find large scale optimal transport maps between discrete probability distributions and have been used to train generative networks (Genvva et al. (2018); (Sanaabi et al.
2018). geometric treatment provided in Gu et al. (2013). The Input Convex Neural Network (ICNN) is used to construct convex Brenier potential for finding optimal transport maps Makkuva cct al. (2020) between continuous distributions and is recently used in population dynamics Bunne al. (2022), which combines the ICNN and Sinkhorn distance methods Amos al. (2022). Despite these successes, most methods are based on the duality formulation and avoids the direct treatment on the Monge problem.
In this paper, we focus on the direct solution of the Monge problem. The Monge problem (Monge (78T)) directly seeks identify the optimal transport maps and is a nonlinear constraint optimiza- tion problem. The major difficulty in solving the problem numerically is that it is nonlinear and includes a constraint that the push-forward distribution is equal the target distribution, which is difficult to implement. Therefore, most optimal transport algorithms avoid directly solving the Monge problem but use the Kantorovich duality (Kantorovich (1422), ), for which the objective func- tion is linear and the transport map is obtained by taking the gradient of the Brenier potential for the quadratic cost. However, these two problems are not always identical (Villani (2009) and is desirable to find direct approach for the Monge problem.
The Monge problem has been solved numerically using optimization based methods with polyno- mial approximations. For example, a Lagrangian penalty method was used to find optimal transport maps approximated by polynomials for Bayesian inference Moselhy & Marzouk (2012) and space discretization was used in Haber et al. (2010) to calculate the Jacobian matrix of the transport maps and transferred the optimization to finite dimensional spaces. However, their approaches are limited number of grids expands exponentially dimensions become large. Consid-
to low dimensions ering the success of deep neural networks in approximating high dimensional data, the integration of classical constraint optimization methods and neural networks holds promise.

---

Under review as a conference paper at ICLR 2023
One successful application of the optimal transport theory to deep learning is the Wasserstein Gener- ative Adversarial Network (WGAN) Arjovsky et al. (2017). However, WGAN only use the optimal transport distance loss function and does not target finding the optimal transport maps. It is desirable to study whether it is possible to lower the transport cost the map learned by WGAN or other networks using the algorithm for finding optimal transport maps.
This paper presents a new approach for finding optimal transport maps between two continuous distributions. We make the following contributions:
We integrate three constraint optimization algorithms including the Standard Lagrangian (SL), the Augmented Lagrangian method (AL) and the Alternating Direction Method of Multipliers (ADMM) with neural networks solve the Monge problem of optimal trans- port with provable guarantees (Theorem
We show that our method is able to find an accurate optimal transport map between Gaus- sian distributions, both theoretically (Theorem and experimentally, Moreover, we apply our method to WGAN and show that our method can find a generative map with lower transport cost while not sacrificing the quality of outputs.
We compare the three algorithms and find the SL algorithm introduces errors but is simple and easy to implement, while AL and ADMM algorithms can find exact results and are more robust, and ADMM gives a lower transport cost in general.
Notations. We use the notations Oa (a..... ,,) Rd and aaxd for the constant d d matrix. The transport cost of map T, which pushes distribution to v, is defined to be Ez~ [|x - T Tx]
2 THE MONGE PROBLEM AS CONSTRAINT OPTIMIZATION
2.1 THE MONGE PROBLEM
Let (X, p), (Y.v) be two separable metric probability spaces. The Monge problem is to find transport map T': Y that realizes the infimum
inf {{ clxr, Tr))d(()) T# = )
where TT denotes the push forward of μ and c X -> RR,, is a Borel measurable function which is lower semicontinuous. In this paper, we simply take the distance |x y1 but our method applies to other distance functions.
The existence the Monge problem is difficult and does not hold always. However, under suitable conditions, for example for continuous distributions without atoms, the existence and uniqueness of the Monge problem is guaranteed (see for example, (Villani 2009. Theorem 5.30). Therefore, here we focus on learning transport maps between continuous distributions. For discrete distributions, one can apply dequantization techniques to transform them to continuous distributions Ho et al. (2019).
2.2 THE MONGE PROBLEM AS CONSTRAINT OPTIMIZATION
In order to solve the Monge problem, set 6 which inputs random samples from the distribution and generates samples representing the target distributions v. As can be seen from the definition, the Monge problem is a constraint optimization problem. However, the constraint T = D/ is a highly nonlinear constraint. In order impose this constraint, we take d--l-) be distance function (such as the Wasserstein distance, the MMD (Gretton et al. (2012)) or the IPM (Müller (1997))) or a probability divergence (such as the Kulbback-Leibler (KL) divergence). The constraint optimization problem reads
generative network, denoted by To with parameter
min Ez || ]
s.t. o(Tag plv) 0.
The objective of this paper is to solve the above problem using techniques from the constraint opti- mization theory (Bertsekas) (0044).
2

---

Under review as a conference paper at ICLR 2023
Since neural network may not fully reveal the target distributions the above problem can be relaxed to
min Eaz [|x Toor
s.t. d(To#1ll)
(3)
When a goes towards zero, we can prove that the solution Te, to the problem [3] converges the The following theorem holds:
solution of the original Monge problem [
Theorem 1. Let μ., V be two probability measures on R2 with finite second moments and are ab- solutely continuous. Let To be given by a neural network with bounded width(each with at least 2d + 2 neurons) and arbitrary depth, and with non-affine activation functions. Suppose for any > O, there exists solution 8 to problem [2] then as 0, Te= - 'T where T is solution of the Monge problem [ Moreover, suprex (T) - T(x)le & Co some constant C' for any
compact subset X E R
Proof of the above theorem follows from the universal approximation theorem (iidger & Lyons
2020 and the existence theorem of the Monge problem (Villani (2009)) and is given in Appendix
A.
2.3 EXAMPLE: THE MONGE PROBLEM FROM GAUSSIAN TO GAUSSIAN
For the when j NV1,, Σn) and (X2,, Σ2) two multivariate normal distributions, the optimal transport map is unique and can be explicitly given by T : H X2 - A (x - -1) with A" Σ (2) Eg | / 1/2 Olkin & Pukelsheim (1822). Taking d DKL be the KL- divergence, prove that the solution problem (2) with T&x Ab (0 {A e Rdxd b Rd))
is the optimal transport map T (Theorem 2 in Appendix A.1
3 CONSTRAINT OPTIMIZATION FOR OPTIMAL TRANSPORT
We propose to leverage three different algorithms to solve the constraint problem (2)
3.1 PENALTY METHOD (OTCOP-P)
Standard Lagrangian (SL). We introduce Lagrangian multiplier À and take the Lagrangian function
as
CsL A) Ez-p [|- Toxl²] + ATT
(4)
Then the solution to the problem [2 is saddle point of the above Lagrangian. By duality theory, for each a 2 O, the problem [3 corresponds to the duality problem ming Cs1 X, 0) for a [0, oo].
Hence we can take suitable À to solve the problem (2) approximately.
According to the Brenier's polar factorization theorem Brenier (1991), the optimal transport map should satisfy V T O, hence we can add an additional term V x T into the above Lagrangian to impose this constraint. We will show experimentally that without this term, this constraint is almost satisfied and will not included in implementations.
Quadratic penalty (QP). Instead of taking d (TT (v). we can take a quadratic penalty loss
Cor(0. p) Eyww [| Tox [2] ((((TTo
(5)
As p goes towards infinity, the constraint violations is penalized with increasing severity. For exam- ple, we can take Pk at the ktt training step to be increased by multiplying by a constant bigger than 1 and parameter 0 can be updated using gradient descent considering p constant.
Convergence. Suppose there exists a global minimizer to the problem 22 and 02 is the exact mini-
mizer of (Q, pk)) and pk 1 co. Then any limit point of the sequence [0x) } is a solution to problem
() Moreover, for any >> 0, there exists a sufficient large K 0, 16x 0-1 < for k 2 K (see Nocedal & Wright 1999. Theorem 17.1)). In addition, without assuming global minimizer, for sequence 8g such that Vo, (op((::)))) (k) + 0, its all limit points 8 satisfy the Karush-Kiut Kuhn-Tucker (KKT) conditions and there exists subsequence such that limg- (pt (pad(To, oo (v) = X", where X
3

---

Under review as a conference paper at ICLR 2023
is the multiplier that satisfies the KKT condition (see Appendix A.4 (444l for the KKT condition and see (Nocedal & Wright 1999. Theorem 17.2) for the proof).
Advantages and disadvantages. The penalty method is simple and easy to implement. However, since the optimal value of the Lagrangian multiplier À is unkown (SL) or the optimal condition for the Lagrangian multiplier p is infinite (QP), the penalty always introduce errors and the exact solution to the problem (2) cannot be reached. Moreover, the Hessian of the Lagrangian oo becomes singular as P goes towards infinity and cause ill-condition problems. These issues can be solved by the methods below, but at the expense of a more computational cost.
3.2 THE AUGMENTED LAGRANGIAN METHON (OCTOP-AL)
In order to overcome the above issues, we can use the augmented Lagrangian method, by taking the loss function as
(O A, p) = Exnn ||0 - Tox: (T (o(Touul)))
(6)
The above function combines standard Lagrangian penalty [4]and quadratic Lagrangian penalty (5} At the kth iteration, fix de. Pl. and solve 0g arg ming Aar (8, Pk.P)).
Pt). After the minimization, we update À by A = As + Pr. (ule). Comparing the KKT conditions of the SL and the AL (see Appendix A.4p implies A + pa (To
X when de is taken close to A Hence the quadratic penalty method that + the infeasibility in 0s will be much smaller. Moreover, for certain choice of the local solution of 2 is strict local minimizer of LAL (0, X, ((Nocedal & Wright 1999 Chapter 17)).
#TTa x (X* x*) /pp. Compared
X//p,
Convergence. One of the nice properties of the AL. method is that for the exact I Lagrangian multiplier X" the solution 8* of the problem 22 is strict minimizer of CAL (0, X X*)p) for all p sufficiently large. The existence of a threshold Is proved under the condition that VøCsz (0*,**) is locally strictly positive ((Nocedal & Wright 1999 Theorem 17.6).. Thus we can take p to be increasing
at each minimizing step and when becomes bigger than some threshold value P. gradient descent methods could find the local minimizer around 0*.
Advanluges and disadvantages. AL method introduces the multiplier estimates and reduces the likelihood that large values p will be needed to obtain good feasibility and accuracy. The method is also simple and easy to implement. However, since this is a min-max method, training may experience oscillations and slower convergence rates.
3.3 ADMM METHOD (OCTOP-ADMM)
The ADMM method blends the decomposition techniques and the AL method and provides an efficient way for constraint optimizations Boyd et al. (201I). Let S be the set S = (To Taw v). problem 22 can be rewritten into the form
min Eapep [|x - Tox|3] Is(To),
where 1 is the indicator function that equals 0 Ta E S and equals co if & S. In order to apply the ADMM method, we rewrite the above problem into the form
min Eg [|x 0,,02
x1 (T(T))),
Top
Tog
In the ADMM method, we alternatively update 02 and 02. First we take 02 to be constant and take the minimization of the above problem over 01, and then we project Og onto the space S. In detail, introduce the loss function
LADMM 10,, Oz p, Exrn [| - x? + 1 10(Tb
+ +A (To, =T,,2) P { S To, 21).
Here A E Rd is the multiplier. The training procedure is given by
1. g/+11 arg ming, LADMM 03, A*, A..p) (assuming la(To = 0):
4

---

Under review as a conference paper at ICLR 2023
2. 62+1 arg ming, oToo
3. A&+1 = A₁ + pl v (T (T (x) T Totti (z)).
Convergence. The convergence of ADMM method only known to hold under convex conditions or for some non-convex problems (Boyd et al. (2011)). Using results of Wang et al. (2019). we can prove the convergence of the ADMM method [8] to
modify the above method by relaxing problem
min Ea [1> - To, x|*] me (o(Te, S)), s.t. 81 82.
(10)
where T7 is the mollifier function which converges to the 8-function as E - O. We show that the ADMM method converges to the KKT points and if the corresponding Lagrangian is Kurdyka- Lojasiewicz function, the ADMM method converges globally to the unique solution (see Appendix A.3 for details). As consequence, there exists convergence sequence of the ADMM method for the problem approximately.
[
Advantage and disadvantages. The advantage of ADMM is that it decomposes the Monge problem into two sub-pna uub-problems: minimizing the transport cost and minimizing the DKL. Compared to the AL method, AMDD solves two decomposed minimization problem and AMDD may also converge faster than the AL method in some situations (Wang et al. (2019). However, the method also solves min-max problem and may facing oscillations in the training process.
3.4 THE DISTANCE
We need a distance/divergencee functional d to compare the generated distribution and the target distribution. Here we require to be able compute using samples.
The KL divergence DKL. When the densities of the target distribution is known, we can use networks of normalizing flow (Rezende & Mohamed (2015) as Ta and the DKL. computed via
(KL(To+Ulp) = Ezz (log ((x) log det Jr, (I) - log (Taz))),
(11)
where Jrg is the Jacobian of the transport map which is able to compute using the normalizing flow networks. When the target densities are unknown. We need to reverse the KI. divergence and take
(KL(v[To#H) = E (log v(x) - log (TT x) log det Jr. (T x)]]
We can drop the first term in the bracket since it is constant. We can take LST. in (4) as
CSL X) = Expp [|x - Tex: [
AE log (uTT x) log det Jr, (T z)].
For the ADMM method, the second step changes to 0 + arg ming En |- log (T," x)
log det Jr, (T-1 In order to apply the AL method, one can use En --V, log μ (T Vo log det Jr, (T (T-1z)] to replace the DKL terms in [6
Test function as multiplier weak form of the constraint Tow v is that for any measurable function f.
I ((Tor))p((s) f f(r (())du(x)
we can introduce a discriminator network fur as Lagrangian multiplier and the Lagrangian becomes
(s w) Ex~pp/x To + Eewz (fr [f
(12)
The augmented Lagrangian then become
(AL(0, w. p) Excu//x To 12] le (Tox)] - Beanifull (z)]
+ (Bon fo (Tox)] Bew (fu)))
(13)
The training procedure is given by
1. g arg ming (9, w. p):
2. #+1 w + Br V w (Tox)) - (
5

---

Under review as a conference paper at ICLR 2023
3. Assign p pp
To apply the ADMM method, we take the loss function to be
LADMM (0,00,,, 1,, W2. P) = Ex-ullx T8, Dw (Toz#1lu)) +B Ifu x) fw, (To, x)] +
P
2 Berulfer (To, x) fwo (To (14)
Then the optimization is decoupled into optimization over 81 to learn the optimal transport map and the optimization over O2 learn the target distribution. Here can take for example Dwy be the loss function of Wasserstein GANs with gradient penalty (Gulrajani et al. (2017)):
O (1))) = Ee~ulf [ (Ter)] - Es-w V fw. (z)) + clea [[I fw (4)ll2 1)²}.
where e is positive constant and is the linear interpolation of To and v. The training procedure given by
1. 0 arg ming, LADMM(PI, 022 wi, wx, p²):
2. (0h+1 2222 1 arg ming, max LADMM(O O2 wi, u2, p²) updated same as
WGAN with gradient penalty.
3. w
k+1 wf Bl V ws fwo (Toz)] EEnn Vwr (z)]
3.5 IMPLEMENTATION OF THE ALGORITHMS
The implementations of the algorithms are given in Algorithm 2, Here we give the implementa- tion when given directly without using discriminator network. For the case when test function is used as multiplier, see section 3,4 and Appendix ?? for details.
Algorithm 1 Solving the Monge problem with the penalty method
Algorithm Solving the Monge problem with the aug- mented Lagrangian method
Input Data: X Y v. Constants: Ao Po.
Input Data: X μ4 ~v
1,
Constants: Ao Pu. >
Training step: 7)-
Training step: 17.
for number of training iterations do
for number of training iterations do
for m steps do с 0 V
for m steps do 0 t- V eSL
end for
(or 0 1) V (QP
Update À + A po(To+Hll)
end for Update p + ap. end for
Update + op.
end for
Algorithm 3 Solving the Monge problem with the ADMM
Input Data: for number of training iterations do
p, Y v, Constants: λo Po. a > Training step: 17.
for mi steps do
0 t 01 COAMMM
end for
for m2 steps do
Bg - 0g - o (
W2 w2 1/ (To,+ (if Dw)
end for
Update X X pdpTo (T, ( To.HH)
end for

---

Under review as a conference paper at ICLR 2023
4 EXPERIMENT
4.1 MULTIVARIATE NORMAL DISTRIBUTIONS
Linear maps. First, consider the optimal transport between multivariate normal distributions = (X1 ) and v/ == [X Mn, Considering the map given by Ta = Ar + b with 0 = { e pddd b E [& J. we prove in Theorem that the solution to problem [2] gives the correct solution to the Monge problem. Indeed, in this case, the SL, AL and ADMM algorithms reduce to optimization of linear objective function constraint by nonlinear function and it could be theoretically analyzed by the constraint optimization theory [Bertsekas] (2014). Here we take X1 02, ΣΙ = I2 and [a, b), (b, al] (0 = {a, b)). Then the problem
X2 02. Σ₂ = [[4, 1). [1, 4]], let Tax- Ax with
reduces to
min 2(1 a)² + 21/, a.b
1 (a² (a2 63)2) + Sa² 4ab + 862 - -30) 0. 15
s.t. DRL DRL(TTply) 5/0 (-15 log
The landscapes of DKL and L (0, 1) as well as the value of ming CSL (0. A) as functions of X, are plotted in Figure [ As can be seen from the figure, the function DKI.(Topp (ule/) has multiple minimizers (red points), whereas the Lagrangian Lsr(0, 1) has a unique global minimizer (blue point). Hence minimizing the Lagrangian LSL helps to find the optimal transport maps by finding the approximate map with the lowest transport cost among all maps that realize the target distributions. More importantly, from the figure, the whole domain by the landscape of DKL, and starting in each piece, gradient descent method will converge to one of the four different points. In contrast, the landscape of the Lagrangian changes dramatically and starting from any point, gradient descent method will only converge to one point.
be divided into four pieces
The Langrangian introduces errors for finite. As can be seen from the leftmost figure, the constraint DKL. = 0 is more relaxed as becomes smaller. However, for large X, the barrier region becomes wider and the gradient descent may converge a local minimizer away from the optimal result. For small A, the gradient descent is easy to find the global minimizer, but since penalty introduces error, the relaxation effect of the constraints can also push the minimizer away from the solution of the Monge problem. The choice of À needs to take this tradeoff into considerations. Training using one layer linear neural network confirms the above analysis.
Figure 1: Graph of the KL error and the loss function of the SP method (left: ming Caz A) as well as the corresponding DKL and transport cost as functions of x; middle: landscape of DKL. as functions of a, b; right: landscape of LSL (6, 1) functions of a, b). Red points: minimizers of the DKL. blue point: minimizer of CsL 1).
Training with neural networks with nonlinear activation function. Next we present our results on the training of Gaussian Gaussian distributions with neural networks with nonlinear activation function. We take X, = 0a - 1a and Σ, Ia, Eg 3Ta + 1dxd- Theoretical analysis gives that the optimal transport distance is 2d (see Theorem 2 Since here we use the DKL as distance function which is always positive, so the QP method behaves similarly as the SL method with different multiplier. Hence the results of the QP method is not presented here. The results are given in Table and the 784D Gaussian is taken for X; 0781, X1 = 2- 1744 with Σ₁ Σg I784 (theoretical result of the optimal transport distance is 784 4). As can be seen from the table, all three algorithms give nice result and learn approximately the optimal transport map with high
7

---

Under review as a conference paper at ICLR 2023
Table I: Training results Gaussian and Gaussian mixtures
benchmark
method
DRL.
Transport cost/d
benchmark
DEL.
Transport cost/d
2D Gaussian
SL
0.002
2.018
78D Gaussian
0.228
1.866
AL
0.002
2.021
0.260
2.290
ADMM
0.003
1.950
0.805
2.250
784D Gaussian
SL
0.365
3,981
2D mixture
0.034
0.048
AL
0.333
4.001
0,021
0,066
ADMM
0.399
3.998
0.059
0.035
accuracy. Training is done by using d width and 10 depth neural network with tanh activation for all cases except for 78D Gaussian, for which a 100 depth neural network is used in order to learn the correlations correctly. Remarkably, the highly correlated distributions high dimensions, our method gives a nice result.
ADMM learns a lower transport cost. From the figure, we can see that the ADMM method learns lower transport cost compared to other methods. This is because during training, minimizing the DkL. between the generated and target distribution may not converge to the solution to the Monge problem, illustrated in the simple 2D case above. Hence, the splitting feature enables ADMM to learn a transport map with lower transport cost. This can also be seen from the learning curve in Figure 4 in Appendix A.5 The transport of the second network (Te,) has higher transport cost when learning the target distribution, whereas the first network (To, has a lower transport cost, while the learned target distribution remains accurate.
4.2 GAUSSIAN GAUSSIAN MIXTURES
We take four component Gaussian mixture, each with variance matrix 0.5Lz, centers lying on the four corners of the square [-1, 1]² and learn the optimal transport map between two dimensional standard Gaussian to this mixture distribution. We plot the Jacobian graph of the learned map and the value of the DxL. between target distribution and the prediction by the network in Figure 2 As can be seen from the figure, solely minimizing the DKL does not fulfil the right directions of the optimal transport maps. Using our method, the learned map is more balanced and approximately satisfies the condition V Txz = O. Note that here we donot include the penalty of V Txx From the learning curves, we confirm the findings above that ADMM learns a lower transport cost than the other methods. Here the transport cost of the second network converges to around 0.07, while the transport cost of the first network converges to about one half of that of the second network.
Figure 2: Graph of the KL emro and the loss function of the minimization of DKL. min DRL. the SP, AL and the ADMM method (from left to right). The DxL curve the decreasing line and the transport cost is the increasing line. For the bottom right figure, orange line is for the first network and blue line is for the second network
AL and ADMM methods more robust. Compared the SL, AL gives with respect to the value of X. For a well chosen X, SL performs as good as AL. However, is not taken properly, the obtained transport cost will be higher or the target distribution is not well
robust result
8

---

Under review as a conference paper at ICLR 2023
realized. For ADMM method, we can see the choice of p affects the training process, but in big range, the choice of p has little effects on the final result (see appendix AA.5 for the graphs indicating this finding). This confirms the benefits of the AL and ADMM method in the literature (Nocedal & Wright 1999); Boyd et al. 2011).
4.3
IMPROVED TRANSPORT COST OF GANS
We the test function as multiplier and test the performance of the algorithms described in section 3.4.
The SL method help overcome the vanishing gradient issue. The constraint fw is Lipschitz-1 function is important in WGAN. Without this constraint, the training of WGAN may face vanishing gradients issue, illustrated in Figure Appendix A.66 However, adding the transport cost the cost function as in [12) the training no longer face this difficulty and the target distribution could be learned. Note that here donot use a penalty term on the discriminator network as the gradient penalty method, the transport cost is only function of the generator network and no regularization of the discriminator network is needed.
All three methods significantly reduce the transport cost of GAN. Compared to the WGAN-GP. which has a transport cost around 3.68, all three methods (SL, AL, ADMM) described section 3.4 show a significantly lower transport cost (0.077 for SL., 0.071 for AL and 0.049 for the ADMM). The transport losses are plotted in Figure 7 in Appendix A.6 However, the SL method sensitive to the choice of X, while the AL and ADMM methods more robust.
MNIST. We also train the WGANs on the MNIST dataset. By using WGAN with SL(* = 1), obtain a transport cost around 1.81 compared to 1.87 with only WGAN. MNIST like samples generated by the learned optimal is plotted in Figure 3 Therefore, our method lowers the transport cost of WGAN while keeping the qualify of the generated distributions.
0 9
/
2
5 5
5 5
6
O
Q Q
Q
Q 2 8 8
0 2
5 3
6
6
Figure 3: MNIST like samples generated by WGAN with SL penalty
CONCLUSION
We have shown that the incorporation of constraint optimization tools provides a direct and efficient way for computing optimal transport maps. By solving the Monge problem directly, our method avoids using special network structures or solving the dual problem. Moreover, applying our method to WGAN shows lower transport cost for the generative networks without sacrificing the quality of the generated data.
REFERENCES
Brandon Amos, Samuel Cohen, Giulia Luise, and levgen Redko. Meta optimal transport. arXiv preprint arXiv:2206.0522 052662, 2022.
Sigurd Angenent, Steven Haker, and Allen Tannenbaum. Minimizing flows for the Monge- Kantorovich problem. SIAM Journal on Mathematical Analysis, 35(1):61-97, 2003.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pp. 214-223. PMLR, 2017.

---

Under review as a conference paper at ICLR 2023
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375 (09)3335-33, 2000.
Dimitri P Bertsekas. Constrained optimization and Lagrange multiplier methods. Academic press, 2014.
Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, Jonathan Eckstein, et al. Distributed optimiza- tion and statistical learning via the alternating direction method of multipliers. Foundations and Trends@ in Machine learning, 301::1112, 2011.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Commu- nications Pure and Applied Mathematics, 44(4)3375-417, 1991.
Charlotte Bunne, Lactitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal trans- port modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pp.
PMLR, 2022.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013.
Tarek A El Moselhy and Youssef M Marzouk. Bayesian inference with optimal maps. Journal of 18-775008 2012.
Computational Physics, 23 (23))2778 5-7850,
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn di- vergences, In International Conference on Artificial Intelligence and Statistics, pp. 1608-1617.
PMLR, 2018.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, (31))723--73,, 2012.
Xianfeng Gu, Feng Luo, Jian Sun, and S-T Yau. Variational principles for minkowski type problems, discrete optimal transport, and discrete monge-ampere equations. arXiv preprint arXiv:1302.5472, 2013.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Im- proved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
Eldad Haber, Tauseef Rehman, and Allen Tannenbaum. An efficient numerical method for the solution of the 1.2 optimal mass transfer problem. SIAM Journal on Scientific Computing, 32(1): 197-211, 2010.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. In International Conference on Machine Learning, pp. 2722-2730. PMLR, 2019.
Leonid V Kantorovich. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS), volume 37, pp. 199-201, 1942.
Patrick Kidger and Terry Lyons Universal approximation with deep narrow networks. In Conference on learning theory, pp. 2306-2327. PMLR, 2020.
Wuchen Li. Penghang Yin, and Stanley Osher. Computations of optimal transport distance with
fisher information regularization. Journal of Scientific Computing, 75(3):1581-1595 2018.
Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lec. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672- 6681. PMLR, 2020.
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pp. 666-704, 1781.
Alfred Müller. Integral probability metrics and their generating classes of functions. Advances in Applied Probability, 29(2):429444 1997.
10

---

Under review as a conference paper at ICLR 2023
Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 1999.
Ingram Olkin and Friedrich Pukelsheim. The distance between two random vectors with given dispersion matrices. Linear Algebra and its Applications, 48:257-263, 1982.
Gabriel Peyré, Marco Cuturi, et Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 116-6))35556077 2019.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. Interna- lional conference machine learning, pp. 1530-1538. PMLR, 2015.
Maziar Sanjabi, Jimmy Ba, Meisam Razaviyayn, and Jason Lee. On the convergence and ro- bustness of training GANs with regularized optimal transport. Advances in Neural Information Processing Systems, 31, 2018.
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv:1711.02283. 2017.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Yu Wang, Wotao Yin, and Jinshan Zeng. Global convergence of ADMM in nonconvex nonsmooth optimization. Journal of Scientific Computing, 78(1):29-63, 2019.
A APPENDIX
A.1 PROOF OF THEOREM 1
The proof of Theorem [ is given below by combing the existence theorem for the Monge problem and the universal approximation of neural networks.
Proof. First given two probability measures μt, satisfying the assumptions of Theorem [| then VVillani 2009 Theorem 5.20) implies that there exists unique deterministic solution 7 to the Monge problem and the map T is Borel map. Then by Lusin's 's theorem, there exists continuous function T e C(R') such that T T any compact subset X' Rd almost everywhere. Since Tp is given by a arbitrary depth neural network, we can use the universal approximation theorem in Kidger & Lyons 2020 conclude:
For any 0, there exists T3 such that supsex [T((x) T Sc.
Therefore, D(T (TTp)l))) = D(T S aT #T)) D(T#1u|##1)) S 2E.
Fix a and take = < a/2, since ( S a. we can get that the solution To: to the problem satisfies
(3p
E B Tog x Ex~all> - Toal < Bewwll Tx ? +
(15)
Let { } be a sequence converging to zero and To, be the solution to the corresponding problem then o(To, #ulp) I as + 0 implies To, To. To # Twm. Since T solves the problem
we have
Ex>μl|> - Tax 2 < Bawull [| To- 20
On the other hand, by taking limit - in 15 we have Pem-pll00000 To- 7002221 2] S Enzzll
| ||
Tox 22
Combing this with the previous inequality implies that Exen | - Tox [2] = fll T
2
By the unique existence of the Monge map, To- T. Moreover, we have by the convergence of To. + Tg-
E Egen, [[ T Enrn [Tor_ x To-x1 Ca.
(16)

---

Under review as a conference paper at ICLR 2023
A.2 PROOF OF THEOREM 2
Theorem 2. Lel μ = [Xo Eo) and v 1, E1) be two multivariate normal distributions
with Xo, X1 IR and En, Eι Raxd Let To : Rd + Bd be an linear operator defined by Ax + b with = (A, and A E Raxd, E Rd_ Then there exists a unique solution to the
following problem
{ 100 Tox du(x), dd))x)
s.1. DRL (KT(TTHull)) - -o].
0 = arg min
(17)
and the corresponding transport map Tg. is the optimal transport map the sense of
Proof. For μr, the linear transformation y Ax also satisfies multivariate normal distri- bution, given by
y y~p- = (AXx, + + AEaA").
Using the formula for KL divergence between multivariate normal distributions, we can get that the KL-divergence between and is
DKL(Plu) = ACoAT) + (X1 (AXo + +) (AXo - b))
+ log det dct(AEg,A7)
(18)
det Σ1
Due to the fact that trB - log det B + 2 0 for any positive definite matrix B > 0 and (X1
(AXo + 2)) (X1 - (AX, + b)) 2 0, (L(p))) is equivalent to
Σj AE = Id. (X1 - (AX, + b)) = 0.
(19)
The objective function in can be calculated by
[(e. |r - Toxx du(())
Jec 1 - (Ax + 6))2d
Ernp [[(I- -0)= - b)) ((I A)x b]]
=tr((T - A)² (I A)B (Nx - 21 (I A)Xo 6
-rr((T A) (T A) + Xoxx -2b² (I A)Xo + 6=b
-tr((1 - (r((I-))(( A)Eo) + +1(1 - 612
(20)
Therefore, the optimization problem [17] becomes
min tr((1 -A) A)Eo) + [1 - A)Xo - -b2,
AER BERd
s.t. AEAT la (X1 (AXo + b)) 0.
(21)
Eliminating b, this is equivalent to
min AERdxd
tr((1 - A) A)Eo) +|X - -2113
s.t.
Σ AE AT Ia-
(22)
Since tr((1 - (I A)E,) = u<Eo)) - Atr(ATIX + ir(AEgAT) = tr(Eo) 2tr(A Eo) + d,
the above problem is equivalent to
max tr(AT (T(AT)o). Eo), Aeaaxa
s.t.
AE = E1.
(23)
Let R = 2 ATM the above problem could be rewritten
RE tr( R 1
(24)
12

---

Under review as a conference paper at ICLR 2023
Note that the optimization is all orthogonal matrix in R The solution is given by
R v where Σ = VP is the polar decomposition of Σ m and is given by V イ and P ( و ) Recalling R
Σ ( ('') we obtain
= A = E {
(25)
and by the constraint (X1 AXo - b) 0, we get
A Xo X1.
(26)
The above formula (A", b") is the same as that obtained by analytical method for the optimal trans- port and is shown to be the unique solution to the Monge problem [T]
A.3 PROOF OF THE CONVERGENCE OF THE ADMM METHOD
Theorem 3. Assume To given by neural network with Lipschitz continuous activation functions and satisfying |/ooll + oo for + co. For sufficiently large the ADMM method generates a bounded sequence that converges the stationary point of the Lagrangian
Ce(01, (0,1,,2.A,p)) = EanN [| -T, + T7 (o(To.=Hl)))
+ A A (To, - To + P (E (To, x -o,, 2))
(27)
2
Proof. We only need to check the conditions in the reference Wang et al. (2019). First we show the objective function is coercive. For bounded 01 02 and 01 > co, then Exen ||z -oo--o,,0 x| +
n- (o(To=# = and thus is coercive. The Lipschitz continuous assumption implies the objective function is also Lipschitz continuous. Moreover, the solution to the subproblem ming, Egwy [| Tg,
Tg,x [=] and ming, (0(TT is also Lipschitz continuous. Thus the con- ditions in Wang et al. (2019) holds and we can conclude that the ADMM method converges.
A.4 THE KKT CONDITIONS
The KKT condition is the conditions for the saddle point of the Lagrangian. The KKT condition for the St. Lagrangian 1
At the saddle point (0* >*), the KKT conditions for [4]
Vo X*) = VoE ( [| - Tox ))) + + +Voooo(Te= = 0.
VxLsr(O" X*) o(To- 0- # u1(1) O.
At each training step (taking P to be constant), the KKT conditions for the QP loss (5) are
VeLop(OR,
V VoEz~a(([Ez ) + PTo (V v)) (Ox) O,
At each training step (taking to be constant), the KKT conditions for the AL loss 66 are
ALLAL(Os, (2, de: pa) = VolL (|x - Tox [] + (As + PxO # (4 = O, VALAL Ak, Pk) (T 11v) 0.
A.5 TRAINING RESULTS
The results of 78D Gaussian is plotted below.
We demonstrate the robustness of AL and ADMM method compared to the SL method. The training curves for these three methods are plotted in Figure
A.6 EXPERIMENT DETAILS
Code for the numerical experiments available at
Table 1 and Figure 2, 3 produced using normalizing flow network with planar transformation layers. The initial and target densities are known and samples are drawn from these distributions. Algorithms described in sections 3. 3.1-3.3 are used. The KL divergence is computed via [11]
13

---

Under review as a conference paper at ICLR 2023
13
Figure 4: Training curves of the three algorithms for 78D Gaussian benchmark (the increasing line the transport cost, and the decreasing line is the DKL- Left: the SL method; middle: the AL method; right: the ADMM method, orange line is for the first network and blue line is for the second network).
035
4.55
24
533
031
030
taations
tastions
0.37
04-
335
835
03-
63
034
03
03
0.
61
431
931
66
0.00
Esutoors
Eauioons
05
335
08-
435
432
111
06-
30
200
250
200
203
26
250
200
Figure 5: Graph of the training curves, increasing line for transport loss and decreasing line for the DKL. Top: SL with À = (le(fef)) X = 10(right) Middle: AL with À = 0.1 (left), X = 10 (right). Bottom: ADMM with A 0.1 (left) and A 10- 1 (right).
For the WGAN example, we use a fully connected neural network with hidden layers width 400. Gradient penalty is used for the WGAN-GP and for the SL. AL and ADMM methods, the alogrithm described in section 3.4 and = is used in the SI. and p = 10-4 is used for AL and 10-5 used for the ADMM.
14

---

Under review as a conference paper at ICLR 2023
33-
Figure 6: Graph of the results of WGAN without gradient penalty and the SL method; Left to the right: the discriminator loss, the generator loss and the generated sample; top: WGAN without gradient penaly, bottom: the SL method (equation 12
Figure 7: Graph of the transport distances. Top (left to right): WGAN-GP, SL, AL and ADMM
15


Reviews:
[{"review_id": "dfbmPD2cth6", "reviewer": "Reviewer_ZQmE", "summary": "", "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": "", "summary_of_the_paper": "Combining traditional numerical methods for solving the Monge equation with deep learning, thus compensating for the limitations of numerical methods for high-dimensional data and the need for deep learning networks based on Kantorovich duality and complex network structures.", "summary_of_the_review": "In general the article would be more valuable if some comparative experiments with traditional numerical algorithms were added, as well as some visual schematics.", "strength_and_weaknesses": "Strength:\n(1) The article has sufficient theoretical derivations and proofs for the three algorithms. \nWeaknesses:\n(1)  Although it is illustrated that traditional numerical methods are not applicable to high-dimensional problems, relevant comparative experiments and graphs are missing.\n(2) Lack of comparative experiments with existing learning-based methods. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "X5sV81zbIJ", "reviewer": "Reviewer_rwsz", "summary": "", "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": "", "summary_of_the_paper": "This work proposes a direct constraint optimization method to solve optimal transport maps using the original Monge's formulation. Three differential algorithms are studied:  the Langrangian multiplier method, the augmented Lagrangian method, and the alternating direction method of multipliers (ADMM). The proposed method improve the accuracy of learned optimal transport maps on high dimensional benchmarks, reduce the regularization effects and better learn the target distributions at a lower transport cost.\n\nThe contributions are \n1. integrate three constraint optimization algorithms including the Standard Lagrangian (SL), the Augmented Lagrangian method (AL) and the Alternating Direction Method of Multipliers (ADMM) with neural networks to solve the Monge problem of optimal transport\nwith provable guarantees (Theorem 1-3).\n2. show that the proposed  method is able to find an accurate optimal transport map between Gaussian distributions, both theoretically (Theorem 2) and experimentally. Moreover, the method is applied to WGAN and show that our method can find a generative map with lower transport cost while not sacrificing the quality of outputs.\n3. Three algorithms are compared and it is found that  the SL algorithm introduces errors but is simple and easy to implement, while AL and ADMM algorithms can find exact results and are more robust, and ADMM gives a lower transport cost in general.", "summary_of_the_review": "This work proposes a direct constraint optimization method to solve optimal transport maps using the original Monge's formulation. Three differential algorithms are studied:  the Langrangian multiplier method, the augmented Lagrangian method, and the alternating direction method of multipliers (ADMM). The proposed method improve the accuracy of learned optimal transport maps on high dimensional benchmarks, reduce the regularization effects and better learn the target distributions at a lower transport cost.\n\nThe paper is well written, the mathematical formulations are rigorous and elegant, the theoretic proofs are convincing. The preliminary numerical experiments support the claims. The work can be further improved : enlarge the scope of theorem 2, add more real world testing experiments. Furthermore, comparisons with other approaches which also solve the problem directly, such as the geometric variational method by Yau et al, which reduce the problem to convex optimization. ", "strength_and_weaknesses": "The strengths of the current work \n1. Propose novel algorithm to solve optimal transport map method directly using Monge's formulation with constraint optimization.\n2. Three different optimization algorithms are formulated, evaluated and compared.\n3. Theoretic proofs for the existence of the solution (Theorem 1 and Theorem 2), and the convergence of ADMM (Theorem 3)\nThe mathematical formulation and proofs are clean, rigorous and thorough.\n\nThe weakness of the current work is that the experiments are conducted on simple datasets (MNIST), it is desirable to test the theoretic results on real world datasets. Another weakness is that by using Yau and Gu's geometric variational algorithm, the problem of finding Monge Optimal transport map is reduced to a convex optimization problem, so the theoretic setup is easier and the optimization can be carried out using Newton's method. It is helpful to compare the proposed method with this convex optimization approach. Another weakness is that the proposed method can not find the singularity of the OT map precisely, since it use soft constraints (mollifiers) which will obscure the boundary of the supports.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6489DmaqxkT", "reviewer": "Reviewer_3wmi", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "In this paper, the authors propose to directly solve the Monge's optimal transport problem under $L^2$ cost function, where the existence and uniqueness of the solution is guaranteed. Specifically, they use a neural network to parameterize the OT map, and try to solve the KL-divergence regularized OT problem proposed in equation (3). Furthermore, different Lagrangian multiplier methods are used to handle the constraints. Finally, the problems are solved by ADMM method. Some toy examples are used to illustrate the performance of the propose method.", "summary_of_the_review": "Generally, the authors need to prove the hidden assumptions and do more experiments to show the performance. ", "strength_and_weaknesses": "Strength: \n- It is interesting to solve the KL regularized OT problem. Combining them together should give us more understanding of the both problems.\n\nWeakness:\n- The method can only solve continuous OT problem supported on the whole Euclidean space, or it will encounter severe convergence problem since the KL divergence can easily be $\\infty$.\n- If the target distribution is discrete like MNIST, how to compute the KL divergence?\n- For the Alg 1,2,3, it is unclear how to conduct the optimization, by sampling?\n- From the current experimental results, either Fig. 4 or Fig. 5, it is hard to say that the proposed methods work well.\n- No implement details.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "0EG_XQuO8jb", "reviewer": "Reviewer_5X5o", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "The paper is considering the problem of learning the optimal transport maps with quadratic cost. In order to do so, they propose to apply a constrained optimization algorithm directly on the Monge formulation. In particular, the marginal constraint T#\\mu = \\nu is formulated as a penalty that involves a distance between T#\\mu and \\nu.  Then, constrained optimization algorithms, such as ADMM, are proposed to obtain the parameters of the map T. The proposed algorithm is explained and illustrated on Gaussian, mixture of Gaussians, and MNISST dataset.  ", "summary_of_the_review": "Although it is a nicely written paper, the proposed approach suffers from the same issues that the paper is trying to avoid. ", "strength_and_weaknesses": "Strength:\n- the paper is well-written\n- the discussion on the Gaussian example is very helpful in understanding the proposed approach, its advantages and disadvantages\n\nWeakness: \n- the paper is motivated by shortcomings of the existing approaches in terms of special neural network structure or regularization, but the proposed approach seems to suffer from the same issues. \n- the most challenging step of the approach is how to compute the penalty and perform the ADMM penalty step. It is not clear how to compute the KL divergence or any other type of distance (W1, IPM,..) in general. Minimizing the penalty term is the subject of all GAN papers and the way to do it exactly suffers from the issues that the paper is trying to avoid. \n   ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject