N EURAL NETWORKS WITH LATE - PHASE WEIGHTS
A BSTRACT
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 ﬁndings are complemented with a theoretical analysis of a noisy quadratic problem which provides a simpliﬁed picture of the late phases of neural network learning.
I NTRODUCTION
Neural networks trained with SGD generalize remarkably well on a wide range of problems.
A classic technique to further improve generalization is to ensemble many such models (Lakshminarayanan 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 complexity, for example by collecting models along an 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 of a single model (Zhang et al., 2015; Pittorino et al., 2020).
We join these efforts and develop a method that ﬁne-tunes the behavior of SGD using late-phase weights: late in training, we replicate a subset of the weights of a neural network and randomly initialize them in a 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 on ensembles with shared parameters (Wen et al., 2020) we explore a 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 (Ioffe & Szegedy, 2015) is a surprisingly simple and effective strategy for improving generalization 1 .
Furthermore, we ﬁnd 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.
M ETHODS AND MODELS 2. L EARNING WITH LATE - PHASE WEIGHTS
Late-phase weights.
To apply our learning algorithm to a given neural network model f w we ﬁrst specify its weights w in terms of two components, base and late-phase ( θ and φ , resp.).
The two components interact according to a weight interaction function w = h ( θ, φ ) .
Base weights are learned throughout the entire training session, and until time step T 0 both θ and φ are learned and treated on equal grounds.
At time step T 0 , a hyperparameter of our algorithm, we introduce K late-phase components Φ = { φ k } Kk =1 , that are learned together with θ until the end.
This procedure yields a late-phase ensemble of K neural networks with parameter sharing: reusing the base weights θ , each late-phase weight φ k deﬁnes a model with parameters w k = h ( θ, φ k ) .
Late-phase weight averaging at test time.
Our ensemble deﬁned by the K late-phase weight conﬁgurations in Φ is kept only during learning.
At test time, we discard the ensemble and obtain a single model by averaging over the K late-phase weight components.
That is, given some input pattern x , we generate a prediction y ( x ) using the averaged model, computed once after learning:
[Equation]
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 ﬁrst learn a base parameter φ 0 from time step t = 0 until T 0 , treating φ 0 as any other base parameter in θ .
Then, at time t = T 0 , each conﬁguration φ k is initialized in the vicinity of φ 0 .
We explore perturbing φ 0 using a symmetric Gaussian noise model,
[Equation]
where (cid:15) k is a standard normal variate of appropriate dimension and σ 0 is a hyperparameter controlling the noise amplitude.
We allow for a φ 0 -dependent normalization factor, which we set so as to ensure layerwise scale-invariance, which helps ﬁnding a single σ 0 that governs the initialization of the entire network.
More concretely, for a given neural network layer l with weights φ ( l ) 0 of dimension D ( l ) , we choose Z ( φ ( l ) 0 ) = [Equation] D
[Equation] l ) / (cid:107) [Equation] ( l
[Equation]
0 (cid:107) .
Our perturbative initialization (Eq. 2) is motivated by ongoing studies of the nonconvex, highdimensional loss functions that arise in deep learning.
Empirical results and theoretical analyses of simpliﬁed models point to the existence of dense clusters of connected solutions with a locallyﬂat 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 evaluated at weight conﬁgurations 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 understood reasons, this appears to be a recurring phenomenon in overparameterized nonlinear problems (Brown & Sethna, 2003; Waterfall et al., 2006).
Based on these observations, we assume that the initial parameter conﬁguration φ 0 can be perturbed in a late phase of learning without leading to mode hopping across the different models w k .
While mode coverage is usually a 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 components, we now present a stochastic algorithm which learns both θ and Φ .
Our algorithm works on the standard stochastic (minibatch) neural network optimization setting (Bottou, 2010).
Given a loss function L ( D , w ) = 1 |D| (cid:80) x ∈D L ( x, w ) to be minimized with respect to the weights w on a set of data D , at every round we randomly sample a subset M from D and optimize instead the stochastic loss L ( M , w ) .
However, in contrast to the standard setting, in late stages of learning ( t > T 0 ) we simultaneously optimize K parameterizations W := { w k | w k = h ( θ, φ k ) } K k =1 , instead of one.
We proceed by iteration over W .
At each step k , we sample a minibatch M k and immediately update the late-phase weights φ k , while accumulating gradients over the shared base weights θ .
Such gradient accumulation has been previously used when learning ensembles (Lee et al., 2015;
Wen et al., 2020) and multi-task models (Rebufﬁ et al., 2017) with shared base parameters.
A single iteration is ﬁnally concluded by changing the base weights in the direction opposite of the accumulated gradient.
We scale the accumulated gradient by γ θ ; setting γ θ = 1 /K recovers the original step size in θ , but other choices are possible.
In particular, we ﬁnd that a large γ θ of unit size is in practice often tolerated, resulting in accelerated learning.
Require: Base weights θ , late-phase weight set Φ , dataset D , gradient scale factor γ θ , loss L
Require:
Training iteration t > Tfor 1 ≤ k ≤ K do M k ← Sample minibatch from D ∆ θ k ← ∇ θ L ( M k , θ, φ k )
φ k ← U φ ( φ k , ∇ φ k L ( M k , θ, φ k ))
θ ← U θ ( θ, γ θ (cid:80) Kk =1 ∆ θ k )
We summarize an iteration of our method in Algorithm 1, where the loss L ( M , θ, φ ) is now seen as a function of θ and φ .
We opt for a general presentation using unspeciﬁed gradient-based update operators
U φ and U θ .
These operators can be set to optimizers of choice.
For instance, our method might beneﬁt from additional noise injection onto parameter updates (Welling & Teh, 2011).
Furthermore, late-phase optimizers need not coincide with the optimizer used in the early phase.
In our work we typically set U φ and
U θ 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.
2. L ATE - PHASE WEIGHT MODELS
As detailed next, we consider a number of distinct late-phase weight models in our experiments.
In particular, we explore weight interaction functions h in which late-phase weights have low dimensionality, to avoid a 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; Ioffe &
Szegedy, 2015) are a staple of current deep neural network models.
Besides standardizing the activity of the layer they are applied to, BatchNorm units introduce a learnable multiplicative (scale) parameter γ and an additive (shift) parameter β .
While being low-dimensional, these additional parameters have large expressive power: it has been shown that learning only γ and β keeping the remaining weights frozen can lead to signiﬁcantly lower loss than when learning random subsets of other weights of matching dimensionality (Frankle et al., 2020; Mudrakarta et al., 2019).
We take the scale and shift parameters of BatchNorm layers as our ﬁrst 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 γ 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 efﬁciently.
More concretely, let us consider an afﬁne transformation layer l which maps an input vector r ( l − 1) to θ ( l ) w r ( l − 1)
+ θ ( l ) b , where the early-phase weight matrix θ ( l ) w and
bias vector θ ( l ) b are already standardized using the respective batch statistics.
For this standard layer, our model introduces a multiplicative interaction between base and late-phase weights, diag ( γ ( l ) ) θ ( l ) w , and an additive interaction between base and late-phase bias parameters, θ ( l ) b + β ( l ) .
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 afﬁne layer l , we deﬁne a late-phase weight matrix with resort to a pair of learnable vectors, φ ( l ) =
u ( l ) v ( l ) T .
Taking the Hadamard product with the base weight matrix yields the effective weights W ( l )
=
φ ( l ) ◦ θ ( l ) .
With this parameterization, we 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 a purely multiplicative manner 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 w of a given target neural network f w based on a weight embedding.
In our framework, we can use a hypernetwork to implement the interaction function w = h ( θ, φ ) directly, with parameters θ corresponding to base weights and embeddings φ to late-phase weights.
We experiment with linear hypernetworks and use the same hypernetwork to produce the weights of multiple layers, following Savarese & Maire (2019); Ha et al.
von Oswald et al.
In this scheme, the weight embedding input speciﬁes the target layer whose parameters are being generated.
More speciﬁcally, the weight matrix for some layer l belonging to a group of layers g which share a hypernetwork is given by W ( g,l )
= θ ( g ) φ ( g,l ) , where θ ( g ) and φ ( g,l ) are appropriatelysized tensors.
Sharing θ ( g ) over a layer group g allows countering an increase in the overall number of parameters.
We parameterize our hypernetworks such that the weight embedding vectors φ ( g,l ) are small, and therefore cheap to ensemble.
Late-phase classiﬁcation layers.
Finally, inspired by Lee et al.
(2015), in classiﬁcation experiments we take the weights of the last linear layer as late-phase weights by default.
In modern neural network architectures these layers do not usually comprise large numbers of parameters, and our architecture explorations indicated that it is typically beneﬁcial to ensemble them.
We therefore include W ( L ) in our late-phase weights φ , where W ( L ) denotes the weights of the ﬁnal layer L .
R ESULTS 3. N OISY QUADRATIC PROBLEM ANALYSIS
Before turning to real-world learning problems, we ﬁrst focus on a simpliﬁed stochastic optimization setup which can be analytically studied.
We consider the noisy quadratic problem (NQP; Schaul et al., 2013; Martens, 2016;
Wu et al., 2018; Zhang et al., 2019a;b), where the goal is to minimize the scalar loss
[Equation]
with respect to w ∈ R n .
In the equation above, w ∗ denotes the target weight vector, which is randomly shifted by a noise variable (cid:15) assumed to follow a Gaussian distribution N (0 , Σ) .
The (constant) Hessian matrix H controls the curvature of the problem.
[Figure]
Despite the simplicity of Eq.
3, the NQP captures a surprising number of empirically-observed aspects of neural network learning (Zhang et al., 2019a).
Here, we motivate its study as a model of late stages of learning, by Taylor expanding the loss around a minimum w ∗ .
Thus, for a sufﬁciently late initialization time T(and small σ 0 ) the NQP is particularly well suited to study our algorithm.
There are three main strategies to improve the expected NQP loss after convergence: ( i) increase the minibatch size B , ( ii) use more members K in an ensemble, and ( iii) decrease the learning rate η (Zhang et al., 2019a).
Our Algorithm 1 combines the ﬁrst two strategies in a non-trivial manner.
First, the gradients for base weights θ are averaged during the inner loop over all ensemble members, corresponding to a minibatch-size rescaling by K .
Second, we introduce K ensemble members, to be averaged in weight space, that only differ in their late-phase weights φ .
In Appendix C, we show analytically that this combination of an increased effective minibatch size for θ and introducing K ensemble members for φ is successful, resulting in a scaling of the expected loss after convergence by 1 K .
This analysis holds for general Σ and H , and for both scalar and hypernetwork multiplicative late-phase weights.
Hence, our approach combines the beneﬁts 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.
Our model includes a multiplicative late-phase weight, w k = θ φ k with φ k ∈ R and θ ∈ R n .
We simulate a standard instance of the NQP, with diagonal Hessian H ii = 1 /i and Σ
= H − 1 (cf.
Zhang et al., 2019a), and report the average loss after convergence.
Hyperparameters are given in 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 a full ensemble of K models trained independently with gradient descent.
3. CIFAR-10/100 EXPERIMENTS
To test the applicability of our method to more realistic problems, we next augment standard neural network models with late-phase weights and examine their performance on the CIFAR-10 and
CIFAR-100 image classiﬁcation benchmarks (Krizhevsky, 2009).
We use standard data preprocessing 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 same amount of data.
Besides SGD (with Nesterov momentum), we also investigate stochastic weight averaging (SWA;
Izmailov et al., 2018), a 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 (cf.
Appendix A).
When learning neural networks with late-phase weights we set U φ and U θ 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 a 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-trained models (a deep ensemble, Lakshminarayanan et al., 2017).
Deep ensembles provide a strong baseline, 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 of an upper baseline.
[Table]
By contrast, augmenting the architectures considered here with late-phase weights results in negligible additional costs during learning (with the exception of hypernetworks, which require additional tensor products) and none during testing.
In principle, a set of independently-trained models yielded by our algorithm can therefore even be used as the basis of a deep ensemble, when the memory and compute budget allows for one.
We present proof-of-concept experiments exploring this option.
Throughout our CIFAR-10/100 experiments we set K = 10 , use a fast base gradient scale factor of γ θ = 1 , and set our late-phase initialization hyperparameters to T
= 120 (measured henceforth in epochs; T 0 = 100 for SWA) and do not use initialization noise, σ 0 = 0 .
These hyperparameters were tuned manually once on CIFAR-100 and then kept ﬁxed unless otherwise 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 Φ .
CIFAR-10.
For CIFAR-10 we focus on the WRN architecture, a high-performance residual network (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.
All algorithms achieve a training error close to zero (cf.
Appendix B).
The resulting predictive accuracies are shown in Table 1.
We ﬁnd that augmenting the WRN 28-10 (a standard WRN conﬁguration) with BatchNorm late-phase weights leads to a systematic improvement in generalization, reducing the gap with a deep ensemble of K = 10 models.
Initializing our ensemble from the onset ( T 0 = 0 ) fails to meet the performance of the base model, reaching only 95 .
68 ± 0 .
23 % (cf.
Appendix 12).
We also investigate initializing a late-phase (full) deep ensemble at T 0 = 120 .
This results in a 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-efﬁciency of our low-dimensional late-phase ensembles which can be trained with as little data as a single model, besides being memory efﬁcient.
In addition, we consider a larger instance of the WRN model (the WRN 28-14), trained forepochs using cutout data augmentation (DeVries & Taylor, 2017), as well as a small convolution neural network without skip connections, cf.
Table 3.
When late-phase weights are employed in combination with SWA, we observe signiﬁcant accuracy gains on the WRN 28-14.
Thus, our latephase weights impose an implicit regularization that is effective on models with many weights.
Similarly, we observe larger gains when training on a random subset of CIFAR-10 with only 10examples (cf.
Appendix B).
[Table]
CIFAR-100.
We next turn to the CIFAR-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 variant (using cutout data augmentation as before) and a 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 our 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-phase hypernetwork weight embeddings yielding essentially no improvements.
Once again, the setting of T 0 =
(onset ensemble learning) fails to match base model performance, ﬁnishing at 80 .
26 ± 0 .
42 % test accuracy.
As for CIFAR-10, a late-phase full deep ensemble only reached intermediate improvements, at 82 .
17 ± 0 .
15 % test accuracy.
Furthermore, a gap towards deep ensembles persists.
This suggests that covering different modes of the loss (Fort et al., 2020) can provide beneﬁts that cannot be captured by ensembling models in a small neighborhood. [Table]
The ﬁnal averaged solutions found with late-phase weights are strong base models to build a deep ensemble of independently-trained networks.
The fact that our algorithm yields a single model allows further pushing the upper bound of what can be achieved when unrestricted full ensemble training is possible.
This improvement comes at no cost compared to a standard deep ensemble.
We train additional neural network architectures restricting our experiments to the BatchNorm latephase weight model, which can be readily implemented without architectural modiﬁcations.
Again, learning with late-phase weights yields a consistent improvement over the baseline, cf.
Table 3.
[Figure] Training epoch [Figure] Notably, SWA can achieve high predictive accuracy with a large constant learning rate (Izmailov et al., 2018).
We reproduce these previous results and show that they improve when learning with late-phase weights, cf
Substantial progress is made both when entering the latephase learning period and when activating SWA.
Out-of-distribution (OOD) generalization.
Deep ensembles are an effective technique for improving the behavior of neural networks in OOD data (Lakshminarayanan et al., 2017).
We ask whether our implicit ensembles modeled during late-phase learning could confer a similar advantage to our ﬁnal averaged model.
Additionally, we evaluate the performance of a late-phase weight ensemble obtained with large initialization noise σ 0 = 0 .
5 (at T 0 = 100 ), skipping the ﬁnal weight averaging step.
This requires integrating predictions over K late-phase ensemble members at test time, y ( x )
= 1 K (cid:80)
Kk =1 y ( x, w k ) .
Unlike standard deep ensembles, training this ensemble is still as cheap as training a single model.
[Table]
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 a WRN 28-10 trained on CIFAR-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 our results to alternative methods with strong uncertainty representation: MC-dropout (Gal & Ghahramani, 2016), SWA-Gaussian (SWAG; Maddox et al., 2019) and BatchEnsemble (Wen et al., 2020).
All three methods require integrating predictions over an ensemble at test time.
We ﬁnd that learning with late-phase weights increases prediction uncertainty in OOD data, allowing for a signiﬁcantly better separation between in and out-of-distribution examples, cf.
Table 4.
The
OOD performance of late-phase BatchNorm weights compares favorably to the alternative methods including deep ensembles, even when using a single weight-averaged 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 σ 0 >(cf.
Appendix B), this method results in the best OOD performance.
[Figure] Despite our improved performance on both predictive accuracy (with late-phase BatchNorm) and OOD discrimination (with late-phase BatchNorm and hypernetwork embeddings), the test set negative log-likelihood (NLL; often used to assess predictive uncertainty, Guo et al., 2017) is surprisingly slightly worse for our solutions.
This is aligned with the ﬁnding that SWA does not always signiﬁcantly reduce NLL, even though predictive accuracy increases (Maddox et al., 2019).
Flatness. Why do our networks generalize better?
Approximate Bayesian inference suggests that ﬂat minima generalize better than sharp minima (Hochreiter & Schmidhuber, 1997a; MacKay, 1992).
Due to symmetries 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 K late-phase weight conﬁgurations can implicitly lead to ﬂatter solutions.
To investigate whether our algorithm ﬁnds ﬂatter minima, we examine a simple ﬂatness score that correlates well with generalization (Pittorino et al., 2020; Jiang et al., 2020).
Concretely, we add multiplicative Gaussian noise
z i ∼ N (0 , w 2 i σ 2 z ) to each weight w i
and then measure the change in the loss δ L
= E z [ L ( w + z ) − L ( w )] .
Our ﬁnal weight conﬁgurations are indeed in ﬂatter regions of weight space according to this measure: δL increases more slowly with σ z for the WRN 28-10 models that are learned with BatchNorm late-phase weights, Fig. 3.
3. I MAGE N ET EXPERIMENTS
[Table]
To investigate whether our gains translate to large-scale learning problems, we train deep residual networks (He et al., 2016) and a densely-connected convolutional network (DenseNet; Huang et al., 2018) on the ImageNet dataset (Russakovsky et al., 2015).
We start from pretrained models and contrast BatchNorm late-phase weight learning to ﬁne-tuning with SGD for 20 epochs, with γ θ = 1 /K and σ 0 =
(cf.
Appendix A).
For simplicity we do not include last-layer weights in Φ .
Fine-tuning with late-phase weights improves the ﬁnal top-1 validation accuracy of this pretrained model signiﬁcantly with only minor training, as seen in Table 5.
These results serve as a proof-ofconcept that existing models can be further improved, taking our late-phase initialization T 0 as the time the previous experimenter stopped training.
In Appendix B, we present additional CIFARexperiments where we apply late-phase learning starting at the suboptimal end-of-training T 0 = 200 , to mimic the pretrained condition.
3. LSTM LANGUAGE MODELING EXPERIMENTS
Finally, we conduct experiments on the language modeling benchmark enwik8 .
To show that the beneﬁts 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.
Section 2.2).
[Table]
Overﬁtting is a 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 outperform previously published state-of-the-art models (Melis et al., 2017).
To avoid this issue, we train models where the number of parameters ( ∼ 1.56M) is drastically smaller than the number of training data points (90M), such that we do not observe any overﬁtting.
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.
We train our LSTM with 500 units for 50 epochs, optimizing every weight with Adam (Kingma & Ba, 2015).
We apply a multiplicative rank-1 matrix elementwise to the recurrent weight matrix.
Interestingly, merely adding the multiplicative parameters to the LSTM (Base) accelerates training and leads to better training and test set performance (measured in bits per character, BPC) with no additional changes to the optimizer (Base + Rank1, Table 6).
Further improvements can be achieved with our late-phase weights.
We generate K = 10 late-phase weight components at epoch 30 with σ 0 = 0 .
35 and set γ θ = 1 .
Additionally, we ﬁnd that SWA (starting at epoch 40) substantially improves all scores, with smaller gains on the models with multiplicative weights.
R ELATED WORK
Our late-phase weights deﬁne an ensemble with the special property that every model shares the same base weights.
Such parameter sharing is 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 a 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 latephase 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 do not ﬁnd the need to introduce a coupling term to our loss function.
Additionally, as we replicate a 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 different 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 et 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 et al., 2016; Ha et al., 2017), and continual learning with biologically-realistic synapses (Kaplanis 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 φ k varies (cf.
Algorithm 1) suggests interpreting the base θ as slow weights and the late-phase Φ as fast weights.
This accumulation is reminiscent of a recent meta-learning algorithm (Zintgraf et al., 2019), which ﬁrst separates parameters into task-shared and task-speciﬁc, and then differentiates through a sequence of accumulated updates performed over the task-speciﬁc parameters (Finn et al., 2017).
Continuing with the fast-slow weight analogy, our averaging over fast weights at the end of learning (Eq. 1) could be thought of as a synaptic consolidation step which integrates the fast weight components onto a slow, persistent form of memory.
C ONCLUSION
We proposed to replicate and learn in parallel a subset of weights in a late phase of neural network learning.
These late-phase weights deﬁne an ensemble of models which share every other weight.
We studied convolutional neural networks, a common recurrent neural network, and a simple quadratic problem.
Surprisingly, across these cases, we found that a small number of appropriately 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 a wide range of existing models, including pretrained ones.
We expect future work to uncover new effective late-phase weight models.
A CKNOWLEDGEMENTS
This work was supported by the Swiss National Science Foundation (B.F.G. CRSII5-173721 and 315230 189251), ETH project funding (B.F.G. ETH-20 19-01), the Human Frontiers Science Program (RGY0072/2019) and funding from the Swiss Data Science Center (B.F.G, C17-18, J.v.
O.
P18-03).
Jo˜ao 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, ˆAngelo Cardoso and Angelika Steger for feedback, Mark van Rossum for discussions on ﬂat minima, Simone Surace for his detailed feedback on Appendix C, and Asier Mujika for providing very useful starter code for our LSTM experiments. A A DDITIONAL IMPLEMENTATION DETAILS
Hypernetwork model.
The base neural network architecture we use when parameterizing our weights using a 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 hypernetwork architectures for WRNs.
We do not use dropout or biases in the convolutional layers.
The parameters of every convolutional layer are hypernetwork-generated, with one hypernetwork per layer group (Table 7).
The remaining parameters, namely those of BatchNorm units and ﬁnal linear layer weights, are non-hypernetwork-generated.
Following Savarese & Maire (2019) we turn off weight decay for the model embeddings and initialize these parameters with a random pseudo-orthogonal initialization over layers.
The hypernetwork parameters are initialized using a standard Kaiming initialization (He et al., 2015). [Table]
Small ConvNet model.
We train a slight modiﬁcation 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 we use rectiﬁed linear units on the hidden layers.
Furthermore, after each such activation, BatchNorm units are inserted.
We optimize the model with SGD and use late-phase BatchNorm weights, with T
= 50 and σ 0 = 0 .
For simplicity of implementation, we do not include the last linear layer in the late-phase weight set Φ .
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 .is
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 a ﬁnal value of 0 .
001 ( 0 .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-phase weights proportional to the ensemble size, 0 .
0005 K .
WRN-28-14: The WRN 28-14 models are trained for 300 epochs on 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 as for the WRN 28-10 model, except that we anneal the learning rate until the 160th epoch.
(SWA)
Require: Base weights θ , dataset D , hyperparameter η , loss L
Require:
Training iteration t
M ← Sample minibatch from D ∆ θ ← ∇ θ L ( M , θ ) θ ← U ( θ, η, ∆ θ ) θ
SWA ← ( t θ SWA + θ ) /
( t + 1) t ← t +
Require: Base weights θ , dataset D , learning rate η , momentum ρ , loss L
M ← Sample minibatch from D ∆ θ ← ∇ θ L ( M , θ + ρ ν )
ν ← ρν − η ∆ θ θ ← θ + ν
Batch normalization units.
Whenever we use SWA, we follow Izmailov et al.
and perform a 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 statistics are computed on the training dataset.
The same transformation is then applied when testing, including to OOD data.
Following a standard procedure (e.g., Zagoruyko & Komodakis, 2016;
He et al., 2016)
we augment our training datasets using random crops (with a 4-pixel padding for
CIFAR) and random horizontal ﬂips.
The ImageNet training dataset is augmented with random horizontal ﬂips, as well as random cropping of size 224, while a centered cropping of size 224 was used on the test set.
Our OOD datasets are resized to ﬁt whenever necessary; we used the resized images made available by Lee et al.
ImageNet experiments.
The pretrained model for the ImageNet experiment is obtained from torchvision’s models subpackage.
We ﬁne-tune the model for 20 additional epochs on ImageNet. Require:
Base weights θ , late-phase weight set Φ , dataset D , gradient scale factor γ θ , learning rate η , ensemble size K , initialization noise σ 0 , initialization time T 0 , number of training iterations T , loss L
Initialization: ˆ K ← 0 , s ← 0 , t ←while t ≤ T do
if t = T 0 then // generate late-phase weights for 1 ≤ k ≤ K do sample (cid:15) ∼ N (0 , 1) φ k ← φ 0 + σ 0 (cid:107) φ 0 (cid:107) (cid:15) //
set range for specialists training ˆ K ← K s ←for s ≤ k ≤ ˆ K do M k ← Sample minibatch from D ∆ θ [Equation]
← ∇
[Equation]
L ( M
[Equation]
, θ, φ k ) φ
[Equation]
←
[Equation]
L ( M
[Equation]
, θ, φ
[Equation]
) t
[Equation]
t +
[Figure]
We use a 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 a batch size of 256 .
We use σ 0 = 0 and K = 10 for our late-phase model.
[Table]
Code forks.
Our hypernetwork 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://github.com/ meliketoy/wide-resnet.pytorch .
LSTM All experiments are conducted using the
Tensorﬂow Python framework (Abadi et al., 2016).
All base weights are initialized uniform in [ − 0 .
whereas the initial rank-1 matrix weights are centered around 1 i.e. [1 − 0 .
01 , 1+0 .
01] to mitigate strong difference in initialization compared to the base model.
We use the Tensorﬂow default values ( β 1 = 0 .
9 , β 2 = 0 .
, (cid:15) =
10 − 8 ) for the
Adam optimiser.
We perform a grid search over σ 0 ∈ [0 , 0 . 5] (in steps of size 0.05) for our LSTM experiments (ﬁxing K = 10 and varying T 0 ∈ { 0 , 30 } ) and obtain the values reported in the main text, T 0 = 30 and σ 0 = 0 . [Figure]
B A DDITIONAL EXPERIMENTS
[Table]
Pretrained CIFAR-100.
We apply our method to a standard WRN 28-10 pretrained on CIFAR100 (i.e., we set T 0 = 200 ) and train for an additional 20 epochs.
At the beginning of the ﬁnetuning, the learning rate is reset to 0.01, then annealed linearly to 0.001 for 10 epochs.
It is then held constant for the remainder of the ﬁne-tuning process.
We observe that augmenting with BatchNorm late-phase weights yields an improved predictive accuracy compared to additional ﬁne-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 ﬁnd that ﬁne-tuning only BatchNorm late-phase weights while keeping all other weights ﬁxed does not even match the Base control.
Together with the ﬁnding that the optimal latephase weight initialization time is at T ∗
= 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 2810 model (without late-phase weights) and start accumulating gradients over K = 10 minibatches at T
= 120 , experimenting both with γ θ = 1 /K and γ θ = 1 .
The models are trained with
SGD using otherwise standard optimization settings.
Both controls fail to improve (even match) the performance of the base model trained without any gradient accumulation.
Sensitivity to T 0 , K and σ 0 .
We present a hyperparameter exploration on the CIFAR-100 dataset using BatchNorm late-phase weights in Tables 8, 11 and 12.
We ﬁnd that our algorithm is largely [Table] [Table] robust to σ 0 when T 0 can be set to its optimal value, which is at 60% of training.
See also Figure 6 for a visualisation of the same data, speciﬁcally the change in mean AUROC score and test set accuracy when changing T 0 .
This result holds also on CIFAR-10, cf.
Table 12.
When starting from a pretrained condition ( T 0 = 200 ), ﬁnite σ 0 leads to a signiﬁcant improvement in performance, cf.
Table 11.
We therefore report results obtained with σ 0 = 0 for every CIFAR and ImageNet experiment in the main text.
The exception to this is the non-averaged (ensemble) late-phase BatchNorm model presented in Table 4, which was optimized for best OOD performance (corresponding to σ 0 = 0 . 5 ). [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
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 a WRN 28-10 on CIFAR-100 with the exact same conﬁguration as for our base model, see above, but include dropout layers as usually done (Zagoruyko & Komodakis, 2016) after the ﬁrst convolution in each residual block.
For a scan over the dropout probability p in this setup, see Table 13.
p = 0 .
2 is reported in the main text - for CIFAR-100 and CIFAR10.
Note that p was only tuned for CIFAR-100.
For the reported results of BatchEnsemble (Wen et al., 2020), we simply execute the code provided by the authors at https://github.com/ google/uncertainty-baselines with their ﬁne-tuned conﬁguration for CIFAR-10/100.
Notably, the authors use a different setup than followed in this manuscript.
First, the WRN 28is trained for 250 epochs (we allow for this increased budget exceptionally for BatchEnsemble), with a multi-step learning rate annealing at [80 , 160 , 180] with a learning rate decay factor of 0 . 2 .
Second, a weight decay of 3 × 10 − 4 is used.
[Table]
[Table]
For the results reported for SWAG (Maddox et al., 2019), we use the code provided by the authors at https://github.com/wjmaddox/swa_gaussian , and the proposed ﬁne-tuned conﬁguration which coincides with the conﬁguration 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 ﬁnal achieved training losses for the base model and when augmenting it with BatchNorm late-phase weights on Table 14, for both CIFAR-10 and CIFAR-100.
Using a fast gradient accumulation scale factor of γ θ = 1 leads to a higher training loss on CIFARthan that of the standard model, but we found this setting crucial to achieve the largest improvement on test set generalization.
CIFAR-10 with a 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://github.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 is computed by averaging the prediction error over the severity levels.
The average of the corruption 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 package 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 a single NVIDIA GeForce
Ti GPU for the experiment.
Results are presented in Table 17.
C T HEORETICAL 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 a minibatch-based optimizer (Schaul et al., 2013; Martens, 2016; Wu et al., 2018; Zhang et al., 2019a;b).
The NQP does a second-order
Taylor expansion of the loss function around the optimum w ∗ and models the minibatch noise as [Table]
[Table]
a random translation (cid:15) of the optimum, while keeping the curvature H the same.
This gives us the following minibatch loss:
[Equation]
with (cid:15) ∼ N ( 0 , Σ) 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 an approximation of the loss function in the ﬁnal phase of learning, where we initialize the late-phase ensemble.
Despite its apparent simplicity, it remains a challenging optimization problem that has important similarities with stochastic minibatch training in deep neural networks (Schaul et al., 2013; Martens, 2016; Wu et al., 2018; Zhang et al., 2019a;b).
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 η
(Schaul et al., 2013; Martens, 2016; Wu et al., 2018; Zhang et al., 2019a;b).
The late-phase weights training combines the two ﬁrst 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-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, we use linear hypernetworks of arbitrary dimension.
The linear hypernetworks parameterize the weights as w = θ e , with θ ∈ R n × d the hypernetwork parameters and e ∈ R
d the embedding vector.
The embedding vectors e are used as late-phase weights ( φ in the main manuscript) to create a latephase ensemble with K members, while using a shared hypernetwork θ as base-weights: w k =
θ e k .
Ultimately, we are interested in the expected risk of the the mean model at steady state:
[Equation]
(5) with ¯ w (cid:44) 1 K (cid:80) k θ e k = θ 1 K (cid:80) k
e k (cid:44) θ ¯ e and ρ ss the steady-state distribution of the parameters.
Note that we cannot put w ∗ = 0 without loss of generality, because the overparameterization of the hypernetworks makes the optimization problem nonlinear.
We start with investigating the discrete time dynamics induced by late-phase learning, after which we 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. D ISCRETE T IME D YNAMICS
As we want to investigate the multiplicative interaction between the shared and late-phase parameters, we substitute w = θ e 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 θ ) and k the index for the ensemble member.
Then we have the following stochastic minibatch loss:
[Equation]
which gives rise to the following parameter updates using late-phase learning with learning rate η and minibatch size B :
[Equation]
The above discrete time dynamics are nonlinear, giving rise to a non-Gaussian parameter distribution ρ .
Hence, it is not possible to characterize these dynamics by the moment-propagating equations of the ﬁrst and second moment as done in Zhang et al.
(2019a;b); Schaul et al.
and Wu et al.
(2018), without having full access of the parameter distribution ρ .
Furthermore, because of the hypernetwork parameterization, we cannot decouple the system of equations, even if H and Σ are 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.
C. C ONTINUOUS T IME D YNAMICS
First, let us deﬁne some compact notations for the various parameters.
[Equation]
[Equation]
[Equation]
where vec ( θ ) concatenates the columns of θ in a vector.
Then the discrete time dynamics (equationand equation 8) can be rewritten as:
[Equation]
15) with
[Equation]
[Equation]
[Equation]
with ⊗ the Kronecker product, I an identity matrix of the appropriate size and 1 a 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:
[Equation]
with D ( x t ) (cid:44) (cid:0) G ( x t )( I ⊗ Σ) G ( x t ) T (cid:1) 0 .
5 and ζ ∼ N (0 , I ) .
Following Liu & Theodorou (2019) and Chaudhari & Soatto (2018), the corresponding continuous-time dynamics are:
[Equation]
20) with W t Brownian motion and β (cid:44)
Bη the inverse temperature.
Note that √ η is incorporated in the noise covariance, such that the correct limit to stochastic continuous time dynamics can be made (Liu & Theodorou, 2019; Chaudhari & Soatto, 2018; but see Yaida, 2018).
For computing the expected loss E [ L t ] of the mean model, we need to have the stochastic dynamics of this loss.
Using the Itˆo lemma (Itˆo, 1951; Liu & Theodorou, 2019), which is an extension of the chain rule in the ordinary calculus to the stochastic setting, we get
[Equation]
(21) with ˜ D (cid:44) (cid:112) 2 β − 1 D ( x t ) for notational simplicity and H L the Hessian of L w.r.t.
x t .
As we are interested in the expected risk (equation 5), we can take the expectation of equation 21 over the parameter distribution ρ t ( x ) to get the dynamics of the ﬁrst moment of the loss (also known as the backward Kolmogorov equation (Kolmogorov, 1931)):
[Equation]
In order to obtain the dynamics of the parameter distribution, the Fokker-Planck equation can be used (Jordan et al., 1998).
However, due to 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
[Equation]
The remainder of 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 ∇L , 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:
[Equation]
[Equation]
By introducing ∆ e k (cid:44) e k − ¯ e and using that (cid:80) k
∆ e k = 0 , we can rewrite F ( x ) as:
[Equation]
with Γ (cid:44) 1 K (cid:80) k ∆ e k ∆ e Tk and ∆ e T (cid:44) [ e T 1 ...
e TK ] .
We see that F is an approximation of the gradient ∇L where the lower block of ∇L 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 ∇L and the upper block is not inﬂuenced by the number of ensemble members K , which we will need later.
The LHS of equation 23 can now be rewritten as:
[Equation]
with M the diagonal matrix of equation 25 (ﬁrst element of the RHS).
The ﬁrst term of the RHS of equation 26 is the expectation of a weighted squared norm of ∇L , while the second term is an approximation error due to the covariance of ∆ e k .
Hence, we see that the LHS of equation 23 can be seen as an approximation of a weighted norm of the gradient ∇L .
By investigating the term ∇L ( x ) T M ∇L ( x ) further, we show that it is closely connected to the loss L .
[Equation]
When comparing to the mean-model loss L = ( ¯ w − w ∗ ) T H ( ¯ w − w ∗ ) we see that the two are tightly connected, both using a weighted distance measure between ¯ w and w ∗ , with only a different weighting.
Taken everything together, we see that we can take the LHS of equation 23 (and hence 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 inﬂuence of the amount of ensemble members on the expected risk.
Zhu et al.
(2018) highlighted this trace quantitiy in equation 23 as a measurement of the escaping efﬁciency out of poor minima.
However, we assume that we are in the ﬁnal 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 H L and D ( x ss ) 2 as main elements, which we structure in block matrices below (for clarity and ease of notation, we drop the subscripts ss ).
[Equation]
[Equation]
) with 1 a matrix or vector of the appropriate size full of ones, ¯ e (cid:44) 1 /K (cid:80) k e k and the rows of
Q ∈ R d × nd given by:
[Equation]
with δ i the i -th column of an appropriately sized identity matrix.
After some intermediate calculations and rearranging of terms, we reach the following expression for the RHS of equation 23: [Equation]
[Equation]
with ˜ E 2 (cid:44) 1 K (cid:80) k
e k e Tk = 1 K EE T 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 1 K in equation 31 translates to a scaling of 1 K of the expected risk of the NQP, following the above argumentation.
Hence, we see 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 k in parameter space to obtain a mean-model, succeeds in scaling the expected loss after convergence inversely by K .
C. NQP WITH FULL ENSEMBLES
As a comparison for the above theoretical results, we also analyze the NQP that uses an ensemble of
K full weight conﬁgurations w k to get a mean model ¯ w , instead of shared weights θ and ensemblemember-speciﬁc weights φ k .
For the case of linear models, the averaging in weight space to obtain a mean model is equivalent to the averaging of the predictions over the ensemble, which is conventionally done using ensembles.
Without loss of generality, we can take w ∗ = 0 (corresponding with a simple reparameterization of w ).
Using equation 4, this results in the following parameter updates for the ensemble members:
[Equation]
The mean model ¯ w (cid:44) 1 K (cid:80) k w k has the following corresponding discrete dynamics:
[Equation]
[Equation]
Exact moment propagating equations.
As this is a discrete linear system with Gaussian noise, the resulting parameter distributions will also be linear and can be fully characterized by the mean and covariance of the parameters.
Taking the expectation and variance of equation 33 results in:
[Equation]
[Equation]
with Σ the covariance matrix of (cid:15) .
For an appropriate η , the above equations converge to the following ﬁxed points at steady-state:
[Equation]
We see that the steady-state covariance of ¯ w and hence of the risk L scales with 1 K ( E ρ ss
[ L ] =
E ρ ss [ ¯ w T H ¯ w ] =
Tr (cid:2) H C ρ ss [ ¯ w ] (cid:3) ).
The expected risk E ρ ss
[ L ] obtained with computationally expensive full ensembles can be seen as a lower limit that we try to reach with the economical ensembles of shared weights θ and
late-phase weights φ k .
Note that for the NQP, increasing the batchsize B has a similar inﬂuence as increasing the number of ensemble members K , as can be seen in equation 37.
Continuous time stochastic dynamics.
We can also do a similar continuous time analysis as
Section C.2 for the case of full ensembles, to better relate it to the results of the late-phase learning with shared parameters.
Following the same approach, we get the following expression for the trace term:
[Equation]
When comparing to equation 31, we see that the economical ensembles with shared parameters reach the same scaling with 1 K 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 parameters.
Experimental details for Fig.
We take the model w = θ φ (i.e., K = 1 ) as our baseline, since this overparameterization could already result in accelerated learning (Arora et al., 2018).
Our parameters are randomly initialized and scaled such that ¯ w has a ﬁxed distance to w ∗ of 1.
Since the
NQP mimics a late phase of learning we set T 0 = 0 .
We study a problem of dimension n = 100 and train the model with gradient descent (without momentum).
To validate the theoretical results, we show in Fig.
1 that the steady-state reached by our method scales inversely with K , similarly to an ensemble of independently-trained models.
We run experiments with K ∈ [2 , 5 , 10 , 15 , 20 , 25] and train every conﬁguration for 2 × 10 7 iterations until convergence.
We average over the last 10 4 weight updates and over 5 different random seeds.