Information Geometry of Orthogonal Initializations and Training

1 INTRODUCTION . 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 can not 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 can not 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 . 2 BACKGROUND . 2.1 FORMAL DESCRIPTION OF THE NETWORK . 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 . 2.2 CRITICAL INITIALIZATIONS . 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 ) . 3 THEORETICAL RESULTS : RELATING THE SPECTRA OF JACOBIAN AND . 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 )

Summary: 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.