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: This paper concentrates on the fast implementation of computing IHVPs (Inverse-Hessian Vector Products, and the expression here means "Inverse-Empirical-Fisher Vector Products") in training DNN, whose high time and memory costs remain a core problem. This paper is mainly based on the WoodFisher method. The WoodFisher uses the SMW (Sherman-Morrison-Woodbury) formula to compute the inverse of Fisher. This paper proposes a new algorithm (M-FAC) to accelerate this process. With two algorithms (static and dynamic, which correspond to gradients from single and multiple epochs), M-FAC can avoid matrices operations and decrease the time and storage costs, and the final time complexity is linear-in-dimension. In real practice, M-FAC outperforms K-FAC and has a similar performance with AdaHessian. Besides, when pruning DNNs, M-FAC also shows an improvement compared with WoodFisher.