From: ICLR.cc/2023/Conference/Program_Chairs
title: Paper Decision
decision: Accept: poster
metareview:_summary,_strengths_and_weaknesses: This paper studies a trainable way of weight averaging (over the history of training) and evaluates the approach on several image datasets. The scheme works both in a single-node and multi-node setting.
Reviewers liked the clear presentation and ideas, and mostly found the results convincing.

We hope the authors will incorporate the several points mentioned by the reviewers in the final version. 
justification_for_why_not_higher_score: Scores don't justify oral presentation yet, in view of the required level of novelty, while results are still solid
justification_for_why_not_lower_score: No major concerns remained
summary_of_AC-reviewer_meeting: We tried to schedule but in the end had to finalize the discussions discussions asynchronously. Two reviewers raised their score, after good discussions with the authors, so we now have solid consensus for this being above the bar.
note_from_PC: if the above contains the word "oral" or "spotlight" please see: "oral" presentation means -> notable-top-5% and "spotlight" means -> notable-top-25%. As stated in our emails, we are disassociating presentation type from AC recommendations

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Thanks to Reviewer zFxA
comment: Thanks again for your time and voluntary work!
We are very glad our response could address your major concerns on muti-node training. 
As you concerned, the sum-one constraint was confusing in our original version. In the revision, we have dropped it directly, since it would not hurt the performance. Actually, we have checked the solutions without the sum-one constraint and found that their sum is very close to 1: e.g., $1.02_{\pm< 0.01}$ on CIFAR-100 with PreResNet-164. We will further clarify this point in the final version.

From: ICLR.cc/2023/Conference/Paper1466/Reviewer_zFxA
title: Further Feedback 
comment: I am not well convinced by the response of Q1. However, Q3 is my most concerning question, which is quite important for practical usage, and the authors' response addresses my concern about this problem. So I raise my score from 5 to 6.

By the way, sorry for the late reply due to some personal stuff.

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: We are still willing to address your concerns before the discussion ending
comment: Dear Program Chairs, Area Chairs, and Reviewers,

We are very grateful for your efforts in reviewing our paper and providing constructive comments, which greatly help us improve the work. We also appreciate the positive feedback from reviewers after our response. As the time is coming to the discussion ending, we are willing to start further discussions and do our best to address the Reviewers' concerns, which might help Area Chairs with better decisions.

We wait in hope that the reviewers could update the review and give us valuable feedback!
    


Best regards,      
Paper 1466 Authors

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Thanks again for the detailed comments and looking forward to your feedback.
comment: Dear Reviewer zFxA,

Thanks again for your thoughtful comments. As the discussion period is close to the end, we would appreciate it if you could kindly let us know whether the response has addressed your concerns. And we are happy to discuss further if you have any questions or concerns.

Authors Paper1466

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Thanks again for the detailed comments and looking forward to your feedback.
comment: Dear Reviewer nodd,

Thanks again for your thoughtful comments. As the discussion period is close to the end, we would appreciate it if you could kindly let us know whether the response has addressed your concerns. And we are happy to discuss further if you have any questions or concerns.

Authors Paper1466

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Thank you for the postive feedback!
comment: We are very grateful for your reading our responce and your reconsideration. We would keep improving our manuscript to involve your insights.

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Continue with the above response (we add more details)
comment: **Q3:**
The way to perform TWA in the parallel distributed training manner is not equivalent to the single-node TWA.

**A3:** Thanks for the insightful question!

We have synchronized the gradients of each node to $\frac{1}{n}\sum\_{i=1}^n\boldsymbol{g}_i$ before the gradient projection and thus the updated rule for single-node and multi-nodes are exactly equivalent. Specifically, let $\boldsymbol{P}=[\boldsymbol{P}\_1,\boldsymbol{P}\_2,\cdots,\boldsymbol{P}\_n]$. With basic block matrix multiplication, we have
$$\boldsymbol{P}\boldsymbol{P}^{\top}\boldsymbol{g}=\frac{1}{n}[\boldsymbol{P}\_1,\boldsymbol{P}\_2,\cdots,\boldsymbol{P}\_n][\boldsymbol{P}\_1^{\top},\boldsymbol{P}\_2^{\top},\cdots,\boldsymbol{P}\_n^{\top}]^{\top}\boldsymbol{g}=\frac{1}{n}\sum\_{i=1}^n\boldsymbol{P}\_i\boldsymbol{P}\_i^{\top}\boldsymbol{g}.$$
For single-node, the updated rule is $\frac{1}{n}\boldsymbol{P}\boldsymbol{P}^{\top}(\sum\_{i=1}^n\boldsymbol{g}\_i)$ as you pointed out. For multi-node, we will first perform an all-reduce operation to synchronize the local gradients of each node to $\frac{1}{n}\sum\_{i=1}^n\boldsymbol{g}\_i$ and then obtain the local gradient projection $\frac{1}{n}\boldsymbol{P}\_i\boldsymbol{P}\_i^{\top}(\sum\_{k=1}^n\boldsymbol{g}\_k)$. We finally synchronize it to $\frac{1}{n}\sum\_{i=1}^n\left[\boldsymbol{P}\_i\boldsymbol{P}\_i^{\top}(\sum\_{k=1}^n\boldsymbol{g}\_k)\right]$ with another all-reduce sum operation. Thus, the two ways are mathematically equivalent under the multiplication of the block matrix: 
$$\frac{1}{n}\boldsymbol{P}\boldsymbol{P}^{\top}(\sum\_{i=1}^n\boldsymbol{g}\_i)=\frac{1}{n}[\boldsymbol{P}\_1,\boldsymbol{P}\_2,\cdots,\boldsymbol{P}\_n][\boldsymbol{P}\_1^{\top},\boldsymbol{P}\_2^{\top},\cdots,\boldsymbol{P}\_n^{\top}]^{\top}(\sum\_{i=1}^n\boldsymbol{g}\_i)=\frac{1}{n}\sum\_{i=1}^n\left[\boldsymbol{P}\_i\boldsymbol{P}\_i^{\top}(\sum\_{k=1}^n\boldsymbol{g}\_k)\right].$$
We note that in this way, the computation for gradient projection is also uniformly distributed to different nodes.

---

Thanks again for your detailed comments! Besides, we respectfully hope the reviewer could reconsider the score after reading our clarifications, and please kindly let us know if there are any additional questions or concerns.

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Continue with the above response (we add more details)
comment: **Q5:**
In subsection 2.1, I don't get why lower training loss by optimisation over to {alpha_i}_{i=1}^n can lead to better generalisation. What is the logic behind it?

**A5:**
Your question is quite essential. Indeed, training loss and test loss is not necessarily positively correlated, which is one of the most challenging problems for analyzing and designing deep learning method. Our main logic for minimizing training loss in subsection 2.1 is that in some cases, equally averaging would cause estimation error (e.g. in the head training stage / the variance of the solution distribution is large), and we hope to correct it for better performance under the supervision of training loss.

Moreover, we consider that minimizing the training loss for averaging the solutions that **not well-trained** (e.g. from the head training stage) is meaningful: the degree of freedom for TWA training is n, which is very small, and thus the averaged model is unlikely to overfitting since the collected solutions themselves are not well-trained. In this sense, minimizing the training loss is expected to attain good generalization. In our experiments, TWA could achieve better performance than SWA with fewer averaging solutions used (Table 2) and significantly reduced generalization gap between training and test accuracies compared with regular training (Table 3, 4).

**Q6:**
No source code available for reproducibility.

**A6:**
Thanks for your reminder. The source code is available in supplementary materials. 

**Q7:**
TWA essentially introduces an addition training procedure after for example SGD. It would be nice to combine SGD and TWA into one training procedure.

**A7:**
Thanks for your discussion. 

We consider that the additional training procedure is an inherent attribute of the SWA series approaches, which needs to first collect the historical solutions and then perform averaging. Practically, we find that it may not be good to switch between the TWA training and SGD training frequently: we observe that continuing SGD training from an SWA/TWA solution will first suffer a severe degradation the performance (we show a finetune experiment from the SWA solution below, with SGD optimizer and learning rate 0.0005, weight decay 0.0001, batch size 128, 10 training epochs). This is perhaps because the solution of SWA/TWA is from averaging and lies in a relatively flat region. Continuing SGD training would escape from that flat region SWA/TWA lies and thus lose the good generalization. This also shows the advantages of TWA training capable of keeping good generalization by constraining the training in tiny subspaces.

| Model | SWA | TWA | SWA (+FT)
| --- | --- | --- | --- |
VGG16 | 74.71 | 75.73 | 74.01
PreResNet-164 | 80.20 | 80.35 | 77.79

---
Thanks again for your detailed comments! Besides, we respectfully hope the reviewer could reconsider the score after reading our clarifications, and please kindly let us know if there are any additional questions or concerns.

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Continue with the above response.
comment: **Q4:**
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.

**A4:**
Thanks for your suggestion.

We add an ablation study to analyze the impact of regularization coefficients in Appendix F of this version. We observe that the regularization term indeed brings improvements but is not that significant: the difference between with/without the regularization term does not exceed 0.2%. This confirms that reducing the number of training variables from D (the number of model parameters) to n (the number of coefficients) itself brings significant regularization effects. Since such regularization term is easy to implement and virtually brings little computation, we include it in our training target.

Datasets | Model | TWA | TWA ($\lambda=0$)
|---|---|---|---|
| | VGG16 | 93.79 | 93.78 
CIFAR-10 | PRN-164 | 95.11 | 95.03
 | | VGG16 | 74.24 | 74.02
CIFAR-100 | PRN-164 | 78.11 | 78.01

For optimizing coefficients in a subspace, we provide an efficient and scalable scheme: 1) Instead of directly optimizing the coefficients, we convert it to subspace training, which could be efficiently realized with matrix production; 2) Further, such matrix operation could be readily fitted into a distributed training scheme and brings minor additional cost compared with regular training. Thus, our scheme is able to handle large-scale problems, in which reducing the training epochs is of much more significance.


**Q5:**
The statements in Section 3 are not well-supported, neither theoretically nor empirically.

**A5:**
Thanks for your comments.

Section 3 is based on the assumption that the solutions to averaged $\boldsymbol{w}_i$ comes from a Gaussian distribution $\mathcal{N} \left( \boldsymbol{\mu}, \boldsymbol{\Sigma} \right)$, which is also the basic assumption in the original SWA paper [1]. We state that if $\boldsymbol{\mu}$ is the center minimum with the ''lowest training loss'', then TWA could approach it further by optimizing the coefficients, and obtain lower expected estimation variance w.r.t. $\boldsymbol{\mu}$, compared with SWA. Honestly, it is hard to characterize the distribution of the solutions and the generalization gap between training data and test data exists. We only give an intuitive explanation on the advantages of TWA. In this version, we have added more discussion on this point.

But empirically, we observe that TWA could improve the SWA solution slightly at the original SWA setting (i.e., SWA_lr=0.05) and obtain significant improvements over SWA when SWA_lr is larger  (i.e., SWA_lr=0.10) where the variance of the distribution correspondingly increases. For example, on CIFAR-100, we achieve improvements of +0.73% on PRN-164 and +4.18% on VGG16 compared with SWA, confirming that optimizing the coefficients indeed helps reduce the expected estimation variance and thereby improve the performance.

[1] Izmailov, Pavel, et al. "Averaging weights leads to wider optima and better generalization." arXiv preprint arXiv:1803.05407 (2018).

Thanks again for your detailed comments! Besides, we respectfully hope the reviewer could reconsider the score after reading our clarifications.

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Thank you for the detailed comments! Please see our responses below.
comment: **Q1:**
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. 

**A1:**
Thanks for pointing out the issue.

Our original goal is to show that TWA provides a "cheap" additional training procedure for better averaging performance, e.g., on CIFAR with VGG in Table 1, SWA averaging 140 epochs of sampled solutions and TWA trains for additional 10 epochs, which is minor. On ImageNet, TWA requires additional one epoch training, but we observe that TWA with 5 epoch solutions (thus consuming a total of 5 + 1 epochs) is better than SWA with 10 epochs. In this version, we have clearly reported the additional epochs (i.e. +10, +1) for TWA in Table 1,2. We also would like to note that TWA starts from the last epoch solution instead of using the SWA solution as the pretraining model in our SWA part experiments for a fair comparison.

To answer your concerns, we first conduct an experiment for TWA with the additional training epochs reduced for sampling on CIFAR-100 (e.g. 140 epochs for SWA while 130 for TWA), for a fair comparison. We observe that this brings a minor influence on the performance since such 10 epochs are relatively small compared with the total samplings.

| Model | SWA | TWA | TWA (-10)
| --- | --- | --- | --- |
VGG16 | 74.71 | 75.73 | 75.77
PreResNet-164 | 80.20 | 80.35 | 80.28


For fine-tuning from the SWA model with the same amount of training epochs as TWA, we use SGD optimizer with a learning rate of 0.0005, batch size 128, weight decay 0.0001, and 10 epochs.
We observe that the performance significantly degrades. This is perhaps because additional training would escape from the original flat region SWA lies and severely degenerate the performance.

| Model | SWA | TWA | SWA (+FT)
| --- | --- | --- | --- |
VGG16 | 74.71 | 75.73 | 74.01
PreResNet-164 | 80.20 | 80.35 | 77.79



**Q2:**
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.

**A2:**
Thanks for your question. 

The number of checkpoints used in shown in Appendix A, where for CIFAR experiments in Table 6, 100 checkpoints are used.

We add a related numerical discussion for Section 4 in Appendix E. Specifically, 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. From the results reported below,
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 efficient and scalable averaging for large-scale problems. 

| GPU | Time | Time | Memory | Memory |
| --- | --- | --- | --- | --- |
|  | SGD | TWA | SGD | TWA |
| 1 | 1638s | 1692s (+3.3%) | 28286.5 MB | 31432.5 MB (+11.1%) |
| 2 | 824s | 862s (+4.6%) | 28382.5 MB | 30092.5 MB (+6.0%) |
| 4 | 420s | 432s (+2.8%) | 28874.5 MB | 29718.5 MB (+2.9%) |


**Q3:**
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.

**A3:**
Thanks for your valuable question!

TWA aims to reduce the training epochs and improve the training efficiency while keeping good performance. SAM requires 2x computation as regular training and thus we are not expected to achieve similar performance as SAM. 

For Lookahead, we apply TWA to the original setting in the Lookahead paper with ResNet-18 model. We use $k=5$ and $\alpha=0.5$ for Lookahead and use the SGD solutions for TWA training.
From the results below, we observe that Lookahead achieves faster convergence than SGD at 110 epochs but is much inferior to that of TWA. We also notice that it does not bring further improvement than SGD at the end of training (200 epochs) in our experiments. Detail training settings are follow this  [repo](https://github.com/uoguelph-mlrg/Cutout).

Datasets | SGD (200) | Lookahead (200) | SGD (110) | Lookahead (110) | TWA (100+10)
--- | --- | --- | --- | --- | ---
CIFAR-10 | 95.15 | 95.11 | 91.48 |  92.58 | 95.06
CIFAR-100 | 77.91 | 77.59 | 68.69 | 71.37 | 77.74

For ImageNet experiments, we demonstrate that TWA could achieve faster convergence than Lookahead with better performance (e.g. 75.70% v.s 75.49% in Section 5.3). 

We consider the direct application of TWA is to improve the training efficiency of regular training by averaging historical solutions, especially for large-scale problems. In this version, we also try NLP tasks and obtain promising results on NLP tasks (please see Appendix G), showing the broad application of TWA.

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Thank you for the detailed comments! Please see our responses below.
comment: Many thanks for your thoughtful comments and valuable suggestions! We have revised our paper following the suggestions and addressed your concerns in the following:

**Q1:** The solution for the problem in Eq.(2) does not satisfy the L1 ball constrain in Eq.(1).

**A1:** Thanks for pointing out the issue. 

We do not explicitly impose the sum-one constraint during the optimization, since 1) by doing so, the feasible set is enlarged and will not hurt the performance, and 2) the implementation for L2 square regularization on coefficients would be easier (details of the derivation are provided in Appendix B). In practice, we examine the sum of the averaging coefficients of the attained solution $\boldsymbol{w}\_{\rm twa}$ on CIFAR-100 with PreResNet-164 model (without sum-one constraint) and observe that
$\sum\_{i=1}^n \alpha\_i=1.02_{\pm< 0.01}$. This is perhaps because the network's performance would be sensitive to a direct scaling over all parameters (i.e. $\boldsymbol{w}\rightarrow k\boldsymbol{w}$), and a model with good performance commonly reaches a coefficient sum very close to one.

We realize that the discussion in our original version is a bit confusing, as pointed out by you and Reviewer nodd. Thus in the revision, we skip introducing this constraint in the target formulation and put the above discussion on the sum-one constraint in Appendix D. We hope this would be more straightforward to elaborate our method and clearer for readers to pick up.


**Q2:** Eq.(A.4) is not quite right.  The gradient w.r.t. $\boldsymbol{w}_{\rm twa}$  is not that in Eq.(A.4).

**A2:**
Thanks for your careful reading and insightful comments.
The discussion in the previous version is indeed not rigorous and brings some ambiguity in elaborating on the optimization. 

More specifically, in order to better optimize the weighted coefficients over the historical solutions, we first orthogonalize $\lbrace \boldsymbol{w}\_i \rbrace_{i=1}^n$ into bases
$\lbrace \boldsymbol{e}\_i \rbrace\_{i=1}^n$, which supports the solution space covering the weighted summation of $\lbrace \boldsymbol{w}\_i \rbrace\_{i=1}^n$. Next, we reformulate the weighted summation over $\lbrace \boldsymbol{w}\_i \rbrace\_{i=1}^n$ into  $\lbrace \boldsymbol{e}\_i\rbrace\_{i=1}^n$, i.e., $\boldsymbol{w}\_{\rm twa}=\sum\_{i=1}^n \beta\_i\boldsymbol{e}\_i$, where the $\ell_2$ norm regularization is imposed on $\beta\_i$. In this way, the optimization to our  $\boldsymbol{w}\_{\rm twa}$ corresponds to the optimization over $\lbrace \beta_i\rbrace \_{i=1}^n$. Accordingly, the previous Eq.(A.1) is refined as 
$$
    \min\_{\boldsymbol{\beta}} \quad 
L(\boldsymbol{\beta})\triangleq\mathbb{E}\_{(\boldsymbol{x},\boldsymbol{y})\sim \mathcal{D}} \left[
    \mathcal{L}  \left ( f \left (\boldsymbol{w}\_{\rm twa}; \boldsymbol{x} \right ), \boldsymbol{y} \right ) 
    \right ]
    +
    \frac{\lambda}{2}\boldsymbol{\beta}^\top\boldsymbol{\beta}, \quad
    \rm{s.t.}\quad \boldsymbol{w}_{\rm twa} = \boldsymbol{P}\boldsymbol{\beta}.
$$

Thus, we can focus on the optimization over $\lbrace \beta_i\rbrace_{i=1}^n$, i.e., $\boldsymbol \beta$. By taking derivatives w.r.t  $\boldsymbol \beta$, we have 
$$
    \frac{\partial L}{\partial \boldsymbol{\beta}}=
    \frac{\partial \boldsymbol{w}\_{\rm twa}}{\partial \boldsymbol{\beta}} \frac{\partial \mathcal{L}(\boldsymbol{w}\_{\rm twa})}{\partial \boldsymbol{w}\_{\rm twa}} + \lambda \boldsymbol{\beta} \\
    =\boldsymbol{P}^\top \frac{\partial \mathcal{L}(\boldsymbol{w}\_{\rm twa})}{\partial \boldsymbol{w}\_{\rm twa}} + \lambda \boldsymbol{\beta},
$$
so that Eq.(A.4) in the previous version was obtained. Now, the optimization proceeds in the $n$-dimensional subspace, and further we can apply the gradient descent to $\boldsymbol \beta$, such that 
$$
     \boldsymbol{\beta}^{(t+1)}=\boldsymbol{\beta}^{(t)}-\eta \left( \boldsymbol{P}^\top \frac{\partial \mathcal{L}(\boldsymbol{w}\_{\rm twa}^{(t)})}{\partial \boldsymbol{w}\_{\rm twa}} + \lambda \boldsymbol{\beta} \right).
$$
Due to $ \boldsymbol{w}\_{\rm twa} = \boldsymbol P \boldsymbol \beta $, we finally update $ \boldsymbol{w}\_{\rm twa}$ with
$$
      \boldsymbol{w}\_{\rm twa}^{(t+1)}=\boldsymbol{w}\_{\rm twa}^{(t)}-\eta \left( \boldsymbol{P}\boldsymbol{P}^\top \frac{\partial \mathcal{L}(\boldsymbol{w}\_{\rm twa}^{(t)})}{\partial \boldsymbol{w}\_{\rm twa}} + \lambda \boldsymbol{P}\boldsymbol{\beta}  \right)
     =\left ( 1 - \eta\lambda \right ) \boldsymbol{w}\_{\rm twa}^{(t)} -\eta \boldsymbol{P}\boldsymbol{P}^\top \frac{\partial \mathcal{L}(\boldsymbol{w}\_{\rm twa}^{(t)})}{\partial \boldsymbol{w}\_{\rm twa}}.
$$
Therefore, the previous Eq.(A.5) is attained as one step of the projection gradient descent method with $\ell_2$ square regularization, as you mentioned in the comment. 

In the revised Appendix B, we refine the derivations in the optimization and also rewrite the relevant parts in the main paper. We hope this version could clarify this issue and resolve your concerns.

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Thank you for the detailed comments! Please see our responses below. 
comment: Many thanks for your thoughtful comments and valuable suggestions! We are glad of your finding our work interesting. We have revised our paper following the suggestions and addressed your concerns in the following:

**Q1:**
I don't think the sum-one constraint in (1) can be naturally satisfied in Algorithm 1.

**A1:**
Thanks for pointing out the issue and providing case explanations. Indeed, the projected gradient $\boldsymbol{P}\boldsymbol{P}^\top \boldsymbol{g}$ lies in the space spanned by $\lbrace \boldsymbol{w}\_i \rbrace\_{i=1}^n$, where the sum-one constraint may not necessarily be satisfied as you concerned. We realize that the discussion in our original version was a bit confusing, as pointed out by you and Reviewer zFxA. In the revision, we skip introducing this constraint for the target formulation (previously Eq.(1)) as it is not used when solving the problem and put the related discussion on it in Appendix D.
By doing so, the feasible set is enlarged without hurting the performance, and the implementation for L2 square regularization on coefficients would be easier (details of the derivations are provided in Appendix B).  

Nonetheless, in practice, we find that a trained $\boldsymbol{w}\_{\rm twa}$ solution (without sum-one constraint) commonly reaches a coefficient sum very close to one. We examine the sum of the averaging coefficients of the attained solution $\boldsymbol{w}\_{\rm twa}$ on CIFAR-100 with PreResNet-164 model and observe that $\sum\_{i=1}^n \alpha_i=1.02\_{\pm< 0.01}$. This is perhaps because the network's performance would be sensitive to a direct scaling over all parameters (i.e. $\boldsymbol{w} \rightarrow k\boldsymbol{w}$), and a model with good performance should conform to the “sum-one” constraints (otherwise the solution could be further optimized by scaling it to one).
We hope that the above explanations and revisions made would be more straightforward to elaborate our method and clearer for readers to pick up.

**Q2:**
The derivation of (4) is questionable.

**A2:**
Thanks for your valuable comment.
Our discussion on this point has an important pre-condition that the distribution center "$\boldsymbol{\mu}$ is the minimum", which is also the starting point of TWA. Since $\boldsymbol{\mu}$ corresponds to the minimum, minimization over the coefficients gives a better estimation than averaging. According to your question, we have modified the manuscript with more discussion on this point.

Empirically, the reduction in expected estimation variance could be validated: we observe that TWA could improve the SWA solution slightly at the original SWA setting (i.e., SWA_lr=0.05) and obtain significant improvements over SWA when SWA_lr is larger  (i.e., SWA_lr=0.10) where the variance of the distribution correspondingly increases. For example, on CIFAR-100, we achieve improvements of +0.73% on PRN-164 and +4.18% on VGG16 compared with SWA, confirming that optimizing the coefficients indeed helps reduce the expected estimation variance and thereby improve the performance.

**Q3:**
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.

**A3:**
Thanks for discussing the practical setting. Indeed, the selection of historical solutions has an impact on the performance. But in practice, we do not need careful manually selecting, since the performance is not that sensitive and the schedule could be pre-given along with the original training schedule of regular training. In fact, we use the same setting, not specifically tuned, for all the reported experiments, e.g. the first 100 epochs for CIFRA and 60 epochs for ImageNet. Also, our TWA could automatically adapt to different training stages (e.g. the early stage) by learnable coefficients. In this sense, we consider our TWA training would not bring much additional effort on manually selecting.

**Q4:**
 Multiple GPUs are needed to allow for parallel computation for gradient projection in TWA, which makes TWA less practical than SWA.

**A4:**
Thanks for your question. 

We use multiple GPUs only in ImageNet experiments, which typically require distributed parallel training. For CIFAR experiments, we only use one GPU for training. Besides, our TWA brings minor additional computation/memory costs compared with regular training, especially in distributed training settings (as shown in Appendix E of the updated version), that is, we could conduct averaging aside the regular training by utilizing the remaining free GPU memory, with little additional effort. Considering that TWA could be easily applied to average the solutions from the head stage of training and result in around 30%-40% of training epochs saving, we think it would be applicable in real-world scenarios.

From: ICLR.cc/2023/Conference/Paper1466/Authors
title: Thank you for the detailed comments! Please see our responses below.
comment: We thank the reviewer for the encouraging and helpful feedback. We have revised our paper following the suggestions and addressed your concerns in the following:

**Q1:**
One straightforward way to make this paper stronger is to try other architectures and NLP datasets.

**A1:**
Thanks for your valuable suggestions!

We have tried TWA to the head training stage with other architectures on CIFAR-100 datasets. We observe that TWA could bring better performance with much fewer training epochs required.

| Model | SGD (200) | SWA(100) | TWA(100+10)|
| --- | --- | --- | --- |
DenseNet121 | 76.42 | 74.09 | **76.92** 
MobileNetV2 | 72.01 | 66.74 | **73.01**
WideResNet-16x8 | 78.39 | 76.94 | **78.69**


For NLP datasets, we try a finetune task with pre-trained models and compare the performance of SWA and TWA.
Specifically, we experiment with The Corpus of Linguistic 
Acceptability (CoLA), a text classification task in the General Language Understanding Evaluation (GLUE[1]) 
benchmark. In experiment, we use a pre-trained BERT[2] model (`bert-base-uncased`) from 
[Hugging Face community](https://huggingface.co/). We fine-tune BERT on CoLA for 3 epochs with AdamW optimizer, learning rate 2e-5
and weight decay 0.0. The model weights at the end of these epochs are collected for SWA and TWA. In TWA, we train 
the fine-tuned model for 1 epoch with learning rate 5e-1 and regularization coefficients $\lambda$=1e-3. From the results below, we observe that TWA could achieve better performance than the competing methods. We add these results in Appendix G. This further demonstrates the broad application of TWA.

| Fine-tune [3] | SWA | TWA |
| --- | --- | --- |
| 56.81 | 59.73 | **60.36** |

[1]: Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., & Bowman, S.R. (2018). GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding. *BlackboxNLP@EMNLP*.

[2]: Devlin, J., Chang, M., Lee, K., & Toutanova, K. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. *ArXiv, abs/1810.04805*.

[3] repo: https://huggingface.co/


**Q2:**
I wonder  if this hypothesis can be tested where SWA is evaluated by varying the number of iterates to average from the end.


**A2:**
Thanks for your suggestion.

Accordingly, we test the performance of SWA with different epochs to start on ImageNet with ResNet-50 model. 
It could be observed 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. 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 automatically find a good set of averaging coefficients and provide better performance. We add these results in Appendix C.
| SWA  (0-60) | SWA (10-60) | SWA (20-60) | SWA (30-60) | SWA (40-60) | SWA (50-60) | TWA (0-60) |
|--- | --- | --- | --- | --- | --- | --- |
67.65 | 70.50 | 72.12 | 74.14 | 75.08 | 75.34 | **75.90**


**Q3:**
If this strategy can help for averaging multiple models at the end of the training.

**A3:**
Thanks for your insightful question. 
The performance of weighted averaging heavily relies on the diversity and the quality of the models. 

These days, we tried the suggested strategy for CIFAR and ImageNet. For CIFAR, averaging multiple models, by SWA or TWA, at the end of the training does not bring further improvements. 
This is perhaps due to that modern DNNs typically overfit the CIFAR training set with almost zero training loss, and the solutions at the end of training are very close to each other. 

For ImageNet, averaging multiple models at the end of training indeed helps, since there is still diversity among models. We use the solutions collected at 61-90 epochs and observe that TWA with additional one-epoch training brings significant accuracy improvements (as shown in Table 5).

| Model | SGD | TWA (+1) |
| --- | --- | -- |
ResNet-18 | 69.82 | **70.37**
ResNet-50 | 75.82 | **76.34**

This shows the potential of applying TWA to average the solutions at the tail training stage or trained by different configurations to attain better performance.

---
Thanks again for your constructive comments! Please kindly let us know if there are any additional questions or concerns.

From: ICLR.cc/2023/Conference/Paper1466/Reviewer_NH8Y
confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
summary_of_the_paper: This paper introduces a trainable way to perform weight averaging and explores it on several image classification datasets.
strength_and_weaknesses: Strengths:
- I like the generalization of traditional weight averaging by explicitly training for it, especially since it empirically helps.
- Authors also consider the computational cost of TWA and how to apply it in a distributed setting.
- The paper does a great job of thoroughly experimenting on large empirical problems.

Weaknesses / Areas of improvement:
- 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.
- 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.
- 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. 
clarity,_quality,_novelty_and_reproducibility: No concerns related to clarity or reproducibility
summary_of_the_review: This work generalizes SWA by training the weights explicitly. The ideas are novel and experiments are thorough.
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.
flag_for_ethics_review: NO.
recommendation: 8: accept, good paper

From: ICLR.cc/2023/Conference/Paper1466/Reviewer_nodd
confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
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.       
strength_and_weaknesses: Strength:  
1) The paper, for the first time, considers training the coefficients for optimally combining the historical solutions instead of  simple averaging of the solutions.   
2)  Different experimental comparison have been conducted showing that TWA has better generalisation than SWA for either head-stage or tail-stage of training.       

Weaknesses: 
(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.    
(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. 
(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.   
(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. 


clarity,_quality,_novelty_and_reproducibility: 
(1) In subsection 2.1,  I don't get why lower training loss by optimisation over to {alpha_i}_{i=1}^n can lead to better generalisation. What is the logic behind it? It is known that Adam can produce low training loss while does not generalize well. 
(2) No source code available for reproducibility.
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. 
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.
flag_for_ethics_review: NO.
recommendation: 5: marginally below the acceptance threshold

From: ICLR.cc/2023/Conference/Paper1466/Reviewer_zFxA
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
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. 
The experiments show the proposed TWA can improve the generalization better than the previous EMA and SWA.
strength_and_weaknesses: Strength
1. The idea is clear and easy to follow. And the proposed TWA proved to be helpful in practice.
2. TWA helps to reduce the training cost while keeping the comparable performance.

Weaknesses
1. 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).''
2. 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).
3. 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. 
clarity,_quality,_novelty_and_reproducibility: The idea of TWA is direct and easy, but it still requires a little ingenuity to make it work in the experiment. From this perspective, this paper is novel. 
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.
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.
flag_for_ethics_review: NO.
recommendation: 6: marginally above the acceptance threshold

From: ICLR.cc/2023/Conference/Paper1466/Reviewer_Tq6j
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
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).

The proposed TWA can outperform SWA to some extent. However, the presented results may have the issues of limited novelty, unfair comparison, etc.
strength_and_weaknesses: # Strength
* The proposed idea is simple and easy to understand.
* Some numerical results on different neural architectures are provided from various aspects to justify the performance gain of TWA.

# Weaknesses
The 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.

1. 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.,
    * taking the SWA model as the initial model and fine-tune the model with the same amount of training epochs as TWA.
    * sampling more models (from these additional training epochs) and performing SWA. 

Without including these baselines, it is hard to justify the benefits of TMA, e.g., in Table 1, Table 2, Figure 5.

2. 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.
3. 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.
4. 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.
5. The statements in Section 3 are not well-supported, neither theoretically nor empirically.

### Post-rebuttal
The reviewer thanks the authors for providing the feedback. Most concerns were addressed, and the reviewer would raise the score to 6.
clarity,_quality,_novelty_and_reproducibility: The manuscript is generally well-structured; however, more experimental details need to be included.

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.
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.
flag_for_ethics_review: NO.
recommendation: 6: marginally above the acceptance threshold