SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To ease interpretation, we again focus on the case of scalar regression. Corollary 3.1 shows that, for small  $\gamma$ , only one singular value grows while the others remain small constants. More precisely, when  $\gamma \sim o(\| \beta_{*}\|)$ , the training model is approximately rank one, with only one spiked singular value. Conversely, when  $\gamma \sim \Omega (\| \beta_{*}\|)$  the low-rank structure disappears. This result perfectly captures the rich learning regime at low initialization where the hidden layer learns the defining feature of the problem, whereas in the lazy regime (large  $\gamma$ ) the singular values of the matrix hardly move and no structure is present in  $W_{1}$ .
Our analysis removes the assumptions of balanced/spectral initialization and whitened data of previous works studying the evolution of singular values [Saxe et al., 2014, Gidel et al., 2019], thus allowing us to reveal the dependence on the scale of initialization.
3.2 Description of the dynamics
So far we have described the structure of the parameters of the neural network at convergence. Here, we show that the dynamics of the singular values of  $\beta$  enjoy a very particular property: it satisfies a mirror flow [Alvarez et al., 2004] with a mirror potential that can be written explicitly.
Theorem 3.2. [Dynamics of the flow] With the same notations as in Theorem 3.1,
(a) Mirror on singular values: The singular values of  $\beta$ , denoted by  $\mathbf{D}_{\beta}$ , follow the mirror flow
$$
\mathrm {d} \nabla \Psi_ {\gamma} (\mathbf {D} _ {\beta}) = - \nabla_ {\mathbf {D} _ {\beta}} \mathcal {L} \mathrm {d} t,
$$
where the potential writes  $\Psi_{\gamma}(\mathbf{D}_{\beta})\coloneqq \mathrm{tr}\left(\frac{1}{2}\mathbf{D}_{\beta}\sinh^{-1}(\mathbf{D}_{\beta} / \gamma) - \sqrt{\mathbf{D}_{\beta}^2 + \gamma^2}\right)$ .
(b) Mirror on  $\beta$ . If  $k = 1$ , the dynamics of  $\beta$  can be characterized as a mirror flow
$$
\mathrm {d} \nabla \psi_ {\gamma} (\beta) = - \left[ \gamma + \sqrt {\| \beta \| ^ {2} + \gamma^ {2}} \right] ^ {1 / 2} \nabla \mathcal {L} (\beta) \mathrm {d} t, \tag {3.3}
$$
where the potential writes  $\psi_{\gamma}(\beta) := \frac{2}{3}\left[\sqrt{\|\beta\|^2 + \gamma^2} + \gamma\right]^{3/2} - 2\gamma\left[\sqrt{\|\beta\|^2 + \gamma^2} + \gamma\right]^{1/2}$ .
Mirror on Singular Values. For vector regression, GF on the parameters induces a continuous-time mirror descent (which we also refer as mirror flow) with the hyperbolic entropy function [Ghai et al., 2020]. This extends Arora et al. [2019b]'s characterization of the evolution of singular values when the initialization is balanced. Note that our result does not fully characterize the evolution of the system, since the characterization of the singular vectors is absent. Still, some interesting comments can be made. In the rich regime in which  $\gamma \rightarrow 0$ , the hyperbolic entropy  $\Psi_{\gamma} \sim (-\ln \gamma)\| \beta \|_*$ . Thus, informally, for small  $\gamma$  the gradient flow on parameters is approximately equivalent to a mirror flow on the nuclear norm. This is reminiscent of the case of diagonal linear networks where such an equivalence is proven rigorously [Pesme and Flammarion, 2023] and is known to lead to an incremental saddle-to-saddle dynamics [Li et al., 2021, Jacot et al., 2022, Berthier, 2022].
In the case of whitened data, i.e.,  $X^{\top}X = I$ , we show that the singular vectors of  $\beta$  are stationary (see Appendix C.11 for details). Therefore, the mirror flow on the singular values characterizes the entire system. In this case, we can even provide an exact expression for the evolution of the singular values
(Appendix C.12) by solving a matrix Ricatti equation [Bittanti et al., 1991]. In the limit of  $\gamma \rightarrow 0$ , we can show that, beyond the case of balanced initialization [Saxe et al., 2014, Gidel et al., 2019], the singular values are learned in a hierarchical manner. When  $\gamma \rightarrow 0$  and with appropriately rescaled time, the limiting trajectory for the  $i^{th}$  singular value  $\sigma_{i,\beta}$  can be seen as the jump process
$$
\sigma_ {i, \boldsymbol {\beta}} \left(\ln \left(\frac {1}{\gamma}\right) t\right) = \sigma_ {i, \boldsymbol {\beta} _ {*}} \mathbb {1} \left(t > \frac {1}{2 \sigma_ {i , \boldsymbol {\beta} _ {*}}}\right),
$$
where  $\sigma_{i,\beta_*}$  is the  $i^{th}$  singular value of  $\beta^*$ . Each singular value is activated at time  $-\ln (\gamma)\left(2\sigma_{i,\beta_*}\right)^{-1}$ . Therefore, we observe an incremental learning process, where the activation begins with the largest singular value and proceeds accordingly.
Mirror descent for scalar regression. The result (b) states that the GF on the parameters  $(W_{1}, W_{2})$  implies a mirror flow on the predictor  $\beta$  with the potential  $\psi_{\delta}$ . To be more precise, the evolution is governed by a mirror flow with the time scaled as a function of  $\|\beta\|$ . This technique of time-warping was proposed in Azulay et al. [2021] for the case of a linear network with a single neuron ( $l = 1$ ) with balanced initializations. In contrast, with a specific initialization shape, we show the existence of a mirror flow for an arbitrary number of neurons and unbalanced initialization of any scale. The existence of a mirror flow is surprising since the reparametrization defining linear networks is not commutative in general [Li et al., 2022]. However, due to the specific initialization we use, this problem can be circumvented by preserving certain commutative properties.
The equivalence with mirror descent enables us to show that  $\mathcal{L}(\beta(t)) = O(1/\gamma t)$  (see Appendix C.8), thus providing a convergence rate for the training loss independent of the conditioning of data, in contrast to Min et al. [2021], Du et al. [2019b]. Note that with decreasing initialization scale  $\gamma$ , the convergence speed diminishes, while according to the results in Theorems 3.1 better implicit bias is achieved. This suggests the existence of a trade-off between optimization and implicit bias already observed in several works [Woodworth et al., 2020], where achieving better quality solutions is linked to slower optimization. In contrast to this behaviour, for the case of balanced initialization [Braun et al., 2022] emphasizes a decoupling between the learning speed and the quality of solutions. Conversely, we stress that in the general setting (e.g., under imbalance) such a decoupling is absent.
3.3 Sketch of the proofs
In this section, we give a short description of the proofs of the main results from the previous sections. The common theme of the following intermediate results is to identify natural invariants of the dynamics, which can be leveraged to understand the hidden mirror structure of the flow.
Lemma 3.1. Consider the dynamics of the gradient flow (2.4) initialized at  $(\mathbf{W}_1(0),\mathbf{W}_2(0)) = (\sqrt{2\gamma} P,0)$ . Let  $\mathbf{Z}_1\coloneqq \mathbf{W}_1P^\top$ $\mathbf{Z}_2\coloneqq \mathbf{PW}_2$  and the residual  $\mathbf{R}\coloneqq X^{\top}(Y - X\mathbf{Z}_{1}\mathbf{Z}_{2})$ , then the evolution of  $(\mathbf{Z}_1,\mathbf{Z}_2)$  is governed by the following ODE
$$
\dot {\boldsymbol {Z}} _ {1} = \boldsymbol {R} \boldsymbol {Z} _ {2} ^ {\top}, \quad \dot {\boldsymbol {Z}} _ {2} = \boldsymbol {Z} _ {1} ^ {\top} \boldsymbol {R}. \tag {3.4}
$$
Furthermore, the dynamics of gradient flow (2.4) is equivalent to (3.4), i.e.,  $(\pmb{W}_1(t),\pmb{W}_2(t)) = (\pmb{Z}_1(t)P,P^\top \pmb{Z}_2(t))$  at any time  $t$ .
Lemma 3.1 derives an equivalent dynamics to equations (2.4). It shows that weights  $\mathbf{W}_1^\top, \mathbf{W}_2$  always stay in the column span of the initialization  $P$ , thus restricting their evolution to a subspace. Going forward, we derive the invariants of the dynamics (3.4).
Lemma 3.2. For the projected matrices given in (3.4), we have the following invariant,
$$
\mathbf {Z} _ {1} ^ {\top} \mathbf {Z} _ {1} - \mathbf {Z} _ {2} \mathbf {Z} _ {2} ^ {\top} = 2 \gamma \mathbf {I}.
$$
This invariant ensures that  $Z_1^\top Z_1$  and  $Z_2Z_2^\top$  commute which is a crucial ingredient in the proofs of Theorems 3.1, 3.2. Now, we derive the evolution of  $\alpha \coloneqq Z_1^{-\top}Z_2$ , which turns out to be the central quantity enabling our result. The lemma below describes certain properties of the evolution of  $\alpha$ .
Lemma 3.3. Let  $\alpha \coloneqq Z_1^{-\top}Z_2$ , we have the following time evolution of parameters:
$$
\dot {\boldsymbol {\alpha}} = \boldsymbol {R} - \boldsymbol {\alpha} \boldsymbol {R} ^ {\top} \boldsymbol {\alpha}, \quad \text {a n d} \quad \boldsymbol {\beta} = \left(1 - \boldsymbol {\alpha} \boldsymbol {\alpha} ^ {\top}\right) ^ {- 1} \boldsymbol {\alpha}.
$$
An outline of the proof of Theorem 3.1. With an ansatz on the potential that it is decomposable in terms of  $Z_{1}, Z_{2}$ , we derive KKT conditions for the constrained optimization problem
$$
\operatorname * {a r g m i n} _ {X Z _ {1} Z _ {2} = Y} \psi_ {1} (Z _ {1}) + \psi_ {2} (Z _ {2}).
$$
Using Lemmas 3.3, we show that  $\alpha$  stays in  $\operatorname{span}(X)$ . We use the isotropic property of the imbalance from Lemma 3.2 to find appropriate functions  $\psi_1, \psi_2$  and finally prove Theorem 3.1. The proofs for theorem 3.1, 3.2 and corollary 3.1 can be found in Appendix B.
4 Further thoughts and perspectives
The previous section provided a deep-dive into the dynamics of the gradient flow, which we complement here with a few steps in the direction of understanding the dynamics with stochastic gradients. We investigate stochastic gradient descent (SGD) by studying its simpler counterpart, label noise gradient descent (LNGD) Blanc et al. [2020].
4.1 The role of noise
It was observed that the noise in stochastic gradient descent has a parameter-dependent shape that induces certain regularization properties [HaoChen et al., 2021, Blanc et al., 2020]. Here, we study the properties of the noise shape induced in the case of parameterization with linear neural networks.
Inspired from the analysis of HaoChen et al. [2021] and the large noise regime described by Pillaud-Vivien et al. [2022] in the context of diagonal neural networks, we design a process driven purely by noise and which carries the same geometric properties as SGD's noise. We consider the scalar  $(k = 1)$  regression problem with  $l = d$  and, through an abuse of notation, denote  $\mathbf{W}_1 = \mathbf{W}$ , and  $\mathbf{W}_2 = \mathbf{a}$ . The noise-driven process which we consider is:
$$
\mathrm {d} \boldsymbol {W} = \left(\mathrm {d} \boldsymbol {B} _ {t}\right) \mathbf {a} ^ {\top}, \quad \mathrm {d} \mathbf {a} = \boldsymbol {W} ^ {\top} \mathrm {d} \boldsymbol {B} _ {t}, \tag {4.1}
$$
where  $B_{t}$  is a  $d$ -dimensional Brownian motion. Details on how this SDE captures the noise of SGD are deferred to Appendix D. We show that, similarly to the rich regime of the gradient flow  $(\gamma \rightarrow 0)$ , this noise also carries a rich spectral bias but for any initialization. Indeed, we have the following result on the SDE dynamics. The proof can be found in Appendix D.
Proposition 4.1. The dynamics (4.1) has the following convergence properties
(a) Variance explosion. The variance of the norms of  $\mathbf{W}$ , a explode, i.e.,
$$
\lim  _ {t \rightarrow \infty} \mathbb {E} \left[\left\| \boldsymbol {W} (t) \right\| ^ {2} \right]\rightarrow \infty , \quad \lim  _ {t \rightarrow \infty} \mathbb {E} \left[\left\| \mathbf {a} (t) \right\| ^ {2} \right]\rightarrow \infty .
$$
(b) Scale divergence. For  $d \geq 5$ , for any  $\alpha > 0$ , we have that,
$$
\lim  _ {t \rightarrow \infty} \mathbb {E} \left[\left\| \boldsymbol {W} (t) \right\| ^ {\alpha} \right]\rightarrow \infty , \quad \lim  _ {t \rightarrow \infty} \mathbb {E} \left[\left\| \mathbf {a} (t) \right\| ^ {\alpha} + \left\| \bar {\mathbf {a}} (t) \right\| ^ {\alpha} \right]\rightarrow \infty .
$$
where  $\bar{\mathbf{a}} := e^{-t} \int_{0}^{t} e^{s} \mathbf{a}(s) \, \mathrm{d}s$  is the exponential moving average of  $\mathbf{a}$ .
(c) Alignment - spectral bias. Denote the  $i^{th}$  row of  $\mathbf{W}$  as  $\mathbf{w}_i$ . Using  $[\mathbf{w}_i, \mathbf{a}] \stackrel{\mathrm{def}}{=} \mathbf{w}_i \mathbf{a}^\top - \mathbf{a} \mathbf{w}_i^\top$ ,
$$
\lim  _ {t \to \infty} \mathbb {E} \left[ | [ \boldsymbol {w} _ {i}, \mathbf {a} ] | \right] \to 0.
$$
For any two vectors  $u, v$ ,  $[u, v]$  denotes the commutator of the vectors: remark that if  $[u, v] = 0$ , then  $u, v$  are aligned, i.e.,  $u = cv$ , for some scalar  $c \in \mathbb{R}$ . First, notice that for  $d = 1$ , the SDE in fact corresponds to the geometric Brownian motion with no drift and the dynamics collapses to zero [Oksendal, 2013]. For dimension  $d \geq 2$ , the proposition states that the system diverges and the weights grow towards infinity. However, despite the fact that the norm grows, the commutator  $[w_i, a]$  goes to zero, indicating that all the rows of  $w_i$  align towards  $a$ . Overall, similarly to the gradient flow in the rich regime, this induces a low rank structure in  $W$ . This phenomenon can be further seen through the evolution of singular values, where the top singular value of  $W$  grows unboundedly, whereas the remaining singular values decay to 0 as depicted in Figure 3a in the Appendix. This sheds some light on how SGD induces a particular parameter-dependent noise which implicitly biases the solutions towards having a low-rank structure of the hidden layer [Andriushchenko et al., 2022].
(b)
Figure 1: (a) Vector regression with orthogonal initialization and scales  $\gamma = 1, 10^{-4}$  and Gaussian initialization with entries from  $N(0, 0.01)$  (b) Scalar regression with Gradient Descent (GD) and Label Noise Gradient Descent (LNGD).
Intricate dynamics in presence of drift. Proposition 4.1 focuses on the process that is solely driven by noise. However, in general SGD also encompasses a drift term which corresponds to the dynamics studied in Section 3. The continuous-time SDE describing the process is
$$
\mathrm {d} \boldsymbol {W} = - \nabla_ {\boldsymbol {W}} \mathcal {L} (\boldsymbol {W}, \mathbf {a}) \mathrm {d} t + \delta (\mathrm {d} \boldsymbol {B} _ {t}) \mathbf {a} ^ {\top}, \quad \mathrm {d} \mathbf {a} = - \nabla_ {\mathbf {a}} \mathcal {L} (\boldsymbol {W}, \mathbf {a}) \mathrm {d} t + \delta \boldsymbol {W} ^ {\top} \mathrm {d} \boldsymbol {B} _ {t},
$$
where  $\delta > 0$  indicates the scale of the noise. The presence of drift quickly complicates the analysis, but intuitively, the noise simplifies the model by inducing a rank reduction, whereas the drift terms prevent the weights from growing unbounded. This noise-driven mechanism relaxes the role of initialization. Empirically this is illustrated in Figure 1b. Gradient descent already exhibits a regularization effect as it increases only one singular value while keeping the others constant. However, gradient descent with label noise [Blanc et al., 2020] enhances this regularization effect by decaying the singular values and promoting low-rank representations. As a result, even for larger initialization scales, we observe the presence of low-rank structures in the hidden layer, unlike in gradient descent. The precise characterization of the this phenomenon is left for future research.
Experiments. We consider a regression problem on synthetic data, with  $n = 5$  samples of Gaussian data in  $\mathbb{R}^{10}$  ( $d = 10$ ) and the labels in  $\mathbb{R}^3$  ( $k = 3$ ) generated by a ground truth  $\beta_* \in \mathbb{R}^{d \times k}$ . We consider a network with width  $l = 200$ . In Figure 1a, we show the evolution of the top-4 singular values of the hidden layer  $W_1$ . We use orthogonal initialization for the network with the two scales of initialization  $\gamma = 1, 10^{-4}$ . Note that, as depicted by Corollary 3.1, for the smaller scale only the first  $k = 3$  singular values are significant in comparison to the remaining  $d - k$  singular values. This shows that the matrix is approximately rank  $k$  and the neurons align along three directions. In contrast, for the larger scale  $\gamma = 1$ , the final weight matrix has rank  $d$ . To complement this, we also consider a Gaussian initialization with variance  $0.01 -$  specifically, we initialize the inner layer with  $d = 10$  Gaussian random vectors in  $\mathbb{R}^l$ . As described when  $l \gg d$ , the initialization is close to the orthogonal initialization. Hence, in this case, we can see that only  $k$  singular values grow and the final model has an approximately rank  $k$  hidden layer. In figure 1b, we depict the time evolution of singular values for GD and LNGD on a scalar regression problem with orthogonal data in  $\mathbb{R}^5$  ( $n, d = 5$ ) and a network with  $l = 200$ . Further details on hyper-parameters can be found in the Appendix.
Extension to non-linear activations. Huh et al. [2023], Andriushchenko et al. [2022] empirically demonstrate a low-rank phenomenon through extensive experiments on deep networks with non-linear activations. However, a comprehensive theoretical comprehension of this behavior remains elusive, despite some efforts addressing these issues [Boursier et al., 2022]. To show that our analysis extends beyond linear activations, we present a toy experiment for ReLU networks (see Figure 2 and further details in Figure 5). Consider a scalar regression problem in a ReLU teacher-student setup. We generate a training set of size 10 sampled from a random Gaussian distribution in  $\mathbb{R}^5$ . The training data  $(x_i, y_i)_{i=1}^{10} \in \mathbb{R}^5 \times \mathbb{R}$  is generated by a teacher ReLU network with 2 neurons  $(w_0, w_1)$ , i.e.,
$$
y _ {i} = a _ {0} \sigma (w _ {0} ^ {\top} x _ {0}) + a _ {1} \sigma (w _ {1} ^ {\top} x _ {1}),
$$
where  $\sigma$  is the ReLU non-linearity. We train a student network with 20 hidden neurons. Note that there are two relevant directions  $w_{0}, w_{1}$  for the student network to learn, therefore we expect the hidden layer to represent these two directions (i.e., a rank 2 hidden layer, and a singular value decomposition with two non-zero singular values). This property is empirically verified in Figure 2. We plot the time evolution of singular values and when initialized at low-scale the network converges to an approximately rank-2 matrix. When initialized at a larger scale, the network weight matrix is high rank and the neurons do not learn the teacher directions.
Perspectives. Learning representations which can be transferred to downstream tasks is a key attribute for the success of deep learning [Bengio et al., 2013, LeCun et al., 2015]. In this work, we present an archetypal problem where for the same predictor in functional space, there exist multiple representations in parameter space, some of which can exhibit a rich structure. This scenario presents a case for going beyond the characterization of implicit bias in the functional space [Parhi and Nowak, 2022] and further studying the implicit bias in the parameter space. Such characterizations facilitate the identification of crucial ingredients in training algorithms that enable effective feature learning.
Limitations and Future Work. This paper tackles the phenomenon of implicit bias, with the aim of furthering the understanding of how neural networks learn in practice. Unfortunately, practical models are highly nonlinear due to their activations and rely on various heuristics to achieve state-of-the-art performance, thus being difficult to grasp mathematically.
This work therefore studies the simplified setting of two-layer linear neural networks. In terms of the assumptions we make, the orthogonality of initialization is only approximately faithful to practical settings where small random weights are used. Nevertheless, we are confident that this requirement can be loosened through a perturbation analysis in the vein of Gidel et al. [2019]. Finally, our dynamical description of the system is yet to be completed in the vector regression case with non-whitened data. A careful set of assumptions is necessary here, and hopefully ones that are weaker than the restricted isometry property used in related works [Li et al., 2018]. Finally, we only partially describe the dynamics in the presence of stochastic noise and giving a full characterization remains a desired objective of future investigations. Further discussion on these aspects is presented in Appendix C.1.
Figure 2: The time evolution of singular values of the hidden layer weights of a 2-layer ReLU network when trained with gradient flow initialized with Gaussian random variables with different scales. We consider a scalar regression problem in a teacher-student setup.
AV is supported by Swiss data science fellowship.
F. Alvarez, J. Bolte, and O. Brahic. Hessian riemannian gradient flows in convex programming. SIAM journal on control and optimization, 43(2):477-501, 2004.
M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion. Sgd with large step sizes learns sparse features. arXiv preprint arXiv:2210.05337, 2022.
S. Arora, N. Cohen, and E. Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In Proceedings of the 35th International Conference on Machine Learning, 2018.
S. Arora, N. Cohen, N. Golowich, and W. Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations, 2019a. URL https://openreview.net/forum?id=SkMQg3C5K7.