Memory-Efficient LLM Training with Online Subspace Descent

ArXiv: 2408.12857

🎯 Pitch

This paper proves that low-rank optimizers like GaLore can converge under arbitary projection update rules—not just expensive periodic SVD—and then shows that replacing SVD with a cheap online PCA step actually improves perplexity over GaLore while cutting subspace update cost by up to 142×. The key insight is a dynamic projection matrix that continuously tracks gradient statistics via gradient descent on the projection itself, making low-rank training both theoretically grounded and practically faster.

---

1. Executive Summary

This paper proposes Online Subspace Descent, a new family of memory-efficient optimizers for large language model training that replaces the expensive singular value decomposition (SVD) used in prior low-rank methods like GaLore with a lightweight online PCA update for the projection matrix. By continuously updating the projection matrix through gradient-based optimization of a PCA objective rather than periodic exact SVD, Online Subspace Descent reduces computational overhead—single-step online PCA can be up to 142× faster than equivalent SVD calls on representative weight tensors—while providing the first convergence guarantee for arbitrary projection matrix update rules under the Hamiltonian descent framework (applicable to common optimizers including Adam, LION, and Adafactor). Pretraining LLaMA models ranging from 60M to 7B parameters on the C4 dataset demonstrates that Online Subspace Descent achieves lower perplexity than state-of-the-art low-rank methods (31.30 vs. 35.52 for GaLore on a 1B model at rank 512) and narrows the gap with full-rank baselines, establishing that continuous online adaptation of the projection subspace can outperform periodic exact decomposition only when the intrinsic rank of the optimization problem is sufficiently high to benefit from timely incorporation of recent gradient information.

2. Context and Motivation

The Core Problem: Memory-Efficient Training Without Sacrificing Convergence

The fundamental tension this paper addresses is deceptively simple: Adam produces the best results for LLM pretraining, but its memory cost is prohibitive for large models. To understand why, recall that standard gradient descent needs only to store the model weights and one gradient buffer during training. Momentum adds a first-order momentum buffer M_t of equal size to the weights. Adam adds two buffers—M_t (first-order momentum) and V_t (second-order momentum / preconditioner)—each the same shape as the model's weight matrices. For a model with $N$ parameters, Adam stores roughly $2N$ additional values in optimizer state, meaning total memory scales as $\sim 3N$ (weights + M_t + V_t), compared to $\sim 2N$ for momentum and $\sim 1N$ for vanilla SGD with gradient accumulation. When $N$ reaches billions (e.g., LLaMA 7B at ~7 billion parameters), this overhead becomes the dominant constraint on which hardware can feasibly train the model.

The problem is not merely that memory-efficient optimizers exist (many do—Adafactor, LION, 8-bit quantization techniques), but that none of them match Adam's convergence quality on LLM pretraining without trading off compute or wall-clock time. Empirically, the field has converged on Adam/AdamW as the optimizer of choice for pretraining large language models because momentum-plus-preconditioning produces better minima and faster convergence on language modeling tasks. The paper's Table 6 (Appendix A.3) illustrates this clearly: on LLaMA 60M, vanilla AdamW8bit achieves 29.77 perplexity, while LION achieves 52.65 and Adafactor 33.45. The gap widens at scale—making Adam-like optimization quality essential for competitive pretraining.

Prior Approaches: Low-Rank Gradient Projection via Periodic SVD

A recent line of work—most prominently GaLore (Zhao et al., 2024) and Sketchy (Feinberg et al., 2024)—proposed an elegant solution: project gradients into a low-rank subspace before feeding them into the optimizer. The key observation driving these methods is that gradients during deep learning training often live in a low-dimensional subspace of the full parameter space (Gur-Ari et al., 2018; Larsen et al., 2021). Instead of maintaining full-size optimizer states M_t and V_t in $\mathbb{R}^{n \times m}$, one can:

1. Periodically compute a projection matrix $P \in \mathbb{R}^{n \times k}$ (where $k \ll n$) from the dominant singular vectors of the gradient.
2. Project the gradient $G_t \in \mathbb{R}^{n \times m}$ into the subspace: $\hat{G}_t = P^\top G_t \in \mathbb{R}^{k \times m}$.
3. Run the optimizer (e.g., Adam) on the reduced-dimension $\hat{G}_t$, maintaining states $\hat{S}_t$ of size $k \times m$ rather than $n \times m$.
4. Map the reduced-dimension update back to full space: $W_{t+1} = W_t + P \cdot \text{update}(\hat{S}_t)$.

This approach reduces optimizer state memory from $O(nm)$ to $O(km)$, achieving significant savings when $k \ll n$. GaLore demonstrated that this works: pretraining LLaMA models with low-rank projected Adam achieved reasonable perplexity while using substantially less GPU memory than full-rank Adam.

Where GaLore and Similar Methods Fall Short

The paper identifies two specific, practical failure modes of the periodic-SVD approach:

Failure 1: Computational overhead from exact SVD becomes prohibitive at scale. GaLore recomputes the projection matrix $P$ via torch.linalg.svd(G_t) every $T$ iterations (e.g., $T=200$). While this is much less frequent than every step, SVD on weight-matrix-scale tensors remains expensive—the paper's Figure 2 shows that torch.svd on typical LLaMA weight tensor shapes can be up to 142× slower than a single gradient step of online PCA on the same tensor. For the 7B model scale, this overhead translates directly to longer training wall-clock time (Table 3: GaLore takes ~9.74 hours vs. ~7.14 hours for Online Subspace Descent, both on the same hardware). The cost of SVD cannot be "hidden" by overlapping with weight updates because the SVD computation time exceeds the update step time by such a large margin.

Failure 2: Stale projection matrices degrade optimization quality. GaLore's projection matrix $P_t$ is set equal to the SVD of $G_{T\lfloor t/T \rfloor}$—the gradient at the most recent periodic recomputation point. As training proceeds, the gradient distribution shifts (the loss landscape changes, different tokens/layers become active). Between recomputation intervals, the same $P$ is used to project all gradients, but this $P$ captures the dominant directions of a past gradient snapshot, not the current ones. This staleness is not merely a theoretical concern: the paper's experimental results (Table 1, Figure 1) show that Online Subspace Descent's continuous updates yield substantially better perplexity than GaLore at the same rank (31.30 vs. 35.52 on LLaMA 1B at rank 512), with the gap growing as model size increases. The hypothesis (Section 5.1) is that larger models have higher intrinsic rank optimization problems, making timely gradient information incorporation more important.

Failure 3: No theoretical understanding of how changing projection matrices interact with optimizer states. This is the deepest gap the paper identifies. In GaLore and similar methods, the projection matrix $P_t$ changes periodically, but the optimizer states $\hat{S}_t$ accumulate information from gradients projected through different $P_\tau$ matrices at different times $\tau \le t$. This creates a subtle entanglement: the momentum buffer $\hat{M}_t$ contains a weighted sum of gradients $\hat{G}_\tau = P_\tau^\top G_\tau$, but when the optimizer later uses $P_t$ to map the update back to full space, the directions accumulated from old $P_\tau$ may be incoherent with the current $P_t$ directions. Prior to this work, no one had provided a convergence guarantee for any method that dynamically changes the projection matrix in the context of general optimizers on non-convex objectives (the paper explicitly notes this in Section 1). Stochastic Subspace Descent (Kozak et al., 2019) provided convergence analysis only for random projection matrices on convex functions—a setting that doesn't capture deep learning training with adaptive optimizers.

This theoretical gap has practical consequences. Without understanding why and under what conditions dynamic projection works, practitioners have no guidance for designing the update rule $\chi_t$ for $P_t$, leaving them to heuristically tune update frequency, rank, and projection method. Worse, there is no framework for predicting whether a given $P_t$ update rule will cause the optimizer to diverge or converge to poor minima—a real concern given the complex interplay between the evolving subspace and the optimizer's momentum state.

How This Paper Positions Itself

The paper positions itself across two axes simultaneously:

Axis 1: A practical improvement over GaLore. On the methodological front, the paper proposes replacing GaLore's periodic, exact SVD with a continuous, gradient-based update of $P_t$ via online PCA. The key technical insight: instead of periodically solving svd(G_t), maintain $P_t$ as an evolving parameter updated by a lightweight optimizer step on the PCA reconstruction loss $L_{G_t}(P) = \|PP^\top \tilde{G}_t - \tilde{G}_t\|^2 + \lambda\|P^\top P - I_{k \times k}\|^2$. Each step of this update is fast enough to execute in parallel with the main weight update (Figure 2), and continuously incorporates the most recent gradient information. This directly addresses both Failure 1 (SVD overhead) and Failure 2 (stale projections).

Axis 2: The first theoretical framework for dynamic subspace optimizers. On the theoretical front, the paper provides a Lyapunov function analysis within the Hamiltonian descent framework (Maddison et al., 2018) that proves convergence for arbitrary update rules of $P_t$—not just online PCA, but any smooth update rule—when applied to optimizers with the Hamiltonian+descent structure (Adam, LION, momentum). This directly addresses Failure 3 (the theoretical gap) and provides the principled foundation that prior work lacked.

The paper is careful to distinguish its approach from LoRA (Hu et al., 2021), which also uses low-rank structures but in a fundamentally different way. LoRA permanently constrains weight updates to live in a fixed low-rank subspace ($W = W_0 + AB$, with $A,B$ low-rank), which is suitable for fine-tuning but too restrictive for pretraining where the model needs to explore the full parameter space. Online Subspace Descent, in contrast, dynamically rotates which subspace is used at each iteration, so that over the course of training, the effective search space is not restricted to any single low-rank manifold—the trajectory can span the full parameter space through sequential projection onto different subspaces. The authors frame this as "analogous to coordinate descent, except in a continuous fashion" (Section 3.2), with $P_t$ determining which coordinate subspace is active at each step.

The paper's main positioning move is to reframe the memory-efficient training problem from "find a single good low-rank approximation of the gradient" (the GaLore/Sketchy approach) to "maintain a continuously adapting subspace that tracks the evolving gradient geometry." This reframing has both practical benefits (lower overhead, better perplexity) and theoretical benefits (the Lyapunov analysis that provides convergence guarantees), and it opens up a design space—what optimizer to use for $P_t$, what rank to select, what regularization to apply—that the paper partially explores in its experiments.

Where the Hamiltonian Descent Framework Fits In

The theoretical contribution builds on an observation that may seem surprising at first glance: the Hamiltonian descent framework was developed to study the continuous-time dynamics of standard optimizers, not memory-efficient ones. The paper's key theoretical move (Section 4.1) is to show that when subspace projection is introduced into the Hamiltonian descent system, the cross terms cancel due to the adjoint property $\langle P_t X, Y \rangle = \langle X, P_t^\top Y \rangle$, regardless of how $P_t$ evolves. This means the Lyapunov function $H(W, \hat{S})$ that guarantees descent for the full-rank optimizer also guarantees descent for the projected version—and this holds for any $P_t$ update rule with a smooth continuous-time limit. The paper is explicit that this is "surprising" (Section 4) because one would naturally expect that an arbitrary dynamically changing projection matrix could interfere with the optimizer's descent dynamics. The adjoint-based cancellation shows that this intuition is wrong for the Hamiltonian+descent family: the projection operator is effectively "invisible" to the Lyapunov function because the same projection appears symmetrically in the weight update (via $P_t$) and the state update (via $P_t^\top$).

This theoretical result has immediate practical implications that the paper exploits: it means practitioners can experiment with any smooth $P_t$ update rule (SGD on the PCA loss, Adam on the PCA loss, even random perturbations around a good basis) and be guaranteed that the optimizer will not diverge due to the projection mechanism itself. The only additional requirement for convergence to local optima (Theorem 4.5) is a mild non-degeneracy condition: $P_t$ should not become stuck in a state where $P_t^\top G = 0$ while $G \neq 0$—which is naturally satisfied by any reasonable PCA-like update that seeks to reconstruct $G_t$.

3. Technical Approach

3.1 Reader Orientation

This is primarily a theoretical analysis paper with a practical algorithm contribution. The core idea is that memory-efficient training methods that project gradients into a low-rank subspace can be made both faster and more theoretically sound by replacing periodic exact SVD of gradients with a continuous gradient-based update of the projection matrix — and that this replacement is provably convergent for a wide class of optimizers due to a previously unnoticed property of the Hamiltonian descent framework.

The system being proposed is a memory-efficient optimizer that maintains a dynamically evolving low-rank projection matrix updated via online PCA. The problem it solves is the memory bottleneck of Adam-style optimizers during LLM pretraining (storing two full-size momentum buffers per parameter), and the "shape" of the solution is a pair of coupled optimization processes: the main weight optimizer runs in a low-dimensional projected space (saving memory), while a lightweight secondary optimizer continuously updates the projection matrix to track the evolving gradient geometry (replacing expensive periodic SVD).

3.2 Big-Picture Architecture (Diagram in Words)

The Online Subspace Descent system has four interacting components:

1. Base optimizer (OptimizerW) — This is the optimizer that updates the model weights $W$. It can be any standard optimizer (Adam, LION, Adafactor) but operates in a reduced-dimensional projected space by working with projected gradients $P_t^\top G_t$ rather than full gradients $G_t$. Its internal states $\hat{S}_t$ (momentum buffers) are therefore of size $k \times m$ rather than $n \times m$, where $k \ll n$ is the projection rank.

2. Projection matrix $P_t \in \mathbb{R}^{n \times k}$ — A tall rectangular matrix that defines the current low-rank subspace. At each iteration, it maps the full gradient $G_t \in \mathbb{R}^{n \times m}$ down to $\hat{G}_t = P_t^\top G_t \in \mathbb{R}^{k \times m}$ (the "projection" step) and maps the reduced-dimension weight update back up to full space via $W_{t+1} = W_t + P_t \cdot \text{update}$ (the "reconstruction" step).

3. Projection matrix optimizer (OptimizerP) — A separate optimizer instance that updates $P_t$ by minimizing the online PCA reconstruction loss $L_{G_t}(P)$. This can be as simple as SGD (gradient descent on the PCA objective) or as sophisticated as Adam with its own momentum states $S^P_t$. The update step for $P_t$ is designed to be fast enough to execute in parallel with the main weight update.

4. Online PCA loss function $L_{G_t}(P)$ — The objective that drives $P_t$ updates. At each step $t$, it measures how well the current projection matrix $P$ can reconstruct the normalized gradient $\tilde{G}_t$ via $PP^\top\tilde{G}_t$, plus an orthonormality penalty.

Information flows through the system in each iteration as follows:

First, the gradient $G_t = \nabla L(W_t)$ is computed from the current mini-batch. Second, this gradient is projected into the current subspace: $\hat{G}_t = P_t^\top G_t$. Third, the base optimizer OptimizerW consumes $\hat{G}_t$, updates its reduced-dimension states $\hat{S}_t$, and produces a reduced-dimension weight update $\hat{\Delta}_{t+1}$. Fourth, this update is mapped back to full space and applied: $W_{t+1} = W_t + \epsilon_t^W (P_t \hat{\Delta}_{t+1} - \lambda_W W_t)$, where weight decay is applied as a standard heuristic. Fifth, in parallel (or sequentially if resources are constrained), the projection matrix optimizer OptimizerP computes the gradient of the online PCA loss with respect to $P_t$, namely $\nabla_P L_{G_t}(P_t)$, and updates $P_{t+1} \gets P_t + \epsilon_t^P(\Delta_t^P - \lambda_P P_t)$.

A critical design choice: the projection matrix $P_t$ is never used as a static constraint on the weights (unlike LoRA, where $W = W_0 + AB$ permanently). Instead, $W_t$ is maintained in full rank, and $P_t$ only determines which coordinate subspace is active during the current update step. Over many iterations with different $P_t$, the effective search space is the full parameter space — the low-rank restriction applies per-update, not globally.

3.3 Roadmap for the Deep Dive

• First, the static subspace descent formulation (Section 3.1 in the paper): understanding the simple case where $P$ never changes, which introduces the core projection mechanism and establishes notation.
• Second, the dynamic extension to Online Subspace Descent (Section 3.2): why we need $P_t$ to change, how the online PCA loss $L_{G_t}(P)$ is defined and minimized, and how Algorithm 1 orchestrates the coupled updates.
• Third, the theoretical difficulty that motivated the analysis (Section 3.3): why dynamic projection creates a non-obvious entanglement between optimizer states and the changing subspace, and why the naive coordinate-descent intuition fails.
• Fourth, the Hamiltonian descent framework (Section 2 in the paper and Section 4): establishing the continuous-time ODE form of common optimizers, the Hamiltonian+descent structure, and why it is the right analytical lens.
• Fifth, the main theoretical result: proving that the Hamiltonian+descent structure is preserved under dynamic subspace projection due to adjoint-based cancellation, and that convergence to local optima follows under a mild non-degeneracy condition.
• Sixth, the generalization to arbitrary linear operators: showing that the framework extends beyond simple left-projection to two-sided projections and general linear operators with adjoints.

3.4 Detailed, Sentence-Based Technical Breakdown

---

Static Subspace Descent: Confining Optimization to a Fixed Low-Rank Space

The paper begins by deriving the memory-efficient optimizer for the simplest case: a fixed projection matrix $P$ that does not change during training. This establishes the core mechanics before introducing the complexity of dynamic updates.

The core idea of static subspace descent: instead of optimizing the full weight matrix $W \in \mathbb{R}^{n \times m}$ directly, impose the low-rank structure $W = P \hat{W}$, where $P \in \mathbb{R}^{n \times k}$ is a fixed projection matrix and $\hat{W} \in \mathbb{R}^{k \times m}$ is the reduced-dimension parameter to be trained. When $k \ll n$, the effective number of trainable parameters drops from $n \times m$ to $k \times m$ (for the weight) plus $n \times k$ (for $P$, which is frozen during training). The crucial memory savings come from the optimizer states: instead of maintaining full-size momentum buffers $M_t, V_t \in \mathbb{R}^{n \times m}$, we maintain reduced-dimension buffers $\hat{M}_t, \hat{V}_t \in \mathbb{R}^{k \times m}$.

Training objective under static projection: the optimization problem becomes

$\min_{\hat{W}} \; \mathcal{L}(P\hat{W})$

where $\mathcal{L}$ is the original loss function (e.g., language modeling cross-entropy). The gradient with respect to $\hat{W}$ is obtained via the chain rule:

$\nabla_{\hat{W}} \mathcal{L}(P\hat{W}) = P^\top \nabla_W \mathcal{L}(W)$

where $\nabla_W \mathcal{L}(W) \in \mathbb{R}^{n \times m}$ is the full-space gradient and $P^\top \in \mathbb{R}^{k \times n}$ projects it into the $k$-dimensional subspace.

What this computes: the full gradient $\nabla_W \mathcal{L}$ (a matrix of partial derivatives, same shape as $W$) is left-multiplied by $P^\top$, producing a $k \times m$ matrix $\nabla_{\hat{W}} \mathcal{L}$ that represents the gradient restricted to the directions spanned by the columns of $P$. This is the projection step: information about gradient components orthogonal to the subspace is discarded.

Why this form matters: the chain-rule derivation $\nabla_{\hat{W}} = P^\top \nabla_W$ is what enables the optimizer to work entirely in the reduced space. Without it, we would need to maintain $\hat{W}$-sized optimizer states AND full-sized gradient buffers, defeating the memory-efficiency purpose. The projection step $P^\top G_t$ is the only place where the full-sized gradient $G_t$ is needed; after projection, all subsequent operations (momentum updates, preconditioning, weight updates) operate on $k \times m$ matrices.

Applying a general optimizer in the static subspace. Given any optimizer that updates parameters via

$W_{t+1} = W_t + \phi_t(S_t), \quad S_t = \psi_t(S_{t-1}, \nabla\mathcal{L}(W_t))$

we can instantiate it in the projected space by substituting $\hat{W}$ for $W$ and $P^\top \nabla\mathcal{L}(W_t)$ for $\nabla\mathcal{L}(W_t)$:

$\hat{W}_{t+1} = \hat{W}_t + \phi_t(\hat{S}_t), \quad \hat{S}_t = \psi_t(\hat{S}_{t-1}, P^\top \nabla\mathcal{L}(W_t))$

Finally, map the update back to the original parameterization $W_t = P\hat{W}_t$ by left-multiplying the update by $P$:

$W_{t+1} = W_t + P \cdot \phi_t(\hat{S}_t)$

This is the key equation: the weight update in full space is $P$ times the reduced-dimension optimizer update. The optimizer states $\hat{S}_t$ track momentum of projected gradients only, so their memory footprint is $O(km)$ rather than $O(nm)$.

What makes this memory-efficient: for a weight matrix of size $n \times m$ and projection rank $k$, the full-rank Adam states require $2nm$ floats. The projected version requires $2km$ floats for $\hat{M}_t, \hat{V}_t$ plus $nk$ floats for storing $P$ (which is small when $k \ll n, m$). For a concrete example: a $4096 \times 4096$ weight matrix with $k=512$ reduces optimizer state memory from $2 \times 16.8\text{M} \approx 33.6\text{M}$ values to $2 \times 2.1\text{M} \approx 4.2\text{M}$ values — roughly an 8× reduction per weight matrix.

The fundamental limitation of static projection. With a fixed $P$, the weight $W$ is permanently constrained to live in the column space of $P$. Regardless of how many training steps are performed, $W$ can never develop components orthogonal to $\text{span}(P)$. This is adequate for fine-tuning (where the pre-trained weights already span the right directions, and only minor adjustments are needed — this is LoRA's domain), but catastrophic for pre-training, where the model must discover weight structures that are not known in advance. The static projection assumption is what the rest of the paper sets out to overcome.

---

Online Subspace Descent: Dynamically Evolving the Projection Matrix

The key innovation of Online Subspace Descent is to allow $P_t$ to change at every iteration while keeping the same per-step projection mechanics. The update rule becomes:

$W_{t+1} = W_t + P_t \cdot \phi_t(\hat{S}_t)$

$\hat{S}_t = \psi_t(\hat{S}_{t-1}, P_t^\top \nabla\mathcal{L}(W_t))$

$P_{t+1} = \chi_t(P_t, W_t, \hat{S}_t)$

where $\chi_t$ is a new update rule for the projection matrix that must be designed. The crucial difference from the static case: $W_t$ is no longer constrained to a fixed subspace because different $P_t$ at different iterations open up different projection directions. Over time, the weight matrix can develop components in any direction that appears in some $P_t$.

The intuition behind dynamic projection (and why it resembles coordinate descent). The authors draw an explicit analogy to coordinate descent in Section 3.2: "This is similar to the update of coordinate descent, except in a continuous fashion." In classical coordinate descent, each iteration picks one coordinate direction and takes a gradient step along it; over many iterations cycling through coordinates, the full-dimensional optimum is reached. Here, each iteration restricts the update to a $k$-dimensional subspace spanned by the columns of $P_t$. By rotating this subspace across iterations, the optimizer can eventually explore all directions in the full parameter space. The projection matrix $P_t$ plays the role of the "active coordinate set" at iteration $t$.

Why this is not obviously correct. The coordinate descent analogy breaks down when the optimizer has state (momentum). In coordinate descent, each step along a coordinate direction is independent: the update in coordinate $i$ does not affect the next step in coordinate $j$. But in Online Subspace Descent, the optimizer states $\hat{S}_t$ accumulate information from gradients projected through all previous $P_\tau$ for $\tau \le t$. The momentum buffer $\hat{M}_t$ at time $t$ contains a weighted sum of $P_\tau^\top G_\tau$ for various $\tau$. If $P_t$ has rotated significantly from earlier $P_\tau$, the directions stored in $\hat{M}_t$ may be incoherent with the directions that $P_t$ will use to map the update back to full space. It is not a priori obvious that this entanglement doesn't cause the optimizer to diverge or stall.

This difficulty is what the authors highlight in Section 3.3: "As $P_t$ changes across time, it is unclear whether the gradient projected to different subspaces $P_\tau$ would be coherent with each other, and useful for future updates that are conducted in different subspaces $P_t$ for $t > \tau$." The resolution of this difficulty via the Hamiltonian descent analysis is the paper's main theoretical contribution, and we will cover it in detail later. But first, we need to establish how $P_t$ is actually updated in practice.

The online PCA objective for updating $P_t$. The paper proposes to update $P_t$ by minimizing a reconstruction loss that measures how well the current projection matrix can capture the current gradient. The objective, defined in Equation 6, is:

$L_{G_t}(P) = \left\|PP^\top \tilde{G}_t - \tilde{G}_t\right\|^2 + \lambda\left\|P^\top P - I_{k \times k}\right\|^2$

where $\tilde{G}_t = G_t / \|G_t\|$ is the normalized gradient (Frobenius norm normalization), $\|\cdot\|$ denotes the Frobenius norm, $I_{k \times k}$ is the $k \times k$ identity matrix, and $\lambda > 0$ is a regularization hyperparameter.

Defining the terms:

• $PP^\top \tilde{G}_t$: the reconstruction of $\tilde{G}_t$ obtained by first projecting into the subspace ($P^\top \tilde{G}_t$) and then mapping back to the original space ($P \cdot$ (result)). This is the best rank-$k$ approximation of $\tilde{G}_t$ achievable with the current $P$.
• $\|PP^\top \tilde{G}_t - \tilde{G}_t\|^2$: the squared Frobenius norm of the reconstruction error. When $P$'s columns span the top-$k$ singular vectors of $\tilde{G}_t$, this error is the sum of squares of the remaining singular values — the PCA reconstruction loss.
• $\|P^\top P - I_{k \times k}\|^2$: the orthonormality penalty. When $P$ has orthonormal columns, $P^\top P = I_{k \times k}$ exactly, and this term vanishes. Without this penalty, the optimization might produce a $P$ with correlated or scaled columns that still achieves low reconstruction error but doesn't properly define an orthogonal projection.
• $\lambda$: controls the tradeoff between reconstruction fidelity and orthonormality. The paper sets $\lambda = 0.1$ for all experiments after determining that results are not sensitive to this value for small models.

What this objective computes, operationally: at each iteration $t$, given the current mini-batch gradient $G_t$, normalize it to unit Frobenius norm to produce $\tilde{G}_t$. Then evaluate how well the current projection matrix $P$ can reconstruct $\tilde{G}_t$ via $PP^\top\tilde{G}_t$. The first term penalizes reconstruction error (pushing $P$ to capture the dominant directions of $G_t$), and the second term penalizes deviations of $P$ from having orthonormal columns (preventing degenerate solutions where columns become parallel or scaled differently).

Why this form (and not something else): the reconstruction loss $\|PP^\top X - X\|^2$ is the standard PCA objective: minimizing it over $P$ yields the top-$k$ singular vectors of $X$ as the global minimum (Baldi & Hornik, 1989, cited in the paper). This connects the online update directly to the SVD used in GaLore — both ultimately aim to find the dominant singular vectors of the gradient. The key difference is that instead of solving this optimization exactly and periodically (via torch.svd), Online Subspace Descent takes a single gradient step per iteration, providing an approximate solution that is cheap to compute and continuously incorporates new gradient information.

The gradient normalization $\tilde{G}_t = G_t / \|G_t\|$ is an important practical detail. Without normalization, the scale of $G_t$ would change dramatically across iterations (due to learning rate schedules, loss landscape variation, gradient clipping), causing the PCA objective to be dominated by scale rather than direction during large-gradient steps and ignored during small-gradient steps. Normalization makes the objective purely about the directional structure of the gradient, which is what matters for subspace identification.

Baldi & Hornik (1989) is cited to support the claim that all stable local minima of $L_{G_t}(P)$ are global minima, and that $P$ is a global minimum if and only if $PP^\top \tilde{G}_t$ forms the optimal rank-$k$ approximation of $\tilde{G}_t$. This is a crucial property because it means there are no "bad" local minima for the optimizer to get stuck in — any first-order stationary point of the PCA loss is actually a correct solution (the top singular subspace). This makes the objective safe for online gradient-based optimization: we don't need to worry about $P_t$ converging to a suboptimal configuration.

The update rule for $P_t$: instead of minimizing $L_{G_t}(P)$ to convergence at each step (which would be as expensive as SVD), the paper performs a single optimizer step:

$P_{t+1} = \texttt{OptimizerP.step}(P_t, \nabla_P L_{G_t}(P_t))$

where $\texttt{OptimizerP.step}$ can be any standard optimizer — gradient descent, momentum, Adam, etc. The paper's experiments (Table 2) show that even SGD on $P_t$ can work well, though Adam is recommended for best results. When using Adam for $P_t$, an additional set of optimizer states $S^P_t$ (momentum and second-order buffers) is maintained for the projection matrix, but these are tiny: $P_t \in \mathbb{R}^{n \times k}$ with $k$ small means the overhead is negligible compared to the savings from not storing full-size optimizer states for $W$.

The complete algorithm (Algorithm 1). The paper presents the full procedure in Algorithm 1, which we can walk through step by step:

Step 0 (initialization): Choose a base optimizer OptimizerW with learning rate $\epsilon_t^W$ and weight decay $\lambda_W$ for the model weights $W_t$. Choose OptimizerP with learning rate $\epsilon_t^P$ and weight decay $\lambda_P$ for the projection matrix $P_t$. Initialize $W_0$, $P_0$, and optimizer states $\hat{S}_0$, $S^P_0$ appropriately. The paper recommends Adam for both optimizers, $\epsilon_t^P = \alpha \epsilon_t^W$ with $\alpha = 5$, and $\lambda_W = \lambda_P$.

Step 1 (gradient computation): Compute the full-size gradient $G_t = \nabla\mathcal{L}(W_t)$ from the current mini-batch. This is the only step that operates on full-sized tensors.

Step 2 (weight update): Project the gradient into the current subspace: $\hat{G}_t = P_t^\top G_t$. Feed $\hat{G}_t$ to OptimizerW, which updates its internal states $\hat{S}_{t-1}$ and produces a reduced-dimension update $\hat{\Delta}_{t+1}$. Apply the update to the weights:

$W_{t+1} = W_t + \epsilon_t^W (P_t \hat{\Delta}_{t+1} - \lambda_W W_t)$

The weight decay term $-\lambda_W W_t$ is applied in the full space as a standard heuristic (not projected).

Step 3 (projection matrix update, performed in parallel with Step 2 when hardware allows): Compute the gradient of the online PCA loss with respect to $P_t$:

$G^P_t = \nabla_P L_{G_t}(P_t)$

Feed $G^P_t$ to OptimizerP, which updates its internal states $S^P_{t-1}$ and produces an update $\Delta^P_t$. Update the projection matrix:

$P_{t+1} = P_t + \epsilon_t^P (\Delta^P_t - \lambda_P P_t)$

What happens in each component and why it matters:

• $\hat{G}_t = P_t^\top G_t$: this single matrix multiplication reduces a $n \times m$ gradient to $k \times m$. All subsequent computations in OptimizerW operate on $k \times m$ tensors. If $n = 4096$ and $k = 512$, this is an 8× reduction in the dimension along which momentum and second-order statistics are tracked.

• $P_t \hat{\Delta}_{t+1}$: this maps the reduced-dimension update back to full space. The update $\hat{\Delta}_{t+1}$ encodes how much to move in each of the $k$ basis directions defined by the columns of $P_t$. Multiplying by $P_t$ computes the linear combination of these basis vectors.

• Why weight decay is $W_{t+1} = W_t + \epsilon_t^W P_t\hat{\Delta}_{t+1} - \epsilon_t^W\lambda_W W_t$ rather than projected: weight decay is meant to pull all weights toward zero uniformly, which is a full-space operation. Projecting it through $P_t$ would restrict weight decay to act only in the current subspace, causing weights in orthogonal directions to drift indefinitely.

The recommended hyperparameters (from the Algorithm 1 remark and Section 5.3):

• Both OptimizerW and OptimizerP should use Adam by default.
• $\epsilon_t^P = \alpha \epsilon_t^W$ with $\alpha = 5$ (the projection matrix learns 5× faster than the weights). This is because the projection matrix needs to track the changing gradient geometry, and using a higher learning rate allows it to adapt quickly.
• $\lambda_W = \lambda_P$ (equal weight decay for both).
• $\lambda = 0.1$ in the PCA loss (the orthonormality regularization; found empirically insensitive for small models).
• For larger models (350M, 1B), the base learning rate $\epsilon_t^W$ should be 0.001 — 10× smaller than what might be used for small models — because "larger learning rates cause unrecoverable spikes and instability" (Section 5.3).
• Rank $k$ should be set as high as memory allows until perplexity reduction saturates (Section 5.2, Figure 3 left).

---

Example Instantiations with Common Optimizers

The paper provides explicit update rules for how common optimizers look when equipped with Online Subspace Descent (Example 3.1) and their continuous-time forms (Examples 4.1-4.3). Understanding these instantiations is essential because they reveal what happens to the optimizer's internal mechanics under projection.

Projected Gradient Descent (the simplest case):

$W_{t+1} = W_t - \epsilon_t P_t P_t^\top G_t$

The term $P_t P_t^\top G_t$ is the projection of the gradient $G_t$ onto the column space of $P_t$. This is equivalent to taking a gradient step restricted to the $k$-dimensional subspace. If $P_t$ has orthonormal columns (as encouraged by the PCA loss), then $P_t P_t^\top$ is exactly the orthogonal projection matrix onto $\text{span}(P_t)$. The effective learning rate in each subspace direction is uniform.

What this computes: multiply $G_t$ by $P_t^\top$ (project into $k$ dimensions), then by $P_t$ (map back to $n$ dimensions). The result is the component of $G_t$ that lies in the subspace — the orthogonal component is zeroed out.

Why this form: the composition $P_t P_t^\top$ is a rank-$k$ symmetric positive semidefinite matrix that serves as a low-rank preconditioner. In the limit where $P_t$'s columns are the top-$k$ eigenvectors of the Hessian, this becomes a low-rank approximation of Newton's method restricted to the high-curvature directions.

Projected Momentum:

$W_{t+1} = W_t - \epsilon_t P_t \hat{M}_t, \quad \hat{M}_t = (1 - \beta) P_t^\top G_t + \beta \hat{M}_{t-1}$

The momentum buffer $\hat{M}_t \in \mathbb{R}^{k \times m}$ tracks an exponential moving average of projected gradients. The key subtlety: $\hat{M}_t$ at step $t$ contains contributions from $P_\tau^\top G_\tau$ for all $\tau \le t$, each projected through a potentially different $P_\tau$. When this momentum is mapped back to full space via $P_t$, the effective update is $P_t \hat{M}_t$, which re-interprets the accumulated momentum through the current subspace lens. This is where the potential for incoherence arises (and where the Hamiltonian analysis later proves that it doesn't cause problems).

Projected Adam:

$W_{t+1} = W_t - \epsilon_t P_t \frac{\hat{M}_t}{\sqrt{\hat{V}_t} + e}$

$\hat{M}_t = (1 - \beta_{1t}) P_t^\top G_t + \beta_{1t} \hat{M}_{t-1}$

$\hat{V}_t = (1 - \beta_{2t}) (P_t^\top G_t)^{\odot 2} + \beta_{2t} \hat{V}_{t-1}$

where $(P_t^\top G_t)^{\odot 2}$ denotes elementwise squaring, $e$ is a small constant for numerical stability, and $\beta_{1t}, \beta_{2t}$ are debiased momentum coefficients. Both the first-order momentum $\hat{M}_t$ and second-order momentum $\hat{V}_t$ are maintained in the $k \times m$ projected space — this is where the memory savings come from. The adaptive learning rate $1 / (\sqrt{\hat{V}_t} + e)$ scales each dimension independently based on historical gradient magnitudes, but crucially, these historical magnitudes were computed in the projected spaces at their respective times, not in the current space.

Projected LION-K:

$W_{t+1} = W_t - \epsilon_t P_t \nabla K(\hat{N}_t)$

$\hat{N}_t = (1 - \beta_1) P_t^\top G_t + \beta_1 \hat{M}_t$

$\hat{M}_t = (1 - \beta_2) P_t^\top G_t + \beta_2 \hat{M}_{t-1}$

where $K$ is a convex function with $\nabla K(0) = 0$ (for standard LION, $K(X) = \|X\|_{1,1}$ and $\nabla K(X) = \text{sign}(X)$). The sign operation is applied elementwise in the projected space, and the result is mapped back to full space via $P_t$.

---

The Theoretical Difficulty That Motivated the Analysis

The paper's Section 3.3 is devoted to explaining why analyzing the convergence of dynamic subspace optimizers is hard. This section is important because it justifies the need for the Hamiltonian descent framework and explains why earlier convergence proofs (e.g., Stochastic Subspace Descent on convex functions) don't extend.

The naive coordinate-descent argument (and why it fails for stateful optimizers). For projected gradient descent (no momentum), the update is $W_{t+1} = W_t - \epsilon_t P_t P_t^\top G_t$. Since $P_t P_t^\top$ is positive semidefinite, the update direction $-P_t P_t^\top G_t$ is always non-increasing in $\mathcal{L}(W)$ for any $P_t$:

$\langle -P_t P_t^\top G_t, G_t \rangle = -\|P_t^\top G_t\|^2 \le 0$

The algorithm is essentially a variant of coordinate descent where $P_t$ defines the active subspace. As long as $P_t^\top G_t = 0$ implies $G_t = 0$ (a non-degeneracy condition that ensures the optimizer doesn't stop early), convergence to a stationary point is straightforward. This is the intuition the authors invoke: the projection matrix just selects which coordinates to descend along at each step, and as long as all directions are eventually covered, we converge.

The momentum entanglement problem. This intuition breaks when the optimizer has state. Consider projected momentum: $\hat{M}_t = (1-\beta) P_t^\top G_t + \beta \hat{M}_{t-1}$. The momentum buffer $\hat{M}_t$ is a weighted sum:

$\hat{M}_t = \sum_{\tau=0}^t \beta^{t-\tau} (1-\beta) P_\tau^\top G_\tau$

When this is mapped back to full space, the update is:

$W_{t+1} = W_t - \epsilon_t P_t \sum_{\tau=0}^t \beta^{t-\tau} (1-\beta) P_\tau^\top G_\tau$

The effective update direction at time $t$ is a linear combination of $P_t P_\tau^\top G_\tau$ for all past $\tau$. If $P_t$ has rotated relative to $P_\tau$, the matrix $P_t P_\tau^\top$ is not the identity — it's a rank-$k$ transformation that maps directions from the old subspace into the current one. These cross-subspace terms have no obvious monotonicity properties, and it's possible in principle for them to point in directions that increase the loss, especially if early gradients from a subspace that no longer matters continue to influence the momentum buffer.

The paper states this difficulty explicitly: "the inertia effect of $S_t$ entangles the different subspaces, making the dynamic behavior fundamentally more complicated than naive coordinate descent where the descent in different subspaces is uncoupled." The word "inertia" is carefully chosen — momentum introduces a memory effect where past gradient information persists in the state, and when the coordinate system (subspace) rotates, this inertia can push in directions that are no longer relevant or even harmful.

Why this is not just a theoretical concern. In practice, if the entanglement caused problems, we would expect Online Subspace Descent to either diverge or plateau at suboptimal points when using momentum-based optimizers like Adam. The paper's experimental results show the opposite — Online Subspace Descent with Adam performs very well. The theoretical analysis is needed to explain why this works despite the intuitive concerns, and to provide conditions under which it is guaranteed to work. This is the gap that Section 4 fills.

---

The Hamiltonian Descent Framework (Background for the Analysis)

The paper's theoretical analysis (Section 4) builds on the Hamiltonian descent framework introduced by Maddison et al. (2018) and further developed by Chen et al. (2023) for analyzing LION and other optimizers. We need to understand this framework to follow the paper's main theoretical result.

The central observation: in the limit of infinitesimal step size (continuous time), many common optimizers can be written as a system of ODEs with a Hamiltonian+descent structure:

$\frac{d}{dt} W_t = \partial_S H(W_t, S_t) - \Phi(\partial_W H(W_t, S_t))$

$\frac{d}{dt} S_t = -\partial_W H(W_t, S_t) - \Psi(\partial_S H(W_t, S_t))$

where:

• $H(W, S)$ is a Hamiltonian (or Lyapunov) function that satisfies $\min_S H(W, S) = \mathcal{L}(W)$ for all $W$. This property means minimizing $\mathcal{L}(W)$ reduces to minimizing the augmented function $H(W, S)$.
• $\partial_W H$ is the partial derivative of $H$ with respect to $W$ (a matrix, same shape as $W$).
• $\partial_S H$ is the partial derivative of $H$ with respect to the optimizer state $S$.
• $\Phi(\cdot)$ and $\Psi(\cdot)$ are monotone mappings satisfying $\langle X, \Phi(X) \rangle \ge 0$ and $\langle X, \Psi(X) \rangle \ge 0$ for all $X$. The notation $\|X\|^2_\Phi = \langle X, \Phi(X) \rangle$ is used.

What the Hamiltonian+descent structure guarantees. The key property, derived in Equation 3 of the paper, is that $H(W_t, S_t)$ is monotonically non-increasing along trajectories:

$\frac{d}{dt} H(W_t, S_t) = \langle \partial_W H_t, \frac{d}{dt} W_t \rangle + \langle \partial_S H_t, \frac{d}{dt} S_t \rangle$

Substituting the ODE forms from above:

$= \langle \partial_W H_t, \partial_S H_t - \Phi(\partial_W H_t) \rangle + \langle \partial_S H_t, -\partial_W H_t - \Psi(\partial_S H_t) \rangle$

Expanding:

$= \langle \partial_W H_t, \partial_S H_t \rangle - \langle \partial_W H_t, \Phi(\partial_W H_t) \rangle - \langle \partial_S H_t, \partial_W H_t \rangle - \langle \partial_S H_t, \Psi(\partial_S H_t) \rangle$

The cross terms $\langle \partial_W H_t, \partial_S H_t \rangle$ and $-\langle \partial_S H_t, \partial_W H_t \rangle$ cancel exactly (because the inner product is symmetric: $\langle A, B \rangle = \langle B, A \rangle$ for real matrices). What remains is:

$= -\langle \partial_W H_t, \Phi(\partial_W H_t) \rangle - \langle \partial_S H_t, \Psi(\partial_S H_t) \rangle = -\|\partial_W H_t\|^2_\Phi - \|\partial_S H_t\|^2_\Psi \le 0$

The inequality holds because $\Phi$ and $\Psi$ are monotone ($\langle X, \Phi(X) \rangle \ge 0$ for all $X$).

Why the cancellation of cross terms is the critical mechanism. In a pure Hamiltonian system ($\Phi = \Psi = 0$), $H$ is conserved — the trajectory preserves the value of $H$. Adding the monotone "descent" terms $\Phi$ and $\Psi$ introduces dissipation: energy flows out of the system, causing $H$ to decrease. The fact that the cross terms $\langle \partial_W H, \partial_S H \rangle$ cancel is not an accident — it's a structural property of the Hamiltonian formalism: the symplectic structure of the Hamiltonian part ensures that these terms sum to zero (this is essentially the statement that Hamiltonian flow preserves the Hamiltonian). The dissipation components then provide a net negative drift.

Why this framework covers Adam, LION, and momentum. The paper provides explicit continuous-time forms and Hamiltonians for these optimizers:

Momentum (Example 2.2):

$\frac{d}{dt} W_t = -M_t, \quad \frac{d}{dt} M_t = a(\nabla\mathcal{L}(W_t) - M_t)$

$H(W, M) = \mathcal{L}(W) + \frac{\|M\|^2}{2a}$

where $a > 0$ is a damping coefficient. This satisfies $\min_M H(W, M) = \mathcal{L}(W)$ (achieved at $M = 0$). The descent follows because:

$\frac{d}{dt} H = \langle \nabla\mathcal{L}(W_t), -M_t \rangle + \frac{1}{a} \langle M_t, a(\nabla\mathcal{L}(W_t) - M_t) \rangle = -\|M_t\|^2 \le 0$

Adam (Example 2.3):

$\frac{d}{dt} W_t = -\frac{M_t}{\sqrt{V_t} + e}, \quad \frac{d}{dt} M_t = a(\nabla\mathcal{L}(W_t) - M_t), \quad \frac{d}{dt} V_t = b(\nabla\mathcal{L}(W_t)^{\odot 2} - V_t)$

$H(W, M, V) = \mathcal{L}(W) + \frac{1}{2a} \left\langle \frac{M}{\sqrt{V} + e}, M \right\rangle$

The condition $a \ge b/4$ ensures $\frac{d}{dt} H \le 0$.

LION-K (Example 2.4):

$\frac{d}{dt} W_t = \nabla K((1-b)M_t - b\nabla\mathcal{L}(W_t)), \quad \frac{d}{dt} M_t = -a(\nabla\mathcal{L}(W_t) + M_t)$

$H(W, M) = a\mathcal{L}(W) + \frac{1}{1-b} K((1-b)M)$

where $K$ is convex with $\nabla K(0) = 0$.

What unifies these cases: in each, $\partial_W H$ is related to the gradient $\nabla\mathcal{L}(W)$ (possibly scaled), and $\partial_S H$ is related to the optimizer state $S$. The Hamiltonian function is constructed so that $\min_S H(W, S) = \mathcal{L}(W)$, meaning the augmented system's energy minimum corresponds exactly to the loss minimum. The descent property $\frac{d}{dt}H \le 0$ ensures that the optimizer monotonically reduces this augmented energy, which forces the loss to also decrease (since $\mathcal{L}(W_t) \le H(W_t, S_t)$ and $H$ is decreasing).

---

The Main Theoretical Result: Hamiltonian+Descent Is Preserved Under Subspace Projection

Section 4.1 contains the paper's central theoretical insight. The claim is that when Online Subspace Descent is applied to any optimizer with the Hamiltonian+descent structure, the same structure is preserved — with the same Lyapunov function $H$ — regardless of how $P_t$ evolves.

The projected continuous-time system. Applying dynamic projection to the Hamiltonian descent ODE (Equation 2 in the paper) yields:

$\frac{d}{dt} W_t = P_t \partial_{\hat{S}} H(W_t, \hat{S}_t) - \Phi(\partial_W H(W_t, \hat{S}_t))$

$\frac{d}{dt} \hat{S}_t = -P_t^\top \partial_W H(W_t, \hat{S}_t) - \Psi(\partial_{\hat{S}} H(W_t, \hat{S}_t))$

$\frac{d}{dt} P_t = \Gamma(P_t, \nabla\mathcal{L}(W_t))$

where $\Gamma$ is an arbitrary update rule for $P_t$ (not specified beyond being sufficiently smooth that the derivatives exist). The key modifications from the unprojected system are:

• In the $W$ equation, $\partial_{\hat{S}} H$ is multiplied by $P_t$ on the left.
• In the $\hat{S}$ equation, $\partial_W H$ is multiplied by $P_t^\top$ on the left.
• The original state $S$ is replaced by the reduced-dimension state $\hat{S}$, and the Hamiltonian is defined on $(W, \hat{S})$.

The Lyapunov derivation. The time derivative of $H(W_t, \hat{S}_t)$ is:

$\frac{d}{dt} H(W_t, \hat{S}_t) = \langle \partial_W H_t, \frac{d}{dt} W_t \rangle + \langle \partial_{\hat{S}} H_t, \frac{d}{dt} \hat{S}_t \rangle$

Substituting the ODE forms:

$= \langle \partial_W H_t, P_t \partial_{\hat{S}} H_t - \Phi(\partial_W H_t) \rangle + \langle \partial_{\hat{S}} H_t, -P_t^\top \partial_W H_t - \Psi(\partial_{\hat{S}} H_t) \rangle$

Expanding:

$= \langle \partial_W H_t, P_t \partial_{\hat{S}} H_t \rangle - \|\partial_W H_t\|^2_\Phi + \langle \partial_{\hat{S}} H_t, -P_t^\top \partial_W H_t \rangle - \|\partial_{\hat{S}} H_t\|^2_\Psi$

Now the critical step: the cross terms are $\langle \partial_W H_t, P_t \partial_{\hat{S}} H_t \rangle$ and $\langle \partial_{\hat{S}} H_t, -P_t^\top \partial_W H_t \rangle$. Using the adjoint property $\langle X, P_t Y \rangle = \langle P_t^\top X, Y \rangle$, we have:

$\langle \partial_W H_t, P_t \partial_{\hat{S}} H_t \rangle = \langle P_t^\top \partial_W H_t, \partial_{\hat{S}} H_t \rangle$

The second cross term is $\langle \partial_{\hat{S}} H_t, -P_t^\top \partial_W H_t \rangle$. Since the inner product is symmetric:

$\langle \partial_{\hat{S}} H_t, -P_t^\top \partial_W H_t \rangle = -\langle P_t^\top \partial_W H_t, \partial_{\hat{S}} H_t \rangle$

These two are exact negatives and cancel, leaving:

$\frac{d}{dt} H(W_t, \hat{S}_t) = -\|\partial_W H_t\|^2_\Phi - \|\partial_{\hat{S}} H_t\|^2_\Psi \le 0$

What this derivation actually shows — the deep reason it works. The cancellation is completely independent of $P_t$ and its evolution rule $\Gamma$. It does not rely on $P_t$ being any particular matrix; it only relies on the algebraic identity $\langle A, P B \rangle = \langle P^\top A, B \rangle$, which holds for any real matrix $P$. The adjoint property ensures that the projection appearing in the $W$ equation ($P_t$) and the projection appearing in the $\hat{S}$ equation ($P_t^\top$) are symmetrically paired, causing their contributions to the energy derivative to cancel exactly.

Why this is surprising (as the paper notes). Intuition would suggest that an arbitrary dynamically changing projection matrix could interfere with the descent dynamics — after all, $P_t$ is potentially rotating, scaling, or even having rank collapse over time. The derivation shows that this intuition is wrong for the Hamiltonian+descent class: the adjoint symmetry is a structural invariant that survives any smooth time-variation of $P_t$. The projection mechanism is effectively "invisible" to the Lyapunov function because it appears symmetrically in the two equations that govern the descent.

What the result does NOT say. The Lyapunov property $\frac{d}{dt} H \le 0$ guarantees that the augmented energy $H(W_t, \hat{S}_t)$ never increases, but it does not guarantee that the rate of decrease is fast, or that the system converges to the global minimum of $\mathcal{L}(W)$ rather than some other point where $\frac{d}{dt}H = 0$. The Lyapunov analysis only provides stability — it ensures the optimizer doesn't diverge — but additional conditions (Section 4.2) are needed to ensure convergence to stationary points of $\mathcal{L}$.

How $\Gamma$ (the $P_t$ update rule) enters the picture. The paper notes (Section 4.1): "In the examples above, although the form of the Hamiltonian $H(W, \hat{S})$ is independent of the update rule of $P_t$, the decreasing rate $\frac{d}{dt} H(W_t, \hat{S}_t)$ depends on $P_t$ in a complicated way through $\hat{M}_t, \hat{V}_t, \hat{G}_t$." The rate of descent — how fast $H$ decreases — depends on $P_t$ through the projected gradient $\hat{G}_t = P_t^\top G_t$ and the projected states $\hat{M}_t, \hat{V}_t$. A good update rule $\Gamma$ that keeps $P_t$ aligned with the dominant gradient directions will make $\hat{G}_t$ capture more of the true gradient, leading to faster descent. A poor choice of $\Gamma$ (e.g., a random $P_t$) would still guarantee $\frac{d}{dt}H \le 0$ but might be very slow.

Worked examples: verifying the cancellation for specific optimizers. The paper provides explicit verifications for momentum, Adam, and LION-K to demonstrate that the adjoint cancellation works in each case.

Momentum + Online Subspace Descent (Example 4.1):

$\frac{d}{dt} W_t = -P_t \hat{M}_t, \quad \frac{d}{dt} \hat{M}_t = a(P_t^\top \nabla\mathcal{L}(W_t) - \hat{M}_t), \quad H(W, \hat{M}) = \mathcal{L}(W) + \frac{\|\hat{M}\|^2}{2a}$

Computing the derivative directly:

$\frac{d}{dt} H = \langle \nabla\mathcal{L}(W_t), -P_t \hat{M}_t \rangle + \frac{1}{a} \langle \hat{M}_t, a(P_t^\top \nabla\mathcal{L}(W_t) - \hat{M}_t) \rangle$

$= -\langle \nabla\mathcal{L}(W_t), P_t \hat{M}_t \rangle + \langle \hat{M}_t, P_t^\top \nabla\mathcal{L}(W_t) \rangle - \|\hat{M}_t\|^2$

Use the adjoint property on the second term: $\langle \hat{M}_t, P_t^\top \nabla\mathcal{L}(W_t) \rangle = \langle P_t \hat{M}_t, \nabla\mathcal{L}(W_t) \rangle = \langle \nabla\mathcal{L}(W_t), P_t \hat{M}_t \rangle$. This exactly cancels the first term, leaving $-\|\hat{M}_t\|^2 \le 0$.

Adam + Online Subspace Descent (Example 4.2): The derivation is more involved due to the second-order moment $\hat{V}_t$. The Hamiltonian is:

$H(W, \hat{M}, \hat{V}) = \mathcal{L}(W) + \frac{1}{2a} \left\langle \frac{\hat{M}}{\sqrt{\hat{V}} + e}, \hat{M} \right\rangle$

The time derivative computation yields:

$\frac{d}{dt} H = -\left\langle \left(1 - \frac{b}{4a}\right) \frac{\sqrt{\hat{V}_t}}{\sqrt{\hat{V}_t} + e}, \frac{\hat{M}_t^{\odot 2}}{\sqrt{\hat{V}_t} + e} \right\rangle - \frac{b}{4a} \left\| \frac{\hat{M}_t \hat{G}_t}{4\sqrt{\hat{V}_t}(\sqrt{\hat{V}_t} + e)} \right\|^2$

The first term is non-positive when $a \ge b/4$ (the standard Adam condition), and the second term is always non-positive. The cross terms involving $P_t$ cancel via the adjoint property in the same way as for momentum. The paper's derivation shows that the subspace projection introduces no new positivity requirements — the same condition $a \ge b/4$ that guarantees descent for full-rank Adam also guarantees descent for projected Adam.

LION-K + Online Subspace Descent (Example 4.3): The derivation uses the convexity of $K$ and the bracket notation $[X; Y]_{\nabla K} = \langle Y, \nabla K(X+Y) - \nabla K(X) \rangle$, which is non-negative for convex $K$. The cross terms again cancel via $\langle G_t, P_t X_t \rangle = \langle P_t^\top G_t, X_t \rangle = \langle \hat{G}_t, X_t \rangle$.

The generalization to arbitrary linear operators (Section 4.3). The paper extends the analysis beyond simple left-projection $P_t$ to general linear operators $\mathcal{P}_t$ with adjoints $\mathcal{P}_t^*$. The system becomes:

$\frac{d}{dt} W_t = \mathcal{P}_t(\partial_{\hat{S}} H) - \Phi(\partial_W H)$

$\frac{d}{dt} \hat{S}_t = -\mathcal{P}_t^*(\partial_W H) - \Psi(\partial_{\hat{S}} H)$

The cancellation now uses $\langle X, \mathcal{P}_t(Y) \rangle = \langle \mathcal{P}_t^*(X), Y \rangle$, which is the definition of the adjoint operator. This generalization covers two-sided projections: if $\mathcal{P}_t(X) = P_t X Q_t$ where $P_t$ projects rows and $Q_t$ projects columns, then the adjoint is $\mathcal{P}_t^*(X) = P_t^\top X Q_t^\top$. This is exactly the two-sided projection used in GaLore (Zhao et al., 2024) for additional memory savings. Other examples include $\mathcal{P}_t(X) = P_t X + X Q_t$, with adjoint $\mathcal{P}_t^*(X) = P_t^\top X + X Q_t^\top$.

What the generalization proves: the Hamiltonian+descent structure is preserved under any linear projection operator, not just left-projection, as long as the adjoint appears symmetrically in the two ODEs. This means the framework applies to a much wider class of memory-efficient methods than just the specific Online Subspace Descent algorithm — it provides a theoretical umbrella for any method that projects gradients through a linear operator and maps updates back through the same operator.

---

Convergence to Local Optima: The Non-Degeneracy Condition

The Lyapunov property $\frac{d}{dt} H \le 0$ ensures stability but does not guarantee convergence to a stationary point of $\mathcal{L}(W)$. The system could, in principle, converge to a point where $\frac{d}{dt}H = 0$ but $\nabla\mathcal{L}(W) \neq 0$ — a "stuck" point where the optimizer has stopped despite not being at a local minimum. Theorem 4.5 provides conditions that rule this out.

Assumption 4.4 (non-degeneracy conditions):

1. $\frac{d}{dt} H(W_t, \hat{S}_t) = 0$ implies $\hat{G}_t = P_t^\top \nabla\mathcal{L}(W_t) = 0$ and $\frac{d}{dt} W_t = 0$.
2. When $G_t \equiv G \neq 0$ (the gradient is constant and non-zero), the set $\{P : P^\top G = 0\}$ is not a positive invariant set of $\frac{d}{dt} P_t = \Gamma(P_t, G_t)$.

What condition (i) means: the optimizer should only stop ($\frac{d}{dt}H = 0$) when the projected gradient is zero. This is easy to verify for momentum, Adam, and LION-K — in each case, $\frac{d}{dt}H = 0$ forces the momentum term $\hat{M}_t$ to be zero, which via the momentum ODE forces $\hat{G}_t = P_t^\top G_t = 0$. Condition (i) essentially says: the optimizer's internal mechanics are such that energy stops decreasing only when there's no gradient signal in the current subspace.

What condition (ii) means: if the true gradient $G$ is non-zero but $P_t$ happens to be in a configuration where $P_t^\top G = 0$ (i.e., the gradient is orthogonal to every column of $P_t$), the update rule $\Gamma$ for $P_t$ must not keep it stuck in that configuration forever. The system must eventually escape, bringing $P_t$ to a configuration where $P_t^\top G \neq 0$, allowing descent to resume. This is satisfied when $P_t$ is updated by a reasonable optimizer of the online PCA loss $L_{G_t}(P)$, because if $P_t^\top G = 0$, then $PP^\top G = 0$, which means $P$ is doing a terrible job of reconstructing $G$ — the PCA loss will be large, and gradient descent on the PCA loss will move $P$ toward the singular vectors of $G$, breaking the orthogonality.

Theorem 4.5 (Convergence to stationary points): Let $(W_t, \hat{S}_t, P_t)_t$ be a bounded solution of the continuous-time system (7). If Assumption 4.4 holds, then all accumulation points of $\{W_t\}$ as $t \to \infty$ are stationary points of $\mathcal{L}(W)$ — that is, $\nabla\mathcal{L}(W) = 0$ at all limit points.

Proof sketch (via LaSalle's invariance principle): LaSalle's invariance principle states that for a dynamical system with a Lyapunov function $H$, all bounded trajectories converge to the largest invariant set $\mathcal{I}$ contained in the set where $\frac{d}{dt}H = 0$. From condition (i), any trajectory in $\mathcal{I}$ must have $\frac{d}{dt}W_t = 0$ (so $W_t$ is constant) and $\hat{G}_t = P_t^\top \nabla\mathcal{L}(W_t) = 0$. Since $W_t$ is constant, $G_t = \nabla\mathcal{L}(W_t)$ is constant — call it $G$. The trajectory satisfies $P_t^\top G = 0$ at all times. But this means $P_t$ is in the set $\{P : P^\top G = 0\}$ for all $t$, which contradicts condition (ii) unless $G = 0$. Therefore, any trajectory in the invariant set must have $\nabla\mathcal{L}(W) = 0$, proving that all limit points are stationary.

What this theorem provides that prior analyses didn't. Before this work, no convergence guarantee existed for any method that dynamically changes the projection matrix for non-convex objectives with stateful optimizers. The theorem fills this gap and does so with remarkably weak conditions — $\Gamma$ can be essentially any smooth update rule, not just the specific online PCA update proposed in the paper. This is what the paper means by "the first convergence guarantee for arbitrary update rules of the projection matrix" (Section 1).

The practical import: the theorem tells practitioners that they can experiment with different $P_t$ update rules (online PCA with SGD, online PCA with Adam, random perturbations around an SVD initialization, etc.) and be guaranteed that the optimizer will not diverge due to the projection mechanism, as long as the non-degeneracy condition is met. The online PCA update with an optimizer on $L_{G_t}(P)$ is one concrete choice that provably satisfies the condition, but it's not the only possible one.

---

Putting It All Together: Why Online Subspace Descent Works in Practice

The theoretical analysis (Sections 4.1–4.2) explains why the algorithm doesn't break, but the practical performance (Sections 5.1–5.5) depends on the specific design choices made in Algorithm 1. Let's connect the theory to the implementation:

Choice 1: online PCA rather than periodic SVD. GaLore's periodic SVD computes the exact top-$k$ singular subspace of the gradient at the recomputation point. Online PCA computes an approximate, continuously updated version. The theory shows that any $P_t$ is acceptable, so the advantage of online PCA is purely practical: it's much faster (142× on some tensor shapes, per Figure 2) and incorporates gradient information from all mini-batches rather than just the recomputation point. The staleness that GaLore suffers between SVD recomputations (where $P_t$ is stuck reflecting old gradient geometry) is eliminated by continuous updates.

Choice 2: using the PCA reconstruction loss $L_{G_t}(P)$ specifically. This loss has the property that its global minima are the top singular vectors of $G_t$ (Baldi & Hornik, 1989). So while online PCA only takes a single gradient step per iteration, the objective it's descending has the same optima as the exact SVD. Over many iterations with slowly changing $G_t$, $P_t$ tracks the dominant gradient subspace closely — not exactly matching the SVD of the current $G_t$, but staying close enough to provide a useful projection.

Choice 3: running OptimizerP in parallel with OptimizerW. The online PCA update step is fast (a single backward pass on a small loss function), making it schedulable to run concurrently with the weight update step. This hides its cost entirely, so the overhead of online subspace descent is effectively the cost of maintaining the $P_t$ optimizer state (negligible) plus the projection matrix multiplications $P_t^\top G_t$ (an $O(nmk)$ matrix multiplication, required by all low-rank methods including GaLore).

Choice 4: Adam as the default optimizer for both $W$ and $P$. The theory guarantees convergence for any optimizer with the Hamiltonian+descent structure, but Adam provides the best practical convergence on language modeling tasks due to its adaptive per-dimension learning rates. Using $\alpha = 5$ (projection matrix learns 5× faster) reflects the intuition that $P_t$ should adapt quickly to the changing gradient landscape while $W_t$ should change more conservatively.

The remaining open question from the theoretical analysis (Section 4.1): the paper notes that "an interesting direction for future investigation is to find optimal rules of $P_t$ to maximize the decreasing rate as an optimal control problem." The current online PCA rule is heuristic — it's fast and empirically effective, but the theory doesn't tell us whether it's optimal in any sense. There may exist better $P_t$ update rules that achieve faster convergence by more intelligently choosing which subspace to descend in at each step, and the Hamiltonian descent framework provides the analytical machinery to study this question.

4. Key Insights and Innovations

Innovation 1: The Hamiltonian+Descent Structure Is Invariant Under Dynamic Subspace Projection — a Previously Unnoticed Cancellation

This is the paper's deepest conceptual contribution, and it is fundamentally a theoretical discovery, not an algorithm. Before this work, the dominant assumption in the field was that dynamically changing the projection matrix in a subspace descent optimizer creates an analytically intractable entanglement between the optimizer's momentum state and the evolving coordinate system. The intuition, which the paper articulates explicitly in Section 3.3, was that momentum accumulates gradient information projected through old subspaces, and when mapped back through a new subspace (via P_t M_t where M_t contains contributions from P_\tau^\top G_\tau for \tau \le t), the resulting cross-subspace terms P_t P_\tau^\top G_\tau could point in directions that increase the loss. This intuition was reasonable enough that no prior work had attempted a general convergence analysis — Stochastic Subspace Descent (Kozak et al., 2019) restricted itself to convex objectives with random (not learned) projection matrices, and GaLore (Zhao et al., 2024) provided no convergence analysis at all.

What the paper discovers is that this intuition is wrong for the Hamiltonian+descent family of optimizers (Adam, LION, momentum, and others). The derivation in Section 4.1 shows that when the continuous-time ODE of any such optimizer is projected through a time-varying P_t, the cross terms that would entangle the subspaces cancel exactly due to the adjoint property ⟨X, P_t Y⟩ = ⟨P_t^\top X, Y⟩. The cancellation does not require P_t to be orthonormal, low-rank, or slowly-varying — it holds for any real matrix P_t at each instant, and it is completely independent of the update rule \Gamma that governs how P_t evolves. The Lyapunov function H(W, \hat{S}) that guarantees descent for the full-rank optimizer is exactly the same Lyapunov function for the projected version, and its time derivative d/dt H(W_t, \hat{S}_t) \le 0 is obtained through the same cancellation of cross terms.

Why is this surprising enough to qualify as a discovery rather than an observation? Because the adjoint-based cancellation is not an algebraic coincidence — it reflects a structural feature of the Hamiltonian formalism that the field had not connected to memory-efficient optimization. In a pure Hamiltonian system, the symplectic structure guarantees that cross terms sum to zero (the Hamiltonian is conserved along the flow). When dissipation terms \Phi and \Psi are added to create the Hamiltonian+descent structure (Maddison et al., 2018), the cross terms continue to cancel, and the dissipation terms provide a net negative drift. The paper's key move is to recognize that left-multiplying the state derivative by P_t in the weight equation and left-multiplying the weight derivative by P_t^\top in the state equation preserves this structure because the adjoint pairing (P_t, P_t^\top) respects the same symmetry as the original Hamiltonian system. The projection operator is effectively invisible to the Lyapunov argument.

The significance of this finding extends well beyond the specific Online Subspace Descent algorithm. It provides a theoretical umbrella for any memory-efficient method that projects gradients through a linear operator and maps updates back through the adjoint — the generalization to arbitrary linear operators \mathcal{P}_t with adjoints \mathcal{P}_t^* (Section 4.3) shows that the result applies to two-sided projections (\mathcal{P}_t(X) = P_t X Q_t), multi-matrix projections, and other linear schemes the field may invent in the future. This transforms subspace descent from a heuristic trick (motivated by the empirical observation that gradients are low-rank) into a provably sound optimization framework with convergence guarantees — the status upgrade that the paper's title emphasizes by calling it a "new family of subspace descent optimizer."

The practical consequence, which the paper exploits, is that practitioners can experiment with any smooth P_t update rule — online PCA, periodic SVD, random perturbations, learned neural networks predicting P_t — and be guaranteed that the projection mechanism itself will not cause divergence. The only additional requirement for convergence to stationary points (Theorem 4.5) is a mild non-degeneracy condition: P_t must not become permanently stuck in a configuration where P_t^\top G = 0 while G \neq 0. This condition is satisfied by any reasonable PCA-like update that seeks to reconstruct the gradient, and it rules out the trivial failure mode where the optimizer stops because the gradient is orthogonal to the subspace rather than because the gradient is zero.

This is a fundamental rather than incremental advance. The paper did not derive a tighter bound or extend an existing proof to a new setting — it identified a previously unnoticed structural invariant that explains why an entire class of methods works despite compelling intuitive reasons to expect failure. The field's prior understanding was that dynamic projection was risky and analytically opaque; the paper's understanding is that it is benign and provably convergent under remarkably weak conditions.

Innovation 2: Reframing the Memory-Efficient Training Problem from "Find a Good Static Low-Rank Approximation" to "Continuously Track the Evolving Gradient Geometry"

The second conceptual move is a reframing of the design problem for memory-efficient optimizers, and it has both practical and theoretical dimensions. The dominant paradigm before this work, exemplified by GaLore (Zhao et al., 2024) and Sketchy (Feinberg et al., 2024), was: at periodic intervals, compute the best low-rank approximation of the gradient (via SVD or a sketching algorithm), and use that fixed subspace for the next T iterations. This is essentially a batch processing view of subspace identification: collect a gradient snapshot, decompose it, and reuse until the next collection point.

The paper reframes the problem as streaming subspace tracking: the optimization process generates a continuous stream of gradient matrices G_t whose dominant subspace slowly rotates as the loss landscape changes. The goal is not to perfectly capture any single G_t but to maintain a projection matrix P_t that stays close to the dominant subspace of the current gradient distribution at all times. This reframing has three consequences, each of which is individually significant:

First, it recasts staleness as the central failure mode of periodic methods. GaLore's P_t is set to the SVD of G_{T\lfloor t/T \rfloor} — the gradient at the most recent recomputation step. Between recomputations (e.g., 200 steps), G_t changes due to mini-batch variation, learning rate decay, and evolving loss landscape geometry, but P_t stays fixed. The paper's experiments (Table 1, Figure 1) provide evidence that this staleness matters: Online Subspace Descent with continuous updates achieves 31.30 perplexity vs. 35.52 for GaLore on LLaMA 1B at rank 512, and the gap grows with model size (the paper hypothesizes that larger models have higher intrinsic-rank optimization problems, making timely gradient information more important). Putting this in terms the field can operationalize: the update frequency of P_t is not a hyperparameter to tune but a fundamental design axis — the analysis says "update as frequently as computationally feasible," and the online PCA approach makes "every iteration" feasible.

Second, it motivates the specific choice of the online PCA loss L_{G_t}(P) by connecting to the SVD's mathematical role. GaLore uses SVD because it produces the optimal rank-k approximation of the gradient. The online PCA loss in Equation 6 has the same global minima — its stable stationary points are exactly the top-k singular vectors of the normalized gradient \tilde{G}_t (Baldi & Hornik, 1989). The reframing says: we don't need the exact optimum at every step (which is what SVD provides at prohibitive cost); we just need to take a step toward the optimum at each iteration, trusting that the gradient subspace changes slowly enough that P_t stays in the vicinity of the true dominant subspace. The orthonormality penalty \lambda\|P^\top P - I\|^2 added to the loss is a direct consequence of this reframing — in a batch-setting SVD, orthonormality is automatic, but gradient-based optimization of a PCA objective can drift toward degenerate P (correlated or scaled columns), so the penalty is needed to keep P_t well-behaved during continuous tracking.

Third, it creates a new design space: the choice of optimizer and hyperparameters for P_t. Under the periodic-SVD paradigm, there is no "optimizer for the subspace" — the subspace is computed exactly and replaced wholesale. Under the streaming paradigm, P_t is a parameter that is optimized over time, which means we must choose: what optimizer updates it? What learning rate relative to the weight optimizer? What regularization? The paper explores this design space empirically (Section 5.3, Section 5.4, Table 2): Adam on P_t with learning rate \alpha = 5\times the weight learning rate works well; SGD on P_t with Adam on weights also works (Table 2, Ours AdamW8bit + SGD: 53.53 perplexity, not as good as full Adam but far better than GaLore's 48.05 on the same 60M model); the orthonormality regularization \lambda = 0.1 is not sensitive (Figure 3 right). The paper does not claim to have optimized this design space, but its reframing of the problem makes the design space visible for the first time.

This is best characterized as a fundamental reframing with incremental practical gains that validate the reframing. The core idea — track rather than batch — is a conceptual shift, not just a new hyperparameter. The practical improvements (lower perplexity, faster wall-clock time) are the empirical signature that the reframing matters, but the reframing's value extends beyond this paper: it suggests that future work on memory-efficient training should focus on streaming algorithms for subspace tracking (online PCA, incremental SVD, sketching with exponential decay) rather than ever-more-accurate periodic decompositions.

Innovation 3: Demonstrating That Test-Time Compute for the Projection Matrix (Online PCA) Can Be Cheaper and Better Than Exact Batch Computation (SVD) — a Counterintuitive Efficiency Result

This innovation is less about theory and more about an empirical finding with engineering significance that challenges a natural assumption. The natural assumption is: if we're going to project gradients into a low-rank subspace, we should use the best low-rank approximation available — exact SVD, computed periodically. Online PCA produces a worse approximation at each individual step (it's a single gradient step, not a converged solution), so it should perform worse for any given rank. This assumption would predict that Online Subspace Descent needs a higher rank to match GaLore's perplexity, or that it simply can't match GaLore at all.

The experimental results in Section 5.1 (Table 1, Figure 1) falsify this assumption. At the same rank (rank 512), Online Subspace Descent achieves substantially better perplexity than GaLore (31.30 vs. 35.52 on LLaMA 1B, sequence length 256, 10K steps). This is not a small margin — it's a 4.22-point perplexity gap, comparable to the gap between GaLore and full-rank AdamW8bit (35.52 vs. 29.40). The finding is consistent across model sizes (60M, 350M, 1B) and grows more pronounced as scale increases (Figure 1, validation perplexity curves).

Why is this counterintuitive? At any single iteration t, GaLore's SVD-based P_t is a better rank-k approximation of the gradient G_t than Online Subspace Descent's online-PCA P_t (by the definition of SVD optimality). If the only thing that mattered was the quality of gradient approximation at each step, GaLore would win. The fact that Online Subspace Descent wins suggests that something matters more than per-step approximation quality. The paper's hypothesis (Section 5.1) is that the "something" is timeliness: online PCA responds immediately to gradient distribution shifts, while GaLore's SVD-based P_t is stale for up to 200 steps between recomputations. In a non-stationary optimization process where the important gradient directions evolve continuously, a slightly worse approximation of the current gradient is more valuable than a better approximation of a past gradient.

The wall-clock time results (Table 3, Figure 2) reinforce the counterintuitive finding. Online PCA is faster per-update (up to 142× faster than SVD on representative tensor shapes, per Figure 2), so Online Subspace Descent trains faster overall (7.14 hours vs. 9.74 hours for 7B model pretraining, Table 3). If SVD's better per-step approximation led to faster convergence (fewer steps needed), GaLore might still win on wall-clock time. The data show the opposite: Online Subspace Descent is both faster per-step and reaches better perplexity at the same step count, meaning the efficiency gain is compound — not just cheaper updates, but more effective ones.

This is an incremental empirical finding with fundamental implications for how the field should think about computation in optimization. It suggests that "compute spent on better approximating the current gradient" and "compute spent on incorporating more recent gradient information" are competing uses of a computational budget, and that the latter can dominate in non-stationary settings. This connects to broader trends in optimization (streaming algorithms, online learning) but makes the case in the specific, high-impact context of LLM pretraining where memory efficiency is a primary bottleneck. The 1.3× wall-clock speedup at 7B scale (Table 3) is practically significant on its own — it directly translates to reduced GPU-hours and cost for pretraining runs.

Innovation 4: Unifying the Convergence Analysis of Subspace Descent Methods Under a Single Lyapunov Framework — and Clarifying What the Field Did and Didn't Know

The paper's theoretical contribution is not just the Hamiltonian+descent invariance result (Innovation 1), but also the unification of that result with a concrete convergence theorem (Theorem 4.5) that specifies exactly what is required for subspace descent optimizers to reach stationary points. This unification resolves a specific, named gap in the literature.

Before this work, the state of theoretical understanding was fragmented along two axes. Axis 1: optimizer type. Stochastic Subspace Descent (Kozak et al., 2019) provided convergence analysis but only for convex objectives and only for random (non-adaptive) projection matrices — no momentum, no Adam, no learned P_t. GaLore and Sketchy provided no convergence analysis at all. Axis 2: projection dynamics. Methods with static projection (LoRA-style fine-tuning) had convergence analysis because the optimizer was effectively running in a fixed low-dimensional space. Methods with dynamic projection had no analysis beyond the naive coordinate-descent intuition (which, as discussed in Innovation 1, fails for stateful optimizers). The paper's contribution is to fill the non-convex + stateful optimizer + dynamic projection cell of this matrix, which is exactly the cell that matters for LLM pretraining.

The unification is powerful because it separates two concerns that were previously conflated:

• Concern 1: Does the projection mechanism break the optimizer's descent? Answer: No — the adjoint cancellation (Section 4.1) guarantees that d/dt H \le 0 for any P_t update rule, any optimizer with the Hamiltonian+descent structure, and any smooth time-variation of P_t. This is a structural invariant that does not depend on convexity, learning rate, rank, or the specific form of \Gamma.

• Concern 2: Does the optimizer converge to a stationary point of the loss, or does it get stuck where P_t^\top G = 0 but G \neq 0? Answer: Convergence to stationary points requires the mild additional condition that P_t does not permanently reside in the nullspace of a non-zero gradient (Assumption 4.4(ii), Theorem 4.5). This is satisfied by any PCA-like update that minimizes reconstruction error — if P_t^\top G = 0, then PP^\top G = 0, which is a terrible reconstruction of G, so the PCA loss gradient will push P_t away from this configuration.

The paper's framing of this as a convergence guarantee for arbitrary update rules (the claim in the abstract and introduction) is deliberate and rhetorically bold. It says: we are providing a guarantee not just for our specific algorithm (online PCA with Adam on P_t) but for an entire class of possible algorithms that differ in how they update P_t. This is stronger than a typical algorithm paper's theoretical contribution (which would only analyze its own update rule) and weaker than a pure theory paper (which might characterize the optimal update rule). The paper is explicitly carving out a middle ground: "here is a framework that certifies a large family of methods as sound; we propose one member of this family (online PCA) and show it works well empirically."

The significance of this unification extends to future research directions that the paper enumerates in its conclusion. Questions like "are there alternative methods for updating projection matrix that could accelerate convergence?" and "what is the impact of weight decay on convergence?" are now foregrounded as natural extensions of the framework rather than open-ended empirical questions. The Lyapunov analysis provides the mathematical language for reasoning about these questions — for example, the paper notes (Section 4.1) that the decreasing rate d/dt H depends on P_t through \hat{M}_t, \hat{V}_t, \hat{G}_t, opening the door to treating the choice of \Gamma as an optimal control problem that maximizes this rate.

This is a fundamental theoretical contribution that both closes a specific gap (no prior analysis for dynamic projection + stateful optimizers) and opens a research program (optimal P_t update rules, extensions to other operator classes, integration with weight decay analysis). The fact that the theorem's conditions are mild (smoothness, boundedness, non-degeneracy of P_t) means the framework is likely to accommodate many future variants without requiring new proofs from scratch — a hallmark of useful theory.

5. Experimental Analysis

Evaluation Methodology

• Dataset. All pretraining experiments use the C4 dataset (Raffel et al., 2019), a large-scale corpus of cleaned web text commonly used for language model pretraining evaluation. The paper does not specify an explicit train/validation/test split size but reports validation perplexity as the primary evaluation signal (Figures 1, 3, Table 1). Downstream evaluation for the 7B models (Table 4) uses the GLUE benchmark: MRPC, RTE, SST2, MNLI, QNLI, and QQP tasks, evaluated via the eval-harness framework.

• Base model(s). All experiments use the LLaMA architecture (Touvron et al., 2023) at four scales: 60M, 350M, 1B, and 7B parameters. The paper states that these scales span "representative" sizes for studying memory-efficient pretraining, with the 60M–1B experiments run on a single NVIDIA A100 GPU and the 7B run on unspecified hardware (likely multi-GPU given wall-clock times in Table 3). The LLaMA architecture is chosen because it is a standard, open, and well-studied model family, allowing direct comparison with GaLore (Zhao et al., 2024) results.

• Metrics.

  • Perplexity: Standard language modeling perplexity on held-out validation data. For most tables (Table 1, 2, 5, 6), perplexity is reported as "the training average of the last 10 steps" (Table 1 caption). Lower is better.
  • GLUE task performance (Table 4): Standard accuracy/F1 metrics for each GLUE task, with the average across tasks reported. Used to evaluate the downstream quality of pretrained checkpoints, not just language modeling loss.
  • Wall-clock time (Table 3): Total training time in hours for the 10K-step 7B pretraining run. Used to quantify the computational overhead advantage of Online Subspace Descent over GaLore.
  • Execution time per operation (Figure 2): Micro-benchmarked execution time of torch.svd vs. a single backward step for the online PCA loss, measured on representative weight tensor shapes from LLaMA architectures across A100 and RTX 3090 GPUs. Used to quantify the per-step speed advantage of online PCA over SVD.

• Baselines. The paper compares against five distinct methods or configurations:

  • Full-rank AdamW8bit (Table 1, 5, 6): Standard AdamW with 8-bit quantization (Dettmers et al., 2021), using full-precision optimizer states without any subspace projection. Serves as the upper bound for convergence quality.
  • GaLore (Zhao et al., 2024): The state-of-the-art low-rank training method that uses periodic SVD to set the projection matrix. Compared across all model sizes and ranks.
  • GaLore with LION/Adafactor (Table 6): GaLore equipped with the LION optimizer and Adafactor optimizer respectively, to test whether GaLore's SVD-based projection generalizes across base optimizers.
  • Standalone LION, Adafactor, AdamW8bit, SGD (Table 6): Full-rank versions of these optimizers without subspace projection, providing baselines for how much the subspace projection itself contributes to memory savings and performance.
  • Majority voting / different optimizers for $P_t$ vs. $W_t$ (Table 2, 6): Ablations where the optimizer for updating $P_t$ differs from the optimizer for updating $W_t$ (e.g., LION for weights + AdamW8bit for $P_t$, SGD for weights + AdamW8bit for $P_t$).

• Generation budget / compute accounting. The paper's primary fairness axis is rank — all low-rank methods are compared at the same projection rank $k$ (e.g., rank 512 for LLaMA 1B, Table 1; rank sweep from 32 to 512 for LLaMA 60M, Table 5). The paper accounts for computational cost through wall-clock time (Table 3) and micro-benchmarks of per-operation execution time (Figure 2), but does not report total FLOPs or GPU memory consumption figures for most experiments (the only memory comparison is qualitative: "consuming a similar amount of GPU memory (8.64 GB v.s 9.01 GB)" for the 1B experiment in Section 5.1). All methods are trained for the same number of steps (10K) to ensure fair step-count comparisons. The paper explicitly notes that online PCA's $P_t$ updates "can be executed in parallel with weight updates, adding no overhead to the training process" (Figure 2 caption), and the wall-clock results in Table 3 bear this out: Online Subspace Descent trains the 7B model in 7.14 hours vs. 9.74 hours for GaLore — a 1.3× speedup — despite both running 10K steps.

• Cross-validation / statistical protocol. No cross-validation or statistical significance testing is reported. The paper sweeps hyperparameters (learning rate, $\alpha$, $\lambda$, rank) on the 60M model (Section 5.2–5.3) and transfers the best settings to larger scales (350M, 1B, 7B). All results are reported as single-run values on the C4 validation set or GLUE evaluation, with no confidence intervals or error bars. The only averaging mentioned is "training average of the last 10 steps" for perplexity in Table 1. This is a notable methodological limitation — see Critical Assessment.

Main Quantitative Results

Search and Iterative Improvement Metrics (GaLore vs. Online Subspace Descent at Scale)

Headline: Online Subspace Descent achieves 31.30 vs. 35.52 perplexity against GaLore on LLaMA 1B with rank 512 — a 4.22-point gap that narrows the distance to full-rank AdamW8bit (29.40) by roughly 69%. The results appear in Table 1 and are consistent across model sizes (Figure 1 shows the validation perplexity curve for LLaMA 1B, sequence length 256, across 10K steps; Table 1 reports final perplexity as the average of the last 10 steps).

Scale-dependent patterns: The perplexity gap between Online Subspace Descent and GaLore grows with model size. At 60M, the gap is small: 56.12 (Ours) vs. 57.03 (GaLore) at rank 512, and 37.41 (Ours) vs. 36.93 (GaLore) at full rank-equivalent (Table 5). At 350M, the gap opens to 43.67 (Ours) vs. 44.34 (GaLore) at rank 512 (Table 1). At 1B, the gap widens substantially to 31.30 (Ours) vs. 35.52 (GaLore) at rank 512 — a 4.22-point difference that represents the largest relative improvement. The paper hypothesizes that this trend occurs because "the higher intrinsic rank of the underlying optimization problem in larger models" makes timely gradient information incorporation (enabled by online PCA) progressively more important than stale-but-exact SVD (Section 5.1).

Sequence length interaction: Figure 1 shows validation perplexity for LLaMA 1B with sequence length 256 across 10K steps. The paper notes that "as model size and sequence length grow, Online Subspace Descent becomes more effective" (Section 5.1), but only the Figure 1 curve (SS 256) is shown in the main results. Appendix A.2 provides the rank sweep for sequence length 1024 (Table 5), reproduced below for the LLaMA 60M case.

Rank sweep results (Table 5, Appendix A.2): On LLaMA 60M with sequence length 1024, the final perplexity improves monotonically with rank for both methods:

Two observations: (1) At rank 128, GaLore (48.05) slightly outperforms Online Subspace Descent (49.01), reversing the usual trend. This is consistent with the scale-dependence hypothesis — the 60M model may have low enough intrinsic rank that exact SVD is adequate. (2) The gap between rank 512 and full rank is small for both methods (37.41 vs. 37.18 for Ours; 36.93 vs. 36.51 for GaLore), indicating that rank 512 largely saturates the available subspace information for this model scale.

Wall-Clock Time and Execution Overhead (Figure 2, Table 3)

Headline: Online Subspace Descent trains LLaMA 7B 1.3× faster than GaLore (7.14 vs. 9.74 hours, Table 3) while achieving lower perplexity (43.72 vs. 51.21), with the speed advantage rooted in the fact that single-step online PCA can be up to 142× faster than equivalent SVD on representative weight tensor shapes (Figure 2).

Micro-benchmark results (Figure 2): The paper measures execution time of torch.svd vs. a single backward pass for the online PCA loss "on matrices of typical shapes in linear layers in the LLaMA 60M to 7B" (Figure 2 caption). The key finding is that the speed ratio grows dramatically with tensor size:

• On smaller tensors (lower portion of Figure 2), the gap is modest (e.g., 5–10×).
• On larger tensors (upper portion, representative of larger models), the gap reaches 142× in the worst case shown.
• The consumer GPU (RTX 3090) and data center GPU (A100) show qualitatively similar patterns.

This micro-benchmark provides the mechanistic explanation for the wall-clock advantage: each GaLore SVD recomputation (every 200 steps) is so expensive that it cannot be hidden by overlapping with weight updates, whereas each online PCA update is cheap enough to execute in parallel with the weight optimization step without becoming a bottleneck. The paper explicitly states: "Online PCA is fast because it is implemented as a single optimization step with respect to a simple loss function. Hence, each step of online PCA can be cleverly scheduled and hidden in the weight optimization step when executed in parallel, whereas SVD is too expensive to be hidden" (Section 5.1).

Wall-clock time at 7B scale (Table 3):

Two key points: (1) The speedup (1.3×) is smaller than the micro-benchmark would suggest for a single SVD vs. single PCA step (up to 142×) because SVD is only called periodically (every 200 steps) while online PCA runs every step. The effective per-step overhead ratio is thus amortized. (2) The perplexity improvement (43.72 vs. 51.21) is substantially larger than the speedup, suggesting that the primary benefit is convergence quality, not just compute savings.

GPU memory: The paper briefly notes that on LLaMA 1B with SS 256, Online Subspace Descent uses 8.64 GB vs. 9.01 GB for GaLore (Section 5.1, body text). This is a modest difference — both methods use comparable memory, which is expected since both maintain $k \times m$ optimizer states and an $n \times k$ projection matrix. The memory advantage of Online Subspace Descent over GaLore is from the optimizer states for $P_t$ being slightly smaller than the SVD workspace, not from fundamental algorithmic differences.

Downstream Task Performance (Table 4)

Headline: Online Subspace Descent pretrained 7B checkpoint outperforms GaLore on 4 of 6 GLUE benchmark tasks, with an average improvement of 0.89 percentage points (0.4951 vs. 0.4862).

The standardized GLUE evaluation results (Table 4) for the 7B model checkpoints are:

The improvements are consistent but uneven: Online Subspace Descent wins on MRPC (+0.0144), SST2 (+0.005), MNLI (+0.0148), QNLI (+0.0196), and QQP (+0.0113), but loses on RTE (-0.0117). The largest absolute gain is on QNLI, a question-answering entailment task. This demonstrates that lower perplexity (Table 3) translates to better downstream task performance, establishing that the gains are not merely an artifact of language modeling metrics — the pretrained representations are genuinely more useful.

An important caveat: these are single-run evaluations of pretrained checkpoints after only 10K steps of training on C4. The GLUE scores are far from converged (e.g., SST2 at 0.52 vs. typical fine-tuned models exceeding 0.90), so these results primarily measure the quality of pretrained representations early in training, not final downstream capability. The consistent improvement over GaLore at this early stage is a valid signal, but the absolute numbers should not be interpreted as indicative of what either method would achieve with full pretraining.

Optimizer Compatibility Experiments (Table 2, Table 6)

Headline: Online Subspace Descent works across multiple base optimizers (LION, Adafactor, AdamW8bit, SGD), with the best results achieved when AdamW8bit updates both $W_t$ and $P_t$, but SGD on $P_t$ with AdamW8bit on $W_t$ still outperforms GaLore (Table 6).

The paper explores an important practical question: is Online Subspace Descent tied to Adam, or does the subspace projection mechanism work with any base optimizer? Table 2 (LLaMA 60M, sequence length 1024) tests LION and Adafactor as base optimizers:

Several patterns emerge:

First, simply swapping both optimizers to LION or Adafactor does not always help. Online Subspace Descent with LION+LION (LION for weights + LION for $P_t$) achieves 57.97 perplexity — worse than GaLore LION (46.90) and much worse than any Adam-based configuration. With Adafactor+Adafactor, perplexity is 47.61 — better than GaLore AdamW8bit (48.05) but worse than GaLore Adafactor (34.32). This suggests that the $P_t$ update is sensitive to the choice of optimizer, and that $P_t$ benefits from adaptive per-dimension learning rates (Adam) more than the weight update does in this setting.

Second, mixing optimizers is effective. The strongest configurations use AdamW8bit for $P_t$ and a different optimizer for $W_t$: LION + AdamW8bit achieves 44.76 perplexity, and Adafactor + AdamW8bit achieves 34.15 — only slightly worse than the best Adafactor-based configuration (GaLore Adafactor at 34.32). This is a significant practical finding: one can use a memory-efficient base optimizer (Adafactor, LION) for the weights while still using Adam for $P_t$ (where the memory overhead is negligible since $P_t$ itself is small), achieving most of the convergence benefit of Adam at reduced memory.

The more comprehensive Table 6 (Appendix A.3) adds additional combinations:

This is a notable result: even SGD on $P_t$ with AdamW8bit on $W_t$ works — it's worse than AdamW8bit+AdamW8bit (49.01), but substantially better than the worst configurations (LION+LION at 57.97). This demonstrates that the projection matrix update rule is robust to the choice of optimizer, consistent with the theoretical finding (Section 4.1) that the Lyapunov property holds for any $\Gamma$. The practical implication is that practitioners with extremely tight memory budgets can use SGD for $P_t$ (zero optimizer state overhead beyond the matrix itself) and still obtain reasonable performance.

Full-rank baselines (Table 6, standalone section): For context, the full-rank optimizers evaluated are LION (52.65), Adafactor (33.45), AdamW8bit (29.77), and SGD (3469.14 — essentially non-convergent). The gap between full-rank AdamW8bit (29.77) and the best Online Subspace Descent configuration at rank 128 on this model (49.01, Table 5) is substantial — roughly 19 points of perplexity — highlighting that low-rank projection at rank 128 with this model size leaves significant convergence quality on the table. The rank sweep in Table 5 shows that increasing rank closes this gap (rank 512 reaches 37.41), but does not eliminate it (full rank reaches 37.18).

Hyperparameter Sensitivity (Figure 3, Section 5.3)

Headline: Rank is the most impactful hyperparameter (higher rank monotonically improves perplexity with diminishing returns); the projection matrix learning rate multiplier $\alpha$ must be kept small to avoid instability; the orthonormality penalty $\lambda$ shows minimal sensitivity for small models.

The paper presents three ablation sweeps on LLaMA 60M with 10K steps (Figure 3):

Rank sweep (Figure 3, left): The loss curves for different rank values show a clear hierarchy — higher ranks converge to lower loss, with the gaps narrowing as rank increases. The paper notes that "the rate of reduction of perplexity decreases as the rank increases, eventually reaching a saturation point" and hypothesizes that "high-frequency tokens can be effectively learned with low-rank training. However, learning lower-frequency tokens requires higher ranks" (Section 5.2). The practical recommendation: "select the highest rank until the perplexity reduction saturates" — essentially, use as much memory budget as available for the rank, since higher rank is never harmful (unlike beam search in the reference example paper, where higher search budget could hurt performance on easy problems due to verifier over-optimization). This is a simpler story: monotonic improvement with rank, just diminishing returns.

$\alpha$ sweep (Figure 3, middle): $\alpha$ controls the relative learning rate of $P_t$ vs. $W_t$ via $\epsilon_t^P = \alpha \epsilon_t^W$. The sweep shows that large $\alpha$ values cause instability (the loss curve for $\alpha = 10$ shows spikes and higher final loss), while $\alpha = 1$ and $\alpha = 5$ produce similar, stable curves. The paper selects $\alpha = 5$ for all experiments, noting that "$\alpha$ must be kept small to avoid instability." This makes intuitive sense: if $P_t$ learns too quickly relative to $W_t$, the projection matrix can overfit to the most recent gradient, potentially causing the optimizer to chase gradient noise rather than signal.

$\lambda$ sweep (Figure 3, right): $\lambda$ controls the orthonormality regularization strength in the PCA loss. The sweep shows nearly identical loss curves for a range of $\lambda$ values for the 60M model. The paper states that "the result is not sensitive to $\lambda$ for small models (60M)" and sets $\lambda = 0.1$ for all subsequent experiments. The low sensitivity is reassuring — it means the PCA loss is well-behaved and the orthonormality penalty is not a fragile hyperparameter that needs careful tuning.

Learning rate for larger models (Section 5.3): The paper separately recommends reducing the base learning rate by 10× for larger models (0.001 for 350M, 1B), as "larger learning rates cause unrecoverable spikes and instability, a general characteristic observed across all methods." This is not specific to Online Subspace Descent — it's a standard practice for training larger models — but the paper flags it as an important practical note.

Ablation Studies and Robustness Checks

• Rank (Table 5, Figure 3 left): Higher rank monotonically improves final perplexity with diminishing returns. On LLaMA 60M SS 1024, perplexity drops from 85.90 (rank 32) to 37.41 (rank 512), with the marginal gain shrinking from ~37 points (32→128) to ~12 points (128→512) to 0.23 points (512→full). The paper draws an interesting connection to token frequency: "high-frequency tokens can be effectively learned with low-rank training, but learning lower-frequency tokens requires higher ranks." This is not experimentally verified (no token-frequency breakdown is provided), but it's a plausible mechanism that would explain why the gains from higher rank saturate — once all token-level gradients have adequate subspace representation, further rank increases provide minimal additional benefit.

• Choice of optimizer for $W_t$ vs. $P_t$ (Table 2, 6): The choice of optimizer for $P_t$ matters more than the choice for $W_t$ in some regimes, with AdamW8bit on $P_t$ consistently yielding the best results regardless of the $W_t$ optimizer. Specifically, Adafactor + AdamW8bit (34.15) nearly matches GaLore Adafactor (34.32) and substantially outperforms Adafactor + Adafactor (47.61). This is a non-obvious finding that has direct practical implications: even when using a memory-efficient base optimizer for the weights, one should still invest the small amount of extra memory needed for Adam on $P_t$. The theoretical framework (Section 4) provides no guidance on this choice (the Lyapunov property holds for any $\Gamma$), so this is purely an empirical finding.

• SGD on $P_t$ (Table 6): Online Subspace Descent with SGD on $P_t$ and AdamW8bit on $W_t$ achieves 53.53 perplexity — substantially worse than Adam on both (49.01) but far better than the non-convergent full-rank SGD baseline (3469.14) and competitive with GaLore AdamW8bit (48.05). This matters because SGD on $P_t$ introduces zero optimizer state overhead for the projection matrix — if memory is so tight that even the small Adam state for $P_t$ is problematic, SGD remains a viable fallback. Combined with the theory, this result also serves as an empirical validation that the convergence guarantee (which does not depend on the choice of $\Gamma$) holds up in practice — the system converges with a suboptimal $P_t$ update rule, just more slowly.

• $\alpha$ (projection matrix learning rate multiplier, Figure 3 middle): Values that are too large ($\alpha = 10$) cause instability (loss spikes), while $\alpha = 1$ and $\alpha = 5$ produce similar stable training. The paper selects $\alpha = 5$ for all experiments. This is consistent with the intuition that $P_t$ should adapt faster than $W_t$ (it tracks the changing gradient geometry, which shifts at a faster timescale than the optimal weights), but not so fast that it overreacts to gradient noise. There is no ablation for $\alpha$ on larger models, leaving open the question of whether the optimal $\alpha$ changes with scale.

• $\lambda$ (orthonormality penalty, Figure 3 right): Insensitive over a reasonable range for the 60M model. The paper sets $\lambda = 0.1$ for all experiments. Without the penalty ($\lambda = 0$), the PCA loss would have degenerate solutions where $P$'s columns become parallel or have varying scales — the penalty prevents this. The low sensitivity is a positive robustness signal: practitioners do not need to carefully tune this hyperparameter. As with $\alpha$, no ablation for $\lambda$ on larger models is provided.

• Sequence length (Figure 1 for SS 256; Table 5 for SS 1024): The paper compares results at different sequence lengths (256 vs. 1024) for the 60M model. At SS 1024 (Table 5), Online Subspace Descent achieves 37.41 vs. GaLore 36.93 at rank 512 — essentially tied, with GaLore slightly ahead. At SS 256 (Table 1, 60M column), the gap is 56.12 vs. 57.03 — Ours slightly ahead, also essentially tied. For the 1B model at SS 256, Ours pulls substantially ahead (31.30 vs. 35.52, Table 1). This pattern suggests that the online PCA advantage is not primarily about sequence length but about model scale — a key robustness check for the "intrinsic rank" hypothesis. A direct sequence length comparison for the 1B model would strengthen this claim, but is not provided.

• Downstream task evaluation for 7B model (Table 4): The paper validates that perplexity improvements translate to downstream task gains, with Online Subspace Descent outperforming GaLore on the GLUE average (0.4951 vs. 0.4862). An important negative result: Online Subspace Descent loses on RTE (0.4901 vs. 0.5018), a 0.0117 deficit. This is not discussed in the paper and may be noise given the single-run nature, but it suggests that the downstream benefits are not uniform across tasks.

• Wall-clock time with SVD vs. online PCA (Figure 2, Table 3): The micro-benchmarks show that online PCA is dramatically faster per-operation (up to 142×) than SVD on representative tensor shapes, and the end-to-end 7B training run confirms a 1.3× overall wall-clock speedup. The gap between micro-benchmark (142× per operation) and end-to-end (1.3× overall) is explained by amortization — SVD runs every 200 steps while online PCA runs every step, so the effective cost difference is 142× / 200 ≈ 0.71× per effective step in the limit, which is close to the observed speedup when accounting for other training overheads (data loading, communication, etc.).

Critical Assessment

Claim from the executive summary: "Online Subspace Descent achieves lower perplexity than state-of-the-art low-rank methods (31.30 vs. 35.52 for GaLore on a 1B model at rank 512)."

This claim is directly supported by Table 1 for the specific configuration tested (LLaMA 1B, sequence length 256, 10K steps, rank 512). However, several qualifications are essential:

First, the claim generalizes inconsistently across model scales. At 60M with SS 1024 (Table 5), Online Subspace Descent (37.41) is actually slightly worse than GaLore (36.93) at rank 512 — the opposite of the headline result. The paper's hypothesis (higher intrinsic rank in larger models favors continuous updates) is plausible but not proven. To demonstrate that this is a scale effect rather than a configuration artifact, the paper would need to show that at 60M, even with extensive hyperparameter tuning, online PCA cannot beat periodic SVD, while at larger scales it consistently does. The current evidence is suggestive but not conclusive: we see one negative result at 60M SS 1024, one near-tie at 60M SS 256, one modest win at 350M, and one large win at 1B. This is consistent with the scale hypothesis, but a single negative example at each small scale would be stronger evidence.

Second, the comparison is at equal rank, but the ranks tested may not be fair. GaLore benefits from exact SVD's ability to perfectly capture the top-$k$ singular vectors of the gradient at the recomputation point. Online Subspace Descent's approximate online PCA may need a higher rank to capture the same amount of gradient information. A fairer comparison might be: what rank does Online Subspace Descent need to match GaLore at rank 512? Or, at equal memory budget (which would be similar since both store $P_t$), which method performs better? The paper's rank sweep (Table 5) shows that at 60M, Online Subspace Descent at rank 512 (37.41) and GaLore at rank 512 (36.93) are close, but Online Subspace Descent at rank 128 (49.01) is worse than GaLore at rank 128 (48.05). This suggests that the rank at which Online Subspace Descent overtakes GaLore may depend on model scale, and the 1B result at rank 512 is the most favorable case shown.

Third, the training duration is short (10K steps) relative to full pretraining. For the 7B model with 10K steps, the perplexity numbers (43.72 vs. 51.21 for GaLore) reflect early-training dynamics, not converged behavior. It is possible that GaLore catches up later in training when the gradient distribution stabilizes and exact periodic SVD becomes more effective relative to continuous online PCA. Conversely, online PCA might maintain its advantage throughout training if gradient geometry continues to evolve. The paper cannot distinguish these possibilities with only 10K-step runs.

Claim from the executive summary: "Online Subspace Descent narrows the gap with full-rank baselines."

Supported with the specific numbers: on LLaMA 1B SS 256, full-rank AdamW8bit achieves 29.40, GaLore achieves 35.52, and Online Subspace Descent achieves 31.30 (Table 1). The gap reduction is from 6.12 points (GaLore vs. full-rank) to 1.90 points (Ours vs. full-rank) — roughly a 69% reduction. This is a substantial closing of the gap.

However, the full-rank baseline (AdamW8bit) is not the only possible point of comparison. The paper does not report memory consumption for the full-rank baseline relative to the low-rank methods, making it difficult to assess whether the perplexity improvement is worth the memory cost. If full-rank AdamW8bit uses 3× the memory of Online Subspace Descent (as one would expect from storing full-size $M_t$ and $V_t$), then closing the gap to within 1.90 perplexity points while using substantially less memory is a genuine achievement. But if the memory savings are modest at rank 512 (because $P_t$ itself and $\hat{S}_t$ are similar in size to full-rank states for some layers), the tradeoff is less compelling. The paper reports only one memory comparison (8.64 GB vs. 9.01 GB for 1B models, Section 5.1), which shows minimal difference between the two low-rank methods but does not include full-rank memory usage.

Claim from the executive summary: "Single-step online PCA can be up to 142× faster than equivalent SVD calls on representative weight tensors."

This is directly supported by Figure 2, which clearly shows the speed ratio growing with tensor size. However, the figure lacks precise numerical labels — the "142×" number appears only in the Figure 2 caption text ("up to 142 times slower"), not on the plot axes. The plot shows execution time on both A100 and RTX 3090, with SVD consistently above PCA, but the specific tensor shapes and exact ratios for each point are not tabulated. The claim "up to 142×" should be understood as the maximum ratio observed across the tested shapes, not an average or typical case. The end-to-end wall-clock speedup (1.3×) provides a more practically meaningful number.

What experiments are missing?

1. Scaled-up training duration experiments. All runs use 10K steps, which is a tiny fraction of typical LLM pretraining (which may use 100B+ tokens, or hundreds of thousands of steps). The relative performance of online PCA vs. periodic SVD may change as training progresses: early in training, the gradient distribution may shift rapidly (favoring online PCA), while later it may stabilize (favoring periodic SVD). Without longer runs, we cannot know whether the 1B result (31.30 vs. 35.52) holds at, say, 100K steps or whether GaLore catches up.

2. Direct memory measurements for all configurations. The paper is about memory-efficient training, yet provides almost no memory usage data. The only memory number is for a single configuration (Section 5.1: 8.64 GB vs. 9.01 GB for the 1B model). There are no memory comparisons for 60M, 350M, or 7B models, no breakdown of memory by component (weights, optimizer states, projection matrix, activations), and no comparison to full-rank baselines. This is a significant omission for a paper whose primary motivation is memory efficiency.

3. Ablation of the SVD recomputation frequency in GaLore. The paper attributes GaLore's inferior performance to staleness between SVD updates, but never varies GaLore's update gap (set to the recommended default of 200). If the staleness hypothesis is correct, reducing GaLore's update gap (e.g., to 50 or 100 steps) should improve GaLore's performance and potentially close the gap with Online Subspace Descent. Without this ablation, we cannot distinguish between "SVD is inherently worse because it's stale" and "the default update gap for GaLore was poorly chosen for this setting."

4. Larger-scale runs (beyond 7B). The paper's scale-dependence hypothesis (larger models benefit more from online PCA) is supported by the 60M → 350M → 1B trend, but the 7B result (Table 3) doesn't include a GaLore comparison at the same rank — only a single perplexity number (43.72 for Ours vs. 51.21 for GaLore) without specifying the rank used or the GaLore configuration details. Testing at 13B or 30B would more strongly validate the claim that online PCA scales better than periodic SVD.

5. Multiple random seeds / statistical significance. All results are single-run. Given the 10K-step training budget, different random seeds could produce different perplexity values, especially for the smaller models where the gap between methods is narrow (e.g., 60M SS 1024 at rank 512: 37.41 vs. 36.93 — this is small enough to potentially reverse with a different seed). Error bars or multiple runs would substantially increase confidence in the reported results.

6. Combination of online PCA with GaLore-style SVD initialization. A natural hybrid would be: initialize $P_0$ with SVD of the first gradient, then use online PCA for continuous updates. This might combine the best of both worlds — exact initialization plus continuous adaptation. The paper does not explore this or any other hybrid.

7. Ablation of the normalization in the PCA loss ($\tilde{G}_t = G_t / \|G_t\|$). The paper normalizes gradients before the PCA loss computation to prevent scale from dominating direction. Removing this normalization and observing the effect on training stability would validate the design choice. This is not ablated.

8. Impact on training dynamics beyond convergence. The paper focuses on final perplexity and downstream task performance, but does not analyze whether online PCA changes the qualitative behavior of training — e.g., whether it reduces loss spike frequency, whether it improves gradient signal-to-noise ratio, or whether it affects the learned representations beyond what perplexity captures.

Where the claims hold conditionally:

• "Lower perplexity than GaLore" holds for the 350M and 1B models at rank 512 with sequence length 256, but not for the 60M model at rank 512 with sequence length 1024 (Table 5: GaLore 36.93 vs. Ours 37.41). The condition appears to be: the model must be large enough (and/or have high enough intrinsic optimization rank) for continuous updates to overcome the approximation error of single-step PCA. The precise threshold is not characterized.

• "1.3× faster wall-clock time" holds for the 7B model trained for 10K steps (Table 3). For smaller models, the speedup may differ because SVD overhead is proportionally smaller relative to other training costs (forward/backward passes). The paper does not report wall-clock times for 60M–1B models, so the speedup should not be assumed to generalize without verification.

• "Up to 142× faster per-operation" is the maximum observed ratio across specific tensor shapes (Figure 2), not an average or typical case. The actual per-operation speedup depends on tensor dimensions and hardware.

Overall assessment: The experiments provide credible evidence that Online Subspace Descent outperforms GaLore on LLM pretraining at the 350M–7B scale in terms of both perplexity and wall-clock time, with the advantage growing with model size. The theoretical framework (Section 4) provides a principled justification for why dynamic projection is sound, and the optimizer compatibility experiments (Tables 2, 6) demonstrate that the approach generalizes across base optimizers. However, the experimental support is limited by short training durations (10K steps), a single model family (LLaMA), single-run reporting without error bars, minimal memory usage data, and missing ablations (GaLore update frequency, normalization in PCA loss, hybrid methods). The central claim that online PCA is better than periodic SVD because it avoids staleness is plausible and supported by the scale-dependent results, but the mechanism is inferred rather than directly tested. A study that varied GaLore's SVD frequency and measured gradient subspace drift over time would provide more direct evidence for the staleness hypothesis. The 7B result (Table 3, 4) is the most practically significant but also the least detailed — the paper does not report the rank, memory, or GaLore configuration for this experiment in sufficient detail for reproduction.

6. Limitations and Trade-offs

Limitation 1: The Convergence Guarantee Requires a Smooth Continuous-Time Limit That May Not Hold for Discrete-Time Training With Large Learning Rates

The assumption or constraint. The paper's central theoretical contribution — that the Hamiltonian+descent structure is preserved under dynamic subspace projection — is derived entirely in continuous time (Section 4.1, Equation 7). The ODE system assumes infinitesimal step sizes, and the Lyapunov argument (d/dt H ≤ 0) relies on derivatives existing everywhere along the trajectory. The discrete-time Algorithm 1 (Section 3.2), which is what practitioners actually implement, uses finite learning rates ε_t^W and ε_t^P, applies weight decay, and operates on stochastic mini-batch gradients rather than full-batch continuous gradients. The paper does not provide a discrete-time convergence analysis, nor does it bound the discretization error between the continuous-time Lyapunov guarantee and the behavior of the discrete algorithm with finite step sizes.

The authors acknowledge this implicitly by framing the analysis as "continuous-time forms provide clearer insights into the asymptotic convergence of the algorithm, abstracting away the choices of step size, discretization, and stochastic errors" (Section 2, Hamiltonian+Descent subsection). But the leap from "the continuous-time system has a Lyapunov function" to "the discrete-time algorithm will converge" is not formally justified in the paper.

The consequence. In the continuous-time analysis, the Lyapunov derivative is strictly non-positive: d/dt H ≤ 0. In discrete time, however, each step of Algorithm 1 may increase the Hamiltonian H(W_t, \hat{S}_t) due to finite step-size effects, even if the continuous-time limit would have guaranteed descent. This is well-known in the optimization literature: discretizations of gradient flows (e.g., forward Euler applied to gradient descent ODEs) can oscillate or diverge when the learning rate exceeds a stability threshold. The paper itself provides suggestive evidence of this failure mode: Section 5.3 notes that "for larger models (350M, 1B), we recommend using a learning rate that is 10 times smaller, specifically 0.001. Larger learning rates cause unrecoverable spikes and instability, a general characteristic observed across all methods." These "unrecoverable spikes" are exactly the behavior one expects when a discretized Lyapunov-stable system is driven outside its stability region by excessive step sizes.

The gap matters practically because practitioners need guidance on maximum safe learning rates for Online Subspace Descent — specifically, how the stability threshold for the coupled (W_t, P_t) system differs from standard full-rank training. The continuous-time theory says the system is stable for any learning rate scheduling in the limit, but the discrete-time reality (as the paper's own hyperparameter guidance acknowledges) is that learning rate choice is critical and poorly chosen rates cause divergence. The recommended value of 0.001 for larger models is empirical; without a discrete-time analysis, there is no principled way to predict stability boundaries when transferring to new model architectures, batch sizes, or datasets.

What evidence exists in the paper. The α-sweep in Figure 3 (middle) shows that large α values (specifically α = 10, which multiplies the projection matrix learning rate relative to the weight learning rate) cause loss spikes and higher final loss compared to α = 1 or α = 5. This is direct evidence of discretization-induced instability: the continuous-time analysis places no restriction on α (any P_t update rule is acceptable), but the discrete implementation shows clear failure at α = 10. The paper's guidance that the base learning rate should be 0.001 for larger models (Section 5.3) is based on observed instability at higher rates, but this tuning was done empirically rather than derived from theory. There is no discrete-time convergence bound or stability condition provided anywhere in the paper.

Mitigation status. The paper does not attempt to address this gap. The authors present the continuous-time analysis as providing a "robust theoretical foundation" (Section 1), but they do not bound the approximation error between the ODE and the discrete algorithm, nor do they derive sufficient conditions on learning rates that guarantee monotonic descent of the Hamiltonian in discrete time. The convergence result (Theorem 4.5) is stated for the continuous-time system (7), and the paper does not claim it holds for Algorithm 1 with finite step sizes. The recommended hyperparameters (Section 5.3, Algorithm 1 remark) are purely empirical — they work on LLaMA models with C4 data, but there is no theoretical guarantee they will work in other settings. This limitation is common to ODE-based optimizer analyses (Maddison et al., 2018; Chen et al., 2023), but the paper does not flag it as such or suggest specific directions for discrete-time analysis.

---

Limitation 2: No Direct Evidence That Staleness (Rather Than Approximation Error) Explains GaLore's Underperformance

The assumption or constraint. The paper's central methodological claim is that online PCA outperforms periodic SVD because continuous updates avoid "stale" projection matrices — that GaLore's P_t reflects the gradient geometry at the last recomputation point (up to 200 steps ago) rather than the current gradient. This staleness hypothesis is stated in Section 5.1: "Online Subspace Descent closes the gap between the state-of-the-art low-rank method and full rank baseline uniformly across different model sizes... We hypothesize that this is due to the higher intrinsic rank of the underlying optimization problem in larger models. Hence, the positive impact on the convergence of the online update of P_t becomes more obvious."

However, the paper never tests this hypothesis directly. The comparison between Online Subspace Descent and GaLore conflates two variables: (1) the update rule (online PCA vs. periodic SVD), and (2) the recency of gradient information (every-step vs. every-200-steps). Online PCA is both more frequent and approximate; GaLore's SVD is both less frequent and exact. If GaLore were re-run with a much shorter SVD interval (e.g., every 10 or 50 steps), would the performance gap close or disappear? If so, staleness is the dominant factor and online PCA is essentially a clever way to achieve frequent updates cheaply. If the gap persists even with frequent SVD, then something else about the online PCA update (its smoothing effect, its stochasticity, its orthonormality penalty) is driving the improvement. The paper cannot distinguish these mechanisms with the experiments provided.

The consequence. Without the missing ablation, practitioners cannot determine the optimal deployment strategy for their specific setting. If staleness is the primary mechanism, then any method that updates P_t frequently (e.g., low-rank SVD approximations, randomized SVD, incremental SVD) should match Online Subspace Descent's performance, and the choice of update rule (PCA loss vs. SVD) is secondary. If online PCA provides benefits beyond just frequency — e.g., smoothing of gradient noise, better orthonormality properties, or implicit regularization from the gradient-based update — then replacing it with frequent approximate SVD would lose these benefits. The paper's results do not tell us which regime we are in.

The staleness hypothesis also makes a testable prediction that the paper does not verify: the gap between Online Subspace Descent and GaLore should grow with the SVD update interval. If GaLore with update gap 50 performs similarly to Online Subspace Descent at the same rank, the staleness explanation is confirmed and the recommendation becomes "use the most frequent updates you can afford," which might not require online PCA at all (frequent low-precision SVD could suffice). Conversely, if even GaLore with update gap 1 (SVD every step, with the associated enormous computational cost) underperforms Online Subspace Descent, then the approximation error of online PCA is actually beneficial — perhaps by smoothing out high-frequency gradient noise — and the mechanism is more subtle than pure staleness.

What evidence exists in the paper. The paper provides two pieces of indirect evidence consistent with (but not probative of) the staleness hypothesis:

1. Scale-dependent gap (Table 1, Figure 1): The perplexity gap between Online Subspace Descent and GaLore grows with model size (60M: near-tie; 350M: ~0.7 point gap; 1B: ~4.2 point gap). The paper argues larger models have higher intrinsic optimization rank, making stale projection more harmful. This is consistent with staleness, but also consistent with online PCA's gradient-based update being better at capturing high-dimensional subspace structure than exact SVD on a single noisy gradient.

2. Wall-clock and per-operation timing (Figure 2, Table 3): Online PCA is dramatically faster per-update, enabling frequent updates without overhead. This explains why frequent updates are feasible with online PCA but not SVD, but does not establish that frequency (rather than the specific update rule) is what matters.

There is no ablation varying GaLore's update frequency, no measurement of how quickly the gradient subspace drifts during training, and no comparison of the subspace alignment between GaLore's and Online Subspace Descent's P_t at corresponding time steps. These would be direct tests of the staleness mechanism.

Mitigation status. Not addressed. The paper states the staleness hypothesis as a plausible explanation (Section 5.1) but does not claim to have verified it. The distinction between "frequent updates matter" and "online PCA specifically matters" is a significant open question that affects both academic understanding (what mechanism drives the improvement?) and practical deployment (what is the simplest sufficient method?). The paper's scope treats this as an empirical demonstration rather than a mechanistic investigation, which is fair but leaves a gap for practitioners who need to know which aspect of Online Subspace Descent to prioritize when adapting it to their use case.

---

Limitation 3: Training Duration Is Too Short to Distinguish Early-Phase From Converged-Phase Dynamics

The assumption or constraint. All pretraining experiments in the paper use 10,000 training steps. For the largest model tested (LLaMA 7B), this represents an extremely short fraction of typical pretraining — full LLaMA pretraining runs for hundreds of thousands to millions of steps. The paper's conclusions about perplexity ranking (Online Subspace Descent outperforms GaLore, narrows gap with full-rank baselines) are based entirely on this early-training window.

The paper does not explicitly discuss this as a limitation, but the context makes the constraint clear: "pretraining LLaMA models ranging from 60M to 7B parameters on the C4 dataset for 10K steps" (Section 5 introduction). The GLUE evaluation results (Table 4) for the 7B model show scores far from typical converged values (e.g., SST2 at 0.52 vs. fine-tuned models routinely exceeding 0.90), confirming that these checkpoints are very early in training.

The consequence. The relative performance of online PCA vs. periodic SVD may change qualitatively over the course of training. Early in training, the loss landscape changes rapidly — weights move far from initialization, gradient magnitudes evolve, and the dominant gradient subspace likely rotates substantially. This is the regime where staleness (Limitation 2) would be most harmful. Later in training, as the model approaches convergence, gradients become smaller and the loss landscape flattens — the gradient subspace may stabilize, making periodic exact SVD more effective and reducing the advantage of continuous updates.

If this is the case, the 10K-step results may overstate Online Subspace Descent's advantage relative to GaLore for full pretraining runs. GaLore might catch up in later training phases when exact SVD on a stable gradient subspace becomes more valuable than approximate but fresh online PCA. Conversely, if gradient subspaces continue to evolve throughout training (due to continued learning of rare token representations, as the paper hypothesizes in Section 5.2), Online Subspace Descent might maintain or even grow its advantage. The current experiments cannot distinguish these scenarios.

The downstream evaluation (Table 4) adds a related concern: the perplexity advantage at 10K steps translates to better GLUE scores, but we don't know whether this advantage persists or whether GaLore-trained models would eventually reach similar downstream performance if trained longer. If GaLore converges more slowly but reaches the same final quality, the practical choice reduces to a wall-clock time tradeoff; if the quality gap is permanent, the choice is clear.

What evidence exists in the paper. The only temporal information comes from the loss curves in Figure 1 (LLaMA 1B validation perplexity across 10K steps) and Figure 3 (LLaMA 60M training loss across 10K steps). Figure 1 shows that the gap between Online Subspace Descent and GaLore appears early (by ~2K steps) and persists through 10K steps, with no sign of closing. The loss curves are still descending at 10K steps — neither method has converged — so the trend could in principle reverse later. There is no evidence either way because training stops at 10K steps.

Mitigation status. Not addressed. The paper does not acknowledge 10K steps as a limitation, does not provide any longer-duration runs (even for the smallest 60M model), and does not discuss whether the results should be expected to generalize to full training durations. The recommended hyperparameters (Section 5.3) are tuned on 10K-step runs and may not be optimal for longer training. This is a significant limitation for practitioners considering adopting Online Subspace Descent for production pretraining, where training typically runs for orders of magnitude more steps. The cost of running 10K-step experiments on a 7B model is noted (Table 3: 7.14 hours), but extending to, say, 100K steps would be ~71 hours — feasible but not done.

---

Limitation 4: Difficulty Estimation and Allocation Are Not Dynamic — The Method Lacks a Mechanism for Adapting Strategy Within a Training Run

The assumption or constraint. The paper presents Online Subspace Descent as a uniform strategy: the same update rule for P_t (online PCA with a fixed optimizer and fixed hyperparameters α, λ, learning rate) is applied at every step of training from initialization to completion. There is no mechanism for dynamically adjusting the rank, the P_t update frequency, the choice of optimizer for P_t, or the relative learning rate α based on the training phase, gradient statistics, or observed loss behavior.

This stands in contrast to the reference example paper on compute-optimal test-time scaling, where the central contribution was adaptive allocation of compute based on estimated problem difficulty. A natural extension of Online Subspace Descent would be to adapt the projection strategy over the course of training: use higher rank or more frequent P_t updates early in training when the gradient geometry shifts rapidly, and reduce rank or update frequency later when the subspace stabilizes. The paper's theoretical framework — which proves convergence for any P_t update rule — actually makes such adaptation straightforward to analyze, but the paper does not explore this.

The consequence. The fixed-strategy approach leaves potential efficiency gains on the table. Early in training, the gradient subspace may be changing so rapidly that even online PCA with α = 5 cannot keep up — a higher adaptation rate might be warranted. Late in training, the projection matrix may be essentially converged, and continuing to update it every step wastes computation that could be redirected to weight optimization (or the rank could be reduced to save memory). Without any adaptation mechanism, Online Subspace Descent applies the same compute budget to P_t updates regardless of whether those updates are providing value.

The paper's hyperparameter recommendations (rank as high as memory allows, α = 5, λ = 0.1, learning rate 0.001 for larger models) are static values that must be chosen before training begins. If the optimal hyperparameters change during training, these fixed choices will be suboptimal in some phases. The rank sweep (Figure 3, left; Table 5) shows that the marginal benefit of higher rank diminishes — at some point, memory spent on additional rank would be better spent elsewhere (e.g., larger batch size, larger model dimension). But without a dynamic rank schedule, the practitioner must commit to a fixed rank for the entire run.

What evidence exists in the paper. There is no direct evidence of this limitation because the paper never varies strategy during a training run. However, the scale-dependent results provide indirect evidence that the "right" strategy depends on the optimization phase. The observation that larger models show a bigger gap between Online Subspace Descent and GaLore (Table 1) suggests that the benefit of continuous updates grows with optimization problem complexity, which typically decreases as training progresses (early training is more complex/higher-rank than later training). This implies that the relative value of online PCA over SVD may diminish late in training, making a static strategy potentially wasteful.

The paper also notes (Section 3.2) that P_t updates "can be done in parallel with that of (W_t, \hat{S}_t) and incurs no slowdown once it is fast enough to not cause a speed bottleneck." But if P_t becomes unnecessary or nearly converged, even parallel computation is wasted energy that could be saved.

Mitigation status. Not addressed, and the paper does not frame this as a limitation or suggest adaptive strategies as future work in its conclusion. The conclusion mentions questions about "alternative methods for updating projection matrix that could accelerate convergence" (Section 7), which implicitly includes the possibility of adaptive methods, but does not specifically call out dynamic rank or dynamic update frequency as directions. This is a missed opportunity given the theoretical framework's flexibility — the Lyapunov analysis in Section 4.1 holds for any Γ, including time-varying and state-dependent update rules, making Online Subspace Descent exceptionally well-suited for adaptive strategies that would be harder to justify theoretically for methods without convergence guarantees.

---

Limitation 5: Memory Usage Measurements Are Essentially Absent From the Paper

The assumption or constraint. The paper's primary motivation is memory efficiency: "Adam produces the best results for LLM pretraining, but its memory cost is prohibitive... One key challenge is to retain the high performance of Adam while enhancing its memory efficiency" (Section 2). The proposed solution reduces optimizer state memory from O(nm) to O(km) per weight matrix, a substantial reduction when k ≪ n. Yet the paper provides almost no actual memory measurements.

The only memory number in the entire paper is a single, brief mention in Section 5.1: "Online Subspace Descent attains significant improvement over GaLore in perplexity, while consuming a similar amount of GPU memory (8.64 GB v.s 9.01 GB)" for the LLaMA 1B with sequence length 256 experiment. There is no:

• Memory usage for 60M, 350M, or 7B models.
• Breakdown of memory by component: model weights, W_t optimizer states, P_t, P_t optimizer states, activations, gradients.
• Comparison to full-rank AdamW8bit memory usage at any model scale.
• Discussion of how memory scales with rank (the main tunable parameter).
• Analysis of memory overhead from maintaining a separate optimizer for P_t.

The consequence. A practitioner reading this paper to decide whether to adopt Online Subspace Descent cannot answer the most basic deployment question: how much memory will this save me compared to full-rank training? The paper makes qualitative arguments (projected states are k × m instead of n × m) but never quantifies the savings for the specific LLaMA configurations tested. This is a critical gap because the memory savings depend heavily on the architecture: in a transformer, not all weight matrices have the same dimensions, and the n × m vs. k × m reduction applies only to the larger dimension while the smaller dimension remains unchanged. For some layers (e.g., attention projections where n and m are similar), the relative savings may be modest; for others (e.g., feed-forward layers where one dimension is 4× the other), they may be substantial. Without per-layer or per-model memory breakdowns, the actual memory reduction is unknown.

Furthermore, the paper recommends using Adam for P_t (Algorithm 1 remark), which means an additional 2 × n × k values for the projection matrix optimizer states (M_t^P and V_t^P). For large n and moderate k, this could be non-negligible. The paper never quantifies this overhead or compares it to the savings from projected W_t states. The finding that SGD on P_t works (Table 6: 53.53 perplexity) is practically important precisely because it eliminates this overhead, but the memory implications are never discussed.

The headline comparison of "8.64 GB vs. 9.01 GB" (Section 5.1) is puzzling: Online Subspace Descent uses less memory than GaLore despite maintaining an additional optimizer state for P_t. This discrepancy is never explained. It could be due to differences in activation checkpointing, gradient accumulation, or batch size, but the paper provides no details. A reader cannot determine whether the 0.37 GB difference is meaningful, reproducible, or an artifact of measurement.

What evidence exists in the paper. Almost none. The single memory datapoint (8.64 GB vs. 9.01 GB) appears in prose without a table or figure, without specifying what is included in the measurement (GPU peak memory? Allocated memory? Including or excluding activations?), and without comparison to full-rank AdamW8bit memory for the same configuration. The paper discusses memory complexity in asymptotic terms (Section 3.1: "much smaller in size compared to W" when k ≪ n) but never converts these asymptotic claims to concrete numbers for the models and ranks actually tested.

The paper also never reports whether the online PCA P_t updates were actually run in parallel with weight updates (as claimed possible in Section 3.2 and Figure 2 caption) or sequentially. If run sequentially, they add to peak memory (gradients for the PCA loss must be materialized). If the parallelism claim is critical to the "no overhead" argument, the implementation details matter.

Mitigation status. Not addressed at all. Memory efficiency is the central motivation (Section 1, Section 2, Section 3), making the absence of systematic memory measurements a significant gap. The paper's conclusion (Section 7) does not mention this as a direction for future work, and the experimental section does not acknowledge the omission. This is the most straightforward-to-fix limitation — adding a table of memory usage by model size, rank, and method would substantially strengthen the paper's practical value.

---

Limitation 6: All Experiments Use a Single Model Architecture (LLaMA) on a Single Dataset (C4) — No Evidence of Generalization to Other Architectures, Domains, or Training Paradigms

The assumption or constraint. Every experiment in the paper uses the LLaMA decoder-only transformer architecture pretrained on the C4 text corpus. The paper does not test Online Subspace Descent on encoder-decoder models (T5), encoder-only models (BERT), vision transformers, diffusion models, or any non-language domain. It also does not test on fine-tuning tasks, only pretraining. The optimizer compatibility experiments (Tables 2, 6) vary the base optimizer (LION, Adafactor, SGD) but keep the architecture and dataset fixed.

The paper does not explicitly acknowledge this as a limitation, but the scope is clear from the experimental section: "pretraining LLaMA models ranging from 60M to 7B parameters on the C4 dataset" (Section 5). The theoretical framework (Section 4) places no restrictions on architecture or data domain — the Hamiltonian+descent analysis applies to any setting where gradients exist — but the empirical validation is narrowly scoped.

The consequence. The paper's central empirical claim — that online PCA outperforms periodic SVD for memory-efficient training — may depend on properties specific to transformer language models trained on text data. Several mechanisms could cause the advantage to shrink or disappear in other settings:

• Gradient rank structure: The low-rank gradient assumption underlying all subspace descent methods (GaLore, Sketchy, Online Subspace Descent) is an empirical observation about deep learning in general (Gur-Ari et al., 2018), but the dynamics of the gradient subspace — how quickly it rotates, how its rank evolves during training — may differ across architectures and domains. Language modeling on C4 involves learning both high-frequency patterns (common token sequences, captured by low-rank projections) and rare tokens (requiring higher rank, as the paper hypothesizes in Section 5.2). In contrast, training on ImageNet with a vision transformer may have different gradient rank profiles, and the relative advantage of continuous vs. periodic updates could shift accordingly.

• Batch size and gradient noise: The online PCA update uses a single mini-batch gradient G_t to compute the PCA loss L_{G_t}(P). If mini-batches are small, G_t is a noisy estimate of the true gradient, and the online PCA update may track the noise rather than the signal. The paper uses batch size 512 (Appendix A.1), which provides moderate gradient signal. For training paradigms with very small batch sizes (e.g., fine-tuning with batch size 8), the online PCA update may become dominated by stochasticity, potentially making periodic SVD (which aggregates gradient information across multiple steps, either explicitly or implicitly through less frequent recomputation) more robust.

• Training duration: As discussed in Limitation 3, the 10K-step regime may not represent full training dynamics. But more importantly, the optimal P_t update strategy may differ between the short training runs in this paper and the long training runs typical in practice. If the gradient subspace stabilizes after 50K–100K steps (as some evidence suggests; Gur-Ari et al., 2018), spending compute on continuous P_t updates for the remaining 400K+ steps would be wasteful — a practitioner would want to switch to a fixed or infrequently-updated projection matrix at that point. The paper's fixed strategy makes no provision for this.

What evidence exists in the paper. The only architectural variation tested is model scale within the LLaMA family (60M, 350M, 1B, 7B). The optimizer compatibility experiments (Tables 2, 6) are the closest the paper comes to domain variation — they show that Online Subspace Descent works with LION, Adafactor, and SGD, not just Adam. This provides evidence that the method is robust to optimizer choice, but says nothing about robustness to architecture or domain.

The theoretical framework (Section 4) provides some reassurance that the method should generalize: the Lyapunov analysis makes no assumptions about the loss landscape beyond differentiability, and the adjoint cancellation works for any linear operator P_t. However, this only guarantees that the optimizer won't diverge — it doesn't guarantee that online PCA will be better than periodic SVD in any given setting. The relative advantage could be domain-specific.

Mitigation status. Not addressed. The paper's conclusion (Section 7) asks "Can this method be applied to problems beyond language modeling?" as an open question for future work, implicitly acknowledging the domain limitation. But this question is framed as an extension rather than as a missing validation of the current claims. The paper does not test on even a single non-language or non-LLaMA setting, making the generalizability of the empirical results entirely unknown. Given that the paper's practical contribution is a training algorithm — not a theoretical result — this narrow empirical scope limits the confidence with which practitioners in other domains (vision, speech, scientific ML) can adopt the method based on this paper alone.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper changes the conversation around memory-efficient LLM training by providing the first theoretical guarantee that dynamic subspace projection is sound, not just a heuristic trick that happens to work. The significance of this shift is not that the field lacked empirical evidence that methods like GaLore work — it had that — but that the mechanisms underlying their success were poorly understood, making it impossible to predict when they would fail, how to improve them, or which design choices were load-bearing versus incidental. The paper's Lyapunov analysis (Section 4.1) resolves this by proving that the Hamiltonian+descent structure — the mathematical property that guarantees convergence for Adam, LION, momentum, and other common optimizers — survives completely unchanged when gradients are projected through a dynamically evolving subspace, regardless of how that subspace is updated. The cancellation that ensures this is structural (adjoint symmetry), not approximate or asymptotic, meaning the guarantee is robust to essentially any smooth P_t update rule.

This is best characterized as a fundamental reframing with practical consequences, not a paradigm shift. The paper does not displace full-rank training, nor does it claim that subspace descent is universally preferable. What it does is convert a class of methods from "promising but theoretically opaque" to "provably convergent under mild conditions." The practical upshot is that researchers and practitioners can now experiment with arbitrary P_t update rules — online PCA, randomized SVD, sketching with exponential forgetting, learned controllers — and be guaranteed that the projection mechanism itself will not cause divergence. Previously, each new update rule would have required its own ad hoc stability analysis (or, more commonly, none at all); now, the analysis is effectively off-the-shelf, requiring only verification that the P_t rule does not permanently trap the optimizer in a degenerate configuration (Assumption 4.4(ii), which is satisfied by any method that seeks to reconstruct the gradient).

The paper also shifts the design focus from per-step optimality to streaming adaptation. GaLore and similar methods implicitly optimize for the quality of gradient reconstruction at a single point in time (exact SVD of the gradient snapshot). The paper's results — online PCA outperforming exact SVD at equal rank despite being a worse approximation at each individual step (31.30 vs. 35.52 perplexity on LLaMA 1B, Table 1) — force a reconceptualization: fresh but approximate gradient information is more valuable than stale but exact information for tracking the evolving optimization landscape. This reframing will likely influence subsequent work to focus on streaming algorithms and online updating schemes rather than ever-more-accurate batch decompositions, and it opens a design space that was previously invisible: what is the optimal tradeoff between accuracy and recency for the projection matrix?

The work also resolves a specific, named contradiction in the prior literature. Stochastic Subspace Descent (Kozak et al., 2019) provided convergence analysis but only for convex objectives with random (non-adaptive) projection matrices. GaLore and Sketchy provided no convergence analysis. The field thus had a gap: dynamic projection for stateful optimizers on non-convex objectives had no theoretical foundation, and practitioners relied entirely on empirical validation. The paper fills this gap by unifying the analysis under the Hamiltonian descent framework, showing that the previous negative intuition — that momentum across changing subspaces creates entanglement that could break convergence — was mistaken. The adjoint cancellation makes the projection operator effectively invisible to the Lyapunov function, and this holds for any optimizer with the Hamiltonian+descent structure, including the Adam, LION, and momentum variants that dominate LLM pretraining.

Finally, the paper makes the case for online PCA as a practical alternative to SVD in the specific, high-impact context of LLM training, where the 142× per-operation speed advantage (Figure 2) translates to a 1.3× end-to-end wall-clock speedup at 7B scale (Table 3) while simultaneously improving convergence quality. This dual advantage — faster and better — is rare in optimization methods and makes the practical case for adoption unusually strong, at least for the settings tested.

Follow-Up Research This Work Enables

Directly testing the staleness hypothesis by varying GaLore's SVD update frequency. The paper attributes GaLore's underperformance to stale projection matrices, but never isolates this variable (Limitation 2 in Section 6). A natural follow-up would train identical LLaMA models with GaLore at update gaps ranging from 1 (SVD every step, at enormous cost) to 500 steps, measuring final perplexity and wall-clock time. If GaLore with frequent updates matches or exceeds Online Subspace Descent, staleness is confirmed as the dominant mechanism, and the practical recommendation becomes "use approximate but frequent updates" (online PCA, incremental SVD, or randomized SVD) rather than specifically requiring online PCA. If GaLore with update gap 1 still underperforms, then something about online PCA's update dynamics — its gradient-based optimization, its orthonormality penalty, or its implicit smoothing — provides benefits beyond recency, and understanding that mechanism becomes a priority.

Characterizing gradient subspace drift rates during LLM pretraining. The paper's central mechanism — that timely gradient information matters more for larger models — is hypothesized but never measured. A follow-up study could instrument training runs to track the principal angle between the dominant gradient subspace at time t and at time t + Δ for various Δ, across model scales (60M → 7B) and training phases (early, middle, late). The key measurements: (a) does subspace drift accelerate or decelerate during training? (b) does drift rate correlate with model size, as the intrinsic-rank hypothesis predicts? (c) what is the half-life of subspace relevance — i.e., after how many steps does a past SVD of the gradient become less predictive of the current gradient than a cheap online PCA update? Answering these questions would convert the paper's qualitative hypothesis into quantitative engineering guidance: given a compute budget for projection updates, allocate it to match the drift timescale.

Adaptive rank and update frequency schedules. The paper's Lyapunov analysis (Section 4.1) holds for any time-varying P_t update rule, but the experiments use a fixed rank and fixed α throughout training. A natural extension — explicitly enabled by the theoretical framework — would be a dynamic rank schedule: high rank early in training when the loss landscape changes rapidly (spending more memory on optimizer states), and progressively lower rank as the gradient subspace stabilizes (freeing memory for larger batch sizes or model dimensions). A concrete experiment: train LLaMA 1B with rank decaying from 512 to 128 over the first 30K steps according to a schedule informed by subspace drift measurements, comparing memory-weighted perplexity against fixed-rank baselines. The ability to prove convergence for a time-varying rank (which is just a special case of time-varying P_t) is a direct consequence of the paper's theoretical contribution and would be much harder to justify for methods without this guarantee.

Distillation of the learned projection matrix into a static projection for deployment. The paper's framework produces a projection matrix P_t that tracks the dominant gradient directions throughout training. At the end of training, the final P_T captures the gradient geometry at the converged optimum. A follow-up could investigate whether this P_T is useful for downstream fine-tuning or inference: can fine-tuning within the subspace spanned by P_T (LoRA-style, fixing P = P_T and only training \hat{W}) match or exceed the performance of randomly initialized low-rank adapters? If so, the training process effectively discovers the "intrinsic task subspace" for free, and Online Subspace Descent could be integrated into a train-once, fine-tune-many workflow where the projection matrix learned during pretraining is reused across many downstream tasks. A concrete experiment: pretrain LLaMA 60M to convergence with Online Subspace Descent, extract P_T, then fine-tune on GLUE tasks using only \hat{W} optimization (static P_T), comparing against standard LoRA with random initialization of both factors.

Extension to non-language domains with different gradient rank profiles. All experiments use LLaMA on C4. The theoretical framework makes no domain-specific assumptions, but the empirical advantage of online PCA over periodic SVD may depend on the gradient rank dynamics of language modeling specifically. A stress-test: apply Online Subspace Descent to vision transformer pretraining on ImageNet-21k, where gradient rank profiles may differ (vision models often show sharper rank decay). If the advantage persists, the method generalizes; if it disappears or reverses, the domain conditions under which continuous updates matter become clearer. A second stress-test: fine-tuning on small datasets (e.g., 1K–10K examples) where gradient subspaces may be dominated by noise. The paper's batch size 512 provides substantial gradient signal; with tiny batches, online PCA may track noise and perform worse than infrequent SVD. Establishing these boundary conditions is essential for practical adoption.

Optimal P_t update rules via the optimal control formulation the paper identifies. Section 4.1 explicitly notes that "an interesting direction for future investigation is to find optimal rules of P_t to maximize the decreasing rate as an optimal control problem." This is a concrete theoretical opening: formulate the choice of Γ in d/dt P_t = Γ(P_t, ∇L(W_t)) as a control policy that maximizes -d/dt H(W_t, \hat{S}_t), subject to a computational budget constraint on Γ. The reward function is the instantaneous energy dissipation rate (which the paper provides closed-form expressions for in Examples 4.1–4.3), and the action space is the set of possible P_t update directions. A follow-up could solve this control problem in simplified settings (quadratic objectives, linear gradient dynamics) and test whether the resulting Γ outperforms online PCA in full LLM pretraining. This would bridge the paper's theoretical framework with its practical algorithm, potentially producing a projection update rule that is provably optimal rather than heuristically motivated.

Practical Applications and Downstream Use Cases

Cost-sensitive LLM pretraining on consumer or single-GPU hardware. The paper's 1.3× wall-clock speedup at 7B scale (Table 3) combined with lower perplexity means that researchers and small organizations with limited GPU budgets can train better models in less time. Concretely: a team with a single 8×A100 node pretraining a 7B model can save ~2.6 hours per 10K steps while getting better convergence. Extended to full pretraining runs (hundreds of thousands of steps), this translates to days of wall-clock time saved and correspondingly lower cloud compute costs. The finding that SGD on P_t works (Table 6, 53.53 perplexity) provides an additional "budget" option: if memory is so tight that even the small Adam state for P_t is problematic, SGD eliminates this overhead entirely while still substantially outperforming GaLore (48.05). This makes online subspace descent deployable in extremely memory-constrained settings (edge devices, older GPUs with limited VRAM) where full Adam states are impossible and even GaLore's SVD overhead is burdensome.

Memory-efficient fine-tuning that reuses pretraining projection structure. Although the paper evaluates only pretraining, the method naturally extends to fine-tuning: initialize P_0 from the final projection matrix of a pretraining run (or a larger model's projection matrix, distilled), then fine-tune \hat{W} with P_t either frozen or slowly updated. This would provide the memory efficiency of LoRA-style adapters with the flexibility of full-model updates (since W_t = W_0 + \sum P_\tau \hat{\Delta}_\tau spans the full parameter space over time). The theoretical guarantee that the system converges regardless of how P_t is updated means the fine-tuning practitioner can experiment freely with different P_t freezing and update schedules without worrying about breaking convergence. A concrete deployment: fine-tune LLaMA 7B on a domain-specific corpus with rank 512, using ~9 GB GPU memory instead of the ~20+ GB required for full-rank Adam, on a single consumer GPU (RTX 3090 with 24 GB). The paper's 1B experiment shows 8.64 GB GPU usage at rank 512 with SS 256; a 7B model with the same rank but longer sequence length should remain feasible on 24 GB cards, making domain adaptation of large models accessible to individual researchers.

Training data generation and self-improvement pipelines with memory-constrained worker nodes. When using LLMs to generate training data for distillation or self-improvement (STaR, ReST^EM), the teacher model is often much larger than the student, and running full-rank Adam on the teacher for fine-tuning between data generation rounds may be memory-prohibitive. Online Subspace Descent enables the teacher to be fine-tuned with substantially reduced memory, allowing the self-improvement loop to run on the same hardware that serves the generation step. Since the theoretical framework guarantees convergence for any base optimizer (not just Adam), the same method can be applied with LION or Adafactor for additional memory savings during the fine-tuning phases, using the mixed-optimizer configurations shown in Tables 2 and 6. The continuous nature of the projection update is particularly well-suited to online self-improvement, where the data distribution shifts each round and the gradient geometry must adapt quickly — exactly the regime where the paper shows online PCA outperforms periodic SVD.