Block Low-Rank Preconditioner with Shared Basis for Stochastic Optimization
Authors: Jui-Nan Yen
Sai Surya Duvvuri
Inderjit S Dhillon
Cho-Jui Hsieh
Venue: NeurIPS 2023 poster (2023)
Keywords: Second Order Optimization
Optimization for deep networks

Abstract: Adaptive methods with non-diagonal preconditioning have shown state-of-the-art results on various tasks. However, their computational complexity and memory requirement makes 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 and ultimately resulting in suboptimal performance. To reduce the time and memory complexity without sacrificing performance, we propose approximating each diagonal block 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 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.

Block Low-Rank Preconditioner with Shared Basis for Stochastic Optimization
Jui-Nan Yen UCLA juinanyen@cs ucla. edu
Sai Surya Duvvuri UT Austin saisurya@cs utexas.
Inderjit S. Dhillon Google and UT Austin inderjit@cs utexas edu
Cho-Jui Hsieh Google and UCLA
chohsieh@cs ucla. edu
Abstract
Adaptive methods with non-diagonal preconditioning have shown state-of-the-an results on various tasks. However, their computational complexity and memory requirement make it challenging scale these methods to modern neural network architectures. To address this challenge, some previous works have adopted block- diagonal preconditioners. However, the memory cost 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 efficiently maintain this shared- -arre 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.
Introduction
Diagonal adaptive methods such as Adagrad for each coordinate of the gradient, which is typically estimated inverse mean square of the histories of the coordinate gradient 1/v (. for gradient vector gt. This is equivalent to multiplying the gradient vector with a diagonal matrix, preconditioner. This approach is adopted in optimizers such Adam 20], RMSProp
10 incorporate an adaptive learning rate schedule
24 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 [6 4 GGT [ which also estimate off- off-diagonal gradient moments such as (gr)). for i * j and develop 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-diagone moment information the optimizer needs to store and the high computational requirement of the inverse operation in computing the preconditioner.
The second-moment matrix E. 9. constituting both diagonal and off-diagonal cross-moments is infeasible to store, however, the gradients from deep-neural networks can reveal structure in this matrix. To simplify this problem we analyze basic neural network block linear transformation X XW, where X E B and W E Rd Xhon and make two important observations: a) the
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

second-moment matrix Σ, 9:97. where E R is a vectorized form of the gradient with respect to W at iteration t, has block-matrix structure with daut blocks of size d x dn 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 (ddaa memory complexity to store all the blocks. To resolve this, a mini-block approach was used previously in (2)) which uses smaller blocks < din- 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 B E Rd x k < dm- and a different residual R(() E IR kx for each block such that { (BR(B B ..... BR (dm) B approximates the diagonal blocks.
We develop an online algorithm to update the shared basis B, for which we derive approxima- tion error guarantees demonstrating the nearness of { ( ..... BR(dw) B to the original matrix. We also this to prove
block-diagonal approximation of the gradient
((VT) regret upper bound in the online convex optimization framework [8] 283. which can
translate to convergence guarantees in non-convex stochastic optimization as in II. We compare the proposed algorithm with state- state-of-the-an second-order optimizers along with Adam, which is widely recognized as the most commonly used first-order optimizer. The evaluation is performed on standard autoencoder benchmark and Transformer benchmark comprising 19.3 million learnable parameters, Compared with other second-order methods, our algorithm utilizes lower memory and time 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.2x lower time, and 1.5x lower memory. In addition this, in the Transformer benchmark, we use 3.5x lower time with less memory to obtain ~ 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.
denote vector and boldfaced upper- case letters such as H to denote a matrix. H., and H,: denote the i-th column and row of the matrix respectively.
Notation. We use boldfaced lower-case letters such
,.,........ Hx) denotes block-diagonal matrix with diagonal blocks Hi,..., Ha. vec( H) flattens a matrix into a vector, while A @ B denotes the Kronecker product of two matrices. || 2,
H |r denote the 2-norm and Frobenious norm of matrix respectively.
2 Proposed Method
We consider training a deep neural network that takes an input vector æ and produces an output f(0, With training set (yi, (y1, =x)"== consisting of n samples and loss function E, the network parameters are learned by minimizing the empirical loss over the training set:
min -(0)==
:- Σ" =1 a()(0
To solve the optimization problem, I0 proposed full matrix variant of Adagrad, which uses preconditioner matrix to partially capture the curvature information. Given the stochastic gradient gi (L where L3 represents the loss of sampled mini-batch, the preconditioner H, and the next iterate 0.41 are computed as follows:
H. H + 9:96" 0 = 0. az ( eI)= g::
(1)
where a₁ is the step size and c is a small constant introduced for numerical stability. As manipulating the full preconditioner matrix H, 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 1 23 have shown that in many scenarios, non-diagonal precondition- ing can lead to superior results. Hence, it is crucial 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, often impractical to employ it directly. Consequently, in line with the majority of existing literature
2

---

(a) The ford-forward (b) The forard layer with ordering H layer with ordering H
(c) Self-attention output
(d) Self-attention output
kernel with ordering H
kernel with ordering H
The second-moment matrices of two different fully connected layers in Transformer (detailed configurations can be found in our experimental section; details of the visualization can be found in Appendix 6.22). As we can see, H shows clearer block diagonal structure compared to H verifying our choice of grouping weights by output neurons.
Figure 1:
16 23 we precondition each layer separately. This effectively imposes block diagonal structure assumption on the preconditioner:
H. Diag H H
where L is the number of layers. From on, will approximate each H separately. Therefore, for notational simplicity, we will omit the layer index and refer to the preconditioner matrix for a single layer H,,
Consider fully connected layer with weights W Rddadda, where dn and deui represent the number of input and output neurons, respectively. In Transformer networks, these dimensions extend to thousands, making it impractical to store the full preconditioner H, which has size din x dontt To address this issue, prior studies have identified block-diagonal structure in H, and applied block-diagonal approximation. Specifically, 7 empirically demonstrated that the preconditioner of fully connected layer exhibits a block-diagonal structure comprising dout blocks, each with a size of dm. this structure, the weights connecting different input neurons to the same output neuron are grouped together block.
We visualize this block-diagonal structure in Figure T Specifically, we examine different fully connected layers of Transformer model and plot the preconditioner using two different orderings: h follows row-major ordering, grouping weights associated with the same input neuron, while H follows column-major ordering, grouping weights associated with the same output neuron. Notably, we observe that H.S demonstrates much stronger block-diagonal structure, confirming our choice of block-diagonal approximation.
Although this particular block-diage 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 X E 1gaxd, as the input of this layer for all the n samples, and the fully connected layer can be written as
Z XW,
where W E X is the weight matrix and Z E RR is the output matrix. Let Dw aveet W) ac and Dz az denote the gradients of the loss function with respect to W and Z, respectively. Additionally, let define the preconditioner as:
8C
(1.1)
(1.)
H
(mm.)
H (m))
where m = dout is the number of blocks. Since H = Dw DTy. we can derive the following lemma:
Lemma Assuming the model is evaluated with a fixed parameter 0, then any two diagonal blocks (6,)), HUS and two corresponding off-diagonal blocks (5.)), H(3) for the preconditioner matrix of a fully connected layer will satisfy the following inequality:
| (8.) + (19)
H (1 + H (.)
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 duu groups and each group has size dm. storing this block-diagonal approximation requires O dour) memory. To illustrate, single fully connected layer in Transformer with din = 4096 and dout 1024 (BERT-large) would occupy over 17GB of memory. To address this challenge, RН] proposed a solution by further dividing each block into mini-blocks. On the other hand, [3] 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 single layer and omit the layer index moving forward. Let H2 R Diag( ( ( H!!)) denotes the block-diagonal approximation of a layer with m blocks with each block containing = d/m variables. For fully connected layers, we use the grouping mentioned the previous subsection, so m = deat and = da. 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:
00th ( R (T RO H(C))
(0) )T. for i ,,, ( E 9j
where k is the rank, B,") Rx and R. b are the eigenvectors and eigenvalues of each block respectively. Using this approach, the memory cost is reduced to Odk)). Nevertheless, this memory cost is still i times larger than that of first-order optimizers, making 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 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:
B.R B,, for i 1,...,
where B, E Roxk is the orthogonal shared basis and I2(") E R denotes the coefficients for each block. Note that here we assume each R!") 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
(d + )) which can approach the cost of first-order methods when k.2 < dn, and k S dout- We now give detailed discussion on this thared-basis block low-rank approximation.
Why can we share basis? To support the idea of this shared-baae 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 fixed parameter In such cases, each diagonal block can be expressed as:
H ( = x ((z)..))) 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 X can be expressed as low-rank decomposition UΣVT using SVD, where U E g E k V E jdnukk we have
H H -VUT (D2)).((((()))))
VCVV
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

---

Accumulated Ratio of Singulor Values
10-
Block0
Block1
Block2
Block3
Block4
02
Blocks
20
43
60
80
160
Percentage Singular
Figure 2: The singular value distribution of the intermediate feature matrix X of 1024 samples after each Transformer block, which verifies the low-rank structure of X in many real-world applications.
How to compute the shared basis in the online setting? Given the previous B (R and the current gradient (g(") m
J=m will describe the procedure for computing the updated approximate
matrices B., (R2 (0) pm 1==1- Let
M B (B
9i ())
be the block to be approximated. To derive the share bases, concatenate all the blocks
M. = [M,) M, (2)
(m)
and compute the SVD derive its left singular vectors. This is equivalent computing the top-k eigenvectors of
A. = M.M. Σ (M,
To efficiently compute the SVD of A,, we adopt the randomized SVD method This approach
enables us to compute the basis without explicitly forming A, or any of the ( (*) We initiate the
process with an initial matrix Y E (k+(k)))) where is the over-sampling constant (we set 0 in experiments). Iteratively, we perform the following update:
Y - QR(AY).
(5)
where QR(X) denotes the QR decomposition matrix X, returning the Q matrix. Using this approach, Y will converge the dominant eigenspace of A,. Additionally, since initialize 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,Y in (). The naive way to calculate A,Y En M (M Y) takes time complexity of ((mbk²) Oddk². However, we can reduce some repeated computations as the blocks share the same basis. Let
p (1) B R. (B.r))
whose time complexity is O(mbk) (dk)). We have
M (M. = B
d/b
R(0)
(B + P
Pe (
Σn ((g (T)T T (g. (.)
Therefore, the time complexity for each power iteration is O(m +bb) +bk2) = +ch+ (the QR decomposition takes Obb?²) time and is included in this complexity).
5

---

Forming the block low-rank preconditioner with shared basis. After obtaining the new shared basis B,, we compute the block deeeedeenddent k-by-k coefficient matrices by
R (
(B.)
M, (8) B,,
which is optimal in terms of minimizing the Frobenius Following the literature of frequent direction approaches [4], enable better online updates:
reconstruction error for each block. we remove the smallest eigenvalue to
RS
minn (R f2 M.
To account for the missing mass in the low-rank approximation, we track scalar value for each block using the recurrence relation:
Pt Pi-1 M B.R, (B.)*1,,
n (3)
which can be efficiently computed using power iteration. Finally, to derive the update direction u (*) we compute
( (B.) + (p: + +c)1) ui = (B.R
(6)
(R + (pi
( +))))) Di ( + (pi (, - -(p)) +
1/2 (
gt
Detailed derivations can be found in Appendix 6.5 Here we follow Shampoo [16] to add the eI before taking the matrix square root, which is slightly different from the Adagrad formulation shown in
Our overall algorithm is summarized in Algorithm The time complexity of the algorithm is dominated by the matrix products in power iterations, (m + dk bb.). The memory complexity is (bk + mk?) for storing the shared basis and the k x is dense coefficient matrices of each block.
For simplicity, the above discussions order tensor parameters, such as the multi-head self-attention layer in Transformers, we adopt a grouping strategy based of the axes. Further discussions can be found in Appendix 6.1
mainly for two two-dimensional matrices. For handling higher-
Algorithm Shared-Basks Low Rank Block-Diagonal Adagrad
Initialize: coefficients R 0, shared bases B 0
for = 1...7 do
Compute stochastic gradient gr
Obtain shared basis B, by running randomized SVD (Eq. 5 for A,
Compute R2) (3)
(B,) M,
B
Compute R(()) R2(1) Omin R
I for 1, 1,...,m
Compute (4) (8)
-i-1 + || (5)
B.R,
(B.) for 1,...,m
Compute the preconditioned update u, -B.R
(1) (B.) + (p) + -)())1//1))) (1) for
I..... m
Update model parameters 0, 0,-1 + a
end for
2.3 Regret Bound Analysis
Following previous works 16
provide a convergence analysis of the proposed algorithm
in the online convex optimization framework [8] 28 In the online convex optimization setting, a parameter 0, E K. is chosen iteratively, where K. is convex decision set. After the decision of 8,, convex loss function ft revealed, which can be selected adversarially. The regret suffered by the algorithm up to step T is defined
Regret, Σn T - min Σ
f((0,)
1=1 fc(0).
OEK
We give the following theorem for proposed method. All the proofs can be found in the Appendix.

---

Theorem Algorithm gives the following regrer
m Regrety S V2D > Tr((H,))112 + v2b > 1=1
(p
where D is the diameter of the constraint set K. and H ΣT= 91 n 10 TT
To bound the second term in Theorem D we further derive the following theorem.
Theorem Algorithm guaranntess
Σn (0)))12 (3m²)
Tr(H
and thus Regrety OVTT).
This shows that our proposed method is optimal asymptotically in terms of T. With some further assumptions, we can further derive bounds of p in terms of the lower eigenvalues of H We defer the results to the Appendix 6.8
2.4 Comparison with Other Methods
We compare the time and space complexity of the proposed algorithm with other second-order methods. We consider the time and memory cost for preconditioning neural network layer of size di x dom- The results are summarized in Table !
As mentioned earlier in Section 2 the full matrix version of Adagrad [CC] poses a significant challenge practice due to the storage and matrix inversion of full preconditioner with dimensions di de To address this issue, previous studies such as 3 22 proposed low-rank approximation for the preconditioner, which reduces the memory cost to (dnn with denoting the rank. However, with this reduced cost, it remains impractical for modern neural network architectures.
Inspired from the observed empirical block-diagonal structure variant of Adagrad and Adam. However, due to the space requirements, storing the block diagonal preconditioner with a block size of b incurs a memory complexity of (db). As result, they are limited to using small block sizes in practice. Similarly, [5] proposes For convolutional layers, they group parameters based on filters, resulting in block size of r² when the kernel size is In fully-connected layers, parameter grouping is based on output neurons, yielding a block size of ((din) and a memory complexity of O(ddm). Additionally, they propose maintaining the average of the block diagonal for larger layers, resulting in an approximation of & R, where R represents the average of all diagonal blocks. This approach resembles one side the Kronecker product methods in the next paragraph, but their research demonstrates some performance loss associated with the approximation.
26
propose a block diagonal
Fisher method.
Another popular approach, such Shampoo 16 and K-fac 2 involves grouping parameters based on both the input and output neurons, resulting in Kronecker product approximation L & R for the preconditioner. In order to enhance the time and space complexity of Shampoo, [ introduces a technique called blocking, which partitions the parameters into smaller blocks. On the other hand, proposes Sketchy as a low-rank variant of Shampoo.
Based the above discussion, except for Sketchy, all other methods exhibit higher memory plexity compared to our method. Under the same k our memory complexity is larger than Sketchy, This is due to the fact that are approximating the individual diagonal blocks, but Sketchy only approximates the average of the diagonal blocks (but in both orders). As a result, Sketchy fails to capture heterogeneity between different diagonal blocks. In our experimental evaluation, we demonstrate that when employing the same memory usage (by varying the value of k), our proposed method outperforms Sketchy.
3 Experimental Results
We evaluate the performance using standard Autoencoder benchmark 27 on the MNIST dataset
and larger Transformer model [30 the Universal Dependencies dataset 24
7

---

Table 1: Time and space complexity comparison for a single parameter matrix of size du x dout.
Algorithm
Time Complexity
Space Complexity
Full Matrix Adagrad
10
de de douck dindlouki d + d dinthaul
d de in d douk din donk de + doot din taut d + dourt doub dak + doak?
SON-FD
Ada-FD
Shampoo [6
Block Shampoo
44
Sketchy |||
dudoouk
Block Diagonal Adam
32
tie dvaab
Ours
don k + dakk? 2 + douk
We adopt A = 32 as the default rank for our methods. For randomized SVD, we set the oversampling parameter to and the number of iterations to The impact of these parameters detailed in our ablation study in Appendix 6.10, Similar to Shampoo, we use the grafting technique [2] in our method. We the grafting type to RMSPROP RASPROP_NORMALIZER) _NORRALLZED.. We compare the proposed optimizer with the widely used Adam optimizer (20) along with the following second-order optimizers:
we adopt the implementation of 4 and keep most default settings with the following differences. Following the setting of [T]. for Shampoo, the preconditioning starts at step 101 for stability and the preconditioners are updated every 10 steps for speed (preconditin tioningiiiiinnniiiii compute steps=10). Furthermore, we explore different block sizes for Shampoo, as suggested by [4. This approach involves partitioning the parameters into smaller blocks to improve both speed and memory utilization. We denote Shampoo with block size x' as Shampoo(z).
Shampoo 16:
Sketchy || this is low-rank approximation version of Shampoo. We follow the setting in their paper to use rank 256 and block size 1024.
Block Adam 32 : we implement the idea of dioco-diagonal approximation with a smaller chose the block size to be a divisor of d Specifically, a block size of 10 is used for the autoencoder benchmark and block size of 8 used for the Transformer benchmark.
block proposed in (32] We group the blocks based on output neurons,
We conduct hyperparameter tuning with random hyperparameter search over the search space defined in the appendix. Note that we search over one minus momentum instead of searching momentum directly. Each hyperparameter is sampled from log-uniform distribution. For the autoencoder benchmark, conduct 180 trials of random search on one NVIDIA RTX 2080Ti GPU with 1IGB memory. For the Transformer benchmark, we conduct 60 trials of random search on one NVIDIA RTX A6000 GPU with 48GB memory.
3.1 Autoencoder Benchmark
This autoencoder benchmark has been commonly used to evaluate various optimizers in the literature [23 3 and follow the same setting in previous papers. Specifically, con- duct 100 epochs of training using the MNIST dataset. The layer sizes of the autoencoder are [1000, 500, 250, 30, 250,500 500, 1000), with total number of 2.83M learnable parameters. We use tanh as the non-linearity. The batch size is 1000. A linear warmup of 5 epochs is used for learning rate scheduling followed by a linear decay
In Table 2 we observe that the proposed method has better performance than all the methods except for Shampoo( 1000). Compared to Shampoo( (1000), our performance is similar while using 3.2x lower time, and 1.5x lower memory. The training loss curve can be seen in Figure
3.2 Transformer Benchmark
For the Transformer benchmark, we adopt the part-of-speech tagging task from the Google flax repository [19], which uses the ancient Greek data from the Universal Dependency data sets 24]. The model has 6 Transformer layers. The embedding dimension size is 5 512 and the MLP dimension size is

---

Table 2: Experimental results on the autoencoder benchmark.
Optimizer
Adam
Block Adam
Shampoo(250)
Shampoo(1000)
Sketchy
Ours
Train Loss
54.66
51.78
51.56
51.22
52.41
51.49
Time (s)
50
230
165
553
1998
140
Memory (MB)
1593
2619
2649
364
2799
1777
Curs
Ours
873
Adam
Adem
Block Adam
654
Adam
Shampco(250)
Shampoc(250)
Train Loss
Shampco(1000)
Train Loss
Shempoo(1C39)
60.0
Sketchy
880
Skatchy
375-
20
60
09
100
200
200
Epoch
Time (5)
(a) Train loss (epoch)
(b) Train loss (time)
Figure 3: Train loss on the autoencoder benchmark. We run Adam for 400 epochs so that it takes
roughly the same time as the other methods.
2048. For the self-attention layers, the number of heads is and the query, key, value dimensions 512. The total number of learnable parameters is 19.3M. We use the default square root learning rate decay scheduling but shorten the training period to 200 epochs and the warmup period 40 epochs. The batch size is 128. By default, Shampoo excludes the embedding layer from preconditioning as one of its dimensions exceeds 4096. We also follow this setting for a fair comparison, even though our method can be used to precondition the embedding layer due to the reduced memory complexity. We adopt the direction obtained from sgd with momentum for the embedding layer.
In Table 3, we can see that our proposed method performs the best. Compared to Shampoo 1024), we 3.5x lower time with less memory to obtain 1.8% improvement the validation error. The training loss curve and the validation accuracy curve can be seen in Figure g where the proposed method outperforms others in terms of both training time and number of training steps.
Related Work
Second order methods. For the optimization of ill-conditioned functions, second order methods often more effective than first order methods. This includes Newton's method and the quasi-Newton methods 21) Subsampled versions of these methods 15 25 problems of a larger scale. However, for deep learning problems, the memory requirements associated with these methods become prohibitive. To address this problem, the majority of existing literature preconditions each layer separately. Kronecker product approximation P3 [6] was proposed to reduce the time and space complexity and has shown promising results on several large-scale training tasks. Further, several other approximation schemes have been proposed to exploit second order information, discussed in Section 2.4
have also been proposed for solving
Table 3: Experimental results the Transformer benchmark.
Optimizer
Adam
Block Adam
Shampoo (512)
Shampoo (1024)
Sketchy
Ours
Validation Accuracy
67.36
70.43
68.79
68.81
68.32
70.63
Time (min)
123
193
325
564
717
161
Memory (MB)
19157
21598
19161
21209
19563
19553

---

Ours
Adam
Adam
Block
Block Adam
Shampoo(512)
Shampeo(512)
Train Loss
Shampoo
Train Loss
Shampoo(1024]
Sketchy
Sketchy
3000
10000
Step
Time (min)
(a) Train loss (step)
(b) Train loss (time)
0.7
9.-
Best Validation Accuracy
0.8
atiin Accuracy
0.5
0.4
Ours
Best Validattion
Ours
Adam
0.3
Block Acam
0.2
Shampoo(S12)
Shampoo(512)
03-
Shampoo (0024)
01-
Shampoo(1024)
Sketchy
Sketchy
0.
2500
$200
710C
10002
12500
29000
11900
Step
(c) Validation accuracy (step)
(d) Validation accuracy (time)
Figure 4: Train loss and the best validation accuracy the Transformer benchmark.
Frequent direction method. Given a stream of vectors, the frequent direction method tains a low-rank sketch matrix to approximate the covariance matrix of these vectors. The frequent direction method is proved to be optimal in space complexity. This suggests we can approximate the covariance matrix decently even when the vectors are given in a streaming manner. The frequent direction method has been used previous work for approximating the preconditioner matrix in stochastic optimization [ 22
S main-
5
Conclusions
We have introduced novel adaptive optimizer that utilizes a shared-basis block low-rank approxi- mation to approximate the full matrix preconditioner. Through evaluations on an autoencoder and Transformer benchmark, our proposed optimizer has demonstrated superior performance compared to existing first and second-order optimizers, while incurring small memory and time overhead. We also provided convergence guarantees for the proposed optimizer.
Limitation
The space complexity of our proposed method scales quadratically with k; the rank used in the block low-rank approximation. While our empirical results indicate that 32 adequate for Autoencoder and Transformer benchmarks examined in our study, this may still pose burden when training large-scale models. Exploring methods to reduce the dependence on while preserving performance will be interesting future direction.
Acknowledgements
We thank the anonymous reviewers for their helpful comments. JNY and CJH are supported in part by NSF under 2008 173, 2048280, 2325121 and 2331966, Cisco and Sony.
10

---

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, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803, 2020.
[3] Ehsan Amid, Rohan Anil, and Manfred Warmuth. Locoprop: Enhancing backprop via local loss optimization. In International Conference on Artificial Intelligence and Statistics, pages 9642. PMLR, 2022.
[4] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint (rxXiv2002.09011 2020.
[5] Achraf Bahamou, Donald Goldfarb, and Yi Ren. mini-block Fisher method for deep neural networks. In International Conference on Artificial Intelligence and Statistics, pages 9191-9220. PMLR, 2023.
[6] Patrick Chen, Hsiang-Fu Yu, Inderjit Dhillon, and Cho-Jui Hsieh. Drone: Data-aware low-rank compression for large nlp models. Advances in neural information processing systems, 34: 29321-29334, 2021.
[7] Ronan Collobert. Large scale machine learning. Technical report, Université de Paris VI, 2004.
[8] Li Deng. The MNIST database of handwritten digit images for machine learning research [best of the web]. JEEE signal processing magazine, 29(6):141-111 2012.
[9] John E Dennis, Jr and Jorge J Moré. Quasi-Newton methods, motivation and theory. SIAM review, 19(1):46-99, 1977.
[10] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[11] Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory- efficient adaptive regularization with frequent directions. arXiv preprint arXi::2302.03644 2023.
[12] Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory- efficient adaptive regularization with frequent directions. arXiv preprint arXiv:2302.0372 2023.
[13] Mina Ghashami, Edo Liberty, Jeff M Phillips, and David P Woodruff. Frequent directions: Simple and deterministic matrix sketching. SIAM Journal on Computing, 45(9):172211792. 2016.
[14] Mina Ghashami, Edo Liberty, Jeff M Phillips, and David Woodruff. Frequent directions: Simple and deterministic matrix sketching. SIAM Journal on Computing, 45(5):172--1792, 2016.
[15] Robert Gower, Donald Goldfarb, and Peter Richtárik. Stochastic block BFGS: Squeezing more curvature out of data. In International Conference on Machine Learning, pages 1869-1878, PMLR, 2016.
[16] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842-1850. PMLR, 2018.
[17] Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review, 53 (2):217-288, 2011.
[18] Elad Hazan et Introduction to online convex optimization. Foundations and Trends in Optimization, 213-4):1573355, 2016.

---

[19] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: neural network library and ecosystem for JAX, 2023. URL http:// http:://gtthbbb.. com www.th....ommgooolele.o coog://l///oo///la
[20] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[21] Dong Liu and Jorge Nocedal. On the limited memory BFGS method for large scale optimiza- tion. Mathematical programming, 45(1-3):503-5222 1989.
[22] Haipeng Luo, Alekh Agarwal, Nicolo Cesa-Bianche and John Langford. Efficient second order online learning by sketching. Advances in Neural Information Processing Systems, 29, 2016.
[23] James Martens and Roger Grosse. Optimizing neural networks with Kronecker- approximate curvature. International conference on machine learning, pages 2408-244 2411.. PMLR, 2015.
[24] Joakim Nivre, Catie-Cathhrine de Marneffe, Filip Ginter, Jan Hajič, Christopher D Manning, Sampo Pyysalo, Sebastian Schuster, Francis Tyers, and Daniel Zeman. Universal dependencies v2: An evergrowing multilingual treebank collection. arXiv preprint arXiv:2004. 2020.
[25] Farbod Kloosta-Khorasann and Michael W Mahoney. Sub-sampled Newton methods. Mathemat- ical Programming, 174:293-326, 2019.
[26] Nicolas Roux, Pierre- Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algorithm. Advances Neural Information Processing Systems, 20, 2007.
[27] Jürgen Schmidhuber. Deep learning in neural networks: overview. Neural networks, 61: 85-117,2011 2015.
[28] Shai Shalev-Shwartz al. Online learning and online convex optimization. Foundations and Trends in Machine Learning, 4(2):107-194, 2012.
[29] T. Tieleman and G. Hinton. Lecture S--..MPPppppp Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning. 2012.
[30] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[31] Yuanyu Wan and Lijun Zhang. Efficient adaptive online learning via frequent directions. IEEE Transactions Pattern Analysis and Machine Intelligence, 44(10):69 10-6923, 2021.
[32] Jihun Yun, Aurelie Lozano, and Eunho Yang. Adablock: SGD with practical block diagonal matrix adaptation for deep learning. In International Conference on Artificial Intelligence and Statistics, pages 2574-2606. PMLR, 2022.
12


Reviews:
[{"review_id": "bT66oMTw4V", "reviewer": "Reviewer_auuW", "summary": "This paper implements an online block version of AdaGrad to capture off-diagonal information in the preconditioner at each block level. A low-rank approximation is used for each block, which is a common approach in the literature to decrease time and memory complexities. To further reduce the memory usage, a shared basis approach across all blocks is proposed, which is justified both theoretically and practically. The proposed optimizer is tested on two benchmarks which are used in previous works.", "questions": "- how does the proposed optimizer behave with smaller ranks  $k \\in \\{1, 2, 4, 8 \\}$? Would it be possible to provide some ablation study for this?\n- in the appendix it is stated that the hyperparameter search was also performed for $1-\\beta_2$, which seems to be a second order momentum parameter. How is this applied to the proposed optimizer? Does $\\beta_2$ decay the outer product of the gradient $g_t$ or is it employed in the momentum applied to the update, such as $v_t^{(i)} = \\beta_2 v_{t-1}^{(i)} + (1-\\beta_2) u_t^{(i)}$?\n- would it be possible to run Sketchy and the proposed optimizer with the same rank to have a fair comparison against them? The reviewer is aware of the relationship between the rank and block size for Sketchy and believes that an ablation study on the block size of it with fixed rank (for Sketchy and the proposed work) should have been performed. (The most important thing is to have similar ranks to be able to compare memory and, of course, block size also matters for Sketchy)\n- can you compare against Adam at least on small BERT models (tiny with 4M params and/or mini with 11M params or even larger variants if time and memory allows) on GLUE-MNLI and SQuAD v1/v2? These are baseline tasks and models for NLP and benchmarking against them would be important. Moreover, in the manuscript at lines 290-292 it is stated that the proposed optimizer can handle large embedding layers and this would be a good test.\n- do the learning rate scheduler and weight decay / L2-regularization (coupled or decoupled) affect the results in any way?\n- based on the formula for the update provided at line 190, the final expression can be rewritten as $u_t^{(i)} = -B_t \\left(R_t^{(i)} + (\\rho_t^{(i)} + \\epsilon)I\\right)^{-1/2}B_t^T g_t^{(i)} - (\\rho_t^{(i)} + \\epsilon)^{-1/2}(I - B_t B_t^T) g_t^{(i)}$. Can you please explain why the term $I - B_t B_t^T$ appears in the final expression of $u_t^{(i)}$?\n- if we assume that the shared base $B_t$ is orthogonal, e.g. $B_t B_t^T = I$, then working out the first expression for the update would result in $u_t^{(i)} = B_t \\left(R_t^{(i)} + (\\rho_t^{(i)} + \\epsilon)I\\right)^{-1/2}B_t^T g_t^{(i)}$. However, this is different from the update in formula $(1)$ and the reviewer believes that the update should have been $B_t \\left[(R_t^{(i)} + \\rho_t^{(i)})^{1/2} + \\epsilon I\\right]^{-1} B_t^T g_t^{(i)}$. It might be possible that the reviewer misses something, can you please clarify why the update $(1)$ is implemented differently?\n\nSince the reviewer cares about the correctness of the formulas, some potential typos that do not influence the decision are highlighted:\n- at the formula of $A_t$ between lines 167 and 168 in the manuscript, the reviewer believes that a transposition operator is missing, even though the matrices $M_t^{(i)}$ are symmetric (this is not stated in the paper). As a consequence, it might also be needed in the formulas between lines 179 and 180 (here, the index in the first sum goes up to $d/b$ instead of $m$, which is essentially the same thing).", "limitations": "The paper has a brief section about limitations. The reviewer would like to add that the lower memory usage and lower training time would have a positive impact in reducing power usage and carbon emissions.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the paper provides convergence analysis in the Online Convex Optimization framework, which might serve as a good starting point for methods that would build on top of this paper\n- the idea of sharing the basis across all blocks is justified theoretically and practically (visually) in section 2.", "weaknesses": "- as stated in the limitations section, the memory complexity scales quadratically with rank k, which might be partially solved by decreasing k. However, since the basis is shared, there might be a lower bound on k\n- the results do not have confidence intervals and the paper does not state how many random seeds were used. This is mandatory especially when the results for each method have similar values\n- in the experiments, Sketchy had rank 256, which is much larger compared to k=32 for your shared-basis approach. The reviewer believes that this is not a fair comparison and this could be fixed with some additional experiments to cover these issues", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review!\n\n> the memory complexity scales quadratically with rank k. \n\nAs stated in the paper, the memory complexity scales quadratically with the rank k. But since each block has a different coefficient matrix, our approach requires a much smaller k compared with Sketchy. As shown in the experiments, our method with k=32 outperforms Sketchy with k=256. \n\nAdditionally, as shown in Figure 4 of the appendix, k=32 gives a decent approximation to the block diagonal. This is the main reason we use k=32 in our experiments. Our further experiment below shows that we can use an even smaller k to obtain decent performance.\n\n> the results do not have confidence intervals and the paper does not state how many random seeds were used\n\nThank you for the advice. We added 10 runs of the best hyperparameter on the transformer benchmark and the results are as follows. \n\n\n|Validation Accuracy |Adam|Block Adam|Ours (rank 32)\n| --- | --- | --- | --- | \n|Mean|67.70|70.36|70.51|\n|Variance|0.16|0.21|0.03|\n\n\n\n> how does the proposed optimizer behave with smaller ranks\n\nFor the autoencoder benchmark,\n\n\n\n||Adam|Ours (rank 1)|Ours (rank 2)|Ours (rank 4)|Ours (rank 8)|Ours (rank 16)|Ours (rank 32)|\n|---|---|---|---|---|---|---|---|\n|Train Loss|54.66|55.29|53.51|52.88|52.42|51.94|51.49|\n\n\nFor the transformer benchmark, we use the best parameters obtained from rank=32, and run it with rank=1 and rank=4 for 10 runs.\n\n|Validation Accuracy |Adam|Block Adam|Ours (rank 1)|Ours (rank 4)|Ours (rank 32)|\n|---|---|---|---|---|---|\n|Mean|67.70|70.36|70.30|70.41|70.51|\n|Variance|0.16|0.21|0.04|0.03|0.03|\n\nThe results suggest that decreasing the k will slightly hurt the performance. Further, the drop on the validation error (for transformer training) is small when decreasing k. \n\n\n> Does $\\beta_2$ decay the outer product of the gradient $g_t$\n\nYes, it decays the outer product of the gradient $g_t$. We will clarify this in the revision.\n\n\n> would it be possible to run Sketchy and the proposed optimizer with the same rank to have a fair comparison against them?\n\nWe added the following experiment to compare our method with Sketchy under the same rank. \n\nFor the autoencoder benchmark,\n\n\n\n|            | Sketchy (block 250, rank 32) | Sketchy (block 1000, rank 32) | Sketchy (block 1000, rank 256) | Ours (rank 32) |\n| ---------- | ---------------------------- | ----------------------------- | ------------------------------ | -------------- |\n| Train Loss | 52.73                        | 52.76                         | 52.41                          | 51.49          |\n\n\n\nFor the transformer benchmark, we use the best parameters obtained from Sketchy rank=128 to run rank=32 for 10 runs.\n\n\n| Validation Accuracy | Sketchy (block 1024, rank 32) | Ours (rank 32) |\n| ------------------- | -------------------------------- | -------------- |\n| Mean                | 67.52                            | 70.51          |\n| Variance            | 0.32                             | 0.03           |\n\nThe results suggest that our method outperforms Sketchy under all these settings. \n\n> do the learning rate scheduler and weight decay / L2-regularization (coupled or decoupled) affect the results in any way?\n\nWe use the default learning rate scheduler in both benchmarks. The auto-encoder benchmark comes with a linear decay and the Transformer benchmark comes with a square root decay. We outperform the other methods on both benchmarks, showing that our method is not sensitive to the learning rate scheduler. \n\n> based on the formula for the update provided at line 190, the final expression can be rewritten as $u_t^{(i)}=-B_t \\left( R_t^{(i)}+(\\rho_t^{(i)}+\\epsilon)I\\right)^{-1/2}B_t^Tg_t^{(i)}-(\\rho_t^{(i)}+\\epsilon)^{-1/2}(I-B_tB_t^T)g_t^{(i)}$\n\nYou are correct that the final expression can be rewritten as above. The term $I-B_tB_t^T$ appears in the final expression of $u_t^{(i)}$ as the projection matrix to the orthogonal space of the shared basis. So the two parts of the update can be seen as the update in the shared basis space (involving both $R_t$ and $\\rho_t^{(i)}+\\epsilon$) and the update in the orthogonal space (only involving $\\rho_t^{(i)}+\\epsilon$). \n\n> However, this is different from the update in formula $(1)$\n\nThanks for pointing this out! Indeed this expression is different from the update in formula $(1)$. The difference between the two is adding $\\epsilon$ before or after the matrix square root. The Adagrad paper uses the latter, while we follow Shampoo to use the former in our implementation. Since both can be implemented efficiently, we will fix this in the writing and compare them in future experiments.\n\n> Since the reviewer cares about the correctness of the formulas, some potential typos that do not influence the decision are highlighted\n\nThanks! We have fixed them in our writing.\n"}, {"review_id": "TUjSJrtNNm", "reviewer": "Reviewer_Z6xV", "summary": "This paper studies how to improve the memory efficiency of adaptive methods with (non-)diagonal preconditioning. The authors 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. The authors also provide theoretical analysis and empirical studies to demonstrate the advantage of the proposed method.\n\n----------------------------\nI have read the author’s rebuttal. Although I feel that my concerns were not fully addressed, I am willing to raise my score after seeing the seriousness of the authors' response, including their careful approach to me and the other reviewers.", "questions": "I was just wondering about the performance of the proposed method in a broader dataset.", "limitations": "(same as Weaknesses, but looking forward to authors' response)\n\n1. The research problem studied in this paper is practical and significant.\n2. Theoretical analysis is rich.\n3. Empirical studies on various datasets.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The research problem studied in this paper is practical and significant.\n2. Theoretical analysis is rich.\n3. Empirical studies on various datasets.", "weaknesses": "1. The document is not very easy to understand.\n2. the proposed methodology is not clearly superior to the baseline methodology.\n3. The authors did not compare the use of multiple modeling architectures", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable feedback!\n\n> The document is not very easy to understand.\n\nWe are sorry for the inconvenience and will further polish the paper.\n\n> The proposed methodology is not clearly superior to the baseline methodology.\n\nOur proposed method outperforms Adam on both benchmarks. On the transformer benchmark, we gain over 3% in validation accuracy over Adam with similar time and memory consumption. Our proposed method beats other baselines with less time and memory consumption.\n\n> The authors did not compare the use of multiple modeling architectures\n\nWe conduct experiments on both the autoencoder and the transformer architectures. We will add more in future experiments.\n"}, {"review_id": "0vpcFAtob4", "reviewer": "Reviewer_ssTg", "summary": "This paper proposes an adaptive optimizer similar to Shampoo and GGT, which considers non-diagonal components of the preconditioner matrix. They reduce the cost of the preconditioner by using a unit-wise block diagonalization on the output units and a further low-rank approximation with a shared basis between the diagonal blocks.", "questions": "What is the effect of removing the smallest eigenvalue from the R matrix? The experiments don't seem to compare with and without this difference.", "limitations": "This paper does not have any potential negative societal impact.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The time and space complexity of preconditioned optimizers is a critical issue, especially when models keep increasing in size at a rapid pace. The proposed method strikes a good balance between accurately representing the preconditioner matrix and its time/space complexity. They compare against Adam, block Adam, Shampoo, and Sketchy on autoencoder and transformer benchmarks, and demonstrate that their", "weaknesses": "The novelty with respect to the existing work [4] and [11] seems incremental, and the improvement over block Adam [32] is marginal (0.2% increase in validation accuracy). The reason why the proposed method shows better validation accuracy than Shampoo is also unclear. Since the proposed method introduces further approximation to Shampoo in order to reduce the time and space complexity, any improvement in validation accuracy is certainly not by design. They choose a fixed rank k=32 for their experiments, but this number seems to be chosen arbitrarily.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable response.\n\n> The novelty with respect to the existing work [4] and [11] seems incremental\n\n\nWe are sorry that our writing might cause some confusion, but our method is very different from the work of Shampoo [4] and Sketchy [11]. Shampoo and Sketchy approximate the second-moment matrix with a Kronecker product approximation so all the diagonal and non-diagonal blocks are the same up to a scalar. In contrast, we focus on the finer approximation of the diagonal blocks – the coefficient matrices ($R_t^{(i)}$) enable the diagonal blocks to be different, while the basis sharing scheme makes the approximation memory efficient. Therefore, the proposed shared basis block low-rank approximation  is totally different from the Kronecker product approximation. We will clarify this in the paper.\n\n> The reason why the proposed method shows better validation accuracy than Shampoo is also unclear. Since the proposed method introduces further approximation to Shampoo\n\nAs explained above, we are not an approximation to Shampoo. Instead, our method can approximate diagonal blocks much better than Shampoo, leading to improved performance. We will clarify this in the paper.\n\n> What is the effect of removing the smallest eigenvalue from the R matrix? The experiments don't seem to compare with and without this difference.\n\nWe follow the frequent direction paper and Sketchy [11] to remove the smallest eigenvalue when applying the frequent direction method. The regret bound can be obtained with and without doing so. Our experiment on the autoencoder benchmark shows that the performances are similar.\n\n\n\n\n||Ours (rank 32)|Ours (rank 32, without subtract)|\n|---|---|---|\n|Train Loss|51.49|51.48|\n\n"}, {"review_id": "QfobOlYMWB", "reviewer": "Reviewer_Yba5", "summary": "Following up on the line of work stemming from AdaGrad, the paper presents a novel adaptive optimizer that uses a block-diagonal preconditioner. The block-diagonal elements of the second-moment matrix are approximated through a low-rank approximation that shares the same basis, computed via randomized SVD, for all the block-diagonal components. Regret bounds are shown to be optimal, and the performance of the optimizer is evaluated against related previous work on an autoencoder and a transformer benchmark.", "questions": "- Could the authors provide the training curves for the transformer benchmark, allowing for a comparison of the generalisation properties of the different optimizers?\n- Would the proposed approach pay off even on convolutional networks, for instance when compared against [5], or against Adam in terms of performance over runtime?\n- How was rank=32 determined? Would performance over steps improve with larger values of $k$? The appendix only provides experiments for lower ranks.\n\n-------------------\nUpdate after rebuttal:\n\nIn view of the discussion of the authors with the other reviewers and the provided new experimental results, which would appear to support the validity of the proposed approach, I keep my score of 7.", "limitations": "Limitations are clearly detailed in the conclusions.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is on average well-written, well-presented, and technically solid. Relatively complex technical content is explained in a clear way, with adequate coverage of background knowledge. \nThe choices to use a block-diagonal second moment matrix and to use a shared low-rank approximation are adequately justified experimentally, and by presenting technical results in a simplified setup (under fixed parameters). Optimal regret bounds are provided, and section 2.3 appears to cover the differences with related previous work in a detailed manner.\nThe experimental results suggest that the proposed algorithm is more effective than more expensive methods even when plotted over optimization steps.", "weaknesses": "I believe the paper would benefit from a wider experimental section, perhaps including convolutional networks or vision transformers (on vision tasks). In addition, I think it would be interesting to study in more detail the regularizing effect of the various optimizers. Information concerning this could be provided by the training curves from the transformer experiment.\n\nAs a very minor comment: I think the presentation could be improved by splitting section 2 into two sections: \"background\" and \"proposed method\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review!\n\n> I believe the paper would benefit from a wider experimental section, perhaps including convolutional networks or vision transformers (on vision tasks). \n\nWe will expand the experimental section with our future experiments.\n\n> As a very minor comment: I think the presentation could be improved by splitting section 2 into two sections: \"background\" and \"proposed method\".\n\nThanks for the suggestion! We will refine the paper following your suggestion.\n\n> Could the authors provide the training curves for the transformer benchmark, allowing for a comparison of the generalization properties of the different optimizers?\n\nThank you for the suggestion. We have added the training curves for the transformer benchmark.\n\n> Would the proposed approach pay off even on convolutional networks, for instance when compared against [5], or against Adam in terms of performance over runtime?\n\nThe Miniblock Fisher paper [5] can be seen as a special case of ours where a mini block size with full-rank shared basis is used for the convolutional layer, and the time complexity will be the same as ours. In their paper, they are outperforming Adam. We will explore in the future experiment whether different block sizes and ranks can further improve the performance.\n\n\n> How was rank=32 determined? Would performance over steps improve with larger values of $k$? The appendix only provides experiments for lower ranks.\n\nAs shown in Figure 4 of the appendix, k=32 gives a decent approximation to the block diagonal. This is the main reason we use k=32 in our experiments.\n\nFor the autoencoder benchmark, the performance over steps improves with larger values of $k$, as shown in the following table. \n\n\n\n\n||Adam|Ours (rank 4)|Ours (rank 32)|Ours (rank 64)|Ours (rank 128)|\n|---|---|---|---|---|---|\n|Train Loss|54.66|52.88|51.49|51.10|50.84|\n\n\n\n\nFor the transformer benchmark, the performance slightly increases, and smaller k still obtains decent performance. Therefore, we think k=32 is sufficient for both settings since it strikes a good tradeoff between performance and memory efficiency. \n\n\n| Validation Accuracy | Adam  | Block Adam | Ours (rank 4) | Ours (rank 32) | Ours (rank 64) |\n| ------------------- | ----- | ---------- | ------------- | -------------- | -------------- |\n| Mean                | 67.70 | 70.36      | 70.41         | 70.51          | 70.58          |\n| Variance            | 0.16  | 0.21       | 0.03          | 0.03           | 0.04           |\n\n\n"}, {"review_id": "K4qffMUbCl", "reviewer": "Reviewer_9J53", "summary": "Computing the full Gram matrix is practically infeasible for large models even if it's done layer-wise. This work first confirms the block structure of each layer's Gram matrix experimentally. Previous work considered approximating each of these blocks as a low-rank matrix to improve the memory and computation requirements. However, Storing the basis for each block remains problematic. Therefore the authors propose using a shared basis for all the blocks, providing an intuition why such a shared basis exists. They compare their method with previous work for training an Autoencoder on MNIST and for training a Transformer. In both cases, they show they can obtain superior or comparable results while being faster and using less memory.", "questions": "1. The method used for computing the shared basis is interesting. I was wondering if this method or similar ones has been used already (in which case I think it would be nice to cite them) or if this is a novel contribution? If the latter, I was wondering if there is a mathematical argument to support using this method or is it possible that some other method would work better in estimating the basis?\n\n2. The authors mentioned that they do not pre-condition the embedding layer to be fair to other methods even though they could. What happens if the embedding layer is included? Is there a limit (though higher than 4096 for Shampoo) on the size of the embedding for your method as well? In any case I think the results would be interesting to have (in addition to the fair comparison you already performed).", "limitations": "The authors briefly discuss the quadratic scaling of the space complexity with rank. I think the discussion could be a bit more thorough and could also include the cubic dependence of computation complexity on the rank. Also, to make it more interpretable, it would be good to mention how slower the algorithm can be for very large models than for example Adam. For example I think using k = 32, it can be 32x slower but helps save memory especially if d[in] and d[out] are both much larger than 1024.\nThe authors also mention that larger models might require larger k. I think including a brief study on this in this work would make it much more clear whether the proposed method is applicable to large-scale models or not.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is written clearly and it is easy to follow and understand. The idea of using a shared basis is interesting and the authors justify the existence of such basis based on the low-rank nature of data. Furthermore, the experiments verify that the method works for training neural networks to some extent. The algorithm is also faster than Adam on the experiments towards the end.", "weaknesses": "Unfortunately, the evaluation of the method remains limited. The auto-encoding task is only tried with MNIST. Given that the time of the proposed method is within 4-5x of Adam, it should be feasible to also run similar experiment on larger datasets such as CIFAR10.  \n\nAnother limitation of the proposed method is that it does not scale well with rank. The ablations stop at k=32 which is the rank used for the experiments. Therefore, it is not clear whether k=32 is already enough to obtain close-to-best performance or if this limitation poses a big challenge and needs to be overcome for the method to be successfully deployed in larger scales.\n \nThe theoretical arguments are also a bit limited. No convergence rate is provided for non-convex settings. Furthermore, I am not sure if the proof of Theorem 1 covers the stochastic case especially since no bound is assumed on the noise.\n\nAs a minor note, I think for the ablation studies in the appendix, using plots would work better than the tables to allow easier qualitative comparison.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable feedback.\n\n> The evaluation of the method remains limited. The auto-encoding task is only tried with MNIST. it should be feasible to also run CIFAR 10.\n\nIn addition to the autoencoder benchmark, we also conduct experiments on transformers. We will further expand the autoencoder experiments in the future.\n\n> Another limitation of the proposed method is that it does not scale well with rank. The ablations stop at k=32 which is the rank used for the experiments. Therefore, it is not clear whether k=32 is already enough to obtain close-to-best performance\n\nAs shown in Figure 4 of the appendix, k=32 already gives a decent approximation to the block diagonal. This is the main reason we use k=32 in our experiments. Our further experiment shows that we can use an even smaller k to obtain decent performance.\n\nFor the autoencoder benchmark,\n\n\n\n\n||Adam|Ours (rank 4)|Ours (rank 32)|Ours (rank 64)|Ours (rank 128)|\n|---|---|---|---|---|---|\n|Train Loss|54.66|52.88|51.49|51.10|50.84|\n\n\n\nFor the transformer benchmark, we use the best parameters obtained from rank=32 and run for 10 runs.\n\n| Validation Accuracy | Adam  | Block Adam | Ours (rank 4) | Ours (rank 32) | Ours (rank 64) |\n| ------------------- | ----- | ---------- | ------------- | -------------- | -------------- |\n| Mean                | 67.70 | 70.36      | 70.41         | 70.51          | 70.58          |\n| Variance            | 0.16  | 0.21       | 0.03          | 0.03           | 0.04           |\n\n> Furthermore, I am not sure if the proof of Theorem 1 covers the stochastic case especially since no bound is assumed on the noise.\n\nWe believe the proof of Theorem 1 covers the stochastic case, as in online convex optimization, the loss function at each step $f_t$ can be seen as a stochastic sample, which can even be chosen adversarially.\n\n> I think for the ablation studies in the appendix, using plots would work better than the tables to allow easier qualitative comparison.\n\nThank you for the suggestion. We have added the plots to the one page pdf for rebuttal.\n\n> The method used for computing the shared basis is interesting. I was wondering if this method or similar ones has been used already.\n\nOur method can be seen as a special case of high order SVD, which guarantees quasi-optimality (the approximation error is optimal up to a factor of the square root of the dimensions, so $\\sqrt{3}$ in our case). We will clarify this in the writing. The application of using high order SVD to obtain shared bases for approximating the preconditioner is our novelty.\n\n> it would be good to mention how slower the algorithm can be for very large models than for example Adam. For example I think using k = 32, it can be 32x slower but helps save memory.\n\nThe time complexity for each training step can be decomposed into two parts: (1) back-propagation for gradient computation, and (2) computing the update based on gradient. For (1), all the optimizers have identical time cost. For (2), our method is faster than other second order methods (as shown in the experiments) but could be k times worse than Adam. However, in large scale training, (1) often dominates since it is proportional to the batch size. Further, for models such as transformers, the gradient computation cost goes quadratically with the sequence length. Consequently, we think our method will not be much slower than Adam for very large models. We will add more descriptions to the paper.\n"}, {"review_id": "cHhfG4zhM0", "reviewer": "Reviewer_BzFF", "summary": "In this paper, the authors introduce a novel method to perform adaptive gradient updates. Similarly to other works in the literature, this method focuses on capturing curvature information by considering an approximation to the second-moment matrix $H_t = \\sum_t g_t g_t^T$, where $g_t$ is the gradient at the t-th optimisation step, and uses this as a preconditioner to apply to the gradient direction before performing the next optimisation step. Given the size of the target matrix $H_t$, using it in its entirety is unfeasible: much of the literature then focuses on slimming its structure by proposing suitable block-diagonal approximations of $H_t$. This work builds upon these approaches.\n\nThe proposed preconditioner is studied for fully-connected layers, and performs the following approximations:\n- Drop all inter-layer dependencies in $H_t$ (that is, consider a block-diagonal approximation where each block $H_t^l$ refers to the parameters in a single layer $l$ - this is already considered in the literature)\n- Further sub-block-diagonalise each of the $H_t^l$: the blocks are picked so to group together parameters affecting the same output variable\n- Consider a low-rank approximation of these sub-blocks. Crucially, the approximation is made on the same manifold for each sub-block within a given $H_t^l$\nThis approach is justified by the following observations:\n- The norm of each pair of diagonal sub-blocks is larger than that of their off-diagonals counterparts (Figure 1 and lemma 1, justifying sub-block-diagonalisation)\n- Each diagonal sub-block referring to the same layer lies on the row-space of the input matrix of the given layer, and this is generally close to low-rank (justifying sharing of low-rank basis)\n\nThe paper then proceeds to analyse the performance of the novel method, both theoretically (providing memory and convergence bounds) and empirically (comparing its performance against available alternatives, mostly in terms of memory usage and time-to-solution). The tests (conducted on two commonly used benchmarks, considering fully-connected layers in a transformer and an auto-encoder) confirm the validity of the method.", "questions": "Some criticism:\n- While I appreciate Lemma 1 in giving an indication of the relative size of diagonal vs non-diagonal entries, I think it requires some caveats, as I find it’s not really sufficient to justify considering only the diagonal part of the matrix. In my opinion, since you’re dropping the *whole* non-diagonal part, the ideal comparison would be between $||H^(i,i)||$ and the whole $\\sum_{j\\neq i} ||H^(j,j)||$; that is, ideally one should provide some indication of block-diagonal dominancy. I don’t think a bound in this form can be established, so I think it would be fair to point this out. (Mind me: I still consider the simplification valid, especially in light of the results in Fig1)\n- How many different versions of Shampoo(x) have you experimented with? I’m asking because in Fig 3 it’s the one method which shows the highest variability with respect to x - and also the largest improvement (especially wrt to time) as x is decreased. It would be nice to include a Shampoo(125 / 256), to check whether this trend continues\n- Why is only the train loss reported in in Table 2 and only the validation accuracy in table 3?\n\nSome curiosity:\n- Have you tried using nonlinearities other than tanh in your MLP benchmark? This is pure speculation, but I was thinking about the observation that $X$ is low-rank (you mention it in Line 149). Since we’re considering gradients, if I got everything right, this matrix should be actually further multiplied by the derivatives of the nonlinearities: I was wondering how this observation would change if different activation functions are considered (or if it would change at all!)\n- I’m quite impressed that initialising $Y$ using information from the previous iteration works so well! You mention in the main paper (Line 175) that just a few iterations are necessary to convergence of the randomised SVD (and indeed you only use one iteration in your experiments), but how quickly does convergence occur? In the ablation study in the appendix (Table 5) you only report the final time/loss results, unfortunately. I ask this, because one might consider simply reusing the bases from the previous iteration, if they’re already good enough, and only update them once every few optimisation steps (as done in Shampoo? - Line 256), to further speed things up: have you investigated this?\n\nSome minor corrections:\n- I found the notation in Sec 2.1 (right before Lemma1) a bit confusing: it’s tricky to keep track of what is a vector, what is a matrix, and what is a vectorisation of a matrix. For example, $D_W$ should be a matrix (since you treat it as such in the proof in the appendix), but here you define it as a vector? Just reporting dimensions explicitly, and ensuring consistency with the notation in the appendix should help with readability\n- Line 34: $g_t$ is a vector in $\\mathbb{R}^{d_{in}d_{out}}$, not a matrix in $\\mathbb{R}^{d_{in}\\times d_{out}}$?\n- In the equation right after Line 162:  I think the $+$ is a typo for $\\approx$?\n- In Line 179, you’re missing a transpose $M_t M_t^T$?\n- In Line 406 in the appendix, you’re missing a transpose $X^TDz$?\n- Appendix 6.8: is $\\Omega$ actually $Y$?\n", "limitations": "A limitation of the method (in terms of memory consumption) is made explicit in the dedicated section in the paper. Still, my impression is that this could be extended (particularly, with comments pertaining the first “weakness”)\n", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The approach is supported and justified by a (reasonably solid) theoretical analysis, which greatly contributes to grounding the intuition behind it\n- The method seems effective, and the experiments conducted are convincing enough in validating its performance - at least compared to other methods which try to approximate $H_t$ in a similar fashion\n- The paper is well-structured, and reads cleanly (especially after some minor proofreading - see comments below)", "weaknesses": "- The approach might be rather limited in its applicability, as it mostly focuses on fully-connected layers: it is unclear to me if/how easily this analysis could be extended to other types of layers (say convolutional, for example). I invite the authors to further comment on this\n- The work might be considered rather incremental: it proposes a possible improvement on an already-existing method. Still, the improvement is noticeable and relevant, at least for the test cases considered\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review!\n\n> The approach might be rather limited in its applicability, as it mostly focuses on fully-connected layers.\n\nWe have experiments on the transformer, which also contains self-attention layers. In general, our method can be used to handle any tensor product operations in neural networks, which is quite universal. For convolutional layers, since each kernel (block) could be quite small, our method will become equivalent to the block-diagonal approximation (e.g., BlockAdam). The contribution of this paper is to provide a way to scale block-diagonal approximation to larger matrices using the proposed shared basis approach. \n\n> While I appreciate Lemma 1 in giving an indication of the relative size of diagonal vs non-diagonal entries, I think it requires some caveats.\n\n> ideally one should provide some indication of block-diagonal dominancy. I don’t think a bound in this form can be established, so I think it would be fair to point this out.\n\nThank you. You are correct that we are not able to establish a bound for block-diagonal dominancy. We will mention this in the paper.\n\n> How many different versions of Shampoo(x) have you experimented with? It would be nice to include a Shampoo(125 / 256), to check whether this trend continues.\n\nIn the experiment, we observed that further decreasing the block size of Shampoo does not give any speed up, so we did not further decrease the block size, as shown in the following table. The potential reason for this could be the amount of parallelism available and the implementation overhead.\n\n\n\n\n|            | Shampoo (256) | Shampoo (512) | Shampoo (1024) |\n| ---------- | ------------- | ------------- | -------------- |\n| Time (min) | 348           | 325           | 564            |\n\n\n\n\n> Why is only the train loss reported in Table 2 and only the validation accuracy in table 3?\n\nWe follow the previous works such as K-fac [23] and Miniblock Fisher [5] to evaluate the train loss on the autoencoder benchmark. We have added the training loss curve for the transformer benchmark to the one page pdf for rebuttal, and our method still outperforms other optimizers in terms of training loss.\n\n> Have you tried using nonlinearities other than tanh in your MLP benchmark?\n\nYes, the transformer benchmark uses the elu non-linearity.\n\n> You mention in the main paper (Line 175) that just a few iterations are necessary to convergence, but how quickly does convergence occur? I ask this, because one might consider simply reusing the bases from the previous iteration, if they’re already good enough, and only update them once every few optimisation steps (as done in Shampoo?)\n\nCurrently we observe that the inner iteration (when setting to 1) is not the major bottleneck, so further reducing this cost may not be useful for our current experiments. However, this is a great suggestion and we will definitely try this idea later. \n\n> Some minor corrections\n\nThank you. We have fixed these in the paper."}]

Author rebuttal:
We thank all the reviewers for the valuable feedback. We have added figures for the ablation study and the training curves for the transformer benchmark.

Meta-review: The submission presented a block-diagonal version of AdaGrad, using a low-rank approximation for each block and a shared basis across blocks; this massively reduces memory complexity, and makes time cost manageable. The authors show good results across standard benchmarks.

The authors and reviewers were active during the discussion, and the authors provided significant additional experimental data which led the reviewers and AC to converge to acceptance. 

I would encourage the authors to incorporate the ample feedback received during the discussion period in the next version of their work.
Decision: Accept (poster)