Exclusive Self-Attention Beyond AdamW: Stability and Generalization under Muon Optimization

Technical Summary: Exclusive Self-Attention Beyond AdamW

Problem Statement
Standard self-attention mechanisms in Transformers allow the output for a token at position $i$ to include a weighted mixture of value vectors, including the token's own value vector ($v_i$). While expressive, this design can create a "representational shortcut" where the attention output remains strongly aligned with $v_i$, even though residual pathways already preserve local token information. This phenomenon, termed "attention similarity bias," was addressed by Zhai [2] through Exclusive Self-Attention (XSA), which explicitly removes the component of the attention output parallel to the token's own value vector.

However, the original XSA study evaluated the method exclusively within a NanoGPT training stack using the AdamW optimizer. The authors of the original work explicitly identified the compatibility of XSA with alternative optimizers, such as Muon, as an open question. This is a critical gap because modern small-scale language model research (e.g., the nanochat stack) increasingly utilizes hybrid optimization strategies (Muon for matrix parameters, AdamW for embeddings) and architectural features like value embeddings and rotary position encoding. It remains unclear whether the geometric constraints of XSA remain stable or beneficial when transferred to these modern, Muon-based training regimes.

Methodology
The study conducts a controlled reproduction and transfer experiment using the nanochat training stack. The methodology involves:

1. Implementation: A parameterized form of XSA is integrated into nanochat. The operation is defined as $z_i = y_i - \alpha \cdot \text{proj}_{v_i}(y_i)$, where $y_i$ is the standard attention output, $v_i$ is the token's value vector, and $\alpha \in [0, 1]$ controls the projection strength. At $\alpha=1$, the self-aligned component is fully removed.
2. Correctness Validation: The implementation is validated numerically to ensure that at $\alpha=0$, the XSA path is equivalent to standard attention. Short training runs verify that baseline and $\alpha=0$ paths exhibit matching early loss behavior.
3. Experimental Setup: Two model depths are evaluated under matched training conditions (same data ordering, optimizer family, learning rate schedule, and context length):
   • d12: 12 layers, hidden dimension 768, 12 heads.
   • d24: 24 layers, hidden dimension 1024, 16 heads.
   • Optimizer: Matrix parameters are optimized with Muon; remaining parameters use AdamW.
   • Data: ClimbMix-400B.
4. Metrics: The study tracks training loss, validation bits-per-byte (bpb), downstream aggregate scores (Base CORE and ChatCORE), and attention diagnostics (cosine alignment between output and self-value, diagonal behavior).

Key Contributions

• Muon-Stack Reproduction: The first integration of XSA into a Muon-based training stack, directly addressing the open question regarding optimizer compatibility left by the original study.
• Parameterized Implementation: A flexible implementation allowing for standard attention, full XSA, and partial projection via a strength parameter $\alpha$, alongside layer-selection controls.
• Correctness and Instrumentation: Validation of the $\alpha=0$ limit and correction of logging artifacts (step-index and smoothing conventions) that previously distorted early training comparisons.
• Two-Depth Evaluation: A comparative analysis at d12 and d24 depths using a comprehensive set of metrics.
• Mechanistic and Systems Analysis: Diagnostics on attention self-alignment and an evaluation of runtime overhead.

Results
The empirical results indicate that XSA is stable under the Muon-based stack but yields mixed downstream benefits depending on model depth and the specific metric:

• Optimization Stability: Training loss trajectories for XSA and baseline models remain closely aligned at both depths. There is no evidence of divergence, NaNs, or instability, confirming that XSA does not disrupt Muon's update geometry.
• Validation Performance (Bits-per-Byte): XSA consistently improves held-out validation bpb:
  • d12: Reduced from 0.8484 to 0.8457 (0.32% relative improvement).
  • d24: Reduced from 0.7193 to 0.7171 (0.31% relative improvement).
• Downstream Metrics (CORE):
  • d24: XSA improves both Base CORE (0.2593 to 0.2606, +0.50%) and ChatCORE (0.3638 to 0.3682, +1.21%).
  • d12: While bpb improves, Base CORE decreases from 0.1602 to 0.1508 (-5.87%).
• Mechanistic Diagnostics: Attention diagnostics show reduced self-alignment and reduced diagonal behavior under XSA, consistent with the intended mechanism of promoting contextual mixing.
• Runtime Overhead: The overhead is characterized as modest. While an initial estimate suggested ~5%, repeated throughput measurements placed the difference near the system-level variability of short runs, suggesting no prohibitive penalty.

Significance and Claims
The paper concludes that Exclusive Self-Attention is a compatible architectural inductive bias for Muon-based training stacks, but its utility is not universal. The primary significance lies in demonstrating that the representational constraint of XSA transfers cleanly to a modern optimization regime without destabilizing training.

However, the author emphasizes that the downstream gains are scale- and metric-dependent. The improvement in validation bpb at both depths suggests a representational benefit, but the degradation of Base CORE at d12 indicates that the benefits of removing self-alignment may depend on model depth, the capacity of residual/value pathways to preserve local information, or the specific evaluation task. The study does not claim XSA is a universally superior replacement for standard attention, but rather a stable, low-disruption constraint that offers benefits in specific contexts (notably at d24 in this study).

The work serves as a robustness check rather than a byte-for-byte reproduction, highlighting the complex interactions between architectural modifications, optimizer geometry, and evaluation metrics in modern language model research.