# ATP: Enabling Fast LLM Serving via Attention on Top Principal Keys

arXiv: https://arxiv.org/abs/2403.02352 · HTML: https://mlsanity.com/paper/2403.02352/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving.
- Builds on 21 of its 34 references that are in the map.
- First posted on arXiv 1 Mar 2024.

## Placed in

- Efficient LLM inference and serving / [KV cache optimization](https://mlsanity.com/inference/kv-cache-optimization/)

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [From Words to Watts: Benchmarking the Energy Costs of Large Language Model Inference](https://mlsanity.com/paper/2310.03003/) (2023)
- [Fast Transformers with Clustered Attention](https://mlsanity.com/paper/2007.04825/) (2020)
- [HyperAttention: Long-context Attention in Near-Linear Time](https://mlsanity.com/paper/2310.05869/) (2023)
- [Sharpness-Aware Minimization Leads to Low-Rank Features](https://mlsanity.com/paper/2305.16292/) (2023)
- [FMMformer: Efficient and Flexible Transformer via Decomposed Near-field and Far-field Attention](https://mlsanity.com/paper/2108.02347/) (2021)
- [3LegRace: Privacy-Preserving DNN Training over TEEs and GPUs](https://mlsanity.com/paper/2110.01229/) (2021)

## Built on by (most cited first, 1 of 1)

- [Eigen Attention: Attention in Low-Rank Space for KV Cache Compression](https://mlsanity.com/paper/2408.05646/) (2024)

## Abstract (arXiv)

We propose a new attention mechanism with linear complexity, ATP, that fixates Attention on Top Principal keys, rather than on each individual token. Particularly, ATP is driven by an important observation that input sequences are typically low-rank, i.e., input sequences can be represented by a few principal bases. Therefore, instead of directly iterating over all the input tokens, ATP transforms inputs into an orthogonal space and computes attention only on the top principal bases (keys). Owing to the observed low-rank structure in input sequences, ATP is able to capture semantic relationships in input sequences with a few principal keys. Furthermore, the attention complexity is reduced from quadratic to linear without incurring a noticeable performance drop. ATP further reduces complexity for other linear layers with low-rank inputs, leading to more speedup compared to prior works that solely target the attention module. Our evaluations on various models (e.g., BERT and Llama) demonstrate that ATP achieves comparable accuracy with much lower computation and memory complexity than the standard attention mechanism. In particular, ATP barely loses accuracy with only 1/2 principal keys, and only incurs around 2% accuracy drops with 1/4 principal keys.