1. Introduction

What is “maximal update parametrization” really about? When training neural networks, we want to (1) maximally update both the weights and activations (or “features”) while (2) keeping the ‘sizes’ or norms of the weights, activations, and gradients under control. The former is so that we can train larger models with fewer resources and the latter is so that our multi-billion dollar training runs would not randomly blow up in our faces during training. Now, bounding only the weights and using only Lipschitz-continuous layers already suffice to control the Lipschitzness of the model; and consequently also controlling the activation and gradient norms. But bounding the weights (e.g. via weight decay) may also discard components of our updates which we could have allocated to more promising directions in the first place. How do we resolve this?

The crux is to consider the geometry of the space where our weights ’live in’ and do steepest descent there. We can control the weight norms by choosing a bounded manifold and using its retraction map to keep weight norms in a comfortable range. And to prevent the retraction map from discarding parts of our updates, we can enforce that the updates be in the tangent space/cone at the current point in the manifold. Finally we can also maximally update the activations by equipping the tangent spaces/cones of the manifold with the $\texttt{RMS} \to \texttt{RMS}$ induced operator norm (or the scaled spectral norm) as discussed by Yang et al. (2024).

Thinking Machines Lab recently published a blog post following the same idea where they discussed why and how to do steepest descent on the Stiefel manifold equipped with the (scaled) spectral norm. But constraining the weights to be in the Stiefel manifold is too tight of a constraint. For one, this halves the effective parameters of the model. We argue that merely enforcing an upper bound on the singular values (and letting them go to zero during training if need be) would suffice and may even be better as it does not harm model expressivity as much as constraining the singular values to all have the same value.

We call this Steepest Descent on the Spectral Ball, and we shall discuss how to get this done in Section 5. But on our way to solving this problem, we needed to solve subproblems such as finding an efficient and GPU/TPU-friendly way to clip the eigenvalues of a (symmetric) matrix in Section 2, and numerically stable methods to compute projectors to eigenbases and singular subspaces in Section 3.2.1, Section 4.1.1, and Section 5.1.3. With the same tools, we also managed to solve the Steepest Descent problem on the Positive Definite Cone in Section 3 and on the Convex Spectrahedron in Section 4. In Section 6.1 we demonstrate that learning rate transfer (and probably also transfer of other hyperparameters) comes naturally when doing steepest descent on such manifolds. And finally, in Section 6.2 we show that our approach yields in larger updates (even after applying the retraction map) and faster grokking.

2. Eigenvalue Clipping

In Fast, Numerically Stable, and Auto-Differentiable Spectral Clipping via Newton-Schulz Iteration, we discussed a novel method for clipping singular values of a matrix requiring only matrix multiplications. Following the same technique, we can also clip the eigenvalues of a (symmetric) matrix efficiently. This can be used to efficiently project matrices onto the positive semidefinite cone, and for capping the eigenvalues to a comfortable range during training.

I have previously communicated this technique to the authors of Factorization-free Orthogonal Projection onto the Positive Semidefinite Cone with Composite Polynomial Filtering as I mistakenly thought their method for projecting onto the positive semidefinite cone was a special case of my prior work. This work, however, does generalize their technique. I recommend reading their paper!

Let $W \in \mathbb{S}^{n}$ where $\mathbb{S}^{n} = \{W \in \mathbb{R}^{n \times n} | W = W^T\}$ is the set of all $n \times n$ real symmetric matrices. Symmetric matrices have real eigenvalues and can be diagonalized by an orthogonal matrix. We define Eigenvalue Clipping as follows:

Definition 1 (Eigenvalue Clipping). Let $W \in \mathbb{S}^{n}$ be a symmetric matrix and $W = Q \Lambda Q^T$ be its eigenvalue decomposition where $\Lambda = \text{diag}(\lambda_1, \ldots, \lambda_n)$ are the eigenvalues of $W$, $\lambda_i \in \mathbb{R}$ for all $i$, and $Q^T Q = I$. Then we define Eigenvalue Clipping as the following matrix function $\texttt{eig\_clip}_{[\lambda_{min}, \lambda_{max}]}: \mathbb{S}^{n} \to \mathbb{S}^{n}$,

$$\begin{equation}\texttt{eig\_clip}_{[\lambda_{min}, \lambda_{max}]}(W) = Q \texttt{clip}_{[\lambda_{min}, \lambda_{max}]}(\Lambda) Q^T\label{1}\end{equation}$$

where $\lambda_{min}, \lambda_{max} \in (-\infty, \infty)$ are hyperparameters that control the minimum and maximum attainable eigenvalues of the resulting matrix and $\texttt{clip}_{[\alpha, \beta]}: \mathbb{R} \to \mathbb{R}$ is applied element-wise on the eigenvalues of $W$,

$$\begin{equation}\texttt{clip}_{[\alpha, \beta]}(x) = \begin{cases} \alpha & \texttt{if } x < \alpha \\ x & \texttt{if } \alpha \leq x \leq \beta \\ \beta & \texttt{if } \beta < x \end{cases}\end{equation}$$

where $\alpha, \beta \in \mathbb{R} \cup \{-\infty, \infty\}$ and $\alpha \leq \beta$.

The naive implementation of this requires computing the eigenvalue decomposition of $W$, which is computationally expensive and requires high numerical precision (typically float32). Instead, we make use of the GPU/TPU-friendly method to compute the matrix sign function $\texttt{msign}$ by Jordan et al. (2024) and the following identity from the previous blog post:

Proposition 2 (Computing $\texttt{clip}$ via $\texttt{sign}$). Let $\alpha, \beta \in \mathbb{R} \cup \{-\infty, \infty\}$ and $\texttt{clip}: \mathbb{R} \to \mathbb{R}$ be the clipping function defined in Definition 1. Then,

$$\begin{equation} \texttt{clip}_{[\alpha, \beta]}(x) = \frac{\alpha + \beta + (\alpha - x)\texttt{sign}(\alpha - x) - (\beta - x)\texttt{sign}(\beta - x)}{2} \label{eq:clipviasign} \end{equation}$$

2.1 Lifting to matrix form

We can lift Equation $\eqref{eq:clipviasign}$ to matrix form as follows:

which we can implement in JAX as follows:

def sym(W: jax.Array) -> jax.Array:
    return (W + W.T) / 2

def eig_clip(W: jax.Array, alpha: float=-1., beta: float=1.) -> jax.Array:
    assert W.shape[0] == W.shape[1], "W must be square"
    W = sym(W)
    I = jnp.eye(W.shape[0], dtype=W.dtype)
    result = (1/2) * (
        (alpha + beta) * I
        + (alpha * I - W) @ _orthogonalize(alpha * I - W)
        - (beta * I - W) @ _orthogonalize(beta * I - W)
    )
    return sym(result)

2.2 Eigenvalue ReLU and orthogonal projection onto the positive semidefinite cone

Suppose we want to bound the eigenvalues of $W$ from below by a minimum value $\alpha$. For $\alpha = 0$, this is equivalent to projecting $W$ onto the positive semidefinite cone which is useful in e.g. finance and quantum mechanics where objects are typically required to be positive semidefinite. We can do this by setting $\beta = +\infty$ in Equation $\eqref{eq:clipviasign}$:

Lifting this to matrix form yields,

which we can implement in JAX as follows:

def eig_relu(W: jax.Array, alpha: float=0.) -> jax.Array:
    W = sym(W)
    I = jnp.eye(W.shape[0], dtype=W.dtype)
    result = (1/2) * (alpha * I + W + (alpha * I - W) @ _orthogonalize(alpha * I - W))
    return sym(result)

For the orthogonal projection onto the positive semidefinite cone, we set $\alpha = 0$:

The last equality follows from $\texttt{msign}(-W) = -\texttt{msign}(W)$. We can then implement this in JAX as follows:

def proj_psd(W: jax.Array) -> jax.Array:
    W = sym(W)
    return sym((1/2) * (W + W @ _orthogonalize(W)))

2.3 Eigenvalue Hardcapping and orthogonal projection onto the negative semidefinite cone

Suppose we have symmetric matrices $W$ as weights in a neural network and we want to guarantee that the weights do not blow up during training. We can do this by capping the eigenvalues of $W$ to a maximum value $\beta$ after each gradient update. To do this, we can set $\alpha = -\infty$ in Equation $\eqref{eq:clipviasign}$:

Lifting this to matrix form yields,

which we can implement in JAX as follows:

def eig_hardcap(W: jax.Array, beta: float=1.) -> jax.Array:
    W = sym(W)
    I = jnp.eye(W.shape[0], dtype=W.dtype)
    result = (1/2) * (beta * I + W - (beta * I - W) @ _orthogonalize(beta * I - W))
    return sym(result)

For the orthogonal projection onto the negative semidefinite cone, we set $\beta = 0$:

which we can implement in JAX as follows:

def proj_nsd(W: jax.Array) -> jax.Array:
    W = sym(W)
    return sym((1/2) * (W - W @ _orthogonalize(W)))

2.4. Stepfun

Stepfun applies the step function on the singular values/eigenvalues of a matrix. As we will discuss in the next sections, this would be useful for e.g. filtering or “picking out” eigenbasis vectors corresponding to eigenvalues in a certain range in a numerically stable way.

You (2025) first devised a implementation for the rectangular case requiring only matrix multiplications. But as can be seen in the figure above, when applied to the symmetric matrix case, it (1) also acts symmetrically to the negative eigenvalues which is not what we want, and (2) requires two (expensive) $\texttt{msign}$ calls. But a simple modification fixes both issues,

As can be seen in the figure above, this implementation applies the step function properly and only requires one $\texttt{msign}$ call.

We can implement this in JAX as follows:

def eig_stepfun(X: jax.Array, alpha=0.) -> jax.Array:
    I = jnp.eye(X.shape[0], dtype=X.dtype)
    return (I + _orthogonalize(X - alpha * I)) / 2.

3. Steepest descent on the PSD cone

3.1. Problem setup

Suppose we want to do steepest descent on the PSD cone under a norm $\|\cdot\|$ chosen a priori. That is, we want to do first-order optimization where we constrain our weights to be positive semidefinite and our weight updates to have bounded norm. As we previously discussed in Heuristic Solutions for Steepest Descent on the Stiefel Manifold, we can do this as follows:

1. Let $W_t \in \mathcal{M}$ be the ‘weight’ parameter at time step $t$. Compute the “raw gradient” $G_t = \nabla f(W_t)$ via e.g. backpropagation.
2. Compute a ‘optimal’ descent direction $A^* \in T_{W_t} \mathcal{M}$ under the norm in the tangent space at $W_t$, $$\begin{equation} A^* = \arg\min_{A \in \mathbb{R}^{m \times n}} \langle G, A \rangle \quad \text{ s.t. } \quad \| A \|_{W_t} \leq \eta,\quad A \in T_{W_t}\mathcal{M}, \end{equation}$$ where $\eta > 0$ is the learning rate.
3. Update the weight in the direction of $A^*$, $$\widetilde{W}_{t+1} \leftarrow W_t + A^*.$$ Note that $\widetilde{W}_{t+1}$ may not be on the manifold $\mathcal{M}$. And so,
4. Retract the result back to the manifold via a retraction map $W_{t+1} \leftarrow \texttt{retract}_{\mathcal{M}}(\widetilde{W}_{t+1})$.

In our case, the manifold is the PSD cone, $\mathcal{M} := \mathbb{S}^n_{+} = \{W \in \mathbb{S}^n : W \succeq 0\}$. And so, we use the $\texttt{proj\_psd}$ function defined in Section 2.2 as our retraction map.

To find an ‘optimal’ descent direction $A^*$, we can, in some cases, use known Linear Minimization Oracles (LMOs) (Pethick et al., 2025). Or, as we discussed in Steepest Descent on Finsler-Structured (Matrix) Manifolds, we can compute an ‘optimal’ descent direction $A^*$ via two orthogonal projection functions: (i) the projection onto the norm ball, $\texttt{proj}_{\| \cdot \|_{W_t} \leq \eta}$, and (ii) the projection onto the tangent space at $W_t$, $\texttt{proj}_{T_{W_t}\mathcal{M}}$.

If we choose the Frobenius norm, then the projection onto the norm ball is simply,

Alternatively, we can also choose to do steepest descent under the $2 \to 2$ induced operator norm. As to why we might want to do this, you need to binge-read my previous blog posts. In short, controlling the $2 \to 2$ induced operator norm of our weights allows us to control the Lipschitzness of our model which has been shown to improve robustness, generalization, and training stability. In this case, we can use the eigenvalue clipping function defined in Section 2.1 to do the projection onto the spectral norm ball,

The tricky part is the projection onto the tangent space/cone at $W_{t} \in \mathbb{S}^n_{+}$, $\texttt{proj}_{T_{W_{t}}\mathbb{S}^n_{+}}$.

3.2. Projection onto the tangent space/cone at a point on the PSD cone

Special Case: $W_{t}$ is an interior point of the PSD cone. That is, $W_{t} \in \mathbb{S}^n_{++} \subset \mathbb{S}^n_{+}$ or, equivalently, $W_{t} \succ 0$. Then the tangent space is the entire space of symmetric matrices,

And the projection onto the tangent space is simply the symmetrization operation $\texttt{sym}(X) = (X + X^T)/2$,

General Case: For any $W_{t} \in \mathbb{S}^n_{+}$, we may no longer have a tangent space but rather a tangent cone. That is, the vectors in the tangent cone still form a closed space, but if $H \in T_{W_{t}} \mathbb{S}^n_{+}$, then $-H$ may not be in $T_{W_{t}} \mathbb{S}^n_{+}$ (Rockafellar & Wets, 2009). And thus, we need to be careful with the directions of our inputs to the projection map. The tangent cone at $W_{t} \in \mathbb{S}^n_{+}$ is given by,

where $U_0 \in \mathbb{R}^{m \times (n-r)}$ is the orthonormal basis for the null space of $W_{t}$ and $r = \texttt{rank}(W_t)$. Note that if $W_{t}$ is full rank (and therefore positive definite), then $U_0 = 0$ and we recover the special case above.

Let $\widehat{X} := \texttt{sym}(X)$, $U = \begin{bmatrix} U_{r} & U_0 \end{bmatrix}$ be the eigenbasis of $W_t$, and $P_0 = U_0 U_0^T$ be the projector onto the null space of $W_{t}$. The projection onto the tangent cone at $W_{t} \in \mathbb{S}^n_{+}$ is given by,

where the fifth equality follows from $I = UU^T$ and the orthogonal invariance of the Frobenius norm, and the second-to-last equality is from the similarity-equivariance of matrix functions that acts entrywise on the eigenvalues/singular values.

In words,

We first symmetrize the input $X$ into $\widehat{X}$ then we subtract the negative eigenvalues of the projection of $\widehat{X}$ onto the null space of $W_{t}$.

3.2.1. Numerically stable computation of the null space projector

Intuitively, to construct the null space projector $P_0$, we can “select” from $Q$ the eigenvectors corresponding to the zero eigenvalues of $W_{t}$ as follows,

where the second line is a relaxation to handle numerical precision issues.

Taking everything together yields,

def project_to_tangent_psd(W: jax.Array, X: jax.Array, tol=1e-3) -> jax.Array:
    P0 = jnp.eye(W.shape[0], dtype=W.dtype) - eig_stepfun(W, tol)
    return jax.lax.cond(
        jnp.rint(jnp.trace(P0)) == 0,
        lambda: sym(X),  # W is an interior point, so tangent space is all symmetric matrices
        lambda: sym((X_ := sym(X)) - proj_nsd(P0 @ X_ @ P0)),
    )

3.2.2. Sanity check

dim = 768
nullity = 128
keys = jax.random.split(jax.random.PRNGKey(0), 2)

W = jax.random.normal(keys[0], (dim, dim))
X = jax.random.normal(keys[1], (dim, dim))

W = sym(W @ W.T)
lam, Q = jnp.linalg.eigh(W)
lam = lam.at[:nullity].set(0)
W = Q @ jnp.diag(lam) @ Q.T

3.3. Update rule for steepest descent on the PSD cone

We now have all the necessary components to do steepest descent under any norm on the PSD cone.

3.3.1. Special case: $W_t$ is an interior point of the PSD cone

As we discussed in the previous section, if $W_t$ is full rank, then the tangent space at that point is simply the space of all symmetric matrices. An interesting coincidence is that Linear Minimization Oracles (LMOs) for common norms derived for the case without the tangency constraint preserve symmetry.

Therefore, it would suffice to symmetrize the “raw gradient” $G_t$ first and then apply the LMO. This guarantees that our updates are indeed on-tangent and maximal (via theory behind LMOs). Our update rule would then be,

3.3.2. General case

In general, LMOs derived for the case without the tangency constraint often ‘send’ its output off-tangent. An example we discussed in previous blog post is $\texttt{msign}$ and the Stiefel manifold. In such cases, we can use either of the following two methods to compute an ‘optimal’ descent direction $A^*$:

1. A heuristic solution such as the one discussed in Heuristic Solutions for Steepest Descent on the Stiefel Manifold where we iteratively apply the projection onto the tangent space and the LMO until convergence. That is, $$\begin{align} W_{t+1} &= \texttt{proj\_psd}\left(W_{t} + \left(\texttt{LMO}_{\|\cdot\|_{W_t} \leq \eta} \circ \texttt{proj}_{T_{W_{t}}\mathbb{S}^n_{+}} \right)^K (-G_t) \right) \end{align}$$ for some integer $K \geq 1$ denoting the number of iterations (typically, $K = 4$ to $8$ suffices; but the iteration can be terminated upon convergence).
2. An exact solution such as the primal-dual hybrid gradient method, $\texttt{pdhg}$, we discussed in Steepest Descent on Finsler-Structured (Matrix) Manifolds, $$\begin{align} W_{t+1} &= \texttt{proj\_psd}(W_{t} + \texttt{pdhg}(W_t, G_t, \texttt{proj}_{\| \cdot \|_{W_t} \leq \eta}, \texttt{proj}_{T_{W_{t}}\mathbb{S}^n_{+}})) \end{align}$$ We can also speed up PDHG by warm-starting the initial iterate $A^0$ with the heuristic above or the solution from the previous time step $A^*_{t-1}$ (in theory, the solutions should not drift too much between time steps if we accumulate the gradients with a momentum rule).

Voilà, we now have an efficient, factorization-free, and GPU/TPU-friendly method for doing steepest descent on the PSD cone under any norm.

4. Steepest descent on the Convex Spectrahedron

Suppose we want to constrain our weights to have eigenvalues bounded within some range $[\alpha, \beta] \subseteq \mathbb{R}$. That is, we “place” our weights on the Convex Spectrahedron,

and do steepest descent there under some norm chosen a priori. For the retraction map, we can use the eigenvalue clipping function defined in Section 2.1,

4.1. Projection onto the tangent space/cone at a point on the Convex Spectrahedron

The tangent cone at $W_t \in \mathcal{K}_{[\alpha, \beta]}$ is generally given by,

where $U_{\alpha}$ and $U_{\beta}$ are the orthonormal bases for the $\alpha$- and $\beta$-eigenspaces of $W_t$, respectively. If $W_t$ is an interior point, that is, $\alpha I \prec W_t \prec \beta I$, then $U_\alpha = U_\beta = 0$ and the tangent space is simply the space of symmetric matrices, $T_{W_t} \mathcal{K}_{(\alpha, \beta)} = \mathbb{S}^n$.

Let $\widehat{X} := \texttt{sym}(X)$, $U := \begin{bmatrix} U_{\beta} & U_{\widetilde{r}} & U_{\alpha} \end{bmatrix}$ be the eigenbasis of $W_{t}$, and $P_\alpha := U_{\alpha}U_{\alpha}^T, P_\beta := U_{\beta}U_{\beta}^T$ be the projectors onto the $\alpha$- and $\beta$-eigenspaces of $W_t$, respectively. Then, following the strategy we discussed in Section 3.2, the projection onto the tangent cone at $W_t \in \mathcal{K}_{[\alpha, \beta]}$ is given by,

or in words,

We first symmetrize the input $X$ into $\widehat{X}$ and then we subtract the negative eigenvalues of the projection of $\widehat{X}$ onto the $\alpha$-eigenspace of $W_t$ and the positive eigenvalues of the projection of $\widehat{X}$ onto the $\beta$-eigenspace of $W_t$.

4.1.1. Numerically stable computation of the eigenspace projectors

As in Section 3.2.1, we can construct the eigenspace projectors $P_\alpha$ and $P_\beta$ as follows,

Likewise, $P_\beta \approx \texttt{eig\_stepfun}(W, \beta - \epsilon)$ for small $\epsilon > 0$.

Taking everything together yields,

def project_to_tangent_convex_spectrahedron(W: jax.Array, X: jax.Array, alpha: float, beta: float, tol=1e-3):
    P_alpha, P_beta = jnp.eye(W.shape[0], dtype=W.dtype) - eig_stepfun(W, alpha+tol), eig_stepfun(W, beta-tol)
    return jax.lax.cond(
        jnp.logical_and(jnp.rint(jnp.trace(P_alpha)) == 0, jnp.rint(jnp.trace(P_beta)) == 0),
        lambda: sym(X),  # W is in the interior, so tangent space is all symmetric matrices
        lambda: sym((X_ := sym(X)) - proj_nsd(P_alpha @ X_ @ P_alpha) - proj_psd(P_beta @ X_ @ P_beta)),
    )

4.2. Update rule for steepest descent on the Convex Spectrahedron

4.2.1. Special case: $W_t$ is an interior point of the Convex Spectrahedron

If $W_t$ is an interior point of the Convex Spectrahedron $\mathcal{K}_{[\alpha, \beta]}$ (that is, $\alpha I \prec W_t \prec \beta I$), then the tangent space at that point is simply the space of all symmetric matrices. Thus, as in the Section 3.3.1, we can use known LMOs that preserve symmetry. Our update rule would then be,

4.2.2. General case

In general, we can use either the heuristic or the PDHG method discussed in Section 3.3.2,

or,

5. Steepest descent on the Spectral Ball

The previous examples are arguably contrived. This example is more practical.

Suppose we no longer constrain our weights to be symmetric, but we still want to bound their spectral norm. That is, we want to do steepest descent on the Spectral Ball,

for some radius $R > 0$. For the retraction map, we can use the GPU/TPU-friendly Spectral Hardcap function discussed in Fast, Numerically Stable, and Auto-Differentiable Spectral Clipping via Newton-Schulz Iteration,

In Appendix A1, we generalize this to steepest descent on the Spectral Band where we bound the singular values within some range $[\alpha, \beta]$ to prevent weights from blowing up or vanishing.

5.1. Projection onto the tangent space/cone at a point on the Spectral Ball

5.1.1. Shortcut via dilation (slower)

The crux is to observe that the singular values of $W_t \in \mathcal{B}_{\|\cdot\|_{2 \to 2} \leq R}$ are $\pm$ the eigenvalues of the block matrix,

where the mapping $\Phi: \mathbb{R}^{m \times n} \to \mathbb{S}^{m+n}$ is an isometry (up to scaling by $\sqrt{2}$) and therefore we can recover the projection via $[\cdot]_{12}$. This allows us to compute the projection onto the tangent cone at $W_t \in \mathcal{B}_{\|\cdot\|_{2 \to 2} \leq R}$ via the projection onto the tangent cone at $\widetilde{W_t} \in \mathcal{K}_{[-R, R]}$,

which we can implement in JAX as follows:

def project_to_tangent_spectral_ball(W: jax.Array, X: jax.Array, R: float, tol=1e-3) -> jax.Array:
    m, n = W.shape
    phi = lambda A: jnp.block([[jnp.zeros((m, m), dtype=A.dtype), A], [A.T, jnp.zeros((n, n), dtype=A.dtype)]])
    return jax.lax.cond(
        _power_iterate(W) < R,  # or jnp.linalg.norm(W, ord=2) < R
        lambda: X,  # W is an interior point, so tangent space is all matrices
        lambda: project_to_tangent_convex_spectrahedron(phi(W), phi(X), -R, R, tol)[:m, m:],
    )

5.1.2. Direct approach (faster)

Similar to the previous sections, the tangent cone at $W_t \in \mathcal{B}_{\|\cdot\|_{2 \to 2} \leq R}$ is generally given by,

where $U_R \in \mathbb{R}^{m \times k}$ and $V_R \in \mathbb{R}^{n \times k}$ are the orthonormal bases for the left and right $R$-singular subspaces of $W_t$ (that is, the singular vectors corresponding to the singular values equal to $R$), respectively, and $k$ is the multiplicity of the singular value $R$. Note that if $W_{t}$ is an interior point, that is, $\| W_t \|_{2 \to 2} < R$, then $U_R = V_R = 0$ and the tangent space is simply the entire space of matrices, $T_{W_t} \mathcal{B}_{\|\cdot\|_{2 \to 2} < R} = \mathbb{R}^{m \times n}$.

Let $U := \begin{bmatrix} U_{< R} & U_R \end{bmatrix}$ and $V := \begin{bmatrix} V_{< R} & V_R \end{bmatrix}$ be the left and right singular bases of $W_{t}$, respectively. Following our strategy in the previous sections then yields the projection onto the tangent cone at $W_t \in \mathcal{B}_{\|\cdot\|_{2 \to 2} \leq R}$,

where $P_{V_{R}} := V_{R} V_{R}^T$ is the projector onto the right $R$-singular subspace of $W_t$, and $J_R := U_R V_R^T$ is the partial isometry corresponding to the $R$-singular subspace of $W_t$. The fourth equality comes from,

5.1.3. Numerically stable computation of the singular subspace projectors

First, note that for $W = U \Sigma V^T$, we have $W_t^T W_t = V \Sigma^2 V^T$. Thus,

And,

Taking everything together yields,

def project_to_tangent_spectral_ball(W: jax.Array, X: jax.Array, R: float, tol=1e-3) -> jax.Array:
    return jax.lax.cond(
        _power_iterate(W) < R - tol,  # or jnp.linalg.norm(W, ord=2) < R - tol
        lambda: X,  # W is an interior point, so tangent space is all matrices
        lambda: X - (J_R := (1./R) * W @ (PV_R := eig_stepfun(W.T @ W / R**2, 1.-tol))) @ proj_psd(sym(J_R.T @ X @ PV_R)),
    )

5.2. Update rule for steepest descent on the Spectral Ball

5.2.1. Special case: $W_t$ is an interior point of the Spectral Ball

If $W_t$ is inside the Spectral Ball, then the tangent space at that point is $\mathbb{R}^{m \times n}$ and thus the projection is simply the identity map. We can use the LMO for any norm without worry,

5.2.2. General case

In general, we can use either the heuristic or the PDHG method discussed in Section 3.3.2,

or,

6. Experiments

In all of our experiments below, we constrain weight updates to have $\texttt{RMS}\to\texttt{RMS}$ induced operator norm $\leq \eta$, where $\eta > 0$ is the learning rate. We then vary the manifold to do steepest descent on, and the dualization strategy to compute the optimal update directions on the respective manifolds. In summary, we use the following maps,

6.1. Learning rate transfer, XOR problem

As a minimal example for learning rate transfer, we train a $2 \to D \to D \to 2$ MLP on the XOR problem for 32 training steps via the Modula library. We then vary the constraint set for the weights and use the PDHG algorithm to compute the optimal update directions on the respective constraint sets.

We also warm-start the initial iterate of PDHG with the 1 step of the Alternating Projections heuristic (which also solves the case when the weight $W_t$ is an interior point).

6.1.1. Positive semidefinite cone

We only constrain the weight for the $D \to D$ linear layer to be positive semidefinite, and the other layers are constrained to the (scaled) Stiefel manifold as discussed in Steepest Descent on Finsler-Structured (Matrix) Manifolds. As can be seen in the Figure above, the optimal learning rates do transfer under our parametrization. However, we cannot impose a Lipschitz bound on the network because the positive semidefinite cone is unbounded.

6.1.2. Steepest descent on the Convex Spectrahedron

Like above, we only constrain the weight for the $D \to D$ linear layer to be in the Convex Spectrahedron, and the other layers are constrained to the (scaled) Stiefel manifold. But here, we can impose a Lipschitz bound on the network because the Convex Spectrahedron is bounded. As can be seen in the Figure above, the optimal learning rates do transfer under our parametrization.

6.1.3. Steepest descent on the Spectral Ball

Unlike the previous two experiments, we now constrain all the weights to be in the unit RMS-RMS ball $\mathcal{B}_{\|\cdot\|_{\texttt{RMS} \to \texttt{RMS}} \leq 1} = \mathcal{B}_{\|\cdot\|_{2 \to 2} \leq \sqrt{m/n}}$. As can be seen in the Figure above, the optimal learning rates do transfer under our parametrization.

6.2. Larger updates accelerate generalization, Addition-Modulo-31 problem

Q: Does having larger updates actually help generalization or merely introduce noise?

Preliminary results show that it is indeed the former.

Here we train an MLP on the Addition-Modulo-31 problem while constraining the weights to be in the $\texttt{RMS}\to\texttt{RMS}$ ball of radius $R = 4$. We use the retraction map discussed in Section 5 to keep the weights bounded. With the Muon optimizer, an equivalent weight decay for such a constraint would be $\lambda = 1/R = 0.25$ (Chen et al., 2025; Pethick et al., 2025; Liu et al., 2025), which is too large and discards too much information from the updates at each step. See Appendix A2 for more details. The problem also only has $961$ data points in total. In all, every update step in this setting matters. Noise that gets added into the updates and any information lost from the weight controls become immediately obvious in the generalization performance.

Aside from test accuracy, we also measure the weight delta between steps,

As can be seen in the Figure above, our dualizers result in around $2\times$ larger weight deltas compared to baseline. This is because the updates our dualizers produce are in the tangent cones and so the retraction map mostly leaves them intact. In contrast, the baseline updates often have (rather large) components that get discarded by the retraction map, reducing the effective learning rate of the update.

Also notice that the 1-step Alternating Projections heuristic (projecting the “raw gradient” to the tangent cone before applying the LMO) results in similarly-sized weight deltas and also groks the problem in roughly the same number of steps as our PDHG dualizer. This suggests that few-step Alternating Projections may already suffice in practice for larger training runs.

6.3. NanoGPT-scale experiments [Under Construction]

How to Cite

@misc{cesista2025rethinkingmupspectralball,
  author = {Franz Louis Cesista},
  title = {Rethinking Maximal Update Parametrization: Steepest Descent on the Spectral Ball},
  year = {2025},
  month = {October},
  day = {15},
  url = {https://leloykun.github.io/ponder/rethinking-mup-spectral-ball/},
}

If you find this post useful, please consider supporting my work by sponsoring me on GitHub:

References

1. Greg Yang, James B. Simon, Jeremy Bernstein (2024). A Spectral Condition for Feature Learning. URL https://arxiv.org/abs/2310.17813
2. Jeremy Bernstein (2025). Modular Manifolds. URL https://thinkingmachines.ai/blog/modular-manifolds/
3. Franz Cesista (2025). Fast, Numerically Stable, and Auto-Differentiable Spectral Clipping via Newton-Schulz Iteration. URL https://leloykun.github.io/ponder/spectral-clipping/
4. Shucheng Kang, Haoyu Han, Antoine Groudiev, Heng Yang (2025). Factorization-free Orthogonal Projection onto the Positive Semidefinite Cone with Composite Polynomial Filtering. URL https://arxiv.org/abs/2507.09165
5. Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein (2024). Muon: An optimizer for hidden layers in neural networks. Available at: https://kellerjordan.github.io/posts/muon/
6. You Jiacheng (2025). X post on Stepfun. URL https://x.com/YouJiacheng/status/1930988035195478303
7. Franz Cesista (2025). Heuristic Solutions for Steepest Descent on the Stiefel Manifold. URL https://leloykun.github.io/ponder/steepest-descent-stiefel/
8. Franz Cesista (2025). Steepest Descent on Finsler-Structured (Matrix) Manifolds. URL https://leloykun.github.io/ponder/steepest-descent-finsler/
9. R. Tyrrell Rockafellar, Roger J-B Wets (2009). Variational Analysis. URL https://sites.math.washington.edu/~rtr/papers/rtr169-VarAnalysis-RockWets.pdf
10. Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, Volkan Cevher (2025). Training Deep Learning Models with Norm-Constrained LMOs. URL https://arxiv.org/abs/2502.07529
11. Jeremy Bernstein (2025). The Modula Docs. URL https://docs.modula.systems/
12. Lizhang Chen, Jonathan Li, Qiang Liu (2025). Muon Optimizes Under Spectral Norm Constraints. URL https://arxiv.org/abs/2506.15054
13. Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang, Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, Zhilin Yang (2025). Muon is Scalable for LLM Training. URL https://arxiv.org/abs/2502.16982
14. Ben Keigwin, Dhruv Pai, Nathan Chen (2025). Gram-Space Manifold Muon. URL https://www.tilderesearch.com/vignettes/gram-space
15. Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, Anima Anandkumar (2018). signSGD: Compressed Optimisation for Non-Convex Problems. URL https://arxiv.org/abs/1802.04434
16. Fabian Schaipp (2024). How to jointly tune learning rate and weight decay for AdamW. URL https://fabian-sp.github.io/posts/2024/02/decoupling/
17. Atli Kosson, Jeremy Welborn, Yang Liu, Martin Jaggi, Xi Chen (2025). Weight Decay may matter more than muP for Learning Rate Transfer in Practice. URL https://arxiv.org/abs/2510.19093
18. Yuandong Tian (2025). Provable Scaling Laws of Feature Emergence from Learning Dynamics of Grokking. URL https://arxiv.org/abs/2509.21519

Appendices

A1. Steepest Descent on the Spectral Band

In Section 5, we discussed steepest descent on the manifold of matrices with bounded-above spectral norm. This prevents weight norms from blowing up, but does not prevent them from collapsing to zero and thus stalling learning anyway. This naturally leads us to steepest descent on the Spectral Band, or the manifold of matrices with bounded-above and bounded-below singular values,

for some inner and outer radii $0 \leq \alpha \leq \beta$.

For the retraction map, we can use the GPU/TPU-friendly Spectral Clip function discussed in Fast, Numerically Stable, and Auto-Differentiable Spectral Clipping via Newton-Schulz Iteration,

A1.1. Projection onto the tangent space/cone at a point on the Spectral Band

Following Section 5.1.2, the tangent cone at a point $W_t \in \mathcal{S}_{[\alpha, \beta]}$ is generally given by,

where $U_{\alpha}$ and $V_{\alpha}$ are the orthonormal bases for the left and right $\alpha$-singular subspaces of $W_t$ (that is, the singular vectors corresponding to the singular values equal to $\alpha$), respectively. Likewise for $U_{\beta}$ and $V_{\beta}$ with respect to the singular value $\beta$. And if $W_t$ is an interior point, that is, $\alpha < \| W_t \|_{2 \to 2} < \beta$, then $U_{\alpha} = V_{\alpha} = U_{\beta} = V_{\beta} = 0$ and the tangent space is simply the entire space of matrices, $T_{W_t} \mathcal{B}_{\alpha < \| \cdot \|_{2 \to 2} < \beta} = \mathbb{R}^{m \times n}$.

As before, let $U := \begin{bmatrix} U_\alpha & U_{\tilde{r}} & U_\beta \end{bmatrix}$ and $V := \begin{bmatrix} V_\alpha & V_{\tilde{r}} & V_\beta \end{bmatrix}$ be the left and right singular bases of $W_{t}$, respectively. Then the projection is,

where $P_{V_{\alpha}}$ and $P_{V_{\beta}}$ are the projectors onto the right $\alpha$- and $\beta$-singular subspaces of $W_t$, respectively, and $J_{\alpha} := U_{\alpha} V_{\alpha}^T$ and $J_{\beta} := U_{\beta} V_{\beta}^T$ are the polar factors of $W_t$ restricted to the respective singular subspaces. We can compute these in a numerically stable way as in Section 5.1.3.

Taking everything together yields,

def project_to_tangent_spectral_band(W: jax.Array, X: jax.Array, alpha: float, beta: float, tol=1e-3) -> jax.Array:
    return jax.lax.cond(
        alpha == 0,
        lambda: project_to_tangent_spectral_ball(W, X, beta, tol),
        lambda: jax.lax.cond(
            jnp.logical_and(
                jnp.rint(jnp.trace(PV_alpha := jnp.eye(W.shape[1]) - eig_stepfun(W.T @ W / alpha**2, 1.+tol))) == 0,
                jnp.rint(jnp.trace(PV_beta := eig_stepfun(W.T @ W / beta**2, 1.-tol))) == 0,  # or spec_norm < beta - tol,
            ),
            lambda: X,  # tangent space is all matrices
            lambda: X
                - (J_alpha := (1./alpha) * W @ PV_alpha) @ proj_nsd(sym(J_alpha.T @ X @ PV_alpha))
                - (J_beta := (1./beta) * W @ PV_beta) @ proj_psd(sym(J_beta.T @ X @ PV_beta)),
        ),
    )

We can then compute the optimal updates $A^*$ as in Section 5.2 via the projection above.

A1.1.1. Sanity check: Stiefel as a special case of the Spectral Band

On the Stiefel manifold $\texttt{St}(m, n) = \{ W \in \mathbb{R}^{m \times n} : W^T W = I_n \}$, the singular values of any $W \in \texttt{St}(m, n)$ are all equal to $1$. Thus, $\texttt{St}(m, n) = \mathcal{S}_{[1, 1]}$ and,

Without loss of generality (up to rotations), we can also choose that $U = W_t$ and $V = I_n$ such that $W_t = UIV^T$. Thus,

which is simply the textbook definition of the tangent space at a point on the Stiefel manifold.

As for the tangent space, we have,

which is, again, the textbook formula for the projection onto the tangent space at a point on the Stiefel manifold.

A1.2. Dual ascent for steepest descent on the Spectral Band

In this work, we compute the optimal updates $A^*$ via PDHG and (orthogonal) projections onto the tangent cones/spaces and the norm balls. However, prior work by Bernstein (2025) and more recently by Keigwin, et al. (2025) use dual ascent methods instead. We will generalize this approach to the Spectral Band and also discuss why we believe PDHG via projections is preferable in this setting.

To recap, our optimization problem is, given a “raw gradient” $G_t \in \mathbb{R}^{m \times n}$ and a choice of norm $\| \cdot \|_{W_t}$ to do steepest descent under, we want to find the optimal update $A^*$ such that,

where the linear maps $L_{\alpha}, L_{\beta} : \mathbb{R}^{m \times n} \to \mathbb{S}^{r}_{\pm}$ and their adjoints are defined as,

such that $\langle L_{\alpha}(X), S_{\alpha} \rangle = \langle X, L_{\alpha}^*(S_{\alpha}) \rangle$ and likewise for $L_{\beta}$.

Restricting the dual states $S_{\alpha}$ to the negative semidefinite cone, $S_{\alpha} \in \mathbb{S}^{r_{\alpha}}_{-}$, and $S_{\beta}$ to the positive semidefinite cone, $S_{\beta} \in \mathbb{S}^{r_{\beta}}_{+}$, yields the Lagrangian,

One can then check that,

And by Sion’s minimax theorem, we can swap the order of minimization and maximization,

First, let us consider the primal minimizer,

where $\texttt{LMO}_{\| \cdot \|_{W_t} \leq \eta}$ is the linear minimization oracle for the norm $\| \cdot \|_{W_t}$ (Pethick et al., 2025). For the $\texttt{RMS} \to \texttt{RMS}$ norm, we have $\texttt{LMO}_{\| \cdot \|_{\texttt{RMS} \to \texttt{RMS}} \leq \eta}(X) = \sqrt{\frac{m}{n}}\eta \cdot \texttt{msign}(X).$

This then yields the dual problem,

where $\| \cdot \|_{W_t}^*$ is the dual norm of $\| \cdot \|_{W_t}$. For the $\texttt{RMS} \to \texttt{RMS}$ norm, we have $\| \cdot \|_{\texttt{RMS} \to \texttt{RMS}}^* \propto \| \cdot \|_{\text{nuc}}$. And by chain rule, the above has supergradients,

We can then do gradient ascent on the dual variables $S_{\alpha}$ and $S_{\beta}$ while projecting them back to their respective cones after each step. Taking everything together then yields,

where $\sigma > 0$ is the dual ascent learning rate. At convergence, we have $A_t \to A^*$.

A1.2.1. Initialization strategy

We can initialize the dual states $S_{\alpha, 0}$ and $S_{\beta, 0}$ as zero matrices. However, notice that the update rule for $A_t$ above is already similar to the 1-step Alternating Projections heuristic we discussed and have shown to be effective in earlier sections.

where,

A1.2.2. JAX implementation

def dual_ascent(
    G: jax.Array,  # R^(m x n)
    L_primal: Callable[[jax.Array], Tuple[jax.Array]],  # R^(m x n) -> K_dual
    L_dual:  Callable[[Tuple[jax.Array]], jax.Array],  # K_dual -> R^(m x n)
    proj_K_dual: Callable[[Tuple[jax.Array]], Tuple[jax.Array]],  # K_dual -> K_dual
    norm_K_dual: Callable[[Tuple[jax.Array]], float],  # K_dual -> R
    lmo: Callable[[jax.Array], jax.Array],  # R^(m x n) -> R^(m x n)
    *,
    max_steps: int=128, sigma: float=1.0,
    rtol: float=1e-3, atol: float=1e-6,
):
    def cond_fn(state):
        S, k, res = state
        return jnp.logical_and(k < max_steps, jnp.logical_and(res > atol, res > rtol * norm_K_dual(S)))

    def body_fn(state):
        S, k, _ = state
        A = -lmo(G + L_dual(S))
        gradS = L_primal(A)
        S_new = proj_K_dual(jax.tree_util.tree_map(lambda s, g: s + sigma / jnp.sqrt(k+1) * g, S, gradS))
        res = norm_K_dual(jax.tree_util.tree_map(lambda s, s_old: s - s_old, S_new, S))
        return S_new, k+1, res

    S = proj_K_dual(L_primal(-G))
    S_final, n_iters, final_res = jax.lax.while_loop(cond_fn, body_fn, (S, 0, jnp.inf))
    A_final = -lmo(G + L_dual(S_final))
    return A_final


def dual_ascent_spectral_band_spectral_norm(
    W: jax.Array, G: jax.Array,
    alpha: float, beta: float, target_norm: float,
    *,
    max_steps: int=128, sigma: float=1.0,
    eig_tol: float=1e-1, rtol: float=1e-3, atol: float=1e-6,
):
    m, n = W.shape
    PV_alpha = jnp.eye(n) - eig_stepfun(W.T @ W / alpha**2, 1.+eig_tol)
    PV_beta  = eig_stepfun(W.T @ W / beta**2, 1.-eig_tol)
    Omega = jax.random.normal(jax.random.PRNGKey(0), (n, n), dtype=W.dtype)
    V_alpha, V_beta = _orthogonalize(PV_alpha @ Omega), _orthogonalize(PV_beta @ Omega)
    U_alpha, U_beta = 1./alpha * W @ V_alpha, 1./beta * W @ V_beta

    L_alpha_primal = lambda A: sym(U_alpha.T @ A @ V_alpha)
    L_beta_primal  = lambda A: sym(U_beta.T @ A @ V_beta)
    L_alpha_dual   = lambda S: U_alpha @ S @ V_alpha.T
    L_beta_dual    = lambda S: U_beta @ S @ V_beta.T

    L_primal    = lambda A: (L_alpha_primal(A), L_beta_primal(A))
    L_dual      = lambda S: L_alpha_dual(S[0]) + L_beta_dual(S[1])
    proj_K_dual = lambda S: (proj_nsd(S[0]), proj_psd(S[1]))
    norm_K_dual = lambda S: jnp.linalg.norm(S[0]) + jnp.linalg.norm(S[1])

    return dual_ascent(
        G,
        L_primal=L_primal,
        L_dual=L_dual,
        proj_K_dual=proj_K_dual,
        norm_K_dual=norm_K_dual,
        lmo=lambda X: target_norm * _orthogonalize(X),
        max_steps=max_steps, sigma=sigma,
        rtol=rtol, atol=atol,
    )

A1.2.3. Sanity check: Stiefel as a special case of the Spectral Band

Pick the $\texttt{RMS} \to \texttt{RMS}$ norm to descend under. For the Stiefel case, we have $\alpha = \beta = 1$, $U_1 = U_2 =: U$, and $V_1 = V_2 =: V$. And WLOG, up to rotations, we can choose $U = W_t$ and $V=I$. Thus,

where $\Lambda_t := S_{\alpha, t} + S_{\beta, t} \in \mathbb{S}^n$. And,

Both match the update rules that Bernstein (2025) previously derived.

A1.2.4. Why PDHG over dual ascent?

Mainly numerical stability.

• First, concerning the internal states. We would still need to maintain states in both cases, but PDHG only needs to maintain projected gradients and the gap between them (which we could even avoid storing explicitly). In contrast, the $S_\alpha$ and $S_\beta$ dual states can grow unboundedly large, especially with a poor choice of dual ascent learning rate $\sigma$.
• Second is that the dual ascent approach requires us to construct $U_\alpha$, $V_\alpha$, $U_\beta$, and $V_\beta$ which leads to additional numerical instability compared to just maintaining the projectors $P_{V_\alpha}$ and $P_{V_\beta}$ as in the PDHG approach.
• Third is that PDHG often converges faster in practice.

For future work, one may explore PDHG on the dual space as well, which may combine the best of both worlds.

A2. Weight Decay as a (manifold) constraint

In Section 6.2, we mentioned that the equivalent to constraining the weight norms to be inside a ball of radius $4$ is to use a weight decay of $1/4$. Here, we shall explain why.

A2.1. Weight norms at equilibrium under weight decay

Prior work has shown that steepest descent under any chosen norm with weight decay $\lambda$ already (secretly) constrains the weights to be inside the corresponding norm ball of radius $1/\lambda$ (Chen et al., 2025; Pethick et al., 2025; Liu et al., 2025).

Proposition 3 (Equilibrium of weight norms with weight decay). Consider the update form,

$$\begin{equation} W_{t+1} = (1 - \eta\lambda) W_t + A_t^* \end{equation}$$

where $\lambda > 0$ is the weight decay term, $\eta > 0$ is the learning rate, and $\| A_t^* \| \leq \eta$ for some norm $\| \cdot \|$ chosen a priori. If $\eta\lambda \leq 1$, then the weight norms are upper bounded as,

$$\lim\sup_{t \to \infty} \| W_{t} \| \to \frac{1}{\lambda}.$$

Furthermore, if the steepest descent updates stabilizes to some constant $A_t^* \to A_*^*$, then the weights also stabilizes to $W_t \to \frac{A_*^*}{\eta\lambda}$.

Proposition 4 (Mechanism behind gradient-weight alignment). Consider the update form,

$$\begin{equation} W_{t+1} = (1 - \eta\lambda) W_t + A_t^* \end{equation}$$

where $\lambda > 0$ is the weight decay term, $\eta > 0$ is the learning rate, and $\| A_t^* \| \leq \eta$ for some norm $\| \cdot \|$ chosen a priori. If $\eta\lambda \leq 1$ and $A_t^* \to A_*^*$ as $t \to \infty$, then the weights converge as,

$$\lim_{t \to \infty} W_{t} \to \frac{A_*^*}{\eta\lambda}.$$

Intuitively, this means that, for example, Muon “pulls” all of the singular values to $\frac{\eta}{\eta\lambda} = \frac{1}{\lambda}$ as weight updates align during training. With the SignSGD optimizer (Bernstein et al., 2018), we would instead expect the weight entries to stabilize to $\pm \frac{1}{\lambda}$. Likewise, since SignSGD is simply AdamW with the momentum parameters set to zero, we would expect AdamW to also stabilize the weight entries to (roughly) $\pm \frac{1}{\lambda}$ as well.

A2.2. Why we may want to hold $\eta\lambda$ constant as we scale

Prior work has also (empirically) shown that if we want hyperparameters to transfer, we should hold the product $\eta\lambda$ constant (1) during hyperparameter search (Schaipp, 2024) and (2) as we scale the model via muP (Kosson et al., 2025). One possible reason why is that this is (roughly) equivalent to holding the relative update size constant with respect to the implicit weight norm constraint radius $R = 1/\lambda$.

Intuitively, this means that the amount of “effort” required to change an “absolutely yes” direction to an “absolutely no” direction remains the same no matter the model width.

A consequence of this is that, as recommended by Kosson et al. (2025), when scaling up the model width via muP, we should also scale up the weight decay by the same factor as we scale down the learning rate. But increasing the weight decay decreases the implicit constraint radius $1/\lambda$, and thus limits the search space of the weights. While this may be harmful in some settings (if the optimal point lies outside the radius), recent work has also shown that deliberately reducing the search space (but not too much) prevents the model from simply memorizing the training data early in training and thereby improves generalization later on in training (Tian, 2025).