Pith. sign in

REVIEW 3 major objections 4 minor 38 references

Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size

T0 review · 3 major / 4 minor · reviewed 2026-08-15 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read As vocabulary grows with width, the optimal embedding-to-hidden learning-rate ratio shifts from Θ(d) to roughly √d, and the paper's Large Vocabulary Parametrization built on this rule improves 1B-scale pretraining.

desk verdict A useful practical rule with solid empirical support, but the √d derivation rests on an independence assumption the model violates; treat the theory as a heuristic and send to review. read the letter →

arxiv 2506.15025 v1 pith:MJNSP2MS submitted 2025-06-17 cs.LG cs.AIcs.CLstat.ML

Soufiane Hayou , Liyuan Liu This is my paper · ORCID
classification cs.LGcs.AIcs.CLstat.ML MSC 68T0760B20
keywords embeddinglayerlearningratemaximalupdateparametrizationvocabularysizescalinglargeregimehyperparametertransferfeatureSignSGDtransformerpretraining
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper argues that the optimal learning-rate parametrization of a large language model changes with vocabulary size, and that modern LLMs sit in a regime where the standard theory mis-scales the embedding layer. Analyzing a minimal embedding-plus-projection linear model with a single SignSGD step, the authors compute how the magnitudes of embedding and projection updates depend on both width $d$ and vocabulary size $m$, under a Zipf-Mandelbrot token-frequency law. They find the dynamics interpolate between the $\mu$P regime (vocabulary fixed) and a Large Vocabulary regime in which the optimal ratio of embedding to hidden learning rate is $\Theta_d(\sqrt{d})$, not the $\Theta_d(d)$ ratio that $\mu$P predicts. If correct, the Large Vocabulary Parametrization (embedding LR $\eta d^{-1/2}$, hidden and output LR $\eta d^{-1}$, all initialization variances $d^{-1}$) should give better loss and perplexity than both standard parametrization and $\mu$P, and the paper reports such gains in small width-vocabulary sweeps and in a 1B-parameter pretraining run.

What carries the argument

The load-bearing object is the one-step feature-learning magnitude of a single token: after one sign-based (SignSGD) update of the embedding–projection network, the output change splits into an embedding contribution $\delta_i^E$ (updating $E$ with $W$ fixed) and a projection contribution $\delta_i^W$ (updating $W$ with $E$ fixed), and the optimality criterion is that both be $\Theta_d(1)$, so neither parameter silently stops learning as width grows. The workhorse computation is the covariance of a sign–Gaussian product, $X = \sum_{j=1}^d \mathrm{sign}(\langle v, W_j\rangle) W_j$ with $v$ a standard Gaussian vector, which evaluates to $\mathrm{Cov}(X) = d I_m + \frac{2}{\pi m} d(d-1) I_m$ through the Stein-type identity $\mathbb{E}[\mathrm{sign}(Z)G] = \sqrt{2/\pi}\,\rho$ for correlated standard Gaussians. The $d I_m$ term reproduces the $\mu$P scalings when $m$ is fixed; the $d(d-1)/m$ term, amplified by the Zipf-Mandelbrot law's $\Theta(m^{-1})$ average squared frequency, produces the $\sqrt{d}$ behavior when $m$ is large.

What would settle it

On the paper's two-layer model with Gaussian-initialized $E$ and $W$, take one sign-based update and measure the mean squared norm of a frequent token's embedding-row update; the theory predicts it grows as $d + 2d(d-1)/(\pi m)$. Vary $m$ over $\{d/4, d, 4d, 16d\}$: if the fitted coefficient of $d(d-1)/m$ departs from $2/\pi$, the covariance identity behind the LV regime is wrong. A complementary check on a full transformer holds the width fixed and quadruples the vocabulary, in which case the optimal embedding-to-hidden LR ratio should move measurably toward $\sqrt{d}$.

Watch

Extended reading notes

Core claim

The central discovery is that vocabulary size changes the width-scaling of feature learning in the embedding layer, and with it the optimal learning-rate ratio. For the two-layer model, the one-step update magnitudes follow $\bar\delta^E_i = \Theta_{m,d}(\eta_E \sigma_W \sqrt{d + 2d(d-1)/(\pi m)})$ and $\bar\delta^W_i = \Theta_{m,d}(\eta_W \sigma_E \sqrt{d + (\alpha_i^2/\bar\alpha^2)\, 2d(d-1)/(\pi m)})$; with $m$ fixed these reduce to the $\mu$P scalings, but when $m$ is large the first collapses to $\Theta(\eta_E \sigma_W \sqrt{d})$ while the second stays $\Theta(\eta_W \sigma_E d)$ for the most frequent tokens. Equalizing the two feature-learning components forces $\mathrm{LR}_{\mathrm{emb}}/\mathrm{LR}_{\mathrm{hidden}} = \Theta_d(\sqrt{d})$ under standard $d^{-1/2}$ initialization, a rule the authors call the $\sqrt{d}$-rule and package as the Large Vocabulary Parametrization (LVP). They validate it with coupled width-vocabulary sweeps and a 1B-parameter model pretrained on 1.75T tokens, where the $\sqrt{d}$ ratio improves training perplexity and Wikitext test perplexity relative to both the equal-ratio baseline and the $\Theta(d)$ ratio prescribed by $\mu$P.

Load-bearing premise

The load-bearing premise is that the initial differences between model outputs and targets behave as independent random noise in every entry, even though the two random matrices producing them make the rows correlated, and that one step of a simplified sign-only optimizer faithfully represents Adam's multi-step behavior.

Editorial extensions

If this is right

  • Under the Large Vocabulary Parametrization, setting $\mathrm{LR}_{\mathrm{emb}}/\mathrm{LR}_{\mathrm{hidden}} \approx \sqrt{d}$ should give lower pretraining loss and perplexity than both the equal-ratio standard recipe and the $\mu$P $\Theta(d)$ ratio whenever $m \gg d$, as in modern LLMs.
  • The optimal embedding learning rate is not width-constant as $\mu$P predicts; in the LV regime it scales roughly as $d^{-1/2}$, so hyperparameter transfer for the embedding layer requires re-tuning the embedding LR constant at each width.
  • As models adopt larger vocabularies relative to width, the gap between $\mu$P's predictions and the optimal embedding LR should widen, making the $\sqrt{d}$-rule more consequential at scale.
  • The hidden- and output-layer rules of $\mu$P (LR $\propto d^{-1}$, initialization variance $d^{-1}$) survive unchanged in the LV regime, so adopting LVP changes only the embedding LR and initialization, not the rest of the model.
  • The 1B-parameter experiments indicate the $\sqrt{d}$ rule converges faster as well as to a better final loss, which translates to token savings in large pretraining runs.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • A direct extension the paper does not state: the optimal ratio exponent should interpolate continuously between $1$ (small $m/d$) and $1/2$ (large $m/d$), so sweeping $m$ at fixed width should trace a curve rather than two separated regimes; the authors' small-model setup can test this immediately.
  • Because the projection-update term in Theorem 1 carries a token-frequency factor $\alpha_i^2/\bar\alpha^2$, the theory implicitly suggests token-frequency-aware embedding learning rates (larger updates for rare tokens) could improve on a single ratio; the paper leaves this untested.
  • The residual-stream coupling that carries the effect is specific to decoder-style transformers, so encoder-decoder or non-residual architectures with separate embeddings and projections may not show the $\sqrt{d}$ regime; training the same model with and without the residual pathway would clarify how general the rule is.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 4 minor

Summary. The paper studies how the optimal ratio between the embedding learning rate and the hidden-layer learning rate scales with model width d when the vocabulary size m is also large. Starting from a linear embedding–projection model trained with SignSGD, the authors derive asymptotic formulas for the per-token feature-learning magnitudes δ^E and δ^W under a Zipf–Mandelbrot token-frequency assumption. For fixed m these formulas reproduce the μP scaling rules; as m grows, the formulas are said to interpolate between the μP regime and a 'Large Vocabulary' regime in which the optimal ratio satisfies LR_emb/LR_hidden = Θ(√d). The paper proposes the 'Large Vocabulary Parametrization' (LVP), with SP-style initialization and embedding LR η d^{-1/2} while hidden/output LRs are η d^{-1}, and reports small-scale sweeps plus a 1B-parameter pretraining experiment showing improved perplexity over the equal-ratio baseline.

Significance. If the central claim holds, the paper addresses a real gap in μP theory for LLMs, where vocabulary size is large and often grows with model scale. The concrete, falsifiable prediction that LR_emb/LR_hidden ≈ Θ(√d) is valuable, and the 1B-token-scale experiment provides practical evidence that this rule helps. The paper is also refreshingly explicit about the limitations of its toy model and about the uncertainty of the optimality criterion. The main strength is a transparent set of covariance computations giving an explicit mechanism by which large vocabulary size changes the embedding-update magnitude. The main weakness is that the proofs rest on an independence assumption that is not justified by the actual product structure of the model, and the bridge from one-step SignSGD on a linear model to Adam-trained transformers is only heuristic.

major comments (3)
  1. [Appendix B, Assumption 2 and Theorems 3–4] As stated, the proof of Theorem 1 does not follow from the model because Assumption 2 is incompatible with the product structure it replaces. In the linear model, the vector v = E_i W − z_i used in B.3.1 is the i-th row of EW − Z, hence a function of W and E_i; it is not an exogenous N(0,1) vector independent of W. Theorem 3 requires that independence to compute the conditional covariance, and the covariance formula (3) relies on it. Similarly, in Theorem 4 the matrix M = D_α(EW − Z) contains E_i, so when the proof conditions on M, E_i is no longer independent of M, and the Stein-lemma correlation ρ_k = M_ik / sqrt(Σ_j M_jk^2) is not justified. If Assumption 2 is intended as an approximation valid when EW is small relative to Z (or in some other limit), that regime must be stated and proved; otherwise the Θ(√d) ratio is not derived from the declared model. This is load-bearing because Theorems 1 and 2 are the only theoretical support for the paper's central claim.
  2. [Section 3 and Section 5] The theoretical result is derived for a single SignSGD step on a model consisting only of an embedding layer and a projection layer, but the central claim and the large-scale experiments concern Adam-trained deep transformers with residual blocks. The sentence in Section 5 that the authors 'expect the results to hold' for multiple steps and for Adam 'since the main ingredient in our analysis is the normalization process' is an expectation, not a proof, and the residual-stream analogy in 'Extension to Transformer Architectures' is qualitative. Because the abstract advertises the Θ(√d) rule for LLM pretraining, the manuscript needs either a formal or a quantitatively justified extension, or a clearly weaker claim about what has been established theoretically versus what is supported only by experiments.
  3. [Section 3.1, definition of optimality] The Θ(√d) ratio is obtained by imposing the feature-learning criterion that both δ_i^E and δ_i^W are Θ(1) while taking initialization variances σ_E = σ_W = d^{-1/2}. This is a design criterion from the μP literature, not a theorem about minimizing the training loss, and Section 5 itself acknowledges that it is 'unclear whether the limit with this scaling is optimal.' The experiments give empirical support, but the theoretical statement that this is the 'optimal' ratio is stronger than the model analysis establishes. Please state explicitly in which sense the ratio is optimal and separate the optimality theorem from the empirical finding.
minor comments (4)
  1. [Section 2, page 4] The sentence asserting 'LRemb/LRhidden = Θ(d^{-1/2})' appears to be inverted; Table 2 and the abstract imply LRemb/LRhidden = Θ(d^{1/2}).
  2. [Section 3.1, footnote 11] The sign function is defined inconsistently: the footnote says S(z_i) = 1 if z_i ≥ 0 and 0 otherwise, while the proof in Appendix B uses S_j ∈ {±1} and the text elsewhere treats it as sign. Please use one definition throughout.
  3. [Figures 1 and 3] The figure captions and axis labels contain corrupted tokens such as 'Step/uni00A01e2' and 'Best/uni00A0(10%)'; these need to be cleaned.
  4. [Throughout] There are several typographical errors, including 'subopotimal' in Section 2 and 'the the embedding layer' in Section 3, that should be corrected in a revision.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the sqrt(d) embedding-LR ratio follows from explicit feature-learning and Zipf-frequency calculations, not from a self-referential fit.

full rationale

The paper's central claim, LR_emb/LR_hidden = Θ(√d) in the large-vocabulary regime, is not obtained by fitting or renaming the target ratio. The derivation defines feature-update components δ_i^E and δ_i^W, computes their average squared norms in Theorems 1 and 2 under Assumption 1 (Zipf-Mandelbrot frequencies) and Assumption 2 (random residual), and then solves the feature-learning balance conditions δ_i^E = Θ(1) and δ_i^W = Θ(1) for η_E/η_W. The √d factor emerges from the asymptotic formulas, in particular the d(d-1)/m terms with m >> d and ᾱ² = Θ(1/m), rather than being put in by hand. The choice σ_W = σ_E = d^{-1/2} is reported as an empirical observation, not as a fit to the LR-ratio prediction; the paper states the dependence explicitly as η_E/η_W = Θ(σ_E/σ_W √d). Changing the initialization scaling would change the constant factor but not reduce the derivation to its conclusion. Self-citations [6, 7, 8, 25] are background on feature-learning limits and are not load-bearing for the LV-regime result. The main caveat, Assumption 2's i.i.d. treatment of EW-Z despite its functional dependence on E and W, is a correctness or approximation gap rather than circularity: it does not assume the conclusion, and the paper's external experiments (Figures 3-5, plus refs [1, 5]) provide independent empirical support. Section 5 candidly lists open issues, including whether the limit is truly optimal and the harder multi-step analysis, which is consistent with a non-circular derivation. No equation or fitted parameter is reused as its own prediction, so no circular step is exhibited.

Assumptions & free parameters 4 free parameters · 6 assumptions · 0 invented entities

The central scaling law rests on an unverified Gaussian-residual assumption, a one-step SignSGD linear model, an empirical Zipf fit, and an empirically chosen SP initialization. The ledger makes these inputs explicit; the '√d rule' is not a fully self-contained derivation from standard transformer dynamics.

free parameters (4)
  • base learning rate η = ≈0.2 (small models)
    Grid-tuned for the d=256 model in Section 4.1 and reused for larger widths; a tunable constant in LVP and µ P.
  • initialization variance exponents σE=σW=d^{-1/2} = d^{-1/2}
    Selected from experiments as best performance, not derived from first principles; this choice drives the ηE/ηW = Θ(√d) conclusion.
  • Zipf exponent a_m = ≈1.01 (m=8192), ≈1.03 (m=32768)
    Fitted to BPE token frequencies on Wikitext2 in Figure 2; assumed bounded between constants >1/2 for Lemma 1.
  • token frequencies α_i = from dataset
    Empirical input distribution used to define ¯α² and the high-frequency token amplification in Theorem 2.
assumptions (6)
  • domain assumption Token frequencies follow the Zipf-Mandelbrot law (Assumption 1)
    Underlies Lemma 1 and the Θ(m) amplification that separates the embedding and projection scaling. Validated only on Wikitext2 BPE tokenizers, then assumed for LLM pretraining at scale.
  • ad hoc to paper Residual at initialization EW−Z has i.i.d. N(0,1) entries (Assumption 2)
    Makes the covariance computation in Theorem 1 tractable. It is not validated and is likely violated because E and W are independent in the product EW, which induces correlation across rows.
  • domain assumption Large-batch infinite-batch approximation (N→∞)
    Used in Section 3 to replace empirical frequencies by α_i. Reasonable for LLM batch sizes of millions, but an approximation that ignores sampling noise in gradient updates.
  • ad hoc to paper SignSGD with one step approximates Adam dynamics
    Makes the theorem tractable. The authors state that the results are expected to hold for Adam and multiple steps because of elementwise normalization, but this is not proven.
  • domain assumption Residual stream preserves the embedding-projection direct interaction in full transformers
    Used to transfer the linear two-layer result to stacked transformer architectures. The authors explicitly acknowledge that a full transformer analysis is not possible with their framework.
  • standard math Standard math: Stein's lemma, rotational invariance, Gaussian moments
    Used in Appendix B to compute covariances of sign-Gaussian products; accepted background.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size." pith.science (2026). https://pith.science/paper/MJNSP2MS

@misc{pith2026250615025,
  author       = {Pith},
  title        = {Pith review of: Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/MJNSP2MS}},
  note         = {Machine review of arXiv:2506.15025}
}
abstract

Pretraining large language models is a costly process. To make this process more efficient, several methods have been proposed to optimize model architecture/parametrization and hardware use. On the parametrization side, $\mu P$ (Maximal Update Parametrization) parametrizes model weights and learning rate (LR) in a way that makes hyperparameters (HPs) transferable with width (embedding dimension): HPs can be tuned for a small model and used for larger models without additional tuning. While $\mu$P showed impressive results in practice, recent empirical studies have reported conflicting observations when applied to LLMs. One limitation of the theory behind $\mu$P is the fact that input dimension (vocabulary size in LLMs) is considered fixed when taking the width to infinity. This is unrealistic since vocabulary size is generally much larger than width in practice. In this work, we provide a theoretical analysis of the effect of vocabulary size on training dynamics, and subsequently show that as vocabulary size increases, the training dynamics \emph{interpolate between the $\mu$P regime and another regime that we call Large Vocab (LV) Regime}, where optimal scaling rules are different from those predicted by $\mu$P. Our analysis reveals that in the LV regime, the optimal embedding LR to hidden LR ratio should roughly scale as $\Theta(\sqrt{width})$, surprisingly close to the empirical findings previously reported in the literature, and different from the $\Theta(width)$ ratio predicted by $\mu$P. We conduct several experiments to validate our theory, and pretrain a 1B model from scratch to show the benefit of our suggested scaling rule for the embedding LR.

Figures

Figures reproduced from arXiv: 2506.15025 by the authors.

Figure 1
Figure 1. Optimal Embed￾ding LR vs model width. Transformer model with vocab size 32768, trained on Wikitext2 for 1e4 steps with Adam. Limitations of µP for LLMs. Recent empirical work [1] has shown that while the Θd(d −1 ) scaling effectively leads to HP transfer for hidden/output LRs, the optimal embedding LR seems to shift as width d increases, contradicting muP’s pre￾dictions of a constant optimal embedding LR across scal… view at source ↗
Figure 2
Figure 2. Verification of the Zipf-Mandelbrot Law assumption for m ∈ {8192, 32768}. In both cases, we train a BPE tokenizer on Wikitext2, tokenize the dataset, and show the frequencies. The Zipf-Mandelbrot Law is a power law frequently observed in ranked data. It is well-studied in probability and statis￾tics, and the coefficient am is usually close to 1 [PITH_FULL_IMAGE:figures/full_fig_p008_2.png] view at source ↗
Figure 3
Figure 3. Optimal ηE vs dmodel. The shaded curves is calculated as the geo￾metric mean of all learning rates that give a train loss within 20% of the best train loss. The average is cal￾culated across learning rate grid and random seeds. We trained a small decoder-only transformer with the follow￾ing configuration: an embedding layer, two hidden layers, and a projection layer. We used Wikitext2 as our training dataset, with a… view at source ↗
Figures from the paper (2 more)
Figure 4
Figure 4. Figure 4: LLM Pre-training Perplexity of 1B Transformer. Note [PITH_FULL_IMAGE:figures/full_fig_p011_4.png]
Figure 5
Figure 5. Figure 5: Perplexity of 1B Transformer on the Wikitext Test Set [PITH_FULL_IMAGE:figures/full_fig_p011_5.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

38 extracted references · 18 canonical work pages

  1. [1]

    Prince, Bj ¨orn Deiseroth, Andres Felipe Cruz-Salinas, Carlo Luschi, Samuel Weinbach, and Douglas Orr

    Charlie Blake, Constantin Eichenberg, Josef Dean, Lukas Balles, Luke Y. Prince, Bj ¨orn Deiseroth, Andres Felipe Cruz-Salinas, Carlo Luschi, Samuel Weinbach, and Douglas Orr. u- µp: The unit-scaled maximal update parametrization, 2025. URL https: //arxiv.org/abs/2407.17465

    arXiv 2025

  2. [2]

    Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit, 2023

    Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit, 2023. URL https://arxiv.org/abs/2309.16620

    arXiv 2023

  3. [3]

    On lazy training in differentiable programming, 2020

    Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming, 2020. URL https://arxiv.org/abs/1812.07956

    arXiv 2020

  4. [4]

    Infinite- width limit of deep linear neural networks

    L ´ena¨ıc Chizat, Maria Colombo, Xavier Fern ´andez-Real, and Alessio Figalli. Infinite- width limit of deep linear neural networks. Communications on Pure and Applied Mathematics, 77(10):3958–4007, 2024. doi: https://doi.org/10.1002/cpa.22200. URL https://onlinelibrary.wiley.com/doi/abs/10.1002/cpa.22200

    doi:10.1002/cpa.22200 2024

  5. [5]

    Alemi, Roman Novak, Peter J

    Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A. Alemi, Roman Novak, Peter J. Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, and Jeffrey Pennington. Scaling exponents across parameterizations and optimizers,

    work page

  6. [6]

    On the infinite-depth limit of finite-width neural networks

    Soufiane Hayou. On the infinite-depth limit of finite-width neural networks. Trans- actions on Machine Learning Research, 2022

    work page 2022

  7. [7]

    On the impact of the ac- tivation function on deep neural networks training

    Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. On the impact of the ac- tivation function on deep neural networks training. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Ma- chine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2672–2680. PMLR, 09–15 Jun 2019. URL h...

    work page 2019

  8. [8]

    Stable resnet

    Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau. Stable resnet. In Arindam Banerjee and Kenji Fukumizu, edi- tors, Proceedings of The 24th International Conference on Artificial Intelligence and Statis- tics, volume 130 of Proceedings of Machine Learning Research, pages 1324–1332. PMLR, 13–15 Apr 2021

    work page 2021

Show all 38 references
  1. [9]

    Delving deep into rec- tifiers: Surpassing human-level performance on imagenet classification, 2015

    Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rec- tifiers: Surpassing human-level performance on imagenet classification, 2015. URL https://arxiv.org/abs/1502.01852

    2015 arXiv

  2. [10]

    Neural tangent kernel: Conver- gence and generalization in neural networks, 2020

    Arthur Jacot, Franck Gabriel, and Cl ´ement Hongler. Neural tangent kernel: Conver- gence and generalization in neural networks, 2020. URL https://arxiv.org/abs/ 1806.07572

    2020 arXiv

  3. [11]

    Muon: An optimizer for hidden layers in neural networks,

    Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks,

  4. [12]

    Kingma and Jimmy Ba

    Diederik P . Kingma and Jimmy Ba. Adam: A method for stochastic optimization,

  5. [13]

    URL https://kellerjordan.github.io/posts/muon/

  6. [14]

    The llama 3 herd of models, 2024

    Llama-Team. The llama 3 herd of models, 2024. URL https://arxiv.org/abs/2407. 21783

    2024

  7. [15]

    Pointer sen- tinel mixture models, 2016

    Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sen- tinel mixture models, 2016

    2016

  8. [16]

    An empirical study of µp learning rate transfer, 2025

    Lucas Lingle. An empirical study of µp learning rate transfer, 2025. URL https: //arxiv.org/abs/2404.05728

    2025 arXiv

  9. [17]

    Poole, S

    B. Poole, S. Lahiri, M. Raghu, J. Sohl-Dickstein, and S. Ganguli. Exponential expres- sivity in deep neural networks through transient chaos. 30th Conference on Neural Information Processing Systems, 2016. 13

    2016

  10. [18]

    Schoenholz, J

    S.S. Schoenholz, J. Gilmer, S. Ganguli, and J. Sohl-Dickstein. Deep information prop- agation. In International Conference on Learning Representations, 2017

    2017

  11. [19]

    Radford M. Neal. Priors for infinite networks. In Bayesian Learning for Neural Net- works, volume 118 of Lecture Notes in Statistics, pages 29–53. Springer New York, 1996. ISBN 978-0-387-94724-2. doi: 10.1007/978-1-4612-0745-0 2

    1996 doi

  12. [20]

    The falcon series of open language models, 2023

    Falcon Team. The falcon series of open language models, 2023. URL https://arxiv. org/abs/2311.16867

    2023 arXiv

  13. [21]

    Gemma 3 technical report, 2025

    Gemma Team. Gemma 3 technical report, 2025. URL https://arxiv.org/abs/2503. 19786

    2025

  14. [22]

    Scaling laws with vocabulary: Larger models deserve larger vocabularies, 2024

    Chaofan Tao, Qian Liu, Longxu Dou, Niklas Muennighoff, Zhongwei Wan, Ping Luo, Min Lin, and Ngai Wong. Scaling laws with vocabulary: Larger models deserve larger vocabularies, 2024. URL https://arxiv.org/abs/2407.13623

    2024 arXiv

  15. [23]

    G. Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation.arXiv preprint arXiv:1902.04760, 2019

    1902 arXiv

  16. [24]

    Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer.arXiv preprint arXiv:2203.03466, 2022

    Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer.arXiv preprint arXiv:2203.03466, 2022

    2022 arXiv

  17. [25]

    Phi-3 technical report: A highly capable language model locally on your phone

    Phi Team. Phi-3 technical report: A highly capable language model locally on your phone. arXiv:2404.14219, 2024

    2024 arXiv

  18. [26]

    How does critical batch size scale in pre-training?,

    Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham Kakade. How does critical batch size scale in pre-training?,

  19. [27]

    Selected Studies of the Principle of Relative Frequency in Language

    George Kingsley Zipf. Selected Studies of the Principle of Relative Frequency in Language. Harvard University Press, Cambridge, MA, 1932. 14 A Infinite-width analysis and µP As the width d grows, model hyperparameters such as initialization variance and learn- ing should be ad...

    1932

  20. [28]

    Tensor programs vi: Fea- ture learning in infinite-depth neural networks, 2023

    Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Fea- ture learning in infinite-depth neural networks, 2023. URL https://arxiv.org/ abs/2310.02244

    2023 arXiv

  21. [32]

    Because Wj∼N (0,Im) and Mj =⟨v,Wj⟩, for anyi∈ [m] the pair (Mj,Wji) is jointly Gaussian with correlationρ := vi ∥v∥

    Conditional mean of one summand. Because Wj∼N (0,Im) and Mj =⟨v,Wj⟩, for anyi∈ [m] the pair (Mj,Wji) is jointly Gaussian with correlationρ := vi ∥v∥. Stein’s lemma above gives EW [ sign(Mj)Wj|v] = r 2 π v ∥v∥

  22. [33]

    Still conditional on v, sign(Mj) is±1 with equal probability, independent of the magnitude of Wj

    Conditional covariance of one summand. Still conditional on v, sign(Mj) is±1 with equal probability, independent of the magnitude of Wj. A direct computation gives CovW (SjWj|v) =Im− 2 π v ∥v∥ v⊤ ∥v∥. Because the rows W1,...,W d are independent, E[X|v] =d q 2 π v ∥v∥, Cov(X|v)...

  23. [34]

    Because the Yj’s are conditionally independent, Ev Cov(X|v) = Ev h dX j=1 Cov(Yj|v) i =d Im− Ev µ(v)µ(v)⊤ | {z } = 2 πmIm = d− 2d πm Im

    First term. Because the Yj’s are conditionally independent, Ev Cov(X|v) = Ev h dX j=1 Cov(Yj|v) i =d Im− Ev µ(v)µ(v)⊤ | {z } = 2 πmIm = d− 2d πm Im. (We used Ev[v/∥v∥v⊤/∥v∥] =Im/m by rotational invariance.)

  24. [35]

    From Step1, E[X|v] =dµ (v), so Covv E[X|v] =d2 Covv µ(v) =d2 2 π Covv v ∥v∥ = 2d2 πmIm

    Sign–correlation term. From Step1, E[X|v] =dµ (v), so Covv E[X|v] =d2 Covv µ(v) =d2 2 π Covv v ∥v∥ = 2d2 πmIm. Adding the two contributions gives Cov(X) = h d− 2d πm + 2d2 πm i Im = h d + 2 πmd(d− 1) i Im, which is the covariance stated in the theorem. 17 B.3.2 Technical Resul...

  25. [36]

    Conditioning on E, each entry of E⊤M a centered Gaussian variable, hence E[S(E⊤M)|E] = 0

    Zero Mean. Conditioning on E, each entry of E⊤M a centered Gaussian variable, hence E[S(E⊤M)|E] = 0. Therefore E[X] = 0

  26. [37]

    Fix a column index k

    Second Moment. Fix a column index k. We have Xk = dX a=1 EiaS(Zak), where, Zak := (E⊤M)ak =EiaMik +Wak, W ak := X j̸=i EjaMjk. SinceXk has zero mean, we have E[X2 k] = E " dX a=1 E2 ia + 2 X a<a′ EiaS(Zak)Eia′S(Za′k) # =d + 2 X a<a′ E [EiaS(Zak)Eia′S(Za′k)]. Fix some a,a′ ∈ {1...

  27. [38]

    Different columns ofM (differentk) are independent, so Cov(X) is diagonal and each coordinate variance is the same as above

    Independence across columns. Different columns ofM (differentk) are independent, so Cov(X) is diagonal and each coordinate variance is the same as above. 19 C Additional Empirical Details C.1 Experimental Setup for Figure 1 • Model arch: embedding layer, 2 hidden layers, proje...

  28. [2017]

    URL https://arxiv.org/abs/1412.6980

    arXiv

  29. [2024]

    URL https://arxiv.org/abs/2407.05872

    arXiv

  30. [2025]

    URL https://arxiv.org/abs/2410.21676

    arXiv

Pith tools

Reviewed August 15, 2026 · model on record in the stance chip above.