A New Perspective on Shampoo's Preconditioner

Introduction

Second-order optimization is a rich research area within deep learning that has seen multiple influential works over the past few decades. Recently, these methods have seen success in practical large scale training runs such as Gemini 1.5 Flash (Gemini Team, 20024) and in academic benchmarks (Dahl et al., 2023). One of the primary challenges in this field arises from the substantial memory and computational demands of traditional second-order methods, such as Adagrad (Duchi et al., 2011b) and Newton’s method. In the context of neural networks, both of these methods require storing and inverting a $|P|\times|P|$ dimensional matrix $H$ (either covariance of the gradients for Adagrad or the Gauss–Newton component of the Hessian for Newton’s method), where $|P|$ represents the number of parameters of the neural network. With modern deep learning architecture scaling to billions of parameters, these requirements make the direct application of these methods impractical. To address this issue, various approaches have been proposed, including Hessian-free optimization (Martens et al., 2010) and efficient approximations of the matrix $H$ (Gupta et al., 2018b; Martens & Grosse, 2015b). These methods aim to leverage second-order information while mitigating the computational and memory overhead.

The class of methods for efficiently approximating the matrix $H$ predominantly involve either a diagonal or a layer-wise Kronecker product approximation of $H$. These choices are motivated by the fact that, compared to maintaining the matrix $H$, both diagonal and layer-wise Kronecker products are significantly more memory-efficient to store and computationally efficient to invert. Two of the most well-known methods that utilize a layer-wise Kronecker product approximation of $H$ are K-FAC (Martens & Grosse, 2015b) and Shampoo (Gupta et al., 2018b).

In this work, we primarily focus on the Shampoo optimizer (Gupta et al., 2018b), which has recently gained increasing attention from the research community. Notably, in a recent benchmark of optimization algorithms proposed for practical neural network training workloads (Dahl et al., 2023), Shampoo appears to outperform all other existing methods. Another recent study, elucidating the Google Ads recommendation search pipeline, revealed that the Google Ads CTR model is trained using the Shampoo optimizer (Anil et al., 2022). Additionally, a recent work (Shi et al., 2023) implemented a distributed data parallel version of Shampoo, demonstrating its superior speed in training ImageNet compared to other methods.

Previously, Shampoo’s approximation was shown to be an upper bound (in spectral norm) on the matrix $H$ (Gupta et al., 2018b). In this work, we make this connection much more precise. Prior research has established the notion of the optimal Kronecker product approximation (in Frobenius norm) of $H$ (Koroko et al., 2023b), which can be obtained numerically using a power iteration scheme. The primary contribution of this work is to theoretically and empirically demonstrate that the square of the approximation used by Shampoo is nearly equivalent to the optimal Kronecker factored approximation of $H$.

The main contributions of the work are summarized below:

We empirically establish that the result of one round of power iteration is very close to the optimal Kronecker factored approximation (see Figure 1), and provide theoretical justification for the same.

For the Hessian based viewpoint of Shampoo (Section 2.1.2), we empirically demonstrate the impact on the Hessian approximation of various practical tricks implemented to make Shampoo more computationally efficient such as averaging gradients over batch (Section 4.1) and using empirical Fisher instead of the actual Fisher (Section 4.2).

Remark. Previous works (Balles et al., 2020; Lin et al., 2024) have explored the question of why Adagrad-based approaches like Adam and Shampoo have an extra square root compared to the Hessian inverse in their update. This alternative question is orthogonal to our contribution. For details, refer Appendix F.

Paper organization. In Section 2, we cover the technical background necessary for understanding this work. In Section 3, we provide a general power iteration scheme for obtaining the optimal Kronecker product approximation of the matrix $H$, and establish the the connection between Shampoo’s approximation and the optimal Kronecker product approximation of $H$. In Section 4, we explore the Hessian approximation viewpoint of Shampoo and empirically study how various practical tricks to make Shampoo more computationally efficient impact the quality of the Hessian approximation. In Section 5, we cover closely related works and conclude with discussing the limitations of the work in Section 6. In Appendix A, we include additional experiments on the ViT architecture and compare with the K-FAC approximation to the Hessian. Detailed related work, proofs, dataset and architecture details have been deferred to the Appendix.

Technical background

Following is a basic lemma about Kronecker products that will be used later

$(A\otimes B)\operatorname{vec}(G)=\operatorname{vec}(BGA^{\top})$.

The original Shampoo (Gupta et al., 2018b) paper introduced its algorithm as an approximation of an online learning algorithm Adagrad (Duchi et al., 2011a). Shampoo can also be interpreted (Anil et al., 2020; Osawa et al., 2023a) as approximating the Gauss–Newton component of the Hessian. Both of these perspectives will be discussed in Section 2.1.1 and 2.1.2 respectively. .

Assume that $G_{1},...,G_{T}$ are matrices of rank at most $r$. Let $g_{t}=\operatorname{vec}(G_{t})$ for all $t$. Then, with $\preccurlyeq$ representing the for any $\epsilon>0$,

In Lemma 2 the matrix $H_{\text{Ada}}=\sum_{t=1}^{T}g_{t}g_{t}^{\top}$ is approximated (ignoring $\epsilon$ and scalar factors) by the the Kronecker product $\left(\sum_{t=1}^{T}G_{t}G_{t}^{\top}\right)^{1/2}\otimes\left(\sum_{t=1}^{T}G_{t}^{\top}G_{t}\right)^{1/2}$. Our main focus will be to study the optimal Kronecker product approximation of the matrix $H_{\text{Ada}}$ and its connection to Shampoo’s approximation (done in Section 3).

1.2 Hessian based perspective of Shampoo

In this section we describe the Hessian approximation viewpoint of Shampoo explored by previous works (Anil et al., 2020; Osawa et al., 2023a) as an alternative to the Adagrad viewpoint described above. Our theoretical and empirical results hold for both viewpoints.

The aim of algorithms such as K-FAC and Shampoo (when viewed from the Hessian perspective) is to do a layerwise Kronecker product approximation of the Fisher matrix $H_{\text{GN}}$. The following lemma establishes the approximation made by Shampoo:

Assume that $G_{x,s}$ are matrices of rank at most $r$. Let $g_{x,s}=\operatorname{vec}(G_{x,s})$ . Then, for any $\epsilon>0$,

In Lemma 2 the matrix on the left hand side is equal to $H_{\text{GN}}$ and the right hand side represents the $H_{\text{GN}}$ approximation made by Shampoo. However, computing this approximation at every step is expensive. So, in practice, Shampoo makes two additional approximations on top.

Thus, if $G_{j}$ and $W_{j}$ represent the batch gradient and weight matrix at iteration $j$, and $\lambda$ is an exponential weighting parameter, then the update of Shampoo is given by

where $L_{j}$ and $R_{j}$ represent the left and right preconditioners maintained by Shampoo, respectively.

Our focus (when viewing Shampoo from the Hessian perspective) will be to study

The optimal Kronecker product approximation of the matrix $H_{\text{GN}}$ and its connection to Shampoo’s approximation (done in Section 3).

The effect of the aforementioned two approximations on the approximation quality (done in Section 4).

2 Optimal Kronecker product approximation

where $\|\cdot\|_{F}$ denotes the Frobenius norm.

Since the optimal rank-1 approximation of a matrix is given by its singular value decomposition (SVD), we conclude:

minimize the Frobenius norm $\|H-L\otimes R\|_{F}$.

Cosine similarity. We will be using cosine similarity between matrices as a metric for approximation. For two matrices $M_{1}$ and $M_{2}$, this refers to $\text{Tr}(M_{1}M_{2}^{\top})/(||M_{1}||_{F}\cdot||M_{2}||_{F})$. A value of 1 indicates perfect alignment, while a value of 0 indicates orthogonality.

Optimal Kronecker product approximation and Shampoo

Loan & Pitsianis (1993) describe an approach to find the optimal Kronecker product approximation of a matrix (with respect to the Frobenius norm). Koroko et al. (2023b) use this approach to find the optimal layer-wise Kronecker product approximation of the hessian matrix for networks without weight sharing. We will now do a general analysis which would also be applicable to neural networks with weight sharing.

Reshaping vectors on both sides into matrices results in:

Our first and main approximation involves replacing the iterative power iteration scheme (Equation 4) with just a single iteration. This leads to the main contribution of our work:

One step of power iteration, starting from the identity, for obtaining the optimal Kronecker product approximation of $H$ is precisely equal to the square of the Shampoo’s approximation of $H$

The initialization for the single iteration will use the identity matrix, i.e., $I_{m}$ and $I_{n}$ for $L$ and $R$, respectively. Thus, we transition from the iterative update equations:

to the simplified single-step expressions:

With the above expression for $L$ and $R$, $L\otimes R$ is precisely equal to the square of the Shampoo’s approximation of $H$ given by the right hand side of Equation 1. ∎

As shown in Figure 1, for various datasets and architectures, this single step of power iteration is very close to the optimal Kronecker product approximation for both $H=H_{\text{GN}}$ (top) and $H=H_{\text{Ada}}$ (bottom). However, we can see that the upper bound proposed by the original Shampoo work (Gupta et al., 2018b) is significantly worse.

One reason why the cosine similarity might be large is that $\hat{H}$ is nearly rank-1 ($\sigma_{1}$ is large); that is, $H$ is closely approximated by a Kronecker product. As illustrated in Figure 1, this assumption does not universally hold. Instead, we propose an alternative explanation for why a single step of power iteration is typically sufficient: the coefficient $\alpha_{1}$ is usually larger than $\alpha_{i}$ for all $i\geq 2$. We begin by providing a theoretical justification for this, followed by empirical evidence from our experiments.

We start by noting that $\alpha_{i}=\operatorname{vec}\left(I_{n}\right)^{T}v_{i}=\text{Tr}(V_{i})$. Now, we will show that using the identity matrix as initialization is a good choice since a) shows it has the maximal dot product with possible top components i.e., PSD matrices (Proposition 2), and b) we expect it to have a small dot product with later components.

$V_{1}$ is a Positive Semi-Definite (PSD) matrix.

Since $V_{1}$ is a PSD matrix we would like to initialize our power iteration with a matrix which is close to all PSD matrices. Now, we will show that identity is the matrix which achieves this, specifically it maximizes the minimum dot product across the set of PSD matrices of unit Frobenius norm.

Consider the set of PSD matrices of unit Frobenius norm of dimension $m$ denoted by $S_{m}$. Then

The previous proposition argues that $I_{m}$ maximizes the worst-case dot product with possible top singular vectors. Now, we argue that its dot product with other singular vectors should be lower.

If $V_{1}$ is positive-definite, then $V_{i}$ for $i\geq 2$ are not PSD.

Therefore, the diagonal elements of $V_{i}$ for $i\geq 2$ need not be positive, and this might lead to cancellations (for $i\geq 2$) in the trace of $V_{i}$ which is equal to $\alpha_{i}$. Hence we expect $\alpha_{i}$’s for $i\geq 2$ to be smaller than $\alpha_{1}$. We now show experiments to demonstrate this in practice. To quantify the benefit of $\alpha_{1}$ usually being larger than $\alpha_{i}$ for $i\geq 2$, we will compare $\frac{\alpha_{1}\sigma_{1}}{\sqrt{\sum_{i}\alpha_{i}^{2}\sigma_{i}^{2}}}$ (for both left and right singular vectors) and $\frac{\sigma_{1}}{\sqrt{\sum_{i}\sigma_{i}^{2}}}$. The latter can be interpreted as the cosine similarity if all $\alpha$’s were equal or as a measure of how close $\hat{H}$ is to being rank 1 since it is equal to the cosine similarity between $u_{1}v_{1}^{T}$ and $\hat{H}$. Thus $\frac{\sigma_{1}}{\sqrt{\sum_{i}\sigma_{i}^{2}}}$ is equal to the “Optimal Kronecker” cosine similarity used in Figure 1. In Figure 2 we track both of these quantities through training and indeed observe that $\frac{\alpha_{1}\sigma_{1}}{\sqrt{\sum_{i}\alpha_{i}^{2}\sigma_{i}^{2}}}$ are significantly closer to 1 than $\frac{\sigma_{1}}{\sqrt{\sum_{i}\sigma_{i}^{2}}}$ for both $H=H_{\text{GN}}$ (top) and $H=H_{\text{Ada}}$ (bottom).

1.2 Exact Kronecker product structure in H𝐻H

Under the assumption that $\hat{H}$ is rank-1,

Let $\hat{H}=\sigma uv^{\top}$, i.e, $H=\sigma U\otimes V$. Let $I_{m}=\text{Tr}(U)U+R_{m}$ and $I_{n}=\text{Tr}(V)V+R_{n}$, where $R_{m}$ and $R_{n}$ are the residual matrices. Now, after one round of power iteration, the left and right estimates provided by Shampoo are given by

Since $H=\hat{H}_{\text{GN}}$ is an $m^{2}\times 1$ matrix for binomial logistic regression, it is rank-1, so the equality in the corollary holds. In other words, the square of Shampoo’s $H_{\text{GN}}$ estimate perfectly correlates with $H_{\text{GN}}$ for binomial logistic regression. This is demonstrated in the first plot of Figure 1.

1.3 Discussion about optimization

Hessian Approximation of Shampoo

From the Hessian approximation viewpoint, the previous section covers the case of using batch size $1$ and sampled labels, as described in Section 2.1.2. To be precise, in Figure 1 top, we consider how well $H_{\text{GN}}$ is correlated with $E_{x,s}[G_{x,s}G_{x,s}^{T}]\otimes E_{x,s}[G_{x,s}^{T}G_{x,s}]$, where $s$ represents that the labels are sampled from the model’s output distribution. On the other hand, as discussed in Section 2.1.2, Shampoo in practice is generally used with arbitrary batch sizes and real labels. We now investigate the effect of these two factors on the Hessian approximation.

The next approximation towards Shampoo is to average the gradient across the batch, i.e., we go from

where $B$ denotes the batch, $\bf{s}$ is the concatenation of $s\sim f(x)$ for all $x\in B$ and $G_{B,\bf{s}}=\frac{1}{|B|}\sum_{x\in B,s=\mathbf{s}[x]}G_{x,s}$ is the batch gradient, with $\mathbf{s}[x]$ representing the sampled label corresponding to $x\in B$.

However, this does lead to a significant improvement in computational complexity by saving up to a factor of batch size.

2 Using real labels instead of sampled labels

As our final approximation we replace using sampled labels $s\sim f(x)$ to using real labels $y$. This approximation, denoted in the literature by empirical Fisher when batch size is 1, has been discussed at length by prior works (Osawa et al., 2023a; Kunstner et al., 2019). The main theoretical argument for why this approximation may work well is that, as we move towards optima, the two quantities converge in the presence of label noise (Grosse, 2021).

In Figure 3 (top), when evaluating $H_{\text{GN}}$ approximation with batch size $1$, we surprisingly find that the approximation quality is good throughout the training. However, unlike the case of sampled labels, the approximation starts to degrade at large batch sizes because the gradients with real labels are not mean 0. The lemma below (Grosse, 2021) shows how this estimator changes with batch size.

Let $B$ denote the batch and $G_{B}=\frac{1}{|B|}\sum_{(x,y)\in B}G_{x,y}$ denote the batch gradient. Then

We note that this approximation has the computational benefit of not requiring another backpropagation with sampled labels; instead, these computations can be done alongside usual training.

Related work

We discuss the related works in detail in Appendix E. Here, we discuss two closely related works: Ren & Goldfarb (2021) and Koroko et al. (2023a).

Ren & Goldfarb (2021) study the Hessian perspective of Shampoo and show that, under the assumption that sampled gradients follow a tensor-normal distribution, the square of the Hessian estimate of Shampoo is perfectly correlated with $H_{\text{GN}}$. We also show the same result under much weaker conditions in Corollary 2. Moreover, in Proposition 1 we show that, in general, the square of the Hessian estimate of Shampoo is closely related to the optimal Kronecker product approximation of $H_{\text{GN}}$. We additionally also study the approximations used by Shampoo to make it computationally efficient (Section 4) and the Adagrad perspective of Shampoo’s preconditioner.

Loan & Pitsianis (1993) develop the theory of optimal Kronecker product approximation of a matrix (in Frobenius norm). Koroko et al. (2023a) use it for finding layer-wise optimal Kronecker product approximation of $H_{\text{GN}}$ for a network without weight sharing. We extend their technique to networks with weight-sharing, and show that the square of the Hessian estimate of Shampoo is nearly equivalent to the optimal Kronecker product approximation of $H_{\text{GN}}$.

Limitations

The main contribution of our work is to show that the square of the Shampoo’s approximation of $H$ (where $H$ refers to either $H_{\text{Ada}}$ or $H_{\text{GN}}$) is nearly equivalent to the optimal Kronecker approximation of $H$. Although we verify this empirically on various datasets and provide theoretical arguments, the gap between them depends on the problem structure. In some of our experiments with ViT architecture (Appendix A), we find that the gap is relatively larger compared to other architectures. Moreover, it remains an open question to understand the conditions (beyond those described in K-FAC Martens & Grosse (2015b)) under which $H$ is expected to be close to a Kronecker product. Again, in some of the experiments with ViTs (Appendix A), we find that the optimal Kronecker product approximation to $H$ is much worse as compared to other architectures.

Acknowledgements

NV and DM are supported by a Simons Investigator Fellowship, NSF grant DMS-2134157, DARPA grant W911NF2010021, and DOE grant DE-SC0022199. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. SK and DM acknowledge funding from the Office of Naval Research under award N00014-22-1-2377 and the National Science Foundation Grant under award #IIS 2229881. LJ acknowledges funding from the National Science Foundation DMS-2134157.

References

Appendix A Additional experimental results

In this subsection, we present the results for a Vision Transformer (ViT) architecture trained on the CIFAR-5m dataset. This architecture features a patch size of 4, a hidden dimension of 512, an MLP dimension of 512, 6 layers, and 8 attention heads.

For these experiments, we utilize three layers from the fourth transformer block: two layers from the MLP (referred to as ’FFN Linear Layer 1’ and ’FFN Linear Layer 2’) and the QK layerThe QK layer is separated from the V part of the layer, following similar decomposition method described by Duvvuri et al. (2024) (referred to as ’Q-K Projection Layer’).

Appendix B Experiments

Datasets and Architectures. We conducted experiments on three datasets: MNIST (LeCun et al., 1998), CIFAR-5M (Nakkiran et al., 2020), and ImageNet (Deng et al., 2009), using logistic regression, ResNet18 (He et al., 2016), and ConvNeXt-T (Liu et al., 2022) architectures, respectively. For MNIST, we subsampled two digits ($\{0,1\}$) and trained a binary classifier.

For MNIST, we used the only layer, i.e, the first layer of the linear classifier for computing the cosine similarities. For Resnet18 and Imagenet, we picked arbitrary layers. In particular, for Resnet 18, we used one of the convolution layers within the first block (’layer1.1.conv1’ in https://pytorch.org/vision/master/_modules/torchvision/models/resnet.html#resnet18). For Imagenet, we used the 1x1 convolutional layer within the 2nd block of convnext-T (’stages.2.1.pwconv1’ in https://pytorch.org/vision/main/models/generated/torchvision.models.convnext_tiny.html#torchvision.models.convnext_tiny).

Cosine similarity estimation for $H_{\text{GN}}$. For estimating the Frobenius norm of $H_{\text{GN}}$, we used the identity:

Hessian-vector products with the Gauss–Newton component were performed using the DeepNetHessian library provided by Papyan (2019).

For estimating the cosine similarity between $H_{\text{GN}}$ and its estimator $\widetilde{H}_{\text{GN}}$, we used the following procedure:

Estimate $\|H_{\text{GN}}\|_{F}$, and calculate $\|\widetilde{H}_{\text{GN}}\|_{F}$.

Define scaled $\widetilde{H}_{\text{GN}}$ as $\widetilde{S}_{\text{GN}}=\frac{\|H_{\text{GN}}\|_{F}}{\|\widetilde{H}_{\text{GN}}\|_{F}}\widetilde{H}_{\text{GN}}$.

$\text{Cos-sim}(H_{\text{GN}},\widetilde{H}_{\text{GN}})=1-\frac{\|H_{\text{GN}}-\widetilde{S}_{\text{GN}}\|_{F}^{2}}{2\|H_{\text{GN}}\|_{F}^{2}}$, where the numerator is again estimated via Hessian-vector products.

Note that in the above procedure, we can exactly calculate $\|\widetilde{H}_{\text{GN}}\|_{F}$ as it is generally of a Kronecker product form with both terms of size $m\times m$ or $n\times n$, where $m\times n$ is the size of a weight matrix.

Cosine similarity estimation for $H_{\text{Ada}}$. We follow a similar recipe as before, but using a difference method for computing the product $H_{\text{Ada}}v$. For a given time $T$, $H_{\text{Ada}}=\sum_{t=1}^{T}g_{t}g_{t}^{\top}$. Thus, $H_{\text{Ada}}v=\sum_{t=1}^{T}(g_{t}^{\top}v)g_{t}$. We maintain this by keeping a running estimate of the quantity for multiple random vectors $v$ during a training run, and use it for estimating the product $H_{\text{Ada}}v$.

Optimal Kronecker method, wherever used was computed with five rounds of power iteration, starting from the identity. For $H=H_{\text{GN}}$, the Hessian approximations $\textit{Shampoo}^{2}$, Shampoo, and K-FAC were done using sampled labels and a batch size of $1$. For $H=H_{\text{Ada}}$ and step $t$, we used gradient enocoutered during the training run in steps $\leq t$.

K-FAC was computed with the “reduce” variant from Eschenhagen et al. (2023).

In Figure 2, the Optimal Kronecker legend represents the cosine similarity between the optimal Kronecker approximation of $H_{\text{GN}}$ and $H_{\text{GN}}$. This is precisely equal to $\frac{\sigma_{1}}{\sqrt{\sum_{i}\sigma_{i}^{2}}}$. Similarly, the label L (resp. R) represents the cosine similarity between the top left (resp. right) singular vector of $\hat{H}_{\text{GN}}$ and the estimate obtained after one round of power iteration starting from $I_{n}$ (resp. $I_{m}$). This is precisely equal to $\frac{\alpha_{1}\sigma_{1}}{\sqrt{\sum_{i}\alpha_{i}^{2}\sigma_{i}^{2}}}$.

In Figure 3 (top), the Hessian approximation is calculated with batch size $1$, i.e, $|B|=1$ in Section 4.2. Similarly, in Figure 3 (bottom), $|B|=256$.

Appendix C Deferred proofs

Consider two PSD matrices $M_{1}$ and $M_{2}$ having the eigenvalue decomposition $M_{1}=\sum\lambda_{1i}q_{1i}q_{1i}^{\top}$ and $M_{2}=\sum\lambda_{2i}q_{2i}q_{2i}^{\top}$. Then

Thus, if $M_{1}$ and $M_{2}$ have unit frobenius norm and $M_{1}$ is positive definite, then $\text{Tr}(M_{1}M_{2})>0$.

Thus, if $V_{1}$ is positive definite, then by orthogonality of successive singular vectors, $V_{i}$ for $i\geq 2$ cannot be positive semi-definite. ∎

Consider the eigendecomposition of any $M\in S_{q}$ given by $\sum_{i=1}^{q}\lambda_{i}v_{i}v_{i}^{\top}$. Denote $L=\{i:\lambda_{i}\leq\frac{1}{\sqrt{q}}\}$. As $\sum\lambda_{i}^{2}=1$, therefore, $|A|\geq 1$. Consider any $j\in A$. Then

As $v_{j}$ is orthogonal to the other eigenvectors. Thus, we can see

Moreover, for the matrix $\frac{1}{\sqrt{q}}I_{q}$, for any matrix $M^{\prime}$,

where $\text{tr}(M^{\prime})$ denotes the trace of the matrix $M^{\prime}$. However, we know $\text{tr}(M^{\prime})=\sum\lambda_{i}\geq 1$ as $\sum\lambda_{i}^{2}=1$. Thus

Note that this is the only matrix with this property as any other matrix will at least have one eigenvalue less than $\frac{1}{\sqrt{q}}$. Thus

Evaluating $G_{B,\bf{s}}G_{B,\bf{s}}^{T}$, we get

Now taking an expectation over batches, we get

Taking the expectation over $B$ on both the sides, we get

Appendix D Technical Background on Hessian

Appendix E Related work

The literature related to second order optimization within deep learning is very rich, with methods that can be broadly classified as Hessian-free and methods based on estimating the preconditioner $H$ (which could refer to either $H_{\text{Ada}}$ or $H_{\text{GN}}$). Hessian-free methods (Martens, 2010) generally tend to approximate the preconditioned step (for Newton’s method) using Hessian vector products, but do not maintain an explicit form of the Hessian. Estimating $H$ (Martens & Grosse, 2015a; Gupta et al., 2018a) methods maintain an explicit form of the preconditioner that could be efficiently stored as well as estimated.

One of the seminal works related to second order optimization within deep learning was the introduction of Hessian-free optimization (Martens, 2010). The work demonstrated the effectiveness of using conjugate gradient (CG) for approximately solving the Newton step on multiple auto-encoder and classifications tasks. Multiple works (Martens & Sutskever, 2011; Cho et al., 2015) have extended this algorithm to other architectures such as recurrent networks and multidimensional neural nets. One of the recent works (Garcia et al., 2023) also takes motivation from this line of work, by approximately using single step CG for every update, along with maintaining a closed form for the inverse of the Hessian, for the single step to be effective.

E.2 Estimating Preconditioner

Given that it is costly to store the entire matrix $H$, various works have tried to estimate layer-wise $H$. KFAC (Martens & Grosse, 2015a) was one of the first work, that went beyond diagonal approximation and made a Kronecker product approximation to layer-wise $H_{\text{GN}}$. It showed that this structure approximately captures the per layer Hessian for MLPs. This approximation was extended to convolutional (Osawa et al., 2019) and recurrent (Martens et al., 2018) architectures. Subsequent works also improved the Hessian approximation, by further fixing the trace (Gao et al., 2021) as well as the diagonal estimates (George et al., 2018; Gao et al., 2020) of the approximation. A recent work (Eschenhagen et al., 2023) also demonstrated that K-FAC can be extended to large-scale training.

From the viewpoint of approximating Adagrad (Duchi et al., 2011b), Gupta et al. (2018a) introduced Shampoo, that also makes a Kronecker product approximation to $H_{\text{Ada}}$. One of the subsequent work (Ren & Goldfarb, 2021) introduced a modification of Shampoo, that was precisely estimating the layer-wise $H_{\text{GN}}$ under certain distributional assumptions. Other works (Anil et al., 2021) introduced a distributed implementation of Shampoo, that has recently shown impressive performance for training large scale networks (Shi et al., 2023). Recently, another paper (Duvvuri et al., 2024) proposed a modification of Shampoo, empirically and theoretically demonstrating that the new estimator approximates $H_{\text{Ada}}$ better than Shampoo’s approximation. Our work shows that the square of Shampoo’s approximation of $H_{\text{Ada}}$ is nearly equivalent to the optimal Kronecker approximation.

Appendix F Comparison with extra square root in Adagrad based approaches

Multiple previous works (Balles et al., 2020; Lin et al., 2024) have tried to address the question of why Adagrad-based approaches like Adam and Shampoo, have an extra square root in their update compared to Hessian inverse in their updates. This question is primarily concerned with the final update to the weights being used in the optimization procedure, once we have approximated the Hessian.

The primary contribution of this work is completely orthogonal to this question. We are addressing the question of optimal Kronecker approximation of the Hessian, and its connection to Shampoo’s Hessian approximation. This is orthogonal to the Hessian power used in the final update.