Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability
Abstract
Traditional analyses of gradient descent with learning rate η show that when the largest eigenvalue of the Hessian of the loss, also known as the sharpness S ( θ ) , is bounded by 2 /η , training is "stable" and the training loss decreases monotonically.
However, Cohen et al.
[7] recently observed two important phenomena.
The first, progressive sharpening , is that the sharpness steadily increases throughout training until it reaches the instability cutoff 2 /η .
The second, edge of stability , is that the sharpness hovers at 2 /η for the remainder of training while the loss non-monotonically decreases.
We demonstrate that, far from being chaotic, the dynamics of gradient descent at the edge of stability can be captured by a cubic Taylor expansion: as the iterates diverge in direction of the top eigenvector of the Hessian due to instability, the cubic term in the local Taylor expansion of the loss function causes the curvature to decrease until stability is restored.
This property, which we call self-stabilization , is a general property of gradient descent and explains its behavior at the edge of stability.
A key consequence of self-stabilization is that gradient descent at the edge of stability implicitly follows projected gradient descent (PGD) under the constraint S ( θ ) ≤ 2 /η .
Our analysis provides precise predictions for the loss, sharpness, and deviation from the PGD trajectory throughout training, which we verify both empirically in a number of standard settings and theoretically under mild conditions.
Our analysis uncovers the mechanism for gradient descent’s implicit bias towards stability.
1. Introduction 1.1. Gradient Descent at the Edge of Stability
Neural networks are often trained using variants of gradient descent such as stochastic gradient descent (SGD) or ADAM [17].
When deciding on an initial learning rate, practitioners draw intuition from the following classical lemma, known as the "descent lemma
Definition
Given a loss function L ( θ ) , the sharpness at θ is defined to be S ( θ ) :
=
λ max ( ∇ 2 L ( θ )) .
When this eigenvalue is unique, the associated eigenvector is denoted by u ( θ ) .
Lemma 2 (Descent Lemma) Assume that S ( θ ) ≤ ℓ for all θ .
If θ t +1 =
θ t − η ∇ L ( θ t ) ,
[Equation]
Here, loss decrease is maximized at η = 1 /ℓ , a popular learning rate criterion.
For any η < 2 /ℓ , the descent lemma guarantees that the loss will decrease, and hence learning rates below 2 /ℓ are considered "stable" while those above 2 /ℓ are considered "unstable."
For quadratic loss functions, this is tight.
Any learning rate above 2 /ℓ provably leads to exponentially increasing loss.
However, it has recently been observed that in neural networks, the descent lemma is not predictive of the optimization dynamics.
Recently, Cohen et al.
[7] observed two interesting phenomena:
Progressive
Sharpening
Throughout most of the optimization trajectory, the gradient of the loss is negatively aligned with the gradient of sharpness, i.e. ∇ L ( θ ) · ∇ S ( θ ) < 0 .
As a result, for any reasonable learning rate η , the sharpness increases throughout training until it reaches S ( θ ) = 2 /η .
Edge of Stability
Once the sharpness reaches 2 /η , it ceases to increase and remains around 2 /η for the rest of training.
Despite the fact that the descent lemma no longer guarantees the loss decreases, the loss still continues to rapidly decrease, albeit non-monotonically.
[Figure]
1.2. Self-stabilization: The Implicit Bias of Instability
In this work we explain the second stage, "edge of stability.
" We identify a new implicit bias of gradient descent which we call self-stabilization .
Self-stabilization is the mechanism by which the sharpness remains bounded around 2 /η , despite the continued force of progressive sharpening, and by which the gradient descent dynamics do not diverge, despite instability.
Traditional non-convex optimization analyses involve Taylor expanding the loss function to second order around θ to prove loss decrease when η ≤ 2 /S ( θ ) .
When this is violated, the iterates diverge exponentially in the top eigenvector direction, u , thus leaving the region in which the loss is locally quadratic.
Understanding the dynamics thus necessitates a cubic Taylor expansion.
Our key insight is that the missing term in the Taylor expansion of the gradient after diverging in the u direction is ∇ 3 L ( θ )( u, u ) , which is exactly equal to the gradient of the sharpness at θ :
Lemma 3 (Self-Stabilization Property)
If the top eigenvalue of ∇ 2 L ( θ ) is unique, then the sharpness S ( θ ) is differentiable at θ and ∇ S ( θ )
= ∇ 3 L ( θ )( u ( θ ) , u ( θ )) .
As the iterates move in the negative gradient direction, this term decreases the sharpness .
So long as the iterates are unstable, the strength of this term grows until the sharpness goes below 2 /η , at which point the iterates in the u direction shrink and the dynamics re-enter the quadratic regime.
This negative feedback loop prevents both the sharpness S ( θ ) and the movement in the top eigenvector direction, u , from growing out of control.
As a consequence, we show that gradient descent implicitly solves the constrained minimization problem : min θ L ( θ ) such that S ( θ ) ≤ 2 /η .
Specifically, if the stable set M is defined by M :=
{ θ : S ( θ ) ≤ 2 /η and ∇ L ( θ ) · u ( θ ) = 0 }
then the gradient descent trajectory { θ t } tracks the following projected gradient descent trajectory { θ † t } which solves the constrained problem [3]:
[Equation]
where
[Equation]
x = u · ( θ − θ ? )
C h a ng e i n s h a r pn e ss y = ∇
S · ( θ − θ ? )
θ ?
Stage 1: Progressive Sharpening x = u · ( θ − θ ? )
θ ?
Stage 2: Blowup x = u · ( θ − θ ? )
− ∇ 3 L ( u , u )=
− ∇ S
θ ? θ ?
Stage 3: Self-Stabilization x = u · ( θ − θ ? )
θ ? θ ?
∇ S
[Figure]
[Figure]
to Stability
Our main contribution is that we explain self-stabilization as a generic property of gradient descent for a large class of loss functions, and provide precise predictions for the loss, sharpness, and deviation from the constrained trajectory { θ † t } throughout training (Section 3).
We verify our predictions by replicating the experiments in Cohen et al.
[7] and show that they model the true gradient descent dynamics (Section 4).
In Appendix D, we prove that under mild conditions (which we verify empirically), our predictions track the true gradient descent dynamics up to higher order error terms.
2. Setup
We denote the loss function by L ∈ C 3 ( R d ) .
Let θ ∈ R
d
follow gradient descent with learning rate η , i.e. θ t +1 :
= θ t − η ∇ L ( θ t ) .
Recall that M :=
{ θ : S ( θ ) ≤ 2 /η and ∇ L ( θ ) · u ( θ ) = 0 } is the set of stable points and proj M :=
arg min θ ′ ∈M
∥
θ − θ ′ ∥ is the orthogonal projection onto M .
For notational simplicity, we will shift time so that θ 0 is the first point such that S (proj M ( θ ))
= 2 /η .
The constrained trajectory θ † is initialized with θ † 0 :=
proj M ( θ 0 ) after which it follows (1).
Our key assumption is the existence of progressive sharpening along the constrained trajectory, which is captured by α ( θ ) .
We also assume that there is a single unstable eigenvalue.
Assumption 1 (Progressive Sharpening) Let α ( θ ) :
= −∇ L ( θ ) · ∇ S ( θ ) .
Then α ( θ † t ) > 0 .
Assumption 2 (Eigengap) For some absolute constant c < 2 we have λ 2 ( ∇ 2 L ( θ † t ))
<
c/η .
3. The Self-stabilization Property of Gradient Descent 3.1. The Four Stages of Edge of Stability: A Heuristic Derivation
In this section we present a heuristic derivation of the self-stabilization property (see Appendix D for the more general analysis).
The derivation proceeds by a cubic Taylor expansion around a fixed reference point θ ⋆ :
= θ † 0 .
For notational simplicity, we will define the following quantities at θ ⋆ :
[Equation]
where α = −∇ L · ∇ S > 0 is the progressive sharpening.
Our heuristic analysis assumes that ∇ S ⊥ u and ∇ L, ∇ S ∈ ker( ∇ 2 L ) , and ignores higher order error terms.
We track the movement in the unstable direction u and the direction of changing sharpness ∇ S , and thus define x t := u · ( θ t − θ ⋆ ) and y t := ∇ S · ( θ t − θ ⋆ ) .
Note that y t is approximately to the change in sharpness from θ ⋆ to θ t , since Taylor expanding the sharpness yields S ( θ t )
≈
S ( θ ⋆ ) + ∇ S · ( θ t − θ ⋆ ) =
2 /η + y t .
At a high level, the mechanism for edge of stability can be described in 4 stages (see Figure 1):
Stage 1: Progressive Sharpening
While x, y are small, ∇ L ( θ t )
≈ ∇ L .
In addition, because ∇ L · ∇ S < 0 , gradient descent naturally increases the sharpness at every step.
In particular,
[Equation]
The sharpness therefore increases linearly with rate ηα .
Stage 2: Blowup
As x t measures the deviation from θ ⋆ in the u direction, the dynamics of x t can be modeled by gradient descent on a quadratic with sharpness S ( θ t )
≈ 2 /η + y t .
In particular, the rule for gradient descent on a quadratic gives
[Equation]
When the sharpness exceeds 2 /η , i.e. when y t > 0 , | x t | begins to grow exponentially.
Stage 3: Self-Stabilization
Once the movement in the u direction is sufficiently large, the loss is no longer locally quadratic.
Understanding the dynamics necessitates a third order Taylor expansion.
The missing cubic term in the Taylor expansion of ∇ L ( θ t ) is ∇ 3 L ( u, u ) x 2 t 2 = ∇ S x 2 t 2 by
Lemma 3.
This biases the optimization trajectory in the −∇ S direction, which decreases sharpness.
Recalling β
= ∥∇ S ∥ 2 , the new update for y becomes:
[Equation]
Therefore once x t > (cid:112) [Equation] α/β , the sharpness begins to decrease and continues to do so until the sharpness goes below 2 /η and the dynamics return to stability.
Stage : Return to Stability
At this point | x t | is still large from stages 1 and 2.
However, the selfstabilization of stage 3 eventually drives the sharpness below 2 /η so that y t < 0 .
Because the rule for gradient descent on a quadratic with sharpness S ( θ t ) = 2 /η
+ y t < 2 /η is x t +1 ≈ − (1+ ηy t ) x t , | x t | begins to shrink exponentially and the process returns to stage 1.
Combining the update for x t , y t in all four stages, we obtain the following simplified dynamics:
[Equation]
and
[Equation]
3.2. Analyzing the simplified dynamics
We now analyze the dynamics in (2).
First, note that x t changes sign at every iteration, and that x t +1 ≈ − x t .
While (2) cannot be modeled by an ODE due to these oscillations, we can instead model | x t | , y t .
We couple the dynamics of | x t | , y t to the following ODE X ( t ) , Y ( t ) : [Equation] and
[Equation]
This system has the unique fixed point ( X, Y ) =
( δ, 0) where δ :=
(cid:112) 2 α/β .
Furthermore, we can prove conservation of a quantity g ,which allows us to explicitly bound the size of the trajectory:
Lemma 4 Let h ( z ) :
=
z − log z − 1 .
Then g ( X ( t ) , Y ( t )) :
= h (cid:16) βX ( t ) 2 2 α (cid:17)
+ Y
[Equation]
[Equation]
is conserved.
Proof
[Equation]
Corollary
For all t , X (0) ≤ X ( t ) ≲ δ (cid:112) log( δ/X (0)) and | Y ( t ) | ≲ (cid:112) α log( δ/X (0)) .
The fluctuations are ˜ O ( √ α ) in the sharpness and ˜ O
( δ ) in the unstable direction.
Moreover, the normalized displacement in the ∇ S direction, is also bounded by ˜ O
( δ ) , so the entire process remains bounded by ˜ O ( δ ) .
Note that the fluctuations increase as α grows, and decrease as β grows.
S
Relationship with the constrained trajectory θ † t (2) determines the displacement θ t − θ ⋆ in the u, ∇ S directions.
However, θ t still evolves in the orthogonal directions by − ηP ⊥ u, ∇ S ∇ L at every step.
This can be interpreted as taking a gradient step of − η ∇ L and then projecting out the ∇ S direction to ensure the sharpness does not change.
Lemma 16, shows that this is precisely the update for θ † t ((1)) up to higher order terms.
Therefore ∥ θ t − θ † t ∥ ≤ ˜ O
( δ ) and that this ˜ O ( δ ) error term is controlled by the self-stabilizing dynamics in (2).
4. Experiments
Appendix D introduces a generalization of eq.
In Figure 2, we replicate the experiments in [7].
We show that these dynamics accurately model gradient descent at the edge of stability and can predict the loss, sharpness, and distance from constrained trajectory.
In addition, while gradient flow diverges linearly from gradient descent, the gradient descent and constrained trajectories remain close throughout training .
See Appendix F for details and Appendix I for more experiments.
5. Discussion Non-Monotonic Loss Decrease
Cohen et al.
[7] observed that, despite non-monotonic fluctuations of the loss, the loss still decreases over long time scales.
Our analysis explains this decrease by showing that gradient descent remains close to the constrained trajectory.
Since this trajectory is stable , it satisfies a descent lemma (Lemma 17), and has monotonically decreasing loss.
Over short time periods, the loss is dominated by the rapid fluctuations of x t described in Section 3.
Over longer time periods, the loss decrease of the constrained trajectory overpowers the bounded fluctuations of x t , leading to an overall loss decrease.
This is reflected in our experiments in Section 4.
Generalization & the Role of Large Learning Rates
Prior work shows that decreasing sharpness of the learned solution [9, 15, 16, 24] and increasing the learning rate [18, 19, 26] are correlated with better generalization.
Our analysis shows that gradient descent implicitly constrains the sharpness to stay near 2 /η , which suggests larger learning improves generalization by reducing the sharpness.
In Figure 3 we confirm that gradient descent generalizes better with large learning rates.
Training Speed
[7] also shows that larger learning rates lead to faster convergence despite instability.
This phenomenon is explained by our analysis.
Gradient descent is coupled to the constrained trajectory which minimizes the loss while constraining movement in the u t , ∇ S ⊥ t directions.
However, the constrained trajectory can still move quickly in the orthogonal directions, using the large learning rate to accelerate convergence.
We demonstrate this empirically in Figure 3.
We defer additional discussion of our work, including a generalization of the predicted dynamics in Appendix H, and the effect of multiple unstable eigenvalues and connections to Sharpness Aware
Minimization [10], warm-up [11], and scale-invariant loss functions [22] to Appendix J.
Future Work
An important question is to understand the dynamics when there are multiple unstable eigenvalues, which we address in Appendix J. Another direction is to understand the global convergence properties at the edge of stability, including convergence to a KKT point of the constrained update (1).
Next, our analysis focused on the edge of stability dynamics but left open the question of why progressive sharpening occurs.
Finally, we would like to understand how selfstabilization interacts with the implicit biases of stochastic -gradient descent [4, 8, 20].
[Figure]
MLP+MSE on CIFAR10, η = 0 .
[Figure]
CNN+MSE on CIFAR10, η = 0 .
[Figure]
Transformer+MSE on SST2, η = 0 .
C A Toy Model for Self-Stabilization
D The Predicted Dynamics and Theoretical Results
D.1 Notation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.
The equations governing edge of stability . . . . . . . . . . . . . . . . . . . . . .
D.3 Coupling Theorem . .
E Definition of the Predicted Dynamics
F Experimental Details
F.
Architectures . .
F.
Data . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Experimental Setup . .
G Empirical Verification of the Assumptions
H
The
Generalized Predicted Dynamics
H.1 Deriving the Generalized Predicted Dynamics . . . . . . . . . . . . . . . . . . . .
H.2 Properties of the Generalized Predicted Dynamics . . . . . . . . . . . . . . . . . .
I Additional Experiments
I.
The Benefit of Large Learning Rates: Training Time and Generalization . . . . . .
I
Experiments with the Generalized Predicted Dynamics . . . . . . . . . . . . . . .
J Additional Discussion
K Proofs
K.1 Properties of the Constrained Trajectory . . . . . . . . . . . . . . . . . . . . . . .
K.2 Proof of Theorem 11 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
K.3 Proof of Auxiliary Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Appendix A. Related Work
Cohen et al.
[7] conducted an extensive empirical study showing progressive sharpening and edge of stability in a wide range of models.
Prior work [28, 29] had also observed that for neural networks full-batch gradient descent reaches instability and the loss is not monotonically decreasing.
Lewkowycz et al.
[18] observed that when the initial sharpness is larger than 2 /η , gradient descent "catapults" into a stable region and converges.
Recent works have sought to provide a theoretical analysis for the edge of stability phenomenon.
Ma et al.
[23] analyzes edge of stability when the loss satisfies a "subquadratic growth" assumption.
Ahn et al.
[1] argues that unstable convergence is possible when there exists a "forward invariant subset" near the set of minimizers.
Arora et al.
[2] analyzes progressive sharpening and the edge of stability phenomenon for normalized gradient descent close to the manifold of global minimizers.
Lyu et al.
[22] uses the edge of stability phenomenon to analyze the effect of normalization layers on sharpness for scale-invariant loss functions.
Chen and Bruna [6] show global convergence despite instability for certain 2D toy problems and in a 1-neuron student-teacher setting.
The concurrent work Li et al.
[21] proves progressive sharpening for a two-layer network and analyzes the edge of stability dynamics through four stages similar to ours using the norm of the output layer as a proxy for sharpness.
Beyond the edge of stability phenomenon itself, prior work has also shown that SGD with large step size or small batch size will lead to a decrease in sharpness [12–14, 16].
Gilmer et al.
[11] also describes connections between edge of stability, learning rate warm-up, and gradient clipping.
At a high level, our proof relies on the idea that oscillations in an unstable direction prescribed by the quadratic approximation of the loss cause a longer term effect arising from the third-order
Taylor expansion of the dynamics.
This overall idea has also been used to analyze the implicit regularization of SGD [4, 8, 20].
In those settings, oscillations come from the stochasticity, while in our setting the oscillations stem from instability.
Appendix B. Notation
We denote by ∇ k L ( θ ) the k -th order derivative of the loss L at θ .
Note that ∇ k L ( θ ) is a symmetric k -tensor in ( R d ) ⊗ k when θ ∈ R d .
For a symmetric k -tensor T , and vectors u 1 , . . .
u
j ∈ R
d we will use T ( u 1 , . . .
, u j ) to denote the tensor contraction of T with u 1 , . .
. , u j , i.e.
[Equation]
We use P u 1 ,...,u k to denote the orthogonal projection onto span( u 1 , . .
. , u k ) and P ⊥ u 1 ,...,u k is the projection onto the corresponding orthogonal complement.
For matrices A 1 , . . .
A k , we define t (cid:89) k
=
A k :=
A 1 . . .
A t and 1 (cid:89) k
= t
A k :=
A t . . .
A 1 .
Appendix C. A Toy Model for Self-Stabilization
For α, β > 0 , consider the function
[Equation]
initialized at the point ( x 0 , 0 , 0) .
Note that the constrained trajectory will follow x † t = 0 , y † t = 0 , z † t = − ηt as it cannot decrease y without increasing the sharpness past 2 /η .
We therefore have:
[Equation]
Note that this satisfies all of the assumptions in Section 3 and it satisfies α
= −∇
L t · ∇ S t = 0 and β = ∥∇ S t ∥ 2 .
This process will then follow (3) in the x, y directions while it tracks the constrained trajectory θ † t moving linearly in the − P ⊥ u, ∇ S ∇ L = [0 , 0 , − 1] direction.
Appendix D. The Predicted Dynamics and Theoretical Results
We now present the equations governing edge of stability for general loss functions.
D.1. Notation
Our general approach Taylor expands the gradient of each iterate θ t around the corresponding iterate θ † t of the constrained trajectory.
We define the following Taylor expansion quantities at θ † t :
Definition 6 (Taylor Expansion Quantities at θ † t )
[Equation]
Furthermore, for any vector-valued function v ( θ ) , we define v ⊥ t :=
P ⊥ u t v ( θ † t ) where P ⊥ u t is the projection onto the orthogonal complement of u t .
We also define the following quantities which govern the dynamics near θ ⋆ t .
Definition 7 Let α t := −∇ L t · ∇ S t , β t :=
(cid:13)(cid:13) ∇ S ⊥ t (cid:13)(cid:13) 2 , and δ t :=
(cid:113)
α t β t .
Furthermore, we define
[Equation]
Recall that α t is the progressive sharpening force, β t is the strength of the stabilization force, and δ t controls the size of the deviations from θ † t and was the fixed point in the x direction in Section 3.2.
The scalars β s → t capture the effect of the interactions between ∇ S and the Hessian.
D.2. The equations governing edge of stability
We now introduce the equations governing edge of stability.
We track the following quantities:
Definition 8 Define
[Equation] [Equation]
, y
[Equation]
S
[Equation]
· v
[Equation]
Our predicted dynamics directly predict the displacement v t and the full definition is deferred to
Appendix E. However, they have a relatively simple form in the u t , ∇ S ⊥ t directions:
Lemma 9 (Predicted Dynamics for x, y ) Let ⋆ v t denote our predicted dynamics (defined in Appendix E).
Letting ⋆ x t = u t · ⋆ v t and ⋆ y t = ∇ S ⊥ t · ⋆ v t , we have
[Equation]
and
[Equation]
Note that when β s → t are constant, our update reduces to the simple case discussed in Section 3, which we analyze fully.
When x t is large, (4) demonstrates that there is a self-stabilization force which acts to decrease y t ; however, unlike in Section 3, the strength of this force changes with t .
D.3. Coupling Theorem
We now show that, under a mild set of assumptions which we verify to hold empirically in Appendix G, the true dynamics are accurately governed by the predicted dynamics.
This lets us use the predicted dynamics to predict the loss, sharpness, and the distance to the constrained trajectory θ † t .
Our errors depend on the unitless quantity ϵ , which we verify is small in Appendix G.
Definition 10 Let ϵ t := η √ α t and ϵ :=
sup t ϵ t .
To control Taylor expansion errors, we require upper bounds on ∇ 3 L and its Lipschitz constant:
Assumption 3 Let ρ 3 , ρ 4 to be the minimum constants such that for all θ , (cid:13)(cid:13) ∇ 3 L ( θ ) (cid:13)(cid:13) op ≤ ρ 3 and ∇ 3 L is ρ 4 -Lipschitz with respect to ∥·∥ op .
Then we assume that ρ 4 = O ( ηρ 23 ) .
Next, we require the following generalization of Assumption 1:
Assumption 4 For all t ,
[Equation]
Finally, we require a set of “non-worst-case" assumptions, which are that the quantities ∇ 2 L, ∇ 3 L, and λ min ( ∇ 2 L ) are nicely behaved in the directions orthogonal to u t , which generalizes the eigengap assumption.
We verify the assumptions on ∇ 2 L and ∇ 3 L empirically in Appendix G.
Assumption 5 For all t and v, w ⊥ u t ,
[Equation]
With these assumptions in place, we can state our main theorem which guarantees ⋆ x, ⋆ y, ⋆ v predict the loss, sharpness, and deviation from the constrained trajectory up to higher order terms:
Theorem 11 Let T :
= O ( ϵ − 1 ) and assume that min t ≤ T | ⋆ x t | ≥ c 1 δ.
Then for any t ≤ T , we have
[Equation]
S
The sharpness is controlled by the slowly evolving quantity ⋆ y t and the period-2 oscillations of ( ∇ S · U ) ⋆ x t .
This combination of gradual and rapid periodic behavior was observed by Cohen et al.
[7] and appears in our experiments.
Theorem 11 also shows that the loss at θ t spikes whenever ⋆ x t is large.
On the other hand, when ⋆ x t is small, L ( θ t ) approaches the loss of the constrained trajectory.
Appendix E. Definition of the Predicted Dynamics
Below, we present the full definition of the predicted dynamics:
Definition 12 (Predicted Dynamics, full) Define ⋆ v 0 = v 0 , and let ⋆ x t =
⋆
[Equation]
· ⋆ v t .
Then
[Equation]
For convenience, we will define the map step t : R d → R d as follows:
Definition
Given a vector v and a timestep t , define step t ( v ) by
[Equation]
where x = u t · v and y = ∇ S ⊥ t · v .
It is easy to see that ⋆ v t +1 =
step t ( ⋆ v t ) .
Proof [Proof of Theorem 9] Defining A t =
( I − η ∇ 2 L t ) P ⊥ u t , we can unfold the recursion in (5) to obtain the following formula for ⋆ v t .
[Equation]
It is then immediate to see that ⋆ x t =
⋆ v t · u t , ⋆ y t = ∇ S ⊥ t · ⋆ v t have the following simple update:
[Equation]
and
[Equation]
where we recall that we have defined
[Equation]
Appendix F. Experimental Details F.1. Architectures
We evaluated our theory on four different architectures.
The 3-layer MLP and CNN are exact copies of the MLP and CNN used in [7].
The MLP has width 200 , the CNN has width 32 , and both are using the swish activation [25].
We also evaluate on a ResNet18 with progressive widths 16 , 32 , 64 , 128 and on a 2-layer Transformer with hidden dimension 64 and two attention heads.
S
F.2. Data
We evaluated our theory on three primary tasks: CIFAR10 multi-class classification with both categorical MSE loss and cross-entropy loss, CIFAR10 binary classification (cats vs dogs) with binary
MSE loss and logistic loss, and SST2 [27] with binary MSE loss and logistic loss.
F.3. Experimental Setup
For every experiment, we tracked the gradient descent dynamics until they reached instability and then began tracking the constrained trajectory, gradient descent, gradient flow, and both our predicted dynamics (Appendix D) and our generalized predicted dynamics (Appendix H).
In addition, we tracked the various quantities on which we made assumptions for Appendix D in order to validate these assumptions.
We also tracked the second eigenvalue of the Hessian at the constrained trajectory throughout training and stopped training once it reached 1 .
9 /η , to ensure the existence of a single unstable eigenvalue.
Finally, as the edge of stability dynamics are very sensitive to small perturbation when | x | is small (see ?? )
, we switched to computing gradients with 64-bit precision after first reaching instability to avoid propagating floating point errors.
Eigenvalues were computed using the LOBPCG sparse eigenvalue solver in JAX [5].
To compute the constrained trajectory, we computed a linearized approximation for proj M inspired by
Lemma 16 along with a Newton step in the u t direction to ensure that ∇ L · u = 0 .
Each linearized approximation step required recomputing the sharpness and top eigenvector and each projection step then consisted of three linearized projection steps, for a total of three eigenvalue computations per projection step.
Our experiments were conducted in JAX [5], using https://github.com/locuslab/ edge-of-stability as a reference for replicating the experimental setup used in [7].
All experiments were conducted on two servers, each with 10 NVIDIA GPUs.
Code is provided in the supplementary material and will be made public on GitHub upon acceptance.
Appendix G. Empirical Verification of the Assumptions
For each of the experimental settings considered (MLP+MSE, CNN+MSE, CNN+Logistic, ResNet18+MSE,
Transformer+MSE, Transformer+Logistic), we plot a number of quantities along the constrained trajectory to verify that the assumptions made in the main text hold.
For each learning rate η we have 8 plots tracking various quantities, which verify the assumptions as follows: Assumption 1 is verified by the 1st plot, ϵ being small is verified by the 2nd plot, Assumption 4 is verified by the 3rd and 4th plots, Assumption 3 is verified by the 5th plot, and Assumption 5 is verified by the last 3 plots.
As described in the experimental setup, training is stopped once the second eigenvalue is 1 .
9 /η , so Assumption 2 always holds with c = 1 . 9 as well.
MLP+MSE on CIFAR
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
CNN+MSE on CIFAR
[Figure]
α t : = − ∇ L t · ∇ S t
[Figure]
ε
[Figure]
η √ α
[Figure]
α t : = − ∇ L t · ∇ S t
[Figure]
ε
[Figure]
= η √ α
[Figure]
: = − ∇ L
[Figure]
· ∇ S t
[Figure]
CNN+Logistic on CIFAR10 (cats vs dogs)
[Figure]
α
[Equation]
t
[Figure]
[Equation]
t :
[Equation]
t
[Equation]
[Figure]
ε t : =
η √ α
[Figure]
α
[Figure]
− ∇ L
[Figure]
· ∇ S t
[Figure]
− ∇ L t · ∇ S
[Figure]
1 .∇
[Equation]
L t ( ∇ L t , ∇ L t ) k ∇ 3 L t k op k ∇ L t k  2000 0 .0 .
.∇
2 L t ( v ⊥ t , v ⊥ t )
k ∇ 2 L t kk v ⊥ t kη = 0 .
ResNet18+MSE on CIFAR10 (cats vs dogs)
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
Transformer+MSE on SST
[Figure]
α
[Figure]
: = − ∇ L t · ∇ S t
[Figure]
ε
[Figure]
η √ α
[Figure]
α t : = − ∇ L t · ∇ S t
[Figure]
ε t : =
η √ α t
[Figure]
− ∇
[Figure]
∇
[Figure]
∇
[Figure]
⊥
[Figure]
Transformer+Logistic on SST
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
α t : = − ∇ L t · ∇ S t
[Figure]
− ∇ L t · ∇ S t
[Figure]
[Equation]
t
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
k
[Figure]
[Equation]
[Figure] [Figure]
[Equation]
[Figure]
Appendix H. The Generalized Predicted Dynamics
Our analysis relies on a cubic Taylor expansion of the gradient.
However, in order for this Taylor expansion to accurately track the gradients we need a bound on the fourth derivative of the loss (Assumption 3).
Section 4 and Appendix G show that this approximation is sufficient to capture the dynamics of gradient descent at the edge of stability for many standard models when the loss criterion is the mean squared error.
However, for certain architectures and loss functions, including
ResNet18 and models trained with the logistic loss, this condition is often violated.
In these situations, the loss function in the top eigenvector direction is either sub-quadratic , meaning that the quadratic Taylor expansion overestimates the loss and sharpness 3 , or super-quadratic , meaning that the quadratic Taylor expansion underestimates the loss and sharpness.
To capture this phenomenon, we derive a more general form of the predicted dynamics which reduces to the standard predicted dynamics in Appendix D when the loss in the top eigenvector direction is approximately quadratic.
In addition, Appendix I shows that the generalized predicted dynamics capture the dynamics of gradient descent at the edge of stability for both mean squared error and cross-entropy in all settings we tested.
H.1. Deriving the Generalized Predicted Dynamics
To derive the generalized predicted dynamics, we will abstract away the dynamics in the top eigenvector direction.
Specifically, for every t we define
[Equation]
We say that L is sub-quadratic at t if F t ( x ) < 0 and super-quadratic if F t ( x ) > 0 .
Note that knowing F t is not sufficient to capture the dynamics in the u t direction.
Specifically,
[Equation]
It is still critically important to track the effect that the movement in the ∇ S ⊥ t direction has on the dynamics of x .
As in Section 3.1, the effect of the movement in the ∇ S ⊥ t direction on the dynamics of x is changing the sharpness by y t .
This gives us the generalized predicted dynamics update:
[Equation]
where x
[Equation]
Note that when F t ( x ) = 0 is exactly quadratic, this reduces to the standard predicted dynamics update in (5).
Note that the update for y is completely unchanged:
Lemma 14 Restricted to the u t , ∇ S t directions, the generalized predicted dynamics v ⋆t imply:
[Equation]
and
[Equation]
The proof is identical to the proof of Theorem 9.
H.2. Properties of the Generalized Predicted Dynamics
Note that due to the sign flipping argument in Appendix K, we can assume that F is an even function as the odd part will only influence the dynamics through additional oscillations of period 2, so throughout the remainder of this section we will assume that F t ( x ) =
F t ( − x ) .
Otherwise, we can simply redefine F by its even part.
Next, note that the fixed point of (10) is still when x t = δ t , regardless of the shape of F t , due to the need to stabilize the ∇ S ⊥ t direction.
This contradicts previous 1-dimensional analyses of edge of stability in which the fixed point in the top eigenvector direction strongly depends on the shape of F t , the loss in the u t direction.
The limiting value of y t can therefore be read from the update for x t .
If ( δ t , y ) is an orbit of period 2 of (10), then
[Equation]
In addition, note that the sharpness can no longer be approximated as S ( θ t )
≈ 2 /η + y t as the sharpness now changes along the u t direction.
In particular, it changes by F ′′ ( x ) so that
[Equation]
Therefore, the limiting sharpness of (10) is
[Equation]
When F t = 0 and the loss is exactly quadratic in the u direction, this update reduces to fixed point predictions in Section 3.1.
One interesting phenomenon observed by Cohen et al.
[7] is that with cross-entropy loss, the sharpness was never exactly 2 /η , but usually hovered above it.
This contradicts the predictions of the standard predicted dynamics which predict that the fixed point has sharpness 0 .
However, using the generalized predicted dynamics (10), we can give a clear explanation.
When the loss is sub-quadratic, e.g. when F t ( x ) = − ρ 4 x 4 24 , we have
[Equation]
so the sharpness will converge to a value below 2 /η .
On the other hand if the loss is super-quadratic, the sharpness converges to a value above 2 /η .
More generally, whether the loss converges to a value above or below 2 /η depends on the sign of F ′′ t ( δ t ) − δ t F ′ t ( δ t ) .
In our experiments in Appendix I, we observed both sub-quadratic and super-quadratic loss functions.
In particular, the loss was usually sub-quadratic when it first reached instability but gradually became super-quadratic as training progressed at the edge of stability.
Appendix I. Additional Experiments I.1. The Benefit of Large Learning Rates: Training Time and Generalization
We trained ResNet18 with full batch gradient descent on the full 50k training set of CIFAR10 with various learning rates, in addition to the commonly proposed learning rate schedule η t := 1 /S ( θ t ) .
We show that despite entering the edge of stability, large learning rates converge much faster.
In addition, due to the self-stabilization effect of gradient descent, the final sharpness is bounded by 2 /η which is smaller for larger learning rates and leads to better generalization (see Figure 3).
Step
T r a i n i ng
L o ss
Step
T r a i n i ng
A cc u r ac y
Step
T e s t
A cc u r ac y η
=
.η
=
0 .η
= 0 .η = 0 .η
= 0 .η
t = 1 / S ( θ t )
MLP+MSE on CIFAR
Gradient Descent
Predicted Dynamics
Constrained Trajectory
Gradient Flow
[Figure]
t o
[Figure]
t a n ce t o θ † η = 0 .
CNN+MSE on CIFAR
Gradient Descent
[Figure]
CNN+Logistic on CIFAR10 (
cats vs dogs
[Figure]
pn
[Figure]
ss
[Figure]
pn
[Figure]
D i s t a n ce t o θ † η = 0 .
ResNet18+MSE on CIFAR10 (cats vs dogs
Gradient Descent
[Figure]
Transformer+MSE on SST
Gradient Descent
Predicted Dynamics
[Figure]
Transformer+Logistic on SST Gradient Descent
[Figure]
η = 0 .
[Figure]
Appendix J. Additional Discussion Multiple Unstable Eigenvalues
Our work focuses on explaining edge of stability in the presence of a single unstable eigenvalue (Assumption 2).
However, Cohen et al.
[7] observed that progressive sharpening appears to apply to all eigenvalues, even after the largest eigenvalue has become unstable.
As a result, all of the top eigenvalues will successively enter edge of stability (see Figure 4).
In particular, Figure 4 shows that the dynamics are fairly well behaved in the period when only a single eigenvalue is unstable, yet appear to be significantly more chaotic when multiple eigenvalues are unstable.
[Figure]
One technical challenge with dealing with multiple eigenvalues is that, when the top eigenvalue is not unique, the sharpness is no longer differentiable and it is unclear how to generalize our analysis.
However, one might expect that gradient descent can still be coupled to projected gradient descent under the non-differentiable constraint S ( θ † T ) ≤ 2 /η .
When there are k unstable eigenvalues, with corresponding eigenvectors u 1 t , . .
, u kt , the constrained update is roughly equivalent to projecting out the subspace span {∇ 3 L t ( u it , u jt ) :
i, j ∈ [ k ] } from the gradient update − η ∇ L t .
Demonstrating self-stabilization thus requires analyzing the dynamics in the subspace span (cid:16) {
[Equation] [Equation]
j ∈ [ k ] } (cid:17) .
We leave investigating the dynamics of multiple unstable eigenvalues for future work.
Connection to Sharpness Aware Minimization (SAM)
Foret et al.
[10] introduced the sharpnessaware minimization (SAM) algorithm, which aims to control sharpness by solving the optimization problem min θ max ∥ δ ∥≤ ϵ L ( θ + δ ) .
This is roughly equivalent to minimizing S ( θ ) over all global minimizers, and thus SAM tries to explicitly minimize the sharpness.
Our analysis shows that gradient descent implicitly minimizes the sharpness, and for a fixed η looks to minimize L ( θ ) subject to S ( θ ) = 2 /η .
Connections to Warmup.
Gilmer et al.
[11] demonstrated that learning rate warmup , which consists of gradually increasing the learning rate, empirically leads to being able to train with a larger learning rate.
The self-stabilization property of gradient descent provides a plausible explanation for this phenomenon.
If too large of an initial learning rate η 0 is chosen (so that S ( θ 0 ) is much greater than 2 /η 0 ), then the iterates may diverge before self stabilization can decrease the sharpness
S to 2 /η 0 .
On the other hand, if the learning rate is chosen that S ( θ 0 ) is only slightly greater than 2 /η 0 , self-stabilization will decrease the sharpness to 2 /η 0 .
Repeatedly increasing the learning rate slightly could then lead to small decreases in sharpness without the iterates diverging, thus allowing training to proceed with a large learning rate.
Connection to Weight Decay and
Sharpness Reduction
Lyu et al.
[22] proved that when the loss function is scale-invariant, gradient descent with weight decay and sufficiently small learning rate converges leads to reduction of the normalized sharpness S ( θ/ ∥ θ ∥ ) .
In fact, the mechanism behind the sharpness reduction is exactly the self-stabilization force described in this paper restricted to the setting in [22].
We present here a heuristic derivation of this equivalence.
Our primary result is that gradient descent solves the constrained problem min θ L ( θ ) such that
S ( θ ) ≤ 2 /η .
To prove equivalence to the sharpness reduction, we will need the following lemma from [22] which follows from the scale invariance of the loss:
[Equation]
Let L λ ( θ ) :
= L ( θ ) + λ 2 ∥ θ ∥ 2 and S λ ( θ ) =
S ( θ ) + λ denote the regularized loss and sharpness respectively and let θ := θ ∥ θ ∥ .
Then we have the following equality between minimization problems: min θ L λ ( θ ) such that
[Equation]
such that
[Equation]
such that
[Equation]
where the last line follows from the scale-invariance of the loss function.
In particular if ηλ is sufficiently small and the dynamics are initialized near a global minimizer of the loss, this will converge to the solution of the constrained problem:
[Equation]
such that
L ( θ ) = 0 .
Appendix K. Proofs K.1. Properties of the Constrained Trajectory
We next prove several nice properties of the constrained trajectory.
Before, we require the following auxiliary lemma, which shows that several quantities are Lipschitz in a neighborhood around the constrained trajectory:
Lemma 15 (Lipschitz Properties)
θ → ∇ L ( θ ) is O ( η − 1 ) -Lipschitz
in each set S t . 2.
θ → ∇ 2 L ( θ ) is ρ 3 -Lipschitz with respect to ∥·∥ 2 . 3.
θ → λ i ( ∇ 2 L ( θ )) is ρ 3 -Lipschitz.
S
θ → u ( θ ) is O ( ηρ 3 ) -Lipschitz in each set S t .
θ → ∇ S ( θ ) is O ( ηρ 23 ) -Lipschitz
in each set S t .
Proof The Lipschitzness of ∇ 2 L ( θ ) follows immediately from the bound (cid:13)(cid:13) ∇ 3 L ( θ ) (cid:13)(cid:13) op ≤ ρ 3 .
Weil’s inequality then immediately implies the desired bound on the Lipschitz constant of the eigenvalues of ∇ 2 L ( θ ) .
Therefore for any t , we have for all θ ∈ S t :
[Equation]
Next, from the derivative of eigenvector formula:
[Equation]
which implies the bound on the Lipschitz constant of u restricted to S t .
Finally, because ∇ S ( θ ) = ∇ 3 L ( θ )( u ( θ ) , u ( θ )) ,
[Equation]
where the second to last inequality follows from the bound on ∥∇ u ( θ ) ∥ 2 restricted to S t and the last inequality follows from Assumption 3.
Lemma 16 (First-order approximation of the constrained trajectory update { θ † t } )
For all t ≤ T ,
[Equation]
and
[Equation]
Proof We will prove by induction that S t = 2 /η for all t .
The base case follows from the definitions of θ 0 , θ † 0 .
Next, assume S ( θ † t ) = 0 for some t ≥ 0 .
Let θ ′ =
θ † t − η ∇ L t .
Then because θ † t ∈ M we have (cid:13)(cid:13)(cid:13) θ † t +1 − θ ′ (cid:13)(cid:13)(cid:13) ≤ (cid:13)(cid:13)(cid:13) θ † t − θ ′ (cid:13)(cid:13)(cid:13) =
η ∥∇ L t ∥ .
Then because θ † t +1 = proj
M ( θ ′ ) , the KKT conditions for this minimization problem imply that there exist x, y with y ≥ 0 such that
[Equation]
Next, note that we can decompose ∇ S t = u t ( ∇ S t · u t ) + ∇ S ⊥ t :
[Equation]
Let s t = ∇ S ⊥ t ∥ ∇ S ⊥ t ∥ .
We can now perform the change of variables
[Equation]
to get
[Equation]
Note that
[Equation]
for sufficiently small ϵ so because (cid:13)(cid:13)(cid:13)
[Equation]
≤ η ∥∇ L t ∥ we have
[Equation]
so x, y =
O ( η ∥∇ L t ∥ ) .
Therefore,
[Equation]
Then Taylor expanding ∇ L t +1 around θ † t gives
[Equation]
so x ′ =
O ( ϵ 2 · η ∥∇ L t ∥ ) .
We can also Taylor expand S t +1 around θ † t and use that S t = 2 /η to get
[Equation]
Now note that for ϵ sufficiently small we have
[Equation]
Therefore if y ′ = 0 , we would have S t +1 > 2 /η which contradicts θ † t +1 ∈ M .
Therefore y ′ >and therefore y > 0 , which by complementary slackness implies S t +1 = 2 /η .
This then implies that
[Equation]
Putting it all together gives
[Equation]
where the last line follows from u t · ∇ L t = 0 .
Lemma 17 (Descent Lemma for θ † )
For all t ≤ T ,
[Equation]
Proof Taylor expanding L ( θ † t +1 ) around L ( θ † t ) and using Lemma 16 gives
[Equation]
Next, note that because γ t = Θ(1) we have ∥∇ L t ∥
=
O ( (cid:13)(cid:13)
[Equation]
P ⊥ u t , ∇ S t ∇ L t (cid:13)(cid:13) (cid:13) ) .Therefore
for ϵ sufficiently small,
[Equation]
Therefore,
[Equation]
which completes the proof.
Corollary 18 Let L ⋆ =
min
θ L ( θ ) .
Then there exists t ≤ T such that
[Equation]
Proof Inductively applying Lemma 17 we have that there exists an absolute constant c such that
[Equation]
which implies that
[Equation]
K.2. Proof of Theorem
We first require the following three lemmas, whose proofs are deferred to Appendix K.3.
Lemma 19 (2-Step Lemma) Let
[Equation]
t
[Equation]
Assume that ∥ v t ∥ ≤ ϵ − 1 δ .
Then
[Equation]
Lemma 20 Assume that there exists constants
c 1 , c 2 such that for all t ≤ T , ∥
⋆ v t ∥ ≤ c 2 δ , | ⋆ x t | ≥ c 1 δ .
Then, for all t ≤ T , we have
[Equation]
Lemma
For t ≤ T , ∥
⋆ v t ∥ ≤ O ( δ ) .
With these lemmas in hand, we can prove Theorem 11.
Proof [Proof of Theorem 11]
First, by Lemma 21, we have ∥ ⋆ v t ∥
≤ O ( δ ) .
Next, by Lemma 20, we have
[Equation]
Next, we Taylor expand to calculate S ( θ t ) :
[Equation]
Finally, we Taylor expand the loss:
[Equation]
where the last line follows from Assumption
[Equation]
K.3. Proof of Auxiliary Lemmas
Proof [Proof of Lemma 19]
Taylor expanding the update for θ t +1 about θ † t , we get
[Equation]
Additionally, recall that the update for θ † t +1 is
[Equation]
Subtracting the previous 2 equations and expanding out ∇ 3 L ( v t , v t ) via the non-worst-case bounds, we obtain
[Equation]
We would first like to compute the magnitude of v t +1 .
[Equation]
Observe that by definition of ϵ and δ , and since ∥ v t ∥ ≤ ϵ − 1 δ
[Equation]
Hence
[Equation]
Note that we can bound
[Equation]
Therefore
[Equation]
the one-step update in the u t direction is:
[Equation]
where we have defined the error term E t as
[Equation]
The update in the v ⊥ direction is
[Equation]
First, observe that
[Equation]
Therefore we can control the first of the error terms as (cid:13)
[Equation]
As for the second error term
[Equation]
we can decompose
[Equation]
[Table]
By Assumption 5, we have (cid:13)(cid:13)
P ⊥ u t ∇ 3 L t ( u t , v ⊥ t ) (cid:13)(cid:13) ≤ O ( ϵρ 3 ∥ v t ∥ ) .
Additionally, (cid:13)(cid:13)(cid:13) P ⊥ u t − P ⊥ u t +1 (cid:13)(cid:13)(cid:13) ≤
O ( ∥ u t − u t +1 ∥ ) .
Therefore
[Equation]
where we used ηρ 3 ∥ v t ∥
=
O (1) .
Altogether
[Equation]
we have
[Equation]
We next compute the two-step update for x t :
[Equation]
We previously obtained ηρ 3 ∥ v t ∥ =
O (1) .
Furthermore,
[Equation]
Hence x
[Equation]
O
[Equation]
The first of these two error terms can be bounded as
[Equation]
As for the second term
[Equation]
we can bound
[Equation]
Additionally, we have
[Equation]
Therefore
[Equation]
Altogether
[Equation]
the two-step update for x t is
[Equation]
Additionally, the two-step update for v ⊥ t is
[Equation]
Define v t +1 = step t ( v t ) , v t +2 = step t +1 ( v t ) , and x i = v i · u i , y i = ∇ S ⊥ i · v i for i ∈ { t + 1 , t + 2 } .
By the definition of step , one sees that
[Equation]
and
[Equation]
The update for x after applying step is
[Equation]
Therefore
[Equation]
Additionally, the update for v ⊥ is
[Equation]
Therefore
[Equation]
Altogether, we get that
[Equation]
as desired.
Proof [Proof of Lemma 20] Define
[Equation]
and define the auxiliary trajectory (cid:98) v by (cid:98) v 0 = v 0 and (cid:98) v t +
= step( (cid:98) v t ) + w t .
I first claim that (cid:98) v t = v t for all even t ≤ T , which we will prove by induction on t .
The base case is given by assumption so assume the result for some even t ≥ 0 .
Then,
[Equation]
which completes the induction.
Next, we will prove by induction that for t ≤ T ,
[Equation]
By definition, (cid:98) v 0 = v 0 =
⋆ v 0 , so the claim is clearly true for t = 0 .
Next, assume the claim holds for t .
If t is even then ∥ w t ∥ = 0 ; otherwise ∥ v t ∥ ≤ 2 c 2 δ
[Equation]
and thus
[Equation]
First observe that
[Equation]
ϵδ ) · (cid:12)(cid:12)(cid:12)(cid:12) (cid:98) x
t − ⋆ x t ⋆ x t (cid:12)(cid:12)(cid:12)(cid:12)
+ O (cid:0) ϵ 2 δ (cid:1)
Next, note that (cid:98) x t +1 ⋆ x
[Equation]
= (
[Equation]
y
[Equation]
(1 + η
[Equation]
y
[Equation]
+
[Equation]
δ )
=
(1 + η ⋆ y
[Equation]
(cid:98) x t + O ( ϵ
[Equation]
δ
[Equation]
+ O (
[Equation]
· (cid:13)(cid:13)(cid:98) v ⊥ t − ⋆ v ⊥ t (cid:13)(cid:13) (1 + η ⋆ y t )
⋆ x t + O ( ϵ 2 δ ) =
(cid:98) x
t ⋆ x t + O (cid:16) ϵ 2 + ϵ δ (cid:13)(cid:13)(cid:13)(cid:98) v ⊥ t − ⋆ v ⊥ t (cid:13)(cid:13)(cid:13)(cid:17) .
Therefore (cid:12)(cid:12)(cid:12)(cid:12) (cid:98) x
t +1 − ⋆ x t +1 ⋆ x t + (cid:12)(cid:12)(cid:12)(cid:12) ≤ (cid:12)(cid:12)(cid:12)(cid:12) (cid:98) x
t − ⋆ x t ⋆ x t (cid:12)(cid:12)(cid:12)(cid:12)
+ O ( ϵ
+ ϵ δ (cid:13)(cid:13)(cid:13)(cid:98) v ⊥ t − ⋆ v ⊥ t (cid:13)(cid:13)(cid:13) ) .
Let d t = max (cid:16)(cid:13)(cid:13)(cid:98) v ⊥ t − ⋆ v ⊥ t (cid:13)(cid:13) , δ (cid:12)(cid:12)(cid:12) (cid:98) x
t − ⋆ x t ⋆ x t (cid:12)(cid:12)(cid:12)(cid:17) .
Then (cid:13)(cid:13)(cid:13)(cid:98) v ⊥ t +1 − ⋆ v ⊥ t +1 (cid:13)(cid:13)(cid:13) ≤ (1 + η (cid:12)(cid:12) λ min ( ∇ 2 L t ) (cid:12)(cid:12) + O ( ϵ ))
d t + O ( ϵ 2 δ ) δ (cid:12)(cid:12)(cid:12)(cid:12) (cid:98) x
t +1 − ⋆ x t +1 ⋆ x t + (cid:12)(cid:12)(cid:12)(cid:12) ≤ (1 + O ( ϵ ))
d t + O ( ϵ 2 δ ) .
Therefore d t +1 ≤ (1 + η (cid:12)(cid:12) λ min ( ∇ 2 L t ) (cid:12)(cid:12) + O ( ϵ ))
d t + O ( ϵ 2 δ ) ≤ (1 + O ( ϵ ))
d t + O ( ϵ 2 δ ) , so for t ≤ T we have d t +1 ≤ O ( ϵδ ) .
Therefore
[Equation]
+
[Equation]
, |
[Equation]
x
[Equation]
−
[Equation]
x t +
[Equation]
| ≤ O ( ϵδ
[Equation]
≤ c 2 δ, so the induction is proven.
Altogether, we get ∥ (cid:98) v t − ⋆ v t ∥ ≤ O ( ϵδ ) for all such t , as desired.
Proof [Proof of Lemma 21] Recall that
[Equation]
and y
[Equation]
t (cid:88) s
[Equation]
β
[Equation]
(cid:20)
[Equation]
Since t ≤
[Equation]
max t | λ min
[Equation]
∇ 2 L t ) | , we have that β s → t = O ( ρ 23 ) , and thus
[Equation]
Therefore
[Equation]
Since t ≤ O ( ϵ − 1 ) , | ⋆ x t | grows by at most a constant factor, and thus | ⋆ x t | ≤ O ( δ ) .
Finally, recall that
[Equation]
By the triangle inequality,
[Equation]
Therefore ∥ ⋆ v t ∥ ≤ O ( δ ) .