Maximal stable learning rate derivation

simple setup

Consider a simple quadratic loss function for a model parameterized by $\theta$, where $H$ is a symmetric positive semi-definite matrix (e.g., the Hessian):

$$L(\theta )=\frac{1}{2}{\theta }^{T}H\theta$$

Expanding this into summation notation:

$$L(\theta )=\frac{1}{2}{\theta }^T\left[\begin{array}{c}{\sum }_i{h}_{1i}{\theta }_i\\ ⋮\end{array}\right]=\frac{1}{2}\sum _j\sum _i{h}_{ji}{\theta }_i{\theta }_j$$

the gradient

Taking the partial derivative with respect to ${\theta }_k$ using the Kronecker delta $\delta$ (and using the product rule):

$$\begin{array}{rl}\frac{\partial L}{\partial {\theta }_k}& =\frac{1}{2}\sum _j\sum _i(h_{ji}{\delta }_{ik}{\theta }_j+h_{ji}{\theta }_i{\delta }_{jk})\\ & =\frac{1}{2}\left(\sum _j{h}_{jk}{\theta }_j+\sum _i{h}_{ki}{\theta }_i\right)\end{array}$$

Because $H$ is symmetric ($H=H^T$), $h_{jk}=h_{kj}$, we can simplify the expression to:

$$\frac{\partial L}{\partial {\theta }_k}=\sum _i{h}_{ki}{\theta }_i=(H\theta {)}_k$$

Thus, the full gradient vector is:

$${\nabla }_{\theta }L=H\theta$$

Taking the gradient of this expression, we can see that ${\nabla }_{\theta }^{2}L$, or the Hessian of $L$, is equal to $H$.

maximal lr under GD

Applying the standard gradient descent update rule with learning rate $\eta$:

$$\theta \to \theta -\eta H\theta$$

Since $H$ is symmetric, we can decompose $H$ using its eigendecomposition $H=U\Sigma U^T$, where $U$ contains orthogonal eigenvectors and $\Sigma$ is a diagonal matrix of eigenvalues ${\sigma }_i$. Substituting this into the update:

$$\theta \to (I-\eta U\Sigma U^T)\theta$$

Since $U{U}^T=I$, we can rewrite the identity matrix:

$$\theta \to (U{U}^T-\eta U\Sigma U^T)\theta$$

Extracting the $U$ and $U^T$ terms,

$$\theta \to U(I-\eta \Sigma )U^T\theta$$

Under this update, we have that

$${\theta }_t=U(I-\eta \Sigma {)}^t{U}^T{\theta }_0$$

where the inner $U$‘s and $U^T$‘s cancel by orthogonality.

Clearly, if $|1-\eta {\sigma }_i|>1$, the norm of ${\theta }_t$ will blow up and not be able to converge. So, for convergence, we need

$$|1-\eta {\sigma }_i|\le 1$$

Observe that under this constraint, ${\theta }_t$ will go to $0$ as $t\to \infty$. For our contrived loss $\frac{1}{2}{\theta }^{T}H\theta$, this is exactly what we want as the global minimum for this loss happens to be $\theta =0$. In a more general setting, this $\theta$ might be more like the error in the parameters, i.e. ${\theta }_{\text{current}}-{\theta }_{\text{target}}$, and we might look at the quadratic approximation for the loss near the minima.

Solving the inequality for $\eta$:

$$-1\le 1-\eta {\sigma }_i\le 1⟹0\le \eta \le \frac{2}{{\sigma }_i}$$

To ensure stability across all dimensions, the learning rate $\eta$ is bounded by the steepest direction (the largest eigenvalue, ${\sigma }_{max}$). The maximal stable learning rate is:

$${\eta }_{max}=\frac{2}{{\sigma }_{max}}$$

The maximal eigenvalue of $H$, which is the Hessian of the loss, is equivalent to the sharpness of the loss. So, the punchline here is that the maximal stable learning rate is inversely proportional to the sharpness of the loss. This punchline holds as long as the Hessian of the loss is constant. However, clearly, for complex nonlinear models, this will not be the case. For neural networks, the “Edge of Stability” (EoS) phenomena suggests that during training, the sharpness of the loss progressively increases (progressive sharpening) until it saturates and hovers just around the value $\frac{2}{\eta }$, the largest it can get without causing instability in training. This can be thought of as a kind of self-correcting, self-regulating behavior in neural nets.

disclaimer: this note was mostly transcribed by Gemini