vit-attention-bench

Spectral theory says ViT attention is low-rank. In practice, pure low-rank fails — and semantic landmark selection wins. This repo lets you measure it yourself.

vitattn is a small library of attention-approximation methods for Vision Transformers — SVD truncation, Nyström, Nyströmformer, Performer (FAVOR+), Linformer, semantic-landmark Nyström — behind one unified AttentionApprox API, plus a reproducible benchmark harness that measures what each approximation actually costs and breaks when you drop it into a pretrained ViT with no retraining.

It exists to interrogate a specific, seductive claim from the efficient-attention literature: learned attention matrices are approximately low-rank, so you can approximate them cheaply. The spectra say that's true. The predictions say it's a trap — until you stop reading the spectrum in aggregate and start reading it per layer.

---

TL;DR — the finding

1. In aggregate, spectral theory looks right. Across all 144 heads of ViT-B/16, attention spectra decay steeply: the top singular value dominates, and the per-layer median effective rank is only ~7–54 out of 197 tokens. On paper, everything is compressible.

2. In practice, pure low-rank breaks the model. Truncating each attention matrix to its Frobenius-optimal rank-64 approximation (the best any rank-64 scheme could ever do, by Eckart–Young–Mirsky) still flips 23% of top-1 predictions vs the exact model. Every full-model Nyström variant we tried — 3-softmax and Williams–Seeger formulations, with/without pseudo-inverse regularisation, uniform and semantic landmarks, m from 32 to 160 — flips between 27% and 99% of predictions. Frobenius error near zero, decisions destroyed.

3. The mechanism: per-head error tracks effective rank almost perfectly. The heads that break the model are the near-full-rank ones in the early layers; the near-one-hot deep heads are trivial. Approximation error vs effective rank has Spearman ρ = +0.94 to +0.99.

4. The redemption: read the spectrum per layer and it tells you exactly where to cut. Approximate only the layers the error profile flags as safe (the shallow layer 0 plus the deep layers 8–11) and prediction flips collapse from ~66% to ~9% — and there, where approximation is actually viable, semantic landmark selection finally beats uniform (9.0% vs 10.4% flips). Approximate the hard middle layers {1–7} instead and you're back to 66%.

What you patch (ViT-B/16, N=197, 512 CIFAR-100 images)	Method	Prediction flips ↓
All 12 layers	SVD rank-64 (Frobenius-optimal ceiling)	23%
All 12 layers	Nyström, m=64	92%
All 12 layers	Nyström, m=64, regularised pinv	66%
Only low-rank layers {0, 8–11}	Nyström, m=64	10.4%
Only low-rank layers {0, 8–11}	Semantic-landmark Nyström, m=64	9.0%
Only high-rank layers {1–7} (inverse control)	Nyström, m=64	66%

Punchline. Spectral theory doesn't lie uniformly — it lies in aggregate. Read per layer, it tells you exactly where approximation is safe.

5. And it survives real accuracy, not just decision-agreement. On real Tiny-ImageNet top-1 (the 182 of 200 classes that exist in the ImageNet-1k head, logits restricted to them, images upscaled 64→224, n≈500), spectrum-guided semantic patching is statistically indistinguishable from the exact baseline — 74.4% vs 74.0% — while approximating the full model collapses to 39.5% (Nyström) / 25.8% (semantic), a 35–48 point drop. Same lesson, now in accuracy points: approximate where the spectrum says it's safe and the model is untouched; approximate everywhere and it falls apart. Details in Results → v0.2.

(Every number above is re-measured by the harness in this repo, not quoted from anywhere — see Reproducibility.)

---

Quickstart

git clone https://github.com/matisdsp/vit-attention-bench
cd vit-attention-bench
pip install -e ".[dev]"        # PyPI package `vitattn` coming at first release

Swap the softmax attention of any pretrained timm ViT for an approximation — no retraining, learned weights untouched:

import torch, timm
from vitattn import patch_vit, NystromAttention, is_patched

# 1. any pretrained timm ViT
model = timm.create_model("vit_base_patch16_224.augreg2_in21k_ft_in1k", pretrained=True).eval()

# 2. reroute the softmax core through an AttentionApprox (fresh instance per block)
patch_vit(model, lambda: NystromAttention(num_landmarks=64))
print("patched:", is_patched(model))          # -> True

# 3. forward pass, identical model API
with torch.no_grad():
    logits = model(torch.randn(1, 3, 224, 224))
print(logits.shape)                            # -> torch.Size([1, 1000])

patch_vit takes a zero-argument factory (not an instance) so every patched block gets its own independent state; unpatch_vit(model) restores the original attention exactly. Pass layer_indices=[...] to patch only selected blocks — that is the whole basis of the spectrum-guided partial-patching result above.

---

Methods

All seven implementations subclass a single narrow contract, AttentionApprox: forward(q, k, v, mask=None, *, return_attn=False) over per-head, already-projected (B, H, N, d) tensors. The return_attn flag gates whether the O(N²) matrix Â is ever materialised, so the latency path pays exactly what native attention pays. Each method declares materializes_attn — whether it can hand back the N×N matrix at all — and the harness skips the approximation-error metric for those that structurally cannot.

Method	Class	Family	Target complexity	materializes_attn	Origin
Exact softmax (baseline)	ExactAttention	fused SDPA	O(N²d)	✅	reference (F.scaled_dot_product_attention)
SVD truncation	SVDTruncatedAttention	low-rank ceiling	O(N²d + N³)	✅	Eckart–Young–Mirsky (1936/1960)
Nyström (vanilla)	NystromAttention	landmark low-rank	O(Nmd)	✅	Williams & Seeger, NeurIPS 2001
Nyströmformer	NystromformerAttention	landmark, sub-quadratic	O(Nmd)	✅	Xiong et al., AAAI 2021 (2102.03902)
Performer (FAVOR+)	PerformerAttention	positive random features	O(NMd)	✅ (diagnostic)	Choromanski et al., ICLR 2021 (2009.14794)
Linformer	LinformerAttention	token-axis projection	O(Nkd)	❌	Wang et al., 2020 (2006.04768)
Semantic Nyström	SemanticNystromAttention	landmark + semantic selection	O(Nmd) + selection	✅	author's MSc thesis idea, re-measured here

Patching with ExactAttention reproduces the original model's logits to floating-point tolerance — the "truth test" that validates the whole patching layer. SVDTruncatedAttention and the diagnostic Nyström formulations form Â on purpose (to measure error), so they are reference instruments, not accelerators; the genuinely sub-quadratic fast paths live in Nyströmformer, Performer, Linformer and semantic-Nyström's fast selector.

---

The benchmark

vitattn bench --config <yaml> patches a pretrained ViT method-by-method and writes one CSV row per method. The measured columns that matter:

• pred_change_rate — the headline metric: fraction of images whose argmax prediction disagrees with the exact model. Train-free, and the honest choice here (see Scope & limitations): an ImageNet-1k head scores ~0.2% top-1 against raw CIFAR-100 labels, so top-1 is uninterpretable — but decision agreement with the exact model is exactly what an approximation is supposed to preserve.
• approx_rel_frobenius — median relative Frobenius error ‖Â − A‖/‖A‖ of the approximated attention matrix vs exact, when the method materialises it.
• latency_p50_ms / latency_p95_ms — wall-clock per forward pass (warmup + repeats).
• top1 / top5 — reported for completeness; see the caveat above.
• flops_total / flops_attention — traced with fvcore.

Two honesty caveats, stated up front:

1. FLOPs (flops_note = sdpa-uncounted). The fused SDPA kernel is invisible to fvcore, so the exact baseline's attention FLOPs are not traced — FLOP counts understate the baseline and should be read as method-to-method deltas, not absolute costs. Latency is the honest cost signal.
2. Latency, and why "efficient" attention loses here. At N=197 tokens a fused SDPA kernel beats every approximation on wall-clock — the sub-quadratic methods are simply not supposed to win at this sequence length, and this benchmark says so plainly rather than hiding it. The point of the harness is accuracy fidelity, not a latency victory lap: it measures what low-rank approximation breaks, and at what rank the breakage starts.

---

Results

(ViT-B/16 augreg2_in21k_ft_in1k, CIFAR-100, 512 images, CPU — the quick protocol.)

Act 1 — the spectra are compressible, but wildly uneven

Median effective rank is low, but individual heads span ~1 to 181 out of 197: near-one-hot in the deep layers, near-full-rank in layers 1–4. Aggregate low-rank hides enormous per-head variance — the first hint that a uniform low-rank budget is the wrong tool.

Act 2 — full-model approximation fails, Frobenius error notwithstanding

Method	Prediction flips ↓	Rel. Frobenius	p50 latency	vs SDPA
Exact (SDPA baseline)	0%	0.000	288 ms	1.0×
SVD truncation, rank 16	60%	0.175	4685 ms	16×
SVD truncation, rank 64	23%	0.063	4546 ms	16×
Nyström, m=32	99%	—	467 ms	1.6×
Nyström, m=64	92%	—	770 ms	2.7×
Nyström, m=64, regularised pinv	66%	0.92	876 ms	3.0×
Nyström, m=160, regularised pinv	27%	0.21	3882 ms	12×

Even the Frobenius-optimal rank-64 truncation flips ~1 prediction in 4 at 0.063 relative error. Nyström only reaches the ballpark of tolerable when m=160 — that is, when it stops being a compression at all (160 of 197 tokens).

Act 3 — spectrum-guided partial patching, and where semantic selection finally wins

Patch only the low-rank / low-error layers {0, 8, 9, 10, 11} and the picture flips:

Layers patched	Method	Prediction flips ↓	Rel. Frobenius
{0, 8–11} (low-rank)	Nyström, m=64	10.4%	0.183
{0, 8–11} (low-rank)	Semantic Nyström, m=64	9.0%	0.104
{1–7} (high-rank, inverse control)	Nyström, m=64	66.2%	1.45

Two things land here. First, which layers you approximate matters far more than the method: the same Nyström kernel goes from 66% flips (hard layers) to 10% (easy layers). Second, once you're operating where approximation is viable, the semantic landmark selection pays off — 9.0% vs 10.4%, at lower Frobenius error too. On the full model that advantage was buried under catastrophic failure; restricting to the safe layers is what surfaces it.

v0.2 — real accuracy, more architectures, and a sequence-length sweep

Three follow-ups turn the CIFAR-100 decision-agreement story into (a) real classification accuracy, (b) a cross-architecture check, and (c) a latency sweep that measures — rather than asserts — why efficient attention doesn't win at these lengths.

Real accuracy (Tiny-ImageNet). The prediction-flips metric exists because an ImageNet-1k head scores ~0.2% top-1 on raw CIFAR-100. Tiny-ImageNet removes that mismatch: 182 of its 200 classes exist in the ImageNet-1k head, so we drop the other 18, restrict logits to the 182, upscale the 64×64 images to 224, and evaluate 512 seed-0 images (n≈500 after the class filter). Now top-1 is a real number:

What you patch (ViT-B/16, Tiny-ImageNet, 182-class head, n≈500)	Real top-1 ↑
Exact baseline	74.0%
SVD rank-64, all 12 layers (Frobenius-optimal ceiling)	69.7%
Nyström m=64, all 12 layers	39.5%
Semantic Nyström m=64, all 12 layers	25.8%
Nyström m=64, only easy layers {0, 8–11}	73.4%
Semantic Nyström m=64, only easy layers {0, 8–11}	74.4%

Approximating the full model sheds 35–48 accuracy points. Restricted to the five easy layers, semantic patching lands at 74.4% vs the baseline's 74.0% — statistically indistinguishable at n≈500 (it also flips only 6.3% of predictions vs 8.2% for uniform landmarks). This is not "beating" the baseline; it is matching it while approximating five of twelve attention layers — which is exactly the claim the whole repo is trying to earn.

Architecture robustness (ViT-S/16, DeiT3-S/16). The failure mode is not a ViT-B artifact — it reproduces and worsens on smaller models (CIFAR-100, all 12 layers, m=64, 512 images, prediction flips):

Method (all 12 layers, m=64)	ViT-S/16	DeiT3-S/16
SVD rank-64 (Frobenius-optimal ceiling)	38.9%	53.1%
Nyström, regularised pinv	98.2%	96.3%
Semantic Nyström, regularised pinv	99.6%	99.6%

The Frobenius-optimal rank-64 ceiling flips 38.9% (ViT-S) and 53.1% (DeiT3) of predictions, versus 22.9% on ViT-B — the smaller the model, the less redundancy there is to give away, so pure low-rank bites harder. Full-model Nyström/semantic variants sit at 96–99.6% flips on both, just as on ViT-B. (The per-architecture easy-layer profile is not re-derived here — the partial-patch map is ViT-B specific; extending it per model is future work.)

Sequence-length sweep (latency-only). Does the latency verdict flip if the sequence grows? Measured over a complete ViT-B forward in eager mode on CPU (batch 4), sweeping N ∈ {197, 577, 785} via image size {224, 384, 448} — speedup vs exact SDPA:

N (image size)	Linformer	Nyströmformer	Performer	Semantic-fast
197 (224²)	0.84×	0.53×	0.62×	0.35×
577 (384²)	0.98×	0.78×	0.70×	0.63×
785 (448²)	0.99×	0.75×	0.72×	0.72×

No method crosses 1.0× at any tested length, but every trend closes the gap monotonically (Nyströmformer 0.53→0.75, semantic-fast 0.35→0.72, Linformer already 0.99× at N=785). The honest read: in eager mode on CPU, over a full ViT-B forward, the fixed overhead of the efficient kernels still dominates their asymptotic savings at these lengths; a linear extrapolation puts the crossover past N>1000. This is the same "short sequences don't favour efficient attention" caveat as before — now measured across length instead of asserted at a single point.

---

Why low-rank theory misleads

The spectral argument is not wrong — it's aggregated at the wrong granularity. Per head, approximation error rises monotonically with effective rank (Spearman ρ = +0.94 to +0.99 across SVD, uniform-Nyström and semantic-Nyström, all measured against the Act-1 ranks):

The trap is that the compressible heads (deep, near-one-hot, low rank) are already cheap and harmless to approximate — they contribute almost nothing to the error budget. The heads that actually decide the output are the high-effective-rank heads in the early layers, and those are precisely the ones a low-rank scheme cannot represent. Average the spectrum over all heads and it looks compressible; the model's decision, though, lives in the tail the average smooths away.

Concretely, that is why the partial-patch control works: the layers that survive approximation {0, 8–11} are the ones dominated by low-rank heads, and the layers that don't {1–7} are the ones carrying the high-rank heads. The effective-rank profile is a per-layer map of where approximation is safe — and reading it that way is the difference between 9% and 66% flips.

📄 The long-form version of this story — how the finding was actually arrived at, including the dead ends — is in The low-rank lie.

---

Scope & limitations

Deliberately narrow, and honest about it:

• Three architectures, not one — but all ViT-family. ViT-B/16 augreg2_in21k_ft_in1k is the primary subject; v0.2 adds ViT-S/16 and DeiT3-S/16, where the full-model failure reproduces and worsens. All three use global self-attention. Swin and other windowed/hierarchical attentions are out of scope: their attention is local per shifted window, not one global N×N matrix, so the aggregate-vs-per-layer spectral story doesn't transfer without redefining the object being approximated — that's a separate study, not a knob to flip here.
• Two datasets; still no full ImageNet-val. CIFAR-100 test (512 images) for decision agreement, plus Tiny-ImageNet val for real top-1 via the 182-class shared-head protocol (n≈500). Enough to establish the effect and now to quote real accuracy; a full ImageNet-1k validation sweep is still not run — these are effect-size measurements, not a leaderboard entry.
• CPU latencies (Apple M-series), eager mode. Wall-clock is machine-, kernel- and mode- specific; the sequence-length sweep is likewise eager CPU. Treat latency as relative signal, and remember the fused SDPA baseline is hard to beat at short sequences.
• Train-free. Attention is swapped with no fine-tuning. Learned/trained variants (e.g. Linformer's trained projection, or fine-tuning the patched model) would very plausibly recover accuracy — measuring that is future work, not a current claim.
• Short-to-moderate sequences (N=197–785). The efficient-attention methods are not expected to win on latency at these lengths, and the v0.2 sweep confirms they don't (crossover extrapolated past N>1000). The interesting axis here is accuracy fidelity vs rank, not throughput.
• Real top-1 where a compatible head exists; prediction-flips otherwise. On Tiny-ImageNet the ImageNet-1k head is directly usable (shared classes), so top-1 is meaningful and reported. On CIFAR-100 there is no compatible head — top-1 against raw labels is uninformative — so agreement with the exact model's decision remains the honest, train-free fidelity metric.

---

Reproducibility

Everything above regenerates from the shipped configs. Fixed seeds, pinned dtype, deterministic landmark selection.

# Act 2 — full-model sweep (SVD / Nyström / semantic), 512 CIFAR-100 images
vitattn bench --config configs/quick_cifar.yaml

# Act 3 — formulations, landmark-count sweep, and the spectrum-guided partial patch
vitattn bench --config configs/act3_formulations.yaml
vitattn bench --config configs/act3_sweep_m.yaml
vitattn bench --config configs/act3_partial_patch.yaml

# v0.2 — real accuracy (Tiny-ImageNet, 182-class head), cross-architecture, seq-length sweep
vitattn bench --config configs/act5_tiny_imagenet.yaml
vitattn bench --config configs/act5_vit_small.yaml
vitattn bench --config configs/act5_deit3_small.yaml
python scripts/act5_seqlen_sweep.py    # latency-only sweep over N ∈ {197, 577, 785} -> docs/figures/act5_seqlen.png

# render the Pareto / heatmap figures from any results CSV
vitattn figures --csv results/quick_cifar.csv

• CSV outputs land in results/, one row per method, each accompanied by a *.config.json capturing the exact sweep that produced it.
• Environment. Python ≥ 3.11, PyTorch ≥ 2.2, timm ≥ 1.0. CI (GitHub Actions) runs ruff + black + pytest on CPU-only PyTorch — see .github/workflows/ci.yml.
• Dev loop. make setup (venv + editable install + pre-commit), make test, make lint.

---

Extending

Adding a new approximation is a single subclass: set materializes_attn and implement _forward_impl(q, k, v, mask, need_attn) -> (out, attn_or_none). The base class owns the return_attn / materializes_attn contract so every method stays drop-in behind patch_vit. See CONTRIBUTING.md for the full walkthrough, the test conventions, and how to wire a new method into a benchmark config.

---

Context & author

Matis Despujols — computer vision, low-rank structure in transformers, and honest evaluation. GitHub @matisdsp.

The semantic-landmark selection idea originates from the author's MSc thesis; every number in this repo is independently re-measured by this benchmark from public data with from-scratch code — none of the thesis's own figures are reproduced here.

Despujols, M. (2025). Spectral Analysis and Low-Rank Approximation of Attention Matrices in Vision Transformers. MSc thesis, KTH Royal Institute of Technology.

Other published work:

• racineai/UI-DETR-1 — UI element detection.
• IEEE BigData 2025.

---

Citation

@software{despujols_vitattn_2026,
  author  = {Despujols, Matis},
  title   = {vit-attention-bench: attention approximation for Vision Transformers
             with a reproducible speed/accuracy benchmark},
  year    = {2026},
  url     = {https://github.com/matisdsp/vit-attention-bench},
  license = {MIT}
}

License

MIT © 2026 Matis Despujols.