When and Why Grouping Attention Heads Accelerates Muon Optimization

• URL: http://arxiv.org/abs/2605.08933v1
• Date: Sat, 09 May 2026 13:06:31 GMT
• Title: When and Why Grouping Attention Heads Accelerates Muon Optimization
• Abstract summary: We study whether Muon should be applied to the full attention projection, to individual heads, or to intermediate head groups.<n>Our analysis reveals a trade-off between the textbf group-wise whitening gain from group-wise updates and an additional update-norm cost caused by replacing full-matrix whitening with group-wise whitening.
• Score: 51.89517622356298
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Muon orthogonalizes matrix updates, but multi-head attention naturally operates at the level of heads. This granularity mismatch raises the question of whether Muon should be applied to the full attention projection, to individual heads, or to intermediate head groups. We study this question through a one-step descent comparison between full-matrix Muon and group-wise Muon. Our analysis reveals a trade-off between the \textbf{group-wise whitening gain} from group-wise updates and the \textbf{grouping-induced norm cost}, an additional update-norm cost caused by replacing full-matrix whitening with group-wise whitening. Motivated by this trade-off, we propose \textbf{Group Muon}, which treats head group size and grouping rule as optimizer hyperparameters. On GPT-2 Small trained on FineWeb, appropriate grouping improves validation loss over both full-QKV Muon and fully head-wise MuonSplit.

Related papers

• Reassessing Muon for Matrix Factorization [9.691906492780353]
  Muon has emerged as a strong contender in large-scale deep learning.<n>We study Muon on a simple, well-understood, and spectrally structured problem: low-rank matrix factorization.<n>We find that Muon does not consistently outperform AdamW in this setting.
  arXiv  Detail & Related papers  (2026-07-14T20:21:38Z)
• Zeta: Dual Whitening for Matrix Optimization via Coordinate-Adaptive Preconditioning [56.24532075189964]
  We propose Zeta, a dual whitening that applies coordinate whitening and spectral whitening in a strictly ordered pipeline.<n>We prove that Zeta matches or surpasses strong baselines across language modeling (0.6B to 8B parameters), mixture-of-experts architectures, and vision tasks.
  arXiv  Detail & Related papers  (2026-06-12T07:10:17Z)
• Muon$^2$: Boosting Muon via Adaptive Second-Moment Preconditioning [18.570226339282296]
  Muon has emerged as a promising computation for large-scale foundation model pre-training.<n>We propose Muon$2$, a memory-efficient factorized variant that preserves most of the gains of Muon$2$ with negligible memory overhead.
  arXiv  Detail & Related papers  (2026-04-11T00:27:40Z)
• Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training [0.0]
  We introduce MUD (MomentUm Decorrelation), a complementary whitening approach that replaces Muon's polar update with a triangular (Cholesky-like) whitening surrogate.<n>In terms of time-to-perplexity, MUD yields consistent 10-50% wall-clock improvements over tuned AdamW and Muon.<n>We also demonstrate training a ESM-2 150M protein language model, where MUD matches Muon-level validation perplexity in significantly less wall-clock time.
  arXiv  Detail & Related papers  (2026-03-18T17:37:31Z)
• Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning [72.62839712454196]
  Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improveimation.<n>We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned.<n>We propose textbfMousse (textbfMuon textbfOptimization textbfUtilizing textbfShampoo's textbfStruct
  arXiv  Detail & Related papers  (2026-03-10T14:03:49Z)
• Muon in Associative Memory Learning: Training Dynamics and Scaling Laws [23.350512542598803]
  We study Muon in a linear associative memory model with softmax retrieval and a hierarchical frequency spectrum over query-answer pairs.<n>We show that Muon mitigates this imbalance, leading to faster and more uniform progress.
  arXiv  Detail & Related papers  (2026-02-05T14:49:40Z)
• The Ky Fan Norms and Beyond: Dual Norms and Combinations for Matrix Optimization [37.169656352055604]
  We introduce a family of Muon-like algorithms we name Fanions, which are closely related to Dion.<n>F-Muon and S-Muon consistently match Muon's performance, while outperforming vanilla Muon on a synthetic linear least squares problem.
  arXiv  Detail & Related papers  (2025-12-10T14:25:45Z)
• Muon is Scalable for LLM Training [50.68746986439438]
  We introduce Moonlight, a Mixture-of-Expert (MoE) model trained with 5.7T tokens using Muon.<n>Our model improves the current frontier, achieving better performance with much fewer training FLOPs compared to prior models.<n>We open-source our distributed Muon implementation that is memory optimal and communication efficient.
  arXiv  Detail & Related papers  (2025-02-24T09:12:29Z)
• GroupedMixer: An Entropy Model with Group-wise Token-Mixers for Learned Image Compression [64.47244912937204]
  We propose a novel transformer-based entropy model called GroupedMixer. GroupedMixer enjoys both faster coding speed and better compression performance than previous transformer-based methods. Experimental results demonstrate that the proposed GroupedMixer yields the state-of-the-art rate-distortion performance with fast compression speed.
  arXiv  Detail & Related papers  (2024-05-02T10:48:22Z)
• Advancing Vision Transformers with Group-Mix Attention [59.585623293856735]
  Group-Mix Attention (GMA) is an advanced replacement for traditional self-attention. GMA simultaneously captures token-to-token, token-to-group, and group-to-group correlations with various group sizes. GroupMixFormer achieves state-of-the-art performance in image classification, object detection, and semantic segmentation.
  arXiv  Detail & Related papers  (2023-11-26T01:25:03Z)
• Just Train Twice: Improving Group Robustness without Training Group Information [101.84574184298006]
  Standard training via empirical risk minimization can produce models that achieve high accuracy on average but low accuracy on certain groups. Prior approaches that achieve high worst-group accuracy, like group distributionally robust optimization (group DRO) require expensive group annotations for each training point. We propose a simple two-stage approach, JTT, that first trains a standard ERM model for several epochs, and then trains a second model that upweights the training examples that the first model misclassified.
  arXiv  Detail & Related papers  (2021-07-19T17:52:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.