Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: 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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
review_title: null

Meta-review: 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
Decision: Accept: poster