argmax centroid

1 Introduction . Many problems in machine learning and statistics involve optimizing a random function such as the empirical loss . This work focuses on approximating the distribution of the optimum points of a random function of interest . Specifically , we consider the following problem : The Argmax Centroid Problem Assume we are given a random function f⇠ ( ✓ ) where ⇠ is a random variable and ✓ is a variable of interest . We are interested in estimating the distribution ⇢⇤ of the minimum points ✓⇠ : = argmin✓ f⇠ ( ✓ ) as we draw ⇠ randomly . In particular , we want to find a set of “ centroid points ” { ✓i } ni=1 , whose empirical measure is close to ⇢⇤ in terms of Wasserstein distance . We assume the minimizer of f⇠ is always unique . This problem can find applications in a variety of machine learning and statistics techniques , including bootstrap ( Efron & Tibshirani , 1994 ) , random MAP ( Hazan et al. , 2013 ) , Thompson sampling ( Russo et al. , 2017 ) , as well as multi-task learning and meta learning . One naive method to approximate ⇢⇤ , which is widely used in the methods above , is to draw an i.i.d . sample of ⇠ and calculate the corresponding argmin points . However , this approach is computationally expensive because calculating each argmin point requires solving an independent optimization problem and the quality of Monte Carlo approximation is poor unless n is very large , especially for large-scale optimization problems in deep learning . In this work , we propose a more efficient centroid approximation to replace the Monte Carlo sampling , in which we explicitly optimize a set of points ✓ = { ✓i } ni=1 to approximate target distribution ⇢⇤ . By carefully choosing the location of each point ✓i ( rather than drawing them randomly ) , our method allows us to obtain a set of points ✓ that are well aligned to approximate the overall argmin distribution ⇢⇤ . Theoretically , under proper conditions , our method can be viewed as minimizing a surrogate of the Wasserstein distance between the empirical distribution of the centroid approximation and the ground-truth distribution . Therefore , given a small budget on sample size , we can estimate the argmin distribution more accurately than using Monte Carlo estimation . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Argmax centroids can find applications in a variety of machine learning problems . In the empirical studies of this work , we mainly focus on meta learning and multi-task learning tasks . For these problems , we learn from a distribution of datasets/tasks/domains and ⇠ denotes a random dataset drawn from some population . Using argmax centroids allows us to obtain an ensemble of centroid models that can capture the uncertainty and variation in different datasets and domains . We test our method on multiple tasks , e.g. , the few-shot image classification , the personalized dialogue system , and the multi-target domain adaptation . For all the above tasks , our method can be easily implemented and enhance the performance of a number of baseline methods , including but not limited to some of the recent state-of-the-art ( SOTA ) methods , e.g . IFSL ( Yue et al. , 2020 ) . 2 Methodology . We introduce our problem description , propose our algorithm to approximate argmin distributions , and study its properties . Problem description Many problems in machine learning can be framed as minimizing or maximizing random functions f⇠ ( ✓ ) where ⇠ is some random variable drawn from a distribution ⇡ on a space ⌅ and ✓ 2 ⇥ ✓ Rd . That is , ✓⇠ = argmin ✓2⇥ f⇠ ( ✓ ) . ( 1 ) We assume the minimum of f⇠ ( ✓ ) is unique . Denote by ⇢⇤ the distribution of ✓⇠ where ⇠ ⇠ ⇡ , which we call the argmin distribution of f . We are interested in approximating ⇢⇤ . The problem in ( 1 ) finds wide applications in machine learning in various different ways , including random MAP ( Hazan et al. , 2013 ) , bootstrapping ( Efron & Tibshirani , 1994 ) , and multi-domain/meta learning as we elaborate in Section 3 . As mentioned in Section 1 , a naive method to approximate ⇢⇤ is to draw an i.i.d . sample { ⇠i } ni=1 from ⇡ and calculate the corresponding argmin points ✓⇠i for each particle respectively . However , the quality of Monte Carlo approximation is poor when n is small ( Kuo & Sloan , 2005 ) . In practice , it is not affordable to draw a large number of argmin points for large-scale optimization problems in deep learning , both because that it requires to solve an optimization for calculating each ✓⇠i , and it is memory-hungry to store a large number of ✓⇠i when the parameter dimension is high . Argmax Centroids We propose a more efficient centroid approximation to replace Monte Carlo sampling , where we explicitly optimize a set of points ✓ = { ✓i } ni=1 2 ⇥n , associated with a set of weights ⌫ = { ⌫i } ni=1 2 Rn , such that the ground-truth target distribution ⇢⇤ can be well approximated by the weighted empirical measure : ⇢̂✓ , ⌫ = nX i=1 ⌫i ✓i , where ✓i is delta measure centered at ✓i , and ⌫ is assumed to take values from the probability simplex on [ n ] : V = { ⌫ : Pn i=1 ⌫i = 1 , ⌫i 0 , 8i 2 [ n ] } . Here [ n ] : = { 1 , 2 , · · · , n } . Ideally , we would like to choose ✓ , ⌫ to minimize certain distance metrics between ⇢̂✓ , ⌫ and ⇢⇤ , a canonical example of which is the p-Wasserstein distance ( p > 0 ) , defined as Wp ( ⇢̂ , ⇢ ⇤ ) = inf µ2⇧ ( ⇢̂ , ⇢⇤ ) E ( ✓ , ✓⇠ ) ⇠µ [ k✓ ✓⇠k p ] 1/p , where ⇧ ( ⇢̂ , ⇢⇤ ) denotes the set of probability measures on ⇥⇥⇥ such that its two marginal distributions on ⇥ are ⇢̂ and ⇢⇤ . We assume that k·k is the standard Euclidean norm in this paper . However , it is expensive to calculate the Wasserstein distance due to the high computational cost of drawing ✓⇠ ⇠ ⇢⇤ , which requires to repetitively solve optimization ( 1 ) . To improve the computational efficiency , we propose to minimize the following surrogate of Wasserstein distance : min ✓2⇥n , ⌫2V Wf ( ⇢̂✓ , ⌫ , ⇡ ) , Wf ( ⇢̂ , ⇡ ) : = inf µ2⇧ ( ⇢̂ , ⇡ ) E ( ✓ , ⇠ ) ⇠µ [ f⇠ ( ✓ ) ] , ( 2 ) where ⇧ ( ⇢̂ , ⇡ ) is the set of probability measures on ⇥⇥ ⌅ , whose marginal distributions on ⇥ and ⌅ equal ⇢̂ and ⇡ , respectively . Here we replace the norm k✓ ✓⇠kp with the function f⇠ ( ✓ ) , so that we do not need to draw ✓⇠ ⇠ ⇢⇤ and hence solve the expensive optimization in ( 1 ) . A key property of Wf is that its global minimum in the space of distributions is achieved by the argmin distribution ⇢⇤ . Hence , we obtain increasingly accurate estimate of ⇢⇤ as we solve ( 2 ) with increasing particle size n. To introduce the result , we start with a general definition of argmin distributions that works for when the minimum of f⇠ is not necessarily unique . Definition 2.1 . Assume f⇤⇠ : = min✓2⇥ f⇠ ( ✓ ) > 1 for every ⇠ 2 ⌅ . We say that a distribution ⇢⇤ on ⇥ is an argmin distribution of f⇠ with ⇠ ⇠ ⇡ , if there exist a coupling measure µ⇤ 2 ⇧ ( ⇢⇤ , ⇡ ) , such that ⇢⇤ ( ✓ 2 argmin f⇠ ) = 1 . Here argmin f⇠ is the set of global minima of f⇠ , that is , argmin f⇠ = { # 2 ⇥ : f⇠ ( # ) = f⇤⇠ } . If the optimum of f⇠ is unique for every ⇠ 2 ⌅ , then there is an unique argmin distribution ⇢⇤ , which is the distribution of ✓⇠ = argmin✓2⇥ f⇠ ( ✓ ) when ⇠ ⇠ ⇡ . Theorem 2.2 . Assume ⇢⇤ is any argmin distribution of f⇠ with ⇠ ⇠ ⇡ . Then the minimum of Wf ( ⇢ , ⇡ ) is achieved by ⇢⇤ , that is , Wf ( ⇢ ⇤ , ⇡ ) = E⇠⇠⇡ min ✓2⇥ f⇠ ( ✓ ) = min ⇢2P Wf ( ⇢ , ⇡ ) , where P is the space of probability distributions on ⇥ . It is not easy to directly minimize the form of Wf in ( 2 ) because it requires to jointly minimize ✓ , ⌫ and the coupling measure µ . However , the key observation below shows that the optimization in both ⌫ and µ can be solved in closed form , yielding a simple objective function on ✓ . Theorem 2.3 . For a fixed ✓ 2 ⇥n , we have min ⌫2V Wf ( ⇢̂✓ , ⌫ , ⇡ ) = E⇠⇠⇡ min i2 [ n ] f⇠ ( ✓i ) , where the minimum on the left hand side is achieved by ⌫⇤ = { ⌫i } ni=1 with ⌫⇤i = E⇠⇠⇡ '' P ( i 2 argmin j2 [ n ] f⇠ ( ✓j ) ) # , 8i 2 [ n ] , ( 3 ) with P ( i 2 argminj2 [ n ] f⇠ ( ✓j ) ) : = 1Z✓ , ⇠ I ( i 2 argminj2 [ n ] f⇠ ( ✓j ) ) and Z✓ , ⇠ the number of elements in the argmin set , i.e. , Z✓ , ⇠ = | argminj2 [ n ] f⇠ ( ✓j ) | . Here the ⌫⇤i in ( 3 ) is the probability that ✓⇤i achieves the minimum value of minj2 [ n ] f⇠ ( ✓⇤j ) , with ties broken randomly with equal probabilities . Therefore , the optimal centroids ✓⇤ and the weights ⌫⇤ are given by ✓⇤ = argmin { ✓i } ni=1 ⇢ L ( ✓ ) : = E⇠⇠⇡ min i2 [ n ] f⇠ ( ✓i ) , ⌫⇤i = E⇠⇠⇡ P ( i 2 arg min j2 [ n ] f⇠ ( ✓ ⇤ j ) ) . ( 4 ) Note that objective in ( 4 ) reduces to the well known K-means objective function if we take f⇠ ( ✓ ) = k✓ ✓⇠k2 , in which case ✓ and ⌫ are the centers and the sizes of the clusters , respectively . Intuitively , minimizing L ( ✓ ) allows us to ensure that the best function value mini2 [ n ] f⇠ ( ✓i ) is small on average . In practice , we can optimize ✓i using stochastic gradient descent , and update ⌫i recursively . See Algorithm 1 for details . At each iteration , we draw a ( set of ) ⇠ ⇠ ⇡ , find the ✓i that attains i = argminj f⇠ ( ✓j ) and update ✓i with gradient descent as displayed in line 4 , Algorithm 1 . The frequency ⌫i is also updated accordingly when a ✓i is updated , see line 5 in Algorithm 1 . Re-sampling During optimization , it is possible that some ✓i are rarely selected and hence its weight ⌫i becomes small and causes a degeneration problem ( i.e. , nearly zero important weight ) . To address this problem , we monitor the effective sample size of the points , defined as neff = ( P i ⌫i ) 2/ ( P i ⌫ 2 i ) . If the effective sample size becomes smaller than a threshold , we re-sample a set of new particles ✓i with replacement from ⇢̂✓ , ⌫ = Pn i=1 ⌫i ✓i and repeat the updates . In practice , we break the ties randomly in argmin to avoid the case when two ✓i remain to be identical throughout the algorithm . Algorithm 1 Main Algorithm : Argmax Centroids for Approximating ⇢⇤ 1 : Input : { ✓i } ni=1 denotes the learnable centroids , ⌫i denotes the updated frequency of ✓i , ✏ denotes the step size for gradient descent , ↵ 2 [ 0 , 1 ] denotes a parameter to control the frequency update and ⌘ is a threshold to control the resampling . 2 : while not convergence do 3 : Find i⇤ = argmini f⇠ ( ✓i ) with ⇠ ⇠ ⇡ . 4 : Update ✓i⇤ ✓i⇤ ✏r✓i⇤ f⇠ ( ✓i⇤ ) , ⌫i⇤ ↵⌫i⇤ + 1 , and ⌫i = ↵⌫i for i 6= i⇤ . 5 : If neff ⌘ , resample { ✓i } ni=1 from ⇢̂✓ , ⌫ / Pn i=1 ⌫i ✓i . 6 : end while Bounds with Wasserstein Because Wf coincides with the Wasserstein distance if f⇠ is a simple quadratic function f⇠ ( ✓ ) = ||✓ ✓⇠||2 , we can provide some simple bounds between Wf and Wasserstein distance by approximating f⇠ with the quadratics using Taylor approximation . Assumption 2.4 . Let ✓⇠ be the minimum of f⇠ ( ✓ ) as in ( 1 ) . Assume there exists h1 , h2 , p1 , p2 2 ( 0,1 ) , such that for any ✓ 2 ⇥ and ⇠ 2 ⌅ , h1 k✓ ✓⇠kp1 f⇠ ( ✓ ) f⇠ ( ✓⇠ ) h2 k✓ ✓⇠kp2 . This assumption holds with p1 = p2 = 2 if f is strongly convex w.r.t . ✓ and 0 < h1/2 min ( f 00⇠ ( ✓ ) ) max ( f 00⇠ ( ✓ ) ) h2/2 < 1 for 8✓ and ⇠ , where min and max are the minimum and maximum eigenvalue ; it also holds with p2 = 1 if f is Lipschitz w.r.t . ✓ with kfkLip = h2 . Theorem 2.5 . Assume Assumption 2.4 holds . 1 ) We have for any ✓ 2 ⇥n and ⌫ 2 V , h1Wp1 ( ⇢̂✓ , ⌫ , ⇢ ⇤ ) p1 Wf ( ⇢̂ , ⇡ ) Wf ( ⇢⇤ , ⇡ ) h2Wp2 ( ⇢✓ , ⌫ , ⇢⇤ ) p2 . 2 ) For the optimal ✓⇤ and ⌫⇤ in ( 4 ) , we have Wp1 ( ⇢̂✓⇤ , ⌫⇤ , ⇢ ⇤ ) p1 h2 h1 inf ✓2⇥n , ⌫2V Wp2 ( ⇢✓ , ⌫ , ⇢ ⇤ ) p2 . ( 5 )

Summary: This paper proposes a method called Argmax Centroids for learning an approximation to the distribution of the minimizer of a random function. The central idea is to learn a set a particles such that the expected distance between the minimizer of a randomly drawn function and the closest particle is small. The method is shown to be connected theoretically to Wasserstein distance. Experiments on few-shot classification, personalized dialogue modeling, and multi-target domain adaptation demonstrate improvements when the technique is applied to baseline models.