SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Aditya Varre
aditya.varre@epfl.ch
Maria-Luiza Vladarean
maria-luiza.vladarean@epfl.ch
Courant Institute of Mathematics, NYU / Flatiron Institute
lpillaudvivien@flatironinstitute.org
Nicolas Flamarion
nicolas.flammarion@epfl.ch
This paper studies the behaviour of two-layer fully connected networks with linear activations trained with gradient flow on the square loss. We show how the optimization process carries an implicit bias on the parameters that depends on the scale of its initialization. The main result of the paper is a variational characterization of the loss minimizers retrieved by the gradient flow for a specific initialization shape. This characterization reveals that, in the small scale initialization regime, the linear neural network's hidden layer is biased toward having a low-rank structure. To complement our results, we showcase a hidden mirror flow that tracks the dynamics of the singular values of the weights matrices and describe their time evolution. We support our findings with numerical experiments illustrating the phenomena.
The most forceful driver of advancements in the field of Machine Learning over the past decades has been the success of deep neural networks. Amongst the striking qualities of these models is the fact that, despite being heavily overparametrized, their optimization consistently yields minima with good generalization properties. A beckoning research direction is thus to unravel the process through which neural networks learn internal representations for a given task [Bengio et al., 2013]. Understanding such phenomena is crucial for lowering the interpretability barrier of these models and developing a principled approach to their training and deployment in practice.
Recent experimental evidence identified one of the likely paths towards achieving these goals as the study of the inherent regularization properties (or implicit biases) of training algorithms [Neyshabur et al., 2014, Zhang et al., 2016]. These observations laid the foundation for a new line of work [see, e.g., Vardi, 2022] whose driving question is which minimum, amongst the many, awaits at the tail end of optimization.
One of the determining factors for the implicit bias of gradient methods is the initialization scale, which controls their operational regime as shown by empirical studies [Chizat et al., 2019]. More precisely, gradient descent with a low-scale initialization is capable of learning rich feature representations from the data. Strikingly, despite overparameterization, the hidden-layer neurons align in the direction of the features [Chizat et al., 2019, Atanasov et al., 2022] and learning of representations reflects in the low-rank structure of the hidden layers. Our work aims to precisely explain this phenomenon and quantify the impact of the initialization scale on feature learning.
Unfortunately, studying such phenomena for the types of neural networks used in practice is mathematically challenging at present due to the non-linearity of their activations. Their less expressive linear counterparts, however, are more tractable and represent a good proxy due to their non-convex
loss landscape and non-linear learning dynamics. Consequently, the study of deep linear networks has received significant amounts of attention over the past years, and spans several important directions, including convergence [Arora et al., 2019a, 2018, Min et al., 2021], learning dynamics [Saxe et al., 2014, Braun et al., 2022] and the implicit bias of optimization algorithms [Azulay et al., 2021]. This work complements these previous approaches by mathematically describing the properties of their parameters at convergence, highlighting the implicit bias phenomenon, and further analyzing the evolution of weight matrices throughout the optimization process.
Specifically, this paper studies overparameterized vector regression problems on two-layer fully-connected linear neural networks. We show the following results when the network is trained with gradient flow (GF).
(i) In Theorem 3.1, we prove that the zero-loss solutions retrieved by the gradient flow are the minimizers of a potential that depends on the initialization scale. Additionally, we provide explicit expressions for the singular values of the hidden layer weights, also as a function of the initialization scale. These characterizations reveal how low-magnitude initialization induces a low-rank structure of the hidden layer.
(ii) In Theorem 3.2, we show that gradient flow on the parameters induces a mirror flow on the singular values. In the specific case of scalar regression, we show that the gradient flow on the weights is equivalent to a mirror flow on the linear predictor. These characterizations give the geometrical structure of the training dynamics of linear neural networks.
(iii) In Proposition 4.1, we design a simple process to analytically describe how stochastic noise in the training algorithm can likewise induce low-rank structures in the weights regardless of the initialization scale.
We proceed by presenting related work in Section 1.1, formalizing the problem setup and assumptions in Section 2, stating and discussing our results in Section 3, and finally, we provide supporting numerical evidence in Section 4.
The first pillar of our work addresses the implicit bias of GF and its stochastic variant in regression problems. One of the hallmarks of bias in this setting is the impact of initialization scale: large initial weights induce a learning regime in which the parameters travel a short distance to convergence and feature learning fails to happen (lazy regime), while small initialization effects a polar opposite behaviour of the system (rich regime) Chizat et al. [2019], Woodworth et al. [2020]. Training and generalization in the lazy regime are well-studied [Jacot et al., 2018, Du et al., 2019a, Arora et al., 2019c, Soltanolkotabi et al., 2017], however this scenario fails to capture the observed behaviour of neural networks in practice [Ghorbani et al., 2019]. While the rich regime more faithfully approximates the feature learning abilities of these models, it is comparatively more challenging to analyze and few results are known. Amongst them are those concerning diagonal linear networks, where a preference towards sparse representations is shown [Woodworth et al., 2020], and a restricted setting of the matrix factorization problem, where the implicit bias leads to low-rank representations [Gunasekar et al., 2017, Arora et al., 2019b, Li et al., 2018]. We similarly study the rich representation learning regime and provide initialization scale-dependent statements on implicit bias for two-layer fully connected linear networks.
Most theoretical results on the implicit bias of GF in overparametrized models rely on the identification of a related mirror flow in a reparametrized space [Gunasekar et al., 2018]. Diagonal linear networks are amenable to this technique and therefore well-studied [Woodworth et al., 2020, Pesme et al., 2021]. For linear fully-connected networks, however, the existence of a mirror flow is not always guaranteed [Li et al., 2022]. To partly alleviate this issue, Azulay et al. [2021] introduce a nonlinear time-rescaling technique and show that for scalar least-squares regression on a two-layer fully connected network with zero-balance initialization, the implicit bias selects low  $\ell_2$ -norm predictors. We prove a similar result under imbalanced initialization controlled by a scale parameter, and characterize the weight matrices independently at convergence, thus presenting a higher-resolution view of the problem.
Other works on linear networks include [Min et al., 2021] where convergence is studied in the presence of weight imbalance and implicit bias results are provided in the functional space; and [Yun
et al., 2021] where tensor networks are studied with the goal of unifying the implicit bias results for linear parameterization. In the case of linear networks, Yun et al. [2021] further show an implicit bias towards minimum  $\ell_2$  linear predictors for vanishingly small initializations. For classification in the case of linear networks, Ji and Telgarsky [2019] show that the weights grow to infinity and the layers of the deep linear network align during the course of optimization. Timor et al. [2023] show a similar phenomenon happens for two-layer ReLU networks.
The second pillar concerns the learning dynamics of linear neural networks. The two-layer case optimized with GF on the square loss has been studied by Fukumizu [1998], Saxe et al. [2014, 2019], Braun et al. [2022]. The common setup of these works is that of zero-balance initialization and whitened data. First, Saxe et al. [2014, 2019] provide expressions for the temporal evolution of singular values of the predictor by assuming decoupled dynamics and a specific data-dependent initialization of the weights. This latter condition is alleviated by the approach of Fukumizu [1998] and Braun et al. [2022], Tarmoun et al. [2021] who solve a matrix Ricatti equation yielding solutions for the weight dynamics in the case where the network initialization has full rank. Finally, Gidel et al. [2019] loosen the whitened data assumption through a perturbation analysis and provide the time-evolution of singular values of the weight matrices. Our work removes the requirement of zero-balanced initialization and full-rank network initialization, and gives formulas for the weights' evolution as a function of the initialization scale. We further provide mirror flows on the weights' singular values and show that components are learned in a hierarchical manner for the case of whitened data.
Related work is further addressed in the following sections, as part of the discussion of results.
2 Preliminaries and problem setup
Notation. Time-dependent variables are written in bold fonts: we drop the  $t$  in  $A(t)$  and simply denote it as  $\mathbf{A}$ . The time derivative of such variables is denoted  $\frac{d}{dt} A(t)$  as  $\dot{\mathbf{A}}$ .
Vector Regression. The set-up is that of standard vector regression problems with inputs  $(x_{1},\ldots ,x_{n})\in \left(\mathbb{R}^{d}\right)^{n}$  and outputs  $(y_{1},\dots y_{n})\in \left(\mathbb{R}^{k}\right)^{n}$  in the so-called overparametrized regime where  $d\geq n$ . Regarding the output dimension, the reader may keep in mind throughout the article that  $k\ll d$ , though the analysis holds for any  $k,d$  pair. In order to learn the input/output rule, we minimize the square loss over a class of parametric models  $\mathcal{H} = \{f_{\boldsymbol{\theta}}(\cdot):\mathbb{R}^d\to \mathbb{R}^k\mid \boldsymbol {\theta}\in \mathbb{R}^p\}$  which we specify in the next paragraph. The train loss therefore can be written as
$$
\mathcal {L} (\boldsymbol {\theta}) = \frac {1}{2 n} \sum_ {i = 1} ^ {n} \left(y _ {i} - f _ {\boldsymbol {\theta}} \left(x _ {i}\right)\right) ^ {2}. \tag {2.1}
$$
Parameterization with a Linear Network. We consider the parametric model of two-layer linear neural networks of width  $l \in \mathbb{N}^*$ : this corresponds to the parametrization  $\pmb{\theta} = (\pmb{W}_1, \pmb{W}_2)$ ,  $\pmb{W}_1 \in \mathbb{R}^{d \times l}$ ,  $\pmb{W}_2 \in \mathbb{R}^{l \times k}$  and  $f_{\pmb{\theta}}(x) = \pmb{W}_2^\top \pmb{W}_1^\top x$ . The model is linear in the input  $x$ , and in terms of expressivity, it is equivalent to the linear class of predictors given by  $f_{\beta}(x) = \beta^\top x$ , with  $\beta = W_1 W_2$ . We henceforth use the symbol  $\beta$  to denote the associated linear predictor of the network. An important consequence of this reparametrization is that the prediction function  $f_{\pmb{\theta}}$  is positively homogeneous of degree 2 in  $\pmb{\theta}$ :  $\forall \lambda \in \mathbb{R}$ , it holds that  $f_{\lambda \pmb{\theta}} = \lambda^2 f_{\pmb{\theta}}$ , as it is the case for two-layer ReLU networks. This property has important consequences in the loss landscape through which  $\pmb{\theta}$  goes.
Train loss. Assume momentarily that  $k = 1$  and denote  $\phi(x) = W_1^\top x \in \mathbb{R}^l$ . It is then clear that the predictor rewrites as  $f_\theta(x) = \langle \phi(x), W_2 \rangle$ . For this reason, we call the hidden layer  $W_1$  the feature layer and the last layer  $W_2$  the weight matrix. We study the overparametrized setting where  $l \gg d$ . Letting  $X^\top \coloneqq [x_1, \ldots, x_n]$  and  $Y^\top \coloneqq [y_1, \ldots, y_n]$ , the loss becomes
$$
\mathcal {L} \left(\boldsymbol {W} _ {1}, \boldsymbol {W} _ {2}\right) = \frac {1}{2 N} \| X \boldsymbol {W} _ {1} \boldsymbol {W} _ {2} - Y \| ^ {2}. \tag {2.2}
$$
For brevity, we ignore the  $N$  in Eq.(2.2) by implicitly rescaling the data as  $(X,Y)\gets (X / \sqrt{N},Y / \sqrt{N})$ .
**Interpolators.** Note that when  $Y \in \operatorname{span}(X)$  and  $X$  is non-degenerate (which occurs with probability one if e.g.,  $X, Y$  are Gaussian and  $d \geq n$ ), there always exists a solution which attains zero loss, i.e.,  $\beta^{*} \in \mathbb{R}^{d \times k}$  such that  $X\beta^{*} = Y$ . We emphasize the fact that there are two levels of overparametrization here: on one hand, when  $d > n$ , the set of zero loss linear predictors  $\mathcal{I}_{\beta} := \{\beta \in \mathbb{R}^{d \times k} \mid X\beta = Y\}$  is typically an affine set of dimension  $(d - n)k$ . On the other hand, since we also reparametrize  $\beta$  as a linear network of width  $l \gg d$ , the manifold of interpolators in the reparametrized space of  $\theta$ , defined by  $\mathcal{I}_{\theta} := \{\theta = (W_1, W_2) \mid W_1W_2 \in \mathcal{I}_{\beta}\}$ , is of dimension  $l(d + k) - nk$ . A natural question, therefore, is to which of these interpolators  $\theta^{*} \in \mathcal{I}_{\theta}$  does a given optimization algorithm converge. This concept is referred to as the implicit bias of an algorithm. The aim of this work is to study that of gradient flow.
Gradient Flow. The dynamics induced in parameter space by running gradient flow on (2.2) is given by
$$
\dot {\boldsymbol {\theta}} = - \nabla_ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta}). \tag {2.3}
$$
We wish to describe the implicit regularization properties of this continuous-time process, which is the vanishing stepsize limit of (stochastic) gradient descent. While the latter algorithms incur additional regularization properties from using non-zero stepsizes [Keskar et al., 2017], the study of GF is an important stepping stone to understanding the implicit bias of gradient-based methods in practice. In terms of  $W_{1}$ ,  $W_{2}$  the dynamics translates to
$$
\dot {\boldsymbol {W}} _ {1} = X ^ {\top} \left(Y - X \boldsymbol {W} _ {1} \boldsymbol {W} _ {2}\right) \boldsymbol {W} _ {2} ^ {\top}, \tag {2.4a}
$$
$$
\dot {\boldsymbol {W}} _ {2} = \boldsymbol {W} _ {1} ^ {\top} X ^ {\top} (Y - X \boldsymbol {W} _ {1} \boldsymbol {W} _ {2}). \tag {2.4b}
$$
We emphasize a crucial point: even if the function  $\beta \rightarrow \| X\beta - Y\|^2$  is convex, its reparametrization in terms of  $W_1$ ,  $W_2$  is not. Non-convexity and non-linearity makes the analysis challenging and a priori it is not even clear whether the time evolution of  $\beta$  can be expressed as a closed system.
Initialization. One of our primary objects of study is the impact of initialization on the behaviour of GF. We describe here our initialization choice, to which we henceforth refer as  $I_{\gamma}$ .
(a) Orthogonal feature layer: We initialize the inner layer such that the rows of  $W_{1}$  are orthogonal and scale with parameter  $\gamma > 0$ . Mathematically, this translates to  $W_{1}(0) = \sqrt{2\gamma} P$  for  $P \in \mathbb{R}^{d \times l}$  in the Stiefel manifold  $V_{d}(\mathbb{R}^{l}) := \{P \in \mathbb{R}^{d \times l} \text{ such that } PP^{\top} = I_{d}\}$ . Initializing with an orthogonal matrix is studied by Pennington et al. [2018], Hu et al. [2020], however from an optimization perspective. Note that when  $l$  is very large, this setting approximates the real-world scenario of initializing the hidden neurons with  $d$  i.i.d. Gaussian vectors in  $\mathbb{R}^{l}$ , which are known to be almost orthogonal.
(b) Zero weight layer: In order to remove any initialization bias from the linear layer, we initialize it at  $\overline{W}_2(0) = 0$ . This can be seen as the limiting case of initializing the weight layer with a very small relative scale  $\overline{\gamma} \ll \gamma$ .
As already mentioned in Section 1.1, existing studies on linear networks assume a "zero-balance initialization", namely that  $\mathbf{W}_1^\top(0)\mathbf{W}_1(0) = \mathbf{W}_2(0)\mathbf{W}_2^\top(0)$  [Saxe et al., 2014, Arora et al., 2019b, Azulay et al., 2021]. This condition introduces the invariant  $\mathbf{W}_1^\top\mathbf{W}_1 = \mathbf{W}_2\mathbf{W}_2^\top$  [Du et al., 2018], which holds for all times  $t \geq 0$ . This balancedness can be seen as a degeneracy assumption on the flow, since it implies that  $\mathbf{W}_1$  has at most rank  $k$  during the entire process, irrespective of the scale of initialization  $\gamma$ . In contrast, we show that depending on  $\gamma$  the feature layer  $\mathbf{W}_1$  is biased (or not) toward a low-rank predictor, thus unveiling a truly rich representation learning regime.
3 Main result: implicit bias and dynamics description
3.1 Implicit bias on the parameters
Non-convex gradient flows are generally not guaranteed to reach global minimizers of the objective and even when they do, such results are difficult to formally prove. Moreover, the existence of many zero-loss solutions with different generalization properties raises the question of which interpolating
network is yielded by training. An elegant answer to such questions is to express the resulting predictor as the optimum among all the possible interpolators of some new, a priori unspecified cost. In addition to the descriptive power of such variational formulations, they express a form of capacity control over the estimator which can be further used to describe its generalization abilities [Bartlett et al., 2020]. The following theorem adds to this series of works, by precisely deriving such a characterization for GF in the setting of linear networks.
Theorem 3.1. Let  $(W_{1}, W_{2})$  be the process that follows the GF equations (2.4a)-(2.4b), initialized according to condition  $I_{\gamma}$ , for some  $\gamma > 0$ . Then
(i) The parameters converge to a global optimum of the loss
$$
\lim _ {t \to \infty} \left(\boldsymbol {W} _ {1} (t), \boldsymbol {W} _ {2} (t)\right) = \left(\boldsymbol {W} _ {1} ^ {\infty}, \boldsymbol {W} _ {2} ^ {\infty}\right) \in \mathcal {I} _ {\boldsymbol {\theta}}.
$$
(ii) The linear predictor  $\beta$  converges to the minimum  $\ell_2$ -norm interpolator
$$
\lim_{t\to \infty}\beta (t) = \operatorname *{argmin}_{\substack{X  \beta = Y}}\left\| \beta \right\|_{F}\stackrel {\mathrm{def}}{=}\beta_{*}.
$$
(iii) We have the following variational characterization of the limiting parameters
$$
\left(\boldsymbol {W} _ {1} ^ {\infty}, \boldsymbol {W} _ {2} ^ {\infty}\right) \in \underset {X \boldsymbol {W} _ {1} \boldsymbol {W} _ {2} = Y} {\operatorname {a r g m i n}} \frac {1}{2} \left\| \boldsymbol {W} _ {2} \right\| _ {F} ^ {2} + \frac {1}{2} \left\| \boldsymbol {W} _ {1} \right\| _ {F} ^ {2} - \gamma \log \left(\det  \left(\boldsymbol {W} _ {1} \boldsymbol {W} _ {1} ^ {\top}\right)\right). \tag {3.1}
$$
Interpretation of the theorem. The theorem is divided into three parts which state that (i) the matrices converge to a zero loss solution, which is a priori non-trivial since the loss in non-convex; (ii) among all the interpolators in  $\mathcal{I}_{\beta}$ ,  $\beta$  converges to the minimum  $\ell_{2}$ -norm interpolator for all  $\gamma > 0$ ; and (iii) among all the interpolators in  $\mathcal{I}_{\theta}$ ,  $(W_{1}, W_{2})$  converge to the ones that minimize a  $\gamma$ -dependent potential. To fully capture the richness of this result, we observe that in the limit of  $\gamma \to 0$ , problem (3.1) informally translates to [Attouch, 1996]
$$
\lim_{\gamma \to 0}\left(\boldsymbol{W}_{1}^{\infty},\boldsymbol{W}_{2}^{\infty}\right)\in \operatorname *{argmin}_{\substack{X\boldsymbol{W}_{1}\boldsymbol{W}_{2} = Y}}\frac{1}{2}\big\|\boldsymbol{W}_{2}\big\|_{F}^{2} + \frac{1}{2}\big\|\boldsymbol{W}_{1}\big\|_{F}^{2}.
$$
This is equivalent, in the space of linear predictors  $\beta$  to the minimum nuclear norm solution  $\beta \in \operatorname{argmin}_{X\beta = Y} \| \beta \|_*$  (which is also the minimum  $\ell_2$ -norm interpolator for the problem we study). We informally derived this interpretation by taking first the limit  $t \to \infty$  and only after  $\gamma \to 0$ . The theorem naturally does not hold if the two limits are reversed, since  $\gamma = 0$  places the initialization at a saddle point of the loss, which is a stationary point of the flow.
With increasing  $\gamma$ , we move towards solutions with a large log det  $(W_{1}W_{1}^{\top})$ , which is a smooth approximation of the rank [Fazel et al., 2003]. Intuitively, this means that solutions with increasing rank are preferred as  $\gamma$  grows. This scale-induced implicit bias is reminiscent of the rich and lazy regimes [Chizat et al., 2019, Woodworth et al., 2020], albeit visible in the space of representations rather than in that of predictors. As such, our result for linear networks is akin to Woodworth et al. [2020]'s, which characterizes the rich and lazy regimes for simpler diagonal linear networks.
Comparison with works on implicit bias of  $\beta$ . Azulay et al. [2021], Min et al. [2021] also study the implicit bias phenomenon in linear networks, however, these results only address the structure of the final predictor  $\beta$  and not that of the factorized problem  $(W_{1}, W_{2})$ . As shown in Theorem 3.1, these works fall short of unveiling all the nuances of the implicit regularization induced by GF in the case of linear networks.
To give a precise example, consider the simplest case of scalar regression  $(k = 1)$  for which both Theorem 3.1 and [Azulay et al., 2021] show that  $\beta$  is biased towards low- $\ell_2$  interpolators. This view is not complete, since there exist many pairs  $(W_{1}, W_{2})$  such that  $W_{1}W_{2} = \beta$ . Theorem 3.1 goes one step further and provides variational characterization of  $(W_{1}, W_{2})$  at convergence. Moreover, it shows that when  $\gamma \rightarrow 0$  all columns of  $W_{1}$  align in the direction of  $\beta$ , thus creating a rank one hidden layer. This is an example of rich representation learning, where  $W_{1}$  is learning the only feature needed to make a prediction.
Implicit bias of the singular values. Proceeding with the description of the spectral bias, we provide a characterization of  $W_{1}$ 's limiting singular values which highlights their dependence on  $\gamma$ .
Corollary 3.1. [Singular values at the limit] Using the same quantities as in Theorem 3.1, denote  $(\sigma_{1}(\pmb{W}_{1}^{\infty}), \ldots, \sigma_{d}(\pmb{W}_{1}^{\infty}))$  and  $(\sigma_{1}(\pmb{\beta}_{*}), \ldots, \sigma_{k}(\pmb{\beta}_{*}))$  the singular values of  $\pmb{W}_{1}^{\infty}, \pmb{\beta}_{*}$  are
$$
\sigma_ {i} \left(\boldsymbol {W} _ {1} ^ {\infty}\right) = \left(\sqrt {\sigma_ {i} \left(\boldsymbol {\beta} _ {*}\right) ^ {2} + \gamma^ {2}} + \gamma\right) ^ {1 / 2}, f o r 1 \leq i \leq k.
$$
$$
\sigma_ {i} \left(\boldsymbol {W} _ {1} ^ {\infty}\right) = (2 \gamma) ^ {1 / 2}, f o r k <   i \leq d.
$$
Discussion. Similar expressions for the singular values of  $W_{2}$  can be derived and are deferred to the Appendix B due to lack of space. In the vanishing initialization limit, only the first  $k$  singular values are activated and  $W_{1}$  resembles a rank  $k$  matrix. Conversely, for large  $\gamma$  all the singular values are approximately equal in scale and  $W_{1}$  resembles an isotropic full-rank matrix.