{"id":"a9633f3d-36f7-4384-9feb-71ed59d3db73","arxiv_id":"2608.03941","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"On Mamba-2 130M, Muon on the output projection alone gives lower validation loss and better token efficiency than Muon on the input projection or on both, across corpora and budgets.","lead":"In Mamba-2 language models, applying the Muon optimizer only to the output projection beats applying it to the input projection or to both, and mostly speeds up training rather than improving the final loss. The finding gives practitioners a concrete rule for where a spectral optimizer helps, though the gain did not show up on downstream tasks.","discovery_kind":"new_application","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Per-regime LR calibration confounds the assignment comparison: a single scalar rate for k=2 may explain why out_proj-only wins; the paper concedes this in Section X.","rationale":"The reader's weakest_assumption identified the per-regime learning-rate calibration as the most fragile point, and I agree. The paper's central claim is that the benefit of Muon is localized to the output projection. This claim is supported by the stable ranking k=3 < k=2 < k=1 < k=0, but that ranking is produced under learning rates that are not equally well tuned for each regime. The k=2 scalar-rate issue is particularly load-bearing because joint assignment is the direct alternative to 'out_proj alone'; if per-group rates allowed k=2 to win, the localization claim collapses. The paper's own Section X caveat confirms this is a live possibility, not a hypothetical. The k=1 off-minimum rate is a smaller concern, since even a 0.013 pilot improvement would not close the 0.085 gap to k=3, but it contributes to the fragility. A single focused experiment—per-group rate tuning for k=2—would settle whether the ranking is a true architectural effect or an LR artifact. I do not see another concern that is more central: the spectral diagnostics are explicitly interpretive and not needed for the primary loss claim; the 50B persistence is single-seed but is a secondary robustness claim; and the downstream parity is honestly reported. Therefore the reader's CONDITIONAL verdict remains appropriate, and I recommend no change. The paper is honest about the caveat, but the abstract's unqualified 'benefit is localized' statement should be tempered until the per-group rate check is run.","tokens_in":19904,"tokens_out":5887,"duration_ms":58861,"concrete_test":"Re-run k=2 with separate Muon learning rates for G_in and G_out, selected from a coarse grid (e.g., η_in, η_out ∈ {5e-3, 1e-2, 2e-2}) on the 2e8-token OWT pilot with 4 seeds; take the best (η_in, η_out) pair and compare its final validation loss at 1e9 tokens on OWT (seeds 42 and 1337) against k=3's reported 3.3019. If the best per-group k=2 still fails to beat k=3, the scalar-rate confound is not the explanation. If it matches or beats 3.3019, the localization claim fails and the paper's ranking is an artifact of the single-rate choice.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim that Muon on G_out alone (k=3) beats both G_in-only (k=1) and joint (k=2) rests on per-regime learning rates chosen from a two-seed pilot grid. The weakest link is k=2: a single Muon rate η=2e-2 is applied to both G_in and G_out, despite the paper's own evidence that these are architecturally heterogeneous (Section III, Appendix B-C). Section X explicitly concedes 'we cannot fully exclude that k=2's under-performance reflects a single scalar rate serving heterogeneous sub-blocks.' If separate per-group rates were tuned, k=2 might match or beat k=3, which would directly falsify 'out_proj alone beats both.' The concern is compounded by the pilot's low seed count and by k=1's rate being deliberately set off its plateau minimum (Appendix A), so the entire per-regime LR table is fragile. Thus the comparison is not a clean test of assignment alone; it confounds assignment with per-regime learning-rate tuning.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper reports a controlled comparison of Muon and AdamW for training Mamba-2 130M, varying only which of the two matrix-valued projection groups (input projection G_in, output projection G_out, or both) are trained with Muon. Four regimes are compared: k=0 (all AdamW), k=1 (Muon on G_in), k=2 (Muon on G_in ∪ G_out), and k=3 (Muon on G_out). Learning rates for AdamW and for each Muon regime are selected from two-seed pilot sweeps at 2e8 tokens, then frozen for full runs at 1e9 and 2.6e9 tokens on OpenWebText and FineWeb-Edu. The main claims are: (i) k=3 achieves the lowest validation loss in all corpus–budget cells (e.g., ΔL = -0.116 vs AdamW at 1B tokens on OpenWebText); (ii) the ordering k=3 < k=2 < k=1 < k=0 is stable across corpora and budgets at the endpoints; (iii) the benefit is mainly token-efficiency, with k=3 matching AdamW's final 1B-token loss at ~5.8e8 tokens; (iv) spectral diagnostics show Muon lowers the condition number of whichever projection it is applied to, but the better-conditioned G_in does not yield the largest loss improvement; and (v) a single-seed run to 5e10 tokens on FineWeb-Edu shows the k=3 vs k=0 gap persists (2.8853 vs 2.9027). The paper is explicitly framed as empirical and includes several negative results (downstream evaluation, pilot-scale ablation reversal).","tokens_in":20234,"tokens_out":5155,"duration_ms":55646,"significance":"If the results are correct, this is a useful empirical data point for optimizer placement in state-space models: the Muon benefit is not uniform across Mamba-2's projection matrices, and the output projection is the most effective target. The paper's protocol is disciplined in several respects: learning rates were frozen before full runs, the AdamW baseline is an external reference, paired seeds are used, a fixed checkpoint grid is specified, and the authors report negative results (no downstream gain, pilot-scale artifacts) rather than only favorable outcomes. The spectral diagnostics, while interpretive, provide a concrete metric-based attempt to explain the effect. However, the study is narrow (130M-parameter model, two shared seeds, no significance tests, single-seed 50B run) and the central localization claim is partly confounded by per-regime learning-rate tuning. The contribution is therefore a plausible and honestly reported empirical finding rather than a definitive or theoretically grounded result.","major_comments":[{"comment":"The paper states that the controlled experimental factor is the Muon assignment and that learning rates are 'held fixed across regimes' (Section IV), but Table VI shows a different Muon learning rate for each regime (η=1e-2 for k=1, 2e-2 for k=2, 1.5e-2 for k=3). The comparison k=3 vs k=2 vs k=1 is therefore a joint comparison of assignment and per-regime learning rate, not an assignment-only comparison. This is load-bearing for the central claim that 'out_proj alone beats both': the abstract and Section X concede for k=2 that a single scalar rate serving heterogeneous sub-blocks could explain its underperformance, but the same caveat applies to the entire ranking. I recommend either (a) adding sensitivity runs that evaluate each assignment at a common Muon learning rate (e.g., η=1.5e-2 for all k), or (b) explicitly reframing the abstract and conclusions to say the result is established","section":"Section IV / Table VI"},{"comment":"The pilot grid uses only two seeds, and the k=1 learning rate was deliberately set at 1.0e-2 instead of the plateau minimum 1.5e-2 because the plateau minimum showed higher seed variance (Appendix A). The difference between these two settings (ΔL=0.013) is smaller than the seed standard deviation at η=1.5e-2 (0.028), so the choice is defensible, but it means the reported k=1 loss is not the best this regime can achieve under the pilot's own criterion. More generally, the rank ordering of the intermediate regimes (k=1 vs k=2) is not stable: at 2.6B tokens the two coincide on OpenWebText and reverse on FineWeb-Edu (Table III). Since the paper's headline claim is the localization to G_out, at minimum the authors should report the full pilot table (all grid points, both seeds) and perform a simple paired significance check (e.g., a paired permutation test across seeds) for the k=3 vs k=2 dif","section":"Appendix A / Section X"},{"comment":"The 5e10-token run is single-seed, and no error bars or other uncertainty estimates are reported for the final gap of 0.017 (2.8853 vs 2.9027). At 130M scale, single-seed loss differences of this magnitude can be within run-to-run noise; the two-seed standard deviations in Table II (0.001–0.010) do not transfer to the much longer run because the LR schedule and data order differ. Additionally, this run only compares k=0 and k=3, so it does not test the localization claim at scale (k=2 and k=1 are absent). I suggest either adding a second seed for the 50B runs or, if that is computationally prohibitive, framing Section VIII explicitly as a single-seed exploratory observation rather than as confirmation of the localization claim.","section":"Section VIII"}],"minor_comments":[{"comment":"The reference AdamW learning rate is given as α0=3e-3 in Section III, but the actual selected rate is α*=3.6e-3 (Appendix A). Clarify the relationship between the reference and the calibrated value to avoid confusion.","section":"Section III / Appendix A"},{"comment":"The bottom row of Table III is labeled 'ΔL(k=0−k=3)' and reports positive numbers (+0.116, +0.063, ...), whereas Table II defines ΔL relative to k=0 and reports negative values for improvements. Use a consistent sign convention, e.g., report ΔL = L_k − L_0 for all rows.","section":"Table III / Section VI-C"},{"comment":"The text 'splitG in andMuon onG out' appears to be missing a space; also consider defining 'split-Muon' formally before first use. Minor presentation issue.","section":"Appendix D"}],"recommendation":"major_revision","confidential_remarks":"The paper is honest and well-structured, with several pre-registration-like precautions (frozen LRs, fixed checkpoint grids, reported negative results). The central issue is the per-regime LR calibration confound, which the authors themselves acknowledge in Section X. If they can add a robustness check with a common Muon LR across regimes (or clearly re-scope the claim), the paper would be acceptable. The 50B single-seed result should be treated as exploratory. No concerns about novelty disclosure or citation practice."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Short version: the paper finds something real and reports it honestly. Under a controlled protocol they vary which Mamba-2 projection groups get Muon: input, output, both, none. On two corpora and two token budgets, out_proj-only (k=3) is best: final validation loss 3.3019 vs 3.4178 for AdamW at 1B tokens on OpenWebText, and it matches AdamW's final loss in about 5.8e8 tokens. The ranking k=3 k=2) is plausible but not as cleanly established.\n\nWho it's for: anyone choosing pre-training optimizers for Mamba-family models, and optimizer researchers tracking where Muon works. It deserves a serious referee; a good reviewer should push for per-group LR sweeps and more seeds on the k=2 vs k=3 comparison before relying on the localization, but the paper should not be desk-rejected.","headline":"A well-run empirical study with a genuinely new localization result—Muon on Mamba-2's output projection helps most—though the comparison is partly confounded by per-regime learning-rate choices and the downstream payoff is flat.","tokens_in":20635,"tokens_out":2023,"would_cite":true,"duration_ms":20327,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Muon's benefit in Mamba-2 is localized to the output projection; better input conditioning does not explain it.","keywords":["Muon","AdamW","Mamba-2","state space models","spectral optimization","Newton-Schulz","token efficiency","optimizer localization"],"falsifier":"Rerun the four regimes on OpenWebText at 1e9 tokens with per-group learning rates for k=2 (separate rates for G_in and G_out) and a denser k=1 sweep around 1.5e-2 using more than two seeds; if k=2 or k=1 then ties or beats k=3 on final validation loss, the output-projection localization claim collapses.","tokens_in":19827,"feed_emoji":"🎯","tokens_out":9192,"duration_ms":82055,"temperature":0.7,"pith_summary":"This paper asks where, if anywhere, the Muon optimizer helps inside a Mamba-2 language model. By varying only which projection matrices receive Muon, it finds the benefit is concentrated in the output projection: Muon on out_proj alone beats Muon on the input projection or on both, and beats the pure AdamW baseline. The advantage is mainly token efficiency — the winning regime matches AdamW's final validation loss at about 58% of the tokens — and the ordering holds on two corpora and two token budgets, persisting well past the compute-optimal point. The paper also closes off the obvious explanation: Muon lowers the condition number of whichever matrix it trains, so the better-conditioned input projection is not the source of the gain. A careful reader would come away with a concrete recipe (apply Muon to the output projection, not the whole block) and a warning that pre-training loss gains at this scale did not carry to zero-shot downstream accuracy.","feed_headline":"Muon on the output projection alone beats Muon on input or both","feed_subtitle":"Muon on the output projection only reaches AdamW's final 1B-token loss in ~0.58B tokens.","key_machinery":"The controlled experimental factor is the assignment of Muon to the two two-dimensional matrix groups in each Mamba-2 block: in_proj (G_in) and out_proj (G_out), giving four regimes k=0..3. Muon's update forms matrix-valued momentum with a Nesterov step, then applies a five-iteration Newton-Schulz polynomial that orthogonalizes the update direction, i.e. steepest descent under the spectral norm; AdamW handles all other parameters. The comparison measures final validation loss, validation-equivalent tokens, and spectral diagnostics (condition number, effective rank, spectral norm) of the two projection groups.","core_discovery":"Muon's benefit in a Mamba-2 130M is localized to the output projection. Assigning Muon only to out_proj (k=3) yields the lowest final validation loss of four regimes: ΔL = −0.116 vs AdamW at 10^9 tokens on OpenWebText (3.3019 vs 3.4178), beating input-only (k=1) and joint (k=2). The ranking k=3 < k=2 < k=1 < k=0 is stable across two corpora and two budgets, and at 5×10^10 tokens k=3 still leads AdamW by 0.017. The gain is token efficiency: k=3 reaches AdamW's final 10^9-token loss in ≈5.8×10^8 tokens. Conditioning does not explain this: Muon conditions whichever projection it trains, yet the best-conditioned input projection does not lower loss, while the best-conditioned output projection d","pith_inferences":["If the proposed mechanism is right — Muon helps a uniform dense projection that writes to the residual stream and hurts a merged row-concatenation of heterogeneous sub-projections — then other architectures with fused QKV or gated projections should show the same pattern; this is testable by repeating the protocol on a Transformer with grouped projection matrices.","The k=2 result may be an artifact of a single learning rate applied across heterogeneous groups; the paper's own split-Muon experiment shows per-slice treatment changes the picture, so a per-group-rate k=2 might close the gap to k=3.","The token-efficiency saving suggests a practical rule for compute-constrained SSM training: spend Muon only where the matrix writes to the residual stream, and validate with a matched-loss token count rather than final loss alone.","Because the 130M zero-shot numbers are near chance, the downstream decoupling should not be taken as evidence that the loss gain is useless at scale; a larger-model replication could plausibly show both better loss and better tasks."],"forward_implications":["For Mamba-2 training, the default should be Muon on out_proj alone; applying Muon to in_proj or to all eligible matrices is worse than the single-group assignment in every corpus-budget cell tested.","At 10^9 tokens, k=3 reaches AdamW's final validation loss in roughly 5.8×10^8 tokens, so localized Muon buys about 40% of the token budget at matched loss.","The k=3 advantage persists when training runs to 5×10^10 tokens (~19× Chinchilla), so it is not an artifact of under-training, though the gap narrows as the budget grows.","Spectral conditioning is not a reliable guide: Muon improves the condition number of whichever projection it trains, but input-projection conditioning does not improve loss, so optimizer gains should be evaluated on loss curves, not spectra.","At 130M, the pre-training loss gap does not transfer to zero-shot tasks (mean accuracy 41.7 vs 42.0), so localized optimizer gains need downstream validation."],"supporting_citations":[{"why":"Defines the Mamba-2 architecture, SSD core, and the 130M reference training configuration that the study trains.","marker":"[6]"},{"why":"Defines Muon's Newton-Schulz orthogonalization and the eligibility rules that restrict it to two-dimensional hidden-layer matrices.","marker":"[2]"},{"why":"Supplies the public Muon reference implementation used for the runs, including Nesterov momentum and quintic coefficients.","marker":"[3]"},{"why":"Supplies the official Mamba implementation from which the parameter groups G_in and G_out and the eligible-matrix sets are taken.","marker":"[1]"},{"why":"Defines AdamW, the baseline optimizer for all non-Muon parameters, with its decoupled weight decay.","marker":"[10]"},{"why":"Reports Muon's Transformer scaling behavior and the shape-dependent update scaling the paper adopts.","marker":"[16]"},{"why":"Provides the optimizer-comparison template of fixed token budgets, explicit learning-rate selection, and a strong baseline that the protocol follows.","marker":"[13]"},{"why":"Establishes the Mamba family and the downstream evaluation protocol used to test whether the loss gain transfers to tasks.","marker":"[5]"}],"fun_headline_variants":["Muon in Mamba-2 works best on output projection alone","Output-only Muon beats joint Muon in Mamba-2 130M","Muon on output proj only speeds Mamba-2 to same loss","Mamba-2: Muon's win is localized — output projection","Token-efficient Muon: output-only in Mamba-2 wins"],"cache_read_input_tokens":2688,"weakest_assumption_plain":"The ranking of the four regimes rests on learning rates picked from a two-seed 2e8-token pilot grid, with k=1 set to 1.0e-2 instead of the plateau minimum 1.5e-2 and k=2 using a single scalar rate for two heterogeneous groups; re-calibration could change which regime wins.","fun_headline_variants_meta":{"raw":{"variants":["Muon in Mamba-2 works best on output projection alone","Output-only Muon beats joint Muon in Mamba-2 130M","Muon on output proj only speeds Mamba-2 to same loss","Mamba-2: Muon's win is localized — output projection","Token-efficient Muon: output-only in Mamba-2 wins"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000488,"raw_usage":{"total_tokens":2243,"prompt_tokens":746,"completion_tokens":1497,"prompt_tokens_details":{"cached_tokens":256},"prompt_cache_hit_tokens":256,"prompt_cache_miss_tokens":490,"completion_tokens_details":{"reasoning_tokens":1400}},"tokens_in":490,"tokens_out":1497,"duration_ms":13017,"temperature":1.0,"reasoning_tokens":1400,"cache_read_input_tokens":256,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-05T05:23:01.078889+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Rerun the four regimes on OpenWebText at 1e9 tokens with per-group learning rates for k=2 (separate rates for G_in and G_out) and a denser k=1 sweep around 1.5e-2 using more than two seeds; if k=2 or k=1 then ties or beats k=3 on final validation loss, the output-projection localization claim collapses.","supporting_citations":[{"cited_title":"Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality,","cited_arxiv_id":null,"evidence_quote":"Defines the Mamba-2 architecture, SSD core, and the 130M reference training configuration that the study trains."},{"cited_title":"Muon: An optimizer for hidden layers in neural networks,","cited_arxiv_id":null,"evidence_quote":"Defines Muon's Newton-Schulz orthogonalization and the eligibility rules that restrict it to two-dimensional hidden-layer matrices."},{"cited_title":"Kellerjordan/Muon,","cited_arxiv_id":null,"evidence_quote":"Supplies the public Muon reference implementation used for the runs, including Nesterov momentum and quintic coefficients."},{"cited_title":"state-spaces/mamba,","cited_arxiv_id":null,"evidence_quote":"Supplies the official Mamba implementation from which the parameter groups G_in and G_out and the eligible-matrix sets are taken."}],"review_version":1}