Spectral-Aware Cautious Weight Decay for Muon
Cautious Weight Decay (CWD) resolves the directional conflict between weight decay and the optimizer’s update using a simple entry-wise mask. However, this approach is misaligned with matrix-aware optimizers such as Muon: its orthogonalized update carries a spectral structure that entry-wise decisions cannot capture, so spectral conflicts induced by weight decay can persist even under CWD, thereby hindering the optimization. We address this mismatch with a geometric formulation: we characterize the decay components spectrally compatible with Muon's update as a polar factor cone, and extract the nearest compatible component by projecting the weight matrix onto this cone. The Euclidean projection admits a clean closed-form solution via symmetric eigendecomposition, and can be efficiently approximated with Newton–Schulz iterations. The resulting Spectral-Aware Cautious Weight Decay (SA-CWD) is a drop-in modification to Muon's weight decay that structurally precludes spectral conflicts, ensuring full directional alignment of the update at the spectral level.
Weight Decay Through the Lens of Multi-Objective Optimization
Decoupled weight decay[1] is a classic optimization technique that plays an important role in large-scale model training. Yet its formulation has remained largely unchanged for years, and it is often adopted without a second thought. However, viewing weight decay through the lens of multi-objective optimization reveals a fundamental directional conflict. From this perspective, training with weight decay can be understood as jointly optimizing two objectives: minimizing the training loss and reducing the parameter norm. Although pursued jointly, these two objectives do not necessarily induce the same descent direction, leading to a potential conflict between them.
Let \(w_t \in \mathbb{R}^{n}\) denote the parameter vector, and let \(u_t\) denote the update direction produced by the optimizer. Standard decoupled weight decay updates the parameters as
\[ w_{t+1} = w_t - \eta_t\,(u_t + \lambda w_t), \]where \(\eta_t\) is the learning rate and \(\lambda\) is the weight decay coefficient.
The optimizer's update \(u_t\) is intended to decrease the training loss, while the decay term \(\lambda w_t\) pulls the parameters toward the origin. When these two directions are misaligned, the consequences go beyond merely slowing down the descent. Let \(g_t\) denote the gradient of the training loss at step \(t\). The update direction of the optimizer satisfies \(\langle g_t, u_t\rangle > 0\), so \(-\eta_t u_t\) is (locally) a descent step on the loss. If the decay term is misaligned with the gradient and its contribution outweighs the optimizer's, the combined direction satisfies \(\langle g_t, u_t + \lambda w_t\rangle < 0\). The update is then no longer a descent step — it performs gradient ascent on the training loss, pushing the loss up rather than down, and can severely impair the optimization process.
Cautious Weight Decay: A Simple Element-wise Modification
Cautious Weight Decay (CWD)[2] addresses precisely this issue with a remarkably simple modification: apply weight decay only to parameter coordinates where the decay direction agrees with the optimizer's descent direction, and skip decay on conflicting coordinates. Its update rule is
\[ w_{t+1} = w_t - \eta_t\left(u_t + \lambda\,\mathbb{I}(u_t \odot w_t \geq 0) \odot w_t\right), \]where \(\odot\) denotes the Hadamard (element-wise) product, and the indicator mask is defined coordinate-wise:
\[ \mathbb{I}(u_t \odot w_t \geq 0)_i = \begin{cases} 1, & u_{t,i}\,w_{t,i} \geq 0, \\ 0, & u_{t,i}\,w_{t,i} < 0. \end{cases} \]To understand CWD more clearly, consider a single parameter coordinate \(i\). The optimizer moves that coordinate by \(-\eta_t u_{t,i}\), while weight decay contributes \(-\eta_t \lambda w_{t,i}\). When \(u_{t,i}\) and \(w_{t,i}\) have the same sign, these two changes point in the same direction; when their signs differ, the decay entry opposes the optimizer's update, so CWD masks it out. In this way, CWD retains exactly the entries of the weight decay term whose signs agree with those of the optimizer's update, while the masked-out entries leave the update unimpeded on the conflicting coordinates.
CWD serves as a drop-in modification: it can be freely incorporated into a wide range of optimizers, such as Adam, Lion, and Muon. The change is lightweight — it introduces only a sign-based mask and no additional hyperparameters.
The Mismatch: Entry-wise CWD vs. Matrix-Aware Optimizers
Despite its generality, CWD defines alignment entirely through coordinate-wise sign comparisons, treating each parameter independently. Whether the optimizer is Adam, Lion, or Muon, the same element-wise rule is applied.
That is where an important distinction emerges. Adam and Lion construct updates through coordinate-wise operations, whereas Muon[3] orthogonalizes the momentum matrix to obtain its update direction. For Muon, the update is not merely a collection of independent coordinate-wise directions; rather, its coordinates collectively define an orthogonal geometric structure at the matrix level.
Applying an element-wise cautious mask to the weight decay term ignores this structure — a clear mismatch between the level at which the mask operates and the level at which the update is constructed. By selectively suppressing individual entries, the mask introduces perturbations that are inconsistent with this geometry, potentially undermining the properties established by orthogonalization. Moreover, entries that appear to conflict when considered individually may still be compatible with the update’s overall matrix-level geometry.
Since Muon’s update is built on the spectral directions extracted by orthogonalization, we next examine how weight decay can create conflicts precisely at this spectral level, which cannot be well resolved by entry-wise CWD. To make this concrete, we present such an example below.
The Spectral Conflict between Muon and Weight Decay
Recall that Muon constructs its update by orthogonalizing the momentum matrix. We illustrate the geometry with wide matrices: let \(M_t\in\mathbb{R}^{m\times n}\) with \(m\le n\) denote the momentum matrix, and consider its polar decomposition
\[ M_t = H_t O_t, \]where \(H_t=(M_tM_t^{\top})^{1/2}\in\mathbb{R}^{m\times m}\) is symmetric positive semidefinite and \(O_t\in\mathbb{R}^{m\times n}\) is the polar factor. Muon uses \(O_t\) as the update, and standard decoupled weight decay then updates the parameters as \[ W_{t+1} = W_t - \eta_t\left(O_t + \lambda W_t\right), \] where \(\eta_t\) is the learning rate and \(\lambda\) is the weight decay coefficient.
The weight decay term \(W_t\) carries its own polar factor, generally unrelated to \(O_t\). Since weight decay enters the update additively, it can dominate the orthogonalized update along some specific spectral direction and reverse the sign of the combined update: the iteration then moves in a direction that conflicts with that prescribed by orthogonalization, resulting in negative progress along that spectral direction. Because entry-wise sign decisions are blind to the global geometry of the matrix, conventional CWD fails to resolve such spectral conflicts. A simple example makes this concrete.
Take the polar factor \[ O_t = \frac{1}{7} \begin{bmatrix} 6 & -2 & -3\\ -2 & 3 & -6\\ -3 & -6 & -2 \end{bmatrix}. \]
Consider the weight matrix \[ W_t = 100 \begin{bmatrix} 10 & -15 & -2\\ -8 & 1 & -15\\ -1 & 15 & -8 \end{bmatrix}. \] Note that \(W_t\) has the same element-wise sign as \(O_t\) in eight out of nine entries. Only the entry at position \((3,2)\) has the opposite sign.
With \(\lambda = 0.01\), the combined update is \[ \begin{aligned} D_t &= O_t + \lambda W_t\\ &= \frac{1}{7} \begin{bmatrix} 6 & -2 & -3\\ -2 & 3 & -6\\ -3 & -6 & -2 \end{bmatrix} + \begin{bmatrix} 10 & -15 & -2\\ -8 & 1 & -15\\ -1 & 15 & -8 \end{bmatrix}. \end{aligned} \]
The polar factor of the combined update is approximately \[ O_{D_t} \approx \begin{bmatrix} 0.7248 & -0.5893 & -0.3568\\ -0.5502 & -0.1835 & -0.8146\\ 0.4146 & 0.7868 & -0.4573 \end{bmatrix}. \] Comparing the corresponding rows of \(O_{D_t}\) and \(O_t\), we obtain \[ \begin{aligned} \left\langle (O_{D_t})_{1,:},\, (O_t)_{1,:} \right\rangle &\approx 0.9426 > 0,\\ \left\langle (O_{D_t})_{2,:},\, (O_t)_{2,:} \right\rangle &\approx 0.7768 > 0,\\ \left\langle (O_{D_t})_{3,:},\, (O_t)_{3,:} \right\rangle &\approx -0.7214 < 0. \end{aligned} \] Thus, the third row reverses its direction relative to the original polar factor, while the first two rows remain aligned.
Now consider applying entry-wise CWD, which removes the weight decay component at the only entry whose sign conflicts with the corresponding entry of \(O_t\). The resulting update is \[ D_t^{\mathrm{CWD}} = O_t + \begin{bmatrix} 10 & -15 & -2\\ -8 & 1 & -15\\ -1 & 0 & -8 \end{bmatrix}. \] Its polar factor is approximately \[ O_{D_t^{\mathrm{CWD}}} \approx \begin{bmatrix} 0.4485 & -0.8787 & -0.1635\\ -0.6839 & -0.2196 & -0.6957\\ 0.5754 & 0.4238 & -0.6995 \end{bmatrix}. \] The corresponding row-wise inner products are \[ \begin{aligned} \langle (O_{D_t^{\mathrm{CWD}}})_{1,:},\, (O_t)_{1,:} \rangle &\approx 0.7056 > 0,\\ \langle (O_{D_t^{\mathrm{CWD}}})_{2,:},\, (O_t)_{2,:} \rangle &\approx 0.6976 > 0,\\ \langle (O_{D_t^{\mathrm{CWD}}})_{3,:},\, (O_t)_{3,:} \rangle &\approx -0.4100 < 0. \end{aligned} \] The directional conflict therefore persists even after the sole element-wise sign conflict has been removed.
This example illustrates a limitation of entry-wise CWD: its decisions are based on individual element-wise signs and do not account for the global geometry of the matrix. Consequently, removing locally conflicting decay components does not necessarily resolve the spectral conflicts induced by the addition of weight decay. Such conflicting spectral directions can pull the parameters in the wrong direction and substantially hinder the optimization process.
The remedy is to design a matrix-aware CWD variant that is aligned with the spectral structure of Muon’s orthogonalized update.
The Polar Factor Cone: What is Kept
To arrive at our new design, we first distill the underlying principle of entry-wise CWD. Stripped to its essence, the cautious rule is a construction by retention: sign agreement serves as the criterion of alignment, and this criterion determines what is kept — the decay components that do not conflict with the optimizer update. This suggests a broader design perspective, in which the focus shifts from how to mask to what deserves to be retained. We apply this perspective to Muon by first designing a retained weight decay term that is consistent with Muon’s orthogonalization geometry.
To operationalize this consistency, we introduce a geometric cone that characterizes the weight decay components eligible for retention. Specifically, we require the retained decay component to share the same polar factor as Muon's update. For wide matrices, this leads to the following polar factor cone \[ \mathcal{K}_{O_t} = \bigl\{ S O_t: \,S = S^{\top}\!,\, S \succeq 0 \bigr\}, \] where \(S\in\mathbb{R}^{m\times m}\) is symmetric positive semidefinite (PSD) and determines the metric component of the retained decay term. Note that we explicitly state the symmetry condition \(S=S^\top\) here for clarity, although it is already implied by \(S\succeq0\).
Why It Is a Cone
The set \(\mathcal{K}_{O_t}\) naturally forms a closed convex cone. It is closed under nonnegative scaling, since \(\alpha S\succeq 0\) for any \(\alpha\ge 0\), and under addition, since the sum of two PSD matrices is also PSD. Geometrically, \(\mathcal{K}_{O_t}\) is the linear image of the PSD cone under the map \(S\mapsto S O_t\). Moreover, because \(O_tO_t^{\top}=I_m\), the mapping \(S\mapsto S O_t\) is injective and preserves the Frobenius norm (\(\|S O_t\|_F=\|S\|_F\)). This isometric property guarantees that the map is bounded below, ensuring that the linear image of the closed PSD cone remains closed.
Why the PSD Constraint Eliminates Spectral Conflicts
We now explain in detail why the construction above eliminates spectral conflicts. By construction, \(S\succeq 0\) ensures that \(S O_t\) admits a polar decomposition with \(O_t\) as a valid polar factor. Since \(S\) is symmetric positive semidefinite, all its eigenvalues \(\sigma_i(S)\) are non-negative. Moreover, using \(O_tO_t^{\top}=I_m\) and the symmetry of \(S\), we obtain \[ (S O_t)(S O_t)^{\top} = S O_tO_t^{\top} S = S^2. \] By definition, the singular values of \(S O_t\) are the square roots of the eigenvalues of \(S^2\), which are exactly the eigenvalues of \(S\). The decay term therefore introduces only non-negative gains along the spectral directions encoded by \(O_t\), without introducing spectral conflicts.
This property extends directly to the combined update: \[ O_t+\lambda S O_t=(I_m+\lambda S)\,O_t. \] For \(\lambda\ge 0\), we have \(I_m+\lambda S\succ 0\). Consequently, the combined update retains \(O_t\) as its polar factor, with singular values \(1+\lambda\sigma_i(S)\ge 1\). Thus, in this construction, spectral conflicts — including the possibility of reversing the polar factor — are structurally precluded.
Euclidean Projection onto the Polar Factor Cone
Having established \(\mathcal{K}_{O_t}\) as the structural safe zone that completely precludes spectral conflicts, we now illustrate how to extract this compatible component in practice. Given the raw weight decay matrix \(W_t\), we wish to retain only the part that is spectrally aligned with Muon's update. Because \(\mathcal{K}_{O_t}\) is a closed convex cone, this operation admits a principled formulation: the Euclidean projection onto the polar factor cone
\[ D_t^\star = \operatorname*{argmin}_{D \in \mathcal{K}_{O_t}} \frac{1}{2} \|D - W_t\|_F^2. \] The closedness of the cone guarantees that the projection is well-defined, and its convexity guarantees uniqueness.
Viewed in this light, the projection acts as a spectral-aware matrix filter: \(D_t^\star\) represents the closest matrix to \(W_t\) that lives entirely within the safe zone, while any off-frame, conflicting components are projected away. Unlike element-wise masks, this approach enforces spectral compatibility directly in matrix space through a principled geometric optimization.
Closed-Form Solution via Projection onto the PSD Cone
The Euclidean projection onto \(\mathcal{K}_{O_t}\) admits a clean closed-form solution, which we derive below. We begin by substituting the cone parameterization \( D = SO_t \) into the Euclidean projection objective. The projection problem then becomes \[ \min_{S \succeq 0}\ \frac{1}{2}\left\| SO_t - W_t \right\|_F^2, \] where \( O_t \in \mathbb{R}^{m \times n} \), \( m \le n \), satisfies \( O_t O_t^\top = I_m \).
Step 1: Expand the Objective
Using the identity \( \|A\|_F^2 = \operatorname{tr}(AA^\top) \), expanding the trace and using the symmetry of \(S\), we obtain
\[ \begin{aligned} \left\| SO_t - W_t \right\|_F^2 &= \operatorname{tr}\!\big( (SO_t - W_t)(SO_t - W_t)^\top \big)\\ &= \operatorname{tr}(SO_t O_t^\top S) - \operatorname{tr}(SO_t W_t^\top) - \operatorname{tr}(W_t O_t^\top S) + \operatorname{tr}(W_t W_t^\top). \end{aligned} \] Then using the invariance of the trace under transposition and the symmetry of \(S\), the two cross terms are equal \[ \operatorname{tr}(SO_tW_t^\top) =\operatorname{tr}((SO_tW_t^\top)^\top) = \operatorname{tr}(W_tO_t^\top S^\top) =\operatorname{tr}(W_tO_t^\top S). \] Moreover, since \( O_t O_t^\top = I_m \) and \( S = S^\top \), \[ \operatorname{tr}(SO_t O_t^\top S) = \operatorname{tr}(S^2). \] Combining these results with the cyclic property of trace \(\operatorname{tr}(W_tO_t^\top S) = \operatorname{tr}(SW_tO_t^\top)\), we arrive at \[ \left\| SO_t - W_t \right\|_F^2 = \operatorname{tr}(S^2) - 2\operatorname{tr}(SW_tO_t^\top) + \left\| W_t \right\|_F^2. \] This expansion isolates the dependence on \( O_t \) in the cross term \( W_t O_t^\top \), an \( m \times m \) matrix that is not necessarily symmetric.Step 2: Extract the Symmetric Component
To simplify notation, let \(A = W_t O_t^\top\). Since every square matrix admits a unique decomposition into symmetric and skew-symmetric components, we can write
\[ A = X_s + X_a, \]where
\[ X_s = \frac{A + A^\top}{2}, \qquad X_a = \frac{A - A^\top}{2}. \]By construction, \(X_s\) is symmetric \((X_s^\top = X_s)\), and \(X_a\) is skew-symmetric \((X_a^\top = -X_a)\). Substituting \(A = X_s + X_a\) into the cross term, we obtain
\[ \operatorname{tr}(SW_tO_t^\top) = \operatorname{tr}(SA) = \operatorname{tr}(SX_s) + \operatorname{tr}(SX_a). \]Since \(S\) is symmetric and \(X_a\) is skew-symmetric, we have
\[ \operatorname{tr}(SX_a) = \operatorname{tr}\big((SX_a)^\top\big) = \operatorname{tr}(X_a^\top S^\top) = -\operatorname{tr}(X_a S) = -\operatorname{tr}(SX_a). \]This directly implies that \(\operatorname{tr}(SX_a) = 0\), and hence
\[ \operatorname{tr}(SA) = \operatorname{tr}(SX_s), \]where
\[ X_s = \frac{W_t O_t^\top + O_t W_t^\top}{2}. \]Thus, although \(W_tO_t^\top\) need not be symmetric, its skew-symmetric component contributes nothing to the cross term. Consequently, the optimization objective depends on \(W_tO_t^\top\) only through its symmetric part \(X_s\).
Step 3: Complete the Square and Reduce to PSD Projection
Substituting the cross-term identity into the expanded objective gives
\[ \begin{aligned} \frac{1}{2}\left\| SO_t - W_t \right\|_F^2 &= \frac{1}{2}\operatorname{tr}(S^2) - \operatorname{tr}(SX_s) + \frac{1}{2}\left\| W_t \right\|_F^2. \end{aligned} \]Since both \(S\) and \(X_s\) are symmetric, we have
\[ \left\| S - X_s \right\|_F^2 = \operatorname{tr}\!\big( (S - X_s)^2 \big) = \operatorname{tr}(S^2) - 2\operatorname{tr}(SX_s) + \operatorname{tr}(X_s^2). \]Rearranging terms yields the equivalence
\[ \frac{1}{2}\operatorname{tr}(S^2) - \operatorname{tr}(SX_s) = \frac{1}{2}\left\| S - X_s \right\|_F^2 - \frac{1}{2}\left\| X_s \right\|_F^2. \]Therefore, the original objective admits the equivalent representation
\[ \frac{1}{2}\left\| SO_t - W_t \right\|_F^2 = \frac{1}{2}\left\| S - X_s \right\|_F^2 + \frac{1}{2}\left\| W_t \right\|_F^2 - \frac{1}{2}\left\| X_s \right\|_F^2. \]The last two terms depend only on the given matrices \(W_t\) and \(O_t\) and are independent of the optimization variable \(S\). Removing these additive constants leaves the minimizers unchanged. Consequently, the original problem reduces to the following equivalent problem
\[\min_{S \succeq 0} \frac{1}{2}\left\| S - X_s \right\|_F^2. \]This is precisely the Euclidean projection of the symmetric matrix \(X_s\) onto the cone of positive semidefinite matrices.
Step 4: Obtain the Closed-Form Solution
The reduced projection problem derived above is a classical problem in matrix computation and optimization, which admits an elegant closed-form solution[4]. For completeness, a detailed derivation of the result is provided in the Appendix.
Specifically, since \(X_s \in \mathbb{R}^{m \times m}\) is symmetric, let its eigendecomposition be
\[ X_s = Q \operatorname{diag}(\mu_1, \ldots, \mu_m) Q^\top, \]where \(Q\) is orthogonal and \(\mu_i\) are the eigenvalues of \(X_s\). The solution to the reduced PSD projection problem is obtained by replacing all negative eigenvalues with zero
\[ S_t^\star = Q \operatorname{diag}\big(\max(\mu_1, 0), \ldots, \max(\mu_m, 0)\big) Q^\top. \]Consequently, the original Euclidean projection onto the polar factor cone is given by
\[ D_t^\star = S_t^\star O_t. \]This is a clean closed-form solution: compute the symmetric cross-alignment matrix \(X_s\), project it onto the PSD cone by eigenvalue thresholding, and multiply the result by \(O_t\). This yields the nearest matrix in \(\mathcal{K}_{O_t}\) to the weight matrix \(W_t\) under the Frobenius norm.
Putting It Together: Spectral-Aware Cautious Weight Decay
Substituting the projected decay direction into the Muon update yields \[ W_{t+1} = W_t - \eta_t \left( O_t + \lambda D_t^{\star} \right). \] We term this approach as Spectral-Aware Cautious Weight Decay (SA-CWD). The complete computational procedure is summarized below.
Spectral-Aware Cautious Weight Decay
# W: parameter matrix (m x n)
# O: Muon polar direction (m x n)
# lr: learning rate, wd: weight decay coefficient
if m <= n:
A = sym(W @ O.T) # (m x m) symmetric matrix
D = project_psd(A) @ O
else:
A = sym(O.T @ W) # (n x n) symmetric matrix
D = O @ project_psd(A)
W = W - lr * (O + wd * D)
The two helper functions used in the above algorithm are implemented as follows. In particular, the PSD projection is computed via direct symmetric eigendecomposition.
def sym(X):
return 0.5 * (X + X.T)
def project_psd(A):
# A is guaranteed to be symmetric
eigvals, eigvecs = torch.linalg.eigh(A)
eigvals = eigvals.clamp_min(0)
return (eigvecs * eigvals.unsqueeze(-2)) @ eigvecs.T
Note that in the above algorithm, we directly present the result for tall matrices. Its derivation follows analogously to that for wide matrices; thus, we omit the proof for brevity. It is also worth noting that in both regimes, the symmetric eigendecomposition is confined to a small square matrix of size \(\min(m, n)\).
Back to the Example: SA-CWD Successfully Resolve the Spectral Conflict
Now, we revisit the localized conflict example to illustrate how SA-CWD successfully addresses the spectral conflict that entry-wise CWD cannot well resolve. We first construct the interaction matrix \[ A=W_tO_t^T= \begin{bmatrix} 1371.43 & -757.14 & 914.29\\ -71.43 & 1557.14 & 685.71\\ -171.43 & 1357.14 & -1014.29 \end{bmatrix}. \] Taking its symmetric part gives \[X_s=\frac{1}{2}(A+A^T)= \begin{bmatrix} 1371.43 & -414.29 & 371.43\\ -414.29 & 1557.14 & 1021.43\\ 371.43 & 1021.43 & -1014.29 \end{bmatrix}. \] Its eigendecomposition yields the eigenvalues \[ \Sigma\approx\operatorname{diag}(2035.14,\;1334.39,\;-1455.24). \] We project onto the PSD cone by clamping negative eigenvalues to zero \[ \Sigma^+=\operatorname{diag}(2035.14,\;1334.39,\;0), \] which further gives \[ S^*_t=Q\Sigma^+Q^T \approx \begin{bmatrix} 1414.02 & -330.28 & 140.96\\ -330.28 & 1722.81 & 566.91\\ 140.96 & 566.91 & 232.69 \end{bmatrix}. \] Mapping the projected interaction back to the weight space yields \[ D^*_t=S^*_tO_t \approx \begin{bmatrix} 1245.98 & -666.38 & -363.19\\ -1018.29 & 346.79 & -1497.12\\ -140.88 & 3.24 & -612.82 \end{bmatrix}. \] For \(\lambda=0.01\), the combined update is \[ D_t^{\mathrm{PSD}}=O_t+\lambda D^*_t \approx \begin{bmatrix} 13.3169 & -6.9495 & -4.0605\\ -10.4686 & 3.8965 & -15.8283\\ -1.8374 & -0.8247 & -6.4139 \end{bmatrix}. \] The row-wise inner products between its polar factor and the original polar factor are approximately \[ \begin{aligned} \langle (O_{D_t^{\mathrm{PSD}}})_{1,:},\, (O_t)_{1,:} \rangle &\approx 0.99999996 > 0,\\ \langle (O_{D_t^{\mathrm{PSD}}})_{2,:},\, (O_t)_{2,:} \rangle &\approx 0.99999993 > 0,\\ \langle (O_{D_t^{\mathrm{PSD}}})_{3,:},\, (O_t)_{3,:} \rangle &\approx 0.99999998 > 0. \end{aligned} \] In fact, this alignment is exact rather than approximate. By construction, the polar factor of \(D_t^{\mathrm{PSD}}\) coincides exactly with \(O_t\). The tiny deviations from \(1.0\) here stem only from floating-point round-off.
In contrast to entry-wise CWD, whose row-wise inner products is \(\{0.71,\ 0.70,\ -0.41\}\), SA-CWD restores positive alignment in all three rows. By projecting the interaction matrix onto the PSD cone in the spectral domain, SA-CWD resolves the spectral conflict globally rather than entry by entry, thereby preserving the direction of every row of the original polar factor throughout the update.A Newton–Schulz Approximation
Although symmetric eigendecomposition is more efficient than general eigendecomposition, it remains computationally expensive, especially when applied repeatedly during training. This motivates us to seek a more efficient approximation to the PSD projection.
Recall that for a symmetric matrix \(A\), its PSD projection admits the equivalent representation
\[ \operatorname{PSD}(A)=\frac{1}{2}(A+(A^2)^{1/2}). \]To see this, consider the eigendecomposition of \(A\). Since \(A \in \mathbb{R}^{m \times m}\) is real symmetric, it admits the following eigendecomposition
\[ A=Q\,\operatorname{diag}(\mu_1,\ldots,\mu_m)\,Q^\top, \]where \(Q\) is an orthogonal matrix and \(\mu_1,\ldots,\mu_m\) are the real eigenvalues. Consequently, we have
\[ A^2=Q\,\operatorname{diag}(\mu_1^2,\ldots,\mu_m^2)\,Q^\top. \]Since \(A^2 \succeq 0\), its unique principal square root is
\[ (A^2)^{1/2} =Q\,\operatorname{diag}\bigl(|\mu_1|,\ldots,|\mu_m|\bigr)\,Q^\top. \]Combining these expressions, we obtain
\[ \begin{aligned} \frac12\left(A+(A^2)^{1/2}\right) &=Q\,\operatorname{diag}\left( \frac{\mu_1+|\mu_1|}{2},\ldots, \frac{\mu_m+|\mu_m|}{2} \right)Q^\top\\[4pt] &=Q\,\operatorname{diag}\bigl( \max(\mu_1,0),\ldots,\max(\mu_m,0) \bigr)Q^\top, \end{aligned} \]where the last equality follows from the scalar identity
\[ \frac{x+|x|}{2}=\max(x,0). \]Since the PSD projection is obtained by replacing all negative eigenvalues of \(A\) with zero while retaining its non-negative eigenvalues, we conclude that
\[ \operatorname{PSD}(A)=\frac12(A+(A^2)^{1/2}). \]This formulation suggests that we can approximate the PSD projection without explicitly computing the eigendecomposition of \(A\): it suffices to approximate the matrix square root of \(A^2\). We employ the Newton–Schulz iteration, which approximates matrix square roots using iterative matrix multiplications and is therefore well suited to GPU execution.
Alternatively, we can avoid computing the square root of \(A^{2}\) by invoking the identity \[ \operatorname{PSD}(A) \;=\; \tfrac{1}{2}\,A\,\bigl(I+\operatorname{sign}(A)\bigr), \label{eq:psd-sign} \] where \(\operatorname{sign}(A)\) denotes the matrix sign function. The sign function admits the Newton–Schulz iteration \[ X_{k+1} \;=\; \tfrac{1}{2}\,X_k\,\bigl(3I - X_k^{2}\bigr).\] The polynomial coefficients in this iteration can be further tuned to improve both convergence speed and numerical stability. We leave the implementation details and experimental comparisons of these different approaches as future work.
Conclusion
In this work, we propose Spectral-Aware Cautious Weight Decay, a principled cautious weight decay mechanism tailored for the spectral optimizer Muon. Our starting point is simple: as optimization techniques evolve to become matrix-aware, their accompanying regularization must evolve accordingly. We demonstrated that naive weight decay can directionally conflict with Muon—flipping the spectral directions of the polar factor and misguiding parameter updates—while entry-wise cautious weight decay fails to resolve such global conflicts. Initial small-scale experiments demonstrate promising results of SA-CWD, with larger-scale evaluations currently underway. We hope this work serves as a stepping stone toward more principled matrix-aware regularizers: as optimizers have learned to be matrix-aware, regularizers must follow suit.
Cite This Article
The paper is in preparation and will be released soon.
References
- I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019. [link]
- L. Chen, J. Li, K. Liang, B. Su, C. Xie, C. Liang, N. Lao, and Q. Liu. Cautious weight decay. In International Conference on Learning Representations (ICLR), 2026. [link]
- K. Jordan. Muon: An optimizer for hidden layers in neural networks. 2024. [link]
- N. J. Higham. Computing a nearest symmetric positive semidefinite matrix. Linear Algebra and its Applications, 103:103–118, 1988. [link] ↩
Appendix: Derivation of the Projection onto the PSD Cone
Consider the optimization problem
\[ \min_{S \succeq 0}\ \frac{1}{2} \left\| S - X \right\|_F^2, \]where \(X \in \mathbb{R}^{m \times m}\) is a real symmetric matrix. Let
\[ X = Q \operatorname{diag}(\mu_1, \ldots, \mu_m)\, Q^\top \]be the eigendecomposition of \(X\), where \(Q \in \mathbb{R}^{m \times m}\) is an orthogonal matrix satisfying \(Q^\top Q = Q Q^\top = I_m\), and \(\mu_1, \ldots, \mu_m\) are the corresponding eigenvalues. Substituting this eigendecomposition into the objective yields
\[ \frac{1}{2} \left\| S - X \right\|_F^2 = \frac{1}{2} \left\| S - Q \operatorname{diag}(\mu)\, Q^\top \right\|_F^2, \]where \(\mu = (\mu_1, \ldots, \mu_m)\). To simplify this expression, we leverage the orthogonal invariance of the Frobenius norm. Recall that for any matrix \(A\) and any orthogonal matrices \(U\) and \(V\), the Frobenius norm satisfies \(\|U \!A V\|_F = \|A\|_F\). This property stems directly from the definition of the Frobenius norm via the trace, and the cyclic property of the trace
\[ \left\| U \!A V \right\|_F^2 = \operatorname{tr}\left( U\! A V V^\top\! A^\top U^\top \right) = \operatorname{tr}\left( U \!A A^\top U^\top \right) = \operatorname{tr}\left( A A^\top \right) = \left\| A \right\|_F^2. \]Multiplying the matrix inside the norm in the objective by \(Q^\top\!\) on the left and \(Q\) on the right leaves the Frobenius norm invariant
\[ \begin{aligned} \left\| S - Q \operatorname{diag}(\mu)\, Q^\top \right\|_F^2 &= \left\| Q^\top\!\left( S - Q \operatorname{diag}(\mu)\, Q^\top \right) Q \right\|_F^2 \\ &= \left\| Q^\top\! S Q - \operatorname{diag}(\mu) \right\|_F^2. \end{aligned} \]This naturally motivates a change of variables by defining \(\widetilde{S} = Q^\top\! S Q\) (or equivalently, \(S = Q \widetilde{S} Q^\top\)). Since \(Q\) is orthogonal, the positive semidefiniteness constraint is preserved: \(S \succeq 0\) if and only if \(\widetilde{S} \succeq 0\). This holds because for any vector \(x \in \mathbb{R}^m\), we have \(x^\top\! S x = (Q^\top x)^\top\! \widetilde{S} (Q^\top x)\), and as \(x\) varies over \(\mathbb{R}^m\), the mapped vector \(Q^\top x\) also spans the entire space \(\mathbb{R}^m\). Consequently, the original problem is equivalently transformed into
\[ \min_{\widetilde{S} \succeq 0}\ \frac{1}{2} \left\| \widetilde{S} - \operatorname{diag}(\mu) \right\|_F^2. \]Expanding the objective gives
\[ \frac{1}{2} \left\| \widetilde{S} - \operatorname{diag}(\mu) \right\|_F^2 = \frac{1}{2} \sum_{i=1}^m (\widetilde{S}_{ii} - \mu_i)^2 + \frac{1}{2} \sum_{i \ne j} \widetilde{S}_{ij}^2. \]Since \( \widetilde{S} \succeq 0 \), its diagonal entries satisfy \( \widetilde{S}_{ii} \ge 0 \). Thus, each diagonal term is minimized over the non-negative domain \( [0, \infty) \). Specifically, there are two cases.
- If \( \mu_i > 0 \), the unconstrained minimizer \( \widetilde{S}_{ii} = \mu_i \) is feasible, yielding a minimum value of \(0\).
- If \( \mu_i \le 0 \), the closest feasible value to \( \mu_i \) is \( \widetilde{S}_{ii} = 0 \), yielding a minimum value of \(\frac{1}{2}\mu_i^2\).
Meanwhile, all off-diagonal contributions are nonnegative. Combining these observations gives the lower bound
\[ \frac{1}{2} \left\| \widetilde{S} - \operatorname{diag}(\mu) \right\|_F^2 \ \ge\ \frac{1}{2} \sum_{i=1}^m \min(\mu_i, 0)^2. \]Now consider the candidate matrix
\[ \widetilde{S}^\star = \operatorname{diag}\bigl( \max(\mu_1, 0), \ldots, \max(\mu_m, 0) \bigr). \]This matrix is positive semidefinite because its diagonal entries are non-negative and its off-diagonal entries vanish. Moreover, it explicitly attains the established lower bound, as every diagonal term is minimized and all off-diagonal contributions are zero. Hence, \(\widetilde{S}^\star\) serves as a global minimizer. Since the objective function is strictly convex and the PSD cone is a convex set, this minimizer is unique.
Mapping back to the original space via \(S^\star = Q \widetilde{S}^\star Q^\top\) yields the unique Euclidean projection
\[ S^\star = Q\,\operatorname{diag}\bigl( \max(\mu_1, 0), \ldots, \max(\mu_m, 0) \bigr)\,Q^\top. \]Consequently, the Euclidean projection onto the PSD cone admits a clean closed-form solution obtained via eigenvalue thresholding, which sets all negative eigenvalues to zero.