Scale-invariant attention
Authors: Ben Anson
Xi Wang
Laurence Aitchison
Venue: NeurIPS 2025 poster ()
Keywords: long context
LLM
Transformer
scale invariance
Primary area: deep_learning
TL;DR: We propose a scale-invariant attention mechanism for transformers and show it improves performance in out-of-distribution long-context settings.

Abstract: One persistent challenge in LLM research is the development of attention mechanisms that are able to generalise from training on shorter contexts to inference on longer contexts. We propose two conditions that we expect all effective long-context attention mechanisms to have: scale-invariant total attention, and scale-invariant attention sparsity. Under a Gaussian assumption, we show that a simple position-dependent transformation of the attention logits is sufficient for these conditions to hold. Experimentally we find that the resulting scale-invariant attention scheme gives considerable benefits in terms of validation loss when zero-shot generalising from training on short contexts to validation on longer contexts, and is effective at long-context retrieval.

Scale-invariant attention
Ben Anson
Xi Wang Department of Computer Science Johns Hopkins University xidulu@gmail com
School of Mathematics
University of Bristol
ben. www.n@bristol.sne... ac.
Laurenee Aitchison
School of Computer Science University of Bristol
laurence wwwchhion@brritttoo.......ooool. ac.
Abstract
One persistent challenge in LLM research is the development of attention mech- anisms that are able to generalise from training on shorter contexts to inference longer contexts. We propose two conditions that expect all effective long- context attention mechanisms to have: scale-invariant total attention, and scale- invariant attention sparsity. Under Gaussian assumption, show that simple deetidend transformation of the attention logits sufficient for these conditions hold. Experimentally tion scheme gives considerable benefits terms of validation loss when zero-shot generalising from training short contexts validation on longer contexts, and is effective at long-context retrieval.
find that the resulting scale-invarian atten-
Introduction
One key challenge in modern LLMs is scaling up context length at inference time, while maintain- ing model performance. We approach this question of length generalisation by considering scale invariance. In particular, we are inspired by the "scale-invariant" statistics of natural images (Van der Schaaf van Hateren, 1996). Scale invariance, for images, is actually highly intuitive, and means that there is structure at all spatial scales. For example, in an image there might be big features that are 100-1000 pixels across, and some small features that are only 1-10 pixels across. In natural images, features at both spatial scales important: you cannot remove features at either scale without radically altering the image (Van der Schaaf & van Hateren, 1996).
For attention over text, instead of pixels, we considered token ranges at different scales:
1-10 tokens in the past (e.g. those in the same sentence),
100 tokens in the past (e.g. those in the same paragraph),
1000-1000 tokens in the past (e.g. those in the same section),
1,000-10,000 tokens in the past (e.g. those in the same document), and
With this in mind, we define scale-invariant attention as any attention scheme that satisfies two desiderata: scale-invariant total attention and scale-invariant attention sparsity.
Scale-invariant total attention the property that the sum of attention weights in each of the above ranges is roughly similar. Intuitively, that means that the model attends to both the local context (c.g. 10-100 tokens ago) while the same time taking account of information from the global context (e.g. 1.00011,0000 tokens ago). Scale-invariant total attention addresses a key issue with attention mechanisms: as the context gets longer, models tend pay more attention distant tokens the
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Attention
Global Entropy
Entropy of bins
latest 100 tokens
LogN
10-1
IID Geussian
10
Scale-invariant
102
5
5
10-3
10-1
15
10
10
Real Logits
10t2
5
5
10-3
10
10?
10
10
10
10c
[10", .101
[10', [10',
110"
.10)
.10)
.109)
105)
10a
10a
10
10
20
[10",
[100
[10*
[10%.
Context size
Context size
Figure 1: Scale-invariant attention controls the entropy without sacrificing attention over the local context. We consider three metrics for attention schemes: (left) the global attention entropy, (middle) entropy within particular ranges of tokens (e.g. (0-100), and (right) total attention the previous 100 tokens. The top row uses IID Gaussian logits, following our theoretical approach in Sec. 3.2. For LogN, the IID logits are multiplied by log N, where is the sequence length and = 0.4. The bottom row uses attention logits sampled from models trained with y-RoPE and 'No scale', LogN, and our scale-invariant transform. With no logit scaling, the attention becomes increasingly diffuuse the context grows (i.e. the distribution over logits has high entropy). LogN scaling reduces the entropy and thus ensures that attention remains sparse even at longer contexts. However, LogN still forfeits the ability attend to the local context (e.g. 100 most recent) tokens. Scale-invariant attention strikes a balance between low entropy and attending over the local context.
expense of the local context (e.g. Fig. 1, right column; the blue and orange lines decay quickly to zero). While scale-invariant total attention doesn't 't entirely eliminate that issue, it does ensure that the attention paid the local context shrinks only very slowly as the context length grows (e.g. Fig. right column; the black lines decay to zero much more slowly).
Total attention tells us the amount of attention certain region receives. However, scale invariance of the total attention does not tell us about how the attention will be distributed among tokens in the region, i.e. whether the attention is spread out among many tokens concentrated onto only few tokens. As the region gets wider (e.g. from 10-100 tokens ago to 10,000-100,000 tokens ago), we might expect attention to spread out over more tokens simply due to the increased number of tokens. This "spreading out" of attention is possibly suboptimal for large contexts (Nakanishi, 2025), and that instead, we should focus attention on only a few of the most relevant tokens. To measure these effects, we use the entropy, which roughly captures the logarithm of the number of tokens attended to.
Scale-invariant attention sparsity captures this notion. In particular, define two kinds of scale- invariant attention sparsity. Strong scale-invariant attention sparsity implies that the number of tokens attended to in each region is constant. For example, if the model attends to 8 tokens in the 10-100 token range, strong scale-invariant attention sparsity says it will also attend to approximately tokens in the 1,000-10,000 token range. Strong scale-invariant attention sparsity implies an extreme increase in sparsity as the context gets longer that may be difficult to achieve with practical attention mechanisms. We therefore also introduce weak scale-invariant attention sparsity, which simply states that the sparsity increases as the context gets longer (i.e. attention is relatively dense in the region from 10-100 tokens, and much sparser in the region from 1,000-10,000, but you still attend to more tokens in the 1,000-10,000 region than the 10-106 100 region).
Our main contributions are:
We introduce the concepts of scale-invariant total attention, and weak and strong scale-invariant attention sparsity desirable properties when attending over long contexts.

---

We derive a simple, scale-ine invee-invariant total attention, and empirically satisfies weak scale-invariant attention sparsity for Gaussian logits.
transformation of attention logits that provably satisfies
We implement scale-invariant attention in conjunction with p-RoPE (Barbero et al., 2024b). We show that our method, 'scale-invariant y-RoPE', exhibits improvements in validation loss both when doing long-context training, and when zero-shot generalising to longer contexts. Our method also matches the performance of the best alternatives in an out-of-distribttion long context 'needle- un-a-haystack task.
2 Related work
Handling long sequences effectively in transformer-based models remains a significant challenge and is of high interest to the deep learning community (Bai et al., 2024; Ye et al., 2024; Jin et al., 2024; Beltagy al., 2020; Ding al., 2023; Munkhdalai et al., 2024; Bulatov et al., 2024; Liu et al., 2023a; Barbero et al., 2024a; Hu et al., 2024). Below, we discuss several strategies in this literature relating work.
Long contexts via the positional encoding: Methods like ALiBi (Press et al., 2021) introduce a static, causal bias directly into the attention logits. ALiBi endows the model with an inductive bias towards recent tokens (Kazemnejad et al., 2023), thus helping with longer contexts. Rotary Position Embeddings (RoPE) (Su et al., 2024) have emerged the most popular position encoding scheme, encoding relative positions. RoPE does not generalise to longer sequences out of the box (Peng et al., 2023), leading to subsequent research improving RoPE's length extrapolation capabilities (Zhu et al., 2023; Wang et al., 2024). Many achieve this by modifying RoPE's frequency spectrum; for example, Positional Interpolation (PI) (Chen et al., 2023), NTK-aware scaling (bloc97, 2023), YaRN (Peng et al., 2023), and by simply increasing the RoPE base 9 parameter (Grattafiori et al., 2024; Gemma Team, 2025; Barbero al., 2024b; Liu al., 2023b).
Entropy Control: Beyond positional information, the properties of the attention distribution itself are crucial, especially in long contexts where attention scores can "smear"/"spread out" across many (SSMax). also known as the 'LogN trick' (Nakanishi, 2025; Chiang & Cholak, 2022; Jianlin, 2021; Bai al., 2023; Llama 4 Team, 2025), addresses this by multiplying the attention logits by s log N. where N is the context length and s is a Icarned scale parameter. This multiplier has the effect of sharpening/focusin the attention distribution. Li et al. (2025) adopt similar approach: controlling the entropy of the attention distribution for better length generalization.
tokens.
These prior works are perhaps the most similar to our approach. However, the key issue with such approaches is that they treat the local context (e.g. previous 10-100 tokens) in the same way as the global context (e.g. 10,000 tokens ago). As the context gets larger, the attention paid to the 100 most recent tokens drops rapidly for the prior approaches, but stays markedly more consistent with inter-invrriant attention (c.g. for LogN specifying how we wanted attention to behave in the local and global contexts (Sec. 3.1) by giving the scale-invariaa total attention and scale-invarian attention sparsity desiderata. This means that e.g. multiplicative bias, whereas our approach has position-dependent
Fig. 1). In contrast, we started off by carefully
LogN has a multiplicative and additive biases for the logits.
Efficient Long-Context Training and Inference: A key observation of Xiong et al. (2023) that continual pretraining on long contexts, after initial pretraining shorter sequences, is often sufficient and much more computationally efficient. Thus, simple strategy for improving long-context performance, given sufficient computational budget, is continual pretraining on longer contexts. This has been adopted widely (Grattafiori ct al., 2024; Gao et al., 2024; Licber et al., 2024; Yang et al., 2024; Cohere Team, 2025; Llama 4 Team, 2025; Liu et al., 2024a). Of course, this strategy considerably increases the complexity and memory cost of the pretraining pipeline, making approaches like ours that can zero-shot generalise very valuable. Furthermore, one might expect long-context training to be far easier (e.g. requiring fewer long-contex training steps for optimal performance) for approaches that already have good long-context performance due to zero-shot generalisation.
For inference on sequences exceeding the trained context length, alternative strategies bypass ap- plying dense attention over the entire context, allowing for 'infinite attention' (Munkhdalai et al., 2024; Liu et al., 2024b; Martins et al., 2021; Chen et al., 2025; Ding et al., 2023). These include maintaining fixed-size attention window (Beltagy et al., 2020), retrieving relevant context tokens
3

---

using heuristics like top-K proximity (Han et al., 2023), or vector similarity search akin to episodic memory systems (Fountas et al., 2024; Xiao et al., 2024). While effective, these approaches operate at higher level, managing context rather than modifying the core attention mechanism's ability to process it directly, which the focus of our work.
Methods
Following FlexAttention (Dong et al., 2024), we define the attention "score" as the dot product of query q and keys K,
S.
QA Kex
Vd
Here, d is the head dimension and À E {I,...,d)) d} indexes the feature, and f indexes the position. We are using unusual form for the sequence indices, but this simplifies notation later. Specifically, we consider a single, fixed query token 9. Then, 2 0 into the keys/values from previous tokens. Critically, t counts backwards from the query token, e.g. t indicates the previous token. Logits are computed by applying an attention modifier function to the score,
L t(
(2)
Here, L is the actual value of the attention logits for the token 1 steps back from the query, while T(S;;; 1) the function used to compute those logits. Thus, I.(S;:t) S, recovers standard attention.
For use later, we define the unnormalised attention weights A. (normalised) attention weights A., and normalisers Z, for a sequence of length T, as,
A: A
T
Å. = exp (Le)
Z
A
3.1 Formal definitions of scale-invariant total attention and attention sparsity
Scale-invariant total attention. In our examples we have considered ranges of tokens, 10-100, 100-1,000, 1,000-10,000, 10,000-100,000, etc.. Scale-invariant total attention is the property that the total attention in each of these ranges is somewhat similar, such that the attention allotted to any one range does not dominate the others. We give formal definition in Def. 3.1, where using 10 gives the ranges from the examples.
Definition 3.1 (Scale-invariant total attention). Consider set of random variables, (Le), represent- ing attention logits. The total unnormalised attention in range t to 12 is,
12-1
1--1
Z
Σ
A
exp(L.).
t=
We say that the total attention is inver-invariant if. for any integer A > 1, the expected total unnormalised attention the range {L, 1,.... ,1 1} is asymptotically constant. That is,
E
1 (1)
t 0c.
(5)
The "0(1)" is big-O notation, and means there exist constants C1, C2 and to such that for all 1 >> l. ci < E [2,A] C2.
Scale-invariant unnormalised attention sparsity. Remember that means that attention should be denser for the local context (c.g. 10-100 tokens ago) and sparser for the global context (c.g. 1,000-10,000) tokens ago. To measure attention sparsity, we consider the number of tokens attended to in a region. To evaluate the number of tokens attended to, one approach to use the entropy over tokens. We measure the sparsity in the region from t to f2, using the entropy for the distribution over tokens in this region,
attention sparsity
4-1
A
A
Σ
2-- D
Áş A.
log zt2 E -1 A.
H =
log
+
(6)
ztz
Z
Z
z

---

Expected entropy
Square expected entropy
1-
30
A-2
ElH
20
10
logutti
logult)
Figure 2: Expected entropy of scale-invariant attention at different scales is sub-logarithmic. Here, we sample sequences of independent standard Gaussian logits, and apply the scale-invariant attention transformation. We estimate the expected entropy in ranges [4, (A). where the size of the range is controlled by (x-axis) and A (line color). We see that this expected entropy measure scales sub-logarithmicaly (left), and with the right plot suggesting a ~ (t) scaling. The dashed lines show best linear fit.
where A, is defined in Eq. (3). Defining the unnormalised negentropy as,
4-1
Σ
A log As
and remembering the definition Z (Eq. 4), we can then write the entropy in range 1 to t2 as,
Nota
- log 23
(8)
Z
Thus, it seems reasonable to consider the behavior of the unnormalised negentropy (Eq. 7), because if Aif and Z are asymptotically constant as f - co, then by Eq. (8) we expect H! to also be asymptotically constant. Thus we define:
Definition 3.2 Scalc-invariaat unnormalised attention sparsity). Consider set of random variables, {Le). representing attention logits. We say that we have scale-invariant unnormalised attention
sparsity if. for any integer A >> 0,
F NMA
O(1)
- 00.
where NIA is the unnormalised negentropy (Eq. 7).
Weak and strong scale-invariant attention sparsity. While the argument above suggests that scale-invariant unnormalised attention sparsity is an important property, ultimately are interested in giving formal definitions of weak and strong attention sparsity.
We define weak from 10-100 100-1,00 1,000 tokens the number of attended tokens grows sublinearly. contrast, standard attention with unscaled logits yields linear growth. Since entropy is roughly the log of the number of attended tokens (a uniform distribution over & tokens has entropy equal to log k), weak sparsity requires sublinear growth in entropy with respect to log(t) (see Definition 3.3).
attention sparsity such that as input lengths increase - for example,
Definition 3.3 (Weak scale-invariant attention sparsity). Consider a set of random variables, {Le}. representing atlention logits. We say that the atlention sparsity is weakly scale-invarians if. for any integer
E [H(A]] o(log
+
(10)
Remember o(log(t)) is 'little-o' notation which means that E [H[A] scales strictly slower than log(!). Strong scale-invariant attention sparsity implies that the number tokens attended to asymptotically constant as go from e.g. the past 10-100 tokens to the past 1,000-10,000.
5

---

Definition 3.4 (Strong representing attention logits. We say that the attention sparsity is strongly seale-invariant if. for any integer A 1,
attention sparsity). Consider a set of random variables, { b
E - as 1
(11)
Le. we expect H to be asymptotically constant as t -> 00.
3.2 What characteristics of the logits are required for scale-invariant attention?
Next, we ask what properties would be sufficient for scale-invariant total attention and for some form of ccale-invariant attention sparsity. Lemma gives ocale-invariant total attention, and Lemma 2 gives scale-invariant unnormalised attention (see Appendix for the proofs).
Lemma 1. Consider set of random variables, {Li}. representing attention logits. Let >> 0 be lengthscale parameter, and a >> a multiplicative constant. If the attention logits satisfy,
E
A
(12)
+/T+1'
then we have scale-invariant total attention (Def. 3.1).
Lemma 2. Consider a set of random variables, {Ia}, representing attention logits. Let T > 0 be lengthscale parameter, and B > 0 multiplicativee constant. If the attention logits satisfy.
log log A.]
3
(13)
t/r+1'
then have scale-invariant unnormalised attention sparsity (Def. 3.2).
To construct an attention mechanism that satisfies Eq. (12) and Eq. (13), we consider a simplified setting in which the logits are marginally Gaussian and arise from taking Gaussian "base logits", L. ~ N (0, 1), and transforming them by multiplying by a and adding a bias, m,,
L a, In + m. N of).
(14)
Our goal is to find m, and o such that scale-invariant total attention and scale-invariant unnormalised attention sparsity hold. In particular that requires.
I ++ E Ag
(15a)
8
E Ae log log Ao] (m, + at am)cm
(15b)
1 +1
where A, = exp (L.). Solving for m, and a,, we have (see Appendix C for details),
a 2 [log(t/T + 1) loga 3/a]].
(16a)
m4 - 3/a.
(16b)
For this solution to be valid, we only require / 2 a log since log(t/r + 1) 2 O when f 2 0. We formally summarise the above results in Theorem I (proof in Appendix F), which tells that this approach does indeed give scale-invariant total attention and scale-invariant unnormalised attention sparsity.
Theorem 1. Suppose attention logits } are marginally Gaussian with mean my and standard deviation as defined by Eq. (16). Assuming a., 7, 0, 3 2 a log a, then we have ccale-invarian total attention and scale-ine inle-invvriant unnormalised attention sparsity.
We therefore propose to scale the logits in real attention using a and m, defined in Eq. (16). As t increases, the variance, af increases as the logarithm of t, while the mean decreases the logarithm t.
Finally, note that while we have proven that we have scale-invariant unnormalised attention with this choice of a, and m, we have not proven that we have strong or weak scale-invariant attention. We therefore checked empirically whether using IID Gaussian logits, scaled by a, and m; gave weak or

---

RoPE
P-RoPE
-LogN=p-RoPE LogNN-NTK -ncale-invrrint P-RoPE (ours)
Train @4K/ Val @44
Train @16k Val @16k
Train @64k/ Val @64k
4000
4500
4000
4500
4000
4500
Step
Step
Step
(a) Scale-invariant attention improves language modelling over range of training context lengths.
Train @4k Val @16k
Train @4k / Val @64k
936k g4ke Vd44k
5.5
3.40
Zoomee Out
Zoomed In
Val Loss
va Loss
3.35
3.30
4.0
3.25
3.5
1000
2000
3000
4000
1000
2000
3000
4000
4000
4500
4000
4500
Step
Step
Step
Step
(b) Scale- invariant attention enables zero-shot long-context generalization of at least 16x.
Figure 3: Validation losses throughout training of 162M parameter GPT-2-like model with different attention mechanisms (our scale-invariant scheme shown in black). NoPE is omitted all but top-left and bottom-left panels to avoid excessive zooming out (due to high loss). The gray dashed line shows baseline of 3.28. The validation loss for many methods increases unison steps ~ 1500 and 3225 (see (b), left) despite aggregating over seeds; this is due to fixed training and validation data ordering.
strong ocale-invariag attention entropy (Fig. 2). We find that H appears to scale with Vlog(t), and hence seems to satisfy weak, but not strong scale-invan invariant attention sparsity.
Hyperparameters. Introducing and my (Eq. 16), appears to have introduced three additional hyperparameters tune. We reduce this one additional hyperparameter, the lengthscale т, by specifying a boundary condition - a3 = 1 and my 0. This boundary condition corresponds to not changing the scale of the local tokens. Substituting into Eq. (16b) gives - +/3/aa. 3//a.. Similarly Eq. (16a) requires that 1 = - log + 28/a. Putting these together, we obtain
0.5
That leaves us with the lengthscale as the only hyperparameter. Note that when is small relative neither af nor change much. Therefore, the timescale sets the size of local region in which attention approximately unscaled. Intuitively, therefore makes sense to choose somewhere in the region of 1-100 tokens. In Appendix H we tried 7 E (10 20-2, 10-1 10 10 102) and find that 10 performs best in practice.
Experiments
In this section, we compare scale-invand invariant attention with other dense attention methods including Dynamic NTK interpolation (RoPE+NTK) (bloc97, 2023), LogN scaling/SSMax (Nakanishi, 2025), P-RoPE (Barbero et al., 2024b), and ALiBi (Press et al., 2021). Our results show that our method, scale-invariant R-RoPE, has uniformly lower validation loss at a variety of training lengths (4k, 16k, 64k). Additionally, scale-invariant R-RoPE demonstrates stronger zero-shot long-context generaliza- tion (e.g. 'Train @4k/Val 16k' and "Train @4k/Val @64k' settings) versus all other methods. Finally, scale-invariant RoRo,E, along with LogN, saturated" neturated"needll-iin -eeeeinnn--hasttcke task.

---

We pretrained GPT-2-style models (Radford et al., 2019) (with QK- -orm, ReLU activations, etc. (Jor- dan et al., 2024a)) from scratch on the Wen dataset (Penedo et al., 2024). using a fixed training data ordering and 10M token validation set. We trained linear layers with Muon (Jordan et al., 2024b), and remaining parameters with Adam. We implemented scale-invariant attention using Flex (Dong et al., 2024).
We trained Iwo models: one with 162M parameters and another with 304M. The 162M one were trained for 4578 steps on 2.4B tokens over range of context lengths (4k, 16k, 64k). The 304M param- eter models were trained only on the best genera-innrriiiingg methods, for 10.9k steps on 10B tokens, at 4k context length. We give further experimental details in Appendix G. For the 162M parameter model, we targeted a validation loss of 3.28. following Karpathy's 's GPT-2 reproduction (Karpathy. 2024: Jordan et al., 2024a), shown in the figures as horizontal grey dashed line.
Long-context performance. We examined in-distribution, long-context performance of the different attention methods by looking at the validation loss for the same context length used for training. Even in this in-distribution setting, scale-invariant p-RoPE shows strong improvements in validation loss at all training lengths, 4k, 16k, and 64k (Fig. 3a).
Length Generalization. We evaluate length generalization by measuring validation loss on long sequences (16k and 64k) when training on 4k tokens. The 'Train @4k/Val @64k' setting in particular represents considerable jump of 16x between train and validation. Table and Fig. 3b report results for the 162M model. In the left panel of the Figure, ALiBi, LogN, and our method substantially outperform other approaches in generalizing to longer sequences. The right panel zooms in and shows that scale-invariant p-RoPH achieves the strongest generalization overall.
In preliminary experiments we tried other scale- invariant methods. We found that the most obvious method, scale-invariant RoPE, did not generalise well to long contexts (see Appendix 1.3). The p-RoPE method is similar to RoPE but excludes www-freguencyhhihh-waeeee position embedding, possibly suggesting that low-frequency components in RoPE interfere with the scale-invariant transformation L, + Ly + m, LogN also demonstrates stronger performance when paired with p- p-RoPE rather than RoPE. We were surprised that RoPE+NTK struggled to generalise in the "Train @ 4k/ / Val @64k' setting, but we believe this can be explained by the training context size: "Train @1 16k / Val @64k' is much better for RoPF+NTT (see Fig. 7 in the Appendix).
components in the
Fig. 4 presents pretraining losses for ALiBi, LogN+p-RoPE, and our method on larger 304M model selected due to their performance on the 162M 'Train @4k/Val @ @64'' task. Scale-invariant R-RoPE maintains its advantage at this larger scale.
Needle in Haystack. The key benefit of scale-invariant attention is that it balances local and sparse global attention. As such, we might worry that long-context retrieval performance might suffer versus other approaches that do not have specific mechanisms to ensure that attention to the local context does not vanish.
To assess whether long-context information retrieval capabilities suffered, we fine-tuned models task (note that fine-tuning on this task is unusual, but we found that
on a prompting alone was not sufficient to perform this task, bigger models/more pretraining would be required). Needle-in-a-haystad (Kamradt, 2023) measures a model's ability to precisely retrieve specific details (needles) from a large body of text (the haystack). Our tasks constructs prompts by concatenating text samples from the C4 dataset (Roberts et al., 2019) and embedding "needles" of the form 'The special magic <city> number is <7 sT_digittmumber>" number>'. We insert three (rather than one) needles uniformly, at random, into each context for more signal per example, with each needle contributing separately to the overall accuracy. A successful retrieval requires the model to output both the city and the associated number correctly. We trained models on sequences of length 4k, and tested at 4k, 16k, and 64k.
Table 2 show that icale-invariant p-RoPE and LogN+p-Roo RoPE perform well, while almost all other methods fail almost completely at 64k context length. Thus, despite focusing more on local context, our method does not seem to have suffered in retrieval performance.
8

---

Table Final mean validation losses (#1 standard error across 3 seeds) for different methods when training with 4k context length on 162M parameter GPT-2-style model. The error bars are small due to consistent training data ordering, and fixed validation set.
Method
Val @ 4k
Val @ 16k
Val @ 64k
RoPE
3.261 1: 0.001
3,936 + 0.010
5.260 +. 0.014
P-RoPE
3,260 + 0.001
3,984 = 0.008
5.735 = 0.085
NoPE
3.397 = 0.000
6.430 0.059
8.125 0.062
RoPE+NTK
3.261 + 0.001
3,703 = 0.007
5,430 = 0.026
YaRN
3.261 + 0.000
3,958 + 0.018
5,353 + 0.065
LogN+RoPE
3.260 + 0.001
3.273 + 0.004
3.378 0.011
LogN+p-RoPE
3,256 + 0.001
3.262 + 0.002
3.317 0.005
LogN+NTK
3.261 + 0.001
3,272 - 0.002
3,394 0.025
ALiBi
3.281 + 0.001
3.272
- 0.001
3,270 0.000
Scale-invarian p-RoPE (ours)
3,244 1 0.001
3,235 1 0.001
3.247 + 0.001
Train @4k Val @4k
Train @4k / Val @16k
Train @4k/ Val @64k
3.3
L p- RoPE
ALBBI
3.2
Scale invarian P-RoPE (ours)
Loss
3.1
3,0
2.9
2000
4000
8000
2000
6000
10000
2000
4000
8000
Step
Step
Step
Figure 4: Validation losses throughout training for 304M parameter model.
Table 2: Mean validation accuracies the needle-in in an-a-haaa haystacc task, 1 standard error. Metrics were calculated over seeds, after 300 steps of fine-tuning.
Method
Val Acc @4k
Val Acc @ 16k
Val Acc @64k
RoPE
0.962 = 0.003
0.000 = 0.000
0.000 = 0.000
p-RoPE
0.966 = 0.001
0,250 = 0.020
0.000 0.000
NoPE
0.964 = 0.000
0.303 = 0.239
0.000 + 0.000
RoPE+NTK
0.962 + 0.001
0.217 + 0.078
0.000 + 0.000
YaRN
0.969 + 0.001
0.000 + 0.000
0.000 + 0.000
LogN+RoPE
0.965 + 0.002
0.276 + 0.010
0.064 0.006
LogN+p-RoPE
0.969 = 0.002
0.962 1 0.003
0.939 0.009
LogN+NTK
0.962 1 0,001
0.253 1 0.015
0.056 1 0.008
ALiBi
0.957 1 0,002
0.020 1 0.001
0.003 1 0.000
Scale-invariant p- RoPE (ours)
0.965 = 0.000
0.969 0.004
0.969 0.005
5 Limitations
In this work, we evaluated our methods by pretraining with 162M and 304M parameter models, and investigated 7B parameter models via continual pretraining in Appendix L.5. Ideally, we would have pretrained from scratch at the multi-billion parameter scale, in line with contemporary state-of-the-an LLMs. While compute resource constraints prevented this, our results, together with the natural theoretical approach, offer no indication that our conclusions would fail to generalise to larger, commercially deployed models.
We focused on scale-invariant R-RoPE with dense attention, but the extension to other settings is promising direction for future investigation.

---

Conclusions
have proposed two desirable properties attention mechanisms, "scale-invariant total attention" and "scale-invariant attention sparsity", and presented straightforward modification to attention logits that enables these properties practice, In our experiments, we found that our scale-invariant attention modification, especially when paired with p-RoPE, substantially improves long-context language modelling performance and gives zero-shot generalisation from training on short contexts to testing long contexts.
Acknowledgements
We sincerely thank the Engineering and Physical Sciences Research Council and the COMPASS CDT for funding Ben Anson at the University of Bristol. Many our experiments used computational resources at the Advanced Computing Research Centre at the University of Bristol, with GPUs generously funded by Dr. Stewart.
References
Bai, J., Bai, S., Chu, Y., Cui, Z., Dang. K., Deng, X., Fan, Y., Ge, W., Han, Y., Huang, E., et al. Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023.
Bai, Y., Lv, X., Zhang, J., He, Y., Qi, J., Hou, L., Tang, J., Dong, Y., and Li, J. Longalign: A recipe for long context alignment of large language models. arXiv preprint arXiv:2401.18058, 2024.
Barbero, F. Banino, A., Kapturowski, S., Kumaran, D., Madeira Araújo, J., Vitvitskyi, O., Pascanu, R., and Veličković, P. Transformers need glasses! information over-squashing in language tasks. Advances in Neural Information Processing Systems, 37:98111 11119492 2024a.
Barbero, F., Vitvitskyi, A., Perivolaropoulos, C., Pascanu, R., and Veličković, P. Round and round we go! what makes rotary positional encodings useful? arXiv preprint arXiv:2410.06205. 2024b.
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150. 2020.
bloc97. Ntk-aware scaled rope allows llama models to have extended (8k+) con- text size without any fine-tuning and minimal perplexity degradation. Reddit, 2023. URL larp__laammmmmmoaa wwwwwoodoos...m tacosssttt_haee_/__o
rope_all aewe_allows..o
Bulatov, A., Kuratov, Y., Kapushev, Y., and Burtsev, M. Beyond attention: breaking the limits of transformer context length with recurrent memory. In Proceedings of the AAAI Conference Artificial Intelligence, 2024.
Chen, J., Peng, S., Luo, D., Yang, F., Wu, R., Li, F., and Chen, X. Edgeinfinite: A memory--fficiintt infinite-context transformer for edge devices. arXiv preprint arXiv:2503.22196, 2025.
Chen, S., Wong. S., Chen, L., and Tian, Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.1559) 2023.
Chiang, D. and Cholak, P. Overcoming theoretical limitation of self-attention. arXiv preprint arXiv:2202. 12172, 2022.
Cohere Team. Command An enterprise-ready large language model. arXiv preprint arXiv:2504.00698, 2025.
Ding, J., Ma, S., Dong, I.., Zhang, X., Huang, S., Wang, W., Zheng, N., Wei, Longnet: Scaling transformens to 1,000,000,000 tokens. arXiv preprint arXiv:2307.024488 2023.
Dong, J., Feng, B., Guessous, D., Liang, Y., and He, H. Flex attention: A programming model for generating optimized attention kernels. arXiv preprint arXiv:2412.05-000 2024.
10

---

Fountas, Z., Benfeghoul, M. A., Oomerjee, A., Christopoulou, F., Lampouras, G., Bou-Ammar, H., and Wang. J. Human-like like episodic memory for infinite context llms. arXiv preprint arXiv:2407.09550, 2024.
Gao, T. Wettig, A., Yen, H., and Chen, D. How to train long-context language models (effectively). arXiv preprint urXi:22410.02660 2024.
Gemma Team. Gemma technical report. arXiv preprint arXiv:2503.19978 2025.
Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A.,
Schelten, A., Vaughan, A., et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
Gu, X., Pang, T., Du, C., Liu, Q.. Zhang, F., Du, C., Wang. Y., and Lin, M. When attention sink emerges in language models: An empirical view. arXiv preprint ar arXiv:2410.10781, 2024.
Han, C., Wang, Q., Peng. H., Xiong, W., Chen, Y., Ji, H., and Wang, S. Lm-infinite: Zero-shot extreme length generalization for large language models. arXiv preprint 2tt.:::33...1333.. 2023.
Hu, Z., Liu, Y., Zhao, J., Wang. S., Wang, Y., Shen, W., Gu, Q. Luu, A. T., Ng. S.-K., Jiang, Z., et al. Longrecipe: Recipe for efficient long context generalization in large language models. arXiv preprint arXiv:2409.00555 2024.
Jianlin, S. Attention's 'ss scale operation from the perspective of entropy invariance, 2021. URL https: //ww:. kexue fm/ www.avesciiveeeeeeeeee
Jin, H., Han, X., Yang. J., Jiang, Z., Liu, Z., Chang, C.-Y., Chen, H., and Hu, X. Llm maybe longlm: Self-extend Ilm context window without tuning. arXiv preprint arXiv:2401.01325, 2024.
Jordan, K., Bernstein, J.. Rappazzo, B., bsskysbarbbsksssciial.. Vlado, B., Jiacheng, Y., Cesista, F., Koszarsky, B., and @ Grad62304977. modded--anogpt: Speedrunning the nanogpt baseline, 2024a. URL https::////.. /gi thub. cwm/ http///lllerJddd/mdddddddo hwtp://mddddddddd nanogpt.
Jordan, K., Jin, Y., Boza, V., Jiacheng, Y., Cesista, F., Newhouse, L.,, and Bernstein, J. Muon: An optimizer for hidden layers neural networks, 2024b. URL httpsttp://lllllerorraaa tithbb.......t.......... io/posts/muon/.
Kamradt, G. Needle in haystack pressure testing 1lms. https://github. htwp:tthuu...mmgmmmmmatc. www.////gaa/////////mmaaa.a...aa
LLMTest LLMTest_Meedle www.eeeeeennaasstaadtt 2023.
Karpathy, A. Discussion #481. https://githi: 2024.
Kazemnejad, A., Padhi, I,, Natesan Ramamurthy, K., Das, P., and Reddy, S. The impact of positional encoding on length generalization in transformers. Advances Neural Information Processing Systems, 36:24892-24922 2023.
Li, K., Kong, Y., Xu, Y., Su, J., Huang, L., Zhang, R., and Zhou, F. Information entropy invariance: Enhancing length extrapolation in attention mechanisms, arXiv preprint arXiv:2501.08570 2025.
Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalee-Shwartz, S., al. Jamba: A hybrid transformer-mambe language model. arXiv preprint (rXiv:2403.19882 2024.
Liu, A., Feng, B., Xue, B., Wang, B., Wu. B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19937, 2024a.
Liu, H., Zaharia, M., and Abbeel, P. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01 2023a.
Liu, X., Yan, H., Zhang, S., An, C., Qiu, X., and Lin, D. Scaling laws of rope-based extrapolation. arXiv preprint arx.v:2310.05209, 2023b.
Liu, X., Li, R., Guo, Q. Liu, 7.,, Song, Y., Lv, K., Yan, H., Li, I.., Liu, Q., and Qiu, X. Reattention: Training free infinite context with finite attention scope. arXiv preprint arXiv:2407.15177 2024b.
11

---

Llama 4 Team. The Ilama 4 herd: The beginning of new era of natively multimodal ai inno- vation. https: / /aii meta. com/blog////oooooooooooocc..owoe 4e-mliimdall-itteeee cessed: 2025-04-08.
2025, Ac-
Martins, P. H., Marinho, Z., and Martins, A. F. oo-former: Infinite Memory Transformer. arXiv preprint arXiv:2109.003:0 2021.
Munkhdalai, T., Faruqui, M., and Gopal, Leave no context behind: Efficient infinite context transformers with infini-attention. arXiv preprint arXiv:2404.07143, 101, 2024.
Nakanishi, K. M. Scalable-softmax is superior for attention. arXiv preprint arXiv:2501. 2025.
Penedo, G., Kydlíček, H.. Lozhkov, A., Mitchell, M., Raffel, C. A., Von Werra, L.,, Wolf, T., et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:3081 30849, 2024.
Peng, B., Quesnelle, J., Fan, H., and Shippole, E. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0007,, 2023.
Press, O., Smith, N. A., and Lewis. M. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv: 108. /2409, 2021.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Roberts, A., Raffel, C., Lee, K., Matena, M., Shazeer, N., Liu, P. J., Narang, S., Li, W., and Zhou, Y. Exploring the limits of transfer learning with a unified text-to-text transformer. Google Research, 2019.
Su, J.. Ahmed, M., Lu, Pan, S., Bo. W., and Liu, Roformer: Enhanced transformer with rotary position embedding. Neurocomputing,
2024.
Team, G., Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L., Mesnard, T., Shahriari, B., Ramé, A., et al. Gemma 2: Improving open language models at practical size. arXiv preprint arXiv:2408.000 000 18, 2024.
Torchtune. Torchtune: Pytorch's finetuning library, April 2024. URL https://github. com/
pytorch/torchtune
Touvron, H., Martin, L.,, Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv 2023.
preprint arXiv:2307.009280
Van der Schaaf, v. A. and van Hateren, v. Modelling the power spectra of natural images: statisties and information. Vision research, (6(17))2755-2777) 2770, 1996.
Wang, S., Kobyzev, I., Lu, P.. Rezagholizadeh, M., and Liu, B. Resonance rope: Improving context length generalization of large language models. arXiv preprint 0tth:2203300007// 2024.
Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R. Funtowicz, M., et al. Transformers: Sffte-of-the art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pp. 38 45, 2020.
Wortsman, M., Liu, J., Xiao, L.,, Everett, K., Alemi, A., Adlam, B., Co-Reyes, J. D., Gur, I., Kumar, A., Novak, R., et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprint arXiv:2309.14433
2023.
Xiao, C., Zhang, P., Han, X., Xiao, G., Lin, Y., Zhang, 7., Liu, 7.,, Han, S., and Sun, M. Infllm: Unveiling the intrinsic capacity of Ilms for understanding extremely long sequences with training- free memory. arXiv e-prints, pp. arXiv-2402, 2024.
Xiong, W., Liu, J., Molybog, I., Zhang, H., Bhargava, P., Hou, R., Martin, L.., Rungta, R., Sankarara- man, K. A., Oguz, B., et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309, 16039, 2023,
12

---

Yang, A.. Yang, B., Zhang, B., Hui, B., Zheng, B., Yu. B., Li, C.,, Liu, D., Huang, F., Wei, H., et al. Qwen2. technical report. arXiv preprint arXiv:2412.15115, 2024.
Ye, J., Xu, H., Liu, H., Hu, A., Yan, M., Qian, Q. Zhang, J., Huang, F., and Zhou, mplug-ow13: Towards long imagensequence understanding in multi-modal large language models. arXiv preprint arXiv:2408. 04840, 2024.
Zhu, D., Yang, N., Wang, L.,, Song, Y., Wu, W., Wei, F., and Li, S. Pose: Efficient context window extension of Ilms via positional skip-wise training. arXiv preprint
2023.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction clearly state our contributions of the notion of scale-invariant attention, a method for implementing it in practice, alongside empirical evidence that it effective.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
See our Limitations section (Sec. 5)-
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiaatinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: All of our theoretical results (Lemmas 1, and 2, and Theorem 1) all clearly state assumptions, and provide proofs and auxiliary results in Appendices B, C, E,
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We give details of the experiments in Section 4, and in Appendix G, which enables reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
15

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide code in the supplementary materials along with instructions to reproduce the main results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr
(oglic/guides/odeeee
oeiiiooolloooooo icy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. cctpublic//aeccccccccccccccc /gwpi/es///////eeeeee /Cae//ooeubbiiicc Sesssuumiiiioos. ssionPolicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide extensive details in Section and Appendix G.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We provide mean validation results + I stderr for the results in Table 1 and Table 2, clearly stating the number of seeds. We do not provide error bars for the plots as that makes them too messy to easily interpreted.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We discuss resources required in Appendix G, under Section G.3.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: / httppssccccpppcccppepepppppssept cew.ieeeenneeeees
Answer: [Yes]
Justification: The authors have reviewed the NeurIPS Code of Ethics, and the research conducted in the paper conforms in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
The paper contains foundational research pertaining to the attention mechanism in neural networks. While this has the potential to improve the LLMs in the long-term, the guidelines suggest is not necessary speculate on such matters.
Justification:
Guidelines:
17

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not release any artifacts at high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We properly credit and reference all assets used in the reference section. We declare all licenses the Appendix.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
18

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. cow datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We do not crowdsource or perform research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification:
We do crowdsource or perform research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
19

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification: LLMs were not used for the core method development in this research, but they used for visualizing data (e.g. creating tables/figures).
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: / //neruurps.... for what should or should not be described.
20

---

A Relating the Indicies in the Paper to Standard Indices
The standard form for attention is,
Sº 1 d
QuKys
(17)
where À indexes the feature, i indexes the query (i.e. the token we're generating now) and j indexes the key (i.e. the token we're attending to). We fix i, and take i = i - j. so
S St-i-s
(18)
where St-inj is given in Eq. 1 in the main text. Then,
L L(S - -) = ((1;tt LI=i-j:
(19)
where Ti=ii- is given Eq. in the main text. Then the unnormalised attention weights, Ä normalised attention weights, A and normalisers, Z are,
Á = exp (L*).
(20)
where, Am=i- is given in Eq. 3 in the main text.
A A
(21)
2
where, Atzin is given in Eq. 3 in the main text.
Z' = A = Z.
(22)
1=1
where, Z is given in Eq. 3 in the main text.
B E x exp(aX)] where X is Gaussian
Assume X ~ ~N(,, then we can compute the expectation in terms of a moment of a Gaussian,
(a-w)
X[ ex
k
20
dx
(23)
2
La
e
202 dx
(24)
V2
k
20 22axt 2200 dx
(25)
2Tozz
0C
th 122-2(( 2(1 yo dx
(26)
V2
22n²
(-- (u-to 20
taj
dx
(27)
V2m22
2,2 op
A
(28)
where X ~ NNw'' 022) NNp ao2, 02). the case k we have,
E X exp(()) = (μ +
+ 02/2).
(29)
In the case a = k = 2, we have,
E [x exp((X)) ((u 20222 + 0²) exp(2 + 202).
(30)
21

---

C Deriving a and my in the Logit Transformation
From Eq. (15), we wish to find ag and my such that,
a
(31)
+ 1
3
(m+ + ot ot))
(32)
1 1
We begin by dividing Eq. (32) by Eq. (31),
8
m4 + dt
(33)
a
which can be rearranged to,
= 4g
2
(34)
and of my.
(35)
Now, taking the log of Eq. (31), we have,
+
1 20 - log ( 1) log
(36)
Define ft = log (# + 1) - log a. Then to solve for af; we substitute my from Eq. (34) into Eq. (36).
[
+ u =
(37)
a
8 1 2 =-fe
(38)
2
2 2
ft + Ø
(39)
To solve for my we substitute of from Eq. (35) into Eq. (36)
(2_
me + 3 a M3
= Ft.
(40)
m
B -fe
(41)
m. -2
Sr +
B 20
(42)
which can also be written as,
m, = -0 + a
(43)
Thus, we have,
af = 2 [log(t/+++)) + log a]
(44)
m, 2 [log(t/T+1) + 3/a log a| 8/a
(45)
2 (l) 3/a + 2 log a.
(46)
D Proofs of Lemmas 1 and 2
Lemma 1. Consider set of random variables, {L.). representing attention logits. Let >> 0 be lengthscale parameter, and a >> a multiplicative constant, If the attention logits satisfy,
E
A.
a
(12)
t/r +1
then have geale-invarian total attention (Def. 3.1).
22

---

Proof.
We want show that E
(1) as 1 - 20. By definition and linearity of expecta-
tion:
A-4
E
1
E
A
Σ
E
A1
(47)
t=
t=(1
Using the given condition E
A1
8/7+1
44-1
4
44-1+4
Σ
t/t+1
a
aT
Σ
1+T
1
a7
ΣE
Σ
K
(48)
t=4
k=1++
Using the standard integral bounds for the harmonic derived by comparison with the integral (see Appendix E):
11 +
hA-1+4 l+
11 11--1+7
In
Σ
In
t +
(49)
t
+
k=
As t1 - 00:
t +
4 1
A
(50)
1 +
& 1+
4
=
(51)
-
t
So, the logarithm terms in both the lower and upper bounds approach In(A). By the Squeeze Theorem:
44-1+T
Σ + In))
(52)
k=
Since a, т, and A are constants, In(A) a positive constant. Thus:
E 1 (n O(1)
(53)
This satisfies the condition for scale-invariant total attention expectation.
Lemma 2. Consider set of random variables, (La,, representing attention logits. Let 0 be lengthscale parameter, and B > 0 multiplicative constant. If the altention logits satisfy,
E Ag log A.]
3
(13)
t/
then we have scale-invariant unnormalised attention sparsity (Def. 3.2).
Proof. We want to show that E
[T
@(1) as t + cc.
By definition and linearity of
expectation:
4
44-1
E
A 41
E
Σ
A log A
Σ
E
Ae log Asl
(54)
t=
t=h
Using the given condition
Ao
Ag
9
1/7+1'
44-1
4
3
E
M4.4 11
Σ
BT
Σ tmti
(55)
1/r+1
L+T
This sum is exactly the same form as in the proof of Lemma 1, just with / instead of Following the same steps using the integral bounds for the harmonic sum derived in Appendix E:
44-1+1
Σ
In(A)
as 4 - 00
(56)
h
A=
Therefore, as 4 3c:
E
N
+ 3r In 9(1)
(57)
This satisfies the condition for scale-invariant unnormalised attention sparsity in expectation.
D
23

---

E Bounds for Harmonic Sums
For the proofs of Lemma and Lemma we need bounds on the partial harmonic sum Σ k where a = f1 + and = 4 - 1 + We can obtain these bounds by comparing the sum to the integral f(x) = 1/x.
Since f()) = 1/2 is decreasing function for x > O, we have,
b+1
dx
1
dx
(58)
a-
Evaluating the integrals gives,
b
+ 1 a
In
In
(59)
k=a
F Proof of Theorem 1
Theorem 1.
Suppose attention logits { are marginally Gaussian with mean my and standard deviation a, defined by Eq. (16). Assuming a, B, > 0, 3 2 a log a, then we have ccale-invariant total attention and scale-ine iale--nvvriint unnormalised attention sparsity.
Proof. We need to show that the given conditions are sufficient for cale-invariiat total attention (Definition 3.1) and scale-invariant unnormalised attention sparsity (Definition 3.2).
1. Scale-invariant total attention: We need to show that E
A
(1) 4 0o. By
Lemma 1, this holds if E [ 7/r-T Since Lt ((n,, a the unnormalised attention weight
A, e follows a log-normal distribution. The expectation of a log-normal variable e where
X N(,, 02) is euto"/ Therefore.
E Ao [[A]] =E [ =
(60)
Substituting the given expressions for m. ot 3/a and at 2[log(1/T + 1) llog a + 3/a]:
A
cl-o2tD/o)+a?7?
(61)
(62)
e hogsl//r4t)-logg--8//a--
(63)
e
(64)
log a-log(t/r++)
(65)
()R())))T))
(66)
(67)
t/r+1 +1
Since this condition matches the requirement of Lemma 1, scale-invariant total attention holds in expectation. The condition 3 2 a log a ensures at 2 0 for all 2 0.
2. Scale-invariant unnormalised attention sparsity: We need show that E [M O(1) as
_ We need to compute E Á, log A [. Using the formula for E [Xe where X (N,, from Appendix B, which
oc. By Lemma this holds if E A, log Ae
1/TFT
(u + o² 12
E L = (m+ + mm
(68)
24

---

Substitute m, -a + 3/0:
E Lie") (-ai +3/a + a am)emtt
(69)
= (BB/a)emetai??
(70)
From the previous step, we know emttu//2 1/TFI
Substituting this:
[ = (3/a) ( (1/T+1)
(71)
g
(72)
t/T
Since this condition matches the requirement of Lemma 2, scale-invariant unnormalised attention sparsity holds.
total attention and scale-invariant Therefore, under the given conditions, have both unnormalised attention sparsity.
G Further experimental details
We provide experimental details here, and we provide the code used in the supplementary materials.
G.1 Pretraining from scratch
Our base model is modded- (Jordan et al., 2024a) variant, which is similar to GPT- 2 (Radford et al., 2019), but has the following main differences from GPT-2: RMSNorm for layer normalization (applied before the attention and MLP blocks, well to the token embeddings and the final output layer), squared ReLU activations, and QK Normalization (RMSNorm on query and key projections). All models used vocabulary size of 50304, with text tokenized with the GPT-2 tokenizer (Wolf et al., 2020). We implemented scale-invariant attention and ALiBi using Atexnttention (Dong al., 2024).
162M parameter model. The smaller model had 12 layers, 768 hidden dimension, and 6 heads. We optimized embedding parameters using Adam with learning rate 0.3, 3 = (0.9, 0.95). We optimized linear layers with Muon, with no weight decay, 7 = 0.02 and momentum 0.95. For remaining parameters (unembedding and LogN scalings, if LogN trick was active), optimized with Adam, using 7 = 0.002, 3 (0.9, 0.95). We trained for 4578 steps. The batch size was 8 x 65536/ Lu. with more gradient accumulation for shorter training lengths. We scheduled the learning rate with a linear schedule for all parameters, with no warmup, constant learning rate for 3270 steps, and linear cooldown for the remaining 1308 steps. We vary the training context length when pretraining with this model, from 4096(4k), to 16384(16k), and 65536(64k).. We validate at 4k, 16k, and 64k sequence lengths every 125 steps.
304M parameter model. For the larger model, we used the same settings as the smaller model, but with the following changes. The larger model had 16 layers, 1024 hidden dimension, and 8 heads. We trained the model for 10900 steps, processing approximately 10B tokens. We trained for 2 accumulation steps
GPUs, with 28 sequences per batch. We scaled learning rates following uParrm (Wortsman et al., 2023), which involves multiplying learning rates of the linear layers by 768/1024, 1004, adjust for changing the model width. Learning rates of other layers (embedding, unembedding, and LogN scalings if active) were not changed. We used a cosine learning rate schedule, with no warmup, and minimum learning rate of ( (at the end of training). All runs the 304M parameter models were with 4096 training sequence length. We validated at 4k, 1 16k, 64k sequence lengths every 250 steps.
RoPE hyperparameters. We used a base of 10, 000 for ROPE, and an effective base of 1024 for the angular frequencies in P-RoPE. For RoPE+NTK scaling (note the scaling applies only when the inference sequence length is longer than the training sequence length), we scaled by (Tant/ Tay ) /d((d-2, where T,, is the training sequence length, Tint is the inference sequence length, and d is the RoPE head dimension (128 for both models).
25

---

Dataset. We used subsets of Web (Penedo et al., 2024) for pretraining from scratch: a 10B token subset for the 162M model, and 100B token subset for the 304M model (from which 10BB tokens were used for its training). We keep the training data ordering fixed, and we keep the validation set identical across all runs to reduce variance.
Seeds. We repeat the experiment for different seeds when training at 4k on the 162M parameter model. We train with seed at 4k on the 304M parameter model, and 16k/64k on the 304M parameter model due to compute limitations.
Compute. We trained the smaller (162M) models single A100 80G GPUs. We trained the 304M models on 4xH100 grace hopper nodes using distributed data parallelism.
G.2 Needle-in-a-haasttck
The Needle in Haystack experiments were conducted by fine-tuning the pre-trained 162M parameter models. We fine-tune with the same learning rate as above, but with 100 warmup, 100 constant, and 100 warmdown steps (decaying y = 0). We use the same optimizers as in the pretraining phase. We train for 300 steps on sequences of length 4096, and batch size 8 with 8 accumulation steps. We repeated with seeds.
The task is to generate responses of the form "city1=needle1 :ttp:eddde:eittyyteeeeeee (ttt::/eeeeelllee....e--eeeeee city3=needee""" where the cities and needles are embedded uniformly at random into samples from C4 in the form "The special magic <city> number is <7 ___iigtt____ T_digitt_umbbr>.. s_tu__uuurrr"""" ". We sample several times from the C4 dataset (concatenating samples) and remove tokens until have the necessary number of tokens - 4096(4k) for training, and 4096(4k), 16384(16k), 65536(64k) for validation. Only the expected response tokens included in the loss (the prompt/context tokens masked). We repeat the task for three different seeds. The validation accuracies presenting in Table are calculated by the proportion of times that cities and numbers are output correctly.
G.3 Resources required to reproduce experiments
To reproduce results, 80G GPUs are required. We used 80G A A100s, and 80G H100 grace hopper nodes. In terms of time taken to execute each experiment type, we give estimates of the resources required:
pretraining 162M parameter model w/o flex attention takes roughly 4/8/22 A100 hours 4k/16k/64k;
pretraining 162M parameter model with flex attention takes roughly 8/16/44 A.00 hours at 4k/16k/64k;
pretraining 304M parameter model takes roughly 36 4xH100 node hours; fine-tuning on needle-in-a-haystee task takes roughly A hours.
To obtain our results, each experiment executed several times. In particular, needle-in-al haa-a-haystack with 3 seeds per method, pretraining at 4k with 162M model is 3 seeds per positional encoding method. The remaining experiments are executed once per positional encoding method. This gives total of 100 4xH100 hours, 550 1xA100 hours. We estimate that very roughly that amount again was spent configuring the experiments correctly, and on preliminary/failed experiments.
H The optimal lengthscale, T, is around 10
By introducing scale-invariant attention, we introduce one extra lengthscale hyperparameter, т, which represents the size of the 'chunks' we attend over. To select we trained at 4k for 7 E (10-2, 1,-1, 100 100010 10 102), and compared validation losses (shown in Fig. 5). Validation performance is very similar amongst different 7 at the training context length (though r 10 is strictly best), but as we extend to out distribution context lengths (64k, 16x the training context length) the benefit of 7 10 becomes clearer.
26

---

Train @4K / Val @4k
Train @4K / Val @16k
Train @4k / Val @64k
3.4
10
t- 10
10
3.3
Vall
3500
4000
4500
3500
4000
4500
3500
4000
4500
Step
Step
Step
Figure 5: Validation losses at 4k (left), 16k (middle), and 64k (right) context lengths, for GPT-2-like model trained with scale-invariant attention for varying т. The models were trained at 4k context length.
Expected entropy
Square expected entropy
20
400
Am 2
300
E[H
E[H
200
100
2
8
2
logulU
logs/(U
Figure 6: Scaling of expected entropy-in-range for standard attention with an independent Gaussian assumption on logits. We calculate expected entropies in the range [t. (A] different and A E {2, 5, 10), with the lengthscale set to 10. The dashed lines show the best linear fit of the data. We empirically find that standard Gaussian logits give logarithmic entropy (left panel).
I Extra Results
L.1 Entropy scaling of regular attention
In the main text, illustrated in Fig. 2 that scale-invarian attention method under Gaussian assumption has sub-logarithmic expected entropy. Fig. 6 empirically shows that expected entropy in range to tA for standard/unscaled attention instead scales logarithmically with
L.2 Pretrain @ 16k
For completeness, we include validation losses when training at 16k in Fig. We see again that scale-invariant p-RoPE outperforms other methods range of validation lengths, with the improvements becoming more noticable at the longest validation length of 64k.
I.3 Alternative scale-invariant attention variants
In Section 3 we presented scale- invariint R-RoPE as our proposed method. In preliminary experiments however, was very natural also consider scale-invariant RoPE and scale-invariant NoPE. We show results when training at 4k in Fig. 8. Scale-invariant RoPE performs almost as well as scale-invariant p-RoPE when evaluating at the training context length, but underperforms more we move 16k and 64k. On the other hand, scale-invariant NoPE underperforms iaveeinnaaiiatt p-RoPE, yet generalises
27

---

ROPE P-RoPE
RoPE+NTK
-LogN=p-RoPE LogNN-NTK -Scale-invarianl P-RoPE (ours)
Train @16k/ Val @4k
Train @16k Val @16k
Train @16k Val @64k
Val Loss
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
Step
Step
Step
Figure 7: Validation losses at 4k / 16k / 64k context lengths, for a 162M parameter GPT-2-style model trained at 16k. NoPE omitted the right-most plot to avoid excessive zooming.
Train @4k Val @4k
Train @4k Val @16k
Train @4k @64k
Scale-invarian Scale-inarriant RoPE
Invariant
1000
2000
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
Step
Figure 8: Validation losses at 16k 64k of icale-invariant NoPE, RoPE, and p-RoPE, for 162M
parameter GPT-2-style model trained at 4k
to long contexts. We hypothesised that scale-invariae invariant RoPE does not generalise due to RoPE's low frequences (i.c. high wavelengths) interfering with the position-dependent scale-invariam transformation, T., ++ Ot Tu + my-
I.4 Infini-Attention
During the review process, attention (Munkhdalai et al., 2024). Since infini-attention is method for compressing the KV-cache, it is slightly different the other methods we compare (which primarily control entropy), and so we include these results separately. See Table 3, which shows that while infini-attention gives long-context generalization, it is not as strong method.
asked compare scale-invariant attention Infini-
L.5 Larger models
We also investigated the ability for scale-invariant attention to generalise to longer contexts in larger models by continual pretraining Llama 2 7B (Touvron al., 2023). We chose Llama 2 for this experiment because is one of the few models that have not already mid-trained at longer context length.
Specifically, we continually pretrained at 4k context length after replacing the default attention mechanism with various other methods, and looked validation loss out-of-distribbtion lengths (4k/16k/64k). We trained using the Torchtune (2024) library, using data from Web (Penedo et al., 2024), with AdamW and learning rate of 2 10-5
We see Table that changing the attention mechanism degraded the loss at the training context length, but the performance of cale-invariant R-RoPE far exceeds the other methods at 16k and 64k.
28

---

Table Final mean validation losses (#1 standard error across 3 seeds) for different methods when training with 4k context length 162M parameter GPT-2--stlle model.
Method
Val 4k
16k
Val
RoPE
3.261 + 0.001
3,936 = 0.010
5.260
0.014
icale-invariant p- p-RoPE (ours)
3.244 + 0.001
3.235 = 0.001
3.247 0.001
Infini-RoPE
3.296 0.003
3.302 + 0.004
3,310 0.008
Infini-p-RoPE
3.295 + 0.003
3.303 + 0.005
3.311 + 0.009
Table 4: Validation performance when fine-tuning Llama-2 7B with different attention methods on 500 tokens. RoPE and RoPE+NTK (denoted *) were not fine-tuned.
Method
Val loss @4k
Val loss @ 16k
Val loss @ @64k
RoPE*
1.968
7.036
8.815
p-RoPE
2.029
3.504
6.800
NoPE
4.152
6.172
7.730
RoPE+NTK*
1,988
3,090
7.523
YaRN
1,957
7,004
8.763
LogN+RoPE
1.966
6,932
8.726
LogN+p-RoPE RoPE
2.049
2.984
6.224
LogN+NTK
1.966
3.063
7,413
ALiBi
2.750
2.745
2.744
Scale-invariant p-RoPE (ours)
2.163
2,193
2.252
Checking Gaussianity of attention logits
In our analysis in Section 3.2 we assume that the unmodified logits (query-key products), L, 's, are standard Gaussian. this section, we empirically verify that the logits Gaussian by looking at QQ-plots.
We consider several sizes of model, including 1B and 8B Llama variants (Grattafiori al., 2024), and Gemma 27B Team et al. (2024). We calculate logits with the introduction paragraph Wikipedia page as the input. In Figs. 9,10,11,1..1314 we show quantiles of {Le (i.c. lower triangular part of the QK matrix) for each layer, aggregated over the input and the attention heads in each layer. Note that do not aggregate the 'beginning of sequence' token, it an outlier attention sink (Gu et al., 2024).
29

---

Layer
Layer 2
Layer
Layer 6
4
Sample quantiles
2
0
0
-2
-2
-2
-2
-4
-4
Layer 8
Layer 10
Layer 12
Layer 14
4
Sample quantiles
2
0
0
-2
-2
-2
-2
-4
-4
-4
-4
-4
-
0
4
2
0
2
4
4
-
2
4
-4
-2
0
2
4
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
Figure 9: Quantile-Quantie (QQ) plots attention logits (without RoPE applied) Llama-1B
model (blue line), with theoretical quantiles of Gaussian shown by the red line.
Layer
Layer 2
Layer
Layer 6
4
Sample quantiles
2
2
0
0
0
-2
-2
-2
-2
-4
-4
Layer 8
Layer 10
Layer 12
Layer 14
4
Sample quantiles
2
0
0
-2
-2
-2
-2
-4
-4
-4
-4
-4
-2
0
4
-4
-2
0
2
4
-4
-
2
4
-4
-2
0
2
4
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
Figure 10: Ctantile-Quantiee (QQ) plots attention logits (with RoPE applied) in Llama- 1B model
(blue line), with theoretical quantiles of Gaussian shown by the red line.
30

---

Layer
Layer 4
Layer
Layer 12
4
Sample quantiles
2
0
0
-2
-2
-2
-2
-4
-4
Layer 16
Layer 20
Layer 24
Layer 27
4
4
Sample quantiles
2
0
0
-2
-2
-2
-2
-4
-4
-4
-4
-4
-
0
4
2
0
2
4
4
-
2
4
-4
-2
0
2
4
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
Figure 11: Quantile-Quanile (QQ) plots attention logits (without RoPE applied) in Llama-8B
model (blue line), with theoretical quantiles of Gaussian shown by the red line.
Layer
Layer 4
Layer
Layer 12
4
Sample quantiles
2
2
0
0
0
-2
-2
-2
-2
-4
-4
Layer 16
Layer 20
Layer 24
Layer 27
4
Sample quantiles
2
0
0
-2
-2
-2
-2
-4
-4
-4
-4
-4
-2
0
4
-4
-2
0
2
4
-4
-
2
4
-4
-2
0
2
Theoratical
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
quantiles
Figure 12: Quantile-Quant (QQ) plots attention logits (with RoPE applied) in Llama-8B model
(blue line), with theoretical quantiles of Gaussian shown by the red line.
31

---

Layer
Layer 4
Layer
Layer 12
4
Sample quantiles
2
2
2
0
0
-2
-2
-2
-2
-4
-4
-4
-4
Layer 16
Layer 20
Layer 24
Layer 28
4
4
4
4
Sample quantiles
2
2
0
0
0
-2
-2
-2
-2
-4
-4
-4
-4
Layer 32
Layer 36
Layer 40
Layer 44
4
Sample quantiles
2
2
2
0
0
-2
-2
-2
-2
-4
-4
-4
-4
-4
2 2
4
-
0
2
4
-4
2
0
2
4
-4
-2
0
2
4
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
Theoratical quantiles
Figure 13: Quantile-Quantee (QQ) plots of attention logits (without RoPE applied) in a Gemma 27B model (blue line), with theoretical quantiles of a Gaussian shown by the red line.
Layer
Layer 4
Layer
Layer 12
4
Sample quantiles
2
2
2
0
-2
-2
-2
-2
-4
-4
-4
-4
Layer 16
Layer 20
Layer 24
Layer 28
4
4
4
Sample quantiles
2
2
0
0
0
-2
-2
-2
-2
-4
-4
-4
-4
Layer 32
Layer 36
Layer 40
Layer 44
4
Sample quantiles
2
2
2
0
0
-2
-2
-2
-2
-4
-4
-4
-4
2 2
4
2
0
2
4
-4
2
0
2
4
4
2
2
4
Theoratical
Theoratical quantiles
Theoratical quantiles
Theoratical
quantiles
quantiles
Figure 14: Quantile- Quaill--uuanilee (QQ) plots of attention logits (with RoPE applied) in a Gemma 2 27B
model (blue line), with theoretical quantiles of Gaussian shown by the red line.
32

---

K
Licenses
This project uses modified version of modded nanopp, https://githnb.com com/
LellerJordannmoddded nanogpt which is MIT licensed.
This project uses 10B subset of the fineweb dataset, http:://hugghhgtace.ov
hatasets/kjjj/fineeee - which MIT licensed.
This project uses a 100B subset of the fineweh dataset, https://hugiingfacecoo gpt2, which is MIT licensed.
This project uses the C4 dataset, https: https://huginggaac. co/dattaetts/kkj//
fineweb100B-gpte which licensed under the Open Civic Data Attribution License (OCD-
BY).
33


Reviews:
[{"review_id": "0eGbWytGXD", "reviewer": "Reviewer_TuDC", "summary": "The paper proposes to transform attention scores to achieve scale invariance. It introduces two notions of scale invariance, namely equal sum and sparsity of attention weights across token ranges. The paper derives a close-form transformation function under some assumptions that enables scale-invariant total attention, while empirically demonstrates that the same transformation also achieves scale-invariant attention sparsity. The experiments demonstrates that the training small-scale models with the proposed attention scheme achieves lower validation loss compared to the baselines, especially when the context is longer than during training. Further, the model performs better than the baselines on long-context retrieval following fine-tuning.", "questions": "N/A", "limitations": "The paper has addressed the limitations of the work.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths:\n- The transformation applied on attention scores is mathematically proven to satisfy the definition of scale invariance under certain assumptions.\n- The empirical results suggest the potential effectiveness of the method.\n\n## Weaknesses:\n- **Motivation:** The choice of scales (powers of 10) appears arbitrary and seems mainly for mathematical convenience. Additionally, it is unclear why a language model should satisfy the scale invariance properties. The paper asserts that scale invariance facilitates generalization on long contexts, but this claim is not substantiated by the toy experiments presented in the paper.\n\n- **Experiments:** The paper lacks solid experiments. Validation loss is the only metric, yet there is no evidence showing its correlation with the accuracy of any downstream tasks. The NIH experiment is also not convincing, as the model is fine-tuned to overfit this very mechanical task. Ideally, one would perform scaling experiments and run a suite of LLM benchmarks to show the method works and scales well. While there may be constraints on computational resources, legitimate experiments are essential to support the claims for acceptance in a peer-reviewed conference. Nonetheless, this work could be a valuable tech report and may gain traction if the method turns out to be effective with community effort.\n\n- **Visualization:** It helps to plot the distribution of attention weights to help visualize scale invariance.", "was_revised": true, "final_justification": "The paper studies scale-invariant attention for long-context understanding. The approach is simple, theoretically grounded, and demonstrates strong results at a limited scale. While it remains unclear whether the method will yield consistent gains as model size grows, the work offers an interesting perspective on how to improve the attention mechanism to better support longer contexts.", "rebuttal": "Thanks for your careful and positive review!\n\n>Motivation: The choice of scales (powers of 10) appears arbitrary and seems mainly for mathematical convenience. \n\nPowers of 10 are only used in the examples we use to build intuition!  They don't appear anywhere in the actual method (i.e. the scaling factors used in Eq. 16).  These scaling factors are derived from Lemma 1 and 2, which again do not use powers of 10.\n\n>Additionally, it is unclear why a language model should satisfy the scale invariance properties. \n\nThis is discussed extensively in the Introduction. In short, it is clear that you don't want the LLM to focus almost exclusively on nearby tokens (1-10 tokens ago).  You also want it to look at more distant tokens.  Additionally, you don't want the LLM to focus almost exclusively on distant tokens (100-1,000,000 tokens ago), at the expense of local tokens.  But that's what standard LLMs do (see Figure 1 right hand column; orange line). You want attention to balance looking at nearby and more distant tokens. And you want the balance to be maintained as the context gets longer. That's what the scale-invariance desiderata give you. While of course the precise optimal choice of desiderata is an important topic for future work, it is clear that any attention mechanism that zero-shot generalises to long-contexts is going to have to balance attention to nearby and more distant tokens, and thus is going to have some type of scale-invariance property.\n\n\n**Experiments**\n\nWe have added further experiments --- continual pretraining on Llama 2 7B at 4k context length, and evaluating the validation loss at longer context sizes (see table below) to demonstrate that scale-invariant attention extends to larger scale models. \n\n| Method                          | Val loss @ 4k | Val loss @ 16k | Val loss @ 64k |\n|--------------------------------|----------|-----------|-----------|\n| RoPE*                          | 1.968    | 7.036     | 8.815     |\n| $p$-RoPE                       | 2.029    | 3.504     | 6.800     |\n| NoPE                           | 4.152    | 6.172     | 7.730     |\n| RoPE+NTK*                      | 1.988    | 3.090     | 7.523     |\n| YaRN                           | 1.957    | 7.004     | 8.763     |\n| LogN+RoPE                      | 1.966    | 6.932     | 8.726     |\n| LogN+$p$-RoPE                  | 2.049    | 2.984     | 6.224     |\n| LogN+NTK                       | 1.966    | 3.063     | 7.413     |\n| ALiBi                          | 2.750    | 2.745     | 2.744     |\n| Scale-invariant $p$-RoPE (ours)| 2.163    | 2.193     | 2.252     |\n\nSpecifically, we continually pretrained on Llama 2 at 4k context length after replacing the default attention mechanism with various other methods, and we looked at validation loss at out of distribution lengths (4k/16k/64k). Changing the attention mechanism degraded the loss at the training context length, but scale-invariant $p$-RoPE far exceeded the other methods. Note that we chose Llama 2 for this continual pretraining experiment because it is one of the few models that have not already mid-trained at a longer context length.\n\n\nWith this addition, our experiments used around 500 H100 hour equivalents. This is a typical scale for an academic paper published at an academic conference, and indeed is similar to the compute used in previous academic papers presented at academic conferences presenting new long-context attention mechanisms (e.g. [1,2,3]).  It is not at all clear we should be demanding authors have foundation-model-lab compute budgets, especially given that these labs are increasingly refusing to publish important innovations.\n\n\n\n> Visualization: It helps to plot the distribution of attention weights to help visualize scale invariance.\n\nWe agree! We will add this for the camera ready.\n\nRefs:\n\n[1]: Press et al., Train short, test long: Attention with linear biases enables input length extrapolation\n\n[2]: Nakanishi, Scalable-softmax is superior for attention\n\n[3]: Li et al., Information entropy invariance: Enhancing length extrapolation in attention mechanisms"}, {"review_id": "NWkiVIYTao", "reviewer": "Reviewer_6x9F", "summary": "This paper proposes \"scale-invariant attention\" to improve the ability of large language models (LLMs) to generalize from short training contexts to longer inference contexts. This method aims for attention to function stably across various scales, from local to global, and is specifically designed to satisfy two desirable properties: \"scale-invariant total attention\" and \"scale-invariant attention sparsity\". Based on a Gaussian assumption, the authors theoretically derive a method to achieve these properties by applying a simple, position-dependent affine transformation to the attention logits. In experiments, the proposed method, when combined with p-RoPE, was demonstrated to outperform prominent existing methods (such as ALiBi and LogN) in terms of validation loss. It showed particularly excellent performance in zero-shot generalization from a short context (4k) to a long context (64k). Furthermore, it maintained high performance in a long-context information retrieval task.", "questions": "- Regarding the IID Gaussian assumption that is central to your theory: while the QQ plot in Appendix J is interesting, do you believe the IID assumption holds even after applying position-dependent encodings like RoPE? How would you predict the effectiveness or performance of your proposed transformation to be affected if this assumption is only partially met?\n- The LogN trick also aims to control the attention distribution for long contexts. Your method differs by using a position-dependent transformation and shows superior results in your experiments. Do you envision any scenarios or task types where the simpler, position-independent LogN trick might be preferable? Or does your position-dependent approach have fundamental advantages that would make it superior in almost all long-context situations?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n- The effectiveness of the proposed method is demonstrated through a multi-faceted and convincing set of experiments, using multiple model sizes, various training/validation context lengths, and a range of strong baseline methods for comparison. The result showing superior performance in a 16x zero-shot generalization task is particularly impressive and highlights the strength of the method.\n- The proposed transformation on logits is a simple affine transformation using the logarithm of the position index, which can be easily integrated into existing Transformer architectures. This makes the research highly practical.\n\n**Major Weaknesses**\n- The theoretical derivation hinges on the assumption that attention scores follow an IID (independent and identically distributed) Gaussian distribution. The paper's justification for this assumption is insufficient. While Appendix J presents a QQ plot for a Llama-3B model, this is just a single case. It is particularly questionable whether logits remain IID with respect to position after applying positional encodings like RoPE. The main text lacks a discussion on how valid this assumption is for real-world models and what the implications would be for the method's performance if the assumption is violated.\n\n**Minor Weaknesses**\n- While Appendix H explores the lengthscale hyperparameter τ, the analysis of why τ=10 is optimal is somewhat superficial. The intuition that it represents the \"size of 'chunks' we attend over\" is helpful, but a more in-depth analysis of the relationship between the optimal τ, model architecture, and data characteristics would have further strengthened the paper.", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thoughtful and positive review!\n\n> The theoretical derivation hinges on the assumption that attention scores follow an IID (independent and identically distributed) Gaussian distribution. The paper's justification for this assumption is insufficient. While Appendix J presents a QQ plot for a Llama-3B model, this is just a single case. It is particularly questionable whether logits remain IID with respect to position after applying positional encodings like RoPE. The main text lacks a discussion on how valid this assumption is for real-world models and what the implications would be for the method's performance if the assumption is violated.\n\nAfter further analysis during the review period, we determined that Theorem 1 does not, in fact, require independence of Gaussians.  It simply requires marginally Gaussian random variables.  This is evident if you look at Lemmas 1+2, which do not require independent Gaussian variables, along with Appendix F, which again does not require independence.\n\n> While Appendix H explores the lengthscale hyperparameter $\\tau$, the analysis of why $\\tau=10$ is optimal is somewhat superficial. The intuition that it represents the \"size of 'chunks' we attend over\" is helpful, but a more in-depth analysis of the relationship between the optimal $\\tau$, model architecture, and data characteristics would have further strengthened the paper.\n\nWe agree that we established the optimal value of $\\tau$ experimentally in a particular setup. \nThus we would recommend that anyone using scale-invariant attention should check that $\\tau=10$ is close-to-optimal in their setting. \nGiven our intuition that $\\tau$ measures the \"size of 'chunks' we attend over\", we would not expect the optimal $\\tau$ to be too far from $10$, even in alternative architectures, but we would emphasise that this does need to be checked.\n\n#### Questions:\n\n> Regarding the IID Gaussian assumption that is central to your theory: while the QQ plot in Appendix J is interesting, do you believe the IID assumption holds even after applying position-dependent encodings like RoPE?\n\nAs discussed above, IID is not necessary.  The QQ plots the QQ plots in Appendix J show attention logits in a Llama-3, to which RoPE position embeddings have been applied, and they clearly show Gaussian marginals.\n\n>The LogN trick also aims to control the attention distribution for long contexts. Your method differs by using a position-dependent transformation and shows superior results in your experiments. Do you envision any scenarios or task types where the simpler, position-independent LogN trick might be preferable? Or does your position-dependent approach have fundamental advantages that would make it superior in almost all long-context situations?\n\nWe had thought that LogN might do better at a needle-in-a-haystack task, where the needle is uniformly distributed across a long context. In this setting, you would expect position-independent LogN scaling to perform better.  In practice, however, we found that scale-invariant p-RoPE performed slightly better at 16k and 64k context lengths (Table 2), but this difference was not statistically significant, so we are only able to say that performance of scale-invariant p-RoPE and LogN was indistinguishable."}, {"review_id": "8CiQU1Xc51", "reviewer": "Reviewer_yBmN", "summary": "This paper introduces a new attention mechanism—scale-invariant attention—aimed at improving large language models’ ability to generalize from training on short contexts to inference on long contexts. The authors formalize two desirable properties for attention over long sequences: scale-invariant total attention, which ensures that attention is evenly distributed across different positional ranges, and scale-invariant attention sparsity, which ensures that attention becomes increasingly sparse over more distant contexts. To achieve these properties, the authors propose a simple transformation of the attention logits using position-dependent multiplicative and additive terms. They provide theoretical proofs under Gaussian assumptions that this transformation satisfies both scale-invariant total attention and (weak) scale-invariant attention sparsity.", "questions": "1. Can the authors compare their method to memory-based approaches or retrieval-augmented LLMs, particularly on long-context retrieval tasks?\n2. Does the position-dependent logit scaling affect optimization dynamics? Empirical data on convergence speed or stability would be helpful.\n3. Can the authors provide more detailed empirical or theoretical analysis on why p-RoPE pairs better with scale-invariant attention than RoPE? Frequency spectrum visualizations could help.\n4. While the study focuses on GPT-style models, could this method extend to encoder-decoder or encoder-only models like BERT or T5?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n1. The paper provides a solid mathematical foundation for its proposed attention mechanism, with clear definitions and provable guarantees for the desired properties.\n2. The proposed logit transformation is lightweight and easy to implement, requiring only position-dependent scaling and biasing of logits.\n3. The method shows strong zero-shot performance without requiring long-context pretraining, which has implications for cost-efficient LLM deployment.\n\n**Weaknesses**\n1. Experiments are conducted on relatively small-scale models (162M and 304M parameters), which may limit confidence in applicability to larger foundation models.\n2. The authors do not analyze whether their logit transformation impacts convergence speed or training stability, which may be useful for practitioners.\n3. The three lines drawn in Figure 2 have low differentiation, so it is recommended that the author use different colors to mark them.\n4. The performance in Table 2 is slightly lower at 4k. It is recommended that the author analyze the reasons.", "was_revised": true, "final_justification": "I will keep my score positive considering all rebuttal and discussions.", "rebuttal": "Thanks for your positive and careful review.\n\n> Experiments are conducted on relatively small-scale models (162M and 304M parameters), which may limit confidence in applicability to larger foundation models.\n\nWe have added further experiments --- continual pretraining on Llama 2 7B at 4k context length, and evaluating the validation loss at longer context sizes (see table below) to demonstrate that scale-invariant attention extends to larger scale models. \n\n\n| Method                          | Val loss @ 4k | Val loss @ 16k | Val loss @ 64k |\n|--------------------------------|----------|-----------|-----------|\n| RoPE*                          | 1.968    | 7.036     | 8.815     |\n| $p$-RoPE                       | 2.029    | 3.504     | 6.800     |\n| NoPE                           | 4.152    | 6.172     | 7.730     |\n| RoPE+NTK*                      | 1.988    | 3.090     | 7.523     |\n| YaRN                           | 1.957    | 7.004     | 8.763     |\n| LogN+RoPE                      | 1.966    | 6.932     | 8.726     |\n| LogN+$p$-RoPE                  | 2.049    | 2.984     | 6.224     |\n| LogN+NTK                       | 1.966    | 3.063     | 7.413     |\n| ALiBi                          | 2.750    | 2.745     | 2.744     |\n| Scale-invariant $p$-RoPE (ours)| 2.163    | 2.193     | 2.252     |\n\nWe continually pretrained on Llama 2 at 4k context length after replacing the default attention mechanism with various other methods, and we looked at validation loss at out of distribution lengths (4k/16k/64k). Changing the attention mechanism degraded the loss at the training context length, but scale-invariant $p$-RoPE performed best out of distribution. Note that we chose Llama 2 for this continual pretraining experiment because it is one of the few models that have not already mid-trained at a longer context length.\n\n> The authors do not analyze whether their logit transformation impacts convergence speed or training stability, which may be useful for practitioners.\n\nFigures 3, 4, 5, 7, 8 all show training curves (i.e. validation loss vs. number of gradient descent steps), and we can clearly see that number of steps required to reach a given loss is improved by our method, and it does not give rise to any indications of a loss-of-stability.\n\n> The three lines drawn in Figure 2 have low differentiation, so it is recommended that the author use different colors to mark them.\n\nFixed! Thanks!\n\n> The performance in Table 2 is slightly lower at 4k. It is recommended that the author analyze the reasons.\n\nThis difference is statistically insignificant at the 1\\% level so all we can say is that the methods are not distinguishable.\n\n\n#### Questions:\n\n> Can the authors compare their method to memory-based approaches or retrieval-augmented LLMs, particularly on long-context retrieval tasks?\n\nLong context (LC) methods and RAG methods are very different methods with very different performance characteristics.  For instance, RAG is thought to perform very poorly for code [1], and LC is generally thought to outperform RAG for QA benchmarks, especially for Wikipedia based tasks [2]. Though of course, RAG is considerably cheaper than LC methods, assuming fast search is available.  Thus, understanding the tradeoffs between these methods is complex, and depends on the tradeoff between desired cost and performance for that context size.\n\nWe consider our work as pushing the frontier of long-context methods.\n\n[1] Nik Pash \"Why I No Longer Recommend RAG for Autonomous Coding Agents\" (Substack blog post) 2025\n\n[2] Xinze Li, Yixin Cao, Yubo Ma, Aixin Sun \"Long Context vs. RAG for LLMs: An Evaluation and Revisits\" (arXiv:2501.01880)\n\n\n> Can the authors provide more detailed empirical or theoretical analysis on why p-RoPE pairs better with scale-invariant attention than RoPE? Frequency spectrum visualizations could help.\n\nIt is understood that $p$-RoPE tends to work better than RoPE in the long-context settings by making semantic attention patterns more robust at long distances [1]. Indeed both DeepSeek and Kimi implement a form of $p$-RoPE via MLA: see in MLA that RoPE is not applied to the compressed keys and queries [2, section 2.1.1]. Further, increasing the base $\\theta$, which is common practice, can be seen as approximating $p$-RoPE. We therefore believe our results on RoPE vs $p$-RoPE are simply in keeping with these more general, and well understood, trends.\n\n> While the study focuses on GPT-style models, could this method extend to encoder-decoder or encoder-only models like BERT or T5?\n\nWe see no reason for our method not to extend to different transformer architectures, or indeed any NN architecture that uses softmax attention.\n\n[1]: Barbero et al., Round and round we go! what makes rotary positional encodings useful? ICLR\n\n[2]: Liu, Aixin, et al. \"Deepseek-v3 technical report.\" arXiv."}, {"review_id": "r1qjSWesK3", "reviewer": "Reviewer_tNeo", "summary": "This paper addresses the challenge of long-context generalization in LLMs by proposing scale-invariant attention, which satisfies two key properties: scale-invariant total attention (the expected sum of unnormalized attention weights over any token range \\([t, t\\Delta)\\) is asymptotically constant) and scale-invariant attention sparsity (attention is denser in local ranges and sparser in global ranges, with entropy growing sub-logarithmically). Under a Gaussian assumption, the authors derive a position-dependent transformation of attention logits (\\(L_t = a_t \\bar{L}_t + m_t\\)) to satisfy these properties. Experimental results show that the proposed \"scale-invariant p-RoPE\" outperforms baselines (e.g., RoPE, LogN, ALiBi) in validation loss across various training/validation context lengths, enables strong zero-shot generalization from short to long contexts, and performs well in the needle-in-a-haystack retrieval task. The main contributions include introducing the two scale-invariant properties, deriving the logit transformation, and validating its effectiveness empirically", "questions": "(1) The paper only evaluates models with 162M and 304M parameters, but contemporary LLMs often have billions of parameters. How does scale-invariant attention perform on larger models (e.g., 10B+ parameters)? Larger models may exhibit different attention behaviors due to increased capacity; verifying scalability is crucial for assessing real-world applicability. For example, do the hyperparameters (e.g., \\(\\tau=10\\)) remain optimal, or do they need adjustment for larger architectures?\n\n(2) The logit transformation is derived under the assumption that base logits (\\(\\bar{L}_t\\)) are IID Gaussian. However, real attention logits in trained models may deviate from Gaussianity. The paper mentions a QQ plot for Llama-3B, but how robust is the proposed method when logits are non-Gaussian (e.g., skewed or with heavier tails)? Would the transformation still satisfy scale-invariant properties, or is additional adaptation needed?\n\n(3) While the paper compares with several baselines, it does not include some recent long-context methods, such as Infini-Attention or Ring Attention. How does scale-invariant attention compare to these methods in terms of performance, computational efficiency, and generalization? Including such comparisons would better situate the work within the broader landscape of long-context LLMs.\n\n(4) The paper mentions that scale-invariant RoPE underperforms scale-invariant p-RoPE due to RoPE’s low-frequency components interfering with the transformation. However, it does not elaborate on why low-frequency components cause this interference. Could this be mitigated by modifying RoPE’s frequency spectrum (e.g., adjusting the base \\(\\theta\\))? A deeper analysis of this interaction would clarify the choice of p-RoPE as the base.\n\n(5) The needle-in-a-haystack task uses fine-tuning, which is unusual (most studies use prompting). The paper states that prompting alone was insufficient, but it does not explain why (e.g., model size, pretraining data). Would larger models or more pretraining enable the proposed method to perform this task with prompting? This information is important for assessing its practical utility in retrieval scenarios without fine-tuning.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "## Strength\n1. The paper rigorously defines scale-invariant total attention and sparsity, provides mathematical proofs (Lemmas 1, 2, and Theorem 1) for the conditions under which these properties hold, and derives a concrete logit transformation, ensuring theoretical soundness.\n2. The proposed method involves a simple position-dependent transformation of attention logits, with only one hyperparameter (\\(\\tau\\)) after simplification, making it easy to implement and integrate with existing frameworks like p-RoPE.\n3. The paper effectively distinguishes its approach from related methods (e.g., LogN’s position-independent scaling vs. its position-dependent transformation), highlighting its unique ability to balance local and global attention.\n\n## Weakness\n1. The paper only validates the proposed method on models with 162M and 304M parameters, like GPT2, while state-of-the-art LLMs typically have billions of parameters. The authors acknowledge this limitation, noting that resource constraints prevented testing on larger models, which raises questions about whether the conclusions generalize to commercially deployed large-scale models.\n2. Although the paper compares with methods like RoPE, LogN, and ALiBi, it omits comparisons with several recent long-context attention mechanisms (e.g., Infini-Attention, Ring Attention) that also aim to address length generalization. This gap makes it harder to situate the proposed method within the broader landscape of current long-context LLM research .\n3. The theoretical derivation of the scale-invariant attention transformation hinges on the assumption that base attention logits (\\(\\bar{L}_t\\)) follow an IID Gaussian distribution. While the paper includes a QQ plot for Llama-3B to partially verify this, it lacks in-depth analysis of how the method performs when logits deviate significantly from Gaussianity (e.g., skewed distributions or heavy tails), which could limit its robustness in real-world scenarios .\n4. The issue that this paper aims to address is significant, while the presentation and experimental desmonstration is poor for me. How the proposed method motivated by studying at entropy? how it affect the design? and why this factor not evaluated at experiments?", "was_revised": true, "final_justification": "The response from authors address my major concerns. So i raise the score.", "rebuttal": "Thanks for your careful and positive review!\n\n> How the proposed method motivated by studying at entropy? how it affect the design? and why this factor not evaluated at experiments?\n\nAs discussed in Sec 3.1, we used the entropy of the distribution defined by the attention weights to measure sparsity.  For a sequence of length $T$, if the entropy is $\\log T$, then the attention must be uniformly distributed across all tokens.  In contrast, if the entropy is $0$, then the attention must be focused on a single token.  Outside of these extreme cases, an entropy of $\\log n$ could be generated by an attention weight of $1/n$ assigned to $n$ tokens and an attention weight of $0$ assigned to all the other tokens (though of course in practice many other distributions will also have an entropy of $\\log n$). \n\nWe evaluated the entropy in Figure 1.\n\n\n\n> For example, do the hyperparameters (e.g., ($\\tau=10$)) remain optimal, or do they need adjustment for larger architectures?\n\nWe agree that we established the optimal value of $\\tau$ experimentally in a particular setup. \nThus we would recommend that anyone using scale-invariant attention should check that $\\tau=10$ is close-to-optimal in their setting. \nGiven our intuition that $\\tau$ measures the \"size of 'chunks' we attend over\", we would not expect the optimal $\\tau$ to be too far from $10$.\n\n\n\n> (1) The paper only evaluates models with 162M and 304M parameters, but contemporary LLMs often have billions of parameters. How does scale-invariant attention perform on larger models (e.g., 10B+ parameters)? Larger models may exhibit different attention behaviors due to increased capacity; verifying scalability is crucial for assessing real-world applicability. For example, do the hyperparameters (e.g., ($\\tau=10$)) remain optimal, or do they need adjustment for larger architectures?\n\nWe have added further experiments --- continual pretraining on Llama 2 7B at 4k context length, and evaluating the validation loss at longer context sizes (see table below) to demonstrate that scale-invariant attention extends to larger scale models. \n\n\n| Method                          | Val loss @ 4k | Val loss @ 16k | Val loss @ 64k |\n|--------------------------------|----------|-----------|-----------|\n| RoPE*                          | 1.968    | 7.036     | 8.815     |\n| $p$-RoPE                       | 2.029    | 3.504     | 6.800     |\n| NoPE                           | 4.152    | 6.172     | 7.730     |\n| RoPE+NTK*                      | 1.988    | 3.090     | 7.523     |\n| YaRN                           | 1.957    | 7.004     | 8.763     |\n| LogN+RoPE                      | 1.966    | 6.932     | 8.726     |\n| LogN+$p$-RoPE                  | 2.049    | 2.984     | 6.224     |\n| LogN+NTK                       | 1.966    | 3.063     | 7.413     |\n| ALiBi                          | 2.750    | 2.745     | 2.744     |\n| Scale-invariant $p$-RoPE (ours)| 2.163    | 2.193     | 2.252     |\n\nWe found that while replacing the attention mechanism had adverse effects @ 4k, scale-invariant attention yielded by far the best results OOD @ 64k. (Note that we chose Llama 2 for this continual pretraining experiment because it is one of the few models that have not already mid-trained at a longer context length).\n\n\n> (2) The logit transformation is derived under the assumption that base logits ($\\bar{L}_t$) are Gaussian. However, real attention logits in trained models may deviate from Gaussianity. The paper mentions a QQ plot for Llama-3B, but how robust is the proposed method when logits are non-Gaussian (e.g., skewed or with heavier tails)? Would the transformation still satisfy scale-invariant properties, or is additional adaptation needed?\n\nTo confirm that the distribution over logits is approximately Gaussian, we have added additional QQ plots for both larger Llama-3 models as well as Gemma models. Given that we were not able to find a realistic setting where the logits deviate from Gaussianity, we were not able to study this possibility (and indeed, it does seem that it would be rare in practice).\n\n> (3) While the paper compares with several baselines, it does not include some recent long-context methods, such as Infini-Attention or Ring Attention. How does scale-invariant attention compare to these methods in terms of performance, computational efficiency, and generalization? Including such comparisons would better situate the work within the broader landscape of long-context LLMs.\n\nInfini-attention is a method for compressing the KV-cache, and is thus slightly different to the other methods we compare to (which primarily control entropy). We nonetheless ran pretraining experiments with Infini-attention, and found that while Infini-attention shows length generalization, it is not as strong as our method. See the table below; we include standard attention (RoPE) for comparison,\n\n| Method                         | Val @ 4k         | Val @ 16k        | Val @ 64k        |\n|-------------------------------|------------------|------------------|------------------|\n| RoPE                          | 3.261 ± 0.001    | 3.936 ± 0.010    | 5.260 ± 0.014    |\n| **Scale-invariant *p*-RoPE** (ours) | **3.244 ± 0.001** | **3.235 ± 0.001** | **3.247 ± 0.001** |\n|-------------------------------|------------------|------------------|------------------|\n| Infini+RoPE                   | 3.296 ± 0.003    | 3.302 ± 0.005    | 3.310 ± 0.010    |\n| Infini+*p*-RoPE               | 3.295 ± 0.003    | 3.303 ± 0.006    | 3.311 ± 0.011    |\n\nRing Attention is a method for distributing the attention computation more efficiently over GPUs/TPUs --- it doesn't actually alter the attention computation at all. Therefore it will not improve long context performance, when coupled e.g. with RoPE, when out of the training distribution.\n\n> (4) The paper mentions that scale-invariant RoPE underperforms scale-invariant p-RoPE due to RoPE’s low-frequency components interfering with the transformation. However, it does not elaborate on why low-frequency components cause this interference. Could this be mitigated by modifying RoPE’s frequency spectrum (e.g., adjusting the base ($\\theta$))? A deeper analysis of this interaction would clarify the choice of p-RoPE as the base.\n\nIt is understood that $p$-RoPE tends to work better than RoPE in the long-context settings by making semantic attention patterns more robust at long distances [4]. Indeed both DeepSeek and Kimi implement a form of $p$-RoPE via MLA: see in MLA that RoPE is not applied to the compressed keys and queries [5, section 2.1.1]. Further, increasing the base $\\theta$, which is common practice, can be seen as approximating $p$-RoPE. We therefore believe our results on RoPE vs $p$-RoPE are simply in keeping with these more general, and well understood, trends.\n\n\n> (5) The needle-in-a-haystack task uses fine-tuning, which is unusual (most studies use prompting). The paper states that prompting alone was insufficient, but it does not explain why (e.g., model size, pretraining data). Would larger models or more pretraining enable the proposed method to perform this task with prompting? This information is important for assessing its practical utility in retrieval scenarios without fine-tuning.\n\nScale-invariant attention makes a relatively minor change to the attention mechanism that enables zero-shot length generalisation.  When keeping the context size at or below the training context size, we'd expect scale-invariant models to perform in almost exactly the same way as standard attention, and that is bourne out in our results (see Val @ 4k columns in Table 1 and Table 2). We found that the networks we pretrained with both standard attention and scale-invariant attention were unable to perform the needle-in-a-haystack task with just prompting, likely due to the small size of the models.  If we were to train a larger model on more tokens, we would expect to be able to get needle-in-a-haystack by prompting with both standard attention and scale-invariant attention.  Further, in the paper, every task we tried showed zero-shot length generalisation for scale-invariant attention and not for standard attention (Table 1, Table2, Figure 3); it is reasonable to expect that pattern to continue for prompted needle-in-a-haystack and on a larger model.\n\n\n[1]: Zhou, Yang, et al. \"GSM-Infinite: How Do Your LLMs Behave over Infinitely Increasing Context Length and Reasoning Complexity?.\" arXiv preprint arXiv:2502.05252 (2025).\n\n[2]: Hong, K., Troynikov, A., & Huber, J. (2025, July). Context Rot: How Increasing Input Tokens Impacts LLM Performance. Chroma Research.\n\n[3]: Li, Tianle, et al. \"Long-context llms struggle with long in-context learning.\" arXiv preprint arXiv:2404.02060 (2024).\n\n[4]: Barbero et al., Round and round we go! what makes rotary positional encodings useful? ICLR\n\n[5]: Liu, Aixin, et al. \"Deepseek-v3 technical report.\" arXiv."}]

Author rebuttal:

Meta-review: This paper proposes scale-invariant attention, a mechanism to improve LLMs' ability to generalize from short training contexts to much longer inference contexts. The authors rigorously define two key properties: scale-invariant total attention and scale-invariant attention sparsity—and mathematically derive a simple, position-dependent logit transformation to satisfy them. The method is empirically validated through a series of experiments, showing strong zero-shot generalization and outperforming baselines like YaRN and LogN on out-of-distribution context lengths.

The paper stands out for its strong theoretical foundation, moving beyond heuristic methods by deriving its approach from first principles. The proposed logit transformation is both theoretically sound and highly practical, as it is simple to implement and integrates easily with existing Transformer architectures. The empirical results, particularly the impressive zero-shot generalization for long-context setting, provide compelling evidence of the method's potential for real-world impact.

The primary initial weakness was the limited scale of the experiments, which was a concern for Reviewer TuDC, Reviewer yBmN, and Reviewer tNeo. This raised questions about the method's applicability to much larger, commercially deployed models. Another point of theoretical weakness is the reliance on the Gaussian assumption for the attention logits, a property that is empirically observed but not yet fully understood, as noted by Reviewer 6x9F and Reviewer tNeo.

I recommend acceptance because the paper presents a significant and well-supported contribution to the field. The authors successfully addressed the most critical concern regarding scalability. The zero-shot generalization performance is compelling and demonstrates the method's potential to enable cost-efficient LLM deployment. 

All reviewers initially raised concerns, primarily focused on the limited experimental scale and the robustness of the theoretical assumptions. In their rebuttal, the authors provided new experiments on a 7B Llama 2 model, which effectively addressed the scalability concerns raised by Reviewer TuDC, Reviewer yBmN, and Reviewer tNeo. The authors also clarified the theoretical assumptions to Reviewer 6x9F and Reviewer tNeo, explaining that their theory only requires "marginally Gaussian" variables. Lastly, the authors addressed the lack of comparisons with Infini-Attention and Ring Attention raised by Reviewer tNeo, providing new results against the former and a clear explanation for why the latter is a different class of method.
Decision: Accept (poster)