Do We Need Adam? Surprisingly Strong and Sparse Reinforcement Learning with SGD in LLMs

Summary

The paper asks a timely question: does the RLVR post-training stage of large language models actually need AdamW, or does plain SGD suffice? The conceptual move is to argue that RL's optimization landscape differs from SFT's in ways that undercut AdamW's two key ingredients — per-parameter adaptivity and momentum — and then to test that argument with a matched optimizer ablation. Two diagnostics motivate the hypothesis: AdamW's $\sqrt{v}$ distribution is roughly $22\times$ narrower in RLVR than in SFT (Figure 2), and the cosine similarity between accumulated momentum and the current gradient collapses from $0.997$ (SFT) to $-0.007$ (RL). Building on this, the authors run a matched grid — AdamW, SGD, SGD+Momentum, RMSProp — across three domains (math, code, RLVE), three base models (Qwen3-1.7B/8B, Llama-3.1-8B), and two RL algorithms (GRPO, PPO), plus an LR sweep on Qwen3-8B and an RLVE extended-training study. Two headlines emerge: SGD matches or slightly outperforms AdamW on aggregate, and SGD's updates are extraordinarily sparse — often 99.9%+ of parameters unchanged — with substantially lower effective rank than AdamW's. This inverts a widely held prior about SGD in transformers, delivers immediate memory savings (about $15.7$ GB on Qwen3-1.7B, of which $13.6$ GB is the analytic optimizer-state saving that I verified independently), and suggests that the effective RLVR problem lives in a very low-dimensional subspace.

The contribution is valuable and, if it holds, has direct practical implications for scaling RLVR. The main limitations are interpretive rather than methodological: several claims outrun what the tables actually show (e.g., "all cases," "fewer than 0.02%," sparsity generalizing to PPO), the LR-sweep peak difference ($0.638$ vs $0.644$) is presented as SGD "outperforming" without any variance information, and the sparsity story rests on a $10^{-5}$ threshold whose relation to bfloat16 precision and to SGD's larger nominal step size is not justified. The two motivating diagnostics ($22\times$ $\sqrt{v}$ ratio, $0.997$ vs $-0.007$ cosines) are single step-50 snapshots on one model whose stability across steps and model families is not shown, and part of the $22\times$ ratio is arithmetically inherited from RL's ~$10\times$ smaller gradient magnitudes rather than reflecting reduced per-parameter heterogeneity. Multiple broken figure references and a scrambled Table 4 make sparsity claims harder to audit than they should be. Finally, an executed replication attempt confirmed that the linked repository is essentially empty (a single 1-byte README.md across all git history) and that the paper's independent variable — the optimizer-selection branch — does not exist in either the authors' artifact or the underlying verl framework; so verifying anything beyond the analytic memory arithmetic requires reconstructing configs, seeds, and analysis scripts from prose.

Strengths

• Timely and useful question. Whether AdamW is truly needed for RLVR is a natural question given recent evidence that RL updates are sparse and geometrically constrained; the paper frames it cleanly and tests it directly.
• Diagnostic-driven hypothesis. The $\sqrt{v}$-standard-deviation comparison between SFT and RLVR ($5.11\times 10^{-6}$ vs $2.29\times 10^{-7}$) and the momentum-vs-gradient cosine analysis ($0.997$ vs $-0.007$) give a concrete mechanistic motivation for the ablations that follow, rather than treating the negative result about AdamW as purely empirical.
• Matched ablation grid. The AdamW / SGD / SGD+Momentum / RMSProp comparison is the right factorial design for disentangling momentum from adaptive learning rate, and the study spans three domains, three base models, and two RL algorithms — a broader sweep than most optimizer papers in this space.
• Robustness to LR tuning. Appendix F sweeps AdamW and SGD over decade grids and shows that SGD remains at least competitive at every reasonable LR; this defuses the natural "you just under-tuned AdamW" objection at coarse resolution.
• Practical payoff, checked. The memory analysis ($2 p d_{\text{optim}}$ theoretical savings, $\sim 13.6$ GB analytical and $\sim 15.7$ GB measured for Qwen3-1.7B) is a clean, useful practitioner-facing takeaway. The analytic half reproduces essentially exactly: measured on the true Qwen3-1.7B parameter count ($p{=}1{,}720{,}574{,}976$), AdamW's 12.00 bytes/param vs SGD's 4.00 gives 13.76 GB of optimizer-state saving, within 1.2% of the paper's rounded value.
• Extended training with RLVE. Using RLVE for 500 steps as an extended-duration stress test is a good design choice — it addresses the natural concern that SGD's competitiveness might collapse over longer horizons.

Weaknesses

• Overstated universal claims. Several summary statements outrun the tables. "SGD consistently outperforms AdamW across all cases" at 3K holds at the per-model mean but is false at the individual-benchmark level (e.g., Qwen3-8B/3K: Math-500 87.4 < 89.2, AMC 57.8 < 60.2, Olympiad 48.3 < 51.1). "Neither momentum nor adaptive learning rates improves performance" (Finding 2) contradicts the paper's own prose noting that RMSProp posts the top mean in several 8K rows. Similarly, the conclusion's blanket "fewer than 0.02%" ignores the paper's own 0.02%–0.46% range in §1.
• Numerical inconsistency between abstract and body. The abstract quotes a "more than $1{,}000\times$" sparsity gap; §1 quotes "nearly $500\times$." Both cannot describe the same headline comparison, and neither the range nor the referent (best cell? average? median?) is disambiguated.
• PPO sparsity is claimed but not tabulated. Finding 4 asserts that SGD's substantial sparsity "generalizes to PPO," but Table 5 contains only benchmark accuracies. The 99.92%/99.98% actor and 61.1%/8% critic numbers appear only in the running text and cannot be cross-checked against a table.
• Sparsity threshold is load-bearing but not stress-tested. The $10^{-5}$ threshold on $\theta_1-\theta_0$ is used to declare 99.99% sparsity, yet its justification ("PyTorch uses $10^{-5}$ as the default tolerance") does not match PyTorch's allclose (its default absolute tolerance is $10^{-8}$; $10^{-5}$ is its relative tolerance). Bfloat16's ULP near typical weight magnitudes is close to this threshold, and SGD's per-parameter step is $\sim 10^{5}$–$10^{6}$ smaller than AdamW's effective per-parameter step at the paper's chosen LRs — precisely where the boundary would separate them. A sensitivity plot at $\{10^{-4}, 10^{-5}, 10^{-6}\}$ would defuse the concern that the SGD/AdamW gap is largely a rounding-boundary artifact.
• LR-sweep "outperforms" claim rests on a $0.006$ gap with no variance. SGD's peak of $0.644$ (LR $1.0$) vs AdamW's $0.638$ (LR $10^{-5}$) is drawn from a single run per LR, and the AdamW grid $\{10^{-7},10^{-6},10^{-5},10^{-4}\}$ leaves the plausible interval $(10^{-5}, 10^{-4})$ untested. "Matches" is well supported by the sweep; "slightly outperforms" is not.
• Motivating diagnostics are single-snapshot on one model. The $22\times$ $\sqrt{v}$ ratio and the $0.997$ vs $-0.007$ cosines are computed at a single training step (50) on Qwen3-1.7B, yet they carry the entire mechanistic motivation for Hypothesis 1. No evidence is offered that they are stable across steps, model families, or domains.
• The $22\times$ ratio may partially reflect gradient-magnitude shift, not per-parameter heterogeneity. Appendix D's own numbers show that RL gradient magnitudes are already ~$12\times$ smaller ($|g|$ mean $3.85\times 10^{-7}$ vs $4.61\times 10^{-6}$) and gradient standard deviations ~$17\times$ smaller. Since $\sqrt{v}$ is a running estimate of gradient magnitude, its standard deviation naturally shrinks with the mean. The load-bearing question is whether the relative heterogeneity (e.g., $\sigma(\sqrt{v})/\mu(\sqrt{v})$) is actually reduced in RL, not just the absolute spread.
• Missing in-house SFT baseline for the central contrast. The paper's headline conceptual claim — SGD works in RL but not in SFT — is only half-tested. AdamW and SGD are run head-to-head in RLVR, but the SFT-side of the contrast is imported from prior work at different scales and setups. A matched SGD-SFT run on Qwen3-1.7B under the same batch size and LR sweep would either strengthen the story dramatically or reveal that the SGD/AdamW gap is smaller than the literature suggests once one uses a modern base model.
• Momentum interpretation. The near-zero cosine ($-0.007$) between accumulated momentum and current gradient in RL supports "momentum is uninformative on average," not the stronger "counter-productive" framing suggested by the sub-section heading; a small clarification would prevent over-reading.
• Sparsity claim attribution. The RLVE "99.8% vs 86.7%" sparsity comparison is cited alongside Table 6, but Table 6 shows only accuracies; the sparsity numbers appear only in prose and should either be tabulated or attributed to the RLVE column of Table 4 (which does contain corresponding values).
• Broken cross-references. At least three "Figure ??" placeholders remain in §5 and §6 — including the one supporting the substantive claim that "sparsity decays much slower in SGD" — leaving key claims formally unsupported in-text.
• Table 4 is hard to parse. The 'Math' header block is duplicated, and the Sparsity/Rank rows are collapsed into a single undifferentiated numeric sequence in which ranks ($\sim 26$) and sparsity percentages ($\sim 90$) are hard to separate visually. Since Table 4 carries the main sparsity/rank story, this presentation issue directly hampers verification.
• Momentum-scaling asymmetry in Table 1. Table 1 defines SGD+Momentum with the standard $m_t = \mu m_{t-1} + g_t$ (no $(1-\mu)$ factor), while AdamW's first moment is $\beta_1 m_{t-1} + (1-\beta_1) g_t$. Later prose treats SGD+Momentum as "AdamW without adaptive LR," which is only true up to a scaling of $1/(1-\beta_1) \approx 10$; the footnote about bias correction/weight decay does not cover this.
• "430M parameters on rank 0" is ambiguous. Appendix D reports capturing optimizer state for "all 430M trainable parameters on rank 0" while the model is Qwen3-1.7B; this reads more like a shard than the full parameter vector, and it matters because the headline $22\times$ $\sigma$ ratio depends on the aggregation set.
• Aggregate wins can mask per-benchmark losses. Under PPO on Qwen3-8B, SGD's higher mean (56.7 vs 54.0) is driven entirely by GPQA (+14.7); on the three math benchmarks SGD loses to AdamW. This should at least be acknowledged, especially since GPQA is arguably out of the RL training distribution.

Reproducibility & code

An executed replication attempt was made. Only one claim produced any evidence; the rest were blocked by an environmental fault (the compute node's CUDA driver could not initialize node-wide, so no GPU training ran). The environmental failure is not the paper's fault, but two paper-side findings emerged from the attempt that materially sharpen the reproducibility picture.

• Released repository is empty — confirmed executably. The linked GitHub repo (SagnikMukherjee/sgd_adam_rlvr) contains a single 1-byte README.md across all of git history — no training entrypoints, no verl configs, no sparsity/effective-rank analysis scripts, no evaluation harness, and no checkpoints. Because verl 0.4.1 hardcodes optim.AdamW in fsdp_workers.py::_build_model_optimizer and exposes no optimizer-name key in ppo_trainer.yaml, the SGD / SGD+Mom / RMSProp / AdamW branch that carries the paper's independent variable does not exist in either the authors' artifact or the underlying framework; a reproducer must invent it from prose. Every headline number in Tables 2–7 and Figures 1–6 therefore rests only on the authors' internal, unreleased code.
• Analytic memory saving reproduces to within 1.2%. Measured on the true Qwen3-1.7B parameter set ($p{=}1{,}720{,}574{,}976$), AdamW allocates 12.00 bytes/param vs SGD's 4.00, giving 13.76 GB of optimizer-state saving — within 1.2% of the paper's rounded 13.6 GB. This part of the memory claim is device-independent and confirmed. The additional 2.1 GB attributed to "reduced communication buffer overhead incurred by FSDP" that yields the 15.7 GB headline was not attempted (requires an FSDP GPU run) and its measurement protocol (tool, FSDP sharding strategy, activation-checkpointing state, batch/sequence configuration) is not documented.
• Nothing else was checked numerically. GRPO math (Table 2), coding (Table 3), sparsity/rank (Table 4), PPO (Table 5), RLVE extended training (Table 6), the LR sweep (Table 7), and Figures 1–6 were all blocked upstream by the same GPU fault. The paper's numbers are therefore neither supported nor refuted by this run; a re-attempt on healthy hardware is needed. Combined with the empty repository, this means the current audit trail for Tables 2–7 is the paper's prose alone.
• Seeds and repetitions. No random seeds and no per-cell repetition counts are reported for Tables 2, 3, 5, 6, or 7. Because much of the argument turns on small gaps ($\Delta$ mean of $0.2$–$2.6$ points at 3K, $0.006$ in the LR sweep), a reproducer cannot distinguish signal from single-seed noise.
• Sparsity/effective-rank measurement. Metric definitions are precise, but the tensor set (embeddings? LM head? RMSNorm? bias tensors?), the training step at which the measurement is taken (end of training implied by Figure 4, not stated), and the aggregation for effective rank (arithmetic mean vs size-weighted mean over per-tensor ranks) are not documented. Reconstructing the analysis from the paper's definitions was feasible during the replication attempt, but all three choices had to be judgment calls and can move the headline 99.99%/26 numbers materially.
• PPO setup. The critic (value network) architecture, initialization, PPO step count, GAE $\lambda$, value-loss weight, and clip range are not stated. These matter especially for the 61.1% vs 8% critic-sparsity comparison, which is unusually sensitive to whether the value head is randomly initialized or LM-initialized and to the number of PPO steps.
• Evaluation harness. Decoding temperature/top-p is specified for coding (temperature 0.2, 10 samples) but not for the math benchmarks (MATH-500, AIME-24/25, AMC, Olympiad, GPQA). Qwen "thinking mode" is enabled without a stated budget. Both are known to move AIME/MATH scores by multiple points.
• Step-50 SFT/RL matching. The $\sqrt{v}$ and momentum-cosine analyses require matching two independent training runs at exactly step 50; the SFT dataset size after filtering to test-score $\geq 0.9$, sequence packing/length, and shuffle order are not documented. Given that the $22\times$ ratio and $0.997$ vs $-0.007$ cosines are the paper's motivating diagnostics, the matching protocol deserves more detail.
• RLVE task scheduling is delegated by reference to Zeng et al. (2025); the specific 260-out-of-400 task subset and difficulty-evolution parameters are not reproduced here.

The qualitative headlines — SGD is competitive with AdamW in RLVR, and its updates are dramatically sparser — are internally coherent, but every numeric table in the paper is currently unaccompanied by the code needed to regenerate it.

Recommended Changes

Essential.

• Populate the released repository. Ship, at minimum, the verl patch or subclass that adds the optimizer-selection branch (SGD / SGD+Momentum / RMSProp / AdamW), the GRPO and PPO configs for each model×domain (matching Tables 2, 3, 5, 6), the SGD/AdamW LR-sweep scripts backing Table 7 and Figure 5, the sparsity/effective-rank analysis notebook (with the exact tensor set, capture step, and aggregation), and the SFT-vs-RL step-50 pipeline used for Figure 2 and the momentum-cosine numbers. This directly resolves the empty-repo issue confirmed in Reproducibility & code, without which the paper's independent variable is not reproducible.
• Add seeds and per-seed variance. Report at least 2–3 seeds per cell (or bootstrap error bars) for Tables 2, 3, 5, 6, and the LR sweep in Table 7, so the reader can distinguish the SGD-vs-AdamW signal from single-seed noise — especially for the AIME-24/25 30-item benchmarks and the $0.638$ vs $0.644$ LR-sweep gap flagged in Weaknesses.
• Reconcile the "$1{,}000\times$" vs "$500\times$" sparsity gap and the "fewer than 0.02%" claim. In the abstract and conclusion, replace the blanket lower-bound language with the paper's own 0.02%–0.46% range and a matching factor range (or state clearly that these are best-case values), addressing the overstated-claims and abstract/body inconsistency Weaknesses.
• Fix overstated universal claims. Reword "SGD consistently outperforms AdamW across all cases" (§4.1) to specify the aggregation level (per-model mean), restore the "consistently" hedge to Finding 2, and either add PPO sparsity numbers to Table 5 or restrict Finding 4's sparsity clause to the GRPO/RLVE runs where sparsity is actually tabulated.
• Add a sparsity-threshold sensitivity analysis. Report the SGD-vs-AdamW sparsity gap at $10^{-4}$, $10^{-5}$, and $10^{-6}$ thresholds and correct the "PyTorch default $10^{-5}$" justification, addressing the load-bearing-threshold weakness. Discuss the interaction between the absolute threshold, bfloat16 precision, and SGD's much smaller per-parameter step at the chosen LR.
• Fix Table 4 layout and broken figure references. Regenerate Table 4 with distinct Sparsity and Rank rows and unambiguous column groupings, and resolve the "Figure ??" placeholders in §5 (SGD-vs-AdamW sparsity trajectory, layerwise sparsity) and §6 (sparsity-decay claim).
• Test the mechanistic diagnostics for robustness. Add a step-trajectory (e.g., steps 10/50/100/200) of the $\sqrt{v}$-std and momentum-cosine numbers, ideally for at least one additional model or domain, and report a CV-normalized version of the $\sqrt{v}$ statistic ($\sigma/\mu$) alongside the raw $\sigma$, so that the $22\times$ ratio is not conflated with RL's overall smaller gradient magnitudes.

Suggested.

• Tighten the AdamW LR grid near its optimum. Add points between $10^{-5}$ and $10^{-4}$ (e.g., $3\times 10^{-5}$, $5\times 10^{-5}$) so that "SGD slightly outperforms tuned AdamW" is not resting on a decade-scale grid.
• Add a matched SGD-SFT baseline. Run SGD (and AdamW) on the same Qwen3-1.7B under an SFT regime with the paper's batch/sequence/LR setup, so the central RL-vs-SFT contrast is tested in-house rather than imported from prior work at different scales.
• Document the sparsity measurement protocol. State the tensor set included, the step at which $\theta_1$ is captured, and the aggregation method for effective rank (arithmetic mean over 2D tensors, or size-weighted).
• Document the PPO setup. Add critic architecture, initialization, step count, GAE $\lambda$, value-loss weight, and clip range to Appendix A.1, and confirm whether critic sparsity is computed at the same threshold and step as the actor.
• Document evaluation decoding for math. State temperature, top-p, sample count, and Qwen thinking-budget per benchmark; without this, math cells in Tables 2, 5, and 6 cannot be reproduced.
• Document peak-memory measurement protocol. State the measurement tool (torch.cuda.max_memory_allocated, NVIDIA-SMI, or FSDP-aware profiler), the FSDP sharding strategy, activation-checkpointing state, and the batch/sequence configuration used to obtain the 15.7 GB headline.
• Clarify the "430M parameters on rank 0" statement in Appendix D/E — say whether $\sigma_{\text{SFT}}=5.11\times 10^{-6}$ vs $\sigma_{\text{RL}}=2.29\times 10^{-7}$ and the $0.997$ vs $-0.007$ cosines are computed over a single FSDP shard or the full parameter vector, and provide a per-layer breakdown so the $0.997$ SFT cosine is not driven by a small subset of parameters.
• Soften "counter-productive" to "uninformative on average" in the momentum sub-heading and around the $-0.007$ cosine, so the framing matches what a near-zero cosine can support.
• Add a per-benchmark note for PPO on Qwen3-8B, acknowledging that SGD's higher mean is driven by the GPQA jump while SGD trails on the in-domain math benchmarks.
• Fix minor wording. "RL suffers less from reduced catastrophic forgetting" reads as a double-negation; "without losing accuracy" for Qwen3-1.7B should be softened given the 1.4-point 8K math gap; and normalize the SGD+Momentum row in Table 1 (or add a note) so its equivalence to "AdamW without adaptive LR" is correct up to the $(1-\beta_1)$ scaling.