# Efficient Matrix-Free Approximations of Second-Order Information, with
Applications to Pruning and Optimization

Elias Frantar
IST Austria
<EMAIL_ADDRESS>
&Eldar Kurtic
IST Austria
<EMAIL_ADDRESS>
&Dan Alistarh
IST Austria
<EMAIL_ADDRESS>
###### 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 can limit 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 as
part of a framework called M-FAC: the first algorithm 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^{2})$ 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^{2})$ for computing the IHVP and $O(dm+m^{3})$ 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 [10] and [18].

##  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 [15] would occupy 2.5 Petabytes of memory. These constraints
have inspired significant work on efficient numerical approximations of the
Hessian for deep neural networks, for example the line of work on the K-FAC
approximation [20, 33, 5, 48], or efficient block-wise approximations [46,
39]. One of the classic approaches, which we focus on in this paper, is the
_empirical Fisher_ [14, 3, 4] approximation to the Hessian, written as:

$\mathbf{H}\simeq\widehat{\mathbf{F}}=\frac{1}{N}\sum_{i=1}^{N}\nabla\ell_{i}\cdot\nabla\ell_{i}^{\top},$
(1)

where $N$ is the number of samples, $\nabla\ell_{i}$ denotes the gradient
w.r.t. the $i$th sample at the given point, and $\cdot$ is the outer product
of individual gradients, which we view as column vectors.

The empirical Fisher approximation is fairly standard, e.g. [14, 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 on the conditions for its applicability [31, 23, 39, 42].

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 [45] 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 [14, 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 [39], 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 competitive pruning results in
terms of accuracy versus average sparsity level. Yet, this approach is limited
by the block-wise approximation: for block size $B$ and dimension $d$, it
requires $\Theta(Bdm)$ time to recursively build the block-wise Fisher
approximation using $m$ gradients, and $\Theta(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 linear-time algorithms are
possible for computing IHVPs in this context.

Contribution. We resolve 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. Concretely, we provide
exact _matrix-free_ algorithms to compute products of the form
$\widehat{\mathbf{F}}^{-1}{\mathbf{v}},$ where $\widehat{\mathbf{F}}^{-1}$ is
the inverse empirical Fisher, and ${\mathbf{v}}$ is an arbitrary vector. We
show that these algorithms can be implemented efficiently at scale, and that
they can lead to more accurate 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
$\theta^{\star}$, 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 [26, 14]. For this, we first
compute $m$ gradients at $\theta^{\star}$, 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 (i.e. IHVPs), we can rewrite the Woodbury recursion such that
we work _exclusively with individual vectors and scalars_ , and never with
full $d\times d$ matrices. Given model dimension $d$ and $m$ gradients, each
defining a rank-one component of the empirical Fisher, the algorithm uses
$O(dm^{2})$ 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.

The Dynamic Algorithm. Next, we adapt this approach to SGD optimization, i.e.
to precondition stochastic gradients by our estimate of the inverse Hessian.
In this case, we use 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 $\Delta\geq 1$
optimization steps. This requires a _dynamic_ representation, allowing
addition and removal of gradients without full recomputation of second-order
statistics.

Our main technical contribution is showing that this can be achieved, with
approximately $O(dm)$ time and space complexity. The key idea is that, for any
set $(\nabla\ell_{j})_{j=1}^{m}$ of $m$ gradients, and any vector
${\mathbf{v}}$, it is possible to represent the corresponding IHVP estimate
$\widehat{\mathbf{F}}^{-1}_{m}{\mathbf{v}}$ as a linear combination of terms
corresponding to individual gradients $\nabla\ell_{j}$ and ${\mathbf{v}}$, of
the form
$\widehat{\mathbf{F}}^{-1}_{m}{\mathbf{v}}=\lambda^{-1}{\mathbf{v}}-\sum^{m}_{j=1}c^{m}_{j}\nabla\ell_{j},\textnormal{
where $\lambda>0$ is a dampening constant.}$

Crucially, we ensure that the coefficients $c^{m}_{j}$ can be computed just
via dot products $\nabla\ell_{j}^{\top}\nabla\ell_{i}$, where $i\leq j$ in the
ordering, and $\nabla\ell_{j}^{\top}{\mathbf{v}}$. Then, to replace a given
gradient $\nabla\ell_{i}$ from this representation, we just have to compute
$m$ scalar products with the new gradient vector, and to update some
intermediate information. Hence, the entire update operation has computational
cost $O(dm+m^{3})$ for adding or removing any gradient from the sliding
window, and $O(dm+m^{2})$ for computing the IHVP.

Implementation and Experiments. We provide an efficient vectorized
implementation for the above algorithms, called _M-FAC_ (for _Matrix-Free
Approximate Curvature_) as well as pruning and optimization code, in PyTorch
[36]. Our implementation includes several additional optimizations, in
particular GPU acceleration with custom CUDA code, and minimizes the cost of
memory transfer between the GPU and main memory via a non-trivial gradient
paging mechanism. For pruning operations, our implementation provides order-
of-magnitude improvements in terms of computation time over the block-wise
approximation of [39] for classic pruning benchmarks such as ResNet50 and
MobileNet on the ImageNet dataset. In turn, this allows us to obtain more
accurate pruned models by exploiting higher numbers of gradients and larger
block sizes. What is more, our preconditioned SGD can be competitive with
state-of-the-art optimizers for deep learning in terms of accuracy, and has
end-to-end computational overhead per step of 5%–55% relative to vanilla SGD.

##  2 Preliminaries and Related Work

General Definitions. We now briefly introduce the setting and notation; we
refer the reader to standard texts [31, 29, 4] for a complete introduction. We
start from the standard setting in which we are given a dataset
$\mathcal{D}=(\mathbf{x}_{i},\mathbf{y}_{i})_{i=1}^{N}$, and wish to identify
a $d$-dimensional model ${\mathbf{\theta}}\in\mathbb{R}^{d}$ to minimize an
empirical loss $L:\mathbb{R}^{d}\rightarrow\mathbb{R}$, defined as
$L(\mathbf{\theta})=\frac{1}{N}\sum_{i=1}^{N}\ell\big{(}\mathbf{y}_{i},f\left(\mathbf{x}_{i};\mathbf{\theta}\right)\big{)}$.
For a twice-differentiable loss $L$, the Hessian is the matrix of second-order
derivatives of $L$ w.r.t. ${\mathbf{\theta}}$, i.e.
$\mathbf{H}=\nabla^{2}_{\mathbf{\theta}}L$. In the _probabilistic_ view, each
input example $\mathbf{x}$ has some probability of being assigned a given
label $\mathbf{y}$. Given input examples ${\mathbf{x}}$ drawn from a
distribution $Q_{{\mathbf{x}}}$ and corresponding outputs drawn from a
conditional distribution $Q_{{\mathbf{y}}|{\mathbf{x}}}$, the goal is to
minimize the distance between the target joint distribution
$Q_{{\mathbf{x}},{\mathbf{y}}}=Q_{{\mathbf{x}}}\,Q_{{\mathbf{y}}|{\mathbf{x}}}$,
and a learned joint distribution
$P_{{\mathbf{x}},{\mathbf{y}}}({\mathbf{\theta}})$, where ${\mathbf{\theta}}$
is the model.

The Fisher Matrix. Assuming the probabilistic view, it can be shown that the
Fisher information matrix $\mathbf{F}$ of the model’s joint distribution
satisfies
$\mathbf{F}=\mathrm{E}_{P_{{\mathbf{x}},{\mathbf{y}}}}\left[-\nabla^{2}_{\mathbf{\theta}}\log
p_{\mathbf{\theta}}({\mathbf{x}},{\mathbf{y}})\right]\,$. If the model’s
output conditional distribution matches the conditional distribution of the
data, then the Fisher and Hessian matrices are in fact equivalent [29].
Roughly, this can be interpreted as saying that, if the model has high
accuracy, one can approximate the Hessian of ${\mathbf{\theta}}$ with the
Fisher matrix. Further, it is sometimes useful to consider an approximation to
the Fisher, where the distribution $P_{{\mathbf{x}},{\mathbf{y}}}$ is replaced
with the empirical training distribution
$\widehat{Q}_{{\mathbf{x}},{\mathbf{y}}}$, leading to the _empirical Fisher
matrix_ $\widehat{\mathbf{F}}$:

$\displaystyle\widehat{\mathbf{F}}=$
$\displaystyle\mathrm{E}_{\widehat{Q}_{{\mathbf{x}}}}\left[\mathrm{E}_{\widehat{Q}_{{\mathbf{y}}|{\mathbf{x}}}}\left[\nabla\log
p_{{\mathbf{\theta}}}({\mathbf{y}}|{\mathbf{x}})\nabla\log
p_{{\mathbf{\theta}}}({\mathbf{y}}|{\mathbf{x}})^{\top}\right]\right]\stackrel{{\scriptstyle}}{{=}}\frac{1}{N}\sum_{i=1}^{N}\underbrace{\nabla\ell\left({\mathbf{y}}_{i},f\left({\mathbf{x}}_{i};{\mathbf{\theta}}\right)\right)}_{\nabla\ell_{i}}\nabla\ell\left({\mathbf{y}}_{i},f\left({\mathbf{x}}_{i};\theta\right)\right)^{\top}.$

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

$\widehat{\mathbf{F}}^{-1}=\widehat{F}_{N}^{-1}=\widehat{F}_{N-1}^{-1}-\frac{\widehat{F}_{N-1}^{-1}\nabla\ell_{N}(\nabla\ell_{N}^{\top}\widehat{F}_{N-1}^{-1})}{N+\nabla\ell_{N}^{\top}\widehat{F}_{N-1}^{-1}\nabla\ell_{N}},$
(2)

where the recursion is over samples, and the base step is
$\widehat{F}_{0}^{-1}=\lambda^{-1}I_{d}$, with $\lambda$ being a small
positive constant. This approach was independently introduced by Hassibi and
Stork [14] for pruning, and Amari [3] for natural gradient. Singh and Alistarh
[39] showed that it can be scaled to DNNs by block-wise approximation, and
that it provides better approximations of the loss than K-FAC-based [43], or
diagonal [41, 7] approximations, leading to more accurate pruning. The main
shortcoming of directly applying (2) is the prohibitive computational cost of
$\Omega(d^{2}m)$, even when reduced to $\Omega(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)$ and $O(m)$, respectively, after $O(dm^{2})$ precomputation. Our
method is numerically equivalent to the direct computation above.

Diagonal Approximations. A common approximation, both for optimization, e.g.
[21, 19] but also for pruning [41], 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 [39]
this provides lower approximation quality relative to both block-wise methods
or K-FAC [43].

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, 35, 48, 43, 25]; however, it
is known that this approximation does not always hold [33]. Another relative
drawback is that the Kronecker factorization only occurs naturally for fully-
connected layers; though, there is work on extensions to other layer types
[12, 32], 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 [34, 28, 30], but the algorithms we present are new. _Hessian-free_
optimization [30] also forgoes the explicit computation of Hessians in favor
of computing an IHVP with a vector ${\mathbf{v}}$. However, this estimation is
performed by iteratively approximating the solution to the linear system
$\mathbf{H}{\mathbf{x}}={\mathbf{v}}$ for some given ${\mathbf{v}}$ without
ever explicitly forming $\mathbf{H}$. One disadvantage of this method in
practice is that it requires tuning and several costly iterations to converge
(for a single $\mathbf{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 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 (the authors take an
SVD) of an $m\times 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 (rather than also computing its square root).
However, 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 considerably 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. Our dynamic
algorithm has per-step cost $O(dm+m^{3}),$ relative to $O(dm^{2}+m^{3})$ for
GGT [2]; this can also result in lower overhead versus GGT, even when the SVD
cost is negligible (e.g. if $m$ is small).

Yao et al. [46] 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 additionally
introduce non-trivial smoothing heuristics. Their approach, called AdaHessian,
has a theoretical per-iteration cost of at least $2\times$ versus SGD, as it
requires a second backward pass over the network; in practice, this cost is
usually higher. 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. [23] 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 [39],
we did not find that this occurs in practice for deep neural networks, as
sample gradients are never zero. The latter reference also 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.

##  3 Algorithm Description

###  3.1 The Static Algorithm

As a warm-up, we first describe the IHVP algorithm for 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
rank-one components. Specifically, we are given $m$ vectors, which we assume
to be the gradient vectors $(\nabla\ell_{i})_{i=1}^{m}$, 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
$\widehat{F}_{m}^{-1}$. The main idea is to rewrite the recursive description
of $\widehat{F}_{m}^{-1}$ such that, after some precomputation, the matrix-
vector-product $\widehat{F}_{m}^{-1}\mathbf{x}$ can be computed efficiently
for any vector $\mathbf{x}$. Using the Sherman-Morrison formula applied to the
partial empirical Fisher matrix $\widehat{F}_{i}$ corresponding to the first
$i$ gradients (and scaled by $1/m$), we obtain the following recursion:

$\displaystyle\widehat{F}_{i}^{-1}{\mathbf{x}}=\widehat{F}_{i-1}^{-1}{\mathbf{x}}-\frac{\widehat{F}_{i-1}^{-1}\nabla\ell_{i}\left(\widehat{F}_{i-1}^{-1}\nabla\ell_{i}\right)^{\top}}{m+\nabla\ell_{i}^{\top}\widehat{F}_{i-1}^{-1}\nabla\ell_{i}}{\mathbf{x}},\textnormal{
with
}\widehat{F}_{0}^{-1}{\mathbf{x}}=\lambda^{-1}I_{d}{\mathbf{x}}=\lambda^{-1}{\mathbf{x}}\textnormal{
and }\lambda>0.$ (3)

For simplicity, let us set
$\mathbf{v_{i}}=\widehat{F}_{i-1}^{-1}\nabla\ell_{i}$ and unroll the Sherman-
Morrison recursion, which gives:

$\displaystyle\widehat{F}_{i}^{-1}{\mathbf{x}}=\widehat{F}_{i-1}^{-1}{\mathbf{x}}-\mathbf{v_{i}}\frac{\mathbf{v_{i}}^{\top}{\mathbf{x}}}{m+\nabla\ell_{i}^{\top}\mathbf{v_{i}}}=\lambda^{-1}{\mathbf{x}}-\sum^{i}_{j=1}\mathbf{v_{j}}\frac{\mathbf{v_{j}}^{\top}{\mathbf{x}}}{m+\nabla\ell_{j}^{\top}\mathbf{v_{j}}}.$
(4)

IHVP Computation. Assuming that we have already computed all the vectors
$\mathbf{v_{j}}$, the above expression can be calculated in $O(dm)$ time
without any intermediate $d\times d$ matrices: first compute the scalar
fractions in (4), and then evaluate the resulting linear combination of
${\mathbf{x}}$ and the $\mathbf{v_{j}}$. Since
$\mathbf{v_{i}}=\widehat{F}_{i-1}^{-1}\nabla\ell_{i}$, it can be computed in
exactly the same way given all $\mathbf{v_{j}}$ for $j<i$. Thus, all $m$
vectors $\mathbf{v_{i}}$ of dimension $d$ can be precomputed in increasing
order of $i$ using $O(dm^{2})$ total time. We can halve the memory usage by
also precomputing all denominators
$q_{i}=m+\nabla\ell_{i}^{\top}\mathbf{v_{i}}$ as this allows us to overwrite
$\nabla\ell_{i}$ with ${\mathbf{v}}_{i}$.

Querying the Inverse. To extract individual elements
$[\widehat{F}_{m}^{-1}]_{ij}$ of the inverse Hessian approximation, we can
write $\mathbf{e_{i}}^{\top}\widehat{F}_{m}^{-1}\mathbf{e_{j}}$ in the form of
(4) where $\mathbf{e_{i}}$ and $\mathbf{e_{j}}$ are indicator vectors

$\mathbf{e_{i}}^{\top}\widehat{F}_{m}^{-1}\mathbf{e_{j}}^{\top}=\mathbf{e_{i}}^{\top}\lambda^{-1}\mathbf{e_{j}}-\sum^{m}_{k=1}\mathbf{e_{i}}^{\top}\mathbf{v_{k}}\frac{\mathbf{v_{k}}^{\top}\mathbf{e_{j}}}{q_{k}}.$
(5)

As $\mathbf{e_{i}}^{\top}\mathbf{v_{k}}$ and
$\mathbf{v_{k}}^{\top}\mathbf{e_{j}}$ 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 $\widehat{F}_{m}^{-1}$ using the
same precomputed $\mathbf{v_{k}}$ and $q_{k}$ as for the efficient calculation
of $\widehat{F}_{m}^{-1}{\mathbf{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
intelligent explicit page swapping between CPU and GPU memory to mitigate
memory transfer costs. Furthermore, the static algorithm can be applied
independently to each block of a block-wise approximation: for block size $B$,
the runtime is reduced to $d/B$. Then, as long as $m<B$ (i.e. number of
gradients smaller than the block size), our method is $B/m$ times faster than
the direct implementation of the Woodbury inverse.

###  3.2 The Dynamic Algorithm

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

$\widehat{F}_{m}^{-1}{\mathbf{x}}=\lambda^{-1}{\mathbf{x}}-\sum^{m}_{j=1}c^{m}_{j}\nabla\ell_{j},$
(6)

where the scalar coefficients $c^{m}_{j}$ can be computed efficiently just
from the scalar products $\nabla\ell_{i}^{\top}\nabla\ell_{j}$ and
$\nabla\ell_{i}^{\top}{\mathbf{x}}$. Then, any gradient $\nabla\ell_{i}$ can
be replaced by updating just $O(m)$ of the stored scalar product values. In
the following, we use
$\mathbf{G}=[\nabla\ell_{1}^{\top};\nabla\ell_{2}^{\top};\dots;\nabla\ell_{m}^{\top}]$
to denote the row-wise $m\times d$ matrix of the $m$ gradients based on which
we wish to compute products between the inverse empirical Fisher and an
arbitrary vector. Further, for each $i$ from 1 to $m$, let $\mathbf{b^{i}}$ be
a vector with $m$ components $(b^{i}_{1},b^{i}_{2},\ldots,b^{i}_{m})$, such
that

$\widehat{F}_{i-1}^{-1}\nabla\ell_{i}=\lambda^{-1}\nabla\ell_{i}-\sum_{j=1}^{i-1}a^{i-1}_{j}\nabla\ell_{j}=\sum_{j=1}^{m}b^{i}_{j}\nabla\ell_{j},$
(7)

i.e. containing the scalar coefficients of
$\widehat{F}_{i-1}^{-1}\nabla\ell_{i}$. This means, $b^{i}_{j}=-a^{i-1}_{j}$
when $j<i$, $b^{i}_{i}=\lambda^{-1}$ and $b^{i}_{j}=0$ otherwise.

Initial Precomputation & Update. The dynamic algorithm maintains three
$m\times m$ matrices: $\mathbf{G}\mathbf{G}^{\top}$, $\mathbf{D}$ and
$\mathbf{B}$. The first is the symmetric matrix
$[\mathbf{G}\mathbf{G}^{\top}]_{ij}=\nabla\ell_{i}^{\top}\nabla\ell_{j}$,
composed of the dot products between gradients. The second is an upper
triangular matrix and stores the precomputed values
$[\mathbf{D}]_{ij}=\nabla\ell_{i}\widehat{F}_{i-1}^{-1}\nabla\ell_{j}$ for
$i\leq j$. The third is the row-wise matrix of the coefficient vectors
$[\mathbf{B}]_{ij}=b^{i}_{j}$, which makes it lower triangular with a diagonal
of $\lambda^{-1}$. We now discuss how to compute those matrices.

The initial setup of the dynamic algorithm begins by evaluating
$\mathbf{G}\mathbf{G}^{\top}$ in a straightforward fashion. Next, $\mathbf{D}$
can be computed for $i\leq j$ according to the following recursion:

$\displaystyle[\mathbf{D}]_{ij}$
$\displaystyle=\left(\ell_{i}^{\top}\widehat{F}_{i-1}^{-1}\nabla\ell_{j}\right)\textnormal{
and
}\left(\nabla\ell_{j}^{\top}\widehat{F}_{0}^{-1}\nabla\ell_{k}\right)=\lambda^{-1}\left[\mathbf{G}\mathbf{G}^{\top}\right]_{jk}$
(8)
$\displaystyle\left(\nabla\ell_{j}^{\top}\widehat{F}_{i}^{-1}\nabla\ell_{k}\right)$
$\displaystyle=\left(\nabla\ell_{j}^{\top}\widehat{F}_{i-1}^{-1}\nabla\ell_{k}\right)-\frac{\left(\nabla\ell_{i}^{\top}\widehat{F}_{i-1}^{-1}\nabla\ell_{j}\right)\left(\nabla\ell_{i}^{\top}\widehat{F}_{i-1}^{-1}\nabla\ell_{k}\right)}{m+\left(\nabla\ell_{i}^{\top}\widehat{F}_{i-1}^{-1}\nabla\ell_{i}\right)}.$
(9)

Given $\mathbf{D}$, we can then conclude the precomputation by calculating
$\mathbf{B}$ for $i\geq j$ recursively as:

$\displaystyle[\mathbf{B}]_{ii}$ $\displaystyle=\lambda^{-1}\,\,\textnormal{
and
}\,\,[\mathbf{B}]_{ij}=-\sum_{k=j}^{i-1}\frac{[\mathbf{D}]_{ki}}{m+[\mathbf{D}]_{kk}}[\mathbf{B}]_{kj}.$
(10)

After the initial setup, gradient $\nabla\ell_{k}$ can be replaced with the
gradient $\nabla\ell^{\prime}_{k}$ by first updating row $k$ of $\mathbf{G}$
and then replacing row and column $k$ in $\mathbf{G}\mathbf{G}^{\top}$ with
$\mathbf{G}\nabla\ell^{\prime}_{k}$. Afterwards, the recomputation of columns
$j\geq k$ of $\mathbf{D}$ and rows $i\geq k$ of $\mathbf{B}$ completes the
update.

Multiplication. Once $\mathbf{G}\mathbf{G}^{\top}$, $\mathbf{D}$ and
$\mathbf{B}$ have been precomputed, one can perform efficient IHVPs of the
form $\widehat{F}_{m}^{-1}\mathbf{x}$ with arbitrary vectors ${\mathbf{x}}$.
This is done by first evaluating $\mathbf{p}=\mathbf{G}\mathbf{x}$ and then
computing all $m$ values $q_{i}$ by the following recursion:

$\displaystyle q_{i}$
$\displaystyle=\frac{\left(\ell_{i}^{\top}\widehat{F}_{i-1}^{-1}\mathbf{x}\right)}{m+[\mathbf{D}]_{ii}}\textnormal{
and
}\left(\nabla\ell_{j}^{\top}\widehat{F}_{0}^{-1}\mathbf{x}\right)=\lambda^{-1}p_{j}.$
(11)
$\displaystyle\left(\nabla\ell_{j}^{\top}\widehat{F}_{i}^{-1}\mathbf{x}\right)$
$\displaystyle=\left(\nabla\ell_{j}^{\top}\widehat{F}_{i-1}^{-1}\mathbf{x}\right)-\frac{[\mathbf{D}]_{ij}}{m+[\mathbf{D}]_{ii}}\left(\nabla\ell_{i}^{\top}\widehat{F}_{i-1}^{-1}\mathbf{x}\right).$
(12)

Eventually, the final result of the IHVP is obtained by:

$\widehat{F}_{m}^{-1}\mathbf{x}=\lambda^{-1}\mathbf{x}-\sum_{j=1}^{m}\left(\sum_{k=j}^{m}q_{k}[\mathbf{B}]_{kj}\right)\nabla\ell_{j}.$
(13)

Complexity Analysis. The dynamic algorithm stores $m$ gradients of dimension
$d$ as well as three $m\times m$ matrices (and a few $m$ vectors) and thus has
an overall memory complexity of $O(dm+m^{2})$.

Next, we analyze the time complexity of all important operations. Initially,
$\mathbf{G}\mathbf{G}^{\top}$ must be computed once, which takes $O(dm^{2})$
time. The recursion of $\mathbf{D}$ has three indices with $1\leq i,j,k\leq m$
and each step takes constant time. Thus, it can be computed in $O(m^{3})$ 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(m^{2})$ space. $\mathbf{B}$ has two indices $1\leq i,j\leq m$ and every
recursion takes $O(m)$ time, meaning that it can also be fully computed in
$O(m^{3})$ through dynamic programming. Hence, the overall initial setup cost
is $O(dm^{2}+m^{3})$.

To replace one gradient with $\nabla\ell^{\prime}$, we have to compute
$\mathbf{G}\nabla\ell^{\prime}$ as well as (partially) recalculate
$\mathbf{D}$ and $\mathbf{B}$, which takes at worst $O(dm+m^{3})$ time. An
IHVP requires two matrix-vector products involving $\mathbf{G}$ and a
recursion with two indices and therefore has a complexity of $O(dm+m^{2})$.

Algorithmic Derivation. The dynamic algorithm can be derived directly from
Theorem 1, which we now state formally and prove.

######  Theorem 1.

Let $(\nabla\ell_{i})_{i=1}^{m}$ be a series of gradients, and, for any index
$i\geq 1$, let $\widehat{F}_{i}=\lambda
I_{d}+\frac{1}{m}\sum_{j=1}^{i}\nabla\ell_{j}\cdot\nabla\ell_{j}^{\top}$ be a
dampened version of the empirical Fisher, where $\lambda>0$ is a small
constant. Then $\widehat{F}_{i}^{-1}{\mathbf{x}}$ can be calculated as:

$\widehat{F}_{i}^{-1}{\mathbf{x}}=\lambda^{-1}{\mathbf{x}}-\sum^{i}_{j=1}c^{i}_{j}\nabla\ell_{j}\,\,\textnormal{
and
}\,\,c^{i}_{j}=\sum_{k=j}^{i}\frac{\left(\nabla\ell_{k}^{\top}\widehat{F}_{k-1}^{-1}{\mathbf{x}}\right)}{m+[\mathbf{D}]_{kk}}[\mathbf{B}]_{kj},$
(14)

where
$[\mathbf{D}]_{kk}=\nabla\ell_{k}^{\top}\widehat{F}_{k-1}^{-1}\nabla\ell_{k}$,
and $[\mathbf{B}]_{kj}$ being defined such that

$\widehat{F}_{k-1}^{-1}\nabla\ell_{k}=\sum_{j=1}^{k}[\mathbf{B}]_{kj}\nabla\ell_{j}\,\,\text{and}\,\,[\mathbf{B}]_{11}=\lambda^{-1}.$
(15)

###### Proof.

The proof makes use of the following two equalities:

$\displaystyle\widehat{F}_{i}^{-1}{\mathbf{x}}$
$\displaystyle=\lambda^{-1}{\mathbf{x}}-\sum^{i}_{j=1}c^{i}_{j}\nabla\ell_{j}$
(16) $\displaystyle\widehat{F}_{i}^{-1}{\mathbf{x}}$
$\displaystyle=\lambda^{-1}{\mathbf{x}}-\sum_{j=1}^{i}\left(\widehat{F}_{j-1}^{-1}\nabla\ell_{j}\right)\frac{\left(\widehat{F}_{j-1}^{-1}\nabla\ell_{j}\right)^{\top}{\mathbf{x}}}{m+[\mathbf{D}]_{ii}},$
(17)

which correspond to the Equation 6 with $m=i$ and the Equation 3 in an
unrolled form, respectively. We begin by setting them equal. Next, we apply
(16) again to $\widehat{F}_{j-1}^{-1}\nabla\ell_{j}$ (naming the corresponding
coefficients $a^{j-1}_{k}$) and reorganize the nested sums to get a new
expression of the form (16). Finally, we simplify the result by using our
definition of $[\mathbf{B}]_{ij}=-a^{i-1}_{j-1}$ for $j<i$ and
$[\mathbf{B}]_{ii}=\lambda^{-1}$ (discussed at the beginning of Section 3.2).
We obtain:

$\displaystyle\widehat{F}_{i}^{-1}{\mathbf{x}}$
$\displaystyle=\lambda^{-1}{\mathbf{x}}-\sum^{i}_{j=1}c^{i}_{j}\nabla\ell_{j}$
(18)
$\displaystyle=\lambda^{-1}{\mathbf{x}}-\sum_{j=1}^{i}\left(\widehat{F}_{j-1}^{-1}\nabla\ell_{j}\right)\frac{\left(\widehat{F}_{j-1}^{-1}\nabla\ell_{j}\right)^{\top}{\mathbf{x}}}{m+[\mathbf{D}]_{ii}}$
(19)
$\displaystyle=\lambda^{-1}{\mathbf{x}}-\sum_{j=1}^{i}\left(\lambda^{-1}\nabla\ell_{j}-\sum_{k=1}^{j-1}a^{j-1}_{k}\nabla\ell_{k}\right)\frac{\left(\nabla\ell_{j}^{\top}\widehat{F}_{j-1}^{-1}{\mathbf{x}}\right)}{m+[\mathbf{D}]_{ii}}$
(20)
$\displaystyle=\lambda^{-1}{\mathbf{x}}-\sum_{j=1}^{i}\left(\frac{\left(\nabla\ell_{j}^{\top}\widehat{F}_{j-1}^{-1}{\mathbf{x}}\right)}{m+[\mathbf{D}]_{ii}}\lambda^{-1}-\sum_{k=j}^{i-1}\frac{\left(\nabla\ell_{k}^{\top}\widehat{F}_{k-1}^{-1}{\mathbf{x}}\right)}{m+[\mathbf{D}]_{kk}}a_{j}^{k-1}\right)\nabla\ell_{j}$
(21)
$\displaystyle=\lambda^{-1}{\mathbf{x}}-\sum_{j=1}^{i}\left(\sum_{k=j}^{i}\frac{\left(\nabla\ell_{k}^{\top}\widehat{F}_{k-1}^{-1}{\mathbf{x}}\right)}{m+[\mathbf{D}]_{kk}}[\mathbf{B}]_{kj}\right)\nabla\ell_{j}.$
(22)

We can directly match coefficients between (18) and (22), as they hold for any
gradient values. This completes the proof. ∎

Equation (14) with $i=m$ is exactly equal to the IHVP computation (13) as the
fraction in the innermost sum corresponds to $q_{k}$. Similarly, index
shifting $i=i-1$ and setting $x=\nabla\ell_{i}$, thus turning $c^{i}_{j}$ into
$-[\mathbf{B}]_{ij}$ as well as the fraction’s numerator to
$[\mathbf{D}]_{ki}$, recovers the precomputation formula of $\mathbf{B}$ given
by (10) for $i>j$. The formulas for $\mathbf{D}$ and $q_{i}$ 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).

###  3.3 Efficient Implementation

Directly implementing the discussed recursive formulas in a modern machine
learning framework would likely 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 and provide a full
implementation in [10]. Specifically, our implementation is able to perform
the calculation of $\mathbf{D}$ and $\mathbf{B}$, i.e. the $O(m^{3})$
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
optimized further). For reference, this is $>10\times$ faster than the highly
optimized $m\times m$ SVD computation done at every step by GGT [2]. We
emphasize that this $O(m^{3})$ computation being very fast in practice is
crucial to reach low overheads, especially when dealing with a sparse model
where the $O(dm)$ matrix-vector products are no longer the bottleneck.

##  4 Experimental Validation

We now validate our algorithms empirically, on pruning and optimization
applications, respectively. Implementations of the algorithms are available at
[10] for the optimizer and one-shot pruning. The framework of [18] implements
a variant of the pruning algorithm, along with an integration into common
PyTorch training flows for gradual pruning.

###  4.1 Application 1: Pruning DNNs using the Static Algorithm

Background. Given a trained model ${\mathbf{\theta}}^{\star}$, the goal of
pruning is to find the weight $\mathbf{\theta}_{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
[26] shows that the “optimal” weight to be removed is the one with the lowest
value of the saliency metric
$\rho(\theta_{i})=\frac{\theta_{i}^{2}}{2\,[\mathbf{\widehat{\mathbf{F}}}^{-1}]_{ii}}$
and proposed to estimate $[\mathbf{\widehat{\mathbf{F}}}^{-1}]_{ii}$ by a
diagonal approximation. The OBS framework [14] observed that the _remaining_
(unpruned) weights should also be updated via the optimal perturbation
$\delta\mathbf{\theta}=-\frac{\mathbf{\theta}_{i}}{[\mathbf{\widehat{\mathbf{F}}}^{-1}]_{ii}}\mathbf{\widehat{\mathbf{F}}}^{-1}\mathbf{e}_{i},$
where $\mathbf{e}_{i}$ is the $i$th basis vector. Our algorithm supports these
operations in linear-in-$d$ time.

Wang et al. [43] raised the valid point that applying the OBS update from
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 [26], the OBS update [14],
and an augmented version of OBS which disentangles the correlations by solving
a corresponding linear system [39]. The results, presented in Appendix C.1,
suggest that the OBD update is quite beneficial even in its approximate form,
and that the effect of correlation is small for unstructured pruning. We use
the augmented version of OBS, solving for correlations when feasible.

Experimental Setup. We prune CNNs (ResNet50 [16] and MobileNet-V1 [17]) on the
ImageNet dataset [38]. The models are standard in the pruning literature [11],
and therefore several strong baselines exist. We pick sparsity levels which
allow us to compare against previous work. Timing experiments are run on a
machine with NVIDIA RTX 2080 Ti GPUs, a 48-core Intel CPU, and 512 GB of RAM.
Following [39], we have used batched gradients (of size 16) as single samples
inside the Fisher approximation. This does not alter the nature of the
results, but reduces variance.

We compare against Global Magnitude Pruning (GMP) [13], Layer-wise Optimal
Brain Surgeon (L-OBS) [7], Soft Threshold Reparametrization (STR) [24], and
WoodFisher [39]. The latter two methods are state-of-the-art for gradual
pruning. WoodFisher is a special case of our method, using a block
approximation. It is numerically equivalent at the same parameter settings,
but has significantly higher computational and storage costs. Relative to it,
we therefore focus on executing at higher parameter values. We compare against
their public implementation. The Appendix contains full hyper-parameters,
additional experiments and ablations with respect to the block size and number
of gradients.

Gradual pruning is performed in incremental steps, according to a polynomial
schedule [49], followed by fine-tuning, where the learning rate is gradually
decreased. We follow the hyper-parameters of [39]: pruning is performed in
steps at 4 epochs apart for the first 24 epochs, followed by fine-tuning. In
the following pruning experiments, we will adopt the standard approach of
measuring _test accuracy_ , which is a more practical metric of interest. In
all cases, the results correlate perfectly with respect to _training accuracy_
, so we omit this latter metric.

Pruning Results. Figure 1(a) 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 for this
experiment following [7].) We note the improved accuracy of global methods
relative to L-OBS and layer-wise magnitude pruning, and the fact that the
Fisher approximation yields consistently better results relative to the 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 [39]. Generally, we found that estimating using more gradients always
improves results, until saturation. However, for some models, smaller block
sizes sometimes yield better results for a fixed number of gradients,
presumably due to gradient scaling effects.

Figure 1(b) 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.

(a) One-shot pruning for ResNet50.

(b) Gradual pruning for MobileNetV1-STR.

Figure 1: Left: Accuracy vs. sparsity for one-shot pruning on
ResNet50/ImageNet. Right: Accuracy during gradual pruning for the
MobileNetV1-STR/ImageNet experiment. The notation M-FAC($X$, $Y$) means we are
basing our approximation on a block size $X$, and $Y$ gradients.

Table 1 presents gradual pruning results for MobileNetV1-STR/ImageNet at 89%
sparsity and ResNet50/ImageNet at 95% sparsity, relative to other state-of-
the-art methods. M-FAC outperforms previous methods in the case of
MobileNetV1-STR by more than 1.1% Top-1 accuracy. Further, despite using
$10\times$ more gradients to estimate the empirical Fisher, the per-step
pruning cost of M-FAC is $>6\times$ lower than WoodFisher’s as can be seen in
Table 2 (for the same parameter settings M-FAC is $>100\times$ faster). The
ResNet50 results show similar gains, but by smaller margins.

Method | MBv1 Top-1 (%) | RN50 Top-1 (%)
---|---|---
Dense baseline | 72.00 | 77.01
STR [24] | 62.10 | 70.40
Global Magnitude | 62.94 | 71.72
WoodFisher [39] | 63.87 | 72.12
M-FAC | 65.01 | 72.32

Table 1: Gradual pruning results (Top-1 accuracy) for MobileNetV1-STR on
ImageNet at 89% sparsity and ResNet50 on ImageNet at 95% sparsity.

Model | $B$ | $m$ | WF | M-FAC
---|---|---|---|---
MBv1-STR | 10k | 400 | 60m | 0.5m
MBv1-STR | 10k | 4k | OOM | 8.7m
MBv1-STR | all | 1k | OOM | 1.2m
ResNet50 | 2k | 400 | 35m | 2.5m
ResNet50 | 10k | 1k | OOM | 14.0m
ResNet50 | all | 512 | OOM | 4.7m

Table 2: Minutes per pruning step compared with WoodFisher (WF). OOM denotes
out of memory on a system with 512GB RAM.

Normalizer-Free Nets. We now examine the “compressibility” of the recently-
proposed normalizer-free nets, which have shown competitive accuracy based on
a similar structure to standard residual networks but without the batch
normalization component [6]. We use the PyTorch re-implementation of [44]. In
Figure 2(a), we provide a relative comparison in terms of pruned accuracy
between a normalizer-free and a regular version of ResNet50. These two
networks have virtually the same number of parameters (approximately 25M);
however, we notice that the normalizer-free variant is significantly “easier”
to prune, in terms of the relative accuracy drop versus the dense baseline. We
conjecture that this is because of the elimination of the BatchNorm layers.
Specifically, when performing large one-shot pruning steps, the BatchNorm
statistics become invalid following a pruning step, which can lead to a
significant loss of accuracy; moreover, their removal may render the Fisher
approximation more accurate, as the model loss is more stable in the local
neighborhood.

YOLOv3. Next, we study the effect of one-shot pruning on the YOLOv3-SPP model
[37] for object detection on the COCO dataset [27]. We use the state-of-the-
art implementation of [9]. We one-shot prune this model using global magnitude
(the only available baseline) and M-FAC with block size 50K and 1K gradients.
(This parameter value would be infeasible for WoodFisher due to storage
costs.) The results in terms of the standard<EMAIL_ADDRESS>metric are provided in
Figure 2(b), and show that this model is quite stable under pruning, and that
M-FAC provides more accurate pruning for this model as well, relative to
magnitude pruning.

(a) Comparison between GMP and M-FAC pruning on the standard ResNet50 and a
Normalizer-Free variant (NF-ResNet50). Observe the significantly better
accuracy given by our method on the normalizer-free variant.

(b) Comparison between GMP and M-FAC in terms of the<EMAIL_ADDRESS>metric for a one-
shot pruned complex model YOLOv3-SPP/COCO. Notice the higher accuracy of the
M-FAC method.

Figure 2: Left: The effectiveness of M-FAC pruning on Normalizer-Free ResNet50
versus the regular variant. Right: One-shot pruning comparison for the
YOLOv3-SPP model on the large-scale object detection COCO dataset.

K-FAC Pruning Comparison. In Figure 3(a) we compare pruning performance
against a pruner which uses K-FAC to estimate second-order statistics (with
and without dampening $\pi$) on a Fully-Connected network and MNIST dataset.
Notice the improved accuracy with M-FAC (and WoodFisher) methods compared to
the K-FAC, even in the setting with a small network on a simple task.

Recomputation Effects. As noted before, our one-shot experiments stretch the
OBD/OBS theory, as this approach implicitly assumes that the Hessian stays
constant across the direction of pruning, which is unlikely to hold for very
large pruning displacements. We can however examine the impact of this effect,
by _recomputing_ the Hessian along the direction of pruning. Figure 3(b) shows
the effect of recomputation for the ResNet50/ImageNet model, for 5
recomputation steps, uniformly across the pruning direction. Notice the
significant increase in accuracy for the resulting sparse models. We note that
the results presented in other experiments do not use recomputation, for
relative fairness to the global magnitude pruning, for which recomputation
would not change the results.

(a) Comparison between K-FAC, WoodFisher and M-FAC in terms of the Top-1 test
accuracy for a one-shot pruned Fully-Connected network.

(b) The effect of recomputation on the accuracy of pruned ResNet50/ImageNet
model. Observe the significantly better accuracy provided by recomputation,
especially at high number of gradients.

Figure 3: Left: One-shot pruning comparison for the Fully-Connected network on
the MNIST dataset. Right: The effect of recomputation along the direction of
pruning for ResNet50/ImageNet.

###  4.2 Application 2: Optimization using the Dynamic Algorithm

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 $\Delta<t$ be the length of the sliding window during
which the estimation of the inverse Hessian is performed. Consider the
iteration
${\mathbf{\theta}}_{t+1}={\mathbf{\theta}}_{t}-\eta_{t}\widehat{F}^{-1}_{[t-\Delta,t]}\nabla\ell_{t},$
where $\eta_{t}$ is the learning rate. Let
$\nabla\ell_{t-\tau+1},\ldots,\nabla\ell_{t}$ be the gradients obtained at
steps $(t-\tau+1),\ldots,{t}$, respectively. Then, we define the
preconditioning matrix as
$\widehat{F}_{[t-\Delta,t]}=\frac{1}{\Delta}\sum_{\tau=0}^{\Delta-1}\nabla\ell_{t-\tau}\cdot\nabla\ell_{t-\tau}^{\top},$
that is, an empirical Fisher matrix generated with respect to a sliding window
over the last $\Delta$ gradients. At each step, our dynamic algorithm first
removes $\nabla\ell_{t-\Delta+1}$ and adds $\nabla\ell_{t}$ to the inverse
estimate, and then computes $\widehat{F}^{-1}_{[t-\Delta,t]}\nabla\ell_{t}$.

We emphasize the fact that this optimizer does not employ momentum, and
therefore only relies on the direct approximation of the curvature.

Setup and Objectives. We now experiment with the above optimizer, implemented
using the dynamic algorithm, without momentum. We compare against standard
optimizers such as SGD with momentum and Adam [19], but also against
approximate second-order methods such as K-FAC [32], GGT [2], and AdaHessian
[46]. 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 (even though this makes overheads harder to compare); for the best
accuracy on the models we considered, we grid-searched learning rates and
weight-decay for each method in turn.

We examine the _loss_ and _test accuracy_ for each method, on the standard
ResNet20 and ResNet32 models [16] on the CIFAR-10 dataset [22] (see also
Figures 4(a) and 4(b)). 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 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.
We provide code, as well as additional information, in the Appendix.

(a) Training loss on ResNet20/CIFAR-10.

(b) Top-1 test accuracy on ResNet20/CIFAR-10.

Figure 4: Comparison of different optimizers on the ResNet20/CIFAR-10 example.

Our first results are presented in Table 3 (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
“best accuracy” settings and also in the same setup as M-FAC; the latter is
given in parentheses. We emphasize the fact that, when available, we used the
hyper-parameter values recommended by the authors for the given model and
dataset; otherwise, for fairness, we perform exactly the same grid search
process for the tuning of each method.

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 much lower than
AdaHessian, on par with K-FAC, but higher than GGT. However, these overheads
are comparable only when GGT uses 5x less gradients (counter-intuitively,
larger window sizes sometimes make results worse for GGT). On the same
settings, M-FAC provides up to 5x lower overhead. We examined this performance
gap in detail by breaking down performance per operation, and determined that
the difference is due to: 1) the better performance of our dynamic algorithm
relative to the SVD computation used by GGT; 2) the difference in complexity
due to the $\Theta(dm^{2})$ term in GGT, versus our $O(dm)$ term. Similarly,
K-FAC’s overhead for the same batch size as M-FAC would be substantially
higher.

Table 3 (right) provides a more realistic scenario, where we run SGD, Adam,
and M-FAC _without tuning or weight-decay_ on Wide Residual Networks (WRN)
[47] 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.
However, we emphasize that these results can probably be improved
significantly by parameter tuning. Another interesting application of M-FAC is
fine-tuning sparse models. We fine-tune 90% sparse MobileNetV1/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 ($\leq 5\%$ overhead) since our algorithm’s complexity
is linear in the model _density_.

| ResNet20 | ResNet32
---|---|---
Method | Acc. | Overhead∗ | Acc. | Overhead∗
SGD | 91.78 | $1.00\times$ | 92.80 | $1.00\times$
Adam | 89.67 | $1.05\times$ | 90.56 | $1.10\times$
KFAC | 91.65 |  $1.21\times$ ($2.35\times$) | 90.09 |  $1.50\times$ ($3.05\times$)
GGT | 88.38 |  $1.05\times$ ($8.35\times$) | 89.14 |  $1.05\times$ ($8.88\times$)
AdaHessian | 92.17 | $3.50\times$ | 92.81 | $3.30\times$
M-FAC | 92.34 | $1.55\times$ | 92.65 | $1.50\times$

Model | SGD | Adam | M-FAC (Overh.)
---|---|---|---
WRN 22-2 | 69.93 | 66.90 | 69.76 ($1.40\times$)
WRN 40-2 | 71.75 | 70.14 |  72.42 ($1.35\times$)
WRN 22-4 | 73.13 | 72.52 |  74.06 ($1.43\times$)
MBv1 Dense | 68.06 | 67.92 |  68.96 ($1.25\times$)
MBv1 Sparse | 64.11 | – |  64.78 ($1.03\times$)
RN50 Sparse | 74.78 | – |  75.10 ($1.05\times$)

Table 3: Left: Comparison of second-order optimizers with individual tuning
and weight decay. Right: Additional experiments with no tuning and weight
decay. Sparse results are from fine-tuning after the last step of gradual
pruning to $\approx 90\%$ sparsity. For the dense models M-FAC uses $m=1024$
and for the sparse ones $m=512$. ∗ For KFAC/GGT we indicate in parentheses the
overhead in a comparable setting to M-FAC (i.e. same batch size/number of
gradients).

##### Wall-clock time comparison.

In Figures 5 and 6 we show the test accuracies during the training of Wide
ResNet (WRN) on the CIFAR-100 dataset and MobileNetV1 on the ImageNet dataset
with SGD, Adam and M-FAC with respect to the wall clock time, when executed on
a single NVIDIA RTX 3090 GPU. (In the case of the largest network, WRN 22-4,
we use a second GPU’s memory to store gradients, but still use a single GPU
for computation.) One can see that, in most plots, M-FAC reaches better
accuracies than SGD already after only slightly higher training time. Further,
we can see that, for the WRN model, the early training accuracy increases
fastest for M-FAC even with respect to the wall-clock time, a phenomenon often
observed with methods which try to leverage approximate second-order
information.

(a) Accuracy versus time for WRN 40-2/CIFAR-100.

(b) Accuracy versus time for MobileNetV1/ImageNet.

Figure 5: Optimization results in accuracy-versus-time format for WRN40-2 and
MobileNetV1.

(a) Accuracy versus time for WRN 22-2/CIFAR-100.

(b) Accuracy versus time for WRN 22-4/CIFAR-100.

Figure 6: Optimization results in accuracy-versus-time format for Wide
Residual Networks (WRN).

##  5 Discussion

We presented static and dynamic linear-in-$d$ 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 practical limitation of our approach is the cost of
storing the $m$ additional gradients. For the static algorithm, we can
leverage system memory to circumvent this limitation, via an efficient paging
mechanism, which we describe in the Appendix. The dynamic algorithm could
parallelize gradient storage, or perform gradient compression [40]. We plan to
investigate this in future work, as well as a larger-scale study of
optimization performance.

### Acknowledgments

This project has received funding from the European Research Council (ERC)
under the European Union’s Horizon 2020 research and innovation programme
(grant agreement No 805223 ScaleML).

## References

  * [1] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al.  Tensorflow: A system for large-scale machine learning.  In 12th USENIX symposium on operating systems design and implementation (OSDI 16), pages 265–283, 2016.
  * [2] Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang.  Efficient full-matrix adaptive regularization.  In International Conference on Machine Learning, pages 102–110. PMLR, 2019.
  * [3] Shun-ichi Amari.  Natural gradient works efficiently in learning.  Neural Computation, 10(2):251–276, 1998.
  * [4] Shun-ichi Amari.  Information geometry and its applications, volume 194.  Springer, 2016.
  * [5] Jimmy Ba, Roger Grosse, and James Martens.  Distributed second-order optimization using kronecker-factored approximations.  2016\.
  * [6] Andrew Brock, Soham De, Samuel L Smith, and Karen Simonyan.  High-performance large-scale image recognition without normalization.  arXiv preprint arXiv:2102.06171, 2021.
  * [7] Xin Dong, Shangyu Chen, and Sinno Jialin Pan.  Learning to prune deep neural networks via layer-wise optimal brain surgeon, 2017.
  * [8] John C. Duchi, Elad Hazan, and Yoram Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  J. Mach. Learn. Res., 12:2121–2159, 2010.
  * [9] Glenn Jocher et al.  YOLOv5 model implementations, AWS, Supervise.ly and YouTube integrations, April 2021.
  * [10] Elias Frantar and Eldar Kurtic.  M-FAC implementation: {https://github.com/IST-DASLab/M-FAC}, 2021\.
  * [11] Trevor Gale, Erich Elsen, and Sara Hooker.  The state of sparsity in deep neural networks, 2019.
  * [12] Roger Grosse and James Martens.  A kronecker-factored approximate fisher matrix for convolution layers, 2016.
  * [13] Masafumi Hagiwara.  A simple and effective method for removal of hidden units and weights.  Neurocomputing, 6(2):207 – 218, 1994.  Backpropagation, Part IV.
  * [14] Babak Hassibi and David G. Stork.  Second order derivatives for network pruning: Optimal brain surgeon.  In Advances in Neural Information Processing Systems 5, [NIPS Conference], page 164–171, San Francisco, CA, USA, 1992. Morgan Kaufmann Publishers Inc.
  * [15] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770–778, 2016.
  * [16] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Jun 2016.
  * [17] Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam.  Mobilenets: Efficient convolutional neural networks for mobile vision applications, 2017.
  * [18] Neural Magic Inc.  Sparseml framework: https://github.com/neuralmagic/sparseml, 2021\.
  * [19] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization, 2014.
  * [20] Diederik P Kingma, Tim Salimans, and Max Welling.  Variational dropout and the local reparameterization trick.  In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28, pages 2575–2583. Curran Associates, Inc., 2015.
  * [21] Shankar Krishnan, Ying Xiao, and Rif A Saurous.  Neumann optimizer: A practical optimization algorithm for deep neural networks.  arXiv preprint arXiv:1712.03298, 2017.
  * [22] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [23] Frederik Kunstner, Philipp Hennig, and Lukas Balles.  Limitations of the empirical fisher approximation for natural gradient descent.  In Advances in Neural Information Processing Systems, pages 4156–4167, 2019.
  * [24] Aditya Kusupati, Vivek Ramanujan, Raghav Somani, Mitchell Wortsman, Prateek Jain, Sham Kakade, and Ali Farhadi.  Soft threshold weight reparameterization for learnable sparsity, 2020\.
  * [25] César Laurent, Thomas George, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent.  An evaluation of fisher approximations beyond kronecker factorization, 2018.
  * [26] Yann Le Cun, John S. Denker, and Sara A. Solla.  Optimal Brain Damage, page 598–605.  Morgan Kaufmann Publishers Inc., San Francisco, CA, USA, 1990.
  * [27] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick.  Microsoft coco: Common objects in context.  In European conference on computer vision, pages 740–755. Springer, 2014.
  * [28] Dong C Liu and Jorge Nocedal.  On the limited memory bfgs method for large scale optimization.  Mathematical programming, 45(1):503–528, 1989.
  * [29] Alexander Ly, Maarten Marsman, Josine Verhagen, Raoul Grasman, and Eric-Jan Wagenmakers.  A tutorial on fisher information, 2017.
  * [30] James Martens.  Deep learning via hessian-free optimization.  In Proceedings of the 27th International Conference on International Conference on Machine Learning, ICML’10, page 735–742, Madison, WI, USA, 2010. Omnipress.
  * [31] James Martens.  New insights and perspectives on the natural gradient method.  arXiv preprint arXiv:1412.1193, 2014.
  * [32] James Martens, Jimmy Ba, and Matt Johnson.  Kronecker-factored curvature approximations for recurrent neural networks.  In International Conference on Learning Representations, 2018.
  * [33] James Martens and Roger Grosse.  Optimizing neural networks with kronecker-factored approximate curvature, 2015.
  * [34] Jorge Nocedal and Stephen Wright.  Numerical optimization.  Springer Science & Business Media, 2006.
  * [35] Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka.  Large-scale distributed second-order optimization using kronecker-factored approximate curvature for deep convolutional neural networks.  2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), Jun 2019.
  * [36] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  arXiv preprint arXiv:1912.01703, 2019.
  * [37] Joseph Redmon and Ali Farhadi.  Yolov3: An incremental improvement.  arXiv, 2018.
  * [38] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  International Journal of Computer Vision, 115(3):211–252, 2015\.
  * [39] Sidak Pal Singh and Dan Alistarh.  Woodfisher: Efficient second-order approximation for neural network compression, 2020.
  * [40] Xu Sun, Xuancheng Ren, Shuming Ma, and Houfeng Wang.  meProp: Sparsified back propagation for accelerated deep learning with reduced overfitting.  In Proceedings of the Thirty-Fourth International Conference on Machine Learning, 2017.
  * [41] Lucas Theis, Iryna Korshunova, Alykhan Tejani, and Ferenc Huszár.  Faster gaze prediction with dense networks and fisher pruning, 2018.
  * [42] Valentin Thomas, Fabian Pedregosa, Bart Merriënboer, Pierre-Antoine Manzagol, Yoshua Bengio, and Nicolas Le Roux.  On the interplay between noise and curvature and its effect on optimization and generalization.  In International Conference on Artificial Intelligence and Statistics, pages 3503–3513. PMLR, 2020.
  * [43] Chaoqi Wang, Roger Grosse, Sanja Fidler, and Guodong Zhang.  Eigendamage: Structured pruning in the kronecker-factored eigenbasis, 2019\.
  * [44] Ross Wightman.  Pytorch image models.  https://github.com/rwightman/pytorch-image-models, 2019.
  * [45] Max A Woodbury.  Inverting modified matrices.  Statistical Research Group, 1950.
  * [46] Zhewei Yao, Amir Gholami, Sheng Shen, Kurt Keutzer, and Michael W Mahoney.  Adahessian: An adaptive second order optimizer for machine learning.  arXiv preprint arXiv:2006.00719, 2020.
  * [47] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  arXiv preprint arXiv:1605.07146, 2016.
  * [48] Wenyuan Zeng and Raquel Urtasun.  MLPrune: Multi-layer pruning for automated neural network compression, 2019.
  * [49] Michael Zhu and Suyog Gupta.  To prune, or not to prune: exploring the efficacy of pruning for model compression, 2017.

###### Contents

  1. 1 Introduction
  2. 2 Preliminaries and Related Work
  3. 3 Algorithm Description
    1. 3.1 The Static Algorithm
    2. 3.2 The Dynamic Algorithm
    3. 3.3 Efficient Implementation
  4. 4 Experimental Validation
    1. 4.1 Application 1: Pruning DNNs using the Static Algorithm
    2. 4.2 Application 2: Optimization using the Dynamic Algorithm
  5. 5 Discussion
  6. A Efficiently Implementing the Dynamic Algorithm
  7. B Efficiently Implementing the Static Algorithm
    1. B.1 Additional Optimizations
  8. C Additional Experimental Results
    1. C.1 Pruning Update Approximation
    2. C.2 Ablation Studies for One-Shot Pruning
      1. C.2.1 ResNet50/ImageNet
      2. C.2.2 MobileNetV1-STR/ImageNet
    3. C.3 Discussion
    4. C.4 Gradual Pruning Experiment Hyperparameters
    5. C.5 Optimization Experiment Hyperparameters
    6. C.6 Cosine Similarity of Descent Directions

Additional Material

##  Appendix A Efficiently Implementing the Dynamic Algorithm

As mentioned in Section 3.2, directly implementing the various recursive
formulas of the dynamic algorithm will most likely be quite slow in practice.
Thus, we now discuss how to develop an efficient practical implementation. The
reader can also find the full code for this implementation of the dynamic
algorithm (and the corresponding optimizer) in PyTorch, available at [10].

##### Setup.

We begin by vectorizing the calculations of $\mathbf{D}$ and $\mathbf{B}$.
Algorithm 1 describes the full setup procedure in PyTorch-like pseudocode.

Algorithm 1 Calculating $\mathbf{D}$ and $\mathbf{B}$ in $O(m^{3})$ time and
$O(m^{2})$ memory, assuming that $\mathbf{G}\mathbf{G^{\top}}$ is already
precomputed.

$\mathbf{D}\leftarrow\lambda^{-1}\mathbf{G}\mathbf{G^{\top}}$

$\mathbf{B}\leftarrow\lambda^{-1}\mathbf{I}_{m\times m}$

for $i\leftarrow 2,3,\dots,m$ do

$\mathbf{D}_{i:,i:}\leftarrow\mathbf{D}_{i:,i:}-\frac{1}{m+D_{i-1,i-1}}\mathbf{D}_{i-1:,i:}^{\top}\mathbf{D}_{i-1:,i:}$

end for

$\mathbf{D}\leftarrow\text{triu}(\mathbf{D})$

for $i\leftarrow 2,3,\dots,m$ do

$\mathbf{B}_{i,:i}\leftarrow(-\mathbf{D}_{:i,i}/(m+\text{diag}(\mathbf{D})_{:i}))^{\top}\mathbf{B}_{:i,:i}$

end for

Initially, $\mathbf{D}$ is the stored scalar product matrix
$\mathbf{G}\mathbf{G}^{\top}$ scaled by $\lambda^{-1}$. Row 1 already has the
correct final values. The algorithm then goes through $m-1$ iterations, each
completing one row but also updating all the rows of higher index. More
concretely, in iteration $i$, we subtract the outer product of the previous
row starting at the element one after the diagonal $\mathbf{D}_{{i-1},i:}$
with itself, scaled by the inverse of $m$ plus the previous diagonal element
$\mathbf{D}_{i-1,i-1}$, from the square sub-matrix with $ii$ as the upper left
corner, i.e. $\mathbf{D}_{i:,i:}$. This process is also visualized in Figure
7, where the already calculated elements are shaded in light blue, the ones
that are currently being updated in darker blue, and the irrelevant ones in
grey. Since the calculation of $\mathbf{D}$, as given here, will also produce
some unnecessary (and incorrect) values below the diagonal, we delete those
after the loop by extracting just the upper triangular part of $\mathbf{D}$.
$\mathbf{B}$ starts off as the identity matrix times $\lambda^{-1}$ (and row 1
is again already done). Next, $\mathbf{B}$ is calculated row by row, where row
$i$ is a negative linear combination of the previous $i-1$ rows (up to index
$i-1$ as $\mathbf{B}$ is lower triangular) $\mathbf{B}_{:i,:i}$ with the
coefficients given by the first $i-1$ elements of column $i$ of $\mathbf{D}$,
i.e. $\mathbf{D}_{:i,i}$, each divided by $m$ plus the corresponding diagonal
element of $\mathbf{D}$, i.e. $m+\text{diag}(\mathbf{D})_{i}$. This process is
also shown in Figure 7.

Figure 7: Visualized calculation of $\mathbf{D}$ and $\mathbf{B}$.

In theory, the algorithm presented so far should be well suited for utilizing
the massive parallel computation capabilities of a GPU. However, a
straightforward implementation in e.g. PyTorch, will not be very fast for
larger values of $m$ as it results in thousands of small kernel launches with
overhead that adds up to a very significant overall slow-down. (Concretely, we
measured $>7\times$ overhead for $m=1024$). This means that, to achieve good
performance, the setup for the coefficient computation of the dynamic
algorithm needs to be implemented with (custom) CUDA kernels that merge
computations to get rid of the excessive kernel launch overhead.

Carefully studying the calculation process of $\mathbf{D}$ reveals that it is
actually very similar to a Cholesky decomposition (concretely, the only
difference is that it does not use the square root of the diagonal element but
the element itself plus $m$). In fact, it is actually exactly equivalent to
taking the upper triangular matrix in the LU-decomposition (with no pivoting)
of $\mathbf{\lambda^{-1}GG^{\top}}+m\mathbf{I}$ and then subtracting
$m\mathbf{I}$ again. This means that it is possible to reuse the highly
optimized LU-decomposition kernel provided by PyTorch to calculate
$\mathbf{D}$ quite efficiently (a custom kernel implementing the modified
Cholesky decomposition would certainly be faster, but we already observed
pretty low overheads with the LU version). For the efficient computation of
$\mathbf{B}$, a custom kernel is needed, which we provide at [10], and
describe next.

The main idea is that we can dramatically reduce the number of separate kernel
launches by computing multiple rows of $\mathbf{B}$ in a single kernel. In
general, we split $\mathbf{B}$ into blocks of $32\times 32$ (this matches the
fact that modern NVIDIA GPUs have 1024 threads per block). Then, we begin by
fully computing all diagonal blocks in parallel. Notice that, due to the lower
triangular structure of $\mathbf{B}$, those are fully independent. Similarly,
each block depends exclusively on the blocks above it (up to the diagonal).
Next, we perform an iterative process which will update all remaining blocks
with respect to the new values of the most recently computed diagonal before
completing the calculation of the blocks adjacent to it. This new diagonal is
then the reference for the next iteration. Figure 8 visualizes this process.
The calculation within the individual blocks is easy to parallelize. For more
details, please refer to our code for this CUDA kernel. Overall, this way of
calculating $\mathbf{B}$ is quite fast; e.g. for $m=1024$, it takes $<2$
milliseconds to execute on an NVIDIA RTX 3090 GPU.

Figure 8: Visualization of the CUDA kernel for calculating $\mathbf{B}$. In
each iteration, all blocks below the previous diagonal (here denoted as 1) are
updated in parallel and those adjacent to the diagonal (here denoted as 2) are
completed to form the reference diagonal for the next iteration.

##### IHVPs

After discussing an efficient implementation of the precomputation steps, we
now focus on actually computing IHVPs. Algorithm 2 presents a vectorized
implementation.

Algorithm 2 Computing $\widehat{F}_{m}^{-1}{\mathbf{x}}$ in $O(dm+m^{2})$
time, assuming that $\mathbf{D}$ and $\mathbf{B}$ have already been
precomputed.

$\mathbf{q}\leftarrow\lambda^{-1}\mathbf{G}{\mathbf{x}}$

$q_{1}\leftarrow q_{1}/(m+\mathbf{D}_{11})$

for $i\leftarrow 2,3,\dots,m$ do

${\mathbf{q}}_{i:}\leftarrow{\mathbf{q}}_{i:}-q_{i-1}\mathbf{D}_{i-1,i:}^{\top}$

$q_{i}\leftarrow q_{i}/(m+\mathbf{D}_{ii})$

end for

return
$\lambda^{-1}{\mathbf{x}}-(({\mathbf{q}}^{\top}\mathbf{B})\mathbf{G})^{\top}$

We note that if $\mathbf{x}=\nabla\ell_{m}$, then
$\mathbf{q}=\mathbf{D}_{:,m}/\text{diag}(\mathbf{D})$, i.e. the last column of
$\mathbf{D}$ divided element-wise by the diagonal values. This can be utilized
for a more efficient combined update-multiply operation, which is essential
for the M-FAC optimizer, where we first update our Hessian estimation with a
gradient $\nabla\ell$ and then return the IVHP with this same $\nabla\ell$.
Further, reusing the same $\mathbf{p}=\mathbf{G}\nabla\ell$ for both updating
$\mathbf{G}\mathbf{G}^{\top}$ as well as calculating $\mathbf{q}$, saves one
redundant $O(dm)$ matrix-vector-product.

##  Appendix B Efficiently Implementing the Static Algorithm

Although the formulas discussed in the main body of the paper have low
theoretical time complexity, a direct implementation does not make the best
use of modern machine learning frameworks and GPUs. Therefore, we now present
more efficient matrix-based versions. We note that a block-wise version would
simply apply the techniques discussed here independently to each block.

Let
$\mathbf{V}=[\mathbf{v_{1}}^{\top};\mathbf{v_{2}}^{\top};\dots;\mathbf{v_{m}}^{\top}]$
be the $m\times d$ row-wise matrix of the $\mathbf{v_{i}}$ and
$\mathbf{q}=(q_{1},q_{2},\dots,q_{m})^{\top}$ the $m$-vector of the $q_{i}$.
Then, $\widehat{F}_{m}^{-1}{\mathbf{x}}$ can be written as follows, where $/$
denotes the element-wise division:

$\widehat{F}_{m}^{-1}{\mathbf{x}}=\lambda^{-1}{\mathbf{x}}-\mathbf{V}^{\top}\left(\left(\mathbf{V}{\mathbf{x}}\right)/{\mathbf{q}}\right).$
(23)

It is also possible to extract one element of each matrix row simultaneously
(in terms of vectorized operations) with $O(dm)$ cost, for example the entire
diagonal. Let $\bm{\pi}({\mathbf{x}})$ be a mapping such that
$\bm{\pi}({\mathbf{x}})_{i}=x_{j}$ where $j$ is the column of the element to
select in row $i$, e.g. $\bm{\pi}({\mathbf{x}})={\mathbf{x}}$ to get the
diagonal, and $\mathbf{e^{\pi}}$ a vector such that $e^{\pi}_{i}=1$ if
$\pi({\mathbf{x}})_{i}={\mathbf{x}}_{i}$ and $e^{\pi}_{i}=0$ otherwise. Using
these definitions, the calculation of the desired result ${\mathbf{y}}$ is
described below, where $\odot$ denotes the element-wise product.

${\mathbf{y}}=\lambda^{-1}\mathbf{e^{\pi}}-\sum^{m}_{i=1}\frac{1}{q_{k}}\mathbf{v_{i}}\odot\bm{\pi}(\mathbf{v_{i}})$
(24)

Our implementation contains several additional optimizations. For instance, we
can efficiently precompute $\mathbf{V}$ and $\mathbf{q}$ by repeatedly
applying (23) to produce the next $\mathbf{v_{i}}$. In the next section, we
additionally discuss several memory-saving optimizations, and in particular
explicit page swapping between CPU and GPU memory for situations where
$\mathbf{V}$ does not fully fit in the GPU memory.

Finally, Algorithm 3 demonstrates, using a Python-like matrix indexing syntax,
how to efficiently precompute $\mathbf{V}$ and $\mathbf{q}$ by repeatedly
applying (23) to produce the next $\mathbf{v_{i}}$. It should be noted that
the row-wise matrix of gradients
$\mathbf{G}=[\nabla\ell_{1}^{\top};\nabla\ell_{2}^{\top};\dots;\nabla\ell_{m}^{\top}]$
is re-purposed as $\mathbf{V}$, a simple way to halve the peak memory
consumption in practice.

Algorithm 3 Precomputation of $\mathbf{V}$ and ${\mathbf{q}}$ given
$\mathbf{G}$.

$\mathbf{V}\leftarrow\mathbf{G}$

${\mathbf{g}}\leftarrow\mathbf{V}_{1,:}^{\top}$

$\mathbf{V}_{1,:}\leftarrow\lambda^{-1}{\mathbf{g}}$

$q_{1}\leftarrow\mathbf{V}_{1,:}\cdot{\mathbf{g}}$

for $i\leftarrow 2,3\dots,m$ do

${\mathbf{g}}\leftarrow\mathbf{V}_{i,:}^{\top}$

$\mathbf{V}_{i,:}\leftarrow\lambda^{-1}{\mathbf{g}}^{\top}-((\mathbf{V}_{:i,:}\cdot{\mathbf{g}})/{\mathbf{q}})^{\top}\mathbf{V}_{:i,:}$

$q_{i}\leftarrow m+\mathbf{V}_{i,:}\cdot{\mathbf{g}}$

end for

###  B.1 Additional Optimizations

As we have seen, obtaining a good Fisher approximation of the Hessian requires
a sizable number of gradients, which means that, for bigger networks, it can
easily happen that the collection of all gradients used for Fisher estimation
does not fully fit into GPU memory. We now discuss how to efficiently handle
such situations, with an implementation that performs explicit swapping of
gradient blocks between GPU and CPU memory (RAM). The most important steps of
the method to be discussed are also visualized in Figure 9.

In general, a simple trick to halve the peak memory consumption of the static
algorithm is repurposing the gradient matrix $\mathbf{G}$ as $\mathbf{V}$,
which is possible as $\nabla\ell_{i}$ is not needed anymore after
${\mathbf{v}}_{i}$ and $q_{i}$ have been calculated. Now, for the explicit
swapping implementation we first split the collection of $m$ gradients into
$k$ blocks / pages $\mathbf{G}^{i}$ containing at most $m/k$ gradients each.
Those are then turned into the corresponding $\mathbf{V}^{i}$ blocks in
increasing order of $i$. Additionally, we maintain a single buffer block
$\mathbf{B}$ of the same size for accumulating intermediate results. All
blocks reside in CPU memory and are only loaded to the GPU as needed,
meanwhile ${\mathbf{q}}$ is small enough to be kept in GPU memory at all
times. To compute $\mathbf{V}^{i}$ we first load block $\mathbf{V}^{1}$ fully
into GPU memory. Then we load the first gradient of $\mathbf{G}^{i}$ denoted
by $\mathbf{g}^{i}_{1}$ and compute the corresponding ${\mathbf{v}}^{1}_{1}$
with respect to the loaded $\mathbf{V}^{1}$, which is afterwards saved in the
first index of the buffer ${\mathbf{b}}_{1}$. After repeating this for all
${\mathbf{g}}^{i}_{j}$, block $\mathbf{V}^{1}$ is swapped with block
$\mathbf{V}^{2}$ and the whole process starts again, accumulating the
resulting ${\mathbf{v}}^{2}_{j}$ into the buffer. Eventually, after handling
$\mathbf{V}^{i-1}$, we can load $\mathbf{B}$ into GPU memory, finish the
calculation of the ${\mathbf{v}}^{i}_{j}$ and store them in $\mathbf{V}^{i}$
(reusing the memory of $\mathbf{G}^{i}$). It should be noted that the loading
of ${\mathbf{g}}^{i}_{j}$ can be parallelized with the calculation of
${\mathbf{g}}^{i}_{j-1}$ and thus costs almost no extra time. Finally, one
wants to choose the number of blocks $k$ to be as small as possible, to
minimize the overhead caused by the $O(k^{2})$ page swaps.

Figure 9: Visualization of the swapping implementation of the static
algorithm, grey blocks have already been fully calculated. Left: $g$ is loaded
from block $\mathbf{G}^{3}$, which allows calculating ${\mathbf{v}}$ using the
GPU block $\mathbf{V}^{1}$ and then accumulating it into the buffer
$\mathbf{B}$. Right: Calculation of $\mathbf{V}^{3}$ is finished using the
completed buffer.

Overall, the implementation described above is effective in practice and
allows scaling up our static algorithm (with only modest overhead over a pure
GPU implementation), in its full non-blockwise version, to a large number of
gradients, even for relatively large networks such as ResNet50.

Lastly, we note that, for a block-wise implementation, where a single block
and all its corresponding gradient parts, easily fit into GPU memory (i.e. the
opposite of the situation described thus far in this section), it can be
beneficial to handle multiple blocks simultaneously via batch-matrix-
multiplies. We provide an implementation that can do this, leading to very
significant speed-ups on lower block-sizes where otherwise most of the time is
spent constantly loading in new memory from the CPU.

##  Appendix C Additional Experimental Results

###  C.1 Pruning Update Approximation

As discussed in the main paper, we use the OBS pruning framework but prune
multiple weights at the same time, which as pointed out by [43], in extreme
cases, is not guaranteed to be better than OBD due to not taking into account
potential correlations. Thus, we now explore how much of a problem this is in
practice. To do that, we compare three different M-FAC variations: not
updating the non-pruned weights at all (i.e. OBD, but with a better Hessian
approximation), jointly applying the OBS update for all pruned weights (i.e.
the OBS approximation) and applying the correct simultaneous update by solving
a large linear system (see [39] for details). We prune a ResNet20/CIFAR-10
model in steps of 10% down to 90% sparsity, using three variations of the full
block size M-FAC with $m=1024$ while recomputing the inverse Hessian
approximation after every pruning step. The recomputations ensure that we have
a reasonably accurate loss approximation at every pruning step and that the
number of weights pruned at each step is small enough to solve the
corresponding linear system. Table 4 provides the results.

Method | 10% | 20% | 30% | 40% | 50% | 60% | 70% | 80% | 90%
---|---|---|---|---|---|---|---|---|---
No update (OBD) | 91.4 | 91.3 | 91.0 | 90.2 | 88.4 | 84.8 | 77.2 | 50.4 | 11.5
Simultaneous OBS | 91.4 | 91.4 | 91.5 | 91.2 | 90.8 | 89.7 | 86.8 | 75.6 | 28.2
Linear solving | 91.5 | 91.4 | 91.4 | 91.2 | 90.7 | 89.7 | 87.5 | 75.7 | 28.1
Table 4: Comparing ResNet20/CIFAR-10 accuracies of different M-FAC variations
at varying levels of sparsity.

The results show a very clear gap in accuracy between the no-update and the
approximate update version. At the same time, there appears to be almost no
difference between the simultaneous OBS and the true OBS update, which has to
solve a (potentially very large) linear system. This suggests that the
simultaneous update approximation done for computational tractability is also
a reasonable choice in terms of achieved accuracy.

###  C.2 Ablation Studies for One-Shot Pruning

To further examine the properties of the local loss approximation, we now
present ablation studies on pretrained ResNet50/ImageNet and
MobileNetV1-STR/ImageNet models. Experiments perform one-shot pruning
according to the OBS metric estimated using M-FAC, with varying block size and
number of gradients. The dampening factor $\lambda$ in these experiments is
set to $10^{-5}$. Following [39], we used batched gradients (of size 16) as
single samples inside the Fisher approximation. (This does not alter results,
but reduces variance.)

The goal of these experiments is to examine two questions: 1) does larger
block size always imply a better approximation and 2) does higher number of
gradients always imply a better approximation? We will see that neither of
these questions have obvious answers, and will discuss a meta-algorithm for
choosing the block size and number of gradients.

The numbers presented for our one-shot experiments are the averages over two
runs. As the variance is very small, we omit error bars.  We sometimes break
the convention and “zoom into” the y axis for visibility.

####  C.2.1 ResNet50/ImageNet

The first set of experiments examines the dependency on the number of
gradients and block size for the ResNet50/ImageNet model. The left subfigure
in Figure 10 shows results for block sizes between $2K$ and all weights, i.e.
$d$, for a fixed number of $1K$ gradients, while the right subfigure shows the
same, but for $2K$ gradients. The first graph presents a fairly unintuitive
finding, i.e. that _lower_ block sizes appear to provide better pruning
accuracy. We analyze this finding in detail and explain it in Section C.3:
roughly, this is due to the fact that gradient entries are scaled by the
gradient norm over the block. As predicted by our analysis, this effect is
mitigated by increasing the number of gradients used for estimation, in which
case block size $10K$ yields the best results, and the performance of full
block size also improves. Please see Section C.3 for the full analysis.

Figure 11 examines the complementary effect, that of number of gradients for a
fixed block size (10K and 100K, respectively). The results suggest that more
gradients help improve the estimation of the Fisher matrix, although we
observe a saturation effect. We also note that, at higher one-shot sparsities
(e.g., 60% in one-shot, not shown for visibility), this effect does not always
hold. However, for such large pruning “perturbations” the pruning assumption
that the Hessian is constant along the direction of pruning is unlikely to
hold, which affects the stability of the results.

Figure 10: An ablation study on the ResNet50/ImageNet model showing the effect
of the block size for a fixed number of gradients.

####  C.2.2 MobileNetV1-STR/ImageNet

The second set of experiments shows the dependency on the number of gradients
and block size for the MobileNetV1-STR/ImageNet model. We use the
implementation and pre-trained weights from [24]. Figure 12 shows results for
block sizes between $2K$ and all weights, i.e. $d$, for a fixed number of
gradients ($1K$ in the left subfigure and $2K$ in the right subfigure). Again,
in both cases, for this model _lower_ block sizes appear to provide an
improved approximation. This is due to the gradient scaling effects which we
discuss in Section C.3. However, the results show that for this compact model
these effects are more prevalent than for the ResNet50 model.

Figure 11: An ablation study on the ResNet50/ImageNet model showing the effect
of the number of gradients for a fixed block size.

Figure 12: An ablation study on the MobileNetV1-STR/ImageNet model showing the
effect of the block size for a fixed number of gradients.

Figure 13 examines the opposite effect, that of the number of gradients for a
fixed block size ($10K$ in the left subfigure and $100K$ in the right
subfigure). We show results for the number of gradients varying between $50$
and $4K$. The results clearly suggest that more gradients help improve the
accuracy, although the improvement appears to saturate, e.g. between $2K$ and
$4K$ gradients.

Figure 13: An ablation study on the MobileNetV1-STR/ImageNet model showing the
effect of the number of gradients for a fixed block size.

###  C.3 Discussion

##### Normalization Effects.

We first discuss our finding that, in the case of some models, e.g.
MobileNetV1, smaller blocks appear to yield better accuracy. This can be
explained by examining the recursive form of the elements of the diagonal
inverse. Specifically, without blocking, we get that the $i$th diagonal
element has the form

$[\widehat{\mathbf{F}}^{-1}]_{ii}\simeq\frac{1}{\lambda}\left[1-\frac{\nabla\ell_{1_{i}}^{2}}{\|\nabla\ell_{1}\|^{2}}-\frac{\left(\nabla\ell_{2_{i}}-\nabla\ell_{1_{i}}\frac{\nabla\ell_{1}^{T}\nabla\ell_{2}}{\|\nabla\ell_{1}\|^{2}}\right)^{2}}{\|\nabla\ell_{2}\|^{2}-\frac{(\nabla\ell_{1}^{T}\nabla\ell_{2})^{2}}{\|\nabla\ell_{1}\|^{2}}}-\ldots\right],$

where $\nabla\ell_{1_{i}}$ and $\nabla\ell_{2_{i}}$ represent the $i$th
element of a gradient. Specifically, notice that, in the case of “full” block
size (“all” or $d$), the $i$th gradient entry is divided by the full gradient
norm, which may cause it to become negligible if the norm is large. In turn,
this leads to essentially the magnitude pruning ranking and update. By
contrast, in the case of smaller blocks, the entry is only divided by the norm
of the gradient over the block, which mitigates this normalization issue.
Similarly, using more gradients also helps, since there are more summands in
the above expression, which allows us to deviate from the magnitude pruning
baseline.

##### Choosing the block size value.

Given this, our procedure for choosing the best block size hyper-parameter
value is to simply perform a one-shot pruning ablation study for this
parameter, usually at the maximum gradient count at which one can reasonably
execute, at a moderate sparsity (e.g., 40% from a pretrained model). We then
choose the best performing block size, in terms of the model’s evaluation
metric. We note that these single pruning steps are quite fast when using our
implementation.

###  C.4 Gradual Pruning Experiment Hyperparameters

We begin by stating hyperparameters for our gradual pruning runs. For both
MobileNetV1-STR and ResNet50-STR, we employ _identical_ hyperparameter values
to WoodFisher [39], so that our results are fully comparable. The only
parameters we change are the ones which pertain to our algorithm, in
particular _block size_ and _number of gradients_ $m$.

For the MobileNetV1-STR gradual experiments, pruning starts from the fully-
trained model used by STR/WoodFisher, whose Top-1 validation accuracy is 72%.
Gradual pruning is then performed over a total of 24 epochs, pruning
additional weights every $4$ epochs, followed by fine-tuning until the next
pruning step. Pruning targets are set following the standard polynomial
schedule of [49, 11]. Unless otherwise stated, SGD is used for fine-tuning.
The learning rate during pruning is $0.005$, momentum is set to $0.9$, and
weight decay is $0.0001$. Following the last pruning step, fine-tuning is
applied until epoch $100$. The initial learning rate value of $0.005$ is
decayed multiplicatively by a factor of $0.92$, every epoch, starting from the
epoch $30$.

For ResNet50-STR gradual experiments, pruning starts from the fully-trained
model used by STR/WoodFisher, whose Top-1 validation accuracy is 77.01%.
Gradual pruning is then performed over a total of $40$ epochs, pruning
additional weights every $5$ epochs, followed by fine-tuning until the next
pruning step. Pruning targets are set following the standard polynomial
schedule of [49, 11]. Unless otherwise stated, SGD is used for fine-tuning.
The learning rate during pruning is $0.005$, momentum is set to $0.9$, and
weight decay is $0.0001$. Following the last pruning step, fine-tuning is
applied until epoch $100$. The initial learning rate value of $0.005$ is
decayed multiplicatively by a factor of $0.6$, every $6$ epochs, starting from
epoch $40$ until epoch $90$.

###  C.5 Optimization Experiment Hyperparameters

##### ResNet20/32

We now discuss the hyperparameters used for the ResNet20/32 comparison between
M-FAC and various first and second order optimizers. Both models are trained
with batch-size 128 for 164 epochs ($\approx 64000$ steps) while dropping the
learning rate by a factor of $0.1$ after $50\%$ and $75\%$ of training. This
is exactly the training setup used by [16]. The best initial learning rate and
weight decay values were determined by starting from the values suggested by
the authors of the respective methods; in addition, we have performed a grid
search around these values to examine whether better results can be obtained.
In particular, Table 5 summarizes the final hyper-parameters settings for SGD,
Adam, AdaHessian and M-FAC.

Method | learning rate | momentum | weight decay
---|---|---|---
SGD | 0.1 | 0.9 | 0.0001
Adam | 0.001 | (0.9, 0.999) | –
AdaHessian | 0.015 | (0.9, 0.999) | 0.003 (RN20), 0.0005 (RN32)
M-FAC ($m=512$) | 0.001 | – | 0.003
Table 5: Hyperparameter settings used for the ResNet20 (RN20) and ResNet32
(RN32) experiments.

All experiments were repeated 5 times with different random seeds and we
report the median of the best test accuracy; the standard deviations were
generally quite low at around $0.2$ percent accuracy.

##### Hyperparameters for GGT and K-FAC.

Unfortunately, GGT and K-FAC did not produce reasonable results in the setup
discussed so far (e.g. achieved poor accuracy or diverged). Thus, for
fairness, we decided to adopt the recommendations of the method authors, even
if those were no longer exactly comparable with the other experiments (e.g.
different batch-size, learning rate schedule, etc.). Further, we also
performed considerably more extensive hyper-parameter searches for these
methods.

For K-FAC, we use the authors’ carefully-tuned parameters for ResNet20, which
they published in their official
repository111https://github.com/tensorflow/kfac. For, ResNet32 we adopt the
same parameters, but use an initial learning rate of $0.075$ and an initial
dampening of $\approx 0.1887$, which we identified to work best for this model
via grid search. The biggest differences of this setup compared to M-FAC are
the batch-size of 1000 and the smoothly exponential decaying learning rate.

For GGT, we use a batch-size of 128 and a gradient window size of $r=100$
(which is $>5\times$ lower than for M-FAC, however we surprisingly found that
GGT’s performance degraded with larger $r$ and thus kept this as the best
value). Further, we use an initial learning rate of $0.05$ and a cosine
decaying schedule with $T=40$, which we found via grid search. All GGT
experiments were performed with the author’s implementation contributed to
TensorFlow.

##### Wide ResNet / MobileNetV1.

Next, we list the exact hyperparameters used for optimizing Wide ResNet (WRN)
and MobileNetV1 models. The corresponding experiments were designed to explore
how well M-FAC performs relative to other methods without any parameter
tuning, i.e. just using reasonable default values. In the case of Wide ResNet,
all models are trained for 200 epochs with batch size 256, where the learning
rate is dropped by a factor $0.1$ after $50\%$ and $75\%$ of training. SGD
uses an initial learning rate of $0.1$ and momentum of $0.9$ (as suggested by
[47]). Adam runs with default parameters, i.e. a learning rate of $0.001$,
$\beta_{1}=0.9$ and $\beta_{2}=0.999$. M-FAC uses $m=1024$ and the standard
learning rate of $0.001$ (but no momentum). No method uses any weight decay in
these experiments. For MobileNetV1, we use exactly the same optimizer settings
and learning rate schedule, but we train only for 100 epochs.

##### Sparse Finetuning.

For the sparse fine-tuning experiments, pruning occurs identically for all
methods. However, we fine-tune the models using either SGD or M-FAC with
$m=512$, for $30$ epochs. Both fine-tuning algorithms are run with exactly the
same hyperparameter values: initial learning rate $0.005$, reduced every $10$
epochs by a factor of $10$ and batch-size 256.

###  C.6 Cosine Similarity of Descent Directions

Finally, in Figure 14, we examine the quality of the sliding window
approximation to the Fisher matrix. We run optimization with M-FAC on
ResNet20/CIFAR-10 (using $m=512$ and otherwise the same hyperparameters as
discussed in the previous section), and every 512 steps we sample $2\times
512$ gradients to produce 2 static estimates of the Fisher matrix, and then we
compare the cosine similarity of the descent direction given by the dynamic
algorithm at this step, $\widehat{\mathbf{F}}^{-1}_{\mathord{\it
dynamic}}\cdot\nabla\ell$ with $\widehat{\mathbf{F}}^{-1}_{\mathord{\it
static1}}\cdot\nabla\ell$ (denoted as dynamic–static) as well as the cosine
similarity between $\widehat{\mathbf{F}}^{-1}_{\mathord{\it
static2}}\cdot\nabla\ell$ and $\widehat{\mathbf{F}}^{-1}_{\mathord{\it
static1}}\cdot\nabla\ell$ (denoted as static–static).

The results show that: 1) the cosine similarities between the sliding window
approximation and the “fresh” approximation are extremely close; 2) they tend
to improve significantly as we advance in the optimization process. Overall,
this validates the sliding window approximation made by our method.

Figure 14: Cosine similarities over the course of training the
ResNet20/CIFAR-10 model. The directions are remarkably close for the dynamic
algorithm, suggesting that the sliding window approximation is valid.