A Practical Guide to Complete(d) Hyperparameter Transfer

1. Thesis

The paper is best read as an attempt to make the following workflow legitimate:

That is stronger than ordinary muP-style transfer in three ways. First, the target is not just a single global learning rate. The paper tunes learning rate, initialization scale, AdamW epsilon, AdamW beta1 and beta2 through their complements, weight decay, and residual multipliers. Second, the tuning can be per-module, so QKV weights, MLP matrices, norms, embeddings, output weights, and depth positions can receive different multipliers. Third, the transfer is not only across model size; it includes batch size and total training duration.

The paper's central idea is not "there is one magic hyperparameter value." It is: there are dimensionless hyperparameters that remain meaningful only after you divide out the predictable scale effects created by architecture, optimizer, batch size, and training horizon. Complete(d)P is the proposed way to divide out those effects.

2. Mental Model and Notation

The paper separates two kinds of quantities:

• Training configuration: the dimensions you intentionally scale: width, depth, batch size, and token count.
• Hyperparameters: optimizer and initialization choices you want to tune once and transfer.

Let a base run have width N0, depth L0, batch size B0, and token count D0. Let a target run have N, L, B, and D. The paper describes scaling through ratios:

A rule like learning_rate *= sqrt(m_B / m_D) means: keep the base hyperparameter value, but before using it at target scale, multiply it by that factor. For per-module hyperparameters, the same idea applies to every parameter group, after also multiplying by that group's module and depth factors.

The Loss-Surface Version

Think of a hyperparameter setting h as a point in a high-dimensional landscape:

Naively, when you change the configuration, the landscape moves. The goal of parameterization is to describe a transformed hyperparameter coordinate h_tilde such that:

This is why hyperparameter transfer is more than "scale the learning rate." It is a claim about the stability of an optimization landscape after a coordinate change.

3. What Complete(d)P Changes

The paper builds on CompleteP, which itself adapts muP/depth-muP ideas to transformers. You said width scaling can be mostly left out, so this section keeps width short and focuses on the conceptual invariants.

Width Initialization: Reconcile the Table with the Spectral Condition

Let an iid-initialized matrix W have shape d_out x d_in and entry standard deviation s. Its spectral norm is typically of order

The spectral condition asks for ||W||_* = Theta(sqrt(d_out / d_in)). Solving for the entry scale gives

Therefore the assembled rule initialization variance *= m_N^(-1) is not a different result. It is the fixed-aspect-ratio specialization: when both hidden dimensions scale by m_N, the min factor is constant and the variance falls by m_N^(-1). For a square hidden matrix this is simply Var(W_ij) = Theta(1 / d). The old table was misleading because it called this merely “finite-activation scaling” and hid the fixed-shape assumption.

Width Learning Rates: Same Spectral Budget, Different Optimizers

The optimizer-independent target is the update norm ||Delta W||_* = Theta(sqrt(d_out / d_in)). A scalar learning-rate rule follows only after specifying the matrix produced before that scalar is applied.

Depth Scaling Has Two Problems

A transformer block is a residual update:

If depth increases but each block keeps making an equally large residual update, the total residual path can become too large. If each block update shrinks too aggressively, the network becomes too close to an identity map and stops learning rich features. Depth parameterizations choose how much each block's contribution shrinks with depth.

Complete(d)P uses a residual multiplier:

where alpha is in the theoretically motivated range [0.5, 1]. The paper reports transfer for alpha = 0.5 and alpha = 1. It also notes that alpha = 0.5 can be slightly better at the largest model in their experiments, consistent with the "feature diversity" motivation from depth-muP.

A Clean Derivation of the Depth Exponents

Here is a derivation that does not use tensor programs. It is not a complete proof for real transformers, but it isolates the same exponent constraints using only residual sums, operator norms, and the feature-learning desideratum.

Part I: keep the forward residual stream stable

First isolate one elementary fact about sums. For one scalar coordinate, write:

If the terms are centered and weakly correlated,

then variances add:

This is the sqrt(L) or random-walk case. Equivalently, for vector residual writes:

If the cross inner products average away, the squared norm is O(L), so the norm is O(sqrt(L)). If the terms instead share an O(1) aligned component, the sum is coherent:

Now apply that fact to a width-normalized residual stack at a fixed training step:

Assume the branch output has elementwise O(1) scale after LayerNorm/RMSNorm and width scaling:

At initialization, if residual branch outputs are centered and weakly correlated across layers, the random-walk estimate gives:

Forward stability asks that feature coordinates stay O(1). Therefore:

If the branch outputs have coherent nonzero drift, use the L case instead:

This is why centering, normalization, zero-mean projections, and bias handling matter. The depth-muP paper uses mean subtraction in its clean model; transformers rely on normalization and approximately centered residual projections. If that centering fails, alpha = 1/2 can be too aggressive.

Part II: preserve feature learning over optimizer steps

At this point we have solved only the forward-scale problem: keep h O(1) as depth grows. The spectral feature-learning condition has a second half: parameter updates should induce an O(1) change in hidden features. We now ask how the feature change caused by one optimizer step should scale when the same model family is made deeper. That second desideratum is still unsatisfied.

To see the remaining problem, suppose we used the residual multiplier L^(-alpha) but made each block's branch output change by Theta(1) RMS per coordinate after one optimizer step. In the coherent accumulation case analyzed below, the output-feature update would scale like:

For alpha = 1/2, this blows up like sqrt(L). For alpha = 1, it is O(1). For alpha > 1, it vanishes. So unless we choose the ODE endpoint alpha = 1, the optimizer update scale must also depend on depth. This is the problem that motivates learning-rate depth scaling.

Do not call it "learning rate" yet; first introduce an unknown exponent for the effective per-block branch-function update:

Here gamma is not an additional architecture parameter and not a result imported from Complete(d)P. It parameterizes the unknown depth dependence of the branch-function update so that we can solve for it. Using Theta, rather than only O, states both an upper and a lower scale; the later nontriviality argument needs the update not to vanish faster than this. After deriving gamma, we will ask how an optimizer should implement the resulting effective update scale.

What the Jacobian is doing in the update recurrence

Start from one residual block, suppressing normalization and attention details:

After one optimizer step, both the block parameters and the block input have changed. Write Delta h_l for the training-induced change in the residual stream at depth l, and Delta theta_l for the parameter update in block l. First-order Taylor expansion gives:

Here J_l means the hidden-state Jacobian of the branch function at the current input:

It is not the parameter-update effect. It describes how a perturbation already present in the residual stream entering block l is transformed by that block before continuing upward. The direct parameter-update effect is instead:

So we need both terms. Delta F_l is the new feature-write injected by block l because its own weights changed, holding its input fixed. J_l Delta h_l is the way block l responds because all earlier blocks changed its input. In words: each layer injects a local update, and later layers transport every earlier local update.

Iterating the recurrence gives the useful mental model:

The product is the transport operator from depth k to the output. If these transport products have O(1) operator norm along the relevant feature-update directions, then the size of Delta h_L is controlled by summing the injected terms a_L Delta F_k. If the transport products explode, depth scaling has failed. If they strongly contract or rotate the injected updates into cancelling directions, the model may be stable but the learned feature update can vanish.

This assumption can fail in both directions. If the products of I + a_L J_l have growing singular values, update perturbations amplify as they move upward and the depth limit is unstable. If the products are strongly contractive, or if transported update directions cancel like centered noise, the learned update may be only sqrt(L)-sized or smaller rather than L-sized. The Complete(d)P depth rule is therefore best read as preserving the coherent feature-learning component, not as a guarantee that every architecture and optimizer implementation has such a component.

The direct contribution of block l to the residual stream after one training step is therefore a vector a_L Delta F_l. Its width-normalized size is:

For a vector v in R^n, ||v||_2 / sqrt(n) equals sqrt((1 / n) * sum_i v_i^2): the root-mean-square coordinate size. When the norm is not dominated by a small number of exceptional coordinates, this is also the scale of a typical element. It is not a claim that every coordinate has exactly that magnitude, nor is it a maximum-coordinate bound.

Unlike random initialization noise, training updates are produced by the same loss and are strongly correlated through the backward pass, then transported through the later residual stack. Under bounded, non-degenerate transport, the relevant deterministic feature-learning component has coherent-sum scaling:

A depth-aware version of the spectral feature-learning condition asks for feature updates at a fixed normalized depth to be O(1) per coordinate: not exploding, and not vanishing. Therefore:

Combine this with the initialization condition alpha >= 1/2. Also require the effective per-branch optimizer update not to grow with depth. Since gamma = 1 - alpha, this gives gamma >= 0, hence alpha <= 1. The resulting interval is:

This is the depth-scaling interval used by Complete(d)P. It also says what the effective per-block update scale must be:

Now connect this to the optimizer. If the parameter-to-branch-output Jacobian has O(1) scale under the width parameterization, and AdamW is kept in the same normalized-update regime, then the simplest way to make Delta F_l scale like L^(alpha - 1) is to scale the hidden learning rate by that factor:

This is where the Complete(d)P table's hidden-weight and hidden-bias/norm learning-rate depth factors come from. They are not assumed at the start of the derivation; they are the implementation of the update-size condition.

The key distinction is signal accumulation versus noise accumulation. Let:

Throughout this guide, the size of a width-n vector means its per-coordinate RMS norm, ||z||_2 / sqrt(n). Raw l2 norm is fine when width is fixed, but ||z||_2 / sqrt(n) is the scale that makes "O(1)" mean O(1) per coordinate.

If the transported updates are independent and centered, the cross terms vanish:

That is the right model for initialization noise, but it is the wrong object to preserve if the goal is learned feature motion. A learned update should contain a descent-aligned component shared across depth:

The centered noise part accumulates like sqrt(L). The aligned signal part accumulates like L. Equivalently, if every z_l has per-coordinate size s, and the average pairwise correlation is rho, then:

With rho = 0, the sum is sqrt(L) * s. With constant positive rho, it is O(L * s). Thus the coherent learned-update assumption is really an assumption about non-vanishing positive cross terms among transported update contributions. If one scaled to preserve only independent centered noise, one would get:

But then any coherent signal component would scale as L * L^(-1/2) = sqrt(L) and explode. Depth-muP and Complete(d)P instead preserve the deterministic learned component:

Arguments for and against the transport assumption

The assumption is not a free theorem for transformers. The defensible arguments are structural and empirical, and each has a clear failure mode.

Tensor Programs VI gives empirical support for the scaling, but not a direct measurement of the transport products above. In simple residual networks, Depth-muP preserves the optimal hidden learning-rate constant across depth, while block-only scaling without LR scaling and standard parameterization show large shifts. TP VI also reports a LayerNorm variant where Depth-muP preserves LR transfer, and feature-diversity experiments supporting the non-degenerate path story. The negative evidence is equally important: ODE-like (alpha, gamma) = (1, 0) can be a feature-learning limit but still shows LR shift, block-depth-2 experiments show failures, and their Megatron transformer study describes the practical scaling as useful but brittle. Complete(d)P supplies more direct modern-transformer evidence, but it should be read as empirical validation of the recipe, not a theorem that every LLM stack satisfies bounded coherent transport.

What Else Needs Depth Scaling?

Learning rate is not the only depth-scaled quantity. The Complete(d)P table uses depth factors in three different places, each solving a different problem.

The SGD row follows directly under an O(1) parameter-to-branch-output Jacobian. The residual multiplier enters the raw gradient once:

Adam divides by a matching gradient scale when epsilon is scaled with it, so its preconditioned update is O(1) and Delta F_Adam ~= eta_Adam. This gives eta_Adam ~= L^(alpha - 1), the factor printed in Complete(d)P. Thus the target feature update is optimizer-independent, but the scalar learning rate that realizes it is not.

Muon realization: orthogonalize, then apply the shape scale

Ideal spectral Muon provides a third realization, but its update is easier to understand if direction, shape normalization, and step size are introduced separately. Let W in R^(p x q), so an input has q coordinates and an output has p. Let M_W be the momentum matrix that Muon will turn into an update direction. This is the matrix being factorized; it is not the weight matrix W.

1. Orthogonalization chooses the direction. Write the reduced singular-value decomposition and its polar factor as

P_W keeps the singular-vector directions of the momentum while replacing every nonzero singular value by one. Thus ||P_W||op=1 whenever M_W is nonzero.

Where the first-order model comes from. Hold every other model parameter fixed and regard the scalar loss as a function of this matrix W. Let G_W=grad_W loss(W), a matrix with the same shape as W. For a candidate displacement Delta W, restrict the multivariable loss to the one-dimensional path

Ordinary one-variable differentiation and the chain rule give

The subscript F means Frobenius. On angle brackets it denotes the matrix inner product; on a norm it denotes the associated Euclidean length of all matrix entries. It is not another function or random variable:

Entry by entry, the derivative says how sensitive the loss is to that entry of W; the Frobenius inner product adds sensitivity times displacement over all entries. The fundamental theorem of calculus integrates this directional derivative along the path:

If the gradient is Lipschitz in Frobenius norm, meaning

then Frobenius Cauchy-Schwarz, |<A,B>F| <= ||A||F ||B||F, bounds the remainder at every point on the path:

Thus the omitted curvature is second order in a sufficiently small displacement. The LMO separates the actual displacement into a scalar step size and a norm-bounded direction, Delta W=eta_M D. Replacing the current gradient by its momentum estimate gives the exact error bookkeeping behind the shorthand first-order model:

The approximation suppresses two different errors:

• Curvature error: R(eta_M D), which grows quadratically with the step under a smoothness assumption.
• Gradient-estimation error: eta_M <G_W-M_W,D>F. Momentum is a lagged, noisy estimate and need not equal or be unbiased for the current gradient.

What the linear minimization oracle asks. For a fixed positive eta_M, only the direction-dependent linear term can be optimized. In the ordinary spectral unit ball, the LMO problem is

"Locally best" therefore has a precise and limited meaning: among directions with the same allowed operator-norm size, choose the one predicting the largest first-order loss decrease. It does not mean the displacement is globally optimal, or that a finite stochastic step must lower the actual loss.

What operator/nuclear norm duality means. The dual of any norm is defined by the largest inner product attainable in its unit ball:

For matrices, the norm dual to operator norm is nuclear norm. Their definitions are

where sigma_i are the nonzero singular values of M_W=U Sigma V^T. The duality can be proved directly from this SVD rather than taken as a black box:

The middle inequality follows because u_i and v_i are unit vectors and operator norm is the largest possible Euclidean amplification of any unit input. Under ||D||op <= 1, no singular-direction contribution can be smaller than -sigma_i. The polar direction reaches every one of those lower bounds simultaneously:

This proves that -P_W solves the spectral LMO. If M_W is rectangular or rank deficient, its action on singular null spaces does not affect the linear objective, so the minimizer need not be unique; the reduced polar factor is the canonical partial-isometry solution. The same calculation with D=+P_W attains +||M_W||nuc, proving the supremum side of the dual-norm identity as well.

The LMO paper's overall interpretation. The norm ball declares what counts as one unit of parameter movement; the LMO spends that unit where the local linear model predicts the most decrease. In fact, the same oracle template produces different familiar normalized optimizers solely by changing that geometry:

Operator norm measures worst-case action on a vector, so the spectral choice budgets how strongly the matrix update can act on any input direction rather than budgeting total entrywise energy. It discards the momentum's singular-value magnitudes and retains its singular subspaces:

Hence the ideal spectral LMO is scale-invariant: LMO(c M_W)=LMO(M_W) for c > 0. Momentum averages noisy gradients to produce M_W; the LMO maps that estimate to a norm-controlled direction; and eta_M restores the separately chosen amount of movement. The unconstrained algorithm in Pethick et al. repeatedly adds these scaled LMO directions even though the norm ball constrains each direction rather than constraining W itself.

2. RMS coordinates determine the rectangular shape factor. Euclidean operator norm compares Euclidean input and output lengths. Feature-learning arguments instead compare typical coordinate sizes, so define

Therefore a unit RMS-to-RMS update budget is the spectral constraint ||D||op <= sqrt(p/q). Solving the same linear minimization problem in this rescaled ball gives

The factor disappears only from the worst-case normalized magnitude, because it exactly converts between Euclidean and RMS coordinates:

3. The polar direction remains in the realized activation action. Let H in R^(N x q) contain representative branch inputs as rows. Holding those inputs fixed while applying the direct matrix update gives

For full-column P_W when p >= q, P_W^T P_W=I_q and kappa_W(H)=1. For full-row P_W when p < q, P_W^T P_W is a rank-p projection. Then kappa_W depends on how the activations align with that projection; with many isotropically oriented activation samples, kappa_W(H)^2 concentrates near p/q. Thus sqrt(p/q) fixes the worst-case RMS-to-RMS operator scale, not the average action on every activation distribution.

4. Only now impose the depth target. The preceding depth argument requires each learned branch-function change to have typical-element size L^(alpha - 1). If this matrix action is the relevant realization of that branch change, its learning rate must satisfy

The familiar Muon depth factor L^(alpha - 1) follows only after measuring that kappa_W=Theta(1) across the depth family. In a nonlinear multi-parameter branch, this matrix probe is not by itself a proof about the complete Delta F_l; the branch-output update and its transport still need to be measured. Unlike Adam, ideal Muon has no second-moment state or Adam epsilon to scale. Its finite Newton-Schulz approximation adds a different obligation: momentum singular values and orthogonalization error must remain in the regime where the nominal polar update is actually realized. The full batch/duration extension is derived in Section 4.9.

So the claim is not “LR is the only other thing necessary.” Once residual scaling fixes forward feature size, the spectral update-size condition forces an effective hidden-update scale. The optimizer determines how LR maps to that scale, and AdamW epsilon must also scale so that Adam remains in the assumed normalized regime.

How AdamW Epsilon and Weight Decay Fit the Same Derivation

For a residual block, the raw gradient of an internal hidden weight contains the residual multiplier L^(-alpha). Ignoring width factors:

AdamW epsilon should be on the same scale as the gradient entering Adam's denominator; otherwise, changing depth changes whether Adam behaves like normalized gradient descent or epsilon-damped small-step descent. This gives the depth factor:

If the gradient scale is multiplied by c, then Adam's first and second moment scales become:

To preserve the same optimizer regime, epsilon must be multiplied by the same factor:

Thus epsilon follows the gradient scale, not the learning-rate scale. Since the hidden-gradient depth scale is L^(-alpha):

If epsilon is left fixed while gradients shrink with depth, epsilon eventually dominates the denominator:

Then the hidden update is smaller than the intended normalized Adam update. For alpha = 1/2, the intended branch-function update is L^(-1/2), but epsilon domination can push it toward L^(-1), making the accumulated feature update vanish.

With epsilon scaled this way, Adam's normalized update direction is O(1). Then the learning-rate factor L^(alpha - 1) makes:

Weight decay: from parameter shrinkage to residual-stream change

Weight decay needs no separate depth multiplier in the table because it is already multiplied by the learning rate in PyTorch AdamW. Its direct parameter-space update is:

This is not yet a hidden-state change. Let K_l^W be the parameter-to-branch-output Jacobian, so K_l^W[W_l] is its directional action along W_l:

The factor L^(-alpha) appears only when this branch-output change is written into the residual stream, because h_(l+1) = h_l + L^(-alpha) F_l. For the linear example F_l(h) = W_l h, K_l^W[W_l] = W_l h = F_l(h). More generally, assume its RMS coordinate size is Theta(1). Then:

The 1 / L scale guarantees that the accumulated decay effect cannot explode under bounded transport, even if all blocks add coherently. It does not guarantee coherent representation-space drift:

Thus a depth-independent lambda needs no additional depth factor for per-step stability. A nonvanishing Theta(1) feature effect still requires alignment after the different blocks' Jacobians; weight shrinkage is coherent in parameter space but need not be coherent in representation space. This is a one-optimizer-step statement. Full-run decay also depends on the number of steps and is handled by the later batch/duration rules.

Connections to a Spectral Feature-Learning Condition

If your mental model is spectral rather than Tensor-Programs-based, the connection is: a good parameterization keeps the update-induced change to representations at a controlled scale. In width, that means weight updates should not create exploding singular-value shifts or vanish into an NTK-like lazy regime. In depth, it means the product/sum of many residual Jacobian perturbations should remain stable as more blocks are inserted.

Complete(d)P can be seen as adding optimizer-scale bookkeeping to that feature-learning condition. It does not merely scale forward activations. It also scales learning rates, AdamW epsilon, and weight decay so the optimizer's update to each tensor has the intended magnitude relative to that tensor's role in the forward pass.

The subtle point is that the spectral condition's Delta h is a training-step update, not the adjacent-depth increment h_l - h_{l-1}. With depth tending to infinity, adjacent residual increments must shrink. Otherwise the residual stack would explode. The desideratum is instead:

The condition alpha + gamma = 1 is exactly the depth version of making Delta_train h O(1). The condition alpha >= 1/2 is the depth version of making h O(1) at initialization. Thus the depth interval is not an extra ad hoc rule; it is what the spectral desiderata become after accounting for the fact that a depth-L residual stream is a sum of L small operators.

There is another important qualifier: this is an update-size condition for a fixed optimizer-step scale. If the training horizon changes, the meaning of one discrete step changes too. For longer horizons, the condition should not be read as "make every additional micro-step move features by O(1)." It should be read as "over a fixed macroscopic unit of training time, the feature motion is O(1)." The batch/duration SDE rules later in the guide are exactly the correction that turns a one-step spectral condition into a training-path condition.

What the Endpoints Mean

Depth-muP argues that alpha = 1/2 is uniquely optimal if the objective is maximal feature diversity. Complete(d)P is more empirical and transformer-focused: it verifies transfer for both alpha = 1/2 and alpha = 1, and treats the whole interval as theoretically justified for stable, nontrivial, faithful depth scaling. Those are compatible claims, not the same claim.

Specific Fixes Relative to CompleteP

These are not cosmetic table edits. Each changes the optimizer coordinates used by CompleteP. The paper's appendix argues from the muP entry scales: hidden activations are Theta(1), hidden activation gradients are Theta(1 / N), and Adam epsilon should track the raw gradient scale seen by the optimizer.

Fix 1: QK-norm parameters are shared across the scaled head axis

In the implementation analyzed by the paper, a QK norm is applied separately inside each attention head, but its affine multiplier vector is shared across heads. Let i index the H heads and let j be a fixed coordinate within a head. The same scalar m_j therefore appears in every operation q_hat_(i,j) = m_j * q_(i,j) + b_j. When model width is N = H * d_head and d_head remains fixed,

More generally, c denotes any positive width-independent limit of H / N. Autograd adds the gradient contribution from every use of the shared parameter. Suppressing batch and token indices, which do not change this width count, define delta_(i,j) := dL / dq_hat_(i,j). The chain rule gives:

Thus one head contributes only Theta(1 / N), but there are Theta(N) uses of the same parameter. Rescale one contribution to order one:

The “empirical average” is therefore the average across heads within this forward/backward pass; it appears because the QK-norm parameter is shared across the axis whose size grows with width. Suppose the X_i are approximately iid across heads and converge to a distribution with finite variance and mean mu. The law of large numbers plus the usual fluctuation scale gives:

The notation Y_N = O_p(N^(-1/2)) means that sqrt(N) * Y_N is bounded in probability. Explicitly: for every epsilon_prob > 0, there is a finite constant M such that, for all sufficiently large N,

Informally, this says the typical random fluctuation is no larger than constant times 1 / sqrt(N); it is not a deterministic bound for every realization. The probability refers to the random ensemble underlying the width limit, such as initialization and sampled data.

Fix 2: input-embedding epsilon follows the active-row gradient

A token lookup does not sum over width: the input index is width independent, and the gradient of an active embedding row is the backpropagated gradient of the embedding output. Under the paper's hidden-gradient scale,

This changes CompleteP's width-independent input-embedding epsilon to m_N^(-1). It is a statement about active coordinates. Sparse token visitation, untouched rows, tied embeddings, and optimizer state initialization can make the distribution across the entire table less clean. The authors report that omitting this correction broke their coordinate checks.

Fix 3: remove the explicit unembedding multiplier by changing coordinates

CompleteP used an explicit forward multiplier a = m_N^(-1):

Define the actual unembedding matrix W_new = a * W_old. Then logits = W_new * h with no special forward multiplier. The initialization transformation is exact:

The learning-rate transformation depends on the optimizer. For Adam, let Q(g; epsilon) = m_hat / (sqrt(v_hat) + epsilon). Scaling every gradient in its history by positive a gives Q(a * g; epsilon) = Q(g; epsilon / a). Hence

Ignoring epsilon or transforming it exactly, direct Adam on W_new therefore uses eta_new = a * eta_old, which gives the table's unembedding LR factor m_N^(-1). Exact coordinate equivalence would also use epsilon_new = epsilon_old / a. Complete(d)P intentionally does not preserve CompleteP's epsilon rule: it argues instead that gradients in the new raw unembedding coordinates are Theta(1), so output-layer epsilon should remain Theta(1).

For SGD the coordinate transformation is different because its update is proportional to the raw gradient:

Do not reuse the Adam folding rule for SGD. Instead, for an actual unembedding matrix with d_out = vocabulary_size and d_in = model_width, the spectral-condition SGD realization gives eta_SGD = Theta(d_out / d_in). With fixed vocabulary size this scales as 1 / d_in, though it follows from the SGD spectral argument rather than Adam scale invariance.

Removing the multiplier is operationally useful because memory-efficient losses such as cut cross-entropy can consume the ordinary unembedding matrix directly without materializing or specially rescaling full logits.

4. Batch-Size Transfer via SDEs

The batch-size part of the paper uses stochastic differential equations: continuous-time approximations to noisy optimizer recurrences. The optimizer-specific scaling argument is only as good as the moment matching and limiting assumptions that lead to its SDE.

Roadmap and Claim Boundary

There are two different claims in this section. Keeping them separate is essential: only the first is an attempt at SDE invariance.

The narrative is: batch covariance gives sigma ~= B^(-1/2); the adaptive recurrence reveals an effective step h = eta^2; preserving the coefficients of that SDE gives the batch rules; only then do we compose with a longer or shorter token horizon.

The One SDE Fact Used Repeatedly

The local bridge from an optimizer update to Brownian motion is derived in foundations section 7. Over a short interval h, an SDE has the Euler-Maruyama form

Its noise covariance is diffusion * diffusion^T * h. Everything below is the algebra needed to rewrite the discrete adaptive update into this shape and identify which hyperparameter combinations appear in its drift, diffusion, and state dynamics.

1. Set Up the Adaptive Approximation

This has two layers. The paper's noise-dominated RMSPropW calculation makes the learning-rate, iteration-count, and PyTorch-style weight-decay factors explicit. Keeping the coupled moment recurrences then forces the beta and epsilon factors. That second layer is what makes this an adaptive-optimizer result.

1a. What changes when batch changes?

Hold the current parameters theta_k fixed for one moment. A minibatch gradient is the true population gradient plus a random estimation error:

Read this from right to left. z_k is standardized fresh noise. The matrix Sigma^(1/2) chooses its parameter-space directions and correlations. The scalar sigma says how large the whole noise term is. Under the conditional batch-covariance law from foundations section 9, sigma ~= 1 / sqrt(B): a larger batch makes the random error smaller.

This is the only place where batch size enters the derivation. The rest asks: after the optimizer normalizes, averages, and decays this noisy gradient, which combinations of hyperparameters must stay unchanged when sigma changes?

1b. What RMSPropW actually computes

RMSPropW has two jobs. First, it builds v, an exponential moving average of squared gradients. Second, it divides the next gradient by sqrt(v) + epsilon, so a coordinate with consistently large gradients receives a smaller normalized update:

All operations are coordinate-wise. beta_2 controls how long v remembers past squared gradients. epsilon prevents a small denominator. The first term in the parameter line is PyTorch-style decoupled weight decay; the second is the adaptive gradient step. We start with RMSPropW because it isolates the adaptive denominator. Adam will add a first-moment average later.

Nothing has been rescaled yet. The next stage chooses a continuous-time unit and introduces only the rescaled variables needed to express these two discrete lines as an SDE.

1c. Where sqrt(v) ~= sigma comes from

Freeze theta temporarily and abbreviate mu = grad_L(theta) and xi_k = Sigma(theta)^(1/2) * z_k. Thus g_k = mu + sigma * xi_k. Write the second-moment state as v_k = sigma^2 * u_k. Substitution into the RMSPropW recurrence gives

The square is elementwise. If mu is O(1), the noise covariance is O(1), and mu / sigma -> 0, then the right-hand side is O(1) whenever u_k is O(1). Starting from v_0 / sigma^2 = u_0 = O(1), the recurrence therefore keeps u_k = O(1) on a fixed SDE-time interval. Consequently, coordinate by coordinate,

The shorthand sqrt(v) ~= sigma hides the direction-dependent O(1) factor sqrt(u_i); it does not say every coordinate has the same denominator. If the parameters are nearly stationary for long enough that the EMA equilibrates, then

The first argument establishes the scale without assuming equilibrium; the second gives a familiar stationary picture. This scaling can fail in a nearly noiseless coordinate, when epsilon dominates, or when the moving average has not been initialized on the corresponding scale.

2. Derive the Effective Time Step

Ignore weight decay for one moment and use sqrt(v_k) = sigma * sqrt(u_k). The adaptive gradient update can be written exactly as

In the noise-dominated regime, the random term xi_k is O(1), so its parameter displacement is O(eta), provided epsilon / sigma stays finite rather than dominating the denominator. The rescaling in section 2a enforces precisely that condition. A Brownian increment has size sqrt(h); matching these sizes gives

This is why adaptive methods use eta^2 as effective SDE time in this calculation. It is not an arbitrary convention and it differs from the usual SGD gradient-flow coordinate h = eta. But an SDE needs a deterministic displacement of order h as well. The mean-gradient term in the display above has size eta / sigma. Requiring it to be O(h = eta^2) gives

More precisely, we keep eta * sigma -> sigma_0 with finite nonzero sigma_0. Then eta / sigma = h / (eta * sigma) ~= h / sigma_0: the mean contributes finite drift, while the standardized noise contributes finite Brownian diffusion. Because eta -> 0 and the product stays nonzero, sigma -> infinity follows. It is not an independent physical claim.

Alternative clocks and noise scalings

The conclusions above are about total movement over a fixed finite SDE interval t in [0, T], not about one discrete update. A compact way to audit the alternatives is to suppose sigma ~= eta^(-a). For a > 0, this uses the noise-dominated denominator scale derived above; for the fixed-noise case a = 0, the denominator is simply O(1). Suppressing those O(1) denominator factors, the normalized parameter noise is O(eta) and the mean-gradient increment is O(eta / sigma) = O(eta^(1 + a)).

So yes: for a < 1, the apparent problem is a genuine blow-up accumulated over a fixed finite interval of the h = eta^2 clock. One can instead use the mean-preserving clock h = eta^(1 + a). When a < 1, its total noise RMS is eta^((1 - a) / 2) -> 0, giving a deterministic ODE. In the important fixed-noise case a = 0, this is simply h = eta: finite drift and vanishing accumulated noise. When a > 1, that mean-preserving clock instead makes the total noise diverge. Thus a = 1 is the boundary at which this simple scaling gives both effects finite on one common time axis.

2a. Rescale only what would otherwise diverge

With h = eta^2, introduce the following quantities. Each is chosen to stay finite or to become a coefficient of the limiting SDE:

Bundle the scaled denominator into

The second equality is just sqrt(v_k) = sigma * sqrt(u_k), multiplied by eta. It is why epsilon appears as eta * epsilon in the SDE coefficients.

2b. Rewrite each part of one update

Use P_k = eta * (sqrt(v_k) + epsilon), h = eta^2, and sqrt(h) = eta. The three pieces of the parameter update become

First distinguish a noisy gradient component from the resulting noisy parameter increment:

The next display does not introduce any new noise. It only rewrites that same parameter increment so the Brownian-sized factor sqrt(h) * z_k is visible:

Adding these three identities gives

The bracketed random variable sqrt(h) * z_k has covariance h * I, so it becomes dW_t in the Brownian limit. Matching conditional mean and covariance gives the RMSPropW SDE

Geometry over a finite SDE horizon

To understand the shape of this process, temporarily freeze the drift and diffusion near one point and write the local SDE as d Theta_t = b * dt + G * dW_t, with D = G * G^T. This constant-coefficient approximation has the exact distribution

The mean moves linearly with t; the covariance ellipse has squared radii proportional to t, hence radii proportional to sqrt(t). Along any unit direction q,

Thus, if drift and noise are comparable in a particular direction at t = 1, then one unit of SDE time is indeed a noisy horizon. But one SDE-time unit is a chosen rescaled unit, not an epoch or a universal amount of training. A ratio of 100 at t = 1 is also compatible with this limit; the asymptotic argument requires finite coefficients, not that their ratio be near one.

The raw-gradient noise dominance and the macroscopic SDE signal-to-noise ratio are different statements. In the simplified one-coordinate, stationary, epsilon_0 = 0 case, u ~= Sigma_(ii), the limiting diffusion has approximately

So the macroscopic SDE drift can be 100 times its diffusion scale over one SDE-time unit while the raw gradient is still noise dominated in the joint limit. For finite training, however, whether that separation is a good approximation remains an empirical question.

We have now exposed the direct parameter-update coefficients: sigma_0, epsilon_0, and lambda / eta. The remaining coefficient c_2 enters through the evolution of the hidden second-moment state u, which we derive next.

3. Track the Second-Moment State

The parameter SDE contains P_t, but P_t is not a fixed matrix. It is the adaptive preconditioner built from the evolving second-moment state:

A large entry of P_t suppresses both deterministic gradient movement and the Brownian kick in that parameter direction; a small entry amplifies both. To preserve the parameter path, we must preserve the law of u_t, not only the displayed Theta equation.

Start from the elementwise RMSPropW recurrence and the definitions of the rescaled state and rate:

Subtract v_k from the first line, then divide by sigma^2:

Finally, insert g_k = grad_L + sigma * Sigma^(1/2) * z_k, so that g_k / sigma = grad_L / sigma + Sigma^(1/2) * z_k. This gives

Every square in this display is elementwise, as in RMSPropW; it is not a matrix square. Be careful about which quantity is averaged: E[Sigma^(1/2) * z_k] = 0. It is the elementwise square whose expectation is diag(Sigma).

To see this, let A = Sigma^(1/2) and y = A * z_k. The i-th coordinate is y_i = sum_j A_(ij) * z_j. Since E[z_k] = 0 and Cov(z_k) = I, E[z_j * z_m] = delta_(jm), where delta_(jm) is one when j = m and zero otherwise. Therefore

Collecting these coordinates gives E[y^2] = diag(A * A^T) = diag(Sigma). This is why the second-moment state drifts toward diag(Sigma), rather than toward a scalar variance or the full covariance matrix.

In the complete limiting family, grad_L / sigma -> 0. Make the stochastic size explicit by writing Y_k = (grad_L / sigma + Sigma^(1/2) * z_k)^2 and R_k = Y_k - E[Y_k | F_k]. Then

Assuming the standardized gradient noise has bounded conditional fourth moments, Var(R_k | F_k) = O(1). The second line is therefore a centered O(h) random kick with conditional variance O(h^2) per optimizer step. Over a fixed SDE interval of length T, there are about T / h such conditionally centered kicks. Their accumulated variance is therefore

Yes: the accumulated variance is O(h), and it vanishes because h = eta^2 -> 0. This is a conditional-variance or martingale calculation; it does not require the whole sequence of increments to be unconditionally independent. Thus u has the deterministic limit

This is not an arbitrary omission of noise: the mean squared-gradient update survives at order dt, while its centered fluctuation is too small to survive as Brownian noise. The parameter update still has Brownian noise because its normalized gradient fluctuation is order sqrt(h).

The limiting state is therefore the joint process X_t = (Theta_t, u_t). Its drift and Brownian diffusion can be written schematically as

This makes c_2 load-bearing: it is the speed at which the denominator adapts relative to Brownian/SDE time. Change c_2, and u_t follows a different path; then P(u_t) differs; then both the drift arrow and covariance ellipse of every later parameter update differ.

Why matching coefficients means matching the continuous process

Imagine a point moving through parameter space. During an interval dt, the drift term gives its mean displacement, while the Brownian term gives a small random kick with a state-dependent covariance ellipse. The matrix P(u_t)^(-1) bends and rescales both: it is the optimizer's local geometry. The u_t dynamics determine how that geometry changes along the path.

If two configurations start from the same joint state and have the same joint drift and diffusion functions, then every sufficiently short interval has the same conditional mean and covariance. Chaining those local steps produces the same limiting path distribution, subject to the usual SDE existence assumptions. If any coefficient changes, the continuous process itself changes; transferred hyperparameters then need not mean the same optimization behavior.

4. Hold the Batch-Only SDE Fixed

Increase batch size by kappa, so sigma_new = sigma / sqrt(kappa). Keeping the coefficients of the same RMSPropW SDE fixed gives:

The last line is why larger-batch training uses fewer iterations at fixed token count. The paper source's displayed k' = kappa * k conflicts with the time-scale equation; the invariant K * eta^2 gives K_new = K / kappa.

5. Add Adam's First-Moment State

Its increment becomes

Thus its limiting SDE is

The same Brownian increment drives Theta and M, because both receive the same noisy gradient. Adam's bias corrections become gamma_i(t) = 1 - exp(-c_i * t). Away from the initial boundary layer, the parameter drift is

Equivalent conventions move a sqrt(gamma_2) factor into the denominator, depending on epsilon placement. The scale invariants are unchanged. Holding c_1 = (1 - beta_1) / eta^2 fixed forces

Malladi et al.'s formal Adam result begins after a positive SDE time t_0, rather than uniformly from the first update. This is a real warm-up limitation: bias-correction factors change rapidly near zero, and the theorem assumes matched states after that boundary layer.

6. Compose Batch and Duration

Let batch and token horizon change by m_B and m_D. The new number of updates is K_new = (m_D / m_B) * K. Requiring the same integration horizon gives

The same invariant bookkeeping yields the table's optimizer factors:

For pure batch scaling (m_D = 1), this preserves the full adaptive SDE: sigma_0, epsilon_0, c_1, c_2, decay drift, and horizon. For duration scaling it preserves the horizon and moment/denominator time scales, but not the identical SDE:

Longer-duration scaling therefore intentionally changes effective noise. It is the paper's iso-horizon hypothesis and empirical rule, not a deduction from batch-only SDE invariance. This is why token-horizon transfer deserves more skepticism than the batch-only square-root rule.

7. Choose the Correct Weight-Decay Semantics

The weight-decay factor is determined by the decay term's SDE drift, so optimizer semantics matter. In PyTorch AdamW the decay increment is -eta * lambda * theta, which becomes -(lambda / eta) * theta * dt when dt = eta^2. The invariant is therefore lambda / eta.

In the original Loshchilov-Hutter convention, called AdamLH in the paper, the decay increment is instead -lambda * theta. In SDE time this is -(lambda / eta^2) * theta * dt, so it preserves a different quantity:

For batch-only scaling, this reads sqrt(kappa) for PyTorch AdamW and kappa for AdamLH. Switching optimizer semantics without changing this factor creates a real drift mismatch.

8. Where This Approximation Can Fail

The SDE argument is a limit argument. It is most credible when the optimizer step is small enough that many discrete steps approximate a smooth noisy path. It becomes less credible when:

• Batch is too large. You take too few optimizer steps; the discrete path no longer approximates the same SDE.
• Learning rate is too large. Euler-Maruyama discretization error becomes a real training effect.
• Gradient noise is not close to batch-averaging noise. Correlated data, packing, dedup artifacts, curriculum, or repeated data can change the noise law.
• Gradient clipping is active. Clipping changes both drift and noise in a state-dependent way not represented by the simple SDE.
• Adam epsilon is numerically irrelevant or numerically dominant. If epsilon is far below optimizer precision, scaling it may not matter. If it dominates the denominator, scaling can matter too much.
• Fused optimizers differ subtly. Bias correction, epsilon placement, decoupled decay semantics, FP32 master weights, and BF16/FP8 optimizer state can all change the effective rule.
• Beta scaling leaves the legal range. If kappa * (1 - beta) is not small, beta can become invalid or qualitatively different. This is a warning that you are outside the asymptotic regime.

9. Muon Batch and Duration: A Finite-Budget Argument

The question is how Muon's learning rate and momentum should change when batch size B or total sample budget T changes. The AdamW SDE argument above cannot simply be relabeled Muon. Ideal Muon applies a polar factor to a momentum matrix, and Polar(c M)=Polar(M) for every positive scalar c. Batch noise changes the singular directions of the update, not merely its magnitude.

The idealized update and its geometry

Use the theorem's indexing convention: first construct m_(k+1), then use it for update k. An implementation that calls this state m_k differs only by an index shift.

This section follows the source and calls the momentum complement alpha. The assembled rule cards call the same quantity delta_mu to distinguish it from the residual-depth exponent.

Here <A,B>F=trace(A^T B) is the Frobenius inner product. If m_(k+1)=U Sigma V^T, one minimizer is D_k=-U V^T: the negative polar factor. Operator/nuclear duality gives

The earlier Muon depth realization derives this direction and then applies architecture-dependent shape and depth scales. This section isolates the dimensionless budget multiplier.

What the finite-budget criterion measures

For spectral updates, ||g_k||nuc is the largest first-order loss decrease available per unit operator-norm displacement. Three related summaries are

The summed descent calculation naturally controls S_avg. This is an ergodic or random-iterate statement: if an iterate is selected uniformly after the run, its expected available first-order descent is S_avg. The two minima are weaker corollaries, not new descriptions of the proof. The parent theorem reports S_Emin; the scaling-law paper displays S_minE.

Thus a small bound certifies that the run can reach and spend steps in low-slope regions. It does not bound E[||g_K||nuc], final loss, held-out loss, or generalization. Those require additional arguments or direct measurement.

Derive the one-step descent inequality

Define the ideal direction selected using the true gradient:

Because D_k minimizes the momentum inner product,

Add and subtract the momentum in the true-gradient inner product:

Operator/nuclear Holder duality says |<A,D>F| <= ||A||nuc ||D||op. Both candidate directions have operator norm at most one, so

The factor two has a concrete origin. One error comes from evaluating the chosen direction with the true gradient; the other comes from momentum choosing a different direction than the true gradient would choose.

Now assume f is L_smooth-smooth in operator/nuclear geometry:

Set U=eta D_k, use ||D_k||op <= 1, and insert the direction inequality:

Unroll and bound the momentum error

The argument now needs four explicit assumptions:

• E[xi_k | history_k]=0: each minibatch gradient is conditionally unbiased.
• E[||xi_k||F^2 | history_k] <= sigma^2/B: effective batch averaging follows a finite-variance 1/B law.
• ||g_i-g_(i+1)||nuc <= L_smooth ||theta_i-theta_(i+1)||op: gradients are Lipschitz in the optimizer's geometry.
• ||A||nuc <= rho ||A||F: Euclidean/Frobenius noise can be converted to the dual norm. For a p x q matrix, the worst-case rho is sqrt(min(p,q)).

Initialize m_0 with a batch-gradient estimate obeying the same variance bound, so E[||m_0-g_0||F^2] <= sigma^2/B.

Let r=1-alpha and e_k=m_(k+1)-g_k. Recursing through the exponential average gives

Each component has a different scaling:

Combining them yields the per-step momentum mismatch bound

Averaging the initialization term over the run uses (1/K) sum_k r^(k+1) <= 1/(alpha K). The factor two in the one-step direction inequality then doubles every momentum-error contribution.

Telescope the loss and assemble the bound

Sum the one-step inequality over k=0,...,K-1. Consecutive losses cancel, and f(theta_K) >= f_inf, so f(theta_0)-f(theta_K) <= Delta_0, where Delta_0=f(theta_0)-f_inf. The left side is eta sum_k E[||g_k||nuc], not a terminal-gradient term. Dividing by eta K therefore gives the trajectory average in the elementary smooth, unregularized result

The cited parent theorem is slightly more general: it allows a convex regularizer R, measures generalized stationarity at the new iterate, and uses F=f+R. Its proof has the same summed-loss backbone, but the theorem reports the weaker best-iterate consequence. As adapted to batch variance sigma^2/B, write that published form as

Optimize learning rate at fixed momentum

Fix a total sample budget T=B K, so K=T/B. This substitution is exact inside the proxy:

With alpha, B, and T fixed, only the first and last terms depend on eta. Define positive constants

Here A/eta is a finite-horizon reach term: steps that are too small cannot spend the initial loss gap within K updates. The term C eta is the certificate's finite-step neighborhood: large steps pay curvature and momentum-staleness costs. Minimizing their sum gives the best guarantee from this proof. It need not minimize measured average stationarity, endpoint stationarity, or loss unless the bound is sufficiently tight and order-preserving for the problem being trained.

Since U_eta also diverges as eta approaches either zero or infinity, its sole derivative-zero point is the unique global minimum:

For m_B=B_new/B and m_D=T_new/T, calibrating the unknown prefactor at one base point gives the candidate

This scalar factor happens to match Complete(d)P's AdamW factor, but the reasons differ. Adam preserves coefficients in an adaptive SDE approximation; this Muon rule minimizes a finite-horizon convergence-bound certificate. It is therefore a sweep prior for optimal retuning, not a path-invariance law.

If batch size is also a decision variable

Substituting eta_star back into the full fixed-momentum proxy retains burn-in and gives

At fixed total budget, increasing B gives cleaner gradients but leaves only K=T/B optimizer steps. Let x=sqrt(B) and a_T=D_alpha/sqrt(T)+H_alpha/T. The batch-dependent proxy is

This is why the increasing deterministic/burn-in side is balanced against the decreasing noise side. The equality is not an extra heuristic; it characterizes the global minimum of this positive a x+c/x proxy. At long horizons, H_alpha/T is lower order than D_alpha/sqrt(T), so

This interior optimum depends on holding T fixed. At fixed K, increasing batch size does not remove optimizer steps, and this proxy improves monotonically with cleaner gradients. In practice B is integer and hardware-constrained; if the continuous optimum lies outside the feasible interval, the constrained optimum is at a boundary.

Optional alternative: tune the momentum horizon

Optional bookkeeping: decoupled weight decay

What must be measured for practical Muon

Sources: non-Euclidean trust-region and Muon theorem, LMO hyperparameter-scaling derivation, non-Euclidean gradient-noise scale, and finite Newton-Schulz spectral scaling.

5. Token-Horizon Transfer

Token horizon is the least settled part of the paper and probably the most important for scaling-law practice. Width and batch transfer try to preserve a training behavior. But increasing total tokens is supposed to improve the model. So what should be invariant?

Two Ways to Train on More Tokens

Suppose you want four times as many tokens. There are two extreme paths:

• More iterations at same batch: optimizer runs four times longer.
• Larger batch at same iterations: optimizer sees four times as much data per step but takes the same number of steps.

The Paper's Observation

Holding batch size fixed and increasing the number of training iterations, the optimal learning rate decays roughly like:

where kappa is the factor by which the training horizon increases. This resembles empirical horizon-scaling results in related work, though reported exponents vary by model size and setup.

The Iso-Horizon Hypothesis

The paper then proposes a stronger interpretation. Rather than saying "longer training simply needs lower LR," they propose an SDE-coordinate prescription for scaling token count:

In table form, this is the case m_B = m_D. Learning rate, epsilon, weight decay, and betas do not change from the batch/duration rule. The model improves because each step sees a cleaner gradient, not because the optimizer travels for longer in continuous time. This keeps the integration horizon K * eta^2 fixed, but does not preserve the identical adaptive SDE: the effective noise coefficient sigma_0 = eta * sigma falls as batch grows. That changed noise level is the intended effect and an empirical hypothesis, not an invariance theorem.

Connection Back to Delta h from Updates

The spectral feature-learning condition says that, at the model-scale coordinate system being considered, a training update should produce an O(1) change in hidden features. That statement is clean when the number of optimizer steps is fixed. It becomes incomplete when we scale the number of steps.

In the Adam/RMSProp SDE picture used by the paper, the effective continuous-time step size scales like the square of the learning rate:

If we increase the number of iterations by kappa at fixed batch size, preserving the same SDE integration horizon asks for:

Since the hidden-feature update from one optimizer step is proportional to the effective optimizer update, the per-step feature change also shrinks:

But the accumulated stochastic motion over the whole run remains O(1) in quadratic variation:

This is a horizon analogue of the spectral condition, but only for the SDE integration scale and simplified noise quadratic variation. It does not show that every feature direction, deterministic drift component, or full parameter distribution is invariant after the time grid is refined.

If instead we scale tokens by increasing batch size with the same number of optimizer steps, then m_B = m_D, the table leaves the learning rate unchanged, and per-step Delta h remains comparable. The improvement comes from lower gradient noise, not from a longer optimizer path.

Why Schedule Shape Matters

The paper runs an additional learning-rate schedule enumeration: 4,842 non-increasing piecewise-constant schedules for a 121M model, across five token horizons from 155M to 1.24B tokens. The best short-horizon schedule was not a prefix of the best long-horizon schedule.

That matters because it says there is no simple greedy rule like: find the best schedule for 155M tokens, then continue it for 310M, then continue again. The final LR creates an optimization bias, but high early LR is needed for progress. The right tradeoff changes with horizon.

6. Per-Module Hyperparameters

The paper's most practically interesting claim is that per-module hyperparameter tuning can produce real gains and still transfer. This is not obvious. Per-module tuning creates a high-dimensional and unstable search problem.

Why Global Hyperparameters Are Suspicious

A transformer contains parameter groups with very different jobs:

• input embeddings map sparse token IDs into representation space;
• Q, K, V matrices control attention geometry and value transport;
• output projections mix heads back into the residual stream;
• MLP matrices control most parameter count and nonlinear feature expansion;
• norm and QK-norm multipliers act more like gates/scales than ordinary matrices;
• unembedding weights translate hidden states into vocabulary logits.

There is no strong reason these should share one learning rate, one weight decay, one Adam epsilon, or one initialization multiplier. Standard training recipes do so largely because the search problem is too large, not because the equality is theoretically justified.

The Depth-Type Kronecker Factorization

Fully independent per-layer, per-module, per-hyperparameter values are too many. The paper uses a factorized parameterization:

Read it as a known base scale times a learned outer product

Exponentiating the displayed equation separates what the transfer theory supplies from what the search must learn. For one hyperparameter family, write H_(m,l) for the value assigned to module type m in residual block l. Then

The first two factors are deterministic transformations prescribed by the model-scale and batch/duration rules. The search only learns the dimensionless corrections a_m and b_l. This matters: the factorization is not claiming that, say, QKV and MLP weights literally need the same learning rate. It says that, after dividing out their known parameterization scale, their remaining deviations can be described by one type effect and one depth effect.

Why it is called a Kronecker factorization

Remove the known scale factors and arrange the learned multipliers as a matrix with module types as rows and residual-block positions as columns:

Thus M is a rank-one outer product. Calling it a Kronecker factorization refers to the vectorized form: a type vector and a depth vector are combined into one vector over every type-by-depth pair. The log form in the paper is useful because products become sums: log M_(m,l) = log a_m + log b_l. Strictly, that additive log matrix is a no-interaction two-way-effects model; the original multiplier matrix M is the rank-one object.

With |M_types| repeated module types and L residual blocks, an unconstrained table has |M_types| * L learned entries. The factorization uses |M_types| + L stored numbers, or |M_types| + L - 1 effective degrees of freedom: multiplying every a_m by a constant c and every b_l by 1 / c leaves every product unchanged. A practical implementation can remove this redundant gauge by, for example, fixing the geometric mean of the depth multipliers to one.

A concrete three-type, four-layer example

Suppose that, after known scale corrections, the learned type and depth factors are

QKV receives twice the base correction at every depth, MLP receives half, and all repeated module types share the same relative depth profile: layer 3 is 20% higher and layer 4 is 10% lower than the middle reference. The factorization lets the search discover both facts without separately tuning all twelve cells.

What it assumes, and what it rules out

The factorization asserts that relative type effects do not change with depth, and relative depth effects do not change with type:

It therefore cannot represent a genuine type-by-depth interaction, such as “QKV should be especially hot only in early layers, while MLP should be especially hot only in late layers.” Such a pattern needs a full table, a low-rank sum of several outer products, or another interaction model. Singleton modules such as an input embedding or unembedding do not have a repeated residual-block depth coordinate, so in practice they naturally use their type and known scale factors without this repeated-block depth profile.

Why normalized depth is used for transfer

At source depth L, the values b_l are samples of an assumed continuous profile b(t) at normalized positions t = l / L. At target depth L_new, the paper samples the linearly interpolated profile at t_new = l_new / L_new. For example, a source layer at the halfway point is meant to correspond to a target layer at the halfway point, not to a target layer with the same absolute index. This is a smoothness assumption about hyperparameters across depth, motivated by the depth ODE/SDE picture; it is not a consequence of that picture alone.

The paper tests both sides of this structural assumption for learning rates. Removing depth multipliers but retaining type-specific rates is worse than the full factorization, so depth position adds useful signal. Starting a more difficult fully uncoupled per-layer search from the factorized optimum produced virtually no further improvement in their experiment. That is encouraging but not a proof of exact factorization: the unconstrained search has many more dimensions and may have missed a better solution.

The Search Problem

The paper reports that the per-module loss landscape has sharp instability cliffs. Different modules become unstable at different learning-rate scales, and module interactions make the stability boundary hard to model. Random search either samples too many divergent runs or misses the optimum if the bounds are too conservative. Standard Gaussian-process Bayesian optimization can struggle because the landscape is non-stationary.

Their practical solution is a trust-region random search:

• optimize in log2 hyperparameter space;
• sample around the best known point inside an infinity-norm box;
• start with radius 1, meaning each sampled multiplier can vary by up to 2x;
• if no improvement happens for 100 trials, shrink the radius by 0.7;
• run asynchronously with up to 100 simultaneous trials.

In the main per-module experiment, they optimize 79 hyperparameters: per-module multipliers for learning rate, weight decay, AdamW epsilon, 1 - beta1, 1 - beta2, initialization standard deviation, and per-depth residual block multipliers.

7. Assembled Transfer Rules

We can now assemble the operational recipe. The preceding sections supplied its ingredients: width/depth parameterization and depth learning-rate/epsilon scaling; the batch SDE rule; duration scaling; and the reason that different module groups need different base hyperparameters. The paper's operational recipe is for PyTorch-style AdamW. The spectral target itself is optimizer independent, so this atlas separately marks the SGD and broader Adam realizations, plus a later conditional Muon extension, rather than presenting every exponent as one theorem.

Four AdamW Examples

8. Empirical Evidence

The paper's empirical case has several layers.

Width/Depth Global Learning Rate Transfer

Complete(d)P is tested for global learning-rate transfer across width and depth. The paper reports transfer for both alpha = 0.5 and alpha = 1, including in a modern transformer variant with QK norms. It also reports that removing QK norms does not reproduce a breakdown observed in CompleteP.

Batch-Size Transfer

The paper compares learning-rate transfer with and without the SDE batch-size rule. With the square-root rule, learning-rate optima align across batch sizes. Without the adjustment, transfer fails. A similar story is shown for weight decay: unadjusted weight decay drifts, while the effective rescaled decay transfers.

Token-Horizon Transfer

Holding batch fixed and increasing training iterations shifts the optimal learning rate downward roughly as 1 / sqrt(kappa). Increasing token count by increasing batch size while holding iteration count fixed gives much cleaner transfer. The paper also fits a scaling-law comparison claiming the full SDE token-horizon rule improves large-scale performance compared with width/depth-only transfer.

Per-Module Transfer

The main result: per-module hyperparameters improve small-scale training and the gains persist after transfer. The paper's figure caption reports:

• search at 50M parameters and 1.6B tokens;
• global HPs require 2.3x longer training to match the small-scale per-module result;
• direct transfer to a roughly 600x larger FLOP budget;
• 1.32x speed-up at large scale plus improved downstream benchmarks.

The body text separately describes transfer to a 1.3B parameter, 26B token scale, about 420x compute, where per-module HPs give a 27% speed-up to reach the global-HP baseline loss. The appendix also describes downstream benchmark comparisons for 7.2B parameter, 143B token runs.

Proxy Scale Matters

The appendix shows that transferring a global learning rate from a very small 58M model to a 483M model incurs a small penalty. Optimizing at 136M or larger seems to remove that penalty for the tested setting. This is exactly what the asymptotic story predicts: too-small proxies may not yet be close enough to the limiting regime for transfer to be clean.

9. Assumptions and Failure Modes

The useful way to read this paper is as a set of invariance assumptions. Each assumption makes predictions, and each has practical ways to fail.

Things I Would Not Expect to Transfer Cleanly Without New Work

• Large architecture changes: dense to MoE, attention to SSM/hybrid, large context-extension changes, or major tokenizer/vocab changes.
• Different optimizer families: Lion, Adafactor, Shampoo, sign-based updates, or schedule-free optimizers need their own coordinate rules. Section 4.9 gives conditional Muon candidates, but not a completed replacement for the AdamW SDE result.
• Heavy regularization changes: dropout, stochastic depth, strong data augmentation, RL objectives, or preference-optimization losses.
• Regimes dominated by numerical precision: FP8 optimizer state, quantized training, extreme epsilon sensitivity, or loss-scaling artifacts.
• Very large batch near the critical-batch limit: the SDE/local-noise picture and the hardware-efficiency objective can pull in different directions.

10. Adjacent Work

Complete(d)P is best seen as part of a developing family of coordinate systems for training recipes. Several adjacent papers prevent overclaiming.

CompleteP

CompleteP is the direct predecessor. It focuses on width/depth transfer in transformers and argues that some parameterizations may transfer hyperparameters while still being too lazy to exploit depth. Complete(d)P inherits that narrative but modifies the table for QK norms, embedding epsilon, output projection handling, batch size, duration, and per-module tuning.

Spectral Condition for Feature Learning

The spectral-condition paper gives the cleanest non-Tensor-Programs route into muP-style reasoning: features should have O(1) coordinates, and their training-step changes should also have O(1) coordinates. Its formal statement is width-focused and hides depth in constants. When depth is also scaled, the same desideratum must be applied to the full residual stack, not to each adjacent residual increment. That is what produces the alpha >= 1/2 and alpha + gamma = 1 constraints derived above.

Tensor Programs VI / Depth-muP

The infinite-depth paper provides the sharper classification: stable initialization requires alpha >= 1/2; stable and nontrivial training requires alpha + gamma = 1; faithfulness gives alpha in [1/2, 1]; and maximal feature diversity singles out alpha = gamma = 1/2. Complete(d)P inherits this depth story but tests it in a modern transformer setting.

Malladi-Style SDE Scaling

The batch-size rules are not invented from scratch. They inherit the square-root Adam/RMSProp scaling logic from SDE approximations to adaptive optimizers. Complete(d)P's main extension here is adding AdamW weight decay and composing the batch rule with model-size and duration rules.

Muon LMO Scaling and Spectral Noise Geometry

A later optimization-theory analysis treats Muon as a spectral-norm linear-minimization-oracle method. Its fixed-momentum proxy predicts eta_mu ~ sqrt(B/T); jointly tuning momentum instead predicts 1-beta_mu ~ B T^(-1/2) and eta_mu ~ B T^(-3/4). This supplies candidate transfer coordinates but holds model size fixed and omits practical schedule and decay effects. Complementary work derives a nuclear-norm gradient-noise scale for spectral descent and shows why Euclidean variance is the wrong diagnostic for polar directions, while explicitly leaving stateful Muon as an open extension. Approximate Newton-Schulz work adds a further scale dependence because momentum singular spectra vary by layer and model size. These results motivate the Muon atlas in Section 4.9 without making it part of the original Complete(d)P claim.

u-muP

u-muP combines maximal-update parameterization with unit scaling. It is a more implementation-oriented route toward stable low-precision training dynamics. It does not try to cover as many axes as Complete(d)P, but it is a reminder that "good transfer" and "good numeric scale" are both necessary in real training.

Scaling Exponents Across Parameterizations and Optimizers

This line of work challenges a simplistic reading of muP uniqueness. It finds that multiple parameterizations can transfer if one allows the right layerwise learning-rate exponents and optimizer details, and it stresses that Adam epsilon can be a first-order issue. This supports one of Complete(d)P's themes: the optimizer coordinate system matters as much as the forward parameterization.

Weight Decay May Matter More Than muP

This critique argues that weight decay can strongly control learning-rate transfer in practice. That does not refute Complete(d)P; it makes the paper's inclusion of weight-decay transfer more important. It does warn against evaluating parameterization schemes while leaving weight decay undertuned or inconsistently scaled.

Scaling Optimal LR Across Token Horizons

This work supports the empirical fact that optimal learning rate changes with token horizon, but fitted exponents vary. Complete(d)P proposes a cleaner SDE-inspired rule; whether that rule dominates under other schedule families remains open.

11. Connection to Scaling Laws

Your earlier concern about Kaplan/Chinchilla-style scaling laws fits naturally here. A scaling law fitted under a fixed or heuristic training recipe estimates the performance of that recipe, not the globally optimized loss over all hyperparameters. If the recipe becomes more or less suboptimal with scale, the fitted law absorbs that drift.

Complete(d)P is partly a way to reduce that confound. If hyperparameters are transferred in a coordinate system where optima are stable, then a scaling experiment is less likely to confuse "larger models need different optimizer settings" with "larger models intrinsically scale this way."

But it does not eliminate the issue. It still fixes:

• a model family;
• a dataset and data order;
• an optimizer family;
• a schedule family, mostly cosine in this paper;
• a set of module groupings;
• a transfer rule for token horizon that is still empirical.

So the right conclusion is: this paper gives a better recipe for measuring scaling under a strong training family. It does not produce architecture-, optimizer-, and data-universal scaling laws.

12. Practical Recipe

A conservative implementation workflow would look like this.

Step 1: Define the Base and Target Exactly

• Use architectural width ratios, not just parameter-count ratios.
• Record whether width scales through head count, head dimension, MLP dimension, or all of them.
• Record depth, batch tokens, total tokens, sequence length, data packing, and accumulation.

Step 2: Audit Optimizer Semantics

• Confirm PyTorch AdamW vs AdamLH-style decay.
• Confirm epsilon placement and whether per-group epsilon is supported.
• For Muon groups, record momentum/Nesterov semantics, rectangular LR adjustment, Newton-Schulz coefficients and iterations, normalization epsilon, and fallback policy.
• Confirm which parameters receive weight decay.
• Confirm gradient clipping behavior and whether it is active in normal runs.
• Log transformed per-group values, not just base values.

Step 3: Tune a Strong Global Baseline

Do not start with 79 per-module variables before the global recipe is credible. Tune global LR, WD, epsilon, betas, initialization scale, residual multiplier, and schedule family at the proxy scale. The paper's baseline global random search used 2048 trials at 50M parameters, which is far more than many informal "baseline" searches.

Step 4: Add Per-Module Multipliers Locally

• Optimize in log-space.
• Start near the global optimum.
• Use trust-region or other locality-biased search.
• Separate divergent trials from merely high-loss trials.
• Prefer a factorized type x depth parameterization before fully uncoupled layer-by-layer values.

Step 5: Transfer, Then Validate at an Intermediate Scale

Even if the final goal is no retuning, a production recipe should include at least one intermediate-scale validation. The paper itself shows that too-small proxy scales can incur a penalty, while a somewhat larger proxy can stabilize the transferred LR.

Step 6: Treat Target Validation as a Guardrail, Not a New Sweep

A practical compromise is to run a very small target-scale guardrail sweep: perhaps a narrow global LR/WD multiplier check around the transferred recipe. If the transferred recipe is outside that narrow optimum region, the transfer assumptions probably failed.

13. Sources

1. Bruno Mlodozeniec, Pierre Ablin, Louis Bethune, Dan Busbridge, Michal Klein, Jason Ramapuram, Marco Cuturi. Completed Hyperparameter Transfer across Modules, Width, Depth, Batch & Duration. arxiv abs, arxiv HTML.
2. Greg Yang et al. Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. NeurIPS PDF.
3. Greg Yang, James B. Simon, Jeremy Bernstein. A Spectral Condition for Feature Learning. arxiv abs, arxiv PDF.
4. Greg Yang, Dingli Yu, Chen Zhu, Soufiane Hayou. Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks. arxiv abs, arxiv PDF.
5. Dey et al. CompleteP / complete parameterization for width-depth hyperparameter transfer. arxiv abs, arxiv HTML.
6. Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, Sanjeev Arora. On the SDEs and Scaling Rules for Adaptive Gradient Algorithms. arxiv abs.
7. Charlie Blake et al. u-muP: The Unit-Scaled Maximal Update Parametrization. arxiv abs, arxiv HTML.
8. Katie Everett et al. Scaling Exponents Across Parameterizations and Optimizers. arxiv abs.
9. Kosson, Welborn, Liu, Jaggi, Chen. Weight Decay may matter more than muP for Learning Rate Transfer in Practice. arxiv abs, arxiv HTML.
10. Scaling Optimal LR Across Token Horizons. arxiv abs, arxiv HTML.
11. Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, Volkan Cevher. Training Deep Learning Models with Norm-Constrained LMOs. arxiv abs, arxiv HTML.
12. Egor Shulgin, Dimitri von Rutte, Tianyue H. Zhang, Niccolo Ajroldi, Bernhard Scholkopf, Antonio Orvieto. Deriving Hyperparameter Scaling Laws via Modern Optimization Theory. arxiv abs, arxiv HTML.
13. Hiroki Naganuma et al. Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent. arxiv abs, arxiv HTML.
14. Gagik Magakyan, Pablo Parrilo, Asuman Ozdaglar. Spectral Scaling Laws of Muon. arxiv abs.
15. Hoffmann et al. Training Compute-Optimal Large Language Models. arxiv abs.
16. Kaplan et al. Scaling Laws for Neural Language Models. arxiv abs.