← Back to homepage

The pains and gains of Transformers

Imagine AI as a collection of exquisite vessels. The Transformer is a common architectural principle that shapes raw materials into diverse forms - be they wine glasses, bowls, vases, or countless others, designed for different purposes. It has been a fundamental building block of Computer Vision (CV) models, Vision-Language Models (VLMs), and Large Language Models (LLMs) since its emergence[1]. While Transformers excel at modeling global dependencies through a fully connected self-attention graph[2], they inherently suffer from the notorious \(O(N^2)\) computational complexity, where N denotes the number of nodes(tokens, or patches) in the fully connected graph.

The computational cost of Transformers comes from two main sources: self-attention, with \(O(N^2)\) complexity in the sequence length \(N\), and feed-forward networks (FFNs, usually two-layer MLPs), with \(O(N \cdot d^2)\) complexity, scaling linearly with the sequence length. These two computational bottlenecks have inspired two complementary directions for improving Transformer efficiency: efficient attention attempts to reform the \(O(N^2)\) attention computation itself, while token pruning directly reduces \(N\).

Efficient Attention

A large body of work has sought to alleviate the quadratic complexity of self-attention. It typically approximates or restructures the attention operation itself, avoiding full pairwise computation. Linear Attention [3] reformulates the self-attention mechanism by expressing the softmax similarity function as a dot product between kernel feature maps, leveraging the associativity of matrix multiplication. This formulation also reveals the intrinsic connection between Transformers and recurrent neural networks (RNNs), enabling efficient autoregressive modeling with faster training and inference. The approximation error is usually non-trivial.

Another branch is FlashAttention[4], which does not approximate attention but instead reorganizes GPU memory access patterns via tiling and recomputation. It reduces memory footprint to \(O(N)\). FlashAttention is an exact computation of attention with no approximation error, and is now the dominant engineering solution in practice, for LLMs and VLMs. Other methods include grouped query attention / multi-query attention (GQA/MQA)[5], where multiple query heads share a single set of key/value projections. This primarily targets the KV cache memory bottleneck at inference time and offers limited training speedup.

Focusing on vision, window-based attentions are able to alleviate the same issue from a different perspective, by restricting self-attention to local regions. When the window or kernel size is relatively small compared to the input resolution, the computational complexity can be significantly reduced. For example, Dilated Neighborhood Attention Transformer(DiNAT) [6] extends Neighborhood Attention [7] into a dilated local attention capable of preserving locality, maintaining translational equivariance, and exponentially expanding the receptive field.

Different from window attention, which is a reduced form of patch-to-patch attention, PaCa-ViT[8] learns to cluster patches into a predefined number of \(M\) clusters that serve as keys and values, relaxing the complexity to linear \(O(N \cdot M)\) and demonstrating greater efficiency than PVT models [9] and Swin [10].

Token Pruning

Another line of research reduces the computational cost by directly decreasing the number of input tokens. This is intuitive for vision tasks because visual information is inherently redundant. The core idea is that not all tokens are equally important, especially across layers, so a subset can be discarded mid-network without materially affecting the final output. A branch of pruning (or merging) methods is criteria-based (e.g., importance or similarity), such as EViT [11], DynamicViT [12], ToMe [13], TokenPooling [14], and SP-ViT [15]. These methods either use attention weights over visual tokens as an importance proxy or rely on a learnable predictor. Some others are clustering-based, represented by ST-ViT [16].

The efficiency gains of token pruning come from several aspects:

In language tasks involving long-context inputs, where the number of tokens can reach millions, additional studies have leveraged the sparsity of attention weights in pretrained LLMs. Representative approaches include heuristic pruning based on attention patterns [17], low-rank decompositions [18], learned thresholding [19], predictive pruning of token indices [20], and entropy-inspired approaches [21] that dynamically allocate computation based on content relevance.

How to choose?

The two approaches are not mutually exclusive. Combining FlashAttention with token pruning has become a natural and increasingly common engineering practice in VLM deployment. Token pruning is particularly well suited for VLMs because visual tokens are inherently much more redundant than text tokens. High-resolution VLMs (e.g., InternVL [22]) can produce more than 1,000 visual tokens per image, creating severe computational pressure. By contrast, text tokens have near-zero redundancy: removing a single word often breaks semantic integrity. This asymmetry leads to a natural division of labor. In pure LLM settings, token pruning is high-risk, and efficient attention (especially FlashAttention) is the safer and more widely adopted solution. In VLM visual encoders, however, token pruning is a highly natural fit.

Flexibility is an additional advantage of token pruning. Some methods, such as ToMe[13], are completely training-free, while some others can be trained independently without modifying the backbone. By contrast, architectural changes to the attention mechanism often require retraining the entire model, introducing substantial computational overhead.

Limitations

Both approaches have their limitations.

For efficient attention, approximation errors introduced by linear attention are often non-negligible. Sparse attention patterns usually rely on hand-designed inductive biases and often have limited transferability. FlashAttention[24] does not reduce the quadratic complexity itself, but accelerates attention by minimizing memory I/O through hardware-aware tiling, making it the most widely used solution in practice today.

For token pruning, pruning decisions are irreversible: once a token is discarded, no downstream layer can recover its complete information. Aggressive pruning can significantly degrade performance on tasks that require fine-grained local details, and almost all existing methods incur a non-zero accuracy drop on most benchmarks, especially at high pruning ratios (e.g., SparseVLM [23]). Moreover, the pruning module itself (e.g., cross-attention or importance predictors) introduces additional overhead that may not be worthwhile for short input sequences.

Overall, token pruning and efficient attention tackle the same problem from orthogonal directions. In VLMs, token pruning has a clear advantage because visual tokens are naturally redundant, allowing simultaneous reductions in both attention and FFN computation. Efficient attention, especially FlashAttention, remains the mainstream solution for LLM backbones, where sequence elements cannot be safely discarded.

References

  1. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., & Houlsby, N. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929.
  2. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention is all you need. Advances in Neural Information Processing Systems, 30.
  3. Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020, November). Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning (pp. 5156–5165). PMLR.
  4. Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems, 35, 16344–16359.
  5. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (pp. 4895–4901).
  6. Hassani, A., & Shi, H. (2022). Dilated neighborhood attention transformer. arXiv preprint arXiv:2209.15001.
  7. Hassani, A., Walton, S., Li, J., Li, S., & Shi, H. (2023). Neighborhood attention transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 6185–6194).
  8. Grainger, R., Paniagua, T., Song, X., Cuntoor, N., Lee, M. W., & Wu, T. (2023). PaCa-ViT: Learning patch-to-cluster attention in vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 18568–18578).
  9. Wang, W., Xie, E., Li, X., Fan, D. P., Song, K., Liang, D., Lu, T., Luo, P., & Shao, L. (2022). PVT v2: Improved baselines with pyramid vision transformer. Computational Visual Media, 8(3), 415–424.
  10. Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., & Guo, B. (2021). Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 10012–10022).
  11. Liang, Y., Ge, C., Tong, Z., Song, Y., Wang, J., & Xie, P. (2022). Not all patches are what you need: Expediting vision transformers via token reorganizations. arXiv preprint arXiv:2202.07800.
  12. Rao, Y., Zhao, W., Liu, B., Lu, J., Zhou, J., & Hsieh, C. J. (2021). DynamicViT: Efficient vision transformers with dynamic token sparsification. Advances in Neural Information Processing Systems, 34, 13937–13949.
  13. Bolya, D., Fu, C. Y., Dai, X., Zhang, P., Feichtenhofer, C., & Hoffman, J. (2022). Token merging: Your ViT but faster. arXiv preprint arXiv:2210.09461.
  14. Marin, D., Chang, J. H. R., Ranjan, A., Prabhu, A., Rastegari, M., & Tuzel, O. (2021). Token pooling in vision transformers. arXiv preprint arXiv:2110.03860.
  15. Kong, Z., Dong, P., Ma, X., Meng, X., Niu, W., Sun, M., Shen, X., Yuan, G., Ren, B., Tang, H., et al. (2022). SPViT: Enabling faster vision transformers via latency-aware soft token pruning. In European Conference on Computer Vision (pp. 620–640). Springer.
  16. Chang, S., Wang, P., Lin, M., Wang, F., Zhang, D. J., Jin, R., & Shou, M. Z. (2023). Making vision transformers efficient from a token sparsification view. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 6195–6205).
  17. Jiang, H., Li, Y., Zhang, C., Wu, Q., Luo, X., Ahn, S., Han, Z., Abdi, A. H., Li, D., Lin, C. Y., et al. (2024). MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention. Advances in Neural Information Processing Systems, 37, 52481–52515.
  18. Wang, S., Li, B. Z., Khabsa, M., Fang, H., & Ma, H. (2020). Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768.
  19. Kim, S., Shen, S., Thorsley, D., Gholami, A., Kwon, W., Hassoun, J., & Keutzer, K. (2022). Learned token pruning for transformers. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (pp. 784–794).
  20. Akhauri, Y., AbouElhamayed, A. F., Gao, Y., Chang, C. C., Jain, N., & Abdelfattah, M. S. (2025). TokenButler: Token importance is predictable. arXiv preprint arXiv:2503.07518.
  21. Xiong, J., Shen, J., Ye, F., Tao, C., Wan, Z., Lu, J., Wu, X., Zheng, C., Guo, Z., Kong, L., et al. (2024). UNComp: Uncertainty-aware long-context compressor for efficient large language model inference. arXiv preprint arXiv:2410.03090.
  22. Chen, Z., Wu, J., Wang, W., Su, W., Chen, G., Xing, S., Zhong, M., Zhang, Q., Zhu, X., Lu, L., et al. (2024). InternVL: Scaling up vision foundation models and aligning for generic visual-linguistic tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 24185–24198).
  23. Zhang, Y., Fan, C. K., Ma, J., Zheng, W., Huang, T., Cheng, K., Gudovskiy, D., Okuno, T., Nakata, Y., Keutzer, K., et al. (2024). SparseVLM: Visual token sparsification for efficient vision-language model inference. arXiv preprint arXiv:2410.04417.
  24. Dao, T. (2024). FlashAttention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (pp. 35549–35562).