Functional Ensemble Distillation
Authors: Coby Penso
Idan Achituve
Ethan Fetaya
Venue: NeurIPS 2022 Accept (2022)
Keywords: Bayesian inference
distillation

Abstract: Bayesian models have many desirable properties, most notable is their ability to generalize from limited data and to properly estimate the uncertainty in their predictions. However, these benefits come at a steep computational cost as Bayesian inference, in most cases, is computationally intractable. One popular approach to alleviate this problem is using a Monte-Carlo estimation with an ensemble of models sampled from the posterior. However, this approach still comes at a significant computational cost, as one needs to store and run multiple models at test time. In this work, we investigate how to best distill an ensemble's predictions using an efficient model. First, we argue that current approaches are limited as they are constrained to classification and the Dirichlet distribution. Second, in many limited data settings, all ensemble members achieve nearly zero training loss, namely, they produce near-identical predictions on the training set which results in sub-optimal distilled models. To address both problems, we propose a novel and general distillation approach, named Functional Ensemble Distillation (FED), and we investigate how to best distill an ensemble in this setting. We find that learning the distilled model via a simple augmentation scheme in the form of mixup  augmentation significantly boosts the performance. We evaluated our method on several tasks and showed that it achieves superior results in both accuracy and uncertainty estimation compared to current approaches.

Functional Ensemble Distillation
Coby Penso Bar-Ilan University, Israel coby- peny-pppsss22lllll com
Idan Achituve Bar-Ilan University, Israel idan. achi tuve@biu.ac.... ac.
Ethan Fetaya Bar-Ilan University, Israel ethan. fwt..feetaabbu..cmm ac. ac
Abstract
Bayesian models have many desirable properties, most notable is their ability to generalize from limited data and to properly estimate the uncertainty in their predictions. However, these benefits come steep computational cost Bayesian inference, in most cases, is computationally intractable. One popular approach alleviate this problem is using Monte-Carde estimation with ensemble of models sampled from the posterior. However, this approach still comes at a significant computational cost, one needs to test time. In this work, we investigate how to best distill an ensemble's 's predictions using an efficient model. First, we argue that current approaches are limited as they are constrained to classification and the Dirichlet distribution. Second, in many limited data settings, all ensemble members achieve nearly zero training loss, namely, they produce near-identical predictions on the training set which results in sub-optimal distilled models. To address both problems, we propose a novel and general distillation approach, named Functional Ensemble Distillation (FED), and we investigate how to best distill an ensemble in this setting. We find that learning the distilled model via a simple augmentation scheme in the form mixup augmentation 43 significantly boosts the performance. We evaluated our method on several tasks and showed that it achieves superior results in both accuracy and uncertainty estimation compared to current approaches.
and run multiple models
Introduction
In parametric Bayesian inference the output is a distribution over model parameters, known as the posterior distribution. This is in contrast to standard approaches which return single model. At test time one needs to marginalize over all models according to the posterior distribution to make prediction. The action averaging over multiple models allows Bayesian approaches to perform well with limited data [ 27 34 42 and properly estimate uncertainties [17 18, 34). The downside that in most cases Bayesian inference is intractable and even methods for approximate inference may be costly. This is especially true for Bayesian neural networks (BNNs), as the dimension of the parameter space for common neural networks is very large
15).
Two main approaches exist for efficient Bayesian inference. The first approximates the posterior distribution (c.g., variational inference and the Laplace approximation) with simple distribution, usually a Gaussian with diagonal covariance. While this can be done efficiently in some cases [10] 20), it is more prone to over-fitting as the approximate distribution might cover only small part of the posterior. A second approach is Monte-Carlo (MC) estimation where we sample M models from the posterior using an approximate sampling algorithm, e.g. Markov Chain Monte- Carlo (MCMC), and estimate the mean prediction by the average over the predictions of the M models. While this can better capture the diversity of the posterior, requires storing and running M models instead of one. One important question that will be the focus of this work is how to maintain the benefits of the Bayesian approach while reducing the computational overhead. is also important to note that an ensemble of M models is commonly used outside of Bayesian settings, e.g. bagging (3), so this line
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

of work has wider implications. Our focus, however, will be in the low-data regime where Bayesian methods are most impactful.
We propose an ensemble distillation 16 26 method that mimics an ensemble of models using a lightweight model. Current approaches mainly consider classification where single prediction is a probability vector, the distilled model then outputs parameters of a Dirichlet distribution that should fit the ensemble predictions. Our approach, which we call Functional Ensemble Distillation (FED), on the other hand tries to learn distribution over functions by transferring the randomness to 6, a noise variable that is sampled from a fixed distribution. We designed our model with two requirements in mind. First, it should be fast to run during test time. This is achieved by being able to run with multiple e values in parallel. Second, the distribution of predictions from our model should match the distribution of predictions from the posterior. This approach allows us more flexibility as it not limited to classification tasks and to the Dirichlet distribution, unlike most previous approaches. Furthermore, our method allows us to compute covariance between predictions which can be useful for further downstream tasks. Consider, for example, the task of object detection for autonomous vehicles. If two nearby detections are made with certain confidence then the probability that that area of space is occupied depends heavily on the correlation between these predictions.
Finally, we investigate how to properly distill from a model ensemble under severe overfitting. It is common for deep networks trained on small datasets such as CIFAR-10 [22] reach almost perfect accuracy on the training set while still generalizing well to the test set. In these cases, the predictions of the models in the ensemble can vary considerably on the test set despite being almost identical on the training set. This means that distilling the ensemble on the training set will not capture the diversity between predictions which the core component of ensemble methods. While some previous works tried to address this issue (26 (9], they either depended on additional unsupervised data that might not be readily available, e.g. medical imaging, or do not scale well. Here we propose simple solution for this problem that works remarkably well creating the mixup augmentation [3]. The mixup procedure generates data that is on the one hand distinct enough from the training set to allow our ensemble to produce diverse predictions, while on the other
auxiliary dataset using
hand, being close enough to so that models that were trained on it work well when applied to the original test data.
To conclude, we make the following novel contributions: (i) We propose ceneral-purposs and flexible method for ensemble distillation that captures well the predictive abilities and the diversity of the original ensemble; (ii) we propose a simple, fast, and scalable method for obtaining a diverse auxiliary dataset for training the distilled model without any new samples using the mixup augmentation; and (iii) we evaluated our method on several image benchmarks and showed that it obtains the best results in terms of accuracy and calibration compared to other ensemble distillation approaches.
2 Background
We first provide a brief introduction to the main components of our model. We use the following notations in the paper. We denote scalars with lower-case letters (e.g., x), and vectors with bold lower-case letters (e.g., x). We assume we are given with dataset D of size N, D = {xx, MiNN where x; E Rdd are the training examples and yi are the corresponding labels or targets.
2.1 Bayesian Inference
In Bayesian inference we are mainly interested in the posterior distribution p(0 x II. p(y [x, p())) where 9 denotes the parameters of the model (c.g., a NN). The posterior tells us how likely 0 is given the observed data. When given new data point x, we will use the entire posterior to calculate the predictive distribution using the rule p(y-|x.. D) = p(y- x., (0 d0. One intuition behind this is that when one has limited data, there are many distinct models that work well on the training set. By averaging over all the possible predictors the final prediction is supported by most of the "weight" in the posterior instead of by single model. This also leads to better uncertainty estimation, as when many different models have different predictions the resulting average prediction is dispersed and uncertain.
In practice, besides a few notable exceptions, computing this integral is computationally intractable. One popular approximate inference solution is Monte-Carlo (MC) estimation. Namely, we estimate
2

---

Ensemble
FED
81
p()(k:81)
q(ylc)1)
92
Ply)):2)
sample Ey5o fq......
q(ylcc2))
Geverator;
W
9m
plyix
q(yl(:
Figure I: (1) Ensemble Forward pass pipeline. (2) FED Forward pass pipeline.
X.; fm). where 6m are (approximate) samples from the posterior p(O Obtaining such samples can be achieved in various ways such as MCMC methods 30 38 44). One problem with this approach that it requires storing and running multiple models to obtain predictive distribution. This may be impractical in many situations such as when making predictions in low resource environments or in real-time systems.
p(y. D) > トΣ p(y.
2.2 Maximum Mean Discrepancy
A big question in machine learning and statistics is how measure distances between distributions, specifically without access to the density function and with only access to samples. One well known measure is the Maximum Mean Discrepancy (MMD) (37 that was used as objective successfully train generative models [23] [4]. Given a family of functions F the MMD score between [fxx][((x)]
two distributions p, q defined as . M ((p q) se [[xx
specific
case in which the MMD has a closed form formula is when F is the unit sphere in a reproducing
kernel Hilbert space (RKHS) defined by kernel function [13. In that case we have
MMD 2)²
Ex.x'>p [x(xx x') 2E 2Exxpxy~ (x(, -yy.' (1(y"")(()))')))
(1)
which we estimate by Monte-Carlo means.
2.3 Mixup Training Procedure
The mixup
4 training procedure is based on the Vicinal Risk Minimization principle 6 The main idea is to enlarge the training set by sampling points in the neighborhood of the given data. (45 proposed simple way to achieve this goal by sampling along the line which connects two data points from the training set. Let (x,; y4). (x, U3) be two training examples, one can generate a new data point (&, 7) by taking their convex combination: X Ax; + (1 A)x;, and Ays + (1 A)y; where À drawn from a Beta distribution with some fixed parameters. Now, standard loss functions, such the cross entropy loss, can be used on (x, W).
3 Related Work
Several previous studies proposed approaches to distill an ensemble of models into a single fast model. Early approaches [6] 21 32 40] focused on capturing only the mean behaviour of the ensemble but did not capture the diversity in it, which may be crucial for uncertainty estimation. For example, in [5] it was first shown that a single network can learn the knowledge of a large ensemble of networks. In [16] a distilled model is learned using convex combination of the ensemble predictions in addition to the standard cross-entropy loss. And, in (21] a student network learned online from a teacher that is trained via MCMC training procedures.
More recent studies proposed alternatives for capturing the diversity as well. Ensemble Distribution Distillation (EnDD) (26) was the first to propose such approach. EnDD builds on the method in 25 which tries to emulate the behaviour of an ensemble with a Dirichlet distribution parameterized by NN. To handle the fact that the ensemble predictions over-confident and have low diversity

---

3x5 conn E4
Sxd carn. 128
3xd came 286
JM0 WWL.OO X02
44 AiD
3i0 512
Softmax
FO
Figure 2: Generator architecture. First, we concatenate 3 random Gaussian channels the input. Second, add Gaussian noise to intermediate features along the ResNet-18 architecture.
on the training set, they proposed to use an unsupervised auxiliary dataset for training the distilled model, and perform temperature annealing. Similarly, [9] also makes the assumption that the predictive distribution follows a Dirichlet distribution. They fit the output to the teacher network by two separate networks, one that learns the predicted class probabilities, i.e. the Dirichlet mean, and another that learns the confidence, i.e. the Dirichlet concentration. Both [26] 9 are suitable for classification tasks only and assume Dirichlet distribution behaviour which may be too restrictive.
We argue that more general solution for the distillation of ensembles is needed.
different line of work [36 [29] generalizes the method in [6 by distilling the ensemble members to task heads in a one-to-one fashion. In [36] the distilled model is partitioned into two parts, a shared
"body" and member-specific "head", one per ensemble member. Each head is trained to copy the prediction of an ensemble member using the standard cross-entropy loss functioo. In [29] the authors propose use BatchEnsemble [39] for weight sharing. They increase the diversity in the ensemble predictions on the training set by randomly perturbing the training data based on three techniques: Gaussian, Output diversified sampling (ODS) [35,, and ConfODS. The ODS perturbations are similar to adversarial perturbations and need to be computed at every iteration. This imposes a computational overhead which renders a slow training process and inability to easily handle large ensembles, unlike our proposed method.
4 Our Method
We take similar approach 19 26 where treat the ensemble distillation conditional generative model. Given an input x and a sampled model parameters we get a distribution p(y|x, 0), so for given
can porw-f the posterior distribution over weights distribution over distributions over y. While this was done in 99 26 there are several limitations with these approaches. First, they use Dirichlet distribution model this distribution over distributions, which is limited to classification. Second, even in classification settings, the family of Dirichlet distributions has fixed expressive power, and there is no evidence that the posterior indeed Dirichlet distribution.
Our approach is to model distribution over functions that both have low memory requirements and that we can parallelize to get multiple predictions quickly. We do this by training a generator neural network with parameters w and added random noise c. The noise can be added, for example, as an additional channels to the input or hidden layers. Our goal is to train the generator network so that the distribution over /3//,, w, c) for ~ p(c) will match that of p(y x, 8) for 0 ~ p(0[D). Namely, each draw of corresponds cost to standard single network and we can parallelize easily as all the linear operations are shared. Computationally it is equivalent running single model a batch instead of single input. An illustration of our method is presented in Figure | We further show in the appendix that for wide range of batch sizes the runtime of our model scales approximately as (# of samples)/ resulting in a considerable speedup. Since we now have a distribution over functions, we call our method Functional Ensemble Distillation (FED).
draw from the posterior. By using this design we get similar memory
We now present our training objective for the generator Iw- Given an ensemble member 0, and batch of B examples, we consider the vectors of predicted probabilities (one per example) as features representing the function parameterized by 0z- Denote by p p(ys [x, 0,) the predictive likelihood for the tt example by the jth function. We represent the function parameterized by 0, by concatenating these vectors to form feature vector p coneat(p' P(g)- We repeat this process for our generator network (only that now each e; represents function) and in similar fashion,

---

Algorithm Generator training
1: input: Predictions of the ensemble on our distillation dataset: {x, { f
# number of data points, M ensemble predictions
2: Initialize Generator Iw
3: while not converged do
4: Sample batch of inputs and model predictions {xi; (pf - Jimmn
5: Sample M latent variables ~ (0)
6: of yw 9w(XX; (Xi; Ej)
7:
q concat (q)
Ip)
8:
P = comcat (p'....
9:
C MMD ((f'))=I:M. (â )=
10: w < w WwL
11: end while
12: return 9w
obtain a representation q concal(qi..... (n) using the same batch. The same process is applied every member of the ensemble. Now we the MMD objective presented in Section 2.2 to score the difference between distribution over functions, using simple kernel over these functions' representations. More formally, the MMD objective given batch of examples is:
MMD ((p)))=1:M, (=))=)))==
M Σ (k(p'.
A'''' .P) -2((', (2)
We present our method in Algorithm 1 In our experiments, we used standard convolutional neural networks where we concatenated 3 random Gaussian channels to the input channels and we add Gaussian noise to intermediate features across the network. We evaluated several options for inserting noise into the generative model and noticed that to learn sufficient diversity, enough noise has to be inserted; otherwise, the generator's diversity suffers and collapses nearly deterministic model. We also found best to set the noise STD to be a learnable parameter (one parameter per input noise layer). See illustration with ResNet18 in Figure Regarding the limitations of our approach, we note that while the run- time is much faster then running the full ensemble, there is still non-negligible overhead compared to single model.
4.1 Auxiliary Dataset
One problem with distilling ensemble models is that the predictions of the ensemble members the training set have a much higher agreement compared to the test set (as can be seen in Table [ Practically this means that distillation of the ensemble predictions on the training set can lead to degenerate solution that returns most cases the same predictions. This leads to worse performance in both accuracy and calibration compared to the original ensemble.
In 26 the authors address this by using an auxiliary unsupervised dataset, we can get diverse ensemble predictions on this dataset and use it to train our distilled model. While this solution performs quite well, in several applications having access to an additional unsupervised dataset not trivial. One example is medical imaging [8, where even unsupervised data requires access to expensive machinery and data that is not readily shared due to privacy concerns. Another example learning from ancient texts [12] where digitized data, labeled or unlabeled, is scarce. We argue that these problems in which the data is limited are exactly the problems where Bayesian methods make the most impact. Therefore, having method that works these types of data is of great importance.
This leads us to seek solution that does not require any additional data, even if it is unsupervised. We investigated several solutions (see Section 5.4) and found one simple solution that worked surprisingly well, the mixup augmentation. The predictions distill from data that more closely resembles the test data. Another important property of the mixup augmentation is that models trained on it transfer well to real images. This is in contrast other
the mixup data are more diverse allowing to
5

---

alternatives that we explored, such as generating an auxiliary dataset using generative adversarial network (GAN), that did not transfer well.
To summarize we generate a synthetic, unlabeled, dataset from the original training set (that was used to train the original ensemble) using the mixup augmentation. We then train our generator network using the predictions of the original ensemble on this dataset, see Algorithm|] [ In Table we show that indeed the ensemble classifier generates more diverse set of predictions on the mixup dataset compared to the actual training set.
Experiments
We evaluated FED against baseline methods on different datasets and using different metrics. We present here the results of the models using the hyperparameters with the best accuracy on the validation set. As a result, for some models the diversity can be improved at the expense of lower accuracy. We will also provide our code for reproducibility /https://githhut com/ cobypenso/ functional, ensemble distillation
Datasets. We evaluated all methods on CIFAR-10, CIFAR-100 [22], and STL-10 datasets. STL-10 and CIFAR- 100 are considered more challenging, the former has relatively small amount of labeled training data (10 classes with 500 labeled examples per class), and the latter has more categories. For all datasets
train/val/test split. The train/val split with ratio 80%:20%.
Compared methods. We compare our method against the following baselines: (1) EnDD (26). Ensemble Distribution Distillation framework, trained with the same train data used to capture the ensemble. (2) EnDD x (26), Ensemble Distribution Distillation framework, trained with the train data and an auxiliary dataset. For CIFAR-10, 10, CIFAR- is the auxiliary dataset. For CIFAR- 100, CIFAR-10 10 is the auxiliary dataset. For STL-10 we used the unlabeled data provided as part of this dataset. We stress that other methods do not have access to this data.
(3) DM [29]. this baseline comes in three variations, Gaussian, ODS, and ConfODS. (4) AMT (9) method which also uses MMD as an objective. (5) Hydra (36), weight sharing method based on [16]. (6) Knowledge Distillation (KD), a naive single model baseline, trained on the mean predictions on the mixup data.
Ensemble. We created the Bayesian ensemble (for distillation) by sampling from a posterior obtained with the Cyclic Stochastic Gradient Hamiltonian Monte Carlo (cSGHMC) method (4). Overall we sampled 120 models. In addition, we show experimented with different ensemble methods and their implications on the distillation. The architecture we choose for the models in the ensemble is ResNet18 with Group Norm 41 and Weight Standardization (33] We did not use Batch normalization since breaks the IID assumption commonly used in Bayesian Inference [7].
We note that sampling from the posterior with eSGHMC, instead of SGD, slightly improves the ensemble diversity on the training set. If our ensemble was sampled using SGD with different seeds the effect of mixup would be even more pronounced.
Training setup. For the baselines, we follow the training strategies suggested in the papers and their code if available, while performing an extra hyperparameter search when training on datasets which did not cover in the original papers, mainly CIFAR- and STL-10. For our method, performed a standard training procedure with Adam optimizer [19], a learning rate scheduler with fixed milestones at epochs (35, 45, 55, 70, 80), and hyperparameter search done over a held-out validation set.
Objective and Kernel choice. We trained our model using the MMD objective while considering various kernels. Specifically, we examined Linear kernel, RBF kernels, and a mixture of RBF kernels. For the RBF kernel, considered different values of length scales, a considerable effect on the results. For CIFAR-10, 10. a mixture of RBF kernels with {2, 10, 20, 50} length scales had the best results. For CIFAR-100, the length scales are {10, 15, 20, 50), and for STL-10 a length scale of 50 works best.
saw that it had
Distilled Model. We distilled using similar ResNet18 with Group Norm and Weight Standardization architecture as the ensemble models. Thus, essentially trying to distill a 120 models ensemble into a single model that has almost the same size as a member in the ensemble, i.e ~ 100x memory improvement. Memory and Timing performance analysis can be found in the Appendix. For the noise sampled and fed into the generative model, we chose simple Gaussian distribution, fed by

---

1.0-
1.0-
Thue Positive Itate
True Positive Rate
True Positive Rale
0.8
0.8
0.8
0.6
0.4
0.4
0.4
0.2
0.2
0.2-
Hydra
0.0
0.2
0.4
0.6
0.6
1.5
0.0 0.2
0.4
0.6
0.8 1.0
0.0 0.2
0.4
0.6 o. 1.0
Positive
(a) STI-10
(b) CIFAR-10
(c) CIFAR-100
Figure 3: Receiver operating characteristic (ROC) curve measuring the performance on in-distribution Vs dott-o--ditributtion data based the mutual information. Best viewed in color.
concatenating it to the input, and by inserting it in intermediate places along the model. Specifically, for the concatenation part, 3 channels of noise and 3 channels of the input are stacked together. For the intermediate noise, the Gaussian noise was added to the features, instead of concatenation, in different places, one after the first convolution layer and the other four after each Block in the ResNet-18 88 architecture (Figure 2). We found that setting the standard deviation of the noise to be learnable parameter as part of optimization, has a slight performance improvement.
Metrics. To evaluate our method, and other baselines, we define the agreement metric as the probability that two distinct elements of our ensemble predict the same label on the same input. In addition, we computed accuracy, and Expected Calibration Error (ECE) [14] reported in percentages. For the agreement metrie, close value to the original ensemble's agreement is good indication that we capture similar behavior.
5.1 Ensemble Analysis
Table I: Ensemble analysis. Test accuracy, agreement & ECE (%) on STL-10, and CIFAR-10/100,
STL-10
CIFAR-10
CIFAR-100
ACC (?)
AGREEMENT
HCE (1)
ACC(())
AGREEMENT
(1)
ACC
AGREEMENT
(1)
Train
100.0
98.23
2.99
99.97
91.85
6.38
99.95
97.53
4.15
94.41
77.79
71.97
mixup
95.20
88.08
88. 56
We first analyze the ensemble we sampled using eSGHMC. To evaluate the diversity between predictions we use the agreement metric. We computed accuracy, agreement and Expected Calibration Error (ECE) [4, 28 over the trainset and testset. We only show agreement on our suggested mixup auxiliary dataset, it does not have discrete labels. We can that the ensemble has almost perfect accuracy and very high agreement on the training set, with a considerable drop in both on the testset. In fact, for CIFAR-100 and STL-10,, the ensemble has essentially diversity over the training set, thus making the trainset inappropriate for the distillation task. The mixup agreement closes part of the gap between the train and test datasets.
5.2 Distillation
We now evaluate FED on learning multi-class classification tasks. While we are not limited to classification tasks, In the following experiment compare approach to the baselines described above. Results in Table that FED outperforms previous methods that do not use an auxiliary dataset by a large margin, especially on the more challenging STL-10 and CIFAR- datasets. This true for both accuracy and ECE. is interesting to note that while EnDD uses additional unlabeled examples from the training distribution (between 50k and 100k), FED is on par and often outperforms this baseline across all metrics, and the only one who outperforms the naive single model distillation baseline (KD)
out of three datasets.
We also show results for distriibtntt (OOD) detection in FigureBJusing Juiin ROC curves (EnDD was excluded due to poor performance), To score the examples used the knowledge uncertainty
7

---

Table 2: Test accuracy, agreement & ECE (%) (# STD) on STL-10 and CIFAR-10/100. Note that the AUX baseline uses additional unsupervised data. Each experiment was repeated 3 times.
STL-10
CIFAR-I0
CIFAR-100
ACC
AGREEMENT
BCR())
ACC (f)
AGREEMENT
(CE
ACC
AGREEMENT
BCE
Ensemble
81.84
$7,85
2.17
94.41
86.08
5.56
71.79
71,97
3.07
KD
80.49:.12
1.1
91,104.21
3.534..22
73,45++00
3.834.64
EaDDarx
80.78:000
99.98
1.20::000
91.05=+1.3
100.0::000
5.63++1.0
70.68::000
100.0:::00 00
21:40:000
EnDD
73.45=1.1
99,881.02
10.62:+1.2
91.011.70
99.99:200
6.391.81
71.16::00
100.01:00 00
14.51=.00
74.21=.13
86.63+3344
15.174.58
90.27+:15
88:
2.48+:944
72.44t.19
85 85.97+-81
4.63=1.0
DMoDs
73.902.65
86.72444.3
15.324..8
89.35.2277
95.51=11.7
2.65::1.2
71.85 36
87. 87.90:1-4.1
4,47=10
73.26::2.2
86.67442.7
14.52443.1
89.55-4.233
95.49:
2.594.29
72.33+ 54
88.73444.7
5.33=111
AMT
54.36=.78
11,31+:000 08
38,97:t
83.81-2566
12.95-2008
64.00+
NIA+N/A
NA+NNA
N/A=N/A
Hydra
16
93.40:2.26
30.68:33.3
91.071188
94.60:2.22
5,3911.6
69.56:2.11
93.80:2.46
29.84210
(Ours)
.36
97.18
3.41 .94 93.61
93.34 +:
+:
87.824 .60
..17
which is obtained by subtracting the aleatoric uncertainty from the total uncertainty 11 26 We used SVHN [3] as OOD data. See appendix for more details. From the figure, FED generally outperforms all baselines here as well.
5.3 Mixup Augmentation as Auxiliary Dataset
Table 3: Ablation EnDD, KD and Fed with and without mixup. Test accuracy, agreement & ECE
(%) STL-10, CIFAR- 10, and CIFAR-100.
STL-10
CIFAR-I0
CIFAR-100
ACC (t) AGREEMENT ECE (L) (CC AGREEMENT ECE (J)
ACC
AGREEMENT ECE (L)
EnDD
72.23
99.89
90.35
14.51 15.86
FnDD mixup
78.40
99.99
7.68
90.83
100.0
5.74
68.78
100.0
KD
16.13
90.78
9.90
3,83
KD mixup
80.49
4.15
91.10
3.53
73.96
13.72
65.
17
5.00
92.58
99.16
68.
97
9.61
FED mixup
80.94
97.16
2.55
93.68
93.71
0.78
74.56
88.20
2.07
As an ablation study, we would like understand the separate effects of the mixup augmentation and our FED approach. Thus, we show results for both FED, KD and EnDD trained on the training set and on the mixup dataset as the auxiliary dataset. Results are detailed in Table We show that for our method the mixup adds significant boost to performance, and also for EnDD and KD on STL10. We also see that even when EnDD is trained with mixup, our approach returns better accuracy. ECE and also has greater diversity. We suspect that this is due to the limited expressive power of the Dirichlet distribution. Additional experiments in the Appendix show that at least part of the lack of diversity can be attributed to the Dirichlet distribution.
5.4 Alternatives for Learning Distilled Model Without Auxiliary Dataset
In this section, we explore an additional attempt to distill the ensemble, this time without access to additional data. This approach was inspired by the Out-Of-Baa (OOB) estimate in bagging. In bagging each model is trained only on a random subset of the training set and the OOB estimate is done by averaging the predictions of all models on those data points that they were not trained on and returning their score (which we refer to as held-out datasets). Using this approach, we can train each member in the ensemble on part of the dataset, and use for distillation only the ensemble members' predictions on the data points that they were not trained on. As an example, consider a single datapoint which is on the held-out dataset for model mi, but is part of the trainset for model m2. To train the distillation model, the ensemble uses only the predictions of model m., but does not use those of model my get the overall output for x. We experimented with two different partitions for held-out datasets: K-fold and bagging.
In K-Fold we first split the data into K = 10 folds, DI...., D10. For each of the 10 folds we define We then sample 12 models using eSGHMC on each D=;
D=j = concal ((D:) dataset. As before we get 120 models and for each fold D; we can use the held-out models, i.e. models trained on D=j, for distillation. Note that the dataset for distillation the entire trainset, but for each data-point, we use only the predictions of 12 models that did not train on it.
8

---

Table 4: Comparison of ensemble techniques STL-10, CIFAR-10, 10, and CIFAR-100.
STL-10
CIHAR-10
CIFAR-100
ACC (t) AGREEMENT ECE (J) ACC(()
AGREEMENT ECE (J)
ACC (t)
AGREEMENT BCE (U)
Ensemble
Train
Test
81.85
87.85
2.17
94.41
86.08
5.56
77.79
71.97
3.07
mixup
95,20
88.08
88.56
10 Fold
Train 96.95
89.17
12:45
99.99
97.78
1.77
99.97
94.08
4,77
Held-out
76.67
91.36
1.42
93.72
93.42
1.21
74.52
78.95
3.67
Test
mixup
85.49
94.08
86.61
10 Bugging
Train
97.32
82.05
14.71
99.70
94.24
3.74
99.48
79.79
14.42
Held-out
76.90
79.03
2.39
91.43
91.43
0.94
73.27
69.32
3.27
Test 78.85
74,04
5.63
93.99
90.03
1.70
74.78
65.55
5.21
mixup
120 Bagging
Train 96.33
77.84
18.85
99.99
94.63
3.59
99.98
79.97
14.08
Held-out
77.79
71.85
8.59
93.85
0.81
74.68
3.77
Test 77.77
70.49
8.57
94.13
91.09
1.12
75.16
67.66
4.02
mixup
73.86
91.52
74.19
Table 5: Comparing FED + mixup and Out-Of-Bag ensembles on a held-out data. Test accuracy, agreement and ECE (%) on STL-10, CIFAR-10, and CIFAR-100.
STL-10
CIFAR-10
CIFAR-100
ACC (1)
AGREEMENT ECE (1) ACC (f) AGREEMENT ECE (1) ACC (1)
AGREEMENT ECE (1)
Bayesian.utzup
80.94
2.55
93.68
93.71
0.78
74.56
2,07
10-Fold
62.92
84.40
5.33
91.00
90.48
0.87
67.73
80.39
1.62
10- ood
63.13
84.41
6.16
90.97
88.24
0.97
67.55
67.27
4.82
120- wagggingn-stt-oom
63.03
82.55
7.96
91 03
88.76
1.02
68.95
69.08
5.64
In bagging each bag samples data points with repetition from a dataset of size N. We explore two alternatives, is sampling 120 bags and sampling for each bag single model using cSGHMC (named 120 Bagging), and the other is to sample 10 bags and from each to sample 12 models using cSGHMC (named 10 Bagging). The held-out dataset for distillation is constructed in the same fashion as described in the K-Fold alternative. The key difference is that for each data point, there is a different number of models that did not use it in training, due to the stochastic partitioning of bags.
Comparing these two approaches the K-fold has an advantage since each ensemble member was trained on exactly 90% of the training data while each ensemble member in bagging is trained on ~ 66%. Bagging (with 120 bags) however has better diversity as when we distill we see a larger variety of ensemble model combinations.
Table shows results for the Bayesian ensemble used throughout this paper and the method described in this section. As one can see the scores on the held-out datasets, is good proxy for the test in terms of calibration, agreement, and accuracy (the lower score due to only averaging small number of ensemble models), but this comes in many cases a price of reduced accuracy of the ensemble on the test set compared to the Bayesian ensemble.
Next, we distill the new ensembles using FED. We show the results of the distilled models using the held-out data, and compare them to our Bayesian ensemble distilled using the mixup dataset. As one can see from our results in Table the Bayesian ensemble distilled using the mixup data is superior. We note that on all ensembles the best distillation results we achieved were by using only mixup, further proving the value of this augmentation scheme. We suspect that distillation works best on the Bayesian ensemble as we have the most diverse examples and labels distill from.
6 Conclusions
In this work, we present FED, novel and general ensemble distillation approach. We show that this approach can distill ensembles while keeping desirable qualities such as accuracy, calibration, and diversity. This is done by maintaining the Bayesian approach of distribution over functions while mimicking the posterior with one that can evaluated efficiently. Additionally, we show how the need for an auxiliary distillation dataset can be easily solved by using mix-up augmentation. Finally, our experiments show that FED achieves state-of-the-an results on various benchmarks in the domain of Ensemble distribution distillation.

---

References
[1] Idan Achituve, Aviv Navon, Yochai Yemini, Gal Chechik, and Ethan Fetaya. GP-Tree: A Gaussian process classifier for few-shot incremental learning. In International Conference Machine Learning, pages 54 63.65. PMLR, 2021.
[2] Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya, Personalized fed- erated learning with Gaussian processes. Neural Information Processing Systems (NeurIPS), 2021.
[3] Leo Breiman. Random forests. Mach. Learn., 2001.
[4] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[5] Cristian Buciluă, Rich Caruana, and Alexandru Niculescu-Mizil. Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 335-541, 2006.
[6] Olivier Chapelle, Jason Weston, Léon Bottou, and Vladimir Vapnik. Vicinal risk minimization. Advances neural information processing systems, 13, 2000.
[7] Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks unsuper- vised feature learning. In Proceedings of the fourteenth international conference artificial intelligence and statistics, pages 215-223, JMLR Workshop and Conference Proceedings, 2011.
[8] Gal Cohen, Hayit Greenspan, and Jacob Goldberger. Learning probabilistic fusion multilabel lesion contours. In IEEE International Symposium on Biomedical Imaging, ISBI, 2020.
[9] Yufei Cui, Wuguannan Yao, Qiao Li, Antoni Chan, and Chun Jason Xue. Accelerating Monte Carlo Bayesian prediction via approximating predictive uncertainty over the simplex. IEEE transactions on neural networks and learning systems, 2020.
[10] Erik Daxherger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux effortless bayesian deep learning. In Neural Information Processing Systems (NeurIPS), 2021.
[11] Stefan Depeweg, Jose-Miguel Hernandez-Lobato, Finale Doshi-Velez, and Steffen Udluft. Decomposition of uncertainty in Bayesian deep learning for efficient and risk-sensitive learning. In International Conference on Machine Learning, pages 1184-199.. PMLR, 2018.
[12] Ethan Fetaya, Yonatan Lifshitz, Elad Aaron, and Shai Gordin. Restoration of fragmentary babylonian texts using recurrent neural networks. PNAS, 2020.
[13] Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander J. Smola. A kernel two-sample test. Mach. Learn. Res., 13:723-773, 2012.
[14] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration modern neural networks. In International Conference Machine Learning, pages 1321-1330. PMLR, 2017.
[15] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[16] Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. Distilling the knowledge a neural network. arXiv preprint 1503::5003.0231, 2(7), 2015.
[17] Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are Bayesian neural network posteriors really like? In International Conference on Machine Learning, pages 4629-4640. PMLR, 2021.
[18] Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? In Neural Information Processing Systems (NeurIPS), 2017.
10

---

[19] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[20] Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparame- terization trick. In Neural Information Processing Systems (NeurIPS), 2015.
[21] Anoop Korattikara Balan, Vivek Rathod, Kevin P Murphy, and Max Welling. Bayesian dark knowledge. Advances in Neural Information Processing Systems, 28, 2015.
[22] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[23] Chun-Liang Li, Wei-Cheng Chang, Yu Cheng, Yiming Yang, and Barnabás Póczos. MMD GAN: towards deeper understanding of moment matching network. In Neural Information Processing Systems (NeurIPS), 2017.
[24] Yujia Li, Kevin Swersky, and Richard S. Zemel. Generative moment matching networks. In International Conference on Machine Learning (ICML), volume 37, 2015.
[25] Andrey Malinin and Mark Gales. Predictive uncertainty estimation via prior networks. Advances in neural information processing systems, 31, 2018.
[26] Andrey Malinin, Bruno Mlodozeniec, and Mark Gales. Ensemble distribution distillation. In International Conference Learning Representations, 2019.
[27] David A. McAllester. Some pac-bayesian theorems. In Conference on Computational Learning Theory (COLT), 1998.
[28] Mahdi Pakdaman Nacini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. Twenty-Ninth AAAI Conference on Artificial Intelligence. 2015.
[29] Giung Nam, Jongmin Yoon, Yoonho Lee, and Juho Lee. Diversity matters when learning from ensembles. Advances in neural information processing systems, 34, 2021.
[30] Radford M. Neal. MCMC using Hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 54:113-162 2010.
[31] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
[32] George Papamakarios. Distilling model knowledge. arXiv preprint arXiv: 2015.
[33] Siyuan Qiao, Huiyu Wang, Chenxi Liu, Wei Shen, and Alan Yuille. Micro-batch training with batch-channel normalization and weight standardization. arXiv preprint 10X.v:1903.10520, 2019.
[34] Jake Snell and Richard Zemel. Bayesian few-shot classification with one-vs-each Pólya-Gamma augmented Gaussian processes. In International Conference on Learning Representations, 2020.
[35] Yusuke Tashiro, Yang Song, and Stefano Ermon. Diversity can be transferred: Output diversifi- cation for white- and black-box attacks. In Neural Information Processing Systems (NeurIPS), 2020.
[36] Linh Tran, Bastiaan S Veeling, Kevin Roth, Jakub Swiatkowski, Joshua V Dillon, Jasper Snoek, Stephan Mandt, Tim Salimans, Sebastian Nowozin, and Rodolphe Jenatton. Hydra: Preserving ensemble diversity for model distillation. arXiv preprint (xXiv:2001.04660 2020.
[37] Martin J. Wainwright. High-Dimensional Statistics: A Non-Asymptotic Viewpoint. Cambridge University Press, 2019.
[38] Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient langevin dynamics.
In International Conference on Machine Learning (ICML), 2011.
11

---

[39] Yeming Wen, Dustin Tran, and Jimmy Ba. Batchensemble: an alternative approach to efficient ensemble and lifelong learning. In International Conference Learning Representations (ICLR), 2020.
[40] Jeremy HM Wong and Mark JF Gales. Multi-task ensembles with teacher-student training, In 2017 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU), pages 84-90. IEEE, 2017.
[41] Yuxin Wu and Kaiming He. Group normalization. In Proceedings the European conference on computer vision (ECCV), pages 3-19, 2018.
[42] Jaesik Yoon, Taesup Kim, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn. Bayesian model-agnostic meta-learning Advances in neural information processing systems, 31, 2018.
[43] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. International Conference on Learning Representations, 2018.
[44] Ruqi Zhang. Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient MCMC for Bayesian deep learning. In International Conference Learning Representations, 2019.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See section
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See section
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes] code included
as a URL.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Full details are in the supplemental material.
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [Yes] Experiments comaring our method to other baselines, done with 3 seeds, thus mean and std values presented 2
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] Provided in the supplementary. assets...
4. If you using existing assets (e.g., code, data, models) or curating/releasing
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "TpJUv4UE6aU", "reviewer": "Reviewer_rtLD", "summary": "The authors propose a method for distilling an ensemble of models. This improves performance (only one model must be stored and used). Rather than making any substantial assumptions about the probabilistic model, the method uses relatively small amounts of added noise and data augmentation.", "questions": "Mostly discussed above:\n- Why noise concat and addition?\n- What actually is the overhead?\n- Augmentation ablation?\n- Why so many ensemble elements? Does it break with fewer?\n- Is there evidence for the claim that ensemble agreement hurts the distilled model?\n- What is the 'Bayesian' baseline compared to the 'Ensemble' baseline?", "limitations": "Seems like there might be an ensemble element number limitation that isn't addressed?", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I liked this paper. It is clear and very well written. I expect it to have significant impact and it is fairly original. Most of my comments are minor and reflect things I feel could have made the arguments even stronger.\n\nThe title/name is 'Functional' but the method doesn't feel super functional to me. Part of the loss is arguably functional, but the method is basically parametric. Is there a clearer way to communicate the method?\n\nln 39 \"using\" and \"use\" redundant.\n\nThere was some overlap of background/intro/related work. Perhaps a little more editing could let you hit your contribution sooner.\n\nI don't understand the motivation for concating 3 Gaussian RVs to the input and adding noise (is it a vector? scalar?) for each block. You mention you considered other things in teh paragraph starting ln 177. Are you able to provide ablations? It isn't a deal breaker, but this felt pretty arbitrary.\n\nIn line 184 you mention non-negligible overhead. Please can you just say what the overhead is? In params and FLOPS/ms?\n\nYou mention in ln 205 that you tried a GAN thing and it didn't work as well. Can you just provide the ablation in the appendix? There's a general theme here that you allude to some work investigating alternatives and it feels like you could maybe just show the work? This could go into Table 3.\n\nIn line 234, a nitpick. It's not that I think you *should* use batchnorm, but I think you are overstating the point. First, Ismailov et al. only say that batch norm makes it harder to clearly interpret the probabilistic model, not that it is wrong. I think that's the correct view---the probabilistic model becomes a bit more complicated but it isn't all of a sudden non-Bayesian. Second, Bayesian inference doesn't make an IID assumption---in fact one of the most important parts of the Bayesian approach is that the data are not distributed at all! So I think this statement is incorrect.\n\nYou use 120 models for your ensemble. That's rather a lot. Typically people doing ensembling might use between 3 and 10. Is there a reason for your choice? Is it perhaps important for your method to have lots of ensemble elements? In that case there might be an unacknowledged limitation of the approach, that during training you have to produce many more ensemble elements in order to get a good distillation (so you are trading off training time for test time). I would like to see some ablations with 5 ensemble members, say, to understand this tradeoff better.\n\nRegarding ensemble agreement, you do show in Table 1 that the ensemble members have high agreement on the training data. But you claim that this leads to problems with the distilled ensemble. I didn't understand you to have provided any evidence for that second claim. Have I missed it?\n\nIn Table 2 there are many apparent errors in what has been made bold. For example, FED is shown bold for ECE on STL despite the fact that it is definitely not as good as EnDD_{AUX} (I understand of course that it doesn't see as much data!). Similarly, in some sense AMT has the 'best' agreement on STL-10 but it is DM_G which is bolded (and moreover the three DM methods are statistically indistinguishable). If you are going to bold things, please can you go through this carefully to make sure that the bold is correct (I don't claim to have caught all the corrections).\n\nThe 'Bayesian' baseline in Table 4 looks like it has overfit quite badly. Do you think that maybe there is a tuning problem and the baseline is a bit weak? Also, from the body of the text it seems like 'Bayesian' is the same as 'Ensemble' from other tables? Or is there a reason you use a different name?", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "RDC7ExHpntN", "reviewer": "Reviewer_m5GE", "summary": "This paper proposes a functional ensemble distillation (FED), a novel approach for distilling an ensemble of deep neural networks into a single model. The main idea is to create a generator in a student network that takes Gaussian noises as additional inputs and emulates samples from the posteriors from them and learns the diverse prediction of ensemble teachers in a nonparametric fashion using Maximum Mean Discrepancy (MMD). Additionally, to resolve the issue of ensemble teachers almost completely agreeing with each other on the training set, the paper further proposes to use mixup augmentation during training. The proposed method is validated on some benchmark image classification tasks.", "questions": "- The paper seems to mainly discuss the setting where an ensemble constructed from cSGHMC are used as a teacher. Does it work well with the typical deep ensembles with a much less number of ensemble models, $M \\in (4~8)$ for instance?\n- Similar to the above question, how does FED perform w.r.t. the different number of ensembles $M$? Does it work well with a much less number of ensemble members? This is particularly important because in real-world scenarios it is often infeasible to use 120 samples during training in a restricted resource setting.\n- I find the results confusing in some sense; most of the time, FED shows higher accuracies and lower ECEs than the baselines, but the agreement values are quite bad, which potentially indicates that the predictions produced by FED are not as diverse as teachers. This contradicts the previous reportings that the diversities in the predictions are actually key to the robustness and better uncertainty quantification in ensembles. Can you elaborate on this?\n- Usually, when it comes to comparing ECEs, it is recommended to first do the temperature scaling and compare the ECEs with the optimal temperatures. Judging from the numbers in the results, the scales of the ECEs are not similar to each other, so I suspect that the models are not temperature-scaled before computing ECEs. Can you further compare the ECEs with calibrated models?\n- In my opinion, ECE is not a good metric on its own to capture various aspects of uncertainty quantification of the model, so it is always beneficial to look at multiple uncertainty metrics such as negative log-likelihoods, brier scores, and out-of-distribution detection performances. Especially, related to the question above, since the proposed model exhibits high accuracy but high agreement values, it would be important to look at these metrics.\n- The baselines can also benefit from mixup augmentation. Have you try applying it for the baselines?\n", "limitations": "- As far as I see from the paper, DM and Hydra use only 8 heads but still distill from 120 samples; this might seriously harm the performance of the baselines, especially because these methods are designed to distill each teacher member to each subnetwork in a one-to-one fashion. If possible, it would be good to see the comparison in the setting where the teachers come with a small number of ensemble members so that one can directly distill DM and Hydra in a one-to-one fashion.\n- As mentioned above, no metrics other than ACCs and ECEs. \n- No experiments on large-scale image classification benchmarks, such as ImageNet. I find the distillation method often suffers when it comes to the scale of ImageNet, so would be good to see the scalability of the proposed methods on such datasets.\n- The main contribution can be understood as a combination of two ingredients; 1) design of the student networks that introduce Gaussian perturbations for inputs and intermediate layers, and 2) learning with MMD. Without an ablation study, it is hard to see the net effect of each component. For instance, we can try learning the proposed model with typical knowledge distillation loss, or try distilling a Hydra architecture with MMD loss.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper is relatively well written and easy to follow.\n- The motivation is well set.\n- The distillation via MMD seems novel (in the ensemble distillation literature, but not in general).\n- The experimental results seem promising.\n\nWeaknesses\n- Only evaluated with accuracy and ECE, so the results do not fully reflect the uncertainty quantification aspect of the models.\n- Somewhat unfair comparison environment for DM (see below).\n- No experiments on larger scale datasets such as TinyImageNet or ImageNet.\n ", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "D_2fCVkIqOF", "reviewer": "Reviewer_i7D4", "summary": "This paper proposes Functional Ensemble Distillation (FED), a new method for distilling (the predictions of) an ensemble of models into a single model. FED trains a generator network which learns a distribution over functions by mapping the input together with a sample of an external random noise variable to the output. As forward passes with the distilled generator model can be made for multiple samples of the random noise variable in parallel, the method is efficient to run at test time. In contrast to previous work, their approach can also quantify the correlations between predictions. Furthermore, the paper provides a solution to the case where all models in the ensemble overfit to the training data and make very similar predictions to achieve zero training loss, such that distilling their predictions becomes sub-optimal: using an auxiliary dataset generated via mixup data augmentation, the individual model predictions can be made much more diverse, which helps distill a more performant model. The paper empirically validates the efficacy of the proposed distillation method in terms of predictive accuracy and calibration compared to other baselines.", "questions": "- Do you have some intuition as for why the OOB methods in Section 5.4 perform substantially worse than the mixup approach? Is it because they simply use less data for training? Or is it because mixup introduces additional regularization that boosts performance?\n- I'm not sure if I fully understand what you mean when saying that FED scales approximately as 1 + (# of epsilon samples) / 35; e.g., for 120 epsilon samples, this would roughly be 1 + 120 / 35 = 4.43; is this referring to a multiplicative factor relative to the inference time of a single deterministic model (i.e. saying that FED with 120 samples is ~4.5x slower than a single model)?\n- Relatedly, in Appendix C.1 you mention that inference with an ensemble of 120 models takes 4.075s wall-clock time. As inference time should grow linearly with the number of ensemble members, I would expect a single model to roughly take 4.075s / 120 = 0.034s for inference. It is then claimed that FED takes only 0.020s with 120 epsilons, which is almost twice as fast as that. How can it be that FED inference time is faster than that of a single deterministic model? Following your claim of the approximate scaling of FED, shouldn't the runtime of FED be much higher than what you report? Or am I misunderstanding something here?\n- Why did you not include the most basic distlliation baseline (i.e. simply the original Hinton et al. distillation procedure) in your evaluation? I believe this would be insightful to some readers.\n", "limitations": "Yes, the authors have adequately addressed this.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Summary of Review**\n\nThe proposed new method to address the important problem of ensemble distillation is clever and empirically shown to significantly boost performance in some settings. While the method seems promising, I see some issues with the experiments (detailed below) that should be addressed to strengthen the work. Overall, I am leaning towards recommending rejection of the paper in its current form, but I am happy to reconsider my judgement if my concerns can be adequately addressed.\n\n**Strengths**\n- The manuscript is overall well written and easy to follow.\n- The work tackles an important problem, namely making ensembles more efficient via distillation.\n- The proposed FED approach for distilling an ensemble using a generator network with an external noise random variable is novel, sound, and intuitively sensible.\n- The proposed method for diversifying the ensemble member predictions in the overfitting regime by generating an auxiliary dataset via mixup augmentation is simple and intuitively appealing.\n- The empirical evaluation is thorough and demonstrates that FED can achieve significant performance gains in terms of predictive accuracy and calibration compared to relevant baselines across several image classification benchmarks.\n\n**Weaknesses**\n- I am not sure if the empirical comparison is fair, as the different methods seem to have different runtime requirements; in particular, you claim that the inference runtime of FED scales as 1 + (# of epsilon samples) / 35, which for 120 samples (i.e. the number you use in the experiments) is roughly 4.5; in my understanding, this means that the inference time of FED is 4.5x slower than that of a single deterministic model (while in contrast, the vanilla ensemble is 120x slower; I requested clarification for this in the Questions below, so please let me know if my understanding is incorrect); if this assumption is correct, this would mean that FED is significantly slower than some of the baselines considered; e.g., in my understanding, EnDD should have a runtime similar to that of a single model (again, please correct me if I'm mistaken); all in all, my concern is that the performance comparison should be conducted more carefully, taking inference runtime into account, to make it more meaningful; to this end, one could e.g. plot the results in 2D with runtime on the x-axis and performance on the y-axis; different methods would then result in different trade-offs on this plot, yielding Pareto curves that would be much more insightful than just single performance numbers; at the end of the day, a practitioner is probably interested in one of two questions: 1) given a fixed runtime budget X, what’s the best performance Y that I can achieve?, or 2) given a desired performance Y, what’s the fastest runtime X to achieve it? to make the comparison even more interesting, it would be great to consider different settings for each individual method for trading off between runtime and performance; I'm not fully sure how to best do this, but e.g. for FED, one could try varying the number of epsilon samples (for a fixed number of ensemble members) to make the method cheaper; conversely, one could try to make the baselines more expensive in more way; this would allow you to compare performance of the methods for a given fixed inference budget.\n- The empirical evaluation could be significantly improved by considering a more diverse set of model architectures and benchmarks; the paper focuses on a single architecture (ResNet18) and three small-to-medium-sized image classification benchmarks (CIFAR-10, CIFAR-100, STL-10); it is not clear if the drawn conclusions would extend to other model families (e.g. transformers), data modalities (e.g. text) and benchmark sizes (e.g. ImageNet); to gain more space for reporting additional major empirical results, one could e.g. move Section 5.4 (on the OOB methods for inducing diversity without an auxiliary dataset) to the appendix -- while the experiment and conclusion in that section are certainly interesting, they do not seem to necessary to have in the main text of the paper due to the comparably poor results.\n- The paper emphasises that, in contrast to previous methods (which are e.g. based on fitting a Dirichlet distribution over predictions), FED is able to capture correlations between outputs, and claims that this is an important feature for many real-world applications (e.g. l. 9-11, l. 52-55, l.149-153); as this is mentioned already prominently in the abstract as a main selling point for the method (l. 9-11), I was expecting at least some empirical evidence for the benefit of this feature; unfortunately, no explicit evidence is provided beyond some vague intuitive written arguments (l. 52-55, l. 149-153). It would significantly strengthen the paper if capturing these covariances could be shown to improve performance in some application. (While I understand that FED is shown to improve upon previous methods, it is not clear to me how this gain can be directly attributed to the aforementioned property.)\n\n**Minor Issues**\n- Table 3: bolding the best values would improve clarity\n- It would be useful to put a brief take-away message into the caption of each table/figure\n- The paper misses the Laplace approximation as another fundamental Bayesian inference approach for neural networks (in addition to VI and MCMC); for a recent overview, see e.g. Daxberger et al., *Laplace Redux -- Effortless Bayesian Deep Learning*, NeurIPS 2021\n- In addition to the runtime comparison requested above, it would also be useful to have a more clear/complete comparison of the memory requirements of all methods; you mention the memory burden of FED in Appendix C.2, and e.g. that of the Hydra baseline in Appendix A; it would be useful to have a table/plot comparing this for all methods considered\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "9iwzNN7ooBz", "reviewer": "Reviewer_a1xo", "summary": "This paper proposes a novel approach for distilling a Bayesian ensemble of models:  \n\n$\\int_{\\theta}p(y | x, \\theta) p(\\theta | D) d\\theta \\approx \\frac{1}{M}\\sum_{\\theta_i \\sim p(\\theta | D) } p(y | x, \\theta_i) $\n\nExisting methods tries to distill the ensemble by fitting a Dirichlet distribution: \n\n$f(x; \\phi) = \\alpha$, where $\\alpha$ is the concentration parameters that fully characterize the Dirichlet distribution \n\nThis paper proposes to learn a generator that can generate predictive distribution like \n$p(y | x, \\epsilon_i)$, where $\\epsilon_i$ are random samples from a fixed distribution. \n\nThis is in spirit very much like the reparameterization trick for variational auto-encoder [1]. \n\nThe inference can be executed by sampling a batch of $\\epsilon$’s and then run the forward function of the learned generator on this batch in parallel. \n\nTo promote the diversity of the ensemble, during training the generator, mixup data augmentation is applied, which is shown to significantly boost the accuracy and diversity. \n\n\n[1] Kingma & Welling, Auto-Encoding Variational Bayes, 2013\n", "questions": "- how does it compare, in terms of accuracy, to a naive baseline where we use multiple teachers (i.e., the ensemble of models) to distill one student model $P(y|x)$? \n \n\n- What's the decision making algorithm given an ensemble of predictions or a Dirichlet distribution? Averaging $p(y | x, \\theta_i)$ or majority vote? for EnDD, do you take the mean of the learned Dirichlet distribution or take samples from that and then take majority vote? If the former, how do you measure agreement metric? \n\n- could you elaborate on how to set the noise STD to be a learnable parameter? \n", "limitations": "n/a ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths \n\n- The proposed method is novel and general purpose, can distill any ensemble of models   \n- The idea of using mixup augmentation for distillation is creative and interesting\n- The experimental results seem strong, with well designed ablation studies on mixup.  \n\nWeaknesses: \n\n- it’s claimed to be a general-purpose distillation method, but only classification tasks are considered in the experiments; it’s not clear how good it performs on regression tasks, and maybe compare with EnDD fitting, e.g., a Gaussian distribution (which seems straightforward to adapt) \n- one advantage of the proposed method is claimed to be able to measure the correlation between two inputs, but seems it’s not discussed how to do that, and not experimentally validated that the correlation is a close approximation of the original ensemble\n- there seems to be large computational overhead during inference since each example need to be extended to a batch with different $\\epsilon$ samples, so inference cost is multiplied by the number of samples, though they could be run in parallel. \n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept