Why Do We Favor Isotropy? An Understanding Based on Steepest Descent

为什么我们偏爱各向同性？基于最速下降的理解 苏剑林 · 2026-01-20 · Machine Learning geometry gradient SGD muon

From data whitening preprocessing in the machine learning era to the many varieties of normalization methods in the deep learning era—BatchNorm, InstanceNorm, LayerNorm, RMSNorm, and so on—these essentially all reflect our preference for "isotropy." Why do we tend to favor isotropic features? What practical benefits does it offer? Many answers can be found for this question—scale alignment, redundancy reduction, decorrelation, and so on—but most of them stay at a surface level.

Recently, while reading the paper The Affine Divergence: Aligning Activation Updates Beyond Normalisation, I came to a new understanding of this question from an optimization perspective, one that I personally think gets fairly close to the essence of the matter. So I'm writing it up to share and discuss with everyone.

Steepest Descent

Let's start from the simplest linear layer

\begin{equation}\boldsymbol{Y} = \boldsymbol{X}\boldsymbol{W}\end{equation}more

where $\boldsymbol{X}\in\mathbb{R}^{b\times d_{in}}$ is the input to the current layer, $\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$ is the weight, and $\boldsymbol{Y}\in\mathbb{R}^{b\times d_{out}}$ is the output. Denoting the loss function as $\mathcal{L}(\boldsymbol{Y}) = \mathcal{L}(\boldsymbol{X}\boldsymbol{W})$, we have

\begin{equation}\frac{\partial \mathcal{L}}{\partial\boldsymbol{W}} = \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}

Taking gradient descent as an example, the update rule is

\begin{equation}\boldsymbol{W}\quad\leftarrow\quad \boldsymbol{W} - \eta\frac{\partial \mathcal{L}}{\partial\boldsymbol{W}} = \boldsymbol{W} - \eta \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}

The basic principle of gradient descent is the well-known fact that "the negative gradient direction is the direction of fastest loss decrease." But this conclusion has a precondition. The most crucial precondition is that the metric being used is the Euclidean norm—if we switch to a different norm, then the steepest direction changes too. We've already discussed this in articles such as Muon Sequel: Why Did We Choose to Try Muon? and Steepest Descent on Manifolds: 1. SGD + Hypersphere.

Switching Perspectives

This article focuses on another, less easily noticed precondition: perspective, or standpoint.

Suppose we agree that "the negative gradient direction is the direction of fastest loss decrease." Then the question is: whose gradient? Some readers might say the parameters' gradient, of course—and this is indeed the standard answer, but it may not be the best answer. Parameters are essentially a byproduct of the model; what we actually care about is not the state of the parameters, but whether the model's input-output mapping matches what we want.

So the change in input-output features is what we should really care about. If we start from the perspective of features, the conclusion turns out to be different. Specifically, as the parameters change from $\boldsymbol{W}$ to $\boldsymbol{W} - \eta \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}$, the change in the output feature $\boldsymbol{Y}$ is

\begin{equation}\Delta \boldsymbol{Y} = \boldsymbol{X}\left(\boldsymbol{W} - \eta \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right) - \boldsymbol{X}\boldsymbol{W} = - \eta \boldsymbol{X} \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}

Following "the negative gradient direction is the direction of fastest loss decrease," if we stand on the side of $\boldsymbol{Y}$, then for the change in $\boldsymbol{Y}$ to make the loss decrease as fast as possible, it should satisfy $\Delta \boldsymbol{Y}\propto -\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}$. Instead, we now have an extra Gram matrix $\boldsymbol{X} \boldsymbol{X}^{\top}$ in the picture, which means the update is not actually moving in the direction of fastest descent.

Isotropy

A naive idea is: wouldn't it be nice if $\boldsymbol{X} \boldsymbol{X}^{\top}$ were, intuitively, equal to the identity matrix (up to some scalar multiple)? We know that $\boldsymbol{X}\in\mathbb{R}^{b\times d_{in}}$, so if $b\leq d_{in}$, then $\boldsymbol{X} \boldsymbol{X}^{\top} = \boldsymbol{I}$ would mean these $b$ vectors form an orthonormal basis. But in practice we usually have $b > d_{in}$, so $\boldsymbol{X} \boldsymbol{X}^{\top} = \boldsymbol{I}$ cannot hold exactly in the mathematical sense.

In this case, the best we can hope for is that these $b$ vectors are distributed as uniformly as possible over the unit hypersphere, so that $\boldsymbol{X} \boldsymbol{X}^{\top} = \boldsymbol{I}$ holds approximately—and this is exactly where isotropy comes from. In other words:

If the input features satisfy isotropy, then the steepest descent on the parameters can approximately coincide with the steepest descent on the features, giving us "two birds with one stone" and substantially improving the model's learning efficiency.

We can further show that if a random variable follows a $d_{in}$-dimensional standard normal distribution, then, while satisfying isotropy, its norm will also concentrate heavily around $\sqrt{d_{in}}$—i.e., it approximately lies on a hypersphere of radius $\sqrt{d_{in}}$. Conversely, if we can standardize the input $\boldsymbol{X}$ into a matrix with zero mean and unit covariance, we consider $\boldsymbol{X} \boldsymbol{X}^{\top} = d_{in}\boldsymbol{I}$ to hold approximately as well—and this operation is exactly whitening.

Normalization Layers

Besides whitening the data beforehand, we can also introduce normalization operations in the middle of the model so that intermediate features approximately satisfy the desired property. More concretely, we try to find a $d_{in}\times d_{in}$ transformation matrix $\boldsymbol{A}$ such that the transformed feature $\boldsymbol{X}\boldsymbol{A}$ satisfies $(\boldsymbol{X}\boldsymbol{A})(\boldsymbol{X}\boldsymbol{A})^{\top}=\boldsymbol{I}$ as closely as possible, i.e.,

\begin{equation}\min_{\boldsymbol{A}} \Vert \boldsymbol{X}\boldsymbol{A}\boldsymbol{A}^{\top}\boldsymbol{X}^{\top} - \boldsymbol{I}\Vert_F\end{equation}

The solution to this problem can be expressed via the pseudoinverse:

\begin{equation}\boldsymbol{A}\boldsymbol{A}^{\top} = \boldsymbol{X}^{\dagger}(\boldsymbol{X}^{\top})^{\dagger} = (\boldsymbol{X}^{\top}\boldsymbol{X})^{-1}\end{equation}

Here we assume $\boldsymbol{X}^{\top}\boldsymbol{X}$ is invertible. From the equation above, we obtain a feasible solution $\boldsymbol{A} = (\boldsymbol{X}^{\top}\boldsymbol{X})^{-1/2}$, and the corresponding transformation is

\begin{equation}\boldsymbol{X}(\boldsymbol{X}^{\top}\boldsymbol{X})^{-1/2}\end{equation}

This is precisely whitening without centering. Since computing $\boldsymbol{X}^{\top}\boldsymbol{X}^{-1/2}$ is relatively expensive, if we instead use a diagonal approximation, this amounts to standardizing each dimension separately—depending on the granularity, this corresponds to BatchNorm, InstanceNorm, and similar operations. If we care more about the "hypersphere," we can instead consider standardizing the norm of each sample separately, which corresponds to LayerNorm, RMSNorm, and similar operations.

Beyond SGD

Interestingly, the conclusion "if the input features satisfy isotropy, then the steepest descent on the parameters approximately coincides with the steepest descent on the features" applies not only to SGD but also to Muon. Consider Muon without momentum, where the update rule is

\begin{equation}\newcommand{msign}{\mathop{\text{msign}}} \boldsymbol{W}\quad\leftarrow\quad \boldsymbol{W} - \eta\msign\left(\frac{\partial \mathcal{L}}{\partial\boldsymbol{W}}\right) = \boldsymbol{W} - \eta \msign\left(\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right)\end{equation}

From $\msign(\boldsymbol{M}) = \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2}$, we have

\begin{equation}\Delta \boldsymbol{Y} = - \eta \boldsymbol{X} \msign\left(\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right) = - \eta \boldsymbol{X} \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}} \left(\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\vphantom{\bigg|}^{\top}\boldsymbol{X}\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right)^{-1/2}\end{equation}

From this we can see that if $\boldsymbol{X}\boldsymbol{X}^{\top}\approx \boldsymbol{I}$, then

\begin{equation}\Delta \boldsymbol{Y} \approx - \eta \frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}} \left(\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\vphantom{\bigg|}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right)^{-1/2} = -\eta\msign\left(\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right)\end{equation}

We know that Muon is steepest descent under the spectral norm, and the equation above means that if the input features satisfy isotropy, then steepest descent under the spectral norm on the parameters approximately coincides with steepest descent under the spectral norm on the features—which is quite elegant! However, this property cannot be reproduced by other optimizers such as SignSGD, and this distinction may be one of the underlying reasons for Muon's outstanding performance.

Summary

In this article we discussed the following question: when does steepest descent at the level of parameters exactly coincide with steepest descent at the level of features? The answer is precisely "isotropy," as stated in the title. From this we arrive at one explanation for why we favor isotropy: it synchronizes steepest descent at these two levels, thereby improving training efficiency.