SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Our overall algorithm is summarized in Algorithm  $\mathbb{I}$ . The time complexity of the algorithm is dominated by the matrix products in power iterations,  $\mathcal{O}(mk^3 + dk + bk^2)$ . The memory complexity is  $\mathcal{O}(bk + mk^2)$  for storing the shared basis and the  $k \times k$  dense coefficient matrices of each block.
For simplicity, the above discussions are mainly for two-dimensional matrices. For handling higher-order 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.
Algorithm 1 Shared-Basis Low Rank Block-Diagonal Adagrad
Initialize: coefficients  $\pmb{R}_0^{(i)} = \mathbf{0}$  , shared bases  $B_0 = 0$    
for  $t = 1\dots T$  do Compute stochastic gradient  $\pmb{g}_{t}$    
Obtain shared basis  $\pmb{B}_{t}$  by running randomized SVD (Eq. 5) for  $\pmb{A}_{t}$    
Compute  $\tilde{\pmb{R}}_{t}^{(i)} = (\pmb {B}_{t})^{T}\pmb{M}_{t}^{(i)}\pmb{B}_{t}$    
Compute  $\pmb{R}_{t}^{(i)} = \tilde{\pmb{R}}_{t}^{(i)} - \sigma_{\mathrm{min}}(\tilde{\pmb{R}}_{t}^{(i)})\pmb{I}$  for  $i = 1,\dots ,m$    
Compute  $\rho_t^{(i)} = \rho_{t - 1}^{(i)} + \| M_t^{(i)} - B_tR_t^{(i)}(B_t)^T\|$  for  $i = 1,\ldots ,m$    
Compute the preconditioned update  $\pmb{u}_{t}^{(i)} = -(B_{t}R_{t}^{(i)}(B_{t})^{T} + (\rho_{t}^{(i)} + \epsilon)\pmb {I})^{-1 / 2}\pmb{g}_{t}^{(i)}$  for  $i = 1,\dots ,m$  Update model parameters  $\theta_{t} = \theta_{t - 1} + \alpha_{t}\pmb{u}_{t}$    
end for
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  $\theta_t \in \mathcal{K}$  is chosen iteratively, where  $\mathcal{K}$  is a convex decision set. After the decision of  $\theta_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
$$
\operatorname {R e g r e t} _ {T} = \sum_ {t = 1} ^ {T} f _ {t} \left(\theta_ {t}\right) - \min  _ {\theta \in \mathcal {K}} \sum_ {t = 1} ^ {T} f _ {t} (\theta).
$$
We give the following theorem for our proposed method. All the proofs can be found in the Appendix.
Theorem 1 Algorithm gives the following regret
$$
R e g r e t _ {T} \leq \sqrt {2} D \sum_ {i = 1} ^ {m} \mathrm {T r} ((\boldsymbol {H} _ {T} ^ {(i)}) ^ {1 / 2}) + \sqrt {2} b \sum_ {i = 1} ^ {m} (\rho_ {T} ^ {(i)}) ^ {1 / 2},
$$
where  $D$  is the diameter of the constraint set  $\mathcal{K}$ , and  $\pmb{H}_T^{(i)} = \sum_{t=1}^T \pmb{g}_t^{(i)} (\pmb{g}_t^{(i)})^T$ .
To bound the second term in Theorem we further derive the following theorem.
Theorem 2 Algorithm guarantees
$$
\sum_ {i = 1} ^ {m} (\rho_ {T} ^ {(i)}) ^ {1 / 2} \leq (3 m ^ {3}) ^ {1 / 4} (\sum_ {i = 1} ^ {m} \mathrm {T r} (\pmb {H} _ {T} ^ {(i)})) ^ {1 / 2},
$$
and thus  $Regret_{T} = \mathcal{O}(\sqrt{T})$ .
This shows that our proposed method is optimal asymptotically in terms of  $T$ . With some further assumptions, we can further derive bounds of  $\rho_T^{(i)}$  in terms of the lower eigenvalues of  $\pmb{H}_T^{(i)}$ . 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_{\mathrm{in}} \times d_{\mathrm{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_{\mathrm{in}}^2 \times d_{\mathrm{out}}^2$ . To address this issue, previous studies such as [31, 22] proposed low-rank approximation for the preconditioner, which reduces the memory cost to  $\mathcal{O}(d_{\mathrm{in}} d_{\mathrm{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  $\mathcal{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  $\mathcal{O}(d_{\mathrm{in}})$  and a memory complexity of  $\mathcal{O}(dd_{\mathrm{in}})$ . Additionally, they propose maintaining the average of the block diagonal for larger layers, resulting in an approximation of  $I \otimes 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 \otimes 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].
Table 1: Time and space complexity comparison for a single parameter matrix of size  $d_{\mathrm{in}} \times d_{\mathrm{out}}$ .
Algorithm	Time Complexity	Space Complexity
Full Matrix Adagrad [10]	d3in d3out	d2in d2out
SON-FD [22]	d1in d1out k2	d1in d1out k
Ada-FD [31]	d1in d1out k2	d1in d1out k
Shampoo [16]	d3in + d3out	d2in + d2out
Block Shampoo [4]	d1in d1out b	d1in d1out
Sketchy [11]	d1in d1out k	d1in k + d1out k
Block Diagonal Adam [32]	d1in d1out b2	d1in d1out b
Ours	d1in d1out k + d1in k2 + d1out k3	d1in k + d1out k2
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 [1]: 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_{\mathrm{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 \times$  lower time, and  $1.5 \times$  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
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	3641	2799	1777
(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 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 \times$  lower time with less memory to obtain  $\sim 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.
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.
Table 3: Experimental results on 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
(b) Train loss (time)
(a) Train loss (step)
(c) Validation accuracy (step)
Figure 4: Train loss and the best validation accuracy on the Transformer benchmark.
(d) Validation accuracy (time)
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].
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.
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.
We thank the anonymous reviewers for their helpful comments. JNY and CJH are supported in part by NSF under 2008173, 2048280, 2325121 and 2331966, Cisco and Sony.
[19] 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.
[20] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[21] Dong C Liu and Jorge Nocedal. On the limited memory BFGS method for large scale optimization. Mathematical programming, 45(1-3):503-528, 1989.
[22] 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.
[23] James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning, pages 2408-2417. PMLR, 2015.
[24] Joakim Nivre, Marie-Catherine de Marneffe, Filip Ginter, Jan Hajic, Christopher D Manning, Sampa Pyysalo, Sebastian Schuster, Francis Tyers, and Daniel Zeman. Universal dependencies v2: An evergrowing multilingual treebank collection. arXiv preprint arXiv:2004.10643, 2020.
[25] Farbod Roosta-Khorasani and Michael W Mahoney. Sub-sampled Newton methods. Mathematical Programming, 174:293-326, 2019.
[26] Nicolas Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algorithm. Advances in Neural Information Processing Systems, 20, 2007.
[27] Jürgen Schmidhuber. Deep learning in neural networks: An overview. Neural networks, 61: 85-117, 2015.
[28] Shai Shalev-Shwartz et 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 6.5—RMSProp: 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 N 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 on Pattern Analysis and Machine Intelligence, 44(10):6910-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.