D EPTH W ITHOUT THE M AGIC : I NDUCTIVE B IAS OF N ATURAL G RADIENT D ESCENT
Anonymous authors
Paper under double-blind review
A BSTRACT
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 classiﬁers 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 ﬁnds 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 ﬂow in deep linear networks for separable classiﬁcation under logistic loss and deep matrix factorization.
Some of our ﬁndings extend to nonlinear neural networks with sufﬁcient but ﬁnite 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.
I NTRODUCTION
Consider two parameter spaces ( W
W 2 ) and two optimization trajectories in each: one EGD, one NGD.
If we map these into the hypothesis space ( H ) then EGD ﬁnds different optima but NGD ﬁnds the same
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 efﬁciently ﬁnd low-rank solutions in matrix completion problems (Arora et al., 2019), sparse solutions in separable classiﬁcation (Gunasekar et al., 2018) or compressed sensing (Vaˇskeviˇcius et al., 2019).
Valle-P´erez et al.
(2018) studied deep neural networks and found evidence that the parameter-hypothesis mapping 1 is biased towards simpler functions as measured by Kolmogorov complexity.
Taken together, these observations and ﬁndings have lead the community to hypothesize that
The parameter-to-hypothesis mapping inﬂuences 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 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 ﬁrst, practical, implication is as follows: if the parameter-to-hypothesis mapping really does play an important role in generalization, then eliminating its inﬂuence 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) efﬁcient 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 classiﬁcation, we show that NGF is invariant under invertible transformations of data (Theorems 1&2) and as a consequence it cannot recover the (cid:96) p large margin solutions that EGD tends to converge to.
• We further show that (in case of separable classiﬁcation) 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 efﬁcient 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.
B ACKGROUND 2. S EPARABLE C LASSIFICATION WITH D EEP L INEAR M ODELS
In this article we consider binary classiﬁcation datasets { ( x n , y n ) , n = 1 , . .
. , N } separable by a homogeneous linear classiﬁer with a positive margin ( i. e. ∃ β ∗ s.t.
y n x (cid:62) n β ∗ ≥ 1 ∀ n ).
(We use the notation X
=
( x 1 · · · x N ) (cid:62) ).
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 (cid:96) 2 large margin classiﬁer in direction, that is
[Equation]
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 β = w 1 (cid:12) w 2 (cid:12) . . .
(cid:12) w L , using (cid:12) to denote elementwise product.
When we adjust parameters w 1 , . .
. , w L through Euclidean gradient descent, β ( t ) converges to the (cid:96)
L large margin separator deﬁned as
[Equation]
A remarkable consequence of this is that unregularized gradient descent can ﬁnd sparse classiﬁers, without any form of explicit regularization.
In fact, this inductive bias is even more sparsity-seeking than the typically used (cid:96) 1 regularization (see e. g. Koh et al., 2007; Tibshirani, 1996).
Figureillustrates this behaviour in a 2D example.
A: direct parametrization, EGD
B: L =4 diagonal network, EGD
C: direct parametrization, NGD
D: L =4 diagonal network, NGD
2. M ATRIX C OMPLETION VIA D EEP M ATRIX F ACTORIZATION
[Figure]
When L ≥ 2 , Euclidean gradient methods also move entries where there is no observation - enabling implicit regularization towards low-rank solutions.
By contrast, and due to invariance, natural gradient methods move only the single entry to match the observation.
The task of matrix completion involves recovering an unknown matrix β ∗ ∈ R D × D from a randomly chosen subset of observed entries 2 .
The problem is clearly underdeﬁned: there are inﬁnitely 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 identiﬁable.
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 ∈ R D × R and V ∈ R R
× 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 β = W 1 W 2 where
W 1 and W 2 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 β = W 1 W 2 · · · W L .
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 preferred 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.
2. N ATURAL GRADIENT DESCENT
In the next section we brieﬂy 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 deﬁne (equipped with a different metric).
More speciﬁcally, let’s say that the parameter of interest is θ , where θ deﬁnes 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 ( θ )
=
(cid:80) Nn =1 l ( θ, x n , y n ) .
Then, NGD is usually deﬁned as
[Equation]
where
[Equation]
is the average Fisher information matrix and η is the step size.
In the above deﬁnition, E Y | X ; θ is taken over the distribution speciﬁed by θ , but distribution with respect to which the expectation E X 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 ﬂow (NGF) the continuous limit of NGD, analogously deﬁned as
[Equation]
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 deﬁne NGF as any trajectory θ t which satisﬁes
[Equation]
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 :
[Equation]
We have seen how in EGD, different parametrization of the same problem leads to drastically different trajectories and optima.
However, NGD with inﬁnitesimally small learning rate (i. e. NGF) always follows the same trajectory in model-space and this ﬁnds 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 ﬂow in w .
Assume that (1) the
Jacobian J =
∂θ t ∂ w t and (2) F ( θ t ) are both full rank for all t .
If w t follows natural gradient ﬂow starting from w 0 then θ t = P ( w t ) follows NGF, i. e. it solves ˙
θ t = − F ( θ t )
+ ∇ θ t L ( X, θ t ) .
N ATURAL GRADIENTS UNDER L OGISTIC L OSS ON S EPARABLE D ATA
We have seen in Section 2.1 that when trained on separable data with the logistic loss EGD tends to converge to large margin classiﬁers.
To illustrate how NGD differs, we ﬁrst 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 β (cid:48)
t = β t ( XA (cid:62) , y ) (the trajectory of NGF on data
XA (cid:62) ).
Then X β = XA (cid:62) β (cid:48) (with the assumption that β and β (cid:48) have equivalent initial conditions).
Proof sketch.
We use the notation s = X β
and s (cid:48)
=
XA (cid:62) β and prove that s t = s (cid:48) 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 (cid:62) , then
A (cid:62) β t ( XA (cid:62) , y ) =
β t ( X, y ) (with the assumption that β and β (cid:48) 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 β (cid:48)
t = β t ( XA (cid:62) , y ) and v (cid:62) =
β (cid:48)(cid:62) A .
Then we prove the following:
[Equation]
Hence we get:
[Equation]
and
[Equation]
So if v and β (cid:48) have the same initialization, then v t = β (cid:48) t .
Full proof can be found in Appendix C.2.
Conclusion.
Let u t ( X, y ) denote the trajectory of X β t , which is the linear function β (cid:62) t x evaluated at each of the datapoints x n .
Then u t ( XA (cid:62) , y ) =
u t ( X, y ) .
Proof.
[Equation]
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 (cid:96) p large-margin behaviour we have seen in EGD.
Remark.
Let A be a D × D invertible transformation and let β ∗ ( X, y ) be the (cid:96) 2 large margin solution, i.e. β ∗ ( X, y )
= argmin (cid:107) β (cid:107) 2 subject to y n β (cid:62) x n ≥ 1 ∀ n .
Then the (cid:96) 2 large margin classiﬁer does not have the invariance property, namely there exists a dataset ( X, y ) and a transformation
A such that A (cid:62) β ∗ t ( XA (cid:62) , y ) (cid:54) =
β ∗ 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 (cid:62) X ) − 1 X (cid:62) 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 ﬂow under logistic loss, the logits s t =
X β t follow an asymptotically linear trajectory with direction vector y .
Remark.
The Fisher information matrix w.r.t.
β is F ( β ) =
X (cid:62) 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 − 1 s and +1 s) 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 ﬁnd qualitatively different solutions in classiﬁcation 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 log N .
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 ﬁnd qualitatively different classiﬁers, 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 ∈ R P , P ≥ N be the parameters of a classiﬁer with logits s = s
( X ; w ) ∈ R N .
If w t follows natural gradient ﬂow on the logistic loss with labels y and the Jacobian J t = ∂ s t ∂ w t is of full rank, then s t 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 β .
[Equation]
[Figure]
u
[Figure]
a c y
[Figure]
there are there are fewer diagonal can reach
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.
3. E XPERIMENTS
In order to validate and illustrate our ﬁndings 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 β = w 1 (cid:12) · · · (cid:12) w L (Gunasekar et al., 2018) for different depth L .
In order to run these experiments we needed to implement an efﬁcient algorithm for computing natural gradients in these models: naively calculating and then inverting the Fisher information matrix is computationally inefﬁcient 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 classiﬁcation 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 ﬁndings of Gunasekar et al.
we expected EGD to ﬁnd the large margin solution when L is low, and the axis-aligned solution when L is sufﬁciently large.
The results in panels a and b of Figure 2 conﬁrm 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 ﬁrst 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 .
Conﬁrming 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 Figurevalidates the predictions of Theorem 4: logits from the model converge to t y .
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 ﬁnds.
[Figure]
M ATRIX C OMPLETION WITH N ATURAL GRADIENT DESCENT
As we have seen in Section 2.
[Figure]
EGD in the deep matrix product parametrization β = W 1 · · · W L converges to low-rank solutions.
However, when L = 1 , i
[Figure]
e. when we run EGD directly on β , the solution we ﬁnd 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 σ 2 n I , where σ 2 n is the observation noise.
The observation noise σ 2 n is assumed a constant, and is inconsequential here as it cancels with the 1 σ 2 n 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 = ∂ w t
∂θ 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 k i , 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 k i,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 ,by moving entries in the ﬁfth row of W 1 or in the ﬁfth column of W 2 .
This, in turn, might result in moving β i, 5 or β 5 ,i as well.
This explains the cross pattern seen in Figure 3 ﬁrst panel in the second row.
This non-identity NTK is what allows generalization to happen as ’information ﬂows’ 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 ( β = W 1 · · · W L ) we calculate the natural gradient w.r.t.
W l as ˜ ∇ W l L = 1 L B (cid:62) l + ˜ ∇ β L A + l , where A l = (cid:81) l − 1 i =1 W i and B l = (cid:81)
Li
=
l +1 W i .
We note that A i and
B i 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 ﬁnds a trivial optimum in deep matrix product parametrizations of varying depth.
We follow the experimental setup of Arora et al.
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.
S UMMARY AND D ISCUSSION
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.
5. O THER RELATED WORK
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 ﬁrst place.
In order to perform experiments that validate our ﬁndings we develop efﬁcient exact natural gradient descent algorithms in overparametrized linear models extending the work of Bernacchia et al.
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 ﬁrst 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 ﬁndings 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 misspeciﬁcation or when the labels are noisy, NGD may have an advantage.
Finally, Wadia et al.
(2021) argued that second order information of the input datawhich 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.
5. Q&A
Q: How about stochastic gradients?
Following Gunasekar et al.
(2017; 2018); Arora et al.
we analysed only full-batch gradient descent.
This allowed us to prove properties of gradient ﬂow, i. e. in the limit of inﬁnitesimally 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 efﬁcient 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.
; 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 inﬂuence 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.
ValleP ´ erez 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 W 1 , . . .
, W L 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 inﬂuence 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 W 1 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 inﬁnitely wide networks, too, where initialization scale determines whether the network behaves like a linear kernel machine, or more like the behaviour we describe in ﬁnite 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 Figurewe 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.
R EPRODUCIBILITY S TATEMENT
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-2F
A U SEFUL LEMMAS
We will need the following lemma in the proof of Theorem 1,
[Equation]
Lemma 1.
If we solve a separable classiﬁcation problem with natrual gradient ﬂow 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 β ):
[Equation]
Proof.
First note that φ ( u ) =
[Equation]
Using Equation (11) we get the following:
[Equation]
Now we get the following:
[Equation]
Hence we get:
[Equation] [Equation]
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 ﬁxed σ n I standard deviation and µ mean, then the Fisher information matrix of the product matrix in matrix factorization is F ( β ) = 1 σ 2 n I .
Proof.
Because of the assumption:
[Equation]
) where θ is the parameters of the model ( µ , σ n ).
[Equation]
therefore we can compute the Fisher as
[Equation]
B E XACT NATURAL GRADIENTS IN LINEAR MODELS B. S IMPLE LINEAR M ODEL LOGISTIC LOSS
To obtain the natural gradient ˜ ∇ β L with respect to β , we have to solve the following linear system:
[Equation]
(18) where F ( β ) is the Fisher information matrix and ∇ β L is the (Euclidean) gradient.
Under the logistic loss the Fisher information matrix becomes
[Equation]
(19) where φ is the logistic sigmoid which is applied elementwise to vector arguments and (cid:12) denotes elementwise product.
The gradient of the logistic loss is as follows:
[Equation]
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 (cid:12) 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:
[Equation]
and the gradient:
[Equation]
As the classiﬁer 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:
[Equation]
obtaining:
[Equation]
where u max is the largest entry of u .
We have thus isolated the term responsible for poor numerical performance into a multiplicative term e − u max 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.
B. D IAGONAL L INEAR N ETWORK UNDER L OGISTIC L OSS
In a diagonal linear network we express β = w 1 (cid:12)· · ·(cid:12) w L .
Here we will discuss how we compute the natural gradient with respect to w l .
We now solve the following (underdetermined) system of linear equations, which we write using using Einstein summation notation:
[Equation]
where J j,l,k = ∂ β j ∂ w l,k is the Jacobian of the mapping from w to β .
In this speciﬁc parametrization, most entries of J is non-zero.
Let’s denote the product of the ﬁrst l − 1 weight vectors as a l and the product of the last L − l − 1 weight vectors as b l so we can have:
[Equation]
Thus, the Jacobian becomes:
[Equation]
Substituting this back, we have to solve the following system of equations:
[Equation] [Equation]
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
B. S EPARABLE CLASSIFICATION
First note that in our model ( ∀ n ∈ { 1 , 2 , ·
[Equation]
So the loss function is ( ∀ n ∈ { 1 , 2 , · · · N } )
[Equation] [Equation]
and
[Equation]
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:
[Equation] [Equation]
The Fisher information matrix is the following:
[Equation]
If we use a diagonal network β = w 1 (cid:12) w 2 (cid:12) · · · (cid:12) w L − 1 (cid:12) w L , where w = (cid:0) w (cid:62)
[Equation]
L
[Equation]
(cid:62) .
The gradient is the following:
[Equation] [Equation]
where J (the Jacobian) is the following
[Equation]
where
[Equation]
The Fisher information matrix is the following:
[Equation]
M ATRIX FACTORIZATION
Before we compute the natural gradient of matrix factorization let us introduce some notations: β = W 1 W 2 . . .
W L , as before and
[Equation]
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 .
[Equation] [Equation]
We use the assumption that J is full rank and because of F ( θ ) =
I is invertible ( J (cid:62) F ( θ ))
J ) − 1 =
J −
[Equation]
F ( θ )
[Equation]
J
[Equation]
Thus, the natural gradient simpliﬁes to
[Equation]
44) and multiplying by J we obtain
[Equation]
We can consider the Jacobian like L consecutive matrices
[Equation]
46) where J i =
∂θ ∂vec ( W i ) , and note that ∇ θ L = vec ( ∇ β L ) and ˜ ∇ w L = vec ( ˜ ∇ W 1 ,W 2 ,...
W L L ) .
Rewrite equation 45:
[Equation]
If we solve the following equation for i = 1 , . . .
L , then the concatenation of vectors vec ( ˜ ∇ W i L ) will solve equation 47 as well.
[Equation]
Let A i
=
W 1 W 2 . . .
W i − 1 and B i
=
W i +
W i
+2 . .
W L and using ⊗ notation for the Kronecker product and utilize the property vec ( ABC )
=
( C (cid:62) ⊗ A ) vec ( X ) we get
[Equation]
thus we need to solve
[Equation]
for vec ( ˜ ∇ W i L ) .
One can do this by exploiting properties of the Kronecker product and using MoorePenrose pseudo-inverses as follows:
[Equation]
We note that when A i and B i 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 deﬁcient case.
In practice we found that this was not necessary for ours experiments.
As a result, in our implementation we use the formula
L B (cid:62) i + ∇ β L A + i to update the factor matrices with the natural gradient.
C P ROOF OF THEOREMS C. P ROOF OF T HEOREM
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 β (cid:48)
t = β t ( XA (cid:62) , y ) (the trajectory of NGF on data
XA (cid:62) ).
Then X β = XA T β (cid:48) .
Proof.
Let s = X β
and s (cid:48)
=
XA T β (cid:48) .
The gradient and the Fisher information matrix are the following (the calculation can be found in Lemma 1).
[Equation]
The exact same can be said about s (cid:48) , so s and s (cid:48) are the solutions of the same differential equations, so if we use the same initialization s t = s (cid:48) t .
C. P ROOF OF T HEOREM
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 (cid:62) , then
A (cid:62) β t ( XA (cid:62) , y ) =
β t ( X, y ) .
Proof.
Let β (cid:48) be the trajectory of NGF on the transformed data:
[Equation]
To run NGF on β (cid:48)
we need its Fisher information matrix.
Note that the Fisher information matrix of linear models with logistic-loss is
[Equation] [Equation]
by Appendix B.1.
Note that in this case the rank of the Fisher information matrix is D , so
[Equation] [Equation]
First, specify ∇ β (cid:48) (cid:96) ( y n β (cid:48)(cid:62)
A x n ) and use the notation v (cid:62) =
β (cid:48)(cid:62)
[Equation]
[Equation]
where J = ∂v (cid:62) ∂ β (cid:48) .
[Equation]
Therefore J = A (cid:62) ⇔ J (cid:62) =
A and
[Equation] [Equation]
We now can continue the computation of the Fisher:
[Equation] [Equation]
Note, that the Fisher of v must be invertible as well from the previous Equation.
Let’s see the NGF on β (cid:48) :
[Equation]
We also have the following (by the Chain Rule):
[Equation]
Now from Equation (60) and (61) we get:
[Equation]
This is the same differential equation as the one β is a solution of.
So if they are initialized the same way v t = β t
( X, y ) , so β t = v t =
A (cid:62) β (cid:48) .
C. P ROOF OF T HEOREM
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 deﬁned by the trajectory of β .
[Equation]
Let’s assume s is 1-dimensional.
In this case s = s , x 1 = 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)
[Equation]
The Fisher information matrix:
[Equation]
Then Equation (63) can be written as:
[Equation]
Now we rescale our data points s.t.
˜
x
= − x , so ˜ s = − s and ˜ y =
− y
= 1 .
Hence we get:
[Equation]
Which can be solved and the solution is
[Equation]
By equation (68) we get the asymptotic behaviour
[Equation]
(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:
[Equation]
And the Fisher information matrix:
[Equation]
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:
[Equation]
− (1 − φ (˜ s N ))
 
=   1 φ (˜ s 1 )φ (˜ s 2 ) .
1 φ (˜ s N )   (73)
Hence we got N independent differential equations which are exactly the same as in the N =case.
So in each dimension ˜ s is asymptotically t
+ c for some c .
Hence ˜ s ≈ t 1 + c , where c ∈ R D is a constant.
So s ≈
t y + c s , where c s ∈ R D is a constant.
D C OUNTEREXAMPLE FOR THE INVARIANCE OF (cid:96) 2 LARGE MARGIN SOLUTION
The counterexample is the following:
A = (cid:18) − 1(cid:19) , y = 1 and X = ( 2 − 3 )
Then β ∗ ( X, y ) =
argmin (cid:107) β (cid:107) 2 subject to 2 β 1 − 3 β 2 ≥ 1 , therefore β ∗ ( X, y )
= (cid:18)−(cid:19) .
Furthermore β ∗ ( XA (cid:62) , y ) =
argmin (cid:107) β (cid:107) 2 subject to − 4 β 1 − 2 β 2 ≥ 1 , therefore β ∗ ( XA (cid:62) , y )
= (cid:18) − 14 0 (cid:19) , but A (cid:62) β ∗ ( XA (cid:62) , y )
= (cid:18) − 14 − 12 (cid:19) (cid:54) =
(cid:18)
− 13 (cid:19) =
β ∗ ( X, y ) .
E P ROOF OF THE STATEMENT ABOUT THE PARAMETRIZATION INVARIANCE OF NGF
Statement.
Let w and θ be two parameter vectors related by the mapping θ = P ( w ) and consider natural gradient ﬂow in w .
Assume that (1) the Jacobian J =
∂θ t ∂ w t and (2) F ( θ t ) are both full rank for all t .
If w t follows natural gradient ﬂow starting from
Proof.
We use that F ( w ) =
J (cid:62) F ( θ ) J which follows from the deﬁnition of F :
[Equation]
The invariance statement follows:
[Equation]
F P ROOF OF THE STATEMENT ABOUT NGD IN MATRIX COMPLETION
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 = ∂ w t
∂θ 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 w t = P ( θ t ) is the solution of:
[Equation]
Which agrees with the EGF with direct parametrization.
G I NVARIANCE 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 (cid:62) X ) − 1 X (cid:62) 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 β (cid:48) for XA (cid:62) , then X β = XA (cid:62) β (cid:48) .
Proof.
Immediately follows from the deﬁnition of the problems: X β = y and XA (cid:62) β (cid:48)
=
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 β (cid:48) for XA (cid:62) , then A (cid:62) β (cid:48)
=
β .
Proof.
[Equation]
H S UPPLEMENTARY F IGURES