Transforming attention outputs improves accuracy in vision transformers
NS-ATTENTION: Newton-Schulz Transformations of Attention Outputs in Vision Transformers
Machine Learning
Summary
Vision transformers use attention to understand images, but this attention can sometimes focus too narrowly on certain parts. The authors tried a mathematical trick called Newton-Schulz iteration to spread out the attention more evenly and make it richer. This adjustment helped models get better at classifying images in tests. It works without adding extra parameters but does slow down the processing a bit.
What this means in practice
- •For computer vision engineers: Improve image classification accuracy by applying Newton-Schulz transformations to attention outputs in vision transformers during model training.
- •For machine learning infrastructure teams: Optimize deployment strategies by balancing increased inference latency against accuracy improvements from Newton-Schulz transformed attention in vision models.
Authors
Xiaohe Jiang, Guoqiang Zhang, Tianjin Huang, Ronghui Mu
Abstract
Newton-Schulz (NS) iteration has recently been used in the Muon optimizer to transform update matrices during the training of large language models. Motivated by its spectral effect, we investigate applying NS directly to Transformer attention representations. We introduce Newton-Schulz Attention (NS-Attn.), a parameter-free transformation applied to the output of each attention head. Each head output is arranged as a feature-by-token matrix and normalized by its Frobenius norm. We then apply a finite NS polynomial step and restore the original norm. The objective is to reduce spectral concentration and increase effective rank before standard head merging and output projection. Across ViT and Swin on CIFAR-10 and CIFAR-100, NS-Attn. improves final-epoch accuracy in all 12 matched-seed comparisons, with mean gains of 0.25--0.83 percentage points. ViT ablations show higher mean accuracy with one iteration than with two. Spectral analysis further shows reduced leading-eigenvalue concentration and increased effective rank. These gains incur additional inference latency.