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

2025-01-24

AI Summary

To address the high computational cost of Transformer self-attention on long sequences and the inherent sequential dependency in causal top-$k$ attention, this paper proposes ZETA: a fully parallelizable causal top-$k$ attention mechanism. Its core innovation is the first integration of the Z-order space-filling curve into attention computation—leveraging low-dimensional key/query projections followed by Z-order linearization to enable distance-aware, efficient sorting while strictly preserving causality. This design overcomes the fundamental limitation of conventional causal top-$k$ attention, which necessitates sequential position-by-position computation. Experiments demonstrate that ZETA matches standard full attention in Multi-Query Associative Recall, and significantly outperforms diverse baselines—including sparse, locality-sensitive, and learned variants—on Long Range Arena and WikiText-103. ZETA thus achieves a favorable trade-off among computational efficiency, modeling accuracy, and theoretical soundness.

Technology Category

• Machine Learning: Causal Learning
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Causality

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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 extbf{Z}-Order Curves for extbf{E}fficient extbf{T}op-$k$ extbf{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 extsc{Multi-Query Associative Recall} task and outperforms attention and its variants on extsc{Long Range Arena} and extsc{WikiText-103} language modeling.

Problem

• Long Sequence Data
• Transformer Efficiency
• Causal Attention

Innovation

• ZETA
• Z-Order Curve
• Efficient Attention Mechanism