Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction

Abstract: Normalization layers (e.g., Batch Normalization, Layer Normalization) were introduced to help with optimization difficulties in very deep nets, but they clearly also help generalization, even in not-so-deep nets. Motivated by the long-held belief that flatter minima lead to better generalization, this paper gives mathematical analysis and supporting experiments suggesting that normalization (together with accompanying weight-decay) encourages GD to reduce the sharpness of loss surface. Here "sharpness" is carefully defined given that the loss is scale-invariant, a known consequence of normalization. Specifically, for a fairly broad class of neural nets with normalization, our theory explains how GD with a finite learning rate enters the so-called Edge of Stability (EoS) regime, and characterizes the trajectory of GD in this regime via a continuous sharpness-reduction flow.

\section{Experiments} \label{sec:add-exp}

In this section,
we provide experiments on matrix completion and CIFAR-10
to validate the main claim in our theory:
GD+WD\xspace on scale-invariant loss persistently reduces spherical sharpness in the EoS regime
(the regime where $2/\tilde{\eta}_t$ roughly equals to the spherical sharpness).
In addition, we also validate whether
the reduction of spherical sharpness indeed leads to better generalization,
and whether experimental phenomena match  our theory.

\subsection{Validation of Sharpness Reduction on Matrix Completion} \label{sec:exp-matlin}

\begin{figure}[tp]
  \includegraphics[width=\textwidth]{fig/matcom-d50-id0.pdf}
  \includegraphics[width=\textwidth]{fig/matcom-d50-id1.pdf}
  \includegraphics[width=\textwidth]{fig/matcom-d50-id2.pdf}
  \caption{
    Overparameterized matrix completion with BN,
    where
    the ground-truth matrix
    ${\bm{M}} \in \mathbb{R}^{50 \times 50}$ is of rank 2,
    and the number of observations varies from $400$ to $1600$.
    See \Cref{fig:matcom-intro} for the case of $800$ observations.
    The test loss starts to decrease significantly
    as soon as the spherical sharpness starts to decrease.
  }
  \label{fig:matcom-d50}
\end{figure}
\begin{figure}[tp]
  \includegraphics[width=\textwidth]{fig/matcom-d100-id1.pdf}
  \includegraphics[width=\textwidth]{fig/matcom-d100-id2.pdf}
  \includegraphics[width=\textwidth]{fig/matcom-d100-id3.pdf}
  \caption{
    Overparameterized matrix completion with BN,
    where
    the ground-truth matrix
    ${\bm{M}} \in \mathbb{R}^{50 \times 50}$ is of rank 2,
    and the number of observations vary from $1600$ to $3200$.
    The test loss starts to decrease significantly
    as soon as the spherical sharpness starts to decrease.
  }
  \label{fig:matcom-d100}
\end{figure}

First, we conduct experiments on matrix completion: there is an unknown low-rank ground-truth matrix ${\bm{M}} \in \mathbb{R}^{d \times d}$,
and only $N$ entries of ${\bm{M}}$ are known. The goal is to recover ${\bm{M}}$ given the known entries.
Matrix completion has been studied via deep learning techniques~(see, e.g., \citep{chi2019nonconvex} for a survey).
Matrix completion can be connected to supervised learning as follows.
Each entry of the matrix can be seen as a data point,
where the $N$ observed entries
constitute the training set.
As in supervised learning,
given a subset of data points,
the goal of matrix completion is to build a model to predict
the rest of the data points.

Here we empirically study solving matrix completion using an overparameterized scale-invariant model.
We can observe that spherical sharpness
is indeed decreasing as soon as the process enters the EoS regime,
and the reduction of spherical sharpness encourages low-rank.

In our experiments,
we generate the ground-truth matrix
${\bm{M}} \in \mathbb{R}^{d \times d}$
as follows.
First, we set $\widetilde{{\bm{M}}} \gets {\bm{U}}_* {\bm{V}}_*^\top$ for
two random matrices ${\bm{U}}_*, {\bm{V}}_* \in \mathbb{R}^{d \times 2}$,
where every entry is sampled uniformly
from $[-1, 1]$.
Then we obtain ${\bm{M}}$ by normalizing $\widetilde{{\bm{M}}}$
so that
the second moment of
the entries is $1$,
i.e.,
${\bm{M}} \gets (d \cdot \normFsm{\widetilde{{\bm{M}}}}^{-1})\widetilde{{\bm{M}}}$.
We uniformly sample $N$
entries of ${\bm{M}}$
to serve as the observations,
and use $\Omega \subseteq [d] \times [d]$
to denote the index set of observed entries.

Matrix completion
has been studied
by a line of works~\citep{gunasekar2017implicit,li2018algorithmic,arora2019implicit,razin2020implicit,li2021towards,stoger2021small}
as a test-bed for the implicit regularization of gradient descent.
More specifically,
they parameterize the target matrix as ${\bm{W}} = {\bm{U}} {\bm{V}}^\top$
where ${\bm{U}}, {\bm{V}} \in \mathbb{R}^{d \times d}$ are two trainable matrices,
and run GD to minimize the squared loss $\mathcal{L}({\bm{U}}, {\bm{V}}) := \frac{1}{N} \sum_{(i, j) \in \Omega} (W_{i,j} - M_{i,j})^2$.
Although there is no explicit constraint on rank,
GD with small random initialization can still exhibit an implicit bias
towards low-rank solutions.

Inspired by this line of works,
we conduct matrix completion experiments to test
if GD+WD\xspace exhibits
the same low-rank bias
in training overparameterized models with BN.
More specifically,
we parameterize the target matrix as ${\bm{W}} = \mathrm{BN}({\bm{U}} {\bm{V}}^\top)$,
where ${\bm{U}}, {\bm{V}} \in \mathbb{R}^{d \times d}$ are two trainable matrices.
Given the observed positions $\Omega$, the output of the model
for a single position $(i, j) \in \Omega$
is
$\frac{\gamma}{\sigma}[{\bm{U}} {\bm{V}}^\top]_{i,j}$,
where $\sigma^2 := \frac{1}{N}\sum_{(i,j) \in \Omega} [{\bm{U}} {\bm{V}}^\top]_{i,j}^2$
is the second moment of $[{\bm{U}} {\bm{V}}^\top]_{i,j}$
over all observed positions,
and $\gamma$ is a rescaling factor.
Multiplying the factor $\frac{\gamma}{\sigma}$
can be seen as doing BN
over observed entries
because
it rescales the output in the same manner as BN.
But the difference is that,
for the sake of simplicity,
we do not subtract the mean.
To ensure the loss to be scale-invariant,
we also fix $\gamma$ to match the second moment of observed entries,
i.e., $\gamma^2 := \frac{1}{N} \sum_{(i, j) \in \Omega} M_{i,j}^2$.

To train this model, we use the standard squared loss,
and we run gradient descent with LR $\hat{\eta} = 0.1$ and WD $\hat{\lambda} = 0.01$ to optimize the loss.
It is obvious that this loss is scale-invariant due to BN.
\begin{align*}
  \mathcal{L}({\bm{U}}, {\bm{V}}) := \frac{1}{N} \sum_{(i, j) \in \Omega} \left(\frac{\gamma}{\sigma} [{\bm{U}} {\bm{V}}^\top]_{i,j} - M_{i,j}\right)^2.
\end{align*}
For an unobserved entry $(i, j)$,
the model
uses the same batch statistics
as for observed entries,
and predicts
$\frac{\gamma}{\sigma}[{\bm{U}} {\bm{V}}^\top]_{i,j}$,
where $\gamma, \sigma$ are the same as above.
So we measure the test loss as
\begin{align*}
  \bar{\mathcal{L}}({\bm{U}}, {\bm{V}}) := \frac{1}{d^2} \sum_{i, j \in [d]} \left(\frac{\gamma}{\sigma} [{\bm{U}} {\bm{V}}^\top]_{i,j} - M_{i,j}\right)^2.
\end{align*}
We note that the loss function has no explicit constraint on rank.
But surprisingly,
in our experiments,
GD+WD\xspace
tends to prefer low-rank solution
as soon as the sharpness-reduction bias starts to
occur.
See \Cref{fig:matcom-intro,fig:matcom-d50}
for experiments on reconstructing
a rank-2 matrix of size $50 \times 50$,
and \Cref{fig:matcom-d100}
for experiments on reconstructing a rank-2 matrix of size $100 \times 100$.
In all these experiments,
we can observe that the train loss first decreases to near zero in a short time,
but the test loss
remains at a high level.
Then after effective LR increases for some more steps,
the dynamics enters the EoS regime,
and as predicted by our theory,
the spherical sharpness starts to decrease.
Meanwhile,
the test loss also starts to decrease significantly.
We notice that the gap between the second and the third largest singular values
of the recovered matrix
is enlarged at the same time,
suggesting that this reduction of spherical sharpness encourages the recovered matrix to be low-rank.

\subsection{Validation of Sharpness Reduction on CIFAR-10} \label{sec:exp-cifar10}

Now we present experiments on CIFAR-10
with crossentropy loss
to validate the sharpness-reduction bias in a more realistic setting.
We run full-batch GD with accompanying WD
on three different architectures:
a scale-invariant variant of VGG-11,
a scale-invariant variant of pre-activation ResNet-20,
and the standard pre-activation ResNet-20.
We fix LR $\hat{\eta} = 0.1$ and WD $\hat{\lambda} = 5 \times 10^{-4}$.
See \Cref{sec:exp-details-cifar10} for more details on the architectures and training procedures.

The experiment for the scale-invariant variants of VGG-11 and ResNet-20 are presented in \Cref{fig:full-cifar,fig:full-cifar-fsi-resnet}.
Here
our ResNet-20 is non-smooth due to the use of ReLU,
but our VGG-11 is smooth because we choose to use Swish activation and mean pooling in this network.
For both the smooth VGG-11 and non-smooth ResNet-20,
it can be seen from the plots that in both cases the spherical sharpness has an overall tendency to decrease over time,
and the test accuracy is increasing accordingly.

\begin{figure}[p]
  \includegraphics[width=\textwidth]{fig/full-cifar-resnet.pdf}
  \vspace{-0.25in}
  \caption{
    In training a scale-invariant ResNet-20 on CIFAR-10 with (full-batch) GD+WD\xspace,
    the spherical sharpness decreases over time.
    $100\%$ training accuracy is achieved after $\sim 4700$ steps (dotted line),
    but as the training continues,
    the test accuracy increases from $60.8\%$ to $77.4\%$.
  }
  \label{fig:full-cifar-fsi-resnet}
\end{figure}

\begin{figure}[p]
  \includegraphics[width=\textwidth]{fig/full-cifar-resnet-bp.pdf}
  \vspace{-0.25in}
  \caption{
    In training the standard pre-activation ResNet-20 on CIFAR-10 with (full-batch) GD+WD\xspace,
    the spherical sharpness is decreasing in the long term run.
    $100\%$ training accuracy is achieved after $\sim 4400$ steps (dotted line),
    but as the training continues,
    the test accuracy increases from $67.5\%$ to $80.4\%$.
    Spherical sharpness is only evaluated for a scale-invariant part of the trainable parameters.
  }
  \label{fig:full-cifar-bp-resnet}
\end{figure}

\begin{figure}[t]
  \includegraphics[width=\textwidth]{fig/full-cifar-vgg-beginning.pdf}
  \vspace{-0.25in}
  \caption{
    The first 1500 steps
    in training a smooth and scale-invariant VGG-11 on CIFAR-10 with GD+WD\xspace (see also \Cref{fig:full-cifar}).
    The spherical sharpness is much larger than $2/\tilde{\eta}_0$ initially
    but is then reduced to the same level as $2/\tilde{\eta}_t$ after a few steps.
  }
  \label{fig:full-cifar-vgg-beginning}
\end{figure}

We note that
this sharpness-reduction bias
actually goes beyond the setting that our theory can directly apply:
the dynamic enters the EoS regime in the very beginning of training,
but our theory
only analyzes the dynamic near a local minimizer manifold.
This is because
that the LR here is not small enough so that ${\bm{\theta}}_t$ can approach to a local minimizer
before the dynamic enters the EoS regime.
In fact,
\Cref{fig:full-cifar-vgg-beginning} shows that
in the scale-invariant VGG-11 experiment,
the initial spherical sharpness is much larger than $2 / \tilde{\eta}_0$,
then after a few steps, the spherical sharpness decreases to a level that is close to $2/\tilde{\eta}_t$
and the dynamic enters the EoS regime.

Besides the scale-invariant models,
we also validate the sharpness-reduction bias on the standard pre-activation ResNet-20,
which is not (fully) scale-invariant but only scale-invariant to a part of its parameters.
For evaluating spherical sharpness,
we compute the partial Hessian only with respect to that part of parameters.
See \Cref{fig:full-cifar-bp-resnet} for the plot and \Cref{sec:exp-details-cifar10} for more experimental details.
Although our theory can only cover scale-invariant models,
we can still observe the sharpness-reduction bias in experiments,
and the test accuracy increases as spherical sharpness decreases.

\subsection{Periodic Behaviors}

A key insight in our proof
is that
the magnitude of oscillation and effective LR evolve periodically in the EoS regime.
Following \Cref{sec:proof-idea},
we can divide the EoS regime into two sub-regimes that occur alternatively in training:
(1)~the sub-regime where $2/\tilde{\eta}_t$ is smaller than $\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$
and the magnitude of oscillation $\abssm{h_t}$ keeps increasing;
and (2)~the sub-regime where $2/\tilde{\eta}_t$ is bigger than $\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$
and the magnitude of oscillation $\abssm{h_t}$ keeps decreasing.
Here ${\bm{\phi}}_t$ is a projection of ${\bm{\theta}}_t$ onto the local minimizer manifold,
$\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$ is the spherical sharpness computed at ${\bm{\phi}}_t$,
and $h_t$ is the inner product between ${\bm{\theta}}_t - {\bm{\phi}}_t$
and the top eigenvector of ${\bm{H}}({\bm{\phi}}_t)$.

\Cref{fig:lin-idea} provides a nice validation of this periodic behavior in linear regression with BN.
We further validate the periodic behavior
in matrix completion by visualizing
$\tilde{\eta}_t$ in \Cref{fig:eos-matcom-intro} around the moment that the dynamic enters the EoS.
We do not plot $\abssm{h_t}$ because computing ${\bm{\phi}}_t$ is inefficient,
but we can observe that the training loss evolves periodically, which signals that $\abssm{h_t}$
evolves periodically as well.

In CIFAR-10 experiments with realistic LR and WD,
the periodic behavior still exists, but it may deviate from the regime that our theory can capture.
In most of our CIFAR-10 experiments,
only the sub-regime where $2/\tilde{\eta}_t < \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$
can be observed,
and the phenomenon that $2/\tilde{\eta}_t$ switches back and forth between being smaller and larger than the spherical sharpness
does not occur anymore.
However, the change of $\tilde{\eta}_t$ and spherical sharpness
still cause to the loss and gradient norm to fluctuate periodically.
See \Cref{fig:full-cifar-vgg-periodic} for details.
This difference is because that
the loss function of a neural net on CIFAR-10 is much less smooth than that of linear regression and matrix completion.
To make the dynamic happen in the regime that our theory describes,
the LR and WD need to be very small
so that ${\bm{\theta}}_t$ is sufficiently close to the minimizer manifold $\mathit{\Gamma}$,
but they cannot be that small in practice due to computational inefficiency.

The periodic behavior can also happen in a macroscopic scale on CIFAR-10.
E.g.,
in \Cref{fig:full-cifar},
the training loss spikes to a high value around step 50,000,
but then recovers to a near-zero value after a few hundred steps.
Such a macroscopic periodic behavior may not always lead to sharpness reduction,
in contrast to
the microscopic periodic behavior of our interest.
But as noted in the extensive empirical study by \citet{lobacheva2021on},
this macroscopic periodic behavior
sometimes leads to better generalization.
They also provide a theoretical analysis for the cause assuming the gradient norm is both lower and upper bounded.
However, their explanation is not completely satisfactory
because the gradient norm does not admit a lower bound when the parameter is close to the minimizer manifold $\mathit{\Gamma}$.
We leave it a future work to look further into the cause of macroscopic periodic behavior and how it helps generalization.

\begin{figure}[t]
  \centering
  \includegraphics[width=\textwidth]{fig/eos-matcom-intro.pdf}

  \caption{
    A closer look at the dynamic in
    the matrix completion experiment (\Cref{fig:matcom-intro})
    around the moment that it enters the EoS regime.
    $2 / \tilde{\eta}_t$
    switches back and forth between being smaller and larger than the spherical sharpness
    (computed as $\lambda^{\mathrm{H}}_1({\bm{\theta}}_t)$ for efficiency),
    which causes the training loss to oscillate periodically.
  }
  \label{fig:eos-matcom-intro}
\end{figure}

\begin{figure}[t]

  \centering
  \includegraphics[width=0.85\textwidth]{fig/full-cifar-vgg-periodic.pdf}

  \caption{
    A closer look at a sample of 200 steps in the middle of the VGG-11 training on CIFAR-10 (\Cref{fig:full-cifar}).
    $2/\tilde{\eta}_t$ is always smaller than the spherical sharpness,
    but the training loss and gradient norm still fluctuate periodically.
  }
  \label{fig:full-cifar-vgg-periodic}
\end{figure}
\begin{figure}[tp]
  \centering
  \includegraphics[width=0.85\textwidth]{fig/full-cifar-vgg-wd-vs-nowd.pdf}
  \vspace{-0.05in}
  \caption{
    In training scale-invariant VGG-11 on CIFAR-10 with full-batch GD (LR $\tilde{\eta} = 0.1$),
    weight decay plays an important role in sharpness reduction.
    When weight decay $\hat{\lambda} = 5 \times 10^{-4}$,
    the spherical sharpness persistently decreases,
    and the test accuracy increases from $69.1\%$ to $84.3\%$ (see also \Cref{fig:full-cifar}).
    But if there is no weight decay,
    the spherical sharpness does not change much after a few initial steps,
    and the test accuracy is stuck at $66.4\%$.
  }
  \label{fig:full-cifar-vgg-wd-vs-nowd}
\end{figure}

\subsection{Ablation Study: Weight Decay}

Now we conduct experiments to study the effects of weight decay (WD)
on sharpness reduction.
It is crucial to have non-zero WD in our theoretical analysis.
Otherwise, it is implied by \Cref{lm:connection-prelim} that the effective LR is monotone decreasing,
but our theory only applies
to the effective LRs that can be viewed as quasi-RMSprop (see \Cref{def:qrms-sche,thm:gdwd-is-quasi-rmsprop}),
in which the effective LR can either increase or decrease as gradient norm changes.
A previous work by \citet{arora2018theoretical}
provides a detailed theoretical analysis in this case
showing that the dynamic always stays in the stable regime
after a few warm-up steps,
and the parameter eventually converges to a stationary point on the unit sphere under standard assumptions in optimization.

We can verify through experiments that
the requirement of WD to be non-zero
is not a technical artifact,
but indeed a necessity in practice to exhibit the sharpness-reduction bias.
We train a scale-invariant VGG-11 following exactly the same hyperparameters and initialization
as \Cref{fig:full-cifar}, except that we set WD $\hat{\lambda}$ to zero.
The result is presented in \Cref{fig:full-cifar-vgg-wd-vs-nowd}.
The spherical sharpness is no longer decreasing
when WD is zero,
which results in a much lower test accuracy than the VGG-11 with WD $5 \times 10^{-4}$.

We further conduct experiments with smaller (but non-zero) WD than our default value $5 \times 10^{-4}$.
For the sake of computational efficiency,
our experiments are conducted on a subset of CIFAR-10 images consisting of 2K images,
which we call CIFAR-10-2k (see \Cref{sec:exp-details-cifar10}).
The result is presented in \Cref{fig:full-cifar2k-vgg-ab-wd},
from which we can observe that
GD exhibits a sharpness-reduction bias
as long as WD is non-zero.
But note that smaller WD leads to
a slower speed of sharpness reduction.
This is expected phenomenon
because
the speed of tracking the sharpness-reduction flow \eqref{eq:main-zeta}
is controlled by
the intrinsic LR $\eta_{\mathrm{in}} := \hat{\eta} \hat{\lambda}$,
and
the smaller the WD, the smaller intrinsic LR (when LR $\hat{\eta}$ is fixed).

\begin{figure}[tp]
  \vspace{-0.1in}
  \centering
  \includegraphics[width=0.85\textwidth]{fig/cifar2k-vgg-ab-wd.pdf}
  \vspace{-0.15in}
  \caption{
    In training scale-invariant VGG-11 on CIFAR-10-2k with full-batch GD (LR $\tilde{\eta} = 0.1$),
    the sharpness-reduction bias occurs as long as the weight decay is non-zero,
    but smaller WD leads to longer training time to reduce the spherical sharpness to the same level of that with larger WD.
    The best test accuracy achieved within
    400,000 steps
    is
    $48.9\%$ when $\hat{\lambda} = 0$,
    $51.1\%$ when $\hat{\lambda} = 5 \times 10^{-5}$,
    $52.5\%$ when $\hat{\lambda} = 10^{-4}$,
    $55.1\%$ when $\hat{\lambda} = 5 \times 10^{-4}$.
  }
  \label{fig:full-cifar2k-vgg-ab-wd}
\end{figure}

\begin{figure}[tp]
  \centering
  \includegraphics[width=\textwidth]{fig/cifar2k-vgg-ab-lrwd.pdf}
  \vspace{-0.15in}
  \caption{
    In training scale-invariant VGG-11 on CIFAR-10-2k with full-batch GD
    and various different initial effective learning rates $\tilde{\eta}_0$.
    We scale learning rate $\hat{\eta}$ and weight decay $\hat{\lambda}$
    inverse proportionally
    while fixing the intrinsic LR $\eta_{\mathrm{in}} = 5 \times 10^{-5}$.
    The best test accuracies achieved within
    200,000 steps
    are similar in all runs,
    with $52.7\%$ for $\hat{\eta} = 10^{-5}$,
    $53.7\%$ for $\hat{\eta} = 10^{-3}$,
    $54.3\%$ for $\hat{\eta} = 0.1$,
    $54.6\%$ for $\hat{\eta} = 10$.
  }
  \label{fig:full-cifar2k-vgg-ab-lrwd}
\end{figure}

\subsection{Ablation Study: Initial Effective Learning Rate}

As we have seen in \Cref{fig:full-cifar-vgg-beginning},
in training the scale-invariant VGG-11 on CIFAR-10,
if $\hat{\eta} = 0.1$ and $\hat{\lambda} = 5 \times 5 \times 10^{-4}$, then
initially the spherical sharpness is much larger than $2/\tilde{\eta}_0$,
and after a few steps, it decreases to the same level as $2/\tilde{\eta}_t$,
thus bringing the dynamic into the EoS regime.

Now we conduct an ablation study on the initial effective LR.
For computational efficiency, we run experiments on CIFAR-10-2k (see \Cref{sec:exp-details-cifar10})
to train our scale-invariant variant of VGG-11.
We can see from \Cref{fig:full-cifar-vgg-wd-vs-nowd}
that setting LR $\hat{\eta} = 0.1$ and WD $\hat{\lambda} = 5 \times 10^{-4}$
leads to a similar phenomenon as above.
Now we change the LR $\hat{\eta}$ while rescaling the WD $\hat{\lambda}$ inverse proportionally
so that the intrinsic LR remains unchanged.
See \Cref{fig:full-cifar2k-vgg-ab-lrwd} for the results.
When the initial effective LR is as small as $10^{-5}$,
the dynamic is in the stable regime initially,
and $2/\tilde{\eta}_t$ falls so slowly
that the dyanmic enters the EoS regime
only after reaching $100\%$ training accuracy.
Then the spherical sharpness starts to persistently decrease,
and the test accuracy increases accordingly.
In the end, all the experiment runs reach nearly the same spherical sharpness and test accuracy,
regardless the initial effective LR.

\section{Experiment Details} \label{sec:exp-details}

All our experiments were conducted on NVIDIA RTX A5000 GPUs.
The longest experiments are the full-batch training experiments on CIFAR-10,
each of which took 7 days to run.

\subsection{Additional Details of CIFAR-10 Experiments} \label{sec:exp-details-cifar10}

Our implementation of full-batch GD is based on the code of \citet{cohen2021gradient}\footnote{\url{https://github.com/locuslab/edge-of-stability}}.
Our implementations of VGGNets~\citep{simonyan2014vgg} and ResNets~\citep{he2016deepres,he2016identityres} on CIFAR-10 are based on a high-starred GitHub repository of
Wei Yang (\href{https://github.com/bearpaw}{bearpaw})\footnote{\url{https://github.com/bearpaw/pytorch-classification}}.
We do not use any data augmentation.
As it is not feasible to do BN over the full dataset
under our GPU memory constraints,
we use ghost batch normalization~\citep{hoffer2017trainlonger} instead,
where we split
the dataset into $50$ ghost batches of size $1000$.

Some of our experiments are conducted on only a subset of CIFAR-10 consisting of $2000$ images,
which we call CIFAR-10-2k.
We construct the dataset by scanning the full CIFAR-10 dataset and taking the first $200$ images
for each of the $10$ classes.
We use the standard batch normalization with full batch for all CIFAR-10-2k experiments.

Three different neural network architectures
are tested in our experiments,
and we refer to them as scale-invariant VGG-11, standard ResNet-20, and scale-invariant ResNet-20 respectively.

\myparagraph{Scale-Invariant VGG-11.} Our scale-invariant VGG-11 architecture is similar to the configuration A of the original VGGNet~\citep{simonyan2014vgg},
but we make the following changes.
We add a BN layer (without affine parameters) between every convolution and activation to introduce scale-invariance,
We use mean pooling instead of max pooling, Swish instead of ReLU to make the training loss smooth.
We replace the final $3$ fully-connected layers with only one fully-connected layer
as in Yang's implementation.
All convolutional and fully-connected layers have no bias terms.
We add a BN layer
after the last fully-connect layer,
in which the affine parameters are fixed to $\gamma = \frac{3}{10} \ln 91 \approx 1.353$, $\beta = 0$ (see \Cref{sec:affine} for discussion).
A visualization of the full architecture is given in \Cref{fig:arch-vgg}, whose output function can be shown to be scale-invariant with respect to the trainable parameters.

\begin{figure}[htbp]
  \centering
  \includegraphics[scale=0.13]{fig/vgg-all.pdf}
  \caption{
    The architecture of
    a scale-invariant variant of VGG-11.
    Each rounded rectangle stands for a layer with trainable parameters,
    including convolutional and fully-connected layers.
    Each hexagon stands for a layer with no trainable parameters,
    including mean pooling, swish, BN without affine parameters, and BN with affine parameters fixed.
  }
  \label{fig:arch-vgg}
\end{figure}

\begin{figure}[htbp]
  \centering
  \begin{subfigure}{\textwidth}
    \centering
    \includegraphics[scale=0.13]{fig/std-preresnet-all.pdf}
    \caption{The architecture of ResNet-20, following Yang's implementation.}
    \label{fig:arch-std-resnet}
  \end{subfigure}
  \par\bigskip
  \begin{subfigure}{\textwidth}
    \centering
    \includegraphics[scale=0.13]{fig/fsi-preresnet-all.pdf}
    \caption{The architecture of a scale-invariant variant of ResNet-20.}
    \label{fig:arch-fsi-resnet}
  \end{subfigure}
  \caption{
    The ResNet architectures used in our experiments.
    Each rounded rectangle stands for a layer with trainable parameters,
    including convolutional and fully-connected layers, and BN with trainable affine parameters.
    Each hexagon stands for a layer with no trainable parameters,
    including mean pooling, ReLU, BN without affine parameters, and BN with affine parameters fixed.
    In Yang's implementation of ResNet-20,
    the output function is scale-invariant to the parameters in layers marked as rounded green rectangles;
    in the scale-invariant variant,
    the output function is scale-invariant to all trainable parameters.
  }
  \label{fig:arch-resnet}
\end{figure}

\myparagraph{Standard ResNet-20.}
We also verify the sharpness-reduction bias on a more realistic setting with ResNet-20 and ReLU.
In some of our experiments, we use Yang's implementation for the pre-activation variant of ResNet-20.
All convolutional layers have no bias, but the final fully-connected layer does have bias terms.
We note that the output function of this architecture
is not scale-invariant to all its trainable parameters.
However, it can be shown that
the output function is indeed scale-invariant to a large part of its parameters.
More specifically,
the trainable parameters can be split into two parts, ${\bm{w}}$ and ${\bm{z}}$,
and the output function remains unchanged if we replace ${\bm{w}}$ with $c{\bm{w}}$ for all $c > 0$.
In our experiments, we measure the spherical sharpness by $\nabla^2_{{\bm{\theta}}} \mathcal{L}({\bm{\theta}}, {\bm{z}})$
while letting ${\bm{\theta}} := \frac{{\bm{w}}}{\normtwosm{{\bm{w}}}}$.
See \Cref{fig:arch-std-resnet} for a visualization of the architecture.

\myparagraph{Scale-Invariant ResNet-20.}
The standard ResNet-20
can be scale-invariant after making minor changes.
First, we remove the affine parameters in all existing BNs.
Then, we also add an extra BN (with trainable affine parameters)
to each shortcut for blocks with downsampling,
following \citet{li2020exp}.
Finally, we add a BN layer
after the last fully-connect layer,
in which the affine parameters are fixed to $\gamma = \frac{3}{10} \ln 91 \approx 1.353$, $\beta = 0$ (see \Cref{sec:affine} for discussion).
A visualization of the full architecture is given in \Cref{fig:arch-fsi-resnet}, whose output function can be shown to be scale-invariant with respect to the trainable parameters.

\subsection{Additional Details of Linear Regression Experiments}

Our experiments on linear regression (\Cref{fig:lin-idea})
follows strictly as \Cref{sec:lin}.
For generating data,
we first sample the ground-truth weight $\vw_{\mathrm{GT}} \in \mathbb{S}^{d-1}$ uniformly,
where $d = 40$,
and we sample the ground-truth bias as $b_{\mathrm{GT}} \sim \mathcal{N}(0, 0.01)$,
Then we sample $n = 20$ points from Gaussian distribution
$\mathcal{N}({\bm{0}}, {\bm{\Sigma}})$,
where ${\bm{\Sigma}} = \mathrm{diag}(1/d, 2/d, 3/d, \dots, 1)$.

To generate \Cref{fig:lin-idea},
we run gradient descent with LR $\hat{\eta} = 0.5$
and WD $\hat{\lambda} = 2 \times 10^{-4}$.
This particular choice of hyperparameters is for better visualization of the periodic behavior
described in \Cref{sec:proof-idea}.
If we enlarge LR or WD, the period will be shortened accordingly.

\subsection{Computing Spherical Sharpness}

To compute the spherical sharpness at a point ${\bm{w}} \in \mathbb{R}^D \setminus \{{\bm{0}}\}$ (\Cref{def:sph-sha}),
we utilize the formula
$\lambda_1(\nabla^2 \mathcal{L}(\frac{{\bm{w}}}{\normtwosm{{\bm{w}}}})) = \normtwosm{{\bm{w}}}^2 \cdot \lambda_1(\nabla^2 \mathcal{L}({\bm{w}}))$,
which can be proved for all scale-invariant loss by simple calculus (see also \Cref{lm:scale-invariant-grad-hess}).
For computing the eigenvalues of $\nabla^2 \mathcal{L}({\bm{w}})$,
we invoke the Lanczos algorithm from SciPy library (\texttt{scipy.sparse.linalg.eigsh}),
where a Hessian-vector product oracle ${\bm{x}} \mapsto \nabla^2 \mathcal{L}({\bm{w}}) {\bm{x}}$ is implemented with PyTorch
and passed to the Lanczos algorithm as a linear operator to find its eigenvalues.

In full-batch training of the full CIFAR-10 dataset,
it is time-consuming to compute even a single Hessian-vector product.
Following~\citet{cohen2021gradient},
we compute Hessian based on
only the first $5000$ training data points.

Note that in our theory
we focus on the spherical sharpness
of minimizers.
However, in practice, it is usually time-consuming
to compute the minimizers exactly,
so in matrix completion and CIFAR-10 experiments,
we compute the spherical sharpness directly
at the current parameter ${\bm{w}}_t$.
But for linear regression experiments (\Cref{fig:lin-idea}),
the spherical sharpness is indeed computed at minimizers
in each step because
our computational power is sufficient in this setting.
Specifically,
we compute the minimizer by
doing projected gradient descent on $\mathbb{S}^{D-1}$
with fixed learning rate $0.005$
until the loss decreases to $10^{-8}$.

\section{Discussion on the Affine Parameters of the Final BN} \label{sec:affine}

In our experiments for scale-invariant models,
we add a BN to the final linear layer and fix the affine parameters to be constants.
Now we analyze its effects on
regression tasks with squared loss and
classification tasks with crossentropy loss,
and discuss how to set the affine parameters reasonably.

We consider a scale-invariant neural net in general while assuming that a BN is put as the last layer.
For input ${\bm{x}}$ and parameter ${\bm{w}}$,
let $(F_1({\bm{x}}; {\bm{w}}), \dots, F_C({\bm{x}}; {\bm{w}})) \in \mathbb{R}^C$ be the output before the final BN.
After the final BN, the output is given by the following function $\Phi_k({\bm{w}}; {\bm{w}})$:
\begin{align*}
    \bar{F}_k({\bm{x}}; {\bm{w}}) &:= \frac{F_k({\bm{x}}; {\bm{w}}) - \mu_k}{\sigma_k}, \\
    \Phi_k({\bm{x}}; {\bm{w}}) &:= \gamma_k \bar{F}_k({\bm{x}}; {\bm{w}}) + \beta_k,
\end{align*}
where $\mu_k$ and $\sigma_k^2$ are mean and variance of $\{ F_k({\bm{x}}_i; {\bm{w}}) \}_{i=1}^{n}$
over the training set.

\subsection{Squared Loss}
For linear regression experiments,
we use squared loss as the loss function,
and we have $C = 1$ because the output is a scalar.
Then as discussed in \Cref{sec:lin}, the mean and variance of $\Phi_1({\bm{x}}; {\bm{w}})$
are always $\beta_1$ and $\gamma_1^2$ regardless of ${\bm{w}}$.
Therefore, to ensure that $\Phi$ has sufficient representation power to express the regression targets $\{y_i\}_{i=1}^{n}$,
we should fix the affine parameters in a way
that $\beta_1$ and $\gamma_1^2$ match with the mean and variance of $\{y_i\}_{i=1}^{n}$.

In the general case when $C \ge 1$, the $k$-th output unit has mean and variance $\beta_k$ and $\gamma_k^2$.
We can argue similarly that
we should fix the affine parameters in a way
that $\beta_k$ and $\gamma_k^2$ match with the mean and variance of $\{y_{i,k}\}_{i=1}^{n}$,
where $y_{i,k}$ stands for the $k$-th coordinate of the regression target for the $i$-th data point.
In other words,
\begin{align*}
  \beta_k \gets \mu_{\mathrm{y}}^{(k)} := \frac{1}{n} \sum_{i=1}^{n} y_{i,k},  & & \gamma_k \gets \sigma_{\mathrm{y}}^{(k)} := \sqrt{\frac{1}{n} \sum_{i=1}^{n} (y_{i,k} - \mu_{\mathrm{y}}^{(k)})^2}.
\end{align*}

\subsection{Crossentropy Loss}
In our CIFAR-10 experiments, we fix all $\gamma_k$ to a constant value $\gamma$,
and fix $\beta_k$ to zero.
We note that fixing these affine parameters in this way for an overparameterized model
is equivalent to adding label smoothing to the training loss in some sense.

The basic idea is to assume the model has sufficient express any function $F$,
then $\Phi_k$ can be any function of mean zero and covariance $\gamma^2$ on the training set.
If the dataset
has $C$ classes and each class has equal numbers of samples,
with some efforts one can show that
the minimum loss is
attained
when
the pre-softmax logit for an input
is $\gamma (C-1)^{1/2}$ for the correct class
and $-\gamma (C-1)^{-1/2}$ for every wrong class.
In this case,
the output probability $P_y({\bm{x}}_i)$ for a class $y$ is
\begin{align} \label{eq:pxy-opt}
  P_y({\bm{x}}_i) = \begin{cases}
    \frac{1}{1 + (C-1)\exp(-\gamma((C-1)^{1/2} + (C-1)^{-1/2}))} & \quad \text{if } y = y_i; \\
    \frac{1}{(C-1) + \exp(\gamma((C-1)^{1/2} + (C-1)^{-1/2}))} & \quad \text{otherwise.} \\
  \end{cases}
\end{align}
Label smoothing is a regularization technique proposed by \citet{szegedy2016rethinking}
to improve the generalization of neural nets
by replacing the one-hot hard labels
with soft labels that assign a probability of $\epsilon/C$ to each wrong class.
To see the connection to label smoothing,
we can let $\epsilon :=
    \frac{C}{(C-1) + \exp(\gamma((C-1)^{1/2} + (C-1)^{-1/2}))}$,
then $P_y({\bm{x}}_i) = \epsilon / C$ if $y \ne y_i$
and $P_y({\bm{x}}_i) = 1 - (1-1/C)\epsilon$ if $y = y_i$.
In our CIFAR-10 experiments,
$\gamma = \frac{3}{10} \ln 91 \approx 1.353$.
Then by simple calculation $\epsilon = 0.1$,
and therefore the neural net trained with $\gamma_k$ and $\beta_k$ fixed to $\gamma$ and $0$
is encouraged to produce
the output probabilities to match with the soft labels.

The minimum loss is non-zero when the affine parameters are fixed.
We can see from \eqref{eq:pxy-opt}
that
the loss attained by any minimizer is
\begin{align}
  \mathcal{L}_{\min} = \ln(1+ (C-1)\exp(-\gamma((C-1)^{1/2} + (C-1)^{-1/2}))).
\end{align}
In making plots for our experiments,
we always subtract the original loss with its theoretical minimum value $\mathcal{L}_{\min}$.
When $\gamma = \frac{3}{10} \ln 91$, this minimum value is $\mathcal{L}_{\min} = \ln \frac{100}{91} \approx 0.0943$.

Choosing $\gamma$ to be $\frac{3}{10} \ln 91$ is mainly because of its connection to
label smoothing with $\epsilon = 0.1$,
which is the value chosen in \citep{szegedy2016rethinking}.
But other choices of $\gamma$ can also lead to the same sharpness-reduction bias in practice;
see \Cref{fig:cifar2k-vgg-ab-final-gamma}.

\begin{figure}[htbp]
  \centering
  \includegraphics[width=\textwidth]{fig/cifar2k-vgg-ab-final-gamma.pdf}
  \vspace{-0.15in}
  \caption{
    Training scale-invariant VGG-11 on CIFAR-10-2k with full-batch GD
    (LR $\hat{\eta} = 0.1$, WD $\hat{\lambda} = 5 \times 10^{-4}$)
    while fixing the rescaling parameter $\gamma$ in the final BN to different values.
    The sharpness-reduction bias can be observed in all cases,
    but for larger $\gamma$ the periodic behavior is more likely to happen at a macroscopic scale.
  }
  \label{fig:cifar2k-vgg-ab-final-gamma}
\end{figure}
\section{Additional Preliminaries} \label{sec:add-prelim}

\subsection{Additional Notations}
We use $\mathrm{cl}({\mathcal{M}})$ to denote the closure of a set ${\mathcal{M}}$.
For ${\bm{\theta}} \in \mathbb{R}^D$ and ${\mathcal{M}} \subseteq \mathbb{R}^D$,
we use $d_2({\bm{\theta}}, {\mathcal{M}}) := \inf\{ \normtwosm{{\bm{\theta}} - {\bm{\phi}}} : {\bm{\phi}} \in {\mathcal{M}} \}$
to denote the $L^2$-distance from ${\bm{\theta}}$ to ${\mathcal{M}}$.
For ${\bm{\theta}} \in \mathbb{R}^D$ and $\epsilon \ge 0$,
$\Ball{\epsilon}{{\bm{\theta}}} := \{ {\bm{\theta}}' \in \mathbb{R}^D : \normtwosm{{\bm{\theta}} - {\bm{\theta}}'} < \epsilon \}$
is the open $\epsilon$-ball centered at ${\bm{\theta}}$.
For a set ${\mathcal{M}} \subseteq \mathbb{R}^D$, ${\mathcal{M}}^{\epsilon} := \bigcup_{{\bm{\theta}} \in {\mathcal{M}}} \Ball{\epsilon}{{\bm{\theta}}}$
is the (open) $\epsilon$-neighborhood of ${\mathcal{M}}$. For a manifold $\mathit{\Gamma}$, we use $\TGa{p}$ and $\NGa{p}$
to denote the tangent and normal space of $\mathit{\Gamma}$ at a point $p \in \mathit{\Gamma}$.

Given a function $f: \mathbb{R}^D \to \mathbb{R}^n$
and vectors ${\bm{x}}, {\bm{v}} \in \mathbb{R}^D$,
we use $\partial f_{{\bm{x}}}[{\bm{v}}]$
to denote the directional derivative
$\partial f_{{\bm{x}}}[{\bm{v}}] := \lim_{t \to 0} \frac{1}{t} f({\bm{x}} + t{\bm{v}})$,
which also equals to the Jacobian of $f$ at ${\bm{x}}$ multiplied with ${\bm{v}}$.
We use $\partial^2 f_{{\bm{x}}}[{\bm{v}}, {\bm{u}}]$ to denote the second-order derivative
$\partial (\partial f_{{\bm{x}}}[{\bm{v}}])_{{\bm{x}}}[{\bm{u}}]$.
For a real-valued function $\mathcal{L}: \mathbb{R}^D \to \mathbb{R}$,
we use $\nabla \mathcal{L}({\bm{\theta}})$ for gradient,
$\nabla^2 \mathcal{L}({\bm{\theta}})$ for Hessian.
For third-order derivatives of $\mathcal{L}$,
we define $\partial^3 \mathcal{L}_{{\bm{\theta}}}[{\bm{v}}, {\bm{u}}] := \partial^2 (\nabla \mathcal{L})_{{\bm{\theta}}}[{\bm{v}}, {\bm{u}}] \in \mathbb{R}^D$.
If $\mathcal{L}$ is $\mathcal{C}^4$-smooth, then the following Taylor expansion holds for $\nabla \mathcal{L}$:
\begin{align*}
  \nabla \mathcal{L}({\bm{\theta}} + {\bm{x}}) = \nabla \mathcal{L}({\bm{\theta}}) + \nabla^2 \mathcal{L}({\bm{\theta}}) {\bm{x}} + \partial^3 \mathcal{L}_{{\bm{\theta}}}[{\bm{x}}, {\bm{x}}] + O(\normtwosm{{\bm{x}}}^3).
\end{align*}

\subsection{Scale-Invariant Functions}

The following lemma summarizes a few important properties of scale-invariant loss that have been
exploited in previous works \citep{arora2018theoretical,li2020exp}.
For completeness, we include a proof here.
\begin{lemma} \label{lm:scale-invariant-grad-hess}
    The following hold for a twice-differentiable scale-invariant function $\mathcal{L}({\bm{w}})$:
    \begin{enumerate}
      \item The gradient is $(-1)$-homogeneous and it is always perpendicular to ${\bm{w}}$, i.e.,
        $\nabla \mathcal{L}(c{\bm{w}}) = c^{-1} \nabla \mathcal{L}({\bm{w}})$ for all $c > 0$ and $\dotp{\nabla \mathcal{L}({\bm{w}})}{{\bm{w}}} = 0$;
      \item The Hessian matrix is $(-2)$-homogeneous, i.e.,
        $\nabla^2 \mathcal{L}(c{\bm{w}}) = c^{-2} \nabla^2 \mathcal{L}({\bm{w}})$ for all $c > 0$.
      \item $\nabla^2 \mathcal{L}({\bm{w}}) {\bm{w}} = -\nabla \mathcal{L}({\bm{w}})$.
    \end{enumerate}
\end{lemma}
\begin{proof}
  Taking gradients with respect to ${\bm{w}}$ on both sides of $\mathcal{L}(c{\bm{w}}) = \mathcal{L}({\bm{w}})$
  gives $\nabla \mathcal{L}(c{\bm{w}}) = c^{-1} \nabla \mathcal{L}({\bm{w}})$.
  Taking gradients again proves that $\nabla^2 \mathcal{L}(c{\bm{w}}) = c^{-2} \nabla^2 \mathcal{L}({\bm{w}})$.

  Taking derivative with respect to $c$ on both sides of $\mathcal{L}(c{\bm{w}}) = \mathcal{L}({\bm{w}})$
  gives $\dotpsm{\nabla \mathcal{L}(c{\bm{w}})}{{\bm{w}}} = 0$.
  Taking $c=1$ gives $\dotpsm{\nabla \mathcal{L}({\bm{w}})}{{\bm{w}}} = 0$.
  Finally, we take gradients with respect to ${\bm{w}}$,
  then $\nabla^2 \mathcal{L}({\bm{w}}) {\bm{w}} + \nabla \mathcal{L}({\bm{w}}) = {\bm{0}}$.
\end{proof}

Now we provide proofs for \Cref{lm:scale-invariant-on-sphere} and \Cref{thm:li-gdwd} in \Cref{sec:prelims}.
\begin{proof}[Proof for \Cref{lm:scale-invariant-on-sphere}]
By \eqref{eq:upd-gdwd} and definition of ${\bm{\theta}}_t$,
\begin{align*}
  {\bm{\theta}}_{t+1} = \Pi({\bm{w}}_{t+1})
  &= \Pi\left((1-\hat{\eta}\hat{\lambda}){\bm{w}}_t - \hat{\eta} \nabla \mathcal{L}({\bm{w}}_t)\right).
\end{align*}
Since $\Pi({\bm{x}}) = \Pi(c{\bm{x}})$ for all $c > 0$,
we can divide ${\bm{w}}_{t+1}$ by $(1 - \hat{\eta} \hat{\lambda}) \normtwosm{{\bm{w}}_t}$ and obtain
\begin{align*}
  {\bm{\theta}}_{t+1}
  &= \Pi\left({\bm{\theta}}_t - \frac{\hat{\eta}}{(1-\hat{\eta}\hat{\lambda})\normtwosm{{\bm{w}}_t}} \nabla \mathcal{L}({\bm{w}}_t)\right).
\end{align*}
Note that $\nabla \mathcal{L}({\bm{w}}_t) = \frac{1}{\normtwosm{{\bm{w}}_t}} \mathcal{L}({\bm{\theta}}_t)$ by \Cref{lm:scale-invariant-grad-hess}.
So we can further rewrite the above formula:
\begin{align*}
  {\bm{\theta}}_{t+1}
  &= \Pi\left({\bm{\theta}}_t - \frac{\hat{\eta}}{(1-\hat{\eta}\hat{\lambda})\normtwosm{{\bm{w}}_t}^2} \nabla \mathcal{L}({\bm{\theta}}_t)\right),
\end{align*}
which proves the lemma by definition of $\tilde{\eta}_t$.
\end{proof}

\begin{proof}[Proof of \Cref{thm:li-gdwd}]
	The first claim is directly from the Lemma D.2 of \citep{li2022robust}. The second claim about $\hat{\eta}$ holds by scrutinizing their proof.
\end{proof}

\subsection{Polyak-\L{}ojasiewicz\xspace Condition}

\begin{definition}[Polyak-\L{}ojasiewicz\xspace]
  For a loss function $\mathcal{L}({\bm{\theta}})$ and a constant $\mu > 0$,
  we say that $\mathcal{L}$ satisfies $\mu$-Polyak-\L{}ojasiewicz\xspace condition
  (or $\mu$-PL for brevity)
  on a set $U$ if
  \[
    \frac{1}{2}\normtwosm{\nabla \mathcal{L}({\bm{\theta}})}^2 \ge \mu \cdot \left( \mathcal{L}({\bm{\theta}}) - \inf_{{\bm{\theta}}' \in U} \mathcal{L}({\bm{\theta}}') \right),
  \]
  for all ${\bm{\theta}} \in U$.
\end{definition}

\subsubsection{Full Space Optimization}

\begin{theorem} \label{thm:rank-to-pl-ful}
  Let $\mathcal{L}: \mathbb{R}^D \to \mathbb{R}$ be a $\mathcal{C}^3$-smooth function,
  and $\mathit{\Gamma}$ be a $\mathcal{C}^1$-smooth, $D_\Gamma$-dimensional, open submanifold of $\mathbb{R}^D$,
  where every ${\bm{\theta}} \in \mathit{\Gamma}$ is a local minimizer of $\mathcal{L}$ and $\mathrm{rank}(\nabla^2 \mathcal{L}({\bm{\theta}})) = D - D_\Gamma$.
  If ${\mathcal{Z}}$ is a compact subset of $\mathit{\Gamma}$,
  then there exist $\epsilon > 0, \mu > 0$ such that
  $\mathrm{cl}({\mathcal{Z}}^{\epsilon}) \cap \mathit{\Gamma}$ is compact and
  $\mathcal{L}$ satisfies $\mu$-PL
  on
  ${\mathcal{Z}}^{\epsilon}$.
\end{theorem}
\begin{proof}
  Since ${\mathcal{Z}}$ is compact and $\mathit{\Gamma}$ is an open submanifold of $\mathbb{R}^D$,
  we can choose a small $\delta > 0$ such that
  ${\mathcal{N}} := \mathit{\Gamma} \cap \mathrm{cl}({\mathcal{Z}}^{\delta})$ is compact.

  It can be shown
  that
  there exists an open neighborhood $U$ of the compact submanifold ${\mathcal{N}}$ such that
  for every ${\bm{\theta}} \in U$,
  the nearest point on ${\mathcal{N}}$,
  $P({\bm{\theta}}) := \argmin\{ \normtwosm{{\bm{\theta}} - {\bm{\phi}}} : {\bm{\phi}} \in {\mathcal{N}} \}$,
  exists and is unique~\citep{foote1984distance}.

  We choose $\epsilon < \delta / 2$
  to be small enough so that $\epsilon < \delta / 2$ and ${\mathcal{Z}}^{\epsilon} \subseteq U$.
  For ${\bm{\theta}} \in {\mathcal{Z}}^{\epsilon}$, $P({\bm{\theta}})$ lies in the interior of the manifold ${\mathcal{N}}$,
  since $d_2(P({\bm{\theta}}), {\mathcal{Z}}) \le \normtwosm{P({\bm{\theta}}) - {\bm{\theta}}} + d({\bm{\theta}}, {\mathcal{Z}}) < \delta$.
  Then it must hold that ${\bm{\theta}} - P({\bm{\theta}}) \in \NGa[{\mathcal{N}}]{P({\bm{\theta}})}$;
  otherwise
  the differential of $\normtwosm{{\bm{\theta}} - {\bm{\phi}}}^2$ on ${\mathcal{N}}$ is non-zero
  at ${\bm{\phi}} = P({\bm{\theta}})$,
  which contradicts to the fact that $P({\bm{\theta}})$ is the nearest point to ${\bm{\theta}}$ on ${\mathcal{N}}$.

  Since $\mathcal{L} \in \mathcal{C}^3$,
  and $\nabla^2 \mathcal{L}({\bm{\theta}})$ is of constant rank $D-D_\Gamma$
  on the compact manifold ${\mathcal{N}}$,
  there exist
  $\lambda_{\min} > 0, \lambda_{\max} > 0$
  such that $\lambda_{D-D_\Gamma}(\nabla^2 \mathcal{L}({\bm{\theta}})) \ge \lambda_{\min}$
  and $\lambda_{1}(\nabla^2 \mathcal{L}({\bm{\theta}})) \le \lambda_{\max}$
  for all ${\bm{\theta}} \in {\mathcal{N}}$.
  Also by $\mathcal{L} \in \mathcal{C}^3$ and compactness of ${\mathcal{N}}$,
  there exists $C_3 > 0$ such that
  the following Taylor expansions hold
  for all ${\bm{\theta}} \in {\mathcal{Z}}^{\epsilon}$,
  \begin{align*}
    \mathcal{L}({\bm{\theta}}) - \mathcal{L}(P({\bm{\theta}})) &\le ({\bm{\theta}} - P({\bm{\theta}}))^\top \nabla^2 \mathcal{L}(P({\bm{\theta}})) ({\bm{\theta}} - P({\bm{\theta}})) + C_3 \normtwosm{{\bm{\theta}} - P({\bm{\theta}})}^3, \\
    &\le (\lambda_{\max} + C_3 \epsilon) \cdot \normtwosm{{\bm{\theta}} - P({\bm{\theta}})}^2. \\
    \normtwosm{\nabla \mathcal{L}({\bm{\theta}})}^2
    &\ge ({\bm{\theta}} - P({\bm{\theta}}))^\top \left(\nabla^2 \mathcal{L}(P({\bm{\theta}}))\right)^2 ({\bm{\theta}} - P({\bm{\theta}})) - C_3 \normtwosm{{\bm{\theta}} - P({\bm{\theta}})}^3 \\
    &\ge (\lambda_{\min}^2 - C_3 \epsilon) \cdot \normtwosm{{\bm{\theta}} - P({\bm{\theta}})}^2.
  \end{align*}
  Then $\mathcal{L}$ satisfies $\mu$-PL on ${\mathcal{Z}}^{\epsilon}$
  for $\mu := \frac{\lambda_{\min}^2 - C_3 \epsilon}{2(\lambda_{\max} + C_3 \epsilon)}$,
  which is positive if we choose $\epsilon$ to be small enough in the beginning.
\end{proof}

\begin{theorem} \label{thm:pl-converge-near-ful}
  If $\vtheta^*$ is a local minimizer of a $\mathcal{C}^2$-smooth function $\mathcal{L}: \mathbb{R}^D \to \mathbb{R}$ and
  $\mathcal{L}$ is $\mu$-PL on an open neighborhood $U$ of $\vtheta^*$,
  then for any ${\bm{\theta}}_0$ sufficiently close to $\vtheta^*$,
  a gradient flow $\frac{\textup{\textrm{d}} {\bm{\theta}}}{\textup{\textrm{d}} t} = -\nabla \mathcal{L}({\bm{\theta}})$
  starting with ${\bm{\theta}}_0$ converges to a point ${\bm{\theta}}_{\infty}$ as
  $t \to +\infty$ and $\normtwosm{{\bm{\theta}}_{\infty} - \vtheta^*} = O(\normtwosm{{\bm{\theta}}_0 - \vtheta^*})$.
\end{theorem}
\begin{proof}
Let $T := \inf\{ t : {\bm{\theta}}(t) \notin U \}$. For all $t < T$,
\begin{align*}
  \frac{\textup{\textrm{d}}}{\textup{\textrm{d}} t} (\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*))^{1/2}
  &= \frac{1}{2} (\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*))^{-1/2} \cdot \dotp{\nabla\mathcal{L}({\bm{\theta}})}{\frac{\textup{\textrm{d}} {\bm{\theta}}}{\textup{\textrm{d}} t}} \\
  &= -\frac{1}{2}(\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*))^{-1/2} \cdot \normtwosm{\nabla\mathcal{L}({\bm{\theta}})} \cdot \normtwo{\frac{\textup{\textrm{d}} {\bm{\theta}}}{\textup{\textrm{d}} t}}.
\end{align*}
Since $(\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*))^{1/2} \le \frac{1}{\sqrt{2\mu}} \normtwosm{\nabla \mathcal{L}({\bm{\theta}})}$,
we have
\begin{align*}
  \frac{\textup{\textrm{d}}}{\textup{\textrm{d}} t} (\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*))^{1/2}
  \le -\frac{\sqrt{2\mu}}{2} \normtwo{\frac{\textup{\textrm{d}} {\bm{\theta}}}{\textup{\textrm{d}} t}}.
\end{align*}
Integrating on both sides proves the following
\[
  \frac{\sqrt{2\mu}}{2}\int_{0}^{T} \normtwo{\frac{\textup{\textrm{d}} {\bm{\theta}}(\tau)}{\textup{\textrm{d}} \tau}} \textup{\textrm{d}} \tau \le
  \sqrt{\mathcal{L}({\bm{\theta}}_0) - \mathcal{L}(\vtheta^*)} = O(\normtwosm{{\bm{\theta}}_0 - \vtheta^*}).
\]
So if $\normtwosm{{\bm{\theta}}_0 - \vtheta^*}$ is small enough, then $T = +\infty$
and ${\bm{\theta}}(t)$ converges to  as $t \to +\infty$.
Moreover, $\normtwosm{{\bm{\theta}}_{\infty} - \vtheta^*} \le \normtwosm{{\bm{\theta}}_{\infty} - {\bm{\theta}}_0} + \normtwosm{{\bm{\theta}}_0 - \vtheta^*} = O(\normtwosm{{\bm{\theta}}_0 - \vtheta^*})$.
\end{proof}

\subsubsection{Spherical Optimization}

\begin{theorem} \label{thm:rank-to-pl-sph}
  Let $\mathcal{L}: \mathbb{R}^D \to \mathbb{R}$ be a $\mathcal{C}^3$-smooth scale-invariant function,
  and $\mathit{\Gamma}$ be a $\mathcal{C}^1$-smooth, $(D_\Gamma - 1)$-dimensional, open submanifold of $\mathbb{S}^{D-1}$,
  where every ${\bm{\theta}} \in \mathit{\Gamma}$ is a local minimizer of $\mathcal{L}$
  on $\mathbb{S}^{D-1}$
  and $\mathrm{rank}(\nabla^2 \mathcal{L}({\bm{\theta}})) = D - D_\Gamma$.
  If ${\mathcal{Z}}$ is a compact subset of $\mathit{\Gamma}$,
  then there exist $\epsilon > 0, \mu > 0$ such that
  $\mathrm{cl}({\mathcal{Z}}^{\epsilon}) \cap \mathit{\Gamma}$ is compact
  and
  $\mathcal{L}$ satisfies $\mu$-PL
  on ${\mathcal{Z}}^{\epsilon} \cap \mathbb{S}^{D-1}$.
\end{theorem}
\begin{proof}
  Let $\mathit{\Gamma}' := \{ \nu {\bm{\theta}} : {\bm{\theta}} \in \mathit{\Gamma}, \nu > 0\}$.
  Then $\mathit{\Gamma}'$ is a $\mathcal{C}^1$-smooth, $D_\Gamma$-dimensional, open submanifold of $\mathbb{R}^D$,
  where every ${\bm{\theta}} \in \mathit{\Gamma}$ is a local minimizer of $\mathcal{L}$ on $\mathbb{R}^D$
  and $\mathrm{rank}(\nabla^2 \mathcal{L}({\bm{\theta}})) = D - D_\Gamma$.
  By \Cref{thm:rank-to-pl-sph},
  there exist $\epsilon > 0, \mu > 0$ such that $\mathcal{L}$ satisfies $\mu$-PL
  on ${\mathcal{Z}}^{\epsilon}$, so it satisfies $\mu$-PL on ${\mathcal{Z}}^{\epsilon} \cap \mathbb{S}^{D-1}$.
\end{proof}

\begin{theorem} \label{thm:pl-converge-near-sph}
  If $\vtheta^*$ is a local minimizer of a $\mathcal{C}^2$-smooth and scale-invariant function $\mathcal{L}: \mathbb{R}^D \setminus \{{\bm{0}}\} \to \mathbb{R}$ and
  $\mathcal{L}$ is $\mu$-PL on an open neighborhood $U$ of $\vtheta^*$ on $\mathbb{S}^{D-1}$,
  then for any ${\bm{\theta}}_0 \in \mathbb{S}^{D-1}$ sufficiently close to $\vtheta^*$,
  a gradient flow $\frac{\textup{\textrm{d}} {\bm{\theta}}}{\textup{\textrm{d}} t} = -\nabla \mathcal{L}({\bm{\theta}})$
  starting with ${\bm{\theta}}_0$ converges to a point ${\bm{\theta}}_{\infty}$ as
  $t \to +\infty$ and $\normtwosm{{\bm{\theta}}_{\infty} - \vtheta^*} = O(\normtwosm{{\bm{\theta}}_0 - \vtheta^*})$.
\end{theorem}
\begin{proof}
  Since $\mathcal{L}$ is $\mu$-PL on $U\subseteq \mathbb{S}^{D-1}$ and scale-invariant, we know $\mathcal{L}$ is $\frac{\mu}{2}$-PL on a open set in $\mathbb{R}^D$,
  $U' = \{{\bm{w}} : \frac{{\bm{w}}}{\norm{{\bm{w}}}_2} \in U, \norm{{\bm{w}}}_2\in [\frac{1}{\sqrt{2}}, \sqrt{2}]\}$.
  The proof is completed by  applying \Cref{thm:pl-converge-near-ful}.
\end{proof}

\section{Supplementary Material for Appendix~\ref{sec:reform-gdwd}}
\label{sec:connection-proof}

\subsection{Proof for Theorem \ref{thm:connect-siwd-rmsprop}} \label{sec:proof-connect-siwd-rmsprop}

\begin{lemma} \label{lm:connection-prelim}
  In the setting of \Cref{thm:connect-siwd-rmsprop},
  \begin{align}
    \normtwosm{{\bm{w}}_{t+1}}^2 &= (1-\hat{\eta} \hat{\lambda})^2 \normtwosm{{\bm{w}}_t}^2 + \frac{\hat{\eta}^2}{\normtwosm{{\bm{w}}_t}^2} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2. \label{eq:connection-w-norm2}
  \end{align}
\end{lemma}
\begin{proof}
Recall that ${\bm{w}}_{t+1} = (1 - \hat{\eta} \hat{\lambda}) {\bm{w}}_t - \hat{\eta} \nabla \mathcal{L}({\bm{w}}_t)$.
By scale-invariance, $\dotpsm{\nabla \mathcal{L}({\bm{w}}_t)}{{\bm{w}}_t} = 0$. Then by Pythagorean theorem (or Gougu Theorem),
\begin{align*}
  \normtwosm{{\bm{w}}_{t+1}}^2 = (1 - \hat{\eta} \hat{\lambda})^2 \normtwosm{{\bm{w}}_t}^2 + \hat{\eta}^2 \normtwosm{\nabla \mathcal{L}({\bm{w}}_t)}^2.
\end{align*}
Since $\nabla \mathcal{L}({\bm{w}}_t) = \frac{1}{\normtwosm{{\bm{w}}_t}} \nabla \mathcal{L}({\bm{\theta}}_t)$ by scale-invariance, we can rewrite
the last term $\hat{\eta}^2 \normtwosm{\nabla \mathcal{L}({\bm{w}}_t)}^2$ as $\frac{\hat{\eta}^2}{\normtwosm{{\bm{w}}_t}^2} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2$,
which implies \eqref{eq:connection-w-norm2}.
\end{proof}

\begin{proof}[Proof for \Cref{thm:connect-siwd-rmsprop}]
Squaring both sides of \eqref{eq:connection-w-norm2}, we have
\begin{align*}
  \normtwosm{{\bm{w}}_{t+1}}^4 &= (1-\hat{\eta} \hat{\lambda})^4 \normtwosm{{\bm{w}}_t}^4 + 2(1-\hat{\eta} \hat{\lambda})^2 \hat{\eta}^2 \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 + \frac{\hat{\eta}^4}{\normtwosm{{\bm{w}}_t}^4} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^4.
\end{align*}
Let $\beta := (1-\hat{\eta}\hat{\lambda})^4$, $\tilde{v}_t := \frac{1}{\tilde{\eta}_t^2} = \frac{(1-\hat{\eta} \hat{\lambda})^2}{\hat{\eta}^2} \normtwosm{{\bm{w}}_t}^4$. Then
\begin{align*}
    \tilde{v}_{t+1} &= \frac{(1-\hat{\eta} \hat{\lambda})^2}{\hat{\eta}^2} \left((1-\hat{\eta} \hat{\lambda})^4 \normtwosm{{\bm{w}}_t}^4 + 2(1-\hat{\eta} \hat{\lambda})^2 \hat{\eta}^2 \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 + \frac{\hat{\eta}^4}{\normtwosm{{\bm{w}}_t}^4} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^4 \right) \\
    &= (1-\hat{\eta} \hat{\lambda})^4 \tilde{v}_t + 2(1-\hat{\eta}\hat{\lambda})^4 \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 + \frac{(1-\hat{\eta}\hat{\lambda})^2\hat{\eta}^2}{\normtwosm{{\bm{w}}_t}^4} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^4 \\
    &= \beta \tilde{v}_t + 2\beta \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 + \frac{\beta}{\tilde{v}_t} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^4,
\end{align*}
where the last equality uses the definition of $\beta$ and $\tilde{v}_t$.

Let $\eta := \sqrt{\frac{1-\beta}{2\beta}}$ and $\bar{g}_t := \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)} / \eta$. Then
\begin{align*}
  \tilde{v}_{t+1} &= \beta \tilde{v}_t + 2\beta \cdot \eta^2 \bar{g}_t^2 + \frac{\beta}{\tilde{v}_t} \cdot \eta^4\bar{g}_t^4 \\
  &= \beta \tilde{v}_t + (1-\beta) \bar{g}_t^2 + \frac{1}{4\beta \tilde{v}_t} (1-\beta)^2 \bar{g}_t^4,
\end{align*}
which is exactly the update rule of GWSI scheduler.
\end{proof}

\subsection{Proof for Theorem \ref{thm:gdwd-is-quasi-rmsprop}} \label{sec:proof-gdwd-is-quasi-rmsprop}

\begin{proof}
We specify the quasi-RMSprop scheduler $\mathcal{H}_{\text{QRMS}}$ as follows.
Let ${\mathcal{P}}_{\mathcal{H}_{\text{QRMS}}} := \{(\eta, \beta) : \beta = 1- 2 \eta^2, \eta \in (0, \frac{1}{\sqrt{2}})\}$.
Given hyperparameters $\eta, \beta$,
we define $\beta' := (1-\frac{1}{4}(1 - \beta))^4$
and $\eta' := \sqrt{(\beta' - 1) / 2}$.
Then $\mathcal{H}_{\text{QRMS}}$ produces the effective LRs as a GWSI scheduler with $(\eta', \beta')$:
\begin{align*}
  \tilde{\eta}_t &\gets \frac{1}{\sqrt{\tilde{v}_t}}, &
  \tilde{v}_{t+1} &\gets \beta' \tilde{v}_t + (1-\beta') \hat{g}_t^2 + \frac{1}{4\beta' \tilde{v}_t} (1-\beta')^2 \hat{g}_t^4, &
  &\text{where} \quad \hat{g}_t := \normtwosm{{\bm{g}}_t} / \eta'.
\end{align*}
When $\eta = \sqrt{2\eta_{\mathrm{in}}}$ and $\beta = 1 - 4\eta_{\mathrm{in}}$,
it is easy to see that $\mathcal{H}_{\text{QRMS}}$ produces the same effective LRs as GD+WD\xspace on scale-invariant functions (\Cref{thm:connect-siwd-rmsprop}).
Now we only need to verify that
$\mathcal{H}_{\text{QRMS}}$ is indeed a quasi-RMSprop scheduler.

When $\beta$ is close enough to $1$, we have $\beta' = \beta + O((1-\beta)^2)$, $\beta' \ge 1/2$, $\eta' = \eta \cdot (1 + O(1-\beta))$.
Let $C_0$ be a constant such that
$\eta / \eta' \le C_0, \abssm{1 - (\eta / \eta')^2} \le C_0(1 - \beta),
1-\beta' \le C_0(1-\beta), \abssm{\beta' - \beta} \le C_0 (1-\beta)^2$.
Let $\bar{g}_t := \normtwosm{{\bm{g}}_t} / \eta$.
Then $\hat{g}_t = (\eta / \eta')\bar{g}_t$, and thus
\begin{align*}
  \hat{g}_t &\le C_0 \bar{g}_t, &
  \abssm{\hat{g}_t^2 - \bar{g}_t^2} &= \abssm{1 - (\eta/\eta')^2} \cdot \bar{g}_t^2 \le C_0 (1-\beta) \bar{g}_t^2.
\end{align*}
We only need to verify that
$\abs{\tilde{v}_{t+1} - \left(\beta \tilde{v}_t + (1-\beta) \bar{g}_t^2 \right)} \le \delta(\tilde{v}_t) \cdot (1-\beta)^2 \cdot P(\bar{g}_t)$
for some continuous function $\delta$ and some polynomial $P$.
\begin{align*}
    & \abs{\tilde{v}_{t+1} - \left(\beta \tilde{v}_t + (1-\beta) \bar{g}_t^2 \right)} \\
    & \qquad \le \abs{\tilde{v}_{t+1} - \left(\beta' \tilde{v}_t + (1-\beta') \hat{g}_t^2 \right)}
    + (1-\beta')\abs{\hat{g}^2_t - \bar{g}_t^2}
    + \abssm{\beta' - \beta} \cdot (\tilde{v}_t + \bar{g}_t^2).
\end{align*}
For the first term,
we have
\[
\abs{\tilde{v}_{t+1} - \left(\beta' \tilde{v}_t + (1-\beta') \hat{g}_t^2 \right)} =
\frac{1}{4 \beta' \tilde{v}_t} (1-\beta')^2 \hat{g}_t^4 \le \frac{C_0^6}{2\tilde{v}_t} \cdot (1-\beta)^2 \cdot \bar{g}_t^4.
\]
For the second and third terms, we have
\begin{align*}
  (1-\beta')\abs{\hat{g}^2_t - \bar{g}_t^2}
  &\le C_0^2(1-\beta)^2 \bar{g}_t^2, &
  \abssm{\beta' - \beta} \cdot (\tilde{v}_t + \bar{g}_t^2)
  &\le C_0(1-\beta)^2 \cdot (\tilde{v}_t + \bar{g}_t^2).
\end{align*}
Finally we can conclude
\begin{align*}
  \abs{\tilde{v}_{t+1} - \left(\beta \tilde{v}_t + (1-\beta) \bar{g}_t^2 \right)}
  &\le
\frac{C_0^6}{2\tilde{v}_t} \cdot (1-\beta)^2 \cdot \bar{g}_t^4
+
  C_0^2(1-\beta)^2 \bar{g}_t^2
+ C_0(1-\beta)^2 \cdot (\tilde{v}_t + \bar{g}_t^2) \\
&\le \left(\frac{C_0^6}{2\tilde{v}_t} + C_0^2 + C_0 (1 + \tilde{v}_t) \right) \cdot (1-\beta)^2 \cdot (\bar{g}_t^4 + \bar{g}_t^2),
\end{align*}
which verifies that $\mathcal{H}_{\text{QRMS}}$ is indeed a quasi-RMSprop scheduler.
\end{proof}

\section{Additional Related Works} \label{sec:add-related}

\myparagraph{SGD Noise Helps to Escape Sharp Minima.}
It has been a folklore that noise in stochastic gradient helps escapes sharp local minima. With the simplification of
assuming the loss is quadratic and treating SGD as its canonical continuous SDE approximation,
\citet{zhu2019anisotropic,xie2021diffusion} showed that anisotropic noise (\emph{e.g.}, the noise covariance is equal to
Hessian) has a better escape efficiency of SGD out of a sharp minimizer,
in comparison with isotropic noise. Under the same
assumptions, \citet{ibayashi2022expescape} proved the exponential escape efficiency without assuming SDE reaches the
stationary distribution. \citet{kleinberg2018alternative} showed that SGD can escape sharp local minima assuming one-point convexity. While all previous escaping analysis of SGD are based on the continuous approximation,
another approach called stability analysis is able to show that SGD cannot converge to sharp local minima when learning
rate is larger than some threshold~\citep{wu2017towards,wu2018howsgd,ma2021linear}.

\myparagraph{Implicit Bias of GD.} There are mainly two types of implicit bias result for GD, where the first type of result applies essentially to the continuous limit of gradient descent, namely gradient flow,
and tolerates error discretization and stochasticity when the learning rate is sufficiently
small~\citep{soudry2018iclrImplicit,soudry2018implicit,lyu2020gradient,ji2020directional,gunasekar2018implicitconv,
gunasekar2018implicit,li2018algorithmic,razin2020implicit,
arora2019implicit,chizat2020implicit,li2021towards,lyu2019gradient,
lyu2021gradient,razin2022implicit,stoger2021small,ge2021understanding}.
The analyses of GD/SGD based on Neural Tangent Kernel (NTK)
also belong to this type~\citep{jacot2018ntk,li2018learning,du2019gradient,arora2019exact,arora2019fine,allenzhu2018convergence,allenzhu2019gobeyond,zou2018stochastic,chizat2019lazy,yang2019scaling}.
This type of result typically relates the
generalization quality to the initialization of GD, and GD in such regimes cannot escape
from bad local minima once reaching there. In contrast, the second type of results, to which the current paper belongs,
fundamentally relies on the discrete nature of gradient descent. For example, \citet{barrett2021implicit} showed that for
small LR, gradient descent is approximately equal to gradient flow minimizing a new objective, \emph{i.e.}, the original
objective plus $\eta$ times squared norm of gradients.
\citet{kong2020multiscale}
studied a special class of ``multiscale'' loss functions,
and they showed that large learning rate
introduces chaos to the dynamics of GD
and provides a mechanism to escape local minima.
\citet{wang2022large} proved that for the matrix factorization problem, GD with a large learning rate has an implicit bias towards
a solution with balanced matrix norms.
Stability analysis~\citep{wu2017towards,wu2018howsgd,mulayoff2021minimastability,ma2021linear} also
belongs to this type.

\myparagraph{Comparison with \citet{arora2022understanding}.} The paper by \citet{arora2022understanding} is probably the most related work among the second type of the implicit bias results.
They assume that there is a smooth function $L$
which satisfies certain regularity conditions around the minimizer manifold (including that Hessian is maximally non-degenerate on the manifold),
and show that running normalized GD on $L$
or GD on $\sqrt{L}$
with sufficiently small LR tracks a deterministic flow
on the minimizer manifold and decreases the largest eigenvalue of $\nabla^2 L$, which has a similar flavor to our result.
However, our setting (scale invariant loss + WD) is more natural and their result and technique do not apply to our
setting because there is no minimizer under their definition, not to mention manifolds. To show the
spherical sharpness decreases, we have to develop new proof techniques,
including connecting the dynamics to a 1-dimensional Hamiltonian system.
Another difference is that our analysis applies to gradient descent directly,
but
the analysis by
\citet{arora2022understanding}
requires injecting stochastic noise to gradients.

\myparagraph{Edge of Stability.}
\citet{cohen2021gradient} provided an extensive empirical study showing that
GD typically occurs at the Edge of Stability (EoS), where the top eigenvalue of Hessian is approximately $2$ / LR
and the descent lemma does not guarantee the loss to decrease.
\citet{ahn2022understanding}
explored the dynamics of GD in the EoS regime
through insightful experiments, and
they attribute the EoS phenomenon to
the lack of flat stationary points near GD trajectory
and the existence of a subset near minima that is forward invariant under GD update;
\citet{ma2022multiscale} proved the EoS phenomenon for a class of loss functions that
are decomposable as a sum of 1-dimensional functions with subquadratic growth.
The aforementioned work by \citet{arora2022understanding}
analyzes normalized GD on $L$ or GD on $\sqrt{L}$ also in the EoS regime,
where the latter case corresponds to a class of loss functions that grow approximately linearly near minima (e.g., the absolute value function $\abssm{x}$).
The evolution of the eigenvalues of Hessian has also been studied for SGD \citep{gilmer2022curvature,jastrzebski2020breakeven,lewkowycz2020large}.
Our work can be seen as a theoretical explanation of the EoS phenomenon
for scale-invariant loss functions with $L^2$-regularization (see also \Cref{sec:eos-cohen}),
which are a more broad and natural class of training loss in deep learning compared with those being studied in previous theoretical analyses.

\section{Proofs for Linear Regression with Batch Normalization} \label{sec:app-lin}

\begin{lemma} \label{lm:lin-bn-H}
    Assume that the regression targets $y_i$ are generated by a linear model.
    For linear regression with BN, the global minimizer manifold of $\mathcal{L}({\bm{w}}) := \frac{1}{n} \sum_{i=1}^{n}(\Phi({\bm{w}}_i; {\bm{w}}, \sigma_{\mathrm{y}}, \mu_{\mathrm{y}}) - y_i)^2$ on the unit sphere is
    \begin{align*}
        \mathit{\Gamma} := \left\{ {\bm{w}} \in \mathbb{S}^{d-1} : \dotp{\frac{{\bm{w}}}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}}}{{\bm{x}}_i - \vmu_{\mathrm{x}}} = \frac{y_i - \mu_{\mathrm{y}}}{\sigma_{\mathrm{y}}} \right\}.
    \end{align*}
    For any global minimizer ${\bm{w}} \in \mathit{\Gamma}$,
    the Hessian matrix ${\bm{H}}({\bm{w}})$ of the loss is given by
    \begin{align*}
        {\bm{H}}({\bm{w}}) = 2\normtwosm{\tilde{\vw}}^2
        \left( \mSigma_{\mathrm{x}} - {\bm{z}} {\bm{z}}^\top \right),
    \end{align*}
    where $\tilde{\vw} := \frac{\sigma_{\mathrm{y}} {\bm{w}}}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}}$ as defined in \eqref{eq:lin-bn-Phi}, and ${\bm{z}} := \frac{1}{n} \sum_{i=1}^{n} \frac{y_i - \mu_{\mathrm{y}}}{\sigma_{\mathrm{y}}} ({\bm{x}}_i - \vmu_{\mathrm{x}})$.
\end{lemma}
\begin{proof}
The model output $\Phi({\bm{x}}; {\bm{w}}, \mu_{\mathrm{y}}, \sigma_{\mathrm{y}})$
can be written as
\[
    \Phi({\bm{x}}; {\bm{w}}, \mu_{\mathrm{y}}, \sigma_{\mathrm{y}}) = \sigma_{\mathrm{y}} \dotp{\frac{{\bm{w}}}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}}}{{\bm{x}}_i - \vmu_{\mathrm{x}}} + \mu_{\mathrm{y}}.
\]
Then it is easy to verify that the global minimizer manifold is $\mathit{\Gamma}$.

Now we compute the Hessian.
    Let $\tilde{\vx}_i := {\bm{x}}_i - \vmu_{\mathrm{x}}$ and $q_i := \frac{y_i - \mu_{\mathrm{y}}}{\sigma_{\mathrm{y}}}$.
Because of the use of squared loss,
on $\mathit{\Gamma}$ the Hessian can be written as
the sum of outer products of gradients:
\begin{align*}
    {\bm{H}}({\bm{w}})
    &= \frac{2}{n} \sum_{i=1}^{n} \nabla_{{\bm{w}}} \Phi({\bm{x}}_i; {\bm{w}}, \sigma_{\mathrm{y}}, \mu_{\mathrm{y}}) \nabla_{{\bm{w}}} \Phi({\bm{x}}_i; {\bm{w}}, \sigma_{\mathrm{y}}, \mu_{\mathrm{y}})^\top.
\end{align*}
For each gradient we have
\begin{align*}
    \nabla_{{\bm{w}}} \Phi({\bm{x}}_i; {\bm{w}}, \sigma_{\mathrm{y}}, \mu_{\mathrm{y}})
    = \frac{\sigma_{\mathrm{y}}}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}} \left( {\bm{I}} - \frac{\mSigma_{\mathrm{x}} {\bm{w}}\vw^\top}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}^2} \right) \tilde{\vx}_i
    = \frac{\sigma_{\mathrm{y}}}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}} \left( {\bm{I}} - \mSigma_{\mathrm{x}} \tilde{\vw} \tilde{\vw}^\top \right) \tilde{\vx}_i
\end{align*}
Note that
    $\mSigma_{\mathrm{x}} \tilde{\vw} = \frac{1}{n} \sum_{i=1}^{n} \tilde{\vx}_i \tilde{\vx}_i^\top \tilde{\vw}
    = \frac{1}{n} \sum_{i=1}^{n} q_i \tilde{\vx}_i =: {\bm{z}}$.
Then we can simplify the gradient by
\begin{align*}
    \nabla_{{\bm{w}}} \Phi({\bm{x}}_i; {\bm{w}}, \sigma_{\mathrm{y}}, \mu_{\mathrm{y}}) &=
    \frac{\sigma_{\mathrm{y}}}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}} \left( \tilde{\vx}_i - q_i {\bm{z}} \right)
\end{align*}
Now we simplify the Hessian.
\begin{align*}
    {\bm{H}}({\bm{w}})
    &= \frac{2}{n} \sum_{i=1}^{n} \nabla_{{\bm{w}}} \Phi({\bm{x}}_i; {\bm{w}}, \sigma_{\mathrm{y}}, \mu_{\mathrm{y}}) \nabla_{{\bm{w}}} \Phi({\bm{x}}_i; {\bm{w}}, \sigma_{\mathrm{y}}, \mu_{\mathrm{y}})^\top \\
    &= \frac{2\sigma_{\mathrm{y}}^2}{n \normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}^2}
    \sum_{i=1}^{n}
    \left( \tilde{\vx}_i - q_i {\bm{z}} \right)\left( \tilde{\vx}_i - q_i {\bm{z}} \right)^\top \\
    &= \frac{2\sigma_{\mathrm{y}}^2}{n \normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}^2}
    \left(
    \sum_{i=1}^{n} \tilde{\vx}_i \tilde{\vx}_i^\top
    - \sum_{i=1}^{n} q_i \tilde{\vx}_i {\bm{z}}^\top
    - \sum_{i=1}^{n} q_i {\bm{z}} \tilde{\vx}_i^\top
    + \sum_{i=1}^{n} q_i^2 {\bm{z}}\vz^\top
    \right) \\
    &= \frac{2\sigma_{\mathrm{y}}^2}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}^2}
    \left(
        \mSigma_{\mathrm{x}}
    - {\bm{z}} {\bm{z}}^\top
    - {\bm{z}} {\bm{z}}^\top
    + {\bm{z}}\vz^\top
    \right) \\
    &= \frac{2\sigma_{\mathrm{y}}^2}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}^2}
    \left(
        \mSigma_{\mathrm{x}}
    - {\bm{z}} {\bm{z}}^\top
    \right).
\end{align*}
We complete the proof by noting that $\frac{\sigma_{\mathrm{y}}^2}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}^2} = \normtwosm{\tilde{\vw}}^2$.
\end{proof}

\begin{proof}[Proof for \Cref{thm:bn-lin-main}]
By \Cref{lm:lin-bn-H},
for ${\bm{w}} \in \mathbb{S}^{D-1}$,
\begin{align*}
    \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\theta}}) = \nabla_{\Gamma} \log \normtwosm{\tilde{\vw}}^2 = \tfrac{1}{\normtwosm{\tilde{\vw}}^2} \nabla_{\Gamma} \normtwosm{\tilde{\vw}}^2.
\end{align*}
By simple calculation,
it can be verified that
the only point that has $\nabla_{\Gamma} \normtwosm{\tilde{\vw}}^2 = {\bm{0}}$
is the unique point $\vw^*$ on $\mathit{\Gamma}$
that is a linear combination of ${\bm{x}}_i - \vmu_{\mathrm{x}}$.
As the shperical sharpness is bounded from below and there is only one stationary point,
the sharpness-reduction flow must converge on $\mathit{\Gamma}$
and the convergence point must be $\vw^*$.

Since $\vw^*$ is a linear combination of ${\bm{x}}_i - \vmu_{\mathrm{x}}$,
the associated $\tilde{\vw}^*$
should be
the least square solution (without bias) of this ``shifted'' dataset: $\{ ({\bm{x}}_i - \vmu_{\mathrm{x}}, y_i - \mu_{\mathrm{y}})\}$.
In other words, $\tilde{\vw}^*$ is the optimal solution of the following constrained optimization problem:
\[
    \min \quad \normtwosm{{\bm{w}}}^2 \quad \text{s.t.} \quad {\bm{w}}^\top ({\bm{x}}_i - \vmu_{\mathrm{x}}) = y_i - \mu_{\mathrm{y}}, \quad \forall i \in [n].
\]
When ${\bm{w}}$ is given in the above optimization problem,
there is only a unique $b$ such that ${\bm{w}}^\top {\bm{x}}_i + b = y_i$.
So we can introduce a bias to this problem without changing the optimal solution:
\[
    \min \quad \normtwosm{{\bm{w}}}^2 \quad \text{s.t.} \quad {\bm{w}}^\top {\bm{x}}_i + b = y_i, \quad \forall i \in [n].
\]
One can also easily see that this $b$ must match with $\tilde{b}$.
Therefore, we can conclude that the sharpness-reduction flow \eqref{eq:main-zeta}
finds to the optimal solution of \eqref{eq:M} at convergence.
\end{proof}

\section{Supplementary Material for Section~\ref{sec:main-result-1}}
\label{sec:app-enter-eos}

\subsection{Proof for Descent Lemma}

\begin{proof}[Proof for \Cref{lm:descent}]
By Taylor expansion,
\begin{align*}
    \mathcal{L}({\bm{\theta}}_{t+1}) = \mathcal{L}({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t))
    &\le \mathcal{L}({\bm{\theta}}_t) - \dotpsm{\nabla \mathcal{L}({\bm{\theta}}_t)}{\tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)}
    + \frac{1}{2}\lambda_{\max}^{(t)} \normtwosm{\tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)}^2 \\
    &= \mathcal{L}({\bm{\theta}}_t) - \tilde{\eta}_t (1 - \tilde{\eta}_t \lambda_{\max}^{(t)} / 2) \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2,
\end{align*}
which proves the lemma.
\end{proof}

\subsection{Proof for Theorem \ref{thm:eos-near-opt}: GD Eventually Enters the EoS Regime}

Let $R := C_0 (\hat{\lambda} \hat{\eta})^{1/2}$ be a radius so that
$\vtheta^*$ is a minimizer of $\mathcal{L}$ on $U := \Ball{R}{{\bm{\theta}}} \cap \mathbb{S}^{D-1}$
and
$\mu$-PL holds within $U$,
where $C_0$ is a large constant to be specified later.
Let $\lambda_{\max} := \sup\{ \lambda^{\mathrm{H}}_1({\bm{\theta}}) : {\bm{\theta}} \in U \}$.
Then we know that $\lambda_{\max} = \lambda^{\mathrm{H}}_1(\vtheta^*) + O(\hat{\lambda} \hat{\eta})$.
Let $T_0$ be the largest number so that ${\bm{\theta}}_t \in U$ for all $t_0 \le t \le T_0$,
We define a potential function $\Psi({\bm{\theta}}) := \sqrt{\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*)}$.
\begin{lemma} \label{lm:Psi-decay}
If ${\bm{\theta}}_t \in U$ and $\tilde{\eta}_t < \frac{2}{\lambda_{\max}}$ for some $t_0 \le t < T_0$, then
\begin{align*}
    \Psi({\bm{\theta}}_t) - \Psi({\bm{\theta}}_{t+1})
    &\ge \frac{\sqrt{2\mu}}{2}(1 - \tilde{\eta}_t \lambda_{\max} / 2) \tilde{\eta}_t \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}.
\end{align*}
\end{lemma}
\begin{proof}
By descent lemma,
\begin{align*}
    \mathcal{L}({\bm{\theta}}_{t+1})
    &\le \mathcal{L}({\bm{\theta}}_t) - \tilde{\eta}_t (1 - \tilde{\eta}_t \lambda_{\max} / 2) \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2.
\end{align*}
Then
\begin{align*}
    \Psi({\bm{\theta}}_t) - \Psi({\bm{\theta}}_{t+1}) = \frac{\mathcal{L}({\bm{\theta}}_t) - \mathcal{L}({\bm{\theta}}_{t+1})}{\Psi({\bm{\theta}}_t) + \Psi({\bm{\theta}}_{t+1})}
    &\ge \frac{(1 - \tilde{\eta}_t \lambda_{\max} / 2) \tilde{\eta}_t \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2}{2 \Psi({\bm{\theta}}_t)}.
\end{align*}
By $\mu$-PL, $\normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)} \ge \sqrt{2\mu} \cdot \Psi({\bm{\theta}}_t)$. Combining these together proves the lemma.
\end{proof}

\begin{lemma} \label{lm:anti-pl}
There exists $C_2 = O(1)$ such that $\normtwosm{\nabla \mathcal{L}({\bm{\theta}})}^2 \le C_2 (\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*))$ for all ${\bm{\theta}} \in U$.
\end{lemma}
\begin{proof}
It is equivalent to give an upper bound for $\sup_{{\bm{\theta}} \in U} \{ G({\bm{\theta}}) \}$, where $G({\bm{\theta}}) := \frac{\normtwosm{\nabla \mathcal{L}({\bm{\theta}})}^2}{\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*)}$.
Since $\mathcal{L} \in \mathcal{C}^2$, $G$ is continuous in its domain.
So it suffices to upper bound $G({\bm{\theta}})$ around every singular point, i.e., around every minimizer of $\mathcal{L}$ on $U$.
And for every minimizer ${\bm{\theta}}' \in U$, we can do Taylor expansions for $\mathcal{L}$ and $\nabla \mathcal{L}$
to show that $G({\bm{\theta}})$ is indeed bounded by $O(1)$ around ${\bm{\theta}}'$.
\end{proof}

\begin{proof}[Proof for \Cref{thm:eos-near-opt}]
Recall that $\tilde{\eta}_{t_0} \le \frac{2}{\rho_2} < \frac{2}{\lambda^{\mathrm{H}}_1(\vtheta^*)}$.
Let $\delta := 2 - \tilde{\eta}_{t_0} \lambda_{\max} \in (0, 2)$
and $T_1$ be the largest number so that $\tilde{\eta}_t \le \frac{2 - \delta / 4}{\lambda_{\max}}$ for all $t_0 \le t \le T_1$.
By \Cref{lm:Psi-decay}, for all $t_0 \le t < \min\{T_0, T_1\}$,
\begin{align*}
    \Psi({\bm{\theta}}_t) - \Psi({\bm{\theta}}_{t+1})
    &\ge \frac{\sqrt{2\mu}}{16} \delta \tilde{\eta}_t \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}.
\end{align*}
Telescoping the sum we have $\frac{\sqrt{2\mu}}{16} \delta \sum_{\tau={t_0}}^{t-1}\tilde{\eta}_{\tau} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_\tau)} \le \Psi({\bm{\theta}}_0)$.
By smoothness of $\mathcal{L}$, $\Psi({\bm{\theta}}_{t_0}) = O(\normtwosm{{\bm{\theta}}_{t_0} - \vtheta^*}) = O((\hat{\lambda} \hat{\eta})^{1/2})$.
So for all $t_0 \le t \le \min\{T_0, T_1\}$,
\[
    \normtwosm{{\bm{\theta}}_t -{\bm{\theta}}_{t_0}} \le
    \sum_{\tau={t_0}}^{t-1}\tilde{\eta}_{\tau} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_{\tau})} = O((\hat{\lambda}\hat{\eta})^{1/2}),
\]
which implies that $T_0 > T_1$ or $T_0 = T_1 = +\infty$ if we choose $C_0$ to be large enough.

By \Cref{thm:connect-siwd-rmsprop},
$\tilde{\eta}_t$ can be seen as the output of a GWSI scheduler with $\beta = (1-\hat{\lambda}\hat{\eta})^4 = 1 - \Theta(\hat{\lambda}\hat{\eta})$ and $\eta = \sqrt{(\beta^{-1}-1)/2} = O((\hat{\lambda}\hat{\eta})^{1/2})$.
Then by the update rule,
\begin{align} \label{eq:eeta-aaa}
    \tilde{\eta}_{t+1}^{-2} = \beta \tilde{\eta}_t^{-2} + (1-\beta) \bar{g}_t^2 + \frac{\tilde{\eta}_t^2}{4\beta} (1-\beta)^2 \bar{g}_t^4,
\qquad \text{where} \qquad \bar{g}_t := \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)} / \eta.
\end{align}
For ${\bm{\theta}} \in U$,
we have $\normtwosm{\nabla \mathcal{L}({\bm{\theta}})}^2 \le C_2(\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*))$
by
\Cref{lm:anti-pl}.
So for all $t_0 \le t \le T_0$, we have
\[
    \bar{g}_t = \frac{1}{\eta} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}
    \le \frac{C_2}{\eta} (\mathcal{L}({\bm{\theta}}_t) - \mathcal{L}(\vtheta^*))
    \le \frac{C_2}{\eta} (\mathcal{L}({\bm{\theta}}_0) - \mathcal{L}(\vtheta^*))
    =O(1).
\]
Then
$\tilde{\eta}_{t+1}^{-2} = \beta \tilde{\eta}_t^{-2} + (1-\beta) \cdot O(1)$,
which implies $\tilde{\eta}_t \ge \Omega(1)$ for some $t = t_0 + O( \frac{1}{1 - \beta} \log(\tilde{\eta}_0^{-2}))$.
Therefore,
we can infer that
it must hold for some steps $t$
that
$\tilde{\eta}_t \in \left[ c_{\min}, \frac{2 - \delta / 4}{\lambda_{\max}} \right]$,
where $c_{\min}$ is some constant.

As $\bar{g}_t \ge 0$,
the update rule \eqref{eq:eeta-aaa} also implies
$\tilde{\eta}_{t+1}^{-2} \ge \beta \tilde{\eta}_t^{-2}$,
or equivalently
$\tilde{\eta}_{t+1} \le (1-\hat{\eta}\hat{\lambda})^{-2} \tilde{\eta}_t$.
This suggests that the number of steps such that
$\tilde{\eta}_t \in \left[ \frac{2 - c}{\lambda_{\max}}, \frac{2 - \delta / 4}{\lambda_{\max}} \right]$,
is at least $\Omega(\sfrac{1}{(\hat{\eta} \hat{\lambda})})$.
When $\tilde{\eta}_t$ does lie in this range, by \Cref{lm:descent} we have
\begin{align*}
    \mathcal{L}({\bm{\theta}}_{t+1}) \le \mathcal{L}({\bm{\theta}}_t) - \frac{1}{8} \delta \tilde{\eta}_t \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2
\end{align*}
Combining with $\mu$-PL gives
\begin{align*}
    \mathcal{L}({\bm{\theta}}_{t+1}) - \mathcal{L}(\vtheta^*)
    &\le (1 - \mu \delta \tilde{\eta}_t / 4) \cdot (\mathcal{L}({\bm{\theta}}_t) - \mathcal{L}(\vtheta^*)) \\
    &\le \left(1 - \mu \delta c_{\min} / 4\right) \cdot (\mathcal{L}({\bm{\theta}}_t) - \mathcal{L}(\vtheta^*)).
\end{align*}
Thus the loss decays by a constant factor in every step (the factor is in $(0, 1)$ as we can choose $c_{\min}$ as small as we want).
As this process lasts for at least $\Omega(\sfrac{1}{(\hat{\eta} \hat{\lambda})})$ steps,
the loss first decreases to $\mathcal{L}(\vtheta^*) + O((\hat{\eta} \hat{\lambda})^{10})$
after $O(\log \frac{1}{\hat{\eta} \hat{\lambda}})$ steps,
then it stays small until $T_1$.

Now we show that $T_1$ is finite.
By \Cref{lm:anti-pl}
and \eqref{eq:eeta-aaa},
when the loss is $\mathcal{L}(\vtheta^*) + O((\hat{\eta} \hat{\lambda})^{10})$
the effective LR steadily grows as $\tilde{\eta}_{t+1} = (1-\hat{\eta}\hat{\lambda})^{-2} \tilde{\eta}_t + o(1)$.
So at some step $t$, it must hold that $\tilde{\eta}_t > \frac{2-\delta/4}{\lambda_{\max}}$, which proves $T_1 < +\infty$.

Let $C_1$ be a large constant,
and $T_2$ be the largest number so that $\tilde{\eta}_t < \frac{2 - 2 C_1 (\hat{\eta} \hat{\lambda})^{1/2}}{\lambda_{\max}}$ for all $t_0 \le t \le T_2$.
By \Cref{lm:Psi-decay}, for all $T_1 \le t < \min\{T_2, T_0\}$,
\[
    \Psi({\bm{\theta}}_t) - \Psi({\bm{\theta}}_{t+1}) \ge \frac{\sqrt{2\mu}}{2} C_1 (\hat{\eta} \hat{\lambda})^{1/2} \tilde{\eta}_t \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}.
\]
Telescoping the sum gives $
(\hat{\eta} \hat{\lambda})^{1/2} \sum_{\tau = T_1}^{t-1} \tilde{\eta}_{\tau} \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_{\tau})} \le O(\Psi({\bm{\theta}}_{T_1})) \le O((\hat{\eta} \hat{\lambda})^5)$,
where the last inequality
is due to $\mathcal{L}({\bm{\theta}}_{T_1}) = \mathcal{L}(\vtheta^*) + O((\hat{\eta} \hat{\lambda})^{10})$.
This shows that $\normtwosm{{\bm{\theta}}_t - {\bm{\theta}}_{T_1}} = O((\hat{\eta}\hat{\lambda})^{4.5})$,
and thus $\normtwosm{{\bm{\theta}}_t - \vtheta^*} = O((\hat{\eta}\hat{\lambda})^{1/2})$ by triangle inequality.
Now we have $T_0 > T_2$ or $T_0 = T_2 = +\infty$ when $C_0$ is chosen to be large enough.
We can finish the proof with a similar argument as for $T_1$ to show that $T_2$ cannot be infinite either.
\end{proof}

\subsection{Connection to the EoS Regime in Cohen et al.'s Definition} \label{sec:eos-cohen}

\citet{cohen2021gradient}
studied the dynamics of GD (without weight decay), i.e., when the loss is $\tilde{\Loss}$,
the update rule is
\begin{align*}
    {\bm{w}}_{t+1} \gets {\bm{w}}_t - \hat{\eta} \nabla \tilde{\Loss}({\bm{w}}_t).
\end{align*}
They define the EoS regime as a regime in which (1) $\lambda_1(\nabla^2\tilde{\Loss}({\bm{w}}_t))$ hovers right at,
or just above $2 / \hat{\eta}$; and (2) the training loss $\tilde{\Loss}({\bm{w}}_t)$ goes up and down over short timescales,
yet still decreases in the long-term run.

For a gradient-based method in general,
$2 / \hat{\eta}$
should be replaced to the maximum sharpness bound that the loss function is guaranteed to decrease through Taylor expansions,
e.g., \citet{cohen2021gradient} derived the bounds exactly for Polyak and Nesterov momentum in Appendix B of their paper.
In our definition,
we view GD+WD\xspace on scale-invariant loss as PGD on $\mathbb{S}^{D-1}$,
and thus we define the EoS regime as the regime where $\tilde{\eta}_t \approx 2 / \lambda_{\max}^{(t)}$, where $\lambda_{\max}^{(t)}$ is the local upper bound of spherical sharpness (\Cref{lm:descent}).
This captures the first condition of \citet{cohen2021gradient}'s definition.
It can then be shown by descent lemma that the loss goes up and down slightly in each step,
but we do not require the loss to decrease in the long-term run because
it is hard to define such a long-term trend in theory.
Also, the loss may not have such a trend in our experiments (see, e.g., \Cref{fig:lin-idea}).

Another view is to write GD+WD\xspace on scale-invariant loss as GD on scale-invariant loss with $L^2$-regularization,
i.e., GD on
$\tilde{\Loss}({\bm{w}}) := \mathcal{L}({\bm{w}}) + (\hat{\lambda}/2) \normtwosm{{\bm{w}}}^2$.
Now we show that $\tilde{\eta}_t \approx 2 / \lambda_{\max}^{(t)}$  is essentially the same as the first condition of \citet{cohen2021gradient}'s definition for small WD $\hat{\lambda}$ and small gradient.
When the gradient is small, ${\bm{\theta}}_t$ does not move far in one step, then $\lambda_{\max}^{(t)} \approx \lambda_1(\nabla^2 \mathcal{L}({\bm{\theta}}_t))$.
By scale-invariance, $\lambda_1(\nabla^2 \mathcal{L}({\bm{\theta}}_t)) = \normtwosm{{\bm{w}}_t}^2 \cdot \lambda_1(\nabla^2 \mathcal{L}({\bm{w}}_t))$ (\Cref{lm:scale-invariant-grad-hess}).
So the condition $\tilde{\eta}_t \approx 2 / \lambda_{\max}^{(t)}$ implies that $\tilde{\eta}_t \normtwosm{{\bm{w}}_t}^2 \approx \frac{2}{\lambda_1(\nabla^2 \mathcal{L}({\bm{\theta}}_t))}$.
Recall that $\tilde{\eta}_t := \frac{\hat{\eta}}{(1 - \hat{\eta}\hat{\lambda}) \normtwosm{{\bm{w}}_t}^2}$. Then we have
\begin{align} \label{eq:eos-near-cohen}
    \frac{\hat{\eta}}{1 - \hat{\eta} \hat{\lambda}} \approx \frac{2}{\lambda_1(\nabla^2 \mathcal{L}({\bm{\theta}}_t))}.
\end{align}
Note that  $\lambda_1(\nabla^2\tilde{\Loss}({\bm{w}}_t)) = \lambda_1(\nabla^2 \mathcal{L}({\bm{w}}_t)) + \hat{\lambda}$.
When $\hat{\lambda}$ is small, we can approximate the LHS of \eqref{eq:eos-near-cohen} as $\hat{\eta}$
and the RHS as $\frac{2}{\lambda_1(\nabla^2 \tilde{\Loss}({\bm{w}}_t))}$,
so $\lambda_1(\nabla^2 \tilde{\Loss}({\bm{w}}_t))$ is around $2 / \hat{\eta}$.

This gives an alternative view showing that our definition corresponds to the first condition in \citet{cohen2021gradient}'s definition.
Again, the second condition is not a part of our definition.
However, our main theorem on sharpness-reduction bias
actually implies that $\tilde{\Loss}$ decreases in the long-term run.
More specifically,
note that $\normtwosm{{\bm{w}}_t}^2 \approx \hat{\eta}/\tilde{\eta}_t \approx \frac{1}{2}\hat{\eta} \lambda_1(\nabla^2\mathcal{L}({\bm{\theta}}_t))$ in the EoS regime.
Then $\normtwosm{{\bm{w}}_t}^2$ is decreasing since GD+WD\xspace has a sharpness-reduction bias (\Cref{thm:gdwd-main}).
Therefore, even though $\mathcal{L}({\bm{\theta}}_t)$ may not be decreasing in the long-term run,
the regularized loss $\tilde{\Loss}({\bm{w}}_t)$ can still have a tendency to decrease as the norm drops.
\section{Case Study: Linear Regression with Batch Normalization} \label{sec:lin}
In this section, we analyze the GD+WD\xspace dynamics on linear regression with Batch Normalization (BN),
as a simple application of our theory.
Let $\{({\bm{x}}_i, y_i)\}_{i=1}^{n}$ be a dataset, where ${\bm{x}}_i \in \mathbb{R}^d$ and $y_i \in \mathbb{R}$ are inputs and regression targets.
We study the over-parameterized case where $d \gg n$,
and we assume that the regression targets are generated by an unknown linear model.

A classic linear model is parameterized by $({\bm{w}}, b) \in \mathbb{R}^d \times \mathbb{R}$ and outputs ${\bm{w}}^\top {\bm{x}} + b$ given input ${\bm{x}}$,
but now we add a BN to the output.
More specifically, we consider a batch-normalized linear model $\Phi({\bm{x}}; {\bm{w}}, \gamma, \beta) := \gamma \cdot \frac{{\bm{w}}^\top {\bm{x}} - \mu_1}{\sigma_1} + \beta$,
where $\mu_1, \sigma_1$ are the mean and variance of $\{{\bm{w}}^\top {\bm{x}}_i\}_{i=1}^{n}$ over the whole dataset\footnote{Note that the batch size is $n$ here as we are running full-batch GD},
and the bias term $b$ is cancelled out due to BN.
Note that $\Phi({\bm{x}}; {\bm{w}}, \gamma, \beta)$ is still a linear function with respect to ${\bm{x}}$.
Let $\vmu_{\mathrm{x}} \in \mathbb{R}^d$ and $\mSigma_{\mathrm{x}} \in \mathbb{R}^{d \times d}$
be the mean and covariance of the input data $\{{\bm{x}}_i\}_{i=1}^{n}$.
Then $\Phi({\bm{x}}; {\bm{w}}, \gamma, \beta)$ can be rewritten as:
\begin{equation} \label{eq:lin-bn-Phi}
    \Phi({\bm{x}}; {\bm{w}}, \gamma, \beta) = \tilde{\vw}^\top {\bm{x}} + \tilde{b},
    \qquad \text{where} \quad \tilde{\vw} := \sfrac{\gamma {\bm{w}}}{\normsm{{\bm{w}}}_{\mSigma_{\mathrm{x}}}}, \quad \tilde{b} := \beta - \tilde{\vw}^\top \vmu_{\mathrm{x}}.
\end{equation}
No matter how ${\bm{w}}$ is set, the output mean and variance of $\Phi$ are always $\beta$ and $\gamma^2$.
To simplify our analysis, we fix $\beta, \gamma$
to be non-trainable constants so that the mean and variance of $\Phi$'s output match with those of $\{y_i\}_{i=1}^{n}$,
that is, we set $\beta = \mu_{\mathrm{y}}$ and $\gamma = \sigma_{\mathrm{y}}$ to be the mean and standard deviation
of $y_i$ over the whole dataset.
Then the training loss is $\mathcal{L}({\bm{w}}) := \frac{1}{n} \sum_{i \in [n]} (\Phi({\bm{x}}_i; {\bm{w}}, \gamma, \beta) - y_i)^2$.

\begin{theorem} \label{thm:bn-lin-main}
    In our setting of linear regression with BN,
    the sharpness-reduction flow ${\bm{\zeta}}$ defined in \eqref{eq:main-zeta}
    converges to the solution $\vw^* \in \mathbb{S}^{d-1}$
    that minimizes sharpness $\lambda^{\mathrm{H}}_1(\vw^*)$ on $\mathit{\Gamma}$, regardless of the initialization.
    Moreover, the coefficients $(\tilde{\vw}, \tilde{b})$ associated with $\vw^*$ (defined in \eqref{eq:lin-bn-Phi})
    are the optimal solution of the following constrained optimization problem \eqref{eq:M}:
    \begin{equation}
        \min \quad \normtwosm{{\bm{w}}}^2 \quad \text{s.t.} \quad {\bm{w}}^\top {\bm{x}}_i + b = y_i, \quad \forall i \in [n]. \tag{M} \label{eq:M}
    \end{equation}
\end{theorem}
At first sight the result may appear trivial because  the intent of WD is to regularize
$L^2$-norm. But this is deceptive because in scale-invariant nets
WD  changes from an explicit regularizer as an implicit one.
This also challenges conventional view of
optimization. GD is usually viewed as a discretization of its continuous counterpart, gradient flow (GF), and
theoretical insight for the discrete update including convergence rate and implicit bias is achieved by analyzing the
continuous counterpart (See \Cref{sec:add-related} for a list). However,
GF does not have the same sharpness-reduction bias
as GD.
As discussed in \citep{li2020exp}, adding WD only performs a time-rescaling on the GF trajectory on scale-invariant loss,
but  does not change the point that GF converge to if we project the trajectory onto the unit sphere.
One can easily show that GF may converge to any zero-loss solution,
but
no matter how small LR is,
GD exhibits the sharpness-reduction bias towards the optimal solution of \eqref{eq:M}.
To our best knowledge, this result is the  first concrete example where even with arbitrarily small LR, GD can still generalize better than GF under natural settings.

\section*{Checklist}

\begin{enumerate}

\item For all authors...
\begin{enumerate}
  \item Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
    \answerYes{}
  \item Did you describe the limitations of your work?
    \answerYes{See \Cref{sec:main-result-2-assumption,sec:con}.}
  \item Did you discuss any potential negative societal impacts of your work?
    \answerNA{We are basically a theoretical work studying the generalization myestery in deep learning. We do not see any negative societal impact.}
  \item Have you read the ethics review guidelines and ensured that your paper conforms to them?
    \answerYes{}
\end{enumerate}

\item If you are including theoretical results...
\begin{enumerate}
  \item Did you state the full set of assumptions of all theoretical results?
    \answerYes{The assumptions of our main theorem for GD+WD\xspace on scale-invariant loss are stated in \Cref{sec:main-result-2-assumption}. For GD/PGD with Quasi-RMSprop schedulers in general, see \Cref{sec:quasi-main-thm}.}
        \item Did you include complete proofs of all theoretical results?
    \answerYes{See \Cref{sec:app-enter-eos,sec:proof-outline-main} for the proofs for our main theorems, \Cref{sec:app-lin} for the proof for the linear example.}
\end{enumerate}

\item If you ran experiments...
\begin{enumerate}
  \item Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
    \answerYes{See our supplementary material.}
  \item Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
    \answerYes{See \Cref{sec:exp-details}.}
        \item Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
    \answerNo{Most of our experiments have only run once due to computational constraints, but we verified the sharpness-reduction bias across various settings.
    }
        \item Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
    \answerYes{See \Cref{sec:exp-details}.}
\end{enumerate}

\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
\begin{enumerate}
  \item If your work uses existing assets, did you cite the creators?
    \answerNA{}
  \item Did you mention the license of the assets?
    \answerNA{}
  \item Did you include any new assets either in the supplemental material or as a URL?
    \answerNA{}
  \item Did you discuss whether and how consent was obtained from people whose data you're using/curating?
    \answerNA{}
  \item Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
    \answerNA{}
\end{enumerate}

\item If you used crowdsourcing or conducted research with human subjects...
\begin{enumerate}
  \item Did you include the full text of instructions given to participants and screenshots, if applicable?
    \answerNA{}
  \item Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
    \answerNA{}
  \item Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
    \answerNA{}
\end{enumerate}

\end{enumerate}

\section{A General Theory for a Broader Class of Adaptive Gradient Methods} \label{sec:qrms}

Our main theorem for GD+WD\xspace on scale-invariant loss, \Cref{thm:gdwd-main},
is actually a corollary of a more general theorem that holds
for any PGD on $\mathbb{S}^{D-1}$
with effective learning rates changing adaptively
according a specific kind of update rules.
We name this kind of update rules as \textit{Quasi-RMSprop Scheduler},
where the name is due to its similarity to RMSprop~\citep{hinton2012rmsprop}.
One of the main reasons that we define this concept is that proving a theorem
for quasi-RMSprop schedulers in general is no harder than proving that only for GD+WD\xspace on scale-invariant loss,
and sometimes the math involved is more simple and elegant when we analyze the dynamics at a higher level through quasi-RMSprop.

In this section, we introduce this key notion, Quasi-RMSprop scheduler.
To motivate it, we first recall the update
rule of RMSprop and present one of its variants with a single scalar learning rate.
Then we formally introduce the notion of quasi-RMSprop schedulers,
which is a class of rules for setting effective learning rates similarly as RMSprop.
Then we
prove a recursive formula for the effective learning rates when GD+WD\xspace optimizes a scale-invariant loss,
and we categorize it as an instance of quasi-RMSprop scheduler.

After introducing this key notion, we then present our general theorem
that holds for any PGD, or even GD, as long as the effective learning rate is set by a quasi-RMSprop scheduler in each step.
Examples include Scalar RMSprop (\Cref{cor:scalar-rmsprop})
and GD+WD\xspace on scale-invariant loss (\Cref{thm:gdwd-main}).

\subsection{Scalar RMSprop and Quasi-RMSprop Scheduler}
The usual RMSprop algorithm maintains a vector ${\bm{v}}_t$ storing the moving
average of the squared gradients for every coordinate, i.e., ${\bm{v}}_{t+1} \gets \beta {\bm{v}}_t + (1-\beta) (\nabla \mathcal{L}({\bm{\theta}}_t))^{\odot 2}$,
where ${\bm{g}}^{\odot 2}$ stands for the vector obtained by squaring ${\bm{g}}$ coordinatewise.
When updating the training parameter, RMSprop divides the usual GD update by the square root of the moving average coordinatewise,
i.e., ${\bm{\theta}}_{t+1} \gets {\bm{\theta}}_t - \frac{\eta}{\sqrt{{\bm{v}}_{t+1} + \epsilon}} \odot \nabla \mathcal{L}({\bm{\theta}}_t)$,
where
$\epsilon$ is a small constant to avoid division by zero.
Here all the addition, division, square root operations are coordinatewise,
and $\odot$ stands for coordinatewise multiplication.

Now we
consider a variant of RMSprop, which we call \textit{Scalar RMSprop}, where
the moving average $v_t$ is maintained as a scalar storing the moving average of the squared norm of gradients,
rather than a vector that stores the moving averages separately for each coordinate.

\begin{definition}[Scalar RMSprop, Standard Form] \label{def:scalar-rms-std}
\textit{Scalar RMSprop} is an iterative method with the following update rule:
\begin{align}
  {\bm{\theta}}_{t+1} &\gets {\bm{\theta}}_t - \frac{\eta}{\sqrt{v_t}} \nabla \mathcal{L}({\bm{\theta}}_t), \label{eq:scalar-rmsprop-theta} \\
  v_{t+1} &\gets \beta v_t + (1-\beta) \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2. \label{eq:scalar-rmsprop-v}
\end{align}
\end{definition}

Besides that $v_t$ is changed from a vector to a scalar, another difference is that the gradient is divided by $\sqrt{v_t}$ in \eqref{eq:scalar-rmsprop-theta},
while in the usual RMSprop it is $\sqrt{{\bm{v}}_{t+1}}$. In fact,
our later analysis also applies if $\sqrt{v_t}$ is changed to $\sqrt{v_{t+1}}$, but the version with $\sqrt{v_t}$
leads to simpler math.

An alternative view of Scalar RMSprop is to regard it as
GD with time-varying learning rate $\tilde{\eta}_t$,
that is,
${\bm{\theta}}_{t+1} \gets {\bm{\theta}}_{t} - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)$,
where $\tilde{\eta}_t$ is the learning rate being used at the $t$-th step, which we call the \textit{effective learning rate}
at step $t$.

We view the effective learning rate here as the output of a \textit{learning rate scheduler}.\footnote{The scheduler here is similar to \texttt{torch.optim.lr\_scheduler} in PyTorch, see \url{https://pytorch.org/docs/stable/optim.html\#how-to-adjust-learning-rate}}
In this view, we call the learning rate scheduler for Scalar RMSprop as \textit{RMSprop scheduler}.
\begin{definition}[Gradient-Based Learning Rate Scheduler]
  A \emph{gradient-based learning rate scheduler} is an algorithm $\mathcal{H}$
  that reads from a stream of vectors ${\bm{g}}_0, {\bm{g}}_1, {\bm{g}}_2, \dots$,
  For all $t \ge 0$, $\mathcal{H}$ outputs a real number $\tilde{\eta}_t := \mathcal{H}({\bm{g}}_0, {\bm{g}}_1, \dots, {\bm{g}}_t)$
  as soon as $\mathcal{H}$ finishes reading the first $t+1$ vectors.
  At each step of gradient descent equipped with a gradient-based learning rate scheduler $\mathcal{H}$,
  the gradients in training are revealed one by one as an input stream to $\mathcal{H}$,
  and $\tilde{\eta}_t$ produced by $\mathcal{H}$ is used as the effective learning rate at step $t$.
\end{definition}
\begin{definition}[RMSprop Scheduler] \label{def:rms-sche}
  Given a constant $\tilde{v}_0 > 0$, a base learning rate $\eta$ and a decay rate $\beta$ as hyperparameters,
  an RMSprop scheduler $\mathcal{H}_{\text{RMS}}$ is a gradient-based learning rate scheduler that
  reads from a stream of vectors ${\bm{g}}_0, {\bm{g}}_1, \dots$,
  and generates effective learning rates $\tilde{\eta}_t := \mathcal{H}_{\text{RMS}}({\bm{g}}_0, \dots, {\bm{g}}_t)$ for all $t \ge 0$
  according to the following recursion:
  \begin{align*}
    \tilde{\eta}_t &\gets \frac{1}{\sqrt{\tilde{v}_t}}, &
    \tilde{v}_{t+1} &\gets \beta \tilde{v}_t + (1-\beta) \bar{g}_t^2, &
    &\text{where} \quad \bar{g}_t := \normtwosm{{\bm{g}}_t} / \eta.
  \end{align*}
\end{definition}
\begin{definition}[Scalar RMSprop, Alternative Form] \label{def:scalar-rms-alt}
  \textit{Scalar RMSprop} is gradient descent with a RMSprop scheduler $\mathcal{H}_{\text{RMS}}$.
  \begin{align}
    {\bm{\theta}}_{t+1} &\gets {\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t),  \qquad \text{where} \qquad \tilde{\eta}_t := \mathcal{H}_{\text{RMS}}(\nabla \mathcal{L}({\bm{\theta}}_0), \dots, \nabla \mathcal{L}({\bm{\theta}}_t)).
  \end{align}
\end{definition}
Note that we use $\tilde{v}_t$ as an internal state of $\mathcal{H}_{\text{RMS}}$ in \Cref{def:rms-sche}.
It is easy to verify that $\tilde{v}_t$ is nothing but a reparameterization of $v_t$: setting $\tilde{v}_t = v_t / \eta^2$ in \Cref{def:rms-sche}
recovers the update rule in \Cref{def:scalar-rms-std}.

Now we introduce the notion of quasi-RMSprop scheduler,
which is a class of gradient-based learning rate schedulers
that have update rules similar to RMSprop scheduler.
At first reading, one can just ignore the details and regard quasi-RMSprop scheduler
as an RMSprop scheduler with negligible perturbations when $\beta \to 1$.
\begin{definition}[Quasi-RMSprop Scheduler] \label{def:qrms-sche}
  A \textit{quasi-RMSprop scheduler} $\mathcal{H}$ is a gradient-based learning rate scheduler
  parameterized by a base learning rate $\eta$ and a decay rate $\beta$,
  satisfying the following properties:
  \begin{enumerate}
  \item $(\eta, \beta)$ is allowed to take value from a hyperparameter space ${\mathcal{P}}_\mathcal{H} \subseteq (0, +\infty) \times (0, 1)$.
  \item There exist thresholds $\eta_{\max}, \beta_{\min}$,
  a continuous function $\delta: (0, +\infty) \to \mathbb{R}$,
  and a polynomial $P: \mathbb{R} \to \mathbb{R}$ such that the following holds.
  If $(\eta, \beta) \in {\mathcal{P}}_{\mathcal{H}}$ and $\eta < \eta_{\max}, \beta > \beta_{\min}$,
  for any input stream ${\bm{g}}_0, {\bm{g}}_1, {\bm{g}}_2, \dots$,
  there exists a sequence of positive real numbers $\tilde{v}_0, \tilde{v}_1, \tilde{v}_2, \dots$
  such that
    $\tilde{\eta}_t:= \mathcal{H}({\bm{g}}_0,\ldots, {\bm{g}}_t)$
  satisfies the following two inequalities for all $t \ge 0$:
  \begin{equation*}
    \begin{aligned}
      \abs{\tilde{\eta}_t - \frac{1}{\sqrt{\tilde{v}_t}}} &\le \delta(\tilde{v}_t) \cdot (1 - \beta) \cdot (1 + \bar{g}_t^2) \\
      \abs{\tilde{v}_{t+1} - \left(\beta \tilde{v}_t + (1-\beta) \bar{g}_t^2 \right)} &\le \delta(\tilde{v}_t) \cdot (1-\beta)^2 \cdot P(\bar{g}_t)
    \end{aligned}
    \qquad \text{where} \quad \bar{g}_t := \normtwosm{{\bm{g}}_t} / \eta.
  \end{equation*}
  The real number $\tilde{v}_t$ is called the \textit{moment estimate} at step $t$
  associated with the input stream and effective learning rates.
  \end{enumerate}
\end{definition}

It is clear that a RMSprop scheduler can be seen as a quasi-RMSprop scheduler with the same hyperparameters $\eta, \beta$,
and $\delta \equiv 0$.

\subsection{Reformulation of GD+WD\xspace on Scale-Invariant Loss via Quasi-RMSprop Scheduler} \label{sec:reform-gdwd}
Now we reformulate GD+WD\xspace on scale-invariant loss as
PGD with a quasi-RMSprop scheduler.
Recall that we say a loss function $\mathcal{L}({\bm{w}})$ is  scale-invariant if $\mathcal{L}(c{\bm{w}}) = \mathcal{L}({\bm{w}})$ for all $c > 0$,
and \Cref{lm:scale-invariant-on-sphere} converts GD+WD\xspace on scale-invariant functions
to Projected Gradient Descent (PGD) on unit sphere,
i.e., ${\bm{\theta}}_{t+1} \gets \Pi({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t))$,
where $\tilde{\eta}_t := \frac{\hat{\eta}}{(1 - \hat{\eta} \hat{\lambda}) \normtwosm{{\bm{w}}_t}^2}$
is the effective learning rate
at step $t$.
However, the evolution of $\tilde{\eta}_t$ over time is unclear unless we know how the parameter norm $\normtwosm{{\bm{w}}_t}$ changes.
Similar as the above analysis for Scalar RMSprop, where we decompose Scalar RMSprop as the GD method and an RMSprop scheduler,
here we abstract the evolution of $\tilde{\eta}_t$ as a learning rate scheduler,
which we call \textit{GWSI scheduler}
(name picked from the initials of \textbf{G}D+\textbf{W}D on \textbf{S}cale-\textbf{I}nvariant loss)
and view GD+WD\xspace on scale-invariant loss as
PGD on $\mathbb{S}^{D-1}$
with effective learning rates being set by a GWSI scheduler.
\begin{definition}[GWSI Scheduler] \label{def:gwsi-sche}
  Given a constant $\tilde{v}_0 > 0$,
  a base learning rate $\eta$
  and a decay rate $\beta$ as hyperparameters,
  a GWSI scheduler $\mathcal{H}_{\text{GWSI}}$ is a gradient-based learning rate scheduler that
  reads from a stream of vectors ${\bm{g}}_0, {\bm{g}}_1, \dots$,
  and generates effective learning rates $\tilde{\eta}_t := \mathcal{H}_{\text{GWSI}}({\bm{g}}_0, \dots, {\bm{g}}_t)$ for all $t \ge 0$
  according to the following recursion:
  \begin{align*}
    \tilde{\eta}_t &\gets \frac{1}{\sqrt{\tilde{v}_t}}, &
    \tilde{v}_{t+1} &\gets \beta \tilde{v}_t + (1-\beta) \bar{g}_t^2 + \frac{1}{4\beta \tilde{v}_t} (1-\beta)^2 \bar{g}_t^4, &
    &\text{where} \quad \bar{g}_t := \normtwosm{{\bm{g}}_t} / \eta.
  \end{align*}
\end{definition}
\begin{theorem} \label{thm:connect-siwd-rmsprop}
  For gradient descent \eqref{eq:upd-gdwd} with learning rate $\hat{\eta} > 0$ and weight decay $\hat{\lambda} > 0$
  on scale-invariant function $\mathcal{L}({\bm{w}})$,
  let ${\bm{\theta}}_t := \frac{{\bm{w}}_t}{\normtwosm{{\bm{w}}_t}}$ be the direction of ${\bm{w}}_t$
  and $\tilde{\eta}_t := \frac{\hat{\eta}}{(1-\hat{\eta}\hat{\lambda}) \normtwosm{{\bm{w}}_t}^2}$ be the effective learning rate at time $t$.
  Then $\tilde{\eta}_t$ evolves exactly the same as
  the GWSI scheduler with hyperparameters $\tilde{v}_0 := \frac{(1-\hat{\eta} \hat{\lambda})^2\normtwosm{{\bm{w}}_0}^4}{\hat{\eta}^2}$,
  $\beta := (1-\hat{\eta} \hat{\lambda})^4$,
  $\eta := \sqrt{(\beta^{-1} - 1)/2}$, and we can write the dynamics of ${\bm{\theta}}_t$ as:
  \begin{align*}
    {\bm{\theta}}_{t+1} &\gets \Pi({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)),  \qquad \text{where} \qquad \tilde{\eta}_t := \mathcal{H}_{\text{GWSI}}(\nabla \mathcal{L}({\bm{\theta}}_0), \dots, \nabla \mathcal{L}({\bm{\theta}}_t)).
  \end{align*}
\end{theorem}
\begin{remark}
  To the best of our knowledge, this particular form of the recursion formula (\Cref{def:gwsi-sche}) for the effective learning rates
  of GD+WD\xspace on scale-invariant loss does not appear in prior works, but some variants have been studied before.
  \citet{arora2018theoretical} derived a similar formula when the weight decay is zero,
  and \citet{li2020reconciling} obtained a Stochastic Differential Equation (SDE) for SGD+WD\xspace on scale-invariant functions, which is essentially a continuous approximation of our formulation
  when $\hat{\eta} \hat{\lambda} \to 0$.
\end{remark}

We defer the proof for \Cref{thm:connect-siwd-rmsprop} to \Cref{sec:proof-connect-siwd-rmsprop}.
The formula of GWSI scheduler clearly resembles RMSprop scheduler:
the only difference is that GWSI scheduler has an extra $\frac{1}{4\beta \tilde{v}_t} (1-\beta)^2 \bar{g}_t^4$ term,
which is negligible when $(1-\beta)^2$ is small.
This intuition is formalized through the definition of quasi-RMSprop.
It can be easily seen that GWSI scheduler is a quasi-RMSprop scheduler with the same hyperparameters $\eta, \beta$.

In fact,
we can use \Cref{thm:connect-siwd-rmsprop} as a basis
to obtain a better way to
write GD+WD\xspace on scale-invariant loss as PGD with a quasi-RMSprop scheduler,
in which $\eta, \beta$ are expressed more simply in terms of the intrinsic learning rate $\eta_{\mathrm{in}} := \hat{\eta} \hat{\lambda}$.
The main idea is that
$\beta = 1 - 4 \eta_{\mathrm{in}} + O(\eta_{\mathrm{in}}^2)$ and $\eta = \sqrt{2 \eta_{\mathrm{in}}} \cdot (1 + O(\eta_{\mathrm{in}}))$
when $\eta_{\mathrm{in}}$ is small,
and we can absorb these approximation errors
into the error bounds in \Cref{def:qrms-sche}.
We defer the details to \Cref{sec:proof-gdwd-is-quasi-rmsprop}.
\begin{theorem} \label{thm:gdwd-is-quasi-rmsprop}
  There exists a quasi-RMSprop scheduler $\mathcal{H}_{\text{QRMS}}$ with
  hyperparameter space $\{(\eta, \beta) : \beta = 1 - 2 \eta^2, \eta \in (0, \frac{1}{\sqrt{2}}) \}$
  such that the following holds
  for GD+WD\xspace on scale-invariant loss.
  If the intrinsic learning rate $\eta_{\mathrm{in}} := \hat{\eta} \hat{\lambda}$ lies in the range $(0, 1/4)$,
  then
  we can set the hyperparameters of~$\mathcal{H}_{\text{QRMS}}$
  to be
  $\eta = \sqrt{2\eta_{\mathrm{in}}}, \beta = 1 - 4\eta_{\mathrm{in}}$
  so that the dynamics of ${\bm{\theta}}_t$ can be written as
  \begin{align*}
    {\bm{\theta}}_{t+1} &\gets \Pi({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)),  \qquad \text{where} \qquad \tilde{\eta}_t := \mathcal{H}_{\text{QRMS}}(\nabla \mathcal{L}({\bm{\theta}}_0), \dots, \nabla \mathcal{L}({\bm{\theta}}_t)).
  \end{align*}
\end{theorem}

\subsection{Main Results for GD/PGD with Quasi-RMSprop Scheduler} \label{sec:quasi-main-thm}
As mentioned in \Cref{sec:proof-idea},
a key step in our analysis is to show the oscillation and periodic behaviors
of the effective learning rates.
But these behaviors can show up in other algorithms
besides GD+WD\xspace on scale-invariant loss.

In the following, we first present \Cref{thm:gdwd-main}, which holds for any PGD on $\mathbb{S}^{D-1}$
with effective learning rates being set by a quasi-RMSprop scheduler.
We focus on the case where $\eta > 0$ is small and $\beta = C_{\mathrm{b}} \eta^2 + O(\eta^4)$
for some constant $C_{\mathrm{b}} > 0$.
Our theorem shows that
if ${\bm{\theta}}_0$ is initially near a local minimizer manifold $\mathit{\Gamma}$
and is in the EoS regime,
then ${\bm{\theta}}_t$ approximately tracks a sharpness-reduction flow defined as follows:
\begin{equation}
    {\bm{\zeta}}(0) = {\bm{\zeta}}_0 \in \mathit{\Gamma}, \qquad \frac{\textup{\textrm{d}}}{\textup{\textrm{d}} \tau} {\bm{\zeta}}(\tau) = -\frac{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau))}{4 + \frac{2}{C_{\mathrm{b}}}\normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau))}^2},
    \label{eq:zeta-general}
\end{equation}
Besides, we also generalize our theorem to any GD with quasi-RMSprop scheduler,
where Scalar RMSprop serves as an important example.

The proof outlines of these theorems are given in \Cref{sec:proof-outline-main},
and the proof details are spread over
\Cref{sec:working-zone-proof,sec:gd-proof,sec:lr-proof,sec:ful-proof,sec:sph-proof,sec:rmsdrift-proof}.

\subsubsection{Spherical Optimization}

Consider Projected Gradient Descent (PGD) on the unit sphere $\mathbb{S}^{D-1}$:
\begin{equation} \label{eq:sph-thm-pgd}
  {\bm{\theta}}_{t+1} \gets \Pi({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)),
\end{equation}
where $\Pi: {\bm{w}} \mapsto \frac{{\bm{w}}}{\normtwosm{{\bm{w}}}}$ is the projection operator,
and the effective learning rate $\tilde{\eta}_t$
is set by a quasi-RMSprop scheduler with base learning rate $\eta > 0$ and decay rate $\beta = 1 - C_{\mathrm{b}} \eta^2 + O(\eta^4)$.

Recall that the dynamic is in the EoS regime
if $\tilde{\eta}_t \approx 2/\lambda_{\max}^{(t)}$
(see \Cref{lm:descent}).
When
${\bm{\theta}}_0$ is around a local minimizer ${\bm{\zeta}}_0$
and $\tilde{\eta}_0$ is generated by a quasi-RMSprop scheduler,
the condition of being in EoS is essentially
$\frac{1}{\sqrt{\tilde{v}_0}} \approx \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\zeta}}_0)}$.
In the following theorem,
we show that PGD with quasi-RMSprop scheduler evolves as \eqref{eq:zeta-general}
if it is initially in the EoS regime.

\myparagraph{Initialization Scheme.}
Given a local minimizer ${\bm{\zeta}}_0 \in \mathit{\Gamma}$ and a hyperparameter $\alpha_0$,
we initialize the initial direction ${\bm{\theta}}_0$ and initial moment estimate $\tilde{v}_0$ as follows:
draw ${\bm{\xi}} \sim \mathcal{N}({\bm{0}}, \sigma_0^2 {\bm{I}} / D)$
from Gaussian and set the direction of ${\bm{w}}_0$ to $\frac{{\bm{\zeta}}_0 + {\bm{\xi}}}{\normtwosm{{\bm{\zeta}}_0 + {\bm{\xi}}}}$,
where $\sigma_0$ can take any value in $[\exp(-\alpha_0^2) \eta, \alpha_0 \eta]$;
then set $\tilde{v}_0$
to be any positive value that satisfies $\abs{\frac{1}{\sqrt{\tilde{v}_0}} - \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\zeta}}_0)}} \le \alpha_0 \eta$.

\begin{theorem} \label{thm:main-sph}
  Under \Cref{ass:man,ass:topeigen},
  for PGD described as \eqref{eq:sph-thm-pgd}
  and initialized as above scheme for some ${\bm{\zeta}}_0 \in \mathit{\Gamma}$ and some $1 \le \alpha_0 \le \eta^{-o(1)}$,
  with probability $1 - O(\alpha_0 \eta \sqrt{\log(1/\eta)})$,
  the trajectory of ${\bm{\theta}}_t$ approximately tracks a sharpness-reduction flow ${\bm{\zeta}}: [0, T] \to \mathit{\Gamma}$
  that starts from ${\bm{\zeta}}_0$ and evolves as the ODE \eqref{eq:zeta-general} (if solution exists),
  in the sense that $\normtwosm{{\bm{\theta}}_t - {\bm{\zeta}}(t \eta^2)} = O(\alpha_0^2 \eta^{1/2} \log(1/\eta))$ for all $0 \le t \le T / \eta^2$.
\end{theorem}

\Cref{thm:gdwd-main} is a direct corollary of \Cref{thm:main-sph}, as GD+WD\xspace on scale-invariant loss can be seen as PGD with GWSI scheduler.
\begin{proof}[Proof for \Cref{thm:gdwd-main}]
\Cref{thm:gdwd-is-quasi-rmsprop} implies that GD+WD\xspace
on scale-invariant loss with LR $\hat{\eta}$
and WD $\hat{\lambda}$
can be seen as a PGD on $\mathbb{S}^{D-1}$
with quasi-RMSprop scheduler,
where the base learning rate and decay rate
of this scheduler is $\eta = \sqrt{2\eta_{\mathrm{in}}}$ and $\beta = 1 - 4 \eta_{\mathrm{in}}$.
Therefore, $\beta = 1 - C_{\mathrm{b}} \eta^2$ for $C_{\mathrm{b}} = 2$.

Now we apply \Cref{thm:main-sph} to prove \Cref{thm:gdwd-main}.
It is easy to translate the initial conditions of \Cref{thm:gdwd-main} to \Cref{thm:main-sph} with $\alpha_0 = \eta^{-o(1)} = \eta_{\mathrm{in}}^{-o(1)}$.
Let $\hat{{\bm{\zeta}}}(t)$
be the sharpness-reduction flow defined as in \eqref{eq:zeta-general} with $C_{\mathrm{b}}=2$ and horizon $\hat{T} := 2T$.
\begin{equation}
    \hat{{\bm{\zeta}}}(0) = {\bm{\zeta}}_0 \in \mathit{\Gamma}, \qquad \frac{\textup{\textrm{d}}}{\textup{\textrm{d}} \tau} \hat{{\bm{\zeta}}}(\tau) = -\frac{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1(\hat{{\bm{\zeta}}}(\tau))}{4 + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1(\hat{{\bm{\zeta}}}(\tau))}^2}.
\end{equation}
Then
with probability $1 - O(\alpha_0 \eta \sqrt{1/\delta}) = 1 - O(\eta^{1-o(1)})$,
${\bm{\theta}}_t$ approximately tracks $\hat{{\bm{\zeta}}}(\tau)$ in the sense that $\normtwosm{{\bm{\theta}}_t - \hat{{\bm{\zeta}}}(t \eta^2)} = O(\eta^{1/2-o(1)})$ for all $0 \le t \le \hat{T} /\eta^2$.
Replacing $\eta$ with $\sqrt{2\eta_{\mathrm{in}}}$ gives
$\normtwosm{{\bm{\theta}}_t - \hat{{\bm{\zeta}}}(2t\eta_{\mathrm{in}})} = O(\eta_{\mathrm{in}}^{1/4-o(1)})$ for all $0 \le t \le T / \eta_{\mathrm{in}}$,
and the success probability becomes $1 - O(\eta_{\mathrm{in}}^{1/2-o(1)})$.
We can finish the proof by noting that
$\hat{{\bm{\zeta}}}(2\tau)$ is just ${\bm{\zeta}}(\tau)$ defined in \eqref{eq:main-zeta}.
\end{proof}

\subsubsection{Full Space Optimization}
Now we present our general theorem for GD on $\mathbb{R}^D$ with quasi-RMSprop scheduler.
Here
GD with quasi-RMSprop scheduler can be written as:
\begin{equation} \label{eq:ful-thm-gd}
  {\bm{\theta}}_{t+1} \gets {\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t),
\end{equation}
where the effective learning rate $\tilde{\eta}_t$
is set by a quasi-RMSprop scheduler with base learning rate $\eta > 0$ and decay rate $\beta = 1 - C_{\mathrm{b}} \eta^2 + O(\eta^4)$.

Similar to the spherical case,
we assume that there is a manifold $\mathit{\Gamma}$
consisting of local minimizers,
but now we are not assuming scale-invariance.
\begin{assumption} \label{ass:man-ful}
  The loss function $\mathcal{L}: \mathbb{R}^D \to \mathbb{R}$ is $\mathcal{C}^4$-smooth.
  $\mathit{\Gamma}$ is a $\mathcal{C}^2$-smooth, $D_\Gamma$-dimensional open submanifold of $\mathbb{S}^{D-1}$
  for some $0 \le D_\Gamma \le D$, where every ${\bm{\theta}} \in \mathit{\Gamma}$ is a local minimizer of $\mathcal{L}$ on $\mathbb{R}^{D}$
  and $\mathrm{rank}({\bm{H}}({\bm{\theta}})) = D - D_\Gamma$.
\end{assumption}

The following assumption is essentially the same as \Cref{ass:topeigen} except that now $\mathit{\Gamma}$ is defined differently.
\begin{assumption} \label{ass:topeigen-ful}
  For all ${\bm{\theta}} \in \mathit{\Gamma}$, $\lambda^{\mathrm{H}}_1({\bm{\theta}}) > \lambda^{\mathrm{H}}_2({\bm{\theta}})$. That is, the top eigenvalue of ${\bm{H}}({\bm{\theta}})$ is unique.
\end{assumption}

Based on \Cref{ass:man-ful,ass:topeigen-ful} above,
we study GD
starting in the EoS regime.
Similar to the spherical case,
we focus on the case where ${\bm{\theta}}_0$ is around a local minimizer ${\bm{\zeta}} \in \mathit{\Gamma}$,
and EoS is the regime in which $\tilde{\eta}_t \approx \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\zeta}}_0)}$,
which
essentially means $\frac{1}{\sqrt{\tilde{v}_0}}$
if the effective LR is set by a quasi-RMSprop scheduler.
In the theorem below, we show that GD with quasi-RMSprop scheduler
tracks the sharpness-reduction flow defined in \eqref{eq:zeta-general}.
Note that the ODE here is the same as the spherical case, but $\mathit{\Gamma}$ is defined differently.

\myparagraph{Initialization Scheme.}
Given a local minimizer ${\bm{\zeta}}_0 \in \mathit{\Gamma}$ and a hyperparameter $\alpha_0$,
we initialize the initial parameter ${\bm{\theta}}_0$ and initial moment estimate $\tilde{v}_0$ as follows:
draw ${\bm{\xi}} \sim \mathcal{N}({\bm{0}}, \sigma_0^2 {\bm{I}} / D)$
from Gaussian and set ${\bm{w}}_0 \gets {\bm{\zeta}}_0 + {\bm{\xi}}$,
where $\sigma_0$ can take any value in $[\exp(-\alpha_0^2) \eta, \alpha_0 \eta]$;
then set $\tilde{v}_0$
to be any positive value that satisfies $\abs{\frac{1}{\sqrt{\tilde{v}_0}} - \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\zeta}}_0)}} \le \alpha_0 \eta$.

\begin{theorem} \label{thm:main-ful}
  Under \Cref{ass:man-ful,ass:topeigen-ful},
  for GD described as \eqref{eq:ful-thm-gd}
  and initialized as above scheme for some ${\bm{\zeta}}_0 \in \mathit{\Gamma}$ and some $1 \le \alpha_0 \le \eta^{-o(1)}$,
  with probability $1 - O(\alpha_0 \eta \sqrt{\log(1/\eta)})$,
  the trajectory of ${\bm{\theta}}_t$ approximately tracks a sharpness-reduction flow ${\bm{\zeta}}: [0, T] \to \mathit{\Gamma}$
  that starts from ${\bm{\zeta}}_0$ and evolves as the ODE \eqref{eq:zeta-general} (if solution exists),
  in the sense that $\normtwosm{{\bm{\theta}}_t - {\bm{\zeta}}(t \eta^2)} = O(\alpha_0^2 \eta^{1/2} \log(1/\eta))$ for all $0 \le t \le T / \eta^2$.
\end{theorem}

A direct corollary is that Scalar RMSprop follows the sharpness-reduction flow~\eqref{eq:zeta-general},
since the RMSprop scheduler is a quasi-RMSprop scheduler.
\begin{corollary} \label{cor:scalar-rmsprop}
  The statement of \Cref{thm:main-ful} holds
  for Scalar RMSprop
  if the decay rate is set to $\beta = 1 - C_{\mathrm{b}}\eta^2 + O(\eta^2)$ for some constant $C_{\mathrm{b}} > 0$.
\end{corollary}

\subsection{Alignment Phase} \label{sec:ful-proof-alignment}

We define $r_t = \normtwosm{{\bm{x}}_t} / \eta$ and the following notations
for this subsection.
\begin{align*}
  \hat{\vx}_{t+1} &:= {\bm{\theta}}_{t+1} - {\bm{\phi}}_t & \hat{u}_{t+1} &:= \tfrac{1}{\eta}( \mu_{t}^2 \tilde{v}_{t+1} - 1)\\
  \hat{h}_{t+1} &:= \tfrac{1}{\eta}\dotpsm{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)}{\hat{\vx}_{t+1}}
\end{align*}
\begin{lemma} \label{lm:ful-gap-2gamin}
  If ${\bm{\theta}}_t \in \gZ^{\epsilon_0}$ at step $t$,
  then $\normtwosm{\mPHnzt({\bm{\phi}}_t) ({\bm{I}} - \mu_t {\bm{H}}_t)} \le 1 - 2 \gamma_{\min}$.
\end{lemma}
\begin{proof}
If $\lambda^{\mathrm{H}}_1({\bm{\phi}}_t), \dots, \lambda^{\mathrm{H}}_D({\bm{\phi}}_t)$ are the eigenvalues of
    ${\bm{H}}({\bm{\phi}}_t)$, then $\left\{ 1 -
    \sfrac{2\lambda^{\mathrm{H}}_i({\bm{\phi}}_t)}{\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)} \right\}_{i=1}^{D}$ are the
    eigenvalues of ${\bm{I}} - \mu_t {\bm{H}}_t$.
By definition of $\gamma_{\min}$, we have
$1 - \sfrac{2\lambda^{\mathrm{H}}_i({\bm{\phi}}_t)}{\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)} \in [-1 + 2\gamma_{\min}, 1 - 2\gamma_{\min}]$
as long as $\lambda^{\mathrm{H}}_i({\bm{\phi}}_t) \ne 0, \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$.
Therefore we have $\normtwosm{\mPHnzt({\bm{\phi}}_t) ({\bm{I}} - \mu_t {\bm{H}}_t)} \le 1 - 2 \gamma_{\min}$.
\end{proof}

\begin{lemma} \label{lm:ful-align-1-step-hat}
  In the setting of \Cref{lm:ful-align-1-step},
  \begin{align}
    \hat{\vx}_{t+1} &= ({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t + O(r_t^2 \eta^2) \label{eq:ful-align-1-step-hatx} \\
    \hat{h}_{t+1} &= -h_t + O(\alpha r_t \eta) \label{eq:ful-align-1-step-hath} \\
    \normtwosm{\mPHnzt({\bm{\phi}}_t) \hat{\vx}_{t+1}} &= (1-1.9\gamma_{\min}) \normtwosm{\mPHnzt({\bm{\phi}}_t) {\bm{x}}_{t}} + O(r_t^2 \eta^2) \label{eq:ful-align-1-step-hatpnzt} \\
    \hat{u}_{t+1} &= u_t + O(\alpha^2 \eta) \label{eq:ful-align-1-step-hatu}
  \end{align}
\end{lemma}
\begin{proof}[Proof for \Cref{lm:ful-align-1-step-hat}]
By \Cref{lm:gd-1-step}, we have the following first-order approximation for $\hat{\vx}_{t+1}$:
\begin{align*}
\hat{\vx}_{t+1} = \left({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}_t)\right) {\bm{x}}_t + O(r_t^2 \eta^2),
\end{align*}
which proves \eqref{eq:ful-align-1-step-hatx}.
Then we can prove \eqref{eq:ful-align-1-step-hath} as follows:
\begin{align*}
  \hat{h}_{t+1} = \dotpsm{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)}{\hat{\vx}_{t+1}}
  &= (1 - \tilde{\eta}_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)) h_t + O(r_t^2 \eta^2) \\
  &= (1 - \mu_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)) h_t + O(\alpha r_t \eta^2) \\
  &= -h_t + O(\alpha r_t \eta^2).
\end{align*}

To prove \eqref{eq:ful-align-1-step-hatpnzt} from \eqref{eq:ful-align-1-step-hatx},
we note that
if $\eta$ is sufficiently small,
then $\tilde{\eta}_t$ is sufficiently close to $\mu_t$.
In this case, by \Cref{lm:ful-gap-2gamin} we have
$\normtwosm{({\bm{I}} - \tilde{\eta}_t{\bm{H}}_t) {\bm{x}}_t} \le (1-1.9\gamma_{\min}) \normtwosm{\mPHnzt({\bm{\phi}}_t) {\bm{x}}_{t}}$.
Combining this with \eqref{eq:ful-align-1-step-hatx} proves
\eqref{eq:ful-align-1-step-hatpnzt}.

Finally, we prove \eqref{eq:ful-align-1-step-hatu}
By Taylor expansion of $\mathcal{L}$ around ${\bm{\phi}}_t$, $\normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)} = O(\normtwosm{{\bm{x}}_t}) = O(r_t \eta)$.
By \Cref{lm:lr-1-step}, we can approximate $\hat{u}_{t+1}$ by
\begin{align*}
    \hat{u}_{t+1}
    &= u_t + C_{\mathrm{b}} \eta (\mu_t^2 \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t) / \eta}^2 - 1) + O(\alpha \eta^2) \\
    &= u_t + C_{\mathrm{b}} \eta (\mu_t^2 \cdot O(r_t^2) - 1) + O(\alpha \eta^2) \\
    &= u_t + O((r_t^2 + 1) \eta) + O(\alpha \eta^2) \\
    &= u_t + O(\alpha^2 \eta),
\end{align*}
which proves \eqref{eq:ful-align-1-step-hatu}.
\end{proof}

\begin{proof}[Proof for \Cref{lm:ful-align-1-step}]
Combining \eqref{eq:ful-align-1-step-hatx} with a Taylor expansion of $\Phi$ gives
\begin{align*}
  {\bm{\phi}}_{t+1} = \Phi({\bm{\phi}}_t + \hat{\vx}_{t+1})
  &= {\bm{\phi}}_t + \partial \Phi_{{\bm{\phi}}_t}[\hat{\vx}_{t+1}] + O(r_t^2\eta^2) \\
  &= {\bm{\phi}}_t + \partial \Phi_{{\bm{\phi}}_t}[({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t){\bm{x}}_t] + O(r_t^2\eta^2).
\end{align*}
By
\Cref{lm:Phi-proj-ful} and \Cref{lm:phiphix-ful},
$\partial \Phi_{{\bm{\phi}}_t}[({\bm{I}} - \mu_t {\bm{H}}_t){\bm{x}}_t] = \partial \Phi_{{\bm{\phi}}_t}[{\bm{x}}_t] = O(r_t^2 \eta^2)$.
So we have ${\bm{\phi}}_{t+1} = {\bm{\phi}}_t + O(r_t^2 \eta^2)$,
which proves \eqref{eq:ful-align-1-step-phi}.
Then ${\bm{x}}_{t+1} = \hat{\vx}_{t+1} + {\bm{\phi}}_{t} - {\bm{\phi}}_{t+1} = \hat{\vx}_{t+1} + O(r_t^2 \eta^2)$,
which proves \eqref{eq:ful-align-1-step-x}.
Finally,
by $\mathcal{C}^1$-smoothness of $\mu({\bm{\phi}}), \vv^{\mathrm{H}}_1({\bm{\phi}}), \mPHnzt({\bm{\phi}})$ on $\mathit{\Gamma}$,
\eqref{eq:ful-align-1-step-hath}, \eqref{eq:ful-align-1-step-hatpnzt}, \eqref{eq:ful-align-1-step-hatu}
imply \eqref{eq:ful-align-1-step-h}, \eqref{eq:ful-align-1-step-pnzt}, \eqref{eq:ful-align-1-step-u} respectively.
\end{proof}

\begin{proof}[Proof for \Cref{thm:ful-alignment-main}]
Let $\delta := C_0 \alpha_0 \eta \sqrt{\log(1/\eta)}$,
where $C_0 = O(1)$ is a sufficiently large constant so that our proof can work.
Let $\alpha_{\max} := \alpha_0 \sqrt{\log(1/\delta)}$.

Note that the initial state is $O(1/\delta)$-misaligned.
So $\normtwosm{{\bm{x}}_0} = O(h_0 \eta / \delta)$.
By \Cref{lm:ful-align-1-step},
for all $t \le T_0 := \left\lceil \frac{1}{\log(1 - 1.9 \gamma_{\min})} \log \frac{1}{\delta} \right\rceil$,
we can prove by induction that
the states are $O(\alpha_{\max})$-bounded, and
\begin{align*}
  \normtwosm{\mPHnzt({\bm{\phi}}_{t}) {\bm{x}}_{t}} &= \Obig{(1-1.9\gamma_{\min})^t \normtwosm{\mPHnzt({\bm{\phi}}_0) {\bm{x}}_0}} \\
  \normtwosm{{\bm{x}}_t} &= O((1 - 1.9\gamma_{\min})^t \normtwosm{{\bm{x}}_0}), \\
  {\bm{\phi}}_t &= {\bm{\phi}}_0 + O(\normtwosm{{\bm{x}}_0}^2) \\
  h_{t} &= (-1)^t h_0 + O(\alpha_{\max} \normtwosm{{\bm{x}}_0})
\end{align*}
At $t = T_0$, $\normtwosm{{\bm{x}}_t} = O(h_0 \eta)$,
${\bm{\phi}}_t = {\bm{\phi}}_0 + O(\normtwosm{{\bm{x}}_0}^2)$,
$\abssm{h_t - (-1)^t h_0} \le O(\alpha_{\max} h_0 \eta / \delta)$,
$\normtwosm{\mPHnzt({\bm{\phi}}_{t}) {\bm{x}}_{t}} = O(h_0 \eta)$.
When $C_0$ is sufficiently large, it also holds that
$\abssm{h_t - (-1)^t h_0} \le \abssm{h_0} / 4$.

After that, for all $T_0 \le t \le T_1 := \left\lceil \frac{1}{\log(1 - 1.9 \gamma_{\min})} \log \frac{1}{\delta h_0 \eta} \right\rceil = O(\alpha_{\max})$,
we can prove by induction that the states are $O(\alpha_{\max})$-bounded and
\begin{align*}
  \normtwosm{\mPHnzt({\bm{\phi}}_{t}) {\bm{x}}_{t}} &= \Obig{(1-1.9\gamma_{\min})^t \normtwosm{\mPHnzt({\bm{\phi}}_0) {\bm{x}}_0} + h_0^2 \eta^2} \\
  \normtwosm{{\bm{x}}_t} &= O(h_0 \eta), \\
  {\bm{\phi}}_t &= {\bm{\phi}}_0 + O(\normtwosm{{\bm{x}}_0}^2 + h_0^2 \eta^2 t) \\
  h_{t} &= (-1)^{t - T_0} h_{T_0} + O(\alpha_{\max} h_0 \eta t)
\end{align*}
At $t = T_1$,
$\normtwosm{\mPHnzt({\bm{\phi}}_{t}) {\bm{x}}_{t}} = O(h_0^2 \eta^2)$,
$\normtwosm{{\bm{x}}_t} = O(h_0 \eta)$,
${\bm{\phi}}_t = {\bm{\phi}}_0 + O(\alpha_{\max}^3 \eta^2) = {\bm{\phi}}_0 + O(\alpha_{\max} \eta)$,
$\abssm{h_t - (-1)^{t - T_1}} \le O(\alpha_{\max}^2 h_0 \eta)$.
When $\eta$ is sufficiently small, it also holds that
$\abssm{h_t - (-1)^{t-T_1} h_{T_1}} \le \abssm{h_0} / 4$.
So $\abssm{h_{T_1} - (-1)^t h_0} \le \abssm{h_0} / 2$.
Putting all these together proves the theorem.
\end{proof}

\subsection{Drifting Phase} \label{sec:ful-proof-drifting}

We define ${\bm{\psi}}_{{\bm{\phi}}}({\bm{x}})$ as in \Cref{lm:gd-2-step}. We abuse the notation to write ${\bm{\psi}}_t({\bm{x}}) = {\bm{\psi}}_{{\bm{\phi}}_t}({\bm{x}})$, that is,
\[
  {\bm{\psi}}_t({\bm{x}}) = \frac{\mu_t}{2} \left( {\bm{U}}_t \partial^3\mathcal{L}_{{\bm{\phi}}_t}[{\bm{x}}, {\bm{x}}] + \partial^3 \mathcal{L}_{{\bm{\phi}}_t}[{\bm{U}}_t {\bm{x}}, {\bm{U}}_t {\bm{x}}] \right).
\]
We also define the following notations for this subsection.
\begin{align*}
  \hat{\vx}_{t+2} &:= {\bm{\theta}}_{t+2} - {\bm{\phi}}_t, & \hat{u}_{t+2} &:= \frac{1}{\eta}(\mu_t^2 \tilde{v}_{t+2} - 1), &
  \hat{h}_{t+2} &:= \frac{1}{\eta}\dotpsm{{\bm{x}}_{t+2}}{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)}.
\end{align*}

\begin{lemma} \label{lm:aligned-hteta}
  If the state $({\bm{\theta}}_t, \tilde{v}_t)$ at step $t$
  is $\eta^{-o(1)}$-bounded and
  $O(\abssm{h_t}\eta)$-misaligned,
  then
  \[
    {\bm{x}}_t = h_t \eta \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(h_t^2 \eta^2).
  \]
\end{lemma}
\begin{proof}
Note that we have the decomposition ${\bm{x}}_t = h_t \eta \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + \mPHz({\bm{\phi}}_t) {\bm{x}}_t + \mPHnzt({\bm{\phi}}_t) {\bm{x}}_t$.
By definition of $O(\abssm{h_t}\eta)$-misaligned state,
$\mPHnzt({\bm{\phi}}_t) {\bm{x}}_t = O(h_t^2 \eta^2)$.
By \Cref{lm:mPHz-x-small}, $\mPHz({\bm{\phi}}_t) {\bm{x}}_t = O(\normtwosm{{\bm{x}}_t}^2)$.
\[
\normtwosm{{\bm{x}}_t} = h_t \eta + O(h_t^2 \eta^2) + O(\normtwosm{{\bm{x}}_t}^2).
\]
Solving this equation gives $\normtwosm{{\bm{x}}_t} = O(\abssm{h_t}\eta)$.
Then $\mPHz({\bm{\phi}}_t) {\bm{x}}_t = O(h_t^2 \eta^2)$,
and therefore we have ${\bm{x}}_t = h_t \eta \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(h_t^2 \eta^2)$.
\end{proof}

\begin{lemma} \label{lm:ful-drift-2-step-hat}
In the setting of \Cref{lm:ful-drift-2-step},
\begin{align}
\hat{\vx}_{t+2} &= ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}_t)({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t - \eta^2 h_t^2 {\bm{\psi}}_t(\vv^{\mathrm{H}}_1({\bm{\phi}}_t)) + O(\alpha h_t^2 \eta^3), \label{eq:ful-drift-2-step-hatx} \\
\tilde{\eta}_{t} &= \mu_t \cdot (1 - \tfrac{1}{2} \eta u_t) + O(\alpha^2 \eta^2),  \label{eq:ful-drift-2-step-eeta-t0}\\
\tilde{\eta}_{t+1} &= \mu_t \cdot (1 - \tfrac{1}{2} \eta u_t) + O(\alpha^2 \eta^2). \label{eq:ful-drift-2-step-eeta-t1} \\
\hat{u}_{t+2}
&= u_t + 8 C_{\mathrm{b}} \eta h_t^2 - 2C_{\mathrm{b}} \eta + O(\alpha(1+h_t^2) \eta^2). \label{eq:ful-drift-2-step-hatu}
\end{align}
\end{lemma}
\begin{proof}
Define
$\bar{g}_t = \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t) / \eta}$
and
$\bar{g}_{t+1} = \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_{t+1}) / \eta}$.
Since the state at step $t$ is $O(\abssm{h_t}\eta)$-misaligned,
$\normtwosm{{\bm{x}}_t} \le O(\abssm{h_t}\eta)$.
By Taylor expansion of $\nabla \mathcal{L}$ around ${\bm{\phi}}_t$,
\begin{align*}
  \nabla \mathcal{L}({\bm{\theta}}_t) = \nabla \mathcal{L}({\bm{\phi}}_t) + \nabla^2 \mathcal{L}({\bm{\phi}}_t) {\bm{x}}_t + O(h_t^2 \eta^2) = {\bm{H}}_t {\bm{x}}_t + O(h_t^2 \eta^2).
\end{align*}
So
$\bar{g}_t = \normtwosm{{\bm{H}}_t {\bm{x}}_t / \eta} + O(h_t^2 \eta) = O(h_t)$.
Then \Cref{lm:lr-1-step} implies \eqref{eq:ful-drift-2-step-eeta-t0} and
the following approximation for $\hat{u}_{t+1}$:
\begin{align*}
  \hat{u}_{t+1} &= u_t + C_{\mathrm{b}} \eta (  \mu_t^2 \bar{g}_t^2 - 1) + O(\alpha \eta^2) \\
  &= u_t + C_{\mathrm{b}} \eta (  \mu_t^2 \normtwosm{{\bm{H}}_t {\bm{x}}_t / \eta}^2 - 1) + O((\alpha + \abssm{h_t}^3) \eta^2).
\end{align*}
As \eqref{eq:ful-drift-2-step-eeta-t0} verifies $\tilde{\eta}_t = \mu_t + O(\alpha \eta)$,
we can use \Cref{lm:gd-1-step} to derive the zeroth-order and first-order approximations for ${\bm{\theta}}_{t+1}$:
${\bm{\theta}}_{t+1} - {\bm{\phi}}_t = O(\abssm{h_t}\eta)$
and
${\bm{\theta}}_{t+1} - {\bm{\phi}}_t = {\bm{U}}_t {\bm{x}}_t + O(\alpha \abssm{h_t} \eta^2)$.
Then by Taylor expansion of $\nabla \mathcal{L}$ around ${\bm{\phi}}_t$ again,
\begin{align*}
  \nabla \mathcal{L}({\bm{\theta}}_{t+1}) = \nabla \mathcal{L}({\bm{\phi}}_t) + \nabla^2 \mathcal{L}({\bm{\phi}}_t) ({\bm{U}}_t{\bm{x}}_t + O(\alpha \abssm{h_t} \eta^2) ) + O(h_t^2\eta^2) = {\bm{H}}_t {\bm{U}}_t {\bm{x}}_t + O(\alpha h_t \eta^2).
\end{align*}
So $\bar{g}_{t+1} = \normtwosm{{\bm{H}}_t {\bm{U}}_t {\bm{x}}_t / \eta} + O(\alpha \abssm{h_t} \eta) = O(\abssm{h_t})$.
Then \Cref{lm:lr-2-step} implies \eqref{eq:ful-drift-2-step-eeta-t1}.
We can further apply \Cref{lm:gd-2-step} to obtain the following:
\begin{align}
\hat{\vx}_{t+2} = ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}_t)({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t - {\bm{\psi}}_t({\bm{x}}_t) + O(\alpha h_t^2 \eta^3).
\label{eq:ful-drift-2-step-hatx-prelim}
\end{align}
Note that \Cref{lm:aligned-hteta} implies that
${\bm{x}}_t = h_t \eta \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(h_t^2 \eta^2)$.
Then by \Cref{lm:ful-psi-formula}
\[
{\bm{\psi}}_t({\bm{x}}_t) = {\bm{\psi}}_t(\eta h_t \vv^{\mathrm{H}}_1({\bm{\phi}}_t)) + O(h_t^2 \eta^2 \cdot \abssm{h_t}\eta)
  = \eta^2 h_t^2 {\bm{\psi}}_t(\vv^{\mathrm{H}}_1({\bm{\phi}}_t)) + O(\abssm{h_t}^3 \eta^3).
\]
Combining this with \eqref{eq:ful-drift-2-step-hatx-prelim} gives \eqref{eq:ful-drift-2-step-hatx}.

Finally, we derive the approximation for $\hat{u}_{t+2}$.
By \Cref{lm:lr-2-step},
\begin{align}
  \hat{u}_{t+2} &= u_t + C_{\mathrm{b}} \eta (  \mu_t^2 \bar{g}_t^2 + \mu_t^2 \bar{g}_{t+1}^2 - 2) + O(\alpha \eta^2).
  \label{eq:ful-drift-2-step-hatu-prelim}
\end{align}
Since ${\bm{x}}_t = h_t \eta \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(h_t^2 \eta^2)$,
for $\bar{g}_t$ we have
\begin{align*}
\bar{g}_t
= \normtwosm{{\bm{H}}_t (\eta h_t \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(h_t^2 \eta^2)) / \eta} + O(h_t^2 \eta)
&= \normtwosm{h_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) \vv^{\mathrm{H}}_1({\bm{\phi}}_t)} + O(h_t^2 \eta) \\
&= \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) \abssm{h_t} + O(h_t^2 \eta).
\end{align*}
Similarly, for $\bar{g}_{t+1}$ we have
\begin{align*}
\bar{g}_{t+1}
= \normtwosm{{\bm{H}}_t {\bm{U}}_t (\eta h_t \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(h_t^2 \eta^2)) / \eta}
&= \normtwo{-h_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) \vv^{\mathrm{H}}_1({\bm{\phi}}_t)} + O(h_t^2 \eta) \\
&= \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) \abssm{h_t} + O(h_t^2 \eta).
\end{align*}
So both $\bar{g}_t^2$ and $\bar{g}_{t+1}^2$ can be approximated by $\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)^2 h_t^2 + O(\abssm{h_t}^3 \eta)$.
Combining this with \eqref{eq:ful-drift-2-step-hatu-prelim} gives
\begin{align*}
\hat{u}_{t+2}
&= u_t + C_{\mathrm{b}} \eta \mu_t^2 (\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)^2 h_t^2 + O(\abssm{h_t}^3 \eta)) + C_{\mathrm{b}} \eta \mu_t^2 (\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)^2 h_t^2 + O(\abssm{h_t}^3 \eta)) \\
& \quad - 2 C_{\mathrm{b}} \eta + \Obig{\alpha(1 + h_t^2)\eta^2} \\
&= u_t + 8 C_{\mathrm{b}} \eta h_t^2 - 2C_{\mathrm{b}} \eta + O(\alpha(1+h_t^2) \eta^2),
\end{align*}
which implies \eqref{eq:ful-drift-2-step-hatu}.
\end{proof}

\begin{lemma} \label{lm:ful-drift-2-step-hat2}
In the setting of \Cref{lm:ful-drift-2-step},
\begin{align}
\hat{h}_{t+2} &= (1 - 2 \eta u_t) h_t  + \Osm{\alpha^2 \abssm{h_t} \eta^2} \label{eq:ful-drift-2-step-hath} \\
\mPHz({\bm{\phi}}_t) \hat{\vx}_{t+2}
&= \mPHz({\bm{\phi}}_t) {\bm{x}}_t - 2 \eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + \Osm{\alpha h_t^2 \eta^3} \label{eq:ful-drift-2-step-hatx-tan} \\
\normtwosm{\mPHnzt({\bm{\phi}}_t) \hat{\vx}_{t+2}}
&\le (1 - 1.9\gamma)^2 \normtwosm{\mPHnzt({\bm{\phi}}_t) {\bm{x}}_t}
+ O(h_t^2 \eta^2) \label{eq:ful-drift-2-step-hatx-w}
\end{align}
\end{lemma}
\begin{proof}
In the following, we derive the approximations from \eqref{eq:ful-drift-2-step-hatx},
\begin{align*}
\hat{\vx}_{t+2} &= ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}_t)({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t - \eta^2 h_t^2 {\bm{\psi}}_t(\vv^{\mathrm{H}}_1({\bm{\phi}}_t)) + O(\alpha h_t^2 \eta^3).
\end{align*}

\myparagraph{Approximation for $\hat{h}_{t+2}$.}
Note that ${\bm{x}}_t = \eta h_t \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(h_t^2\eta^2)$
since the state at step $t$ is $O(\abssm{h_t}\eta)$-misaligned.
For $\hat{h}_{t+2}$, we have
\begin{align*}
\hat{h}_{t+2} &= \tfrac{1}{\eta}\left( \dotpsm{{\bm{x}}_t}{({\bm{I}} - \tilde{\eta}_{t+1} {\bm{H}}_t)({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) \vv^{\mathrm{H}}_1({\bm{\phi}}_t)}
- \eta^2 h_t^2 \dotpsm{{\bm{\psi}}_t(\vv^{\mathrm{H}}_1({\bm{\phi}}_t))}{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)} + O(\alpha h_t^2\eta^3) \right) \\
&= (1 - \tilde{\eta}_{t+1} \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)) (1 - \tilde{\eta}_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)) h_t + O(\alpha h_t^2 \eta^2),
\end{align*}
where we use
the fact that $\dotpsm{{\bm{\psi}}_t(\vv^{\mathrm{H}}_1({\bm{\phi}}_t))}{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)} = 0$ by \Cref{lm:ful-psi-formula}.

By \eqref{eq:ful-drift-2-step-eeta-t0}, \eqref{eq:ful-drift-2-step-eeta-t1},
$\tilde{\eta}_{\tau} = \mu_t \cdot (1 - \tfrac{1}{2} \eta u_t) + O(\alpha^2 \eta^2)$ for $\tau \in \{t, t+1\}$.
Note that $\mu_t \cdot \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) = 2$.
Then for $\tau \in \{t, t+1\}$,
\begin{align*}
1 - \tilde{\eta}_{\tau} \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)
&= 1 - \mu_t \cdot (1 - \tfrac{1}{2} \eta u_t) \cdot \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + O(\alpha^2 \eta^2) \\
&= 1 - 2 \cdot (1 - \tfrac{1}{2} \eta u_t)) + O(\alpha^2 \eta^2) \\
&= -1 + \eta u_t + O(\alpha^2 \eta^2).
\end{align*}
Then $(1 - \tilde{\eta}_{t+1} \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)) (1 - \tilde{\eta}_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t))$ can be approximated by
\[
  (1 - \tilde{\eta}_{t+1} \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)) (1 - \tilde{\eta}_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)) = (-1 + \eta u_t + O(\alpha^2 \eta^2))^2 = 1 - 2 \eta u_t + O(\alpha^2 \eta^2).
\]
Therefore, we have
$\hat{h}_{t+2} = \left(1 - 2 \eta u_t + O(\alpha^2 \eta^2)\right) h_t + O(\alpha h_t^2 \eta^2)$,
which implies \eqref{eq:ful-drift-2-step-hath}.

\myparagraph{Approximation for $\mPHz({\bm{\phi}}_t) \hat{\vx}_{t+2}$.}
For $\mPHz({\bm{\phi}}_t) \hat{\vx}_{t+2}$, we have
\begin{align*}
\mPHz({\bm{\phi}}_t) \hat{\vx}_{t+2}
&= \mPHz({\bm{\phi}}_t) ({\bm{I}} - \tilde{\eta}_{t+1} {\bm{H}}_t)({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t
- \eta^2 h_t^2 \mPHz({\bm{\phi}}_t) {\bm{\psi}}_t(\vv^{\mathrm{H}}_1({\bm{\phi}}_t)) + O(\alpha h_t^2 \eta^3) \\
&= \mPHz({\bm{\phi}}_t) ({\bm{I}} - \tilde{\eta}_{t+1} {\bm{H}}_t)({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t
- 2\eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + O(\alpha h_t^2 \eta^3),
\end{align*}
where we use
the fact that
$\mPHz({\bm{\phi}}_t) {\bm{\psi}}_t(\vv^{\mathrm{H}}_1({\bm{\phi}}_t)) = 2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$
by \Cref{lm:ful-psi-formula}.
To obtain \eqref{eq:ful-drift-2-step-hatx-tan}, we only need to note that
$\mPHz({\bm{\phi}}_t) ({\bm{I}} - \tilde{\eta}_{t+1} {\bm{H}}_t)({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t = \mPHz({\bm{\phi}}_t){\bm{x}}_t$.

\myparagraph{Approximation for $\normtwosm{\mPHnzt({\bm{\phi}}_t) \hat{\vx}_{t+2}}$.}
To approximate $\normtwosm{\mPHnzt({\bm{\phi}}_t) \hat{\vx}_{t+2}}$,
we note that
if $\eta$ is sufficiently small,
then $\tilde{\eta}_t, \tilde{\eta}_{t+1}$ are sufficiently close to $\mu_t$.
In this case, by
\Cref{lm:ful-gap-2gamin}
we have
$\normtwosm{\mPHnzt({\bm{\phi}}_t)({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}_t)({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t} \le (1 - 1.9 \gamma_{\min})^2 \normtwosm{\mPHnzt({\bm{\phi}}_t){\bm{x}}_t}$.
Combining this with \eqref{eq:ful-drift-2-step-hatx} proves \eqref{eq:ful-drift-2-step-hatx-w}.
\end{proof}

\begin{lemma} \label{lm:ful-drift-2-step-phi}
  In the setting of \Cref{lm:ful-drift-2-step},
  the approximation \eqref{eq:ful-drift-2-step-phi} holds for ${\bm{\phi}}_{t+2}$ and
  \begin{align}
    \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_{t+2}) &= \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)
    - 2 \eta^2 h_t^2 \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2 + \Osm{\alpha h_t^2 \eta^3}.\label{eq:ful-drift-2-step-loglam}
  \end{align}
\end{lemma}
\begin{proof}
By Taylor expansion,
\begin{align}
  {\bm{\phi}}_t &= \Phi({\bm{\phi}}_t) + \partial \Phi_{{\bm{\phi}}_t}[{\bm{x}}_t] + \frac{1}{2} \partial^2 \Phi_{{\bm{\phi}}_t}[{\bm{x}}_t, {\bm{x}}_t] + O(\abssm{h_t}^3 \eta^3). \label{eq:ful-drift-2-step-phi-taylor1} \\
  {\bm{\phi}}_{t+2} &= \Phi({\bm{\phi}}_t) + \partial \Phi_{{\bm{\phi}}_t}[\hat{\vx}_{t+2}] + \frac{1}{2} \partial^2 \Phi_{{\bm{\phi}}_t}[\hat{\vx}_{t+2}, \hat{\vx}_{t+2}] + O(\abssm{h_t}^3 \eta^3). \label{eq:ful-drift-2-step-phi-taylor2}
\end{align}
By definition, $\Phi({\bm{\phi}}_t) = {\bm{\phi}}_t$. By \eqref{eq:ful-drift-2-step-hatx-tan}
and \Cref{lm:Phi-proj-ful},
\begin{align*}
\partial\Phi_{{\bm{\phi}}_t}[\hat{\vx}_{t+2}]
&= \partial \Phi_{{\bm{\phi}}_t}[{\bm{x}}_t] - 2 \eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + \Osm{\alpha h_t^2 \eta^3}.
\end{align*}
Note that ${\bm{x}}_t = \eta h_t \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(h_t^2\eta^2)$
since the state at step $t$ is $O(\abssm{h_t}\eta)$-misaligned.
Also note that the identity
$\partial^2 \Phi_{{\bm{\phi}}_t}[\vv^{\mathrm{H}}_1({\bm{\phi}}_t), \vv^{\mathrm{H}}_1({\bm{\phi}}_t)] = {\bm{0}}$
holds
by \Cref{lm:Phi-2nd-normal-ful}. Then
\begin{align*}
\partial^2 \Phi_{{\bm{\phi}}_t}[{\bm{x}}_t, {\bm{x}}_t]
&= \partial^2 \Phi_{{\bm{\phi}}_t}[\eta h_t \vv^{\mathrm{H}}_1({\bm{\phi}}_t), \eta h_t \vv^{\mathrm{H}}_1({\bm{\phi}}_t)] + O(\abssm{h_t}^3 \eta^3) \\
&= O(\abssm{h_t}^3 \eta^3).
\end{align*}
Similarly, we have
$\partial^2 \Phi_{{\bm{\phi}}_t}[\hat{\vx}_{t+2}, \hat{\vx}_{t+2}] = O(\abssm{h_t}^3 \eta^3)$
since
\eqref{eq:ful-drift-2-step-hatx-tan} implies $\normtwosm{\mPHz({\bm{\phi}}_t) \hat{\vx}_{t+2}} = O(h_t^2 \eta^2)$
and
\eqref{eq:ful-drift-2-step-hatx-w} implies $\normtwosm{\mPHnzt({\bm{\phi}}_t) \hat{\vx}_{t+2}} = O(h_t^2 \eta^2)$.

Now we can prove \eqref{eq:ful-drift-2-step-phi} by
subtracting \eqref{eq:ful-drift-2-step-phi-taylor2} with \eqref{eq:ful-drift-2-step-phi-taylor1}:
\begin{align*}
  {\bm{\phi}}_{t+2} - {\bm{\phi}}_t
  &= \left(\partial \Phi_{{\bm{\phi}}_t}[\hat{\vx}_{t+2}] - \partial \Phi_{{\bm{\phi}}_t}[{\bm{x}}_t]\right)
    + \tfrac{1}{2} \left( \partial^2 \Phi_{{\bm{\phi}}_t}[\hat{\vx}_{t+2}, \hat{\vx}_{t+2}] - \partial^2 \Phi_{{\bm{\phi}}_t}[{\bm{x}}_t, {\bm{x}}_t] \right) + O(\abssm{h_t}^3 \eta^3) \\
  &= \left(
  - 2 \eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + \Osm{\alpha h_t^2 \eta^3}
  \right) + O(\abssm{h_t}^3 \eta^3) + O(\abssm{h_t}^3 \eta^3) \\
  &= - 2 \eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + \Osm{\alpha h_t^2 \eta^3}.
\end{align*}
Finally, for $\log \lambda^{\mathrm{H}}_1({\bm{\phi}}_{t+2})$ we have
\begin{align*}
  \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_{t+2}) - \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)
  &= \dotpsm{\nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}{{\bm{\phi}}_{t+2} - {\bm{\phi}}_t} + O((h_t^2 \eta^2)^2) \\
  &= \dotpsm{\nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}{-2 \eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)} + \Osm{\alpha h_t^2 \eta^3 + h_t^4 \eta^4} \\
  &= -2 \eta^2 h_t^2 \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2 + \Osm{\alpha h_t^2 \eta^3},
\end{align*}
which proves \eqref{eq:ful-drift-2-step-loglam}.
\end{proof}

\begin{proof}[Proof for \Cref{lm:ful-drift-2-step}]
We have verified \eqref{eq:ful-drift-2-step-phi} in \Cref{lm:ful-drift-2-step-phi}.
By \eqref{eq:ful-drift-2-step-phi} and definitions of $\hat{\vx}_{t+2}$ and ${\bm{x}}_{t+2}$, we have
\begin{align*}
  {\bm{x}}_{t+2} - \hat{\vx}_{t+2} = {\bm{\phi}}_t - {\bm{\phi}}_{t+2} = 2 \eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + \Osm{\alpha h_t^2 \eta^3}.
\end{align*}
And we can write ${\bm{x}}_{t+2} - \hat{\vx}_{t+2} = {\bm{\phi}}_t - {\bm{\phi}}_{t+2} = \Osm{h_t^2 \eta^2}$ as a loose approximation.

\myparagraph{Approximation for $h_{t+2}$.}
For $h_{t+2}$, we have
\begin{align*}
  h_{t+2} - \hat{h}_{t+2}
  &= \dotpsm{{\bm{x}}_{t+2}}{\vv^{\mathrm{H}}_1({\bm{\phi}}_{t+2})} - \dotpsm{\hat{\vx}_{t+2}}{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)} \\
  &= \dotpsm{{\bm{x}}_{t+2}}{\vv^{\mathrm{H}}_1({\bm{\phi}}_{t+2}) - \vv^{\mathrm{H}}_1({\bm{\phi}}_t)} + \dotpsm{{\bm{x}}_{t+2} - \hat{\vx}_{t+2}}{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)}\\
  &= O(\abssm{h_t}\eta) \cdot O(\normtwosm{{\bm{\phi}}_{t+2} - {\bm{\phi}}_t}) + 2\eta^2 h_t^2 \dotpsm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)} + O(\alpha h_t^2 \eta^3) \\
  &= O(\abssm{h_t}^3 \eta^3) + 0 + \Osm{\alpha h_t^2 \eta^3} \\
  &= \Osm{\alpha h_t^2 \eta^3},
\end{align*}
where the fourth equality is due to
$\dotpsm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)} = 0$ and
$\normtwosm{{\bm{\phi}}_{t+2} - {\bm{\phi}}_t} = O(h_t^2 \eta^2)$.
Combining this with \eqref{eq:ful-drift-2-step-hath} proves the claimed approximation \eqref{eq:ful-drift-2-step-h}.

\myparagraph{Approximation for $\normtwosm{\mPHnzt({\bm{\phi}}_{t+2}) {\bm{x}}_{t+2}}$.}
For $\mPHnzt({\bm{\phi}}_{t+2}) {\bm{x}}_{t+2}$, we have
\begin{align*}
  \mPHnzt({\bm{\phi}}_{t+2}) {\bm{x}}_{t+2} -
  \mPHnzt({\bm{\phi}}_t) \hat{\vx}_{t+2}
  &=
  (\mPHnzt({\bm{\phi}}_{t+2}) - \mPHnzt({\bm{\phi}}_t)) {\bm{x}}_{t+2}
  + \mPHnzt({\bm{\phi}}_t) ({\bm{x}}_{t+2} -\hat{\vx}_{t+2}) \\
  &=
  \Osm{\normtwosm{{\bm{\phi}}_{t+2} - {\bm{\phi}}_t}} \cdot O(\abssm{h_t}\eta)
  + O(\normtwosm{{\bm{x}}_{t+2} -\hat{\vx}_{t+2}}) \\
  &=
  O(h_t^2 \eta^2) \cdot O(\abssm{h_t}\eta)
  + O(h_t^2 \eta^2) \\
  &= O(h_t^2 \eta^2),
\end{align*}
where the third equality is due to
${\bm{x}}_{t+2} - \hat{\vx}_{t+2} = {\bm{\phi}}_t - {\bm{\phi}}_{t+2} = O(h_t^2 \eta^2)$.
Combining this with \eqref{eq:ful-drift-2-step-hatx-w} proves the claimed approximation \eqref{eq:ful-drift-2-step-x-w}.

\myparagraph{Approximation for $u_{t+2}$.}
Now we prove the formula for $u_{t+2}$.
Note that ${\bm{\phi}}_{t+2} - {\bm{\phi}}_t = O(h_t^2 \eta^2)$ implies that $\mu_{t+2} - \mu_t = O(h_t^2 \eta^2)$.
Then we have
\begin{align*}
  \eta\left(u_{t+2} - \hat{u}_{t+2}\right) &= (\mu_{t+2}^2 - \mu_t^2) \tilde{v}_{t+2} \\
  &= (\mu_{t+2} - \mu_t) (2 \mu_t + O(h_t^2 \eta^2)) \cdot (\tfrac{1}{\mu_t^2} + O(\alpha \eta)) \\
  &= (\mu_{t+2} - \mu_t) \left( 2 \mu_t \cdot \tfrac{1}{\mu_t^2} + O(\alpha \eta) \right) \\
  &= \tfrac{2}{\mu_t}(\mu_{t+2} - \mu_t) + O(\alpha h_t^2 \eta^3).
\end{align*}
Note that $\log \mu_{t+2} - \log \mu_t = \log\left(1 + \frac{1}{\mu_t}(\mu_{t+2} - \mu_t)\right) = \tfrac{1}{\mu_t}(\mu_{t+2} - \mu_t) + O((h_t^2 \eta^2)^2)$.
By \eqref{eq:ful-drift-2-step-loglam},
\begin{align*}
  \log \mu_{t+2} - \log \mu_t = \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) - \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_{t+2})
  = 2 \eta^2 h_t^2 \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2 + O(\alpha h_t^2 \eta^3).
\end{align*}
Combining these together gives the following approximation for $\frac{1}{\mu_t}(\mu_{t+2} - \mu_t)$:
\begin{align*}
  \tfrac{1}{\mu_t}(\mu_{t+2} - \mu_t) &= \log \mu_{t+2} - \log \mu_t + O(h_t^4 \eta^4) \\
  &= 2 \eta^2 h_t^2 \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2 + \Osm{\alpha h_t^2 \eta^3}.
\end{align*}
So $\eta \left(u_{t+2} - \hat{u}_{t+2}\right) = 4 \eta^2 h_t^2 \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2 + O(\alpha h_t^2 \eta^3)$.
Then by \eqref{eq:ful-drift-2-step-hatu}, we have
\begin{align*}
  u_{t+2} &= \hat{u}_{t+2} + 4 \eta h_t^2 \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2 + \Osm{\alpha h_t^2 \eta^2} \\
  &= u_t + 8 \eta C_{\mathrm{b}} h_t^2 + 4 \eta h_t^2 \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2 - 2 \eta C_{\mathrm{b}} + \Obig{\alpha(1+h_t^2)\eta^2} \\
  &= u_t + 4 \eta h_t^2 (2 C_{\mathrm{b}} + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2) - 2 \eta C_{\mathrm{b}} + \Osm{\alpha(1+h_t^2)\eta^2},
\end{align*}
which proves the claimed approximation \eqref{eq:ful-drift-2-step-u}.
\end{proof}
\section{Reduction to RMS-drift Process: The Case of Full Space Optimization} \label{sec:ful-proof}
In this section, we let $\{({\bm{\theta}}_t, v_t)\}_{t \ge 0}$ be a trajectory of gradient descent
with quasi-RMSprop scheduler,
and let $\eta, \beta$ be the base learning rate and decay rate.
We follow the notations and terminologies in \Cref{sec:proof-outline-main}.

\subsection{Good Initialization} \label{sec:ful-good-init}

\begin{proof}
[Proof for \Cref{lm:ful-good-init}]
Let $\hat{\vx}_0 := {\bm{\theta}}_0 - {\bm{\zeta}}_0$
and $r = \normtwosm{\hat{\vx}_0}$ for short.
By Gaussian concentration and anti-concentration,
with probability $1- \delta$, the following holds:
\begin{align*}
&\Omega(\delta \sigma_0) \le r \le O(\sigma_0 \sqrt{\log(1/\delta)}), \\
&\abssm{\dotpsm{\hat{\vx}_0}{\vv^{\mathrm{H}}_1({\bm{\zeta}}_0)}} \ge \Omega(\delta r).
\end{align*}
By Taylor expansion of $\Phi$ and \Cref{lm:Phi-proj-ful},
\begin{align*}
{\bm{\phi}}_0 &= {\bm{\zeta}}_0 + \partial \Phi_{{\bm{\zeta}}_0}[\hat{\vx}_0] + O(r^2) \\
&= {\bm{\zeta}}_0 + \mPHz({\bm{\zeta}}_0) \hat{\vx}_0 + O(r^2).
\end{align*}
Thus we can approximate ${\bm{x}}_0$ by
\begin{align*}
    {\bm{x}}_0 = \hat{\vx}_0 + {\bm{\zeta}}_0 - {\bm{\phi}}_0 = ({\bm{I}} - \mPHz) \hat{\vx}_0 + O(r^2).
\end{align*}
Now we give a lower bound for $\abssm{h_0} \eta$:
\begin{align*}
\abssm{h_0} \eta = \abssm{\dotpsm{{\bm{x}}_0}{\vv^{\mathrm{H}}_1({\bm{\phi}}_0)}}
&\ge
\abssm{\dotpsm{\hat{\vx}_0}{\vv^{\mathrm{H}}_1({\bm{\zeta}}_0)}}
- \abssm{\dotpsm{\hat{\vx}_0 - {\bm{x}}_0}{\vv^{\mathrm{H}}_1({\bm{\zeta}}_0)}}
- \normtwosm{\hat{\vx}_0} \cdot \normtwosm{\vv^{\mathrm{H}}_1({\bm{\zeta}}_0) - \vv^{\mathrm{H}}_1({\bm{\phi}}_0)} \\
&\ge \Omega(\delta r) - O(r^2) - O(r^2) \\
&\ge \Omega(r) \cdot (\Omega(\delta) - O(r)).
\end{align*}
Since $r = O(\sigma_0 \sqrt{\log(1/\delta)})$,
we can choose $\delta := C_0 \alpha_0 \eta \sqrt{\log(1/\eta)}$
with a large enough $C_0$ such that the above inequality gives $\abssm{h_0} \eta \ge \Omega(\delta r)$.

Now we verify the conditions claimed in the lemma statement.
First,
we can see from the following that the initial state is $\Osm{\alpha_0 \sqrt{\log(1/\delta)}}$-bounded:
\begin{align*}
    \normtwosm{{\bm{x}}_0} &\le \normtwosm{({\bm{I}} - \mPHz) \hat{\vx}_0} + O(r^2) \le O(r) \le O(\alpha_0 \eta \sqrt{\log(1/\delta)}) \\
    \abssm{u_0} &= \tfrac{1}{\eta} \abssm{\mu_t^2 \tilde{v}_0 - 1} \le O(\alpha_0).
\end{align*}
It is also $O(\alpha_0 + \sqrt{\log(1/\delta)})$-deviated
since
\[
    \normtwosm{{\bm{x}}_0} \ge \abssm{h_0} \eta \ge \Omega(\delta r) \ge \Omega(\delta^2 \exp(-\alpha_0^2)) \ge \eta \exp\left(-\Obig{\alpha_0 + \sqrt{\log(1/\delta)}}^2\right).
\]
Next, we verify that the initial state is $O(1/\delta)$-misaligned:
\[
    \normtwosm{\mPHnzt({\bm{\phi}}_0) {\bm{x}}_0} \le \normtwosm{{\bm{x}}_0} \le O(r) \le O(1/\delta) \cdot \abssm{h_0} \eta.
\]
Finally, $\normtwosm{{\bm{\phi}}_0 - {\bm{\zeta}}_0} \le O(r) \le O(\alpha_0 \eta \sqrt{\log(1/\delta)})$.
\end{proof}
\section{Lemmas for Gradient Descent} \label{sec:gd-proof}

\subsection{Full Space Optimization}

\begin{lemma} \label{lm:gd-1-step}
  Under \Cref{ass:man-ful},
  consider one step of gradient descent
  ${\bm{\theta}}_{t+1} = {\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)$
  with effective learning rate $\tilde{\eta}_t$.
  Let ${\bm{\phi}} \in \gZ^{\epsilon_0} \cap \manGa$ be a local minimizer.
  If for some parameters $\alpha = \eta^{-o(1)}$ and $r = O(\alpha)$,
  $\normtwosm{{\bm{\theta}}_t - {\bm{\phi}}} = O(r\eta)$
  and $\tilde{\eta}_t = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)$,
  then after one step, we have the following approximations for ${\bm{\theta}}_{t+1}$:
  \begin{enumerate}
    \item Zeroth-order approximation:
      \begin{align*}
        {\bm{\theta}}_{t+1} - {\bm{\phi}} = O(r \eta).
      \end{align*}
    \item First-order approximation:
      \begin{align*}
        {\bm{\theta}}_{t+1} - {\bm{\phi}}
        &= \left({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})\right) ({\bm{\theta}}_t - {\bm{\phi}}) + O(r^2 \eta^2) \\
        &= \left({\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})\right) ({\bm{\theta}}_t - {\bm{\phi}}) + O(\alpha r \eta^2).
      \end{align*}
    \item Second-order approximation:
      \begin{align*}
        {\bm{\theta}}_{t+1} - {\bm{\phi}}
        &= \left({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})\right) ({\bm{\theta}}_t - {\bm{\phi}}) - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{\theta}}_t - {\bm{\phi}}, {\bm{\theta}}_t - {\bm{\phi}}] + O(\alpha r^2 \eta^3).
      \end{align*}
  \end{enumerate}
\end{lemma}
\begin{proof}
  For the zeroth-order approximation, we expand $\mathcal{L}({\bm{\theta}}_t)$ around ${\bm{\phi}}$ using Taylor expansion:
  \begin{align*}
    \nabla \mathcal{L}({\bm{\theta}}_t) &= \nabla \mathcal{L}({\bm{\phi}}) + O(r \eta) = O(r\eta).
  \end{align*}
  So ${\bm{\theta}}_{t+1} - {\bm{\phi}} = {\bm{\theta}}_t - {\bm{\phi}} - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t) = O(r\eta)$.

  For the first-order approximation, we expand $\mathcal{L}({\bm{\theta}}_t)$ around ${\bm{\phi}}$ using Taylor expansion:
  \begin{align*}
    \nabla \mathcal{L}({\bm{\theta}}_t) &= \nabla \mathcal{L}({\bm{\phi}}) + {\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}}) + O((r \eta)^2) \\
    &= {\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}}) + O(r^2 \eta^2).
  \end{align*}
  So ${\bm{\theta}}_{t+1} - {\bm{\phi}} = {\bm{\theta}}_t - {\bm{\phi}} - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t) = ({\bm{I}} - \tilde{\eta}_t{\bm{H}}({\bm{\phi}}))({\bm{\theta}}_t - {\bm{\phi}}) + O(r^2\eta^2)$.

  For the second-order approximation, we again use Taylor expansion to expand $\mathcal{L}({\bm{\theta}}_t)$:
  \begin{align*}
    \nabla \mathcal{L}({\bm{\theta}}_t) &= \nabla \mathcal{L}({\bm{\phi}}) + {\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}}) + \frac{1}{2}\partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{\theta}}_t - {\bm{\phi}}, {\bm{\theta}}_t - {\bm{\phi}}] + O((r \eta)^3) \\
    &= {\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}}) + \frac{1}{2}\partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{\theta}}_t - {\bm{\phi}}, {\bm{\theta}}_t - {\bm{\phi}}] + O(r^3 \eta^3).
  \end{align*}
  So we have
  \begin{align*}
    {\bm{\theta}}_{t+1} - {\bm{\phi}} &= {\bm{\theta}}_t - {\bm{\phi}} - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t) \\
    &= ({\bm{I}} - \tilde{\eta}_t{\bm{H}}({\bm{\phi}}))({\bm{\theta}}_t - {\bm{\phi}}) - \frac{\tilde{\eta}_t}{2}\partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{\theta}}_t - {\bm{\phi}}, {\bm{\theta}}_t - {\bm{\phi}}] + O(r^3 \eta^3) \\
    &= ({\bm{I}} - \tilde{\eta}_t{\bm{H}}({\bm{\phi}}))({\bm{\theta}}_t - {\bm{\phi}}) - \left(\tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)\right) \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{\theta}}_t - {\bm{\phi}}, {\bm{\theta}}_t - {\bm{\phi}}] + O(r^3 \eta^3) \\
    &= ({\bm{I}} - \tilde{\eta}_t{\bm{H}}({\bm{\phi}}))({\bm{\theta}}_t - {\bm{\phi}}) - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{\theta}}_t - {\bm{\phi}}, {\bm{\theta}}_t - {\bm{\phi}}] + O(\alpha r^2 \eta^3),
  \end{align*}
  where the last equality uses $r = O(\alpha)$.
\end{proof}

\begin{lemma} \label{lm:gd-2-step}
  Under \Cref{ass:man-ful},
  consider two steps of gradient descent
  with effective learning rates $\tilde{\eta}_t$ and $\tilde{\eta}_{t+1}$:
  \begin{align*}
    {\bm{\theta}}_{t+1} &= {\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t), \\
    {\bm{\theta}}_{t+2} &= {\bm{\theta}}_{t+1} - \tilde{\eta}_{t+1} \nabla \mathcal{L}({\bm{\theta}}_{t+1}).
  \end{align*}
  Let ${\bm{\phi}} \in \gZ^{\epsilon_0} \cap \manGa$ be a local minimizer.
  If for some parameters $\alpha = \eta^{-o(1)}$ and $r = O(\alpha)$,
  $\normtwosm{{\bm{\theta}}_t - {\bm{\phi}}} = O(r\eta)$,
  $\tilde{\eta}_t = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)$,
  and $\tilde{\eta}_{t+1} = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)$,
  then after two steps,
  \[
    {\bm{\theta}}_{t+2} - {\bm{\phi}} = ({\bm{I}} - \tilde{\eta}_{t+1} {\bm{H}}({\bm{\phi}}))({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})) ({\bm{\theta}}_t - {\bm{\phi}}) - {\bm{\psi}}_{{\bm{\phi}}}({\bm{\theta}}_t - {\bm{\phi}}) + O(\alpha r^2 \eta^3),
  \]
  where ${\bm{\psi}}_{{\bm{\phi}}}(\hat{\vx}) := \frac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})}\left( {\bm{U}}_{{\bm{\phi}}} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}, \hat{\vx}] + \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{U}}_{\bm{\phi}} \hat{\vx}, {\bm{U}}_{\bm{\phi}} \hat{\vx}] \right)$,
  ${\bm{U}}_{{\bm{\phi}}} := {\bm{I}} - \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})$.
\end{lemma}
\begin{proof}
Let $\hat{\vx}_{\tau} := {\bm{\theta}}_{\tau} - {\bm{\phi}}$ for all $\tau \in \{t, t+1, t+2\}$.
By \Cref{lm:gd-1-step}, we have
$\hat{\vx}_{t+1} = O(r \eta)$ and the following first-order and second-order approximations for $\hat{\vx}_{t+1}$:
\begin{align*}
    \hat{\vx}_{t+1} &= {\bm{U}}_{\bm{\phi}} \hat{\vx}_t + O(\alpha r \eta^2). \\
    \hat{\vx}_{t+1} &= ({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})) \hat{\vx}_t - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}_t, \hat{\vx}_t] + O(\alpha r^2 \eta^3).
\end{align*}
Note that $\tilde{\eta}_{t+1} = \tilde{\eta}_t + O((1+r^2) \eta^2) = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)$.
By \Cref{lm:gd-1-step} again, we have the following second-order approximation for $\hat{\vx}_{t+2}$:
\begin{align*}
    \hat{\vx}_{t+2} &= ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}({\bm{\phi}})) \hat{\vx}_{t+1} - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}_{t+1}, \hat{\vx}_{t+1}] + O(\alpha r^2\eta^3).
\end{align*}
Now we combine the two steps together.
\begin{align*}
    \hat{\vx}_{t+2} &= ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}({\bm{\phi}}))\left(({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})) \hat{\vx}_t - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}_t, \hat{\vx}_t] + O(\alpha r^2 \eta^3)\right) \\
    & \qquad - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{U}}_{\bm{\phi}} \hat{\vx}_t + O(\alpha r \eta^2), {\bm{U}}_{\bm{\phi}} \hat{\vx}_t + O(\alpha r \eta^2)] + O(\alpha r^2 \eta^3) \\
    &= ({\bm{I}} - \tilde{\eta}_{t+1} {\bm{H}}({\bm{\phi}}))({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})) \hat{\vx}_t - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{U}}_{\bm{\phi}} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}_t, \hat{\vx}_t] \\
    & \qquad - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{U}}_{\bm{\phi}} \hat{\vx}_t, {\bm{U}}_{\bm{\phi}} \hat{\vx}_t] + O(\alpha r^2 \eta^3) \\
    &= ({\bm{I}} - \tilde{\eta}_{t+1} {\bm{H}}({\bm{\phi}}))({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})) \hat{\vx}_t - {\bm{\psi}}_{{\bm{\phi}}}(\hat{\vx}_t) + O(\alpha r^2 \eta^3),
\end{align*}
where the second equality uses ${\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}({\bm{\phi}}) = {\bm{U}}_{\bm{\phi}} + O(\alpha \eta)$.
\end{proof}

The following lemma characterizes the function ${\bm{\psi}}_{{\bm{\phi}}}(\hat{\vx})$
in \Cref{lm:gd-2-step} when $\hat{\vx}$ is the top eigenvector of ${\bm{H}}({\bm{\phi}})$.
We will use this property later in \Cref{sec:ful-proof-drifting}.
\begin{lemma} \label{lm:ful-psi-formula}
  Under \Cref{ass:man-ful,ass:topeigen-ful}, for ${\bm{\phi}} \in \mathit{\Gamma}$,
\begin{align*}
  {\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}})) = (2{\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})) \nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}),
\end{align*}
where ${\bm{\psi}}_{{\bm{\phi}}}$ is defined as in \Cref{lm:gd-2-step}.
Moreover,
\begin{align*}
  \dotpsm{{\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}}))}{\vv^{\mathrm{H}}_1({\bm{\phi}})} &= 0, \\
  \mPHz({\bm{\phi}}) {\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}})) &= 2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}).
\end{align*}
\end{lemma}
\begin{proof}
Note that ${\bm{U}}_{{\bm{\phi}}} \vv^{\mathrm{H}}_1({\bm{\phi}}) = - \vv^{\mathrm{H}}_1({\bm{\phi}})$. Then we can rewrite
${\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}}))$ as follows:
\begin{align*}
  {\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}}))
  &= \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \left( {\bm{U}}_{{\bm{\phi}}} \partial^3\mathcal{L}_{{\bm{\phi}}}[\vv^{\mathrm{H}}_1({\bm{\phi}}), \vv^{\mathrm{H}}_1({\bm{\phi}})] + \partial^3 \mathcal{L}_{{\bm{\phi}}}[-\vv^{\mathrm{H}}_1({\bm{\phi}}), -\vv^{\mathrm{H}}_1({\bm{\phi}})] \right) \\
  &= \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} (2{\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})) \partial^3\mathcal{L}_{{\bm{\phi}}}[\vv^{\mathrm{H}}_1({\bm{\phi}}), \vv^{\mathrm{H}}_1({\bm{\phi}})] \\
  &= \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} (2{\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})) \nabla \lambda^{\mathrm{H}}_1({\bm{\phi}}) \\
  &= (2{\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})) \nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}).
\end{align*}
Since $(2{\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})) \vv^{\mathrm{H}}_1({\bm{\phi}}) = 2\vv^{\mathrm{H}}_1({\bm{\phi}}) - 2\vv^{\mathrm{H}}_1({\bm{\phi}}) = {\bm{0}}$,
we have $\dotpsm{{\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}}))}{\vv^{\mathrm{H}}_1({\bm{\phi}})} = 0$.

The projection matrix
onto the tangent space of $\mathit{\Gamma}$ at ${\bm{\phi}}$
equals to $\mPHz({\bm{\phi}})$
by \Cref{lm:Phi-proj-ful}.
Then
\begin{align*}
  \mPHz({\bm{\phi}}) {\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}}))
  = \mPHz({\bm{\phi}})
    (2{\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})) \nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}})
  &= (2\mPHz({\bm{\phi}}) - {\bm{0}}) \nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}})
  \\
  &= 2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}),
\end{align*}
which proves the lemma.
\end{proof}

\subsection{Spherical Optimization}

\begin{lemma} \label{lm:pgd-1-step}
  Under \Cref{ass:man},
  consider one step of projected gradient descent on $\mathbb{S}^{D-1}$,
  ${\bm{\theta}}_{t+1} = \Pi({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t))$
  with effective learning rate $\tilde{\eta}_t$.
  Let ${\bm{\phi}} \in \gZ^{\epsilon_0} \cap \manGa$ be a local minimizer.
  If for some parameters $\alpha = \eta^{-o(1)}$ and $r = O(\alpha)$,
  $\normtwosm{{\bm{\theta}}_t - {\bm{\phi}}} = O(r\eta)$
  and $\tilde{\eta}_t = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)$,
  then after one step, we have the following approximations for ${\bm{\theta}}_{t+1}$:
  \begin{enumerate}
    \item Zeroth-order approximation:
      \begin{align*}
        {\bm{\theta}}_{t+1} - {\bm{\phi}} = O(r \eta).
      \end{align*}
    \item First-order approximation:
      \begin{align*}
        {\bm{\theta}}_{t+1} - {\bm{\phi}}
        &= \left({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})\right) ({\bm{\theta}}_t - {\bm{\phi}}) + O(r^2 \eta^2) \\
        &= \left({\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})\right) ({\bm{\theta}}_t - {\bm{\phi}}) + O(\alpha r \eta^2).
      \end{align*}
    \item Second-order approximation:
      \begin{align*}
        {\bm{\theta}}_{t+1} - {\bm{\phi}}
        &= \left({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})\right) ({\bm{\theta}}_t - {\bm{\phi}}) - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{\theta}}_t - {\bm{\phi}}, {\bm{\theta}}_t - {\bm{\phi}}] \\
        & \qquad - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 {\bm{\phi}} + O(\alpha r^2 \eta^3).
      \end{align*}
  \end{enumerate}
\end{lemma}
\begin{proof}
  Let $\hat{\vtheta}_{t+1} = {\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)$.
  Then ${\bm{\theta}}_{t+1} = \frac{\hat{\vtheta}_{t+1}}{\normtwosm{\hat{\vtheta}_{t+1}}}$.

  Since $\nabla \mathcal{L}({\bm{\theta}}_t)$ is perpendicular to ${\bm{\theta}}_t$,
  $\normtwosm{\hat{\vtheta}_{t+1}} = \sqrt{\normtwosm{{\bm{\theta}}_t}^2 + \tilde{\eta}_t^2 \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 }$.
  By Taylor expansion, $\nabla \mathcal{L}({\bm{\theta}}_t) = O(r \eta)$.
  Then the norm of $\hat{\vtheta}_{t+1}$ can be estimated by
  \begin{align*}
    \normtwosm{\hat{\vtheta}_{t+1}} = \sqrt{\normtwosm{{\bm{\theta}}_t}^2 + \tilde{\eta}_t^2 \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 }
    = \sqrt{1 + O(r^2\eta^2) } = 1 + O(r^2 \eta^2).
  \end{align*}
  Then ${\bm{\theta}}_{t+1} = \frac{\hat{\vtheta}_{t+1}}{\normtwosm{\hat{\vtheta}_{t+1}}} = \hat{\vtheta}_{t+1} + O(r^2 \eta^2)$.
  Combining with \Cref{lm:gd-1-step} proves the zeroth- and first-order approximations.

  To prove the second-order approximation, we need a tighter estimate for the norm of $\hat{\vtheta}_{t+1}$.
  By Taylor expansion, we have $\nabla \mathcal{L}({\bm{\theta}}_t) = {\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}}) + O(r^2 \eta^2)$.
  Squaring both sides gives $\normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 = \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 + O(r^3 \eta^3)$.
  Then we have
  \begin{align*}
    \normtwosm{\hat{\vtheta}_{t+1}} &= \sqrt{\normtwosm{{\bm{\theta}}_t}^2 + \tilde{\eta}_t^2 \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 } \\
    &= \sqrt{1 + \left(\tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)\right)^2 (\normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 + O(r^3 \eta^3)) } \\
    &= \sqrt{1 + \tfrac{4}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 + O(\alpha r^2 \eta^3) } \\
    &= \sqrt{1 + \tfrac{4}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2} + O(\alpha r^2 \eta^3) \\
    &= 1 + \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 + O(\alpha r^2 \eta^3).
  \end{align*}
  So ${\bm{\theta}}_{t+1}$ can be estimated by
  \begin{align*}
    {\bm{\theta}}_{t+1} = \frac{\hat{\vtheta}_{t+1}}{\normtwosm{\hat{\vtheta}_{t+1}}}
    &= \frac{\hat{\vtheta}_t}{1 + \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 + O(\alpha r^2 \eta^3)} \\
    &= \left(1 - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2\right) \hat{\vtheta}_t + O(\alpha r^2 \eta^3).
  \end{align*}
  Then ${\bm{\theta}}_{t+1} - {\bm{\phi}}$ can be estimated by
  \begin{align*}
    {\bm{\theta}}_{t+1} - {\bm{\phi}}
    &= \left(1 - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2\right)
    \hat{\vtheta}_t - {\bm{\phi}} + O(\alpha r^2 \eta^3) \\
    &= \left(1 - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2\right) (\hat{\vtheta}_t - {\bm{\phi}}) \\
    & \qquad - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 {\bm{\phi}} + O(\alpha r^2 \eta^3) \\
    &= (\hat{\vtheta}_t - {\bm{\phi}}) + O(r^3 \eta^3) \\
    & \qquad - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 {\bm{\phi}} + O(\alpha r^2 \eta^3) \\
    &= \left({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})\right) ({\bm{\theta}}_t - {\bm{\phi}}) - \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{\theta}}_t - {\bm{\phi}}, {\bm{\theta}}_t - {\bm{\phi}}] \\
    & \qquad - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) ({\bm{\theta}}_t - {\bm{\phi}})}^2 {\bm{\phi}} + O(\alpha r^2 \eta^3),
  \end{align*}
  where the last equality uses the second-order approximation in \Cref{lm:gd-1-step}.
\end{proof}

\begin{lemma} \label{lm:pgd-2-step}
  Under \Cref{ass:man},
  consider two steps of projectioned gradient descent on $\mathbb{S}^{D-1}$
  with effective learning rates $\tilde{\eta}_t$ and $\tilde{\eta}_{t+1}$:
  \begin{align*}
    {\bm{\theta}}_{t+1} &= \Pi({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)), \\
    {\bm{\theta}}_{t+2} &= \Pi({\bm{\theta}}_{t+1} - \tilde{\eta}_{t+1} \nabla \mathcal{L}({\bm{\theta}}_{t+1})).
  \end{align*}
  Let ${\bm{\phi}} \in \gZ^{\epsilon_0} \cap \manGa$ be a local minimizer.
  If for some parameters $\alpha = \eta^{-o(1)}$ and $r = O(\alpha)$,
  $\normtwosm{{\bm{\theta}}_t - {\bm{\phi}}} = O(r\eta)$,
  $\tilde{\eta}_t = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)$,
  and $\tilde{\eta}_{t+1} = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)$,
  then after two steps,
  \[
    {\bm{\theta}}_{t+2} - {\bm{\phi}} = ({\bm{I}} - \tilde{\eta}_{t+1} {\bm{H}}({\bm{\phi}}))({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})) ({\bm{\theta}}_t - {\bm{\phi}}) - {\bm{\psi}}_{{\bm{\phi}}}({\bm{\theta}}_t - {\bm{\phi}}) + O(\alpha r^2 \eta^3),
  \]
  where
  \begin{align*}
    {\bm{\psi}}_{{\bm{\phi}}}(\hat{\vx}) &:=
    \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})}\left( {\bm{U}}_{{\bm{\phi}}} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}, \hat{\vx}] + \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{U}}_{\bm{\phi}} \hat{\vx}, {\bm{U}}_{\bm{\phi}} \hat{\vx}] \right) \\
    & \qquad\qquad + \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}\left( \normtwosm{{\bm{H}}({\bm{\phi}}) \hat{\vx}}^2 + \normtwosm{{\bm{H}}({\bm{\phi}}) {\bm{U}}_{\bm{\phi}} \hat{\vx}}^2 \right) {\bm{\phi}}, \\
    {\bm{U}}_{{\bm{\phi}}} &:= {\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}}).
  \end{align*}
\end{lemma}
\begin{proof}
Let $\hat{\vx}_{\tau} := {\bm{\theta}}_{\tau} - {\bm{\phi}}$ for all $\tau \in \{t, t+1, t+2\}$.
Let $\widehat{{\bm{\psi}}}_{{\bm{\phi}}}({\bm{x}})$ be the following function:
\begin{align*}
  \widehat{{\bm{\psi}}}_{{\bm{\phi}}}(\hat{\vx}) := \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}, \hat{\vx}] + \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) \hat{\vx}}^2 {\bm{\phi}}.
\end{align*}
By \Cref{lm:pgd-1-step}, we have
$\hat{\vx}_{t+1} = O(r \eta)$ and the following first-order and second-order approximations for $\hat{\vx}_{t+1}$:
\begin{align*}
    \hat{\vx}_{t+1} &= {\bm{U}}_{\bm{\phi}} \hat{\vx}_t + O(\alpha r \eta^2). \\
    \hat{\vx}_{t+1} &= ({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})) \hat{\vx}_t - \widehat{{\bm{\psi}}}_{{\bm{\phi}}}(\hat{\vx}_t) + O(\alpha r^2 \eta^3).
\end{align*}
Note that $\tilde{\eta}_{t+1} = \tilde{\eta}_t + O((1+r^2) \eta^2) = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} + O(\alpha \eta)$.
By \Cref{lm:pgd-1-step} again, we have the following second-order approximation for $\hat{\vx}_{t+2}$:
\begin{align*}
    \hat{\vx}_{t+2} &= ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}({\bm{\phi}})) \hat{\vx}_{t+1} - \widehat{{\bm{\psi}}}_{{\bm{\phi}}}(\hat{\vx}_{t+1}) + O(\alpha r^2\eta^3).
\end{align*}
Now we combine the two steps together.
\begin{align*}
    \hat{\vx}_{t+2} &= ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}({\bm{\phi}}))\left(({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}})) \hat{\vx}_t - \widehat{{\bm{\psi}}}_{{\bm{\phi}}}(\hat{\vx}_{t}) + O(\alpha r^2 \eta^3)\right) \\
    & \qquad - \widehat{{\bm{\psi}}}_{{\bm{\phi}}}( {\bm{U}}_{\bm{\phi}} \hat{\vx}_t + O(\alpha r \eta^2) ) + O(\alpha r^2 \eta^3) \\
    &= ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}({\bm{\phi}}))({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}) \hat{\vx}_t - ({\bm{U}}_{\bm{\phi}} + O(\alpha \eta))\widehat{{\bm{\psi}}}_{{\bm{\phi}}}(\hat{\vx}_{t}) + O(\alpha r^2 \eta^3) \\
    & \qquad - \left(\widehat{{\bm{\psi}}}_{{\bm{\phi}}}( {\bm{U}}_{\bm{\phi}} \hat{\vx}_t) + O(\alpha r^2 \eta^3)\right) + O(\alpha r^2 \eta^3) \\
    &= ({\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}({\bm{\phi}}))({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}) \hat{\vx}_t
    - \left({\bm{U}}_{\bm{\phi}} \widehat{{\bm{\psi}}}_{{\bm{\phi}}}(\hat{\vx}_{t}) + \widehat{{\bm{\psi}}}_{{\bm{\phi}}}( {\bm{U}}_{\bm{\phi}} \hat{\vx}_t) \right)
    + O(\alpha r^2 \eta^3).
\end{align*}
where the second equality uses ${\bm{I}} - \tilde{\eta}_{t+1}{\bm{H}}({\bm{\phi}}) = {\bm{U}}_{\bm{\phi}} + O(\alpha \eta)$.
Finally, we note that
\begin{align*}
  {\bm{U}}_{\bm{\phi}} \widehat{{\bm{\psi}}}_{{\bm{\phi}}}(\hat{\vx}) + \widehat{{\bm{\psi}}}_{{\bm{\phi}}}( {\bm{U}}_{\bm{\phi}} \hat{\vx}) &=
  \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{U}}_{\bm{\phi}} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}, \hat{\vx}] + \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) \hat{\vx}}^2 {\bm{U}}_{\bm{\phi}} {\bm{\phi}} \\
  & \qquad + \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{U}}_{\bm{\phi}} \hat{\vx}, {\bm{U}}_{\bm{\phi}} \hat{\vx}] + \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}  \normtwosm{{\bm{H}}({\bm{\phi}}) {\bm{U}}_{\bm{\phi}} \hat{\vx}}^2 {\bm{\phi}} \\
  &= \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})}\left( {\bm{U}}_{{\bm{\phi}}} \partial^3 \mathcal{L}_{{\bm{\phi}}}[\hat{\vx}, \hat{\vx}] + \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{U}}_{\bm{\phi}} \hat{\vx}, {\bm{U}}_{\bm{\phi}} \hat{\vx}] \right) \\
  & \qquad + \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2}\left( \normtwosm{{\bm{H}}({\bm{\phi}}) \hat{\vx}}^2 + \normtwosm{{\bm{H}}({\bm{\phi}}) {\bm{U}}_{\bm{\phi}} \hat{\vx}}^2 \right) {\bm{\phi}} \\
  &= {\bm{\psi}}_{{\bm{\phi}}}(\hat{\vx}),
\end{align*}
where the second equality uses the fact that ${\bm{U}}_{{\bm{\phi}}} = ({\bm{I}} - \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})) {\bm{\phi}} = {\bm{\phi}} + \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \nabla \mathcal{L}({\bm{\phi}}) = {\bm{\phi}}$ (\Cref{lm:scale-invariant-grad-hess}).
\end{proof}

\begin{lemma} \label{lm:sph-psi-formula}
  Under \Cref{ass:man,ass:topeigen},
  for ${\bm{\phi}} \in \mathit{\Gamma}$,
  \begin{align*}
    {\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}})) = (2{\bm{I}} - \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}_t) \nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}) + 4{\bm{\phi}},
  \end{align*}
  where ${\bm{\psi}}_{{\bm{\phi}}}$ is defined as in \Cref{lm:pgd-2-step}. Moreover,
\begin{align}
  \dotpsm{{\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}}))}{\vv^{\mathrm{H}}_1({\bm{\phi}})} &= 0, \label{eq:sph-psi-proj-v1} \\
  \mPHz({\bm{\phi}}) {\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}})) &= 2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}). \label{eq:sph-psi-proj-H0}
\end{align}
\end{lemma}
\begin{proof}
  Let ${\bm{V}}_{{\bm{\phi}}} := 2{\bm{I}} - \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} {\bm{H}}({\bm{\phi}})$.
    Using a similar argument as in \Cref{lm:ful-psi-formula},
    \[
      \tfrac{1}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} \left( {\bm{U}}_{{\bm{\phi}}} \partial^3\mathcal{L}_{{\bm{\phi}}}[\vv^{\mathrm{H}}_1({\bm{\phi}}), \vv^{\mathrm{H}}_1({\bm{\phi}})] + \partial^3 \mathcal{L}_{{\bm{\phi}}}[-\vv^{\mathrm{H}}_1({\bm{\phi}}), -\vv^{\mathrm{H}}_1({\bm{\phi}})] \right)
      = {\bm{V}}_{{\bm{\phi}}} \nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}).
    \]
    Also notice that
    \begin{align*}
    \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})^2} \left( \normtwosm{{\bm{H}}({\bm{\phi}}) \vv^{\mathrm{H}}_1({\bm{\phi}})}^2 + \normtwosm{{\bm{H}}({\bm{\phi}}) {\bm{U}}_{{\bm{\phi}}} \vv^{\mathrm{H}}_1({\bm{\phi}})}^2 \right) {\bm{\phi}}
    = \tfrac{2}{\lambda^{\mathrm{H}}_1({\bm{\phi}})} ( \lambda^{\mathrm{H}}_1({\bm{\phi}})^2 + \lambda^{\mathrm{H}}_1({\bm{\phi}})^2 ) {\bm{\phi}} = 4 {\bm{\phi}}.
    \end{align*}
    Combining these together proves
    ${\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}})) = {\bm{V}}_{{\bm{\phi}}} \nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}) + 4{\bm{\phi}}$.

    To obtain the last two equations \eqref{eq:sph-psi-proj-v1}, \eqref{eq:sph-psi-proj-H0},
    we first note that
    ${\bm{\phi}}$ is a $0$-eigenvector of ${\bm{H}}({\bm{\phi}})$
    since
    ${\bm{H}}({\bm{\phi}}) {\bm{\phi}} = -\nabla \mathcal{L}({\bm{\phi}}) = {\bm{0}}$
    by \Cref{lm:scale-invariant-grad-hess}.
    So we have $\dotpsm{{\bm{\phi}}}{\vv^{\mathrm{H}}_1({\bm{\phi}})} = 0$ and $\mPHz({\bm{\phi}}) {\bm{\phi}} = {\bm{\phi}}$.

    To prove \eqref{eq:sph-psi-proj-v1},
    we note that
    ${\bm{V}}_{{\bm{\phi}}} \vv^{\mathrm{H}}_1({\bm{\phi}}) = {\bm{0}}$.
    Then
    \begin{align*}
      \dotpsm{{\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}}))}{\vv^{\mathrm{H}}_1({\bm{\phi}})} = \dotpsm{\nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}})}{{\bm{V}}_{{\bm{\phi}}}\vv^{\mathrm{H}}_1({\bm{\phi}})} + 4\dotpsm{{\bm{\phi}}}{\vv^{\mathrm{H}}_1({\bm{\phi}})} = 0.
    \end{align*}

    To prove \eqref{eq:sph-psi-proj-H0},
    first we note that $\mPHz({\bm{\phi}}){\bm{V}}_{{\bm{\phi}}} = 2{\bm{I}}$ and $\mPHz({\bm{\phi}}) {\bm{\phi}} = {\bm{\phi}}$, which implies
    \begin{align*}
      \mPHz({\bm{\phi}}) {\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}})) = 2\mPHz({\bm{\phi}}) \nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}) + 4 {\bm{\phi}}.
    \end{align*}
    Then
    by \Cref{lm:Phi-proj-sph},
    we can decompose
    $\mPHz({\bm{\phi}})\nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}})$
    as a component parallel to ${\bm{\phi}}$ and a component perpendicular to ${\bm{\phi}}$:
    \begin{align*}
      \mPHz({\bm{\phi}})\nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}}) &=
      ({\bm{I}} - {\bm{\phi}} {\bm{\phi}}^\top)\mPHz({\bm{\phi}})\nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}})
      + \dotpsm{\nabla \log \lambda^{\mathrm{H}}_1({\bm{\phi}})}{{\bm{\phi}}} {\bm{\phi}} \\
      &=
      \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}})
      + \left.\frac{\partial}{\partial c} \log \lambda^{\mathrm{H}}_1(c{\bm{\phi}})\right|_{c=1}.
    \end{align*}
    For the second term, we note that
    $\lambda^{\mathrm{H}}_1(c {\bm{\phi}}) = c^{-2} \lambda^{\mathrm{H}}_1({\bm{\phi}})$ by scale-invariance,
    and thus we have $\frac{\partial}{\partial c}\log \lambda^{\mathrm{H}}_1(c {\bm{\phi}}) = -2 / c$
    and $\left.\frac{\partial}{\partial c} \log \lambda^{\mathrm{H}}_1(c{\bm{\phi}})\right|_{c=1}= -2$.
    Combining all these together gives
    \begin{align*}
      \mPHz({\bm{\phi}}) {\bm{\psi}}_{{\bm{\phi}}}(\vv^{\mathrm{H}}_1({\bm{\phi}})) =
      2\left(\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}})
      - 2 {\bm{\phi}}\right)
      + 4 {\bm{\phi}} = 2\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}),
    \end{align*}
    which proves \eqref{eq:sph-psi-proj-H0}.
\end{proof}

\section{Lemmas for Quasi-RMSprop Schedulers} \label{sec:lr-proof}

\begin{lemma} \label{lm:lr-1-step}
    Given gradients $\{{\bm{g}}_t\}_{t \ge 0}$,
    let $\tilde{\eta}_0, \tilde{\eta}_1, \tilde{\eta}_2, \dots$ be the effective learning rates produced by
    a quasi-RMSprop scheduler with base learning rate $\eta$
    and decay rate $\beta$,
    and let $\{\tilde{v}_t\}_{t \ge 0}$ be the corresponding moment estimating sequence.
    Consider the case of $\eta = o(1)$, $\beta = 1 - C_{\mathrm{b}} \eta^2 + O(\eta^4)$ for some $C_{\mathrm{b}} = \Theta(1)$.
    Fix $\mu = \Theta(1)$ and define $\hat{u}_t := \frac{1}{\eta}(\mu^2\tilde{v}_t - 1)$ for every $t \ge 0$.
    If $\normtwosm{{\bm{g}}_t} \le \alpha \eta$ and $\abssm{\mu^2 \tilde{v}_t - 1} \le \alpha \eta$
    for some $\alpha = (\tfrac{1}{\eta})^{o(1)}$ at step $t$,
    then the following holds
    \begin{align}
        \tilde{\eta}_t &= \mu \cdot \left(1 - \tfrac{1}{2}\eta \hat{u}_t\right) + O(\alpha^2 \eta^2). \label{eq:lr-1-step-eta} \\
        \hat{u}_{t+1} &= \hat{u}_t + C_{\mathrm{b}} \eta (  \mu^2 \bar{g}_t^2 - 1) + O(\alpha \eta^2).\label{eq:lr-1-step-u}
    \end{align}
\end{lemma}
\begin{proof}
    By definition of RMSprop-like learning rate scheduler,
    \begin{align}
        \tilde{\eta}_t         &= \frac{1}{\sqrt{\tilde{v}_t}} + O(\alpha^2 \eta^2) \label{eq:lr-1-step-a} \\
        \tilde{v}_{t+1} &= \beta \tilde{v}_t + (1-\beta) \bar{g}_t^2  + O(\mathrm{poly}(\alpha) \eta^4). \label{eq:lr-1-step-b}
    \end{align}
    By definition of $\hat{u}_t$,
    we can express $\tilde{v}_t$ as
    $\tilde{v}_t = \frac{1}{\mu^2}(1 + \eta \hat{u}_t)$.
    Combining this with \eqref{eq:lr-1-step-a} proves \eqref{eq:lr-1-step-eta}:
    \begin{align*}
        \tilde{\eta}_t
        = \frac{1}{\sqrt{\frac{1}{\mu^2}(1 + \eta \hat{u}_t)}}
        = \frac{\mu}{\sqrt{1 + \eta \hat{u}_t}}
        &= \mu \cdot \left(1 - \tfrac{1}{2}\eta \hat{u}_t + O(\alpha^2 \eta^2)\right) \\
        &= \mu \cdot \left(1 - \tfrac{1}{2}\eta \hat{u}_t\right) + O(\alpha^2 \eta^2).
    \end{align*}
    By substituting $\frac{1}{\mu^2}(1+\eta \hat{u}_t)$ for $\tilde{v}_t$ in \eqref{eq:lr-1-step-b} we have
    \begin{align*}
        \frac{1}{\mu^2}(1+\eta \hat{u}_{t+1}) =
        \beta \cdot \frac{1}{\mu^2}(1+\eta \hat{u}_t) + (1-\beta) \bar{g}_t^2  + O(\mathrm{poly}(\alpha) \eta^4).
    \end{align*}
    Multiplying $\mu^2$ and subtracting $1$ on both sides gives
    \begin{align*}
        \eta \hat{u}_{t+1}
        &= -(1 - \beta) + \beta \eta \hat{u}_t + (1-\beta) \mu^2 \bar{g}_t^2  + O(\mathrm{poly}(\alpha) \eta^4) \\
        &= \eta \hat{u}_t + (1 - \beta)\left(  \mu^2 \bar{g}_t^2 - \eta \hat{u}_t - 1\right) + O(\mathrm{poly}(\alpha) \eta^4).
    \end{align*}
    Now we divide $\eta$ on both sides. Then we have
    \begin{align*}
        \hat{u}_{t+1}
        &= \hat{u}_t + \left(C_{\mathrm{b}} \eta + O(\eta^3)\right) \left(  \mu^2 \bar{g}_t^2 - \eta \hat{u}_t - 1\right) + O(\mathrm{poly}(\alpha) \eta^3) \\
        &= \hat{u}_t + C_{\mathrm{b}} \eta \left(  \mu^2 \bar{g}_t^2 - \eta \hat{u}_t - 1\right) + O(\mathrm{poly}(\alpha) \eta^3) \\
        &= \hat{u}_t + C_{\mathrm{b}} \eta (  \mu^2 \bar{g}_t^2 - 1) + O(\alpha \eta^2),
    \end{align*}
    which proves \eqref{eq:lr-1-step-u}.
\end{proof}

\begin{lemma} \label{lm:lr-2-step}
    In the setting of \Cref{lm:lr-1-step},
    If $\normtwosm{{\bm{g}}_t} \le \alpha \eta$, $\normtwosm{{\bm{g}}_{t+1}} \le \alpha \eta$,
    and $\abssm{\mu^2 \tilde{v}_t - 1} \le \alpha \eta$
    for some $\alpha = (\tfrac{1}{\eta})^{o(1)}$ at step $t$,
    then the following holds
    \begin{align}
        \tilde{\eta}_{t+1} &= \mu \cdot \left(1 - \tfrac{1}{2}\eta \hat{u}_t\right) + O(\alpha^2 \eta^2). \label{eq:lr-2-step-eta}\\
        \hat{u}_{t+2} &= \hat{u}_t + C_{\mathrm{b}} \eta (  \mu^2 \bar{g}_t^2 + \mu^2 \bar{g}_{t+1}^2 - 2) + O(\alpha \eta^2). \label{eq:lr-2-step-u}
    \end{align}
\end{lemma}
\begin{proof}
    Note that \eqref{eq:lr-1-step-u} implies $\hat{u}_{t+1} = O(\alpha)$. Applying the
    inequalities \eqref{eq:lr-1-step-eta} and \eqref{eq:lr-1-step-u}
    to step $t+1$ proves \eqref{eq:lr-2-step-eta} and \eqref{eq:lr-2-step-u}.
\end{proof}
\section{Introduction}
Training modern deep neural nets crucially rely on normalization layers to make the training process less sensitive to hyperparameters and initialization.
Popular normalization layers include Batch Normalization (BN)~\citep{ioffe2015batch},
Layer Normalization~\citep{ba2016layer}, etc.
Normalization layers amount to a reparametrization of the neural net
so that
the loss becomes invariant to scale of most parameters
(and with a minor change, to {\em all} parameters): $\mathcal{L}(c{\bm{w}}) = \mathcal{L}({\bm{w}})$
for all scalings $c>0$~\citep{ioffe2015batch,arora2018theoretical,li2020exp}.
The current paper
assumes
this scale-invariance for all parameters and analyzes the trajectory of gradient descent with \textit{weight decay} (WD):
\begin{equation} \label{eq:upd-gdwd}
  {\bm{w}}_{t+1} \gets (1 - \hat{\eta} \hat{\lambda}){\bm{w}}_t - \hat{\eta} \nabla \mathcal{L}({\bm{w}}_t).
\end{equation}
Use of WD may appear nonsensical at first sight because traditionally it is used to penalize large parameter norm,
which of course is inconsequential for scale-invariant loss --- one can scale down the parameter norm arbitrarily without changing the loss value.
However, the scale of the parameter \textit{does} matter for gradient and Hessian,
and thus WD can affect the training dynamics.
In particular, simple calculus shows $\nabla \mathcal{L}({\bm{w}}) = \frac{1}{\normtwosm{{\bm{w}}}} \nabla
\mathcal{L}(\frac{{\bm{w}}}{\normtwosm{{\bm{w}}}})$ and $\nabla^2 \mathcal{L}({\bm{w}}) = \frac{1}{\normtwosm{{\bm{w}}}^2} \nabla^2
\mathcal{L}(\frac{{\bm{w}}}{\normtwosm{{\bm{w}}}})$, so WD is in effect trying to increase the norm of gradient and Hessian in training.
This makes the training dynamics very different from unnormalized nets,
and requires revisiting classical convergence analyses, as was done in various papers \citep{arora2018theoretical,li2020exp,li2022robust}.
More notably, such a change in dynamics also yields a very different solution at the end
--- in particular a solution that generalizes better~\citep{zhang2018three}.

The goal of the current paper is to improve mathematical understanding of
how normalization together with WD
can improve generalization.
While this may arise from many places,
we focus on
sharpness-based generalization measures,
and exhibit settings where normalization provably induces an implicit bias
to persistently reduce sharpness in the training process, which we call the \textit{sharpness-reduction bias}.
See \Cref{fig:matcom-intro,fig:full-cifar} for experiments on matrix completion (with BN)
and CIFAR-10.

\begin{figure}[t]
  \includegraphics[width=\textwidth]{fig/matcom-intro.pdf}
  \vspace{-0.25in}
  \caption{
    Experiment on overparameterized matrix completion with Batch Normalization.
    Given $800$ ($32\%$) entries $\Omega$ of a rank-$2$ matrix ${\bm{M}} \in \mathbb{R}^{50 \times 50}$,
    use GD+WD\xspace to optimize the loss $\mathcal{L}({\bm{U}}, {\bm{V}}) := \frac{1}{\abssm{\Omega}} \sum_{(i, j) \in \Omega} (\mathrm{BN}([{\bm{U}} {\bm{V}}^\top]_{i, j}) - M_{i,j})^2$,
    where ${\bm{U}}, {\bm{V}} \in \mathbb{R}^{50 \times 50}$ (thus no explicit constraint on rank).
    Starting from step $\sim 2000$, spherical sharpness drops significantly (\textbf{b}), which
    encourages low-rank~(\textbf{d}) and
    causes the test loss (MSE of all entries)
    to decrease from $1.12$ to $0.013$ (\textbf{a}). See also \Cref{sec:exp-matlin}.
  }
  \vspace{-0.2in}
  \label{fig:matcom-intro}
\end{figure}

\begin{figure}[t]
  \includegraphics[width=\textwidth]{fig/full-cifar.pdf}
  \vspace{-0.25in}
  \caption{
    Training a smooth and scale-invariant VGG-11 on CIFAR-10 for a longer time with (full-batch) GD+WD\xspace leads to lower sharpness and higher test accuracy.
    BN is added after every linear layer to ensure scale-invariance.
    $100\%$ training accuracy is achieved after $\sim 680$ steps (dotted line),
    but as the training continues for $~47k$ steps,
    the spherical sharpness keeps decreasing (\textbf{b})
    and the test accuracy increases from $69.1\%$ to $72.0\%$ (\textbf{a}).
    Then the training exhibits destabilization but the test accuracy is further boosted to $84.3\%$.
    See also \Cref{sec:exp-cifar10}.
  }
  \label{fig:full-cifar}
  \vspace{-0.1in}
\end{figure}

It is long believed that flatter minima generalize better~\citep{hochreiter1997flat,keskar2017large,neyshabur2017exploring},
but the notion of sharpness/flatness makes sense only if it is carefully defined
in consideration of various symmetries in neural nets.
One of the most straightforward measures of sharpness is
the maximum eigenvalue of Hessian, namely $\lambda_1(\nabla^2 \mathcal{L}({\bm{w}}_t))$.
But for normalized nets,
this sharpness measure
is vulnerable to weight rescaling, because one can scale the weight norm
to make a minimizer arbitrarily flat~\citep{dinh2017sharp}.
Also, this sharpness measure may not be decreasing with the number of training steps:
an empirical study by \citet{cohen2021gradient}
shows that for various neural nets (including normalized nets),
GD has an overwhelming tendency to persistently increase $\lambda_1(\nabla^2 \mathcal{L}({\bm{w}}_t))$
until it reaches the \textit{Edge of Stability (EoS) regime},
a regime where $\lambda_1(\nabla^2 \mathcal{L}({\bm{w}}_t))$ stabilizes around $2 / \hat{\eta}$ ($\hat{\eta}$ is the learning rate).
See also Figure \hyperref[fig:full-cifar]{\ref*{fig:full-cifar}c}.

The sharpness measure we use in this paper takes the scale-invariance property into account.
We note that techniques from previous works~\citep{mcallester2003simplified,neyshabur2017exploring,foret2021sam}
can be easily adopted here to establish a
PAC-Bayes bound on the test error, where our sharpness measure appears as an additive term (see \Cref{sec:pac-bayes}).
\begin{definition}[Spherical Sharpness] \label{def:sph-sha}
  For a scale-invariant loss $\mathcal{L}({\bm{w}})$ (i.e., $\mathcal{L}(c{\bm{w}}) = \mathcal{L}({\bm{w}})$ for all $c > 0$),
  the spherical sharpness at ${\bm{w}} \in \mathbb{R}^D$ is defined by $\lambda_1(\nabla^2 \mathcal{L}(\frac{{\bm{w}}}{\normtwosm{{\bm{w}}}}))$,
  the maximum eigenvalue of the Hessian matrix after projecting ${\bm{w}}$ onto the unit sphere.
\end{definition}

\subsection{Our Contributions}
In this work, we study the aforementioned sharpness-reduction bias of
gradient descent with weight decay (GD+WD\xspace),
assuming the loss $\mathcal{L}$ is scale-invariant due to the presence of normalization.
For constant learning rate $\hat{\eta}$ and weight decay $\hat{\lambda}$,
we can rewrite gradient descent \eqref{eq:upd-gdwd} on scale-invariant loss
equivalently as Projected Gradient Descent (PGD)
on the unit sphere with adaptive learning rates, ${\bm{\theta}}_{t+1} \gets \Pi({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t))$,
where ${\bm{\theta}}_t := \frac{{\bm{w}}_t}{\normtwosm{{\bm{w}}_t}}$ is the direction of ${\bm{w}}_t$,
and $\tilde{\eta}_t$ is the ``effective'' learning rate at step $t$ (defined in Section~\ref{sec:prelims}).
We call $\tilde{\eta}_t$ adaptive because it can be shown that $\tilde{\eta}_t$
increases when gradient is small and decreases when gradient is large,
which resembles the behaviors of adaptive gradient methods such as RMSprop~\citep{hinton2012rmsprop}.
Our main contributions are as follows:
\begin{enumerate}
  \item We theoretically show that once ${\bm{\theta}}_t$ reaches a point near the manifold of minimizers of $\mathcal{L}$, the
  effective learning rate $\tilde{\eta}_t$ will keep increasing until
  $2 / \tilde{\eta}_t$ roughly equals to the spherical sharpness,
  or in other words, gradient descent enters the EoS regime (\Cref{sec:main-result-1}).
  \item In the EoS regime, we show that for gradient descent with a small (but finite) learning rate,
  ${\bm{\theta}}_t$ oscillates around the manifold and moves approximately along a sharpness-reduction flow,
  which is a gradient flow for minimizing spherical sharpness on the manifold (with gradient-dependent learning rate) (\Cref{sec:main-result-2}).
  \item As an application of our theory,
  we show that for linear regression with BN, GD+WD\xspace finds the minimizer
  that corresponds to the linear model with minimum weight norm,
  which
  looks surprisingly the same as the conventional effect of WD
  but is achieved through the completely different sharpness-reduction mechanism
  (\Cref{sec:lin}).
  \item We experimentally verified the sharpness-reduction bias phenomena predicted by our theorem and its benefits to
  generalization on CIFAR-10 with scale-invariant VGG-11 and ResNet-20, as well as matrix completion with BN
  (\Cref{sec:add-exp}).
  \item We generalize our theoretical results of sharpness-reduction bias to a broader class of adaptive gradient methods,
  most notably a variant of RMSprop with scalar learning rate (\Cref{sec:qrms}). Our proof technique is novel and may have independent interest to the ML community.
\end{enumerate}

\myparagraph{Technical Contribution.}
The main challenge to establish our theorem is that we need to analyze the implicit bias of GD in the EoS regime. In particular, we crucially rely on step size being finite --- this is
in sharp contrast to many previous works on implicit bias of GD~\citep{soudry2018iclrImplicit,soudry2018implicit,lyu2020gradient,ji2020directional,gunasekar2018implicitconv, gunasekar2018implicit,li2018algorithmic,razin2020implicit,
arora2019implicit,chizat2020implicit,li2021towards,lyu2019gradient,
lyu2021gradient,razin2022implicit,stoger2021small,ge2021understanding} where the same bias exists at infinitesimal LR.
Our analysis is inspired by a previous line of works~\citep{blanc2020ou,damian2021label,li2022what} showing that label noise can
drive SGD to move on the minimizer manifold along the direction of minimizing the trace of Hessian,
but the key difference here is that we do not have any stochastic gradient noise.
Instead, we study the EoS regime and exhibit that GD oscillates around the minimizer manifold by connecting it to power method.
We show that this oscillation is a driving power that pushes the parameter to move on the manifold,
and analyze the speed of this movement by modeling
two key parameters of the dynamics
as a 1-dimensional Hamiltonian system.
To the best of our knowledge,
our work is the first one that
theoretically proves a sharpness measure to decrease in standard training processes with gradient descent,
without any additional regularization such as label noise~\citep{blanc2020ou,damian2021label,li2022what},
or using non-standard variants of gradient descent update rule,
e.g., using normalized GD or non-smooth wrappings on the loss function~\citep{arora2022understanding}.

\input{related}

\input{prelim}

\input{scale-invariant}

\input{applications}

\section{Conclusions and Future Work} \label{sec:con}
We exhibited settings  where
gradient descent has an implicit bias
to reduce spherical sharpness
in training neural nets with normalization layers and weight decay,
and we verified experimentally this sharpness-reduction bias predicted by our theorem
as well as its generalization benefit  on CIFAR-10.

Our theoretical analysis applies to dynamics around a minimizer manifold
and requires a small (but finite) learning rate
so that we can show that the parameter oscillates locally
and
approximately tracks a sharpness-reduction flow.
We note that in practice a decrease in
 spherical sharpness is observed even with
 moderate LR and even before getting close to a minimizer manifold.
Explaining these phenomena is left for future work.
Now we list some other future directions. First is to
generalize our results to SGD,
where the sharpness measure may not be the spherical sharpness
and could depend on the structure of gradient noise.
Second, to understand benefit of reducing spherical sharpness on specific tasks,
e.g., why does reducing spherical sharpness encourage low-rank
on matrix completion with BN (\Cref{fig:matcom-intro})?
Third, to study sharpness-reduction bias for neural net architectures that are not scale-invariant on all parameters (e.g., with certain unnormalized layers).

\section{PAC-Bayes Bounds Based on Spherical Sharpness} \label{sec:pac-bayes}

In this section we give the PAC-Bayes bound for generalization error using spherical sharpness~(\Cref{def:sph-sha}). We will start with our setting and then recap the classic PAC-Bayes theorem in \citep{mcallester2003simplified}. The main result in this section is \Cref{thm:spherical_Sharpness_pac_bayes}.

\paragraph{Setting.} Let $\ell({\bm{w}},{\bm{z}})$ be the loss of parameter ${\bm{w}}$ on data point ${\bm{z}}$ and assume  $\ell_{\max}=
\sup_{z\in\DatZ,{\bm{w}}\in \mathbb{S}^{D-1}} \ell({\bm{w}},{\bm{z}})<\infty$. Let $ \DatS := \{ {\bm{z}}_i \}_{i=1}^{n}$ where ${\bm{z}}_i$ are
sampled independently. Different to the previous notation, we use $\mathcal{L}_\DatS$ to denote the empirical loss on training
dataset $\DatS$ where we run optimization algorithms and $\mathcal{L}$ to denote the population loss. Let
$\thirdL{\mathcal{L}_\DatS}:= \sup_{{\bm{\theta}}\in\mathbb{S}^{D-1}} \normtwosm{\nabla^3  \mathcal{L}_\DatS({\bm{\theta}})}$. Since
$\mathbb{S}^{D-1}$ is compact and  $\normtwosm{\nabla^3 \mathcal{L}_\DatS(\cdot)}$ is continuous,  $\thirdL{\mathcal{L}_\DatS}$ is finite.

\begin{theorem}[PAC-Bayes theorem~\citep{mcallester2003simplified}]\label{thm:PAC_bayes_theorem}
Given any distribution $P$  on $\mathbb{R}^D$, with at least $1-\delta$ probability over the randomness of the dataset $\DatS$, for all distribution $Q$ on  $\mathbb{R}^D$, it holds that
\begin{align*}\label{eq:pac_bayes}
\mathbb{E}_{{\bm{w}}\sim Q} \mathcal{L}_{S}({\bm{w}}) - \mathbb{E}_{{\bm{w}}\sim Q}\mathcal{L}({\bm{w}})	 \le     \ell_{\max} \sqrt{\frac{D_{\mathrm{KL}}(Q \| P) + \ln(n/\delta)}{2(n-1)}}.
\end{align*}

\end{theorem}

Now we are ready to state the main theorem in this section, \Cref{thm:spherical_Sharpness_pac_bayes}, which shows that small spherical sharpness $\lambda_1(\nabla^2 \mathcal{L}_S({\bm{\theta}}))$ leads to small generalization error.
\begin{theorem}\label{thm:spherical_Sharpness_pac_bayes} For any $\sigma \le \frac{1}{2 + 2\sqrt{(\ln n)/D}}$, with at
	least $1-\delta$ probability over the randomness of the dataset $\DatS$, where  $\DatS := \{ {\bm{z}}_i \}_{i=1}^{n}$ and
	every ${\bm{z}}_i$ is sampled independently, for any ${\bm{\theta}} \in \mathbb{S}^{D-1}$,
\begin{align*}
&\mathbb{E}_{{\bm{\epsilon}}\sim \mathcal{N}(\bm{0}, \sigma^2{\bm{I}}_D / D)} \mathcal{L}({\bm{\theta}}+{\bm{\epsilon}}) - \mathcal{L}_\DatS({\bm{\theta}})\\
& \quad \le\frac{\sigma^2}{2} \lambda_1 (\nabla^2\mathcal{L}_\DatS({\bm{\theta}}))
+ \frac{16\sigma^3}{3}\thirdL{\mathcal{L}_\DatS} (1 + ((\ln n)/D)^{1.5})
+ \ell_{\max} \sqrt{\frac{D/\sigma^2 + 2\ln(n/\delta)}{n-1}},
\end{align*}
Thus with the standard assumption in \citep{foret2021sam} that $\mathbb{E}_{{\bm{\epsilon}}\sim \mathcal{N}(\bm{0}, \sigma^2{\bm{I}}_D/D)} \mathcal{L}({\bm{\theta}}+{\bm{\epsilon}}) \ge \mathcal{L}({\bm{\theta}})$, we have the same upper bound for $\mathcal{L}({\bm{\theta}}) - \mathcal{L}_{\DatS}({\bm{\theta}})$.
\end{theorem}

To prove \Cref{thm:spherical_Sharpness_pac_bayes}, we will first need the following lemma. The proof is standard so it is omitted.
\begin{lemma}
Let $Q = \mathcal{N}({\bm{\mu}}_Q,\sigma^2_Q {\bm{I}}_D)$ and $P = \mathcal{N}({\bm{\mu}}_P,\sigma^2_P{\bm{I}}_D)$, we have that
	\begin{align*}
		D_{KL}(Q\|P)  = \frac{1}{2}\left[ \frac{D \sigma_Q^2 + \norm{{\bm{\mu}}_P-{\bm{\mu}}_Q}_2^2}{\sigma_P^2}-D+D\log\left( \frac{\sigma^2_P}{\sigma^2_Q}\right)\right]
	\end{align*}
\end{lemma}

\begin{proof}[Proof of \Cref{thm:spherical_Sharpness_pac_bayes}]
Let $Q := \mathcal{N}({\bm{\theta}}, \sigma^2 {\bm{I}}_D / D)$, $P := \mathcal{N}({\bm{0}}, \sigma^2 {\bm{I}}_D / D)$.
Then $D_{KL}(Q\|P) = \frac{D}{2\sigma^2}$.
By \Cref{thm:PAC_bayes_theorem} we have
\begin{align*}
\mathbb{E}_{{\bm{\epsilon}}\sim \mathcal{N}(\bm{0}, \sigma^2{\bm{I}}_D / D)} \mathcal{L}({\bm{\theta}}+{\bm{\epsilon}}) - \mathbb{E}_{{\bm{\epsilon}}\sim \mathcal{N}(\bm{0}, \sigma^2{\bm{I}}_D / D)} \mathcal{L}_{S}({\bm{\theta}}+{\bm{\epsilon}}) \le \ell_{\max}\sqrt{\frac{\frac{D}{2\sigma^2 } + \ln(n/\delta)}{2(n-1)}}.
\end{align*}
Let $h(\sigma):= (1+\sqrt{(\ln n) / D})\sigma$. By assumption it holds that $h(\sigma)\le \frac{1}{2}$. Thus by Lemma 1 in \cite{laurent2000adaptive}, we have for any positive $t$:
\begin{align}
	\mathbb{P}[\norm{{\bm{\epsilon}}}_2^2-\sigma^2\ge 2\sigma^2 \sqrt{t/D} +2t\sigma^2 / D] \le \exp(-t).
\end{align}
Therefore, with probability $1-1/\sqrt{n}$, we have
\begin{align*}
	\norm{{\bm{\epsilon}}}_2^2 \le \sigma^2\left(1 + 2 \sqrt{(\ln \sqrt{n}) / D} + 2 \cdot (\ln \sqrt{n}) / D\right)\le \sigma^2\left(1+\sqrt{(\ln n)/D}\right)^2\le h(\sigma)^2.
\end{align*}
Thus for any ${\bm{\theta}}$ with $\norm{{\bm{\theta}}}_2=1$, we have
\begin{align*}
	\mathbb{E}_{{\bm{\epsilon}}\sim \mathcal{N}(\bm{0}, \sigma^2{\bm{I}}_D / D)} \mathcal{L}_{S}({\bm{\theta}}+{\bm{\epsilon}}) \le \ell_{\max}/\sqrt{n} + \mathbb{E}_{{\bm{\epsilon}}\sim \mathcal{N}(\bm{0}, \sigma^2{\bm{I}}_D / D)} \left[ \mathcal{L}_{S}({\bm{\theta}}+{\bm{\epsilon}}) \onec{\norm{{\bm{\epsilon}}}_2\le h(\sigma)} \right].
\end{align*}

By Taylor expansion, $$\mathcal{L}_\DatS({\bm{\theta}}+{\bm{\epsilon}})\le \mathcal{L}_\DatS({\bm{\theta}}) + \dotp{{\bm{\epsilon}}}{\nabla \mathcal{L}_\DatS({\bm{\theta}})}+\frac{1}{2}\dotp{{\bm{\epsilon}}}{\nabla^2 \mathcal{L}_\DatS({\bm{\theta}}){\bm{\epsilon}}} + \frac{1}{6}\sup_{\lambda\in [0,1]} \dotp{\nabla^3 \mathcal{L}_\DatS({\bm{\theta}}+\lambda{\bm{\epsilon}})}{{\bm{\epsilon}}^{\otimes 3}}.$$

Note that $\norm{{\bm{\theta}}+\lambda {\bm{\epsilon}}}_2 \ge \norm{{\bm{\theta}}}_2 - \lambda\norm{{\bm{\epsilon}}}_2 \ge \frac{1}{2}$ for all $\lambda\in [0,\frac{1}{2}]$,
it holds that $\norm{\nabla^3 \mathcal{L}_\DatS({\bm{\theta}}+\lambda{\bm{\epsilon}})}= \norm{{\bm{\theta}}+\lambda {\bm{\epsilon}}}_2^{-3}\norm{\nabla^3 \mathcal{L}_\DatS( \frac{{\bm{\theta}}+\lambda{\bm{\epsilon}}}{\norm{{\bm{\theta}}+\lambda{\bm{\epsilon}}}_2})}_2 \le 8 \thirdL{\mathcal{L}_\DatS} $, we have
\begin{align*}
	\mathbb{E}_{{\bm{\epsilon}}\sim \mathcal{N}(\bm{0}, \sigma^2{\bm{I}}_D / D)} \left[\mathcal{L}_{S}({\bm{\theta}}+{\bm{\epsilon}}) \onec{\norm{{\bm{\epsilon}}}\le h(\sigma)}\right]
	&\le \mathcal{L}_\DatS({\bm{\theta}})
	+ \frac{\sigma^2}{2D} \Tr[\nabla^2 \mathcal{L}_\DatS({\bm{\theta}})] + \frac{4}{3}\thirdL{\mathcal{L}_\DatS} h(\sigma)^3 \\
	&\le \mathcal{L}_\DatS({\bm{\theta}})
	+ \frac{\sigma^2}{2} \lambda_1(\nabla^2 \mathcal{L}_\DatS({\bm{\theta}})) + \frac{4}{3}\thirdL{\mathcal{L}_\DatS} h(\sigma)^3.
\end{align*}
Thus we conclude that
	\begin{align*}
&~\mathbb{E}_{{\bm{\epsilon}}\sim \mathcal{N}(\bm{0}, \sigma^2{\bm{I}}_D / D)} \mathcal{L}({\bm{\theta}}+{\bm{\epsilon}}) - \mathcal{L}_\DatS({\bm{\theta}})\\
\le &~\frac{\sigma^2}{2}\lambda_1(\nabla^2 \mathcal{L}_\DatS({\bm{\theta}})) + \frac{4}{3}\thirdL{\mathcal{L}_\DatS} h(\sigma)^3 +
\ell_{\max} \left(\sqrt{\frac{\frac{D}{2\sigma^2 } + \ln(n/\delta)}{2(n-1)}} +1 \right)\\
\le &~\frac{\sigma^2}{2} \lambda_1 (\nabla^2\mathcal{L}_\DatS({\bm{\theta}}))  + \frac{16\sigma^3}{3}\thirdL{\mathcal{L}_\DatS} (1 + (\ln n / D)^{1.5}) +
\ell_{\max} \sqrt{\frac{D/\sigma^2 + 2\ln(n/\delta)}{n-1}},
	\end{align*}
which completes the proof.
\end{proof}
\section{Preliminaries}
\label{sec:prelims}

Let $\mathbb{S}^{D-1} := \{ {\bm{\theta}} \in \mathbb{R}^D : \normtwosm{{\bm{\theta}}} = 1\}$ be the unit
sphere equipped with subspace topology.
We say a loss function $\mathcal{L}({\bm{w}})$ defined on $\mathbb{R}^{D} \setminus \{{\bm{0}}\}$ is
\textit{scale-invariant} if $\mathcal{L}(c{\bm{w}}) = \mathcal{L}({\bm{w}})$ for all $c > 0$.
In other words, the loss value does not change with the parameter norm.
For a differentiable scale-invariant function $\mathcal{L}({\bm{w}})$,
the gradient is $(-1)$-homogeneous and it is always perpendicular to ${\bm{w}}$, i.e.,
$\nabla \mathcal{L}(c{\bm{w}}) = c^{-1} \nabla \mathcal{L}({\bm{w}})$ for all $c > 0$ and $\dotp{\nabla \mathcal{L}({\bm{w}})}{{\bm{w}}} = 0$.
The focus of this paper is the dynamics of GD+WD\xspace on scale-invariant loss.
\eqref{eq:upd-gdwd} gives the update rule
for learning rate (LR) $\hat{\eta}$ and weight decay (WD) $\hat{\lambda}$.
We use ${\bm{\theta}}_t := \frac{{\bm{w}}_t}{\normtwosm{{\bm{w}}_t}}$ to denote the projection of ${\bm{w}}_t$ onto the unit sphere at step $t$.
A number of works \citep{van2017l2,hoffer2018norm,arora2018theoretical} propose to use the ``effective learning rate'' $\frac{\hat{\eta}}{\normtwosm{{\bm{w}}_t}^2}$
to measure the update efficiency of ${\bm{\theta}}_t$ on $\mathbb{S}^{D-1}$.
Inspired by this notion, we write GD+WD\xspace on scale-invariant loss
as a specific kind of Projected Gradient Descent (PGD),
and define $\tilde{\eta}_t := \frac{\hat{\eta}}{(1-\hat{\eta}\hat{\lambda})\normtwosm{{\bm{w}}_t}^2}$
to be the effective learning rate with slight abuse of terminology.
\begin{lemma} \label{lm:scale-invariant-on-sphere}
    When the parameters ${\bm{w}}_t$ are updated as \eqref{eq:upd-gdwd}, ${\bm{\theta}}_{t}$ satisfies the following equation:
    \begin{equation} \label{eq:theta-update-simple}
        {\bm{\theta}}_{t+1} = \Pi({\bm{\theta}}_t - \tilde{\eta}_t \nabla \mathcal{L}({\bm{\theta}}_t)),
    \end{equation}
    where $\tilde{\eta}_t := \frac{\hat{\eta}}{(1-\hat{\eta}\hat{\lambda})\normtwosm{{\bm{w}}_t}^2}$ is called the \textbf{effective learning rate} at step $t$,
    and $\Pi: {\bm{w}} \mapsto \frac{{\bm{w}}}{\normtwosm{{\bm{w}}}}$ is
    the projection operator that projects any vector onto the unit sphere.
\end{lemma}
The convergence rate of GD+WD\xspace has been analyzed by \citet{li2022robust}.
Here we present a variant of their theorem
that bounds both the gradient and effective LR.
\begin{theorem}[Variant of Theorem D.2 in \citep{li2022robust}] \label{thm:li-gdwd}
    Let $\mathcal{L}({\bm{w}})$ be a scale-invariant loss function and
    $\rho_2 := \sup\{ \normtwosm{\nabla^2 \mathcal{L}({\bm{w}})} : {\bm{w}} \in \mathbb{S}^{D-1} \}$ be the smoothness constant of $\mathcal{L}$ restricted on the unit sphere.
    For GD+WD\xspace \eqref{eq:upd-gdwd} with $\hat{\eta} \hat{\lambda} \le 1/2$ and $\tilde{\eta}_0 \le \frac{1}{\pi^2 \rho_2(1-\hat{\eta}\hat{\lambda})}$, let
    $T_0 := \left\lceil \frac{1}{2 \hat{\eta} \hat{\lambda}}  \ln \frac{\normtwosm{{\bm{w}}_0}^2}{\rho_2 \pi^2 \hat{\eta}}  \right\rceil$
    steps,
    there must exist $0 \le t \le T_0$ such that
    $\normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2 \le 8\pi^4\rho_2^2\hat{\lambda}\hat{\eta}$ and $\tilde{\eta}_t \le \frac{2}{\pi^2 \rho_2 (1-\hat{\eta}\hat{\lambda})}$.
\end{theorem}

\section{Proof Outlines of Our Theorems on Sharpness Reduction} \label{sec:proof-outline-main}

In this section,
we give proof outlines of \Cref{thm:main-ful,thm:main-sph}.
The main proof idea
for the spherical case
is stated in \Cref{sec:proof-idea},
but technically it is easier
to state and prove the lemmas for full space optimization.
Therefore, we present the full details for the full space case
and omit certain details for the spherical case if they are similar
to the full space case.

As mentioned in \Cref{sec:proof-idea},
a key ingredient in our proof is to show that
the period-2 oscillation drives the parameter to move along the manifold.
For both full space and spherical cases,
we project ${\bm{\theta}}_t$
onto the manifold $\mathit{\Gamma}$
with a carefully-defined projection function, ${\bm{\phi}}_t := \Phi({\bm{\theta}}_t)$.
Then we show that
after a period of oscillation ($2$ steps),
the projection drifts from ${\bm{\phi}}_t$ to a new position ${\bm{\phi}}_{t+2}$
approximately along the direction of $\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$.
To analyze the speed of each drift,
we show that
the oscillation can be tracked with
two variables $h \in \mathbb{R}$, $u \in \mathbb{R}$,
where
$h$ is related to the displacement of ${\bm{\theta}}_t$ from the manifold,
$u$ is related to the closeness of the current dynamic to the edge of stability.
We formally define
a discrete process called \textit{RMS-drift process}
and show that
the oscillations in both full space and spherical cases
can be regarded as RMS-drift processes.

The rest of the section is organized as follows.
In \Cref{sec:outline-notation} we introduce some additional notations.
In \Cref{sec:outline-rmsdrift-def} we formally introduce the concept
of RMS-drift process.
In \Cref{sec:outline-ful} we show how to
reduce our analysis in the full space case
into studying an RMS-drift process.
In \Cref{sec:outline-sph} we show how to
reduce our analysis in the spherical case
into studying an RMS-drift process.
Finally, in \Cref{sec:outline-rmsdrift-analysis}
we analyze the RMS-drift process
and show that the projections of parameters in training can be tracked with
the sharpness-reduction flow defined in \eqref{eq:zeta-general}.

\subsection{Additional Notations} \label{sec:outline-notation}

We follow the notations in \Cref{sec:add-prelim}.
We also need some additional notations in this section.
It is implied by
the uniqueness of the top eigenvalue
(\Cref{ass:topeigen-ful} or \Cref{ass:topeigen})
and $\mathcal{L} \in \mathcal{C}^4$
that $\lambda^{\mathrm{H}}_1({\bm{\theta}})$ is $\mathcal{C}^2$-smooth on $\mathit{\Gamma}$,
and we can construct a $\mathcal{C}^2$-smooth function $\vv^{\mathrm{H}}_1({\bm{\theta}})$
such that
$\vv^{\mathrm{H}}_1({\bm{\theta}})$ is a unit top eigenvector of ${\bm{H}}({\bm{\theta}})$ on $\mathit{\Gamma}$.
Let $\mPHz({\bm{\theta}})$ be the projection matrix onto the null space of ${\bm{H}}_t$,
and $\mPHnzt({\bm{\theta}})$ be the projection matrix onto the space spanned by the eigenvectors of ${\bm{H}}_t$
corresponding to non-zero and non-top eigenvalues ($\lambda \ne 0, \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$).
For any local minimizer
${\bm{\theta}} \in \mathit{\Gamma}$,
we define $\gamma({\bm{\theta}}) := \frac{1}{\lambda^{\mathrm{H}}_1({\bm{\theta}})} \min\left\{\lambda^{\mathrm{H}}_1({\bm{\theta}}) -
\lambda^{\mathrm{H}}_2({\bm{\theta}}), \lambda^{\mathrm{H}}_{D-D_\Gamma}({\bm{\theta}})\right\} $ to be the
relative eigenvalue gap (between $\lambda^{\mathrm{H}}_1$ and $\lambda^{\mathrm{H}}_2$, or between $\lambda^{\mathrm{H}}_{D- D_\Gamma}$ and $0$).
Then $\mPHz({\bm{\theta}}), \mPHnzt({\bm{\theta}}), \gamma({\bm{\theta}})$ are all $\mathcal{C}^2$-smooth on $\mathit{\Gamma}$.
Finally, we define $\mu({\bm{\theta}}) := \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\theta}})}$.

We define $\Phi({\bm{\theta}})$ to be the convergence point of a gradient flow of $\mathcal{L}$ starting from ${\bm{\theta}}$.
That is, the limit of $\tilde{{\bm{\theta}}}(t)$ as $t \to +\infty$ when
$\tilde{{\bm{\theta}}}(t)$ is described by the following ODE:
\[
  \frac{\textup{\textrm{d}} \tilde{{\bm{\theta}}}}{\textup{\textrm{d}} t} = -\nabla \mathcal{L}(\tilde{{\bm{\theta}}}), \quad \text{where} \quad \tilde{{\bm{\theta}}}(0) = {\bm{\theta}}.
\]
We leave $\Phi({\bm{\theta}})$ undefined if the ODE does not converge to any point.

For GD/PGD with quasi-RMSprop scheduler (\Cref{def:qrms-sche}),
the state at step $t$ can be written as $({\bm{\theta}}_t, \tilde{v}_t)$,
where ${\bm{\theta}}_t$ is the trainable parameter in GD/PGD and $\tilde{v}_t$ is the moment estimate.
Whenever $\Phi({\bm{\theta}}_t)$ exists and is in $\mathit{\Gamma}$,
we define
\begin{align*}
  {\bm{\phi}}_t &:= \Phi({\bm{\theta}}_t) & {\bm{H}}_t &:= {\bm{H}}({\bm{\phi}}_t) \\
  \mu_t &:= \mu({\bm{\phi}}_t) & {\bm{U}}_t &:= {\bm{I}} - \mu_t {\bm{H}}_t \\
  {\bm{x}}_t &:= {\bm{\theta}}_t - {\bm{\phi}}_t
\end{align*}
Most importantly, we define two hidden variables $(h_t, u_t)$ as follows:
\begin{align}
  h_t &:= \tfrac{1}{\eta}\dotpsm{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)}{{\bm{x}}_t}, &
  u_t &:= \tfrac{1}{\eta}( \mu_t^2 \tilde{v}_t- 1).
  \label{eq:hu-def}
\end{align}
Note that the $h_t$ defined here {\em differs} with that in \Cref{sec:proof-idea} by a factor of $1/\eta$.
We introduce this factor for the sake of convenience,
and we will use the definition with the factor $1/\eta$
only in our theoretical analysis.

\subsection{RMS-drift Process: Introduction} \label{sec:outline-rmsdrift-def}

\begin{definition}
  A \textit{drift state} is described by a tuple $S = (h, u, {\bm{\phi}})$ in the drift state space $\mathcal{S} := \mathbb{R} \times \mathbb{R} \times \mathit{\Gamma}$.
  We say that $S$ is \textit{$\alpha$-bounded} if $\max\{\abssm{h}, \abssm{u}\} \le \alpha$.
\end{definition}
\begin{definition}
  Given two drift states $S_t = (h_t, u_t, {\bm{\phi}}_t)$, $S_{t+2} = (h_{t+2}, u_{t+2},
  {\bm{\phi}}_{t+2})$ in the drift state space $\mathcal{S}$,
  for learning rate $\eta > 0$ and hyperparameter $C_{\mathrm{b}} > 0$,
  we say that the transition $S_t \to S_{t+2}$
  is a \textit{$C_0$-RMS-drift transition}
  if for all $\alpha \ge 1$,
  as long as $S_t$ is $\alpha$-bounded,
  $S_{t+2}$ is close to an auxiliary state $S'_{t+2} := (h'_{t+2}, u'_{t+2}, {\bm{\phi}}'_{t+2})$
  in the following sense:
  \begin{align*}
    h'_{t+2} &:= (1 - 2\eta u_t) h_t,  &
    \abssm{h_{t+2} - h'_{t+2}} &\le C_0 \alpha^2 \abssm{h_t} \eta^2, \\
    u'_{t+2} &:= u_t + 4 \eta h_t^2 (2 C_{\mathrm{b}} + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2) - 2\eta C_{\mathrm{b}}, &
    \abssm{u_{t+2} - u'_{t+2}} &\le C_0 \alpha (1 + h_t^2) \eta^2,  \\
    {\bm{\phi}}'_{t+2} &:= {\bm{\phi}}_t - 2 \eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t), &
    \normtwosm{{\bm{\phi}}_{t+2} - {\bm{\phi}}'_{t+2}} &\le C_0 \alpha h_t^2 \eta^3.
  \end{align*}
  For a sequence of states $S_0, S_2, S_4, \dots, S_{2M}$, we say it is a $C_0$-\textit{RMS-drift process} if $S_t \to S_{t+2}$ is a
  $C_0$-RMS-drift transition for all even
  numbers $0 \le t < 2M$.
\end{definition}

In our analysis of GD/PGD with quasi-RMSprop scheduler,
we can rewrite the dynamics as RMS-drift processes (after a few warm-up steps),
where
${\bm{\phi}}_t := \Phi({\bm{\theta}}_t)$ is the gradient flow projection
of the parameter at step $t$ onto $\mathit{\Gamma}$,
and $h_t,u_t$ are two hidden variables defined in \eqref{eq:hu-def}.
This RMS-drift process
serves as an abstraction of the original dynamics
that contains the minimal but sufficient amount of information
so that we can compute the continuous approximation for the trajectory of ${\bm{\phi}}_t$.

In RMS-drift process, ${\bm{\phi}}_t$ evolves as gradient descent on $\mathit{\Gamma}$ for minimizing $\log \lambda^{\mathrm{H}}_1({\bm{\theta}})$,
and the corresponding learning rate is changing with $h_t$ per step.
To obtain the final flow approximation \eqref{eq:zeta-general},
we need to sum up the learning rates over time.

An intuitive way to understand RMS-drift process is to use the following first-order continuous approximation with time scaling $(h(\tau),u(\tau),{\bm{\phi}}(\tau)) \approx (h_{\tau/\eta}, u_{\tau/\eta}, {\bm{\phi}}_{\tau/\eta})$ and ignore all the second order terms $O(\eta^2)$:
\begin{align*}
  \frac{\textup{\textrm{d}} h}{\textup{\textrm{d}} \tau} &= - u h, & \frac{\textup{\textrm{d}} u}{\textup{\textrm{d}} \tau} &= 2h^2(2C_{\mathrm{b}} + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}})}^2) - C_{\mathrm{b}},
  & \frac{\textup{\textrm{d}} {\bm{\phi}}}{\textup{\textrm{d}} \tau} &= {\bm{0}}.
\end{align*}
This approximation gives an important insight:
$h_t$ and $u_t$ are changing much faster than ${\bm{\phi}}_t$
when $\eta$ is small.
Therefore, we can analyze this first-order approximation
to obtain an average value of $2\eta^2h_t^2$,
and use this average value as the ``effective'' learning rate
in the flow approximation.

In fact, $(h, u)$ forms a $1$-dimensional Hamiltonian system after a coordinate transformation
and evolves periodically in the above ODE.
To see this,
we can compute the time derivatives of $(\log \abssm{h}, -u)$ while letting $K := \sqrt{2C_{\mathrm{b}} + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}})}^2}$:
\begin{align} \label{eq:qrms-cont}
  \frac{\textup{\textrm{d}} \log \abssm{h}}{\textup{\textrm{d}} \tau} &= -u, & \frac{\textup{\textrm{d}} (-u)}{\textup{\textrm{d}} \tau} &= -\left(2K^2 h^2 - C_{\mathrm{b}}\right).
\end{align}
Consider a unit-mass particle in the system with position $x(t)$, velocity $v(t) = x'(t)$ and acceleration $a(t) = v'(t)$.
Suppose that this system has the potential energy
$U(x) := K^2e^{2x} - C_{\mathrm{b}} x$.
Then we can see that the position and velocity of this particle evolve exactly the same as $(\log \abssm{h}, -u)$!
This also shows that $(\log \abssm{h}, -u)$
evolves periodically because $1$-dimensional Hamiltonian system with unimodal potential energy must be periodic.
See also \Cref{fig:qrms-vis}.

\begin{figure}[t]
  \centering
  \includegraphics[width=\textwidth]{fig/qrms-vis.pdf}
  \vspace{-0.2in}
  \caption{
    A visualization of the dynamical system described by \eqref{eq:qrms-cont} with $K=2, C_{\mathrm{b}}=1$.
    This system can be associated with a physical system in which a unit-mass particle
    moves in a potential well $U(x) = K^2 e^{2x} - C_{\mathrm{b}} x$ without any energy loss; see (a).
    $\log \abssm{h}$ can be seen as the position of the particle,
    and $-u$ can be seen as the velocity.
    This system must be periodic because of the conservation of energy,
    and it can be seen clearly from its phase portrait (b).
  }
  \label{fig:qrms-vis}
\end{figure}

In \Cref{sec:outline-rmsdrift-analysis} below, we will use this observation
to obtain the flow approximation nicely.
But now
we first outline how to reduce the original dynamics to an RMS-drift process.

\subsection{Reduction to RMS-drift Process: The Case of Full Space Optimization} \label{sec:outline-ful}

Now we outline how to reduce the dynamics to an RMS-drift process
in the setting of \Cref{thm:main-ful},
the main theorem for the case of full-space optimization.

\subsubsection{Construction of Working Zones}

Let ${\mathcal{Z}} := \{ {\bm{\zeta}}(t) : t \in [0, T] \} \subseteq \mathit{\Gamma}$
be the set of points passed by the sharpness-reduction flow~\eqref{eq:zeta-general}.
Inspired by \citet{arora2022understanding}, we construct a two-level nested working zone
$(\gZ^{\epsilon_0}, \gZ^{\epsilon_1})$,
where $\gZ^{\epsilon_0}, \gZ^{\epsilon_1}$ are essentially the $\epsilon_0$- and $\epsilon_1$-neighborhoods of ${\mathcal{Z}}$
with $\epsilon_0, \epsilon_1$ carefully chosen.
In our later analysis, we will ensure that ${\bm{\theta}}_t$ is always in $\gZ^{\epsilon_0}$,
and its gradient flow projection ${\bm{\phi}}_t$ is always well-defined and lies in $\gZ^{\epsilon_1} \cap \manGa$.

The following lemma shows some important properties of the working zone in our construction.
We defer the proof to \Cref{sec:working-zone-proof-construction}.
\begin{lemma}[Working Zone Lemma] \label{lm:ful-working-zone}
  There exist $0 < \epsilon_0 < \epsilon_1$ such that $\gZ^{\epsilon_0}, \gZ^{\epsilon_1}$ satisfy the following:
  \begin{enumerate}
    \item $\mathrm{cl}(\gZ^{\epsilon_1}) \cap \mathit{\Gamma}$ is compact;
    \item $\mathcal{L}$ satisfies $\mu_{\mathtt{PL}}$-PL on $\gZ^{\epsilon_1}$ for some $\mu_{\mathtt{PL}} > 0$;
    \item $\Phi$ is well-defined on $\gZ^{\epsilon_0}$, and $\Phi({\bm{\theta}}) \in \gZ^{\epsilon_1} \cap \manGa$ for all ${\bm{\theta}} \in \gZ^{\epsilon_0}$;
    \item $\Phi$ is $\mathcal{C}^3$-smooth on $\gZ^{\epsilon_0}$;
    \item $\gamma({\bm{\theta}}) \ge \gamma_{\min}$ holds uniformly on $\gZ^{\epsilon_1} \cap \manGa$ for some $\gamma_{\min} > 0$.
  \end{enumerate}
\end{lemma}

It can be seen from \Cref{lm:ful-working-zone} that
${\bm{\phi}}_t, {\bm{H}}_t, \mu_t, {\bm{U}}_t, {\bm{x}}_t, h_t, u_t$
are all well-defined
as long as ${\bm{\theta}}_t \in \gZ^{\epsilon_0}$.
Now we define some useful notions in the EoS regime.
A state is $\alpha$-bounded
if ${\bm{\theta}}_t$ is close to $\mathit{\Gamma}$
and the dynamic is in the EoS regime due to $\tilde{v}_t \approx \mu_t^{-2}$.
A state is $\alpha$-deviated if ${\bm{\theta}}_t$
is not too close to $\mathit{\Gamma}$.
\begin{definition}[$\alpha$-Bounded State]
  We say that the state $({\bm{\theta}}_t, \tilde{v}_t)$ at some step $t$
  is $\alpha$-bounded if
  ${\bm{\theta}}_t \in \gZ^{\epsilon_0}$,
  $\normtwosm{{\bm{x}}_t} \le \alpha \eta$ and $\abssm{u_t} \le \alpha$.
\end{definition}
\begin{definition}[$\alpha$-Deviated State]
  We say that the state $({\bm{\theta}}_t, \tilde{v}_t)$ at some step $t$
  is $\alpha$-deviated (from the manifold $\mathit{\Gamma}$) if
  $\normtwosm{{\bm{x}}_t} \ge \eta\exp(-\alpha^2)$ or
  ${\bm{\theta}}_t \notin \gZ^{\epsilon_0}$.
\end{definition}

Next, we define
a quantitative measurement
for how much ${\bm{x}}_t$ aligns to the top eigenvector of ${\bm{H}}_t$.
Note that ${\bm{x}}_t$ can be decomposed into the projections
onto the top eigenspace, the null space, and the space spanned by non-zero and non-top eigenvectors;
we can write this decomposition as
${\bm{x}}_t = h_t \eta \vv^{\mathrm{H}}_1({\bm{\phi}}_t) + \mPHz({\bm{\phi}}_t) {\bm{x}}_t + \mPHnzt({\bm{\phi}}_t) {\bm{x}}_t$.
The following lemma shows that $\mPHnzt({\bm{\phi}}_t)$ is always negligible,
so we characterize the alignment to the top eigenvector only through comparing
$\eta h_t \vv^{\mathrm{H}}_1({\bm{\phi}}_t)$ and $\mPHnzt({\bm{\phi}}_t) {\bm{x}}_t$.
\begin{lemma} \label{lm:mPHz-x-small}
  At any step $t$, if ${\bm{\theta}}_t \in \gZ^{\epsilon_0}$, then $\normtwosm{\mPHz({\bm{\phi}}_t) {\bm{x}}_t} = O(\normtwosm{{\bm{x}}_t}^2)$.
\end{lemma}
\begin{proof}
  Direct consequence of \Cref{lm:phiphix-ful}.
\end{proof}

\begin{definition}[$p$-Misaligned State]
  We say that the state $({\bm{\theta}}_t, \tilde{v}_t)$ at some step $t$
  is at most $p$-misaligned (to the top eigenvector) if
  ${\bm{\theta}}_t \in \gZ^{\epsilon_0}$,
  $\normtwosm{\mPHnzt({\bm{\phi}}_t) {\bm{x}}_t} \le p \cdot h_t \eta$.
\end{definition}

\subsubsection{Good Initialization}

First, we show that the initialization satisfies some desirable properties
with high probability.
The proof is deferred to \Cref{sec:ful-good-init}.
\begin{lemma} \label{lm:ful-good-init}
  There exists $\delta = O(\alpha_0 \eta \sqrt{\log(1/\eta)})$ such that
  the following holds.
  With probability $1 - \delta$,
  the initial state is $\Osm{\alpha_0\sqrt{\log(1/\delta)}}$-bounded,
  $O(\alpha_0 + \sqrt{\log(1/\delta)})$-deviated,
  at most $O(1/\delta)$-misaligned,
  and satisfies $\normtwosm{{\bm{\phi}}_0 - {\bm{\zeta}}_0} \le O(\alpha_0\eta \sqrt{\log(1 / \delta)})$.
\end{lemma}

\subsubsection{Alignment Phase}
At initialization, the state is not well-aligned to the top eigenvector.
But in the following, we show that it becomes at most $O(\abssm{h_t}\eta)$-misaligned
after only $\eta^{-o(1)}$ steps.
We defer the proofs to \Cref{sec:ful-proof-alignment}.

The key lemma is the following, which gives good approximations for various important variables.
\begin{lemma} \label{lm:ful-align-1-step}
  For small enough base learning rate $\eta$,
  at any step $t$,
  if the state $({\bm{\theta}}_t, \tilde{v}_t)$ is $\alpha$-bounded
  for some $1 \le \alpha \le \eta^{-o(1)}$,
  then ${\bm{\theta}}_{t+1} \in \gZ^{\epsilon_0}$, and
  \begin{align}
    {\bm{\phi}}_{t+1} &= {\bm{\phi}}_t + O(\normtwosm{{\bm{x}}_t}^2)  \label{eq:ful-align-1-step-phi} \\
    {\bm{x}}_{t+1} &= ({\bm{I}} - \tilde{\eta}_t {\bm{H}}_t) {\bm{x}}_t + O(\normtwosm{{\bm{x}}_t}^2) \label{eq:ful-align-1-step-x} \\
    h_{t+1} &= -h_t + O(\alpha \normtwosm{{\bm{x}}_t}) \label{eq:ful-align-1-step-h} \\
    \normtwosm{\mPHnzt({\bm{\phi}}_{t+1}) {\bm{x}}_{t+1}} &\le (1 - 1.9\gamma_{\min}) \normtwosm{\mPHnzt({\bm{\phi}}_t) {\bm{x}}_t} + O(\normtwosm{{\bm{x}}_t}^2) \label{eq:ful-align-1-step-pnzt} \\
    u_{t+1} &= u_t + O(\alpha^2 \eta) \label{eq:ful-align-1-step-u}
  \end{align}
\end{lemma}

Applying the above lemma through an induction proves the following theorem.
\begin{theorem} \label{thm:ful-alignment-main}
  There exists $\delta = O(\alpha_0 \eta \sqrt{\log(1/\eta)})$
  and
  $T_1 = O(\log \frac{1}{\eta} + \alpha_0 \sqrt{\log(1/\delta)})$
  such that the following holds.
  Let $\alpha_{\max} := \alpha_0 \sqrt{\log(1/\delta)}$.
  If the initial state is $O(\alpha_{\max})$-bounded,
  $O(\alpha_{\max})$-deviated,
  at most $O(1/\delta)$-misaligned,
  and satisfies $\normtwosm{{\bm{\phi}}_0 - {\bm{\zeta}}_0} \le O(\alpha_{\max}\eta)$,
  then
  at step $t = T_1$,
  the state is
  at most $O(\abssm{h_t}\eta)$-misaligned
  while still being
  $O(\alpha_{\max})$-bounded,
  $O(\alpha_{\max})$-deviated,
  and satisfying $\normtwosm{{\bm{\phi}}_t - {\bm{\zeta}}_0} \le O(\alpha_{\max} \eta)$.
\end{theorem}

\subsubsection{Drifting Phase}

After the alignment phase, the state is now at most $O(\abssm{h_t}\eta)$-misaligned.
Then we have the following lemma showing that $(h_t, u_t, {\bm{\phi}}_t)$ evolves
as an $O(1)$-RMS-drift process.
We defer the proof to \Cref{sec:ful-proof-drifting}.

\begin{lemma} \label{lm:ful-drift-2-step}
  For small enough base learning rate $\eta$,
  at any step $t$,
  if
  for some $1 \le \alpha \le \eta^{-o(1)}$,
  the state $({\bm{\theta}}_t, \tilde{v}_t)$ is $\alpha$-bounded
  and at most $O(\abssm{h_t}\eta)$-misaligned,
  then ${\bm{\theta}}_{t+2} \in \gZ^{\epsilon_0}$, and
  \begin{align}
    h_{t+2} &= (1 - 2 \eta u_t) h_t + \Osm{\alpha^2 \abssm{h_t}\eta^2}, \label{eq:ful-drift-2-step-h} \\
    \normtwosm{\mPHnzt({\bm{\phi}}_{t+2}) {\bm{x}}_{t+2}} &\le (1- 1.9\gamma_{\min})^2 \normtwosm{\mPHnzt({\bm{\phi}}_t) {\bm{x}}_t} + O(h_t^2 \eta^2), \label{eq:ful-drift-2-step-x-w} \\
    u_{t+2} &= u_t + 4 \eta h_t^2 (2 C_{\mathrm{b}} + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)}^2) - 2 \eta C_{\mathrm{b}} + \Osm{\alpha(1+h_t^2)\eta^2},  \label{eq:ful-drift-2-step-u} \\
    {\bm{\phi}}_{t+2} &= {\bm{\phi}}_t
    - 2 \eta^2 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + \Osm{\alpha h_t^2 \eta^3}. \label{eq:ful-drift-2-step-phi}
  \end{align}
  In other words, $(h_t, u_t) \to (h_{t+2}, u_{t+2})$ is an $O(1)$-RMS-drift transition
  if the state at step $t$ is $\alpha$-bounded and at most $O(\abssm{h_t}\eta)$-misaligned.
\end{lemma}

\subsection{Reduction to RMS-drift Process: The Case of Spherical Optimization} \label{sec:outline-sph}

Now we outline how to reduce the dynamics to an RMS-drift process
in the setting of \Cref{thm:main-ful},
the main theorem for the case of spherical optimization.
The basic logic is the same as the case of full space optimization,
so we only list the new lemma and theorem statements here.
The proofs in this section are deferred to
\Cref{sec:working-zone-proof-construction,sec:sph-proof}.

\subsubsection{Construction of Working Zones}

We still define ${\mathcal{Z}} := \{ {\bm{\zeta}}(t) : t \in [0, T] \} \subseteq \mathit{\Gamma}$.
The construction of working zone becomes the following.
We defer the proof to \Cref{sec:working-zone-proof-construction}.
\begin{lemma}[Working Zone Lemma] \label{lm:sph-working-zone}
  There exist $0 < \epsilon_0 < \epsilon_1$ such that $\gZ^{\epsilon_0}, \gZ^{\epsilon_1}$ satisfy the following:
  \begin{enumerate}
    \item $\mathrm{cl}(\gZ^{\epsilon_1}) \cap \mathit{\Gamma}$ is compact;
    \item $\mathcal{L}$ satisfies $\mu_{\mathtt{PL}}$-PL on $\gZ^{\epsilon_1}$ for some $\mu_{\mathtt{PL}} > 0$;
    \item $\Phi$ is well-defined on $\gZ^{\epsilon_0} \cap \sphS^{D-1}$ and $\Phi({\bm{\theta}}) \in \gZ^{\epsilon_1} \cap \manGa$ for all ${\bm{\theta}} \in \gZ^{\epsilon_0} \cap \sphS^{D-1}$;
    \item $\Phi$ is $\mathcal{C}^3$-smooth on $\gZ^{\epsilon_1} \cap \sphS^{D-1}$;
    \item $\gamma({\bm{\theta}}) \ge \gamma_{\min}$ holds uniformly on $\gZ^{\epsilon_1} \cap \manGa$ for some $\gamma_{\min} > 0$.
  \end{enumerate}
\end{lemma}

In the working zone,
we continue to define $\alpha$-bounded states,
$\alpha$-deviated states,
at most $p$-misaligned states
following the same definitions
as the full space case except that the definition of working zone is changed.

\begin{definition}[$\alpha$-Bounded State]
  We say that the state $({\bm{\theta}}_t, \tilde{v}_t)$ at some step $t$
  is $\alpha$-bounded if
  ${\bm{\theta}}_t \in \gZ^{\epsilon_0} \cap \sphS^{D-1}$,
  $\normtwosm{{\bm{x}}_t} \le \alpha \eta$ and $\abssm{u_t} \le \alpha$.
\end{definition}
\begin{definition}[$\alpha$-Deviated State]
  We say that the state $({\bm{\theta}}_t, \tilde{v}_t)$ at some step $t$
  is $\alpha$-deviated (from the manifold $\mathit{\Gamma}$) if
  $\normtwosm{{\bm{x}}_t} \ge \eta\exp(-\alpha^2)$ or
  ${\bm{\theta}}_t \notin \gZ^{\epsilon_0} \cap \sphS^{D-1}$.
\end{definition}

\begin{lemma} \label{lm:mPHz-x-small-sph}
  At any step $t$, if ${\bm{\theta}}_t \in \gZ^{\epsilon_0} \cap \sphS^{D-1}$, then $\normtwosm{\mPHz({\bm{\phi}}_t) {\bm{x}}_t} = O(\normtwosm{{\bm{x}}_t}^2)$.
\end{lemma}
\begin{proof}
  Direct consequence of \Cref{lm:phiphix-sph}.
\end{proof}
\begin{definition}[$p$-Misaligned State]
  We say that the state $({\bm{\theta}}_t, \tilde{v}_t)$ at some step $t$
  is at most $p$-misaligned (to the top eigenvector) if
  ${\bm{\theta}}_t \in \gZ^{\epsilon_0} \cap \sphS^{D-1}$,
  $\normtwosm{\mPHnzt({\bm{\phi}}_t) {\bm{x}}_t} \le p \cdot h_t \eta$.
\end{definition}

\subsubsection{Good Initialization}

\begin{lemma} \label{lm:sph-good-init}
  The same statement as \Cref{lm:ful-good-init} holds for the spherical case.
\end{lemma}

\subsubsection{Alignment Phase}

\begin{lemma} \label{lm:sph-align-1-step}
  For small enough base learning rate $\eta$,
  at any step $t$,
  if the state $({\bm{\theta}}_t, \tilde{v}_t)$ is $\alpha$-bounded
  for some $1 \le \alpha \le \eta^{-o(1)}$,
  then ${\bm{\theta}}_{t+1} \in \gZ^{\epsilon_0} \cap \sphS^{D-1}$, and
  \eqref{eq:ful-align-1-step-phi} to \eqref{eq:ful-align-1-step-u}
  in the full space case
  continue to hold in the spherical case.
\end{lemma}

\begin{theorem} \label{thm:sph-alignment-main}
    The same statement as \Cref{thm:ful-alignment-main} holds for the spherical case.
\end{theorem}

\subsubsection{Drifting Phase}

\begin{lemma}\label{lm:sph-drift-2-step}
  For small enough base learning rate $\eta$,
  at any step $t$,
  if
  for some $1 \le \alpha \le \eta^{-o(1)}$,
  the state $({\bm{\theta}}_t, \tilde{v}_t)$ is $\alpha$-bounded
  and at most $O(\abssm{h_t}\eta)$-misaligned,
  then ${\bm{\theta}}_{t+2} \in \gZ^{\epsilon_0} \cap \sphS^{D-1}$, and
  \eqref{eq:ful-drift-2-step-h} to \eqref{eq:ful-drift-2-step-phi}
  in the full space case
  continue to hold in the spherical case.
\end{lemma}

\subsection{RMS-drift Process: Analysis} \label{sec:outline-rmsdrift-analysis}

Now we outline how to obtain the final flow approximation \eqref{eq:zeta-general}
from the RMS-drift process.
We say that an RMS-drift process $S_0, \dots, S_{2M}$ is \textit{in the working zone}
if ${\bm{\phi}}_t \in \gZ^{\epsilon_1} \cap \manGa$
for all even numbers $0 \le t \le 2M$.
We focus on RMS-drift processes in the working zone,
and later we will show that the RMS-drift processes of interest
are indeed in the working zone.

First, we define a potential function that resembles the total energy (or Hamiltonian)
in physics.
\begin{definition}[Energy]
For a drift state $S = (h, u, {\bm{\phi}})$, we define the energy $E(S)$ as follows:
\[
  E(S) := \frac{1}{2} u^2 + (2C_{\mathrm{b}} + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}})}^2) h^2 + C_{\mathrm{b}} \log \frac{1}{\abssm{h}}.
\]
\end{definition}
If the energy is bounded by $\alpha^2$ at some step $t$,
then it is easy to see that the state at step $t$
is $O(\alpha)$-bounded and $O(\alpha)$-deviated.

The first key lemma is the conservation of energy in RMS-drift process,
which shows that the energy is preserved for $O(1/\eta^2)$ steps.
The proof is deferred to \Cref{sec:rmsdrift-proof-energy}.
\begin{theorem} \label{thm:rmsdrift-ec}
  For an $O(1)$-RMSdrift process $S_0, \dots, S_{2M}$ in the working zone,
  if $E(S_0) \le \alpha^2$ for some parameter $1 \le \alpha \le \eta^{-o(1)}$
  and $M = O(1/\eta^2)$,
  then $E(S_t) = O(\alpha^2)$ for all even numbers $0 \le t \le 2M$.
\end{theorem}

Next, we show that as long as the states are $\eta^{-o(1)}$-bounded,
the RMS-drift process tracks the sharpness-reduction flow nicely.
We defer the proof to \Cref{sec:rmsdrift-flow}.
\begin{theorem} \label{thm:rmsdrift-flow}
  Let ${\bm{\zeta}}: [0, T] \mapsto \gZ^{\epsilon_1} \cap \manGa$
  be a sharpness-reduction flow defined in \eqref{eq:zeta-general}.
  For an $O(1)$-RMS-drift process $S_0, \dots, S_{2M}$ in the working zone, where $M := \lfloor \frac{T}{2\eta^2} \rfloor$,
  if $\normtwosm{{\bm{\phi}}_0 - {\bm{\zeta}}(0)} \le O(\alpha^2 \eta^{1/2})$,
  and $S_t$ is $O(\alpha)$-bounded for all even numbers $0 \le t \le 2M$,
  then $\normtwosm{{\bm{\phi}}_t - {\bm{\zeta}}(t\eta^2)} \le O(\alpha^2 \eta^{1/2})$
  for all even numbers $0 \le t \le 2M$.
\end{theorem}

\subsection{Finalizing Proofs}

\begin{proof}[Proof for \Cref{thm:main-ful}]
First, we use \Cref{lm:ful-good-init}
to ensure a good initialization.
Then we apply \Cref{thm:ful-alignment-main}
to show that
the state at $t = T_1$
is at most $O(\abssm{h_t}\eta)$-misaligned,
$O(\alpha_{\max})$-bounded,
$O(\alpha_{\max})$-deviated,
and satisfies $\normtwosm{{\bm{\phi}}_t - {\bm{\zeta}}_0} \le O(\alpha_{\max} \eta)$.
By smoothness of ${\bm{\zeta}}$,
we also have $\normtwosm{{\bm{\phi}}_t - {\bm{\zeta}}(t \eta^2)}$.
Then we
consider the even-indexed and odd-indexed steps separately
and do an induction for each of them to show that the later dynamics
(1)~has bounded energy $O(\alpha_{\max}^2)$ (\Cref{thm:rmsdrift-ec});
(2)~stays at most $O(\abssm{h_t}\eta)$-misaligned (\Cref{lm:ful-drift-2-step});
(3)~follows the flow (\Cref{thm:rmsdrift-flow}).
The proof is done when the induction proceeds to $T / \eta^2$.
\end{proof}

\begin{proof}[Proof for \Cref{thm:main-sph}]
Same as above but we invoke the spherical version of lemmas and theorems.
\end{proof}

\section{Related Works} \label{sec:related}

\myparagraph{Flatness and Generalization.}  It has been long believed that minima locate in a flat valley generalize
better~\citep{hochreiter1997flat}. \citet{li2018visualizing} first visualized the loss landscape of neural networks and
found ``sharp'' minima generalizes worse. \citet{keskar2017large,wu2017towards,jastrzkebski2017three} empirically
verified the positive correlation between flatness and generalization.
\citet{neyshabur2017exploring} gave a theoretical
explanation for generalization benefit of low sharpness using PAC-Bayes theory~\citep{mcallester2003simplified}.
In light of this, \citet{foret2021sam} proposed SAM algorithm to improve the generalization of SGD by minimizing the
sharpness of the loss. However, the definition of sharpness remains ambiguous in face of the positive homogeneity and
invariance in ReLU networks~\citep{dinh2017sharp}, as networks with different sharpness may represent the same function.
Towards resolving this challenge, multiple definitions of scale-invariant sharpness have been
proposed~\citep{yi2019positively,yi2019bninvariant,tsuzuku20normalized,rangamani2021invariant}.
\citet{kwon2021asam} derived new algorithms with better generalization by defining new sharpness notions
that are aware of positive homogeneity and invariance.
\citet{he2019asymmetric} argued that the local minima of modern deep
networks are more than being flat or sharp and could be asymmetric valleys.
The theoretical implication of flatness has also been explored specifically for two-layer nets~\citep{blanc2020ou,mulayoff2021minimastability,haochen2021shape,li2022what,ding2022flat}
and deep linear nets~\citep{mulayoff2020unique}.

\myparagraph{Understanding Normalization Layers.} BN is proposed in~\citep{ioffe2015batch} with the original motivation being to reduce Internal Covariate Shift (ICS).
\citet{santurkar2018does} challenged this view by arguing that the effectiveness of BN comes from a smoothing effect on the training objective
that allows for faster training.
Several existing works have reported that optimization with normalization and weight decay can be far away from classical optimization theory.
Specifically, for training scale-invariant neural nets with weight decay,
\citet{li2020exp} showed that one can train the net to small loss even with learning rates exponentially increasing;
\citet{li2020reconciling} showed that the dynamic of (full-batch) GD may leave the stable regime;
\citet{lobacheva2021on} empirically showed that the training loss can exhibit a periodic behavior that sometimes improves the test accuracy. However, despite a lot of effort has been devoted to understanding the optimization and generalization benefits of normalization in various specific settings, \emph{e.g.}, \citep{van2017l2, bjorck2018understanding,zhang2018three,luo2019towards,de2020identity,lin2021spectral,daneshmand2021batch,lubana2021beyond,lange2022preconditioning,ma2021riemannian}, we still lack a complete and rigorous analysis for the role of normalization in general, especially in terms of generalization.

\section{Analysis of RMS-drift Process} \label{sec:rmsdrift-proof}

In this section, we provide proofs for theorems in \Cref{sec:outline-rmsdrift-analysis}.
For convenience, we define
$R({\bm{\theta}}) = \nabla \log \lambda^{\mathrm{H}}_1({\bm{\theta}})$,
$K_t := \sqrt{2C_{\mathrm{b}} + \normtwosm{\nabla_{\Gamma} R({\bm{\phi}}_t)}^2}$.
Then a $C_0$-RMS-drift transition
$S_t \to S_{t+2}$
can be written as
\begin{align*}
  h'_{t+2} &:= (1 - 2\eta u_t) h_t,  &
  \abssm{h_{t+2} - h'_{t+2}} &\le C_0 \alpha^2 \abssm{h_t} \eta^2, \\
  u'_{t+2} &:= u_t + 4 \eta K_t^2 h_t^2 - 2\eta C_{\mathrm{b}}, &
  \abssm{u_{t+2} - u'_{t+2}} &\le C_0 \alpha (1 + h_t^2) \eta^2,  \\
  {\bm{\phi}}'_{t+2} &:= {\bm{\phi}}_t - 2 \eta^2 h_t^2 \nabla_{\Gamma} R({\bm{\phi}}_t), &
  \normtwosm{{\bm{\phi}}_{t+2} - {\bm{\phi}}'_{t+2}} &\le C_0 \alpha h_t^2 \eta^3.
\end{align*}

\subsection{Conservation of Energy} \label{sec:rmsdrift-proof-energy}

To establish the conservation of energy,
we first compute the change in energy after one transition.
\begin{lemma} \label{lm:e-cons}
  Given two drift states $S_0 = (h_0, u_0, {\bm{\phi}}_0)$ and $S_2 = (h_2,
  u_2, {\bm{\phi}}_2)$ in the working zone,
  for learning rate $\eta$ and hyperparameter $C_{\mathrm{b}} > 0$,
  if $S_0$ is $\alpha$-bounded for some $1 \le \alpha \le \eta^{-o(1)}$,
  and
  $S_0 \to S_2$ is a $C_0$-RMSdrift transition,
  then
  \[
  	 E(S_2) - E(S_0) = \Obig{\alpha^2(1+h_0^2) \eta^2} = \begin{cases}
       O(\alpha^2 \eta^2) & \quad \abssm{h_0} \le 2, \\
       O(\alpha^2 h_0^2 \eta^2) & \quad \abssm{h_0} > 2. \\
     \end{cases}
  \]
\end{lemma}
\begin{proof}[Proof for \Cref{lm:e-cons}]
  $h_0$ and $h_2$ have the same sign when $\eta$ is small enough.
  We can decompose $E(S_2) - E(S_0)$ as follows:
	\begin{align*}
		E(S_2) - E(S_0) &= \underbrace{\frac{1}{2} (u_2^2 - u_0^2)}_{=: \delta_1}
    + \underbrace{K_0^2(h_2^2 - h_0^2)}_{=: \delta_2}
	+ \underbrace{(\normtwosm{\nabla_{\Gamma} R({\bm{\phi}}_2)}^2 - \normtwosm{\nabla_{\Gamma} R({\bm{\phi}}_0)}^2) h_2^2}_{=: \delta_3} + \underbrace{C_{\mathrm{b}} \log \frac{h_2}{h_0}}_{=: \delta_4}.
	\end{align*}
	Now we bound each error term. For $\delta_1$ and $\delta_2$, we use the formula $a^2 - b^2 = 2b(a-b) + (a-b)^2$:
	\begin{align*}
		\delta_1 &= u_0 (u_2 - u_0) + \tfrac{1}{2}(u_2 - u_0)^2 \\
		&= u_0 \left(4 \eta h_0^2K_0^2 - 2 \eta C_{\mathrm{b}} + O(\alpha(1+h_0^2) \eta^2)\right) + \Obig{(h_0^2 \eta)^2}. \\
		&= 4 \eta u_0 h_0^2K_0^2 - 2 \eta u_0 C_{\mathrm{b}} + O(\alpha^2(1+h_0^2) \eta^2). \\
		\delta_2 &= K_0^2\left(2h_0 (h_2 - h_0) + (h_2 - h_0)^2\right) \\
		&= K_0^2 \left(\left(-4\eta u_0 h_0^2 + O(\alpha^2 h_0^2 \eta^2) \right) + \Obig{(\alpha h_0 \eta)^2}\right) \\
		&= -4\eta u_0 h_0^2 K_0^2 + \Osm{\alpha^2 h_0^2 \eta^2}.
	\end{align*}
	For $\delta_3$, we use the Lipschitzness of $\nabla_{\Gamma} R({\bm{\theta}})$:
	\begin{align*}
		\delta_3 &= O(\normtwosm{{\bm{\phi}}_2 - {\bm{\phi}}_0}) \cdot h_2^2 = O(h_0^2\eta^2) \cdot O(h_0^2) = O(h_0^4 \eta^2).
	\end{align*}
	For $\delta_4$, note that $\log(1 + z) \le z + O(z^2)$ when $z = o(1)$. Then
	\begin{align*}
		\delta_4 = -C_{\mathrm{b}} \log\left(1 + \frac{h_2 - h_0}{h_0}\right) &= -C_{\mathrm{b}} \cdot \frac{h_2 - h_0}{h_0} + O(\alpha^2 \eta^2) \\
		&= 2\eta u_0 C_{\mathrm{b}} + O(\alpha^2 \eta^2).
	\end{align*}
	Adding $\delta_1, \delta_2, \delta_3, \delta_4$ together gives
	\begin{align*}
		E(S_2) - E(S_0) = & +  4 \eta u_0 h_0^2K_0^2 & &- 2 \eta u_0 C_{\mathrm{b}} & &+ O(\alpha^2(1+h_0^2) \eta^2) \\
		& - 4\eta u_0 h_0^2 K_0^2 & & & & + O(\alpha^2 h_0^2 \eta^2) \\
		& & & & & +O(h_0^4 \eta^2) \\
		& & & +2\eta C_{\mathrm{b}} u_0 & & + O(\alpha^2 \eta^2).
	\end{align*}
  So $E(S_2) - E(S_0) = O(\alpha^2(1+h_0^2) \eta^2)$.
\end{proof}

To sum up the energy change over time, we need
the following lemma.
\begin{lemma} \label{lm:sum-h2}
  For any $M = o((\alpha \eta)^{-2})$,
  if
  $S_0, \dots, S_{2M}$ is an $O(1)$-RMS-drift process in the working zone,
  and $S_t$ is $\alpha$-bounded for all even numbers $t \le 2M$, then
  \[
    \sum_{m = 0}^{M - 1} h_{2m}^2 = \frac{C_{\mathrm{b}}}{2 K_0^2} M + O(\alpha^2 \eta^2 M^2 + \alpha / \eta).
  \]
\end{lemma}
\begin{proof}[Proof for \Cref{lm:sum-h2}]
  By the update rule of $u_t$, we have
  \begin{align*}
    u_{2M} - u_0 = \sum_{m=0}^{M-1} u_{2m+2} - u_{2m} &= \sum_{m=0}^{M-1} \left(4\eta K_{2m}^2 h_{2m}^2 - 2 \eta C_{\mathrm{b}} + O(\alpha^3 \eta^2)\right).
  \end{align*}
  Since $S_{2m}$ is $\alpha$-bounded for all $m \le M$, ${\bm{\phi}}_{2m} - {\bm{\phi}}_0 = O(\alpha^2 \eta^2 M)$.
  By smoothness of $R$, $\normtwosm{\nabla_{\Gamma} R({\bm{\phi}}_{2m})}^2 = \normtwosm{\nabla_{\Gamma} R({\bm{\phi}}_0)}^2 + O(\alpha^2 \eta^2 M)$,
  then $K_{2m} = K_0 + O(\alpha^2 \eta^2 M)$.
  Since $S_0$ and $S_{2M}$ are $\alpha$-bounded, $u_T - u_0 = O(\alpha)$.
  Combining all these together,
  \begin{align*}
    O(\alpha) &= \sum_{m=0}^{M-1} \left(4\eta \left(K_0^2 + O(\alpha^2 \eta^2 M)\right) h_{2m}^2 - 2 \eta C_{\mathrm{b}} +O(\alpha^3 \eta^2) \right)
  \end{align*}
  Let $Q := \sum_{m=0}^{M-1} h_{2m}^2$.
  Then we have
  \begin{align*}
    O(\alpha) = 4\eta \left(K_0^2 + O(\alpha^2 \eta^2 M)\right) Q - 2\eta C_{\mathrm{b}} M + O(\alpha^3 \eta^2 M).
  \end{align*}
  Rearranging the terms while noting that $\alpha^3 \eta^2 M = o(\alpha)$, we have
  \begin{align*}
    4\eta \left(K_0^2 + O(\alpha^2 \eta^2 M)\right) Q = 2\eta C_{\mathrm{b}} M + O(\alpha).
  \end{align*}
  So we can estimate $Q$ by
  \begin{align*}
    Q &= \frac{2C_{\mathrm{b}} M + O(\alpha / \eta)}{4K_0^2 + O(\alpha^2 \eta^2 M)} = \frac{C_{\mathrm{b}}}{2 K_0^2} M + O(\alpha^2 \eta^2 M^2 + \alpha / \eta),
  \end{align*}
  which completes the proof.
\end{proof}

\begin{lemma} \label{lm:rmsdrift-ec-15}
  For an $O(1)$-RMSdrift process $S_0, \dots, S_{2M}$ in the working zone,
  if $E(S_0) \le \alpha^2$ for some parameter $1 \le \alpha \le \eta^{-o(1)}$
  and $M = O(1/\eta^{1.5})$,
  then $E(S_t) = E(S_0) + O(\alpha^2 \eta^{0.5})$ for all even numbers $0 \le t \le 2M$.
\end{lemma}
\begin{proof}[Proof for \Cref{lm:rmsdrift-ec-15}]
    We do a bootstrap.
    As $E(S_0) \le O(\alpha^2)$, we can leverage \Cref{lm:e-cons} to prove by induction that $E(S_t) \le E(S_0) + O(\alpha^4 \eta^2 M)$
    for all $0 \le t \le 2M$.
    Then we apply \Cref{lm:e-cons} again. For all $N \le M$,
    \begin{align*}
        E(S_{2N}) - E(S_0) &= \sum_{m = 0}^{N-1} O(\alpha^2 (1 + h_{2m}^2) \eta^2) \\
        &= O(\alpha^2\eta^2) \left( N + \sum_{m = 0}^{N-1} h_{2m}^2\right) \\
        &\le O(\alpha^2\eta^2) \left( N + O(N + \alpha^2 \eta^2 N^2 + \alpha / \eta)\right) \\
        &\le O(\alpha^2\eta^{0.5}).
    \end{align*}
    where the third line uses \Cref{lm:sum-h2}.
\end{proof}

\begin{proof}[Proof for \Cref{thm:rmsdrift-ec}]
  We group the $M$ transitions into $O(1/\eta^{0.5})$
  segments of length $O(1/\eta^{1.5})$.
  We can do an induction with
  \Cref{lm:rmsdrift-ec-15} applied
  on each segment to show that for all $t$ in the $k$-th segment,
  $E(S_t) \le (1 + O(\eta^{0.5}))^k E(S_0)$.
  Noting that $k = O(1/\eta^{0.5})$ and $(1 + O(\eta^{0.5}))^{O(1/\eta^{0.5})} = O(1)$ finishes the proof.
\end{proof}

\subsection{Flow Approximation} \label{sec:rmsdrift-flow}

\begin{lemma} \label{lm:rmsdrift-flow-15}
  In the setting of \Cref{thm:rmsdrift-flow} but with $M = \Theta(1/\eta^{1.5})$,
  if $S_t$ is $\alpha$-bounded for all even numbers $0 \le t \le 2M$, then
  \[
    \normtwosm{{\bm{\phi}}_t - {\bm{\zeta}}(t\eta^2)} \le O(\alpha^2 \eta).
  \]
\end{lemma}
\begin{proof}[Proof for \Cref{lm:rmsdrift-flow-15}]
  When $S_t$ is $\alpha$-bounded, $\normtwosm{{\bm{\phi}}_{t+2} - {\bm{\phi}}_t} =
  O(\alpha^2 \eta^2)$, so $\normtwosm{{\bm{\phi}}_t - {\bm{\phi}}_0} = O(\alpha^2
  \eta^2 M) = O(\alpha^2 \eta^{0.5})$.
  For every $N \le M$
  we have
  \begin{align*}
    {\bm{\phi}}_{2N} - {\bm{\phi}}_0 = \sum_{m=0}^{N-1} {\bm{\phi}}_{2m+2} - {\bm{\phi}}_{2m}
    &= \sum_{m=0}^{N-1} \left(-2\eta^2h_{2m}^2 \nabla_{\Gamma} R({\bm{\phi}}_{2m}) + O(\alpha^3 \eta^3) \right) \\
    &= \sum_{m=0}^{N-1} \left(-2\eta^2h_{2m}^2 (\nabla_{\Gamma} R({\bm{\phi}}_0) + O(\alpha^2 \eta^{0.5})) + O(\alpha^3 \eta^3) \right) \\
    &= \sum_{m=0}^{N-1} \left(-2\eta^2h_{2m}^2 \nabla_{\Gamma} R({\bm{\phi}}_0) + O(\alpha^2 \eta^{2.5}) \right) \\
    &= -\underbrace{2 \eta^2 \left(\sum_{m = 0}^{N-1} h_{2m}^2\right)}_{=: \delta} \nabla_{\Gamma} R({\bm{\phi}}_0) + O(\alpha^2 \eta).
  \end{align*}
  By \Cref{lm:sum-h2}, we have
  \begin{align*}
    \delta = 2\eta^2 \cdot \left(\frac{C_{\mathrm{b}}}{2K_0^2} N + O(\alpha^2 \eta^2 N^2 + \alpha / \eta)\right)
    &= \frac{C_{\mathrm{b}}}{K_0^2} \eta^2 N + 2\eta^2 \cdot O(\alpha^2 / \eta) \\
    &= \frac{C_{\mathrm{b}}}{K_0^2} \eta^2 N + O(\alpha^2 \eta).
  \end{align*}
  Note that $\frac{C_{\mathrm{b}}^2}{2 K_0^2} \nabla_{\Gamma} R({\bm{\phi}}_0) = \frac{\textup{\textrm{d}}}{\textup{\textrm{d}} t} {\bm{\zeta}}(0) + O(\alpha^2 \eta^{1/2})$. Then we have
  \begin{align*}
    {\bm{\phi}}_{2N} &= {\bm{\phi}}_0 -\delta \nabla_{\Gamma} R({\bm{\phi}}_0) + O(\alpha^2 \eta) \\
    &= {\bm{\phi}}_0 - (2\eta^2 N) \left(\frac{\textup{\textrm{d}}}{\textup{\textrm{d}} t}{\bm{\zeta}}({\bm{\phi}}_0) + O(\alpha^2 \eta^{1/2})\right) + O(\alpha^2 \eta) \\
    &= {\bm{\zeta}}(2N\eta^2) + \Obig{(\eta^2 N)^2} + O(\alpha^2 \eta) \\
    &= {\bm{\zeta}}(2N\eta^2) + O(\alpha^2 \eta),
  \end{align*}
  where the third equality uses the smoothness of ${\bm{\zeta}}(t)$.
\end{proof}

\begin{proof}[Proof for \Cref{thm:rmsdrift-flow}]
  We group the $M$ transitions into $O(1/\eta^{1/2})$
  segments of length $O(1/\eta^{1.5})$.
  Then we leverage
  \Cref{lm:rmsdrift-flow-15}
  and do an induction to show that $\normtwosm{{\bm{\phi}}_t - {\bm{\zeta}}(t\eta^2)} \le O(\alpha^2 \eta^{1/2})$ for all even numbers $0 \le t \le 2M$.
\end{proof}

\section{GD+WD\xspace on Scale-Invariant Loss Functions} \label{sec:main-result}

This section  analyzes GD+WD\xspace \eqref{eq:upd-gdwd} on a
scale-invariant loss $\mathcal{L}({\bm{w}})$, in particular what happens after the approximate convergence of Theorem~\ref{thm:li-gdwd}.
We use ${\bm{w}}_t$ to denote the trainable parameter at step $t$
and ${\bm{\theta}}_t := \frac{{\bm{w}}_t}{\normtwosm{{\bm{w}}_t}}$ to denote its projection onto $\mathbb{S}^{D-1}$.
\Cref{sec:main-result-1} analyzes the dynamics in the stable regime, where loss is guaranteed to decrease monotonically,
and \Cref{thm:li-gdwd} suggests ${\bm{w}}_t$ can get close to a local minimizer at some time $t_0$.
We show that the effective LR is increasing after $t_0$, causing GD+WD\xspace to eventually leave this stable regime
and enters a new regime which we call the Edge of Stability (EoS).
In \Cref{sec:main-result-2}, we establish our main theorem,
which connects the dynamics of ${\bm{w}}_t$ in the EoS regime to a sharpness-reduction flow.

\subsection{GD+WD\xspace Eventually Leaves the Stable Regime} \label{sec:main-result-1}

A standard step for analyzing optimization methods is to do Taylor expansion locally for the loss function,
and show that how the optimization method decreases the loss using a \textit{descent lemma}.
In our case of scale-invariant loss functions,
we use ${\bm{H}}({\bm{w}}) := \nabla^2 \mathcal{L}({\bm{w}}) \in \mathbb{R}^{D \times D}$ to denote the Hessian matrix of $\mathcal{L}$ at ${\bm{w}} \in \mathbb{R}^D$,
and $\lambda^{\mathrm{H}}_1({\bm{w}}) := \lambda_1({\bm{H}}({\bm{w}}))$ to denote the top eigenvalue of ${\bm{H}}({\bm{w}})$.
\begin{lemma}[Descent Lemma] \label{lm:descent}
  For scale-invariant loss $\mathcal{L}({\bm{w}})$,
  at step $t$ of GD+WD\xspace we have
  \begin{align*}
    \mathcal{L}({\bm{\theta}}_{t+1}) \le \mathcal{L}({\bm{\theta}}_t) - \tilde{\eta}_t (1 - \tilde{\eta}_t \lambda_{\max}^{(t)} / 2) \normtwosm{\nabla \mathcal{L}({\bm{\theta}}_t)}^2.
  \end{align*}
  where $\lambda_{\max}^{(t)} := \sup_{\alpha \in [0, \tilde{\eta}_t]} \left\{ \lambda^{\mathrm{H}}_1({\bm{\theta}}_t - \alpha \nabla \mathcal{L}({\bm{\theta}}_t)) \right\}$ is an upper bound of spherical sharpness locally.
\end{lemma}
This descent lemma shows that
the training loss $\mathcal{L}({\bm{\theta}}_t)$ is decreasing
as long as the effective LR $\tilde{\eta}_t$ is smaller than $2/\lambda_{\max}^{(t)}$,
We call the regime of $\tilde{\eta}_t < 2 / \lambda_{\max}^{(t)}$ as the \textit{stable regime} of GD+WD\xspace.
If $\tilde{\eta}_t \approx 2 / \lambda_{\max}^{(t)}$ with a small difference, then we call it as the \textit{Edge of Stability (EoS) regime}.
This EoS regime is conceptually the same as that defined by \citet{cohen2021gradient}; see \Cref{sec:eos-cohen} for discussion.

Fix an initial point ${\bm{w}}_0 \in \mathbb{R}^D \setminus \{{\bm{0}}\}$.
Now we aim to characterize the dynamics of GD+WD\xspace when LR $\hat{\eta}$ and WD $\hat{\lambda}$ are small enough.
\Cref{thm:li-gdwd} shows that
for some $t_0 \le T_0$,
 $\normtwosm{\nabla \mathcal{L}({\bm{\theta}}_{t_0})}^2 \le O(\hat{\lambda} \hat{\eta})$ and $\tilde{\eta}_{t_0} \le \frac{1}{\pi^2 \rho_2} < \frac{2}{\rho_2}$,
which means ${\bm{\theta}}_{t_0}$ is  an approximate first-order stationary point of $\mathcal{L}$ on the unit sphere.
This does not guarantee that ${\bm{\theta}}_{t_0}$ is close to any global minimizer,
but in practice the training loss rarely gets stuck at a non-optimal value
when the model is overparameterized~\citep{lee2016minimizers,panageas2017only,lee2017first,zhang2017rethinking}.
We are thus motivated to study the case where
${\bm{\theta}}_{t_0}$ not only has small gradient
$\normtwosm{\nabla \mathcal{L}({\bm{\theta}}_{t_0})}^2 \le O(\hat{\lambda} \hat{\eta})$
but also is close to a local minimizer $\vtheta^* \in \mathbb{S}^{D-1}$ of $\mathcal{L}$
in the sense that $\normtwosm{{\bm{\theta}}_{t_0} - \vtheta^*} \le O((\hat{\lambda} \hat{\eta})^{1/2})$
(note that the latter implies the former assuming smoothness).

\begin{wrapfigure}{r}{0.39\textwidth}
\vspace{-0.1in}
  \centering
    \begin{subfigure}[b]{0.19\textwidth}
      \centering
      \includegraphics[width=\textwidth]{fig/gdwd-a1.png}
      \caption{}
      \label{fig:norm-smaller}
    \end{subfigure}%
    \begin{subfigure}[b]{0.19\textwidth}
      \centering
      \includegraphics[width=\textwidth]{fig/gdwd-a2.png}
      \caption{}
      \label{fig:norm-bigger}
    \end{subfigure}
    \begin{subfigure}[b]{0.38\textwidth}
      \includegraphics[width=\textwidth]{fig/ex1.png}
      \caption{}
      \label{fig:ex3d}
    \end{subfigure}
  \caption{\textbf{(a), (b)}: The norm of ${\bm{w}}_t$ decreases when gradient is small and increases when gradient is large.
  \textbf{(c)}: The trajectory of ${\bm{\theta}}_t$ on a 3D scale-invariant loss function.
  Darker color means lower loss on the unit sphere,
  and points in the black line are minimizers
  (see \Cref{sec:3d}).
  In the end, ${\bm{\theta}}_t$ approaches the flattest one (red star).}
  \label{fig:ex}
  \vspace{-0.7in}
\end{wrapfigure}

As the gradient is small
near the local minimizer $\vtheta^*$,
starting from step $t_0$,
the norm of ${\bm{w}}_t$ decreases due to the effect of WD. See \Cref{fig:norm-smaller}.
Since the effective LR
is inversely proportional
to $\normtwosm{{\bm{w}}_t}^2$,
this leads to the effective LR to increase.
Then \Cref{thm:eos-near-opt} will
show that the GD+WD\xspace dynamic eventually leaves the stable regime at some time $t_1 > t_0$,
and enters the EoS regime where $\tilde{\eta}_t \approx 2/\lambda_{\max}^{(t)}$.

To establish \Cref{thm:eos-near-opt}, we need to assume that
$\mathcal{L}$ satisfies Polyak-\L{}ojasiewicz\xspace (PL) condition locally,
which is a standard regularity condition in the optimization literature to ease theoretical analysis
around a minimizer.
Intuitively, PL condition guarantees that the gradient grows faster than a quadratic function as we move a parameter ${\bm{\theta}}$ away from $\vtheta^*$.
Note that PL condition is strictly weaker than convexity as the function can still be non-convex under PL condition~(see, e.g., \citep{karimi2016pl}).
\begin{definition}[Polyak-\L{}ojasiewicz\xspace Condition] \label{def:pl}
  For a scale-invariant loss $\mathcal{L}({\bm{w}})$ and $\mu > 0$, we say that $\mathcal{L}$ satisfies $\mu$-Polyak-\L{}ojasiewicz\xspace condition
  (or $\mu$-PL)
  locally around a local minimizer $\vtheta^*$ on $\mathbb{S}^{D-1}$ if
  for some neighborhood $U \subseteq \mathbb{S}^{D-1}$ of $\vtheta^*$,
  $\forall {\bm{\theta}} \in U: \frac{1}{2}\normtwosm{\nabla \mathcal{L}({\bm{\theta}})}^2 \ge \mu \cdot (\mathcal{L}({\bm{\theta}}) - \mathcal{L}(\vtheta^*))$.
\end{definition}
\begin{theorem} \label{thm:eos-near-opt}
  Let $\mathcal{L}({\bm{w}})$ be a $\mathcal{C}^2$-smooth scale-invariant loss that satisfies $\mu$-PL around a local minimizer $\vtheta^*$ on the unit sphere,
  and $\rho_2 := \sup\{ \normtwosm{\nabla^2 \mathcal{L}({\bm{w}})} : {\bm{w}} \in \mathbb{S}^{D-1}\}$.
  For GD+WD\xspace on $\mathcal{L}({\bm{w}})$ with learning rate $\hat{\eta}$ and weight decay $\hat{\lambda}$,
  if at some step $t_0$, $\normtwosm{{\bm{\theta}}_{t_0} - \vtheta^*} \le O((\hat{\lambda} \hat{\eta})^{1/2})$
  and $\tilde{\eta}_{t_0} \le \frac{2}{\rho_2} < \frac{2}{\lambda^{\mathrm{H}}_1(\vtheta^*)}$,
  and if $\hat{\lambda} \hat{\eta}$ is small enough,
  then there exists a time $t_1 > t_0$ such that $\normtwosm{{\bm{\theta}}_{t_1} - \vtheta^*} = O((\hat{\lambda} \hat{\eta})^{1/2})$
  and $\tilde{\eta}_{t_1} = \frac{2}{\lambda^{\mathrm{H}}_1(\vtheta^*)} + O((\hat{\lambda} \hat{\eta})^{1/2})$.
\end{theorem}

\subsection{Dynamics at the Edge of Stability} \label{sec:main-result-2}
From the analysis in the previous subsection, we know that ${\bm{\theta}}_t$ can get close to a local minimizer $\vtheta^*$
and enter the EoS regime at some step $t_1$.
But what happens after $t_1$?

\Cref{fig:ex3d} gives a warm-up example on a 3D scale-invariant loss $\mathcal{L}: \mathbb{R}^3 \to \mathbb{R}$,
where the black line is a manifold $\mathit{\Gamma}$ consisting of all the minimizers.
In training with GD+WD\xspace,
${\bm{\theta}}_t$ first goes close to a local minimizer ${\bm{\zeta}}_0$,
then \Cref{thm:eos-near-opt} suggests
that WD causes the effective LR to steadily increase until the EoS regime is reached.
Now something interesting happens ---
${\bm{\theta}}_t$ moves a bit away from ${\bm{\zeta}}_0$ and starts to oscillate
around the manifold $\mathit{\Gamma}$.
This oscillation
is not completely perpendicular to $\mathit{\Gamma}$ but actually forms a small angle that pushes
${\bm{\theta}}_t$ to move downward persistently until ${\bm{\theta}}_t$ approaches the minimizer ${\bm{\zeta}}_*$ denoted in the plot.

For a general scale-invariant loss $\mathcal{L}: \mathbb{R}^D \to \mathbb{R}$,
which minimizer does ${\bm{\theta}}_t$ move towards?
In this work, we consider the setting where there is a manifold $\mathit{\Gamma}$ consisting only of local minimizers (but not necessarily all of them).
We show that
${\bm{\theta}}_t$ always oscillates around the manifold once it approaches to the manifold and enters the EoS regime,
and meanwhile ${\bm{\theta}}_t$ keeps moving in a direction of reducing spherical sharpness.

\subsubsection{Assumptions} \label{sec:main-result-2-assumption}

Now we formally introduce our main assumption on the local minimizer manifold $\mathit{\Gamma}$.
\begin{assumption} \label{ass:man}
  The loss function $\mathcal{L}: \mathbb{R}^D \setminus \{{\bm{0}}\} \to \mathbb{R}$ is $\mathcal{C}^4$-smooth and scale-invariant.
  $\mathit{\Gamma}$ is a $\mathcal{C}^2$-smooth, $(D_\Gamma - 1)$-dimensional open submanifold of $\mathbb{S}^{D-1}$
  for some $0 \le D_\Gamma < D$, where every ${\bm{\theta}} \in \mathit{\Gamma}$ is a local minimizer of $\mathcal{L}$ on $\mathbb{S}^{D-1}$
  and $\mathrm{rank}({\bm{H}}({\bm{\theta}})) = D - D_\Gamma$.
\end{assumption}
Scale-invariance has become a standard assumption
in studying neural nets with normalization layers~\citep{li2020exp,li2020reconciling,lobacheva2021on}.
For VGG and ResNet, the scale-invariance can be ensured
after making minor changes to the architectures (see \Cref{sec:exp-details-cifar10}).
The training loss $\mathcal{L}$ may not be smooth if the activation is ReLU,
but lately it has become clear that differentiable activations
such as Swish~\citep{ramachandran2017swish}, GeLU~\citep{hendrycks2016gelu}
can perform equally well.
Swish is indeed used in our VGG-11 experiments (\Cref{fig:full-cifar}),
but ResNet with ReLU activation also exhibits a sharpness-reduction bias empirically (see \Cref{sec:exp-cifar10}).

For any local minimizer ${\bm{\theta}} \in \mathit{\Gamma}$,
the eigenvalues $\lambda^{\mathrm{H}}_k({\bm{\theta}})$ must be non-negative.
And $\lambda^{\mathrm{H}}_k({\bm{\theta}}) = 0$ for all $D - D_\Gamma < k \le D$,
since $\mathit{\Gamma}$ is of dimension $D_\Gamma - 1$.
The condition $\mathrm{rank}({\bm{H}}({\bm{\theta}})) = D - D_\Gamma$ ensures that
the Hessian is maximally non-degenerate on $\mathit{\Gamma}$,
which also appears as a key assumption
in previous works \citep{li2022what,arora2022understanding,fehrman2020convergence}.
This condition simplifies the calculus on $\mathit{\Gamma}$ in our analysis as it ensures that the null space of the matrix ${\bm{H}}({\bm{\theta}})$ equals to the tangent space of $\mathit{\Gamma}$ at ${\bm{\theta}} \in \mathit{\Gamma}$.
It is also closely related to PL condition (\Cref{def:pl}) as
\Cref{ass:man}
implies that $\mathcal{L}({\bm{\theta}})$ satisfies $\mu$-PL (for some $\mu > 0$) locally around every ${\bm{\theta}} \in \mathit{\Gamma}$ on the unit sphere (\citet{arora2022understanding}, Lemma B.3).

To ease our analysis, we also need the following regularity condition to ensure that the largest eigenvalue is unique.
In our experiments, sharpness reduction happens even when the multiplicity of the top eigenvalue is more than $1$, but we leave the analysis of that case to future work.
\begin{assumption} \label{ass:topeigen}
  For all ${\bm{\theta}} \in \mathit{\Gamma}$, $\lambda^{\mathrm{H}}_1({\bm{\theta}}) > \lambda^{\mathrm{H}}_2({\bm{\theta}})$. That is, the top eigenvalue of ${\bm{H}}({\bm{\theta}})$ is unique.
\end{assumption}

\subsubsection{Main Theorem} \label{sec:main-thm-stat}
First, we define $\eta_{\mathrm{in}} := \hat{\eta} \hat{\lambda}$ as the intrinsic learning rate (name from \citet{li2020reconciling}) for convenience.
As suggested in \Cref{thm:li-gdwd,thm:eos-near-opt},
${\bm{\theta}}_t$ can get close to a local minimizer and be in the EoS regime at some step $t_1$:
if ${\bm{\zeta}}_0$ is the local minimizer, then $\normtwosm{{\bm{\theta}}_{t_1} - {\bm{\zeta}}_0} = O(\eta_{\mathrm{in}}^{1/2})$
and $\tilde{\eta}_{t_1} = \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\zeta}}_0)} + O(\eta_{\mathrm{in}}^{1/2})$.
In our main theorem, we start our analysis from step $t_1$ while setting $t_1 = 0$ WLOG (otherwise we can shift the step numbers).
We connect GD+WD\xspace in the EoS regime
to the following gradient flow~\eqref{eq:main-zeta} on the manifold $\mathit{\Gamma}$ minimizing spherical sharpness (with gradient-dependent learning rate),
and show that one step of GD+WD\xspace tracks a time interval of $\eta_{\mathrm{in}}$ in the gradient flow.
\begin{equation}
    {\bm{\zeta}}(0) = {\bm{\zeta}}_0 \in \mathit{\Gamma}, \qquad \frac{\textup{\textrm{d}}}{\textup{\textrm{d}} \tau} {\bm{\zeta}}(\tau) = -\frac{2\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau))}{4 + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau))}^2}.
    \label{eq:main-zeta}
\end{equation}
Here we use the notation $\nabla_{\Gamma} R({\bm{\theta}})$
for any $R: \mathbb{R}^D \to \mathbb{R}$
to denote the projection of $\nabla R({\bm{\theta}})$ onto the tangent space $\TGa{{\bm{\theta}}}$ at ${\bm{\theta}} \in \mathit{\Gamma}$.
${\bm{\zeta}}(\tau)$ reduces sharpness as it moves in direction of
the negative gradient of $\log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau))$ on $\mathit{\Gamma}$.
A simple chain rule shows how fast the spherical sharpness decreases:
\begin{equation*}
  \frac{\textup{\textrm{d}}}{\textup{\textrm{d}} t} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau)) = -\frac{2 \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau))}^2}{4 + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau))}^2}
  \approx \begin{cases}
    -\frac{1}{2} \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(\tau))}^2 &~~~\text{for small gradient}; \\
    -2 &~~~\text{for large gradient}.
  \end{cases}
\end{equation*}
Note that it is not enough to just assume that ${\bm{\theta}}_0$ is close to ${\bm{\zeta}}_0$.
If ${\bm{\theta}}_0 = {\bm{\zeta}}_0$ holds exactly,
then the subsequent dynamic of ${\bm{w}}_t$ is described by ${\bm{w}}_t = (1 - \hat{\eta} \hat{\lambda})^t {\bm{w}}_0$
with direction unchanged.
There are also some other bad initial directions of ${\bm{w}}_0$ that may not lead to the sharpness-reduction bias.
This motivates us to do a smoothed analysis for the initial direction:
the initial direction is ${\bm{\zeta}}$ with tiny random perturbation,
where the perturbation scale is allowed to vary from $\exp(-\eta_{\mathrm{in}}^{-o(1)})$ to $\eta_{\mathrm{in}}^{1/2-o(1)}$,
and we show that a good initial direction is met
with high probability as $\eta_{\mathrm{in}} \to 0$.%
\footnote{Here $\eta_{\mathrm{in}}^{-o(1)}$ can be constant, $O(\log(1/\eta_{\mathrm{in}}))$, or $O(\mathrm{polylog}(1/\eta_{\mathrm{in}}))$, but not $\eta_{\mathrm{in}}^{-\epsilon}$ if $\epsilon > 0$ is a constant. As mentioned later, this need for random initialization is very similar to the one needed in power method for computing eigenvalues.}
Alternatively, one can regard it as a modeling
of the tiny random noise in GD+WD\xspace due to the precision errors in floating-point operations.
See Figure \hyperref[fig:lin-idea]{\ref{fig:lin-idea}b};
the training loss can never be exactly zero in practice.

\myparagraph{Initialization Scheme.}
Given a local minimizer ${\bm{\zeta}}_0 \in \mathit{\Gamma}$,
we initialize ${\bm{w}}_0 \in \mathbb{R}^D \setminus \{{\bm{0}}\}$ as follows:
draw
${\bm{\xi}} \sim \mathcal{N}({\bm{0}}, \sigma_0^2 {\bm{I}} / D)$
from Gaussian and set the direction of ${\bm{w}}_0$ to $\frac{{\bm{\zeta}}_0 + {\bm{\xi}}}{\normtwosm{{\bm{\zeta}}_0 + {\bm{\xi}}}}$,
where $\sigma_0$ can take any value in $[\exp(-\eta_{\mathrm{in}}^{-o(1)}), \eta_{\mathrm{in}}^{1/2-o(1)}]$;
then set the parameter norm $\normtwosm{{\bm{w}}_0}$
to be any value that satisfies $\abs{\tilde{\eta}_0 - \frac{2}{\lambda^{\mathrm{H}}_1({\bm{\zeta}}_0)}} \le \eta_{\mathrm{in}}^{1/2-o(1)}$,
where $\tilde{\eta}_0 := \frac{\hat{\eta}}{(1-\hat{\eta}\hat{\lambda}) \normtwosm{{\bm{w}}_0}^2}$ is the effective LR for the first step.

\begin{theorem} \label{thm:gdwd-main}
  Under \Cref{ass:man,ass:topeigen},
  for GD+WD\xspace \eqref{eq:upd-gdwd}
  with sufficiently small
  intrinsic learning rate $\eta_{\mathrm{in}} := \hat{\eta} \hat{\lambda}$,
  if we follow the above initialization scheme for some ${\bm{\zeta}}_0 \in \mathit{\Gamma}$,
  then with probability $1 - O(\eta_{\mathrm{in}}^{1/2-o(1)})$,
  the trajectory of ${\bm{\theta}}_t := \frac{{\bm{w}}_t}{\normtwosm{{\bm{w}}_t}}$ approximately tracks a sharpness-reduction flow
  ${\bm{\zeta}}: [0, T] \to \mathit{\Gamma}$ that starts from ${\bm{\zeta}}_0$ and evolves as the ODE \eqref{eq:main-zeta} up to time $T$ (if solution exists),
  in the sense that $\normtwosm{{\bm{\theta}}_t - {\bm{\zeta}}(t \eta_{\mathrm{in}})} = O(\eta_{\mathrm{in}}^{1/4 - o(1)})$ for all $0 \le t \le T / \eta_{\mathrm{in}}$.
\end{theorem}

\begin{remark}[Magnitude of Oscillation]
  As suggested by \Cref{fig:ex3d}, ${\bm{\theta}}_t$ actually oscillates
  around the manifold.
  But according to our analysis,
  the magnitude of oscillation is as small as $O(\eta_{\mathrm{in}}^{1/2-o(1)})$,
  so it is absorbed into our final bound $O(\eta_{\mathrm{in}}^{1/4-o(1)})$
  for the distance between ${\bm{\theta}}_t$ and ${\bm{\zeta}}(t \eta_{\mathrm{in}})$.
\end{remark}

\subsubsection{Proof Idea} \label{sec:proof-idea}
Throughout our proof, we view GD+WD\xspace for ${\bm{w}}_t$ as a PGD for ${\bm{\theta}}_t$
with effective LR $\tilde{\eta}_t$ (\Cref{lm:scale-invariant-on-sphere}).
To track ${\bm{\theta}}_t$ with ${\bm{\zeta}}(t \eta_{\mathrm{in}})$,
for each step $t$, we construct a local minimizer ${\bm{\phi}}_t \in \mathit{\Gamma}$
that serves as the ``projection'' of ${\bm{\theta}}_t$ onto the manifold $\mathit{\Gamma}$,
in the sense that the displacement ${\bm{x}}_t := {\bm{\theta}}_t - {\bm{\phi}}_t$
is approximately perpendicular to the tangent space of $\mathit{\Gamma}$ at ${\bm{\phi}}_t$.
Our entire proof works through induction.
According to the initial conditions,
the dynamic is initially in the EoS regime: $\normtwosm{{\bm{x}}_t} \le \eta_{\mathrm{in}}^{1/2-o(1)}$
and $\abssm{\tilde{\eta}_t - 2/\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)} \le \eta_{\mathrm{in}}^{1/2-o(1)}$ at $t = 0$.
In our induction, we maintain
the induction hypothesis that these two EoS conditions continue to hold for all $t \ge 0$.

\myparagraph{Period-Two Oscillation.}
A key insight in our proof is that
after a few initial steps,
${\bm{\theta}}_t$ is oscillating around ${\bm{\phi}}_t$
along the $\pm \vv^{\mathrm{H}}_1({\bm{\theta}})$ directions,
where
$\vv^{\mathrm{H}}_1({\bm{\theta}})$
is a unit top eigenvector of ${\bm{H}}({\bm{\theta}})$
and is chosen in a way
that $\vv^{\mathrm{H}}_1({\bm{\theta}})$ is continuous on $\mathit{\Gamma}$.
More specifically,
${\bm{x}}_t =
h_t\vv^{\mathrm{H}}_1({\bm{\phi}}_t) + O(\normtwosm{{\bm{x}}_t}^2)$
for $h_t := \dotpsm{{\bm{x}}_t}{\vv^{\mathrm{H}}_1({\bm{\phi}}_t)}$.
The oscillation is of period 2: $h_t > 0$ when $t$ is even and $h_t < 0$ when $t$ is odd.
See Figure \hyperref[fig:lin-idea]{\ref*{fig:lin-idea}d} for an example.

This oscillation can be connected to a power method for the matrix ${\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}_t)$.
In the EoS regime,
we can approximate ${\bm{\theta}}_{t+1}$ (when ${\bm{x}}_t$ is small)
as
${\bm{\theta}}_{t+1}
  = \Pi({\bm{\theta}}_t - \tilde{\eta}_t\nabla \mathcal{L}({\bm{\theta}}_t))
  \approx \Pi({\bm{\theta}}_t - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}_t) {\bm{x}}_t)
  \approx {\bm{\theta}}_t - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}_t) {\bm{x}}_t$
by Taylor expansions
of $\nabla \mathcal{L}$
and $\Pi: \mathbb{R}^D \to \mathbb{S}^{D-1}$.
We can further show that ${\bm{\phi}}_{t+1} \approx {\bm{\phi}}_t$
due to our choice of projections.
Then the connection to power method is shown below:
\[
  {\bm{x}}_{t+1} \approx {\bm{\theta}}_{t+1} - {\bm{\phi}}_t \approx ({\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}_t)) {\bm{x}}_t.
\]
By simple linear algebra,
$\vv^{\mathrm{H}}_1({\bm{\phi}}_t)$
is an eigenvector of ${\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}_t)$,
associated with eigenvalue $1 - \tilde{\eta}_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) \approx -1$.
The remaining eigenvalues are $\{ 1 - \tilde{\eta}_t \lambda^{\mathrm{H}}_i({\bm{\phi}}_t) \}_{i=2}^{D}$,
where $\lambda^{\mathrm{H}}_i({\bm{\phi}}_t)$ is the $i$-th largest eigenvalue of ${\bm{H}}({\bm{\theta}}_t)$,
and they lie in the range $(-1, 1]$ since $\lambda^{\mathrm{H}}_i({\bm{\phi}}_t) \in [0, \lambda^{\mathrm{H}}_1({\bm{\phi}}_t))$.
Using a similar analysis to power method, we show that ${\bm{x}}_t$ quickly aligns to the direction of $\pm \vv^{\mathrm{H}}_1({\bm{\phi}}_t)$ after a few initial steps,
as the corresponding eigenvalue
has approximately the largest absolute value.\footnote{
  Our construction of ${\bm{\phi}}_t$
  ensures that ${\bm{x}}_t$ only
  has a small overlap with the $1$-eigenspace of ${\bm{I}} - \tilde{\eta}_t {\bm{H}}({\bm{\phi}}_t)$,
  so ${\bm{x}}_t$ can only align to $\pm \vv^{\mathrm{H}}_1({\bm{\phi}}_t)$.
}

To formally establish the above result, we need
a tiny initial alignment
between ${\bm{x}}_0$ and $\vv^{\mathrm{H}}_1({\bm{\phi}}_0)$,
just as the initial condition in power method.
This is where we need the initial random perturbation.

\begin{figure}[t]
  \centering
  \includegraphics[width=\textwidth]{fig/lin-idea.pdf}
  \caption{
    Illustration of the oscillation and periodic behaviors of GD+WD\xspace on linear regression with BN
    (see \Cref{sec:proof-idea,sec:lin}).
    The training loss decreases to $\approx 10^{-14}$ in the first 1k steps
    and achieves test loss $0.26$.
    Starting from step $\sim$ 1k, the dynamic enters the EoS regime.
    \textbf{(a).}~The test loss decreases to $0.16$ as a distance measure to the flattest solution \eqref{eq:M} decreases towards~$0$;
    \textbf{(b).}~The training loss oscillates around $\sim 10^{-4}$ in the EoS regime;
    \textbf{(c).}~$2/\tilde{\eta}_t$
    switches back and forth between being smaller and larger than $\lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$;
    \textbf{(d).}~The parameter oscillates around the minimizer manifold along the top eigenvector direction,
    and the magnitude of oscillation $\abssm{h_t}$ rises and falls periodically.
  }
  \label{fig:lin-idea}
  \vspace{-0.15in}
\end{figure}

\myparagraph{Oscillation Drives ${\bm{\phi}}_t$ to Move.}
This period-two oscillation is the driving power to push ${\bm{\phi}}_t$
to move on the manifold.
The main idea here is to
realize that the oscillation direction deviates slightly from the direction of $\pm \vv^{\mathrm{H}}_1({\bm{\phi}}_t)$
if we use a higher-order approximation.
We specifically use the Taylor approximation to show that
this deviation leads
${\bm{\phi}}_t$ to move slightly on $\mathit{\Gamma}$:
after each cycle of oscillation, ${\bm{\phi}}_{t+2} \approx {\bm{\phi}}_t - 4 h_t^2 \nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) + O(\eta_{\mathrm{in}}^{1.5-o(1)})$,
which resembles two steps of gradient descent on $\mathit{\Gamma}$
to minimize the logarithm of spherical sharpness with learning rate $2h_t^2$,

\myparagraph{Periodic Behavior of $h_t$ and $\tilde{\eta}_t$.}
It remains to analyze the dynamics of $h_t$
so that we can know how fast the sharpness reduction is.
Our analysis is inspired by an empirical study from~\citet{lobacheva2021on},
which reveals a periodic behavior of gradients and effective learning rates in training normalized nets
with weight decay.
In our theoretical setting, we capture
this periodic behavior by showing that $h_t$ and $\tilde{\eta}_t$ \textit{do} evolve periodically.
See Figures \hyperref[fig:lin-idea]{\ref*{fig:lin-idea}c} and \hyperref[fig:lin-idea]{\ref*{fig:lin-idea}d} for an example.

The key is that $\tilde{\eta}_t$ changes as an adaptive gradient method:
$\tilde{\eta}_t$ increases when gradient is small and decreases when gradient is large
(due to the effect of WD; see \Cref{fig:norm-smaller,fig:norm-bigger}),
and in our case the gradient norm scales as $\abssm{h_t}$ since $\nabla \mathcal{L}({\bm{\theta}}_t) \approx h_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t) \vv^{\mathrm{H}}_1({\bm{\phi}}_t)$.
According to our power method approximation, $h_{t+2} \approx (1 - \tilde{\eta}_t \lambda^{\mathrm{H}}_1({\bm{\phi}}_t))^2 h_t$,
so $\abssm{h_t}$ decreases when $\tilde{\eta}_t < 2 / \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$.
But $\abssm{h_t}$ cannot decrease forever, since $\tilde{\eta}_t$ increases when $\abssm{h_t}$ is sufficiently small.
When $\tilde{\eta}_t$ rises to over $2 / \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$,
$\abssm{h_t}$ changes from decreasing to increasing according to our approximation.
But $h_t$ cannot increase indefinitely either, since
$\tilde{\eta}_t$ decreases when $\abssm{h_t}$ is sufficiently large.
A period is finally finished when $\tilde{\eta}_t$ drops below $2 / \lambda^{\mathrm{H}}_1({\bm{\phi}}_t)$.

In our theoretical analysis, we connect this periodic behavior with a 1-dimensional Hamiltonian system (see \Cref{sec:outline-rmsdrift-def}),
and show that
$2 h_t^2$ in each step can be approximated by its average value in the period
without incurring a large error.
Further calculations show that this average value
is approximately
$\frac{2 \eta_{\mathrm{in}}}{4 + \normtwosm{\nabla_{\Gamma} \log \lambda^{\mathrm{H}}_1({\bm{\zeta}}(t\eta_{\mathrm{in}}))}}$,
the learning rate
in the flow~\eqref{eq:main-zeta} multiplied with $\eta_{\mathrm{in}}$.
We can therefore conclude that each step of ${\bm{\phi}}_t$ (or ${\bm{\theta}}_t$)
tracks a time interval of $\eta_{\mathrm{in}}$ in the flow.

\myparagraph{Extensions.} We note that this periodic behavior is not limited to GD+WD\xspace on scale-invariant loss,
as the above intuitive argument holds as long as
the effective LR changes adaptively with respect to gradient change.
Based on this intuition,
an important notion called \textit{Quasi-RMSprop scheduler}
is proposed.
For a PGD method, a learning rate scheduler is a rule for
changing the effective LRs in each step,
and Quasi-RMSprop is a specific class of schedulers we defined,
including the way that the effective LRs change
in GD+WD\xspace on scale-invariant loss (if viewed as PGD).
Our proof is done in a unified way that
works as long as the effective LRs are changed according to a Quasi-RMSprop scheduler.
As a by-product, a similar theorem can be proved for GD (without projection) on non-scale-invariant loss,
if a Quasi-RMSprop scheduler sets the LR in each step.
For example, we can extend our analysis to RMSprop with a scalar learning rate.
See \Cref{sec:qrms}.

\section{Details of the 3D Example} \label{sec:3d}

In this section, we give more details for \Cref{fig:ex3d}.

The loss $\mathcal{L}({\bm{w}})$ is constructed as follows.
First, we define the following scale-invariant function:
\begin{align*}
    F(x, y, z) := 2 - \frac{x+y}{\sqrt{x^2-xy+y^2}}.
\end{align*}
By taking gradient on $\mathbb{S}^{2}$, one can easily see that
the minimum is attained when $(x, y)$ points to $(1, 1)$ in direction.
In other words, the minimizer manifold of $F$
is $\mathit{\Gamma} := \{ (x, y, z) \in \mathbb{S}^2 : x = y > 0 \}$.

Then we fix an orthogonal matrix ${\bm{U}}$ (generated randomly)
and define $\mathcal{L}: \mathbb{R}^3 \setminus \{ {\bm{0}} \} \to \mathbb{R}, {\bm{w}} \mapsto F({\bm{U}}{\bm{w}})$,
i.e., the function $F$ after an orthogonal transformation.
For plotting the figure,
we transform the coordinates back to the domain of $F$.

The initial point is ${\bm{w}}_0 = (0.3, 1.3, 1.2)$ in the domain of $F$.
We run gradient descent on $\mathcal{L}$ with LR $\hat{\eta} = 0.5$
and WD $\hat{\lambda} = 0.08$.
It can be seen from the figure that
${\bm{\theta}}_t$ does not stop moving after reaching ${\bm{\zeta}}_0$.
The point that ${\bm{\theta}}_t$ eventually oscillate around is ${\bm{\zeta}}_* = (\frac{1}{\sqrt{2}}, \frac{1}{\sqrt{2}}, 0)$.

One can check that the Hessian matrix of $F$ at $(x, y, z) \in \mathit{\Gamma}$ is
\begin{align*}
H(x, y, z) = \frac{3}{x^2 + y^2}
    \begin{bmatrix}
    1 & -1 & 0 \\
    -1 & 1 & 0 \\
    0 & 0 & 0
\end{bmatrix} = \frac{3}{1 - z^2}
    \begin{bmatrix}
    1 & -1 & 0 \\
    -1 & 1 & 0 \\
    0 & 0 & 0
\end{bmatrix}.
\end{align*}
Therefore, the spherical sharpness is controlled by $\abssm{z}$.
The smaller the absolute value of $z$,
the flatter the minimizer.
And the flattest one is
${\bm{\zeta}}_*$,
which has $z$-coordinate being zero.
This matches with our theory of sharpness-reduction bias
as GD+WD\xspace moves along $\mathit{\Gamma}$ and oscillates near ${\bm{\zeta}}_*$ in the end.
\section{Reduction to RMS-drift Process: The Case of Spherical Optimization} \label{sec:sph-proof}

In this section, we let $\{({\bm{\theta}}_t, v_t)\}_{t \ge 0}$ be a trajectory of projected gradient descent
with quasi-RMSprop scheduler,
and let $\eta, \beta$ be the base learning rate and decay rate.
We follow the notations and terminologies in \Cref{sec:proof-outline-main}.

As the analysis in the spherical case is nearly the same
as the full space case,
we only discuss the difference here.

\subsection{Good Initialization} \label{sec:sph-good-init}
\begin{proof}[Proof for \Cref{lm:sph-good-init}]
Using a similar argument as in \Cref{lm:ful-good-init},
we know that the lemma holds if there is no projection
in the random initialization.
But the projection only leads to an error of order $O(r^2)$,
so the lemma holds.
\end{proof}

\subsection{Alignment Phase} \label{sec:sph-proof-alignment}

\begin{proof}[Proofs for \Cref{lm:sph-align-1-step,thm:sph-alignment-main}]
The proof is essentially the same as
\Cref{lm:ful-align-1-step,thm:ful-alignment-main}.
To see this,
we only need to note that
we only have used a linear approximation
of the update rule
with error $O(r_t^2)$,
and the linear approximation remains unchanged if we add a projection operator
(\Cref{lm:pgd-1-step}).
\end{proof}

\subsection{Drifting Phase}
We define ${\bm{\psi}}_{{\bm{\phi}}}({\bm{x}})$ as in \Cref{lm:pgd-2-step}. We abuse the notation to write ${\bm{\psi}}_t({\bm{x}}) = {\bm{\psi}}_{{\bm{\phi}}_t}({\bm{x}})$, that is,
\begin{equation} \label{eq:sph-psi}
    {\bm{\psi}}_t({\bm{x}}) :=
    \frac{\mu_t}{2}\left( {\bm{U}}_t \partial^3 \mathcal{L}_{{\bm{\phi}}_t}[{\bm{x}}, {\bm{x}}] + \partial^3 \mathcal{L}_{{\bm{\phi}}}[{\bm{U}}_t {\bm{x}}, {\bm{U}}_t {\bm{x}}] \right)
    + \frac{\mu_t^2}{2} \left( \normtwosm{{\bm{H}}_t {\bm{x}}}^2 + \normtwosm{{\bm{H}}_t {\bm{U}}_t {\bm{x}}}^2 \right) {\bm{\phi}}_t.
\end{equation}

\begin{lemma}
In the setting of \Cref{lm:sph-drift-2-step},
the same statement as \Cref{lm:ful-drift-2-step-hat,lm:ful-drift-2-step-hat2} holds,
where ${\bm{\psi}}_t$ is interpreted as \eqref{eq:sph-psi}.
\end{lemma}
\begin{proof}
  We can follow the argument in the proof for \Cref{lm:ful-drift-2-step-hat,lm:ful-drift-2-step-hat2},
  but now we are using \Cref{lm:pgd-2-step,lm:sph-psi-formula} to establish the proof.
\end{proof}

\begin{lemma}
In the setting of \Cref{lm:sph-drift-2-step},
  the same statement as \Cref{lm:ful-drift-2-step-phi} holds.
\end{lemma}
\begin{proof}
  The argument is the same as \Cref{lm:ful-drift-2-step-phi},
  but we apply \Cref{lm:Phi-proj-sph,lm:Phi-2nd-normal-sph} in doing Taylor expansion
  for $\Phi({\bm{\theta}})$.
\end{proof}

\begin{proof}[Proof for \Cref{lm:sph-drift-2-step}]
Same as \Cref{lm:ful-drift-2-step}
but we invoke the spherical version of lemmas.
\end{proof}

\section{Lemmas for Working Zones}
\label{sec:working-zone-proof}

\subsection{Construction of Working Zones} \label{sec:working-zone-proof-construction}

\begin{proof}[Proof for \Cref{lm:ful-working-zone}]
By \Cref{ass:man-ful} and \Cref{thm:rank-to-pl-ful},
there exists $\epsilon_1$ such that Items 1 and 2 hold.
By \Cref{thm:pl-converge-near-ful},
we can choose $\epsilon_0' > 0$ small enough so that
any gradient flow starting in ${\mathcal{Z}}^{\epsilon_0'}$
moves at most $O(\epsilon_0')$ in distance
and converges in $\gZ^{\epsilon_1}$.
We can further combine this with
the results by \citet{falconer1983differentiation}
to show that $\Phi$ is $\mathcal{C}^{3}$-smooth in a neighborhood of ${\mathcal{Z}}$.
Then we can choose $\epsilon_0 > 0$ small enough
so that $\epsilon_0 < \epsilon_0'$ and $\gZ^{\epsilon_0}$ is a subset of that neighborhood,
which ensures Items 3 and 4.
Finally, Item 5 is directly implied by the compactness of $\mathrm{cl}(\gZ^{\epsilon_1}) \cap \mathit{\Gamma}$.
\end{proof}

\begin{proof}[Proof for \Cref{lm:sph-working-zone}]
We can use a similar argument as above but with
\Cref{thm:rank-to-pl-sph,thm:pl-converge-near-sph}.
\end{proof}

\subsection{Gradient Flow Projection}

In the working zone, $\Phi({\bm{\theta}})$ is well-defined and $\mathcal{C}^3$-smooth.
Now we highlight some useful properties.

\subsubsection{Full Space Optimization}

The following two lemmas are from \citet{li2022what}.
\begin{lemma} \label{lm:Phi-proj-ful}
  For ${\bm{\phi}} \in \gZ^{\epsilon_0} \cap \manGa$, $\partial \Phi_{{\bm{\phi}}}[{\bm{x}}] = \mPHz({\bm{\phi}}) {\bm{x}}$,
  which also
  equals to the projection of ${\bm{x}}$ onto the tangent space $\TGa{{\bm{\phi}}}$ at ${\bm{\phi}}$.
\end{lemma}

\begin{lemma} \label{lm:Phi-2nd-normal-ful}
  For ${\bm{\phi}} \in \gZ^{\epsilon_0} \cap \manGa$ and ${\bm{x}} \in \NGa{{\bm{\phi}}}$, $\partial^2 \Phi_{{\bm{\phi}}}[{\bm{x}}, {\bm{x}}] = {\bm{0}}$.
\end{lemma}

The following lemma can be proved by Taylor expansion.
\begin{lemma} \label{lm:phiphix-ful}
  For ${\bm{\theta}} \in \gZ^{\epsilon_0}$ and ${\bm{\phi}} = \Phi({\bm{\theta}})$, then $\normtwosm{\partial \Phi_{{\bm{\phi}}}[{\bm{\theta}} - {\bm{\phi}}]} \le O(\normtwosm{{\bm{\theta}} - {\bm{\phi}}}^2)$.
\end{lemma}
\begin{proof}
  When $\normtwosm{{\bm{\theta}} - {\bm{\phi}}}$ is small enough,
  the linear interpolation of ${\bm{\theta}}$ and ${\bm{\phi}}$ lies in $\gZ^{\epsilon_1}$.
  By Taylor expansion,
  $\Phi({\bm{\theta}}) = \Phi({\bm{\phi}}) + \partial \Phi_{{\bm{\phi}}}[{\bm{\theta}} - {\bm{\phi}}] + O(\normtwosm{{\bm{\theta}} - {\bm{\phi}}}^2)$.
  In fact, $\Phi({\bm{\theta}}) = {\bm{\phi}} = \Phi({\bm{\phi}})$. So we can conclude $\partial \Phi_{{\bm{\phi}}}[{\bm{\theta}} - {\bm{\phi}}] = O(\normtwosm{{\bm{\theta}} - {\bm{\phi}}}^2)$.
\end{proof}

\subsubsection{Spherical Optimization}

In the spherical case,
we have the following lemmas that highly resemble those in the full space case.
All these lemmas can be proved in the same manner:
we can first define $\mathit{\Gamma}' := \{ \nu {\bm{\theta}} : {\bm{\theta}} \in \mathit{\Gamma}, \nu \in (1/2, 2)\}$
and apply the counterparts of these lemmas in the full space case;
then translate the results back to the spherical case.
\begin{lemma} \label{lm:Phi-proj-sph}
  For ${\bm{\phi}} \in \gZ^{\epsilon_0} \cap \manGa$, $({\bm{I}} - {\bm{\phi}}\vphi^\top)\partial \Phi_{{\bm{\phi}}}[{\bm{x}}] = ({\bm{I}} - {\bm{\phi}}\vphi^\top) \mPHz {\bm{x}}$,
  which also equals to the projection of ${\bm{x}}$ onto the tangent space $\TGa{{\bm{\phi}}}$ at ${\bm{\phi}}$.
\end{lemma}

\begin{lemma} \label{lm:Phi-2nd-normal-sph}
  For ${\bm{\phi}} \in \gZ^{\epsilon_0} \cap \manGa$ and ${\bm{x}} \in \NGa{{\bm{\phi}}}$, if $\dotpsm{{\bm{x}}}{{\bm{\phi}}} = 0$, then $\partial^2 \Phi_{{\bm{\phi}}}[{\bm{x}}, {\bm{x}}] = {\bm{0}}$.
\end{lemma}

\begin{lemma} \label{lm:phiphix-sph}
  For ${\bm{\theta}} \in \gZ^{\epsilon_0} \cap \sphS^{D-1}$ and ${\bm{\phi}} = \Phi({\bm{\theta}})$, then $\normtwosm{\partial \Phi_{{\bm{\phi}}}[{\bm{\theta}} - {\bm{\phi}}]} \le O(\normtwosm{{\bm{\theta}} - {\bm{\phi}}}^2)$.
\end{lemma}

Machine Learning (cs.LG); Artificial Intelligence (cs.AI)