Muon优化器在视觉Transformer中表现更优,尤其配合强数据增强时效果显著。
Muon in Vision Transformers: Optimizer-Recipe Interactions and Gradient Spectra

- Muon通过更广谱的梯度分布提升训练效果,尤其在深层MLP块中防止模式坍缩。
- 在Pl@ntNet数据集上,Muon相比AdamW提升显著,长尾分类性能改善明显。
- 适合追求高精度视觉模型的开发者,尤其关注数据增强与优化器协同设计者。
Muon是一种新型矩阵感知优化器,在Transformer训练中表现优异,但其在视觉Transformer(ViTs)中的行为尚不明确。本研究在ImageNet-100和Pl@ntNet-300K数据集上,对比Muon与AdamW在标准视觉训练配方(含mixup、cutmix、标签平滑、随机增强与擦除)下的表现。结果显示,Muon持续优于AdamW,尤其在长尾分布的Pl@ntNet宏平均准确率上提升显著。该优势具有配方依赖性:当使用更强的数据增强时,Muon收益远高于AdamW。通过分析视觉变压器中矩阵梯度的奇异值结构发现,若移除强数据增强,Muon训练后期会出现梯度矩阵的谱集中与模式坍缩,主要集中在深层MLP-down模块。而在固定完整增强配方下,最显著差异出现在QKV梯度中:AdamW的能量集中于更窄的奇异基,而Muon则扩散至更多奇异模式。因此,Muon在ViTs中的优势本质是优化器与训练配方的协同作用。此外,我们在图像分割和掩码自编码器任务中也验证了Muon在所有设置下的优越性。
原文摘要 · Abstract (English)
Muon is a recently developed matrix-aware optimizer that has shown strong results in transformer training, but its behavior in vision transformers (ViTs) is not yet well understood. We study Muon for ViT training, largely on ImageNet-100 and Pl@ntNet-300K, comparing against AdamW under standard vision recipes involving mixup, cutmix, smoothing, and random augmentation and erasing. Muon consistently outperforms AdamW, with especially large gains on long-tailed Pl@ntNet macro top-1. These gains are also recipe-dependent, where Muon benefits much more than AdamW from advanced and significant data augmentation techniques. To understand this interaction, we analyze the singular-value structure of matrix gradients throughout the ViT. Within Muon training runs, removing heavy data augmentation induces a late-training spectral concentration and mode collapse in gradient matrices, primarily in deep MLP-down blocks. Under a fixed "full" augmentation recipe, the clearest Muon-AdamW contrast appears instead in QKV gradients, where AdamW gradient energy remains concentrated in a much narrower basis while Muon spreads energy across substantially more singular modes. Muon in ViTs is therefore best understood as an optimizer-recipe interaction. Under a fixed recipe, Muon differs from AdamW most clearly in attention projections, where its gradients consist of a broader spectral basis. Within Muon, a full training recipe is important for preventing late spectral concentration and mode collapse in deep feedforward blocks. We further demonstrate efficacy in training ViTs on image segmentation and masked autoencoder models, where Muon outperforms AdamW in all settings considered.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。