Lookbehind Optimizer: k steps back, 1 step forward

We thank the reviewer for their comments. We believe there are a few misconceptions that we address below:

“The idea of multi-step SAM is not new and has been explored before”

The primary aim of our work is not to introduce the concept of multiple ascent step SAM, but rather to make it work well in practice. By enhancing SAM’s maximization and minimization steps, our approach demonstrated significant performance improvements over vanilla SAM and ASAM, as well as additional baseline methods.

“Although the interpolation step makes this work different from the previous work, I think that is a marginal contribution.”

Beyond (gradient combination and) linear interpolation, we also introduce a mechanism to analytically control $\alpha$ during training (Section 6). By considering the angle between the first gathered gradient and the final update direction, we successfully remove the need to finetune $\alpha$, while maintaining performance. We show that our method leads to improved performance in a myriad of tasks: generalization performance, model robustness, and lifelong learning. We believe these contributions are substantial rather than marginal.

“The complete ImageNet should be in the numerical studies”

We want to clarify that our experiments used ImageNet-1k with an image size of 224x224, following common practice in the literature. We added this clarification in Appendix A.3 in the manuscript.

“Transformer-based models need to be added”

SAM has been successfully applied to Transformers [1] and we should expect Lookbehind to further improve SAM’s performance, as presented throughout our paper. While we understand the current desire to use Transformers in the community, we do not think this is a critical issue with our manuscript that warrants a reason for rejection.

[1] Chen et al. When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations. ICLR 2022.

We thank the reviewer for their constructive feedback and concrete suggestions to better articulate our motivation and contribution. Below we explain how we address the reviewer’s concerns.

“The paper ought to mention SAM algorithm in their title to give a clear view, such as Lookbehind SAM Optimizer or Lookbehind SAM.”

Thank you for bringing attention to this oversight. Initially, our choice of title was influenced by the parallel with Lookahead, aiming for consistency in titles. However, we acknowledge the merit of including SAM explicitly for a clearer representation of the paper's content. In response to this valuable feedback, we have updated the title to "Lookbehind-SAM: k steps back, 1 step forward." This modification aims to maintain brevity while ensuring a direct connection with the SAM algorithm. We appreciate the suggestion and believe this title better conveys the essence of our work.

“The authors claim that the proposed method could reduce the variance derived from the multiple ascent steps. Why we need to reduce such variance is not clearly demonstrated.”

We thank the reviewer for this opportunity to clarify this important point. In the context of Multistep-SAM, the objective is to enhance the solution of the maximization part of SAM’s objective. However, the substantial departure of the gradient collected far away from the original point can adversely impact the original minimization objective, resulting in a suboptimal loss-sharpness trade-off, as illustrated in Figure 1.

Lookbehind addresses this issue by collecting and combining gradients along the ascent step, effectively reducing the variance of these gradients. This reduction in variance contributes to an improved loss-sharpness trade-off, ultimately leading to enhanced performance. We have incorporated additional explanations and clarifications in Section 3 of the manuscript to further elucidate this critical aspect. We hope this provides a clearer understanding of the motivation behind our approach.

“A single parameter update requires performing multiple backwards propagation to calculate the ascent gradient. In other words, the utility efficiency of training samples is quite low”

We agree with the reviewer’s concern about the increased computational overhead associated with multiple backwards propagation for a single parameter update, which in fact applies to any Multistep Sam method. Note that Lookbehind entails the same number of gradient computations as Multistep-SAM/ASAM, yet consistently surpasses the latter in performance. Through this lens, although the utility efficiency of training samples is the same for both methods, the utility efficacy is higher with Lookbehind.

“the authors have not trained the models to achieve comparable results when using vanilla SAM in their baseline on Cifar dataset.”

We appreciate your concern about the strength of the baselines and have taken steps to address it. We conducted additional experiments in line with the training setup from ASAM's paper [1], by applying a cosine learning rate scheduler and label smoothing to WRN-28-10 trained on CIFAR-100. The results, detailed in Appendix A.8 (Table 8), demonstrate that Lookbehind achieves an error rate of $16.28$ and $16.00$ for SAM and ASAM, respectively. These performances convincingly outperform the error rates of $16.58$ and $16.32$ reported for vanilla SAM and ASAM in [1]. We hope these additional results provide a more robust basis for comparison.

[1] Kwon et al. ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks. ICML 2021.

“Many advanced SAM variants have not mentioned in their experiments”

Thank you for your suggestion. It’s important to note that Lookbehind is a versatile wrapper that can be used with any SAM variant. While our focus has been on the original SAM algorithm and the adaptive variant (ASAM), our results indicate that Lookbehind can be applied to further improve the loss-sharpness trade-off and generalization performance across a range of SAM variants.

Thank you for your response and the consideration of the score adjustment. We kindly note that the reviewer’s original score of 3 is still reflected on Openreview, and we appreciate your attention to this matter.

Regarding the concern about computational overhead, we recognize its significance, but it’s essential to clarify that this issue is inherent to any multiple ascent-step SAM algorithm, not specific to our method. Our paper contributes to the active research on multiple ascent-step SAM, focusing on addressing the performance challenge arising from a poor sharpness-loss tradeoff. In response to the reviewer’s feedback, we’ve added a note in the conclusion, underscoring the importance of exploring the computational efficiency of multistep SAM.

While testing Lookbehind with a broader range of SAM variants is an interesting suggestion, it’s crucial to note that Lookbehind is not just a new SAM variant. Instead, it serves as a versatile wrapper applicable to any sharpness-aware minimization method, aiming to enhance their sharpness-loss tradeoff. We firmly believe our results highlight the effectiveness of Lookbehind in this context.

We thank the reviewer for their accurate summary of our work. We also thank the reviewer for their questions and concrete suggestions to clarify our contribution and to improve both the content and presentation.

“Can the authors compare Lookbehind against averaging the ascent gradients baseline ?”

Thank you for this suggestion. We have added additional comparisons between Lookbehind and Multistep-SAM/ASAM with averaging the ascent gradients. Results are shown in Appendix A7 (Table 7) of the updated manuscript. We observe that despite showing some improvements over Multistep-SAM/ASAM, both variants of Lookbehind (with and without adaptive $\alpha$) outperform these baselines.

“More detailed descriptions are needed for Figure 2 and Figure 11.”

Thank you for pointing this out. We added a detailed description of Figure 11 and presented the key differences between Figures 2 and 11 in Appendix A.2.

“The decay term can be omitted for simplification.”

Thank you for this suggestion. We removed the decay term from all equations and kept the mention of using weight decay in Appendix A.3 where we describe the training details.

“While leveraging multiple ascent steps can improve over the original SAM/ASAM, a prior study [2] shows that the inner gradient ascent can be calculated periodically while maintaining similar performance to the conventional SAM (i.e. it is redundant to compute the ascent gradient at every step). Can the author elaborate more on this?”

Lookbehind is a general wrapper for any sharpness-aware minimization method. Therefore, more efficient SAM variants such as re-using the inner gradient ascent steps are orthogonal to our work and may be used together to improve training efficiency.

“Can the authors compare Lookbehind against SAM/ASAM at different training budgets?”

We understand and share the reviewer’s concern. Indeed, we already briefly discuss this in Appendix A.1.1, where we study the average number of epochs at which SAM and ASAM reach top validation performance. Our findings suggest that increasing training time would not be a viable way to further improve performance in our setup.

In response to the reviewer’s recommendation, we conducted additional experiments by replicating the ASAM setup. We trained a WRN-28-10 model on CIFAR-100 using SAM and ASAM for an extended period of 400 epochs instead of 200. The results, detailed in Appendix A.8, Table 8, show that the 400-epoch models achieved similar generalization performance compared to the original 200 epochs. Specifically, SAM and ASAM attained $83.37_{\pm .12}$ and $83.67_{\pm .12}$ validation accuracy at an average epoch of $194.8_{\pm 2.78}$ and $202.8_{\pm 8.23}$, respectively. This supports our earlier observation that the models’ generalization performance was already saturated by the original 200 epochs. We hope this additional analysis addresses the reviewer’s concern.

Thank you for your comments. We are grateful for your positive feedback and overall excitement regarding our work. We address specific questions below:

“Can the fast weights (updated in line 8 of Algorithm 1) be approximately updated using any standard optimization algorithm like SGD or Adam?”

Thank you for this important question. Yes, the fast weights updated in line 8 of Algorithm 1 can be approximately updated using standard optimization algorithms like SGD or Adam. Lookbehind is a wrapper for sharpness-aware minimization methods, and these, in turn, are a wrapper for standard optimization algorithms like SGD or Adam. Although SAM and ASAM have been primarily used with SGD, both methods have also been successfully combined with Adam, showing improvements over vanilla Adam [1].

[1] Kwon et al. ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks. ICML 2021.

“Why not conduct an analysis of the sensitivity of Lookbehind to the step size of the fast weights?”

We appreciate your suggestion regarding the analysis of Lookbehind's sensitivity to the step size of the fast weights. However, we maintain a consistent inner learning rate for all methods. In the context of our specific method, the influence of the inner step size can be managed through the parameter $\alpha$, which can be considered as compensation for variations in the learning rate applied to the fast weights.

“What is the relation between Lookbehind and Anderson acceleration?”

Thank you for bringing up this interesting question. As far as we understand, the goals of Lookbehind and Anderson acceleration appear to be distinct. Anderson acceleration is typically employed as a convergence acceleration method, whereas Lookbehind serves as a variance-reduction method, enhancing the efficiency of the maximization and minimization components of SAM’s objective. However, we acknowledge that we were not familiar with Anderson acceleration before this discussion and are open to further conversations on this topic.

Minor fixes: We changed the text where appropriate to address your comments - thank you for your corrections.