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
Venue: ICLR 2025 Poster (2025)
Keywords: Transformer
In-context learning
Long Context
long-range Transformer
Efficient Transformer
Primary area: foundation or frontier models, including LLMs
TL;DR: A novel sparse attention architecture for efficient parallel top-k attention search in one-dimensional space

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.

Published as a conference paper at ICLR 2025
ZETA: LEVERAGING Z-ORDER CURVES FOR EFFICIENT TOP-k ATTENTION
Qiuhao Zeng**
Jerry Huang 90
Peng Lu
Gezheng Xu
Boxing Chen
Charles Ling**
Boyu Wang**
University of Western Ontario
Université de Montréal
Mila
Noah's s Ark Lab
* Vector Institute
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 most relevant tokens and achieves per- formance comparable vanilla self-attention while significantly reducing space and computational demands. However, causal masks require the current query token only attend to past tokens, preventing 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. We first theoretically show that the choice of key and query dimensions involves trade-off between the curse of dimensionality and the preservation of relative distances after projection. 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 the synthetic MULTI-QUERY ASSOCIATIVE RECALL task and outperforms attention and its variants on LONG RANGE ARENA and WIKITEXT- 103 language modeling.
INTRODUCTION
Transformers Vaswani et al. 2017 have become indispensable for sequence modeling across var-
ious domains OpenAA et al. 2024 Zeng et al. 2024a b) 2023] Fang et al.] (202;) 2025). including
natural language processing (NLP) DDellin et al. 2019 Brown et al.) 2020; OpenAl et al. 2024 Jiang al. 2024), computer vision (Dosovitskiy et al. 202 Ramesh et al. 2rooks Brooks et al.
2024). etc. The foundation of Transformer models is the self-attention mechanism. This mecha- nism (Bahdanau al.] 2015), inspired by recurrent neural networks (RNNs) and their ability con- struct representations from all elements in a sequence, has revolutionized numerous fields, enabling breakthroughs in tasks such as language modeling (Radford al. 2019). machine translation (Otl) et al. 2018), text generation (Brown et al. 2020). image classification (Touvron et al. 2021] and video generation (Brooks al.) 2024).. However, self-attention has quadratic complexity both memory and computation as the sequence length N increases, which presents significant challenge when scaling to long sequences (Child et al., 2019; [Betagy et al.) (020)). This makes the direct ap- plication of self-attention in large-scale problems computationally prohibitive for many real-world applications, particularly when long sequences are involved (Tay et al., (021).
Recent advances have explored strategies mitigate the inefficiencies of vanilla self-attention. One such approach is top-k attention, which focuses computation on subset of the most relevant to- kens, significantly reducing memory and computation costs while maintaining competitive perfor-
mance (Kitaev et al. 2020 Gupta et al. 2021 Bertsch et al. 2023 Mao et al. 2024). However,
"Correspondigg author: Boyu Wang.

---

Published as a conference paper at ICLR 2025
existing top-k attention methods (Kitaev al. 2020 Zhuoran al. 2021 typically apply causal masking after selecting the top-k tokens, causing earlier tokens (i < N) to often attend to nothing as their top-k relevant tokens may include future tokens that are masked out by the causal mask. Alternatively, [ao et al. (2024) process input tokens by token to exclude masked-out keys from the k-nearest neighbors search, preventing future tokens from influencing currently generated ones. Consequently, with causal masks, current top-k attention approaches fail to fully leverage the paral- lel computation abilities of modern accelerators, limiting their efficiency in Long-sequence modeling.
To overcome the limitations of existing top-k attention methods, we introduce ZETA, a novel model designed to search for the top-k tokens within chunked one-dimensional dineesnnnsi sorted key sequences projected via Z-order curves. Specifically, our approach strikes a balance between mitigating the "curse of dimensionality" and preserving the relative distance of token represntations after projection by carefully selecting a lower dimensionality for keys and queries in contrast to values. This reduction allows the query and key to map to a one-dimensional space using Z-order curves, preserving proximity. As a result, ZETA efficiently performs top-k token selection in paral- lel within this one-dimensional space with accelerators. Additionally, since the aforementioned top- k search is based on the Euclidean metric for low dimensional data, directly applying the traditional dot product based softmax function is not appropriate. To address this, we propose an Adaptive Cauchy-Softue -oottmx mechanism that replaces the exponential function in the attention operation with a trainable Cauchy kernel (Billingsley 1986). This enables dynamic adjustment of receptive fields across layers. providing greater flexibility in capturing both short and long- Longerange dependencies.
Extensive empirical evaluations show that ZETA matches the performance of standard self-attention on the Associative Recall task and consistently outperforms existing attention variants on the Long- Range Arena (LRA) and WikiText- 103 datasets. We summarize our key contributions as follows:
Efficient Parallel Top-k Attention: We introduce ZETA, a novel model that enables top-k attention to operate in parallel across entire sequences, significantly improving training and inference efficiency with time complexity of O(N log N).
Dimensionality Selection for Key and Query Pairs: We theoretically show that the di- mensionality of keys and queries decides the trade-off between the curse of dimensionality and the preservation of relative distances for keys and queries.
Z-order Curve Integration: By leveraging Z-order curves, we enable efficient top-k to- ken selection in one-dimensional space, allowing the use of parallel sorting algorithms on GPUs for faster attention computation.
Adaptive Cauchy-Softman Mechanism: We introduce Adaptive Cauchy-Softmax, a Soft- max variant with trainable parameters based on the Cauchy kernel, dynamically adjusting receptive fields to enhance attention's 's flexibility.
2 RELATED WORKS
The Transformer architecture (Vaswani et al. 2017) is foundational for se- quence modeling, but its quadratic complexity limits efficiency with long sequences. Various ef- ficient variants (Tay et al.] 2022] 2020) Chen et al. [021] Qin et al. 20226] Zhang et al. 2024) have been proposed as alternatives, mainly categorized into sparse, low-rank, and linear transform- ers. Sparse transformers, such as BigBird (Zaheer et al.) 2020) and Longformer (Beltagy et al.) 2020), restrict attention to local windows or global tokens achieve linear complexity. SparseAx- aal (Ho et al. (020) further enhances this by combining sparse attention with axial mechanisms for high-dimensional inputs. Reformer (Kitaev al., (020) locality-sensitie hashing to handle variable-length sequences efficiently. Low-rank transformers like Linformer (Wang et al. (020) re- duce the attention matrix to lower-dimensional space, reducing memory and computation costs. Linear transformers such as Performer (Choromanski et al. (221 use kernel-based approximations for linear-time complexity, while Nyströmformer (Xiong et al. 2021) leverages Nyström decompo- sition for near-linear performance.
Efficient Transformer
Top-k Attention (Gupta et al. 2021) falls under the category of sparse attention, reducing atten- tion complexity by selecting only the top-k most relevant tokens at each layer. thereby focusing computational resources on the most critical interactions. Unlimiformer (Bertsch et al.] 2023) en- ables transformers to handle arbitrarily long sequences by chunking inputs and using a retrieval
2

---

Published as a conference paper at ICLR 2025
mechanism to attend to relevant past contexts. Similarly, IceFormer (Mao al. 2024 improves transformer efficiency by integrating a k-nnarest-neighhor (KNN) search mechanism that focuses on the KNN results the most relevant tokens during inference, bypassing the need compute the full attention matrix. However, with causal masks, these approaches can not compute the outputs of long sequences in parallel, making them less efficient for training models from scratch by not fully exploiting the parallel computing power of GPUs. In contrast, ZETA performs KNN-based searches for relevant tokens in parallel across the entire sequence on GPUs using chunking tech- niques, enabling efficient training and inference with a time and space complexity of (NN log N).
METHODOLOGY
3.1 PRELIMINARIES
Attention has proven to be a fundamental building block in modern deep learning, particularly in natural language processing and sequence modeling tasks. It allows a model to focus on specific parts of the input sequence, thereby capturing dependencies within the sequence more effectively. In the standard formulation of attention (Vaswani et al., 2017), given a set of queries Q E IR N xxda, IR N Nxdy, the attention scores are computed as
keys K R. Rg.xdy, and values V
Attention QQ K, V) = softmax (QK") lvdkk) V.
(1)
where dx, do and dv is the dimension of the keys, queries and values, respectively. It is common to use the same value for all three.
Top-k Attention approaches aim to enhance the efficiency and flexibility of traditional self-attention mechanisms by focusing attention to only the most relevant tokens in a sequence. Both methods reduce the computational overhead associated with self-attention, especially for long sequences, by selectively attending to the most important tokens rather than computing attention scores across the entire sequence. This process lowers the computational complexity from O(N2) in traditional self- attention to ONN k), where N is the sequence length and k is the number of selected tokens. The top-k highest scores are selected for each query:
1 = {i qK. 2 2x
(2)
where I represents the set of indices corresponding to the top-k highest attention scores for the query q, 7) the threshold defined as the k-th highest attention score. The attention is then re- stricted to the tokens in this subset, reducing the number of operations required for long sequences. Specifically, the attention score for the top-k tokens is recalculated using the self-attention mecha- nism but limited to the selected indices:
Attention (g, K, K,V)
Σ iE
softmax (gK7 / dx V;
(3)
where q denotes the query vector, K; and V, denote the key and value vectors corresponding to the top-k relevant tokens, respectively.
Z-order Curves (Dugundji 1966). also known as Morton codes, provide a way to map multi- dimensional data into one--imensione, space while preserving locality, whereas other dimension- ality reduction methods (Abdi & Williams 2010 McInnes et al. 2018) are not designed for map- ping data to a ID space. This approach is valuable in tasks that require efficient spatial indexing or key-query matching, such as attention. By maintaining the relative proximity of data points after projection, Zordeer curves ensure that points that are close together the original multi-dimmensonall space remain close in the projected one-dimensional space.
The Z-order curve interleaves the binary representations of each coordinate in a multi-dimennsionll point. For point in the d-dimensional space with coordinates x (21, (2 x; is a binary number, the Zoorder curve computes a scalar value Z by interleaving the bits of each coordinate. Given the binary representation of I; as bu ... bin the Zorder curve is expressed as:
I.)) where each
Z by b2i bai b122222 ... baz ... ban .ban
3

---

Published as a conference paper at ICLR 2025
where n refers to the number of bits used represent each coordinate 2; in its binary form. Through this interleaving of bits, the Z-order curve creates a scalar value that allows efficient sorting or indexing of points while approximately maintaining their original spatial relationships.
The primary advantage of Z-order curves is their ability to preserve locality. In other words, nearby points in the original multi-dimensional space have similar Z-values in the projected one- dimensional space. This property enables efficient search and selection processes in attention mech- anisms or spatial indexing, where key-query pairs can be processed more efficiently in the one- dimensional space without significantly losing the locality information from the higher-dimennionall space. Z-order curves are designed to preserve locality, and hence not suitable for dot- product similarity measures, which does not reflect locality.
3.2 SEARCHING FOR THE TOP-k ATTENDED TOKENS IN ONE- ONE-DIMENSIONAL SPACE
Since we project key and query vectors into a one-dimensional space using Z-order curves, using large dæ can still distort locality (as shown in Figure 3) and compromise the preservation of relative distances. Thus we ask whether dx can be reduced in a way such that even after mapping to one dimension, relative distances between tokens are maintained. Importantly, the key and query dimensions dæ and de do not have to match the dimension of the values dy. This is because dy should remain large to capture more semantic information, as seen with Gaussian distributions, where higher dimensionality increases the measure of information entropy (Cover & Thomas 2006). Hence, as long as the relative distances between queries and keys are preserved, dx and do can be reduced. The following theoretical analysis provides insights into the selection of dg.
3.2.1 THEORETICAL ANALYSIS ON dx
We first introduce the in high-dimensiooal space can be projected into a much lower-dimensional subspace using random projections while approximately preserving the pairwise distances between the points. Since random projections can preserve locality, this provides justification for setting a smaller dx with trainable projection functions for keys and queries, which could also preserve locality.
Lemma (Johnson al. 1986), which states that data
Lemma 3.1. (Johnson-Lindenstrruus Lemma) For any 0 positive integer such that d = 22 mm Then for any set x of m points in RD, there exists a map R + I2d such that for all 24, E X,
and any integer m, let d be
(1 (1- - 2,1 S (1((x)) f(r,)1 < +)]la, x2,1
(5)
The following assumption then provides a mathematical depiction that attention weights are con- strained within an m--dimensional simplex, and the learnable similarity function I outputs the atten- tion scores, ensuring the most relevant tokens are emphasized during the information aggregation process. This reflects the primary goal of attention: to aggregate critical information for more accu- rate predictions.
Assumption 3.2. Let a E Am-1 be an element of the m-dimensional simplex, defined as Am-1 - {ar E RT aa; 2 0, Ei 0; 1]. Assume that hattn equipped with can achieve an optimal learnable similarity critic function r, where the attention scores are given by = softmax (T (ff( f(x))) such that I' is trained to be optimal to have the minimal expected risk: min, | Szi a) yll. where hatin denotes the attention- based hypothesis, x is the input, and Sz is the context.
Theorem 3.3 highlights the importance of choosing dx carefully, as it controls for a trade-off be- tween locality preservation and the curse of dimensionality. Larger dx allow for more detailed feature capture at the cost of the high-dimensional curse, leading to increased complexity. On the other hand, smaller dx loses locality between tokens, which is crucial for efficient query. The bounds provided give valuable insights into the underlying mechanisms of attention and can guide future designs of more efficient attention models. For simplicity. assume WLOG that keys and queries share the same projection functions, as Kitaev et al. (2020).
Theorem 3.3. Let X E Rdd, y E R. and D be a distribution over X x y for which the conditional probability function, h dd - RR,, is a ipsscchiit function. Let h denote hypothesis, and hattn
4

---

Published as a conference paper at ICLR 2025
denote the one-layer attention model aggregate the predictions of sample set S ~ DD predict for another i..d sample x. Specifically, here we assume the same linear map for a key mapping fk and a query mapping fy as Rd [-B, B|dx where B is the bound of projection of by f. and we assume the value mapping fu be the Bayes optimal rule ht* = =E hypothesis that minimizes [r)(h) over all functions. Then, the expected risk of the regression tasks (h) = (abx)>/llh(())) - yll with h as battn can be upper bounded by
( X = æ/ which is the
[D((assea] S LD(h') 4lcy/d Bm- 1/(xx+1) Dm
C Clam
dx
This indicates that dx should be carefully chosen rather than simply being set equal to dy (as is common practice). Empirically, we show that the dimension of query and key can decreased without degrading performance, more specifically in Figure 2b
Secciin 4.2 which we later discuss.
3.2.2 TOP-k SEARCH IN ONE DIMENSION
As the sequence length becomes extremely large, iterating through the entire context history to search for the top & tokens becomes infeasible. Ideally, searching should be efficient and aim to achieve the optimal time complexity of ON log N), similar to general sorting problems with arbi- trary inputs. To achieve this, we map both keys and queries into a -ie-dimesioona space using the Z-order curve and sort them with a sorting operation that can be executed in parallel on an accelera- tor, e.g. the torch. sort operator in Py'Torch (Paszke et al. 2019). The insertion position of query in the key sequence can then be found using a binary search (e.g. with torch. gearchsortedy allowing us to retrieve the top-k attended tokens using window centered on the insertion position.
Specifically, the key and query dimensions are set to be significantly smaller than that of the values, i.e dx de < dv. While values need a high dimensionality carry rich semantic information, keys and queries primarily serve to preserve relative distances, which can be achieved with much lower dimensionality as argued by Lemma [3.]] To facilitate the fast retrieval of queries from keys, we leverage sorting, which can be efficiently parallelized, after dimensional space via the Z-order curve. We define Q, K E R Bxxxxxx Xxdx, where B is the batch size and is the sequence length. The Z-order transformation applied as follows:
queries and keys into one-
Q= = Z-order(Q),, K: (K)
where Q and K2 are the one-dimensional representations of Q and K, respectively.
With causal masks, directly collecting the top-k tokens from the entire sequence is not guaranteed to have plenty of tokens to make inference. For instance, for query at position 32 in sequence of length 2048, selecting the top 16 tokens from the entire sequence followed by causal masking would leave approximately 32 16 = tokens to attend to, effectively leaving tokens available for the query. Consequently, the query at this position would not attend to anything, rendering top-k attention ineffective. To enable parallel k-nearest neighbors (kNN) searching considering causal masks, we first sort the Z-order keys and divide them into chunks. For query in the m-th chunk (where m [ and M is the chunk size), restrict search to the first chunks, indexing the original unsorted keys from 0 to m x M - in the sorted list. This ensures that future keys j E (j : m x M) are excluded, in accordance with the causal mask requirements. This process is performed in parallel for every query.
2045
Next, perform the nearest neighbor search in these one-dimensional Z-order spaces. For each query, the insertion position is first found using a binary search and followed by selecting the nearest keys using a window of size K centered around the insertion point collect top-k tokens as Ig. denoting the indices of the k-nearest neighbors. This ensures efficient and accurate retrieval while maintaining the constraints imposed by causal masking.
3.3 ADAPTIVE CACCHY-SOFTMAX
Top-k token searching relies similarity metric between data points, and prefer the Euclidean metric for top-k attention with small dx for two main reasons. First, as illus- trated in Figure I Fuclidean distance is more effective for low-dimensional data in top-k
5

---

Published as a conference paper at ICLR 2025
methods: it reliably identifies the correct class in one-dimensional classification tasks, lead- ing to accurate predictions, whereas the dot product can misleading (samples in the misclassified area will be classified as "+1" using the dot-product metric). Second, k-NN search is typ- ically based on the Euclidean metric, while using the dot-product requires normalization that loses token magnitudes.
Class: -1
Class: +1
Figure 1: Illustration of attention using Eu- clidean distance vs. dot product. Eu- clidean distance correctly classifies points into classes =1, while the dot product leads to a misclassified area.
To better align with the Euclidean measure for low- dimensional representations, we propose the Adap- tive Cauchy Softmax function, to replace the expo- nential function in traditional attention mechanisms. The Cauchy kernel, with its heavier tails, ensures that distant tokens retain influence, overcoming the limitations of the exponential function, which suppresses distant tokens (Shen al. 2023). This allows the attention mechanism to capture both local and global dependencies, with the shape of the kernel determining how key vectors influence the query. Specifically, the Adaptive Souchy-Softmax between query vector q and keys K is computed as:
7 11 K; + ?
[lqq K; + 7
softmax. K) ΣjEr, 2 [1lq K,
(6)
+ +2
Ejer, 1|| - K, + 7
where y is a trainable parameter that controls the shape of the distribution. By training task-specific y for each attention layer, the model adjusts receptive fields dynamically. We define 72 as the output of sigmoid function applied to trainable parameter, ensure range of [0. 1]. with smaller values sharpening attention and improving focus on relevant inputs (Qin et al.] 2022a Zhang et al.) (024). while larger values allow for smoother attention. The adaptive Cauchy softmax effectively handles long-range dependencies, preventing entropy collapse (Zhai et al. 2023) or explosion (Zhang et al., 2024) and adaptively balancing attention across the sequence.
3.4 SPARSE ATTENTION WITH Z-ORDER CURVE FOR EFFICIENT KNN RETRIEVAL.
Instead of calculating full attention scores, which is computationally expensive and memory in- tensive, we compute sparse attention scores by leveraging the Z-order curve for efficient nearest neighbor retrieval. ZETA is then computed as below, according to E'quatin
Atlention (Q, (2,K = Σ softmax. (Q, V/
(7)
iEL
where K, and V. are the corresponding keys and values for the indices Ig. As result of the sparsity of top-k, most of the tokens will not join the predictions, which stops the gradient back pook-propagated through the low-probabilitics tokens and fails to leverage this current prediction's s information. We append the mean vector of the history tokens to the top-k tokens matrix using cumsum function in Ig. to keep the attention from assigning zero probability, which can be regarded smoothing n-gram language model (Juratsky & Martin 2024).
4 EXPERIMENTAL RESULTS
We evaluate ZETA's 's performance on several aspects: ZETA's ability to solve the synthetic MULTI- QUERY ASSOCIATIVE RECALL (MQAR) task (Arora et al., (024a),, long sequence modeling abil- ity the LONG RANGE ARENA (LRA) benchmark and auto-regressive language modeling on WIKITEXT-I0S: Then we conduct extensive analysis experiments: an ablation study examining the influence of dimensionality on attention model performance (Section 4.2), ablations various Euclidcan- Softmax operators (Section 4.3), the empirical results locality preservation using Z-order curves (Section 4.4) and an ablation study over the number of A in ZETA (Section 4.5).
4.1 EMPIRICAL VALIDATION
Associative Recall. Associative recall tasks (Arora et al. 2024a) have been popular for testing the ability of language models to look up information in their context. Broadly, they involve feeding

---

Published as a conference paper at ICLR 2025
0.8
0.9
Attention
0.95
0.8
0.8
8
0.7
0.90
0.7
Basod
0.6
k=
ZETA
0
Negad Eudikkaan
0.5
21
0.4
Coucty Softmax
0.4
0.3
0.75
42
Inverse Buclldean
32
128
256
32
128
256
32
128
256
Model dimension
Model dimension
Model dimension
Figure 2: Experiments on Associative Recall: (a) Model Accuracy (b) Performance of Transformer with varying dx across different model dimensions; even with low dx, the model achieves near- perfect performance (c) Comparison different Buclidean-basee Softmax operators across varying key-query dimensions dæ (d) Ablation on k in ZETA.
auto-regressive model pairs of key-value associations and then prompting the model to produce the correct completion upon being shown a previously seen key. The MULTI-QUERY ASSOCIATIVE RECALL (MQAR) task is a particular formulation of this task that requires the model to memorize multiple associations (Arora al. 2024a). We evaluate the performance Associative Recall task, a classical grueence-o--eqqenee task that requires the model to recall the first token associated with a target token after processing long sequence. The task is tests the ability of models to capture long-range dependencies and to maintain information over time.
various models on the
We compare the performance of four models: a vanilla Transformer, Performer [Choromanski al, 2021). BASED (Arora et al. 2024b), and ZETA, with different model dimensions (32, 64, 128, and 256). As illustrated in Figure 2a accuracy increases as the model dimension grows. Attention and Based models show strong performance with higher dimensions, achieving nearly perfect accuracy for dimensions larger than 64. ZETAfollows a similar trend and achieves competitive performance, especially for larger model dimensions, with perfect accuracy at dimension 256. In contrast, the Performer struggles, showing significantly lower accuracy across all dimensions.
Long Range Arena (LRA). The LONG RANGE ARENA (LRA) benchmark (Tay et al., 2021) is a comprehensive suite designed to evaluate the performance of models on long sequence tasks. It includes tasks that span across multiple domains, such as natural language processing, image classi- fication, and mathematical reasoning. LRA focuses sequence classification, challenging models to efficiently process longer input sequences while capturing long-range dependencies, providing an ideal testbed for Transformer models and their efficient variants.
LRA consists of five key tasks: LISTOPS, TEXT, RETRIEVAL, IMAGE, and PATHFINDER. Each task evaluates different aspects of long-range de- pendency handling such the ability to handle mathematical reasoning tasks on long sequences of operations (LiSTOPs), capture dependencies over long textual inputs (TEXT), retrieve relevant ele- ments from long sequence (RETRIEVAL) and cap- ture spatial dependencies (IMAGE). PATHFINDER presents difficult problem where models must dis- tinguish between connected and disconnected paths within maze-like patterns. modified version of PATHFINDER, called PATHFINDER-XX X, is also in- cluded where the patterns are presented larger image 2256x256 compared to 32x32) but has yet to be solved by existing attention-basee methods.
Table 1: Test perplexity (lower is better) on WikrTExT-103.
Model
Params
Test PPL
Vanilla Transformer
125M
26.2
Performer Reformer
125M
26.8
125M
25.6
AFT-conv
125M
28.2
Linear Transformer
125M
30.2
RFA-Gause FFA--Gausiaa
125M
27.5
CosFormer
125M
23.1
ZETA
124M
26.3
We evaluate various Transformer-based models, including several linear and efficient variants, trained from scratch on the LRA sequence classification tasks. For each model, adopt the same hyperparameter settings provided by the official LRA benchmark (Tay et al. (02T)) to ensure a fair comparison. Results summarized in Table 2 which compares the performance of the models across all five tasks, along with their average accuracy, showing that ZETA significantly outper- forms other attention-based models.

---

Published as a conference paper at ICLR 2025
Table 2: Results of the Transformer and various variants on LRA. We consistently outperform the next closest competitor (Zhu & Soricut 202I).
Model
ListOps
Text
Retrieval
Image
Pathfinder
Average
Transformer
36.37
64.27
57.46
42.44
71.40
54.39
Reformer
37.27
56.10
53.40
38.07
68.50
50.67
Sparse Trans
17.07
63.58
59.59
44.24
71.71
51.24
Sinkhorn Trans
33.67
61.20
53.83
41.23
67.45
51.29
Linformer
35.70
53.94
52.27
38.56
76.34
51.36
BigBird
36.05
64.02
59.29
40.83
74.87
55.01
Linear Trans.
16.13
65.90
53.09
43.40
75.30
50.76
Performer
18.01
65.40
53.82
42.77
77.05
51.41
Nyströmformer
41.28
58.38
65.40
37,54
71.76
54.87
H- I-Transformer-ID
49.53
78.69
63,99
46.05
68.78
Top-k Attention
38.12
63.72
59.14
X
53.66
IceFormer
41.53
60.01
66.02
40.46
74.42
56.49
cosFormer
37.90
63.41
61.36
43.17
70.33
55,23
Skyformer
39.25
64.70
82.06
40.77
70.73
59,50
Hedgehog
37.15
64.60
82.24
40.15
74.16
59.66
ZETA
42.52
64.52
77.92
64.39
68.20
63.51
evaluate several models the WIKITEXT- 103 (Merity et al.) (017). a widely used benchmark for language modeling containing 100 million tokens extracted from high-quality Wikipedia articles characterized by large vo- cabulary and long-range dependencies. This makes it challenging benchmark for testing model's ability to predict the next token sequence. We use perplexity (PPL) as the primary evaluation metric, where lower scores indicate better ability to capture the sequential structure within natural language text. Table 1 shows a vanilla Transformed to achieve a test perplexity of 26.2. How- ever, linear approximation models such as the Linear Transformer (Qin et al.) (022a) struggle to compete, with higher perplexity values of 30.2 on the test set. The table further compares several other models, including efficient attention mechanisms like Performer (Choromanski et al. 2021), Reformer (Kitaev et al., 2020), and CosFormer (Qin et al., 2022b). Notably, CosFormer achieves the lowest perplexity the test set with score of 23.1, outperforming all other models. Reformer also shows competitive results, achieving a perplexity of 25.6, improving on the Vanilla Transformer. ZETA achieves perplexity of 26.3, comparable the Vanilla Transformer.
Autoregressive Language Modeling. Furthermore,
The results highlight the trade-offs between using conventional transformers, linear transformers, and models adopting approximate attention mechanisms like ZETA. It reinforces the importance of balancing computational efficiency with model performance, particularly in the context of long- sequence language modeling tasks, especially the information necessary solve task becomes sparsely located within long contexts.
4.2 EFFECT OF VARYING dx ON ASSOCIATIVE RECALL TASK
We next evaluate the effect of different key-query dimensions d& on the Transformer model's per- formance for MQAR. The model dimensions varied between {32, 64, 128, 256} while adjusting dx to values {1, 2, 3, 8, 32}. As shown in Figure 2b the performance remains near-perfect even with low dx values, such dx = 2. The model achieves close to 100% accuracy across all model dimensions, except for the smallest dimension (dx = 1), where performance slightly drops for lower model dimensions. This demonstrates that the Transformer is capable of handling long-range dependencies in the Associative Recall task, even with relatively low key-query dimensionality.
These results suggest that reducing dx does not significantly impair the model's ability recall in- formation in sequence tasks. In fact, maintaining low dæ can provide computational savings with- sacrificing performance, especially when model dimensions are large. This indicates that while random projections-suce as those used in the Johnson-Lindenstrauus Lemma-approxing approximately pre- serve distances, trainable projection networks fü and fg can better adapt to task ssecifccc data and more effectively retain locality even with low dg. For instance, by setting dx as low as 3, we
We auto-regressive Transformer based Biderman al. 2023 for comparision.
8

---

Published as a conference paper at ICLR 2025
reduce it from the typical head dimension (normally 32, i.c. feature dimension 512 with heads). We can further mitigate information loss by configuring fs and fg as two-layer neural networks rather than single-layer ones.
4.3 PERFORMANCE OF EUCLIDEAN- BALLLDEANN--EEDD SOFTMAX OPERATORS
We further evaluate the performance of transformers with various Euclidean-based Softmax op- erators on MQAR. Specifically, we compare Negative Euclidean, Cauchy Softmax (our proposed method), and Inverse Euclidean operators. The goal of this experiment is to test how these different formulations of Softmax handle varying key-query dimensions dæ in terms of accuracy.
As shown in Figure 2c the proposed Cauchy Softmax consistently outperforms the other operators across all values of dk. It achieves near-perfect accuracy for dg 2 2, whereas Negative Euclidean shows a drop in accuracy for lower dx values. Inverse Euclidean, while performing comparably at higher dimensions, struggles significantly at lower values dx (e.g., dx = 1).
These results highlight the advantage of using the Cauchy distribution for smaller dx. as it allows for better handling of long-range dependencies and achieves more stable performance across various key-query dimensions. The heavier tails of the Cauchy distribution enable distant tokens to retain non-negligible influence, which crucial for tasks like Associative Recall where long-range token relationships are important.
4.4 LOCALITY PRESERVATION AFTER Z-ORDER CURVE PROJECTION
Next, we evaluate how well Z-order curve projections preserve locality across different dimensions and sam-
1.0
N-512
0.9
ple sizes. Specifically, we test the locality preserva-
N=1024
tion by measuring the overlap between the top-64 nearest neighbors before and after projection, with sample sizes
0.8
N=2048
0.7
N (512, 1024, 2048).
0.6
0.5
Figure shows the relationship between locality preser- vation and the dimensionality dx- dx increases, the overlap between the top-64 nearest neighbors dimin- ishes for all sample sizes, indicating a decrease in locality preservation. Lower dx values exhibit higher level of locality preservation across all sample sizes. However, for larger sample sizes, such N 2048, the drop in locality preservation is more pronounced as the dimen- sionality increases. We select dx for ZETA.
0.4
0.3
4
6
8
dk Figure 3: The effect of dimensional-
ity reduction before and after Z-order curves projection on locality preserva- tion for different sample sizes.
These results highlight the importance of choosing an appropriate dx for maintaining locality, es- pecially for larger datasets where higher dimensionality can lead distortions in spatial proximity after a projection.
4.5 THE EFFECTS OF 1 IN ZETA
As an ablation study, we explore the effect of varying & in the attention mechanism of ZETA. The goal of this experiment is to analyze how different values of influence the model's performance on the associative recall task across different model dimensions.
Figure 2d shows ZETA to achieve near pear-perfect accuracy across all model dimensions (32, 64, 128, and 256) for different values of k ranging from 16 48. In most of our experiments, we set k = 32, as it provides good balance between performance and computational efficiency. Interestingly, there is little variation in accuracy between different values of k; indicating that ZETA is robust to changes in this parameter.
4.6 EFFICIENCY BENCHMARKING
In order to better understand the effectiveness ZETA, we further conduct an experimental study to demonstrate its computational efficiency in comparison to existing attention methods. In particular,

---

Published as a conference paper at ICLR 2025
compare with naive attention implementation from PyTorch (based on Vaswani et al. (2017) as well as an 10-aware Attth-Attntion (Dao et al. on Triton.
2022 Dao 2024). Our implementation is based
Table 3: Time (in milliseconds) for different operations to compute for fixed-sized batch of varying sequence length. Our method outperforms naive attention implementation across all lengths while also outperforming Flash-Attention by a signficant margin as the sequence length increases.
Method
Torch Attention
Mamba
Flash Attention
ZETA
Input Length
FND
FWD+BND
FWD
FND
FWD
FWD+BWD
FWD
FWD+BND
4096
44.3
117.9
7.1
14.0
29.2
5.6
8192
0OM
0OM
11.8
23.0
12.8
111.5
11.0
76.4
16384
0OM
0OM
23.5
45.7
50.4
437.7
21.7
152.6
32768
0OM
0OM
47.3
91.8
198.2
1733.5
43.0
304.8
65536
0OM
0OM
94.0
183.7
805.3
7044.1
85.8
608.2
Table 4: Memory consumption (in MB) for different operations compute for fixed-sized batch of varying sequence length. Our method outperforms a naive attention implementation across all
lengths while only marginally trailing highly optimized Flash Attention implementation.
Method
Torch Attention
Mamba
Flash Attention
ZETA
Input Length
FMD
FWD+BWD
FND
FWD+BWD
FWD
FWD+3WD
FND
FWD+BWD
4096
17268.1 I
25972.1
574.2
632.2
886.1
1784.1
1314,1
1926.1
8192
00M
00M
904.2
1020.2
1528.1
3324.1
2382.1
3606.1
16384
00M
0OM
1564.2
1776.2
2812.1
6404.1
4520.1
6968.2
32768
00M
0OM
2884.2
3200.2
5380.1
12564.1
8796.2
13692.2
65536
00M
OOM
5524.2
6048.2
10516.1
24884.
17348.2
27140.3
Table 3indicates the time required for both a forward pass as well forward-backwe pass using our efficient ZETA implementation as well as the aforementioned attention implementations. We observe that our implementation significantly outspeeds a naive implementation of attention and do not suffer from out of memory issues while also outperforming Aatentntttntion for long sequences, with widening gap the sequence length increases. This indicates both the computational effi- ciency of our method as well as serves as an empirical validation of the ONN log N) complexity of ZETA which previously justify theoretically. Furthermore, if we compare with Mamba (Gu & pass
Dao 2024). we demonstrate that ZETA has faster forward pass while the maintains similar relative performance the sequence length increases.
Table meanwhile shows that ZETA uses less memory than a naive attention implementation while also only slightly utilizing more memory than a highly optimized Flash Attention implementation. Nevertheless, in comparison to a sequence model such as Mamba, all attention models predictably use more memory due to the use of softmax-type operations.
CONCLUSION
In this paper, we presented ZETA, model designed enhance the efficiency of top-k attention by leveraging Z-order curves for parallel token selection in one-dimensional space, reducing both time and space complexity to ONN log N). By carefully selecting the dimensionality of key and query pairs, ZETA effectively preserves relative distances, improving both locality and computational effi- ciency. Our comprehensive experiments on synthetic associative recall, LRA, and WIKITEXT- 103 demonstrate that ZETA consistently matches or outperforms traditional attention mechanisms, mak- ing it particularly well-suited for long-sequence tasks that demand scalability and efficiency. Addi- tionally, the introduction of the Adaptive Cauchy-Softmax mechanism enhances ZETA's 's flexibility, enabling it to handle long-range dependencies more effectively and efficiently. Overall, ZETA offers robust, scalable, and efficient solution for sequence modeling, combining adaptive token selection with dynamic softmax to optimize performance across a range of tasks and datasets.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We appreciate constructive feedback from anonymous reviewers and meta-reviewers. This work is supported by the Natural Sciences and Engineering Research Council of Canada (NSERC), Discov- cry Grants program.
REFERENCES
Hervé Abdi and Lynne J Williams. Principal component analysis. Wiley interdisciplinary reviews: computational statistics, 2(4):433-459, 2010.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Re. Zoology: Measuring and improving recall in efficient language models. In The Twelfth International Conference on Learning Representations, 2024a. URL /ttps://openrevieee
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Re. Simple linear attention language models balance the recall-throughuut tradeoff. In Forty-first International Conference on Machine Learning, 2024b.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations, 2015.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint wwXiv:2004.05//55 2020.
Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew R. Gormley. Unlimiformer: Long-range transformers with unlimited length imput. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https::t//// /opp:rs//ppeeeeweew wwwp:eee/ooooommdn//eemeeo
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Patrick Billingsley. Probability and Measure. John Wiley and Sons, second edition, 1986.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng. Ricky Wang, and Aditya Ramesh. Video generation models as world simulators, 2024. URL https:://opeaaicc compeeeaaccom//eeercchh/
sideo-generation-nee modens-odllaaa as wots--owdd-smmlatooorrro
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN
Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. Skyformer: Remodel self-attention with gaussian kernel and nyström method. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL httph:tt/ppeeeeiiiiiie/ het/foormmidd=p2cGG7ggkkk
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: /904. 10509, 2019.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021. URL https: wtwpee/wwwwweeeeeeroiddddddeeeee

---

Published as a conference paper at ICLR 2025
Thomas M. Cover and Joy A. Thomas. Elements of information theory. John Wiley & Sons, 2006.
Tri Dao. FlashAttention-2:: Faster attention with better parallelism and work partitioning. In Inter-
national Conference on Learning Representations (ICIR), 2024.
Tri Dao, Daniel Y. Fu, Stefano Ermon. Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Process- ing Systems (NeurIPS), 2022.
Jacob Devlin, Ming-Wey Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4188 Minneapolis, Minnesota, June 2019. Association for Computa- 1423. URL htttps::/aclaanholggy////// wrw.//ooo.oo////910
tional Linguistics. doi:
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. International Conference on Learning Representations, 2021. URL. https: ///penreview.net//eeeeeee wwt.:ntt//orummmiiiiiccrrrhoe
James Dugundji. Topology. Allyn and Bacon, 1966.
Ruiyi Fang, Liangjin Wen, Zhao Kang, and Jianzhuang Liu. Structure-presevvve graph represen- tation learning. In 2022 IEEE International Conference on Data Mining (ICDM), pp. 927-932.
IEEE, 2022.
Ruiyi Fang, Bingheng Li, Zhao Kang, Qiuhao Zeng, Ruizhi Pu, Nima Hosseini Dashthayaz, Boyu Wang, and Charles Ling. On the benefits of attribute-driven graph domain adaptation. The Thir- teenth International Conference on Learning Representations, 2025.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In http://pppeneevieeeeee/err/rrree
First Conference on Language Modeling, 2024. URL
LEEskw11VY VY2
Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-effiiient transformers via top-k attention. In Nafise Sadat Moosavi, Iryna Gurevych, Angela Fan, Thomas Wolf, Yufang Hou, Ana Marasović, and Sujith Ravi (eds.), Proceedings of the Second Work- shop on Simple and Efficient Natural Language Processing, pp. 39-52, Virtual, November 2021. Association for Computational Linguistics. doi: 10. 110.663w//20011sssaanpplll. URL bttps://aallattoooooy.rrr/// sus.s
sustainlp-1.5
Jonathan Ho, Nal Kalehbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidi- mensional transformers, 2020. URL https:///operreview
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Emma Bou Hanna, Florian Bressand, Gi- anna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux. Pierre Stock, Sandeep Subramanian, Sophia Yang. Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang. Timothée Lacroix, and William El Sayed. Mixtral of experts. CoRR, abs/2401.04088, 2024. doi: 10.48550/ARXIV.400..00 044088. URL https://doi org/ 10. 48550/ 2401 04088
William B Johnson, Joram Lindenstrauss, and Gideon Schechtman. Extensions of lipschitz maps into banach spaces. Israel Journal of Mathematics, 54(2):129-138, 1986.
Daniel Jurafsky and James H. Martin. Speech and Language Processing: An Introduction Nat- ural Language Processing, Computational Linguistics,, and Speech Recognition with Language Models. Pearson, 3rd edition, 2024. URL wtt.httpss//webbtttfordddd ed//~jurafsky///aaaaaeee slp3/
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In
International Conference on Learning Representations, 2020. URL |htpss: http://///p/eeeeeeeeeeeeee net/ "
12

---

Published as a conference paper at ICLR 2025
Yuzhen Mao, Martin Ester, and Ke Li. Iceformer: Accelerated inference with long- long-sequence trans- formers on CPUs. In The Twelfth International Conference on Learning Representations, 2024. URL https://openrewiww.netttte
Leland McInnes, John Healy, Nathaniel Saul, and Lukas Grossberger. Umap: Uniform manifold approximation and projection. The Journal of Open Source Software, 3(29):861, 2018.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mix-
ture models. In International Conference Learning Representations, 2017. URL https:
//openreview. net.
//openreview.ne//forriiee www.ee/ormmidddyyy2aeeae
OpenAl, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Floren- cia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman. Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Moham- mad Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Semadeet-Shapirg Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brock- man, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichael, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux, Thomas Degry. Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, Tyna Eloundou, David Farhi, Liam Fedus, Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gib- son, Vik Goel, Tarun Gogineni, Gabriel Goh, Rapha Gontijo-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo. Chris Hal- lacy, Jesse Han, Jeff Harris, Yuchen He, Mike Hcaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey, Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga, Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang, Roger Jiang, Haozhun Jin, Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, Lukasz Kaiser, Ali Ka- mali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim. Yongjik Kim, Jan Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Lukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger, Vishal Kuo, Michael Lampe, Ikai Lan, Teddy Lee, Jan Leike, Jade Leung, Daniel Levy, Chak Ming Li, Rachel Lim, Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning. Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer. Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mchta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reicchiro Nakano, Ra- jeev Nayak, Arvind Neelakantan, Richard Ngo. Hyeonwoo Noh, Long Ouyang. Cullen O' Keefe, Jakub Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng, Adam Perelman, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny, Michelle Pokrass, Vitchyr H. Pong, Tolly Powell, Alethea Power, Boris Power, Elizabeth Prochl, Raul Puri, Alec Radford, Jack Rae, Aditya Ramesh, Cameron Raymond, Francis Real, Kendra Rimbach, Carl Ross. Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish Sastry. Heather Schmidt, David Schnurr, John Schulman, Daniel Sel- sam, Kyla Sheppard, Toki Sherbakov, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky, Yang Song, Natalie Staudacher. Felipe Petroski Such, Natalie Summers. Ilya Sutskever, Jie Tang, Nikolas Tezak, Madeleine B. Thompson, Phil Tillet, Amin Tootoonchian. Elizabeth Tseng, Pre- ston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vi- jayvergiya, Chelsea Voss, Carroll Wainwright, Justin Jay Wang, Alvin Wang, Ben Wang. Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng, Lilian Weng, Matt Wiethoff, Dave Willner, Clemens Winter, Samuel Wolrich, Hannah Wong, Lauren Work- man, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph. Gpt-4 technical report, 2024. URL https://arxiv. wwwlvvorggggee wttp/b//////../../// 08774
13

---

Published as a conference paper at ICLR 2025
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli Scaling neural machine translation. In Ondřej Bojar, Rajen Chatterjee, Christian Federmann, Mark Fishel, Yvette Graham, Barry Haddow, Matthias Huck, Antonio Jimeno Yepes, Philipp Koehn, Christof Monz, Matteo Negri, Aurélie Névéol, Mariana Neves, Matt Post, Lucia Specia, Marco Turchi, and Karin Verspoor (eds.), Proceedings of the Third Conference on Machine Translation: Research Papers. pp. 1-9, Brussels, Belgium, October 2018. Association for Computational Linguistics. doi: 10.18653/v1/ W18-6301. URL httttttsss//aaaattttttat
wrg/w18 6301
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chil- amkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An impera- tive style, high-performance deep learning library. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' d'Alché-Bus Buc Emily B. Fox, and Roman Garnett (eds.), Ad- vances Neural Information Processing Systems 32: Annual Conference on Neural In- formation Processing Systems 2019, NeurlPS 2019, December 8-14, 20/9, Vancouver, BC, Canada, pp. 8024 8035, 2019. URL httt:ss://poocceeinngeeeooo neurips. wttps/ccpppppe//////e/e Ddbca288fee7f922f 01 html
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.). Proceedings of the 2022 Conference on Empirical Methods in Natural Language Process- ing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pp. 7025-7041. Association for Computational Linguistics, 2022a. doi: URL https://doi org//0/ 18653/v1/2022 wwt.::222222mmmmmmmemmmmmm manlp-main...
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Ling- peng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In International Conference on Learning Representations, 2022b. URL http:::/ppearrvvie..eeeffooorom B18C0rx2Up4
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners, 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and llya Sutskever. Zero-shot text-to-image generation. In Marina Mcila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 8821-8831. PMLR, 2021. URL http:hhttp/ppcceeeeeeee
mlr.ppres/vv / 39t/am/eh21a....mm htmll
Shai Shalev-Shwartz and Shai Ben-David. Understanding Machine Learning: From Theory to Algorithms. Cambridge University Press, 2014.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang. and Jiang Bian. A study on relu and softmax in transformer. CoRR, 0bs/2300.0064.1. 2023. doi: 10.4555/ARXX12220206661 URL https://doi. 0rg/10.0// 18550/arXivv... wttp:///wrivv..2.....m 26461
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse Sinkhorn attention. In Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9438-9447. PMLR, 13-18 Jul 2020.
mlr wwwpp/ressvvveee http////9/ttyyaaaaa 20ml html.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena A benchmark for efficient transformers. In International Conference on Learning Representations, 2021. URL https: 7/openreview.net//ee wttp://neeet/orrmm//////eeeeeeee
Yi Tay, Mostafa Dehghani. Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Compur. Surv., 55(6), December 2022. ISSN 0360-0300. doi: 10.1145/353081. URL https://doi.. org/10. 1145/3539811
14

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers & distillation through attention. In Ma- rina Mcila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 10347-10357. PMLR, 18-24 Jul 2021. httths:s::/ppooe
mlr wttss////t/t/tuuroo2222
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URLhhtpps//prrocediing 3f5ee243547deee1 hbtp///eeeeeebbbcccccccoacaaaaae
Sinong Wang, Belinda 7. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04766, 2020.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty. Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystrômformer: A nyström- based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, 2021.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, and Amr Ahmed. Big bird: Trans- formers for longer sequences. 13:17283-17299 2020. URL https::/ppooceeddngggt hetreedddinnnnniniis..e 31
cc/paper /illa//pppp//0200//faa///2222222 2aaaaa/22888.0111111121111211112 Pab .dff
Qiuhao Zeng, Wei Wang, Fan Zhou, Charles Ling, and Boyu Wang. Foresee what you will learn:
data augmentation for domain generalization non-stationary environment. In Proceedings of the AAAI conference on artificial intelligence, volume 37, pp. 11147-11155, 2023.
QIUHAO Zeng. Long-Kai Huang, Qi Chen, Charles X Ling, and Boyu Wang. Towards understand- ing evolving patterns sequential data. Advances in Neural Information Processing Systems, 37: 132747-132773, 2024a.
QIUHAO Zeng, Changjian Shui, Long-Kat Kai Huang, Peng Liu, Xi Chen, Charles Ling, and Boyu Wang. Latent trajectory learning for limited timestamps under distribution shift over time. In The Twelfth International Conference on Learning Representations, 2024b. URL https: www.:n//oooooddddiiinoiooo
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M. Susskind. Stabilizing transformer training by preventing atten- tion entropy collapse. In Andreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 40770-40800 PMLR, 23-29 Jul 2023. URL https: //tppp:::/p/ http::///occceen...... www.eeeeiiiin.........eeemeeme
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re.
The hedgehog & the porcu- pine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference on Learning Representations, 2024. URL
Zhenhai Zhu and Radu Soricut. H H-transformer- ID: Fast one-dimensional hierarchical attention for sequences.
Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Pro- ceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Pa- pers), pp. 3801-3815, Online, August 2021. Association for Computational Linguistics. doi: 10. 1/255/ wttp:///222..a1oonno..oo acll-long.294, URL https::/aclantthooggy org/ 2021 acl-long. .24....9999)
Shen Zhuoran, Zhang Mingyuan, Zhao Haiyu, Yi Shuai, and Li Hongsheng. Efficient attention: Attention with linear complexities. In 202/ IEEE Winter Conference on Applications of Computer Vision ( (WACV), 2021.
15

---

Published as a conference paper at ICLR 2025
A THEORETICAL ANALYSIS
The recent paper, Reformer, proposed the Shared-QK Transformer that shares the linear projection layer for keys and queries. reduces the number parameters and memory space while not affect- ing the model performance. https//arxiv..rraas//000.... 00445a very simple but efficient technique.
Lemma A.1 UJhnson-Lindensttrun positive integer such that d 2 m Then for any set z of points in R. there exists a map
Lemma). For any 0 < < 1 and any integer m, let d be
TRD - Rdd such that for all 24, 2.g E Ex,
(8) Assumption A.2. Let a E Am-1 be element of the m-dimensional simplex, defined as 1}. Assume that lattn equipped with can achieve
- -->)|<=== =,51² 2 (1,( [(x)))1 (1 (1++))==; xyll²
Am-1 & fa E R" a; 2 2. an optimal learnable similarity function r, where the attention scores are given by a softmax(T(fe(rr))..)))) (x))), such that r is trained to be optimal to have the minimal expected risk: min, |/aat (x, Szia) a)) yll. where hattn denotes the crention-based hypothesis, x is the input, and Sg is the context.
Lemma A.3. Let E Rm, 4h1 : Rm - R and vh2 Rm -> R. Assume that h (a) S h2(a)) Then min, hi (a) S ming (a).
Proof. Let 01 = arg min, (a) and 02 = arg min, h2(a). Then, h (12) S h2(02) due to condition h (a) hg(a)). We also have h, (a,) I (a2) due to condition 0j achieving the minimum of h1, To sum up, we have hi ((())) S h2(02).
Lemma A.4. Shalev-Shwarzz & Ben-David 2014) Let Cl,..., Cr be a collection of subsets of some domain set, X. Let S be a sequence of m points sampled i.d. according to some probability distribution, D, over X. Then,
EsnDm
Σ
P[C.]
me
Proof. From the linearity of expectation, we can rewrite
Es
P(C]|
(A
EC
Next, for each we have
Es [Anss= Ps|C. ns = = "(C.""" P|C.Im
Combining the preceding two equations, get
Es
Σ
P[C.] <
P[C.j--(C.Jmm
F(C.Jm < <mmxxP(CCJe P P[C.Imm
Finally, by elementary calculus, max de ma
concluding the proof.
me
Theorem A.S. Let X Rd VERD and be distribution over for which the conditional probability function, h - +RD, is a a --Lipschit.functionn. function. Let h denote hypothesis, and hattn denote the one- one-layer attention model to aggregate the predictions of sample set S Dm to predict for another i.i.d sample x. Specifically, here assume the same linear map for key mapping fx B, B) de where B the bound of projection of x by f. and assume the value mapping fo be the Bayes optimal rule h E X ær), which is the
and a query mapping fg as f Rd
16

---

Published as a conference paper at ICLR 2025
hypothesis that minimizes Lb(h) over all functions. Then, the expected risk of the regression tasks (p(h) = (a.y)>> (/())))))((())>>>>)))))))))))>>) - yll with as hattn can be upper bounded by
EswDm [LD Lp(h") + 4lcydx Bm 1/(dx+1)
dx
Proof. Es [L (h)] is the root square (RMSE) between the prediction and conditioned on a sampled set S and an additional example (I, 3), such that the label of T (z) is different from y. In other words, additional example, (z, y). It follows that
first sample m examples, S₂ {mm, according Dx, and an
Eg (LD (hatua)) Es,, DDT DD By min h (x;a) - vll
Rs, DD and Dx B Dyjx min ogh (xx))))) vll
< Ess~Dy
am an a By
min
osl ash - 11 + (1)(x)) vll
= Lp(h") + Es. and Dn E Dyw min
(a," ht" (x)],,
where ag is the attention between x; and The inequality follows the Cauchy-Schwan inequality and Lemma A.3
Es x~ x~ Dx Dyjx min
(x.) (x))l S Es xx Dyjx min
a (x;) (x:) (x))l
1=1
llf - f(x))l
S Es, DDp DD BynDyyx D,
1. min Σ
a [ - z||
S Bs Ey
min Σ ai
C m
dx
kis 9
Es, DD DD Byy-
I-min
Cl m dx
The first inequality follows from the Cauchy-Schwarl inequality: the second inequality follows from the L-Lipschitzness of h"; the third inequality follows from the Johnson-Lindenstrauss Lemma, where define in the Johnson-LLindensrraue Lemma as k; =x(x)) and g fzx). It is obvious that
In where C is constant;
dx
min Σ
aall gll ([ 9ll
where (z,(g) is the closest & to q. Thus we have
m
I (9) 91l
ES aw Dy min Σ
a (xT) - (x) ES Dw E EynDyyx
1.
C In
Fix some ( = 2B 2B/T. For some integer T, let r = Tdx and C1,... Cr be the cover of the set X using boxes of length Ç. Namely, for every (E1,..., Ea) [T] there exists set C of the form {k: Vjj ky E [2 - 1)/T, 2BC;/T]}
For each k. q in the same box we have ||k - gll S Bd Otherwise, ||k - all S Vd Therefore,
17

---

Published as a conference paper at ICLR 2025
Es... (()))())) - d|l
Es
U
U Li
C
vdxB+P
U
C
BCVdk
ions
and by combining Lemma A.4 with the trivial bound
PI
U CI &c,ns,
we get that
Ex s. [[]k (2) a1l S ddxB
( ( + +)) me
Since the number of boxes (1/{)ax, it follows that
12xx <-
Es (m) - a/ll S v d
+ +5)
me
Setting Ç = 2m 1/(xx+ +1 and noting that
2dx ç-d + 2dx 2 dx m dx (dxx+)
+ 2m -1/
=m 1/(dx+1)
+2 S 4m -1(dk+()++))
then combining the preceding with previous results, we obtain
Ess (hattt))) s ID(hh) + 4leydK Bm +1(dx++)
C In
dx
B AN ILLUSTRATIVE EXAMPLE OF EFFICIENT TOP-k SELECTION VIA Z-ORDER CURVE CHUNKING
In ZETA, the chunking process is crucial step for efficient key-query matching in large-scale attention mechanisms. The process begins by projecting the high-dimensional keys into a one- dimensional space using Z-order curves, which preserve spatial locality. These projections create linear representation of the keys, as shown in the second row of the figure.
After projection, the keys radix sorted in O(N) into ascending order of their Z-order integer val- ues, enabling efficient binary searching in O(N) for the most relevant keys in the one-dimensional space. Sorting ensures that keys that spatially close in the original multi-dimensional space remain close in the projected space, making the retrieval process computationally efficient.
Next, the sorted keys are partitioned into chunks. Each chunk contains fixed number of keys, and queries are matched with the keys within their respective chunks. This chunk-based structure facilitates parallel processing, where each query can efficiently search for its top-k nearest keys within a local subset of the sorted key space. For instance, in the figure, the 5th query retrieves its top-4 keys from the chunks containing its most relevant keys.
This chunking approach significantly reduces the computational overhead compared to searching the entire key space for each query, while leveraging the locality-preserving properties of Z-order curves. The result is a scalable and efficient mechanism for top-k selection in ZETA, ensuring both speed and accuracy for attention-bascl tasks.
The pseudo-code in Algorithm I outlines the ZETA Top-k Attention mechanism, which combines Z-order curve projections with chunk-based sorting to efficiently identify and retrieve the top-k nearest neighbors while maintaining causal constraints. It provides structured approach to reduce computational overhead by limiting the search space to relevant chunks, ensuring both efficiency and adherence to sequence masking requirements.
18

---

Published as a conference paper at ICLR 2025
Keys:
Z-Order Curve Projections:
2
5
8
9
11
12
Sorted Keys:
11
12
Chunk Sorted Keys:
Fürst Chunk (1-4):
5
2
5-th Query
Second Chunk (1-8)-
3
2
8
7
1
6
4
5
Top-4 Keys
Third Chunk (1-12):
11
12
Figure 4: Illustration of the chunking process in ZETA: Keys projected into one-dimensional space using Z-order curves, sorted, and partitioned into chunks for efficient retrieval of top-k keys for each query.
C EXPERIMENTS DETAILS
The ZETA model configuration generally involves setting the number of chunks values such as 4, 8, 16, 32 depending on the sequence length, which provides a flexible way to handle different input scales effectively. This chunking strategy facilitates parallelism in processing, allowing for ef- ficient memory use and computational speedup during attention operations. The hidden dimension, dy, is typically set to 256 or 512 with attention heads when working with LRA datasets. However, for larger and more complex datasets such as WIKITEXT- 1 03, the hidden dimension is increased to dy = 768 with 12 attention heads to ensure that the model has sufficient capacity to learn intricate long-range dependencies effectively. Additionally, the dimensions of keys and queries are kept sig- nificantly lower at dx = dQ 3. which aids reducing the computational burden and mitigates the "curse of dimensionality" while still preserving enough information for efficient attention com- putation. This choice of dimensions strikes a balance between model efficiency and effectiveness, making ZETA well-suited for long-sequence modeling tasks.
19

---

Published as a conference paper at ICLR 2025
Algorithm 1: ZETA Top-k Attention Using Z-order Curves and Chunking
Data: Keys of size M d, sequence length N, chunk size M, query index i, top-k value Result: Top-k nearest keys for each query
Step 1: Z-order Curve Projection;
for each key vector in do
Project each d-dimensional key vector into a one-dimensional representation using Z-order curve projection:
end
Step 2:
Sorting of Keys;
Sort the projected one-dimensional keys in ascending order;
Step 3: Chunk Division;
Divide the sorted keys into multiple chunks based on the chunk size M;
Step 4: Chunk-wise Causal Masking;
for each query at position do
Determine chunk index m Search only in the first m chunks for the top-k keys to attend to;
J 1
Exclude keys from positions >> m x M;
end
Step 5: Nearest Neighbor Search; for each query do
Perform nearest neighbor search within the selected chunks to find top-k tokens; Collect the indices of the nearest neighbors I, while maintaining causal constraints; end
return Top-k nearest keys for each query
D 1/O-AWARE ZETA OPTIMIZED WITH TRITON
Our Triton implementation of ZETA focuses on improving the efficiency of sparse attention through customized kernel programming. We leverage Triton create specialized GPU kernels for top-k sparse attention. The sparse topk attention kernel and its corresponding backward pass kernel ttteetttttttttiioonnooeo woo.baacwwrrddkkk d_ennel are implemented using the Triton JIT (Just-In-Time compiler. This approach allows for significant speedup by optimizing memory access patterns and reducing I/O overhead during the computation. The kernel is tuned to different configurations, like "block size" and warps",, which directly influence how GPU resources Especially, we compute the mean vector of history tokens in the block of the current kernel, instead of computing global mean vectors, which effectively reduce overheads. he @tri soriton, autotune decorator is used to evaluate multiple kernel configurations for optimal performance, making sure that GPU resources well-utilized depending on sequence length and other parameters.
allocated.
One key challenge addressed is efficient indexing for large tensors in the backward and forward computations. In Triton, indexing is achieved via program IDs, t1. program.id() that are used to identify which part of the workload is being computed by each block or thread, ensuring that paral- lelism is effectively exploited. The Triton kernel computes Cauchy Softmax for top-k KV pairs for each query, employing a specialized kernel to access only the most relevant & keys during the atten- tion process. This reduces the computational complexity compared to full attention mechanism, and Triton's low-level bit manipulation operations retrieval.
and tl. store) are used for fast data
The sparse Attention Mechanism is computed by considering only the top-k keys per query, which significantly reduces the computational load. The sparse_topk_at wewttttettionneeeeeeeee involves effi- cient gathering of keys and values based top-k indices. The indices pre-computed in sorted order, which enables efficient retrieval without scanning the entire key space.
The Triton kernel also includes custom backward functions to handle the gradient flow. The back- ward kernel sparse topk attention backward kernel computes gradients for each parameter in- volved in the sparse attention operation, including 4, k, and the learable parameter yy Triton's
20

---

Published as a conference paper at ICLR 2025
tl.atomic_ade is used to accumulate gradients, ensuring that all updates to shared memory locations are synchronized.
During the forward pass, intermediate values such as the Euclidean distances and normalization constants are stored. These values are reused in the backward pass, which reduces redundant com- putations and accelerates the training process.
By using Triton, we managed to reduce the I/O overhead that traditional PyTorch operations faced, especially during backward computations. Furthermore, used fused kernels to mitigate the over- head associated with multiple indexing operations. This fusion helps in reducing the number of kernel launches, which translates reduced latency and faster execution, Triton allows more control over memory coalescing and efficient block-wise operations.
Overall, the Triton-based implementation in ZETA allows for a more scalable sparse attention mech- anism that retains the benefits of locality preservation through Z-order Curves while significantly reducing computational and I/O bottlenecks. This makes the ZETA attention more suitable for long- sequence tasks where traditional transformers are too resource-intensee
E
GRADIENT DERIVATION FOR THE BACKWARD PASS OF TRITON ZETA
E. UPDATED ATTENTION MECHANISM WITH SPARSE ATTENTION
We introduce sparse attention by computing values and keys from an index set I₃ of top-k tokens specific to each query q,. The unnormalized attention scores are defined as:
1 for j 1; Sij Dij +E'
(9)
where:
Dij = ||q: kyy
== 22 is trainable scalar parameter
Define:
dij Dig +
(10)
The steps of the attention mechanism are:
1. Compute Pairwise Distances:
Dij = |qq k2l?: for E I
(11)
2. Compute Unnormalized Attention Scores:
1
Sij 5ij"
for j E I
(12)
3. Normalize Attention Weights:
Z, = Σ Sij
(13)
jET,
Sij
Aij Z,
for E I
(14)
4. Compute Output:
0, = Σ A
(15)
jEL:
Our goal is to compute the gradients of the output o with respect to q, kj. Vj, and E, considering the sparse attention.
21

---

Published as a conference paper at ICLR 2025
E.2 GRADIENTS WITH RESPECT TO Q, K, V, AND
The gradient of the output o, with respect to Vj is:
{ Aijy 0v; 0,
do:
ifj I; otherwise
(16)
Derivative of Dij with respect to q, and ky:
0D; 2(q: - kj), for
(17)
0q:
0Dj -2( - -ky), j 4
(18)
0k, 10
otherwise
Derivative of Siy:
First, compute the derivative of Sij with respect to diy:
aSg 0ō.
for j E 4:
(19)
0
Compute the derivative of dy with respect to Dij and
06.j 1, for
(20)
0Di
0dg 1, for E T
(21)
0=
Now, compute the derivative of Siy with respect to Dis and
aSg
1
for j E I
(22)
aD;
8
aS.j
for j E 7
(23)
dc
82
Derivative of S.y with respect to q, and k,:
as.
2(q: kj)
for j E T
(24)
dq:
aSg
2(q:
k,)
if 1;
87
(25)
0k,
0,
otherwise
Derivative of Z, with respect to S.y:
az;
1, for j E I
(26)
aS
Derivative of A, with respect to Sty²
0 AAi
Z, Sij
for j E 1
(27)
as
Z
22

---

Published as a conference paper at ICLR 2025
Derivative of Ay with respect to Sa for # j:
Sky z
OA, a5a
if I otherwise
(28)
Gradient of o, with respect to Ss:
0o:
0Ai
Σ
(29)
aS
aSa
jEL
VI O,
for E I
(30)
Zi
Gradient with Respect q:
0o;
0o; aSa
(31)
Ōq.
dSa
Ōq.
IET.
Vi O2
(q: ki)
= -
(32)
Z
82
1E1.
Gradient with Respect to ky:
do;
2 12
(q--k,)
if E
8,
(33)
0k,
0,
otherwise
The total gradient with respect to kj is:
aL = Σ
ÖL
V, - O;
(q: k,)
(34)
ok
0o:
Z:
82
Gradient with Respect to 8:
0o:
do:
05a
(35)
0=
aSa
TET,
(36)
2, 82
1E1
The total gradient with respect to is:
aT.
aL
Vi O, 1
Σ
(37)
Ō=
0o: ich.
Z,
8%
E.3 SUMMARY OF GRADIENT COMPUTATIONS
1. Compute Dij, dejo Sij, Zig Aij, and o;:
23

---

Published as a conference paper at ICLR 2025
Dú ||q kyll. for j I
(38)
5Ú Dij te, for j E I
(39)
1
S dij
for j E 11
(40)
Z, Σ
Si
(41)
jEl:
Sij
Ao Zi
for j 4
(42)
0, = Σ
A
(43)
jel:
2. Compute Gradient with Respect to V:
ØL
Σ
OL
Aij
(44)
0v,
Do:
3. Compute Gradient with Respect Q:
aL 0L
V, 0:
(q: ky)
(45)
dq.
Z,
8
jEL,
4. Compute Gradient with Respect to K:
aL Σ
2 OL
Vj
(q. kj)
(46)
8o:
Z₁
6
5. Compute Gradient with Respect to
aL
Σ
ØL
Vj O,
(47)
8E
0o:
Z,
jEL,
E.4 IMPLEMENTATION NOTES
Sparse Attention Index Set 7,:
is the set of indices that query qi attends to, which is collected from Z-order Curve
projected sequences.
The attention computations and gradient updates are performed only over E T.
Numerical Stability:
-The addition of ensures that &g 0f preventing division by zero.
Ensure that remains positive during training, possibly by parameterizing cxp(#) where 0 is unconstrained.
Efficient Computation:
Utilize sparse matrix representations handle the index sets efficiently.
Use vectorized operations and appropriate masking perform computations only over the valid indices.
F
LIMITATIONS
Given that our method a top-k attention mechanism, there are some shared limitations between our method and that of prior work that deals with attention, such as there still potentially being higher chances to ignore attention to important information (with low attention scores) than full attention methods given the use of only the top-k tokens.
24

---

Published as a conference paper at ICLR 2025
G ADDITIONAL EXPERIMENTS
G.1 ABLATION ON ATTENTION PERFORMANCE WITH VARYING dx ON LRA
We expand on the ablation study presented in Figure 2(b), focusing on the ListOps and Image tasks from the Long Range Arena (LRA) benchmark. Specifically, we examine the impact of varying the dimensionality of the keys and queries (dx) on attention performance. The results are summarized in Table 5
Our experimental findings indicate that the performance remains relatively consistent for dæ 3, whereas noticeable decline is observed for dx < 3. This supports our hypothesis that, unlike value vectors, the keys and queries predominantly encode positional information rather than intricate se- mantic features. As such, reducing dk to small value allows effectively lower computational costs without incurring a significant loss performance. This insight guided our selection of dR values in subsequent experiments.
dx
2
3
32
64
128
ListOps
31.04
34.79
36.06
36.
36.32
36.37
Image
37.6
40.23
42.64
42.72
42.51
42.44
Table 5: Performance metrics for different values of dk.
G.2 PERFORMANCE OF ZETA USING DIFFERENT SIMILARITY METRIC
We utilize Euclidean distance for k-nearest neighbor (k-NN) searches to identify the top-k attended tokens, as it is particularly well-suited for this purpose. In contrast, dot- dot-product similarity cannot be directly employed for k-NN searches without normalization, as highlighted in (Mao et al.) (024). Our experimental results also indicate that normalized dot- dot-product similarity performs worse than Euclidean distance, as demonstrated in additional MQAR experiments below.
Specifically, Euclidean distance-basse methods, such Negative Euclidean with traditional softmax and Cauchy Softmax, consistently outperform dot-product- based methods for top-k attention when using small dimensionality (dg < 4), which is the setting adopted in ZETA.
dy
3
Negative Euclidean
64.6
99.4
99.4
99.3
Inverse Euclidean
22.9
81.3
99.6
99.9
Cauchy Softmax
74.5
99.6
99.5
99.2
Normalized Dot Prod
62.6
92.6
99.3
99.1
Table 6: Performance using various different similarity metrics
25


Reviews:
[{"review_id": "sRSXtKd7wp", "reviewer": "Reviewer_36fb", "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.", "questions": "Why choose ASSOCIATIVE RECALL in section 4.2 ?\nTesting 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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "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.\nQuality:  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.\nClarity : The paper is well-organized and accessible, with clear explanations of techniques like Adaptive Cauchy-Softmax and Z-order projection, aiding reader comprehension.\nSignificance : 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:\nThe 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.\n\nUnclear Transferability of \\gamma in Adaptive Cauchy-Softmax:\nThe 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?\nLack of Formal Justification for Euclidean Distance Claim:\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UUJXBCbWc4", "reviewer": "Reviewer_pVHs", "summary": "Paper is well written, with solid experimental work and theoretical analysis.\nAuthors 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.", "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'\nCould 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?\n\n\nCould you please share computation efficiency measurements (for example average training step time) for baseline transformer and ZETA?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "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.\nOn top of this paper this paper investigate dimensionality of keys and queries the trade-off between the curse of dimensionality\nand 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1YMG0sPQdI", "reviewer": "Reviewer_WjdT", "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.\n\nThe 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.\n\nKNN 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.\n\nAssuming 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.\n\nAs 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.", "questions": "1. Can you include more details on the chunking part. Also explain where efficiency is coming from.\n2. 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.\n3. Can you discuss the limitations of your method.\n4. Can you provide results on Path-X task in LRA?\n\n~~I am willing to change my score if the authors can address these concerns.~~  \n**[Post author rebuttal]**: The authors have largely addressed my concerns and I have increased my score.\n\n\nTypos:\n- Equation 6: \"(q-Ki)\" should be \"||q-Ki||\"\n- Line 232: Period missing after Lemma 3.1", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is interesting, and the trick of projecting to 1D for query-key similarity computation is intuitive.\n\n2. Significant and comprehensive theoretical justification is provided for dimensionality reduction via random projection and is well-supported via detailed proofs.\n\n3. 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.\n\n4. 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.\n\n5. 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.\n\n2. 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.\n\n3. 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.\n\n4. Benchmarking of speedup / memory usage is not provided.\n\n5. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DaWHdC93gn", "reviewer": "Reviewer_xJi4", "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.", "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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* The paper is well-written and the method is believable.\n* 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.\n  - In section 3.2.1, line 202, the authors mentioned: \"a lower $d_K$ loses locality between tokens, which is crucial for efficient query\".\n  - 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.\"\n  - There seems to be a contradiction between these two statements, so the effect of $d_K$ needs clarification. \n* Can the authors clarify the binary expansion? What is the precision of the original d-dimensional vector?\n  - Eq. (4) introduced Z-order curves, which rearrange a d-dimensional vector into a one-dimensional binary representation.\n  - 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. \n  - Will this result in an overflow? If not, how did you set the precision of the original d-dimensional vector?\n* Can the authors provide a comparison of the time and space usages for ZETA model and the attention-based models?\n  - 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)