1: Introduction

Attention is a test-time regressor

Attention is a retrieval mechanism. Given a query token, the model has a context of key-value pairs and returns a weighted combination of values, where the weights reflect how well each key matches the query. In softmax attention, those weights are always non-negative and sum to one, so the output is always a convex combination of the value vectors in context. You can concentrate weight on one value or spread it across many, but the output stays inside the convex hull of what's there.

This is nonparametric regression over the KV pairs [3]. Each forward pass fits a function to the context → keys are inputs, values are labels, the query is the test point. Different attention designs correspond to different choices of what function class to fit.

Softmax attention [4] is the Nadaraya–Watson estimator [5] [6]. At each query, it fits a constant to the nearby data by taking a kernel-weighted average of the values. The hypothesis space contains only constant functions, with one built fresh for each query.

Linear attention [7] sits at the other extreme. It fits a single global linear function to all the data, which means the same $W$ must serve every query. These design choices, particularly the choice of hypothesis space, fundamentally impact the associative memory capacity of each mechanism.

Limitations of Linear & Softmax Attention

The two dominant attention families sit at opposite ends of the regression spectrum, and each has a characteristic failure mode.

Linear attention fits a single global affine function $f(x)=Wx+b$ to the entire context. The same function must serve every query, which means it cannot adapt its prediction based on where in the context the query is looking. If the true relationship between keys and values varies across the sequence, or is nonlinear, the global fit must be wrong somewhere. This is misspecification error whereby the hypothesis class is too restrictive to faithfully perform query-dependent retrieval.

Softmax attention avoids misspecification by building a fresh local estimate at each query but the local constant fit has its own problem. Because the output is a convex combination of the values in the context window, it can only interpolate between them and it cannot extrapolate outside the data it has already seen.

This is boundary bias [2], a known property of the Nadaraya–Watson estimator. When the query sits near the edge of the key distribution, where keys are dense on one side and sparse on the other, the kernel-weighted average is systematically biased inward. The error is not a consequence of finite data or insufficient training, but structural and inherent to local constant regression.

2: Parallax & Local Linear Attention

Local linear attention (LLA) [2] addresses the global-linear misspecification and local-constant boundary-bias failures by fitting a local linear function at each query instead of a local constant. As a result, it has strictly smaller integrated MSE than both Softmax and Linear Attention (Appendix A). We provide three perspectives on LLA:

Canonical Formulation

LLA fits a local linear estimator $\mathcal F(\boldsymbol q_i) = \{\boldsymbol b + \boldsymbol W(\boldsymbol x - \boldsymbol q_i)\}$ with the softmax kernel weight $w_{ij} = \exp(\boldsymbol q_i^\top \boldsymbol k_j / h)$ and ridge regularization $\lambda \|\boldsymbol W\|_F^2$. Let $\boldsymbol z_{ij} = \boldsymbol k_j - \boldsymbol q_i$, $\omega_i = \sum_{j\le i} w_{ij}$, $\boldsymbol\mu_i = \sum_{j\le i} w_{ij} \boldsymbol z_{ij}$, and $\boldsymbol\Sigma_i = \sum_{j\le i} w_{ij} \boldsymbol z_{ij} \boldsymbol z_{ij}^\top + \lambda \boldsymbol I$. The prediction at $\boldsymbol q_i$ is:

LLA is a query-centered second-order correction to softmax attention that leverages the geometry of key vectors around the query. It produces a better prediction when the keys are not uniformly distributed under the softmax weighting, e.g. at the boundaries.

Generalization of MesaNet & Attention

LLA can also be read as constructing query-dependent states through the kernel, in contrast to the global states in MesaNet [8]. LLA can degenerate to both mechanisms by tuning $\lambda$ and $h$:

Additive Correction to Softmax

The most useful formulation rewrites LLA as softmax attention plus a correction term. Write $\boldsymbol p_i = \mathrm{softmax}(\boldsymbol K_i \boldsymbol q_i / h)$, $t_{ij} = \boldsymbol\rho_i^{\star\top} \boldsymbol z_{ij}$, and $\bar t_i = \mathbb E_{\boldsymbol p_i}[t_{ij}]$. Then

where $\boldsymbol\Sigma_{KV}^{(i)} = \mathbb E_{\boldsymbol p_i}[(\boldsymbol v_j - \bar{\boldsymbol v}_i)(\boldsymbol k_j - \bar{\boldsymbol k}_i)^\top]$ is the softmax-weighted KV-covariance and $\eta_i = \bar t_i / (1 - \bar t_i)$ is the boundary amplification factor.

The quantity $\eta_i$ is non-negative and measures the Mahalanobis distance from the query to the key center under $\boldsymbol p_i$. Intuitively, if $\eta_i \approx 0$, the query is close to the weighted key center and the correction becomes pure covariance; if $\eta_i \gg 1$, the query is close to the weighted key boundary and the correction is amplified to compensate for the boundary bias. The full derivation can be found in Appendix B.

Difficulties in Scaling LLA

LLA has the right theoretical properties but is difficult to use at scale. The exact forward requires solving a linear system $\boldsymbol\Sigma_i\boldsymbol{x} = \boldsymbol{\mu}_i$ independently at every query position, typically with a parallel conjugate gradient (CG) solver. Three problems follow.

• Intensive I/O. The CG iteration requires $TLd$ memory accesses in the forward pass (where $T$ is the iteration number), dominating the $2Ld$ memory access of Softmax Attention.
• Regularization–expressiveness tradeoff. Large ridge regularization coefficient $\lambda$ ensures $\boldsymbol\Sigma_i \succ 0$ but drives $\boldsymbol\rho_i^\star \to \boldsymbol 0$, collapsing LLA back to Softmax. Small $\lambda$ preserves the correction but risks ill-conditioning and instability. We find it nontrivial to balance the tradeoff in practical pretraining settings.
• Low-precision incompatibility. The stability of CG is sensitive to the precision format, while modern hardware and computation primitives are increasingly shaped around reduced precision.

Parallax

Parallax keeps the local-linear principle and removes the solve in favor of learned parametrization.

Parallax eliminates the per-query solve of $\boldsymbol\rho_i^\star$ by learning a direct mapping from the layer input. Let $\boldsymbol x_i$ be the input to the layer, and we parameterize

where $\boldsymbol W_R \in \mathbb R^{d_{\text{qk}} \times d}$ is a learnable projection matrix.

Parallax additionally sets $\eta_i = 0$ to remove the boundary amplification, yielding the forward equation

Empirically, we find removing $\eta_i$ is necessary to avoid collapse. Because $\boldsymbol\rho_i = \boldsymbol W_R \boldsymbol x_i$ is no longer the constrained solution of the exact LLA, the denominator can approach zero or flip sign, causing $\eta$ to explode.

Parallax in one line: the softmax output minus a learned projection of the local KV covariance.

Perspectives on Parallax

Wide bandwidth limit. Taking $h \to \infty$ makes the kernel weight $w_{ij} = \exp(\boldsymbol q_i^\top \boldsymbol k_j / h) \to 1$ uniformly over all $j$. The local softmax-weighted statistics become global running averages, and the three nonparametric mechanisms reduce to the affine variants:

All three share the output template (OLS regression with intercept) and differ only in the probe. Softmax uses a zero probe (returning just the optimal intercept), Parallax uses a learned probe, and LLA solves for the optimal one. Standard linear attention and MesaNet are the same template with the intercept dropped, which collapses the centered moments to their raw counterparts.

This framing also clarifies the dual role of the query across the family. In softmax attention $q_i$ shapes the kernel weights and defines where attention concentrates. In the linear attention family, what is conventionally called the query is actually the probe → a directional readout from the recurrent state. In MesaNet and LLA the probe is fully determined by other statistics, while in Parallax it is learned.

Strong regularization limit. As $\lambda \to \infty$ or $\|\boldsymbol\rho_i\| \to 0$, the probe is suppressed and the intercept dominates the output. Both Parallax and LLA degenerate identically to Softmax Attention. The same parametrization axis explains the relationship between Parallax and LLA, just as MesaNet differs from Linear Attention by probe preconditioning.

3: Hardware-aligned Implementation

One extra branch, no extra I/O

The Parallax forward has the form

which involves the softmax output and a softmax-weighted covariance. Both terms depend on the same KV pairs and the question is whether we can compute them together in a single streaming pass, without materializing the full attention matrix which is the same constraint FlashAttention [13] solves for standard softmax.

The answer is yes. We can expand the Parallax forward into a form that streams over KV blocks:

This decomposes into two parallel branches that process the same KV tiles. The softmax branch maintains the usual FA running state $(\mathbf m_r, \mathbf d_1, \mathbf O_1)$, and the covariance branch maintains $(\mathbf d_2, \mathbf O_2)$. In each loop, the covariance branch uses the same $\mathbf K_c, \mathbf V_c$ to compute the unnormalized scores $\mathbf S_2 = \mathbf R_r \mathbf K_c^\top$, fuses them with the softmax weights as $\mathbf P_2 = \mathbf P_1 \odot \mathbf S_2$, and then accumulates $\mathbf O_2 = \mathbf P_2 \mathbf V_c$ alongside $\mathbf O_1$. The final output combines the two running sums.

Both branches share the online maximum $\mathbf m_r$, the rescaling factor $\boldsymbol\alpha$, and the $\mathbf K_c, \mathbf V_c$ tiles. Therefore Parallax does not require extra I/O in each iteration and, as a result, roughly doubles the arithmetic intensity over standard FA.

Optimized Decode Kernel

We prototype the Parallax decode kernel in CuTeDSL [16] on NVIDIA Hopper GPUs. The design exploits the structural property that the two branches share the same KV stream, so on I/O-bound decode workloads it consumes essentially no additional HBM traffic. We report further kernel optimizations in Appendix C.

Profiling. We profile the prototype kernel against FA2 [14] and FA3 [15] on H200 GPUs at BF16 precision. Because Parallax doubles the arithmetic intensity, FA cannot be matched on both FLOPs and HBM traffic simultaneously. We therefore report two settings:

• Compute-matched setting. Parallax uses $d_h = 64$ such that it matches the FLOPs of FA at $d_h = 128$.
• I/O-matched setting. Both kernels use $d_h = 128$ and Parallax doubles the compute of FA with the same HBM traffic.

The prototype Parallax kernel matches or outperforms FA across all configurations in both settings.

4: Empirical Results

We empirically validate Parallax on both synthetic and language modeling benchmarks. We compare Parallax against Softmax Attention (Attn, Transformer) [4], Mamba [12], Gated DeltaNet (GDN) [10], MesaNet (Mesa) [8], and Kimi DeltaAttention (KDA) [11].

Synthetic: MAD benchmark

We evaluate Parallax on the MAD-Benchmark [20], which consists of six synthetic tasks designed to evaluate the core ability of sequence mixers. All models follow a two-layer architecture with sequence mixer and MLP blocks interleaved. Different from prior work, all models are trained with the Muon optimizer [17].

Parallax consistently improves on the recall-oriented tasks (ICR, FCR, NCR, SC) while remaining competitive on the compression and memorization tasks (CMP, MEM), and attains the highest overall accuracy.

To further showcase the advantage of Parallax under more challenging recall conditions, we synthesize an additional set of harder tasks by scaling up the KV pairs and the sequence length on ICR, NCR, and SC, with vocabulary size and context length stressed up to 512 and 2048 respectively. Parallax retains accuracy as the difficulty grows, whereas other baselines degrade dramatically, most visibly on SC at the longest context lengths.

Small-Scale Pretraining

We adopt the Qwen-3 architecture [21] (which applies RMSNorm [27] to $\boldsymbol q$ and $\boldsymbol k$ vectors) as implemented in TorchTitan [22], and additionally apply RMSNorm to $\boldsymbol\rho$ in each Parallax layer.

At the 0.6B scale we also include KDA, GDN, and two controlled experiment baselines:

• Parameter-matched Transformer (Transformer^†). Parallax introduces extra parameters from the $\boldsymbol W_R$ projection. Transformer^† adds the same number of parameters to the Transformer baseline by increasing the query head count in GQA [23], which maintains KV size.
• Compute-matched Parallax (Parallax^†). Since Parallax doubles the arithmetic intensity compared to FA, Parallax^† halves the head dimension to strictly match the attention layer compute. The reduced parameter count is rebalanced by increasing the FFN dimension to match the total parameter count of the standard Parallax.

• Does the gain come from extra parameters? Transformer^† with matched parameters closes only a small fraction of the gap to Parallax, confirming that the improvement is not simply from additional parameters in query-like projections.
• Does the gain require extra attention compute? Parallax^† with matched attention compute significantly outperforms baseline Transformer and Transformer^†, ruling out the added compute as a necessary condition.

These results provide strong evidence that the gain is driven by the mechanism itself.

The training perplexity curves under different optimizers and schedulers show that the Muon [17]-model curves exhibit a substantial gap between Parallax and the Transformer, consistent with the downstream evaluation. However, the advantage shrinks markedly or even disappears under AdamW [19]. The performance difference indicates a strong optimizer-architecture interaction, which the next section analyzes.

Mechanistic Analysis

Attention Sinks

Though Parallax attention weights $s_{ij}$ sum to 1, unlike softmax attention they can go negative and exceed 1 in magnitude. There are three concrete consequences:

1. Score range. Negative weights let the model actively subtract irrelevant value components rather than just de-emphasize them. Under Muon, extreme weights reach $\pm 40$ in the deepest layers, which is consistent with the COR growth.
2. Attention sink. Softmax Attention notoriously dumps mass on the first token in the attention sink phenomenon [25]. Parallax substantially reduces this in both the base softmax $p_{ij}$ and the combined weights $s_{ij}$ — the correction branch appears to absorb the routing role the sink was serving.
3. Entropy. Parallax's base softmax is consistently more diffuse than a standard Transformer's. The softmax handles broad aggregation; the correction branch handles fine-grained discrimination.

5: Muon [17] Unlocks Parallax's Capacity

Probe Sensitivity

In Parallax, the probe $\rho_i = \boldsymbol W_R \boldsymbol x_i$ is learned instead of solved as in standard LLA. Only the probe's component aligned with the optimal $\boldsymbol\rho_i^\star$ contributes to the correction and the rest is noise. If the probe is misaligned or norm-suppressed, the covariance term collapses and Parallax degenerates to Softmax Attention. We found probe-collapse is a strong function of the choice of optimizer and spectral dynamics.

To measure collapse, we track three quantities per layer:

• the correction-to-output ratio (COR) which is how large the correction is relative to the softmax output
• KV correlation (Corr)
• Covariance-probe alignment (CPA)

Diagnostics

COR grows with depth under both optimizers. Muon reaches above 8 in the deepest layers, AdamW stays below 4. Muon also produces higher $\|\texttt{Corr}\|_F$ and higher CPA, and the probe under Muon is both larger and better-aimed.

Interestingly, random initialization has uniformly high COR. Training suppresses the correction in early layers and amplifies it in deep ones, which barely recovers initialization level under AdamW and exceeds it under Muon.

Is AdamW simply using a different scale convention, or does it genuinely fail to use the correction branch? To test this, we add a learnable sigmoid gate $\rho_i = \sigma(\boldsymbol w_g^\top \boldsymbol x_i) \cdot \boldsymbol W_R \boldsymbol x_i$. Under Muon, the gate opens and converges to the same loss as the ungated run, but under AdamW, it stabilizes at 0.26 and matches Transformer performance. AdamW learns to suppress the correction, independent of scaling.

Muon produces higher stable rank on the probe-related matrices ($\boldsymbol W_R$, $\boldsymbol W_{RK}$), but the architecture and optimizer effects are separable.

Optimizer effect. $\boldsymbol W_Q$, $\boldsymbol W_K$, and $\boldsymbol W_{QK}$ are identical between Parallax and Transformer under every optimizer, and the attention scoring structure is purely optimizer-determined. $\boldsymbol W_R$ is the outlier as it has the largest optimizer sensitivity of any projection, larger than $\boldsymbol W_Q$ or $\boldsymbol W_K$. $\boldsymbol W_{RK}$ inherits this sensitivity and enjoys high rank under Muon but not under AdamW. CPA is higher under Muon since a high-rank RK circuit gives $\boldsymbol\rho$ more directions to align with.

Architecture effect. $\boldsymbol W_V$, $\boldsymbol W_O$, and $\boldsymbol W_{OV}$ are consistently higher rank under Parallax than Transformer, across all optimizers. The correction branch enriches the value pathway regardless of how the probe is trained.

Architecture-Optimizer Codesign

Muon's benefit to Parallax is a clear example of architecture-optimizer codesign. To validate generality, we run modded-NanoGPT experiments with four optimizers (Muon, Aurora, NorMuon, SOAP), each paired with Parallax (+Parallax, solid) and standard attention (dashed). Across all four optimizers, Parallax reaches lower val loss at matched step count. The Parallax benefit is not Muon-specific, though the magnitude varies between different optimizers.

The recipe-design space remains largely unexplored. Muon with warmup-stable-decay scheduling is likely suboptimal, and these preliminary NanoGPT results suggest that more advanced optimizers may further amplify Parallax's advantage.

6: Inflection

Three things excite us about this work.

1. Principled Improvement upon Softmax Attention. LLA is a provably principled improvement upon softmax attention's local constant estimator, correcting the structural boundary-bias pathology. Parallax inherits the local-linear correction principle and empirically realizes its advantages in pretraining, manifesting in strong gains across synthetic and real-world settings.

2. Architecture-optimizer Codesign. The clearest finding in this paper is that Muon counterfactually unlocks Parallax through stable optimization of $W_R$. This is, to our knowledge, the first empirical demonstration of strong architecture-optimizer codesign for an attention mechanism. It definitely won't be the last.

3. Parallax can be up-trained from Softmax. With $W_R = 0$, Parallax is exactly Softmax Attention. Any pretrained Transformer converts to Parallax by adding one weight matrix and fine-tuning or running continued pretraining.

4. Hardware Efficiency. Parallax doubles arithmetic intensity without extra I/O. The Hopper decode kernel already matches or beats FA2/FA3.

Future Work

The most pressing open question is why Muon unlocks $W_R$ and AdamW fails. We have some empirical anatomy, but not a strong theoretical justification.

At scale, the doubled arithmetic intensity opens new recipes for trading head dimension against head count against FFN width, and this path is uncharted. Parallax is also structurally compatible with MLA, sliding window, and block-sparse attention.

Finally, the affine structure here applies to the whole family of sequence mixers. DeltaNet [9] [10], MesaNet [8], and Linear Attention [7] all drop the intercept by construction, and could benefit from the Parallax construction.

Open Source

We release the kernel and a detailed technical report.

• Paper: arxiv.org/abs/2605.29157
• Code: github.com/yifei-zuo/Parallax

@article{zuo2026parallax,
  title   = {Parallax: Parameterized Local Linear Attention for Language Modeling},
  author  = {Zuo, Yifei and Pai, Dhruv and Zeng, Zhichen and Dewulf, Alec
             and Hu, Shuming and Wang, Zhaoran},
  year    = {2026},
  url     = {https://arxiv.org/abs/2605.29157}
}

Appendix

A: Bias-variance Benefit of LLA

B: LLA as an additive correction — the derivation

Starting from the LLA forward

divide numerator and denominator by $\omega_i$ and set $p_{ij} = w_{ij}/\omega_i$. Using $\boldsymbol\mu_i^\top \boldsymbol\rho_i^\star / \omega_i = \mathbb E_{\boldsymbol p_i}[t_{ij}] = \bar t_i$,

Expanding the numerator with $\boldsymbol z_{ij} = \boldsymbol k_j - \boldsymbol q_i$ and the cross-moment factorization $\mathbb E_{\boldsymbol p_i}[\boldsymbol v_j \boldsymbol k_j^\top] = \boldsymbol\Sigma_{KV}^{(i)} + \bar{\boldsymbol v}_i \bar{\boldsymbol k}_i^\top$:

So $\mathbb E_{\boldsymbol p_i}[(1 - t_{ij})\boldsymbol v_j] = (1 - \bar t_i)\boldsymbol o_i^{\textsf{SA}} - \boldsymbol\Sigma_{KV}^{(i)}\boldsymbol\rho_i^\star$. Dividing by $1 - \bar t_i$ and identifying $1/(1 - \bar t_i) = 1 + \eta_i$ gives the additive form

Setting $\eta_i = 0$ and replacing the exact solve with the learned probe $\boldsymbol\rho_i = \boldsymbol W_R \boldsymbol x_i$ recovers the Parallax forward used throughout the post.

C: Hopper decode kernel: three optimizations

The CuTeDSL prototype on H200 hits the matches-or-beats result against FA2/FA3 through three Hopper-specific tricks:

1. WGMMA sharing. Each CTA packs $\mathbf Q_r$ and $\mathbf R_r$ into a single shared-memory tile, with Q in the first row, R in the second, and issues one WGMMA that emits $\mathbf S_1$ and $\mathbf S_2$ jointly in the same accumulator. After $\mathbf P_1$ is produced, $\mathbf P_2 = \mathbf P_1 \odot \mathbf S_2$ is built in registers and stacked with $\mathbf P_1$ in shared memory. The PV WGMMA then emits $\mathbf O_1$ and $\mathbf O_2$ jointly. The cost is one extra row of register accumulators per CTA, no extra HBM traffic, and this is what lets the second branch ride along for free.
2. Persistent split over the KV loop. Decoding only presents $BH$ query rows, often well below H200's 132 SMs. A persistent grid of $(B, H, S)$ CTAs is launched, with $S$ CTAs sharing a $(B, H)$ partition of the $\lceil L/\mathcal B_c \rceil$ tile loop. $S$ is chosen so the launch fits one wave and rounds to a power of two, enabling vectorized cross-split reduction.
3. In-kernel reduction. Each CTA writes its unnormalized partials to a small fp32 HBM workspace and atomically increments a per-$(B, H)$ counter. The CTA that observes the final increment is elected the merger: it reads the $S$ partials, runs log-sum-exp rescaling in fp32, evaluates $(1 + d_2/d_1)\, O_1/d_1 - O_2/d_1$, and writes the output row in the same kernel launch. A compile-time branch on $S = 1$ skips the workspace round-trip entirely on short-context shapes.

D: Arithmetic intensity

The key property of the streaming algorithm is that it increases the arithmetic intensity ($\texttt{AI}$) over FA, defined as the ratio of floating point operations (FLOPs) to high-bandwidth memory (HBM) traffic in bytes. Write $L_q$ and $L_{kv}$ as the query and KV sequence lengths respectively and $d_h$ as the head dimension,

where $n_r = \lceil L_q / \mathcal B_r \rceil$ is the number of query row blocks. In the regime where $n_r L_{kv} \gg L_q$, Parallax roughly doubles the arithmetic intensity by adding more compute while reusing the same KV stream. The shift toward a more compute-bound operator is what makes decoding a target for kernel-level optimization on modern hardware.

E: Boundary amplification is non-negative

The proof that $\eta_i \ge 0$ exposes a Mahalanobis structure inside $\boldsymbol\Sigma_i$. Decomposing via the variance identity gives a rank-one split

where $\bar{\boldsymbol z}_i = \mathbb E_{\boldsymbol p_i}[\boldsymbol z_{ij}]$. Sherman–Morrison inverts $\boldsymbol\Sigma_i$ and, using $\boldsymbol\mu_i = \omega_i \bar{\boldsymbol z}_i$, yields

Substituting back through $\bar t_i = \bar{\boldsymbol z}_i^\top \boldsymbol\rho_i^\star$,

So $\eta_i$ equals $\omega_i$ times the squared Mahalanobis distance from $\boldsymbol q_i$ to the weighted key mean $\bar{\boldsymbol k}_i$ under the metric $\boldsymbol A_i^{-1}$. Queries near the weighted key center give $\eta_i \approx 0$ (pure covariance correction); queries near the boundary give $\eta_i \gg 1$ (amplified correction). This is the geometric content sitting under "boundary amplification."

F: Parallax backward (closed form)

The backward admits the same row-tile / column-tile streaming structure as the FA backward. With upstream gradient $\mathrm{d}\boldsymbol O_i$, define the row scalars and per-token quantities

and the two channel coefficients

The gradients are then

The forward caches $(\boldsymbol o_i, \bar{\boldsymbol v}_i, \bar t_i, \omega_i, m_i)$ per row — only $d + 1$ extra values beyond the FA cache. Because the row-sum and column-sum directions are opposite, the backward splits into a row-tile pass (loads $\mathbf Q_r, \mathbf R_r, \mathrm{d}\mathbf O_r$, streams over column blocks of $\mathbf K, \mathbf V$, accumulates $\mathrm{d}\mathbf Q_r, \mathrm{d}\mathbf R_r$) and a column-tile pass (loads $\mathbf K_c, \mathbf V_c$, streams over row blocks of $\mathbf Q, \mathbf R, \mathrm{d}\mathbf O$ in reverse order, accumulates $\mathrm{d}\mathbf K_c, \mathrm{d}\mathbf V_c$).

G: Weight-decay annealing during WSD decay

While Muon delivers a substantial advantage over AdamW, we do not claim Muon with WSD is the optimal combination for Parallax. The advantage of Parallax shrinks during the final linear decay phase of the WSD schedule, and weight norms shrink throughout the decay phase, which may partially explain the shrinkage. We ablate this with weight-decay annealing (WDA), which replaces the constant decay $\lambda$ with a schedule $\lambda(t) = \lambda (1 - t)^\gamma$. WDA effectively mitigates the weight-norm shrinkage and reaches lower final training loss than the standard recipe, with the improvement monotone in $\gamma$ over the range tested.

The relative effective step in parameter space is

Under standard WSD, both numerator and denominator shrink together over the decay window, partially preserving $\Delta_t$. The Parallax advantage shrinks during decay, and $\|\boldsymbol W_R\|_F$ also shrinks throughout. This correlation suggests the relative-step collapse may be the mechanism.

WDA replaces the constant weight decay $\lambda$ with $\lambda(t) = \lambda(1 - t)^\gamma$ over $t \in [0, 1]$, where $\gamma = 0$ recovers standard WSD and $\gamma > 0$ suppresses decay more aggressively toward the end of training. Tested at $\gamma \in \{0.5, 1, 2\}$ at 0.6B scale, WDA holds $\|\boldsymbol W_t\|_F$ - higher $\gamma$, larger final norms - and yields lower final training loss than the standard recipe, monotone in $\gamma$, with $\gamma = 2$ giving the largest gain. But the gap closes in the second half of the decay stage at similar step counts across $\gamma$: as $\eta_t \to 0$ with the denominator now held by WDA, $\Delta_t$ collapses faster than under standard WSD. WDA partially mitigates the issue without fully resolving it, which is preliminary evidence that Muon + WSD is not the optimal recipe for Parallax in its current form, and an open recipe-design question.