Source-linked AI summary
NS-ATTENTION: Newton-Schulz Transformations of Attention Outputs in Vision Transformers
Xiaohe Jiang, Guoqiang Zhang, Tianjin Huang, Ronghui Mu
TL;DR
The paper asks whether redistributing spectral energy in attention-head outputs can improve image classification. It introduces NS-Attn., a parameter-free finite Newton–Schulz transformation applied to each head’s output. Across ViT and Swin on CIFAR-10 and CIFAR-100, accuracy improves in all 12 matched-seed comparisons, while the transformation adds inference latency.
Problem
The paper investigates whether reshaping the spectral energy of attention-head outputs can improve image classification.
Source: §1 · p. 1Method
NS-Attn. applies a finite Newton–Schulz polynomial transformation to each attention-head output, restoring its Frobenius norm before standard head merging and output projection.
Source: §1 · p. 1, §2.2 · p. 2Results
Mean final-epoch accuracy improves by 0.25–0.83 percentage points across four ViT and Swin settings, with gains in all 12 matched-seed comparisons; trained models also show less concentrated spectra and higher effective rank.
Source: §1 · p. 1Takeaways & Limitations
The spectral changes accompany classification gains but do not establish that spectral flattening causes them.
Source: §4 · p. 4Additional forward latency makes the practical value dependent on application accuracy–latency requirements, and larger-scale datasets and equal-compute evaluations would help establish the benefits’ scope.
Source: §4 · p. 4, §3.5 · p. 4
Abstract
from arXiv · showhide
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.
1. INTRODUCTION
The paper investigates whether reshaping the spectral energy of attention-head outputs can improve image classification. It introduces NS-Attn. and evaluates it across ViT and Swin on CIFAR-10 and CIFAR-100.
1. INTRODUCTION
- The study targets a flatter normalized eigenvalue spectrum and higher effective rank by reducing the dominance of a few directions in each head’s output.Squared singular values describe how energy is distributed across spectral directions.Source: §1 · p. 1
- Finite NS-style matrix polynomials modify relative singular-value magnitudes through matrix multiplications, motivating their use to redistribute attention-representation energy.The paper contrasts this with uniform rescaling and exact orthogonalization, and treats iteration count as a design choice.Source: §1 · p. 1
- Prior work transforms optimizer updates with Muon, uses matrix iterations for covariance pooling and feature whitening, or applies NS to state-space input injections.Within attention, prior approaches mix heads [6], condition projection weights, or construct an orthogonal operator from queries and keys [8]; NS-Attn. instead transforms aggregated representations.Source: §1 · p. 1
- NS-Attn. applies a finite NS polynomial to each feature-by-token attention-head output before merging, restores its Frobenius norm, and adds no trainable parameters.The paper describes this as the first application of NS-style iterations directly to per-head Transformer attention outputs.Source: §1 · p. 1
- Mean final-epoch accuracy improves by 0.25–0.83 percentage points across four settings, with gains in all 12 matched-seed comparisons.The analysis also finds less concentrated output spectra and higher effective rank than Standard Attn.Source: §1 · p. 1
2. NEWTON–SCHULZ TRANSFORMATIONS OF ATTENTION OUTPUTS
NS-Attn. reshapes each attention head’s output as a feature-by-token matrix and transforms each matrix independently before head merging and output projection. It normalizes the matrix, applies a finite polynomial transformation, and restores the original norm.
2.1. Attention Output and Feature View
- Each sample-head attention output is treated as a feature-by-token matrix, with rows representing feature responses across tokens; matrices are transformed independently without mixing heads.For Swin, samples count windows across the batch and tokens are those within each window.Source: §2.1 · p. 1, Figure 1 · p. 2
2.2. Newton–Schulz Transformation
- Each matrix is first normalized by its Frobenius norm, with ε = 10−6 guarding against division by very small norms.The method forms a Gram matrix from the normalized matrix and applies one Muon-inspired polynomial step using a = 3.4445, b = −4.7750, and c = 2.0315.Source: §2.2 · p. 2
- After the polynomial step, the method normalizes the result and restores the original norm; with stabilizing clamps inactive, it preserves matrix norm while changing the singular-value distribution.The paper calls this a finite Muon-inspired transformation rather than exact orthogonalization.Source: §2.2 · p. 2
- The transformed output retains the input shape and passes through existing head merging and output projection, with no trainable parameters added.The main configuration uses K = 1; the ablation compares it with K = 2, and normalization and norm restoration each occur once outside the iteration loop.Source: §2.2 · p. 2
3. EXPERIMENTS AND ANALYSIS
Across ViT and Swin on CIFAR-10 and CIFAR-100, NS-Attn. improves classification accuracy in every matched-seed comparison. The experiments also examine iteration count, spectral effects, and the added inference latency.
3.1. Experimental Setup
- ViT and Swin were evaluated on CIFAR-10 and CIFAR-100, with final-epoch Top-1 accuracy averaged over three matched seeds after 100 training epochs.NS-Attn. was applied in every attention block; the models used AdamW and the stated augmentation and learning-rate schedule.Source: §3.1 · p. 3
3.2. Classification Performance
- NS-Attn. improves mean final-epoch accuracy in all four backbone–dataset settings and all 12 matched-seed comparisons.Both backbones show larger absolute gains on CIFAR-100 than on CIFAR-10.Source: §3.2 · p. 3, Table 1 · p. 3
3.3. Ablation Study
- On ViT, one NS iteration achieves higher mean accuracy than two on both datasets, supporting one iteration as the main configuration.The tested epsilon thresholds from 10^-8 to 10^-2 were insensitive in audited final-checkpoint evaluations; neither clamp activated in those evaluations or the monitored early-training run.Source: §3.3 · p. 3, Table 2 · p. 3
3.4. Spectral Analysis
- NS-Attn. yields less concentrated spectra and higher effective rank for both ViT and Swin on both datasets.ViT mean p1 / mean reff: CIFAR-10, 0.72 to 0.62 / 2.53 to 2.91; CIFAR-100, 0.72 to 0.61 / 2.49 to 3.00. Swin: CIFAR-10, 0.78 to 0.64 / 2.21 to 3.04; CIFAR-100, 0.79 to 0.61 / 2.19 to 3.13.Source: §3.4 · p. 4, Figure 2 · p. 4
3.5. Computational Cost
- Measured GPU forward latency rises by 41.0% for ViT and 103.7% for Swin, creating an accuracy–latency trade-off.The reported table covers batch size 128; the relative overhead is larger for Swin.Source: §3.5 · p. 4, Table 3 · p. 4
4. DISCUSSION AND LIMITATIONS
NS-Attn. redistributes spectral energy while preserving total energy when numerical clamps are inactive, but the measured spectral changes do not establish why classification gains occur. Benefits appear across both evaluated backbones, with additional latency and evaluation scope limiting practical conclusions.
4. DISCUSSION AND LIMITATIONS
- When numerical clamps are inactive, restoring the input Frobenius norm preserves total spectral energy while the polynomial redistributes it across directions.The observed smaller leading-eigenvalue fraction and higher effective rank accompany classification gains but do not establish that spectral flattening causes them.Source: §4 · p. 4
- Classification improvements and aggregate eigenvalue trends extend to both ViT and Swin Transformer, but the transformation adds forward latency.Its practical value depends on application accuracy–latency requirements.Source: §4 · p. 4
- Larger-scale dataset evaluation and comparisons under equal compute budgets would help establish the scope of the observed benefits.Source: §4 · p. 4
5. CONCLUSION
The conclusion presents NS-Attn. as a parameter-free finite Newton–Schulz transformation of each head’s attention output. One iteration improves mean final-epoch accuracy across evaluated ViT and Swin settings, alongside less concentrated spectra and higher effective rank.
5. CONCLUSION
- NS-Attn. applies a finite Newton–Schulz-style transformation to each head’s attention output without adding trainable parameters.Source: §5 · p. 4
- One iteration improves mean final-epoch accuracy across the evaluated ViT and Swin settings.Source: §5 · p. 4
- Eigenvalue analysis of trained ViT and Swin Transformer models shows less concentrated output spectra and higher effective rank.Source: §5 · p. 4