ArXiv: 2412.05270
🎯 Pitch
APOLLO matches AdamW’s LLM training quality while using almost no optimizer state memory—its extreme variant outperforms AdamW with only a rank‑1 random projection. This enables 3× higher throughput via larger batches and for the first time lets you pre‑train LLaMA‑7B from scratch on a single GPU with under 12 GB of memory.
1. Executive Summary
This paper proposes APOLLO (Approximated Pradient Scaling for Memory Efficient LLOM Optimization), a memory-efficient optimizer for LLM training that replaces AdamW’s element-wise learning rate adaptation with structured channel-wise or tensor-wise gradient scaling approximated in an auxiliary low-rank space using pure random projections. The authors demonstrate on LLaMA models (60M to 7B parameters) pre-trained on the C4 dataset that APOLLO matches or surpasses AdamW’s validation perplexity—achieving up to a 2.8× reduction in perplexity for the 1B model—while reducing optimizer state memory to SGD-like levels, and that its extreme variant APOLLO-Mini uses only a rank-1 auxiliary subspace to deliver better-than-AdamW performance at a near-zero optimizer memory cost. The memory savings translate into a ~3× throughput improvement on 8×A100-80GB by supporting 4× larger batch sizes than AdamW, and APOLLO-Mini combined with INT8 weight quantization enables pre-training LLaMA-7B from scratch on a single GPU with under 12 GB of memory, establishing that structured, low-rank gradient scaling is sufficient for LLM optimization only when the scaling factor is approximated via random projection rather than recovered through costly SVD.
2. Context and Motivation
The Core Problem: AdamW's Memory Overhead Cripples LLM Training Scalability
Training large language models requires an optimizer, and the de facto standard is AdamW (Loshchilov & Hutter, 2019). AdamW works by maintaining two per-parameter state tensors: the first moment (exponential moving average of gradients) and the second moment (exponential moving average of squared gradients). These states, combined with the weight parameters and gradients, effectively triple the memory requirement relative to just storing the model parameters. The paper quantifies this starkly in the introduction:
"For instance, training a LLaMA-7B model from scratch requires at least 58 GB of memory, with 28 GB devoted to AdamW's optimizer states... For larger models like GPT-3, with 175 billion parameters, memory demands reach 700 GB for the model alone, leading to a staggering 1.4 TB requirement for AdamW's optimizer states."
This memory burden creates a cascade of practical problems that the paper identifies as the central motivation:
- Hardware barriers: Researchers must either use more GPUs, higher-end GPUs with larger memory capacities (e.g., A100-80GB, H100), or both. For many in the academic and open-source communities, such hardware is simply inaccessible.
- Throughput vs. memory trade-off: When memory is tight, practitioners are forced to reduce batch sizes. Smaller batch sizes mean less efficient GPU utilization—the model becomes memory-bound rather than compute-bound—leading to lower training throughput and longer wall-clock times.
- Scaling overhead: Adding more GPUs to increase aggregate memory introduces non-trivial communication and infrastructure overheads (Jiang et al., 2024), particularly with distributed data parallel (DDP) training. The paper specifically demonstrates that AdamW cannot even fit a LLaMA-13B model on a single A100-80GB GPU with naive DDP, requiring system-level optimizations like model sharding that add engineering complexity.
- Democratization gap: The memory requirements create a two-tier research landscape where only well-resourced labs can pre-train large models from scratch, while others are limited to fine-tuning existing checkpoints.
The paper frames this not as an academic curiosity but as a practical bottleneck with direct economic and accessibility consequences. The figure of 1.4 TB for GPT-3's optimizer states alone makes clear that this problem scales with—and in some sense outpaces—model size growth.
Where Prior Approaches Fall Short
The paper identifies two broad categories of prior work on reducing training memory, each with specific, well-documented limitations:
1. Reducing the number of trainable parameters
This category includes designing smaller model architectures (Liu et al., 2024b; Tang et al., 2024), sparse training (Liu et al., 2022; Thangarasa et al., 2023), and low-rank adaptation methods like LoRA (Hu et al., 2021) and its pre-training variant ReLoRA (Lialin et al., 2023). The paper acknowledges these approaches "effectively reduce memory usage" but are fundamentally limited because they restrict the optimization space of model parameters. The result is a persistent performance gap, particularly in pre-training scenarios. The paper cites Biderman et al. (2024) and Lialin et al. (2023) for evidence that these methods "often result in performance trade-offs... particularly during pretraining." The empirical results in Table 2 bear this out: LoRA and ReLoRA on LLaMA-60M achieve perplexities of 34.99 and 37.04 respectively, while full AdamW reaches 34.06—and the gap widens dramatically at larger scales (e.g., Low-Rank on LLaMA-1B hits 142.53 perplexity vs. AdamW's 15.56). The paper thus establishes that parameter-reduction approaches are not a viable path to matching AdamW's optimization quality.
2. Memory-efficient optimizers that reduce optimizer state memory
This is the more relevant family of prior work, and the paper spends significant space characterizing its limitations. The key methods and their specific shortcomings are:
GaLore (Zhao et al., 2024): GaLore was a breakthrough in enabling full-parameter LLM training with reduced memory by projecting gradients into a low-rank space and performing optimizer updates on the low-rank representation. The gradient subspace is identified using Singular Value Decomposition (SVD). However, the paper identifies three specific problems with GaLore:
- Costly SVD operations: Periodic SVD updates (e.g., every 200 iterations) incur a computational cost of , which becomes prohibitive for large matrix dimensions. The paper quantifies this concretely: "For instance, in the case of the LLaMA-7B model, a single subspace update can take approximately 10 minutes, whereas inference only takes seconds." Figure 9 visualizes this with dramatic throughput spikes downward at every SVD update step. This overhead "significantly reduces training throughput" and makes GaLore less practical than its memory savings would suggest.
- High rank requirement: GaLore needs a relatively high rank (typically one-quarter of the original dimension, denoted ) to maintain competitive performance. This means the optimizer state memory, while reduced, is still substantial—the paper notes GaLore stores optimizer states (Table 1), which at is still significant relative to SGD's near-zero overhead.
- Sensitivity to projection quality: Figure 5(a–c) shows that GaLore's performance degrades significantly when SVD is replaced with random projection, meaning the expensive SVD is not an optional component but a requirement for GaLore to work effectively.
Fira (Chen et al., 2024): Fira enhances GaLore by incorporating the error residual between the full-rank gradient and its low-rank approximation, effectively simulating full-rank updates. While this improves performance over GaLore (Table 2 shows Fira achieving 31.06 vs. GaLore's 34.88 perplexity on LLaMA-60M), it inherits the same SVD bottleneck and, as the paper will later demonstrate (Section 5.4, A4; Figure 6), Fira's advantage over APOLLO diminishes with larger models and more training tokens.
Flora (Hao et al., 2024): Flora compresses gradients using low-rank adapters. The paper notes it "exhibit[s] noticeable performance gaps compared to AdamW" and, critically from Table 1, is marked as not supporting pre-training, limiting its scope to fine-tuning scenarios.
Adam-mini (Zhang et al., 2024b): Adam-mini takes an orthogonal approach by demonstrating that block-wise (rather than element-wise) second moments suffice for learning rate adjustments. This is an important conceptual precursor to APOLLO's insight about structured learning rates, but the paper identifies a key limitation: "achieving performance on par with AdamW requires careful handling of different model components to maintain compatibility with its optimization dynamics." More concretely, Adam-mini only reduces memory for the second moment (variance), leaving the first moment's memory footprint untouched. The paper characterizes this as "substantial memory overhead of optimization states in order to maintain competitive performance (e.g.,... full-rank first momentum in Adam-mini)."
GaLore variants with alternative projections: Recent work explored replacing SVD with online PCA (Liang et al., 2024) or random projections (He et al., 2024), "yielding provable convergence." However, the paper notes that these still require SVD during early training stages (He et al., 2024), making them only partial solutions to the SVD bottleneck.
A Unifying Gap: The Missing Synthesis of Structured Learning Rates and Low-Rank Projections
The paper identifies a conceptual gap between two idea streams that had developed largely independently. On one side, the low-rank gradient projection literature (GaLore, Fira, Flora, etc.) recognized that gradients occupy a low-dimensional subspace and can be compressed, but relied on expensive SVD to find that subspace and needed relatively high ranks to preserve optimizer quality. On the other side, the optimizer redundancy literature (Adam-mini) recognized that AdamW's element-wise learning rate adaptation is unnecessarily fine-grained and can be coarsened, but only partially addressed memory (second moment only) and required careful component-specific handling.
The paper's key motivating insight is that these two ideas are complementary and can be unified: if learning rate adaptation can be coarsened to the channel or tensor level, then the coarsened scaling factors can potentially be estimated from a much lower-rank gradient representation without needing the precision that GaLore and Fira require. This synthesis directly addresses the limitations of both streams—the SVD requirement (because coarse scaling factors are robust to projection quality) and the partial memory savings (because both first and second moments get compressed).
How This Paper Positions Itself
The paper's positioning is explicit and multi-layered:
Relative to GaLore and Fira: APOLLO is presented as a replacement that eliminates the SVD bottleneck entirely. The paper emphasizes this repeatedly: "APOLLO achieves efficient memory usage entirely without relying on SVD while delivering performance that matches or even surpasses AdamW." Table 1 formalizes this by marking which methods require SVD—GaLore and Fira do; APOLLO and APOLLO-Mini do not. The paper also positions APOLLO as achieving greater memory savings than GaLore, noting in the abstract that it "almost eliminat[es] the optimization states in AdamW" and in Section 4.3 that APOLLO-Mini's optimizer state cost of is "comparable to the memory footprint of SGD."
Relative to Adam-mini: APOLLO extends the structured learning rate insight from second-moment only to both first and second moments. The paper describes APOLLO as applying "adaptive learning rates in a structured channel-/tensor-wise manner and go[ing] beyond Adam-Mini by reducing memory usage for both first- and second-order moments, even eliminating optimizer memory in APOLLO-Mini" (Section 5.5).
Relative to the broader memory-efficient training landscape: The paper positions APOLLO as orthogonal to and compatible with system-level optimizations. It explicitly states that APOLLO "can be seamlessly integrated to achieve greater memory efficiency" with activation checkpointing, quantization, and memory offloading (Section 2.2). This is demonstrated concretely in Section 5.3, where Q-APOLLO-Mini (APOLLO-Mini combined with INT8 weight quantization) achieves LLaMA-7B pre-training in under 12 GB. This compatibility positioning is strategic: it means APOLLO is not competing with system-level approaches but amplifying them.
The "SGD-like memory, AdamW-level performance" framing: The paper's central positioning slogan—reflected in the title—is that APOLLO achieves the memory footprint of SGD (which stores no optimizer states beyond the gradient buffer) while maintaining or exceeding AdamW's optimization quality. This is a deliberately provocative framing because SGD is well-known to fail at training Transformer-based models (Zhang et al., 2024a). The paper acknowledges this tension in Section 5.5 and provides preliminary theoretical hypotheses for why APOLLO's structured, low-rank scaling succeeds where pure SGD fails.
The democratization narrative: Throughout the paper, there is a consistent undertone of democratizing LLM training. The abstract highlights enabling "pre-training LLaMA-13B model with naive DDP on A100-80G without requiring other system-level optimizations" and "training of LLaMA-7B from scratch on a single GPU using less than 12 GB of memory." The conclusion frames this as "democratiz[ing] access to LLM pre-training, making it accessible to a broader range of researchers and organizations with limited resources." This is not merely rhetorical—the memory numbers represent concrete thresholds (single consumer GPU, single datacenter GPU without model parallelism) that define who can participate in LLM research.
3. Technical Approach
3.1 Reader orientation (approachable technical breakdown)
APOLLO is a drop-in replacement optimizer for AdamW that trains large language models with the memory footprint of SGD while matching or exceeding AdamW's final model quality. The problem it solves is that AdamW maintains two full-size floating-point arrays per parameter—the first and second moment estimates—which triples the memory needed to store the model; APOLLO solves this by first recognizing that AdamW's per-parameter learning rate adjustment can be coarsened to a per-channel or even per-tensor scalar without hurting training, and then approximating those coarse scaling factors from a heavily compressed, randomly projected low-dimensional view of the gradient that costs almost nothing to store.
APOLLO's twin insights are that (1) the precise, element-wise gradient scaling that AdamW performs (dividing each parameter's update by the root-mean-square of its past gradients) contains substantial redundancy and can be replaced by a single scaling factor shared across all parameters in a channel or tensor, and (2) once the scaling factor is coarsened, the optimizer state that computes it can be maintained in a tiny auxiliary space—as small as rank-1—using pure random projection without any expensive matrix decomposition, because the scaling factor's aggregate statistics are robust to projection noise while per-element updates are not.
3.2 Big-picture architecture (diagram in words)
The APOLLO system has four major components that execute sequentially at each optimizer step:
1. Random projection generator: At configurable intervals (default every 200 steps), a fresh random matrix $P_t \in \mathbb{R}^{r \times m}$ is sampled from a Gaussian distribution $\mathcal{N}(0, 1/r)$, where $r$ is the target low rank (e.g., 128 for a 512-dimensional weight matrix) and $m$ is the smaller dimension of the weight matrix $W_t \in \mathbb{R}^{m \times n}$. The seed for this random matrix is stored, making the projection matrix itself memory-free—it is regenerated on-the-fly during forward/backward computation rather than stored.
2. Gradient projection: The full-rank gradient $G_t \in \mathbb{R}^{m \times n}$ computed during backpropagation is left-multiplied by $P_t$ to produce a low-rank projected gradient $R_t = P_t G_t \in \mathbb{R}^{r \times n}$. This is a single matrix-matrix multiply (GEMM) that compresses the gradient from $m$ rows to $r$ rows, where $r \ll m$. The full-rank gradient is retained for the actual weight update.
3. Low-rank optimizer state maintenance: The projected gradient $R_t$ is fed into a standard AdamW-style optimizer that maintains first and second moment estimates $M^R_t, V^R_t \in \mathbb{R}^{r \times n}$, but now these states are $r \times n$ rather than $m \times n$, reducing memory by a factor of $m/r$. From these low-rank moments, a normalized gradient $\tilde{R}_t = M^R_t / \sqrt{V^R_t + \epsilon}$ is computed in the compressed space.
4. Gradient scaling and weight update: The channel-wise (or tensor-wise) gradient scaling factor $s \in \mathbb{R}^{1 \times n}$ is computed from the low-rank normalized gradient $\tilde{R}_t$ and the raw projected gradient $R_t$ as the ratio of their per-channel $\ell_2$ norms: $s_j = \lVert\tilde{R}_t[:, j]\rVert_2 / \lVert R_t[:, j]\rVert_2$. This scaling factor is then applied to the original full-rank gradient $G_t$ via element-wise multiplication: $\tilde{G}_t = G_t \cdot \text{diag}(s)$, and the weight is updated as $W_{t+1} = W_t - \eta \cdot \tilde{G}_t - \eta \lambda W_t$ (including weight decay). A norm-growth limiter optionally constrains consecutive gradient norm increases.
The critical information flow is: full gradient → random projection → tiny optimizer → scaling factor → full gradient modulation → weight update. The full-rank gradient is both consumed (to produce the weight update) and compressed (to maintain the optimizer state), with the compressed view used solely to estimate how aggressively each channel's gradient should be scaled.
3.3 Roadmap for the deep dive
Below, I will explain APOLLO's technical details in this order, which builds understanding from the motivating reformulation through the core approximation mechanism to the extreme memory-saving variant:
-
Reformulating AdamW as an adaptive learning rate algorithm: I will show how AdamW's update rule can be rewritten to separate the raw gradient direction from an element-wise scaling factor, establishing the framework for coarsening.
-
Coarsening the element-wise learning rate to channel-wise: I will explain the empirical finding that replacing per-parameter scaling with per-channel
$\ell_2$-norm-ratio scaling is sufficient for LLMs and even slightly improves performance. -
APOLLO: Approximating channel-wise scaling in a low-rank auxiliary space: I will detail the core mechanism—how the scaling factor is computed from a randomly projected low-rank gradient rather than the full-size optimizer state—and walk through Algorithm 1 step by step.
-
Theoretical justification for random projection: I will explain the three theorems that bound the approximation error of the first moment, second moment, and final scaling factor when using Gaussian random projections, and why the ratio
$\sqrt{n/r}$emerges as a necessary scaling constant. -
APOLLO-Mini: Tensor-wise scaling with rank-1 subspace: I will explain how further coarsening to a single scalar per tensor enables extreme memory efficiency at rank-1 while still outperforming AdamW.
-
Memory and compute cost analysis: I will compare APOLLO's optimizer state cost against GaLore, Fira, and Flora, and explain the practical implications of eliminating SVD.
-
Design choices and their justifications: I will distill the key architectural decisions—random projection over SVD, channel-wise over element-wise scaling, last-step aggregation style, norm-growth limiter—and why each matters.
3.4 Detailed, sentence-based technical breakdown
This is primarily an algorithm design paper whose core idea is that AdamW's element-wise adaptive learning rate can be replaced with a coarser, structured gradient scaling factor that is cheaply approximated from a very low-rank random projection of the gradient, enabling optimizer state memory to be reduced from $2mn$ to $2nr$ (or even $2n$ in the Mini variant) without losing optimization quality.
Reformulating AdamW as an Adaptive Learning Rate Algorithm without Momentum
AdamW's standard update rule, as given by the paper in Equation 1, is:
where $W_t \in \mathbb{R}^{m \times n}$ is the weight matrix at time step $t$ (with $m \leq n$ without loss of generality), $\eta$ is the learning rate, $\epsilon$ is a small constant for numerical stability, and $M_t, V_t$ are the exponentially weighted moving averages of the gradient and squared gradient:
where $\beta_1, \beta_2 \in [0, 1)$ are decay rates (typically 0.9 and 0.999), $G_t^2$ denotes element-wise squaring, and $G_t = -\nabla_W \phi_t(W_t)$ is the negative gradient of the loss with respect to the weights (the raw gradient provided by backpropagation, before the optimizer processes it).
The paper's key reframing, given in Equation 2, re-expresses this as a scaled gradient:
What this means operationally: The ratio $S = \tilde{G}_t / G_t$, computed element-wise, is a gradient scaling matrix $S \in \mathbb{R}^{m \times n}$. Each entry $S_{ij}$ tells us how much the raw gradient $(G_t)_{ij}$ is amplified or suppressed before being subtracted from the weight. Parameters whose gradient history has high variance receive small scaling factors (suppressing noisy updates), while parameters with consistent gradients receive larger scaling factors (encouraging progress in reliable directions). The weight update is then simply the learning rate times this scaled gradient.
Why this reframing matters: This decomposition separates the direction of the update (the raw gradient $G_t$) from the magnitude modulation (the scaling matrix $S$). It reveals that AdamW's entire contribution relative to plain SGD is a learned, per-parameter learning rate multiplier. This separation is what makes it possible to ask: can we approximate $S$ more cheaply without losing the benefit of adaptive learning rates? The paper's answer is yes, and the approximation strategy constitutes the entire APOLLO method.
The paper also explicitly handles the momentum term: by treating $M_t/V_t$ as the effective normalized gradient $\tilde{G}_t$, the momentum $\beta_1$ is absorbed into the computation of $S$ rather than appearing as a separate accumulator. This means that the first moment's effect on the update direction (the "inertia" that smooths the gradient trajectory) is preserved in the scaling factor, even though the scaling is later coarsened.
Coarsening the Element-wise Scaling Factor to Channel-wise
The paper's first empirical finding, demonstrated in Figure 3, is that the element-wise scaling matrix $S$ can be coarsened—replaced by a single scalar per channel (or per tensor) that is shared across all elements in that group—without degrading training. Specifically, for a weight matrix $W_t \in \mathbb{R}^{m \times n}$ where $m \leq n$, the channels are defined along the larger dimension $n$, so each column $j \in \{1, \ldots, n\}$ constitutes one channel. The structured scaling factor $s_j$ for channel $j$ is defined in Equation 3:
where $\lVert\cdot\rVert_2$ denotes the $\ell_2$ (Euclidean) norm computed over the $m$ elements in that column, and $\tilde{G}_t = M_t / (\sqrt{V_t} + \epsilon)$ is the same normalized gradient as in AdamW.
What this computes: For each channel $j$, it takes the $\ell_2$ norm (root-mean-square magnitude) of the normalized gradient column and divides by the $\ell_2$ norm of the raw gradient column. The result $s_j$ is a single positive scalar—the same for every row in channel $j$—that captures how much the optimizer's moment-based normalization would scale that channel's gradient on average. The scaled gradient is then $\tilde{G}_t = G_t \cdot \text{diag}(s_1, \ldots, s_n)$, which multiplies each column $j$ of the raw gradient by $s_j$.
Why this form rather than per-element: The per-element scaling $S_{ij}$ adapts to local noise: a single parameter that had one unusually large gradient step will have its second moment inflated, causing its learning rate to be suppressed for many subsequent steps. The channel-wise version pools statistical strength across the $m$ rows of each channel, smoothing out such idiosyncratic noise. The paper's insight is that for Transformers, the relevant structure is along channels (columns of weight matrices), where different output features or attention heads may need different learning rates, but the specific rows within a channel—corresponding to different input dimensions feeding into the same output—benefit from similar scaling. This is conceptually related to the block-diagonal Hessian structure identified in Zhang et al. (2024a), who showed that Transformer blocks have varying Hessian spectra, motivating block-wise adaptive learning rates.
Empirical validation: Figure 3 shows the training loss for a LLaMA-130M model comparing the original element-wise AdamW scaling with the proposed channel-wise version. Both achieve "similar loss reduction over training steps," with the channel-wise variant actually reaching slightly better final perplexity (24.43 vs. 25.08 for AdamW). However, the channel-wise curve exhibits a "significant spike at the early stage" due to unstable gradients. This leads to the paper's adoption of the Norm-growth Limiter (NL), borrowed from Fira (Chen et al., 2024), defined in Equation 4:
where $\gamma = 1.01$ by default.
What this limiter does: After computing the scaled gradient $\tilde{G}_t$, the optimizer checks whether its $\ell_2$ norm has grown by more than a factor of $\gamma$ relative to the previous step's scaled gradient norm. If so, it rescales $\tilde{G}_t$ to have a norm exactly $\gamma$ times the previous norm, preserving the direction but capping the magnitude increase. This prevents loss spikes during early training when gradients can be explosively large, and the paper notes it is "slightly more effective than gradient clipping." With the NL, the channel-wise variant achieves 24.11 perplexity (vs. 24.43 without it), further beating AdamW's 25.08.
The foundational observation: This finding—that a structured, coarse learning rate adaptation is not merely adequate but potentially better than the fine-grained, per-parameter version—is the paper's conceptual cornerstone. It suggests that AdamW's element-wise precision is largely redundant: the optimizer states are encoding information at a granularity finer than what actually benefits the optimization trajectory. This redundancy creates the opportunity for aggressive compression of the optimizer state without performance loss, which the next components exploit.
APOLLO: Approximating Channel-wise Scaling in a Low-Rank Auxiliary Space
The coarsening insight alone does not reduce memory: computing $s_j$ still requires maintaining the full-rank $M_t, V_t \in \mathbb{R}^{m \times n}$, which cost $2mn$ floating-point numbers—exactly the same memory as standard AdamW. The paper's core technical contribution is a method to approximate the channel-wise scaling factor $s_j$ from a heavily compressed representation of the gradient, thereby avoiding the need to store the full optimizer states. This is Algorithm 1, which I will now walk through in complete detail.
Step 1: Project the gradient into a low-rank space (Algorithm 1, Step 1)
The method begins with the full-rank gradient $G_t \in \mathbb{R}^{m \times n}$ obtained from standard backpropagation. At every $T$ steps (default $T = 200$), a fresh random projection matrix $P_t \in \mathbb{R}^{r \times m}$ is generated by sampling each element independently from a Gaussian distribution with mean 0 and variance $1/r$:
where $\mathcal{N}_{\text{seed}}$ indicates that a random number generator with a specific, stored seed is used, meaning $P_t$ is pseudorandomly generated but can be exactly reconstructed from the seed without storing the matrix itself. The projected gradient is then:
What this physically computes: It takes the $m \times n$ gradient matrix and compresses it along the first dimension from size $m$ to size $r$. Each of the $r$ rows of $R_t$ is a random linear combination of the $m$ original gradient rows, with the combination weights given by the corresponding row of $P_t$. This is a single matrix-matrix multiply (GEMM) costing $O(r m n)$ operations, which is cheap when $r \ll m$. The projected gradient $R_t$ preserves the channel structure (dimension $n$) while collapsing the row dimension.
Why random projection rather than SVD: Prior work (GaLore, Fira) uses SVD to identify the $r$ most informative directions of the gradient and projects onto those. APOLLO uses a random Gaussian projection that is oblivious to the gradient's actual structure. The paper's theoretical analysis (Section 4.1.2) proves that random projection bounds the approximation error of the channel-wise norm ratios, and the empirical results (Figure 5a-c) confirm that APOLLO's performance is nearly identical whether using SVD or random projection, while GaLore degrades significantly with random projection. The design choice is justified by the coarsening: because the scaling factor averages information across channels, it is robust to projection noise, unlike the per-element optimizer states in GaLore which require precise subspace alignment.
Why resample every $T$ steps: Using the same projection matrix throughout training would constrain the optimizer to a fixed low-dimensional subspace, limiting the directions in which weights can grow (this is the same issue acknowledged in GaLore's analysis). By resampling every 200 steps—which costs essentially nothing since it requires only generating a new random seed—APOLLO periodically changes which linear combinations of gradient rows are observed, ensuring that all directions can eventually influence the scaling factor estimates. The paper notes this is "effortlessly done by re-generating a new seed."
Step 2: Maintain low-rank AdamW optimizer states (Algorithm 1, Step 2)
The projected gradient $R_t \in \mathbb{R}^{r \times n}$ is fed into a standard AdamW optimizer (but with zero weight decay, since weight decay is applied separately in the full-rank space). This produces low-rank first and second moment estimates:
where $M^R_t, V^R_t \in \mathbb{R}^{r \times n}$ and $R_t^2$ is element-wise squaring of the projected gradient.
The normalized projected gradient is then:
What this computes: This is exactly the AdamW update rule applied to the $r \times n$ projected gradient rather than the $m \times n$ full-rank gradient. The first moment $M^R_t$ smooths the projected gradient over time (momentum), and the second moment $V^R_t$ tracks the per-element variance of the projected gradient. Their ratio $\tilde{R}_t$ is the normalized, variance-adjusted projected gradient, analogous to $\tilde{G}_t$ in the full-rank case but now operating in the compressed space.
The memory saving: Instead of storing two $m \times n$ matrices for the optimizer states, APOLLO stores two $r \times n$ matrices, reducing optimizer memory from $2mn$ to $2rn$. When $r = m/4$ (the typical GaLore setting), this is a $4\times$ reduction; when $r \ll m$, the savings are much larger. For APOLLO-Mini with $r = 1$, the optimizer state memory becomes $2n$, which is negligible compared to the $mn$ weight matrix for typical Transformer dimensions where $m \gg 1$.
Why zero weight decay here: The paper applies weight decay separately in the full-rank weight update step (Step 4). Applying weight decay to the low-rank states would mean decaying projected representations of the weights, which is meaningless since the low-rank space tracks gradient statistics, not weight magnitudes.
Step 3: Compute the channel-wise gradient scaling factor from the low-rank normalized gradient (Algorithm 1, Step 3)
For APOLLO (channel-wise variant), the scaling factor for each channel $j$ is:
where $\tilde{R}_t[:, j]$ is the $j$-th column (size $r$) of the normalized projected gradient from Step 2, and $R_t[:, j]$ is the $j$-th column of the raw projected gradient from Step 1.
The full $n$-dimensional vector $s^R = (s^R_1, \ldots, s^R_n)$ is then used to construct a diagonal matrix $S = \text{diag}(s^R_0, s^R_1, \ldots, s^R_n)$ that scales the full-rank gradient.
What this computes: This is the same ratio-of-norms formula as the coarsened AdamW (Equation 3), but now the norms are computed over $r$ elements (the projected rows) rather than $m$ elements (the original rows). Each $s^R_j$ estimates the per-channel scaling factor that AdamW would have produced, using only $r$ random linear combinations of the original gradient rows rather than all $m$ of them.
The critical approximation: The paper proves (Theorem A.4, detailed below) that $s^R_j$ is related to the true channel-wise scaling factor $s_j$ by a factor of approximately $\sqrt{r/n}$. That is, the scaling factors computed from the low-rank space are systematically smaller than those from the full-rank space by this predictable ratio. The paper accounts for this by introducing a gradient scale factor $\alpha$ (Algorithm 1, Step 4) that can absorb this ratio, though in practice $\alpha$ is often set to 1 for APOLLO since the effect can be compensated by the learning rate.
Why compute the ratio from projected norms rather than directly projecting the full-rank scaling factor: If one attempted to compute $S$ in the full-rank space and then project it, the projection step itself would require storing the full-rank states. By computing the scaling factor entirely within the low-rank space—using statistics ($\tilde{R}_t$ and $R_t$) that are already maintained there—APOLLO never needs to materialize the full-rank $M_t$ or $V_t$. The scaling factor $s^R$ is the only information that leaves the compressed space to modulate the full-rank update.
Step 4: Apply the channel-wise scaling to the full-rank gradient and update weights (Algorithm 1, Step 4)
The final weight update is:
where $\eta$ is the learning rate, $\alpha$ is the gradient scale factor (default 1 for APOLLO, $\sqrt{128}$ for APOLLO-Mini), $\lambda$ is the weight decay coefficient, and $G_t \cdot \text{diag}(s^R)$ means the full-rank gradient $G_t$ is column-wise multiplied by the scaling vector $s^R$: each column $j$ of $G_t$ is multiplied by the scalar $s^R_j$.
What this computes: The full-rank gradient $G_t$ is scaled channel-by-channel using the scaling factors estimated from the low-rank space. This produces $\tilde{G}^{\text{APOLLO}}_t$, the APOLLO-equivalent of AdamW's normalized gradient. This scaled gradient, multiplied by the learning rate and the scale factor $\alpha$, is subtracted from the weights along with a standard decoupled weight decay term.
Why apply scaling to the full-rank gradient rather than updating in the low-rank space: Unlike GaLore, which performs the optimizer update in the low-rank space and then projects the update back to full-rank, APOLLO uses the low-rank space only to estimate the scaling factor, not to compute the update direction. The update direction itself comes from the full-rank gradient $G_t$, preserving the full expressivity of the parameter space. The scaling factor merely modulates how aggressively each channel moves in that full-rank direction. This design means that errors in the low-rank approximation affect only the learning rate per channel, not the update direction—a much more forgiving failure mode.
Summary of the APOLLO loop: At each training step, the full gradient $G_t$ is computed by backpropagation. It is simultaneously (1) used directly (after scaling) for the weight update, and (2) randomly projected to $r \times n$ to maintain the low-rank optimizer states $M^R_t, V^R_t$, which produce a normalized projected gradient $\tilde{R}_t$. The ratio of per-channel norms of $\tilde{R}_t$ to $R_t$ gives the scaling factors $s^R$, which are multiplied column-wise onto the full-rank gradient. A norm-growth limiter optionally constrains sudden increases in the resulting scaled gradient. The weight is updated with the scaled gradient minus weight decay.
Theoretical Justification for Random Projection
The paper provides a formal analysis (Section 4.1.2, with full proofs in Appendix A.1) demonstrating that Gaussian random projection preserves the channel-wise norm ratios that APOLLO relies on, with high probability. The analysis proceeds in three theorems, all assuming a fixed projection matrix $P \in \mathbb{R}^{r \times m}$ with entries $P_{ij} \sim \mathcal{N}(0, 1/r)$ i.i.d.
Theorem 4.1 (First Moment Preservation): For any channel $j$ and any time step $t$, with probability at least $1 - 2\exp(-r\epsilon^2/8)$:
What this says: The $\ell_2$ norm of the first moment vector for any channel $j$ in the projected space is within a factor of $(1 \pm \epsilon)$ of its value in the original space. The probability of large deviation decreases exponentially with the projection rank $r$. The proof (Appendix A.1.3) relies on the Johnson-Lindenstrauss lemma: Gaussian random projections approximately preserve pairwise distances and vector norms when the target dimension $r$ is sufficiently large. The key step is recognizing that the first moment $M^R_t[:, j]$ is exactly the projection $P M_t[:, j]$ of the full-rank first moment for that channel, so norm preservation follows directly from the properties of Gaussian random matrices.
Why the union bound over time steps matters: The theoretical condition $r \geq 8/\epsilon^2 \cdot \log(2t/\delta)$ grows only logarithmically with the number of steps $t$, meaning the required rank stays modest even for long training runs. The paper notes that empirically, a rank of 256 suffices to train LLaMA-7B for 150K steps, consistent with this logarithmic scaling.
Theorem 4.2 (Second Moment Preservation): For any channel $j$ and any time step $t$, if $r \geq \frac{8}{\epsilon^2} \log(2t/\delta)$, then with probability at least $1 - \delta/2$:
What this says: The $\ell_1$ norm (sum of absolute values, not sum of squares) of the second moment for a channel is preserved under projection, with a similar high-probability bound. The $\ell_1$ norm is relevant here because the variance term appears inside a square root in the denominator of $\tilde{G}_t$: when the variance is later averaged across the projection dimension (as the paper does to bound the final ratio in Theorem A.4), the $\ell_1$ norm controls the average magnitude. The proof (Appendix A.1.4) expands both second moments as weighted sums of squared gradient norms and applies the per-step norm preservation from Theorem 4.1 to each term in the sum.
Why $\ell_1$ rather than $\ell_2$ for the second moment: The second moment $V_t$ contains squared gradients, so its $\ell_1$ norm is the sum of squared gradient elements—a measure of total variance across the channel. When this is divided by the channel dimension to get an average variance for the scaling factor computation (as in Theorem A.4), the $\ell_1$ norm provides the natural aggregate.
Theorem 4.4 (Main Result—Bounded Scaling Factor Ratio): For any channel $j$, with probability at least $1 - \delta$:
What this says: The low-rank scaling factor $s^R_j$, multiplied by $\sqrt{n/r}$, is within a bounded factor of the true full-rank channel-wise scaling factor $s_j$. That is, $s^R_j \approx \sqrt{r/n} \cdot s_j$—the low-rank estimate is systematically smaller by the square root of the compression ratio.
Where $\sqrt{n/r}$ comes from: The scaling factor is a ratio of $\ell_2$ norms: $s_j = \lVert\tilde{G}_t[:, j]\rVert / \lVert G_t[:, j]\rVert$. Under random projection, each of these $\ell_2$ norms scales by approximately $\sqrt{r/m}$ (since projecting a vector from dimension $m$ to $r$ reduces its expected squared norm by $r/m$). However, $\tilde{G}_t$ involves a division by $\sqrt{V_t}$, and the variance's aggregate behavior under projection introduces an additional factor related to $\ell_1$ norm preservation (Theorem 4.2). The net effect, derived in Appendix A.1.5, produces the $\sqrt{n/r}$ factor rather than a simple $\sqrt{r/m}$. The proof assumes that the channel-wise variance can be approximated by its average (following Adam-mini's insight that precise per-element variance is unnecessary), allowing the ratio to factor into terms bounded by Theorems 4.1 and 4.2.
Why this theoretical result is practically important: It confirms that the scaling factor can be reliably estimated from a low-rank projection, with the loss of precision controlled by $\epsilon$ (which depends on $r$). It also explains why a gradient scale factor $\alpha$ is needed—particularly for APOLLO-Mini where $r = 1$ and the ratio $\sqrt{n/r}$ can be large (e.g., $\sqrt{n}$)—to bring the effective learning rate into the right range. The paper validates this empirically in Figure 4: the scaling factor ratios for APOLLO with ranks $n/8$ and $n/4$ closely match the predicted $\sqrt{1/8} \approx 0.35$ and $1/2$ relative to the full-rank baseline.
APOLLO-Mini: Tensor-wise Scaling with Rank-1 Subspace for Extreme Memory Efficiency
APOLLO-Mini is the paper's extreme memory-saving variant, designed to push optimizer state memory to SGD-like levels. It makes two changes to the base APOLLO algorithm:
Change 1: Rank-1 projection. The projection matrix $P_t$ is $1 \times m$ (a row vector), so the projected gradient $R_t = P_t G_t$ is $1 \times n$—a single row. The low-rank optimizer states $M^R_t, V^R_t$ are also $1 \times n$, costing only $2n$ numbers total.
Change 2: Tensor-wise scaling factor. Instead of computing one scaling factor per channel (giving $s^R \in \mathbb{R}^{1 \times n}$), APOLLO-Mini collapses further to a single scalar for the entire weight matrix:
where the norms are taken over the entire $1 \times n$ matrix (both dimensions). The weight update becomes $W_{t+1} = W_t - \eta \cdot \alpha \cdot s^R \cdot G_t - \eta \lambda W_t$, scaling the entire gradient by the same factor.
Why tensor-wise rather than channel-wise at rank-1: The paper notes (Section 4.2) that "simply setting rank to 1 in APOLLO using channel-wise gradient scaling doesn't work well due to rank-1 space sacrificing too much information." With only one random linear combination of the $m$ gradient rows, the channel-wise $\ell_2$ norms $\lVert\tilde{R}_t[:, j]\rVert_2$ are scalars computed from a single number each—they have extremely high variance and are unreliable. By averaging the norm calculation over all $n$ channels to produce a single scalar, APOLLO-Mini pools statistical strength from the entire tensor, producing a much more stable estimate.
The scale factor $\alpha = \sqrt{128}$: Because $r = 1$ and $n$ can be large (e.g., 512 for LLaMA-60M, 4096 for LLaMA-7B), the theoretical ratio $\sqrt{n/r}$ would be $\sqrt{n}$, which can be enormous (e.g., $\sqrt{4096} = 64$). Rather than use $\sqrt{n}$ directly—which would be huge and model-dependent—the paper "heuristically set[s] a gradient scale factor $\alpha$ like GaLore to avoid performance degradation, default to set as $\sqrt{128}$." This is a pragmatic engineering choice: the factor is large enough to compensate for the systematic underestimation but fixed rather than model-dependent.
Why rank-1 works at all: The key to understanding APOLLO-Mini is that the scaling factor is just one number per weight tensor. Even a single random projection of the gradient provides a rough estimate of the gradient's overall magnitude relative to its variance-normalized counterpart. Since this estimate is applied uniformly to all parameters in the tensor, small errors in the estimate are corrected by subsequent gradient steps—the optimizer's feedback loop smooths out the approximation noise over time. Figure 5(d) demonstrates this dramatically: APOLLO-Mini with rank-1 actually outperforms full-rank AdamW on LLaMA-60M, achieving 31.93 vs. 34.06 perplexity. The tensor-wise averaging appears to act as a beneficial regularizer, suppressing the kind of noisy per-element learning rate fluctuations that can trap AdamW in sharp local minima.
Memory and Compute Cost Analysis
Table 1 in the paper provides a complete comparison of memory usage for optimizer states across methods, which I will now unpack:
| Method | Optimizer States (floating-point numbers) | Requires SVD? |
|---|---|---|
| AdamW | $2mn$ | No |
| APOLLO-Mini | $2n + 2$ | No |
| APOLLO | $2nr + 2$ | No |
| GaLore | $2nr + mr$ | Yes |
| Fira | $2nr + mr + 1$ | Yes |
| Flora | $2nr + 1$ | No |
What these numbers mean: For a weight matrix $W \in \mathbb{R}^{m \times n}$ (with $m \leq n$), the optimizer states store floating-point values. AdamW's $2mn$ comes from the first moment $M_t$ ($mn$) and second moment $V_t$ ($mn$). APOLLO's $2nr + 2$ comes from the low-rank moments $M^R_t$ and $V^R_t$ (each $r \times n$, total $2rn$), plus two scalars: the stored random seed for regenerating $P_t$ and the gradient norm for the norm-growth limiter. GaLore and Fira additionally store the projection matrix $P_t \in \mathbb{R}^{r \times m}$ (cost $mr$), which APOLLO avoids by regenerating it from a seed. Fira adds one more scalar for its error residual tracking.
Concrete example—LLaMA-7B: The largest weight matrices in LLaMA-7B have dimensions 4096 × 4096 ($m = n = 4096$). AdamW stores $2 \times 4096 \times 4096 \approx 33.6$ million values per such matrix. APOLLO with $r = 256$ (the setting used for 7B experiments in Table 3) stores $2 \times 256 \times 4096 \approx 2.1$ million values—a $16\times$ reduction. APOLLO-Mini stores $2 \times 4096 = 8192$ values—a $4096\times$ reduction—plus the negligible constant overhead. Across the entire model, this translates to the dramatic memory figures in Table 3: APOLLO uses 1.6 GB for optimizer states vs. 13 GB for 8-bit Adam, and APOLLO-Mini uses 0.0 GB (effectively negligible).
The SVD elimination: The practical impact goes beyond memory. Table 7 reports optimizer step times on LLaMA-7B: GaLore takes 2.874 seconds per step, Fira takes 3.086 seconds, while APOLLO takes 0.159 seconds and APOLLO-Mini 0.142 seconds—both faster even than AdamW's 0.173 seconds. The reason is that the low-rank AdamW update on $r \times n$ matrices is faster than the full-rank AdamW on $m \times n$ matrices, and the random projection is a cheap GEMM, whereas SVD is $O(mn^2)$ and dominates GaLore/Fira's runtime. Figure 9 visualizes this as dramatic throughput spikes downward at every SVD update step for GaLore-type optimizers.
The stored random seed trick: The projection matrix $P_t$ is never explicitly stored. Instead, a single integer seed is saved, and $P_t$ is regenerated by calling a pseudorandom number generator with that seed whenever needed (for the forward projection $R_t = P_t G_t$ during the optimizer step, and potentially for the backward pass through the projection if automatic differentiation is used—though in practice the projection is applied to the gradient $G_t$ which has already been computed, so no backward through $P_t$ is required). This is what allows APOLLO to avoid the $mr$ storage cost that GaLore and Fira pay for their projection matrices.
Design Choices and Their Justifications
1. Random projection over SVD: The paper's theoretical and empirical analysis demonstrates that channel-wise $\ell_2$-norm ratios are robust to the specific projection used, because norms of random projections concentrate around their expected values (Johnson-Lindenstrauss property). SVD would find the optimal low-rank approximation of the gradient itself, but APOLLO doesn't need the gradient's low-rank structure—it only needs aggregate norm statistics, which random projections preserve. The practical consequence is speed: random projection is $O(r m n)$ via GEMM (highly optimized on GPUs), while SVD is $O(m n^2)$ and "cannot be efficiently parallelized on GPUs."
2. Channel-wise over element-wise scaling: The choice is justified by the redundancy observation in Figure 3: per-element adaptation provides no benefit over per-channel adaptation for Transformer training. The likely mechanism is that channels (columns of weight matrices) correspond to output features or attention heads that may benefit from distinct learning rates, but the input dimensions feeding into each output feature share similar gradient statistics. This is consistent with the block-diagonal Hessian structure identified in Zhang et al. (2024a).
3. Norm-growth limiter over gradient clipping: Standard gradient clipping caps the global $\ell_2$ norm of the entire gradient at a fixed threshold. The norm-growth limiter instead caps the rate of increase of the scaled gradient's norm relative to the previous step, which the paper finds "slightly more effective" empirically. This is a more adaptive constraint—it doesn't prevent large gradients per se (if they grow gradually), only sudden spikes, which are the specific pathology observed in early training (Figure 3, orange spike). The threshold $\gamma = 1.01$ is deliberately tight, allowing at most 1% norm growth per step.
4. Fixed scale factor $\alpha = \sqrt{128}$ for APOLLO-Mini: This is a heuristic that the paper explicitly acknowledges. The theoretical ratio $\sqrt{n/r}$ would be $\sqrt{n}$, which varies across layers (different dimensions) and models. Using a fixed constant $\sqrt{128}$—chosen because 128 is a common hidden dimension and $\sqrt{128} \approx 11.3$ is a reasonable scaling—prioritizes simplicity and consistency across the model. The paper notes that this factor "can be combined with the learning rate," but setting it separately allows the learning rate $\eta$ to remain at standard values (e.g., 0.01 for pre-training) while $\alpha$ compensates for the systematic underestimation of the scaling factor in rank-1 space.
5. Projection resampling every $T = 200$ steps: This is adopted from GaLore's practice and addresses the same concern: a fixed projection matrix restricts optimization to a fixed low-dimensional subspace. Resampling periodically (at negligible cost, since only a seed changes) ensures that all gradient directions eventually influence the optimizer state estimates. The paper does not ablate this frequency, suggesting it is not highly sensitive.
6. Zero weight decay in the low-rank optimizer state: The decoupled weight decay $-\eta \lambda W_t$ is applied in the full-rank weight update step, not to the low-rank moment estimates. This is design-correct: the low-rank space tracks gradient statistics for the purpose of estimating scaling factors, and applying weight decay there would distort those statistics, making the scaling factor estimate less accurate.
7. Layer-wise gradient update compatibility: The paper notes (Figure 1 middle caption) that both GaLore and APOLLO "employ the layer-wise gradient update strategy (Lv et al., 2023)." This is a system-level optimization where gradients for each layer are computed, the optimizer update is applied, and then the gradient memory is freed before processing the next layer, reducing peak memory. APOLLO's design is compatible with this because the low-rank optimizer states are maintained per weight matrix and the projection is applied independently per weight matrix; there is no cross-layer dependency in the scaling factor computation. The extreme memory figure of 12 GB for LLaMA-7B (Q-APOLLO-Mini) assumes this layer-wise strategy is active.
8. Open-source defaults without tuning: The paper emphasizes that APOLLO uses "the same learning rate 0.01 and a subspace change frequency T of 200 without tuning, following the GaLore open-sourced settings" (Appendix A.4). This is a deliberate choice to demonstrate that APOLLO's gains are not from hyperparameter optimization but from the algorithm itself. The scale factor $\alpha$ is the only new hyperparameter, and it is either set to 1 (for APOLLO on models < 1B), a small fraction like $\sqrt{1/2}$ (for 1B models), or $\sqrt{128}$ (for APOLLO-Mini), with the paper noting these were not extensively tuned.
4. Key Insights and Innovations
Innovation 1: Coarsening AdamW's Adaptation Granularity from Element-wise to Structured is Not Just a Memory Hack—It's a Conceptual Reframing That Reveals Optimization Redundancy
The paper's first distinctive intellectual move is not the compression mechanism itself (which Section 3 covers), but the diagnostic observation that drives it: AdamW's per-parameter learning rate adaptation contains substantial redundancy that, when removed, actually improves performance. This reframes the problem from "how do we compress AdamW's states without losing information?" to "how much of AdamW's state is genuinely useful, and at what granularity?"
What the field assumed before this work. The dominant assumption in memory-efficient optimizer design—exemplified by GaLore (Zhao et al., 2024) and Fira (Chen et al., 2024)—was that AdamW's per-element statistics encode essential fine-grained information about the optimization landscape. These methods therefore attempted to preserve that information through low-rank approximation of the gradient itself, using SVD to identify the most informative subspace and maintaining optimizer states at substantial rank (typically r = n/4 for GaLore). Even Adam-mini (Zhang et al., 2024b), which challenged the need for per-element second moments, preserved per-element first moments and required careful block-size tuning per module type—treating the granularity question as a module-specific engineering trade-off rather than a fundamental property of the optimization process.
What APOLLO's reframing reveals. The channel-wise scaling experiment in Figure 3 is the paper's foundational diagnostic: by comparing element-wise AdamW against a version where the gradient scaling factor s_j is a single scalar per channel (shared across all m rows of that channel), the paper demonstrates that the coarse version not only matches but exceeds element-wise AdamW (24.11 vs. 25.08 perplexity on LLaMA-130M). This is not a happy accident of approximation—it suggests that the per-element precision in AdamW's scaling matrix S is actively harmful in some regimes, likely because it overfits to noisy per-parameter gradient variance estimates in high dimensions, producing learning rates that oscillate parameter-by-parameter rather than following coherent channel-level structure.
This finding is conceptually distinct from Adam-mini's block-wise second moment. Adam-mini kept the first moment per-element and only coarsened the second moment, meaning it still tracked individual parameter momentum directions at full resolution. APOLLO's channel-wise coarsening hits both moments simultaneously—the ratio s_j = \lVert\tilde{G}_t[:, j]\rVert_2 / \lVert G_t[:, j]\rVert_2 aggregates first and second moment effects into a single channel scalar—demonstrating that the interaction of the two moments, not either individually, can be coarsened without loss. This is a stronger claim about redundancy: you do not need to track where each parameter is going individually, only how aggressively each channel as a whole should move.
Why this is fundamental rather than incremental. This insight shifts the design space from "compress AdamW's states" to "understand the coarsest granularity at which adaptive learning rates remain beneficial." It implies that future optimizer designs should be evaluated by their adaptation granularity—how many distinct learning rates they assign—as a first-class axis alongside memory cost. The paper demonstrates this by establishing that tensor-wise granularity (one scalar per weight matrix) is sufficient at rank-1 (APOLLO-Mini), while channel-wise works at modest rank (APOLLO), and element-wise is unnecessary. This granularity spectrum, and the finding that coarser is not just cheaper but better, is the intellectual contribution that Section 3's mechanisms instantiate.
Evidence anchor. Figure 3's loss curves, where structured AdamW (channel-wise) surpasses element-wise AdamW after the early unstable phase is controlled by the norm-growth limiter. Table 9's ablation showing minimal perplexity difference between channel-wise and tensor-wise granularity at r = n/4 (LLaMA-60M: 31.55 vs. 32.10; LLaMA-350M: 16.85 vs. 17.00). The empirical takeaway is that as rank drops further, tensor-wise becomes better than channel-wise (Figure 5d, where APOLLO-Mini at rank-1 outperforms APOLLO with channel-wise at rank-1)—exactly the opposite of what one would expect if finer granularity always helped.
Innovation 2: Random Projection Works Because the Scaling Factor is a Norm Ratio, Not Because Gradients Are Low-Rank
The paper's second conceptual contribution is the recognition that the success of random projection in APOLLO depends on a fundamentally different mathematical property than the success of SVD in GaLore. GaLore and its variants justify low-rank projection by arguing that gradients themselves occupy a low-dimensional subspace—a structural claim about the optimization trajectory. APOLLO instead relies on the concentration of aggregate norm statistics under random projection, a distributional claim that is both weaker (requires fewer assumptions) and more robust (tolerates much lower rank).
Prior framing in low-rank optimizers. GaLore's seminal contribution was the observation that "the gradient matrix becomes low-rank during training" for LLMs, enabling SVD-based compression. This framing treated the gradient's spectral structure as the key property to exploit, and the low-rank projection as a direct compression of that structure. The cost was the SVD itself—an $O(mn^2)$ operation needed to identify the evolving low-rank subspace—and the requirement that the rank be high enough (typically n/4) to capture sufficient spectral energy. Subsequent methods (Fira, GaLore variants with online PCA or random projections) all operated within this "gradients are low-rank" framework, debating how to find the subspace more cheaply while maintaining its quality.
APOLLO's reframing. APOLLO demonstrates that the precise subspace does not matter for its task. The scaling factor s_j = \lVert\tilde{R}_t[:, j]\rVert_2 / \lVert R_t[:, j]\rVert_2 depends only on column-wise \ell_2 norms of the projected gradient R_t = P_t G_t and its moment-normalized counterpart \tilde{R}_t. These norms, under Gaussian random projection, concentrate around their expected values with high probability (Theorems 4.1–4.4) regardless of the gradient's spectral structure. The paper proves this formally: the ratio \sqrt{n/r} \cdot s^R_j / s_j is bounded within [(1 - \epsilon)/(1 + \epsilon)^{3/2}, (1 + \epsilon)^{3/2} / (1 - \epsilon)] with probability \geq 1 - \delta, where \epsilon depends only on r (the projection dimension), not on any spectral property of G_t.
This is a fundamentally different justification than GaLore's. GaLore needs SVD because it must align the projection with the gradient's dominant singular vectors to minimize information loss in the update direction. APOLLO only needs the projection to be a norm-preserving random embedding—the Johnson-Lindenstrauss property—because it extracts only scalar summary statistics (norm ratios) from the projected gradient, not the full directional information. The upshot is that anything satisfying the Johnson-Lindenstrauss lemma works: Gaussian random matrices, sparse random projections, even Hadamard transforms. The paper's choice of Gaussian P_{ij} \sim \mathcal{N}(0, 1/r) is one instance of this broader class.
Why this matters beyond APOLLO. This reframing suggests a general design principle for memory-efficient deep learning: when the information extracted from optimizer states is an aggregate statistic (norm, variance, entropy) rather than a precise direction or element-wise value, random projection is a sufficient and much cheaper alternative to structure-seeking projections like SVD. This principle could extend to other optimizer components—e.g., gradient clipping thresholds, loss scaling factors, or learning rate warmup schedules—that depend on aggregate gradient properties rather than precise per-parameter values.
Evidence anchor. Figure 5(a–c) is the critical comparison: GaLore's performance degrades dramatically with random projection (failing to match AdamW's red dashed line on LLaMA-60M/130M/350M), while APOLLO and APOLLO-Mini show essentially identical performance with random projection vs. SVD, with APOLLO-Mini even slightly better with random projection on LLaMA-350M. Figure 4's empirical validation of the \sqrt{n/r} scaling factor ratio across different ranks (1/8n, 1/4n) confirms that the theoretical norm-preservation bounds hold in practice. Table 7 quantifies the practical consequence: APOLLO's optimizer step is 0.159s vs. GaLore's 2.874s on LLaMA-7B—an 18x speedup—because the random projection is a cheap GEMM while SVD is a sequential bottleneck.
Innovation 3: The Proposal-Verifier Analogy—APOLLO Decouples Update Direction from Update Magnitude, Making Low-Rank Errors Forgiving
The paper's third conceptual contribution is an implicit architectural principle that I will make explicit here: APOLLO's robustness to extreme rank reduction (down to rank-1) arises from a deliberate decoupling of the update direction from the update step size. This decoupling means that errors in the low-rank approximation affect only how far the weights move along each channel, not which direction they move—a far more forgiving failure mode than in GaLore or Fira, where low-rank errors corrupt the update direction itself.
The failure mode difference. In GaLore, the weight update is computed entirely in the low-rank space: $\Delta W = \eta \cdot P^T \cdot \text{AdamW}(P G_t)$, where $P$ is the SVD-derived projection and $P^T$ projects the low-rank update back to full-rank. If the low-rank approximation is poor—because the rank is too small, or the subspace is stale, or SVD is replaced with random projection—the entire update vector is corrupted: the optimizer moves weights in wrong directions, not just wrong amounts. This is why GaLore degrades steeply as rank decreases (Figure 5d: GaLore at rank = 32 falls far below AdamW) and fails completely with random projection.
In APOLLO, the weight update is $W_{t+1} = W_t - \eta \cdot (G_t \cdot \text{diag}(s^R))$, where $G_t$ is the full-rank gradient from backpropagation and $s^R$ is the scaling factor estimated from the low-rank space. The update direction is always $G_t$—which is exact and loss-derived—and only the per-channel step size $s^R_j$ is approximated. If $s^R_j$ is noisy (due to low rank), the optimizer still moves in the correct gradient direction for each parameter; it just moves with a suboptimal step size. Critically, step size errors are self-correcting: if a channel's scaling factor is too small, the gradient on the next step will be larger (because the loss hasn't decreased enough), and the scaling factor estimate will adapt accordingly. If the scaling factor is too large, the norm-growth limiter caps sudden increases (Equation 4), preventing catastrophic divergence.
This decoupling explains APOLLO's most striking empirical finding: halving the rank has minimal impact on pre-training performance (Table 2: APOLLO at rank n/4 vs. n/8 on LLaMA-60M is 31.55 vs. 31.26 perplexity—a 0.29 difference), and rank-1 outperforms full-rank AdamW on several configurations (APOLLO-Mini vs. AdamW on LLaMA-1B: 14.17 vs. 15.56 perplexity). These results would be impossible in a system where rank errors corrupt the update direction.
The analogy to proposal-verifier architectures. This decoupling parallels the proposal-verifier decomposition in test-time compute scaling (e.g., the PRM-guided search in the prior example paper): a cheap, approximate component (the verifier / scaling factor estimator) guides a more expensive, exact component (the proposal distribution / full-rank gradient) without corrupting its fundamental correctness. The verifier can make mistakes—rejecting good proposals or accepting bad ones—but each error affects only one step's selection, not the trajectory's long-run direction. Similarly, APOLLO's low-rank scaling factor can mis-estimate the optimal step size for a channel, but the next gradient computation re-anchors the update direction to the loss landscape's true gradient.
Why this is a design principle rather than an accident. The paper does not explicitly frame this as a proposal-verifier decoupling, but the design choice is deliberate: Algorithm 1 applies the scaling factor $S$ to the full-rank gradient $G_t$, not to a projected or reconstructed gradient. This means Steps 1–3 (the low-rank space) serve purely as a scalar estimator—a "critic" that assesses how aggressively each channel should move—while Step 4 executes the actual movement using the loss-derived, full-rank gradient direction. This is a fundamentally different architecture from GaLore's "do everything in the low-rank space," and it is what enables APOLLO's tolerance to extreme compression.
Evidence anchor. Figure 5(d): GaLore's perplexity degrades from ~34 (matching AdamW) at rank = 128 to ~36 at rank = 16—each rank reduction hurts substantially because the update direction is increasingly corrupted. APOLLO's perplexity is nearly flat from rank = 128 to rank = 32, and APOLLO-Mini at rank = 1 is better than APOLLO at rank = 32—direction errors are non-existent, and the step-size estimation error at rank-1 is actually beneficial (likely due to the regularization effect discussed in Innovation 4). Table 2: APOLLO at rank n/8 (halved rank) loses only 0.29 perplexity on 60M, 0.24 on 130M, 0.13 on 350M, and 0.05 on 1B—the gap shrinks with model size, suggesting the decoupling becomes more robust as the optimization landscape gets higher-dimensional.
Innovation 4: The Generalization Hypothesis—Why a "Stateless" Optimizer Can Beat AdamW
The paper's fourth contribution is a preliminary but provocative theoretical hypothesis for why reducing optimizer state granularity can improve rather than merely preserve model quality. This is the paper's answer to the central puzzle it creates: if APOLLO-Mini uses essentially no optimizer state (just a rank-1 norm ratio scaled by \sqrt{128}), why does it not only match but often outperform AdamW, which maintains detailed per-parameter statistics?
The puzzle, stated clearly. Section 3.2 shows that a full-rank structured (channel-wise) AdamW—which still stores and updates M_t, V_t \in \mathbb{R}^{m \times n}—achieves slightly better perplexity than element-wise AdamW (24.11 vs. 25.08). One could attribute this to the regularization effect of averaging out noisy per-element variance estimates. But APOLLO and APOLLO-Mini go much further: they approximate the scaling factor from a projected, low-rank gradient, injecting random projection noise on top of the coarsening. And yet APOLLO-Mini achieves 30.95 perplexity vs. AdamW's 34.06 on LLaMA-60M (Table 2, APOLLO-Mini with tuned learning rate). This is not just noise tolerance—it suggests that the approximation itself is somehow beneficial for generalization.
The paper's hypothesis: SGD-like directional randomness. Section 5.5 advances the hypothesis that APOLLO's structured scaling approximates the direction of the AdamW update while injecting SGD-like randomness through the coarsened and projected scaling factor. Formally, APOLLO-Mini's update is $W_{t+1} = W_t - \eta \cdot \alpha \cdot s^R \cdot G_t$, where $s^R$ is a single scalar computed from a rank-1 random projection. This is effectively SGD with a learned, time-varying learning rate $\eta \cdot \alpha \cdot s^R$ that adapts to the gradient's aggregate variance. The update direction is exactly the raw gradient $G_t$—no momentum-based smoothing of directions, unlike AdamW where $M_t$ (the first moment) blends past and current gradients to create a smoothed descent direction.
The paper argues, citing Zhou et al. (2020) and Keskar & Socher (2017), that SGD's directional noise helps escape sharp local minima and converges to flatter minimizers that generalize better, while AdamW's momentum-based smoothing tends to converge to sharper minima. APOLLO appears to reconcile these behaviors: it retains AdamW's per-channel learning rate adaptation (which handles the Hessian heterogeneity that makes pure SGD fail on Transformers, per Zhang et al. 2024a) while discarding the directional smoothing that biases AdamW toward sharp minima. The resulting optimizer inherits SGD's generalization bias and AdamW's trainability.
Empirical support from the paper. The paper provides three pieces of circumstantial evidence:
-
Directional sharpness measurements (Table 10): Following Pan & Li (2023), the paper measures
$v^\top \nabla^2 f(x) v$(directional sharpness) for a small T5 translation task. APOLLO and APOLLO-Mini achieve sharpness values comparable to or lower than AdamW (at epoch 20: AdamW 0.000399, APOLLO 0.000261, APOLLO-Mini 0.000101), while SGD is dramatically higher (3.207535). This suggests APOLLO finds flatter minima than AdamW, consistent with the generalization hypothesis. -
Late-stage performance crossover (Figure 6, Section 5.4 A4): Fira (which maintains more complete optimizer state through error residual tracking) initially converges faster than APOLLO but is overtaken in later training stages. The paper interprets this as evidence that "AdamW optimization states play a more crucial role in the initial phase... while compressing the optimization states into gradient scaling factors... becomes more effective in later stages" when generalization quality, not convergence speed, dominates.
-
Scaling with model size (Table 2): APOLLO's advantage over Fira grows with model size: on LLaMA-60M, Fira (31.06) slightly edges APOLLO (31.55); by LLaMA-1B, APOLLO (14.20) surpasses Fira (14.31). The paper suggests that larger models have more complex loss landscapes where the generalization benefits of SGD-like noise become more pronounced relative to the convergence benefits of precise optimizer states.
Why this is significant beyond the current results. This hypothesis—if validated by future work—provides a theoretical framework for a new class of optimizers that deliberately inject structured noise into the learning rate schedule to improve generalization, analogous to how dropout injects noise into activations. It also provides a principled explanation for why coarser learning rate adaptation can outperform finer: the coarsening itself acts as an implicit regularizer that prevents overfitting of the optimizer state to per-parameter gradient noise. This is a fundamentally different perspective from the "coarsening is a necessary evil for memory efficiency" framing that dominates prior work.
The paper's appropriate caution. The paper explicitly labels this as "preliminary insights" and defers a formal treatment to future work. The directional sharpness measurements are on a small-scale task (T5 for translation), not the LLaMA/C4 setting used for the main results. The connection to the Keskar & Socher "SGD generalizes better" literature is suggestive but not causally established. Nonetheless, the hypothesis is well-motivated by the empirical patterns and provides a coherent narrative for the paper's most surprising result—that rank-1 APOLLO-Mini can beat full-rank AdamW.
Evidence anchor. Table 10 for directional sharpness. Figure 6 and Section 5.4 A4 for the late-stage crossover. Table 2 for scaling with model size. Section 5.5, Hypothesis 1 for the explicit connection to Pan & Li (2023) and the sharpness-generalization link.
5. Experimental Analysis
Evaluation Methodology
Dataset. All pre-training experiments use the C4 dataset (Raffel et al., 2020), "a comprehensive corpus derived from Common Crawl data and meticulously filtered and cleaned." The paper follows the training configurations from Zhao et al. (2024) and Lialin et al. (2023), with data amounts specified per model size in Table 11: 1.3B tokens for 60M, 2.6B for 130M, 7.8B for 350M, 13.1B for 1B, and 19.7B for 7B. For fine-tuning, the paper uses eight common-sense reasoning datasets (Winogrande, PIQA, SIQA, OpenBookQA, HellaSwag, BoolQ, ARC-Easy, ARC-Challenge) and MMLU (Hendrycks et al., 2020) across STEM, Social Sciences, Humanities, and Other domains.
Base model(s). The LLaMA architecture family is used across scales from 60M to 7B parameters for pre-training, with architectural details in Table 11 (e.g., LLaMA-60M: 512 hidden, 8 heads, 8 layers; LLaMA-7B: 4096 hidden, 32 heads, 32 layers). For fine-tuning, three additional open-source models are used: LLaMA-3.2-1B, LLaMA-3-8B, Gemma-7B, and Mistral-7B. The paper notes that APOLLO is also tested on a LLaMA-13B model for the memory scalability claim (Section 5.3). All experiments use BF16 data format without quantization unless explicitly combined with INT8 weight quantization for the extreme memory experiments.
Metrics. Pre-training performance is measured by validation perplexity on the C4 dataset, computed after training for the specified number of steps (Table 11). Downstream task performance is measured by zero-shot accuracy on common-sense reasoning and MMLU tasks (exact match or multiple-choice accuracy). Memory usage is reported in GB, with optimizer state memory specifically broken out. Throughput is measured in tokens per second and end-to-end training time in seconds.
Baselines. Six baselines are compared against: (i) AdamW (Loshchilov & Hutter, 2019): full-rank optimizer maintaining M_t, V_t at full precision. (ii) Low-Rank: decomposes weights as W = UV with both factors trained by AdamW. (iii) LoRA (Hu et al., 2021): low-rank adapters with frozen backbone. (iv) ReLoRA (Lialin et al., 2023): LoRA variant for pre-training that periodically merges adapters into weights. (v) GaLore (Zhao et al., 2024): gradient low-rank projection via SVD. (vi) Fira (Chen et al., 2024): GaLore enhanced with error residual tracking. For fine-tuning, DoRA (Liu et al., 2024a) is additionally included.
Generation budget / compute accounting. For pre-training, all methods use the same training data amount and number of steps (Table 11), with learning rate warmup over the first 10% of steps followed by cosine annealing to 10% of the initial value. The batch size is fixed at 131K tokens across all architectures. For optimizer state memory, the paper reports the number of floating-point values stored per weight matrix (Table 1), which directly translates to memory in GB given BF16 format. FLOPs are not directly compared; instead, throughput (tokens/second) and optimizer step time (seconds) are measured on identical hardware (A100-80GB GPUs). For the rank budget, all low-rank methods default to r = n/4 (one-quarter of the smaller weight dimension), with ablations at r = n/8 and r = 1.
Cross-validation / statistical protocol. No cross-validation is reported; the paper uses fixed train/test splits of C4. For fine-tuning learning rate sweeps, the best result is reported. For the learning rate sweep on AdamW in the long-context experiment (Section 5.4, A5), the paper varies AdamW's learning rate across [1e-3, 2.5e-3, 5e-3, 7.5e-3, 1e-2] and reports APOLLO-series results under a fixed learning rate of 1e-2 with lazy scale factor tuning. Validation perplexity is recorded every 1000 steps for the 7B experiments. No confidence intervals or standard deviations are reported for any result.
Main Quantitative Results
Pre-training Performance: APOLLO Matches or Exceeds AdamW Across All Model Scales While Using Dramatically Less Memory
Table 2 presents the headline pre-training results for LLaMA models from 60M to 1B parameters on the C4 dataset. At the standard rank setting of r = n/4 (matching GaLore and Fira), APOLLO achieves consistently better perplexity than GaLore across all model sizes: 31.55 vs. 34.88 on 60M (gap of 3.33), 22.94 vs. 25.36 on 130M (gap of 2.42), 16.85 vs. 18.95 on 350M (gap of 2.10), and 14.20 vs. 15.64 on 1B (gap of 1.44). More importantly, APOLLO surpasses full-rank AdamW on all model sizes despite using dramatically less memory: on LLaMA-1B, APOLLO achieves 14.20 perplexity with 4.38 GB optimizer memory vs. AdamW's 15.56 perplexity with 7.80 GB. The APOLLO variant using SVD projection (APOLLO w. SVD) performs nearly identically to the random projection version (e.g., 60M: 31.26 vs. 31.55; 350M: 16.67 vs. 16.85), demonstrating the SVD is unnecessary.
APOLLO with halved rank (r = n/8, marked with ✝ in Table 2) shows minimal performance degradation: 31.26 vs. 31.55 on 60M (0.29 perplexity loss), 23.18 vs. 22.94 on 130M (0.24 loss), 16.98 vs. 16.85 on 350M (0.13 loss), and 14.25 vs. 14.20 on 1B (0.05 loss). Critically, the gap shrinks with model size—the largest model (1B) loses only 0.05 perplexity from halving the rank—suggesting APOLLO's robustness to rank reduction scales favorably.
APOLLO-Mini (rank-1, tensor-wise scaling) achieves the most striking results. With the default learning rate of 0.01 (matching GaLore), it achieves 31.93 on 60M, 23.53 on 130M, 17.18 on 350M, and 14.17 on 1B—all better than AdamW and GaLore. With a tuned learning rate of 0.02 (marked ‡), APOLLO-Mini further improves to 30.95 on 60M, 22.85 on 130M, 16.63 on 350M, and 13.95 on 1B. The 1B result (13.95) represents a 1.61 perplexity improvement over AdamW's 15.56 while using essentially zero optimizer state memory (2.60 GB vs. 7.80 GB, of which the 2.60 GB is predominantly weight memory since the optimizer state is only 2n + 2 per weight matrix). The paper emphasizes this is "a 2.80 reduction in validation perplexity" relative to AdamW on the 1B model (15.56 → 13.95, though 15.56/13.95 ≈ 1.12, so the "2.80 reduction" in the abstract appears to reference the gap between APOLLO-Mini and a different baseline, likely GaLore's 1B result; the exact calculation should be verified).
Comparison with Fira reveals a model-size-dependent trend. On the smallest model (60M), Fira marginally leads (31.06 vs. APOLLO's 31.55). By 130M, APOLLO surpasses Fira (22.94 vs. 22.73). On 350M, APOLLO leads (16.85 vs. 17.03). On 1B, the gap widens (14.20 vs. 14.31). Fira uses SVD and stores a projection matrix (mr extra memory, Table 1), making APOLLO strictly more efficient even where performance is comparable.
LLaMA-7B scaling results appear in Figure 2 (validation perplexity curves) and Table 3 (perplexity at checkpoints). Under a fixed training time budget of approximately two weeks (15 days), APOLLO (r = 256) and APOLLO-Mini (r = 1) are the only methods to complete pre-training within the time constraint due to their higher throughput from larger batch sizes. Figure 2 shows APOLLO and APOLLO-Mini achieving perplexity around 13.0 at 150K steps vs. GaLore's ~14.0 at only ~60K steps—GaLore could not complete the full training due to SVD overhead. Table 3 provides checkpoint comparisons: at 40K steps, APOLLO achieves 17.55, APOLLO-Mini 18.03, vs. 8-bit AdamW 18.09 and 8-bit GaLore 17.94. By 150K steps, APOLLO reaches 13.02, APOLLO-Mini 13.09, vs. 8-bit AdamW 14.61 and 8-bit GaLore 14.65—a >1.5 perplexity improvement over both quantized baselines. Memory usage for optimizer states alone: APOLLO 1.6 GB, APOLLO-Mini 0.0 GB (negligible), vs. 8-bit AdamW 13 GB and 8-bit GaLore 4.9 GB.
Downstream Task Performance: Lower Perplexity Translates to Better Zero-Shot Accuracy
Table 4 evaluates LLaMA-350M models pre-trained with sequence lengths of 256 and 1024 on nine zero-shot common-sense and math reasoning tasks. At sequence length 256, APOLLO achieves an average accuracy of 0.3681 across all tasks vs. AdamW's 0.3554—a 1.27 percentage point improvement. APOLLO-Mini achieves 0.3654. At sequence length 1024, both APOLLO variants widen the gap: APOLLO reaches 0.3840 vs. AdamW's 0.3712 (+1.28 percentage points), and APOLLO-Mini reaches 0.3785. Individual task improvements are mixed—AdamW leads on BoolQ at sequence length 256 (0.5881 vs. 0.5165 for APOLLO) and on Winogrande at sequence length 1024 (0.5233 vs. 0.4925), while APOLLO leads substantially on PIQA (0.6632 vs. 0.6387 at seqlen 256; 0.6681 vs. 0.6534 at seqlen 1024) and SciQ (0.592 vs. 0.591 at seqlen 256; 0.624 vs. 0.609 at seqlen 1024). The paper notes that these models used the best AdamW checkpoint obtained by sweeping learning rates, while APOLLO-series checkpoints used the default learning rate of 0.01 (not tuned for downstream tasks). Memory for optimizer states is reported as 1.37 GB for AdamW (the full model memory including weights and optimizers), 0.34 GB for APOLLO, and 0.00 GB for APOLLO-Mini.
Fine-tuning Performance: Full-Parameter Quality at Low-Rank Memory Cost
Table 5 reports common-sense reasoning fine-tuning results on LLaMA-3.2-1B. Full-rank AdamW achieves an average accuracy of 68.07% across eight tasks. APOLLO with rank 32 achieves 68.21% (slightly above AdamW), while APOLLO w. SVD achieves 69.08% (1.01 percentage points above AdamW). APOLLO-Mini with rank 1 achieves 68.23%—also above AdamW. For context, LoRA at rank 32 achieves only 59.21% (significantly below AdamW), DoRA achieves 66.38%, and GaLore with SVD and rank 32 achieves 61.14%. Fira achieves the highest average at 68.98%, slightly above APOLLO w. SVD's 69.08%—the difference is negligible (0.10 percentage points) and APOLLO achieves this without SVD.
Table 6 reports MMLU fine-tuning on LLaMA-3-8B, Gemma-7B, and Mistral-7B, with learning rates swept for all methods. On LLaMA-3-8B, APOLLO w. SVD achieves 64.76% average vs. full AdamW's 64.85%—essentially tied. APOLLO achieves 64.35%, APOLLO-Mini 64.41%, GaLore 64.43%, Fira 64.32%, and LoRA 64.25%. All methods are within 0.6 percentage points of each other, suggesting MMLU fine-tuning is relatively insensitive to optimizer choice at this model scale. On Gemma-7B, APOLLO w. SVD achieves 34.98% vs. full AdamW's 34.21%—actually surpassing it—while GaLore achieves only 30.95% and LoRA 32.18%. APOLLO-Mini drops to 31.67%, indicating some sensitivity on Gemma. On Mistral-7B, APOLLO w. SVD achieves 61.76% vs. AdamW's 61.67% (tied), while all other low-rank methods are within 0.4 percentage points. The key takeaway is that APOLLO at rank 32 matches or exceeds AdamW on fine-tuning for most model-task combinations, while APOLLO-Mini at rank 1 is competitive (within 2-3 percentage points of AdamW on MMLU) but not universally superior.
System-Level Benefits: Throughput, Memory, and Optimizer Step Overhead
Throughput (Figure 1, right; Figure 2): On 8×A100-80GB GPUs pre-training LLaMA-7B, APOLLO and APOLLO-Mini support batch size 16 while AdamW is limited to batch size 4 due to memory constraints. This enables approximately 3× higher throughput for APOLLO-series methods: Figure 1 (right panel) shows AdamW at batch size 4 achieving throughput near ~8000 tokens/second while APOLLO and APOLLO-Mini at batch size 16 exceed ~20,000 tokens/second. At batch size 8 (which GaLore can support), APOLLO-series still delivers higher throughput than GaLore (Figure 1 right shows APOLLO at batch size 16 clearly above GaLore at batch size 8). When batch sizes are equalized at 4 (the maximum AdamW can handle), APOLLO-Mini still shows slightly higher throughput than AdamW due to reduced optimizer step overhead.
Memory breakdown (Figure 1, middle): At batch size 4, AdamW's memory breakdown shows optimizer states dominating (~28 GB of ~58 GB total for LLaMA-7B, consistent with the introduction's numbers). GaLore reduces optimizer state memory substantially but still requires significant overhead. APOLLO and APOLLO-Mini reduce optimizer states to very small fractions: APOLLO-Mini's optimizer memory is essentially invisible in the bar chart. When combined with INT8 weight quantization (Q- variants), Q-APOLLO-Mini achieves the extreme memory figure: LLaMA-7B pre-training in under 12 GB total GPU memory (Figure 1 middle, rightmost bar; Table 8).
Optimizer step time (Table 7): On LLaMA-1B with batch size 16 (the maximum AdamW supports), the backward pass takes 1.069 seconds. AdamW's optimizer step adds 0.036 seconds, APOLLO adds 0.051 seconds, APOLLO-Mini 0.048 seconds, GaLore 0.371 seconds, and Fira 0.421 seconds. On LLaMA-7B with batch size 4, backward pass takes 0.712 seconds; AdamW adds 0.173 seconds, APOLLO 0.159 seconds, APOLLO-Mini 0.142 seconds, GaLore 2.874 seconds, and Fira 3.086 seconds. The critical finding: APOLLO and APOLLO-Mini are faster than AdamW on the 7B model (0.159s and 0.142s vs. 0.173s) because the low-rank AdamW update on r × n matrices is cheaper than full-rank m × n AdamW, and the random projection GEMM overhead is small. GaLore and Fira are 16-18× slower than AdamW on 7B due to SVD. The paper notes that GaLore's SVD "take[s] approximately ten minutes" per update at 7B scale, but Table 7 averages over 400 steps with subspace updates every 200 steps, explaining the ~3 second average (two SVD calls amortized over 200 steps).
Weight quantization combination (Table 8): Integrating INT8 weight quantization (following Q-GaLore, Zhang et al., 2024c) with a group size of 128: Q-APOLLO achieves 31.97 perplexity on 60M (vs. 31.55 without quantization), 24.16 on 130M (vs. 22.94), and 18.79 on 350M (vs. 16.85). Q-APOLLO-Mini achieves 33.05 on 60M (vs. 31.93), 24.70 on 130M (vs. 23.84), and 18.90 on 350M (vs. 17.18). Quantization does introduce a performance gap (e.g., 16.85 → 18.79 on 350M for APOLLO), but the quantized versions still outperform full-precision AdamW (34.06, 25.08, 18.80 on 60M/130M/350M respectively) while dramatically reducing total memory: Q-APOLLO-Mini uses 0.06 GB on 60M, 0.12 GB on 130M, and 0.35 GB on 350M for optimizer states (the memory column in Table 8 appears to report optimizer state memory only, not total memory; the paper's claim of "pre-training of a LLaMA-7B model using just 12 GB of memory" with Q-APOLLO-Mini assumes layer-wise gradient updates).
LLaMA-13B pre-training on single A100-80GB: The paper states this is enabled by APOLLO-Mini "without requiring other system-level optimizations, such as model sharding" but does not provide perplexity numbers for this configuration—it is presented as a feasibility demonstration rather than a performance benchmark.
Ablation Studies and Robustness Checks
Random projection vs. SVD (Figure 5a–c): On LLaMA-60M, 130M, and 350M pre-training, GaLore's performance degrades significantly when SVD is replaced with random projection—it fails to match the AdamW baseline (red dashed line) in all cases. In contrast, APOLLO and APOLLO-Mini show nearly identical performance with random projection vs. SVD. On LLaMA-350M, APOLLO-Mini with random projection actually outperforms the SVD variant. This is the single most important ablation for establishing that APOLLO's mechanism (norm ratio estimation) does not depend on finding the optimal low-rank subspace. The paper also provides the SVD variant numbers in Table 2: APOLLO w. SVD achieves 31.26 vs. APOLLO 31.55 on 60M (0.29 difference), 22.84 vs. 22.94 on 130M (0.10), 16.67 vs. 16.85 on 350M (0.18)—well within the noise of different random seeds.
Rank sensitivity (Figure 5d): On LLaMA-60M, GaLore's perplexity degrades monotonically as rank decreases: at rank = 128 (n/4), it matches AdamW (~34); at rank = 64, perplexity rises to ~35; at rank = 32, to ~36; and at rank = 16, to ~37. Fira follows a similar trajectory but with better absolute performance. APOLLO is remarkably flat across ranks: from rank = 128 down to rank = 32, perplexity changes by less than 1 point. APOLLO-Mini at rank = 1 actually achieves better perplexity than APOLLO at rank = 32 or 64, validating the claim that tensor-wise averaging at rank-1 is more effective than channel-wise scaling at very low rank. The red dashed line (AdamW full-rank) is comfortably above the APOLLO-Mini point, confirming that rank-1 outperforms full-rank AdamW.
Scaling factor granularity (Table 9): On LLaMA-60M/130M/350M at rank = n/4, the difference between channel-wise and tensor-wise scaling is small. Channel-wise APOLLO: 31.55, 22.94, 16.85. Tensor-wise APOLLO: 32.10, 23.82, 17.00. The gap is 0.55, 0.88, and 0.15 perplexity respectively—all well within the range where both variants outperform AdamW (34.06, 25.08, 18.80). This supports the claim that tensor-wise scaling is sufficient at modest rank. However, at rank-1 (Figure 5d), channel-wise fails while tensor-wise (APOLLO-Mini) succeeds—the granularity interacts with rank in a non-trivial way that the paper attributes to variance in the channel-wise norm estimates at extreme low rank.
Training trajectory comparison with Fira (Figure 6, Section 5.4 A4): On LLaMA-350M with 60K training steps, Fira converges faster in early training (lower perplexity at steps 0-20K), but APOLLO gradually catches up and crosses over at approximately step 35K, maintaining lower perplexity through the end of training. The zoomed-in panels (early, middle, late) show this crossover clearly. The paper also quantifies this on LLaMA-130M with extended training: at 20K steps, Fira achieves 22.73 vs. APOLLO's 22.84; at 30K steps, Fira achieves 21.69 vs. APOLLO's 21.71—the gap narrows from 0.11 to 0.02 perplexity. These results suggest that APOLLO benefits more from additional training tokens than Fira, consistent with the optimization dynamics hypothesis (Section 5.5) that the "SGD-like" update becomes advantageous in later stages.
Long-context training (Figure 7, Section 5.4 A5): On LLaMA-350M with sequence length 1024 (4× longer than GaLore's standard 256), both APOLLO and APOLLO-Mini outperform AdamW across the training trajectory, with the gap widening in later steps. This experiment uses a stronger AdamW baseline (learning rate swept across [1e-3, 2.5e-3, 5e-3, 7.5e-3, 1e-2]) while APOLLO-series uses fixed learning rate 1e-2 with lazy scale factor tuning. The paper notes this demonstrates APOLLO's suitability for "industrial LLM pre-training settings(long sequences and extensive training tokens)."
Directional sharpness (Table 10, Section 5.5): On a T5 model for machine translation (following Pan & Li, 2023), directional sharpness—the quantity v^T ∇^2 f(x) v where v is the normalized update direction—is measured across optimizers. At epoch 2: SGD 1.96, Adam 0.0092, APOLLO 0.0060, APOLLO-Mini 0.0040. At epoch 20: SGD 3.21, Adam 0.00040, APOLLO 0.00026, APOLLO-Mini 0.00010. Both APOLLO variants achieve lower sharpness than Adam at every epoch, with APOLLO-Mini being the lowest. This is offered as a mechanistic explanation for why APOLLO's "SGD-like" update generalizes better: it finds flatter minima. The paper appropriately caveats that this is on a small-scale task, not the LLaMA/C4 setting.
APOLLO-Mini scale factor sensitivity (Section 5.4 A5): For long-context training, APOLLO-Mini's scale factor is varied in [\sqrt{128}, \sqrt{256}, \sqrt{384}] while APOLLO's is varied in [\sqrt{1}, \sqrt{2}, \sqrt{3}]. The paper does not report the sensitivity of final perplexity to these choices, only that the reported values are the best from this sweep. This leaves open the question of how much tuning is required for the scale factor in practice.
Norm-growth limiter effect (Figure 3): On LLaMA-130M, channel-wise AdamW without the norm-growth limiter shows a large loss spike in early training (orange curve), while adding the limiter (green curve) eliminates the spike and leads to better final perplexity (24.11 vs. 24.43 without NL, vs. 25.08 for element-wise AdamW). The ablation establishes that the NL is necessary for stable training with structured scaling, particularly in early stages when gradients can be explosive.
Compatibility with layer-wise gradient updates: Not presented as an ablation, but the memory numbers in Figure 1 (middle) and the 12 GB LLaMA-7B claim both assume the layer-wise gradient update strategy (Lv et al., 2023). This means the memory numbers reported are peak memory with layer-wise updates enabled, not the memory required to store all optimizer states simultaneously. The paper does not report memory without layer-wise updates, making this an implicit assumption in all system-level benefit claims.
Critical Assessment
Do the experiments support the claim that APOLLO achieves "AdamW-level performance" with "SGD-like memory"?
The performance claim: largely supported, with important nuances. Table 2 is the central evidence: APOLLO and APOLLO-Mini achieve better perplexity than AdamW on all model sizes from 60M to 1B. The strongest result is APOLLO-Mini (tuned) on LLaMA-1B: 13.95 vs. 15.56 perplexity. This is unambiguous—APOLLO-Mini uses a rank-1 auxiliary space (essentially no optimizer state beyond the gradient buffer) and beats full-rank AdamW. The downstream task results (Table 4) corroborate that the perplexity improvement translates to real accuracy gains. However, three qualifications apply. First, all results are on a single dataset (C4) and architecture family (LLaMA). The paper appropriately acknowledges this as a limitation in Section 6 ("we leave a formal [generalization analysis] as the future work"). Second, the learning rate for APOLLO-Mini was tuned (0.02 vs. the default 0.01) to achieve the strongest results in Table 2; the paper transparently marks this with ‡, but the claim "APOLLO-Mini beats AdamW" in the abstract and introduction is based on the tuned variant, not the default. The gap between tuned and default APOLLO-Mini is non-trivial: 30.95 vs. 31.93 on 60M (0.98 perplexity). Third, the AdamW baseline uses a single learning rate (0.01, Appendix A.4) without an equivalent sweep. For the long-context experiment (Section 5.4 A5), the paper does sweep AdamW's learning rate and still shows APOLLO-series superiority, but this sweep is not reported for the main Table 2 results. A more rigorous baseline would sweep AdamW's learning rate on each model size to ensure the comparison is against the best possible AdamW configuration.
The memory claim: supported, with the layer-wise update caveat. Table 1 and the concrete numbers in Table 3 (APOLLO 1.6 GB, APOLLO-Mini 0.0 GB optimizer states vs. 8-bit AdamW 13 GB) make a compelling case. The "SGD-like memory" characterization is quantitatively justified: SGDM (SGD with momentum) stores one momentum buffer of size mn, while APOLLO-Mini stores 2n + 2 values per matrix. For a 4096 × 4096 matrix, that's 8194 values vs. 16.8 million—APOLLO-Mini's optimizer state is indeed closer to nothing than to AdamW. However, the paper's memory numbers assume the layer-wise gradient update strategy (Lv et al., 2023), which is a system-level optimization orthogonal to APOLLO. The peak memory without this strategy would be higher (all optimizer states for all layers simultaneously). The paper does not report this number, meaning the 12 GB LLaMA-7B claim is a joint result of APOLLO-Mini + layer-wise updates + INT8 quantization, not APOLLO-Mini alone. The contribution is real—APOLLO-Mini makes the layer-wise strategy viable by reducing per-layer optimizer state to nearly zero—but the "SGD-like memory" framing should be understood as applying to the per-weight-matrix optimizer state, not necessarily the total training memory without system-level optimizations.
Do the experiments support the claim that APOLLO eliminates the need for SVD and achieves higher throughput than GaLore?
Strongly supported. Figure 5(a–c) is the most important ablation in the paper: it demonstrates that APOLLO's performance is nearly identical with random projection vs. SVD, while GaLore's collapses with random projection. This is a clean, well-controlled comparison—the same models, same data, same training steps, varying only the projection method. Table 7 quantifies the practical consequence: APOLLO's optimizer step on LLaMA-7B is 0.159 seconds vs. GaLore's 2.874 seconds—an 18× speedup at the optimizer step level. The throughput comparison in Figure 1 (right) combines this with the batch size advantage: APOLLO at batch size 16 processes ~3× more tokens/second than AdamW at batch size 4. GaLore at batch size 8 sits between them. The missing comparison is APOLLO at batch size 8 vs. GaLore at batch size 8—this would isolate the optimizer step overhead from the batch size effect. Since the paper shows APOLLO can support larger batches, the fair comparison is at each method's maximum batch size (which is what throughput measures), but a comparison at equal batch size would clarify how much of the gain is from the optimizer vs. from increased parallelism.
Do the experiments support the claim that APOLLO enables training on previously infeasible hardware configurations?
Partially supported, with missing performance data for the headline claims. The most dramatic claim—LLaMA-7B pre-training in 12 GB—is presented as a memory breakdown (Figure 1, middle) and Table 8 (memory numbers), but no perplexity or training curve is reported for the Q-APOLLO-Mini LLaMA-7B configuration. The paper shows Q-APOLLO-Mini perplexity for models up to 350M (Table 8), where it achieves 18.90 vs. AdamW's full-precision 18.80—a small degradation. Extrapolating to 7B is reasonable but unverified. Similarly, the LLaMA-13B on single A100-80GB claim is stated in Section 5.3 without any training results. These are "feasibility demonstrations" rather than validated training runs, and their placement as major contributions (abstract, Section 5.3 conclusion) overstates what was experimentally verified.
Do the experiments support the directional sharpness hypothesis for APOLLO's generalization advantage?
Suggestive but not rigorous. The directional sharpness measurements (Table 10) are on a different model (T5), different task (translation), and different scale (small, unspecified T5 variant) than the main LLaMA/C4 experiments. The paper acknowledges this as "preliminary insights" and defers formal treatment to future work, which is appropriate. However, the claim in Section 5.5 that APOLLO "reconciles SGD's generalization benefits with AdamW's convergence speed" is an interpretation that goes beyond the evidence. The paper demonstrates that APOLLO generalizes better than AdamW on LLaMA/C4, but the causal link to SGD-like directional noise is not experimentally established—there is no ablation that varies the amount of directional noise while controlling for other factors, and no measurement of sharpness on the LLaMA models themselves. The sharpness hypothesis is a plausible and well-motivated narrative, but it remains a hypothesis.
What experiments would strengthen the paper?
1. AdamW learning rate sweep on main Table 2. The paper sweeps AdamW's learning rate for the long-context experiment but not for the main pre-training results. This is a standard baseline-strengthening practice that would eliminate the possibility that APOLLO's advantage is partly from a suboptimal AdamW learning rate.
2. Multiple random seeds / confidence intervals. No result in the paper reports standard deviations. The training of LLaMA models involves substantial randomness (weight initialization, data order, dropout), and the differences between methods are sometimes small (e.g., APOLLO 14.20 vs. Fira 14.31 on 1B). Without error bars, it is unclear which differences are statistically significant.
3. Training curve and perplexity for the headline memory claims. Training LLaMA-7B with Q-APOLLO-Mini and reporting the validation perplexity curve (as was done for full-precision APOLLO in Figure 2) would convert the 12 GB claim from a feasibility demonstration to a performance benchmark. Similarly, reporting any training result for LLaMA-13B with APOLLO-Mini would support that claim.
4. Comparison with Adam-mini and other concurrent work. The paper positions itself relative to Adam-mini conceptually (Section 5.5) but provides no empirical comparison. Adam-mini reduces only second-moment memory; APOLLO reduces both moments, so a direct memory-vs-perplexity comparison would strengthen the case that APOLLO's unified approach is superior.
5. Ablation on the norm-growth limiter's sensitivity. The paper uses γ = 1.01 from Fira without ablating this value. Given that the limiter is necessary for stable training (Figure 3), understanding how sensitive APOLLO is to this threshold would inform practical adoption. Similarly, the scale factor α for APOLLO-Mini is set to √128 without systematic study of alternatives beyond the brief sweep mentioned in Section 5.4 A5.
6. Evaluation on a non-LLaMA architecture (e.g., GPT-2, Pythia, or a vision transformer). The paper argues APOLLO's design principles are general, but all experiments use LLaMA. Testing on at least one other architecture would begin to establish generality.
Overall assessment
The experiments provide strong evidence for the paper's core technical claims: (1) channel-wise/tensor-wise gradient scaling suffices for LLM training (Figure 3, Table 9), (2) this scaling can be accurately approximated from a randomly projected low-rank gradient (Figure 5a–c, Figure 4), (3) the resulting optimizer matches or exceeds AdamW's pre-training perplexity while using dramatically less memory (Table 2), and (4) the memory savings translate to real throughput gains and enable larger batch sizes (Figure 1, Figure 2). The main results are internally consistent across model scales (60M to 7B) and show a pattern where APOLLO's advantage often grows with model size and training duration (Section 5.4 A4, Table 3).
The principal weakness is the gap between the validated results and the strongest rhetorical claims. The "12 GB LLaMA-7B" and "LLaMA-13B on single A100" claims are presented as major contributions but lack corresponding training curves or perplexity numbers. The directional sharpness hypothesis for APOLLO's success is measured on a different task and model family, providing suggestive but not dispositive evidence. The single-dataset, single-architecture evaluation leaves generality an open question.
These weaknesses do not undermine the paper's central contribution—a novel optimizer that eliminates SVD, achieves SGD-like memory, and matches AdamW's performance on LLaMA pre-training—but they mean the paper should be read as establishing APOLLO's promise on LLaMA/C4 rather than proving its universality across all LLM training scenarios. The strongest, most rigorously supported claim is: on LLaMA models from 60M to 7B pre-trained on C4, APOLLO with random projection matches or exceeds AdamW's validation perplexity and downstream accuracy while reducing optimizer state memory by 16–4096× and enabling 3× higher throughput through larger batch sizes. Everything beyond this is extrapolation that awaits further validation.
6. Limitations and Trade-offs
The Cost of Difficulty Estimation Is Unaccounted for and Dominates the Compute Budget
The assumption or constraint. In the prior paper analysis framework, difficulty estimation was identified as a major unaccounted cost. Applying the same lens here: APOLLO requires no explicit difficulty estimation, but it introduces a different form of unaccounted overhead—the random projection and auxiliary optimizer state maintenance steps that occur at every training iteration. While the paper meticulously reports optimizer memory savings (Table 1, Table 3), it does not report the wall-clock time overhead of the APOLLO mechanism relative to plain SGD at equal batch sizes. The optimizer step time comparison in Table 7 shows APOLLO (0.159s) is faster than AdamW (0.173s) and dramatically faster than GaLore (2.874s) on LLaMA-7B, but this comparison is made at batch size 4—the maximum AdamW supports. The paper's headline throughput gains (3× over AdamW) come primarily from APOLLO's ability to use 4× larger batch sizes (16 vs. 4), not from the optimizer step being faster per se. The paper does not isolate the per-iteration overhead of the random projection and low-rank optimizer update relative to a baseline that already supports large batch sizes (e.g., plain SGD with momentum).
The consequence. A practitioner with a fixed compute budget who is not memory-constrained (e.g., training a modestly sized model that fits comfortably in GPU memory with AdamW) might see no throughput benefit from APOLLO and potentially a small slowdown. The paper's throughput gains are contingent on being memory-bound with AdamW—a common but not universal scenario. If a model already fits with AdamW at the practitioner's desired batch size, switching to APOLLO would add the overhead of Steps 1-3 in Algorithm 1 (random projection GEMM, low-rank AdamW update, scaling factor computation, norm-growth limiter) without the countervailing benefit of enabling a larger batch. The paper does not provide an ablation showing APOLLO's throughput at the same batch size as AdamW for all model scales—Table 7 provides this at batch sizes 16 (1B) and 4 (7B), showing APOLLO is slightly faster than AdamW on 7B (0.159s vs. 0.173s) but slightly slower on 1B (0.051s vs. 0.036s). The 1B result suggests the overhead is not always negligible and depends on the ratio of projection cost to optimizer state update cost.
What evidence exists in the paper. Table 7 provides the optimizer step time breakdown at identical batch sizes. On LLaMA-1B with batch size 16, the backward pass dominates (1.069s) and APOLLO's optimizer step (0.051s) is 0.015s slower than AdamW's (0.036s)—a 42% relative increase in optimizer step time, though a negligible 1.4% increase in total step time. On LLaMA-7B with batch size 4, APOLLO is actually faster (0.159s vs. 0.173s), which the paper attributes to the reduced cost of maintaining low-rank rather than full-rank moments. The crossover point where APOLLO's optimizer step becomes cheaper than AdamW's is not characterized. The per-iteration compute cost of the random projection (a dense GEMM of size r × m × n) relative to the gradient computation itself is also not analyzed—as model dimensions grow, this GEMM cost scales as O(r m n), which could become non-trivial for very wide layers.
Mitigation status. The paper partially addresses this by noting that the APOLLO series "incur minimal overhead by relying only on cheap random projections" (Section 5.3) and that "projecting gradients into a low-rank space reduces the overhead of maintaining first-order and second-order moments" (Table 7 discussion). The finding that APOLLO is faster than AdamW on 7B suggests the crossover favors APOLLO at practical LLM scales. However, the paper does not provide a systematic study of how the per-iteration overhead scales with model dimension, rank, and batch size, leaving practitioners to benchmark this for their specific configuration.
Generality Is Unproven: All Results Are on LLaMA Architectures and the C4 Dataset
The assumption or constraint. Every experiment in the paper—pre-training from 60M to 7B, fine-tuning on common-sense reasoning and MMLU, throughput measurements, memory breakdowns—uses the LLaMA architecture family. The pre-training data is exclusively C4. The paper does not conduct experiments on any other architecture (GPT-2/3, Pythia, OPT, Falcon, MPT), any other modality (vision transformers, speech models, multi-modal architectures), or any other pre-training dataset (The Pile, RedPajama, RefinedWeb, multilingual corpora). The paper's theoretical justification (Theorems 4.1–4.4) makes no architectural assumptions—the bounds on norm preservation under random projection apply to any matrix—but the empirical validation is entirely within the LLaMA/C4 ecosystem. The paper acknowledges this implicitly by stating "we leave a formal [generalization analysis] as the future work" (Section 5.5, regarding the sharpness hypothesis), but this caveat is not extended to the broader architecture/data scope.
The consequence. Several design choices in APOLLO may have architecture-specific interactions that are invisible in the current evaluation. LLaMA uses grouped-query attention (in larger variants) and SwiGLU activations in the feed-forward layers, producing weight matrices with specific dimensionalities (e.g., the gate and up projections in SwiGLU are typically hidden_dim × intermediate_dim). APOLLO's definition of "channel" as the larger dimension of each weight matrix (Section 3.2) has different implications depending on whether the matrix is tall (m ≪ n), square (m ≈ n), or wide (m » n). In LLaMA's attention projections (Q, K, V, O), the dimensions are hidden_dim × (num_heads × head_dim) or vice versa, creating a mix of shapes. The paper does not analyze whether APOLLO's performance is uniform across these shape regimes or whether certain layer types benefit disproportionately. On a different architecture with different weight matrix shapes—e.g., a convolutional network where "channels" have a different physical meaning, or an encoder-decoder model with cross-attention—the channel-wise scaling assumption may interact differently with the optimization dynamics.
The C4 dataset, while large and diverse, is English-only and web-scraped with standard filtering. Models trained on code-heavy corpora (e.g., The Stack), multilingual data, or domain-specific text (scientific articles, legal documents) may have different gradient statistics—potentially different Hessian spectra, different levels of gradient sparsity, or different channel-wise variance patterns—that affect APOLLO's scaling factor approximation quality. The paper provides no evidence that the √n/r scaling factor relationship holds outside the C4/LLaMA regime.
What evidence exists in the paper. The paper is transparent about its experimental scope: the abstract states "We conduct extensive experiments across different model architectures and tasks," but "different model architectures" refers to LLaMA at different scales (60M through 7B, plus LLaMA-3 variants for fine-tuning), not architecturally distinct families. The fine-tuning experiments (Tables 5 and 6) expand to LLaMA-3.2-1B, LLaMA-3-8B, Gemma-7B, and Mistral-7B—four model families—but only for fine-tuning, not pre-training. Gemma and Mistral have architectural differences from LLaMA (e.g., Gemma uses GeGLU activations; Mistral uses sliding window attention and grouped-query attention), and APOLLO generally performs well on these, though APOLLO-Mini drops on Gemma-7B MMLU (31.67% vs. AdamW's 34.21%—a 2.54 percentage point gap). This is a hint that APOLLO-Mini's extreme compression may interact with architectural choices, but the evidence is thin (one data point) and confined to fine-tuning.
Mitigation status. Not addressed. The paper presents APOLLO as a general-purpose LLM optimizer and makes no caveat about architecture dependence. The title ("SGD-like Memory, AdamW-level Performance") and abstract do not qualify the scope. The Gemma-7B APOLLO-Mini result is noted in passing but not discussed as a potential generalization limitation. No ablation studies vary architectural features (e.g., activation function, normalization placement, attention type) to test robustness.
The Norm-Growth Limiter and Scale Factor α Introduce Un-tuned Hyperparameters That the Paper Does Not Systematically Study
The assumption or constraint. APOLLO introduces two new hyperparameters beyond standard AdamW: the norm-growth limiter threshold γ (set to 1.01, inherited from Fira without tuning) and the gradient scale factor α (set to 1 for APOLLO on models < 1B, 0.4–0.7 for larger models, and √128 for APOLLO-Mini). While the paper emphasizes that APOLLO "runs using the same learning rate 0.01 and a subspace change frequency T of 200 without tuning, following the GaLore open-sourced settings" (Appendix A.4), this understates the tuning burden: the scale factor α and the norm-growth limiter are new knobs that prior work does not have. The paper sweeps α for APOLLO-Mini in the long-context experiment (Section 5.4 A5: values of √128, √256, √384) and for APOLLO on the 1B model (values of √1/2 for high-rank, 0.4 for SVD variant), but does not report the sensitivity of final perplexity to these choices. The norm-growth limiter threshold γ = 1.01 is used in all experiments without a single ablation varying it.
The consequence. A practitioner adopting APOLLO must either trust the paper's default settings (which were chosen based on LLaMA/C4 experiments and may not transfer to other architectures, datasets, or model scales) or conduct their own hyperparameter sweep over α and γ. The α parameter is particularly subtle: it absorbs the theoretical √n/r ratio from Theorem A.4, but n varies across layers in a Transformer (attention layers have different dimensions than MLP layers), meaning a single global α may be suboptimal for some layers. The paper's use of a fixed √128 for APOLLO-Mini is a heuristic that works for the LLaMA models tested (where hidden dimensions range from 512 to 4096), but for a model with hidden dimension 8192 or 256, the optimal α may differ substantially. The paper does not provide guidance on how to set α for new model configurations beyond "it can be combined with the learning rate" (Section 4.1.2), which implies that sweeping the learning rate jointly with α is necessary—increasing the tuning burden compared to AdamW, which has a single primary hyperparameter (learning rate) once β₁, β₂ are fixed.
The norm-growth limiter's importance is demonstrated in Figure 3, where training without it produces a "significant spike at the early stage" and worse final perplexity (24.43 vs. 24.11 with NL). This means γ = 1.01 is load-bearing for APOLLO's stability, and the paper provides no evidence that this specific value (which is extremely tight—allowing only 1% norm growth per step) is optimal or even safe across different training regimes. A practitioner training with a larger learning rate, different warmup schedule, or on noisier data might need a different γ. The paper's silence on sensitivity leaves this as a potential failure mode.
What evidence exists in the paper. The norm-growth limiter's necessity is demonstrated (Figure 3), but its sensitivity is not. The scale factor α is adjusted for APOLLO on 1B models (Appendix A.4: α = √1/2 for high-rank, α = 0.4 for SVD) and for APOLLO-Mini in fine-tuning (Appendix A.5.2: α = √4 for LLaMA-3-8B and Mistral-7B, α = 1 for Gemma-7B, noted as "it exhibits higher sensitivity during fine-tuning"). These ad-hoc adjustments demonstrate that α is not a "set once" parameter—it varies with model size, architecture, and training regime. The long-context experiment (Section 5.4 A5) sweeps α but does not report the range of resulting perplexities, only that the best was selected. The subspace update frequency T = 200 is adopted from GaLore without ablation, though the paper notes that resampling is "effortlessly done" and likely insensitive since it only costs generating a new random seed.
Mitigation status. Partially addressed through the open-source release (code is available, so practitioners can reproduce the sweeps) and the explicit acknowledgment that α can be combined with the learning rate. However, the paper does not provide a principled method for setting α and γ a priori based on model properties, nor does it characterize the sensitivity of final performance to these values. The Gemma-7B fine-tuning note ("higher sensitivity") is the only acknowledgment that these hyperparameters may interact with architecture, and it is buried in Appendix A.5.2 without further analysis. This leaves the practical tuning cost of APOLLO higher than the paper's "no tuning needed" framing suggests.
The Difficulty Estimation Analogy: APOLLO Does Not Address the Training-Serving Mismatch—What Works for Pre-training May Differ from What Works for Inference or Fine-tuning
The assumption or constraint. APOLLO is designed and evaluated as a pre-training and fine-tuning optimizer. The paper demonstrates that models trained with APOLLO achieve lower perplexity and better downstream accuracy than AdamW-trained models (Tables 2, 4). However, the paper does not evaluate whether the optimization trajectory induced by APOLLO—with its structured, coarsened learning rate adaptation and SGD-like directional noise—produces models that behave differently at inference time beyond the accuracy metrics measured. The directional sharpness hypothesis (Section 5.5) suggests APOLLO converges to flatter minima, which is generally associated with better generalization. But flatter minima may also have other properties—e.g., different calibration, different sensitivity to input perturbations, different behavior under quantization or pruning—that are not captured by perplexity or zero-shot accuracy and that matter for deployed models.
The consequence. APOLLO's "SGD-like" update direction (raw gradient G_t rather than a momentum-smoothed direction) means that the effective optimization trajectory is noisier than AdamW's, even though the per-channel step sizes are adapted. This noisier trajectory might produce models with different loss landscape geometry—potentially beneficial for generalization (as the paper hypothesizes) but potentially detrimental for other desiderata. For example, models trained with noisier optimizers sometimes exhibit worse calibration (overconfidence on incorrect predictions) because the noise prevents the model from settling into well-calibrated confidence estimates. They may also be more sensitive to weight perturbations, which matters for post-training quantization—ironically, one of the paper's headline features (Q-APOLLO-Mini). The paper evaluates Q-APOLLO-Mini's perplexity (Table 8) but does not evaluate whether APOLLO-trained models quantize differently than AdamW-trained models—e.g., whether the weight distribution after APOLLO training has different outlier characteristics or different per-channel variance patterns that affect INT8 quantization error.
Additionally, the paper's fine-tuning results (Tables 5, 6) evaluate APOLLO starting from pre-trained checkpoints that were themselves trained with AdamW (the standard publicly available weights). This means the fine-tuning experiments test APOLLO as a fine-tuning optimizer, not the end-to-end pipeline of pre-training with APOLLO and then fine-tuning or deploying. The interaction between APOLLO pre-training and subsequent fine-tuning (with either APOLLO or AdamW) is unexplored. A model pre-trained with APOLLO-Mini's rank-1 scaling might have weight matrices with different spectral properties than one pre-trained with AdamW, which could affect how well LoRA or other PEFT methods adapt those weights.
What evidence exists in the paper. The paper provides zero-shot accuracy (Table 4) and MMLU fine-tuning results (Table 6), which cover the most common deployment scenarios. However, it provides no analysis of model calibration, robustness to input perturbations, sensitivity to quantization (beyond the Q-APOLLO perplexity numbers which evaluate quantization during training, not post-training quantization of an APOLLO-trained model), or compatibility with inference-time optimizations like KV-cache compression or speculative decoding. The directional sharpness measurements (Table 10) are the only geometric characterization of the trained models, and they are on a different architecture and task. The paper does not analyze the singular value spectrum, weight norm distributions, or activation statistics of APOLLO-trained vs. AdamW-trained models, which would reveal whether the optimization trajectory leaves artifacts that matter for deployment.
Mitigation status. Not addressed. The paper treats training performance (perplexity, downstream accuracy) as the sole metric of optimizer quality, which is standard in the optimizer literature but incomplete for LLM deployment pipelines. The paper does not acknowledge this gap. The deployment-focused claims (throughput, memory, single-GPU training) are about training efficiency, not inference-time model quality beyond accuracy. A practitioner adopting APOLLO for pre-training would need to independently verify that the resulting models behave as expected under their specific inference or fine-tuning pipeline.
The Theoretical Bounds Assume a Fixed Projection Matrix, but Resampling Every T Steps Breaks the Guarantees
The assumption or constraint. The paper's theoretical analysis (Theorems 4.1–4.4, Appendix A.1) provides high-probability bounds on the scaling factor approximation error under the assumption that the projection matrix P_t is fixed across time steps. The proofs rely on expanding the first and second moments as exponentially weighted sums over historical gradients and applying the norm preservation property of P to each term. If P changes at step T, the historical terms in the moment estimates M^R_t and V^R_t were computed with a different projection matrix, and the simple decomposition M^R_t = P M_t (used in Theorem 4.2's proof) no longer holds. The paper acknowledges this tension in the remark following Theorem 4.2: "Here, we assume the projection matrix is fixed over time step t. GaLore also derives their theorem with the same assumption. However, as acknowledged in GaLore, using the same projection matrix for the entire training may limit the directions in which the weights can grow. Therefore, empirically, as in GaLore, we can periodically resample P over T iterations to introduce new directions."
The consequence. When the projection matrix is resampled every T = 200 steps, the low-rank moment estimates M^R_t and V^R_t are hybrids—they contain contributions from multiple different projection matrices, each of which was applied to gradients from different time windows. The norm of M^R_t is no longer guaranteed to be close to the norm of P_current M_t because M^R_t is not a simple projection of the full-rank moment under a single P. The practical consequence is that the scaling factor s^R_j after a projection update may be systematically biased for the first few steps, until the exponential averaging has given sufficient weight to gradients observed under the new projection. The paper does not analyze this "cold-start" effect after resampling, nor does it study whether the optimal resampling frequency T should depend on the decay rates β₁, β₂ (which control how quickly the moment estimates forget old projections). With default β₁ = 0.9, the half-life of the first moment is approximately 6.6 steps—meaning after 200 steps, the contribution of gradients from before the last resampling is negligible. With β₂ = 0.999, the half-life is approximately 693 steps—meaning the second moment still contains substantial information from up to three resampling periods ago, all computed with different projection matrices. This is a theoretical gap: the proofs do not cover the actual algorithm with periodic resampling, and the interaction between resampling frequency and the second moment's long memory is not characterized.
What evidence exists in the paper. The empirical results strongly suggest this is not a practical problem: APOLLO with random projection and T = 200 performs as well as APOLLO with SVD (Figure 5a–c) and better than GaLore with the same resampling schedule. The paper also varies T implicitly by using different training lengths (10K to 150K steps) with the same T = 200, meaning the ratio of resampling events to total steps varies, and APOLLO still performs well. However, this evidence is correlational—it shows that the algorithm works, not that the theoretical guarantees hold. A practitioner who wanted to use a much shorter resampling period (e.g., T = 10) or much longer one (e.g., T = 1000) would be operating outside the empirically validated regime. The paper provides no ablation of T, stating only that it is "set to 200 by default" and that "empirically, as in GaLore, we can periodically resample P."
Mitigation status. The paper is transparent about the gap between theory and practice in the Theorem 4.2 remark, which is good scholarship. However, it offers no analysis of the resampling cold-start effect, no characterization of how the moment estimates behave immediately after resampling, and no sensitivity study of T. The remark that "we can simply re-sample P from the Gaussian distribution by changing the random seed" (Appendix A.1.3) treats the discrepancy as resolved by empirical success, but it leaves open the question of whether the current theoretical framework can be extended to the time-varying projection case. A rigorous analysis would need to bound the error from mixing projection matrices in the exponential moving averages—a non-trivial extension that the paper does not attempt.
The Rank-1 APOLLO-Mini "Beats AdamW" Claim Relies on a Single Tuned Learning Rate and a Heuristic Scale Factor
The assumption or constraint. APOLLO-Mini's strongest results in Table 2—the entries marked with ‡ that show APOLLO-Mini outperforming both AdamW and APOLLO at all model scales—use a tuned learning rate of 0.02 rather than the default 0.01 used for all other methods (including the base APOLLO-Mini entries without ‡). The paper states: "Here, we report the APOLLO-Mini with a learning rate of 0.02 (not 0.01 in GaLore), achieving stronger results, marked with ‡." Additionally, APOLLO-Mini uses the heuristic scale factor α = √128 (Section 4.2), which is "heuristically set... like GaLore to avoid performance degradation." The tuning of two hyperparameters (learning rate and scale factor) for APOLLO-Mini is not matched by equivalent tuning of AdamW's learning rate in Table 2—AdamW uses a single learning rate of 0.01 throughout, following the GaLore settings.
The consequence. The headline result that APOLLO-Mini (SGD-level memory) beats AdamW (full optimizer state) is confounded by unequal hyperparameter tuning effort. The gap between tuned APOLLO-Mini and default APOLLO-Mini is substantial: on LLaMA-60M, tuning improves perplexity from 31.93 to 30.95 (a 0.98 gap); on LLaMA-130M, from 23.53 to 22.85 (0.68); on LLaMA-350M, from 17.18 to 16.63 (0.55); on LLaMA-1B, from 14.17 to 13.95 (0.22). These gaps are large enough to change the ranking: default APOLLO-Mini (31.93) is worse than APOLLO (31.55) on 60M, while tuned APOLLO-Mini (30.95) is better. If AdamW's learning rate were similarly swept—and the long-context experiment in Section 5.4 A5 shows the paper is willing to sweep AdamW's learning rate when probing a specific question—would the gap narrow or close? The paper does not answer this for the main pre-training results. The long-context experiment (Figure 7) does compare against a learning-rate-swept AdamW and still shows APOLLO-Mini winning, but this is a different training regime (sequence length 1024 vs. 256 for Table 2), and the AdamW sweep in that experiment is explicitly to "establish a strong baseline," which raises the question of why the same strong baseline wasn't used for the main results.
The heuristic scale factor α = √128 compounds this issue. The theoretical ratio is √n/r (Theorem A.4), which would be √n for the rank-1 case. For a 4096-dimensional weight matrix, that's √4096 = 64; for a 512-dimensional matrix, it's √512 ≈ 22.6. Using a fixed √128 ≈ 11.3 is a compromise that is suboptimal for both small and large matrices. The paper does not report what performance would be with the theoretically motivated √n (per-layer scale factors) vs. the heuristic fixed √128. The long-context experiment sweeps over √128, √256, √384 and selects the best, but the sensitivity of final perplexity to this sweep is not quantified—we only know the best was selected, not whether the range of outcomes is 0.1 perplexity or 2.0 perplexity. This makes it difficult for a practitioner to assess how carefully α must be tuned for their specific model.
What evidence exists in the paper. The paper is transparent about the tuning: the ‡ marker is explicitly defined in the Table 2 caption. The long-context experiment (Section 5.4 A5) explicitly states that APOLLO-Mini's scale factor was varied. The Appendix A.4 notes that α is adjusted for the 1B model and that APOLLO-Mini uses √128. This transparency is commendable. However, the paper's abstract and introduction make the unqualified claim that "APOLLO-Mini... achieves SGD-level memory cost while outperforming AdamW" without noting the tuning caveat. The abstract does mention tuning for the ‡ results but the rhetorical framing ("outperforms AdamW") doesn't distinguish between the default and tuned variants.
Mitigation status. Partially addressed by the transparency in Table 2's caption and Appendix A.4, and by the long-context experiment which includes a learning-rate-swept AdamW baseline. However, the main Table 2 results—which are the paper's most important quantitative evidence—still compare a tuned APOLLO-Mini against an un-tuned AdamW. A fair comparison would either tune both methods' primary hyperparameters (learning rate for AdamW; learning rate and α for APOLLO-Mini) or present both with their defaults. The long-context experiment demonstrates that the conclusion holds with a tuned AdamW in that specific setting, but a similar experiment for the standard pre-training setting (Table 2 parameters) would significantly strengthen the claim. The paper's reliance on "following the GaLore open-sourced settings" as justification for not tuning AdamW is pragmatic but methodologically weak—GaLore's choice of learning rate for AdamW was itself not necessarily optimal, and inheriting this choice propagates any baseline weakness.
7. Implications and Future Directions
How This Work Changes the Landscape
APOLLO reframes the memory-efficient LLM training problem from "how do we compress AdamW's states without losing information?" to "how much of AdamW's state is actually necessary, and at what granularity?" This shift is more than a new algorithm—it is a diagnostic finding that the element-wise precision in AdamW's learning rate adaptation is largely redundant for Transformer training, and that a coarser, structured scaling factor estimated from a heavily compressed gradient view is not just adequate but often superior.
The magnitude of this shift is substantial but bounded. It is not a paradigm shift in optimization theory—the paper does not introduce new convergence proofs or a fundamentally new class of optimizers—but it is a practical reframing with immediate engineering consequences. Prior to APOLLO, the dominant approach to memory-efficient LLM training (GaLore, Fira, and their variants) treated gradient low-rank structure as the property to exploit and SVD as the necessary tool. This created a tension: memory savings from low-rank projection were offset by the computational cost of SVD and the need for relatively high ranks. The field was implicitly searching for cheaper alternatives to SVD (online PCA, lazy updates, random projections with SVD warm-start) while still operating within the "gradients are low-rank" framework.
APOLLO demonstrates that this framework was solving a harder problem than necessary. By showing that random projection suffices—and that the critical property is norm preservation under random embedding, not spectral alignment—APOLLO eliminates the SVD bottleneck entirely without sacrificing optimization quality. This makes several prior research directions less attractive:
- Further optimization of SVD for training: Incremental SVD, randomized SVD, and batched SVD for gradient subspace updates are now solutions to a problem that APOLLO avoids altogether. The paper's empirical demonstration that GaLore's performance collapses with random projection while APOLLO's does not (Figure 5a–c) means SVD is a requirement of the GaLore approach, not a general necessity for low-rank optimization. Researchers working on faster SVD approximations for training should consider whether the real bottleneck was the SVD-dependent design, not the SVD implementation.
- High-rank requirements for low-rank optimizers: GaLore needs
r = n/4to match AdamW; APOLLO withr = n/8loses only 0.05 perplexity on LLaMA-1B (Table 2). This suggests that prior work's rank requirements were driven by the need to preserve gradient direction in the low-rank space, not by an inherent requirement of adaptive optimization. The design principle that APOLLO establishes—decouple update direction from update step size, and only approximate the latter—makes the rank requirement much less stringent. - Component-specific optimizer designs: Adam-mini's block-wise second moment requires "careful handling of different model components" (Section 2), with different block sizes for attention vs. MLP layers. APOLLO's channel-wise and tensor-wise scaling is applied uniformly across all weight matrices regardless of their role, and the paper shows it works across all LLaMA components without per-module tuning. This suggests that the complexity of Adam-mini's design was compensating for the fact that it only coarsened one of the two moments; coarsening both simultaneously eliminates the need for component-specific handling.
The paper also resolves a latent tension in the memory-efficient training literature. Prior work had fragmented into two streams that seemed complementary but were never unified: (1) low-rank gradient projection (GaLore et al.), which compressed optimizer states by working in a smaller space, and (2) optimizer state redundancy reduction (Adam-mini), which coarsened the granularity of adaptation. APOLLO synthesizes these by showing that coarsening enables more aggressive compression: once the scaling factor is structured (channel-wise or tensor-wise), it can be estimated from a much lower-rank gradient representation than the gradient direction itself requires. This synthesis is likely to redirect the field toward joint granularity-rank optimization rather than treating them as independent axes.
The most provocative implication is the generalization hypothesis (Section 5.5): that coarsening the learning rate adaptation and injecting SGD-like directional noise actually improves final model quality by biasing optimization toward flatter minima. If this hypothesis holds up under broader validation, it would invert the conventional wisdom that adaptive optimizers are strictly better for Transformers and that memory efficiency comes at a performance cost. APOLLO-Mini's perplexity results (Table 2) and the directional sharpness measurements (Table 10) are preliminary evidence, but the claim that "less optimizer state can mean better generalization" would, if validated, change how the field thinks about the role of optimizer memory in deep learning—not as a necessary evil to be compressed, but as a potential source of harmful overfitting to gradient noise.
However, the work does not change the landscape for problems outside the base model's capability range. The paper's findings are confined to LLaMA architectures pre-trained on C4; there is no evidence that APOLLO's structured scaling applies to other architectures (convolutional networks, mixture-of-experts, state-space models) or modalities (vision, speech, multi-modal). The theoretical bounds (Theorems 4.1–4.4) are architecture-agnostic, but the empirical validation is narrow. Researchers working on non-Transformer architectures or non-language domains should view APOLLO as a promising hypothesis to test, not a proven solution.
Follow-Up Research This Work Enables
1. Cheap, per-layer difficulty estimation for dynamic rank allocation during training. APOLLO's core mechanism estimates a scaling factor from a low-rank projection, and the theoretical bounds (Theorem A.4) show that the estimation error depends on the projection rank r. The paper uses a fixed global rank for all layers, but the required rank for accurate scaling factor estimation likely varies across layers and across training stages—some layers may have more structured gradient statistics and tolerate lower rank, while others need higher rank. A natural follow-up is to dynamically allocate rank per layer based on the observed variance of the scaling factor estimate. Concretely: maintain a running estimate of the variance of s^R_j over recent steps for each layer; if the variance is low, reduce the rank for that layer (saving memory and compute); if the variance spikes (suggesting the current rank is insufficient), increase it. This would produce a compute-optimal test-time equivalent for training: allocate the "rank budget" where it provides the most benefit. A strong follow-up would measure the Pareto frontier of total optimizer state memory vs. validation perplexity for dynamic vs. static rank allocation on LLaMA-7B, and characterize whether the optimal rank distribution correlates with known layer properties (e.g., early vs. late layers, attention vs. MLP, or the Hessian spectral properties identified in Zhang et al., 2024a).
2. Systematic evaluation of APOLLO-trained model properties beyond perplexity: calibration, robustness, and quantizability. The paper demonstrates that APOLLO achieves lower perplexity and better zero-shot accuracy than AdamW (Table 4), but provides no analysis of other model properties that matter for deployment. Given the directional sharpness hypothesis—that APOLLO converges to flatter minima with different loss landscape geometry—there may be downstream consequences for model behavior. A strong follow-up would evaluate: (a) Calibration: expected calibration error (ECE) on common-sense QA tasks for APOLLO-trained vs. AdamW-trained LLaMA models. Flatter minima are sometimes associated with worse calibration because the model is less "committed" to its predictions. (b) Adversarial robustness: accuracy under input perturbations (e.g., contrast sets, adversarial attacks on BoolQ or HellaSwag) to test whether APOLLO's SGD-like directional noise produces more brittle or more robust decision boundaries. (c) Post-training quantization sensitivity: the paper already combines APOLLO with quantization during training (Q-APOLLO-Mini, Table 8), but does not evaluate how an APOLLO-trained full-precision model behaves under post-training quantization (PTQ). If APOLLO produces weights with different outlier characteristics or per-channel variance patterns, PTQ error could differ from AdamW-trained models. This follow-up would directly inform the deployment claims the paper makes—if APOLLO-trained models quantize poorly, the Q-APOLLO approach (quantization-aware training) becomes a requirement, not an optional enhancement.
3. Testing the generalization hypothesis through controlled optimization trajectory interventions. Section 5.5 proposes that APOLLO's advantage comes from injecting SGD-like directional noise while retaining AdamW-like per-channel learning rate adaptation. This hypothesis is testable through a controlled experiment: start with APOLLO, then progressively reduce the directional noise by blending the raw gradient G_t with the AdamW-style momentum-smoothed gradient M_t in the update direction, while keeping the channel-wise scaling factor estimation identical. If the hypothesis is correct, reducing directional noise should improve convergence speed (lower perplexity early in training, similar to Fira's early advantage in Figure 6) but worsen final generalization (higher perplexity at convergence, sharper minima as measured by directional sharpness). Conversely, increasing directional noise beyond APOLLO's default (e.g., by adding explicit Gaussian noise to the update direction) should further flatten the minima but eventually hurt convergence. A strong follow-up would sweep the interpolation ratio λ where the update direction is (1 - λ) * G_t + λ * M_t and the scaling factor is APOLLO's s^R, measuring both the training loss trajectory and the final model's directional sharpness (following the protocol from Table 10 but on LLaMA models). This would provide causal evidence for the generalization hypothesis and characterize the optimal noise level—a finding that would be independently useful beyond APOLLO.
4. Extending APOLLO to non-LLaMA architectures and long-context regimes to identify failure modes. The paper's experiments are confined to LLaMA architectures (60M–7B) and C4, with a single long-context experiment (Section 5.4 A5) at sequence length 1024. Two stress-tests would clarify the boundaries of APOLLO's applicability. First, evaluate on a non-LLaMA Transformer: GPT-2 (pre-layer norm, no SwiGLU, no RoPE) or a vision transformer (ViT on ImageNet) would test whether the channel-wise scaling assumption is tied to LLaMA's specific weight matrix shapes and activation functions. The SwiGLU feed-forward layers in LLaMA have gate_proj and up_proj matrices where the "channels" correspond to intermediate dimensions with gating; this structure may interact with APOLLO's column-wise scaling in ways that a standard ReLU MLP does not. Second, stress-test extreme sequence lengths: pre-train LLaMA-350M with sequence lengths of 2048, 4096, and 8192, measuring whether APOLLO's scaling factor estimation degrades with longer sequences. Longer sequences produce gradients with different variance structures (more tokens to average over), potentially making the low-rank s^R estimate noisier or more biased. If APOLLO-Mini's rank-1 tensor-wise scaling fails at very long contexts but APOLLO with modest rank succeeds, that defines a practical boundary for the extreme memory-saving variant. Each stress-test should report validation perplexity, throughput relative to AdamW, and optimizer state memory, with the explicit goal of finding where APOLLO stops working.
5. Developing a principled hyperparameter selection method for the scale factor α and norm-growth limiter threshold γ. APOLLO introduces two new hyperparameters—the gradient scale factor α and the norm-growth limiter threshold γ—that the paper tunes heuristically or inherits from prior work without sensitivity analysis. The theoretical ratio α ∝ √(n/r) (Theorem A.4) provides a functional form but not a coefficient: should α be exactly √(n/r), or √(n/r) times a constant, or something else? A systematic empirical study on LLaMA-350M would sweep α as a function of √(n/r) for each layer individually (since n varies across layers) and measure the resulting per-layer gradient norm statistics over training. If the per-layer optimal α closely tracks √(n/r), that validates the theory and provides a recipe for setting α on new architectures without tuning. Conversely, if the optimal α deviates systematically—e.g., attention layers need a different multiplier than MLP layers—that reveals where the theoretical assumptions break down. Similarly, sweeping γ from 1.001 to 1.1 and measuring the tradeoff between training stability (absence of loss spikes) and convergence speed would replace the inherited γ = 1.01 with an empirically justified choice. A strong follow-up would produce a simple, theory-grounded rule (e.g., "set α = √(n/r) for all layers and γ = 1.01 for models with hidden dimension ≥ 1024; tune γ lower for smaller models") that a practitioner can apply without sweeping.
6. APOLLO for continual pre-training and domain adaptation: does the reduced optimizer state forget faster? The paper evaluates APOLLO for pre-training from scratch and for fine-tuning on downstream tasks, but does not evaluate the intermediate regime of continual pre-training—taking an existing pre-trained model (trained with AdamW) and continuing training on new data with APOLLO. This is a practically important scenario: many practitioners fine-tune or continually pre-train models on domain-specific corpora. APOLLO's low-rank optimizer state would need to "warm-start" from the existing model's gradient statistics, but the projection matrix P_t has no relationship to the original AdamW states—the first and second moments in the low-rank space start from zero. This means APOLLO would initially have no meaningful scaling factor estimates and would need to accumulate them from scratch. A strong follow-up would compare: (a) continual pre-training LLaMA-7B on a domain-specific corpus (e.g., PubMed, code) starting from a public checkpoint, using APOLLO vs. AdamW vs. GaLore, measuring both final domain perplexity and knowledge retention (perplexity on the original C4 or general benchmarks). If APOLLO's cold-start period causes it to "forget" more of the original knowledge before the scaling factor estimates stabilize, that is a significant practical limitation. The experiment would also test whether initializing the low-rank moments from a brief warmup period (e.g., running AdamW for 100 steps to establish baseline statistics, then projecting those into the low-rank space to initialize M^R_0, V^R_0) mitigates the problem.
Practical Applications and Downstream Use Cases
1. Democratized LLM pre-training on consumer and low-end datacenter GPUs. The paper's headline memory figure—LLaMA-7B pre-training in under 12 GB with Q-APOLLO-Mini—crosses a critical threshold: this is the memory capacity of widely available consumer GPUs (RTX 3060 12 GB, RTX 4070 12 GB, RTX 3080 12 GB) and older datacenter GPUs (Tesla K80 24 GB, though split across two dies, and P40 24 GB). The practical implication is that a researcher with a single consumer GPU can run full-parameter pre-training experiments on 7B-scale models from scratch, something that previously required multiple datacenter GPUs with model parallelism. The throughput will be low—the paper does not report tokens/second for this configuration—but the capability is what matters for iterative research: testing architectural changes, new training objectives, or domain-specific pre-training without waiting for cluster allocation. The paper's finding that APOLLO-Mini with a tuned learning rate (0.02) reaches 13.95 perplexity on LLaMA-1B vs. AdamW's 15.56 (Table 2) suggests that the memory savings do not come at a model quality cost at that scale; extrapolating to 7B requires validation but the trend across model sizes (60M through 1B) is consistent.
2. Accelerating large-scale distributed pre-training by replacing optimizer-state-bound communication with compute-bound batch scaling. In distributed training with data parallelism, the optimizer state memory per GPU often determines the maximum micro-batch size, which in turn affects communication frequency (gradient accumulation steps) and throughput. APOLLO's 16× reduction in optimizer state memory (from 2mn to 2rn with r = n/4, or up to 4096× with rank-1) means that for a fixed GPU memory budget, practitioners can increase the micro-batch size—reducing the number of gradient accumulation steps and therefore the frequency of all-reduce communication. The paper demonstrates this concretely: on 8×A100-80GB, AdamW is limited to micro-batch size 4 while APOLLO supports size 16, enabling ~3× higher throughput (Figure 1, right). For very large-scale training (hundreds or thousands of GPUs), this translates to proportionally fewer communication rounds per global batch, reducing the impact of network bottlenecks and stragglers. The benefit is most pronounced when training is optimizer-state-bound rather than activation-bound—a regime common for wide Transformers with modest sequence lengths. The paper's LLaMA-7B experiment (Figure 2) shows that APOLLO completes 150K training steps within a fixed two-week budget while GaLore and AdamW do not, demonstrating that the throughput advantage compounds over full training runs.
3. Full-parameter fine-tuning with near-LoRA memory costs for multi-task and continual learning scenarios. The fine-tuning results (Tables 5 and 6) show APOLLO with rank 32 matching or exceeding full AdamW fine-tuning accuracy while using 2nr + 2 optimizer states per weight matrix—for a LLaMA-7B with n = 4096 and r = 32, this is roughly 0.26 million values per matrix vs. 33.6 million for AdamW. This makes full-parameter fine-tuning feasible in memory-constrained environments where previously only parameter-efficient methods (LoRA, adapters) could fit. The practical use case is sequential fine-tuning on multiple tasks: a practitioner can maintain separate APOLLO optimizer states (which are tiny) for each fine-tuning run without accumulating AdamW's full 2mn states per task. Alternatively, APOLLO-Mini's rank-1 memory cost is so low that a practitioner could run many parallel fine-tuning experiments in separate processes on the same GPU, each at full parameter resolution, enabling hyperparameter sweeps that would otherwise require sequential runs or reduced model fidelity. The paper's finding that APOLLO-Mini achieves 68.23% average accuracy on common-sense reasoning vs. AdamW's 68.07% (Table 5, LLaMA-3.2-1B) and 64.41% on MMLU vs. AdamW's 64.85% (Table 6, LLaMA-3-8B) suggests the memory-accuracy tradeoff is essentially flat for these tasks—the practitioner pays almost no performance cost for the memory savings.
4. Enabling on-device or edge fine-tuning for personalization and privacy-sensitive applications. The combination of APOLLO-Mini with INT8 quantization pushes LLaMA-7B training memory to 12 GB—approaching the range of edge devices and high-end laptops (MacBook Pro with M3 Max at 36 GB unified memory, high-end workstations with 24 GB GPUs, or even CPU-only training with sufficient RAM). This opens the door to on-device full-parameter fine-tuning for personalization: a user's local data (documents, emails, writing style) can fine-tune the full model without uploading data to a cloud service, preserving privacy. The practical obstacle has always been that fine-tuning a 7B model with AdamW requires ~58 GB of memory, which exceeds any consumer device; parameter-efficient methods like LoRA reduce memory but restrict the model's capacity to adapt. APOLLO-Mini removes the optimizer memory bottleneck entirely, making the model weights and activations the dominant memory consumers. With 4-bit weight quantization (not explored in the paper but a natural extension), LLaMA-7B fine-tuning could potentially fit in 8 GB of memory—the capacity of a laptop CPU. The paper's finding that APOLLO-Mini fine-tuning quality is close to full AdamW (Table 5: 68.23% vs. 68.07% average; Table 6: 64.41% vs. 64.85%) suggests the practical quality loss would be small. The open research question is whether the reduced directional smoothing in APOLLO-Mini affects personalization quality—does the SGD-like noise help or hurt adaptation to small, domain-specific datasets?
When to Prefer This Method
The paper explicitly positions APOLLO against AdamW and GaLore across memory, throughput, and performance axes, providing clear tradeoffs that I summarize as decision rules rather than a generic matrix:
-
Prefer APOLLO (with random projection, rank
n/4ton/8) when: You are pre-training or fine-tuning a Transformer-based LLM, you are memory-bound on your GPU(s) with AdamW at your desired batch size, and you want to at least match AdamW's model quality. APOLLO delivers the same or better perplexity as AdamW (Table 2: 14.20 vs. 15.56 on LLaMA-1B) while using 4–8× less optimizer memory and enabling 2–4× larger batch sizes. The APOLLO default settings (learning rate 0.01,T = 200,α = 1for models < 1B) require minimal tuning. APOLLO has no SVD overhead, making it strictly faster than GaLore/Fira at equal batch size (Table 7: 0.159s vs. 2.874s per step on LLaMA-7B). -
Prefer APOLLO-Mini (rank-1, tensor-wise scaling) when: You need extreme memory efficiency—SGD-level optimizer state—and are willing to tune the learning rate (
0.02instead of0.01in the paper's experiments) and scale factor (α = √128default, but may need per-model adjustment as indicated by the Gemma-7B sensitivity noted in Appendix A.5.2). APOLLO-Mini enables training configurations that are physically impossible with AdamW (LLaMA-13B on single A100-80GB without model sharding; LLaMA-7B in 12 GB with quantization). The model quality is competitive with or better than AdamW on the LLaMA/C4 setting (Table 2: APOLLO-Mini tuned achieves 13.95 vs. AdamW 15.56 on LLaMA-1B), but the evidence is thinner for other architectures and tasks, and the scale factorαmay require task-specific tuning. -
Still prefer AdamW when: You are not memory-bound, you are training a model architecture or on a dataset substantially different from LLaMA/C4 (the paper provides no generalization evidence), or you need the strongest possible convergence speed in early training (Figure 6 shows Fira, which maintains more complete optimizer state, converges faster initially than APOLLO). If your model or training regime is sensitive to optimizer hyperparameters and you cannot afford to sweep
αandγ, AdamW's single primary hyperparameter (learning rate) has well-characterized behavior across a wide range of settings. -
Still prefer GaLore or Fira when: You need the absolute best pre-training performance on small-to-medium models with limited training tokens, where Fira's early convergence advantage (Figure 6, Section 5.4 A4) may matter more than APOLLO's late-stage crossover. The paper shows Fira leading on LLaMA-60M (31.06 vs. APOLLO's 31.55) and being caught or passed only at larger scales and longer training. If you are training a model smaller than ~350M parameters with a limited token budget and can tolerate the SVD overhead, Fira may be preferable. Note that this advantage narrows and eventually reverses with larger models and more training (as the paper quantifies for LLaMA-130M at 20K vs. 30K steps), so the decision is scale-dependent.
