Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling
Authors: Xiang Hu
Zhihao Teng
Jun Zhao
Wei Wu
Kewei Tu
Venue: ICML 2025 poster (2025)
Keywords: long-context language modeling
retrieval-based LM
length generalization
attention mechanism
Primary area: deep_learning->attention_mechanisms
TL;DR: A novel attention mechanism with length generalization ability achieves perfect accuracy in passkey retrieval over 16M context length.

Abstract: Despite the success of Transformers, handling longer contexts remains challenging due to the limited length generalization and quadratic complexity of self-attention, which often requires post-training with a larger attention window, significantly increasing computational and memory costs. In this paper, we propose a novel attention mechanism based on dynamic context, Grouped Cross Attention (GCA), which can generalize to 1000 $\times$ the pre-training context length while maintaining the ability to access distant information with a constant attention window size. For a given input sequence, we split it into chunks and use each chunk to retrieve top-$k$ relevant past chunks for subsequent text generation. 
Specifically, unlike most previous works that use an off-the-shelf retriever, our key innovation allows the retriever to learn how to retrieve past chunks that better minimize the auto-regressive loss of subsequent tokens in an end-to-end manner, which adapts better to causal language models.
Such a mechanism accommodates retrieved chunks with a fixed-size attention window to achieve long-range information access, significantly reducing computational and memory costs during training and inference. 
Experiments show that GCA-based models achieve near-perfect accuracy in passkey retrieval for 16M context lengths, which is $1000 \times$ the training length.

Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling
Xiang Hu Zhihao Teng² 2 Jun Zhao
Wei Wu
Kewei Tu 2
Abstract
Touvron et al., 2023; Dubey et al., 2024). While Transform- ers excel in representational power, they still struggle to process longer contexts as expanding the attention window beyond the pre-trained size may disrupt attention distribu- tion (Wang et al., 2024). Some methods use a sliding win- dow mechanism (Xiao et al., 2024) with a fixed attention field to achieve extrapolation, but they fail to capture long- distance context dependencies outside the attention window. Thus, a main challenge for the attention mechanism lies in how to achieve length generalization but still maintain random access to long-range information. To mitigate the issue, most long-range LMs (Reid et al., 2024) resort to expanding the attention window (Liu et al., 2023) during post-training, significantly increasing training and inference costs due to the quadratic complexity of self-attention.
Despite the success of Transformers, handling longer contexts remains challenging due to the limited length generalization and quadratic com- plexity of self-attention. Transformers often re- quire post-training with larger attention win- dow, significantly increasing computational and memory costs. In this paper, we propose a novel attention mechanism based on dynamic context, Grouped Cross Attention (GCA), which can gen- eralize to 1000 x the pre-training context length while maintaining the ability to access distant in- formation with a constant attention window size. For given input sequence, we split into chunks and use each chunk to retrieve top-k relevant past chunks for subsequent text generation. Specifi- cally, unlike most previous works that use an off- the-shelf retriever, our key innovation allows the retriever to learn how to retrieve past chunks that better minimize the auto-regressive loss of sub- sequent tokens in an end-to-end manner. Such a mechanism accommodates retrieved chunks with fixed-size attention window to achieve long- range information access, significantly reducing computational and memory costs during training and inference. Experiments show that GCA-based models achieve near- pera-ptefect accuracy passkey retrieval for 16M context lengths, which is 1000 the training length.
In this work, inspired by retrieval- based language models (RLM) (Asai et al., 2023), we introduce an efficient atten- tion mechanism that achieves both length generalization and long-range information accessibility. Typically, RLMs (Ru- bin & Berant, 2024) divide an input sequence into chunks, retrieve relevant ones from the history for the current input, and then integrate the retrieved chunks into the decoder to predict subsequent tokens. By choosing top-k chunks as "dynamic context", RLMs can maintain the ability to access distant information without expanding the attention win- dow. However, most RLMs (Lewis et al., 2020; Borgeaud et al., 2022) rely on separately pre-trained retrievers, whose retrieved chunks are not necessarily helpful to the causal LMs. Although a straightforward approach is training the retriever end-to-end with the auto-regressive loss, it is rarely explored. The main difficulty is that the relevance scores computed by the retriever do not participate in the next token prediction, thus cannot receive gradient backpropagation from the regressiee loss. To overeome the difficulty, we propose Grouped Pross-AAtention (GCA), a novel attention mechanism which integrates chunk-to-chunk retrieval an inherent ability, thus the retriever can learn to retrieve past chunks that most effectively reduce the auto-regressive loss of subsequent tokens, which we refer to as causal retrieval.
1. Introduction
Transformers (Vaswani al., 2017), serving as the backbone of large language models (LLM), have demonstrated excep- tional performance across a wide range of natural language processing tasks (Brown et al., 2020; Achiam et al., 2023;
'Ant Group ShanggaaTTec University FFudan University. Cor- respondence to: Wei Wu <congyue wwe.yuewwwwaaae wwwantgroup.coms Kewei Tu @tukw@shtt wwwnghaitech.cdu.......
Specifically, GCA enables the relevance scores to participate in the next token prediction in a differentiable way and can be understood chunk-wise analogy of token-wise self-
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Lenral-Generrlizaalle Attention via Causal Retrieval
attention. In self-attention, considering the next token pre- diction in causal Transformers, self-attention scores could be viewed as the relevance scores of the current token to past tokens. These scores serve as weights to fuse information gathered from past tokens to predict the next token. Analo- gously, we divide the input sequence into chunks and use the relevance scores between the current and past chunks as weights fuse past information for the next chunk predic- tion. A detailed comparison between previous works and GCA is depicted in Figure 1. By appending GCA after self- attention in Transformer layers, we introduce Differentiable Retrieval-basee Transformers (DRT), enabling pre-trate tr-ininiing from scratch with context lengths up to 64K. To make pre- training efficient, we sample top-k past chunks according to the relevance scores for each chunk to perform GCA, along with fixed-size sliding window self-attention (Child al., 2019), achieving linear complexity for the entire input se- quence's s attention operations. During inference, we offload hidden states of past chunks to CPU memory and reload them when retrieved. It introduces additional memory-swap but largely reduces memory footprint.
2. Related works
Relation to RPT & Landmark Attention. There are two long-range LMs closely related to ours. One of them is Retrieval-Pretraines Transformer (RPT) (Rubin & Berant, 2024). The key difference between DRT and RPT is the training approach of the retriever. During data sata-preparation, for cach chunk, RPT picks relevant past chunks by using BM25 (Robertson & Zaragoza, 2009), concatenates them with the current chunk, and evaluates them by a reference LMM like Pythia 1.4B (Biderman et al., 2023). The past chunks that increase the probability of the next chunk are identified as 'gold chunks' to train RPT's retriever. This process requires that the reference LM be pre-trained on the target dataset beforehand, significantly limiting scala- bility and flexibility. In contrast, GCA integrate retrieval into the attention mechanism thus can be pre-trained jointly end-to-end without reliance on any external data or mod- els. Landmark Attention (LA) (Mohtashami & Jaggi, 2023) is another close work. LA is pre-trained with short con- texts but capable of handling long contexts during inference. It addresses long-range language modeling by modifying self-attention KV Cache During inference, each token, at each layer, selects top-k chunks based on token-to-chunk attention scores and appends their key and value vectors to the current KV cache of self-attention. The token-to-chunk attention scores are trained in an end-to-end manner with grouped softmax technique. However, it has to perform top-k chunk selection per token, per layer, which incurs sigmificant extra costs during inference. Moreover, it fails to extrapolate on longer context length. Our method combines the chunk-retrieval and grouped softmax ideas, resolving the aforementioned issues while balancing training efficiency and inference performance.
To fairly compare GCA with other attention mechanisms, we pre-train all models from scratch and evaluate them on tasks such as long-range language modeling, summariza- tion, and the demonstrate that DRT significantly outperforms all base- lines with comparable tre-training costs and much lower inference costs. Notably, in the NIAH test, DRT maintains perfect accuracy on inputs up to 16M tokens. More inter- estingly, case studies on the arXiv-math dataset validate the causal retrieval ability in DRT, which retrieves lemmas, variants, or functions defined distantly but used in the next chunk. These findings suggest that GCA has the potential to be a fundamental component in long-range LMs. Overall, our main contributions are:
(NIAH) tests. The results
Attention & Length Generalization. The length general- ization issue of Transformers stems from the self-attention mechanism, which is usually pre-trained with in-domain position encoding and window size. Consequently, directly inputting longer sequences can lead to out-of-domain is- sues. Some works mitigate this issue by positional interpo- lation (bloc97, 2023; Chen et al., 2023; Peng et al., 2024). However, these solutions still generalize only to limited extent. Even without position encoding, expanding the at- tention window during inference still fails to extrapolate due to distraction of attention (Wang et al., 2024). To ad- dress the discrepancy in attention window between train- ing and inference, some methods (Child et al., 2019: Xiao et al., 2024) employ sliding window attention. However, this sacrifices the model's ability capture long-range depen- dencies. Furthermore, even with a fixed attention window, Landmark Attention still encounters length generalization issues. Therefore, developing an attention mechanism that can generalize to various lengths while maintaining long- range information accessibility remains challenging.
1. We propose novel attention mechanism called Grouped Cross-Attention (GCA) and its hardware-aware imple- mentation, which simultaneously achieves length gener- alization and long-range random-access.!
2. Building upon GCA, we introduce Differentiable Retrieval-based Transformers (DRT), which is fast and memory-effieien in both pre-training and inference long contexts.
3. We conduct comprehensive experiments and demon- strate the promising results of GCA. To the best of our knowledge, GCA is the first attention mechanism that can achieve perfect passkey retrieval with 16M context length, 1000 x the pre-training length.
'The code is released https: http://ggtttttcco wontoooggoootttttooooooooo--e
2

---

Lenett-Generrlizable Attention via Causal Retrieval
BP
relevance scores:
next token prediction
Top-k
retrieve
BP X
G1
chunk selection
C
0+1
Oc+1
group,
groupz
o
Qenkki PeocK
QutKii 1. softmax
Vi 2 0
PeorKy softmax
vi te o o
sof tmax
softmax
ce+1
Ge+1
G2+1
(1) Chunked Cross-Attention
(b) Grouped Cross-Attention
Figure 1: Comparing previous works with GCA. Consider current chunk e with its past chunk relevance scores rř, where k E (1, (1.....t. 1}, r and r/ are the top two. In this example, each chunk contains tokens, whose query, key, and value vectors are denoted which as Q; V. V. (a) In previous work, information from retrieved chunks is fused into LM decoders via Chunked Pross-AAtentionn relevance scores are merely used for chunk selection. Thus, the loss can not back-propagate to the scores. (b) GCA separately applies with the two chunks, yielding intermediate outputs O²+1 and O/+11 The softmaxed relevance scores serve weights fuse these intermediates into LM decoders and thus can receive hack propaaationnn from the loss.
Long-Range Language Modeling. Various methods have been proposed improve long-range language modeling. A straightforward approach is extending the attention window and conducting post-training 2024; Reid et al., 2024; Qwen et al., 2025). However, it comes with significant additional training cost and still has length generalization issues. One line of research is intro- ducing memorization to Transformers via recurrence. Many works (Dai al., 2019; Burtsev & Sapunov, 2020; Martins al., 2022; Hutchins et al., 2022; Munkhdalai et al., 2024) compress past information into fixed-sized vectors. How- ever, these methods often sacrifice the flexibility to attend to arbitrary past tokens. Mcanwhile, other works focus on maintaining random-access flexibility of attention. Memo- rizing Transformers (Wu et al., 2022) appends retrieved past keys and values to the current attention segment via k-NN search, but they do not proppproe gradients to them. CEPE (Yen et al., 2024) parallelly encode long-conext by chunks and fuses them into the decoder via cross-attention. During the training process, the decoder parameters are fixed, and only the encoder is adjusted. Recently, state- space models (Gu & Dao, 2023; Dao & Gu, 2024), RNN models (Beck al., 2024) and their variants (Nunez al., 2024) provide new architecture alternatives, with compa- rable performance to Transformers but much lower cost for inference. A notable distinction of our work is able to achieve perfect past information retrieval for input length 1000 x the pre-training length.
Attention (CCA) after self-attention to fuse information from retrieved chunks, which a retriever is merely used to pick relevant chunks. Our approach makes the retriever learnable by replacing CCA with GCA, which represents the key innovation of this work. The novelty of GCA lies in using relevance scores to fuse information from retrieved chunks for LM decoders, enabling the retriever to adaptively learn to select the best past chunks for predicting subsequent tokens, guided by the auto auto-regressive loss. This section details the model architecture, training, and inference.
longer context (Bai et al.,
3.1. Model Architecture
DRT is composed of N Transformer- layers. Similar to RETRO (Borgeaud et al., 2022), the input sequence of DRT is equally divided into chunks. Formally, given sequence x = [Ir 1 I2, xr. with L tokens, we divide the sequence into 1 chunks, where S is the chunk size, denoted as (ca, C22 .... CL/S), where x; E C/t/s] Similar to Landmark Attention, we insert a special token LMK at the end each chunk, which summarizes the preceding content via self-attention.
Forward Pass. Figure 2 illustrates the forward pass of a token DRT. DRT layers are bifurcated into upper and lower sections like in RPT (Rubin Berant, 2024). The key differences are the introduction of GCA and the further division of the upper layers into G' groups, enabling learning to retrieve the fly and adaptive multiple retrievals. The lower layers comprise standard Transformer layers while each upper layer has an additional GCA module after self- attention. In the forward pass, the chunk hidden states output by the lower layers, besides being fed to the upper layers, are also fed into a bi-directional Transformer encoder,
3. Methodology
A typical architecture of RLMs (Borgeaud al., 2022; Yen al., 2024; Rubin & Berant, 2024) appends Chunked Cross-
3

---

Length-GGnneralizabi Attention via Causal Retrieval
GCA at the 1-th layer
1-th layer
Upper Decoder X21
XG
top-k chunk rep.
softmax
N
C
CAC C,
20
W1. (0.13)
3
CA( C..
20
1.3
Wa,6 (0.29);
20
Feed Forward
4
CA Co, ....
20
W4.6 (0.09);
fusion
Grouped CA
Cs
CAC C.. 20
Wa.s (0.49)
Self-Attn.
Retrieve
relevance scores
4
=288
2.8
Lower Decoder
C2
12
=112
Top
3.6
Feed Forward
2.4
Encoder
4
=41
4.1
Self-Attn.
Retrieve
Inputs:
X2
LMK
X4
X6.
LMX
X7
Xg
LMK
X10
X11
X12
X13
X15
X,6
X17
X18
LMK
X19
X
Figure 2: The illustration shows how the current chunk es retrieves past chunks for better next token prediction of 22e0 in the next chunk Cr. The landmarks's hidden representation is used to compute relevance scores with past chunks, selecting the top four. The hidden states (CA) with tokens in each separate chunk. The chunk wise CA outputs, O2n. are
of Izu at the (-th layer, h20. perform fused via weighted sum, whose weights are softmaxed relevance scores.
which further contextualizes the representations with inner- chunk positional encoding, yielding C₁ R 5xd and Ię E Idd shared across all upper layers, where d is the hidden state dimension. At the g-th upper decoder group, chunk c retrieves the top-k relevant past chunks for its next chunk:
and is defined as follows:
N
N
8(1)
[11
2
2G
In GCA, perform Cross-Attentinn (CA) separately and fuse results via relevance scores:
rg. (W)) Wel Vd
=op-kp-((>>""" r0-t-1))
0 CA(H Cs. Cx), k C
(xp()
Here, bo E R represents the landmark representation out- put by the previous decoder layer of the g-th group, which accumulates all information from groups 1 to g 1. This enables he to perform multiple retrievals based retrieved information fused by previous layers. W W, Rdxd apply linear transformations on their respective inputs. r represents the causal relevance score of c ez C contains the indices of past chunks with top-k relevance scores. The retrieved chunks are shared among the subsequent layers within the same group. The upper layers apply GCA to fuse retrieved information into the decoder.
(0),
Exeeogh) exp().(0.0)
(2)
0
Σ Eufitiu Oli
A = Norm(H) O'41).
Here g(l) converts the layer index to the group index and C₁ E Sxd represents token representations the k-th retrieved chunk. O E g($+1)xd represents the informa- tion that 5 + tokens chunk c++1 gather from past chunk w x(), is the normalized relevance score after softmax, serving the weight of ++1 for information fusion. The final fused results of GCA is O
Grouped Cross-Attention. For the /-th layer, let H and H E (f denote the batched token represen- tations in the (t + 1)-th chunk before and after GCA. The function g(l) is used to map layer index a group index
Since C shared across layers, we use the same K, linear transformations across layers to compact model pa- rameters and reduce memory footprint. For each head h, we

---

Length-Generralizable Attention via Causal Retrieval
have CA CA(H(+1: Ok. Ck)) defined as:
largely reduces the complexity but maintains the random- access flexibility.
Softmaxı Uh (H )V.(CC) Vdh.
Memory Offloading. During the inference stage of Trans- formers, the default memory cost for KV Cache is ( Ld). To reduce GPU memory usage, we can offload past chunk representations to CPU memory. This results in a spatial complexity of the GPU memory usage O 5 + 2 Kdd + NWd) d) during inference. Here, Ld is the memory footprint of landmark representations, while the remaining terms ac- count for the GCA and sliding-window KV cache. Although each retrieval involves gathering representations from CPU memory and transferring them to the GPU, this operation oc- curs G times every S tokens. Therefore, the cost of memory exchange from chunk retrievals is minimal.
Here, dn is per-head dimension, and Qh Kn, Vo, are linear transformations per head, where 0% varies across layers and Kn, VA are layer-shared. We use softmax-off-by-nn to allow tokens in the current chunk ignore any retrieved tokens. The final CA outputs are concatenated vectors from all heads. It is worth noting that GCA is easy integrate with FlashAttentions as detailed in the pseudo-code in Appendix B.
GCA vs CCA. A key distinction between GCA and CCA lies in how softmax is applied cross-attention matrices as shown in Figure ((a)(b). In CCA, all retrieved chunks are concatenated and softmax is directly applied to the whole attention matrix to fuse token-level information. Notably, relevance scores are entirely excluded from the process. In contrast, GCA applies softmax to each chunk's ' attention matrix separately. This modification allows information to be gathered from each chunk separately. The softmaxed relevance scores then serve soft choices (Hu et al., 2021; 2022), participating in the next token prediction thus receiv- ing back-propagated gradients.
4. Experiments
We compare GCA with other attention mechanisms in lan- guage modeling, real-word downstream tasks, and RULER benchmark (Hsieh et al., 2024). Notably, we include the closely related works RPT and Landmark Attention. To en- sure faimess, we train all models from scratch with similar configurations.
4.1. Experimental Setup
3.2. Training& & Inference
4.1.1. DATASETS
The pre-training objective of DRT the next token predic- tion, but there are certain details as discussed below. We also give detailed time complexity analysis of training.
PG19. PG19 (Rae et al., 2020) is a language modeling benchmark widely used to evaluate long-range text under- standing capabilities of models. It includes a collection of full-length books from Project Gutenberg across wide range of genres, styles, and topics. The dataset is particu- larly useful for evaluating models' abilities to capture de- pendencies and context over long sequences of text.
Gumbel Top-k sampling. The core idea of self- supervised retrieval is making candidate chunks compete with each other by softmaxing relevance scores as weights. The weights of the chunks contributing most to the next chunk prediction are enhanced while the weights of the rest suppressed. To balance exploration and exploitation, we sample chunks based on relevance scores instead of always picking the top-k, enabling highly relevant chunks to be more likely chosen while still considering lower-scoring ones. A simple trick is to add Gumbel noise (Gumbel, 1954) to the raw scores before the top-k operation. Importantly, this noise doesn't affect subsequent operations.
ArXiv-math. ArXiv-math is corpus consisting of math- ematical papers from arXiv. Characterized by a sustained coherence over long distance, the corpus requires models' capability of correctly referring to long-range history infor- mation and using them effectively for predictions. We use the preprocessed corpus from (Azerbayev et al., 2023).
Time Complexity. Our approach reduces training com- plexity by compressing quadratic operations. Vanilla Trans- formers have a complexity of ON L²) for full self-attention. In DRT, we encode chunks with S tokens into landmark representations, performing chunk-wise full attention to gr)) for G groups of
4.1.2. MODELS
DRT re A DRT consists of 12 Transformer de- coder layers, divided into lower and upper layers, with upper layers further split into G groups. The sliding window size is set to W=512, the chunk size is set to S=64, and 8 chunks are retrieved for GCA, resulting in an attention field of 512 (8x x 64). We implement hardware-aware GCA based on Triton (Tillet et al., 2019). As we employ various param- eter settings across different experiments, further details are provided Appendix
compute relevance scores within (C 12 upper layers. By employing diding-windoo attention and top-k retrieval, we scale down self-attention and GCA costs to ( 52 + N LKS + NLW), where K is the retrieved chunk number and Wi is the window size, K, W < L.. This
2
5

---

Length-Generalizable Attention via Causal Retrieval
attn. win.
PG19]
ArXiv-math, test
PG19
ArXiv-mathu
Model
Time
valid
test
valid
valid
test
valid
test
Train length=16k, eval length 2K
128M
350M
BaseL.M (Sliding Window+Alibi)
512
15.00
14.10
3.31
3.31
12.98
12.09
3.15
3.15
1.03x
768
14.86
13.96
3.24
3.24
12.86
11.98
3.04
3.04
+2 layers
1.15x
704
14.71
13.92
3.06
3.06
12.78
11.89
3.08
3.08
RPT contiieeer impl.)
2.5x
8
512
14.81
13.92
3.24
3.24
Landmark Attn.
1.5x
768
14.41
13.40
3.17
3.16
12.52
11.64
2.91
2.91
Block Recurrent TFM
2x
768
15.99
15.00
3.33
3.32
DRT retrioaalx
1.22x
512
14.65
13.78
3.24
3.24
12.90
12.02
3.05
3.05
DRT x2
1.24x
512
14.56
13.69
3.22
3.22
12.66
11.78
3.01
3.01
Train length=16K, eval length = 16k
BaseLM
1x
512
14.55
13.68
3.06
3.06
12.57
11.70
2.91
2.91
(Sliding window+Alibi)
1.03x
768
14.36
13.49
2.95
2.95
12.40
11.55
2.76
2.76
+2 layers
1.15x
658
14.23
13.37
2.95
2.94
12.33
11.47
2.80
2.80
RPT comrrever impl.)
2.5x
512
14.39
13.52
2.93
2.92
Landmark Attn.
1.5x
4
768
14.10
13.21
3.02
3.02
12.18
11.25
2.70
2.70
Block Recurrent TPM
2x
768
15.59
14.60
3.14
3.14
DRT retricvalx1
1.22x
8
512
14.05
13.21
2.89
2.89
12.37
11.48
2.74
2.74
DRT 2
1.24x
512
14.02
13.18
2.85
2.85
12.12
11.22
2.68
2.68
BaseLMM
512
14.50
13.64
3.05
3.04
12.52
11.67
2.88
2.88
(Sliding Window+Alibi)
1.03x
768
14.30
13.46
2.93
2.92
12.35
11.52
2.74
2.73
+2 layers
1.15x
658
14.18
13.34
2.93
2.92
12.29
11.43
2.78
2.77
RPT cour impl.)
2.5x
8
512
14.35
13.49
2.91
2.91
Landmark Attn.
1.5x
4
768
14.19
13.33
3.07
3.07
12.17
11.24
2.75
2.75
Block Recurrent TFM
2x
768
15.61
14.56
3.13
3.12
DRT 1
1.22x
512
14.01
13.19
2.85
2.85
12.33
11.45
2.71
2.71
DRT x2
1.24x
512
13.98
13.16
2.81
2.81
12.08
11.19
2.65
2.65
Ablation studies (Train length=16K, eval length = 16k)
-wo Triton
1.45x
8
512
-w/o gumbel top-k
1.22x
8
512
14.36
13.46
2.90
2.90
w/ contriever
2.5x
512
14.55
13.69
3.06
3.06
12.57
11.70
2.91
2.91
w/ random retriever
1.22x
512
14.53
13.67
3.05
3.05
12.56
11.69
2.90
2.90
Table 1: Perplexity for all datasets. We highlight the best results in bold and underline the second best.
Base LM. Our base LM is based on the implementation of TinyLlama (Zhang et al., 2024) combined with Flash Atten- tion2 (Dao, 2024) enabling ALiBi (Press et al., 2022) and sliding window attention (Child et al., 2019). We compare models against the baseline across various configurations. One configuration involves 12 layers with a sliding window of 512 tokens, aligning with the DRT sliding window size. Another configuration of 12 layers with a 768-token slid- ing window ensures the same attention field coverage, as 12 768 12 512 6 512(GCA). The strongest base- line, with 14 layers and 658-token sliding window, has parameter count comparable to our DRT while maintaining similar total attention field across all 12 layers, calculated as 658 x 14 ~ 12x x 512 + 512.
the retriever is not released, we reimplement RPT in Py- Torch and replace the retriever with Contriever (Izacard et al., 2022). Elaborations be found in Appendix E.
Landmark Attn. We use the official Llama-like imple- mentation of Landmark Att.. Similar to Base LM, we extend the length the self-attention range from 512 to 768 to ensure shares the same attention field as DRT.
Bacck-Reeurrent TFM. Since the official implementation of Block-Reccurrett Transformer is also based on JAX, we utilized a Py PyTorch implementation to ensure all baselines are running with the same framework.
Retrieval- Pretraiied Transformer (RPT). Since the of- ficial implementation JAX and the code for distilling
/tttps:/ http:::///ithbb.... httptthhww...ommciiieeeeaaeeaeee peeerreeettrrraaffooe-----ooooo
6

---

Leneaa-Gennrrliiaable Attention via Causal Retrieval
Ablations. w/o Triton: A naively implemented version of GCA without Triton. w/o gumbel top-k: The architecture is exactly the same DRT, while the only difference lies in the training process by eliminating the gumbel noise when selecting the top-k chunks. w/ contriever: We replace the relevance scores in GCA by using an off-the-shelf retriever. w/ random retriever: We pick top-k chunks randomly but still keep GCA learnable.
et al., 2018), NIAH tests akin to RULER. The details for the NIAH tests are described in Appendix D. Then we analyze the inference cost, the relationship between training time and context length, the training throughput, and the extrap- olation capability of DRT. In the inference cost analysis, We skip RPT because it has a similar cost to DRT. In the extrapolation experiments, we utilize the pre-trained models described in $ 4.2 to assess their perplexity on longer inputs.
4.2. Long-range Language Modeling
Training time per batch&step different lengths
PPL PG19
15.2
landmark
In this section, we evaluate DRT against baselines in long- range language modeling on PG19 and arXiv-math, and report their respective perplexities. All models are pre- trained with the same attention field and 16K context by default, except for baselines that cannot efficiently pre- train on long contexts. To ensure faimess, we adjusted these baselines. Detailed hyper-parameteed are provided in Appendix A.
15.0-
RPT outs
10
14.0
15 14.8
14.4
10
512
26 de
16t
32k
64c
14.0
16k
32x
64k
train Len,/ooenss
evall Lon./tokens
(a)
(b)
Results. From Table 1, we have several observations. Firstly, DRT outperforms all baselines where the evaluation length exceeds 16K. While DRT performs retrieval for G times every 64 tokens, LA performs retrieval at every token and in every layer, potentially offering better random access flexibility. However, DRT still surpasses LA on longer in- puts. This is likely because LA follows a "train short, test long" approach due to the need for full attention during pre- training, whereas DRT can be directly pre-trained on longer sequences thanks to GCA. DRT can randomly access dis- tant contexts during pre-training with fixed attention field, allowing it to better utilize long-range information during pre-training with negligible extra training costs. Secondly, in terms of parameter efficiency, we compared against Base LM with two additional layers in the decoder stack. It can be observed that on shorter evaluation lengths, the baseline has slight advantage. However, with longer context, our model consistently leads. This experiment suggests precisely retrieving semantic knowledge from long con- text may be more beneficial for improving the language model than increasing model parameters. Thirdly, multiple retrievals yield positive gains with a low marginal cost. al- lows upper layers to access more diverse chunks and enables further retrieval based on previous retrieval results. Finally, ablation studies show that all the training techniques we add bring positive improvements. In conclusion, the above results fully demonstrate that the GCA module can indeed bring effective gains in modeling long texts, and it is more advantageous compared to other baselines.
Figure 3: Training speed, extrapolation ability
Results. From Table 2, we observe that DRT significantly outperforms all baselines in the summarization tasks, val- idating its capability to effectively utilize long contexts. Notably, in the single NIAH test, DRT maintains 100% performance even with context length of up to 16 mil- lion tokens, demonstrating its strong length generalization ability in long-context scenarios. Furthermore, in the 2-hop NIAH test, DRT x performs comparably to Landmark Attn at context lengths below 16K tokens while success- fully extrapolating to longer context lengths beyond 64K to- kens. Additionally, DRTretievaxx2 significantly outperforms DRT 1e confirming our hypothesis that conducting causal retrieval every G-layer contributes to scenarios re- quiring multiple retrievals.
Table 3 shows DRT significantly outperforms Landmark Attn in terms of memory footprint and inference speed. The main overhead for Landmark Attn arises from modifica- tions the self-attention KV cache and gathering tensors from memory offloaded to the CPU. In DRT, thanks the chunk-wise retrieval, we perform retrieval only once ev- ery 64 tokens, which is 1 1/(12 x 64) of the corresponding operation in Landmark Atn.
From Figure 3(a), it can be observed that the total training time increases linearly with the increment of the sequence length. In the extrapolation experiments, both BaseLM and DRT perform well as shown in Figure 3(b). However, Land- mark Attn fails to extrapolate at longer eval lengths. While both methods use fixed attention window and dynamically select past chunks to fill it, why can GCA extrapolate over 1000 times with stable PPL and perfect passkey retrieval, while Landmark Attn fails? We speculate a major reason lies
4.3. Downstream task evaluation & Efficiency Analysis
In this section, we fine-tune all baselines under the same configurations and evaluate them against downstream tasks, including summarization (Nallapati et al., 2016; Narayan
7

---

Length-GGnneralizabb Attention via Causal Retrieval
Models
Single
NIAHt
Retricval
128M
4K
8K
16K
32K
64K
128K
256K
16M
Base LM (+2 layers)
15.37
8.30
3.89
2.13
0.0
Block Recurrent TFM
13,96
7.60
6.01
2.13
4.29
CPTconsrieerr
fixed
11.66
6.71
4.24
1.42
2.86
Landmark Attn.
adaptive
99.82
97.74
97.88
96.45
0.00
0.00
DRT 1erieall1
adaptive
98.50
98.76
98.59
100.00
100.00
100.00
100.00
100.00
DRT n2riovalx2
adaptive
99.65
99.47
99.65
99.99
100.00
100.00
100.00
100.00
For reference
Yarn-base (7B)
100.00
100.00
99.65
99.82
100.00
42.32
CEPE (7B)
100.00
100.00
55.45
Models
XSumt
CNN/Daily Mailf
2-hop
NIAH 1
128M
R-1
R-2
R-L
R-1
R-2
R-L
1K
4K
16K
64K
256K
4M
BaseLM
29.43
8.04
23.26
32.38
12.57
22.61
3.60
1.15
0.71
0.0
0.0
+2 layers
29.74
8.26
23.48
34.14
14.09
23.60
16.60
5.83
1.06
0.0
0.0
Landmark Attn.
27.98
6.96
21.99
34.06
13.80
23.77
90.82
88.35
86.41
0.0
0.0
1 1ouievalx11
DRT
30.30
8.59
23.92
36.27
15.88
25.08
41.07
33.39
39.93
38.57
35.29
34.29
DRT retrieaaxx2
30.39
8.64
23.98
36.39
15.96
25.15
88.52
84.45
86.21
81.43
94.11
79.41
Table The performances various models on summarization tasks akin RULER.
Prompt #tokens
Generated #tokens
w/ cpu offload
w/o cpu offload
time/token,
mem. cost.
time/token.
mem. cost.
16K
128
160.9x
1.98x
4.16x
32x
Landmark Attn. Basc LM
48K
128
163.1x
2.98x
4.25x
96x
16K
128
2.85x
Block Recurrent TFM / Base LM
48K
128
2.85x
2x
16K
128
1,25x
1.54x 1.62x
1.06
4.08
DRTineeat Base LM
48K
128
1.27
1.08x
9.41 x
Table 3: The inference time per token and memory footprint ratio compared to the Base LM (12 layers with 512 sliding window), with lower values indicating better performance.
Model
Throughput Memory
Throughput opparams. 760M, ctx-len=16K
Memory
Throughput #purams.=l
Memory 8tx-len-8K
Throughput Npurams.=31
Memory etx-len=4K
Mpurams, 330M, etx-len=32K
Transformersw
2.91 x
1.31 10"
50G
6.56 x 10
50G
3.41 x 10"
72G
Transformerym.an
4.44 10
50G
4.00 10
50G
3.50 x 10*
72G
2.85 10*
72G
BRT
10*
10
56G
4.02 10*
59G
0OM
RPTou impl
1.13 10
57G
3.17 10a
59G
2.73 10*
64G
1.45 10a
80G
DRT DRTumeei
2.32 10
56G
1.06 10
59G
5.52 10*
63G
2.93 10"
80G
DRT ae
2.28 x 10
57G
1.05
10
59G
5.52 x 10*
63G
2.93 x 10*
81G
Table 4: Training throughput of models with parameters ranging from 350M to 3B.
in the different objectives for optimizing retrievers. Land- mark Attn optimizes for predicting the next token based on retrieved chunks, aiming for maximizing P( (R(x()) where ((F<)) denotes retrieved chunks. In contrast, GCA retrieves past chunks for the next chunk prediction, aim- ing to maximize P( ()) A single token offers limited information, which may result in false positive re- trievals thus leading to inefficacy in larger retrieval spaces. While chunks provide richer semantic information, enabling generalization larger search spaces.
in Figure 4. When retrieving past chunks, the results not only include the definition of prepositions referenced in the current chunk but also lemmas to be used in the next chunk. This validates the idea of causal retrieval, allowing to not only retrieve semantically similar content but also information that better predicts the next chunk. More case studies can be found in Appendix C.
5. Conclusion & Future works
In this study, we introduce mechanism that maintains efficient long-range information access. Notably, it achieves perfect accuracy in passkey retrieval on 16M contexts-a first in the field. The core in-
attention
4.4. Case Studies
By analyzing DRT's retrieval results on the arXiv-math dataset, we find some intriguing cases. A case given
8

---

Lengal-Generalizaab Attention via Causal Retrieval
We start first with the following lemma which is useful to establish the Quillen equivalence. [begin(lemma) (ebbee
Let M be
symmetric monoidal model category that combinatorial and left proper. Assume that the transferred (natural) model structure com exists. Let C between usual commutative...
morphism
weak equivalence between co-Segal categories is just a level-wise weak equivalence. (propn plabel-poo-eta-kk--ooo---
For any F coms, the canonical map F + [F] is an equivalence in comsepe i.c, it's kb([)-local equivalence comsep (whence in comse). )
Thanks to Proposition /
we know that 7) F -> [F is always ko(T)-local equivalence. Then by 3-for-2
that is kb(T)- local equivalence if and only (l()) is. Now thanks to
Lemma vre((lem-reflect-equivv we know that ...
Figure 4: the case above, retrieved 2nd best chunk describes proposition which appears in the current chunk. retrieved best chunk and retrieved 3rd best chunk are adjacent chunks which introduce the lemma used in the next chunk.
novation lies in the Grouped Cross-Attention (GCA), which makes relevance scores learnable by using them to fuse information retrieved by the current chunk for next chunk prediction. Combined with Gumbel top-k sampling, this ap- proach enables the pre-training of LMs efficiently on context lengths extending up to 64K tokens.
In future work, will explore self-supervised causal re- trieval from vast amounts of tokens outside the context. Meanwhile, we will combine structured representations (Hu et al., 2024b;a) to achieve multi-granular retrieval.
Impact Statement
This paper presents a novel attention mechanism that aims to address the length generalization issue of L.Mss. There are many potential societal consequences of our work associated with LLMs. Beyond LLMs, we feel no other consequences must be highlighted here.

---

Leneth-GGeneralizable Attention via Causal Retrieval
References
commenuss///1222515//aaarrrrrgr rope_
a awlow_11amm_nne wataaaaaodesssseeeeeeeeee s_toddleesto_eeee....
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S.,
Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Ruther- ford, E., Millican, K., van den Driessche, G., Lespiau, J., Damoc, B., Clark, A., de Las Casas, D., Guy, A., Menick, J., Ring, R., Hennigan, T., Huang, S., Mag- giore, L., Jones, C., Cassirer, A., Brock, A., Pagamini, M., Irving, G., Vinyals, O., Osindero, S., Simonyan, K., Rae, J. W., Elsen, E., and Sifre, L. Improving language models by retrieving from trillions of tokens, In Chaudhuri, K., Jegelka, S., Song, L., Szepesvári, C., Niu, G., and Sabato, S. (eds.), International Con- ference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Pro- ceedings of Machine Learning Research, pp. 2206- 2240. PMLR, 2022. URL https:/ h/tpsss:s///ccciinggs.... www.s/rasseeee /vl http:/vv12/bbboe/ wwtpb/ogeaeue.....mo..
Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Asai, A., Min, S., Zhong, Z., and Chen, D. Retrieval- based language models and applications. In Chen, Y.- N. V., Margot. M., and Reddy, S. (eds.), Proceedings of the 61st Annual Meeting of the Association for Com- putational Linguistics (Volume 6: Tutorial Abstracts), pp. 41-46, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: (0.18653/11/2023. acl-tutorials.6. URL attps://aclanthologyy.ou 2023.ac1-tutorien
Azerbayev, Z., Ayers, E., and Piotrowski, B. Proof- pile: A pre-training dataset of mathematical text, 2023. URL https://huggingfaae coceee.oo/aaeeee cotaatsetss/////// hoski nson- potooncceeeerrrooo-ppiii
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan,
J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https:
Bai, Y., Lv, X., Zhang, J., He, Y., Qi, J., Hou, I.., Tang. J. Dong, Y. and Li, J. Longalign: A recipe for long context alignment of large language models. In Al- Onaizan, Y., Bansal, M., and Chen, Y. (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, Miami, Florida, USA, November 12-16, 2024, pp. 1376-1395 Association for Computational Linguis- tics, 2024. URL https:://///..... aawlaaaalannto...... . /204. findings-emn.l 74.
Beck, M., Pöppel, K., Spanring, M., Auer, A., Prud- nikova, O., Kopp, M., Klambauer, G., Brandstetter, J. and Hochreiter, S. xlstm: Extended long short- term memory. CoRR, abs/2405.04517, 2024. doi: 10. 48550/ARXIV.24005 044177 URL https://doi...... /rg// 10.48550/ (wt/5550/xxiv.... 04517.
/ //proceedings http://pppp.ccpppp///////aaaaaaa 145/cod6btcco49
http///411111111e... htw.ccaaarr...........aaaaaaaaa. html.
Burtsev, M. S. and Sapunov, G. V. Memory transformer. CoRR, abs/2006. 11527, 2020. URL https://arxiv. 115277
Biderman, S., Schoelkopf, H., Anthony, Q. G., Bradley, H., O' O'Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., Skowron, A., Sutawika, L., and van der Wal, O Pythia: A suite for ana- lyzing large language models across training and scal- ing. In Krause, A., Brunskill, E., Cho, K., Engel- hardt, B., Sabato, S., and Scarlett, J. (eds.), Interna- tional Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 2397- 2430. PMLR, 2023. URL https: //t:::/pproooo pptppeeecceeen... w.r.press/vv000
Chen, S., Wong, S., Chen, L., and Tian, Y. Extending context window of large language models via positional interpolation. CoRR, abs/2306. 15595, 2023. doi: 10. 48550/ARXIV.2306.15550 URL attps://doi.org/ 10. 00.44850/arrKa http://wwwwv......... 15595.
Child, R., Gray, S., Radford, A., and Sutskever, Gener- ating long sequences with sparse transformers. CoRR, abs/1904. 10509, 2019. URL http://arxiv.org/ abs 1104 10509.
nttm.
Dai, Z., Yang, Z., Yang. Y., Carbonell, J. G., Le, Q. V., and Salakhutdinov, R. Transformer-xi: Attentive language models beyond a fixed-length context. In Korhonen, A., Traum, D. R., and Mårquez, L. (eds.), Proceedings of the 57th Conference of the Association for Computa- tional Linguistics, ACL. 2019, Florence, Italy, July 28-
bloc97. Ntk-aware scaled rope allows Ilama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation, 2023.
URL https: / /www. reddi t.
10

---

Lenggh-GGeneralizabbi Attention via Causal Retrieval
August 2, 2019, Volume I: Long Papers, pp. 2978-2988. Association for Computational Linguistics, 2019. doi: 10. 10.18663/V1//11-1235 1285. URL
Language Processing, EMNLP 2022, Abu Dhabi, United
Arab Emirates, December 1, 2022, pp. 2809-2821, Association for Computational Linguistics, 2022. doi: URL https: /00 10.1116633vvv /2022 .mmn p-main... maii.... 181.
1865
10. //doi //doi.org/111
Dao, T. Flashattention-2: Faster attention with better par- allelism and work partitioning. In The Twelfth Interna- tional Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL htttps://opeereve// net/ forum?id- m2n2Xyh98c.
Hu, X., Ji, P., Zhu, Q. Wu, W., and Tu, K. Generative pretrained structured transformers: Unsupervised syn- tactic language models at scale. In Ku, L., Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd An- nual Meeting of the Association for Computational Lin- guistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August I 11-16, 2024, pp. 2640-2657. Associa- tion for Computational Linguistics, 2024a. URL https: / //acl anthology wrg//024.2222 .ccl long. 1 145.
Dao. Т. and Gu, A. Transformers ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning, ICMI. 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https:// openreview wwt.
Hu, X., Zhu, Q., Tu, K., and Wu, W. Augmenting trans- formers with recursively composed multi muulti-rained repre- sentations. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, OpenReview.neh 2024b. URL ht https: //openevveww net wttt://orummmt/m www.ooooumiiiuuuuoooo (tg.ddu--u.9977......
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Ilama herd of models. arXiv preprint arXiv:2407.21783, 2024.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. CoRR, abs/2312.007522 2023. doi: 10.48550ARXXV.2312.0077
Hutchins, D., Schlag, I., Wu, Y., Dyer, E.,, and Neyshabur, B. Block-recurrent transformers. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
URL sotttp:://oo. org/11 /rgg11..0 10. 48550/ wwt.:iiii.......... 2312. 00752.
Gumbel, E. J. Statistical theory of extreme values and some practical applications: a series of lectures, volume 33. US Government Printing Office, 1954.
Izacard, G., Caron, M., Hosseini, L., Riedel, S., Bojanowski, P., Joulin, A., and Grave, E. Unsupervised dense infor- mation retrieval with contrastive learning. Trans. Mach. Learn, Res., 2022, 2022. URL httpss:/epenreviio/ net/forun's Pid= (at1:iin111111111
Hsieh, C.-P., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., and Ginsburg, B. RULER: What's 's the real context size of your long-context language models? In First Conference on Language Modeling, 2024. URL https: //openreview net. for.m?id-kIIEoeeee
Hu, X., Mi, H., Wen, Z., Wang, Y., Su, Y., Zheng, J., and de Melo, G. R2D2: recursive transformer based on differentiable tree for interpretable hierarchical lan- guage modeling. In Zong, C., Xia, F., Li, W., and Nay- igli, R. (eds.), Proceedings of the 59th Annual Meet- ing of the Association for Computational Linguistics and the 11th International Joint Conference Natu- ral Language Processing, ACL/IJCNLP 2021, (Volume 1:
Lewis, P. S. H., Perez, E., Piktus, A., Petroni, F. Karpukhin, V. Goyal, N., Küttler, H., Lewis, M.. Yih, W., Rocktäschel, T., Riedel, S., and Kiela, D. Retrieval-augmeened generation for knowledge-intensivee NLP tasks. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Infor- mation Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020. URL https: / //proceedi www.nenriipeeee www.oceedinngs.nnrrss....... cwwprrppp.cc/ppppe//0/00000 6b4 4444:339000111000002294888888888 Abtec. html.
Long Papers), Virtual Event, August 1-6, 2021, pp. 4897- 4908. Association for Computational Linguistics, 2021. doi: 10. //dai , org/111/1/11 100.////11110 0. 8853 /2021 aal..aacl-10 1ogg0.0..31111111377323
URL https:
Hu, X., Mi, H., Li, L.,, and de Melo, G. Fast-r2d2: A pre- trained recursive neural network based on pruned CKY for grammar induction and text representation. In Gold- berg, Y., Kozareva, Z., and Zhang, Y. (eds.), Proceedings of the 2022 Conference on Empirical Methods Natural
Liu, H., Zaharia, M., and Abbeel, P. Ring attention with blockwise transformers for near-infinite context. CoRR, abs/2310.01889 2023. doi: 10.48550/ARXIV.2310. 01889. URL https: http://doo..gggo.o 4ttp:///0......rxx....... 2310.01889.

---

Leneah-GGnnrralizablc Attention via Causal Retrieval
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- larization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019, URL https: // //openrevien net wwwee//forumideeeeeeeee (on.rummmminikkkki.... Cqy7.
Peng, B., Quesnelle, J., Fan, H., and Shippole, E. Yarn: Efficient context window extension of large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: //openzcview http:/lwww..tteteeeeeeeeeeeeeeee
Martins, P. H., Marinho, Z., and Martins, A. F. T. cormor::r: Infinite memory transformer. In Muresan, S., Nakov, P., and Villavicencio. A. (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), ACL 2022, Dublin, Ireland, May 22-27. 2022. pp. 5468-5485. Association for Computational Linguistics, 2022. doi: 10.18653/V1/
Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https: //penzzvieww.c wwwp:/ewwwwwneeteeeeeeeeee wtt.:e/eoooooocccccee00eeeeee.e.
2022.ACL-LLON.3355. URL https:/ / http//di..o/// org/10. ong.375.
Qwen, Yang, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Li, C.,, Liu, D., Huang, E, Wei, H., Lin, H., Yang, J., Tu, J., Zhang, J., Yang, J. Yang, J., Zhou, J., Lin, J., Dang, K., Lu, K.. Bao, K., Yang, K., Yu, L., Li, M., Xue, M., Zhang, P., Zhu, Q. Men, R., Lin, R., Li, T., Tang, T., Xia, T., Ren, X., Ren, X., Fan, Y., Su, Y., Zhang. Y., Wan, Y., Liu, Y., Cui, Z., Zhang, Z., and Qiu, Z. Qwen2.5 technical report, 2025. URL https: //arxiv. .rgg abs 221222222222 15115.
Mohtashami, A. and Jaggi, M. Random-access infinite context length for transformers. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Radford, A. Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsuper- vised multitask learners. 2019. URL attps://api. semanticscholar.com wwwpe/hhlla...ogorrpsssttt. 160025533.
Munkhdalai, T., Faruqui, M., and Gopal, S. Leave no context behind: Efficient infinite context transformers with infini-attention. CoRR, Abs/2404.07144, 2024. doi: 10.45550/ARXXV.244000001 07143.
URL https::/dddo.
org/ 10. 48550/ http//r.ivvv.o/ 2404 071433
Rae, J. W., Potapenko, A., Jayakumar, S. M., Hillier, C., and Lillicrap, T. P. Compressive transformers for long- range sequence modelling. In International Conference on Learning Representationss 2020. URL https:// openreview not wooomrodd-dyyy
Nallapati, R., Zhou, B., dos Santos, C., Guülçehre, Ç., and Xiang, B. Abstractive text summarization using sequence- to-sequence RNNs and beyond. In Riezler, S. and Gold- berg, Y. (eds.), Proceedings of the 20th SIGNLL Confer- ence on Computational Natural Language Learning, pp. 280-290, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 10. 10.18653/v1/KK6- 1028. URL http::ttt////ett http:p::///aaaaahoooagee org/ K1 x16-1028.
Reid, M., Savinov, N., Teplyashin, D., Lepikhin, D., Lilli- crap. T. P., Alayrac, J., Soricut, R., Lazaridou. A., Firat, O., Schrittwieser, J., Antonoglou, I., Anil, R., Borgeaud, S., Dai, A. M., Millican, K., Dyer, E., Glaese, M., Sottiaux, T., Lee, B., Viola, F,, Reynolds, M., Xu, Y., Molloy, J., Chen, J., Isard, M., Barham, P., Hennigan, T., Mcllroy, R., Johnson, M., Schalkwyk, J., Collins, E., Rutherford, E., Moreira, E., Ayoub, K., Goel, M.. Meyer, C., Thornton, G., Yang, Z., Michalewski, H., Ab- bas, Z., Schucher, N., Anand, A., Ives, R., Keeling, J., Lenc, K., Haykal, S., Shakeri, S., Shyam, P., Chowd- hery, A., Ring, R., Spencer, S., Sezener, E., and et al. Gemini L.5: Unlocking multimodal understanding across millions of tokens of context. CoRR, abs/24433005530 2024. doi: 10-48550/ARX1V22400055555 URL ht https: //doi orrg/ 10. 48550/ar55i10 2403. 22.224433333 055330.
Narayan, S., Cohen, S. B., and Lapata, M. Don't 't give me the details, just the summary! twpee- convolutional neural networks for extreme summarization. In Riloff, E., Chiang. D., Hockenmaier, J., and Tsujii, J. (eds.), Pro- ceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 1797-1801 Brus- sels, Belgium, October-November 2018. Association for Computational Linguistics. doi: 10. 10111855/1/D1... 1206. URL htt / http::/:////..mmmooooooooo////// wtwp:/looyyyg//11-1110001
Nunez, E., Zancato, L., Bowman, B., Golatkar, A., Xia, W., and Soatto, S. Expansion span: Combining fading mem- ory and retrieval in hybrid state space models. CoRR, 0as/2412.13328 2024. doi: 10.48550/ARXIV22412..
Robertson, S. E. and Zaragoza, H. The probabilistic rel- evance framework: BM25 and beyond. Found. Trends Inf. Retr. 3(4):333-389 2009. doi: 10. 10.1561/11000000099 URL wwt.tttppp.o.. //wpi//////////// http:://////////////////oo/ /10.//1..10 1.....66 500000019.
13328. URL https::/// htt////oi..oo orc.oorg/1...o arxiv. 2412.13328
12

---

Length-Generalizabbe Attention via Causal Retrieval
Rubin, O. and Berant, J. Retrieval-pretrainee transformer:
OpenReview.net, 2024. URL https: https::/openrevieet net/Foorum21--6775537
Long-range language modeling with self-retrieval. Trans.
Assoc. Comput. Linguistics, 12:1197-1213, 2024. doi: ( -00693. URL https:/ //doi.///////// /rr/ 20.1162/tac1 _a__6693.
Yen, H., Gao, T., and Chen, D. Long-context language modeling with parallel context encoding. In Ku, L., Mar- tins, A., and Srikumar, V. (eds.), Proceedings the 62nd Annual Meeting of the Association for Computational Lin- guistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 1 11-16, 2024, pp. 2188-2610. Associa- tion for Computational Linguistics, 2024. URL https: / www.aaattoooog..org/222222 acl long. .142.
Tillet, P., Kung, H., and Cox, D. D. Triton: an in- termediate language and compiler for tiled neural net- work computations. In Mattson, T., Muzahid, A., and Solar-Lezama, A. (eds.). Proceedings of the 3rd ACM SIGPIAN International Workshop on Machine Learn- ing and Programming Languages, MAPL@ PLDI 2019, Phoenix, AZ, USA, June 22, 2019, pp. 10-19, ACM, 2019. doi: 10. 10.1145/33155088332993.. URL https: //doi wtt. 100 33/
Zhang, P., Zeng, G., Wang, T., and Lu, W. Tinyllama: An open-source small language model, 2024.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint (0XY:2237..09288, 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N.. Kaiser, L., and Polosukhin, I. Attention is all you need. In Guyon, I., von Luxburg, U. Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Pro- cessing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-6008, 2017. URL https: //proceedings... //prossedng.s..uriis...... htt.:erwwpp.ccccpppp//////aaaaaa 3fee 4 3eeeee2445dddee9111 aaao html.
Wang, J., Ji, T., Wu, Y., Yan, H., Gui, T., Zhang, Q. Huang, X., and Wang, X. Length generalization of causal trans- formers without position encoding. In Ku, L.., Martins, A. and Srikumar, V. (eds.), Findings the Associa- tion for Computational Linguistics, ACL 2024, Bangkok, Thailand and virtual meeting. August 11-16, 2024, pp. 14024-144 14040. Association for Computational Linguis- tics, 2024. doi: 10. 10.:8653/1/20224FFDDIN
834. URL https:/ http::/t.....r...111..//// findings-acl.es
Wu, Y., Rabe, M. N., Hutchins, D., and Szegedy, C. Memorizing transformers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event. April 25-29, 2022. OpenReview.net, 2022. URL Tz J=x2Reonff-
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Effi- cient streaming language models with attention sinks. In The Twelfth International Conference on Learning Repre- sentations, ICLR 2024, Vienna, Austria, May 7-11, 2024.
13

---

Length-GGnneralizablc Attention via Causal Retrieval
A. Hyper-parameters
Long-Range Language Modeling. We employ Llama-like architecture (Touvron et al., 2023) featuring 12-layer, decoder-only transformer with 12 heads per layer (64 dimensions each), an embedding dimension of 768, and an FFN size of 2048. Training utilizes the AdamW optimizer (Loshchilov & Hutter, 2019) with 81 0.9 and [2 0.95, and weight decay factor of 0.001. We used base learning rate 2 10-3 for all our experiments with warmup stage that 2% of the whole training and applied cosine scheduler with final learning rate being 4 x 10-1, We used GPT-2's (Radford training with 6 over at 8 Nvidia A100 GPUs. We train all models with an effective batch size of 2 tokens for 60K steps resulting in a total training budget of 32.2 billion tokens. We train Base LM, RPT and DRT on each dataset with a context length of 16K tokens. Due to Landmark Attention doesn't t support sliding- window attention, the model is pre-trianed with full self-attention with a context length of 768. Due to Block Recurrent Transformer cannot fully paralleled, which takes 5x wall-clock training time with 16K context length, we pre-train it with context length of 4K.
et al., 2019) tokenizer. We used muixed-preciste
B. Hardware-aware GCA psuedo-code
Algorithm 1 FLASHGGA forward pass
Require: Matrices Q E xd K, V RRXNNN in HBM, vector E R in HBM, block sizes Be, Br. blocks QI.... ........ Qr, size B, x d each, and divide K, V in to x To blocks where T,
1: Divide Q into T, = B.
Nxx Be
K1,1,..., Kx.T. and V Vil..... VK,T. of size B, x d each.
2: Divide the output O E R Naxi into T, blocks O.,..., Or, of size B, d each, and divide the logsumexp L X into T, K blocks ........, LTT,KK K
B, each.
3: Divide the output O' R R N, into T, blocks .,,,...., OK.T, of size K x B, cach.
4: for S is do
5: Load Q; from HBM to on-chip SRAM.
6: Load wne from HBM to on-chip SRAM.
7: for 1 S k < K do
8: On chip, initialize ! (0) xa E E Brxd e (0) E ER.r,m, .m = (-oc) E CR.
9:
for 1 < j < T. do
10:
Load Ku,, k, from HBM to on-chip SRAM.
On chip, compute 5(5) Q Q, E B..
11:
12:
On chip, compute m. (5) max(m (5m ()-1) rowmax(S) E R p000 exp(s!") mm()) E B.xB.. (pointwise).
(0) = -
(9) g 1)-1) + rowsum (P()) R
On chip, compute ((1) diag(em end for
m (2-1) - () p
13:
14:
15:
On chip, compute O',, (1( O
16:
On chip, compute O < O, + we O
17:
Write O., to HBM.
On chip, compute Li, = m (Te) + log ! (Te)).
18:
19:
Write Li, to HBM.
20:
end for
21:
Write O, HBM as the i-th block of O.
22:
end for
23:
Return the output O and the logsumexp L.
14

---

Length-GGnneralizabb Attention via Causal Retrieval
Algorithm 2 FLASHGCA Backward Pass
Require:
Matrices Q, 0,0 RNxx K, VE BKXNEoKe xd L E RN, RKxN,xd in HBM, vector w E RR in HBM, block sizes Be, B,.
1: Divide Q into T,
blocks QI..... Qr, of size B, x d each, and divide K, V in to K x To. where Te blocks K1.1,..., Kx.r, and VI.,.... VK.T. of size Be d cach.
N B-
Divide into T. blocks O,,... Oz, of size B, divide L into T. x K blocks L..,...., LT,.K of size B, each.
each, divide dO into T. blocks dO.,... dOr, size d each, and
3: Initialize dQ= NO) xd HBB and divide into T. blocks dQ 1 dQr, of size B, x each. Divide dK, dV xx. K x T. blocks 1.1, dKk,T, and i......... dV K.T,, of size B. d cach. Initialize dW (0)r, x HBM. 4: Compute D = rowsum (dO O') E ER NN* (pointwise multiply). write D to HBM and divide it into T, blocks DI...., Dr, of
size B, each.
5: for 11 k S K do
6: Load wt from HBM to on-chip SRAM.
7: for 1 j S T. do
8: Load Ka.j. Va.j from HBM on-chip SRAM.
Initialize dK pJj = (0) B. dV (0) xd, dWa.. = (0) on SRAM.
10:
for S i T,
11:
Load Q ,0,dO.d0 dQ., Dr from HBM to on-chip SRAM.
12:
Load Li,k from HBM to on-chip SRAM.
13:
On chip, compute s = Q,K E B.
On chip. compute p(d) eSp(S.) LiN) B..B.
14:
15:
On chip, compute dV k. dV k. + (wiP dO, E Bexd On chip, compute dp
16:
do.v. E R B.
On chip, compute dW.a L rowsum (P()) dp!))
18:
On chip, compute ds()) wx,((") (dp()) D..)) B.
19:
Write dW., to HBM.
20:
Load dQ, from HBM to SRAM, then on chip, update dQ, + dQ, + dss K, E R. B..xd write back to HBM.
21:
On chip, compute a dK + ds( (k.,) Q R
22: end for
23:
Write dK &V.., to HBM.
24:
end for
25:
end for
26:
dW dWW uWsuum(dim = 0)
27: Return dQ, dK, dV, dW.
15

---

Leneth-GGnnrralizablc Attention via Causal Retrieval
C. More case studies
...An alternate approach is given below in Corollary \ref(cor:ifffy). Along the way we obtain more information about the eigenfunctions, which leads directly to an explicit formula for Um (æ; oo) see \eqref(conjsum2) and (eqref (Bkmexplicit). As a increases, the derivatives of um (a; a) remain bounded, and so to ensure that the interior condition in /eqref( Um Um (a)ai) converges to infinity....
continues to hold, the values
...must converge to 0 as converges to infinity. Our first corollary of Theorem \ref(thmmmin))
is that
these values converge to 0 at the same rate for each node zt. (corin(cor) Mabel an overall normalization factor, for each 2 0 and 1 S m
Up to
1, the values of the eigenfunctions...
...To obtain the limiting eigenfinctions, which we denote by \label [def:uminfty) ) (m( limo- Um 0z; one can use the fact that (mm))) - ntt for 1 1 m S n to obtain...
...the eigenvalues (4m(o) for 1 S m S ni all converge lan n²x² as o tends to infinity. (Note that this is consistent with our implicit expression for the eigenvalues (a) from Theorem ref tthm:main) .) From Corollary /ref(cornoode) this ensures that Um (1kk o) con- verges to zero as a tends to infinity. This means that Um (2; oo) (defined in \eqref( (efrumfdef is proportional..
Figure 5: In the case above, retrieved top-1 chunk introduces the definition used in the target chunk, while the adjacent retrieved 3rd best chunk and retrieved 4th best chunk both cover the same variants as those appear in the target chunk. Retrieved 2nd best chunk contains the theorem and corollary used in the query chunk.
denote the projection Arpuc > \epue. (begin \label (lemma:m circ pi=2n} The complex and real moment maps for GC are related by μ² o end (lemma) |begin Many of our computations...
x(W([U)) w(r(u)) where w(p) denotes the w-limit set of the negative gradient flow start- ing from (propp (etrrroooppbbiiippnnerree (proof Applying Lemma (eff(mmmmamm cire pi=2n] we have < grad||n/|| [[o]] [p) >= 4 ....
Figure 6: the case above, retrieved top-1 chunk introduces the lemma used in the target chunk.
16

---

Length-Generalizzbbe Attention via Causal Retrieval
They are not the same: see Section \ref(sec:15). To establish Theorem \ref(thm:ABn) it suffices to prove it for B(n); the estimate for A(n) then follows from the linear relation A(n) = log Gn + B(n) (from \eqref(eqm:GABx) combined with the asymptotic estimate for G(n) in \egref(epmlogg--ssymp)
The main contributionin the sum B(n) comes from those primes p having p vñi, whose key property
that expomential sum methods yield alternative unconditional estimates for A(n,z), B(n,w) and log G(n,)), which are nontrivial when x = o(m), and apply for x > vū. These estimates improve on the estimates of our main theorems for certain ranges
exponents Up (Gn) as a difference of quantitics given by statistics of the base p radix expansion of integers up to n. (see Theorem /reffthmeexplicit) ) Summing over p S yields a formula log G(n,z) - A(n,,)) - B(n,x) x) involving nonnegative arithmetiifunctionn A(n. x) and B(n,z)...
The implied constant in the O-notation does not depend on a. {end(thm} The limit func- tion fa(a)) is pictured in Figure [fg. The function lies strictly above the diagonal line B (1-7)a; note that in \eqref(eqmGAAB) in its relation to log G(n, it appears with negative sign, consistent with fa(a)...
Figure 7: In the case above, retrieved top-1 chunk and retrieved 3rd best chunk adjacent, which mentions the same equation as target chunk. retrieved 2nd best chunk and retrieved 4th best chunk both mention log G(n, 2) which also appears in target chunk.
D. The details for the NIAH test
In all evaluations conducted for the NIAH tests, we fine-tune all models using checkpoints derived from PG19, employing the same set of synthetic data. The number of fine-tuning steps is set to one-tenth of the total steps used during pre-training. while all other hyperparameters are kept constant. Examples of the synthetic data utilized for each task presented in the table 5. Specificallyy, we pad the input tokens to ensure that the landmark token can be inserted before "is" in the question.
Task
Example
(essays)...
The passkey is: {tokens).
Single NIAH
What is the passkey? The passkey is {tokens}.
(essays)...
DEF ttokens 5)-> (tokens 6)
DEF tooens_2 23>>> 2
2-hop NIAH with noises
DEF ttokens _4 >ttokess
DEF [tokens_ |1-> 1ttokesess_) ...
The path from {tokens_1] | is: [tokens 2}, {tokens.3)
Table 5: Task examples for the two NIAH tests.
E. Discussions about the RPT RPTCCottiver
Contriver baseline
In the original RPT, reference LM used pre-prepare target chunks for each chunk, as discussed in the related works. Compared to using Contriever as the retrieval module, the original method offers stronger causal retrieval capabilities. However, since the code for retriever distillation in the original RPT is not released and the approach is costly and less flexible, we opt to use Contriever instead. RPT can be considered a fusion of RETRO and RPT. It retrieves past chunks in a manner similar to RETRO and integrates the retrieved information in the style of RPT. The retrieval process involves dividing every 64 tokens into a chunk, encoding them with Contriever to obtain chunk representations, and then retrieving past chunks based on cosine similarity with the current chunk.
17


Reviews:
[{"review_id": "7glpSXvRyz", "reviewer": "Reviewer_gMz8", "summary": "This paper introduces grouped cross attention (GCA), where the model learns to retrieve past chunks of tokens to reduce the prediction error on future tokens.  The model is trained end-to-end to retrieve relevant chunks, and thus does not depend on a fixed retriever.  The GCA modules are appended after attention modules in a transformer architecture.  They offload the hidden states of past chunks to CPU to save memory, which enables them to significantly reduce memory overhead relative to competing methods and enable large contexts.  Their throughput is similar to a simple sliding window attention baseline.  They evaluate, PG19, arXiv math, single key NIAH, and summarization tasks.  They compare against baselines RPT, landmark attention, and Block recurrent TFM.\n\n## Update after rebuttal\n\nThe authors addressed my concerns during the rebuttal stage.  I maintain my score and positive assessment of the paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your professional and constructive review.\n\n**W1. They only evaluate on single NIAH/It would help to see evals on the full RULER suite including multi-key retrieval,**\n\nBesides the single NIAH, we also evaluated the variable tracking task in Table 2, which is the most challenging and comprehensive among RULER NIAH tasks. Our setup includes 6 variables and 2-hop assignments, covering both multi-key and multi-hop scenarios. Results confirm that our length generalization capability remains effective even under these complex conditions.\n\nIn addition, we are also willing to report the results of multi-key passkey retrieval, which contains 12 keys.\n| | 64K | 256K | 1M | 4M\n|-|-|-|-|-|\nTRMsw | 1.2 | 0.0 | - | - |\nLMK | 0.0 | - | - | - |\nDRTretreivalx1| 85.71 | 88.24 | 87.5 | 87.50\n\n**Q3//W2/W4. Can you report throughput, memory overhead, and scaling curves for larger models.**\n\nSure, here are the results:\n\n| Model | Throughput (350M, 32K) | Memory（350M) | Throughput (760M, 16K) | Memory(760M)  | Throughput (1.5B, 8K)     | Memory(1.5B)  | Throughput (3B, 4K)     | Memory (3B)     |\n|-|-|-|-|-|-|-|-|-|\n| Transformers$_\\text{SW}$  | 2.91e5 tokens/s | 50G | 1.31e5 tokens/s | 50G | 6.56e4 tokens/s | 50G  | 3.41e4 tokens/s | 72G  |\n| Transformers$_\\text{fullattn}$| 4.44e4 tokens/s | 50G | 4.00e4 tokens/s | 50G | 3.50e4 tokens/s | 50G | 2.85e4 tokens/s  | 72G  |\n| BRT | 7.94e4 tokens/s | 59G  | 5.70e4 tokens/s | 56G  | 4.02e4 tokens/s | 59G | -  | OOM |\n| RPT$_{\\text{Our impl}}$| 1.13e4 tokens/s | 57G | 5.17e4 tokens/s | 59G | 2.73e4 tokens/s | 64G | 1.45e4 tokens/s | 80G |\n| DRT$retrieval \\times 1$| 2.32e5 tokens/s  | 56G | 1.06e5 tokens/s | 59G | 5.52e4 tokens/s | 63G | 2.93e4 tokens/s | 80G |\n| DRT$retrieval \\times 2$| 2.28e5 tokens/s  | 57G | 1.05e5 tokens/s | 59G | 5.52e4 tokens/s | 63G  | 2.93e4 tokens/s | 81G |\n\nHere, we did not enable FSDP or gradient checkpointing. If enabled, the 3B model should achieve over 8K context length.\nCompared to full attention, our approach shows significant speed advantages at 8K and above. Compared to sliding window, the additional overhead is around 20\\%, but while gaining the ability to capture ultra-long-range information. Overall, DRT is scalable.\n\nIf we have the opportunity to submit a camera-ready version, we will include the corresponding scaling curve for this data.\n\n**Q1. Can you explain how you might select the groups for larger models and how that would affect your results?**\n\nBefore applying the softmax off-by-one in Equation(3), we observed that increasing the number of groups resulted in a slight increase in perplexity. However, after applying the softmax off-by-one, PPL generally decreases as the number of groups increases. We suppose that softmax off-by-one enables GCA to disregard noisy retrieved chunks, thereby enhancing performance.\nIncreasing groups essentially expands the attention fields because each retrieval may access different chunks. However, the marginal gain diminishes, and the memory footprint also rises with the increase in the groups. Therefore, for larger models, we would likely still set the group number between 1 and 2, balancing computational cost and performance.\n\n**Q2. Can you provide a generic analysis of the time to offload the chunks to CPU based on the GPU to CPU communication speed to demonstrate scalability?**\n\nThe offloading time from GPU to CPU is proportional to the KV cache size, which is determined by (batch_size, chunk_num, chunk_size, hidden_size, groups). Note that this is independent of the number of layers, as the gca KV is shared within the same group. \n\nAssuming a large model with a hidden size of 4096 and bfloat16 precision (2 bytes per value), for batch size = 1, group = 1, chunk size = 64, and input length = 1M tokens, the total KV cache size is:\n1 * (1M // 64) * 64 * 4096 * 1 * 2(bf16) = 8192 MB = 8 GB\nWith typical GPU-to-CPU communication speeds exceeding 20 GB/s, all required KV caches for GCA can be offloaded in less than 0.4 seconds, which is fully scalable.\n\n**Missing Reference:**\n\nIf we have the opportunity, we will include these references in the camera-ready version.\n\nBy the way, NSA was not released on arXiv at the time of our submission. During the review period, we evaluated its extrapolation capabilities for the S-NIAH task based on (https://github.com/fla-org/native-sparse-attention) and found that its extrapolation ability is still limited. The results are as follows:\n\n| | 16K | 64K | 256K | 1M\n|-|-|-|-|-|\n|NSA | 98.3 | 50.8 | 18.4 | oom |\n\nThis result shows that GCA still exhibits significant advantages in terms of length generalization.\n\nThank you once again for your thoughtful and professional review. Your comments have been invaluable in helping us refine our work and improve our paper.\nWe sincerely hope that our responses have addressed your concerns, and we would be truly grateful if you could kindly offer us more support.", "claims_and_evidence": "Claims:\n* \"GCA is the first attention mechanism that can achieve perfect passkey retrieval with 16M context length, 1000× the pre-training length\"\n* \"DRT significantly outperforms all baselines with comparable pre-training costs and much lower\ninference costs\"\n* In discussion: \"GCA can extrapolate over 1000 times with stable PPL and perfect passkey retrieval, but Landmark Attn\nfails\". \n\nEvidence:\n* This claim is verified in Table 2 where they get 100% on single NIAH at 16M context length.\n* The perplexity evals are verified in Table 1 where they beat the baselines by a small margin for the large contexts.  In Table 2 they significantly beat the baselines on NIAH.  The throughput and memory overhead is documented in Figure 3 and Table 3.  Table 3 demonstrates they have significantly smaller memory overhead compared to landmark attention and have similar throughput to the sliding window baseline.  It would help to show the throughput and memory overhead of the other baselines RPT, and block recurrent TFM.\n* Extrapolation to 1000x is tested on NIAH but not on PPL (which is only evaluated up to 32K), can you please clarify or correct this?", "methods_and_evaluation_criteria": "The evaluation on RULER is pertinent for long-context, and language modeling on PG19 is standard for such works.\n\nIt would help to see evals on the full RULER suite including multi-key retrieval, as they only evaluate on single NIAH.", "theoretical_claims": "This is an empirical work, there are no theoretical claims to justify.", "experimental_designs_or_analyses": "They train 128M and 350M TinyLlama based models on 32B tokens, which is past Chinchilla optimal and is thus reasonable.  It would be helpful to show results with larger models.  While training larger models on 32B tokens may take too long with your 8 A100 setup, you could test throughput, memory overhead and plot scaling curves for the initial set of tokens for larger models to show that your method is scalable.", "relation_to_broader_scientific_literature": "This work relates to broader literature on chunk retrieval and long-context language modeling.  They position themselves among related works such as RETRO, Landmark attention, RPT, Transformer-XL, Infini-attention,", "essential_references_not_discussed": "They seemed to have missed the reference [1], which is similar to Transformer-XL in that they segment the input into chunks and maintain summary tokens to summarize prior tokens.  I will also mention concurrent works [2], [3] for the authors benefit to add to the camera-ready, but I do not expect comparisons against these works as they are concurrent so I am not penalizing the authors for this in the review.  [2] appeared within 1 month of the submission deadline and [3] appeared after the submission deadline.  [2], [3], summarize prior chunks into summary tokens which are used to retrieve chunks which are then used for cross-attention.  [3] notably also introduces optimized kernels, similar to the present submission.\n\n[1] Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth\nGopal. Leave no context behind: Efficient infinite context transformers with infini-attention. arXiv preprint\narXiv:2404.07143, 2024\n\n[2] Elvis Nunez, Luca Zancato, Benjamin Bowman, Aditya Golatkar, Wei Xia, and Stefano Soatto. Expansion Span: Combining Fading Memory and Retrieval in Hybrid State Space Models. arXiv:2412.13328, 2024\n\n[3] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. arXiv preprint arXiv:2502.11089, 2025.", "other_strengths_and_weaknesses": "Strengths:\n* They achieve perfect recall on single NIAH up to 16M context length, while achieving reasonable perplexity up to 2x the pre-training context length\n* They provide hardware-aware psuedocode for their Triton implementation of FlashGCA forward and backward\n* The case studies on retrieving relevant llemas and definitions for math theorems on arXiv math is quite interesting and illuminating.\n\nWeaknesses:\n* They only evaluate on single NIAH\n* The models are still relatively small 128M and 350M.  It would help to have some experiments with larger models.  The authors state that they trained with 8 A100 GPUs in bf16 precision.  In bf16 with Adam you should be able to train up to around a 3B model without model parallel assuming 40GB of HBM for an A100.  Although, it may take a long time to train on the full 32B tokens.  Perhaps some smaller scale fine-tuning experiments or other simplifications could be made to get signal at larger scale.  In any case, you don't need a large number of tokens in order to time the throughput and calculate the memory overhead relative to other methods.  This should be included in the camera-ready.\n* Larger models will take up more GPU memory, so the 16M context result will likely not apply (as GPU memory will be exhausted quicker by the KV cache), however it is still a great proof-of-concept\n* They are lacking throughput and memory calculations for the other baselines RPT, Block recurrent TFM.  Adding some additional curves to Figure 3a would be helpful.", "other_comments_or_suggestions": "* It would be of great help to the community if you could provide the Triton code in the camera ready to supplement the pseudocode.\n* Can you report throughput, memory overhead, and scaling curves for larger models (even if you don't train on the full 32B tokens)\n\nSuggested edits:\n\nLine 212: \"the K, V **liner** transformations\" --> \"the K, V **linear** transformations\"\n\nLine 402: \"confirming our hypothesis that conducting causal retrieval every **G-layer** contributes to scenarios requiring multiple retrievals\" --> \"confirming our hypothesis that conducting causal retrieval every **G layers** contributes to scenarios requiring multiple retrievals\"", "questions_for_authors": "* It seems you only try setting the number of groups to 1 or 2.  Can you explain how you might select the groups for larger models and how that would affect your results?\n* How would your method scale to LLMs with billions of parameters?  Can you provide a generic analysis of the time to offload the chunks to CPU based on the GPU to CPU communication speed to demonstrate scalability?\n* Can you provide throughput, memory overhead, and scaling curves for larger models?", "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": "5oqs6XwoVu", "reviewer": "Reviewer_wqhZ", "summary": "The paper introduces a new attention mechanism to integrate dynamic context called Grouped cross attention (GCA). GCA helps maintaining long term dependencies during sequence generation enabling long range information access and length generalization. GCA integrates chunk to chunk retrieval to learn and retrieve past chunks to reduce autoregressive generation loss for next time step prediction in the decoding process. Based on GCA in transformers layers, the authors introduce differentiable retrieval based transformers (DRT) to enable pretraining across longer sequences and discuss hardware aware implementation. DRT showcases performance gains over very long range sequence generation via empirical evaluation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your review comments and support for this work!\n\nAlthough the improvement on perplexity is relatively marginal, the performance of passkey retrieval in Table 2 is quite significant, especially when the context length far exceeds that of the pre-training stage. This demonstrates extraordinary extrapolation capabilities, which none of the existing attention mechanisms possess.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "The work proposes the DRT model architecture using GCA attention mechanism. Experiments are conducted to compare DRT on various datasets and tasks and comparisons are made w.r.t SOTA literature in the field. Standard evaluation metrics corresponding to each dataset/task is used.", "theoretical_claims": "No", "experimental_designs_or_analyses": "DRT is compared to SOTA literature in the field and showcases marginal performance gains compared to existing literature across various datasets and sequence generation tasks. Ablation studies are performed to strengthen the claims.", "relation_to_broader_scientific_literature": "The authors present a novel attention mechanism to preserved information across long range sequences and improve language modeling empirical performance across various sequence generation tasks. The broader scientific community would benefit from the results and model architecture presented by the study.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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": "D6XNhX17Tk", "reviewer": "Reviewer_PLe3", "summary": "This paper proposes an attention mechanism, Grouped Cross-Attention (GCA), to improve long-context language modeling. By integrating a retrieval mechanism directly into the attention computation, GCA allows Transformers to generalize to significantly longer contexts while maintaining computational efficiency. The authors also introduce the Differentiable Retrieval-based Transformer (DRT), which utilizes GCA to retrieve and integrate relevant past chunks dynamically. Experimental results demonstrate that DRT achieves superior performance in long-context modeling tasks.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for reviewing our manuscript.\n\n**W2. For instance, the representation of chunk-wise CA outputs in Figure 2's caption does not match the notation used in Equation (2).**   \n\nIf the inconsistency refers to the subscripts, it can be easily fixed by rewriting $O_{t+1,k}^l$ to $O_{t+1}^{l,k}$. Besides, eq(2) is the batch version of Figure 2.\n\n**Cmt 1. Figure 2 is very confusing..**\n\nThanks for your comment.\n\nIn fact, using landmarks for retrieval can be regarded as a special case when group = 1. Note that $l_k$ is essentially $h_k^0$ passed through the encoder.\nTo make it easier to understand, we can additionally explain when G=1,  $r_t^k=l_t^\\top l_k / \\sqrt{d}$. As the retrieval is conducted only once and does not involve the use of landmark representations from different layers. \n\n**W3. Regarding parameter justification**\n\nAs hyperparameters cannot be exhaustively studied, we follow the setups from previous works ([1][2][3]), where a chunk size of 64 is commonly chosen, along with a 2^n chunk number. We set the chunk number to 8 for better alignment with landmark attention and sliding window, which can theoretically make the attention fields consistent as described in lines 312-319.\n\n**W4. Regarding the difference in perplexity between the PG19 dataset and the ArXiv-math**\n\nOur results are consistent with all previous paper works ([1][2][4][5]).\nThe ArXiv-math dataset is a collection of mathematical papers that involve a significant amount of logical symbolic reasoning. The patterns within this dataset are more structured and regular, leading to a lower perplexity.\n\n**W5(a) Since the GCA method integrates retrieval, comparing its perplexity with baseline methods might be unfair.**\n\nPlease note that our baselines, like RPT and Landmark Attention, also integrate retrievers into the training process, while Block Recurrent Transformer employs recurrent mechanisms to maintain long-range memory. They are designed for long-context. Thus, it's not unfair to compare with them. \n\nMoreover, we also evaluate downstream task performance in Table 2. \n\n**W5(b) It would be better to test and compare on more practical tasks (such as LongBench benchmark).**\n\nWe believe that LongBench is not well-suited for evaluating small models. \nIn fact, most results of small models on LongBench are even inferior to random guessing among four options. Here we report results on LongBench v2 for reference, following the evaluation method used in the cloze task presented in this work (https://arxiv.org/abs/2406.07887):\n\n|Model|Overall | Easy | Hard | Short | Medium | Long |\n|-|-|-|-|-|-|-|\n|TFM /sw | 24.0 | 22.9 | 23.2 | 27.2 | 20.0 | 26.9 |\n|Block Recurrent TFM | 22.8 | 20.8 | 25.1 | 19.4 | 22.8 | 25.9 |\n|Landmark | 24.3 | 23.4 | 25.7 | 28.3 | 19.1 | 25.0 |\n|RPT | 23.3 | 25.5 | 22.2 | 21.1 | 25.6 | 22.2 |\n|DRTretrieval×1 | 24.8 | 28.6 | 23.8 | 21.1 | 22.8 | 27.8 |\n|DRTretrieval×2| 25.6 | 24.0 | 25.7 | 27.2 | 25.1 | 25.9 |\n\n**Q1. Could you provide a more detailed explanation of C_k and l_k?**\n\nTo elaborate, each chunk contains S tokens and 1 landmark token. These tokens are fed into the bi-encoder, where the representations of the S tokens correspond to $C_k$, while the representation of the single landmark token corresponds to $l_k$.\n\n**Q2. Equation (2): What normalization method does Norm() represent?**\n\nBoth RMSNorm and LayerNorm are valid options. In our implementation, DRT is built upon LLama, where RMSNorm is employed.\n\n**Q3: Line 324: What is the \"off-the-shelf retriever\" used by the authors?**\n\nAt lines 323–324, it states: \"w/ contriever: We replace the relevance scores in GCA by using an off-the-shelf retriever.\" From this, the \"off-the-shelf retriever\" refers to Contriever.\n\n**Q4: The paper only tests cases where the group number (G) is 1 and 2. ... what is their impact on performance?**\n\nBefore applying the softmax off-by-one in Equation(3), we observed that increasing the number of groups resulted in a slight increase in perplexity. However, after applying the softmax off-by-one, PPL generally decreases as the number of groups increases. We suppose that softmax off-by-one enables GCA to disregard noisy retrieved chunks, thereby enhancing performance.\n\n**Q5: Figure 3: Which variant of the model is used here—DRTretrieval×1 or DRTretrieval×2?**\n\nWe use DRTretrieval×1 by default for all experiments.\n\nReferences:\n\n[1] LANDMARK https://arxiv.org/abs/2305.16300, \n\n[2] RPT https://aclanthology.org/2024.tacl-1.66/,\n\n[3] RETRO https://arxiv.org/abs/2112.04426\n\n[4] BRT https://arxiv.org/abs/2203.07852,\n\n[5] cepe https://arxiv.org/abs/2402.16617,\n\n\nIn summary, all clarity issues can be fixed with minor adjustments. We follow previous work for choosing hyperparameters, fairly compare with competitive baselines, and use standard evaluation metrics in the long-context area. We hope our replies address your concerns. We would greatly appreciate it if you could consider increasing the score.", "claims_and_evidence": "While the paper presents some evidence for its claims, these significant gaps in explanation, methodology, and evaluation make it difficult to conclude that the claims are fully supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "While the proposed method addresses a relevant problem and uses some standard evaluation approaches, there are significant gaps in the methodological clarity and evaluation comprehensiveness that limit the conclusiveness of the paper's findings. The evaluation would be stronger with more diverse metrics, standard benchmarks, and clearer explanation of parameter choices and performance variations across datasets.", "theoretical_claims": "There don't appear to be explicit theoretical proofs or claims in the paper. The paper primarily focused on introducing a new attention mechanism (GCA) and a model architecture (DRT) with empirical evaluations rather than presenting mathematical proofs or theoretical guarantees.", "experimental_designs_or_analyses": "There are significant concerns about the experimental methodology, particularly regarding parameter justification, comparison fairness, and incomplete exploration of model variations. These issues potentially undermine the strength and generalizability of the paper's findings, suggesting that while the core approach may be promising, the experimental validation falls short of convincingly demonstrating its effectiveness.", "relation_to_broader_scientific_literature": "The paper tackles the important challenge of extending Transformer models to handle longer contexts, which is a recognized limitation of standard attention mechanisms. This connects to the broader literature on efficient Transformers and long-context modeling.", "essential_references_not_discussed": "The reference is good.", "other_strengths_and_weaknesses": "Strength\n1. The paper introduces a useful and effective method for handling long-context sequences, which is a crucial problem in Transformer-based language modeling.\n2. GCA incorporates the relevance scores of retrieval chunks into the LM decoder, allowing the retriever to adaptively learn how to retrieve through training for predicting subsequent tokens. This manner for retrieval-based language models is novel and reasonable.\n﻿\nWeakness：\n1.Some annotations and explanations are not sufficiently clear, making the proposed method difficult to understand. See Questions 1, 2, and 3 for details.\n2.The notation in the equations is inconsistent. For instance, the representation of chunk-wise CA outputs in Figure 2 caption does not match the notation used in Equation (2). Please unify the notation for clarity.\n3. In practice, a chunk size of 64 and a retrieval number of 8 are chosen as parameters for model training in this paper. What is the rationale behind selecting these specific parameters? Are these preferred settings consistent across different tasks? Are there any ablation studies that address this issue?\n4. In Table 1, there is a noticeable difference in perplexity between the PG19 dataset and the ArXiv-math dataset. What accounts for this discrepancy?\n5. Since the GCA method integrates retriever into the training process of predicting subsequent tokens, comparing its perplexity with baseline methods might be unfair. It would be better to test and compare on more practical tasks (such as LongBench benchmark).", "other_comments_or_suggestions": "1. Figure 2 is very confusing. Based on the figure, it appears that chunk c_6 retrieves the top-k relevant past chunks for its next chunk using only landmark representation l_6 and previous chunks’ landmark representations encoded by the bi-directional Transformer encoder. However, according to Equation (1), the causal relevance score is computed using the landmark representation output by the previous decoder layer of the g-th group. The concept of grouping (g-th group) is only introduced in the upper decoder. This discrepancy makes Figure 2 very misleading, making the whole methodology difficult to understand. I strongly suggest that the authors provide further clarification on this aspect and improve the figure to enhance readability.", "questions_for_authors": "1. Line 163: Could you provide a more detailed explanation of C_k and l_k? The paper only gives a vague description. Based on my understanding: C_k represents the token representations of chunk k, encoded by the bi-directional Transformer encoder. L_k is the landmark representation of chunk k, which (if I haven’t missed anything) is only briefly mentioned in Figure 2 caption. Could you clarify if this interpretation is correct?\n2. Equation (2): What normalization method does Norm() represent? The paper does not specify this.\n3. Line 324: What is the \"off-the-shelf retriever\" used by the authors? Please specify.\n4. The paper only tests cases where the group number (G) is 1 and 2. Have you experimented with a wider range of group numbers? If so, what is their impact on performance?\n5. Figure 3: Which variant of the model is used here—DRTretrieval×1 or DRTretrieval×2? Please clarify.", "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:

Meta-review: This paper proposes Grouped Cross Attention (GCA), a novel attention mechanism that enables Transformers to generalize to context lengths 1000× longer than those seen during training while maintaining a constant attention window size. By learning to retrieve the most relevant past chunks end-to-end, GCA significantly reduces memory and computational costs while achieving strong long-range information access and near-perfect performance on passkey retrieval tasks with up to 16M token contexts.

In general, the idea of the paper is interesting and the experimental results support the claim. Nonetheless, the reviewers have a few concerns about the clarity of the proposed method’s presentation, the limited evaluation on small models and a single dataset, missing comparisons with similar or concurrent works, lack of throughput and memory benchmarks against baselines, and the absence of ablation studies and justification for key hyperparameter choices. During the rebuttal, the authors attempted to address the reviewers’ questions, but their responses, particularly to reviewer PLe3, did not fully resolve the concerns.
Decision: Accept (poster)