{"id":"a0807acb-a874-4446-b87e-e4911586d2b4","arxiv_id":"2412.13148","paper_version":3,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":5,"one_line_summary":"SWAN, a stateless optimizer combining gradient normalization and whitening, matches or beats Adam on LLaMA pretraining through 1.3B parameters with roughly half the memory and reported 2x token efficiency.","lead":"A Microsoft Research team proposes SWAN, an optimizer that pre-processes each raw gradient with row-wise normalization and whitening, removing the per-parameter momentum and variance states Adam maintains. On LLaMA-style pretraining from 60M to 1.3B parameters, they report matching or beating Adam's validation perplexity while using roughly half the memory and fewer tokens.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The 2x token-speedup claim rests on an inconsistent Adam baseline: Table 2 mixes official and reproduced numbers, the two disagree by more than SWAN's own edge at 1.3B, and no code is released to disambiguate.","rationale":"The reader's weakest assumption concerns the single-mini-batch whitening estimate used to approximate curvature. That is a real limitation, but the empirical headline can hold even if that theoretical justification is imperfect. The more immediate threat to the central claim is that the 2x token-speedup numbers are not anchored to a single, reproducible Adam baseline: Table 2 and the reproduced results in Appendix J disagree substantially at 350M and 1.3B, and the paper does not release code or checkpoints to resolve the discrepancy. This is a fairness and reproducibility concern about the main result, not a disagreement with the community consensus. The method may well work, and the existing CONDITIONAL verdict is appropriate: the claim should be verified with fully re-run baselines before acceptance-grade confidence is warranted.","tokens_in":34395,"tokens_out":10398,"duration_ms":103865,"concrete_test":"Re-run the 1.3B C4 pretraining setup for Adam, SWAN-0, SWAN-2, Galore, and Apollo-mini in a single codebase with identical token budgets and LR schedules. Compute the steps-to-perplexity ratio R(P)=S_Adam(P)/S_SWAN(P) using the official Adam curve that matches the cited 15.56 final perplexity; also report the same ratio using the reproduced Adam curve. If R at P=15.56 falls below 2x for either reference, the central '2x speedup' claim is not supported as stated.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central empirical claim, that SWAN reaches the same evaluation perplexity as Adam with roughly half the tokens, depends on which Adam curve is used as the reference. Section 6.1 says all baseline results are 'directly taken from the official numbers' in Zhao et al. (2024a) and Zhu et al. (2024), while Appendix J reports that the authors' own reproductions 'generally obtain slightly worse results for ... Adam for larger models.' Table 2 shows the concrete tension: at 350M, reproduced Adam is 19.24 vs. cited Adam 18.80; at 1.3B, reproduced Adam is 16.44 vs. cited Adam 15.56; but at 60M, reproduced Adam is 33.02 vs. cited 34.06. The paper uses official numbers in Table 2 and in the speedup figures, but Figure 4 says the training curves are reproduced, so it is unclear which Adam generates the '>2.3x' and '>2.4x' speedups. At 1.3B, SWAN-0's final perplexity (15.13) is only 0.43 lower than the cited Adam (15.56), while the reproduced Adam differs from the cited Adam by 0.88. A baseline shift of that size is larger than the method's entire reported edge over the quoted baseline, so the headline speedup magnitude is not robust to a consistent, fully re-run comparison.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes SWAN, a stateless optimizer that pre-processes the instantaneous SGD gradient with row-wise normalization (GradNorm) followed by whitening (GradWhitening), optionally accelerated by a diagonal-substitution Newton-Schulz scheme (NSDS). The authors claim that SWAN matches the memory footprint of SGD while achieving comparable or better validation perplexity than Adam on LLaMA-style pretraining, with up to 2x faster convergence in tokens on 350M and 1.3B models. The paper includes a theoretical analysis of GradNorm and GradWhitening under a simplified transformer block, plus experiments on 60M-1.3B models, ablations, and throughput measurements.","tokens_in":34812,"tokens_out":4818,"duration_ms":48910,"significance":"If the central claims hold, SWAN would be a practically significant contribution to memory-efficient LLM training: it achieves roughly 50% total memory reduction versus Adam on the measured setups, with no optimizer states for the linear layers to which it is applied. The paper is also visibly careful in several respects: it runs ablations for both components, studies the effect of learning-rate warmup, reports throughput with and without model parallelism, and attempts to connect the design to a theoretical learning-dynamics analysis. The empirical claim is nevertheless strongly tied to the choice of Adam baseline, and the paper's own reproduced Adam numbers differ from the quoted official numbers by more than SWAN's edge over the quoted baseline at 1.3B. The stateless framing is also broader than what is actually implemented, since non-linear parameters still use Adam.","major_comments":[{"comment":"The headline '2x speedup' comparison is not based on a single consistent Adam baseline. Table 2 reports two Adam rows, 'Adam' (33.02/24.44/19.24/16.44) and 'Adam (cited)' (34.06/25.08/18.80/15.56), while Appendix J states that the authors' own reproductions 'generally obtain slightly worse results for ... Adam for larger models.' At 1.3B, the reproduced Adam differs from the cited Adam by 0.88 perplexity, while SWAN-0's edge over the cited Adam is only 0.43 (15.13 vs 15.56). The speedup row '>2.4X' is therefore not robust to replacing quoted baselines with consistently re-run baselines. Please report speedups against a fully re-run, identically configured Adam baseline, or explicitly quantify the sensitivity of the speedup to the choice of Adam reference.","section":"Section 6.1, Table 2; Appendix J"},{"comment":"The claim that SWAN 'eliminates the need to store any optimizer states' is overstated as stated. Appendix J says SWAN is applied only to all linear projection weights in transformer blocks and that 'the rest of the non-linear parameters still uses Adam as the default choice.' Thus the full training run still maintains Adam optimizer states for a subset of parameters, and the end-to-end memory numbers in Figure 1(c) reflect that hybrid setup. The abstract and title should be scoped to 'stateless for linear layers' or the remaining states should be eliminated as well.","section":"Section 6.1, Appendix J; Abstract"},{"comment":"The theoretical justification for GradWhitening as a curvature-correcting update relies on Assumption 1 (H = I⊗Htilde) and on replacing the expectation E[vec(G)vec(G)^T] with the single-mini-batch sample GG^T. The paper does not provide evidence that this single-sample estimate is accurate for the Hessian/FIM blocks of real transformers; Proposition 1 is derived under the simplified transformer block and full-batch dynamics. Given that this assumption is the main link between GradWhitening and second-order optimization, I ask the authors to either add a direct empirical check of the approximation (e.g., correlation between GG^T and Hessian blocks during a real run) or explicitly weaken the claim to 'a heuristic motivated by Kronecker structure.'","section":"Section 5.3, Assumption 1; Proposition 1"},{"comment":"The NSDS acceleration is critical to the throughput claim, but it is presented as a heuristic whose selection is not documented in a falsifiable way. Appendix B.3 shows that 1 NSDS iteration makes the loss diverge, while 2 iterations work, and the step size beta=0.4 is chosen empirically. Since SWAN‡ is the variant used for the headline speedup, please report the sensitivity of the final perplexity and throughput to the NSDS iteration count and beta, or at least state that the reported numbers hold only for the specific values used.","section":"Section 4.2, Algorithm 2; Appendix B.3"}],"minor_comments":[{"comment":"The definition of s appears to be missing a square on the gradient entries: it should be s = sqrt(1/n * sum_j G_{:,j}^2), not s = sqrt(1/n * sum_j G_{:,j}).","section":"Equation (1)"},{"comment":"There is a typo in the related-work discussion: 'making the optimizer fully statless' should be 'stateless.'","section":"Section 2"},{"comment":"The memory-usage formulas for Apollo and other entries are garbled in the typeset table; please reformat them so the reader can verify the comparison.","section":"Table 1"},{"comment":"The counterfactual additive curve is defined using a set of PPL thresholds P_i, but it is not specified how these thresholds are chosen and how many are used; please clarify the exact procedure and include the number of thresholds N.","section":"Section 6.3, Equations (7)-(10)"}],"recommendation":"major_revision","confidential_remarks":"The core idea is interesting and the memory savings are real as implemented, but the main empirical claim needs a consistent baseline comparison. I would not reject if the authors can provide a fully re-run Adam baseline under identical conditions, clarify the hybrid nature of the state removal, and temper the theory where it relies on a single-sample curvature estimate. The paper would also benefit from releasing code to disambiguate the official versus reproduced numbers."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Ma et al. propose SWAN, a stateless optimizer that composes row-wise gradient normalization and gradient whitening (orthogonalization) and applies it to linear layers, with a diagonal-substitution Newton-Schulz scheme to keep the cost near Adam's. That composition, without momentum states, is new relative to Muon and SGD-Sal, and the empirical story is compelling: on LLaMA pretraining at 350M and 1.3B, SWAN reaches target perplexity in roughly half the tokens and uses about half the memory. The ablations are clean; removing either operator degrades performance. The paper is also upfront that a full theory is open and that its GradNorm/GradWhitening analysis rests on strong assumptions.\n\nThe main soft spot is the baseline bookkeeping. Table 2 mixes official quoted numbers (GaLore, Apollo) with the authors' own reproduced Adam, and those two disagree: at 350M, reproduced Adam is 19.24 vs cited 18.80; at 1.3B, 16.44 vs 15.56. The paper says it compared with official numbers because reproductions were worse for larger models. That is a favorable choice, and at 1.3B the gap between reproduced and cited Adam (0.88) is larger than the edge of the weakest SWAN variant over the cited baseline (0.43). The strongest variant (SWAN-double-dagger, 13.56) still beats a reproduced Adam, so the idea survives, but the magnitude of the speedup claim is not robust until all baselines are re-run in the same code. No code is released, which makes this harder to chase.\n\nTwo minor quibbles. First, 'stateless' is qualified: SWAN is applied only to linear layers; the non-linear parameters still use Adam. That is a small fraction of memory but not zero state. Second, the whitening estimates the curvature block from a single mini-batch (GG^T), which is a strong assumption; the paper acknowledges this but the theory is post-hoc and not load-bearing.\n\nWho is this for? People working on memory-efficient optimizers and those who care about whether Adam is irreplaceable. It deserves a serious referee; a good revision would re-run Adam and Apollo baselines under a single harness, release code, and reframe the memory claim to be exact about which parameters are optimized by what. I would not desk-reject it.","headline":"Solid idea, honest paper, but the 2x speedup is only as good as the Adam baseline you choose; needs a fully re-run, code-released revision.","tokens_in":35303,"tokens_out":5032,"would_cite":true,"duration_ms":44066,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"A stateless optimizer that pre-processes each SGD gradient with row-wise normalization and whitening claims Adam-level LLM performance at half the memory.","keywords":["stateless optimizer","gradient whitening","gradient normalization","LLM pre-training","memory-efficient training","Newton-Schulz iteration","SGD","Adam-free training"],"falsifier":"Run the 350M or 1.3B pre-training twice: once with the full SWAN pipeline and once with GradWhitening replaced by element-wise diagonal whitening, or with the whitening matrix computed from a different mini-batch than the gradient. If the 2x token-speedup over Adam persists, the claim that matrix-level whitening counteracts local curvature is not the operative mechanism.","tokens_in":34162,"feed_emoji":"⚡","tokens_out":8264,"duration_ms":70454,"temperature":0.7,"pith_summary":"SWAN is a stochastic optimizer that pre-processes each raw mini-batch gradient with two stateless matrix operations, row-wise normalization and whitening, and then applies a plain SGD update. The paper's central claim is that this suffices to replicate the benefits of Adam for pre-training large language models without storing any optimizer state. Empirically, SWAN uses the same memory as SGD, roughly half Adam's end-to-end footprint, and reaches the same validation perplexity as Adam while consuming about half as many tokens on 350M- and 1.3B-parameter models. The authors argue the two operations respectively stabilize the time-varying covariance of transformer gradients and counteract the local curvature of the loss landscape.","feed_headline":"Plain SGD plus whitening matches Adam with half the memory","feed_subtitle":"It reaches Adam-level perplexity on roughly half the tokens, with SGD-level memory.","key_machinery":"The central object is the composed operator $\\text{SWAN}(G) = \\text{GradWhitening}(\\text{GradNorm}(G))$ applied to the current gradient matrix $G \\in \\mathbb{R}^{m \\times n}$, followed by a norm-preserving rescale. GradNorm($G$) divides each row by its root-mean-square across columns, while GradWhitening($G$) = $(GG^\\top)^{-1/2} G$ is the closest orthogonal matrix to $G$, so the update uses only current-batch statistics. The fast variant replaces the exact inverse square root with a diagonal-substitution Newton\\textendash{}Schulz iteration whose per-step cost is $O(m^2)$ instead of $O(m^3)$, bringing raw throughput close to Adam's without distributed computation.","core_discovery":"On its own terms, the discovery is that the memory and convergence advantages of adaptive optimizers can be obtained by instantaneous, stateless pre-processing of SGD gradients. GradNorm standardizes each row of the gradient matrix by its root-mean-square across output dimensions, and GradWhitening multiplies the normalized gradient by $(GG^\\top)^{-1/2}$, which orthogonalizes the gradient and is derived as a non-diagonal second-order preconditioner under the assumption that the local Hessian has the Kronecker form $I \\otimes \\tilde H$. No running averages or other state variables are kept. On LLaMA-style models from 60M to 1.3B parameters trained on the C4 dataset, SWAN matches or beats Adam's validation perplexity while cutting total memory by roughly half, and the 350M and 1.3B runs reach Adam's perplexity in roughly half the tokens.","pith_inferences":["If single-batch whitening is a reliable curvature proxy, the same preprocessing could be composed with momentum or with other stateless gradient transforms; the paper explicitly leaves the composition effect open.","The quadratic analysis implies SWAN-type updates should transfer to other ill-conditioned, non-convex problems where Adam's diagonal preconditioner is a poor fit, though only toy Rastrigin tests are reported.","The additive-versus-multiplicative speedup evidence suggests SWAN's token advantage over Adam may grow with model scale, so running at 7B-13B would be the natural next check of the scaling claim.","Because GradWhitening cost depends on the smaller dimension of each weight matrix, very wide layers may need a block-wise or low-rank whitening variant, which the paper does not develop."],"forward_implications":["LLM pre-training can run with zero optimizer states, cutting total end-to-end memory by roughly 50% compared with Adam and up to about 70% when combined with per-layer training techniques.","At 350M and 1.3B scale, the token budget needed to reach a fixed evaluation perplexity is roughly halved, a about 2x speedup in steps or tokens.","The diagonal-substitution Newton\\textendash{}Schulz scheme makes SWAN's raw throughput comparable to Adam's without distributed computation, so the token-efficiency gain translates into roughly 2x effective throughput.","The same hyperparameters work across model sizes, and a no-warmup, Adam-like setting already outperforms Adam in perplexity, suggesting the optimizer can be dropped in without extensive tuning.","Ablations show that removing either GradNorm or GradWhitening degrades final perplexity, so the combination, not simply a larger effective learning rate, carries the result."],"supporting_citations":[{"why":"Defines Adam, the stateful optimizer whose memory footprint and convergence SWAN is measured against.","marker":"(Kingma & Ba, 2015)"},{"why":"Provides the LLaMA/C4 pre-training setup, baseline numbers, and the gradient-dynamics lemma used in the proof of Theorem 1.","marker":"(Zhao et al., 2024a)"},{"why":"Defines the simplified transformer block whose asymptotic dynamics underlie Theorem 1 and Proposition 1.","marker":"(Tian et al., 2023)"},{"why":"Establishes the momentum-plus-whitening baseline that isolates the contribution of GradNorm and motivates orthogonalizing gradients.","marker":"(Jordan et al., 2024)"},{"why":"Supplies the Apollo/Apollo-mini baselines and official perplexity numbers used in the comparison table.","marker":"(Zhu et al., 2024)"},{"why":"Introduces the equal-token-budget evaluation protocol that justifies quoting speedups in tokens to reach a perplexity.","marker":"(Kaddour et al., 2024)"},{"why":"Supplies the Newton\\textendash{}Schulz iteration for the matrix square-root inverse that GradWhitening and the NSDS variant accelerate.","marker":"(Huang et al., 2019)"},{"why":"Gives the condition-number lower bounds for GD and Adam that the GradWhitening convergence analysis improves upon.","marker":"(Zhang et al., 2024a)"}],"fun_headline_variants":["SGD with normalization and whitening matches Adam at half memory","Stateless SGD with whitening halves LLM training memory","SWAN: stateless SGD matches Adam with half the memory","Adam-level LLM training without optimizer states"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that a single mini-batch's gradient outer product $G G^\\top$ is a reliable proxy for the local curvature of the loss, so whitening with its inverse square root genuinely counteracts the Hessian rather than just rescaling noise; the paper specifically assumes the Hessian has the shared block-diagonal form $I \\otimes \\tilde H$.","fun_headline_variants_meta":{"raw":{"variants":["SGD with normalization and whitening matches Adam at half memory","Stateless SGD with whitening halves LLM training memory","SWAN: stateless SGD matches Adam with half the memory","Adam-level LLM training without optimizer states"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000657,"raw_usage":{"total_tokens":3038,"prompt_tokens":1007,"completion_tokens":2031,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":623,"completion_tokens_details":{"reasoning_tokens":1965}},"tokens_in":623,"tokens_out":2031,"duration_ms":13345,"temperature":1.0,"reasoning_tokens":1965,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T13:22:53.343736+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run the 350M or 1.3B pre-training twice: once with the full SWAN pipeline and once with GradWhitening replaced by element-wise diagonal whitening, or with the whitening matrix computed from a different mini-batch than the gradient. If the 2x token-speedup over Adam persists, the claim that matrix-level whitening counteracts local curvature is not the operative mechanism.","supporting_citations":[{"cited_title":"Muon: An optimizer for hidden layers in neural networks, 2024","cited_arxiv_id":null,"evidence_quote":"Establishes the momentum-plus-whitening baseline that isolates the contribution of GradNorm and motivates orthogonalizing gradients."},{"cited_title":"No train no gain: Revisiting efficient training algorithms for transformer-based language models","cited_arxiv_id":null,"evidence_quote":"Introduces the equal-token-budget evaluation protocol that justifies quoting speedups in tokens to reach a perplexity."},{"cited_title":"Iterative normalization: Beyond standardization towards efficient whitening","cited_arxiv_id":null,"evidence_quote":"Supplies the Newton\\textendash{}Schulz iteration for the matrix square-root inverse that GradWhitening and the NSDS variant accelerate."}],"review_version":1}