M-FAC: Efficient Matrix-Free Approximations of Second-Order Information

1 Introduction . Given the recent success and increasing impact of deep learning , there has been significant work on improving the fundamental technical tools underpinning its progress . One such tool is the ability to estimate the local geometry of the loss function for deep models , which often comes in the form of estimates for second-order ( Hessian ) information . Such information is critical in several settings , such as neural network optimization and pruning . Directly using Hessian information in the context of deep learning is infeasible : for example , just storing the Hessian matrix for the standard ResNet50 model [ 14 ] would occupy 2.5 Petabytes . These constraints have inspired significant work on efficient numerical approximations of the Hessian for deep neural networks , such as the line of work on the K-FAC approximation [ 19 , 31 , 5 , 49 ] , or efficient block-wise approximations [ 47 , 37 ] . One of the classic approaches , which we focus on in this paper , is the empirical Fisher [ 13 , 3 , 4 ] approximation to the Hessian , written as : H ' F̂ = 1 N N∑ i=1 ∇ ` i · ∇ ` > i , ( 1 ) where N is the number of samples , ∇ ` i denotes the gradient w.r.t . the ith sample at the given point , and · is the outer product of individual gradients , which we view as column vectors . The empirical Fisher approximation is fairly standard , e.g . [ 13 , 3 , 4 ] , and has been recognized to be useful in a variety of practical settings where exact estimation of the Hessian or of the true Fisher information matrix is not feasible . At the same time , there is still active research in the community 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . on the conditions for its applicability [ 29 , 21 , 37 , 40 ] . A useful property of this approximation is that it also allows to estimate the inverse of the Hessian , which is essential in many applications . Specifically , the fact that the empirical Fisher can be written as a sum of rank-one matrices allows the use of the Woodbury-Sherman-Morrison inversion formula [ 46 ] to exactly compute its inverse , by recursively integrating terms corresponding to each gradient into the inverse . ( Please see Equation 2 for an exact derivation . ) This approach was independently proposed by [ 13 , 3 ] , for pruning and optimization , respectively , where it was validated on small networks , with hundreds of weights . The idea was adapted to deep neural networks ( DNNs ) by [ 37 ] , through approximation of the inverse in small diagonal blocks . The authors show improved approximation quality for the Hessian inverse relative to a simple diagonal approximation , and that this leads to state-of-the-art pruning results in terms of accuracy . Yet , this approach is limited by the block-wise approximation : for block size B and dimension d , it requires Θ ( Bdm ) time to recursively build the block-wise Fisher approximation using m gradients , and Θ ( Bd ) time and memory for computing the Inverse-Hessian-Vector-Products ( IHVPs ) necessary for estimating pruning statistics . Clearly , the ideal B = d case is still intractable at scale , and generally it is still unknown whether efficient algorithms are possible for computing IHVPs in this context . Contribution . We address this question by introducing two efficient algorithms for computing IHVPs under the empirical Fisher approximation , with computational and storage costs that are linear in the dimension d of the model , without the need for block-wise approximations , assuming that the number of samples m in the approximation is constant . Concretely , we provide exact matrix-free algorithms to compute products of the form F̂−1v , where F̂−1 is the inverse empirical Fisher , and v is an arbitrary vector . We show that these algorithms can be implemented efficiently , and that they can match or improve the state-of-the-art results for both neural network pruning and optimization . The Static Algorithm . Our first algorithm assumes a static scenario , which is standard in neural network pruning : we are given a fully-trained model θ ? , for which we wish to estimate IHVPs and diagonal elements of the inverse Hessian , in order to determine the “ optimal ” pruning update using e.g . the Optimal Brain Surgeon ( OBS ) framework [ 24 , 13 ] . For this , we first compute m gradients at θ ? , which we will use to estimate IHVPs via the empirical Fisher and compute pruning statistics . The main idea is that , since we only wish to compute products between the inverse and an arbitrary vector ( IHVPs ) , we can rewrite the Woodbury recursion such that we work exclusively with individual vectors and scalars , and never with full d× d or d×B matrices . Given model dimension d and m gradients , each defining a rank-one component of the empirical Fisher , the algorithm uses O ( dm2 ) pre-computation time , and will have O ( dm ) cost for exactly computing the IHVP . Further , we can specialize the algorithm to directly query elements of the Hessian inverse , at a cost of O ( m ) time per element . This provides efficient , linear-in-d implementations for all operations required by the OBS pruning framework . Finally , we note that the static algorithm can also be applied in a block-wise manner without any change in the total compute and memory costs . The Dynamic Algorithm . Our main contribution is in extending this idea to preconditioned SGD optimization , i.e . to precondition stochastic gradients by our estimate of the inverse Hessian . We start from the classic idea of bootstrapping the approximation by leveraging previous gradients : the preconditioner at time t is built from gradients obtained during a “ sliding window ” over the last ∆ ≥ 1 optimization steps . This requires a dynamic representation , allowing addition and removal of gradients without full recomputation of second-order statistics . We show that this can be achieved , with approximately O ( dm ) time and space complexity . The key idea is that , for any ordered set ( ∇ ` j ) mj=1 of m gradients , and any vector v , it is possible to represent the corresponding IHVP estimate F̂−1m v as a linear combination of terms corresponding to individual gradients ∇ ` j and v , of the form F̂−1m v = λ−1v −∑m j=1 c m j ∇ ` j , where λ > 0 is a dampening constant . Crucially , we ensure that the coefficients cmj can be computed just via dot products∇ ` > i ∇ ` j , where i ≤ j in the ordering , and ∇ ` > j v. Then , to replace a given gradient ∇ ` i from this representation , we just have to compute m scalar products with the new gradient vector ( as well as update some intermediate information ) . Hence , the entire update operation has computational cost O ( dm+m3 ) for replacing any gradient in the sliding window , and O ( dm+m2 ) for computing the IHVP . Implementation and Experiments . We provide efficient vectorized implementations for the above algorithms , called M-FAC , for Matrix-Free Approximate Curvature . Specifically , M-FAC consists of Pytorch [ 34 ] implementations of a pruning and optimization library . Our implementation introduces several additional optimizations , in particular GPU acceleration via custom CUDA kernels , and minimizes the cost of memory transfer between the GPU and main memory via memory paging . For pruning , our implementation provides order-of-magnitude improvements over the block-wise approximation of [ 37 ] for classic benchmarks such as pruning ResNet50 and MobileNet on the ImageNet dataset . This allows us to obtain more accurate sparse models by exploring higher parameter settings and increasing the total number of pruning steps , while remaining practical in terms of memory and compute even for larger models . What is more , our preconditioned SGD ( even without momentum ) can be competitive in terms of validation accuracy with state-of-the-art optimizers on models of moderate size , including compact vision architectures and Transformer language models [ 42 ] . Its computational overheads are of 5 % –55 % relative to vanilla SGD on standard CNN architectures . 2 Preliminaries and Related Work . General Definitions . We now briefly introduce the setting and notation ; we refer the reader to standard texts [ 29 , 27 , 4 ] for a complete introduction . We start from the standard setting in which we are given a dataset D = ( xi , yi ) Ni=1 , and wish to identify a d-dimensional model θ ∈ Rd to minimize an empirical loss L : Rd → R , defined as L ( θ ) = 1N ∑N i=1 ` ( yi , f ( xi ; θ ) ) . For a twice-differentiable loss L , the Hessian is the matrix of second-order derivatives of L w.r.t . θ , i.e . H = ∇2θL . In the probabilistic view , each input example x has some probability of being assigned a given label y . Given input examples x drawn from a distribution Qx and corresponding outputs drawn from a conditional distribution Qy|x , the goal is to minimize the distance between the target joint distribution Qx , y = QxQy|x , and a learned joint distribution Px , y ( θ ) , where θ is the model . The Fisher Matrix . Assuming the probabilistic view , it can be shown that the Fisher information matrix F of the model ’ s joint distribution satisfies F = EPx , y [ −∇2θ log px , y ( θ ) ] where px , y ( θ ) is the density function . If the model ’ s output conditional distribution matches the conditional distribution of the data , then the Fisher and Hessian matrices are in fact equivalent [ 27 ] . Roughly , this means that , if the model has high accuracy , we can approximate the Hessian of L at θ with the Fisher matrix . It is sometimes useful to consider an approximation to the Fisher , where the distribution Px , y is replaced with the empirical training distribution Q̂x , y , leading to the empirical Fisher matrix : F̂ =EQ̂x [ EQ̂y|x [ ∇ log py|x ( θ ) ∇ log py|x ( θ ) > ] ] = 1 N N∑ i=1 ∇ ` ( yi , f ( xi ; θ ) ) ︸ ︷︷ ︸ ∇ ` i ∇ ` ( yi , f ( xi ; θ ) ) > . The Woodbury-Sherman-Morrison Trick . The fact that this approximation is a sum of rank-one matrices has the benefit that it allows efficient exact computation of the Fisher inverse . Specifically , one can apply the classic Woodbury-Sherman-Morrison formula to compute the inverse recursively as F̂−1 = F̂−1N = F̂ −1 N−1 − F̂−1N−1∇ ` N ( ∇ ` > N F̂ −1 N−1 ) N +∇ ` > N F̂ −1 N−1∇ ` N , ( 2 ) where the recursion is over samples , and the base step is F̂−10 = λ −1Id , with λ being a small positive constant . This approach was independently introduced by Hassibi and Stork [ 13 ] for pruning , and Amari [ 3 ] for natural gradient . Singh and Alistarh [ 37 ] showed that it can be scaled to DNNs by blockwise approximation , and that it provides better approximations of the loss than K-FAC-based [ 44 ] , or diagonal [ 39 ] approximations [ 7 ] , leading to more accurate pruning . The main shortcoming of directly applying ( 2 ) is the prohibitive computational cost of Ω ( d2m ) , even when reduced to Ω ( dBm ) by B-block-wise approximation . Our method proposes a matrix-free approach for exactly calculating IHVPs and querying individual elements of the inverse Hessian , of cost O ( dm ) after O ( dm2 ) precomputation . Our method is numerically equivalent to the direct computation above . Diagonal Approximations . A common approximation , both for optimization , e.g . [ 20 , 18 ] but also for pruning [ 39 ] , is to assume that the Fisher matrix is diagonal . In our notation , this method has setup cost O ( dm ) , and diagonal query cost O ( d ) . However , as evident from the experimental results of [ 37 ] this provides lower approximation quality relative to both block-wise methods or K-FAC [ 44 ] . K-FAC Approximations . This approach observes that the entries of the true Fisher corresponding to blocks “ between ” two layers , which can be written as the expectation of a Kronecker product between two matrices , can be approximated as the Kronecker product of the expectations of those two matrices ( reversing the order between the product and the expectation ) . This approximation has been leveraged for both pruning and for optimization , and it allows efficient computation of the inverse [ 5 , 33 , 49 , 44 , 23 ] ; however , it is known to not always hold [ 31 ] . Another relative drawback is that the Kronecker factorization only occurs naturally for fully-connected layers ; but there is work on extensions to other layer types [ 11 , 30 ] , via additional approximations . We compare against K-FAC-based optimizers and pruners , and find that our method yields better results . Additional Approaches . Our approach is similar in spirit to matrix-free methods [ 32 , 25 , 28 ] , but the algorithms we present are new . Hessian-free optimization [ 28 ] also forgoes the explicit computation of Hessians in favor of computing an IHVP with a vector v. However , this estimation is performed by iteratively approximating the solution to the linear system Hx = v for some given x without ever explicitly forming H. One disadvantage of this method in practice is that it requires tuning and several very costly iterations to converge ( for a single v ) , as the underlying Hessian can be ill-conditioned . The L-OBS pruning method [ 7 ] approximates second-order information by defining independent layer-wise objectives , which allows the direct approximation of the layer-wise Hessians via a carefully-crafted block structure . In contrast , our approach allows for fully-global Hessian estimation and it also yields better pruning results at scale . Full-matrix adaptive regularization has similar goals , but in the context of adaptive optimizers [ 8 ] . Agarwal et al . [ 2 ] proposed GGT , which allows the efficient computation of the inverse square root of the low-rank matrix resulting from the sum of gradient outer products over a sliding window . At its core , this procedure requires an eigen-decomposition ( implemented via SVD ) of an m×m matrix at every time step , which is reasonably efficient for small values of m. Our dynamic algorithm solves a similar , but slightly simpler problem , as we only want to invert the matrix , without computing its square root . We do not perform any eigen-decompositions ; instead , we carefully maintain intermediate information that allows an efficient explicit computation of the scalar coefficients in Equation 6 . At the end of Section 3.2 , we discuss how this approach is more efficient in practice and allows executing at larger window sizes ( with small overhead ) , which leads to improved model accuracy as we show in our experiments . Additionally , our dynamic algorithm has per-step cost O ( dm + m3 ) , relative to O ( dm2 + m3 ) reported by GGT [ 2 ] ; this can result in lower overhead versus GGT , even when the SVD cost is negligible ( e.g . if m is small ) . Yao et al . [ 47 ] recently provided an alternative method for approximating the diagonal of the inverse Hessian , using Hutchinson ’ s randomized algorithm for estimating the diagonal . To mitigate the variance , the authors introduce non-trivial smoothing heuristics . Their approach has theoretical per-iteration cost of at least 2× versus SGD , as it requires a second backward pass over the network ; and this cost is usually higher in practice . Experimentally , our algorithm often matches their accuracy , and , unlike AdaHessian , its cost depends exclusively on the model size , irrespective of the underlying structure . Thus , as we show in the experimental section , our average practical overhead is less than 50 % for dense models , and less than 10 % for sparse ones . Approximation Quality . Künstner et al . [ 21 ] performed an in-depth analysis of the empirical Fisher , making the point that , in theory , the approximation could become meaningless if all sample gradients are zero . Similarly to [ 37 ] , we did not find that this occurs in practice for deep neural networks , as sample gradients are never zero . The latter reference provides detailed comparisons of diagonal , K-FAC , and other approximations , and finds that the empirical Fisher can provide competitive approximation quality for DNNs . Further , they demonstrate that better loss approximation implies better accuracy for neural network pruning . We therefore do not repeat their loss analysis , and mainly compare methods in terms of application performance .

Summary: In this work, the authors propose efficient techniques for computing inverse of a matrix that is represented as a sum of rank-1 matrices. The basic idea is to use recursively well-known Sherman-Morrison formula for computing inverse of rank-1 update. The authors show that their technique can be applied for computing inverse of the empirical estimate of the Fisher Matrix and thus can be used in optimization algorithms and for so-called prunning of a trained model.