SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 6: A rough estimate of memory requirement comparisons of different optimizers tested across benchmarks.
Benchmark	# model parameters	K-FAC	Shampoo	FishLeg	Eva	Adam	SGD+Momentum	RMSprop	tds-SONew
Autoencoder	n=1.4M	5.56n	6.56n	4.28n	n	2n	n	n	3n
GraphNetwork	n=3.5M	8.6n	10.6n	4.8n	n	2n	n	n	3n
Vision Transformer	n=22M	6.4n	7.2n	3.7n	n	2n	n	n	n
Language Model	n=1.3B	5.6n	6.6n	3.3n	n	2n	n	n	3n
A.4.2 Memory Requirements
We present a list of approximate memory requirements of different optimizers across different benchmarks in Table 6. Note that for K-FAC and Shampoo, because preconditioner is updated once only a few steps, they require storing the latest computed preconditioners as well along with the statistics, causing even higher memory overhead.
A.4.3 Hyperparameter search space
We provide the hyperparameter search space for experiments presented in Section 5. We search over  $2k$  hyperparameters for each Autoencoder experiment using a Bayesian Optimization package. The search ranges are: first order momentum term  $\beta_{1} \in [1e - 1, 0.999]$ , second order momentum term  $\beta_{2} \in [1e - 1, 0.999]$ , learning rate  $\in [1e - 7, 1e - 1]$ ,  $\epsilon \in [1e - 10, 1e - 1]$ . We give the optimal hyperparameter value for each experiment in Table 12. For VIT and GraphNetwork benchmark, we search  $\beta_{1}, \beta_{2} \in [0.1, 0.999]$ ,  $lr \in [1e - 5, 1e - 1]$ ,  $\epsilon \in [1e - 9, 1e - 4]$ , weight decay  $\in [1e - 5, 1.0]$ , learning rate warmup  $\in [2\%, 5\%, 10\%] * \text{total_train_steps}$ , dropout  $\in [00, 0.1]$ , label smoothing over  $\{0.0, 0.1, 0.2\}$ . We use cosine learning rate schedule. Batch size was kept = 1024, and 512 for Vision Transformer, and GraphNetwork respectively. We sweep over 200 hyperparameters in the search space for all the optimizers.
For rfdSON [37], there's no  $\epsilon$  hyperparameter. In addition to the remaining hyperparameters, we tune  $\alpha \in \{1e - 5,1.0\}$  (plays similar role as  $\epsilon$ ) and  $\mu_t \in [1e - 5,0.1]$ .
For LLM [47] benchmark, we only tune the learning rate  $\in [1e - 2, 1e - 3, 1e - 4]$  while keeping the rest of the hyperparams as constant. This is due to the high cost of running experiments hence we only tune the most important hyperparameter. For Adafactor [45], we use factored=False, decay method=adam,  $\beta_{1} = 0.9$ , weight decay  $= 1e - 3$ , decay factor  $= 0.99$ , and gradient clipping  $= 1.0$ .
A.4.4 Additional Experiments
VIT and GraphNetwork Benchmarks: In Figure 5 we plot the training loss curves of runs corresponding to the best validation runs in Figure 1. Furthermore, from an optimization point of view, we plot the best train loss runs in Figure 6 got by searching over 200 hyperparameters. We find that tridiag-SONew is  $9\%$  and  $80\%$  relatively better in ViT and GraphNetwork benchmark respectively (Figure 6), compared to Adam (the next best memory efficient baseline).
Autoencoder float32 and bfloat16 experiments: We provide curves of all the baselines and SONew in Figure 4(a) and the corresponding numbers in Table 7 for float32 experiments.
To test numerical stability of SONew and compare it with other algorithms in low precision regime, we also conduct bfloat16 experiments on the Autoencoder benchmark (Table 8). We notice that SONew undergoes the least degradation. Tridiagonal-sparsity SONew CE loss increases by only 0.21 absolute difference (from 51.72 in float32 (7) to 51.93), whereas Shampoo and Adam incur 0.70 loss increase. It's worthwhile to note that SONew performs better than all first order methods while taking similar time and linear memory, whereas while Shampoo performs marginally better, it is  $22 \times$  slower
Table 7: float32 experiments on Autoencoder benchmark. We observe that diag-SONew performs the best among all first order methods while taking similar time. tridiag and band-4 perform significantly better than first order methods while requiring similar linear space and time. Shampoo performs best but takes  $\mathcal{O}(d_1^3 + d_2^3)$  time for computing preconditioner of a linear layer of size  $d_1 \times d_2$ , whereas our methods take  $\mathcal{O}(d_1 d_2)$  time, as mentioned in Section 3.3. rfdSON takes similar space as SONew but performs considerably worse.
Optimizer	First Order Methods
SGD	Nesterov	Adagrad	Momentum	RMSProp	Adam	diag-SONew
Train CE loss	67.654	59.087	54.393	58.651	53.330	53.591	53.025
Time(s)	62	102	62	67	62	62	63
Optimizer	Second Order Methods
Shampoo(20)	rfdSON(1)	rfdSON(4)	tridiag-SONew	band-4-SONew
Train CE loss	50.702	53.56	52.97	51.723	51.357
Time(s)	371	85	300	70	260
Table 8: bffloat16 experiments on Autoencoder benchmark to test the numerical stability of SONew and robustness of Algorithm 3. We notice that diag-SONew degrades only marginally (0.26 absolute difference) compared to float32 performance. tridiag-SONew and band-4-SONew holds similar observations as well. Shampoo performs the best but has a considerable drop (0.70) in performance compared to float32 due to using matrix inverse, and is slower due to its cubic time complexity for computing preconditioners. Shampoo implementation uses 16-bit quantization to make it work in 16-bit setting, leading to further slowdown. Hence the running time in bffloat16 is even higher than in float32.
Optimizer	First Order Methods
SGD	Nesterov	Adagrad	Momentum	RMSProp	Adam	diag-SONew
Train CE loss	80.454	72.975	68.854	70.053	53.743	54.328	53.29
Train time(s)	36	43	37	36	37	38	44
Optimizer	Second Order Methods
Shampoo(20)	rfdSON(1)	rfdSON(4)	tridiag-SONew	band-4-SONew
Train CE loss	51.401	57.42	55.53	51.937	51.84
Train time(s)	1245	80	284	55	230
(a) float32 - autoencoder
(b) bfloat16 - autoencoder
Figure 4: Training curves of all the baselines for Autoencoder benchmark (a) float32 training (b) bfloat16 training
than tridiagonal-SONew. The corresponding loss curves are given in Figure 4(b).
Note: In the main paper, our reported numbers for rfdSON on Autoencoder benchmark in Table 2 for float32 experiments are erraneous. Please consider the numbers provided in Table 7 and the corresponding curve in Figure 4(a). Note that there's no qualitative change in the results and none of the claims made in the paper are affected. SONew is still significantly better than rfdSON. We also meticulously checked all other experiments, and they do not have any errors.
Autoencoder on KFAC, FishLeg, Eva: For completeness, We compare SONew against KFAC [38], FishLeg [20], and Eva [52] on Autoencoder benchmark as used in their official implementation.
(a) VIT train CE loss
(b) GraphNetwork train CE loss
Figure 5: Train loss corresponding to the best validation runs in Figure 1 (a) VIT benchmark (b) GraphNetwork benchmark. We observe that tridiag-SONew match or perform better than Adam.
(a) Best VIT train CE loss
(b) Best GraphNetwork train CE loss
Figure 6: Best train loss achieved during hyperparam tuning. (a) VIT benchmark (b) GraphNetwork benchmark. We observe that tridiag-SONew significantly outperforms Adam, while being comparable or better than shampoo.
Figure 7: Autoencoder benchmark run using Pytorch on KFAC, FishLeg, Eva, and tridiag-SONew. We notice that tridiag-SONew beats all other baselines by a large margin.
The main difference is their implementation uses ReLU activation compared to Tanh that we used for all our Autoencoder experiments. As these baselines done have JAX implementation, we use their official PyTorch implementation and run tridiag-SONew in PyTorch as well. Hyperparameter search is conducted for SONew similar to as reported above, over learning rate,  $\beta_{1},\beta_{2}$ , and  $\epsilon$ . For KFAC and Eva, rather than  $\beta_{2}$ , damping factor is tuned over  $[1e - 5,10]$  (default value specified is 0.03). kl Clip is tuned as well over  $[1e - 5,1.0]$ . Preconditioners are updates once every 15 iterations to have same wall clock time as other baselines and SONew. For FishLeg, auxiliary learning rate is tuned  $\in [1e - 7,1e - 1]$  and damping  $\in [1e - 5,1.0]$ . All other hyperparameters are tuned similar to SONew. Eva is trained for 100 epochs, and for other methods we change number of epochs such that
(a) Validation CE loss
Figure 8: We observe mixed results for 248M parameter language model benchmark. This is possibly because 48 trials were insufficient for optimal tuning. We leave thorough tuning and investigating into the above observation as future work.
(b) Train CE loss
each experiment takes same amount of time. Each optimizer is tuned using 600 hyperparameters. The results are in Figure 7, where notice that tridiag-SONew beats all the baselines by a large margin.
Adam vs SONew on 248M Language Model: We conduct an additional experiment on a 248M parameter transformer architecture [50] language model. The model is trained on WikiText103, introduced in [39]. We train the model for 3 epochs, having 8M tokens per epoch with a batch size of 8k tokens. We search over 48 hyperparameters, tuning learning rate  $\in 2e - 2$ ,  $1e - 2$ ,  $5e - 3$ ,  $1e - 3$ ,  $\beta_{2} \in 0.99$ ,  $0.999$ , weight decay  $\in 0.0$ ,  $0.1$ , and  $\epsilon \in 1e - 10$ ,  $1e - 8$ ,  $1e - 6$ , while fixing  $\beta_{1} = 0.9$ . Validation and training loss is given in Figure 8. Our observations indicate mixed results. While tds-SONew exhibits superior validation performance, Adam outperforms in training metrics. We believe that the 48 trials might have been insufficient for optimal tuning. It's conceivable that with further trials, SONew's training loss could surpass that of Adam. We leave this line of investigation for future work.
A.4.5 Convex experiments
As our regret bound applies to convex optimization, we compare SONew to rfdSON [37], another recent memory-efficient second-order Newton method. We follow [37] for the experiment setup - each dataset is split randomly in  $70\% / 30\%$  train and test set. Mean squared loss is used. For tridiag-SONew, we use a total of  $2 \times d$  space for  $d$  parameters. Hence, for fair comparison we show rfdSON with  $m = 2$ . Since the code isn't open sourced, we implemented it ourselves. In order to show reproducibility with respect to the reported numbers in [37], we include results with  $m = 5$  as well. We see in the Table 9 that tridiag-SONew consistently matches or outperforms rfdSON across all 3 benchmarks. Each experiment was run for 20 epochs and we report the best model's performance on test set.
Table 9: Comparison of rfdSON and tridiag-SONew in convex setting on three datasets. We optimize least square loss  $\sum_{t}(y_{t} - w^{T}x_{t})^{2}$  where  $w$  is the learnable parameter and  $(x_{t},y_{t})$  is the  $t^{th}$  training point. Reported numbers is the accuracy on the test set.
Table 10: (a) Dataset stats
Dataset	# total points	dimension
a9a	32,561	123
gisette	6000	5000
mnist	11791	780
Table 11: (b) RFD-SON vs tridiag-SONew
Dataset	RFD-SON, m=2	RFD-SON, m=5	tridiag-SONew
a9a	83.3	83.6	84.6
gisette	96.1	96.2	96.6
mnist	93.2	94.5	96.5
Table 12: Optimal hyperparams for Autoencoder Benchmark
Table 13: (a) float32 experiments optimal hyperparameters
Baseline	β1	β2	ε	lr
SGD	0.99	0.91	8.37e-9	1.17e-2
Nesterov	0.914	0.90	3.88e-10	5.74e-3
Adagrad	0.95	0.90	9.96e-7	1.82e-2
Momentum	0.9	0.99	1e-5	6.89e-3
RMSProp	0.9	0.9	1e-10	4.61e-4
Adam	0.9	0.94	1.65e-6	3.75e-3
Diag-SONew	0.88	0.95	4.63e-6	1.18e-3
Shampoo	0.9	0.95	9.6e-9	3.70e-3
tridiag	0.9	0.96	1.3e-6	8.60e-3
band-4	0.88	0.95	1.5e-3	5.53e-3
Table 14: (b) bfloat16 experiments optimal hyperparameters
Baseline	β1	β2	ε	lr
SGD	0.96	0.98	2.80e-2	1.35e-2
Nesterov	0.914	0.945	8.48e-9	6.19e-3
Adagrad	0.95	0.93	2.44e-5	2.53e-2
Momentum	0.9	0.99	0.1	7.77e-3
RMSProp	0.9	0.9	2.53e-10	4.83e-4
Adam	0.9	0.94	3.03e-10	3.45e-3
Diag-SONew	0.9	0.95	4.07e-6	8.50e-3
Shampoo	0.85	0.806	6.58e-4	5.03e-3
ztridiag	0.83	0.954	1.78e-6	7.83e-3
band-4	0.9	0.96	1.52e-6	4.53e-3