# Error Feedback Can Accurately
Compress Preconditioners

Ionut-Vlad Modoranu
IST Austria
<EMAIL_ADDRESS>
&Aleksei Kalinov
IST Austria
<EMAIL_ADDRESS>
&Eldar Kurtic
IST Austria
<EMAIL_ADDRESS>
&Elias Frantar
IST Austria
<EMAIL_ADDRESS>
&Dan Alistarh
IST Austria
<EMAIL_ADDRESS>Corresponding author.

###### Abstract

Leveraging second-order information about the loss at the scale of deep
networks is one of the main lines of approach for improving the performance of
current optimizers for deep learning. Yet, existing approaches for accurate
full-matrix preconditioning, such as Full-Matrix Adagrad (GGT) or Matrix-Free
Approximate Curvature (M-FAC) suffer from massive storage costs when applied
even to small-scale models, as they must store a sliding window of gradients,
whose memory requirements are multiplicative in the model dimension. In this
paper, we address this issue via a novel and efficient error-feedback
technique that can be applied to compress preconditioners by up to two orders
of magnitude in practice, without loss of convergence. Specifically, our
approach compresses the gradient information via sparsification or low-rank
compression _before_ it is fed into the preconditioner, feeding the
compression error back into future iterations. Experiments on deep neural
networks show that this approach can compress full-matrix preconditioners to
up to 99% sparsity without accuracy loss, effectively removing the memory
overhead of full-matrix preconditioners such as GGT and M-FAC. Our code is
available at https://github.com/IST-DASLab/EFCP.

##  1 Introduction

The remarkable success of stochastic gradient descent (SGD)-based optimizers
in deep learning motivated a long line of research for accelerated
preconditioned variants that can still scale to the massive parameter and
dataset sizes common in the area. Among these, optimizers based on adaptive
regularization, such as Adagrad [10] and Adam [18] are extremely well-
established, with myriads of extensions, e.g. [2, 1, 38, 14].

Yet, most work in this area restricts preconditioning of the descent direction
to _a diagonal matrix_. By contrast, a promising—but considerably less
developed—direction employs what we will call _full-matrix preconditioning_ ,
i.e. pre-multiplying the descent direction with a full matrix, whose structure
is justified either via adaptive regularization [1], or via approximations of
natural gradient such as the Fisher approximation [5, 12].

Yet, existing implementations of full-matrix preconditioning are currently
impractical given their massive memory costs. Interestingly, although they
have different theoretical justifications, both full-matrix adaptive
preconditioning methods such as GGT [1] and full-matrix natural-gradient
methods such as M-FAC [12] run into the same key practical limitation: each
maintains a “sliding window” of several past gradients, employed in estimating
the preconditioner at each step. Even for small-scale models, maintaining the
gradient history results in unmanageable total cost of $\Theta(md)$ memory,
where $d$ is the model dimension, usually in the millions, and $m$ is the size
of the gradient history, usually of size 100-1000. For example, a standard
recipe for fine-tuning the medium-scale BERT-base model using the M-FAC
preconditioner with standard settings requires more than 450GB of GPU RAM,
unavailable on any single GPU. Thus, while showing promising results in terms
of accuracy, full-matrix preconditioning is currently impractical at scale due
to memory constraints.

#### Contribution.

In this paper, we present a new algorithmic approach for reducing the massive
memory cost of full-matrix preconditioning at the scale of DNNs, called Error
Feedback for Compressed Preconditioning (EFCP), which applies to both adaptive
(Adagrad/GGT [1]) and natural-gradient (M-FAC [12]) full-matrix
preconditioners. In a nutshell, EFCP works by compressing the gradient window
either via high sparsity or extremely low rank. In practice, we are able to
compress the gradient history by one to two orders of magnitude _without
impacting accuracy_ , effectively removing the memory limitations of full-
matrix preconditioning.

Our results are enabled by two technical innovations. On the conceptual side,
the surprising observation behind our work is that the sizable error in the
preconditioner induced by gradient sparsification or low-rank projection can
in fact be handled via an error feedback mechanism applied to the gradients
_before they are stored_ into the gradient window. For both variants, we
provide efficient compression-aware data structures which can leverage
compression for memory savings, and can apply to both M-FAC and GGT.

On the practical side, our main contribution is an efficient GPU
implementation for EFCP in the case of high sparsity. Specifically, this is
enabled by a novel GPU implementation of a _dynamic sparse ring-buffer_ data
structure, in which each element is a sparse gradient tensor, and which
enables efficient sparse multiplication against new gradients, as needed for
preconditioner computations, as well as fast gradient insertions and removals,
as required by the sliding window semantics.

We validate our implementation experimentally on standard vision
(ResNet/ImageNet) and language modeling (BERT/GLUE) tasks. Results show that
EFCP enables us to significantly reduce the memory overhead of full-matrix
preconditioning, to the point where it is in the same range as the memory
required by standard optimizers such as SGD with Momentum or Adam.

We show this extensively by testing the EFCP implementation in the context of
the M-FAC optimizer, compressing gradients to $\sim 99\%$ sparsity (a variant
we call Sparse MFAC), as well as via smaller-scale experiments for low-rank
compression. Moreover, we show that our implementation is also effective for
the GGT full-matrix Adagrad optimizer. For example, for the standard ResNet-18
model [16] trained from scratch on ImageNet [30], the Sparse MFAC variant with
full-matrix preconditioning uses a maximum of 22.5 GB of GPU RAM, compared to
21.3 GB used by SGD with momentum with the same parameters. At the same time,
Sparse MFAC outperforms the well-tuned SGD baseline by approximately 1% Top-1
test accuracy, recovering the accuracy of Dense M-FAC in the same setup (which
uses more than 56GB GPU RAM in this setting). We provide additional validation
for Sparse M-FAC in the context of training ViT [9] and BERT [8] models as
well, and perform ablations w.r.t. sparsity and model size.

In summary, for the first time, EFCP enables experimenting with full-matrix
preconditioner methods within the memory capacity of a single GPU in
reasonable time, with negligible accuracy loss relative to full
preconditioning.

##  2 Background and Related Work

###  2.1 General Setting

We consider a standard classification setting, in which we are given a
training dataset $\mathcal{D}=\\{(x_{i},y_{i})_{i=1}^{N}\\}$ containing $N$
i.i.d. samples $x_{i}$ and their associated labels $y_{i}$. We suppose input
pairs $(x,y)$ are drawn from a true, unknown distribution $p_{true}\triangleq
p_{t}$ having the density $p_{t}(x,y)=p_{t}(x)p_{t}(y|x)$ and our model
$\theta\in\mathbb{R}^{d}$ parameterizes a probabilistic distribution
$p_{\theta}(y|x)\triangleq p(y|f(x,\theta))$ with density
$p_{\theta}(x,y)=p_{t}(x)p_{\theta}(y|x)$. Our aim is to minimize the
“difference” between the true conditional distribution $p_{t}(y|x)$ and the
modelled conditional distribution $p_{\theta}(y|x)$, measured in some
appropriate metric. If we pick this to be the KL divergence between the two
distributions, we obtain the following loss function:

$KL(p_{t}(x,y)||p_{\theta}(x,y))=\int
p_{t}(x,y)\log\frac{p_{t}(x,y)}{p_{\theta}(x,y)}dxdy=\mathbb{E}_{p_{t}(x)}KL(p_{t}(y|x)||p_{\theta}(y|x))]$
(1)

In practice, we only have a finite number of samples $x\sim p_{t}(x)$ and no
information about the density function. In this case, we are using the
empirical training distribution $\hat{p}_{t}$ instead of $p_{t}$, which is
given by the samples in $\mathcal{D}$ and thus we can define the loss as

$L(\theta)\triangleq
KL(p_{t}(x,y)||p_{\theta}(x,y))=\mathbb{E}_{\hat{p}_{t}(x)}[KL(p_{t}(y|x)||p_{\theta}(y|x))]\approx-\frac{1}{N}\sum_{i=1}^{N}\log
p_{\theta}(y_{i}|x_{i}),$ (2)

which is the standard objective function minimized for the maximum likelihood
estimation. The loss $L:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is assumed to be
continuous and twice differentiable.

In practice, the loss $L(\theta)$ defined above is often minimized using
variants of Stochastic Gradient Descent (SGD)
$\theta_{t+1}\leftarrow\theta_{t}-\eta_{t}g_{t}$, where
$g_{t}\triangleq\nabla_{\theta}L(\theta_{t})$ and $\eta_{t}\in\mathbb{R}$ is
the learning rate. However, theoretical results suggest that convergence and
end-results can be improved by _preconditioning_ this gradient with a matrix
$C$ which incorporates information about the geometry of the loss landscape,
leading to the parameter update rule

#### Natural Gradient.

A popular approach for estimating a good preconditioner is Natural Gradient
[4], which adapts to the information geometry of the loss landscape by
measuring the discrepancy between the true distribution and the modelled
distribution in terms of KL-divergence in the distribution space, as opposed
to using the Euclidean distance as is the case for regular SGD. Then, the
preconditioner $C$ is the Fisher Information Matrix (FIM), defined as
$F(\theta)\triangleq\frac{1}{N}\sum_{i=1}^{N}\mathbb{E}_{p_{\theta}(y|x_{i})}[\nabla_{\theta}\log
p_{\theta}(y|x_{i})\nabla_{\theta}\log p_{\theta}(y|x_{i})^{T}].$ In this
context, one can show that the FIM is the expected Hessian with respect to the
log-likelihood. Since the exact FIM is difficult to compute due to the
expectation term in its definition, different efficient approximations of it
exist in the literature. For neural networks, K-FAC [13, 25] provides a
diagonal and tri-diagonal approximation of FIM by sampling a limited number of
$y$ values from $p_{\theta}(y|x_{i})$ in order to compute the expectation.

Another popular approximation to the FIM is the Empirical Fisher, defined as
$\hat{F}(\theta)\triangleq\frac{1}{N}\sum_{i=1}^{N}\nabla_{\theta}\log
p_{\theta}(y_{i}|x_{i})\nabla_{\theta}\log p_{\theta}(y_{i}|x_{i})^{T}$.
Instead of sample $y$’s from model’s predictive distribution, it uses the
associated label $y_{i}$. It is important to note that the empirical Fisher is
a good approximation of the FIM if $p_{t}(y|x)\approx p_{\theta}(y|x)$, e.g.
when the model has good accuracy. The M-FAC method [12] (Section 2.2) provides
an efficient implementation for this approach.

#### Adaptive Regularization.

A different justification for obtaining good preconditioners is via _adaptive
regularization_ , arising from online convex optimization. Broadly, the idea
is to adapt the learning rate corresponding to each parameter in order to
minimize regret, defined as the cumulative difference between the optimizer’s
performance and the best decision in hindsight, by taking into account the
history of gradients observed up to a given point. By and large, the community
has focused on diagonal preconditioners, which allow computationally-efficient
and extremely popular algorithms, such as Adam [18] and AdaGrad [10]. The GGT
result [1] (Section 2.2) showed that results can be improved by taking
“second-order” weight correlations into account, leading to a full-matrix,
$d\times d$ preconditioner, and provided an efficient approximation of this
preconditioner.

###  2.2 Efficient Full-Matrix Preconditioners

#### Overview.

We now provide a detailed description of the M-FAC full-matrix preconditioner
[12], which will serve as the main illustration of our preconditioner
compression method in the paper body. A fully-worked-out example for the GGT
preconditioner [1] is also provided in the Appendix. A common point of both
methods is that they store a history of the past $m$ gradients in the matrix

$G=[g_{t-(m-1)}^{T},g_{t-(m-2)}^{T},...,g_{t-1}^{T},g_{t}^{T}]\in\mathbb{R}^{m\times
d},$ (3)

which can be represented as a ring buffer, updated at each step by replacing
the oldest gradient $g_{t-(m-1)}$ with the most recently-generated one
$g_{t}$. (The main difference between M-FAC and GGT comes from the way the
ring buffer $G$ is used to approximate the preconditioner.) To preserve
performance, the buffer $G_{t}$ must be stored in the GPU memory. For best
accuracy, both methods require storing a large window of $m$ gradients, where
$m$ is usually between $100$ and $1000$, each of which are of size $d$ (the
model dimension). Thus, the memory cost of each method becomes intractable in
large-scale settings. Next, we briefly describe the algorithmic implementation
of M-FAC.

#### The M-FAC Optimizer.

M-FAC provides an efficient estimator for Inverse-Hessian-Vector-Products
(IHVPs) for the special case where the Hessian is approximated via the
Empirical Fisher approximation
$\hat{F}=\frac{1}{m}\sum_{i=1}^{m}g_{i}g_{i}^{T}$, where the gradients
$g_{t}\triangleq\nabla_{\theta}L(\theta_{t})$ are stored in the matrix $G$,
defined above. The first idea is to use the classic Woodbury-Sherman-Morrison
(WSM) formula to integrate new gradient information as rank-1 updates into the
current inverse. This leads to the following recursion to compute the IVHPs
with an arbitrary vector $x\in\mathbb{R}^{d}$ as

$\hat{F}_{t+1}^{-1}x=\left(\hat{F}_{t}+\frac{1}{m}g_{t}g_{t}^{T}\right)^{-1}x\\\
\stackrel{{\scriptstyle
WSM}}{{=\joinrel=\joinrel=}}\hat{F}_{t}^{-1}x-\frac{\hat{F}_{t}^{-1}g_{t}(\hat{F}_{t}^{-1}g_{t})^{T}}{m+g_{t}^{T}\hat{F}_{t}^{-1}g_{t}}x\\\
=\frac{1}{\lambda}x-\sum_{k=1}^{t}\frac{\hat{F}_{k}^{-1}g_{k}(\hat{F}_{k}^{-1}g_{k}^{T})}{m+g_{k}^{T}\hat{F}_{k}^{-1}g_{k}}x,$
(4)

where $F_{0}=\lambda I_{d}$ and $\lambda>0$ is a damping parameter.

The key technical idea behind M-FAC is the authors’ non-trivial observation
that IHVPs can be expressed as a linear combination between the past gradients
in the sliding window, where the coefficients are efficiently computed based
only on scalar products $g_{i}^{T}g_{j}$ and $g_{i}^{T}x$, leading to the
following expression for the IHVPs as

$\hat{F}_{m}^{-1}x=\frac{1}{\lambda}x-\sum_{k=1}^{m}c_{k}g_{k}.$ (5)

This allows M-FAC to efficiently add or remove gradients from the sliding
window, done by computing the corresponding coefficients $c_{k}$, and then
making the resulting update into the formula for $\hat{F}_{m}^{-1}x$.

A key aspect regarding M-FAC is the interaction with regularization. Recent
work [40] explains that weight decay and $L_{2}$-regularization coincide for
SGD, but are different for preconditioned methods. Simply inputting
$g_{t}=\nabla_{\theta}L(\theta)+\gamma\theta$ into the M-FAC preconditioner
invalidates the empirical Fisher definition due to the additional term
$\gamma\theta$. Based on this, and differring from M-FAC [12], we will
decouple the regularization such that the empirical Fisher definition holds
and use the update
$\theta_{t+1}=(1-\gamma\eta_{t})\theta_{t}-\eta_{t}\hat{F}_{t}^{-1}\nabla_{\theta}L(\theta_{t})$,
where $\gamma>0$ is the weight decay parameter.

#### Data Structures and Operations.

Next, we briefly summarize the main operations in the original M-FAC optimizer
(denoted next by Dense M-FAC). We detail the challenges that appear when the
internal buffer $G$ is sparse, as well as our GPU-efficient solution, in the
further sections. First, notice that the main operations in M-FAC are the
Scalar Products (SP) and the Linear Combination of Gradients (LCG), which are
needed to express the IHVPs as a linear combination.

#### Scalar Products (SP).

This first operation corresponds to the pairwise scalar products of rows in
$G\in\mathbb{R}^{m\times d}$ stored in an intermediary matrix
$S=GG^{T}\in\mathbb{R}^{m\times m}$, where $S_{ij}=g_{i}^{T}g_{j}$. When the
buffer $G$ is updated (e.g. the oldest gradient is replaced by a new gradient
$g\in\mathbb{R}^{d}$ at row $i$), the $i^{th}$ row and $i^{th}$ column of
matrix $S$ are updated as $S_{i,:}=S_{:,i}=G\cdot g\in\mathbb{R}^{m}$. This
operation can be easily performed on GPU using frameworks such as Pytorch.

#### Linear Combination of Gradients (LCG).

The second main operation of M-FAC is the linear combination of gradients,
having the coefficients vector $c\in\mathbb{R}^{m}$. For a _dense_ buffer $G$,
the LCG is easily computed by multiplying each component $c_{i}$ with the
$i^{th}$ row of $G$ and then computing the sum of all rows in $G$. This
operation can also be easily performed in PyTorch on GPU.

###  2.3 Related Methods

In the Appendix, we provide a full discussion of GGT [1], as well as its
compressed version, which can be seen as a full-matrix version of AdaGrad
optimizer. More generally, obtaining efficient approximations of second-order
(curvature) information in deep learning is an active area. The _diagonal_
approximation of the Hessian [19, 18] is very popular; however, it is known to
provide lower quality relative to both block-wise methods or K-FAC [37, 32,
12].

K-FAC provides a block diagonal approximation of the FIM, and allows efficient
computation of the inverse [6, 29, 39, 37, 21]; however, it is known that its
prerequisites do not always hold [26]. _Hessian-free_ optimization [24]
forgoes the explicit computation of Hessians in favor of computing an IHVP,
but may require several “inner” iterations to converge at each step.
AdaHessian [38] provided an alternative approximation of the inverse Hessian
diagonal, using Hutchinson’s randomized algorithm for estimating the diagonal,
which requires smoothing heuristics, and has at least 2x theoretical iteration
cost vs SGD, which in practice is usually around 3x. Shampoo [14] provides an
efficient heuristic approximation of GGT, reducing memory cost by maintaining
preconditioner estimates per model tensor, using an efficient iterative method
to compute preconditioner inverses.

#### The Error Feedback Mechanism.

The error feedback mechanism has been studied extensively in the context of
gradient compression for standard SGD, in the absence of complex
preconditioning. The first analyses of sparsified SGD with error feedback were
provided by [33, 3] under additional assumptions, while [17, 34, 27] provided
additions to the analysis. More recently, a very complex argument by [23]
analyzed convergence of diagonal Adagrad with error feedback. We show that
their analysis extends to a diagonal version of our algorithm.

##  3 Method

Algorithm 1 EFCP: Error Feedback for Accurate Compressed Full-Matrix
Preconditioning

1:Parameters: $T=$steps count; $m=$ gradient count;

2:$\xi_{0}=0_{d}\in\mathbb{R}^{d}$ $\triangleright$ initialize error feedback

3:for each step $t\in\\{1,2,...T\\}$ do

4: $g_{t}\leftarrow\nabla_{\theta}L(\theta_{t})$ $\triangleright$ get gradient

5: $a_{t}\leftarrow\xi_{t-1}+g_{t}$ $\triangleright$ error feedback

6: $c_{t}\leftarrow\textsc{Compress}(a_{t})$ $\triangleright$ compress the
accumulator

7: $\xi_{t}\leftarrow a_{t}-c_{t}$ $\triangleright$ update the error

8: $u_{t}\leftarrow\mathcal{A}({\color[rgb]{0,0,1}c_{t}})$ $\triangleright$
update $G$ and precondition using $\mathcal{A}$

9: $\theta_{t}\leftarrow\theta_{t-1}-\eta_{t}u_{t}$ $\triangleright$ update
model parameters

10:end for

#### Overview.

Our general approach for compressing the gradient history required by full-
matrix preconditioners is described in Algorithm 1. Conceptually, we assume a
standard gradient-based optimization setup with an oracle $g_{t}$ at each step
$t$, augmented by a preconditioning algorithm $\mathcal{A}$, based on a
gradient sliding window $G$, whose large memory cost we want to reduce. We
propose to do this by feeding _compressed_ gradients into the preconditioner
estimate maintained by algorithm $\mathcal{A}$. Yet, doing so directly by
projecting the gradients onto a sparse or low-rank support diverges.

Instead, we apply _error feedback_ [31, 3, 17] to gradients _before they are
fed into the preconditioner data structure_ , using an error accumulator
vector $\xi_{t}\in\mathbb{R}^{d}$, initially set to zero. At each step, we
obtain the gradient $g_{t}$ of the model parameters w.r.t. the loss, after
which we feed back the error $\xi_{t-1}$ into the current gradient $g_{t}$,
obtaining the _accumulator_ $a_{t}$. Next, the accumulator $a_{t}$ is
compressed, using e.g. sparsity or low-rank, to obtain a compressed
representation $c_{t}$. The compression error $\xi_{t}$ is updated to the
difference between the accumulator $a_{t}$ and its compressed version $c_{t}$.
Finally, the compressed gradient $c_{t}$ is fed directly to the
preconditioning algorithm $\mathcal{A}$.

Notice that, other than the three lines corresponding to error feedback (blue
in the algorithm), the optimization loop is the same as the uncompressed one.
However, we emphasize that the internal buffer of the preconditioned optimizer
will store a sequence of $m$ representations
$[c_{t-(m-1)},c_{t-(m-2)},...,c_{t-1},c_{t}]$, each of which is highly-
compressed, leading to memory savings.

While this blueprint appears simple, two key questions remain: (1) _how can we
leverage the compressed gradient window for actual memory and computational
speedups on complex real-world preconditioners?_ and (2) _why does this
approach preserve convergence?_ We address them below.

###  3.1 Compressing Preconditioners via Sparsity

We now expand on accumulator compression via sparsity, specifically by using
Top-$k$ sparsification and error feedback. Specifically, the preconditioner
data structure will only store $k\ll d$ entries from each accumulated
gradient, where $k$ is around $1\%$ of the full dimension $d$. As we will show
experimentally, this will essentially remove the memory overhead of full-
matrix preconditioning.

We detail our approach for the Sparse M-FAC algorithm; the Sparse GGT variant
is slightly simpler and is described in Appendix D due to space limitations.
For Sparse M-FAC, we describe the changes required by sparsity for the buffer
$G$ and the Scalar Product (SP) and Linear Combination of Gradients (LCG)
operations. A key concern will be re-interpreting the mathematical formulation
in order to be able to leverage fast CUDA kernels, which is notoriously
challenging in the context of sparsity.

#### Sparsity format for $G$.

To store the gradient history $G$, we need a sparse matrix-like data structure
for storing gradients that fulfills the following two properties at the same
time: i) it allows easy replacement of any row (corresponding to adding and
removing gradients) and ii) it allows computing the scalar product (SP) and
linear combination of gradients (LCG) operations efficiently. To the best of
our knowledge, no such data structure is known in the literature.

Specifically, for Sparse M-FAC, we want to input sparse accumulators $c_{t}$
of size $k\ll d$ into the buffer $G$. Instead of storing
$G\in\mathbb{R}^{m\times d}$ densely, we will store
$G=(\mathcal{I},\mathcal{V})$, where $\mathcal{I}\in\mathbb{N}^{m\times k}$
will store the indices returned by Top-$k$ and
$\mathcal{V}\in\mathbb{R}^{m\times k}$ will store the corresponding values.
When the $j^{th}$ row of $G$ is updated with the new indices-values pair
$(i_{j},v_{j})$, we replace the $j^{th}$ row in $\mathcal{I}$ and
$\mathcal{V}$ with $i_{j}$ and $v_{j}$, respectively. Breaking $G$ into two
matrices of indices and values means that $G$ loses the matrix structure in
the dense sense and the SP and LCG operations must be redefined because we
cannot use the basic PyTorch operations for matrix multiplications anymore.
Below we describe our solution for these two operations that involve CUDA
specific kernels to efficiently perform the computations on a GPU for large
models, but first we would like to provide details about the sparsity pattern
which we can leverage to obtain more efficient kernels for our particular
problem.

#### Sparsity Pattern.

One observation is that, due to the algorithm structure, each row in
$G=(\mathcal{I},\mathcal{V})$ has the same sparsity. To leverage this, at each
step we compress the accumulator $a_{t}$ using Top-$k$ applied in blocks of
(large) fixed size $B$. We do this to have selected values uniformly chosen in
order to speed up the LCG operation (details below). We express $k$ as a
percentage of total number of weights $d$ and then convert it to a particular
integer value that represents the density of values used in the Top-$k$ call.
For example, in most of our experiments we will have $k=1\%$, meaning that for
each block of size $B$ we will have roughly $B/100$ non-zero values. This way
we know that for any block of size $B$ we have the same number of weights in a
specific interval. We provide more details about choosing $B$ in the section
where we describe the LCG operation.

#### Format for $\mathcal{V}$.

Starting with the CUDA capability 8.0, the Nvidia GPUs have bfloat16 support,
which is a half precision data type optimized for Deep Learning. We found that
our optimizer still converges when we use bfloat16 instead of float32. This is
particularly useful when automatic mixed precision (AMP) is activated because
it speeds up the overall training and has a lower memory footprint (both for
the model’s state and for the optimizer’s state).

#### Sparse Scalar Products.

This operation is commonly known as sparse-matrix vector multiplication (SPMV)
and several implementations already exist for a few sparsity representations
(e.g. CSR, COO), but unfortunately they cannot be directly used for our
particular choice for the matrix $G=(\mathcal{I},\mathcal{V})$, each of size
$m\times k$. Thus, we can write a kernel that leverages the properties of our
buffer as follows. We use $m$ thread blocks, each containing 1024 threads. One
thread block processes one row and has a shared memory array of 1024 floats to
store the dot product computed on a subset of the row. We perform memory
coalesced reads for the indices and values from the global memory. In the end,
we perform parallel reduction over the shared memory buffer to compute the sum
(the final value for the dot product) in logarithmic complexity. We show our
implementation of the SP operation in the Appendix C.

#### Sparse LCG.

This operation leverages the block sparsity to efficiently compute the linear
combination of gradients. Given that we apply sparsity in blocks of size $B$,
we know that in every block we will always have a fixed number of values (and
indices) selected by Top-$k$ (e.g., for $B=10000$ and $k=1\%$, we will have
$100$ values and indices in any interval $[n\cdot B,(n+1)\cdot
B]\cap\mathbb{N}$, with $n\in\mathbb{N}$. This allows us to use a shared
memory block to accumulate the linear combination of gradients across a slice
of size $m\times B$ from $G=(\mathcal{I},\mathcal{V})$ and write the content
of the shared memory at the correct offsets in the global memory that stores
the result (here, we use the information that the indices are in a specific
interval). Similar to the SP case, we also perform memory coalesced reads for
the indices and values from the global memory only once for each component. We
provide explanations about how we choose the thread blocks dimensions in the
next paragraph and we show our implementation of the LCG operation in the
Appendix C.

#### Thread Blocks for LCG.

Choosing the right number of thread blocks is critical for the performance of
the Sparse LCG, since it needs to balance memory transfer and compute load. To
address this, we provide an automatic procedure for choosing this parameter,
based on the GPU capabilities. Specifically, given a number of threads and
maximum available GPU shared memory size, we determine the total number of
thread blocks to be used to allow for full utilization, avoiding waiting in
the queue at each SM.

#### Efficiency Improvements.

These technical data structure optimizations are critical for good
performance. Relative to a naive Pytorch implementation of the same approach,
which stores the gradients as sparse but performs the operations in dense
format, the above approach is 1.5-2x faster and uses 20-30% less memory.

#### Memory Savings.

Compressing the buffer using sparsity yields $\tilde{O}(mk)$ space complexity,
where $k$ is the target density of the gradients after Top-$k$ compression,
ignoring logarithmic factors due to indexing. In practice, we obtain stable
convergence for $k=d/100$, which translates into practical space savings of
approximately $20$x (relative to GGT) and $30$x (relative to M-FAC). We
provide exact numbers in the experimental section.

###  3.2 Low-Rank Compression of Preconditioners

In Appendix E, we discuss low-rank compression of the gradients used in the
preconditioner. Specifically, this requires re-implementing the Compress step
in Algorithm 1 via a low-rank approximation of the gradient, which we do using
an efficient variant of power iteration [35]. Then, we again consider the
M-FAC algorithm and reformulate its algebraic operations to leverage low-rank
matrices for memory and computational savings.

##  4 Experiments

We now validate our results experimentally. In the main text, we focus on
comparing Sparse M-FAC with other optimizers, as it proved to be the most
practical variant of EFCP. We show results for low-rank M-FAC and Sparse GGT
in the Appendix.

In the following, we validate Sparse MFAC (S-MFAC) experimentally as an
optimizer for training in the context of standard vision and language
modelling tasks. Specifically, we integrate our S-MFAC optimizer in ASDL [28],
a benchmarking library for second-order optimizers, but also examine larger-
scale tasks such as ImageNet training of ResNet-18 and compare S-MFAC with
AdamW and Dense M-FAC (D-MFAC) on BERT models on GLUE tasks. We report
training loss, test accuracy and total running times and top memory usage for
the entire training process.

#### Notations.

We use the notation $E$ for number of epochs, $\eta$ for learning rate,
$\gamma$ for weight decay, $B$ for batch size, $m$ for the number of gradients
(sliding window size), $k=1\%$ the gradient density. We use the prefixes D for
Dense and S for Sparse in the context of adaptive optimizers M-FAC. For the
gradient window size $m$, we use the standard settings for D-MFAC [12], which
is $m=1000$ unless otherwise specified.

#### The ASDL Benchmark.

In the first experiment, we integrate S-MFAC into the ASDL benchmarking
framework for second-order optimizers [28], for the task of fine-tuning a
ViT-T model pre-trained on ImageNet. We report the test accuracy for the run
that has largest validation accuracy. Table 1 shows the test accuracy of our
S-MFAC optimizer compared to the other optimizers reported in the paper. (We
had to remove P-SGD from the comparison due to errors raised by the
implementation during experiments.) We observe that S-MFAC is competitive with
other preconditioned gradient methods, outperforming SGD, AdamW, K-FAC and
SENG, and matching the highest accuracy obtained on this task (98%) with
affordable time and memory overheads relative to SGD.

Table 1:  Test accuracy of S-MFAC for ResNet-18 and ViT-Tiny on CIFAR-10. The accuracies for the other optimizers are obtained from ASDL [28]. We also report relative time and memory overheads w.r.t. SGD. Optimizer | ViT-T | Time | Memory
---|---|---|---
SGD | 97.8 | $1\times$ | $1\times$
AdamW | 97.9 | $1\times$ | $1.01\times$
K-FAC(1mc) | 97.4 | $1.12\times$ | $1.13\times$
SENG | 97.7 | $35\times$ | $7.31\times$
Shampoo | 98.0 | $1.6\times$ | $1.07\times$
S-MFAC | 98.0 | $1.63\times$ | $1.15\times$

#### ImageNet/ResNet-18.

Next, we move to a more complex vision task, by training ResNet-18 on ImageNet
[7] from scratch using a highly-tuned version of SGD (with momentum) from the
FFCV [22] library, as well as S/D-MFAC. All hyper-parameters are provided in
the Appendix A. Figure 1 shows the validation accuracy curves for this
experiment, while Figure 2 shows the best achieved training loss and
validation accuracy for each method, including a detailed sparsity sweep below
$1\%$ density for S-MFAC. (The loss curves show similar trends, and are
provided in the Appendix.)

We begin by noting that SGD obtains training loss $2.237$ with accuracy
$69.12\%$. D-MFAC reaches $2.164$ loss and significantly higher Top-1
validation accuracy, $70.03\%$. Despite slightly higher training loss than
D-MFAC ($2.185$), S-MFAC has comparable $70.01\%$ validation accuracy, which
is very strong for this model and training recipe. Thus, even for this
challenging task, compressing the preconditioner essentially recovers the
accuracy of the dense M-FAC optimizer. When we examine the behavior of S-MFAC
as we decrease gradient density, we observe that, remarkably, S-MFAC with
$k=0.17\%$ (only $\approx 20k$ entries!) still achieves comparable accuracy to
SGD, and that accuracy decreases gradually. However, the resulting
preconditioned gradients are not necessarily very sparse.

To examine memory, we ran the experiments on a single A100 GPU and recorded
the max memory usage throughout the entire process via the nvidia-smi tool. We
find that SGD uses $12.8$ GB and S-MFAC $14.3$ GB, while D-MFAC in the same
setting needs $59$ GB.

Figure 1:  Validation (Top-1) accuracy for ResNet18 training from scratch on
ImageNet-1K. M-FAC and 99%-Sparse M-FAC reach approximately 1% higher
validation accuracy relative to SGD using the well-tuned FFCV recipe [22].

Figure 2:  Training loss and validation (Top-1) accuracy for ResNet18 training from scratch on ImageNet-1K under a detailed sparsity sweep, between $1\%$ and $0.04\%$ density. Dense M-FAC outperforms SGD in this setting by about $1\%$, which is recovered by S-MFAC at 1% density. Even with $0.17\%$ gradient density, M-FAC slightly outperforms SGD in terms of training loss and test accuracy.  Density $k$ | Train Loss | Test Accuracy
---|---|---
D-MFAC (100%) | 2.164 | 70.03%
117k (1.00%) | 2.185 | 70.01%
100k (0.85%) | 2.183 | 70.18%
75k (0.64%) | 2.187 | 69.99%
50k (0.42%) | 2.190 | 69.68%
25k (0.21%) | 2.206 | 69.43%
20k (0.17%) | 2.217 | 69.26%
15k (0.12%) | 2.230 | 69.05%
10k (0.08%) | 2.246 | 68.84%
5k (0.04%) | 2.300 | 68.27%
SGD | 2.237 | 69.12%

#### BERT/GLUE-MNLI.

Finally, we test our S-MFAC implementation for BERT-Tiny/Mini/Base models on
the MNLI task from the GLUE benchmark [36], comparing against the Adam
optimizer with optimized default settings from Hugging Face [8], which
reproduce or improve upon the original paper [8]. Dense M-FAC with a window
size of $m=1024$ can only be executed on the tiny and mini model variants,
requiring $18$ GB and $45$ GB memory, respectively, and runs out of memory
(OOM) for the base model (the required memory usage for the ring buffer $G$
alone is 417GB).

The results in Table 2 show that our S-MFAC implementation not only recovers
the performance of D-MFAC on BERT-tiny/mini, but also allows us to run on
BERT-base, with superior validation accuracy relative to Adam using standard
hyper-parameters (e.g., learning rate $\eta=2e-5$). For the tiny model S-MFAC
achieves almost $2\%$ higher accuracy compared to D-MFAC and close to $5\%$
higher validation accuracy compared to Adam. The former improvement can be
justified by the regularization effect of sparsity relative to D-MFAC, as this
is a fine-tuning task, whereas the latter is due to leveraging second-order
information.

At the same time we observe that, accuracy differences become smaller as model
size increases. This behavior is also reported in the original D-MFAC [12]
optimizer for the same GLUE/MNLI task. While this is another indication that
the sparse implementation preserves the behavior of the dense counterpart, it
suggests that the benefits of second-order information diminish in the case of
over-parametrized fine-tuning tasks.

Table 2:  Evaluation accuracy for BERT-Tiny/Mini/Base on GLUE-MNLI. OOM stands for Out of Memory error. | T | M | B
---|---|---|---
Adam | 66.86 | 74.77 | 84.6
S-MFAC | 71.06 | 76.40 | 84.86
D-MFAC | 69.26 | 76.39 | OOM

#### Running times and memory usage.

Finally, we examine running times and memory usage for the algorithm. The
results across our experiments are shown in Table 3. We note that, in the
current implementation, both D-MFAC and our optimizer have a constant
additional cost relative to the highly-optimized Adam or SGD implementations
in Pytorch, due to the additional data structures employed. This cost can be
reduced via optimizations.

Specifically, our current implementation was primarily optimized for residual
networks. This is reflected in the runtime and memory numbers: on ResNet-18
($11.7$M params), D-MFAC requires $4.6\times$ more memory compared to SGD with
momentum and $25\%$ larger running time. Our S-MFAC implementation diminishes
these costs, bringing both the running time and memory usage to be almost
negligible compared to SGD. Specifically, the memory usage is $25\%$ of
D-MFAC.

For BERT-tiny ($4$M params), D-MFAC uses $14\times$ more memory and BERT-mini
($11$M params) uses $15\times$ more memory compared to S-MFAC. Further, we
find that EFCP is extremely effective for large models like BERT-Base ($110$M
params). In this case, running D-MFAC would raise an out of memory (OOM) error
(just the dense ring buffer used to store the gradient window requires 417 GB
of memory under standard parameters).

Although our primary concern in experiments is memory cost, we observe that,
in terms of running time, our algorithm has essentially no overhead on
ResNet-18. On BERT models, there is a notable runtime overhead, which is due
in part to the fact that our CUDA kernels are not specifically adapted to BERT
models. We expect this overhead to be reduced via tuning and further
optimizations, which we plan to investigate in future work. However, we note
that these single-GPU runtimes are very manageable already.

To decrease memory usage and running time, one can try reducing the number of
gradients $m$. However, it is clear from prior work [1, 12] that large values
of $m$ are needed to provide better accuracy results, and specifically to
improve over standard baseline. This is theoretically motivated in part by the
fact that the rank of matrix $G$ is at most $m$, which can be extremely low
compared to the rank of Hessian matrix for large models and the approximation
would be very poor for low values of $m$. Thus, large models require large
number of gradients $m$, which renders our technique effective. To our
knowledge, our approach allows full-matrix adaptive optimizers to be run on a
consumer GPU for the first time. Appendix B contains a detailed theoretical
analysis for memory savings.

Table 3:  Running times and memory usages for ImageNet/ResNet-18 and GLUE/BERT reported on an NVIDIA A100 GPU with 82GB RAM. OOM stands for Out Of Memory. Model | Optimizer | Memory | Training Time
---|---|---|---
| SGD | 12.8 GB | 4h 10m
ResNet-18 | S-MFAC | 14.3 GB | 4h 21m
| D-MFAC | 58.5 GB | 5h 33m
| Adam | 0.8 GB | 11m
BERT-Tiny | S-MFAC | 1.3 GB | 1h 5m
| D-MFAC | 17.9 GB | 1h 20m
| Adam | 1.3 GB | 21m
BERT-Mini | S-MFAC | 2.9 GB | 1h 5m
| D-MFAC | 44.9 GB | 1h 50m
| Adam | 7.1 GB | 1h 45m
BERT-Base | S-MFAC | 14.8 GB | 3h 20m
| D-MFAC | OOM | -

##  5 Conclusions and future work

We have provided a versatile new tool for addressing the memory costs of
adaptive full-matrix preconditioning, in the form of a series of methods which
compress the gradient history, by leveraging an instance of the error feedback
mechanism. We complement this mechanism via efficient sparse/low-rank
algorithms for maintaining existing preconditioners, and have shown
experimentally that our approach can essentially provide lossless compression
when applied via sparsity, with remarkably large compression rates. We provide
CUDA implementations for a sparse data structure that 1) allows easy
replacement of any row (adding/replacing new compressed representations of
error feedback accumulators and 2) allows computing the SP and LCG operations
efficiently which, to the best of our knowledge, is the first such data
structure known in the literature.

In future work, we plan to further investigate theoretical justification for
our method. Obtaining general bounds appears extremely challenging, as even in
the case of standard SGD understanding error feedback required significant
technical effort, e.g. [3, 17]. Another direction of future work is exploring
larger-scale validation of our method for additional architectures and tasks,
and for additional preconditioning mechanisms.

## Acknowledgements

The authors thank Adrian Vladu, Alexandra Peste, Mher Safaryan and Elias
Frantar for their valuable feedback, the IT department from Institute of
Science and Technology Austria for the hardware support and Weights and Biases
for the infrastructure to track all our experiments.

## References

  * [1] 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.
  * [2] Naman Agarwal, Brian Bullins, and Elad Hazan.  Second-order stochastic optimization for machine learning in linear time, 2016.
  * [3] Dan Alistarh, Torsten Hoefler, Mikael Johansson, Nikola Konstantinov, Sarit Khirirat, and Cédric Renggli.  The convergence of sparsified gradient methods.  In Advances in Neural Information Processing Systems, pages 5973–5983, 2018.
  * [4] S.-I. Amari.  Natural gradient works efficiently in learning, 1998.
  * [5] Shun-ichi Amari.  Natural gradient works efficiently in learning.  Neural Computation, 10(2):251–276, 1998.
  * [6] Jimmy Ba, Roger Grosse, and James Martens.  Distributed second-order optimization using kronecker-factored approximations.  https://openreview.net/forum?id=SkkTMpjex, 2016.
  * [7] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
  * [8] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding, 2018.
  * [9] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  In Proceedings of the Ninth International Conference on Learning Representations, 2021.
  * [10] 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.
  * [11] Matthias Fey et al.  Pytorch-sparse library.  https://github.com/rusty1s/pytorch_sparse.
  * [12] Elias Frantar, Eldar Kurtic, and Dan Alistarh.  M-fac: Efficient matrix-free approximations of second-order information, 2021.
  * [13] Roger Grosse and James Martens.  A kronecker-factored approximate fisher matrix for convolution layers, 2016.
  * [14] Vineet Gupta, Tomer Koren, and Yoram Singer.  Shampoo: Preconditioned stochastic tensor optimization.  In International Conference on Machine Learning, pages 1842–1850. PMLR, 2018.
  * [15] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition, 2015.
  * [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] Sai Praneeth Karimireddy, Quentin Rebjock, Sebastian U Stich, and Martin Jaggi.  Error feedback fixes SignSGD and other gradient compression schemes.  In Proceedings of the Thirty-sixth International Conference on Machine Learning, pages 3252–3261, 2019.
  * [18] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization, 2014.
  * [19] Shankar Krishnan, Ying Xiao, and Rif A Saurous.  Neumann optimizer: A practical optimization algorithm for deep neural networks.  arXiv preprint arXiv:1712.03298, 2017.
  * [20] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton.  Cifar-10 (canadian institute for advanced research).  https://www.cs.toronto.edu/ kriz/cifar.html, 2009.
  * [21] César Laurent, Thomas George, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent.  An evaluation of fisher approximations beyond kronecker factorization, 2018.
  * [22] Guillaume Leclerc, Andrew Ilyas, Logan Engstrom, Sung Min Park, Hadi Salman, and Aleksander Madry.  FFCV: Accelerating training by removing data bottlenecks.  https://github.com/libffcv/ffcv/, 2022.
  * [23] Xiaoyun Li, Belhal Karimi, and Ping Li.  On distributed adaptive optimization with gradient compression.  arXiv preprint arXiv:2205.05632, 2022.
  * [24] 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.
  * [25] James Martens, Jimmy Ba, and Matt Johnson.  Kronecker-factored curvature approximations for recurrent neural networks.  In International Conference on Learning Representations, 2018.
  * [26] James Martens and Roger Grosse.  Optimizing neural networks with kronecker-factored approximate curvature, 2015.
  * [27] Giorgi Nadiradze, Ilia Markov, Bapi Chatterjee, Vyacheslav Kungurtsev, and Dan Alistarh.  Elastic consistency: A general consistency model for distributed stochastic gradient descent.  arXiv preprint arXiv:2001.05918, 2020.
  * [28] Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler.  Asdl: A unified interface for gradient preconditioning in pytorch, 2023\.
  * [29] 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.
  * [30] 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\.
  * [31] Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu.  1-bit stochastic gradient descent and its application to data-parallel distributed training of speech DNNs.  In Fifteenth Annual Conference of the International Speech Communication Association, 2014.
  * [32] Sidak Pal Singh and Dan Alistarh.  Woodfisher: Efficient second-order approximation for neural network compression, 2020.
  * [33] Sebastian U Stich, Jean-Baptiste Cordonnier, and Martin Jaggi.  Sparsified sgd with memory.  Advances in Neural Information Processing Systems, 31, 2018.
  * [34] Sebastian U Stich and Sai Praneeth Karimireddy.  The error-feedback framework: Better rates for sgd with delayed gradients and compressed communication.  arXiv preprint arXiv:1909.05350, 2019.
  * [35] Thijs Vogels, Sai Praneeth Karimireddy, and Martin Jaggi.  PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization.  In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * [36] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman.  GLUE: A multi-task benchmark and analysis platform for natural language understanding.  In Proceedings of the Seventh International Conference on Learning Representations, 2019.
  * [37] Chaoqi Wang, Roger Grosse, Sanja Fidler, and Guodong Zhang.  Eigendamage: Structured pruning in the kronecker-factored eigenbasis, 2019\.
  * [38] 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.
  * [39] Wenyuan Zeng and Raquel Urtasun.  MLPrune: Multi-layer pruning for automated neural network compression, 2019.
  * [40] Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse.  Three mechanisms of weight decay regularization, 2018.

##  Appendix A Hyper-parameters

We restate the notations that we introduced in the main body of our work. We
use notation $\mu$ for momentum, $\eta$ for the initial learning rate,
$\gamma$ for weight decay, $\lambda$ for M-FAC dampening, $k$ for gradient
density, $\rho$ for rank, $E$ for epochs, $B$ for batch size.

###  A.1 ASDL

We perform hyper-parameter tuning for CIFAR-10 and ViT-Tiny, as described in
the appendix section of the ASDL [28] paper. Concretely, the general grids we
use are $\\{32,128,512,2048\\}$ for the batch size,
$\\{{3e-1,1e-1,3e-2,1e-2,3e-3,1e-3}\\}$ for the learning rate, $E=20$ epochs,
we clip the preconditioned gradient norm to $10$ and perform grid search for
the dampening parameter $\lambda$ over the grid
$\\{1e-2,5e-3,1e-3,5e-4,1e-4,5e-5,1e-5,5e-6,1e-6\\}$ for S-MFAC. In
particular, ViT-Tiny uses $\gamma=1e-4$.

###  A.2 FFCV

For the ImageNet dataset and ResNet-18 model, we use $E=88$ epochs, batch size
$B=1024$, linear decay schedule for the learning rate $\eta$. In particular,
we set the weight decay $\gamma=1e-5$ for convolution and fully connected
parameters and $\gamma=0$ for batch-normalization parameters, thus fixing the
issue in the main FFCV repository. We use linear learning rate decay for all
three optimizers and disable image rescaling. We present the final hyper-
parameters in Table 6.

###  A.3 BERT/GLUE

We use the HuggingFace repository and plug in our D/S-MFAC optimizers. We
compare against default Adam baselines as follows. For GLUE/MNLI we use
$E=3,\eta=2e-5,\gamma=0$. D-MFAC uses the default parameters from the original
paper, e.g. $\eta=1e-4,\lambda=1e-6,\gamma=0$ and we adopt the same values for
S-MFAC for BERT-Tiny/Mini. For BERT-Base we use $\lambda=5e-5$ for S-MFAC.

###  A.4 CIFAR-10/RN-20

For D/S-GGT we use window size $m=100,k=1\%$ and the grid in Table 4. We show
the corresponding pairs $(\eta,\epsilon)$ for the best runs extracted from the
search grids. In these experiments we use $E=164,B=128$ and decay the learning
rate by a factor of $\gamma_{\eta}=0.1$ at 50% and 75% of training. The final
hyper-parameters for each method are summarized in table 5.

Table 4:  Search grid for learning rate $\eta$ and dampening $\epsilon$ for S/D-GGT. $\eta,\epsilon$ | 1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1
---|---
$\eta_{D},\epsilon_{D}$ | loss: (1, 1e-6), accuracy: (1, 1e-5)
$\eta_{S},\epsilon_{S}$ | loss: (1, 1e-6), accuracy: (1, 1e-5)
Table 5:  Final hyper-parameters for CIFAR-10 / ResNet-20. SGD | $\eta=0.1,\mu=0.9,\gamma=5e-4$
---|---
D-GGT | $m=100,\eta=0.1,\epsilon=1e-5,\gamma=0$
S-GGT | $m=100,\eta=0.1,\epsilon=1e-5,\gamma=0$
D-MFAC | $m=1024,\eta=1e-3,\lambda=1e-6,\mu=0,\gamma=1e-4$
S-MFAC | $k=1\%,m=1024,\eta=1e-3,\lambda=1e-4,\mu=0,\gamma=1e-4$
LR-MFAC | $\rho=4,m=1024,\eta=1e-3,\lambda=1e-6,\mu=0,\gamma=0$
Table 6:  Hyper-parameters for SGD and S/D-MFAC for FFCV/ImageNet. SGD | $\eta=0.5,\mu=0.9$
---|---
D-MFAC | $\eta=1e-3,\lambda=1e-6,\mu=0$
S-MFAC |  $\eta=1e-3,\lambda=1e-7,\mu=0$, $k=1\%$

##  Appendix B Memory savings.

In this section we detail the theoretical memory savings for the sparse
preconditioners by computing the total memory in bytes and the ratio of dense
over sparse preconditioners to describe de memory savings.

###  B.1 Sparse M-FAC vs Dense M-FAC

For M-FAC we will consider that $m=1024$ and $k=d/100$.

Dense M-FAC stores a matrix $G\in\mathbb{R}^{m\times d}$ of 32 bit floats,
having a total of $4md$ bytes, resulting in a theoretical memory footprint of
$4096d$ bytes.

On the other hand, Sparse M-FAC stores multiple tensors:

  * •

$\mathcal{I}\in\mathbb{R}^{m\times k}$ of 32 bit integers for the indices

  * •

$\mathcal{V}\in\mathbb{R}^{m\times k}$ of 32 bit floats or 16 bit bfloat16s
for the values

  * •

$\xi\in\mathbb{R}^{d}$ of 32 bit floats for the error feedback

  * •

$r_{m}\in\mathbb{R}^{m}$ of 32 bit floats to store the SP result

  * •

$r_{d}\in\mathbb{R}^{d}$ of 32 bit floats to store the LCG result

When $\mathcal{V}$ holds 16 bit bfloat16s, the memory footprint for
$\mathcal{I},\mathcal{V},\xi,r_{m},r_{d}$ altogether is
$4mk+2mk+4d+4m+4d=0.06md+8d+4m=61.44d+8d+4m=69.44d+4m\approx 70d$.

When $\mathcal{V}$ holds 32 bit floats, the memory footprint for
$\mathcal{I},\mathcal{V},\xi,r_{m},r_{d}$ altogether is
$4mk+4mk+4d+4m+4d=0.08md+8d+4m=81.92d+8d+4m=89.92d+4m\approx 90d$.

In the end, the memory savings for M-FAC is $\frac{4096d}{90d}\approx
45.5\times$ for float and $\frac{4096d}{70d}\approx 58.5\times$ for bfloat16.

###  B.2 Sparse GGT vs Dense GGT

For GGT we will consider that $m=100$ and $k=d/100$.

Dense GGT stores a matrix $G\in\mathbb{R}^{m\times d}$ of 32 bit floats,
having a total of $4md$ bytes, resulting in a theoretical memory footprint of
$400d$ bytes.

Sparse GGT stores the following tensors:

  * •

$\mathcal{I}\in\mathbb{R}^{m\times k}$ of 32 bit integers for the indices

  * •

$\mathcal{V}\in\mathbb{R}^{m\times k}$ of 32 bit floats for the values

  * •

$\xi\in\mathbb{R}^{d}$ of 32 bit floats for the error feedback

  * •

$r_{m}\in\mathbb{R}^{m}$ of 32 bit floats to store the SP result

The theoretical memory footprint for $\mathcal{I},\mathcal{V},\xi,r_{m}$
altogether is $4mk+4mk+4d+4m=8mk+4d+4m=0.08md+4d+4m=8d+4d+4m=12d+4m\approx
12d$. In the end, the memory savings for GGT is $\frac{400d}{12d}\approx
33.3\times$.

##  Appendix C CUDA Kernels

In this section we describe our efficient CUDA kernels for large models in
algorithmic form. We reiterate that M-FAC has two operations that have to be
reimplemented when using a sparse ring buffer $G=(\mathcal{I},\mathcal{V})$,
the scalar products (SP) and linear combination of gradients (LCG). In the
next paragraphs we provide a brief algorithmic description for SP and LCG.

###  C.1 SP Kernel

For this operation we have to compute the scalar product between each row in
$G$ and a general gradient $g$. The sparse matrix $G$ is splitted into two
matrices, the indices matrix $\mathcal{I}\in\mathbb{N}^{m\times k}$ and the
values matrix $\mathcal{I}\in\mathbb{R}^{m\times k}$. We use 1024 thread
blocks, each containing 1024 threads and one thread block will process one
row. In the algorithm we will consider that $\mathcal{I}$ and $\mathcal{V}$
are 2-dimensional, but in practice they are 1-dimensional. We use
$\omega\in\mathbb{R}^{m}$ to store result of the SP operation in Algorithm 2

Algorithm 2 SP kernel (CUDA)

1:Input: $m,k,g,\mathcal{I},\mathcal{V},\omega$

2:$mem$ is a shared memory buffer with $T$ components

3:$\sigma\leftarrow 0$ $\triangleright$ stores the partial dot product in
current thread

4:$Bid\leftarrow$ block id $\triangleright$ gives the row index in
$\mathcal{I}$ and $\mathcal{V}$

5:$Tid\leftarrow$ thread id

6:$T\leftarrow$ number of threads in the block

7:$r\leftarrow Bid$ $\triangleright$ row index

8:for $j\in\\{T_{id}<k,T\\}$ do $\triangleright$ $j$ from $T_{id}$ to $k$ with
step $T$ (coalesced memory access)

9: $i\leftarrow\mathcal{I}_{r,j}$

10: $v\leftarrow\mathcal{V}_{r,j}$

11: $\sigma\leftarrow\sigma+v\cdot g_{i}$ $\triangleright$ accumulate partial
dot product

12:end for

13:$mem_{Tid}\leftarrow\sigma$

14:parallel_reduce(mem, T) $\triangleright$ logarithmic sum: result is in
$mem_{0}$

15:if $Tid=0$ then

16: $\omega_{r}\leftarrow mem_{0}$

17:end if

###  C.2 LCG Kernel

The LCG operation is more complex than the SP operation and requires providing
more details. We will start from the compression operator from line 6 of
Algorithm 1, which actually applies Top-$k$ on the accumulator $a_{t}$,
returning a compressed representation $c_{t}$ (not to be confused with the
vector $c$ that stores the coefficients of the linear combination) - the
$c_{t}$ is placed on a specific row of $G=(\mathcal{I},\mathcal{V})$ at each
iteration $t$.

The Top-$k$ operator is applied in blocks of size $B_{d}$ (the subscript $d$
means that this block is part of a $d$-dimensional vector), which is chosen
automatically to be the maximum number of floating point values that can fit
in the entire shared memory space for a thread block. For example, the A100
GPU from Nvidia has CUDA capability 8.0 and the maximum amount of shared
memory per block is 163KB, meaning that we can store $B_{d}=41728$ float32
values in the shared memory. The value $k$ is inputted as a percentage at
runtime and represents the density percentage of $c_{t}$. We convert this
percentage to an actual density in each block $B_{d}$. For example, for
$k=1\%$ we will have around $B_{k}=B_{d}/100$ (we will further skip some
implementation details).

The key detail of our LCG kernel is we use the shared memory of size $B_{d}$
to accumulate partial linear combinations over slices of size $(m,B_{d})$ of
the sparse matrix $G$ and then write the result to the output
$\omega\in\mathbb{R}^{d}$. This is possible because we apply Top-$k$ operator
in blocks of size $B_{d}$ and we know the ends of the interval where we write
the output to. Each thread block can process multiple slices of size
$(m,B_{d})$, depending on how large the model is. The pseudocode of the LCG
kernel is provided in Algorithm 3.

Algorithm 3 LCG kernel (CUDA)

1:Input: $B_{d},B_{k},d,m,k,c,\mathcal{I},\mathcal{V},\omega$

2:$B\leftarrow$ total number of thread blocks

3:$Bid\leftarrow$ block id $\triangleright$ gives the row index in
$\mathcal{I}$ and $\mathcal{V}$

4:$Tid\leftarrow$ thread id

5:$T\leftarrow$ number of threads in the block

6:$mem$ is a shared memory buffer with $T$ components

7:$mem\leftarrow 0$ $\triangleright$ initialize shared memory (the LCG
accumulator) with zero

8:synchronize threads $\triangleright$ wait for all threads to finish
initializing

9:$n_{B_{d}}\leftarrow\lceil\frac{d}{B_{d}}\rceil$ $\triangleright$ how many
blocks of size $B_{d}$ we have

10:$workload\leftarrow\lceil\frac{n_{B_{d}}}{B}\rceil$ $\triangleright$ how
many blocks of size $B_{d}$ each thread block processes

11:for $i\in\\{0<workload\\}$ do $\triangleright$ iterate through all slices
$(m,B_{k})$ (workload) for the curent thread block

12: $start_{d}\leftarrow B_{d}\cdot B_{id}\cdot workload+B_{d}\cdot i$
$\triangleright$ where slice starts in $g$

13: $end_{d}\leftarrow min(d,start_{d}+B_{d})$ $\triangleright$ where slice
ends in $g$

14: $start_{k}\leftarrow B_{k}\cdot B_{id}\cdot workload+B_{k}\cdot i$
$\triangleright$ where slice starts in $G$

15: $end_{k}\leftarrow min(B_{k},start_{k}+B_{k})$ $\triangleright$ where
slice ends in $G$

16: for $row\in\\{0<workload\\}$ do $\triangleright$ iterate all rows in the
slice $(m,B_{k})$ between $(start_{k},end_{k})$

17: for $col\in\\{start_{k}+T_{id},min(end_{k},end_{d}),T\\}$ do
$\triangleright$ jump $T$ steps (coalesced memory access)

18: $i\leftarrow\mathcal{I}_{r,col}$ $\triangleright$ read index

19: $v\leftarrow\mathcal{V}_{r,col}$ $\triangleright$ read value

20: $mem_{i-start_{d}}+=c_{r}\cdot v$ $\triangleright$ multiply the
coefficient of current row with value

21: end for$\triangleright$ finished processing one row

22: synchronize threads

23: end for$\triangleright$ finished processing all rows

24: $\triangleright$ finished processing a slice of size $(m,B_{k})$, now
write result

25: for $i\in\\{T_{id}<min(B_{d},d-start_{d}),T\\}$ do $\triangleright$ write
result (coalesced memory access)

26: $\omega_{i+start_{d}}\leftarrow mem_{i}$ $\triangleright$ write result
from $mem$ to output $\omega$

27: $mem_{i}\leftarrow 0$ $\triangleright$ zerorize shared memory to prepare
for the next step

28: end for

29: synchronize threads

30:end for

##  Appendix D Sparse GGT.

In principle, the same sparsity format can be used for GGT as in the case of
M-FAC. However, we would like to provide a slightly different approach just to
show that an existing sparsity format is not completely efficient with respect
to memory since they store redundant information. Our Sparse GGT uses the same
sparsity format as S-MFAC and here we want to emphasize some drawbacks of COO
format.

At a high level, we wish to leverage sparsity in the logical gradient history
matrix $G\in\mathbb{R}^{m\times d}$. For this, we store it in the COOrdinate
list sparse format, which requires storing matrix
$\mathcal{V}\in\mathbb{R}^{m\times k}$ for values, matrix
$\mathcal{R}\in\mathbb{Z}^{m\times k}$ for rows and
$\mathcal{C}\in\mathbb{Z}^{m\times k}$ for columns. All these matrices can be
stored in vector format, with dimension $m\cdot k$. This will be compatible
with standard sparse operations, such as the Sparse Matrix-Matrix (SPMM)
method in the PyTorch Sparse library [11].

We will use the tuple $(\mathcal{R},\mathcal{C},\mathcal{V})$ to refer to the
sparse buffer $G$. We distinguish between $\mathcal{V}$, which is the internal
state of the sparse buffer (holding the $k$ most significant values from
$a_{t}$) and $\mathcal{V}_{t}$, which is the output of Top-$k$ operator at a
single step. With slight abuse of notation, we call the compressed accumulator
$c_{t}=(\mathcal{V}_{t},\mathcal{I}_{t})$ a sparse gradient, even though it is
an accumulation of previous gradients. The matrix $\mathcal{C}$ is created in
such a way that each row contains the index of that row (e.g. [0,…,0; 1,…,1;…;
m-1,…,m-1]) and is represented as a single vector where the rows are
concatenated (here is where we have redundancy, e.g. duplicated data - the
values in this tensor never change).

Let $i\in\\{0,...,m-1\\}$ be an index indicating the location where the new
gradient will be placed in the buffer $(\mathcal{R},\mathcal{C},\mathcal{V})$.
At each step, we receive the accumulator $a_{t}$, the most significant $k$
values $\mathcal{V}_{t}$ and their corresponding indices $\mathcal{I}_{t}$.
The expression $a_{t}[\mathcal{I}_{t}]\in\mathbb{R}^{d}$ will represent a
“densified” version of a sparse vector, which contains zeros at indices
$j\notin\mathcal{I}_{t}$ (sparse tensor which also contains the zeros - does
not provide any memory reduction).

We integrate the values and indices in the internal buffers $\mathcal{V}$ and
$\mathcal{I}$ in row $i$ using the indices $\alpha$ and $\beta$. Since the
sparse $G\in\mathbb{R}^{m\times k}$ and we have to compute the scalar product
$\delta=G^{T}x$, we can easily transpose the sparse $G$ by swapping the rows
and columns in the SPMM call. The scalar products matrix $G^{T}G$ is updated
by overwriting the $i^{th}$ row and column by $\delta$. At this point we can
perform eigenvalue decomposition of $G^{T}G$ and finally compute the update
direction $u_{t}$ by preconditioning the sparse accumulator
$a_{t}[\mathcal{I}_{t}]$. The full procedure is presented in Algorithm 4.

#### GGT results for CIFAR-10/ResNet-20.

Our experimental setup covers ResNet-20 [15] ($d=0.27$M) for CIFAR-10 [20]. We
train models from scratch using SGD with momentum, dense and sparse GGT (99%
sparsity) and M-FAC and low-rank M-FAC (rank $4$). We decouple the weight
decay $\gamma$ from the gradient, and apply it in the optimizer, which allows
us to compare the loss values. The results in Table 7 show validation accuracy
on the ResNet20/CIFAR-10 experiment. We observe that 1) Dense(D) M-FAC reaches
similar accuracy to SGD, whereas Dense(D) GGT achieves lower accuracy in this
task; 2) Sparse(S) M-FAC reaches the same accuracy as the dense variant (and
SGD); 3) Low-Rank(LR) M-FAC with rank $4$ reaches < 0.8% lower accuracy
relative to the best results.

Table 7:  Results for ResNet-20 on CIFAR10. The first row reports the training loss and second row validation accuracy. The number in parentheses indicate stdev (for training loss, this is within $0.001$). | SGD | D-MFAC | S-MFAC | D-GGT | S-GGT | LR-MFAC
---|---|---|---|---|---|---
Loss | 0.032 | 0.017 | 0.017 | 0.108 | 0.097 | 0.021
Acc. | 92.16 ( .10) | 92.12 (.13) | 92.25 (.21) | 87.77 (.22) | 87.73 ( .34) | 91.43 (.30)
Algorithm 4 Detailed Sparse GGT Implementation

1:Variables: $\mathcal{V}\leftarrow 0\in\mathbb{R}^{mk}$,
$\mathcal{R}\leftarrow 0\in\mathbb{Z}^{mk}$, $i\leftarrow 0$,
$G^{T}G\leftarrow 0\in\mathbb{R}^{m\times m}$

2:Operators: $\mathcal{C}\leftarrow vec([r\cdot
1_{k}]_{r=0}^{m-1})\in\mathbb{Z}^{mk}$

3:procedure SparseGGT-
Step($t,a_{t}\in\mathbb{R}^{d},\mathcal{V}_{t}\in\mathbb{R}^{k},\mathcal{I}_{t}\in\mathbb{Z}^{k}$)

4: $\alpha\leftarrow i\cdot
k,\;\beta\leftarrow\alpha+k,\;\mathcal{V}_{\alpha:\beta}\leftarrow\mathcal{V}_{t},\;\mathcal{R}_{\alpha:\beta}\leftarrow\mathcal{I}_{t}$
$\triangleright$ integrate values and indices

5:
$\delta\leftarrow\textsc{SPMM}(Sparse(r=\mathcal{C},c=\mathcal{R},v=\mathcal{V}),Dense(a_{t}[\mathcal{I}_{t}]))$
$\triangleright$ compute $G^{T}a_{t}[\mathcal{I}_{t}]$

6: $row_{i}(G^{T}G)\leftarrow\delta,\;col_{i}(G^{T}G)\leftarrow\delta$
$\triangleright$ update dot products in row and col $i$

7: $G^{T}G\leftarrow V\Sigma_{m}^{2}V^{T},\;i\leftarrow(i+1)\;\mod\;m$
$\triangleright$ eigenvalue decomposition & index update

8:
$U_{m}\leftarrow\textsc{SPMM}(Sparse(r=\mathcal{R},c=\mathcal{C},v=\mathcal{V}),Dense(V\sqrt{\Sigma_{m}}^{\dagger})$

9:
$u_{t}\leftarrow\frac{1}{\epsilon}(a_{t}[\mathcal{I}_{t}])+U_{m}\left[(\Sigma_{m}+\epsilon
I_{m})^{-1}-\frac{1}{\epsilon}I_{m}\right]U_{m}^{T}(a_{t}[\mathcal{I}_{t}])$
$\triangleright$ compute GGT direction

10: return $u_{t}$

11:end procedure

##  Appendix E Compressing Preconditioners via Low-Rank Approximation and
Error Feedback

In this section, we discuss an alternative approach for compressed
preconditioners, via low-rank compression of the gradients. Specifically, we
will implement the Compress step in Algorithm 1 via a low-rank approximation
of the gradient, using an efficient variant of power iteration [35]. Then, we
consider the M-FAC algorithm, and reformulate its operations to leverage low-
rank matrix operations for memory savings. We start by describing how the
algorithm works a single layer, and then present the global algorithm, which
computes preconditioning for the whole model.

#### Low-Rank M-FAC for a Single Layer.

Assume that the gradient at a given layer of the model is represented by the
$s$-dimensional tensor $g\in\mathbb{R}^{p_{1}\times p_{2}\times\ldots\times
p_{s}}$. To compress the tensor, we use a variant of power iteration [35] by
firstly unfolding the tensor into a matrix
$\tilde{g}\in\mathbb{R}^{p_{1}\times p_{2:s}}$, where
$p_{2:s}=\prod_{i=2}^{s}p_{i}$. This matrix is then iteratively multiplied by
a rank $\rho$ matrix $Q\in\mathbb{R}^{p_{2:s}\times\rho}$, which is reused
from the previous iteration, to obtain the left decomposition matrix
$P\in\mathbb{R}^{p_{1}\times\rho}$. After orthogonalization of $P$, the
updated right matrix $Q$ is obtained from the relation $\tilde{g}=PQ^{T}$.
Notice that $P$ and $Q$ have $(p_{1}+p_{2:s})\rho$ elements which can be
significantly smaller than $p_{1}\cdot p_{2:s}$ for smaller ranks $\rho$. This
compression procedure is outlined in Algorithm 5.

In order to “introduce” the current gradient $g$ into the M-FAC preconditioner
data structure, we firstly need to compute an inner product between previously
stored gradients and $g$ (see Equation 5).

In matrix representation the product between two gradients corresponds to the
Frobenius inner product. Specifically, for a given low-rank representation of
gradient matrices $\tilde{g_{i}}=P_{i}Q_{i}^{T}$ and
$\tilde{g_{j}}=P_{j}Q_{j}^{T}$, the Frobenius product can be written as
$\operatorname{Tr}(\tilde{g_{i}}^{T}\tilde{g_{j}})=\operatorname{Tr}(Q_{i}P_{i}^{T}P_{j}Q_{j}^{T})=\operatorname{Tr}((P_{j}^{T}P_{i})^{T}(Q_{j}^{T}Q_{i}))$.
Notice that matrix products $P_{j}^{T}P_{i}$ and $Q_{j}^{T}Q_{i}$ are only of
size $\rho$ by $\rho$, so we compute them explicitly and then perform
summation over their element-wise products to get the final result. As full
gradients do not appear in this expression, we store just left and right
gradients $P_{i}$ and $Q_{i}$ and update the M-FAC inner product matrix as
needed.

The final preconditioned gradient update is computed as a weighted sum of
previous gradients. The weights are computed via dynamic programming on top of
the inner product values as in the original M-FAC. The final weighted sum of
gradients is obtained by sequentially adding gradients reconstructed from
their low-rank representation $g_{i}=P_{j}Q_{j}^{T}$.

1:procedure PowerCompress(gradient tensor $g$, previous $Q$) $\triangleright$
applied for each layer

2: $\tilde{g}\leftarrow g.\textsc{Reshape}(p_{1},p_{2:s})$

3: $P\leftarrow\tilde{g}Q$

4: $P\leftarrow\textsc{Orthogonalize(P)}$

5: $Q\leftarrow\tilde{g}^{T}P$

6: return $(P,Q)$

7:end procedure

Algorithm 5 Power iteration compression algorithm [35]

1:INPUT: per-layer compressed (($P^{1}$, $Q^{1}$), …,($P^{L}$, $Q^{L}$))
gradients

2:STATE: per-layer ring buffers $G^{\ell}$ with compressed gradient history;
full inner products $G^{T}G$

3:procedure Low Rank M-FAC(($P^{1}$, $Q^{1}$), …,($P^{L}$, $Q^{L}$))

4: $\delta\leftarrow$ zero vector of size $m$ $\triangleright$ new full inner
product

5: for each layer $\ell\in\\{1,2,...,L\\}$ do $\triangleright$ aggregate inner
products across the whole model

6: Update $G^{\ell}$ with current gradient $(P^{\ell},Q^{\ell})$

7: for each time step $\tau\in\\{t-m+1,...,t\\}$ do

8: $\delta_{t-\tau}\leftarrow$ low-rank inner product
$Tr((g_{\tau}^{\ell})^{T}(P^{\ell}(Q^{\ell})^{T}))$

9: end for

10: end for

11: Update row and column in $G^{T}G$ from $\delta$

12: M-FAC: compute summation weights $w\in\mathbb{R}^{m}$ from $G^{T}G$

13: for each layer $\ell\in\\{1,2,...,L\\}$ do

14:
$u_{t}^{\ell}\leftarrow\epsilon^{-1}P^{\ell}(Q^{\ell})^{T}+\sum_{\tau=t-m+1}^{t}w_{\tau}g_{\tau}$
$\triangleright$ aggregate gradients from the ring buffer

15: end for

16: return vector of per-layer updates $u_{t}$

17:end procedure

Algorithm 6 Low Rank M-FAC preconditioning

#### Global Low-Rank M-FAC.

To obtain an equivalent algorithm to the original M-FAC, independent per-layer
steps are not sufficient as we would lose cross-layer correlation information.
Instead, we use linearity of the inner product to compute the final result
efficiently. In the original version of M-FAC, a current complete gradient
from a $L$-layered model is represented as a concatenation of flattened per-
layer gradients
$\widehat{g}=\begin{pmatrix}\widehat{g}^{1}&\widehat{g}^{2}&\ldots&\widehat{g}^{L}\end{pmatrix}^{T}$,
where $\hat{g}^{i}$ is the gradient from layer $i$. In this, representation an
inner product between two gradients $\widehat{g}_{i}$ and $\widehat{g}_{j}$
can be expressed in a block form as
$\left(\widehat{g}_{i},\widehat{g}_{j}\right)=\sum_{\ell=1}^{L}\left(\widehat{g}_{i}^{\ell},\widehat{g}_{j}^{\ell}\right)$.
This dot product is then used to update the inner product matrix $G^{T}G$.

In the context of our efficient scheme for single-layer low-rank update, we
replace each inner product block by its low-rank counterpart to obtain the
equivalent values. Recall that the computation of summation weights in the
M-FAC needs access only to these inner products, so this part of the algorithm
remains the same. The final gradient update is computed per-layer and we rely
on linearity of the sum again to share the summation weights and dampening
coefficient $\epsilon$ between layers. For more details refer to the
pseudocode in Algorithm 6.

#### Results for Low-Rank M-FAC.

We present and discuss the experiments for Low-Rank M-FAC on
CIFAR-10/ResNet-20 in Table 7 and hyper-parameters in Appendix A. Next, we
present an ablation study of low rank decomposition on both CIFAR-10/ResNet-20
and GLUE/MNLI for BERT-tiny.

To determine how the size of the decomposition rank $\rho$ affects the
effectiveness of the low rank M-FAC algorithm, we perform a series of
experiments on ResNet-20 and BERT models. For ResNet-20 we use CIFAR-10
dataset and the same parameters as in the main experiments. For BERT family we
use tiny, mini and base variants and evaluate them on MNLI task from GLUE
benchmark. For both setups we vary the rank from $1$ to $8$ and report the
final test accuracy in the Table 8.

We observe that the decomposition rank significantly affects smaller BERT
models but barely influences BERT-base and RN-20 model results. While the
exact reason for such behaviour is not obvious, we hypothesise that for RN-20
the decomposition is effective even at rank $1$ due to gradient tensors
already having a small first dimension, resulting in small compression error.
We also note that Rank $8$ BERT-base run diverged, requiring more
investigation into the stability of the training.

Table 8: Test accuracy on CIFAR-10 dataset for ResNet-20 (RN-20) and on MNLI for BERT-tiny, -mini and -base, trained with Low-Rank M-FAC with different decomposition ranks $\rho$. rank $\rho$ | RN-20 | TINY | MINI | BASE
---|---|---|---|---
1 | 91.40 | 57.34 | 66.88 | 84.42
2 | 91.19 | 63.49 | 68.76 | 84.44
4 | 91.43 | 66.44 | 70.93 | 84.45
8 | 91.13 | 67.90 | 72.22 | —