Year: 2020

Abstract:
Recently mean field theory has been successfully used to analyze properties of wide, random neural networks. It gave rise to a prescriptive theory for initializing feed-forward neural networks with orthogonal weights, which ensures that both the forward propagated activations and the backpropagated gradients are near `2 isometries and as a consequence training is orders of magnitude faster. Despite strong empirical performance, the mechanisms by which critical initializations confer an advantage in the optimization of deep neural networks are poorly understood. Here we show a novel connection between the maximum curvature of the optimization landscape (gradient smoothness) as measured by the Fisher information matrix (FIM) and the spectral radius of the input-output Jacobian, which partially explains why more isometric networks can train much faster. Furthermore, given that orthogonal weights are necessary to ensure that gradient norms are approximately preserved at initialization, we experimentally investigate the benefits of maintaining orthogonality throughout training, and we conclude that manifold optimization of weights performs well regardless of the smoothness of the gradients. Moreover, we observe a surprising yet robust behavior of highly isometric initializations — even though such networks have a lower FIM condition number at initialization, and therefore by analogy to convex functions should be easier to optimize, experimentally they prove to be much harder to train with stochastic gradient descent. We conjecture the FIM condition number plays a non-trivial role in the optimization.

heading: 1 INTRODUCTION
text: Deep neural networks (DNN) have shown tremendous success in computer vision problems, speech recognition, amortized probabilistic inference, and the modelling of neural data. Despite their performance, DNNs face obstacles in their practical application, which stem from both the excessive computational cost of running gradient descent for a large number of epochs, as well as the inherent brittleness of gradient descent applied to very deep models. A number of heuristic approaches such as batch normalization, weight normalization and residual connections (He et al., 2016; Ioffe & Szegedy, 2015; Salimans & Kingma, 2016) have emerged in an attempt to address these trainability issues. Recently mean field theory has been successful in developing a more principled analysis of gradients of neural networks, and has become the basis for a new random initialization principle. The mean field approach postulates that in the limit of infinitely wide random weight matrices, the distribution of pre-activations converges weakly to a Gaussian. Using this approach, a series of works proposed to initialize the networks in such a way that for each layer the input-output Jacobian has mean singular values of 1 (Schoenholz et al., 2017). This requirement was further strengthened to suggest that the spectrum of singular values of the input-output Jacobian should concentrate on 1, and it was shown that this can only be achieved with random orthogonal weight matrices.
Under these conditions the backpropagated gradients are bounded in `2 norm (Pennington et al., 2017) irrespective of depth, i.e., they neither vanish nor explode. It was shown experimentally in (Pennington et al., 2017; Xiao et al., 2018; Chen et al., 2018) that networks with these critical initial conditions train orders of magnitude faster than networks with arbitrary initializations. The empirical success invites questions from an optimization perspective on how the spectrum of the hidden layer
input-output Jacobian relates to notions of curvature of the parameters space, and consequentially to convergence rate. The largest effective (initial) step size η0 for stochastic gradient descent is inversely proportional to the local gradient smoothness M (Bottou et al., 2018; Boyd & Vandenberghe, 2004). Intuitively, the gradient step can be at most as large as the fastest change in the parameter landscape. Recent attempts have been made to analyze the mean field geometry of the optimization using the Fisher information matrix (FIM) (Amari et al., 2019; Karakida et al., 2019). The theoretical and practical appeal of measuring curvature with the FIM is due to among other reasons the fact that the FIM is necessarily positive (semi-)definite even for non-convex objectives, and due to it its intimate relationship with the Hessian matrix. Karakida et al. (2019) derived an upper bound on the maximum eigenvalue, however this bound is not satisfactory since it is agnostic of the entire spectrum of singular values and therefore cannot differentiate between Gaussian and orthogonal weight initalizations.
In this paper, we develop a new bound on the parameter curvature M given the maximum eigenvalue of the Fisher information λmax(G) which holds for random neural networks with both Gaussian and orthogonal weights. We derive this quantity to inspect the relation between the singular value distribution of the input-output Jacobian and locally maximal curvature of the parameter space. We use this result to probe different orthogonal, nearly isometric initializations, and observe that, broadly speaking, networks with a smaller initial curvature train faster and generalize better, as expected. However, consistent with a previous report (Pennington et al., 2018), we also observe highly isometric networks perform worse despite having a slowly varying loss landscape ( i.e. small initial λmax(G)). We conjecture that the long term optimization behavior is depends on trivially on the smallest eigenvalue m and therefore surprisingly there is a sweetspot with the condition number being mM > 1.
We then investigate whether constraining the spectrum of the Jacobian matrix of each layer affects optimization rate. We do so by training networks using Riemannian optimization to constrain their weights to be orthogonal, or nearly orthogonal and we find that manifold constrained networks are insensitive to the maximal curvature at the beginning of training unlike the unconstrained gradient descent (hereafter “Euclidean”). In particular, we observe that the advantage conferred by optimizing over manifolds cannot be explained by the improvement of the gradient smoothness as measured by λmax(G). Finally, we observe that contrary to Bansal et al. (2018)’s results Euclidean networks with a carefully designed initialization reduce the test misclassification error at approximately the same rate as their manifold constrained counterparts, and overall attain a higher accuracy.

heading: 2 BACKGROUND

heading: 2.1 FORMAL DESCRIPTION OF THE NETWORK
text: Following (Pennington et al., 2017; 2018; Schoenholz et al., 2017), we consider a feed-forward, fully connected neural network with L hidden layers. Each layer l ∈ {1, . . . , L} is given as a recursion of the form
xl = φ(hl), hl = Wlxl−1 + bl (1)
where xl are the activations, hl are the pre-activations, Wl ∈ RN l×N l−1 are the weight matrices, bl are the bias vectors, and φ(·) is the activation function. The input is denoted as x0. The output layer of the network computes ŷ = g−1(hg) where g is the link function of some generalized linear model (GLM) and hg = WgxL + bg . The hidden layer input-output Jacobian matrix Jx L
x0 is,
Jx L x0 , ∂xL
∂x0 = L∏ l=1 DlWl (2)
where Dl is a diagonal matrix with entries Dli,i = φ ′(hli). As pointed out in (Pennington et al., 2017; Schoenholz et al., 2017), the conditioning of the Jacobian matrix affects the conditioning of the back-propagated gradients for all layers.

heading: 2.2 CRITICAL INITIALIZATIONS
text: Extending the classic result on the Gaussian process limit for wide layer width obtained by (Neal, 1996), recent work (Matthews et al., 2018; Lee et al., 2018) has shown that for deep untrained networks with elements of their weight matrices Wi,j drawn from a Gaussian distribution N (0, σ 2 W
N l )
the empirical distribution of the pre-activations hl converges weakly to a Gaussian distribution N (0, qlI) for each layer l in the limit of the width N → ∞. Similarly, it has been postulated that random orthogonal matrices scaled by σW give rise to the same limit. Under this mean-field condition, the variance of the pre-activation distribution ql is recursively given by,
ql = σ2W
∫ φ (√ ql−1h )
dµ(h) + σ2b (3)
where µ(h) denotes the standard Gaussian measure ∫
dh√ 2π
exp (−h 2
2 ) and σ 2 b denotes the variance
of the Gaussian distributed biases (Schoenholz et al., 2017). The variance of the first layer preactivations q1 depends on `2 norm squared of inputs q1 = σ2W N1 ∥∥p (x0)∥∥2 2
+ σ2b. The recursion defined in equation 3 has a fixed point
q∗ = σ2W
∫ φ (√ q∗h )
dµ(h) + σ2b (4)
which can be satisfied for all layers by appropriately choosing σW, σb and scaling the input x0. To permit the mean field analysis of backpropagated signals, the authors (Schoenholz et al., 2017; Pennington et al., 2017; 2018; Karakida et al., 2019) further assume the propagated activations and back propagated gradients to be independent. Specifically, Assumption 1. [Mean field assumptions] (i) limN→∞ h
d−→ N (0, q∗) (ii) limN→∞ Cov [ Jgxi+1h i,Jgxj+1h j ] = 0 for all i 6= j
Under this assumption, the authors (Schoenholz et al., 2017; Pennington et al., 2017) analyze distributions of singular values of Jacobian matrices between different layers in terms of a small number of parameters, with the calculations of the backpropagated signals proceeding in a selfsame fashion as calculations for the forward propagation of activations. The corollaries of Assumption 1 and condition in equation 4 is that φ′(hl) for 1 ≤ l ≤ L are i.i.d. In order to ensure that JxLx0 is well conditioned, (Pennington et al., 2017) require that in addition to the variance of pre-activation being constant for all layers, two additional constraints be met. Firstly, they require that the mean square singular value of DW for each layer has a certain value in expectation.
χ = 1 N E [ Tr [ (DW)>DW ]] = σ2W ∫ [ φ′( √ q∗h) ]2 dµ(h) (5)
Given that the mean squared singular value of the Jacobian matrix Jx L
x0 is (χ) L, setting χ = 1
corresponds to a critical initialization where the gradients are asymptotically stable as L → ∞. Secondly, they require that the maximal squared singular value s2max of the Jacobian J xL
x0 be bounded. Pennington et al. (2017) showed that for weights with Gaussian distributed elements, the maximal singular value increases linearly in depth even if the network is initialized with χ = 1. Fortunately, for orthogonal weights, the maximal singular value smax is bounded even as L→∞ (Pennington et al., 2018).

heading: 3 THEORETICAL RESULTS: RELATING THE SPECTRA OF JACOBIAN AND
text: FISHER INFORMATION MATRICES
To better understand the geometry of the optimization landscape, we wish to put a Lipschitz bound on the gradient, which in turn gives an upper bound on the largest step size of any first order optimization algorithm. For a general objective function f , the condition is equivalent to
‖∇f(x)−∇f(x′)‖2 ≤M‖x− x′‖2 for all x, x′ ⊂ S ⊆ Rd
The Lipschitz constant ensures that the gradient doesn’t change arbitrarily fast with respect to x, x′, and therefore ∇f defines a descent direction for the objective over a distance M . In general estimating the Lipschitz constant is NP-hard (Kunstner et al., 2019), therefore we seek to find local measures of curvature along the optimization trajectory. As we will show below the approximate gradient smoothness is tractable for randomly initialized neural networks. The analytical study of Hessians of random neural networks started with (Pennington & Bahri, 2017), but was limited to shallow architectures. Subsequent work by Amari et al. (2019) and Karakida et al. (2019) on
second order geometry of random networks shares much of the spirit of the current work, in that it proposes to replace the possibly indefinite Hessian with the related Fisher information matrix as a measure of curvature. The Fisher information matrix plays a fundamental role in the geometry of probabilistic models under the Kullback-Leibler divergence loss — it defines a (local) Riemannian metric, which in turn defines distances on the manifold of probability distributions generated by the model. Notably, the FIM does not define a unique metric on this statistical manifold, and alternative notions of intrinsic curvature can be derived by replacing the Kullback-Leibler divergence with the 2-Wasserstein distance (Li & Montúfar, 2018). Moreover, since the Fisher information matrix bears a special relation to the Hessian it can also be seen as defining an approximate curvature matrix for second order optimization. Recall that the FIM is defined as
Definition. Fisher Information Matrix G , Epθ(y|x0) [ Ep(x0) [ ∇θ log pθ(y|x0)∇θ log pθ(y|x0)> ]] (6)
= Epθ(y|x0) [ Ep(x0) [ Jh g> θ ∇2hgLJh g θ ]] = Epθ(y|x0) [ Ep(x0) [ H− ∑ k ∇xgLk∇2θhgk ]] (7)
where L denotes the loss and hg is the output layer. The relation between the Hessian and Fisher Information matrices is apparent from equation 7, showing that the Hessian H is a quadratic form of the Jacobian matrices plus the possibly indefinite matrix of second derivatives with respect to parameters.
Our goal is to express the gradient smoothness using the results of the previous section. Given equation 7 we can derive an analytical approximation to the Lipschitz bound using the results from the previous section; i.e. we will express the expected maximum eigenvalue of the random Fisher information matrix in terms of the expected maximum singular value of the Jacobian Jh L
h1 . To do so, let us consider the output of a multilayer perceptron as defining a conditional probability distribution pθ(y|x0), where Θ = {vec(W1), . . . , vec(WL),b1, . . . ,bL} is the set of all hidden layer parameters, and θ is the column vector containing the concatenation of all the parameters in Θ. As observed by Martens & Grosse (2015) the Fisher of a multilayer network naturally has a block structure, with each corresponding to the weights and biases of each layer. These blocks with respect to parameter vectors a, b ∈ Θ can further be expressed as
Ga,b = J hg a >HgJ hg b (8)
where the final layer Hessian Hg is defined as∇2hg log pθ(y|x0). We can re-express the outer product of the score function∇hg log pθ(y|x0) as the second derivative of the log-likelihood (see equation 6), provided it satisfies certain technical conditions. What is important for us is that all canonical link function for generalized linear models, like the softmax function and the identity function allow this re-writing, and that this re-writing allows us drop the conditional expectation with respect to pθ ( y|x0 ) . The Jacobians in equation 8 can be computed iteratively. Importantly the Jacobian from the output layer to the a-th parameter block is just the product of diagonal activations and weight matrices multiplied by the Jacobian from the α-th layer to the a-th parameter. We define these matrices of partial derivatives of the α-th layer pre-activations with respect to the layer specific parameters separately for Wα and bα as:
Jh α a = x α−1> ⊗ I for a = vec(Wα) (9) Jh α
a = I for a = b α (10)
Under the infinitesmally weak correlation assumption (see Assumption 1), we can further simplify the expression for the blocks of the Fisher information matrix equation 8.
Lemma 1. The expected blocks with respect to weight matrices for all layers α, β 6= 1 are Gvec(Wα),vec(Wβ) = E [ xα−1xβ−1 >]⊗ E [Jhghα>HgJhhβ] (11) Lemma 2. The expected blocks with respect to a weight matrix Wα and a bias vector bβ are
Gvec(Wα),bβ = E [ xα−1 > ⊗ I ] E [ Jh g hα > HgJ hg hβ ] (12)
The crucial observation here is that in the mean-field limit the expectation of the product of activations xα−1, xβ−1 is either zero or rank 1 for activations in different layers. The case when both activations are in the same layer is trivially taken care of by our mean-field assumptions — the term is equal to the 2nd non-central moment, i.e. the covariance plus potentially a rank one mean term.
Now, leveraging Lemmas 1 and 2 we derive a block diagonal approximation which in turn allows us to bound the maximum eigenvalue λmax(G). In doing so we will use a corollary of the block Gershgorin theorem.
Proposition 1 ((informal) Block Gershgorin theorem). The maximum eigenvalue λmax(G) is contained in a union of disks centered around the maximal eigenvalue of each diagonal block with radia equal to the sum of the singular values of the off-diagonal terms.
For a rigorous statement of the theorem see Appendix A.1. It is noteworthy that block-diagonal approximations have been crucial to the application of Fisher Information matrices as preconditioners in stochastic second order methods (Botev et al., 2017; Martens & Grosse, 2015). These methods were motivated by practical performance in their choice of number of diagonal blocks used for preconditioning. Under the mean-field assumptions we are able to show computable bounds on the error in approximating the spectrum of the Fisher information matrix. The proposition 1 suggest a simple, easily computable way to bound the expected maximal eigenvalue of the Fisher information matrix—choose the block with the largest eigenvalue and expected spectral radia for the corresponding off diagonal terms. We do so by making an auxiliary assumption:
Assumption 2. The maximum singular value of Jh g hα monotonically increases as α ↓ 1.
We motivate this assumption in a twofold fashion: firstly the work done by Pennington et al. (2017; 2018) shows that the spectral edge, i.e. the maximal, non-negative singular value in the support of of the spectral distribution increases with depth, secondly it has been commonly observed in numerical experiments that very deep neural networks have ill conditioned gradients.
Under this assumption it is sufficient to study the maximal singular value of blocks of the Fisher information matrix with respect to vec(W1), b1 and the spectral norms of its corresponding offdiagonal blocks. We define functions Σmax of each block as upper bounds on the spectral bounds of the respective block. The specific values are given in the following Lemma: Lemma 3. The maximum expected singular values of the off-diagonal blocks ∀β 6= 1 are bounded by Σmax(·) : for weight-to-weight blocks
E [ σmax ( Gvec(W1),vec(Wβ) )] ≤ Σmax ( Gvec(W1),vec(Wβ) ) (13)
, √ Nβ |E [φ(h)]| ∥∥E [x0]∥∥ 2 E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (14)
for weight-to-bias blocks E [ σmax ( Gvec(W1),bβ )] ≤ Σmax ( Gvec(W1),bβ ) (15)
, |E [φ(h)]|E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (16)
and for bias-to-bias blocks E [ σmax ( Gb1,bβ )] ≤ Σmax ( Gb1,bβ ) , E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (17)

heading: For proof see Appendix A.2
text: Note that the expectations for layers > 1 is over random networks realizations and averaged over data x0; i.e. they are taken with respect to the Gaussian measure, whereas the expectation for first layer weights is taken with respect to the empirical distribution of x0 (see equation 4). Depending on the choice of q∗ and therefore implicitly both the rescaling of x0 and the values of E[φ(h)] the singular values of the weight blocks might dominate those associated with biases dominate — compare equation 14 and equation 17. Theorem (Bound on the Fisher Information Eigenvalues). If ∥∥E [x0]∥∥ 2 ≤ 1 then eigenvalue associated with b1 will dominate, giving an upper bound on λmax(G)
E [λmax(G)] ≤ E [ σmax ( Gb1,b1 )] + Σmax ( Gb1,vec(W1) ) + ∑ β>1 Σmax ( Gb1,bβ ) + Σmax ( Gvec(b1),vec(Wβ)
) otherwise the maximal eigenvalue of the FIM is bounded by
E [λmax(G)] ≤ E [ σmax ( Gvec(W1),vec(W1) )] + Σmax ( Gb1,vec(W1) ) + ∑ β>1 Σmax ( Gvec(W1),bβ ) + Σmax ( Gvec(W1),vec(Wβ)
) Moreover, it is interesting to note two things. Firstly, E [ σmax ( Jh g
h1
)] factor appear in
all the above summands. Secondly, we can bound σmax for the diagonal blocks with E [λmax (Hg)]E [ σmax ( Jh g
h1
)]2 . These two fact reveal that the FIM maximum eigenvalue is up-
per upperbounded by a quadratic function of the spectral radius of the input-output Jacobian; i.e. λmax (G) is O ( σ2max ( Jh g h1 )) . For proof see Appendix A.4
The functional form of the bound is essentially quadratic in E [ σmax(J hg h1 ) ]
since the term appears in the summand as with powers at most two. This result shows that the strong smoothness, given by the maximum eigenvalue of the FIM, is proportional to the squared maximum singular value of the input-output Jacobian (see Fig. 2). Moreover, the bound essentially depends on q∗ via the expectation E[φ(h)], through Jhgh1 and implicitly through Hg. For regression problems this dependence is monotonically increasing in q∗ (Pennington et al., 2018; 2017) since Hg is just the identity. However, this does not hold for all generalized linear models since λmax(Hg) is not necessarily a monotonically increasing function of the pre-activation variance at layer hg. We demonstrate this in the case of softmax regression in the Appendix A.3. Finally, to obtain a specific bound on λmax(G) we might consider bounding each E [ σmax(J hg hα) ]
appearing in Theorem 3 in terms of its Frobenius norm. The corresponding result is the eigenvalue bound derived by Karakida et al. (2019).

heading: 3.1 NUMERICAL EXPERIMENTS
text: We first verify the bound derived in section 3 for networks with random orthogonal weights. We then numerically investigate the behavior of the maximum FIM eigenvalue during training with particular attention being paid to the possible benefits of maintaining orthogonality or near orthogonality during optimization in relation to unconstrained networks. Following Pennington et al. (2017) we trained a 200 layer tanh network on CIFAR-10 and SVHN1 and we set the width of each layer to be N = 400 and chose the σW, σb in such a way to ensure that mean singular value of the input-output Jacobian concentrates on 1 but s2max varies as a function of q
∗ (see Fig. 2). We considered four different critical initializations with q∗ = [ 10−4, 164 , 1 2 , 8 ] , which differ both in spread of the singular values as well as in the resulting training speed and final test accuracy as reported by (Pennington et al., 2017). In the main text we focus on the smaller values since those networks should be closer to being isometric and therefore, by our theory, ought to train better. The remaining two networks with q∗ = [ 1 2 , 8 ] are presented in the Appendix A.5. To test how enforcing strict orthogonality or near orthogonality affects convergence speed and the maximum eigenvalues of the Fisher information matrix, we trained Stiefel and Oblique constrained networks and compared them to the unconstrained “Euclidean” network described in (Pennington et al., 2017). We used a Riemannian version of ADAM (Kingma & Ba, 2015). When performing gradient descent on non-Euclidean manifolds, we split the variables into three groups: (1) Euclidean variables (e.g. the weights of the classifier layer, biases), (2) non-negative scaling σW both optimized using the regular version of ADAM, and (3) manifold variables optimized using Riemannian ADAM. The initial learning rates for all the groups, as well as the non-orthogonality penalty (see 43) for Oblique networks were chosen via Bayesian optimization, maximizing validation set accuracy after 50 epochs. All networks were trained with a minibatch size of 1000. We trained 5 networks of each kind, and collected eigenvalue and singular value statistics every 5 epochs, from the first to the fiftieth, and then after the hundredth and two hundredth epochs.
Based on the bound on the maximum eigenvalue of the Fisher information matrix derived in Section 3, we predicted that at initialization λmax(G) should covary with σ2max(J hg
x0 ). We tested our prediction using the empirical Fisher information matrix (Kunstner et al., 2019) and we find a significant correlation between the two (Pearson coefficient ρ = 0.64). We additionally test the relation for mean-squared error loss, and observe that our theoretical predictions hold better in this case — the discrepancy can be attributed to the effect of the Hessian of the output layer, and its non-monotonic relation between preactivation variance and maximum eigenvalue of the Hessian of the GLM layer. The numerical values are presented in Fig. 2. Additionally we see that both the maximum singular value and maximum eigenvalue increase monotonically as a function of q∗. Motivated by the previous work by (Saxe et al., 2014) showing depth independent learning dynamics in linear orthogonal networks, we included 5 instantiations of this model in the comparison. The input to the linear
1Code available on: https://github.com/PiotrSokol/info-geom
network was normalized the same way as the critical, non-linear networks with q∗ = 1/64. The deep linear networks had a substantially larger λmax(G) than its non-linear counterparts initialized with identically scaled input (Fig. 2). Having established a connection between q∗ the maximum singular value of the hidden layer input-output Jacobian and the maximum eigenvalue of the Fisher information, we investigate the effects of initialization on subsequent optimization. As reported by (Pennington et al., 2017), the learning speed and generalization peak at intermediate values of q∗ ≈ 10−0.5. This result is counter intuitive given that the maximum eigenvalue of the Fisher information matrix, much like that of the Hessian in convex optimization, upper bounds the maximal learning rate (Boyd & Vandenberghe, 2004; Bottou et al., 2018). To gain insight into the effects of the choice of q∗ on the convergence rate, we trained the Euclidean networks and estimated the local values of λmax during optimization. At the same time we asked whether we can effectively control the two aforesaid quantities by constraining the weights of each layer to be orthogonal or near orthogonal. To this end we trained Stiefel and Oblique networks and recorded the same statistics. We present training results in Fig. 1, where it can be seen that Euclidean networks with q∗ ≈ 9× 10−4 perform worse with respect to training loss and test accuracy than those initialized with q∗ = 1/64. On the other hand, manifold constrained networks are insensitive to the choice of q∗. Moreover, Stiefel and Oblique networks perform marginally worse on the test set compared to the Euclidean network with q∗ = 1/64, despite attaining a lower training loss. We observe that reduced performance of Euclidean networks initialized with q∗ ≈ 9×10−4 may partially be explained by their rapid increase in λmax(G) within the initial 5 epochs of optimization (see Fig. 3 in the Appendix). While all networks undergo this rapid increase, it is most pronounced for Euclidean networks with q∗ ≈ 9× 10−4. The increase λmax(G) correlates with the inflection point in the training loss curve that can be seen in the inset of Fig. 1. Interestingly, the manifold constrained networks optimize efficiently despite differences in λmax(G) (Kohler et al., 2019).

heading: 4 DISCUSSION
text: Critical orthogonal initializations have proven tremendously successful in rapidly training very deep neural networks (Pennington et al., 2017; Chen et al., 2018; Pennington et al., 2018; Xiao et al., 2018). Despite their elegant derivation drawing on methods from free probability and mean field theory, they did not offer a clear optimization perspective on the mechanisms driving their success. With this work we complement the understanding of critical orthogonal initializations by showing that the maximum eigenvalue of the Fisher information matrix, and consequentially the local gradient smoothness is proportional to the maximum singular value of the input-output Jacobian. This gives an information geometric account of why the step size and training speed depend on q∗
via its effect on E [ smax(J hL h1 ) ] . We observed in numerical experiments that the paradoxical results reported in (Pennington et al., 2017) whereby training speed and generalization attains a maximum for q∗ = 10−0.5 can potentially be explained by a rapid increase of the maximum eigenvalue of the FIM during training for the networks initialized with Jacobians closer to being isometric (i.e., smaller q∗). This increase effectively limits the learning rate during the early phase of optimization
and highlights the need to analyze the trajectories of training rather than just initializations. Notably these effects persist under extensive search for the optimal learning rate, indicating the increased difficulty in training highly isometric neural networks reflects a qualitative change of the optimization problem. We conjecture that this phenomenon relates to the stability of descent direction in response to small parameter perturbations. An analysis of the cubic term (Riemmanian connection) would give insight into the initial change in curvature at initialization, however that currently seems infeasible both numerically and analytically. In lieu of that, let us consider the stability of neural networks linearized about their initial parameters. A number of works on overparametrized networks have derived stability conditions for the linearization both in infinitesmal regimes and also over potentially infinitely many gradient updates (Chizat et al., 2019; Lee et al., 2019). The crucial property governing this stability is ensuring that the relative change in parameters is small in proportion to the norm of the Gram matrix of the gradients of the networks. The latter curiously has a spectrum that coincides with that of the Fisher information matrix under mean-squared loss (see A.7). Other arguments about the stability of the linearization have proposed bounds that depend non-trivially on a function of both the largest and smallest eigenvalues of this Gram matrix (Lee et al., 2019). It is therefore tempting to understand the behavior of the lowest eigenvalue of this curvature matrix. Lower bounds on the smallest eigenvalue are typically much harder to obtain, however Karakida et al. (2019) showed that the smallest eigenvalue of the Fisher scales roughly reciprocally with the maximal one. This might imply that a low condition number λmax(G0)λmax(G0) may be undesirable, and a degree of anisotropy is necessary for the Fisher information matrix at initialization to be predictive of training performance. We observe that orthogonal linear networks experience a similar rapid increase in the maximum eigenvalue of their FIM, while having an FIM condition number that isO(1), up to applying similarity transform to the input data (see Appendix A.8).
Finally, we compared manifold constrained networks with the Euclidean network, each evaluated with two initial values of q∗. From these experiments we draw the conclusion that manifold constrained networks are less sensitive to the initial strong smoothness, unlike their Euclidean counterparts. Furthermore, we observe that the rate at which Stiefel and Oblique networks decrease training loss is not dependent on their gradient smoothness, a result which is consistent with the recent analysis of (Kohler et al., 2019).

heading: A APPENDIX

heading: A.1 BLOCK GERSHGORIN THEOREM
text: In Section 3, we considered a block diagonal approximation to the Fisher information matrix and derived an upper bound on the spectral norm for all the blocks. Using the properties of the off-diagonal blocks, we can get a more accurate estimate of the maximal eigenvalue of the Fisher information might be. First, let us consider an arbitrarily partitioned matrix A ∈ RN×N , with spectrum λ(A) The partitioning is done with respect to the set
π = {pj}Lj=0 (18) with the elements of the set satisfying 0 < p1 < p2 < . . . < pL = N . Then each block of the matrix Ai,j is a potentially rectangular matrix in R(pi−pi−1)×(pj−pj−1). We assume that Ai,i is self-adjoint for all i.
Let us define a disk as C(c, r) , { λ : ‖c− λ‖ ≤ r } . (19)
The theorem as presented in (Tretter, 2008) shows that the eigenvalues of λ(A) are contained in a union of Gershgorin disks defined as follows
λ(A) ⊂ L⋃ i=1 { pi−pi−1⋃ k=1 C λk(Aii), L∑ j=1,j 6=i smax(Ai,j) } (20) where the inner union is over a set disks for each eigenvalue of the block diagonal Ai,i while the outer union is over the L blocks in A. The radius of the disk is constant for every eigenvalue in the ith diagonal block Ai,i and is given by the sum of singular values of the off diagonal blocks. Therefore, the largest eigenvalue of A lies in
λmax(A) ⊂ L⋃ i=1 C λmax(Aii), L∑ j=1,j 6=i smax(Ai,j)  (21) Now suppose A is positive semidefinite. This suggests a strategy for upper-bounding the maximal eigenvalue of A— picking the disk containing the largest (non-negative) value will upper-bound the maximal eigenvalue of A:
λmax (A) ≤ max i λmax(Aii) + L∑ j=1,j 6=i smax(Ai,j)  (22)

heading: A.2 DERIVATION OF THE EXPECTED SINGULAR VALUES
text: We present a derivation of the bounds presented in Lemma 3. The singular value bound for the weight-to-weight blocks is
σmax ( Gvec(W1),vec(Wβ) ) = E [ σmax ( φ(h)1x0> )] ⊗ E [ σmax ( Jh g> h1 HgJ hg hβ )] (23)
= √ Nβ |E [φ(h)]| ∥∥E [x0]∥∥ 2 E [ σmax ( Jh g> h1 HgJ hg hβ )] (24)
where φ(h)1 is a vector of ones times φ(h)
≤ √ Nβ |E [φ(h)]| ∥∥E [x0]∥∥ 2 E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ ]) (25)
The singular value bound for the weight-to-bias blocks is σmax ( Gvec(W1),bβ ) ≤ E [ σmax ( x0> ⊗ I )] E [ σmax ( Jh g hα >HgJ hg hβ )] (26)
= ∥∥E [x0]∥∥ 2 (E [ σmax J hg> h1 HgJ hg hβ )] (27)
≤ ∥∥E [x0]∥∥ 2 E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (28)
σmax ( Gb1,vec(Wβ) ) ≤ E [ σmax ( xβ−1> ⊗ I )] (29)
= √ Nβ |E [φ(h)]|E [ σmax ( Jh g> h1 HgJ hg hβ )] (30)
≤ √ Nβ |E [φ(h)]|E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (31)
The singular value bound for the weight-to-bias blocks is σmax ( Gb1,bβ ) = E [ σmax ( Jh g> h1 HgJ hg hβ )] (32)
≤ E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (33)

heading: A.3 MONTECARLO ESTIMATE OF MAXIMUM EIGENVALUE OF THE HESSIAN OF THE OUTPUT LAYER FOR 10 WAY SOFTMAX CLASSIFICATION
text: Variance of preactivation q*
D is tr
ib ut
io n
of λ
m a x ( H
g )
Figure 4: Distribution of λmax(Hg) as a function of q∗: In general, increasing the variance of the distribution of hg does not result in a monotonic increase in the spectral radius of the Hessian of the GLM layer. We plot the distribution of the maximum eigenvalues as a function of the variance of the softmax layer obtained from factorizing 10,000 random matrices.

heading: A.4 PROOF OF THEOREM
text: Using the results from Section A.2, we observe that all the terms off diagonal for the diagonal block Gvec(W1),vec(W1) are of the form
|E [φ(h)]| ∥∥E [x0]∥∥ 2 E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (34)
and ∥∥E [x0]∥∥ 2 E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (35)
while the diagonal term is λmax(E [ x0ᵀx0 ] )E [ σmax ( Jh g> h1 )] E [σmax (Hg)]E [ σmax ( Jh g hβ )] (36)
then by Gershgorin theorem the eigenvalues for the Gvec(W1),vec(W1) are bounded by λmax(Gvec(W1),vec(W1)) ≤ E [ σmax ( Jh g> h1 )] E [σmax (Hg)]
× λmax(E [x0ᵀx0])E [σmax (Jhg>h1 )]+∑ β (1 + |E[φ(h)]|) ∥∥E [x0]∥∥ 2 E [ σmax ( Jh g hβ )] (37)
Similarly, for the Gb1,b1 block, we observe that λmax(Gb1,b1) ≤ ( E [ σmax ( Jh g> h1 )] E [σmax (Hg)] ) ×
× ∑ β (1 + √ Nβ |E[φ(h)]|) ∥∥E [x0]∥∥ 2 E [ σmax ( Jh g hβ )] (38)

heading: A.5 ADDITIONAL FIGURES

heading: CIFAR-10 performance

heading: A.6 MANIFOLD OPTIMIZATION
text: The potentially non-convex constraint set constitutes a Riemannian manifold, when it is locally isomorphic to Rn, differentiable and endowed with a suitable (Riemannian) metric, which allows us to measure distances in the tangent space and consequentially also define distances on the manifold. There is considerable freedom in choosing a Riemannian metric; here we consider the metric inherited from the Euclidean embedding space which is defined as 〈W,W′〉 , Tr(W′>W). To optimize a cost function with respect to parameters lying in a non-Euclidean manifold we must define a descent direction. This is done by defining a manifold equivalent of the directional derivative. An intuitive approach replaces the movement along a vector t with movement along a geodesic curve γ(t), which lies in the manifold and connects two points W,W′ ∈M such that γ(0) = W, γ(1) = W′. The derivative of an arbitrary smooth function f(γ(t)) with respect to t then defines a tangent vector for each t.
Tangent vector ξW is a tangent vector at W if ξW satisfies γ(0) = W and
ξW , df(γ(t))
dt ∣∣∣∣ t=0 , γ′(0)f (39)
The set of all tangents toM at W is referred to as the tangent space toM at W and is denoted by TWM. The geodesic importantly is then specified by a constant velocity curve γ′′(t) = 0 with initial velocity ξW. To perform a gradient step, we must then move along ξW while respecting the manifold constraint. This is achieved by applying the exponential map defined as ExpW(ξW) , γ(1), which moves W to another point W′ along the geodesic. While certain manifolds, such as the Oblique manifold, have efficient closed-form exponential maps, for general Riemannian manifolds, the computation of the exponential map involves numerical solution to a non-linear ordinary differential equation (Absil et al., 2007). An efficient alternative to numerical integration is given by an orthogonal projection onto the manifold. This projection is formally referred to as a retraction RtW : TWM→M. Finally, gradient methods using Polyak (heavy ball) momentum (e.g. ADAM (Kingma & Ba, 2015)) require the iterative updating of terms which naturally lie in the tangent space. The parallel translation Tζ(ξ) : TM ⊕ TM → TM generalizes vector composition from Euclidean to non-Euclidean manifolds, by moving the tangent ξ along the geodesic with initial velocity ζ ∈ T and endpoint W′, and then projecting the resulting vector onto the tangent space TW′M. As with the exponential map,
parallel transport T may require the solution of non-linear ordinary differential equation. To alleviate the computational burden, we consider vector transport as an effective, projection-like solution to the parallel translation problem. We overload the notation and also denote it as T , highlighting the similar role that the two mappings share. Technically, the geodesics and consequentially the exponential map, retraction as well as transport T depend on the choice of the Riemannian metric. Putting the equations together the updating scheme for Riemannian stochastic gradient descent on the manifold is
Wt+1 = ΠWt(−ηt gradf) (40) where Π is either the exponential map Exp or the retraction Rt and gradf is the gradient of the function f(W) lying in the tangent space TWM.

heading: A.6.1 OPTIMIZING OVER THE OBLIQUE MANIFOLD
text: (Cho & Lee, 2017) proposed an updating scheme for optimizing neural networks where the weights of each layer are constrained to lie in the oblique manifold Ob(p, n). Using the fact that the manifold itself is a product of p unit-norm spherical manifolds, they derived an efficient, closedform Riemannian gradient descent updating scheme. In particular the optimization simplifies to the optimization over Ob(1, n) for each column wi∈{1,...,p} of W.
Oblique gradient The gradient gradf of the cost function f with respect to the weights lying in Ob(1, n) is given as a projection of the Euclidean gradient Gradfonto the tangent at w
gradf = Gradf − (w>Gradf)w (41)
Oblique exponential map The exponential map Expw moving w to w′ along a geodesic with initial velocity ξw
Expw = ξw cos(‖w‖) + w
‖w‖ sin(‖w‖) (42)
Oblique parallel translation The parallel translation T moves the tangent vector ξw along the geodesic with initial velocity ζw
Tζw(ξw) = ξw − ζw ‖ζw‖ ((1− cos(‖ζw‖)) + w sin(‖ζw‖)) ζw ‖ζw‖ >ξw
Orthogonal penalty which enforces the weight matrices to be near orthogonal.
ρ(λ,W) = λ
2
∥∥W>W − I∥∥2 F
(43)

heading: A.6.2 OPTIMIZING OVER THE STIEFEL MANIFOLD
text: Optimization over Stiefel manifolds in the context of neural networks has been studied by (Harandi & Fernando, 2016; Wisdom et al., 2016; Vorontsov et al., 2017). Unlike (Wisdom et al., 2016; Vorontsov et al., 2017) we propose the parametrization using the Euclidean metric, which results in a different definition of vector transport.
Stiefel gradient The gradient gradf of the cost function f with respect to the weights lying in St(p, n) is given as a projection of the Euclidean gradient Gradfonto the tangent at W (Edelman et al., 1999; Absil et al., 2007)
gradf = (I−WW>)Gradf + 1 2 W ( W>Gradf −Gradf>W ) Stiefel retraction The retraction RtW(ξW) for the Stiefel manifold is given by the Q factor of the QR decomposition (Absil et al., 2007).
RtW(ξW) = qf(W + ξW) (44)
Stiefel vector transport The vector transport T moves the tangent vector ξw along the geodesic with initial velocity ζw for W ∈ St(p, n) endowed with the Euclidean metric.
Tζw(ξw) = ( I−YY> ) ξW + 1 2 Y ( Y>ξW − ξW>Y ) (45)
where Y , RtW(ζW). It is easy to see that the transport T consists of a retraction of tangent ζW followed by the orthogonal projection of ηW at RtW(ζW).

heading: A.6.3 OPTIMIZING OVER NON-COMPACT MANIFOLDS
text: The critical weight initialization yielding a singular spectrum of the Jacobian tightly concentrating on 1 implies that a substantial fraction of the pre-activations lie in expectation in the linear regime of the squashing nonlinearity and as a consequence the network acts quasi-linearly. To relax this constraint during training we allow the scales of the manifold constrained weights to vary. We chose to represent the weights as a product of a scaling diagonal matrix and a matrix belonging to the manifold. Then the optimization of each layer consists in the optimization of the two variables in the product. In this work we only consider isotropic scalings, but the method generalizes easily to the use of any invertible square matrix.
It is interesting to contrast our approach with projected gradient descent with a spectral prox function, or penalizing the Jacobian norm through back propagation. PGD with projections onto compact spectral balls requires the user to pre-specify the desired spectral radius. On the other hand, both our approach as well as penalizing the Jacobian norm can be thought of penalizing decreases in function smoothness in an adaptive way. Finally, our approach is naturally amenable to incorporating a data-dependent penalty — this would allow us to smoothly vary the spectral radia of the weight matrices.

heading: A.7 FIM AND NTK HAVE THE SAME SPECTRUM
text: The empirical Neural Tangent Kernel (NTK) is defined as
Θ̂t,i,j , J hg θ J hg> θ (46)
which gives a Ng|D| by Ng|D| kernel matrix. By comparison the empirical Fisher information matrix with a Gaussian likelihood is ∑
i∈|D|
Jh g> θ J hg θ (47)
To see that the spectra of these two coincide consider the third order tensor underlying both Jh g
h1i for i ∈ 1 . . . |D|, additionally consider and unfolding A with dimensions |θ| by Ng|D|; i.e. we construct a matrix with dimension of number of parameters by number of outputs times number of data points. Then
G = A>A (48)
Θ̂ = AA> (49) (50)
and their spectra trivially coincide.

heading: A.8 LINEAR NEURAL NETWORKS WITH ORTHOGONAL INITIALIZATIONS

Review summaries:
This paper analyses the training behavior of wide networks and argues orthogonal initialization helps the training. They suggest projections to the manifold of orthogonal weights during training and provide analysis. Their main result seems to be a bound on the eigen-values of the Fisher information matrix for wide networks (Theorem on pg 6). In their experiments they train Stiefel and Oblique networks as examples of manifold constrained networks and claim they converge faster than unconstrained networks.

This paper formulates a connection between the Fisher information matrix (FIM) and the spectral radius of the input-output Jacobian in neural networks. This results derive the eigenvalues' bound to theoretically study the convergence of several networks.  Here the upper bound further improves the upper bound of FIM derived in (Karakida et al., 2018).