Paper: Curvature-Informed SGD via General Purpose Lie-Group Preconditioners (ICLR-2024)

Reviewer comment (Presentation):
First question is that in Proposition 2.1, the authors refer to equation 6. But there is no equation 6 in the main part of the paper. Does the authors refer to the equation 6 in the appendix? Then the authors should put this equation in the main part of the submission.

Author response:
Yes, you can click it to bring you to the eq. 6 in the appendix. We will try to squeeze that equation into the main body to improve reading experience.

Reviewer comment (Theory):
In Corollary 2.1.1 the authors mention the linear convergence rate, but what is the exact or explicit upper bound for this linear convergence rate? The authors should explicitly present this linear convergence rate not just present an asymptotical convergence result.

Author response:
Corollary 2.1.1 is just the immediate results of proposition 2.1 and text book results. We just want to show that PSGD recovers SGD. One can follow the standard technique for SGD to give a specific convergence rate for PSGD for each specific setting. We have no intent to elaborate in this direction.

Reviewer comment (Theory):
In Corollary 2.1.2, the authors claimed that the algorithm can cover the quadratical convergence rate of Newton's method. However, the inequality in this corollary only shows the linear convergence. Because $\|g_t\| \geq 2\alpha(L(\theta_t) - L(\theta_*))$, we can get that $L(\theta_{t + 1}) - L(\theta_*) \leq (1 - \frac{\alpha^2}{\beta^2})(L(\theta_t) - L(\theta_*))$. However, this is only linear convergence rate. It's not quadratical convergence rate as Newton's method. And the linear convergence rate is even worse than the gradient descent of $1 - \frac{\alpha}{\beta}$.

Author response:
Again, corollary 2.1.2 is just the immediate results of proposition 2.1 and text book results. We just show that PSGD recovers Newton's method. The Newton method does not always have quadratic convergence, it only converges quadratically around the basin of attraction. One can follow the standard textbook technique to decide the size of that basin of attraction. We do not claim any stronger results, no intention to elaborate in this direction.

Still, we apologize that the presentation of corollary 2.1.2 can be improved to avoid such possible misunderstanding. Will revise it.

Reviewer comment (Theory):
There is no overall theorem for the convergence rate of the proposed algorithm. We need a specific detailed theoretical convergence rate of the PSGD to make comparison with other methods.

Author response:
Again, Proposition 2.1 says that P must converge to inv(abs(Hessian)). Then one can follow the standard techniques to give a specific convergence rate for PSGD in each case, by specifying settings like a) Lipschitz constant, 2) learning rate,  3) distance between initial guess to final solution, 4) alpha, 5) beta in corollary 2.1.2, .... 

PSGD just provides inv(abs(Hessian)) as a preconditioner, and thus gives you a learning rate free method, and recovery of Newton method for convex problems. We do not intend to prove any stronger results in corollary 2.1.1 and 2.1.2, also not the focus on the paper. Still, corollary 2.1.1 and 2.1.2 are not nonsense as not all preconditioners can recover the Newton method in a convex setting.

Reviewer comment (Theory):
The authors didn't give any theoretical analysis for the step size. They only presented some empirical results for the step size choices.

Author response:
All the step sizes are normalized and scaleless in theory, not empirically. It’s completely a learning rate free method. On Lie group, you update as  (I + step * Generator) * Q, and thus you know how much you updated by checking ||step * Generator||. For parameters, all the eigenvalues of P * Hessian have unit absolute value, and again, you know how much you have updated the parameters as ||delta preconditioned_gradient|| = ||delta parameters||. Nevertheless, a scheduler is needed in a stochastic setting, and we do need to tune the scheduler.

Reviewer comment (Theory):
From the presentation of section 2.3, the "Lie group" only refers to orthogonal groups and general linear groups, which is a much more limited case than general Lie groups.

Author response:
1, Only limited to orthogonal or general linear group? 

No. We do not limited to orthogonal group. We only use orthogonal group to illustrate that coordinate change consists of rotation and scaling.  

Yes. We assume a general linear group to prove Proposition 2.1. 

No. The two proposed preconditioners (matrix-free and low-rank approximation) are not limited to either orthogonal or general linear group. They all are sparse, and have linear complexity (the same order as SGD and Adam). The orthogonal and general linear group have quadratic complexity, which means that they need 10000000 billion elements for preconditioning a network with 100 millions parameters!

Reviewer comment (Writing):
A brief introduction on Lie groups is needed, at least in appendix, this can be part of section D.1. Since gradient descent on manifold has been mentioned as well, it is necessary to define it formally, especially the form that is used in the Lie groups of this paper.

Author response:
2, Introduction of Lie group. We do have some background materials on Lie groups in appendix D and G. Readers with knowledge of group generators should find them easy to follow; we hope that readers knowing the definition of Lie group can follow them as well.

Reviewer comment (Presentation):
Same issue happens in experiments. There are much more experiments than a theory oriented paper, but it is hard to follow the whole experiment sections due to current presentation.

Author response:
3, Much more experimental results than a theory oriented paper. 

We do not intend to have a theory oriented paper. The general theory of PSGD is established in previous Li’s work (the ICLR 2018 one). However, that work only considers affine groups and thus have three limitations: a) not all network blocks can be reparameterized as affine transforms, although many, say CNN and attention, can; b) no interaction among affine transform preconditioners of different layers or blocks; c) it is inconvenient to use the affine group preconditiioner as it requires to reparameterize all blocks into their affine transform forms. We fully disclosed these backgrounds in Section I and II. One major contribution of this paper is the low-rank approximation (LRA) preconditioner. The math for LRA preconditioner on Lie group is much heavier than in the Euclidean space. This makes our paper full of math equations, still we do not intend to make it a theory oriented paper.          

Readers could just focus on the results listed in the main body. We provide so much more experimental results because even up to today, most people still think that 2nd order optimizers only work for small scale toy problems (although true for BFGS, LM-BFGS, HF, CG, …). We really hope that all these readily reproducible experimental results in this paper could let people give 2nd order optimizers a second look.

Reviewer comment (Writing):
Inconsistence of terminology. In sentences, "Theorem" is used but in presentation of the results formally, "Claims" have been used for no reason.

Author response:
4, Inconsistency of terminology. 

Apology for the confusion. We will correct them. Our intentions for these terms are:
Proposition) important theoretical results;
Corollary) immediate results from propositions and known results;
Claims) can be proved with straightforward algebra operations, although still could be math heavy.

Reviewer comment (Experiments):
In each experiment, what is the specific constraints?

Author response:
5, Constraints for each experiment. 

We guess the reviewer is inquiring about resource constraints. We have details in the appendix, and code for reproduction. A rough idea of the complexity of LRA method (the other one is cheaper): a) 1.2 x wall clock time of SGD; b) 2~3 x memory consumption of SGD; c) less tuning time as all the step sizes are normalized.

Reviewer comment (Theory):
In proposition 2.1, how is the step-size set?

Author response:
6, Step size in proposition 2.1. 

On the Lie group, Q is updated as (I + step * Generator) * Q. To make an infinitesimal update, we need || step * Generator || << 1. Then we can choose the step as a_normalized_step / || Generator ||, where a_normalized_step << 1, say 0.01. This way of normalization does not always work. But here it well works (reason: see algorithm 2 in page 5, the two terms of nabla_U are homogeneous). Also, the spectral norm estimation is tight (within 3 dB, see D.1.2.2). So no implementation difficulties.

Reviewer comment (Theory):
Contribution of theoretical analysis can be stronger, there are sufficently many works on optimization on manifolds, what is the main challenging of this paper compare to other related works? Is that possible to leverage other works including rates to enhence the results of this paper?

Author response:
7, Stronger theoretical results? Agree that the theoretical results (corollary 2.1.1 and 2.1.2) are pretty plain, mainly for the completeness of the paper. Our focus is on the practical side: how to construct the LRA preconditioner and how effective it is on large scale ML problems. Difficulties for theoretical studies are from two sides. a) the PSGD preconditioner fitting criterion does not have a closed-from rank-2 preconditioner update rule even in the Euclidean space (BFGS criterion has, but does not work for large scale ML problems as it ignores gradient noises). b) Update on Lie group is nice as we can normalize the step size and have all the nice properties of Lie group. But, Lie group is a manifold plus extra constraints: the way we update the preconditioner cannot break the closeness of the Lie bracket. Still, a good direction worthy closer check.

Reviewer comment (Novelty):
Theoretical frameworks is mostly ihnerited from that in Xi-Lin Li's papers, especially one that is titled "PRECONDITIONER ON MATRIX LIE GROUP FOR SGD" and contains the majority of ideas ensuring success of an approach of the paper under review.

Author response:
Yes, criterion for preconditioner estimation and the idea of Lie group fitting are from Li’s previous work. We fully disclosed these details in Section I and II. 

But, the affine Lie group preconditioner developed in Li’s ICLR work have three limitations: a) not all network blocks can be reparameterized as affine transforms, although many, say CNN and attention, can; b) no interaction among affine transform preconditioners of different layers or blocks; c) it is inconvenient to use the affine group preconditiioner as it requires to reparameterize all blocks into their affine transform forms. 

So this paper proposed two black box style preconditioners. One can use them as easily as the plain SGD. Notably, the low-rank approximation (LRA) preconditioner is of significant theoretical importance (LRA tends to fit real world patterns well) and practical importance (makes PSGD have the same order of complexity as SGD, equally easy to use).

Reviewer comment (Theory):
Can you provide convergence rate guarantees which do not include conditioning number? The presence of condition number in rate makes rate similar to that of SGD without preconditioning. Taking into account, that purpose of preconditioning is in fact getting rid of conditioning number impact, theory seems to fail (in non-convex case).

Author response:
Such a picture makes intuitive sense, but still many hinders in math. PSGD just gives a learning rate free method as P*H has unit absolute eigenvalues, but loss still could be arbitrarily rough, like a large Lipschitz constant, and could be full of saddle points for nonconvex problems. A convergence rate independent of condition number would be extremely strong! Even for the standard Newton algorithm on strictly convex problems, text book results can only show that the Newton method converges quadratically around a small basins of attraction, generally no globally quadratic convergence. Away from the attractor, Newton method with reduced step size also converges linearly as gradient descent.     

To clarify, Proposition 2.1 is completely ours, stating that P must converge to inv(abs(Hessian)). Corollary 2.1.1 and 2.1.2 are the immediate results of Proposition 2.1 and established results from convex optimization. Corollary 2.1.1 says that PSGD is no worse than SGD for nonconvex problems, and Corollary 2.1.2 says that PSGD recovers the Newton method for convex problems. We claim no stronger results. Sounds a little pessimistic in theory, but in practice, a good preconditioner helps a lot for sure as shown in our empirical results. At least, less tuning effort due to normalized learning rate.

Reviewer comment (Novelty):
Another criterion from Lin's original papers leads to Fisher matrix as an optimal preconditioner. Can you explain why this important case was left outside of the framework you propose? This case is interesting because it has no quasi-Newton property, but is greedy-optimal for quaite natural criterion at the same time. Whithout considering both cases, your approach seems to be just a detalisation of particular case of Lin's approach.

Author response:
It’s true that we only consider the Newton-type preconditioner. This paper benchmarks PSGD on a wide range of problems, from classic convex optimization to reinforcement learning. The Fisher-type is for sure irrelevant for some of the problems. Nevertheless, we agree with the reviewer that this is a direction worthy of a closer look. Luckily, both types share the same math. The same LRA tool also applies to the Fisher-type.    

Why Fisher-type is not good for general optimizations. a) the Fisher-type works only if the Fisher information matrix can be defined. Unfortunately, this is not always the case. Many regression and reinforcement learning problems do not use divergence or xentropy as losses. b) Let’s narrow down to problems with Fisher information matrix defined. Then, certain regularizations, say L2, could be absorbed by adding a damping term to the Fisher matrix. Still, many other regularizations, say norm of gradients, are incompatible with the use of Fisher information matrix. c) Let’s put aside the regularization term. Then we face two choices for the Fisher: empirical Fisher or the true Fisher. Using the empirical Fisher could be problematic. Using the true Fisher makes the Fisher-type method reduce to the Gaussian-Newton method for pdf from the exponential family, which eventually is an approximation of the Newton-type method by discarding the cross-terms in the Hessian.

Reviewer comment (Writing):
Unfortunately, since I am not familiar with Lie groups, it is rather hard for me to fully understand the contribution of the work. If you could provide more background on this topic (maybe even in Appendix), it would be very helpful.

Author response:
1, Background on Lie groups or PSGD?

Lie group) We put some background material on Lie groups in appendix D and G. Nevertheless, this is a conference paper, we keep them short. Maybe particle physicists have nice Lie group introduction materials.  

PSGD) The general theory of PSGD and PSGD on the affine Lie group are not new (Li’s ICLR 2018 work). We disclosed them fully in Section I and II. But, PSGD on affine Lie group have three limitations: a) not all network blocks can be reparameterized as affine transforms, although many, say CNN and attention, can; b) no interaction among affine transform preconditioners of different layers or blocks; c) it is inconvenient to use the affine group preconditiioner as it requires to reparameterize all blocks into their affine transform forms.

Reviewer comment (Presentation):
It is rather weird to see theoretical convergence results in the Background section. If this is not your result, please provide the citation of the work, where from you have taken this result (and other theorems/corollaries). If this is your result, please move it to other section with results.

Author response:
3, Putting convergence results in the background? 

Proposition 2 is completely ours. Corollary 2.1.1 and 2.1.2 are the immediate results of Proposition 2 and established results from optimization theory. We will move them into Section 3.