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: The paper proposed a method that optimizes a set of centroid points that approximates the argmax distribution p*. Rather than using the Monte Carlo sampling that draws samples randomly, their approach choose the location of each points that approximate target distribution p*.  This approach can be an alternative for bootstrap and can be applied to deep learning applications. The paper showed the effectiveness of the proposed method on few shot image classification, personalized dialogue systems and multi-target domain adaptation. The proposed algorithm boosts the SOTA performance for few-shot classification and meta learning tasks.