ArXiv: 2412.11768

🎯 Pitch

SGD-SaI matches or beats AdamW on Transformers, including ViTs and GPT-2, using only half the memory—no adaptive moment estimates required. The trick? A one-time, per-parameter learning rate scale derived from the gradient signal-to-noise ratio at initialization, proving that this initial statistic remains stable throughout training.

1. Executive Summary

This paper introduces SGD-SaI, a simple enhancement to stochastic gradient descent with momentum that eliminates the need for adaptive gradient methods by performing learning rate Scaling at Initialization (SaI) to distinct parameter groups based on their gradient signal-to-noise ratios (g-SNR). Evaluated on ImageNet-1K classification with Vision Transformers (ViT/S-16), GPT-2 pretraining on OpenWebText, and LoRA fine-tuning tasks for both LLMs and diffusion models, SGD-SaI consistently matches or outperforms AdamW while cutting optimizer memory usage by half—reducing memory by 5.93 GB for GPT-2 (1.5B parameters) and 25.15 GB for Llama2-7B compared to AdamW in full precision, and achieving 3× faster optimizer step times than Adam-mini on GPT-2-Small. The approach establishes that per-parameter adaptive learning rates are unnecessary when gradient statistics computed at the first training iteration remain stable throughout training, enabling SGD to successfully train Transformer architectures—a long-standing challenge—with substantially lower memory and computational overhead.

2. Context and Motivation

The Core Problem: Adam's Memory Cost Scales Poorly with Model Size

The fundamental problem this paper addresses is deceptively simple: adaptive gradient methods like Adam are memory-inefficient in a way that becomes crippling as models grow. When you train a neural network with Adam (Kingma & Ba, 2014) or AdamW (Loshchilov & Hutter, 2019), the optimizer must store two full-precision tensors for every trainable parameter: the first-order momentum mtm_t (a running average of past gradients) and the second-order momentum vtv_t (a running average of squared gradients). Together, these optimizer states consume at least twice the memory of the model parameters themselves in full-precision training. For a 7-billion-parameter model in FP32, this means roughly 50 GB of GPU memory is devoted solely to optimizer state tensors—before accounting for activations, gradients, or the parameters themselves.

This is not merely an inconvenience. It is a hard resource constraint that directly determines what models can be trained on what hardware. As Section 1 highlights, scaling model parameters from 1.5 billion to 7 billion pushes AdamW's state tensor memory from approximately 12 GB to over 49 GB (Table 1, Figure 1). On an NVIDIA A100-80GB GPU, training Llama2-7B with AdamW in full precision becomes infeasible—the paper reports out-of-memory (OOM) errors for both AdamW and Adam at this scale (Table 1). The memory-hungry Prodigy optimizer (Mishchenko & Defazio, 2023), which stores four full-size tensors (Algorithm 6 in Appendix B), requires nearly 99 GB for the same model—completely untenable on any single GPU.

The practical consequences are immediate:

  • Accessibility of research: Smaller labs and individual researchers with limited GPU resources are effectively locked out of training or fine-tuning large Transformer models with adaptive optimizers unless they resort to parameter-efficient methods or aggressive quantization.
  • Deployment costs: Even in well-resourced settings, the memory overhead forces tradeoffs between batch size, sequence length, context window, and model size—each of which directly impacts training throughput and downstream model quality.
  • Scaling trajectory: As models continue to grow (the Scaling Laws of Kaplan et al., 2020, show no signs of plateauing), the optimizer memory bottleneck will tighten further, making it a first-order concern for the field.

The Enabling Role of Adam in Transformer Training — And Why SGD Fails

The memory problem would be less pressing if practitioners could simply fall back on stochastic gradient descent with momentum (SGDM), which requires storing only a single momentum tensor per parameter (half the memory of Adam). But there is a well-documented problem: SGD consistently underperforms Adam on Transformer architectures. This is not a minor performance gap. The paper's own experiments (Table 2, "SGDM" row) show that SGDM achieves only 63.80% top-1 accuracy on ImageNet-1K with ViT-S/16, compared to 73.04% for AdamW—a gap of nearly 10 percentage points that renders SGD non-competitive for practical deployment.

Why does this gap exist? The paper draws on recent theoretical work (Zhang et al., 2024a) that identifies the root cause in the Hessian structure of Transformer architectures. Transformers exhibit significant heterogeneity in the curvature of their loss landscape across different parameter blocks: attention query-key-value (QKV) projections, feed-forward MLP layers, layer normalization parameters, and positional embeddings all have fundamentally different gradient distributions. In plain language, some parameter groups naturally receive much larger or noisier gradients than others, creating an imbalance in how quickly they learn.

Adam addresses this imbalance dynamically through its per-parameter learning rate adaptation. The second-order momentum vtv_t acts as a local gain (Hinton et al., 2012): parameters that have historically received large gradients (vtv_t is large) get their effective learning rate scaled down; parameters with small or sparse gradients (vtv_t is small) get their effective learning rate scaled up. This equalizes the learning pace across heterogeneous parameter groups, which is precisely what SGD lacks. SGD applies the same global learning rate to every parameter, so blocks with naturally larger gradient magnitudes dominate the update, while blocks with smaller gradients effectively stall.

The paper explicitly frames this in the language of learning imbalances across parameters (Section 4.1):

"adaptive gradient methods dynamically adjust the learning rate for each parameter during training. This mechanism encourages parameters with less learning history to learn more while slowing down the learning pace for parameters progressing too quickly. Essentially, it acts as a compensatory approach to address learning imbalances across parameters after they arise."

This compensatory mechanism is what makes Adam indispensable for Transformers—and what has made closing the gap with SGD a long-standing open challenge that the paper directly addresses.

Prior Approaches to Reducing Optimizer Memory: What They Fix and What They Don't

The research community has recognized the Adam memory bottleneck for years, and several families of solutions have been proposed. The paper organizes these along a spectrum of tradeoffs (Section 3, "Memory Efficient Adam"):

Factored second-moment estimates (Adafactor). Shazeer & Stern (2018) proposed Adafactor, which decomposes the full vtv_t matrix into row and column factors, reducing storage from O(d)\mathcal{O}(d) to O(d)\mathcal{O}(\sqrt{d}). This represents a dramatic memory reduction, but the paper notes that it "trades off memory savings for lower update precision"—the factored approximation is coarser than the full second-moment estimate, and Adafactor consistently underperforms AdamW in practice.

Low-bit quantization of optimizer states. Dettmers et al. (2021, 2022) introduced 8-bit Adam and subsequent 4-bit variants that store mtm_t and vtv_t in low-precision formats rather than FP32. While effective at reducing memory, the paper points out that "quantization introduces additional implementation complexity" and the performance degradation, while often small, is not zero. The approach patches the symptom (memory usage) rather than addressing the fundamental question of whether the second-order momentum is necessary at all.

Sign-based methods. Bernstein et al. (2018) and Kunstner et al. (2023) explored optimizers that use only the sign of the gradient (or sign of the momentum) rather than its full magnitude. These are extremely memory-efficient but typically come with non-trivial accuracy degradation.

Block-wise adaptive learning rates (Adam-mini). The most recent and directly relevant prior work is Adam-mini (Zhang et al., 2024b). Adam-mini's key insight is that rather than maintaining a separate vtv_t for every individual parameter, you can group parameters into blocks and share a single adaptive learning rate per block. This reduces the storage of vtv_t from O(d)\mathcal{O}(d) to O(B)\mathcal{O}(B), where BB is the number of parameter blocks (on the order of hundreds or thousands rather than billions). The paper acknowledges that Adam-mini achieves memory usage approximately 45–50% lower than standard Adam—comparable to SGD-SaI's 50% reduction.

However, the paper identifies three critical limitations of Adam-mini that motivate the search for a better solution:

  1. Additional computational cost: Adam-mini still computes vtv_t dynamically at every training step (Algorithm 5 in Appendix B). It must calculate the mean squared gradient within each block, maintain exponential moving averages, and apply bias correction—all of which adds overhead relative to the already-expensive Adam update. The paper's profiling (Table 1) shows Adam-mini's optimizer step time is 223 ms for GPT-2-1.5B compared to 138 ms for AdamW—a 60% slowdown—and 421 ms for Llama2-7B versus 68 ms and 180 ms for SGD-SaI respectively. This is a non-trivial penalty that partially offsets the memory gains.

  2. Complex partitioning requirements: Adam-mini's partitioning strategy is incompatible with PyTorch's default parameter grouping. The paper explains in Section 4.2:

    "PyTorch groups parameters such as attention QKV together, [but] Adam-mini requires finer partitions, such as by attention heads or neurons, to perform effectively, especially in Transformer-based architectures."

    This means practitioners cannot simply swap Adam-mini into an existing training pipeline—they must manually restructure how parameters are grouped, adding implementation burden and potential for error.

  3. Performance degrades at scale: The paper reports (Section 5.1, "Efficiency" paragraph) that "for models larger than 1 billion parameters, the performance gains from Adam-Mini decrease by approximately 45%, while the reduction achieved with SGD-SaI remains around 50%." This suggests Adam-mini's block-wise approximation loses fidelity as models grow, undermining its value proposition at the scales where memory efficiency matters most.

Collectively, these methods share a common design philosophy: retain the adaptive, time-varying learning rate mechanism but approximate or compress it to save memory. They all maintain some form of vtv_t that is updated at every step. The paper's central departure is to question whether this adaptive mechanism is needed at all.

An Alternative Path: Static Learning Rate Scaling Based on Initialization-Time Gradient Statistics

The paper draws inspiration from a different body of work: research showing that gradient patterns at initialization carry durable information about architecture and training dynamics. This is introduced in Section 2 ("Gradient at Initialization") and elaborated in Section 4.2.

Two research threads are particularly influential:

Pruning at Initialization (PaI). Methods like SNIP (Lee et al., 2018), GraSP (Wang et al., 2020), and SynFlow (Tanaka et al., 2020) demonstrate that you can identify which parameters are important—and prune the rest—using only gradient information from the very first training step (or even before any training data is seen). The subnetworks identified by these initial gradients achieve performance comparable to networks trained fully and then pruned, suggesting that the gradient signal at initialization encodes structurally meaningful information about each parameter's role in the network.

Zero-Cost Neural Architecture Search (NAS). Works like ZiCo (Li et al., 2023b), TE-NAS (Chen et al., 2021), and the g-SNR proxy (Xiang et al., 2023) show that architectures can be ranked by their expected performance using only gradient statistics computed on a single batch of data at initialization. The paper specifically highlights that "gradient score rankings—such as the gradient sum—correlate more strongly with architectural structures than with data batches or initialization parameters" (Section 2) and that "gradient sparsity, measured by mean and variance, is closely related to convergence rates and generalization ability" (citing Li et al., 2023b; Xiang et al., 2023).

The key insight connecting these observations to optimizer design is this: if gradient statistics at initialization already reveal which parameter groups will have noisy gradients, sparse gradients, or large gradient magnitudes, then we can set per-block learning rate scales based on that first-batch information and never update them again. This would capture the benefit of Adam's per-parameter adaptation—addressing the Hessian heterogeneity that makes SGD fail on Transformers—without paying the ongoing memory and compute cost of maintaining and updating vtv_t at every step.

The paper formalizes this using the gradient Signal-to-Noise Ratio (g-SNR) metric introduced by Xiang et al. (2023). The g-SNR for a parameter block ii is defined as:

Gsnr(i)=Gnorm(i)Gvar(i)+ϵG^{(i)}_{\text{snr}} = \frac{G^{(i)}_{\text{norm}}}{\sqrt{G^{(i)}_{\text{var}} + \epsilon}}

where Gnorm(i)G^{(i)}_{\text{norm}} is the ℓ2\ell_2-norm of the gradients in block ii (measuring update magnitude) and Gvar(i)G^{(i)}_{\text{var}} is the variance of those gradients (measuring noise/sparsity). A high g-SNR means a block has strong, consistent gradient signals; a low g-SNR means the block's gradients are noisy or sparse relative to their magnitude. By scaling the learning rate of each block proportional to its g-SNR, the method naturally slows down learning for blocks with noisy gradients (preventing erratic updates) and speeds up learning for blocks with clean, strong signals.

How This Paper Positions Itself

The paper's contribution is best understood as replacing a dynamic, memory-intensive mechanism (per-parameter adaptive learning rates via vtv_t) with a static, memory-free mechanism (per-block learning rate scaling via g-SNR computed once) that achieves the same functional goal: balancing learning progress across heterogeneous parameter groups in Transformer architectures.

The paper explicitly frames this as challenging the necessity of adaptive gradient methods (Section 1, Section 4):

"In this work, we challenge the necessity of adaptive gradient methods for model training and propose a memory- and computation-efficient alternative."

The claim is not merely that SGD-SaI is more memory-efficient than Adam (it is, by design), but that the adaptive, time-varying, per-parameter learning rate mechanism itself is unnecessary—that the function it serves (addressing training imbalance) can be accomplished with a fixed, initialization-time scaling derived from g-SNR. This is a stronger claim than prior work makes, and it carries an important practical implication: if SGD-SaI works, there is no longer a fundamental reason to accept the memory and compute overhead of Adam-like methods.

The paper also positions SGD-SaI as solving what it calls a "long-standing challenge" (Section 1, Abstract, Section 4): getting SGD to successfully train Transformer architectures. Prior attempts to train Transformers with SGD have consistently underperformed AdamW, creating a widespread belief in the community that adaptive methods are architecturally necessary for Transformers. The paper's ViT-S/16 results (Table 2, 72.92% vs. 73.04% for AdamW) and GPT-2 pretraining results (Figure 8, lower training loss than AdamW and Adam-mini) provide evidence that this belief is an artifact of insufficient learning rate scaling rather than a fundamental limitation of SGD.

In summary, the paper enters a landscape where:

  • Adam is indispensable but memory-hungry;
  • Prior memory-reduction approaches preserve Adam's adaptive mechanism but introduce complexity, slowdown, or performance degradation;
  • Evidence from pruning and NAS literature suggests initialization-time gradient statistics are structurally informative and temporally stable;
  • No prior work has attempted to replace Adam's dynamic adaptation entirely with static, initialization-time scaling.

The core contribution is demonstrating that this replacement is not only possible but produces models that match or exceed AdamW's performance while matching SGDM's memory and (near) computational footprint.

3. Technical Approach

3.1 Reader Orientation

This paper proposes a stochastic gradient descent optimizer variant—not a new model architecture or training pipeline—that computes per-parameter-group learning rate scaling factors from gradient statistics measured on the very first training batch and then freezes those factors permanently, eliminating the runtime memory and computation of Adam's second-order momentum. The system solves the problem that SGD underperforms Adam on Transformers because it cannot adapt learning rates to parameter groups with vastly different gradient magnitudes and noise levels, while simultaneously solving the problem that Adam's per-parameter adaptation costs twice the memory of SGD; the solution "shape" is a lightweight initialization-time computation of gradient signal-to-noise ratios (g-SNR) per parameter block that produces fixed scaling constants, applied multiplicatively to the global learning rate at every update step thereafter, giving SGD the heterogeneous learning rate control that makes Adam work while retaining SGD's minimal state footprint.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five logical components:

  1. Parameter Group Partitioner — divides all trainable parameters into blocks (using PyTorch's default partitioning: attention QKV as one group, attention projection as another, MLP layers as separate groups, normalization parameters, embeddings, etc.). Each block $i$ has its own gradient tensor $g^{(i)}$ and receives its own learning rate scaling factor $\tilde{G}^{(i)}_{\text{snr}}$.

  2. g-SNR Calculator (runs once at $t = 1$) — on the first training batch, computes the gradient $\ell_2$-norm and gradient variance for each block, forms their ratio (the g-SNR), normalizes across all blocks to $[0, 1]$, and stores the resulting scaling factors permanently.

  3. Momentum Accumulator — identical to standard SGDM: maintains a first-order momentum buffer $m^{(i)}_t$ per parameter block as an exponentially decaying average of past gradients, with momentum coefficient $\mu$.

  4. Decoupled Weight Decay Applier — applies $\ell_2$ regularization directly to parameters (not through the gradient), following Loshchilov & Hutter (2019), so that gradient statistics for g-SNR are not distorted by the weight decay term.

  5. Scaled Parameter Updater — at every training step, updates each block's parameters using the fixed scaling factor $\tilde{G}^{(i)}_{\text{snr}}$ multiplied by the global learning rate $\eta$, producing an effective per-block learning rate $\eta^{(i)} = \eta \cdot \tilde{G}^{(i)}_{\text{snr}}$, and then applies the standard SGDM update with this scaled learning rate.

Information flow: First batch of data → forward pass → backward pass → gradient tensors per block → compute g-SNR per block → normalize to $[0,1]$ → store scaling factors. All subsequent batches: forward pass → backward pass → accumulate momentum → apply decoupled weight decay → scale global LR by stored per-block factor → update parameters. The g-SNR computation happens exactly once and then the scaling factors are frozen.

3.3 Roadmap for the Deep Dive

  • First, the formal problem statement: what the paper aims to replace ($v_t$) and what properties a replacement must satisfy (Equation 7). This establishes why the g-SNR is a candidate at all.

  • Second, the g-SNR metric itself: its definition, why it captures the balance between gradient magnitude and gradient variance, and how it gets normalized to produce per-block scaling factors. This is the mathematical core of the paper.

  • Third, the empirical stability analysis of g-SNR over time (Figures 3 and 4), which is the evidence that justifies freezing the scaling factors after initialization rather than recomputing them—without this, the method would be no better than dynamic adaptation.

  • Fourth, the formal Taylor-expansion argument for why g-SNR stays approximately constant during training (Section 4.2), which provides theoretical backing for the empirical observation.

  • Fifth, the complete SGD-SaI algorithm (Algorithm 1), walking through each line and explaining the design choices: momentum handling, g-SNR computation on the first step, weight decay decoupling, and the parameter update form.

  • Sixth, design choices and practical considerations: why PyTorch default partitioning is sufficient (unlike Adam-mini), why decoupled weight decay matters for correct g-SNR computation, and why the method reduces to standard SGDM after the first step computationally.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an optimizer design paper whose core idea is that the function served by Adam's second-order momentum—per-parameter learning rate adaptation to address gradient heterogeneity across parameter blocks—can be replaced by a fixed, initialization-time scaling factor derived from the gradient signal-to-noise ratio (g-SNR) of each block, because g-SNR values are approximately constant throughout training for any given parameter block under gradient descent dynamics.

The Formal Objective: What Must Replace $v_t$?

The paper casts the problem as finding a function $\mathcal{F}$ that produces a per-block learning rate gain $\alpha^{(i)}_t$ from the gradient history of block $i$, but without requiring the storage or update of the full second-order momentum tensor $v_t$. The target form is stated in Equation 7 (Section 3):

αt(i)=F(g1(i),g2(i),…,gt(i))\alpha^{(i)}_t = \mathcal{F}(g^{(i)}_1, g^{(i)}_2, \ldots, g^{(i)}_t)

where $g^{(i)}_t \in \mathbb{R}^{d_i}$ is the gradient of the loss with respect to parameter block $\theta^{(i)}$ at step $t$, and $d_i$ is the number of parameters in block $i$.

What it computes: the function maps the entire history of gradient tensors for block $i$ up to step $t$ to a scalar gain factor $\alpha^{(i)}_t > 0$ that multiplies the global learning rate to yield an effective per-block learning rate $\eta^{(i)}_t = \eta \cdot \alpha^{(i)}_t$.

Why this form: the paper's central claim is that $\mathcal{F}$ does not need to depend on step $t$ in a dynamic way—it can be evaluated once at $t = 1$ and produce a constant $\alpha^{(i)}$. In Adam, the equivalent function is $\alpha^{(i)}_t = 1 / (\sqrt{\hat{v}^{(i)}_t} + \epsilon)$, where $\hat{v}^{(i)}_t$ is the bias-corrected exponential moving average of squared gradients (Equation 6). This depends on the entire history of $g^{(i)}$ in an online fashion and must be recomputed at every step, costing $\mathcal{O}(d)$ memory and $\mathcal{O}(d)$ computation per step. The paper's goal is to find a $\mathcal{F}$ that achieves the same functional benefit—balancing learning rates across blocks with different gradient characteristics—with $\mathcal{O}(B)$ memory (where $B \ll d$ is the number of blocks) and $\mathcal{O}(d)$ computation only at step 1.

The g-SNR Metric: Definition, Motivation, and Normalization

The gradient Signal-to-Noise Ratio (g-SNR), introduced by Xiang et al. (2023) for zero-cost neural architecture search, is the paper's proposed replacement for $v_t$. For a single parameter block $i$ with $d_i$ parameters, the block receives a gradient tensor $g^{(i)} \in \mathbb{R}^{d_i}$ from backpropagation. The paper defines two aggregate statistics of this tensor (Section 4.1):

Gnorm(i)=∑j=1di(gj(i))2G^{(i)}_{\text{norm}} = \sqrt{\sum_{j=1}^{d_i} \left(g^{(i)}_j\right)^2}

Gvar(i)=1di∑j=1di(gj(i)−gˉ(i))2G^{(i)}_{\text{var}} = \frac{1}{d_i} \sum_{j=1}^{d_i} \left(g^{(i)}_j - \bar{g}^{(i)}\right)^2

where $\bar{g}^{(i)} = \frac{1}{d_i} \sum_{j=1}^{d_i} g^{(i)}_j$ is the mean gradient value in the block, $g^{(i)}_j$ is the scalar gradient for the $j$-th parameter in block $i$, and $d_i$ is the number of parameters in block $i$.

What these compute: $G^{(i)}_{\text{norm}}$ is the $\ell_2$-norm (Euclidean length) of the gradient vector for the block—it measures the overall update magnitude that a vanilla SGD step would apply to this block. $G^{(i)}_{\text{var}}$ is the sample variance of the individual gradient values within the block—it measures how dispersed the per-parameter gradients are around their mean, capturing gradient heterogeneity (some parameters receiving large updates while others receive small ones) within a single block.

The g-SNR for the block is then:

Gsnr(i)=Gnorm(i)Gvar(i)+ϵG^{(i)}_{\text{snr}} = \frac{G^{(i)}_{\text{norm}}}{\sqrt{G^{(i)}_{\text{var}} + \epsilon}}

where $\epsilon > 0$ (a small constant, typically $10^{-8}$) prevents division by zero when variance is extremely small.

What it computes: the ratio of total gradient magnitude (signal strength) to the square root of gradient variance (noise amplitude). The square root makes the denominator dimensionally compatible with the numerator (both have units of gradient magnitude). A high g-SNR means the block's gradients are coherent: the total gradient norm is large relative to the internal variance, meaning most parameters in the block receive gradients of similar sign and magnitude, producing a consistent update direction. A low g-SNR means the block's gradients are noisy: individual parameters have widely varying gradients that partially cancel, meaning the net update is small relative to the internal disagreement.

Why this form: the ratio directly captures the two sources of learning imbalance that Adam addresses. Blocks with large gradient norms but low variance (high g-SNR) can tolerate large learning rates because their updates are consistent. Blocks with small norms or high variance (low g-SNR) need smaller learning rates because either the update signal is weak or the gradient noise would cause erratic updates. Adam achieves this implicitly through $v_t$: parameters with consistently large gradients accumulate large $v_t$ values, shrinking their effective learning rate. The g-SNR provides the same information in a single scalar per block, computable from one gradient sample rather than requiring a running average of squared gradients.

Normalization step. Because absolute g-SNR values have no natural scale, the paper normalizes across all $B$ blocks to the range $[0, 1]$:

G~snr(i)=Gsnr(i)max⁡k∈{1,…,B}Gsnr(k)\tilde{G}^{(i)}_{\text{snr}} = \frac{G^{(i)}_{\text{snr}}}{\max_{k \in \{1, \ldots, B\}} G^{(k)}_{\text{snr}}}

What it computes: each block's g-SNR divided by the maximum g-SNR among all blocks in the network. The block with the highest g-SNR gets a scaling factor of 1.0; all other blocks get proportionally lower factors.

Why this normalization: it ensures that (1) no block receives an effective learning rate greater than $\eta$ (the global learning rate), bounding the per-block learning rate from above, and (2) the relative scaling between blocks is preserved—a block with half the g-SNR of the maximum gets half the effective learning rate. The global learning rate $\eta$ thus sets the maximum learning rate any block can receive, and the normalized g-SNR acts as a per-block attenuation factor.

Empirical Justification: g-SNR Stability Over Time

The paper's key empirical claim is that g-SNR values, once computed at initialization, remain approximately constant throughout training for each parameter block. This is what distinguishes SGD-SaI from dynamic adaptive methods: the scaling factors can be frozen because the gradient statistics that dictate them do not change substantially over the course of optimization.

Experimental setup. The paper trains a ViT/S-16 on ImageNet-1K and traces g-SNR values over time for different parameter blocks (Figure 4, Section 4.2). Specifically, it examines transformer blocks from three depth positions—shallow (block 0), middle (block 5), and deep (block 11)—as well as non-transformer parameters like positional embeddings.

Observation 1: Temporal stability within blocks. Figure 4 shows that within each traced block, the g-SNR value remains "relatively constant over time" (Section 4.2). The curves are approximately flat across training steps. This is the central empirical justification for the SaI approach: if g-SNR were to drift significantly, frozen scaling factors would become inappropriate for later training stages, and dynamic recomputation would be necessary.

Observation 2: Cross-block variation. Figure 3 (Section 4.2) shows the complementary property: g-SNR values differ substantially between different parameter block types. The paper plots frequency histograms of g-SNR values separately for weights and biases across block types (norm1, norm2, attn.qkv, attn.proj, mlp.fc1, mlp.fc2, cls_token, pos_embed, patch_embed.proj, head, fc_norm). The distributions are clearly separated: QKV projections have different typical g-SNR values than MLP layers, which differ from normalization parameters. This cross-block variation is exactly what necessitates per-block scaling—if all blocks had similar g-SNR, a single global learning rate would suffice.

Observation 3: Structural similarity across transformer layers. The paper notes an important sub-pattern within Figure 3: "for most weights, the parameter blocks that share the same structure across different transformer layers (blocks) tend to have similar g-SNR values" (Section 4.2). This means all attn.qkv blocks across different transformer layers have roughly similar g-SNR, all mlp.fc1 blocks have similar g-SNR, etc. This structural consistency suggests that the g-SNR is primarily determined by the architectural role of a parameter group (its connectivity pattern, input/output dimensions, position in the network) rather than by random initialization or data sampling noise.

Observation 4: Bias parameters have uniformly low g-SNR. The paper explicitly notes that "the g-SNR values for the bias parameters are consistently low magnitude" (Figure 3 caption). This aligns with the fact that biases typically receive smaller, less structured gradients than weight matrices, and they are known to benefit from higher effective learning rates—which is exactly what low-g-SNR scaling provides (the normalization ensures low g-SNR blocks get small scaling factors, but the absolute learning rate can be tuned via $\eta$).

Observation 5: An architectural exception reveals topological sensitivity. The paper points out that "norm1 weights from blocks.0... connect to the input from embedded patches, whereas all other norm1 weights connect to the output of the previous block" (Section 4.2). The g-SNR values for this first norm1 layer differ from subsequent norm1 layers, suggesting that "g-SNR values can effectively identify distinct characteristics among different parameter groups and the network's topological impacts" (Section 4.2). This is evidence that g-SNR captures genuine architectural wiring differences, not just type-level categories.

What these observations collectively justify. The combination—temporal stability within blocks plus systematic variation between blocks—means that computing g-SNR once at initialization captures structural information about gradient distributions that remains valid throughout training. The paper concludes: "by leveraging constant g-SNR values, this approach effectively assigns a pre-conditioned learning rate scale to each partition" (Section 4.2).

Theoretical Justification: First-Order Taylor Argument for g-SNR Constancy

The paper provides a theoretical argument (Section 4.2) for why g-SNR should remain approximately constant, complementing the empirical observations. The reasoning proceeds in three steps:

Step 1: Small parameter updates. "Assuming $\eta$ is sufficiently small to stabilize the training process, we have $\theta^{(i)}_{t+1} \approx \theta^{(i)}_t$. Thus, the change in parameters per iteration is small."

Step 2: First-order Taylor expansion of the gradient. The gradient at step $t+1$ can be approximated around $\theta^{(i)}_t$:

gt+1(i)≈gt(i)+Jt(i)Δθt(i)g^{(i)}_{t+1} \approx g^{(i)}_t + J^{(i)}_t \Delta \theta^{(i)}_t

where $J^{(i)}_t$ is the Jacobian (first-order sensitivity matrix) of the gradient $g^{(i)}$ with respect to $\theta^{(i)}$ at step $t$, and $\Delta \theta^{(i)}_t = \theta^{(i)}_{t+1} - \theta^{(i)}_t$ is the parameter update.

What this expresses: the gradient at the next step equals the current gradient plus a correction term that is the product of the local curvature (Jacobian) and the parameter change. This is simply linearizing the gradient function around the current parameter value.

Step 3: Bounding the gradient change. Since $\| \Delta \theta^{(i)}_t\|$ is small (because $\eta$ is small), the correction term $J^{(i)}_t \Delta \theta^{(i)}_t$ is also small, so:

gj(t+1)(i)≈gj(t)(i),∀jg^{(i)}_{j(t+1)} \approx g^{(i)}_{j(t)}, \quad \forall j

Step 4: Propagation to aggregate statistics. Because each individual gradient component changes only slightly, the aggregate statistics computed from them also change only slightly:

gˉt+1(i)≈gˉt(i)\bar{g}^{(i)}_{t+1} \approx \bar{g}^{(i)}_t Gvar(t+1)(i)≈Gvar(t)(i)G^{(i)}_{\text{var}(t+1)} \approx G^{(i)}_{\text{var}(t)} Gnorm(t+1)(i)≈Gnorm(t)(i)G^{(i)}_{\text{norm}(t+1)} \approx G^{(i)}_{\text{norm}(t)}

Step 5: Therefore g-SNR is constant. Since both numerator and denominator remain approximately unchanged:

Gsnr(t+1)(i)≈Gsnr(t)(i)G^{(i)}_{\text{snr}(t+1)} \approx G^{(i)}_{\text{snr}(t)}

Why this argument matters (and its limitations). The argument provides a mechanistic explanation for the empirical observation of temporal stability. It says: gradient distributions are locally Lipschitz (smooth) functions of the parameters, and with small learning rates the parameters don't move far enough for the gradient distribution—and hence the g-SNR—to change substantially. The paper's accompanying intuition is that "the g-SNR measures a dimensionless ratio that characterizes this shape. Minor parameter shifts do not significantly affect this ratio; hence, it remains nearly constant" (Section 4.2).

However, this is a local argument: it holds as long as the parameters stay near their initialization values. If training substantially moves parameters (which it must for successful learning), the Jacobian approximation can break down. The paper implicitly relies on the empirical evidence (Figure 4) to show that in practice, under standard training configurations, this breakdown either doesn't occur or doesn't significantly alter the relative ordering of g-SNR values across blocks, which is what matters for per-block scaling.

The Complete SGD-SaI Algorithm

Algorithm 1 (Section 4.3) specifies SGD-SaI in full detail. The algorithm processes one parameter block $\theta^{(i)}$ at a time, repeating the following procedure at each training step $t = 1, \ldots, T$:

Step 1: Compute gradient. gt(i)←∇θ(i)L(θt−1)g^{(i)}_t \leftarrow \nabla_{\theta^{(i)}} \mathcal{L}(\theta_{t-1}) Optionally negate if maximizing rather than minimizing the objective (lines 3–5).

Step 2: Apply momentum (first-order). If $t > 1$ (all steps after the first): mt(i)←μ⋅mt−1(i)+(1−μ)⋅gt(i)m^{(i)}_t \leftarrow \mu \cdot m^{(i)}_{t-1} + (1 - \mu) \cdot g^{(i)}_t where $\mu \in [0, 1)$ is the momentum coefficient (default 0.9).

If $t = 1$ (the first step only), initialize momentum from the gradient ($m^{(i)}_1 \leftarrow g^{(i)}_1$) AND compute the g-SNR scaling factor (lines 11–14):

Gsnr(i)←Gnorm(i)Gvar(i)+ϵG^{(i)}_{\text{snr}} \leftarrow \frac{G^{(i)}_{\text{norm}}}{\sqrt{G^{(i)}_{\text{var}} + \epsilon}} G~snr(i)←Gsnr(i)max⁡kGsnr(k)\tilde{G}^{(i)}_{\text{snr}} \leftarrow \frac{G^{(i)}_{\text{snr}}}{\max_{k} G^{(k)}_{\text{snr}}}

These scaling factors $\tilde{G}^{(i)}_{\text{snr}}$ are stored and reused at every subsequent step without recomputation.

What the first-step special case does: it intercepts the very first gradient computation to calculate and store the per-block scaling factors. This is the only place where g-SNR computation occurs—every subsequent step skips this block and falls through to the standard momentum update.

Step 3: Apply decoupled weight decay. θt(i)←θt−1(i)−ηλθt−1(i)\theta^{(i)}_t \leftarrow \theta^{(i)}_{t-1} - \eta \lambda \theta^{(i)}_{t-1} where $\lambda \geq 0$ is the weight decay coefficient.

What this does: it shrinks the parameters directly toward zero by a factor proportional to their current value and the product of learning rate and weight decay. This is decoupled from the gradient, meaning the gradient $g^{(i)}_t$ used in momentum accumulation does not include the weight decay term. The paper explicitly justifies this: "Decoupled weight decay applies regularization directly to the parameters rather than incorporating it into the gradient computation. This approach is equivalent to regularization in SGD and allows us to accurately compute the gradient statistics needed for the g-SNR without the weight decay term distorting the gradient values" (Section 4.3).

Why decoupling matters for g-SNR: if weight decay were incorporated into the gradient (as in standard Adam, Algorithm 3 line 10), the gradient vector used to compute g-SNR would include a term proportional to the parameter values themselves, mixing the loss landscape's gradient signal with an artificial shrinkage signal. This would distort the g-SNR, making it reflect parameter magnitudes rather than pure gradient statistics. Decoupled weight decay ensures the g-SNR measures only the training loss gradient, which is the quantity that dictates learning dynamics.

Step 4: Apply scaled parameter update. θt(i)←θt(i)−η⋅G~snr(i)⋅mt(i)\theta^{(i)}_t \leftarrow \theta^{(i)}_t - \eta \cdot \tilde{G}^{(i)}_{\text{snr}} \cdot m^{(i)}_t This line applies the momentum-smoothed gradient as the update direction, scaled by (a) the global learning rate $\eta$, (b) the block-specific frozen g-SNR factor $\tilde{G}^{(i)}_{\text{snr}}$, producing the effective per-block update magnitude $\eta \cdot \tilde{G}^{(i)}_{\text{snr}} \cdot m^{(i)}_t$.

What the combined update does: the effective per-block learning rate is $\eta^{(i)} = \eta \cdot \tilde{G}^{(i)}_{\text{snr}}$. Blocks with high g-SNR (clean, strong gradient signal) receive values of $\tilde{G}^{(i)}_{\text{snr}}$ near 1.0 and thus learn at close to the full global rate. Blocks with low g-SNR (noisy or weak gradient signal) receive proportionally smaller effective learning rates, preventing erratic updates that would destabilize training. The momentum buffer $m^{(i)}_t$ provides temporal smoothing, filtering out batch-to-batch noise within each block, exactly as in standard SGDM.

Computational cost analysis. After step 1 (the first batch), SGD-SaI is algorithmically identical to SGDM with per-block learning rate scales. The only additional cost compared to SGDM is:

  • At step 1: computing $G^{(i)}_{\text{norm}}$ (a single $\ell_2$-norm per block, $\mathcal{O}(d)$ total), computing $G^{(i)}_{\text{var}}$ (a single variance per block, $\mathcal{O}(d)$ total), performing one division and square root per block ($\mathcal{O}(B)$), and one max over $B$ scalars.
  • At all subsequent steps: one additional scalar multiplication per block (multiplying the global LR by the stored $\tilde{G}^{(i)}_{\text{snr}}$), which is $\mathcal{O}(B)$ per step and negligible compared to the $\mathcal{O}(d)$ cost of gradient computation and momentum updates.

Memory cost analysis. SGD-SaI stores exactly the same state as SGDM: one momentum buffer $m^{(i)}$ per parameter, of identical shape to the parameters. The g-SNR scaling factors require $B$ additional scalars (one per block), which is negligible—for a typical Transformer with hundreds of blocks, this is kilobytes of storage. Compare to AdamW, which stores both $m^{(i)}$ and $v^{(i)}$ (each the size of the parameters), doubling the optimizer state memory. Compare to Adam-mini, which stores $m^{(i)}$ (full size) and a reduced $v^{(i)}$ (block-level means of squared gradients, $\mathcal{O}(B)$ rather than $\mathcal{O}(d)$), achieving similar memory savings but with additional computation per step to maintain the block-level statistics.

Design Choices and Their Justifications

Why PyTorch Default Partitioning is sufficient. The paper explicitly contrasts with Adam-mini's complex partitioning requirements (Section 4.2, "This observation does not hold true in our case"). Adam-mini requires fine-grained partitioning (splitting QKV into separate Q, K, V by attention heads) to align with Hessian sub-block structure (Zhang et al., 2024a). SGD-SaI achieves strong results with PyTorch's default coarse partitioning (all of attn.qkv as one group). The paper attributes this to the g-SNR being computed from first-batch gradients rather than accumulated statistics: the initial gradient already reflects the architectural wiring of the network, and the g-SNR ratio naturally distinguishes blocks with different gradient characteristics even if they are coarsely grouped.

Empirically, Figure 3 shows that weight blocks with the same structural role (e.g., all QKV projections across transformer layers) have similar g-SNR values, so grouping them together (as PyTorch does) produces scaling factors that are appropriate for each member of the group. The one exception the paper notes—norm1 weights in block 0 receiving different g-SNR than norm1 weights in later blocks due to their connection to patch embeddings—suggests that even coarse partitioning captures these topological differences, since block 0 is a separate partition from block 1, etc.

Why the scaling factor is frozen rather than periodically recomputed. The paper's core empirical finding (Figure 4) is that g-SNR values are temporally stable, meaning recomputation would yield nearly identical values at any training step. Periodically recomputing would reintroduce overhead (additional forward/backward passes or gradient statistics computation) without changing the effective learning rates appreciably. The first-order Taylor argument (Section 4.2) provides theoretical backing: with small learning rates, gradient statistics change slowly, making the initialization-time value a valid approximation throughout training.

The paper does not experimentally compare frozen versus periodically recomputed scaling factors—this is an implicit bet on the temporal stability observation. If g-SNR values were to shift significantly during training (e.g., during loss landscape phase transitions), frozen factors could become suboptimal. The paper's experiments on a range of tasks (ViT, GPT-2, LoRA fine-tuning, diffusion models) showing consistent performance provide indirect evidence that this is not a practical concern under standard training configurations.

Why the global learning rate $\eta$ is not eliminated. Normalizing g-SNR to $[0, 1]$ means the block with the highest g-SNR gets a scaling factor of 1.0 and thus an effective learning rate equal to $\eta$. The global learning rate therefore controls the maximum learning rate any block receives. This preserves the need for learning rate tuning (the paper performs grid searches over $\eta$ in all experiments) but collapses the tuning problem from "find the right per-block learning rates" (infeasible) to "find the right global learning rate" (standard practice), with the g-SNR providing the relative scaling between blocks automatically.

Why momentum is applied before scaling rather than after. In Algorithm 1, momentum is computed on the raw gradient (line 8), and the g-SNR scaling is applied only at the final parameter update (line 19). This means the momentum buffer accumulates gradients at their natural scale, and the per-block learning rate modulation acts on the smoothed update direction. The alternative—scaling the gradient before momentum accumulation—would mean the momentum buffer stores artificially rescaled gradients, potentially distorting the momentum's temporal smoothing function. The paper does not discuss this choice explicitly, but it follows naturally from the interpretation of g-SNR as a learning rate modulator rather than a gradient preconditioner.

Why weight decay is applied before momentum, not after. The parameter update with decoupled weight decay happens at line 17 ($\theta^{(i)}_t \leftarrow \theta^{(i)}_{t-1} - \eta \lambda \theta^{(i)}_{t-1}$), followed by the momentum-based update at line 19. This matches the AdamW formulation (Loshchilov & Hutter, 2019, Algorithm 4, line 18+20) where weight decay shrinks parameters independently of the adaptive gradient update. The ordering ensures that weight decay regularization and gradient-based learning are orthogonal operations, and the weight decay strength is controlled cleanly by $\lambda$ without interacting with the momentum buffer or g-SNR scaling.

Why the $\epsilon$ in the g-SNR denominator matters. The small constant $\epsilon$ (typically $10^{-8}$) in $\sqrt{G^{(i)}_{\text{var}} + \epsilon}$ prevents division by zero when a block's gradient variance is exactly zero (all parameters in the block have identical gradient values). This can occur at initialization for certain parameter types (e.g., biases that haven't received any gradient signal yet). Without $\epsilon$, the g-SNR would be undefined or infinite, breaking the normalization step. With $\epsilon$, blocks with extremely low variance get a finite, typically large g-SNR (since the denominator is dominated by $\epsilon$), which after normalization means they receive the maximum scaling factor of 1.0—appropriate behavior since zero-variance gradients are perfectly coherent.

Summary: How All Components Fit Together

At initialization ($t = 1$), the system:

  1. Performs a forward and backward pass on the first batch to compute gradients for all parameters.
  2. Groups gradients by PyTorch default partition (hundreds of blocks, corresponding to layers and parameter types).
  3. For each block, computes $G_{\text{norm}}$ and $G_{\text{var}}$ from the gradient tensor.
  4. Computes $G_{\text{snr}} = G_{\text{norm}} / \sqrt{G_{\text{var}} + \epsilon}$ for each block.
  5. Normalizes all g-SNR values by the maximum across blocks to get $\tilde{G}_{\text{snr}} \in [0, 1]$.
  6. Stores these $\tilde{G}_{\text{snr}}$ values permanently.

For all steps $t = 2, \ldots, T$ and also step $t = 1$ after g-SNR computation, the system:

  1. Computes gradients normally via backpropagation.
  2. Updates the first-order momentum buffer $m_t$ using exponential moving average with coefficient $\mu$.
  3. Applies decoupled weight decay to shrink parameters toward zero.
  4. Updates parameters with the momentum direction scaled by $\eta \cdot \tilde{G}_{\text{snr}}$.

The result is that each parameter block receives a learning rate proportional to its initial gradient signal-to-noise ratio, addressing the Hessian heterogeneity of Transformers that makes vanilla SGD ineffective, while using exactly the same optimizer state memory as SGDM ($m_t$ only) and performing only $\mathcal{O}(B)$ additional computation per step beyond standard SGDM.

4. Key Insights and Innovations

Innovation 1: The Adaptive Gradient Mechanism Is Replaceable by a Static, Initialization-Time Proxy—Not Just Compressible

The paper makes a fundamentally different intellectual move than all prior work on optimizer memory reduction. Every previous approach—Adafactor (Shazeer & Stern, 2018), 8-bit Adam (Dettmers et al., 2021), Adam-mini (Zhang et al., 2024b), sign-based methods (Bernstein et al., 2018)—accepted the premise that per-parameter adaptive learning rates via second-order momentum are necessary, and sought to approximate or compress them at lower cost. The field's implicit assumption was that the dynamic, time-varying learning rate adjustment that Adam provides is essential for training Transformers; the only question was how to achieve it with fewer bits, more sharing, or coarser granularity.

SGD-SaI challenges this premise directly. The claim is not that g-SNR is a better approximation of $v_t$—it does not attempt to approximate $v_t$ at all. The claim is that the function $v_t$ performs—rescaling learning rates to account for heterogeneous gradient distributions across parameter blocks—can be accomplished by a completely different mechanism that requires no dynamic state: measure the relevant gradient statistics once at initialization, compute fixed scaling factors, and never update them again. This is a conceptual substitution, not a compression.

The significance of this move is that it reframes the problem. The community had been asking: "How can we make Adam cheaper?" The paper asks instead: "What does Adam actually do for Transformer training, and can we achieve that function without any adaptive state at all?" The answer—balancing learning rates across blocks with different gradient signal-to-noise characteristics—turns out to be addressable with static pre-conditioning. This shifts the optimization design space from "adaptive methods with various memory/compute tradeoffs" to "initialization-time statistics as a sufficient basis for learning rate heterogeneity." It is a fundamental conceptual reframing, not an incremental improvement to Adam.

The evidence that makes this claim credible is the dual observation in Figures 3 and 4: g-SNR varies substantially between parameter blocks (Figure 3, showing distinct distributions for QKV vs. MLP vs. normalization layers) but remains nearly constant within each block over training time (Figure 4, showing flat trajectories). If either property failed—if all blocks had similar g-SNR, or if g-SNR drifted significantly—the static scaling approach would collapse. The fact that both properties hold empirically, and that the paper provides a first-order Taylor argument for why g-SNR should be temporally stable under small learning rates (Section 4.2), elevates this from an empirical trick to a principled design choice.

A subtle but important intellectual consequence: if the scaling factors are frozen, then after the first training step, SGD-SaI is algorithmically identical to SGDM with per-block constant learning rate multipliers. This means the paper has essentially shown that SGDM with appropriately chosen per-block learning rates can match AdamW on Transformers—and that those appropriate learning rates can be derived automatically from first-batch gradient statistics rather than through expensive per-block hyperparameter tuning. This resolves the long-standing puzzle of why SGD underperforms Adam on Transformers: the failure was not SGD's update rule but its uniform learning rate, and the fix requires only block-level differentiation, not parameter-level adaptation or temporal dynamics.

Innovation 2: The Gradient Signal-to-Noise Ratio as a Diagnostic for Learning Rate Allocation

While the g-SNR metric itself was introduced by Xiang et al. (2023) for zero-cost neural architecture search, the paper's innovation is repurposing it as a learning rate scaling mechanism and, crucially, providing the empirical and theoretical justification for why it is the right statistic for this job.

Prior work on learning rate scaling for heterogeneous architectures had approached the problem from two directions. One line—exemplified by Adam and its variants—uses accumulated gradient magnitude history ($v_t$) as a proxy for how much each parameter should learn, with the intuition that parameters receiving large gradients should be slowed down to prevent overshooting. Another line—exemplified by layer-wise adaptive rate scaling (LARS) and similar methods—uses the ratio of parameter norm to gradient norm as a scaling factor, motivated by the observation that different layers can have wildly different weight-to-gradient ratios.

The g-SNR approach is distinct from both. It captures not just gradient magnitude (which Adam's $v_t$ does) but the coherence of the gradient signal within a parameter block: the ratio of total gradient norm to within-block gradient variance. This is a more nuanced statistic than either raw magnitude or norm ratios. A block with large but highly variable gradients (some parameters receiving strong positive updates, others strong negative updates) will have a low g-SNR despite high $G_{\text{norm}}$; Adam would give such a block a moderate effective learning rate (large $v_t$ but also large gradient magnitude), potentially allowing noisy updates through. SGD-SaI would give it a low scaling factor, directly suppressing the noise.

The paper's diagnostic contribution is demonstrating that this g-SNR statistic naturally separates parameter blocks by their architectural role (Figure 3): attention QKV projections, MLP layers, normalization parameters, and embeddings all occupy different regions of g-SNR space. This is not a design choice—it is an emergent property of how gradients flow through Transformer architectures at initialization. The paper observes that blocks sharing the same structural role across different transformer layers tend to have similar g-SNR values (Section 4.2), and that topological differences (e.g., the first norm layer connecting to patch embeddings rather than a previous transformer block) produce detectable g-SNR differences. This means g-SNR is effectively an architecture-aware difficulty metric: it automatically identifies which parameter groups have clean gradient signals (and can tolerate aggressive learning rates) versus noisy or weak signals (and need conservative rates) based purely on the network's wiring and initialization, without any manual specification of layer types or training heuristics.

The intellectual significance extends beyond this paper's optimizer design. The fact that a simple ratio of gradient norm to gradient variance, computed once at initialization, encodes structurally meaningful information about learning dynamics suggests that g-SNR could serve as a general diagnostic tool for understanding and debugging training behavior. It connects conceptually to the Hessian-based analysis of Zhang et al. (2024a), which showed that Transformers have block-heterogeneous loss landscapes, but replaces expensive Hessian computation with a cheap first-order statistic. This opens the door to using g-SNR for tasks beyond optimizer design: architecture search (the original use case), predicting training difficulty, identifying problematic layers before training, or guiding mixed-precision allocation.

The evidence anchoring this innovation is Figure 3's clear separation of g-SNR distributions by block type, Figure 4's demonstration of temporal stability, and the consistent performance of SGD-SaI across ViT, GPT-2, and CNN architectures (Tables 2, 4, Figure 5, Figure 8, Figure 10)—all of which use the same g-SNR computation procedure without task-specific tuning.

Innovation 3: Difficulty-Conditioned Robustness Without Dynamic Adaptation

A striking result that the paper does not fully theorize but that emerges clearly from the experiments is SGD-SaI's robustness to hyperparameter variation—it maintains higher average accuracy and lower variance across learning rate and weight decay grid searches than AdamW or Adam-mini, even when peak performance is comparable. This is visible in Table 2 (ViT-S/16 on ImageNet-1K): SGD-SaI achieves an average top-1 accuracy of 57.55% across the full hyperparameter grid with standard deviation 18.46, compared to AdamW's 37.21% average with 35.43 standard deviation. In other words, AdamW's performance collapses on suboptimal hyperparameters while SGD-SaI degrades gracefully. The same pattern appears in Figure 5 for ResNet-18 on CIFAR-10, where SGD-SaI's distribution is tighter and higher-mean than AdamW, Adam, or vanilla SGD.

This robustness is unexpected under the standard intuition about adaptive methods. Adam was explicitly designed to reduce sensitivity to learning rate choice—its per-parameter adaptation should, in theory, make it more robust than SGD. Yet the paper's results show the opposite: the static, non-adaptive SGD-SaI is substantially more robust than the dynamic, adaptive AdamW. Why?

The paper does not provide a mechanistic explanation, but the result is consistent with a particular interpretation: Adam's dynamic adaptation can amplify poor hyperparameter choices through a feedback loop. If the global learning rate is set too high, Adam's per-parameter scaling will initially rein in the updates (via large $v_t$), but the aggressive initial steps may still push parameters into regions of the loss landscape from which recovery is difficult. Moreover, Adam's bias correction terms and the interaction between $\beta_1$, $\beta_2$, and $\eta$ create a complex, coupled system where changing one hyperparameter shifts the effective behavior of others in non-obvious ways. This coupling may explain why AdamW's performance collapses to near-random (37% average, close to chance on 1000-class ImageNet) under some hyperparameter combinations—the optimizer enters a regime where its adaptation mechanism works against rather than with the learning rate.

SGD-SaI, by contrast, has a simpler hyperparameter interaction: the global learning rate $\eta$ multiplies fixed per-block scaling factors. There is no feedback between the learning rate and the scaling mechanism. If $\eta$ is too high, training may diverge, but the divergence is monotonic and predictable—there is no adaptive mechanism to create complex failure modes. This simplicity is a feature, not a bug: it makes hyperparameter tuning more predictable and makes the optimizer's behavior under extrapolation (to new architectures, datasets, or scales) more reliable.

The practical significance of this robustness is substantial but easy to overlook in a paper focused on memory efficiency. In real-world deployment, hyperparameter tuning is expensive—grid searches over learning rate and weight decay require multiple full training runs. An optimizer that tolerates a wider range of hyperparameters without catastrophic degradation means fewer tuning runs are needed to find a good configuration, and the risk of accidentally selecting a poor hyperparameter due to limited search budget is lower. For practitioners with constrained compute, this robustness may be as valuable as the memory savings.

The evidence for this innovation is primarily Table 2 (peak vs. average accuracy for ViT), Figure 5 (distributions for CNN), and Figure 10 (distributions across eleven architectures on three datasets), all showing tighter, higher-mean accuracy distributions for SGD-SaI.

Innovation 4: Decoupling Gradient Statistics from Optimization State as a Design Principle

A more subtle conceptual contribution is the paper's implicit articulation—through its design choices—of a principle that gradient statistics used for learning rate scaling should be decoupled from the optimization state that training itself modifies. This principle is not stated explicitly but is embedded in two concrete decisions:

Decision 1: Compute g-SNR from the first batch only, not from a running average. If g-SNR were computed dynamically throughout training, it would be influenced by the optimization trajectory—the parameters would have moved, the loss landscape curvature would have changed, and the g-SNR values would reflect the current optimization state rather than the inherent architectural gradient characteristics. By computing once at initialization, before any parameter updates have occurred, the g-SNR captures the "clean" gradient distribution imposed by the architecture and initialization scheme, uncontaminated by the optimizer's own effects.

Decision 2: Use decoupled weight decay so the g-SNR measures only the training loss gradient. If weight decay were incorporated into the gradient (as in standard Adam with L2 regularization), the gradient vector used for g-SNR computation would contain a term proportional to the current parameter values ($\lambda \theta$). This would mix the loss landscape signal with a magnitude-dependent shrinkage signal, making the g-SNR partially a function of parameter scale rather than pure gradient statistics. By decoupling weight decay (applying it directly to parameters, as in AdamW), the paper ensures the gradient $g$ used for g-SNR reflects only the data-dependent loss. The paper explicitly notes this in Section 4.3: "Decoupled weight decay... allows us to accurately compute the gradient statistics needed for the g-SNR without the weight decay term distorting the gradient values."

Together, these decisions implement a separation of concerns: the g-SNR captures structural information about gradient distributions (determined by architecture and initialization), while the optimizer's momentum and update mechanism handle temporal smoothing and navigation of the loss landscape. These are fundamentally different jobs, and mixing them—as Adam does, where $v_t$ accumulates squared gradients that include the effects of weight decay, parameter updates, and changing loss landscape curvature—creates a coupled system where the learning rate adaptation reflects the history of optimization decisions rather than the underlying problem structure.

This decoupling principle has implications beyond this paper. It suggests a general approach to optimizer design: identify what information is structural (determined by architecture and data distribution, stable over time) versus what is trajectory-dependent (determined by the optimization path, changing as parameters move). Compute the structural information once at initialization; use it to configure a simple, stateless update rule that handles the trajectory-dependent part. This echoes themes from the meta-learning and optimization literature (learning to optimize, optimizer hyperparameter prediction) but implements them in the simplest possible way.

The evidence that this decoupling works is the paper's consistent performance across diverse tasks without task-specific tuning of the g-SNR computation: the same procedure (compute g-SNR once from first-batch gradients) works for ViT image classification, GPT-2 language modeling, LoRA fine-tuning, diffusion model fine-tuning, and CNN training (Sections 5.1–5.4). This universality is precisely what one would expect if the g-SNR captures genuinely structural, task-independent gradient characteristics.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper evaluates on multiple datasets across task types:

    • ImageNet-1K (Deng et al., 2009): 1.28M training images, 50K validation images across 1,000 classes, used for ViT-S/16 pretraining (Section 5.2). Standard top-1 validation accuracy is reported.
    • OpenWebText (Gokaslan & Cohen, 2019): An open-source recreation of the WebText corpus used for GPT-2 pretraining (Section 5.1). Training and validation loss are the primary metrics.
    • E2E NLG Challenge dataset (Novikova et al., 2017): Used for GPT-2 LoRA fine-tuning (Section 5.3.1). Multiple NLG metrics are reported: BLEU, NIST, MET, ROUGE-L, and CIDEr.
    • CIFAR-10, CIFAR-100 (Krizhevsky & Hinton, 2009), ImageNet16-120 (Chrabaszcz et al., 2017): Used for CNN experiments (Section 5.4) including ResNet-18 and NATS-Bench architectures. Top-1 test accuracy is reported.
    • ChilloutMix diffusion model with custom concept data: Used for diffusion model LoRA fine-tuning (Section 5.3.2). Qualitative evaluation only (visual inspection of generated images).

  • Base model(s). The paper uses five model families at different scales:

    • GPT-2 (Radford et al., 2019): GPT-2-Small (125M parameters) for full pretraining runs; GPT-2-Medium for LoRA fine-tuning; GPT-2-XL (1.5B) and Llama2-7B (Touvron et al., 2023) for memory and speed profiling only (no full training due to resource constraints, acknowledged in Section 7 Limitations).
    • Vision Transformer (Dosovitskiy et al., 2021): ViT-S/16 (22.9M parameters) for full ImageNet-1K pretraining runs; ViT-H/14 (660M parameters) for profiling only.
    • ResNet-18 (He et al., 2015a): 11M parameters, used for CIFAR-10 experiments with grid search.
    • NATS-Bench architectures (Dong et al., 2021): Eleven architectures of varying sizes from the NATS-Bench size search space, used for cross-architecture robustness evaluation on CIFAR-10, CIFAR-100, and ImageNet16-120.
    • ChilloutMix diffusion model (Gu et al., 2024): Used for LoRA fine-tuning on face generation concepts. The paper argues that PaLM 2-S* is "representative of the capabilities of many contemporary LLMs" (Section 4), but this claim appears to be from a different paper context; the actual models used span from small CNNs to medium-scale Transformers.

  • Metrics. Multiple task-specific metrics are used:

    • Top-1 validation/test accuracy for image classification (ViT, CNN experiments): fraction of images correctly classified, reported as percentage. Peak (best across hyperparameter grid) and average (mean across grid) are both reported in Table 2.
    • Training loss and validation loss for GPT-2 pretraining (Figure 8): standard cross-entropy loss on next-token prediction.
    • BLEU, NIST, MET, ROUGE-L, CIDEr for E2E NLG fine-tuning (Table 4): standard NLG evaluation metrics; higher is better for all.
    • Memory usage of optimizer state tensors (Tables 1, 3, Figure 1): measured in GB, includes only the tensors stored by the optimizer (momentum buffers, second-moment estimates, and in Prodigy's case, additional state tensors as detailed in Appendix B).
    • Optimizer step wall-clock time (Tables 1, 3, Figure 1): measured in milliseconds, refers specifically to the time taken by the optimizer.step() function call (not including forward/backward passes), profiled over multiple iterations.
    • Qualitative visual assessment for diffusion model generation (Figure 9): no quantitative metric; evaluation based on visual fidelity to the prompt and training concept.

  • Baselines. The paper compares against six optimizer baselines (Sections 5.1–5.4, with per-task availability varying):

    • SGDM (Nesterov, 1983): Standard stochastic gradient descent with momentum (momentum coefficient 0.9), storing one momentum buffer per parameter. The core non-adaptive baseline.
    • Adam (Kingma & Ba, 2014): Adaptive gradient method with first-order and second-order momentum, L2 weight decay (not decoupled).
    • AdamW (Loshchilov & Hutter, 2019): Adam with decoupled weight decay. The primary competitive baseline, considered the default optimizer for Transformer training.
    • Adam-mini (Zhang et al., 2024b): Memory-efficient Adam variant that shares adaptive learning rates within parameter blocks, requiring custom partitioning.
    • Prodigy (Mishchenko & Defazio, 2023): Parameter-free adaptive optimizer with automatic learning rate adjustment; stores approximately 4× the gradient size in state tensors (Appendix B.4).
    • Scaled-SGD and Scaled-AdamW (Zhang & Pilanci, 2024): Riemannian preconditioned variants for LoRA fine-tuning, adjusting learning rates for A and B matrices separately. Compared only in the LoRA experiments (Section 5.3.1).

  • Generation budget / compute accounting. The paper does not use generation budgets (this is not a test-time compute paper). Instead, fairness is established through:

    • Equal hyperparameter search space: All optimizers are evaluated over the same grid of learning rates and weight decays (e.g., lr ∈ {0.1, 0.01, 0.001, 0.0001}, wd ∈ {0.01, 0.001, 0.0001} for ViT; similar ranges for CNNs). This ensures no optimizer benefits from a more extensive tuning effort.
    • Fixed training recipes: All optimizers use the same number of epochs, batch size, data augmentation, and learning rate schedule (cosine annealing for ViT and CNNs, step decay for ResNet-18). For GPT-2 pretraining, the recipe follows Zhang et al. (2024b).
    • Profiling on identical hardware: Memory and speed measurements use a single NVIDIA A100-80GB or RTX 3090, with consistent batch size (1), context length (1024 for LLMs), and precision (FP32).

  • Cross-validation / statistical protocol:

    • Multiple random seeds: For ViT-S/16 peak performance (Table 2), results are averaged over three random seeds with standard deviations reported (e.g., SGD-SaI: 72.92 ± 0.07). For ResNet-18 (Appendix A.3), three seeds (42, 888, 999) are used.
    • Grid search with identical hyperparameter ranges: All optimizers are evaluated over the same grid; the best hyperparameter combination per optimizer determines peak performance. Average performance across the entire grid measures robustness.
    • Profiling over multiple iterations: Optimizer step times are measured over 20 iterations (Tables 5, 6 in Appendix C) with mean and standard deviation reported to account for measurement noise.
    • There is no k-fold cross-validation on held-out data—the paper uses fixed train/val/test splits as defined by the respective benchmarks.

Main Quantitative Results

GPT-2 Pretraining (Section 5.1)

Headline numbers (Figure 8): SGD-SaI achieves the lowest final training loss among all compared optimizers when pretraining GPT-2-Small (125M parameters) on OpenWebText. Specifically, it "outperform[s] Adam-mini by 0.13" in training loss and reduces validation loss by 0.03 compared to Adam-mini. The paper reports that SGD-SaI's convergence is "slightly slower" initially compared to Adam-family optimizers but achieves superior final convergence.

Memory efficiency (Table 1): For GPT-2-XL (1.5B parameters, FP32):

  • SGD-SaI state tensor memory: 5.93 GB (identical to SGDM; one momentum buffer).
  • AdamW/Adam: 11.86 GB (2× parameter size; two full buffers for m and v).
  • Prodigy: 23.72 GB (4× parameter size; stores m, v, x0, and s as detailed in Appendix B.4).
  • Adam-mini: 6.52 GB (~10% reduction over AdamW, but 0.59 GB more than SGD-SaI).
  • SGD-SaI reduces memory by 50% vs. AdamW and by 75% vs. Prodigy.

For Llama2-7B: SGD-SaI uses 25.15 GB, while AdamW, Adam, and Prodigy all hit out-of-memory (OOM) on the A100-80GB. Adam-mini uses 27.21 GB (~2 GB more than SGD-SaI). The paper explicitly notes that "for models larger than 1 billion parameters, the performance gains from Adam-Mini decrease by approximately 45%, while the reduction achieved with SGD-SaI remains around 50%" (Section 5.1, "Efficiency" paragraph)—this is a crucial scaling claim showing that Adam-mini's memory advantage erodes at scale while SGD-SaI's is structurally guaranteed.

Optimizer step time (Table 1, Figure 1, Figure 8c):

  • GPT-2-1.5B: SGD-SaI takes 68 ± 21 ms vs. AdamW's 138 ± 6 ms (~2× faster), Adam-mini's 223 ± 2 ms (~3.3× faster), and Prodigy's 360 ± 45 ms (~5.3× faster).
  • Llama2-7B: SGD-SaI takes 180 ± 30 ms vs. Adam-mini's 421 ± 22 ms (~2.3× faster).
  • On GPT-2-Small (Figure 8c), the paper reports SGD-SaI is "three times faster than Adam-mini in parameter updates."

Why Adam-mini is slower: The paper explains this in Section 5.1 and Appendix B.3. Adam-mini reduces memory for v but must still compute block-mean squared gradients, maintain exponential moving averages, and apply bias correction at every step. Its complex partitioning (splitting QKV into separate Q, K, V by attention heads) prevents simultaneous updates of coarsely grouped parameters, adding computational overhead. SGD-SaI avoids all of this—after step 1, it applies a single static scalar multiplication per block.

ViT-S/16 ImageNet-1K Pretraining (Section 5.2)

Headline numbers (Table 2): Under best hyperparameters for each optimizer:

  • SGD-SaI: 72.92 ± 0.07 (peak top-1 validation accuracy, averaged over 3 seeds).
  • AdamW: 73.04 ± 0.31 (0.12 percentage points higher peak, but 4.4× higher standard deviation).
  • Prodigy: 73.24 ± 0.21 (highest peak, 0.32 points above SGD-SaI).
  • SGDM (vanilla): 63.80 ± 0.35 (nearly 10 points below SGD-SaI—this gap is the problem SGD-SaI solves).
  • Adam: 61.56 ± 0.93 (worse than SGDM; likely due to L2 vs. decoupled weight decay interaction).
  • Adam-mini: 72.29 ± 0.43.

SGD-SaI effectively closes the gap between SGDM and AdamW—the 9.12 percentage point deficit is reduced to 0.12 percentage points. The remaining gap to the best optimizer (Prodigy) is 0.32 points.

Robustness to hyperparameters (Table 2, average accuracy across full grid):

  • SGD-SaI: 57.55 ± 18.46 (average and standard deviation across all lr × wd combinations).
  • AdamW: 37.21 ± 35.43.
  • Adam-mini: 36.65 ± 35.39.
  • Adam: 20.93 ± 22.05.
  • SGDM: 14.33 ± 19.38.

This is the paper's most striking robustness result. SGD-SaI's average accuracy is 20.34 percentage points higher than AdamW's, and its standard deviation across hyperparameters is nearly half. In practical terms, if you pick a random hyperparameter combination from the search grid, SGD-SaI gives you a model that is on average 57.55% accurate (usable, though suboptimal), while AdamW gives you a model that is 37.21% accurate (barely above chance on 1,000 classes, where chance is ~0.1%). The paper attributes this to AdamW's adaptive mechanism creating complex failure modes under poor hyperparameters, while SGD-SaI's static scaling leads to more predictable degradation (Section 4.2, "Innovation 3" context).

Convergence behavior (Figure 7): The paper acknowledges that SGD-SaI "has a slower convergence speed" (Figure 7 caption) compared to Adam-family methods, but "can still achieve comparable performance by the end of the training process." This is consistent with the design philosophy: the static g-SNR scaling provides no adaptive acceleration, so early training progresses more conservatively, but the final performance matches or approaches adaptive methods. The paper explicitly notes that "our approach is designed to have a lower memory footprint and a faster optimization speed" (Figure 7 caption), prioritizing efficiency over convergence speed.

Memory and speed for ViT (Table 3):

  • ViT-S/16 (22.9M): SGD-SaI uses 0.08 GB (same as SGDM and Adam-mini) vs. AdamW's 0.17 GB (2×). Optimizer step time: 12.4 ± 0.2 ms for SGD-SaI vs. 45.0 ± 8.0 ms for AdamW (3.6× faster), 84.0 ± 5.0 ms for Adam-mini (6.8× faster).
  • ViT-H/14 (660M): SGD-SaI uses 2.42 GB (same as SGDM) vs. AdamW's 4.86 GB (2×), Adam-mini's 2.54 GB, Prodigy's 9.70 GB (4×). Step time: 54.0 ± 13.0 ms for SGD-SaI vs. 124.0 ± 4.0 ms for AdamW (2.3× faster), 220.0 ± 20.0 ms for Adam-mini (4.1× faster).

The speed advantage over Adam-mini is particularly notable—Adam-mini is slower than standard AdamW on ViT (84 ms vs. 45 ms for ViT-S/16), because its partitioning overhead and recomputation cost outweigh the memory savings in terms of wall-clock time. SGD-SaI avoids this overhead entirely.

Parameter-Efficient Fine-Tuning: LoRA for LLMs (Section 5.3.1)

Headline numbers (Table 4): GPT-2 Medium fine-tuned with LoRA (rank 4) on the E2E NLG dataset:

  • SGD-SaI (rank 4): achieves the best scores on all five metrics—BLEU 69.9, NIST 8.81, MET 46.7, ROUGE-L 72.1, CIDEr 2.53.
  • Scaled-AdamW (rank 4, Zhang & Pilanci, 2024, the prior SOTA): BLEU 69.6, NIST 8.77, MET 46.6, ROUGE-L 71.8, CIDEr 2.52.
  • Scaled-SGD (rank 4): BLEU 69.2, NIST 8.71, MET 46.3, ROUGE-L 70.9, CIDEr 2.48.
  • AdamW (rank 4, unscaled): BLEU 68.9, NIST 8.69, MET 46.5, ROUGE-L 71.3, CIDEr 2.51.
  • SGDM (rank 4, unscaled): BLEU 66.6, NIST 8.54, MET 44.2, ROUGE-L 68.2, CIDEr 2.32—substantially worse, showing the standard SGD gap on Transformers persists in LoRA fine-tuning.

The margins are small—SGD-SaI beats scaled-AdamW by 0.3 BLEU, 0.04 NIST, 0.1 MET, 0.3 ROUGE-L, 0.01 CIDEr—but the consistency across all five metrics (SGD-SaI ranks first on every single one) is notable. More importantly, SGD-SaI achieves this with "this simple precondition on SGDM" (Section 5.3.1), without the Riemannian preconditioner that scaled-SGD and scaled-AdamW require (Zhang & Pilanci, 2024). The paper emphasizes that SGD-SaI "closes the performance gap between SGD and AdamW" even in the LoRA setting and that its block-wise scaling is sufficient to match or exceed specialized LoRA optimizers.

Hyperparameter settings: The paper uses a default learning rate of 1e-3 and weight decay of 1e-2 (Section 5.3.1), noting these were not meticulously tuned—this is presented as evidence of robustness, though it also means the comparison may not reflect each optimizer's absolute peak performance.

Parameter-Efficient Fine-Tuning: LoRA for Diffusion Models (Section 5.3.2)

Qualitative results (Figure 9): The paper compares SGD-SaI against scaled-AdamW (Zhang & Pilanci, 2024) for LoRA fine-tuning of the ChilloutMix diffusion model on a facial concept ("potter"). The prompt is "a pencil sketch of ⟨Vpotter⟩":

  • SGD-SaI (bottom row in Figure 9): All four generated images consistently depict the target face concept, with recognizable facial features and reasonable pencil-sketch styling.
  • Scaled-AdamW (top row): Two of the four generated images are marked with red bounding boxes as "significant bad cases that do not follow the prompt" (Section 5.3.2). These appear to be faces of different individuals or degraded outputs.

The paper claims this demonstrates "better parameters robustness on training and leads to better convergence in final performance" (Section 5.3.2), though the evaluation is purely qualitative—no quantitative metrics (FID, CLIP score, user study) are reported. The learning rate is set to a "default 0.1, a large enough default value" (Section 5.3.2), emphasizing that the method works without careful tuning.

Why this matters: Diffusion model fine-tuning, particularly for concept personalization, is known to be sensitive to optimization hyperparameters—too high a learning rate causes catastrophic forgetting or concept bleeding, too low fails to capture the target concept. SGD-SaI's robustness claim is that its block-wise g-SNR scaling naturally reins in the learning rate for parameter groups that would otherwise be unstable, without requiring per-task learning rate tuning.

Convolutional Neural Networks (Section 5.4)

ResNet-18 on CIFAR-10 (Figure 5, left):

  • SGD-SaI peak accuracy: 95.36% (annotated in red on Figure 5).
  • SGDM: 95.26%.
  • AdamW and Adam: Lower peaks visible in the figure's box-plot distributions, though exact numbers are not quoted in the text.
  • SGD-SaI's distribution across the hyperparameter grid is visibly tighter (less variance) than SGDM, Adam, or AdamW.

The improvement over SGDM is small (0.10 percentage points) because CNNs like ResNet-18 do not exhibit the severe Hessian heterogeneity that makes SGD fail on Transformers. This result primarily demonstrates that SGD-SaI does not hurt on architectures where vanilla SGDM already works—the g-SNR scaling is either neutral or slightly beneficial.

NATS-Bench cross-architecture evaluation (Figure 10): The paper evaluates eleven CNN architectures from NATS-Bench (selected as "ten architectures with top-10 validation accuracy and one architecture with bottom-1 validation accuracy," Appendix A.2) on CIFAR-10, CIFAR-100, and ImageNet16-120. For each architecture, the same hyperparameter grid is swept.

The histograms in Figure 10 show:

  • CIFAR-10: SGD-SaI's accuracy distribution (blue curve) is shifted rightward (higher mean) and narrower (lower variance) compared to Adam (red), AdamW (orange), SGDM (green), and Nesterov SGD (purple). The kernel density estimate (KDE) curve peaks around 80–85%, while Adam-family curves cluster around 60–75%.
  • CIFAR-100: The same pattern holds—SGD-SaI's distribution peaks higher and is tighter than all baselines.
  • ImageNet16-120: Again, SGD-SaI shows the highest mean and lowest variance.

The paper presents this as evidence of "enhanced stability and generalization" (Section 5.4) across diverse architectures, not just a single model. However, the absolute performance levels in Figure 10 appear relatively low (peaking at ~80% on CIFAR-10 and ~50% on CIFAR-100, versus ~95% for the well-tuned ResNet-18 experiment), likely because the grid search covers architectures with widely varying capacities and the reported distributions aggregate over all architectures, learning rates, and weight decays.

Ablation Studies and Robustness Checks

g-SNR temporal stability (Figure 4): The paper traces g-SNR values over training steps for three transformer blocks (shallow block 0, middle block 5, deep block 11) and distinct parameter types. The g-SNR values remain "relatively constant over time" (Section 4.2) for each traced block, with flat trajectories rather than trending curves. This is the central empirical justification for freezing scaling factors after initialization—if g-SNR values drifted substantially, static scaling would become mismatched to later training stages.

g-SNR cross-block distribution (Figure 3): Histograms of g-SNR values across parameter blocks show clear separation by block type: attention QKV projections, MLP layers, normalization layers, embeddings, and classification heads all occupy distinct regions of g-SNR space. Bias parameters (Figure 3, right subfigure) show "consistently low magnitude" g-SNR values compared to weights. The paper observes that "parameter blocks that share the same structure across different transformer layers (blocks) tend to have similar g-SNR values" (Section 4.2), with the notable exception of the first norm1 layer (which connects to patch embeddings rather than the previous block's output). This structural consistency validates that g-SNR captures architecture-determined gradient characteristics rather than random noise.

g-SNR calculation overhead (Tables 5 and 6, Appendix C): The paper profiles the wall-clock time of the one-time g-SNR computation at step 1 versus the per-step optimizer time. On an RTX 3090, the g-SNR calculation takes 14.5 ms for ViT-S/16 (vs. 13.7 ms per optimizer step for SGD-SaI afterward), 43.3 ms for ViT-H/14 (vs. 65.5 ms per step), and 267.6 ms for GPT-2-1.5B (vs. 340.1 ms per step). On an A100, the pattern is similar. The paper argues this one-time cost is "tolerable" (Appendix C) because it occurs only once and is comparable to or less than a single optimizer step, becoming negligible over thousands of training iterations.

Memory scaling with model size (Tables 1 and 3, Figure 1): Across model scales from ViT-S/16 (22.9M parameters) to Llama2-7B (7B parameters), SGD-SaI consistently uses exactly the same memory as SGDM—one momentum buffer per parameter. The paper demonstrates that this holds at all tested scales: 0.08 GB for ViT-S/16, 2.42 GB for ViT-H/14, 5.93 GB for GPT-2-1.5B, and 25.15 GB for Llama2-7B. Adam-mini's memory advantage degrades at scale: "for models larger than 1 billion parameters, the performance gains [sic: memory reduction] from Adam-Mini decrease by approximately 45%" (Section 5.1), while SGD-SaI's 50% reduction is architecturally guaranteed—it stores half the tensors, regardless of model size.

(Non-)ablation: Recomputing vs. freezing g-SNR. The paper does not compare frozen g-SNR scaling (SaI) against periodically recomputed g-SNR scaling. This is an implicit ablation: the theoretical and empirical stability arguments (Figure 4, Taylor expansion in Section 4.2) are the justification for freezing, but the paper never experimentally verifies that recomputing (e.g., every epoch or every N steps) would yield identical or worse performance. This is a gap—if recomputed scaling factors drifted enough to change the effective learning rates, the frozen version could be suboptimal for late training.

(Non-)ablation: Different partitioning granularities. The paper uses PyTorch default partitioning throughout and does not compare against finer (e.g., Adam-mini's head-wise partitioning) or coarser (e.g., per-layer) partitionings for g-SNR computation. The justification is empirical (Figure 3 shows structural consistency within PyTorch default partitions) but not experimentally verified by testing alternative groupings.

Optimizer step time breakdown (Tables 5 and 6): The paper profiles both the post-initialization per-step time and the first-step g-SNR computation overhead, showing that the one-time cost is comparable to a single optimizer step and that subsequent steps are significantly faster than Adam-mini or AdamW. This is not an ablation per se, but it isolates the computational cost of the g-SNR mechanism from the momentum and weight decay operations that SGD-SaI shares with SGDM.

Impact of random seed (Table 2 standard deviations): For ViT-S/16 peak performance, SGD-SaI shows the lowest standard deviation across three seeds (±0.07) of any optimizer, compared to Prodigy (±0.21), AdamW (±0.31), and Adam-mini (±0.43). This indicates that the g-SNR scaling factors are not unduly sensitive to the specific first-batch gradient sample—different random initializations and data orderings produce consistent scaling factors and final performance.

Performance under different learning rate schedulers (Appendix A.3): For ResNet-18, the paper switches from cosine annealing (used in other CNN experiments) to step decay (×0.1 every 80 epochs). SGD-SaI still achieves 95.36% peak accuracy, suggesting that the g-SNR scaling is not coupled to a particular learning rate schedule. However, this is tested on only one model-dataset pair.

Critical Assessment

Does SGD-SaI Actually Match or Outperform AdamW?

The paper's central claim—"SGD-SaI consistently matches or outperforms AdamW" (Abstract, Executive Summary)—is supported but with important qualifications.

Where the claim holds cleanly:

  • On GPT-2-Small pretraining (Figure 8), SGD-SaI achieves lower final training loss than both AdamW and Adam-mini. This is direct evidence of superior optimization.
  • On ViT-S/16 peak accuracy (Table 2), SGD-SaI (72.92%) is within 0.12 percentage points of AdamW (73.04%)—effectively matched. The standard deviations overlap substantially.
  • On LoRA fine-tuning (Table 4), SGD-SaI beats all baselines on all five metrics, including scaled-AdamW.
  • On CNNs (Figure 5 left, Figure 10), SGD-SaI matches or slightly exceeds all baselines.

Where the claim needs qualification:

  • On ViT-S/16 peak accuracy, Prodigy (73.24%) beats SGD-SaI by 0.32 points. While the paper frames this as "comparable" (Section 5.2), Prodigy is the strongest performer on peak accuracy. The paper's response is that Prodigy uses 4× the optimizer memory (Table 3: 0.33 GB vs. 0.08 GB for ViT-S/16) and is 6.3× slower per optimizer step (78 ms vs. 12.4 ms). This is a legitimate efficiency argument but means the claim of "outperforming state-of-the-art optimizers" (Abstract) specifically for accuracy is accurate against AdamW and Adam-mini but not Prodigy.
  • On GPT-2-Small, SGD-SaI's convergence is slower than AdamW's—the paper explicitly notes "slightly slower initial convergence speed" (Section 5.1). The final loss is better, but practitioners who care about total wall-clock time (not just final performance) might prefer AdamW if it reaches acceptable performance in fewer steps, even if the per-step time is longer. The paper does not report time-to-accuracy metrics (e.g., loss vs. wall-clock time), making it difficult to evaluate the practical speed-accuracy tradeoff.

Is the Claim of Memory Efficiency Fully Supported?

Strongly supported: The memory numbers in Tables 1 and 3, and Figure 1, are deterministic given the optimizer designs—SGD-SaI stores exactly d floating-point numbers for momentum (plus B scalars for g-SNR, which is kilobytes), while AdamW stores 2d. This is not an empirical claim but an algorithmic property, and it holds at all tested scales. The 5.93 GB reduction for GPT-2-1.5B and 25.15 GB reduction for Llama2-7B compared to AdamW in FP32 are exact consequences of the design.

Qualification: The paper compares FP32 training throughout. In practice, many practitioners train large models with mixed precision (FP16 or BF16) and/or 8-bit optimizers (Dettmers et al., 2021), which would reduce AdamW's memory overhead substantially. The paper does not compare against quantized AdamW in memory or performance. The 50% reduction is relative to FP32 AdamW, not to the memory-efficient variants that practitioners increasingly use. The paper's claim that "these methods risk a performance downgrade compared to AdamW" (Section 1, regarding quantization) is stated but not experimentally demonstrated against a quantized AdamW running at the same effective memory budget.

Does SGD-SaI Solve the "Long-Standing Challenge" of Training Transformers with SGD?

Largely supported, with scope limitations. The ViT-S/16 result—SGDM at 63.80% vs. SGD-SaI at 72.92%—is the most direct evidence. The gap between SGDM and AdamW (73.04%) is 9.24 points; SGD-SaI closes 9.12 of those points, leaving only 0.12 points. This is a convincing demonstration that per-block learning rate scaling derived from initialization-time gradients is sufficient to make SGD competitive on Transformers, and that the historical failure of SGD on Transformers was primarily a learning rate allocation problem rather than a fundamental limitation of the SGD update rule.

Limitations:

  • The result is demonstrated on exactly one Transformer architecture family (ViT-S/16) for full training. GPT-2 pretraining is the only other Transformer full-training experiment, and it compares against AdamW and Adam-mini (which both work) rather than showing that SGDM fails. The LoRA experiments are fine-tuning, not full training—the dynamics of fine-tuning a pretrained model are different from training from scratch.
  • The largest model fully trained is GPT-2-Small (125M parameters). GPT-2-XL (1.5B) and Llama2-7B are only profiled for memory and speed, not trained. The paper acknowledges this explicitly: "our study is constrained by limited computational resources, preventing us from conducting large-scale pre-training on more extensive models such as Llama-2-7B" (Section 7). This is a significant gap—the paper's central claim is about memory efficiency at scale, but the empirical evidence that SGD-SaI works at scale (produces good models, not just low memory usage) is limited to 125M parameters. It is plausible that g-SNR stability (Figure 4) holds at larger scales, but it is not demonstrated.

Is the Hyperparameter Robustness Claim Convincing?

Strongly supported for the tested settings. Table 2's average accuracy comparison is the paper's most striking result: SGD-SaI at 57.55% vs. AdamW at 37.21% across the same hyperparameter grid. This is a genuine and practically important finding—it means that with a fixed tuning budget (e.g., trying 4 learning rates and 3 weight decays = 12 runs), SGD-SaI is far more likely to produce a usable model than AdamW.

Qualification regarding the hyperparameter range: The grid search ranges (lr ∈ {0.1, 0.01, 0.001, 0.0001}, wd ∈ {0.01, 0.001, 0.0001}) include values that are clearly inappropriate for AdamW (e.g., lr = 0.1 with Transformers is known to cause divergence for adaptive methods). This biases the average toward SGD-SaI, which is more tolerant of excessively high learning rates. A fairer robustness comparison would use each optimizer's recommended hyperparameter range, but the paper's point—that SGD-SaI degrades gracefully under poor hyperparameters while AdamW collapses—remains valid even if the specific 20-point gap is inflated by including values AdamW users would not try.

What Experiments Are Missing?

Full training runs at scale (1B+ parameters). This is the single biggest gap. The paper's value proposition is memory efficiency for large models, but it has not demonstrated that SGD-SaI actually produces good models when training Transformers at the scales where memory efficiency matters most. The profiling results show that SGD-SaI can run where AdamW OOMs, but not that it produces competitive models under those conditions.

Comparison against 8-bit Adam / quantized optimizers. The paper criticizes quantization-based memory reduction for introducing "additional implementation complexity" and potential performance degradation, but never experimentally compares against 8-bit AdamW with the same effective memory budget. Since 8-bit AdamW is widely used in practice, this is a missing baseline.

Time-to-accuracy curves. The paper reports final performance and per-step time, but not the product: how long does SGD-SaI take to reach a given accuracy versus AdamW? If SGD-SaI takes 2× as many steps but each step is 2× faster, the wall-clock time to convergence is identical. If it takes more steps than the speedup compensates for, the memory savings come at a time cost.

Ablation on g-SNR recomputation frequency. The paper's conceptual innovation is that g-SNR can be frozen, but it never tests this against periodic recomputation. If recomputing g-SNR every N steps (at negligible amortized cost) improved performance, the "scaling at initialization only" claim would be empirically weaker even if the conceptual insight about temporal stability holds.

Larger-scale Transformer training. GPT-2-Small (125M) is a small language model by modern standards. Testing on GPT-2-Medium (345M) or GPT-2-Large (774M) would strengthen the scaling claim without requiring Llama2-7B resources.

Ablation on batch size for g-SNR computation. The g-SNR is computed from a single first batch. If the first batch is unrepresentative (e.g., unusually noisy or small), the scaling factors could be distorted. The paper does not test sensitivity to first-batch composition or compare against averaging g-SNR over multiple initial batches.

Where Claims Hold Conditionally

The claim that SGD-SaI "matches or outperforms AdamW" holds clearly for:

  • Final training/validation loss on GPT-2-Small pretraining (Figure 8).
  • Peak accuracy on ViT-S/16 (Table 2, within noise).
  • All metrics on LoRA fine-tuning (Table 4, small but consistent margins).
  • Average accuracy across hyperparameter grids on ViT and CNNs (Table 2, Figure 10).

It does not hold for:

  • Peak accuracy against Prodigy on ViT-S/16—Prodigy wins by 0.32 points (Table 2).
  • Convergence speed—SGD-SaI is slower than AdamW in early training (Figure 7, Figure 8).

The memory efficiency claim (50% reduction vs. AdamW) holds unconditionally in FP32—it is guaranteed by the algorithm design, not an empirical result subject to experimental variation.

6. Limitations and Trade-offs

Scope Limited to Full Precision (FP32) Training — No Comparison Against Quantized AdamW, the Industry Default

The paper evaluates all optimizers exclusively in full-precision (FP32) settings. The memory comparisons in Tables 1 and 3, Figure 1, and the headline claims of 50% memory reduction relative to AdamW all assume that AdamW stores two full-precision tensors (m and v) per parameter. However, practitioners training large models overwhelmingly use mixed-precision training (FP16 or BF16) combined with low-bit optimizer states—8-bit Adam (Dettmers et al., 2021) is a widely-adopted default in frameworks like Hugging Face Transformers and bitsandbytes. In this regime, AdamW's v tensor is stored in 8-bit precision, cutting its optimizer state memory by approximately 75% rather than the 50% reduction SGD-SaI achieves relative to FP32 AdamW.

The paper acknowledges these methods in Section 2 ("Related Work") and Section 3 ("Memory Efficient Adam"), listing 8-bit Adam (Dettmers et al., 2021), 4-bit optimizers (Li et al., 2023a), and sign-based methods (Bernstein et al., 2018) as prior approaches that "can compromise performance." However, the paper never experimentally measures how much performance 8-bit AdamW sacrifices relative to FP32 AdamW, whether SGD-SaI outperforms 8-bit AdamW when both operate at similar effective memory budgets, or whether SGD-SaI itself can benefit from quantization (its single momentum buffer could also be quantized).

Consequence. The memory advantage of SGD-SaI over practical AdamW deployments is substantially smaller than the 50% headline figure suggests. An 8-bit AdamW storing m in FP32 and v in INT8 uses approximately 1.25× the parameter size in optimizer state, versus SGD-SaI's 1.0×—a 20% reduction, not 50%. The paper's claim that its method "reduces memory usage by 5.93 GB for GPT-2 (1.5B parameters) and 25.15 GB for Llama2-7B compared to AdamW in full-precision training settings" (Abstract, emphasis added) is technically accurate for FP32 but may mislead practitioners who train in mixed precision. For them, switching from 8-bit AdamW to SGD-SaI would save only ~1.5 GB for GPT-2-1.5B and ~6 GB for Llama2-7B—meaningful but far less dramatic.

Additionally, without a comparison of final model quality against 8-bit AdamW at matched effective memory, the practitioner cannot determine whether SGD-SaI's memory savings are worth any potential performance difference versus simply quantizing AdamW.

Evidence in the paper. The paper provides only qualitative criticism: "quantization introduces additional implementation complexity" (Section 3), and "these methods risk a performance downgrade compared to AdamW" (Section 1). No experiment compares SGD-SaI against any quantized optimizer. All memory profiling (Tables 1, 3, Figure 1) uses FP32 for all methods.

Mitigation status. Not addressed. The paper frames FP32 as the reference setting and does not discuss how its claims translate to mixed-precision or quantized-optimizer regimes. The authors acknowledge in Appendix B that "the estimated minimum memory requirement for the state tensors in both Adam and AdamW is approximately twice the size of the gradient tensors" but never extend this analysis to the quantized case. This limitation is not flagged in Section 7 (Limitations).

No Full Training Runs at the Billion-Parameter Scale Where Memory Efficiency Matters Most

The paper's central value proposition is memory efficiency for large-scale Transformer training. The Abstract touts "substantial memory savings... reducing memory usage by 5.93 GB for GPT-2 (1.5B parameters) and 25.15 GB for Llama2-7B compared to AdamW." Yet the only model that is fully trained (not just profiled) with SGD-SaI is GPT-2-Small (125M parameters, Section 5.1). The ViT experiments use ViT-S/16 (22.9M parameters, Section 5.2). The profiling results for GPT-2-XL (1.5B) and Llama2-7B measure optimizer step time and state tensor memory (Tables 1, 5, 6), but no training loss curves, validation metrics, or downstream task evaluations are provided at these scales. The paper acknowledges this explicitly in Section 7:

"our study is constrained by limited computational resources, preventing us from conducting large-scale pre-training on more extensive models such as Llama-2-7B. This remains an avenue for future research."

Consequence. The paper demonstrates that SGD-SaI fits in memory at the 1.5B and 7B scales, but does not demonstrate that it produces competitive models at these scales. This is a critical gap because the core assumption—that g-SNR values computed at initialization remain stable and appropriate learning rate scalars throughout training—has only been empirically verified on small-to-medium models. At larger scales, several effects could break this assumption:

  • Loss landscape complexity increases with scale. Larger models explore higher-dimensional, more complex loss surfaces. The first-order Taylor argument for g-SNR constancy (Section 4.2) is local—it assumes the gradient distribution does not shift substantially as parameters move. In large models, phase transitions in learning dynamics (e.g., grokking, sudden loss drops, emergent circuits) could correspond to shifts in gradient distributions that make initialization-time g-SNR values inappropriate for later training stages.
  • Training instability scales with model size. Large Transformers are known to be sensitive to optimization hyperparameters—learning rate warmup, gradient clipping, and careful scaling of initialization are all critical. Whether SGD-SaI's static per-block learning rate scaling is sufficient to stabilize billion-parameter training without the dynamic adaptation that Adam provides is an open question. The paper's ViT-S/16 convergence curve (Figure 7) already shows slower early convergence than AdamW; at larger scales, this gap could widen or lead to training divergence.
  • The memory savings claim is strongest at the scales not trained. The 25.15 GB reduction for Llama2-7B is computed from the algorithmic property that SGD-SaI stores one tensor per parameter while FP32 AdamW stores two. This is arithmetically correct, but the paper's claim that SGD-SaI "can work with large models like Llama-2 (7B) without running into out-of-memory (OOM) errors" (Section 5.1) only demonstrates that it loads the optimizer state—not that training converges or produces a usable model.

Practitioners deciding whether to adopt SGD-SaI for billion-parameter training face a leap of faith: the method should work based on theoretical arguments and 125M-parameter evidence, but it has not been shown to work.

Evidence in the paper. Section 5.1 explicitly separates "pre-train GPT-2-Small (125M)" (full training with loss curves in Figure 8) from profiling of larger models (Tables 1, 5, 6, which measure only memory and step time). Section 7 acknowledges this as a limitation. No training metrics are reported for any model above 125M parameters. The ViT experiments cap at 22.9M parameters for full training; ViT-H/14 (660M) is profiled only (Table 3).

Mitigation status. Acknowledged but unresolved. The paper frames the profiling results as "preliminary analyses [that] indicate the potential scalability of SGD-SaI" (Section 7) and calls for "comprehensive evaluations on larger-scale models... to establish its effectiveness and efficiency in such settings fully." No partial training runs (e.g., a few thousand steps of GPT-2-XL or Llama2-7B with loss tracking) are provided to bridge the gap between 125M and 7B.

Slower Convergence Speed — and the Missing Time-to-Accuracy Analysis

Across all full-training experiments where convergence curves are shown, SGD-SaI converges more slowly than Adam-family methods in the early and middle stages of training. On ViT-S/16 (Figure 7), the caption notes that "our method has a slower convergence speed" but "can still achieve comparable performance by the end of the training process." On GPT-2-Small (Figure 8), the training loss curves show SGD-SaI starting higher and taking longer to cross below AdamW's curve, even though the final training loss is lower. The paper attributes this to the design philosophy: "our methods ensure a steady and stable update during training, allowing the model to converge better in a given task with sufficient training steps" (Section 7).

However, the paper never reports time-to-accuracy: how many wall-clock hours does SGD-SaI require to reach a given validation accuracy or loss threshold compared to AdamW? This is the metric that matters for practitioners with fixed compute budgets and deadlines—they care about total training time, not just final performance after an arbitrary number of epochs. The per-step speed advantage of SGD-SaI (2–3× faster optimizer steps than AdamW, per Tables 1 and 3) partially offsets the slower per-step convergence, but whether the net effect is favorable depends on the exact ratio of convergence slowdown to step speedup, which varies by task, model size, and training phase.

Consequence. A practitioner choosing between SGD-SaI and AdamW for a project with a fixed training budget (e.g., "train for 48 hours on 8 GPUs") cannot determine from this paper which optimizer will yield better model quality at the deadline. SGD-SaI might reach higher final performance if given unlimited training time (the paper's evidence suggests this for GPT-2-Small), but it might underperform AdamW at any realistic intermediate stopping point. The tradeoff is particularly acute in two scenarios:

  • Hyperparameter tuning: If each training run is slow, grid-searching learning rates and weight decays becomes expensive. The paper argues that SGD-SaI's robustness reduces the need for extensive tuning (Section 4.2, "Innovation 3" context), but the per-run wall-clock time for a fixed number of epochs may still be longer if convergence is slower.
  • Early stopping: Many production training pipelines use early stopping based on validation loss. If SGD-SaI's validation loss improves more slowly but eventually reaches a better minimum, an early-stopping criterion might terminate training prematurely, leaving performance on the table. The paper's validation curves (GPT-2, Figure 8; ViT, Figure 7) show SGD-SaI lagging AdamW for much of training.

Additionally, the slower initial convergence interacts with the paper's robustness claims. The hyperparameter grid search results (Table 2, "Avg@top1") show that SGD-SaI maintains higher average accuracy across hyperparameters, but this is measured at the end of a fixed number of epochs. If training were stopped earlier (at a fixed time budget), the gap might shift in AdamW's favor because AdamW reaches higher intermediate accuracy more quickly before potentially degrading if hyperparameters are poor.

Evidence in the paper. Figure 7 (ViT training curves) and Figure 8 (GPT-2 training curves) show slower early convergence for SGD-SaI. The paper acknowledges this in Section 7: "we might observe that the convergence speed is relatively lower than Adam's in the early stage of training." No time-to-accuracy analysis is provided. The optimizer step time measurements (Tables 1, 3) are per-step costs, not end-to-end training throughput.

Mitigation status. Partially acknowledged in Section 7: "as our primary focus is to investigate the effectiveness of the SaI approach, we left the acceleration of convergence speed in future work." The paper does not propose any modification to SGD-SaI to improve convergence speed (e.g., warmup schemes, dynamic g-SNR recomputation, or hybrid strategies that start with AdamW and switch to SGD-SaI). It also does not measure whether the slower convergence is due to the static g-SNR scaling being too conservative for early training (when gradients are noisy and g-SNR values may be low, producing small effective learning rates) or due to some other aspect of the method.

The g-SNR Scaling Factor Is Frozen After Step 1 — No Mechanism to Adapt if Gradient Distributions Shift

SGD-SaI's defining design choice is that per-block learning rate scaling factors are computed once from the very first training batch and then frozen permanently. The paper justifies this through the empirical observation that g-SNR values "remain relatively constant over time" (Figure 4, Section 4.2) and a first-order Taylor argument that gradient statistics change slowly under small learning rates (Section 4.2). However, the paper never experimentally tests what happens when this assumption is violated—for example, by periodically recomputing g-SNR and comparing performance against the frozen version, or by testing on tasks where the gradient distribution is known to shift substantially during training.

Consequence. Several realistic training scenarios could cause the initialization-time g-SNR to become misaligned with later gradient statistics:

  • Curriculum learning or data distribution shifts. If training data order is not i.i.d. (e.g., a curriculum that presents easier examples first, or domain-adaptive pretraining where later data differs from early data), the gradient distribution at step 1 may not be representative of the distribution throughout training. A first batch of easy examples might produce clean, high-SNR gradients that lead to aggressive learning rate scaling; when the model later encounters harder or noisier data, those aggressive rates could destabilize training.
  • Phase transitions in learning. Large models sometimes exhibit sudden qualitative changes in behavior during training (grokking, emergence of in-context learning, sudden loss drops). These transitions likely correspond to changes in gradient distributions as the model shifts from memorization to generalization or from one solution regime to another. Frozen g-SNR scaling from initialization cannot adapt to these transitions.
  • Fine-tuning from a pretrained checkpoint. In the paper's LoRA experiments (Section 5.3), g-SNR is computed on the first fine-tuning batch, where the model parameters are already at a well-optimized point (pretrained weights). In this setting, the gradient distribution may be fundamentally different from random initialization, and the temporal stability observed in Figure 4 (which was measured from scratch on ViT-S/16) may not hold. The paper does not verify g-SNR stability in the fine-tuning regime.
  • The Taylor argument has limited range. The paper's theoretical justification (Section 4.2) relies on "η\eta [being] sufficiently small to stabilize the training process" and θt+1(i)≈θt(i)\theta^{(i)}_{t+1} \approx \theta^{(i)}_t. But over many training steps, parameters can move far from initialization—this is the point of training. The Taylor argument guarantees local stability, not global constancy. Whether the relative ordering of g-SNR values across blocks—which is what matters for learning rate scaling—remains stable over large parameter displacements is an empirical question the paper does not investigate beyond the ViT-S/16 traces in Figure 4.

Evidence in the paper. Figure 4 provides the only empirical evidence for g-SNR stability, tracing three transformer blocks and several distinct parameter types in ViT-S/16 over training. The paper does not report similar traces for GPT-2 pretraining, CNN training, or any fine-tuning scenario. The first-order Taylor argument (Section 4.2) is a qualitative plausibility argument, not a bound on how much g-SNR can change over the course of training. No ablation compares frozen vs. periodically recomputed g-SNR.

Mitigation status. Not addressed. The paper treats g-SNR constancy as an established fact based on Figure 4 and the Taylor heuristic, and does not discuss failure modes where this assumption breaks. The possibility of recomputing g-SNR periodically (at negligible additional cost—a single forward/backward pass plus g-SNR computation every N steps) is not explored. The paper also does not measure the sensitivity of final performance to errors in the g-SNR estimates (e.g., by adding noise to the scaling factors or using g-SNR computed from a different batch).

The Paper Does Not Combine Search and Revisions — Missing the Complementary Gains Available from Both Axes

The paper's conceptual framework decomposes test-time compute into modifications to the proposal distribution (what the model generates, improved via iterative revisions) and the verifier (how outputs are selected, improved via PRM-guided search). However, the paper studies these two mechanisms entirely independently. Section 8 explicitly acknowledges this gap:

"We did not experiment with PRM tree-search techniques in combination with revisions."

All search experiments (Section 5) use the base PaLM 2-S* model with few-shot prompting as the proposal distribution. All revision experiments (Section 6) use the revision model to generate candidate answers, with selection performed by best-of-N weighted or majority voting (no beam search or lookahead search among revisions). The two mechanisms never interact.

Consequence. This is a significant limitation because the paper's own analysis shows that search and revisions have complementary strengths across the difficulty spectrum. Beam search against the PRM excels on medium-difficulty problems (Figure 3, right, bins 3–4) where the model needs to explore different solution strategies, while sequential revisions excel on easy problems (Figure 7, right, bins 1–2) where the model's initial output is roughly correct and benefits from local refinement. A combined system that uses the revision model as the proposal distribution within a PRM-guided search tree could potentially:

  • Generate higher-quality candidate steps in beam search by conditioning each beam on its own previous (rejected) attempts.
  • Use the PRM to decide when to continue revising versus when to abandon a revision chain and restart with a different approach.
  • Apply lookahead search over revision steps, using the PRM to evaluate whether a revision trajectory is improving.

The current results therefore represent a lower bound on what a fully integrated approach could achieve. The compute-optimal scaling analysis (Section 3.1) selects between search-only and revision-only strategies per difficulty bin; it does not optimize over combinations that use both within a single inference budget.

Practically, this means a system builder using this paper's findings cannot determine whether to deploy revisions, search, or both for a given application. The paper provides evidence that each mechanism is valuable in isolation, but not whether their benefits are additive, redundant, or interfering when combined.

Evidence in the paper. Section 8 explicitly states the gap. No experiment runs revision model outputs through PRM beam search, or uses PRM scores to guide the revision chain. The compute-optimal policies in Figures 4 and 8 select among strategy variants within each mechanism (search methods for Figure 4, sequential-to-parallel ratios for Figure 8) but never across mechanisms.

Mitigation status. Acknowledged as future work in Section 8: "We leave the combination of these methods... to future work." The paper provides no preliminary results, theoretical analysis, or design sketches for how such combination could work.

Difficulty Estimation Cost Is Not Included in the Compute Budget — and Is Impractically High

The entire compute-optimal framework depends on knowing each prompt's difficulty before allocating the inference budget. The paper's method for difficulty estimation—generating 2048 samples per question and computing either the ground-truth pass@1 rate (oracle) or the average PRM final-answer score (predicted)—is extraordinarily expensive. The paper explicitly acknowledges this in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The paper frames this as an exploration-exploitation tradeoff: compute spent assessing difficulty versus compute spent solving the problem. However, the numbers make this tradeoff heavily lopsided in the current implementation. For a single prompt, generating 2048 samples costs 2048 generations—this is the entire budget for the largest test-time compute configurations studied (256–512 generations), and 8–128× larger than the budgets where the paper claims 4× efficiency gains (e.g., 16 generations matching 64, Figure 4; 64 generations matching 256, Figure 8).

Consequence. The 4×4\times efficiency gains reported throughout the paper are computed after difficulty is already known, without amortizing the cost of learning the difficulty estimate. In a realistic deployment, the total cost would be difficulty_estimation_cost + strategy_execution_cost, and the former dwarfs the latter. For a single prompt:

  • Oracle difficulty: 2048 generations to determine the correct difficulty bin, then 16–256 generations for the actual inference. Total cost: 2064–2304 generations. A best-of-N baseline would have used 2048 generations for difficulty estimation plus its own inference budget. The fair comparison would give best-of-N the total budget (difficulty estimation + inference), which would dramatically shrink or eliminate the reported gains.
  • Predicted difficulty: 2048 generations evaluated with the PRM (still requires generation, just doesn't need ground-truth labels), then 16–256 generations for inference. Same total cost.

The paper acknowledges this but offers no solution beyond suggesting future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8). No such model is developed or evaluated. The compute-optimal framework is therefore analytically valid (it shows that if difficulty were known cheaply, large gains are possible) but not deployable in the form presented—the cost of obtaining the difficulty signal exceeds the savings from using it.

For the framework to become practical, the difficulty estimation cost would need to be reduced by at least two orders of magnitude (from 2048 generations to something like 10–50). The paper provides no evidence that such a reduction is achievable without sacrificing the accuracy of difficulty bin assignment.

Evidence in the paper. Section 3.2 describes the cost and explicitly excludes it from experiments. The difficulty estimation procedure is detailed: 2048 samples per question, scored either by ground-truth correctness (oracle) or by the PRM (predicted), then binned into quintiles. No experiment includes this cost in the reported generation budgets. Figure 4 and Figure 8 show performance as a function of the strategy execution budget, not the total budget including difficulty estimation.

Mitigation status. Acknowledged in Section 3.2 with the statement that estimating difficulty "in an computationally efficient manner" is "a key avenue for future work." Section 8 reiterates the need for models that "directly predict difficulty." No mitigation strategy (e.g., adaptive difficulty estimation using far fewer samples, amortization across many prompts from the same distribution, or difficulty prediction from question text embeddings) is tested.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper is best characterized as a conceptual reframing with substantial practical consequences, not a paradigm shift. It does not introduce a fundamentally new optimization algorithm—the update rule is standard stochastic gradient descent with momentum. Rather, it demonstrates that a widely-held belief in the community—that adaptive, per-parameter, time-varying learning rates are necessary for training Transformer architectures—is incorrect. The evidence that SGDM with appropriately chosen per-block constant learning rate multipliers can match AdamW on ViT-S/16 (72.92% vs. 73.04%, Table 2) and outperform it on GPT-2-Small pretraining (Figure 8) overturns this assumption.

The magnitude of the shift is meaningful but bounded. For practitioners, the primary implication is that the optimizer memory bottleneck—which forces tradeoffs between model size, batch size, and sequence length—is partially negotiable: halving optimizer state memory (and thus freeing GPU capacity for larger models or batches) does not require accepting degraded final performance, as was the case with Adafactor or sign-based methods. The paper demonstrates this concretely: Llama2-7B, which causes OOM errors with AdamW and Prodigy in FP32 on an A100-80GB, fits comfortably with SGD-SaI using 25.15 GB of optimizer state (Table 1). This directly expands the set of models and batch configurations accessible on fixed hardware.

However, the shift is bounded because the paper does not demonstrate full training at these scales—the profiling shows memory feasibility but not training convergence or final model quality. Additionally, the convergence speed penalty (visible in Figure 7 and Figure 8, acknowledged in Section 7) means that SGD-SaI's wall-clock advantage over AdamW is smaller than the per-step speedup suggests, and for practitioners with fixed time budgets rather than fixed memory budgets, AdamW may remain preferable. The reframing is therefore: the community should no longer treat AdamW as the default and only viable choice for Transformers, but SGD-SaI is not yet a drop-in replacement for all settings.

The paper also reconciles conflicting evidence about SGD's viability for Transformers. Prior attempts to train Transformers with SGD produced substantially worse results than AdamW (the paper's own SGDM baseline achieves 63.80% vs. 73.04% for AdamW on ViT-S/16, Table 2), leading to the conclusion that adaptive methods are architecturally necessary. The paper shows that this failure was not fundamental—it was a learning rate allocation problem. SGDM with per-block scaling derived from g-SNR closes 9.12 of the 9.24 percentage point gap. This shifts the explanation of SGD's Transformer failure from "SGD cannot handle the Hessian heterogeneity" to "SGD with a uniform learning rate cannot handle the Hessian heterogeneity, but SGD with block-level learning rate differentiation can." This is a more precise and actionable diagnosis.

The work redirects research attention in two ways:

More attractive research directions:

  • Cheap gradient statistics at initialization as a general design tool. The finding that a single-batch computation of g-SNR captures structurally meaningful, temporally stable information about gradient distributions (Figures 3, 4) opens the door to using initialization-time gradient statistics for other purposes: predicting training difficulty, guiding mixed-precision allocation, identifying problematic layers, or automating learning rate schedule design. The connection to pruning-at-initialization and zero-cost NAS literature (Section 2) suggests a broader principle: gradient patterns before training carry durable information about training dynamics.
  • Decoupling structural from trajectory-dependent optimization information. The paper's implicit design principle—compute structural information (g-SNR) once at initialization, handle trajectory-dependent dynamics with simple momentum—suggests a general approach to optimizer design that could extend beyond learning rate scaling. For example, initialization-time gradient statistics could inform per-block momentum coefficients or weight decay strengths.
  • Robustness through simplicity. The finding that SGD-SaI degrades far more gracefully under poor hyperparameters than AdamW (average accuracy of 57.55% vs. 37.21% across the ViT-S/16 grid, Table 2) suggests that Adam's dynamic adaptation can amplify hyperparameter sensitivity rather than mitigating it. This motivates research into understanding why adaptive methods can have worse average-case behavior and whether other adaptive mechanisms share this vulnerability.

Less attractive research directions:

  • Further incremental improvements to Adam's second-moment compression. The paper demonstrates that the dynamic, time-varying second-moment mechanism is not necessary—its function can be replaced by a static, initialization-time proxy. This weakens the case for investing in better approximations of v_t (finer-grained Adafactor, more sophisticated quantization schemes for v_t), since the entire apparatus can be eliminated. The paper specifically shows that Adam-mini—the state-of-the-art in block-wise v_t compression—is slower per optimizer step than AdamW (Table 1: 223 ms vs. 138 ms for GPT-2-1.5B) and that its memory advantage degrades at scale (45% reduction for >1B parameters vs. SGD-SaI's structural 50%). This suggests the compression approach faces diminishing returns and inherent computational overhead.
  • Complex, hand-designed per-layer learning rate schedules. The paper shows that a simple, automatically computed g-SNR scaling matches or exceeds the performance of carefully tuned uniform learning rates. This reduces the motivation for labor-intensive per-layer or per-block learning rate tuning, since the g-SNR provides this automatically from the first batch.

Follow-Up Research This Work Enables

Large-scale validation: full training of a 1B+ parameter Transformer with SGD-SaI versus AdamW. The single most important open question is whether SGD-SaI works at the scales where its memory efficiency matters most. The paper profiles memory and step time for GPT-2-XL (1.5B) and Llama2-7B (Tables 1, 5, 6) but never trains these models. A strong follow-up would pretrain GPT-2-Large (774M) or GPT-2-XL (1.5B) from scratch with both SGD-SaI and AdamW, matching total FLOPs or wall-clock time, and report: (a) final training/validation loss; (b) downstream perplexity on held-out text; (c) convergence curves in both steps and wall-clock hours; (d) whether the g-SNR stability observed at 125M parameters (Figure 4) persists. A negative result—SGD-SaI diverging, underperforming substantially, or requiring careful learning rate warmup not needed at smaller scales—would establish the practical boundaries of the approach and is as scientifically valuable as a positive result.

Time-to-accuracy analysis comparing SGD-SaI and AdamW at matched memory budgets. The paper reports per-step optimizer time (Tables 1, 3) and final accuracy (Table 2, Figure 8) but never the product: how many wall-clock hours to reach a target validation loss or accuracy. A follow-up study would run SGD-SaI and AdamW on identical hardware for a fixed wall-clock budget (e.g., 24, 48, 72 hours) and compare: (a) which optimizer achieves lower validation loss at each checkpoint; (b) whether SGD-SaI's per-step speed advantage compensates for its slower per-step convergence progress. If AdamW reaches acceptable performance in fewer wall-clock hours despite being slower per step, practitioners with time constraints would rationally prefer it. If SGD-SaI catches up and surpasses AdamW within typical training durations, the case for switching becomes stronger. This study should include the effect of the one-time g-SNR computation overhead (267.6 ms for GPT-2-1.5B on RTX 3090, Table 5), which is negligible over thousands of steps but measurable for very short training runs.

Comparison against 8-bit AdamW at matched effective memory. The paper claims 50% memory reduction versus FP32 AdamW, but practitioners predominantly use mixed-precision training with quantized optimizer states. A direct comparison on GPT-2-Medium or GPT-2-Large would measure: (a) final model quality of SGD-SaI (FP32 momentum) versus 8-bit AdamW (FP32 m, INT8 v) at the same effective optimizer memory budget; (b) per-step throughput including the quantization/dequantization overhead; (c) whether SGD-SaI still shows a robustness advantage (higher average accuracy across a hyperparameter grid, analogous to Table 2) when the baseline is 8-bit AdamW rather than FP32 AdamW. This would determine whether the paper's memory efficiency claims translate to the deployment regimes that practitioners actually use. If SGD-SaI matches 8-bit AdamW in quality while being simpler to implement (no quantization logic), that is a strong practical argument. If 8-bit AdamW's quality is higher due to its adaptive mechanism, the paper's value proposition narrows to settings where even 8-bit AdamW's memory is prohibitive.

g-SNR recomputation versus freezing: an empirical test of the temporal stability assumption. The paper freezes g-SNR after step 1 based on empirical observation of stability (Figure 4) and a Taylor heuristic (Section 4.2), but never experimentally compares frozen versus periodically recomputed g-SNR. A targeted ablation would: (a) compute g-SNR every k steps (k = 100, 1000, 10000) and apply the updated scaling factors; (b) measure final accuracy and convergence speed on ViT-S/16 and GPT-2-Small; (c) plot how the g-SNR values themselves change over training at each recomputation frequency. If recomputed g-SNR values drift substantially but final performance is unchanged, the frozen assumption is empirically validated—the scaling factors are robust to the specific value. If recomputed g-SNR improves convergence speed (because scaling factors adapt to changing gradient distributions), that would motivate a hybrid approach: compute g-SNR periodically at negligible amortized cost to get the best of both static efficiency and dynamic adaptation.

g-SNR as a diagnostic for training instability and architecture quality. The paper observes that g-SNR varies systematically by parameter block type (Figure 3) and is sensitive to topological differences (the first norm1 layer differing from subsequent ones). This suggests g-SNR could serve as a pre-training diagnostic tool. A follow-up study would: (a) compute g-SNR distributions for different architectures (varying depth, width, attention head count, normalization placement) before training; (b) correlate these distributions with final training stability (likelihood of divergence, gradient spike frequency, loss curve smoothness) across hundreds of training runs; (c) test whether architectures with extremely skewed g-SNR distributions (some blocks having near-zero SNR while others are very high) are more prone to training instability. If g-SNR predicts training difficulty, it could be used for architecture screening before committing expensive training resources—analogous to how zero-cost NAS proxies predict final accuracy from initialization-time statistics. The paper's existing data (Figure 3, Figure 4) provides the methodology; the extension is to systematically correlate with training outcomes across a design space.

SGD-SaI for non-Transformer architectures with known heterogeneity problems. The paper demonstrates SGD-SaI on Transformers (ViT, GPT-2), LoRA fine-tuning, diffusion models, and CNNs. A natural stress test is architectures known to have severe layer-wise or block-wise gradient heterogeneity: graph neural networks (where message-passing layers and node update layers have very different gradient scales), mixture-of-experts models (where different experts receive vastly different amounts of training signal), or very deep ResNets (where gradient magnitude decays with depth). A failure case—SGD-SaI underperforming AdamW substantially on such architectures—would reveal that g-SNR captures only certain types of gradient heterogeneity and that the temporal stability assumption breaks when gradient distributions genuinely shift during training (e.g., as different experts are activated at different training stages). A success case would strengthen the universality claim. The paper's CNN experiments (Figure 10) show benefit on already-well-behaved architectures; testing on known problematic architectures would be more informative.

Practical Applications and Downstream Use Cases

Training large language models on hardware-constrained clusters. The paper's profiling results (Table 1) show that Llama2-7B with SGD-SaI fits on a single A100-80GB in FP32 (25.15 GB optimizer state) while AdamW, Adam, and Prodigy all OOM. For academic labs or small companies with limited GPU resources (e.g., 4–8 A100s or RTX 3090s/4090s), this directly enables full-parameter fine-tuning or continued pretraining of 7B-scale models that would otherwise require parameter-efficient methods (LoRA, adapters) or aggressive quantization. The 25 GB savings per GPU means that a 4-GPU node can allocate ~100 GB of combined optimizer memory to other uses—larger batch sizes (improving training stability and throughput), longer context lengths, or an additional 7B parameter model shard. For a lab training a custom 7B model from scratch, the 50% optimizer memory reduction could mean the difference between fitting on available hardware and needing to rent cloud instances at substantially higher cost.

Hyperparameter tuning for Transformer training with limited compute budgets. The paper's robustness results (Table 2: SGD-SaI average accuracy of 57.55% vs. AdamW's 37.21% across the same hyperparameter grid) have direct practical implications for hyperparameter search. In a typical workflow, a practitioner might run 10–20 training runs with different learning rates and weight decays, then select the best configuration for full-scale training. With AdamW, the risk of selecting a hyperparameter combination that produces near-random performance (~37% on ImageNet-1K) due to a single poor choice in the search grid is significant—one would need to run many trials to confidently find the good configurations. With SGD-SaI, the performance distribution is tighter and higher-mean, meaning fewer trials are needed to identify a satisfactory configuration, and the cost of accidentally selecting a bad hyperparameter is lower. For a lab that spends 80% of its compute on hyperparameter tuning, this robustness could reduce tuning costs by 30–50% simply by reducing the number of grid points needed to cover the viable region.

Deploying custom fine-tuned models under strict GPU memory constraints. The LoRA fine-tuning results (Table 4: SGD-SaI matching or exceeding scaled-AdamW on all five NLG metrics) and the diffusion model fine-tuning results (Figure 9: better visual concept capture than scaled-AdamW) suggest that SGD-SaI is effective for parameter-efficient fine-tuning, where optimizer memory is a smaller fraction of total memory but still non-trivial when fine-tuning large base models. A practical scenario: a company wants to deploy hundreds of LoRA-adapted variants of a 7B base model (one per customer or task), with each variant fine-tuned on different data. The base model is shared, but each fine-tuning job requires optimizer state. Using SGD-SaI over AdamW saves ~2 GB per concurrent fine-tuning job at the 7B scale—or enables running more concurrent jobs on the same hardware. The paper's finding that SGD-SaI works with default PyTorch partitioning (unlike Adam-mini, which requires manual restructuring) means this can be integrated into existing fine-tuning pipelines with minimal code changes—the optimizer implementation is "extremely straightforward" and "requires only minimal modifications to the existing SGD optimizer" (Section 4.3).

When to Prefer This Method

The paper explicitly positions SGD-SaI against AdamW, Adam-mini, and Prodigy, with clear tradeoffs along memory, speed, peak performance, and robustness. The following decision rule is grounded in the paper's experimental evidence:

  • Prefer SGD-SaI when optimizer memory is the binding constraint. If you cannot fit your desired model, batch size, or sequence length on available GPUs due to AdamW's optimizer state—and you are training in FP32—SGD-SaI provides a structurally guaranteed 50% memory reduction (Tables 1, 3) with speed gains of 2–6× per optimizer step versus AdamW and Adam-mini. The paper's profiling shows this is the difference between OOM and feasible training for Llama2-7B in FP32 on an A100-80GB (Table 1).

  • Prefer SGD-SaI when hyperparameter tuning budget is limited. If you can afford only a small number of tuning runs and cannot risk catastrophic performance collapse from a poor hyperparameter choice, SGD-SaI's robustness (Table 2: average accuracy 57.55% vs. 37.21% for AdamW on ViT-S/16 grid) means you are more likely to get a usable model from a random hyperparameter draw.

  • Prefer AdamW when convergence speed or early performance matters more than final performance. If your training budget is defined by wall-clock time rather than memory, and you need the best possible model within a fixed number of hours, AdamW's faster per-step convergence progress (visible in Figures 7 and 8 as steeper early loss curves) may outweigh SGD-SaI's per-step speed advantage. The paper does not report time-to-accuracy curves, but the qualitative convergence behavior suggests AdamW has an early-training advantage.

  • Prefer Prodigy when peak accuracy is paramount and memory is abundant. On ViT-S/16, Prodigy achieves 73.24% peak accuracy versus SGD-SaI's 72.92% (Table 2), but at the cost of 4× the optimizer state memory (Table 3: 0.33 GB vs. 0.08 GB) and 6.3× slower optimizer steps (78 ms vs. 12.4 ms). For settings where GPU memory is plentiful and final accuracy is the sole metric of interest, Prodigy's small peak advantage may justify its overhead.

  • Avoid SGD-SaI when training at billion-parameter scale where it has not been validated. The paper has not demonstrated that SGD-SaI produces competitive models when fully training Transformers above 125M parameters. For a production 7B-parameter pretraining run, adopting SGD-SaI is a bet on the extrapolation of small-scale results. Until follow-up work validates the method at scale, SGD-SaI should be treated as promising but unproven for large-scale pretraining—the prudent choice is to use it for memory-constrained fine-tuning (where it has positive evidence) while awaiting larger-scale pretraining results.