SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lookaround Optimizer:  $k$  steps around, 1 step average
Jiangtao Zhang $^{1}$ , Shunyu Liu $^{1}$ , Jie Song $^{1,\dagger}$ , Tongtian Zhu $^{1}$ , Zhengqi Xu $^{1}$ , Mingli Song $^{1,2}$ $^{1}$ Zhejiang University,  $^{2}$ Hangzhou City University  
{zhjgtao, liushunyu, sjie, raiden, xuzhengqi, brooksong}@zju.edu.cn
Weight Average (WA) is an active research topic due to its simplicity in assembling 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.
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.
WA has 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, WA easily breaks down owing to the diversity-locality 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
Figure 1: Test set loss landscape. (Left) Around step for diversity. (Right) Average step for locality.
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.
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 motivate 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.
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.
Require: Initial parameters  $\phi_0$ , objective function  $\mathcal{L}$ , data augmentation list AUG of size  $d$ , synchronization period  $k$ , optimizer  $A$ , dataset  $\mathcal{D}$ , numbers of training epochs  $E$ .
for epoch  $= 1$  to  $E$  do
Synchronize parameters
for  $j = 1,2,\ldots ,d$  do
$$
\theta_ {t, j, 0} \leftarrow \phi_ {t - 1}
$$
Around Step: Independent model training.
for  $i = 1,2,\ldots ,k$  do
sample minibatch of data  $B\sim \mathcal{D}$
$$
\theta_ {t, j, i} \leftarrow \theta_ {t, j, i - 1} + A (\mathcal {L}, \theta_ {t, j, i - 1}, \mathrm {A U G} _ {j} (B))
$$
Average Step: Weight averaging.
Compute average weight  $\theta_{t,*,k} \gets \frac{1}{d}\sum_{j=1}^{d}\theta_{t,j,k}$
Perform update  $\phi_t\gets \theta_{t,*,k}$
return parameters  $\phi$
multiple intervals of length  $k$ , with each interval consisting of an "around step" and an "average step". Let  $\phi_0$  denote the initial weights. At  $t^{th}$  round of Lookaround, the weights are updated from  $\phi_{t-1}$  to  $\phi_t$  according to the around step and the average step. The resulting weights  $\phi_t$  then serve as the starting point for the subsequent round.
Around Step. Given the optimizer  $A$ , the objective function  $\mathcal{L}$ , and the stochastic mini-batch of data  $B$  from train dataset  $\mathcal{D}$ . In the around step,  $d$  different models are independently trained under  $d$
different data augmentations  $\mathrm{AUG} = \{\mathrm{AUG}_1,\dots ,\mathrm{AUG}_d\}$  for  $k$  batch steps, as follows:
$$
\theta_ {t, j, k} = \theta_ {t, j, k - 1} + A (\mathscr {L}, \theta_ {t, j, k - 1}, \mathrm {A U G} _ {j} (B)). \tag {1}
$$
Thanks to the data augmentations, each model is trained in period  $k$  to a diverse location in the loss landscape scattered around  $\phi_{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  $\theta_{t,i,k}$ , the weights of each independently trained model to obtain  $\phi_t$  for the next around step as follows:
$$
\phi_ {t} = \frac {1}{d} \sum_ {i = 1} ^ {d} \theta_ {t, i, k}. \tag {2}
$$
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.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  $\hat{\mathcal{L}} (\theta) = \frac{1}{2} (\theta -\mathbf{c})^T\mathbf{A}(\theta -\mathbf{c})$ . Following Zhang et al. [49], we assume that the noise vector  $\mathbf{c}\sim \mathcal{N}(\theta^{*},\Sigma)$ , both  $\mathbf{A}$  and  $\boldsymbol{\Sigma}$  are diagonal matrices, and that the optimal solution  $\theta^{*} = \mathbf{0}$ . We denote  $a_{i}$  and  $\sigma_i^2$  as the  $i$ -th elements on the diagonal of  $\mathbf{A}$  and  $\boldsymbol{\Sigma}$ , respectively. As the  $i$ -th element of the noise vector  $\mathbf{c}$  satisfies  $\mathbb{E}[c_i^2 ] = (\theta_i^*)^2 +\sigma_i^2$ , the expected loss of the iterates  $\theta_t$  can be written as follows:
$$
\mathcal {L} \left(\theta_ {t}\right) = \mathbb {E} \left[ \hat {\mathcal {L}} \left(\theta_ {t}\right) \right] = \frac {1}{2} \sum_ {i} a _ {i} \left(\mathbb {E} \left[ \theta_ {t, i} \right] ^ {2} + \mathbb {V} \left[ \theta_ {t, i} \right] + \sigma_ {i} ^ {2}\right), \tag {3}
$$
where  $\theta_{t,i}$  represents the  $i^{th}$  item of the parameter  $\theta_t$ . The limiting risk of SGD, Lookahead and Lookaround are compared in the following Proposition by unwrapping  $\mathbb{E}[\theta_t]$  and  $\mathbb{V}[\theta_t]$  in Equation 3.
Proposition 1 (Steady-state risk). Let  $0 < \gamma < 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:
$$
V _ {S G D} ^ {*} = \frac {\gamma^ {2} \mathbf {A} ^ {2} \Sigma^ {2}}{\mathbf {I} - (\mathbf {I} - \gamma \mathbf {A}) ^ {2}}, \tag {4}
$$
$$
V _ {\text {L o o k a h e a d}} ^ {*} = \underbrace {\frac {\alpha^ {2} (\mathbf {I} - (\mathbf {I} - \gamma \mathbf {A}) ^ {2 k})}{\alpha^ {2} (\mathbf {I} - (\mathbf {I} - \gamma \mathbf {A}) ^ {2 k}) + 2 \alpha (1 - \alpha) (\mathbf {I} - (\mathbf {I} - \gamma \mathbf {A}) ^ {k})}} _ {\leqslant \mathbf {I}, i f \alpha \in (0, 1)} V _ {S G D} ^ {*}, \tag {5}
$$
$$
V _ {\text {L o o k a r o u n d}} ^ {*} = \underbrace {\frac {\alpha^ {2} \left(\mathbf {I} - \left(\mathbf {I} - \gamma \mathbf {A}\right) ^ {2 k}\right) + 2 \alpha (1 - \alpha) \left(\mathbf {I} - \left(\mathbf {I} - \gamma \mathbf {A}\right) ^ {k}\right)}{\alpha^ {2} \left(d \mathbf {I} - (d - 1) (\mathbf {I} - \gamma \mathbf {A}) ^ {2 k}\right)}} _ {\leqslant \mathbf {I}, i f d \geqslant 3 a n d \alpha \in [ 1 / 2 , 1)} V _ {\text {L o o k a h e a d}} ^ {*}, \tag {6}
$$
respectively, where  $\alpha$  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 \geqslant 3$  and  $\alpha \in [1/2, 1)$ . 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  $\rho$  characterizes the speed in quadratic function at which the independent variable converges to the optimal solution, satisfying  $||\theta_t - \theta^* || \leqslant \rho^t ||\theta_0 - \theta^* ||$ . 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
Figure 2: Convergence rate on quadratics of varying condition number. We fix the step  $k = 20$  for Lookahead and Lookaround, and fix the CM factor  $\beta = 0.99$ .
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.
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
Table 1: Test set accuracy under training procedure with random initialized models. In this table, all models are trained for same amount of time at  $32 \times 32$  resolution. (To provide a more comprehensive view of the data, we substitute the Top5 metric for CIFAR10 with the NLL loss.)
Method	VGG19	ResNet50	ResNet101	ResNet152	ResNext50
Top1	NLL	Top1	NLL	Top1	NLL	Top1	NLL	Top1	NLL
CIFAR10	SGDM	93.92	0.26	95.96	0.16	96.16	0.16	96.28	0.16	95.72	0.17
SWA	94.89	0.21	96.42	0.14	96.34	0.14	96.82	0.13	96.09	0.16
SWAD	93.23	0.29	95.39	0.19	94.49	0.22	95.00	0.20	93.89	0.26
Lookahead	94.72	0.23	96.38	0.16	96.61	0.15	96.46	0.15	96.47	0.15
Ours	94.44	0.25	96.59	0.14	96.73	0.14	97.02	0.14	96.70	0.13
Method	VGG19	ResNet50	ResNet101	ResNet152	ResNext50
Top1	Top5	Top1	Top5	Top1	Top5	Top1	Top5	Top1	Top5
CIFAR100	SGDM	73.84	91.09	79.61	95.21	79.91	95.54	80.16	95.49	79.10	94.86
SWA	73.62	90.89	80.17	95.56	80.53	95.39	80.86	95.43	79.14	94.96
SWAD	73.37	89.93	80.19	95.46	79.92	95.09	80.00	95.36	79.27	95.17
Lookahead	74.02	90.68	80.06	95.24	81.14	95.84	81.36	95.77	80.12	95.26
Ours	74.29	91.15	81.60	95.99	81.97	96.15	82.22	96.40	81.14	96.12
Figure 3: Training loss and Top-1 accuracy curves of various networks on CIFAR100 under different optimization methods.
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.
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
Figure 4: Training loss and Top-1 accuracy curves of ResNet50 (left) and ViT-B/16 (right) on CIFAR100 under different optimization methods.