Summary
The paper presents Online Subspace Descent, a memory-efficient modification applicable to a wide class of gradient-descent based algorithms where low-rank projections of gradients can be employed to reduce the memory overhead. Contrary to recent techniques such as GaLore, which require infrequent but costly updates of the projection matrix, the proposed method updates the projection matrix in an on-line manner, and proves the convergence of the resulting algorithm.
Strengths
Memory-efficient optimizers are crucial to enable LLM training/fine-tuning outside of huge datacenters; as such, this is a potential high-significance paper. While I think the idea of gradually updating the projection matrix is pretty much "the obvious thing to do", I really like that the authors didn't just propose "Adam[Add-some-random-letter-here]", but instead took a step back and showed that the same principle applies to a wide class of optimization algorithms.
Weaknesses
* There are some grammatical errors and typos, I think (not a native speaker), but not to the degree that they hamper the understanding of the text (e.g., I think it should be: update rules of _the_ projection matrix; in the subspaces that dynamically _change_, How Should P be _updated_, is not a problem for optimizers that _yield_)
l.1ß6: `Note that the update of P t can be done in parallel with that of 106 (W t, ˆ St), and incurs no slow down once it is fast enough to not cause a speed bottleneck` Generally, Adam updates are memory-bound operations, so I'm not sure if you can schedule a second operation concurrently without slowing down the update. Also, the same could be said for GaLore, you run you SVD in parallel to the main network; if it takes multiple steps, you could just wait till it is done, and update only afterwards, so GaLore can be made as parallel as the algorithm that is proposed here, I think.
l.119: `In this work, we propose to update P t in a continuous online fashion that incorporates the most 119 recent gradient information in a timely fashion, without calling linear algebra routines.`
Matrix multiplications are linear algebra routines, so I think most P updates will involve some linear algebra.
I think the description around l. 119-128 could do with some rewriting. The sentence `we propose to update P t at time t by minimizing the following PCA objective` is misleading, because P_t is not actually updated like that; instead P_t is one step along some optimizer that decreases this objective. Second, `Instead of minimizing [...] exactly, to retain computational efficiency, we propose` seems to be misleading, too; in an earlier paragraph, it was GaLore was criticized for using only a single gradient information for its update. This sentence suggests that if you had enough compute, you would do the same here, only due to computational constraints do you end up with something that integrates multiple gradients, which I think goes against the message you want to make in the paper.
l. 170: `There is no requirement on Γ here, besides that the derivative in (8) should exist` Γ does not appear in (8), so it isn't clear which derivative should exist.
l. 223: `The typical Pytorch implementation of SVD can be up to 142 times slower than running a single step online PCA on representative weight tensors from LLaMA architectures.` Given that GaLore runs SVD only every 300 steps, this would make GaLore comutationally more efficient than the proposed method.
Questions
It has been observed that scaling up LLMs is notrivial also from a numerical stability point of view, e.g., in quantization, starting with 7B models, outliers become much more of a problem than they are for smaller models. As such, it would be great to see this method applied to larger models, at least to 7B, ideally even larger. If compute is a problem, maybe scaling the amount of training steps less than would be optimal could help, to at least have some proof-of-concept for the method at larger model scales.
l.145: `this is a mild condition that can be satisfied e.g., when P_t is updated by (online) PCA on G_t` Is this really true, in the generality that is suggested here? If so, I'd like to see a formal argument. The paper only fixes the objective eq. (6), but leaves the choice of optimizes to the user, and my intuition says it should be possible to construct an example where minimizing (6) with a momentum-based optimizer could lead to $<G_t, P_t>$ being zero without $G_t$ being zero.
l 158, I'm not sure what "once" is supposed to mean here. Should it just be "if"?
Regarding the convergence guarantee for arbitrary $P_t$, I think it would be illuminating to consider what happens in case of an adversarial choice; for example, for SGD with momentum, choose $P_t = min <G_t, P_t^T P M_t>$, i.e., choose P so that it projects the momentum buffer _against_ the direction of the current gradient. Why does this choice of P not break convergence?
l. 235
> In conclusion, given sufficient time and resources, higher ranks yield better performance for Online Subspace Descent. It is recommended to select the highest rank until the perplexity reduction saturates.
Can this reliably be selected a priori, i.e., does picking a rank optional for running optimization for, e.g., 10k steps, also result in optimal rank for training until "convergence"
Limitations
Evaluations for production-size LLMs are missing.