Improving Generalization and Convergence by Enhancing Implicit Regularization
Authors: Mingze Wang
Jinbo Wang
Haotian He
Zilin Wang
Guanhua Huang
Feiyu Xiong
Zhiyu li
Weinan E
Lei Wu
Venue: NeurIPS 2024 poster (2024)
Keywords: implicit regularization
convergence
generalization
deep learning
Primary area: optimization_for_deep_networks
TL;DR: We propose an algorithmic framework, Implicit Regularization Enhancement, which can improve both generalization and convergence of deep neural networks..

Abstract: In this work, we propose an Implicit Regularization Enhancement (IRE) framework to accelerate the discovery of flat solutions in deep learning, thereby improving generalization and convergence. 
Specifically, IRE decouples the dynamics of flat and sharp directions, which boosts the sharpness reduction along flat directions while maintaining the training stability in sharp directions. We show that IRE can be practically incorporated with *generic base optimizers* without introducing significant computational overload. Experiments show that IRE consistently improves the generalization performance for image classification tasks across a variety of benchmark datasets (CIFAR-10/100, ImageNet) and models (ResNets and ViTs). 
Surprisingly, IRE also  achieves a $2\times$ *speed-up* compared to AdamW in the pre-training of Llama models (of sizes ranging from 60M to 229M) on datasets including Wikitext-103, Minipile, and Openwebtext. Moreover, we provide theoretical guarantees, showing that IRE can substantially accelerate the convergence towards flat minima in Sharpness-aware Minimization (SAM).

Improving Generalization and Convergence by Enhancing Implicit Regularization
Mingze Wang 3
Jinbo Wang
Haotian Hel.3
Zilin Wang' Guanhua Huang 5.6
Feiyu Xiong?
Zhiyu Li Weinan E 1,2,3.4
Lei Wu 12tt
School of Mathematical Sciences, Peking University
Center for Machine Learning Research, Peking University
Institute for Advanced Algorithms Research (Shanghai) AAI for Science Institute
School of Data Science. University of Science and Technology of China *ByteDance Research
(mingzewang, wangjinbo, hactianhe, wangzilin]@stu, pkt..kkuuede guanhnshuangdmail. ustc..du..c. (xiongfy, lizy) cn {weinan, 1eiwu] waw.maahhpk.edu..... edu.
Abstract
In this work, propose an Implicit Regularization Enhancement (IRE) framework to accelerate discovery flat solutions deep learning, thereby improving generalization and convergence. Specifically, IRE decouples the dynamics of flat and sharp directions, which boosts the sharpness reduction along flat directions while maintaining the training stability in sharp directions. We show that IRE can be practically incorporated with generic base optimizers without introduc- ing significant computational overload. Experiments show that IRE consistently improves the generalization performance for image classification tasks across a variety of benchmark datasets CCFFAR-101000 ImageNet) and models (ResNets and ViTs). Surprisingly, IRE also achieves 2x speed-up compared to Adam W the pre-training of Llama models (of sizes ranging from 60M 229M) datasets including Wikitext-103 Minipile, and Openwebtext, Moreover, we provide theo- retical guarantees, showing that IRE can substantially accelerate the convergence aware aware minimization (SAM).
towards flat minima in sharpness
Introduction
Deep learning has achieved remarkable success across a variety of fields, including computer vision, scientific computing, and artificial intelligence. The core challenge in deep learning lies in how to train deep neural networks (DNNs) efficiently to achieve superior performance. Understanding and improving the generalization and convergence commonly-used optimizers, such stochastic gradient descent (SGD) (Robbins and Monro, 1951; Rumelhart et al., 1986), in deep learning is crucial for both theoretical research and practical applications.
Notably, optimizers often exhibit a preference for certain solutions in training DNNs. For instance, SGD and its variants consistently converge solutions that generalize well, even when DNNs are highly over-parameterized and there are many solutions that generalize poorly. This phenomenon is referred to as implicit regularization in the literature (Neyshabur et al., 2014; Zhang et al., 2017).
The most popular explanation for implicit regularization is that SGD and its variants tend to converge to flat minima (Keskar al., 2016; Wu et al., 2017), and flat minima generalize better (Hochreiter and
Correspondence to: Mingze Wang and Lei Wu.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Schmidhuber, 1997; Jiang et al., 2019). However, the process of this implicit sharpness regularization occurs at a very slow pace, as demonstrated in works such as Blanc al. (2020), Li et al. (2022), and Ma et al. (2022). Consequently, practitioners often use large learning rate (L.R) and extend the training time even when the loss no longer decreases, ensuring the convergence to flatter minima (He et al., 2016; Goyal et al., 2017; Hoffer et al., 2017). Nevertheless, the largest allowable LR is constrained by the need to maintain training stability. In addition, Foret et al. (2021) proposed SAM, which aims to explicitly regularize sharpness during training and has achieved superior performance across a variety of tasks.
Our contributions can be summarized as follows:
We propose an Implicit Regularization Enhancement (IRE) framework to speed up the con- vergence towards flatter minima. As suggested by works like Blanc et al. (2020), Li al. (2022) and Ma et al. (2022), the implicit sharpness reduction often occurs at a very slow pace, along flat directions. Inspired by this picture, IRE particularly accelerates the dynamics along flat directions, while keeping sharp directions' dynamics unchanged. As such, IRE can boost the implicit sharpness reduction substantially without hurting training stability. For detailed illustration of this mechanism, we refer to Section 2.
We then provide a practical IRE framework, which can be efficiently incorporated with generic base optimizers. We evaluate the performance of this practical IRE in both vision and language tasks. For vision tasks, IRE consistently improves the generalization performance of popular optimizers like SGD, Adam, and SAM in classifying the CIFAR-10/100 and ImageNet datasets with ResNets (He et al., 2016) and vision transformers (ViTs) (Dosovitskiy et al., 2020). For language modelling, we consider the pre-training of Llama models (Touvron et al., 2023) of various sizes, finding that IRE surprisingly can accelerate the pre-training convergence. Specifically, we observe a remarkable 2.0x speedup compared to AdamW in the scenarios we examined, despite IRE being primarily motivated to speed up the convergence to flat solutions.
Lastly, we provide theoretical guarantees showing that IRE can achieves a ((L/p)-time acceler- ation over the base SAM algorithm in minimizing the trace of Hessian, where p E (0, 1) is a small hyperparameter in SAM.
1.1 Related works
Implicit sharpness regularization. There have been extensive attempts to explain the mystery of implicit regularization deep learning (see the survey by Vardi (2023) and references therein). Here, we focus on works related to implicit sharpness regularization. Wu et al. (2018; 2022) and Ma and Ying (2021) provided an explanation of implicit sharpness regularization from dynamical stability perspective. Moreover, in-depth analysis of SGD dynamics near global minima shows that the SGD noise (Blane et al., 2020; Li et al., 2022; Ma et al., 2022; Damian et al., 2021) and the edge of stability (EoS)-driven (Wu et al., 2018; Cohen et al., 2021) oscillations (Even et al., 2024) can drive SGD/GD towards flatter minima. Additional studies explored how training components, including learning rate and batch size (Jastrzębski et al., 2017), normalization (Lyu et al., 2022), cyclic LR (Wang and Wu, 2023), influence this sharpness regularization. Furthermore, some works have provided theoretical evidence explaining the superior generalization of flat minima for neural networks (Ma and Ying, 2021; Mulayoff et al., 2021; Wu and Su, 2023; Gatmiry et al., 2023; Wen et al., 2023b). Our work is inspired by this line of research, aiming to boost implicit sharpness regularization by decoupling the dynamics along flat and sharp directions.
Sharpness-awan minimization. IRE shares the same motivation as SAM in enhancing sharpness regularization, although their specific approaches differ significantly. It is worth noting that the per-step computational cost of SAM is twice that of base optimizers. Consequently, there have been numerous attempts to reduce the computational cost of SAM (Kwon et al., 2021; Liu et al., 2022: Du et al., 2021; Mi et al., 2022; Mueller et al., 2024). In contrast, the per-step computational cost of IRE is only approximately 1.1 times that of base optimizers (see Table 5). Moreover, we provide both theoretical and experimental evidence demonstrating that the mechanism of IRE in boosting sharpness regularization is nearly orthogonal to that of SAM.
Optimizers for large language model (LLM) pre-training. (Momentum) SGD (Sutskever et al., 2013; Nesterov, 1983) and its adaptive variants like Adagrad (Duchi et al., 201 RMSProp (Tieleman, 2012), and Adam (Kingma and Ba, 2014) have been widely used in DNN training. Despite the efforts

---

in designing better adaptive gradient methods (Liu et al., 2019a; Luo et al., 2019; Heo et al., 2020; Zhuang et al., 2020; Xie al., 2022b;a), AdamW (Adama((ecoouppllee weight decay) (Loshchilov and Hutter, 2017) has become the default optimizer in LLM pre-training. Recently, Chen et al. (2024) discovered Lion by searching the space of adaptive first-order optimizers; Liu et al. (2024) introduced Sophia, scalable second-order optimizer. In this paper, we instead empirically demonstrate that IRE can accelerate the convergence of AdamW in the pre-training Llama models.
1.2 Notations
Throughout this paper, let L: - R30 be the function of total loss, where p denotes the number of model parameters. For a C² -u-smmmanioold MM in TR,, we denote the tangent space of M at A E M as T,M, which is linear subspace in R For f (M) and 0 M, let V ( denote the Riemannian gradient, where To. Iº + Rº denotes the orthogonal projection to M. For symmetric matrix E _2xp, its eigen pairs are denoted as {(Ai, [,))e with the
Q MM (
order A >> ... > Ap We use (A) = Ei= un to denote the projection operator onto span (u...... U,) Denote N(H, Σ) as the Gaussian distribution with mean μ and covariance matrix Σ, and USS as the uniform distribution over set S. Given a vector h = (hi....: Ap)) let
|h| (lh1)..... I/) We denote by the all-ones vector. We will use standard big-O notations like O()), (()). and (-)) to hide constants.
2 An Illustrative Example Motivating IRE
In this section, we provide an illustration of how the dynamics along flat directions can reduce the sharpness (curvatures along sharp directions) and how IRE can accelerate this sharpness reduction. To this end, we consider the following phenomenological problem:
min L( = H(u)u/2, DERP
where E Rimm u E RR-m, and vec(u, 2) E Re. We assume H()) E (TR--m and inf, Amiu (H(u)) > O. Then, the minimizers of (() form m-dim manifold M = {(us, 2) 0}
and the Hessian at any 0 E M is given by V=L(0) c H(u)) For clarity, we shall call u and the flat and sharp directions, respectively.
Example 2.1. The loss landscape of fitting zero labels with two-layer neural networks (2LNNs) exhibits exactly the form (1). Let f(z; 0) = a o(x; W) be a 2LNN with 0 = vec(W, a). Then (()) [[ 0) //2 a Eg (d(:; W) W) la/2 =: a HWWaa/2
For breviety, we further assume Huu) diag A(u)) with X(u) (u). Am (u)). ) In this case, the GD dynamics can be naturally decomposed into the flat and sharp directions as follows
m
U++1 ui
17
2 (),(),)
2
Vt+1 (1-7)(u,) (u:() Vr:
where: 2 denotes the element-wise multiplication of two vectors.
The implicit sharpness regularization. From Eq. (2), we can see that 1) the flat direction u, dynamics monotonically reduces the sharpness A(u) as long as v, is nonzero; 2) the sharp direction v,'s dynamics determines the speed of sharpness reduction. The larger |vl is, the faster the curvature X(u) decreases. Particularly, when near convergence, we have |uyl = o(1) and thus the implicit sharpness reduction is very slow during the late phase of GD. Figure la provides a visualization of this slow implicit sharpness reduction.
Accelerating the sharpness reduction. Inspired by the above analysis, we can accelerate the sharp- ness reduction by speeding up the flat directions' dynamics. To this end, there are two approaches:
Naively increasing the global learning rate 7) (fail), Increasing 7) accelerates the dynamies of Ue but the largest allowed 7) is constrained by curvatures of sharpest directions. In GD (2), to
3

---

mara masidd M
CD [em-
0.5
0.01
the fltter aeliien
ietur
0.00
0.03
-0.01
-05
-0.01
0,500
0,552
1,000
0.5
0.0
0.5
1.0
0.2
0.4
0.5
(a) GD with 17 1.
(b) GD with
2.
(c) IRE with various R's v.s. GD with 1y
Figure I: A 2-d example of (I): (u,, = (1 + + 2 The gray arrows denote to the minima manifold M ((u, : 0}. where the smaller the ru]. the flatter the minimizer. The red marker highlights the flattest minimizer (0, 0). (a) The dynamics of GD (7) = 1), which moves slowly towards flatter minima it converges. (b) The dynamies of GD (1) = 2), which diverges due to excessively large n, (c) The behavior of
our IRE approach with varying R's v.s. GD (n 1). Is is shown that IRE can significantly accelerate the uy
dynamics, almost reaching the flattest minimum (0, 0) when taking a very large
maintain training stability, 1 must be smaller than 2/ max; (uee)). Otherwise, V, 's dynamics will blow up. As illustrated in Figure 1b, setting r 2 leads divergence, whereas 7) ensures convergence.
Increasing only the flat directions' learning rate (our approach, IRE). Specifically, for GD (2), the GD-IRE dynamics is given by
ux+1 U (1 k) o (1))).
(3)
0++1 (1 (1-1) s)/((()))))
2 De:
where K >> O controls the enhancement strength. In GD-IRE (3), u,'s dynamics is (1 + +x) faster than that of GD (2). Notably, the sharp directions' dynamics (v2) unchanged. The choice of x only needs to maintain the stability of flat directions' dynamics, for which, we can always take significantly large enhance the sharpness regularization. As demonstrated in Figure Ic, IRE with larger always find flatter minima.
Remark 2.2 (The generality). It is worth noting that similar implicit sharpness regularization also holds for SGD (Ma et al., 2022; Li et al., 2022) and SAM (Wen et al., 2023a). In this section, we focus on the above toy model and GD mainly for illustration. In Appendix B. we provide an analogous illustrative analysis of how IRE accelerates the sharpness reduction of SGD. In Section 5, we further provide theoretical evidence to show that IRE can boost the implicit sharpness regularization of SAM.
A Practical Framework of Implementing IRE
Although the preceding illustration of IRE is for GD, in practice, we can incorporate IRE with any base optimizers. Specifically, for a generic update: 0.41 = O, - rg:,, the corresponding IRE modification is given by
0 = 1 n/gg + kPigt).
where x denotes the enhancement strength and P. R + R' projects gi into the flat directions of the landscape. The flat directions and corresponding projection operator P. can be estimated using the Hessian information.
However, estimating the full Hessian matrix V ((L( E Rexp is computationally infeasible. Con- sequently, we propose to use only the diagonal Hessian diag(V² L(O))) E TR? to estimate P,. Let h, R. be estimate of the diagonal Hessian. Then, we perform the projection follows
{
if (lh,l); < Topemali (1h.),)) otherwise
Pig: n4 ④ 9t, with (n+))
(5)
where E (0, and Top Topemall ( [/eel.7 7) returns the 1p-7|-th smallest value in |hs Note that n, E RR denotes a mask vector and the above approximate projection essentially masks the top-(1 (1 sharp coordinates out. As such, the projection (5) will retain the top-ry flat coordinates. Noticing that in DNNs, there are much more flat directions than sharp directions (Yao et al., 2020), we thus often use 7 0.5 in practice.

---

Algorithm 1: Practical IRE (A practical framework of implementing IRE)
Input: 0g T. K, learning rate {ncce. warm-up time Tw, IRE hyperparams: and y (0.5. 1):
for f == Tw: ... T do
Compute the original update direction 9: according to the base optimizer;
(t Tw) mod = 0 then
Estimate the diagonal Hessian h, E IRP using Eq. (6);
Update the mask n, E Rº using Eq. (5);
else
nu: = 1
0i+1 0: Tte (grt kn, O gt
Output: OT.
A light- estimator of the diagonal Hessian. Let E(.,-) be the cross-entropy loss. Given an input data æ E Rdx and label y E RR,, let the model's prediction be f(æ; 0) E IRd. The Fisher (Gauss-Newton) matrix F()) is widely acknowledged to be a good approximation of the Hessian, particularly near minima. Thus, we can estimate the diagonal Hessian by h. diag F(O.)). which has been widely used in deep learning optimization (Martens and Grosse, 2015; Grosse and Martens, 2016; George et al., 2018; Mi et al., 2022; Liu et al., 2024). Given an input batch {(xs) ) the empirical diagonal Fisher is given by diag (()) Vf(f 0); ÿ)), where ÿn ~ softmax (0; xt)). However, as noted by Liu et al. (2024), imple- menting this estimator is computationally expensive due to the need to calculate B single-batch gradi- ents. Liu et al. (2024) proposed a more convenient estimator diag( Fen (0)), only requires computing the mini-batch gradient VLg (0) = 1 Σ (x); 0): 9b) with ys softmax(f(ax) 0)):
T Σ (e(f(z;; 0): 9o
h diag (0)
( O (L)
(6)
According to Liu et al. (2024), this estimator is unbiased estimate of the empirical diagonal Fisher, i.e., Eg [diagg Fett (0)) = E, (diag(f(())) For more discussions on the efficiency of this estimator, please refer to (Liu et al., 2024, Section 2). Additionally, for squared loss, one can simply use Fisher the estimator (Liu et al., 2024).
The practical IRE and computational efficiency. The practical IRE is summarized in Algorithm 1, which is notably lightweight. The estimation of h, using (6) requires computational resources roughly equivalent to one back- proppppppaattion. Consequently, by setting K = 10 in Algorithm I (estimating the projection every 10 steps), the average per-step computational load of IRE is only 1.1 times that of the base optimizer. This claim can be empirically validated as shown in Table 5.
4 Experiments
In this section, we evaluate how IRE performs when incorporating with various base optimizers. Specifically, examine the incorporation of IRE with SGD (SGD-IRE), SAM (SAM-IRE), and AdamW (AdmIRE) across vision and language tasks.
4.1 Image classification
4.1.1 Validating our motivation
To show that IRE can accelerate the sharpness reduction, we train WideResNet-16- (Zagoruyko and Komodakis, 2016) on CIFAR-10 dataset (Krizhevsky and Hinton, 2009) by SAM-IRE (with K = 10, varying i and 7). Here, we incoporate IRE into SAM starting from the 30-th epochs. We vary fy {0.8, 0.9, 0.95) and и E {0, 2, 5, 10). Regarding the learning rate (LR), both constant LR and decayed LR are considered. The sharpness is measured by Tr(V ((0)). Further experimental details can be found in Appendix C.
As depicted in Fig. 2(a), SAM-IRE (with constant LR) consistently finds flatter solutions compared to SAM and higher always leads to flatter minima. Additionally, SAM-IRE also shows robustness to
5

---

variations of For SAM-IRE with decayed LR (Fig. 2(b)), SAM-IRE still consistently finds flatter solutions than SAM. Notably, flatter solutions correlate positively with lower training loss and higher test accuracy (Fig. 2(c,d)).
sharpness pness
shar pinsiss
training
accuracy (%6
17988
1665
1478
1478
7774
7760
76/66
2581
0.095
0.037
0.095
0.016
48
44:05
94.55
50%
1600
906
2758
0.035
9.021
94.12
68%
1728
1002
1241
1200
2774
2725
2503
2005
-00
0.009
0.533
9.019
0.013
50
94119
94.73
95.27
0.0
2.0
5.0
10.0
0.0
2.0
5.0
10.0
0.0
2.0
5.0
16.0
0.0
2.0
5.0
16.0
(a) sharpness, constant LR. (b) sharpness, decayed LR. (c) train loss, decayed LR. (d) test acc, decayed LRR
Figure 2: Training WR--16-8 CIFAR-10 by SAM-IRE with varying y, R. Particularly, the case of Rt correspond the standard SAM.
0
4.1.2 IRE can consistently improve generalization
Convolutional Neural Networks (CNNs). In this experiment, we first consider the classification of (1FAR-( with WideRes (Zagoruyko and Komodakis, 2016) and ResNet-56 (He al., 2016). Both SGD and SAM optimizers are adopted. All the experiments use base data augmen- tation and label smoothing. For SGD-IRE/SAM-IRE, we fix = 10, and tune hyperparameters y and K via a grid search over E {0.99, 0.9, 0.8} and K E {1, 21. The total epochs are set to 100 for CIFAR- 10 and 200 for CIFAR-100, and we switch from SGD/SAM to SGD-IRE/SAM-IR when the training loss approaches 0.1. The other experimental details are deferred Appendix C and the results are shown in Table 1.
Secondly, we evaluate IRE for training ResNet-50 on ImageNet (Deng et al., 2009). The experimental details are deferred to Appendix and the results are shown in Table
Vision Transformers (ViTs). We also examine the impact of IRE on generalization of ViT-T and ViT-S (Dosovitskiy al., 2020) on CIFARION 100. The default optimizers used are AdamW and SAM (Mueller et al., 2024), Strong data augmentations (basic + AutoAugment) are utilized. The total epochs set to 200 and we switch from Adam AdamW/SAM
AdmIRE/SAM-IRE when the training we fix K = 10, and tune hyperparameters and x via a grid search over y E {0.99, 0.9, 0.8) and λ E {20, 50}. Other experimental details are deferred to Appendix C. The results are shown in Table 3.
loss approaches 0.5. For AdmIRE/SAM-
Additionally, evaluate IRE for training VTT-S ImageNet. The experimental details are deferred to Appendix C. and the results are shown in Table 4.
As demonstrated in Table 1, 2, 3, and 4, IRE consistently improves generalization of SGD, AdamW and SAM across all settings examined.
Table 1: WRN-28 28- 10/ResNet-56 on CIFAR- 10/100.
WRN-28-10
ResNet-56
CIFAR-10
100
CIFAR-I0
CIFAR- 100
SGD
95.84
80.81
93.49
72.81
SGD-IRE
96.24 (+0.40)
81.49 (+0.68)
93.78 (+0.29)
73.78 (+0.97)
SAM
96.58
83.05
94.05
75.54
SAM-IRE
96.70 (+0.12)
83.50 (+0.45)
94.46 (+0.41)
75.86 (+0.32)
Table 2: ResNet-50 on ImageNet.
Table 3: VIT-T/S on CIFAR- 100.
Top-1
Top-5
ViT-T
V
SGD
76.81
93.31
AdamW AdmIRE
63.90
65.43
SGD-IRE
77.04 (+0.23)
93.58 (+0.27)
67.05 (43.15)
68.39 (+2.96)
SAM
77.47
93.90
SAM
64.25 67.33 (+3.08)
66.93
SAM-IRE
77.92 (+0.45)
94.12 (+0.22)
SAM-IRE
70.47 (+3.54)

---

Table 4: ViT-S on ImageNet.
Top-1
Top-5
AdamW
78.7
94.0
AdmIRE (R 2,7 0.6)
79.0 (+0.3)
94.3 (+0.3)
AdmIRE (x = 2, 0.8)
79.1 (+0.4)
94.2 (+0.2)
4.2 Large language model pre-training
We now evaluate IRE in the pre-training of decoder-only large language models (LLMs). Following the training protocol of Llama models, employ the AdamW optimizer with hyperparameters 31 = 0.9, 82 = 0.95 and weight decay À = 0.1 (Touvron et al., 2023). The learning rate strategy includes warm-up up phase followed by cosine decay scheduler, capped ir max. In each experiment, we tune ir 1r_max only for Adam W and use it also for AdmIRE, for which the IRE is activated at the end of warm-up phase.
4.2.1 Computational efficiency and hyperparameter robustness
The first experiment conducted to verify both the computational efficiency and the robustness of hyperparameters (y, k) in IRE for pre-training tasks. Specifically, we train 2-layer decoder-only Transformer (8M) on the Wikitext-2 dataset (4.3M) (Merity et al., 2016) by Adam and AdmIRE (with K 10 and varying 7, K). The total training duration is 100k steps, including a 3k-step warm-up phase.
First, we tune 1r_max in Adam W, identifying the optimal 1r_ sr_max=6e-4. Subsequently, we train both AdamW and Ad- mIRE using this 1r ir_max.
Table 5: Wall-clock time 1 A800.
Algorithm AdamW
time (/step)
Computational efficiency. As shown in Table 5, AdmIRE with K = 10 (estimating the projection mask every 10 steps) computationally efficient: the average time per step of Ad- mIRE is only 1.12 times that of AdamW, corresponding to the theoretical estimation (1.1 times).
0.165s
AdmIRE
0.185s
Robustness to hyperparameters. Figure 3 shows that Ad- mIRE, with varying y and K, consistently speeds up the pre- training. Remarkably, with the best configuration, AdmIRE can achieves 5.4x speedup than well-tuned AdamW.
training loss
5.4x Speedup
Ade Re
Ade Say 1.01
More experimental details and results are deferred to Ap- pendix C.
25k
50k
100k
Num of steps
Figure 3: Transformer on wikitext-2.
4.2.2 Experiments on Llama models
Llama (Touvron et al., 2023), a popular open LLM, exhibits remarkable capabilities across general domains. In this section, examine the performance of AdmIRE in training Llama models of various sizes across various datasets:
Llama (60M) wikitext-103 (0.5G). Wikitext-108 (Merity et al., 2016) serves as a standard language modeling benchmark for pre-training, which contains 103M training tokens from 28K articles, with an average length of 3.6K tokens per article.
Llama (119M) on minipile (6G). Minipile (Kaddour, 2023), 6GB subset of the deduplicated Pile (825GB) (Gao et al., 2020) presents a highly diverse text corpus. Given its diversity, training on minipile poses more challenges and potential instabilities for optimizers compared to Wikitext-103.
Llama (229M) openwebtext (38G). Openwebtext (Gokaslan and Cohen, 2019), an open- source recreation of the WebText corpus, is extensively utilized for LLM pre-training such as RoBERTa (Liu et al., 2019b) and GPT-2 (Radford et al., 2019).
Additionally, gradient clipping is adopted maintain the training stability (Pascanu et al., 2012).
First, we tune Ar_max in AdamW for each of the three experiments, separately. The optimal 1r max
7

---

identified for these three experiments is all 6e-4. Then, both AdamW and AdmIRE are trained using this optimal 1r ___max. For more details, please refer to Appendix C.
AdmIRE is 2x faster than AdamW. The results are reported in Figure 4. We can see that AdmIRE consistently achieves a 2.1 speedup compared with well-tuned A for all three cases.
Llama (60M) wikitext-100
Llama (119M) on minipile
Llama (229M) on openwebtext
validation locs
validaaaa locs
Adanth (o0)
validationn loss
Adriii D30t 6.7
Adr 011 22)
1
num of Steps
num of Steps
num of steps
Figure 4: AdmIRE outperforms AdamW in the pre-train pre-training of Llama models.
Notice that the primary motivation behind IRE is to speed up the sharpness reduction, which only requires to increase learning rate along completely flat (zero-curvature) directions. However, practical implementation may also increase the learning rate along directions with small but non-zero curvatures, which can further speed up loss convergence. A thorough explanation for the significant acceleration provided by this approach is left for future research.
We further assess the sharpness reduction capability of IRE for LLM pre- pre-trainnn training. Specifically, we com- pare the sharpness of solutions, Tr( (()),, found by Adam W/AdmIRE during pre-training of Llama (60M) on wiki-103 dataset (corresponding to Figure 4 (left)). The results shown in Table 6 demonstrate that AdamIRE not only achieves the same loss only half the iterations required by AdamW, but also the solutions found by AdmIRE are significantly flatter than that found by AdamW.
Table 6: Comparison of the sharpness of the solu- tions found by Adam W.dmIR..
Adam
AdmIRE
training steps
100k
50k
final (())
2.47
2.47
final (T(V"C(e)
120.41
88.86
Recently, Liu al. (2023) revealed a strong correlation between the sharpness and downstream task
performance, suggesting that for models with the same pre-training loss, flatter solutions yield better performance on downstream tasks. Based on this, we hypothesize that the solutions found by IRE may also have better performance in downstream tasks, which we leave to future work.
5 Theoretical Guarantees for IRE on SAMs
Both empirical (Foret et al., 2021) and theoretical (Wen al., 2023a) studies have validated that SAM algorithms exhibit superior sharpness regularization compared to (S)GD. In this section, we provide theoretical analysis demonstrating that IRE can further enhance the sharpness regularization of SAM algorithms substantially.
5.1 Theoretical setups
Recall that L()) = + (0) denote the total loss, where (.) is the loss on the i-th data.
Without loss of generality, we assume ming ((0) = O. We further make the following assumption:
Assumption 5.1 (Manifold of minimizers). Assume that L E (RPP), M == arg ming L()) is (p- - m)-dim C -usummmaiioold in R2 for some m E [p], and rank ((0)) = m for any O E M.
The above connectivity assumption on the manifold of minimizers M has been empirically veri- fied in works such as Draxler et al. (2018) and Garipov et al. (2018), and theoretically supported in Cooper (2018). This assumption is also widely used the theoretical analysis of implicit regular- ization (Fehrman al., 2020; Li et al., 2022; Arora al., 2022; Wen et al., 2023a).
Besides, we introduce the following definitions to characterize the dynamics of gradient flow (GF) near the minima manifold M, which is also used in the related works above.
8

---

Definition 5.2 (Limiting map of GF), Consider the GF: do())
V (()(t)) starting from ( = 0.
dt
Denote by (()) := lim (t) the limiting map of this GF.
Definition 5.3 (Attraction set of M). Let U be the attraction set of .M under GF, i.e., GF starting in U converges some point in M. Formally, U == (0 E RR (0))= M).
As proven in (Arora et al., 2022, Lemma B.15), Assumption 5.1 ensures that U (in Definition 5.3) is open and (-) (in Definition 5.2) C² on U.
5.2 Theoretical results
The stochastic SAM (Foret et al., 2021) is given by
V (O.) standard SAM: 0.+1 0. L,,, 0, [VL,, where it U(n]])
The generalization capability of standard SAM can be bounded by the average sharpness, _ave pÉ/IlElI) (Foret et al., 2021). This leads researchers to also explore average
(0 1 SAM (Wen et al., 2023a; Zhu al., 2023; Ujváry et al., 2022), which minimizes Lave;
average SAM: 0,+1 = 0 ( + p where & IO.,.)).
(8)
Two-phase algorithms. Our theoretical focus is on how IRE accelerates the sharpness reduction of SAM (7) and (8) near the minima manifold M. Thus, we analyze the two-phase algorithms.
Specifically, let the initialization Og E U. In Phase I (L & Ti), we employ GF L(,,) to ensure that the loss decreases sufficiently; then in Phase II (Ti < & Ti TT = T), we incorporate IRE into the standard / average SAM.
or
Effective dynamics: sharpness regularization. The implicit reg- ularization of SAMs can be modeled using effective dynamies. In Phase II, 0, are close the manifold of minimizers M and let 21 : (0,) E M. Then, the effective dynamics is given by {2 revealing how SAMs explore the manifold of mini- mizers M. Particularly, Wen et al. (2023a) showed that the effective dynamics of standard/average SAM are both
((2))//
tangent space T.M
VII VI
(O.)
minima manifold M
E [ = Nen V Tr [(z)//2 + o(Te)))
(9)
which minimizes the trace of Hessian on .M. The difference between the standard SAM (7) and average SAM (8) lies in the effective learning rate (LR) nett's 's visual illustration of some quantities in (9) is provided in the figure above.
Summary of our theoretical results. In this section, show that incorporating IRE into SAMs can significantly increase the effective LR Nett in (9) while maintaining the same training stability as SAMs. In Table we present the effective LR for SAMs and the SAM-IRES. We see clearly that IRE can accelerate the sharpness reduction by a non-trivial factor for both standard and average SAM.
Table 7: Comparison of the implicit regularization strength of SAMs w/o IRE.
Algorithm
Effective LR: Tefff
average SAM (8)
n /p (Thm 5.5)
IRE average SAM (8)
1/.5 (Thm 5.5)
standard SAM (7)
Tpp (Thm 5.6; Wen et al. (2023a))
IRE+ + standard SAM (7)
no (Thm 5.6)
Remark 5.4 (The mechanism of IRE's success). The success of SAM-IRE follows the same mecha- nism illustrated in Section 2. The key fact that IRE only increases the LR along flat directions has two implications: 1) It does not change the trend of implicit regularization in Eq. (9) but accelerates SAMs' effective dynamics by a factor of (1 + k); 2) Since the LR is only increased along flat directions, can be set substantially large without hurting the training stability, because the dynamics in sharp directions remain unchanged. Specifically, we theoretically justify in SAM-IRE, & can be selected as large 1/p.

---

5.2.1 IRE on average SAM: An (01//n0.5) acceleration
We first consider IRE average SAM. Let Ti be the hitting time: Ti := in > 0: 1|0; $
O(vnp)). When running GF Phase I, Definition 5.3 guarantees < co. Thus, the starting of Phase II, |00, ((0.) (T)))) Furthermore, the following result holds for Phase II.
Theorem 5.5 (IRE on average SAM). Suppose Assumption 5.1 holds. If = ((1) and p ((vT)
in SAM (8). K & 1/p. and Pi P pm ( (V))( in IRE (4), then with high probability at least 1 1-7i exp (1/ (1) -pp p))) 110- ((0))1l ((vūp) holds for all Ti < t T. Furthermore,
the effective dynamics of 2 == (00,) E M satisfies:
E [ = (1 + x)m0²
V Tr (2 - 0
P
Note that p ((v and can be large 1/p. Consequently, the effective LR of minimizing the trace of Hessian can be selected as large as Nell (x + 1)mpmm /p O17) 5 /p). contrast, that of average SAM is at most (n /p). The proof of Theorem 5.5 can be found in Appendix D.
5.2.2 IRE standard SAM: An ((1/p) acceleration
This subsection delves into IRE standard SAM (7), which is more widely used and often yields better performance than average SAM (8). However, since standard SAM (7) requires stochastic gradients (0) (i [n]). need additional assumption regarding the features on the manifold (see Setting E. which is commonly used in the literature (Du et al., 2018; 2019; Li et al., 2022; Arora et al., 2022; Wen et al., 2023a). We defer it to Appendix E due space constraints. Under this Setting, Assumption 5.1 holds naturally with m = n.
During Phase I of GF, Definition 5.3 ensures that there exists ( < xc such that |10, +(0.) On p) for any a E [0, 1). We define Ti as the hitting time: Ti := inf > 0: 100, 2b (0,) ||
Oln a *)}. Then the following result holds for Phase II, whose proof can be founded in Appendix E.
Theorem 5.6 (IRE on standard SAM). Under Setting E.I, if 1). P 0(1) in SAM (7), K < 1/p. and P. P (V)((0)) in IRE (4), then with high probability least 1 Tń exp (-1(1/n)))
110 - =(0))ll = O(1²~ap holds for all Ti & t & T. Moreover, the effective dynamics &t
(0.) E M satisfies:
[z+11] = 2: (1 + k) m) M / Tr V [(z))/22 + + ((x + 1-1)mppp n/--))
Taking 6 = 0 and a 0 recovers the result established in Wen et al. (2023a). However, A can be as large as 1/p. where IRE provides a (11/p) acceleration over the standard SAM.
Conclusion
In this work, we propose a novel IRE framework to enhance the implicit sharpness regularization of base optimizers. Experiments demonstrate that IRE only consistently improves generalization but also accelerates loss convergence in the pre-training of Llama models of various sizes. The code is available at https: //gtthub/ com/ wtom////////////////..t///m (0nz9/IRE-aae L-RE-allore tew-tthh-faareee wotbmffrrewook......
For future work, there are two urgent directions: 1) understanding why IRE can accelerate conver- gence, which may require studying the interplay between IRE and the Edge of Stability (EoS) (Wu et al., 2018; Jastrzębski et al., 2017; Cohen et al., 2021); and 2) conducting a larger-scale investi- gation into the acceleration of AdmIRE compared to AdamW in LLM pre-training, as well as the downstream performance of the I.LMs pre-trained by AdmIRE.
Acknowledgments
Lei Wu is supported by the National Key R&D Program of China (No. 2022YFA1008200 and National Natural Science Foundation of China (No. 12288101). Mingze Wang supported in part by the National Key Basic Research Program of China (No. 2015 2015CB856000). We thank Dr. Hongkang Yang, Liu Ziyin, Liming Liu, Zehao Lin, Hao Wu, and Kai Chen for helpful discussions and anonymous reviewers for their valuable suggestions.
10

---

References
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In International Conference on Machine Learning, pages 948-1024.
PMLR, 2022. 8, 9, 10, 21, 22, 23
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an Eestein-Uhleeneek like process. In Conference on learning theory, pages 483-513. PMLR, 2020. 2
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in Neural Information Processing Systems, 36, 2024. 3, 20
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, pages 1305-1338. PMLR, 2020. 17
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. International Conference on Learning Representations, 2021. 2, 10
Yaim Cooper. The loss landscape of overparameterized neural networks. arXiv preprint arXiv: 1804. 10200, 2018. 8
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe Le, and Ruslan Salakhutdi- nov. Transformer-XI: Attentive language models beyond fixed-length context. arXiv preprint arXiv:1901.02860, 2019. 21
Alex Damian, Tengyu Ma. and Jason D Lee. Label noise SGD provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449-27461, 2021. 2. 18
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition,
pages 248-255. leee, 2009. 6
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 19999, 2020. 2, 6
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pages 1309-1318 PMLR, 2018. 8
Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent YF Tan. Efficient sharpness-aware minimization for improved training of neural networks. arXiv preprint arXiv:21 10.03/ 141, 2021. 2
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pages 1675-1685. PMLR, 2019. 10
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.00000 2018. 10
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011. 2
Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (S)GD over diagonal linear networks: Implicit regularisation, large stepsizes and edge of stability. arXiv preprint arXiv:2302.08982, 2023, 18
11

---

Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (S) GD over diagonal linear networks: Implicit bias, large stepsizes and edge of stability. Advances Neural Information Processing Systems, 36, 2024. 2
Benjamin Fehrman, Benjamin Gess, and Arnuff Jentzen. Convergence rates for the stochastic gradient descent method for non-convex objective functions. Journal of Machine Learning Research, 21, 2020. 8
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. International Conference on Learning Representations, 2021. 2, 19, 20
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00022
2020.
Timur Garipov, Pavel Izmailov, Dmitri Podoprikhin, Dmitry P Vetrov, and Andrew Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information processing systems, 31, 2018.
Khashayar Gatmiry, Zhiyuan Li, Ching- Chuang, Sashank Reddi, Tengyu Ma, and Stefanie Jegelka. The inductive bias of flatness regularization for deep matrix factorization. arXiv preprint arXiv:2306.1322 13239, 2023. 2
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a Kronecker-factoree eigenbasis. Advances in Neural Information Processing Systems, 31, 2018. 5
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. httpp://kkylion007,, wwt.rr...ttu........ io/
OpeniebTexttCoppus, 2019. 7
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in hour. arXiv preprint arXiv: 1706.0277 2017. 2
Roger Grosse and James Martens. A Kronecker-factored approximate Fisher matrix for convolution layers. International Conference on Machine Learning, pages 573-582. PMLR, 2016. 5
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018. 17
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016. 2,
Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung-Woo Woo Ha. Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights. arXiv preprint Stt::20006082.17. 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42 1997.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the general- ization gap in large batch training of neural networks. arXiv preprint arXiv:1705.08741, 2017.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. arXiv preprint 1711:0711 2017. 2, 10
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. International Conference on Learning Representations, 2019a.
Ziwei and Matus Telgarsky. Risk and parameter convergence of logistic regression. Conference on Learning Theory, 2019b. 17
12

---

Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33:17176-17186 2020. 17
Ziwei Ji and Matus Telgarsky. Characterizing the implicit bias via primal-dual analysis. In
Algoritumic Learning Theory, pages 774-804. PMLR, 2021. 17
Ziwei Ji, Miroslav Dudík, Robert E Schapire, and Matus Telgarsky. Gradient descent follows the regularization path for general losses. In Conference on Learning Theory, pages 2109-2136. PMLR, 2020. 17
Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration. In International Conference Machine Learning, pages 4860-4869. PMLR, 2021. 17
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. International Conference on Learning Representations, 2019. 2
Jean Kaddour. The MiniPile challenge for data-efficient language models. arXiv preprint
arXiv:2304.08442, 2023." 7
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2016.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint
141.69041,.6980, 2014. 2
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images, 2009. URL https://www cs toronto. edu/"kriz/cifar cifar. html.
Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi- honogenooo neural networks. International Conference on Learning Representations, 2023. 17
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. ASAM: Adaptive sharpness- aware minimization for inveriinvrriant learning of deep neural networks. In International Confer- ence on Machine Learning, pages 5905-5914. PMLR, 2021. 2.20
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pages 2101-2110 10. PMLR, 2017. 18
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms Mathematical foundations. The Journal of Machine Learning Research, 20 (1):1474-1520, 2019. 18
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss?-a mathematical framework. International Conference on Learning Representations, 2022. 2, 4, 8, 10, 18
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. In International Conference on Machine Learning. pages 22188 22188-22214. PMLR, 2023. 8
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. International Conference on Learning Representations, 2024. 5, 21
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance the adaptive learning rate and beyond. arXiv preprint 1xXiv:1908.0320 2019a. 3
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11999 1169, 2019b.
13

---

Yong Liu, Siqi Mai. Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-ar awwpe minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12360-12370, 2022. 2
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 10X1v:1711.05101, 2017.
Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint 100.v:190,.08843, 2019. 3
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv: 0906009069999 2019. 17
Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34, 2021. 17
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35: 14689-34708, 2022. 2
Chao Ma and Lexing Ying. On linear stability of SGD and imput-smoothness of neural networks. Advances in Neural Information Processing Systems, 34: 14:16805-1681 2021. 2
Chao Ma, Daniel Kunin, Lei Wu, and Lexing Ying. Beyond the quadratic approximation: The multiscale structure of neural network loss landscapes. Journal of Machine Leurning, 1(3): 247-267, 2022. 2, 4, 18
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning, pages 2408 2417. PMLR, 2015. 5
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint (60.:16609008443 2016.
Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness-awan aware minimization stronger: A sparsified perturbation approach. Advances in Neural Information Processing Systems, 35:30950-30962. 2022. 2.
Maximilian Mueller, Tiffany Vlaar, David Rolnick, and Matthias Hein. Normalization layers are all that sharpness-aware minimization needs. Advances in Neural Information Processing Systems, 36, 2024. 2, 6, 20
Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of minima stability: A view from function space. Advances in Neural Information Processing Systems, 34: 177491 2021.
Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro. and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models. In International Conference on Machine Learning, pages 4683-4692 PMLR, 2019a. 17
Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. The 22nd International Conference on Artificial Intelligence and Statistics, pages 1420-34288 PMLR, 2019b.
17
Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with fixed learning rate. In The 22nd International Conference Artificial Intelligence and Statistics, pages 3051-3059. PMLR, 2019c. 17
Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry. Implicit bias of the step size in linear diagonal neural networks. In International Conference on Machine Learning, pages 16270-16295. PMLR, 2022. 18
Yurii Nesterov. A method of solving convex programming problem with convergence rate o(1/22). Doklady Akademii Nauk SSSR, 269(3):543, 1983. 2
14

---

Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization deep learning. arXiv preprint arXiv:1412.6614 2014. 1
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. Understanding the exploding gradient problem. CoRR, abx/1211.5063 2(417):1, 2012.
Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. Advances in Neural Information Processing Systems, 2023. 18
Scott Pesme, Loucas Pillaud- Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of SGD for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:2921 1-22920,, 2021. 18
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. 7
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407, 1951.
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. Nature, 323(6088):533-536, 1986.
Noam Shazeer. GLU variants improve transformer. arXiv preprint arXiv:2002.05202 2020. 21
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learming Research, 19(1): 2822-2878, 2018 17
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024. 21
llya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance initialization and momentum in deep learning. In International conference machine learning, pages 1139- 1147. PMLR, 2013. 2
Tijmen Tieleman. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26, 2012.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. LLaMA: Open and efficient foundation language models. arXiv preprint arXiv:2302. 139771112,1111,, 2023. 2, 7, 20, 21
Szilvia Ujváry, Zsigmond Telek, Anna Kerekes, Anna Mészáros, and Ferene Huszár. Rethinking sharpness-awan minimization variational inference. arXiv preprint arXiv:2210. 10452, 2022.
Gal Vardi. On the implicit bias in deep-learning algorithms. Communications of the ACM, 66(6):
86-93, 2023. 2, 17
Gal Vardi, Ohad Shamir, and Nati Srebro. On margin maximization in linear and ReLU networks. Advances in Neural Information Processing Systems, 35:37024-37036, 2022. 17
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017. 20
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018. 38
Guanghui Wang. Rafael Hanashiro, Etash Guha, and Jacob Abernethy. On accelerated perceptrons and beyond. arXiv preprint arXiv:2210.09377, 2022. 17
Mingze Wang and Chao Ma. Understanding multi-phase optimization dynamics and rich nonlinear behaviors of ReLU networks. Advances in Neural Information Processing Systems, 2023. 17
15

---

Mingze Wang and Lei Wu. The noise geometry of stochastic gradient descent: A quantitative and analytical characterization. arXiv preprint arXiv:2310.00692, 2023. 2
Mingze Wang, Zeping Min, and Lei Wu. Achieving margin maximization exponentially fast via progressive norm rescaling, International Conference on Machine Learning. 2024. 17
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-awars minimization minimizes sharpness? In The Eleventh International Conference on Learning Representations, 2023a. 4, 8, 9, 10, 32, 33, 34
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization. Advances in Neural Information Processing Systems, 2023b.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673 PMLR, 2020. 18
Lei Wu and Weijie J Su. The implicit regularization of dynamical stability in stochastic gradient descent. In The 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 37656-37684, PMLR, 2023. 2
Lei Wu, Zhanxing Zhu, and Weinan E. Towards understanding generalization of deep learning: Perspective of loss landscapes. arXiv preprini arXiv: 10239, 2017.
Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over- per-parameteriiee learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31:8279- 8288, 2018. 2, 10
Lei Wu, Mingze Wang, and Weijie J Su. The alignment property of SGD noise and how it helps select flat minima: A stability analysis. Advances in Neural Information Processing Systems, 35: 4680-4693, 2022. 2
Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. arXiv preprint 20Xiv:2208.06677, 2022a.
Zeke Xie, Xinrui Wang, Huishuai Zhang. Issci Sato, and Masashi Sugiyama. Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In International conference on machine learming, pages 24430-24459, PMLR, 2022b. 3
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the Hessian. In 2020 IEEE international conference on big data (Big data), pages 581-590. IEEE, 2020. 4
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint 17X1v:1605.07146. 2016. 5,6 6
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019. 21
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017. 1
Tongtian Zhu, Fengxiang He, Kaixuan Chen, Mingli Song, and Dacheng Tao. Decentralized SGD and average-direction SAM are asymptotically equivalent. In International Conference Machine Learning, pages 43005-43036. PMLR, 2023. 9
Juntang Zhuang, Tommy Tang, Yifan Ding, Sckhar C Tatikonda, Nicha Dvornek, Xenophon Pa- pademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 10188899118888 2020.
16

---

Appendix
Other Related Works
17
Proofs for SGD Section
18
Experimental Details
19
C.1 Experimental details in Section 4.1
19
C.2 Experimental details in Section 4.2
20
Proofs in Section 5.2.1
D.I Preliminary Lemmas
21
D.2 Proof of Theorem 5.5
24
Proofs in Section 5.2.2
32
E.1 Preliminary Lemmas
32
E.2 Proof of Theorem 5.6
34
F Useful Inequalities
37
A Other Related Works
Beyond the implicit sharpness regularization, numerous other attempts have explored implicit biases in deep learning algorithms Vardi (2023). Among these, a popular research line is the max-margin bias, which suggests that optimizers favor the solutions with large margin, which generalizes well. Soudry et al. (2018 8) showed that GD converges to max-margin solutions under exponentially-tailed loss linearly separable data, albeit with extremely slow rate ( log I). Furthermore, Nacson et al. (2019c) studied this bias for SGD, Ji and Telgarsky (2019b) investigated linearly non-separable data, and Ji al. (2020) analyzed the effects of the tail behavior of loss functions.
Other Implicit Biases.
To achieve faster margin maximization, Nacson et al. (2019b); Ji and Telgarsky (2021) demonstrated that GD with aggressively loss-scaled step sizes can achieve a faster polynomial rate of (1/t). Building this, Ji et al. (2021); Wang al. (2022) proposed momentum-base based gradient methods which achieve a rate of (1)tt). Recently, Wang et al. (2024) established that the polynomial rates for most previous algorithms tight, and proposed progressive rescaling gradient descent method that achieves margin maximization exponentially fast (Oe-R(t)) -R(t)..
The margin- maximization bias has also been studied for nonlinear models. and Telgarsky (2019a); Gunasekar et al. (2018) 8) examined deep linear networks, while Chizat and Bach (2020) focused on wide two-layer ReLU networks. Notably, Nacson al. (2019a); Lyu and Li (2019); Ji and Telgarsky (2020) demonstrated that for general homogeneous networks, Gradient Flow (GF) and GD converge to solutions corresponding the KKT point of the max-margin problem. Recently, Kunin et al. (2023) extended this analysis to honsgihomogeneous networks. For two-layer (leaky-)Rel.U neural networks, Lyu al. (2021); Vardi et al. (2022); Wang and Ma (2023) examined whether the convergent KKT point of GF correspond to global optima of the max-margin problem.
17

---

Future work could investigate whether the IRE framework can also enhance the margin maximization bias, although it is primarily designed for enhancing the implicit sharpness regularization.
Additionally, Woodworth et al. (2020); Pesme et al. (2021); Nacson et al. (2022); Pesme and Flam- marion (2023); Even et al. (2023) conducted fine-grained analyses of training dynamics, examining how initialization and step size impact (SJGD's minima selection in linear diagonal networks.
B Proofs for SGD in Section 2
For the example in Section 2, we have studied the implicit sharpness regularization of GD dynamics and how IRE enhances the implicit regularization of GD. In this Section, we illustrate that, for this example, similar results hold for SGD dynamics.
SDE Modelling of SGD. We consider SGD approximated by SDE (Li et al., 2017; 2019; 2022) with noise covariance Σ: d0, = - ( dt + V/ We consider that the noise covariance
aligns with the Hessian near minima, i.c., E()) = O0 0 as the label noise (Damian et al., 2021; Li et al., 2022). Then, the SDE above can be rewritten as d
o²h(u)) (where o >> 0 is scalar), such
(/Vh(u))/2 dt (VnoPh(u.)dw.)) (r(((u))
0
(10)
Implicit Sharpness Regularization. Intuitively, when i/ close O the speed of u, is much slower than et due to oz & vr. Following Ma et al. (2022), when this speed separation is large, d always equilibrium given u,. Solving the Ornstein Uhlenbeck Ornstein-Ullenbeckk process about vt, we know the equilibrium distribution of vg is U.a (O,, and hence the dynamics u, (along the manifold) is
du//dt = Evr 3 vh(u:)/2] Vh2uu)/2.
(1I)
This derivation clearly shows the slow "effective dynamics" u, along the manifold is a gradi- ent flow minimizing the sharpness h(-). When SGD minimizes the loss, it also minimzes the sharpness implicitly, that is to say, SGD has the following implicit sharpness regularization: ming Tr(v=((0) s.t. 0.
Generalization and Optimization benefits of the sharpness regularization. In terms of generalization, as discussed in related works, a common view is that flat minima generalize well, which has been proved in a large number of previous works. In addition, in terms optimization, after SGD reaches the equilibrium too ~ (0, n the loss near the flat minimum is smaller because Ee [L(0)] Ev,. [h(u)) /22 ho²h(u)/2 L(u))
Q. How can we enhance the implicit sharpness regularization of SGD?
A. Accelerating SGD's 's slow "effective dynamics" u, along the minima manifold.
Implicit Regularization Enhancement (IRE) by accelerating the effective dynamics along minima manifold. First, it is worth noting that naively increasing the learning rate cannot achieve our aim, because increasing 7) will influence the dynamic stability of 0ú and the equilibrium V Therefore, we need to accelerate the effective dynamics u, without affecting the dynamics of v;. Another main point is that SGD's 's effective dynamics u, can naturally minimize the sharpness implicitly, so we only need to enhance this property. To achieve this, only need to correct the update direction in (10) from - ((e,)) to ((e) + kPM (L().)) where PM is the projection matrix to the manifold M and is a scalar. Using this new algorithm, the SDE dynamics corresponds to
d ---((1+ k)) Vh(u))())) at (Vno=h(u.Jaw.)) vch(us))
-
(12)
Comparing (10) and (12), the dynamics of ot are the same, so they attain the same equilibrium distribution U.e NN As for the effective dynamics along manifold, (10) corresponds to the form:
du:/dt - -E [11 + 2 Vh(u:)/22 - -(1+ //,u))/2
(13)
18

---

(1+k) times Enhancement. Comparing (13) and (11), it is clear that our new algorithm can enhance
implicit sharpness regularization (1 + x) times faster than the original SGD.
C Experimental Details
This section describes the experimental details in Section 4.
C.1 Experimental details in Section 4.1
C.1.1 Experimental details in Section 4.1.1
We train WideResNet-1668 16-8 on CIFAR-10 dataset by SAM-IRE (with 10, varying x and 7)- The experiments employ basic data augmentations and 0.1 label smoothing, as outlined by Foret et al. (2021). The mini-batch size is set to 128, the weight decay is set to 5c-4, and the in SAM is to 0.05, in Foret et al. (2021). To evaluate the implicit flatness regularization of SAM itself, the momentum is set to 0.O. Regarding the learning rate (Ir), we evaluate for both constant (within our theoretical framework) and decayed Ir (common in practice though not covered by our theory). In the experiment in Fig 2 (a), a fixed Ir 0.1 is used. In the experiment in Fig 2 (d), a step-decayed schedule is employed, starting 0.1 and reducing Ir by factor of 5 at epoch 20, 50, 80. We transit from SAM to SAM-IRE at the 30th epoch out of 100 total epochs. We test in 0.8, 0.9, 0.95, and k in (original SAM), 2, 5, 10.
The flatness measure, Tr Tr))),(()), is approximated by the trace of Fisher (O)) Specifically, we use diag Fea (f)) in (6) for the estimate because Eyldiag(Fatt (0)) = (d (())) and thus, (0))) Ey (diag Pan (0)) Moreover, the first ">" above takes "=" when ((0) = O.
Tr Tr(V²L(0)) (T( = ET
In this section, all experiments were conducted using a single A800 GPU.
C.1.2 Experimental details in Section 4.1.2
We first evaluate the impact of IRE on general- ization of baseline models (WideRes.Nt-225 10 and ResNet-56) and default optimizers (SGD and For the base optimizers, SGD and SAM, cosine learning rate decay is
Experiments for CNNs on CIFAR-1O/CIER-
SAM) on CIFAR-(10. adopted with an initial Ir 0.1. For other training components, (2021): basic data augmentations and 0.1 label smoothing: for both SGD and SAM, the momentum
follow the settings in Foret et al. set 0.9, the batch size set to 128, and the weight decay to 0.05 for CIFAR-I0 and 0.1 for CIFAR-100 The total epochs is set to 100 for CIFAR-10 and 200 for CIFAR- 100, and we switch from SGD/SAM to SGD-IRE/SAM-IRS when the training loss approaches 0.1. For SGD-IRESSAM-IRE, we fix K = 10, and tune hyperparameters 7 and k via a grid search over y E {0.99, 0.9, 0.8) and K E {1, 2). The results are reported in Table 1.
5e-4; for SAM, p set
Experiments without finely tuned hyperparameters. A high sensitivity to the choice of hyperpa- rameters would make method less practical. To demonstrate that our IRE performs even when K. ? are not finely uned, we conduct experiments using fixed = 0.99, Λ = 1, under the same settings as described above.. The results (averaged over random seeds) reported in Table 8.
Table 8: Results for SGD-IRE and SAM-IRE on (WideRtssNet288-10 ResNet-56) on CIFAR- (10, 100), using fixed y 0.99, K in IRF.
WideResNet-28-10
ResNet-56
CIFAR-10
CIFAR-100
CIFAR-10
LIFAR-100
SGD
95.93
80.77
93.80
72.72
SGD-IRE
96.13 (+0.20)
81.12 (+0.35)
14.94(+0.14)
72.93(+0.21)
SAM
96.73
83.22
94.58
75.25
SAM-IRE IRE
96.75 (+0.02)
83.40 (+0.19)
94.65 (+0.07)
75.49 (+0.24)
19

---

Experiments for CNNs on ImageNet. We also examine the impact of IRE on generalization of ResNet-50 and default optimizers (SGD and SAM) on on ImageNet. Following Foret et al. (2021) and Kwon et al. (2021), we use basic data augmentations and 0.1 label smoothing. For the base optimizers, SGD and SAM, we also follow the settings in Kwon et al. (2021): the momentum is set to 0.9; cosine learning rate decay is adopted with an initial Ir 0.2; the batch size is set to 1024; the weight decay is set to 1e-4; for SAM, p is set to 0.05, The total epochs is set to 200, and we switch from SGD/SAM to SGD-IRE/SAM-IRE when the training loss approaches 1.5. For SGD-IRE/SAM-IRE we fix K = 10, and tune hyperparameters Y and K via a grid search over 7 E {0.8, 0.6) and K E {2, 4). The results are reported in Table 2.
Experiments for ViTs on CIFAR-100. We examine the impact of IRE on generalization of ViT-T and ViT-S on CIFAR-100. We follow the settings in Mueller et al. (2024): the default optimizers used are AdamW and SAM, with cosine Ir decay to 0 starting from an initial le-4; the weight decay is 5e-4; batch size is 64; strong data augmentations (basic + Automen))ett) are utilized; 0.1 for SAM. The total epochs are set to 200, and we switch from Adam W/SAM to AdmIRE/SAM-IRE when the training loss approaches 0.5. For AdmIRE/SAM-IRE, we fix K 10, and tune hyperparameters Y and x via a grid search over r E {0.99, 0.9, 3.8} and 6 E {00, 50}. The results are reported in Table 3.
Experiments for ViTs on ImageNet. We also examine the impact of IRE on generalization of ViT-S/16 on ImageNet. We follow the settings in Chen et al. (2024); RandAugment and Mixup with
= 0.5 are utilized; the default optimizer used is AdamW with hyperparameters B1 = 0.9, 82 = 0.999 and weight decay x= = 1.0; the learning rate strategy integrates a warm-up phase followed by a cosine decay scheduler with 1r__la______ ___ax=3e-3;; batch size 4096; the total training duration 300 epochs, including 30 warm-up epochs; For AdmIRE, we switch from Adam to AdmIRE at epoch 100 and examine different 7, K. The results are reported in Table 4.
In this section, the experiments on and the experiments on ImageNet were conducted using 4 A800 GPUs.
were conducted using a single A800 GPU,
C.2 Experimental details in Section 4.2
C.2.1 Experimental details in Section 4.2.1
We train 2-layer decoder-only Transformer (8M parameters) using absolute positional encod- ings (Vaswani et al., 2017), with 8 heads in each layer and a hidden size of 128, on the wikitext-2 dataset (4.3M) by AdamW and AdmIRE (with K 10 and varying 7, x). The (max) sequence length is set to 512, and the batch size is set to 32. The experiments in this section are conducted on A800.
For the optimizer AdamW, we use the hyperparameters 81 0.9, B2 0.95 and weight decay = 0.1, as suggested in Touvron et al. (2023). The total training duration is 100,000 steps, including 3,000 warm-up steps followed by a cosine decay scheduler with 1r_max and 1r 1r_min=1r_max/3
First, we tune 1r 2r_max Adam W from {1. (1.5e-4. 3e-4, 6e-4, 1. 2e-3, 8. 3e-3). The results, shown in Figure 5 (left), identify the optimal 1r max==ax--44- -4. We also use the optimal 1r__aa___ ___ of 6e-4 for AdmIRE, for which the IRE is enable at the end of warm-up phase.
The results are reported in Figure 3.
Transformer (8M) on wikitext-2
Liama (00M) on wikitest-103
30
1z maamm Ser4
20. m.3m3.33mmmmm 0e-4
1t muc De-4
muc
ma..aamm... 0e-4
it..mael 2o
traising loss
2n. 2e
volidation loos
2e-3
Je. m Ownd
Num
steps
num steps
Figure 5: The results for tuning ar_max in Adam W.
20

---

C.2.2 Experimental details in Section 4.2.2
Llama (Touvron et al., 2023) a decode-only Transformer architecture using Rotary Positional Encoding (RoPE) (Su et al., 2024), Swish-Gated Linear Unit (SwiGLU) (Shazeer, 2020), and Root mean square layer normalization (RMSNorm) (Zhang and Sennrich, 2019). The experiments in this section examine the performance of AdmIRE in training Llama models with different sizes. For implementation, we utilize the Llama code available on huggingface. The experiments are conducted on 4 H800.
For the optimizer AdamW, we use the well-tuned hyperparameters 31 0.9, 3 0.95 and weight decay À = 0.1 for LLama (Touvron al., 2023). The learning rate strategy integrates a warm-up phase followed by cosine decay scheduler with 1r_max and 1r frr_m==lr_maas _ma/200. Additionally, is used with gradient clipping 1.0 to maintain the training stability.
In each experiment, we tune the optimal Ir ar_max for AdamW and then use it also for AdmIRE, for which the IRE is enable at the end of warm-up phase.
Llama (60M)
wikitext-103 (0.5G). We train 16-layer Llama model, with 10 heads in each layer and a hidden size of 410, on the wikitext- 103 dataset. The (max) sequence length set to 150, and the batch size is set to 240, following Dai et al. (2019). The total training duration is 50,000 or 100,000 steps, including 500 warm-up steps. First, we tune 1r 1r_max in Adam from (3e-4, 6e-4, 1. 1.2e-3,
max
1. identifying the optimal 1r_max 6e-4. The experimental results are very similar to Figure 5 (right), so we will not show them again. Then, both Adam W and AdmIRE are trained using the optimal 1r_max.
Llama (119M) on minipile (6G). We train a 6-layer Llama model, with 12 heads in each layer and a hidden size of 768, the minipile dataset. The (max) sequence length is set to 512, and the batch size is set to 300. The total training duration is 30,000 or 60,000 steps, including 300 warm-up steps. First, we tune 1r_ 1r_max in AdamW from (3e-4, 6e-4, 2e-3. 1. 3..-3), identifying the optimal Ir_max 6e-4. (The results are very similar to Figure 5 (right), so we do not show them repeatly.) Then, both AdamW and AdmIRE are trained using the optimal 1r 1r_max.
Llama (229M) on openwebtext (38G). We train a 16-layer Llama model, with 12 heads in each layer and hidden size of 768, on the openwebtext dataset. The (max) sequence length is set to 1024, and the batch size is set to 480, following nanoGPT and Liu et al. (2024). The total training duration is 50,000 100,000 steps, including 1,000 warm-up steps. First, we tune 1r__ma___ __max in AdamW from (3e-4, 6e-4, 1. 2e-3, 3, 1. 8e-3). identifying the optimal ir ___a 6e-4. (The results are very similar to Figure 5 (right),
do not show them repearly.) Then, both Adam W and AdmIRE are trained using the optimal 1r ir_max.
D Proofs in Section 5.2.1
Additional Notations. For the proofs in Section 5, some additional notations are used. For any set K IPP and a constant R O, we denote [(K; ; R) := {0 IR' dist K) < RJ. (...) represents the standard Euclidean inner product between two vectors. ||-|| denotes the l2 norm of a vector the spectral norm of a matrix, whereas || || denotes the Frobenius norm of a matrix.
D.1 Preliminary Lemmas
Lemma D.1 (Arora et al. (2022), Lemma B.2). Under Assumption 5. any compact set K C there exist absolute constants R1, p > such that
(i) [BK; Ri) UU;
(ii) (()) is L-PL (defined in Def F.1) on B(K; R:);
(iii) inf Am (V²L(0)) > μ
RcB(K:R1)
21

---

We further define the following absolute constants on R1):
B2 = sup (E)(K:
VV-C(0)
33 :=:= (E3(K:
sup
:=
sup
V-L(0)
(E
V==
inf (E3(K;
Am (V2((0))
:=
sup (EB(K:RL)
Lemma D.2 (Key properties of (-) (Arora et al., 2022)). Under Assumption 5.1,
For any U, (0()) O.
For any (+( P 1m (V2
Lemma D.3 (Continuity of P 1:p): Under Assumption 5.1, there exists absolute constants R2, Cp O such that for any 0 E ((KKR2).
|| m 1xp 2 ((0)) p:p (V2(((00)) < -00-000000 (0011.
Proof of Lemma D.3. Let the orthogonal decomposition of (²C( and V [((())) be (² Eh= sinny (A1 2
Ap) and V (O) = / (J1 > ... > 1pp). respectively.
By Lemma D.I, for any 0 E (BK; R1), holds that V V>L(+(0)) 33 2())
We choose R₂ R1 A 4%% Then for any B(K; R2),
" (' < 8s 0 - ((0))l 4
Consequently, by Lemma F.2, we can bound the gap of eigenvalues: for any A E [p].
po < V2C(0) V- <
4
Noticing (0) I, by Lemma D.I, holds that lar pm 2 μ and pm+1 we can obtain the bounds of {Ax
Up 0. Thus,
34 Au 2 Am 2 um - 1m Mml 2
for all & m,
4
ds & Am+1 & pm+11 + |m - Mm+il < /
for all k 3 m + 1,
4.
For simplicity, we denote Uwp = (u1, Um). Ubowom = (um+1, . Up). Vio = (v),"", Um)) Viottom (Um+1,"' Up)
By Lemma F3, we can bound the gap between the subspaces:
Unottmm ( V V2C(()) (Z Vap
| Vio 1p
3u
Lemma F.6
bottom
[(2((0) V V²((((0)) Vopli
U bottom
V22((0) - V²C(O(0))
Voopl
(m (p m))
| - V=
According to the definition of 1mp (-), it holds that
F Σ
(V²
Us u
U Vottom U bottom
P
P Pm+1p(V V22(((0))
Σ
URUY
Voottom V bottom
k=m+1
22

---

Noticing the relationship
| l:p (()(()) Pmp (V=(((()))
< ((())))))
- lp (V=(((())) (
= || bottom
Vootto Vottom 12- - 2(p m) 2Tr (U Uboutm Voottom V Vootoo))
-2 Tr (Ubotion Ubottom
(I bottom V bottom)) )) 2 Tr (Ubottom U Vop Viog)
-2 Ubottm VipllFF
we obtain the bound:
(P=tip((=C(()) - - (P) < [ Vupllr
2m2 A (p-m) 1 -- < 2v (m A (p - (p-m)) Ba 110 2(0)]l.
P
To summarize, we only need to choose the constants Rg = Rí A # and (p 2v2(mA(p-m)) (p)m to ensure this lemma holds.
43g
Proof Notations. Now we introduce some additional useful notations in the proof in this section.
First, choose R == (R, A R2)/2, where R₁ is defined in Lemma D.I and R2 is defined in Lemma D.3. Let μ be the PL constant on B(K; R). Moreover, we use the following notations:
82 = sup DER(K:R)
|
Bg
sup OER(K;R)
|
A4
sup (ER(K:R)
V+L(0)|
;
inf 0E3(K:1)
Am (V²C(0)): So = sup (ER(K;R)
(14)
Pm+tip P Pmttip P V2
Cp
sup -E
|0 (0)]
Ensured by Lemma D.1 and D.3, these quantities are all absolute constants in (0, +00). Moreover, without loss of generality, we can assume that 31, 32, Bs, Ço- Sp > 1 and μ & v <
Lemma D.4 (Connections between para norm, grad norm, and loss). For any 0 B BB; R) > holds that:
(para norm v.s. grud norm) μ [[V]
< 0 ((0))l < 32 |[
(grad norm v.s. loss) (0) < | (0) < 283 L(0);
(loss v.s. para norm) 0 (0) 2 < (
82 0 (0)
Proof of Lemma D.4. This lemma is a corollary of local PL and smoothness (Lemma D.1). For the three lower bounds, please refer to the proof of Lemma B.6 in Arora et al. (2022). Then utilizing these lower bounds and the smoothness 32, the upper bounds hold naturally.
Lemma D.5. For all 0 B):
< ( (
(1)))
||Pn=1,(VCCCOO)VVO))) < 0 0(1 -0()))²)
Let p and 1- 0 p E B(K: R), then
V + pل)ه ||VC())] + p32;
1793
P (V p < (10 (0)))) 0 x(0)l)
2
23

---

Proof of Lemma D.5.
< (1200//0||||0P77710/000000010000 +(0) + --&l10 +())
=0+0(10-20011;
PPn+tr(V2210))V<(() (,20000000000||00002777070070000 + (
< (|0p://0660000000000000000000000 - (0))))
+b))++ < [V( + p3::
Pm ) (V
Pm+1:: (1 (/n (1 0²3s
Pm+1pp
<0 =111--00)))+ + ( (0))1) 0²333
D.2 Proof of Theorem 5.5
D.2.1 Proof of Keeping Moving Near Minimizers for SAM
We first give the proof for SAM about "keeping moving near minimizers", which provides important insights into the proof for SAM-IRE.
Recalling (8), the update rule of average SAM is:
0: + P ||E/|| ).l where & ~ NOO,
E.
i 0,
V
Let the p in SAM satisfy:
p = O(vT).
For simplicity, denote
4
C
C2 = 83.
E
J
Notice C(Om) < Cimp², we have the following upper bound for the probability
P ( T Tu), ((0,) > 2C/n0²)
T+Tn-1
P (C(((++) > 2Cipp; Vse [T, ((0,) 2Cinp²).
For each term E [Ti, Ti Tn 1). it can be bounded by:
P (C(0r+1) > 2 v E [Ti, ((O.) < 2Cinp²)
<P (C(0+41) : Cinp²)
1-1
P (C(00++) > 2Crnp; ((O.) < Cimp?: Vr e (s 1.t).Cinp < ((0,) < 20/10²)
==Ii
(((++)) > 20 ((0,) < Cino²)
Σ P (C(0++1) > 2Cmmp²: V Vr is 1, Cinp" & C(0.) < 2Cinp²
((0.) ((0.) Cinp?).
24

---

For simplicity, we denote
PAI1: (C(00+)) 2Cino| Cino²),
P1+1, := (((++))) > 2Cinp; V [s+ 1,t], Cimp² < ((0,) 20in ((0.) Cinp²) SE [T 1].
Step I. Bounding Pi+1.
From C(0.) Cinp², we have 0, ((,)) < / ((,,) O ((>x))) thus
||: + po: ( < |0, (0,)) + p ((vno) + O(p) < R.
which means 0, + pv: E B(K; RR. Furthermore,
10:11- - (0.) < 10: (
pv;)
<T()T))) + V (0, + pm)) < + 32mp
2 [(0.) + 32np & O(vnn) 00n" O((p)) & R,
<0(/np) + +
which implies 0 E BB (K:; R). Consequently, we have the following quadratic upper bound: [(0,+1) (0, V L (0, + pv:))
<(())) 1 (VC(),), 7L (0, + pvi)) + 12 IV (0, + por)l?
<C(O.) " || 1/ (VC(O.), (2C(0.)..) np²Bs Bs
./332 (IVL((0) + pB22
|VC(O)]
2
2
n
<C(O)) - /. /| 7o (V2(0.), V2L(0)).;
VL((0]] +1233 (IVC + +1283)
2
n0²33 3/
<C(0)) - // ((0))
|V((0)] + +17 32
2
<C(O,) - (pan) + n0²B, B3
|VC(0.)] + +077pp
<C()) (n882 + np A3
2,32
VC(e.) + 011²t0²>
2
<Cmp + 017 p²) + O 30inp²/2.
Thus, we obtain
PAIIA = (C(0++1)
(+11=P((((0+))) > 2Cnnp² [(0.) < Cinp²) 0.
Step II. Bounding Pit1, for E [Ti, t 1).
We prove this step under the condition L(O.) Cinp². Define process (X, 1411 X [(0+11),
X1+1 = [10,,1), Cf & X, = L(O,) < 2Cimp²
Xx, -2mp2 else
It is clear that
PAL1. < (Xt++ > 2Cino²).
Then key step is to prove the following claims about the process (X,).
Claim I. X, C2rnp² is super-martingale From the definition of X,, we only need to prove that if Cip' < L(O,) & mmpp,, then E [C((r++)] & L(O,) - p2
If Cimp² & L(O+) & 20 similar to Step I, it is clear that 0 E B(K; R). Applying the quadratic upper bound. it holds that:
C(0,+1) C C(O, VL (0, + po,))
25

---

<C((+) - // po) 1/992
1 (0, + pot)1?
2
<C(O.) / VV <
oo (V
(2C(()))))
n 33 |V + 3/ (IV2 1 +
2
Taking the expectation, we have:
E (C(0++1) <L(O,) - [V
|V + +/ (IV + +0203)
2
3 <C(O,) /|V2(0,)| (1 VC((,))
no²ps B3
npp
<L(O,) 4"' (V
|V (T((0,)1++)7))
1/
2
+ +3?2p33. p²
2
From ((0,) Cinp² and p= (()T)) holds that
|[VC(,,)| 22p((0,) > V2C1uu/p > 43zp²
Therefore, we have:
5
E [[((0++1)] <C( 8
[ )1I p ((0,)
10
10
<C(O,) 8
Cipn + +n2p2 120233
8
<C(O,) S2n2p' p² -((,)) - C2m1pp
Claim II. Xxy1 X, C2n² is O + 1/²/2 pp/ppppppp (1t/))-su--ussssii From the defi- nition of X,, we only need to prove for the case Cinp² & [(0,) < 2C
If Cinp² & <(()) 2Cipp' then X, ((0,) and X,+1 L(0++1). Similar Step I, is clear that 0 E B(K; R).
Applying the smoothness, we have:
[(0++1) = C (0, V/L (0,, + pv+))
=C(O,) - (VC(0,), VC + po,)) +0 (1² (V + po:))²)
=(())) ", - (VC())))))
+ +0 (np" =0
+2pp²
+ = (0 +0 (0))
=L(0,) - np (VC Vz +0(1p07
which implies:
|| -x, X, 1212p1l < -1C(0++1)- ((0.))l + +|(221²p"
(VE(()))
+ +0(7pp²
<mp (3C(0,) |K
( -
11 >))1 +0(1000
1/2
50 1/0120² 02 K [4-100
((22 8 4/
) +0(7Pp)
> )
Lemma E5
(1223 // /VD) + O(1P))
With the preparation of Claim I and Claim II, we can use the Azuma-Hoefiding inequality (Lemma F.4 (ii)): for any Q > 0, it holds that
Q2
P (X - X.+1 + (1 02 > Q) < 2exp
(
2(t 2(t- - 55) (O(1)3/2 / 1/2 + 17 p))))
26

---

As proved in Claim I, Xe41 [(0++1) < C due to [(0.) & Cynp². Therefore, by
choosing Q = (t S)C2np "inp + +2 = (t-ssC2nf p² Cipp. we have
P141,8 & PP (Xe+1 > 20111// (111
X X + (1 - ()-00 (1 -1- (1-)) + 1 "cmm²)
<P
2
(11-8) + Cimp²)
<xexp
2(1- 8) (0(1)/2 12 + 12222))²
4(t s)C2n' C
(-)
( ))
1
cxp
< 2cxp
-S2
4(t - (0(r)))))
17
Therefore, obtain the union bound:
7+ Σ
P ((tt [Ti, Ti + Tn), ((0.) > 2Cinp²)
P141+
P41s
T
Pe-1. & THe exp (-12
Hence, with probability at least
T7 exp
expe
(-2
n+ )) for any E (Ti, Ti Tu],
1
||0 10:- (00,1l < 2 ((0)) <2
Ci
43
T O())p).
P
D.2.2 Proof of Moving Near Minimizers for SAM-IRE
We prove "keeping moving near minimizers" for SAM-IRE. The proof outline for SAM-IRE the same SAM. However, the key non-trivial difference is that the IRE term will hardly cause loss instability since IRE only perturbs the parameters in the flat directions.
Under the conditions in Theorem 5.5, the update rule of IRE on average SAM is:
& |&/1 1ppPm+1i ( C 0, + |
&
0 = 0, nV- 02 +
where & ~ I0,,
Let the p in SAM and the & in IRE satisfy:
p p= O(T)) K - 1 P
(15)
For simplicity, we denote
452 432
U1 = &
P(O.) == Pm+1pp P (V²L(0))),
C1
C2 = B3
Following the proof for SAM, we denote
PIL1A := (C(Be+1) 2Cinp² ((0)) Cimp²)
P P= :PL:: > 2Cinp²:
Vr Is + 1, f), Cinp² < L((,) < 21npp ((0.)) < CinD"), SE [Ti, -.]. 1).
and it holds that
T++Tn-1
P Pate [Ti T > 2Cinp²)
<
Σ
P + Σ PiL1s
27

---

Step I. Bounding P4+1.
From [(0.)= Cimp², we have |0: - - < / ((a)) O((Tp), thus
| + po: ((.)) & (0.))l + -0((pp)) + O(p) R
which means 0, + pr; E B(K; R). Furthermore,
00+1 - (
00 (00,) - +
((())) | po:)
+ nn
Lemma D.5
+ - 82np O (map²)
+0
(np)
;0 (np) < R
which implies 1 E B(K; R). Consequently, we have the following quadratic upper bound:
[(0(+1) L (0: V + por) mx VC + pv.))
<C(O,) - (
+ - pv,))
+ 1²332
2 VVC +(00))))))=0000 <.(),) " (V)(O,),
+ - (2)
poi))
+ +/ [ (0, + pm: |000==
<(())) / np32 VC
3xp
<mp -mpp(VC(0)))P0))))200000) K [I((O)VC((O)]
+ 1/ (1IV2(0)) + -
<L(0.) + 7032 |10
kmpp? 33
+ 2
|P(0))V2(0))1= ( ((VC())) + +32)) +
(P
Lemma D.5
((0.) + Onp (l) -0mp))))000000mlo-00000000010110
+0 +00mmi) <kпp² 10. +00)
<Cinp² + =(nno²) < Cimp².
Thus, we obtain
P PP (C(0)+1)
(PC(0+)(((((+)) > 20'np² ((0.) Cinp²) 0.
Step II. Bounding 1r+ for s E [T. 1-f--1] 1].
We prove this step under the condition L(O,) < Cinp². Define a process (X, 1441, X L(0++1).
X,+1= = [L(0,+1), Crnp & X, = L(0,) < 2Cn
X,, Comppp else
is clear that
P.LLa < P > 2Cino²)
Then our key step is to prove the following two claims about the process {X,}-
- Claim I. X, - C2rnp² is super-martingale. From the definition of X,, we only need to prove that if Cimp" L(O,) & 2Cino², then ((((0,+1) < ((0,) p²
28

---

If Cinp² & (())) < 2Cnnp', similar to Step I, it is clear that 0,+1 E B(K; R). Applying the quadratic upper bound. holds that:
[(0r+1) C (0, VL (0, + pu,) mxP((.)) (0, + pv,))
<C(0,) - (V VL (0, + pw,) - KP(0.)VL
+ pv-))
3/ [V. + po.) ( + poz) +
2
<C(0,) - -n (VC(O,,,, V2 (0, + pu,) (x7()))) P(O.)/C + po,))
+ 1/ (1 -p000~++,00111)))000))))))))))) | (O, + po,))l
- (VC(O,), (20(), + no B3 |VC(0,)]
2
Bar²
Lu 11
- (,C0:0))))-mmp(((CC))00PP0/01110 + km 2
[IP(0-)V<(0-)]
+ +7
Taking the expectation, we have:
[[((0,+1)]
- 1V + np²Bs
2
3zp²
+ KT) 2 |P(00) + 3
1 + poo))
< <C(0,) - ( / + 7
3B
83p2
VL(0,)] + Kл
[I((O))
2
2
+ + B/ (.- + 28 +B27 R² IP(()) pp
Lemma D.5
37
|C mppp 33
((0,)
|C(00)] np²)
4
+ 2202n2pp + 32m12r² x22 (O(
B3 Papp))
37 <C(O,) 4
|V np²p3s 33
+ p² + 3203n222'' + 00720')
+
From Cinp² < ((O,) < 20 npp and -= ( holds that
||VC(0,)| (,, v 2C 4.3pp²
Moreover, recall < 1/p. Therefore, we have the upper bound:
E[((0++1)
<C(O,) 37
|VC(0,) + 70333 VV((+)) + 2 +
2
50
<C(O,)
|VC(0,) "1² + +28 3233022ppppppppp (
8
10
<C(O,) (0,) + +
<C(0,) 10
33
B28
1/ + 283 + B233ip + o
1
<L(0,) - 1/ = -((,) - 22
Claim II. X X, C2n22 p² is On72 p² + 1/2 1////)-uub-Gausiii From the defi-
nition of X,, we only need to prove for the case Cinp² < [(O,) < 2C1np².
If Cimp² & <((O)) & 2Cmpp2, then X, = C(O,) and X X = L(0,+1). Similar to Step I, it is clear that 0,+1 E R).
29

---

Applying the smoothness, we have:
VLL ((0,+1) = L (0, -0, ( + pv,) - OKP(O, VLL + pv,))
=L(0,)- (V (0, + + ==P(0,) (0, + pp,))
-
+ (0.+ + (L, + pot))l²)
= -C(),) -1tp1//2000011mmmm//00000000/000 + O (VV[()))))
- mll
P(
+ +0 (ne'² [P(O,)
+ 17 (0, xP(O,, + po:))l
In the same way the proof of Claim I, it holds that
n |V 21 p²) nu² [VC = 01122²), nMpp
17
+ pv,) po,))++ (0, - =
Thus,
-(0++))-((0))
= - wwt.://......000((((1)eeee.---10 (00)))
which implies:
|| | - X+ - C212p21
< - C222
v22
+ + )
110, 0
006'))))))))))))))))))))))))
1/2
V=[(0,)VC((0))
<np
/11
+nmp v2 (= |
+0
v
1/2
Lemma 0 12/2
(= 12 "*) >,,,.)))
>11I) +
+0 ?2/2p?
Lemma F.5
0 (13/2 vp) + P1)
With the preparation of Claim I and Claim II, we can use the Azuma-Hoefiding inequality (Lemma E44 (i)): for any Q > 0, it holds that
Q2
P P(X++1 - X.+1 + (t - s) > Q) < <2exp 2( - s) (O / + 17² p2))
As proved in Claim I, X,41 C(0 < Cmmp² due to [(O,) & Cinp². Therefore, by choosing Q = - S2n02 Cmmp + 220mnp2 (1 (1- 22- 1² Cimpp, have
P441 < P (X++1 > 2Cm10²)
X(+1 - Xst1 + (t- - -D=s (t- (C=sCC212p2 p² + 1
cinp))
<P
2
((t - S) + Cnnpp)
<2exp
2(t - s) (0(r3/2 0² /p1 12 + n2 )²
4(1 - s)C212 p² Crupp
& 2exp (-0 (-0(11))
<xexp
4(t- 8) (0
-
30

---

Therefore, we obtain the union bound:
T++T=-
[-1
P( [T, Tnn), > 20 &
P141I+ +
Pit1.
T=
5
T++Tu-
P P < Tn exp (-2 1 (-
Hence, with probability at least 1 1 exp (-1 (-- for any [Ti, Tull
100 100--100 (0.))l < c()))) & <
C
4,2
Vnp
Vnp O(vTp).
1
D.2.3 Proof of the Effective Dynamics
(-1)
(1++)) (++++)))) for any t [Ti, Ti
We have proved that with high probability least 1 Ti exp
Tul, |10, ((0)) (()))). Then prove this theorem when the above event occurs. For any E [Ti, Ti Tn]:
[0.+1 0:1l
+ 100 pox)l map |PO.) + O(nmp²)
<
Lemma D.5
+0mm))))")))))))))m)))))0))0)- ((1)).
Then by Taylor's expansion,
((0++1) - ((,) (0,))) (0.+1 - 0.) + (10.+1 -1:+0
-0.1'- 0.11²)
=- =-n0b60.)V + prL) - (m + por) + O3)2p²)
For the term (0(0.)VL(() + pv;) and + pv,), using Taylor's expansion, have
(0( p7m)
-(0(0.VCC()) (200))) 8² VT + +00))
-00(0.)V 2(0 +
2
00(0))P(..VVV(0. po))
0²
=00(0))((..VV((00)) + -.00 + [1000)00000TTTT((2220000 +067
2
1²
-30(0))P(O.VV2(()) + P00
(..00))) + +0(0²).
2
Taking the expectation (about v,), we have
Tr (V)))
1 0,
E U C()))))
=
P
Additionally, using Lemma D.2 and Taylor's expansion, we have:
(0(0.)VL(0)) = O;
@4(0.) (2(((0)) + 0 x(0.)) 02(((0.) + O(vTp));
31

---

- -
(II()=((0)) = ( - x - + O(yī))
==0
Combining the results above, we obtain:
((++)))))
-0,) no (0.)
n1²
kmP? (V
(4(0,) Ir ((O)))
2p
2p
(n + 1)mp²
=#(0.) 0 (17 kp²)
O(pp²)
2p
(k + 1)mp²
=0(0.) 0 (1²"p²)
0T
(0:
0 O(p)))
2p
=0(()) - (n + (V 00)
E Proofs in Section 5.2.2
Setting E.1. Consider the empirical risk minimization min [()) E/= (0), where () ((f:(0) vi) is the loss on the i-th data (æ,, yi). Let fi()) and E(.,-) be Cs. Suppose all global minimizers interpolate the training dataset, i.e., L( ming L()) implies f((0") = yi for all i [n]. We denote the minima manifold by . (0 fi(0) yı, Vi [n]]. Moreover, assume ('((..))
that
8y ji-y 0 and the feature matrix (V); ()) E RPxn is full-rank at O E M.
Lemma E.2 (Theorem 5.2 in Wen et al. (2023a)). Under Setting E. I, Assumption 5.1 holds with m n.
E.1 Preliminary Lemmas
Similar to the proofs for Section 5.2.1, we need the following similar preliminary lemmas. Lemma E.3. Under Setting f. any compact set K L, there exist absolute constants R1, / 0 such that
(i) B(K; U.;
(ii) (.(-) (i E (m)) and L()) are p-PL B(K; RT);
(iii) inf An (V²L())) > μt (ER(K;R1)
inf \ (V)L:(0)) > 44 Vi E [n].
(.)(K;R))
We further define the following absolute constants on B(K; R):
B2 = sup (E)(
||
max sup
L (0)
:Ein) (E
(ER(K;R)) || Bg = sup
V
max sup sEjn (E
[
v:=
inf (EB(K:R))
Am (V2C(0))
^ min inf NEn (E
A1 (.
|V4(0) sup
sup
(EB(K:R1)
(E3(K;
32

---

Lemma E.4 (Wen et al. (2023a)). Under Assumption 5.1,
For any 0 E U, (0 O.
For any 0 E I, ( = Pn+1:p (V2 and (+(0)V2L(() = 0.
For any 0 T (0(0))0(L()))))) (0) 0, vi E [n].
Lemma E.5. Under Setting E. there exists absolute constants R2: Cr such that for any 0 E [BK; R2).
11+ ( (0)) - Pn (V=L((()) < Sr 0 E(0)11-
Proof Notations. Now we introduce some additional useful notations in the proof in this section.
First, we choose R == [R1 A R2)/2, where R1 is defined Lemma E.3 and R2 is defined in Lemma E.5. Let μ be the PL constant BB (K; R). Moreover, use the following notations:
82 = sup
V max sup sEinl (E
V2 v22,(0)]
(EB(K;M)
83
sup (EB(K;R)
|
max sup sE|m) (E vv'
84 =
sup (EB(K;R)
max sup
sE|n) (EB(K:R) |v
inf (EB(K;R)
Am (ViC(0))
min (E(n) (E)(K;
inf
\1 (V>
PR+Ip ( (()) Pu+lp (2
Ço = sup (EB(K:R)
v+)(0)]|; SP
sup (EB(K;
0 (())
(16)
Ensured by Lemma D.1 and D.3, these quantities are all absolute constants in (0, +00). Moreover, without loss of generality, we can assume that 31, B2, B3, Co. <p > 1 and μ & P < 1.
Then have the following lemmas, similar Lemma D.4 and D.5. Lemma E.6. For any E B(K; R), holds that:
(para norm v.S. grad norm) ( < - ||0 < B2 (V P VVL( < 10 (())l & 32 (O)], vi E [n].
(grad norm v.s. loss) 2p) < |V < 25 (()): ( & [V)((0)) 2 (,))) Vi E [n].
(loss v.s. para norm) ||0 1 (00)) < L < 24
0 +(0) # Φ(0))1 < L) <
c 0 20)) R vi [n].
Lemma E.7. For all 0 e B(K; R), Vi E |n].
(V
(B ||P
(,000000000)) <|(00 -()))
& <0(10-(00))
LLt p > and v e -r 1 + ~ B(K; then
[V pu)ll < [V Vie [n];;
1²83 vi P (VCC)))) p < 0 (1e +000)) (())))
+0
[n]:
33

---

Lemma E.8 (Lemma H.9 in Wen et al. (2023a)). For any absolute constant C > 0, there exist
absolute constant C C2 > O such that: if 0, E [(K; R) and Cipp < 10, ((0,))l I| & Cp. then holds that:
B, 0//1// 4 1(0++1/2)| < 100 - ((,)) C2mp_
Lemma E.9 (Lemma H.10 in Wen al. (2023a)). For any absolute constant C > there exists absoulte constant C3 such that: if E B(K; R) and ||0 4(0,)l || & Cp, then we have that
0 1/2 - + (0.-1/2)
|0: - -(66,)1 < Csmp.
Now we fix the positive number C' and use the absolute constants C22, C3, defined Lemma E.8 and Lemma E.9.
E.2 Proof of Theorem 5.6
E.2.1 Proof of Moving Near Minimizers for SAM-IRE
This proof is similar the proof for Section 5.2.1.
Under the conditions Theorem 5.6, the update rule of IRE on standard SAM
VC,
01+1 =0; 18---11 -V, 0 P IVC, (0.)1
VLi, (0:) (: + P LVV,, (0:)
- (V²L(0.)) VL..
where in U([x)).
Let the k in IRE satisfy
< 1/p.
Additionally, we fix constant E =0,1) in the proof. For simplicity, denote
= VL, (0.)
P(O.) = Pattop (V²L(0.))
VL. (0.)
and
0.+1/2 = VL., (0, + pvr):
0++1 = 0 0,+1/2 p(0) Li, (0, + prt)
Additionally, we denote
Pe+1,e := IP (10041 (0.+1))) 2C1)-=0
+(O.)ll < -(..)l
<Cn'-ap)
F1t1, ::([10++ - +00++1)) > 20111--a
( Cr 1- & 110 - 2(,)) < 20' 1-a 10, -(0))) < CvT10). s sE [T, T-,--11 1].
Then the following bound holds naturally:
T+
P (Bte [T. Tul. |i (0,) > 201)--a) <
P Σ P
T=
Step I. Bounding 1r
From 0, - <b(0))) = O171 10p)) thus
||0 + pv; (0,,)) || < 0, - 00p)) <R,
34

---

which means 0, + pod E B(K; R). Using Lemma E.7 and K & 1/p. we can estimate:
| 0 0/41/2| = mp 1VP() (0: + po,)ll
(0:) + 83
1/p
(0.1)+
// 0²)
11pp + 1
83
0²) B37Mp2
C2
2
2(1 + (1
npp
|0c+1 ((0++)) - 200 +(0++1/2))]
C2
< + (1+ | - 0c41/211 &
2 mp.
Then we have the following bound:
0 +(0(+1)]l
- (0,())))) + + -(00)+ 12)
- (0, ((.))
+ +100-11 ((0++1) - (0t+1/2 ((0++1/2))
Lemma E.9
C1n Cpyl-ap- 0 | ((0 (0.+1/2 (0
x (mp)) < 301 1) -
2
Thus, we obtain
PitL -M (100 (0,,+1) > 201)² pl 10, -0(0)) < Cn' Cnn'ap) 0.
Step II. Bounding Pit1,s for E [Ti, ( 1].
We prove this step under the condition |0, - -((0)) < Crt² "p. Define a process (X, |+1
Xs+1
10,+1 ((++1))l
{ [10 11,+
if ir
X, = -, - ((0)) & 201)--ap
X++1
X C211/2, else
It is clear that
Pitt, & &P > 20m''=p)
Then our key step is to prove the following two claims about the process {X,)-
Claim I. X,, 22 is super-martingale. From the definition of X,, only need to p < X, = 0, (0,))] < 2Cy)²-0,, p, then E [ - (00++1)l < |0, ( - C2np/2
prove that if -
If Cnp1-ap p < X, = |00 - ((0-)l < -Cp similar to Step I, it holds that 0r+1 BKK; R) and 0 - ((0++1) (0++1/2 - 4 ((0++//))) < C2 mp. Moreover,
||0,+1 ((0++))
10/41/2 - < (0, (0++1/2)] + |041 $ - (0,+1/2 (0,+1/2))
C2 2 mp..
0/+1/22 (0.+1/22) +
Taking the expectation and using Lemma E.8, we have
C2
E 00,1 - ((0++1)ll < | - (0(0+112)|| +
2 mp
C₂ np 10, ((0.)
C₂
1, *(0,)) C2np
np.
35

---

Claim II. X++1 - X, + C2mp/2 is ((np)) From the definition of X,, we only
need prove for the case Cn1'-ap "
X, = 100,- +((0)) || & 2Cr
If Cnit-ap p & X, - 110, ((0)] < 2C1 3C1)--app p. we have 0,+1 E (BK; Rt)
and |0 +(0++1) (0++1/2 - + ((0-11/2))
< 2 np. Combining this result and
Lemma E.9, we have
||0 ((0++1)) Il - 110- +((+)]||
| $(0,+1)ll - ||0,-1/2 - ((0,+11/2) + | -
|0, $(00) |||
(0r+1 - +(0++1)) - (0,+1/2 - $(0,+1/2)) + 0 - +(0,+1/2)] - |0 + ((+))]
C2
2 np + Canp = ((p)).
With the preparation of Claim and Claim II, we can use the Azuma-Hoeffeding inequality: for any Q > 0, it holds that
Q²
P PX X,+I + (t- s)m2np/2 > Q) < exp
2(t - ([t=s)((pp) ()7)?)?))
As proved in Claim I, Xx+1 = 110,+1 - +(0++1) & 3Cnl-o0 p due to 10, -(0))l E
Cnp-aa. Therefore, by choosing Q = (t - S2C2pp/2 Cn p + 201) (t
s)C2np/2 + Cn'-a p//,, we have
P141, F (X1+1 > 20111=ap)
C2
C
<P
(x+1 - Xs+1 + (t S2pppp (1 8)
np + 2 "1
2
((1 C2np1 + p/2)
<2exp
2(t- - s)0(1pp2)
(1 (1--8)0 C2mp -x)
(-1 7pp
<2exp
< 2exp
4(t - A(t=s)0(r²) ()???)? p2)
Therefore, obtain the union bound:
T
1-1
P [i, + Tu), |0, +(0.))l >
<
Σ
Poile P1t1,s
Σ
t=
==Ii
7+Th-1 [-1 Σ Peis < Ti (-Q(((11)))))))
==
E.2.2 Proof of the Effective Dynamics
By proof above, with probability at least 1 - Tí exp (-22 (-12 (1) for any t [Ti, Ti + Tul.
100 ((0))ll = O17) 1-a p). Then we prove this theorem when the above event occurs.
Due to |10, +0,))l = O(1''Onn''00 p). we have:
||0 0.1l
< [VL, (0, + po:) + TH (IP.)) po.))
Lemmu E.7
"7 VL. (0.) + O(p)) + ((()))
O(pp).
Then by Taylor's expansion,
2(0++1) (00)) = 04(0)) - 0,) + +0 (110-1 - 0:11²) = 17 L., (0. + pox) (0.) (0: + prs)) + (1)²p²) For the term (0 V., (0, + pvt) and (00(0)) (0. + pv:). using Taylor's expansion and Lemma E.7, have
(O, (0. + pv))
36

---

_ 6 [0000))TTTT00332,000000 + 0
(0000))22.000-000000000000000, +
- (0000000,,,0000)0000000000000000
[0.00 +0(0))
((0))) (.,
(L, 00 ((),)))
0² (
VC,,
(.(
+ 2
V'C, [VC,
Vo,
+ +0 10: -
+
p²
D²),
2
and
(6 pv;)
p² 2 -=
=0(10, - 0 + -(0))))) + - (1)00)))(.,)V(((VVvv.. + +0
-0(1-~00² 1² 0 (
(.(,0) T
(L(()((0))
[Q(,, (o(..) VL. (0(0)))
+ +0 10, (()) - =(p²)
0. 0.) ( (V Li (+()))) pp))
Combining the results above, we obtain:
((0++1) -b.)) - (1+ ?0² 2 ( (00)..))) + 00mr7 + knp²).
Additionally, taking the expectation, we have:
E, [a (70
Therefore, we obtain
E, 10 (1 k) np²
2 +0000000 + kmp?)
= - (1 + k) np
20 20 (000000)VVV
h.o.t.. h..ot.
F Useful Inequalities
Definition F.1 (μ-PL), Let p > 0 be constant. A function is p-PL in set U iff [V((0)] 2p(C(()) - inf (())), vo E U.
oeu
Lemma E.2 (Weyl Theorem), Let A, B E RexP be symmetric with eigenvalues X > ... > A, and
#1 2 Mp respectively, then for any ls E (p), it holds that
D& uk < A
37

---

Lemma F.3 (Davis-Kahan sin(@) theorem). Let A, B E RexP be symmetric matrices. Denote their orthogonal decomposition as A = E T E + E Eg and B Fi + F2T,, F2 with (Ei, E2) and (D1, D2) orthogonal. the eigenvalues in A are contained in an interval (a, b), and the eigenvalues of 2 are excluded from the interval (a 6,b + 6) for some ó > 0, then for any unitarily invariant norm ||--1,,
F2 (A- (A-B)EIl) |1
F2 E1
8
Lemma F.4 (Azuma-Hoeffdingg Hoeftding Inequality). Suppose (_nn JEE super-martingale.
(i) (Bounded martingale difference). If & Xi+1 X; & B, then for any 1. >> 0, we have:
P (X,, - Xo > < 2exp
2n +33
(ii) Stub-Gaussian martingale difference). If Xi+1 X, is of -u-sub-Gausse Gaussian, then for any
n,, t O, we have:
t2
P (Xx, - -xu 3 t) < 2exp
2 2 0?
Lemma F.5. Let v E RP. g N(0, 1). Then there exists an absolute constant c > 0 such that for any 1 0,
(K
P
9
4e cptt
|||||
||9||
Proof of Lemma F.5. From P54 in Vershynin (2018), there exists an absolute constant >> 0 Ouchh Without loss of generality, we can assume *
that for any > [e]
< 4e opt
Igll
Then have:
1
(el, g) g >t & 4cpp
P
IP
|||||
g
vp
Lemma F.6.
||AB]] 1p & ||||
F-
38

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We believe that the abstract and introduction reflect the contributions and scope the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
In Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
39

---

Justification:
In Section 2 and 5; Appendix B, D, E, and F.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We believe that all of the experimental results are reproducible in work.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
40

---

Answer: [Yes]
Justification: In atttp:///githie http://////RREaaaeeae watt--aagrgritt-rrrreoooooo framework.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /CodeSubmic / //nips. wttpccccpppccccceeececccccccccct / saissssillllo onPolicy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: In Section 4 and Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: In Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
41

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Appendix C.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttp///eerrp/.ee/ee cctppblic//aaecccccccccccccc wtw.:/tthhhiiieeeeeeeeeeeeeeeeee
Answer: [Yes]
Justification: Ethics.
We have confirmed that the research conducted with the NeurIPS Code of
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
42

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
43

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
44


Reviews:
[{"review_id": "ek7mrqvImz", "reviewer": "Reviewer_69nS", "summary": "The authors propose IRE to enhance the implicit regularization of base optimizers, thereby improving the generalization and convergence in deep learning. IRE decouples the dynamics of flat and sharp directions, reducing sharpness along flat directions while maintaining stability in sharp directions. The paper provides theoretical evidence that IRE can substantially expedite convergence towards flat minima in SAM.", "questions": "1. Could IRE help in the finetuning phase of LLMs, Including both convergence speed and convergence position properties?", "limitations": "1. The author did not provide information on the computing resources they used.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. IRE's ability to integrate with existing optimizers without major modifications makes it easily adoptable in current systems.\nPerformance Improvement: Empirical results show that IRE enhances the generalization capabilities of popular optimizers across multiple tasks and datasets.\n2. In the pre-training of large language models, IRE has demonstrated a significant acceleration in convergence speed. \n3. The paper offers a theoretical foundation for IRE's effectiveness in minimizing the trace of the Hessian, reinforcing its practical applications.", "weaknesses": "1. The improvements of IRE in Table 1 and 7 are not as significant as expected. It makes one wonder about its usefulness for CNN networks and its oversensitivity to hyperparameters\n2. A big concern for me is that the experiment in Figure 3 does not appear to have converged yet, the rate of convergence in the early part of an experiment does not equate to the rate of convergence throughout the training process, as well as a more detailed analysis of the performance of the model after convergence should have been added to demonstrate that the final position of convergence is good.\n3. Judging from the code, the cost of each training step is at least twice that of SGD.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s recognition of our work and helpful comments. Below, we offer detailed responses to the reviewers questions.\n\n- **W1.** The improvements of IRE in Table 1 and 7 are not as significant as expected. It makes one wonder about its usefulness for CNN networks and its oversensitivity to hyperparameters.\n\n  **Response:** We thank the reviewer for this insightful comment.\n  - **CNNs vs ViTs.**  It is important to emphasize that **CNNs are designed  with a strong image prior**, which limits the potential for further improvement through regularization. Consequently, it is not unexpected that the improvement achieved by IRE for CNNs is less significant. This aligns  with  observations in previous regularization techniques, such as ASAM [1] and SAM-ON [2], which similarly shows limited improvements for CNNs. For instance, ASAM improves SAM by only 0.26\\% points when training WRN-28-10 on CIFAR-100, and SAM-ON improves SAM by just 0.17\\% points when training ResNet-50 on ImageNet. In contrast, ViTs, which have a much weaker image prior, benefit more from regularization, as shown in Table 3.\n  - **Sensitivity to hyperparameters.** 1) For image classification tasks, the optimal  hyperparameters for IRE do vary depending the dataset and model. Nevertheless, Table 7 demonstrates that  IRE **consistently** improves performance even without tuning the hyper-parameters, although the improvement is modest. 2) In contrast, when training Llama models, IRE with a fixed hyperparameter $\\gamma=0.6$ performs effectively across varying model and data sizes, as shown in Fig.3. \n\n- **W2.** A big concern for me is that the experiment in Figure 3 does not appear to have converged yet, the rate of convergence in the early part of an experiment does not equate to the rate of convergence throughout the training process, as well as a more detailed analysis of the performance of the model after convergence should have been added to demonstrate that the final position of convergence is good.\n\n  **Response:** We thank the reviewer for raising this concern.\n  - **Loss convergence in LLM pre-training.** It is important to note that in LLM pre-training, the training loss never decrease significantly due to the huge size of  dataset. For example, as reported in [3], the official Llama pre-training achieves only  a final loss around 1.55, despite the model size being 65B.\n  - **Final loss v.s. performance.** In LLM pre-training, it is widely observed that  the performance on downstream tasks is predominantly determined by the final pre-training loss and has little correlation with other pre-training components such as model types. We refer to [4] for a detailed analysis of this issue. Therefore,  the primary focus in LLM pre-training is to reduce the final loss  as much as possible, given specific computational and data resources.\n  - **Final loss.** Our final loss aligns with  expectations based on our model scale. For instance, for our baseline, our LLama (229M) was trained on openwebtext by AdamW for 100k steps, resulting in a final loss of 2.835. Fig.1(d) in [5] show that GPT models of varying sizes were trained on openwebtext by AdamW for 100k steps, achieving losses of 2.915 for GPT-small (125M) and 2.695 for GPT-middle (355M). Our loss of 2.835  for Llama (229M) (the baseline) is consistent with these results. Notably, our proposed AdmIRE achieves the same loss of 2.835 with only 50k iterations, which is half the number of iterations  required by AdamW.\n    \n- **W3.** Judging from the code, the cost of each training step is at least twice that of SGD.\n\n  **Response:** We  clarify that the cost of our IRE takes twice that of base optimizers only during the steps of estimating the diagonal Hessian and updating the mask. However, these steps are triggered only once every $K$ steps (see Alg. 1). In all our experiments, we set $K=10$ (see l.168-l.172), so the average cost of IRE is only **1.1 times** that of the base optimizer. Table 4 further verify this estimate empirically.\n\n- **Q1.** Could IRE help in the finetuning phase of LLMs, Including both convergence speed and convergence position properties?\n\n  **Response:** We thank the reviewer for this constructive question. To address this, we have conducted a new Supervised Fine-Tuning (SFT) experiment. Specifically, we finetune pretrained Llama2-7B with LoRA on Stanford's Alpaca dataset. The results, shown in the **attached PDF** in our **Global Response to All Reviewers**, demonstrate that IRE can **accelerate the convergence in SFT**. This is consistent with the advantage of IRE in pretraining. Unfortunately, we do not have enough time to evaluate the performance of the convergent models on downstream tasks. In the revised version, we will include complete results for SFT.\n\n\n- **L1.** The author did not provide information on the computing resources they used.\n\n  **Response:** Thank you for this reminder. For Section 4.1 (image classification), the experiments on Cifar-10/100 were conducted using a single A800 GPU, while the experiments on ImageNet were conducted using 4 A800 GPUs. Details regarding the computing resources for the experiments in Section 4.2 are provided in Appendix B (l.664 and l.678).\n\n\n[1] Kwon et al. ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks. (ICML 2021)\n\n[2] Muller et al. Normalization Layers Are All That Sharpness-Aware Minimization Needs. (NeurIPS 2023)\n\n[3] Touvron et al., LLaMA: Open and Efficient Foundation Language Models. (2023)\n\n[4] Du et al. Understanding Emergent Abilities of Language Models from the Loss Perspective. (2024)\n\n[5] Liu et al. Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. (ICLR 2024)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qBT6tQbhx6", "reviewer": "Reviewer_QWSd", "summary": "This paper proposes an Implicit Regularization Enhancement (IRE) framework to accelerate the convergence of optimization algorithms towards flat minima in deep learning, thereby improving generalization and convergence. The key idea behind IRE is to decouple the dynamics along flat and sharp directions in the optimization landscape. It speeds up the dynamics along flat directions while keeping the dynamics in sharp directions unchanged. This allows IRE to substantially boost the implicit sharpness reduction of the base optimizer without hurting training stability.\n\nThe authors provide a practical way to efficiently incorporate IRE with generic base optimizers without introducing significant computational overhead. Extensive experiments across image classification and language modeling tasks demonstrate that IRE consistently improves the generalization performance of base optimizers like SGD, Adam, and SAM. Surprisingly, IRE also achieves a 2x speedup compared to a well-tuned AdamW optimizer in pre-training Llama language models of various sizes.\n\nFurthermore, the paper provides theoretical guarantees showing that IRE can achieve a substantial acceleration over the base SAM algorithm in minimizing the trace of the Hessian matrix, which measures the flatness of the loss landscape. \n\nIn summary, the key contributions are:\n\n1. The IRE framework that enhances implicit regularization by decoupling and accelerating optimization dynamics along flat directions\n2. A practical and efficient implementation of IRE that can be incorporated into generic base optimizers \n3. Extensive empirical results showing IRE improves generalization and speeds up convergence across vision and language tasks\n4. Theoretical analysis proving IRE can substantially accelerate the sharpness reduction of the SAM optimizer", "questions": "1. **Comparison with related works:**\n   - Question: How does IRE differ from and improve upon other techniques that aim to accelerate convergence to flat minima or reduce the computational cost of sharpness-aware minimization, such as those proposed by Kwon et al. (2021) and Liu et al. (2022)?\n   - Suggestion: Provide a more detailed discussion comparing and contrasting IRE with these related approaches to better highlight the novelty and advantages of IRE.\n\n2. **Limitations and failure cases:**\n   - Questions: Are there scenarios where the diagonal Hessian approximation might be less effective or lead to suboptimal results? How sensitive is IRE to the choice of hyperparameters (e.g., $\\lambda$ and $\\gamma$)? Are there any particular types of models, datasets, or tasks where IRE may not provide significant benefits or even hurt performance?\n   - Suggestion: Include a more in-depth discussion of the potential limitations and failure cases of IRE to help readers better understand its scope and applicability.\n\n3. **Experimental setup and implementation details:**\n   - Question: What are the specific hyperparameter tuning ranges for $\\lambda$ and $\\gamma$ used in the experiments, the architectures of the models (e.g., ResNet and ViT variants), and the data augmentation and preprocessing techniques applied?\n   - Suggestion: Include these key details in the main text to make the experiments more self-contained and easier to interpret and reproduce.\n\n4. **Comparisons with state-of-the-art optimizers:**\n   - Question: How does IRE perform compared to other recent optimization approaches, such as adaptive gradient methods (e.g., Adagrad, Adam), second-order methods (e.g., K-FAC, Shampoo), and other regularization techniques (e.g., weight decay, dropout, label smoothing)?\n   - Suggestion: Include comparisons with a broader range of state-of-the-art optimizers and regularization techniques to comprehensively understand IRE's performance and potential advantages.\n\n5. **Downstream performance evaluation:**\n   - Question: How does the downstream performance of models trained with IRE compare to those trained with other optimizers when fine-tuned on benchmark tasks like language understanding, question answering, or text generation?\n   - Suggestion: If feasible within the rebuttal period, evaluate the downstream performance of the pre-trained Llama models to provide insights into the practical impact of IRE on model quality and generalization. If not feasible, discuss this limitation and propose it as a direction for future work.\n\n6. **Theoretical analysis:**\n   - Questions: What is the practical significance of Assumption 5.1 (manifold of minimizers) and how does it relate to the empirical observations in deep learning? Can you provide a high-level interpretation of the non-asymptotic bounds and their dependence on the hyperparameters (e.g., $\\eta$, $\\rho$, $\\lambda$)?\n   - Suggestion: Provide more intuitive explanations of the key assumptions and their implications in the theoretical analysis. Additionally, explore the theoretical connections between IRE and other optimization techniques (e.g., momentum, adaptive methods) to provide further insights into its behavior and potential extensions.\n\nAddressing these questions and incorporating the suggestions in the rebuttal or a revised version of the paper would help to strengthen the work and provide a more comprehensive understanding of the proposed IRE framework, its novelty, effectiveness, and impact.", "limitations": "1. **Limitations:** \n   - Add a dedicated subsection discussing the assumptions made in the theoretical analysis, potential scenarios where IRE may not provide benefits or even hurt performance, and the sensitivity of IRE to hyperparameter choices.\n\n2. **Potential negative societal impact:** \n   - Include a subsection discussing the environmental cost of training LLMs, the potential misuse of LLMs for generating harmful content, and the possible widening of the gap between well-resourced and under-resourced research groups due to IRE's computational advantages. \n   - Propose mitigation strategies or areas for future research to address these concerns.\n\n3. **Reproducibility and transparency:** \n   - Provide clear instructions for reproducing the experiments, make the code and pre-trained models publicly available, and discuss any limitations or challenges in reproducibility.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "## Originality\nThe key idea behind the proposed Implicit Regularization Enhancement (IRE) framework - decoupling the optimization dynamics along flat and sharp directions to accelerate convergence to flat minima - appears to be quite novel and creative. To the best of my knowledge, this specific approach has not been explored in the optimization and generalization literature before. The authors motivate IRE through an intuitive illustrative example, showing how selectively increasing the learning rate along flat directions can speed up the implicit regularization. Adapting this core idea into a practical algorithm that can be efficiently incorporated into existing optimizers is also a notable contribution.\n\nHowever, the paper could benefit from a more detailed discussion comparing and contrasting IRE with other approaches that aim to improve sharpness-aware minimization, such as those reducing the computational cost of SAM (e.g., Kwon et al., 2021; Liu et al., 2022). This would help further highlight the novelty and uniqueness of IRE.\n\n## Quality\nThe paper presents a solid mix of conceptual explanations, practical algorithm design, extensive experiments, and theoretical analysis. The illustrative example in Section 2 provides a clear and intuitive understanding of the mechanism behind IRE. The practical implementation of IRE in Section 3, particularly the efficient approximation of the projection operator using diagonal Hessian estimates, demonstrates the authors' attention to computational efficiency. \n\nThe experimental results on image classification and language modeling tasks are extensive and convincing, showing consistent improvements in generalization performance across various models, datasets, and base optimizers. The surprising 2x speedup over AdamW in Llama pre-training is particularly impressive and warrants further investigation.\n\nThe theoretical analysis in Section 5, proving the substantial acceleration of IRE over SAM in minimizing the trace of Hessian, adds rigor to the empirical findings. The proofs leverage reasonable assumptions and provide non-asymptotic guarantees.\n\nHowever, the paper could be strengthened by providing more insights and discussions on the potential limitations and failure cases of IRE. For example, are there scenarios where the diagonal Hessian approximation might be less effective? How sensitive is IRE to the choice of hyperparameters (e.g., $\\lambda$ and $\\gamma$)?\n\n## Clarity\nOverall, the paper is well-structured and clearly written. The main ideas, algorithms, and results are presented in a logical flow, making it easy for readers to follow. The use of figures (e.g., Figure 1) and illustrative examples enhances the clarity of the exposition. The mathematical notations and definitions are introduced appropriately and used consistently throughout the paper.\n\nOne area that could be improved is the description of the experimental setup and implementation details. While the paper provides references to the appendix for more details, including some key information (e.g., hyperparameter tuning ranges, model architectures) in the main text would make the experiments more self-contained and easier to interpret.\n\n## Significance\nThe proposed IRE framework has the potential to make a significant impact in the field of deep learning optimization and generalization. By providing a principled and efficient way to accelerate convergence to flat minima, IRE can lead to models with better generalization performance and faster training times. The consistent improvements demonstrated across a range of vision and language tasks suggest that IRE could be widely applicable.\n\nMoreover, the theoretical analysis of IRE's acceleration over SAM opens up new avenues for understanding and improving sharpness-aware optimization. The non-asymptotic guarantees provide a solid foundation for further theoretical investigations.\n\nThe 2x speedup achieved by IRE in Llama pre-training is particularly significant, given the computational challenges in training large language models. If these speedups can be reliably reproduced and scaled to larger models and datasets, IRE could meaningfully contribute to advancing the state-of-the-art in language model training.\n\nTo fully assess the significance of IRE, it would be valuable to see more comparisons with other state-of-the-art optimizers and regularization techniques. Additionally, evaluating the downstream performance of models trained with IRE (e.g., fine-tuning Llama on benchmarks) would provide a more comprehensive understanding of its impact.", "weaknesses": "### Comparison with Related Works\nWhile the paper presents a novel approach to enhancing implicit regularization, it could benefit from a more detailed comparison with related works. The authors should discuss how IRE differs from and improves upon other techniques that aim to accelerate convergence to flat minima or reduce the computational cost of sharpness-aware minimization (e.g., Kwon et al., 2021; Liu et al., 2022). This would help to better highlight the novelty and advantages of IRE.\n\n### Limitations and Failure Cases\nThe paper could be strengthened by providing a more in-depth discussion of the potential limitations and failure cases of IRE. For example:\n- Are there scenarios where the diagonal Hessian approximation might be less effective or lead to suboptimal results?\n- How sensitive is IRE to the choice of hyperparameters (e.g., $\\lambda$ and $\\gamma$)? Is there a risk of instability or divergence if these hyperparameters are not tuned properly?\n- Are there any particular types of models, datasets, or tasks where IRE may not provide significant benefits or even hurt performance?\n\nAddressing these questions would help readers better understand the scope and applicability of IRE.\n\n### Experimental Setup and Implementation Details\nThe description of the experimental setup and implementation details could be improved. While the appendix contains additional information, it would be helpful to include key details in the main text, such as:\n- The specific hyperparameter tuning ranges for $\\lambda$ and $\\gamma$ used in the experiments\n- The architectures of the models used (e.g., ResNet and ViT variants)\n- The data augmentation and preprocessing techniques applied\n\nIncluding these details would make the experiments more self-contained and easier to interpret and reproduce.\n\n### Comparisons with State-of-the-Art Optimizers\nTo fully demonstrate the significance of IRE, it would be valuable to include comparisons with a broader range of state-of-the-art optimizers and regularization techniques. While the paper shows consistent improvements over SGD, AdamW, and SAM, it would be informative to see how IRE performs compared to other recent approaches, such as:\n- Adaptive gradient methods like Adagrad (Duchi et al., 2011), Adam (Kingma and Ba, 2014), and their variants\n- Second-order optimization methods like K-FAC (Martens and Grosse, 2015) and Shampoo (Gupta et al., 2018)\n- Other regularization techniques like weight decay, dropout, and label smoothing\n\nThese comparisons would provide a more comprehensive understanding of IRE's performance and potential advantages over existing methods.\n\n### Downstream Performance Evaluation\nWhile the paper demonstrates impressive speedups in Llama pre-training, it would be informative to evaluate the downstream performance of the models trained with IRE. For example, fine-tuning the pre-trained Llama models on benchmark tasks like language understanding, question answering, or text generation would provide insights into the practical impact of IRE on model quality and generalization.\n\nHowever, given the computational cost and time constraints of the rebuttal period, it may not be feasible to conduct extensive downstream evaluations. In this case, the authors could discuss this limitation and propose it as a direction for future work.\n\n### Theoretical Analysis\nThe theoretical analysis in Section 5 provides valuable guarantees for IRE's acceleration over SAM. However, the paper could benefit from a more intuitive explanation of the key assumptions and their implications. For example:\n- Discussing the practical significance of Assumption 5.1 (manifold of minimizers) and how it relates to the empirical observations in deep learning\n- Providing a high-level interpretation of the non-asymptotic bounds and their dependence on the hyperparameters (e.g., $\\eta$, $\\rho$, $\\lambda$)\n\nAdditionally, exploring the theoretical connections between IRE and other optimization techniques (e.g., momentum, adaptive methods) could provide further insights into its behavior and potential extensions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's recognition of our work and helpful comments. Below, we offer detailed responses to the reviewer’s questions:\n\n- **W\\&Q1. Comparison with related works.** \n\n  **Response:** We thank the reviewer for this question and will provide a more detailed comparison with related algorithms in the revised version. First, we clarify that the mechanism behind IRE in accelerating sharpness reduction is fundamentally different from that of SAM. Moreover, IRE can be integrated with generic base optimizers, including SGD, Adam, and SAM. In contrast, the SAM variants, including Kwon et al. (2021) and Liu et al. (2022), are specifically designed to improve SAM.\n    \n\n- **W\\&Q2. Limitations and failure cases.** \n\n  **Response:** Thank the reviewer for the insightful question. We acknowledge  that diagonal Hessian may not always provide accurate second-order information for identifying the flat (zero-curvature) directions. Recall that IRE iterates as $\\theta_{t+1} = \\theta_t - \\eta (1+\\kappa P) g_t$, where $P$ represents the projection to flat directions. When $\\kappa=0$, it recovers the base optimizer. For problems where  diagonal Hessian provides relatively accurate second-order information,  a large $\\kappa$ can be used. Conversely, for problems where diagonal Hessian is less informative, we can set $\\kappa$ to be small or even to zero. This approach allows to control over the influence of whether the diagonal Hessian approximation is effective. This idea is analogous to the damping technique adopted in Newton's methods with an approximate Hessian.   Specifically, we tune $\\kappa\\in$\\{1,2\\}  for CNNs, $\\kappa\\in$\\{2,3,4\\} for Llama pre-training, and $\\kappa\\in$\\{20,50\\} for ViTs.\n\n- **W\\&Q3. Experimental setup and implementation details.** \n\n  **Response:** Thank the reviewer for this question. For image classification tasks, the models including ResNets and ViT-T/S follows the ones in Muller et al. (2023). For LLM pre-training, we use the Llama models from HuggingFace. For image data pre-processing, we apply the default normalization technique: *transforms.Normalize(mean, std)*. The search ranges of hyperparameters, as well as the data augmentation methods are specified in Appendix B. In the revised version, we will carefully incorporate all these experimental details in the main text.\n\n\n- **W\\&Q4. Comparisons with state-of-the-art optimizers.** \n\n  **Response:** Thank the reviewer for this question. We clarify that IRE is a generic framework for boosting optimizer's implicit bias and can be integrated with general optimizers. Due to space and resource limit, we have only compared IRE with the popular algorithm including SGD, SAM and AdamW. Specifically, we show that IRE can improve the generalization in image classification task and convergence of AdamW in LLM pre-training. As for other optimizers mentioned by the reviewer, Adagrad has been replaced by Adam(W) in most fields; second-order algorithms like K-FAC and Shampoo is computational expensive and it is less  popular in training large models.\n\n  About regularization, IRE is proposed to enhance implicit regularization and thus, orthogonal to those explicit regularization, including weight decay, dropout and label smoothing. It is worth mentioning that our experiments incorporate various regularization techniques such as weight decay and label smoothing. Please refer to the Appendix for more details. \n\n- **W\\&Q5. Downstream performance evaluation.**\n\n  **Response:** Thank the reviewer for this constructive comment. It is valuable to explore the performance of the pretrained models in downstream tasks. However, due to the limited time for response, we could not fully evaluate the downstream tasks for pretrained models, but we conducted preliminary experiments to explore this issue.\n  - Notably, \"Liu et al. (2023) Same pretraining loss...\" points out a strong correlation between the flatness and downstream task performance among models. Specifically, it implies that *for models with the same pre-training loss, flatter solutions yield better performance on downstream tasks.*\n  - Inspired by this observation, we evaluated the flatness, measured by the trace of Hessian, of models pre-trained by AdamW and AdmIRE. Due to time constraints, we only focus on the experiments in Fig.3 (left), i.e., training Llama (60M) on wiki-103 dataset using either AdamW or AdmIRE. The experimental results, shown in the **attached PDF** in our **Response to All Reviewers**, demonstrate that AdmIRE not only achieves the same loss in *only half the iterations* required by AdamW, but also the solutions found by AdmIRE are *significantly flatter* than that found by AdamW.\n\n  In the revised version, we will include a comprehensive evaluation of downstream tasks.\n\n\n- **W\\&Q6. Theoretical analysis.**\n\n  **Response:** Thanks the reviewer for the question.\n  - We clarify that Assumption 5.1 essentially only assume: 1) the loss $\\mathcal{L}(\\cdot)$ is  $C^4$ smooth and 2) the global minima are connected. For over-parametrized models, this connectivity assumption on the minima manifold has been empirically verified in works such as Draxler et al. (2018) and Garipov et al. (2018), and theoretically supported in Cooper (2018). We refer to lines 267-270 of our manuscript for more details.\n  - Our main theoretical results are explained in lines 286-369, where we provide an intuitive interpretation: IRE accelerates the \"effective dynamics\" of base optimizers in flat directions, achieving faster reduction of sharpness; moreover, IRE maintains stability by not affecting movement along sharp directions. This dual effect ensures IRE can improve generalization performance without compromising training stability.\n  - As for other optimization techniques, our experiments have incorporated momentum, weight decay, and the adaptive method (AdamW). However, these techniques poses new challenges to theoretical analysis, which we leave as future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IIIK7ovueh", "reviewer": "Reviewer_Jrgs", "summary": "This paper proposes a new optimization scheme for deep learning.  The idea is to periodically estimate the Hessian diagonal and use a larger learning rate on the smaller-curvature parameters.   The paper argues that this algorithm enhances the implicit curvature regularization of the optimization algorithm while not degrading stability (which is what would happen if you used a larger learning rate on all parameters).  The paper first motivates the scheme by analyzing a toy setting.  The paper then shows empirically that adding this scheme to SGD and to SAM results in better generalizing and lower-curvature solutions, where curvature is quantified using Hessian trace.   The paper then shows that adding this scheme to Adam accelerates the _optimization_ of Llama transformers, and leaves an explanation why to future work.  Finally, the paper theoretically analyzes this scheme for SAM in the manifold-of-global-minima setting, and proves that this scheme accelerates the drift towards flatter regions of the manifold.", "questions": "Shouldn't this algorithm also speed up _optimization_?  The algorithm uses a larger learning rate on some coordinates, which should also have an optimization effect.  (This would not be visible in the analysis setting of manifold-of-global-minima because no optimization is occurring in that setting.)\n\nHow would you define IRE at the most abstract level?  If we define an abstract version of preconditioned gradient descent as $\\theta_{t+1} = \\theta_t - P g_t$ (where vanilla GD with learning rate $\\eta$ corresponds to $P = \\eta I$), is the idea that it's better to use $P_2$ rather than $P_1$ provided that $P_2 \\succeq P_1$ and both are locally stable i.e. $\\lambda_{\\max}(P H) \\le 2$ for $P \\in $ {$P_1, P_2$}?", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the idea is novel, and the experiments are promising, so this paper could stimulate significant research in the future.", "weaknesses": "One weakness is that the motivation comes from the stylized example in section 2, which may be unrealistic for general deep learning optimization.  For example, the idea of \"sharp coordinates\" only makes sense in this stylized example.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the appreciation of our work and insightful comments. Below, we address the reviewer’s questions in detail.\n\n- **W1.** One weakness is that the motivation comes from the stylized example in section 2, which may be unrealistic for general deep learning optimization. For example, the idea of \"sharp coordinates\" only makes sense in this stylized example.\n\n  **Response:** We thank the reviewer for this insightful question.\n\n  - **Clarification.** We clarify that the idea behind IRE is to distinguish between \"sharp/flat directions\" rather than \"sharp/flat coordinates\". The  use of a diagonal Hessian in Section 2  is only for simplification purpose. In fact,  the theoretical analysis in Section 5 considers a more general setting  without assuming a diagonal Hessian (see Assumption 5.1). \n        \n  - **Practical Considerations.** In the practical implementation described in Section 3, we indeed use the diagonal Hessian to approximate the full Hessian, aiming to reduce computational costs. The diagonal Hessian is commonly acknowledged as a reasonable approximation of the full Hessian in deep neural networks and is frequently employed in the design of deep learning optimizers, as noted in references [1][2]. For a more detailed discussion, please refer to lines 156-162 in our manuscript.\n\n- **Q1.** Shouldn't this algorithm also speed up optimization? The algorithm uses a larger learning rate on some coordinates, which should also have an optimization effect. (This would not be visible in the analysis setting of manifold-of-global-minima because no optimization is occurring in that setting.)\n\n  **Response:** We thank the reviewer for this insightful comment and will include a   discussion on this issue in our revised version. The motivation of IRE is to speed up the sharpness reduction, which only requires to increase learning rate along completely flat (zero-curvature) directions. However, practical implementation may also increase  learning rate along directions with small but non-zero curvatures, which can further speed up loss convergence. However, explaining why  this approach can provide a significant acceleration is  unclear.\n\n- **Q2.** How would you define IRE at the most abstract level? If we define an abstract version of preconditioned gradient descent as $\\theta_{t+1}=\\theta_t-Pg_t$ (where vanilla GD with learning rate $\\eta$ corresponds to $P=\\eta I$), is the idea that it's better to use $P_2$ rather than $P_1$ provided that $P_2\\succeq P_1$ and both are locally stable i.e. $\\lambda_{\\max}(PH)<2$ for $P\\in\\{P_1,P_2\\}$?\n\n  **Response:** We are grateful to the reviewer for sharing this high-level perspective, which can help  explain the acceleration of loss convergence observed with IRE. Given a base optimizer with update $\\theta_{t+1}=\\theta_t-P_1g_t$, we aim to improve it by selecting a better $P_2=QP_1$, while ensuring $\\lambda_{\\max}(P_2 H)<2$ to maintain training stability. Let $P_{flat}$ be the projection matrix into the flat (zero-curvature) directions. In our IRE approach, we choose $P_2=(I+\\mu P_{flat}) P_1$, which meets the condition $P_2\\succeq P_1$ mentioned by the reviewer. However, we remark that our original motivation of  this specific $P_2$ is to accelerate effective dynamics on the manifold of global minima, thereby accelerating the sharpness reduction.\n\n\n[1] George et al. Fast approximate natural gradient descent in a Kronecker-factored eigenbasis. (NeurIPS 2018)\n\n[2] Liu et al. Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. (ICLR 2024)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
### **Global Response to All Reviewers.**

- We express our sincere gratitude to all reviewers for appreciating our results, i.e,
  - A novel algorithm framework (IRE), which can improve both generalization and optimization, by enhancing the implicit regularization.
  - Experimentally, IRE consistently improves generalization performance for image classification. Remarkably, IRE achieves a $2\times$ speed-up in pre-training Llama models.
  - Theoretically, we demonstrate that IRE substantially accelerate convergence towards flat minima in sharpness-aware minimization (SAM).

- We also thank appreciate the reviewers for their valuable comments and suggestions for improving our paper. In our revised version, we will correct all typos, provide complete experimental settings and results, and incorporate the discussions with the reviewers.

- **New experiments.** To further explore the performance of IRE in various settings, we have conducted 3 new experiments. As suggested by Reviewer 69nS and QWSd, we examined: 
  -  (i) Whether IRE can accelerate the convergence of AdamW in the finetuning phase;
  - (ii) During the pre-training phase, whether the solution found by AdmIRE has better properties than that found by AdamW.

  Additionally, to supplement our results on generalization, we  conducted an additional experiment:
  - (iii) Whether IRE can improve the generalization performance of AdamW when training ViT on ImageNet.

  All of these results are reported in the **attached PDF**, which further verifies remarkable performance of IRE under more settings. 

- We have addressed each concern raised by the reviewers through separate responses provided below.

Meta-review: This paper proposes a method for decoupling the optimization dynamics arising from sharp and flat directions, amplifying the dynamics along flat directions. This allows the optimization algorithm to navigate along flat directions faster without jeopardizing the optimization's stability due to sharp directions. The resulting algorithm is shown to improve both generalization performance and convergence speed, supported by empirical and theoretical evidence.

The paper received a detailed review, in which reviewers found the paper well-written and the proposed idea original with the potential for a significant impact. The reviewers also raised questions and concerns about the computational cost of the algorithm's steps, related works, empirical assessment, and the relevance of the stylized example in section 2 to deep learning. The authors' response to this feedback was clear and complete enough that reviewers did not need to ask follow-up questions. All reviewers rated the paper as accept. In agreement with them, I believe this paper has an interesting idea with potential impact, and I recommend acceptance.
Decision: Accept (poster)