Research map / Architectures and efficiency / Attention and transformers
ZETA: Leveraging $Z$-order Curves for Efficient Top-$k$ Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ZETA, a top-k attention method that projects low-dimensional queries and keys onto a 1D Z-order curve for parallel sorting-based retrieval under causal masks, paired with an adaptive Cauchy-softmax.
Problem
Top-k attention with causal masks cannot search for relevant past tokens in parallel across a sequence, which limits training efficiency for long sequences.
Approach
Reduce query/key dimensionality (d_K=d_Q=3, much smaller than value dimension), map them to 1D with Z-order curves, sort and chunk keys so each query retrieves its top-k neighbors in parallel under causal masking, and replace softmax with a trainable Cauchy kernel (Adaptive Cauchy-Softmax).
Key results
- ZETA reaches 63.51 average accuracy on Long Range Arena, the best among the compared models (Transformer 54.39).
- ZETA achieves 26.3 test perplexity on WikiText-103 with 124M parameters, comparable to the vanilla Transformer (26.2) and behind CosFormer (23.1).
- ZETA matches standard attention on MQAR, reaching perfect accuracy at model dimension 256.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ZETA instead of H-Transformer-1D | LRA models · LRA average | accuracy: 61.41 -> 63.51 | Helped |
| Use ZETA instead of H-Transformer-1D | LRA models · LRA Image | accuracy: 46.05 -> 64.39 | Helped |
| Use ZETA instead of H-Transformer-1D | LRA models · LRA Text | accuracy: 78.69 -> 64.52 | Hurt |
| Use ZETA instead of Vanilla Transformer | WikiText-103 language models · WikiText-103 test | perplexity: 26.2 -> 26.3 | Hurt |
| Reduce key-query dimension to d_K = 2 (ablation) | Transformer · MQAR | accuracy | No clear effect |
| Use d_K = 1 instead of larger key-query dimensions (ablation) | Transformer · MQAR | accuracy | Hurt |
| Use Cauchy Softmax instead of Negative Euclidean or Inverse Euclidean (ablation) | Transformer · MQAR | accuracy | Helped |
| Vary k from 16 to 48 (ablation) | ZETA · MQAR | accuracy | No clear effect |
| Increase key-query dimension d_K before Z-order projection (ablation) | ZETA · Locality preservation, top-64 neighbor overlap | overlap | Hurt |
| Use ZETA instead of Flash Attention | attention implementation · sequence length 65536 forward | time (milliseconds): 805.3 -> 85.8 | Helped |
| Use ZETA instead of Flash Attention | attention implementation · sequence length 65536 forward + backward | time (milliseconds): 7044.1 -> 608.2 | Helped |
| Use ZETA instead of Flash Attention | attention implementation · sequence length 65536 forward | memory consumption (MB): 10516.1 -> 17348.2 | Hurt |
Surprising findings
- On WikiText-103, ZETA (26.3) is essentially tied with the vanilla Transformer and does not beat CosFormer (23.1), despite strong LRA results.
- Reducing the key/query dimension to as low as d_K=2 or 3 leaves MQAR accuracy near perfect.
Limitations
- Causal chunked search and Z-order projection rely on low d_K, and the paper argues the trade-off theoretically but gives a loose bound.
- On WikiText-103 ZETA does not outperform the best baselines (CosFormer 23.1).
- Uses more memory than Flash Attention in the efficiency benchmark.
Experimental setup
Datasets: MQAR Long Range Arena WikiText-103
Benchmarks: MQAR Long Range Arena WikiText-103
Metrics: accuracy perplexity latency (ms) memory (MB)
Baselines: Vanilla Transformer Performer Based Reformer CosFormer Flash Attention Mamba IceFormer
Builds on: Top-k attention Reformer Johnson-Lindenstrauss lemma Based
Closest papers by meaning
- Loki: Low-rank Keys for Efficient Sparse Attention (NeurIPS 2024 · similarity 0.86)
- QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.86)
- Tensor Product Attention Is All You Need (NeurIPS 2025 · similarity 0.86)
- FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference (ICLR 2025 · similarity 0.85)
- Star Attention: Efficient LLM Inference over Long Sequences (ICML 2025 · similarity 0.85)
- RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.