Higher-Order MuP: A More Concise yet More Sophisticated Spectral Condition Scaling

Preliminaries

As the name suggests, the Spectral Condition is related to the Spectral Norm, and its starting point is a basic inequality of the spectral norm: \[\begin{equation} \Vert\boldsymbol{x}\boldsymbol{W}\Vert_2\leq \Vert\boldsymbol{x}\Vert_2 \Vert\boldsymbol{W}\Vert_2\label{neq:spec-2} \end{equation}\] where \(\boldsymbol{x}\in\mathbb{R}^{d_{in}}, \boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}\). As for \(\Vert\cdot\Vert_2\), we may call it the “\(2\)-norm.” For \(\boldsymbol{x}\) and \(\boldsymbol{x}\boldsymbol{W}\), which are both vectors, the \(2\)-norm is the vector magnitude; \(\boldsymbol{W}\) is a matrix, and its \(2\)-norm is also called the spectral norm, which equals the smallest constant \(C\) such that \(\Vert\boldsymbol{x}\boldsymbol{W}\Vert_2\leq C\Vert\boldsymbol{x}\Vert_2\) always holds. In other words, the above inequality is a direct corollary of the definition of the spectral norm and requires no additional proof.

For more on the spectral norm, readers may consult posts such as Lipschitz Constraints in Deep Learning: Generalization and Generative Models and The Road to Low-Rank Approximation (II): SVD; we will not elaborate here. A matrix also has a simpler \(F\)-norm, which is a straightforward generalization of vector magnitude: \[\begin{equation} \Vert \boldsymbol{W}\Vert_F = \sqrt{\sum_{i=1}^{d_{in}}\sum_{j=1}^{d_{out}}W_{i,j}^2} \end{equation}\] From the perspective of singular values, the spectral norm equals the largest singular value of the matrix, while the \(F\)-norm equals the square root of the sum of squares of all singular values. Similarly, we can define the “Nuclear Norm,” which equals the sum of all singular values: \[\begin{equation} \Vert \boldsymbol{W}\Vert_* = \sum_{i=1}^{\min(d_{in}, d_{out})} \sigma_i \end{equation}\] Matrix norms that can be expressed via singular values, such as the spectral norm, the \(F\)-norm, and the nuclear norm, are all instances of the Schatten-\(p\) norms. Finally, let us define the RMS (Root Mean Square), a variant of vector magnitude: \[\begin{equation} \Vert\boldsymbol{x}\Vert_{RMS} = \sqrt{\frac{1}{d_{in}}\sum_{i=1}^{d_{in}} x_i^2} = \frac{1}{\sqrt{d_{in}}}\Vert \boldsymbol{x}\Vert_2 \end{equation}\] To generalize it to matrices, we have \(\Vert\boldsymbol{W}\Vert_{RMS} = \Vert \boldsymbol{W}\Vert_F/\sqrt{d_{in} d_{out}}\). In fact, the name is self-explanatory: vector magnitude or matrix \(F\)-norm can be called “Root Sum Square,” and RMS simply replaces Sum with Mean; it mainly serves as an indicator of the average scale of vector or matrix entries. Now, substituting RMS into inequality \(\eqref{neq:spec-2}\), we obtain \[\begin{equation} \Vert\boldsymbol{x}\boldsymbol{W}\Vert_{RMS}\leq \sqrt{\frac{d_{in}}{d_{out}}}\Vert\boldsymbol{x}\Vert_{RMS} \Vert\boldsymbol{W}\Vert_2\label{neq:spec-rms} \end{equation}\]

Desired Properties

Our earlier approach to deriving MuP was to carefully analyze the forms of forward propagation, backpropagation, the loss increment, and feature change, and to adjust initialization and learning rates to achieve their scale invariance. The Spectral Condition, after distilling the essence, found that only two points—forward propagation and feature change—are needed.

Put simply, the Spectral Condition expects that the output and the increment of every layer are scale-invariant. How should we understand this statement? If we briefly denote each layer as \(\boldsymbol{x}_k= f(\boldsymbol{x}_{k-1}; \boldsymbol{W}_k)\), then this statement can be translated as “we expect every \(\Vert\boldsymbol{x}_k\Vert_{RMS}\) and \(\Vert\Delta\boldsymbol{x}_k\Vert_{RMS}\) to be \(\mathcal{\Theta}(1)\)” (where \(\mathcal{\Theta}\) is Big Theta Notation):

1. \(\Vert\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\) is easy to understand: it represents the stability of forward propagation, and the derivation in the previous article also included this requirement;

2. \(\Delta\boldsymbol{x}_k\) denotes the change in \(\boldsymbol{x}_k\) caused by parameter changes, so \(\Vert\Delta\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\) merges the requirements of backpropagation and feature change.

Some readers may wonder: shouldn’t there at least also be a “loss increment” requirement? Not necessary. In fact, we can prove that if every layer’s \(\Vert\boldsymbol{x}_k\Vert_{RMS}\) and \(\Vert\Delta\boldsymbol{x}_k\Vert_{RMS}\) are \(\mathcal{\Theta}(1)\), then \(\Delta\mathcal{L}\) is automatically \(\mathcal{\Theta}(1)\). This is the first beautiful aspect of the Spectral Condition idea: it reduces the four conditions needed in the original MuP derivation down to two, reducing the number of analytical steps.

The proof is not difficult. The key here is that we assume every layer satisfies \(\Vert\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\) and \(\Vert\Delta\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\), so the last layer naturally satisfies them too. Suppose the model has \(K\) layers in total, and the per-sample loss function is \(\ell\); then it is a function of \(\boldsymbol{x}_K\), i.e., \(\ell(\boldsymbol{x}_K)\). For simplicity we omit the label input here, since it is not a variable for the analysis below.

By assumption, \(\Vert\boldsymbol{x}_K\Vert_{RMS}\) is \(\mathcal{\Theta}(1)\), so \(\ell(\boldsymbol{x}_K)\) is naturally \(\mathcal{\Theta}(1)\); and since \(\Vert\Delta\boldsymbol{x}_K\Vert_{RMS}\) is \(\mathcal{\Theta}(1)\), it follows that \(\Vert\boldsymbol{x}_K + \Delta\boldsymbol{x}_K\Vert_{RMS}\leq \Vert\boldsymbol{x}_K\Vert_{RMS} + \Vert\Delta\boldsymbol{x}_K\Vert_{RMS}\) is also \(\mathcal{\Theta}(1)\), and hence \(\ell(\boldsymbol{x}_K + \Delta\boldsymbol{x}_K)\) is \(\mathcal{\Theta}(1)\). Therefore \[\begin{equation} \Delta \ell = \ell(\boldsymbol{x}_K + \Delta\boldsymbol{x}_K) - \ell(\boldsymbol{x}_K) = \mathcal{\Theta}(1) \end{equation}\] So the per-sample loss increment \(\Delta \ell\) is \(\mathcal{\Theta}(1)\), and \(\Delta\mathcal{L}\) is the average of all the \(\Delta \ell\), so it too is \(\mathcal{\Theta}(1)\). This proves that \(\Vert\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\) and \(\Vert\Delta\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\) automatically imply \(\Delta\mathcal{L}=\mathcal{\Theta}(1)\). The principle, in plain words, is that \(\Delta\mathcal{L}\) is a function of the last layer’s output and its increment; once those are stable, \(\Delta\mathcal{L}\) is naturally stable.

The Spectral Condition

Next, we look at how to satisfy the two desired properties. Since neural networks are dominated by matrix multiplication, we first consider the simplest linear layer \(\boldsymbol{x}_k = \boldsymbol{x}_{k-1} \boldsymbol{W}_k\), where \(\boldsymbol{W}_k\in\mathbb{R}^{d_{k-1}\times d_k}\). To satisfy the condition \(\Vert\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\), the Spectral Condition does not follow the traditional initialization analysis of assuming i.i.d. distributions and computing expectations and variances; instead it directly applies inequality \(\eqref{neq:spec-rms}\): \[\begin{equation} \Vert\boldsymbol{x}_k\Vert_{RMS}\leq \sqrt{\frac{d_{k-1}}{d_k}}\Vert\boldsymbol{x}_{k-1}\Vert_{RMS}\, \Vert\boldsymbol{W}_k\Vert_2 \end{equation}\] Note that this inequality can attain equality, and in a certain sense it is the tightest possible. So if the input \(\Vert\boldsymbol{x}_{k-1}\Vert_{RMS}\) is already \(\mathcal{\Theta}(1)\), then to make the output \(\Vert\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\), we need \[\begin{equation} \sqrt{\frac{d_{k-1}}{d_k}}\Vert\boldsymbol{W}_k\Vert_2 = \mathcal{\Theta}(1)\quad\Rightarrow\quad \Vert\boldsymbol{W}_k\Vert_2 = \mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\label{eq:spec-c1} \end{equation}\] This gives the first Spectral Condition—a requirement on the spectral norm of \(\boldsymbol{W}_k\). It is unrelated to initialization and distributional assumptions; it is purely the result of analysis and algebra. This is, in the author’s view, the second beautiful aspect of the Spectral Condition—it simplifies the analysis. Of course, we have omitted the basics of spectral norms here; including them, the total length might not be shorter than the analysis under distributional assumptions, but distributional assumptions are ultimately quite limited in scope and less flexible than the algebraic framework here.

After analyzing \(\Vert\boldsymbol{x}_k\Vert_{RMS}\), we turn to \(\Vert\Delta\boldsymbol{x}_k\Vert_{RMS}\). The increment \(\Delta\boldsymbol{x}_k\) has two sources: first, the parameter changes from \(\boldsymbol{W}_k\) to \(\boldsymbol{W}_k+\Delta \boldsymbol{W}_k\); second, the input \(\boldsymbol{x}_{k-1}\) changes to \(\boldsymbol{x}_{k-1} + \Delta\boldsymbol{x}_{k-1}\) due to parameter changes earlier. Hence \[\begin{equation} \begin{aligned} \Delta\boldsymbol{x}_k =&\, (\boldsymbol{x}_{k-1} + \Delta\boldsymbol{x}_{k-1})(\boldsymbol{W}_k+\Delta \boldsymbol{W}_k) - \boldsymbol{x}_{k-1}\boldsymbol{W}_k \\[5pt] =&\, \boldsymbol{x}_{k-1} (\Delta \boldsymbol{W}_k) + (\Delta\boldsymbol{x}_{k-1})\boldsymbol{W}_k + (\Delta\boldsymbol{x}_{k-1})(\Delta \boldsymbol{W}_k) \end{aligned} \end{equation}\] So \[\begin{equation} \begin{aligned} \Vert\Delta\boldsymbol{x}_k\Vert_{RMS} =&\, \Vert\boldsymbol{x}_{k-1} (\Delta \boldsymbol{W}_k) + (\Delta\boldsymbol{x}_{k-1})\boldsymbol{W}_k + (\Delta\boldsymbol{x}_{k-1})(\Delta \boldsymbol{W}_k)\Vert_{RMS} \\[5pt] \leq&\, \Vert\boldsymbol{x}_{k-1} (\Delta \boldsymbol{W}_k)\Vert_{RMS} + \Vert(\Delta\boldsymbol{x}_{k-1})\boldsymbol{W}_k\Vert_{RMS} + \Vert(\Delta\boldsymbol{x}_{k-1})(\Delta \boldsymbol{W}_k)\Vert_{RMS} \\[5pt] \leq&\, \sqrt{\frac{d_{k-1}}{d_k}}\left({\begin{gathered}\Vert\boldsymbol{x}_{k-1}\Vert_{RMS}\,\Vert\Delta \boldsymbol{W}_k\Vert_2 + \Vert\Delta\boldsymbol{x}_{k-1}\Vert_{RMS}\,\Vert \boldsymbol{W}_k\Vert_2 \\[5pt] + \Vert\Delta\boldsymbol{x}_{k-1}\Vert_{RMS}\,\Vert\Delta \boldsymbol{W}_k\Vert_2\end{gathered}} \right) \end{aligned} \end{equation}\] Analyzing term by term: \[\begin{equation} \underbrace{\Vert\boldsymbol{x}_{k-1}\Vert_{RMS}}_{\mathcal{\Theta}(1)}\,\Vert\Delta \boldsymbol{W}_k\Vert_2 + \underbrace{\Vert\Delta\boldsymbol{x}_{k-1}\Vert_{RMS}}_{\mathcal{\Theta}(1)}\,\underbrace{\Vert \boldsymbol{W}_k\Vert_2}_{\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)} + \underbrace{\Vert\Delta\boldsymbol{x}_{k-1}\Vert_{RMS}}_{\mathcal{\Theta}(1)}\,\Vert\Delta \boldsymbol{W}_k\Vert_2 \end{equation}\] From this we see that for \(\Vert\Delta\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1)\), we need \[\begin{equation} \Vert\Delta\boldsymbol{W}_k\Vert_2 = \mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\label{eq:spec-c2} \end{equation}\] This is the second Spectral Condition—a requirement on the spectral norm of \(\Delta\boldsymbol{W}_k\).

The above analysis did not consider nonlinearity. In fact, as long as the activation function is element-wise and its derivative is bounded by some constant (common activations such as ReLU, Sigmoid, and Tanh all satisfy this), the result is the same even when nonlinear activation functions are considered. This is what the analysis in the previous article described as “the influence of the activation function is scale-irrelevant.” If readers are still uneasy, they may derive it themselves.

Spectral Normalization

Now we have two spectral conditions \(\eqref{eq:spec-c1}\) and \(\eqref{eq:spec-c2}\); next we need to see how to design the model itself and its optimization so that both conditions are satisfied.

Note that \(\boldsymbol{W}_k\) and \(\Delta \boldsymbol{W}_k\) are both matrices. The standard way to make a matrix satisfy a spectral norm condition is usually Spectral Normalization (SN), and this case is no exception. First, to make the initialized \(\boldsymbol{W}_k\) satisfy \(\Vert\boldsymbol{W}_k\Vert_2=\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\), we can pick any initialization matrix \(\boldsymbol{W}_k'\) and then apply spectral normalization: \[\begin{equation} \boldsymbol{W}_k = \sigma\sqrt{\frac{d_k}{d_{k-1}}}\frac{\boldsymbol{W}_k'}{\Vert\boldsymbol{W}_k'\Vert_2} \end{equation}\] Here \(\sigma > 0\) is a scale-invariant constant. Similarly, for the update \(\boldsymbol{\Phi}_k\) given by any optimizer, we can reconstruct \(\Delta \boldsymbol{W}_k\) via spectral normalization: \[\begin{equation} \Delta \boldsymbol{W}_k = \eta\sqrt{\frac{d_k}{d_{k-1}}}\frac{\boldsymbol{\Phi}_k}{\Vert\boldsymbol{\Phi}_k\Vert_2} \end{equation}\] where \(\eta > 0\) is also a scale-invariant constant (the learning rate). In this way, every step satisfies \(\Vert\Delta\boldsymbol{W}_k\Vert_2=\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\). Since the spectral norms of the initialization and of every update both satisfy \(\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\), \(\Vert\boldsymbol{W}_k\Vert_2\) satisfies \(\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\) throughout, thereby satisfying both spectral conditions.

At this point some readers may ask: considering only the stability of initialization and increments, can that really guarantee the stability of \(\boldsymbol{W}_k\)? Could it not happen that \(\Vert\boldsymbol{W}_k\Vert_{RMS}\to\infty\)? The answer is: yes, it could. The \(\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\) here emphasizes the relationship with the model scale (currently mainly the width); it does not rule out the possibility of training collapse caused by other ill-suited hyperparameter settings. What it expresses is that with this setup, even if collapse occurs, the cause is unrelated to the change of scale.

Singular Value Clipping

To implement the spectral norm conditions, besides the standard method of spectral normalization, we can also consider Singular Value Clipping (referred to below as “SVC”). This section is the author’s own supplement and does not appear in the original paper, but it can explain some interesting results.

From the perspective of singular values, spectral normalization rescales the largest singular value to \(1\) and rescales all other singular values proportionally. Singular value clipping is in a sense more lenient: it only sets singular values greater than \(1\) to \(1\), but leaves those already less than or equal to \(1\) unchanged: \[\begin{equation} \mathop{\text{SVC}}(\boldsymbol{W}) = \boldsymbol{U}\min(\boldsymbol{\Lambda},1)\boldsymbol{V}^{\top},\qquad \boldsymbol{U},\boldsymbol{\Lambda},\boldsymbol{V}^{\top} = \mathop{\text{SVD}}(\boldsymbol{W}) \end{equation}\] By comparison, spectral normalization is \(\mathop{\text{SN}}(\boldsymbol{W})=\boldsymbol{U}(\boldsymbol{\Lambda}/\max(\boldsymbol{\Lambda}))\boldsymbol{V}^{\top}\). Replacing spectral normalization with singular value clipping, we obtain \[\begin{equation} \boldsymbol{W}_k = \sigma\sqrt{\frac{d_k}{d_{k-1}}}\mathop{\text{SVC}}(\boldsymbol{W}_k'), \qquad \Delta \boldsymbol{W}_k = \eta\sqrt{\frac{d_k}{d_{k-1}}}\mathop{\text{SVC}}(\boldsymbol{\Phi}_k) \end{equation}\] A shortcoming of singular value clipping is that it only guarantees the clipped spectral norm equals \(1\) when at least one singular value is greater than or equal to \(1\). If that is not satisfied, we can consider multiplying by some \(\lambda > 0\) and then clipping, i.e., using \(\mathop{\text{SVC}}(\lambda\boldsymbol{W})\) instead. However, different scaling factors yield different results, and it is not easy to determine an appropriate one. But we can consider a limiting version: \[\begin{equation} \lim_{\lambda\to\infty} \mathop{\text{SVC}}(\lambda\boldsymbol{W}) = \mathop{\text{msign}}(\boldsymbol{W}) \end{equation}\] Here \(\mathop{\text{msign}}\) is exactly the matrix version of the sign function, \(\mathop{\text{msign}}\), from Muon (see Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices). Replacing spectral normalization or singular value clipping with \(\mathop{\text{msign}}\), we get \[\begin{equation} \Delta \boldsymbol{W}_k = \eta\sqrt{\frac{d_k}{d_{k-1}}}\mathop{\text{msign}}(\boldsymbol{\Phi}_k) \end{equation}\] In this way we have actually obtained a generalized Muon optimizer: standard Muon applies \(\mathop{\text{msign}}\) to the momentum, whereas this allows us to apply \(\mathop{\text{msign}}\) to the update produced by any existing optimizer. Coincidentally, some time ago someone on Twitter actually ran an experiment applying \(\mathop{\text{msign}}\) to the Adam update (they called it “Mudamw,” link), and found the results slightly better than Muon, as shown below:

After seeing this, we tried it on small models ourselves and found we could reproduce a similar conclusion! So it is quite possible that applying \(\mathop{\text{msign}}\) to the updates of existing optimizers can yield better results. The feasibility of such an operation would be very hard to explain under the original Muon framework, but here, by understanding it as singular value clipping (in its limiting form) applied to the update, this result follows naturally.

Approximate Estimation

It is generally believed that operations related to SVD (singular value decomposition), such as spectral normalization, singular value clipping, or \(\mathop{\text{msign}}\), are rather expensive, so we would still like to find simpler forms. Since our goal is only to find scaling laws across model scales, further simplification is indeed possible.

(Note: in fact, our Moonlight work shows that, if implemented well, even applying \(\mathop{\text{msign}}\) at every update step incurs very limited extra cost; so the content of this section, at present, is more about exploring explicit scaling laws than about saving computation.)

First, initialization. Initialization is a one-time thing, so a slightly larger computation cost is not really a problem; thus the earlier scheme of random initialization followed by spectral normalization / singular value clipping / \(\mathop{\text{msign}}\) can be retained. If one still wants to perfect it further, one can use a statistical result: a \(d_{k-1}\times d_k\) matrix independently sampled from a standard normal distribution has a largest singular value of approximately \(\sqrt{d_{k-1}} + \sqrt{d_k}\). This means that if we change the sampling standard deviation to \[\begin{equation} \sigma_k = \mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}(\sqrt{d_{k-1}} + \sqrt{d_k})^{-1}\right) = \mathcal{\Theta}\left(\sqrt{\frac{1}{d_{k-1}}\min\left(1, \frac{d_k}{d_{k-1}}\right)}\right) \label{eq:spec-std} \end{equation}\] then the requirement \(\Vert\boldsymbol{W}_k\Vert_2=\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\) can be satisfied at the initialization stage. For a proof of this statistical result, readers may consult High-Dimensional Probability and the Marchenko–Pastur law; we will not expand on it here.

Next we consider the update, which is relatively more troublesome, because the spectral norm of an arbitrary update \(\boldsymbol{\Phi}_k\) is not so easy to estimate. Here we need to use an empirical conclusion: the gradient matrices of parameters are usually low-rank. This low-rank property is not necessarily mathematically exact low rank; it means that the largest few singular values (their number being unrelated to the model scale) are significantly larger than the rest, making low-rank approximation applicable. This is also the theoretical basis for various LoRA optimizations.

A direct corollary of this empirical assumption is the approximate equivalence of the spectral norm and the nuclear norm: the spectral norm is the largest singular value, the nuclear norm is the sum of all singular values, and under the above assumption the nuclear norm is approximately the sum of the largest few singular values, so the two are at least of the same order, i.e., \(\mathcal{\Theta}(\Vert\nabla_{\boldsymbol{W}_k}\mathcal{L}\Vert_2)=\mathcal{\Theta}(\Vert\nabla_{\boldsymbol{W}_k}\mathcal{L}\Vert_*)\). Next we use the relationship between \(\Delta\mathcal{L}\) and \(\Delta\boldsymbol{W}_k\): \[\begin{equation} \Delta\mathcal{L} \approx \sum_k \langle \Delta\boldsymbol{W}_k, \nabla_{\boldsymbol{W}_k}\mathcal{L}\rangle_F \leq \sum_k \Vert\Delta\boldsymbol{W}_k\Vert_2\, \Vert\nabla_{\boldsymbol{W}_k}\mathcal{L}\Vert_* \end{equation}\] Here \(\langle\cdot,\cdot\rangle_F\) is the \(F\)-inner product, i.e., the inner product treating flattened matrices as vectors; as for the inequality sign, it is due to the classical matrix norm inequality \(\langle\boldsymbol{A},\boldsymbol{B}\rangle_F \leq \Vert\boldsymbol{A}\Vert_2\, \Vert\boldsymbol{B}\Vert_*\), which is analogous to Hölder’s inequality; in fact, it is exactly what we proved in Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices when deriving Muon. Based on the above and combined with the low-rank assumption on gradients, we have \[\begin{equation} \Delta\mathcal{L} \sim \sum_k \mathcal{\Theta}(\Vert\Delta\boldsymbol{W}_k\Vert_2\, \Vert\nabla_{\boldsymbol{W}_k}\mathcal{L}\Vert_*) \sim \sum_k \mathcal{\Theta}(\Vert\Delta\boldsymbol{W}_k\Vert_2\, \Vert\nabla_{\boldsymbol{W}_k}\mathcal{L}\Vert_2) \end{equation}\] Remember that we already proved earlier that, under the two spectral conditions, \(\Delta\mathcal{L}=\mathcal{\Theta}(1)\) necessarily holds. Combining this with the above, we find that when \(\Vert\Delta\boldsymbol{W}_k\Vert_2=\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\) we have \[\begin{equation} \mathcal{\Theta}(\Vert\nabla_{\boldsymbol{W}_k}\mathcal{L}\Vert_2) = \mathcal{\Theta}\left(\sqrt{\frac{d_{k-1}}{d_k}}\right)\label{eq:grad-norm} \end{equation}\] This is an important estimate of the order of magnitude of the gradient. It follows directly from the two spectral conditions and avoids explicit gradient computation. This is the third beautiful aspect of the Spectral Condition: it enables us to obtain such estimates without computing gradient expressions via the chain rule.

Learning Rate Strategy

Applying the estimate \(\eqref{eq:grad-norm}\) to SGD, i.e., \(\Delta \boldsymbol{W}_k = -\eta_k \nabla_{\boldsymbol{W}_k}\mathcal{L}\), by \(\eqref{eq:grad-norm}\) we have \(\Vert\nabla_{\boldsymbol{W}_k}\mathcal{L}\Vert_2=\mathcal{\Theta}\left(\sqrt{\frac{d_{k-1}}{d_k}}\right)\). To achieve the goal \(\Vert\Delta\boldsymbol{W}_k\Vert_2=\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\), we need \[\begin{equation} \eta_k = \mathcal{\Theta}\left(\frac{d_k}{d_{k-1}}\right)\label{eq:sgd-eta} \end{equation}\]

As for Adam, we still use the SignSGD approximation \(\Delta \boldsymbol{W}_k = -\eta_k \mathop{\text{sign}}(\nabla_{\boldsymbol{W}_k}\mathcal{L})\). Since the entries of \(\mathop{\text{sign}}\) are generally \(\pm 1\), we have \(\Vert\mathop{\text{sign}}(\nabla_{\boldsymbol{W}_k}\mathcal{L})\Vert_F = \mathcal{\Theta}(\sqrt{d_{k-1} d_k})\); and since \(\mathop{\text{sign}}\) is an element-wise operation that generally does not significantly increase rank, we believe \(\mathop{\text{sign}}(\nabla_{\boldsymbol{W}_k}\mathcal{L})\) is low-rank just like \(\nabla_{\boldsymbol{W}_k}\mathcal{L}\). Thus, similarly to the nuclear norm, the \(F\)-norm and the spectral norm will be of the same order, i.e., \(\Vert\mathop{\text{sign}}(\nabla_{\boldsymbol{W}_k}\mathcal{L})\Vert_2 = \mathcal{\Theta}(\sqrt{d_{k-1} d_k})\).

Therefore, to achieve the goal \(\Vert\Delta\boldsymbol{W}_k\Vert_2=\mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right)\), we need \[\begin{equation} \eta_k = \mathcal{\Theta}\left(\frac{1}{d_{k-1}}\right)\label{eq:adam-eta} \end{equation}\]

Now we can compare the results of the Spectral Condition with MuP. MuP assumes we want to build a model \(\mathbb{R}^{d_{in}}\mapsto\mathbb{R}^{d_{out}}\), and divides the model into three parts: first a \(d_{in}\times d\) matrix projects the input to \(d\) dimensions; then modeling is performed in the \(d\)-dimensional space, with parameters all being \(d\times d\) square matrices; finally a \(d\times d_{out}\) matrix produces the \(d_{out}\)-dimensional output. Correspondingly, MuP’s conclusions are also divided into three parts: input, intermediate, and output.

Regarding initialization, MuP’s input variance is \(1/d_{in}\), output variance is \(1/d^2\), and the remaining parameters have variance \(1/d\), while the Spectral Condition’s result is just a single formula \(\eqref{eq:spec-std}\). But upon closer inspection, we find that formula \(\eqref{eq:spec-std}\) already covers all three cases of MuP: setting the input, intermediate, and output matrix sizes as \(d_{in}\times d\), \(d\times d\), \(d\times d_{out}\) and substituting into \(\eqref{eq:spec-std}\), we obtain \[\begin{equation} \begin{aligned} \sigma_{in}^2 =&\, \mathcal{\Theta}\left(\frac{1}{d_{in}}\min\left(1, \frac{d}{d_{in}}\right)\right) = \mathcal{\Theta}\left(\frac{1}{d_{in}}\right) \\ \sigma_k^2 =&\, \mathcal{\Theta}\left(\frac{1}{d}\min\left(1, \frac{d}{d}\right)\right) = \mathcal{\Theta}\left(\frac{1}{d}\right) \\ \sigma_{out}^2 =&\, \mathcal{\Theta}\left(\frac{1}{d}\min\left(1, \frac{d_{out}}{d}\right)\right) = \mathcal{\Theta}\left(\frac{1}{d^2}\right) \end{aligned} \qquad(d\to\infty) \end{equation}\] Some readers may wonder why we only consider \(d\to\infty\): because \(d_{in},d_{out}\) are task-dependent numbers, effectively constants; the only variable model scale is \(d\). Since MuP studies asymptotic laws of hyperparameters with respect to model scale, these are all simplified laws referring to the regime where \(d\) is sufficiently large.

Regarding learning rates: for SGD, MuP’s input learning rate is \(d\), output learning rate is \(1/d\), and the remaining parameters have learning rate \(1\)—note these relations are proportionalities rather than equalities—and the Spectral Condition’s result \(\eqref{eq:sgd-eta}\) likewise covers all three cases. Similarly, for Adam, MuP’s input learning rate is \(1\), output learning rate is \(1/d\), and the remaining parameters have learning rate \(1/d\); the Spectral Condition again describes all three cases with the single formula \(\eqref{eq:adam-eta}\).

Thus, the Spectral Condition, in a way (that in the author’s view) is simpler, arrives at more concise results, and these more concise results are in fact richer in practical implication than MuP’s, because they do not impose overly strong assumptions on the model architecture or parameter shapes. For this reason, the author calls the Spectral Condition a higher-order version of MuP.

Article Summary

This article introduced the upgraded version of MuP—the Spectral Condition. It starts from inequalities related to the spectral norm to analyze the conditions for stable model training, and in a more convenient manner obtains results richer than MuP’s.

\[\begin{equation*} \left\{\begin{aligned} &\,\text{Desired properties:}\left\{\begin{aligned} &\,\Vert\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1) \\[5pt] &\,\Vert\Delta\boldsymbol{x}_k\Vert_{RMS}=\mathcal{\Theta}(1) \end{aligned}\right. \\[10pt] &\,\text{Spectral conditions:}\left\{\begin{aligned} &\,\Vert\boldsymbol{W}_k\Vert_2 = \mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right) \\[5pt] &\,\Vert\Delta\boldsymbol{W}_k\Vert_2 = \mathcal{\Theta}\left(\sqrt{\frac{d_k}{d_{k-1}}}\right) \end{aligned}\right. \\[10pt] &\,\text{Implementations:}\left\{\begin{aligned} &\,\text{Spectral normalization:}\left\{\begin{aligned} &\,\boldsymbol{W}_k = \sigma\sqrt{\frac{d_k}{d_{k-1}}}\frac{\boldsymbol{W}_k'}{\Vert\boldsymbol{W}_k'\Vert_2} \\[5pt] &\,\Delta \boldsymbol{W}_k = \eta\sqrt{\frac{d_k}{d_{k-1}}}\frac{\boldsymbol{\Phi}_k}{\Vert\boldsymbol{\Phi}_k\Vert_2} \end{aligned}\right. \\[10pt] &\,\text{Singular value clipping:}\left\{\begin{aligned} &\,\boldsymbol{W}_k = \sigma\sqrt{\frac{d_k}{d_{k-1}}}\mathop{\text{SVC}}(\boldsymbol{W}_k')\xrightarrow{\text{limit}} \sigma\sqrt{\frac{d_k}{d_{k-1}}}\mathop{\text{msign}}(\boldsymbol{W}_k')\\[5pt] &\,\Delta \boldsymbol{W}_k = \eta\sqrt{\frac{d_k}{d_{k-1}}}\mathop{\text{SVC}}(\boldsymbol{\Phi}_k)\xrightarrow{\text{limit}} \eta\sqrt{\frac{d_k}{d_{k-1}}}\mathop{\text{msign}}(\boldsymbol{\Phi}_k) \end{aligned}\right. \\[10pt] &\,\text{Approximate estimation:}\left\{\begin{aligned} &\,\sigma_k = \mathcal{\Theta}\left(\sqrt{\frac{1}{d_{k-1}}\min\left(1, \frac{d_k}{d_{k-1}}\right)}\right) \\[5pt] &\,\eta_k = \left\{\begin{aligned} &\,\text{SGD: }\mathcal{\Theta}\left(\frac{d_k}{d_{k-1}}\right) \\[5pt] &\,\text{Adam: }\mathcal{\Theta}\left(\frac{1}{d_{k-1}}\right) \end{aligned}\right. \end{aligned}\right. \\[10pt] \end{aligned}\right. \end{aligned}\right. \end{equation*}\]

When reprinting, please include this article’s address: https://kexue.fm/archives/10795

For more detailed reprinting matters, please refer to: Scientific Spaces FAQ