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

TL;DR

ZETA addresses the quadratic cost of self-attention for long sequences by introducing a locality-preserving, one-dimensional projection of small-dimension keys/queries via $Z$-order curves, enabling parallel top-$k$ token retrieval under causal masking. It pairs this with Adaptive Cauchy-Softmax to robustly capture short- and long-range dependencies in a low-dimensional space, achieving $\mathcal{O}(N \log N)$ time/space complexity. Across MQAR, Long Range Arena, and WikiText-103, ZETA matches or surpasses baseline attention variants, while providing substantial efficiency gains on long sequences. The approach holds practical impact for scalable sequence modeling on GPUs, combining efficient kNN-style retrieval with dynamic softmax to adapt receptive fields across layers.

Abstract

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 \textbf{Z}-Order Curves for \textbf{E}fficient \textbf{T}op-$k$ \textbf{A}ttention, to enable parallel querying of past tokens for entire sequences. % in both space and time complexity of $\mathcal{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 \emph{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 \textsc{Multi-Query Associative Recall} task and outperforms attention and its variants on \textsc{Long Range Arena} and \textsc{WikiText-103} language modeling.

Figures (4)

• Figure 1: Illustration of attention using Euclidean distance vs. dot product. Euclidean distance correctly classifies points into classes $\pm$1, while the dot product leads to a misclassified area.
• Figure 2: Experiments on Associative Recall: (a) Model Accuracy (b) Performance of Transformer with varying $d_K$ across different model dimensions; even with low $d_K$, the model achieves near-perfect performance (c) Comparison of different Euclidean-based Softmax operators across varying key-query dimensions $d_K$ (d) Ablation on $k$ in ZETA.
• Figure 3: The effect of dimensionality reduction before and after $Z$-order curves projection on locality preservation for different sample sizes.
• Figure 4: Illustration of the chunking process in ZETA: Keys are projected into one-dimensional space using $Z$-order curves, sorted, and partitioned into chunks for efficient retrieval of top-$k$ keys for each query.

Theorems & Definitions (9)

• Lemma 3.1
• Theorem 3.3
• Lemma A.1: Johnson–Lindenstrauss Lemma
• Lemma A.3
• proof
• Lemma A.4
• proof
• Theorem A.5
• proof