{'fraction_non_alphanumeric': 0.04056688430237955, 'fraction_numerical': 0.021258161339233222, 'mean_word_length': 4.789681117878897, 'pattern_counts': {'":': 0, '<': 2, '<?xml version=': 0, '>': 4, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Improving Ensemble Distillation With Weight Averaging and Diversifying Perturbation

Ensembles of deep neural networks have demonstrated superior performance, but their heavy computational cost hinders applying them for resource-limited environments. It motivates distilling knowledge from the ensemble teacher into a smaller student network, and there are two important design choices for this ensemble distillation: 1) how to construct the student network, and 2) what data should be shown during training. In this paper, we propose a weight averaging technique where a student with multiple subnetworks is trained to absorb the functional diversity of ensemble teachers, but then those subnetworks are properly averaged for inference, giving a single student network with no additional inference cost. We also propose a perturbation strategy that seeks inputs from which the diversities of teachers can be better transferred to the student. Combining these two, our method significantly improves upon previous methods on various image classification tasks.


Introduction
Deep Ensemble (DE; Lakshminarayanan et al., 2017) averages outputs of multiple models of the same architecture trained with the same data. Despite being simple to implement, DE achieves state-of-the-art performances for various tasks, serving as an oracle for many algorithms to compare against. However, the computational cost of DE scales linearly with the number of models involved in the ensemble, both for training and inference. Especially, the inference cost often becomes critical for a real-world scenario where both memory and time budget is limited. 1 Korea Advanced Institute of Science and Technology, Daejeon, Korea 2 Naver, Korea 3 AITRICS, Seoul, South Korea. Correspondence to: Giung Nam <giung@kaist.ac.kr>, Juho Lee <juholee@kaist.ac.kr>.
Proceedings of the 39 th International Conference on Machine Learning, Baltimore, Maryland, USA, PMLR 162, 2022. Copyright 2022 by the author(s). Knowledge Distillation (KD; Hinton et al., 2015) is a method to transfer knowledge from a large teacher network to a smaller student network. The heavy inference cost of DE thus naturally motivates applying KD to reduce it, where a DE is set as a teacher and a single neural network is introduced to be set as a student network. This task of distilling an ensemble teacher network (or distilling from multiple teachers if we treat each ensemble member as a teacher) is often called ensemble distillation and has recently been studied actively Mariet et al., 2021;Nam et al., 2021;Ryabinin et al., 2021;Du et al., 2020).
For a successful ensemble distillation, one should carefully choose the architecture for a student network. The most straightforward choice would be using a single neural network having the same architecture as the teacher, but this usually yields suboptimal results due to the limited flexibility of the student network. Another choice is to use a student network having subnetworks, where the subnetworks share most of the parameters but have a small number of individual parameters, e.g., rank-one factors (Wen et al., 2020;Mariet et al., 2021) or multiple classification heads . The ensemble distillation with subnetworks are reported to improve performance upon vanilla ensemble distillation (Mariet et al., 2021;Tran et al., 2020;Nam et al., 2021), but they usually require additional computational costs for inference. For instance, the inference cost of a student network having subnetworks defined with rank-one factors scales linearly to the number of subnetworks.
Another important choice for an ensemble distillation algorithm is the training data perturbation strategy. Recently, Nam et al. (2021) studied the importance of diversities in ensemble distillation. When ensemble teachers achieve nearzero train error, the outputs of ensemble teachers for a training data point would be nearly identical, so the ensemble distillation with normal training data would not effectively transfer diversities of the ensemble teachers to students. To this end, Nam et al. (2021) proposed to perturb training data to output-diversifying directions (Tashiro et al., 2020) and use them for distillation. While this indeed improves performance, the perturbation strategy proposed in Nam et al. (2021) only considers teachers without considering how a student would react to such perturbed data. Student arXiv:2206.15047v1 [cs.LG] 30 Jun 2022 networks are typically less flexible than teacher networks, so a perturbation increasing diversities of teachers may act in an unexpected way when applied to students. Hence, we were motivated to consider both student diversities and teacher diversities into account when designing a perturbation strategy.
In this paper, we propose a novel ensemble distillation method that resolves both of the above-mentioned limitations. Our first contribution is the new way of constructing a student network; we propose to distill with a student with multiple subnetworks during training, but average the subnetwork weights later for inference to get a single student network as a result. As a prototype, we apply this idea to BatchEnsemble (BE; Wen et al., 2020), where the subnetworks are differentiated with multiplicative rank-one factors. Specifically, for such a student network, we propose a training scheme that encourages the subnetwork parameters to stay within the same low-loss region so that the averaged prediction does not degrade the performance while maximally absorbing the diversities transferred from teachers. We show that under a representative training scheme on image classification, averaging those rank-one factors does not degrade predictive performance. The second contribution is a novel perturbation strategy improving upon the one proposed in Nam et al. (2021). Unlike the previous method, our perturbation method considers both students and teachers. More specifically, we find the "weak points" of the student networks by seeking inputs on which the student subnetworks agree with each other (low diversity), and at the same time, find the inputs on which the teachers disagree (high diversity). With this perturbation considering both students and teachers, our method effectively transfers diversities of teachers to students. To demonstrate the effectiveness of our method, we compare ours with previous methods on various image classification benchmarks. We find that ours achieve significantly improved predictive accuracy and uncertainty calibration results without increasing inference cost.

Setup
The problem we address in this paper is the K-way classification problem; a neural network F : R D → R K takes D-dimensional inputs x (i.e., images) and makes predictions about corresponding outputs y (i.e., class label) with K-dimensional logits F(x). We denote the output probabilities of the model F for a given input x as for k = 1, ..., K. Here, we introduce a single scale parameter τ > 0 for temperature scaling which will be used both for training and evaluation procedures.

Ensemble Distillation
Let {T 1 , ..., T M } be a set of pre-trained teachers, and S θ be a student. In a vanilla ensemble distillation, using the knowledge distillation (KD; Hinton et al., 2015) framework, the student tries to mimic the probabilistic outputs of an ensemble of teachers under the given temperature τ by minimizing the averaged KD loss, which is equivalent to minimizing where H[·, ·] computes the cross-entropy between two probability vectors. Note that this vanilla approach minimizes the discrepancy between the student predictions and the mean predictions of the ensemble teachers. As a result, the diversities among ensemble teachers are removed by mean operation and hardly transferred to the student.

BatchEnsemble and one-to-one distillation
BatchEnsemble (BE; Wen et al., 2020) is a parameterefficient way to ensemble deep neural networks; each member of the ensemble is constructed in the low-rank subspace with rank-one factors, instead of the full parameter space. With a slight abuse of notation, while a DE would have full set of parameters {θ 1 , . . . , θ M }, BE introduces a shared parameter θ and a set of rank-one matrices {r 1 s 1 , . . . , r m s m }, and construct m th subnetwork parameter as θ • r m s m , where • denotes the Hadamard product. Based on BE, Mariet et al. (2021) proposed a one-to-one ensemble distillation scheme, where each BE subnetwork is trying to mimic single ensemble member in one-to-one fashion. Instead of learning the mean prediction of the teachers, the one-to-one distillation minimizes That is, each subnetwork copies a member from the ensembles in a one-to-one way. The training procedure for BE ensemble distillation is summarized in Algorithm 1.

Ensemble distillation with diversifying perturbation
When the ensemble teachers are flexible enough to achieve zero-train error, their responses to a training input would be nearly identical, so a student network distilled from them would not be exposed to the diversities of the teachers. To resolve this, Nam et al. (2021) proposes to perturb training inputs with the Output Diversified Sampling (ODS; Tashiro et al., 2020) that encourages ensemble teachers disagree with each other. The ODS for an input x is computed as Compute loss for the m th subnetwork:

5:
Update rank-one factors: Update shared parameters: where w denotes a random guidance vector sampled from the K-dimensional uniform distribution with zero means. Intuitively, the ODS perturbation seeks the direction in the input space to make the output follow the random guidance vector w, and thus drives ensemble members to produce diverse outputs 1 . Nam et al. (2021) demonstrated that a BE student trained with ODS perturbation shows significantly improved performance. They also proposed an improved version of ODS perturbation called ConfODS, where the ODS perturbations are scaled by the confidence of teacher predictions.

Improved Ensemble Distillation
3.1. LatentBE : a weight averaged BE student Utilizing the subnetwork structure, a BE student can capture the diversities of ensemble teachers, but this comes at a cost of increased inference time. To get an output from a BE network, one should execute forward passes M times from scratch since the different rank-one subnetworks do not share hidden layer responses. Hence, although BE significantly reduces the number of parameters compared to DE, its computation cost for inference is identical to that of DE. This is definitely undesirable, especially for distillation where we want a cheap student network applicable for real-world applications. In this section, we propose a novel ensemble distillation framework to circumvent this limitation, where a student network has the same inference cost as a single neural network yet maintains the flexibility to well absorb the diversity of teachers.
Based on the ensemble distillation with BE, we propose a novel framework entitled LatentBE, in a sense that the rankone factors defining BE are averaged out for inference just as the latent variables are marginalized out for probabilistic Algorithm 2 Ensemble distillation with LatentBE + diversifying perturbation Require: Temperature τ , learning rate η, rank-one weight decay parameter λ, perturbation step size γ. 1: Initialize rank-one factors to ones: r m ← 1 and s m ← 1 for m = 1, ..., M . 2: while not converged do 3: Sample an input x from the train split.

6:
for m = 1, ..., M do 7: Compute loss for the m th subnetwork: Update rank-one factors: inference. Specifically, we employ a BE student having multiple subnetworks, and do one-to-one ensemble distillation similarly to Algorithm 1, but after training, compute the weight average of the rank-one factors to construct a single student network with parameter After this weight averaging, the inference cost remains the same as that of a single neural network. The idea of weight averaging was first considered in Stochastic Weight Averaging (SWA; Izmailov et al., 2018), where the parameters collected from a single learning trajectory is averaged to construct a better generalizing model. The key observation in SWA is that due to the choice of the specific learning rate scheduling, the parameters in a learning trajectory remain in a wide low-loss region in the loss surface, so averaging them leads to a single robust model. Our LatentBE shares some spirits with SWA but has crucial differences: 1) LatentBE utilizes the diverse subnetworks from the guidance of ensemble teachers, which brings diversity and performance gain of ensemble distillation to weight average, and 2) instead of special learning rates schedule of SWA, LatentBE enables weight average with the rank-one factors of BE.
The key for making LatentBE successful is, as in SWA, to keep the subnetwork parameters stay in the same lowloss region. For this, 1) we initialized all the rank-one (a) BE-2 (b) LatentBE-2 Figure 1. A schematic diagram depicting BE and LatentBE. Here, − → d1 and − → d2 denote learning directions obtained from two different teachers. For the BE, (a) two subnetwork parameters explore different modes with low-rank subspaces, while for the LatentBE, (b) two subnetwork parameters expand a low-rank subspace.
factors {r m , s m } M m=1 to be one vectors, and 2) set Gaussian prior with one vector mean to those rank-one factors. By doing these, all the rank-one factors start from a similar location and gradually split into individual factors, but they do not deviate too much from each other due to the weightdecay effect driven by the prior. Also, other than these two, we do not require careful learning rate scheduling as in SWA, and this is presumably due to the fact that we are only differentiating rank-one factors for subnetworks with a large body of shared parameters while SWA averages entire parameters.
The rank-one factors for BE and LatentBE play different roles in ensemble distillation. BE aims to find a subspace for each rank-one factor and spread those subspaces to different modes. On the other hand, LatentBE seeks for a flat minima in which all the subspaces defined by the rank-one factors are embedded, and then the rank-one factors are trained to "stretch" the subspace area by letting the subnetworks follow different teacher directions ( Figure 1).

A better perturbation strategy
As we discussed earlier, Nam et al. (2021) utilizes perturbation strategies, ODS and ConfODS, to improve diversity transfer in ensemble distillation. It is an innovative approach for ensemble distillation but does not consider the student networks, especially their diversities. Since there have been several works controlling teacher networks based on the status of student networks for KD (Jin et al., 2019;Mirzadeh et al., 2020), we further conjecture that the ensemble distillation can also be improved with a perturbation strategy considering both teachers and students. Thus, we propose a novel perturbation scheme that considers both student and teacher diversities for more effective diversity transfer.
For a given input x, we measure the functional diversity of the ensemble of {F 1 , ..., F M } by averaging pairwise KL-divergence between probabilistic outputs from different members, . From this, we denote the student and teacher diversities as TDiv ( where S m := S θ•rms m . We suggest perturbing input to the direction that minimizes the student diversity while the teacher diversity is maximized, The intuition behind this perturbation is as follows: ideally, we want the student subnetworks to learn the diverse outputs of ensemble teachers almost everywhere in the input space. Hence, we first introduce negative student diversity term "− SDiv(x)" to find a point that student subnetworks have low diversities. At the same time, as we originally intended, we use teacher diversity term "TDiv(x)" to improve diversity transfer of ensemble teachers. The combined perturbation thus finds an input point that maximizes the diversity gap between teachers and student subnetworks and gives a strong learning signal to correct it.
In practice, exactly computing (7) and (8) would be costly, so we use stochastic approximations of them where the student and teacher diversities are computed for a randomly selected pair. That is, we pick i, j ∼ {1, . . . , M } uniformly and compute and define the perturbation aŝ We also found that blocking the gradient flow through one of the teachers or students stabilizes the training, where sg(·) denotes the stop grad operation, for example, .detach() in PyTorch library.
One thing to note here is that we are directly measuring divergences between teachers or student subnetworks to get perturbations, unlike the ODS-based perturbation proposed in Nam et al. (2021). An ODS computed from a specific network, in principle, does not guarantee the output diversification of other networks. Hence, the ODS perturbation computed from a single teacher, as suggested in Nam et al. (2021), does not guarantee the diversities among teacher outputs. Nam et al. (2021) argued that this issue can be circumvented by assuming transferability of teacher networks, where we assume that the gradients of ensemble teachers are similar to each other, so an ODS computed from a specific teacher generalizes to the other teachers, driving overall diversities among teacher outputs as a result. However, the transferability does not always hold for which an ODS perturbation fails to properly bring diversities. On the other hand, our perturbation directly minimizing or maximizing KL divergences does not require transferability of gradients.

Improved ensemble distillation algorithm
Our final ensemble distillation algorithm combines two ingredients discussed so far; LatentBE and novel diversifying perturbation. The only overhead during training is the procedure of computing the diversifying perturbation which requires additional forward and backward passes through teacher networks. Our algorithm is summarized in Algorithm 2, with highlights on the part different from the vanilla one-to-one ensemble distillation with BE.

Related Works
Ensembles Recent works have shown that an ensemble of deep neural networks can achieve superior performance both in terms of prediction accuracy and uncertainty estimation (Lakshminarayanan et al., 2017;Ovadia et al., 2019). The power of the ensemble comes from the diversity among ensemble members, and there have been several works to enhance it, e.g., constructing ensembles with varying hyperparameters (Wenzel et al., 2020), or architectures (Zaidi et al., 2021), or reducing conditional redundancy (Rame & Cord, 2021), or introducing kernelized repulsion (D'Angelo & Fortuin, 2021).
Ensemble distillation The seminal work of Hinton et al. (2015) has already shown the effectiveness of the ensemble distillation. It can be further enhanced by considering the diversity inside the ensemble teacher, e.g., dynamically assign weights to teachers (Du et al., 2020), or treating predictions from teachers as a set of samples from an implicit distribution (Malinin et al., 2020;Ryabinin et al., 2021), or amplifying the diversity via input perturbations (Nam et al., 2021). Besides, several existing approaches propose to use a student having subnetworks which can represent the diversity in predictions Mariet et al., 2021;Nam et al., 2021). However, their resulting students have additional costs for inference and defeat their ends that reduce the computational cost of the ensemble.
• How does the proposed perturbation strategy affect the training of LatentBE? -Section 5.2.
• Does our ensemble distillation algorithm improves performance both in terms of predictive accuracy and uncertainty calibration? -Section 5.3.
Please refer to Appendix B for the training details including data augmentation, learning rate schedules, and other hyperparameter settings.

Subspaces of LatentBE
In order to investigate the subspaces defined by the subnetwork parameters {θ • (r m s m )} M m=1 , we first consider the case of M = 2 models where the subspace forms a simple line. More precisely, we parameterize the line passing through two subnetwork parameters as {θ t | t ∈ R}, where θ t = (1 − t) θ • (r 1 s 1 ) + t θ • (r 2 s 2 ) .
(15) Figure 2 shows how prediction error and negative loglikelihood vary along the line subspace. The main difference between BE and LatentBE students is the presence of a loss barrier between two end-points. The LatentBE student does not have a barrier while the BE student does, as we have depicted in Figure 1. This difference enables the weight averaging of subnetwork parameters for LatentBE. Notably, as shown in (Figure 2, right), the averaged parameter effectively improves negative log-likelihood on the test data, which is consistent with the findings Izmailov et al. (2018) and recent study on the neural network subspaces (Wortsman et al., 2021). These improvements in performance, as we will show later in Section 5.3, confirm the validity of our weight averaging strategy for ensemble distillation.

Diversification effects of perturbations
Using on the LatentBE-2 model discussed on the previous section, Figure 3 shows the functional diversity (defined in Equation (6)) between two end-points, when distilled with and without ConfODS (Nam et al., 2021) and ours (i.e.,  Figure 3. Function diversity in predictions from two subnetwork parameters of the LatentBE-2 student on the train (left) and test (middle) data measured for a training run. We also visualize the resulting student's test NLL (right). LatentBE-2 students are distilled from the DE-2 teacher for WRN28x4 on CIFAR-100 with the stated input perturbations. TDiv-SDiv). As one can see from the figures, ours better diversifies the subnetworks while maintaining lower test NLL values. Figure 4 shows the effects of the perturbation strategies during training. For this, we measure the changes in the function diversities of the teacher networks and student subnetworks due to a perturbation ε during a training procedure: Again, our intuition behind the perturbation strategy proposed in Section 3.2 is, pinpointing inputs that should be diversified for students while diversifying learning signals from teachers. In other words, the perturbation that decreases Equation (17) while increasing Equation (16) will be helpful for ensemble distillation.
Figure 4 clearly shows that our proposed perturbation (i.e., TDiv -Sdiv) accomplishes our goal to increase teacher diversities and decrease student diversities. Especially, the student diversities significantly drop during the early stage of training, but gradually increase as the training progress. We conjecture that this is because after enough training the students are diversified for a wide range of inputs, so the effect of −Sdiv perturbation gradually decreases. Although ConfODS exhibits high diversification effects on teachers as intended, it has no significant effect on the students.
As an ablation study, we compared the efficacy of the perturbation strategies in terms of the actual classification performance. Table 1 shows that increasing student diversities (TDiv-SDiv) clearly improves the performance compared to the ones not considering the student diversities (ConfODS or TDiv).

Results on CIFAR-10/100
We compare ours to the existing ensemble distillation methods. Here, we consider baselines using a single student network: KD (Hinton et al., 2015), AE-KD (Du et al., 2020), and Proxy-EnD 2 (Ryabinin et al., 2021). As suggested in Ashukha et al. (2020), we report both original and calibrated metrics for NLL and ECE. See Appendix B.4 for the details in distillation methods, and Appendix B.3 for the definitions of evaluation metrics.
The results for CIFAR-10 and CIFAR-100 are presented in Tables 2 and 3. LatentBE consistently outperforms the baselines both in terms of accuracy and uncertainty estimates, and our perturbation further boosts up the performance of LatentBE. We note that our method gets better as the number of teachers increases, indicating that ours effectively transfers diversities from multiple teachers.
Moreover, a benefit of our approach is that it consistently outperforms the vanilla KD even when the number of teachers M is small. On the other hand, Ryabinin et al. (2021) suffers when M is small because the estimation for Dirichlet parameters required for the distillation become inaccurate. This gives an advantage to our method under resource limited setting where training large number of ensemble teachers are intractable.  (Hinton et al., 2015) 72.49±0.38 1.492±0.039 0.202±0.004 AE-KD (Du et al., 2020) 72.06±0.49 1.540±0.046 0.207±0.005 Proxy-EnD 2 (Ryabinin et al., 2021) (Hinton et al., 2015) 48.79±0.16 3.410±0.056 0.341±0.005 AE-KD (Du et al., 2020) 48.86±0.12 3.407±0.066 0.340±0.007 Proxy-EnD 2 (Ryabinin et al., 2021)  Robustness to common corruptions We further evaluate our method on corrupted CIFAR datasets to verify its robustness under common corruptions (Hendrycks & Dietterich, 2019). Table 4 reports the evaluation metrics averaged over all types of corruptions and intensities and shows that ours are better calibrated than the existing baselines. Out-of-distribution data A well-calibrated classifier should be uncertain for out-of-distribution data while being certain for in-distribution data. Following Lakshminarayanan et al. (2017), we draw histograms depicting the distribution of the predictive entropy for in-distribution (i.e., CIFAR-10/100) and out-of-distribution data (i.e., SVHN) in Figures 5 and 6. It shows that ours are more uncertain about unseen classes, behaving similarly to ensemble teachers. Refer to Appendix A.2 for more results with other out-of-distribution datasets.

Runtime analysis
Tables 2 and 3 show that LatentBE is competitive or even better for some metrics than the one-to-one distillation methods with BE students (Mariet et al., 2021;Nam et al., 2021). This is quite remarkable since BE requires additional inference cost while LatentBE doesn't. More specifically, Table 5 reports the runtimes of BE and LatentBE on the same single GeForce RTX 3090 setting. We measure the wall-clock time in hours for training and milliseconds per image (ms/img) for inference. In principle, BE requires multiple forward passes for inference (i.e., 0.390 ms/img), but this can be avoided by parallelized inference (0.288 ms/img). Still, the parallelization requires additional cost from handling larger mini-batches, which is still significantly larger than a single model inference time (0.107 ms/img).

Results on TinyImageNet and ImageNet-1k
We verify the scalability of our proposed method on largescale datasets, including TinyImageNet and ImageNet-1k (Russakovsky et al., 2015). Tables 6 and 7 show that ours outperform the baselines for both datasets.

Conclusion
In this paper, we proposed a novel ensemble distillation algorithm improving both prediction accuracy and uncertainty calibration without increasing inference cost. We first presented LatentBE, where the rank-one factors of BEs are trained in a one-to-one way, but later weight-averaged for inference. We showed that under a suitable training scheme, the subspaces defined by the rank-one factors of BE remain in a flat minimum, and weight averaging those rank-one factors thus yields a robust single student model. We further presented a novel perturbation strategy for ensemble distillation that decreases student diversities and increases teacher diversities at the same time. By training with inputs perturbed in that way, we can effectively enhance the diversities of students. Our ensemble distillation algorithm combining these two achieved remarkable performance on various image classification tasks.

A. Additional Results
A.1. Baseline results for CIFAR-10/100 In Table 8, we report the baseline results for the experiments on CIFAR-10/100: (1) evaluation results for DE teachers distilling knowledge into students, (2) performance of the single model trained with the classical cross-entropy loss without distillation, and (3) performance of the single model distilled from DE-1.

A.2. Further experiments on predictive uncertainty
For the experiments on predictive uncertainty, we consider SVHN (Netzer et al., 2011), LSUN (Yu et al., 2015), and TinyImageNet (Russakovsky et al., 2015) as out-of-distribution data. Here, LSUN and TinyImageNet images are downscaled into 32 × 32 × 3. Figures 7 and 8 further provide the predictive uncertainty results on LSUN and TinyImageNet. Again, our approach exhibits higher predictive uncertainty on out-of-distribution examples than existing baselines.

B.1. Datasets
CIFAR-10/100 The dataset is available at https://www.cs.toronto.edu/˜kriz/cifar.html. It consists of 50,000 train examples and 10,000 test examples from 10/100 classes, with images size of 32 × 32 × 3. In this paper, the last 5,000 examples of the train split are used as the validation split for computing calibrated metrics. We follow the standard data augmentation policy (He et al., 2016)   ImageNet-1k It consists of 1,281,167 train examples, 50,000 validation examples and 100,000 test images from 1,000 classes. Since the labels of the official test set are not publicly available, we only report the evaluation results on the validation set. We follow the standard data augmentation policy from PyTorch which consists of random cropping with an images size of 224 × 224 × 3 and random horizontal flipping 3 . Throughout experiments on ImageNet-1k classification, we use the ResNet-50 (R50) network introduced in He et al. (2016).

B.2. Training details
All images are standardized by subtracting the per-channel mean and dividing the result by the per-channel standard deviation. We use SGD optimizer with Nesterov momentum 0.9, and a single-cycle cosine annealing learning rate schedule with a linear warm-up, i.e., the learning rate starts from 0.01 × base lr and reaches base lr after the first 5 epochs, and is decayed by the single-cycle cosine annealing learning rate schedule. More precisely, (1) for CIFAR-10/100, we run 200 epochs on a single machine with batch size 128 and base lr = 0.1, (2) for TinyImageNet, we run 80 epochs on four machines with the total batch size 128 and base lr = 0.1, and (3) for ImageNet-1k, we run 100 epochs on eight machines with the total batch size 256 and base lr = 0.1. We also apply the weight decay (Krogh & Hertz, 1991) to regularize training; the weight decay coefficient is set to be 0.0005 for CIFAR-10/100 and TinyImageNet, and 0.0001 for ImageNet-1k.

B.3. Evaluation
The problem we address in this paper is the K-way classification problem; a neural network f : R D → R K takes Ddimensional inputs x (i.e., images) and makes predictions about outputs y (i.e., class label) with K-dimensional logits. We denote the output probabilities of the model f for a given input x as , for k = 1, ..., K.
Standard metrics We evaluate the following standard metrics of the model f on the dataset D: • ACC (accuracy; higher is better): where [·] denotes the Iverson bracket.
• NLL (negative log-likelihood; lower is better): • ECE (expected calibration error; lower is better): where {B 1 , ..., B L } is a partition of D, where B l = {(x, y) ∈ D | max k p (k) f (x) ∈ ((l − 1)/L, l/L]}. Here, the difference between the accuracy and mean confidence of predictions represents the calibration gap for each bin. We fixed L = 15 for all evaluation results in this paper.
Calibrated metrics Temperature scaling softens output probabilities of the model with a single scale parameter τ > 0, and it can be used for calibrating probabilistic models without affecting the model's prediction accuracy (Guo et al., 2017). We define the temperature scaled output probabilities of the model f for a given input x as p (k) f (x; τ ) = exp f (k) (x)/τ K j=1 exp f (j) (x)/τ , for k = 1, ..., K.
Following Ashukha et al. (2020), we also evaluate the calibrated metrics which are computed using the temperature scaled outputs. Specifically, we first find the optimal temperature which minimizes the NLL on the validation split D valid , i.e., and then we can compute the following calibrated metrics: • cNLL (calibrated negative log-likelihood): cNLL(f , D, τ * ) = − (x,y)∈D y log p