On amortizing convex conjugates for optimal transport
Authors: Brandon Amos
Venue: ICLR 2023 poster (2023)
Keywords: optimal transport
wasserstein-2
convex conjugate
c-transform
amortized optimization

Abstract: This paper focuses on computing the convex conjugate operation that arises when solving Euclidean Wasserstein-2 optimal transport problems. This conjugation, which is also referred to as the Legendre-Fenchel conjugate or c-transform,is considered difficult to compute and in practice,Wasserstein-2 methods are limited by not being able to exactly conjugate the dual potentials in continuous space. To overcome this, the computation of the conjugate can be approximated with amortized optimization, which learns a model to predict the conjugate. I show that combining amortized approximations to the conjugate with a solver for fine-tuning significantly improves the quality of transport maps learned for the Wasserstein-2 benchmark by Korotin et al. (2021a) and is able to model many 2-dimensional couplings and flows considered in the literature. All of the baselines, methods, and solvers in this paper are available at http://github.com/facebookresearch/w2ot.

Published as a conference paper at ICLR 2023
ON AMORTIZING CONVEX CONJUGATES FOR OPTIMAL TRANSPORT
Brandon Amos
Meta AI
ABSTRACT
This paper focuses on computing the convex conjugate operation that arises when solving Euclidean Wasserstein-2 optimal transport problems. This conjugation, which is also referred to as the Legendre--Tnnceel conjugate or c-transform, is considered difficult to compute and in practice, Wasserstein-2 methods are lim- ited by not being able to exactly conjugate the dual potentials in continuous space. To overcome this, the computation of the conjugate can be approximated with amortized optimization, which learns a model to predict the conjugate. show that combining amortized approximations the conjugate with solver for fine-tuning significantly improves the quality of transport maps learned for the Wasserstein-2 benchmark by Korotin al. (2021a) and is able to model many 2-dimensional couplings and flows considered in the literature. All of the base- lines, methods, and solvers in this paper are available at http::/ql
INTRODUCTION
Optimal transportation (Villani, 2009; Ambrosio, 2003; Santambrogio, 2015; Peyré et al., 2019) thriving area of research that provides way of connecting and transporting between proba- bility measures. While optimal transport between discrete measures is well-understood, e.g. with Sinkhorn distances (Cuturi, 2013), optimal transport between continuous measures an open re- search topic actively being investigated (Genevay et al., 2016; Seguy et al., 2017; Taghvaei and Jalali, 2019; Korotin et al., 2019; Makkuva et al., 2020; Fan et al., 2021; Asadulaev et al., 2022). Continuous OT has applications in generative modeling (Arjovsky al., 2017; Petzka et al., 2017; Wu et al., 2018; Liu et al., 2019; Cao et al., 2019; Leygonie et al., 2019), domain adaptation (Luo et al., 2018; Shen et al., 2018; Xie et al., 2019), barycenter computation (Li al., 2020; Fan al., 2020; Korotin et al., 2021b), and biology (Bunne et al., 2021; 2022; Lübeck et al., 2022).
This paper focuses on estimating the Wasserstein-2 transport map between measures a and / in Euclidean space, i.e. supp(a) = -Supp(3) = TR" with the Euclidean distance as the transport cost. The Wasserstein-2 transport map, T R + R is the solution to Monge's 's primal formulation:
i arginf E | - T(x)
T)) z~a
where T(a, 3): (T: 33 is the set of admissible couplings and the push-forward operator # is defined by TB) = (B)) for measure a, measurable map T. and all measurable sets B. 7 exists and is unique under general settings, e.g. as in Santambrogio (2015, Theorem 1.17), and is often difficult to solve because of the coupling constraints T. Almost every computational method instead solves the Kantorovich dual, e.g. as formulated in Villani (2009, $5) and Peyré al. (2019, $2.5). This paper focuses the dual associated with the negative inner product cost (Villani, 2009, eq. 5.12), which introduces a dual potential function f : IR" + IR and solves:
i E arg - EE [f(x)) fel+(o) z~a
E E [f)( g~3
where I. (a:) is the space of measurable functions that are convex conjugate, Segendre-Fennhhee transform, of a function f defined by:
over a and f the
(w) inf ( with objective Jy(2;y) f(x) - (xx ().)). ZEX
(3)

---

Published as a conference paper at ICLR 2023
{yy) denotes optimal solution to eq. (3). Even though the eq. (2) searches over functions in L' (a), the optimal dual potential f is convex (Villani, 2009, theorem 5.10). When one of the measures has a density, Brenier (1991, theorem 3.1) and McCann (1995) relate f to an optimal transport map T for the primal problem in eq. (1) with (xx) = v=f(2). and the inverse to the transport map is given by 7- (v) - V g.))
stream of foundational papers have proposed methods to approximate the dual potential f with a neural network and learn it by optimizing eq. (2): Taghvaci and Jalali (2019); Korotin al. (2019); Makkuva et al. (2020) parameterize f as an input-convex neural network (Amos al., 2017), which can universally represent any convex function with enough capacity (Huang et al., 2020). Other works explore parameterizing f as non-convex neural network (Nhan Dam et al.,
2019; Korotin et al., 2021a; Rout et al., 2021).
Efficiently solving the conjugation operation in eq. (3) is the key computational challenge to solving the Kantorovich dual in eq. (2) and is an important design choice. Exactly computing the conjugate as done in Taghvaci and Jalali (2019) is considered computationally challenging and approximating it as in Korotin et al. (2019); Makkuva et al. (2020); Nhan Dam et al. (2019); Korotin et al. (2021a); Rout et al. (2021) may be instable. Korotin et al. (2021a) fortifies this observation:
The [exact conjugate] solver is slow since each optimization step solves a hard subproblem for computing [the conjugate]. [Solvers that approximate the conjugate] are also hard to optimize: they either diverge from the start or diverge after converging nearly-optimal saddle point.
In contrast to these statements on the difficulty of exactly estimating the conjugate operation, I will show in this paper that computing the (near-)exact conjugate easy. My key insight is that the approximate, i.e. amortized, conjugation methods can be combined with fine-tuning procedure using the approximate solution as a starting point. Sect. 3 discusses the amortization design choices and sect. 3.2.2 presents a new amortization perspective the cycle consistency term used in Wasserstein-2 generative networks (Korotin et al., 2019), which was previously not seen in this way. Sect. 5 shows that amortizing and fine-tuning the conjugate results in state-of-the-art performance in all of the tasks proposed in the Wasserstein-2 benchmark by Korotin et al. (2021a). Amortization with fine-tuning also nicely models synthetic settings (sect. 6), including for learning fingleblock potential flow without using the likelihood.
2 LEARNING DUAL POTENTIALS: A CONJUGATION PERSPECTIVE
This section reviews the standard methods of learning parameterized dual potentials to solve eq. (2). The first step is to re-cast the Kantorovich dual problem eq. (2) as being over a parametric family of potentials fe with parameter as an input-comee convex neural network (Amos et al., 2017) or a more general non-convex neural network. Taghvaei and Jalali (2019); Makkuva et al. (2020) have laid the foundations for optimizing the parametric potentials for the dual objective with
max (0) where ((0) == E [fa
E (v)] = E [fo(z]++ E (*v)), (4) yrー
grog
where is the objective to the conjugate optimization problem in eq. (3), x(y) is the solution to the convex conjugate, and eq. (4) assumes finite solution to eq. (2) exists and replaces the sup with max. Taghvaci and Jalali (2019) show that the model can be learned, i.e. the optimal parameters can be found, by taking gradient steps of the dual with respect to the parameters of the potential, ie. using V.V. This derivative going through the loss and conjugation operation can be obtained by applying Danskin's 's envelope theorem (Danskin, 1966; Bertsekas, 1971) 1) and results in only needing derivatives of the potential:
VoV(O) Vo [- E [b(())] E. JJ (+,)) T~D
w~B
(5)
[V)fo(=)] + E. (ofo((()))
a~a
w~3
where i(y) is not differentiated through.
Assumption A standard assumption is that the conjugate is smooth with a well-defined arg min. This has been shown to hold when f is strongly convex. e.g. in Kakade et al. (2009), or when f is essentially strictly convex (Rockafellar, 2015, theorem 26.3).
2

---

Published as a conference paper at ICLR 2023
In practice, assumption 1 in not guaranteed, e.g. non-convex potentials may have parameterization that results in the conjugate taking infinite values in regions. The dual objective in eq. (2) and eq. (4) discourage the conjugate from diverging the supremum involves the negation of the conjugate.
Remark 1 In eq. (4), the dual potential associated with the measure's constraints is the central object that is parameterized and learned, and the dual potential associated with the 3 measure is given by the conjugate f and does not require separately learning. Because of the symmetry of eq. (1), the order can also be reversed as in Korotin et al. (2021b) so that the duals associated with the measure are the ones directly parameterized, but we will not consider doing this. Potentials associated with both measures can also be purameterized and we will next see that it is the most natural to think about the model associated with the conjugate as an amortization model.
The dual objective V can be upper-bounded by replacing ž with any approximation because any sub-optimal solution to the conjugation objective provides an upper-bound to the true objective, iee J(E(y): S J(r;y) y) for all x. In practice, maximizing a loose upper-bound can cause significant divergence issues the potential can start over-optimizing the objective.
Remark 2
Computing the updates the dual potential's parameters in eq. (5) is a well-defined machine learn- ing setup given parameterization of the potential fa. but is often computationally bottlenecked by the conjugate operation. Because of this bottleneck. many existing work resorts to amortizing the conjugate by predicting the solution with a model o I overview the design choices behind amortizing the conjugate in sect. 3, and then go on in sect. 4 to show that it is reasonable to fine-tune the amortized predictions with an explicit solver CONJUGATE(S Tinit = (v)). Algorithm summarizes how to learn a dual potential with an amortized and fine-tuned conjugate.
AMORTIZING CONVEX CONJUGATES: MODELING AND LOSSES
This section scopes to predicting an approximate solution to the conjugate optimization problem in eq. (3). This is an instance of amortized optimization methods which predict the solution to family of optimization problems that are repeatedly solved (Shu, 2017: Chen et al., 2021; Amos, 2022). Amortization is sensible here because the conjugate is repeatedly solved for y 3 every time the dual V from eq. (4) is evaluated across a batch. Using the basic setup from Amos (2022), I call a prediction the solution of eq. (3) the amortization model (v). which is parameterized by some p.. The goal is to make the amortization model's prediction match the true conjugate solution, i.e. To( H(y), for samples y B. In other words, amortization uses model simultaneously solve all of the conjugate optimization problems. There are two main design choices: sect. 3.1 discusses parameterizing the amortization model and sect. 3.2 overviews amortization losses.
3.1 PARAMETERIZING A CONJUGATE AMORTIZATION MODEL
The amortization model (y) maps point y 3 to solution to the conjugate in eq. (3), i.c. i.p : R + Rn and the goal is for (v) x ă(y). In this paper, I take standard potential models further described in app. B and keep them fixed to ablate across the amortization loss and fine-tuning choices. The main categories are:
1. R.p : R - Th directly maps to the solution of eq. (3) with a multilayer perceptron (MLP) as in Nhan Dam al. (2019), or a U-Net (Ronneberger et al., 2015) for image-based trans- port. These are also used in parts of Korotin et al. (2021a).
2. S.p = V Vyg is the gradiens of a function g. TR" - TR. Korotin et al. (2019); Makkuva al. (2020) parameterize Ie
input-convex neural network, and some methods of Korotin et al. (2021a) parametrize ge as a ResNet (He et al., 2016). This is well-motivated because the arg min of convex conjugate is the derivative, i.e. z(y) v,)
3

---

Published as a conference paper at ICLR 2023
Algorithm Learning Wasserstein-2 dual potentials with amortized and fine-tuned conjugation
Inputs: Measures a and / to couple, initial dual potential fa, and initial amortization model T.p
while unconverged do
Sample batches {x,} ~ and {y, 3 indexed by j E [N]
Obtain the amortized prediction of the conjugate (1z)
Fine-tume the prediction by numerically solving T(y;) CONJUGATE(F Vj, Winit =x11))))
Update the potential with a gradient estimate of the dual in eq. (5), i.e. VoV
Update the amortization model with gradient estimate of loss from sect, 3, i.e. v end while
return optimal dual potentials fo and conjugate amortization model Gip
3.2 CONJUGATE AMORTIZATION LOSS CHOICES
4
|| x (v)
Regression
We now turn the design choice of what loss to optimize so that the conjugate amortization model é best-pred west-predicts the solution to the conjugate. In all cases, the loss dif- ferentiable and yp is optimized with a gradient-based op- timizer. present an amortization perspective of methods not previously presented as amortization methods, which useful to help think about improving the amortized pre- and exact solvers in sect. 4.
J,(2; y) Objective
xx (v)
α V J
Cycle
dictions with the Figure illustrates the main loss choices.
Figure 1: Conjugate amortization losses.
3.2.1 OBJECTIVE-BAN BASEDD AMORTIZATION
Nhan Dam et al. (2019) propose to make the amortized prediction optimal on the conjugation ob- jective Jj from eq. (3) across samples from 3, i.e.:
min (V)) where (ob)(() == E Jy (T(T)
(6)
gmo
We refer to Lobj as objective-based amortization and solve eq. (6) by taking gradient steps V Vooobi using a Monte-Carlo estimate of the expectation.
The maximin method proposed in Makkuva et al. (2020, theorem 3.3) is equivalent to maximizing an upper-bound the dual loss V with respect 0 of a potential fo and minimizing the objective-based amortization loss Lobi with respect to yp of an amortization model ā.p := V Their formulation replaces the exact conjugate in eq. (4) with an approximation Tep i.e.:
Remark 3
max min (M sp) where "0. = - srwax
E J, [fo 4-gg"
(. (v) v)..
(7)
#
Makkuva et al. (2020) propose to optimize VMM with gradient ascent-descent steps. For optimizing 8. VMDa (f..p) up) is upper bound on the true dual objective V(0) as discussed in remark 2 with equality if and only if T.p = T. Evaluating the inner optimization step is exactly the objective- based amortization update, i.e., (0)
= V (T
Remark Suboptimal predictions of the conjugate often leads to a divergent upper bound on V()). Makkuva et al. (2020, algorithu 1) propose to fix this by running more updates the amortization model. In sect. 4, I propose fine-tune fne-tuning as an alternative to obtain a near-exace exact conjugates.
3.2.2 FIRST-ORDER OPTIMALITY AMORTIZATION: CYCLE CONSISTENCY AND W2GN
An alternative to optimizing the dual objective directly as in eq. (6) is to optimize for the first-order optimality condition. Eq. (3) is an unconstrained minimization problem, so the first-order optimality condition is that the derivative of the objective is zero, i.e. 5 y) vzf(æ) y 0. The conjugate amortization model can be optimized for the residual norm of this condition with
min (4p) where Eeycte(()) == E. ( 8~
EE, [Vas((,, v1.. (8)
4

---

Published as a conference paper at ICLR 2023
Remark 5 W2GN (Korotin et al., 2019) is equivalent to maximizing upper-bound to the dual loss V with respect to of a potential fo and minimizing the first-order amortization loss Leycle with respect to ip of conjugate amortization model 3. == Vg.- Korotin et al. (2019) originally motivated the cycle consistency term from the use in cross-domain generative modeling Zhu et al. (2017) and eq. (8) shows alternative way of deriving the eycle consistency term by amortizing the first-order optimality conditions of the conjugate.
Remark The formulation in Korotin et al. (2019) does not disconnect So when optimizing the cycle loss in eq. (8). From an amortization perspective, this performs amortization by updating fo to have a solution closer to T.g rather than the usual amortization setting of updating 3. to make a pre- diction closer to the solution of fo- In my experiments, updating So with the amortization term seems to help when not fine-tuning the conjugate to be exact, but not when using the exact conjugates.
Remark 7 Korotin et al. (2019) and followup papers such as Korolin et al. (2021b) state that they do not perform maximin optimization as in eq. from Makkuva el al. (2020) because they replace the inner optimization of the conjugate with an approximation. I disagree that the main distinction between these methods should be based on their formulation as a maximin optimization problem. I instead propose that the main difference between their losses is how they amortize the convex conjugate: Makkuva et al. (2020) use the objective- -based loss in eq. (6) while Korotin et al. (2019) use the first-order optimality condition (eq. (8)). Sect. 5 shows that adding fine-tuning and exact
conjugates to both of these methods makes their performance match in most cases.
Remark 8 Optimizing for the first-order optimality conditions may not be ideal for non-convex conjugate objectives as inflection points with near-zero derivative may not be global minimum of eg. (3). The left and right regions fig. illustrate this.
3.2.3 REGRESSION BBSSED AMORTIZATION
The previous objective and first-order amortization methods locally refine the model's prediction ing local derivative information. The conjugate amortization model can also be trained by regressing onto ground-truth solutions when they available, i.e.
min Freg (P)) where Lreg(() == E (T - H1y)
This regression loss is the most useful when approximations to the conjugate are computationally easy to obtain, e.g. with a method described in sect. 4. freg gives the amortization model informa- tion about where the globally optimal solution is rather than requiring it to only locally search over the conjugate's 's objective J.
NUMERICAL SOLVERS FOR EXACT CONJUGATES AND FINE TUNING
Algorithm 2 CONJUGATE(F y, Zinit)
In the Euclidean Wasserstein-2 setting, the conjugation oper- ation in eq. (3) is continuous and unconstrained optimiza- tion problem over a possibly non-convex potential f. It is usually implemented with method using first-order infor- mation for the update in algorithm 2, such as:
I f- Fini while unconverged do Update x with V=J{(<;;)
end while
return optimal H(y)= =
1. Adam (Kingma and Ba, 2014) is an adaptive first-order optimizer for high-dimensional optimization problems and is used for the exact conjugations in Korotin et al. (2021a).
Note: Adam here is for algorithm 2 and is not performing parameter optimization.
2. L-BFGS (Liu and Nocedal, 1989) is a quasi-Newton method for optimizing unconstrained convex functions. App. A discusses more implementation details behind setting up L- BFGS efficiently to run on the batches of optimization problems considered here. Choos- ing the line search method is the most crucial part the conditional nature of some line searches may be prohibitive over batches. Table shows that an Armijo search often works well to obtain approximate solutions.
5

---

Published as a conference paper at ICLR 2023
5 EXPERIMENTAL RESULTS ON THE WASSERSTEIN-2 BENCHMARK
I have focused most of the experimental investigations the W benchmark (Korotin et al., 202 because it provides concrete evaluation setting with established baselines for learning potentials for Euclidean Wasserstein-2 optimal transport. The tasks in the benchmark have known (ground-truit) truth) optimal transport maps and include transporting between: 1) nigh-diimnssinal (HD) mixtures of Gaussians. and 2) samples from generative models trained on CelebA (Liu al., 2015). The main evaluation metric is the unexplained variance percentage (C²-UVP) metric from (Korotin et al., 2019), which compares a candidate map T to the ground truth map T with:
-2-UVP(T;a, a. B) == 100 || - ( (Varta////a//// (B)%.)
(10)
In all of the experimental results, I report the final -VVUVP evaluated with 16384 samples at the end of training, and average the results over 10 trials. App. C further details the experimental setup. My original motivation for running these experiments was understand how ablating the amortization losses and fine-tuning options impacts the final UVP performance of the learned potential.
The main experimental takeaway of this paper is that fine-tuning the amortized conjugate with solver significantly improves the learned transport maps. Tables 1 and 2 report that amortizing and fine-tuning the conjugate improves the CV-UVP performance by a factor of 1.8 to 4.4 over the previously best-known results on the benchmark. App. C.3 shows that the conjugate can often be fine-tuned within 100ms per batch of 1024 examples on an NVIDIA Tesla V100 GPU, fig. 2 and app. C.2 compare Adam and L-BFGS for solving the conjugation. The following remarks further summarize the results from these experiments:
With fine-tuning, the choice of regression objective-based amortization doesn't 't sig- nificantly impact the L2-UVP of the final potential. This is because fine-tuning is usually able to find the optimal conjugates from the predicted starting points.
Remark 9
Remark 10 My re-implementation of W2GN (Korotin et al., 2019), which uses cycle consistency amortization with no fine-tuning often outperforms the results reported Korotin et al. (2021a). This is likely due to differences in the base potential and conjugate amortization models.
Remark 11 Cycle consistency sometimes provides difficult starting points for the fine-tuning meth- ods, especially for L-BFGS. When learning non-convex potentials, this poor performance is likely related to the fact that Newton methods are known to be difficult for saddle points (Dauphin et al., 2014). Combining cycle consistency, which tries find point where the derivative is zero, with L-BFGS, which also tries to find point where the derivative is zero, results in finding suboptimal inflection points of the potential rather than the true minimizer.
The performance of the methods using objective- bhased amortization without fine- tuning, as done Taghvaei and Jalali (2019), are lower than the performance reported Korotin et al. (2021a). This is because I do not run multiple inner updates to update the conjugate amor- tization model. instead advocate for fine-tuning the conjugate predictions with a known solver, eliminating the need for hyper-parameter of the number of inner iterations that needs to be del- icately tuned to make sure the amortized prediction alone does not diverge too much from the true conjugate.
Remark 12
D 64
D D= 128
D 256
1.00
Normalized Jive
0.75
0.50
0.25
0.00
10
20
40
50.0
10
20
30
40
50.0
20
30
10
Solver Iseration
Solver Trenation
Salver Treeation
Amortized Initialization + L-BSSS
L-BFGS
Adam
Figure 2: Conjugate solver convergence on the HD benchmarks with an ICNN potential.

---

Published as a conference paper at ICLR 2023
Table I: Comparison of -V-UPP on the high- dibe-siinnsi tasks from the Wasserstein-2 benchmark by Korotin et al. (2021a), where *[the gray tags] denote their results. report the mean and stan- dard deviation across 10 trials. Fine-tuning the amortized prediction with L-BFGS or Adam consistently improves the quality of the learned potential.
Buselines from Korotin et al. (2021a)
Amortization loss
Conjugate solver
17
2
11
21
16
11 m 32
11
64
12 - 128
X 256
[[22]
Cycle
None
0.1
0.7
2.6
3.3
6.0
7.2
2.0
2.7
[MMv
None
0.2
1.4
2.2
[MMv2]
Objective
None
0.1
0.68
2.2
3.1
5.3
10.1
3.2
2.7
[MMI
Objective
Nore
0.1
0.3
0.9
2.2
4.2
3.2
3.1
4.1
Potential model: input convex network deseribed app.
Amortization model: describeed spp.
Amertization
Conjugate solver
n 2
21
0 126
Cycle
None
0:28 +0.00
0.90 +0.11
2.23
+0.20
3.03
+0.05
5:32 =0.14
8.79 +:0.16
5.66
+0-55
4.34 -..14
Objective
None
0,27 10.09
0,78
10:12
1.78
1.0.26
2,00
10.11
>1000
2-00
>100
>100
Cycle
L.BPGS
+0.00
0.77
=0:1
+00228
#0.14
2.02 =0:10
+0089
#0:30
5.90
Olyectiiee
L-HHUS
0:25 +0.00
0.79
=0:12
1.63 +:0.3
1.12
10.11
1.92
:0:19
4.40
10070
1.64 20.11
2:24
=00.1
Regression
L-BFGS
0.26 10.09
0.78
10:12
1.61 1.022
1.14 10.12
1.93 10:20
4.41
10.74
1.09
10:11
2.21 10.15
Cycle
+0.00
0.79
+0.1
+0.29
+0.12
1.95 +0.31
+0.02
1.88 +0.25
Obycctive
Adam
0:26 10.00
0.79
20.14
1.62
1.0.1
1.08
10.14
1.89
00.10
4.23
1.0.7
1.59
10.12
1.99
10.15
Regressine
Adam
0.35 +0.07
0.81
+0.12
1.61
+:0.52
1,00
+0.0
1.85
-0.30
4.42
+-0.05
1.63
+0.08
1.00
+0.16
Potential model: non coneec network (MLP) described upp.
Amortization model:
described app.
Amortization
Conjugate solver
16
0 126
Cyele Objective
Nore
0.05 10:00
0.35 20001
1.51 10.08
>000
>100
>-00
>000
>100
None
>100
>100
>100
>000
>100
>100
>100
>100
Cyele
L.RPGS
5100
>100
>100
5100
>100
>100
5100
Objective
L-BPGS
0.03 10.00
0.22
10.01
0.60
10.03
0.80
10.11
2.09
10:31
2.08
100.0
0.67
10:05
0.59
10.00
Regression
L-BKGS
0.05 10.00
0.22
20.01
0.61
10.04
0.77
10.10 10
1.97
=0:38
2.08
10.50
0.67
20.05
0.65
20007
Cycle
40.0
0.89
#2.82
>100
>100
Objectivee
Adam
0.06 +0.01
0.26
=0.02
0.63
+0.07
0.81
+0.10
1.09
=0.32
2:21
+0.32
0.77
+0005
0.66 +0.07
Regression
Adam
0,22 10.01
0,28
10.00
0:61
1.007
0,80 10
2.07
10.38
2.37
1.0.08
0.77
10.06
0.75 1.000
Improvement factor over prior
3.3
1.8
Table 2: Comparison of C3-UVP on the CelebA64 tasks from the Wasserstein-2 benchmark by Korotin et al. (2021a), where *tthe gray tags] denote their results. I report the mean and standard deviation across 10 trials. Fine-tuning the amortized prediction with L-BFGS or Adam con- sistently improves the quality of the learned potential. The ConvICNN64 and ResNet potential models are from Korotin et al. (2021a), and app. B.5 describes the (non-convex) ConvNet model.
Amortization loss
Conjugate solver
Potential Model
Early Generator
Mid Generator
Late Generator
[[2]]
Cycle Objective
None
ConvICNN64
1.7
0.25
"[MM]
None
ResNet
2.2
0,9
0.53
*[MM-R]]
Objective
None
ResNet
1.4
0.4
0.22
Cycle
None
ConyNet
>100
26.50 4-60.14
0.29 4-0.59
Objective
None
ConvNet
>100
0.29 +0.15
0.69 4:0.90
Cycle
Adam
ConvNet
0.65 20.02
0.21 10.00
0.11 1:0.01
Cycle
L-BFGS
ConvNet
0.62 10.01
0.20 10000
0.09 1:0.0
Objective
Adam
ConyNet
0.65 +0002
0.21 +0.00
0.11 +:0.05
Objective
L-BFGS
ConvNet
0.61 =0.01
0.20 #0.00
0.09 2:000
Regression
Adam
ConvNet
0.66 2.001
0.21 4:0.0
0.12 00.00
Regression
L-BPGS
ConvNet
0.62
40000
0.09
Improvement factor over prior work
2.3
2.0
2.4
the reversed direction from Korotin et al. (2021a), i.e. the potential model associated with the 3 measure
6 DEMONSTRATIONS ON SYNTHETIC DATA
I lastly demonstrate the stability of amortization and fine-tuning as described in algorithm I to learn optimal transport maps between many 2d synthetic settings considered in the literature. In all of these settings, I instantiate ICNN and MLP architectures and use regression- based amortization with L-BFGS fine-tuning. Figures 3 to 5 show the settings considered in Makkuva al. (2020) and Rout et al. (2021), and fig. shows the conjugate objective landscapes. Figure 7 shows maps learned on synthetic settings from Huang et al. (2020). App. D contains more experimental details here.
7

---

Published as a conference paper at ICLR 2023
Samples from 3
Samples from a Push-forward 3 a
Transport paths
Figure 3: Learned transport maps synthetic settings from Rout et al. (2021).
Remark 13 Optimizing the dual in eq. (4) is an alternative to the maximum likelihood training in Huang al. (2020) for potential flows. While maximum likelihood training requires the density of one of the measures, optimizing the dual only requires samples from the measures. This makes easy to compute flows such as the bottom one fig. even though is difficult specify the density.
7 RELATED WORK
Numerical conjugation. Brenier (1989); Lucet (1996; 1997); Gardiner and Lucet (2013); Trienis (2007); Jacobs and Léger (2020); Vacher and Vialard (2021) also use computational methods for nu- merical conjugation, which also have applications in solving Hamilton-Jacobi via discretizations. These methods typically consider conjugating univariate and bivariate functions by discretizing the space, which make them challenging to apply in the settings from Korotin al. (2021a) that we report in sect. 5: we are able to conjugate dual potentials in up to 256- 36--dimeniional spaces for the HD tasks and -028-diimnsionall spaces for the CelebA64 tasks. Taking grid-based discretization of the space with 10 locations in each dimension would result in (12228) 10 grid points in the CelebA64 task. Garcia et al. (2023) amortizes the conjugate operation predict the natural gradient update, which is related to the amortized proximal optimization setting in Bae et al. (2022).
Burger's equation
Learning solutions to problems. Dinitz et al. (2021); Khodak et al. (2022); Amos et al. (2022) amortize and learn the solutions to OT and matching problems by predicting the optimal duals given input measures. These approaches are complimentary to this paper as they amortize the solation to the dual in eq. (2) while this paper amortizes the conjugate subproblem in eq. (3) that is repeatedly computed when solving a single OT problem.
CONCLUSIONS, FUTURE DIRECTIONS, AND LIMITATIONS
This paper explores the of amortization and fine-tuning for computing convex conjugates. The methodological insights and amortization perspective may directly transfer to many other applications and extensions of Euclidean Wasserstein-2 optimal transport, including for computing barycenters (Korotin et al., 2021b), Wasserstein gradient flows (Alvarez-Melis et al., 2021; Mokrov al., 2021), or cellular trajectories (Bunne al., 2021). Many of the key amortization and fine-tuning concepts from here will transfer beyond the Euclidean Wasserstein-2 setting, e.g. the more general c-transform arising non- -ucliiedd optimal transport (Sei, 2013; Cohen al., 2021; Rezende and Racanière, 2021) or for the Moreau envelope computation, which can be decomposed into a term that involve the convex conjugate as described in Rockafellar and Wets (2009, ex. 11.26) and Lucet (2006, sect. 2).
Limitations. The most significant limitation in the field of estimating Euclidean Wasserstein-2 optimal transport maps is the lack of convergence guarantees. The parameter optimization problem in eq. (4) is always non-convex, even when using input-convex neural networks. have shown that improved conjugate estimations significantly improve the stability when the base potential model is properly set up, but all methods are sensitive to the potential model's hyper-parametee parameters. found that small changes to the activation type initial learning rate cause method converge.
8

---

Published as a conference paper at ICLR 2023
f1+ #
(V (V))) #8
Potential f
Conjugate f
Input samples
Ground-truth target samples
Push-forward samples
Transport paths
Figure 4: Learned potentials on settings considered in Makkuva et al. (2020).
< 01 -
+
Figure 5: Mesh grid 9 warped by the conjugate potential flow V from the top setting of fig. 4.
y +v()
Figure 6: Sample conjugation landscapes J(;; y) of the top setting of fig. 4. The inverse transport map V, (v) #(y) is obtained by minimizing J, which is convex optimization problem. The contour shows Jax y) filtered to not display a color for values above J ( y).
8
-1-
(Vf))
O
o
Figure Single-block potential flows on synthetic settings considered in Huang al. (2020).

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
I would like to thank Max Balandat, Ricky Chen, Samuel Cohen, Marco Cuturi, Carles Domingo- Enrich, Yaron Lipman, Max Nickel, Misha Khodak, Aram- eaeeaale Pooladian, Mike Rabbat, Adriana Romero Soriano, Mark Tygert, and Lin Xiao, for insightful comments and discussions. The core set of tools in Python (Van Rossum and Drake Jr, 1995; Oliphant, 2007) enabled this work, including Hydra (Yadan, 2019), JAX (Bradbury et al., 2018), Flax (Heek et al., 2020), Matplorlib (Hunter, 2007), numpy (Oliphant, 2006; Van Der Walt al., 201 I), and pandas (McKinney, 2012).
REFERENCES
David Alvarez-Melis, Yair Schiff, and Youssef Mrouch. Optimizing fuucciooals on the space of probabilities with input convex neural networks. arXiv preprint arXiv:2166.00744, 2021.
Luigi Ambrosio. Lecture notes optimal transport problems. In Mathematical aspects of evolving interfaces, pages 1-52. Springer, 2003.
Brandon Amos. Tutorial on amortized optimization for learning to optimize over continuous domains, arXiv preprint arXiv:2202.006666 2022.
Brandon Amos, Lci Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pages 146-155. PMLR, 2017.
Brandon Amos, Samuel Cohen, Giulia Luise, and Tevgen Redko. Meta optimal transport. arXiv preprint arXiv:2266.05262 2022.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Inter- national conference on machine learning, pages 214-223. PMLR, 2017.
Larry Armijo. Minimization of functions having lipschitz continuous first partial derivatives. Pacific Journal of mathematics, 16(1);1-3, 1966.
Arip Asadulaev. Alexander Korotin, Vage Egiazarian, and Evgeny Burnaev. Neural optimal transport with general cost functionals. arXiv preprint 2iv: 15403, 2022.
Juhan Bac, Paul Vicol. Jeff Z HaoChen, and Roger Grosse. Amortized proximal optimization. arXiv preprint arXiv:2203.0008: 00089, 2022.
Jonathan Barron. Continuously differentiable exponential linear units. arXiv preprint urXiv:1704.07483, 2017.
Albert Berahas, Jorge Nocedal, and Martin Takáe. A multi-batch I-bfgs method for machine learning. Ad- vances in Neural Information Processing Systems, 29, 2016.
Dimitri P Bertsekas. Control of uncertain systems with set-membership description of the uncertainty. PhD thesis, Massachusetts Institute of Technology. 1971.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnne and Qiao Zhang. JAX: composable trans- formations of Nummy programs. GitHub, 2018. URL http: htgggbtthhughh//ggg//gggg com/ woo//ooglll google/////ooooooooooooooooo jax.
Yann Brenier. Un algorithme rapide pour le calcul de transformées de legendre-fenchel discretes. Comptes rendus de l'Académie des sciences. Série 1, Mathématique, 308(20):587-589 1989.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375-417, 1991.
Charles George Broyden. The convergence of class of double-rank minimization algorithms 1. general con- siderations. IMA Journal of Applied Mathematics, 6(1):76-90, 1970.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong- Van Lehmann, Lucas Pelk- mans, Andreas Krause, and Gunnar Ratsch. Learning single-cell perturbation responses using neural optimal transport. bioRxiv, 2021.
Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised training of conditional monge maps. Ad- vaRces Neural Information Processing Systems, 2022.
10

---

Published as a conference paper at ICLR 2023
Jiezhang Cao, Langyuan Mo. Yifan Zhang, Kui Jia, Chunhua Shen, and Mingkui Tan. Multi-marginal wasser- stein gan. Advances in Neural Information Processing Systems, 32, 2019.
Tianlong Chen, Xianhan Chen, Wuyang Chen, Howard Heaton, Jialin Liu, 7hangyang Wang, and Wotao Yin. .earming to optimize: primer and benchmark. arXiv preprint arXiv:2 12828 2021.
Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint 15 2015.
Samuel Cohen, Brandon Amos, and Yaron Lipman. Riemannian convex potential maps. In International Conference on Machine Learning, pages 2028-2038. PMLR, 2021.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural informa- rion processing systems, 26:2292-2300, 2013.
John M Danskin. The theory of max-min, with applications. SIAM Journal on Applied Mathematics, 14(4): 641-664, 1966.
Yann N Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio. Identifying and attacking the saddle point problem in high-dimensional non-convex optimization. Advances in neural information processing systems, 27, 2014.
Michael Dinitz, Sungjin Im. Thomas Lavastida, Benjamin Moseley, and Sergei Vassilvitskii. Faster matchings via learned duals. Advances in neural information processing systems, 54:10393-10406, 2021.
Jiaojiao Fan, Amirhossein Taghvaci, and Yongxin Chen. Scalable computations of wasserstein barycenter via input convex neural networks. arXiv preprint 2xXi:2007.04446) 2020.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Yongxin Chen, and Haomin Zhou. Scalable computation of monge maps with general costs. arXiv preprint wwxiv:2206.0388 2021.
Roger Fletcher. A new approach to variable metric algorithms. The computer journaaa 13(3):317-322, 1970.
Sara Fridovich Meciiichh-KKklaae and Benjamin Recht. Approximately exact line search. arXiv preprint arXiv:201 2020.
Jezabel Garcia, Federica Freddi, Stathi Fotiadis, Maolin Li, Sattar Vakili, Alberto Bernacchia, and Guillaume Hennequin. Fisher-logendre (fishleg) optimization of deep neural networks. International Conference on Learning Representations, 2023. URL https:/ http:::://///////////eeee oppprevieeeeeeeeeeeeeeeeee net. forum?
Bryan Gardiner and Yves Lucet. Computing the conjugate of convex piecewise quaar-quadratic bivariate functions. Mathematical Programming, 139(1):161-18))) 2013.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems, 29, 2016.
Donald Goldfarb. A family of variable-metric methods derived by variational means. Mathematics of compu- tation, 24(109):23- 26, 1970.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pages 630-645. Springer, 2016.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee Flax: A neural network library and ecosystem for JAX. GitHub, 2020. URL http:
//github. wtt.c/////oooo gaonn//gool////oooooooooo flax.
Chin-Wat Huang, Ricky TQ Chen, Christos Tsirigotis, and Aaron Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. arXiv preprint arXiv:2012.05942,, 2020.
John D Hunter. Matplotlib: A 2d graphics environment. Computing in science & engineering, 9(3)990, 2007.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pages 448-456 456. PMLR, 2015.
Matt Jacobs and Flavien Léger. A fast approach to optimal transport: The back-and-forth method. Numerische Mathematik, 146(3):513- 544, 2020.
Sham M Kakade, Shai Shalev-Shwariz, and Ambuj Tewari. On the duality of strong convexity and strong smoothness: Learning applications and matrix regularization. TTIC, 2009.
11

---

Published as a conference paper at ICLR 2023
Mikhail Khodak, Maria-Florina Balcan, Ameet Talwalkar, and Sergei Vassilvitskii.. Learning predictions for algorithms with predictions. arXiv preprint arXiv:200.009312, 2022.
Diederik P Kingma and Jimmy Ba. (rXiv:/412.666 2014.
Adam: A method for stochastic optimization. arXiv preprint
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions, Advances in neural information processing systems, 31, 2018.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. arXiv preprint arXiv:1909. 13082, 2019.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Alexander Filippov, and Evgeny Bur- nacv. Do neural optimal transport solvers work? a continuous wasserstein-2 benehmark. arXiv preprint arXiv:2106.0I4 009444 2021a.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev, Continuous wasserstein-2 barycenter estimation without minimax optimization. arXiv preprint arXiv:2102.01755 2021b.
Jacob Leygonie, Jennifer She, Amjad Almahairi, Sai Rajeswar, and Aaron Courville. Adversarial computation of optimal transport maps. arXiv preprint 1966911906.0991,, 2019.
Dong-Hui Hui Li and Masao Fukushima. On the global convergence of the bfgs method for nonconvex uncon- strained optimization problems. SIAM Journal on Optimization, 11(4): 1054,,0040 2001.
Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin Solomon. Continuous regularized wasserstein barycenters. arXiv preprint arXiv:2008. 12534, 2020.
Dong Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, 45(1)=503-528, 1989.
Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein gan with quadratic transport cost. In Proceed- ings of the IEEE/CVF international conference computer vision, pages 4832 12244841, 2019.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceed- ings of the IEEE international conference on computer vision, pages 3730-3738, 2015.
Frederike Lübeck, Charlotte Bunne, Gabricle Gut, Jacobo Sarabia del Castillo, Lucas Pelkmans, and David Alvarez-Melis. Neural unbalanced optimal transport via cycle-consistent semi-couplings. arXiv preprint arXiv:2209.1500 2022.
Yves Lucet. A fast computational algorithm for the legendre-fenchel transform. Computational Optimization and Applications, (():27757, 1996.
Yves Lucet. Faster than the fast legendre transform, the thear-time legendre transform. Numerical Algorithmms 16(2):171-185, 1997.
Yves Lucet. Fast moreau envelope computation Numerical algorithms. Numerical Algorithms, 43(3):235- 249, 2006.
Yun Luo, Yi-Yang Zhang, Wei-L.ong Zheng, and Bao-Liang Lu. Wgan domain adaptation for eeg-hased emotion recognition. In International Conference on Neural Information Processing, pages 275-286. Springer, 2018.
Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. International Conference on Machine Learning, pages 6672 6681. PMLR, 2020.
Robert J McCann. Existence and uniqueness of monotone measure-preserving maps. Journal, 80(2):309-33) -23, 1995.
Duke Mathematical
Wes McKinney. Python for data analysis: Data wrangling with Pandas, NumPy, and IPython. O'Reilly Media, Inc.", 2012.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev. Large- scale wasserstein gradient flows. Advances Neural Information Processing Systems, 34:15243 -52566, 2021.
Quan Hoang Nhan Dam, Trung Le, Tu Dinh Nguyen, Hung Bui, and Dinh Phung. Threoplayer Wasserstein GAN via amortised duality. In Proc. the the 28th Ins. Joint Conf. on Artificial Intelligence (IJCAI), 2019.
Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 1999.
12

---

Published as a conference paper at ICLR 2023
Travis E Oliphant. guide NumPy, volume 1. Trelgol Publishing USA, 2006.
Travis E Oliphant. Python for scientific computing. Computing in Science & Engineering, 9(3):10-20, 2007.
Henning Petzka, Asja Fischer, and Denis Lukovnicov. On the regularization wasserstein gans, arXiv preprint (:Xiv:/709.088800 2017.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foun- dations and Trends in Machine Learning, 11(5-6):355 607, 2019.
Danilo J Rezende and Séhastien Racanière. Implicit riemannian concave potential maps. arXiv preprint 0rXiv:2/10.00 10.00288, 2021.
R Tyrelll Rockafellar and Roger J-B Wets. Variational analysis, volume 317. Springer Science & Business Media, 2009.
Ralph Tyrell Rockafellar. Convex analysis. In Conver analysis. Princeton university press, 2015.
Olaf Ronneherger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and compider-assister intervention, pages 234-241. Springer, 2015.
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. arXiv preprint arXiv:21 10.02999, 2021.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blon- del. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv: /71 1.02283, 2017.
Tomonari Sei. jacobian inequality for gradient maps the sphere and its application to directional statistics. Communications in Statistics- Theory and Methods, 42(14):52552542 2013.
David F Shanno. Conditioning of quasi-newton methods for function minimization. Mathematics of computa- tion, 1):141):647-- -566, 1970.
Jian Shen, Yanru Qu, Weinan Zhang, and Yong Yu. Wasserstein distance guided representation learning for domain adaptation. In Thirty-second AAAI conference artificial intelligence, 2018
Jack Sherman and Winifred J Morrison. Adjustment of an inverse matrix corresponding change in one element of given matrix. The Annals of Mathematical Statistics, 21(1):124 127, 1950.
Rui
Shu.
Amortized Optimization.
iatieiiooniiiiiiiino amort wedd-ppptimmee Lon 2017. Accessed: 2020-02-02.
Amirhossein Taghvaei and Amin Jalali. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv preprint 1902009902.07197,, 2019.
Michael Joseph Trienis. Computational convex analysis: From continuous deformation to finite corvex inte- gration. PhD thesis, University of British Columbia, 2007.
Adrien Vacher and François-Xavier Vialard. Convex transport potential selection with semi-duul criterion. arXiv preprint 17Xiv:2112.07255, 2011.
Stefan Van Der Walt. S Chris Colbert, and Gael Varoquaux. The numpy arry: a structure efficient numerical computation. Computing Science & Engineering, 13(2):22,
Guido Van Rossum and Fred I. Drake Jr. Python reference manual. Centrum Wiskunde Informatica Amsterdam, 1995.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Zhong Wan, Shuai Huang, and Xiao Dong Zheng. New cautious bfgs algorithm based on modified armijo-type line search. Journal of Inequalities and Applications, 2012(1):1-10, 2012.
Philip Wolfe. Convergence conditions for ascent methods. SIAM review, 11(2):26-235, 1969.
Max Woodbury. Inverting modified matrices. Statistical Rescarch Group, 1950.
Jiqing Wu, Zhiwu Huang, Janine Thoma, Dinesh Acharya, and Lue Van Gool. Wasserstein divergence for gans. In Proceedings of the European Conference Computer Vision (ECCV), pages 653-668, 2018.
13

---

Published as a conference paper at ICLR 2023
Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computa- tion of large scale optimal transport. International Conference on Machine Learning, pages 6882-6892. PMLR, 2019.
Onry Yadan. Hydra framework for elegandy configuring complex applications. Github, 2019. URL
https: / thub. com/ waacebookeeeerrcchyyyy.....
Jun Yan Zhu, Taesung Park, Phillip Isola, Alexei Efros. Unpaired image to-too--onn image translation using cycle-consistent adversarial networks, In Proceedings of the IEEE international conference on computer vision, pages 2223-2232 2017.
14

---

Published as a conference paper at ICLR 2023
Algorithm 3 The Setarn-lllecchr-Golftttte Nocedal and Wright (1999, alg. 6.1).
(BFGS) method to solve eq. (11) as presented
Inputs: Function J to optimize, initial iterate zy and Hessian approximation Bo
k 4-0 0
while unconverged do
Compute the search direction p4 = - JJ
Set Tk+1 = T3 + O&&P where og; is computed from a line search from app. A.2 Compute B₁ with the update in eq. (12)
+ k
end while
return optimal solution T3 R ă
Backtracking Armijo search
Parallel Armijo search
g(a)
Sequentially find g(a) 2
Evaluate many g(a) in parallel, pick best
Figure 8: Visualization of backtracking and parallel line searches to solve eq. (17).
A L-BFGS OVERVIEW AND LINE SEARCH DETAILS
The conjugate optimization problem in eq. (3) is an unconstrained convex optimization problem for convex potentials, which is setting BFGS (Broyden, 1970; Fletcher, 1970; Goldfarb, 1970; Shanno, 1970) and L-BFGS (Liu and Nocedal, 1989) methods thrive The default strong Wolfe line search methods in the Jax and JaxOpt L-BFGS implementations may take long time to solve a batch of optimization problems. Without efficiently setting the line search method, some of the Wasserstein-2 benchmark experiments in app. C that ran in few hours would have otherwise taken a month to run. This section provides a brief overview of BFGS methods and shows that an Armijo line search can be the most efficient computing the conjugate.
A.1 BACKGROUND BEGS METHODS
Nocedal and Wright (1999, alg. 6.1) is standard reference for BFGS methods and extensions and the key steps are summarized algorithm for solving an optimization problem of the form
# arg min (xx
(11)
where J : IR" - R is possibly non-convex and twice continuously differentiable. The method iteratively finds a solution & by 1) maintaining an approximate Hessian around the current iterate, L.e. B, & {zJ(xx), 2) computing approximate Hessian, 3) updating the iterate with Ik+1 = Iy + C&Pk. where ag is found with a line search, and 4) updating the Hessian approximation with the Mormann-Morrion -ooobbron-Woduurry formula (Woodbury, 1950; Sherman and Morrison, 1950)
approximate Newton step Pk - B , (xx) using the
B++1 = B BxS S B
ykUk
(12)
8 BYSK
Dg 8k
where Vk (I (zJ( and 8 24+1 TA. Instead of estimating Bş and inverting it in every iteration, most implementations maintain a direct approximation to the inverse Hessian H₂ B 1 The limited-memory version of BFGS (L-BFGS) in Liu and Nocedal (1989) propose to replace the inverse Hessian approximation as matrix Hę with the sequence of vectors [yk, defining the updates Hy and never requires instantiating the full approximation.
15

---

Published as a conference paper at ICLR 2023
Algorithm 4 Backtracking Armijo line search to solve eq. (17)
Inputs: Iterate T2, search direction Pk. decay control parameter C₁, initial
< Cinit
while J [zx ayPk)) J(m)) + caaypk V (1r) do
a Ta
end while
return a satisfying the Armijo condition in eq. (17).
Algorithm 5 Parallel Armijo line search to solve eq. (17)
Inputs:
Iterate The search direction pk. decay т, control parameter c₁, initial oo, #evaluations M
Compute candidate step lengths Om T for E [M]
Evaluate the line search condition g(mm) from eq. (16) in parallel if all g(am) 0 then
Error: No acceptable step found
else
return largest satisfying eq. (17), i.e. max om subject to g(am) 0
end
A.2 LINE SEARCHES
The line search find the step size og for the iterate update Tk+11 xx + O&PE is often done with:
1. a Wolfe line search (Wolfe, 1969) to satisfy the conditions:
Jaz a&PE) s ((xx) + crap v.J((a))
Pr V J(n a.px) S c2p2 T (z)(
(13)
where < ci < 02 < 1,
2. a strong Wolfe line search to satisfy the conditions:
J(r + оaPP) S J(xx) + croupk (xx(
(14)
|P I + a.P=)) S c2 P2 [xx()
This is often found via the zoom procedure from Nocedal and Wright (1999, algorithm 3.5). 3. an Armijo line search (Armijo, 1966) to satisfy the first condition:
J(r + (gPp)) S (1x)) + c10.P (zx))
(15)
For notational simplicity, we can also write the Armijo condition as:
g(a) == f(xx) + cropi (xxx)) - f(x + opp) 20
(16)
Remark 14 The strong Wolfe line search is the most commonly used line search for L-BFGS as guarantees that the resulting update to the Hessian eq. (12) stays positive definite, but the Armijo line search may be more efficient as it does not involve re-evaluating the derivative of the objective. Unfortunately. an iterate obtained by an Armijo line search may not satisfy the curvature condition y% 8% > 0 that ensure the Hessiun update stays positive definite while a step satisfying the strong Wolfe conditions provably does (Nocedal and Wright, 1999, page 143). Nonetheless, Armijo searches are still combined with BFGS and can be guarded by only updating the Hessian approximation if y Sy > O, or a modification thereof, as described Li and Fukushima (2001); Wan et al. (2012); Fridovich Keil and Recht (2020) and Berahas al. (2016, sect. 3.2).
16

---

Published as a conference paper at ICLR 2023
Table Runtime and number of L-BFGS iterations for line search methods converge to solution I of eq. (11) for conjugating the trained ICNN potential on the 256-dimensional HD benchmark from sect. 5 with batch of 1024 samples and a tolerance of [ I 0.1, starting from the amortized prediction. The Wolfe and Armijo line search methods use standard values of c₁ 10" and C2 = 0.9, all backtracking options use a decay factor of = 2/3 with M 15 evaluations, and the runtimes are averaged over 10 trials on an NVIDIA Tesla V100 GPU. In this setting, Armijo line searches without many conditionals or gradient evaluations consistently take the shortest time.
Base L-BFGS
Line search
Runtime (ms)
# Iterations
Jax
Strong Wolfe Zoom (default)
4803.40
6.21
Backtracking Armijo
156.69
8.56
Parallel Armijo
119.41
8.56
JaxOpt
Strong Wolfe Zoom (default)
776.52
7.78
Backtracking Strong Wolfe
233.39
7.97
Backtracking Wolfe
225.30
8.65
Backtracking Armijo
154.90
8.58
A.2.1 THE ARMIJO LINE SEARCH
The Armijo line search can be written as the optimization problem
ar (x Ph) = max subject to J( + apx) S J(xx) + crapp V (
(17)
Eq. (17) is typically solved as shown in algorithm 4 and fig. 8 by setting a decay factor T and iteratively decreasing candidate step length until the condition is satisfied.
When the objective J can be efficiently evaluated parallel on a GPU, and when solving many batches of optimization problems concurrently, e.g. with vmap, the backtracking Armijo line search described in algorithm 4, and the Wolfe line search described in Nocedal and Wright (1999, alg. 7.5), are computationally slowed down by serial and conditional operations. These issues arise from: 1) the sequential nature of the line search, and 2) the fact that the line search may run for different number of iterations for every optimization problem in the batch. Wolfe line searches such as No- cedal and Wright (1999, alg. 7.5) have other conditionals scoping the search interval that cause the line search to perform potentially different operations for every optimization problem the batch.
I propose parallel Armijo line search in algorithm 5, which is also visualized in fig. 8, to remove serial and conditional operations to improve the computation of the line search on the GPU for solving batches of optimization problems. The key idea is to instantiate many possible step sizes, evaluate them all at once, and then select the largest a'm satisfying the Armijo condition g(am) 2 O.
The parallel line search may unnecessarity evaluate more candidate steps sizes than the sequential line search, but on GPU architectures this may not be very detrimental the performance because additional parallel function evaluations are computationally cheap. Furthermore, when solving batch of N optimization problems with M line search evaluations, Lee. when using vmap on the line search or optimizer, the parallel line search in algorithm 5 can efficiently evaluate N M candidate step lengths in tandem on GPU and then select the best for each element in the batch.
Remark 15
Remark 16 A potential concern with this parallel line search is that may not find a step size satisfying the Armijo condition if M is not set to be low enough. While this may be significant issue for when high-precision solves are needed, have found in practice for the Euclidean Wasserstein-2 conjugates that taking M 10 line search evaluations frequently finds solution.
A.2.2 COMPARING LINE SEARCH METHODS
Table 3 takes trained ICNN potential and isolates the comparison between L-BFGS runtimes to only changing the linesearch methods. This the same optimization procedure and batch size used for all of the training runs on the Wasserstein-2 benchmark. Despite the concerns in remark 14 about the Armijo line search resulting in slower convergence and an indefinite Hessian approximations, the Armijo line searches are consistently able to solve the batch of optimization problems the fastest.
17

---

Published as a conference paper at ICLR 2023
B MODEL DEFINITIONS AND PRETRAINING
All of the potential and conjugate amortization models in this paper can be implemented 30055 lines of readable Jax code with Flax (Heek et al., 2020). They are included in the w2ot /models directory of the code here, and reproduced here to precisely define them.
B.1 PRETRAINING AND INITIALIZATION
Following Korotin et al. (2021a), every experimental setting has a pre-training phase so that the po- tentials and amortization maps approximate the identity mapping. i.e. V v=fo(x) &e and (v) 2 y.
B.2 InitNN: NON-CONVEX NEURAL NETWORK AMORTIZATION MODEL Sipp
Remark 17 The passthrough on line 18 is helpful for learning an identity initialization.
class InitNN (nn. Module) dim_hidden: Sequence [int] act: str 'elu'
conpact
def _call_(sse
assert ndid n_nnput x. shape[[[1] [-1]
act_fn fn layers. (self.
12
z x
13
for n hidden in self.din hidden:
14
Wx nn. Dense (n_hidden, use use_bias-Truel fct_fn (xx(2))
16
17
Rx nn. Dense (n input, use_bies-True)
18
z = x
Wx (z) Encourage identity initialization.
return
18

---

Published as a conference paper at ICLR 2023
B.3 ICNN: INPUT-CONVEX NEURAL NETWORK POTENTIAL fo
actnorm is the activation normalization layer from Kingma and Dhariwal (2018), which was also used in the ICNN potentials in Huang et al. (2020) and normalizes the activations at initialization follow a normal distribution.
Remark 18 Applying an activation the ouppu on line 41 is helpful to lower-bound the otherwise unconstrained potential and adds stability to the training.
Remark 19 The final quadratic on line 46 makes it easy to initialize the potential to the identity.
Remark 20 This ICNN does not the quadratic activations proposed in Korotin et al. (2019, Appendix B.I). While did not heavily experiment with them, table shows that this ICNN architec- ture without the quadratic activations performs better than the results originally reported in Korotin et al. (2021a) which use an ICNN architecture with the quadratic activations.
class ICNN (nn. Module):
dim_hidden: Sequence [int]
act: str 'elu'
actnorm: bool True
def petup(self):
kernel_init nn. on.intiaalizers maattaallarss.vrranne.caallll
1., "fan_in", "truncated_normall"" num hum_hidden
w_ list()
12
for range (1, num sam_hideen):
13
w w_ append (ayerssdaayrr..Psiiiveeeeee
self dim hidden|il, kernel wnit=kernee (nitt)
28.8.00 append (layers. www.s..ossiivvee Positiee (1, self. w_ w_28
16
17
18
w_ = list
for i in range (num_hiddes) hidddmiddd
w_xs. append Densee)
self.dim_hiddenlll) use see_bias=True,
22
23
x_x8_appee append Inn.Dense (1, use_bias-True, self self.w_xs w_xs
)
26
27
8nn. compact
def call (self,
30
assert x. ndii
31
n iinput x. [shape[-1]
32
act_fn fn layers act (self.act)
act En (self.
)x))
35
for Wz, wx in zip(self. [:-]], self.x_xs[1:-1]
36
z (z (x)
37
if Layers.
39
z act_
40
41
act_ (self.nnss ____s[=1] (z) + self.w_xst-1] (x))
y jqp.equeeze (y, -1)
44
log alpha self.param
45
'log_ Llog_alpha', , nn www.iiiiiiiiis....mm constant (0), (1)
46
+= jjppeepp alohaalpha +0, 5+ut botthls.aatcccooo (x, x)
48
return
19

---

Published as a conference paper at ICLR 2023
B.4 PotentialNN. NON-CONVEX NEURAL NETWORK (MLP) POTENTIAL fo
Remark 21 Consistent with remarks 18 and 19, applying an activation to the last layer (line 18) and adding final quadratic term (line 24) helps this non-convex potential model too.
class PNN (oduued
dim_hidden: Sequence [int] act: str = 'elu'
@nn. enn.compact
def _all__ (self, x)
assert x. ndni = 2
n_nnput shappel-11
10
act_ fn layers.get_ace (self.act)
11
12
x
for h__iidden in welf.dim_hidee
14
Wx Dense (n hidden, use_bias=True)
15
z act_ (Wx
16
wx nense (1, use_bias-Irue) act fn (xx(2))
19
20
z = jnp. squeezezz, -1)
log_alpha self.
23
'log_alpha" nn iniitaaliss constant (0), (1)
24
2 += 0.5+ jnp. exp (log_alpha) wtt.ilssaatccootttoooooooooo (x, x)
25
return
20

---

Published as a conference paper at ICLR 2023
B.5 ConvPotential: NON-CONVEX CONVOLUTIONAL POTENTIAL fo
Remark 22 I was not able to easily add batch normalization (loffe and Szegedy, 2015) to this potential. In contrast to standard use cases of batch normalization that only call into batch- normalized model once over samples from single distribution, the dual objective in eq. (4) calls into the potential multiple times estimate Exva fe(z) and Eg fo(ā))), which also involve internally solving the conjugate optimization problem in eq. (3) to obtain ă. This makes it not clear what training and evaluation statistics batch normalization should use when computing the dual objective. One choice could be to only use the statistics induced from the samples æ ~
class ConvPotential Module) act: str 'elu'
mean = jnp. array 410.485, 0.456, 0, 406]] stá árray [0.229, 0 0.224, 0.225]
@nn. compact def
(self, x)
assert x.ndim - 2 Images should be flattened num num_batch shapeel
11
12
x flat x Save for taking the quadratic at the end.
13
Reshape and renormalize
15
X x.reshape(-1, 3, 64, 64) transpooe (0, 2, 3, 1)
16
x (x 1.)/2.
17
x = (x-self.mean) self.std
y
20
act_fn layers.get_ace (self.act)
21
22
conv - nn Conv (128, kernel_size- 4], strides-2)
y act_ fn (conv (
25
cony nn. Conv (128, ternel_sizer 14, 4), strides=2)
26
y = act f (conv (
conv Conv (256, y fc (conv (y))
kernel_size-(4,A4 strides-2)
29
30
31
conv = nn. Conv (512, y act_fn (conv (
kernel_size-1 [4, 4], strides-2)
34
conv nn. Conv (1024, kernel_size= 14,4), strides=2)
35
y = act fn (conv )
36
conv nn. Conv
1, kernel ze- 1izee[22) [2, sadding=/VALID> strides=1)
39
y act_fn (conv (y))
40
y-squeeze ( [112,31)
assert y. shape (num mutchb
45
46
log_alpha self colfparamt
45
'log_alpha', nn nn. im.iniiiaaizerss constant (0), (])
y 0.5+ jnp..inp.co jnp exp alpha) sutils. batch (x flat, f_faat)
48
return y
21

---

Published as a conference paper at ICLR 2023
C ADDITIONAL WASSERSTEIN-2 BENCHMARK EXPERIMENT DETAILS
C.I HYPER- PARAMETRS
Tables and 5 detail the main hyper-parameteee for the Wasserstein-2 benchmark experiments. I tried to keep these consistent with the choices from Korotin et al. (2021a), e.g. using the same batch sizes, number of training iterations, and hidden layer sizes for the potential.
All experiments use the same settings for the conjugate solvers: The conjugate solvers stop early all dimensions the iterates change by less than 0.1, and otherwise run for a maximum 100 iterations. The line search parameters for the parailel Armijo search in algorithm 5 for L-BFGS are to decay the steps with a base of = 1.5 and to search M 10 step sizes. With the Adam conjugate solver, use the default = [0.9, 0.999] with initial learning rate of 0.1 with cosine annealing schedule to decrease 10
Table 4: Hyper-parameters for the D-dimensional Wasserstein-2 benchmark experiments
Name
Value
potential model fo fo hidden layer sizes conjugate amortization model ž activation functions
ICNN PotentialMo
[max(21) 64), max(21, 64), max(D,aa 32)] InitNN (dim_hidden- [512, 512]) ELU (Clevert et al., 2015)
training iterations
250000
optimizer
Adam with cosine annealing (a=le-4)
initial learning rate
5e-4
Adam 3
[0.5, 0.5]
batch size 1024
Table Hyper-paramettrs for the CelebA64 Wasserstein-2 benchmark experiments
Name
Value
potential model fo ConvPotential
conjugate amortization model क..- V activation functions
Gradient of ConvPotential
ELU (Clevert et al., 2015) 50000
number of training iterations
optimizer
Adam with cosine annealing (a=le-4)
initial Icarning rate
le-3
Adam
[0.5, 0.5]
batch size
64
22

---

Published as a conference paper at ICLR 2023
C.2 CONVERGENCE OF L-BFGS AND ADAM FOR SOLVING THE CONJUGATE
Figure 9 shows that with a non-convex potential, many of the initial amortized predictions are sub- optimal and difficult for the L--BGS to improve upon. This indicates that the amortized predictions may be in parts of the space that are difficult recover from and suggests future avenue of work better characterizing and recovering from this behavior. L-BFGS converges fast to an optimal solu- tion in fig. 10 while Adam often gets stuck suboptimal solutions.
64
D 128
256
1.00
Norndled J(a
0.75
0:25
0.00
10
20
40
30
0
10
30
40
30.0
20
30
Salver Iseration
Solver Iseraaton
Salver Iteration
L-BFGS
Amortization L-BFGS
Adam
Figure 9: Conjugate solver convergence on the HD benchmarks with a NN potential.
Early
Mid
Late
1.00
Nornslize J(vgg))
0.75
0.50
0.25
2
10.0
10.0
Solver Iseration
Solver Iseration
Solver Ierration
L-BFGS
Amortization L-BFGS
Adam
Figure 10: Conjugate solver convergence on the CelebA64 benchmark
C.3 ADDITIONAL RUNTIMES AND CONJUGATION INFORMATION
Tables and 7 contain additional experimental information with the:
1. wall-clock time for the entire training run measured on an NVIDIA Tesla V100 GPU,
2. number of conjugation iterations from algorithm 2 for the conjugate solver to converge at the end of training after warm-starting from the conjugate amortization model's pre- diction, and
3. runtime for the conjugate solver to converge on a batch of instances, set to the batch size used during training, i.e. 1024 for the HD benchmark and 64 for the CelebA64 benchmark.
These results give an idea of how much additional time is spent fine-tuning. On the HD benchmark, fine-tuning takes between xe 10-50ms per batch. The overall wall clock time may take 2-3 times longer than the training runs without fine-tuning, but are able to find significantly better solutions. On the CelebA64 benchmarks, the conjugation time impacts the overall runtime even less hecause, especially in the "Mid" and "Late" settings the transport maps here close to being the identity mapping and are easy to conjugate.
Remark 23 Some settings immediately diverged to an irrecoverable state providing a C2-UPP of 10P, including runs using the objective-based and cycle amortization losses without fine-tuning. I early-stopped those experiments and do not report the runtimes or conjugation times here, as the few minutes that the objective-base amortization experiments took to diverge is not very interesting or comparable the times of the experiments that converged.
23

---

Published as a conference paper at ICLR 2023
Table 6: Additional runtime and conjugation information for the HD benchmark. These report the median time from the converged runs.
Runtinc (bours)
Final conjugation itce
Conj runtime (scconds)
Amertization loss
Conjugate sulver
64
128
256
64
128
256
128
256
ICNN
Cycle
None
1.10
1.24
1.58
IONN
Cycle
L-BRGS
4.39
7.76
21.14
14.32
19.38
70.23
0.05
0.10
0.29
ICNN
Objective
L-BFGS
4.70
0.05
ICNN
Regression
L-DFGS
2.74
4.30
12.53
5.64
6.34
8.73
0.02
0.04
0.11
ICNN
Cycle
Adam
2.09
2.80
0.86
37.23
51.34
97.21
0.02
0.03
0.06
IONN
Objective
Adam
2,23
2.99
4.94
29.44
55.05
0.02
0.03
IONN
Regression
Adam
2.09
2,86
4,81
29.51
38.20
54.21
0.02
0.03
0.06
NN
Objective
L-BEGS
2.31
3.34
6.05
5.05
6.50
5.77
0.02
0.03
0.05
NN
Regression
L.BGGS
2.28
3.21
5.76
5.21
5.01
0.02
0.02
NN
Objecrive
Adam
1,61
2.18
3,70
27.84
36.81
45.82
0.01
0.02
0.04
NN
Regression
Adam
2.51
3.79
31.46
40.31
0.01
Table 7: Additional runtime and conjugation information for the CelebA64 benchmark
Runtime (hours)
canj Mid
Conj Early
runtime (seconds)
Amortization loss
Conjugate solver
Model
Direction
Early
Mid
Late
Early
Late
Mid
Late
Objective
None
Conv
Forward
3.28
3.41
3,28
Cycle
Conv
Forward
0.94
4.17
4.17
Cycle
Adam
Conv
Forward
5.44
4.78
3.57
18.23
2,02
2,80
0.15
0.11
0.04
Cycle
L-BFGS
Conv
Forward
6.38
3.79
3.74
5.18
2,00
2.00
0.22
0.04
0.04
Objective
Adam
Conv
Forward
5.40
4.85
3.47
19.87
1.83
1.79
0.17
0.13
0.03
Objective
L-BFGS
Conv
Forward
6.09
3.86
3.70
4.48
2.01
2.00
0.21
0.06
0.04
Regression
Adam L-BFGS
Cony Cony
Forward
5.44
4.86
3.51
22.12
2.84
1.01
0.19
0.14
0.02
Regression
Forward
5.96
3.77
3.66
4.55
2.01
2.02
0.22
0.04
0.04
D ADDITIONAL 2D DEMONSTRATION EXPERIMENT DETAILS
Table details the main hyper- syper-parameters for the synthetic benchmark experiments, and fig. 11
shows additional conjugation landscapes.
Remark 24 I found leaky ReLU activations on the potential model to work better in these low- dimensional settings than ELU activations, which work better in the HD benchmark settings. I do not have a strong explanation for this but found the LReLU capable of performing shurper transitions in the space, e.g. the sharp boundaries shown in fig. 4. One reason that the ELU potentials could perform better on the benchmark settings is that the ground-trute truth transport maps in the benchmark, described in Korotin al. (202 la, Appendix B) use an ICNN with CELU activations (Barron, 2017) which may be easier to recover with potential models that use ELU activations.
I trained convex and non-convex potentials on every synthetic setting and show the results from the best- rest-performing potential model, which are:
Makkuva et al. (2020): an ICNN. This setting originally considered convex potentials, and the non-convex potentials tried training on these settings diverged,
Rout et al. (2021): a non-convex potential (an MLP). This setting also originally considered MLP and couldn't 't find an ICNN that accurately transports between the highly curved and concentrated parts of the measures.
Huang et al. (2020): non-convex potential (an MLP). In contrast to the ICNNs origi- nally used, I found that an MLP works better when learned with the OT dual. Almost every setting in Huang et al. (2020) requires composing multiple blocks ICNNs, which means the flow will not necessary be the optimal transport flow, while the non-convex MLP potential am using here estimates the optimal transport map between the measures.
All of the synthetic settings use the L-BFGS conjugate solver set to obtain slightly higher precision solves than in the Wasserstein-2 benchmark. The conjugate solver stops early if all dimensions of the iterates change by less than 0.001, and otherwise run for maximum of 100 iterations. The line search parameters for the parallel Armijo search in algorithm 5 for L-BFGS are to decay the steps with base of 7 1.5 and to search M 30 step sizes.
24

---

Published as a conference paper at ICLR 2023
y *z(
Figure 11: Sample conjugation landscapes ((;;)) of the bottom setting of fig. 4. The inverse trans- port map (y) x(y) obtained by minimizing J, which is a convex optimization problem. The contour shows [z; y) filtered to not display a color for values above J y).
Table 8: Hyper-parameters for the synthetic experiments
Name
Value
potential model fe fo hidden layer sizes conjugate amortization model P
ICNN or PotentialNM [128, 128]
InitNN (dim_hiddene) www.hhiddenniic (512, 512]) Leaky ReLU with slope 0.2
activation functions
training iterations
50000
optimizer
Adam with cosine annealing (a=le-4)
initial learning rate Adam 33
5e-4
[0.5, 0.5]
batch size
10000
D.1 NON-CONVEX REGIONS THE LEARNED POTENTIALS
Brenier's theorem (Brenier, 1991) shows that the known Wasserstein-2 optimal transport map asso- ciated with the negative inner product cost is the gradient of convex function, i.e. T(x) (=f( Because of this, optimizing over convex potentials is theoretically nice and also results in a convex and easy conjugate optimization problem in eq. (3) to compute f*. The input-convex property is usually enforced by constraining all of the weights of the network to be positive in every layer ex- cept the first. Unfortunately, in practice, the positivity constraints of a convex potential may be prohibitive and not easy to optimize over and result in sub-optimal transport maps. In other words, the parameter optimization problem over the input-convex model is still non-convex and may be exasperated by the input-convex constraints. Due to these limitations, non-convex potentials are ap- pealing as their parameter space less constrained and may therefore be easier search over. And in practice, this has been shown to be true, e.g. the main results in table I show that non-convex potential significantly outperforms the convex potential. However, non-convex potentials can re- sult in non-convex conjugate optimization problems in eq. (3) that can cause significant numerical instabilities and an exploding upper-bound on the dual objective.
Figure 12 illustrates small non-convex region arising in learned non-convex potential. While the non-convex region mostly does not impact the transport map in this case, they can easily blow up and make the dual optimization problem challenging. In contrast, the ICNN-based convex potential provably retains convexity and keeps this region nicely flat, but the constraints on the parameter space may hinder the performance.
25

---

Published as a conference paper at ICLR 2023
Interpolation from a non-convex potential (an MLP)
((1- f) + +
(V/*)²
Interpolation from a convex potential (an ICNN)
( ((1-t) -)- -
(Vf*)
Non-convex potential contours
Convex potential contours
Figure 12: Convex and non-convex potentials trained on the same transport task.
D.2 INTERPOLATIONS ON SYNTIETIC SETTINGS FROM ROUT ET AL. (2021)
26


Reviews:
[{"review_id": "xJG5zwOf2o", "reviewer": "Reviewer_eQGj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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 results in this paper apply to the computation of optimal transport maps for the Wasserstein-2 distance in Euclidean space. Under weak assumptions, these maps are given by gradients of convex functions. Moreover, the convex function and its Fenchel-Légendre conjungation are a solution to the dual problem. The paper presents different frameworks for performing the conjugation via a differentiable optimization problem. Moreover, the author also provides an efficient and paralelizable \"Armijo-type\" algorithm for performing the conjugation. Experiments suggest that this approach leads to significant improvements in the computation of optimal transport maps. \n\n ", "summary_of_the_review": "The paper presents novel methods for the conjugation step in dual Wasserstein-2 optimal transport. Experiments suggest that the methods are good, but the paper is a bit hard to read, and leaves some questions regarding tuning unresolved. ", "strength_and_weaknesses": "*Strenghts*\n\n- The experimental results look promising. \n- The paralellized Armijo is a cool idea. \n\n*Weaknesses*\n\n- I found the writing a bit difficult for a non-expert (eg. terms like \"amortization\" are never defined, and is only really explained in the code in the appendix).  \n- There's no theory to back up the methods.  \n- Some parts of notation are not well defined (eg. the letter for the dimension changes from n to D in the text).\n- It is not clear to me how the method should be tuned in the abscence of a ground truth. \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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "t3JjIJA4dgT", "reviewer": "Reviewer_7Bdg", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Monge's primal Wasserstein-$2$ transport map optimization problem admits a (Kantorovich) dual with appealing properties which enable the use of several learning/optimization algorithms to (approximately) solve the problem. However, the dual objective function includes a term involving the convex conjugate of the potential function (which is the dual optimization variable), which is hard to compute in many settings, and even impossible in some others. The author proposes taking an \"amortized optimization\" view for the conjugate computation problem within the iterations of estimating the potential function (e.g., while approximating it with NNs). The amortized optimization view is shown to be useful since the potential function estimation stage will involve solving the conjugate sequentially/iteratively.\n\n**Note:** My score increased from \"6\" to \"8\" after the rebuttal.", "summary_of_the_review": "The paper is in general written well with a very accurate literature review, the topic studied is very interesting with high relevance to both optimal transport and optimization communities. Some of the results that are discussed in detail are already well-known (if not they are easy to derive as there is no new theory, but rather a combination of known methods), the paper is written in a subjective/informal way, and in general, the paper looks like a review paper rather than a novel one. I recommend \"marginally above acceptance threshold\", but I do not have strong positive feelings either. I decided in favor of acceptance as the paper provides a complete set of reusable algorithms, which the community can benefit from.", "strength_and_weaknesses": "**Strength:** Firstly, I would like to thank the author for such a clear motivation, literature review, numerical experiments, implementation, and discussions. The paper is very \"fun\" to read, and I had a good time reading it. The focus on computing the conjugate is very essential to the relevant community. \n\n**Weaknesses:** \n*I list my major and minor concerns that hold me from being *strictly* positive about this paper. However, I still think the paper has some interesting ideas, therefore recommending an acceptance. I am also going to stay active during the discussion period in case the author has updates or questions regarding my review.*\n\n***Major Weaknesses/Questions***\n- Most of the remarks and optimization-related discussions are very redundant and they do not add much to the paper. For example, in Remark 8 it is discussed that the first-order conditions would not give a global solution for nonconvex functions. Do we even need this? I do not think any reader who wouldn't already know this will understand even page 1 of the paper.\n- The tone of the author may be taken as subjective/informal or even offensive for some people. There are several **bold highlighted** sentences that are subjective and these may make the paper look like some course notes or informal discussion. Remark 7 is a good example. Also, Remark 6 \"in my experiments, updating .. hurts ...\" is not formal.\n- Some results are not very new, but the way they are presented may sound like a breakthrough from the tone. I would like to kindly ask the author whether Section 4.1 adds anything new. Isn't this just a parallel line search implementation?\n- I would be happier to see further optimization formality. Validity of $\\min$ instead of $\\inf$ is typically not argued. The existence of feasible solutions for, for example, duality arguments, is not discussed.\n- Some problems are claimed to be \"easy\", for example, \"**computing the exact conjugate is easy**\". Is it really an easy problem? Or is that a numerical observation?\n- Sometimes there are overly positive claims for simple/standard evidence. For example, in Section 3.1, the selection of a gradient for the amortization of a function is said to be well-motivated because the $\\arg\\min$ of the conjugate is also a derivative.\n- Finally, a general question to the author: The OT setting looks relevant because we need to know the conjugate and because we would apply gradient-based learning to solve approximations of (2), but in general, this solution technique is a way of estimating the conjugate of a potential function. I am wondering, is there anything else that makes the underlying OT setting unique? Moreover, amortized optimization is already being used a lot in similar settings, may I please ask the author what makes this setting special? Apologies if I am not seeing something obvious, and again, many thanks for the work.\n \n***Minor Weaknesses/Questions***\n- Abstract has \"an amortized approximations\". \n- Whenever $:= \\arg\\min$ notation is used, could the author please discuss that the $\\arg\\min$ set is a singleton? Otherwise, $\\in \\arg\\min$ would perhaps be more suitable.\n- \"solving\" the conjugation sounds a little confusing in Section 1, especially since $f$ is an optimization function. Maybe the context behind the parametrization of $f$ should be discussed first.\n- Page 2: \"the the exact conjugate\" (typo)\n- Equation (4): $f_\\theta^\\star(x)$ has a typo I believe. \n- Equation (4): $J_{f_\\theta} (x(y))$ -> this shorthand is not explained yet (shorthand for $J_{f_\\theta} (x(y); y)$\n- When Algorithm is first mentioned in Section 1, the \"initial dual potential $f_\\theta$\" is not explained yet (which appears in the algorithm). Similarly, the same algorithm uses $\\tilde{x}_{\\varphi}(y_j)$ which is not yet discussed. Also, the same algorithm, maybe $N$ can be shown as an input, too?\n- Remark 1: \"$W_2$ coupling\" is not abbreviated yet. Section 3.1: \"MLP\" not abbreviated.\n- Remark 2: Perhaps the author should discuss the feasibility $x \\in \\mathcal{X}$, too (in order to make a weak-duality sort of argument)?\n- Last paragraph of Section 2: \"by predicting the solution with a model $\\tilde{x}_{\\varphi}(y_j)$\". Solution of \"what\" is not clear, and $\\tilde{x}_{\\varphi}(y_j)$ is not introduced.\n- Section 3.1, element 1: \"directly maps to the solution by\" -> please also state \"solution of what\"\n- Section 3.2.1: \"as optimal as possible\" -> this is not a very usual terminology\n- I am curious: how does the approach mentioned right after (6) relate to the standard SAA-like techniques in Stochastic Optimization?\n- Page 4: \".. et al ... proposes\" -> 'propose'\n- Remark 7: \"... state that they are not a maximin optimization procedure\" -> this is not clear to me\n- Question: Could the author please discuss how someone would still be interested in 3.2.3 even when the true solution to the conjugate problem is known? (for example, when we use NN based approximations for the potential?)\n- Page 12, Nhan Dam et al. (2019) citation -> please capitalize \"GAN\".\n- Conclusions: The sentence that starts with \"In non-Euclidean ....\" is hard to grasp.", "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": "oSe6-bfwBzw", "reviewer": "Reviewer_G2nL", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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 paper suggests fine-tuning a learnable conjugate map when solving the Wasserstein-2 distance dual formulation (where the potential map is parameterized using a neural network) can lead to significant improvement of the resulting transport map as demonstrated on the Wasserstein-2 benchmark by Korotin et al. (2021a). To facilitate the optimization during fine-tuning, a new parallel line search for L-BFGS is implemented in Jax which is much faster than Jax's default line search method.\n", "summary_of_the_review": "Overall I think this is a good paper, despite not having a lot of novelty as mentioned above. The experiments are carefully done. OT practitioners can all learn something useful from reading this paper. \n", "strength_and_weaknesses": "## Strengths:\n* The problem of computing the conjugate map effectively is a central one in neural OT.\n* The proposed fine-tuning step is simple in concept/implementation yet very effective.\n* Careful software engineering is done to obtain good experimental results.\n* Lots of good empirical advice from the remarks in the paper for OT practitioners.\n\n\n## Weaknesses:\n* The novelty of the proposed method is quite limited. There is no new methodological formulation or new theoretical results. To me, it seems mostly like a nice engineering hack (which works quite well).\n* The novelty of the new line search method is also limited. The proposed line search improvement is a straightforward way to convert a classical sequential algorithm involving conditional branches to a batched version.\n* In a couple of the remarks there are some new perspectives (e.g. Remark 5, 7) of viewing existing formulation in different ways, but I do not find these new perspectives reveal particularly interesting new insights.\n* The paper does not reveal a new understanding of why certain combinations of parts of the algorithms work better than others. Some heuristics are provided but could be better explained (e.g. Remark 8 and Figure 1).\n\n## Detailed comments:\n* When citing Brenier's theorem below (3), it should be mentioned that it only holds when one distribution has a density (which I believe is the assumption for this paper and most of the related works).\n* The message of Remark 8 and Figure 1 is not clear to me. Does it just mean solving (9) is easier than non-convex (6) and (8)?\n* In Table 2, ICNN performs very poorly compared to MLP, even with fine-tuning. What could be an explanation for this, other than ICNN is just hard to train as commonly understood? I figure fine-tuning using regression loss should make ICNN easier to train.\n\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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "4QYzJgsCok", "reviewer": "Reviewer_dkEA", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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 studies efficient method for computing convex conjugate arising in wasserstein OT problem. Although exactly and approximately computing so is believed to be hard in prior work, this work proposes a new method based on amortized approximation scheme could be used to computing the exact conjugate. This amortization with fine-tuning also has favorable practical performance, improving over all Wasserstein-2 benchmark tasks, and produce stable results on synthetic dataset.", "summary_of_the_review": "I am unfamiliar with this area and especially for the computational perspective for 2-wasserstein OT. The paper reads like trying to address an important problem, and the practical performance is quite compelling. My slight concerns are regarding the novelty of the algorithmic idea, lack of support in theory, and a potential for a better presentation. I'd love to hear from the author's perspective before giving a final assessment from my side.", "strength_and_weaknesses": "This paper studies an important problem in OT and provides interesting idea and compelling practical implementation of the proposed method. The experiment part looks solid and convincing. The theoretical part feels less convincing for the following reasons:\n- If not mistaken, it feels the key idea of amortizing largely follows the amortizing optimization framework in Amos [2022] and other prior work, and the main part of Section 3 is mostly restating this framework in the particular setting of OT. Most new designs tailored to the setting are discussed in Section 4, but most of it is also still using well known optimizer like Adam and L-BFGS to directly solve the Conjugate function. In that case, the contribution of the theoretical part of the paper reads unclear to me.\n- Though I am unsure possible or not, it may be helpful to have at least some convergence guarantee, stability analysis, or even hard instance which explicitly shows why *all methods* need to be sensitive to model's hyper-parameters.\n- There are many remarks in the paper, some of them seem a bit vague and in particular, require a very good knowledge on previous work to be able to understand, e.g. Remark 6, Remark 11. It may be helpful to discuss their formulations first and then distinction to this work in detail in the main text.", "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": {}}]

Author rebuttal:

Decision: Accept: poster