Attention is known to be sparse in vision transformers. In particular, we analyze the unequal relevance of attention-matrix columns to the final computation and show that this phenomenon becomes increasingly pronounced with network depth. This can be interpreted as saying that tokens have an absolute importance for the attention process, independently of queries. Inspired by this, we propose a reformulation of the attention formula that makes this phenomenon explicit by adding an input-dependent bias to the query-key dot product. We show that this new term gives each token a scalar value indicating its anticipated influence on other tokens during the attention process, and can thus serve as a proxy for pruning attention maps. Across multiple transformer model sizes, this allows us to substantially reduce the computational cost of attention, in terms of both FLOPs and inference latency, while retaining most of the original classification accuracy. We publicly released the code and models at https://github.com/EIDOSLAB/AttentionWithBias.
Figures & tables
Figure 1: Bias-based column pruning. The learned bias provides a query-independent importance score for each token, allowing low-bias columns to be removed before computing keys and values. Attention is then evaluated only over the retained set K , reducing its size from N×N to N×∣K∣ .
Figure 2: Bias-based pruning pipeline. The importance score b is computed first (shaded) and selects the kept set K . All queries are kept, but keys and values are computed only for the kept tokens, so the pruned columns are never materialized.
Figure 3: Attention maps at different depths of the pretrained DeiT-S model. Vertical lines naturally emerge in the last layers.
Figure 4: Decomposition of the attention logits at different network depths.
Figure 5: Learned bias b versus the per-column mean of the query-key logits meanq(QK⊤) , at three depths. The correlation grows with depth ( r reported per panel).
Figure 6: Layer-wise Pearson correlation between the learned bias and the mean attention received by each column.
Model
Pruning schedule
Accuracy
FLOPs (G)
Latency (ms)
Top-1 (%)
Top-5 (%)
DeiT-Ti
-
72.13
91.14
199.4
47.92±0.10
DeiT-Ti b
✗
72.66
91.35
209.1
49.04±0.12
DeiT-Ti b
0.9↘0.1
69.10
89.73
122.1
32.52±0.13
DeiT-S
-
79.83
94.95
485.8
92.05±0.14
DeiT-S b
✗
80.19
94.99
505.4
94.47±0.16
Table 1: Accuracy and efficiency of bias-based attention pruning on ImageNet-1k. End-to-end inference latency, including bias computation, top- k selection, index gathering, and all remaining model operations, is measured with the math attention backend on a single 1024×1024 image on one NVIDIA A40 GPU.
Model
Pruning schedule
COCO accuracy
Backbone only
Full detector
AP Box
AP Mask
GFLOPs
Latency (ms)
GFLOPs
Latency (ms)
DeiT-S
–
43.3
38.7
485.60
67.53±0.11
1171.68
110.37±0.69
DeiT-S b
0.9↘0.1
42.7
38.1
331.17
51.36±0.09
1017.25
94.48±0.52
Table 2: Object detection and instance segmentation on COCO val2017 . Both models use Mask R-CNN with a DeiT-S backbone and the same ViTDet-style training recipe. GFLOPs and end-to-end latency are reported for the isolated backbone and the complete detector. Latency is measured in full precision with the math attention backend on one NVIDIA A40 GPU using a single 1024×1024 image.
Figure 7: Accuracy-efficiency trade-off for ToME [ 2 ] , our bias-based column pruning (star), and their combination on DeiT-S. The dashed line denotes the original DeiT-S accuracy.
Method
Pruning
Head
Top-1
FLOPs
Latency
schedule
pruning
(%)
(G)
(ms)
DeiT-S b (ours)
0.9↘0.1
–
78.55
331.3
63.08±0.13
Structured pruning
–
✗
76.80
303.0
88.13±0.09
Structured pruning
–
✓
76.26
301.3
63.56±0.10
Table 3: Comparison with 30%ℓ2 -norm structured pruning on DeiT-S. Latency is measured with the math attention backend at 1024×1024 resolution on one NVIDIA A40 GPU.
Training pruning
Inference pruning
Accuracy
Top-1 (%)
Top-5 (%)
✗
✗
80.19
94.99
✗
✓
76.90
93.01
✓
✗
77.89
94.03
✓
✓
78.55
94.31
Table 4: Effect of applying top- k pruning during training and inference on ImageNet-1k. All configurations use DeiT-S b and the 0.9↘0.1 keep-ratio schedule whenever pruning is enabled.
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
Epochs
300
Batch size
1 024
Optimizer
AdamW
Learning rate
10−3
LR decay
cosine
Warmup epochs
5
Appendix
Table 5: Training recipe used for the ImageNet100 experiments. RandAugment “2/9” denotes 2 operations at magnitude 9, following the torchvision implementation. Mixup ( α=0.2 ) and Cutmix ( α=1.0 ) intensities follow the PyTorch convention: α is the parameter of the symmetric Beta distribution from which the mixing coefficient is sampled.
Model
Pruning schedule
Accuracy
FLOPs (G)
Latency (ms)
Top-1 (%)
Top-5 (%)
DeiT-Ti
-
72.13
91.14
199.4
47.92±0.10
DeiT-Ti b
0.9↘0.5
71.47
90.55
153.0
42.22±0.07
DeiT-Ti b
0.9↘0.3
70.32
90.11
137.6
37.29±0.05
DeiT-Ti b
0.9↘0.1
69.10
89.73
122.1
32.52±0.13
Appendix
Table 6: Accuracy and efficiency of bias-based attention pruning on ImageNet-1k with different pruning schedules. Inference latency is measured with the math attention backend on a single 1024×1024 image on one NVIDIA A40 GPU.
Method
Selection rule
Pruning ratio (%)
Accuracy
Top-1 (%)
Top-5 (%)
No pruning
–
0
80.04
93.30
Probability threshold
τ=0.8
49.55
79.66
93.18
Fixed top- k
0.5
50.00
79.70
92.30
Layer-wise top- k
0.9↘0.1
50.00
80.48
93.18
Appendix
Table 7: Ablation of the column-selection strategy on ImageNet100 using ViT-S b . The pruning ratio denotes the fraction of removed attention columns.
Model
Probability threshold
Pruning ratio (%)
Accuracy
Top-1 (%)
Top-5 (%)
ViT-Ti
–
0
79.08
92.90
ViT-Ti b
✗
0
78.40
92.82
ViT-Ti b
τ=0.8
48.00
78.56
93.14
ViT-S
–
0
80.12
92.54
ViT-S b
✗
0
80.04
93.30
Appendix
Table 8: Results of the probability-threshold pruning strategy on ImageNet100. The pruning ratio denotes the fraction of removed attention columns and is measured at the end of training.
Model
Top- k value
Pruning ratio (%)
Accuracy
Top-1 (%)
Top-5 (%)
ViT-Ti
–
0
79.08
92.90
ViT-Ti b
✗
0
78.40
92.82
ViT-Ti b
0.9↘0.1
50
78.74
92.94
ViT-Ti b
0.5
50
78.68
92.80
ViT-S
–
0
80.12
92.54
Appendix
Table 9: Results of the top- k pruning strategy on ImageNet100. The pruning ratio denotes the fraction of removed attention columns. The top- k value is expressed as the fraction of tokens that are kept.
Model
Query-key-free blocks
Accuracy
Top-1 (%)
Top-5 (%)
ViT-S
–
80.12
92.54
ViT-S b
✗
80.04
93.30
ViT-S b
4
80.28
92.82
ViT-S b
6
79.30
92.72
ViT-S b
8
79.22
92.08
Appendix
Table 10: ImageNet100 accuracy when replacing the query–key term with the broadcast bias in the last transformer blocks. “Query-key-free blocks” denote the number of final blocks using only 1b⊤ .
math backend
efficient backend
Model
Original
Bias-pruned
Original
Bias-pruned
DeiT-Ti
47.92±0.10
32.52±0.13(1.47×)
33.60±0.30
20.83±0.20(1.61×)
DeiT-S
92.05±0.14
63.08±0.13(1.46×)
63.22±0.09
39.88±0.09(1.59×)
DeiT-B
195.40±0.33
139.96±0.12(1.40×)
140.59±0.16
93.61±0.16(1.50×)
ViT-L (MAE)
556.92±1.40
406.98±0.60(1.37×)
417.46±0.79
288.55±0.37(1.45×)
Appendix
Table 11: Inference latency (ms) of the original and bias-pruned models using the math and efficient attention backends. Measurements use a single 1024×1024 image on one NVIDIA A40 GPU. Values in parentheses report the speedup of the pruned model over its corresponding original model.
Model
Original
Bias-pruned
DeiT-Ti
1,886.12±53.98
966.00±86.12(1.95×)
DeiT-S
2,325.13±49.94
1,369.32±75.35(1.70×)
DeiT-B
6,770.29±70.85
3,891.03±252.29(1.74×)
ViT-L (MAE)
21,187.74±440.54
11,807.39±166.21(1.79×)
Appendix
Table 12: CPU inference latency (ms) of the original and bias-pruned models. Measurements use one 1024×1024 image on an AMD EPYC 7402 processor with 12 CPU threads and the math attention backend. We report the mean and standard deviation over 20 runs after 5 warm-up iterations. Parentheses indicate the speedup over the corresponding original model.
Figure 8: Global versus query-dependent column selection in the original pretrained DeiT-S, averaged over all attention heads and 200 ImageNet-1K validation images. (a–c) Retained attention mass at different keep ratios. (d) Oracle gap Δ=Cquery−Cglobal across all evaluated keep ratios.
Figure 9: Layer-wise Pearson correlation between the learned bias and the pre-softmax per-column mean meanq(QK⊤) . The correlation grows steadily with depth, reaching r≈0.98 in the final layer, confirming that the bias serves as an effective and cheap surrogate of column importance.
Figure 10: Attention maps for DeiT-Ti at layers {1,3,4,8,9,11} (columns), head 2, across several ImageNet validation images (rows). In all cases the attention evolves from a diagonal, locally-concentrated pattern in early layers to strongly column-sparse maps in the final blocks, where a small number of tokens attract the attention of virtually all queries.
Figure 11: Attention maps for DeiT-S at layers {1,3,4,8,9,11} (columns), head 4, across several ImageNet validation images (rows). The transition from diagonal to vertical attention structure is visible and consistent across all images, confirming that column sparsity is a systematic property of the attention mechanism rather than an image-specific artifact.
Figure 12: Attention maps for DeiT-B at layers {1,3,4,8,9,11} (columns), head 7, across several ImageNet validation images (rows). The same depth-wise evolution observed in DeiT-Ti and DeiT-S is reproduced here, showing that column sparsity is not specific to a particular model width.
Figure 13: Attention maps for the MAE-pretrained ViT-L at layers {3,7,11,15,19,22} (columns), head 11, across several ImageNet validation images (rows). Despite its greater depth and different pretraining strategy, ViT-L exhibits the same transition from predominantly diagonal attention in early blocks to increasingly column-sparse patterns in deeper blocks.
Figure 14: Decomposition of the attention logits ( QK⊤ , b , QK⊤+b ) for DeiT-Ti b at layers {1,3,5,7,9,11} , head 2. In the early layers, QK⊤ dominates with a strong diagonal pattern. With increasing depth the bias b progressively captures the vertical structures, until in the final layers it accounts for most of the visible structure in the combined logits QK⊤+b .
Figure 15: Decomposition of the attention logits for DeiT-S b at layers {1,3,5,7,9,11} , head 2. The bias term increasingly dominates the combined logits with depth, consistently with Fig. 4 and Fig. 5 of the main paper.
Figure 16: Decomposition of the attention logits for DeiT-B b at layers {1,3,5,7,9,11} , head 2. The depth-wise takeover of the bias term is visible across all three model sizes, confirming that the learned bias captures a universal property of deep vision-transformer attention.
Figure 17: Decomposition of the attention logits ( QK⊤ , b , and QK⊤+b ) for MAE-pretrained ViT-L b at layers {3,7,11,15,19,22} , head 2. The model is fine-tuned with the bias term but without top- k pruning. As depth increases, the bias progressively captures the vertical, query-independent structure of the combined logits, reproducing the behavior observed for DeiT.
Figure 18: Scatter plots of the learned bias b versus meanq(QK⊤)k for DeiT-Ti b at layers {1,3,5,7,9,11} (columns), head 2, across multiple validation images (rows). The Pearson correlation r reported per panel grows steadily with depth and remains consistent across images, confirming that the bias provides a reliable surrogate across different input images of column importance.
Figure 19: Scatter plots for DeiT-S b at layers {1,3,5,7,9,11} , head 2, across multiple validation images. The correlation reaches r≈0.98 in the final layer, independently of the input image, showing that the bias approximates the query-averaged column importance with high fidelity in the deeper blocks.
Figure 20: Scatter plots for DeiT-B b at layers {1,3,5,7,9,11} , head 2, across multiple validation images. The depth-wise growth of r is reproduced in the larger model, confirming that the relationship between the bias and the query-averaged column importance is a property of the attention mechanism rather than of a specific model scale.
Figure 21: Learned bias b versus the per-column mean meanq(QK⊤)k for MAE-pretrained ViT-L b at layers {3,7,11,15,19,22} (columns), head 2, across multiple ImageNet validation images (rows). The Pearson correlation reported in each panel becomes stronger with depth, showing that the bias also learns a reliable proxy for query-averaged column importance in ViT-L.
Transformer attention lets models relate tokens by weighting information from other positions, enabling context-dependent sequence and multimodal processing.
Transformers usually expose one inference cost per trained model, while deployed systems often need multiple cost-quality operating points. We study Budgeted Attention Allocation, a monotone head-gating mechanism conditioned on a requested attention budget. Dense warm-starting is important for stability: on a robust synthetic sequence task, one budgeted model reaches 99.7% accuracy at 0.303 estimated attention cost and 100.0% accuracy at 0.504 cost. On held-out AG News with a custom word-level transformer, hard-gate adaptation turns soft cost control into measured single-thread CPU speed, reaching 82.1% accuracy with 1.28x speedup at budget 0.50. In pretrained BERT-Mini AG News, budgeted structural pruning reaches 87.6% accuracy with 1.20x speedup at budget 0.50; a validation-ranked zero-shot dense post-hoc structural baseline reaches 86.1%, and one recovery epoch raises that per-budget specialist to 87.9%. On DBpedia14, BERT-Mini budgeted gates reach 97.4% at exact budget 0.50 versus 96.6% for dense full attention. Static fixed-budget gates and recovered dense specialists remain strong. The contribution is therefore not universal dominance, but a reproducible feasibility study of one controllable checkpoint across budgets that can trade attention cost for accuracy and be converted into measured structural speedups on small CPU benchmarks.
Pretrained transformers spend a large share of their inference compute inside self-attention, yet the per-head query-key interaction is structurally low-rank by construction. We ask a diagnostic question: across layers and architectures, which heads actually use the full rank budget on real data, and which do not? We propose SVD-Attention, a post-hoc method that decomposes the per-head interaction matrix Mh and scores every singular mode by both its spectral magnitude and its activation variance across tokens on a small calibration set. The score keeps the modes the head can amplify and that real data actually activates, and discards the rest; the resulting factorization is exact for the retained modes and needs no finetuning. Across seven pretrained models and five task types, the diagnostic finds a task-dependent pattern: redundancy sits in the early to middle layers for sentiment and image classification and for generation, and in the late layers for dense prediction and extractive question answering. It is also predictive: truncating the layers it flags costs at most 1.75 F1, while truncating the same number of layers elsewhere in the same model costs up to 50 F1. SVD-Attention is an analysis tool rather than a compression method: acting on it reduces inference FLOPs by 0.8-1.5% on its own and by 8-11% when composed with weight pruning, with 0-4% change in task quality.
Debopam Sanyal, Hongjie Chen, Riya Elizabeth Kuriyan +5
Georgia Institute of Technology · Dolby Laboratories
We study efficient reasoning under tight compute. We ask how to make structured, correct decisions without increasing test time cost. We add two training only components to small and medium Transformers that also transfer to broader differentiable optimizers. First, a length aware attention prior built via fuzzy regime position alignment, RPA, yields a normalized pre softmax bias that guides attention like a structured regularizer while adding no new inference parameters. Second, a minimal gain aware controller, Guardian, nudges attention sharpness only when validation improvements warrant it, following a two timescale policy gradient view of nonconvex optimization. It is disabled at inference. A KL perspective shows softmax of z plus log pi as MAP with KL regularization, grounding the prior in a principled objective. Under strict compute parity on WikiText 2, we reduce validation cross entropy while matching baseline latency and memory. At inference, we add a precomputed, cached prior B of T as a single additive bias per head. The controller does not run. In practice, this incurs negligible overhead, a cached bias add per head, with no measurable p50 latency shift. Our results suggest that length aware priors and late phase gain control preserve scarce improvements, especially in long span, noisy logit regimes, while keeping test time costs effectively unchanged.