Improving Generalization and Convergence by Enhancing Implicit Regularization

Authors: mingze  wang
jinbo  wang
haotian  he
zilin  wang
guanhua  huang
feiyu  xiong
zhiyu  li
lei  wu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 14:02:57

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

Main paper: 
 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.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ębski et al., 2017) , normalization (Lyu et al., 2022) , cyclic LR (Wang and Wu, 2023) , influence this sharpness regularization. Furthermore, some works have provided theoretical evidence explaining the superior generalization of flat minima for neural networks (Ma and Ying, 2021; Mulayoff et al., 2021; Wu and Su, 2023; Gatmiry et al., 2023; Wen et al., 2023b) . Our work is inspired by this line of research, aiming to boost implicit sharpness regularization by decoupling the dynamics along flat and sharp directions.Sharpness-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.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. 
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:EQUATIONwherev ∈ 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). ThenL(θ) = 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 assumeH(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 followsEQUATIONwhere ⊙ 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 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 byEQUATIONwhere κ > 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. 
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 byEQUATIONwhere κ 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 followsEQUATIONwhere γ ∈ (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.Algorithm 1: Practical IRE (A practical framework of implementing IRE)Input: θ 0 , T , K, learning rate {η t } t , warm-up time T w , IRE hyperparams: κ ⩾ 0 and γ ∈ (0.5, 1) ;for t = T w , • • • , T -1 doCompute the original update direction g t according to the base optimizer; if (t -T w ) mod K = 0 then Estimate the diagonal Hessian h t ∈ R p using Eq. ( 6); Update the mask n t ∈ R p using Eq. ( 5); else n t = n t-1θ t+1 = θ t -η t g t + κn t ⊙ g t ;Output: θ T .A light-weight estimator of the diagonal Hessian. Let ℓ(•, •) be the cross-entropy loss. Given an input data x ∈ R dx and label y ∈ R dy , let the model's prediction be f (x; θ) ∈ R dy . 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 ; θ); ŷb )⊙ ∇ℓ(f (x b ; θ); ŷb ), where ŷ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( Feff (θ)), only requires computing the mini-batch gradient ∇ LBEQUATIONAccording to Liu et al. (2024) , this estimator is an unbiased estimate of the empirical diagonal Fisher, i.e., E ŷ [diag( Feff (θ))] = E ŷ [diag( F (θ))]. For more discussions on the efficiency of this estimator, please refer to (Liu et al., 2024, Section 2) . Additionally, for squared loss, one can simply use Fisher 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 . 
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.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. 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 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.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 Ad-mIRE 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 Ad-mIRE is only 1.12 times that of AdamW, corresponding to the theoretical estimation (1.1 times). Robustness to hyperparameters. Figure 3 shows that Ad-mIRE, 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.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 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 100knum of steps 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. 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. 
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.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) dt = -∇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 .The stochastic SAM (Foret et al., 2021) is given by standard SAM:EQUATIONThe 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 :average SAM:EQUATIONTwo-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 dt = -∇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.minima manifold M tangent space T zt M -∇Tr[∇ 2 L(z t )/2] -∇ M Tr[∇ 2 L(z t )/2] θ t z t = Φ(θ t )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=TI+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 bothEQUATIONwhich 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/ρ.An Ω(1/η 0.5 ) accelerationWe 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. in SAM (8), κ ⩽ 1/ρ , and P t = P m+1:p (∇ 2 L(θ t )) in IRE (4), then with high probability at least1 -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:E ξt [z t+1 ] = z t - (1 + κ)ηρ 2 p ∇ M Tr ∇ 2 L(z t )/2 + O(η 3/2 ρ 2 ).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.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:E it [z t+1 ] = z t -(1 + κ)ηρ 2 ∇ M Tr ∇ 2 L(z t )/2 + O (κ + 1)ηρ 2 (ρ + η 1-α ) .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. 
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ębski 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. B Proofs for SGD in Section 2 
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] Justification: In Appendix C.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). 
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). 
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, 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). 
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 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. 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. • 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. 
 
 Appendix: 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 )dt + ηΣ(θ t )dW 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 asEQUATION)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) isEQUATIONThis 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 becauseE 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 toEQUATIONComparing ( 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:EQUATION)(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.This section describes the experimental details in Section 4. 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 The flatness measure, Tr(∇ 2 L(θ)), is approximated by the trace of Fisher Tr(F (θ)). Specifically, we use diag( Feff (θ)) in ( 6) for the estimate becauseE ŷ [diag( Feff (θ))] = E ŷ [diag( F (θ))] and thus,Tr(∇ 2 L(θ)) ≈ E ŷ [Tr( F (θ))] = E ŷ [Tr(diag( F (θ)))] = E ŷ [diag( Feff (θ))].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.2Experiments 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. 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 . 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.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.1Additional 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.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;R1) λ m ∇ 2 L(θ) ⩾ µ.Proof of Lemma D.5.P m+1:p (∇ 2 L(θ))∇ 2 L(θ) ⩽ P m+1:p (∇ 2 L(Φ(θ)))∇ 2 L(Φ(θ)) + ζ P β 2 ∥θ -Φ(θ)∥ + β 3 ∥θ -Φ(θ)∥ =0 + O (∥θ -Φ(θ)∥) ; P m+1:p (∇ 2 L(θ))∇L(θ) ⩽ P m+1:p (∇ 2 L(θ))∇ 2 L(Φ(θ))(Φ(θ) -θ) + O ∥θ -Φ(θ)∥ 2 ⩽ P m+1:p (∇ 2 L(Φ(θ)))∇ 2 L(Φ(θ))(Φ(θ) -θ) + O ∥θ -Φ(θ)∥ 2 = O ∥θ -Φ(θ)∥ 2 ;∥∇L(θ + ρv)∥ ⩽ ∥∇L(θ)∥ + ρβ 2 ;P m+1:p (∇ 2 L(θ))∇L(θ + ρv) ⩽ P m+1:p (∇ 2 L(θ))∇L(θ) + ρ P m+1:p (∇ 2 L(θ))∇ 2 L(θ)v + ρ 2 β 3 2 ⩽O ∥θ -Φ(θ)∥ 2 + O (ρ ∥θ -Φ(θ)∥) + ρ 2 β 3 2 .D.2 Proof of Theorem 5.5We 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:θ t+1 = θ t -η∇L θ t + ρ ξ t ∥ξ t ∥, where ξ t ∼ N (0, I).Let the ρ in SAM satisfy:ρ = O( √ η).For simplicity, we denotev t := ξ t ∥ξ t ∥ , C 1 := 4β 3 2 µ , C 2 := β 3 2 .Notice L(θ TI ) < C 1 ηρ 2 , we have the following upper bound for the probabilityP ∃ t ∈ [T I , T I + T II ], L(θ t ) ⩾ 2C 1 ηρ 2 ⩽ TI+TII-1 t=TI P L(θ t+1 ) ⩾ 2C 1 ηρ 2 ; ∀ s ∈ [T I , t], L(θ s ) < 2C 1 ηρ 2 .For each term t ∈ [T I , T I + T II -1], it can be bounded by:P L(θ t+1 ) ⩾ 2C 1 ηρ 2 ; ∀ s ∈ [T I , t], L(θ s ) < 2C 1 ηρ 2 ⩽P L(θ t+1 ) ⩾ 2C 1 ηρ 2 ; L(θ t ) < C 1 ηρ 2 + t-1 s=TI P L(θ t+1 ) ⩾ 2C 1 ηρ 2 ; L(θ s ) < C 1 ηρ 2 ; ∀ τ ∈ [s + 1, t], C 1 ηρ 2 ⩽ L(θ τ ) < 2C 1 ηρ 2 ⩽P L(θ t+1 ) ⩾ 2C 1 ηρ 2 L(θ t ) < C 1 ηρ 2 + t-1 s=TI P L(θ t+1 ) ⩾ 2C 1 ηρ 2 ; ∀ τ ∈ [s + 1, t], C 1 ηρ 2 ⩽ L(θ τ ) < 2C 1 ηρ 2 L(θ s ) < C 1 ηρ 2 .For simplicity, we denoteP t+1,t := P L(θ t+1 ) ⩾ 2C 1 ηρ 2 L(θ t ) < C 1 ηρ 2 , P t+1,s := P L(θ t+1 ) ⩾ 2C 1 ηρ 2 ; ∀ τ ∈ [s + 1, t], C 1 ηρ 2 ⩽ L(θ τ ) < 2C 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 ∥θ t + ρv t -Φ(θ t )∥ ⩽ ∥θ t -Φ(θ t )∥ + ρ = O( √ ηρ) + O(ρ) < R, which means θ t + ρv t ∈ B(K; R). Furthermore, ∥θ t+1 -Φ(θ t )∥ ⩽ ∥θ t -Φ(θ t )∥ + η ∥∇L (θ t + ρv t )∥ ⩽O( √ ηρ) + η ∥∇L (θ t + ρv t )∥ ⩽ O( √ ηρ) + η ∥∇L (θ t )∥ + β 2 ηρ ⩽O( √ ηρ) + η 2β 2 2 µ L(θ t ) + β 2 ηρ ⩽ O( √ ηρ) + O(η 3/2 ρ) + O(ηρ) ⩽ R,which implies θ t+1 ∈ B(K; R). Consequently, we have the following quadratic upper bound:L(θ t+1 ) = L (θ t -η∇L (θ t + ρv t )) ⩽L(θ t ) -η ⟨∇L(θ t ), ∇L (θ t + ρv t )⟩ + η 2 β 2 2 ∥∇L (θ t + ρv t )∥ 2 ⩽L(θ t ) -η ∥∇L(θ t )∥ 2 -ηρ ∇L(θ t ), ∇ 2 L(θ t )v t + ηρ 2 β 3 2 ∥∇L(θ)∥ + η 2 β 2 2 (∥∇L(θ t )∥ + ρβ 2 ) 2 ⩽L(θ t ) -η ∥∇L(θ t )∥ 2 -ηρ ∇L(θ t ), ∇ 2 L(θ t )v t + ηρ 2 β 3 2 ∥∇L(θ)∥ + η 2 β 2 ∥∇L(θ t )∥ 2 + ρ 2 β 2 2 ⩽L(θ t ) + ηρ ∇ 2 L(θ t )∇L(θ t ) + ηρ 2 β 3 2 ∥∇L(θ t )∥ + η 2 β 2 ρ 2 β 2 2 ⩽L(θ t ) + ηρβ 2 + ηρ 2 β 3 2 ∥∇L(θ t )∥ + O(η 2 ρ 2 ) ⩽L(θ t ) + ηρβ 2 + ηρ 2 β 3 2 2β 2 2 µ L(θ t ) + O(η 2 ρ 2 ) ⩽C 1 ηρ 2 + O(η 3/2 ρ 2 ) + O(η 2 ρ 2 ) < 3C 1 ηρ 2 /2.Thus, we obtainP t+1,t = P L(θ t+1 ) ⩾ 2C 1 ηρ 2 L(θ t ) < C 1 ηρ 2 = 0. • Step II. Bounding P t+1,s for s ∈ [T I , t -1].We prove this step under the conditionEQUATIONIt is clear thatP t+1,s ⩽ P X t+1 ⩾ 2C 1 ηρ 2 .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 ifC 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 1 ηρ 2 , then E [L(θ τ +1 )] ⩽ L(θ τ ) -C 2 η 2 ρ 2 . If C 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 1 ηρ 2 , similar toStep I, it is clear that θ τ +1 ∈ B(K; R). Applying the quadratic upper bound, it holds that:L(θ τ +1 ) = L (θ τ -η∇L (θ τ + ρv τ )) ⩽L(θ τ ) -η ⟨∇L(θ τ ), ∇L (θ τ + ρv τ )⟩ + η 2 β 2 2 ∥∇L (θ τ + ρv τ )∥ 2 ⩽L(θ τ ) -η ∥∇L(θ τ )∥ 2 -ηρ ∇L(θ τ ), ∇ 2 L(θ τ )v t + ηρ 2 β 3 2 ∥∇L(θ τ )∥ + η 2 β 2 ∥∇L(θ τ )∥ 2 + ρ 2 β 2 2 .Taking the expectation, we have:E [L(θ τ +1 )] ⩽L(θ τ ) -η ∥∇L(θ τ )∥ 2 + ηρ 2 β 3 2 ∥∇L(θ τ )∥ + η 2 β 2 ∥∇L(θ τ )∥ 2 + ρ 2 β 2 2 ⩽L(θ τ ) - 3 4 η ∥∇L(θ τ )∥ 2 + ηρ 2 β 3 2 ∥∇L(θ τ )∥ + η 2 ρ 2 β 3 2 ⩽L(θ τ ) -η ∥∇L(θ τ )∥ 3 4 ∥∇L(θ τ )∥ - ρ 2 β 3 2 + η 2 ρ 2 β 3 2 . From L(θ τ ) ⩾ C 1 ηρ 2 and ρ = O( √ η), it holds that ∥∇L(θ τ )∥ ⩾ 2µL(θ t ) ⩾ 2C 1 µ √ ηρ ⩾ 4β 3 ρ 2 .Therefore, we have:E [L(θ τ +1 )] ⩽L(θ τ ) - 5 8 η ∥∇L(θ τ )∥ 2 + η 2 ρ 2 β 3 2 ⩽L(θ τ ) - 10 8 ηµL(θ τ ) + η 2 ρ 2 β 3 2 ⩽ L(θ τ ) - 10 8 C 1 µη 2 ρ 2 + η 2 ρ 2 β 3 2 ⩽L(θ τ ) -β 3 2 η 2 ρ 2 = L(θ τ ) -C 2 η 2 ρ 2 . -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 caseC 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 1 ηρ 2 . If C 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 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:L(θ τ +1 ) = L (θ τ -η∇L (θ τ + ρv τ )) =L(θ τ ) -η ⟨∇L(θ τ ), ∇L (θ τ + ρv τ )⟩ + O η 2 ∥∇L (θ τ + ρv τ )∥ 2 =L(θ τ ) -η ∥∇L(θ τ )∥ 2 -ηρ ∇L(θ τ ), ∇ 2 L(θ τ )v t + O ηρ 2 ∥∇L(θ τ )∥ + O η 2 ∥∇L (θ τ )∥ 2 + η 2 ρ 2 =L(θ τ ) + O η 2 ρ 2 -ηρ ∇L(θ τ ), ∇ 2 L(θ τ )v t + O ηρ 2 √ ηρ + O η 3 ρ 2 + η 2 ρ 2 =L(θ τ ) -ηρ ∇L(θ τ ), ∇ 2 L(θ τ )v t + O(η 2 ρ 2 ),which implies:X τ +1 -X τ -C 2 η 2 ρ 2 ψ2 ⩽ ∥L(θ τ +1 ) -L(θ τ )∥ ψ2 + C 2 η 2 ρ 2 ψ2 ⩽ ηρ ∇L(θ τ ), ∇ 2 L(θ τ )v t ψ2 + O(η 2 ρ 2 ) ⩽ηρ ∇ 2 L(θ τ )∇L(θ τ ) ∇ 2 L(θ τ )∇L(θ τ ) ∥∇ 2 L(θ τ )∇L(θ τ )∥ , v t ψ2 + O(η 2 ρ 2 ) ⩽O η 3/2 ρ 2 ∇ 2 L(θ τ )∇L(θ τ ) ∥∇ 2 L(θ τ )∇L(θ τ )∥ , v t ψ2 + O(η 2 ρ 2 ) Lemma F.5 ⩽ O η 3/2 ρ 2 / √ p + O(η 2 ρ 2 ).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 thatP X t+1 -X s+1 + (t -s)C 2 η 2 ρ 2 > Q ⩽ 2 exp - Q 2 2(t -s) O(η 3/2 ρ 2 /p 1/2 + η 2 ρ 2 ) 2 .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 + 2C 1 ηρ 2 = (t -s)C 2 η 2 ρ 2 + 1 2 C 1 ηρ 2 , we have P t+1,s ⩽ P X t+1 ⩾ 2C 1 ηρ 2 ⩽P X t+1 -X s+1 + (t -s)C 2 η 2 ρ 2 > (t -s)C 2 η 2 ρ 2 + 1 2 C 1 ηρ 2 ⩽2 exp - (t -s)C 2 η 2 ρ 2 + 1 2 C 1 ηρ 2 2 2(t -s) O(η 3/2 ρ 2 /p 1/2 + η 2 ρ 2 ) 2 ⩽2 exp - 4(t -s)C 2 η 2 ρ 2 • 1 2 C 1 ηρ 2 4(t -s) (O(η 3 ρ 4 /p + η 4 ρ 4 )) ⩽ 2 exp -Ω 1 η + p -1.Therefore, we obtain the union bound:P ∃ t ∈ [T I , T I + T II ], L(θ t ) ⩾ 2C 1 ηρ 2 ⩽ TI+TII-1 t=T I P t+1,t + t-1 s=TI P t+1,s ⩽ TI+TII-1 t=T I t-1 s=TI P t+1,s ⩽ T 2 II exp -Ω 1 η + p -1 .Hence, with probability at least 1 -T 2 II exp -Ω1 η+p -1 , for any t ∈ [T I , T I + T II ], ∥θ t -Φ(θ t )∥ ⩽ 2 µ L(θ t ) ⩽ 2 C 1 µ √ ηρ = 4β 3/2 2 µ √ ηρ = O( √ ηρ).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:θ t+1 = θ t -η∇L θ t + ρ ξ t ∥ξ t ∥ -ηκP m+1:p (∇ 2 L(θ t ))∇L θ t + ρ ξ t ∥ξ t ∥, where ξ t ∼ N (0, I).Let the ρ in SAM and the κ in IRE satisfy:EQUATION)For simplicity, we denotev t := ξ t ∥ξ t ∥ , P (θ t ) := P m+1:p (∇ 2 L(θ t )), C 1 = 4β 3 2 µ ∨ 4β 2 β 2 3 µ , C 2 = β 3 2Following the proof for SAM, we denoteP t+1,t := P L(θ t+1 ) ⩾ 2C 1 ηρ 2 L(θ t ) < C 1 ηρ 2 , P t+1,s := P L(θ t+1 ) ⩾ 2C 1 ηρ 2 ; ∀ τ ∈ [s + 1, t], C 1 ηρ 2 ⩽ L(θ τ ) < 2C 1 ηρ 2 L(θ s ) < C 1 ηρ 2 , s ∈ [T I , t -1].and it holds thatP ∃ t ∈ [T I , T I + T II ], L(θ t ) ⩾ 2C 1 ηρ 2 ⩽ TI+TII-1 t=TI P t+1,t + t-1 s=TI P t+1,s . • Step I. Bounding P t+1,t . From L(θ t ) ⩽ C 1 ηρ 2 , we have ∥θ t -Φ(θ t )∥ ⩽ 2 µ L(θ t ) = O( √ ηρ), thus ∥θ t + ρv t -Φ(θ t )∥ ⩽ ∥θ t -Φ(θ t )∥ + ρ = O( √ ηρ) + O(ρ) < R, which means θ t + ρv t ∈ B(K; R). Furthermore, ∥θ t+1 -Φ(θ t )∥ ⩽ ∥θ t -Φ(θ t )∥ + η ∥∇L(θ t + ρv t )∥ + ηκ ∥P (θ t )∇L(θ t + ρv t )∥ Lemma D.5 ⩽ O( √ ηρ) + η ∥∇L (θ t )∥ + β 2 ηρ + O ηκρ 2 ⩽O( √ ηρ) + O(η 3/2 ρ) + O (ηρ) + O (ηρ) ⩽O( √ ηρ) + O(η 3/2 ρ) + O (ηρ) ⩽ R,which implies θ t+1 ∈ B(K; R). Consequently, we have the following quadratic upper bound:L(θ t+1 ) = L (θ t -η∇L (θ t + ρv t ) -ηκP (θ t )∇L (θ t + ρv t )) ⩽L(θ t ) -η ⟨∇L(θ t ), ∇L (θ t + ρv t ) + κP (θ t )∇L (θ t + ρv t )⟩ + η 2 β 2 2 ∥∇L (θ t + ρv t ) + κP (θ t )∇L (θ t + ρv t )∥ 2 ⩽L(θ t ) -η ⟨∇L(θ t ), ∇L (θ t + ρv t )⟩ -κη ⟨∇L(θ t ), P (θ t )∇L (θ t + ρv t )⟩ + η 2 β 2 ∥∇L (θ t + ρv t )∥ 2 + κ 2 ∥P (θ t )∇L (θ t + ρv t )∥ 2 ⩽L(θ t ) -η ∥∇L(θ t )∥ 2 + ηρβ 2 ∥∇L(θ t )∥ -κη ⟨∇L(θ t ), P (θ t )∇L(θ t )⟩ -κηρ ∇L(θ t ), P (θ t )∇ 2 L(θ t )v t + κη β 3 ρ 2 2 ∥P (θ t )∇L(θ t )∥ + η 2 β 2 (∥∇L(θ t )∥ + ρβ 2 ) 2 + κ 2 ∥P (θ t )∇L (θ t + ρv t )∥ 2 ⩽L(θ t ) + ηρβ 2 ∥∇L(θ t )∥ + κηρ P (θ t )∇ 2 L(θ t )∇L(θ t ) + κηρ 2 β 3 2 ∥P (θ t )∇L(θ t )∥ + η 2 β 2 (∥∇L(θ t )∥ + ρβ 2 ) 2 + κ 2 ∥P (θ t )∇L (θ t + ρv t )∥ 2 Lemma D.5 ⩽ L(θ t ) + O(ηρ ∥∇L(θ t )∥) + O(κηρ ∥θ t -Φ(θ t )∥ ∥∇L(θ t )∥) + O κηρ 2 ∥θ t -Φ(θ t )∥ 2 + O η 2 ρ 2 + O η 2 κ 2 ρ 4 ⩽C 1 ηρ 2 + o(ηρ 2 ) ⩽ 3 2 C 1 ηρ 2 .Thus, we obtainP t+1,t = P L(θ t+1 ) ⩾ 2C 1 ηρ 2 L(θ t ) < C 1 ηρ 2 = 0. • Step II. Bounding P t+1,s for s ∈ [T I , t -1].We prove this step under the conditionEQUATIONIt is clear thatP t+1,s ⩽ P X t+1 ⩾ 2C 1 ηρ 2 .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 ifC 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 1 ηρ 2 , then E[L(θ τ +1 )] ⩽ L(θ τ ) -C 2 η 2 ρ 2 . If C 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 1 ηρ , similar toStep I, it is clear that θ τ +1 ∈ B(K; R). Applying the quadratic upper bound, it holds that:L(θ τ +1 ) = L (θ τ -η∇L (θ τ + ρv τ ) -ηκP (θ τ )∇L (θ τ + ρv τ )) ⩽L(θ τ ) -η ⟨∇L(θ τ ), ∇L (θ τ + ρv τ ) + κP (θ t )∇L (θ τ + ρv τ )⟩ + η 2 β 2 2 ∥∇L (θ τ + ρv τ ) + κP (θ τ )∇L (θ τ + ρv τ )∥ 2 ⩽L(θ τ ) -η ⟨∇L(θ τ ), ∇L (θ τ + ρv τ )⟩ -κη ⟨∇L(θ τ ), P (θ t )∇L (θ τ + ρv τ )⟩ + η 2 β 2 ∥∇L (θ τ + ρv τ )∥ 2 + κ 2 ∥P (θ τ )∇L (θ τ + ρv τ )∥ 2 ⩽L(θ τ ) -η ∥∇L(θ τ )∥ 2 -ηρ ∇L(θ τ ), ∇ 2 L(θ t )v τ + ηρ 2 β 3 2 ∥∇L(θ τ )∥ -κη ⟨∇L(θ τ ), P (θ τ )∇L(θ τ )⟩ -κηρ ∇L(θ τ ), P (θ τ )∇ 2 L(θ τ )v τ + κη β 3 ρ 2 2 ∥P (θ τ )∇L(θ τ )∥ + η 2 β 2 (∥∇L(θ τ )∥ + ρβ 2 ) 2 + κ 2 ∥P (θ τ )∇L (θ τ + ρv τ )∥ 2 .Taking the expectation, we have:E[L(θ τ +1 )] ⩽L(θ τ ) -η ∥∇L(θ τ )∥ 2 + ηρ 2 β 3 2 ∥∇L(θ τ )∥ -κη ⟨∇L(θ τ ), P (θ τ )∇L(θ τ )⟩ + κη β 3 ρ 2 2 ∥P (θ τ )∇L(θ τ )∥ + η 2 β 2 (∥∇L(θ τ )∥ + ρβ 2 ) 2 + κ 2 ∥P (θ τ )∇L (θ τ + ρv τ )∥ 2 ⩽L(θ τ ) -η ∥∇L(θ τ )∥ 2 + ηρ 2 β 3 2 ∥∇L(θ τ )∥ + κη β 3 ρ 2 2 ∥P (θ τ )∇L(θ τ )∥ + 2η 2 β 2 ∥∇L(θ τ )∥ 2 + 2β 3 2 η 2 ρ 2 + β 2 η 2 κ 2 ∥P (θ τ )∇L (θ τ + ρv τ )∥ 2 Lemma D.5 ⩽ L(θ τ ) - 3η 4 ∥∇L(θ τ )∥ 2 + ηρ 2 β 3 2 ∥∇L(θ τ )∥ + O κηρ 2 • ηρ 2 + 2β 3 2 η 2 ρ 2 + β 2 η 2 κ 2 O( √ ηρ 2 ) + β 3 2 ρ 2 2 ⩽L(θ τ ) - 3η 4 ∥∇L(θ τ )∥ 2 + ηρ 2 β 3 2 ∥∇L(θ τ )∥ + 2β 3 2 η 2 ρ 2 + β 2 β 2 3 η 2 κ 2 ρ 4 + o(η 2 ρ 2 ). From C 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 1 ηρ 2 and ρ = O( √ η), it holds that ∥∇L(θ τ )∥ ⩾ 2µL(θ τ ) ⩾ 2C 1 µ √ ηρ ⩾ 4β 3 ρ 2 .Moreover, recall κ ⩽ 1/ρ. Therefore, we have the upper bound:E[L(θ τ +1 )] ⩽L(θ τ ) - 3η 4 ∥∇L(θ τ )∥ 2 + ηρ 2 β 3 2 ∥∇L(θ τ )∥ + 2β 3 2 η 2 ρ 2 + β 2 β 2 3 η 2 κ 2 ρ 4 + o(η 2 ρ 2 ) ⩽L(θ τ ) - 5η 8 ∥∇L(θ τ )∥ 2 + 2β 3 2 η 2 ρ 2 + β 2 β 2 3 η 2 κ 2 ρ 4 + o(η 2 ρ 2 ) ⩽L(θ τ ) - 10 8 ηµL(θ τ ) + 2β 3 2 η 2 ρ 2 + β 2 β 2 3 η 2 ρ 2 + o(η 2 ρ 2 ).⩽L(θ τ ) -10 2β 3 2 µ ∨ β 2 β 2 3 µ η 2 ρ 2 + 2β 3 2 η 2 ρ 2 + β 2 β 2 3 η 2 ρ 2 + o(η 2 ρ 2 ) ⩽L(θ τ ) -β 3 2 η 2 ρ 2 = L(θ τ ) -C 2 η 2 ρ 2 . -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 caseC 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 1 ηρ 2 . If C 1 ηρ 2 ⩽ L(θ τ ) ⩽ 2C 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:L(θ τ +1 ) L (θ τ -η∇L (θ τ + ρv τ ) -ηκP (θ τ )∇L (θ τ + ρv τ )) =L(θ τ ) -η ⟨∇L(θ τ ), ∇L (θ τ + ρv τ ) + κP (θ t )∇L (θ τ + ρv τ )⟩ + η 2 ∥∇L (θ τ + ρv τ ) + κP (θ τ )∇L (θ τ + ρv τ )∥ 2 =L(θ τ ) -η ∥∇L(θ τ )∥ 2 -ηρ ∇L(θ τ ), ∇ 2 L(θ τ )v τ + O ηρ 2 ∥∇L(θ τ )∥ -ηκ ⟨∇L(θ τ ), P (θ τ )∇L(θ τ )⟩ -ηκρ ∇L(θ τ ), P (θ τ )∇ 2 L(θ τ )v t + O ηκρ 2 ∥P (θ τ )∇L(θ τ )∥ + η 2 ∥∇L (θ τ + ρv τ ) + κP (θ τ )∇L (θ τ + ρv τ )∥ 2 .In the same way as the proof of Claim I, it holds thatη ∥∇L(θ τ )∥ 2 = O(η 2 ρ 2 ), ηρ 2 ∥∇L(θ τ )∥ = O(η 2 ρ 2 ), ηκρ 2 ∥P (θ τ )∇L(θ τ )∥ = O η 2 ρ 3 , η 2 ∥∇L (θ τ + ρv τ ) + κP (θ τ )∇L (θ τ + ρv τ )∥ 2 = O η 2 ρ 2Thus,L(θ τ +1 ) -L(θ τ ) = -ηρ ∇L(θ τ ), ∇ 2 L(θ τ )v t -ηκρ ∇L(θ τ ), P (θ τ )∇ 2 L(θ τ )v t + O(η 2 ρ 2 ),which implies:X τ +1 -X τ -C 2 η 2 ρ 2 ψ2 ⩽ ∥L(θ τ +1 ) -L(θ τ )∥ ψ2 + C 2 η 2 ρ 2 ψ2 ⩽ηρ ∇L(θ τ ), ∇ 2 L(θ τ )v t ψ2 + ηκρ ∇L(θ τ ), P (θ τ )∇ 2 L(θ τ )v t ψ2 + O(η 2 ρ 2 ) ⩽ηρ ∇ 2 L(θ τ )∇L(θ τ ) ∇ 2 L(θ τ )∇L(θ τ ) ∥∇ 2 L(θ τ )∇L(θ τ )∥ , v t ψ2 + ηκρ ∇ 2 L(θ τ )P (θ τ )∇L(θ τ ) ∇ 2 L(θ τ )P (θ τ )∇L(θ τ ) ∥∇ 2 L(θ τ )P (θ τ )∇L(θ τ )∥ , v t ψ2 + O(η 2 ρ 2 ) Lemma D.5 ⩽ O η 3/2 ρ 2 ∇ 2 L(θ τ )∇L(θ τ ) ∥∇ 2 L(θ τ )∇L(θ τ )∥ , v t ψ2 + O η 3/2 ρ 2 ∇ 2 L(θ τ )P (θ τ )∇L(θ τ ) ∥∇ 2 L(θ τ )P (θ τ )∇L(θ τ )∥ , v t ψ2 + O(η 2 ρ 2 ) Lemma F.5 ⩽ O η 3/2 ρ 2 / √ p + O(η 2 ρ 2 ).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 thatP X t+1 -X s+1 + (t -s)C 2 η 2 ρ 2 > Q ⩽ 2 exp - Q 2 2(t -s) O(η 3/2 ρ 2 /p 1/2 + η 2 ρ 2 ) 2 .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 + 2C 1 ηρ 2 = (t -s)C 2 η 2 ρ 2 + 1 2 C 1 ηρ 2 , we have P t+1,s ⩽ P X t+1 ⩾ 2C 1 ηρ 2 ⩽P X t+1 -X s+1 + (t -s)C 2 η 2 ρ 2 > (t -s)C 2 η 2 ρ 2 + 1 2 C 1 ηρ 2 ⩽2 exp - (t -s)C 2 η 2 ρ 2 + 1 2 C 1 ηρ 2 2 2(t -s) O(η 3/2 ρ 2 /p 1/2 + η 2 ρ 2 ) 2 ⩽2 exp - 4(t -s)C 2 η 2 ρ 2 • 1 2 C 1 ηρ 2 4(t -s) (O(η 3 ρ 4 /p + η 4 ρ 4 )) ⩽ 2 exp -Ω 1 η + p -1.Therefore, we obtain the union bound:P ∃ t ∈ [T I , I + T II ], L(θ t ) ⩾ 2C 1 ηρ 2 ⩽ TI+TII-1 t=T I P t+1,t + t-1 s=TI P t+1,s ⩽ TI+TII-1 t=T I t-1 s=TI P t+1,s ⩽ T 2 II exp -Ω 1 η + p -1 .Hence, with probability at least 1 -T 2 II exp -Ω1 η+p -1 , for any t ∈ [T I , T I + T II ], ∥θ t -Φ(θ t )∥ ⩽ 2 µ L(θ t ) ⩽ 2 C 1 µ √ ηρ = 4β 3/2 2 µ √ ηρ = O( √ ηρ).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 ],∥θ t+1 -θ t ∥ ⩽η ∥∇L(θ t + ρv t )∥ + ηκ ∥P (θ t )∇L(θ t + ρv t )∥ ⩽η ∥∇L(θ t )∥ + O(ηρ) + ηκ ∥P (θ t )∇L(θ t )∥ + ηκρ P (θ t )∇ 2 L(θ t ) + O(ηκρ 2 ) Lemma D.5 ⩽ O(η 3/2 ρ) + O(ηρ) + O(η 3/2 ρ) + O(η 3/2 ρ 2 ) + O(ηρ 2 ) + O(ηρ) = O(ηρ).Then by Taylor's expansion,Φ(θ t+1 ) -Φ(θ t ) = ∂Φ(θ t ) (θ t+1 -θ t ) + O ∥θ t+1 -θ t ∥ 2 = -η∂Φ(θ t )∇L(θ t + ρv t ) -ηκ∂Φ(θ t )P (θ t )∇L(θ t + ρv t ) + O(η 2 ρ 2 ).For the term ∂Φ(θ t )∇L(θ t + ρv t ) and ∂Φ(θ t )P (θ t )∇L(θ t + ρv t ), using Taylor's expansion, we have∂Φ(θ t )∇L(θ t + ρv t ) =∂Φ(θ t )∇L(θ t ) + ρ∂Φ(θ t )∇ 2 L(θ t )v t + ρ 2 2 ∂Φ(θ t )∇ Tr v t ∇ 2 L(θ t )v ⊤ t + O(ρ 3 ) =∂Φ(θ t )∇L(θ t ) + ρ∂Φ(θ t )∇ 2 L(θ t )v t + ρ 2 2 ∂Φ(θ t )∇ v ⊤ t ∇ 2 L(θ t )v t + O(ρ 3 ), ∂Φ(θ t )P (θ t )∇L(θ t + ρv t ) =∂Φ(θ t )P (θ t )∇L(θ t ) + ρ∂Φ(θ t )P (θ t )∇ 2 L(θ t )v t + ρ 2 2 ∂Φ(θ t )P (θ t )∇ Tr v t ∇ 2 L(θ t )v ⊤ t + O(ρ 3 ) =∂Φ(θ t )P (θ t )∇L(θ t ) + ρ∂Φ(θ t )P (θ t )∇ 2 L(θ t )v t + ρ 2 2 ∂Φ(θ t )P (θ t )∇ v ⊤ t ∇ 2 L(θ t )v t + O(ρ 3 ).Taking the expectation (about v t ), we haveE[v t ] = 0, E v ⊤ t ∇ 2 L(θ t )v t = Tr ∇ 2 L(θ t ) p .Additionally, using Lemma D.2 and Taylor's expansion, we have:∂Φ(θ t )∇L(θ t ) = 0; ∂Φ(θ t ) = ∂Φ(Φ(θ t )) + ∥θ t -Φ(θ t )∥ = ∂Φ(Φ(θ t )) + O( √ ηρ); ∂Φ(θ t )P (θ t ) = ∂Φ(Φ(θ t ))P (Φ(θ t )) + ∥θ t -t )∥ = ∂Φ(Φ(θ t )) + O( √ ηρ); ∇ Tr ∇ 2 L(θ t ) = ∇ Tr ∇ 2 L(Φ(θ t )) + ∥θ t -Φ(θ t )∥ = ∇ Tr ∇ 2 L(Φ(θ t )) + O( √ ηρ); ∂Φ(θ t )P (θ t )∇L(θ t ) =∂Φ(Φ(θ t ))P (Φ(θ t ))∇L(θ t ) + O (∥θ t -Φ(θ t )∥ ∥∇L(θ t )∥) =∂Φ(Φ(θ t ))∇L(θ t ) + O ηρ 2 = 0 + O(ηρ 2 ).Combining the results above, we obtain:E[Φ(θ t+1 )] =Φ(θ t ) -η∂Φ(θ t )∇L(θ t ) -ηκ∂Φ(θ t )P (θ t )∇L(θ t ) - ηρ 2 2p ∂Φ(θ t )∇ Tr ∇ 2 L(θ t ) - κηρ 2 2p ∂Φ(θ t )P (θ t )∇ Tr ∇ 2 L(θ t ) + O(ηρ 3 ) =Φ(θ t ) + O η 2 κρ 2 - (κ + 1)ηρ 2 2p ∂Φ(θ t )∇ Tr ∇ 2 L(θ t ) + O(η 3/2 ρ 3 ) + O(κη 3/2 ρ 3 ) + O(ηρ 3 ) =Φ(θ t ) + O η 2 κρ 2 - (κ + 1)ηρ 2 2p ∂Φ(θ t )∇ Tr ∇ 2 L(Φ(θ t )) + O(η 3/2 ρ 3 ) + O(κη 3/2 ρ 3 ) + O(ηρ 3 ) =Φ(θ t ) -(κ + 1)ηρ 2 ∂Φ(θ t )∇ Tr ∇ 2 L(Φ(θ t ))/2p + O(η 3/2 ρ 2 ).E Proofs in Section 5.2.2Setting 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 (•)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• (i) B(K; R 1 ) ⊂ U ; • (ii) L i (•) (i ∈ [n]) and L(•) are µ-PL on B(K; R 1 ); • (iii) inf θ∈B(K;R1) λ n ∇ 2 L(θ) ⩾ µ; inf θ∈B(K;R1) λ 1 ∇ 2 L i (θ) ⩾ µ, ∀i ∈ [n].We further define the following absolute constants on B(K; R):β 2 := sup θ∈B(K;R1) ∇ 2 L(θ) ∨ max i∈[n] sup θ∈B(K;R1)∇ 2 L i (θ) ;β 3 := sup θ∈B(K;R1) ∇ 3 L(θ) ∨ max i∈[n] sup θ∈B(K;R1) ∇ 3 L i (θ) ; ν := inf θ∈B(K;R1) λ m ∇ 2 L(θ) ∧ min i∈[n] inf θ∈B(K;R1) λ 1 ∇ 2 L i (θ) ; ζ Φ 1 := sup θ∈B(K;R1) ∥∇Φ(θ)∥ , ζ Φ 2 := sup θ∈B(K;R1) ∇ 2 Φ(θ) .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 ),P n+1:p (∇ 2 L(θ)) -P n+1:p (∇ 2 L(Φ(θ))) ⩽ ζ P ∥θ -Φ(θ)∥ .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:β 2 := sup θ∈B(K;R) ∇ 2 L(θ) ∨ max i∈[n] sup θ∈B(K;R)∇ 2 L i (θ) ;EQUATION)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 (∥θ -Φ(θ)∥) ; • P n+1:p (∇ 2 L(θ))∇L(θ) , P n+1:p (∇ 2 L(θ))∇L i (θ) , ∥∂Φ(θ)∇L i (θ)∥ ⩽ O ∥θ -Φ(θ)∥ 2 ; • Let ρ > 0 and v ∈ S p-1 . If θ + ρv ∈ B(K; R), then ∥∇L i (θ + ρv)∥ ⩽ ∥∇L i (θ)∥ + ρβ 2 , ∀i ∈ [n]; ; P n+1:p (∇ 2 L(θ))∇L i (θ + ρv) ⩽ O ∥θ -Φ(θ)∥ 2 + O (ρ ∥θ -Φ(θ)∥) + ρ 2 β 3 2 , ∀i ∈ [n];Lemma E.8 (Lemma H.9 in Wen et al. (2023a) ). For any absolute constant C > 0, there exist absolute C 1 , C 2 > 0 such that: if θ t ∈ B(K; R) and C 1 ηρ ⩽ ∥θ t -Φ(θ t )∥ ⩽ Cρ, then it holds that:E it θ t+1/2 -Φ(θ t+1/2 ) ⩽ ∥θ t -Φ(θ t )∥ -C 2 ηρ.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θ t+1/2 -Φ(θ t+1/2 ) -∥θ t -Φ(θ t )∥ ⩽ C 3 ηρ.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.6This 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θ t+1 =θ t -η∇L it θ t + ρ ∇L it (θ t ) ∥∇L it (θ t )∥ -ηκP n+1:p ∇ 2 L(θ t ) ∇L it θ t + ρ ∇L it (θ t ) ∥∇L it (θ t )∥ , where i t ∼ U([n]).Let the κ in IRE satisfy κ ⩽ 1/ρ.Additionally, we fix a constant α ∈ (0, 1) in the proof.For simplicity, we denotev t := ∇L it (θ t )∥∇L it (θ t )∥ , P (θ t ) := P n+1:p ∇ 2 L(θ t ) ; and θ t+1/2 =θ t -∇L it (θ t + ρv t ) ;θ t+1 =θ t+1/2 -κηP (θ t )∇L it (θ t + ρv t ) .Additionally, we denoteP t+1,t := P ∥θ t+1 -Φ(θ t+1 )∥ ⩾ 2Cη 1-α ρ ∥θ t -Φ(θ t )∥ < Cη 1-α ρ , P t+1,s :=P ∥θ t+1 -Φ(θ t+1 )∥ ⩾ 2Cη 1-α ρ; ∀τ ∈ [s + 1, t], Cη 1-α ρ ⩽ ∥θ τ -Φ(θ τ )∥ < 2Cη 1-α ρ ∥θ s -Φ(θ s )∥ < C √ ηρ , s ∈ [T I , t -1].Then the following bound holds naturally:P ∃ t ∈ [T I , T I + T II ], ∥θ t -Φ(θ t )∥ ⩾ 2Cη 1-α ρ ⩽ TI+TII-1 t=TI P t+1,t +t-1 s=TI P t+1,s .• Step I. Bounding P t+1,t .From ∥θ t -Φ(θ t )∥ = O(η 1-α ρ), thus∥θ t + ρv t -Φ(θ t )∥ ⩽ ∥θ t -Φ(θ t )∥ + ρ = O(η 1-α ρ) + O(ρ) < R,which means θ t + ρv t ∈ B(K; R). Using Lemma E.7 and κ ⩽ 1/ρ, we canθ t+1 -θ t+1/2 = ηκ ∥∇P (θ t )∇L it (θ t + ρv t )∥ ⩽ηκ ∥∇P (θ t )∇L it (θ t )∥ + ρ ∇P (θ t )∇ 2 L it (θ t ) + β 3 2 ρ 2 ⩽ηκ ηρ 2 + η 1-α ρ 2 + β 3 2 ρ 2 ⩽ β 3 ηκρ 2 ⩽ C 2 2(1 + ζ Φ 1 )ηρ,θ t+1 -Φ(θ t+1 ) -(θ t+1/2 -Φ(θ t+1/2 )) ⩽(1 + ζ Φ 1 ) θ t+1 -θ t+1/2 ⩽ C 2 2 ηρ.Then we have the following bound:∥θ t+1 -Φ(θ t+1 )∥ ⩽ ∥θ t -Φ(θ t )∥ + θ t+1/2 -Φ(θ t+1/2 ) -(θ t -Φ(θ t )) + θ t+1 -Φ(θ t+1 ) -(θ t+1/2 -Φ(θ t+1/2 )) Lemma E.9 ⩽ C 1 η 1-α ρ + O(ηρ) + θ t+1 -Φ(θ t+1 ) -(θ t+1/2 -Φ(θ t+1/2 )) ⩽C 1 η 1-α ρ + O(ηρ) + O(ηρ) < 3C 1 2 η 1-α ρ.Thus, we obtainP t+1,t = P ∥θ τ +1 -Φ(θ τ +1 )∥ ⩾ 2Cη 1-α ρ ∥θ t -Φ(θ t )∥ < Cη 1-α ρ = 0. • 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 :EQUATIONIt is clear thatP t+1,s ⩽ P X t+1 ⩾ 2Cη 1-α ρ .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-α ρ ⩽ X τ = ∥θ τ -Φ(θ τ )∥ ⩽ 2Cη 1-α ρ, then E ∥θ τ +1 -Φ(θ τ +1 )∥ ⩽ ∥θ τ -Φ(θ τ )∥ -C 2 ηρ/2. If Cη 1-α ρ ⩽ X τ = ∥θ τ -Φ(θ τ )∥ ⩽ 2Cη 1-α ρ, similar to Step I, it holds that θ τ +1 ∈ B(K; R) and θ t+1 -Φ(θ τ +1 ) -(θ τ +1/2 -Φ(θ τ +1/2 )) ⩽ C2 2 ηρ. Moreover, ∥θ τ +1 -Φ(θ τ +1 )∥ ⩽ θ τ +1/2 -Φ(θ τ +1/2 ) + θ τ +1 -Φ(θ τ +1 ) -(θ τ +1/2 -Φ(θ τ +1/2 )) ⩽ θ τ +1/2 -Φ(θ τ +1/2 ) + C 2 2 ηρ.Taking the expectation and using Lemma E.8, we haveE ∥θ τ +1 -Φ(θ τ +1 )∥ ⩽ E θ τ +1/2 -Φ(θ τ +1/2 ) + C 2 2 ηρ ⩽ ∥θ τ -Φ(θ τ )∥ -C 2 ηρ + C 2 2 ηρ = ∥θ τ -Φ(θ τ )∥ - C 2 2 ηρ. -Claim II. X τ +1 -X τ + C 2 ηρ/2 is O(ηρ)-bounded.From the definition of τ , we only need to prove for the case Cη1-α ρ ⩽ X τ = ∥θ τ -Φ(θ τ )∥ ⩽ 2Cη 1-α ρ. If Cη 1-α ρ ⩽ X τ = ∥θ τ -Φ(θ τ )∥ ⩽ 2Cη 1-α ρ, we have θ τ +1 ∈ B(K; R) and θ t+1 -Φ(θ t+1 ) -(θ t+1/2 -Φ(θ t+1/2 )) ⩽ C2 2 ηρ.Combining this result and Lemma E.9, we have|∥θ τ +1 -Φ(θ τ +1 )∥ -∥θ τ -Φ(θ τ )∥| ⩽ ∥θ τ +1 -Φ(θ τ +1 )∥ -θ τ +1/2 -Φ(θ τ +1/2 ) + θ τ +1/2 -Φ(θ τ +1/2 ) -∥θ τ -Φ(θ τ )∥ ⩽ (θ τ +1 -Φ(θ τ +1 )) -(θ τ +1/2 -Φ(θ τ +1/2 )) + θ τ +1/2 -Φ(θ τ +1/2 ) -∥θ τ -Φ(θ τ )∥ ⩽ C 2 2 ηρ + C 3 ηρ = O(ηρ).With the preparation of Claim I and Claim II, we can use the Azuma-Hoeffeding inequality: for any Q > 0, it holds thatP (X t+1 -X s+1 + (t -s)C 2 ηρ/2 > Q) ⩽ 2 exp - Q 2 2(t -s)O(η 2 ρ 2 ).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-α ρ + 2Cη 1-α ρ = (t - s)C 2 ηρ/2 + Cη 1-α ρ/2, we have P t+1,s ⩽ P X t+1 ⩾ 2Cη 1-α ρ ⩽P X t+1 -X s+1 + (t -s)C 2 ηρ/2 > (t -s) C 2 2 ηρ + C 2 η 1-α ρ ⩽2 exp - (t -s)C 2 ηρ/2 + Cη 1-α ρ/2 2 2(t -s)O(η 2 ρ 2 ) ⩽2 exp - (t -s)C 2 ηρ • Cη 1-α ρ 4(t -s)O(η 2 ρ 2 ) ⩽ 2 exp -Ω 1 η α .Therefore, we obtain the union bound:P ∃ t ∈ [T I , T I + T II ], ∥θ t -Φ(θ t )∥ ⩾ 2Cη 1-α ρ ⩽ TI+TII-1 t=T I P t+1,t + t-1 s=TI P t+1,s ⩽ TI+TII-1 t=T I t-1 s=TI P t+1,s ⩽ T 2 II exp (-Ω (1/η α )) .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.t -Φ(θ t )∥ = O(η 1-α ρ), we have: ∥θ t+1 -θ t ∥ ⩽η ∥∇L it (θ t + ρv t )∥ + ηκ ∥P (θ t )∇L it (θ t + ρv t )∥ Lemma E.7 ⩽ η ∥∇L it (θ t )∥ + O(ηρ) + O(ηκρ 2 ) ⩽ O(ηρ).Then by Taylor's expansion,Φ(θ t+1 ) -Φ(θ t ) = ∂Φ(θ t )(θ t+1 -θ t ) + O ∥θ t+1 -θ t ∥ 2 = -η∂Φ(θ t )∇L it (θ t + ρv t ) -ηκ∂Φ(θ t )P (θ t )∇L it (θ t + ρv t ) + O(η 2 ρ 2 ).For the term ∂Φ(θ t )∇L it (θ t + ρv t ) and ∂Φ(θ t )P (θ t )∇L it (θ t + ρv t ), using Taylor's expansion and Lemma E.7, we haveEQUATION∂Φ(θ t )P (θ t )∇L it (θ t + ρv t ) =∂Φ(θ t )P (θ t )∇L it (θ t ) + ρ∂Φ(θ t )P (θ t )∇ 2 L it (θ t )v t + ρ 2 2 ∂Φ(θ t )P (θ t )∇ Tr v t ∇ 2 L it (θ t )v ⊤ t + O(ρ 3 ) =O(∥θ t -Φ(θ t )∥ 2 ) + O(ρ ∥θ t -Φ(θ t )∥) + ρ 2 2 ∂Φ(θ t )P (θ t )∇ v ⊤ t ∇ 2 L it (θ t )v t + O(ρ 3 ) =O(η 1-α ρ 2 ) + ρ 2 2 ∂Φ(Φ(θ t ))∇ ∇L it (Φ(θ t )) ∥∇L it (Φ(θ t ))∥ ⊤ ∇ 2 L it (Φ(θ t )) ∇L it (Φ(θ t )) ∥∇L it (Φ(θ t ))∥ + O(ρ 2 ∥θ t -Φ(θ t )∥) + O(ρ 3 ) = ρ 2 2 ∂Φ(Φ(θ t ))∇λ 1 ∇ 2 L it (Φ(θ t )) + O(η 1-α ρ 2 + ρ 3 ), .Combining the results above, we obtain:Φ(θ t+1 ) =Φ(θ t ) -(1 + κ) ηρ 2 2 ∂Φ(Φ(θ t ))∇λ 1 ∇ 2 L it (Φ(θ t )) + O(κη 2-α ρ 2 + κηρ 3 ).Additionally, taking the expectation, we have:E it ∂Φ(Φ(θ t ))∇λ 1 ∇ 2 L it (Φ(θ t )) = ∂Φ(Φ(θ t ))∇ Tr ∇ 2 L(Φ(θ t )) .Therefore, we obtainE it [Φ(θ t+1 )] =Φ(θ t ) -(1 + κ) ηρ 2 2 ∂Φ(Φ(θ t ))∇ Tr ∇ 2 L(Φ(θ t )) + O(κη 2-α ρ 2 + κηρ 3 ) =Φ(θ t ) -(1 + κ) ηρ 2 2 ∂Φ(Φ(θ t ))∇ Tr ∇ 2 L(Φ(θ t )) + h.o.t.. 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 |λ k -µ k | ⩽ ∥A -B∥ . Lemma F.3 (Davis-Kahan sin(Θ) theorem). Let A, B ∈ R p×p be symmetric matrices. Denote their orthogonal 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 ∥•∥ ⋆ ,F ⊤ 2 E 1 ⋆ ⩽ F ⊤ 2 (A -B)E 1 ⋆ δ .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:P (X n -X 0 ⩾ t) ⩽ 2 exp - t 2 2n(α + β) 2 . • (ii) (Sub-Gaussian martingale difference). If X i+1 -X i is σ 2i -sub-Gaussian, then for any n, t > 0, we have:P (X n -X 0 ⩾ t) ⩽ 2 exp - t 2 2 n i=1 σ 2 i . 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,P v ∥v∥ , g ∥g∥ ⩾ t ⩽ 4e -cpt 2 .Proof of Lemma F.5. From P54 in Vershynin ( 2018), there exists an absolute constant c > 0 such that for any t > 0, P |⟨e1,g⟩| ∥g∥ ⩾ t ⩽ 4e -cpt 2 . Without loss of generality, we can assume v ̸ = 0. Then we have:P v ∥v∥ , g ∥g∥ ⩾ t √ p = P ⟨e 1 , g⟩ ∥g∥ ⩾ t ⩽ 4e -cpt 2 .Lemma F.6. ∥AB∥ F ⩽ ∥A∥ ∥B∥ F .• It is OK to report 1-sigma error bars, but one should state it. The authors should 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.† Correspondence to: Mingze Wang and Lei Wu. 


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

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

Reviews: review: paper_summary: paper_summary: The authors propose IRE to enhance the implicit regularization of base optimizers, thereby improving the generalization and convergence in deep learning. IRE decouples the dynamics of flat and sharp directions, reducing sharpness along flat directions while maintaining stability in sharp directions. The paper provides theoretical evidence that IRE can substantially expedite convergence towards flat minima in SAM.
strength_weakness: strength_weakness: 1. The improvements of IRE in Table 1 and 7 are not as significant as expected. It makes one wonder about its usefulness for CNN networks and its oversensitivity to hyperparameters
2. 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.
3. Judging from the code, the cost of each training step is at least twice that of SGD.
questions: questions: 1. Could IRE help in the finetuning phase of LLMs, Including both convergence speed and convergence position properties?
limitations: limitations: 1. The author did not provide information on the computing resources they used.
paper_summary: paper_summary: This paper proposes an Implicit Regularization Enhancement (IRE) framework to accelerate the convergence of optimization algorithms towards flat minima in deep learning, thereby improving generalization and convergence. The key idea behind IRE is to decouple the dynamics along flat and sharp directions in the optimization landscape. It speeds up the dynamics along flat directions while keeping the dynamics in sharp directions unchanged. This allows IRE to substantially boost the implicit sharpness reduction of the base optimizer without hurting training stability.

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

Furthermore, 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. 

In summary, the key contributions are:

1. The IRE framework that enhances implicit regularization by decoupling and accelerating optimization dynamics along flat directions
2. A practical and efficient implementation of IRE that can be incorporated into generic base optimizers 
3. Extensive empirical results showing IRE improves generalization and speeds up convergence across vision and language tasks
4. Theoretical analysis proving IRE can substantially accelerate the sharpness reduction of the SAM optimizer
strength_weakness: strength_weakness: ### Comparison with Related Works
While 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.

### Limitations and Failure Cases
The paper could be strengthened by providing a more in-depth discussion of the potential limitations and failure cases of IRE. For example:
- 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$)? Is there a risk of instability or divergence if these hyperparameters are not tuned properly?
- Are there any particular types of models, datasets, or tasks where IRE may not provide significant benefits or even hurt performance?

Addressing these questions would help readers better understand the scope and applicability of IRE.

### Experimental Setup and Implementation Details
The 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:
- The specific hyperparameter tuning ranges for $\lambda$ and $\gamma$ used in the experiments
- The architectures of the models used (e.g., ResNet and ViT variants)
- The data augmentation and preprocessing techniques applied

Including these details would make the experiments more self-contained and easier to interpret and reproduce.

### Comparisons with State-of-the-Art Optimizers
To 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:
- Adaptive gradient methods like Adagrad (Duchi et al., 2011), Adam (Kingma and Ba, 2014), and their variants
- Second-order optimization methods like K-FAC (Martens and Grosse, 2015) and Shampoo (Gupta et al., 2018)
- Other regularization techniques like weight decay, dropout, and label smoothing

These comparisons would provide a more comprehensive understanding of IRE's performance and potential advantages over existing methods.

### Downstream Performance Evaluation
While 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.

However, 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.

### Theoretical Analysis
The 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:
- Discussing the practical significance of Assumption 5.1 (manifold of minimizers) and how it relates to the empirical observations in deep learning
- Providing a high-level interpretation of the non-asymptotic bounds and their dependence on the hyperparameters (e.g., $\eta$, $\rho$, $\lambda$)

Additionally, 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.
questions: questions: 1. **Comparison with related works:**
   - 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)?
   - Suggestion: Provide a more detailed discussion comparing and contrasting IRE with these related approaches to better highlight the novelty and advantages of IRE.

2. **Limitations and failure cases:**
   - 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?
   - 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.

3. **Experimental setup and implementation details:**
   - 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?
   - Suggestion: Include these key details in the main text to make the experiments more self-contained and easier to interpret and reproduce.

4. **Comparisons with state-of-the-art optimizers:**
   - 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)?
   - 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.

5. **Downstream performance evaluation:**
   - 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?
   - 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.

6. **Theoretical analysis:**
   - 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$)?
   - 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.

Addressing these questions and incorporating the suggestions in the rebuttal or a revised version of the paper would help to strengthen the work and provide a more comprehensive understanding of the proposed IRE framework, its novelty, effectiveness, and impact.
limitations: limitations: 1. **Limitations:** 
   - 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.

2. **Potential negative societal impact:** 
   - 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. 
   - Propose mitigation strategies or areas for future research to address these concerns.

3. **Reproducibility and transparency:** 
   - Provide clear instructions for reproducing the experiments, make the code and pre-trained models publicly available, and discuss any limitations or challenges in reproducibility.
paper_summary: paper_summary: This paper proposes a new optimization scheme for deep learning.  The idea is to periodically estimate the Hessian diagonal and use a larger learning rate on the smaller-curvature parameters.   The paper argues that this algorithm enhances the implicit curvature regularization of the optimization algorithm while not degrading stability (which is what would happen if you used a larger learning rate on all parameters).  The paper first motivates the scheme by analyzing a toy setting.  The paper then shows empirically that adding this scheme to SGD and to SAM results in better generalizing and lower-curvature solutions, where curvature is quantified using Hessian trace.   The paper then shows that adding this scheme to Adam accelerates the _optimization_ of Llama transformers, and leaves an explanation why to future work.  Finally, the paper theoretically analyzes this scheme for SAM in the manifold-of-global-minima setting, and proves that this scheme accelerates the drift towards flatter regions of the manifold.
strength_weakness: strength_weakness: One weakness is that the motivation comes from the stylized example in section 2, which may be unrealistic for general deep learning optimization.  For example, the idea of "sharp coordinates" only makes sense in this stylized example.
questions: questions: Shouldn't this algorithm also speed up _optimization_?  The algorithm uses a larger learning rate on some coordinates, which should also have an optimization effect.  (This would not be visible in the analysis setting of manifold-of-global-minima because no optimization is occurring in that setting.)

How would you define IRE at the most abstract level?  If we define an abstract version of preconditioned gradient descent as $\theta_{t+1} = \theta_t - P g_t$ (where vanilla GD with learning rate $\eta$ corresponds to $P = \eta I$), is the idea that it's better to use $P_2$ rather than $P_1$ provided that $P_2 \succeq P_1$ and both are locally stable i.e. $\lambda_{\max}(P H) \le 2$ for $P \in $ {$P_1, P_2$}?
limitations: limitations: Yes

Comments: title: Thanks for the response
reviewer response
comment: Thanks for responding to my questions; they generally address my concerns. I don't have any further questions, and I still support the work's acceptance; I'm keeping my original positive score to reflect my assessment.
Thanks, I don't have any more questions.
Thanks again for your valuable time and effort in reviewing our work!

We are wondering if our responses address your questions or concerns.

We are happy to try to address any other comments in the time remaining.
Thanks again for your valuable time and effort in reviewing our work!

We are wondering if our responses and new experiments address your questions or concerns.

We are happy to try to address any other comments in the time remaining.
Thanks again for your valuable time and effort in reviewing our work!

We are wondering if our responses and new experiments address your questions or concerns.

We are happy to try to address any other comments in the time remaining.