A Computationally Efficient Sparsified Online Newton Method

Summary

This paper proposes SONew -- a sparsified online Newton method, which uses the LogDet matrix divergence measure to sparsify the preconditioner. The LogDet divergence regularizes by matching the largest eigenvalues of the preconditioner matrix between iterates. It also leads to a convex optimization problem since the LogDet divergence is convex in its first argument. This leads to a simple and cost effective update rule for the preconditioner matrix. By minimizng the LogDet divergence they impose a sparsity constraint while requiring the sparse preconditioner to remain close to the full-matrix preconditioner in terms of LogDet divergence. They demonstrate the capability of their optimizer on autoencoders, vision transformers, graph neural networks, and large language models.

Strengths

SONew shows faster convergence compared to first order optimizers like Adam and AdaFactor. SONew can be applied to models sizes where Shampoo runs out of memory, while showing competitive convergence and generalization to these second-order optimizers. When using bfloat16, SONew has the least degradation in performance compared to all other optimizers.

Weaknesses

The experiments can be improved. Table 1 only shows the final value of the train CE loss and the time to achieve it, but without looking at the training curve it is difficult to see whether SONew is converging faster or simply reaching a lower loss but slowly. What is the criteria they used to stop the training? Is the learning rate schedule tuned for each optimizer? For the autoencoder experiments they compare with Adam, RMSProp, Shampoo(20), and diagonal, tridiagonal, and band-4 versions of SONew. For the GraphNetwork experiment they compare with Shampoo, Adam, and rfdSON. For LLMs they only compare with AdaFactor. This inconsistency makes it difficult see the effect of changing the dataset and model on all these different optimizers.

Questions

How robust is SONew to perturbation in the hyperparameters? Does it require less hyperparameter tuning compared to Adam?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

In the end, we need to run many hyperparameter trials to get any optimizer to achieve optimal convergence and generalization. Therefore, an optimizer that converges twice as fast, but requires more than twice the amount of hyperparameter tuning, is unlikely to be adopted by the community. When proposing a novel optimizer, there should be more emphasis on the robustness to change in the hyperparameters, model, and dataset. This is even more critical for LLMs, where we cannot afford to run many hyperparameter trials. The criteria for a good optimizer should be this kind of robustness, and not the convergence.

Summary

This paper strives to present a second-order optimization approach, called Sparsified Online Newton (SONew). SONew solves the regret minimization via introducing the LogDet matrix divergence. Using specific sparsified structures i.e. diagonal or banded sparsity, the proposed method could achieve considerable computational complexity. Finally, the authors show the effectiveness of the proposed method by performing the MNIST task with autoencoder, ImageNet task with ViT, OGBG-molpcba with GraphNetwork and language models.

Strengths

***Strengths*** 1. The paper is easy to follow and clearly written, although the layout of the paper is arranged too narrow. 2. The proposed method is mathematically sound and reasonable. The paper presents novel insight of online Newton methods in regards to regret minimization with LogDet divergence of the preconditioning matrix, and proposes an elegant way solve it via specific sparsified structures, and gives mathematical analysis in regards to the regret bound, which btw is O(sqrt(T)) the same with many related methods, and numerical stability. 3. Given the large computation cost and memory demand of the second-order approach, the paper shows some promising practical results in regards to the proposed methods.

Weaknesses

***Weakness*** 1. The authors seem to have largely changed the spacing between sections. The current spacing presented in the paper would be much more narrow than that in the latex template. Slight changes in the latex template could be accecptable, but such large changes could not be available. I could understand that the authors try to present as much content as possible. But the writing seems to be quite wordy. For example Abstract contains 24 lines, which btw is near 15 lines for other papers in common. I understand sometimes it is necessary to give a long abstract, but 16 lines in this paper are used to describe the comprehensive results of the propose method. Also, these results are described again in the Introduction. The authors should make the paper more succinct to fit the latex template. 2. For eq(3), the authors use LogDet matrix divergence to force the $X_{t}^{-1}$ and $X^{-1}_{t-1}$ to be "close". However, how the such closeness in LogDet matrix divergence could ensure the final term which is a complicated matrix multiplication in Eq3 to be safely omitted? Could the authors provide more details? 3. The authors could compare with more second-order approaches, such as K-FAC [36], K-BFGS [21], SON [33]. 4. It would be highly recommended that the authors would release the code. Due to lack of details in the experiment, it may be hard to reproduce the reported results. For example, what is this large language model with 1B parameters used in the experiment? 5. Line 135. The authors seem to omit the item $c_t$ in the presented formula, despite that this item will not make influence regarding this optimization.

Questions

See weakness.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

I have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on the learning algorithm. Still, it is highly encouraged to add corresponding discussions.

Summary

The paper proposes a new second-order method, named SONew, to optimize convex problems theoretically and conducts experiments on non-convex deep learning training tasks, which show SONew trains faster than first-order Adam and second-order fsdSON and Shampoo in some scenarios like the ViT benchmarks. The key idea of SONew is an online Newton method that uses an approximate LogDet divergence with a sparisified preconditioner. The sparse approximation of the LogDet divergence allows the efficient computation of second-order information so that SONew is compute-efficient enjoying exploiting some second-order information to accelerate the training. Experiments were conducted on various DNNs like Autoencoders, ViT, GNNs, and LLMs, showing SONew’s possible acceleration over some existing first-order methods and second-order methods.

Strengths

1. The proposed sparsified approximation of LogDet divergence is a new direction to reduce the computing cost using second-order information in optimizing DNNs. 2. A theoretical analysis is provided to give an insight for the sparsified preconditioner.

Weaknesses

1. The theoretical analysis is particularly for convex problems, so it may not hold as a DNN optimizer under non-convex optimization. 2. How to choose the sparse graph is unknown. The paper doesn’t conclude what sparse graphs should be used for different DNNs. 3. Experimental results show that SONew is worse than existing second-order methods like Shampoo (e.g., Table 1). In addition, some recent compute- and memory-efficient second-order methods (e.g., M-FAC [A] and Eva [B]) were not included in the comparison. [A] M-FAC: Efficient Matrix-Free Approximations of Second-Order Information, NeurIPS 2021. [B] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023.

Questions

1. In the experimental results in Section 5, some results show that SONew is obviously worse than Shampoo, e.g., Fig. 1(b), Fig. 2(b) in terms of iterations to loss (or accuracy). How to conclude that SONew achieves “up to 30% faster convergence, 3.4% relative improvement in validation performance, and 80% relative improvement in training loss” in the abstract? How about comparing to existing compute-efficient and/or memory-efficient second-order methods like M-KFAC and Eva? 2. Given a deep neural network for training, how we can choose tridiag-SONew, band-4-SONew, band-8-SONew, or band-$n$-SONew to achieve good training performance? 3. The paper provides a regret bound for SONew in the online convex optimization framework, but the optimizer is particularly used in DNNs which are known non-convex problems. How is the established bound used in DNN optimization? [A] M-FAC: Efficient Matrix-Free Approximations of Second-Order Information, NeurIPS 2021. [B] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

See above.

Summary

This paper proposes a memory-efficient and scalable second-order optimizer derived from the LogDet matrix divergence measure combined with sparsity constraints. The resulting preconditioner admits a banded/tridiagonal structure found via an efficient algorithm.

Strengths

- The paper is well-written and fairly easy to follow. The steps are clear and logically laid out. - I find the derivation of the preconditioner from the LogDet divergence interesting, even though variations of the particular form has been explored before. - The theoretical framework around the proposed method looks sound and the analysis of numerical stability is appreciated.

Weaknesses

- The banded/tridiagonal structure for which an efficient solution exists might not be natural, depending on the model architecture. A block-diagonal pattern such as in KFAC that models inter-layer dependencies would potentially be more suitable. - The paper is positioned as a practical, scalable optimizer for large-scale deep learning tasks. As such, I am not completely convinced by the experimental results. See breakdown below. - First, comparison with KFAC in terms of generalization and wall-clock performance would be important, as it is one of the very few higher order optimizers in deep learning that has shown adoption in practice. - Based on the presented experiments, it is hard to tell when one would pick the proposed optimizer over a simple Adam. Training loss value is not necessarily relevant, as second-order optimizers are known to often quickly converge to local optima with low loss but poor generalization. The plot showing validation error rate does not depict how this translates to classification accuracy. In case of minor/negligible differences one would pick Adam that offers less added complexity and faster wall-clock performance. Overall, there could be more focus on wall-clock performance, as performance with respect to number of steps is not that relevant in practice.

Questions

- How does the method compare with KFAC focusing on generalization on standard tasks explored in the paper and with respect to wall-clock time? - What is the practical advantage of the proposed method compared with well-established optimizers such as Adam? Can the proposed method accelerate training speed, in terms of time to achieve SOTA accuracies? If not, how much does generalization suffer compared with SOTA first-order optimizers? - Minor comments: - The notation for indexing the matrix in Eq. 12 is confusing. - Personal preference: I wouldn't use subjective terms such as embarassingly parallelizable or beautiful paper.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

Some limitations are discussed.

Summary

This paper introduces an online Newton method for optimization of convex losses, with the goal of applying it to large neural networks. In particular, the algorithm is derived by optimizing a preconditioning matrix on a regret bound and using the logdet divergence as a regularizer. It introduces a sparsification procedure to alleviate the computational burden of computing and inverting the full pre-conditioning matrix. With a fixed spasification structure (banded) of the preconditioning matrix, the problem regularized by the logdet divergence admits solutions that can be computed efficiently. The method is applied to a variety of (non-convex) benchmarks and compared with other optimizers.

Strengths

1) Fast optimization of neural networks is very important, to possibly reduce the amount of resources necessary for training neural networks. 2) Theoretical derivation of the algorithm is clear and sound.

Weaknesses

1) The derivation of the algorithm is based on the assumption of a convex loss function. Most of previous work discussed in the "Related work" section is about other online Newton methods, also designed for convex problems. However, all numerical experiments are shown on non-convex problems. This is problematic, because if the claim is that the method should be used for second-order optimization of non-convex losses, then it should be compared with the corresponding benchmarks: KFAC, K-BFGS, FishLeg. Indeed, the paper acknowledges that KFAC uses a similar sparsification structure (tridiagonal), but no empirical comparison is provided. On the other hand, if the claim is that the method should be used for convex optimization, then it should be clarified what are the (theoretical and empirical) advantages of the algorithm with respect to other online Newton methods. 2) Some of the statements made in the paper are misleading. For example: i) The algorithms provided on page 5 are not actually the algorithms used for the numerical experiments. "Adam grafting" is used on top of those algorithm, but no justification is given, also in the light of the proposed theory and derivations. My assumption is that none of the theory developed holds after including "Adam grafting" in the algorithm, and a description of the full algorithm should be given instead of just giving a reference to another paper. ii) The following statement in section 5 is not true: "We compare SONew against widely used first order methods including SGD, SGD with Momentum, Nesterov, Adagrad, Adam, and Rmsprop." The first four of the six optimizers mentioned in the list are actually not compared with. iii) The abstract says that the proposed method is slower than first-order methods, but no results are shown in the results about this. In fact, figures show that the proposed method is faster in terms of number of steps, but what is the significance of this results if the method is overall slower? 3) The results of numerical experiments are underwhelming. As I already mentioned, no comparison with other second-order methods is given, besides Shampoo, and comparison with simple and strong benchmarks is also not provided. For example, SGD with momentum is never compared with, and neither Adam nor SGD are compared with in the case of LLM in figure 3. Furthermore, it is not clear if the method provides a reliable advantage. Only single training runs are shown, and fluctuations look pretty large in all figures. Averaging over multiple training runs would help understanding how reliable is the method. Finally, training curves are shown only vs epochs or number of steps, no results in terms of wall clock time are provided, therefore it is hard to tell whether the method provides any real advantage.

Questions

See above

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

See above

Summary

This manuscript proposes SONew, a second-order algorithm that uses sparsity in the preconditioner to distil to linear time and space complexity. Different sparsity patterns (tridiagonal, banded, etc.) allow generalization of the algorithm for training deep neural networks almost as efficiently as first-order methods. SONew achieves up to 30% faster convergence compared to Adam and improves validation performance when trained for the same number of iterations when training with the Adam optimizer.

Strengths

• The paper provides a strong theoretical foundation for sparsifying the preconditioning matrix, making a clear connection between regret minimization and LogDet divergence of the preconditioning matrix • The paper provides an extensive suite of experiments indicating the computational efficiency and performance gains of their method on large scale model training • The authors provide theoretical guarantees of the numerical stability of their method, which is a serious concern for existing second-order approximation OCO algorithms • The paper is clearly written

Weaknesses

The paper does a good job providing both theoretical and empirical justification for their method. However, there are a few details about the experimental setup and bottlenecks for baselines that would improve the manuscript. The limitations section should also be expanded to explain under what conditions SONew would not be optimal or theoretical guarantees would fail. Details provided below in the questions section:

Questions

• The authors compare Shampoo with RMSProp grafting but use Adam grafting for their method. Adam grafting would yield a more accurate step size compared to RMSProp. It would be beneficial to add a baseline of Shampoo with Adam grafting • Fig. 1a mentinos that hyperparameter search could not be performed on Shampoo on 16GB TPU. Could the precision be reduced (bfloat16) to run shampoo?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors need to expand on the limitation of their method and point out specific cases where their method may be disadvantaged compared to existing methods.

Summary

This paper propose a memory-efficient second-order algorithm named Sparsified Online Newton (SONew) method. By introducing the LogDet matrix divergence for regret analysis, the authors derives an online Newton update rule. Then the authors introduces structured sparsity patterns on the preconditioner matrix on the preconditioner matrix to reduces storage and computational complexity. They adopt the proposed methods on several benchmarks and show SOTA performances comparing to other optimization methods.

Strengths

1. The writing and organization of the paper are very good, the notation representation is accurate. 2. The method proposed in the paper is reasonable, and the paper applies the proposed method on some benchmarks, achieving better results compared to other optimizers.

Weaknesses

1. As an optimization method with lower resource usage, the authors only mention the theoretical reduction in space complexity, and the actual decrease in training time compared to other optimizers. However, they do not provide actual comparisons of GPU memory usage. 2. The authors evaluate the proposed optimization method on deep autoencoders, Vision Transformers, and Graph Neural Networks. However, I think that as a newly proposed optimization method, the authors need to conduct experiments on more benchmarks and more network architectures to verify its general applicability, such as results of more CNN methods on ImageNet, and generative models in computer vision like Diffusion Models.

Questions

1. Could the author provide the actual GPU memory cost and comparing the proposed method with other methods? 2. Could the author provide more experiments to show the generalization of the proposed methods?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The author does not discuss the limitations and potential negative societal impact of their work. I consider one of the limitations of the proposed method is that it can not be guaranteed that this method can achieve better performance compared to other optimizers on more datasets and network architectures.

Thanks for the authors' kindly response.

Thanks for the authors' kindly response. Overall, I think the proposed method is interesting and promising, and may have further impact in regards to the area of learning methods. The authors promise to make revisions in their final version to address the concerns. I recommend for acceptance.

Review

Thanks for the authors' kindly response. Overall, I think the proposed method is interesting and promising, and may have further impact in regards to the area of learning methods. The authors promise to make revisions in their final version to address the concerns. I recommend for acceptance. I appreciate the authors' response. Regarding the points above 1. Using Adam grafting with shampoo would help ablate the importance of grafting type to different optimizer settings 2. mixed precision training (including w/ bfloat16) is the norm in LLM training. this would substantially improve comparisons regarding memory 3. it's helpful that the authors listed an additional limitation, but it is important that the final version has a deeper dive into the limitations and potential impact of these limitations. b/c of these points i am maintaining my score

Official Comment by Authors

We hope that the rebuttal clarifies questions raised by the reviewer. Please let us know if we can add more to justify the contributions and significance of our work. We kindly request you to reconsider your rating if our responses were sufficiently able to address your concerns.

Official Comment by Authors

We hope that the rebuttal clarifies questions raised by the reviewer. Please let us know if we can add more to justify the contributions and significance of our work. We kindly request you to reconsider your rating if our responses were sufficiently able to address your concerns.

Official Comment by Authors

We hope that the rebuttal clarifies questions raised by the reviewer. Please let us know if we can add more to justify the contributions and significance of our work. We kindly request you to reconsider your rating if our responses were sufficiently able to address your concerns.

Thank you for the clarification

The authors have addressed my concerns in their rebuttal. I am bumping up my score to an accept.