Multi-Head Normalization for Wide Vision Transformers

Scaling the width of Vision Transformers (ViT) often triggerssevere instability issues. This work identifies that these instabilities stemfrom the “high-activation domination” phenomenon frequently encoun-tered during feature normalization in high-dimensional spaces. Specifi-cally, the norm of a high-dimensional vector can be easily dominated by afew outlier elements with disproportionately large values. Upon normal-ization, the remaining elements are suppressed toward zero, leading tosignificant gradient propagation issues. To address this, this work intro-duces Multi-Head RMSNorm (MH-RMSNorm), which partitions a high-dimensional feature vector into multiple chunks and applies RMSNormto each segment independently. This simple, drop-in replacement sig-nificantly enhances the training stability and predictive performance ofhigh-dimensional ViTs. For example, Our 1280-dimensional ViT achieves85.5% top-1 accuracy on ImageNet-1K, whereas baseline models usingstandard LayerNorm or RMSNorm suffer from complete training col-lapse; when applied to DiT on ImageNet-256, our method achieves a1.95 FID, significantly outperforming standard ViT counterparts.