SlowGuess/NeurIPS_2023_OCR huggingface dataset:
D.K. and F.C. initiated the project. F.C. formulated an initial analysis on stochastic collapse and is primarily responsible for the experiments associated with deep neural networks. D.K. is primarily responsible for the linear teacher-student analysis and the original observation that SGD can collapse to a saddle-point introduced by symmetry. A.Y. primarily contributed to the theoretical formulations and proofs, which include the concept of invariant sets and conditions for stochastic attractivity. S.G. advised throughout the work and provided funding for computation. All of the authors have worked on writing the manuscript.
[1] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Commun. ACM, 64(3):107-115, feb 2021. ISSN 0001-0782. doi: 10.1145/3446776. URL https://doi.org/10.1145/3446776.
[2] Fengxiang He, Tongliang Liu, and Dacheng Tao. Control batch size and learning rate to generalize well: Theoretical and empirical evidence. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/dc6a70712a252123c40d2adba6a11d84-Paper.pdf.
[3] David Barrett and Benoit Dherin. Implicit gradient regularization. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=3q5IqUrkcF.
[4] Jonas Geiping, Micah Goldblum, Phil Pope, Michael Moeller, and Tom Goldstein. Stochastic training is not necessary for generalization. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=ZBESeIUB5k.
[5] Daniel Kunin, Javier Sagastuy-Brena, Lauren Gillespie, Eshed Margalit, Hidenori Tanaka, Surya Ganguli, and Daniel LK Yamins. Limiting dynamics of sgd: Modified loss, phase space oscillations, and anomalous diffusion. arXiv preprint arXiv:2107.09133, 2021.
[6] Jeff Z. HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Mikhail Belkin and Samory Kpotufe, editors, Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 2315-2357. PMLR, 15-19 Aug 2021. URL https://proceedings.mlr.press/v134/haochen21a.html.
[7] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 483–513. PMLR, 09–12 Jul 2020. URL https://proceedings.mlr.press/v125/blanc20a.html.
[8] Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 27449-27461. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/e6af401c28c1790eaef7d55c92ab6ab6-Paper.pdf.
[51] Andrew K. Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=ryfMLoCqtQ.
[52] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborova. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/cab070d53bd0d200746fb852a922064a-Paper.pdf.
[53] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
[54] Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural Networks, 2(1):53-58, 1989. ISSN 0893-6080. doi: https://doi.org/10.1016/0893-6080(89)90014-2. URL https://www.sciencedirect.com/science/article/pii/0893608089900142.
[55] Kenji Fukumizu and Shun-ichi Amari. Local minima and plateaus in hierarchical structures of multilayer perceptrons. Neural networks, 13(3):317-327, 2000.
[56] Shun-ichi Amari, Hyeyoung Park, and Tomoko Ozeki. Singularities affect dynamics of learning in neuromanifolds. Neural computation, 18(5):1007-1065, 2006.
[57] Christopher DiMattina and Kechen Zhang. How to modify a neural network gradually without changing its input-output functionality. Neural computation, 22(1):1-47, 2010.
[58] Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
[59] Florent Benaych-Georges and Raj Rao Nadakuditi. The singular values and vectors of low rank perturbations of large rectangular random matrices. Journal of Multivariate Analysis, 111:120–135, 2012. ISSN 0047-259X. doi: https://doi.org/10.1016/j.jmva.2012.04.019. URL https://www.sciencedirect.com/science/article/pii/S0047259X12001108.
[60] Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In 2007 15th European Signal Processing Conference, pages 606-610, 2007.
SGD's noise covariance shape matters. Our work is closely related to HaoChen et al. [6] work that showed how parameter-dependent noise introduces an implicit bias towards local minima with smaller variance, while spherical Gaussian noise does not. In particular, their work conjectured that the parameter dependent noise introduced by SGD has an implicit bias effect "toward parameters where the noise covariance is smaller". They studied this effect in two-layer quadratically-parameterized model introduced by [15]. Our work generalizes this effect to a broad range of models by introducing invariant sets and demonstrates that this implicit bias towards these regions is regulated by a competition between the curvature of the full-batch loss and the magnitude of the noise.
SGD can converge to a local maximum. Our work is closely related to a recent work studying failure modes of SGD [28] and a concurrent follow up work studying the stability of SGD near a fixed point [29]. Ziyin et al. [28] studied the behavior of SGD in one-dimensional quadratic and quartic losses and showed that SGD can converge to local maxima, escape from saddle-points slowly, and prefer sharp minima due to the gradient noise. A strength of their analysis is that they directly derive their results from discrete SGD updates, while we use a continuous formulation. However, their analysis is limited to one-dimensional scenarios and is used to illustrate behaviors of SGD they implicitly deemed unfavorable. In our work, we describe the mechanism of stochastic collapse generally and connect it to a simplicity bias that can benefit generalization. Concurrent to our work, Ziyin et al. [29] proposes a new notion of stability for SGD near a fixed point called probabilistic stability, which overlaps with our work in discussing stochastic collapse for single neurons but not for permutation invariance. They spotlight SGD learning phases using this stability concept, while we explore how stochastic gradients incline networks toward simpler subnetwork invariants.
SGD learns sparse features. Our work is closely related to Andriushchenko et al. [18] recent work studying a sparsity bias of SGD with large learning rates. The novelty of their work can be summarized as: 1. They revealed that large-step SGD dynamics have effective slow dynamics with multiplicative noise during loss stabilization. 2. They theoretically analyzed diagonal linear networks, showing that the SDE has implicit bias toward sparser representations. 3. They conjectured that deep nonlinear networks also show this phenomenon, which is supported by the empirical results. 4. They argued that SGD first learns sparse features and then fits the data after the step size is decreased. With this foundation laid, we would like to emphasize two fundamental aspects that set our work apart:
First, our analysis goes far beyond the basic understanding of diagonal linear networks, offering a broader perspective by introducing the invariant sets with a more general theorem (Theorem 4.2). In their work, the authors conjectured an implicit bias towards sparser features beyond their diagonal linear model, leaving this exploration open as an "exciting avenue for future work". We believe that our paper makes substantial progress along this avenue by understanding the implicit bias towards sparse representations as stochastic collapse to the invariant sets. Furthermore, our theory is applicable to general deep neural nets. We also provide a theoretical framework that predicts the quantitative conditions for this attractiveness in a general setting. While applying this theorem to arbitrary neural nets is a challenging future work, our work utilized this theorem to quantitatively analyze the collapsing threshold of simple models. Second, our results shed light on the collapsing phenomena of weight vectors, a perspective that contrasts with their emphasis on neuron similarity based on activation patterns. Given that weight vector similarities often lead to similarity in activation patterns, our findings suggest that neural networks adhere to a stronger condition.
B Further Discussion on Invariant Sets of SGD
In our work, we start by describing subsets of parameter space that are unmodified by SGD, which we term invariant sets. Several works have explored geometric properties of the loss landscape through the lens of symmetry and invariance. [35] demonstrated how the continuous symmetries of the loss function create level sets of the loss and impose geometric constraints on the gradients. [36] and [37] showed how permutation symmetries in the hidden neurons can lead to many potentially isolated minima with saddle points between them. [38] studied the learning dynamics of gradient flow near singular regions of parameter space, which for a multilayer perceptron, include permutation and sign invariance. A series of works [54-57] have leveraged the hierarchical structures of neural networks along with properties of their activation functions to identify critical points of the loss landscape.
B.1 Proof of Proposition 3.1 and Proposition 3.2
Proof.  $A$  is affine and  $0 \in A$ , thus it suffices to show that  $\frac{\partial \ell}{\partial \theta} \in A$  and is independent of the training data. Denote the activation at layer  $l$  as  $h^{(l)} = \sigma \left(w^{(l)^T} h^{(l-1)} + b^{(l)}\right)$ . The gradients associated with the outgoing weights are,
$$
\begin{array}{l} \frac {\partial \ell}{\partial w _ {\text {o u t} , i} ^ {(l + 1)}} = \frac {\partial \ell}{\partial h ^ {(l + 1)}} \odot \frac {\partial h ^ {(l + 1)}}{\partial w _ {\text {o u t} , i} ^ {(l + 1)}} \\ = \frac {\partial \ell}{\partial h ^ {(l + 1)}} \odot \sigma^ {\prime} \left(w _ {\text {o u t}, i} ^ {(l + 1)} h _ {i} ^ {(l)} + b _ {i} ^ {(l + 1)}\right) \sigma \left(w _ {\text {i n}, i} ^ {(l)} ^ {T} h ^ {(l - 1)} + b _ {i} ^ {(l)}\right). \tag {9} \\ \end{array}
$$
If  $\theta \in A$ ,  $\frac{\partial \ell}{\partial w_{\mathrm{out},p}^{(l + 1)}} = 0$ . The gradients associated with the incoming weights are,
$$
\begin{array}{l} \frac {\partial \ell}{\partial w _ {\text {i n} , i} ^ {(l)}} = \left\langle \frac {\partial \ell}{\partial h ^ {(l + 1)}}, \frac {\partial h ^ {(l + 1)}}{\partial h _ {i} ^ {(l)}} \right\rangle \frac {\partial h _ {i} ^ {(l)}}{\partial w _ {\text {o u t} , i} ^ {(l)}} \\ = \left\langle \frac {\partial \ell}{\partial h ^ {(l + 1)}}, \left(\sigma^ {\prime} \left(w _ {\text {o u t}, i} ^ {(l + 1)} h _ {i} ^ {(l)} + b ^ {(l + 1)}\right) \odot w _ {\text {o u t}, i} ^ {(l + 1)}\right) \right\rangle \sigma^ {\prime} \left(w _ {\text {i n}, i} ^ {(l) T} h ^ {(l - 1)} + b _ {i} ^ {(l)}\right) \odot h ^ {(l - 1)}. \tag {10} \\ \end{array}
$$
If  $\theta \in A$ ,  $\frac{\partial\ell}{\partial w_{\mathrm{in},p}^{(l)}} = 0$ . The gradients associated with the biases are,
$$
\begin{array}{l} \frac {\partial \ell}{\partial b _ {i} ^ {(l)}} = \left\langle \frac {\partial \ell}{\partial h ^ {(l + 1)}}, \frac {\partial h ^ {(l + 1)}}{\partial h _ {i} ^ {(l)}} \right\rangle \frac {\partial h _ {i} ^ {(l)}}{\partial b _ {i} ^ {(l)}} \\ = \left\langle \frac {\partial \ell}{\partial h ^ {(l + 1)}}, \left(\sigma^ {\prime} \left(w _ {\text {o u t}, i} ^ {(l + 1)} h _ {i} ^ {(l)} + b ^ {(l + 1)}\right) \odot w _ {\text {o u t}, i} ^ {(l + 1)}\right) \right\rangle \sigma^ {\prime} \left(w _ {\text {i n}, i} ^ {(l)} ^ {T} h ^ {(l - 1)} + b _ {i} ^ {(l)}\right). \tag {11} \\ \end{array}
$$
If  $\theta \in A$ ,  $\frac{\partial \ell}{\partial b_p^{(l)}} = 0$ . Therefore,  $\frac{\partial \ell}{\partial \theta} \in A$ .
Proof. We can directly see from Eq. 9-11 that if  $\theta \in A$ ,  $\frac{\partial\ell}{\partial w_{\mathrm{out},p}^{(l + 1)}} = \frac{\partial\ell}{\partial w_{\mathrm{out},q}^{(l + 1)}}, \frac{\partial\ell}{\partial w_{\mathrm{in},p}^{(l)}} = \frac{\partial\ell}{\partial w_{\mathrm{in},q}^{(l)}}$  and  $\frac{\partial\ell}{\partial b_p^{(l)}} = \frac{\partial\ell}{\partial b_q^{(l)}}$ . Therefore,  $\frac{\partial\ell}{\partial\theta} \in A$ .
B.2 Proof of Theorem 3.1
Proof. It suffices to show that  $\nabla \ell(\theta; x, y) \in A$  for all  $\theta \in A$  and training data. In the following, we use the notation of  $\ell(\theta)$  as a shorthand for  $\ell(\theta; x, y)$ . First, we show  $\nabla \ell_{Q}(\theta) = \nabla \ell(\theta)$  for  $\theta \in A$ , where  $\ell_{Q}(\theta) = \ell(Q\theta)$ . The derivative of  $\ell_{Q}$  along an arbitrary unit vector  $n$  is given by
$$
\nabla_ {n} \ell_ {Q} (\theta) = \lim  _ {h \rightarrow 0} \frac {\ell (Q (\theta + h n)) - \ell (Q \theta)}{h} = \lim  _ {h \rightarrow 0} \frac {\ell (\theta + h n) - \ell (\theta)}{h} + \frac {\ell (Q (\theta + h n)) - \ell (\theta + h n)}{h}.
$$
If  $n$  is a tangent vector of  $A$ , the second term is zero. If  $n$  is a normal vector of  $A$ , then by the approximate  $Q$ -symmetry assumption, the second term converges to zero. These two facts imply that
$\nabla \ell_{Q}(\theta) = \nabla \ell (\theta)$ . Then, for any  $\theta \in A$ ,  $\nabla \ell (\theta) = \nabla \ell_{Q}(\theta) = Q^{\intercal}\nabla \ell (Q\theta)$ . Since  $Q$  is symmetric  $(Q = Q^{\intercal})$  or orthogonal  $(QQ^{\intercal} = I_d)$ , we obtain
$$
Q \nabla \ell (\theta ; x, y) = \nabla \ell (\theta ; x, y),
$$
which implies  $\nabla \ell (\theta ;x,y)\in A$ . Hence, regardless of the sequence of mini-batches, starting from  $\theta^{(0)}\in A$ , the SGD trajectory remains within  $A$ , demonstrating  $A$  is an invariant set.
Theorem 3.1 provides an alternative proof of Proposition 3.1 and Proposition 3.2, simply identify a symmetric or orthogonal matrix  $Q$  that generates these sets. For the sign invariant sets,  $Q$  is the identity matrix with negative diagonals restricted to the index of the parameters associated with the invariant set. For the permutation invariant set,  $Q$  is the identity matrix with the diagonal blocks for the parameters associated with the two permutatable neurons swapped to being off diagonal.
B.3 Additional Invariant Sets
Beyond the classes of invariant sets already presented, there can exist other invariant sets depending on the structure of the architecture. For example, for a network using a softmax layer, the columns of the weight matrix and the vector of biases immediately preceding this layer can be shifted by a constant without changing the outputs. As a result, the dynamics for these parameters are constrained to an affine space with a constant sum determined by their initialization. This subspace is an invariant set. Similarly, if the training data has a non-trivial kernel space, as in the case of overparameterized linear regression, then the rows of the first-layer's weight matrix can be translated by any vector in the kernel space without changing the output of the network. Again, this constrains the learning dynamics to an affine space, which is an invariant set. In general, any translational symmetry should generate invariant sets.
C Invariant sets of SGF
C.1 Definition of invariant sets for continuous processes
Here, we provide the formal definition of invariant sets for a continuous process and a formal statement of Proposition 4.1.
Definition C.1 (Invariant sets of continuous processes). For a given stochastic process  $\{\theta_t \in \mathbb{R}^d : t \geq 0\}$ , a Borel-measurable set  $A \subset \mathbb{R}^d$  is defined as invariant if for every initial point  $\theta_0$  in  $A$ , the probability that the process stays in  $A$  for all  $t \geq 0$  is 1, i.e.,
$$
\mathbb {P} \left[ \theta_ {t} \in A f o r a n y t \geq 0 \mid \theta_ {0} \in A \right] = 1. \tag {12}
$$
C.2 Proof of Proposition 4.1
Proposition 4.1 (Formal). Consider an SGD process  $\{\theta^{(t)}\in \mathbb{R}^d:t\in \mathbb{N}\}$  given by Eq. 2, where the gradients  $\nabla \ell (\cdot ;x_i,y_i):\mathbb{R}^d\to \mathbb{R}^d$  in Eq. 2 are Lipschitz continuous and bounded for any  $i\in [n]$ . If  $A\subseteq \mathbb{R}^d$  is an affine invariant set of SGD  $\{\theta^{(t)}:t\in \mathbb{N}\}$ , then  $A$  is also an invariant set of an SGF process  $\{\theta_t:t\geq 0\}$  given by Eq. 3.
Proof. By assumption,  $\nabla \ell (\theta ;x_i,y_i)$  is  $L$ -Lipschitz in  $\theta$  and  $\| \nabla \ell (\theta ;x_i,y_i)\| \leq L$  with some constant  $L > 0$ . Then we can see that  $\nabla \mathcal{L}$  is also Lipschitz continuous and bounded as follows:
$$
\| \nabla \mathcal {L} (\theta_ {1}) - \nabla \mathcal {L} (\theta_ {2}) \| = \frac {1}{n} \left\| \sum_ {i \in [ n ]} \left(\nabla \ell \left(\theta_ {1}; x _ {i}, y _ {i}\right) - \nabla \ell \left(\theta_ {2}; x _ {i}, y _ {i}\right)\right) \right\| \leq L \| \theta_ {1} - \theta_ {2} \|
$$
$$
\| \nabla \mathcal {L} (\theta) \| \leq \frac {1}{n} \sum_ {i} \| \ell (\theta ; x _ {i}, y _ {i}) \| \leq L.
$$
Let  $\bar{\ell} (\theta ,x_i,y_i) = \ell (\theta ,x_i,y_i) - \mathcal{L}(\theta ,x_i,y_i)$ .  $\nabla \bar{\ell} (\cdot ,x_i,y_i)$  is clearly  $2L$ -Lipschitz continuous and is bounded by  $2L$  for any  $i\in [n]$ . Then,
$$
\begin{array}{l} \left\| D _ {l m} \left(\theta_ {1}\right) - D _ {l m} \left(\theta_ {2}\right) \right\| \\ = \frac {1}{n} \left\| \sum_ {i \in [ n ]} \left(\partial_ {l} \bar {\ell} \left(\theta_ {1}; x _ {i}, y _ {i}\right) \partial_ {m} \bar {\ell} \left(\theta_ {1}; x _ {i}, y _ {i}\right) - \partial_ {l} \bar {\ell} \left(\theta_ {2}; x _ {i}, y _ {i}\right) \partial_ {m} \bar {\ell} \left(\theta_ {2}; x _ {i}, y _ {i}\right)\right) \right\| \\ \leq \frac {1}{n} \sum_ {i \in [ n ]} \left[ 2 L \left(\left| \partial_ {l} \bar {\ell} \left(\theta_ {1}; x _ {i}, y _ {i}\right) \right| + \left| \partial_ {m} \bar {\ell} \left(\theta_ {2}; x _ {i}, y _ {i}\right) \right|\right) \| \theta_ {1} - \theta_ {2} \| ^ {2} \right] \\ \leq 4 L ^ {2} \| \theta_ {1} - \theta_ {2} \| ^ {2}. \\ \end{array}
$$
This implies that  $\| D(\theta_1) - D(\theta_2)\|_{\infty}\leq 4L^2\| \theta_1 - \theta_2\| ^2$  , where  $\| \cdot \|_{\infty}$  denotes the operator norm. On the other hand, the Powers-Stormer inequality gives,
$$
\left\| \sqrt {D (\theta_ {1})} - \sqrt {D (\theta_ {2})} \right\| _ {\mathrm {F}} \leq \sqrt {d} \sqrt {\left\| D (\theta_ {1}) - D (\theta_ {2}) \right\| _ {1}} \leq d \sqrt {\left\| D (\theta_ {1}) - D (\theta_ {2}) \right\| _ {\infty}},
$$
where  $\| \cdot \|_{\mathrm{F}}$  and  $\| \cdot \|_{1}$  denote the Frobenius norm and Schatten 1-norm, respectively. Thus,  $\sqrt{D(\theta)}$  is Lipschitz continuous with respect to the Frobenius norm. This establishes the existence of a unique strong solution to the SDE in Eq. 3. It now suffices to show that the SDE has a solution  $\{\theta_t : t \geq 0\}$  satisfying Eq. 12. Let  $P: \mathbb{R}^d \to \mathbb{R}^d$  be the projection operator to the affine subset  $A$ , and we construct a projected SDE:
$$
d \tilde {\theta} _ {t} = - P \nabla \mathcal {L} (\tilde {\theta} _ {t}) d t + \sqrt {\frac {\eta}{\beta}} P \Sigma (\tilde {\theta} _ {t}) d B _ {t}.
$$
By construction, this projected SDE also has a unique solution  $\tilde{\theta}_t$ , which clearly satisfies the condition Eq. 12. Moreover, given that  $A$  is an invariant set of the original SGD process, we have  $P\nabla \ell (\theta ;x_i,y_i) = \nabla \ell (\theta ;x_i,y_i)$  for all  $i\in \mathcal{B}$  and  $x\in A$ . This means that  $\{\tilde{\theta}_t:t\geq 0\}$  satisfies Eq.12 almost surely. Hence, we conclude that  $A$  is also an invariant set of the solution to Eq.3.
D SDE formulations of learning dynamics with gradient noise
D.1 The relationship between SGD and SGF
Recall the SGD update rule from Eq. 2.
$$
\theta^ {(t + 1)} = \theta^ {(t)} - \frac {\eta}{\beta} \sum_ {i \in \mathcal {B} ^ {(t)}} \nabla_ {\theta} \ell \left(\theta^ {(t)}; x _ {i}, y _ {i}\right). \tag {13}
$$
We factorize the right-hand side into the negative full-batch gradient and the deviation from the full-batch gradient,
$$
\underbrace {\theta^ {(t + 1)} = \theta^ {(t)} - \eta \nabla \mathcal {L} \left(\theta^ {(t)}\right)} _ {\text {g r a d i e n t d e s c e n t}} - \sqrt {\eta} \underbrace {\left(\frac {\sqrt {\eta}}{\beta} \sum_ {i \in \mathcal {B}} \nabla_ {\theta} \ell \left(\theta^ {(t)} ; x _ {i} , y _ {i}\right) - \sqrt {\eta} \nabla \mathcal {L} \left(\theta^ {(t)}\right)\right)} _ {\text {g r a d i e n t n o i s e}}. \tag {14}
$$
Let  $\xi_{\mathcal{B}}(\theta)$  represent the gradient noise at  $\theta$ . Assuming we sample with replacement, the gradient noise can be disentangled into a sum of i.i.d. per-sample gradient noises:  $\xi_{\mathcal{B}}(\theta) = \frac{\sqrt{\eta}}{\beta} \sum_{i \in \mathcal{B}} \xi_i$  where  $\xi_i = \nabla \ell(\theta; x_i, y_i) - \nabla \mathcal{L}(\theta)$ . It is easy to check that  $\mathbb{E}[\xi_i(\theta)] = 0$  for all  $\theta$ . Assuming that the gradient noises are gaussian random variables, Eq. 14 turns into
$$
\Delta \theta^ {(t)} = - \eta \nabla \mathcal {L} \left(\theta^ {(t)}\right) + \sqrt {2 D \left(\theta^ {(t)}\right)} \Delta W, \tag {15}
$$
where  $\Delta W\sim \mathcal{N}(0,\eta I)$  and the diffusion matrix  $D(\theta^{(t)}) = \frac{1}{2}\mathrm{Var}(\xi_{\mathcal{B}}) = \frac{\eta}{2\beta}\mathrm{Var}(\xi_{i})$  . Eq. 15 is a Euler-Maruyama discretization of the following continuous-time stochastic differential equation, which we denote as stochastic gradient flow (SGF),
$$
d \theta_ {t} = - \nabla \mathcal {L} \left(\theta_ {t}\right) d t + \sqrt {2 D \left(\theta_ {t}\right)} d B _ {t}, \quad \theta_ {0} = \theta^ {(0)}. \tag {16}
$$
The diffusion matrix can be expressed as:
$$
D \left(\theta_ {t}\right) = \frac {\eta}{2 \beta} \left(\frac {1}{N} \sum_ {i \in [ N ]} \left[ \nabla_ {\theta} \ell \left(\theta_ {t}; x _ {i}, y _ {i}\right) \nabla_ {\theta} \ell \left(\theta_ {t}; x _ {i}, y _ {i}\right) ^ {T} \right] - \nabla \mathcal {L} \left(\theta_ {t}\right) \nabla \mathcal {L} \left(\theta_ {t}\right) ^ {T}\right). \tag {17}
$$
The diffusion matrix can be decomposed into a product of a parameter-independent magnitude term  $D_{m}$  and a parameter-dependent shape term  $D_{s}(\theta)$  such that  $D(\theta) = D_{m}D_{s}(\theta)$ , where
$$
D _ {m} = \frac {\eta}{2 \beta} \tag {18}
$$
$$
D _ {s} (\theta) = \frac {1}{N} \sum_ {i \in [ N ]} \left[ \nabla_ {\theta} \ell \left(\theta_ {t}; x _ {i}, y _ {i}\right) \nabla_ {\theta} \ell \left(\theta_ {t}; x _ {i}, y _ {i}\right) ^ {T} \right] - \nabla \mathcal {L} \left(\theta_ {t}\right) \nabla \mathcal {L} \left(\theta_ {t}\right) ^ {T} \tag {19}
$$