Original: OTCOP: L EARNING OPTIMAL TRANSPORT MAP VIA CONSTRAINT OPTIMIZATION
Revised: OTCOP: L EARNING OPTIMAL TRANSPORT MAP VIA CONSTRAINT OPTIMIZATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: The approximation power of the neural network makes it an ideal tool to learn optimal transport maps.
Revised: The approximation power of the neural network makes it an ideal tool to learn optimal transport maps.

Original: However, existing methods are mostly based on the Kantorovich duality and require regularizations and/or special network structures.
Revised: However, existing methods are mostly based on the Kantorovich duality and require regularizations and/or special network structures.

Original: In this paper, we propose a direct constraint optimization algorithm for the computation of optimal transport maps based on the Monge formulation.
Revised: In this paper, we propose a direct constraint optimization algorithm for the computation of optimal transport maps based on the Monge formulation.

Original: We solve this constraint optimization problem by using three different methods: the Langrangian multiplier method, the augmented Lagrangian method, and the alternating direction method of multipliers (ADMM).
Revised: We solve this constraint optimization problem by using three different methods: the Langrangian multiplier method, the augmented Lagrangian method, and the alternating direction method of multipliers (ADMM).

Original: We demonstrate a significant accuracy of learned optimal transport maps on high dimensional benchmarks.
Revised: We demonstrate a significant accuracy of learned optimal transport maps on high dimensional benchmarks.

Original: Moreover, we show that our methods reduce the regularization effects and accurately learn the target distributions at a lower transport cost.
Revised: Moreover, we show that our methods reduce the regularization effects and accurately learn the target distributions at a lower transport cost.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: There has been a great interest in applying modern machine learning techniques for finding optimal transport maps between two distributions.
Revised: There has been a great interest in applying modern machine learning techniques for finding optimal transport maps between two distributions.

Original: Different from traditional computational methods that solve PDEs for optimal transport maps (Benamou & Brenier (2000); Angenent et al.
Revised: Different from traditional computational methods that solve PDEs for optimal transport maps (Benamou & Brenier (2000); Angenent et al.

Original: ; Li et al.
Revised: ; Li et al.

Original: (2018)), modern machine learning techniques aim to solve the problem directly by optimizations.
Revised: (2018)), modern machine learning techniques aim to solve the problem directly by optimizations.

Original: The Sinkhorn Distance method Cuturi (2013); Peyr´e et al.
Revised: The Sinkhorn Distance method Cuturi (2013); Peyr´e et al.

Original: (2019), the regularized OT dual Seguy et al.
Revised: (2019), the regularized OT dual Seguy et al.

Original: (2017) have been used to find large scale optimal transport maps between discrete probability distributions and have been used to train generative networks Genevay et al.
Revised: (2017) have been used to find large scale optimal transport maps between discrete probability distributions and have been used to train generative networks Genevay et al.

Original: ; Sanjabi et al.
Revised: ; Sanjabi et al.

Original: The Input Convex Neural Network (ICNN) is used to construct a convex Brenier potential for finding optimal transport maps Makkuva et al.
Revised: The Input Convex Neural Network (ICNN) is used to construct a convex Brenier potential for finding optimal transport maps Makkuva et al.

Original: (2020) between continuous distributions
Revised: (2020) between continuous distributions

Original: and is recently used in population dynamics Bunne et al.
Revised: and is recently used in population dynamics Bunne et al.

Original: (2022), which combines the ICNN and Sinkhorn distance methods Amos et al.
Revised: (2022), which combines the ICNN and Sinkhorn distance methods Amos et al.

Original: Despite these successes, most methods are based on the duality formulation and avoids the direct treatment on the Monge problem.
Revised: Despite these successes, most methods are based on the duality formulation and avoids the direct treatment on the Monge problem.

Original: In this paper, we focus on the direct solution of the Monge problem.
Revised: In this paper, we focus on the direct solution of the Monge problem.

Original: The Monge problem (Monge (1781)) directly seeks to identify the optimal transport maps and is a nonlinear constraint optimization problem.
Revised: The Monge problem (Monge (1781)) directly seeks to identify the optimal transport maps and is a nonlinear constraint optimization problem.

Original: The major difficulty in solving the problem numerically is that it is nonlinear and includes a constraint that the push-forward distribution is equal to the target distribution, which is difficult to implement.
Revised: The major difficulty in solving the problem numerically is that it is nonlinear and includes a constraint that the push-forward distribution is equal to the target distribution, which is difficult to implement.

Original: Therefore, most optimal transport algorithms avoid directly solving the
Revised: Therefore, most optimal transport algorithms avoid directly solving the

Original: Monge problem but use the Kantorovich duality (Kantorovich (1942)), for which the objective function is linear and the transport map is obtained by taking the gradient of the Brenier potential for the quadratic cost.
Revised: Monge problem but use the Kantorovich duality (Kantorovich (1942)), for which the objective function is linear and the transport map is obtained by taking the gradient of the Brenier potential for the quadratic cost.

Original: However, these two problems are not always identical (Villani (2009)) and it is desirable to find a direct approach for the Monge problem.
Revised: However, these two problems are not always identical (Villani (2009)) and it is desirable to find a direct approach for the Monge problem.

Original: The Monge problem has been solved numerically using optimization based methods with polynomial approximations.
Revised: The Monge problem has been solved numerically using optimization based methods with polynomial approximations.

Original: For example, a Lagrangian penalty method was used to find optimal transport maps approximated by polynomials for Bayesian inference El Moselhy & Marzouk (2012) and space discretization was used in Haber et al.
Revised: For example, a Lagrangian penalty method was used to find optimal transport maps approximated by polynomials for Bayesian inference El Moselhy & Marzouk (2012) and space discretization was used in Haber et al.

Original: (2010) to calculate the Jacobian matrix of the transport maps and transferred the optimization to finite dimensional spaces.
Revised: (2010) to calculate the Jacobian matrix of the transport maps and transferred the optimization to finite dimensional spaces.

Original: However, their approaches are limited to low dimensions as number of grids expands exponentially as dimensions become large.
Revised: However, their approaches are limited to low dimensions as number of grids expands exponentially as dimensions become large.

Original: Considering the success of deep neural networks in approximating high dimensional data, the integration of classical constraint optimization methods and neural networks holds a promise.
Revised: Considering the success of deep neural networks in approximating high dimensional data, the integration of classical constraint optimization methods and neural networks holds a promise.

Original: One successful application of the optimal transport theory to deep learning is the Wasserstein Generative Adversarial Network (WGAN)
Revised: One successful application of the optimal transport theory to deep learning is the Wasserstein Generative Adversarial Network (WGAN)

Original: Arjovsky et al.
Revised: Arjovsky et al.

Original: However, WGAN only use the optimal transport distance as a loss function and does not target at finding the optimal transport maps.
Revised: However, WGAN only use the optimal transport distance as a loss function and does not target at finding the optimal transport maps.

Original: It is desirable to study whether it is possible to lower the transport cost of the map learned by WGAN or other networks using the algorithm for finding optimal transport maps.
Revised: It is desirable to study whether it is possible to lower the transport cost of the map learned by WGAN or other networks using the algorithm for finding optimal transport maps.

Original: This paper presents a new approach for finding optimal transport maps between two continuous distributions.
Revised: This paper presents a new approach for finding optimal transport maps between two continuous distributions.

Original: We make the following contributions:
Revised: We make the following contributions:

Original: • 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 to solve the Monge problem of optimal transport with provable guarantees (Theorem 1-3).
Revised: • 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 to solve the Monge problem of optimal transport with provable guarantees (Theorem 1-3).

Original: • We show that our method is able to find an accurate optimal transport map between Gaussian distributions, both theoretically (Theorem 2) and experimentally.
Revised: • We show that our method is able to find an accurate optimal transport map between Gaussian distributions, both theoretically (Theorem 2) and experimentally.

Original: 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.
Revised: 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.

Original: • 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.
Revised: • 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.

Original: Notations.
Revised: Notations.

Original: We use the notations α d = ( α, · · · , α ) ∈ R d and α d × d for the constant d × d matrix.
Revised: We use the notations α d = ( α, · · · , α ) ∈ R d and α d × d for the constant d × d matrix.

Original: The transport cost of a map T , which pushes distribution µ to ν , is defined to be E x ∼ µ [ | x − Tx | 2 ] .
Revised: The transport cost of a map T , which pushes distribution µ to ν , is defined to be E x ∼ µ [ | x − Tx | 2 ] .

Original:  T HE M ONGE PROBLEM AS CONSTRAINT OPTIMIZATION 2. T HE M ONGE PROBLEM
Revised:  T HE M ONGE PROBLEM AS CONSTRAINT OPTIMIZATION 2. T HE M ONGE PROBLEM

Original: Let ( X, µ ) , ( Y, ν ) be two separable metric probability spaces.
Revised: Let ( X, µ ) , ( Y, ν ) be two separable metric probability spaces.

Original: The Monge problem is to find a transport map T : X (cid:55)→ Y that realizes the infimum
Revised: The Monge problem is to find a transport map T : X (cid:55)→ Y that realizes the infimum

Original: [Equation]
Revised: [Equation]

Original: (cid:26)(cid:90)
Revised: (cid:26)(cid:90)

Original: [Equation]
Revised: [Equation]

Original: (1) where T # µ denotes the push forward of µ and c : X × Y → R + is a Borel measurable function which is lower semicontinuous.
Revised: (1) where T # µ denotes the push forward of µ and c : X × Y → R + is a Borel measurable function which is lower semicontinuous.

Original: In this paper, we simply take the distance | x − y | 

Revised: In this paper, we simply take the distance | x − y |

Original: but our method applies to other distance functions.
Revised: but our method applies to other distance functions.

Original: The existence of the Monge problem is difficult and does not hold always.
Revised: The existence of the Monge problem is difficult and does not hold always.

Original: 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).
Revised: 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).

Original: Therefore, here we focus on learning transport maps between continuous distributions.
Revised: Therefore, here we focus on learning transport maps between continuous distributions.

Original: For discrete distributions, one can apply dequantization techniques to transform them to continuous distributions Ho et al.
Revised: For discrete distributions, one can apply dequantization techniques to transform them to continuous distributions Ho et al.

Original: 2. T HE M ONGE PROBLEM AS CONSTRAINT OPTIMIZATION
Revised: 2. T HE M ONGE PROBLEM AS CONSTRAINT OPTIMIZATION

Original: In order to solve the Monge problem, we use a generative network, denoted by T θ with parameter set θ , which inputs random samples from the distribution µ and generates samples representing the target distributions ν .
Revised: In order to solve the Monge problem, we use a generative network, denoted by T θ with parameter set θ , which inputs random samples from the distribution µ and generates samples representing the target distributions ν .

Original: As can be seen from the definition, the Monge problem is a constraint optimization problem.
Revised: As can be seen from the definition, the Monge problem is a constraint optimization problem.

Original: However, the constraint T # µ = ν is a highly nonlinear constraint.
Revised: However, the constraint T # µ = ν is a highly nonlinear constraint.

Original: In order to impose this constraint, we take d ( ·|· ) to be a distance function (such as the Wasserstein distance, the MMD (Gretton et al.
Revised: In order to impose this constraint, we take d ( ·|· ) to be a distance function (such as the Wasserstein distance, the MMD (Gretton et al.

Original: or the IPM (M¨uller (1997))) or a probability divergence (such as the Kullback–Leibler (KL) divergence).
Revised: or the IPM (M¨uller (1997))) or a probability divergence (such as the Kullback–Leibler (KL) divergence).

Original: The constraint optimization problem reads as
Revised: The constraint optimization problem reads as

Original: [Equation]
Revised: [Equation]

Original: The objective of this paper is to solve the above problem using techniques from the constraint optimization theory (Bertsekas (2014)).
Revised: The objective of this paper is to solve the above problem using techniques from the constraint optimization theory (Bertsekas (2014)).

Original: Since a neural network may not fully reveal the target distributions ν , the above problem can be relaxed to
Revised: Since a neural network may not fully reveal the target distributions ν , the above problem can be relaxed to

Original: [Equation]
Revised: [Equation]

Original: When α goes towards zero, we can prove that the solution T θ α to the problem (3) converges to the solution of the original Monge problem (1).
Revised: When α goes towards zero, we can prove that the solution T θ α to the problem (3) converges to the solution of the original Monge problem (1).

Original: The following theorem holds:
Revised: The following theorem holds:

Original: Theorem 1.
Revised: Theorem 1.

Original: Let µ, ν be two probability measures on R d with finite second moments and are absolutely continuous.
Revised: Let µ, ν be two probability measures on R d with finite second moments and are absolutely continuous.

Original: Let T θ be given by a neural network with bounded width(each with at least 2 d + 2 neurons) and arbitrary depth, and with non-affine activation functions.
Revised: Let T θ be given by a neural network with bounded width(each with at least 2 d + 2 neurons) and arbitrary depth, and with non-affine activation functions.

Original: Suppose for any α > 0 , there exists a solution θ ∗ α to problem (2), then as α → 0 , T θ ∗ α → T where T is a solution of the Monge problem (1).
Revised: Suppose for any α > 0 , there exists a solution θ ∗ α to problem (2), then as α → 0 , T θ ∗ α → T where T is a solution of the Monge problem (1).

Original: Moreover, sup x ∈X | T α ( x ) − T ( x ) | C ≤ Cα for some constant C for any compact subset X ∈ R d .
Revised: Moreover, sup x ∈X | T α ( x ) − T ( x ) | C ≤ Cα for some constant C for any compact subset X ∈ R d .

Original: Proof of the above theorem follows from the universal approximation theorem (Kidger & Lyons (2020)) and the existence theorem of the Monge problem (Villani (2009)), and is given in Appendix
Revised: Proof of the above theorem follows from the universal approximation theorem (Kidger & Lyons (2020)) and the existence theorem of the Monge problem (Villani (2009)), and is given in Appendix

Original: A.1.
Revised: A.1.

Original: 2. E XAMPLE : THE M ONGE PROBLEM FROM G AUSSIAN TO G AUSSIAN
Revised: 2. E XAMPLE : THE M ONGE PROBLEM FROM G AUSSIAN TO G AUSSIAN

Original: For the case when µ ∼ N ( X 1 , Σ 1 ) and ν ∈ N ( X 2 , Σ 2 ) are two multivariate normal distributions, the optimal transport map is unique and can be explicitly given by T ∗ : x (cid:55)→ X 2 + A ∗ ( x − X 1 ) with
Revised: For the case when µ ∼ N ( X 1 , Σ 1 ) and ν ∈ N ( X 2 , Σ 2 ) are two multivariate normal distributions, the optimal transport map is unique and can be explicitly given by T ∗ : x (cid:55)→ X 2 + A ∗ ( x − X 1 ) with

Original: A ∗ =
Revised: A ∗ =

Original: Σ − 1 / 2(
Revised: Σ − 1 / 2(

Original: [Equation]
Revised: [Equation]

Original: (Olkin & Pukelsheim (1982)).
Revised: (Olkin & Pukelsheim (1982)).

Original: Taking d = D KL to be the KLdivergence, we prove that the solution to problem (2) with T θ x
Revised: Taking d = D KL to be the KLdivergence, we prove that the solution to problem (2) with T θ x

Original: =
Revised: =

Original: Ax + b ( θ = { A ∈ R d × d , b ∈ R d } ) is the optimal transport map T ∗ (Theorem 2 in Appendix A.1).
Revised: Ax + b ( θ = { A ∈ R d × d , b ∈ R d } ) is the optimal transport map T ∗ (Theorem 2 in Appendix A.1).

Original:  C ONSTRAINT OPTIMIZATION FOR OPTIMAL TRANSPORT
Revised:  C ONSTRAINT OPTIMIZATION FOR OPTIMAL TRANSPORT

Original: We propose to leverage three different algorithms to solve the constraint problem (2)
Revised: We propose to leverage three different algorithms to solve the constraint problem (2)

Original: . 3. P ENALTY METHOD (OTCOP-P) Standard
Revised: . 3. P ENALTY METHOD (OTCOP-P) Standard

Original: Lagrangian (SL).
Revised: Lagrangian (SL).

Original: We introduce a Lagrangian multiplier λ and take the Lagrangian function as
Revised: We introduce a Lagrangian multiplier λ and take the Lagrangian function as

Original: [Equation]
Revised: [Equation]

Original: Then the solution to the problem (2) is a saddle point of the above Lagrangian.
Revised: Then the solution to the problem (2) is a saddle point of the above Lagrangian.

Original: By duality theory, for each α ≥ 0 , the problem (3) corresponds to the duality problem min θ L SL ( θ, λ, 0) for a λ ∈ [0 , ∞ ] .
Revised: By duality theory, for each α ≥ 0 , the problem (3) corresponds to the duality problem min θ L SL ( θ, λ, 0) for a λ ∈ [0 , ∞ ] .

Original: Hence we can take a suitable λ to solve the problem (2) approximately.
Revised: Hence we can take a suitable λ to solve the problem (2) approximately.

Original: According to the Brenier’s polar factorization theorem Brenier (1991), the optimal transport map should satisfy ∇ × T = 0 , hence we can add an additional term |∇ × T | 2 into the above Lagrangian to impose this constraint.
Revised: According to the Brenier’s polar factorization theorem Brenier (1991), the optimal transport map should satisfy ∇ × T = 0 , hence we can add an additional term |∇ × T | 2 into the above Lagrangian to impose this constraint.

Original: We will show experimentally that without this term, this constraint is almost satisfied and we will not included in our implementations.
Revised: We will show experimentally that without this term, this constraint is almost satisfied and we will not included in our implementations.

Original: Quadratic penalty (QP).
Revised: Quadratic penalty (QP).

Original: Instead of taking d ( T θ # | ν ) , we can take a quadratic penalty loss
Revised: Instead of taking d ( T θ # | ν ) , we can take a quadratic penalty loss

Original: [Equation]
Revised: [Equation]

Original: As ρ goes towards infinity, the constraint violations is penalized with increasing severity.
Revised: As ρ goes towards infinity, the constraint violations is penalized with increasing severity.

Original: For example, we can take ρ k at the k th training step to be increased by multiplying by a constant bigger than 1 and parameter θ can be updated using gradient descent considering ρ as a constant.
Revised: For example, we can take ρ k at the k th training step to be increased by multiplying by a constant bigger than 1 and parameter θ can be updated using gradient descent considering ρ as a constant.

Original: Convergence.
Revised: Convergence.

Original: Suppose there exists a global minimizer to the problem (2), and θ k is the exact minimizer of L QP ( θ, ρ k ) and ρ k ↑ ∞ .
Revised: Suppose there exists a global minimizer to the problem (2), and θ k is the exact minimizer of L QP ( θ, ρ k ) and ρ k ↑ ∞ .

Original: Then any limit point of the sequence { θ k } is a solution to problem (2).
Revised: Then any limit point of the sequence { θ k } is a solution to problem (2).

Original: Moreover, for any ε > 0 , there exists a sufficient large K > 0 , | θ k − θ ∗ | ≤ ε for k ≥ K (see (Nocedal & Wright, 1999, Theorem 17.1)).
Revised: Moreover, for any ε > 0 , there exists a sufficient large K > 0 , | θ k − θ ∗ | ≤ ε for k ≥ K (see (Nocedal & Wright, 1999, Theorem 17.1)).

Original: In addition, without assuming a global minimizer, for a sequence θ k such that ∇ θ k L QP ( θ k ; ρ k ) → 0 , its all limit points θ ∗ satisfy the Karush–Kuhn–Tucker (KKT) conditions and there exists a subsequence such that lim k →∞ ( ρ k d ( T θ k # | ν ) = λ ∗ , where λ ∗ is the multiplier that satisfies the KKT condition (see Appendix A.4 for the KKT condition and see (Nocedal & Wright, 1999, Theorem 17.2) for the proof).
Revised: In addition, without assuming a global minimizer, for a sequence θ k such that ∇ θ k L QP ( θ k ; ρ k ) → 0 , its all limit points θ ∗ satisfy the Karush–Kuhn–Tucker (KKT) conditions and there exists a subsequence such that lim k →∞ ( ρ k d ( T θ k # | ν ) = λ ∗ , where λ ∗ is the multiplier that satisfies the KKT condition (see Appendix A.4 for the KKT condition and see (Nocedal & Wright, 1999, Theorem 17.2) for the proof).

Original: Advantages and disadvantages.
Revised: Advantages and disadvantages.

Original: The penalty method is simple and easy to implement.
Revised: The penalty method is simple and easy to implement.

Original: However, since the optimal value of the Lagrangian multiplier λ is unkown (SL) or the optimal condition for the Lagrangian multiplier ρ is infinite (QP), the penalty always introduce errors and the exact solution to the problem (2) cannot be reached.
Revised: However, since the optimal value of the Lagrangian multiplier λ is unkown (SL) or the optimal condition for the Lagrangian multiplier ρ is infinite (QP), the penalty always introduce errors and the exact solution to the problem (2) cannot be reached.

Original: Moreover, the Hessian of the Lagrangian ∇ 2 θθ L QP becomes singular as ρ goes towards infinity and cause ill-condition problems.
Revised: Moreover, the Hessian of the Lagrangian ∇ 2 θθ L QP becomes singular as ρ goes towards infinity and cause ill-condition problems.

Original: These issues can be solved by the methods below, but at the expense of a more computational cost.
Revised: These issues can be solved by the methods below, but at the expense of a more computational cost.

Original: 3. T HE AUGMENTED L AGRANGIAN METHOD (OCTOP-AL)
Revised: 3. T HE AUGMENTED L AGRANGIAN METHOD (OCTOP-AL)

Original: In order to overcome the above issues, we can use the augmented Lagrangian method, by taking the loss function as
Revised: In order to overcome the above issues, we can use the augmented Lagrangian method, by taking the loss function as

Original: [Equation]
Revised: [Equation]

Original:  The above function combines standard Lagrangian penalty (4) and quadratic Lagrangian penalty (5).
Revised: [Equation] The above function combines standard Lagrangian penalty (4) and quadratic Lagrangian penalty (5).

Original: At the k th iteration, fix λ k , ρ k and solve θ k = arg min θ L AL ( θ, λ k , ρ k ) .
Revised: At the k th iteration, fix λ k , ρ k and solve θ k = arg min θ L AL ( θ, λ k , ρ k ) .

Original: After the minimization, we update λ by λ k
Revised: After the minimization, we update λ by λ k

Original: +

Revised: +

Original: =
Revised: =

Original: λ k
Revised: λ k

Original: + ρ k d
Revised: + ρ k d

Original: ( T θ k # µ | ν ) .
Revised: ( T θ k # µ | ν ) .

Original: Comparing the KKT conditions of the SL and the
Revised: Comparing the KKT conditions of the SL and the

Original: AL (see Appendix A.4) implies λ k + ρ k d
Revised: AL (see Appendix A.4) implies λ k + ρ k d

Original: ( T θ k # µ | ν ) ≈ λ ∗ when λ k is taken close to λ ∗ .
Revised: ( T θ k # µ | ν ) ≈ λ ∗ when λ k is taken close to λ ∗ .

Original: Hence d ( T θ k # µ | ν ) ≈
Revised: Hence d ( T θ k # µ | ν ) ≈

Original: ( λ ∗ − λ k ) /ρ .
Revised: ( λ ∗ − λ k ) /ρ .

Original: Compared to the quadratic penalty method that d ( T θ k # µ | ν ) ≈ λ ∗ /ρ , the infeasibility in θ k will be much smaller.
Revised: Compared to the quadratic penalty method that d ( T θ k # µ | ν ) ≈ λ ∗ /ρ , the infeasibility in θ k will be much smaller.

Original: Moreover, for certain choice of ρ , the local solution of (2) is a strict local minimizer of L AL ( θ, λ, ρ ) ((Nocedal & Wright, 1999, Chapter 17)).
Revised: Moreover, for certain choice of ρ , the local solution of (2) is a strict local minimizer of L AL ( θ, λ, ρ ) ((Nocedal & Wright, 1999, Chapter 17)).

Original: Convergence.
Revised: Convergence.

Original: One of the nice properties of the AL method is that for the exact Lagrangian multiplier λ ∗ , the solution θ ∗ of the problem (2) is a strict minimizer of L AL ( θ, λ ∗ , ρ ) for all ρ sufficiently large.
Revised: One of the nice properties of the AL method is that for the exact Lagrangian multiplier λ ∗ , the solution θ ∗ of the problem (2) is a strict minimizer of L AL ( θ, λ ∗ , ρ ) for all ρ sufficiently large.

Original: The existence of a threshold is proved under the condition that ∇ 2 θ L SL ( θ ∗ , λ ∗ ) is locally strictly positive ((Nocedal & Wright, 1999, Theorem 17.6)).
Revised: The existence of a threshold is proved under the condition that ∇ 2 θ L SL ( θ ∗ , λ ∗ ) is locally strictly positive ((Nocedal & Wright, 1999, Theorem 17.6)).

Original: Thus we can take ρ to be increasing at each minimizing step and when ρ becomes bigger than some threshold value ¯ ρ , gradient descent methods could find the local minimizer around θ ∗ .
Revised: Thus we can take ρ to be increasing at each minimizing step and when ρ becomes bigger than some threshold value ¯ ρ , gradient descent methods could find the local minimizer around θ ∗ .

Original: Advantages and disadvantages.
Revised: Advantages and disadvantages.

Original: AL method introduces the multiplier estimates and reduces the likelihood that large values of ρ will be needed to obtain good feasibility and accuracy.
Revised: AL method introduces the multiplier estimates and reduces the likelihood that large values of ρ will be needed to obtain good feasibility and accuracy.

Original: The method is also simple and easy to implement.
Revised: The method is also simple and easy to implement.

Original: However, since this is a min-max method, training may experience oscillations and slower convergence rates.
Revised: However, since this is a min-max method, training may experience oscillations and slower convergence rates.

Original: 3. ADMM METHOD (OCTOP-ADMM)
Revised: 3. ADMM METHOD (OCTOP-ADMM)

Original: The ADMM method blends the decomposition techniques and the AL method and provides an efficient way for constraint optimizations Boyd et al.
Revised: The ADMM method blends the decomposition techniques and the AL method and provides an efficient way for constraint optimizations Boyd et al.

Original: Let S be the set S
Revised: Let S be the set S

Original: = { T θ : T θ # µ = ν } , problem (2) can be rewritten into the form
Revised: = { T θ : T θ # µ = ν } , problem (2) can be rewritten into the form

Original: [Equation]
Revised: [Equation]

Original: where 1 S is the indicator function that equals 0 if T θ ∈ S and equals ∞ if T θ ̸∈ S .
Revised: where 1 S is the indicator function that equals 0 if T θ ∈ S and equals ∞ if T θ ̸∈ S .

Original: In order to apply the ADMM method, we rewrite the above problem into the form
Revised: In order to apply the ADMM method, we rewrite the above problem into the form

Original: [Equation]
Revised: [Equation]

Original: In the ADMM method, we alternatively update θ 1 and θ 2 .
Revised: In the ADMM method, we alternatively update θ 1 and θ 2 .

Original: First we take θ 2 to be constant and take the minimization of the above problem over θ 1 , and then we project θ 2 onto the space S .
Revised: First we take θ 2 to be constant and take the minimization of the above problem over θ 1 , and then we project θ 2 onto the space S .

Original: In detail, we introduce the loss function
Revised: In detail, we introduce the loss function

Original: [Equation]
Revised: [Equation]

Original: Here Λ ∈ R d is the multiplier.
Revised: Here Λ ∈ R d is the multiplier.

Original: The training procedure is given by
Revised: The training procedure is given by

Original: θ k 1 = arg min
Revised: θ k + = arg min

Original: θ 1 L ADMM ( θ 1 , θ k 2 , Λ k , ρ ) (assuming 1 d ( T θ 2# µ | ν )=

Revised: θ 1 L ADMM ( θ 1 , θ k 2 , Λ k , ρ ) (assuming 1 d ( T θ 2# µ | ν )=

Original: = 0 ); 2.
Revised: = 0 ); 2.

Original: θ k 2 = arg min
Revised: θ k + = arg min

Original: θ 2 d ( T θ 2 # µ | ν ) ; 3.
Revised: θ 2 d ( T θ 2 # µ | ν ) ; 3.

Original: Λ k
Revised: Λ k

Original: +

Revised: +

Original: =
Revised: =

Original: Λ k
Revised: Λ k

Original: + ρ E x ∼ µ ( T θ k +11 # ( x ) − T θ k [-+
-]{++1 
+}
Revised: + ρ E x ∼ µ ( T θ k +11 # ( x ) − T θ k [-+
-]{++1 
+}

Original: # ( x )) .
Revised: # ( x )) .

Original: Convergence.
Revised: Convergence.

Original: The convergence of ADMM method is only known to hold under convex conditions or for some non-convex problems (Boyd et al.
Revised: The convergence of ADMM method is only known to hold under convex conditions or for some non-convex problems (Boyd et al.

Original: Using results of Wang et al.
Revised: Using results of Wang et al.

Original: (2019), we can prove the convergence of the ADMM method if we modify the above method by relaxing problem (8) to
Revised: (2019), we can prove the convergence of the ADMM method if we modify the above method by relaxing problem (8) to

Original: [Equation]
Revised: [Equation]

Original: where η ε is the mollifier function which converges to the δ -function as ε → 0 .
Revised: where η ε is the mollifier function which converges to the δ -function as ε → 0 .

Original: We show that the
Revised: We show that the

Original: ADMM method converges to the KKT points and if the corresponding Lagrangian is a KurdykaŁojasiewicz function, the ADMM method converges globally to the unique solution (see Appendix
Revised: ADMM method converges to the KKT points and if the corresponding Lagrangian is a KurdykaŁojasiewicz function, the ADMM method converges globally to the unique solution (see Appendix

Original: A.3 for details).
Revised: A.3 for details).

Original: As a consequence, there exists a convergence sequence of the ADMM method for the problem (8) approximately.
Revised: As a consequence, there exists a convergence sequence of the ADMM method for the problem (8) approximately.

Original: Advantage and disadvantages.
Revised: Advantage and disadvantages.

Original: The advantage of ADMM is that it decomposes the Monge problem into two sub-problems: minimizing the transport cost and minimizing the D KL .
Revised: The advantage of ADMM is that it decomposes the Monge problem into two sub-problems: minimizing the transport cost and minimizing the D KL .

Original: Compared to the
Revised: Compared to the

Original: AL method, AMDD solves two decomposed minimization problem and AMDD may also converge faster than the AL method in some situations (Wang et al.
Revised: AL method, AMDD solves two decomposed minimization problem and AMDD may also converge faster than the AL method in some situations (Wang et al.

Original: However, the method also solves a min-max problem and may facing oscillations in the training process.
Revised: However, the method also solves a min-max problem and may facing oscillations in the training process.

Original: 3. I MPLEMENTATION OF THE ALGORITHMS
Revised: 3. T HE DISTANCE d

Revised: We need a distance/divergence functional d to compare the generated distribution and the target distribution. Here we require d to be able to compute using samples.

Revised: The KL divergence D KL .

Revised: When the densities of the target distribution is known, we can use networks of normalizing flow (Rezende & Mohamed (2015)) as T θ and the D KL is computed via

Revised: [Equation]

Revised: [Equation]

Revised: where J T θ is the Jacobian of the transport map which is able to compute using the normalizing flow networks.

Revised: When the target densities are unknown.

Revised: We need to reverse the KL divergence and take

Revised: [Equation]

Revised: We can drop the first term in the bracket since it is a constant.

Revised: We can take L SL in (4) as

Revised: [Equation]

Revised: For the ADMM method, the second step changes to θ k 2 ← arg min θ E x ∼ ν [ − log µ ( T − 1 θ x ) − log det J T

Revised: θ

Revised: ( T − 1 x )] .

Revised: In order to apply the AL method, one can use E x ∼ ν

Revised: [ −∇ θ log µ ( T − 1 θ x ) − ∇ θ log det J T θ ( T − 1 x )] to replace the D KL terms in (6).

Revised: Test function as multiplier.

Revised: A weak form of the constraint T θ # µ = ν is that for any measurable function f ,

Revised: [Equation]

Revised: we can introduce a discriminator network f w as Lagrangian multiplier and the Lagrangian becomes

Revised: [Equation]

Revised: [Equation]

Revised: The augmented Lagrangian then become

Revised: [Equation]

Revised: The training procedure is given by

Revised: θ k = arg min θ L AL ( θ, w, ρ ) ; 2.

Revised: w k +1 =

Revised: w k

Revised: + ρ k E x ∼ µ [ ∇ w f w ( T θ x )] − E z ∼

Revised: ν

Revised: [ ∇ w f w ( z )] ; 3.

Revised: Assign ρ k +1 ≥ ρ k .

Revised: To apply the ADMM method

Revised: [Equation]

Revised: we take the loss function to be

Revised: [Equation]

Revised: Then the optimization is decoupled into optimization over θ 1 to learn the optimal transport map and the optimization over θ 2 to learn the target distribution.

Revised: Here we can take for example d w 2 to be the loss function of Wasserstein GANs with gradient penalty (Gulrajani et al.

Revised: [Equation]

Original: The implementations of the algorithms are given in Algorithm 1, 2, 3. When the distributions µ and νare given, we may use the divergence D KL and the metric d = D KL . 
Revised: where c is a positive constant and γ is the linear interpolation of T θ # µ and ν . The training procedure is given by

Original: For the general case, we can take d = d w to be a discriminator network with parameter set w .
Revised: [Equation] same as [Equation] 3. I MPLEMENTATION OF THE ALGORITHMS

Original: The discriminator network is updated each training step after the parameter of the generative network is updated.
Revised: The implementations of the algorithms are given in Algorithm 1, 2, 3. Here we give the implementation when d is given directly without using a discriminator network. For the case when test function is used as multiplier, see section 3.4 and Appendix ?? for details.

Original: Algorithm 1 Solving the Monge problem with the penalty method
Revised: Algorithm 1 Solving the Monge problem with the penalty method

Original: Input Data: X ∼ µ, Y ∼ ν ,
Revised: Input Data: X ∼ µ, Y ∼ ν ,

Original: Constants: λ 0 , ρ 0 , α > 1 ,
Revised: Constants: λ 0 , ρ 0 , α > 1 ,

Original: Training step: η . for number of training iterations dofor m steps do θ ← θ − η ∇ θ L SL (or θ ← θ − η ∇ θ L QP ) end for
Revised: Training step: η . for number of training iterations dofor m steps do θ ← θ − η ∇ θ L SL (or θ ← θ − η ∇ θ L QP ) end for

Original: Update ρ ← αρ .
Revised: Update ρ ← αρ .

Original: end for
Revised: end for

Original: Algorithm 2 Solving the Monge problem with the augmented Lagrangian method
Revised: Algorithm 2 Solving the Monge problem with the augmented Lagrangian method

Original: Input Data: X ∼ µ, Y ∼ ν
Revised: Input Data: X ∼ µ, Y ∼ ν

Original: Constants: λ 0 , ρ 0 , α > 1 ,
Revised: Constants: λ 0 , ρ 0 , α > 1 ,

Original: Training step: η . for number of training iterations do for m steps do θ ← θ − η ∇ θ L ALw ← w + η ∇ w L AL (if d = d w ) end for
Revised: Training step: η . for number of training iterations do for m steps do θ ← θ − η ∇ θ L AL end for

Original: Update λ ← λ + ρ d ( T θ # µ | ν )
Revised: Update λ ← λ + ρ d ( T θ # µ | ν )

Original: Update ρ ← αρ .
Revised: Update ρ ← αρ .

Original: end for
Revised: end for

Original: Algorithm 

Revised: Algorithm

Original: Solving the Monge problem with the ADMM
Revised: Solving the Monge problem with the ADMM

Original: Input Data: X ∼ µ, Y ∼ ν , Constants: λ 0 , ρ 0 , α > 1 , Training step: η . for number of training iterations do for m 1 steps do θ 1 ← θ 1 − η ∇ θ L ADMM ( θ 1 , θ 2 , λ, ρ ) end for for m 2 steps do θ 2 ← θ 2 − η ∇ θ d ( T θ 2 # µ | ν ) w 2 ← w 2 + η ∇ w d
Revised: Input Data: X ∼ µ, Y ∼ ν , Constants: λ 0 , ρ 0 , α > 1 , Training step: η . for number of training iterations do for m 1 steps do θ 1 ← θ 1 − η ∇ θ L ADMM ( θ 1 , θ 2 , λ, ρ ) end for for m 2 steps do θ 2 ← θ 2 − η ∇ θ d ( T θ 2 # µ | ν ) w 2 ← w 2 + η ∇ w d

Original: ( T θ 2 # µ | ν ) (if d = d w ) end for
Revised: ( T θ 2 # µ | ν ) (if d = d w ) end for

Original: Update λ ← λ + ρ d ( T θ 1 # µ | T θ 2 # µ ) end for
Revised: Update λ ← λ + ρ d ( T θ 1 # µ | T θ 2 # µ ) end for

Original:  E XPERIMENT 4. M ULTIVARIATE NORMAL DISTRIBUTIONS
Revised:  E XPERIMENT 4. M ULTIVARIATE NORMAL DISTRIBUTIONS

Original: Linear maps.
Revised: Linear maps.

Original: First, we consider the optimal transport between multivariate normal distributions µ = N ( X 1 , Σ 1 ) and ν = N ( X 2 , Σ 2 ) .
Revised: First, we consider the optimal transport between multivariate normal distributions µ = N ( X 1 , Σ 1 ) and ν = N ( X 2 , Σ 2 ) .

Original: Considering the map given by T θ =
Revised: Considering the map given by T θ =

Original: Ax + b with θ =
Revised: Ax + b with θ =

Original: { A ∈
Revised: { A ∈

Original: R d × d , b ∈ R d } , we prove in Theorem 2 that the solution to problem (2) gives the correct solution to the Monge problem.
Revised: R d × d , b ∈ R d } , we prove in Theorem 2 that the solution to problem (2) gives the correct solution to the Monge problem.

Original: Indeed, in this case, the SL, AL and ADMM algorithms reduce to optimization of linear objective function constraint by a nonlinear function and it could be theoretically analyzed by the constraint optimization theory Bertsekas (2014).
Revised: Indeed, in this case, the SL, AL and ADMM algorithms reduce to optimization of linear objective function constraint by a nonlinear function and it could be theoretically analyzed by the constraint optimization theory Bertsekas (2014).

Original: Here we take X 1 = 0 2 , Σ 1 =
Revised: Here we take X 1 = 0 2 , Σ 1 =

Original: I 2 and
Revised: I 2 and

Original: X 2 = 0 2 , Σ 2 =
Revised: X 2 = 0 2 , Σ 2 =

Original: [[4 , 1] , [1 , 4]] , let T θ x
Revised: [[4 , 1] , [1 , 4]] , let T θ x

Original: = Ax with A =
Revised: = Ax with A =

Original: [[ a, b ] , [ b, a ]] ( θ = { a, b } ).
Revised: [[ a, b ] , [ b, a ]] ( θ = { a, b } ).

Original: Then the problem (2) reduces to
Revised: Then the problem (2) reduces to

Original: [Equation]
Revised: [Equation]

Original: The landscapes of D KL and L SL ( θ, λ = 1) as well as the value of min θ L SL ( θ, λ ) as functions of λ , are plotted in Figure 1.
Revised: The landscapes of D KL and L SL ( θ, λ = 1) as well as the value of min θ L SL ( θ, λ ) as functions of λ , are plotted in Figure 1.

Original: As can be seen from the figure, the function D KL ( T θ # µ | ν ) has multiple minimizers (red points), whereas the Lagrangian L SP ( θ, 1) has a unique global minimizer (blue point).
Revised: As can be seen from the figure, the function D KL ( T θ # µ | ν ) has multiple minimizers (red points), whereas the Lagrangian L SP ( θ, 1) has a unique global minimizer (blue point).

Original: Hence minimizing the Lagrangian L SL 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.
Revised: Hence minimizing the Lagrangian L SL 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.

Original: More importantly, from the figure, the whole domain can be divided into four pieces by the landscape of D KL , and starting in each piece, gradient descent method will converge to one of the four different points.
Revised: More importantly, from the figure, the whole domain can be divided into four pieces by the landscape of D KL , and starting in each piece, gradient descent method will converge to one of the four different points.

Original: In contrast, the landscape of the Lagrangian changes dramatically and starting from any point, gradient descent method will only converge to one point.
Revised: In contrast, the landscape of the Lagrangian changes dramatically and starting from any point, gradient descent method will only converge to one point.

Original: The Langrangian introduces errors for λ finite.
Revised: The Langrangian introduces errors for λ finite.

Original: As can be seen from the leftmost figure, the constraint
Revised: As can be seen from the leftmost figure, the constraint

Original: D KL = 0 is more relaxed as λ becomes smaller.
Revised: D KL = 0 is more relaxed as λ becomes smaller.

Original: However, for large λ , the barrier region becomes wider and the gradient descent may converge to a local minimizer away from the optimal result.
Revised: However, for large λ , the barrier region becomes wider and the gradient descent may converge to a local minimizer away from the optimal result.

Original: For small λ , 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
Revised: For small λ , 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

Original: Monge problem.
Revised: Monge problem.

Original: The choice of λ needs to take this tradeoff into considerations.
Revised: The choice of λ needs to take this tradeoff into considerations.

Original: Training using a one layer linear neural network confirms the above analysis.
Revised: Training using a one layer linear neural network confirms the above analysis.

Original: ●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■
Revised: ●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●●▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■

Original: KL
Revised: KL

Original: Loss
Revised: Loss

Original: Transportcost    10 λ 0.
Revised: Transportcost    10 λ 0.

Original: Training with neural networks with nonlinear activation function.
Revised: Training with neural networks with nonlinear activation function.

Original: Next we present our results on the training of Gaussian to Gaussian distributions with neural networks with nonlinear activation function.
Revised: Next we present our results on the training of Gaussian to Gaussian distributions with neural networks with nonlinear activation function.

Original: We take X 1 = 0 d , X 1 = 1 d and Σ 1 = I d , Σ 2 = 3 I d + 1 d × d .
Revised: We take X 1 = 0 d , X 1 = 1 d and Σ 1 = I d , Σ 2 = 3 I d + 1 d × d .

Original: Theoretical analysis gives that the optimal transport distance is 2 d (see Theorem 2).
Revised: Theoretical analysis gives that the optimal transport distance is 2 d (see Theorem 2).

Original: Since here we use the D KL as distance function which is always positive, so the QP method behaves similarly as the SL method with different multiplier.
Revised: Since here we use the D KL as distance function which is always positive, so the QP method behaves similarly as the SL method with different multiplier.

Original: Hence the results of the QP method is not presented here.
Revised: Hence the results of the QP method is not presented here.

Original: The results are given in Table 1.
Revised: The results are given in Table 1.

Original: and the 784D Gaussian is taken for X 1 = 0 784 , X 1 = 2 · 1 784 with Σ 1 = Σ 2 = I(theoretical
Revised: and the 784D Gaussian is taken for X 1 = 0 784 , X 1 = 2 · 1 784 with Σ 1 = Σ 2 = I(theoretical

Original: result of the optimal transport distance is 784 ∗ 4 ).
Revised: result of the optimal transport distance is 784 ∗ 4 ).

Original: As can be seen from the table, all three algorithms give nice result and learn approximately the optimal transport map with a high accuracy.
Revised: As can be seen from the table, all three algorithms give nice result and learn approximately the optimal transport map with a high accuracy.

Original: Training is done by using a 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.
Revised: Training is done by using a 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.

Original: Remarkably, for the highly correlated distributions in high dimensions, our method gives a nice result.
Revised: Remarkably, for the highly correlated distributions in high dimensions, our method gives a nice result.

Original: ADMM learns a lower transport cost.
Revised: ADMM learns a lower transport cost.

Original: From the figure, we can see that the ADMM method learns a lower transport cost compared to other methods.
Revised: From the figure, we can see that the ADMM method learns a lower transport cost compared to other methods.

Original: This is because during training, minimizing the
Revised: This is because during training, minimizing the

Original: [Table]
Revised: [Table]

Original: D KL between the generated and target distribution may not converge to the solution to the Monge problem, as illustrated in the simple 2D case above.
Revised: D KL between the generated and target distribution may not converge to the solution to the Monge problem, as illustrated in the simple 2D case above.

Original: Hence, the splitting feature enables ADMM to learn a transport map with lower transport cost.
Revised: Hence, the splitting feature enables ADMM to learn a transport map with lower transport cost.

Original: This can also be seen from the learning curve in Figure 2.
Revised: This can also be seen from the learning curve in Figure 4 in Appendix A.5.

Original: The transport of the second network ( T θ 2 ) has a higher transport cost when learning the target distribution, whereas the first network ( T θ 1 has a lower transport cost, while the learned target distribution remains accurate.
Revised: The transport of the second network ( T θ 2 ) has a higher transport cost when learning the target distribution, whereas the first network ( T θ 1 has a lower transport cost, while the learned target distribution remains accurate.

Original: 4. G AUSSIAN TO G AUSSIAN MIXTURES
Revised: 4. G AUSSIAN TO G AUSSIAN MIXTURES

Original: We take a four component Gaussian mixture, each with variance matrix 0 .
Revised: We take a four component Gaussian mixture, each with variance matrix 0 .

Original: 5 I 2 , centers lying on the four corners of the square [ − 1 , 1] 2 and learn the optimal transport map between two dimensional standard Gaussian to this mixture distribution.
Revised: 5 I 2 , centers lying on the four corners of the square [ − 1 , 1] 2 and learn the optimal transport map between two dimensional standard Gaussian to this mixture distribution.

Original: We plot the Jacobian graph of the learned map and the value of the D KL between target distribution and the prediction by the network in Figure 3.
Revised: We plot the Jacobian graph of the learned map and the value of the D KL between target distribution and the prediction by the network in Figure 2.

Original: As can be seen from the figure, solely minimizing the D KL does not fulfil the right directions of the optimal transport maps.
Revised: As can be seen from the figure, solely minimizing the D KL does not fulfil the right directions of the optimal transport maps.

Original: Using our method, the learned map is more balanced and approximately satisfies the condition ∇× Tx = 0 .
Revised: Using our method, the learned map is more balanced and approximately satisfies the condition ∇× Tx = 0 .

Original: Note that here we donot include the penalty of |∇× Tx | 2 .
Revised: Note that here we donot include the penalty of |∇× Tx | 2 .

Original: From the learning curves, we confirm the findings above that ADMM learns a lower transport cost than the other methods.
Revised: From the learning curves, we confirm the findings above that ADMM learns a lower transport cost than the other methods.

Original: 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.
Revised: 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.

Original: AL and ADMM methods are more robust.
Revised: AL and ADMM methods are more robust.

Original: Compared to the SL, AL gives a more robust result with respect to the value of λ .
Revised: Compared to the SL, AL gives a more robust result with respect to the value of λ .

Original: For a well chosen λ , SL performs as good as AL.
Revised: For a well chosen λ , SL performs as good as AL.

Original: However, if λ is not taken properly, the obtained transport cost will be higher or the target distribution is not well realized.
Revised: However, if λ is not taken properly, the obtained transport cost will be higher or the target distribution is not well realized.

Original: For ADMM method, we can see the choice of ρ affects the training process, but in a big range, the choice of ρ has little effects on the final result (see appendix A.5 for the graphs indicating this finding).
Revised: For ADMM method, we can see the choice of ρ affects the training process, but in a big range, the choice of ρ has little effects on the final result (see appendix A.5 for the graphs indicating this finding).

Original: This confirms the benefits of the AL and ADMM method in the literature (Nocedal &
Revised: This confirms the benefits of the AL and ADMM method in the literature (Nocedal &

Original: Wright (1999); Boyd et al. (2011)).
Revised: Wright (1999); Boyd et al. (2011)).

Original: 4. I MPROVED TRANSPORT COST OF GAN S
Revised: 4. I MPROVED TRANSPORT COST OF GAN S

Original: We test our method on the Wasserstein GANs and show our method lowers the transport cost of 
Revised: We use 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.

Original: WGANs.
Revised: The constraint f w is a Lipschitzfunction is important in WGAN.

Original: Gaussian mixtures.
Revised: Without this constraint, the training of WGAN may face vanishing gradients issue, as illustrated in Figure 6 in Appendix A.6. However, adding the transport cost to the cost function as in (12), the training no longer face this difficulty and the target distribution could be learned.

Original: First we  use  the Wasserstein GANs to learn the Gaussian mixture distributions described in the previous subsection. We train a GAN with a  generator  and  discriminator both with 4 linear layers, each hidden layer with width 400 and activation function ReLU.
Revised: Note that here we donot use a penalty term on the discriminator network as the gradient penalty method, the transport cost is only a 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 in section 3.

Original:  The generated samples are plotted in Figure 4. As can be seen from the results, classical WGAN gives a transport cost 2.5054, and theWGAN with SL  and ADMM give significantly lower values, with 0.0915 and 0.0758, respectively.
Revised: 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 is sensitive to the choice of λ , while the AL and ADMM methods are more robust.

Original: MNIST.
Revised: MNIST.

Original: We also train the WGANs on the  dataset.
Revised: We also train the WGANs on the MNIST dataset.

Original: By using WGAN with SL( λ = 1 ), we obtain a transport cost around 1.81 compared to 1.87 with only WGAN.
Revised: By using WGAN with SL( λ = 1 ), we obtain a transport cost around 1.81 compared to 1.87 with only WGAN.

Original: MNIST like samples generated by the learned optimal is plotted in Figure 5.
Revised: MNIST like samples generated by the learned optimal is plotted in Figure 3.

Original: Therefore, our method lowers the transport cost of WGAN while keeping the qualify of the generated distributions.
Revised: Therefore, our method lowers the transport cost of WGAN while keeping the qualify of the generated distributions.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We have shown that the incorporation of constraint optimization tools provides a direct and efficient way for computing optimal transport maps.
Revised: We have shown that the incorporation of constraint optimization tools provides a direct and efficient way for computing optimal transport maps.

Original: By solving the Monge problem directly, our method avoids using special network structures or solving the dual problem.
Revised: By solving the Monge problem directly, our method avoids using special network structures or solving the dual problem.

Original: Moreover, applying our method to WGAN shows a lower transport cost for the generative networks without sacrificing the quality of the generated data.
Revised: Moreover, applying our method to WGAN shows a lower transport cost for the generative networks without sacrificing the quality of the generated data.

Original: A A PPENDIX A. P ROOF OF T HEOREM
Revised: A A PPENDIX A. P ROOF OF T HEOREM

Original: The proof of Theorem 1 is given below by combing the existence theorem for the Monge problem and the universal approximation of neural networks.
Revised: The proof of Theorem 1 is given below by combing the existence theorem for the Monge problem and the universal approximation of neural networks.

Original: Proof.
Revised: Proof.

Original: First given two probability measures µ, ν satisfying the assumptions of Theorem 1, then (Villani, 2009, Theorem 5.20) implies that there exists a unique deterministic solution T to the
Revised: First given two probability measures µ, ν satisfying the assumptions of Theorem 1, then (Villani, 2009, Theorem 5.20) implies that there exists a unique deterministic solution T to the

Original: Monge problem and the map T is a Borel map.
Revised: Monge problem and the map T is a Borel map.

Original: Then by Lusin’s theorem, there exists a continuous function ˜ T ∈ C ( R d )
Revised: Then by Lusin’s theorem, there exists a continuous function ˜ T ∈ C ( R d )

Original: such that ˜ T
Revised: such that ˜ T

Original: =
Revised: =

Original: T on any compact subset X ⊂ R d almost everywhere.
Revised: T on any compact subset X ⊂ R d almost everywhere.

Original: Since
Revised: Since

Original: T θ is given by a arbitrary depth neural network, we can use the universal approximation theorem in
Revised: T θ is given by a arbitrary depth neural network, we can use the universal approximation theorem in

Original: Kidger & Lyons (2020) to conclude:
Revised: Kidger & Lyons (2020) to conclude:

Original: For any ε > 0 , there exists T ˆ θ such that sup x ∈X | T ˆ θ ( x ) − ˜ T | ≤ ε .
Revised: For any ε > 0 , there exists T ˆ θ such that sup x ∈X | T ˆ θ ( x ) − ˜ T | ≤ ε .

Original: Therefore, d ( T ˆ θ # µ | ν ) =
Revised: Therefore, d ( T ˆ θ # µ | ν ) =

Original: d ( T ˆ θ # µ | T # µ ) ≤ d ( T ˆ θ # µ | ˜ T # µ ) + d ( T # µ | ˜ T # µ ) ≤ 2 ε .
Revised: d ( T ˆ θ # µ | T # µ ) ≤ d ( T ˆ θ # µ | ˜ T # µ ) + d ( T # µ | ˜ T # µ ) ≤ 2 ε .

Original: Fix α and take ε < α/ 2 , since d ( T ˆ θ # µ | ν ) ≤
Revised: Fix α and take ε < α/ 2 , since d ( T ˆ θ # µ | ν ) ≤

Original: α , we can get that the solution T
Revised: α , we can get that the solution T

Original: θ ∗ α to the problem (3) satisfies
Revised: θ ∗ α to the problem (3) satisfies

Original: [Equation]
Revised: [Equation]

Original: Let { α k } be a sequence converging to zero and T θ k be the solution to the corresponding problem (3), then d ( T θ k # µ | ν ) ≤ α k → 0 implies T θ k # µ → T θ ∗ # µ = T # µ .
Revised: Let { α k } be a sequence converging to zero and T θ k be the solution to the corresponding problem (3), then d ( T θ k # µ | ν ) ≤ α k → 0 implies T θ k # µ → T θ ∗ # µ = T # µ .

Original: Since T solves the problem (1), we have
Revised: Since T solves the problem (1), we have

Original: [Equation]
Revised: [Equation]

Original: On the other hand, by taking limit α → 0 in (11), we have E x ∼ µ [ | x − T θ ∗ x | 2 ] ≤ E x ∼ µ [ || x − T θ x | 2 ] .
Revised: On the other hand, by taking limit α → 0 in (15), we have E x ∼ µ [ | x − T θ ∗ x | 2 ] ≤ E x ∼ µ [ || x − T θ x | 2 ] .

Original: Combing this with the previous inequality implies that E x ∼ µ [ | x − T θ x | 2 ] =
Revised: Combing this with the previous inequality implies that E x ∼ µ [ | x − T θ x | 2 ] =

Original: E x ∼ µ [ || x − T θ ∗ x | 2 ] .
Revised: E x ∼ µ [ || x − T θ ∗ x | 2 ] .

Original: By the unique existence of the Monge map, T θ ∗ = T .
Revised: By the unique existence of the Monge map, T θ ∗ = T .

Original: Moreover, we have by the convergence of
Revised: Moreover, we have by the convergence of

Original: T θ α → T θ ∗
Revised: T θ α → T θ ∗

Original: [Equation]
Revised: [Equation]

Original: A. P ROOF OF T HEOREM
Revised: A. P ROOF OF T HEOREM

Original: Theorem 2.
Revised: Theorem 2.

Original: Let µ = N ( X 0 , Σ 0 ) and ν = N ( X 1 , Σ 1 ) be two multivariate normal distributions with X 0 , X 1 ∈ R d and Σ 0 , Σ 1 ∈ R d × d .
Revised: Let µ = N ( X 0 , Σ 0 ) and ν = N ( X 1 , Σ 1 ) be two multivariate normal distributions with X 0 , X 1 ∈ R d and Σ 0 , Σ 1 ∈ R d × d .

Original: Let T θ : R d (cid:55)→ R d be an linear operator defined by y
Revised: Let T θ : R d (cid:55)→ R d be an linear operator defined by y

Original: =
Revised: =

Original: Ax + b with θ =
Revised: Ax + b with θ =

Original: ( A, b ) and A ∈ R d × d , b ∈ R d .
Revised: ( A, b ) and A ∈ R d × d , b ∈ R d .

Original: Then there exists a unique solution to the following problem
Revised: Then there exists a unique solution to the following problem

Original: [Equation]
Revised: [Equation]

Original: and the corresponding transport map T θ ∗ is the optimal transport map in the sense of (1).
Revised: and the corresponding transport map T θ ∗ is the optimal transport map in the sense of (1).

Original: Proof.
Revised: Proof.

Original: For x ∼ µ , the linear transformation y
Revised: For x ∼ µ , the linear transformation y

Original: =
Revised: =

Original: Ax + b also satisfies a multivariate normal distribution, given by
Revised: Ax + b also satisfies a multivariate normal distribution, given by

Original: [Equation]
Revised: [Equation]

Original: Using the formula for KL divergence between two multivariate normal distributions, we can get that the KL-divergence between ρ and ν is
Revised: Using the formula for KL divergence between two multivariate normal distributions [Equation] we can get that the KL-divergence between ρ and ν is

Original: [Equation]
Revised: [Equation]

Original: Due to the fact that tr B − log det B + d ≥ 0 for any positive definite matrix B > 0 and ( X 1 − (AX 0 + b )) T Σ − 1 1 ( X 1 − ( AX 0 + b )) ≥ 0 , KL ( ρ | ν ) = 0 is equivalent to
Revised: Due to the fact that tr B − log det B + d ≥ 0 for any positive definite matrix B > 0 and ( X 1 − (

Original: [Equation]
Revised: [Equation]

Revised: ≥

Original: [Equation]
Revised: [Equation]

Original: The objective function in (13) can be calculated by
Revised: , KL ( [Equation] | ν ) = 0 is equivalent to

Original: [Equation]
Revised: [Equation]

Revised: The objective function in (17) can be calculated by

Original: [Equation]
Revised: [Equation]

Original: Therefore, the optimization problem (13) becomes 
Revised: Therefore, the optimization problem (17) becomes min

Original: [Equation]
Revised: [Equation]

Original: Eliminating b , this is equivalent to
Revised: Eliminating b , this is equivalent to

Original: [Equation]
Revised: [Equation]

Original: Since tr (( I − A ) T ( I − A )Σ 
Revised: Since tr (( I − A ) T ( I − A )Σ 0 ) =

Original: [Equation]

Original: ) = tr (Σ [Equation] [Equation]
Revised: tr (Σ 0 ) − 2 tr ( A T Σ 0 )

Original:  the above problem is equivalent to
Revised: + tr ( A Σ 0 A T ) = tr (Σ 0 ) − 2 tr ( A T Σ 0 ) + d , the above problem is equivalent to

Original: [Equation]
Revised: [Equation]

Original: Let R = Σ 12 

Revised: Let R = Σ 12

Original: A T Σ − 12 1 , the above problem could be rewritten as
Revised: A T Σ − 12 1 , the above problem could be rewritten as

Original: [Equation]
Revised: [Equation]

Original: Note that the optimization is over all orthogonal matrix in R d × d .
Revised: Note that the optimization is over all orthogonal matrix in R d × d .

Original: The solution is given by
Revised: The solution is given by

Original: R ∗ =
Revised: R ∗ =

Original: V T where Σ 12 1 Σ 12 0 =
Revised: V T where Σ 12 1 Σ 12 0 =

Original: V P is the polar decomposition of Σ 12 1 Σ 12 0 and is given by V = Σ 1 2 1 Σ 12 0 ((Σ 12 1 Σ 12 0 ) T Σ 12 1 Σ 12 0 ) − [-
-]
Revised: V P is the polar decomposition of Σ 12 1 Σ 12 0 and is given by V = Σ 12 1 Σ 12 0 ((Σ 12 1 Σ 12 0 ) T Σ 12 1 Σ 12 0 ) − [-
-]12 =

Original: = Σ 1 2 1 Σ 12 0 (Σ 12 0 Σ 1 Σ 12 0 ) − 12 and P = (Σ 1 2 0 Σ 1 Σ 12 0 )
Revised: Σ 12 1 Σ 12 0 (Σ 12 0 Σ 1 Σ 12 0 ) − 12 and P = (Σ 12 0 Σ 1 Σ 12 0 )

Original: 12 Recalling R ∗ =
Revised: 12 Recalling R ∗ =

Original: Σ 12 0 ( A ∗ ) T Σ − 12 1 , we obtain
Revised: Σ 12 0 ( A ∗ ) T Σ − 12 1 , we obtain

Original: [Equation]
Revised: [Equation]

Original: (21) and by the constraint ( X 1 − AX 0 − b ) = 0 , we get
Revised: (25) and by the constraint ( X 1 − AX 0 − b ) = 0 , we get

Original: [Equation]
Revised: [Equation]

Original: The above formula ( A ∗ , b ∗ ) is thesame as that obtained by analytical method for the optimal transport and is shown to be the unique solution to the Monge problem (1).
Revised: The above formula ( A ∗ , b ∗ ) is the

Original: A. P ROOF OF THE CONVERGENCE OF THE ADMM METHOD
Revised: A. P ROOF OF THE CONVERGENCE OF THE ADMM METHOD

Original: Theorem 3.
Revised: Theorem 3.

Original: Assume T θ is given by a neural network with Lipschitz continuous activation functions and satisfying ∥ T θ ∥ → ∞ for θ → ∞ .
Revised: Assume T θ is given by a neural network with Lipschitz continuous activation functions and satisfying ∥ T θ ∥ → ∞ for θ → ∞ .

Original: For sufficiently large ρ , the ADMM method generates a bounded sequence that converges to the stationary point of the Lagrangian
Revised: For sufficiently large ρ , the ADMM method generates a bounded sequence that converges to the stationary point of the Lagrangian

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We only need to check the conditions in the reference Wang et al.
Revised: We only need to check the conditions in the reference Wang et al.

Original: First we show the objective function is coercive.
Revised: First we show the objective function is coercive.

Original: For bounded θ 1 = θ 2 and θ 1 → ∞ , then E x ∼ µ (cid:2) | x − T θ 1 x | 2 (cid:3) + η ε ( d ( T θ 2 # µ | ν ))
Revised: For bounded θ 1 = θ 2 and θ 1 → ∞ , then E x ∼ µ (cid:2) | x − T θ 1 x | 2 (cid:3) + η ε ( d ( T θ 2 # µ | ν ))

Original: = ∞ and thus is coercive.
Revised: = ∞ and thus is coercive.

Original: The Lipschitz continuous assumption implies the objective function is also Lipschitz continuous.
Revised: The Lipschitz continuous assumption implies the objective function is also Lipschitz continuous.

Original: Moreover, the solution to the subproblem min θ 1 E x ∼ µ (cid:2) | x − T θ 1 x | 2 (cid:3) and min θ 2 η ε ( d ( T θ 2 # µ | ν )) is also Lipschitz continuous.
Revised: Moreover, the solution to the subproblem min θ 1 E x ∼ µ (cid:2) | x − T θ 1 x | 2 (cid:3) and min θ 2 η ε ( d ( T θ 2 # µ | ν )) is also Lipschitz continuous.

Original: Thus the conditions in Wang et al.
Revised: Thus the conditions in Wang et al.

Original: (2019) holds and we can conclude that the ADMM method converges.
Revised: (2019) holds and we can conclude that the ADMM method converges.

Original: A. T HE KKT CONDITIONS
Revised: A. T HE KKT CONDITIONS

Original: The KKT condition is the conditions for the saddle point of the Lagrangian.
Revised: The KKT condition is the conditions for the saddle point of the Lagrangian.

Original: The KKT condition for the SL Lagrangian (4).
Revised: The KKT condition for the SL Lagrangian (4).

Original: At the saddle point ( θ ∗ , λ ∗ ) , the KKT conditions for (4) are
Revised: At the saddle point ( θ ∗ , λ ∗ ) , the KKT conditions for (4) are

Original: [Equation]
Revised: [Equation]

Original:  At each training step (taking ρ to be constant), the KKT conditions for the QP loss (5) are 
Revised: ( T θ # µ | ν ))( θ ∗ ) = 0 , ∇ λ L SP ( θ ∗ , λ ∗ ) = d ( T θ ∗ # µ | ν ) = 0 . At each training step (taking ρ to be constant), the KKT conditions for the QP loss (5) are ∇ θ L QP ( θ k , ρ k ) = ∇ θ E x ∼ µ ( (cid:2) | x − T θ x

Original: [Equation]
Revised: | 2 (cid:3) ) ( θ k ) + ρ k d ( T θ k # µ | ν )( ∇ θ d ( T θ # µ | ν )) ( θ k ) = 0 ,

Original: At each training step (taking ρ to be constant), the KKT conditions for the AL loss (6) are 
Revised: At each training step (taking ρ to be constant), the KKT conditions for the AL loss (6) are (cid:26) ∇ θ L AL

Original: [Equation]
Revised: ( θ k , λ k , ρ k ) = ∇ θ E x ∼ µ ( (cid:2) | x − T θ x | 2 (cid:3) )( θ k ) + ( λ k + ρ k d ( T θ k # µ | ν ))( ∇ θ d ( T θ # µ | ν )) ( θ k ) = 0 , ∇ λ L AL ( θ k , λ k , ρ k ) =

Revised: d ( T θ k µ | ν ) = 0 .

Original: A. T RAINING RESULTSWITH DIFFERENT PARAMETERS
Revised: A. T RAINING RESULTS

Revised: The results of 78D Gaussian is plotted below.

Original: We demonstrate the robustness of AL and ADMM method compared to the SL method.
Revised: We demonstrate the robustness of AL and ADMM method compared to the SL method.

Original: The training curves for these three methods are plotted in Figure 
Revised: The training curves for these three methods are plotted in Figure 5.

Revised: A. E XPERIMENT DETAILS

Revised: Table 1 and Figure 2, 3 are produced using a normalizing flow network with planar transformation layers.

Revised: The initial and target densities are known and samples are drawn from these distributions.

Revised: Algorithms described in sections 3.1-3.3 are used.

Revised: The KL divergence is computed via (11).

Revised: For the WGAN example, we use a fully connected neural network with hidden layers of width 400.

Revised: Gradient penalty is used for the WGAN-GP and for the SL, AL and ADMM methods, the alogrithm is described in section 3.4 and λ = 1 is used in the SL and ρ = 10 − 4 is used for AL and ρ = 10 −is used for the ADMM.