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 paper proposes an algorithm for the minimization of functionals of probability distributions. The authors derive a convergence rate in the case of a smooth and PL function (under the geometry given by the Wasserstein distance of order 2). They also run some synthetic experiments of non-parametric estimation