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

Authors: QIUHAO Zeng, Jerry Huang, Peng Lu, Gezheng Xu, Boxing Chen, Charles Ling, Boyu Wang
Primary area: foundation or frontier models, including LLMs
Keywords: Transformer, In-context learning, Long Context, long-range Transformer, Efficient Transformer

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 existing top-$k$ attention methods 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. 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 leveraging 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 synthetic tasks Associative Recall and outperforms attention and its variants on Long-Range Arena and WikiText-103 language modeling.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1729221512420

thread_type:
official_review

<thread_object>
creator: reviewer_xji4

<content>
summary:
This paper introduced ZETA, a model to enhance the top-k attention using Z-order curves, reducing the complexity from N^2 to N log N. ZETA's query & key dimensions are much smaller than the value dimension. This is to ensure the speed and locality preservation of the Z-order curves. The authors also replaced the Softmax operation with the Cauchy-Softmax to enhance the performance of ZETA. Experiments showed that ZETA matched the performance of the attention models on Associative Recall while outperforming the attention models on Long-Range Arena and WikiText-103 language modeling.

strengths:
* The paper is well-written and the method is believable.
* Compared with the attention models, the ZETA model matches the performance on Associative Recall while outperforming on Long-Range Arena and WikiText-103 language modeling.

weaknesses:
* The effect of $d_K$ needs clarification.
  - In section 3.2.1, line 202, the authors mentioned: "a lower $d_K$ loses locality between tokens, which is crucial for efficient query".
  - However, in section 4.4, line 454, the authors mentioned that "Lower $d_K$ values exhibit a higher level of locality preservation across all sample sizes."
  - There seems to be a contradiction between these two statements, so the effect of $d_K$ needs clarification. 
* Can the authors clarify the binary expansion? What is the precision of the original d-dimensional vector?
  - Eq. (4) introduced Z-order curves, which rearrange a d-dimensional vector into a one-dimensional binary representation.
  - Suppose d=3. Suppose the precision of the numbers in the d-dimensional vector is 32 bits (e.g., the float32). Then the resulting Z-order curves will have 3*32=96 bits. 
  - Will this result in an overflow? If not, how did you set the precision of the original d-dimensional vector?
* Can the authors provide a comparison of the time and space usages for ZETA model and the attention-based models?
  - The authors claimed that ZETA reduces the time and space complexity to N log N. So I assume that ZETA runs faster and uses less memory. However, there is no actual number on this. So, the run-time metrics on time and space are needed.

questions:
* The recent state-of-the-art LLMs have at least 100k context length. Gemini 1.5 Pro even supports 2M context length. How does this impact the selection of your $d_K$ and other hyperparameters? Will $d_K$ go to 2 or even 1?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730477734347

thread_type:
official_review

<thread_object>
creator: reviewer_wjdt

<content>
summary:
The authors study an important and timely problem of addressing the quadratic complexity of dot-product attention in Transformers and propose leveraging sparsity by attending to only k most similar keys.

The authors aim to preserve query-key distances ||query_i - key_j|| after projecting queries and keys to a single scalar. To do so, they first reduce the dimension via a random projection per the Johnson-Lindenstrauss lemma.

KNN in d-dim is further reduced to sorting in 1D via the following trick: Let q and k be d-dim vectors. Let the binary representation of the i'th element q[i] be (q_i_31,...q_i_0). Then we can represent q as a number N(q) having a large binary representation (q_0_31,...,q_d-1_31,...q_0_0,...,q_d-1_0). Then |N(q)-N(k)| is likely to preserve the ordering of ||q-k|| as distances are more influenced by more significant bits of each element. This is not always true as q_0_31,...,q_d-1_31,... imposes an ordering on the dimensions whereas ||q-k|| is symmetric. Also note that most data types are limited to 64 bits.

Assuming that this does preserve the pairwise distances of queries and keys, we can sort the key scores N(k1),..,N(kL) and for each q_i use binary search to determine the most similar keys to attend to. The authors point out that for causal LM, we need to exclude future keys from being considered, and propose to use a chunking mechanism where positions are chunked and future key chunks are excluded from consideration.

As the above method is suitable for ||q-k|| and not <q,k>, the authors propose replacing dot-product attention with Cauchy attention where similarity of q and k is computed as 1 / (||q-k||^2 + gamma^2) instead of exp(<q,k> / sqrt(d)). Recent work on linear attention has shown this kernel to be performant.

strengths:
1. The proposed method is interesting, and the trick of projecting to 1D for query-key similarity computation is intuitive.

2. Significant and comprehensive theoretical justification is provided for dimensionality reduction via random projection and is well-supported via detailed proofs.

3. The authors show that this method works well on the multi-query associative recall task that has been used to justify the in-context learning abilities of Transformers. On language modeling on Wikitext-103, the method also works as well as vanilla Transformer.

4. The authors demonstrate that this method outperforms full dot-product attention on Long Range Arena, which is a popular benchmark for stress testing long-range abilities across various modalities.

5. The authors have included a pytorch implementation of their method. I appreciate this and will take it positively into account in my evaluation.

weaknesses:
1. Distances are not preserved even for modest key dimensions, which might be an issue when scaling up. It is not clear if this method can be used as a drop-in replacement for existing pretrained models due to the change in attention similarity.

2. The authors claim that smaller key dimensions do not affect performance. This is not a well-established fact - more comprehensive experiments are required to justify this claim.

3. The chunking part of keys and queries could have been more clearly written, or pseudocode of the method in an appendix would have helped. I looked at the provided torch code in the supplementary but this part looks involved.

4. Benchmarking of speedup / memory usage is not provided.

5. SOTA performance on LRA benchmark is much higher. The performance on the Path-X task is not included, and the authors do not mention this exclusion.

questions:
1. Can you include more details on the chunking part. Also explain where efficiency is coming from.
2. Can you provide resource usage results especially for long-sequences. Also include it for flash attention and other methods such as SSMs which perform better on the tasks that you include the paper.
3. Can you discuss the limitations of your method.
4. Can you provide results on Path-X task in LRA?

~~I am willing to change my score if the authors can address these concerns.~~  
**[Post author rebuttal]**: The authors have largely addressed my concerns and I have increased my score.


Typos:
- Equation 6: "(q-Ki)" should be "||q-Ki||"
- Line 232: Period missing after Lemma 3.1
</content>
</thread_object>

</thread>
<thread>
creation date:
1730688202917

thread_type:
official_review

<thread_object>
creator: reviewer_pvhs

<content>
summary:
Paper is well written, with solid experimental work and theoretical analysis.
Authors introduce Efficient Parallel Top-k Attention by utilising Z-order Curves.  As this implies approximation, thorough experimental evaluation is performed to on several datasets including Long Range Arena, WikiText, MQAR to estimate computation efficiency vs model performance  tradeoffs. On top of this of paper investigate effects of Dimensionality Selection for Key and Query Pairs and Adaptive Cauchy-Softmax Mechanism which help to carefully select tradeoffs.

strengths:
Paper address important area of computation cost (quadratic run time) of classical self-attention, which may be significant for long context. Novel Efficient Parallel Top-k Attention by utilising Z-order Curve was proposed, carefully analysed and evaluated.
On top of this paper this paper investigate dimensionality of keys and queries the trade-off between the curse of dimensionality
and the preservation of relative distances for keys and queries and introduces Adaptive Cauchy-Softmax which allows dynamic adjusting of receptive fields to enhance attention’s flexibility.

weaknesses:
- Evaluation is done on very small scale models (125M), so it could be that they won't translate to large size models.
- While Top-K should take O(N log N) and potentially should be more computation efficient,  it is not clear how much was gained in practice.  Also providing some more details about experiments (what is number of chunks, etc) would be helpful.

questions:
'The insertion position of a query in the key sequence can then be found using a binary search (e.g. with torch.searchsorted), allowing us to retrieve the top-k attended tokens using a window centered on the insertion position'
Could you please clarify this. As far as I understand to find nearest k neighbours,  inspecting several segments of Z-order curve should be necessary or some false positives would be included?


Could you please share computation efficiency measurements (for example average training step time) for baseline transformer and ZETA?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699380185

thread_type:
official_review

<thread_object>
creator: reviewer_36fb

<content>
summary:
This paper introduces ZETA, a model leveraging Z-order curves and Adaptive Cauchy-Softmax for efficient top-k attention, specifically aimed at handling long-sequence tasks. The authors propose a novel approach to reduce dimensionality for keys and queries, preserving locality while enabling faster token retrieval through Z-order-based projections. Their experimental results demonstrate competitive or superior performance on various benchmarks, suggesting ZETA’s potential as an efficient alternative for long-range attention tasks.

strengths:
Originality : The paper introduces a novel use of Z-order curves for dimensionality reduction in attention mechanisms, adding a fresh approach to efficient top-k attention.
Quality:  Methodologically sound with theoretical insights like the Johnson–Lindenstrauss Lemma for dimensionality reduction. Experimental results across benchmarks support the model’s performance, though more tests could strengthen findings.
Clarity : The paper is well-organized and accessible, with clear explanations of techniques like Adaptive Cauchy-Softmax and Z-order projection, aiding reader comprehension.
Significance : ZETA shows potential in handling long sequences more efficiently, which could benefit applications in NLP and large-scale modeling. Further testing on diverse tasks would solidify its impact.

weaknesses:
bout the experiment about  d_k​ in Experiments:
The authors test a narrow range of dkd_kdk​ values (1, 2, 3, 8, and 32), which may not reflect real-world applications where larger dimensions are common. Why weren’t higher dkd_kdk​ values explored? Expanding this range, particularly with larger model dimensions, could yield more generalizable insights.

Unclear Transferability of \gamma in Adaptive Cauchy-Softmax:
The transferability of the trainable parameter γ\gammaγ in Adaptive Cauchy-Softmax across different tasks or datasets is unclear. Could the authors offer evidence or discussion on its robustness beyond the tested contexts?
Lack of Formal Justification for Euclidean Distance Claim:
The claim that Euclidean distance is more effective for low-dimensional top-k methods lacks formal support. Providing a proof or theoretical justification would enhance the validity of this argument.

questions:
Why choose ASSOCIATIVE RECALL in section 4.2 ?
Testing different d_k values on more Long Range Arena (LRA) tasks would provide a fuller picture of task-specific impacts on dimensionality selection. Could the authors evaluate dkd_kdk​ across a wider range of LRA tasks to strengthen their findings?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732260273248

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response + Revision Summary

comment:
We gratefully acknowledge all the reviewers for their insightful comments. We are happy that the reviewers describe our work as **interesting/original** [Reviewers 36fb, pVHs, WjdT], **significant/important** [Reviewers 36fb, pVHs, WjdT], **methodologically and theoretically sound** [Reviewers 36fb, pVHs, WjdT, xJi4]. We also appreciate that most found our work **clear and well-written** [Reviewers 36fb, pVHs, xJi4]. We would also like to specifically thank Reviewer WjdT for acknowledging our theoretical proofs as well as making the added effort to understand the code implementation that was attached.

To summarize our revisions and rebuttals, we have

- We have included tables (provided below) measuring speed (in milliseconds) and memory consumption of our implementation in rebutalls and a revised version of our manuscript in a new **Section 4.6**. (To pVHs & WjdT & xJi4)
- Provided detailed explanations of our **Triton implementation of ZETA** in **Appendix D and E**, including both the forward pass and the gradient derivations for the backward pass of effective top-$k$ indexing and Cauchy softmax. Additionally, include the Triton-based version of ZETA in the supplementary materials.
- Provided additional insight into the choice of $d_k$ through additional experiments in **Appendix G.1**. (To 36fb & WjdT & xJi4)
- Provided justification for design choices, such as our Adaptive Cauchy Softmax in **Appendix G.2**. (To 36fb).
- Clarified points that some reviewers found unclear, such as justifying our claim regarding Euclidean Distance (To 36fb) and clarifying our methodology such as the chunking mechanism, which is also added an illustrative example in **Appendix B** (to pVHs & WjdT).


Furthermore, while we have addressed the following point in individual rebuttals, we would like to again acknowledge the common question regarding the efficiency of our method. As this was a common point raised by all reviewers, we want to acknowledge that we have taken this important point into account and wish to again provide these results in this general response.

| Length | Torch Attention (FWD) | Torch Attention (FWD+BWD) | Mamba (FWD) | Mamba (FWD+BWD) | Flash Attention (FWD) | Flash Attention (FWD+BWD) | ZETA (FWD) | ZETA (FWD+BWD) | 
| ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- |------------- | ------------- |
| 4096 | 44.3 | 117.9 | 7.1 | 14.0| 3.4 | 29.2 | 5.6 | 38.2 |
| 8192 | OOM | OOM |11.8 |23.0|  12.8 | 111.5 | 11.0 | 76.4 |
| 16384 | OOM  | OOM | 23.5 | 45.7 | 50.4 | 437.7 | 21.7 | 152.6 |
| 32768 | OOM  | OOM | 47.3 | 91.8 | 198.2 | 1733.5 | 43.0 | 304.8 |
| 65536 | OOM  | OOM | 94.0 | 183.7 | 805.3 | 7044.1 | 85.8 | 608.2 |

where *OOM* stands for out-of-memory. This in particular demonstrates that our implementation is significantly more efficient that both a naive implementation of attention as well as the latest Flash Attention release. We in particular want to highlight the fact that our method does not suffer from OOM issues while becoming significantly more efficient than Flash Attention as the sequence length increases, validating our claim regarding ZETA theoretical time complexity. Compared to Mamba, ZETA costs slightly less time in forward. 

GPU MEM Consumption Results are provided in MB:
| Length | Torch Attention (FWD) | Torch Attention (FWD+BWD) | Mamba (FWD) | Mamba (FWD+BWD) | Flash Attention (FWD) | Flash Attention (FWD+BWD) | ZETA (FWD) | ZETA (FWD+BWD) | 
| ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- |------------- | ------------- |
| 4096 | 17268.1 | 25972.1 | 574.2 | 632.2 | 886.1 | 1784.1 | 1314.1 | 1926.1 |
| 8192 | OOM | OOM | 904.2 | 1020.2 | 1528.1 | 3324.1 | 2382.1 | 3606.1 |
| 16384 | OOM | OOM | 1564.2 | 1776.2 | 2812.1 | 6404.1 | 4520.1 | 6968.2 |
| 32768 | OOM | OOM  | 2884.2 | 3200.2| 5380.1| 12564.1| 8796.2| 13692.2 |
| 65536 | OOM | OOM |5524.2| 6048.2| 10516.1| 24884.1| 17348.2| 27140.3 |

From this table, ZETA outperforms the standard PyTorch implementation of attention by a large margin. Admittedly, ZETA consumes slightly more memory than the latest Flash Attention, but in totality we believe our results here demonstrate ZETA to be as good if not better than Flash Attention. 

We would also acknowledge that significant improvements are likely further possible for ZETA as we (the authors) are not as experienced with system-level programming such as CUDA and Triton and therefore our implementation is likely not as mature as a library such as Flash Attention. However, we would like to highlight our efforts in making our implementation as strong as possible and hope that this table is sufficient to leave the reviewers with the same impression, which we hope can be taken into consideration.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes ZETA, a novel method for efficient parallel top-k attention using Z-order curves to successfully reduce computational complexity, enabling better scaling to long sequences while maintaining or improving performance compared to the standard attention mechanism.
Reviewers were consistently positive about the paper's technical contributions and extensive empirical validation, noting clear innovations in efficient attention, giving progress in an area of high interest and relevance for recent AI models.
All 4 reviewers recommend acceptance.