Neural networks with late-phase weights
Authors: Johannes von Oswald
Seijin Kobayashi
Joao Sacramento
Alexander Meulemans
Christian Henning
Benjamin F Grewe
Venue: ICLR 2021 Poster (2021)

Abstract: The largely successful method of training neural networks is to learn their weights using some variant of stochastic gradient descent (SGD). Here, we show that the solutions found by SGD can be further improved by ensembling a subset of the weights in late stages of learning. At the end of learning, we obtain back a single model by taking a spatial average in weight space. To avoid incurring increased computational costs, we investigate a family of low-dimensional late-phase weight models which interact multiplicatively with the remaining parameters. Our results show that augmenting standard models with late-phase weights improves generalization in established benchmarks such as CIFAR-10/100, ImageNet and enwik8. These findings are complemented with a theoretical analysis of a noisy quadratic problem which provides a simplified picture of the late phases of neural network learning.

Published as a conference paper at ICLR 2021
NEURAL NETWORKS WITH LATE-PHASE WEIGHTS
Johannes von Oswald Seijin Kobayashi João Sacramento*, Alexander Meulemans, Christian Henning, Benjamin F. Grewe
equul contribution
Institute of Neuroinformatics
University of Zürich and ETH Zürich
Zürich, Switzerland
(voswaldj, sei jink, r joao, meeuleemahemiing,,Beee hgh.eeeee.gg.....eeeeeeee...... ch
ABSTRACT
The largely successful method of training neural networks is to learn their weights using some variant of stochastic gradient descent (SGD). Here, we show that the solutions found by SGD can be further improved by ensembling subset of the weights in late stages of learning. At the end of learning, we obtain back a single model by taking a spatial average in weight space. To avoid incurring increased computational costs, we investigate family of low-dimensional late-phase weight models which interact multiplicatively with the remaining parameters. Our re- sults show that augmenting standard models with late-phase weights improves generalization in established benchmarks such as (IFAR-I0/100 ImageNet and enwik8. These findings are complemented with a theoretical analysis of noisy quadratic problem which provides simplified picture of the late phases of neural network learning.
INTRODUCTION
Neural networks trained with SGD generalize remarkably well on wide range of problems. A classic technique to further improve generalization is to ensemble many such models (Lakshmi- narayanan et al., 2017). At test time, the predictions made by each model are combined, usually through a simple average. Although largely successful, this technique is costly both during learning and inference. This has prompted the development of ensembling methods with reduced complex- ity, for example by collecting models along optimization path generated by SGD (Huang et al., 2017), by performing interpolations in weight space (Garipov et al., 2018), or by tying a subset of the weights over the ensemble (Lee et al., 2015; Wen et al., 2020).
An alternative line of work explores the use of ensembles to guide the optimization a single model (Zhang et al., 2015; Pittorino al., 2020). We join these efforts and develop method that fine-tunes the behavior of SGD using late-phase weights: late in training, we replicate a subset the weights of neural network and randomly initialize them in small neighborhood. Together with the stochasticity inherent to SGD, this initialization encourages the late-phase weights to explore the loss landscape. As the late-phase weights explore, the shared weights accumulate gradients. After training we collapse this implicit ensemble into a single model by averaging in weight space.
Building upon recent work ensembles with shared parameters (Wen et al., 2020) explore family of late-phase weight models involving multiplicative interactions (Jayakumar et al., 2020). We focus on low-dimensional late-phase models that can be ensembled with negligible overhead. Our experiments reveal that replicating the ubiquitous batch normalization layers (loffe & Szegedy, 2015) is a surprisingly simple and effective strategy for improving generalization' Furthermore, we find that late-phase weights can be combined with stochastic weight averaging (Izmailov et al., 2018), a complementary method that has been shown to greatly improve generalization.
We provide code to reproduce experiments at https::/gg thub. com/ sei j no-oobayshh 1 Tate- tatte-phasewwa ghts

---

Published as a conference paper at ICLR 2021
2 METHODS AND MODELS
2.1 LEARNING WITH PHASE WEIGHTS
Late-phase weights. To apply our learning algorithm to given neural network model fw we first specify its weights w in terms of two components, base and late-phase (0 and o, resp.). The two components interact according to a weight interaction function w = h(0, Φ). Base weights are learned throughout the entire training session, and until time step To both 6 and & are learned and treated on equal grounds. At time step To, a hyperparameter of our algorithm, we introduce K late-phase components +- ) that are learned together with 0 until the end.
This procedure yields a late-phase ensemble of neural networks with parameter sharing: reusing the base weights f,, each late-phase weight &R defines a model with parameters tii h(0,, ox).
Late-phase weight averaging at test time. Our ensemble defined by the K late-phase weight configurations in & is kept only during learning. At test time, discard the ensemble and obtain a single model by averaging over the K late-phase weight components. That is, given some input pattern x, generate prediction ya:) using the averaged model, computed once after learning:
K
1
y(z) fT(().
W
h
0.
Φk
K
Hence, the complexity of inference is independent of K, and equivalent to that of the original model.
Late-phase weight initialization. We initialize our late-phase weights from a reference base weight. We first learn a base parameter Φ from time step & 0 until To, treating Φn as any other base parameter in 0. Then, at time t = To, each configuration Φ is initialized in the vicinity of Φu. We explore perturbing øu using symmetric Gaussian noise model,
do
Φ = Φu +
Ck: Z(oo)
where Eμ is standard normal variate of appropriate dimension and do is hyperparameter con- trolling the noise amplitude. We allow for @o-dependent normalization factor, which we set so as to ensure layerwise scale-invariance, which helps finding a single do that governs the initializa- tion of the entire network. More concretely, for given neural network layer I with weights & of dimension D(!), we choose Z( /DD/119 |l
Our perturbative initialization (Eq. 2) is motivated by ongoing studies of the nonconvex, high- dimensional loss functions that arise in deep learning. Empirical results and theoretical analyses of simplified models point to the existence of dense clusters of connected solutions with a locally- flat geometry (Hochreiter & Schmidhuber, 1997a) that are accessible by SGD (Huang et al., 2017; Garipov et al., 2018; Baldassi et al., 2020). Indeed, the eigenspectrum of the loss Hessian evalu- ated at weight configurations found by SGD reveals a large number of directions of low curvature (Keskar et al., 2017: Chaudhari et al., 2019; Sagun et al., 2018). For not yet completely under- stood reasons, this appears to be recurring phenomenon in overparameterized nonlinear problems (Brown& & Sethna, 2003; Waterfall et al., 2006).
Based on these observations, we assume that the initial parameter configuration do can be perturbed in a late phase of learning without leading to mode hopping across the different models wx. While mode coverage is usually sought after property when learning neural network ensembles (Fort et al., 2020), here it would preclude us from taking the averaged model at the end of learning (Eq. 1).
Stochastic learning algorithm. Having decomposed our weights into base and late-phase com- ponents, we now present stochastic algorithm which learns both @ and 4. Our algorithm works on the standard stochastic (minibatch) neural network optimization setting (Bottou, 2010). Given a loss function L(D, w) [d] ΣeED L(x, w) to be minimized with respect to the weights w on set of data D. at every round we randomly sample subset M from D and optimize instead the stochastic loss L(M, w). However, in contrast to the standard setting, in late stages learning (f > Ta) we simultaneously optimize K parameterizations W:= (wk Wk h(O, Φk) K k= instead of one.
2

---

Published as a conference paper at ICLR 2021
We proceed by iteration over W. At each step k, we sample a minibatch Ma and immediately update the late-phase weights Øk, while accumulating gradients over the shared base weights 8. Such gradient accumulation has been previously used when learning ensembles (Lee et al., 2015; Wen et al., 2020) and multi-task models (Rebuffi et al., 2017) with shared base parameters. A single iteration is finally concluded by changing the base weights in the direction opposite of the accumulated gradient. We scale the accumulated gradient by 78; setting " = 1/ K recovers the original step size in 0, but other choices are possible. In particular, we find that large 79 of unit size is in practice often tolerated, resulting in accelerated learning.
We summarize an iteration of our method in Algo- rithm 1, where the loss L(M, f Φ) is now seen as function of 0 and d) We opt for a general presenta- tion using unspecified gradient-based update operators Ug and Ug. These operators can be set to optimiz- ers of choice. For instance, our method might benefit from additional noise injection onto parameter updates (Welling & Teh, 2011), Furthermore, late-phase opti- mizers need not coincide with the optimizer used in the early phase. In our work we typically set Ue and Un to a single step of SGD with Nesterov momentum (Nesterov, 2004), and explore Adam (Kingma Ba, 2015) and plain SGD in a smaller set of experiments.
Algorithm 1: Late-phase learning
Require: Base weights A, late-phase weight set Φ dataset D, gradient scale factor na. loss L Require: Training iteration 1>> To
for 1 S k S. K do
Mr + Sample minibatch from D
A0% + V L(MMM 0, Φx)
ΦR x Ua(or, V L(ME: B, @x))
0 - Uol0, 70 Σx A0A)
2.2 LATE-PHASE WEIGITT MODELS
As detailed next, we consider a number of distinct late-phase weight models in our experiments. In particular, we explore weight interaction functions which late-phase weights have low di- mensionality, to avoid large increase in complexity with the ensemble size K. To counteract this reduced dimensionality, we make extensive use of multiplicative base-late weight interactions. This design choice is motivated by the large expressive power of multiplicative interactions despite low dimensionality, which has been demonstrated in a wide range of settings (Jayakumar et al., 2020).
Late-phase batch normalization layers. Batch normalization layers (BatchNorm; loffe & Szegedy, 2015) are a staple of current deep neural network models. Besides standardizing the ac- tivity of the layer they are applied to, BatchNorm units introduce learnable multiplicative (scale) parameter and an additive (shift) parameter 3. While being low-dimensional, these additional parameters have large expressive power: it has been shown that learning only 7 and keeping the remaining weights frozen can lead to significantly lower loss than when learning random subsets of other weights of matching dimensionality (Frankle et al., 2020; Mudrakarta al., 2019).
We take the scale and shift parameters of BatchNorm layers as our first choice of late-phase weights; the base weights are the remaining parameters of the model. Batch statistics are also individually estimated for each model in W. This late-phase weight parameterization is motivated by (i) the expressive power of y and | discussed above, and by (ii) practical considerations, as BatchNorm layers are generally already present in feedforward neural network models, and are otherwise easy to implement efficiently.
More concretely, let us consider an affine transformation layer which maps an input vector (1-1) to 0i) ,((-1) + 0 where the early-phase weight matrix dL and bias vector 0/9 are already stan- dardized using the respective batch statistics. For this standard layer, our model introduces a multi- plicative interaction between base and late- aate- weights, (i)g(> o!! and an additive interaction between base and late-phase bias parameters, 0²¹) + g())
Late-phase rank-1 matrix weights. We also study a closely related late-phase weight model, where existing weight matrices the base components, as before are multiplied elementwise by rank-1 matrices (Wen et al., 2020). For a given affine layer I, we define a late-phase weight matrix with resort pair of learnable vectors, e(¹) (! ,0)T. Taking the Hadamard product with the base weight matrix yields the effective weights (i)) g(00) g0)
3

---

Published as a conference paper at ICLR 2021
With this parameterization, recover the ensemble proposed by Wen et al. (2020), except that here it is generated late in training using our perturbative initialization (Eq, 2). Unlike BatchNorm layers, which include the shift parameter, rank-1 late-phase weights interact in purely multiplicative man- ner with base weights. We study this model since it is easy to implement on neural networks which do not feature BatchNorm layers, such as standard long short-term memories (LSTMs; Hochreiter & Schmidhuber, 1997b).
Hypernetworks with late-phase weight embeddings. Additionally, we generalize the late-phase weight models described above using hypernetworks (Ha et al., 2017). A hypernetwork generates the parameters of given target neural network fu based weight embedding. In our frame- work, we can use a hypernetwork to implement the interaction function w = h(0, o) directly, with late-phase weights.
parameters 0 corresponding base weights and embeddings
We experiment with linear hypernetworks and use the same hypernetwork to produce the weights of multiple layers, following Savarese & Maire (2019); Ha et al. (2017); von Oswald et al. (2020). In this scheme, the weight embedding input specifies the target layer whose parameters are being generated. More specifically, the weight matrix for some layer belonging to a group of layers which share hypernetwork is given by W(a.)) ((a) g(a.), where ((a) and ((a.)) appropriately- sized tensors. Sharing ((9) over layer group allows countering increase in the overall number of parameters. We parameterize our hypernetworks such that the weight embedding vectors g(9.)) are small, and therefore cheap to ensemble.
Late-phase classification layers. Finally, inspired by Lee et al. (2015), in classification experi- ments we take the weights of the last linear layer as late-phase weights by default. modern neural network architectures these layers do not usually comprise large numbers of parameters, and our architecture explorations indicated that it is typically beneficial ensemble them. We therefore include (L) in our late-phase weights o, where (L)) denotes the weights of the final layer L.
3 RESULTS
3.1 NOISY QUADRATIC PROBLEM ANALYSIS
Before tuming to real-world learning problems, first focus simplified stochastic optimization setup which can be analytically studied. We consider the noisy quadratic problem (NQP; Schaul al., 2013; Martens, 2016; Wu et al., 2018; Zhang al., 2019a;b), where the goal is to minimize the scalar loss
1 5 ww w" +e)²
TTTTT H (w w" +e)
with respect to R". In the equation above, wn denotes the target weight vector, which is randomly shifted by a noise variable assumed to follow a Gaussian distribution W/ E). The (constant) Hessian matrix H controls the curvature of the problem.
Despite the simplicity of Eq. 3, the NQP captures surprising number of empiricalll--obeeved aspects of neural network learning (Zhang et al., 2019a). Here, motivate its study as a model of late stages of learn- ing, by Taylor expanding the loss around a minimum w". Thus, for sufficiently late initialization time To (and small on) the NQP is particularly well suited to study our algorithm.
Ours
Ensemble
0.004
Loss
0.002
5
10
15
20
25
There are three main strategies to improve the ex- pected NQP loss after convergence: (i) increase the minibatch size B, (ii) use more members K in en- semble, and (iii) decrease the learning rate 77 (Zhang al., 2019a). Our Algorithm 1 combines the first two strategies in a non non-trivial manner. First. the gradients for base weights 0 are averaged during the in- loop over all ensemble members, corresponding minibatch-size rescaling by K. Second,
K
Figure 1: Steady-state loss for varying K, of multiplicative late-phase weights (Ours) compared to an ensemble of models.
4

---

Published as a conference paper at ICLR 2021
introduce K ensemble members, to be averaged in weight space, that only differ in their late- late-phase weights d.
In Appendix C. we show analytically that this combination of increased effective minibatch size for and introducing ensemble members for Φ is successful, resulting in a scaling of the expected loss after convergence by K This analysis holds for general E and H, and for both scalar and hypernetwork multiplicative late- thase weights. Hence, our approach combines the benefits of an increased effective minibatch size and of ensembling, while yielding a single model after training.
We present a numerical validation of this theoretical result in Fig. 1. Our model includes a multi- plicative late-phase weight, Wa 0 Φb with ΦR E RR and R" We simulate standard instance of the NQP, with diagonal Hessian Hμ = 1/i and Σ = H (cf. Zhang et al., 2019at 2019a), and report the average loss after convergence. Hyperparameters are given Appendix C. As predicted by the theory, the loss falls as ~ 1/ K with increasing ensemble size K, and our algorithm performs on par with full ensemble of models trained independently with gradient descent.
3.2 CIFAR- 10/100 EXPERIMENTS
To test the applicability of our method to more realistic problems, next augment standard neu- ral network models with late-phase weights and examine their performance on the CIFAR- 10 and CIFAR- 100 image classification benchmarks (Krizhevsky, 2009). We use standard data preprocess- ing methods (cf. Appendix A) and train our models for 200 epochs from random initializations, except when noted otherwise. All evaluated methods are trained using the amount of data.
Besides SGD (with Nesterov momentum), we also investigate stochastic weight averaging (SWA; Izmailov et al., 2018), recent reincarnation of Polyak averaging (Polyak & Juditsky, 1992) that can strongly improve neural network generalization. For completeness, we present pseudocode for SWA in Algorithm 2 and SGD with Nesterov momentum in Algorithm 3 (ef. Appendix A). When learning neural networks with late-phase weights we set Us and Ue to one step of SGD (or SGD wrapped inside SWA).
We compare our method to dropout (Srivastava et al., 2014), a popular regularization method that can improve generalization in neural networks. Like our approach, dropout produces single model at the end of training. We also consider its Monte Carlo variant (MC-dropout; Gal & Ghahramani, 2016), and the recently proposed BatchEnsemble (Wen et al., 2020). This method generates an ensemble using rank-1 matrices as described in Section 2.2. Predictions still need to be averaged over multiple models, but this averaging step can be parallelized in modern hardware.
Additionally, we report single-seed results obtained with an ensemble of K independently models (a deep ensemble, Lakshminarayanan et al., 2017). Deep ensembles provide a strong base- line, at the expense of large computational and memory costs. Therefore, they are not directly comparable to the other methods considered here, and serve the purpose an upper baseline.
By contrast, augmenting the architectures con- sidered here with late-phase weights results in negligible additional costs during learn ing (with the exception of hypernetworks, which require additional tensor products) and none during testing. In principle, a set of independentty-trained models yielded by our algorithm can therefore even be used as the basis of a deep ensemble, when the mem- ory and compute budget allows for one. We present proof-of-concept experiments explor- ing this option.
Table CIFAR-10, WRN 28-10. Mean + std. over 5 seeds. Late-phase BatchNorm (LPBN).
Model
Test acc. (%)
Base (SGD)
96. 96.16 10.122
Dropout (SGD)
96.02=0.06
MC-Dropout (SGD)
96.03 0.009
BatchEnsemble (SGD)
96.19=0.18
Late- (SGD)
96.46-0.15
Base (SWA) Late-phase (SWA)
96,48-1..0
96.81=0.07
Throughout our CIFAR-10/100 experiments we set K 10, use fast base gradient scale fac- tor of 7a= = 1, and set our late- thase initializa- tion hyperparameters to To henceforth in epochs; To = 100 for SWA) and do not use initialization noise, op = 0. These hyperparameters were tuned manually once CIFAR-100 and then kept fixed unless otherwise
Deep ensemble (SGD)
96.91
Deep ensemble (LPBN, SGD)
96.99
120 (measured
5

---

Published as a conference paper at ICLR 2021
noted. We use standard learning rate scheduling, optimized for SGD and SWA on the base model (cf. Appendices A and B). Last-layer weights are included by default in our late-phase weight set P.
CIFAR-10. For CIFAR-10 we focus on the WRN architecture, a work (WRN; Zagoruyko & Komodakis, 2016) which features BatchNorm layers. Taking advantage of this we implement a late-phase weight model consisting of BatchNorm shift and scale parameters.
residual net-
All algorithms achieve training error close zero (cf. Appendix B). The resulting predictive accuracies are shown in Table 1. We find that augmenting the WRN 28-10 (a standard WRN config- uration) with BatchNorm late-phase weights leads systematic improvement in generalization, reducing the gap with a deep ensemble of K 10 models. Initializing our ensemble from the onset (Tu 0) fails meet the performance of the base model, reaching only 95.63 + 0.23% (ef. Appendix 12).
We also investigate initializing a late-phase (full) deep ensemble at Tn 120. This results in test set accuracy of 96.32 + 0.09%, in between late-phase BatchNorm weights and no late-phase weights at all. This speaks to the data-efficiency of our low-dimensional late-phase ensembles which can be trained with as little data single model, besides being memory efficient.
In addition, we consider larger instance of the WRN model (the WRN 28-14), trained for 300 epochs using cutout data augmentation (DeVries & Taylor, 2017), as well as small convolution neural network without skip connections, cf. Table 3. When late-phase weights are employed in combination with SWA, observe significant accuracy gains the WRN 28- 14. Thus, our late- phase weights impose an implicit regularization that is effective on models with many weights. Similarly, we observe larger gains when training random subset of CIFAR-10 with only 10a examples (cf. Appendix B).
Table 2: Mcan CIFAR- 100 test set accuracy (%) + std. over 5 seeds, WRN 28-10. Different late- phase weight augmentations are compared to the base architecture and to an upper bound consisting of ensemble of models. Deep ens. stands for deep ensemble, LPBN for late-phase BatchNorm.
Base
BatchNorm
Hypernetwork
Deep ens.
Deep (LPBN)
SGD
81.35
82.87+0.22
81.55=0.31 82.01 10.17
84.09
84.69
SWA
82.46
83.06 10.08
83.62
CIFAR-100. We next turn to the CIFAR-100 100 dataset, which has 10-fold less examples per class and more room for improvements. We study the WRN 28-10, as well as the larger WRN 28-14 vari- ant (using cutout data augmentation before) and PyramidNet (Han et al., 2017) with ShakeDrop regularization (Yamada et al., 2019). The latter are trained for 300 epochs.
Predictive accuracy is again highest for neural networks with late-phase weights, trained with SGD or SWA, cf. Table 2. We observe that the simplest BatchNorm late-phase weight model reaches the highest accuracy, with late- late-phase hypernetwork weight embeddings yielding essentially im- provements. Once again, the setting of To = 0 (onset ensemble learning) fails to match base model performance, finishing 80.26 + 0.42% test accuracy. As for CIFAR- 10, late-phase full deep en- semble only reached intermediate improvements, at 82. 12.1710.159 test accuracy. Furthermore, gap towards deep ensembles persists. This suggests that covering different modes of the loss (Fort al., 2020) can provide benefits that cannot be captured by ensembling models in a small neighborhood.
The final averaged solutions found with late-phase weights are strong base models to build a deep ensemble of independently-traiee
Table 3: Additional architectures, CIFAR-10 (C10) and CIFAR-100 (C100). Mean test set acc. + std. over seeds (%). Late-phase BatchNorm weights.
networks. The fact that our algorithm yields a single model allows further pushing the upper bound of what can be achieved when un- restricted full ensemble training is possi- ble. This improvement comes at no cost compared to a standard deep ensemble.
Base
Late-phase
C10 ConvNet (SGD)
77.41 10.23
77.94
C10 WRN 28- (SWA)
96.75 40.05
97.45 97.45=0010
C100 WRN 28-14 14 (SWA)
84.01 10.20
85.00
C100 PyramidNet (SGD)
84.04
84.35

---

Published as a conference paper at ICLR 2021
We train additional neural network architectures restricting experiments to the BatchNorm late- phase weight model, which can be readily implemented without architectural modifications. Again, learning with late-phase weights yields consistent improvement over the baseline, cf. Table 3.
Notably, SWA can achieve high predictive accuracy with large constant learning rate (Izmailov al., 2018). We reproduce these previous results and show that they im-
Test acc. (%)
prove when learning with late-phase weights, cf. Fig. 2.
70
Substantial progress is made both when entering the late- phase learning period and when activating SWA.
K=1
60
K=10
-ut-of-distribution (OOD) generalization. Deep en- sembles effective technique for improving the behavior of neural networks in OOD data (Lakshmi- narayanan et al., 2017). We ask whether our implicit en- sembles modeled during late-phase learning could confer a similar advantage our final averaged model.
120
160
200
Training epoch
Figure WRN 28-10, CIFAR-100, constant learning rate SWA (activated at epoch 150). With BatchNorm late- phase weights (K=10, initialized epoch 120) and without (K=1). Mean acc. (%) + std. over 5 seeds.
Additionally, we evaluate the performance of late-phase weight ensemble obtained with large initialization noise Øo = 0.5 (at To 100), skipping the final weight averaging step. This requires integrating predic- tions over K late-phase ensemble members at test time, y(z) * Σk=1 y(z, wy). Unlike standard deep ensembles, training this ensemble is still as cheap as training a single model.
Table 4: CIFAR-100, 1R.-28-10, uncertainty representation results. Mean + std. over seeds (except deep ensembles). This first group of methods yield single model; the second group requires test-time averaging over models while training as efficiently as K=1; the last group are full deep ensembles which require training K=10 models from scratch (Deep ens.). We report in-distribution test set acc. (%) and negative log- logllikel (NLL), and in-distribution vs. out-of- out-oo-distribution (OOD) discrimination performance (average AUROC over four OOD datasets, see main text).
Test acc. (%)
Test NLL
OOD
Base (SGD)
81.35#00116
0.7400
0.8015#000189
Dropout (Mean) (SGD)
81.31 10.20
0.7736 10.0025
0.8022
Late-phase Hypernetwork (SGD)
81.55 00.32
0.8327=00.0m 20.0066
0.8209
Late-phase BatchNorm (SGD)
82.8740.14
0.7542 +0.0076
0.8360+t0.0116
MC-Dropout (SGD)
81.55=+0.11
0.7105
1..105==00000 00.000000000 0026
0.8225=t0.0185
SWAG (SWA)
12.12
0.6189 10.0036
0.8283110.0274
BatchEnsemble (SGD)
81.25 81.25
0.7691 10.0048
0.8285
0.8285:00.0155 L0000189
Late-phase BatchNorm (SGD, non-averaged)
82.71 10.10
0.75
2 000009
0.8624
00662410.0000 E0.0094
Deep ens. (SGD)
84.09
0.5942
0,8312
Deep ens. (Late-phase BatchNorm, SGD)
84.69
0.6712
0.8575
We draw novel images from a collection of datasets (SVHN, Netzer et al. (2011); LSUN, Yu et al. (2015); Tiny ImageNet; CIFAR- 10) and present them to WRN 28- 10 trained on CIFAR-100 100. We use Shannon's entropy (Cover Thomas, 2006) to measure the uncertainty in the output predictive distribution, which should be high for OOD and low for CIFAR- 100 data. Overall performance is summarized using the area under the receiver operating characteristics curve (AUROC), averaged over all datasets. We report per-dataset results in Appendix B (Table 16) alongside experiments measuring robustness to corruptions in the input data (Hendrycks Dietterich, 2019).
We compare results to alternative methods with strong uncertainty representation: MC-dropout (Gal & Ghahramani, 2016), SWA-Gaussian (SWAG: Maddox et al., 2019) and BatchEnsemble (Wen al., 2020). All three methods require integrating predictions over ensemble at test time.
We find that learning with late- thase weights increases prediction uncertainty in OOD data, allowing for significantly better separation between in and out-of-distribution examples, ef. Table 4. The

---

Published as a conference paper at ICLR 2021
OOD performance of late-phase BatchNorm weights compares favorably to the alternative methods including deep ensembles, even when using a single weight-averagee model, while maintaining high predictive accuracy. Remarkably, keeping the late-phase BatchNorm ensemble at test time allows reaching the highest OOD performance throughout. Paired with non-zero initialization noise dg > (cf. Appendix B), this method results in the best OOD performance.
Despite our improved performance on both predictive ac- curacy (with late-phase BatchNorm) and OOD discrim- ination (with late-phase BatchNorm and hypernetwork embeddings), the test set negative log- ligllikii (NLL; often used to assess predictive uncertainty. Guo et al., 2017) surprisingly slightly worse for our solutions. This is aligned with the finding that SWA does not al- ways significantly reduce NLL, even though predictive accuracy increases (Maddox et al., 2019).
15
K=1
K=10
8L/L
10
5
0.00
0.05
0.10
0z
Figure 3: Flatness score. Mean score + std. over 5 seeds, WRN 10,10, CIFAR- 100, SGD, with and without Batch- Norm late-phase weights. Slower in- crease with ag is better.
Flatness. Why our networks generalize better? Ap-
proximate Bayesian inference suggests that flat min- ima generalize better than sharp minima (Hochreiter & Schmidhuber, 1997a; MacKay, 1992). Due to symme- tries that are present in neural networks there is some debate surrounding this argument (Dinh et al., 2017), but current evidence seems favorable (Jiang et al., 2020).
We hypothesize that sharing base weights over late-phase weight configurations can implicitly lead to flatter solutions. To investigate whether algorithm finds flatter minima, examine a simple flatness score that correlates well with generalization (Pittorino et al., 2020; Jiang et al., 2020). Concretely, we add multiplicative Gaussian noise 2; ~ (O, wf of) to each weight we and then measure the change in the loss 8L = Ex[[(w + z) [ww)]. Our final weight configurations are indeed in flatter regions of weight space according to this measure: őL increases more slowly with for the WRN 28- 10 models that are learned with BatchNorm late-phase weights, Fig. 3.
3.3 IMAGENET EXPERIMENTS
To investigate whether our gains trans- late to large-scale learning problems, we train deep residual networks (He et al., 2016) and a densely-connected convolu- tional network (DenseNet; Huang al., 2018) on the ImageNet dataset (Rus- sakovsky al., 2015). We start from pre- trained models and contrast BatchNorm late-phase weight learning to fine-tuning with SGD for 20 epochs, with 7e 1/K and og = 0 (cf. Appendix A). For simplicity we do not include last-layer weights in Φ.
Table 5: Validation set acc. (%) on ImageNet. Mean + std. over 5 seeds. BatchNorm late-phase and baseline trained for 20 epochs with SGD.
Initial
Base
Late-phase
ResNet-50
76.15
76.62
76.87+0.03
ResNet-152
78.31
78.37
78.77
DenseNet-161
77.65
78. 78.1720..01
78.31 $0.01
Fine-tuning with late-phase weights improves the final top-1 validation accuracy of this pretrained model significantly with only minor training, as seen in Table 5. These results serve as a proof-of- concept that existing models be further improved, taking time the previous experimenter stopped training. In Appendix B, we present additional CIFAR-100 experiments where we apply late-phase learning starting at the suboptimal end-of-training To 200, to mimic the pretrained condition.
late late-phase initialization To as the
3.4 LSTM LANGUAGE MODELING EXPERIMENTS
Finally, we conduct experiments on the language modeling benchmark enwik8. To show that the benefits of late-phase weights extend to recurrent neural networks, we augment a standard LSTM with multiplicative late-phase weights consisting of rank-1 matrices (Wen et al., 2020, cf. Sec- tion 2.2).
8

---

Published as a conference paper at ICLR 2021
Overfitting is major issue when training LSTMs. Recent studies have shown that by leveraging vast amounts of computation and smart black-box optimizers (Golovin et al., 2017), properly regularized LSTMs can outper- form previously published state-of-the-art mod- els (Melis al., 2017). To avoid this issue, we train models where the number of parameters (~1.56M) is drastically smaller than the num- ber of training data points (90M), such that we do not observe any overfitting. Thus, we do not apply any regularization. This helps minimize the effects of hyperparameter tuning. Our only hyperparameter is the learning rate (0.001 here), which we tune via grid search to maximize base model performance.
Table 6: enwik8 results measured bits per
character (BPC), LSTM with 500 units. Mean over 5 seeds, with std. o < 0.01 for all results.
Model
Train
Test
Test (SWA)
Base
1.570
1.695
1.626
Base + Rank1
1.524
1.663
1.616
Late-phase Rank1
1.522
1.633
1.615
We train our LSTM with 500 units for 50 epochs, optimizing every weight with Adam (Kingma & Ba, 2015). We apply a multiplicative rank-l matrix elementwise to the recurrent weight matrix. Interestingly, merely adding the multiplicative parameters the LSTM (Base) accelerates training and leads to better training and test set performance (measured in bits per character, BPC) with no additional changes the optimizer (Base + Rank1, Table 6). Further improvements with our late-phase weights. We generate K = 10 late-phase weight components at epoch 30 with To 0.35 and set 78 Additionally, we find that SWA (starting epoch 40) substantially improves all scores, with smaller gains on the models with multiplicative weights.
be achieved
RELATED WORK
Our late-phase weights define an ensemble with the special property that every model shares the same base weights. Such parameter sharing an established method for ensembling neural networks while controlling for the memory and time complexity of learning (Lee et al., 2015). In designing our late-phase weight models, we draw directly from recent work which proposes sharing set of base parameters over K rank-1 matrices (Wen et al., 2020) or K heads (Lee et al., 2015).
The elastic averaging SGD algorithm learns K neural networks in parallel, coupled through an additional central model (EASGD; Zhang et al., 2015). Like our algorithm, EASGD often yields solutions which generalize better than those found by standard SGD (Pittorino et al., 2020). Our late- phase weight learning is intimately related to EASGD, as we optimize the performance of a central model through an ensemble. However, thanks to parameter sharing and late-phase ensembling, we replicate
do not find the need to introduce coupling term to our loss function. Additionally, small number of parameters only, the complexity of our algorithm is greatly reduced in comparison to EASGD, which requires learning a full ensemble of models.
Splitting the weights of a neural network into a set of fast and slow components which vary on dif- ferent timescales is a classic technique (Hinton & Plaut, 1987; Schmidhuber, 1992) that has proven useful in a wide range of problems. This list includes applications to few-shot learning (Munkhdalai & Yu, 2017; Nichol al., 2018; Perez et al., 2018: Zintgraf et al., 2019; Flennerhag et al., 2020), optimization (Zhang et al.. 2019b; Chaudhari et al., 2019), improving recurrent neural networks (Ba al., 2016; Ha al., 2017). and continual learning with biologically
synapses (Kapla- nis et al., 2018; Leimer et al., 2019), to name a few. Although there is no explicit separation of timescales in our weight components, the update accumulation in as Φu varies (cf. Algorithm 1) suggests interpreting the base 0 as slow weights and the late-phase & as fast weights.
This accumulation is reminiscent first separates parameters into task-shared and task-specific, and then differentiates through a se- quence of accumulated updates performed over the tasc-fpe parameters (Finn et al., 2017). Continuing with the fast-slow weight analogy, our averaging over fast weights at the end of learn- ing (Eq. 1) could be thought of as a synaptic consolidation step which integrates the fast weight components onto a slow, persistent form of memory.
recent meta-learning algorithm (Zintgraf et al., 2019), which

---

Published as a conference paper at ICLR 2021
5 CONCLUSION
We proposed to replicate and learn in parallel subset of weights in a late phase of neural net- work learning. These late-phase weights define an ensemble of models which share every other weight. We studied convolutional neural networks, common recurrent neural network, and a sim- ple quadratie problem. Surprisingly, across these cases, we found that a small number of appropri- ately chosen such weights can quickly guide SGD towards solutions that generalize well. Most of our experiments relied on BatchNorm late-phase weights, making our method easy to implement in wide range of existing models, including pretrained ones. We expect future work to uncover new effective late-phase weight models.
ACKNOWLEDGEMEETT
This work was supported by the Swiss National Science Foundation (B.EG. CRSI15-17377 and 315230_189251) ETH project funding (B.F.G. ETH-20 19-01), the Human Frontiers Science Pro- gram (RGY0072/2019) and funding from the Swiss Data Science Center (B.E.G, C17-18, J.v.O. P18-03). João Sacramento was supported by an Ambizione grant (PZ00P3.186027) from the Swiss National Science Foundation. We would like to thank Nicolas Zucchet, Simon Schug, Xu He, Ángelo Cardoso and Angelika Steger for feedback, Mark van Rossum for discussions on flat min- ima, Simone Surace for his detailed feedback on Appendix C, and Asier Mujika for providing very useful starter code for our LSTM experiments.
REFERENCES
Martin Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: a system for large-scale machine learning. In /2th USENIX Symposium on Operating Systems Design and Implementa- fion, 2016.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: implicit acceleration by overparameterizatton. In International Conference Machine Learning, 2018.
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin lonescu. Using fast weights to attend to the recent past. In Advances in Neural Information Processing Systems 29. 2016.
Carlo Baldassi, Fabrizio Pittorino, and Riccardo Zecchina. Shaping the learning landscape in neural networks around wide flat minima. Proceedings of the National Academy of Sciences, 117(1): 161-170, January 2020.
Léon Bottou. Large- Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT'2010, pp. 177-186. Springer, 2010.
Kevin S. Brown and James P. Sethna. Statistical mechanical approaches to models with many poorly known parameters. Physical Review E, 68(2):011944, August 2003.
Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks. In Information Theory and Applications Workshop (ITA). IEEE, 2018
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SCH SGDD: Biasing gradi- ent descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019
1240:1 2019.
Thomas M. Cover and Joy A. Thomas. Elements of Information Theory. Iitte--ntersciencee, USA, 2006.
10

---

Published as a conference paper at ICLR 2021
Terrance VeViies and Graham W. Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: November 2017.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. International Conference on Machine Learning, pp. 1019-1028, August 2017.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, pp. 1126 1135, July 2017.
Sebastian Flennerhag, Andrei A. Rusu, Razvan Pascanu, Francesco Visin, Hujun Yin, and Raia Hadsell. Meta-learning with warped gradient descent. In International Conference on Learning Representations, 2020.
Stanislav Fort, Huiyi Hu, and Balaji Lakshminarayanan. Deep ensembles: loss landscape perspec- tive. arXiv preprint arXiv: June 2020.
Jonathan Frankle, David J. Schwab, and Ari S. Morcos. Training batchnorm and only batchnorm: on the expressive power of random features in CNNs. arXiv preprint orXiv:2003.0015 June 2020.
Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: representing model uncertainty in deep learning. In International Conference on Machine Learning, pp. 1059, June 2016.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of DNNs. In Advances in Neural Information Processing Systems 31. 2018.
Daniel Golovin, Benjamin Solnik, Subhodeep Moitra, Greg Kochanski, John Karro, and D. Sculley. Google vizier: a service for black-box optimization. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. 2017.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International Conference on Machine Learning, 2017.
David Ha, Andrew M. Dai, and Quoc V. Le. Hypernetworks. In International Conference on Learning Representations, 2017.
Dongyoon Han, Jiwhan Kim, and Junmo Kim. Deep pyramidal residual networks. In Conference on Computer Vision and Pattern Recognition, 2017. arXiv: 1610.02915.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: surpassing human-level performance on ImageNet classification. arXiv preprint (rcXiv::502.018:: February 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. pp. 770-778, 2016.
Dan Hendrycks and Thomas Dictterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. In International Conference Learning Representations, 2019.
Geoffrey E Hinton and David C Plaut. Using fast weights to deblur old memories. In Proceedings of the ninth annual conference of the Cognitive Science Society, pp. 177-186, 1987.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, January 1997a.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9()): 1735-1780, November 1997b.
Gao Huang, Yixuan Li. Geoff Pleiss, Zhuang Liu, John E. Hopcroft, and Kilian Q. Weinberger. Snapshot ensembles: train 1, get M for free. In International Conference on Learning Represen- fations, March 2017.
11

---

Published as a conference paper at ICLR 2021
Gao Huang, Zhuang Liu, Laurens van der Maaten, and Kilian Q. Weinberger. Densely connected convolutional networks. arXiv preprint arXiv: 1608. 06993, January 2018.
Sergey loffe and Christian Szegedy. Batch normalization: accelerating deep network training by reducing internal covariate shift. In International Conference Machine Learning, 2015.
Kiyosi Itô. On stochastic differential equations. Number 4. American Mathematical Soc., 1951.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Uncertainty in Artificial Intelligence, 2018.
Siddhant M. Jayakumar, Wojciech M. Czarnecki, Jacob Menick, Jonathan Schwarz, Jack Rae, Si- mon Osindero, Yee Whye Teh, Tim Harley, and Razvan Pascanu. Multiplicative interactions and where find them. In International Conference on Learning Representations, 2020.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker-Planck equation. SIAM Journal on Mathematical Analysis, (9(1):1-17, 1998. Publisher: SIAM.
Christos Kaplanis, Murray Shanahan, and Claudia Clopath. Continual reinforcement learning with complex synapses. In International Conference Machine Learning. June 2018.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: generalization gap and sharp minima. In International Conference Learning Representations, 2017.
Diederik P. Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Andrey Nikolaevich Kolmogorov. On analytical methods probability theory. Math. Ann, 104: 415-458, 1931.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems 30. 2017.
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied document recog- nition. Proceedings of the IEEE, 86(1
November 1998.
Kimin Lee, Kibok Lee, Honglak Lec, and Jinwoo Shin. A simple unified framework for detecting diittr--iistiittioonn samples and adversarial attacks. Advances in Neural Information Process- ing Systems 31. 2018.
Stefan Lee, Senthil Purushwalkam, Michael Cogswell, David Crandall, and Dhruv Batra. Why M heads are better than one: training diverse ensemble of deep networks. arXiv preprint: arXiv:1511.06314 November 2015.
Pascal Leimer, Michael Herzog, and Walter Senn. Synaptic weight decay with selective consolida- tion enables fast learning without catastrophic forgetting. bioRxiv, pp. 613265, April 2019.
Guan-Horng Liu and Evangelos Theodorou. Deep learning theory review: An optimal control and dynamical systems perspective. arXiv preprint arXiv: 10920, 2019.
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In Inter- national Conference on Learning Representations, 2017.
David J. C. MacKay. A practical Bayesian framework for backpropagation networks. Neural Com- putation, 4(3):488 472, 1992.
12

---

Published as a conference paper at ICLR 2021
Wesley J Maddox. Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for Bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems 32. 2019.
James Martens. Second-orter optimization for neural networks. PhD thesis, University of Toronto, 2016.
Gábor Melis, Chris Dyer, and Phil Blunsom. On the state of the art of evaluation in neural language models. arXiv preprint arXiv:170.00558, November 2017.
Pramod Kaushik Mudrakarta, Mark Sandler, Andrey Zhmoginov, and Andrew Howard. K for the price of I: parameter-efficient multi-task and transfer learning. In International Conference on Learning Representations, 2019.
Tsendsuren Munkhdalai and Hong Yu. Meta networks. In International Conference on Machine Learning, June 2017.
Yurii Nesterov. Introductory lectures on convex optimization: a basic course. Springer US, 2004.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning, 2011.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint arXiv:1803.0290 022999, October 2018.
Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: visual reasoning with a general conditioning layer. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, pp. 3942-3951, 2018.
Fabrizio Pittorino, Carlo Lucibello, Christoph Feinauer, Enrico M. Malatesta, Gabricle Perugini, Carlo Baldassi, Matteo Negri, Elizaveta Demyanenko, and Riccardo Zecchina. Entropic gradient descent algorithms and wide flat minima. arXiv preprint arXiv:2006.07897 June 2020.
Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4):838-855, 1992.
Sollvstre-AAlviee Rebuffi, Hakan Bilen, and Andrea Vedaldi. Learning multiple visual domains with residual adapters. Advances in Neural Information Processing Systems 30, pp. 506-516. 2017.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma. Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael S. Bernstein, Alexander C, Berg, and Fei-Fei Li. ImageNet large scale visual recognition challenge. International Journal of Computer Vision, 115(3):211 1-222 2015.
Levent Sagun, Utku Evci, V. Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the Hessian of wver-parmeetrized neural networks. arXiv preprint arXiv:1706. 144.0.4400.4444, May 2018.
Pedro Savarese and Michael Maire. Learning implicitly recurrent CNNs through parameter sharing. In International Conference on Learning Representations, 2019.
Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In International Con- ference Machine Learning, pp. 343-351, 2013.
Jürgen Schmidhuber. Learning to control fast-weight memories: an alternative to dynamic recurrent networks. Neural Computation, 4(1): January 1992.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting, The Journal of Machine Learning Research, 199)::199--1988 2014.
Johannes von Oswald, Christian Henning, João Sacramento, and Benjamin F. Grewe. Continual learning with hypernetworks. In International Conference on Learning Representations, 2020.
13

---

Published as a conference paper at ICLR 2021
Joshua J. Waterfall. Fergal P. Casey, Ryan N. Gutenkunst, Kevin S. Brown, Christopher R. My- ers, Piet W. Brouwer, Veit Elser, and James P. Sethna. Sloppy-model universality class and the Vandermonde matrix. Physical Review Letters, 97(15))15060, October 2006.
Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient Langevin dynamics. In International Conference Machine Learning, 2011.
Yeming Wen, Dustin Tran, and Jimmy Ba. BatchEnsemble: an alternative approach to efficient ensemble and lifelong learning. In International Conference on Learning Representations, 2020.
Yuhuai Wu, Mengye Ren, Renjie Liao, and Roger Grosse Understanding short-horizon bias in stochastic meta-optimization. International Conference on Learning Representations, March 2018.
Sho Yaida. cettuaationdissppptiin relations for stochastic gradient descent. arXiv preprint 1rXivv1810.00000 2018.
Yoshihiro Yamada, Masakazu Iwamura, Takuya Akiba, and Koichi Kise. Shakedrop regularization for deep residual learning. IEEE Access, 7:186126-1861366 2019.
Fisher Yu, Yinda Zhang, Shuran Song, Ari Seff, and Jianxiong Xiao. LSUN: construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv:1506.03:30 03366, 2015.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedings of the British Machine Vision Conference, 2016.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter which batch sizes? Insights from a noisy quadratic model. In Advances in Neural Information Processing Systems 32, pp. 8196-8207. 2019a.
Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead Optimizer: k steps forward, step back. In Advances in Neural Information Processing Systems 32, pp. 9597-9608. 2019b.
Sixin Zhang, Anna E Choromanska, and Yann LeCun. Deep learning with elastic averaging SGD. In Advances in Neural information Processing Systems, pp. 685-693, 2015.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise stochastic gradient descent: its behavior of escaping from sharp minima and regularization effects. arXiv preprint 1XX3::88030001900 2018
Luisa Zintgraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning In International Conference on Machine Learning, pp. 7693-7702,
May 2019.
A ADDITIONAL IMPLEMENTATION DETAILS
Hypernetwork model. The base neural network architecture we use when parameterizing our weights using hypernetwork is identical to the WRN 28- 10 described by Zagoruyko & Komodakis (2016). Our hypernetwork implementation closely follows Savarese Maire (2019), who studied high-performing linear hypeeeetwrrk architectures for WRNs. We do not use dropout or biases in the convolutional layers. The parameters of every convolutional layer are hypemetwork hypemetwork-generateed generated, with one hypernetwork per layer group (Table 7). The remaining parameters, namely those of BatchNorm units and final linear layer weights, are
Snypernetworr-geneeeee
Following Savarese & Maire (2019) turn off weight decay for the model embeddings and initial- ize these parameters with a random pseudo-orthogome initialization over layers. The hypernetwork parameters initialized using standard Kaiming initialization (He al., 2015).
14

---

Published as a conference paper at ICLR 2021
Table 7: Specification of the hypernetwork used each convolutional layer of the WRN, indexed by its depth in the network. A depth marked by refers to the residual connection spanning across the specified layers. The characteristics of each layer described in the format input-channels x [kernel-size] output-channels under Conv-layer. Layers within the same group generated by the same hypernetwork. Each hypernetwork has a unique parameter tensor of shape Hner-PS, which, when multiplied by layer and weight embedding of shape Emb-PS and reshaped appropriately, generates the primary network parameter of shape Base-PS.
Depth
Conv-layer
Base-PS
Layer group
Hnet-PS
Emb-PS
3x[3x3]x16
[16, 3, 3]
[16, 3, 3, 3, 10]
[10, 1]
160
[160, 3, 16]
1
[160, 3, 16,
[7, 1]
3
x60x[3x3]
[160, 3, 160] 2
[160, 3, 80, 14]
[14, 2]
4
160x
1160,3,3, 160]
2
[160, 3, 14]
[14.2]
5
160x[3x3] 160
1160,3,3, 160]
2
3,3,,,3,,80 14]
[14,2]
6
160x[3x3] 160
3,3, 160] 2
3,3, 80, 14]
[14.2]
7
160x[3x3] 160
[160, 3,3,1660
2
3, 14]
[14,2]
8
160x(3x3]x 160
[160, 3, 160]
2
[160, 3, 80, 14]
[14, 2]
1603
[160, 3, 160]
2
[160, 3, 80, 14]
[14, 2]
10
160x
[320, 3, 160]
3
[320, 3, 3, 160, 14]
[14,
11
xx3
[320, 3, 320]
3
[320, 3, 160, 14]
[14, 2]
12
x3)
[320, 3,
3
[320, 3, 160, 14]
[14.2]
13
x
[320, 3,
3
[320, 3,3, 160, 14]
[14.2]
14
320x[3x33x320
[320, 3,
3
[320, 3,3, 160, 14]
[14.2]
15
320x(3x33xx320
[320, 3, 3, 320]
3
[320, 3, 160, 14]
[14,2]
16
320x(3xx3xx320
[320, 3, 3, 320]
3
[320, 3, 3, 160, 14)
[14,2]
17
x303
[320, 3, 3, 320]
3
[320, 3, 3, 160, 14]
[14, 2]
18
320x
[640, 3, 3, 320]
4
[640, 3, 14]
[14,
19
640x(3x3]x644)
[640, 3, 640]
4
[640, 0,30,, 3, 14]
[14,2]
20
640x(3x3)x64D
[640, 3, 640]
4
[640, 3,3,320, 14]
[14.2]
21
640x(3x3]x640
[640, 3,3, 640]
4
3,3,320 14]
[14.2]
22
640x(3x3]x640
[640, 3,3, 640]
4
[640, 3,0, 14]
[14.2]
23
640x(3x3)x640
[640, 3, 640]
4
[640, 3, 14]
[14,2]
640x(3x3)x640
[640, 3, 640)
4
[640,
14)
[14,2]
640x(3x3)x640
[640, 3, 3, 640) 4
[640, 3, 3, 320, 14]
[14, 2]
2-+4*
16x[1x1]x160
[160,1, 1, 1, 16]
5
[160, 1, 1, 16, 7]
[7. 1]
10-+12#
160x[1x1]x320
[320. 1, 1, 160]
6
[320, 1, 1, 160, 7]
[7.
18 *0
xx0xx1x1]
[640, 1, 320]
7
[640, 1, 320,
[7, 1]
15

---

Published as a conference paper at ICLR 2021
Small ConvNet model. We train a slight modification of the classic LeNet-5 (Lecun et al., 1998) for 200 epochs on CIFAR-10. Both convolutional and fully-connected layers are left unchanged, but use rectified linear units the hidden layers. Furthermore, after each such activation, Batch- Norm units are inserted. We optimize the model with SGID and use late-phase BatchNorm weights, with To 50 and do 0.5. For simplicity of implementation, we do not include the last linear layer in the late-phase weight set 2.
Optimization. We optimize the cross-entropy loss, using either SGD with Nesterov momentum (0.9) or SGD with Nesterov momentum (0.9) wrapped inside SWA. LSTM: Our LSTM experiments use Adam with constant learning rate 0.001, batch size 128, and no regularizers such as weight decay or dropout. WRN-28-10: For our WRN experiments on the CIFAR datasets we use the learning rate annealing schedule of Izmailov et al. (2018), according to which an initial learning rate of 0.1 linearly decreased at every epoch from the end of the 100th epoch (80th for SWA) to the end of the 180th epoch (144th for SWA; SWA is activated at epoch 160), when final value of 0.001 (0.05 for SWA) is reached. Our optimizers use Nesterov momentum (set to 0.9), a batch size of 128 and weight decay (set to 0.0005). On CIFAR-100 (SGD) we set the weight decay of late- late-phase weights proportional to the ensemble size, 0.0005K.. WRN-28-14: The WRN 28-14 models are trained for 300 epochs CIFAR-100. The learning rate is initialized at 0.1, then annealed to 0.05 from the 80th epoch to the 240th epoch. SWA is activated at epoch 160. All other hyperparameters are identical to those of WRN 28-10. ConvNet: Same for the WRN 28-10 10 model, except that anneal the learning rate until the 160th epoch.
Algorithm 2: Stochastic weight averaging (SWA)
Algorithm 3: SGD with Nesterov mo- mentum
Require: Base weights 0, dataset D,
Require:
Base weights 0 dataset D, learning rate momentum
hyperparameter y), loss &
Require: Training iteration t
loss L
M Sample minibatch from D
M < Sample minibatch from DD
AO < Va C(M, 0)
A < C(M, + pv)
0 U 17, 20)
V pv A0
SSWA+ + (t (swA + (++++)
0 t 0 +
Figure 4: Pseudocode for a single parameter update for SWA and SGD with Nesterov momentum, the two main optimizers used in our experiments. These are either used standalone, or Ug and Ug in Algorithm 1 (main text). U in Algorithm 2 (SWA) serves as placeholder for a parameter update rule such as SGD (with Nesterov momentum) or Adam. Training iteration ł is counted from the activation of SWA in Algorithm 1.
Batch normalization units. Whenever we use SWA, we follow Izmailov al. (2018) and per- form full pass over the training set to re-estimate BatchNorm unit statistics before testing. This correction is required since the online BatchNorm mean and variance estimates track the activations produced with the raw (non-averaged) weights during training, while the averaged solution is the one used when predicting at test time.
Data augmentation and preprocessing. On both CIFAR and ImageNet datasets, all images are normalized channelwise by subtracting the mean and dividing by the standard deviation; both statis- ties are computed on the training dataset. The same transformation then applied when testing, including to OOD data. Following standard procedure (e.g., Zagoruyko & Komodakis, 2016; He et al., 2016) we augment our training datasets using random crops (with 4-pixel padding for CIFAR) and random horizontal flips. The ImageNet training dataset is augmented with random hor- izontal flips, as well as random cropping of size 224, while a centered cropping of size 224 was used on the test set. Our OOD datasets resized to fit whenever necessary; we used the resized images made available by Lee et al. (2018).
ImageNet experiments. The pretrained model for the ImageNet experiment is obtained from torchvision's 's models subpackage. We fine-tune the model for 20 additional epochs ImageNet.
16

---

Published as a conference paper at ICLR 2021
Algorithm 4: Late-phase learning
Require: Base weights 0, late-phase weight set Φ, dataset D, gradient scale factor "/o learning rate y, ensemble size K, initialization noise do. initialization time To, number of training iterations T, loss
Initialization: K - 0, - 0,1+1
while f T do
if = To then
I/ generate late-phase weights
for 1 1 k v do
sample (0,
Φ% + @ +
To
db
1/ set range for specialists training
K K
++1
for 8 < k S K do
M₁ t- Sample minihatch from D
A0 Ve & 0, dx)
Φh t ok 11 V on L(MR, 0, db)
t x t+1
0 4 0 7
Ek= k= 202
K
Figure 5: Complete pseudocode for an entire training session using late-phase weights. To avoid notational clutter T, To and are measured numbers of minibatches consumed. In the paper, we measure To and T in epochs. For simplicity, SGD (without momentum) and Φh of dimension 1. Other optimization algorithms (e.g., Algorithm 2 Algorithm can be used to replace Uş and Up. as described in Algorithm 1. Note that we increase inside the inner loop. This highlights (i) that every specialist parameter is trained only on 1/ K data samples after To compared to 6, and (ii) that we count every minibatch drawn from the data to compare fairly to algorithms without an inner loop.
present the case where Ug and Up are set to plain
We use multistep learning rate scheduler, starting at 0.001 then decreasing at the 10th epoch to 0.0001. We use SGD with momentum (set to 0.9) and weight decay (set to 0.0001) as our optimizer, with batch size of 256. We use op and K 10 for late-phase model.
Code forks. Our hypernetwrk implementation was inspired by the code made publicly available by Savarese & Maire (2019). Our implementation of SWA was adapted from the code accompanying the work of Izmailov et al. (2018), now available on the torchcontrib Python package. The SWAG method was evaluated directly using the code provided by the authors (Maddox et al., 2019). We used the same base WRN model as Maddox et al. (2019), which can be retrieved from https: / /githhb. com/ me meliketoy weckeetoywwieeeee http:/wddderreeeeeeeeeeeee resnet pytorch.
Table 8: CIFAR-100 set accuracy (%) depending on different values of K for WRN 10, SGD. Mean t std. over 5 seeds.
K
Test acc. (%)
81.35 16
5
82.44+0.22
10
82.8770.22
15
83.01 #0.27
20
82.86
LSTM All experiments are conducted using the Tensorflow Python framework (Abadi et al., 2016).
All base weights are initialized uniform [-0.01, 0.01] whereas the initial rank-1 matrix weights are centered around 1 i.e. [1 -0001 1.000 to mitigate strong difference in initialization compared
to the base model. We use the Tensorflow default values (3) = 0.9, 32 = 0., e = 10 -8 for the Adam optimiser. We perform grid search over dg E [0, 0.5] (in steps of size 0.05) for our LSTM experiments (fixing K 10 and varying To E {0, 30}}) and obtain the values reported in the main
text, To 30 and do 0.35.
17

---

Published as a conference paper at ICLR 2021
0.850
Accuracy (%)
80
0.825
AUROC
0.800
Test set
AUROC
75
Base Test set
Base AUROC
0.775
0
40
80
120
160
200
To
Figure 6: Sensitivity analysis of To- Mean AUROC score (OOD) and test set accuracy for different values of To for WRN 28-10, CIFAR-100, SGD, with BatchNorm late-phase weights.
B ADDITIONAL EXPERIMENTS
Table 9: Applying late-phase weights pretrained WRN 28-10, CIFAR-100, SGD, Mean + std. over 5 seeds.
Model
Test acc. (%)
Initial
81.35=0.16
Base
81.47=0.11
Late-phase BatchNorm
82.02 10.12
Late-phase BatchNorm, frozen base weight 81.50=0.220
Pretrained CIFAR-100. We apply method standard WRN 28-10 pretrained CIFAR- 100 (i.e., we set To 200) and train for an additional 20 epochs. At the beginning of the fine- tuning, the learning rate reset 0.01, then annealed linearly to 0.001 for 10 epochs. It is then held constant for the remainder of the fine-tuning process. We observe that augmenting with BatchNorm late-phase weights yields an improved predictive accuracy compared additional fine-tuning with SGD (Base), cf. Table 9. Both methods improve over the initial baseline (Initial), including the base model. This can be explained by the optimization restart and the accompanying spike in the learning rate introduced by our scheduler (Loshchilov & Hutter, 2017).
Importantly, we find that fine-tuning only BatchNorm late-phase weights while keeping all other weights fixed does not even match the Base control. Together with the finding that the optimal late- phase weight initialization time is at Tj 120 (when learning for 200 epochs), this result speaks to the importance of jointly optimizing both base and late-phase weights through our Algorithm 1.
Gradient accumulation control. Here we show that the improved generalization we report in the main text is not merely due to gradient accumulation over larger batches. We take our base WRN 28- 10 model (without late-phase weights) and start accumulating gradients over K = 10 minibatches To 120, experimenting both with 79 1/K and to 1. The models are trained with SGD using otherwise standard optimization settings. Both controls fail improve (even match) the performance of the base model trained without any gradient accumulation.
Sensitivity to To, K and ag. We present hyperparameter exploration on the CIFAR-100 dataset using BatchNorm late-phase weights in Tables 8, 11 and 12. We find that our algorithm is largely
Table 10: Gradient accumulation control, CIFAR-100, WRN 28- 10, SGD. Mean + std. over 5 seeds.
Model
Test acc. (%)
Base (SGD)
81.35=0.16
Base gradient accumulation (70 = 1)
80.76 2.22
Base gradient accumulation (70 = 1/K)
80.34
18

---

Published as a conference paper at ICLR 2021
Table 11: CIFAR-100 test set accuracy (%) depending on different values of do for WRN 28-10 SGD with late-phase BatchNorm weights (LPBN). Mean + std. over 5 seeds.
CIFAR-100
CIFAR-100
CIFAR-100 (LPBN, pretrained)
(LPBN)
(LPBN, non-averaged)
d
Test acc. (%)
OOD
Test acc. (%)
OOD
Test acc. (%)
00D
0
82,87+0.22
0.833
83,20=00.20
0.854#0.017
81.70
0.803 00.017
0.25
82.77=0.19
0.836#0.012
82.68=0.32
0.861 #0.013
82.02=0012
0.808 00.
0.5
82.78
18.18
0.83740.011
82.71
0.862
81.15 10.29
0.797 0.0007
0.75
82,41
2:020
0.839:10.012
82,43 10.15
0.855 18.5110.013
1.0
81.52 1.00
0.84040.017
82,38 82.38-00.15
1..08400.014
robust to do when To can be set to its optimal value, which is at 60% of training. See also Figure for visualisation of the same data, specifically the change mean AUROC score and test set accuracy when changing To. This result holds also on CIFAR-10, ef. Table 12. When starting from pretrained condition (To = 200), finite To leads to a significant improvement in performance, cf. Ta- ble 11. We therefore report results obtained with de for every CIFAR and ImageNet experiment in the main text. The exception to this is the non non-avergged (ensemble) late late-phase BatchNorm model presented in Table 4, which was optimized for best OOD performance (corresponding to on 0.5).
Table 13: Performance of a WRN 28-10 on 100 with different dropout probability p. For MC-dropout we average over 10 different samples. Mean + std. over 5 seeds.
p
Test acc. (%)
Test NLL
OOD
Dropout
0.1
40.13
0.7476-t0.005
0,8031 1000064
Dropout
0.2
81.31
0.7736:40.0025
0.8022#000290
Dropout
0.3
80.93
0.8342
0.7833+00.0239
MC-Dropout
0.1
81.51+0.14
0.7197 0.7197=#.0054 0005-1
0.8 0.8149=00.00**
MC-Dropout
0.2
81.55
0.7105 10.0026
0.82251:0.040
MC-Dropout
0.3
81.36 10.31
4.00009
0.8040
Related work. Here we provide details for the training setups of alternative methods we compare against in the main text. For the results reported for dropout (Srivastava et al., 2014) and MC-dropout (Gal & Ghahramani, 2016), we simply train WRN 28-10 on CIFAR-I00 with the exact same configu- ration for our base model, see above, but include dropout layers as usually done (Zagoruyko & Ko- modakis, 2016) after the first convolution in each residual block. For a scan over the dropout prob- ability p this setup, see Table 13. p 0.2 is re- ported in the main text for CIFAR- 100 and CIFAR- 10. Note that p was only tuned for CIFAR-100.
Table 12: CIFAR-10 and CIFAR-100 test set accuracy (%) depending different late phase timing To for WRN 28-10, SGD. Mean
+ std. over 5 seeds.
To
CIFAR-10
CIFAR-100
95.68-10.23
74.38=10.71
40
96.34
79.69=#0.11
60
96.42=0.10
80.53=0.21
80
96.50=t0111
01.72
100
96.45
82.48=40.21
120
96.48=00.20
82.87=00.22
For the reported results of BatchEnsemble (Wen al., 2020), we simply execute the code pro- vided by the authors at https::/githiei google/ bonggee/nneeteee fuunerttiinne wwg.-baaeeinng fine-tuned configuration for CIFAR-10/100. No- tably, the authors use a different setup than followed in this manuscript. First, the WRN 28-10 is trained for 250 epochs (we allow for this increased budget exceptionally for BatchEnsemble),
140
96.26
02.5
160
96.23
81.41 1.0.11
com/
180
96.25
81.43
1.0.2
inteerttintyy--beeiinee lnnes with their
200
96.16
81.35 #:0.6
with a multi-step learning rate annealing at [80, 160, 180] with a learning rate decay factor of 0.2.
Second, weight decay of 3 x 10-1 is used.
19

---

Published as a conference paper at ICLR 2021
Table 14: Final training set loss on CIFAR datasets, WRN 20-10, SGD. Mean std. over 5 seeds.
Training loss
CIFAR-10
base
1.00100...000 0.0019=00.000
CIFAR-10 late-phase BatchNorm
CIFAR-100 base
0.0024+0000011 0.0267=#0.000
CIFAR-100 late-phase BatchNorm
Table 15: Performance of models trained on a reduced CIFAR-10 training set and evaluated on the full CIFAR-10 test set. Mean + std. over seeds.
Model
Test acc. (%)
Base (SGD)
88.98 90.18
Late-phase BN (SGD)
89.58 10.19
For the results reported for SWAG (Maddox et al., 2019), we use the code provided by the authors and the proposed fine-tuned config-
at https::/// / thub. com/w /ttp:///w/o///////////oo//////// ddox/ uration which coincides with the configuration used to obtain all CIFAR-100 results reported in this manuscript, except for BatchEnsembles (see above). We report results for SWAG after training on 200 epochs for fair comparison.
Training losses. We provide the final achieved training losses for the base model and when aug- menting with BatchNorm late-phase weights on Table 14, for both CIFAR- 10 and CIFAR-100 Us- ing a fast gradient accumulation scale factor of T leads to a higher training loss on CIFAR-100 than that of the standard model, but found this setting crucial to achieve the largest improvement on test set generalization.
CIFAR-10 with reduced training set. Here we evaluate the performance of our method on a reduced training set of CIFAR-10. We randomly pick 10000 training data out of the 50000 available, and use this new set to train different models. After training, the models are evaluated on the standard CIFAR-10 test set. Results are shown in Table 15.
Detailed OOD results and mean corruption error (mCE) experiments. In order to test the robustness of late-phase weights against input data corruption, we used the corruptions and dataset proposed by Hendrycks Dietterich (2019), freely available at https::/// /gi thub. . com/ hendrycks/ robustness The authors propose 15 noise sources such as random Gaussian noise, spatter or contrast changes to deform the input data and report the model test set accuracy on the corrupted dataset under 5 severity levels (noise strengths). For each source noise, its corruption error computed by averaging the prediction error over the severity levels. The average of the cor- ruption error of all 15 noises gives us the Mean Corruption Error (mCE). See Table 16 for the mCE computed on the corrupted CIFAR-100 dataset.
Training run time. Here we compare the training run time of our method with the baseline. The result was computed in Python 3.7, using the automatic differentiation and GPU acceleration pack- age PyTorch (version 1.4.0). We used the standard datasets (including training and test splits) as provided by the torchvision package unless stated otherwise. We used single NVIDIA GeForce 2080 Ti GPU for the experiment. Results are presented in Table 17.
C THEORETICAL ANALYSIS OF THE NOISY QUADRATIC PROBLEM
In this section, we consider a noisy quadratic problem (NQP) that can be theoretically analyzed and that captures important characteristics of the stochasticity of minibatch-based optimizer (Schaul et al., 2013; Martens, 2016; Wu et al., 2018: Zhang et al., 2019a;b). The NQP does a second-orde -rder Taylor expansion of the loss function around the optimum w and models the minibatch noise as
20

---

Published as a conference paper at ICLR 2021
Table 16: OOD performance measured by the AUROC, and robustness measured by the Mean Corruption Error (mCE). We train the models on CIFAR-100 and attempt to discriminate test set images from novel ones drawn from the SVHN, LSUN, Tiny ImageNet (TIN) and CIFAR-10 dataset. The mCE value is the average across 75 different corruptions from the LLFAR-100-4 dataset. LPBN and LP HNET stand respectively for late-phase BatchNorm and late-phase hypernetwork.
SVHN
LSUN
TIN
CIFAR-10
mCE
Base
0.814 #0.004
0.798
0.776=t0.03$
0,818=0.000
47.84 47.84=0.41
LPBN
0.831 #0.021
0.862
0.838+400023
0.814#0.002
45.59-t0.25
LPBN (non-avg.)
0,877 0.0008
0.883
0.863 10.023
0,827
46.2 46.211..299 1.029
LP HNET
0.815 0.002
0..02
0.816 10.027
0,81 100002
47 47.84 1.0.42
Dropout (Mean)
0.792 0.0093
0.807
0.788 00.04
0,822 000003
48 48,97420.35
MC-Dropout
0.806 080082
0.842
0,81 1817:000041
0.824=0.003
48, +90.96
SWAG
0.824=0.01
0.839
0.805
0.816#0.0004
BatchEnsemble
0.848 0.0020
0,828
0.820-0.0030
0.829
Deep ens.
0.839
0.836 0.884
0.812
0.839 0.834
44.21
Deep ens. (LPBN)
0.855
0,856
43.15
Table 17: Training time in seconds and hours on CIFAR-10 for 200 epochs on a single NVIDIA GeForce 2080 Ti GPU.
Model
seconds
hours
Base (SGD)
17714
4.92
Late-phase BN (SGD)
17772
4.94
a random translation e of the optimum, while keeping the curvature H the same. This gives us the following minibatch loss:
w c) H(W w'
VB c)
(4)
VB
with c ~ (O, Σ) and B the minibatch size. Note that we use boldface notation for vectors in this analysis for notational clarity. The NQP can be seen as approximation of the loss function in the final phase of learning, where we initialize the late-phase ensemble. Despite its apparent simplicity, remains a challenging optimization problem that has important similarities with stochastic mini- batch training in deep neural networks (Schaul et al., 2013; Martens, 2016; Wu et al., 2018; Zhang al., 20;),;)) For the simple loss landscape of the NQP, there are three main strategies to improve the expected loss after convergence: (i) increase the mini-batch size B (Zhang et al., 2019a). (ii) use more members K in an ensemble (c.f. Section C.3 and (iii) decrease the learning rate 1 (Schaul et al., 2013; Martens, 2016; Wu et al., 2018; Zhang et al., 2019a;b). The late-phase weights training combines the two first strategies in a non-trivial manner by (i) averaging over the base-weights gradients for all ensemble members and (ii) averaging the late-phase weights in parameter space to obtain a mean-model. The goal of this theoretical analysis is to show that the expected loss after convergence scales inversely with the number of late- late-phase ensemble members K, which indicates that the non-trivial combination of the two strategies is successful.
To model the multiplicative weight interaction between late-phase weights and base weights, linear hypernetworks of arbitrary dimension. The linear hypernetworks parameterize the weights Oe. with 0 E Raxd the hypernetwork parameters and E E the embedding vector. The embedding vectors e are used as late-phase weights (ø in the main manuscript) to create a late- phase ensemble with K members, while using a shared hypernetwork Ultimately, we are interested in the expected risk of the the mean model at steady state:
base-weights: Wa deg.
[[[(ss)] = E
w*)" w")l
(S)
with - Σ fc% p Σx ex & Ae and Pas the steady-state distribution of the parameters. Note that we cannot put w = without loss of generality, because the overparameterization the hypernetworks makes the optimization problem nonlinear.
21

---

Published as a conference paper at ICLR 2021
We start with investigating the discrete time dynamics induced by late-phase learning, after which derive the corresponding continuous time dynamics to be able to use the rich stochastic dynamical systems literature for analyzing the resulting nonlinear stochastic dynamical system.
C.1 DISCRETE TIME DYNAMICS
As we want to investigate the multiplicative interaction between the shared and late-phase parame- ters, substitute Oe into equation 4, instead of computing a new Taylor approximation in the hypernetwork parameter space. Let us take t as the index for the outer loop (updating 0) and & the index for the ensemble member. Then we have the following stochastic minibatch loss:
((.,)) (0())) (1) -
())))) H( ( e
w
(t.,))
(6)
VB
vB
which gives rise to the following parameter updates using late-phase learning with learning rate n) and minibatch size B:
( = ((t)
"
1
Σ &
o( w")ek (1)
Σ K
H H
e
(oT
VBK
K
e =
(++1)
(1) n H( (0) - ww)+ + 7
o He (B.A)
(8)
VB
The above discrete time dynamics nonlinear, giving rise to non Gannaassi parameter distribution P. Hence, it is not possible to characterize these dynamics by the mooment-propagating equations of the first and second moment as done in Zhang al. (2019a;b); Schaul et al. (2013) and Wu et al. (2018), without having full access of the parameter distribution hypernetwork parameterization, diagonal, which is a common approach in the literature. Therefore, we investigate the corresponding continuous time dynamics, such that we can use the rich literature on stochastic dynamical systems.
Furthermore, because of the cannot decouple the system of equations, even if H and E
C.2 CONTINUOUS TIME DYNAMICS
First, let us define some compact notations for the various parameters.
e, - le)
(t)T (T
(9)
E, - le
(1)
(10)
K
0, _ vec(O,)
(11)
10. ITTTT
(12)
6, A le (t.1)T (T
(13)
(14)
where vec(f) concatenates the columns of A in a vector. Then the discrete time dynamics (equation and equation 8) can be rewritten as:
" G(xi)e:
X = xx -n((xs) +
(15)
/B
with
w') (T (0) Ho.)) 1 & (0f Hw')
1 Eg (c 10 & H) H) (B,c) (0)
F(x,)= A
(16)
+ F4 s H 1 (0
G(x,
(17)
(18)
with the Kronecker product, I identity matrix of the appropriate size and vector full of ones of the appropriate size. As a linear transformation of Gaussian variables remains a Gaussian variable, we can rewrite eq. equation 15 as follows:
X X, F)()))
D(x:)G=
(19)
VB
22

---

Published as a conference paper at ICLR 2021
with (2xe) (I @ ())())
0.5 and Ç NN(0,))) I). Following Liu & Theodorou (2019)
and Chaudhari & Soatto (2018), the corresponding continuous-timg time dynamics are:
dx, dxx + 2 D(x)dd
(20)
with W, Brownian motion and B - 2B the inverse temperature. Note that v is incorporated in the noise covariance, such that the correct limit to stochastic continuous time dynamics can be
7)
made (Liu & Theodorou, 2019; Chaudhari & Soatto, 2018; but see Yaida, 2018). For computing the expected loss E[C,] of the mean model, we need to have the stochastic dynamics of this loss. Using the Itō lemma (Itō, 1951; Liu & Theodorou, 2019), which is an extension of the chain rule in the ordinary calculus to the stochastic setting, we get
d.((x)) = VL(())) F(x,) + 1 [ di + [VL(x))
(21)
with δ A (xx-I for notational simplicity and He the Hessian of w.r.t. Xx. As we are take the expectation of equation 21 over the
interested in the expected risk (equation 5), parameter distribution (x(x to get the dynamics of the first moment of the loss (also known as the backward Kolmogorov equation (Kolmogorov, 1931)):
d. [c(x)) = Ep. V[(x))" + I [T((DHH)
(22)
In order to obtain the dynamics of the parameter distribution, the Fokker-Plande equation can be used (Jordan et al., 1998). However, due the nonlinear nature of the stochastic dynamical system, the distribution is non-Gaussian and it is not possible (to our best knowledge) to obtain an analytical solution for equation 22. Nevertheless, we can still gain important insights by investigating the steady-state of equation 22. After convergence, the left-hand side (LHS) is expected to be zero. Hence, we have that
Epss (x = Ep [T(DDHe]]]
(33)
The remainder our arguments is structured as follows. First, we will show that the left-hand-side (LHS) of equation 23 is the expectation of an approximation of a weighted norm of the gradient VL, after which we will connect this norm to the loss L of the mean model. Second, we will investigate the RHS to show that the late-phase learning with ensembles lowers the expected risk of the NQP at steady-state. For clarity and ease of notation, we will drop the ss subscripts. The gradient of the mean-model loss is given by:
[ (e (0e w*) K HHC w¹)
VL(x)
(24)
By introducing Ae A e₁ è and using that Σx Del- = 0, we can rewrite F(x) as:
I 1O
0 KT
[ & H) (1 8 (0 HO) Ae
F(x)=
v((x)+
(25)
with r & 1 K E, Ac and ACT e [oT...o%el We see that F is an approximation of the gradient L where the lower block of VL is scaled by K. Importantly, the lower block of the second element of the RHS of equation 25 (the approximation error) will disappear when taking the inner product with VL and the upper block is not influenced by the number of ensemble members K, which will need later. The LHS of equation 23 can now be rewritten as:
E 8 [x(x)* F(x)] E [V (x)) MVL(x)) + Ept HTT w"'e"']
(26)
with M the diagonal matrix of equation 25 (first element of the RHS). The first term of the RHS of equation 26 the expectation of weighted squared norm of VC. while the second term is an approximation error due to the covariance of Áek. Hence, we see that the LHS of equation 23 can be seen as an approximation of a weighted norm of the gradient VL. By investigating the term v L(x) MV [(x) further, we show that it is closely connected to the loss L.
V MVL()) (w- (' TH H00T H)(w w")
(27)
23

---

Published as a conference paper at ICLR 2021
When comparing to the mean-model loss L = -w W(W w") tightly connected, both using a weighted distance measure between w and w", with only a different weighting. Taken everything together, also the RHS) as a rough proxy for the expected risk under the steady-state distribution (equation 5), which will be important to investigate the influence of the amount of ensemble members the expected risk. Zhu et al. (2018) highlighted this trace quantitiy in equation 23 as measurement of the escaping efficiency out of poor minima. However, we assume that we are in the final valley of convergence (emphasized by this convex NQP), so now this interpretation does not hold and the quantity should be considered as a proxy measurement of the width of the steady-state parameter distribution around the minimum. The trace quantity has Hc and (xxs) as main elements, which we structure in block matrices below (for clarity and ease of notation, we drop the subscripts 88).
that the two
that we take the LHS of equation 23 (and hence
(ee") 201 11 H₂
1 @ Q² K2 1 & HB
(28)
18Q
K EEE*)
@ (HEH) (0 HEH)
E a (HSHO) I& (0 HΣHO)
D(x)² = G(T @ E)
(29)
E
R
with 1 a matrix or vector of the appropriate size full of ones, e & 1/K Ex ek and the rows of Q E Rdxxd given by:
Q ( + 0 a 6. @ (w H),
(30)
with ò; the i-th column of an appropriately sized identity matrix. After some intermediate calcula- tions and rearranging of terms, we reach the following expression for the RHS of equation 23:
E [Tr [Tr = K3
E [T[E"co"]] [+ + Epe, [Tr @ (HEHIQ)] +
2
...Tr( . @ (O"HΣ)))Q") + Tr HOHHO HO) HI
(31)
K E C.C 1 EET Note that everything between the big brackets in the RHS is independent of K in expectation. Hence, we see that the RHS of equation 23 scales inversely by K, exactly as the case for full ensembles (see Section C.3). Importantly, the approximation errors in equation 25 are independent of K, hence, the found scaling of к in equation 31 translates to K that the non-trivial combination of (i) averaging over the base-weights gradients for all ensemble members and (ii) averaging the late-phase weights eş parameter space to obtain a mean-model, succeeds in scaling the expected loss after convergence inversely by K.
with E
scaling of 1 of the expected risk of the NQP, following the above argumentation. Hence,
C.3 NQP WITH FULL ENSEMBLES
As comparison for the above theoretical results, we also analyze the NQP that uses an ensemble of K full weight configurations W& to get mean model w, instead of shared weights 0 and ensemble- member-speeific weights dg. For the case of linear models, the averaging in weight space to obtain mean model is equivalent to the averaging of the predictions over the ensemble, which is conven- tionally done using ensembles. Without loss of generality, we can take w = 0 (corresponding with simple reparameterization of w). Using equation 4, this results in the following parameter updates for the ensemble members:
W
(++1) = = (T H (
He (C.)
(32)
VB
The mean model w A K Ee Wg has the following corresponding discrete dynamics:
w (++1) = (1 ( +
KVB H Σ (tA.
(33)
K
Exact moment propagating equations. As this is discrete linear system with Gaussian noise, the resulting parameter distributions will also be linear and can be fully characterized by the mean
24

---

Published as a conference paper at ICLR 2021
and covariance of the parameters. Taking the expectation and variance of equation 33 results in:
E [1 = ( |
(34)
rf
C C[w(++) = (I T ( (I- n
HEH
(35)
KB
with Σ the covariance matrix of e. For an appropriate 77, the above equations converge to the follow- ing fixed points at steady-state:
E
[w
(36)
17
vec ( pea
w
II (I- & (I nH) vee (HEH) KB
(37)
We see that the steady-state covariance of w and hence of the risk scales with से (E,-1(4) = Ep- [w" Hw] = Tr| [W]]= The expected risk Ep- obtained with computationally expen- sive full ensembles can be seen as lower limit that we try to reach with the economical ensembles of shared weights @ and late-phase weights Φg. Note that for the NQP, increasing the batchsize B has a similar influence as increasing the number of ensemble members K, as can be seen equation 37.
Continuous time stochastic dynamics. We can also do similar continuous time analysis as Section C.2 for the case full ensembles, to better relate to the results of the late-phase learning with shared parameters. Following the same approach, we get the following expression for the trace term:
2 Ep- [Tr[De
TT| (1 @ (HHH)
K (1 & H)]
(38)
KB TH|E
(39)
When comparing to equation 31, we see that the economical ensembles with shared parameters reach the same scaling with R as a result of ensembling, however, some extra terms that vanish asymptotically for big K appear as a result of the interplay between shared and late-phase parame- ters.
Experimental details for Fig. 1. We take the model w = 0 (i.e., K - 1) as our baseline, since this overparameterization could already result accelerated learning (Arora al., 2018). Our parameters are randomly initialized and scaled such that u has a fixed distance to w" of 1. Since the NQP mimics late phase of learning we To O. We study problem of dimension n 100 and train the model with gradient descent (without momentum).
To validate the theoretical results, we show in Fig. that the steady-state reached by our method scales inversely with K, similarly to an ensemble of independently-trainee models. We run ex- periments with E [2, 10, 15, 20, 25) and train every configuration for 2 x 10" iterations until convergence. We average over the last 10a weight updates and over 5 different random seeds.
25


Reviews:
[{"review_id": "YdZYlJKa-a", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThe paper proposes a method to improve solutions found by SGD by ensembling subsets of weights in late-phase. A family of low-dimensional late-phase methods are analyzed and shown to improve generalization in CIFAR-10/100, ImageNet and enwik8. Authors also analyze the method in more tractable noisy quadratic settings. \n\nContribution of the authors is that rather obtaining ensemble they utilize efficient ensemble to guide SGD training and ultimately obtain a single model. \n\nReason for score:\n\nWhile the paper discusses efficient ways of utilizing late-phase weight ensemble and improving SGD training, the demonstrated benefit is not significant enough for practitioners to pursue the method. Without strong practical application potential, merit of the proposed method is weak since it does not obviously elucidate some aspects of neural network training. \n\nPros:\n\nThe paper is clearly written and easy to understand the proposed method is. It is well structured that helps to improve the clarity. \n\nProposed method tackles a significant problem in the standard ensemble method in which both training/inference computation can be quite costly. The paper’s method only ensembles subset of weights therefore added training cost is minimal and since inference is done on averaged weight, it becomes essentially a single model.\n\nAmong various late-phase schemes, BatchNorm late-phase seems to work well which is widely used among vision models so easily applicable. Also since late-phase can be applied post-pretraining, it can be used to improve pre-trained models.  \n\nAs far as I can tell various experimental conditions are very well controlled and thoughtfully designed.\n\nCons:\nThe idea of weight averaging is not so novel as duly noted by the authors.\n\nMain question arises for the paper is whether the proposed method is worth the effort. While all experiments show that the proposed method improves the baseline somewhat, deep ensemble baselines remain strong. Also quoted difference between methods does not mean statistically significant effect (see Vincent Vanhoucke’s article on reporting significant figures https://towardsdatascience.com/digit-significance-in-machine-learning-dea05dd6b85b). According to this article, results reported in Table 1, CIFAR-10 in WRN, a significant figure with a 10k test set should be around 0.2% and differences between different methods are at best marginal. This can be applied to most tables and except for Deep Ensemble’s improvement other differences are not very significant.\n\t\nI wonder as discussed by the authors, this is due to mostly the benefit of ensembles is\nthrough incorporating different modes as argued in [Fort et al., 2020] rather than a single mode. I imagine a single mode ensemble could be beneficial when variance within the mode is large, however for models considered by the authors seem to have small model variance which minimizes effect of technique utilizing single mode. \n\nWhile \\sigma_0 and T_0 are hyperparameters of the algorithm, no good way to determine it is explained. \n\nThe role of section 3.1 is not clear. For one thing, the legend in Figure 1 is confusing where the role of non-integer K is mysterious to me. I would suggest clarifying what the message of the section would be in context of understanding late-phase weight models. \n\n\nNits and additional feedback:\n\nAnonymized link is neither there in the main paper or included as supplementary material. If the authors intended to include the code, this is a note that code can not be found to the reviewers.\n\nFor models that do not use BatchNorm, I believe most interest to practitioners would be using Transformer based models. I wonder if rank-1 late-phase or LayerNorm late-phase would show improvements in this case. \n\nWas “Late-phase classification layers” ever evaluated or discussed in the main paper? I find some discussion on the appendix but seem to be missing in the main text. \n\n---\nI thank the authors for their hard work addressing issues raised by the reviewers.\n\nAuthors have answered many issues pointed out (by improved performance and showing robustness to hyperparameters) and I've increased my score from 5 to 6, and support accepting the paper. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "W3KnwJKuqj2", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work suggests a variant of ensembling that is more compute-efficient. Specifically, it involves forking an ensemble only in the late stage of training, and forming this ensemble via a \"low-dimentional\" family. That is, instead of maintaining independent networks, maintain only \"low-rank\"-style perturbations of the base network (for various instanciations of \"low-rank\").\nThe experimental results are somewhat limited, but appear to be competitive with current efficient-ensembling approaches like SWA/SWAG. The absolute improvement of this method is not very large (<0.3% on CIFAR, <0.2% on imagenet), and there is a large gap to Deep Ensembles. I weakly recommend acceptance, because the method appears promising for future work, and the experiments seem correct. \n\nThere is also a theory section included, though I am generally unconvinced by results in such simple toy examples.\n(such settings can usually be contrived to exhibit any desired behavior)\n\nWeaknesses:\n- The experimental section would be greatly strengthened by additional experiments for different models and settings. There are only 2 architectures tested on CIFAR-10, for example. It would also be informative to see the performance of these methods in \"harder\" settings -- for example, CIFAR-10 with fewer train samples.\n- The OOD uncertainty results could be expanded. Uncertainty estimation and robustness are some of the most relevant practical uses of ensemble methods, so it is especially important to evaluate ensembles in this context. Currently aggregate results are shown in Table 4, but it would be good to explicitly see, for example: how the performance of this method degrades with increasing CIFAR-10C corruption severity, as opposed to Deep Ensembles. Also, reporting the Mean Corruption Error (mCE) for each dataset individually will allow standard comparison to prior methods.\n\nComments which do not affect the score:\nIt seems that starting the ensembling at a \"late phase\" in training is the main contribution of this work. This could be applied to any ensemble method, and you propose several explicit instantiations. It could help to focus the writing in terms of this contribution, and also to further investigate the role of T0 (the time at which ensembling starts).\n\n---\nEdit after rebuttal: Increased score from 6 to 7.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "7vdGSI01_iN", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "To improve the generalization performance of SGD methods, \nthis paper proposes to use an efficient ensemble-like approach \nwhich computes an average of an ensemble of SGD weights\nwhen retrained from some late-phase of SGD dynamics. \nThis idea is different to most recent ensemble-based approaches which \naim to average the predictions of the models. \n\nThe paper focuses on some specific layers of neural networks \nin order to apply the late-phase training. \nThe batch normalization layers are shown to be \nsimple and effective. Some other layers are also analyzed, \nincluding a recently introduced rank-1 multiplicative matrix \nweights idea for full-connected layers. \nSection 3 presents the numerical results and show that the generalization of SGD \nis more-or-less improved on various benchmarks.\nExplanation of why the generalization is improved in relation with the flatness\nof energy landscape is also discussed.  \n\nI find that this approach is quite sensitive the choice of the hyper-parameters, \nsuch as the beginning of the late-phase T0, and the noise perturbation sigma0. \nIt is written in Section 2.1 that in practice … sigma0>0 yields a set of models …\nthis results in improved final generalization. However, in the result of ImageNet in Section 3.3, \nthe sigma0 equals to 0. Thus, it is not conclusive that sigma0>0 is better. \nAs the improvement in Section 3.3 seems marginal compared to the baseline and the \nstandard deviation, it thus does not fully support the effectiveness of the batch normalization layers. \nI would recommend using some other dataset or models, \nbut with a more consistent set of hyper-parameters. \n\nIn terms of writing, I would recommend to write out\nthe full algorithm of Alg. 1 or at least in the Appendix, \nincluding the variant of the SGD momentum and Adam. \nThe SWA is also worth writing out clearly, which is not clear to the reader. \nIs the DeepEnsemble result in Table 1 from SGD or SWA? \nThis is not clear from the text.\n\nOverall, I think both the methodology and the writing need to be improved. \n\n##\nThe revisions made by the authors have addressed all my concerns. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "2G9emdue6Pa", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n### Summary\nThe authors propose late-phase weights, a method of updating the weights near the end of training via a splitting and ensembling mechanism. They analyze the benefits in the noisy quadratic setting. The method improves validation performance on a range of image recognition tasks and on enwiki8.\n\n### Comments\n* The weight interaction functions $h$ should be more explicitly defined rather than just described in text.\n* The paper is overall well written and flows smoothly.\n* I think there should be more discussion on the choice of $T_0$. For example, in table 1, why does SGD perform worse when $T_0=0$? It would be good to get a sense of robustness to this hyperparameter.\n* Good results on CIFAR. Late-phase weights are shown to boost performance over SGD and to be complementary with SWA. There are some benefits in the OOD setting as well.\n\n### Recommendation / Justification\nI vote to accept the paper. The idea is interesting, well-motivated, and seems straightforward to incorporate into existing pipelines. However, the improvements seems modest in some settings (e.g. ImageNet) and for the best performance, it seems like we should still stick to Deep Ensembles.\n\n### Questions\n* On the ImageNet experiments, what is the validation accuracy of the pre-trained model?\n* Can you comment on the computaional and memory complexity of your algorithm versus vanilla SGD?\n* In the comparisons between late phase weights and SGD, do both algorithms consume the same amount of data? If so, this would be good to mention.\n* Could the entire network be treated as \"late-phase weights\"? Would this help performance?\n\n### Minor comments\n* I would consider alluding to possible choices of the weight interaction functions $h$ when it is first introduced at the start of 2.1.\n* In Algorithm 1: How does the loss function consume three inputs? This is different from when it is initially described.  \n* It's a bit unclear what is being compared in Figure 2.\n\n\n(increased score from 6 to 7)\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to learn an ensemble of weights given a set of base weights from some point late in normal training. The authors apply this approach to a number of configurations and find modest performance improvements for normal test settings and larger improvements for out of distribution settings. While reviewers had some concerns about the size of the improvement relative to baselines, all reviewers agreed that the proposed method is interesting and will likely impact future work, especially given the new experiments provided by the authors. I recommend that the paper be accepted. 
Decision: Accept (Poster)