On the Variance of the Fisher Information for Deep Learning

1 Introduction . The Fisher information is one of the most fundamental concepts in statistical machine learning . Intuitively , it measures the amount of information carried by a single random observation when the underlying model varies along certain directions in the parameter space : if such a variation does not change the underlying model , then a corresponding observation contains zero ( Fisher ) information and is non-informative regarding the varied parameter . Parameter estimation is impossible in this case . Otherwise , if the variation significantly changes the model and has large information , then an observation is informative and the parameter estimation can be more efficient as compared to parameters with small Fisher information . In machine learning , this basic concept is useful for defining intrinsic structures of the parameter space , measuring model complexity , and performing gradient-based optimization . Given a statistical model that is specified by a parametric form p ( z |θ ) and a continuous domain θ ∈ M , the Fisher information matrix ( FIM ) is a 2D tensor varying with θ ∈ M , given by I ( θ ) = Ep ( z | θ ) ( ∂ℓ ∂θ ∂ℓ ∂θ⊤ ) , ( 1 ) where Ep ( z | θ ) ( · ) , or simply Ep ( · ) if the model p is clear from the context , denotes the expectation w.r.t . p ( z |θ ) , and ℓ : = log p ( z |θ ) is the log-likelihood function . All vectors are column vectors throughout this paper . Under weak conditions ( see Lemma 5.3 in Lehmann and Casella [ 16 ] for the univariate case ) , the FIM has the equivalent expression I ( θ ) = Ep ( z | θ ) ( −∂2ℓ/∂θ∂θ⊤ ) . Given N i.i.d . observations z1 , . . . , zN , these two equivalent expressions of the FIM lead to two different estimators Î1 ( θ ) = 1 N N∑ i=1 ( ∂ℓi ∂θ ∂ℓi ∂θ⊤ ) and Î2 ( θ ) = 1 N N∑ i=1 ( − ∂ 2ℓi ∂θ∂θ⊤ ) , ( 2 ) 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . where ℓi : = log p ( zi |θ ) is the log-likelihood of the i ’ th observation zi . The notations Î1 ( θ ) and Î2 ( θ ) are abused for simplicity as they depend on both θ and the random observations zi . These estimators are universal and independent to the parametric form p ( z |θ ) . They are expressed in terms of the 1st- or 2nd-order derivatives of the log-likelihood . Usually , we already have these derivatives to perform gradient-based learning . Therefore , we can save computational cost and reuse these derivatives to estimate the Fisher information , which in turn can be useful , e.g. , to perform natural gradient optimization [ 1 , 25 ] . Estimating the FIM is especially meaningful for deep learning , where the computational overhead of the exact FIM can be significant . It is straightforward from the law of large numbers and the central limit theorem that both estimators in Eq . ( 2 ) are unbiased and consistent . This is formally stated as follows . Proposition 1 . Ep ( z | θ ) ( Î1 ( θ ) ) = Ep ( z | θ ) ( Î2 ( θ ) ) = I ( θ ) . ∀ϵ > 0 , lim N→∞ Prob ( ∥∥∥Î1 ( θ ) − I ( θ ) ∥∥∥ F + ∥∥∥Î2 ( θ ) − I ( θ ) ∥∥∥ F > ϵ ) = 0 , where Prob ( · ) denotes the probability of the parameter statement being true and ∥ · ∥F is the Frobenius norm of a tensor ( with ∥ · ∥2 as the regular vector L2-norm ) The Fisher information can be zero for non-regular models or infinite [ 7 ] . However , these properties may not be preserved by the empirical estimators . How far can Î1 ( θ ) and Î2 ( θ ) deviate from the “ true FIM ” I ( θ ) , and how fast can they converge to I ( θ ) as the number of observations increases ? To answer these questions , it is natural to think of the variance of Î1 ( θ ) and Î2 ( θ ) . For example , an estimator with a large variance means the estimation does not accurately reflect I ( θ ) ; and any procedure depending on the FIM consequently suffers from the estimation error . Through studying the variance , we can control the estimation quality and reliably perform subsequent measurements or algorithms based on the FIM . Towards this direction , we made the following contributions that will unfold in the following Sections 2 to 4 : • We review and rediscover two equivalent expression of the FIM in the context of deep feed-forward networks ( Section 2 ) ; • We give in closed form the variance ( extending to meaningful upper bounds ) and discuss the convergence rate of the estimators Î1 ( θ ) and Î2 ( θ ) ( Section 3 ) ; • We analyze how the 1st- and 2nd-order derivatives of the neural network can affect the estimation of the FIM ( Section 4 ) . We discuss related work in Section 5 and conclude in Section 6 . 2 Feed-forward Networks with Exponential Family Output . This section realizes the concept of Fisher information in a feed-forward network with exponentially family output and explains why its estimators are useful in theory and practice . Consider supervised learning with a neural network . The underlying statistical model is p ( z |θ ) = p ( x ) p ( y |x , θ ) , where z = ( x , y ) , the random variable x represents features , and y is the target variable . The marginal distribution p ( x ) is parameter-free , usually fixed as the empirical distribution p ( x ) = 1M ∑M i=1 δ ( x−xi ) w.r.t . a set of observations { xi } Mi=1 , where δ ( · ) is the Dirac delta . In this paper , we consider w.l.o.g . M = 1 as the FIM w.r.t . observations { xi } Mi=1 is simply the average over FIMs of each individual observation . All results generalize to multiple observations by taking the empirical average . The predictor p ( y |x , θ ) is a neural network with parameters θ = { Wl−1 } Ll=1 and exponential family output units , given by p ( y |x ) = exp ( t⊤ ( y ) hL − F ( hL ) ) , hL = WL−1h̄L−1 , hl = σ ( Wl−1h̄l−1 ) , ( l = 1 , . . . , L− 1 ) h̄l = ( h ⊤ l , 1 ) ⊤ , h0 = x , ( 3 ) where t ( y ) is the sufficient statistics of the prediction model , F ( h ) = log ∫ exp ( t⊤ ( y ) h ) dy is the log-partition function , and σ : ℜ → ℜ is an element-wise non-linear activation function . Moreover , Wl is a nl+1 × ( nl + 1 ) matrix , representing the neural network parameters ( weights and biases ) in the l ’ th layer , where nl : = dim ( hl ) denotes the size of layer l. We use W−l for the nl+1 × nl weight matrix without the bias terms , obtained by removing the last column of Wl . hl is a learned representation of the input x . All intermediate variables hl are extended to include a constant scalar 1 in h̄l , so that a linear layer can simply be expressed as Wlh̄l . The last layer ’ s output hL with dimensionality nL specifies the natural parameter of the exponential family . We need the following Lemma which gives the FIM w.r.t . hL , which is a nL × nL matrix in simple closed form for commonly used probability distributions . Lemma 2 . For the neural network model specified in Eq . ( 3 ) , I ( hL ) = Cov ( t ( y ) ) = ∂η ∂hL , where Cov ( · ) denotes the covariance matrix w.r.t . p ( y |x , θ ) , η : = η ( hL ) : = ∂F/∂hL is the expectation parameters , and the vector-vector-derivative ∂η/∂hL denotes the Jacobian matrix of the mapping hL → η . The derivatives of the log-likelihood ℓ ( θ ) : = log p ( x , y |θ ) characterize its landscape and are essential to compute the FIM . By Eq . ( 3 ) , the score function ( gradient of ℓ ) is ∂ℓ ∂θ = ( ∂hL ∂θ ) ⊤ ( t ( y ) − η ( hL ) ) = ∂haL ∂θ ( ta − ηa ) . ( 4 ) In this paper , we mix the usual Σ-notation of summation with the Einstein notation : in the same term , an index appearing in both upper- and lower-positions indicates a sum over this index . For example , tah a = ∑ a taha . Hence , in our equations , upper- and lower-indexes have the same meaning : both ha and ha mean the a ’ th element of h. For convenience and consistency , we take quantities w.r.t . θ as upper indexed and other quantities as lower indexed , i.e. , Iij ( θ ) versus Iij ( hL ) . This mixed representation of sums helps to simplify our expressions without causing confusion . From Eq . ( 4 ) and Lemma 2 , the Hessian of ℓ is given by ∂2ℓ ∂θ∂θ⊤ = ( ta − ηa ) ∂2haL ∂θ∂θ⊤ − ∂h a L ∂θ ∂ηa ∂θ⊤ = ( ta − ηa ) ∂2haL ∂θ∂θ⊤ − ∂h a L ∂θ Iab ( hL ) ∂hbL ∂θ⊤ . ( 5 ) Similar to the case of a general statistical model , the FIM is equivalent to the expectation of the Hessian of −ℓ as long as the activation function is smooth enough . Theorem 3 . Consider the neural network model in Eq . ( 3 ) . For any activation function σ ∈ C2 ( ℜ ) ( both σ′ ( z ) and σ′′ ( z ) exist and are continuous ) , we have I ( θ ) = Ep ( − ∂ 2ℓ ∂θ∂θ⊤ ) . Remark 3.1 . ReLU networks do not have this equivalent expression as ReLU ( z ) is not differentiable at z = 0 . Through the definition of the FIM , or alternatively its equivalent formula in Theorem 3 , we arrive at the same expression I ( θ ) = ( ∂hL ∂θ ) ⊤ I ( hL ) ∂hL ∂θ⊤ = ∂haL ∂θ Iab ( hL ) ∂hbL ∂θ⊤ . ( 6 ) Equation ( 6 ) takes the form of a generalized Gauss-Newton matrix [ 18 ] . This general expression of the FIM has been known in the literature [ 24 , 25 ] . Under weak conditions , I ( θ ) is a pullback metric [ 30 ] of I ( hL ) in Lemma 2 associated with the mapping θ → hL . To compute I ( θ ) in closed form , one need to first compute the Jacobian matrix of size nL × dim ( θ ) then perform the matrix multiplication in Eq . ( 6 ) . The naive algorithm to evaluate Eq . ( 6 ) has a computational complexity of O ( n2L dim ( θ ) +nL dim 2 ( θ ) ) , where the term O ( nL dim2 ( θ ) ) is dominant as dim ( θ ) ≫ nL in deep architectures . Once the parameter θ is updated , the FIM has to be recomputed . This is infeasible in practice for large networks where dim ( θ ) can be millions or billions . The two estimators Î1 ( θ ) and Î2 ( θ ) in Eq . ( 2 ) provide a computationally inexpensive way to estimate the FIM . Given θ and x , one can draw i.i.d . samples y1 , . . . , yN ∼ p ( y |x , θ ) . Both ∂ℓi/∂θ and ∂2ℓi/∂θ∂θ⊤ can be evaluated directly through auto-differentiation ( AD ) that is highly optimized for modern GPUs . For Î1 ( θ ) , we already have ∂ℓi/∂θ to perform gradient descent . For Î2 ( θ ) , efficient methods to compute the Hessian are implemented in AD frameworks such as PyTorch [ 26 ] . Using these derivatives , the computational cost only scales with the number N of samples but does not scale with nL . We rarely need the full FIM of size dim ( θ ) ×dim ( θ ) . Most of the time , only its diagonal blocks are needed , where each block corresponds to a subset of parameters , e.g. , the neural network weights of a particular layer . Therefore the computation of both estimators can be further reduced . If p ( y |x , θ ) has the parametric form in Eq . ( 3 ) , from Eqs . ( 4 ) and ( 5 ) , the FIM estimators become Î1 ( θ ) = ∂haL ∂θ · 1 N N∑ i=1 ( ta ( yi ) − ηa ) ( tb ( yi ) − ηb ) · ∂hbL ∂θ⊤ , ( 7 ) Î2 ( θ ) = ( ηa − 1 N N∑ i=1 ta ( yi ) ) ∂2haL ∂θ∂θ⊤ + ∂haL ∂θ Iab ( hL ) ∂hbL ∂θ⊤ . ( 8 ) Recall that the notation of Î1 ( θ ) and Î2 ( θ ) is abused as they depend on x and y1 · · ·yN . Notably , in Eq . ( 7 ) , Î1 ( θ ) is expressed in terms of the Jacobian matrix of the mapping θ → hL and the empirical variance of the minimal sufficient statistic t ( yi ) of the output exponential family . In Eq . ( 8 ) , Î2 ( θ ) depends on both the Jacobian and Hessian of θ → hL and the empirical average of t ( yi ) . The second term on the right-hand side ( RHS ) of Eq . ( 8 ) is exactly the FIM , and therefore the first term serves as a bias term . Eqs . ( 7 ) and ( 8 ) are only for the case with exponential family output . If the output units belong to non-exponential families , e.g. , a statistical mixture model , one falls back to the general formulae , i.e. , Eq . ( 2 ) for the FIM . As an application of the Fisher information , the Cramér-Rao lower bound ( CRLB ) states that any unbiased estimator θ̂ of the parameters θ satisfies Cov ( θ̂ ) ⪰ [ I ( θ ) ] −1 . For example , in Lemma 2 , the FIM is w.r.t . the output of the neural network . As such , I ( hL ) can be used to study the estimation covariance of hL based on random samples y1 · · ·yN drawn from p ( y |x , θ ) . Similarly for Eq . ( 6 ) , we can consider unbiased estimators of the weights of the neural network . In any case , to apply the CRLB , one needs an accurate estimation of I ( θ ) . If the scale of I ( θ ) is relatively small when compared to its covariance , its estimation Î ( θ ) is more likely to be a small positive value ( or even worse , zero or negative ) . The empirical computation of the CRLB is not meaningful in this case .

Summary: Update: I thank the authors for their response, and I will keep the score as is. It would be nice to have a simple empirical analysis in the final supplemental material.   The paper discusses the covariance of two FIM estimators, which is useful in deep learning where the exact FIM over the full dataset can be hard to compute. One is based on the variance of the score (type 1) and the other is based on the Hessian of the relative entropy (type 2). Comparing these covariances can depend on the exponential family used in the model, type 1 has a higher variance for Normal and Poisson whereas type 2 has a higher variance for Bernoulli. Convergence rate results are given which is $O(N^{-0.5))$ with constant determined by the variances. These results provide some theoretical justification as to which estimator of FIM to use when using it to perform natural gradient in deep learning.