Original: T RAINABLE W EIGHT A VERAGING : E FFICIENT T RAINING BY O PTIMIZING H ISTORICAL S OLUTIONS
Revised: T RAINABLE W EIGHT A VERAGING : E FFICIENT T RAINING BY O PTIMIZING H ISTORICAL S OLUTIONS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Stochastic gradient descent (SGD) and its variants are considered as the de-facto methods to train deep neural networks (DNNs).
Revised: Stochastic gradient descent (SGD) and its variants are considered as the de-facto methods to train deep neural networks (DNNs).

Original: While recent improvements to
Revised: While recent improvements to

Original: 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.
Revised: 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.

Original: Recently, an interesting attempt is stochastic weight averaging (SWA), which signiﬁcantly improves the generalization by simply averaging the solutions at the tail stage of training.
Revised: Recently, an interesting attempt is stochastic weight averaging (SWA), which signiﬁcantly improves the generalization by simply averaging the solutions at the tail stage of training.

Original: In this paper, we realize that the averaging coefﬁcients 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.
Revised: In this paper, we realize that the averaging coefﬁcients 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.

Original: TWA has much greater ﬂexibility and can be applied to the head stage of training to achieve training efﬁciency while preserving good generalization capability.
Revised: TWA has much greater ﬂexibility and can be applied to the head stage of training to achieve training efﬁciency while preserving good generalization capability.

Original: Further, we propose a distributed training scheme to resolve the memory burden of large-scale training with efﬁcient parallel computation.
Revised: Further, we propose a distributed training scheme to resolve the memory burden of large-scale training with efﬁcient parallel computation.

Original: In the extensive numerical experiments, (i)
Revised: In the extensive numerical experiments, (i)

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Training deep neural networks (DNNs) usually requires a large amount of time.
Revised: Training deep neural networks (DNNs) usually requires a large amount of time.

Original: As the sizes of models and datasets grow larger, more efﬁcient optimization methods together with better performance are increasingly demanded.
Revised: As the sizes of models and datasets grow larger, more efﬁcient optimization methods together with better performance are increasingly demanded.

Original: In the existing works, great efforts have been made to improve the efﬁciency of stochastic gradient descent (SGD) and its variants, which mainly focus on adaptive learning rates (Duchi et al., 2011; Zeiler, 2012; Kingma & Ba, 2015; Loshchilov & Hutter, 2019; Yao et al., 2021; Heo et al., 2021) or accelerated schemes (Polyak, 1964; Nesterov, 1983; 1988; 2003).
Revised: In the existing works, great efforts have been made to improve the efﬁciency of stochastic gradient descent (SGD) and its variants, which mainly focus on adaptive learning rates (Duchi et al., 2011; Zeiler, 2012; Kingma & Ba, 2015; Loshchilov & Hutter, 2019; Yao et al., 2021; Heo et al., 2021) or accelerated schemes (Polyak, 1964; Nesterov, 1983; 1988; 2003).

Original: As an iterative descent method, SGD generates a series of solutions during optimization.
Revised: As an iterative descent method, SGD generates a series of solutions during optimization.

Original: These historical solutions provide dynamic information about the training and have brought many interesting perspectives, e.g., trajectory (Li et al., 2022), landscape (Garipov et al., 2018), to name a few.
Revised: These historical solutions provide dynamic information about the training and have brought many interesting perspectives, e.g., trajectory (Li et al., 2022), landscape (Garipov et al., 2018), to name a few.

Original: In fact, they can also be utilized to improve the training performance, resulting in the so-called stochastic weight averaging (SWA) (Izmailov et al., 2018), which shows signiﬁcantly better generalization by simply averaging the tail stage explorations of SGD.
Revised: In fact, they can also be utilized to improve the training performance, resulting in the so-called stochastic weight averaging (SWA) (Izmailov et al., 2018), which shows signiﬁcantly better generalization by simply averaging the tail stage explorations of SGD.

Original: A similar idea could be found in Szegedy et al.
Revised: A similar idea could be found in Szegedy et al.

Original: (2016), which designs an exponential moving average (EMA) and considers that a heuristic strategy could be better than equivalently averaging.
Revised: (2016), which designs an exponential moving average (EMA) and considers that a heuristic strategy could be better than equivalently averaging.

Original: The success of SWA and EMA encourages more indepth investigations on the roles of historical solutions obtained during the training (Athiwaratkun et al., 2018; Nikishin et al., 2018; Yang et al., 2019).
Revised: The success of SWA and EMA encourages more indepth investigations on the roles of historical solutions obtained during the training (Athiwaratkun et al., 2018; Nikishin et al., 2018; Yang et al., 2019).

Original: In this paper, our main purpose is to utilize historical solutions by optimizing them, rather than using ﬁxed averaging (e.g., SWA) or a heuristic combination (e.g., EMA).
Revised: In this paper, our main purpose is to utilize historical solutions by optimizing them, rather than using ﬁxed averaging (e.g., SWA) or a heuristic combination (e.g., EMA).

Original: With such an optimized averaging scheme, we can achieve higher accuracy using only the solutions in a relatively early state (i.e. the head stage).
Revised: With such an optimized averaging scheme, we can achieve higher accuracy using only the solutions in a relatively early state (i.e. the head stage).

Original: In other words, we speed up the training and meanwhile improve the performance.
Revised: In other words, we speed up the training and meanwhile improve the performance.

Original: The idea of utilizing these early solutions in DNNs’ training is mainly inspired by two facts.
Revised: The idea of utilizing these early solutions in DNNs’ training is mainly inspired by two facts.

Original: On the one hand, high test accuracy commonly starts appearing at an early stage.
Revised: On the one hand, high test accuracy commonly starts appearing at an early stage.

Original: For example, a PreAct
Revised: For example, a PreAct

Original: ResNet-164 model (He et al., 2016) achieves over 80% test accuracy within 10 training epochs on
Revised: ResNet-164 model (He et al., 2016) achieves over 80% test accuracy within 10 training epochs on

Original: CIFAR-10 (Krizhevsky & Hinton, 2009), while requiring to complete the whole 200 epochs to reach its ﬁnal 95% accuracy.
Revised: CIFAR-10 (Krizhevsky & Hinton, 2009), while requiring to complete the whole 200 epochs to reach its ﬁnal 95% accuracy.

Original: This observation also coincides with the recent ﬁndings that the key connec- 𝐰𝐰𝐰𝐰 swa 𝐰 twa −𝜂𝑷
Revised: This observation also coincides with the recent ﬁndings that the key connec- 𝐰𝐰𝐰𝐰 swa 𝐰 twa −𝜂𝑷

Original: ( 𝑷 ⊤ ∇𝐿 ( 𝐰 ))
Revised: ( 𝑷 ⊤ ∇𝐿 ( 𝐰 ))

Original: Input: Sampled weights f w i g ni =1 , Batch size b , Loss function L : W (cid:2) X (cid:2) Y ! R + , Learning rate η .
Revised: Input: Sampled weights { w i } ni =1 , Batch size b , Loss function L : W × X × Y → R + , Learning rate η .

Original: Output: Model trained with TWA
Revised: Output: Model trained with TWA

Original: Orthogonalize f w i g ni =1 to f e i g ni =1 ;
Revised: Orthogonalize { w i } ni =1 to { e i } ni =1 ;

Original: Initialize w (0)twa , t = 0 , P = [ e 1 , e 2 , (cid:1) (cid:1) (cid:1) , e n ] ; while not converged do
Revised: Initialize w (0)twa , t = 0 , P = [ e 1 , e 2 , · · · , e n ] ; while not converged do

Original: Sample batch data: B = f ( x k , y k ) g bk =1 ;
Revised: Sample batch data: B = { ( x k , y k ) } bk =1 ;

Original: Compute gradient: g = r w L B ( w ( t ) twa ) ;
Revised: Compute gradient: g = ∇ w L B ( w ( t ) twa ) ;

Original: Update weights: w ( t +1)
Revised: Update weights: w ( t +1)

Original: twa = w ( t ) twa (cid:0) η P
Revised: twa = w ( t ) twa − η P

Original: ( P ⊤ g ) ;
Revised: ( P ⊤ g ) ;

Original: t = t + 1 ; end while
Revised: t = t + 1 ; end while

Original: Return w ( t ) twa tivity patterns of DNNs emerge early in training (You et al., 2020; Frankle et al., 2020), indicating a well-explored solution space formed.
Revised: Return w ( t ) twa tivity patterns of DNNs emerge early in training (You et al., 2020; Frankle et al., 2020), indicating a well-explored solution space formed.

Original: On the other hand, simply averaging the solutions collected at the SWA stage immediately provides a huge accuracy improvement, e.g. over 16% on CIFARwith Wide ResNet-28-10 (Zagoruyko & Komodakis, 2016) than before averaging (Izmailov et al., 2018).
Revised: On the other hand, simply averaging the solutions collected at the SWA stage immediately provides a huge accuracy improvement, e.g. over 16% on CIFARwith Wide ResNet-28-10 (Zagoruyko & Komodakis, 2016) than before averaging (Izmailov et al., 2018).

Original: These facts point out a promising direction that sufﬁciently utilizing these early explorations may be capable of quickly composing the ﬁnal solution while obtaining good accuracy.
Revised: These facts point out a promising direction that sufﬁciently utilizing these early explorations may be capable of quickly composing the ﬁnal solution while obtaining good accuracy.

Original: As the model parameters go through a rapid evolution at the early stage of training, a simple averaging strategy with ﬁxed weighting coefﬁcients as in SWA and EMA can result in large estimation errors.
Revised: As the model parameters go through a rapid evolution at the early stage of training, a simple averaging strategy with ﬁxed weighting coefﬁcients as in SWA and EMA can result in large estimation errors.

Original: We introduce a Trainable Weight Averaging (TWA), which allows explicit adjustments for the averaging coefﬁcients in a trainable manner.
Revised: We introduce a Trainable Weight Averaging (TWA), which allows explicit adjustments for the averaging coefﬁcients in a trainable manner.

Original: Speciﬁcally, we construct a subspace that contains all sampled solutions during the training and then conduct efﬁcient optimization therein.
Revised: Speciﬁcally, we construct a subspace that contains all sampled solutions during the training and then conduct efﬁcient optimization therein.

Original: 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 coefﬁcients regardless of the quality of sampled solutions and largely reduce the estimation errors.
Revised: 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 coefﬁcients regardless of the quality of sampled solutions and largely reduce the estimation errors.

Original: The proposed optimization scheme is essentially the gradient projection onto a tiny subspace.
Revised: The proposed optimization scheme is essentially the gradient projection onto a tiny subspace.

Original: 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 to overﬁtting, the latter explains that better training accuracy of TWA over SWA or EMA can lead to better test accuracy.
Revised: 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 to overﬁtting, the latter explains that better training accuracy of TWA over SWA or EMA can lead to better test accuracy.

Original: In extensive experiments with various network architectures on different tasks, we reach superior performance with TWA applied to the head stage of training.
Revised: In extensive experiments with various network architectures on different tasks, we reach superior performance with TWA applied to the head stage of training.

Original: For instance, we attain 1 . 5 (cid:24) 2 .
Revised: For instance, we attain 1 . 5 ∼ 2 .

Original: 2% accuracy improvement on CIFAR-100 and 0 .
Revised: 2% accuracy improvement on CIFAR-100 and 0 .

Original: 1% on ImageNet with over 40% and 30% training epochs reduced, respectively, compared with the regular training.
Revised: 1% on ImageNet with over 40% and 30% training epochs reduced, respectively, compared with the regular training.

Original: In summary, we make the following contributions:
Revised: In summary, we make the following contributions:

Original: • We propose Trainable Weight Averaging (TWA) that allows the averaging coefﬁcients determined in a trainable manner instead of a pre-deﬁned strategy.
Revised: • We propose Trainable Weight Averaging (TWA) that allows the averaging coefﬁcients determined in a trainable manner instead of a pre-deﬁned strategy.

Original: It brings consistent improvements over SWA or EMA with reduced estimation error.
Revised: It brings consistent improvements over SWA or EMA with reduced estimation error.

Original: • We successfully apply TWA to the head stage of training, resulting in a 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.
Revised: • We successfully apply TWA to the head stage of training, resulting in a 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.

Original: • Our TWA is easy to implement and can be ﬂexibly plugged into different stages of training to bring consistent improvements.
Revised: • Our TWA is easy to implement and can be ﬂexibly plugged into different stages of training to bring consistent improvements.

Original: It provides a new scheme for achieving efﬁcient DNNs’ training by sufﬁciently utilizing historical explorations.
Revised: It provides a new scheme for achieving efﬁcient DNNs’ training by sufﬁciently utilizing historical explorations.

Original:  M ETHOD
Revised:  M ETHOD

Original: In this section, we ﬁrst formulate the optimization target of TWA.
Revised: In this section, we ﬁrst formulate the optimization target of TWA.

Original: Then a detailed training algorithm is introduced, which consists of two phases: Schmidt orthogonalization and projected optimization.
Revised: Then a detailed training algorithm is introduced, which consists of two phases: Schmidt orthogonalization and projected optimization.

Original: Note in this paper, the model’s weights are aligned as a vector, i.e., w 2 R D , where D denotes the number of parameters.
Revised: Note in this paper, the model’s weights are aligned as a vector, i.e., w ∈ R D , where D denotes the number of parameters.

Original: 2. O PTIMIZATION T ARGET
Revised: 2. O PTIMIZATION T ARGET

Original: In SWA,the weight averaging is simply given by w swa = 1 n ∑ ni =1 w i , where n solutions of the networkare collected at the tail of trainingand are equally weighted.
Revised: In SWA, weight averaging is simply given by w swa = 1 n ∑ ni =1 w i , where n solutions of the network collected at the tail of training are equally weighted.

Original: Such an averaging strategy has been proved quite effective with better generalization capability.
Revised: Such an averaging strategy has been proven quite effective with better generalization ability.

Original: However, it is too arbitrary to claim that averaging is a perfect solution, which motivates some heuristic modiﬁcations on the weights, e.g., EMA.
Revised: However, equally averaging could not always be a perfect solution, which motivates some heuristic modiﬁcations on weighting strategy, e.g., EMA.

Original: However, both SWA and EMA are inadequate at the head stage, due to the fact thatthe early historical solutions have not stepped into a stationary distribution, and ﬁxed averaging strategies will result in a large estimation error.
Revised: Both SWA and EMA are ﬁxed 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.

Original: Therefore, we propose to train the averaging coefﬁcients with the hope of reducing the estimation error.
Revised: In this paper, we propose to optimize the averaging coefﬁcients with the hope of reducing the estimation error.

Original: Speciﬁcally, the set of possible TWA solutions, i.e., w twa , can be represented as follows:
Revised: Speciﬁcally, the set of possible TWA solutions considered, i.e., w twa , can be represented as follows:

Original: [Equation]
Revised: [Equation]

Revised: The weight vectors between consecutive solutions could have a high cosine similarity.

Revised: To decouple them and for better optimization, we will further orthogonalize { w i } ni =1 and ﬁnd a set of orthogonal bases { e i } ni =1 to support the solution space,

Revised: i

Original:  Then we search for a good solution  in this set by optimizing the following problem,
Revised: [Equation] e [Equation] Then we search for a good solution w twa in A by optimizing the following problem,

Original: [Equation]
Revised: [Equation]

Original: where L ( (cid:1) , (cid:1) ) is the loss function as in regular training and the second term is a regularizationwith a pre-given coefﬁcient λ > 0 .
Revised: where L ( · , · ) is the loss function as in regular training and the second term is a regularization coefﬁcient λ > 0 .

Original: Obviously, SWA ( w swa ) is actually a special solution of (2) with a ﬁxed weight α i = 1 /n , which is not necessarily always a good solution. In the same point of view, EMA which tries to heuristically tune the weights is also a special solution of (2) without optimization.
Revised: Note that both SWA and EMA are special solutions of (2) without optimization.

Original: Optimizing over α i brings beneﬁts in the view of training loss, and a good generalization capability can also be expected: in regular training, the number of optimization variables is D , which is very large, but in (2), there are only n averaging coefﬁcients f α i g ni =1 to be optimized.
Revised: Optimizing over β i brings beneﬁts in the view of training loss, and a good generalization ability could also be expected: in regular training, the number of optimization variables is D , which is very large, but in (2), there are only n averaging coefﬁcients { β i } ni =1 to be optimized.

Original: The signiﬁcantreduction in dimensionality leads to better generalization.
Revised: The signiﬁcant dimensionality reduction could beneﬁt better generalization.

Original: Further, we control the model complexity by introducing the ℓ 2 -norm regularization term.

Original: Notice that weights with good performance will have a quite similar norm ( k w k 2 ).

Original: Thus, the ﬁrst term in (2) that minimizes the training loss has already implicitly required the norm of w twa to be close as w i , i.e., the sum-one constraint in (1) can be naturally satisﬁed and we do not require explicitly imposing it in optimizing (2).

Original: 2. T RAINING A LGORITHM
Revised: 2. T RAINING A LGORITHM

Original: Instead of directly optimizing the averaging coefﬁcients, we notice that there exists a bijection between the coefﬁcient space f α
Revised: Instead of directly optimizing β

Original: i g ni =1 2 R n and the parameter space R D , i.e., each set of theaveraging coefﬁcients is uniquely mapped to one point of the parameter space, which forms a complete subspace (with dimensionality n ).
Revised: i , we note that there exists a bijection between the coefﬁcient space { β i } ni =1 ∈ R n and the parameter space R D , i.e., each set of the coefﬁcients is uniquely mapped to one point in the parameter space, which forms a complete subspace (with dimensionality n ).

Original: We could instead optimize these coefﬁcients in such a subspace.
Revised: We could alternatively optimize these coefﬁcients in such a subspace.

Original: We ﬁrst focus on constructing such a subspace in which the optimization proceeds.

Original: Speciﬁcally, we need to ﬁnd a set of orthogonal bases f e i g ni =1 to span the subspace that covers f w i g ni 
Revised: We ﬁrst focus on ﬁnding a set of orthogonal bases { e i } ni =1 to span the subspace that covers { w i } ni =1 .

Original: =1 .

Original: This is a standard Schmidt orthogonalization, which sequentially takes the following steps: {
Revised: This is a standard Schmidt orthogonalization, which sequentially takes the following steps: {

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: We then initialize the w (0)twa as one point in the subspace (e.g. 1 n ∑ ni =1 w i ), and optimize the network’s parameters therein.
Revised: We then initialize the w (0)twa as one point in the subspace (e.g. 1 n ∑ ni =1 w i ), and optimize the network’s parameters therein.

Original: Let P =
Revised: Let P =

Original: [ e 1 , e 2 , (cid:1) (cid:1) (cid:1) , e n ] , such optimization can be easily achieved by projecting the gradient onto the subspace via projection matrix P P ⊤ .
Revised: [ e 1 , e 2 , · · · , e n ] , such optimization can be easily achieved by projecting the gradient onto the subspace via projection matrix P P ⊤ .

Original: We summarize the detailed training procedures of the proposed TWA in Algorithm 1 with an intuitive illustration in Figure 1.
Revised: We summarize the detailed training procedures of the proposed TWA in Algorithm 1 with an intuitive illustration in Figure 1.

Original: Regularization implementation The regularization in (2) can be integrated into subspace training as a weight decay in a very easy implementation.

Original: Details are elaborated in Appendix B.
Revised: The detailed implementation is described in Appendix B.

Original:  O PTIMIZATION P ROVIDES B ETTER F LEXIBILITY
Revised:  O PTIMIZATION P ROVIDES B ETTER F LEXIBILITY

Original: The key difference between SWA and TWA is that the averaging coefﬁcients in TWA are determined in a trainable manner, or more precisely, are data-dependent.
Revised: The key difference between SWA and TWA is that the averaging coefﬁcients in TWA are determined in a trainable manner, or more precisely, are data-dependent.

Original: This potentially enables more precise estimation for the center minima and better tolerance for the outliers that are not aware by SWA.
Revised: This potentially enables more precise estimation for the center minima and better tolerance for the outliers that are not aware by SWA.

Original: Notice that in the view of coefﬁcient optimization, there is no essential difference between SWA and
Revised: Notice that in the view of coefﬁcient optimization, there is no essential difference between SWA and

Original: EMA, which both provide speciﬁc and data-independent solutions.
Revised: EMA, which both provide speciﬁc and data-independent solutions.

Original: Thus, we in the following only compare TWA with SWA.
Revised: Thus, we in the following only compare TWA with SWA.

Original: Mandt et al.
Revised: Mandt et al.

Original: (2017) demonstrated that under appropriate assumptions, running SGD with a constant learning rate is equivalent to sampling from a stationary Gaussian distribution, and the variance of the distribution is controlled by the learning rate.
Revised: (2017) demonstrated that under appropriate assumptions, running SGD with a constant learning rate is equivalent to sampling from a stationary Gaussian distribution, and the variance of the distribution is controlled by the learning rate.

Original: Accordingly, we assume the solutions at the tail stage of SGD training are sampled from a Gaussian distribution N ( µ , Σ ) centered at the minimum µ with covariance Σ .
Revised: Accordingly, we assume the solutions at the tail stage of SGD training are sampled from a Gaussian distribution N ( µ , Σ ) centered at the minimum µ with covariance Σ .

Original: Approximately, the sampled solutions f w i g ni =1 are independent random variables from N ( µ , Σ ) , as long as there are sufﬁcient iterations between adjacent samplings.
Revised: Approximately, the sampled solutions { w i } ni =1 are independent random variables from N ( µ , Σ ) , as long as there are sufﬁcient iterations between adjacent samplings.

Original: SWA and TWA provide two estimators for the minimum µ , i.e., w swa and w twa .
Revised: SWA and TWA provide two estimators for the minimum µ , i.e., w swa and w twa .

Original: As an averaged solution, w swa has statistically better estimation than any single solution due to the effect of variance reduction, while w twa is approaching the center by minimizing the training loss.
Revised: As an averaged solution, w swa has statistically better estimation than any single solution due to the effect of variance reduction, while w twa is approaching the center by minimizing the training loss.

Original: As long as the training loss serves asa meaningful supervision, w twa could approach µ better than w swa .
Revised: As long as the training loss serves as meaningful supervision (which holds under the typical assumption that µ is the center minimum with lowest training loss Mandt et al.

Revised: (2017); Izmailov et al.

Original: Both w swa and w twa are unbiased estimators for µ , and w twa has a lower expected variance:
Revised: , w twa could approach µ better than w swa with the posterior optimization for averaging coefﬁcients. In this regard, w twa could have a lower expected variance:

Original: [Equation]
Revised: [Equation]

Original: The advantages of optimizing averaging coefﬁcients could be more prominent in the head stage of training, where the weights are going through a rapid evolution.
Revised: The advantages of optimizing averaging coefﬁcients could be more prominent in the head stage of training, where the weights are going through a rapid evolution.

Original: A simple averaging strategy as
Revised: A simple averaging strategy as

Original: SWA could introduce a large estimation error (as illustrated in Table 3 and 4), while TWA enables correcting it to a smaller estimation error via optimization.
Revised: SWA could introduce a large estimation error (as illustrated in Table 3 and 4), while TWA enables correcting it to a smaller estimation error via optimization.

Original: In fact, TWA provides much more ﬂexibility to sufﬁciently utilize historical solutions and produces an optimized solution adaptively.
Revised: In fact, TWA provides much more ﬂexibility to sufﬁciently utilize historical solutions and produces an optimized solution adaptively.

Original:  A N E FFICIENT I MPLEMENTATION FOR D ISTRIBUTED T RAINING
Revised:  A N E FFICIENT I MPLEMENTATION FOR D ISTRIBUTED T RAINING

Original: The above discussion shows promising improvement by optimizing the historical solutions.
Revised: The above discussion shows promising improvement by optimizing the historical solutions.

Original: The only issue one may worry about is the burden in storage (the additional time complexity is small as shown in Table 6).
Revised: The only issue one may worry about is the burden in storage (the additional time complexity is small as shown in Table 6).

Original: During optimization, TWA requires the projection matrix P involving dozens or hundreds of historical weights, which indeed poses a challenge for large models on storage burden.
Revised: During optimization, TWA requires the projection matrix P involving dozens or hundreds of historical weights, which indeed poses a challenge for large models on storage burden.

Original: It is preferable to locate P in GPUs to enable efﬁcient matrix operations.
Revised: It is preferable to locate P in GPUs to enable efﬁcient matrix operations.

Original: However, the size of
Revised: However, the size of

Original: P increases as the model becomes larger, making it prohibitive to store in a single GPU.
Revised: P increases as the model becomes larger, making it prohibitive to store in a single GPU.

Original: To cope with this, we design an efﬁcient scheme with parallel distributed training to enable a) partition of the memory burden of P into multiple GPUs and b) efﬁcient parallel computation of gradient projection.
Revised: To cope with this, we design an efﬁcient scheme with parallel distributed training to enable a) partition of the memory burden of P into multiple GPUs and b) efﬁcient parallel computation of gradient projection.

Original: As a result, we successfully optimize more than 900 historical solution coefﬁcients for ResNeton ImageNet task by 4 v100 GPUs.
Revised: As a result, we successfully optimize more than 900 historical solution coefﬁcients for ResNeton ImageNet task by 4 v100 GPUs.

Original: In our experiments, we use at most 300 historical solutions and there is still available space for larger tasks.
Revised: In our experiments, we use at most 300 historical solutions and there is still available space for larger tasks.

Original: Suppose that there are k GPUs for parallel training.
Revised: Suppose that there are k GPUs for parallel training.

Original: We ﬁrst uniformly divide P into k sub-matrices as P =
Revised: We ﬁrst uniformly divide P into k sub-matrices as P =

Original: [ P 1 , P 2 , (cid:1) (cid:1) (cid:1) , P k ] , where each GPU stores a local sub-matrix P i , i = 1 , (cid:1) (cid:1) (cid:1) , k .
Revised: [ P 1 , P 2 , · · · , P k ] , where each GPU stores a local sub-matrix P i , i = 1 , · · · , k .

Original: Recall that for an iteration in distributed training, each GPU computes a local gradient g i and synchronizes it with other GPUs to obtain the global gradient through an efﬁcient all-reduce operation (Rabenseifner, 2004).
Revised: Recall that for an iteration in distributed training, each GPU computes a local gradient g i and synchronizes it with other GPUs to obtain the global gradient through an efﬁcient all-reduce operation (Rabenseifner, 2004).

Original: We mimic such a process for gradient projection: the local projected gradient P i P ⊤ i g is ﬁrstly computed in each card and then synchronized with others to obtain the global projected gradient with another all-reduce operation.
Revised: We mimic such a process for gradient projection: the local projected gradient P i P ⊤ i g is ﬁrstly computed in each card and then synchronized with others to obtain the global projected gradient with another all-reduce operation.

Original: We illustrate such a process in Figure 2.
Revised: We illustrate such a process in Figure 2.

Original: For averaging n historical solutions with per size B , the memory burden for each GPU card is reduced to d n/k e B , while the computation of gradient projection is also reduced to O ( d n/k e D )
Revised: For averaging n historical solutions with per size B , the memory burden for each GPU card is reduced to ⌈ n/k ⌉ B , while the computation of gradient projection is also reduced to O ( ⌈ n/k ⌉ D )

Original: ( D is the number of the parameters).
Revised: ( D is the number of the parameters).

Original: Hence, we can achieve efﬁcient TWA training by making full use of the remaining memory of each GPU aside the forward/backward training. 
Revised: Hence, we can achieve efﬁcient TWA training by making full use of the remaining memory of each GPU aside the forward/backward training. [Figure]

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, a series of numerical experiments are conducted, demonstrating the effectiveness of our proposed TWA for fast convergence and better performance.
Revised: In this section, a series of numerical experiments are conducted, demonstrating the effectiveness of our proposed TWA for fast convergence and better performance.

Original: First, we show that TWA improves
Revised: First, we show that TWA improves

Original: SWA in the existing SWA settings, i.e., used in the tail stage of training.
Revised: SWA in the existing SWA settings, i.e., used in the tail stage of training.

Original: Second, we apply TWA to the head stage of training, which brings signiﬁcant efﬁciency improvements together with better performance.
Revised: Second, we apply TWA to the head stage of training, which brings signiﬁcant efﬁciency improvements together with better performance.

Original: Then, we visualize loss / accuracy surfaces to demonstrate the improvements of TWA.
Revised: Then, we visualize loss / accuracy surfaces to demonstrate the improvements of TWA.

Original: 5. E XPERIMENTAL SETTINGS
Revised: 5. E XPERIMENTAL SETTINGS

Original: Datasets.
Revised: Datasets.

Original: We experiment over three benchmark image datasets, i.e., CIFAR-10, CIFAR-
Revised: We experiment over three benchmark image datasets, i.e., CIFAR-10, CIFAR-

Original: (Krizhevsky & Hinton, 2009), and ImageNet (Deng et al., 2009).
Revised: (Krizhevsky & Hinton, 2009), and ImageNet (Deng et al., 2009).

Original: Following prior works (Izmailov et al., 2018; Yang et al., 2019), we apply standard preprocessing for experiments on CIFAR datasets, and adopt the preprocessing and data augmentation procedures in the public Pytorch example on
Revised: Following prior works (Izmailov et al., 2018; Yang et al., 2019), we apply standard preprocessing for experiments on CIFAR datasets, and adopt the preprocessing and data augmentation procedures in the public Pytorch example on

Original: ImageNet (Paszke et al., 2017).
Revised: ImageNet (Paszke et al., 2017).

Original: Architectures.
Revised: Architectures.

Original: We use two representative architectures, VGG-16 (Simonyan & Zisserman, 2014) and PreAct ResNet-164 (He et al., 2016) on CIFAR experiments.
Revised: We use two representative architectures, VGG-16 (Simonyan & Zisserman, 2014) and PreAct ResNet-164 (He et al., 2016) on CIFAR experiments.

Original: For ImageNet, we use ResNetand ResNet-50 (He et al., 2016).
Revised: For ImageNet, we use ResNetand ResNet-50 (He et al., 2016).

Original: Training.
Revised: Training.

Original: The main body of experiments contains two parts: (1) for the tail stage of training, we use the same hyper-parameters as in SWA (Izmailov et al., 2018) and then a larger tail learning rate is also tried.
Revised: The main body of experiments contains two parts: (1) for the tail stage of training, we use the same hyper-parameters as in SWA (Izmailov et al., 2018) and then a larger tail learning rate is also tried.

Original: (2) for the head stage of training, we adopt the standard training protocol with a step-wise learning rate.
Revised: (2) for the head stage of training, we adopt the standard training protocol with a step-wise learning rate.

Original: For CIFAR, we run all experiments with 3 seeds and report the mean test accuracy.
Revised: For CIFAR, we run all experiments with 3 seeds and report the mean test accuracy.

Original: We use SGD optimizer with momentum 0 .
Revised: We use SGD optimizer with momentum 0 .

Original: 9 , weight decay 10 − 4 , and batch size 128 .
Revised: 9 , weight decay 10 − 4 , and batch size 128 .

Original: We train the models for 200 epochs with an initial learning rate 0 .
Revised: We train the models for 200 epochs with an initial learning rate 0 .

Original: 1 and decay it by 10 at the 100th and the 150th epochs.
Revised: 1 and decay it by 10 at the 100th and the 150th epochs.

Original: For ImageNet, we follow ofﬁcial PyTorch implementation 1 .
Revised: For ImageNet, we follow ofﬁcial PyTorch implementation 1 .

Original: For TWA, we sample solutions once after each epoch training for CIFAR and uniformly sample 5 times per epoch for ImageNet.
Revised: For TWA, we sample solutions once after each epoch training for CIFAR and uniformly sample 5 times per epoch for ImageNet.

Original: We use a scaled learning rate (Figure 4), which takes 10 epochs of training for CIFAR and 2 epochs for ImageNet for fast convergence.
Revised: We use a scaled learning rate (Figure 4), which takes 10 epochs of training for CIFAR and 2 epochs for ImageNet for fast convergence.

Original: The regularization coefﬁcient λ defaults to 10 − 5 .
Revised: The regularization coefﬁcient λ defaults to 10 − 5 .

Original: More details (including the number of the historical solutions used) could be found in Appendix A.
Revised: More details (including the number of the historical solutions used) could be found in Appendix A.

Original: 5. I MPROVING SWA SOLUTIONS
Revised: 5. I MPROVING SWA SOLUTIONS

Original: In this part, we compare SWA and TWA in the original SWA settings.
Revised: In this part, we compare SWA and TWA in the original SWA settings.

Original: Speciﬁcally, TWA and SWA use the same weights sampled from the tail stage of training.
Revised: Speciﬁcally, TWA and SWA use the same weights sampled from the tail stage of training.

Original: For CIFAR, we try two different tail learning rates: 0.05, the recommended one in (Izmailov et al., 2018), and 0.10, a larger one for the case with greater variance.
Revised: For CIFAR, we try two different tail learning rates: 0.05, the recommended one in (Izmailov et al., 2018), and 0.10, a larger one for the case with greater variance.

Original: The results in Table 1 show that TWA brings consistent improvements over SWA.
Revised: The results in Table 1 show that TWA brings consistent improvements over SWA.

Original: Especially when the learning rate is not well-tuned, SWA’s performance suffers a distinctive drop, but TWA is less sensitive since the estimation error could be well controlled through training.
Revised: Especially when the learning rate is not well-tuned, SWA’s performance suffers a distinctive drop, but TWA is less sensitive since the estimation error could be well controlled through training.

Original: For instance, in the case of CIFAR-100 with VGG-16 and a tail learning rate of 0.10, the estimation error of SWA substantially increases while TWA achieves a signiﬁcant accuracy 
Revised: For instance, in the case of CIFAR-100 with VGG-16 and a tail learning rate of 0.10, the estimation error of SWA substantially increases while TWA achieves a signiﬁcant accuracy im- [Figure] [Table]

Original: - provement, i.e., 4 .
Revised: provement, i.e., 4 .

Original: 18% , over SWA.
Revised: 18% , over SWA.

Original: Note that for a fair comparison, TWA starts from the last sampled weights, not the averaged solution of SWA.
Revised: Note that for a fair comparison, TWA starts from the last sampled weights, not the averaged solution of SWA.

Original: [Table]
Revised: [Table]

Original: On ImageNet, we experiment with ResNet-18/50 (He et al., 2016).
Revised: On ImageNet, we experiment with ResNet-18/50 (He et al., 2016).

Original: Following (Izmailov et al., 2018; Yang et al., 2019), we start from pre-trained models in torchvision.models and collect
Revised: Following (Izmailov et al., 2018; Yang et al., 2019)

Original: [Figure]
Revised: [Table]

Original:  model weights by running SGD optimizer up to 10 epochs (with a constant learning rate 0 . 
Revised: we start from pre-trained models in torchvision.models and collect model weights by running SGD optimizer up to 10 epochs (with a constant learning rate 0 . 005 ).

Original: In Table 2, we report the test accuracy and observe that with more sampling epochs, both TWA and SWA achieve better performance.
Revised: In Table [Table] we report the test accuracy and observe that with more sampling epochs, both TWA and SWA achieve better performance.

Original: Notably, TWA performs better than SWA by 0 .
Revised: Notably, TWA performs better than SWA by 0 .

Original: 1 (cid:24) 0 .
Revised: 1 ∼ 0 .

Original: 3% , and such improvements are more obvious in the “ 5 EPOCHS ” case.
Revised: 3% , and such improvements are more obvious in the “ 5 EPOCHS ” case.

Original: For example, using 5 epochs of sampled weights, TWA achieves 70 .
Revised: For example, using 5 epochs of sampled weights, TWA achieves 70 .

Original: 23% accuracy with ResNet-18 and 76 .
Revised: 23% accuracy with ResNet-18 and 76 .

Original: 78% accuracy with
Revised: 78% accuracy with

Original: ResNet-50, outperforming the SWA counterparts with 10 epochs.
Revised: ResNet-50, outperforming the SWA counterparts with 10 epochs.

Original: 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 coefﬁcients.[Figure] [Table] [Figure]
Revised: 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 coefﬁcients.

Original: i n

Original: [Figure]
Revised: [Table]

Original: a t

Original: [Figure]

Original: [Figure]

Original: 5. E FFICIENT TRAINING AND BETTER GENERALIZATION
Revised: 5. E FFICIENT TRAINING AND BETTER GENERALIZATION

Original: In the head stage of training, SWA usually fails due to the large estimation variance from fastevolving solutions and large learning rate.
Revised: In the head stage of training [Table] SWA usually fails due to the large estimation variance from fastevolving solutions and large learning rate.

Original: 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.
Revised: 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.

Original: If so, it is promising to simultaneously attain generalization improvements and training efﬁciency.
Revised: If so, it is promising to simultaneously attain generalization improvements and training efﬁciency.

Original: We ﬁrst investigate the experiments on CIFAR datasets.
Revised: We ﬁrst investigate the experiments on CIFAR datasets.

Original: The original training schedule containsepochs
Revised: The original training schedule containsepochs

Original: and we take the ﬁrst 100 epoch explorations for TWA.
Revised: and we take the ﬁrst 100 epoch explorations for TWA.

Original: The results are given in Table 3.
Revised: The results are given in Table 3.

Original: It can be observed that TWA achieves better performance compared to the regular SGD training with a signiﬁcantly reduced generalization gap.
Revised: It can be observed that TWA achieves better performance compared to the regular SGD training with a signiﬁcantly reduced generalization gap.

Original: For instance, we attain 1 .
Revised: For instance, we attain 1 .

Original: 52% accuracy improvement on CIFAR-100 with VGG-16 while the generalization gap is reduced by 9 . 56% .
Revised: 52% accuracy improvement on CIFAR-100 with VGG-16 while the generalization gap is reduced by 9 . 56% .

Original: 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 overﬁtting problems and harm the generalization.
Revised: 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 overﬁtting problems and harm the generalization.

Original: By comparisons, we also apply SWA to average these samples, which shows degraded performance due to the existence of estimation error.
Revised: By comparisons, we also apply SWA to average these samples, which shows degraded performance due to the existence of estimation error.

Original: Apart from the good performance in accuracy, TWA also manifests its great potential in improving the training efﬁciency: we use onlyepochs to complete the convergence, while the regular SGD needs 100 epochs.
Revised: Apart from the good performance in accuracy, TWA also manifests its great potential in improving the training efﬁciency: we use onlyepochs to complete the convergence, while the regular SGD needs 100 epochs.

Original: As TWA and SGD have nearly the same computation overhead per epoch, the time-saving is around 45% in TWA.
Revised: As TWA and SGD have nearly the same computation overhead per epoch, the time-saving is around 45% in TWA.

Original: [Table]
Revised: [Table]

Original: Generally, utilizing more epochs of explorations can provide a better estimation for the center minimum and hence lead to better performance.
Revised: Generally, utilizing more epochs of explorations can provide a better estimation for the center minimum and hence lead to better performance.

Original: Then
Revised: Then

Original: [Table]
Revised: [Table]

Original: we also study the impact of different averaging epochs
Revised: we also study the impact of different averaging epochs

Original: [Table]
Revised: [Table]

Original: and the results are illustrated in Figure
Revised: and the results are illustrated in Figure

Original: [Table]
Revised: [Table]

Original: where the ﬁnal accuracy of SGD and the accuracy reached by SGD before averaging are also given for reference.
Revised: where the ﬁnal accuracy of SGD and the accuracy reached by SGD before averaging are also given for reference.

Original: It could be observed that the model’s performance is consistently improved with more epochs of explorations.
Revised: It could be observed that the model’s performance is consistently improved with more epochs of explorations.

Original: Notably, although each historical solution in a relatively short period of explorations is not good, satisﬁed solutions have already emerged in the subspace spanned by these solutions.
Revised: Notably, although each historical solution in a relatively short period of explorations is not good, satisﬁed solutions have already emerged in the subspace spanned by these solutions.

Original: Then through proper optimization in subspace, TWA could ﬁnd them out, e.g., on CIFAR-100 with PreAct ResNet-164 model, averaging over 50 epochs via TWA has already matched the ﬁnal performance of regular SGD training.
Revised: Then through proper optimization in subspace, TWA could ﬁnd them out, e.g., on CIFAR-100 with PreAct ResNet-164 model, averaging over 50 epochs via TWA has already matched the ﬁnal performance of regular SGD training.

Original: [Table]
Revised: [Table]

Original: For ImageNet, the efforts required for each epoch training are much greater, and hence efﬁcient methods to reduce the training epochs are highly desirable.
Revised: For ImageNet, the efforts required for each epoch training are much greater, and hence efﬁcient methods to reduce the training epochs are highly desirable.

Original: The comparison results of SGD/SWA/TWA are shown in Table 4.
Revised: The comparison results of SGD/SWA/TWA are shown in Table 4.

Original: Besides the reduced generalization gap, TWA takes only 2 epochs to average the historical solutions of the ﬁrst 60 epochs, reaching comparable or even better performance than regular SGD training with 90 epochs.
Revised: Besides the reduced generalization gap, TWA takes only 2 epochs to average the historical solutions of the ﬁrst 60 epochs, reaching comparable or even better performance than regular SGD training with 90 epochs.

Original: For comparison, Lookahead (Zhang et al., 2019) is another advanced optimizer recently proposed for improving convergence and reported 75 .
Revised: For comparison, Lookahead (Zhang et al., 2019) is another advanced optimizer recently proposed for improving convergence and reported 75 .

Original: 49% accuracy at the 60th epoch
Revised: 49% accuracy at the 60th epoch

Original: (Table 2 in Zhang et al.
Revised: (Table 2 in Zhang et al.

Original: with an aggressive learning rate decay (i.e., the learning rate is decayed at the 30th, 48th, and 58th epochs), while our TWA reaches 75 .
Revised: with an aggressive learning rate decay (i.e., the learning rate is decayed at the 30th, 48th, and 58th epochs), while our TWA reaches 75 .

Original: 70% with the same budget but simply using the conventional decay.
Revised: 70% with the same budget but simply using the conventional decay.

Original: [Table]
Revised: [Table]

Original: TWA is very ﬂexible and can be readily applied to different training stages, and we also conduct an experiment by averaging the solutions of the ﬁnal training period (i.e. 61-90 epochs) and simply performing TWA for one epoch training, as presented in Table 5.
Revised: TWA is very ﬂexible and can be readily applied to different training stages, and we also conduct an experiment by averaging the solutions of the ﬁnal training period (i.e. 61-90 epochs) and simply performing TWA for one epoch training, as presented in Table 5.

Original: Such cheap training still shows to bring significant improvements (e.g. +0 .
Revised: Such cheap training still shows to bring significant improvements (e.g. +0 .

Original: 51% on ResNet-50 for
Revised: 51% on ResNet-50 for

Original: ImageNet).
Revised: ImageNet).

Original: Thus, TWA can serve as an effective approach for composing a better ﬁnal solution.
Revised: Thus, TWA can serve as an effective approach for composing a better ﬁnal solution.

Original: Scaled learning rate The optimization of TWA is conducted in a very low-dimensional space, which also suppresses the sensitivity of the learning rate.
Revised: Scaled learning rate The optimization of TWA is conducted in a very low-dimensional space, which also suppresses the sensitivity of the learning rate.

Original: In fact, we can allow a very large learning rate to accelerate the training.
Revised: In fact, we can allow a very large learning rate to accelerate the training.

Original: Thereby, we design a scaled learning rate , which linearly scales up the learning rate and reduces the training epochs accordingly, as shown in Figure 4.
Revised: Thereby, we design a scaled learning rate , which linearly scales up the learning rate and reduces the training epochs accordingly, as shown in Figure 4.

Original: Within an appropriate range, scaling the learning rate largely speeds up the convergence without affecting the ﬁnal performance.
Revised: Within an appropriate range, scaling the learning rate largely speeds up the convergence without affecting the ﬁnal performance.

Original: For example, with the learning rate of 4, TWA approaches the ﬁnal accuracy with only 1 epoch and converges within 5 epochs.
Revised: For example, with the learning rate of 4, TWA approaches the ﬁnal accuracy with only 1 epoch and converges within 5 epochs.

Original: [Table]
Revised: [Table]

Original: 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.
Revised: Wall-clock time comparison In the above experiments, we report the number of training epochs, since the wallclock time per epoch for SGD and TWA is similar.

Original: As an illustration, Table 6 provides the wallclock time of training PreAct ResNet-164 for CIFAR100 on one Nvidia Geforce GTX 2080 TI.
Revised: As an illustration, Table 6 provides the wall-clock time of training PreAct ResNet-164 for CIFAR-100 on one Nvidia Geforce GTX 2080 TI.

Original: [Figure]
Revised: [Figure]

Original: Comparison with EMA Exponential moving averages (EMA) is an alternative to SWA, which averages the model weights along the training trajectory by employing an exponential decay.
Revised: Comparison with EMA EMA serves as an alternative to SWA, which averages the model weights along the training trajectory with exponential decay.

Original: It requires a hyper-parameter γ to control the averaging horizons.
Revised: It requires a hyper-parameter γ to control the averaging horizons.

Original: As a manually pre-deﬁned averaging strategy, EMA could be sensitive to learning rate, datasets or architectures, etc.
Revised: As a manually deﬁned averaging strategy, EMA could be sensitive to learning rates, datasets, architectures, etc.

Original: Here we compare the  EMA with SWA and TWA in the head stage of training, where we try γ = 0 .
Revised: Here we compare the performance of EMA with SWA and TWA in the head stage of training, where we try γ = 0 .

Original: The results are illustrated in Figure 5.
Revised: The results are illustrated in Figure 5.

Original: We observe that the performance of EMA varies signiﬁcantly with different choices of γ .
Revised: We observe that the performance of EMA varies signiﬁcantly with different choices of γ .

Original: It could perform notably better than SWA in the early stage of training as more weights are paid to the latest solutions.
Revised: It could perform notably better than SWA in the early stage of training as more weight is paid to the latest solutions.

Original: We note that both EMA and SWA are heuristic strategies for adapting differentstage of training  (essentially can be viewed asa particular solution of TWA), and TWA could  achievesigniﬁcantly better performance via optimization.
Revised: Note that both EMA and SWA are ﬁxed averaging strategies for adapting different training stages (essentially could be viewed as particular solutions of TWA). By optimizing the averaging coefﬁcients, TWA could consistently achieve better performance.

Original: [Figure]
Revised: [Figure]

Original: Optimized Averaging Coefﬁcients We visualize the averaging coefﬁcients α i learned by TWA in Figure 6.
Revised: Optimized Averaging Coefﬁcients In Figure 6, we visualize the averaging coefﬁcients α i learned by TWA.

Original:  We observe all historical solutions  contribute to the ﬁnal solutions.
Revised: Detailed derivation could be found in Appendix D. We observe all historical solutions could contribute to the ﬁnal solutions.

Original: Solutions from the latter training stage are attached with more importance as expected.
Revised: Solutions from the latter training stage are attached with more importance as expected.

Original: Different from the hand-craft averaging strategies like SWA or EMA, such averaging coefﬁcients enable to take full advantage of the historical solutions through delicate optimization and better adapt to the training dynamics.
Revised: Different from the ﬁx averaging strategies like SWA or EMA, such averaging coefﬁcients enable to take full advantage of the historical solutions through delicate optimization and better adapt to the training dynamics.

Original: 5. L ANDSCAPE VISUALIZATION
Revised: 5. L ANDSCAPE VISUALIZATION

Original: Following (Garipov et al., 2018; Izmailov et al., 2018), we visualize the training loss and test error surfaces of SWA and TWA in Figure 7 on CIFAR-100 with PreAct ResNet-164.
Revised: Following (Garipov et al., 2018; Izmailov et al., 2018), we visualize the training loss and test error surfaces of SWA and TWA in Figure 7 on CIFAR-100 with PreAct ResNet-164.

Original: We set the SGD solution after 125 training epochs as the origin and plot the TWA and SWA solutions on the plane.
Revised: We set the SGD solution after 125 training epochs as the origin and plot the TWA and SWA solutions on the plane.

Original: For the case with a default learning rate of 0 . 05 , TWA achieves slightly better test accuracy with lower training loss.
Revised: For the case with a default learning rate of 0 . 05 , TWA achieves slightly better test accuracy with lower training loss.

Original: This shows that in the subspace, minimizing the training loss is meaningful and results in lower test errors.
Revised: This shows that in the subspace, minimizing the training loss is meaningful and results in lower test errors.

Original: Especially for the case with a larger learning rate of 0 . 10 , the superiority of TWA over SWA is more signiﬁcant (over 0 .
Revised: Especially for the case with a larger learning rate of 0 . 10 , the superiority of TWA over SWA is more signiﬁcant (over 0 .

Original: 7% improvement on test accuracy), since the variance grows larger and the variance reduction effect of TWA becomes more obvious.
Revised: 7% improvement on test accuracy), since the variance grows larger and the variance reduction effect of TWA becomes more obvious.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Improving the model’s generalization capability is of great importance and has received wide attention.
Revised: Improving the model’s generalization capability is of great importance and has received wide attention.

Original: The recent efforts mainly focus on two aspects: (1) proper regularization terms to search for more ﬂat minimum (Keskar et al., 2016;
Revised: The recent efforts mainly focus on two aspects: (1) proper regularization terms to search for more ﬂat minimum (Keskar et al., 2016;

Original: Li et al., 2018), such as weight decay (Krogh & Hertz, 1991), dropout (Srivastava et al., 2014), label smoothing (Szegedy et al., 2016), Shake-Shake (Gastaldi, 2017), MixUp (Zhang et al., 2018), SAM (Foret et al., 2020) and AMP (Zheng et al., 2021); (2) effective data augmentation to diversify the dataset, such as Cutout (DeVries & Taylor, 2017), AutoAug-
Revised: Li et al., 2018), such as weight decay (Krogh & Hertz, 1991), dropout (Srivastava et al., 2014), label smoothing (Szegedy et al., 2016), Shake-Shake (Gastaldi, 2017), MixUp (Zhang et al., 2018), SAM (Foret et al., 2020) and AMP (Zheng et al., 2021); (2) effective data augmentation to diversify the dataset, such as Cutout (DeVries & Taylor, 2017), AutoAugment (Cubuk et al., 2019) and RandAugment (Cubuk et al., 2020). Different from these techniques,

Original: [Figure]
Revised: [Figure]

Original: ment (Cubuk et al., 2019) and RandAugment (Cubuk et al., 2020).

Original: 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.
Revised: we improve the generalization ability by constraining the training in a low-dimensional subspace spanned by historical explorations, which regularizes the model complexity.

Original: We note that TWA is orthogonal to these methods
Revised: We note that TWA is orthogonal to these methods

Original: [Figure]
Revised: [Figure]

Original: and it is promising to combine them for boosted improvements.
Revised: and it is promising to combine them for boosted improvements.

Original: A lot of efforts have been made to speed up the DNNs’ training.
Revised: A lot of efforts have been made to speed up the DNNs’ training.

Original: Apart from the well-known methods on adaptive learning rates, e.g. Adam (Kingma & Ba, 2015) and accelerated schemes, e.g. Nesterov momentum Nesterov (1983), a new method is proposed in Zhang et al.
Revised: Apart from the well-known methods on adaptive learning rates, e.g. Adam (Kingma & Ba, 2015) and accelerated schemes, e.g. Nesterov momentum Nesterov (1983), a new method is proposed in Zhang et al.

Original: (2019) proposed, where a look-ahead search direction generated by another “fast” optimizer is utilized, achieving faster convergence and better learning stability.
Revised: (2019) proposed, where a look-ahead search direction generated by another “fast” optimizer is utilized, achieving faster convergence and better learning stability.

Original: Goyal et al.
Revised: Goyal et al.

Original: (2017) adapted a large mini-batch to speed up the training and introduced a scaling rule for adjusting the learning rates.
Revised: (2017) adapted a large mini-batch to speed up the training and introduced a scaling rule for adjusting the learning rates.

Original: In this paper, we realize training efﬁciency by sufﬁciently utilizing the historical solutions of DNNs’ training and conducting training in a subspace with substantially reduced dimensions.
Revised: In this paper, we realize training efﬁciency by sufﬁciently utilizing the historical solutions of DNNs’ training and conducting training in a subspace with substantially reduced dimensions.

Original: For utilizing historical explorations, SWA (Izmailov et al., 2018) adopts a simple averaging strategy at the tail of training.
Revised: For utilizing historical explorations, SWA (Izmailov et al., 2018) adopts a simple averaging strategy at the tail of training.

Original: Cha et al.
Revised: Cha et al.

Original: (2021) extended it to the domain generalization task with a dense and overﬁt-aware stochastic weight sampling strategy.
Revised: (2021) extended it to the domain generalization task with a dense and overﬁt-aware stochastic weight sampling strategy.

Original: We ﬁrstly propose to utilize the explorations at the head stage of training to achieve training efﬁciency.
Revised: We ﬁrstly propose to utilize the explorations at the head stage of training to achieve training efﬁciency.

Original: Exponentially decaying running average (Hunter, 1986; Szegedy et al., 2016) is a common technique adopted by practitioners.
Revised: Exponentially decaying running average (Hunter, 1986; Szegedy et al., 2016) is a common technique adopted by practitioners.

Original: It requires a manually set averaging horizon and generally performs comparably as SGD (Izmailov et al., 2018).
Revised: It requires a manually set averaging horizon and generally performs comparably as SGD (Izmailov et al., 2018).

Original: Another closely related work is model soups (Wortsman et al., 2022), which improves the model performance by averaging the weights from different ﬁne-tuning conﬁgurations in a greedy order.
Revised: Another closely related work is model soups (Wortsman et al., 2022), which improves the model performance by averaging the weights from different ﬁne-tuning conﬁgurations in a greedy order.

Original: We differ in that the historical solutions are from one single conﬁguration.
Revised: We differ in that the historical solutions are from one single conﬁguration.

Original: We mainly focus on improving training efﬁciency and optimizing the averaging coefﬁcients in a trainable manner.
Revised: We mainly focus on improving training efﬁciency and optimizing the averaging coefﬁcients in a trainable manner.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we propose TWA, a novel training algorithm that optimizes the averaging coefﬁcients of historical solutions in DNNs’ training to achieve efﬁciency and better performance.
Revised: In this work, we propose TWA, a novel training algorithm that optimizes the averaging coefﬁcients of historical solutions in DNNs’ training to achieve efﬁciency and better performance.

Original: It differs from the manually set averaging strategies as SWA or EMA and manifests better adaptation to different stages of training.
Revised: It differs from the manually set averaging strategies as SWA or EMA and manifests better adaptation to different stages of training.

Original: We further design a parallel framework for large-scale training with efﬁciency in memory and computation.
Revised: We further design a parallel framework for large-scale training with efﬁciency in memory and computation.

Original: Extensive experiments demonstrate the superior performance of TWA on benchmark computer vision tasks with various architectures.
Revised: Extensive experiments demonstrate the superior performance of TWA on benchmark computer vision tasks with various architectures.

Original: A T RAINING D ETAILS
Revised: A T RAINING D ETAILS

Original: For SWA experiments, we replicate the SWA baseline by using the publicly released implementation of Izmailov et al.
Revised: For SWA experiments, we replicate the SWA baseline by using the publicly released implementation of Izmailov et al.

Original: We use VGG-16 architecture with batch normalization for a uniﬁed learning rate setting as PreAct ResNet-164.
Revised: We use VGG-16 architecture with batch normalization for a uniﬁed learning rate setting as PreAct ResNet-164.

Original: For ImageNet experiments, we follow ofﬁcial PyTorch implementation.
Revised: For ImageNet experiments, we follow ofﬁcial PyTorch implementation.

Original: 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 to 2 on CIFAR).
Revised: 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 to 2 on CIFAR).

Original: CIFAR experiments are performed on one Nvidia Geforce GTX 2080 TI GPU, while ImageNet experiments are on four NVIDIA Tesla A100.
Revised: CIFAR experiments are performed on one Nvidia Geforce GTX 2080 TI GPU, while ImageNet experiments are on four NVIDIA Tesla A100.

Original: The number of historical solutions optimized by TWA is presented in Table A1.
Revised: The number of historical solutions optimized by TWA is presented in Table A1.

Original: [Table]
Revised: [Table]

Original: We now disclose the speciﬁc hyper-parameters in the following.
Revised: We now disclose the speciﬁc hyper-parameters in the following.

Original: A. SWA TRAINING A.1. CIFAR
Revised: A. SWA TRAINING A.1. CIFAR

Original: We use the same schedule and hyper-parameters as in Izmailov et al.
Revised: We use the same schedule and hyper-parameters as in Izmailov et al.

Original: For VGG-16, we use weight decay of 5 (cid:2) 10 − 4 and train the model for 300 epochs with weight averaging at 161 toepochs.
Revised: For VGG-16, we use weight decay of 5 × 10 − 4 and train the model for 300 epochs with weight averaging at 161 toepochs.

Original: For PreAct ResNet-164, we use weight decay of 3 (cid:2) 10 − 4 and train the model forepochs with weight averaging at 126 to 225 epochs.
Revised: For PreAct ResNet-164, we use weight decay of 3 × 10 − 4 and train the model forepochs with weight averaging at 126 to 225 epochs.

Original: For TWA training, we use the same weights as SWA and initialize w twa as the last checkpoint (i.e.
Revised: For TWA training, we use the same weights as SWA and initialize w twa as the last checkpoint (i.e.

Original: 300 / 225 epochs).
Revised: 300 / 225 epochs).

Original: 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.
Revised: 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.

Original: The regularization coefﬁcient λ is set to 5 (cid:2) 10 − 5 .
Revised: The regularization coefﬁcient λ is set to 5 × 10 − 5 .

Original: A.1. I MAGE N ET
Revised: A.1. I MAGE N ET

Original: Following Izmailov et al.
Revised: Following Izmailov et al.

Original: ; Yang et al.
Revised: ; Yang et al.

Original: (2019), we start from pre-trained models (they are from torchvision.models ) and collect weights by running SGD optimizer up to 10 epochs (with a constant learning rate 0.005, weight decay 1 (cid:2) 10 − 4 ).
Revised: (2019), we start from pre-trained models (they are from torchvision.models ) and collect weights by running SGD optimizer up to 10 epochs (with a constant learning rate 0.005, weight decay 1 × 10 − 4 ).

Original: We uniformly sample the solutions 5 times per epoch.
Revised: We uniformly sample the solutions 5 times per epoch.

Original: For TWA training, we use the same weights as SWA and initialize w twa as the pre-trained model.
Revised: For TWA training, we use the same weights as SWA and initialize w twa as the pre-trained model.

Original: 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 0.
Revised: 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 0.

Original: The regularization coefﬁcient λ is set to 1 (cid:2) 10 − 5 .
Revised: The regularization coefﬁcient λ is set to 1 × 10 − 5 .

Original: A. R EGULAR TRAINING A.2. CIFAR
Revised: A. R EGULAR TRAINING A.2. CIFAR

Original: For regular training, we train the models for 200 epochs with an initial learning rate of 0.1 and decay it by 10 at the 100th and the 150th epoch.
Revised: For regular training, we train the models for 200 epochs with an initial learning rate of 0.1 and decay it by 10 at the 100th and the 150th epoch.

Original: We use SGD optimizer with momentum 0.9, weight decay 1 (cid:2) 10 − 4 , and batch size 128 by convention.
Revised: We use SGD optimizer with momentum 0.9, weight decay 1 × 10 − 4 , and batch size 128 by convention.

Original: For TWA training, we initialize w twa as 1 n ∑ ni =1 w i , i.e., the center of sampled solutions.
Revised: For TWA training, we initialize w twa as 1 n ∑ ni =1 w i , i.e., the center of sampled solutions.

Original: 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.
Revised: 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.

Original: The regularization coefﬁcient λ is set to 1 (cid:2) 10 − 5 .
Revised: The regularization coefﬁcient λ is set to 1 × 10 − 5 .

Original: A.2. I MAGE N ET
Revised: A.2. I MAGE N ET

Original: We follow the training protocol described in He et al.
Revised: We follow the training protocol described in He et al.

Original: Speciﬁcally, we train the models for 90 epochs with an initial learning rate of 0.1 and decay it by a factor of 10 every 30 epochs.
Revised: Speciﬁcally, we train the models for 90 epochs with an initial learning rate of 0.1 and decay it by a factor of 10 every 30 epochs.

Original: We use
Revised: We use

Original: SGD optimizer with momentum 0.9, weight decay 1 (cid:2) 10 − 4 , and batch size 256 .
Revised: SGD optimizer with momentum 0.9, weight decay 1 × 10 − 4 , and batch size 256 .

Original: For TWA training, we uniformly sample solutions 5 times per epoch and initialize w twa asn ∑ ni =1 w i .
Revised: For TWA training, we uniformly sample solutions 5 times per epoch and initialize w twa asn ∑ ni =1 w i .

Original: We train the models for 2 epochs with a learning rate of 0.3 and 0.03, respectively.
Revised: We train the models for 2 epochs with a learning rate of 0.3 and 0.03, respectively.

Original: For the extra one epoch training, we use the same training protocol as in subsection A.1.2, i [Table] e., linearly decaying the learning rate from 0.03 to 0.
Revised: For the extra one epoch training, we use the same training protocol as in subsection A.1.2, i.e., linearly decaying the learning rate from 0.03 to 0.

Original: The regularization coefﬁcient λ is set to 1 (cid:2) 10 − 5 .
Revised: The regularization coefﬁcient λ is set to 1 × 10 − 5 .

Original: A.2.  R ESULTS ON A DAM O PTIMIZER
Revised: A.2.  R ESULTS ON A DAM O PTIMIZER

Original: Adam (Kingma & Ba, 2015) is another mainstream optimizer with adaptive gradient descent, which enjoys fast convergence and insensitivity to the initial learning rate.
Revised: Adam (Kingma & Ba, 2015) is another mainstream optimizer with adaptive gradient descent, which enjoys fast convergence and insensitivity to the initial learning rate.

Original: Here, we apply TWA to the solutions generated by Adam optimizer and the results are in Table A2.
Revised: Here, we apply TWA to the solutions generated by Adam optimizer and the results are in Table A2.

Original: The training settings are the same as in Table 3 with default β 1 = 0 .
Revised: The training settings are the same as in Table 3 with default β 1 = 0 .

Original: 9 and β 2 = 0 .
Revised: 9 and β 2 = 0 .

Original: 999 for Adam.
Revised: 999 for Adam.

Original: TWA is trained for 5 epochs with a intial learning rate 3.
Revised: TWA is trained for 5 epochs with a intial learning rate 3.

Original: We observe that it could similarly bring generalization improvement and training efﬁciency
Revised: We observe that it could similarly bring generalization improvement and training efﬁciency.

Original: [Table]
Revised: [Table]

Original: [Table]

Original: B I MPLEMENTATION
Revised: B I MPLEMENTATION

Original: The optimization target for TWA is given by
Revised: Let β = [ β 1 , β 2 , · · · , β n ] ⊤ ∈ R n , P = [ e 1 , e 2 , · · · , e n ] ∈ R D × n , the optimization target for TWA can be formulated as,

Original: [Table]
Revised: [Equation]

Revised: We short the loss term E ( x , y )

Revised: ∼D [ L ( f ( w ; x ) , y )] to L ( w ) and the gradient w.r.t.

Revised: β can be written as,

Original: [Equation]
Revised: [Equation]

Original: [-where A = f α 1 w 1 + α 2 w 2 + (cid:1) (cid:1) (cid:1) + α n w n j α 
-]
Revised: [-where A = f α 1 w 1 + α 2 w 2 + (cid:1) (cid:1) (cid:1) + α n w n j α 
-]Let η be the learning rate. We could optimize (A.1) with the following gradient descent:

Original: + α 2 + (cid:1) (cid:1) (cid:1)
Revised: [Equation]

Original: + α n = 1 g is the possible solution set of w twa .
Revised: Since w twa = P β , we have the corresponding update in the parameter space: [Equation]

Original: We have discussed in the main paper that the averaging coefﬁcients could be alternatively optimized in the parameter space, and the sum-one constrain can be naturally satisﬁed inoptimizing the left term of (A.1) (the empirical loss on training data).
Revised: As β does not explicitly appear in (A.6), we could treat the coefﬁcient β as an implicit variable. In practice, we optimize (A.1) by directly updating the model weights w twa with weight decay λ , which is an optimization in the reduced subspace with projection matrix P P ⊤ .

Original: Hence, we do not need to explicitly impose the constraint and the solution set A could be relaxed to span f w
Revised: We test the performance of SWA with different starting epochs to average on ImageNet with the

Original: i g ni =1 .
Revised: ResNet-50 model.

Original: In the following, we ﬁrst convert the gradients w.r.t.
Revised: We 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.

Original:  averaging coefﬁcients into the parameter space.
Revised: We also notice that TWA (with 0-60 epoch solutions) consistently outperforms the SWA wherever the averaging begins, conﬁrming that TWA could automatically ﬁnd a good set of averaging coefﬁcients and provide better performance. D D ISCUSSION ON THE S UM - ONE C ONSTRAINT Let α = [ α 1 , α 2 , · · · , α n ] ⊤ ∈ R n , W = [ w 1 , w 2 , · · · , w n ] ∈ R D × n , we have w twa =

Original: The partial derivative of the loss function L w.r.t. the averaging coefﬁcient α 
Revised: W α . We multiply W ⊤ on the both sides, i.e., W ⊤ w twa = W ⊤ W α , and could obtain

Revised: [Table]

Original: i is
Revised: α

Original: [Equation]
Revised: [Equation]

Original: Since there exists a bijection betweenf αi g ni =1 and span f w i
Revised: Further, we could establish the relation between α and β : α = ( [Equation]

Original: g ni =1 , mapping the above derivatives into the parameter space and we have
Revised: −

Original: [Equation]
Revised: [Equation]

Original: Let d L d w be the gradient of {+averaging coefﬁcients ∑ ni =
 α i to be 1 since+} the empirical loss w.r.t.
Revised: ⊤ P β , since they are the coordinates of w twa under two different set of bases. In the solution set A = { α 1 w 1 + α 2 w 2 + · · · + α n w n | α i ∈ R } , we do not explicitly require that the sum of {+averaging coefﬁcients ∑ ni =
 α i to be 1 since+} the network’s performance would be sensitive to a direct scaling over all parameters (i.e. k w ), that is, a good solution in A will inherently have a coefﬁcient sum very close to 1.

Original:  the model parameters w  and P = 
Revised: We verify the sum of the averaging coefﬁcients of the attained solution w twa on CIFAR-100 with PreResNet-164 model and observe that ∑ ni =1 α i = 1 . 02 ± < 0 .

Original: [ e 1 , e 2 , (cid:1) (cid:1) (cid:1) , e n ] be the orthogonal basis vectors of the subspace span f w
Revised: We numerically measure the averaged epoch training time and memory burden for SGD and TWA in the DDP training setting. Speciﬁcally, we experiment with the ResNet-50 model on ImageNet and use 1, 2, and 4 GPUs with a batch size of 256 per GPU and a total of 300 historical solutions. The experiments are conducted on NVIDIA Tesla A100 40G GPUs.

Original: i g ni
Revised: From the results reported in Table E4, we observe that TWA brings minor additional costs, e.g. +2.8% on time cost and +2.9% on 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 efﬁcient and scalable averaging for large-scale problems.

Original: =1 .
Revised: [Table]

Original: To conduct optimization in subspace, we project the gradient d L d w (in full parameter space) onto the subspace via the projection matrix P P ⊤ and obtain the complete gradient of loss function L w.r.t.
Revised: We conduct an ablation study in Table F5 to analyze the impact of the regularization coefﬁcient λ . We observe that such regularization brings improvements but is not signiﬁcant. This is because the main regularization effects come from the signiﬁcant decrease of training variables,

Original: w twa as following:
Revised: i [Table] 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 it in our method.

Original: [Equation]
Revised: [Table]

Original: Thus, we could optimize (A.1) with  the following gradient descent:
Revised: For NLP datasets, we try a ﬁnetune task with pre-trained models and compare the performance of SWA and TWA. Speciﬁcally, we experiment with The Corpus of Linguistic Acceptability (CoLA), a text classiﬁcation task in the General Language Understanding Evaluation (GLUE, Wang et al.

Original: [Equation]
Revised: (2018)) benchmark. In experiment, we use a pre-trained BERT (Devlin et al., 2018) model ( bert-base-uncased ) from Hugging Face community 2 . We ﬁne-tune BERT on CoLA forepochs with AdamW optimizer Loshchilov & Hutter (2017), learning rate 2e-5, and weight decay 0.0. The model weights at the end of these epochs are collected for SWA and TWA.

Original: where one could see clearly that the ℓ 2 -norm regularizationon averaging coefﬁcients is equivalent to weight decay on model parameters in reduced subspace.
Revised: In TWA, we train the ﬁne-tuned model for 1 epoch with a learning rate of 0.5 and regularization coefﬁcients λ = 0 . From the results below, we observe that TWA could achieve better performance than the competing methods. This further demonstrates the broad application of TWA. [Table]