Switch EMA: A Free Lunch for Better Flatness and Sharpness

Authors: Siyuan Li, Juanxi Tian, Zicheng Liu, Ge Wang, Zedong Wang, Weiyang Jin, Di Wu, Tao Lin, Cheng Tan, Yang Liu, Baigui Sun, Stan Z. Li
Primary area: optimization
Keywords: Regularization, Exponential Moving Average, Weight Averaging, Optimization, Optimizer

Abstract:
Exponential Moving Average (EMA) is a widely used weight averaging (WA) regularization to learn flat optima for better generalizations without extra cost in deep neural network (DNN) optimization. Despite achieving better flatness, existing WA methods might fall into worse final performances or require extra test-time computations. This work unveils the full potential of EMA with a single line of modification, i.e., switching the EMA parameters to the original model after each epoch, dubbed as Switch EMA (SEMA). From both theoretical and empirical aspects, we demonstrate that SEMA can help DNNs to reach generalization optima that better trade-off between flatness and sharpness.
To verify the effectiveness of SEMA, we conduct comparison experiments with discriminative, generative, and regression tasks on vision and language datasets, including image classification, self-supervised learning, object detection and segmentation, image generation, video prediction, attribute regression, and language modeling. Comprehensive results with popular optimizers and networks show that SEMA is a free lunch for DNN training by improving performances and boosting convergence speeds.

Decision: reject
Reviewer scores: 5, 8, 6, 3

Reviews and discussion:
<thread>
creation date:
1729809713532

thread_type:
official_review

<thread_object>
creator: reviewer_cagp

<content>
summary:
This paper presents a new weight averaging approach SEMA that seeks to optimize both flatness and sharpness of optima when training NNs. SEMA relies on a simple change to EMA: the EMA parameters are switched to the original model parameters after K epochs. The authors provide theoretical and empirical validation of SEMA: SEMA is evaluated across a very large number of tasks against other weight averaging techniques.

strengths:
* The proposed change in EMA is sensible and simple; the proposed intuitions (balancing flatness and sharpness) for the change are well motivated.
* The authors propose some theoretical motivation to SEMA, including variance analysis, convergence analysis and error bounds.
* The evaluation is extremely thorough and covers an extremely large number of tasks, architectures.

weaknesses:
* Despite all these positive comments that show that the execution of the work was well lead, I feel that the proposed change is relatively minor to justify an acceptance to ICLR. the change is in practice a single line change, and does not significantly change the mechanisms behind EMA. The improvement to the models is relatively modest, although this may justify using SEMA as an alternative to EMA to gain a little bit in performance. On this point, I feel that including some variance across training seeds could at least provide some more information to appreciate the improvement.

questions:
* Could the authors report variance over several seeds on some well chosen datasets to show that the improvements are at least statistically significant ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730100348072

thread_type:
official_review

<thread_object>
creator: reviewer_v8md

<content>
summary:
This paper proposes a novel optimization method called SEMA, which incorporates flatness and sharpness by switching between fast and slow models at the end of each training epoch. SEMA results in faster convergence, better performance, and smoother decision-making. Extensive experiments have been conducted to demonstrate the effectiveness of SEMA in improving model performance across various scenarios, along with comprehensive and well-founded theoretical proofs.

strengths:
1. The paper is very well written and it was easy for me to follow.
2. The paper focuses on the key issues of optimizers and creatively improves existing methods, presenting a novel approach.
3. The paper demonstrates the effectiveness of the method through extensive experiments and theoretical analysis.

weaknesses:
1. The legend in the upper right corner of each curve plot in Figure 2 is quite small, requiring significant magnification to be readable.
2. There is no clear explanation of the data distribution visualization formats used in Figure 4 and Figure A.1.

questions:
Why does the baseline in Figure 4 have a decision boundary represented by two lines? If the decision boundary is like this, achieving an accuracy of 0.8 seems somewhat unreasonable.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730357336392

thread_type:
official_review

<thread_object>
creator: reviewer_wryd

<content>
summary:
This paper proposes a simple yet powerful algorithm, switch exponential moving average (SEMA), to improve the stability, convergence, and generalization performance of optimization of deep neural networks using a family of stochastic gradient descent methods. SEMA is a very simple method that updates the training parameters with EMA parameters once per epoch for existing EMAs. The paper conceives the algorithm based on experimental observations but also theoretically shows that it improves optimization stability, convergence speed, and generalization error bounds. In addition, the paper evaluates SEMA on various tasks, domains, architecture, training algorithms, and optimization algorithm combinations using SEMA and confirms that it improves the baseline in all cases.

strengths:
- The paper proposes an practical algorithm that is widely applicable and stably improves the performance of optimization of neural networks using stochastic gradient descent.
- The paper investigates the theoretical properties have been well analyzed to investigate why a simple EMA improvement is effective.
- The effectiveness of the proposed algorithm is evaluated through extensive experiments.

weaknesses:
- Although the experiments provided in the paper are basically extensive, further verification of the smoother decision boundary by SEMA discussed in Sec. 3.2 would be beneficial to the reader. The paper mentions the performance against fine-grained discrimination in L268, but in fact, the main focus of the experiments is on coarse-grained evaluation. For example, the performance on FGVCAircraft [a], CUB [b], Stanford Cars [c], etc., which are often used to evaluate fine-grained classification, and the evaluation of adversarial robustness, as shown in [d], would be useful.
- Since the standard deviation of the experimental results is not given, it is not possible to evaluate how stable SEMA actually is.

[a] Maji, Subhransu, et al. "Fine-grained visual classification of aircraft." arXiv preprint arXiv:1306.5151 (2013).

[b] Caltech-UCSD Birds-200-2011 (CUB-200-2011) https://www.vision.caltech.edu/datasets/cub_200_2011/

[c] Krause, Jonathan, et al. "Collecting a large-scale dataset of fine-grained cars." (2013).

[d] Zhang, Yihao, et al. "On the duality between sharpness-aware minimization and adversarial training.", ICML 2024.

questions:
Nothing to ask. Please see the weaknesses section and address the concerns.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730615848902

thread_type:
official_review

<thread_object>
creator: reviewer_agnx

<content>
summary:
The paper proposes SEMA methods, which apply exponential moving average (EMA) periodically at specific intervals to improve model "sharpness" and "flatness." The authors provide theoretical analysis and present extensive experimental results across various domains, including image classification, self-supervised learning, object detection, video prediction, and language modeling.

strengths:
The authors conduct comprehensive experiments in diverse areas such as image classification, self-supervised learning, object detection, video prediction, and language modeling, highlighting the advantages of the proposed method.

weaknesses:
Although the proposed method demonstrates superior performance, I believe this paper is not yet ready for publication for the following several reasons.

## **Sharpness vs. Flatness?**
The concepts of sharpness and flatness are unclear and deviate from conventional terminology. The authors argue that sharpness measures the depth of local minima while flatness assesses their width. However, I believe that when evaluating local minima in terms of depth and width, it is essential to consider the relative scale between these factors, such as the eigenvalues of the Hessian of the loss function. 

## **Concerns over Statistical Validity**
The results lack standard deviation values, and with only a maximum of three trials conducted, it is challenging to determine the statistical significance of the findings. This limited result makes it difficult to trust the robustness of the conclusions.

## **Some Definitions of Mathematical Terminologoies are Missing**
While I understand that these may seem like minor details, I believe they are crucial for the reader's comprehension and should be addressed.
- In eq(1), the definitions of $\Theta^{\text{OPT}}$ and $\Theta^{\text{EMA}}$ are missing
- In Proposition 1, the exact definition of variance is missing e.g. $V^{(t)}_{\text{SGD}}:= \mathbb{V}(\Theta_t^{\text{SGD}})$
- In Proposition 1, small $\eta$ condition is missing. If $\frac{2}{\eta}$ is smaller than the top eigenvalue of $A$, it seems that Banach's fixed point theorem cannot be applied.
- Typos in lines 1082-1083, 1122-1123, 1133: subscripts on $\mathcal{E}$ are missing

I hope the author clarifies these points in the next revision.

questions:
- Can the author provide the mathematical definition for sharpness and flatness considered in this paper?
- Why does (7) implies $ E_{\text{SEMA}}<E_{\text{EMA}}<E_{\text{SGD}}$ ?
 It seems that (7) only provides upper bounds for the three terms, and I believe that comparing only upper bounds does not support the conclusion that $E_{\text{SEMA}}<E_{\text{EMA}}<E_{\text{SGD}}$.
- In (6), which terms are hidden in the $\propto$ notation?If the hidden constant is inconsistent across iterations $t$,I believe that Proposition 2 does not guarantee convergence, even for convex functions, as the hidden constant could be too large for certain iterates.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732222817115

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely thank all reviewers for their efforts and detailed reviews. We deeply appreciate the time and effort you have invested in providing constructive feedback and thoughtful questions that will significantly improve our manuscript. We wanted to respond to SEMA's core motivation and some of the reviewer's underlying concerns:

**1. Challenges in Optimizer Improvement**

While most optimization algorithms aim to enhance the performance and generalization capabilities of deep neural networks, achieving plug-and-play functionality, performance gains, zero additional computational overhead, and generality simultaneously remains an intractable challenge. Many sophisticated and seemingly novel approaches attempt to overcome the limitations imposed by algorithmic foundations, yet they often fail to achieve a comprehensive balance. Certain improvements may compromise performance gains to minimize computational overhead, while others introduce negligible additional computations but lack general applicability. Consequently, existing optimizer enhancements can only satisfy a subset of the desired advantages, lacking generality. Therefore, we aim to advance DNN optimization from an alternative perspective, striving to achieve an optimal trade-off among these four desiderata to the greatest extent possible.

**2. The Concept of a "Free Lunch" in DNN Optimization**

It is crucial to emphasize that our proposed SEMA is neither a new optimizer nor a simple weighted average regularization. Instead, it is a simple, novel, plug-and-play EMA-based regularization technique. Unlike similar optimizers (e.g., SAM), wrappers (e.g., LookAhead), and regularizers (e.g., EMA and SWA), SEMA's advantage lies in its ability to circumvent the aforementioned conflicting objectives of DNN optimization improvements – the trade-off between performance gains, additional training costs, plug-and-play capability, and generalization – in the simplest and most innovative manner possible. The only cost is that, in certain scenarios, it may lead to slightly diminished performance gains, which can be compensated by adjusting two hyper-parameters (switching interval and momentum). These points underscore SEMA's achievement of a comprehensive balance, rendering it a veritable "free lunch" solution.

**3. The Latest Revision**

In addition to addressing these key points, we have taken the following steps to improve our manuscript. The key points of revision are highlighted in $\color{magenta}magenta$ color.

- **Theoretical Analysis and Visualization**: We have thoroughly reviewed and optimized the theoretical sections of the manuscript to enhance clarity and comprehensibility. We have also incorporated additional visualizations to better illustrate the concepts and results, thereby improving the overall readability of the manuscript.

- **Detailed Experimental Results**: We have expanded the experimental results section to include additional metrics that are of particular interest to the reviewers, such as standard deviations. This provides a more comprehensive understanding of the variability and robustness of our findings.

We are committed to further enhancing the manuscript in the next revision by:

- **Re-evaluating and Optimizing Theoretical Sections**: We will conduct a thorough re-evaluation of the theoretical sections to ensure they are as clear and robust as possible. This will involve refining the explanations and ensuring that the theoretical underpinnings are well-supported and easy to follow.

- **Enhancing Visualizations**: We will add more visual aids to the manuscript to make the concepts and results more accessible. This includes creating additional figures, charts, and diagrams that will help readers better understand the key points and findings.

- **Incorporating Detailed Statistical Analysis**: We will provide a more detailed statistical analysis of the experimental results, including the inclusion of standard deviations and other relevant metrics that are of particular interest to the reviewers. This will give a more comprehensive view of the variability and robustness of our findings.

We believe these improvements would address the concerns raised by the reviewers and significantly enhance the quality and impact of our manuscript. Since the discussion period only lasts until November 26 and we are approaching this deadline, we would like to discuss these with you during this time and would be happy to provide more information based on your further feedback. If you are satisfied with our responses, please consider updating your score. We are looking forward to hearing back from you!

Best regards,

Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732811474361

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Encouraging Final Check and Feedback

comment:
Dear Esteemed Reviewers,

We hope this message finds you well. We are writing to express our profound gratitude for the invaluable feedback and thoughtful discussions we have engaged in over the past few weeks. Your insights have been pivotal in significantly elevating the quality of our manuscript, and the improvements we have made are substantial. We believe the current version of the manuscript now stands as a testament to the collaborative effort and dedication we have all invested.

We understand that we are currently unable to submit further revisions, but we want to assure you that we remain fully committed to considering any additional suggestions or concerns you may raise. Should you find the revised manuscript satisfactory, we kindly request that you consider increasing the score accordingly. We are also eager to continue our dialogue and are open to further discussions to address any lingering doubts or questions you might have.

Moreover, we wish to emphasize that we are willing to make further refinements to the manuscript based on any subsequent discussions. Your continued engagement and feedback are of utmost importance to us, and we are dedicated to ensuring that our work meets the highest standards of quality and relevance.

Your participation and the time and effort you have invested in reviewing our work are deeply appreciated. We are hopeful that this revised manuscript will gain the recognition it deserves within the community and contribute meaningfully to the field.

Thank you once again for your unwavering support and constructive feedback. We look forward to the possibility of further collaboration and discussion, and we remain committed to refining our work to meet your expectations.

Warmest regards,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents the switch exponential moving average (SEMA) to better train deep-net models. The high-level idea is to switch between using the original model parameter vs. the SEMA parameter at a selected interval. The paper demonstrates experiments on a large set of domains. Overall, the reviewers value the simplicity of the algorithm and the comprehensiveness of the experiment. However, there were concerns regarding the writing and the significance of the results. Two of the reviewers found the paper to be below the acceptance bar after the discussion period. The main concern is with the limited novelty, pointed out by AGNx, and Cagp, of the approach, the relationship to the loss of landscape geometry, and the statistical significance of the result. 

The AC is weighted toward the novelty less, however, do believe there is a slight gap between the motivation of the proposed method and the loss of landscape geometry. Furthermore, the authors' response regarding statistical significance is not adequate for the AC. It is unclear how the authors concluded significance, i.e., what statistical test was performed. Additionally, based on the AC's experience,  I have never gotten a zero standard deviation over multiple experiments with different random seeds (or even under the same seed, under a distributed setting). The experiments would be more convincing if state-of-the-art could be achieved on the benchmarks. For example, the COCO experiment is far from the state-of-the-art (SOTA), since the optimization algorithm can easily be used with the latest vision models, it would be better to show the performance gain to push the SOTA. Lastly, just a nitpicking from the AC, when presenting the gains in tables, it's better to report the difference from the **best compared** method rather than the basic model.