SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jui-Nan Yen
juinanyen@cs.ucla.edu
Sai Surya Duvvuri
saisurya@cs.utexas.edu
Inderjit S. Dhillon
Google and UT Austin
inderjit@cs.utexas.edu
Google and UCLA
Adaptive methods with non-diagonal preconditioning have shown state-of-the-art results on various tasks. However, their computational complexity and memory requirement make it challenging to scale these methods to modern neural network architectures. To address this challenge, some previous works have adopted block-diagonal preconditioners. However, the memory cost of storing the block-diagonal matrix remains substantial, leading to the use of smaller block sizes that ultimately leads to suboptimal performance. To reduce the time and memory complexity without sacrificing performance, we propose approximating the diagonal blocks of the second moment matrix by low-rank matrices and enforcing the same basis for the blocks within each layer. We provide theoretical justification for such basis sharing and design an algorithm to efficiently maintain this shared-basis block low-rank approximation during training. Our results on a deep autoencoder and a Transformer benchmark demonstrate that the proposed method outperforms first-order methods with slightly more time and memory usage, while also achieving competitive or superior performance compared to other second-order methods with less time and memory usage.
Diagonal adaptive methods such as Adagrad [10] incorporate an adaptive learning rate schedule for each coordinate of the gradient, which is typically estimated as inverse root mean square of the histories of the coordinate gradient  $1 / \sqrt{\sum_{t}(\pmb{g}_{t})_{i}^{2}}$ , for a gradient vector  $\pmb{g}_{t}$ . This is equivalent to multiplying the gradient vector with a diagonal matrix, a preconditioner. This approach is adopted in optimizers such as Adam [20], RMSProp [29], which became the preferred approach to training deep neural networks due to their robustness to the learning rate parameter. Several methods later developed preconditioners such as Shampoo [16, 4], GGT [11], which also estimate off-diagonal gradient moments such as  $\sum_{t}(\pmb{g}_{t})_{i}(\pmb{g}_{t})_{j}$  for  $i \neq j$  and develop a positive definite preconditioner, while outperforming the diagonal methods in both training and validation performance. However, this approach can have high memory requirements due to additional off-diagonal moment information the optimizer needs to store and the high computational requirement of the inverse operation in computing the preconditioner.
The second-moment matrix  $\sum_{t} g_{t} g_{t}^{T}$  constituting both diagonal and off-diagonal cross-moments is infeasible to store, however, the gradients from deep-neural networks can reveal a structure in this matrix. To simplify this problem we analyze a basic neural network block - linear transformation  $X \rightarrow XW$ , where  $X \in \mathbb{R}^{n \times d_{\mathrm{in}}}$  and  $W \in \mathbb{R}^{d_{\mathrm{in}} \times d_{\mathrm{out}}}$  and make two important observations: a) the
second-moment matrix  $\sum_{t} g_{t} g_{t}^{T}$ , where  $g_{t} \in \mathbb{R}^{d_{\mathrm{in}} d_{\mathrm{out}}}$  is a vectorized form of the gradient with respect to  $\pmb{W}$  at iteration  $t$ , has a block-matrix structure with  $d_{\mathrm{out}}$  blocks of size  $d_{\mathrm{in}} \times d_{\mathrm{in}}$ , with high norm diagonal blocks; b) all the diagonal blocks share similar column and row-space with low dimension. Using a block-diagonal approximation directly can still be infeasible due to  $\mathcal{O}(d_{\mathrm{out}} d_{\mathrm{in}}^2)$  memory complexity to store all the blocks. To resolve this, a mini-block approach was used previously in [32], which uses smaller blocks  $b < d_{\mathrm{in}}$ . However, this approximation can discard important within-layer off-diagonal second-moment information. To alleviate this issue we use observation b) and compute a common low-rank basis  $\pmb{B} \in \mathbb{R}^{d_{\mathrm{in}} \times k}$ ,  $k \ll d_{\mathrm{in}}$ , and a different residual  $\pmb{R}^{(i)} \in \mathbb{R}^{k \times k}$  for each block such that  $\{\pmb{B R}^{(1)} \pmb{B}^{\top}, \dots, \pmb{B R}^{(d_{\mathrm{out}})} \pmb{B}^{\top}\}$  approximates the diagonal blocks.
We develop an online algorithm to update the shared basis  $B$ , for which we derive approximation error guarantees demonstrating the nearness of  $\{ BR^{(1)}B^{\top}, \ldots, BR^{(d_{\mathrm{out}})}B^{\top} \}$  to the original block-diagonal approximation of the gradient second-moment matrix. We also use this to prove a  $\mathcal{O}(\sqrt{T})$  regret upper bound in the online convex optimization framework [18, 28], which can translate to convergence guarantees in non-convex stochastic optimization as in [1]. We compare the proposed algorithm with state-of-the-art second-order optimizers along with Adam, which is widely recognized as the most commonly used first-order optimizer. The evaluation is performed on a standard autoencoder benchmark and a Transformer benchmark comprising 19.3 million learnable parameters. Compared with other second-order methods, our algorithm utilizes lower memory and time to perform updates at every iteration while showing better or similar validation and training performance. Specifically, in the Autoencoder benchmark, we show similar training performance as Shampoo while using  $3.2 \times$  lower time, and  $1.5 \times$  lower memory. In addition to this, in the Transformer benchmark, we use  $3.5 \times$  lower time with less memory to obtain  $\sim 1.8\%$  improvement in validation error. Further, we are able to achieve better validation errors compared with Adam, even under the same training time budget.
Notation. We use boldfaced lower-case letters such as  $\pmb{x}$  to denote a vector and boldfaced upper-case letters such as  $\pmb{H}$  to denote a matrix.  $H_{i,i}$  and  $H_{i,j}$  denote the  $i$ -th column and row of the matrix respectively.  $\mathrm{Diag}(H_1, \ldots, H_k)$  denotes a block-diagonal matrix with diagonal blocks  $H_1, \ldots, H_k$ .  $\mathrm{vec}(H)$  flattens a matrix into a vector, while  $A \otimes B$  denotes the Kronecker product of two matrices.  $\| H \|_2, \| H \|_F$  denote the 2-norm and Frobenius norm of a matrix respectively.
2 Proposed Method
We consider training a deep neural network that takes an input vector  $\pmb{x}$  and produces an output  $f(\pmb{\theta},\pmb{x})$ . With a training set  $(\pmb{y}_i,\pmb{x}_i)_{i=1}^n$  consisting of  $n$  samples and a loss function  $\ell$ , the network parameters  $\pmb{\theta}$  are learned by minimizing the empirical loss  $\mathcal{L}$  over the training set:
$$
\min  _ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta}) := \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (f (\boldsymbol {\theta}, \boldsymbol {x} _ {i}), \boldsymbol {y} _ {i}).
$$
To solve the optimization problem, [10] proposed a full matrix variant of Adagrad, which uses a preconditioner matrix to partially capture the curvature information. Given the stochastic gradient  $\pmb{g}_t = \nabla \mathcal{L}_B(\pmb{\theta}_t)$  where  $\mathcal{L}_B$  represents the loss of a sampled mini-batch, the preconditioner  $H_t$  and the next iterate  $\pmb{\theta}_{t+1}$  are computed as follows:
$$
\boldsymbol {H} _ {t} = \boldsymbol {H} _ {t - 1} + \boldsymbol {g} _ {t} \boldsymbol {g} _ {t} ^ {T}, \quad \boldsymbol {\theta} _ {t + 1} = \boldsymbol {\theta} _ {t} - \alpha_ {t} \left(\boldsymbol {H} _ {t} ^ {1 / 2} + \epsilon \boldsymbol {I}\right) ^ {- 1} \boldsymbol {g} _ {t}, \tag {1}
$$
where  $\alpha_{t}$  is the step size and  $\epsilon$  is a small constant introduced for numerical stability. As manipulating the full preconditioner matrix  $H_{t}$  is computationally prohibitive, commonly used adaptive optimizers often rely on only the diagonal elements of the preconditioner matrix. Examples include the element-wise version of Adagrad [10] and many of its extensions such as Adam [20].
However, previous studies [16, 23] have shown that in many scenarios, non-diagonal preconditioning can lead to superior results. Hence, it is crucial to design an approximation of the preconditioner that captures more information compared to diagonal Adagrad, while still maintaining practicality for neural network training. This is the primary focus of our paper.
2.1 Block-diagonal Structure of the Preconditioner
Due to the substantial size of the full preconditioner matrix in deep neural networks, it is often impractical to employ it directly. Consequently, in line with the majority of existing literature
(a) The feed-forward layer with ordering  $H_{t}^{\mathrm{in}}$
(b) The feed-forward layer with ordering  $H_{t}^{\mathrm{out}}$
Figure 1: The second-moment matrices of two different fully connected layers in a Transformer (detailed configurations can be found in our experimental section; details of the visualization can be found in Appendix [6.2]). As we can see,  $H_{t}^{\mathrm{out}}$  shows a clearer block diagonal structure compared to  $H_{t}^{\mathrm{in}}$ , verifying our choice of grouping weights by output neurons.
(c) Self-attention output kernel with ordering  $H_{t}^{\mathrm{in}}$
(d) Self-attention output kernel with ordering  $H_{t}^{\mathrm{out}}$
[16,23,II], we precondition each layer separately. This effectively imposes a block diagonal structure assumption on the preconditioner:
$$
\boldsymbol {H} _ {t} = \operatorname {D i a g} \left(\boldsymbol {H} _ {t} ^ {1}, \dots , \boldsymbol {H} _ {t} ^ {L}\right),
$$
where  $L$  is the number of layers. From now on, we will approximate each  $\pmb{H}_t^l$  separately. Therefore, for notational simplicity, we will omit the layer index and refer to the preconditioner matrix for a single layer as  $\pmb{H}_t$ .
Consider a fully connected layer with weights  $W \in \mathbb{R}^{d_{\mathrm{in}} \times d_{\mathrm{out}}}$ , where  $d_{\mathrm{in}}$  and  $d_{\mathrm{out}}$  represent the number of input and output neurons, respectively. In Transformer networks, these dimensions can extend to thousands, making it impractical to store the full preconditioner  $H_t$  which has size  $d_{\mathrm{in}} d_{\mathrm{out}} \times d_{\mathrm{in}} d_{\mathrm{out}}$ . To address this issue, prior studies have identified a block-diagonal structure in  $H_t$  and applied a block-diagonal approximation. Specifically, [5,7] empirically demonstrated that the preconditioner of a fully connected layer exhibits a block-diagonal structure comprising  $d_{\mathrm{out}}$  blocks, each with a size of  $d_{\mathrm{in}}$ . In this structure, the weights connecting different input neurons to the same output neuron are grouped together as a block.
We visualize this block-diagonal structure in Figure 1. Specifically, we examine different fully connected layers of a Transformer model and plot the preconditioner using two different orderings:  $H_{t}^{\mathrm{in}}$  follows row-major ordering, grouping weights associated with the same input neuron, while  $H_{t}^{\mathrm{out}}$  follows column-major ordering, grouping weights associated with the same output neuron. Notably, we observe that  $H_{t}^{\mathrm{out}}$  demonstrates a much stronger block-diagonal structure, confirming our choice of block-diagonal approximation.
Although this particular block-diagonal structure has been observed in the literature, to the best of our knowledge, the underlying reason for this block-diagonal structure has not been analyzed mathematically. Here we show that the block-diagonal structure exists based on the following analysis. Assume we compute the preconditioner with a fixed parameter, then we can denote  $\mathbf{X} \in \mathbb{R}^{n \times d_{\mathrm{in}}}$  as the input of this layer for all the  $n$  samples, and the fully connected layer can be written as
$$
Z = X W,
$$
where  $\mathbf{W} \in \mathbb{R}^{d_{\mathrm{in}} \times d_{\mathrm{out}}}$  is the weight matrix and  $\mathbf{Z} \in \mathbb{R}^{n \times d_{\mathrm{out}}}$  is the output matrix. Let  $D_{\mathbf{W}} = \frac{\partial \mathcal{L}}{\partial \mathrm{vec}(\mathbf{W})}$  and  $D_{\mathbf{Z}} = \frac{\partial \mathcal{L}}{\partial \mathbf{Z}}$  denote the gradients of the loss function with respect to  $\mathbf{W}$  and  $\mathbf{Z}$ , respectively. Additionally, let us define the preconditioner as:
$$
\boldsymbol {H} = \left[ \begin{array}{c c c} \boldsymbol {H} ^ {(1, 1)} & \dots & \boldsymbol {H} ^ {(1, m)} \\ \vdots & \ddots & \vdots \\ \boldsymbol {H} ^ {(m, 1)} & \dots & \boldsymbol {H} ^ {(m, m)} \end{array} \right], \tag {2}
$$
where  $m = d_{\mathrm{out}}$  is the number of blocks. Since  $\pmb {H} = D_{\pmb{W}}D_{\pmb{W}}^{T}$  , we can derive the following lemma:
Lemma 1 Assuming the model is evaluated with a fixed parameter  $\theta$ , then any two diagonal blocks  $\pmb{H}^{(i,i)}$ ,  $\pmb{H}^{(j,j)}$  and two corresponding off-diagonal blocks  $\pmb{H}^{(i,j)}$ ,  $\pmb{H}^{(j,i)}$  for the preconditioner matrix of a fully connected layer will satisfy the following inequality:
$$
\left\| \boldsymbol {H} ^ {(i, i)} \right\| _ {F} ^ {2} + \left\| \boldsymbol {H} ^ {(j, j)} \right\| _ {F} ^ {2} \geq \left\| \boldsymbol {H} ^ {(i, j)} \right\| _ {F} ^ {2} + \left\| \boldsymbol {H} ^ {(j, i)} \right\| _ {F} ^ {2}. \tag {3}
$$
The proof is deferred to the appendix. Based on this lemma, we can conclude that the diagonal blocks have a higher Frobenius norm compared with the corresponding off-diagonal blocks.
However, directly using this block-diagonal structure is impractical for large-scale problems. Since there are  $d_{\mathrm{out}}$  groups and each group has size  $d_{\mathrm{in}}$ , storing this block-diagonal approximation requires  $O(d_{\mathrm{in}}^2 d_{\mathrm{out}})$  memory. To illustrate, a single fully connected layer in a Transformer with  $d_{\mathrm{in}} = 4096$  and  $d_{\mathrm{out}} = 1024$  (BERT-large) would occupy over 17GB of memory. To address this challenge, [32] proposed a solution by further dividing each block into mini-blocks. On the other hand, [5] explored an alternative approach that enforces all diagonal blocks to be identical for the larger layers. However, they demonstrated that this method adversely affects performance in practice.
2.2 Block Low-rank Approximation with Shared Basis
To simplify notation, we will concentrate on a single layer and omit the layer index moving forward. Let  $\pmb{H}_t \approx \mathrm{Diag}(\pmb{H}_t^{(1)}, \dots, \pmb{H}_t^{(m)})$  denotes the block-diagonal approximation of a layer with  $m$  blocks with each block containing  $b = d / m$  variables. For fully connected layers, we use the grouping mentioned in the previous subsection, so  $m = d_{\mathrm{out}}$  and  $b = d_{\mathrm{in}}$ , while in general, we let  $m$  be one of the dimensions for other weight tensors. To reduce the memory requirement, we consider an alternative approach to form low-rank approximations for each diagonal block:
$$
\boldsymbol {B} _ {t} ^ {(i)} \boldsymbol {R} _ {t} ^ {(i)} (\boldsymbol {B} _ {t} ^ {(i)}) ^ {T} \approx \boldsymbol {H} _ {t} ^ {(i)} = \sum_ {j = 1} ^ {t} \boldsymbol {g} _ {j} ^ {(i)} (\boldsymbol {g} _ {j} ^ {(i)}) ^ {T}, \quad \text {f o r} i = 1, \ldots , m,
$$
where  $k$  is the rank,  $B_{t}^{(i)} \in \mathbb{R}^{b \times k}$  and  $R_{t}^{(i)} \in \mathbb{R}^{k \times k}$  are the eigenvectors and eigenvalues of each block respectively. Using this approach, the memory cost is reduced to  $\mathcal{O}(dk)$ . Nevertheless, this memory cost is still  $k$  times larger than that of first-order optimizers, making it impractical for large-scale training tasks.
We introduce a novel shared-basis approximation scheme to further reduce memory costs. Our key insight is that the diagonal blocks share a similar subspace, which motivates our algorithm to use a joint basis for all the diagonal blocks, thereby reducing memory requirements. To achieve this, we propose the following shared-basis approximation for all the blocks within a single layer:
$$
\boldsymbol {H} _ {t} ^ {(i)} \approx \boldsymbol {B} _ {t} \boldsymbol {R} _ {t} ^ {(i)} \boldsymbol {B} _ {t} ^ {T}, \quad \text {f o r} i = 1, \dots , m, \tag {4}
$$
where  $B_{t} \in \mathbb{R}^{b \times k}$  is the orthogonal shared basis and  $R_{t}^{(i)} \in \mathbb{R}^{k \times k}$  denotes the coefficients for each block. Note that here we assume each  $R_{t}^{(i)}$  is a dense matrix to fully capture the information of each block. This shared-basis approach will reduce the memory complexity to  $O(bk + mk^2) = O(d_{\mathrm{in}}k + d_{\mathrm{out}}k^2)$  which can approach the cost of first-order methods when  $k^2 \leq d_{\mathrm{in}}$  and  $k \leq d_{\mathrm{out}}$ . We now give a detailed discussion on this shared-basis block low-rank approximation.
Why can we share basis? To support the idea of this shared-basis approximation, we focus on the analysis of fully connected layers and examine the low-rank approximation of each block. Following the analysis presented in Section 2.1 we consider the form of the preconditioner when evaluated on a fixed parameter  $\theta$ . In such cases, each diagonal block can be expressed as:
$$
\boldsymbol {H} ^ {(i)} = \boldsymbol {X} ^ {T} (D _ {\boldsymbol {Z}}) _ {:, i} ((D _ {\boldsymbol {Z}}) _ {:, i}) ^ {T} \boldsymbol {X}.
$$
Therefore, all the diagonal blocks lie in the row space of  $X$ , the input matrix of a certain layer. Additionally, we observe that in many applications,  $X$  tends to occupy a relatively low-rank space. Figure 2 shows the singular value distributions of the intermediate feature  $X$  of the Transformer. For all layers, the top  $10\%$  of the rank accounts for over  $80\%$  of the singular values, which suggests  $X$  is low rank. This property has also been observed in the literature [6].
If  $\mathbf{X}$  can be expressed as low-rank decomposition  $\mathbf{U}\mathbf{\Sigma}\mathbf{V}^T$  using SVD, where  $\mathbf{U} \in \mathbb{R}^{n\times k}$ ,  $\mathbf{\Sigma} \in \mathbb{R}^{k\times k}$ ,  $\mathbf{V} \in \mathbb{R}^{d_{\mathrm{in}}\times k}$ , we have
$$
\boldsymbol {H} _ {t} ^ {(i, i)} = \boldsymbol {V} \boldsymbol {\Sigma} \boldsymbol {U} ^ {T} (D _ {\boldsymbol {Z}}) _ {:, i} ((D _ {\boldsymbol {Z}}) _ {:, i}) ^ {T} \boldsymbol {U} \boldsymbol {\Sigma} \boldsymbol {V} ^ {T} = \boldsymbol {V} \boldsymbol {C} _ {i} \boldsymbol {V} ^ {T}
$$
for every  $i$ , which provides an intuition as to why we can share bases for different blocks. We further verify the performance of our proposed basis-sharing approximation is only slightly worse than forming an independent SVD of each diagonal block. We present these experiments in Appendix 6.3.
Figure 2: The singular value distribution of the intermediate feature matrix  $\mathbf{X}$  of 1024 samples after each Transformer block, which verifies the low-rank structure of  $\mathbf{X}$  in many real-world applications.
How to compute the shared basis in the online setting? Given the previous  $B_{t-1}, \{R_{t-1}^{(i)}\}_{i=1}^m$  and the current gradient  $\{g_t^{(i)}\}_{i=1}^m$ , we will describe the procedure for computing the updated approximate matrices  $B_t, \{R_t^{(i)}\}_{i=1}^m$ . Let
$$
\boldsymbol {M} _ {t} ^ {(i)} \equiv \boldsymbol {B} _ {t - 1} \boldsymbol {R} _ {t - 1} ^ {(i)} (\boldsymbol {B} _ {t - 1}) ^ {T} + \boldsymbol {g} _ {t} ^ {(i)} (\boldsymbol {g} _ {t} ^ {(i)}) ^ {T}
$$
be the block to be approximated. To derive the share bases, we concatenate all the blocks
$$
\bar {M} _ {t} = \left[ M _ {t} ^ {(1)}, M _ {t} ^ {(2)}, \dots , M _ {t} ^ {(m)} \right]
$$
and compute the SVD to derive its left singular vectors. This is equivalent to computing the top- $k$  eigenvectors of
$$
\boldsymbol {A} _ {t} \equiv \bar {\boldsymbol {M}} _ {t} \bar {\boldsymbol {M}} _ {t} ^ {T} = \sum_ {i = 1} ^ {m} \boldsymbol {M} _ {t} ^ {(i)} (\boldsymbol {M} _ {t} ^ {(i)}) ^ {T}.
$$
To efficiently compute the SVD of  $\mathbf{A}_t$ , we adopt the randomized SVD method [17]. This approach enables us to compute the basis without explicitly forming  $\mathbf{A}_t$  or any of the  $M_t^{(i)}$ . We initiate the process with an initial matrix  $\mathbf{Y} \in \mathbb{R}^{b \times (k + p)}$ , where  $p$  is the over-sampling constant (we set  $p = 0$  in experiments). Iteratively, we perform the following update:
$$
\boldsymbol {Y} \leftarrow \operatorname {Q R} \left(\boldsymbol {A} _ {t} \boldsymbol {Y}\right), \tag {5}
$$
where  $\mathrm{QR}(\mathbf{X})$  denotes the QR decomposition of matrix  $\mathbf{X}$ , returning the Q matrix. Using this approach,  $\mathbf{Y}$  will converge to the dominant eigenspace of  $\mathbf{A}_t$ . Additionally, since we can initialize  $\mathbf{Y}$  with the previous basis, only a few iterations are necessary in practice to obtain an accurate estimation.
Further, we have an improved way to compute the matrix product  $A_{t}Y$  in (5). The naive way to calculate  $A_{t}Y = \sum_{i=1}^{m} M_{t}^{(i)}(M_{t}^{(i)}Y)$  takes time complexity of  $\mathcal{O}(mbk^2) = \mathcal{O}(dk^2)$ . However, we can reduce some repeated computations as the blocks share the same basis. Let
$$
\boldsymbol {p} _ {t} ^ {(i)} \equiv \boldsymbol {B} _ {t - 1} \boldsymbol {R} _ {t - 1} ^ {(i)} (\boldsymbol {B} _ {t - 1}) ^ {T} \boldsymbol {g} _ {t} ^ {(i)},
$$
whose time complexity is  $\mathcal{O}(mbk) = \mathcal{O}(dk)$ . We have
$$
\begin{array}{l} \sum_ {i = 1} ^ {m} \boldsymbol {M} _ {t} ^ {(i)} (\boldsymbol {M} _ {t} ^ {(i)}) ^ {T} \boldsymbol {Y} = \boldsymbol {B} _ {t - 1} \left(\sum_ {i = 1} ^ {d / b} \boldsymbol {R} _ {t - 1} ^ {(i)} (\boldsymbol {R} _ {t - 1} ^ {(i)}) ^ {T}\right) (\boldsymbol {B} _ {t - 1}) ^ {T} \boldsymbol {Y} + \sum_ {i = 1} ^ {m} \boldsymbol {p} _ {t} ^ {(i)} (\boldsymbol {g} _ {t} ^ {(i)}) ^ {T} \boldsymbol {Y} \\ + \sum_ {i = 1} ^ {m} \boldsymbol {g} _ {t} ^ {(i)} (\boldsymbol {p} _ {t} ^ {(i)}) ^ {T} \boldsymbol {Y} + \sum_ {i = 1} ^ {m} \boldsymbol {g} _ {t} ^ {(i)} ((\boldsymbol {g} _ {t} ^ {(i)}) ^ {T} \boldsymbol {g} _ {t} ^ {(i)}) (\boldsymbol {g} _ {t} ^ {(i)}) ^ {T} \boldsymbol {Y}. \\ \end{array}
$$
Therefore, the time complexity for each power iteration is  $\mathcal{O}(m(k^3 +bk) + bk^2) = \mathcal{O}(mk^3 +dk + bk^2)$  (the QR decomposition takes  $\mathcal{O}(bk^2)$  time and is included in this complexity).
Forming the block low-rank preconditioner with shared basis. After obtaining the new shared basis  $B_{t}$ , we compute the block-dependent  $k$ -by-  $k$  coefficient matrices by
$$
\tilde {\boldsymbol {R}} _ {t} ^ {(i)} = (\boldsymbol {B} _ {t}) ^ {T} \boldsymbol {M} _ {t} ^ {(i)} \boldsymbol {B} _ {t},
$$
which is optimal in terms of minimizing the Frobenius norm reconstruction error for each block. Following the literature of frequent direction approaches [14], we remove the smallest eigenvalue to enable better online updates:
$$
\boldsymbol {R} _ {t} ^ {(i)} = \tilde {\boldsymbol {R}} _ {t} ^ {(i)} - \sigma_ {\min } (\tilde {\boldsymbol {R}} _ {t} ^ {(i)}) \boldsymbol {I}.
$$
To account for the missing mass in the low-rank approximation, we track a scalar value for each block using the recurrence relation:
$$
\rho_ {t} ^ {(i)} = \rho_ {t - 1} ^ {(i)} + \| \pmb {M} _ {t} ^ {(i)} - \pmb {B} _ {t} \pmb {R} _ {t} ^ {(i)} (\pmb {B} _ {t}) ^ {T} \|,
$$
which can be efficiently computed using power iteration. Finally, to derive the update direction  $\pmb{u}_t^{(i)}$ , we compute
$$
\begin{array}{l} \boldsymbol {u} _ {t} ^ {(i)} = - \left(\boldsymbol {B} _ {t} \boldsymbol {R} _ {t} ^ {(i)} \left(\boldsymbol {B} _ {t}\right) ^ {T} + \left(\rho_ {t} ^ {(i)} + \epsilon\right) \boldsymbol {I}\right) ^ {- 1 / 2} \boldsymbol {g} _ {t} ^ {(i)} \tag {6} \\ = - \boldsymbol {B} _ {t} (\boldsymbol {R} _ {t} ^ {(i)} + (\rho_ {t} ^ {(i)} + \epsilon) \boldsymbol {I}) ^ {- 1 / 2} (\boldsymbol {B} _ {t}) ^ {T} \boldsymbol {g} _ {t} ^ {(i)} + (\rho_ {t} ^ {(i)} + \epsilon) ^ {- 1 / 2} \boldsymbol {B} _ {t} (\boldsymbol {B} _ {t}) ^ {T} \boldsymbol {g} _ {t} ^ {(i)} - (\rho_ {t} ^ {(i)} + \epsilon) ^ {- 1 / 2} \boldsymbol {g} _ {t} ^ {(i)}. \\ \end{array}
$$
Detailed derivations can be found in Appendix 6.5. Here we follow Shampoo [16] to add the  $\epsilon I$  before taking the matrix square root, which is slightly different from the Adagrad formulation shown in (1).