Title: Lookaround Optimizer: $k$ steps around, 1 step average
Venue: NeurIPS 2023
Accepted: yes
Keywords: Deep Learning

Computer Vision

Mode Connectivity

Weight Average
Abstract: Weight Average (WA) is an active research topic due to its simplicity in ensembling deep networks and the effectiveness in promoting generalization. Existing weight average approaches, however, are often carried out along only one training trajectory in a post-hoc manner (i.e., the weights are averaged after the entire training process is finished), which significantly degrades the diversity between networks and thus impairs the effectiveness. In this paper, inspired by weight average, we propose Lookaround, a straightforward yet effective SGD-based optimizer leading to flatter minima with better generalization. Specifically, Lookaround iterates two steps during the whole training period: the around step and the average step. In each iteration, 1) the around step starts from a common point and trains multiple networks simultaneously, each on transformed data by a different data augmentation, and 2) the average step averages these trained networks to get the averaged network, which serves as the starting point for the next iteration. The around step improves the functionality diversity while the average step guarantees the weight locality of these networks during the whole training, which is essential for WA to work. We theoretically explain the superiority of Lookaround by convergence analysis, and make extensive experiments to evaluate Lookaround on popular benchmarks including CIFAR and ImageNet with both CNNs and ViTs, demonstrating clear superiority over state-of-the-arts. Our code is available at https://github.com/Ardcy/Lookaround.

Lookaround Optimizer: k steps around, 1 step average

Abstract

Weight Average (WA) is an active research topic due to its simplicity in ensembling deep networks and the effectiveness in promoting generalization. Existing weight average approaches, however, are often carried out along only one training trajectory in a post-hoc manner ( i.e. , the weights are averaged after the entire training process is finished), which significantly degrades the diversity between networks and thus impairs the effectiveness. In this paper, inspired by weight average, we propose Lookaround , a straightforward yet effective SGD-based optimizer leading to flatter minima with better generalization. Specifically, Lookaround iterates two steps during the whole training period: the around step and the average step. In each iteration, 1) the around step starts from a common point and trains multiple networks simultaneously, each on transformed data by a different data augmentation, and 2) the average step averages these trained networks to get the averaged network, which serves as the starting point for the next iteration. The around step improves the functionality diversity while the average step guarantees the weight locality of these networks during the whole training, which is essential for WA to work. We theoretically explain the superiority of Lookaround by convergence analysis, and make extensive experiments to evaluate Lookaround on popular benchmarks including CIFAR and ImageNet with both CNNs and ViTs, demonstrating clear superiority over state-of-the-arts. Our code is available at https://github.com/Ardcy/Lookaround .

1 Introduction

Recent research on the geometry of loss landscapes in deep neural networks has demonstrated Linear Mode Connectivity (LMC): two neural networks, if trained similarly on the same data starting from some common initialization, are linearly connected to each other across a path with near-constant loss [33, 9]. LMC reveals that neural network loss minima are not isolated points in the parameter space, but essentially forms a connected manifold [8, 35]. It has recently been attracting increasing attention from research communities, attributed to its great potential into motivating new tools to more efficiently reuse trained networks. A simple but effective strategy is Weight Average (WA), which directly averages the network weights of multiple trained [46] or sampled [19, 3] networks along the training trajectories for flatter minima and thus better generalization.

WAhas become a popular approach in various fields for its efficiency in parameters and effectiveness in performance. For example, Izmailov et al. [19] show that Stochastic Weight Averaging (SWA) of multiple points along the trajectory of SGD, with a cyclical or constant learning rate, leads to better generalization than conventional training. Following that, Cha et al. [3] propose Stochastic Weight Averaging Densely (SWAD), i.e. , averaging weights for every iteration in a region of low validation loss , to capture flatter minima and thus achieve superior out-of-domain generalizability. Recently, Wortsman et al. [46] propose averaging weights of multiple fine-tuned models with various hyperparameters, coined Model Soups , to improve accuracy without increasing inference time.

Generally speaking, WA guides the model points around the loss basin into the interior of the loss basin to find a flatter solution, which is shown to approximate ensembling the predictions [19].

Albeit striking results achieved in some cases, WAeasily breaks down owing to the diversitylocality conflicts. In one end, when WA is carried out within only one optimization trajectory after the training converges (as SWA does), the function diversity of the sampled checkpoints whose weights are averaged is limited, impair-

𝜃

!"!#

𝜃

$

𝜃

%

𝜃

&

𝜃

'()"

𝜃

$

𝜃

%

𝜃

&

ing the benefits of WA in ensembling. In the other end, if WA is conducted between models trained independently, WA can result in a completely invalid model due to the large barrier between different local minima (violating the locality principle that solutions should be in the same loss basin and closed to each other in parameter space). Prior WA approaches are either limited in function diversity or easily violating the locality principle, hindering its effectiveness and application in practice.

In this work, we propose Lookaround , a straightforward yet effective SGD-based optimizer to balance the both. Unlike prior works where WA often carried out along only one training trajectory in a post-hoc manner, Lookaround adopts an iterative weight average strategy throughout the whole training period to constantly make the diversity-locality compromise. Specifically, in each iteration, Lookaround consists of two steps: the around step and the average step, as shown in Figure 1. The around step starts from a common point and trains multiple networks simultaneously, each on transformed data by a different data augmentation, resulting in higher function diversity between trained networks. The average step averages these trained networks to get the averaged network, which serves as the starting point for the next iteration, guarantees the weight locality of these networks during the whole training. Lookaround iteratively repeats the two steps constantly, leading the averaged model to a flatter minima and thus enjoying better generalization. We theoretically explain the superiority of Lookaround by expected risk and convergence analysis, and make extensive experiments to evaluate Lookaround on popular benchmarks including CIFAR and ImageNet with both CNNs and ViTs, demonstrating clear superiority over state-of-the-arts.

The contributions of this work can be summarized as follows.

We propose Lookaround , a SGD-based optimizer that enjoys a diversity-in-locality exploration strategy. To seek flat minima within the loss basins, Lookaround iteratively averages the trained networks starting from one common checkpoint with various data augmentation.

We theoretically analyze the convergence of different competitive optimizers in the quadratic case, and prove that our Lookaround show lower limiting expected risk and faster convergence.

Extensive experiments conducted on the CIFAR and ImageNet benchmarks demonstrate that various CNNs and ViTs equipped with the proposed Lookahead optimizer can yield results superior to the state-of-the-art counterparts, especially the multi-model ensemble methods.

2 Related Work

Mode Connectivity. Mode connectivity is an intriguing phenomenon in deep learning where independently trained neural networks can have their minima connected by a low-loss path on the energy loss landscape [12, 14]. Some early works [12, 8] find that different independently trained networks can be connected by curves or multiple polylines in low-loss basins. Based on this, some later works [35, 11] show that when training from pretrained weights, the model stays in the same basin of the loss landscape, and different instances of such models can be connected by a straight line. Such findings motivates weight averaging, which averages the weights of different models to obtain higher performance. Recently, weight averaging has been widely used in the fields of neural language processing [21, 32, 39] and multi-objective optimization [33]. We aim to use model connectivity to give our network access to lower positions in the loss basin. In this paper, we propose a simple yet effective method to demonstrate the practical significance of this study.

Weight Averaging. While weight averaging shares some similarities with the aggregation operations in Graph Neural Networks [24, 45, 23, 22], its primary application is distinct. In the context of

neural architectures, weight averaging typically represents an overall averaging of weights across different models with the same architecture, diverging from the node-based aggregation in GNNs. Historically, the general idea of weight averaging can trace its roots back to convex optimization [41, 38]. And with the breakthrough of loss landscape visualization technology [12, 28], many applications have applied this idea to neural networks [12, 49, 34]. The concept of weight averaging has been integrated into single-trajectory training [12, 19, 49] and has seen extensive application in distributed setups [29, 31, 50]. Specifically, [19] found that the models trained with SGD optimizer often fell on the edge of a flat basin. Based on that, they propose the SWA method, which uses one-time weight averaging of multiple network points after the entire training process is finished, leading the network down to lower locations in the basin. The later work [49] proposes an optimizer that continuously uses weight averaging to update model weight to improve performance and robustness. Averaging weight during training is more effective than the averaging after training. Unlike single-trajectory optimization strategy, the recent Model Soups method [46] is inspired by mode connectivity, which starts from a typical pretrained weight and averages final fine-tuned models to improve generalizability. These methods inspire us to employ weight averaging method with multi-data augmentations to improve convergence and generalizability.

Ensemble. The ensemble learning is a traditional technology that combines multiple model outputs to achieve better robustness and performance [6, 1, 2, 16, 27]. Ensemble methods usually need to train different models to reduce variance and improve the prediction effect. However, in recent years, some methods [18, 12] can obtain different model checkpoints to conduct the ensemble in a single trajectory of model training, which significantly reduces the training time. Note that these methods all require separate inference through each model, which adds to the calculation cost. In contrast, our method does not require additional inference calculations.

3 Method

In this section, we present our optimization method Lookaround and provide an in-depth analysis of its properties. In Section 3.1, we provide a detailed description of the Lookaround optimizer. In Section 3.2, we analyze the expected risk of Lookaround, then thoroughly study its convergence at various levels and compare it with the existing optimizer Lookahead [49] and SGD [40]. Intuitively, our proposed optimizer seeks flat minima along the training trajectory by iteratively looking for the diverse model candidates around the trained points, hence called Lookaround. The pseudo-code of Lookaround is provided in Algorithm 1.

3.1 Lookaround Optimizer

In this subsection, we introduce the details of Lookaround. In Lookaround, the training process is divided into

Algorithm 1 Lookaround Optimizer.

multiple intervals of length k , with each interval consisting of an "around step" and an "average step". Let ϕ 0 denote the initial weights. At t th round of Lookaround, the weights are updated from ϕ t ´ 1 to ϕ t according to the around step and the average step. The resulting weights ϕ t then serve as the starting point for the subsequent round.

Around Step. Given the optimizer A , the objective function L , and the stochastic mini-batch of data B from train dataset D . In the around step, d different models are independently trained under d

different data augmentations AUG ' t AUG 1 , . . . , AUG d u for k batch steps, as follows:

Thanks to the data augmentations, each model is trained in period k to a diverse location in the loss landscape scattered around ϕ t ´ 1 . The Around Step allows optimizing the region surrounding each model iterate, which helps the search of flatter minima. Moreover, note that as the d different models are trained independently, we can use parallel computing to speed up the whole training process.

Average Step. In the average step, we simply average θ t,i,k , the weights of each independently trained model to obtain ϕ t for the next around step as follows:

It has been observed that models with the same initialization and trained on different data augmentations exhibit similar loss basins [35], with the models scattered around the edges of these basins. Therefore, incorporating a simple averaging technique into the training process may facilitate the convergence of the models towards lower-loss regions. In Appendix A, we demonstrate how this technique effectively guides the model to the interior of the loss basin.

3.2 Theoretical Analysis

3.2.1 Noisy Quadratic Analysis

The quadratic noise function has been commonly adopted as an effective base model to analyze optimization algorithms [42, 49, 47, 25], where the noise incorporates stochasticity introduced by mini-batch sampling. In this subsection, we analyze the steady-state risk of our proposed optimizer on the quadratic noise function to gain insights into the performance of Lookaround, and compare the result with those obtained using SGD and Lookahead.

The quadratic noise model is defined as ˆ L p θ q ' 1 2 p θ ´ c q T A p θ ´ c q . Following Zhang et al. [49], we assume that the noise vector c ' N p θ ˚ , Σ q , both A and Σ are diagonal matrices, and that the optimal solution θ ˚ ' 0 . We denote a i and σ 2 i as the i -th elements on the diagonal of A and Σ , respectively. As the i -th element of the noise vector c satisfies E r c 2 i s ' p θ ˚ i q 2 ` σ 2 i , the expected loss of the iterates θ t can be written as follows:

where θ t,i represents the i th item of the parameter θ t . The limiting risk of SGD, Lookahead and Lookaround are compared in the following Proposition by unwrapping E r θ t s and V r θ t s in Equation 3.

Proposition 1 (Steady-state risk) . Let 0 ă γ ă 1 { L be the learning rate satisfying L ' max i a i . One can obtain that, in the noisy quadratic setup, the variance of the iterates obtained by SGD, Lookahead [49] and Lookaround converge to the following matrix:

I

ď

,

if

α

Pp

0

,

1

q

respectively, where α denotes the average weight factor of models with varying trajectory points, as described in [49].

Proposition 1 implies that the steady-state variance matrix of Lookaround is element-wise smaller than those of SGD and Lookahead if d ě 3 and α P r 1 { 2 , 1 q . Moreover, Proposition 1 shows that increasing the number of data augmentation methods d yields smaller limiting variance and thus lower expected loss (see equation 3), which guarantees good generalizability in real-world scenarios. The proof is deferred to Appendix B.

3.2.2 Convergence on Deterministic Quadratic function

In this section, we analyze the convergence of Lookaround and Lookahead in the noise-free quadratic setting, which is important for studying the convergence of optimization algorithms [13, 37, 44].

Convergence rate ρ characterizes the speed in quadratic function at which the independent variable converges to the optimal solution, satisfying || θ t ´ θ ˚ || ď ρ t || θ 0 ´ θ ˚ || . In order to calculate the convergence rate, we model the optimization process, and treating this function as a linear dynamical system allows us to calculate the convergence rate, as in [30]. The value of the convergence rate will be determined by the eigenvalues of the dynamic transition equation, and we leave the calculation of this part to show in the Appendix B.2.

Given the emphasized significance of the condition number in paper [37], we conducted extensive experiments under a series of conditional numbers.

We evaluated the convergence rate of Lookaround, Lookahead, and Classical Momentum (CM) under different condition numbers from 10 1 to 10 7 , which are shown in Figure 2. The blue dashed line represents the optimal convergence rate that an algorithm can achieve in the absence of any constraints. A value closer to the blue dashed line indicates a faster convergence rate. CM can achieve the optimal convergence rate on a critical condition number [30]. To the left of this critical point, there exist complex eigenvalues that correspond to oscillations. In this regime, CM exhibits a flat convergence rate, which is known as "under-damped" behavior [30]. We show that Lookaround and Lookahead are faster than CM in the under-damped regime. Lookaround converges faster and is more stable than Lookahead in the low condition num-

10 1

10 2

10 3

10 4

10 5

10 6

10 7

Condition Number

10 -1

10 -2

10 -3

10 -4

1 - Convergence Rate

ρ

Optimal

CM

Lookahead

Lookaround

ber case, and realizes comparable performance with Lookahead in the high condition number case. In subsequent experiments, we show that Lookaround achieves fast and stable convergence in training deep neural networks on popular benchmarks.

4 Experiments

In Section 4.1 and Section 4.2, we present experimental verification on different tasks, including both random initialization and finetuning tasks on both ViTs and CNNs architectures to validate the effectiveness of the Lookaround optimizer. In Section 4.4, we compare the proposed method with ensemble learning methods that require multiple models. In Section 4.5, we conduct ablation experiments on different components of the Lookaround method to explore their contributions. In Section 4.6, we analyze the parameter robustness of Lookaround and further investigate its performance in the micro-domain.

Within a single epoch, both the proposed Lookaround and the competitors undergo training on an identical times the data augmentations. With such a setup, we guarantee consistency in the data volume utilized by each method, thereby ensuring fair comparisons in terms of computation.

4.1 Random Initialization

Random initialization is a standard model initialization method in deep learning. In this subsection, we verify the performance of our algorithm on various networks with randomly initialized weights based on CIFAR and ImageNet datasets.

4.1.1 CIFAR 10 and CIFAR 100

We conduct our experiments on CIFAR10 [26] and CIFAR100 [26] datasets. Both CIFAR10 and CIFAR100 datasets have 60,000 images, 50,000 of which are used for training and 10,000 for validation. We use SGDM [36] as our baseline, and we compare our method with SWA [19], SWAD [3], Lookahead [49]. We validate different methods on multiple network architectures such as VGG19 [43], ResNet50 [17], ResNet101 [17], ResNet152 [17], ResNeXt50 [48] and all methods are

0

50

100

150

200

0

1

2

3

4

Train Loss

VGG19

SGDM

Lookahead

Lookaround

0

50

100

150

200

0

1

2

3

4

ResNet50

SGDM

Lookahead

Lookaround

0

50

100

150

200

0

1

2

3

4

ResNeXt50

SGDM

Lookahead

Lookaround

0

50

100

150

200

Epoch

40

50

60

70

Top-1 Accuracy(%)

SGDM

Lookahead

Lookaround

0

50

100

150

200

Epoch

40

50

60

70

80

SGDM

Lookahead

Lookaround

0

50

100

150

200

Epoch

40

50

60

70

80

SGDM

Lookahead

Lookaround

performed on well-validated parameters. For each epoch in all experiments, we use the same mixed data for training using three kinds of data augmentation: random horizontal flip, random vertical flip, and RandAugment [5]. This ensures that the comparison between different methods is fair. For the training process, we use a discount factor of size 0.2 to decay the initial learning rate 0.1 at 60, 120, 160 epochs. Please refer to Appendix C.1.1 for more specific details.

We present our results in Figure 3 and Table 1, which indicates that our method achieves stable performance improvement on different networks. Compared to the SGDM optimizer, both Lookaround and Lookahead demonstrate faster convergence rates and higher accuracy. However, what's noteworthy is that, with training losses similar to Lookahead, Lookaround not only achieves a higher accuracy but also has better stability. This suggests that within the same training duration, we have found a more optimal optimization path. Meanwhile, although the SWA and SWAD methods adopt the same optimization process as the SGDM, they achieve higher performance via multiple samplings in the training trajectory combined with an average step. In contrast to SWA and SWAD, the multi-averaging Lookaround can achieve consistent and superior performance across different networks.

4.1.2 ImageNet

The ImageNet dataset is widely used to test the model performance. It has a training set of 1.28 million images and a validation set of 50,000 images with 1,000 classes. We train the model

0

10

20

30

40

50

Epoch

0

1

2

3

Train Loss

ResNet50

SGDM

Lookahead

Lookaround

0

10

20

30

40

50

Epoch

65

70

75

80

85

Top-1 Accuracy(%)

ResNet50

SGDM

Lookahead

Lookaround

0

10

20

30

40

50

Epoch

0

1

2

Train Loss

ViT-B/16

Adam

Lookahead

Lookaround

0

10

20

30

40

50

Epoch

80

85

90

Top-1 Accuracy(%)

ViT-B/16

Adam

Lookahead

Lookaround

on the training set and observe the model performance on the validation set. Based on standard hyperparameter settings, we train for 100 epochs and use a multi-step scheduler to adjust the learning rate on 30, 60 and 90 epochs with a discount factor of size 0.1. We present our results in Table 2. We get similar conclusions on the ImageNet as on CIFAR: higher test set accuracy and stronger generalization performance, which coincides with the theoretical results in the previous section.

4.2 Finetuning

The transfer learning technique, which involves pretraining a model on a large related dataset and then finetuning on a smaller target dataset of interest, has become a widely used and powerful approach for creating high-quality models across various tasks.

To explore the scalability of our Lookaround method, we apply our method to scenarios that require pre-training, including CNN or ViT [7] architecture. Unlike the previous experiments, we adopt more appropriate hyperparameters for its training in the ViT experiments, including the Adam optimizer and cosine annealing scheduler structure. Our results are shown in Table 3 and Figure 4. Under the ResNet architecture, our proposed Lookaround method remains stable and performs better, while the SGDM method shows greater fluctuations.

In contrast, the Lookaround methods based on weight averaging do not produce such a phenomenon, which partly demonstrates the hypothesis that average weight can reduce the convergence variance. In our comparison experiments with the other optimizations, we found that the results of ViT-B/16 under different optimizers were better than those reported in the original paper [7]. This is mainly due to our use of more data augmentation during the training process, aiming for a fair comparison with Lookaround. This result also indicates that Lookaround can still achieve better performance under improved baseline conditions.

4.3 Compared with Sharpness-Aware Minimization

We draw a comparison between Lookaround and Sharpness-Aware Minimization (SAM) [10], an algorithm designed to improve generalization by steering model parameters towards flatter loss regions. 1 The comparative results are presented in Table 4.

The results seem to indicate that, under the default parameters of SAM, this method is more suitable for higher resolutions, while Lookaround achieves the best performance at lower resolutions. For the finetuning, neither SAM nor Lookaround consistently emerges as the top choice when used individually. However, when combined, they often enhance model performance significantly. Thus, for higher accuracy, we recommend using both the SAM and Lookaround optimizers concurrently. Both are orthogonal to each other. By combining the two methods, we can achieve superior performance.

4.4 Compared with Ensemble Method

Ensemble methods are traditionally used to further improve inference performance by combining the outputs of multiple different models. We compare our Lookaround method with classical Ensemble methods (Logit Ensemble and Snapshot Ensemble [18]), and the results are shown in Table 5.

The performance of the single model obtained by our training exceeds the performance of the ensemble under multiple models. Moreover, the model obtained by our method is not a simple superposition of Ensemble models. We take the single model obtained by us with the three models obtained by Logit Ensemble to Logit Ensemble, and get a more significant performance improvement. This indi-

cates that the model we obtained is orthogonal to the three models obtained by Logit Ensemble and can be combined to produce better performance.

4.5 Ablation Study

In this subsection, we investigate the impact of different components of Lookaround on its performance, including data Augmentation methods and weight averaging. As shown in Table 6,

1

2

5

10

20

50

100 200 5001000

Steps(k)

80.5

81.0

81.5

Top-1 Accuracy(%)

Random Initialization

0

25

50

75

100

125

150

175

200

Epoch

50

60

70

80

Top-1 accuracy(%)

Finetuning

SGDM

(Ours)Small-k(5)

(Ours)Large-k(391)

0

50

100

150

Iteration

45

50

55

60

65

Top-1 Acc (%)

H-trajectory

H-net

R-trajectory

R-net

M-trajectory

Mean-net

0

50

100

150

Iteration

60.0

62.5

65.0

67.5

70.0

72.5

Top-1 Acc (%)

lr=0.02

0

50

100

150

Iteration

72

73

74

75

Top-1 Acc (%)

lr=0.004

when we consider Data Augmentation (DA) or Weight Averaging (WA) independently, their improvement effects may be relatively small. However, when we combine them, the synergistic effect between them can lead to more significant overall improvement. Specifically, we found that the improvement effects of DA and WA show a nonlinear trend in the interaction case, indicating a certain synergistic effect between them. Under the interaction, Data augmentation guides the model into different loss basins, and Weight Averaging guides different models into lower loss basins, resulting in higher performance improvement and a qualitative leap.

4.6 Additional Analysis

Robustness of the Parameters. For the number of data augmentations, we train the model with different numbers of data augmentations, and the results are shown in Table 7. The results indicate that an increase in the number of data augmentations results in improved network performance. In

order to reduce the training time, only three data augmentation methods are selected in this paper for experimental verification at different levels. Note that the methods compared with ours use the same incremental data set, so the comparison is fair.

For different choices of k , the number of around steps, we select different intervals for full verification, and the results are shown in Figure 5 (Left). Let us look at Large-k (391) versus Small-k(5) in Figure 5 (Right) (391 is determined by the number of batches in a single epoch). We find that large steps (green line) can achieve higher accuracy and stability in the early stage of training, while small steps can achieve higher accuracy in the later stage. In future research, we may explore a learning rate scheduler suitable for Lookaround to help it achieve the best performance.

We visualize the training trajectory accuracy of the sub-models in Lookaround in Figure 6. Under different learning rates, the model, after averaging the weights, can maintain a good baseline and continuously improve the performance of the model under the low variance, and the sub-models also obtain more beneficial training effects due to the weight averaging.

5 Conclusion

In this paper, we present the Lookaround optimizer, an SGD-based optimizer that employs a diversityin-locality exploration strategy to find solutions with high generalization ability in the loss basin. We theoretically prove that Lookaround can obtain lower variance and faster convergence speed. Experimental results on CIFAR and ImageNet and across various network backbones show that, our proposed Lookaround significantly improves the training stability and achieves better generalization performance comparable to the single-trajectory WA optimizers and ensemble methods.

Limitation. Lookaround is limited by the additional cost of network searching using different data augmentation, resulting in a longer training time proportional to the the number of trained networks. Notably, Lookaround incurs the same inference time as other optimization methods. Considering the significantly superior performance, we believe the extra time consumption is worthwhile in many realworld scenarios where the training time is not so concerned as the inference time. Moreover, training time can be reduced due to the parallelizability of the "around step" mechanism in Lookaround.

References

Eric Bauer and Ron Kohavi. An empirical comparison of voting classification algorithms: Bagging, boosting, and variants. Machine Learning , 1999.

Leo Breiman. Bagging predictors. Machine Learning archive , 1996.

Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Conference on Neural Information Processing Systems , 2021.

Ekin D. Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V. Le. Randaugment: Practical automated data augmentation with a reduced search space. Conference on Neural Information Processing Systems , 2019.

Ekin D. Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V. Le. Randaugment: Practical automated data augmentation with a reduced search space. In IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops , 2020.

Thomas G. Dietterich. Ensemble methods in machine learning. multiple classifier systems , 2000.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, and S. Gelly. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2021.

Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred A. Hamprecht. Essentially no barriers in neural network energy landscape. International Conference on Machine Learning , 2018.

R. Entezari, H. Sedghi, O. Saukh, and B. Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. International Conference on Learning Representations , 2022.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. International Conference on Learning Representations , 2021.

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. International Conference on Machine Learning , 2019.

Timur Garipov, Pavel Izmailov, D. A. Podoprikhin, Dmitry Vetrov, and Andrew Gordon Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Conference on Neural Information Processing Systems , 2018.

Gabriel Goh. Why momentum really works. Distill , 2017.

A. Gotmare, N. S. Keskar, C. Xiong, and R. Socher. Using mode connectivity for loss landscape analysis. arXiv: Learning , 2018.

Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations , 2020.

Trevor Hastie, Jerome H. Friedman, and Robert Tibshirani. The Elements of Statistical Learning: Data Mining, Inference, and Prediction . Springer Series in Statistics. 2001.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE/CVF Computer Vision and Pattern Recognition Conference , 2016.

Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E. Hopcroft, and Kilian Q. Weinberger. Snapshot ensembles: Train 1, get M for free. In International Conference on Learning Representations , 2017.

Pavel Izmailov, D. A. Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. Conference on Uncertainty in Artificial Intelligence , 2018.

Samyak Jain, Sravanti Addepalli, Pawan Kumar Sahu, Priyam Dey, and R Venkatesh Babu. Dart: Diversifyaggregate-repeat training improves generalization of neural networks. In IEEE/CVF Computer Vision and Pattern Recognition Conference , 2023.

Sébastien Jean, Kyunghyun Cho, Roland Memisevic, and Yoshua Bengio. On using very large target vocabulary for neural machine translation. Annual Meeting of the Association for Computational Linguistics , 2014.

Yongcheng Jing. Efficient Representation Learning With Graph Neural Networks . PhD thesis, 2023.

Yongcheng Jing, Yiding Yang, Xinchao Wang, Mingli Song, and Dacheng Tao. Meta-aggregator: Learning to aggregate for 1-bit graph neural networks. In IEEE International Conference on Computer Vision , 2021.

Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations , 2016.

Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International Conference on Machine Learning , 2017.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Handbook of Systemic Autoimmune Diseases , 2009.

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Conference on Neural Information Processing Systems , 2017.

Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Conference on Neural Information Processing Systems , 2018.

Xiangru Lian, Ce Zhang, Huan Zhang, Cho-Jui Hsieh, Wei Zhang, and Ji Liu. Can decentralized algorithms outperform centralized algorithms? a case study for decentralized parallel stochastic gradient descent. In Advances in Neural Information Processing Systems , 2017.

James Lucas, Shengyang Sun, Richard Zemel, and Roger Grosse. Aggregated momentum: Stability through passive damping. In International Conference on Learning Representations , 2018.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-Efficient Learning of Deep Networks from Decentralized Data. In International Conference on Artificial Intelligence and Statistics , 2017.

Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing lstm language models. Learning , 2017.

Seyed Iman Mirzadeh, Mehrdad Farajtabar, Dilan Gorur, Razvan Pascanu, and Hassan Ghasemzadeh. Linear mode connectivity in multitask and continual learning. Learning , 2020.

Giung Nam, Hyungi Lee, Byeongho Heo, and Juho Lee. Improving ensemble distillation with weight avearging and diversifying perturbation. In International Conference on Machine Learning , 2022.

Behnam Neyshabur, Hanie Sedghi, and Chiyuan Zhang. What is being transferred in transfer learning. Conference on Neural Information Processing Systems , 2020.

Q. Ning. On the momentum term in gradient descent learning algorithms. Neural Networks , 1999.

Brendan O'donoghue and Emmanuel Candes. Adaptive restart for accelerated gradient schemes. Foundations of computational mathematics , 15:715-732, 2015.

Boris T. Polyak and Anatoli Juditsky. Acceleration of stochastic approximation by averaging. Siam Journal on Control and Optimization , 1992.

Yujia Qin, Cheng Qian, Jing Yi, Weize Chen, Yankai Lin, Xu Han, Zhiyuan Liu, Maosong Sun, and Jie Zhou. Exploring mode connectivity for pre-trained language models. In Conference on Empirical Methods in Natural Language Processing , 2022.

Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics , pages 400-407, 1951.

David Ruppert. Efficient estimations from a slowly convergent robbins-monro process. Technical report, Cornell University Operations Research and Industrial Engineering, 1988.

Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In International Conference on Machine Learning , 2013.

Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. IEEE/CVF Computer Vision and Pattern Recognition Conference , 2014.

Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International Conference on Machine Learning , pages 1139-1147, 2013.

Petar Veliˇ ckovi´ c, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations , 2018.

Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo Lopes, Ari S. Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning , 2022.

Yuhuai Wu, Mengye Ren, Renjie Liao, and Roger Grosse. Understanding short-horizon bias in stochastic meta-optimization. arXiv preprint arXiv:1803.02021 , 2018.

Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual transformations for deep neural networks. IEEE/CVF Computer Vision and Pattern Recognition Conference , 2016.

Michael R. Zhang, James Lucas, Jimmy Ba, and Geoffrey E. Hinton. Lookahead optimizer: k steps forward, 1 step back. Conference on Neural Information Processing Systems , 2019.

Tongtian Zhu, Fengxiang He, Lan Zhang, Zhengyang Niu, Mingli Song, and Dacheng Tao. Topology-aware generalization of decentralized SGD. In International Conference on Machine Learning , 2022.

Tongtian Zhu, Fengxiang He, Kaixuan Chen, Mingli Song, and Dacheng Tao. Decentralized SGD and average-direction SAM are asymptotically equivalent. In International Conference on Machine Learning , 2023.

A Exploration in the Loss Landscape

To demonstrate the effectiveness of Lookaround during training, we set up an experiment and visualize the loss landscape of the models under different data augmentations in Figure 7. By observing the loss landscape, we can gain a clearer understanding of the role played by the weight averaging at different stages during the training process.

Training Process and Parameter Settings. We first train a ResNet50 on the CIFAR100 dataset. The learning rate is initialized to 0.1 and decay at 60, 120, and 160 epochs using a MultiStepLR scheduler with a decay factor 0.2. The batch size is set to 128, we use stochastic gradient descent with momentum to optimize the model and use random crops and random vertical flips augmentation to enhance the training datasets. We use model checkpoints at epochs 50, 110, and 150 as our three pretrained models(V-network). The three pretrained models correspond to learning rates of 0.1, 0.02 and 0.004, respectively. Using these pretrained models as the starting point, we finetune the model with 1, 10, 100, 1000, 10000 iterations under the corresponding learning rate and the setting of random horizontal flipping (H-network) or RandAugment (R-network).

Visualization Method. We use the visualization method in Garipov et al. [12] to plot the loss landscape. In this method, the weights of the three models are flattened respectively as one-dimensional vectors w v , w h , w r , and then two orthogonal vectors are calculated between the three vectors as the X-axis direction and the Y-axis direction: u ' p w h ´ w v q , v ' p w r ´ w v , w h ´ w v q{} w h ´ w v } 2 ¨ p w h ´ w v q . Then the normalized vectors ˆ u ' u {} u } , ˆ v ' v {} v } foam an orthonormal basis in the plane contain w v , w h , w r . Then a point P with coordinates p x, y q in the plane would be given by P ' w v ` x ¨ ˆ u ` y ¨ ˆ v .

Discussion and Inspiring. Under different learning rates and different around steps k , Lookaround has the tendency to lead the model trained on different data augmentation to the near-constant loss manifold. In such circumstances, the "average step" can lead the model into the center of the loss basin to get a lower test loss. However, weight averaging does not necessarily work in all cases. For example, the network obtained after weight averaging gets a larger loss under a large learning rate with a large around step (e.g., lr ' 0 . 1 , k ' 10000 ). In this case, the model is located on a peak between different basins rather than in the center of a basin. Moreover, in the case of around step k ' 1 , the weight averaging also does not achieve better performance. Nevertheless, such extreme cases do not prevent weight averaging from being a useful tool to speed up the training process. The center of the basin in loss landscape, which requires multi-step gradient descent to reach, can be reached by only one weight averaging step. At other learning rates and around steps k , the models after weight averaging all result in a lower loss than the individual model. Such phenomena encourage us to explore more methods to find more optimal solutions in the loss landscape in the future.

B Steady-state and Convergence Analysis of Lookaround

We use quadratic functions to analyze the steady-state and convergence analysis of Lookaround. First, we present the proof of Proposition 2.

Proposition 2 (Steady-state risk) . Let 0 ă γ ă 1 { L be the learning rate satisfying L ' max i a i . One can obtain that in the noisy quadratic setup, the variances of SGD, Lookahead [49] and Lookaround converge to the following fixed matrix:

Learning Rate

lr

=0.1

around steps

k

=1

V-network

R-network

H-network

Network obtained through "average step"

around steps

k

=10

around steps

k

=100

around steps

k

=1000

around steps

k

=10000

Learning Rate

lr

=0.02

Learning Rate

lr

=0.004

0.8

1.2

1.6

2.0

0.8

1.2

1.6

2.0

0.8

1.2

1.6

2.0

respectively, where α denotes the average weight factor of models with varying trajectory points, as described in [49].

From Wu et al. [47], we can obtain the following dynamics of SGD with learning rate γ :

Lemma 1. The expectation and variance of lookaround have the following iterates:

ϕ t yields ϕ t ` 1 by performing an around step and an average step.

Proof. The expected iterate follows from SGD:

For the variance of ϕ t ` 1 , we can break it down into two parts as follows:

The covariance of the different models can be calculated in the following way:

After permuting and regrouping again, we can obtain the iterate with respect to the variance.

The proof is now complete.

Remark. From Equation 10, the expectation term for ϕ in Lookaround eventually converges to 0, as does Lookahead and SGD.

From Zhang et al. [49], we have the following analysis about the variance fixed point of lookahead with learning rate γ and weight factor α , which represents the average weight factor of models with different trajectory points in the Lookahead optimizer, which is generally (0, 1):

We now derive the fixed point of the variance, proceed with the proof of Proposition 1:

Proof.

According to Equation 12, we can deduce that

The proof is now complete.

B.1 Comparing the dynamics of Lookahead

We now proceed with the proof for the range of constraints variable α in Equation 9. When d ě 3 , and α P r 0 . 5 , 1 q , the Lookaround method can obtain a smaller variance than the Lookahead method:

Proof. Let B ' p I ´ γ A q k , due to 0 ă γ ă 1 { L, L ' max i a i , so we can have B ď I . Substituting the matrix B into the expressions for the variance fixed point relation of Lookaround and Lookahead:

Let the coefficient matrix be denoted as C, when d ě 3 , for each diagonal element C ii of C, we can scale the denominator of this expression as follows:

Then, we can derive the range of α by restricting the right-hand side expression as follows:

As 0 ď B ii ď 1 , we can multiply both sides of the inequality by the denominator, then rearrange and combine like terms to obtain the following form:

Skipping the detailed steps, we can obtain α ě 0 . 5 by solving the quadratic equation. Therefore, in the case where α P r 0 . 5 , 1 q and d ě 3 ( α ă 1 is subject to Lookahead's settings), the coefficient matrix C ď I , so the convergence speed of Lookaround is slower than Lookahead.

B.2 Deterministic quadratic convergence

Our method samples data from multiple data augmentation strategies, which can be analogously viewed as averaging the sampling of historical trajectories during the convergence analysis of quadratic functions. Thus, for weight averaging, we select model points that align with each point in the k -step trajectory. From this perspective, we examine and compare the convergence rates of Lookaround and Lookahead.

We first show the state transition equation for the classical momentum method in quadratic functions:

Here, v stands for the momentum term. We can generalize this to matrix form:

Thus, given the initial θ 0 , we can obtain the convergence rate with respect to θ by the maximum eigenvalue of the matrix. Referring to Zhang et al. [49] and Lucas et al. [30], we obtain the state transition matrix regarding our algorithm as follows:

where L, B and T denote the average weight matrix, the single-step transfer matrix and the position transformation matrix respectively:

After specifying the appropriate parameters and performing matrix multiplication to obtain the state transition matrix, the convergence rate ρ can be obtained by calculating the largest eigenvalue of the matrix. Note that since this linear dynamical system corresponds to k updates, we finally have to compute the k th root of the eigenvalues to recover the correct convergence bounds.

C Experimental Detail

C.1 Random Initialization

C.1.1 CIFAR10 and CIFAR100

Data augmentation details. For the CIFAR10 dataset, we use [RandomCrop + ˚ ] for data augmentation, and for the CIFAR100 dataset, we use [RandomCrop + ˚ + RandomRotation] for data augmentation. ˚ can be replaced by three different data augmentation methods of random horizontal flip, random vertical flip, or RandAugment [4].

Training details. For the CIFAR10 and CIFAR100 datasets, we have applied some common settings. The initial learning rate is set to 0.1, and the batch size is set to 128. Additionally, a warm-up phase of 1 epoch is implemented. Subsequently, different learning rate schedulers are used based on the specific dataset. For the CIFAR100 dataset, we utilize the MultiStepLR scheduler. The learning rate is decayed at the 60, 120, and 160 epochs, with a decay factor of 0.2. For the CIFAR10 dataset, we employ the CosineAnnealingLR learning rate scheduler. In comparison with other optimizers or optimization methods, we use the default settings of the other methods in the corresponding papers.

C.1.2 ImageNet.

For the ImageNet dataset, our data augmentation strategy involves [RandomResizedCrop + ˚ + RandomRotation]. Here, ˚ can be substituted by one of the three data augmentation methods: random horizontal flip, random vertical flip, or RandAugment [4]. We adopt the following training settings: an initial learning rate of 0.1 and a batch size of 256. The model undergoes training for a total of 100 epochs. Additionally, we make use of the MultiStepLR scheduler, which decays the learning rate by a factor of 0.1 at the 30th, 60th, and 90th epochs.

C.2 Finetuning

In this stage, all images are resized to 224 x 224 pixels to match the input size of the pretrained model. All models use the ImageNet-1k pretrained weights sourced from the PyTorch library. Throughout the finetuning process, we employed three data augmentation methods: random horizontal flip, RandAugment, and AutoAugment.

For the training of ResNet50, we performed a grid search as shown in Table 8 to determine the optimal learning rate and weight decay. Every training process lasted for 50 epochs with a batch size

set to 128. For the training of ViT-B/16, we used the Adam optimizer, starting with a learning rate of 0.00001. We set β 1 to 0.9, β 2 to 0.999, and applied a weight decay of 0.01. To optimize memory usage, the batch size was set to 64.

C.3 Compared with Sharpness-Aware Minimization

We adopt the default hyperparameters from the origin paper of SAM [10] and combine them with the optimal hyperparameters found during random initialization or finetuning to conduct comparative experiments between SAM and Lookaround.

C.4 Compared with Ensemble Method

We compare Lookaround with Logit Ensemble and Snapshot Ensemble [18]. In the setting of Logit Ensemble, we train multiple models separately using different data augmentation methods, and then average the outputs of these models for prediction. However, this approach requires more inference time. In the setting of Snapshot Ensemble, we use the CosineAnnealingWarmRestarts learning rate scheduler to collect four snapshots during the training process. Then, we average the outputs of these different snapshots for prediction. This approach also requires more inference time.

Besides, we conducted a full experimental validation of Logit Ensemble and SWA under a new dataset Stanford Cars, as shown in Table 9. Under the Stanford Cars dataset, Lookaround still performs better than Logit Ensemble. Additionally, adding SWA into the training process of Lookaround can further improve performance. Under the limit of single-model inference, Lookaround+SWA achieves a Top-1 accuracy of 91.04%, exceeding the optimal 88.11% (SWA with AutoAugment).

C.5 Ablation Study

In our ablation experiments, we evaluate the individual contributions of two components: Data Augmentation (DA) and Weight Averaging (WA) to the effectiveness of the Lookaround. Below are the specific settings for different ablation experiments.

For the experiments without both DA and WA, three models were trained independently, each utilizing a distinct data augmentation strategy, for a duration of 200 epochs. After training, we identified

and reported the highest accuracy achieved by these models on the test set. In the experiments that employed only DA, we trained a single model using all data, combining the three data augmentation strategies, without the use of weight averaging. For the experiments with only WA, three models were trained separately for 200 epochs each, employing different data augmentation strategies. Following this, we selected the model that showcased the best performance on the test set.

C.6 Additional Analysis

In the robustness experiments with the number of Data Augmentation methods, the six data augmentation methods are given as: RandomVerticalFlip, RandomHorizontalFlip, RandAugment, AutoAugment, RandomPerspective, RandomEqualize. All the Augmentation methods are from the PyTorch library. When more data augmentation methods are used, the training time will be correspondingly increased in our method. Therefore, in this paper's main experiments, we only select three data augmentation methods for comparison to reduce the time consumption.

D Contemporaneous Work

After completing our work, we discovered similar work at CVPR 2023 within the same year. Inspired by similar motivations, the DART [20] method proposed by Samyak Jain et al. utilizes a similar training structure that includes different data augmentation training of multiple sub-models, with weight averaging occurring at intervals.

Methodology. DART starts weight averaging until the second half training process and adopts a sparser weight averaging strategy. Meanwhile, the DART scheme uses the ERM [15] strategy to initialize the model. In contrast, Lookaround maintains a consistent and dense weight averaging strategy throughout its training, making it more user-friendly in real-world applications.

Theory. DART proves that the average while training is more robust from the perspective of noise feature. On the other hand, Lookaround proves that it can get lower expected loss under the setting of quadratic noise function.

Experiments. DART validates its effectiveness in domain generalization. Lookaround is tested in the scenarios of finetuning and training from scratch.

In conclusion, although sharing a similar idea, two methods differ substantially in theoretical foundations and experimental approaches. Both methodologies are complementary and furnish significant insights for the broader research community.

And we compared Lookaround with DART with the same data augmentation, the results are shown in Table 10. Our results were better than DART, the reason behind this may come from the fact that we used more frequent "average steps" during training and started "average steps" earlier in the cycle. More frequent "average steps"

Table 10: Comparison with the DART method with ResNet18 on CIFAR10 and CIFAR100. The results of DART : are taken from reference [20].

are more conducive to the training of Lookaround, as referenced in the main body of the paper.

E Relationship with Model Soups

Model Soups [46] is a framework for finetuning a common pretrained model using different hyperparameters and then averaging the weights of different finetuned models to improve model performance and generalization. In the finetuning settings, we conducted an in-depth comparison between Lookaround and Model Soups, with the related results presented in Table 11. The "Update BN" suggestion originated from the reviewer's feedback. Given that we employed a CNN architecture, the averaging of means and variances in the batch normalization (BN) layers of different models within Model Soups might lead to distortions, failing to accurately represent the statistical characteristics of the true hidden layer inputs. To address this, we passed the averaged model through

the entire training set using forward propagation to individually update the BN layer's running mean and running var.

The experimental results indicate that the Lookaround method, employing only three types of data augmentation, still remains comparable in performance to the Model Soups method, which utilizes 18 types of data augmentation. This further attests to Lookaround's generalization capability even under limited data augmentation conditions. Additionally, we observed that the Lookaround method is applicable to tasks with random initialization, while the Model Soups, due to its prerequisite of pretrained weights, cannot be employed in such scenarios.

F Computational Complexity and Comparison

The Lookaround optimizer can be viewed as maintaining multiple backups of model weights during the optimization process. Each backup is only trained under its corresponding data augmentation strategy (forward and backward propagation) and updates the parameters using weight averaging after multiple iterations.

In general, if the time for an iteration is defined as Ω , and the time to perform a weight average is ω , given that the dataset contains B batches, the standard time complexity of SGDM to complete an epoch is O p B Ω q . The time complexity of Lookaround is O p dB Ω ` Bω d q . Given that ω is much smaller than Ω , this can be approximated as O p d Ω B q .

In our experiments (corresponding to Table 1 to 5 in the paper), to establish fair comparisons in computational among different approaches, we adopt the same augmentations for the competitors. Specifically, within a single epoch, both the proposed Lookaround and the competitors undergo training on an identical d times the data augmentations. With such a setup, we guarantee consistency in the data volume utilized by each method, thereby ensuring fair comparisons in terms of computation.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents Lookaround, a novel optimizer for weight average ensembling (WA). Unlike existing approaches that perform weight averaging post-training, Lookaround adopts a two-step process throughout the training period. In each iteration, the \"around\" step trains multiple networks simultaneously on transformed data using different augmentations, while the \"average\" step combines these networks to obtain an averaged network as the starting point for the next iteration.The approach demonstrates clear superiority over state-of-the-art methods in extensive experiments on CIFAR and ImageNet datasets using both traditional CNNs and Vision Transformers. The paper provides theoretical justification and commits to open science by making the code publicly available.\n\nSTRENGTHS:\n1. `Innovative Approach`: The introduction of Lookaround as a straightforward and effective optimizer for weight average ensembling brings a novel perspective to the field. The two-step process during training enhances network diversity and preserves weight locality, addressing the limitations of post-hoc weight averaging approaches.\n\n2. `Theoretical Justification`: The paper offers strong theoretical support for the superiority of Lookaround through convergence analysis. Proposition 1 provides insights into the convergence variance of Lookaround in comparison to typical SGD and Lookahead optimizers. This rigorous theoretical analysis enhances the credibility and significance of the proposed method, demonstrating a solid foundation for its effectiveness and performance.\n\n3. `Extensive Experimental Validation`: The authors conduct extensive experiments on popular benchmarks, including CIFAR and ImageNet, with both traditional CNNs and Vision Transformers (ViTs). The clear superiority of Lookaround over state-of-the-art methods on these datasets demonstrates its effectiveness and applicability.\n\n4. `Commitment to Open Science`: The authors state their intention to make the code publicly available, fostering reproducibility and enabling further research in the field.\n\nWEAKNESSES:\n1. `Computational Complexity and Comparison`: The paper lacks a comprehensive discussion of the computational complexity introduced by the Lookaround optimizer. As Lookaround incorporates an additional around step during training, it is crucial to assess its computational requirements and compare them to existing optimization methods. A detailed analysis of the computational trade-offs, including runtime and memory usage, would provide a more comprehensive understanding of Lookaround's practical applicability and scalability.\n\n2. `Limited Experiments on Datasets Beyond Image Classification`: The paper primarily focuses on evaluating Lookaround on image classification tasks using datasets like CIFAR and ImageNet. However, to establish the versatility and effectiveness of Lookaround, it is important to explore its performance on datasets beyond image classification. Conducting experiments on diverse tasks such as language modeling or segmentation/detection tasks would demonstrate the generalizability of Lookaround across various domains and provide a more comprehensive evaluation of its performance.\n\nQUESTIONS:\n1. With the Lookaround optimizer requiring k around steps and 1 average step, I am unclear about how the epoch and iteration are defined in this setting. Are they counted based on the average step or the around step? Could you provide clarification on how the epoch and iteration are structured in Lookaround?\n\n2. Considering that each batch is augmented k times in Lookaround, I am curious to understand how this approach differs from simply applying `repeated augmentations`[A] to each batch. What distinguishes the effect of k augmentations and 1 step update in Lookaround versus performing each step update individually and then averaging all checkpoints? In other words, what advantages does Lookaround's approach provide in terms of diversity and ensemble performance compared to the alternative method of averaging checkpoints from individual updates?\n\n[A] Hoffer E, Ben-Nun T, Hubara I, et al. Augment your batch: Improving generalization through instance repetition[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020: 8129-8138.\n\nLIMITATIONS:\nAs discussed in the limitation section \"additional cost of network searching using different data augmentation, resulting in a longer training time proportional to the the number of trained networks\", the running time is a important concern.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

comments: SUMMARY:
This work provides a new optimizer, "Lookaround optimizer," which is built upon a previous proposed "Lookahead optimizer" [40]. By incorporating data augmentation, this work shows an improved convergence rate under low condition numbers. It also empirically shows some improvement in classification accuracy under several datasets.

STRENGTHS:
This work proposes a new optimizer and shows consistent (despite little) improvement in various settings. The supplementary material shows a solid derivation.

WEAKNESSES:
1. The major concern lies in the limited improvement over the "Lookahead optimizer." The improvement in Figure 2 and Table 1-3 all seems to be subtle. Since the proposed methods see more augmented data in each epoch than others, it is questionable whether the improvement actually comes from these additionally augmented data. Also, the idea itself is almost identical to the Lookahead optimizer, so the novelty is rather limited.

2. Minor mistakes: In table 1, row CIFAR 10, column ResNext50 Top5 results, the best results happen in Lookahead (99.96), not the proposed method (99.95).

QUESTIONS:
What is the d (number of data augmentation) used in the experiment? I cannot find it in the context.

LIMITATIONS:
The author addresses that the data augmentation process consumes extra time, which is likely to be the main limitation.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new optimization algorithm named Lookaround, which draws inspiration from the recent achievements of weight averaging (WA) techniques in deep learning. The proposed Lookaround optimizer looks around nearby points by performing multiple gradient computations for a given training input using different data augmentations and averages them to obtain better generalizing solutions.\n\nSTRENGTHS:\nThe theoretical analysis in this paper largely relies on Zhang et al. (2019), including noisy quadratic analysis and deterministic quadratic convergence. While it does not introduce a novel form of analysis, it is significant as it establishes a theoretical foundation within the existing framework.\n\nFor experiments, the baseline comprises commonly employed optimization techniques, such as SWA (Izmailov et al., 2018) and SAM (Foret et al., 2021), which are widely accepted as standard approaches in the field. In addition to convolutional neural networks of various scales, the authors also considered experiments on vision transformers.\n\nWEAKNESSES:\nExperimental issues:\n\n* __There are no error bars.__  \nWhile the authors answered \"yes\" for \"Error Bars\", the paper only provide a single value across the tables and figures. It is unclear how many experiments were conducted to derive those values. It would be preferable if the authors also included averages accompanied by standard deviations to provide a more comprehensive representation of the experimental results.\n\n* __The outcomes obtained from the ImageNet experiments appear to be strange.__  \nDespite the authors' efforts to demonstrate the scalability of the proposed algorithm, the Top-1 accuracy of 72.27% reported in Table 2 appears to be comparatively low. After reviewing Appendix C.1.2, it seems that the experimental setup follows the PyTorch convention (https://github.com/pytorch/examples/tree/main/imagenet), except for some additional augmentations. It is widely recognized that the ResNet-50 model typically achieves an accuracy of approximately 76% on the ImageNet dataset using this standard setup, which significantly differs from the 72.27% accuracy reported by the authors. It would be beneficial to provide clarification on the reasons behind the considerable performance drop observed in the ImageNet results.\n\nPractical issues:\n\n* __Excessive training costs incurred by the proposed algorithm.__  \nThe Lookaround algorithm, as proposed, demands $d$ times the number of forward and backward passes for each optimization step. This is considerably higher compared to SAM (Foret et al., 2021), which only requires twice the number of passes. One could argue that the training epoch is effectively enlarged by a factor of $d$ and this is the actual reason for the performance improvements. It would be valuable to present results akin to Table 2 in Foret et al. (2021), that is, exploring the impact of increasing the number of training epochs through an ablation study.\n\nQUESTIONS:\nThere are several issues in the experimental results of the current version of the paper as I mentioned above, and I would respectfully suggest the followings that could make the paper solid:\n\n* __The absence of SAM in Table 4.__  \nIs there a specific reason for excluding SAM from the results presented in Table 4?\n\n* __The disparity in results for ResNet50 on CIFAR100 between Tables 1 and 6.__  \nTable 1 presents Top-1/5 accuracies of 81.60/95.99, while these values are not reflected in Table 6.\n\n* __Top-5 accuracy does not provide a meaningful comparison in the main tables.__  \nIt is suggested to exclude the Top-5 accuracy metric from the main results presented in Tables 1 and 4 for CIFAR10/100.  Specifically, in the case of CIFAR-10, evaluation metrics where all techniques achieve 99.9% accuracy hold little significance. Instead, including the Top-5 accuracy in Table 2 would be more appropriate, as it provides a more meaningful performance evaluation for the ImageNet experiments.\n\n* __It would be nice to provide uncertainty metrics.__  \nWhile the Top-1 accuracy is the main evaluation metric for classification, it would be advantageous to incorporate additional metrics that can effectively demonstrate the superiority of the proposed algorithm. For instance, including the negative log-likelihood (NLL) as a measure of in-domain uncertainty would provide valuable insights into the algorithm's performance.\n\n* __It would be beneficial to offer an analysis of the training costs involved.__  \nSeveral algorithms, including the proposed one, impose additional computational burdens, which could impact their practicality (e.g., SAM requires performing double forward and backward passes; SWA requires an additional model copy in memory). Providing detailed information about the specific computational requirements of each algorithm would greatly benefit future researchers and engineers. The most straightforward analysis is measuring wall-clock time for training. It would be nice to replace the meaningless \"#param.\" column with the training runtime in Table 4.\n\nLIMITATIONS:\nThe authors stated the additional training cost incurred by the proposed algorithm.\n\n---\n__References:__  \nZhang et al. (2019). _Lookahead optimizer: k steps forward, 1 step back_.  \nIzmailov et al. (2018). _Averaging weights leads to wider optima and better generalization_.  \nForet et al. (2021). _Sharpness-aware minimization for efficiently improving generalization_.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes Lookaround, a new optimization method that incorporates weight averaging into the optimization process. The algorithm consists of two steps: 1) the around step launches several parallel runs of gradient descent led by different data augmentations, 2) the average step does weight averaging of the networks obtained by these parallel runs. These two steps are repeated along the whole training process, with the result of the average step being a starting point for the next around step. The proposed method is similar to Stochastic Weight Averaging and Model Soups, though it does the averaging during training rather than at the end of training. Lookaround shows strong results compared to the existing optimization methods (SGD with momentum, AdamW, Lookahead, SAM).\n\nSTRENGTHS:\n1. The idea of using weight averaging during optimization with various augmentations is novel and leads to a better generalization of neural networks.\n2. This paper has a theoretical analysis of the quadratic noise setup which demonstrates faster convergence of Lookaround compared to SGD and Lookahead.\n\nWEAKNESSES:\nDespite the fact that the paper proposes a novel and interesting method, in my opinion, it is not quite ready for publishing in its current form. I believe that addressing the following concerns could help to make it a much stronger submission.\n1. The text of the paper is hard to follow, contains some vague parts and numerous typos.\n    - The theoretical reasoning is heavily inspired by the Lookahead paper, which makes it impossible to understand the theory without reading the original paper (e.g., I could not understand Proposition 1 and the definition of $\\alpha$ without reading the Lookahead paper). Moreover, there is no remark that equations 4 and 5 are derived in the Lookahead paper. \n    - Line 5-6: weight averaging and ensembles are mixed up. \n    - Line 313: it is not clear how gradient boosting is related to this setup.\n2. I have found two contemporaneous works published on arxiv (https://arxiv.org/abs/2302.14685, https://arxiv.org/abs/2304.03094) that propose approaches very similar to Lookaround. I believe they need to be at least discussed in the related work section.\n3. The experimental part of the paper raises some questions and is not convincing enough, in my opinion.\n    - The main results (Table 1) lack standard deviations, even though they are claimed in the OpenReview form. \n    - The baselines seem too weak, i.e., the Lookahead paper reports >75% ImageNet test accuracy on ResNet-50 for SGD baseline, while this paper shows 72.27%. \n    - The augmentation policy during baseline training is not clear. If the same set of augmentations as for Lookaround is used, then this may be the reason for the bad quality of the baselines. \n    - Despite the paper proposing a new optimization method, it lacks experiments on a wider choice of architectures. Table 1 illustrates only VGG-19 and four networks from the ResNet family. It would be beneficial to add experiments on more modern convolutional architectures and more extensive experiments on image transformers.\n    - The paper lacks a fair comparison to logit ensembles. If I understood correctly, different models of the ensemble are trained with different Lookaround augmentations. This leads to the suboptimal quality of each model and, as a result, to a suboptimal ensemble.\n    - The comparison to model soups in Appendix D is not fair as well, i. e., the original paper fine-tunes models with various hyperparameters and does a greedy search for the combination with the best validation accuracy. Near zero accuracy of the average of $\\theta_1$ and $\\theta_2$ may be due to improper fine-tuning hyperparameters. Moreover, this is an important baseline, and it would be better to move this comparison to the main part of the paper.\n    - A proper ablation study of augmentations is required. Is it possible to train Lookaround without augmentations for the around step “branches” (e.g., the only source of randomness is batch ordering)? What if each branch utilizes the same set of augmentations? Why are the mentioned augmentations used (and not some other ones, e. g. vertical flip seems to be a strange augmentation)? Probably, some of these setups are covered in Section 4.4, but it is not clear from the text.\n    - Top-5 accuracy is redundant and makes it difficult to read the tables.\n4. Minor issues/typos:\n    - Line 25: the sentence is about LMC, but the citation [7] leads to the paper about permutations (Entezari et al., 2022). \n    - Line 112-113: we provide\n    - Line 147-148: repetition\n    - Line 156: we analyze\n    - Line 162: it seems that it should be $\\mathbb{E} [c_i^2]$\n    - Line 242: CIFAR\n\nQUESTIONS:\n1. What is the augmentation policy during baseline training? \n2. I do not understand the experimental setup in the ablation study (Section 4.4). Is the training without DA and WA similar to the regular training without augmentations? Is the training without DA but with WA similar to Lookaround, where the ordering of batches is the only source of randomness between branches?\n3. Section 3.2.2, convergence on deterministic quadratic functions. How branches of Lookaround are different from each other if there is no randomness in gradient steps?\n4. Could the authors elaborate on how the Lookaround method relates to the methods proposed in the contemporaneous works (https://arxiv.org/abs/2302.14685, https://arxiv.org/abs/2304.03094)?\n\nLIMITATIONS:\nThe main limitation of the proposed method is the increased training budget, which is highlighted by the authors. However, there is no fair comparison to networks trained for a larger number of epochs.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis article introduces the Lookaround Optimizer, a novel optimization algorithm for deep neural networks. The Lookaround Optimizer is based on the idea of lookaround, which involves maintaining two sets of weights. The first set of weights is updated using the standard gradient descent algorithm, while the second set of weights is updated using the averaged gradients of the first weight. The Lookaround Optimizer has been shown to improve the generalization performance of deep neural networks, and it outperforms other state-of-the-art optimization algorithms on a variety of benchmark datasets. The authors also provide theoretical analysis of the Lookaround Optimizer. Overall, the Lookaround Optimizer is a promising approach for improving the training of deep neural networks.\n\nSTRENGTHS:\n- Lookaound Optimizer performs better than other state-of-the-art optimization algorithms (Lookahead, SAM) on multiple benchmark datasets.\n- The proposed optimizer seems to be model-free and can be applied to various computer vision scenarios.\n- The paper is well-written and easy to understand.\n\nWEAKNESSES:\n- While the authors have demonstrated fascinating performance on benchmarks, the technical innovation of this paper is limited. Similar ideas that utilize multiple models have already in proposed in varioius area, especially in meta learning. \n- With that being said, while the authors have mentioned the connections to these related topics, they did not compare with some of them. For example, I think at least the authors should compare the performance of Lookaround Optimizer with Model Soups (i.e., train multiple copies of models using the same initialization, the same order of data, but with different augmentation), and other model merging techniques like Model Ratatouille. \n- The training time will be a bottleneck for applying this method as the authors  acknowledged. \n- Another limitation is that, the augmentation seems to be a must for this method. Therefore, it seems that this method is not applicable (at least not straightforward) when we try to apply it on other modality, such as language and speech.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nThe authors have addressed the limitation.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFlatness-aware optimizers have gained significant attention in the field of research for training deep neural networks that are robust. Weight Averaging (WA) is a popular approach to finding solutions within the flat regions of the loss surface. However, previous WA methods have two limitations. First, when WA is performed within a single optimization trajectory after training convergence, resulting in limited functional diversity among the averaged members. Second, when WA is performed across diverse modes, the averaged weights may negatively impact performance. To address these challenges, this paper introduces a new method called Lookaround Optimizer, which aims to overcome these limitations.\n\nSTRENGTHS:\nOriginality\n\n- Proposed optimizer seems to be relatively simple and straightforward but still attractive.\n- This article presents a theoretical analysis regarding the variance of the steady state.\n\nClarity\n\n- The paper is written effectively, ensuring high accessibility for readers.\n- Methods are simple and easy to follow.\n\nWEAKNESSES:\nExperiments\n\n- Regarding the experiments conducted, the authors acknowledged in the Limitation section that their proposed method incurs additional training costs (I think the proposed optimizer requires nearly $d$ times more computational resources when averaging $d$ models). Therefore, for a fair comparison, it would be necessary for other baseline methods to undergo additional training epochs as well. However, it appears that the authors used the same number of training epochs for all the experiments.\n\n- The overall results of the experiments are rather counterintuitive. Particularly in Table 3, the reported accuracy of the Deep Ensemble method being lower than the accuracy of a single solution trained using the proposed method is quite unexpected. Similarly, in Table 1, the reported performance of flatness-aware optimizers such as SWA [1] and SAM [2] being lower than SGDM contradicts the knowledge and expectations of the research community.\n\nReferences\n\n[1] Pavel Izmailov, D. A. Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. Conference on Uncertainty in Artificial Intelligence, 2018.\n\n[2] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. International Conference on Learning Representations, 2021.\n\nQUESTIONS:\nExperiments\n\n- The response regarding concerns addressed in the Weakness section will be critical to the final decision.\n\nLIMITATIONS:\nLimitations are adequately addressed in the Limitation section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-12"}