# ZETA: Leveraging Z-order Curves for Efficient Top-k Attention

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 53 references that are in the map.
- First posted on arXiv 24 Jan 2025.

## Builds on (most cited first, 30 of 40)

- [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)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attention](https://mlsanity.com/paper/2102.03902/) (2021)
- [Efficient Attention: Attention with Linear Complexities](https://mlsanity.com/paper/1812.01243/) (2018)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [cosFormer: Rethinking Softmax in Attention](https://mlsanity.com/paper/2202.08791/) (2022)
- [Simple linear attention language models balance the recall-throughput tradeoff](https://mlsanity.com/paper/2402.18668/) (2024)
- [Scaling Neural Machine Translation](https://mlsanity.com/paper/1806.00187/) (2018)
- [Zoology: Measuring and Improving Recall in Efficient Language Models](https://mlsanity.com/paper/2312.04927/) (2023)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)

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

- [Long-Context Generalization with Sparse Attention](https://mlsanity.com/paper/2506.16640/) (2025)
- [Mamba Modulation: On the Length Generalization of Mamba](https://mlsanity.com/paper/2509.19633/) (2025)
- [Homophily Enhanced Graph Domain Adaptation](https://mlsanity.com/paper/2505.20089/) (2025)
- [Full spectrum Unlearnable Examples via Spectral Equalization](https://mlsanity.com/paper/2606.26719/) (2026)
- [Z-Order Transformer for Feed-Forward Gaussian Splatting](https://mlsanity.com/paper/2605.13465/) (2026)
- [Balancing Computation Load and Representation Expressivity in Parallel Hybrid Neural Networks](https://mlsanity.com/paper/2505.19472/) (2025)
- [k-Maximum Inner Product Attention for Graph Transformers and the Expressive Power of GraphGPS](https://mlsanity.com/paper/2604.03815/) (2026)
- [ARM: Attention with Routed-Memory for Learnable Sparse Control](https://mlsanity.com/paper/2609.24417/) (2026)
- [DimPO: Dimensionality Reduction for Attention using Preference Optimization](https://mlsanity.com/paper/2609.32579/) (2026)

## Abstract (arXiv)

Over recent years, the Transformer has become a fundamental building block for sequence modeling architectures. Yet at its core is the use of self-attention, whose memory and computational cost grow quadratically with the sequence length N, rendering it prohibitively expensive for long sequences. A promising approach is top-k attention, which selects only the k most relevant tokens and achieves performance comparable to vanilla self-attention while significantly reducing space and computational demands. However, causal masks require the current query token to only attend to past tokens, preventing the existing top-k attention method from efficiently searching for the most relevant tokens in parallel, thereby limiting training efficiency. In this work, we propose ZETA, leveraging Z-Order Curves for Efficient Top-k Attention, to enable parallel querying of past tokens for entire sequences. % in both space and time complexity of O(N log N). We first theoretically show that the choice of key and query dimensions involves a trade-off between the curse of dimensionality and the preservation of relative distances after projection. In light of this insight, we propose reducing the dimensionality of keys and queries in contrast to values and further leverage Z-order curves to map low-dimensional keys and queries into one-dimensional space, which permits parallel sorting, thereby largely improving the efficiency for top-k token selection. Experimental results demonstrate that ZETA matches the performance of standard attention on the synthetic textscMulti-Query Associative Recall task and outperforms attention and its variants on textscLong Range Arena and textscWikiText-103 language modeling.