argmax centroid
Authors: Chengyue Gong
Mao Ye
qiang liu
Venue: NeurIPS 2021 Poster (2021)
Keywords: distribution approximation
Monte Carlo
approximate inference
multi-domain learning
meta learning
few-shot learning

Abstract: We propose a general method to construct centroid approximation for the distribution of maximum points of a random function (a.k.a. argmax distribution), which finds broad applications in machine learning. Our method optimizes a set of centroid points to compactly approximate the argmax distribution with a simple objective function, without explicitly drawing exact samples from the argmax distribution. Theoretically, the argmax centroid method can be shown to minimize a surrogate of Wasserstein distance between the ground-truth argmax distribution and the centroid approximation under proper conditions. We demonstrate the applicability and effectiveness of our method on a variety of real-world multi-task learning applications, including few-shot image classification, personalized dialogue systems and multi-target domain adaptation.

Argmax Centroids: with Applications to Multi-domain Learning
Chengyue Gong
Mao Ye
Qiang Liu
Computer Science Department, The University of Texas at Austin {cygong17, my21, qiang) utexas. edu
Abstract
We propose general method to construct centroid approximation for the dis- tribution of maximum points a random function (a.k.a. argmax distribution), which finds broad applications in machine learning. Our method optimizes centroid points to compactly approximate the argmax distribution with simple objective function, without explicitly drawing exact samples from the argmax distribution. Theoretically, the argmax centroid method can be shown to minimize surrogate of Wasserstein distance between the ground-truth argmax distribution and the centroid approximation under proper conditions. We demonstrate the applicability and effectiveness of our method on variety of real-world multi- task learning applications, including few-shot image classification, personalized dialogue systems and multi-target domain adaptation.
of
Introduction
Many problems in machine learning and statisties 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 fe where is a random variable and f is variable of interest. We are interested in estimating the distribution p" of the minimum points 0g := arg ming (0) as we draw & randomly. In particular, we want to find of "centroid points" {0:)==> whose empirical measure close to p² in terms of Wasserstein distance. We assume the minimizer of Je is always unique.
This problem can find applications in variety of machine learning and statistics techniques, including bootstrap (Eiron& & Tibshirani] (994), random MAP (Hazan et al. 2013), Thompson sampling Russo ct al. 2017). well multi-task learning and meta learning. One naive method to approximate p", 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 is very large, especially for large-scale optimization problems in deep learning.
In this work, we propose more efficient centroid approximation replace the Monte Carlo sampling, in which we explicitly optimize a set of points = { to approximate target distribution p². By carefully choosing the location of each point 0, (rather than drawing them randomly), allows us to obtain a set of points that are well aligned to approximate the overall argmin distribution p" Theoretically, under proper conditions, of the Wasserstein distance between the empirical distribution of the centroid approximation and the ground-truth distribution. Therefore, given small budget sample size, estimate the argmin distribution more accurately than using Monte Carlo estimation.
method
method be viewed minimizing surrogate
35th Conference on Neural Information Processing Systems (NcurIPS 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/domalle 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, c.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 baseline methods, including but not limited to some of the recent state-of-the of-the-art art (SOTA) methods, e.g. IFSL (Yue al. 2020
and & denotes random dataset
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 maxi- mizing random functions (e( where is some random variable drawn from a distribution ı space and f E EOCR" That
0g arg min fe( 0
We assume the minimum of f0) is unique. Denote by p" the distribution of 0g where & ~ which we call the argmin distribution of f. We are interested in approximating p². The problem in { 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 +
As mentioned in Section a naive method approximate p² is draw an i.i.d. sample {=))== from π and calculate the corresponding argmin points gg. for each particle respectively. However, the quality of Monte Carlo approximation is poor when n is small (Kuo & Sloan| 2005). practice, 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 fz., and is memory-hungry to store large number of 0g. when the parameter dimension is high.
Argmax Centroids We propose more efficient centroid approximation replace Monte Carlo sampling. where we explicitly optimize a set of points 0 = {0, (0.)%_. E O", associated with set of weights v = {u:)²= E RR",, such that the ground-truth target distribution p² can be well approximated by the weighted empirical measure:
Po.
where 5a, is delta measure centered at 0, and v assumed to take values from the probability simplex on [n]: V = {v: Ex= 14 1, 14 2 0, Vi E [n]]]. Here [n] == {1, ... ,n].
Ideally, we would like choose 0, v to minimize certain distance metries between سع and p², canonical example of which is the p- Wasserstein distance (p > 0), defined as
W Wp p") pEIl(p,0))
inf 1(0,0 - [| 1/P/
where IT p²) denotes the set of probability measures on θ x θ such that its two marginal distribu- tions on ② are and p". We assume that ||-|| 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 0g p", which requires to repetitively solve optimization ① To improve the computational efficiency, we propose to minimize the following surrogate of Wasserstein distance:
min PEV".PEV
Wp(po T).
Wy(2, T) = inf B(,,))= [fc( well(a.m)
(2)
where IT(2, m) the set of probability measures on E, whose marginal distributions on θ and = equal å and respectively. Here replace the norm ||0 0gll² with the function fo fe()), so that

---

do not need to draw g p" and hence solve the expensive optimization in T. A key property of W is that its global minimum in the space distributions is achieved by the argmin distribution p². Hence, we obtain increasingly accurate estimate of p" as we solve ? with increasing particle size
To introduce the result, start with general definition of argmin distributions that works for when the minimum of fe is not necessarily unique.
Definition 2.1. Assume f == minsee fc(0) -o for every E. We say that distribution p" on θ is an argmin distribution of with & n. if there exist a coupling measure p² I1pp',, T), such that p² (0 E arg min fe) 1. Here arg min Je is the set of global minima of Se, that is, arg min fg { E O: [e((0= [:).
If the optimum of f& is unique for every & E E, then there is an unique argmin distribution p", which is the distribution of 0g = arg mingee fe(0) when & π.
Theorem 2.2. Assume p¹ is any argmin distribution of Se with и. Then the minimum of W (p( achieved by p", that is,
W BENT
min fe( 1E8
min Wx(p,
PEP
where P the space of probability distributions on O
It is not easy to directly minimize the form of W in ( because it requires to jointly minimize 0, and the coupling measure μ. However, the key observation below shows that the optimization in both and y can be solved in closed form, yielding a simple objective function on 0.
Theorem 2.3. For a fixed O", we have
min W (pe T) Bewn
min fe fe(0.)) [iEmn
wEV
where the minimum on the left hand side achieved by v² = {ui) with
v Ee
P(i arg fo (e,)) vi [n], jE(n]
(3)
with Pi E arg minje(m] fe() = Zo. e arg minje(m] fe() and Ze. e the number of elements the argmin set, i.e., Zo.c = arg minje(x) fe(0,))-
Here the of in is the probability that 07 achieves the minimum value min fc(07), with ties broken randomly with equal probabilities.
Therefore, the optimal centroids 0* and the weights D² are given by
0 = arg min 10
((() = E
min fc fe(0.) f=(0))) sEin]
1// = EęNa P(i E arg min fe(09))
(4)
JE
Note that objective in 4 reduces the well known K K-means objective function we take (0)
||e 0g 12 in which case and v are the centers and the sizes of the clusters, respectively. Intuitively, minimizing L(0) allows to ensure that the best function value min;e["] fe().) is small on average.
In practice, we can optimize 0; using stochastic gradient descent, and update v; recursively. See Algorithm [ for details. At each iteration, we draw (set of) & ~ find the 0; that attains i = arg min, fs (0,) and update 0; with gradient descent as displayed in line 4, Algorithm I The frequency vs is also updated accordingly when a 0 is updated, see line 5 in Algorithm [
During optimization, is possible that some 6 are rarely selected and hence its weight i/ becomes small and causes degeneration problem (i.e., nearly zero important weight). To address this problem, we monitor the effective sample size of the points, defined as Thf = (E, (x) xf). If the effective sample size becomes smaller than threshold, we re-sample set of new particles 0 with replacement from ~0.0 Σn_ padg, and repeat the updates. In practice, we break the ties randomly in argmin to avoid the case when two 0 remain to be identical throughout the algorithm.
Re-sampling
3

---

Algorithm Main Algorithm: Argmax Centroids for Approximating p²
Input: {0:)== denotes the learnable centroids, V4 denotes the updated frequency 0,, e denotes the step size for gradient descent, a E [0, 1] denotes parameter to control the frequency update and 17 is threshold to control the resampling.
2: while not convergence do
3: Find i, arg min; fe:) with E~ π.
4: Update 0z, + 0g. VVo, [e(Oi.), Vi, + awy, + and vs au; for i i * in.
5: If Beff I 17, resample [0: | from Pavv x Eim U;Og..
6: end while
Bounds with Wasserstein Because Wy coincides with the Wasserstein distance if fe a simple quadratic function (0) = 110 0gll?. we can provide some simple bounds between W, and Wasserstein distance by approximating fe with the quadratics using Taylor approximation.
Assumption 2.4. Let Og be the minimum of fe(0) as in I (0, oo), such that for any 0 E è and & E,
Assume there exists h1, h2, P1, P2
hi 10 0el1 < (0) fe() h2 10 00
This assumption holds with P1 = P2 = 2 if f strongly convex w.r.t. 0 and 0 < h2/2 s Amin (f (0)) S Amaz and maximum eigenvalue; it also holds with p2 1 if is Lipschitz w.r.t. 0 with Il h2.
( (0) S h < co for v0 and &, where Amin and Xmaz are the minimum
Theorem 2.5. Assume Assumption 2.4 holds.
1) We have for any 0 E tt and E V,
h1 pp (Po.: pr) WW π) - Wy(o", TT) S p2 P²)
2) For the optimal 0* and v² in we have
l22 inf
W (po- p² p')p <
(Po,e pr)
(5)
bn B0 WEV
3 Applications and Related Works
Approximating argmax distribution can be found as a key component of many machine learning methods. This includes methods like bootstrap (Ye & Liu 2021). Thompson sampling, and random MAP, in which the argmax distribution plays a key role by design, and as well as novel applications from meta learning and multi-task learning, which we study extensively in Section 4
Learning with Ensembles of Datasets Approximat- ing the argmax distribution finds natural application in meta learning and multiitassk/multi -omain learning, in which need to learn from an ensemble or distri- bution of datasets. Denote by D a dataset on which we want to learn model parameter 0 with a loss function {(0, D). Instead of having a single dataset, we consider the case when we have a large number of datasets {D:} which we assume is drawn from a distribution P.
Training Datnset
Models
Test Dataset
D
Dr
e,
Seleet Best A,
20
24
Align
25
Train
D6
Test
We aim to find set of models (0, )a where exists least one good model for every dataset D with small n. Denote by 0D arg min, E(O, ()) and p the distribu- tion of 9n as ~ P. Then the problem can framed finding {0. to approximate p" and hence be framed into
An illustration of our selection process for both training and evaluation. Test dataset D refers to the training or vali- dation subset of the coming test task D.
Figure 1:
min (0:)
EDaP
min ((0 D) iE(n)
(6)

---

In the test phase, given new coming dataset D, which we assume includes a validation set D with true labels, we select the best model B, among the centroids that minimizes the loss i, arg [m] E(O,, D). The pipeline is shown in Figure The idea here is to prepare a pool of candidate models to best cover all different scenarios (according to probability 8p ~ p²), so that we can select the best model for any new random task drawn from p". This is different from typical ensemble learning which averages the output of multiple models during evaluation; instead, we use the best single model 0;, selected from the pool during evaluation. The objective E and how to optimize is problem-dependen For example, meta learning, is updated with gradient descent and meta gradient. The details about the formulation of l and its update rule is shown in Section $ for different tasks. The method in ( can also be used in bootstrapping, in which we are given a dataset Do of true interest and we draw random datasets D ~ P by resampling from Do with replacement. Using the argmax centroids can yield better approximation with small n than naive random sampling that calculates @D from randomly drawn D ~ P. The theoretical properties and efficient algorithms specified to bootstrap desires an independent treatment, which we explore separately in another work. Another related method is multiple choice learning (MCL) (Guzman-Rivers et al. 2012 Yu et al. 2018 Lee et al., 2016), which is an ensemble learning method using a similar objective function but with very different motivation and settings. MCL can be viewed minimizing ( but assuming each D to be a single data example rather than a dataset, with the goal of learning a set of models and ensemble them for single dataset. During evaluation, the outputs of the multiple models from MCL are averaged to give the final output, rather than selecting the best one as we suggest above based on the perspective of approximating p²
Multi-fask Learning Jointly learning models from multiple tasks has attracted long-term attention in machine learning community. In multi-task learning (MTL), there are multiple learning tasks and most of them (or all of them) are assumed to be related to each other. MTL algorithms can be grouped into two main approaches: the feature learning methods and task relation learning methods. Since tasks are related,
intuitive to assume that different tasks share some common feature representation, and thus feature learning approaches focus on learning the common representation with regularizations (Yang et al. 2009)), structure designing (Caruana 1997). etc. The other approach focuses learning task relations. The task relations learnt from data automatically with a given
prior (Yu et al. 2005 Xue et al. 2007; Zhang & Yeung 2012), matrix dccomposition (Chen et al.
2012) or clustering (Thrun & O O'Sullivan 1996). Recently meta learning (Finn et al. 2017), lifelong
learning (hrun 1998) and continual learning (Zenke et al. 2017) are proposed to solve MTL under
different conditions.
Our method is similar to learning task relation with clustering, Different from previous works which cluster the tasks with latent representation Zhou & Zhao 2015), learnable relation matrix (Zhou
et al. 201 T Kang et al. 2011 or side information (Caruana 1997), , we automatically assign a model
to certain domain/task based on its loss.
Experiments
In our experiments, we start with a toy example on which we show that our algorithm approximates the argmax distribution better than simple Monte Carlo estimation (Section 4.1 We then apply our method to standard few-shot learning benchmarks, in which parameters are learned on thousands of domains (Section 4.2) We further apply our method to personalized dialogue systems (Peng et al.) 2020) (Section 4.3) and multiple target domain adaptation (Yu et al. 2018) (Section 4.4) to verify
whether our proposed method can be applied to different real-worls applications. In all experiments, we set the replacement controller r)- = 1.2 and a = 0.5 for Algorithm
4.1 Toy Examples
We start with a toy example to verify that our method can provide a better approximation of pa than Monte Carlo estimation with the same number of samples. We set fe( = |e E|| so that 0g and p² = म. We set π to be d-dimension Gaussian mixture model with five modes. Note that our method is equivalent to gradient descent variant of K-means in this case. Figure 2 demonstrates the two-dimension case, from which we can see that argmax centroids captures the different modes in the distribution and are aligned optimally to get better approximation then random sampling.
5

---

Random
Log (Wass)
100 Randoms
1000 Optinized
A 100 Optinized
Sample Size (n) (d)
(a)
(b)
(c)
Figure 2: 2D toy example. Figure (a), (b) and (c) show the estimated centroids {0: (the red dots) and contour (green) the true distribution p². Figure(d) shows the 2- 2-Wasserstein distance (p = 2) between the empirical measure PO ," of the centroids (centroid approximation) and the true distribution '² w.r.t. the sample size n, when we vary the dimension in {10, 100J. 'Log (Wass)' denotes log-scale Wasserstein distance.
d) shows the 2-Wasserstein distance (p = 2) from the empirical measure po- of the centroids and the true distribution p² (calculated by drawing a large sample from p'), when the dimension d of f varies. We can see that for both 100-dimension and 10-dimension cases, our method yields lower Wasserstein distance than random sampling for different
Figure 2
Method
Transductive
Backbone Conv4
1-shot (Acc)
5-shot (Acc)
MAML (Finn et al. 2017
4..7
64.6 1.2
Tian et al. 2020)
ResNet12
64.8:10.6
82.1
Baseline (Dhilonn et al. 2020
WRN-28-10
68.1
80.4
SIB (Hu et al. 2020)
WRN-28-10
70.0:E0.6
80.0:E0.3
Ours
WRN-28-10
71.2:1:0.6
81.1
IFSI. Yue et al. 2020
WRN-28-100
73.2+-0.3
0..9+0.4
+ Ours
WRN-28-10
73.6++0.4
83.3++0.3
Table 1: Few-shot learning results averaged over 3 trials Mini-ImageNee
Method
Transductive
Backbone
1-shot (Acc)
5-shot (Acc)
MAML (Finn et al. 2017)
Cony4
(Tian et al. (0200)
ResNet12
71,5++0.2
86.0.10.5
Baseline (Dhillon et 2020
WRN-28-10
72.9+00.1
86.2:10.5
SIB (uu et al. 2020)
WRN-28-10
81.2+:0.6
87.1
+ Ours
WRN-28-10
81.8++0.6
87.6++0.4
IFSL Yuue et al. 2020)
WRN-28-10
82.4
88.3
+ Ours
WRN-28-10
82.7-0.5
88.5
Table 2: Few-shot Learning results averaged over trials on TieredImageNet.
4.2 Few-Shot Supervised Learning
We now apply our method to few-shot image classification using the formulation in @ in combination with model agnostic meta learning (MAML) (Finn al. 2017).
To introduce the setting, let D = (Drrain Dieat) be a random dataset (or task) drawn P, which includes training and test set. Let 0' A(0, D be the result of training the parameter taain starting from 0, where A denotes the training algorithm, which is typically one step or few steps of stochastic gradient descent with the training set. In MAML, we optimize 0 so that the peformance of = A(B, Drmin) is maximized on D The MAML loss f(0, D) = L(AA(B, D"min), Diear,, with the hope rapidly solve new coming task after learning several other similar tasks.
tasks D when the discrepancy between tasks is large. By plugging the MAML loss into ( we obtain new generalization MAML with which we can prepare pool of candidate models {0,}2_ (the argmax centroids), out of which select the best one during testing. Note that this is different from the
Standard MAML learns single model 0, which may not work well for

---

existing ensemble MAML methods such as ge Lee et al. 2019; Yoon et al. 2018), which averages (instead of selecting the best among) the outputs of multiple models.
Baselines We compare and implement method based the recent state-of-the-art few-shot learning methods (c.g Hu et al. 2020; Yue et al. 2020 [Wang et al. 2019). These methods train network consisting of backbone network shared by all the datasets (tasks), which is fed into classification heads specialized to individual datasets. During training, we first train a backbone network on all the training datasets in the training tasks using typical method, and then use meta- learning style training to tune the classification heads while keeping the backbone frozen.
We implement proposed argmax centroids upon recently proposed few-learning methods, SIB and IFSL.. These recently proposed methods achieve state-of-the-ar results on supervised few-shot image classification benchmarks [] SIB (Hu et al. 2020) proposes learn synthetic gradient for new coming tasks, and IFSL Yue et al. 2020) finds out that pre-trained knowledge essentially a confounder that causes spurious correlations between the sample features and class labels in support set. They further propose regularization terms to eliminate the mismatch.
Method ResNet-18
MAML
Ensemble Yonn 2018
Bayesian Yoon 2818
Ours $7.41.2
56.6
56. E.3
57.0E.4
WRN-28-10
4
57.4+
57.7
58.1+.4
Table 3: Comparison to ensemble MAML methods on Mini-ImageNet. We demonstrate the accuracy for one-shot classification with ResNet-10 and WRN-28-10 We set the number of heads to 16 while sharing the backbone parameters. 'Ensemble' and 'Bayesian' denotes ensemble MAML and Bayesian MAML, respectively.
All these baselines and other recent works (e.g. Hu et al. 2020 Yue et 2020; Wang et al. 2019
Dhillon al,, 2020) focus on learning a better single model by introducing strong regularization methods. For the baselines, we focus on, a clear disadvantage of SIB and IFSL is that they share one single for all the datasets and may not work well once the discrepancy between task is large. By applying method, we train a fixed number of classification heads, which are treated the argmax centroids. This allows us to pick the best head and further finetune it on the training set for each test task during testing. We simply combine these baseline approaches with our algorithm to obtain further improvements. To save computation cost in practice, we evaluate A(e, D Drain) by one-step gradient descent and then select 0, with lowest loss among O.
Datasets Standard benchmarks of few-shot classification are chosen for experiments. We evaluate all the baselines and our algorithms on two subsets of ImageNet, Mini-ImageNet and TieredImageNet (Sun et al. 2019). Mini- ImageNet contains 64 classes for training, 16 for validation and 20 for test. Train- ing/validation/tess tasks are sampled from the 64/16/20 classes. TieredImageNee (Sun al.) 2019) is much larger compared to minilma- geNet and other benchmarks, with 608 classes and 1300 samples per class. These classes are partitioned into 351997/160 disjoint sets for train/val/test to achieve larger domain differ- ence between training and testing During train- ing, the images in Mini- ImageNet and tieredImageNet are resized to 80 and 84x84, respectively. Each sampled dataset D is 5-class classification, with 1-shot/5-shot images per class for training and 15 images per class for test. When evaluating the trained models, 2000 test tasks are sampled from the test set, and the accuracy on Diesr of all the test tasks are reported for comparison.
Beseline
Test Accuracy
73
Solect
71
70
16 Number of Heads n
8
32
64
Figure 3: Training n different heads for the few-
shot learning task MinilmageNNet with SIB (Hu al.
2020). 'Baseline' denotes training one model creating {0. by dropout and select- ing {0,)2_ during evaluation.
https://fer--hhtt.yyl.e

---

Implementation Details For each baseline method, we create small number of additional param- eters. During the second training stage, given a backbone, e.g., Wide-ResNett28..
Zagoruyko & Komodakis, 2016 we use additional parameters {0, J." to re-weight the final-layer features. In the experiments, for few-shot learning based on SIB and IFSL, we set n 16. We only introduce a small number of additional parameters (0.01M parameters) compared to the model size (22M parameters).
For both the training and testing, we keep all the settings (e.g. optimizer, learning rate, number of tasks) the same for each baseline method. We first train the feature extractors using the standard cross- entropy loss on the base classes (i.e., the classes in the training set in all training tasks). Secondly, training the heads following the standard meta learning pipeline. Finally, we test the trained models on sampled test tasks. For each test task, we do few gradient descent steps on the training set. The number of iterations follows the baselines.
Results Table 3 shows the comparison between our proposed method and standard ensemble MAML methods. We implement method upon several recently proposed strong baselines. As demonstrated in Table [ and 2 our algorithm consistently improves all the baseline methods in all settings, which suggests that our proposed algorithm is general method for few-shot image classification. For example, on Mini-ImageNet, we improve the accuracy of SIB by large margin for both I-shot (70.0% 71.2%) and 5-shot (80.0% 81.1%) 1.) classification. For IFSL, we can also boost the performance under most of the settings.
Comparison to Ensemble Methods Here, we follow the setting and ensemble MAMI. YYonn at. 2018) where the backbone is not pretrained on all the training images and is jointly training during the meta learning step. We compare two popular ensemble baselines, ensemble MAML and Bayesian MAML (Yoon et al. 2018). Ensemble MAML trains models independently and take an average of the outputs of all the models during evaluation. Bayesian MAML (Yoon et al. 2018) uses SVGD Liu & Wang 2016) to jointly update n models during training and ensemble them during evaluation. Instead of choosing the best model for a test domain like our method, the ensemble methods average the outputs of all the models. From Table 3 than both ensemble approaches on one-shot classification on Mini- Mini-ImageNer with different neural architectures. On ResNet- 18 and WRN-28-10 10, method achieves better results than ensemble methods.
that method is better
Ablation Study In Figure we change n and make selection based on training/test/validatton loss to have deeper understanding of our algorithm. We notice that, if we do not have a separate validation set, we may need to select the best model on the training set during evaluation, which cause overfitting. We investigate this issue in Figure S in which we show the results when we select the best models using training, validation and, test set, respectively. In order to do the selection on the validation set, additionally draw 5 images per class for each test task. The result shows that, in case, both selecting on training and validation set is worse than the "oracle" way (i.c., selecting on the test set and evaluating on the test set). When the number of heads increases, the performance of selecting on training set drops and over-fitting happens. However, in practice, one cannot get access to the test set during training. Therefore, selecting the validation set is an alternative choice.
4.3 Personalized Dialogue System
Problem Set We consider learning personalized chit-chat dialogue agents for making chat-bots more consistent for each user (e.g. Zhang et al., 2018 Madotto et al. 2019h Tian et al. 2020).
Previous approaches usually learning persona similarity from persona descriptions (Zhang al. 2018). Different from these methods, current personalized dialogue systems focus on learning quickly adapt new personas by using few samples. In this problem, each dataset D = Driin Disrr ~ P contains several dia- logues of user. A dialogue contains a set of utterances, and the objective is to predict the next utterance given previous utterances. We use the MAML objective for training.
Method PAML
PPL 4
BLEU 1
Madorto et al. 2019
46.35
0.77
Ours
43.94
0.91
Table 4: Results of automatic evaluation on per- sonalized dialogue generation. We report case- insensitive BLEU and PPL.
8

---

Baselines We use a recently-proposed meta learning dialogue system (Madotto et al. 2019b) as the baseline. We measure the difference between the ground-truth output and model output evaluate the model performance by PPL. (Perplexity) and case-insensitive BLEU score (Papineni et al. 2002).
Implementation Details Before training. the datasets sampled from P is split into meta- meta-training veta-validation and meta-test dataset. The meta-learner learns how to learn by training and evaluating on the meta-training set. Meta-learning hyper-parameters are tuned on the meta-validation set. The meta-test set measures generalization on new, unseen tasks. We use the BERT-base (Devlin et al.] 2019) model as the backbone and tune the parameters in the backbone together with the final softmax layer. For meta training,
Adam KKngmaa & Ba 2014) with learning rate 10-3, 10 for inner and outer loop training, respectively. During the evaluation, for all the models, we used beam search with beam size 4 and length penalty 1.2. These paramm-arreetters are selected based on the performance of the meta-validation dataset.
For our method, create 16 softmax head (i.e. a linear transform followed by softmax layer), share the backbone parameters and select one head for each persona with the lowest loss on the test set for this persona during training, and select the softmax head with the lowest loss on the training set during evaluation.
Result Table shows that our proposed algorithm beats the meta-learning baseline on both the PPL and BLEU measure. We improve the BLEU from 0.77 to 0.91 I and PPL. from 46.35 to 43.94, respectively.
4.4 Multi-Target Domain Adaptation
Domain adaptation is a powerful approach for learning under distributional shift between training and test data (e.g. Wu al., 2018 Tzeng al. 2017). Here, we show that our proposed algorithm can also be applied to multi-
Test Loss Discriminaton OossO
1.2
Loss
target domain adaptation YYu et al. 2018) using the dataset proposed in Peng et al. (201)). For multi-targer domain adaptation, one source domain and m multi-target domain is given. In this experiment, we use the classical domain adaptation method proposed by Tzeng et al.] (2017). in which an adversarial net is used to learn the adaptation from the latent space in the source domain.
0.9
0.6
5
15 20
25
Number of models
We first pretrain a CNN using labeled image examples from the source domain and denote the backbone g. Next, for the target domain, we perform adversarial adap- tation by learning a map after g such that discriminator that sees source examples g(-) and target examples ((9(-) cannot reliably predict their domain label. During the evaluation, target images are mapped with backbone g
Figure Illustration of the per-domain test loss, and per domain negative dis- criminator loss. For the multi-domain adaptation task, when the number of models increases, the negative discrimi- nator loss always increases while the test loss first decreases and then increases.
followed by the learned map to the shared feature space
and classified by the final-layer classifier trained on the source domain.
Our Method After pretraining a backbone on the source domain, We consider to use n adversarial nets {0:) } domain adaptors to match the representation of the source domain and target domains. Each 02 = (Ø,, 7x) contains two part of parameters where Φb is a feature map and 7/ a discriminator. 6; is an MLP whose input and output have the same dimension and y maps the input to one- dimension scalar followed by a sigmoid function.
Denote = = g(x) the features extracted by the backbone and S the source domain, E(0,, D) equals to
-((((=))) --N-)))))))()) E. === { log (zg - (x)))
- (a=) log (()) (o(()((()1(1(())))
is an adversarial loss which forces the discriminator to make uniform prediction.

---

Baselines We compare our method with three simple baselines. One is using shared adversarial domain adaptors for all the target domains. Another using n different domain adaptors for m target domains. The third baseline is randomly assigning domains with models for each iteration during training and average all the outputs for all models during evaluation.
Experiment Settings We use the DomainNet dataset proposed by Peng et al. (011). DomainNet contains styles, 345 classes and 569,010 images. We create subset of DomainNet to do experiments. We sample 25 classes in the same style and sample 300 images for each class. For each sample subset (domain), we use 5 different data shifts (e.g. brightness, fog, snow, MotionBlur, Gaussian noise) (Hendrycks et al. 2019) to create new domains. Finally, obtain 30 different domains. We randomly select one domain data as the source domain and the rest as target domains during experiments. For the neural architecture, the source domain, a three-layer MLP as "A and another three-layer MLP as o;. During training and testing, we resize the image into 64 64 size. During training, we use standard data augmentation (mirror and flip).
use standard ResNet-18 (He et al. 2016) to train
Method
Adaptor
m Adaptor
Ensemble
Ours
Acc 1
3..
41.6
41.3=0.7
44.3 0.6
Table 5: Results multi-target domain adaptation. The reported accuracy are averaged over 5 trials. "Ensemble" denotes randomly assigning models with domains during training and ensembling all the models during evaluation.
Results As demonstrated in Table 5 compared sharing one domain adaptor for all the domains, using a domain-specific domain adaptor improves the accuracy from 39.5% to 41.8%. With our proposed method, we can further improve the 41.8% accuracy to 44.3%. In Table [ we also demonstrate that randomly aligning different domains hurts the performance. It indicates that in this setting our method is much better than ensemble multiple models. We further set up an ablation study on the number of heads. As shown in Figure g similar to the results in Figure 3 increasing the number of heads can always improve the training loss (i.e., discriminator loss in 7 Here, increasing makes the negative discriminator loss higher indicates that it mixes the two domains successfully. However, the test loss increases once too large. It demonstrates that our argmax centroids method implicitly serves as a regularization by sharing the same linear map á for multiple similar domains. Sharing the same linear map increases the number of training data for each 0 and therefore avoids overfitting,
5 Conclusion
In this work, we propose a method to estimate the argmax distribution. It can be an alternative for related methods, and we further apply to many multi-domain and multi-task learning tasks. Aligning similar domains into same model, the proposed algorithm boost the SOTA performance on several few-shot classification and meta learning tasks. In the concurrent work (Ye & Liu 202 we use the centroid learning idea to improve the bootstrap for data uncertainty quantification, in which we establish a rigorous theoretically result showing that the centroid objective function is uniformly close to Wasserstein distance through the whole optimization trajectory induced by modified gradient descent algorithm. For the next step, we will consider the connection with diversifying models and other topics. We will also apply our method to more meaningful and challenging tasks, e.g. lifelong learning, continual learning.
Acknowledgements The work is conducted in the statistical learning and AI group in computer science department at UT Austin, which is supported part by CAREER-184642) SenSE-2037267, EAGER-204 1327, and Office of Navy Research, and NSF AÍ Institute for Foundations of Machine Learning (IFML.. We would like thank the anonymous reviewers for their thoughtful comments and efforts towards improving our manuscript.
References
Caruana, Rich. Multitask learning. Machine learning, 28(1):41-75, 1997.
10

---

Chen, Jianhui, Liu, Ji, and Ye, Jieping. Learning incoherent sparse and low-rank patterns from multiple tasks. ACM Transactions on Knowledge Discovery from Data, 5(4):1-31, 2012
Devlin, Jacob, Chang, Ming-Wei, Lee, Kenton, and Toutanova, Kristina. Bert: Pre-training of deep bidirectional transformers for language understanding. North American Chapter of the Association for Computational Linguistics, 2019.
Dhillon, Guneet S, Chaudhari, Pratik, Ravichandran, Avinash, and Soatto, Stefano. A baseline for few-shot image classification. International Conference on Learning Representations, 2020.
Efron, Bradley and Tibshirani, Robert J. An introduction to the bootstrap. CRC press, 1994.
Finn, Chelsea, Abbeel, Pieter, and Levine, Sergey. Model-agnostic meta-learning for fast adaptation of deep networks. International Conference on Machine Learning, 2017.
Guzman Guzman-Rivera, Abner, Batra, Dhruv, and Kohli, Pushmeet. Multiple choice learning: Learning to produce multiple structured outputs. In Advances in neural information processing systems, volume 1, pp. 3. Citeseer, 2012.
Hazan, Tamir, Maji, Subhransu, Keshet, Joseph, and Jaakkola, Tommi. Learning efficient random maximum a-posteriori predictors with non-decomposable loss functions. Advances in neural information processing systems, 2013.
He, Kaiming, Zhang, Xiangyu, Ren, Shaoqing, and Sun, Jian. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Hendrycks, Dan, Mu, Norman, Cubuk, Ekin D, Zoph, Barret, Gilmer, Justin, and Lakshminarayanan, Balaji. Augmix: A simple data processing method improve robustness and uncertainty, arXiv preprint arXiv:19/2.0278. 2019.
Hu, Shell Xu, Moreno, Pablo G, Xiao, Yang, Shen, Xi, Obozinski, Guillaume, Lawrence, Neil D, and Damianou, Andreas. Empirical bayes transductive meta-learning with synthetic gradients. International Conference on Learning Representations, 2020.
Kang, Zhuoliang, Grauman, Kristen, and Sha, Fei. Learning with whom to share in multi-task feature learning. International Conference on Machine Learning, 2011.
Kingma, Diederik P and Ba, Jimmy, Adam: A method for stochastic optimization. International Conference on Learning Representations, 2014.
Kuo, Frances Y and Sloan, Ian H. Lifting the curse of dimensionality. Notices of the AMS, 52(11): 1320-1328, 2005.
Lee, Hae Beom, Nam, Taewook, Yang, Eunho, and Hwang, Sung Ju. Meta dropout: Learning to perturb latent features for generalization. In International Conference on Learning Representations, 2019.
Lee, Stefan, Purushwalkam, Senthil, Cogswell, Michael, Ranjan, Viresh, Crandall, David, and Batra, Dhruv. Stochastic multiple choice learning for training diverse deep ensembles. Advances in neural information processing systems, 2016.
Liu, Qiang and Wang, Dilin. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 2016.
Madotto, Andrea, Lin, Zhaojiang, Wu, Chien- Shenn,,,eng,,,, and Fung. Pascale. Personalizing dialogue agents via meta-learning. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp- 5454-5459, 2019a.
Madotto, Andrea, Lin, Zhaojiang, Wu, Chien-Sheng and Fung, Pascale. Personalizing dialogue agents via meta-learning. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 5454-5459, 2019b.

---

Papineni, Kishore. Roukos, Salim, Ward, Todd, and Zhu, Wei-Jing, Bleu: method for automatic evaluation of machine translation. In Annual Meeting of the Association for Computational Linguistics, pp. 311-318, 2002.
Peng, Baolin, Zhu, Chenguang, Li, Chumyuan, Li, Xiujun, Li, Jinchao, Zeng, Michael, and Gao, Jianfeng. Few-shot natural language generation for task-oriented dialog. Conference on Empirical Methods in Natural Language Processing, 2020.
Peng, Xingchao, Bai, Qinxun, Xia, Xide, Huang, Zijun, Saenko, Kate, and Wang, Bo. Moment matching for multi-source domain adaptation. In Proceedings the IEEE International Conference on Computer Vision, pp. 1406-144 1416-115, 2019.
Russo, Daniel, Van Roy, Benjamin, Kazerouni, Abbas, Osband, lan, and Wen, Zheng. A tutorial on thompson sampling. arXiv preprint arXiv:1707.0203: 2017.
Sun, Qianru, Liu, Yaoyao, Chua, Tat-Seng, and Schiele, Bernt. Meta-transfer learning for few-shot learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 403-412, 2019.
Thrun, Sebastian. Lifelong learning algorithms. Learning to learn, pp. 181-209. Springer, 1998.
Thrun, Sebastian and O' O'Slllia,,, Joseph. Discovering structure in multiple learning tasks: The to algorithm. In International Conference Machine Learning, volume 96, pp. 489-497 1996.
Tian, Yonglong, Wang, Yue, Krishnan, Dilip, Tenenbaum, Joshua B, and Isola, Phillip. Rethinking few-shot image classification: good embedding is all you need? Proceedings of the European Conferenceon Computer Vision. 2020.
Tzeng, Eric, Hoffman, Judy, Saenko, Kate, and Darrell, Trevor. Adversarial discriminative domain adaptation, In Proceedings of the IEEE conference on computer vision and patlern recognition, pp. 7167-7176, 2017.
Wang, Yan, Chao, Wei-Lun, Weinberger, Kilian Q, and van der Maaten, Laurens. Simpleshot: Revisiting nearest-neighbor classification for few-shot learning. Advances in neural information processing systems, 2019.
Wu, Zuxuan, Han, Xintong, Lin, Yen-Liang, Gokhan Uzunbas, Mustafa, Goldstein, Tom, Nam Lim, Ser, and Davis, Larry S. Dean: Dual channel-wise alignment networks for unsupervised scene adaptation. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 518-534, 2018
Xue, Ya, Liao, Xuejun, Carin, Lawrence, and Krishnapuram, Balaji. Multi-task learning for classifi- cation with dirichlet process priors. Journal of Machine Learning Research, 8(1), 2007.
Yang, Xiaolin, Kim, Seyoung, and Xing, Eric. Heterogeneous multitask learning with joint sparsity constraints. Advances in neural information processing systems, 22:2151-2155 2009.
Ye, Mao and Liu, Qiang. Centroid approximation for bootstrap. arXiv preprint arxiv:2110.08720. 2021.
Yoon, Jaesik, Kim, Taesup, Dia, Ousmane, Kim, Sungwoong, Bengio, Yoshua, and Ahn, Sungjin. Bayesian model-agnostic meta-learning. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 7343-7353, 2018.
Yu, Huanhuan, Hu, Menglei, and Chen, Songcan. Multi-target unsupervised domain adaptation without exactly shared categories. Advances in neural information processing systems, 2018.
Yu, Kai, Tresp, Volker, and Schwaighofer, Anton. Learning gaussian processes from multiple tasks. In International conference Machine learning, pp. 1012-1019, 2005.
Yue, Zhongqi, Zhang, Hanwang, Sun, Qianru, and Hua, Xian-Sheng. Interventional few-shot learning. Advances in neural information processing systems, 2020.
Zagoruyko, Sergey and Komodakis, Nikos, Wide residual networks. BMVC, 2016.
12

---

Zenke, Friedemann, Poole, Ben, and Ganguli, Surya. Continual learning through synaptic intelligence. In International Conference on Machine Learning, pp. 3987-3995, PMLR, 2017.
Zhang, Saizheng, Dinan, Emily, Urbanek, Jack, Szlam, Arthur, Kiela, Douwe, and Weston, Jason. Personalizing dialogue agents: I have dog, do you have pets too? Association for Computational Linguistics, 2018.
Zhang, Yu and Yeung, Dit- Yan. A convex formulation for learning task relationships in multi-task learning. Association for Uncertainty in Artificial Intelligence, 2012.
Zhou, Jiay,, Chen, Jianhui, and Ye, Jieping. Clustered multi-task learning via alternating structure optimization. Advances in neural information processing systems, 201 1:702, 2011.
Zhou, Qiang and Zhao, Qi. Flexible clustered multi-task learning by learning representative tasks.
IEFE transactions on pattern analysis and machine intelligence, 38(2):266-278, 2015.
13


Reviews:
[{"review_id": "xJzC2-tAeFy", "reviewer": "Reviewer_ZWfd", "summary": "The paper proposes a generic method for approximating the argmax distribution of a random function using optimized centroids without explicitly sampling from the argmax distribution. While the closest application for the method is in meta-learning and multi-task learning, the proposed method can find application as an alternative to bootstrap among other applications. In short, for a random function $f_\\xi(\\theta)$ (think of $\\xi$ as data, $\\theta$ as parameters, and $f_\\xi(\\theta)$ as the loss function), the method finds a set of parameters $\\boldsymbol{\\theta}:=\\\\{ \\theta_i\\\\}_{i=1}^n$ \nsuch that it minimizes \n$E_\\{\\xi\\sim\\pi\\}\\[ min_\\{i\\in\\\\{1,...,n\\\\}\\} f_\\xi(\\theta_i)\\]$. During inference and for a new task, the authors first find the appropriate parameters for the task (based on a validation set) from the set $\\boldsymbol{\\theta}$ and then use the corresponding model to solve the task. The method is shown to minimize a bound in Wasserstein distance between the approximated argmax distribution and the real distribution. Finally, the authors show the method's superior performance in various numerical experiments spanning few-shot supervised learning, personalized dialogue systems, and multi-target domain adaptation compared to various baselines. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**\n\n* The paper provides a unifying and interesting perspective for various multi-task learning applications by studying the argmax distribution of a random function\n* The method can be interpreted as an implicit clustering for tasks, where similar tasks are assigned to the same model parameters\n* By assigning different parameters to different 'clusters of tasks,' the proposed method addresses one of the main shortcomings of meta-learning approaches (like MAML), which assume similarity/homogeneity of training tasks. \n* The authors provide extensive numerical experiments to demonstrate the practicality of the proposed approach.\n\n**Weaknesses** \n\n* argmin and argmax are used interchangeably. While this is not a major issue, picking one and sticking with it throughout the paper will improve consistency.\n* The paper does not report any insights regarding the computational cost of the approach, specifically in comparison with meta-learning approaches like MAML and Reptile (which is more efficient). I think it might be safe to say that for $n=1$ the method is technically a classic meta-learning approach, but does the computation linearly scales in $n$?\n* The model selection required for solving new tasks seems to be constraining. When the argmax distribution is approximated on supervised learning tasks, but we don't have any training/val/test labels for the new tasks, how can one find the right model for the new task?\n\n**Summary** \n\nThe paper provides an interesting perspective on estimating the argmax distribution of a random function, which finds application in many machine learning problems involving learning multiple tasks/datasets. The paper is well written (minus a few typos here and there) and is interesting. The authors show that the proposed approach theoretically minimizes a Wasserstein bound between their approximation and the true distribution of argmax.  Moreover, the method is validated on various multi-task learning problems. While the paper is in good shape, it lacks in certain aspects. For instance, it doesn't give a sense of the computational scalability of the algorithm. Also, the task clustering aspect is exciting. The authors could have provided the tasks that share a model to give a sense to the reader about the implicit modeling of task similarities. I think the paper is interesting to the broad ML community, and despite its mentioned deficiencies, it stands above the acceptance threshold.\n\n**Minor editorial comments**\nPlease proofread the paper and address the typos. Here are some that I caught:\n* line 242: \"in and ensemble\" should be \"in an ensemble\" \n* line 356: \"increase\" should be \"increases\"", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors state that \"Our algorithm slightly increases the time cost and memory cost in the experiments,\" however no rigorous computational analysis is provided in the paper. ", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hrs"}, "extra_text": {}}, {"review_id": "nv4R_pN--AE", "reviewer": "Reviewer_pFrh", "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.\n\n", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1- The paper is well written and it well motivates the application of approximating the argmax distribution in different deep learning tasks.\n\n2- Proposed approach optimizes a set of centroid points that approximates the argmax distribution rather than random sampling, and with a toy example, it shows that their approach yields lower Wasserstein distance than random sampling.\n\n3- The effectiveness of their method is evaluated on multi-task learning applications, including few shot image classification, personalized dialogue systems and multi-target domain adaptation. Their result shows that the proposed algorithm improves the baseline methods in different settings.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "1-  The paper proposed a centroid-like approximation to replace Monte Carlo sampling. At each iteration, the algorithm draws a set of random variables, finds \\theta_i, and updates \\theta_i with gradient descent. It also keep track of frequency which later on used for the estimation of distribution p*. This approach is more sample efficient than the simple Monte Carlo sampling but each step is expensive. My concern is that it may need less samples but eventually takes same time to estimate the p*. To have a better comparison, I suggest not only showing how Wasserstein distance is changing by increasing the number of samples but also the time and memory cost. \n\n2- The proposed algorithm uses the best single model selected from the pool during evaluation. Specifically, It trains a fixed number of classification heads (the argmax centroids) and it picks the best head and further finetune it on the training set for each test task during testing.It may lead to a high computation cost in comparison to existing methods. Adding some results/datapoints should make it clear if this approach is practical.\n\n3- Existing ensemble learning averages the output of multiple models, while this approach picks the best single model. I am wondering if relaxing the best to average would reduce the computation cost!\n\n4- As discussed in the paper, if some \\theta_i is rarely selected, it's important weight gets close to zero. The paper addressed this problem with resampling a set of new \\theta_i with replacement and repeating the updates. But if the total samples are small, this approach never detects those rare samples and may end up with a very biased estimation of the argmax distribution. In experimental results some of these details are missing. If the sample size is very large then the simple Monte Carlo may also work! \n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 hours"}, "extra_text": {}}, {"review_id": "n308jnsb3wQ", "reviewer": "Reviewer_Zrxe", "summary": "This paper proposes a natural and simple (yet general) method for approximating the argmax distribution through optimizing a set of centroid points. It is also shown theoretically that under suitable assumptions such an approximation minimizes a bound based on the Wasserstein distance with the actual argmax distribution. Extensive experimentation is carried to support the usefulness of the method. The data include a toy dataset, a few shot learning task, a dialogue system task and a domain adaptation task.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The central object of interest in this paper is the argmax distribution, which appears frequently in different contexts in machine learning. More specifically, given a random function f_\\xi(\\theta) with \\xi representing a random variable and \\theta the variables of interest, we want to estimate a distribution \\rho^\\ast of the optimal points when \\xi is drawn randomly. This problem is approached in quite a natural manner: By seeking a set of centroid points from \\theta (optimized for explicitly) such that they are close to \\rho^\\ast. Such an approach permits avoiding expensive Monte Carlo sampling based methods. The MC approximations are also poor when the number of sampled points are small. \n\nThe solution to the above problem is quite simple and straightforward. The objective function and method is summarized in equations (2) and (3). We optimize for the \\theta (which are also assigned importance weight so as to approximate \\rho^\\ast. As might be expected, the objective is quite similar to the k-means objective, except that the distortion function is for distributions, and in particular for the argmax distribution.  The whole procedure is summarized in algorithm 1. Further, it is shown theoretically that the method can be seen as minimizing a bound on the Wasserstein distance between the empirical measure of the centroids and the true distribution (rho^\\ast above). Given the usage of importance weights, and a simple distortion based objective, this is simple to see. \n\nThen a set of experiments are presented that show that the above procedure indeed leads to gains. The first experiment is a toy dataset composed of a mixture of gaussians. It is shown that the method yields a smaller Wasserstein distance with the true distribution with few number of samples as compared to a monte carlo approximation. Then is a few-shot task akin to that in the MAML (and following) papers. Using some recently proposed methods as the baseline, the centroid approximation is shown to have consistent performance improvements. Similar results are shown for a standard personalized dialogue systems task and a domain adaptation task. \n\nIn summary the paper proposes a very simple and natural method for approximating the argmax distribution. The theoretical contributions are also straightforward. From a purely technical perspective (in the sense of the tools used, originality of the method etc), the contributions are also straightforward. From my perspective the main contribution of the paper is experimental -- showing that such a simple method will lead to performance gains in a variety of tasks, which I believe is sufficiently corroborated. \n\n\nMinor comments:\n- Line 48: \"theoretically and show its property.\" --> this sentence is a little bit awkward, please reword this. \n- Line 68: \"important weight\" --> importance weight.\n- Line 77: \"nearly zero important weight\" --> same as above", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are discussed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}, {"review_id": "Z9AKWLTZazm", "reviewer": "Reviewer_kFgf", "summary": "\n\n\nThis paper proposes optimizing centroid points to compactly approximate the argmax distribution with a simple objective function. The proposed method is theoretically minimizing a bound of Wasserstein distance between the empirical distribution of the centroids and the ground-truth distribution. Argmax centroids have many applications for machine learning tasks; the author validates the proposed method on multiple meta-learning and mult-task learning tasks.\n\n", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "My main issue is with the empirical evaluation. I am not an expert in centroid approximation, so I leave that to other reviewers to judge the originalities and soundness of the proposed method.\n\n- For the personalized dialog experiments, the results are not consistent with the original paper. Which metric is case-insensitive? Only BLEU (as described in line 279) or both (described in the title of table 4)? \n\n- Since multi-task learning is a natural application of argmax centroid, it will be interesting to extend the experiments with multi-task learning on SuperGLUE using pre-traineds model to test the adaptability of the proposed framework. \n\n- Some aspects of the paper's presentation can be improved. E.g., Line 343, incorrect ref. \"discriminator loss in (??)\", confusing references on table 2 and table 3 in section 4.2.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "-", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12"}, "extra_text": {}}, {"review_id": "YHAhMA7wceH", "reviewer": "Reviewer_a626", "summary": "A method is presented in the paper to approximate the distribution of argmax of functions depending on additional random variables. The idea is to use a set to represent centroids to approximate the argmax distribution, and in training iterations and test phase assign each realization one member of the set. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "+:\n\nThe proposed method is interesting, and a theoretical interpretation is provided in the paper. The experiments on setups like few-shot image classification and multi-target domain adaptation show that integrating the proposed procedure with existing methods can improve their performance. Additionally, an ablation study is done comparing the method with some ensemble methods showing reasonable improvement.\n\n\n-:\n\nThe first part of the paper leads the reader to expect that with higher n better representation of the argmax distribution and better results can be achieved, while the results are not consistent in that regard. In my opinion, the authors need to further discuss this in the paper.\n\nThe experiments need to be strengthened by an ablation study to show the effect of selecting a model at test phase vs. the training procedure to form the models, for example, by picking best model for other ensemble approaches.\n\nWhile the current experiments are valuable, in my view it is a bit unfortunate that the focus is only on multi-task/domain learning. The motivation and initial discussions about the method are very general, for example, bootstrap and Thompson sampling are mentioned several times to motivate the approach, but no experiment is performed on them. The authors suggest the method can be applicable to a single task setup, so an evaluation of it can be added. \n\nThe experimental details do not seem enough to easily reproduce the results and some details like training/run time and compute resources are not provided. The authors can provide more details in the supplementary.\n\n-------------------------------\nAfter the rebuttal, I still think the ablation study and additional experiments will be helpful in supporting the claims. After looking through the paper again and assuming the authors will incorporate reviewers’ suggestions, I think the paper can be useful to the community, especially its way of considering the meta-learning and multi-task settings. Hence, I have increased the score.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.25"}, "extra_text": {}}, {"review_id": "Bj8qDcrBHBq", "reviewer": "Reviewer_TLpj", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The formulation of the problem (learning the distribution of a random function's minimizer) encompasses a wide range of applications and should be of broad interest. The paper is overall clear and well-written.\n\nThe experiments were thoughtfully chosen to illustrate the appeal of the method. It can be applied in a straightforward manner to various different choices of model and random functions. The authors' choice to view meta-learning and multi-target domain adaptation from a random function perspective was insightful. Unlike the dialogue personalization and domain adaptation experiments, the few-shot learning results did not show a large win for Argmax Centroids. However, the fact that Argmax Centroids is versatile enough to be applied to various models is a point in its favor.\n\nDespite the aforementioned strengths, there are some weaknesses of this paper. Firstly, the related work does not connect Argmax Centroids well to prior work. In particular, the connection to Perturb-and-MAP models (which also consider optimization of random functions) and particle filters (which also perform inference with parameter samples) could be expanded upon.\n\nI also have concerns about some details in §2. The first is the upper bound in the proof of Theorem 2.2 (2nd line of the proof). For the lower bound, the min is moved outside of $|| \\theta_i - \\theta_\\xi ||^{p_1}$, which makes sense. But moving the min outside in the same way for $|| \\theta_i - \\theta_\\xi ||^{p_2}$ possibly makes the bound tighter and thus does not follow from Assumption 2.1. Similarly, in the 3rd line of the proof, the order of the expectation and the min is changed. In general, the expectation of a minimum will not be the same as the minimum of an expectation. Finally, is there a citation for the expression for the optimal $\\nu$ in the fourth line of the proof? As an aside, the definition of Wasserstein should probably be stated somewhere in §2. \n\nll. 256-258: The need for selecting the best model on the training set does not make much sense to me. Isn't the $\\theta_i$ selected based on the support set for the current episode?\n\nOverall, the aim of the paper is well-motivated and the method is simple enough to implement yet has the potential for broad impact. However, my concerns about related work and the proof of Theorem 2.2 lead me to not recommend acceptance in the paper's current state. I hope these points will be addressed during the discussion phase.\n\n=====\nAfter reading the authors' comments, my concerns about Theorem 2.2 have been addressed, assuming the authors include a derivation of lines 3-4 of the proof in the appendix.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors do not touch upon negative societal impact. However, the method is fairly general and therefore it seems quite difficult to predict what negative impacts there will be.", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper is well motivated by the problem of learning the distribution of a random function's minimizer)  which has a wide range of applications. The paper made a solid contribution to this important problem by proposing an interesting and novel method (argmax centroids) with theoretical guarantees and experimental validations. During the rebuttal, concerns regarding the proofs of theorem 2.2 were addressed.  
The final version should address the following points as well as other suggestions made by the reviewers: 

1. correcting these typos and highlighting the contribution in the introduction part.
2. Discussion of the limitation of choosing the best model using a labeled set of test data 
3. providing detailed proofs of the folklore result on K-means and Wasserstein
Decision: Accept (Poster)