Title: Improving Generalization and Convergence by Enhancing Implicit Regularization
Venue: NeurIPS 2024
Accepted: yes
Keywords: implicit regularization

convergence

generalization

deep learning
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

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 × 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).

1 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 of 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 to 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 et al., 2016; Wu et al., 2017), and flat minima generalize better (Hochreiter and

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 et al. (2020), Li et al. (2022), and Ma et al. (2022). Consequently, practitioners often use a large learning rate (LR) 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 convergence towards flatter minima. As suggested by works like Blanc et al. (2020), Li et 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 a 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 . 0 × 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 Θ(1 /ρ ) -time acceleration over the base SAM algorithm in minimizing the trace of Hessian, where ρ ∈ (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 in 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 a dynamical stability perspective. Moreover, in-depth analysis of SGD dynamics near global minima shows that the SGD noise (Blanc 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˛ ebski 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-aware 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., 2011), 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 et al., 2022b;a), AdamW(Adam+decoupled 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, a scalable second-order optimizer. In this paper, we instead empirically demonstrate that IRE can accelerate the convergence of AdamW in the pre-training of Llama models.

1.2 Notations

Throughout this paper, let L : R p ↦→ R ⩾ 0 be the function of total loss, where p denotes the number of model parameters. For a C 2 -submanifold M in R p , we denote the tangent space of M at θ ∈ M as T θ M , which is a linear subspace in R p . For f ∈ C 1 ( M ) and θ ∈ M , let ∇ M f ( θ ) := Q T θ M ∇ f ( θ denote the Riemannian gradient, where Q T θ M : R p ↦→ R p denotes the orthogonal projection to T θ M . For a symmetric matrix A ∈ R p × p , its eigen pairs are denoted as { ( λ i , u i ) } i ∈ [ p ] with the order λ 1 ⩾ · · · ⩾ λ p . We use P i : j ( A ) = ∑ j k = i u k u ⊤ k to denote the projection operator onto span { u i , . . . , u j } . Denote N ( µ , Σ) as the Gaussian distribution with mean µ and covariance matrix Σ , and U ( S ) as the uniform distribution over a set S . Given a vector h = ( h 1 , . . . , h p ) , let | h | = ( | h 1 | , . . . , | h p | ) . We denote by 1 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:

where v ∈ R m , u ∈ R p -m , and θ = vec ( u , v ) ∈ R p . We assume H ( · ) ∈ C 2 ( R p -m ) and inf u λ min ( H ( u )) > 0 . Then, the minimizers of L ( · ) form a m -dim manifold M = { ( u , v ) : v = 0 } and the Hessian at any θ ∈ M is given by ∇ 2 L ( θ ) = ( 0 0 0 H ( u ) ) . For clarity, we shall call u and v 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 ( x ; θ ) = a ⊤ ϕ ( x ; W ) be a 2LNN with θ = vec ( W, a ) . Then L ( θ ) = E ( x ,y ) [( f ( x ; θ ) -y ) 2 ] / 2 = a ⊤ E x [ ϕ ( x ; W ) ϕ ( x ; W ) ⊤ ] a / 2 =: a ⊤ H ( W ) a / 2 .

For breviety, we further assume H ( u ) = diag( λ ( u )) with λ ( u ) = ( λ 1 ( u ) , · · · , λ m ( u )) . In this case, the GD dynamics can be naturally decomposed into the flat and sharp directions as follows

where ⊙ denotes the element-wise multiplication of two vectors.

The implicit sharpness regularization. From Eq. (2), we can see that 1) the flat direction u t 's dynamics monotonically reduces the sharpness λ ( u ) as long as v t is nonzero; 2) the sharp direction v t 's dynamics determines the speed of sharpness reduction. The larger | v t | is, the faster the curvature λ ( u ) decreases. Particularly, when near convergence, we have | v t | = o (1) and thus the implicit sharpness reduction is very slow during the late phase of GD. Figure 1a provides a visualization of this slow implicit sharpness reduction.

Accelerating the sharpness reduction. Inspired by the above analysis, we can accelerate the sharpness reduction by speeding up the flat directions' dynamics. To this end, there are two approaches:

Naively increasing the global learning rate η (fail). Increasing η accelerates the dynamics of u t , but the largest allowed η is constrained by curvatures of sharpest directions. In GD (2), to

)

0

.

990

0

.

995

1

.

000

-

0

.

01

0

.

00

0

.

01

M

η

= 1

)

the flattest minimum

initialization

θ

0

(a) GD with

η

= 1

.

-

0

.

5

0

.

0

0

.

5

1

.

0

1

.

5

-

0

.

5

0

.

0

0

.

5

η

= 2

)

(b) GD with

η

= 2

.

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

-

0

.

01

0

.

00

0

.

01

IRE (

κ

= 9

e

1

)

IRE (

κ

= 9

e

2

)

IRE (

κ

= 9

e

3

)

η

= 1

)

the flattest minimum

(c) IRE with various

κ

's v.s. GD with

η

= 1

.

maintain training stability, η must be smaller than 2 / max i λ i ( u t ) . Otherwise, v t 's dynamics will blow up. As illustrated in Figure 1b, setting η = 2 leads to divergence, whereas η = 1 ensures convergence.

Increasing only the flat directions' learning rate (our approach, IRE) . Specifically, for GD (2), the GD-IRE dynamics is given by

where κ > 0 controls the enhancement strength. In GD-IRE (3), u t 's dynamics is (1 + κ ) faster than that of GD (2). Notably, the sharp directions' dynamics ( v t ) are unchanged. The choice of κ only needs to maintain the stability of flat directions' dynamics, for which, we can always take a significantly large κ to enhance the sharpness regularization. As demonstrated in Figure 1c, 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.

3 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: θ t +1 = θ t -η g t , the corresponding IRE modification is given by

where κ denotes the enhancement strength and P t : R p → R p projects g t into the flat directions of the landscape. The flat directions and corresponding projection operator P t can be estimated using the Hessian information.

However, estimating the full Hessian matrix ∇ 2 L ( θ t ) ∈ R p × p is computationally infeasible. Consequently, we propose to use only the diagonal Hessian diag( ∇ 2 L ( θ t )) ∈ R p to estimate P t . Let h t ∈ R p be an estimate of the diagonal Hessian. Then, we perform the projection as follows

where γ ∈ (0 , 1) and Top small ( | h t | , γ ) returns the ⌊ p · γ ⌋ -th smallest value in | h t | . Note that n t ∈ R p denotes a mask vector and the above approximate projection essentially masks the top(1 -γ ) sharp coordinates out. As such, the projection (5) will retain the topγ flat coordinates . Noticing that in DNNs, there are much more flat directions than sharp directions (Yao et al., 2020), we thus often use γ > 0 . 5 in practice.

A light-weight estimator of the diagonal Hessian. Let ℓ ( · , · ) be the cross-entropy loss. Given an input data x ∈ R d x and label y ∈ R d y , let the model's prediction be f ( x ; θ ) ∈ R d y . 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 t = diag( F ( θ t )) , 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 { ( x b , y b ) } B b =1 , the empirical diagonal Fisher is given by diag( ˆ F ( θ )) = 1 B ∑ B b =1 ∇ ℓ ( f ( x b ; θ ); ˆ y b ) ⊙ ∇ ℓ ( f ( x b ; θ ); ˆ y b ) , where ˆ y b ∼ softmax( f ( θ ; x b )) . However, as noted by Liu et al. (2024), implementing this estimator is computationally expensive due to the need to calculate B single-batch gradients. Liu et al. (2024) proposed a more convenient estimator diag( ˆ F eff ( θ )) , only requires computing the mini-batch gradient ∇ ˆ L B ( θ ) = 1 B ∑ B b =1 ∇ ℓ ( f ( x b ; θ ); ˆ y b ) with ˆ y b ∼ softmax( f ( x b ; θ )) :

According to Liu et al. (2024), this estimator is an unbiased estimate of the empirical diagonal Fisher, i.e., E ˆ y [diag( ˆ F eff ( θ ))] = E ˆ y [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 as 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 t using (6) requires computational resources roughly equivalent to one back-propagation. Consequently, by setting K = 10 in Algorithm 1 (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, we 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-8 (Zagoruyko and Komodakis, 2016) on CIFAR-10 dataset (Krizhevsky and Hinton, 2009) by SAM-IRE (with K = 10 , varying κ and γ ). Here, we incoporate IRE into SAM starting from the 30 -th epochs. We vary γ ∈ { 0 . 8 , 0 . 9 , 0 . 95 } and κ ∈ { 0 , 2 , 5 , 10 } . Regarding the learning rate (LR), both constant LR and decayed LR are considered. The sharpness is measured by Tr( ∇ 2 L ( θ )) . 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

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)).

0.0

2.0

5.0

10.0

κ

80%

90%

95%

γ

1768

1685

1478

1428

1768

1610

1392

1214

1768

1602

1241

1206

sharpness

0.0

2.0

5.0

10.0

κ

80%

90%

95%

γ

2774

2760

2626

2581

2774

2758

2609

2553

2774

2735

2588

2365

sharpness

0.0

2.0

5.0

10.0

κ

80%

90%

95%

γ

0.066

0.037

0.023

0.016

0.066

0.035

0.021

0.014

0.066

0.033

0.019

0.013

training loss

0.0

2.0

5.0

10.0

κ

80%

90%

95%

γ

93.48

94.06

94.55

94.97

93.48

94.12

94.74

95.10

93.48

94.19

94.71

95.27

test accuracy (

%

)

Figure 2: Training WRN-16-8 on CIFAR-10 by SAM-IRE with varying γ, κ . Particularly, the case of κ = 0 correspond to the standard SAM.

4.1.2 IRE can consistently improve generalization

Convolutional Neural Networks (CNNs). In this experiment, we first consider the classification of CIFAR-{10,100} with WideResNet-28-10 (Zagoruyko and Komodakis, 2016) and ResNet-56 (He et al., 2016). Both SGD and SAM optimizers are adopted. All the experiments use base data augmentation and label smoothing. For SGD-IRE/SAM-IRE, we fix K = 10 , and tune hyperparameters γ and κ via a grid search over γ ∈ { 0 . 99 , 0 . 9 , 0 . 8 } and κ ∈ { 1 , 2 } . 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-IRE when the training loss approaches 0.1. The other experimental details are deferred to 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 C and the results are shown in Table 2.

Vision Transformers (ViTs). We also examine the impact of IRE on generalization of ViT-T and ViT-S (Dosovitskiy et al., 2020) on CIFAR100. The default optimizers used are AdamW and SAM(Mueller et al., 2024). Strong data augmentations (basic + AutoAugment) are utilized. The total epochs are set to 200 and we switch from AdamW/SAM to AdmIRE/SAM-IRE when the training loss approaches 0 . 5 . For AdmIRE/SAM-IRE, we fix K = 10 , and tune hyperparameters γ and κ via a grid search over γ ∈ { 0 . 99 , 0 . 9 , 0 . 8 } and κ ∈ { 20 , 50 } . Other experimental details are deferred to Appendix C. The results are shown in Table 3.

Additionally, we evaluate IRE for training ViT-S on 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.

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, we employ the AdamW optimizer with hyperparameters β 1 = 0 . 9 , β 2 = 0 . 95 and weight decay λ = 0 . 1 (Touvron et al., 2023). The learning rate strategy includes a warm-up phase followed by a cosine decay scheduler, capped at lr_max . In each experiment, we tune lr_max only for AdamW 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 is conducted to verify both the computational efficiency and the robustness of hyperparameters ( γ, κ ) in IRE for pre-training tasks. Specifically, we train a 2-layer decoder-only Transformer (8M) on the Wikitext-2 dataset (4.3M) (Merity et al., 2016) by AdamW and AdmIRE (with K = 10 and varying γ, κ ). The total training duration is 100k steps, including a 3k-step warm-up phase.

First, we tune lr_max in AdamW, identifying the optimal lr_max = 6e-4 . Subsequently, we train both AdamW and AdmIRE using this lr_max .

Computational efficiency. As shown in Table 5, AdmIRE with K = 10 (estimating the projection mask every 10 steps) is computationally efficient: the average time per step of AdmIRE is only 1 . 12 times that of AdamW, corresponding to the theoretical estimation ( 1 . 1 times).

Robustness to hyperparameters. Figure 3 shows that AdmIRE, with varying γ and κ , consistently speeds up the pretraining. Remarkably, with the best configuration, AdmIRE can achieves 5.4 × speedup than well-tuned AdamW.

More experimental details and results are deferred to Appendix C.

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, we examine the performance of AdmIRE in training Llama models of various sizes across various datasets:

Llama (60M) on wikitext-103 (0.5G). Wikitext-103 (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), a 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) on openwebtext (38G). Openwebtext (Gokaslan and Cohen, 2019), an opensource 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 to maintain the training stability (Pascanu et al., 2012). First, we tune lr_max in AdamW for each of the three experiments, separately. The optimal lr_max

3k

25k

50k

75k

100k

Num of steps

2

.

5

3

.

0

training loss

5

.

4

×

Speedup

AdamW

AdmIRE(

κ

= 2

,γ

=0

.

99

)

AdmIRE(

κ

= 2

,γ

=0

.

9

)

AdmIRE(

κ

= 2

,γ

=0

.

8

)

AdmIRE(

κ

= 5

,γ

=0

.

99

)

AdmIRE(

κ

= 5

,γ

=0

.

9

)

AdmIRE(

κ

= 5

,γ

=0

.

8

)

identified for these three experiments is all 6e-4 . Then, both AdamW and AdmIRE are trained using this optimal lr_max . For more details, please refer to Appendix C.

AdmIRE is 2 × 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 AdamW for all three cases.

0

25k

50k

75k

100k

num of steps

2

.

4

2

.

6

2

.

8

3

.

0

3

.

2

validation loss

2

.

1

×

Speedup

Llama (60M) on wikitext-103

AdamW(50k)

AdamW(100k)

AdmIRE(50k)(

κ

=2

,γ

=0

.

7

)

AdmIRE(50k)(

κ

=3

,γ

=0

.

6

)

0

50k

100k

0

6e-4

50k

100k

0

15k

30k

45k

60k

num of steps

2

.

70

2

.

75

2

.

80

2

.

85

2

.

90

2

.

95

3

.

00

validation loss

2

.

3

×

Speedup

Llama (119M) on minipile

AdamW(30k)

AdamW(60k)

AdmIRE(30k)(

κ

=2

,γ

=0

.

6

)

AdmIRE(30k)(

κ

=5

,γ

=0

.

5

)

AdmIRE(30k)(

κ

=4

,γ

=0

.

6

)

0

30k

60k

0

6e-4

30k

60k

0

25k

50k

75k

100k

num of steps

2

.

8

2

.

9

3

.

0

3

.

1

3

.

2

validation loss

2

.

1

×

Speedup

Llama (229M) on openwebtext

AdamW(50k)

AdamW(100k)

AdmIRE(50k)(

κ

=4

,γ

=0

.

6

)

0

50k

100k

0

6e-4

50k

100k

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-training. Specifically, we compare the sharpness of solutions, Tr( ∇ 2 L ( θ )) , found by AdamW/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 in only half the iterations required by AdamW, but also the solutions found by AdmIRE are significantly flatter than that found by AdamW.

Recently, Liu et 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 et al., 2023a) studies have validated that SAM algorithms exhibit superior sharpness regularization compared to (S)GD. In this section, we provide a theoretical analysis demonstrating that IRE can further enhance the sharpness regularization of SAM algorithms substantially.

5.1 Theoretical setups

Recall that L ( θ ) := 1 n ∑ n i =1 L i ( θ ) denote the total loss, where L i ( θ ) is the loss on the i -th data. Without loss of generality, we assume min θ L ( θ ) = 0 . We further make the following assumption:

Assumption 5.1 (Manifold of minimizers) . Assume that L ∈ C 4 ( R p ) , M := arg min θ L ( θ ) is a ( p -m )-dim C 2 -submanifold in R p for some m ∈ [ p ] , and rank( ∇ 2 L ( θ )) = m for any θ ∈ M .

The above connectivity assumption on the manifold of minimizers M has been empirically verified 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 in the theoretical analysis of implicit regularization (Fehrman et al., 2020; Li et al., 2022; Arora et 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.

Definition 5.2 (Limiting map of GF) . Consider the GF: d θ ( t ) d t = -∇L ( θ ( t )) starting from θ (0) = θ . Denote by Φ( θ ) := lim t →∞ θ ( 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 to some point in M . Formally, U := { θ ∈ R p : Φ( θ ) ∈ 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) is C 2 on U .

5.2 Theoretical results

The stochastic SAM (Foret et al., 2021) is given by

The generalization capability of standard SAM can be bounded by the average sharpness, L avg ( θ ) := E ξ ∼N ( 0 ,I ) L ( θ + ρ ξ / ∥ ξ ∥ ) (Foret et al., 2021). This leads researchers to also explore average SAM (Wen et al., 2023a; Zhu et al., 2023; Ujváry et al., 2022), which minimizes L avg :

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 θ 0 ∈ U . In Phase I ( t ⩽ T I ), we employ GF d θ t d t = -∇L ( θ t ) to ensure that the loss decreases sufficiently; then in Phase II ( T I < t ⩽ T I + T II := T ), we incorporate IRE into the standard / average SAM.

Effective dynamics: sharpness regularization. The implicit regularization of SAMs can be modeled using effective dynamics. In Phase II, θ t are close the manifold of minimizers M and let z t := Φ( θ t ) ∈ M . Then, the effective dynamics is given by { z t } T t = T I +1 , revealing how SAMs explore the manifold of minimizers M . Particularly, Wen et al. (2023a) showed that the effective dynamics of standard/average SAM are both

tangent space

T

z

t

M

-∇

Tr[

∇

2

L

(

z

t

)

/

2]

-∇M

Tr[

∇

2

L

(

z

t

)

/

2]

θ

t

zt

= Φ(

θ

t

)

M

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) η eff 's. A visual illustration of some quantities in (9) is provided in the figure above.

Summary of our theoretical results. In this section, we show that incorporating IRE into SAMs can significantly increase the effective LR η eff in (9) while maintaining the same training stability as SAMs. In Table 7, 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.

Remark 5.4 (The mechanism of IRE's success) . The success of SAM-IRE follows the same mechanism 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 + κ ) ; 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 as 1 /ρ .

5.2.1 IRE on average SAM: An Ω(1 /η 0 . 5 ) acceleration

We first consider IRE on average SAM. Let T I be the hitting time: T I := inf { t ⩾ 0 : ∥ θ t -Φ( θ t ) ∥ = O ( √ ηρ ) } . When running GF in Phase I, Definition 5.3 guarantees T I < ∞ . Thus, at the starting of Phase II, ∥ θ t -Φ( θ t ) ∥ = O ( √ ηρ ) . Furthermore, the following result holds for Phase II.

Theorem 5.5 (IRE on average SAM) . Suppose Assumption 5.1 holds. If η = O (1) and ρ = O ( √ η ) in SAM (8) , κ ⩽ 1 /ρ , and P t = P m +1: p ( ∇ 2 L ( θ t )) in IRE (4) , then with high probability at least 1 -T 2 II exp ( -Ω ( 1 / ( η + p -1 ))) , ∥ θ t -Φ( θ t ) ∥ = O ( √ ηρ ) holds for all T I ⩽ t ⩽ T . Furthermore, the effective dynamics of z t := Φ( θ t ) ∈ M satisfies:

Note that ρ = O ( √ η ) and κ can be as large as 1 /ρ . Consequently, the effective LR of minimizing the trace of Hessian can be selected as large as η eff = ( κ +1) ηρ 2 /p = O ( η 1 . 5 /p ) . In contrast, that of average SAM is at most O ( η 2 /p ) . The proof of Theorem 5.5 can be found in Appendix D.

5.2.2 IRE on standard SAM: An Ω(1 /ρ ) acceleration

This subsection delves into IRE on 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 ∇L i ( θ ) ( i ∈ [ n ] ), we need an additional assumption regarding the features on the manifold (see Setting E.1), 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 to 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 t < ∞ such that ∥ θ t -Φ( θ t ) ∥ = O ( η 1 -α ρ ) for any α ∈ [0 , 1) . We define T I as the hitting time: T I := inf { t ⩾ 0 : ∥ θ t -Φ( θ t ) ∥ = O ( η 1 -α ρ ) } . 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.1, if η, ρ = O (1) in SAM (7) , κ ⩽ 1 /ρ , and P t = P n +1: p ( ∇ 2 L ( θ t )) in IRE (4) , then with high probability at least 1 -T 2 II exp( -Ω(1 /η α )) , ∥ θ t -Φ( θ t ) ∥ = O ( η 1 -α ρ ) holds for all T I ⩽ t ⩽ T . Moreover, the effective dynamics z t = Φ( θ t ) ∈ M satisfies:

Taking κ = 0 and α = 0 recovers the result established in Wen et al. (2023a). However, κ can be as large as 1 /ρ , where IRE provides a Θ(1 /ρ ) -time acceleration over the standard SAM.

6 Conclusion

In this work, we propose a novel IRE framework to enhance the implicit sharpness regularization of base optimizers. Experiments demonstrate that IRE not 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://github.com/wmz9/IRE-algorithm-framework .

For future work, there are two urgent directions: 1) understanding why IRE can accelerate convergence, which may require studying the interplay between IRE and the Edge of Stability (EoS) (Wu et al., 2018; Jastrz˛ ebski et al., 2017; Cohen et al., 2021); and 2) conducting a larger-scale investigation into the acceleration of AdmIRE compared to AdamW in LLM pre-training, as well as the downstream performance of the LLMs 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 is supported in part by the National Key Basic Research Program of China (No. 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.

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 Ornstein-Uhlenbeck 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, Quoc V Le, and Ruslan Salakhutdinov. Transformer-XL: Attentive language models beyond a 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. Ieee, 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.11929 , 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:2110.03141 , 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.02054 , 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

Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (S) GD over diagonal linear networks: Implicit bias, large stepsizes and edge of stability. Advances in Neural Information Processing Systems , 36, 2024. 2

Benjamin Fehrman, Benjamin Gess, and Arnulf 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, 8, 9, 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.00027 , 2020. 7

Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information processing systems , 31, 2018. 8

Khashayar Gatmiry, Zhiyuan Li, Ching-Yao Chuang, Sashank Reddi, Tengyu Ma, and Stefanie Jegelka. The inductive bias of flatness regularization for deep matrix factorization. arXiv preprint arXiv:2306.13239 , 2023. 2

Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a Kronecker-factored eigenbasis. Advances in Neural Information Processing Systems , 31, 2018. 5

Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http://Skylion007.github.io/ OpenWebTextCorpus , 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 1 hour. arXiv preprint arXiv:1706.02677 , 2017. 2

Roger Grosse and James Martens. A Kronecker-factored approximate Fisher matrix for convolution layers. In International Conference on Machine Learning , pages 573-582. PMLR, 2016. 5

Suriya Gunasekar, Jason D 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, 6

Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung-Woo Ha. Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights. arXiv preprint arXiv:2006.08217 , 2020. 3

Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation , 9(1):1-42, 1997. 1

Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. arXiv preprint arXiv:1705.08741 , 2017. 2

Stanisław Jastrz˛ ebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. arXiv preprint arXiv:1711.04623 , 2017. 2, 10

Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. International Conference on Learning Representations , 2019a. 17

Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. Conference on Learning Theory , 2019b. 17

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 a primal-dual analysis. In Algorithmic Learning Theory , pages 772-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 on Machine Learning , pages 4860-4869. PMLR, 2021. 17

Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In 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. 1

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.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.html . 5

Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi-homogeneous neural networks. International Conference on Learning Representations , 2023. 17

Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. ASAM: Adaptive sharpnessaware minimization for scale-invariant learning of deep neural networks. In International Conference 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. PMLR, 2017. 18

Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms I: 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-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. 3, 5, 21

Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint arXiv:1908.03265 , 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.11692 , 2019b. 7

Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-aware 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 arXiv:1711.05101 , 2017. 3

Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint arXiv:1902.09843 , 2019. 3

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.05890 , 2019. 17

Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent on 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 normalization layers: Sharpness reduction. Advances in Neural Information Processing Systems , 35: 34689-34708, 2022. 2

Chao Ma and Lexing Ying. On linear stability of SGD and input-smoothness of neural networks. Advances in Neural Information Processing Systems , 34:16805-16817, 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 Learning , 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 arXiv:1609.07843 , 2016. 7

Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness-aware minimization stronger: A sparsified perturbation approach. Advances in Neural Information Processing Systems , 35:30950-30962, 2022. 2, 5

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:17749-17761, 2021. 2

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. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 3420-3428. PMLR, 2019b. 17

Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with a fixed learning rate. In The 22nd International Conference on 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 a convex programming problem with convergence rate o (1 /k 2 ) . Doklady Akademii Nauk SSSR , 269(3):543, 1983. 2

Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614 , 2014. 1

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. Understanding the exploding gradient problem. CoRR, abs/1211.5063 , 2(417):1, 2012. 7

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-Vivien, and Nicolas Flammarion. Implicit bias of SGD for diagonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems , 34:29218-29230, 2021. 18

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019. 7

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

David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. Nature , 323(6088):533-536, 1986. 1

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 Learning 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

Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning , pages 11391147. 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. 2

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.13971 , 2023. 2, 7, 20, 21

Szilvia Ujváry, Zsigmond Telek, Anna Kerekes, Anna Mészáros, and Ferenc Huszár. Rethinking sharpness-aware minimization as variational inference. arXiv preprint arXiv:2210.10452 , 2022. 9

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, Łukasz Kaiser, and Illia Polosukhin. Attention is 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.09371 , 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

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-aware 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. 2

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 on 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 preprint arXiv:1706.10239 , 2017. 1

Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over-parameterized 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 arXiv:2208.06677 , 2022a. 3

Zeke Xie, Xinrui Wang, Huishuai Zhang, Issei Sato, and Masashi Sugiyama. Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In International conference on machine learning , 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 arXiv:1605.07146 , 2016. 5, 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. Understanding 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 on Machine Learning , pages 43005-43036. PMLR, 2023. 9

Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems , 33:18795-18806, 2020. 3

Appendix

A Other Related Works

Other Implicit Biases. 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) showed that GD converges to max-margin solutions under exponentially-tailed loss on linearly separable data, albeit with an extremely slow rate O (1 / log t ) . Furthermore, Nacson et al. (2019c) studied this bias for SGD, Ji and Telgarsky (2019b) investigated linearly non-separable data, and Ji et al. (2020) analyzed the effects of the tail behavior of loss functions.

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 O (1 /t ) . Building on this, Ji et al. (2021); Wang et al. (2022) proposed momentum-based gradient methods which achieve a rate of O (1 /t 2 ) . Recently, Wang et al. (2024) established that the polynomial rates for most previous algorithms are tight, and proposed a progressive rescaling gradient descent method that achieves margin maximization exponentially fast O ( e -Ω( t ) ) .

The margin-maximization bias has also been studied for nonlinear models. Ji and Telgarsky (2019a); Gunasekar et al. (2018) examined deep linear networks, while Chizat and Bach (2020) focused on wide two-layer ReLU networks. Notably, Nacson et 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 quasi-homogeneous networks. For two-layer (leaky-)ReLU neural networks, Lyu et 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.

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 Flammarion (2023); Even et al. (2023) conducted fine-grained analyses of training dynamics, examining how initialization and step size impact (S)GD'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 Σ : d θ t = -∇L ( θ t ) d t + √ η Σ( θ t ) d W t . We consider that the noise covariance aligns with the Hessian near minima, i.e., Σ( θ ) = ( 0 d 0 0 σ 2 h ( u ) ) (where σ > 0 is a scalar), such as the label noise (Damian et al., 2021; Li et al., 2022). Then, the SDE above can be rewritten as

Implicit Sharpness Regularization. Intuitively, when v t is close to 0 , the speed of u t is much slower than v t due to v 2 t ≪ v t . Following Ma et al. (2022), when this speed separation is large, v t is always at equilibrium given u t . Solving the Ornstein-Uhlenbeck process about v t , we know the equilibrium distribution of v t is v ∞ ∼ N (0 , ησ 2 ) , and hence the dynamics u t (along the manifold) is

This derivation clearly shows the slow 'effective dynamics' u t along the manifold is a gradient 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 : min θ Tr( ∇ 2 L ( θ )) s . t . L ( θ ) ≈ 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 of optimization, after SGD reaches the equilibrium v ∞ ∼ N (0 , ησ 2 ) , the loss near the flat minimum is smaller because E v ∞ [ L ( θ )] = E v ∞ [ h ( u ) v 2 ∞ / 2] = ησ 2 h ( u ) / 2 ∝ Tr( ∇ 2 L ( u )) .

Q. How can we enhance the implicit sharpness regularization of SGD? A. Accelerating SGD's slow 'effective dynamics' u t 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 η will influence the dynamic stability of v t and the equilibrium v ∞ . Therefore, we need to accelerate the effective dynamics u t without affecting the dynamics of v t . Another main point is that SGD's effective dynamics u t can naturally minimize the sharpness implicitly, so we only need to enhance this property. To achieve this, we only need to correct the update direction in (10) from -∇L ( θ t ) to -( ∇L ( θ t ) + κP M ∇L ( θ t )) , where P M is the projection matrix to the manifold M and κ is a scalar. Using this new algorithm, the SDE dynamics corresponds to

Comparing (10) and (12), the dynamics of v t are the same, so they attain the same equilibrium distribution v ∞ ∼ N (0 , ησ 2 ) . As for the effective dynamics along manifold, (10) corresponds to the form:

(1+ κ ) times Enhancement. Comparing (13) and (11), it is clear that our new algorithm can enhance implicit sharpness regularization (1 + κ ) 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-16-8 on CIFAR-10 dataset by SAM-IRE (with K = 10 , varying κ and γ ). 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 5e-4, and the ρ in SAM is to 0.05, as in Foret et al. (2021). To evaluate the implicit flatness regularization of SAM itself, the momentum is set to 0.0. Regarding the learning rate (lr), we evaluate for both constant lr (within our theoretical framework) and decayed lr (common in practice though not covered by our theory). In the experiment in Fig 2 (a), a fixed lr 0 . 1 is used. In the experiment in Fig 2 (b)(c)(d), a step-decayed lr schedule is employed, starting at 0 . 1 and reducing lr by a 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 κ in 0 (original SAM), 2 , 5 , 10 .

The flatness measure, Tr( ∇ 2 L ( θ )) , is approximated by the trace of Fisher Tr( F ( θ )) . Specifically, we use diag( ˆ F eff ( θ )) in (6) for the estimate because E ˆ y [diag( ˆ F eff ( θ ))] = E ˆ y [diag( ˆ F ( θ ))] and thus,

Moreover, the first ' ≈ ' above takes ' = ' when L ( θ ) = 0 .

In this section, all experiments were conducted using a single A800 GPU.

C.1.2 Experimental details in Section 4.1.2

Experiments for CNNs on CIFAR-10/CIFAR-100. We first evaluate the impact of IRE on generalization of baseline models (WideResNet-28-10 and ResNet-56) and default optimizers (SGD and SAM) on CIFAR-{10,100}. For the base optimizers, SGD and SAM, cosine learning rate decay is adopted with an initial lr 0 . 1 . For other training components, we follow the settings in Foret et al. (2021): basic data augmentations and 0 . 1 label smoothing; for both SGD and SAM, the momentum is set to 0 . 9 , the batch size is set to 128 , and the weight decay is set to 5e-4; for SAM, ρ is set to 0 . 05 for CIFAR-10 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-IRE when the training loss approaches 0.1. For SGD-IRE/SAM-IRE, we fix K = 10 , and tune hyperparameters γ and κ via a grid search over γ ∈ { 0 . 99 , 0 . 9 , 0 . 8 } and κ ∈ { 1 , 2 } . The results are reported in Table 1.

Experiments without finely tuned hyperparameters. A high sensitivity to the choice of hyperparameters would make a method less practical. To demonstrate that our IRE performs even when κ, γ are not finely tuned , we conduct experiments using fixed γ = 0 . 99 , κ = 1 , under the same settings as described above.. The results (averaged over 3 random seeds) are reported in Table 8.

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 lr 0 . 2 ; the batch size is set to 1024; the weight decay is set to 1e-4; for SAM, ρ 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 γ and κ via a grid search over γ ∈ { 0 . 8 , 0 . 6 } and κ ∈ { 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 lr decay to 0 starting from an initial lr 1e-4; the weight decay is 5e-4; batch size is 64 ; strong data augmentations (basic + AutoAugment) are utilized; ρ = 0 . 1 for SAM. The total epochs are set to 200 , and we switch from AdamW/SAM to AdmIRE/SAM-IRE when the training loss approaches 0 . 5 . For AdmIRE/SAM-IRE, we fix K = 10 , and tune hyperparameters γ and κ via a grid search over γ ∈ { 0 . 99 , 0 . 9 , 0 . 8 } and κ ∈ { 20 , 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 β 1 = 0 . 9 , β 2 = 0 . 999 and weight decay λ = 1 . 0 ; the learning rate strategy integrates a warm-up phase followed by a cosine decay scheduler with lr_max =3e-3; batch size is 4096 ; the total training duration is 300 epochs, including 30 warm-up epochs; For AdmIRE, we switch from AdamW to AdmIRE at epoch 100 and examine different γ, κ . The results are reported in Table 4.

In this section, the experiments on CIFAR-10/CIFAR-100 were conducted using a single A800 GPU, and the experiments on ImageNet were conducted using 4 A800 GPUs.

C.2 Experimental details in Section 4.2

C.2.1 Experimental details in Section 4.2.1

We train a 2-layer decoder-only Transformer (8M parameters) using absolute positional encodings (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 γ, κ ). The (max) sequence length is set to 512, and the batch size is set to 32. The experiments in this section are conducted on 1 A800.

For the optimizer AdamW, we use the hyperparameters β 1 = 0 . 9 , β 2 = 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 lr_max and lr_min = lr_max / 20 .

First, we tune lr_max in AdamW from { 1.5e-4 , 3e-4 , 6e-4 , 1.2e-3 , 1.8e-3 , 3e-3 } . The results, shown in Figure 5 (left), identify the optimal lr_max = 6e-4 . We also use the optimal lr_max 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.

0k

25k

50k

75k

100k

4

6

8

10

training loss

Transformer (8M) on wikitext-2

lr_max

=

1.5e-4

lr_max

=

3.0e-4

lr_max

=

6.0e-4

lr_max

=

1.2e-3

lr_max

=

1.8e-3

lr_max

=

3.0e-3

0

25k

50k

4

6

8

10

validation loss

Llama (60M) on wikitext-103

lr_max

=3.0e-4

lr_max

=6.0e-4

lr_max

=1.2e-3

lr_max

=1.8e-3

Num of steps num of steps

C.2.2 Experimental details in Section 4.2.2

Llama (Touvron et al., 2023) is 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 β 1 = 0 . 9 , β 2 = 0 . 95 and weight decay λ = 0 . 1 for LLama (Touvron et al., 2023). The learning rate strategy integrates a warm-up phase followed by a cosine decay scheduler with lr_max and lr_min = lr_max / 20 . Additionally, it is used with gradient clipping 1.0 to maintain the training stability.

In each experiment, we tune the optimal lr_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) on wikitext-103 (0.5G). We train a 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 is 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 lr_max in AdamW from { 3e-4 , 6e-4 , 1.2e-3 , 1.8e-3 } , identifying the optimal lr_max 6e-4 . The experimental results are very similar to Figure 5 (right), so we will not show them again. Then, both AdamW and AdmIRE are trained using the optimal lr_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, on 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 lr_max in AdamW from { 3e-4 , 6e-4 , 1.2e-3 , 1.8e-3 } , identifying the optimal lr_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 lr_max .

Llama (229M) on openwebtext (38G). We train a 16-layer Llama model, with 12 heads in each layer and a 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 or 100,000 steps, including 1,000 warm-up steps. First, we tune lr_max in AdamW from { 3e-4 , 6e-4 , 1.2e-3 , 1.8e-3 } , identifying the optimal lr_max 6e-4 . (The results are very similar to Figure 5 (right), so we do not show them repearly.) Then, both AdamW and AdmIRE are trained using the optimal lr_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 ⊂ R p and a constant R > 0 , we denote B ( K ; R ) := { θ ∈ R p : dist( θ ; K ) ⩽ R } . ⟨· , ·⟩ represents the standard Euclidean inner product between two vectors. ∥·∥ denotes the ℓ 2 norm of a vector or the spectral norm of a matrix, whereas ∥ · ∥ F denotes the Frobenius norm of a matrix.

D.1 Preliminary Lemmas

Lemma D.1 (Arora et al. (2022), Lemma B.2) . Under Assumption 5.1, for any compact set K ⊂ Γ , there exist absolute constants R 1 , µ > 0 such that

(i) B ( K ; R 1 ) ⊂ U ;

(ii) L ( · ) is µ -PL (defined in Def F .1) on B ( K ; R 1 ) ;

(iii) inf θ ∈ B ( K ; R 1 ) λ m ( ∇ 2 L ( θ ) ) ⩾ µ .

We further define the following absolute constants on B ( K ; R 1 ) :

Lemma D.2 (Key properties of Φ( · ) (Arora et al., 2022)) . Under Assumption 5.1,

For any θ ∈ U , ∂ Φ( θ ) ∇L ( θ ) = 0 .

For any θ ∈ Γ , ∂ Φ( θ ) = P m +1: p ( ∇ 2 L ( θ )) .

LemmaD.3 (Continuity of P m +1: p ) . Under Assumption 5.1, there exists absolute constants R 2 , ζ P > 0 such that for any θ ∈ B ( K ; R 2 ) ,

Proof of Lemma D.3.

By Lemma D.1, for any θ ∈ B ( K ; R 1 ) , it holds that ∥ ∥ ∇ 2 L ( θ ) -∇ 2 L (Φ( θ )) ∥ ∥ ⩽ β 3 ∥ θ -Φ( θ ) ∥ . We choose R 2 := R 1 ∧ µ 4 β 3 . Then for any θ ∈ B ( K ; R 2 ) ,

Let the orthogonal decomposition of ∇ 2 L ( θ ) and ∇ 2 L (Φ( θ )) be ∇ 2 L ( θ ) = ∑ p k =1 λ k u k u ⊤ k ( λ 1 ⩾ · · · ⩾ λ p ) and ∇ 2 L (Φ( θ )) = ∑ p k =1 µ k v k v ⊤ k ( µ 1 ⩾ · · · ⩾ µ p ), respectively.

Consequently, by Lemma F.2, we can bound the gap of eigenvalues: for any k ∈ [ p ] ,

Noticing Φ( θ ) ∈ Γ , by Lemma D.1, it holds that µ 1 ⩾ µ m ⩾ µ and µ m +1 = · · · = µ p = 0 . Thus, we can obtain the bounds of { λ k } k :

For simplicity, we denote U top := ( u 1 , · · · , u m ) , U bottom := ( u m +1 , · · · , u p ) , V top := ( v 1 , · · · , v m ) , V bottom := ( v m +1 , · · · , v p ) .

By Lemma F.3, we can bound the gap between the subspaces:

According to the definition of P m +1: p ( · ) , it holds that

Noticing the relationship

we obtain the bound:

To summarize, we only need to choose the constants R 2 = R 1 ∧ µ 4 β 3 and ζ P = 2 √ 2( m ∧ ( p -m )) β 3 µ to ensure this lemma holds.

Proof Notations. Now we introduce some additional useful notations in the proof in this section.

First, we choose R := ( R 1 ∧ R 2 ) / 2 , where R 1 is defined in Lemma D.1 and R 2 is defined in Lemma D.3. Let µ be the PL constant on B ( K ; R ) . Moreover, we use the following notations:

Ensured by Lemma D.1 and D.3, these quantities are all absolute constants in (0 , + ∞ ) . Moreover, without loss of generality, we can assume that β 1 , β 2 , β 3 , ζ Φ , ζ P > 1 and µ ⩽ ν < 1 .

Lemma D.4 (Connections between para norm, grad norm, and loss) . For any θ ∈ B ( K ; R ) > 0 , it holds that:

(para norm v.s. grad norm) µ ∥∇L ( θ ) ∥ ⩽ ∥ θ -Φ( θ ) ∥ ⩽ β 2 ∥∇L ( θ ) ∥ ;

(grad norm v.s. loss) 2 µ L ( θ ) ⩽ ∥∇L ( θ ) ∥ 2 ⩽ 2 β 2 2 µ L ( θ ) ;

(loss v.s. para norm) µ 2 ∥ θ -Φ( θ ) ∥ 2 ⩽ L ( θ ) ⩽ β 2 2 2 µ ∥ θ -Φ( θ ) ∥ 2 .

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 β 2 , the upper bounds hold naturally.

Lemma D.5. For all θ ∈ B ( K ; R ) ,

∥ ∥ P m +1: p ( ∇ 2 L ( θ )) ∇ 2 L ( θ ) ∥ ∥ ⩽ O ( ∥ θ -Φ( θ ) ∥ ) ;

∥ ∥ P m +1: p ( ∇ 2 L ( θ )) ∇L ( θ ) ∥ ∥ ⩽ O ( ∥ θ -Φ( θ ) ∥ 2 ) ;

Let ρ > 0 and v ∈ S p -1 . If θ + ρ v ∈ B ( K ; R ) , then

Proof of Lemma D.5.

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:

Let the ρ in SAM satisfy:

For simplicity, we denote

Notice L ( θ T I ) < C 1 ηρ 2 , we have the following upper bound for the probability

For each term t ∈ [ T I , T I + T II -1] , it can be bounded by:

For simplicity, we denote

P t +1 ,s := P ( L ( θ t +1 ) ⩾ 2 C 1 ηρ 2 ; ∀ τ ∈ [ s +1 , t ] , C 1 ηρ 2 ⩽ L ( θ τ ) < 2 C 1 ηρ 2 ∣ ∣ ∣ L ( θ s ) < C 1 ηρ 2 ) , s ∈ [ T I , t -1] .

· Step I. Bounding P t +1 ,t .

From L ( θ t ) ⩽ C 1 ηρ 2 , we have ∥ θ t -Φ( θ t ) ∥ ⩽ √ 2 µ L ( θ t ) = O ( √ ηρ µ ) , thus which means θ t + ρ v t ∈ B ( K ; R ) . Furthermore,

which implies θ t +1 ∈ B ( K ; R ) . Consequently, we have the following quadratic upper bound:

Thus, we obtain

· Step II. Bounding P t +1 ,s for s ∈ [ T I , t -1] .

It is clear that

Then our key step is to prove the following two claims about the process { X τ } .

-Claim I. X τ -C 2 τηρ 2 is a super-martingale. From the definition of X τ , we only need to prove that if C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 , then E [ L ( θ τ +1 )] ⩽ L ( θ τ ) -C 2 η 2 ρ 2 . If C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 , similar to Step I, it is clear that θ τ +1 ∈ B ( K ; R ) . Applying the quadratic upper bound, it holds that:

We prove this step under the condition L ( θ s ) < C 1 ηρ 2 . Define a process { X τ } t +1 τ = s : X s +1 = L ( θ s +1 ) ,

.

Taking the expectation, we have:

From L ( θ τ ) ⩾ C 1 ηρ 2 and ρ = O ( √ η ) , it holds that

Therefore, we have:

-Claim II. X τ +1 -X τ + C 2 η 2 ρ 2 is O ( η 2 ρ 2 + η 3 / 2 ρ 2 /p 1 / 2 ) -sub-Gaussian. From the definition of X τ , we only need to prove for the case C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 . If C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 , then X τ = L ( θ τ ) and X τ +1 = L ( θ τ +1 ) . Similar to Step I, it is clear that θ τ +1 ∈ B ( K ; R ) .

Applying the smoothness, we have:

which implies:

With the preparation of Claim I and Claim II, we can use the Azuma-Hoeffding inequality (Lemma F.4 (ii)): for any Q > 0 , it holds that

As proved in Claim I, X s +1 = L ( θ s +1 ) ⩽ 3 2 C 1 ηρ 2 due to L ( θ s ) ⩽ C 1 ηρ 2 . Therefore, by choosing Q = ( t -s ) C 2 η 2 ρ 2 -3 2 C 1 ηρ 2 +2 C 1 ηρ 2 = ( t -s ) C 2 η 2 ρ 2 + 1 2 C 1 ηρ 2 , we have

Therefore, we obtain the union bound:

Hence, with probability at least 1 -T 2 II exp ( -Ω ( 1 η + p -1 )) , for any t ∈ [ T I , T I + T II ] ,

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 is the same as 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:

Let the ρ in SAM and the κ in IRE satisfy:

For simplicity, we denote

Following the proof for SAM, we denote

and it holds that

Step I. Bounding P t +1 ,t .

From L ( θ t ) ⩽ C 1 ηρ 2 , we have ∥ θ t -Φ( θ t ) ∥ ⩽ √ 2 µ L ( θ t ) = O ( √ ηρ ) , thus

which means θ t + ρ v t ∈ B ( K ; R ) . Furthermore,

which implies θ t +1 ∈ B ( K ; R ) . Consequently, we have the following quadratic upper bound:

Lemma D.5

Thus, we obtain

Step II. Bounding P t +1 ,s for s [ T I , t 1] .

We prove this step under the condition L ( θ s ) < C 1 ηρ 2 . Define a process { X τ } t +1 τ = s : X s +1 = L ( θ s +1 ) ,

∈ -

.

It is clear that

Then our key step is to prove the following two claims about the process { X τ } .

-Claim I. X τ -C 2 τηρ 2 is a super-martingale. From the definition of X τ , we only need to prove that if C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 , then E [ L ( θ τ +1 )] ⩽ L ( θ τ ) -C 2 η 2 ρ 2 .

If C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 , similar to Step I, it is clear that θ τ +1 ∈ B ( K ; R ) . Applying the quadratic upper bound, it holds that:

Taking the expectation, we have:

From C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 and ρ = O ( √ η ) , it holds that

Moreover, recall κ ⩽ 1 /ρ . Therefore, we have the upper bound:

-Claim II. X τ +1 -X τ + C 2 η 2 ρ 2 is O ( η 2 ρ 2 + η 3 / 2 ρ 2 /p 1 / 2 ) -sub-Gaussian. From the definition of X τ , we only need to prove for the case C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 . If C 1 ηρ 2 ⩽ L ( θ τ ) ⩽ 2 C 1 ηρ 2 , then X τ = L ( θ τ ) and X τ +1 = L ( θ τ +1 ) . Similar to Step I, it is clear that θ τ +1 ∈ B ( K ; R ) .

Applying the smoothness, we have:

In the same way as the proof of Claim I, it holds that

Thus,

which implies:

With the preparation of Claim I and Claim II, we can use the Azuma-Hoeffding inequality (Lemma F.4 (ii)): for any Q > 0 , it holds that

As proved in Claim I, X s +1 = L ( θ s +1 ) ⩽ 3 2 C 1 ηρ 2 due to L ( θ s ) ⩽ C 1 ηρ 2 . Therefore, by choosing Q = ( t -s ) C 2 η 2 ρ 2 -3 2 C 1 ηρ 2 +2 C 1 ηρ 2 = ( t -s ) C 2 η 2 ρ 2 + 1 2 C 1 ηρ 2 , we have

Therefore, we obtain the union bound:

Hence, with probability at least 1 -T 2 II exp ( -Ω ( 1 η + p -1 )) , for any t ∈ [ T I , T I + T II ] ,

D.2.3 Proof of the Effective Dynamics

We have proved that with high probability at least 1 -T 2 II exp ( -Ω ( 1 η + p -1 )) , for any t ∈ [ T I , T I + T II ] , ∥ θ t -Φ( θ t ) ∥ = O ( √ ηρ ) . Then we prove this theorem when the above event occurs. For any t ∈ [ T I , T I + T II ] ,

Then by Taylor's expansion,

For the term ∂ Φ( θ t ) ∇L ( θ t + ρ v t ) and ∂ Φ( θ t ) P ( θ t ) ∇L ( θ t + ρ v t ) , using Taylor's expansion, we have

Taking the expectation (about v t ), we have

Additionally, using Lemma D.2 and Taylor's expansion, we have:

Combining the results above, we obtain:

E Proofs in Section 5.2.2

Setting E.1. Consider the empirical risk minimization min : L ( θ ) = 1 n ∑ n i =1 L i ( θ ) , where L i ( θ ) = ℓ ( f i ( θ ) , y i ) is the loss on the i -th data ( x i , y i ) . Let f i ( · ) and ℓ ( · , · ) be C 4 . Suppose all global minimizers interpolate the training dataset, i.e., L ( θ ⋆ ) = min θ L ( θ ) implies f i ( θ ⋆ ) = y i for all i ∈ [ n ] . We denote the minima manifold by M = { θ : f i ( θ ) = y i , ∀ i ∈ [ n ] } . Moreover, we assume that ∂ 2 ℓ (ˆ y,y ) ∂ ˆ y 2 | ˆ y = y > 0 and the feature matrix ( ∇ f i ( θ ) , · · · , ∇ f n ( θ )) ∈ R p × n is full-rank at θ ∈ M . Lemma E.2 (Theorem 5.2 in Wen et al. (2023a)) . Under Setting E.1, 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 E.1, for any compact set K ∈ Γ , there exist absolute constants R 1 , µ > 0 such that

We further define the following absolute constants on B ( K ; R ) :

Lemma E.4 (Wen et al. (2023a)) . Under Assumption 5.1,

For any θ ∈ U , ∂ Φ( θ ) ∇L ( θ ) = 0 .

For any θ ∈ Γ , ∂ Φ( θ ) = P n +1: p ( ∇ 2 L ( θ )) and ∂ Φ( θ ) ∇ 2 L ( θ ) = 0 .

For any θ ∈ Γ , ∂ Φ( θ ) ∇ 2 L i ( θ ) = 0 , ∀ i ∈ [ n ] .

Lemma E.5. Under Setting E.1, there exists absolute constants R 2 , ζ P > 0 such that for any θ ∈ B ( K ; R 2 ) ,

Proof Notations. Now we introduce some additional useful notations in the proof in this section.

First, we choose R := ( R 1 ∧ R 2 ) / 2 , where R 1 is defined in Lemma E.3 and R 2 is defined in Lemma E.5. Let µ be the PL constant on B ( K ; R ) . Moreover, we use the following notations:

Ensured by Lemma D.1 and D.3, these quantities are all absolute constants in (0 , + ∞ ) . Moreover, without loss of generality, we can assume that β 1 , β 2 , β 3 , ζ Φ , ζ P > 1 and µ ⩽ ν < 1 .

Then we have the following two lemmas, similar to Lemma D.4 and D.5.

Lemma E.6. For any θ ∈ B ( K ; R ) , it holds that:

(para norm v.s. grad norm) µ ∥∇L ( θ ) ∥ ⩽ ∥ θ -Φ( θ ) ∥ ⩽ β 2 ∥∇L ( θ ) ∥ ; µ ∥∇L i ( θ ) ∥ ⩽ ∥ θ -Φ( θ ) ∥ ⩽ β 2 ∥∇L i ( θ ) ∥ , ∀ i ∈ [ n ] .

(grad norm v.s. loss) 2 µ L ( θ ) ⩽ ∥∇L ( θ ) ∥ 2 ⩽ 2 β 2 2 µ L ( θ ) ; 2 µ L i ( θ ) ⩽ ∥∇L i ( θ ) ∥ 2 ⩽ 2 β 2 2 µ L i ( θ ) , ∀ i ∈ [ n ] .

(loss v.s. para norm) µ 2 ∥ θ -Φ( θ ) ∥ 2 ⩽ L ( θ ) ⩽ β 2 2 2 µ ∥ θ -Φ( θ ) ∥ 2 ; µ 2 ∥ θ -Φ( θ ) ∥ 2 ⩽ L i ( θ ) ⩽ β 2 2 2 µ ∥ θ -Φ( θ ) ∥ 2 , ∀ i ∈ [ n ] .

Lemma E.7. For all θ ∈ B ( K ; R ) , ∀ i ∈ [ n ] ,

∥ ∥ P n +1: p ( ∇ 2 L ( θ )) ∇ 2 L ( θ ) ∥ ∥ , ∥ ∥ P n +1: p ( ∇ 2 L ( θ )) ∇ 2 L i ( θ ) ∥ ∥ , ∥ ∥ ∂ Φ( θ ) ∇ 2 L i ( θ ) ∥ ∥ ⩽ O ( ∥ θ -Φ( θ ) ∥ ) ;

•

Let ρ > 0 and v ∈ S p -1 . If θ + ρ v ∈ B ( K ; R ) , then ∥∇L i ( θ + ρ v ) ∥ ⩽ ∥∇L i ( θ ) ∥ + ρβ 2 , ∀ i ∈ [ n ]; ;

Lemma E.8 (Lemma H.9 in Wen et al. (2023a)) . For any absolute constant C > 0 , there exist absolute constant C 1 , C 2 > 0 such that: if θ t ∈ B ( K ; R ) and C 1 ηρ ⩽ ∥ θ t -Φ( θ t ) ∥ ⩽ Cρ , then it holds that:

Lemma E.9 (Lemma H.10 in Wen et al. (2023a)) . For any absolute constant C > 0 , there exists absoulte constant C 3 such that: if θ t ∈ B ( K ; R ) and ∥ θ t -Φ( θ t ) ∥ ⩽ Cρ , then we have that

Now we fix the positive number C = 1 and use the absolute constants C 2 , C 3 , defined in 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 to the proof for Section 5.2.1.

Under the conditions in Theorem 5.6, the update rule of IRE on standard SAM is

Let the κ in IRE satisfy

Additionally, we fix a constant α ∈ (0 , 1) in the proof.

For simplicity, we denote

and

Additionally, we denote

Then the following bound holds naturally:

Step I. Bounding P t +1 ,t .

From ∥ θ t -Φ( θ t ) ∥ = O ( η 1 -α ρ ) , thus

which means θ t + ρ v t ∈ B ( K ; R ) . Using Lemma E.7 and κ ⩽ 1 /ρ , we can estimate:

Then we have the following bound:

Thus, we obtain

Step II. Bounding P t +1 ,s for s [ T I , t 1] .

We prove this step under the condition ∥ θ s -Φ( θ s ) ∥ < Cη 1 -α ρ . Define a process { X τ } t +1 τ = s X s +1 = ∥ θ s +1 -Φ( θ s +1 ) ∥ ,

∈ -:

It is clear that

.

Then our key step is to prove the following two claims about the process { X τ } .

∥ -∥ -If Cη 1 -α ρ ⩽ X τ = ∥ θ τ -Φ( θ τ ) ∥ ⩽ 2 Cη 1 -α ρ , similar to Step I, it holds that θ τ +1 ∈ B ( K ; R ) and ∥ ∥ θ t +1 -Φ( θ τ +1 ) -( θ τ +1 / 2 -Φ( θ τ +1 / 2 )) ∥ ∥ ⩽ C 2 2 ηρ . Moreover,

-Claim I. X τ -C 2 τηρ/ 2 is a super-martingale. From the definition of X τ , we only need to prove that if Cη 1 -α ρ ⩽ X τ = ∥ θ τ -Φ( θ τ ) ∥ ⩽ 2 Cη 1 -α ρ , then E ∥ θ τ +1 -Φ( θ τ +1 ) ∥ ⩽ θ τ Φ( θ τ ) C 2 ηρ/ 2 .

Taking the expectation and using Lemma E.8, we have

-Claim II. X τ +1 -X τ + C 2 ηρ/ 2 is O ( ηρ ) -bounded. From the definition of X τ , we only need to prove for the case Cη 1 -α ρ ⩽ X τ = ∥ θ τ -Φ( θ τ ) ∥ ⩽ 2 Cη 1 -α ρ . If Cη 1 -α ρ ⩽ X τ = ∥ θ τ -Φ( θ τ ) ∥ ⩽ 2 Cη 1 -α ρ , we have θ τ +1 ∈ B ( K ; R ) and ∥ ∥ θ t +1 -Φ( θ t +1 ) -( θ t +1 / 2 -Φ( θ t +1 / 2 )) ∥ ∥ ⩽ C 2 2 ηρ . Combining this result and Lemma E.9, we have

With the preparation of Claim I and Claim II, we can use the Azuma-Hoeffeding inequality: for any Q > 0 , it holds that

As proved in Claim I, X s +1 = ∥ θ s +1 -Φ( θ s +1 ) ∥ ⩽ 3 2 Cη 1 -α ρ due to ∥ θ s -Φ( θ s ) ∥ ⩽ Cη 1 -α ρ . Therefore, by choosing Q = ( t -s ) C 2 ηρ/ 2 -3 2 Cη 1 -α ρ + 2 Cη 1 -α ρ = ( t -s ) C 2 ηρ/ 2 + Cη 1 -α ρ/ 2 , we have

Therefore, we obtain the union bound:

E.2.2 Proof of the Effective Dynamics

By our proof above, with probability at least 1 -T 2 II exp( -Ω(1 /η α )) , for any t ∈ [ T I , T I + T II ] , ∥ θ t -Φ( θ t ) ∥ = O ( η 1 -α ρ ) . Then we prove this theorem when the above event occurs.

Due to Φ( ) = ( η 1 -α ρ )

Then by Taylor's expansion,

For the term ∂ Φ( θ t ) ∇L i t ( θ t + ρ v t ) and ∂ Φ( θ t ) P ( θ t ) ∇L i t ( θ t + ρ v t ) , using Taylor's expansion and Lemma E.7, we have

and

Combining the results above, we obtain:

Additionally, taking the expectation, we have:

Therefore, we obtain

F Useful Inequalities

Definition F.1 ( µ -PL) . Let µ > 0 be a constant. A function L is µ -PL in a set U iff ∥∇L ( θ ) ∥ 2 ⩾ 2 µ ( L ( θ ) -inf θ ∈ U L ( θ )) , ∀ θ ∈ U .

Lemma F.2 (Weyl Theorem) . Let A , B ∈ R p × p be symmetric with eigenvalues λ 1 ⩾ · · · ⩾ λ p and µ 1 ⩾ · · · ⩾ µ p respectively, then for any k ∈ [ p ] , it holds that

Lemma F.3 (Davis-Kahan sin(Θ) theorem) . Let A , B ∈ R p × p be symmetric matrices. Denote their orthogonal decomposition as A = E 1 Λ 1 E ⊤ 1 + E 2 Λ 2 E ⊤ 2 and B = F 1 Γ 1 F ⊤ 1 + F 2 Γ 2 F ⊤ 2 with ( E 1 , E 2 ) and ( D 1 , D 2 ) orthogonal. If the eigenvalues in Λ 1 are contained in an interval ( a, b ) , and the eigenvalues of Γ 2 are excluded from the interval ( a -δ, b + δ ) for some δ > 0 , then for any unitarily invariant norm ∥·∥ ⋆ ,

Lemma F.4 (Azuma-Hoeffding Inequality) . Suppose { X n } n ∈ N is a super-martingale.

(i) (Bounded martingale difference). If -α ⩽ X i +1 -X i ⩽ β , then for any n, t > 0 , we have:

(ii) (Sub-Gaussian martingale difference). If X i +1 -X i is σ 2 i -sub-Gaussian, then for any n, t > 0 , we have:

Lemma F.5. Let v ∈ R p . Let g ∼ N ( 0 , I ) . Then there exists an absolute constant c > 0 such that for any t > 0 ,

Proof of Lemma F.5. From P54 in Vershynin (2018), there exists an absolute constant c > 0 such that for any t > 0 , P ( |⟨ e 1 , g ⟩| ∥ g ∥ ⩾ t ) ⩽ 4 e -cpt 2 . Without loss of generality, we can assume v = 0 . Then we have:

̸

Lemma F.6. ∥ AB ∥ F ⩽ ∥ A ∥ ∥ B ∥ F .

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 of the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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 or 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: In Section 6.

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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We believe that all of the experimental results are reproducible in our 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 a 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 it 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 to 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 a 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 submissions 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 it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or 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. In the case of closed-source models, it 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 to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: In https://github.com/wmz9/IRE-algorithm-framework .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a 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 as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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]

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, confidence intervals, or statistical significance tests, at least for the experiments that support the 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, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 or failed experiments that didn't make it 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://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have confirmed that the research is conducted with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 or 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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 a best faith effort.

12. Licenses 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 not use 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 a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

If this information is not available online, the authors are encouraged to reach out to the asset'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/model 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 as details about compensation (if any)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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 not involve crowdsourcing nor 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nI think the idea is novel, and the experiments are promising, so this paper could stimulate significant research in the future.\n\nWEAKNESSES:\nOne 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\nQUESTIONS:\nShouldn'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$}?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\n\nSTRENGTHS:\n## 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.\n\nWEAKNESSES:\n### 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.\n\nQUESTIONS:\n1. **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.\n\nLIMITATIONS:\n1. **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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. Could IRE help in the finetuning phase of LLMs, Including both convergence speed and convergence position properties?\n\nLIMITATIONS:\n1. The author did not provide information on the computing resources they used.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}