Beyond MuP: 1. Three Characteristics of a Good Model

MuP之上：1. 好模型的三个特征 苏剑林 · 2025-10-21 · Machine Learning optimization optimizer scaling law MuP

I wonder if anyone else has noticed an interesting detail: both Muon and MuP start with "Mu," but the two "Mu"s have completely different origins — the former is "MomentUm Orthogonalized by Newton-Schulz," the latter is "Maximal Update Parametrization." Yet the two are indeed deeply connected. In other words, Muon and MuP started from entirely different premises but ultimately converged on the same direction, and even unintentionally ended up with similar names — as if it were truly "meant to be."

Anyway, back to business. Through a series of fortunate coincidences, I happened to learn about Muon and MuP around the same time, which greatly deepened my understanding of model optimization and got me thinking about more fundamental principles underlying it. After some trial and error, I've gathered a few modest insights, which I'd like to share here.

Preface

In terms of chronological order, MuP came before Muon, but my own learning order was the reverse — I learned Muon first, then MuP. Looking back, this turned out to be a rather good order to learn them in.

In Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices, Muon Sequel: Why Did We Choose to Try Muon?, and other posts, we described Muon as "steepest descent under a spectral norm constraint," while the MuP line of work happens to explain "why we need a spectral norm constraint" — the two fit together perfectly.

I should clarify here that when we say "MuP," it actually has two meanings: first, the one introduced in A First Look at MuP: Cross-Model-Scale Transfer Laws for Hyperparameters, which belongs to the Tensor Programs line of work — we'll call this "elementary MuP"; second, the one introduced in Higher-Order MuP: A Simpler yet Cleverer Spectral-Condition Scaling, which we'll call "higher-order MuP." The latter arrives at richer conclusions than elementary MuP via a more concise route — both are the work of Greg Yang (hats off to him).

Unless otherwise noted, "MuP" in this post refers to "higher-order MuP." In fact, this series, which I'm calling "Beyond MuP," consists of a series of reflections and extensions built on top of higher-order MuP. However, some readers may only be familiar with the "elementary MuP" from the Tensor Programs series, so at first glance they might wonder how MuP could possibly answer the question of "why we need a spectral norm."

In any case, I'll try to make this series self-contained as much as possible. So although we'll mention many related papers and blog posts along the way, readers don't need to read all of them in depth.

Steady Yet Fast

Back to the main thread. As the first post in this series, the task here is to nail down the core objective — more specifically, to think through "what kind of model do we actually want" and "how do we go about training such a model."

Intuitively, as long as a model shows no signs of collapse, we can just keep training it until it converges to a result we're satisfied with; on top of that, we naturally try to find ways to make it converge faster. So in the end, it all comes down to two things: "stability" and "speed" — or, put another way, "seeking speed while staying stable." So how do we judge whether a model is stable? This inevitably requires monitoring various "internal medicine metrics" — the more we monitor, the more likely we are to expose problems.

That said, this post doesn't intend to enumerate every possible internal metric, but instead tries to identify the most essential, or most necessary, few conditions. To this end, let's first define a concept — RMS (Root Mean Square): given $\boldsymbol{x}=(x_1,x_2,\cdots,x_d)\in\mathbb{R}^d$, we define

\begin{equation}\Vert\boldsymbol{x}\Vert_{RMS} = \sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2} = \frac{\Vert\boldsymbol{x}\Vert_2}{\sqrt{d}}\end{equation}

This represents the average scale per element, differing from the vector's norm $\Vert\boldsymbol{x}\Vert_2$ by a factor of $\sqrt{d}$.

Some readers might ask: since it only differs by a constant factor, why not just look at the norm directly instead of defining a new concept? There are a few considerations here — for instance, RMSNorm is commonly used, and RMS is easier to intuit than the norm. But there's also an important reason: most activation functions operate element-wise, so we need to examine and control the average scale per element, in order to ensure that activation functions play a similar role across different models.

Three Conditions

With the notation for RMS in hand, we can now write down what I consider the three most essential conditions for stably training a good model:

\begin{align} &\text{forward stability:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c1}\\[5pt] &\text{dependency stability:}\quad\max_{\boldsymbol{x}_1,\boldsymbol{x}_2} \frac{\Vert \boldsymbol{f}(\boldsymbol{x}_1;\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x}_2;\boldsymbol{\omega})\Vert_{RMS}}{\Vert\boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{RMS}} = \Theta(1) \label{eq:c2}\\[5pt] &\text{update stability:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega} + \Delta\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c3} \end{align}

Here $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$ represents a family of models parametrized by $\mathbb{R}^{d_{in}}\mapsto \mathbb{R}^{d_{out}}$, with input $\boldsymbol{x}\in\mathbb{R}^{d_{in}}$ and output $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\in\mathbb{R}^{d_{out}}$, and $\boldsymbol{\omega}$ is the model parameter, which could be a scalar, vector, matrix, etc. $\Theta$ is "Big Theta Notation." Here, $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$ can be a single layer, a block composed of several layers, or even the entire model — in principle, the coarser the granularity, the looser (or, one might say, the more accurate) the resulting constraint, but the harder it is to solve for $\max$. So this depends on how well we can actually compute $\max$.

Note that we just said $\boldsymbol{x}\in\mathbb{R}^{d_{in}}$, but it need not span the entire $\mathbb{R}^{d_{in}}$. For example, the input to an embedding layer is a token ID, or more precisely, a one-hot vector in $\mathbb{R}^{|V|}$ (where $|V|$ is the vocabulary size), and it can only be a one-hot vector — it can never fill out the whole space $\mathbb{R}^{|V|}$. So when computing $\max$, we should only take it over these one-hot inputs; computing it over the entire space would actually be inaccurate.

Analyzing Each Condition

Among the three equations, equation $\eqref{eq:c1}$ is probably the easiest to understand — it represents the stability of the forward computation. After taking $\max$ over $\boldsymbol{x}$, the only remaining variable is $\boldsymbol{\omega}$, so this is a constraint on $\boldsymbol{\omega}$. Note that, in general, we don't restrict the range of $\boldsymbol{x}$, so by default $\boldsymbol{x}\in\mathbb{R}^{d_{in}}$, meaning the maximum might not even exist — for instance, for nonzero $\boldsymbol{W}$ we have $\max\limits_{\boldsymbol{x}}\Vert \boldsymbol{x}\boldsymbol{W}\Vert_{RMS}\to\infty$. To guarantee the existence of a maximum, we typically need to add some normalization operation, such as:

\begin{align} &\text{In Norm:}\quad \mathop{\text{Norm}}(\boldsymbol{x})\boldsymbol{W} \\[5pt] &\text{Out Norm:}\quad \mathop{\text{Norm}}(\boldsymbol{x}\boldsymbol{W}) \end{align}

where $\mathop{\text{Norm}}(\boldsymbol{x})=\boldsymbol{x}/\Vert\boldsymbol{x}\Vert_{RMS}$, i.e., RMS Norm — though other kinds of normalization could also be considered. So condition $\eqref{eq:c1}$ also implicitly imposes certain requirements on the model architecture.

The last equation, $\eqref{eq:c3}$, should also be fairly easy to understand — it measures how much a change in the parameters affects the output. That is, we want every parameter update to produce a correspondingly smooth change in the model's predictions. Since parameter changes only occur during training, this represents our expectation of training stability, and we can use it to guide the choice of optimizer hyperparameters, or even to design new optimizers based directly on it.

Specifically, after taking $\max$ over $\boldsymbol{x}$, the result of equation $\eqref{eq:c3}$ becomes a constraint on $\boldsymbol{\omega}$ and $\Delta\boldsymbol{\omega}$. Compared with $\eqref{eq:c1}$, it mainly introduces one extra term, $\Delta\boldsymbol{\omega}$, so what it primarily constrains is the increment $\Delta\boldsymbol{\omega}$. As we'll see later, this constraint can be used to tune the initialization, learning rate, and so on, or even combined directly with the "steepest descent" idea to derive the form of $\Delta\boldsymbol{\omega}$.

Stable Dependence

Condition $\eqref{eq:c2}$ deserves a dedicated discussion of its own — it describes the degree to which $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$ depends on the input $\boldsymbol{x}$.

We want that when the input changes by $\Delta\boldsymbol{x} = \boldsymbol{x}_2 - \boldsymbol{x}_1$, the output also changes by roughly the same order of magnitude, $\Theta(\Vert\Delta\boldsymbol{x}\Vert_{RMS})$. This carries two implications: first, the output change should be at most on the order of $\Vert\Delta\boldsymbol{x}\Vert_{RMS}$, which requires the model to be smooth in some sense — it shouldn't produce sudden, huge jumps; second, the output change should be at least on the order of $\Vert\Delta\boldsymbol{x}\Vert_{RMS}$ (for at least some inputs), meaning the model genuinely depends on $\boldsymbol{x}$ rather than treating it as irrelevant.

The first implication is probably easy to grasp; as for the second, consider a simple example: $f(x;\omega)=x\times\omega\times 0 + 1$. This "model" is certainly stable in the forward pass (it's always equal to 1), but it doesn't depend on $x$ at all, so condition $\eqref{eq:c2}$ can't be satisfied — meaning it's not a good model. Put simply, dependence stability means a good model should reasonably depend on its input.

On the other hand, one can show that dependence stability actually also covers the stability of backpropagation. When $\Delta\boldsymbol{x}\to \boldsymbol{0}$, we have the approximate expansion $\boldsymbol{f}(\boldsymbol{x} + \Delta\boldsymbol{x};\boldsymbol{\omega}) \approx \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega}) + \Delta\boldsymbol{x}\cdot\nabla_{\boldsymbol{x}}\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$, and therefore

\begin{equation}\frac{\Vert\boldsymbol{f}(\boldsymbol{x} + \Delta\boldsymbol{x};\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS}}{\Vert\Delta\boldsymbol{x}\Vert_{RMS}} \approx \frac{\Vert\Delta\boldsymbol{x}\cdot\nabla_{\boldsymbol{x}}\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS}}{\Vert\Delta\boldsymbol{x}\Vert_{RMS}} \leq \sqrt{\frac{d_{in}}{d_{out}}}\Vert\nabla_{\boldsymbol{x}}\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_2\end{equation}

where $\nabla_{\boldsymbol{x}}\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$ denotes the Jacobian matrix of $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$, and $\Vert\nabla_{\boldsymbol{x}}\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_2$ is its spectral norm. This shows that condition $\eqref{eq:c2}$ also imposes a requirement on the spectral norm of the model's Jacobian, and since the Jacobian is exactly what drives gradient computation via the chain rule during backpropagation, this condition implicitly also demands stability of backpropagation.

Some Reflections

To sum up, the three conditions in equation $\eqref{eq:c1},\eqref{eq:c2},\eqref{eq:c3}$ jointly incorporate considerations from model architecture, initialization, optimization, and backpropagation. It's hard to argue that any one of them could be dropped, so I believe all three are necessary. That said, there are still some details worth discussing about these three conditions — for instance, the choice between $\max$ and $\mathbb{E}$.

In the equations as written, we "eliminated" $\boldsymbol{x}$ by taking $\max$, leaving equations that only involve $\boldsymbol{\omega}$ and $\Delta\boldsymbol{\omega}$. This might raise a question for some readers, since a more intuitive approach might be to take the mathematical expectation $\mathbb{E}_{\boldsymbol{x}}$ instead. Why $\max$ and not $\mathbb{E}$? There are several reasons. First, computing $\max$ only requires specifying the domain of $\boldsymbol{x}$, whereas computing $\mathbb{E}$ requires specifying the distribution of $\boldsymbol{x}$ — different distributions give different results, and pinning down this distribution accurately is far from trivial.

Second, $\max$ has the advantage of being invariant under monotonic transformations, a property $\mathbb{E}$ lacks. For example, for $\max$ we have the identity $(\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS})^2 = \max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS}^2$, meaning that taking $\max$ of either $\Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS}$ or $\Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS}^2$ is essentially equivalent. But this is not the case for $\mathbb{E}$: the expectation of $\Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS}$ and that of $\Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS}^2$ generally differ in computational difficulty, and their results need not be related at all.

So $\max$ is simpler both conceptually and in terms of its properties. One possible concern is whether $\max$ is too strict — akin to a "sufficient but not necessary" condition? Actually, $\max$ is just an informal way of putting it; mathematically it's called the "supremum ($\sup$)," where "supremum" indicates that this value is tight and attainable. In practice, though, the mean and the maximum are usually of the same order of magnitude, and since our goal is only $\Theta(1)$, the distinction rarely matters much. In contrast, $\max$ accounts for extreme cases too, maximizing the guarantee of training stability, which is especially important for training large models such as LLMs.

In fact, elementary MuP — that is, the Tensor Programs line of work — is a series of analyses built on $\mathbb{E}$, whereas higher-order MuP, like the present post, is built on $\max$. In hindsight, analyses based on $\mathbb{E}$ fall short of higher-order MuP based on $\max$ in terms of computational simplicity and generality of the results — which in turn corroborates the effectiveness of $\max$.

Summary

Starting from this post, I'll be sharing some top-down understanding of model optimization, extending and building on the "higher-order MuP" framework covered previously. As the first post in the series, we've mainly laid out three fundamental conditions for model stability — or, put another way, three characteristics of a good model — which will serve as the cornerstone for the calculations and analyses to come.