Anti-Correlated Noise in Epoch-Based Stochastic Gradient Descent: Implications for Weight Variances

## Abstract

Abstract Stochastic gradient descent (SGD) has become a cornerstone of neural network optimization, yet the noise introduced by SGD is often assumed to be uncorrelated over time, despite the ubiquity of epoch-based training. In this work, we challenge this assumption and investigate the effects of epoch-based noise correlations on the stationary distribution of discrete-time SGD with momentum, limited to a quadratic loss. Our main contributions are twofold: first, we calculate the exact autocorrelation of the noise for training in epochs under the assumption that the noise is independent of small fluctuations in the weight vector, and find that SGD noise is anti-correlated in time. Second, we explore the influence of these anti-correlations on SGD dynamics. We find that for directions with a curvature greater than a hyperparameter-dependent crossover value, the results for uncorrelated noise are recovered. However, for relatively flat directions, the weight variance is significantly reduced, and our variance prediction leads to a considerable reduction in loss fluctuations as compared to the constant weight variance assumption.

## 1 Introduction

Initially developed to address the challenges of computational efficiency in neural networks, stochastic gradient descent (SGD) has exhibited exceptional effectiveness in managing large datasets compared to the full-gradient methods []. It has since garnered widespread acclaim in the machine learning domain [], with applications spanning image recognition [, , ], natural language processing [, ], and mastering complex games beyond human capabilities []. Alongside its numerous variants [, , ], SGD remains the cornerstone of neural network optimization.

SGD’s success can be attributed to several key properties, such as rapid escape from saddle points [] and its capacity to circumvent "bad" local minima, instead locating broad minima that generally lead to superior generalization [, , , , , ]. This is often ascribed to anisotropic gradient noise [, , , , , , ]. Nonetheless, recent empirical research posits that even full gradient descent, with minor adjustments, can achieve generalization performance comparable to that of SGD [].

To deepen our understanding, multiple studies have investigated the limiting dynamics of neural network weights during the latter stages of training [, ]. Of particular interest is the behavior of weight fluctuations in proximity to a minimum of the loss function [, , ]. Several authors found empirical evidence that the covariance matrix $\bf C$ of SGD is proportional to the Hessian matrix $\bf H$ of the loss function [, , , , , ]. Consequently, theory posits that the stationary covariance matrix of weights $\bf\Sigma$ exhibits isotropy for sufficiently small learning rates [, , ]. Nevertheless, a recent empirical investigation [] identified profound anisotropy in $\bf\Sigma$.

In this work, we delve into both theoretical and empirical analyses of weight fluctuations during the later stages of training, accounting for the emergence of anti-correlations in the noise produced by SGD, which stem from the prevalent epoch-based learning schedule. As a result of these anti-correlations, we discover that the covariance matrix $\bf\Sigma$ displays anisotropy in a subspace of weight directions corresponding to Hessian eigenvectors with small eigenvalues, while maintaining the isotropy of $\bf\Sigma$ in directions associated with Hessian eigenvectors possessing large eigenvalues.

### Our contributions

We make the following contributions:

- •
We uncover generic anti-correlations in SGD noise, resulting from the common practice of drawing examples without replacement. These anti-correlations lead to a smaller-than-anticipated weight variance in some directions.
- •
We calculate the autocorrelation function of SGD noise, assuming that the noise is independent of small fluctuations in the weight vector over time. Given that the noise sums to zero over an epoch where each example is shown once, we find anti-correlations of SGD noise.
- •
Building on the computed autocorrelation function, we develop a theory that elucidates the relationships between variances for different Hessian eigenvectors. The weight space can be partitioned into two groups of Hessian eigenvectors, with the well-known isotropic variance prediction holding for eigenvectors whose Hessian eigenvalues exceed a hyperparameter-dependent crossover value $\lambda_{\textrm{cross}}$. For eigenvectors with smaller eigenvalues, the variance is proportional to the Hessian eigenvalue, smaller than the isotropic value. This theory accounts for SGD’s discrete nature, avoiding continuous-time approximations and incorporating heavy-ball momentum.
- •
For a given Hessian eigenvector, an intrinsic correlation time of update steps exists, introduced by the loss landscape and proportional to $1/\lambda_{i}$, where $\lambda_{i}$ represents the corresponding Hessian eigenvalue. Additionally, a fixed noise correlation time $\tau_{\textrm{SGD}}$ is present, on the order of one epoch and direction-independent. The smaller of these two timescales governs the dynamic behavior for a given Hessian eigenvector, determining the actual correlation time $\tau_{i}$ of the update steps. As weight variances are proportional to $\tau_{i}$, the two distinct variance-curvature relationships arise.
- •
Our theoretical predictions are validated through the analysis of a neural network’s training within a subspace of its top Hessian eigenvectors. Utilizing Hessian eigenvectors proves advantageous over the principal component analysis of the weight trajectory used in a prior empirical study [], where finite size effects induced significant artifacts.
- •
We investigate the effect of weight variance on training error in neural networks, focusing on Hessian eigenvectors and weight fluctuations. Our analysis, based on a quadratic loss model with Gaussian weight fluctuations, demonstrates that
the reduction of weight fluctuations due to noise anti-correlations leads to a significant reduction of loss fluctuations. In a case study with the LeNet architecture, analyzing the 5,000 largest Hessian eigenvalues, we observe a 62% reduction in loss fluctuation compared to models assuming constant weight variance. This highlights the importance of considering a broader range of Hessian eigendirections in understanding and optimizing neural network training dynamics.

## 2 Background

We consider a neural network characterized by its weight vector, $\bm{\mathbf{\theta}}\in\mathbb{R}^{d}$. The network is trained on a set of $N$ training examples, each denoted by $x_{n}$, where $n$ ranges from 1 to $N$. The loss function, defined as ${L(\bm{\mathbf{\theta}})\coloneqq\frac{1}{N}\sum_{i=1}^{N}l(\bm{\mathbf{\theta
}},x_{n})}$, represents the average of individual losses incurred for each training example $l(\bm{\mathbf{\theta}},x_{n})$.

It is common practice to add some kind of momentum to the SGD algorithm when training a network as it leads to faster convergence, has a smoothing effect, and helps in escaping local minima. Therefore, to keep the analysis general, we consider a training process that employs stochastic gradient descent augmented with Heavy-ball momentum. This approach updates the network parameters according to the following rules:

$$ $\displaystyle\bm{\mathbf{g}}_{k}(\bm{\mathbf{\theta}})$ $\displaystyle=\frac{1}{S}\sum_{n\in\mathcal{B}_{k}}\bm{\mathbf{\nabla}}l(\bm{ \mathbf{\theta}},x_{n})\ ,$ (1a) $\displaystyle\bm{\mathbf{v}}_{k}$ $\displaystyle=-\eta\bm{\mathbf{g}}_{k}(\bm{\mathbf{\theta}}_{k-1})+\beta\bm{ \mathbf{v}}_{k-1}\ ,$ (1b) $\displaystyle\bm{\mathbf{\theta}}_{k}$ $\displaystyle=\bm{\mathbf{\theta}}_{k-1}+\bm{\mathbf{v}}_{k}\ .$ (1c) $$

Here, $k$ signifies the discrete update step index, $\eta$ is the learning rate, and $\beta$ is the momentum parameter. The stochastic gradient at each step is computed with respect to a batch of $S\ll N$ random examples. Each batch is denoted by $\mathcal{B}_{k}=\{n_{1},...,n_{S}\}$, where $n_{j}\in\{1,\dots,N\}$.
The training process is structured into epochs. During each epoch, every training example is used exactly once, implying that the examples are drawn without replacement and do not recur within the same epoch.

In the realm of SGD as opposed to full gradient descent, we introduce noise, denoted as $\vec{\delta g}_{k}(\vec{\theta})\coloneqq\vec{g}_{k}(\vec{\theta})-\vec{\nabla
}L(\vec{\theta})$, with a covariance matrix $\bm{\mathbf{C}}(\vec{\theta})\coloneqq\textrm{cov}\bigl{(}\vec{\delta g}_{k}(
\vec{\theta}),\vec{\delta g}_{k}(\vec{\theta})\bigr{)}$. The noise covariance matrix is known to be proportional to the gradient sample covariance matrix $\vec{C_{0}}(\vec{\theta})\coloneqq\frac{1}{N-1}\sum_{n=1}^{N}\vec{\nabla}\left
[l(\bm{\mathbf{\theta}},x_{n})-L(\vec{\theta})\right]\vec{\nabla}^{\top}\left[
l(\bm{\mathbf{\theta}},x_{n})-L(\vec{\theta})\right]$. A detailed derivation of this relation, particularly for the case of drawing without replacement, is provided in Appendix E.

Our primary focus lies on the asymptotic or limiting covariance matrix of the weights $\bm{\mathbf{\Sigma}}\coloneqq\textrm{cov}\bigl{(}\vec{\theta}_{k},\vec{\theta}
_{k}\bigr{)}$. That means we are interested in the covariance computed over an infinite run of SGD optimization, and not the covariance for a fixed update step $k$ computed over multiple runs of SGD optimization (see Appendix A).

To better understand and explain the behavior of the weight variances, we further examine the covariance matrix of the velocities $\bm{\mathbf{\Sigma_{v}}}\coloneqq\textrm{cov}\bigl{(}\vec{v}_{k},\vec{v}_{k}
\bigr{)}$, hypothesize that these two matrices commute, and proceed to explore their ratio $\bm{\mathbf{\Sigma}}/\bm{\mathbf{\Sigma_{v}}}$. The eigenvalues of this matrix ratio are denoted as $\tau_{i}$ because, under general assumptions, they equate to the velocity correlation time of the corresponding eigendirection (see Section 4.2).

## 3 Related Work

### 3.1 Hessian Noise Approximation

The equivalence between the gradient sample covariance $\vec{C_{0}}$ and the Hessian matrix of the loss $\vec{H}$ is an approximation that frequently appears in the literature [, , ]. Numerous theoretical arguments have indicated that when the output of a neural network closely matches the example labels, these two matrices should be similar [, , , , ]. However, even a slight deviation between network predictions and labels can theoretically disrupt this relationship, as highlighted by Thomas et al. []. Despite this, empirical observations suggest a strong alignment between the gradient sample covariance and the Hessian matrix near a minimum.

Zhang et al. [] pursued this line of thought and evaluated the assumption numerically. They examined both matrices in the context of a particular basis that presents both high and low curvature across different directions. Their findings indicated a close match between curvature and gradient variance in a given direction for a convolutional image recognition network, and a reasonably good relationship for a transformer model.

Meanwhile, Thomas et al. [] presented both a theoretical argument and empirical evidence across different architectures for image recognition. Although they did not discover an exact match between the two matrices, they did observe a proportionality between them, indicated by a high cosine similarity.

Xie et al. [] also conducted an investigation with an image recognition network. In the eigenspace of the Hessian matrix, they plotted entries within a specific interval against corresponding entries from the gradient sample covariance and found a close match.

For our theoretical considerations, it is crucial to assume that both matrices commute, $\left[\bm{\mathbf{C_{0}}},\bm{\mathbf{H}}\right]=0$. Additionally, we hypothesize that $\bm{\mathbf{C_{0}}}\propto\bm{\mathbf{H}}$, to derive power law predictions for the weight variance.

### 3.2 Limiting Dynamics and Weight Fluctuations

Various studies have scrutinized the limiting dynamics, often modeling SGD as a stochastic differential equation (SDE). Commonly, researchers such as Mandt et al. [] and Jastrzębski et al. [] approximate the loss near a minimum as a quadratic function and present the SDE as a multivariate Ornstein-Uhlenbeck (OU) process. This process proposes a stationary weight distribution with Gaussian weight fluctuations. Jastrzębski et al. [] further assume the Hessian noise approximation, observing under these conditions that the weight fluctuations are isotropic. Kunin et al. [] who also incorporate momentum into their analysis, predict and empirically verify isotropic weight fluctuations. Chaudhari & Soatto [] also investigate the SDE but without the assumption of a quadratic loss, nor that it reached equilibrium. They gain insights via the Fokker-Plank equation.

Alternatively, some studies derive relations from a stationarity assumption instead of a continuous time approximation [, , ]. Yaida [] assumes that the weight trajectory follows a stationary distribution and derives general fluctuation-dissipation relations from that. Liu et al. [] go further to assume a quadratic loss function, leading them to derive exact relations for the weight variance of SGD with momentum. If the additional Hessian noise approximation is made, their results also predict the weight variance to be approximately isotropic, except in directions where the product of learning rate and Hessian eigenvalue is significantly high.

Such computed weight variances are explicitly applied in various contexts, such as computing the escape rate from a minimum or assessing the approximation error in SGD, which captures the additional training error attributed to noise [].

Feng & Tu [] present a phenomenological theory based on their empirical findings, which also account for flat directions, unlike Kunin et al. []. They describe a general inverse variance-flatness relation, analyzing the weight trajectory of different image recognition networks via principal component analysis. They discovered a power law relationship between the curvature of the loss and the weight variance $\sigma_{\theta,i}^{2}$ in any given direction, where a higher curvature corresponds to a higher variance. They also observed that both the velocity variance $\sigma_{v,i}^{2}$ and the correlation time $\tau_{i}$ are larger for higher curvatures.

In our approach, we do not make the continuous time approximation but base our results on the assumption that the weights adhere to a stationary distribution near a quadratic minimum.

## 4 Theory

### 4.1 Autocorrelation of the noise

Figure: Figure 1: Autocorrelations of the SGD noise observed over a span of 20 epochs, equivalent to 20,000 update steps. This data is collected from a later phase in the training process. The autocorrelation is projected onto 5,000 Hessian eigenvectors, and the result is averaged. The theoretical prediction Equation 2 is also displayed along with a $2\sigma$-interval, where $\sigma$ represents the expected standard deviation of the SGD noise. The zero-point correlation is omitted as it is inherently equal to one.
Refer to caption: extracted/5732104/Figs/top_eigvecs/autocorrelation.jpeg

Autocorrelation of the noise means we are considering the correlation between two noise terms, but they stem from different update steps. In the case of SGD, when we sample the examples without replacement while keeping the weight vector $\bm{\mathbf{\theta}}$ constant, there are inherent anti-correlations in the noise. Here, we provide a brief conceptual motivation followed by the exact result where we direct the reader to Appendix E for a comprehensive derivation.

Suppose the ratio $M\coloneqq N/S$ is an integer, where $N$ represents the total number of examples, and $S$ is the batch size. This condition allows for the even distribution of examples among batches, ensuring each example is presented exactly once per epoch. Given a fixed weight vector $\bm{\mathbf{\theta}}$ and computing every gradient $\bm{\mathbf{g}}_{k}(\vec{\theta})$ for one epoch, it is evident that their mean equates to the total gradient $\bm{\mathbf{\nabla}}L(\bm{\mathbf{\theta}})$. Consequently, the sum of the noise components $\bm{\mathbf{\delta g}}_{k}(\bm{\mathbf{\theta}})=\vec{g}_{k}(\vec{\theta})-
\vec{\nabla}L(\vec{\theta})$ for one epoch must be zero. This condition implies that two noise components within the same epoch are inherently anti-correlated.

This anti-correlation between two noise terms of the same epoch can therefore be expressed in terms of the the equal-time correlation, multiplied by a factor of $\frac{1}{M-1}$ with $M$ being the number of batches per epoch. Next, we consider the probability that two batches $k$ and $k+h$ with an update step difference $h$ are from the same training epoch. This probability is given by the factor $\frac{M-|h|}{M}$ for $|h|\leq M$, and zero otherwise. By combining with the above anti-correlation, we arrive at the correlation formula:

###### Theorem 4.1 .

If the total number of examples $N$ is an integer multiple of the batch size $S$ and the parameters $\bm{\mathbf{\theta}}$ of a network are kept fixed, then the autocorrelation formula for the gradient noise of an epoch-based learning schedule, where the examples for one epoch are drawn without replacement, is given by

$$ $\displaystyle{\rm cov}\bigl{(}\vec{\delta g}_{k}(\vec{\theta}),\vec{\delta g}_ {k+h}(\vec{\theta})\bigr{)}=$ $\displaystyle\bm{\mathbf{C}}(\vec{\theta})\cdot\left(\delta_{h,0}-\mathbf{1}_{ \{1,...,M\}}(|h|)\frac{M-|h|}{M(M-1)}\right)\ ,$ (2) $$

where $M\coloneqq N/S$ signifies the number of batches per epoch.

In the above theorem $\mathbf{1}_{A}(k)$ represents the indicator function over the set A, which is one for $k\in A$ and zero otherwise and $\delta_{i,j}$ represents the Kronecker delta. The actual noise autocorrelation is illustrated in Figure 1, with the experimental details elaborated in Section 5.2. The complete calculation is available in Appendix E, highlighting the relationship $\bm{\mathbf{C}}(\vec{\theta})=(1/S)\cdot(1-S/N)\cdot\bm{\mathbf{C_{0}}}(\vec{
\theta})$ with the gradient sample covariance matrix $\bm{\mathbf{C_{0}}}$.

It is important to note that the above formula is only applicable for a static weight vector. During training, weights alter with each update step, which could potentially modify this relationship. Nevertheless, if the weights remain relatively consistent over one epoch, this correlation approximately holds true. In a simple quadratic model one can also construct a noise which would be independent of the weight vector (see Appendix B), for which then this correlation holds true without restrictions on the weights. Either way, for the later stages of training, such a noise constancy assumption seems to be a good approximation, as evidenced by the close fit of the data to the theory in Figure 1.

Further numerical investigations show, that even in the case where $N$ is not an integer multiple of $S$, the later deduced relations for the variances still hold (see Appendix J).

###### Theorem 4.1 .

### 4.2 Correlation time definition

To gain a more comprehensive understanding of the weight variance behavior, we additionally scrutinize the velocity variance and the ratio between them. We label this ratio, scaled by a factor of two, as the correlation time $\tau_{i}$ (see Equation 3). This definition aligns with that of the velocity correlation time, hence justifying the nomenclature. The equivalence stands under general assumptions that are satisfied in our problem setup described in Section 4.3 (see Section C.4).

The assumptions encompass: (i) Existence and finiteness of $\bm{\mathbf{\Sigma}}\coloneqq$ cov$(\bm{\mathbf{\theta}}_{k},\bm{\mathbf{\theta}}_{k})$, $\bm{\mathbf{\Sigma_{v}}}\coloneqq$ cov$(\bm{\mathbf{v}}_{k},\bm{\mathbf{v}}_{k})$, and $\langle\bm{\mathbf{\theta}}\rangle$. (ii) $\lim_{n\to\infty}\textrm{cov}\bigr{(}\bm{\mathbf{\theta}}_{k},\bm{\mathbf{
\theta}}_{k+n}\bigr{)}=0$. (iii) $\lim_{n\to\infty}n\cdot\textrm{cov}\bigr{(}\bm{\mathbf{\theta}}_{k},\bm{
\mathbf{\theta}}_{k+n}-\bm{\mathbf{\theta}}_{k+n+1}\bigr{)}=0$. For example, the latter two assumptions hold true if the weight correlation function decays as $\textrm{cov}\bigr{(}\bm{\mathbf{\theta}}_{k},\bm{\mathbf{\theta}}_{k+n}\bigr{)
}\propto n^{-2}$ or faster.

###### Theorem 4.2 .

Under the mentioned assumptions, it holds that

$$ $\displaystyle\tau_{i}\coloneqq\frac{2\sigma_{\theta,i}^{2}}{\sigma_{v,i}^{2}}= \frac{\sum_{n=1}^{\infty}n\cdot{\rm cov}\bigl{(}v_{k,i},v_{k+n,i}\bigr{)}}{ \sum_{n=1}^{\infty}{\rm cov}\bigl{(}v_{k,i},v_{k+n,i}\bigr{)}}$ (3) $$

justifying the label correlation time for this variance ratio(^1^11In a continuous time setting, the definition would necessitate $\tau_{i}\coloneqq\frac{2\sigma_{\theta,i}^{2}}{\sigma_{v,i}^{2}t_{0}}$, involving a microscopic timescale $t_{0}$ to assign $\tau_{i}$ as a timescale.).

Here $\sigma_{\theta,i}^{2}$ and $\sigma_{v,i}^{2}$ are eigenvalues of $\bm{\mathbf{\Sigma}}$ and $\bm{\mathbf{\Sigma_{v}}}$ respectively for a shared eigenvector(^2^22The relation still holds in the case where $[\bm{\mathbf{\Sigma}},\bm{\mathbf{\Sigma_{v}}}]\neq 0$ and $\bm{\mathbf{p}}_{i}$ is not a shared eigenvector but just any vector onto which we project the dynamics. In this case, $\sigma_{\theta,i}^{2}$ and $\sigma_{v,i}^{2}$ would just be the variances of the weight and the velocity in the given direction but no longer necessarily eigenvalues of $\bm{\mathbf{\Sigma}}$ and $\bm{\mathbf{\Sigma_{v}}}$.) $\bm{\mathbf{p_{i}}}$, with $\theta_{k,i}\coloneqq\vec{\theta}_{k}\cdot\vec{p_{i}}$ and $v_{k,i}\coloneqq\vec{v}_{k}\cdot\vec{p_{i}}$. The derivation is presented in Appendix D.

###### Theorem 4.2 .

### 4.3 Variance for late training phase

In light of the autocorrelation of the noise calculated earlier, our next objective is to compute the expected weight and velocity variances at a later stage of the training. To describe the conditions of this phase, we adopt the following assumptions.

##### Assumption 1: Quadratic Approximation

We postulate that we have reached a minimum point of the loss function, which can be adequately represented with a quadratic form as $L(\vec{\theta})=L_{0}+\frac{1}{2}(\vec{\theta}-\vec{\theta_{*}})^{\top}\bm{
\mathbf{H}}(\vec{\theta}-\vec{\theta_{*}})$. We can set $L_{0}$ and $\vec{\theta_{*}}$ to zero without any loss of generality, which simplifies to:

$$ $\displaystyle L(\vec{\theta})=\frac{1}{2}\vec{\theta}^{\top}\bm{\mathbf{H}} \vec{\theta}\ .$ (4) $$

##### Assumption 2: Spatially Independent Noise

We presume that the noise terms are independent of the weight vector in this late phase, represented as $\vec{\delta g}_{k}(\vec{\theta})\equiv\vec{\delta g}_{k}$.

This assumption leads to the conclusion that the noise covariance matrix is also spatially independent, $\bm{\mathbf{C}}(\vec{\theta})\equiv\bm{\mathbf{C}}$, and that the actual autocorrelation of the SGD noise now follows the relation previously calculated in Equation 2. For further motivation we refer to Appendix B.

##### Assumption 3: Hessian Noise Approximation

We assume that the covariance of the noise commutes with the Hessian matrix, as discussed in Section 3.1,

$$ $\displaystyle\left[\bm{\mathbf{C}},\bm{\mathbf{H}}\right]=0\ .$ (5) $$

Moreover, we assume that $0\leq\beta<1$ and $0<\eta\lambda_{i}<2(1+\beta)$ for all eigenvalues $\lambda_{i}$ of $\vec{H}$. If these conditions are not met, the weight fluctuations would diverge.

##### Calculation for one eigenvalue

With the previously stated assumptions in place, the covariance matrices $\bm{\mathbf{\Sigma}}$ and $\bm{\mathbf{\Sigma_{v}}}$ commute with $\bm{\mathbf{C}}$, $\bm{\mathbf{H}}$, and with each other (see Section C.1). As a result, they all share a common eigenbasis(^3^33The relations calculated in the following paragraphs still hold in the case $\left[\bm{\mathbf{C}},\bm{\mathbf{H}}\right]\neq 0$ and where $\vec{p_{i}}$ are just the eigenvectors of the Hessian matrix. The calculated weight and velocity variances, however, are no longer eigenvalues of the corresponding covariance matrices. In Appendix M is shown what happens when $\bm{\mathbf{C}}$ and $\bm{\mathbf{H}}$ do not commute.) $\vec{p_{i}}$, with $i=1,...,d$, which facilitates the computation of the expected variance. We will outline the most important steps here, while details can be found in Section C.2.

For a given common eigenvector $\vec{p_{i}}$ we project the relevant variables onto this vector. This yields the projected weight $\theta_{k,i}\coloneqq\vec{p_{i}}\cdot\vec{\theta}_{k}$, velocity $v_{k,i}\coloneqq\vec{p_{i}}\cdot\vec{v}_{k}$, and noise term $\delta g_{k,i}\coloneqq\vec{p_{i}}\cdot\vec{\delta g}_{k}$ at the update step $k$. Correspondingly, we define the eigenvalues for the common eigenvector $\bm{\mathbf{p_{i}}}$ as $\lambda_{i}$ for $\vec{H}$, $\sigma_{\theta,i}^{2}$ for $\vec{\Sigma}$, $\sigma_{v,i}^{2}$ for $\vec{\Sigma_{v}}$, and $\sigma_{\delta g,i}^{2}$ for $\vec{C}$. We denote the number of batches per epoch as $M=N/S$, presuming it is an integer.

By introducing the vector $\bm{\mathbf{x}}_{k,i}\coloneqq\bigl{(}\theta_{k,i}\quad\theta_{k-1,i}\bigr{)}^
{\top}$ that contains not only the current weight variable but also the weight variable with a one-step time lag we can write the update equation as:

$$ $\displaystyle\bm{\mathbf{x}}_{k,i}$ $\displaystyle=\vec{D_{i}}\vec{x}_{k-1,i}-\eta\delta g_{k,i}\vec{e}_{1}$ (6) $$

where $\bm{\mathbf{e}}_{1}\coloneqq\bigl{(}1\quad 0\bigr{)}^{\top}$ and the matrix $\vec{D_{i}}$ governs the deterministic part of the update, its explicit expression being:

$$ $\displaystyle\vec{D_{i}}$ $\displaystyle\coloneqq\begin{pmatrix}1+\beta-\eta\lambda_{i}&-\beta\\ 1&0\end{pmatrix}\ .$ (7) $$

Further, we define

$$ $\displaystyle\vec{E_{i}}\coloneqq\vec{D_{i}}\,{\rm cov}\left(\bm{\mathbf{x}}_{ k-1,i},\delta g_{k,i}\vec{e}_{1}\right)\ .$ (8) $$

This term encapsulates the correlation between the current weight variable and the noise term of the next update step. Typically, noise terms are assumed to be temporally uncorrelated, which would render $\vec{E_{i}}$ null. However, given the anti-correlated nature of the noise, we find a non-zero $\vec{E_{i}}$. An explicit expression of $\vec{E_{i}}$ can be found in Section C.2.

###### Theorem 4.3 .

With the above assumptions and definitions the following relation for the weight and velocity variances hold:

$$ $\displaystyle\begin{pmatrix}\sigma_{\theta,i}^{2}\\ \sigma_{v,i}^{2}\end{pmatrix}$ $\displaystyle=\eta^{2}\sigma_{\delta g,i}^{2}\vec{F_{i}}\left[\vec{e}_{1}- \left(\bm{\mathbf{E_{i}}}+\vec{E_{i}}^{\top}\right)\vec{e}_{1}\right]\ ,$ (9) $$

where the matrix $\bm{\mathbf{F_{i}}}$ is explicitly expressed as:

$$ $\displaystyle\bm{\mathbf{F_{i}}}$ $\displaystyle=\frac{1}{(1-\beta)\bigl{(}2(1+\beta)-\eta\lambda_{i})\bigr{)}} \begin{pmatrix}\frac{1+\beta}{\eta\lambda_{i}}&\frac{2\beta(\eta\lambda_{i}-1- \beta)}{\eta\lambda_{i}}\\ 2&2(\eta\lambda_{i}-2)\end{pmatrix}\,.$ (10) $$

The calculations can be found in Section C.2. The exact relation Equation 9 can be easily evaluated numerically but it can also be approximated by assuming that $M\gg 1/(1-\beta)$, which implies that the correlation time induced by momentum is substantially shorter than one epoch. Consequently, two distinct regimes of Hessian eigenvalues emerge, separated by $\lambda_{\textrm{cross}}\coloneqq 3(1-\beta)/(\eta M)$. For each of these regimes, specific simplifications apply. Notably, at $\lambda_{\textrm{cross}}$, both approximations converge.

###### Corollary 4.4 .

Relations for large Hessian eigenvalues: For Hessian eigenvectors with eigenvalues $\lambda_{i}>\lambda_{\textrm{cross}}$ and when $M\gg 1/(1-\beta)$, the effects of noise anti-correlations are minimal. Consequently, we can use the following approximate relationships, which also hold true in the absence of correlations:

$$ $\displaystyle\sigma_{\theta,i}^{2}$ $\displaystyle\approx\frac{\eta^{2}\sigma_{\delta g,i}^{2}}{(1-\beta)\bigl{(}2( 1+\beta)-\eta\lambda_{i}\bigr{)}}\cdot\frac{1+\beta}{\eta\lambda_{i}}\ ,$ (11a) $\displaystyle\sigma_{v,i}^{2}$ $\displaystyle\approx\frac{\eta^{2}\sigma_{\delta g,i}^{2}}{(1-\beta)\bigl{(}2( 1+\beta)-\eta\lambda_{i}\bigr{)}}\cdot 2\ ,$ (11b) $\displaystyle\tau_{i}$ $\displaystyle\approx\frac{1+\beta}{\eta\lambda_{i}}\ .$ (11c) $$

The detailed derivation of these formulas is presented in Section C.3.

It is frequently the case that the product $\eta\lambda_{i}$ is considerably less than one, enabling us to further simplify the prefactor of the variances. Assuming the noise covariance matrix is proportional to the Hessian matrix, such that $\sigma_{\delta g,i}^{2}\propto\lambda_{i}$, we derive the following power laws for the variances: $\sigma_{\theta,i}^{2}\propto\text{const.}$ and $\sigma_{v,i}^{2}\propto\lambda_{i}$. As a result, in the subspace spanned by Hessian eigenvectors with eigenvalues $\lambda_{i}>\lambda_{\textrm{cross}}$, our theory predicts an isotropic weight variance $\bm{\mathbf{\Sigma}}$.

###### Corollary 4.5 .

Relations for small Hessian eigenvalues: In the case of Hessian eigenvectors associated with eigenvalues $\lambda_{i}<\lambda_{\textrm{cross}}$ and under the condition that $M\gg 1/(1-\beta)$, the noise anti-correlation significantly modifies the outcome. We can express the approximate relationships as follows:

$$ $\displaystyle\sigma_{\theta,i}^{2}$ $\displaystyle\approx\frac{\eta^{2}\sigma_{\delta g,i}^{2}}{2(1-\beta)(1+\beta) }\cdot\frac{M}{3}\frac{1+\beta}{1-\beta}\ ,$ (12a) $\displaystyle\sigma_{v,i}^{2}$ $\displaystyle\approx\frac{\eta^{2}\sigma_{\delta g,i}^{2}}{2(1-\beta)(1+\beta) }\cdot 2\ ,$ (12b) $\displaystyle\tau_{i}$ $\displaystyle\approx\frac{M}{3}\frac{1+\beta}{1-\beta}\ .$ (12c) $$

The derivation of these formulas is provided in Section C.3.

Figure: (a) Variances
Refer to caption: extracted/5732104/Figs/top_eigvecs/variance.jpeg

If we once again assume that the noise covariance matrix is proportional to the Hessian matrix, such that $\sigma_{\delta g,i}^{2}\propto\lambda_{i}$, we obtain the following power laws for the variances: $\sigma_{\theta,i}^{2}\propto\lambda_{i}$ and $\sigma_{v,i}^{2}\propto\lambda_{i}$. This implies that in the subspace spanned by Hessian eigenvectors with eigenvalues $\lambda_{i}<\lambda_{\textrm{cross}}$, the weight variance $\bm{\mathbf{\Sigma}}$ is not isotropic but proportional to the Hessian matrix $\bm{\mathbf{H}}$.

It is s noteworthy that the correlation time $\tau_{i}\approx\frac{M}{3}\frac{1+\beta}{1-\beta}\eqqcolon\tau_{\textrm{SGD}}$ in this subspace is independent of the Hessian eigenvalue. Moreover, $\tau_{\textrm{SGD}}$ is equivalent to the correlation time of the noise $M/3$, up to a factor that depends on momentum.

###### Theorem 4.3 .

###### Corollary 4.4 .

###### Corollary 4.5 .

## 5 Numerics

### 5.1 Analysis Setup

In order to corroborate our theoretical findings, we have conducted a small-scale experiment. We have trained a LeNet architecture, similar to the one described in [], using the Cifar10 dataset []. LeNet is a compact convolutional network comprised of two convolutional layers followed by three dense layers. The network comprises approximately 137,000 parameters. As our loss function, we employed Cross Entropy, along with an L2 regularization with a prefactor of $10^{-4}$.
Here we present results for a single seed and specific hyperparameters. However, we have also performed tests with different seeds and combinations of hyperparameters, all of which showed comparable qualitative behavior (see Appendix J). Furthermore, in Appendix L we studied a ResNet architecture [], a different and more modern network, where we obtained similar results.

We used SGD to train the network for 100 epochs, employing an exponential learning rate schedule that reduces the learning rate by a factor of $0.98$ each epoch. The initial learning rate is set at $5\cdot 10^{-3}$, which eventually reduces to approximately $7\cdot 10^{-4}$ after 100 epochs. The momentum parameter and the minibatch size $S$ are set to $0.9$ and $50$, respectively, which results in a thousand minibatches per epoch, $M=1000$. This setup achieves 100% training accuracy and 63% testing accuracy. The evolution of loss and accuracy during training can be seen in Appendix I. We then compute the variances right after the initial schedule over a period of 20 additional epochs, equivalent to 20,000 update steps. Throughout this analysis period, the learning rate is maintained at $7\cdot 10^{-4}$ and the recorded weights are designated by $\vec{\theta}_{k}$, with $k=K,\dots,K+T$ and $T=20{,}000$.

Given the impracticability of obtaining the full covariance matrix for all weights and biases over this period due to the excessive memory requirements, we limit our analysis to a specific subspace. We approximate the five thousand largest eigenvalues and their associated eigenvectors of the Hessian matrix $\vec{H}(\vec{\theta}_{K})$, drawn from the roughly 137,000 total, using the Lanczos algorithm []. Here, $\vec{\theta}_{K}$ represents the weights at the beginning of the analysis period. To compute the Hessian-vector products required for the Lanczos algorithm, we employ the resource-efficient Pearlmutter trick []. The eigenvectors of the Hessian matrix are represented by $\bm{\mathbf{p_{i}}}$, and the projected weights by $\theta_{k,i}=\vec{\theta}_{k}\cdot\vec{p_{i}}$. The variances are computed exclusively for these particular directions. The distribution of the approximated 5,000 eigenvalues is illustrated in Appendix G.

### 5.2 Noise Autocorrelations

We scrutinize the correlations of noise by recording both the minibatch gradient $\vec{g}_{k}(\vec{\theta}_{k})$
and the total gradient $\vec{\nabla}L(\vec{\theta}_{k})$ at each update step throughout the analysis period, enabling us to capture the actual noise term $\bm{\mathbf{\delta g}}_{k}(\vec{\theta}_{k})$. All these are projected onto the approximated Hessian eigenvectors.

The theoretical prediction for the anti-correlation of the noise is proportional to the inverse of the number of batches per epoch, in our case on the order of $10^{-3}$. To extract the predicted relationship from the fluctuating data, we compute the autocorrelation of the noise term for each individual Hessian eigenvector. We then proceed to average these results across the 5,000 approximated eigenvectors. Figure 1 provides a visual representation of this analysis, showcasing a strong alignment between the empirical autocorrelation of noise and the prediction derived from our theory. This consistency can be interpreted as a validation of our assumption that the noise is spatially independent.

### 5.3 Variances and Correlation time

Previous studies have observed that network weights continue to traverse the parameter space even after the loss appears to have stabilized [, , ]. This behavior persists despite the use of L2 regularization and implies that the recorded weights, $\vec{\theta}_{k}$, do not settle into a stationary distribution. Notably, however, over the course of the 20 epochs under scrutiny, the weight movement, excluding the SGD noise, appears to be approximately linear in time. This suggests that the mean velocity $\bm{\mathbf{\bar{v}}}\coloneqq\langle\bm{\mathbf{v}}_{k}\rangle$ is substantial compared to the SGD noise. To isolate this ongoing movement and uncover the underlying structure, we redefine $\vec{\theta}_{k}$ and $\vec{v}_{k}$ by subtracting the mean velocity. This results in $\vec{\theta}_{k}^{\textrm{(s)}}\coloneqq\vec{\theta}_{k}-\vec{\bar{v}}\cdot k$ and $\vec{v}_{k}^{\textrm{(s)}}\coloneqq\vec{v}_{k}-\vec{\bar{v}}$. We then compute the variances of these redefined values, $\vec{\theta}_{k}^{\textrm{(s)}}$ and $\vec{v}_{k}^{\textrm{(s)}}$, which exhibit a more stationary distribution.

Again, we limit our variance calculations to the directions of the 5,000 approximated Hessian eigenvectors. In the two different regimes of Hessian eigenvalues, either greater or lesser than the crossover value $\lambda_{\textrm{cross}}$, the weight and velocity variance closely follow the respective power law predictions from our theory (see upper panel of Figure 2). The slight discrepancy, where the predicted exponent of one does not lie within the error bars, may arise from minor deviations in the noise covariance from the Hessian approximation $\vec{C}\propto\vec{H}$.
The calculated correlation time, derived from the ratio between the weight and velocity variance, aligns reasonably well with our theoretical predictions (see lower panel of Figure 2). This correlation time prediction remains independent of the exact relation between $\vec{C}$ and $\vec{H}$, thereby providing a more general result.

## 6 Limitations

As highlighted in Section 5.3, the weight distribution is not entirely stationary; the weights exhibit a finite mean drift. To fully understand the underlying weight variances, it is essential to account for this mean velocity. However, for time windows extending beyond the measured 20 epochs, the mean velocity’s variability increases. This change can cause the average weight movement to deviate from linearity, potentially leading to higher variances in flat directions than predicted.

In general, the late phase of training under scrutiny may not be representative of the entire training process. However, we argue that exploring such a regime, especially where a quadratic approximation of the loss seems feasible, can provide important insights into the dynamics of SGD, and in particular its stability around a minimum. Furthermore, it is an important starting point for our research, since such a situation can be treated analytically.

## 7 Discussion

We provide an intuitive interpretation of our results through the correlation time associated with different Hessian eigenvectors and corroborate this timescale empirically. A natural question then emerges about how this correlation time relates to the total duration of the analysis period. Indeed, the maximum correlation time, $\tau_{\textrm{SGD}}$, predicted for the flat directions is roughly seven epochs, comparable to the recorded 20-epoch period. However, the outcomes differ considerably when we sample the examples without replacement during training (see Appendix H). Hence, we infer that the observed correlation time behavior stems from the correlations resulting from sampling without replacement.

Here, we primarily focus on the weight covariance denoted as $\bf\Sigma$, treating the correlation time $\tau_{i}$ and the velocity variance $\bf\Sigma_{v}$ as auxiliary variables. Our assumption is that the velocity covariance $\bf\Sigma_{v}$ is directly proportional to the noise covariance $\bf C$. However, the behavior of the weight covariance $\bf\Sigma$ is more complex. This complexity arises from the autocorrelation time, which links the two variances. Specifically, in the context of a single eigendirection, the relationship is expressed as $\sigma_{\theta,i}^{2}=\frac{1}{2}\tau_{i}\sigma_{v,i}^{2}$. A detailed examination of these individual quantities offers further insights.

Additionally, our findings for the three quantities contrast with a previous empirical study []. In Appendix F, we clarify how the different analysis method used in that study impacts the results due to finite size effects.

In this manuscript, we study the influence of smaller Hessian eigenvalues and corresponding eigendirections, which are often overlooked in discussions about the characteristics of a minimum in optimization landscapes. Previous research primarily focuses on the role of larger Hessian eigenvalues in defining crucial properties of a minimum [, ]. Our work raises the question of the potential impact of the lower-than-expected weight variance in flat directions, which we argue are significant in this scenarios. We propose that this reduced variance might enhance the stability of a minimum, particularly given the substantial difference in the number of flat and steep directions, as illustrated in Appendix G. This suggests that the observed reduction in variance helps explain why SGD exploration predominantly focuses on a limited subspace of Hessian eigenvectors with higher curvature, as reported in prior studies [, ].

To further substantiate this argument, we examine the impact of weight variance on the additional training error. In a simplified scenario with a quadratic loss for a Hessian eigenvector direction with eigenvalue $\lambda_{i}$ and Gaussian weight fluctuations of variance $\sigma_{\theta,i}^{2}$, the expected loss variance due to fluctuations is $l_{{\rm fluct},i}=\frac{1}{2}\lambda_{i}\sigma_{\theta,i}^{2}$. Although the contribution of flat directions to the total fluctuation of the loss $l_{\rm fluct}=\sum_{i}l_{{\rm fluct},i}$ is reduced by their small Hessian eigenvalues, this is counterbalanced by their abundance.
Analyzing the 5,000 largest Hessian eigenvalues of the LeNet, we find that our variance prediction leads to a 62% reduction in loss fluctuation compared to the constant weight variance assumption, despite only affecting the smaller Hessian eigenvalues.
We note that this effect would only be stronger if we also considered the remaining 132,000 Hessian eigendirections.

Furthermore, we argue that flat directions are significant for generalization, a point elaborated in Appendix K. Our investigation of LeNet’s test accuracy, following the decomposition of neural network weights in the Hessian eigenbasis, reveals compelling insights. Omitting the projection of the weight vector onto the top 35 eigenvectors only reduces test accuracy from 63% to 54%. Excluding the top 5,000 eigenvectors, as identified in our analysis, leads to a further drop to 44%. The accuracy only falls to 10%, akin to random guessing, after discarding around 7,000 eigenvectors. These findings underscore the relevance of small Hessian eigenvalues and weight variances for controlling fluctuations in the magnitude of weights, thereby affirming the importance of weight fluctuations in these flat directions.

##### Conclusion

Our exploration of anti-correlations in SGD noise, which result from drawing examples without replacement, reveals a lower-than-expected weight variance in Hessian eigendirections with eigenvalues smaller than the crossover value $\lambda_{\textrm{cross}}$. The introduction of the concepts of intrinsic correlation time, shaped by the loss landscape, and constant noise correlation time, further illuminates the dynamics of SGD optimization.

One potential implication of our findings relates to the role of decreased weight variance in flat directions for SGD’s stability. This observed decrease implies a reduction in the exploration of the weight space along these directions. This diminished exploration could boost the stability of SGD, as it encourages the network to concentrate on fine-tuning parameters along a handful of directions characterized by higher curvature and weight variances.

## Appendix A Definition of Limiting Quantities

When we speak of a covariance matrix or an average in the main text and in the following sections of the appendix, we mean the limiting average or the limiting covariance, unless otherwise specified. In other words, we are interested in the average of a quantity over one infinite run of SGD optimization, not the mean value for a fixed update step $k$ averaged over multiple runs of SGD optimization. With this in mind, we define the covariance matrix of two quantities $\bm{\mathbf{a}}_{k}$ and $\bm{\mathbf{b}}_{k}$ as

$$ $\displaystyle\textrm{cov}(\bm{\mathbf{a}}_{k},\bm{\mathbf{b}}_{k})\coloneqq \left\langle\left(\bm{\mathbf{a}}_{k}-\langle\bm{\mathbf{a}}_{k}\rangle_{k} \right)\left(\bm{\mathbf{b}}_{k}-\langle\bm{\mathbf{b}}_{k}\rangle_{k}\right)^ {\top}\right\rangle_{k}$ (13) $$

and the limiting average is defined as

$$ $\displaystyle\langle\bm{\mathbf{a}}_{k}\rangle_{k}=\lim_{K\to\infty}\frac{1}{K +1}\sum_{k=k_{0}}^{k_{0}+K}\bm{\mathbf{a}}_{k}\ .$ (14) $$

When possible, we will suppress $k$ and denote the average as $\langle\cdot\rangle$. The average is independent of the starting value $k_{0}$, therefore we can shift indices within the average, meaning $\langle\bm{\mathbf{a}}_{k}\rangle=\langle\bm{\mathbf{a}}_{k+l}\rangle$ for any $l\in\mathbb{Z}$.

To see this we take any integer $l\in\mathbb{Z}$ and instead of adding it to the index $k$ we can also subtract it from the starting value $k_{0}$ and then separate the sum into two sums,

$$ $\displaystyle\langle\bm{\mathbf{a}}_{k+l}\rangle$ $\displaystyle=\lim_{K\to\infty}\frac{1}{K+1}\sum_{k=k_{0}}^{k_{0}+K}\bm{ \mathbf{a}}_{k+l}$ $\displaystyle=\lim_{K\to\infty}\frac{1}{K+1}\sum_{k=k_{0}-l}^{k_{0}-l+K}\bm{ \mathbf{a}}_{k}$ $\displaystyle=\lim_{K\to\infty}\frac{1}{K+1}\sum_{k=k_{0}-l}^{k_{0}-1}\bm{ \mathbf{a}}_{k}+\lim_{K\to\infty}\frac{1}{K+1}\sum_{k=k_{0}}^{k_{0}-l+K}\bm{ \mathbf{a}}_{k}\ .$ (15) $$

The first sum is independent of $K$ except for the factor $\frac{1}{K+1}$, so the limit of the first part is zero. The second part of the limit can be rearranged as follows,

$$ $\displaystyle\lim_{K\to\infty}\frac{1}{K+1}\sum_{k=k_{0}}^{k_{0}-l+K}\bm{ \mathbf{a}}_{k}$ $\displaystyle=\lim_{K\to\infty}\frac{K-l+1}{K+1}\frac{1}{K-l+1}\sum_{k=k_{0}}^ {k_{0}-l+K}\bm{\mathbf{a}}_{k}$ $\displaystyle=\lim_{\tilde{K}\to\infty}\frac{1}{\tilde{K}+1}\sum_{k=k_{0}}^{k_ {0}+\tilde{K}}\bm{\mathbf{a}}_{k}$ (16) $\displaystyle\eqqcolon\langle\bm{\mathbf{a}}_{k}\rangle\ ,$ (17) $$

where we used $\lim_{K\to\infty}\frac{K-l+1}{K+1}=1$ and renamed $K-l$ in the second to last step. All together this gives us the desired relation $\langle\bm{\mathbf{a}}_{k}\rangle=\langle\bm{\mathbf{a}}_{k+l}\rangle$.

If one were to consider the covariance for a fixed update step $k$ averaged over multiple runs of SGD optimization, it is possible that this covariance could depend on the index $k$, but this is not our case of interest.

## Appendix B Weight independent Noise

It is important to note that Equation 2, the formula for the autocorrelation of the noise described in Section 4.1 and derived in Appendix E, is strictly speaking only valid for a static weight vector. During training, the weights change with each update step, which could potentially change this relationship. However, if the noise terms were independent of the weight vector, the formula would still hold. A simple case where this assumption of weight independent noise holds is a model where not only the total loss can be described by a quadratic function,

$$ $\displaystyle L(\bm{\mathbf{\theta}})=\frac{1}{2}\bm{\mathbf{\theta}}^{\top} \bm{\mathbf{H}}\bm{\mathbf{\theta}}+{\rm const.}\ ,$ (18) $$

but the loss for a single example is described by a quadratic function as well, with the same Hessian,

$$ $\displaystyle l(\bm{\mathbf{\theta}},\bm{\mathbf{x}}_{n})=\frac{1}{2}(\bm{ \mathbf{\theta}}-\bm{\mathbf{\mu}}_{n})^{\top}\bm{\mathbf{H}}(\bm{\mathbf{ \theta}}-\bm{\mathbf{\mu}}_{n})+{\rm const.}\ .$ (19) $$

However, in this model, the minimum of the loss function for this single example is offset by an example-dependent vector. This offset results in a noise term that is independent of the current weight vector for each individual example,

$$ $\displaystyle\bm{\mathbf{\nabla}}\bigl{(}l(\bm{\mathbf{\theta}},\bm{\mathbf{x} }_{n})-L(\bm{\mathbf{\theta}})\bigr{)}$ $\displaystyle=-\bm{\mathbf{H}}\bm{\mathbf{\mu}}_{n}$ $\displaystyle\neq f(\bm{\mathbf{\theta}})\ .$ (20) $$

Therefore, also the noise introduced by SGD in this model would be independent of the current weight vector. For the later stages of training, such an assumption for the noise could be a good approximation, as evidenced by the close fit of the data to the theory in Figure 1.

## Appendix C Variance Calculation

### C.1 Commutativity of the covariance matrices

In this section, we show that if $\left[\bm{\mathbf{C}},\bm{\mathbf{H}}\right]=0$ also $\bm{\mathbf{\Sigma}}$ and $\bm{\mathbf{\Sigma_{v}}}$ will commute with $\bm{\mathbf{C}}$, with $\bm{\mathbf{H}}$ and with each other. We make the assumptions one to three from Section 4.3 and therefore the SGD update equations become

$$ $\displaystyle\vec{v}_{k}$ $\displaystyle=-\eta\vec{H}\vec{\theta}_{k-1}+\beta\vec{v}_{k-1}-\eta\vec{ \delta g}_{k}\ ,$ (21) $\displaystyle\bm{\mathbf{\theta}}_{k}$ $\displaystyle=\left(\vec{1}-\eta\vec{H}\right)\vec{\theta}_{k-1}+\beta\vec{v}_ {k-1}-\eta\vec{\delta g}_{k}\ ,$ (22) $$

which can be rewritten by using the vector $\vec{y}_{k}\coloneqq\bigl{(}\bm{\mathbf{\theta}}_{k}\quad\vec{v}_{k}\bigr{)}^{\top}$, combining both the current weight and velocity variable, to be

$$ $\displaystyle\vec{y}_{k+1}=\vec{X}\vec{y}_{k}-\vec{z}_{k+1}\ .$ (23) $$

Here, $\vec{z}_{k}\coloneqq\bigl{(}\eta\vec{\delta g}_{k}\quad\eta\vec{\delta g}_{k}
\bigr{)}^{\top}$ contains the current noise term, and the matrix governing the deterministic part of the update is defined to be

$$ $\displaystyle\vec{X}\coloneqq\begin{pmatrix}\vec{1}-\eta\vec{H}&\beta\vec{1}\\ -\eta\vec{H}&\beta\vec{1}\end{pmatrix}\ .$ (24) $$

By iteratively applying Equation 23 we obtain

$$ $\displaystyle\vec{y}_{k+h}=\vec{X}^{h}\vec{y}_{k}-\sum_{i=1}^{h}\vec{X}^{h-i} \vec{z}_{k+i}\ .$ (25) $$

Under the assumption $0\leq\beta<1$ and $0<\eta\lambda_{i}<2(1+\beta)$, for all eigenvalues $\lambda_{i}$ of $\vec{H}$, the magnitude of the eigenvalues of $\bm{\mathbf{X}}$ will be less than one. It is straightforward to show this relation for the eigenvalues of $\vec{X}$ by using the eigenbasis of $\vec{H}$. Therefore,

$$ $\displaystyle\lim_{h\rightarrow\infty}\vec{X}^{h}\vec{y}_{k}=0.$ (26) $$

As we can shift the index in the weight variance, $h$ can be chosen arbitrarily large, which yields the following relation for the covariance

$$ $\displaystyle\langle\vec{y}_{k}\vec{y}_{k}^{\top}\rangle=\lim_{h\rightarrow \infty}\sum_{i,j=1}^{h}\vec{X}^{h-i}\langle\vec{z}_{k+i}\vec{z}_{k+j}^{\top} \rangle\left(\vec{X}^{h-j}\right)^{\top}.$ (27) $$

Because Equation 25 together with Equation 26 implies $\langle\bm{\mathbf{y}}_{k}\rangle=0$ and therefore $\langle\vec{\theta}_{k}\rangle=0$ and $\langle\vec{v}_{k}\rangle=0$, the left hand side of Equation 27 contains the covariance matrices of interest,

$$ $\displaystyle\langle\vec{y}_{k}\vec{y}_{k}^{\top}\rangle=\begin{pmatrix}\vec{ \Sigma}&\langle\vec{\theta}_{k}\vec{v}_{k}^{\top}\rangle\\ \langle\vec{v}_{k}\vec{\theta}_{k}^{\top}\rangle&\vec{\Sigma_{v}}\end{pmatrix}\ .$ (28) $$

From Equation 27 we can also infer that $\langle\vec{y}_{k}\vec{y}_{k}^{\top}\rangle$ is finite as the magnitude of the eigenvalues of $\bm{\mathbf{X}}$ is less than one. Consequently, by Equation 28, the covariance matrices $\bm{\mathbf{\Sigma}}$ and $\vec{\Sigma_{v}}$ are finite as well. The average over the noise terms $\vec{z}_{k}$ on the right hand side of Equation 27 is by assumption equal to

$$ $\displaystyle\langle\vec{z}_{k+i}\vec{z}_{k+j}^{\top}\rangle=\eta^{2}\left( \delta_{i,j}-\mathbf{1}_{\{1,...,M\}}(|i-j|)\frac{M-|i-j|}{M(M-1)}\right)\cdot \begin{pmatrix}\vec{C}&\vec{C}\\ \vec{C}&\vec{C}\end{pmatrix}\ ,$ (29) $$

from which it follows that for any finite $h$ every matrix entry of the two by two super matrix on the right hand side of Equation 27 is a function of $\bm{\mathbf{C}}$ and $\bm{\mathbf{H}}$. Therefore, when considering the limit $h\rightarrow\infty$, $\left[\bm{\mathbf{C}},\bm{\mathbf{H}}\right]=0$ implies that $\vec{\Sigma}$ and $\vec{\Sigma_{v}}$ will also commute with $\bm{\mathbf{C}}$, with $\bm{\mathbf{H}}$ and with each other.

### C.2 Proof of the variance formula for one specific eigenvalue

Since $\bm{\mathbf{\Sigma}}$ and $\bm{\mathbf{\Sigma_{v}}}$ will commute with $\bm{\mathbf{C}}$, with $\bm{\mathbf{H}}$ and with each other, it is sufficient to prove the one dimensional case. For the multidimensional case simply apply the proof in the direction of each common eigenvector individually.
The expectation values discussed below are computed with respect to the asymptotic distributions of $\theta$ and $v$, since we are only interested in the asymptotic behavior of training. We want to find $\sigma_{\theta}^{2}\coloneqq\langle\theta_{k}\theta_{k}\rangle$ and $\sigma_{v}^{2}\coloneqq\langle v_{k}v_{k}\rangle$. We assume $0\leq\beta<1$ and $0<\eta\lambda<2(1+\beta)$ where $\lambda$ is the hessian eigenvalue.

The equations describing SGD in one dimension are:

$$ $\displaystyle g_{k}(\theta)$ $\displaystyle=\frac{\partial}{\partial\theta}L(\theta)+\delta g_{k}(\theta)$ (30) $\displaystyle v_{k}$ $\displaystyle=-\eta g_{k}(\theta_{k-1})+\beta v_{k-1}$ (31) $\displaystyle\theta_{k}$ $\displaystyle=\theta_{k-1}+v_{k}\ .$ (32) $$

Our remaining assumptions can then be described the following way

$$ $\displaystyle L(\theta)$ $\displaystyle=\frac{1}{2}\theta\lambda\theta$ (33) $\displaystyle\delta g_{k}(\theta)$ $\displaystyle=\delta g_{k}\textrm{ }\dots\textrm{ is independent of }\theta$ (34) $\displaystyle\Rightarrow\langle\delta g_{k}\delta g_{k+h}\rangle$ $\displaystyle=\sigma_{\delta g}^{2}\left(\delta_{h,0}-\mathbf{1}_{\{1,...,M\}} (|h|)\frac{M-|h|}{M(M-1)}\right)$ (35) $\displaystyle\sigma_{\delta g}^{2}$ $\displaystyle\coloneqq\langle\delta g_{k}\delta g_{k}\rangle\ .$ (36) $$

With these assumptions the update equations can be described by a discrete stochastic linear equation of second order

$$ $\displaystyle\theta_{k}=(1+\beta-\eta\lambda)\theta_{k-1}-\beta\theta_{k-2}- \eta\delta g_{k}$ (37) $$

which can be rewritten into matrix form as follows

$$ $\displaystyle\bm{\mathbf{x}}_{k}$ $\displaystyle=\vec{D}\vec{x}_{k-1}-\eta\delta g_{k}\vec{e}_{1}$ (38) $\displaystyle\bm{\mathbf{x}}_{k}$ $\displaystyle\coloneqq\begin{pmatrix}\theta_{k}\\ \theta_{k-1}\end{pmatrix}$ (39) $\displaystyle\bm{\mathbf{e}}_{1}$ $\displaystyle\coloneqq\begin{pmatrix}1\\ 0\end{pmatrix}$ (40) $\displaystyle\vec{D}$ $\displaystyle\coloneqq\begin{pmatrix}1+\beta-\eta\lambda&-\beta\\ 1&0\end{pmatrix}\ .$ (41) $$

We are now interested in the following covariance matrix

$$ $\displaystyle\vec{\tilde{\Sigma}}$ $\displaystyle\coloneqq\left\langle\vec{x}_{k}\vec{x}_{k}^{\top}\right\rangle$ $\displaystyle=\begin{pmatrix}\sigma_{\theta}^{2}&\langle\theta_{k}\theta_{k-1} \rangle\\ \langle\theta_{k}\theta_{k-1}\rangle&\sigma_{\theta}^{2}\end{pmatrix}$ (42) $$

where the second equality is due to the fact that $\langle\theta_{k}\theta_{k}\rangle=\langle\theta_{k-1}\theta_{k-1}\rangle$.
As we are interested in the asymptotic covariance, this expectation value is independent of any finite shift of the index $k$. By inserting Equation 38 into $\left\langle\vec{x}_{k}\vec{x}_{k}^{\top}\right\rangle$ we arrive at the following equality

$$ $\displaystyle\left\langle\vec{x}_{k}\vec{x}_{k}^{\top}\right\rangle=\vec{D} \left\langle\vec{x}_{k-1}\vec{x}_{k-1}^{\top}\right\rangle\vec{D}^{\top}+\eta^ {2}\left\langle\delta g_{k}\delta g_{k}\right\rangle\vec{e}_{1}\vec{e}_{1}^{ \top}-\eta\left(\vec{D}\left\langle\bm{\mathbf{x}}_{k-1}\delta g_{k}\right \rangle\vec{e}_{1}^{\top}+\left(\vec{D}\left\langle\bm{\mathbf{x}}_{k-1}\delta g _{k}\right\rangle\vec{e}_{1}^{\top}\right)^{\top}\right)$ (43) $$

which can be simplified to the equivalent equation

$$ $\displaystyle\vec{\tilde{\Sigma}}-\vec{D}\vec{\tilde{\Sigma}}\vec{D}^{\top}= \eta^{2}\sigma_{\delta g}^{2}\vec{e}_{1}\vec{e}_{1}^{\top}-\eta\left(\vec{D} \left\langle\bm{\mathbf{x}}_{k-1}\delta g_{k}\right\rangle\vec{e}_{1}^{\top}+ \left(\vec{D}\left\langle\bm{\mathbf{x}}_{k-1}\delta g_{k}\right\rangle\vec{e} _{1}^{\top}\right)^{\top}\right)\ .$ (44) $$

If we apply the left-hand side on the vector $\vec{e}_{1}$, it can be expressed as

$$ $\displaystyle\left[\vec{\tilde{\Sigma}}-\vec{D}\vec{\tilde{\Sigma}}\vec{D}^{ \top}\right]\vec{e}_{1}$ $\displaystyle=\vec{F}_{1}^{-1}\vec{\tilde{\Sigma}}\vec{e}_{1}$ (45) $\displaystyle\vec{F}_{1}^{-1}$ $\displaystyle\coloneqq\begin{pmatrix}\eta\lambda(2-\eta\lambda)-2\beta(1+\beta -\eta\lambda)&2\beta(1+\beta-\eta\lambda)\\ -(1+\beta-\eta\lambda)&1+\beta\end{pmatrix}\ .$ (46) $$

Also notice $v_{k}=\theta_{k}-\theta_{k-1}$ and therefore

$$ $\displaystyle\sigma_{v}^{2}=2\sigma_{\theta}^{2}-2\langle\theta_{k}\theta_{k-1 }\rangle\ ,$ (47) $$

again due to the fact that the expectation value does not depend on k. Hence, the variances can then be expressed as

$$ $\displaystyle\begin{pmatrix}\sigma_{\theta}^{2}\\ \sigma_{v}^{2}\end{pmatrix}$ $\displaystyle=\vec{F}_{2}\vec{\tilde{\Sigma}}\vec{e}_{1}$ (48) $\displaystyle\vec{F}_{2}$ $\displaystyle\coloneqq\begin{pmatrix}1&0\\ 2&-2\end{pmatrix}\ .$ (49) $$

We define the matrix $\vec{F}\coloneqq\vec{F}_{2}\vec{F}_{1}$. By applying both sides of Equation 44 to the vector $\vec{e_{1}}$, then multiplying by the matrix $\vec{F}$ from the left and using Equations 45 and 48 we obtain

$$ $\displaystyle\begin{pmatrix}\sigma_{\theta}^{2}\\ \sigma_{v}^{2}\end{pmatrix}=\vec{F}\left[\eta^{2}\sigma_{\delta g}^{2}\vec{e}_ {1}\vec{e}_{1}^{\top}-\eta\left(\vec{D}\left\langle\bm{\mathbf{x}}_{k-1}\delta g _{k}\right\rangle\vec{e}_{1}^{\top}+\left(\vec{D}\left\langle\bm{\mathbf{x}}_{ k-1}\delta g_{k}\right\rangle\vec{e}_{1}^{\top}\right)^{\top}\right)\right] \vec{e}_{1}\ .$ (50) $$

with

$$ $\displaystyle\bm{\mathbf{F}}=\frac{1}{(1-\beta)\bigl{(}2(1+\beta)-\eta\lambda) \bigr{)}}\begin{pmatrix}\frac{1+\beta}{\eta\lambda}&\frac{2\beta(\eta\lambda-1 -\beta)}{\eta\lambda}\\ 2&2(\eta\lambda-2)\end{pmatrix}\ .$ (51) $$

To simplify Equation 50 further we go back to Equation 38 and iterate it to obtain

$$ $\displaystyle\bm{\mathbf{x}}_{k}=\vec{D}^{n}\bm{\mathbf{x}}_{k-n}-\eta\sum_{h= 0}^{n-1}\vec{D}^{h}\vec{e}_{1}\delta g_{k-h}\ .$ (52) $$

We note that $\left\langle\bm{\mathbf{x}}_{k-n}\delta g_{k}\right\rangle=0$ for $n\geq M$. The correlation between noise terms separated by at least one epoch vanishes, and $\bm{\mathbf{x}}_{k}$ only depends on past noise terms. By setting $n=M$ we find

$$ $\displaystyle\left\langle\bm{\mathbf{x}}_{k-1}\delta g_{k}\right\rangle$ $\displaystyle=\vec{D}^{M}\left\langle\bm{\mathbf{x}}_{k-1-M}\delta g_{k}\right \rangle-\eta\sum_{h=0}^{M-1}\vec{D}^{h}\vec{e}_{1}\left\langle\delta g_{k} \delta g_{k-1-h}\right\rangle$ $\displaystyle=-\eta\sigma_{\delta g}^{2}\sum_{h=0}^{M-1}\vec{D}^{h}\vec{e}_{1} \left(-\frac{M-(h+1)}{M(M-1)}\right)\ ,$ (53) $$

where the assumption about the correlation of the noise terms, Equation 35, was inserted for the last line. Equation 53 is a sum of a finite geometric series and a derivative of that which can be simplified to

$$ $\displaystyle\left\langle\bm{\mathbf{x}}_{k-1}\delta g_{k}\right\rangle$ $\displaystyle=\eta\sigma_{\delta g}^{2}\frac{\bm{\mathbf{D}}^{M}+(\vec{1}-\vec {D})M-\vec{1}}{(\vec{1}-\vec{D})^{2}M(M-1)}\,\vec{e}_{1}\ .$ (54) $$

Substituting this result back into Equation 50 yields

$$ $\displaystyle\begin{pmatrix}\sigma_{\theta}^{2}\\ \sigma_{v}^{2}\end{pmatrix}$ $\displaystyle=\eta^{2}\sigma_{\delta g}^{2}\vec{F}\left[\vec{e}_{1}-\left(\bm{ \mathbf{E}}+\vec{E}^{\top}\right)\vec{e}_{1}\right]$ (55) $$

with the definition

$$ $\displaystyle\vec{E}$ $\displaystyle\coloneqq\vec{D}\frac{\bm{\mathbf{D}}^{M}+(\vec{1}-\vec{D})M-\vec {1}}{(\vec{1}-\vec{D})^{2}M(M-1)}\,\vec{e}_{1}\vec{e}_{1}^{\top}\ .$ (56) $$

With Equation 55 we have arrived at the exact formula for the variances which can easily be evaluated numerically.

### C.3 Approximation of the exact formula

It is possible to approximate the exact result for the variance assuming small or large eigenvalues, respectively. For that, it is necessary to approximate $\bm{\mathbf{D}}^{M}\bm{\mathbf{e_{1}}}$. To do so, we will use the the following eigendecomposition of $\bm{\mathbf{D}}$

$$ $\displaystyle\vec{D}$ $\displaystyle=\vec{Q}\vec{\Lambda}\vec{Q}^{-1}$ (57) $\displaystyle\vec{\Lambda}$ $\displaystyle=\begin{pmatrix}\Lambda_{+}&0\\ 0&\Lambda_{-}\end{pmatrix}$ (58) $\displaystyle\vec{Q}$ $\displaystyle=\begin{pmatrix}\Lambda_{+}&\Lambda_{-}\\ 1&1\end{pmatrix}$ (59) $\displaystyle\vec{Q}^{-1}$ $\displaystyle=\frac{1}{\Lambda_{+}-\Lambda_{-}}\begin{pmatrix}1&-\Lambda_{-}\\ -1&\Lambda_{+}\end{pmatrix}$ (60) $\displaystyle\Lambda_{\pm}$ $\displaystyle=\frac{1}{2}\left(1+\beta-\eta\lambda\pm s\right)$ (61) $\displaystyle s$ $\displaystyle\coloneqq\sqrt{\left(1-\beta\right)^{2}-\eta\lambda\bigl{(}2(1+ \beta)-\eta\lambda\bigr{)}}$ (62) $$

It is straightforward to show that the magnitude of the eigenvalues of $\vec{D}$ is strictly smaller than one, $|\Lambda_{\pm}|<1$, under the conditions $0<\eta\lambda<2(1+\beta)$ and $0\leq\beta<1$.

#### Large hessian eigenvalues

$$ $\displaystyle\sigma_{\theta}^{2}$ $\displaystyle\approx\frac{\eta^{2}\sigma_{\delta g}^{2}}{(1-\beta)(2(1+\beta)- \eta\lambda)}\cdot\frac{1+\beta}{\eta\lambda}$ (63) $\displaystyle\sigma_{v}^{2}$ $\displaystyle\approx\frac{\eta^{2}\sigma_{\delta g}^{2}}{(1-\beta)(2(1+\beta)- \eta\lambda)}\cdot 2$ (64) $$

We will show that this approximation for large hessian eigenvalues is valid under the assumption $M(\eta\lambda)^{2}\gg 1$ where $M$ is the number of batches per epoch. However, numerical studies indicate that these relations also hold under the previously mentioned assumptions of $\frac{M\eta\lambda}{1-\beta}\gg 1$, equivalent to $\lambda\gtrsim\lambda_{\textrm{cross}}$, and $M(1-\beta)\gg 1$.

Inserting the eigendecomposition of $\bm{\mathbf{D}}$ into the expression $\vec{D}^{M}\vec{e_{1}}$ yields

$$ $\displaystyle\vec{D}^{M}\vec{e_{1}}=\begin{pmatrix}y_{M+1}\\ y_{M}\end{pmatrix}$ (65) $\displaystyle y_{M}\coloneqq\frac{\Lambda_{+}^{M}-\Lambda_{-}^{M}}{\Lambda_{+} -\Lambda_{-}}.$ (66) $$

From the definition of $y_{M}$ one sees that

$$ $\displaystyle y_{M}=\frac{\Lambda_{+}+\Lambda_{-}}{2}\,y_{M-1}+\frac{\Lambda_{ +}^{M-1}+\Lambda_{-}^{M-1}}{2}\ ,$ (67) $$

and by using $|\Lambda_{\pm}|<1$ as well as $y_{0}=0$ one can show iteratively that

$$ $\displaystyle|y_{M}|\leq M+1.$ (68) $$

Therefore, we have

$$ $\displaystyle\left\lVert\vec{D}^{M}\vec{e_{1}}\right\rVert_{\infty}\leq M+1$ (69) $$

where $\left\lVert\cdot\right\rVert_{\infty}$ is denoting the maximum norm $\left\lVert\bm{\mathbf{x}}\right\rVert_{\infty}\coloneqq\max_{i}|x_{i}|$ for a vector $\vec{x}$ or its induced matrix norm $\left\lVert\bm{\mathbf{A}}\right\rVert_{\infty}\coloneqq\max_{i}\sum_{j}|a_{ij}|$ for a matrix $\vec{A}$.

Explicit calculations show that

$$ $\displaystyle\left\lVert(\bm{\mathbf{1}}-\vec{D})^{-1}\right\rVert_{\infty} \leq\frac{4}{\eta\lambda}$ (70) $$

under the assumption that $0\leq\beta<1$ and $0<\eta\lambda<2(1+\beta)$. From here it is straightforward to show that

$$ $\displaystyle\left\lVert\left(\vec{E}+\vec{E}^{\top}\right)\vec{e_{1}}\right \rVert_{\infty}\leq\frac{\tilde{c}}{M(\eta\lambda)^{2}}$ (71) $$

where $\tilde{c}$ is a factor of order unity under the constraints $0\leq\beta<1$ and $0<\eta\lambda<2(1+\beta)$. By substituting this result back into Equation 55 one directly sees that a comparison to the approximation yields

$$ $\displaystyle\left\lvert 1-\frac{\sigma_{\theta}^{2}}{\sigma_{\theta,\textrm{ large}}^{2}}\right\rvert\leq\frac{c_{1}}{M(\eta\lambda)^{2}}\ ,$ (72) $\displaystyle\left\lvert 1-\frac{\sigma_{v}^{2}}{\sigma_{v,\textrm{large}}^{2} }\right\rvert\leq\frac{c_{2}}{M(\eta\lambda)^{2}}\ ,$ (73) $$

where $c_{1}$ and $c_{2}$ are again of order unity and the approximation is defined as

$$ $\displaystyle\begin{pmatrix}\sigma_{\theta,\textrm{large}}^{2}\\ \sigma_{v,\textrm{large}}^{2}\end{pmatrix}$ $\displaystyle\coloneqq\eta^{2}\sigma_{\delta g}^{2}\vec{F}\vec{e}_{1}$ $\displaystyle=\frac{\eta^{2}\sigma_{\delta g}^{2}}{(1-\beta)(2(1+\beta)-\eta \lambda)}\cdot\begin{pmatrix}\frac{1+\beta}{\eta\lambda}\\ 2\end{pmatrix}.$ (74) $$

Interestingly, one can see that the approximation for large hessian eigenvalues is equivalent to the result we would obtain if we assumed there was no autocorrelation of the noise to begin with.

In the case where the stricter assumption is not true, $M(\eta\lambda)^{2}<1$, but the numerically obtained conditions still hold, $\lambda\gtrsim\lambda_{\textrm{cross}}$ and $M(1-\beta)\gg 1$, it occurs that $\left\lVert\left(\vec{E}+\vec{E}^{\top}\right)\vec{e_{1}}\right\rVert_{\infty}$ is no longer small. But in that case, $\vec{F}\left(\vec{E}+\vec{E}^{\top}\right)\vec{e_{1}}$ can still be neglected compared to $\vec{F}\vec{e_{1}}$, as numerical experiments show.

#### Small hessian eigenvalues

To obtain the relations for small hessian eigenvalues, we perform a Taylor expansion with respect to $\lambda$ with the help of computer algebra. We neglect the terms which are at least of order lambda. Numerical study indicates that these relations hold under the mentioned assumption of $\lambda\lesssim\lambda_{\textrm{cross}}$ and $M(1-\beta)\gg 1$.

It is straightforward but lengthy to obtain the following expression using the eigendecomposition of $\vec{D}$

$$ $\displaystyle\begin{pmatrix}\sigma_{\theta}^{2}\\ \sigma_{v}^{2}\end{pmatrix}=\frac{\eta^{2}\sigma_{\delta g}^{2}}{2(1-\beta)(1+ \beta)}\cdot\begin{pmatrix}\frac{M}{3}\frac{1+\beta}{1-\beta}+\mathcal{O}( \lambda)\\ 2+\mathcal{O}(\lambda)\end{pmatrix}$ (75) $$

where the zeroth order terms are simplified under approximation $M(1-\beta)\gg 1$.

### C.4 Satisfying the assumptions of the correlation time relation

In this section we want to show that the weight and velocity variances resulting from stochastic gradient descent as described above and in Section 4.3 satisfies the necessary assumptions (i) to (iii) of Section 4.2 such that the velocity correlation time is equal to $\tau_{i}=2\sigma_{\theta,i}^{2}/\sigma_{v,i}^{2}$. Validity of assumption (i) existence and finiteness of $\bm{\mathbf{\Sigma}}\coloneqq$ cov$(\bm{\mathbf{\theta}}_{k},\bm{\mathbf{\theta}}_{k})$, $\bm{\mathbf{\Sigma_{v}}}\coloneqq$ cov$(\bm{\mathbf{v}}_{k},\bm{\mathbf{v}}_{k})$, and $\langle\bm{\mathbf{\theta}}\rangle$ can be inferred from the calculation presented in Section C.1. Therefore, we concentrate on assumption (ii) $\lim_{n\to\infty}\textrm{cov}\bigr{(}\bm{\mathbf{\theta}}_{k},\bm{\mathbf{
\theta}}_{k+n}\bigr{)}=0$ and (iii) $\lim_{n\to\infty}n\cdot\textrm{cov}\bigr{(}\bm{\mathbf{\theta}}_{k},\bm{
\mathbf{\theta}}_{k+n}-\bm{\mathbf{\theta}}_{k+n+1}\bigr{)}=0$. We consider the one dimensional case, but the extension to the multidimensional case is straightforward. Additionally $\langle\vec{\theta}\rangle=0$ (see Section C.1) and, therefore, the remaining two assumptions (ii) and (iii) can be written as $\lim_{m\rightarrow\infty}{\langle\theta_{k}\theta_{k+m}\rangle}=0$ and $\lim_{m\rightarrow\infty}{m\bigl{(}\langle\theta_{k}\theta_{k+m}\rangle-
\langle\theta_{k}\theta_{k+m+1}\rangle\bigr{)}}=0$.

We will now show that for stochastic gradient descent under the assumptions of Section 4.3 the more restrictive relation $\lim_{m\rightarrow\infty}{m\langle\theta_{k}\theta_{k+m}\rangle}=0$ is satisfied, from which follows (ii) and (iii). Following Section C.2 and using the same notation, we have the relation

$$ $\displaystyle\langle\vec{x}_{k}\vec{x}_{k-m}^{\top}\rangle=\bm{\mathbf{D}} \langle\vec{x}_{k-1}\vec{x}_{k-m}^{\top}\rangle-\eta\vec{e_{1}}\langle\delta g _{k}\vec{x}_{k-m}^{\top}\rangle\ ,$ (76) $$

where $\vec{x}_{k}\coloneqq\bigl{(}\theta_{k}\quad\theta_{k-1}\bigr{)}^{\top}$. For $m>M$, with $M$ being the number of batches per epoch, the correlation with the noise term on the right hand side of Equation 76 is equal to zero as discussed in Section C.2. By iterating Equation 76, for $m>M$ we have

$$ $\displaystyle\langle\vec{x}_{k}\vec{x}_{k-m}^{\top}\rangle$ $\displaystyle=\bm{\mathbf{D}}^{m-M-1}\langle\vec{x}_{k-m+M+1}\vec{x}_{k-m}^{ \top}\rangle$ $\displaystyle=\bm{\mathbf{D}}^{m-M-1}\langle\vec{x}_{k}\vec{x}_{k-M-1}^{\top} \rangle\ .$ (77) $$

As described in Section C.2, the magnitude of both eigenvalues of $\vec{D}$ is strictly smaller than one. This implies that there exists a matrix norm $\left\lVert\cdot\right\rVert_{\vec{D}}$ such that $\left\lVert\vec{D}\right\rVert_{\vec{D}}<1$ from which one can deduce

$$ $\displaystyle\left\lVert\langle\vec{x}_{k}\vec{x}_{k-m}^{\top}\rangle\right \rVert_{\vec{D}}\leq\left\lVert\bm{\mathbf{D}}\right\rVert_{\vec{D}}^{m-M-1} \cdot\left\lVert\langle\vec{x}_{k}\vec{x}_{k-M-1}^{\top}\rangle\right\rVert_{ \vec{D}}\ .$ (78) $$

Taking the limit of $m\to\infty$ we obtain

$$ $\displaystyle\lim_{m\to\infty}m\left\lVert\langle\vec{x}_{k}\vec{x}_{k-m}^{ \top}\rangle\right\rVert_{\vec{D}}$ $\displaystyle\leq\textrm{const}\cdot\lim_{m\to\infty}m\left\lVert\bm{\mathbf{D }}\right\rVert_{\vec{D}}^{m-M-1}$ $\displaystyle=0\ ,$ (79) $$

and because

$$ $\displaystyle\langle\vec{x}_{k}\vec{x}_{k-m}^{\top}\rangle=\begin{pmatrix} \langle\theta_{k}\theta_{k-m}\rangle&\langle\theta_{k}\theta_{k-m-1}\rangle\\ \langle\theta_{k-1}\theta_{k-m}\rangle&\langle\theta_{k-1}\theta_{k-m-1} \rangle\end{pmatrix}$ (80) $$

we finally find

$$ $\displaystyle\lim_{m\rightarrow\infty}{m\langle\theta_{k}\theta_{k-m}\rangle}$ $\displaystyle=0$ $\displaystyle\Rightarrow\lim_{m\rightarrow\infty}{m\langle\theta_{k}\theta_{k+ m}\rangle}$ $\displaystyle=0\ .$ (81) $$

## Appendix D Calculation of the Correlation Time Relation

We want to prove the relation

$$ $\displaystyle\frac{2\sigma_{\theta,i}^{2}}{\sigma_{v,i}^{2}}=\frac{\sum_{n=1}^ {\infty}n\,\langle v_{k,i}v_{k+n,i}\rangle}{\sum_{n=1}^{\infty}\langle v_{k,i} v_{k+n,i}\rangle}\ .$ (82) $$

We assume that that $\langle\vec{\theta}\rangle$, $\bm{\mathbf{\Sigma_{\theta}}}$ and $\bm{\mathbf{\Sigma_{v}}}$ exist and are finite. Without loss of generality, let $\langle\vec{\theta}\rangle=0$. We consider only the one-dimensional case. For the multidimensional case, simply apply the proof in the direction of any basis vector individually. Note, that the relation still holds if $[\bm{\mathbf{\Sigma}},\bm{\mathbf{\Sigma_{v}}}]\neq 0$. In this case, $\sigma_{\theta,i}^{2}$ and $\sigma_{v,i}^{2}$ would just be the variances of the weight and the velocity in the given direction but no longer necessarily eigenvalues of $\bm{\mathbf{\Sigma}}$ and $\bm{\mathbf{\Sigma_{v}}}$.

The remaining two assumptions (ii) and (iii) of Section 4.2 can now be written as

$$ $\displaystyle\lim\limits_{m\rightarrow\infty}{\langle\theta_{k}\theta_{k+m}\rangle}$ $\displaystyle=0$ (83) $\displaystyle\lim\limits_{m\rightarrow\infty}{m\bigl{(}\langle\theta_{k}\theta _{k+m}\rangle-\langle\theta_{k}\theta_{k+m+1}\rangle\bigr{)}}$ $\displaystyle=0\ .$ (84) $$

We begin the proof with the following chain of equations

$$ $\displaystyle\sigma_{\theta}^{2}$ $\displaystyle=\left\langle\theta_{k}^{2}\right\rangle$ $\displaystyle=\left\langle\left(\theta_{k}-\theta_{k+J}+\theta_{k+J}\right)^{2 }\right\rangle$ $\displaystyle=\left\langle\left(\theta_{k}-\theta_{k+J}\right)^{2}\right \rangle-2\left\langle\theta_{k+J}^{2}\right\rangle+2\left\langle\theta_{k} \theta_{k+J}\right\rangle+\left\langle\theta_{k+J}^{2}\right\rangle\ ,$ (85) $$

which holds for any integer $J$. We have $\left\langle\theta_{k+J}^{2}\right\rangle=\left\langle\theta_{k}^{2}\right\rangle$ since the expectation value cannot depend on $k$. Additionally, by definition we have $v_{k}=\theta_{k}-\theta_{k-1}$ which yields

$$ $\displaystyle\theta_{k}-\theta_{k+J}=\sum_{i=1}^{J}v_{k+i}\ .$ (86) $$

Therefore, we can rewrite Equation 85 as follows

$$ $\displaystyle 2\sigma_{\theta}^{2}$ $\displaystyle=2\left\langle\theta_{k}\theta_{k+J}\right\rangle+\sum_{i,j=1}^{J }\langle v_{k+i}v_{k+j}\rangle$ $\displaystyle=2\left\langle\theta_{k}\theta_{k+J}\right\rangle+\sum_{i,j=1}^{J }\langle v_{k}v_{k+j-i}\rangle$ $\displaystyle=2\left\langle\theta_{k}\theta_{k+J}\right\rangle+\sum_{m=0}^{J-1 }\sum_{n=-m}^{m}\langle v_{k}v_{k+n}\rangle\ ,$ (87) $$

where we first shifted the index within the expectation value and then restructured the sum by defining $m\coloneqq\max(i,j)-1$ and $n\coloneqq j-i$. We now take the limit of $J\rightarrow\infty$ and because of Equation 83 and the assumption of a finite $\sigma_{\theta}^{2}$ we have

$$ $\displaystyle\sum_{m=0}^{\infty}\sum_{n=-m}^{m}\langle v_{k}v_{k+n}\rangle$ $\displaystyle<\infty$ (88) $\displaystyle\Rightarrow\sum_{n=-\infty}^{\infty}\langle v_{k}v_{k+n}\rangle$ $\displaystyle=0\ .$ (89) $$

We note that $\langle v_{k}v_{k+n}\rangle=\langle v_{k}v_{k-n}\rangle$ because we can shift the index, and the two factors commute. Substituting this relation into Equation 89 yields

$$ $\displaystyle\sum_{n=1}^{\infty}\langle v_{k}v_{k+n}\rangle$ $\displaystyle=-\frac{1}{2}\langle v_{k}v_{k}\rangle$ $\displaystyle=-\frac{1}{2}\sigma_{v}^{2}\ .$ (90) $$

For the second part of the proof we will start again with $v_{k}=\theta_{k}-\theta_{k-1}$ and the following sum

$$ $\displaystyle\sum_{n=1}^{m}n\langle v_{k}v_{k+n}\rangle$ $\displaystyle=\sum_{n=1}^{m}n\bigl{(}2\langle\theta_{k}\theta_{k+n}\rangle- \langle\theta_{k-1}\theta_{k+n}\rangle-\langle\theta_{k}\theta_{k+n-1}\rangle \bigr{)}$ $\displaystyle=-\langle\theta_{k}\theta_{k}\rangle+\langle\theta_{k}\theta_{k+m }\rangle+m\bigl{(}\langle\theta_{k}\theta_{k+m}\rangle-\langle\theta_{k}\theta _{k+m+1}\rangle\bigr{)}\ ,$ (91) $$

where nearly all terms cancel each other again due to the fact that we can shift the index within the expectation value. By taking the limit $m\rightarrow\infty$ and using the assumptions (ii) and (iii) (Equations 83 and 84) we have

$$ $\displaystyle\sum_{n=1}^{\infty}n\langle v_{k}v_{k+n}\rangle=-\langle\theta_{k }\theta_{k}\rangle\ .$ (92) $$

Finally, by dividing Equation 92 by Equation 90 we arrive at the final expression

$$ $\displaystyle\frac{2\sigma_{\theta,i}^{2}}{\sigma_{v,i}^{2}}=\frac{\sum_{n=1}^ {\infty}n\,\langle v_{k}v_{k+n}\rangle}{\sum_{n=1}^{\infty}\langle v_{k}v_{k+n }\rangle}\ .$ (93) $$

## Appendix E Calculation of the Noise Autocorrelation

We want to calculate the autocorrelation function of epoch-based SGD for a fixed weight vector $\bm{\mathbf{\theta}}$ and under the assumption that the total number of examples is an integer multiple of the number of examples per batch. For that we repeat the following definitions:

$$ $\displaystyle\bm{\mathbf{\delta g}}_{k}(\bm{\mathbf{\theta}})$ $\displaystyle\coloneqq\frac{1}{S}\sum_{n\in\mathcal{B}_{k}}\bm{\mathbf{\nabla} }\bigl{(}l(\bm{\mathbf{\theta}},x_{n})-L(\bm{\mathbf{\theta}})\bigr{)}$ (94) $\displaystyle\mathcal{B}_{k}=\{n_{1},...,n_{S}\}$ $\displaystyle\dots\text{ batch of step $k$, sampeling without replacement within epoch}$ (95) $\displaystyle n_{j}\in\{1,\dots,N\}$ (96) $\displaystyle N$ $\displaystyle\dots\text{ total number of examples}$ (97) $\displaystyle S$ $\displaystyle\dots\text{ number of examples per batch}$ (98) $$

We can rewrite the noise terms as follows:

$$ $\displaystyle\bm{\mathbf{\delta g}}_{k}(\bm{\mathbf{\theta}})$ $\displaystyle=\frac{1}{S}\sum_{n\in\mathcal{B}_{k}}\bm{\mathbf{\delta g_{e}}}( n,\bm{\mathbf{\theta}})$ $\displaystyle=\frac{1}{S}\sum_{n=1}^{N}\bm{\mathbf{\delta g_{e}}}(n,\bm{ \mathbf{\theta}})s_{k}^{n}$ (99) $\displaystyle s_{k}^{n}$ $\displaystyle\coloneqq\mathbf{1}_{\mathcal{B}_{k}}(n)$ $\displaystyle=\begin{cases}1\ \text{ if }~{}n\in\mathcal{B}_{k}\\ 0\ \text{ if }~{}n\notin\mathcal{B}_{k}\end{cases}$ (100) $\displaystyle\bm{\mathbf{\delta g_{e}}}(n,\bm{\mathbf{\theta}})$ $\displaystyle\coloneqq\bm{\mathbf{\nabla}}\bigl{(}l(\bm{\mathbf{\theta}},x_{n} )-L(\bm{\mathbf{\theta}})\bigr{)}\ .$ (101) $$

Let $h\geq 0$ be fixed. The correlation matrix can be expressed as

$$ $\displaystyle\textrm{cov}\bigl{(}\vec{\delta g}_{k}(\vec{\theta}),\vec{\delta g }_{k+h}(\vec{\theta})\bigr{)}$ $\displaystyle=\mathbb{E}\left[\bm{\mathbf{\delta g}}_{k}(\bm{\mathbf{\theta}}) \bm{\mathbf{\delta g}}_{k+h}(\bm{\mathbf{\theta}})^{\top}\right]$ $\displaystyle=\frac{1}{S^{2}}\sum_{n,\tilde{n}=1}^{N}\bm{\mathbf{\delta g_{e}} }(n,\bm{\mathbf{\theta}})\bm{\mathbf{\delta g_{e}}}(\tilde{n},\bm{\mathbf{ \theta}})^{\top}\mathbb{E}\left[s_{k}^{n}s_{k+h}^{\tilde{n}}\right]\ .$ (102) $$

The expectation value of $s_{k}^{n}=\mathbf{1}_{\mathcal{B}_{k}}(n)$ is the probability that example $n$ is part of batch $k$. Because every example is equally likely to appear in a given batch, this probability is equal to $S/N$.

$$ $\displaystyle\mathbb{E}\left[s_{k}^{n}\right]$ $\displaystyle=\textrm{P}\!\left(s_{k}^{n}=1\right)$ $\displaystyle=\frac{S}{N}\ .$ (103) $$

Similarly we can calculate the desired correlation:

$$ $\displaystyle\mathbb{E}\left[s_{k}^{n}s_{k+h}^{\tilde{n}}\right]$ $\displaystyle=\textrm{P}\!\left(s_{k}^{n}=1,s_{k+h}^{\tilde{n}}=1\right)$ $\displaystyle=\textrm{P}\!\left(s_{k}^{n}=1\right)\textrm{P}\!\left(s_{k+h}^{ \tilde{n}}=1\mid s_{k}^{n}=1\right)$ $\displaystyle=\frac{S}{N}\,\textrm{P}\!\left(s_{k+h}^{\tilde{n}}=1\mid s_{k}^{ n}=1\right)\ .$ (104) $$

The last term can be split up into different probabilities for different values of h. We can also distinguish the case where the two steps $k$ and $k+h$ are within the same epoch ($\textrm{ep}(k)=\textrm{ep}(k+h)$) or in different epochs ($\textrm{ep}(k)\neq\textrm{ep}(k+h)$).

$$ $\displaystyle\textrm{P}\!\left(s_{k+h}^{\tilde{n}}=1\mid s_{k}^{n}=1\right)$ $\displaystyle=\delta_{h,0}\cdot\textrm{P}\!\left(s_{k}^{\tilde{n}}=1\mid s_{k} ^{n}=1\right)+(1-\delta_{h,0})\,\cdot$ $\displaystyle\quad\Bigl{[}\textrm{P}\bigl{(}\textrm{ep}(k)=\textrm{ep}(k+h) \bigr{)}\textrm{P}\!\left(s_{k+h}^{\tilde{n}}=1\mid s_{k}^{n}=1,\textrm{ep}(k) =\textrm{ep}(k+h)\right)+$ $\displaystyle\quad\ \textrm{P}\bigl{(}\textrm{ep}(k)\neq\textrm{ep}(k+h)\bigr{ )}\textrm{P}\!\left(s_{k+h}^{\tilde{n}}=1\mid s_{k}^{n}=1,\textrm{ep}(k)\neq \textrm{ep}(k+h)\right)\Bigr{]}\ .$ (105) $$

The first term of the right hand side of Equation 105 is the probability that a given example occurs in a batch, assuming that we already know one of the examples of that batch.

$$ $\displaystyle\textrm{P}\!\left(s_{k}^{\tilde{n}}=1\mid s_{k}^{n}=1\right)$ $\displaystyle=\textrm{P}\!\left(s_{k}^{n}=1\mid s_{k}^{n}=1\right)\cdot\, \delta_{n,\tilde{n}}+\textrm{P}\!\left(s_{k}^{\tilde{n}}=1\mid s_{k}^{n}=1,n \neq\tilde{n}\right)\,(1-\delta_{n,\tilde{n}})$ $\displaystyle=1\cdot\,\delta_{n,\tilde{n}}+\frac{S-1}{N-1}\,(1-\delta_{n, \tilde{n}})$ $\displaystyle=\frac{N-S}{N-1}\,\delta_{n,\tilde{n}}+\textrm{const.}$ (106) $$

The second term of Equation 105 is multiplied by $(1-\delta_{h,0})$. Therefore, we assume $h\geq 1$ for the following argument. That is, we want to know the probabilities under the assumption that we are comparing examples from different batches. If the two batches are still from the same epoch, examples cannot repeat as the total number of examples is an integer multiple of the number of examples per batch and because of that every example is shown only once per epoch. Therefore, for $h\geq 1$ holds:

$$ $\displaystyle\textrm{P}\!\left(s_{k+h}^{\tilde{n}}=1\mid s_{k}^{n}=1,\textrm{ ep}(k)=\textrm{ep}(k+h)\right)$ $\displaystyle=0\cdot\delta_{n,\tilde{n}}+\frac{S}{N-1}\,(1-\delta_{n,\tilde{n}})$ $\displaystyle=-\frac{S}{N-1}\delta_{n,\tilde{n}}+\textrm{const.}$ (107) $$

If we consider batches from different epochs, the probability becomes independent of the given examples:

$$ $\displaystyle\textrm{P}\!\left(s_{k+h}^{\tilde{n}}=1\mid s_{k}^{n}=1,\textrm{ ep}(k)\neq\textrm{ep}(k+h)\right)$ $\displaystyle=\frac{S}{N}$ $\displaystyle=\textrm{const.}$ (108) $$

Lastly, we need to know the probability that two given batches $k$ and $k+h$ are from the same epoch:

$$ $\displaystyle\textrm{P}\bigl{(}\textrm{ep}(k)=\textrm{ep}(k+h)\bigr{)}=\mathbf {1}_{\{1,\dots,M\}}(h)\,\frac{M-h}{M}\ ,$ (109) $$

where $M=N/S$ is again the number of batches per epoch.

We can now combine all derived probabilities and arrive at the following relation:

$$ $\displaystyle\mathbb{E}\left[s_{k}^{n}s_{k+h}^{\tilde{n}}\right]$ $\displaystyle=\delta_{n,\tilde{n}}\,\frac{S}{N}\frac{N-S}{N-1}\left(\delta_{h, 0}-\mathbf{1}_{\{1,\dots,M\}}(h)\frac{S}{N-S}\frac{M-h}{M}\right)+\textrm{ const.}$ $\displaystyle=\delta_{n,\tilde{n}}\,S^{2}\left(\frac{1}{S}-\frac{1}{N}\right) \frac{}{}\frac{1}{N-1}\left(\delta_{h,0}-\mathbf{1}_{\{1,\dots,M\}}(h)\frac{M- h}{M\,(M-1)}\right)+\textrm{const.}$ (110) $$

If we now also consider negative values for h, the expression depends only on the absolute value of h due to symmetry.

By using the following two helpful relations:

$$ $\displaystyle\sum_{n,\tilde{n}=1}^{N}\bm{\mathbf{\delta g_{e}}}(n,\bm{\mathbf{ \theta}})\bm{\mathbf{\delta g_{e}}}(\tilde{n},\bm{\mathbf{\theta}})^{\top}\, \delta_{n,\tilde{n}}$ $\displaystyle=\sum_{n=1}^{N}\bm{\mathbf{\delta g_{e}}}(n,\bm{\mathbf{\theta}}) \bm{\mathbf{\delta g_{e}}}(n,\bm{\mathbf{\theta}})^{\top}$ $\displaystyle\eqqcolon(N-1)\,\bm{\mathbf{C}}_{0}(\bm{\mathbf{\theta}})\ ,$ (111) $\displaystyle\sum_{n,\tilde{n}=1}^{N}\bm{\mathbf{\delta g_{e}}}(n,\bm{\mathbf{ \theta}})\bm{\mathbf{\delta g_{e}}}(\tilde{n},\bm{\mathbf{\theta}})^{\top}\cdot 1$ $\displaystyle=\left(\sum_{n=1}^{N}\bm{\mathbf{\delta g_{e}}}(n,\bm{\mathbf{ \theta}})\right)\left(\sum_{n=1}^{N}\bm{\mathbf{\delta g_{e}}}(n,\bm{\mathbf{ \theta}})^{\top}\right)$ $\displaystyle=\bm{\mathbf{\nabla}}\bigl{(}L(\bm{\mathbf{\theta}})-L(\bm{ \mathbf{\theta}})\bigr{)}\bm{\mathbf{\nabla}}^{\top}\bigl{(}L(\bm{\mathbf{ \theta}})-L(\bm{\mathbf{\theta}})\bigr{)}$ $\displaystyle=0\ ,$ (112) $$

we can insert the expectation value $\mathbb{E}\left[s_{k}^{n}s_{k+h}^{\tilde{n}}\right]$ into Equation 102 and arrive at the final expression:

$$ $\displaystyle\textrm{cov}\bigl{[}\vec{\delta g}_{k}(\vec{\theta}),\vec{\delta g }_{k+h}(\vec{\theta})\bigr{]}$ $\displaystyle=\textrm{cov}\bigl{[}\vec{\delta g}_{k}(\vec{\theta}),\vec{\delta g }_{k}(\vec{\theta})\bigr{]}\cdot\left(\delta_{h,0}-\mathbf{1}_{\{1,...,M\}}(|h |)\frac{M-|h|}{M(M-1)}\right),$ (113) $\displaystyle\textrm{cov}\bigl{[}\vec{\delta g}_{k}(\vec{\theta}),\vec{\delta g }_{k}(\vec{\theta})\bigr{]}$ $\displaystyle=\left(\frac{1}{S}-\frac{1}{N}\right)\bm{\mathbf{C}}_{0}(\bm{ \mathbf{\theta}})\ .$ (114) $$

## Appendix F Comparison with principal component analysis

Our approach to analysis sets itself apart from that of Feng & Tu [] principally in the selection of the basis $\{\vec{p_{i}},,i=1,\dots,d\}$ used for examining the weights. While they employ the principal components of the weight series - the eigenvectors of $\vec{\Sigma}$ - we use the eigenvectors of the hessian matrix $\vec{H}(\vec{\theta_{K}})$ computed at the beginning of the analysis period.

This choice enables us to directly create plots of variances and correlation time against the hessian eigenvalue for each corresponding direction. Feng & Tu devised a landscape-dependent flatness parameter $F_{i}$ for every direction $\vec{p_{i}}$. However, with the assistance of the second derivative $F_{i}\approx\left(\partial^{2}L(\vec{\theta})/\partial\theta_{i}^{2}\right)^{-
\frac{1}{2}}$, where $\theta_{i}=\vec{\theta}\cdot\vec{p_{i}}$, this parameter can be approximated, provided this second derivative retains a sufficiently positive value. Hence, in the eigenbasis of the hessian matrix, the flatness parameter can be approximated as $F_{i}\approx\lambda_{i}^{-\frac{1}{2}}$, facilitating comparability between our analysis and that of Feng & Tu.

The principal component basis, as used by Feng & Tu, holds a distinct advantage. For our analysis, we needed to eliminate the near-linear trajectory of the weights by deducting the mean velocity. However, in Feng & Tu’s analysis, this movement is automatically subsumed in the first principal component due to its pronounced variance. Hence, there’s no necessity for additional subtraction of this drift in the weight covariance eigenbasis.

Yet, the weight covariance eigenbasis has a significant shortcoming: it yields artifacts. This is because $\vec{\Sigma}$ is calculated as an average over a finite data set, skewing its eigenvalues from the anticipated distribution. Consequently, the resultant eigenvectors may not align perfectly with the expected ones. This issue is further exacerbated due to the high dimensionality of the underlying space.

The artifact issue becomes evident in Figure 3, which displays synthetic data generated through stochastic gradient descent within an isotropic quadratic potential coupled with isotropic noise. With 2,500 dimensions, the model mirrors the scale of a layer in the fully connected neural network that Feng & Tu investigated. The weight series comprises 12,000 steps, which correspond to ten epochs of training this network. Analyzing this data with the weight covariance eigenbasis seemingly suggests anisotropic variance and correlation time. However, if the data is inspected without any basis change, both the variance and correlation time appear isotropically distributed as anticipated.

To navigate around this key issue associated with the eigenbasis of $\vec{\Sigma}$, we adopted the eigenvectors of the Hessian matrix. Unlike $\vec{\Sigma}$, the Hessian is not computed as an average over update steps but can, in theory, be precisely calculated for any given weight vector. Consequently, the Hessian matrix does not suffer from finite size effects. The difference between these two bases for actual data is visible in Figure 4. Here, we analyzed only the weights of the first convolutional layer of the LeNet from the main text to ensure comparability with Feng & Tu’s results. In this specific comparison, the network was trained without weight decay. Due to this and the fact that we are only investigating the weights of one layer, $\lambda_{\textrm{cross}}$ is significantly larger than all Hessian eigenvalues. As a result, when analyzing in the eigenbasis of the Hessian matrix related to this layer, both the variance and the correlation time align well with the prediction for smaller Hessian eigenvalues.

However, analyzing in the eigenbasis of the weight covariance matrix, the correlation time appears heavily dependent on the second derivative of the loss in the given direction. Additionally, the relationship between the weight variance and the second derivative shifts and more closely aligns with Feng & Tu’s results as the power law exponent is significantly larger than one. The first principal component, which Feng & Tu referred to as the drift mode, stands out due to its unusually long correlation time. This is to be expected, as this is the direction in which the weights are moving at an approximately constant velocity.

Figure: Figure 3: Comparison of weight and velocity fluctuations for synthetic data analyzed in two different bases. We define the variance in weight, $\sigma_{\theta,i}^{2}$, as $\vec{p_{i}}^{\top}\vec{\Sigma}\vec{p_{i}}$, and the variance in velocity, $\sigma_{v,i}^{2}$, as $\vec{p_{i}}^{\top}\vec{\Sigma_{v}}\vec{p_{i}}$. The correlation time, $\tau_{i}$, is given by $2\sigma_{\theta,i}^{2}/\sigma_{v,i}^{2}$. The synthetic data was generated by simulating SGD for 12,000 steps within a 2,500-dimensional space featuring an isotropic quadratic potential and isotropic noise. In the original basis analysis, both the variance and correlation time, as expected, retain isotropy. However, when the analysis is conducted in the eigenbasis of the weight covariance matrix, a pronounced anisotropy emerges.
Refer to caption: extracted/5732104/Figs/artificial_data/var_theta.jpeg

Figure: Figure 4: Comparison of weight and velocity variances for all 450 weights of the first convolutional layer of the LeNet, as discussed in the main text, analyzed in two different bases. In order to facilitate a more directly comparable analysis to Feng & Tu [], the network was trained without weight decay for this specific analysis and the analysis period was limited to 10 epochs, as opposed to the usual 20 epochs. The columns represent different bases: for the left column $\vec{p_{i}}$ are the eigenvectors of $\vec{\Sigma}$ and for the right column $\vec{p_{i}}$ are the eigenvectors of $\vec{H}$. The mean velocity was subtracted in the right column. The rows illustrate the weight and velocity variance, $\sigma_{\theta,i}^{2}=\vec{p_{i}}^{\top}\vec{\Sigma}\vec{p_{i}}$, $\sigma_{v,i}^{2}=\vec{p_{i}}^{\top}\vec{\Sigma_{v}}\vec{p_{i}}$ (top row), and the correlation time $\tau_{i}=2\sigma_{\theta,i}^{2}/\sigma_{v,i}^{2}$ (bottom row). The second derivative of the corresponding direction is depicted on the x-axis, $\lambda_{i}=\vec{p_{i}}^{\top}\vec{H}\vec{p_{i}}$. The top row solid lines indicate fit regions for a linear fit. For the $\vec{H}$ eigenbasis, the respective exponent of the power law relation is $1.018\pm 0.008$ for weight variance and $1.017\pm 0.002$ for velocity variance with a $2\sigma$-error. For the $\vec{\Sigma}$ eigenbasis, the corresponding exponent is $1.537\pm 0.012$ for weight variance and $1.134\pm 0.002$ for velocity variance.
Refer to caption: extracted/5732104/Figs/lyr_0/var_pca.jpeg

## Appendix G Hessian Eigenvalue Density

Figure: Figure 5: The distribution of the approximated 5,000 Hessian eigenvalues of the LeNet discussed in the main text. The inset shows that the smallest approximated eigenvalue has a magnitude of about 0.005.
Refer to caption: extracted/5732104/Figs/ICML_Hessian_density_inset_LeNet.jpeg

## Appendix H Drawing with replacement

To confirm that the results obtained are indeed affected by the correlations present in SGD noise, due to the epoch-based learning strategy, we reapply the analysis described in the main text. In this instance, however, we deviate from our previous method of choosing examples for each batch within an epoch without replacement. Instead, we select examples with replacement from the complete pool of examples for every batch. This modification during the analysis period allows a more complete assessment of the impact of correlations on the derived results.

Figure 6 offers clear visual proof that when examples are selected with replacement, the previously noted anti-correlations within the SGD noise vanish. This observation confirms our hypothesis that the anti-correlations mentioned in the main text are indeed an outcome of the epoch-based learning technique. Consequently, we can predict that this change will influence the behaviour of the weight and velocity variance. As previously discussed, the theoretical results we have achieved for Hessian eigenvectors with eigenvalues exceeding $\lambda_{\textrm{cross}}$ conform to what one would predict in the absence of any correlation within the noise. Therefore, when examples are drawn with replacement, we anticipate the weight variance to be isotropic in all directions, while the velocity variance should remain unchanged.

Figure: Figure 6: Autocorrelations of the SGD noise compared for drawing examples without replacement (left) and with replacement (right).
Refer to caption: extracted/5732104/Figs/with_replacement/correlation_NR.jpeg

Upon reviewing Figure 7, it is clear that the velocity variance stays unchanged as predicted. However, while the weight variance remains constant for a broader subset of Hessian eigenvalues, it reduces for extremely small eigenvalues. Likewise, the correlation time is still limited for these minuscule Hessian eigenvalues. These deviations can be attributed to the finite time frame of the analysis period, comprising 20,000 update steps. This limited time window sets a cap on the maximum correlation time, consequently leading to a decreased weight variance for these small Hessian eigenvalues. Despite this, it is noteworthy that this maximum correlation time is still roughly one order of magnitude longer than the maximum correlation time induced by the correlations arising from the epoch-based learning approach.

Figure: Figure 7: Relationship between Hessian eigenvalues and the variances of weights and velocities, as well as correlation times. For the left column the examples are drawn in epochs without replacement and for the right column the examples are drawn with replacement.
Refer to caption: extracted/5732104/Figs/with_replacement/variance_NR.jpeg

## Appendix I Loss and Accuracy during Training

Figure: Figure 8: The evolution of the loss (left) and accuracy (right) during training of LeNet described in the main text. The statistics are shown for both training and test set. For the first 100 epochs, the exponential learning rate decay was used, and for the last 20 epochs, the learning rate was fixed at the final value of the exponential decay.
Refer to caption: extracted/5732104/Figs/added_appendix/ICML_loss_evolution.jpeg

## Appendix J Testing different Hyperparameters

In this section we examine the dependence of the theoretical predictions on the three hyperparameters learning rate $\eta$, momentum $\beta$ and batch size $S$. For this, we train the LeNet again for 100 epochs, using an exponential learning rate schedule that reduces the learning rate by a factor of $0.98$ every epoch and afterwards we perform the numerical analysis as described in the main text.

However, we now train the network several times, always varying one of the hyperparameters while keeping the other two fixed. If not varied, the momentum was set to 0.90 and the batch size was set to 64. To ensure that training is always successful and 100% training accuracy is achieved, the initial learning rate was set to 0.005 when the batch size is varied and to 0.02 when the momentum is varied. Five different values are examined for each hyperparameter. To investigate the dependencies on the learning rate, the values 0.005, 0.01, 0.02, 0.03, and 0.04 were used for training. For momentum, the values 0.00, 0.50, 0.75, 0.90, and 0.95 were examined, and for batch size, the values 32, 50, 64, 100, and 128 were examined.

In addition, the training was repeated for five different seeds for each hyperparameter combination in order to obtain reliable results. This results in a relatively high computational cost. To reduce this, for the analysis in this section the weight and velocity variances are examined only in the subspace of the 2,000 largest Hessian eigenvalues and associated eigenvectors.

Figure 9 shows as an example the weight and velocity variances as well as the correlation times for different values of the batch size for one training seed each. It can be seen that the theory is not only valid for the hyperparameter combination from the previous section, but is also generally applicable for different hyperparameters. In particular, we see that there is still good agreement with the theory even if the strictly necessary condition for the theoretical derivation of the noise autocorrelation, that the number of batches per epoch $M=N/S$ is an integer, is not met.

Figure: Figure 9: Testing the LeNet training with different hyperparameters. Here the relationship between the Hessian eigenvalues and the variances and correlation times for varying batch size is shown as an example. The momentum was set to 0.90 and the initial learning rate was set to 0.005.
Refer to caption: extracted/5732104/Figs/added_appendix/hyperparameter_single_seed_b.jpeg

To further examine the predictions of the theory for the hyperparametric dependencies, we now focus on the two quantities of the maximum correlation time $\tau_{\rm SGD}$ and the Hessian eigenvalue crossover value $\lambda_{\rm cross}$ and recall the theoretical predictions for these quantities:

$$ $\displaystyle\tau_{\textrm{SGD}}$ $\displaystyle=\frac{N}{3S}\frac{1+\beta}{1-\beta}\ ,$ (115a) $\displaystyle\lambda_{\textrm{cross}}$ $\displaystyle=\frac{3S(1-\beta)}{\eta N}\ ,$ (115b) $$

where $N$ is the number of examples in the training data set.

For the evaluation of the dependence of these variables on the hyperparameters, they were determined as follows for the various hyperparameter combinations using the data from the respective correlation time plot. For the maximum correlation time $\tau_{\textrm{SGD}}$, the average of all correlation times was taken for which the corresponding Hessian eigenvalue is smaller than the theoretical crossover value. However, the result for the numerically determined maximum correlation time is not significantly different when simply taking the average of all determined correlation times for a hyperparameter combination, since only very few Hessian eigenvalues are larger than the crossover value.

For the numerical determination of the crossover value $\lambda_{\rm cross}$, a linear function was first fitted to the correlation times of the 20 largest Hessian eigenvalues in the log-log plot of the correlation times against the Hessian eigenvalues. In this region of the first 20 values, the correlation time is always clearly dependent on the Hessian eigenvalue and does not yet belong to the region of constant correlation times. The intersection of the fitted line with the numerically determined maximum correlation time $\tau_{\textrm{SGD}}$ is then taken as the crossover value $\lambda_{\rm cross}$. If the numerically determined correlation times follow the theory exactly, then the correlation times determined in this way for $\tau_{\textrm{SGD}}$ and $\lambda_{\rm cross}$ would also follow the theory accurately.

And indeed, Figure 10 shows a good agreement between the theory and the numerically determined values, although it should be noted that the deviations are larger than the random fluctuations between the different seeds.

Figure: Figure 10: The empirically determined maximum correlation time, $\tau_{\textrm{SGD}}$, and the empirically determined crossover value $\lambda_{\rm cross}$ for the training of the LeNet with different hyperparameters, averaged over five different seeds for each set of hyperparameters. The predictions of our theory are shown as a solid line. The fluctuations between different seeds are smaller than the marker size and are therefore not included in the figure.
Refer to caption: extracted/5732104/Figs/added_appendix/hyperparameter_dependence_thesis.jpeg

## Appendix K Test Accuracy and the Hessian Eigendecomposition of the Weight Vector

Recent research [, ] suggests that the characteristics of a minimum in a neural network’s loss landscape are largely determined by the large Hessian eigenvalues and their corresponding eigendirections. However, our findings indicate that flat directions, which correspond to smaller eigenvalues, should not be overlooked. When we decompose a network’s weight vector in the Hessian eigenbasis, it becomes evident that both the projections onto high curvature directions and those onto relatively flat directions are important.

In Figure 11, we demonstrate this with the LeNet model. We analyzed the model’s test accuracy in relation to the exclusion of projections onto Hessian eigenvectors, starting with those possessing the largest eigenvalues. After calculating the top 10,000 Hessian eigenvalues and eigenvectors out of a total of 137,000, we observed that omitting the projection onto the top 35 eigenvectors (those above the crossover value $\lambda_{\rm cross}$) resulted in a decline in test accuracy from 63% to 54%. Discarding the top 5,000 eigenvectors, as analyzed in our main study, further reduced accuracy to 44%. It was only after excluding around 7,000 eigenvectors that the accuracy plummeted to 10%, akin to random guessing. This underscores the significance of the weight vector’s orientation relative to flat Hessian eigendirections for maintaining test accuracy, highlighting the potential impact of weight fluctuations in these directions.

The pronounced accuracy drop observed after removing around 6,000 eigenvectors may be attributed to the unexpectedly large projection of the weight vector onto these eigenvectors, which is apparent from Fig. Figure 11.

Figure: Figure 11: The top panel shows the test accuracy of the LeNet discussed in the main text and how it changes as we cumulatively discard projections onto more and more Hessian eigenvectors, starting with those with the largest eigenvalues. There is a sharp drop in accuracy for the first few eigenvectors, and then only after more than 6,000 discarded projections out of a total of 137,000 does the accuracy drop rapidly down to 10%, which is equivalent to random guessing. For comparison, the middle panel shows the corresponding Hessian eigenvalues and compares them to the crossover value $\lambda_{\rm cross}$, where the lower than expected weight variances hold only for eigendirections with eigenvalues smaller than this crossover value. It is in the range of eigenvalues significantly smaller than the crossover value that the accuracy drops sharply to 10%. The bottom panel shows the decomposition of the weight vector $\bf\theta$ in the Hessian eigenbasis. One explanation for the rapid drop in accuracy could be that the projection of the weights onto the corresponding eigenvectors is considerably larger than one would expect if the weight vector was a random vector. The expected $2\sigma$ interval for a random vector is displayed in red.
Refer to caption: extracted/5732104/Figs/added_appendix/ICML_Hessian_accuracy.jpeg

## Appendix L Different Network Architecture

To further confirm our theoretical predictions within trained networks, in this section we turn to a more modern architecture. Instead of the previously used LeNet network, we examine the ResNet-20 network []. It is a convolutional network with significantly more convolutional layers than LeNet. It also uses residual blocks with residual connections, which allows for deeper network structures. As our loss function, we again employed Cross Entropy, along with an L2 regularization with a prefactor of $10^{-4}$ and we did not use batch normalization. The number of layers, which is already indicated in the name with 20, is significantly higher than in the LeNet with just five layers. With approximately 272,000 parameters, the ResNet-20 also has significantly more parameters and the computational cost is significantly higher.

Therefore, in this section we limit ourselves to examining the weight and velocity variances in the subspace of the 400 largest Hessian eigenvalues and associated eigenvectors. The network was trained with SGD for 100 epochs using the same exponential learning rate schedule as before, with a learning rate of $5\cdot 10^{-3}$, a momentum parameter of $0.9$, and a minibatch size of $50$. This setup achieves 100% training accuracy and 73% testing accuracy. In Figure 12 one can observe a good agreement between the theory predictions for the variances and correlation times and the numerical observations.

Figure: Figure 12: Relationship between Hessian eigenvalues and the variances and correlation times for a ResNet-20 trained on CIFAR10. The mean velocity of the weight trajectory was subtracted. The solid lines in the left panel indicate the regions utilized for a linear fit. The exponents resulting from the power law relationship are $0.820\pm 0.099$ for weight variance and $0.965\pm 0.005$ for velocity variance, with a $2\sigma$-error. The analysis was performed for the 400 largest Hessian eigenvalues and corresponding eigendirections.
Refer to caption: extracted/5732104/Figs/added_appendix/ResNet_variance.jpeg

## Appendix M Effects of Non-Commutativity

The exact variance equation of Theorem 4.3 described in Section 4.3 and the calculation shown in Appendix C are to some extent still valid even if the previously mentioned assumption $[{\bf C},{\bf H}]\neq 0$ is not given. In particular, the result for the correlation time holds independently of the commutativity assumption. However, the calculated weight and velocity variances are no longer eigenvalues of the corresponding covariance matrices, but only the variances in the directions of the chosen eigenvector of the Hessian matrix.

Assuming that ${\bf C}$ and ${\bf H}$ do not necessarily commute, we can still project the update equations onto an arbitrary Hessian eigenvector $\bm{\mathbf{p}}_{i}$ with eigenvalue $\lambda_{i}$, which gives us

$$ $\displaystyle v_{k,i}$ $\displaystyle=-\eta\lambda_{i}\theta_{k-1,i}+\beta v_{k,i}-\eta\bm{\mathbf{p}} _{i}\cdot\bm{\mathbf{\delta g}}_{k}\ ,$ (116a) $\displaystyle\theta_{k,i}$ $\displaystyle=(1-\eta\lambda_{i})\theta_{k-1,i}+\beta v_{k,i}-\eta\bm{\mathbf{ p}}_{i}\cdot\bm{\mathbf{\delta g}}_{k}\ .$ (116b) $$

Since $\bm{\mathbf{p}}_{i}\cdot\bm{\mathbf{\delta g}}_{k}$ is independent of weights and velocity, these two equations are decoupled for each individual Hessian eigenvector. As $\bm{\mathbf{p}}_{i}\cdot\bm{\mathbf{\delta g}}_{k}$ still follows the proposed anti-correlation, the calculations of Appendix C can be performed similarly. Therefore, the theory prediction for the correlation time $\tau_{i}$, defined as the ratio between the weight and the velocity variance in the eigendirection $\bm{\mathbf{p}}_{i}$ of the Hessian, is still valid. However, the weight and the velocity variance in the eigendirection $\bm{\mathbf{p}}_{i}$ are no longer eigenvalues of the covariance matrices $\bf\Sigma$ and $\bf\Sigma_{v}$ if $\bm{\mathbf{p}}_{i}$ is not also an eigenvector of ${\bf C}$.

To better understand the consequences of a situation where ${\bf C}$ and ${\bf H}$ do not commute, we perform an SGD simulation in an artificial quadratic loss landscape where we introduce epoch-based noise with anti-correlations as described in Section 4.1. However, we choose the Hessian matrix $\bf H$ of the loss and the gradient noise covariance matrix $\bf C$ to be non-commuting. The results are shown in Figure 13.

When $\bf H$ and $\bf C$ are approximately proportional to each other, to an extent reported in the literature [], the results still show good agreement with our theory and there are only small deviations for very small Hessian eigenvalues. If there is no similarity between $\bf H$ and $\bf C$, then the results for the variances no longer follow our theory, but the prediction for the correlation time $\tau_{i}$ still shows good agreement with the numerical data, as expected.

Figure: Figure 13: Variances and correlation times of an SGD trajectory in an artificial quadratic potential with epoch-based noise. The quantities are examined in the Hessian eigenbasis and plotted against the corresponding eigenvalue. For the top two panels, the covariance matrix of the noise $\vec{C}$ is equal to the Hessian $\vec{H}$ plus a random matrix $\vec{M}$ of small magnitude ($\vec{M}=\frac{1}{d}\vec{X}\vec{X}^{\top}$ with $\vec{X}\in\mathbb{R}^{d\times d}$ and $X_{ij}\sim\mathcal{N}(0,\,\sigma^{2})$, $\sigma=0.02$, $d=500$). The cosine similarity between $\vec{C}$ and $\vec{H}$ is 0.96 which is a value that has been found empirically in real networks before []. Here, our theory for the variances still holds except below a certain small eigenvalue. For the bottom two panels, the Hessian and the covariance matrix are two independent random matrices. There is no longer a relationship between the Hessian eigenvalue and the velocity covariance. However, the theoretical prediction for the correlation time still holds.
Refer to caption: extracted/5732104/Figs/added_appendix/mix_H_C_variance.jpeg

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: 1. In line 319, it is claimed that if these conditions are not met, the weight fluctuations would diverge. How to see that?

2. Is the momentum necessary in your arguments? Do the results still hold for vanilla SGD?
soundness: 2
strengths: 1. Understanding SGD noise is an important topic in the optimization. This paper analyzes the SGD noise by connecting the covariance matrix to the Hessian of the loss and further connects the flat directions to the weight variance. I believe these insights are nice.

2. The paper is well-written and the logic is easy to follow. The experiments are nicely done to verify the theory.
summary: This paper studies the correlation of gradient noise in the later training phase of SGD. They show that SGD noise will be anti-correlated over time, assuming the weights do not change. They also show that weight variance is small in flat directions, i.e., eigendirection corresponding to small eigenvalues of loss hessian. This result is obtained by assuming quadratic loss. They further implement experiments to verify their theoretical results.
weaknesses: 1. I am not sure if it is reasonable to assume the fixed weights to analyze the gradient noise. It is almost trivial to obtain Theorem 4.1 under such an assumption since some probability arguments are sufficient. Additionally, if the weights are fixed, or almost fixed at the end of training, the effect of gradient noise is minimal. It should make more sense to understand the gradient noise at the beginning of training, where the weights change significantly. 

2. About the assumptions in Section 4.3, I do not see why Hessian and noise covariance commute in general. It seems like they commute when the loss is quadratic in weights, which is assumed in Assumption 1. Then analyzing quadratic loss seems less interesting and I believe there should have been many results regarding this setting. Could you comment on this?

3. I suggest that the authors avoid using $\gg$ or $\approx$ in mathematical statements. These symbols are not precise, and it is unclear which terms are considered small. Furthermore, simply stating the expressions (Eqs. (8) and (9)) makes it difficult to understand how the variance of weights and velocity changes. Adding some discussion would be beneficial.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: - I'm a bit confused about the statement, "Next, we consider the probability that two batches $k$ and $k + h$, separated by $h$ update steps, belong to the same epoch". I understood this as follows: given two update step indexes $k$ and $k + h$, what is the probability of the corresponding batches belonging to the same epoch? I don't see any randomness in this statement. Yes, the batches are random, but if the indices are deterministically given, then the two indices are in the same epoch if there exists a $i \geq 1$ such that $M (i - 1) < k < k + h \leq M i$. I guess the randomness is supposed to be w.r.t. the randomness of the batches being sampled, but 
  - In my head, $\frac{M - |h|}{M}$ is the probability of two deterministic batches being in the same epoch when we uniformly randomly allocate their indices such that they are separated by $h$ steps.
  - Maybe I'm completely missing (or misunderstanding) something here, so please feel free to correct me here!
  - As a clarification, you are considering RandomReshuffling, where at each epoch, an ordered partition of $[N]$ is sampled at random, right? The author should consider making this precise, as there is another variant of without replacement sampling, namely, ShuffleOnce, where a single shuffle takes place at the beginning and the same batches are used in the same order for all the epochs.

- Although I agree that weights remain approximately constant near minima, theoretically, could one use appropriate perturbation arguments (e.g., Taylor expansion) to provide a more complete theory of anti-correlation with varying states? Or would such more intricate analyses not give any useful insights and thus be unnecessary?

- Does the current analysis extend to ShuffleOnce, or at least empirically?

- [low priority, did not affect my initial evaluation] If time allows, can authors try the similar experiments for small scale transformer?

- Section 4.2 is IMHO too sudden without proper motivation. Why are we suddenly interested in the weight and velocity variance? Why consider their ratio? What is the intuition of the quantity *correlation time*? Should it be understood as, 'after the correlation time, the weights and velocities are somehow correlated'...?
Moreover, Section 4.2 refers to the setup described in Section 4.3, making me think that for the sake of organization, Section 4.2 should come after 4.3?
soundness: 2
strengths: - An important gap in the SGD dynamics literature was considered by considering batched SGD without replacement in finite learning rate, showing explicit geometry-dependent implicit bias compared to replacement.
- The theory is general enough to include momentum SGD.
- Sufficient experimental results supporting the theories
summary: The paper analyzes the behavior of SGD(m) without replacement (RandomReshuffle) and its implicit bias when the iterates reach a basin of some minimum. Theoretically, they derive the exact anti-correlation of the gradient noise for static weight. Then, under additional assumptions, they prove a connection between such anticorrelations and the variance of the weight fluctuations. Qualitatively speaking, this implies that the weight variance is decreased *only* in the flat directions, showing that without replacement provides a benign implicit bias towards flat minima. Their theoretical results are complemented with various experiments.
weaknesses: - IMHO, the biggest point preventing me from giving a higher score (along with my questions below) is that *some (not all)* of this paper's main conclusions seem to overlap with the preprint [1] for SGD without replacement (without momentum). Given that the authors could provide a satisfactory answer, I'm open to raising my score further.
To my knowledge, [1] showed that in-expectation, for both **ShuffleOnce** and **RandomReshuffle**,
    - SGD without replacement = SGD with replacement along larger curvature "+" Shrinking the tr(cov of gradients) along flatter directions (see their Theorem 1; "+" means decoupling)
    - Actually, SGD without replacement adds an implicit regularize that penalizes a weighted tr(cov of gradients) over single data points (see their Theorem 3)
    - In all their derivations, they do not make any restrictive assumptions (e.g., quadratic loss, anti-correlated noise for changing state, $[C, H] = 0$, etc)

- (continuing from above) All in all, given how the problem that [1] is tackling (behavior of SGD without replacement) and the qualitative general conclusion (greater curvature is similar to uncorrelated noise (with replacement), and flatter curvature decreases the fluctuation) both seem similar, I'm curious on what additional intuitions and results this paper adds compared to [1], and how the results presented here are related to [1]. I feel this should be much more emphasized, given that this paper relies on several assumptions that [1] does not consider: quadratic loss, anti-correlated noise even under dynamically changing weights, and $[C, H] = 0$. Especially the last assumption is critical for all of the theories presented in this paper, which allows for the use of common eigenbasis.

- (continuing from above) Immediately I see two advantages: one is that this paper's framework encompasses momentum and that this has more realistic experiments supporting the theories. I'm looking for something like "under such additional assumptions, our paper provides a much more precise characterization of ..... compared to [1], which can only say that ....".

- [writing] Overall, the writing and organization should be improved. Many important discussions have been completely relegated to the Appendix. Space-wise, maybe move/reduce Sec 3?

- [writing] The authors should consider making all the assumptions explicit in the main text and collecting them in an orderly fashion, instead of using phrases such as "under general assumptions, stated in Appendix C" or ""With the above assumptions".


[1] https://arxiv.org/abs/2312.16143

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: While the overall quality of this paper is good, there are still several questions I would like to ask.

1. The main theorem is derived on the basis of three fundamental assumptions 1-3. Assumption 1 is a commonly adopted one if the discussion is limited to a local minimum. However, **Assumption 2** is newly invoked by the authors on the state-dependence of the noise. The validity of this assumption is discussed through mere words and a "look" at the experiment done in Appendix H.  It lacks a rigorous demonstration of when this assumption can be adopted and how large deviations it may lead to the main theorems. Also, the experiment and discussion provided in Appendix H are extremely specific to the very task, and a detailed connection to Assumption 2 is appreciated.

2. **Assumption 3** is about the commutation relation between the noise covariance and the Hession. The authors claim that "this assumption is not strictly necessary but it simplifies the analysis". They provide evidence in Appendix L that under the circumstances they are considering, these 2 matrices are almost aligned with each other. In Appendix L they calculate a 0.82 cosine similarity between them and claim that "it seems to be a good approximation" in the final sentence. The details in this calculation are not provided and a 0.82 similarity does not seem satisfactory to claim that two things are alike in a common sense.  Besides, "it seems to be a good approximation" sounds like the authors themselves are not very confident in this statement and it is not a rigorous manner to make an approximation. I would like the authors to elaborate more on this assumption. Could the authors provide more details about the cosine similarity of two matrices? What level of similarity would be considered sufficient? Is it possible to provide a quantitative analysis of the influence of this cosine similarity on the validity of Approximation 3? Finally, a minor question, if this assumption is "not strictly necessary", is it possible to just discard it and establish a more general theory?

3. The experiments are performed on CIFAR10 using the LeNet and the details are provided in Appendix H. From **Figure 8**, it seems that the model is severely overfitting and the authors make no comments about overfitting. Is overfitting related to any assumptions made in the paper? I'm curious about that is this severely overfitting situation chosen intentionally. If yes, why? If not, what will happen to the results if the model does not overfit?
soundness: 3
strengths: This paper is well-organized and well-written. The definitions and assumptions are made clear. The auto-correlation function of the gradient noise between different time steps is an interesting topic. The theoretical findings are solid upon the assumptions and experiments are thorough.
summary: The authors study the dynamics of SGD in a discrete-time regime while sampling without replacement which leads to correlations in the noise between different time steps. By invoking certain assumptions, they derive the correlation function of the noise and claim that the anti-correlations in the noise may cause better generalization ability.
weaknesses: See **Questions**.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 4
questions: These questions are lower-priority. Answers to them are primarily towards understanding the implicit viewpoint of the paper

1. What happens in the case of a loss function like $L(x) = x^4$ which is not well-described by a quadratic approximation near its minimum?
2. Orvieto et al. 2022 consider adding noise which is independent from the data rather than data-dependent SGD noise. How can we know that the anti-correlations in that kind of noise behave the same way as anti-correlations in SGD noise?
3. Is it true that sampling without replacement is better than increasing batch size, assuming that both result in the same overall gradient variance? 
4. How do you justify subtracting the weight drift at late times for experiments when your assumptions don't require this?
soundness: 3
strengths: This paper is clearly written with a transparent analysis. Mathematical analysis is generally paired with sufficient words to explain the meaning of the relevant equations, and is well-motivated in itself. They perform a deep analysis in a relatively common setting of sampling without replacement, which should carry over (at least heuristically) to all such algorithms. Though the experiments are at a small-scale this is not an issue in my opinion as the result is mathematically robust, and the experiments are illustrative. Because the setting is well-circumscribed, the applicability of this work could be tested on a case-by-case basis.
summary: This paper explores the consequences of anti-correlations of update steps which arise due to selection without replacement in SGD. The paper rests on the assumption that we are at the end of training, or that we are near enough a minimum of the loss that the landscape is quadratic. The analysis begins by characterizing the different-time covariance of the gradients in terms of the covariance at equal times, and verifying numerically that their formula is correct. They use this to derive the related formulas for the covariance of the weights, and the covariance of the updates to the weights. 

They define a correlation time, $\tau_i$ as the ratio of these two covariances, which becomes the central object of study. Using the results from before, they calculate $\tau_i$ in terms of the momentum hyperparameter, $\beta$ and hessian eigenvalue, $\lambda$, which reveals two distinct limits. When the eigenvalue is large they find that the correlation time decreases as $\lambda^{-1}$ while it is constant when $\lambda$ is small. These results are all confirmed numerically for the top 5000 eigenvalues of the spectrum of a LeNet model on CIFAR10.
weaknesses: The primary weakness, which is why I did not recommend this paper for acceptance, is my concern about the relevance of this work to the practical setting. In my reading of the paper I did not understand the motivation behind studying the end point of training as a limit. Due to this, the contribution of this paper is limited due to the specific setting considered. I am glad to increase my rating if the authors could sufficiently clarify the following two questions:

1. How can we understand finite training time behavior from this analysis?
2. In the setting where some $\lambda_i = 0$ exactly, how can we make sense of the constant weight assumption?

Meta-review: This paper studies the anti-correlation of gradient noise for iterates of epoch-based (without-replacement) SGD and its implications to the variance of weight fluctuation at the end of training. Under the assumption that the noise of SGD is static (which holds e.g., when the noise is evaluated at the same parameter value for the entire epoch), the paper shows that the gradient noise of different iterates within a single epoch is inherently anti-correlated, and then moves on to use this result to consider the fluctuation of weights and velocity. The main message is that anti-correlation of noise reduces the fluctuation of weights along flat directions compared to with-replacement SGD. The authors also claim through experiments that this reduced fluctuation leads to better generalization performance of models trained via without-replacement SGD than with-replacement SGD.

All reviewers agreed that the paper is well-written, the considered problem is of interest, and the presented experiments align well with theory. However, at the same time, all reviewers expressed concerns about the strong assumption that the noise is static, which is essentially the same as assuming constant weight value over one entire epoch. The reviewers and I found the authors’ justification somewhat insufficient. Although this assumption may be approximately true in some situations, it is deemed that this assumption in its current form is too strong and simplifies the analysis too much. Consider training a ReLU network; due to non-smoothness, the gradients can change abruptly due to a very small change in the weight value. Hence, I believe that the analysis of anti-correlation needs to be made more rigorous, through an analysis that takes within-epoch updates into account.

Overall, although the paper tackles an important problem and offers good insights, the theoretical analysis is based on strong assumption(s) that is hard to justify. I recommend the authors to revise their theory to remove restrictive assumptions, and I believe doing so will significantly improve the paper. At this time, I recommend rejection.

Decision: Reject