Stochastic Gradient Descent-Induced Drift of Representation in a
  Two-Layer Neural Network
Representational drift refers to over-time changes in neural activation
accompanied by a stable task performance. Despite being observed in the brain
and in artificial networks, the mechanisms of drift and its implications are
not fully understood. Motivated by recent experimental findings of
stimulus-dependent drift in the piriform cortex, we use theory and simulations
to study this phenomenon in a two-layer linear feedforward network.
Specifically, in a continual online learning scenario, we study the drift
induced by the noise inherent in the Stochastic Gradient Descent (SGD). By
decomposing the learning dynamics into the normal and tangent spaces of the
minimum-loss manifold, we show the former corresponds to a finite variance
fluctuation, while the latter could be considered as an effective diffusion
process on the manifold. We analytically compute the fluctuation and the
diffusion coefficients for the stimuli representations in the hidden layer as
functions of network parameters and input distribution. Further, consistent
with experiments, we show that the drift rate is slower for a more frequently
presented stimulus. Overall, our analysis yields a theoretical framework for
better understanding of the drift phenomenon in biological and artificial
neural networks.
Farhad Pashakhanloo, Alexei Koulakov
# Stochastic Gradient Descent-induced drift of representation in a two-layer neural network

###### Abstract

Representational drift refers to over-time changes in neural activation accompanied by a stable task performance. Despite being observed in the brain and in artificial networks, the mechanisms of drift and its implications are not fully understood. Motivated by recent experimental findings of stimulus-dependent drift in the piriform cortex, we use theory and simulations to study this phenomenon in a two-layer linear feedforward network. Specifically, in a continual learning scenario, we study the drift induced by the noise inherent in the Stochastic Gradient Descent (SGD). By decomposing the learning dynamics into the normal and tangent spaces of the minimum-loss manifold, we show the former correspond to a finite variance fluctuation, while the latter could be considered as an effective diffusion process on the manifold. We analytically compute the fluctuation and the diffusion coefficients for the stimuli representations in the hidden layer as a function of network parameters and input distribution. Further, consistent with experiments, we show that the drift rate is slower for a more frequently presented stimulus. Overall, our analysis yields a theoretical framework for better understanding of the drift phenomenon in biological and artificial neural networks.

Machine Learning, ICML, Neural Networks, Neural Networks, Neural Networks, Neural Networks

## 1 Introduction

Representational drift has been observed across different parts of the nervous system, such as in the hippocampus (Ziv et al., 2013), sensorimotor (Rule et al., 2019), and the visual systems (Deitch et al., 2021; Marks & Goard, 2021). A recent study in the piriform cortex, (Schoonover et al., 2021), a brain structure processing information about smells, demonstrated drift of odorant representations despite stable odor identification. Such drift was characterized by a gradual and across-days decay in the self-similarity of the stimulus representation. Additionally, in the same study, the drift was shown to be stimulus dependent. It was shown that a more familiar stimulus drifts at a smaller rate (Schoonover et al., 2021).

The mechanisms and implications of the drift in the brain and artificial neural networks are still under investigation (see recent reviews by Masset et al. (2022), Rule et al. (2019), and Driscoll et al. (2022)). Recent modeling studies have shown that drift could happen in the presence of synaptic or other types of noise (Qin et al., 2023; Aitken et al., 2022). Using simulations of neural networks, Aitken et al. (2022) showed different noise types injected during training could lead to drift with qualitatively different patterns and geometries. However, no theoretical consideration of the drift was provided in that study. Qin et al. (2023) studied the drift in a network with a similarity-matching objective and a Hebbian/anti-Hebbian learning rule. They showed that noisy synaptic updates could lead to a random-walk exploration of the solution space. Other mechanisms have been suggested on how a stable readout could be performed despite an evolving population code (Rule & O'Leary, 2022; Rule et al., 2020; Kalle Kossio et al., 2021).

One important source of noise during learning for both natural and artificial neural networks is the sampling noise that arises from the stochasticity in observing the data. It is therefore natural to ask if and how this type of noise can lead to drift, and whether it can explain the stimulus-dependency of the rate of the drift observed in experiments. Here, we aim to answer these questions in a two layer linear neural network model that undergoes online continual learning via Stochastic Gradient Descent (SGD). In addition to being common in machine learning, feed-forward networks are a reasonable first approximation to sensory (olfactory) processing in the brain. Specifically, we found the two layer linear network to be one of the simplest models that enables studying the representational drift in the hidden layer, and yet allows for analytical tractability.
## 2 Model and Theory

A multilayer network including \(L\) layers can be described by a set of weight matrices \(\mathbf{W}^{(l)}\), where index \(l\) enumerates the individual layers. The network can be represented by a vector in a vector space constructed from non-commuting weight matrices:

\[\mathbf{\theta}=\left(\mathbf{W}^{(1)},\mathbf{W}^{(2)},...,\mathbf{W}^{(L)}\right). \tag{1}\]

Our goal is to study the dynamics of learning and the drift in this space for a two-layer neural network described in Section 2.1. We derive the manifold of stable minimum-loss (Section 2.2), and study its first order differential geometry (Section 2.3). We use the Euclidean inner product in this space, which means for two vectors \(\mathbf{\theta}_{1}\) and \(\mathbf{\theta}_{2}\), we have:

\[\mathbf{\theta}_{1}^{T}\mathbf{\theta}_{2}=\sum_{l}\mathrm{tr}\Big{(}\mathbf{W}_{1}^{(l) T}\bar{\mathbf{W}}_{2}^{(l)}\Big{)}. \tag{2}\]

By characterizing the movements normal and tangential to the manifold (Sections 2.4.1 and 2.4.2), we derive an effective diffusion process on the manifold and calculate the corresponding diffusion ("drift") rates for the representations (Section 2.4.3). Finally, in Sections 3 and 4 we study the dependency of the drift on the input statistics in two cases of an isotropic Gaussian stimuli, and a case with a frequently presented stimulus. Our results are further validated by numerical simulations.

### Neural network model

Our model consists of a linear feed-forward neural network with an expansive hidden layer, as shown in Figure 1. The input to the network represents the external stimulus (\(\mathbf{x}_{n\times 1}\)), the hidden layer activation is considered as the representation of the stimulus (\(\mathbf{h}_{p\times 1}\), e.g. neural activities in the piriform cortex), and the output is the task outcome (\(\mathbf{y}_{m\times 1}\), e.g. the percept or an identity of an stimulus). The two weight matrices are \(\mathbf{U}=\mathbf{W}^{(1)}\) and \(\mathbf{W}=\mathbf{W}^{(2)}\) respectively, and the predicted output is \(\mathbf{\hat{y}}=\mathbf{W}\mathbf{U}\mathbf{x}\). Additionally, in the parameter space the network is denoted by \(\mathbf{\theta}=(\mathbf{U},\mathbf{W})\).

We consider a continual online learning scenario in which the network sees one sample at a time taken independently from a fixed data distribution. The objective function consist of a Mean Squared Error (MSE) loss and L2-regularization. Hence, the sample loss becomes:

\[l(\mathbf{x},\mathbf{y};\mathbf{\theta})=\frac{1}{2}\|\mathbf{y}-\mathbf{W}\mathbf{U} \mathbf{x}\|^{2}+\frac{\gamma}{2}\|\mathbf{W}\|_{F}^{2}+\frac{\gamma}{2}\|\mathbf{U}\|_{F }^{2}, \tag{3}\]

where \(\gamma\) is the regularization coefficient. Further, we assume \(\mathbf{y}=\mathbf{x}\), which means the goal of the network is to learn the identity mapping from the input to the output. Note this essentially becomes an autoencoder but with an expansive hidden layer (\(p\geq n=m\)). Finally, the learning occurs via the SGD with a minibatch size of one. The update equation upon observing sample \(\mathbf{x}\) is:

\[\Delta\mathbf{\theta}=-\eta\mathbf{g}(\mathbf{x};\mathbf{\theta}), \tag{4}\]

where \(\mathbf{g}(\mathbf{x};\mathbf{\theta})=\nabla_{\mathbf{\theta}}l(\mathbf{x};\mathbf{\theta})\) is the gradient vector, and \(\eta\) is the learning rate. In terms of the weight matrices, this is equivalent to \(\Delta\mathbf{W}=-\eta\nabla_{\mathbf{W}}l\) and \(\Delta\mathbf{U}=-\eta\nabla_{\mathbf{U}}l\).

### Degeneracy and the manifold of solutions

One of the conditions that make representational drift possible is the redundancy of the parameters in achieving an optimal task performance. In our model, this could be identified by the rotational symmetry in the network, as for any orthogonal \(\mathbf{Q}_{p\times p}\), the transformation \(\mathbf{\tilde{W}}\rightarrow\mathbf{\tilde{W}}\mathbf{Q}\) and \(\mathbf{\tilde{U}}\rightarrow\mathbf{Q}^{-1}\mathbf{\tilde{U}}\) leaves the loss unchanged. We aim to study whether, and how the space of redundant parameters are explored due to the online learning stochasticity, and from that, characterize the rate and the geometry of the drift for the representations.

We define the _manifold of solutions_, \(\mathcal{M}\), to be the set of stable critical points in the parameter space. This manifold represents the redundancy in the model, as all the points on it have the same expected loss value \(L(\mathbf{\theta})=\langle l(\mathbf{\theta})\rangle_{x}\), and hence are equally preferable (note that \(\langle.\rangle_{x}\) is the expectation over the input distribution). We analytically derive \(\mathcal{M}\) in the following theorem as a function of the input distribution, and for the rest of the paper refer to it as the _manifold_.

**Theorem 2.1**.: (Manifold) _The manifold of solutions for learning the identity map (\(\mathbf{y}=\mathbf{x}\)) is:_

\[\mathcal{M}:\{\tilde{\mathbf{\theta}}=(\tilde{\mathbf{U}},\tilde{\mathbf{W}} )\mid\tilde{\mathbf{W}}\tilde{\mathbf{W}}^{T}=\mathbf{I}_{n}-\gamma\mathbf{\Sigma}_{\mathbf{x}}^{ -1},\tilde{\mathbf{U}}=\tilde{\mathbf{W}}^{T}\}, \tag{5}\]

_where \(\mathbf{\Sigma}_{\mathbf{x}}=\langle\mathbf{x}\mathbf{x}^{T}\rangle_{x}\) is the second order moment of the input distribution with eigenvalues that are assumed to be greater than \(\gamma\)._

The above can be proven by first finding the critical points of the expected loss satisfying \(\nabla_{\mathbf{U}}L=\nabla_{\mathbf{W}}L=0\), and from those finding the stable solutions (see Section A.1 in the Appendix).

Figure 1: (left) Neural network model. (right) Manifold of minimum-loss (\(\mathcal{M}\)) in the parameter space.
### Differential geometry of the manifold

To study the learning dynamics near the manifold, we will characterize its first order differential geometry. This is done by finding the local tangent and normal spaces (\(T_{\tilde{\mathbf{\theta}}}\mathcal{M}\) and \(N_{\tilde{\mathbf{\theta}}}\mathcal{M}\), respectively), as shown in the following lemmas. Note the results of the lemmas are not directly used in the rest of the main paper, and we only mention them here for completeness (see also Appendix Section A.2).

**Lemma 2.2**.: (Tangent space) _The local tangent space to the manifold at point \((\tilde{\mathbf{W}}^{T},\tilde{\mathbf{W}})\) is spanned by vectors \(\mathbf{t}=(\mathbf{T}_{\mathbf{W}}^{T},\mathbf{T}_{\mathbf{W}})\) with \(\mathbf{T}_{\mathbf{W}}=\tilde{\mathbf{W}}\tilde{\mathbf{W}}^{T}\mathbf{\Omega}\tilde{\mathbf{W}}+\mathbf{K }\tilde{\mathbf{W}}_{\perp}\), where \(\mathbf{\Omega}\) is an arbitrary \(n\times n\) skew-symmetric matrix, \(\mathbf{K}\) is an arbitrary \(n\times(p-n)\) matrix, and \(\tilde{\mathbf{W}}_{\perp}\) is a full-rank \((p-n)\times p\) matrix whose rows are orthogonal to rows of \(\tilde{\mathbf{W}}\)._

With the definition of the inner product in Eq.2, the normal space can be defined, which is described in the next lemma.

**Lemma 2.3**.: (Normal space) _The local normal space to the manifold at point \((\tilde{\mathbf{W}}^{T},\tilde{\mathbf{W}})\) is spanned by vectors \(\mathbf{n}=(\mathbf{M},\mathbf{S}\tilde{\mathbf{W}}-\mathbf{M}^{T})\), where \(\mathbf{S}\) is an arbitrary \(n\times n\) symmetric matrix and \(\mathbf{M}\) is an arbitrary \(p\times n\) matrix._

We will use the above lemmas to find tangential and normal projection operators, \(\Pi_{T}(.)\) and \(\Pi_{N}(.)\), which project an arbitrary vector to the tangent and normal space of the manifold respectively (see lemma A.1 in the Appendix).

### Learning dynamics near the manifold and emergence of drift

In this section, we will study the stochastic dynamics of learning in an online learning scenario in which the data are sampled from a fixed distribution. We assume a sufficient training has passed, and hence we are predominantly on or close to the manifold being continuously nudged around due to the SGD noise. An arbitrary point \(\mathbf{\theta}\) near the manifold can be represented as:

\[\mathbf{\theta}=\tilde{\mathbf{\theta}}+\mathbf{\theta}_{N} \tag{6}\]

where \(\tilde{\mathbf{\theta}}\) is the closest point on the manifold to \(\mathbf{\theta}\), and \(\mathbf{\theta}_{N}\in N_{\tilde{\mathbf{\theta}}}\mathcal{M}\) is the deviation from the manifold in the normal space. In the rest of this section, we will find update equations for \(\mathbf{\theta}_{N}\) and \(\tilde{\mathbf{\theta}}\) respectively, and describe their dynamics.

#### 2.4.1 Fluctuation in the normal space

We can find an update equation for \(\mathbf{\theta}_{N}\) by projecting the two sides of the SGD update (Eq.4) into the normal space. For small learning rates and over long times, this equation can be approximated with a continuous-time stochastic differential equation (SDE):

\[d\mathbf{\theta}_{N}=-\mathbf{H}\mathbf{\theta}_{N}dt-\sqrt{\eta}\,\mathbf{C}d\mathbf{B}_{t}. \tag{7}\]

In the above, \(\mathbf{H}_{k,l\,(\in[1,2np])}=\frac{\partial L}{\partial\mathbf{\theta}_{k}\partial _{t}}|_{\tilde{\mathbf{\theta}}}\) is the Hessian, \(\mathbf{B}_{t}\) is the standard multi-dimensional Brownian motion (Wiener process), and \(\mathbf{C}\approx\langle\tilde{\mathbf{g}}\tilde{\mathbf{g}}^{T}\rangle_{x}^{1/2}\), where \(\tilde{\mathbf{g}}\) is the gradient at point \(\tilde{\mathbf{\theta}}\) on the manifold. In general, the gradient (\(\mathbf{g}\)) can be analytically calculated at any point by differentiating the loss function with respect to the weight matrices (see Eq.29 in the Appendix). On the manifold we have:

\[\tilde{\mathbf{g}}(\mathbf{x}):=\mathbf{g}(\mathbf{x};\tilde{\mathbf{\theta}})=(\tilde{\mathbf{W}}^{T }\mathbf{Z}_{\mathbf{x}},\mathbf{Z}_{\mathbf{x}}\tilde{\mathbf{W}}), \tag{8}\]

where \(\mathbf{Z}_{\mathbf{x}}=\gamma(\mathbf{I}_{n}-\mathbf{\Sigma}_{n}^{-1}\mathbf{x}\mathbf{x}^{T})\). Note that \(\mathbf{C}\) in Eq.7 is in general a function of \(\mathbf{\theta}_{N}\). But in deriving the SDE, we approximated it with its value on the manifold (\(\mathbf{\theta}_{N}=0\)), which is justified for small deviations (see Section A.3.1 for the derivation of the SDE).

Since the Hessian is positive semidefinite on the manifold (see Section A.1.3), the process defined by the SDE in Eq.7 is a mean-reverting process known as the multi-dimensional Ornstein-Uhlenbeck (OU) stochastic process (Gardiner et al. (1985)). OU has a stationary solution with zero mean and a finite variance (see Section A.3.2 in the Appendix). Hence, we refer to the deviations in the normal space as _fluctuation_. We represent the fluctuation in a basis constructed from the eigenvectors of the Hessian with positive eigenvalues:

\[\mathbf{\rho}=\mathbf{N}^{T}\mathbf{\theta}_{N},\quad\text{where }\mathbf{N}=[\mathbf{n}_{1}|\, \mathbf{n}_{2}|\,...\,|\mathbf{n}_{K}] \tag{9}\] \[\text{and }\mathbf{H}\mathbf{n}_{k}=\lambda_{k}\mathbf{n}_{k},\ \lambda_{k}>0\]

Here \(\rho_{k}\) represents the deviation from the manifold along the Hessian eigenvector \(\mathbf{n}_{k}\). As shown in Section A.3.2 in the Appendix, the stationary solution of \(\mathbf{\rho}\) satisfies \(\langle\rho_{k}\rangle=0\), and has the covariance:

\[\langle\rho_{k}\rho_{l}\rangle=\frac{\eta}{\lambda_{k}+\lambda_{l}}\langle\mathbf{n }_{k}^{T}\tilde{\mathbf{g}}(\mathbf{x})\,\mathbf{n}_{l}^{T}\tilde{\mathbf{g}}(\mathbf{x})\rangle_{ x}. \tag{10}\]

As expected, since \(\tilde{\mathbf{g}}\) is the driver of the fluctuations in the SDE, the covariance depends on the the projection of \(\tilde{\mathbf{g}}\) on \(\mathbf{n}_{k}\). However, it turns out that \(\tilde{\mathbf{g}}\) has no projection on a subspace of the Hessian eigenspace, and hence for the purpose of fluctuations, only a subset of \(\{\mathbf{n}_{k}\}\) are relevant. This is described in the next proposition.

**Proposition 2.4**.: (Hessian eigenspace) _If \((\mathbf{v}_{i},s_{i})\) are the eigenvector/eigenvalue pairs of \(\mathbf{\Sigma}_{\mathbf{x}}=\langle\mathbf{x}\mathbf{x}^{T}\rangle\) with \(s_{1}\geqslant..\geqslant s_{n}>\gamma\), the Hessian eigenvectors \(\mathbf{n}_{k}\) along which there is non-zero fluctuation (i.e. \(\mathbf{n}_{k}^{T}\tilde{\mathbf{g}}\neq 0\)) correspond to \((\mathbf{v}_{i},\mathbf{v}_{j})\) pairs for \(i,j\in[1,n]\) via:_

\[\mathbf{n}_{k}\equiv\mathbf{n}_{ij}=(\tilde{\mathbf{W}}^{T}\mathbf{Z}_{ij},\mathbf{Z}_{ij}\tilde{ \mathbf{W}}), \tag{11}\]

_where \(\mathbf{Z}_{ij}=C_{ij}(\kappa_{ij}\mathbf{v}_{i}\mathbf{v}_{j}^{T}+\mathbf{v}_{j}\mathbf{v}_{i}^{T})\). Here, \(\kappa_{ij}=\text{sgn}(i-j)(\sqrt{1+b^{2}}-b)\) for \(b=(\frac{1}{\gamma}-\frac{s_{i}+s_{j}}{2s_{i}s_{j}})|s_{i}-s_{j}|\), and \(C_{ij}=[(2-\gamma/s_{i}-\gamma/s_{j})(1+\kappa_{ij}^{2})]^{-\frac{1}{2}}\) is a normalization constant ensuring \(\mathbf{n}_{k}^{T}\mathbf{n}_{k}=1\). The corresponding Hessian eigenvalues are \(\lambda_{ii}=2(s_{i}-\gamma)\) and \(\lambda_{ij(i\neq j)}=2s_{i}-\gamma(s_{i}/s_{j}+\kappa_{ij})\). (see proof in Section A.4.)_
The above proposition relates the relevant directions in the normal space to the corresponding \((ij)\) indices of the input covariance eigenvectors (this makes \(k\) a composite index). Subsequently, the components of the fluctuation covariance can be written as \(\langle\rho_{k}\rho_{l}\rangle\equiv\langle\rho_{ij}\rho_{pq}\rangle\) (see Section A.5.1 in the Appendix for derivation of the components). The remark below provides an interpretation of movements along different \(\mathbf{n}_{k}\) in the representation space.

_Remark 2.5_.: (Fluctuation of the representations): On the manifold, the representations of all possible stimuli form an n-dimensional ellipsoid embedded in \(\mathbb{R}^{p}\). The main axes of this ellipsoid are \(\widetilde{\mathbf{h}}_{i}=\tilde{\mathbf{W}}^{T}\mathbf{v}_{i}\) for \(i\in[1,n]\), with norms \(|\tilde{\mathbf{h}}_{i}|=(1-\gamma/s_{i})^{1/2}\). By replacing \(\Delta\mathbf{U}\) from the previous proposition in \(\Delta\mathbf{h}_{i}=\Delta\mathbf{U}\mathbf{v}_{i}\), we can find how deviation along different \(\mathbf{n}_{k}\) leads to different deformation modes for the ellipsoid. Specifically, we can see that deviation along the eigenvector \(\mathbf{n}_{ii}\) changes \(\tilde{\mathbf{h}}_{i}\) radially (changing its norm), while leaving other (orthogonal) axes unchanged. Similarly, deviation along \(\mathbf{n}_{ij}\) for \(i\neq j\) moves \(\tilde{\mathbf{h}}_{i}\) and \(\tilde{\mathbf{h}}_{j}\) toward or away from each other, and leaves other axes intact. The latter shows that the fluctuations can change not only the lengths but also the pairwise angles of the representation vectors. Note however that the average of these changes is zero over time.

We take the variance of \(|\tilde{\mathbf{h}}_{i}|\) as a measure of fluctuation for the representations:

\[\sigma_{i}^{2}:=\text{var}(|\mathbf{h}_{i}(t)|)=\frac{1}{2}\langle\rho_{ii}^{2} \rangle=\frac{\eta\gamma^{2}}{4s_{i}}(\frac{\langle x_{i}^{4}\rangle_{x}}{s_{ i}^{2}}-1) \tag{12}\]

where \(x_{i}:=\mathbf{v}_{i}^{T}\mathbf{x}\) (see Section A.5.2 for derivation). As shown later in the paper, this has an excellent match to numerically measured fluctuations.

#### 2.4.2 Tangential updates

Over a step of learning update, the displacement in the parameter space (\(\Delta\mathbf{\theta}\)) has a corresponding projected movement on the manifold (\(\Delta\tilde{\mathbf{\theta}}\)). For small learning rates and by avoiding the curvature effect, we have:

\[\Delta\tilde{\mathbf{\theta}}=-\eta\mathbf{g}_{T}(\mathbf{x};\tilde{\mathbf{\theta}}+\mathbf{N} \mathbf{\rho}), \tag{13}\]

where \(\mathbf{g}_{T}=\Pi_{T}(\mathbf{g})\) is the projection of the gradient vector onto the tangent space. In the next theorem, we calculate this projection and find its action on the representations.

**Theorem 2.6**.: (Tangential update) _For a normal deviation \(\mathbf{\rho}=\sum_{k}\rho_{k}\mathbf{n}_{k}\) from point \(\tilde{\mathbf{\theta}}=(\tilde{\mathbf{W}}^{T},\tilde{\mathbf{W}})\) on the manifold, the tangential projection of the gradient is:_

\[\mathbf{g}_{T}(\mathbf{x};\tilde{\mathbf{\theta}}+\mathbf{N}\mathbf{\rho})=(\mathbf{G}_{ \mathbf{U}_{T}},\mathbf{G}_{\mathbf{U}_{T}}^{T}),\quad\text{for} \tag{14}\] \[\mathbf{G}_{\mathbf{U}_{T}}=\tilde{\mathbf{W}}^{T}(\tilde{\mathbf{W}}\tilde{\mathbf{ W}}^{T})^{-\frac{1}{2}}\sum_{k}\rho_{k}\mathbf{\mathcal{G}}_{k}^{\,:}(\tilde{\mathbf{W}} \tilde{\mathbf{W}}^{T})^{\frac{1}{2}}+\mathcal{O}(|\mathbf{\rho}|^{2}),\]

_where \(\mathbf{\mathcal{G}}\) is a rank-3 tensor with components \(\mathbf{\mathcal{G}}_{k}^{s,r}\equiv\mathbf{\mathcal{G}}_{ij}^{s,r}\) that are defined as:_

\[\mathbf{\mathcal{G}}_{ij}^{s,r}(\mathbf{x})=\frac{\gamma C_{ij}\sqrt{ \omega_{r}\omega_{s}}}{2(\omega_{r}+\omega_{s})}[x_{j}(S_{sj}^{ij}x_{s}\delta _{i}^{r}-S_{rj}^{ij}x_{r}\delta_{i}^{s})\\ +x_{i}(S_{is}^{ij}x_{s}\delta_{j}^{r}-S_{ir}^{ij}x_{r}\delta_{j}^{ s})], \tag{15}\]

_where \(\delta_{s}^{r}\) is the Kronecker delta function, \(\omega_{i}:=1-\gamma/s_{i}\), and \(S_{rs}^{ij}:=\kappa_{ij}/s_{r}+1/s_{s}\) (\(r,s,i,j\in[1,n]\)). The action of the tangential gradient in the representation space is a small rigid-body rotation around the origin. In this rotation, the representation vectors stay within the column-space of \(\tilde{\mathbf{W}}^{T}\), which itself stays fixed during the tangential updates. Additionally, the angular displacement of the representation \(\mathbf{h}_{s}(=\mathbf{U}\mathbf{v}_{s})\) toward \(\mathbf{h}_{r}(=\mathbf{U}\mathbf{v}_{r})\) is:_

\[\Delta\varphi_{sr}(\mathbf{x};\tilde{\mathbf{\theta}}+\mathbf{N}\mathbf{\rho})=\eta\sum_{k=1}^{ K}\rho_{k}\mathbf{\mathcal{G}}_{k}^{s,r}(\mathbf{x})+\mathcal{O}(|\mathbf{\rho}|^{2}). \tag{16}\]

In deriving the above, the gradient was approximated to the first order of the deviation from the manifold, and subsequently the tangential projector operator was applied (see Section A.6 of the Appendix). Here \(\mathbf{\mathcal{G}}\) is a tensor that, for a given direction in the normal space (corresponding to \(k\equiv ij\)), returns a rotation generator via the skew-symmetric matrix \(\mathbf{\mathcal{G}}_{k}^{\,:\,:\,:\,:}\) (note \(G_{k}^{s,r}=-G_{k}^{r,s}\)).

#### 2.4.3 Drift as a diffusion process

As we saw in Section 2.4.1, the mean-reverting property of the gradient confines the normal deviations near the manifold (Fig.2 - Left). The movements in the tangent space, however, face no such mean-reverting property and hence can diffuse freely on the manifold in a random-walk fashion. As theorem 2.6 showed, the tangential component of the gradient is proportional to the deviation from the manifold (\(\mathbf{\rho}\)), such that it vanishes on the manifold (see the schematics in the middle and right panels of Fig.2). We can find an effective diffusion process for the movements on the manifold, which because of the above correspondence, is expected to depend on the statistics of the normal deviations. Specifically, over large timescales, Eq.13 can be approximated as the following continuous-time SDE, describing the evolution of the the parameter vector on the manifold:

\[d\tilde{\mathbf{\theta}}=(2\mathbf{D}_{\theta}/\eta)^{\frac{1}{2}}\,d\mathbf{B}_{t}. \tag{17}\]

In the above, \(\mathbf{D}_{\theta}=(1/2)\langle\Delta\tilde{\mathbf{\theta}}\Delta\tilde{\mathbf{\theta}}^{T} \rangle_{x,\rho}\) is the diffusion tensor for the parameters (measured over a training step), and \(\mathbf{B}_{t}\) is multi-dimensional Brownian motion. The above SDE is equivalent to the diffusion equation \(\partial c/\partial t=\mathbf{D}_{\theta}\nabla^{2}c\). The above diffusion process is manifested as a rotational diffusion in the representation space. To fully characterize this, it's sufficient to measure the pairwise diffusion rates
between the axes of the ellipsoid \(\{\mathbf{h}_{s}\}_{s\in[1,n]}\). Similar to the definition of \(\mathbf{D}_{\theta}\) above, we define \(D_{sr}\) to be the diffusion rate between the two representations \(\mathbf{h}_{s}\) and \(\mathbf{h}_{r}\) based on the mean squared of the angular displacement between the two vectors, i.e.:

\[D_{sr}:=\frac{1}{2}\langle\Delta\varphi_{sr}^{2}\rangle_{x,\rho}=\frac{\eta^{2 }}{2}\sum_{k,l=1}^{K}\langle\rho_{k}\rho_{l}\rangle\,\bar{\mathbf{G}}_{k,l}^{s,r}, \tag{18}\]

where we defined \(\bar{\mathbf{\bar{Q}}}_{k,l}^{s,r}:=\langle\mathbf{\mathcal{G}}_{k}^{s,r}\mathbf{\mathcal{G }}_{l}^{s,r}\rangle_{x}\). The right-hand side results from replacing \(\Delta\varphi_{sr}\) from Eq.16 and taking the average. The total diffusion for the representation \(\mathbf{h}_{s}\) can be derived by summing up the diffusion rates along different directions: \(D_{s}=\sum_{r=1}^{n}D_{sr}\).

Eq.18 suggests the diffusion rate between two representation vectors is an aggregation of terms resulting from deviations in different directions in the normal space. In this equation, \(\eta^{2}\bar{\mathbf{\bar{Q}}}_{k,l}\) can be thought of as the diffusion per direction (or more accurately direction pair \(k,l\)) in the normal space, and \(\langle\rho_{k}\rho_{l}\rangle\) is the covariance of the fluctuations associated with those directions. As both \(\bar{\mathbf{\bar{Q}}}\) and \(\langle\mathbf{\rho}\mathbf{\rho}^{T}\rangle\) depend on the fourth-moment of the stimuli distribution, the diffusion is also a function of the fourth-moment of the input distribution, as we'll see later in the paper.

### Numerical Simulations

Alongside the analytical derivations, we also performed numerical simulations in which we measured the drift in a neural network undergoing continual SGD training. We first numerically validated the equations of the manifold by verifying that after enough training steps the network gets close enough to the theoretical manifold. Next, to measure the drift, we initialized multiple (\(>10^{4}\)) realizations of the network, all starting from a fixed point on the manifold but undergoing different SGD sampling schemes during the training. Following a transitory phase, we studied the over-time trajectories of the hidden layer activation for different trial stimuli. This was done by measuring the fluctuation (variance) of the representation norm, and the angular displacements. The diffusion coefficient was estimated as half of a linear fit slope to the mean squared angular displacements aggregated from all the realizations.

## 3 Drift under Gaussian stimuli

In this section, we present complete analytical results for a case where the stimuli are drawn randomly and independently from a standard n-dimensional Gaussian distribution, i.e. \(x_{i}\sim\mathcal{N}(0,1)\). Since the input covariance is identity, \(\{\mathbf{v}_{i}\}\) form an orthonormal basis for \(\mathbb{R}^{n}\), all corresponding to the eigenvalue \(s_{i}=\langle x_{i}^{2}\rangle=1\). Additionally, we have \(\langle x_{i}x_{j}x_{p}x_{q}\rangle=\delta_{ij}^{pq}+\delta_{ij}^{qp}+\delta_ {ij}^{qp}\). The Hessian eigenspace can be found from proposition 2.4, where we can show \(\kappa_{ij(i>j)}=1\), \(\kappa_{ii}=0\), and \(\kappa_{ij(i<j)}=-1\), for \(i,j\in[1,n]\). Hence the eigenspace corresponds to three sets with \(\mathbf{Z}_{ij(i>j)}=\frac{1}{2\sqrt{\omega}}(\mathbf{v}_{i}\mathbf{v}_{j}^{T}+\mathbf{v}_{j} \mathbf{v}_{i}^{T})\), \(\mathbf{Z}_{ii}=\frac{1}{\sqrt{2\omega}}\mathbf{v}_{i}\mathbf{v}_{i}^{T}\), and \(\mathbf{Z}_{ij(i<j)}=\frac{1}{2\sqrt{\omega}}(-\mathbf{v}_{i}\mathbf{v}_{j}^{T}+\mathbf{v}_{j} \mathbf{v}_{i}^{T})\), with eigenvalues \(\lambda_{ij(i\geqslant j)}=2(1-\gamma)\) and \(\lambda_{ij(i<j)}=2-\gamma\) respectively (\(\omega:=1-\gamma/s_{i}\)).

The components of the fluctuation matrix can be calculated from Eq.10 as:

\[\langle\rho_{ij}^{2}\rangle=\eta\gamma^{2},\;\text{for}\;i\geqslant j, \tag{19}\]

where the rest of the components are zero (see Section A.5.1). Hence, the fluctuation in the representation norm of an arbitrary unit-length stimulus (Eq.12) becomes:

\[\sigma_{s}^{2}=\frac{\eta\gamma^{2}}{2}. \tag{20}\]

To find the diffusion coefficients, we will have to first calculate the components of the \(\mathbf{\mathcal{G}}\) tensor from Eq.15, to have:

\[\mathbf{\mathcal{G}}_{ii}^{s,r}=\frac{\gamma}{2\sqrt{2\omega}}x_{i}(x _{s}\delta_{i}^{r}-x_{r}\delta_{i}^{s}), \tag{21}\] \[\mathbf{\mathcal{G}}_{ij(i>j)}^{s,r}=\frac{\gamma}{4\sqrt{\omega}}[x _{j}(x_{s}\delta_{i}^{r}-x_{r}\delta_{i}^{s})+x_{i}(x_{s}\delta_{j}^{r}-x_{r} \delta_{j}^{s})]\]

and \(\mathbf{\mathcal{G}}_{ij(i<j)}^{s,r}=0\). Subsequently, the coefficients \(\bar{\mathbf{\bar{Q}}}_{ij,pq}^{s,r}\) can be found by taking the averages of the product terms from the above, leading to:

\[\bar{\mathbf{\bar{Q}}}_{ii,ii}^{s,r}=\frac{\gamma^{2}}{8\omega}(\delta _{i}^{r}+\delta_{i}^{s}),\,\bar{\mathbf{\bar{Q}}}_{ii,jj(i\neq j)}^{s,r}=\frac{- \gamma^{2}}{8\omega}(\delta_{ij}^{rs}+\delta_{ij}^{sr})\] \[\bar{\mathbf{\bar{Q}}}_{ij,ij(i>j)}^{s,r}=\frac{\gamma^{2}}{16 \omega}(\delta_{i}^{r}+\delta_{i}^{s}+\delta_{j}^{r}+\delta_{j}^{s}+2\delta_{ ij}^{rs}+2\delta_{ij}^{sr})\]

and the other components are zero. Replacing these and the fluctuation terms in Eq.18 results in the diffusion rate between the representations of two arbitrary \(s\) and \(r(\neq s)\):

\[D_{sr} =\frac{\eta^{3}\gamma^{2}}{2}(\sum_{i=1}^{n}\bar{\mathbf{\bar{Q}}}_{ii,ii}^{s,r}+\sum_{i,j=1,i>j}^{n}\bar{\mathbf{\bar{Q}}}_{ij,ij}^{s,r})\] \[=\frac{\eta^{3}\gamma^{4}}{2(1-\gamma)}(\frac{1}{4}+\frac{n}{8})= \frac{1}{16}\frac{\eta^{3}\gamma^{4}}{1-\gamma}(n+2) \tag{22}\]

Note the summation was performed only over indices where the fluctuation covariance was non-zero. Since the diffusion is isotropic, the total diffusion for representation

Figure 2: Schematics showing (left) the probability distribution of the fluctuations outside the manifold, (middle) the gradient vectors on and near the manifold, and (right) the tangential component of the gradient.
becomes:

\[D_{s}=(n-1)D_{sr}=\frac{1}{16}\frac{\eta^{3}\gamma^{4}}{1-\gamma}(n-1)(n+2) \tag{23}\]

## 4 Drift under a frequently presented stimulus

As discussed in the previous section, an SGD-induced drift could occur even with an isotropic background stimuli. In this section, we will study how the presence of a frequent stimulus in the environment influences the drift. We consider a case where in addition to the background Gaussian stimuli, there is a relatively more frequent stimulus, \(\mathbf{a}\), that is presented with probability \(\alpha\), i.e.:

\[\mathbf{x}=\left\{\begin{array}{ll}\mathbf{a}&\text{Pr}=\alpha\\ \mathcal{N}(\mathbf{0},\mathbf{I}_{n})&\text{Pr}=1-\alpha\end{array}\right. \tag{24}\]

(note the previous case is equivalent to \(\alpha=0\)). Without loss of generality, we take \(\mathbf{a}\) to be along the first axis. The second and fourth moments of the input distribution become \(\langle x_{i}x_{j}\rangle=\alpha|\mathbf{a}|^{2}\delta^{11}_{ij}+(1-\alpha)\delta^ {j}_{i}\) and \(\langle x_{i}x_{j}x_{p}x_{q}\rangle=\alpha|\mathbf{a}|^{4}\delta^{111}_{ijpq}+(1- \alpha)(\delta^{iq}_{ip}+\delta^{pq}_{ij}+\delta^{pq}_{ij})\) respectively. For simplicity, we will also assume \(|\mathbf{a}|=1\). The eigenvalues of \(\mathbf{\Sigma}_{\mathbf{x}}\) will be \(s_{a}=1\) and \(s_{b}=1-\alpha\), corresponding to eigenvectors \(\mathbf{v}_{a}=\mathbf{a}\), and orthonormal vectors \(\{\mathbf{v}_{b}\}\perp\mathbf{a}\) for \(b\in[2,n]\) respectively. The fluctuation of representation norm for unit-length stimuli \(\mathbf{a}\) and \(\mathbf{b}\) (\(\perp\mathbf{a}\)) can be found by replacing \(\langle x_{a}^{4}\rangle=3-2\alpha\) and \(\langle x_{b}^{4}\rangle=3(1-\alpha)\) in Eq.12, to get:

\[\sigma_{a}^{2}=\frac{\eta\gamma^{2}}{2}(1-\alpha),\quad\sigma_{b}^{2}=\frac{ \eta\gamma^{2}}{4}\frac{2+\alpha}{(1-\alpha)^{2}}. \tag{25}\]

It is easy to show that \(\sigma_{a}^{2}\leqslant\sigma_{b}^{2}\) irrespective of \(\alpha\), which suggests the fluctuation of the frequent stimulus representation is smaller than that of a background stimulus.

Acknowledging the symmetry within the space of background stimuli, we only need to find two diffusion coefficients \(D_{ab}=D_{ba}\) and \(D_{bc}=D_{cb}\) for \(a=1\) and any \(b\neq c\in[2,n]\) (without loss of generality, we take \(b=2\) and \(c=3\)). To calculate these coefficients, we need to perform the summation in Eq.18 which is taken over \((ij,pq)\) indices. Despite many terms being zero due to the nature of the \(\mathbf{\mathcal{G}}\) tensor and the fourth moment of the input, the full summation could lead to bulky equations which we avoid here and only present approximate results. Specifically, for large \(n\), many of the terms in the summation can be ignored (Section A.7.1). Additionally, we first consider the regime of small \(\alpha\), which is a perturbation to the previous Gaussian stimuli case for which \(\alpha\) was zero. As a first order correction, it is sufficient to sum over indices for which the unperturbed case had non-zero \(\langle\rho_{ij}\rho_{pq}\rangle\) or \(\mathbf{\widetilde{\mathcal{G}}}^{s,r}\) components. This limits the summation to indices \((ij,ij)_{i>j}\) for which either \(i\) or \(j\) is equal to \(r\) or \(s\) (see detail in Section A.7.1). If we take \(d(>c)\) to be an index within the background subspace (i.e. \(d=4\)), the diffusion summings simplify as below:

\[D_{ba} \approx\frac{n\eta^{2}}{2}(\langle\rho_{db}^{2}\rangle\mathbf{ \widetilde{\mathcal{G}}}^{a,b}_{db,db}+\langle\rho_{da}^{2}\rangle\mathbf{ \widetilde{\mathcal{G}}}^{a,b}_{da,da})\quad(\alpha\ll 1,n\gg 1)\] \[\approx\frac{n\eta^{3}\gamma^{4}}{32}[(1+2\alpha)(1+(1+\frac{ \alpha}{2})(1+\alpha)]\] \[=\frac{n\eta^{3}\gamma^{4}}{16}(1+\frac{7\alpha}{4})\] \[D_{bc} \approx\frac{n\eta^{2}}{2}(\langle\rho_{db}^{2}\rangle\mathbf{ \widetilde{\mathcal{G}}}^{c,b}_{db,db}+\langle\rho_{dc}^{2}\rangle\mathbf{ \widetilde{\mathcal{G}}}^{c,b}_{dc,dc})\] \[\approx\frac{n\eta^{3}\gamma^{4}}{32}[(1+2\alpha)(1+\alpha)+(1+2 \alpha)(1+\alpha)]\] \[=\frac{n\eta^{3}\gamma^{4}}{16}(1+3\alpha) \tag{26}\]

In the above, we replaced the following quantities calculated for small \(\alpha\) and \(\gamma\) (see Section A.7.3):

\[\langle\rho_{db}^{2}\rangle=\langle\rho_{dc}^{2}\rangle=\eta\gamma^{2}(1+2 \alpha),\quad\langle\rho_{da}^{2}\rangle=\eta\gamma^{2}(1+\frac{\alpha}{2})\]

\[\mathbf{\widetilde{\mathcal{G}}}^{a,b}_{db,db}=\frac{\gamma^{2}}{16},\quad\mathbf{ \widetilde{\mathcal{G}}}^{a,b}_{da,da}=\mathbf{\widetilde{\mathcal{G}}}^{c,b}_{db, db}=\mathbf{\widetilde{\mathcal{G}}}^{c,b}_{dc,dc}=\frac{\gamma^{2}}{16}(1+\alpha)\]

Total diffusion coefficients for stimuli \(a\) and \(b\) become:

\[D_{a} =(n-1)D_{ab}\approx\frac{n^{2}\eta^{3}\gamma^{4}}{16}(1+\frac{7 \alpha}{4})\quad(\alpha\ll 1,n\gg 1)\] \[D_{b} =D_{ba}+(n-2)D_{bc}\approx\frac{n^{2}\eta^{3}\gamma^{4}}{16}(1+3 \alpha). \tag{27}\]

which shows \(D_{b}\geq D_{a}\).

In Section A.7.2, we also perform analytical derivation of the diffusion under \(n\gg 1\) and \(\alpha\gg\gamma\) (specifically, with respect to the eigenvalues we assume \(s_{a},s_{b},s_{a}-s_{b}\gg\gamma\)). The results of those calculations are as follows:

\[D_{a} \approx\frac{n^{2}\eta^{3}\gamma^{4}\langle x_{a}^{2}x_{c}^{2} \rangle\langle x_{b}^{2}x_{c}^{2}\rangle}{128s_{b}^{5}}[1+3\frac{s_{b}}{s_{a}} +2(\frac{s_{b}}{s_{a}})^{2}+\frac{4(\frac{s_{b}}{s_{a}})^{2}}{1+\frac{s_{b}}{ s_{a}}}]\] \[=\frac{n^{2}\eta^{3}\gamma^{4}}{16(1-\alpha)^{3}}[\frac{1-\frac {7}{4}\alpha+\frac{15}{16}\alpha^{2}-\frac{1}{8}\alpha^{3}}{1-\frac{\alpha}{ 2}}]\]

\[D_{b} \approx\frac{n^{2}\eta^{3}\gamma^{4}\langle x_{b}^{2}x_{c}^{2} \rangle^{2}}{16s_{b}^{5}}=\frac{n^{2}\eta^{3}\gamma^{4}}{16(1-\alpha)^{3}} \quad(\alpha\gg\gamma,n\gg 1) \tag{28}\]

where in the equalities we replaced \(s_{a}=1\), \(s_{b}=1-\alpha\), and \(\langle x_{a}^{2}x_{c}^{2}\rangle=\langle x_{b}^{2}x_{c}^{2}\rangle=1-\alpha\). The term inside the second line brackets in \(D_{a}\) is always smaller than one. Hence, we have \(D_{b}\geq D_{a}\).

In Figure 3, we plot the diffusion and fluctuations coefficients for \(\alpha\in[0,0.6]\) for stimuli \(a\) and \(b\). We see that, consistent with the above results, the frequent stimulus drifts at a relatively slower rate and has a smaller fluctuation irrespective of \(\alpha\). Additionally, there is an excellent match
between the theoretical and the simulations results for both fluctuations and diffusion coefficients. Finally, in the bottom panel of the same figure, we plotted the trajectories of the representations over time for \(n=p=3\). The lower fluctuation and diffusion rates for the more frequent stimulus could be visually observed from the smaller point cloud for this stimulus.

## 5 Discussion

In a two-layer neural network model of the olfactory system, we show, using theory and simulations, that the stochasticity in SGD online learning could result in a drift of stimuli representation over time, even after the training is complete and no measurable change in the performance is observed. We analytically demonstrate the dependency of the drift on the input distribution and, in particular, show that a frequently presented stimuli drifts at a relatively slower rate. This finding is consistent with experimental observations in the piriform cortex (Schoonov et al., 2021).

We studied the learning dynamics in the high-dimensional space of network parameters. In this space, drift can be considered as any movement tangential to the manifold of minimum-loss, the effects of which aggregate over time to create an effective diffusion process. Orthogonal to this is the fluctuation outside the manifold that is determined by the mean-reverting property of the gradient. For the tangential gradient to have a non-zero value, an orthogonal deviation from the manifold was necessary (see Eq.14 and Figure 2). In a way, this makes the diffusion on the manifold a second order phenomenon, and that explains why the amount of diffusion depends on the fluctuation covariance (Eq.18). We showed diffusion (random walk) as a mechanism by which the representational drift can happen. In this sense it is different from the term drift in Physics.

In the representation space, the effects of the fluctuation and tangential movements are deformations and rigid-body rotations of the space respectively. If we consider the representation of a given stimulus over time, its trajectory essentially consists of two parts: a random-walk movement on a high-dimensional sphere, and simultaneously, a mean-reverting fluctuating process that causes deviations on and outside the sphere (i.e. changing norms and pair-wise angles of representations). This can be observed from the bottom panel in Figure 3. The lower diffusion rate observed for a more frequent stimulus suggests that the rigid-body rotation is such that on average it rotates the frequent stimulus to a lesser degree. In three dimensions, this could be imagined by, for example, the axis of rotation being on average closer to the representation of the frequent stimulus. The lower diffusion and fluctuation for the more frequent stimulus could be traced to the mechanism by which the gradient tends to preserve the representation and the output of a more rehearsed stimuli to a greater extent compared to other stimuli.

Our work relates closely to the recent study by Qin et al. (2023). They studied the drift in a one-layer neural network model of sensory systems with a similarity matching objective and a Hebbian/anti-Hebbian learning rule. Similarly to our results, they showed that the representations undergo a rotational diffusion process. This is not surprising as both of our models have objective functions with a degenerate solution that has a rotational symmetry. In our case, the L2 regularization creates that symmetry. Qin et al. (2023) also demonstrated a stimulus dependency of the diffusion coefficient. In their case the drift for a given eigenvector direction is inversely related to its eigenvalue. Additionally, in that study, the drift is produced by injected synaptic noise and goes away when this noise is close to zero. Here, we used a two-layer neural network model with the MSE loss objective to study drift under the condition of no synaptic noise. We showed that the sampling noise due to the SGD stochasticity and under no synaptic noise is enough to demonstrate the stimulus dependency of the drift rate.

We were able to analytically solve the drift in a simple linear case. Although each layer on such a network is linear, the resulting output is a product of the network weights. As a result, linear networks, despite their apparent simplicity, have nonlinear learning trajectories that could lead to rather com

Figure 3: Plots of (left) diffusion, and (right) fluctuation for representations of a frequent and a background stimulus, denoted by \(a\) and \(b\) respectively. \(\alpha\) is the probability of frequent stimulus, and \(m=n=10,p=20\), \(\gamma=0.04,\eta=0.005\), \(|\mathbf{a}|=1\). (bottom) Over time history of representations for three trial stimuli after \(2.2\times 10^{5}\) training steps. \(n=p=3\), \(\gamma=0.1\), \(\eta=0.1\), \(|\mathbf{a}|=1\) and \(\alpha=0.5\), and \(\mathbf{c}\) is another background stimulus.
plicated and interesting behaviors, especially as the number of layers increases (Saxe et al., 2013; Li & Sompolinsky, 2021). Our framework could be applied to more general cases as long as the manifold represents the redundancy in the model. The detail of the drift, however, may depend on the specific features of the setup such as the noise type, etc. Future steps of our work may include nonlinear and deep networks with recurrent connections in addition to a non-stationary data-distribution.

## Acknowledgements

We are grateful for fundings from the Swartz Foundation, and the National Institute of Health (U19NS112953-01).
