Trainable Weight Averaging: Efficient Training by Optimizing Historical Solutions
Authors: Tao Li
Zhehao Huang
Qinghua Tao
Yingwen Wu
Xiaolin Huang
Venue: ICLR 2023 poster (2023)
Keywords: efficient training
weight averaging
optimization

Abstract: Stochastic gradient descent (SGD) and its variants are considered as the de-facto methods to train deep neural networks (DNNs). While recent improvements to SGD mainly focus on the descent algorithm itself, few works pay attention to utilizing the historical solutions---as an iterative method, SGD has gone through substantial explorations before convergence. Recently, an interesting attempt is stochastic weight averaging (SWA), which significantly improves the generalization by simply averaging the solutions at the tail stage of training. In this paper, we realize that the averaging coefficients could be determined in a trainable manner and propose Trainable Weight Averaging (TWA), a novel optimization method in the reduced subspace spanned by historical solutions. TWA has much greater flexibility and can be applied to the head stage of training to achieve training efficiency while preserving good generalization capability. Further, we propose a distributed training scheme to resolve the memory burden of large-scale training with efficient parallel computation. In the extensive numerical experiments, (i) TWA achieves consistent improvements over SWA with less sensitivity to learning rate; (ii) applying TWA in the head stage of training largely speeds up the convergence, resulting in over $40\%$ time saving on CIFAR and $30\%$ on ImageNet with improved generalization compared with regular training.

Published as a conference paper at ICLR 2023
TRAINABLE WEIGHT AVERAGING: EFFICIENT TRAIN- ING BY OPTIMIZING HISTORICAL SOLUTIONS
Tao Li', Zhehao Huang', Qinghua Tao?, Yingwen Wu' & Xiaolin Huang"!
Department of Automation, Shanghai Jiao Tong University PESAT-STADIUS, KU Leuven
ABSTRACT
Stochastic gradient descent (SGD) and its variants are considered as the de-facto methods to train deep neural networks (DNNs). While recent improvements SGD mainly focus on the descent algorithm itself, few works pay attention to utilizing the historical solutions substantial explorations before convergence. Recently, an interesting attempt is stochastic weight averaging (SWA), which significantly improves the generaliza- tion by simply averaging the solutions at the tail stage of training. In this paper, we realize that the averaging coefficients could be determined in a trainable man- ner and propose Trainable Weight Averaging (TWA), a novel optimization method in the reduced subspace spanned by historical solutions. TWA has much greater flexibility and can be applied to the head stage of training to achieve training ef- ficiency while preserving good generalization capability. Further, we propose a distributed training scheme to resolve the memory burden of large-scale training with efficient parallel computation.
iterative method, SGD has gone through
the extensive numerical experiments, (i) TWA achieves consistent improvements over SWA with less sensitivity to learn- ing rate; (ii) applying TWA in the head stage of training largely speeds up the convergence, resulting in over 40% time saving CIFAR and 30% on ImageNet with improved generalization compared with regular training. The code of imple- mentation is available https: //cpps://qaae thub. com/ nb nbIt/TWA
INTRODUCTION
Training deep neural networks (DNNs) usually requires a large amount of time. As the sizes of models and datasets grow larger, more efficient optimization methods together with better perfor- mance are increasingly demanded. In the existing works, great efforts have been made to improve the efficiency of stochastic gradient descent (SGD) and its variants, which mainly focus on adaptive learning rates (Duchi et all 2011: Zeiler, 2012: Kingma & Ba, 2015% Loshchilov & Hutter, 2019; Yad alall 202I: Heo all 202T) or accelerated schemes (Polyak] 1%64; Nesterov, 1983; 1988; 2003). As an iterative descent method, SGD generates series of solutions during optimization. These his- torical solutions provide dynamic information about the training and have brought many interesting perspectives, c.g., trajectory (Liciall 2022), landscape (Garipoy et all M0TN). to name a few. In fact, they can also be utilized to improve the training performance, resulting in the so-called stochastic weight averaging (SWA) (Izmailov et all 2018), which shows significantly better generalization by simply averaging the tail stage explorations of SGD. A similar idea could be found in Szegedy et al. (2016), which designs an exponential moving average (EMA) and considers that a heuristic strategy could be better than equivalently averaging. The success of SWA and EMA encourages more in- depth investigations on the roles of historical solutions obtained during the training (Athhwaratkunn ctal. 2018: Nikishun ct all 201: Yang et al. 2019).
In this paper, our main purpose is to utilize historical solutions by optimizing them, rather than using fixed averaging (e.g., SWA) or a heuristic combination (e.g.. EMA). With such an optimized averag- ing scheme, we can achieve higher accuracy using only the solutions in a relatively early state (i.e. the head stage). In other words, we speed up the training and meanwhile improve the performance. The idea of utilizing these early solutions in DNNs training is mainly inspired by two facts. On the
"Corresponding author (xianlinhung@@sjjuucdd....

---

Published as a conference paper at ICLR 2023
Input:
Sampled weights {w; f= Batch size b, Loss
function L: W x X' x ) + R4, Learning rate Output: Model trained with TWA
Orthogonalize {w; } to feilin
Wi
Initialize u (wwa 0. P
(0)
(e1, 62, 62, ...
en):
while not converged do
Sample batch data: B = { yh yb))
(W)
W
Sw2
Compute gradient: g = V w LBs ( (t)
twal
W3
Update weights: w(wa!! w wowa P
(1)
t= t + 1; end while Return w twa.
Figure 1: TWA intuition.
Algorithm 1: TWA algorithm.
one hand, high test accuracy commonly starts appearing at an early stage. For example, PreAct ResNet- 164 model (He et all. 2016) achieves over 80% test accuracy within 10 training epochs on CIFAR-10 10 (Krizhevsky Hinton, 2X19, while requiring to complete the whole 200 epochs reach its final 95% accuracy. This observation also coincides with the recent findings that the key connec- tivity patterns of DNNs emerge early in training aLl 2020, Frankle et all 2020), indicating a well-explored solution space formed. On the other hand, simply averaging the solutions collected at the SWA stage immediately provides a huge accuracy improvement, e.g. over 16% on CIFAR-100 with Wide ResNet-28- 10 (Zagoruyko & KomodakIs, 2016) than before averaging (Izmailoy et all 2018). These facts point out a promising direction that sufficiently utilizing these early explorations may be capable of quickly composing the final solution while obtaining good accuracy.
As the model parameters go through a rapid evolution at the early stage of training, a simple aver- aging strategy with fixed weighting coefficients as in SWA and EMA can result in large estimation errors. We introduce Trainable Weight Averaging (TWA), which allows explicit adjustments for the averaging coefficients tains all sampled solutions during the training and then conduct efficient optimization therein. As optimization in such a subspace takes into account all possible averaging choices, we are able to adaptively search for a good set of averaging coefficients regardless of the quality of sampled solu- tions and largely reduce the estimation errors. The proposed optimization scheme is essentially the gradient projection onto tiny subspace. Hence, the degree of freedom for training is substantially reduced from the original millions to dozens or hundreds (equal to the dimension of the subspace), making TWA enjoy fast convergence and meanwhile immune overfitting, the latter explains that better training accuracy TWA over SWA or EMA can lead to better test accuracy. In extensive experiments with various network architectures with TWA applied to the head stage of training. For instance, we attain 1.5 ~ 2.2% accuracy im- provement CIFAR-100 and 0. 1% ImageNet with over 40% and 30% training epochs reduced, respectively, compared with the regular training.
trainable manner. Specifically, we construct subspace that con-
different tasks, reach superior performance
In summary, we make the following contributions:
We propose Trainable Weight Averaging (TWA) that allows the averaging coefficients de- termined in trainable manner instead of pre-defined strategy. It brings consistent im- provements over SWA or EMA with reduced estimation error.
We successfully apply TWA to the head stage of training, resulting in great time saving (e.g. over 40% on CIFAR and 30% on ImageNet) compared to regular training along with improved performance and reduced generalization gap.
Our TWA is easy to implement and can be flexibly plugged into different stages of training to bring consistent improvements. It provides a new scheme for achicving efficient DNNs' training by sufficiently utilizing historical explorations.
2
METHOD
In this section, we first formulate the optimization target of TWA. Then a detailed training algorithm is introduced, which consists of two phases: Schmidt orthogonalization and projected optimization.
2

---

Published as a conference paper at ICLR 2023
Note in this paper, the model's weights are aligned vector, Le., E IR,, where D denotes the number parameters.
2.1 OPTIMIZATION TARGET
In SWA, weight averaging is simply given by Wswa = w; where n solutions of the network collected at the tail of training are equally weighted. Such an averaging strategy has been proven quite effective with improved generalization ability. However, equally averaging could not always be perfect solution, which motivates some heuristic modifications Both SWA and EMA are fixed averaging strategies, which may not adequately adapt to the head stage of training and would result in estimation errors, due to the fact that early historical solutions have not stepped into a stationary distribution.
weighting strategy, e.g., EMA.
In this paper, propose optimize the averaging coefficients of different weights with the hope of reducing the corresponding estimation error. Specifically, the set of possible TWA solutions considered, i.e., Wtwas can be represented as follows:
A = { [1 w1 + 022 + Og 3n, E [R]
The weight vectors between consecutive solutions could have a high cosine similarity. To decouple them and for better optimization, we will further orthogonalize [w:} } and find of orthogonal bases {e: to support the solution space, i.c., A = {B1e1 + B2e2 + ..... +..4 Bnen, B& E R}. Then we search for good solution Wtwa by optimizing the following problem.
min
ElayinD [C (f (Wtwai 2) x), +
3
P1, Sn
(2)
1=1
s.t. Wtwa = iier + Aze2 + ... Bnen.
where -(..-) is the loss function as in regular training and the second term is regularization coeffi-
cient O. Note that both SWA and EMA are special solutions of (L)) without optimization.
Optimizing over 3: brings benefits in the view of training loss, and good generalization ability could also be expected: regular training, the number of optimization variables is D, which is very large, but in (L). there only averaging coefficients { { to be optimized. The significant dimensionality reduction could benefit better generalization.
2.2 TRAINING ALGORITHM
Instead of directly optimizing 8,, we that there exists a bijection between the coefficient space {5,) 1 E R and the parameter space RR,, i.e., each set of the coefficients is uniquely mapped to one point in the parameter space, which forms complete subspace (with dimensionality n). We could alternatively optimize these coefficients in such a subspace.
We first focus on finding a set of orthogonal bases {e; to span the subspace that covers (w:) This is standard Schmidt orthogonalization, which sequentially takes the following steps:
{ e Wt (wx Te,) e1 (wx Tex) e2 es es e
(we Te-1) ek-1.
(3)
We then initialize the un (0) as one point in the subspace (e.g. Σ w;), and optimize the net- work's parameters therein. Let P [e1, e2, en), such optimization can be easily achieved by projecting the gradient onto the subspace via projection matrix PP We summarize the detailed training procedures of the proposed TWA in Algorithm with an intuitive illustration in Figure The detailed implementation is described in Appendix B
OPTIMIZATION PROVIDES BETTER FLEXIBILITY
The key difference between SWA and TWA is that the averaging coefficients in TWA are determined trainable manner, or more precisely, are data-dependent. This potentially enables more precise estimation for the center minima and better tolerance for the outliers that are not aware by SWA.
3

---

Published as a conference paper at ICLR 2023
GPU1
P1
Ps.
Ps.0
Pt.P.Pio
PL.PP's
6PU
02
P
Pa.PoPio
Pa: PP
:
PA
Pa: Ak
P
Pu.PxPE0
Pa.PPT
Divide
Local Gradient
All-reduce
Local Projection
All-reduse
Figure 2: An efficient parallel scheme for distributed training.
Notice that in the view of coefficient optimization, there is no essential difference between SWA and EMA, which both provide specific and data-independent solutions. Thus, we the following only compare TWA with SWA.
Mand ct al (200 Z demonstrated that under appropriate assumptions, running SGD with a constant learning rate is equivalent to sampling from stationary Gaussian distribution, and the variance of the distribution is controlled by the learning rate. Accordingly, we assume the solutions at the tail stage of SGD training sampled from Gaussian distribution N (11, Σ) centered at the minimum with covariance E. Approximately, the sampled solutions (w, )= are independent random vari- ables from N (1, E) as long as there are sufficient iterations between adjacent samplings. SWA and TWA provide two estimators for the minimum i.c., Wewa and Wtwa. As an averaged solution, Wawa has statistically better estimation than any single solution due the effect of variance reduc- tion, while Weaa is approaching the center by minimizing the training loss. As long as the training loss serves as meaningful supervision (which holds under the typical assumption that fu is the center minimum with lowest training loss Mandt all (2017); Izmailoy et all (2018)), Wrwa could approach fu better than ussn with the posterior optimization for averaging coefficients. In this regard, Wtwa could have a lower expected variance: E (I|Wiwa =1 2 I ( Wawa
m112).
The advantages of optimizing averaging coefficients could be more prominent in the head stage of training. where the weights are going through rapid evolution. A simple averaging strategy as SWA could introduce a large estimation eтror (as illustrated in Table B and {)) while TWA enables correcting it to smaller estimation eтror via optimization. In fact, TWA provides much more flexibility to sufficiently utilize historical solutions and produces an optimized solution adaptively.
AN EFFICIENT IMPLEMENTATION FOR DISTRIBUTED TRAINING
The above discussion shows promising improvement by optimizing the historical solutions. The only issue one may worry about is the burden in storage (the additional time complexity is small as shown in Table R) During optimization, TWA requires the projection matrix P involving dozens hundreds of historical weights, which indeed poses a challenge for large models on storage burden. locate P in GPUs enable efficient matrix operations. However, the size of
It is preferable P increases as the model becomes larger, making it prohibitive to store in single GPU, To cope with this, we design an efficient scheme with parallel distributed training to enable a) partition the memory burden of P into multiple GPUs and b) efficient parallel computation of gradient projection. As result, we successfully optimize more than 900 historical solution coefficients for ResNet-50 on ImageNet task by 4 v100 GPUs. In our experiments, we use at most 300 historical solutions and there is still available space larger tasks.
Suppose that there are 1a GPUs for parallel training. We first uniformly divide P into k sub-matrices as P [P,, P2, ... Ps). where each GPU stores a local sub-matrix P, 1,....k Recall that for iteration in distributed training, each GPU computes a local gradient g: and synchro- nizes it with other GPUs to obtain the global gradient through an efficient all-reduce operation (Rabenseither, 2013). We mimic such process for gradient projection: the local projected gradi- ent P.P. is firstly computed in each card and then synchronized with others to obtain the global projected gradient with another all-reduce operation. We illustrate such process in Figure E1
For averaging n historical solutions with per size B, the memory burden for each GPU card is reduced to [n/k] /k] B, while the computation of gradient projection is also reduced to ([nn/k) D) (D
4

---

Published as a conference paper at ICLR 2023
CIFAR-20
PreResWeet
ProResNet
Accuracy (%5
90
70
60-
9mail
815
5G0
Bnw
TWA
40
TWA
500
32D
80
10
50
100
10
50
100
10
100 10
50
100
Epochs
Epochs
Epochs
Epochs
Figure 3: Performance comparisons on before and after TWA w.I.t different epochs of weights used. "SGD final" indicates the accuracy reached by regular SGD training and "TWA" the corresponding accuracy reached by Algorithm II with these epochs of weights. The final accuracy of SGID training is plotted for reference. TWA dramatically lifts the SGD accuracy and outperforms the final accuracy of SGD within 100 epochs. The experiments are repeated over 3 trials.
the number of the parameters). Hence, we can achieve efficient TWA training by making full use of the remaining memory of each GPU aside the forward/backward training.
EXPERIMENTS
In this section, a series of numerical experiments are conducted, demonstrating the effectiveness of our proposed TWA for fast convergence and better performance. First, show that TWA improves SWA in the existing SWA settings, i.c., used in the tail stage of training. Second, we apply TWA the head stage of training, which brings significant efficiency improvements together with better performance. Then, we visualize loss accuracy surfaces to demonstrate the improvements of TWA.
5.1 EXPERIMENTAL SETTINGS
Datasets. We experiment over three benchmark image datasets, i.e., CIFAR-10, CIFAR-100 (Knzhevsky & Hinton, 2009),, and ImageNet (Deng et al., 2009). Following prior works (imaally atall 2018 Yang et all, 2019) we apply standard preprocessing for experiments on CIFAR datasets, and adopt the preprocessing and data augmentation procedures in the public Pytorch example on
ImageNet (Paszke.cLal) 20177).
Architectures. We use two representative architectures, VGG-16 (Smonnyan & Zisserman, 2IEE) and PreAct ResNet-164 (Heetall 2016) on CIFAR experiments. For ImageNet, we use ResNet-18 and ResNet-50 (He et all 2016).
Training. The main body of experiments contains two parts: (1) for the tail stage of training, we use the same Syyper-parameee
in SWA (Izmarloy ci all 201) and then larger tail learning rate is also tried. (2) for the head stage of training, we adopt the standard training protocol with a step-wise learning rate. For CIFAR, we all experiments with 3 seeds and report the mean test accuracy. We use SGD optimizer with momentum 0.9, weight decay 10-4, and batch size 128. We train the models for 200 epochs with an initial learning rate 0.1 and decay by 10 at the 100th and the 150th epochs. For ImageNet, we follow official PyTorch implementation!! For TWA, we sample solutions once after each epoch training for CIFAR and uniformly sample 5 times per epoch for ImageNet. We use a scaled learning rate (Figure B) which takes 10 epochs of training for CIFAR and 2 epochs for ImageNet for fast convergence. The regularization coefficient À defaults to 10 More details (including the number of the historical solutions used) could be found in Appendix A
5.2 IMPROVING SWA SOLUTIONS
In this part, we compare SWA and TWA in the original SWA settings. Specifically, TWA and SWA use the same weights sampled from the tail stage of training. For CIFAR, we try two different tail learning rates: 0.05, the recommended one in (izaallyy et all 201), and 0. 10, a larger one for the case with greater variance. The results in Table M show that TWA brings consistent improvements
Available httph:t////// htspa://githtb.. httphhwwww.ooooptwcohhchccecceee http:///tthtcccoo/o/ooooo/////// Secmmplls/ee
magener
5

---

Published as a conference paper at ICLR 2023
Condersed leaming
VEG16 CIFAR-10
PreResNet-1 64 CHAR-100
Irei epachsl
Leaming rate
Ife2 cpachs)
Ire4 15 spocha]
Accuracy
93.6
Accuracy
93.4
77
apecis) Ined 15 epochs)
In=1 apochs) I- spochal
93.2
Inud 15 epochs)
93.0
76
20
Epachs
Epochs
Scaled learning rate schedules with different scaling factors; Middle and Right: Test accuracy curves of TWA w.r.t. different schedules CIFAR-I0/100. Training in subspace shows high robustness to scaled learning rate, which enlarges the learning rate and reduces the corresponding training epochs. this way, TWA achieves very fast convergence.
Figure 4: Left:
over SWA. Especially when the learning rate is not well-tuned, SWA's performance suffers a dis- tinctive drop, but TWA is less sensitive since the estimation enror could be well controlled through training. For instance, in the case of CIFAR- 100 with VGG-16 and tail learning rate of 0.10, the estimation error of SWA substantially increases while TWA achieves significant accuracy im- provement, i.e., 4.18%, over SWA. Note that for a fair comparison, TWA starts from the last sampled weights, not the averaged solution SWA.
Table 1: Test accuracy (%) on 1IFAR-10/109 for tail training with different learning rates
SWA LR 0.05
SWA_LR = 0.10
DATASET
MODEL
SWA
TWA (+10)
SWA
TWA (+10)
VGG16
94.01 0.04
94.16 0.14
91.03 0.14
92.41 0.15
CIFAR-10
PRERESNET-164
95.58 0.09
95.65 + 0.13
91.58 0.45
92.61 0.09
VGG16
74.71
75.73 0.18
65.52 0.25
69.70 0.45
CIFAR-100
PRERESNET-164
80.20 0.41
80.35 = 0.27
78.14 = 0.48
78.87 + 0.27
On ImageNet, we experiment with ResNet-18/50 (He el all 2016). Following (Izmarloy et all 2018: Yang et all 20T),, start from pre-trained models in Eorchvisi models and collect model weights by running SGD optimizer up to 10 epochs (with a constant learning rate 0.005). In Table Q report the test accuracy and observe that with more sampling epochs, both TWA and SWA achieve better performance. Notably, TWA performs better than SWA by 0.1 ~ 0.3%, and such improvements of sampled weights, TWA achieves 70.23% accuracy with ResNet-18 and 76.78% accuracy with ResNet-50, outperforming the SWA counterparts with 10 epochs. This indicates that TWA requires fewer historical solution samples to achieve a comparable or even better performance than SWA, due to its ability to reduce the estimation variance with optimized averaging coefficients.
more obvious in the "5 EPOCHS"
For example, using epochs
Table 2: 1 accuracy (%) on ImageNet for tail training with different averaging epochs
EPOCHS
10 EPOCHS
MODEL
PRETRAINED
SWA
TWA (+1)
SWA
TWA (+1)
RESNET-18
69.76 76.13
70.02
70.28 76.78
70.12 76.74
70.32
RESNET-50
76.62
76.93
5.3 EFFICIENT TRAINING AND BETTER GENERALIZATION
In the head stage of training, SWA usually fails due to the large estimation variance from fast- evolving solutions and large learning rate. Since TWA could reduce the variance and be less sensitive to the learning rate, it can also be expected to work well in this stage. If so, it is promising to simultaneously attain generalization improvements and training efficiency.

---

Published as a conference paper at ICLR 2023
We first investigate the experiments CIFAR datasets. The original training schedule contains 200 epochs and we take the first 100 epoch explorations for TWA. The results are given in Table BL It can be observed that TWA achieves better performance compared to the regular SGD training with a significantly reduced generalization gap. For instance, we attain 1.52% accuracy improvement on CIFAR-100 100 with VGG-16 while the generalization gap is reduced by 9.56%. This suggests that a better solution could already be composed using these historical solutions without further training by more delicate learning rates, which instead may bring overfitting problems and harm the generalization. By comparisons, we also apply SWA to average these samples, which shows degraded performance due to the existence of estimation error. Apart from the good performance in accuracy, TWA also manifests its great potential in improving the training efficiency: we use only 10 epochs to complete the convergence, while the regular SGD needs 100 epochs. As TWA and SGD have nearly the same computation overhead per epoch, the time-saving is around 45% in TWA.
Table 3: Test accuracy (%) and generalization gap (%) on CIFAR- 10/100 for the head training
SGD (200 EPOCHS)
SWA (100 EPOCHS)
TWA (100 + 10 EPOCHS)
DATASET
MODEL
ACCURACY
GAP
ACCURACY
ACCURACY
GAP
VGG16
93.54 0.11
6.42
92.40 0.08
93.79 0.18
5.59 ( 0.83) 4.11 (4 0.75)
CIFAR-10
PRERESNET-164
95.11 - 0.17
4.86
92.52 = 0.04
95.19 0.04
VGG16
72.72 + 0.17
26.70
69.18 + 0.24
74.24 0.24
17.14 (4 9.50)
CIFAR-100
PRERESNET-164
75.85 + 0.18
24.10
73.36 + 0.34
78.11 + 0.19
20.02 (1 4.08)
Generally, utilizing more epochs of explorations can provide better estimation for the center min- imum and hence lead to better performance. Then, we also study the impact of different averaging epochs, and the results are illustrated in Figure & where the final accuracy of SGD and the accuracy reached by SGD before averaging are also given for reference. It could be observed that the model's performance is consistently improved with more epochs of explorations. Notably, although each historical solution in relatively short period of explorations is not good, satisfied solutions have already emerged in the subspace spanned by these solutions. Then through proper optimization in subspace, TWA could find them out, e.g., on CIFAR-100 with PreAct ResNet-166 model, averaging over 50 epochs via TWA has already matched the final performance of regular SGD training.
Table 4: 1op-1 accuracy (%) and generalization gap (%) on ImageNet for head training
SGD (90 EPOCHS)
SWA (60 BPOCHS)
TWA (60 + EPOCHS)
MODEL
ACCURACY
GAP
ACCURACY
ACCURACY
GAP
RESNET-18
69.82
-1.59
62.19
69.82
-2.36 (1 0.77)
RESNET-50
75.82
0.25
67.66
75.90
-0.68 (4 0.93)
For ImageNet, the efforts required for each epoch training are much greater, and hence efficient meth- ods to reduce the training epochs are highly desirable. The comparison results of SGD/SWA/TW/ are shown in Table E Besides the reduced generalization gap. TWA takes only 2 epochs to average the historical solutions of the first 60 epochs, reaching comparable even better performance than regular SGD training with 90 epochs. For comparison, Lookahead (Zhang et all 201) is another advanced optimizer recently proposed for improving convergence and reported 75.49% accuracy at the 60th epoch (Table 2 in Zhang et all (2019) with an aggressive learning rate decay (i.e., the learning rate is decayed the 30th, 48th, and 58th epochs), while our TWA reaches 70% with the same budget but simply using the conventional decay.
TWA is very flexible and can be readily applied to different training stages, and we also conduct an ex- periment by averaging the solutions of the final train- ing period (i.e. 61-90 epochs) and simply performing TWA for one epoch training, as presented in Table S Such cheap training still shows to bring significant improvements (c.g. +0.51% on ResNet-50 for Ima-
Table 5: Top-1 accuracy (%) on ImageNet for tail training and short TWA
MODEL SGD
TWA (+1 EPOCH)
RESNET-18 69.82
70.37 76.34
RESNET-50 75.82
geNet). Thus, TWA can serve as effective approach for composing better final solution.
7

---

Published as a conference paper at ICLR 2023
Scaled learning rate The optimization of TWA is conducted in very low-dimensional space, which also suppresses the sensitivity of the learning rate. In fact, we can allow a very large learning rate to accelerate the training. Thereby, we design scaled learning rate, which linearly scales up the learning rate and reduces the training epochs accordingly, as shown in Figure E Within an appropriate range, scaling the learning rate largely speeds up the convergence without affecting the final performance. For example, with the learning rate of 4, TWA approaches the final accuracy with only epoch and converges within 5 epochs.
Wall-clock time comparison In the above experiments, we report the number of training epochs, since the wall- clock time per epoch for SGD and TWA is similar. As an illustration, Table E provides the wall-clock time of train- ing PreAct 1664 for CIFAR-100 on one Nvidia Geforce GTX 2080 TI.
Table 6: Wall-clock time per epoch
OPTIMIZER
TIME PER EPOCH
SGD TWA
59.04S 60.20s
Comparison with EMA EMA serves an alternative to SWA, which averages the model weights along the training trajectory with exponential decay. It requires hyper-parameter y to control the averaging horizons. As a manually defined averaging strategy, EMA could be sensi- tive to learning rates, datasets, architectures, etc. Here we
PreResNet-164 on CIFAR-100
80
Accurary (96)
70-
TWA
TMC:
compare the performance of EMA with SWA and TWA in
5MA
LMC:
the head stage of training. where we try = 0.99/0.999. The results are illustrated in Figure B We observe that the performance of EMA varies significantly with differ- ent choices of ?.. It could perform notably better than SWA in the early stage of training as more weight is paid to the latest solutions. Note that both EMA and SWA are fixed averaging strategies for adapting different training stages (essentially could be viewed as particular solutions of TWA). By optimizing the averaging coefficients, TWA could consistently achieve better performance.
SCD
50
20
60
80
100
Epachs
Figure 5: Comparisons with EMA.
Optimized Averaging Coefficients In Figure & we visualize the averaging coefficients oş learned by TWA. Detailed derivation could be found in Appendix D We observe all historical solutions could contribute to the
Prelessdei-166 on CIFAN-100
8:14
Cooffidens
6.00
final solutions. Solutions from the latter training stage are attached with more importance as expected. Dif-
-802
ferent from the fix averaging strategies like SWA or EMA, such averaging coefficients enable to take full ad- vantage of the historical solutions through delicate opti- mization and better adapt to the training dynamics.
100
Spochs
Figure 6: Averaging coefficients of TWA.
5.4 LANDSCAPE VISUALIZATION
Following (arrpov et al. 2IT: Izmanlov et all 2018). visualize the training loss and test error surfaces of SWA and TWA in Figure @ on CIFAR- 100 with PreAct ResNet-164. We set the SGD solution after 125 training epochs the origin and plot the TWA and SWA solutions
the plane. For the case with a default learning rate of 0.05, TWA achieves slightly better test accuracy with lower training loss. This shows that in the subspace, minimizing the training loss is meaningful and results in lower test errors. Especially for the case with larger learning rate of 0.10, the superiority of TWA over SWA is more significant (over 0.7% improvement on test accuracy), since the variance grows larger and the variance reduction effect of TWA becomes more obvious.
6 RELATED WORK
Improving the model's 's generalization capability is of great importance and has received wide atten- tion. The recent efforts mainly focus two aspects: (1) proper regularization terms to search for more flat minimum (Keskar all 2016; Lictall 2018), such as weight decay (Krogh & Hertz, 1991).
dropout (Srivastava et all 2014), label smoothing (Szegedy all, 2016), Shake-Shake (Gastald
8

---

Published as a conference paper at ICLR 2023
Training loss
Test error (%6
12
524
10
10
232
27.1
514
33.6
-2
(a) SWA 0.05
Training
error (5)
24
1.35
42.0
5/1
312
+2
(b) SWA LR
0.10
Figure 7: Train loss and test surface of TWA and SWA with different SWA LR.
2017), MixUp (Zhang et all 2018), SAM etll 2020) and AMP (Zheng et al, 2021);; (2) effec- tive data augmentation to diversify the dataset, such as Cutout De & Taylor, 2017), AutoAug- ment (Cubuu et al, 2019) and RandAugment (Cubuk et all (020). Different from these techniques, we improve the generalization ability by constraining the training in a low-dimensional subspace spanned by historical explorations, which regularizes the model complexity. We note that TWA is orthogonal to these methods, and it promising to combine them for boosted improvements.
A lot of efforts have been made to speed up the DNNs' training. Apart from the well-known methods on adaptive learning rates, e.g. Adam (Kingma & Ba, 2015) and accelerated schemes, e.g. Nesterov momentum Nesteroxi (1983), new method is proposed in [Zhang all (2II)) proposed, where a look-ahead search direction generated by another "fast" optimizer is utilized, achieving faster convergence and better learning stability. Goyal et al. (2017) adapted large mini-batch to speed up the training and introduced scaling rule for adjusting the learning rates. In this paper, we realize training efficiency by sufficiently utilizing the historical solutions of DNNs' training and conducting training in subspace with substantially reduced dimensions.
For utilizing historical explorations, SWA (Izmailov et all 2018) adopts a simple averaging strategy at the tail of training. Cha et all (2021) extended it to the domain generalization task with dense and overfit-aware stochastic weight sampling strategy. We firstly propose to utilize the explorations at the head stage of training to achieve training efficiency. Exponentially decaying running average (Hunter, 1986 Szegedy et all 2016) a common technique adopted by practitioners. It requires manually set averaging horizon and generally performs comparably SGD (Izmailov et all 2018). Another closely related work is model soups (Wortsman et all, 2022), which improves the model performance by averaging the weights from different fine-tuning configurations in a greedy order. We differ in that the historical solutions are from one single configuration. We mainly focus on improving training efficiency and optimizing the averaging coefficients in trainable manner.
CONCLUSION
In this work, we propose TWA, a novel training algorithm that optimizes the averaging coefficients of historical solutions in DNNs' training to achieve efficiency and better performance. It differs from the manually set averaging strategies as SWA or EMA and manifests better adaptation to different stages of training. We further design a parallel framework for large-scale training with efficiency in memory and computation. Extensive experiments demonstrate the superior performance TWA on benchmark computer vision tasks with various architectures.

---

Published as a conference paper at ICLR 2023
ACKNOW
We are very grateful for anonymous reviewers for the valuable feedback on the paper. We thank Minqi Chen at Huawei Technologies for the great support. The research leading these results has received funding from National Natural Science Foundation of China (61977046), Shanghai Science and Technology Program (2251 1056000, and Shanghai Municipal Science and Technology Major Project (2021SHZDZ200102)
REFERENCES
Ben Athiwaratkun, Mare Finzi, Pavel Izmailov, and Andrew Gordon Wilson. There are many con- sistent explanations of unlabeled data: Why you should average. In International Conference on Learning Representations (ICLR), 2018.
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lec, and Sungrae Park. Swad: Domain generalization by seeking flat minima. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoe V Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition (CVPR), pp. 113-123, 2019.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with reduced search space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pp. 702-703, 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: large-scale hierarchical image database. In Proceedings of the IEEE Conference Computer Vision and Puttern Recognition (CVPR), pp. 248-255, 2009.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04880 2018.
Terrance Dries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: 04552, 2017.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research (JMLR), 12(7), 2011.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- tations (ICLR), 2020.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In International Conference Machine Learning (ICML), 2020.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in Neural Information Processing Systems (NeurIPS), 2018.
Xavier Gastaldi. Shake-shake regularization of 3-branch residual networks. In Workshop Track Proceedings in International Conference Learning Representations (ICLR), 2017.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: 1766:0700600677 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IÉEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
10

---

Published as a conference paper at ICLR 2023
Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung- Woo Ha. Adamp: Slowing down the slowdown for momentum optimiz- on scale-invariant weights. In International Conference Learning Representations (ICLR),
2021.
J Stuart Hunter. The exponentially weighted moving average. Journal of Quality Technology, 18(4): 203-210, 1986.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.054007 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. arXiv preprint arXiv:16609 04866, 2016.
Diederik P. Kingma and Jimmy Lei Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations (ICLR), 2015.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Tech- nical Report, 2009.
Anders Krogh and John Hertz. simple weight decay can improve generalization. Advances in Neural Information Processing Systems (NeurIPS), 1991.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. Advances in Neural Information Processing Systems (NeurIPS), 2018.
Tao Li, Lei Tan, Zhehao Huang, Qinghua Tao, Yipeng Liu, and Xiaolin Huang. Low dimensional trajectory hypothesis is true: Dnns be trained in tiny subspaces. IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI), 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1000501711.05101, 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations (ICLR), 2019.
Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent approximate bayesian inference. Journal of Machine Learning Research (JMLR), 18:1-35, 2017.
Yurii Nesterov. On approach to the construction of optimal methods of minimization of smooth convex functions. Ekonomika Mateaticheskie Metody, 24(3):509-517, 1988.
Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2003.
Yurii E Nesterov. A method for solving the convex programming problem with convergence rate (1/k2 2). In Dokl. akad. nauk Sssr, volume 269, pp. 543-547, 1983.
Exgenii Nikishin, Pavel Izmailov, Ben Athiwaratkun, Dmitrii Podoprikhin, Timur Garipov, Pavel Shvechikov, Dmitry Vetrov, and Andrew Gordon Wilson. Improving stability in deep reinforce- ment learning with weight averaging. In Uncertainty in Artificial Intelligence Workshop Un- certainty in Deep learning, 2018
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr Compu- tational Mathematics and Mathematical Physics, 4(5):1-17, 1964.
Rolf Rabenseifner. Optimization of collective reduction operations. In International Conference on Computational Science, pp. 1-9. Springer, 2004.
11

---

Published as a conference paper at ICLR 2023
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint vrXiv:1409.1555 2014.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. Journal of Machine Learning Research (JMIR), 15(1): (51):1992-1995, 2014.
Christian Szegedy. Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition (CVPR), pp. 2818-2822 2016.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461, 2018.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning (ICML), 2022.
Guandao Yang, Tianyi Zhang, Polina Kirichenko, Junwen Bai, Andrew Gordon Wilson, and Chris De Sa. Swalp: Stochastic weight averaging in low precision training. In International Conference on Machine Learning (ICML), pp.
PMLR, 2019.
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael W. Mahoney. ADAHESSIAN: an adaptive second order optimizer for machine learning. In Thirty-Fifth AAAI Conference Artificial Intelligence, AAAI, 2021.
Haoran You, Chaojian Li, Pengfei Xu, Yonggan Fu, Yue Wang, Xiaohan Chen, Richard G Baraniuk, Zhangyang Wang, and Yingyan Lin. Drawing early-bird tickets: Towards more efficient training of deep networks. In International Conference on Learning Representations (ICIR), 2020.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Richard C. Wilson, Edwin R. Hancock, and William A. P. Smith (eds.), British Machine Vision Conference (BMVC), 2016.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv: 2012.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations (ICLR), 2018.
Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead optimizer: k steps forward, step back. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8156-8165, 2021.
A TRAINING DETAILS
For SWA experiments, we replicate the SWA baseline by using the publicly released implementa- tion of Izmailoy at (2018). We use VGG-16 architecture with batch normalization for a unified learning rate setting PreAct ResNet-164. For ImageNet experiments, implementatron. We use a scaled learning rate for TWA training with 20x and 30x factors on CIFAR and ImageNet, respectively (e.g. the original learning rate of 0.1 is scaled up 2 on CIFAR). CIFAR experiments are performed on one Nvidia Geforce GTX 2080 TI GPU, while ImageNet experiments four NVIDIA Tesla A100. The number of historical solutions optimized by TWA is presented in Table A
follow official PyTorch
We now disclose the specific hyper-parameters in the following.
12

---

Published as a conference paper at ICLR 2023
Table Al: The number of historical solutions optimized by TWA
DATASETS
SAMPLING EPOCHS
TIMES EPOCH
#NUM
CIFAR
100
100
IMAGENET
60
5
300
A.1 SWA TRAINING
A.1.1 CIFAR
We use the same schedule and hyper-pann parer-parameters as in Izmailoy el (2018). For VGG-16, 16, we use weight decay of 5 10 and train the model for 300 epochs with weight averaging 161 300 epochs. For PreAct ResNet-164, we use weight decay of 3 x 10-4 and train the model for 225 epochs with weight averaging at 126 225 epochs.
For TWA training, we use the same weights as SWA and initialize Wowa as the last checkpoint (i.c. 300 225 epochs). We train the models for 10 epochs with an initial learning rate of 2 and decay by 10 at the 5th and 8th epochs. The regularization coefficient À is set to 5 x 10-5
A.1.2 IMAGENET
Following imailov et all (2018); Yang et all (2011), we start from pre-trained models (they are from Eerchv siinn mode (s) and collect weights by rumning SGD optimizer up 10 epochs (with constant learning rate 0.005, weight decay x 10-1). We uniformly sample the solutions 5 times
per epoch.
For TWA training, we use the same weights SWA and initialize Wtwa as the pre-trained model. For ImageNet, there are many iterations in one epoch, and hence we conduct TWA training for one epoch, in which we linearly decay the learning rate from 0.03 to O. The regularization coefficient À set to x 10-5
A.2 REGULAR TRAINING
A.2.1 CIFAR
For regular training, we train the models for 200 epochs with an initial learning rate of 0.1 and decay by 10 at the 100th and the 150th epoch. We use SGD optimizer with momentum 0.9, weight decay 1 x 10-4, and batch size 128 by convention.
For TWA training, we initialize Wowa as E wi, i.e., the center of sampled solutions. We train the models for 10 epochs with an initial learning rate of 2 and decay it by 10 at the 5th and 8th epochs. The regularization coefficient À to 1 10-5
A.2.2 IMAGENET
We follow the training protocol described He (2016). Specifically, we train the models for 90 epochs with an initial learning rate of 0.1 and decay it by factor of 10 every 30 epochs. We use SGD optimizer with momentum 0.9, weight decay x 10-4 and batch size 256.
uniformly sample solutions 5 times per epoch and initialize Wtwa as 1 DE u. We train the models for epochs with a learning rate of 0.3 and 0.03, respectively. For the extra epoch training, use the same training protocol in subsection A. i.c. linearly decaying the learning rate from 0.03 The regularization coefficient is set x -
For TWA training,
A.2.3 RESULTS ON ADAM OPTIMIZER
Adam (Kinmma & Ba 201)) is another mainstream optimizer with adaptive gradient descent, which enjoys fast convergence and insensitivity the initial learning rate. Here, apply TWA to the solutions generated by Adam optimizer and the results are in Table A2I The training settings are the same in Table E with default 0.9 and 8 0.999 for Adam. TWA is trained for 5 epochs
13

---

Published as a conference paper at ICLR 2023
with intial learning rate 3. We observe that it could similarly bring generalization improvement and training efficiency.
Table A2: Test accuracy (%) and generalization gap (%) CIFAR-10/100 with Adam Optimizer
ADAM (200 EPOCHS)
SWA (100 EPOCHS)
TWA (100 + 5 EPOCHS)
DATASET
MODEL
ACCURACY
GAP
ACCURACY
ACCURACY
GAP
VGG16
93.60 = 0.12
6.31
92.39 = 0.07
93.63 + 0.11
5.74 (4 0.57) 3.85 (1 1.04)
CIFAR-10
PRERESNET-164
95.09 + 0.11
4.89
92.65 + 0.08
95.16 + 0.09
VGG16
72.77 1 0.13
26.67
68.74 + 0.05
74.07 1 0.07
17.35 (1 9.32)
CIFAR-100 100 PRERESNET-164
76.23 = 0.14
23.71
73.09 = 0.13
77.96 + 0.14
19.62 (1 4.09)
B IMPLEMENTATION
Let B = 82, ...... ... , ; ER" P (er, .z, eea] E Rxn, the optimization target for TWA can be formulated as,
A y)] =
min L(3) e E( (f (Wewai x)
TT B.
(A.1)
s.t. Utwa Pg.
short the loss term E (m [C (f (w: x) y)l to [(w) and the gradient w.r.t. 3 can be written as,
0L Swowa OL(Wowa) XB
(A.2)
03
0g
Owewa
p T OL (Wow)) X
(A.3)
Owtwa
Let be the learning rate. We could optimize
AA with the following gradient descent:
@l(W(wa)
g(++1) = =
P
Wwaw)
XB
(A.4)
Owewn
Since Wewa PB, we have the corresponding update in the parameter space:
Wewa
(1+1)
wiia
(1)
17
PP
(((w))) twa) Owewn
(t)
+ XP3
(A.5)
( (2)
(1 17) w (t) Wtwa
pP
(A.6)
Owiwa
As 3 does explicitly appear in (A.f), we could treat the coefficient as an implicit variable. In practice, we optimize (A.I) by directly updating the model weights Wiwa with weight decay X, which is optimization in the reduced subspace with projection matrix PP!
C SWA WITH DIFFERENT STARTING EPOCHS
We test the performance of SWA with different starting epochs to average on ImageNet with the ResNet-50 model. observe that the performance of SWA gradually becomes better with the relatively latter stage of solutions averaged, showing that SWA could not adapt well to the head stage of training where the solutions are fast-evolving. Hence, a good solution for SWA may require manually selecting which period to average. We also notice that TWA (with 0-60 epoch solutions) consistently outperforms the SWA wherever the averaging begins, confirming that TWA could auto- matically find good set of averaging coefficients and provide better performance.
14

---

Published as a conference paper at ICLR 2023
SWA Epoch 0-60
SWA Epoch 10-60
SWA Epoch 20-60
SWA Epoch 30-60
SWA Epoch 40-60
SWA Epoch 50-60
TWA
Epoch 0-60
67.66
70.50
72.12
74.14
75.08
75.34
75.90
Table C3: SWA with different starting epochs.
D DISCUSSION ON THE SUM-ONE CONSTRAINT
Let [a a2,,... a2,.02
Opp R" W [w) W2.... : W.. DDx, we have Wowa Wa. We multiply W on the both sides, i.c., W Wtwa = W Wa, and could obtain (WTW) W) W Wtwa- Further, could establish the relation between and B: a (W W) W P3, since they are the coordinates of Wtwa under two different set of bases.
In the solution set A = {a1w] + a2 + ... + 0t Un a, R], we do not explicitly require that the sum of averaging coefficients Σ a; to be since the network's performance would be sensitive to direct scaling over all parameters (i.e. kwr), that is, good solution in A will inherently have a coefficient sum very close to We verify the sum of the averaging coefficients of the attained solution Urwa on CIFAR-100 with PreResNet-166 model and observe that Σ ag 1.02+<0.01-
E NUMERICAL DISCUSSION FOR DDP TRAINING
We numerically measure the averaged epoch training time and memory burden for SGD and TWA in the DDP training setting. Specifically, we experiment with the ResNet-50 model on ImageNet and use 1, and 4 GPUs with batch size of 256 per GPU and total of 300 historical solutions. The experiments are conducted on NVIDIA Tesla A 40G GPUs. From the results reported in Table
20 we observe that TWA brings minor additional costs, e.g. +2.8% time cost and +2.9% memory burden with 4 GPUs, compared with regular SGD training. The additional memory burden becomes even minor with more GPUs. This shows that TWA could provide efficient and scalable averaging for large-scale problems.
Table E4: Time and memory comparisons of SGD and TWA with DDP training.
Time
Memory
#GPUs
SGD
TWA
SGD
TWA
1638s
1692s
(43.3%)
28286.5 MB
31432.5 MB (+11.1%)
824s
862s (+4.6%)
28382.5 MB
30092.5 MB (+6.0%)
420s
432s (+2.8%)
28874.5 MB
29718.5 MB (+2.9%)
ABLATION STUDY
We conduct ablation study in Table analyze the impact of the regularization coefficient A. We observe that such regularization brings improvements but is not significant. This is because the main regularization effects come from the significant decrease of training variables, i.e., regular training has D variables but TWA contains only n. Since such regularization is easy to implement and virtually brings little training cost, we include our method.
15

---

Published as a conference paper at ICLR 2023
Table F5: Ablation studies on the regularization coefficient A.
Datasets
Model
TWA
TWA (A
VGG16
93.79 + 0,18
93.78 0.07
CIFAR-10
PreActResNet-644
95.11 = 0.04
95.03 = 0.09
VGGI6
74.24 + 0.24
74.02 0.16
CIFAR-100
PreAcRReNNet-64
78.11 + 0.19
78.01 + 0.22
G ADDITIONAL RESULTS ON NLP TASKS
For NLP datasets, we try finetune task with pre-trained models and compare the performance of SWA and TWA. Specifically, we experiment with The Corpus of Linguistic Acceptability (CoLA), text classification task in the General Language Understanding Evaluation (GLUE, Wang et al (2018) benchmark. In experiment, we use a pre-trained BERT (Devlin et all 2013) model (bert-base-unczsed) from Hugging Face community We fine-tune BERT CoLA for epochs with AdamW optimizer Loshchilov & Hutter (201Z), learning rate 2e-5, and weight decay 0.0. The model weights at the end of these epochs are collected for SWA and TWA. TWA, train the fine-tuned model for epoch with a learning rate of 0.5 and regularization coefficients 0.001. From the results below, we observe that TWA could achieve better performance than the competing methods. This further demonstrates the broad application of TWA.
Fine-tune
SWA
TWA
56.81
59.73
60.36
Table G6: Fine-tune results on CoLA.
Availlabe http:::/huugigggggiate
16


Reviews:
[{"review_id": "xjtQHMmjb9", "reviewer": "Reviewer_Tq6j", "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": "The paper proposes to perform trainable weight aggregation on the historical models, so as to speed up the convergence while preserving good generalization capability (over SWA).\n\nThe proposed TWA can outperform SWA to some extent. However, the presented results may have the issues of limited novelty, unfair comparison, etc.", "summary_of_the_review": "Due to the missing (1) experimental details, (2) some strong baselines, and (3) limited novelty, currently, the reviewer cannot suggest the acceptance of the paper.", "strength_and_weaknesses": "# Strength\n* The proposed idea is simple and easy to understand.\n* Some numerical results on different neural architectures are provided from various aspects to justify the performance gain of TWA.\n\n# Weaknesses\nThe main weak points of the draft are the missing baselines and experimental details, making the current numerical results hard to convince readers about the advance of the TWA.\n\n1. The main difference between TWA and SWA lies in the aggregation coefficients on a fixed set of the sampled weights, where SWA uses static weights while TWA involves additional training epochs to learn coefficients. However, the number of training epochs used for TWA---different from the one used for the pretraining---did not take into account for a fair comparison. As we can witness e.g. from Table 1 and Table 2, the performance gain of TWA over SWA is quite limited, and it is unclear if the performance again is coming from the extra training or from the design of TWA. Some strong baselines needs to be included to here, e.g.,\n    * taking the SWA model as the initial model and fine-tune the model with the same amount of training epochs as TWA.\n    * sampling more models (from these additional training epochs) and performing SWA. \n\nWithout including these baselines, it is hard to justify the benefits of TMA, e.g., in Table 1, Table 2, Figure 5.\n\n2. The setup used in `Wall-clock time comparison` is unclear, e.g. the number of model checkpoints. Besides, there is no related numerical discussion for sec 4.\n3. A practical usage guideline of TWA currently is missing, and it is unclear to the reviewer why this strategy should be used, rather than other strategies, e.g., SAM, lookahead.\n4. The ablation study on the design choices of TWA should be included, e.g., the impact of regularization term, the importance of optimizing coefficients in a subspace. Otherwise, the novelty of TWA cannot be identified.\n5. The statements in Section 3 are not well-supported, neither theoretically nor empirically.\n\n### Post-rebuttal\nThe reviewer thanks the authors for providing the feedback. Most concerns were addressed, and the reviewer would raise the score to 6.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "qvppaQZNZks", "reviewer": "Reviewer_zFxA", "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": "This submission provides a trainable weight averaging strategy to allow an average of the historical solutions/weights in a trainable manner instead of a pre-defined way, such as SWA and EMA. \nThe experiments show the proposed TWA can improve the generalization better than the previous EMA and SWA.", "summary_of_the_review": "Overall, I hold a positive attitude towards TWA that can rise in practical performance without too much additional effort. However, there are some confusing points in this paper. If the author can handle them well, especially the multi-node one, I would be happy to raise my score.", "strength_and_weaknesses": "Strength\n1. The idea is clear and easy to follow. And the proposed TWA proved to be helpful in practice.\n2. TWA helps to reduce the training cost while keeping the comparable performance.\n\nWeaknesses\n1. The gap between Eq.(1) and Eq.(2) is significant. In most cases, the solution for the problem in Eq.(2) does not satisfy the L1 ball constrain in Eq.(1). The author may refine the claim ``the sum-one constraint in (1) can be naturally satisfied, and we do not require explicitly imposing it in optimizing (2).''\n2. Eq.(A.4) is not quite right. From Eq.(A.1), the gradient w.r.t. $w_{twa}$ is not that in Eq.(A.4). There may include some heuristic derivation, but this part is the crucial aspect of this submission, and we may prefer a more religious presentation way. Eq.(A.5) is one step of the projection gradient descent methods with L2 square regularization (there still exists some gap).\n3. The way to perform TWA in the parallel distributed training manner is not equivalent to the single-node TWA. For single-node, the updated rule is $\\frac{1}{n}PP^\\top (\\sum_i^n g_i)$ while for multi-nodes, the the updated rule is $\\frac{1}{n} (\\sum_i^n P_i P_i^\\top g_i)$, where $P = [P_1,P_2,\\cdots,P_n]$. The two rules are not the same. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "X8zEFaT7eY", "reviewer": "Reviewer_NH8Y", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a trainable way to perform weight averaging and explores it on several image classification datasets.", "summary_of_the_review": "This work generalizes SWA by training the weights explicitly. The ideas are novel and experiments are thorough.", "strength_and_weaknesses": "Strengths:\n- I like the generalization of traditional weight averaging by explicitly training for it, especially since it empirically helps.\n- Authors also consider the computational cost of TWA and how to apply it in a distributed setting.\n- The paper does a great job of thoroughly experimenting on large empirical problems.\n\nWeaknesses / Areas of improvement:\n- One straightforward way to make this paper stronger is to try other architectures and NLP datasets. This way these gains can be illustrated to be more general than currently explored.\n- Section 5.3: the authors mention that one reason SWA is inferior is \"In the head stage of training, SWA usually fails due to the large estimation variance from fast evolving solutions and large learning rate\". I wonder if this hypothesis can be tested where SWA is evaluated by varying the number of iterates to average from the end.\n- Right now the paper tackles weighted averaging of iterates across time. I wonder if this strategy can help for averaging multiple models at the end of the training. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0F1WlzDc_Y", "reviewer": "Reviewer_nodd", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Suppose a set of historical solutions of a DNN model are obtained by applying SGD over a sequence of epochs, the paper considers finding the optimal combination of the historical solutions by minimizing the training loss subject to an equality constraint. The proposed optimisation process is referred to as trainable weight averaging (TWA) in comparison to the stochastic weight averaging (SWA) technique. Design of TWA involves applying the standard Schmidt orthononalization to the set of historical solutions, which is used for gradient projection within a subspace spanned by the set historical solutions. Different experimental comparison have been conducted showing that TWA has better generalisation than SWA.       ", "summary_of_the_review": "Overall, it is interesting to make the coefficients for combining the historical solutions trainable.  The presentation needs to be improved, for example, showing the correctness of (4) and justifying that the sum-one constraint is satisfied or not. The method requires multiple GPUs to enable training, which makes it less practical. TWA essentially introduces an addition training procedure after for example SGD. It would be nice to combine SGD and TWA into one training procedure. ", "strength_and_weaknesses": "Strength:  \n1) The paper, for the first time, considers training the coefficients for optimally combining the historical solutions instead of  simple averaging of the solutions.   \n2)  Different experimental comparison have been conducted showing that TWA has better generalisation than SWA for either head-stage or tail-stage of training.       \n\nWeaknesses: \n(1) I don't think the sum-one constraint in (1) can be naturally satisfied in Algorithm 1. Consider a simple example in 3d space where w1= [1,0,0] and w2=[0,1,0]. In this case e1=w1 and e2=w2. The projected gradient PP^Tg, in principle, could be any point in the x-y domain. As a result, w_{twa} may not always be on the line between w1 and w2.    \n(2)  The derivation of (4) is questionable. If the expectation is with respect to the Guassian distribution N(mu, Sigma). Then I think W_swa has lowest variance. \n(3) I think one inconvenience for applying TWA in the head-stage of the training process is that one needs to manually select the historical solutions. Solutions from early epochs may not work well.   \n(3)  Multiple GPUs are needed to allow for parallel computation for gradient projection in TWA due to high dimensionality of the DNN model and large number of  involved historical solutions, which makes TWA less practical than SWA. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster