Title: Block Low-Rank Preconditioner with Shared Basis for Stochastic Optimization
Venue: NeurIPS 2023
Accepted: yes
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

Abstract

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 blockdiagonal 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 firstorder 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.

1 Introduction

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 / √∑ t ( g t ) 2 i , for a gradient vector 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 [1], which also estimate off-diagonal gradient moments such as ∑ t ( g t ) i ( g t ) j for i = 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.

/negationslash

The second-moment matrix ∑ 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 → XW , where X ∈ R n × d in and W ∈ R d in × d out and make two important observations: a) the

second-moment matrix ∑ t g t g T t , where g t ∈ R d in d out is a vectorized form of the gradient with respect to W at iteration t , has a block-matrix structure with d out blocks of size d in × d 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 O ( d out d 2 in ) 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 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 B ∈ R d in × k , k /lessmuch d in, and a different residual R ( i ) ∈ R k × k for each block such that { BR (1) B /latticetop , . . . , BR ( d out ) B /latticetop } 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 /latticetop , . . . , BR ( d out ) B /latticetop } to the original block-diagonal approximation of the gradient second-moment matrix. We also use this to prove a O ( √ 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 × lower time, and 1 . 5 × lower memory. In addition to this, in the Transformer benchmark, we use 3 . 5 × 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.

Notation. We use boldfaced lower-case letters such as x to denote a vector and boldfaced uppercase letters such as H to denote a matrix. H : ,i and H i, : denote the i -th column and row of the matrix respectively. Diag ( H 1 , . . . , H k ) denotes a block-diagonal matrix with diagonal blocks H 1 , . . . , H k . vec ( H ) flattens a matrix into a vector, while A ⊗ B denotes the Kronecker product of two matrices. ‖ H ‖ 2 , ‖ H ‖ F denote the 2-norm and Frobenious norm of a matrix respectively.

2 Proposed Method

We consider training a deep neural network that takes an input vector x and produces an output f ( θ , x ) . With a training set ( y i , x i ) n i =1 consisting of n samples and a loss function /lscript , the network parameters θ are learned by minimizing the empirical loss L over the training set:

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 g t = ∇ L B ( θ t ) where L B represents the loss of a sampled mini-batch, the preconditioner H t and the next iterate θ t +1 are computed as follows:

where α t is the step size and /epsilon1 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 elementwise version of Adagrad [10] and many of its extensions such as Adam [20].

However, previous studies [16, 23, 1] 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

in

t

(b)

The

feed-forward

layer with ordering

H

out

t

(c)

Self-attention

output

kernel with ordering

H

in

t

(d)

Self-attention output

kernel with ordering

H

out

t

0

4.0

128

- 3.5

128

128

6

256

- 3.0

256 -

256

5

256

384

- 2.5

384 -

384

4

384

512

640

- 2.0

- 1.5

640 -

3

512

640

3

512

640 -

768

- 1.0

768

768

768

896

0.5

968

968

968

1024

128

256

0#9

1024

1024

128

992

968

1024

1024

离

968

1024

1024

128

992

968

1024

[16, 23, 1], we precondition each layer separately. This effectively imposes a block diagonal structure assumption on the preconditioner:

where L is the number of layers. From now on, we will approximate each H l t separately. Therefore, for notational simplicity, we will omit the layer index and refer to the preconditioner matrix for a single layer as H t .

Consider a fully connected layer with weights W ∈ R d in × d out , where d in and d 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 in d out × d in d 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 out blocks, each with a size of d 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 in t follows row-major ordering, grouping weights associated with the same input neuron, while H out t follows column-major ordering, grouping weights associated with the same output neuron. Notably, we observe that H out t 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 X ∈ R n × d in as the input of this layer for all the n samples, and the fully connected layer can be written as

where W ∈ R d in × d out is the weight matrix and Z ∈ R n × d out is the output matrix. Let D W = ∂ L ∂ vec ( W ) and D Z = ∂ L ∂ Z denote the gradients of the loss function with respect to W and Z , respectively. Additionally, let us define the preconditioner as:

where m = d out is the number of blocks. Since H = D W D T W , we can derive the following lemma:

Lemma 1 Assuming the model is evaluated with a fixed parameter θ , then any two diagonal blocks H ( i,i ) , H ( j,j ) and two corresponding off-diagonal blocks H ( i,j ) , H ( j,i ) for the preconditioner matrix of a fully connected layer will satisfy the following inequality:

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 out groups and each group has size d in, storing this block-diagonal approximation requires O ( d 2 in d out ) memory. To illustrate, a single fully connected layer in a Transformer with d in = 4096 and d 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 H t ≈ Diag ( H (1) t , . . . , H ( m ) t ) 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 out and b = d 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:

where k is the rank, B ( i ) t ∈ R b × k and R ( i ) t ∈ R k × k are the eigenvectors and eigenvalues of each block respectively. Using this approach, the memory cost is reduced to 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:

where B t ∈ R b × k is the orthogonal shared basis and R ( i ) t ∈ R k × k denotes the coefficients for each block. Note that here we assume each R ( i ) t 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 in k + d out k 2 ) which can approach the cost of first-order methods when k 2 ≤ d in and k ≤ d 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 θ . In such cases, each diagonal block can be expressed as:

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 Σ V T using SVD, where U ∈ R n × k , Σ ∈ R k × k , V ∈ R d in × k , we have

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.

Singular Values

1.0

0.8

0.6

AccumulatedRatio

Blockl

0.4

Block2

Block3

Block4

0.2

Block5

0

20

40

60

80

100

Howto compute the shared basis in the online setting? Given the previous B t -1 , { R ( i ) t -1 } i m =1 and the current gradient { g ( i ) t } i m =1 , we will describe the procedure for computing the updated approximate matrices B t , { R ( i ) t } i m =1 . Let

be the block to be approximated. To derive the share bases, we concatenate all the blocks

and compute the SVD to derive its left singular vectors. This is equivalent to computing the topk eigenvectors of

To efficiently compute the SVD of A t , we adopt the randomized SVD method [17]. This approach enables us to compute the basis without explicitly forming A t or any of the M ( i ) t . We initiate the process with an initial matrix Y ∈ R b × ( k + p ) , where p is the over-sampling constant (we set p = 0 in experiments). Iteratively, we perform the following update:

where QR ( X ) denotes the QR decomposition of matrix X , returning the Q matrix. Using this approach, Y will converge to the dominant eigenspace of A t . Additionally, since we can 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 t Y in (5). The naive way to calculate A t Y = ∑ m i =1 M ( i ) t ( M ( i ) t Y ) takes time complexity of O ( mbk 2 ) = O ( dk 2 ) . However, we can reduce some repeated computations as the blocks share the same basis. Let

whose time complexity is O ( mbk ) = O ( dk ) . We have

Therefore, the time complexity for each power iteration is O ( m ( k 3 + bk )+ bk 2 ) = O ( mk 3 + dk + bk 2 ) (the QR decomposition takes 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 -byk coefficient matrices by

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:

To account for the missing mass in the low-rank approximation, we track a scalar value for each block using the recurrence relation:

which can be efficiently computed using power iteration. Finally, to derive the update direction u ( i ) t , we compute

Detailed derivations can be found in Appendix 6.5. Here we follow Shampoo [16] to add the /epsilon1 I before taking the matrix square root, which is slightly different from the Adagrad formulation shown in (1).

Our overall algorithm is summarized in Algorithm 1. The time complexity of the algorithm is dominated by the matrix products in power iterations, O ( mk 3 + dk + bk 2 ) . The memory complexity is O ( bk + mk 2 ) for storing the shared basis and the k × k dense coefficient matrices of each block.

For simplicity, the above discussions are mainly for two-dimensional matrices. For handling higherorder tensor parameters, such as the multi-head self-attention layer in Transformers, we adopt a grouping strategy based on one of the axes. Further discussions can be found in Appendix 6.1.

2.3 Regret Bound Analysis

Following previous works [16, 11], we provide a convergence analysis of the proposed algorithm in the online convex optimization framework [18, 28]. In the online convex optimization setting, a parameter θ t ∈ K is chosen iteratively, where K is a convex decision set. After the decision of θ t , a convex loss function f t is revealed, which can be selected adversarially. The regret suffered by the algorithm up to step T is defined as

We give the following theorem for our proposed method. All the proofs can be found in the Appendix.

Theorem 1 Algorithm 1 gives the following regret

To bound the second term in Theorem 1, we further derive the following theorem.

where D is the diameter of the constraint set K , and H ( i ) T = ∑ T t =1 g ( i ) t ( g ( i ) t ) T .

Theorem 2 Algorithm 1 guarantees

T O

This shows that our proposed method is optimal asymptotically in terms of T . With some further assumptions, we can further derive bounds of ρ ( i ) T in terms of the lower eigenvalues of H ( i ) T . 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 a neural network layer of size d in × d out. The results are summarized in Table 1.

As mentioned earlier in Section 2, the full matrix version of Adagrad [10] poses a significant challenge in practice due to the storage and matrix inversion of a full preconditioner with dimensions d 2 in × d 2 out . To address this issue, previous studies such as [31, 22] proposed low-rank approximation for the preconditioner, which reduces the memory cost to O ( d in d out k ) , with k denoting the rank. However, even with this reduced cost, it remains impractical for modern neural network architectures.

Inspired from the observed empirical block-diagonal structure [7, 26], [32] propose a block diagonal 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 O ( db ) . As a result, they are limited to using small block sizes in practice. Similarly, [5] proposes a mini-block Fisher method. For convolutional layers, they group parameters based on filters, resulting in a block size of r 2 when the kernel size is r . In fully-connected layers, parameter grouping is based on output neurons, yielding a block size of O ( d in ) and a memory complexity of O ( dd in ) . Additionally, they propose maintaining the average of the block diagonal for larger layers, resulting in an approximation of I ⊗ R , where R represents the average of all diagonal blocks. This approach resembles one side of the Kronecker product methods in the next paragraph, but their research demonstrates some performance loss associated with the approximation.

Another popular approach, such as Shampoo [16] and K-fac [23], involves grouping parameters based on both the input and output neurons, resulting in a Kronecker product approximation L ⊗ R for the preconditioner. In order to enhance the time and space complexity of Shampoo, [4] introduces a technique called blocking, which partitions the parameters into smaller blocks. On the other hand, [11] proposes Sketchy as a low-rank variant of Shampoo.

Based on the above discussion, except for Sketchy, all other methods exhibit higher memory complexity compared to our method. Under the same k , our memory complexity is larger than Sketchy. This is due to the fact that we 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 a standard Autoencoder benchmark [27] on the MNIST dataset [8] and a larger Transformer model [30] on the Universal Dependencies dataset [24].

We adopt k = 32 as the default rank for our methods. For randomized SVD, we set the oversampling parameter to 0 and the number of iterations to 1 . The impact of these parameters is detailed in our ablation study in Appendix 6.10. Similar to Shampoo, we use the grafting technique [2] in our method. We set the grafting type to RMSPROP_NORMALIZED . We compare the proposed optimizer with the widely used Adam optimizer [20] along with the following second-order optimizers:

Shampoo [16]: we adopt the implementation of [4] and keep most of the default settings with the following differences. Following the setting of [11], for Shampoo, the preconditioning starts at step 101 for stability and the preconditioners are updated every 10 steps for speed ( preconditioning_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( x ).

Sketchy [11]: this is a 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 block-diagonal approximation with a smaller block proposed in [32]. We group the blocks based on output neurons, so we chose the block size to be a divisor of d in. Specifically, a block size of 10 is used for the autoencoder benchmark and a block size of 8 is used for the Transformer benchmark.

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 a log-uniform distribution. For the autoencoder benchmark, we conduct 180 trials of random search on one NVIDIA RTX 2080Ti GPU with 11GB 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 we follow the same setting in previous papers. Specifically, we conduct 100 epochs of training using the MNIST dataset. The layer sizes of the autoencoder are [1000 , 500 , 250 , 30 , 250 , 500 , 1000] , with a 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 to 0.

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 . 2 × lower time, and 1 . 5 × lower memory. The training loss curve can be seen in Figure 3.

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 512 and the MLP dimension size is

(a) Train loss (epoch)

(b) Train loss (time)

70.0

Ours

67.5

Adam

65.0 -

EI

Shampoo(250)

5

S

62.5

Shampoo(1000)

Sketchy

Train

60.0

57.5

55.0

52.5

50.0

0

20

40

60

80

100

Epoch

70.0

Ours

67.5

Adam

65.0

Shampoo(250)

5

S

62.5

Sketchy

Shampoo(1000)

rain

60.0

57.5

55.0

52.5

50.0

0

50

100

150

200

250

300

Time (s)

2048. For the self-attention layers, the number of heads is 8, and the query, key, value dimensions are 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 to 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 use 3 . 5 × lower time with less memory to obtain ∼ 1 . 8% improvement in the validation error. The training loss curve and the validation accuracy curve can be seen in Figure 4, where the proposed method outperforms others in terms of both raw training time and number of training steps.

4 Related Work

Second order methods. For the optimization of ill-conditioned functions, second order methods are often more effective than first order methods. This includes Newton's method and the quasi-Newton methods [9, 21]. Subsampled versions of these methods [15, 25] have also been proposed for solving 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 [23, 16] 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, as discussed in Section 2.4.

(a) Train loss (step)

(b) Train loss (time)

(c) Validation accuracy (step)

(d) Validation accuracy (time)

Ours

Adam

6

Block Adam

5

Shampoo(512)

Train Loss

Shampoo(1024)

4

Sketchy

.......

3

2

1

0

0

2500

5000

7500

10000

12500

15000

17500

Step

0.7

0.6

0.5

Best Validation

0.4

Ours

0.3

Adam

Block Adam

0.2

Shampoo(512)

0.1

Shampoo(1024)

Sketchy

0.0

0

2500

5000

7500

10000

12500

15000

17500

Step

7

Ours

Adam

6

5

Shampoo(512)

SS07

Shampoo(1024)

Train

4

Sketchy

3

1

0

25

50

75

100

125

150

175

200

Time (min)

0.7

0.6

0.5

BestValidation

0.4

Ours

0.3

Adam

0.2

Shampoo(512)

B

0.1

Shampoo(1024)

........

Sketchy

0.0

0

25

50

75

100

125

150

175

200

Time (min)

Frequent direction method. Given a stream of vectors, the frequent direction method [13] maintains 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 in previous work for approximating the preconditioner matrix in stochastic optimization [11, 22, 31, 12].

5 Conclusions

We have introduced a novel adaptive optimizer that utilizes a shared-basis block low-rank approximation to approximate the full matrix preconditioner. Through evaluations on an autoencoder and a 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 k = 32 is adequate for the Autoencoder and Transformer benchmarks examined in our study, this may still pose a burden when training large-scale models. Exploring methods to reduce the dependence on k while preserving performance will be an interesting future direction.

References

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.

Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803 , 2020.

Ehsan Amid, Rohan Anil, and Manfred Warmuth. Locoprop: Enhancing backprop via local loss optimization. In International Conference on Artificial Intelligence and Statistics , pages 9626-9642. PMLR, 2022.

Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018 , 2020.

Achraf Bahamou, Donald Goldfarb, and Yi Ren. A mini-block Fisher method for deep neural networks. In International Conference on Artificial Intelligence and Statistics , pages 9191-9220. PMLR, 2023.

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.

Ronan Collobert. Large scale machine learning. Technical report, Université de Paris VI, 2004.

Li Deng. The MNIST database of handwritten digit images for machine learning research [best of the web]. IEEE signal processing magazine , 29(6):141-142, 2012.

John E Dennis, Jr and Jorge J Moré. Quasi-Newton methods, motivation and theory. SIAM review , 19(1):46-89, 1977.

John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memoryefficient adaptive regularization with frequent directions. arXiv preprint arXiv:2302.03764 , 2023.

Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memoryefficient adaptive regularization with frequent directions. arXiv preprint arXiv:2302.03764 , 2023.

Mina Ghashami, Edo Liberty, Jeff M Phillips, and David P Woodruff. Frequent directions: Simple and deterministic matrix sketching. SIAM Journal on Computing , 45(5):1762-1792, 2016.

Mina Ghashami, Edo Liberty, Jeff M Phillips, and David P Woodruff. Frequent directions: Simple and deterministic matrix sketching. SIAM Journal on Computing , 45(5):1762-1792, 2016.

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.

Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning , pages 1842-1850. PMLR, 2018.

Nathan Halko, Per-Gunnar Martinsson, and Joel A Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review , 53 (2):217-288, 2011.

Elad Hazan et al. Introduction to online convex optimization. Foundations and Trends® in Optimization , 2(3-4):157-325, 2016.

Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL http://github.com/google/flax .

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Dong C Liu and Jorge Nocedal. On the limited memory BFGS method for large scale optimization. Mathematical programming , 45(1-3):503-528, 1989.

Haipeng Luo, Alekh Agarwal, Nicolo Cesa-Bianchi, and John Langford. Efficient second order online learning by sketching. Advances in Neural Information Processing Systems , 29, 2016.

James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning , pages 2408-2417. PMLR, 2015.

Joakim Nivre, Marie-Catherine de Marneffe, Filip Ginter, Jan Hajiˇ c, Christopher D Manning, Sampo Pyysalo, Sebastian Schuster, Francis Tyers, and Daniel Zeman. Universal dependencies v2: An evergrowing multilingual treebank collection. arXiv preprint arXiv:2004.10643 , 2020.

Farbod Roosta-Khorasani and Michael W Mahoney. Sub-sampled Newton methods. Mathematical Programming , 174:293-326, 2019.

Nicolas Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algorithm. Advances in Neural Information Processing Systems , 20, 2007.

Jürgen Schmidhuber. Deep learning in neural networks: An overview. Neural networks , 61: 85-117, 2015.

Shai Shalev-Shwartz et al. Online learning and online convex optimization. Foundations and Trends® in Machine Learning , 4(2):107-194, 2012.

T. Tieleman and G. Hinton. Lecture 6.5-RMSProp: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning. 2012.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Yuanyu Wan and Lijun Zhang. Efficient adaptive online learning via frequent directions. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44(10):6910-6923, 2021.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n- 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)\n\nWEAKNESSES:\n- 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\nQUESTIONS:\nSome 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\nLIMITATIONS:\nA 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”)", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nComputing 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nUnfortunately, 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.\n\nQUESTIONS:\n1. 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).\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFollowing 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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\".\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nLimitations are clearly detailed in the conclusions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-22"}

comments: 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.

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-24

comments: 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.

----------------------------
I 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.

STRENGTHS:
1. The research problem studied in this paper is practical and significant.
2. Theoretical analysis is rich.
3. Empirical studies on various datasets.

WEAKNESSES:
1. The document is not very easy to understand.
2. the proposed methodology is not clearly superior to the baseline methodology.
3. The authors did not compare the use of multiple modeling architectures

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)

1. The research problem studied in this paper is practical and significant.
2. Theoretical analysis is rich.
3. Empirical studies on various datasets.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- 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).\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}