Year: 2021

Abstract:
In gradient descent, changing how we parametrize the model can lead to drastically different optimization trajectories, giving rise a surprising range of meaningful inductive biases: identifying sparse classifiers or reconstructing low-rank matrices without explicit regularization. This implicit regularization has been hypothesised to be a contributing factor to good generalization in deep learning. However, natural gradient descent is approximately invariant to reparameterization, it always follows the same trajectory and finds the same optimum. The question naturally arises: What happens if we eliminate the role of parameterization, which solution will be found, what new properties occur? We characterize the behaviour of natural gradient flow in deep linear networks for separable classification under logistic loss and deep matrix factorization. Some of our findings extend to nonlinear neural networks with sufficient but finite over-parametrization. We demonstrate that there exist learning problems where natural gradient descent fails to generalize, while gradient descent with the right architecture peforms well.

text: 1 INTRODUCTION
There is plenty of empirical evidence that the choice of network architecture is an important determinant of the success of deep learning (He et al., 2015; Vaswani et al., 2017). The empirical observations are now supported by theoretical work into the role that parameter-to-hypothesis mapping plays in determining inductive biases of gradient-based learning. Unregularized gradient descent can efficiently find low-rank solutions in matrix completion problems (Arora et al., 2019), sparse solutions in separable classification (Gunasekar et al., 2018) or compressed sensing (Vaškevičius et al., 2019). Valle-Pérez et al. (2018) studied deep neural networks and found evidence that the parameter-hypothesis mapping1 is biased towards simpler functions as measured by Kolmogorov complexity. Taken together, these observations and findings have lead the community to hypothesize that
The parameter-to-hypothesis mapping influences the inductive biases of gradient-based learning and may play an important role in generalization.
In parallel to improving architectures, considerable research was done to improve optimization algorithms for deep learning, with a focus on faster convergence and robustness to hyperparameters. Among the most advanced optimization methods are natural gradient descent (NGD) techniques. An intuitive motivation for NGD is that it improves convergence by implicitly lifting the problem from parameter-space, where the loss is non-convex and poorly behaved to the Riemannian manifold
1The mapping between the parameter space and the set of hypotheses as seen on Figure 1
of hypotheses, where the loss is better behaved. From the perspective of inductive biases, the most interesting aspect of NGD is its approximate invariance to reparametrization.
Natural gradient descent eliminates the effect of parameter-to-hypothesis mapping.
These two observations invite questions about the nature of inductive biases in NGD as well as the role of parametrization-dependence in generalization. The first, practical, implication is as follows: if the parameter-to-hypothesis mapping really does play an important role in generalization, then eliminating its influence on the optimization path may be undesirable, and consequently the pursuit of implementing exact NGD in deep architectures may be counterproductive. Secondly, studying the behaviour of NGD in various models and tasks may give us new insights about the importance of parametrization, and could perhaps offer a way to experimentally or theoretically test hypotheses.
In this paper we study the inductive bias of natural gradient descent in deep linear models. These models are particularly suited for our analysis because (a) efficient algorithms exist to calculate exact natural gradients which is otherwise computationally intractable and (b) the inductive biases of Euclidean gradient descent (EGD) in these models have been thoroughly studied and understood.
We make the following contributions: • In linear classification, we show that NGF is invariant under invertible transformations of
data (Theorems 1&2) and as a consequence it cannot recover the `p large margin solutions that EGD tends to converge to. • We further show that (in case of separable classification) when the number of parameters exceeds the number of datapoints, NGF interpolates training labels in a way similar to ordinary least squares or ridgeless regression (Theorems 3&4). • We demonstrate experimentally that there exist learning problems where NGD can not reach good generalization performance, while EGD with the right architecture can succeed. • To perform experiments, we extended the work of Bernacchia et al. (2018) to derive efficient and numerically stable algorithms for calculating exact natural gradients in diagonal networks (Gunasekar et al., 2018) and deep matrix factorization (Arora et al., 2019).
Before stating our main theoretical and experimental results we review some relevant background on parametrization-dependent implicit regularization and natural gradients.

heading: 2 BACKGROUND

heading: 2.1 SEPARABLE CLASSIFICATION WITH DEEP LINEAR MODELS
text: In this article we consider binary classification datasets {(xn, yn), n = 1, . . . , N} separable by a homogeneous linear classifier with a positive margin ( i. e. ∃β∗ s.t. ynx>nβ
∗ ≥ 1 ∀n). (We use the notation X = (x1 · · ·xN )>). In such situation β∗ is not unique and there may be many separating hyperplanes which all achieve 0 training loss - it is up to the inductive biases of the learning algorithm to select one. Soudry et al. (2017) studied the dynamics of unregularized Euclidean gradient descent on logistic loss and found that the iterate β(t) converges to the well-known `2 large margin classifier in direction, that is
limt→∞ β(t) |β(t)| = β∗`2 |β∗`2 | where β∗`2 = argmin β∈RD ||β||2 s.t. ynx>nβ ≥ 1 ∀n.
Importantly, Gunasekar et al. (2018) later showed that this behaviour changes if the gradient descent is performed on a different parametrization. In this paper we will focus on L-layer linear diagonal networks (Gunasekar et al., 2018), where β = w1 w2 . . . wL, using to denote elementwise product. When we adjust parameters w1, . . . ,wL through Euclidean gradient descent, β(t) converges to the ` 2
L large margin separator defined as
limt→∞ β(t) |β(t)| = β∗diag |β∗diag| where β∗diag = argmin β∈RD ||β|| 2 L s.t. ynx>nβ ≥ 1 ∀n.
A remarkable consequence of this is that unregularized gradient descent can find sparse classifiers, without any form of explicit regularization. In fact, this inductive bias is even more sparsity-seeking than the typically used `1 regularization (see e. g. Koh et al., 2007; Tibshirani, 1996). Figure 2 illustrates this behaviour in a 2D example.

heading: 2.2 MATRIX COMPLETION VIA DEEP MATRIX FACTORIZATION
text: The task of matrix completion involves recovering an unknown matrix β∗ ∈ RD×D from a randomly chosen subset of observed entries2. The problem is clearly underdefined: there are infinitely many matrices that match the observed entries. It is common to make additional assumptions about β∗, most commonly that that it has low rank, under which it becomes identifiable.
One approach to matrix completion under the low-rank assumption is based on explicit regularization (e.g. nuclear norm) which leads to a convex optimization problem. Another common approach is matrix factorization using an underparametrized representation β = UV where the sizes of U ∈ RD×R and V ∈ RR×D are restricted to ensure β’s rank is at most R. Learning then proceeds by minimizing the non-convex mean-squared reconstruction error in U, V via gradient descent.
Remarkably, Gunasekar et al. (2017) showed that the gradient-based matrix factorization method tends to converge to low-rank solutions even in the overparametrized setting, i.e. when β = W1W2 where W1 and W2 are full square matrices, without any explicit regularization. This was later extended by Arora et al. (2019), who studied the deep matrix product parametrization of the form β =W1W2 · · ·WL. Arora et al. (2019) ran experiments for different matrix completion tasks varying initialization, depth and number of observations and compared them to minimum nuclear norm solution. When the number of observed entries is large gradient descent in deep matrix factorization models tended to the minimum nuclear norm solution. However, in the interesting case of fewer observed entries, the behaviour was different. Gradient descent pre-
ferred solutions with lower effective rank at the expense of higher nuclear norm. From the evolution of the singular values of β they also concluded that the implicit regularization is towards low rank that becomes stronger as depth grows.
2to simplify presentation we assume the matrices are square, but our arguments hold more generally.

heading: 2.3 NATURAL GRADIENT DESCENT
text: In the next section we briefly introduce some notation and key properties of natural gradient descent (NGD, Amari, 1997; Pascanu & Bengio, 2013). Intuitively, one can think of NGD as a gradient descent method, but not in the Euclidean space (with the Euclidean metric) of parameters, but instead on the Riemannian manifold of probabilistic models the parameters define (equipped with a different metric). More specifically, let’s say that the parameter of interest is θ, where θ defines a probabilistic model p(y|x, θ). We assume that we wish to minimize the log loss under this model, i. e. l(θ,x, y) = − log p(y|x, θ) and L(θ) = ∑N n=1 l(θ,xn, yn). Then, NGD is usually defined as
θ(t+ 1) = θ(t)− ηF−1(θ)∇θL(θ), where (1) F (θ) = EX [EY |X;θ[∇θL(θ)∇>θ L(θ)]] (2)
is the average Fisher information matrix and η is the step size. In the above definition, EY |X;θ is taken over the distribution specified by θ, but distribution with respect to which the expectation EX is calculated can be arbitrarily chosen. In this article we use the empirical distribution of training data, though other choices are possible (Pascanu & Bengio, 2013). We will also consider natural gradient flow (NGF) the continuous limit of NGD, analogously defined as
θ̇ = −F−1(θ)∇θL(θ). (3)
We also note, that F (θ) is not generally invertible, and indeed it will not be in some of the cases we will consider. Therefore, it is more correct to define NGF as any trajectory θt which satisfies
F (θ)θ̇ = −∇θL(θ). (4)
The natural gradient direction is thus only unique within the eigenspace of F (θ). Of all natural gradient directions, one common choice is to use the Moore-Penrose pseudoinverse of F :
θ̇ = −F+(θ)∇θL(θ). (5)
We have seen how in EGD, different parametrization of the same problem leads to drastically different trajectories and optima. However, NGD with infinitesimally small learning rate (i. e. NGF) always follows the same trajectory in model-space and this finds the same optimum, irrespective of how it is parametrized, provided that the parametrization is smooth and locally invertible. Below we formally state this property Amari (1997), alongside a short proof for illustration. Statement (Invariance of NGF under reparametrization). Let w and θ be two parameter vectors related by the mapping θ = P(w) and consider natural gradient flow in w. Assume that (1) the Jacobian J = ∂θt∂wt and (2) F (θt) are both full rank for all t. If wt follows natural gradient flow starting from w0 then θt = P(wt) follows NGF, i. e. it solves θ̇t = −F (θt)+∇θtL(X, θt).

heading: 3 NATURAL GRADIENTS UNDER LOGISTIC LOSS ON SEPARABLE DATA
text: We have seen in Section 2.1 that when trained on separable data with the logistic loss EGD tends to converge to large margin classifiers. To illustrate how NGD differs, we first prove an invariance property which, as we will see, rules out large margin behaviour. We state this property separately when N < D and when N ≥ D in the theorems that follow. We denote the number of data points with N and the number of input features with D. Theorem 1. Let’s assume, that N < D, X is full rank and A is an invertible D × D matrix. Let βt = βt(X,y) be the trajectory of NGF and β ′ t = βt(XA
>,y) (the trajectory of NGF on data XA>). ThenXβ = XA>β′ (with the assumption that β and β′ have equivalent initial conditions). Proof sketch. We use the notation s = Xβ and s′ = XA>β and prove that st = s′t. The full proof can be found in Appendix C.1. Theorem 2. Let βt(X,y) be the trajectory of NGF and let A be a D × D invertible transformation. If N ≥ D, X has full rank and we consider NGF on the transformed data XA>, then A>βt(XA >,y) = βt(X,y) (with the assumption that β and β ′ have equivalent initial conditions). Remark. When N ≥ D and X is full rank, the size of F (β) is D ×D and its rank is D, therefore the Fisher information matrix of β is invertible.
Proof sketch. First let’s say β′t = βt(XA >,y) and v> = β′>A. Then we prove the following: ∇β′L(ynβ′>Axn) = A∇vL(ynv>xn) and F (β′) = AF (v)A>. (6) Hence we get:
β̇ ′ = F (β′)−1∇β′L(β′) and v̇ = F (v)−1∇vL(v). (7)
So if v and β′ have the same initialization, then vt = β′t. Full proof can be found in Appendix C.2. Conclusion. Let ut(X,y) denote the trajectory ofXβt, which is the linear function β > t x evaluated at each of the datapoints xn. Then ut(XA>,y) = ut(X,y).
Proof. ut(XA>,y) = XA>βt(XA >,y) = Xβt(X,y) = ut(X,y)
One special case of this invariance property is invariance to scaling the dimensions of input data (when A is diagonal). Imagine we scale any dimension by a constant a, NGF counteracts it by scaling the corresponding coordinate of β by a−1. We see now why this rules out characterising implicit regularization of NGD as minimizing non-data-dependent norms of β. In particular, it rules out the `p large-margin behaviour we have seen in EGD. Remark. Let A be a D ×D invertible transformation and let β∗(X,y) be the `2 large margin solution, i.e. β∗(X,y) = argmin ‖β‖2 subject to ynβ>xn ≥ 1 ∀n. Then the `2 large margin classifier does not have the invariance property, namely there exists a dataset (X, y) and a transformation A such that A>β∗t (XA >,y) 6= β∗t (X,y). We include a proof by counterexample in Appendix D.
Having ruled out norm-based implicit regularization, it’s natural to consider other statistical methods that exhibit invariance under invertible data transformations. One candidate is ridge-less regression or ordinary least squares (OLS), whose parameter is given by the formula βOLS = (X
>X)−1X>y. As it turns out, the connection between NGD in linear regression and the OLS estimate run deeper than sharing this invariance property. Theorem 3. If N < D and X is full rank, if parameters βt of a linear model follow natural gradient flow under logistic loss, the logits st = Xβt follow an asymptotically linear trajectory with direction vector y. Remark. The Fisher information matrix w.r.t. β is F (β) = X>D(β)X , where D(β) is diagonal with positive elements on the diagonal. We see, that rank(F ) = rank(X) ≤ N , so F is singular, thus several NGF paths are possible. When rank(X) = N , β has D − N degrees of freedom and we did describe β on N dimensions. That’s why we consider s instead of β.
This Theorem follows from the more general Theorem 4 which we will state later.
Informally, when we have more parameters than datapoints, NGD discovers a solution that interpolates the training labels y (encoded as −1s and +1s) perfectly just like ordinary least squares does in this case. Furthermore, if one uses the Moore-Penrose pseudoinverse to calculate the descent direction, i. e. Eqn. (5), then βt converges in direction to the OLS parameter.
In general cases, OLS interpolation and large-margin (LM) methods find qualitatively different solutions in classification tasks. While the LM solution is typically a linear combination of a small subset of training data (the support vectors), in OLS all datapoints are support vectors. As shown in (Hsu et al., 2020), under some conditions this difference disappears in the highly overparametrised regime - when D > N logN . An implication of Theorem 3 is that this phenomenon, known as support vector proliferation, occurs in NGF when D > N . Thus there is a regime where NGF and EGF find qualitatively different classifiers, with different generalisation properties (Hsu et al., 2020).
Theorem 3 provided useful in the context of linear models but it turns out it is relatively straightforward to extend this to a result which holds for non-linear overparametrized models as well. Theorem 4. Let w ∈ RP , P ≥ N be the parameters of a classifier with logits s = s(X;w) ∈ RN . If wt follows natural gradient flow on the logistic loss with labels y and the Jacobian Jt = ∂st∂wt is of full rank, then st grows asymptotically linearly with direction vector y. Remark. If our network is linear J = X , so Theorem 3 is a special case of Theorem 4 indeed. Proof sketch. The main idea is that, since J is full rank, by parametrization invariance of NGD the trajectory of s is determined by the trajectory of the corresponding β.
ṡ = −F−1(s)∇sL(s) (8)
Then we can calculate F (s) which turns out to be diagonal, so we have N independent differential equations. We solve them to get the result. The details of the proof can be found in the Appendix C.

heading: 3.1 EXPERIMENTS
text: In order to validate and illustrate our findings we have run two main simulations, with results presented in Figures 2 and 4. In both experiments we considered the direct parametrization β = w and the diagonal parametrization β = w1 · · · wL (Gunasekar et al., 2018) for different depth L. In order to run these experiments we needed to implement an efficient algorithm for computing natural gradients in these models: naively calculating and then inverting the Fisher information matrix is computationally inefficient and numerically unstable. We therefore developed an algorithm that exploits the structure in the Fisher information matrix, extending the work of Bernacchia et al. (2018) for diagonal networks. The details of our algorithms can be found in Appendix B.2.
In Experiment 1 we illustrated EGD and NGD in a 2D toy classification dataset. Positive and negative classes were generated such that they are separable by the the axis-aligned separator, but there exists a non-axis-aligned separator with a higher margin. Based on the findings of Gunasekar et al. (2018) we expected EGD to find the large margin solution when L is low, and the axis-aligned solution when L is sufficiently large. The results in panels a and b of Figure 2 confirm these predictions. Figure 2c-d illustrate the parametrization-independence of NGD: it converges to the same solution irrespective of parametrization. The solution is different from both the EGD solutions.
In Experiment 2 we focused on generalization performance. We generated a 1000-dimensional dataset with standard Gaussian X , and a sparse ground-truth separator whose first 20 components were set to 1, the rest were 0. Methods with explicit or implicit regularization towards sparse solutions should enjoy good generalization even when N < D. Confirming our expectations, we observed that EGD in diagonal parametrizations (L = 2, L = 3) performed best on this task. The deeper diagonal model (L = 4) was on par with the shallow solution, we expect that our 2 million EGD steps were simply not long enough for the implicit regularization to kick in (Moroshko et al., 2020). The NGD solution on the other hand completely fails to generalize when N < D and does relatively poorly even as N > D. This catastrophic performance is remedied by averaging the Fisher information on a larger dataset - i. .e. using the population Fisher (Amari et al., 2020), but even this variant of NGD fails to match the performance of EGD. The middle panel of Figure 4 validates the predictions of Theorem 4: logits from the model converge to ty. Finally, the right-hand panels of Figure 4 show that NGD was unable to identify the sparse structure, which the diagonal model infers best, and even the shallow model approximately finds.

heading: 4 MATRIX COMPLETION WITH NATURAL GRADIENT DESCENT
text: As we have seen in Section 2.2, EGD in the deep matrix product parametrization β = W1 · · ·WL converges to low-rank solutions. However, when L = 1, i.e. when we run EGD directly on β, the solution we find is trivial: entries of β where we have observation will converge to the observed value, while other entries won’t move. Due to parameter-invariance, NGD cannot differentiate between parametrizations of different depth, it is natural to expect that it will fail the same way as EGD does when L = 1. Let’s look at NGD in matrix completion.
In matrix completion we minimize the squared reconstruction error, which corresponds to the log loss in an isotropic Gaussian observation model with β as mean. In a Gaussian model, the Fisher Information Matrix of β becomes F (β) = 1σ2n I , where σ 2 n is the observation noise. The observation noise σ2n is assumed a constant, and is inconsequential here as it cancels with the 1 σ2n
term in the log loss. Consequently, without loss of generality, we can consider F (β) the identity.
Statement. Let’s apply NGF for the problem of matrix completion. EGF in the direct parametrization (β = w) is equivalent to NGF under any parametrization θ for which J = ∂wt∂θt is full rank.
The proof of the statement can be found in Appendix E. This implies that NGF will completely fail to generalize, i. e. make an accurate prediction of any unobserved entry of the matrix.
Figure 3 illustrates the key property of the dynamics which allows EGD to generalize in deeper parametrizations. Each panel shows values of the neural tangent kernel (NTK) (Jacot et al., 2018), its equivalent object for NGF called the natural NTK (Rudner et al., 2019), or their discretized versions. For matrix factorization the NTK K (θ) is a (D × D) × (D × D) tensor which depends on the parameters θ where ki, j, k, l(θ) measures how much the entry βi,j moves in reaction to a negative loss gradient w.r.t. βk,l. In these visualizations, we set D = 11, and we plot the heatmap of ki,j,5,5. We can see that when we parametrize β directly, the NTK is simply the identity, only the entry β5,5 moves. However, when L = 2, EGD can now respond to the gradient signal at β5,5 by moving entries in the fifth row of W1 or in the fifth column of W2. This, in turn, might result in moving βi,5 or β5,i as well. This explains the cross pattern seen in Figure 3 first panel in the second row. This non-identity NTK is what allows generalization to happen as ’information flows’ from observations to unobserved entries of β. However, in NGF, the natural NTK remains the identity irrespective of parametrization. This is true even in the approximately invariant NGD.
For our Matrix Factorization experiments we had to develope a scalable and numerically stable algorithm for computing the natural gradient. We did this by extending the algorithm of Bernacchia et al. (2018) to matrix factorization. Exploiting the structure of the Jacobian in the deep matrix product parametrization (β = W1 · · ·WL) we calculate the natural gradient w.r.t. Wl as ∇̃WlL = 1LB > l +∇̃βLA+l , where Al = ∏l−1 i=1Wi and Bl = ∏L i=l+1Wi. We note that Ai and
Bi are matrices that are readily computed during the forward and backward pass of reverse-mode automatic differentiation of the loss. The details of the derivation can be found in Appendix B.4.
Using this algorithm, in Figure 5 we experimentally verify that NGD finds a trivial optimum in deep matrix product parametrizations of varying depth. We follow the experimental setup of Arora et al. (2019) and reproduce their results for EGD. We performed an extensive grid search of hyperparameters and found no setting where NGD would achieve non-trivial performance.

heading: 5 SUMMARY AND DISCUSSION
text: Inductive biases of gradient-based learning are driven to a large extent by the way we parametrize our hypothesis. Natural gradient descent (NGD), on the other hand, ignores the parametrization and implicitly optimizes over the manifold of hypothesis. This invited the question whether NGD exhibits any of the useful implicit regularization that EGD has been shown to have. We characterized the behaviour of NGD over logistic loss, and found that in the overparametrized regime, NGD converges to the ordinary least squares interpolant of training labels. This is in contrast with the large-margin-type behaviour EGD exhibits. In experiments we found that in the models we studied, NGD fails to generalize as well as EGD with the right parametrization.

heading: 5.1 OTHER RELATED WORK
text: Approximate NGD algorithms: Since exact NGD is computationally prohibitive, a great deal of research has been devoted to developing approximate NGD algorithms for deep leaning: K-FAC Grosse & Martens (2016); Martens & Grosse (2015) exploits the approximately Kronecker structure of the Fisher information matrix, while, while Bernacchia et al. (2018) start from exact gradient descent in linear neural networks and then apply the formula verbatim to the non-linear case. Another line of work aims at improving the invariance properties of NGD algorithms bringing them closer to ideal of NGF (Song et al., 2018; Luk & Grosse, 2018). Our motivation differs in that are not focused on designing better NGD algorithms, instead we raise the question whether closer approximation of NGF is desirable in the first place. In order to perform experiments that validate our findings we develop efficient exact natural gradient descent algorithms in overparametrized linear models extending the work of Bernacchia et al. (2018).
Convergence Rates for NGD: The main reason for using NGD in deep learning is the intuitive notion it might speed up convergence by virtue of being invariant to parametrization (Amari, 1997; Pascanu & Bengio, 2013; Martens, 2014). This intuition is backed up by theory: Amari (1998) proved fast convergence on a quadratic loss; Bernacchia et al. (2018) proved fast convergence for deep linear models under quadratic loss; more recently, Zhang et al. (2021) gave a proof of fast convergence which holds for a broad class of overparametrized networks and also extends to KFAC; Rudner et al. (2019) analysed NGD in the neural tangent kernel (NTK) regime. Our work differs in that our primary interest is not whether NGD converges fast, but to better understand and illustrate possible trade-offs between fast convergence and generalization.
Generalization of NGD: Wilson et al. (2017) were the first to propose that faster convergence may come at the cost of diminished generalization performance in deep learning. Much like our work, Wilson et al. (2017) provided illustrative examples where different methods reach qualitatively different solutions. They focused on adaptive learning rate algorithms like Adam, but due to the connections between Adam and the empirical Fisher information, one might speculate that their findings would extend to NGD as well Zhang et al. (2019) argued against the notion that NGD may not generalize well, and supported their argument with a generalization bound which holds for both NGD and EGD. However, generalization bounds often fail to predict the empirically observed performance of deep learning (Jiang et al., 2019, see e. g.). In a setting most closely resembling our work Amari et al. (2020) studied generalisation of preconditioned GD for minimising squared loss and found that the optimal preconditioner depends on several factors: EGD generalises better for clean labels, but in scenarios like misspecification or when the labels are noisy, NGD may have an advantage. Finally, Wadia et al. (2021) argued that second order information of the input data - which some second-order optimisation methods can’t utilize well, is key to good generalisation in some neural network architectures. This general connection is related to our Theorems 1 and 2.

heading: 5.2 Q&A
text: Q: How about stochastic gradients? Following Gunasekar et al. (2017; 2018); Arora et al. (2019) we analysed only full-batch gradient descent. This allowed us to prove properties of gradient flow, i. e. in
the limit of infinitesimally small learning rates, which is not a meaningful limit in SGD. This line of work demonstrates that useful inductive biases exist in gradient-based learning even in the absence of gradient noise. Indeed, recent empirical evidence suggests that stochasticity may not be necessary for good generalization in deep networks (see e. g. Geiping et al., 2021). In practice, we expect the question of generalization to be complex, with multiple factors like stochasticity or parametrizationdependence playing a role. We propose that analysing NGD is a useful tool in understanding this complex interplay, as it acts as a form of ablation by eliminating parametrization-dependence.
Q: Does this mean NGD does not generalize well? Not necessarily. We show that there are cases where it does not, but it is possible that in other situations the inductive biases of NGD are more helpful than those of EGD + parametrization, especially when trained on large data. Intuitively, our theorems suggest that NGD may be too efficient at minimising the training loss at the cost of poorer generalisation. However, in our experiments we saw that averaging the Fisher information matrix over test data may remedy this, which would be in line with the practical recommendation of Pascanu & Bengio (2013). Empirical evidence for generalization in exact NGD is sparse due to the computational cost. Some works report good test performance using approximate methods (Grosse & Martens, 2016; Bernacchia et al., 2018) or small models (Pascanu & Bengio, 2013), but since the focus in these works was on demonstrating the usefulness of new methods, it is questionable how thorough these comparisons were. Zhang et al. (2019); Amari et al. (2020) studied generalisation of natural gradient methods theoretically in limited settings and provided some empirical evidence to support their claims. A systematic empirical investigation similar to (Wilson et al., 2017) may be more informative on this question.
Q: Does initialization play a role? Changing the parametrization may influence generalisation in at least three ways: (1) initialization, (2) training dynamics, and (3) constraining the hypothesis space. As weights are often initialized from a parameter-wise independent distribution, these may give rise to a non-trivial and parameter-dependent initial distribution in hypothesis-space. VallePérez et al. (2018) argued that in deep networks, this manifests as a form of simplicity bias. In our models, initialisation has a simlicity bias, too: if matrices W1, . . . ,WL are drawn from an isotropic Gaussian, their product β will be effectively low-rank with an increasing probability as L increases. By replacing EGD by NGD, we only eliminate the influence of parametrization on training dynamics, but the effects of initialization remain. It is therefore important to disentangle relative importance of initialisation (1), and parameter-dependent dynamics (2). To this end, we designed a set of additional experiments, where we controlled the effect of initialization separately from the effects through dynamics. We initialised deep matrix factorisation models by drawing each component matrix W1 as a product of independent Gaussian matrices, then ran EGD. Thus, we were able to create models behaving like a L = 6 layer model at initialization but L = 2 layer model during training. We found that the effect of initialization on generalization performance was negligible compared to the effects of training dynamics (Figure 5.c), at least in deep linear models. We further note that initialization plays a very important role in the limit of infinitely wide networks, too, where initialization scale determines whether the network behaves like a linear kernel machine, or more like the behaviour we describe in finite networks here (Woodworth et al., 2020).
Q: What if you calculate Fisher information on test data? Pascanu & Bengio (2013) noted that in deep learning, averaging the Fisher information over test data, rather than training data seemingly improves performance. In our theorems and experiments we assume averaging over the training data, sometimes referred to as the sample Fisher information (see e. g. Amari et al., 2020) as this makes our proofs tractable. In our high-dimensional sparse classivication experiment in Figure 4 we tested the performance of NGD when the Fisher information is averaged over a large number of samples, called the population Fisher, and we found that generalisation performance improved, but still did not match that of EGD, especially when sparsity-inducing diagonal parametrisations are used.
Q: What about other forms of natural gradients? In addition to the Fisher-Rao natural gradients that we consider here, there are other forms of natural gradients, such as those based on the Wasserstein metric (Li & Montufar, 2018; Arbel et al., 2019). When considering this broader family of natural gradient descent, it is natural to ask if the choice of metric may give rise to different inductive biases in NGD similarly to how different parametrizations effect EGD differently. We think this is a fertile area for future research.

heading: REPRODUCIBILITY STATEMENT
text: Python code to reproduce our results (including all Figures except Figure 1) can be found in the following (anonymized) git repository which contains unit tests and documentation: https://anonymous.4open.science/r/deeplinear-2F10

heading: A USEFUL LEMMAS
text: We will need the following lemma in the proof of Theorem 1,4. Lemma 1. If we solve a separable classification problem with natrual gradient flow with separator β and output s, then the gradient and the Fisher information matrix are the following (in case of linear network this means s = Xβ):
[∇sL(s)]i = −yi(1− φ(yisi)) (9)
[F (s)]i,j = δi,jφ(si)(1− φ(si) (10)
Proof. First note that φ(u) = 11+e−u and φ(−u) = 1− φ(u) = e−u 1+e−u .
[∇sL(s)]i = ∂L ∂si = ∂ ∂si N∑ n=1 log(1 + e−ynsn) = −yie−yisi 1 + e−yisi = −yi(1− φ(yisi)) (11)
Using Equation (11) we get the following:
[F (s)]i,j = [Ey[∇sL(s)∇>s L(s)]]i,j = [Ey[yi(1− φ(yisi))yj(1− φ(yjsj))]]i,j =
=
{ Ey[(1− φ(yisi))2] if i = j
Eyi [yi(1− φ(yisi))]Eyj [yj(1− φ(yjsj))] if i 6= j =
=
{ φ(si)(1− φ(si)) if i = j
Eyi [yi(1− φ(yi))]Eyj [yj(1− φ(yj))] if i 6= j
(12)
Now we get the following:
Eyi [yi(1− φ(yisi))] = φ(si)(1− φ(si))− (1− φ(si))(1− φ(−si)) = 0 (13)
Hence we get: [F (s)]i,j = δi,jφ(si)(1− φ(si)). (14)
The next lemma is essential in all computation connected to matrix completion with matrix factorization. Lemma 2. If we assume that the product matrix β comes from a Gaussian distribution with fixed σnI standard deviation and µ mean, then the Fisher information matrix of the product matrix in matrix factorization is F (β) = 1σ2n I .
Proof. Because of the assumption:
p(X|θ) = N (X|µ, σnI), (15) where θ is the parameters of the model (µ, σn).
∇θ log p(X|θ) = ∇θ log ( 1 σn √ 2π e − (X−µ)
2
2σ2n ) = ∇µ(log( 1
σn √ 2π
)− (X − µ) 2
2σ2n ) = (X − µ) σ2n ,
(16) therefore we can compute the Fisher as F (β) = EX∼p(X|θ) [ ∇θ log p(X|θ)∇>θ log p(X|θ) ] = EX∼p(X|θ) [ (X − µ)(X − µ)>
σ4n
] = σ2nI
σ4n =
1
σ2n I
(17)

heading: B EXACT NATURAL GRADIENTS IN LINEAR MODELS

heading: B.1 SIMPLE LINEAR MODEL LOGISTIC LOSS
text: To obtain the natural gradient ∇̃βL with respect to β, we have to solve the following linear system:
F (β)∇̃βL = ∇βL, (18) where F (β) is the Fisher information matrix and∇βL is the (Euclidean) gradient. Under the logistic loss the Fisher information matrix becomes
F (β) = X> diag[φ(Xβ) φ(−Xβ)]X, (19) where φ is the logistic sigmoid which is applied elementwise to vector arguments and denotes elementwise product. The gradient of the logistic loss is as follows:
∇βL = −(y X)>φ(−(y X)β) (20) Mathematically, we could use these expressions and solve the linear system Equation (18), however, this would be potentially numerically unstable for reasons outlined below. Let’s introduce the notation X̃ = y X and u = X̃β to simplify the formulæ. Due to symmetry, in the Fisher information all occurrences of X can be replaced by X̃ . This gives rise to the following expressions for the Fisher information matrix: F (β) = X̃T diag[φ(u) φ(−u)]X̃ (21) and the gradient:
∇βL = −X̃φ(−u). (22) As the classifier gets better, components of u increase and diverges to +∞. As a consequence both F (β) and∇βL are expected to become small, from the term φ(−u). This could lead to issues with numerical stability. To solve this, we rewrite both using following identity:
φ(−u) = 1 1 + eu
= e−u
1 + e−u = e−uφ(u) (23)
obtaining:
F (β) = e−umaxX̃T diag[e−u+umaxφ2(u)]X̃ (24)
∇βL = −e−umaxX̃e−u+umaxφ(u), (25) where umax is the largest entry of u. We have thus isolated the term responsible for poor numerical performance into a multiplicative term e−umax which we can simply leave out when solving the linear system. The remaining terms are well-behaved even as u increases, provided that the difference between elements of u is not too large.

heading: B.2 DIAGONAL LINEAR NETWORK UNDER LOGISTIC LOSS
text: In a diagonal linear network we express β = w1 · · · wL. Here we will discuss how we compute the natural gradient with respect to wl.
We now solve the following (underdetermined) system of linear equations, which we write using using Einstein summation notation:
F (β)i,jJj,l,k∇̃wl,kL = ∇βiL, (26)
where Jj,l,k = ∂βj ∂wl,k
is the Jacobian of the mapping from w to β. In this specific parametrization, most entries of J is non-zero. Let’s denote the product of the first l− 1 weight vectors as al and the product of the last L− l − 1 weight vectors as bl so we can have:
βi = w1,i · · ·wl−1,i︸ ︷︷ ︸ al,i wl,iwl+1,i · · ·wL,i︸ ︷︷ ︸ bl,i = al,iwl,ibl,i. (27)
Thus, the Jacobian becomes:
Ji,l,k =
{ al,ibl,i if i = j
0 if i 6= j (28)
Substituting this back, we have to solve the following system of equations:
F (β)i,jJj,l,k∇̃wl,kL = ∇βiL (29) F (β)i,jal,jbl,j∇̃wl,jL = ∇βiL. (30)
(31)
To ensure numerical stability, we use the same trick as in B.2.
Since the above system of equations is underdetermined, we could choose different solutions. In our experiments we used the pytorch.linalg.lstsq least squares solver which finds the solution with the lowest `2 norm.

heading: B.3 SEPARABLE CLASSIFICATION
text: First note that in our model (∀n ∈ {1, 2, · · ·N}, φ(s) = 11+e−s )
p(yn = 1|xn,β) = 1
1 + e−ynx > nβ
= φ(−ynx>nβ)
p(yn = −1|xn,β) = 1− 1
1 + e−ynx > nβ
= 1− φ(−ynx>nβ) (32)
So the loss function is (∀n ∈ {1, 2, · · ·N})
`(ynx > nβ) = log(1 + e −ynx>nβ) (33)
and
L(β) = N∑ n=1 log(1 + e−ynx > nβ) (34)
Until now this did not depend on the parametrization. Now look at the parametrizations we used in our article. If we use a fully connected network the gradient is the following:
∇βL(β) = N∑ n=1 ∇β log(1 + e−ynx > nβ) = N∑ n=1 −ynxne−ynx > nβ 1 + e−ynx > nβ =
= N∑ n=1 −ynxn 1 + eynx > nβ = N∑ n=1 −ynxn(1− φ(ynx>nβ)).
(35)
The Fisher information matrix is the following:
F (β) = EX [EY |X [∇β`(−ynx>nβ)∇>β `(−ynx>nβ)]] =
= 1
N N∑ n=1 EY |X [xnx > n (1− φ(ynx>nβ))2] =
= 1
N N∑ n=1 xnx > n (φ(x > nβ)(1− φ(x>nβ))2 + (1− φ(x>nβ))(1− φ(−x>nβ))2) =
= 1
N N∑ n=1 xnx > n (φ(x > nβ)(1− φ(x>nβ))2 + (1− φ(x>nβ))φ2(x>nβ)) =
= 1
N N∑ n=1 xnx > n φ(x > nβ)(1− φ(x>nβ))
(36)
If we use a diagonal network β = w1 w2 · · · wL−1 wL, where w = ( w>1 w > 2 · · · w>L )> . The gradient is the following:
∇βL = J>∇wL (37)
where J (the Jacobian) is the following J = ( ∂β ∂w1 · · · ∂β∂wL ) . (38) where [ ∂β
∂wn ] i,j = ∂βi ∂[wn]j = δi,j N∏ k=1,k 6=i [wk]i (39)
The Fisher information matrix is the following:
F (w) = J>F (β)J (40)

heading: B.4 MATRIX FACTORIZATION
text: Before we compute the natural gradient of matrix factorization let us introduce some notations: β =W1W2 . . .WL, as before and
θ = vec(β), (41)
w = vec(W1,W2, . . . ,WL), (42)
where vec vectorizes the matrices to obtain a column vector. θ is a reparametrization of w, so θ = P(w) and let J = ∂θ∂w . With this notation, let’s compute the natural gradient with respect to the parametrization w.
∇̃wL = F (w)−1∇wL = (J>F (θ))J)−1(J>∇θL) = J−1F (θ)−1∇θL (43)
We use the assumption that J is full rank and because of F (θ) = I is invertible (J>F (θ))J)−1 = J−1F (θ)−1J−>. Thus, the natural gradient simplifies to
∇̃wL = J−1∇θL (44)
and multiplying by J we obtain J∇̃wL = ∇θL. (45)
We can consider the Jacobian like L consecutive matrices
J = [J1J2 . . . JL] (46)
where Ji = ∂θ∂vec(Wi) , and note that∇θL = vec(∇βL) and ∇̃wL = vec(∇̃W1,W2,...WLL). Rewrite equation 45:
Jvec(∇̃W1,W2,...WLL) = vec(∇βL). (47)
If we solve the following equation for i = 1, . . . L, then the concatenation of vectors vec(∇̃WiL) will solve equation 47 as well.
Jivec(∇̃WiL) = 1
L vec(∇βL) (48)
Let Ai = W1W2 . . .Wi−1 and Bi = Wi+1Wi+2 . . .WL and using ⊗ notation for the Kronecker product and utilize the property vec(ABC) = (C> ⊗A)vec(X) we get
Ji = ∂vec(β) ∂vec(Wi) = ∂vec(AiWiBi) ∂vec(Wi) = ∂(B>i ⊗Ai)vec(Wi) ∂vec(Wi) = B>i ⊗Ai, (49)
thus we need to solve (B>i ⊗Ai)vec(∇̃WiL) = 1
L vec(∇βL) (50)
for vec(∇̃WiL). One can do this by exploiting properties of the Kronecker product and using MoorePenrose pseudo-inverses as follows:
vec(∇̃WiL) = 1
L (B>i + ⊗A+i )vec(∇βL) = 1 L (B>i +∇βLA+i ) (51)
We note that when Ai and Bi are near full-rank, using the pseudoinverses may not be numerically stable. Fausett & Fulton (1994) instead proposed a solution based on QR decomposition, and even discussed an approach which extends to the rank deficient case. In practice we found that this was not necessary for ours experiments. As a result, in our implementation we use the formula 1 LB > i +∇βLA+i to update the factor matrices with the natural gradient.

heading: C PROOF OF THEOREMS

heading: C.1 PROOF OF THEOREM 1
text: Statement. Let’s assume, that N < D, X is full rank and A is an invertible D × D matrix. Let βt = βt(X,y) be the trajectory of NGF and β ′ t = βt(XA
>,y) (the trajectory of NGF on data XA>). Then Xβ = XATβ′.
Proof. Let s = Xβ and s′ = XATβ′. The gradient and the Fisher information matrix are the following (the calculation can be found in Lemma 1).
[∇sL(s)]i = −yi(1− φ(yisi)) [F (s)]i,j = δi,jφ(si)(1− φ(si))
(52)
The exact same can be said about s′, so s and s′ are the solutions of the same differential equations, so if we use the same initialization st = s′t.

heading: C.2 PROOF OF THEOREM 2
text: Statement. Let βt(X,y) be the trajectory of NGF and let A be a D × D invertible transformation. If N ≥ D, X has full rank and we consider NGF on the transformed data XA>, then A>βt(XA >,y) = βt(X,y).
Proof. Let β′ be the trajectory of NGF on the transformed data:
β′t = βt(XA >,y). (53)
To run NGF on β′ we need its Fisher information matrix. Note that the Fisher information matrix of linear models with logistic-loss is
F (β) = X>diag[φ(Xβ) φ(−Xβ)]X. (54) by Appendix B.1. Note that in this case the rank of the Fisher information matrix is D, so it is invertible. Same is true for β′. Let’s compute the Fisher information matrix of β′.
F (β′) = 1
N N∑ n=1 Eyn|Axn [∇β′`(ynβ ′>Axn)∇>β′`(ynβ ′>Axn)] (55)
First, specify∇β′`(ynβ′>Axn) and use the notation v> = β′>A.
∇β′`(ynβ′>Axn) = J>∇v>`(ynv>xn) (56)
where J = ∂v >
∂β′ .
Ji,j = ∂vi ∂β′j = ∂ ∑d k=1 β ′ kAk,i ∂β′j = Aj,i (57)
Therefore J = A> ⇔ J> = A and
∇β′`(ynβ′>Axn) = A∇v`(ynv>xn). (58)
We now can continue the computation of the Fisher:
F (β′) = 1
N N∑ n=1 Eyn|Axn [A∇v>`(ynv >xn)∇>v>`(ynv >xn)A >] =
= A( 1
N N∑ n=1 Eyn|xn [∇v>`(ynv >xn)∇>v>`(ynv >xn)])A > = AF (v)A>.
(59)
Note, that the Fisher of v must be invertible as well from the previous Equation. Let’s see the NGF on β′:
β̇ ′ = −F (β′)−1∇β′L(β′>XAT , y) = −(AF (v)A>)−1 N∑ n=1 ∇β′`(ynβ′>Axn) =
= −(A>)−1F (v)−1A−1A N∑ n=1 ∇v`(ynv>xn) = −(A>)−1F (v)−1∇vL(v)
(60)
We also have the following (by the Chain Rule):
v̇ = Jβ̇ ′ = A>β̇ ′ (61)
Now from Equation (60) and (61) we get:
v̇ = F (v)−1∇vL(v) (62)
This is the same differential equation as the one β is a solution of. So if they are initialized the same way vt = βt(X,y), so βt = vt = A >β′.

heading: C.3 PROOF OF THEOREM 4
text: Statement. If N < D and β is the separator. Let s be the output. Let the Jacobi matrix from β to s be J = ∂s∂β . If J is full rank s is asymptotically linear with direction vector y.
Proof. First let’s note that by the invariance property of NGF the trajectory of s is defined by the trajectory of β. ṡ = −F−1(s)∇sL(s) (63) Let’s assume s is 1-dimensional. In this case s = s, x1 = x and y = y can be used since we have only one data point. To solve equation (63) we need the gradient and the Fisher information matrix which are the following (the calculation can be found in the Appendix B.2)
∇sL(s) = −y(1− φ(ys)) (64)
The Fisher information matrix:
F (s) = φ(s)(1− φ(s)) (65)
Then Equation (63) can be written as:
ṡ = y(1− φ(ys)) φ(s)(1− φ(s))
(66)
Now we rescale our data points s.t. x̃ = −x, so s̃ = −s and ỹ = −y = 1. Hence we get: ∂s̃
∂t =
1
φ(s̃) (67)
Which can be solved and the solution is
log(1 + es̃) = t+ c ⇐⇒ s̃ = log(et+c − 1) (68) By equation (68) we get the asymptotic behaviour
lim t→∞
s̃
t+ c = lim t→∞ log(et+c − 1) t+ c = (69)
= lim t→∞
et+c
et+c − 1 = lim t→∞
1
1− e−(t+c) = 1 (70)
(From (69) to (70) we use L’Hopital Rule). Hence we proved Theorem 4. for N = 1. Now let’s assume, that N > 1. Now we write down the gradient again: [∇sL(s)]i = −yi(1− φ(yisi)) (71) And the Fisher information matrix:
[F (s)]i,j = δi,jφ(si)(1− φ(si)) (72) So now if we substitute in Equation (71) and Equation (72) to Equation (63). We can rescale, so ỹi = 1 ∀i as we did in the previous case. Hence we get the following:
∂s̃ ∂t = −  1 φ(s̃1)(1−φ(s̃1)) 0 · · · 0 0 1φ(s̃2)(1−φ(s̃2)) · · · 0 ... ... . . .
... 0 0 · · · 1φ(s̃N )(1−φ(s̃N ))
  −(1− φ(s̃1)) −(1− φ(s̃2))
... −(1− φ(s̃N ))
 =  1 φ(s̃1) 1 φ(s̃2)
... 1
φ(s̃N )  (73)
Hence we got N independent differential equations which are exactly the same as in the N = 1 case. So in each dimension s̃ is asymptotically t+ c for some c. Hence s̃ ≈ t1+ c, where c ∈ RD is a constant. So s ≈ ty + cs, where cs ∈ RD is a constant.

heading: D COUNTEREXAMPLE FOR THE INVARIANCE OF `2 LARGE MARGIN SOLUTION
text: The counterexample is the following:
A = ( 1 2 −1 0 ) , y = 1 and X = (2 −3)
Then β∗(X, y) = argmin‖β‖2 subject to 2β1 − 3β2 ≥ 1, therefore β ∗(X, y) = ( 0 − 13 ) . Further-
more β∗(XA>, y) = argmin‖β‖2 subject to −4β1− 2β2 ≥ 1, therefore β ∗(XA>, y) = ( − 14 0 ) ,
but A>β∗(XA>, y) = ( − 14 − 12 ) 6= ( 0 − 13 ) = β∗(X, y).

heading: E PROOF OF THE STATEMENT ABOUT THE PARAMETRIZATION INVARIANCE
text: OF NGF
Statement. Let w and θ be two parameter vectors related by the mapping θ = P(w) and consider natural gradient flow in w. Assume that (1) the Jacobian J = ∂θt∂wt and (2) F (θt) are both full rank for all t. If wt follows natural gradient flow starting from w0 then θt = P(wt) follows NGF, i. e. it solves θ̇t = −F (θt)+∇θtL(X, θt).
Proof. We use that F (w) = J>F (θ)J which follows from the definition of F :
F (w) = EX [∇wL(X,w)∇>wL(X,w)] = EX [J>∇θL(X, θ)∇>θ L(X, θ)J ] = J>F (θ)J
The invariance statement follows:
θ̇ = ˙(P(wt)) = Jẇt = −JF (wt)+∇wtL(X,wt) = = −JJ+F (θt)+(J>)+J>∇θtL(X, θt) = −F (θt)+∇θtL(X, θt).

heading: F PROOF OF THE STATEMENT ABOUT NGD IN MATRIX COMPLETION
text: Statement. Let’s apply NGF for the problem of matrix completion. EGF in the direct parametrization (β = w) is equivalent to NGF under any parametrization θ for which J = ∂wt∂θt is full rank.
Proof. First let’s consider a parametrization θ s.t. the direct parametrization β = w (= P(θ)) and J = ∂w∂θ is full rank. Then by the invariance property if θt is the solution of the NGF with the arbitrary parametrization, then wt = P(θt) is the solution of:
ẇ = −∇wL(w)
Which agrees with the EGF with direct parametrization.
G INVARIANCE PROPERTY OF OLS
We show the same transformation invariance property for OLS that we showed in Theorem 1,2 for NGF. Again, we split the problem into two cases: N < D and N ≥ D. Note that for the problem Xβ = y the Ordinary least squares solution is β = (X>X)−1X>y if the columns of X is linearly independent. Statement. Let N < D, A is an invertible D ×D matrix. If β is the solution of the Ordenary least squares problem for the matrix X and β′ for XA>, then Xβ = XA>β′.
Proof. Immediately follows from the definition of the problems: Xβ = y and XA>β′ = y.
Statement. Let N ≥ D, X has full rank and A is an invertible D×D matrix. If β is the solution of the OLS problem for the matrix X and β′ for XA>, then A>β′ = β.

heading: Proof.
text: A>β′ = A>((XA>)>XA>)−1(XA>)>y = A>A>−1(X>X)−1A−1AX>y = β
H SUPPLEMENTARY FIGURES

Review summaries:
The paper studies inductive bias of natural gradient descent (NGD) in deep linear networks. Utilizing _reparameterization invariance_ properties of NGD, the authors attempt to eliminate the role of parameterization and isolate the effect of solution found by gradient descent.  The paper studies gradient flow dynamics for 1) separable classification under logistic loss and 2) deep matrix factorization. One of the contributions stated is identifying learning problems where NGD fails to generalize while GD with the right architecture performs well. 

This paper studies the inductive bias of natural gradient flow in deep linear networks and in matrix factorization. They show that the solution to the empirical risk minimization problem using NGF is invariant to pixel permutation, and conclude that the desirable property of Euclidean gradient descent under logistic loss of finding the minimum l2 norm solution does not hold (theorem 1&2).  They also contribute an efficient NGD algorithm for deep diagonal linear networks.

This paper studies the role of parameterization in determining inductive biases and influencing generalization in deep learning by studying natural gradient descent.  As the authors review, natural gradient descent is invariant to reparameterization.  This implies that the trajectory of NGD is in some sense ablating the role of parameterization allowing the authors to (1) determine the importance of parameterization with gradient descent in generalization, and (2) study the inductive biases unique to NGD.  The authors consider two settings both theoretically and empirically: separable classification with deep linear models and matrix completion via deep matrix factorization.