Modified Frank Wolfe in Probability Space

1 Introduction . Problems in artificial intelligence , statistics , and optimization often find a common root as an infinite dimensional optimization problem in the form inf { J ( µ ) : µ ∈ P ( Rd ) } , ( 1 ) for the space P ( Rd ) of Borel probability measures over Rd . In recent years , quantitative statistical and algorithmic treatments of these formulations have produced insights into modern computational methods– resulting in novel approaches to difficult , open problems . Recent works in robust optimization [ 6 , 44 , 54 , 55 ] , probabilistic fairness [ 57 , 53 ] , reinforcement learning [ 63 , 64 ] , and generative adversarial networks [ 42 , 18 , 19 ] highlight these gains and are linked by the following theme : problems in the form of ( 1 ) provide access to rich infinite dimensional structure that sidesteps brittle artifacts of finite dimensional formulations . This paper provides the construction and analysis of a modified Frank-Wolfe algorithm for ( 1 ) that operates from this infinite dimensional perspective and yields concrete convergence and complexity guarantees for a sub-family of problems ( 1 ) which are well-behaved with respect to the Wasserstein distance of order 2 ( see Algorithm 1 and Theorem 1 ) . Under canonical conditions of smoothness and convexity we recover linear rates of convergence while , even for functionals which exhibit low degrees of smoothness and for conditions that go beyond convexity , we recover sublinear rates that are to be expected from finite dimensional analogues [ 36 ] ( see Section 2.2 ) . The vanilla Frank-Wolfe method can not work in probability space , in general , since the planar derivative ( i.e . the first variation also known as the influence function ) can be unbounded when distributions do not have compact support . To overcome this issue , we conduct a natural modification 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to introduce a tractable , local , linear constraint– inspired by efforts in DRO [ 44 , 6 , 25 , 54 ] . Specifically , the modified Frank-Wolfe step admits the prototypical DRO formulation as below , sup { ∫ f dµ : Dc ( µ , µ0 ) ≤ δ } , ( 2 ) whereDc ( µ , µ0 ) is the optimal transport cost between µ and µ0 ( a reference measure ) under some cost function c. The form of ( 2 ) , itself , immediately suggests the basis of an infinite dimensional FrankWolfe procedure since it provides a “ linear ” objective subject to a local , “ trust-region ” constraint– centered at µ0 . The relevance of ( 2 ) in distributionally robust optimization ( DRO ) and mathematical finance has resulted in a multitude of computational schemes [ 44 , 34 , 37 , 54 , 61 ] for solving ( 2 ) . However , hitherto , such works have failed to consider ( 2 ) within the context of a general variational method for ( 1 ) . What makes these efforts notable , within the scope of this work , is that they emphasize that the solution of ( 2 ) can be highly non-trivial . Indeed , without particular assumptions , ( 2 ) can disguise an computationally hard problem– despite being convex in a Banach sense on P ( Rd ) . Even in the case where the cost is the squared Euclidean norm c ( x , y ) = ‖x− y‖2 ( the case of primary concern for this work ) , computational trouble can lie dormant– an artifact of inherently difficult problems in unconstrained optimization [ 14 ] . This should not be surprising , however , given specters of computational hardness dating back to early formulations of DRO [ 22 ] . To resolve these issues , we also provide a novel analysis of techniques from distributionally robust optimization ( DRO ) which illustrate how such formulations can be used , in a computationally tractable way , to construct first-order , variational methods . By localizing the problem to a Wasserstein ball , the new Local Linear Minimization Oracle ( LLMO ) that we make in this paper not only makes the vanilla Frank-Wolfe problem sensible , but it also renders the problem computationally tractable via finite dimensional convex optimization . 1.1 Previous work . Distributionally robust optimization To navigate such pitfalls , one can consider particular instances of ( 2 ) where the objective and constraints are sufficiently structured to preclude computational intractability and permit solution via methods adapted to the provided structure . Early work with this line [ 27 , 22 , 60 ] , has recently been supplemented by approaches [ 13 , 26 , 7 , 46 , 37 , 67 , 35 ] which focus directly on DRO formulations from particular contexts in machine learning and operations research . Unfortunately , the techniques offered by these efforts require assumptions which are too restrictive for this work . These assumptions typically relate to a specific form for the objective function or constraints in ( 2 ) ( e.g . linear/convex functions/piecewise-convex objectives or constraints with support or density requirements , see [ 31 , 66 , 44 , 34 , 65 , 4 , 58 ] for additional examples ) . In this instance , such limitations preclude their applicability since , in general , a “ gradient object ” for a functional J ( see Section 2 ) need not satisfy these conditions . A second , more relevant , approach to compute DRO problems ( 2 ) is to restrict the level of robustness δ for which the problem is solved . This technique has been used in works such as [ 6 , 54 ] and we apply this principle in a similar spirit to [ 54 ] . In that work , smoothness of the objective in ( 2 ) is used to , qualitatively , argue that a sufficiently small δ provides a computationally-tractable optimization problem . In contrast , we provide quantification of the level of robustness required to achieve this tractability and we demonstrate that this level robustness is sufficiently large to achieve canonical convergence rates for an infinite-dimensional Frank-Wolfe algorithm . Variational methods Although formulation of a Frank-Wolfe method for ( 1 ) ( with quantitative bounds on complexity and convergence ) has not appeared in previous literature , certain , tangentiallyrelated variational methods offer conceptual similarity . The first of these methods is [ 41 ] , which draws similar inspiration from finite-dimensional Frank-Wolfe procedures . However , the notion of first-order variation in [ 41 ] appears to be induced by the total variation distance– requiring compact support assumptions for the problem . Alternatively , we exploit Wasserstein geometry to provide a weaker notion of first-order variation , namely , Wasserstein differentiability– allowing us to eliminate restrictions to compact support . Hence , by balancing tractability of the analysis and fidelity of the procedure , the proposed algorithm scheme can be applied a broad class of computational examples . Indeed , [ 41 ] only discussed the Sinkhorn barycenter problem and it is still unclear how it can be applied to our setting . The second , more closely related effort , is [ 39 ] where similar , infinitedimensional conditions ( Section 2.2 ) are used to study particle-based methods for computing Nash equilibria of zero-sum games . While the setting and procedures developed in this work are completely different , we note that , as a special case , our Frank-Wolfe method can also produce a particle-based optimization procedure . This hints at possible insightful connections with other particle methods [ 40 , 24 , 11 , 10 ] which are left for future consideration . 2 Wasserstein geometry . This work considers the problem min ν∈P2 ( Rd ) J ( ν ) ( 3 ) for functionals J : P2 ( Rd ) → R̄ over the subset of Borel measures P ( Rd ) defined by P2 ( Rd ) : = { µ ∈ P ( Rd ) : ∫ Rd ‖x‖2 dµ ( x ) < ∞ } ( 4 ) In particular , we consider functionals J that are differentiable ( Definition 1 ) with respect to Wasserstein distance of order 2 W2 ( µ , ν ) : = inf γ∈Π ( µ , ν ) ∫ Rd×Rd ‖x− y‖2 dγ ( x , y ) ( 5 ) where Π ( µ , ν ) is the set of all joint couplings with marginals µ , ν ∈ P2 ( Rd ) . Common examples of functionals that can be cast within this framework are as follows . Example 1 ( Divergences ) . For any convex , lower-semicontinuous function f : R+ → R such that f ( 1 ) = 0 , one can consider a “ f -divergence ” of the form J ( µ ) : = Df ( µ||ν ) = ∫ Rd f ( dµ dν ) dν . ( 6 ) For instance , if f ( t ) = t log t , ( 6 ) reduces to Kullback-Leibler divergence DKL . Example 2 ( Integral Probability Metrics ) . For a set of real valued functions F on Rd one can define the discrepancy J ( µ ) : = IPM ( µ , ν ) = sup f∈F ∫ Rd f dµ− ∫ Rd f dν ( 7 ) for µ , ν ∈ P ( Rd ) , where ν is a fixed , reference measure . Such discrepancies are termed Integral Probability Metrics ( IPMs ) , although they may not strictly satisfy the requirements of a metric– say , by failing to distinguish all pairs of measures . Instead , for a pair of measures µ , ν ∈ P ( Rd ) , IPMs can be interpreted as measuring the extent to which µ and ν differ on functions in F– or , rather , measuring the extent to which µ and ν can be distinguished by F . Concretely , consider F = { f : ‖f‖H ≤ 1 } where H is a reproducing kernel hilbert space ( RKHS ) . In this case , one obtains the dual formulation of Maximum Mean Discrepancy ( MMD ) . 2.1 Properties and differentiability for Wasserstein space . Under the Wasserstein distance , P2 ( Rd ) is a Polish space [ 59 ] and , more importantly , it is a geodesic space . That is , for every µ , ν ∈ P2 ( Rd ) , there exists a constant-speed geodesic curve µt : [ 0 , 1 ] → P2 ( Rd ) where µ0 = µ , µ1 = ν and W ( µt , µs ) = |t− s|W ( µ0 , µ1 ) ( 8 ) Moreover , there is a bijection between constant-speed geodesics and optimal transport plans [ 1 , Theorem 7.2.2 ] . Every geodesic corresponds to a unique , optimal transport plan γ ∈ Π ( µ , ν ) such that µt = ( ( 1− t ) x+ ty ) # γ where W 2 ( µ , ν ) = ∫ Rd×Rd ‖x− y‖2 dγ ( x , y ) , ( 9 ) and µt is the distribution of the random variable ( 1 − t ) X + tY with the pair ( X , Y ) following distribution γ. Conversely , every optimal transport plan gives rise to a unique geodesic via ( 9 ) . Since our Frank-Wolfe method minimizes a sequence of linear approximations , one must define the notion of a gradient ( of a functional J ) to be compatible with respect to the geometry of these geodesics . In particular , as P2 ( Rd ) is curved underW ( see Appendix A ) , gradients must be defined in terms of a selection in an appropriate cotangent bundle . For Wasserstein space , this cotangent bundle ( denoted CoTanP2 ( Rd ) ) is essentially the set of vector fields on Rd that can be approximated by gradients of smooth functions ( see Appendix A for details ) . This results in the following definition . Definition 1 ( Wasserstein differentiability ) . Let S be a geodesically convex set ; that is , µt ∈ S for any geodesic µt between µ , ν ∈ S. A functional J : P2 ( Rd ) → R̄ is Wasserstein differentiable on S if there is a map F : P2 ( Rd ) → CoTanP2 ( Rd ) such that for all µ , ν ∈ S and a geodesic µt : [ 0 , 1 ] → P2 ( Rd ) between µ and ν , one has lim α→0+ J ( µα ) − J ( µ ) α = ∫ Rd×Rd F ( µ ; x ) T ( y − x ) dγ ( x , y ) , ( 10 ) where γ is the unique optimal transport plan ( 9 ) corresponding to µt . Note that F ( µ ; x ) = ( F ( µ ) ) ( x ) provides an aesthetic way of representing the evaluation at x ∈ Rd of the output of F at µ . The map F is called the Wasserstein derivative of J . Remark 1 . The description of differentiability provided by Definition 1 falls within the general framework of metric derivatives and Wasserstein gradient flows , largely codified in [ 1 ] . This framework is now a well-established component of the theory of Wasserstein spaces , while the relation ( 10 ) , itself , presents only a narrow structuring of ideas from this framework . Definition 1 , however , is often how works in statistical and algorithmic fields interact with this broader area [ 56 , 16 , 38 , 39 ] . Moreover , this literature demonstrates the most motivating feature of ( 10 ) : a large number of functionals of interest for machine learning and statistical inference exhibit Wasserstein gradients in the sense of ( 10 ) . The curious reader is referred to [ 1 , 52 , 9 ] for precise statements of conditions under which ( 10 ) is guaranteed . However , F is intimately relative to the Gateaux differential for J [ 52 , 56 ] . Recall , the Gateaux differential for a functional J exists when there is an appropriate , dual space D∗ on a closed subspace D ⊆ P ( Rd ) such that 〈dJ ( µ ) , ν − µ〉 = lim α→0+ J ( µ+ α ( ν − µ ) ) − J ( µ ) α ( 11 ) for some dJ ( µ ) ∈ D∗ and all µ in some set S such that S−S ⊆ D. In instances where the Gauteaux differential dJ ( µ ) exists , the Wasserstein derivative F will usually exist [ 39 ] and be given by∇dJ ( µ ) . Here , we use the finite dimensional gradient operator∇ ( · ) formally , and omit a rigorous exposition on this operation in the context of CoTanP2 ( Rd ) . It should be noted that computation of the Wasserstein derivative might be difficult . Indeed , for a J in a variational form such as ( 7 ) , computation of the Wasserstein derivative is equivalent to finding a witness function that achieves the supremum [ 52 ] . In the case of a pathological sets ( in ( 7 ) ) , such a task might be intractable . To resolve this issue , and to simplify our treatment , this work utilizes the existence of an oracle for the computation of a Wasserstein gradient . This oracle permits a unified description of our Frank-Wolfe algorithm and abstracts away variation in functional-specific computational cost . Recall that a function in the Hölder space C1 ( Rd ) is called L-smooth if has L-Lipschitz gradients . Definition 2 ( Wasserstein Derivative Oracle ) . Let J : P2 ( Rd ) → R be a Wasserstein differentiable functional on a set S with Wasserstein derivative F : P2 ( Rd ) → CoTanP2 ( Rd ) . A L-smooth Wasserstein derivative oracle over S is an oracle which , given sample access to a distribution µ ∈ S and an error parameter , returns an L-smooth function φ̂µ ∈ C1 ( Rd ) satisfying∥∥∥∇φ̂µ − F ( µ ) ∥∥∥ L2 ( µ ) ≤ ( 12 ) where ‖·‖L2 ( µ ) is the canonical norm on the space L2 ( µ ) of square integrable functions with respect to µ ∈ P ( Rd ) . In this work , the output of this oracle is represented as Θ ( µ , ) . Remark 2 . The qualification that the Wasserstein derivative oracle return an L-smooth function is necessary to exclude the , aforementioned , possibility of a pathological Wasserstein derivative– that would be intractable for use in a computational procedure . In some ways , this is representative of the fact that the cotangent space CoTan ( µ ) at a point µ is too large . Such a condition is common in other variational methods [ 3 , 19 , 64 , 20 ] and is relatively superficial– when coupled with the degree of approximation afforded by . Via smoothing techniques [ 51 , 11 , 38 ] , functionals can often be assumed to have Wasserstein derivatives which are C1 ( Rd ) or are well-approximable by C1 ( Rd ) functions .

Summary: This work investigates optimization in infinite dimensional functionals of probability measures. The authors first define a notion of gradient of functional using Wasserstein gradient flows. Then, the authors propose a Frank-Wolfe procedure in which gradient steps can be computed using finite dimensional and convex optimization methods. The authors show that their algorithm is guaranteed to converge under some assumptions by choosing proper step sizes. Finally, the authors apply their method to two non-parametric estimation problems.