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

Abstract:
Efficiently approximating local curvature information of the loss function is a key tool for optimization and compression of deep neural networks. Yet, most existing methods to approximate second-order information have high computational or storage costs, which limits their practicality. In this work, we investigate matrix-free, linear-time approaches for estimating Inverse-Hessian Vector Products (IHVPs) for the case when the Hessian can be approximated as a sum of rank-one matrices, as in the classic approximation of the Hessian by the empirical Fisher matrix. We propose two new algorithms: the first is tailored towards network compression and can compute the IHVP for dimension d, if the Hessian is given as a sum of m rank-one matrices, using O(dm) precomputation, O(dm) cost for computing the IHVP, and query cost O(m) for any single element of the inverse Hessian. The second algorithm targets an optimization setting, where we wish to compute the product between the inverse Hessian, estimated over a sliding window of optimization steps, and a given gradient direction, as required for preconditioned SGD. We give an algorithm with cost O(dm+m) for computing the IHVP and O(dm+m) for adding or removing any gradient from the sliding window. These two algorithms yield state-of-the-art results for network pruning and optimization with lower computational overhead relative to existing second-order methods. Implementations are available at [9] and [17].

heading: 1 Introduction
text: 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.

heading: 2 Preliminaries and Related Work
text: 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.

heading: 3 Algorithm Descriptions

heading: 3.1 The Static Algorithm
text: As a warm-up, we first describe the IHVP algorithm given a static set of gradients. While our notation is customized for the empirical Fisher, all our techniques are applicable to any matrix that can be written as the sum of m rank-one components. Specifically, we are given m vectors, which we assume to be the gradient vectors (∇`i)mi=1, and must compute quantities related to the inverse of the matrix resulting from the sum (more precisely, the average) of their outer products, which we denote by F̂−1m . The main idea is to rewrite the recursive description of F̂ −1 m such that, after some precomputation, the matrix-vector-product F̂−1m x can be computed efficiently for any vector x. Using
the Sherman-Morrison formula applied to the partial empirical Fisher matrix F̂i corresponding to the first i gradients (and scaled by 1/m), we obtain the following recursion:
F̂−1i x = F̂ −1 i−1x−
F̂−1i−1∇`i(F̂ −1 i−1∇`i)>
m+∇`>i F̂ −1 i−1∇`i
x, with F̂−10 x = λ −1Idx = λ −1x and λ > 0. (3)
For simplicity, let us set vi = F̂−1i−1∇`i and unroll the Sherman-Morrison recursion, which gives:
F̂−1i x = F̂ −1 i−1x− vi
vi >x
m+∇`>i vi = λ−1x− i∑ j=1 vj vj >x m+∇`>j vj . (4)
IHVP Computation. Assuming that we have already computed all the vectors vj, the above expression can be calculated in O(dm) time without any intermediate d× d matrices: first compute the scalar fractions in (4), and then evaluate the resulting linear combination of x and vj. Since vi = F̂ −1 i−1∇`i, it can be computed in exactly the same way given all vj for j < i. Thus, all m vectors vi of dimension d can be precomputed in increasing order of i using O(dm2) total time. As an additional optimization (to cut the required memory in half), we also precompute all qj = m+∇`>j vj.
Querying the Inverse. To extract individual elements [F̂−1m ]ij of the inverse Hessian approximation, we can write ei>F̂−1m ej in the form of (4) where ei and ej are indicator vectors
ei >F̂−1m ej = ei >λ−1ej − m∑ k=1 ei >vk vk >ej qk . (5)
As ei>vk and vk>ej can both be realized as constant time indexing operations, the above turns into a sum over m scalars. Hence, our method admits O(m) access to any element of F̂−1m using the same precomputed vk and qk as for the efficient calculation of F̂−1m x. Additional Optimizations. The algorithm admits a fast vectorized implementation, and several optimizations, which we describe in the Appendix. For example, we perform several memory-saving optimizations, as well as explicit page swapping between CPU and GPU memory to mitigate the gradient transfer costs. Furthermore, the static algorithm can be applied independently to each block of a block-wise approximation: for block size B, the computation and memory costs per block are reduced by B/d, but since there are now d/B blocks, the overall costs will stay the same irrespective of B. Thus, as long as m < B, our method is B/m times faster (and less memory intense) than the direct implementation of the Woodbury inverse.

heading: 3.2 The Dynamic Algorithm
text: We now describe the dynamic algorithm, which assumes that gradients arrive in an online fashion and must be integrated into the Fisher estimate. We first present the algorithm itself, i.e. the setup / update / IHVP computations, and perform a complexity analysis. Next, we show how the algorithm can be derived and finally we conclude with notes on an efficient practical implementation. High-Level View. The main idea of the dynamic algorithm is to write the IHVP as
F̂−1m x = λ −1x− m∑ j=1 cmj ∇`j , (6)
where the scalar coefficients cmj can be computed efficiently from just the scalar products ∇`>i ∇`j and ∇`>i x. Then, any gradient ∇`i can be replaced by updating just O(m) of the stored scalar product values. In the following, we use G = [∇`>1 ;∇`>2 ; . . . ;∇`>m] to denote the row-wise m× d matrix of the m gradients for which we wish to compute the inverse empirical Fisher. Further, for each i from 1 to m, let bi be a vector with m components (bi1, b i 2, . . . , b i m), such that
F̂−1i−1∇`i = λ −1∇`i − i−1∑ j=1 ai−1j ∇`j = m∑ j=1 bij∇`j , (7)
i.e. containing the scalar coefficients of F̂−1i−1∇`i. This means that bij = −a i−1 j when j < i, b i i = λ −1, and bij = 0 otherwise. Initial Precomputation & Update. The dynamic algorithm maintains three m × m matrices: GG>, D and B. The first is the symmetric gradient scalar product matrix [GG>]ij = ∇`>i ∇`j .
The second is an upper triangular matrix and stores the precomputed values [D]ij = ∇`>i F̂ −1 i−1∇`j for i ≤ j. The third is the row-wise matrix of the coefficient vectors [B]ij = bij , which makes it lower triangular with a diagonal of λ−1. We now discuss how to compute those matrices. The initial setup of the dynamic algorithm begins by evaluating GG> in a straightforward fashion. Next, D can be computed for i ≤ j according to the following recursion:
[D]ij = (∇`>i F̂−1i−1∇`j) and (∇` > j F̂ −1 0 ∇`k) = λ−1[GG>]jk (8)
(∇`>j F̂−1i ∇`k) = (∇` > j F̂ −1 i−1∇`k)−
(∇`>j F̂ −1 i−1∇`i)(∇`>i F̂ −1 i−1∇`k)
m+ (∇`>i F̂ −1 i−1∇`i)
. (9)
Given D, we can then conclude the precomputation by calculating B for i ≥ j recursively as:
[B]ii = λ −1 and [B]ij = − i−1∑ k=j [D]ki m+ [D]kk [B]kj . (10)
After the initial setup, gradient ∇`k can be replaced with gradient ∇`′k by first updating row k of G and then replacing row and column k in GG> with G∇`′k. Afterwards, the recomputation of columns j ≥ k of D and rows i ≥ k of B completes the update. Multiplication. Once GG>, D and B have been precomputed, one can perform efficient IHVPs of the form F̂−1m x with arbitrary vectors x. This is done by first evaluating p = Gx and then computing all m values qi by the following recursion:
qi = (∇`>i F̂ −1 i−1x)
m+ [D]ii and (∇`>j F̂−10 x) = λ−1pj (11)
(∇`>j F̂−1i x) = (∇` > j F̂ −1 i−1x)− [D]ij m+ [D]ii (∇`>i F̂−1i−1x). (12)
Eventually, the final result of the IHVP is obtained by:
F̂−1m x = λ −1x− m∑ j=1 ( m∑ k=j qk[B]kj)∇`j . (13)
Complexity Analysis. The dynamic algorithm stores m gradients of dimension d as well as three m×m matrices, and thus has an overall memory complexity of O(dm+m2). Next, we analyze the time complexity of all important operations. Initially, GG> must be computed once, which takes O(dm2) time. Then, the recursion of D has three indices with 1 ≤ i, j, k ≤ m and each step takes constant time. Thus, it can be computed in O(m3) time with dynamic programming. Further, since values for index i depend only on values for index i− 1, the dynamic programming can be implemented in O(m2) space. B has two indices 1 ≤ i, j ≤ m and every recursion takes O(m) time, meaning that it can also be fully computed in O(m3) through dynamic programming. Hence, the overall initial setup cost is O(dm2 +m3). To replace one gradient with∇`′, we have to compute G∇`′ as well as (partially) recalculate D and B, which takes at worst O(dm+m3) time. An IHVP requires two matrix-vector products involving G and a recursion with two indices and therefore has a complexity of O(dm+m2). Algorithmic Derivation. The dynamic algorithm can be derived directly from Theorem 1, which we state here in a simplified form (using the definitions of B and D), and prove in the Appendix.
Theorem 1. Let F̂i = λI + 1/m · ∑i j=1∇`j∇`>j , then F̂ −1 i x can be calculated as:
F̂−1i x = λ −1x− i∑ j=1 cij∇`j and cij = i∑ k=j (∇`>k F̂ −1 k−1x) m+ [D]kk [B]kj . (14)
Equation (14) with i = m is exactly equal to the IHVP computation (13) as the fraction in the innermost sum corresponds to qk. Similarly, index shifting i = i − 1 and setting x = ∇`i, thus turning cij into −[B]ij as well as the fraction’s numerator to [D]ki, recovers the precomputation formula of B given by (10) for i > j. The formulas for D and qi follow directly from an expansion of the Woodbury formula followed by an appropriate recursive evaluation that avoids any matrix / vector operations except in the base case (we indicate the corresponding recursive calls by brackets).
Efficient Implementation. Directly implementing the discussed recursive formulas in a modern machine learning framework would result in very slow code. Fortunately, it is possible to implement all computations required for the dynamic algorithm very efficiently on a GPU. We describe how to do this in the Supplementary Material and provide a full implementation [9]. Specifically, we provide complete sample code that is able to perform the calculation of D and B, i.e. the O(m3) component of the overall update cost, in < 10 milliseconds (on an NVIDIA RTX 2080 Ti) for values of m as high as 1024 (and the code can still be further optimized). For reference, this is > 10× faster than the highly-optimized m×m SVD computation done at every step by GGT [2]. We emphasize that this O(m3) computation being very fast in practice is crucial to reach low overheads, especially when dealing with a models where the O(dm) matrix-vector products are not the bottleneck.

heading: 4 Experimental Validation

heading: 4.1 Application 1: Pruning DNNs using the Static Algorithm
text: Background. Given a trained model θ?, the goal of pruning is to find the weight θi, or the set of weights, whose setting to 0 would lead to a minimal increase in the training loss. Under a local quadratic approximation of the loss function, the OBD framework [24] shows that the “optimal” weight to be removed is the one with the lowest value of the saliency metric ρ(θi) =
θ2i 2 [F̂−1]ii and
proposed to estimate [F̂−1]ii by diagonal approximation. The OBS framework [13] observed that the remaining weights θ should also be updated via the optimal perturbation δθ = −θiF̂−1ei/[F̂−1]ii, where ei is the ith basis vector. Our algorithm efficiently supports both these operations. Wang et al. [44] raised the valid point that applying the OBS update above to multiple weights being removed at once may be incorrect, since it ignores possible correlations between those weights. We considered this point in detail, comparing results between OBD pruning [24], the OBS update [13], and an augmented version of OBS which disentangles the correlations by solving a corresponding linear system [37]. The results, presented in the Appendix, suggest that the OBS update is quite beneficial even in its approximate form, and that the effect of correlation is small for unstructured pruning. In fact, we find that the OBS pruning mask in each step is usually very similar to the one obtained by simple magnitude pruning, suggesting that the final accuracy improvements are primarily due to the updates of the remaining weights, facilitating better recovery. Experimental Setup. We prune CNNs (ResNet-50 [15] and MobileNet-V1 [16]) on the ImageNet dataset [36]. These models are standard in the pruning literature [10], and therefore several strong baselines exist. Timing experiments are run on a machine with NVIDIA RTX 2080 Ti GPUs, a 48-core Intel CPU, and 512 GB of RAM. Following [37], we used batched gradients (of size 16) as single samples inside the Fisher approximation. This does not alter results, but reduces variance. We compare against Global Magnitude Pruning (GMP) [12], Layer-wise Optimal Brain Surgeon (L-OBS) [7], Soft Threshold Reparametrization (STR) [22], and WoodFisher (WF) [37]. The latter two methods are state-of-the-art for gradual pruning. WF is numerically equivalent to our method at the same parameter settings, but has significantly higher computational and storage cost. Relative to it, we therefore focus on executing at higher parameter values. We compare against their public implementation. The Appendix contains full hyper-parameters, ablations with respect to block size and number of gradients, and a comparison with K-FAC pruning in a simpler setting [49]. We first perform a one-shot comparison to evaluate the “raw” per step pruning performance of M-FAC relative to other methods. Next, we evaluate how simply increasing M-FAC parameter values (but keeping everything else exactly the same as in WF’s experiments [37] which follow [50]) improves over state-of-the-art gradual pruning results by WF, GMP and STR. Finally, we demonstrate how M-FAC’s high per-step efficiency can be utilized to craft practical pruning schedules with little computational overhead relative to GMP and only moderate extra memory consumption, which at the same time yield significantly better results than WF’s state-of-the-art numbers. Oneshot Pruning. Figure 1a shows the Top-5 accuracy of one-shot ResNet50 pruned models for L-OBS, GMP, WoodFisher, and M-FAC, where our method is executed with a parameter setting that is infeasible for WoodFisher due to the required computational and storage costs. (We use Top-5 accuracy to compare with [7].) We note the improved accuracy of global methods relative to layer-wise, and that the Fisher approximation yields consistently better results relative to global magnitude. This suggests that a better approximation of the Fisher inverse also yields better pruning results, and is in line with the loss approximation study of [37].
Generally, we found that estimation using more gradients always improves results, until saturation. Interestingly, increasing the block size does not necessarily improve results, for some models, smaller block sizes sometimes yield better results for a fixed number of gradients. Figure 1b examines the effect of the improvement in one-shot pruning accuracy at each step of pruning, on the final result, when pruning MobileNetV1-STR gradually to 89% sparsity. WoodFisher uses block size 10K, and m = 400 gradients, while M-FAC uses the same block size but m = 4000 gradients. Executing WoodFisher with the same parameters would be extremely slow. Note the gap in one-shot accuracy following each pruning step: this translates in the final accuracy gap of more than 1% even after the extensive fine-tuning phase. Gradual Pruning Comparisons. Table 1 presents gradual pruning results for MobileNetV1STR/ImageNet at 89% sparsity and ResNet50/ImageNet at 95% sparsity, relative to other state-of-theart methods. M-FAC outperforms previous methods in the case of MobileNetV1-STR by more than 1.5% Top-1 accuracy. Further, despite using 10× more gradients to estimate the empirical Fisher, the per-step pruning cost of M-FAC is > 6× lower than WoodFisher’s as can be seen in Table 2 (for the same parameters settings M-FAC is > 100× faster). The ResNet50 results show smaller gains.
Practical Pruning. Due to the computational cost of a WF step, the previous schedules used only a small number of pruning steps. However, M-FAC makes it possible to prune much more frequently. In this case, the cost of extracting and storing several thousand gradients may still be a bottleneck.
We find that M-FAC already performs well with low parameter settings, especially when coupled with recomputations (i.e. resampling gradients and rebuilding the M-FAC approximation after partial pruning steps). We now perform MobileNet and ResNet50 gradual pruning runs using M-FAC with block size B = 128, 64 gradients and 16 recomputations per step. A single M-FAC pass takes less than a few seconds at these settings. Details about the pruning schedule can be found in the Appendix. This is a practical setup, which can scale to larger models. At the same time, Table 3 shows that leads to significantly improved pruning results, sometimes even by several percent accuracy, compared to WF and the M-FAC results in Table 1.

heading: 4.2 Application 2: Optimization using the Dynamic Algorithm
text: Matrix-Free Preconditioned SGD. The dynamic algorithm can provide an efficient implementation for the following variant of SGD. Let t be the current iteration index, and ∆ < t be the length of the sliding window during which the estimation of the inverse Hessian is performed. Consider the iteration θt+1 = θt − ηtF̂−1[t−∆,t]∇`t, where ηt is the learning rate. Let ∇`t−τ+1, . . . ,∇`t be the gradients obtained at steps (t−τ+1), . . . , t, respectively. Then, we define the preconditioning matrix as F̂[t−∆,t] = λ−1Id + 1/∆ · ∑∆−1 τ=0 ∇`t−τ · ∇`>t−τ , that is, an empirical Fisher matrix generated with respect to the sliding window of ∆ gradients. At each step, our dynamic algorithm replaces ∇`t−∆+1 with ∇`t in the inverse estimate, and then computes F̂−1[t−∆,t]∇`t. This corresponds to full-matrix natural gradient SGD [3] under the extra assumption that gradients do not change too quickly during training. This assumption has been validated in the K-FAC literature [5] and we also provide additional experimental justification in the Appendix. Finally, to better study the effect of the preconditionier, we do not apply momentum to M-FAC in any of the experiments below. Initial Test. We begin by verifying that the M-FAC optimizer actually behaves as one would expect from a second-order method, i.e. learning faster than standard first-order techniques, especially in the early epochs. For that purpose, we run M-FAC, SGD with and without momentum and Adam on WideResNet 40-2 / CIFAR10 with a compressed 50 epoch schedule. The corresponding Figures 1c and 1d indeed showcase the expected behavior. Interestingly, they also show that SGD without momentum performs poorly in this setup, which further highlights the major impact of the M-FAC preconditioner, which seems to be able to make up for the lack of momentum. Second Order Comparison. Next, we work with the common ResNet20 and ResNet32 models [15] and compare against standard optimizers such as SGD with momentum, Adam [18] and AdamW [26], but also against approximate second-order methods such as K-FAC [30], GGT [2], and AdaHessian [47]. For the former two, we use the implementations contributed to TensorFlow [1] by the authors, while for the latter we use the authors’ PyTorch implementation. For fairness, we always follow the settings recommended by the authors, although this makes overheads harder to compare: we used specific learning rates, and grid-searched the weight-decay values for each method in turn. Since implementations have slightly different framework requirements, our main performance metric is overhead over SGD, measured in each algorithm’s environment, on an NVIDIA RTX 3090 GPU for PyTorch and a Titan RTX for TensorFlow. For all second-order methods, we compute the preconditioner at each step—reducing the frequency reduces overheads proportionally for each method, but introduces an additional hyper-parameter and generally results in lower final accuracy. We provide additional information in the Supplementary. Our first results are presented in Table 4 (left), where we examine the best Top-1 test accuracy obtained by each method, as well as its overhead relative to SGD. For K-FAC, the best accuracy is obtained with batch size 1000, while all other methods operate on batch size 128. Similarly, GGT’s best results are achieved with window size m = 100, while for M-FAC larger window sizes improve results; we use m = 512. Due to these discrepancies, we measure overheads with the “best accuracy” setting for each method but also in the same setup as M-FAC; the latter is given in brackets.
In terms of accuracy, M-FAC scores highest in the ResNet20 experiment, and comes third, to tuned SGD and AdaHessian (performing almost the same), on ResNet32. The overhead of M-FAC is around 50%, which is lower than AdaHessian, on par with K-FAC, but higher than GGT. However, these overheads are comparable only when GGT uses 5× less gradients (surprisingly, larger window sizes yield worse results for GGT). On the same settings, M-FAC provides up to 5x lower overhead. This is due to (a) better performance of our dynamic algorithm relative to SVD and (b) the Θ(dm2) term in GGT, versus our O(dm) term. Image Classification. Table 4 (right) provides a scenario where we run SGD with momentum, Adam, and M-FAC without tuning or weight-decay on Wide Residual Networks (WRN) [48] for CIFAR-100, as well as on MobileNetV1/ ImageNet, and examine test accuracy. M-FAC achieves the highest accuracy on most runs, even on ImageNet. We emphasize that these results can probably be improved by parameter tuning. Another interesting application is fine-tuning sparse models. We fine-tune 90% sparse MobileNet / ResNet50 models for 30 epochs starting from the last gradual pruning step and find that M-FAC reaches higher final accuracy than SGD in both cases. At the same time, it is only negligibly slower (≤ 5% overhead) since algorithm’s complexity (compute and memory) is linear in the model density. At last, we note that M-FAC achieving higher test accuracy is generally well correlated with a lower final training loss. Language Modelling using Transformers. Finally, we test the M-FAC optimizer for smaller Transformer models. We use the default values m = 1024 gradients (unless the dataset has less than this number of total samples), dampening λ = 10−6, learning rate 10−4 and no weight decay. We then train BERT [6] tiny and mini models [41] on SQuADv2 [35] and the GLUE benchmark suite [43], comparing against HuggingFace’s [45] Adam baseline. For M-FAC, we use exactly the same training setup and only replace the optimizer. The results in Table 5 show that M-FAC performs better than the Adam baseline on almost all tasks, on several even by considerable margins. We provide more detailed results in the Appendix, including hyperparameters and a comparison with AdamW.

heading: 5 Discussion
text: We presented static and dynamic algorithms for computing IHVPs when the Hessian matrix can be approximated by a sum of rank-1 matrices. We used the classic empirical Fisher approximation, but our results can apply more broadly. The main limitation is the cost of storing m additional gradients. For the static algorithm, we can efficiently leverage system memory, as described in the Appendix. The dynamic algorithm could parallelize gradient storage, or perform gradient compression [38]. We plan to investigate this in future work and perform a larger-scale study of optimization performance.

heading: 6 Acknowledgements
text: We gratefully acknowledge funding the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement No 805223 ScaleML), as well as computational support from Amazon Web Services (AWS) EC2.

Review summaries:
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. 

The authors propose a compute-/memory- efficient implementation of inverse empirical Fisher vector products (IEFVPs) for deep neural networks, called M-FAC, based on the Woodbury-Sherman-Morrison inversion formula. They use IEFVPs as estimations of inverse Hessian vector products (IHVPs) used in model compression (e.g., Optimal Brain Damage) and second-order optimization. M-FAC is similar to the previous method, WoodFisher (Singh and Alistarh, 2020), which uses the WSM inversion formula to calculate inverse empirical Fisher. Unlike WoodFisher (calculates matrices), M-FAC directly calculates IEFVPs (vectors) and thus requires less compute/memory cost. This allows M-FAC to calculates IEFVPs with the full empirical Fisher and a larger number of samples ($m$) than WoodFisher.  - The static algorithm (with a fixed set of samples) of M-FAC is tested on the one-shot/gradual pruning with trained CNNs on ImageNet. It shows a better accuracy-sparsity trade-off than WoodFisher and other magnitude-based methods.  - The dynamic algorithm (with the fixed size sliding window of samples) of M-FAC is tested on training CNNs on CIFAR-10 and fine-tuning sparse CNNs on CIFAR-100 and ImageNet with preconditioned SGD. It shows a better computational overhead (vs. SGD) than existing second-order optimization methods in training and better accuracy than first-order methods in fine-tuning.

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.