O N AMORTIZING CONVEX CONJUGATES FOR OPTIMAL TRANSPORT
Anonymous authors
Paper under double-blind review
A BSTRACT
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 difﬁcult 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 ﬁne
-tuning signiﬁcantly improves the quality of transport maps learned for the Wassersteinbenchmark by Korotin et al.
(2021a)
and is able to model many 2-dimensional couplings and ﬂows considered in the literature.
I NTRODUCTION
Optimal transportation (Villani, 2009; Ambrosio, 2003; Santambrogio, 2015; Peyré et al., 2019) is a thriving area of research that provides a way of connecting and transporting between probability measures.
While optimal transport between discrete measures is well-understood, e.g. with Sinkhorn distances (Cuturi, 2013), optimal transport between continuous measures is an open research 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 et 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 et al., 2020; Fan et 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 α and β in
Euclidean space, i.e. supp( α ) =
supp( β ) =
R n with the Euclidean distance as the transport cost.
The Wasserstein-2 transport map , ˘ T : R n → R n , is the solution to Monge’s primal formulation :
[Equation]
where T ( α, β ) :
= { T : T # α = β } is the set of admissible couplings and the push-forward operator # is deﬁned by T # α ( B ) :
= α ( T − 1 ( B )) for a measure α , measurable map T , and all measurable sets B .
˘ T exists and is unique under general settings, e.g. as in Santambrogio (2015, Theorem 1.17), and is often difﬁcult 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é et al.
This paper focuses on the dual associated with the negative inner product cost (Villani, 2009, eq. 5.12), which introduces a dual potential function f : R n → R and solves:
[Equation]
where L 1 ( α ) is the space of measurable functions that are Lebesgue-integrable over α and f (cid:63) is the convex conjugate , or Legendre-Fenchel transform, of a function f deﬁned by:
[Equation]
J f ( x ; y ) with objective J f ( x ; y ) :
= f ( x ) − (cid:104) x, y (cid:105) . ˘ x ( y ) denotes an optimal solution to eq.
Even though the eq.
(2) searches over functions in L 1 ( α ) , 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 ˘ T ( x ) = ∇ x “ f ( x ) , and the inverse to the transport map is given by ˘ T − 1 ( y ) =
∇ y “ f (cid:63) ( y ) .
A stream of foundational papers have proposed methods to approximate the dual potential f with a neural network and learn it by optimizing eq.
(2): Taghvaei and Jalali (2019); Korotin et al.
; Makkuva et al.
(2020) parameterize f as an input-convex neural network (Amos et al., 2017), which can universally represent any convex function with enough capacity (Huang et al., 2020).
Other works explore parameterizing f as a non-convex neural network (Nhan Dam et al., 2019; Korotin et al., 2021a; Rout et al., 2021).
Efﬁciently 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 Taghvaei and Jalali (2019) is considered computationally challenging and approximating it as in Korotin et al.
; Makkuva et al.
(2020); Nhan Dam et al.
; Korotin et al.
(2021a)
; Rout et al. (2021) may be instable.
Korotin et al.
(2021a) fortiﬁes 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 to nearly-optimal saddle point.
In contrast to these statements on the difﬁculty of exactly estimating the conjugate operation,
I will show in this paper that computing the (near-)exact conjugate is easy.
My key insight is that the approximate, i.e. amortized , conjugation methods can be combined with a ﬁne-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 on 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 ﬁne-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 ﬁne-tuning also nicely models synthetic settings (sect. 6), including for learning a single-block potential ﬂow without using the likelihood.
L EARNING DUAL POTENTIALS : A CONJUGATION PERSPECTIVE
This section reviews the standard methods of learning parameterized dual potentials to solve eq.
The ﬁrst step is to re-cast the Kantorovich dual problem eq.
(2) as being over a parametric family of potentials f θ with parameter θ as an input-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 θ V ( θ ) where V (
[Equation]
where J is the objective to the conjugate optimization problem in eq.
(3), ˘ x ( y ) is the solution to the convex conjugate, and eq.
(4) assumes a ﬁnite solution to eq.
(2) exists and replaces the sup with a max .
Taghvaei 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, i.e. using ∇ θ V .
This derivative going through the loss and conjugation operation can be obtained by applying Danskin’s envelope theorem (Danskin, 1966; Bertsekas, 1971) and results in only needing derivatives of the potential:
[Equation]
where ˘ x ( y ) is not differentiated through.
Assumption 1 A standard assumption is that the conjugate is smooth with a well-deﬁned 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).
In practice, assumption 1 in not guaranteed, e.g. non-convex potentials may have a parameterization that results in the conjugate taking inﬁnite values in regions.
The dual objective in eq.
(2) and eq.
(4) discourage the conjugate from diverging as the supremum involves the negation of the conjugate.
Remark 1 In eq.
(4) , the dual potential f associated with the α measure’s constraints is the central object that is parameterized and learned, and the dual potential associated with the β measure is given by the conjugate f (cid:63) and does not require separately learning.
Because of the symmetry of W 2 couplings, 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 parameterized and we will next see that it is the most natural to think about the model associated with the conjugate as an amortization model.
Remark
The dual objective V can be upper-bounded by replacing ˘ x with any approximation because any sub-optimal solution to the conjugation objective provides an upper-bound to the true objective, i.e. J (˘ x ( y ); y )
≤ J ( x ; y ) for all x .
In practice, maximizing a loose upper-bound can cause signiﬁcant divergence issues as the potential can start over-optimizing the objective.
Computing the updates to the dual potential’s parameters in eq.
(5) is a well-deﬁned machine learning setup given a parameterization of the potential f θ , 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 ˜ x φ ( y ) .
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 ﬁne-tune the amortized predictions with an explicit solver CONJUGATE ( f, y, x init =
˜ x φ ( y )) .
Algorithmsummarizes how to learn a dual potential with an amortized and ﬁne-tuned conjugate.
A MORTIZING CONVEX CONJUGATES : MODELING AND LOSSES
This section scopes to predicting an approximate solution to the conjugate optimization problem in eq.
This is an instance of amortized optimization methods which predict the solution to a 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 ∼ β every time the dual V from eq.
(4) is evaluated across a batch.
Using the basic setup from Amos (2022), I call a prediction to the solution of eq.
(3) the amortization model ˜ x ϕ ( y ) , which is parameterized by some ϕ .
The goal is to make the amortization model’s prediction match the true conjugate solution, i.e.
˜ x φ ( y ) ≈
˘ x ( y ) , for samples y ∼ β .
In other words, amortization uses a model to simultaneously solve all of the conjugate optimization problems.
There are two main design choices: sect.
discusses parameterizing the amortization model and sect.
3.2 overviews amortization losses .
3. P ARAMETERIZING A CONJUGATE AMORTIZATION MODEL
The amortization model ˜ x ϕ ( y ) maps a point y ∼ β to a solution to the dual, i.e. ˜ x ϕ : R n → R n .
In this paper, I take standard potential models further described in app.
B and keep them ﬁxed to ablate across the amortization loss and ﬁne-tuning choices.
The main categories are:
˜ x ϕ : R n → R n directly maps to the solution by an MLP as in Nhan Dam et al.
or a U-Net (Ronneberger et al., 2015) for image-based transport.
These are also used in parts of Korotin et al.
(2021a).
˜ x ϕ = ∇ y g ϕ is the gradient of a function g ϕ : R n → R .
Korotin et al.
; Makkuva et al.
(2020) parameterize g ϕ as an input-convex neural network, and some methods of Korotin et al.
(2021a) parametrize g ϕ as a ResNet (He et al., 2016).
This is well-motivated because the arg min of a convex conjugate is the derivative, i.e. ˘ x ( y )
= ∇ y f (cid:63) ( y ) .
: Measures α and β to couple, initial dual potential f θ , and initial amortization model ˜ x ϕ while unconverged do Sample batches { x j } ∼ α and { y j } ∼ β indexed by j ∈ [ N ] Obtain the amortized prediction of the conjugate ˜ x ϕ ( y j ) Fine-tune the prediction by numerically solving ˘ x ( y j ) =
CONJUGATE ( f, y j , x init =
˜ x ϕ ( y j ))
Update the potential with a gradient estimate of the dual in eq.
(5), i.e. ∇ θ V
Update the amortization model with a gradient estimate of a loss from sect.
3, i.e. ∇ ϕ L end while return optimal dual potentials f θ and amortization model ˜ x ϕ
3. C ONJUGATE AMORTIZATION LOSS CHOICES
−  x  ∝ (cid:107)∇ J f ( x ) (cid:107)
Cycle
J f ( x ; y )
Objective (cid:107) x − x (cid:63) ( y ) (cid:107)
Regression x (cid:63) ( y )
We now turn to the design choice of what loss to optimize so that the amortization model ˜ x
ϕ best-predicts the solution to the conjugate.
In all cases, the loss is differentiable and ϕ is optimized with a gradient-based optimizer.
I present an amortization perspective of methods not previously presented as amortization methods, which is useful to help think about improving the amortized predictions with the ﬁne-tuning and exact solvers in sect.
Figure 1 illustrates the main loss choices.
3.2. O BJECTIVE - BASED AMORTIZATION
Nhan Dam et al.
(2019) propose to make the amortized prediction as optimal as possible on the conjugation objective J f from eq.
(3) across samples from
[Equation]
We refer to L obj as objective-based amortization and solve eq.
(6) by taking gradient steps ∇ ϕ L obj using a Monte-Carlo estimate of the expectation.
Remark
The maximin method proposed in Makkuva et al.
(2020, theorem 3.3) is equivalent to maximizing an upper-bound to the dual loss V with respect to θ of a potential f θ and minimizing the objective-based amortization loss L obj with respect to ϕ of an amortization model ˜ x
ϕ := ∇ g ϕ .
Their maximin formulation replaces the exact conjugate ˘ x in eq.
(4) with an approximation ˜ x ϕ , i.e.: [Equation] θ
[Equation] ϕ )
[Equation]
V MM ( θ, ϕ ) [Equation]
Makkuva et al.
(2020) proposes to optimize V MM with gradient ascent-descent steps.
For optimizing θ , V MM ( θ, ϕ ) is an upper bound on the true dual objective V ( θ ) as discussed in remark 2 with equality if and only if ˜ x ϕ = ˘ x .
Evaluating the inner optimization step is exactly the objectivebased amortization update, i.e., ∇ ϕ V MM ( θ, ϕ ) =
∇ ϕ L obj ( ϕ ) =
∇ ϕ J f θ (˜ x ϕ ( y ); y ) .
Remark 4 Suboptimal predictions of the conjugate often leads to a divergent upper bound on V ( θ ) .
Makkuva et al.
(2020, algorithm 1) propose to ﬁx this by running more updates on the amortization model.
In sect.
4, I propose ﬁne-tuning as an alternative to obtain the exact conjugates.
3.2. F IRST - ORDER OPTIMALITY AMORTIZATION : CYCLE CONSISTENCY AND W2GN
An alternative to optimizing the dual objective directly as in eq.
(6) is to optimize for the ﬁrst-order optimality condition .
Eq.
(3) is an unconstrained minimization problem, so the ﬁrst-order optimality condition is that the derivative of the objective is zero, i.e. ∇ x J f ( x ; y )
=
∇ x f ( x ) − y
= 0 .
The amortization model can be optimized for the residual norm of this condition with min ϕ L cycle ( ϕ ) where L cycle ( ϕ ) :=
E y ∼ β
(cid:107)∇ x J f (˜ x ϕ ( y ); y ) (cid:107) 2
=
E y [Equation]
(cid:107)∇ x f (˜ x ϕ ( y )) − y (cid:107) 2 2 .
Remark 5 W2GN (Korotin et al., 2019) is equivalent to maximizing an upper-bound to the dual loss V with respect to θ of a potential f θ and minimizing the ﬁrst-order amortization loss L cycle with respect to ϕ of an amortization model ˜ x
ϕ := ∇ g ϕ .
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 an alternative way of deriving the cycle consistency term by amortizing the ﬁrstorder optimality conditions of the conjugate.
Remark 6 The formulation in Korotin et al.
(2019) does not disconnect f θ when optimizing the cycle loss in eq.
From an amortization perspective, this performs amortization by updating f θ to have a solution closer to ˜ x ϕ rather than the usual amortization setting of updating ˜ x
ϕ to make a prediction closer to the solution of f θ .
In my experiments, updating f θ with the amortization term seems to help when not ﬁne-tuning the conjugate to be exact, but not when using the exact conjugates.
Remark 7 Korotin et al.
and followup papers such as Korotin et al.
(2021b) state that they are not a maximin optimization procedure as in eq.
(7) from Makkuva et 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 ﬁrst-order optimality condition (eq.
Sect.
5 shows that adding ﬁne-tuning and exact conjugates to both of these methods makes their performance match in most cases.
Remark 8 Optimizing for the ﬁrst-order optimality conditions may not be ideal for non-convex conjugate objectives as inﬂection points with a near-zero derivative may not be a global minimum of eq.
The left and right regions of ﬁg. 1 illustrate this.
3.2. R EGRESSION - BASED AMORTIZATION
The previous objective and ﬁrst-order amortization methods locally reﬁne the model’s prediction using local derivative information.
The amortization model can also be trained by regressing onto ground-truth solutions when they are available, i.e.
[Equation]
reg
[Equation]
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.
L reg gives the amortization model information about where the globally optimal solution is rather than requiring it to only locally search over the conjugate’s objective J .
N UMERICAL SOLVERS FOR EXACT CONJUGATES AND FINE TUNING
Algorithm 2 CONJUGATE ( f, y, x init ) x ← x init while unconverged do
Update x with ∇ x J f ( x ; y ) end while return optimal ˘ x ( y ) =
x
In the Euclidean Wasserstein-2 setting, the conjugation operation in eq.
(3) is a continuous and unconstrained optimization problem over a possibly non-convex potential f .
It is usually implemented with a method using ﬁrst-order information for the update in algorithm 2, such as:
Adam (Kingma and Ba, 2014) is an adaptive ﬁrst-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.
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 LBFGS efﬁciently to run on the batches of optimization problems considered here.
Choosing the line search method is the most crucial part as the conditional nature of some line searches may be prohibitive over batches.
Table 3 shows that an Armijo search often works well to obtain approximate solutions.
4. R ELATED WORK ON 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 numerical conjugation, which also have applications in solving Hamilton-Jacobi or Burger’s equation 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 et al.
(2021a) that we report in sect.
5: we are able to conjugate dual potentials in up to 256-dimensional spaces for the HD tasks and 12228-dimensional spaces for the CelebA64 tasks.
Taking a grid-based discretization of the space with 10 locations in each dimension would result in (12228) 10 grid points in the CelebA64 task.
E XPERIMENTAL RESULTS ON THE W ASSERSTEIN -2 BENCHMARK
I have focused most of the experimental investigations on the Wasserstein-2 benchmark (Korotin et al., 2021a) because it provides a concrete evaluation setting with established baselines for learning potentials for Euclidean Wasserstein-2 optimal transport.
The tasks in the benchmark have known (ground-truth) optimal transport maps and include transporting between: 1) high-dimensional (HD) mixtures of Gaussians, and 2) samples from generative models trained on CelebA (Liu et al., 2015).
The main evaluation metric is the unexplained variance percentage ( L 2 -UVP) metric from (Korotin et al., 2019), which compares a candidate map T to the ground truth map T (cid:63) with:
[Equation]
In all of the experimental results, I report the ﬁnal L 2 -UVP 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 to understand how ablating the amortization losses and ﬁne-tuning options impacts the ﬁnal L 2 -UVP performance of the learned potential.
The main experimental takeaway of this paper is that ﬁne-tuning the amortized conjugate with a solver signiﬁcantly improves the learned transport maps.
Tables 1 and 2 report that amortizing and ﬁne-tuning the conjugate improves the L 2 -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 ﬁne-tuned within 100ms per batch of 1024 examples on an NVIDIA Tesla V100 GPU, ﬁg. 2 and app.
C.2 compare Adam and L-BFGS for solving the conjugation.
The following remarks further summarize the results from these experiments:
Remark 9 With ﬁne-tuning, the choice of regression or objective-based amortization doesn’t signiﬁcantly impact the L 2 -UVP of the ﬁnal potential.
This is because ﬁne-tuning is usually able to ﬁnd the optimal conjugates from the predicted starting points.
Remark 10 My re-implementation of W2GN (Korotin et al., 2019), which uses cycle consistency amortization with no ﬁne-tuning, often outperforms the results reported in Korotin et al.
(2021a).
This is likely due to differences in the base potential and amortization models.
Remark 11 Cycle consistency sometimes provides difﬁcult starting points for the ﬁne-tuning methods, 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 difﬁcult for saddle points
(Dauphin et al., 2014).
Combining cycle consistency, which tries to ﬁnd a point where the derivative is zero, with
L-BFGS, which also tries to ﬁnd a point where the derivative is zero, results in ﬁnding suboptimal inﬂection points of the potential rather than the true minimizer.
Remark 12 The performance of the methods using objective-based amortization without ﬁnetuning, as done in Taghvaei and Jalali (2019), are lower than the performance reported in Korotin et al.
(2021a).
This is because I do not run multiple inner updates to update the amortization model.
I instead advocate for ﬁne-tuning the conjugate predictions with a known solver, eliminating the need for a hyper-parameter of the number of inner iterations that needs to be delicately tuned to make sure the amortized prediction alone does not diverge too much from the true conjugate. [Table] [Table] None
[Table] Cycle [Table] Objective
[Table]
Objective
[Table]
Regression
[Table]
the reversed direction from Korotin et al.
(2021a), i.e. the potential model is associated with the β measure
[Figure]
[Equation] [Figure]  D EMONSTRATIONS ON SYNTHETIC DATA I lastly demonstrate the stability of amortization and ﬁne-tuning as described in algorithm 1 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 ﬁne-tuning.
Figures 3 to 5 show the settings considered in Makkuva et al.
(2020) and
Rout et al.
(2021), and
ﬁg. 6 shows the conjugate objective landscapes.
Figure 7 shows maps learned on synthetic settings from Huang et al.
App.
D contains more experimental details here.
Remark 13 Optimizing the dual in eq.
(4) is an alternative to the maximum likelihood training in
Huang et al.
(2020) for potential ﬂows.
While maximum likelihood training requires the density of one of the measures, optimizing the dual only requires samples from the measures.
This makes it easy to compute ﬂows such as the bottom one of ﬁg. 7, even though it is difﬁcult to specify the density.
C ONCLUSIONS , FUTURE DIRECTIONS , AND LIMITATIONS
This paper explores the use of amortization and ﬁne-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 ﬂows (Alvarez-Melis et al., 2021;
Mokrov et al., 2021), or cellular trajectories (Bunne et al., 2021).
And, many of the key amortization and ﬁne-tuning concepts from here will transfer beyond the Euclidean Wasserstein-2 setting.
In nonEuclidean optimal transport where a more general c -transform arises, e.g. in Sei (2013); Cohen et al.
; Rezende and Racanière (2021), or for amortizing 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 signiﬁcant limitation in the ﬁeld of estimating Euclidean Wassersteinoptimal 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.
I have shown that improved conjugate estimations signiﬁcantly improve the stability when the base potential model is properly set up, but all methods are sensitive to the potential model’s hyper-parameters.
I found that small changes to the activation type or initial learning rate can cause no method to converge.
[Figure]
β
Potential f
Conjugate f ∗
[Figure] [Figure]
[Figure] [Figure]
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 a setting BFGS (Broyden, 1970; Fletcher, 1970; Goldfarb, 1970;
Shanno, 1970) and L-BFGS (Liu and Nocedal, 1989) methods thrive in.
The default strong Wolfe line search methods in the Jax and JaxOpt L-BFGS implementations may take a long time to solve a batch of optimization problems. Without efﬁciently setting the line search method, some of the Wasserstein-2 benchmark experiments in app.
C that ran in a 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 efﬁcient at computing the conjugate.
A. B ACKGROUND ON BFGS METHODS
Nocedal and Wright (1999, alg. 6.1) is a standard reference for BFGS methods and extensions and the key steps are summarized in algorithm 3 for solving an optimization problem of the form
[Equation]
where J : R n → R is possibly non-convex and twice continuously differentiable.
The method iteratively ﬁnds a solution ˘ x by 1) maintaining an approximate Hessian around the current iterate, i.e. B k ≈ ∇ 2 J ( x k ) , 2) computing an approximate Newton step p k = − B − 1 k ∇ x J k ( x k ) using the approximate Hessian, 3) updating the iterate with x k +
= x k + α k p k , where α k is found with a line search , and 4) updating the Hessian approximation with the Sherman–Morrison–Woodbury formula (Woodbury, 1950; Sherman and Morrison, 1950)
[Equation]
where y k = ∇ x J ( x k +1 ) − ∇ x J ( x k ) and s k
= x k +1 − x k .
Instead of estimating B k and inverting it in every iteration, most implementations maintain a direct approximation to the inverse Hessian
H k :=
[-The line search to ﬁnd the step size α k for the iterate update x k +
 = x k + α k p-]B − 1 k .
The limited-memory version of BFGS (L-BFGS) in Liu and Nocedal (1989) proposes to replace the inverse Hessian approximation as a matrix H k with the sequence of vectors [ y k , s k ] deﬁning the updates to H k and never requires instantiating the full n × n approximation.
Inputs: Iterate x k , search direction p k , decay τ , control parameter c 1 , initial αα ← α init while J ( x k + α j p k ) >
J ( x k ) + c 1 α j p (cid:62) k ∇ x J
( x k ) do α ← τα end while return α satisfying the Armijo condition in eq.
Algorithm
Inputs: Iterate x k , search direction p k , decay τ , control parameter c 1 , initial α 0 , #evaluations M
Compute candidate step lengths α m = τ − m for m ∈
[ M ]
Evaluate the line search condition g ( α m ) from eq.
(16) in parallel if all g ( α m ) < 0 then return Error: No acceptable step found elsereturn Largest α satisfying eq.
(17), i.e. max α m subject to g ( α m ) >end if
[-Remark 
-]{+The line search to ﬁnd the step size α k for the iterate update x k +
 = x k + α k p k is often done with: 1.+}
a Wolfe line search (Wolfe, 1969) to satisfy the conditions:
[Equation]
where 0 < c
1 < c 2 < 1 , 2.
a strong Wolfe line search to satisfy the conditions:
[Equation]
This is often found via the
zoom procedure from Nocedal and Wright (1999, algorithm 3.5).
an Armijo line search (Armijo, 1966) to satisfy the ﬁrst condition:
[Equation]
For notational simplicity, we can also write the Armijo condition as:
[Equation]
{+Remark 
+}
The strong Wolfe line search is the most commonly used line search for L-BFGS as it guarantees that the resulting update to the Hessian in eq.
(12) stays positive deﬁnite, but the Armijo line search may be more efﬁcient 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 (cid:62)
k s k > 0 that ensure the Hessian update stays positive deﬁnite 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 (cid:62) k s k > 0 , or a modiﬁcation thereof, as described in Li and Fukushima (2001); Wan et al.
; Fridovich-Keil and Recht (2020) and Berahas et al.
(2016, sect. 3.2).
[Table]
T HE A RMIJO LINE SEARCH
The Armijo line search can be written as the optimization problem α k ( x k , p k ) = max α subject to J
[Equation]
[Equation]
Eq.
(17) is typically solved as shown in algorithm 4 and ﬁg. 8 by setting a decay factor τ and iteratively decreasing a candidate step length α until the condition is satisﬁed.
When the objective J can be efﬁciently evaluated in parallel on a GPU [Table] 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 a different number of iterations for every optimization problem in the batch.
Wolfe line searches such as Nocedal 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 in the batch.
I propose a parallel Armijo line search in algorithm 5, which is also visualized in ﬁg. 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 α m satisfying the Armijo condition g ( α m ) ≥ 0 .
Remark 15 The parallel line search may unnecessarily evaluate more candidate steps sizes than the sequential line search, but on GPU architectures this may not be very detrimental to the performance because additional parallel function evaluations are computationally cheap.
Furthermore, when solving a batch of N optimization problems with M line search evaluations, i.e. when using vmap on the line search or optimizer, the parallel line search in algorithm 5 can efﬁciently evaluate NM candidate step lengths in tandem on a GPU and then select the best for each element in the batch.
Remark 16 A potential concern with this parallel line search is that it may not ﬁnd a step size satisfying the Armijo condition if M is not set to be low enough.
While this may be a signiﬁcant issue for when high-precision solves are needed, I have found in practice for the Euclidean Wassersteinconjugates that taking M = 10 line search evaluations frequently ﬁnds a solution.
A.2. C OMPARING LINE SEARCH METHODS
Table 3 takes a trained ICNN potential and isolates the comparison between L-BFGS runtimes to only changing the linesearch methods.
This is 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 indeﬁnite Hessian approximations, the Armijo line searches are consistently able to solve the batch of optimization problems the fastest.
B M ODEL DEFINITIONS AND PRETRAINING
All of the potential and amortization models in this paper can be implemented in ≈ 30 – 50 lines of readable Jax code with Flax (Heek et al., 2020).
They are included in the w2ot/models directory of the code, and reproduced here to precisely deﬁne them.
B. P RETRAINING AND INITIALIZATION
Following Korotin et al.
(2021a), every experimental setting has a pre-training phase so that the potentials and amortization maps approximate the identity mapping, i.e. ∇ x f θ ( x ) ≈ x and ˜ x ϕ ( y ) ≈ y .
B. InitNN : N ON - CONVEX NEURAL NETWORK AMORTIZATION MODEL ˜ x ϕ
Remark 17 The passthrough on line 18 is helpful for learning an identity initialization.
1 class InitNN(nn.
Module): dim_hidden: Sequence[ int ] ’elu’ 4@ compact 19return z
ICNN : I NPUT - CONVEX NEURAL NETWORK POTENTIAL f θ actnorm is the activation normalization layer from Kingma and Dhariwal (2018), which was also used in the ICNN potentials in Huang et al.
and normalizes the activations at initialization to follow a normal distribution.
Remark 18 Applying an activation to the output on line 41 is helpful to lower-bound the otherwise unconstrained potential and adds stability to the training.
Remark
The ﬁnal quadratic on line 46 makes it easy to initialize the potential to the identity.
Remark
This ICNN does not use the quadratic activations proposed in Korotin et al.
Appendix B.1).
While I did not heavily experiment with them, table 1 shows that this ICNN architecture 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.
_bias
True, kernel ) 23w_xs.append(
.Dense use bias ) self.w w 26 @nn.compact def __call__ self, x): assert x.ndim
n_input = x.shape[-1] act_fn layers.get
(self.act) 33z
for in
if alpha ’log
47return y
B. PotentialNN : N ON - CONVEX NEURAL NETWORK (MLP) POTENTIAL f θ
Remark 21 Consistent with remarks 18 and 19, applying an activation to the last layer (line 18) and adding a ﬁnal quadratic term (line 24) helps this non-convex potential model too.
class PotentialNN(nn.
Module): dim_hidden: Sequence[ int ] act: str = ’elu’ 4@nn.compact def call , x):
25return z B. ConvPotential : N ON - CONVEX CONVOLUTIONAL POTENTIAL f θ
{+Remark 
+}
I was not able to easily add batch normalization (Ioffe and Szegedy, 2015) to this potential.
In contrast to standard use cases of batch normalization that only call into a batchnormalized model once over samples from a single distribution, the dual objective in eq.
(4) calls into the potential multiple times to estimate E x ∼ α f θ ( x ) and E y ∼ β f θ (˘ x ( y )) , which also involve internally solving the conjugate optimization problem in eq.
(3) to obtain ˘ x .
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 x ∼ α .
, kernel_size=[4,4], strides=2) y = act_fn(conv(y))
27conv = nn.
Conv(256, kernel_size=[4,4], strides=2) y = act_fn(conv(y))
30conv = nn.
Conv(512, kernel_size=[4,4], strides=2) y = act_fn(conv(y))
33conv =
nn.
Conv(1024, kernel_size=[4,4], strides=2)
y
= act_fn(conv(y))
36conv = nn.
Conv( 1, kernel_size=[2,2], padding=’VALID’, strides=1)
y = act_fn(conv(y))
y
= y.squeeze([1,2,3]) 41assert
y.shape =
= (num_batch,) 43log_alpha = self.param( ’log_alpha’, nn.initializers.constant(0), []) y += 0.5*jnp.exp(log_alpha)*utils.batch_dot(x_flat, x_flat) 47return
y
C A DDITIONAL W ASSERSTEIN -2 BENCHMARK EXPERIMENT DETAILS C. H YPER - PARAMETERS
Tables 4 and 5 detail the main hyper-parameters for the Wasserstein-2 benchmark experiments.
I tried to keep these consistent with the choices from Korotin et al.
(2021a)
, e
[Table]
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 if all dimensions of the iterates change by less than 0. [Table] and otherwise run for a maximum ofiterations.
The line search parameters for the parallel 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, I use the default β = [0 . 9 , 0 . 999] with an initial learning rate of 0.1 with a cosine annealing schedule to decrease it to 10 − 5 .
[Table]
[Table] [Table]
C
[Figure]
C ONVERGENCE OF L-BFGS AND A DAM FOR SOLVING THE CONJUGATE
Figure 9 shows that with a non-convex potential, many of the initial amortized predictions are suboptimal and difﬁcult for the L-BFGS to improve upon.
This indicates that the amortized predictions may be in parts of the space that are difﬁcult to recover from and suggests a future avenue of work better characterizing and recovering from this behavior.
L-BFGS converges fast to an optimal solution in ﬁg. 10 while Adam often gets stuck at suboptimal solutions
[Figure]
[Figure]
A DDITIONAL RUNTIMES AND CONJUGATION INFORMATION
Tables 6 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 it from the amortization model’s prediction, 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 ﬁne-tuning.
On the HD benchmark, ﬁne-tuning takes between ≈ 10 – 50 ms per batch.
The overall wall clock time may take ≈ 2 – 3 times longer than the training runs without ﬁne-tuning, but are able to ﬁnd signiﬁcantly better solutions.
On the CelebA64 benchmarks, the conjugation time impacts the overall runtime even less because, especially in the “Mid” and “Late” settings as the transport maps here are close to being the identity mapping and are easy to conjugate.
Remark 23 Some settings immediately diverged to an irrecoverable state providing a L 2 -UVP of 10 9 , including runs using the objective-based and cycle amortization losses without ﬁne-tuning.
I early-stopped those experiments and do not report the runtimes or conjugation times here, as the few minutes that the objective-based amortization experiments took to diverge is not very interesting or comparable to the times of the experiments that converged.
[Table] [Table]
D A DDITIONAL 2 D DEMONSTRATION EXPERIMENT DETAILS
Table 8 details the main hyper-parameters for the synthetic benchmark experiments, and ﬁg.shows additional conjugation landscapes
[-Remark 
-][Table]
Remark 24 I found leaky ReLU activations on the potential model to work better in these lowdimensional 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 sharper transitions in the space, e.g. the sharp boundaries shown in ﬁg. 4.
One reason that the ELU potentials could perform better on the benchmark settings is that the ground-truth transport maps in the benchmark, described in Korotin et al.
(2021a, Appendix B.1), 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-performing potential model, which are:
• Makkuva et al.
(2020): an ICNN.
This setting originally considered convex potentials, and the non-convex potentials I tried training on these settings diverged, • Rout et al.
(2021): a non-convex potential (an MLP).
This setting also originally considered an MLP and I couldn’t ﬁnd an ICNN that accurately transports between the highly curved and concentrated parts of the measures.
• Huang et al.
(2020): a non-convex potential (an MLP).
In contrast to the ICNNs originally 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 of ICNNs, which means the ﬂow will not necessary be the optimal transport ﬂow, while the non-convex MLP potential I 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 a 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 a base of τ = 1 . 5 and to search M = 30 step sizes.
[Figure]
[Table]
D. N ON - CONVEX REGIONS IN THE LEARNED POTENTIALS
Brenier’s theorem (Brenier, 1991) shows that the known Wasserstein-2 optimal transport map associated with the negative inner product cost is the gradient of a convex function, i.e. ˘ T ( x )
= ∇ x “ f ( x ) .
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 (cid:63) .
The input-convex property is usually enforced by constraining all of the weights of the network to be positive in every layer except the ﬁrst.
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 [Table] 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 appealing as their parameter space is less constrained and may therefore be easier to search over.
And in practice, this has been shown to be true, e.g. the main results in table 1 show that a non-convex potential signiﬁcantly outperforms the convex potential.
However, non-convex potentials can result in non-convex conjugate optimization problems in eq.
(3) that can cause signiﬁcant numerical instabilities and an exploding upper-bound on the dual objective.
Figure 12 illustrates a small non-convex region arising in a 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 ﬂat, but the constraints on the parameter space may hinder the performance.
Interpolation from a non-convex potential (an MLP)
[Figure]
( ∇ f (cid:63) ) # β ← ((1 − t ) I + t ∇ f (cid:63) )
# β →
Interpolation from a convex potential (an ICNN)
[Figure]
( ∇ f (cid:63) ) # β ← ((1 − t ) I + t ∇ f (cid:63) )
# β →
[Figure]
D. I NTERPOLATIONS ON SYNTHETIC SETTINGS