When Precision Meets Position: BFloat16 Breaks Down RoPE in Long-Context Training
Authors: Haonan Wang
Qian Liu
Chao Du
Tongyao Zhu
Cunxiao Du
Kenji Kawaguchi
Tianyu Pang
Venue: Accepted by TMLR ()

Abstract: Extending context window sizes allows large language models (LLMs) to process longer sequences and handle more complex tasks. Rotary Positional Embedding (RoPE) has become the de facto standard due to its relative positional encoding properties that benefit long-context training. However, we observe that using RoPE with BFloat16 format results in numerical issues, causing it to deviate from its intended relative positional encoding, especially in long-context scenarios. This issue arises from BFloat16's limited precision and accumulates as context length increases, with the first token contributing significantly to this problem. Despite its limitations, BFloat16 remains desirable for its computational efficiency, particularly given the substantial memory overhead required to extend the context window.  To improve long-context training under BFloat16, we develop AnchorAttention, a plug-and-play attention method that enhances long-context capabilities, and speeds up training. AnchorAttention reduces unnecessary attention computations, maintains semantic coherence, and boosts computational efficiency by treating the first token as a shared anchor with a consistent position ID, making it visible to all documents within the training context. Experiments on three types of LLMs demonstrate that AnchorAttention significantly improves long-context performance and reduces training time by over 50\% compared to standard full attention mechanisms, while preserving the original LLM's capabilities on general tasks.

Published in Transactions on Machine Learning Research (02/2025)
When Precision Meets Position: BF Breaks Down RoPE in Long-Context Training
Haonan Wang'
hoonan San eus.edu
National University of Singapore
Qian Liu
liugian &seaa
Sea AT Lab, Singapore
Chao Du Sea AT Lab, Singapore
duchoo &sea com
Tongyao Zhu
tongyao zhu nus... edu
National University of Singapore
Cunxiao Du Sea AT Lab, Singopore
ducr@sea.com com
Kenji Kawaguchi
kenjo@comp.rrus www.pmmppp..uuueo.oooom www..du............
National University of Singapore
Tianyu Pangl Sea AT Lab, Singopore
tiaryeppnn@@eea.com
Reviewed OpenReview: https://openreviiet net/ forum? id= f 23
Abstract
Extending context window sizes allows large language models (LLMs) to process longer quences and handle more complex tasks. Rotary Positional Embedding (RoPE) has become the de facto standard due to its relative positional encoding properties that benefit long- context training. However, we observe that using RoPE with BFloat16 format results in numerical issues, causing especially in long-context scenarios. This issue arises from BFloat16's limited precision and accumulates context length increascs, with the first token contributing significantly to this problem. Despile its limitations, BFloat16 remains desirable for its computational efficiency, particularly given the substantial memory overhead required to extend the context window. To improve long-context training under BFloat16, we develop Anchor Attention, a plug- and-play attention method that enhances long-context capabilities, and speeds up training. AuchorAtien Attention reduces unnecessary attention computations, maintains semantic coher- ence, and boosts computational efficiency by treating the first token as shared anchor with consistent position ID, making visible to all documents within the training context. Ex- periments three types of LLMs demonstrate that AnchorAttentio significantly improves long-context performance and reduces training time by over 50% compared to standard full attention mechanisms, while preserving the original LLM's capabilities on general tasks H
deviale from its intended relative positional encoding,
Work doue during Haonan Wang's interuship at Sca Al Lab. 1Correspondence to Tianyu Pang. AnchorContent: The implementation AnchorAttention supports popular madels, using Flaentoten and Flex Attention, and available httpp://githubb.com wawtoom/aaaan////aam Ancann/CCnnnncoCCeeeeen

---

Published in Transactions on Machine Learning Research (02/2025)
Attention Difference Shift
Attention Difference per Token 1st Token Difference over Longer Lengths
35
30
Attention Score Difference
Attention Score Difference
Attention Logit Difference
25
20
0.8
4.0
15
0.6
10
Pretrained (BFlcat16)
Random (8flost16)
Pretrained (Floet32)
0.2
3.0
0.0
12
100
200 1000
2000
2000
4000
8000
2000
6000
Shift Index
Token
Sequence Length
Figure 1: Effects of positional shifts on attention computations under different settings. Left: Attention difference D (Eq. 4 plotted against varying positional shift A (with Å₂ = 16 fixed). Pretrained models under BFloat16 (blue line) exhibit significant discrepancics compared to Float32 (yellow line) and random initialization (green line), indicating that the relative positional encoding property of RoPE is broken un- der BFloat16 and that pretraining amplifies this effect. Middle: Per-token attention differences wetn A 0 and A2 16, highlighting the first token accounts for most of the attention difference observed. Right: Attention logit difference (Eq. 5 for the first token as sequence length increases, showing increased discrepancies with longer sequences.
1 Introduction
In recent years, natural language processing has seen surge in models that handle increasingly long se- quence (Yang et al. 2024 Dubey et al. 2024 Jiang et al., 2023 Team et al.] 2024). A 128K token context window allows large language models (LLMs) to handle complex lasks such as multi-document question answering [Wang et al.) 2024a). repository-level code comprehension (Jimenez et al. 2024), and many-shot learning by capturing long-range dependencies (Agarwal et al. 2024), leading to more coherent and contex- tually relevant outputs (Mazumder Liu 2022).
In order to achieve long-context capabilities in LLMs, Rotary Position Embedding (RoPE) (Su et al. 2021
have emerged the dominant backbone for positional encoding (Liu et al. 2023b). The success of RoPE is often attributed to its trigonometric (rotational) properties and its relative positional encoding, which enable models to avoid out-of-distribution (OOD) rotation angles (Wang et al. 2024b Chen et al. 2023a 2023 Men et al. 2024). Thcorctically, by adjusting the rotary frequencios, exlended conlext position IDs (OOD) can be mapped to in-distribution ranges that have been sufficiently trained. In addition, due its relative positional encoding properties, interpolating into the original familiar range allows the model to recognize the relative positions of input tokens in the extended context. As result, relatively small amount of additional training over long contexts enables LLMs to adapt to extended context lengths (Zhao et al. 2024a Fu et al. 2024 Zhang 2023). However, minimal long-context training remains challenging due to the quadratic increase in GPU memory consumption with context length (Xiong ct al. 2023). To address this issue, Brain Floating Point (BFloat16) (Wang & Kanwar 2019 commonly used
Peng et al. 2023 LocaILLaMA
during pre-training, is also adopted in the long-context training phase. Its use reduces memory bandwidth requirements without significantly impacting model accuracy KKlamkar et al. 2019), making idcal for managing the computational demands of long-context models.
Despite the computational advantages of BFloat16, we have identified a critical issue: when combined with BFloat16, the relative positional encoding properties of RoPE are broken, especially in long-context scenarios. As shown in Figure [ this breakdown occurs because of BFloat16's limited precision. As the training window size increases, numerical errors accumulate, exacerbating the issue and resulting a more substantial discrepancy. In contrast, this degradation disappears when using Float32, which maintains the integrity of RoPE's relative positional encoding. Our empirical observations confirm that this breakdown diminishes the benefits RoPE offers for long-context, training.
2

---

Published in Transactions on Machine Learning Research (02/2025)
L
L
L
d1
d2
d3
d1
d2
d3
A
d
d2
d3
1
4
8
2
8
Figure 2: Illustrations of different attention paradigms. Left: Standard intra-document attention. Middle: Our improved version, intra-document attention with position ID reset per document. Right: AnchorAt. tention incorporating a shared anchor token, A.
To improve long-context training, we investigated the source of the breakdown in RoPE's relative posi- tional encoding and obscrved that the first token in the attention window contributes most significantly to this deviation. However, current methods do not explicitly consider this issue in their design. In the full attention mechanism, each token attends to all previous ones, resulting in cumulative deviations the window size grows. Standard intra-document attention (Figure 2 left) uses cross-document masking, effectively converting the large window into multiple smaller ones. This approach introduces multiple first tokens for each small window, and assigning different position IDs to these first tokens causes inconsistencies in the model's positional understanding. Empirically, we found that simply resetting the position IDs to maintain consistency across windows (our version of intra-document allention, Figure 2 middle) improves longecontext performance. This confirms that position ID inconsistencies are key issue. Yet, resetting position IDs introduces
problem: the model can only learn the full spectrum of rotational angles when processing data sequences that reach exceed the maximum context length. Building on these insights, we propose Anchor Attention versatile and plug-and-ple attention method that enhances long-contert context capabilities while accelerating the training process. As illustrated in Figure 2 of AnchorAttention lics in treating the first token as a shared anchor: we always assign the first position ID, making it visible documents are invisible to each other. By having all documents share the same initial token, AuchorAt- tention eliminates inconsistencies while allowing the model learn the full rotational span from sequences shorter than the context length. Additionally, AnchorAttention reduces the number of tokens involved in attention computations by not attending to all previous tokens, which helps prevent the rolling accumula- tion of numerical errors. Experimental results demonstrate that training with AnchorAttention consistently outperforms full attention, standard intra-document attention, and our improved intra-document attention on the long-context benchmark RULER, across lengths from 8K to 128K. On real-world long- contect bench- marks like LongBench, Anchor Attention improves in-context lcarning performance while largely preserving the model's capabilities on general tasks such MMLU and HellaSwag.
(Right), the core innovation
documents within the context window while ensuring that tokens from different
In summary, the main contributions of this paper are follows:
We found that the relative properties of RoPE are compromised under BFloat16 precision.
We identified that the first. loken of sequence contributes to the deviation of RoPE's relative properties, which should be preserved in theory. Moreover, this deviation becomes more pronounced with larger training window sizes.
Based on these observations, we introduce a practical approach, AnchorAttention for long-context continuous training, which improves the model's ability to handle long contexts, utilizes less than 50% of
3

---

the typical base value for current LLMs is Published in Transactions on Machine Learning Research (02/2025)
the training time required by standard attention training, and requires minimal modifications to existing training pipelines.
2 Discrepancies in RoPE's Relative Positional Encoding
2.1 Background of Rotary Position Embedding (RoPE)
Modern LLMs are primarily based on the Transformer (Vaswani et al. 2017 architecture. One of its core components the attention mechanism, which can be written as:
Aij- q kj ATTN(X) softmax + M)*),
(1)
(2)
where the attention logit matrix A A (A,,) RTXT, the query and key vectors take the form q, Wox; and kj W WXX,, respectively. Wq. WK E Rdxd are the query and key matrices, and X,, x, represent the matrix, ATTN(X), 1/vd introduced by Vaswani et al. (2017) to simplify notation. The causal mask, M E RTxT where M -o if i < and M, kj = otherwise, prevents each token from attending to future tokens by setting the upper triangular part of A to negative infinity.
i-th and j-th token of X R In the attention
ignore the scaling factor
Transformer architectures employ positional embeddings to encode sequence order information, which is essential for processing sequential data. RoPE (Su et al. 2021) are the currently most widely adopted encodings, especially in LLMs. RoPE acts on the query and key by splitting them into 2-dimensional chunks and rotating each chunk at different froquency, which applies the rotation matrix into the calculation of the altention logit in Eq. 1 which can be written as:
Computed within FlashAttention2, BFloat16
Aú (Ri,og;) Rzoky of Ri-i,okg- 2i Rm,oki
Applied before FlashAttention2, Float32
where m = j --jiii is the relative distance of and j. We leave the details of rotation matrix R in Appendix A RoPE achieves efficient relative position encoding by implementing it absolute position encoding. This is done by applying rotation matrix directly to the key and query vectors, specifically Rg,g:: and Ry.ok,, mirroring the operation of absolute position encoding. Note, the rotation matrix is applied outside the FlashAttention2 (Dao 2024) module using Float32 precision. But, the inner product of the key and query occurs within the FlashAttention2 module and must be cast to BFloat16, as required by FlashAttentio22. And, the selection of rotation angles satisfies 0: base 2i/d 22/d 10,000.
2.2 BFloat16 Disrupts RoPE's Relative Properties, Especially with Long Context
The RoPE is theoretically designed to depend solely on the relative positional distance m j i bet ween tokens. A key implication is that adding constant positional shift A every position index should not alter the attention computation. Formally, it can be expressed as:
A(444414) = (H.44.9g) T (Ry = q. RG (11A), = q Rm, kj = Asj.
(3)
where Aij is the attention logit between positions i and j, Rm, is the rotation matrix corresponding to the relative distance m (derivation in Appendix Ab. However, key question is, does this relative property hold in practice?
We empirically investigate the impact of positional shifts attention computations. In order clarify our analysis, we redefine the attention computation A (X, A), representing the attention matrix produced by the h-th head in the 1-th transformer layer when a positional shift is applied. Specifically, if the original position indices are (0, 1, 1...,, L I-11,, after applying the shift they become A -
4

---

Published in Transactions on Machine Learning Research (02/2025)
We measure the difference in attention due to two different positional shifts A1 and A2 using the following metric:
ΣΣ
ATTN (x, A (X.
1.
This metric measures the cumulative difference attention scores across all layers and heads. And the normalization vector
T-T ....: is applied element-wise (denoted by @) to account for the varying number of clements in lower triangular matricos duc to causal masking.
Experimental Setup. In Figure [ (left), we set A2 = 16 and vary A over the list [0, 2, 4, 6, 9, 10, 12, 14, 15, 17, 18, 20, 22, 50, 100, 200, 500, 1000, 2000]. The blue line represents results obtained using the pretrained parameters of LLaMA-2-7B with BFloat16 precision. The yellow line shows results using the same pretrained parameters, but converted to Float32 precision. The green line corresponds to results when using random initialized parameters (initialized with kaiming. uniform kaiming.umiiform, in PyTorch by default) in BFloat16 precision. For each case, compute the difference D, averaging over 50 text sequences, each with length of T = 4096. Note, we shop the A₁ = 16 to improve visualization. As expected, the difference is zero when both A and D are equal. More detailed visualizations and discussions are provided in the Appendix
Results Discussion. We observed that when using BFloat16 precision, the positional shift A affects the attention computations of the pretrained LLaMA-2-7B. In contrast, when all computations are performed with Float32 precision, this effect disappears (as shown by the blue green lincs). Additionally, when comparing pretrained paramelers to randomly initialized ones, both under BFloat16 16 precision (blue vs. yellow lines), we found that pretraining amplifies the discrepancy. These findings suggest that under BFloat16 precision, the RoPE deviates (slightly) from its theorctically claimed relative positional encoding, and this property breaks down after the model pretrained.
Furthermore, we delve into the details of the attention difference between A and A₂ 16 for each individual token (i.c., without summing over fow in metric for per-token study). We visualize these per-token differences in Figure I (middle). The results clearly indicate that the first token contributes most significantly to the attention difference. When the first token is excluded, the remaining tokens retain the relative positional property.
In the previous experiments, quence length affects the attention logit difference, we extended our study to sequence lengths of [64, 128, 256, 512, 1024, 2048, 4096, 8192]. In this trial, we measured the attention logits instead of the at- tention scores because the softmax operation over varying sequence lengths would render the results for the first token incomparable across different lengths. For clarity, we define the attention logit difference as:
fixed the sequence length T 4, 096. To investigate how se-
Diogit T
A ( - A/t (A2))
(5)
1A 11
where (A) is the attention logit matrix with positional shift A for the l-th layer's 's h-th head. We keep A = 0 and A₂ = 16 to measure the metric Diogit averaged over 50 sequences with different context lengths. Figure [ (right) summarizes these results. We observe that as the sequence length increases, the attention logit difference for the first token also increases. This indicates that longer sequences exacerbate the discrepancics in the first token.
Summary of the Section: Under BFloat16 precision, the relative property of RoPE is broken. The discrepancies are primarily sourced from the first token, which significantly contribudes to the break- down. Additionally, as the context length increases, the impact of these discrepancies becomes more pronounced.
5

---

Published in Transactions on Machine Learning Research (02/2025)
3 AnchorAttention
In the previous section, we demonstrated that BFloat16 precision compromises the relative positional en- coding of RoPE,, with discrepancies becoming more pronounced as sequence length increases. heless BFloat16 remains desirable in practice due to its computational efficiency, which is especially crucial for long-context extensions. The most effective strategy in these sccnarios involves continuing pre-training on long-context data (Fu el al. 2024 Gao et al. 2024). However, this approach is challenging because attention mechanisms incur a quadratic computational cost relative to sequence length. To mitigate the increasing errors with longer sequences while retaining the advantages of BFloat16, we propose Anchor Attention.
3.1 Training CongtCon Models with Small Window Sizes
Training a long-context model typically requires processing long sequences. However, this direct way exac- erbates precision errors under BFloat16, as shown by Figure | (Right) with the number of processed tokens increasing, the discrepancy will exacerbale. Reducing the sequence length could mitigate the precision issues, but it seems contradictory to train long-context model with short window sizes. To reconcile this, re- visit techniques from the literature and we found that the intra-document attention (Zhao et al. 2024b) that mask out cross-documett attention (as shown in Figure 2 Left) can be used to train a long-contex: model with less sequence length compared with full attention. Besides, empirically, the intra-document attention has been successfully applied in open-source LLM, like LLaMA-3 series. recent work by Gao et al. (2024) also verifies the effectiveness of intra-docunent attention show that masking out attention across document boundaries improves both the short and long-contess performance.
3.2 Does the Discrepancy in RoPE under BFloat16 Impact Long- Performance?
Previously, we identified deviation the relative positional encoding of RoPE when utilizing BFloat16, as evidenced by differences in allention scores. This observation prompts the question: does this discrepancy significantly affect long-context performance? Especially, considering intra-document attention appears to mitigate discrepancies in attention scores, the impact of BFloat16 might be negligible on long-context ca- pabilities, when the model is trained with intra-document attention. To investigate this, we compare two types of position indices. The first assigns position IDs continuously from the start to the end of the se- quence (Figure Left), following the common practice in intra-document attention Zhao et al. (2024b): Gao et al. (0244.. The second resets the position index to within each document (Figure 2 Middle). Theoreti- cally, both positioning schemes should yield identical results. This comparison is used to assess whether the deviation in relative positional encoding further leads mcasurable differences in long-context performance.
Experimental Setup. We conducted experiments by training a 128K model based on the ALaMA-2-77 architecture using the Slimpajama dataset with intra-document attention and evalu- ated them on the widely used long-contexl evaluation bench- mark RULER. These experiments compared scenarios with and without resetting the position IDs, aiming to assess the impact of position ID assignment on model performance. We defer the training and evaluation protocols to Section 5
RULER Performance 128K Model
90
Doc Attention
Intra Doc Attention Reset
85
Performance (%)
80
75
Results Discussion. The results are presented in Figure We observe that resetting position IDs consistently enhances the model's long-context performance. In contrast, the con- ventional approach, which not resetting position IDs (Figure 2 Left), results in inferior performance. This performance differ- ence may arise from assigning different position IDs to the first. token of each document, introducing inconsistencies that can potentially confuse the model. By resetting the position IDs, we maintain a consistent positional relationship (first position ID is on the first token). Note, this analysis is based on the
70
65
4K
8K
16K
32K
64K
128K
Figure 3: Resetting position 1Ds improves per- formance, contradicting theoretical predictions of RoPE

---

Published in Transactions on Machine Learning Research (02/2025)
hypothesis that RoPE functions in some way absolute positional encoding mechanism under BFloat16 precision. Further research is necessary to rigorously validate this assumption and fully understand its implications.
3.3 A Little Goes A Long Way: A Shared Anchor Token for Efficient Long- Context Training
Restarting position IDs enhances model performance but introduces a significant. drawback: the model can only learn the full spectrum of rotational angles when processing sequences that reach or exceed the context length. This limitation hinders the model's ability generalize to longer context length scenarios because, as we increase the context window size, collecting sufficient long sequences to fill the entire context window becomes impractical due to the scarcity of such lengthy data.
To address this challenge, propose Anchor Attention, an attention mechanism that introduces a shared common anchor token visible to all documents within the long context window. Additionally, redundant attention connections across documents are ignored, allowing the model to focus on coherent information. An illustraatin of AnchorAttenting is provided in Figure (Right).
In AnchorAttention, the anchor token A serves as the common starting point for all documents. This design is motivated by our observation that the first token is responsible for deviations from relative positional encoding in RoPE; except for the first token, the subsequent tokens can preserve the desired relative positional encoding properties. By introducing shared anchor token with fixed position ID (the starting ID), the model is no longer confused about the relationship between the beginning of a document and the starting position ID. Specifically, we designate the begimning-of-sequeen Since the <bos> token typically lacks explicit semantic meaning, any precision errors caused by BF concentrated on this token have minimal impact on overall model performance.
token, <bos>, as this common anchor.
Furthermore, with the shared anchor, we eliminate the need to restart position IDs to ensure that the model correctly interprets the relationship between the beginning of sentences and the starting position ID. This allows us to use continuous position IDs from the start to the end of the window length (the non-resctting scheme, as shown in Figure 2 Right). Consequently, a full spectrum of rolational angles can be trained in each iteration. Compared to the intra-document attention approach, our method does not heavily depend on abundant long sequences that can fill the entire training window length train the full rotational span. This reduces the dependency on collecting and upsampling long-sequence data.
Summary of the Section:
1. Devialions RoPE's relative positional encoding observed in atlention scores under BFlont16 pre- cision also affect the model's ability handle long sequences effectively.
2. Resetting position IDs within attannattnnion mechanisms improves long-contert performance by maintaining consistent positional relationships.
3. A Little Goes Long Way: Introducing a shared anchor token provides a simple yet powerful solution for contecnntett training.
4 Long-Context Extension Protocol
In this section, we focus on the optimized training strategics and robust evaluation metrics, detailing the se- lection of base models, the long-context training dataset, specific training configurations, and the benchmark setup.
4.1 Base Model, Dataset and Training Configuration
Base Models. In our experiments, we primarily use the L1aMA-2-7B model as the base model. To assess the effectiveness of our proposed methods across different pretrained models, we also evaluate Llama-3-88 (Dubey et al. 2024), Qwen-1.5 (1.8B) Yang et al. 2024), and Mistral-7B-vo.3 3 (Jiang al.

---

Published in Transactions on Machine Learning Research (02/2025)
2023 in Section 5.3 These models represent range of architectures and base model pretraining paradigms, allowing us to test the generality of our approach.
Dataset. We use the SlimPajann dataset (Soboleva et al. 2023) for long-context training, an open-source replication of the LLaMA pretraining data mixture (Touvron et al. 2023 The dataset includes diverse text sources, such as Common Crawl (CC), C4, GitHub, ArXiv, Books, Wikipedia (Wiki), and StackExc In addition the original Sajmaajama, we apply the data mixture method from Fuu et al. (2024), which upsamples long sequences within each source while maintaining the overall distribution of the original dataset. This upsampling aims to better expose the model to mong-conextt scenarios during training.
In our experiments, we sample billion tokens from both the original and upsampled SlimPajama datasets. We refer to the datasets with sequence lengths of 64K and 128K lokens as SlimPajama- and MimPaajaa--28K, respectively. The upsampled dataset is denoted as UpSampledMix-122 Detailed statistics and discussions on data enginecring are provided in Appendix C
Training Configuration. Our training hyperparameters are primarily based on Zhang 2023 All models are trained on 8 NVIDIA A100 GPUs. We set the learning rate to 2 x 10 and use the AdamW optimizer with weight decay of 0.1, $1 0.9, and 32 0.95. Each model is trained for 2000 steps, which corresponds to approximately epoch over the 2 billion token dataset. The batch size is set to 8, equating to 0.5 million tokens per batch for 64K context and million tokens for 128K context lengths. Table summarizes the training configurations.
Table 1: The training Configuration.
Long-context Continuous Training
Data
UpSempledMix / SlimPajemal22K/ SimPajama661KK
[JSSapleedMixx128K
58% CC, 20% C4, 7% GitHub, 6% ArXiv, 5% Books, 4% Wiki, 2% StackExchange
SlimPajama-
53% CC, 27% C4, 5% GitHub, 5% ArXiv, 4% Books, 3% Wiki, 3% StackExchange
SlimPajama-64a 64K:
54% CC, 25% C4, 5% ArXiv, 5% GitHub, 4% Books, 3% Wiki, 3% StackExchange
Model
Initialization:
Llsma-2-7B Llama-3-8B/ / D0wen-1.5-L8BB Mistral-7B-v0.3 16K: x 10", 64K: 5 x 10", 128K: 1 10"
RoPE:
Attention:
Full attention/ Intra-doc attention / Iutra-doc attention with Reset Anchor AnchorAttention with Tag
Optim.
AdamW (weight docay
U.1, 81 = 0.9, Ag 0.95)
LR:
2e
steps (0.5M token 64K, 1M tokens for 128K)
Steps:
Batch size:
4.2 Controllable Study with Meaningful Evaluations
Measuring Long-Context Ability with Appropriate Metrics. To conclusively evaluate the effective- ness of our proposed attention mechanism in enhancing the long-context capabilities of base models, is important to use robust and suitable evaluation metries.
Perplexity (PPL.) is commonly used to evaluate long-context language models; however, recent studies ques- tion its reliability for assessing long-text understanding, Hu et al. (0044) show that PPL poorly correlates with a model's ability to comprehend long-ran range dependencies because primarily measures local informa- tion capture. Additionally, Fang al. (2024) provide empirical evidence that PPT. overlooks key tokens crucial for undorstanding long-context inputs, leading to unrcliable assessments. Morcover, Gao al. (2024) find that while increasing the amount of long data improves PPL, exclusively using long data can significantly degrade downstream contecontext performance. In line with these observations, PPL remains unchanged after the initial training steps, performance on the RULER benchmark (Hsich) et al. 2024) continues to improve (in Figure further indicating that PPL may not adequately reflect enhancements in long-context performance.
also found that while
Existing benchmarks for long-context language models such as HELMET Yen et al. 2024 Long- Bench (Bai al.) 22023), LongBench- Cite (Zhang al.) (0244)), InfiniteBench [Zhag et al. 2024c). and NoCha (Karpinska et al. 2024 heavily rely instrontion-following abilities, making them unsuitable for
8

---

Published in Transactions on Machine Learning Research (02/2025)
evaluating long-context base models. Moreover, tasks in LongBench (Bai et al. 2023 addressed with a context. length of 16K tokens. While 16K tokens represent medium-lenglh context, this does not fully challenge models like ours that can handle 64K or even 128K tokens. Despite these limita- tions, we still evaluated our models on LongBench because includes real-world long-context tasks, and its lew-shot in-conlext learning (TCL) lasks are appropriate for assessing base models. Notably, the effectiveness of few-shot ICL in evaluating base models is also observed concurrent work by Gao et al. (2024).
be adequately
To more thoroughly assess model's long-context capabilities, employ the RULER benchmark (Hsich) et al., 2024), which is specifically designed for extensive context lengths. RULER evaluates abilities such (1) locating specific data within vast content, (2) tracing relationships across broad contexts, (3) aggregating and quantifying dispersed information, and (4) distinguishing relevant details from extraneous information in complex queries. The benchmark includes tasks across various Lategories-Neette -Needle-in-aa a-tatsack (NIAH), Variable Tracing (VT), Common and Frequent Words (Aggrogation), and Question Answering (QA)- each
targeting different aspects of long-context processing that more accurately represent the capabilities of models like ours.
Perplexity and RULER Performance over Steps
83
120
100
81
RUL Performance
Perplexity
80
80
60
78
40
Perplexity RULER Performance
77
76
0
100
500
1000
1250
1500
1750
2000
2190
2200
2300
2000
2010
2000
2000
2000
2000
Steps
Figure 4: RULER performance varies during long-context training, recommend reporting the averaged
RULER performance rather than just the final training step. PPL remains unchanged after the first several steps, failing to reflect improvements in long-context ability.
However, unlike the official evaluation, which measures overall performance across all 13 tasks, we exclude two tasks (and NIAH-Mulikky 3 and Common Word Extraction) for the LLaMA-2-7B model. This exclusion is based on the assumption that long-context training primarily extends a model's existing abilities from its original context window to longer sequences, and on the observation that LLaMA-2-78 struggled with these two tasks even within its original 4K context window. It is therefore unreasonable to expect the model to acquire new skills in extended contexts that it could not perform within its original window. (Detailed results for both LLaMA-2-7B and LLaMA-2-7B-Chat across all 13 tasks along with a comprehensive discus- sion are provided in Appendix D Note that for other advanced models (LLaMA-3- --B Mistral-7B-vo. Qwen-1.5-1.89
the cross-architecture evaluation section (Section 5.3 we use all 13 tasks.
Moreover, unlike previous studies that report RULER performance single checkpoint, we suggest report- ing the average RULER performance across five checkpoints, saving a checkpoint every 10 steps over the last 50 training steps. Although retraining the model five times with different random seeds is another reasonable method for averaging performance, it would significantly increase the experimental cost. To demonstrato the necessity of averaging over multiple checkpoints, we evaluate the RULER performance of LLaMA-2-7B trained with standard full attention on the original SlimPajama with a 16K window size. As shown in Figure 4 the performance on the 16K RULER benchmark exhibits fluctuations during training. To prevent cherry-picking specific checkpoints, we suggests that future research report the averaged performance across several checkpoints.
9

---

Published in Transactions on Machine Learning Research (02/2025)
LLLM General Ability. To ensure that continual pretraining preserves the models' core abilities within their initial context length, we use the MMLU (Hendryeks et al. 2021 and HellaSwag (Zellers et al. 2019) benchmarks. MMLU assesses knowledge and reasoning across 57 diverse subjects, including STEM, human- itics, social sciences, philosophy, law, and modicine, providing a broad mcasure of general understanding. HellaSwag lests commonsense reasoning by having models choose the most plausible continuation of a sce- nario, challenging their contextual prediction skills. These benchmarks are standard choices for evaluating foundational abilities, and we observe minimal performance differences between the base and chat models, supporting their suitability for evaluating base model capabilities.
4.3 Optimizing Positional Embeddings for Long- Context Training
To find the optimal positional embedding for long-context training, trained LLaMA-2-7B on the SlimPajama-644 aam dataset, chunked into sequences of 64K tokens. We evaluated three methods: RoPE (Su let al. 2021b, NTK-aware {LocaILLaMA 2023), and YaRN (Peng et al., 2023). Following Men et al. 2024), we varied the base value in the embedding formula 0, = base- 22yd (Fquation G corresponding to the rope- theta paramcter in Transformers (Wolf et al. 2020). NTK-aware and YaRN depend on the relative scale 8 = Tnew//Torigh where Torigin is the original context size (4K context length) and Tnew is the extended size (64K context length), so we set s = 16. Following (Fu et al.) 2024), which suggests that 1B tokens suffice for learning long-context abilitics, we use 1B tokens in each training to balance training time and the evaluation of multiple models.
Table 2 shows the performance from 4K to 64K tokens the RULER benchmark. Vanilla RoPE with suitable base value performs best within the training context length, outperforming NTK-aware and YaRN methods at 64K context length. To measure performance beyond the training context length (generalization) we evaluated the models on the 128K RULER benchmark (Table YaRN embeddings generalize better to unseen longer contexts; while vanilla RoPE performs well within the trained length, its performance declines at longer contexts. Recognizing the importance of the RoPE base value, we further studied rope- theta selection by training models with a 32K context length to save computational resources. As shown in Table 3. performance peaks at specific theta value, with marginal gains beyond that point and slight declines when increasing theta further, aligned with observations from (Liu et al. 2023b Men et al. 2024). Based on these findings, we adapt base theta values with context lengths: 1M for 32K, 5M for 64K, 10M for 128K, and 50M for 256K. This approach ensures optimal performance across different context lengths.
Table 2: Long-context models trained with vanilla RoPE outperform those with NTK or YaRN when evaluated within the exlended context window size. The Free" indicates models evaluated directly with- out training.
Table 3: An appropriate base is crucial for long- context training; using values that are too large too small can degrade performance.
Base
32K
16K
8K
4K
Avg.
500
8,52
16.43
21.18
37.06
20,80
10K
15.21
28.79
43.61
83.03
42.66
200K
68.91
76.67
85.64
88.00
79.81
Base
64K
32K
16K
8K
4K
600K
74.61
81.79
85.03
89.36
82.70
900K
75.05
81.62
85.22
88.65
82.64
Froe*
RoPE 10K
85.29
5M
76.09
77.15
85.66
87.00
81.48
YaRN 10K
0.18
1.09
6.02
18.38
29.06
1B
73.97
76.48
81.09
82.33
78.47
RoPE
1M
66.89
70.79
77.02
84.64
90.42
Table 4: Evaluating beyond the extended context window, YaRN show advantages.
5M
68.81
72.68
81.03
85.79
89.28
10M
69.43
71.01
81.01
85.28
88.13
Base
128K
64K
YaRN
1M
64.22
65.69
73.28
79.98
86.02
RoPE
IM
35.93
(-30.95)
66.88
5M
63.48
64.47
75.29
79.90
84.01
5M
53.98
(-14.83)
68.81
10M
58.91
62.03
70.68
75.62
83.04
10K
27.25
(-22.84)
50.09
IM
57.07
(-7.15)
64.22
YaRN
10K
47.32
53.05
60.62
72.42
81.61
5M
53.30
(-11.18)
64.48
NIK
IM
62.47
70.49
75.87
83.48
86.41
10M
50.85
(-8.06)
58,91
5M
61.45
68.52
75.90
78.47
85.17
10K
24.31
(-23.01)
47.32
NTK
1M
35.53
(-26.94)
62.47
10M
58.75
63.95
72.27
77.43
83,99
5M
41.87
(-19.58)
61.45
10

---

Published in Transactions on Machine Learning Research (02/2025)
Summary of the Section:
1. The RULER benchmark is effective for evaluating the long-contert training of base models. But, it is advisable to average performance over last several checkpoints to mitigate variance.
2. (Vanilla) RoPE is effective, but you need to carefully select the base volue.
Vanilla RoPE with properly chosen base value excels within the training contert length but may underperform beyond it.
YaRN offers better generalization to contezts longer than the training window.
good base value for ALL (RoPE-baseds positional embeddings is critical to achieve optimal performance; values that are too large or too small can degrade performance.
5 Experimental Results
5.1 AnchorAttention Performance on RULER
We continued pretraining the LLaMA-2-7B model three datasets SlimPajama-64K SlimPajama-1255 and "pSSaplleMMix-128KK Attention, and our proposed AnchorAttention We then evaluated their performance on the RULER bench- mark (Table 5 We observed that resetting positional IDs generally improved the performance of Intra- Document Attention, except some cases where the differences were negligible (e.g., 4K tokens for the SlimPajama-6e 64K and UpSampledMix--225 128K datasets). Our proposed Auchor formed Full Attention and Intra-Document Attention across all datasets and sequence lengths, particularly excelling at longer contexts. On all three datasets, AnchorAtlention achieved top scores at every length, consistently outperforming other baseline methods. We also observed that the UpSampledMix- 128K dataset improves model performance when training with Full Attention and Intra-Document Attention mechanisms. However, when using our proposed Atteenntt for long-context training, the performance gap between models trained on SlimPajama-122s and UpSampledMMx--15 potential for AnchorAttention to reduce dependency on carefully upsampled data, simplifying future data preparation processos.
varions attention mechanisms, including Full Attention, Aatra-Document
consistently outper-
significantly reduced. This suggests the
5.2
What Works and Doesn't in AnchorAttention: Domain Tagging and Interleaved Chunks
We further investigate two data utilization strategies that may enhance the performance of AnchorAttention: domain tagging and interleaved chunks, as illustrated in Figure 5
Domain Tagging: This strategy prepends each training lext sequence with a domain identifier (e.g., "Wikipedia" or "CommonCrawl"", masking out the loss from these tags during training. This allows the model recognize and potentially prioritize information from specific domains. Previous studies (Allen-Zhu
& Li 2024 Zhang et al. 2024b) suggest that domain tagging can optimize knowledge storage and aid in selective learning, retaining information while avoiding conflicts.
Documents are segmented into multiple chunks at random split points, shuffled, and recombined into new sequences, with the original order chunks within each document preserved (e.g., the second chunk of a document always appears after the first chunk in the newly organized data sequence). Zhao al. (2024a) employed this technique to generate synthetic long-context data, effectively training models to handle extended contexts.
Interleaved Chunks:
In Table 5 our experimental results show that incorporating interleaved chunks (Anchor.Atentiio + Inter- leaved Chunks) consistently degrades performance compared to the base AnchorAttennion. To investigate whether this is due lo an incompalibility between AnchorAutention and interleaved chunks, we conducted additional experiments using nutra-document attention with interleaved chunks on the SlimPajama--64K and SlimPajama- 128K datascts. This combination also resulted in worse performance, even falling below that of the baseline Full Attention method. Notably, previous work (Zhao al.) 2024a) demonstrated that inter- leaved chunks combined with Full Attention can improve performance. Therefore, hypothesize that the
11

---

Published in Transactions on Machine Learning Research (02/2025)
L
L
L
d
T2
d2
d
d2
d
d
A
d1
dg
dg
db
3
3
5
3
5
6
7
9
Figure 5: Illustrations of domain tagging and interleaved chunks. Left: AnchorAtienteo with domain tagging, where Ji denotes the domain of document d1. Middle: Intra-document attention with interleaved chunks; documents are split into shuffled, interleaved chunks, preserving the original order within each document. Right: Attennin with interleaved chunks.
cross-document attention masking is incompatible with interleaved chunks. Our proposed AnchorAttention gencrally improves performance over the bascline Full Attention but less effective than using raw data, the coherence of documents preserved, with Anthontten Regarding domain tagging, adding domain tags + Tag) does not consistently improve upon the base AnchorAttention the SlimP PaimPajama-64K dataset, it achieves slightly higher score at 64K tokens (73.88 vs. 73.25), but at 32K tokens, it performs slightly worse. Similar trends are observed on the SlimPajama-122 128K and UpSampledMix- 128K datasets;
For example,
+ Tag sometimes surpasses AnchorAttention at specific token lengths, but the base AnchorAttention generally performs better overall. Overall, those observations suggest that incor- porating domain tagging can sometimes improve long-context performance, interleaved chunks consistently degrade performance when used with rross-document attention masking.
5.3 Cross-Modes Evaluation of AnchorAttention for Long-Conted Performance
To evaluate the generalizability of AnchorAttention across various model architectures, we assessed its long- context performance on multiple pretrained models. As shown in Table D our proposed AnchorAtten- tion mechanism consistently outperforms the standard Full Atlention across different models and context lengths, particularly with longer sequences. For all three models, LLaMA-3-8B, Mistral-7B-vo... and Qwen-1.5-1.83, AnchorAttention yields higher scores across most sequence lengths. Note, the Qwen series does not utilize bos token, case of the LLaMA-3-8B model with context length of 128K tokens, Anchor Attention achieves score of 51.49 compared to 34.02 with Full Attention. Incorporating domain tags AAtentotttention + Tag) some- times further improves performance, as observed with the "iattral-BB-v.. model at 128K tokens, where the score increases to 49.61. These results confirm that AnchorAttente effectively enhances long-conte context performance across different model architectures, demonstrating its robustness and broad applicability.
the eos token the shared anchor in implementation. In the
5.4 General Performance on Medium- and Short-Conter. Tasks
Enhancing long-context capabilities in language models often introduces a tradc-off with performance on medium- and short-contexl lasks, as highlighted by previous research (Xiong et al.) 2023). Our objective is not only to improve long-context understanding but also to ensure that our models retain strong performance
on tasks with shorter contexts. To evaluate this balance, we tested our models on LongBench (a modium- context evaluation), as well HellaSwag and MMLU (short-context benchmarks). Table summarizes these results. The findings show that models trained with Anchor Allention AnchorAllention better preserve general abilities
12

---

Published in Transactions on Machine Learning Research (02/2025)
Table 5: Results on 64K and 128K Tokens Datasets. Highest scores across all methods are shown boldface. Within the Intra-Doc Attention category, the higher scores are underlined. AnchorAttention and its variants, outperforming other methods, are highlighted with a background color.
Attention Mechanism
128K
64K
32K
16K
8K
4K
SlimP SiimPajama-44
Full Attention
66.40
71.78
77.63
83.86
89.84
Intra-Doc Attention
69.97
74.70
79.15
83.50
89.62
Resel
70.03
74.18
80.27
84.51
89.52
+ Interleaved Chunks
60.59
66.52
71.70
79.70
84.71
AnchorAttention
73.25
75.97
82.91
85.48
90.69
+ Tog
73.88
74.21
82.46
85.13
89.93
+ Interleaved Chunks
66.77
69.73
77.81
85.35
89.31
128K
Full Attention
62.75
70.56
71.38
81.65
83.61
88.85
Intra-Doe Attention
64.31
70.87
72.07
82.60
84.11
88.98
+ Reset
65.75
73.34
73.30
82.82
84.43
90.01
+ Interleaved Chunks
53.74
61.08
65.51
75.25
80.59
82.71
AnthorAttention
66.15
77.69
74.28
83.67
86.41
90.60
+ Tag
65.46
74.67
75.77
83.07
84.07
89.09
UpSampledMiz-Pee 128K
Full Attention
63.70
71.45
72.69
82.57
84.55
90.08
Intra-Doc Attention
63.96
74.52
76.53
82.46
86.61
90.35
+ Reset
64.10
74.55
77.73
82.82
87.16
89.98
65.24
76.11
79.51
86.54
87.43
90.44
+ Tag
66.85
73.52
77.18
81.62
84.90
89.01
Table 6: Attention Mechanism Performance Across Different Models and Token Sizes
Attention Mechanism
128K
64K
32K
16K
8K
4K
LLaMA-3-8A
Full Áttention
34.02
61.80
72.09
79.99
82.43
83.68
AnchorAttention
51.49
70.99
83.06
86.90
88.09
88.72
+ Tag
49.67
70.37
84.14
87.13
88.36
88.97
Mistral-7B-vo.3
Full Attention
45.64
49.05
54.49
64.06
69.99
72.80
47.46
61.26
68.53
73.47
76.06
78.94
+ Tag
49.61
56.80
64.13
69.47
74.65
77.34
Qwen-1.5-1.8B
Full Attention
33.56
41.77
47.01
56.15
61.33
67.26
Anchor
34.32
44.31
48.63
56.90
62.62
68.61
Tag
35.84
43.91
50.70
57.39
61.96
67.41
from the pretraining stage compared to those using full attention intra-document attention. Importantly, even when compared to the original LLaMA-2-7B, models with Anchor remain competitive short- context tasks. On the HellaSwag dataset, the Anchor model trained on SlimPajama-64I, achieves a score of 70.78, closely matching the original LLaMA-2-7B's score of 71.39. On the MMLU benchmark, the Anthottttention + Tag variant attains a score of 42.85, nearing the baseline of 46.66. These results indicate that while significantly improving long-context capabilitics, AnchorAffenttin effectively maintains performance on medium- and short-context tasks without substantial trade-offs. Additionally, incorporating domain tagging effectively preserves the general abilities of large language models.
13

---

Published in Transactions on Machine Learning Research (02/2025)
Table 7: Results on LongBench ICL, HellaSwag, and MMLU datasets.
Attention Mechanism
LongBench ICL
HellaSwag
MMLU
LLaMA-2-7B
6.22
71.39
46.66
Slim (annnuajjmm-644
Full Attention
62.51
68.50
33.93
Intra-Doc Attention
62.79
71.01
36.94
+ Reset
63.76
70.12
37.92
65.38
70.78
40.32
+ Tag
66.02
69.10
40.67
SlimPajama-1295
Full Attention
50.72
69.46
37.93
Intra-Doc Attention
51.22
69.93
39.49
+ Reset
50.07
69.88
37.42
AnchorAttention
51.85
70.51
41.63
+ Tag
51.89
70.37
42.85
UpSumpledMMe- 128K
Full Attention
48.96
67.62
40.58
Intra-Doc Attention
49.51
70.86
41.27
+ Reset
50.18
70.97
40.79
Anchor Attention
50.17
70.11
41.15
+ Tag
50.70
68.97
42.03
5.5 Infrastructure and Engineering
Efficient implementation is erucial for training long-context language models. We introduce AnchorContezl, codebase that offers the AnchorAttention mechanism compatible with various models, including the LLaMA series, Mistral series, and Qwen2 series. To enhance compatibility with a wider range of frameworks, it provides two computational engine options: FlexAttention (which will be natively supported in PyTorch 2.5.0) and FlashAttention (currently the most commonly used).
To ensure the reliability our experimental conclusions, we focused on comparing our method with existing techniques in terms of accuracy, speed, and ease of integration:
Accuracy. Distributed training can introduce numerical errors that affect model outputs, making it essential to assess whether method maintains numerical consistency across different computational selups. We evaluated the numerical accuracy of our method by measuring the differences in logits (the outputs of the forward pass) when processing the same training data sequence with a context length of 32K on 8 A100 GPUs. Specifically, we compared the outputs of models using: 1. Altentttt alone (the baseline, without distributed training), 2. Zigzng- Riizag-Ring atlention (from the Context implementation (Zhang 2023 in a distributed setting, 3. Our AnchorContest based on sequence parallelism with DeepSpeed- Ulysses (Jacobs
et al. 2023 in distributed setting.
As shown Table g the Zigzag-Ring allention exhib- ited a maximum logits difference of up to 0.75 compared to the bascline, indicating numerical discrepancies intro- duced by the distributed compulation. In contrast, AnchorContent inplementation showed zero difference in logits when compared to using F'ashAAtentionn without any distributed training, highlighting its mumerical accu-
Table 8: Our distributed computation achieves zero logits difference over 32K sequence length.
Zigzag-Ring (Easy Context)
Our Impl. (AnchorContext)
Full Attn Anchor Attn
0.75
0
0
racy.
Speed and Efficiency. Our method is not only accurate but also significantly faster. By integrating FlashAttention2 and optimizing the attention mechanism within our Contert framework, we achieve higher GPU utilization and faster training times. Figure illustrates the estimated number of days required to process billion tokens various context lengths. It demonstrates that Anchor Allention substantially
14

---

Published in Transactions on Machine Learning Research (02/2025)
reduces training time compared to Full Attention when using the same sequence parallelism and DeepSpeed- Clysses configurations, showcasing the efficiency gains of our approach.
Ease of Integration. Designed with practicality in mind, our muchorAttention seamlessly integrates into other training codebase that using FlashAttention (Dao 2024) and Hugging Face transformers (Wolf et al., 2020). The flexibility of our AnchorContert approach allows for effortless adoption, enabling researchers incorporate it to much substantial modifications to their codebase.
Supporting Advanced Experiments. Moreover, the combination of our method with FlexAttention mechanisms can support more advanced experiments. For example, it facilitates flexible attention masks needed by interleaved chunks, as illustrated in Figure G This flexibility empowers researchers to explore novel ideas and push the boundaries of long-context models.
GPU: 80G A100 8
GPU: 80G A100
GPU: 40G A100
GPU: 40G A100
10.0
Full Attention
Full Acionionn
Full Atteniion ArchorAtterrinn
Full Attantion
ArchorAttsmiioon
And Anchenkttetionnn
Anc MectorAttrtion
7.5
Days
5.0
2.5
0.0
700K
BODK
900K
350K
400K
450K
250K
300K
350K
100K
150K
200K
Context Length
Context Length
Context Length
Context Length
Figure 6: Estimated training time required to process 1 billion tokens at various context lengths using dif- ferent attention mechanisms. Our AnchorAttention reduce more than 50% of time needed by Full Attention.
6 Related Work
Positional Embedding. Positional Embedding is essential for Transformers (Vaswani et al. 2017 to capture the sequential order of imput tokens. Early methods used absolute or learned positional embed--
dings (Vaswani al.] [2017 Devlin 2018). but they often struggled to generalize to longer sequences. Rel-
ative positional embeddings (Shaw al. 2018 Ke et al.
2021 were introduced to handle variable-length
sequences more effectively. Recently, Rotary Positional Embedding (RoPE) QSu et al.) 2021 has become
popular in LLMs (Touvron et. al. 2023
Dubey et al.
2024
Jiang el. al. 2023 Yang el. al.] 2024 Abdin
et al., 2024) for its ability to extend inference context
length with minimal fine-tuning (Liu et al. 2023b;
Chen ct al. 2023a LocaILLaMA] 2023 Peng et al.
2023
Men et al. 2024 Fu et al. 2024). Building on
RoPE, methods like Position Interpolation (Chen el al.]
2023a), NTK Interpolation
2023)
YaRN (Peng et al. 2023), Resonance RoPE
al.
2024b). and CLEX (Chen et al. 2024] have been
proposed to enhance long-context capabilities. In a different direction, we investigate how the precision of BFloat16 could compromise the relative position properties of RoPE.
Extending Language Model Context Lengths. Although training free methods have shown promise in extending context lengths Xiao et al.
(2023) ; Han et al. (2023): Ruoss et al. (2023), recent studies indi- cate that dedicating additional training to handle long contexts yields significantly better results Fu et al.
(2024; Xiong et al. (2023); Gao et al. (2024). The long-context training poses challenges: the quadratic computational complexity and performance degradation beyond the pre-training context length. To address these issues, various approaches have been proposed. Some methods extend context windows by training
2023 Men et al. 2024), adjusting the RoPE base to enable models to handle longer contexts with minimal training. Other strategics manip-
with modified RoPE (Chen et al. 2023a LocalLLaMA 2023 Peng et al.
ulate attention patterns lo efficiently handle extended contexts (Chen et al. 2023b; Xiao et al.) 2023 2024
Bertsch et al.] 2024 Jin et al.) 2024 Ge et al. 2024 Yin et al. 2024), redesigning the attention mechanism to prioritize computational resources on the most relevant parts of the input, thus roducing overhead for longer sequences. Additional efforts reduce the computational complexity of attention through techniques
like sparse attention [Lou et al. 2024 Ge et al.] 2024 group query attention (Ainslie al. 2023), with
15

---

Published in Transactions on Machine Learning Research (02/2025)
optimize implementation Dao 2024). To further overcome hardware limitations, methods like Sequence Parallelism [Li et al. 2021 dist ribute sequences across multiple devices, while Ring Attention (Liu et al. 2023a) and DeepSpeed
(Jacobs et al. 2023) enhance efficiency by optimizing communication strate- gies when procossing segmented sequences. Our work adopts the contimous long-context training approach with efficient implementation, but specifically focuses on atlention design to mitigale issues in RoPE caused by BF precision.
7 Conclusion
In this paper, we identified a critical issue where the combination of Rotary Position Embedding (RoPE) and BF precision breaks the relative positional encoding propertics of RoPE, particularly in long- context scenarios. Our analysis revealed that the first token in the sequence contributes significantly this breakdown, and the problem worsens the training window size increases due to cumulative numerical errors. To address this, proposed AnchorAttention token as a shared anchor across all documents within the context window. By assigning the same position ID to this anchor token and ensuring it is visible to all documents while keeping tokens from different. documents invisible to each other, Attention maintains consistency in positional encoding. This approach not only preserves the integrity of RoPE's relative positional properties but also reduces the number of tokens involved in attention computations, mitigating the accumulation of numerical errors. Our experiments demonstrated that Anchor consistently outperforms full attention and standard intra-document attention methods on long-context benchmarks like RULER, across context lengths ranging from 8K to 128K tokens. Additionally, real-world long-context benchmarks such as LongBench, Anchor Attention improved in-context learning performance while largely preserving the model's capabilities on general tasks like MMLU and HellaSwag. Importantly, AnchorAttention requires minimal modifications to existing training pipelines and reduces training time by more than 50% compared to standard allention training.
novel attention mechanism that treats the first
8 Future Work
In our empirical analysis, we measured the average attention difference over multiple sequences. We also examined the per-sample attention differences, with visualizations provided in the Appendix B The differ- ences were consistent across samples, and we observed that even a simple function could approximate the difference. Notably, the first token contributes the most this discrepancy, leading us to hypothesize that the position IDD of the first token acts as absolute position. Future work is needed to more rigorously investigate the properties of the first position and its impact on positional encoding. Additionally, given the observed significance of the first token in an input sequence, we speculate potential relationship between this phenomenón and the well-documentedd atiention ion sinks and massive activation {Xiao et al. 2023 Han
et al. 2023 Gu et al. 2024 Guo et al. 2024a Sun et al.] 2024 Guo et al. 2024b). Further research is re-
quired to explore this connection and to understand how it might affect attention mechanisms long-context language models.
Broader Impact Statement
The development of large language models (LLMs) to effectively handle long-context sequences. This advancement has the potential to significantly improve applications that require understanding and processing extensive textual data, such as financial document analysis, long-form content generation, and multi-document summarization. By enhancing the ability of LLMs to maintain coherence and context over long sequences, we contribute to the creation of more capable and contextually aware Al systems.
addresses a critical challenge in natural language processing by enabling
Limitations
Although we try ablate the major components of our proposed attention mechanism, due to resource limitations, we cannot exhaust all aspects, such study its effectiveness for pretraining, optimization hy-
16

---

Published in Transactions on Machine Learning Research (02/2025)
perparameters and additional data mixtures. We also limit ourselves to the 10B-scale model size regime with 2B tokens, which may limit the generalizability of our findings.
References
Marah Abdin, Jyoti Aneja, Hany Awadalla, and al. Phi-3 technical report: A highly capable language model locally on your phone, 2024. URL https:
hogg/bbs/2455 14219
Rishabh Agarwal, Avi Singh, Lei M Zhang, Bernd Bohnet, Stephanie Chan, Ankesh Anand, Zaheer Ab- bas, Azade Nova, John D Co-Reyes, Eric Chu, al. Many-shot in-context learning. ar preprint WiXiv:2404.11016 2024.
Joshua Ainslic, James Leo- heo-Thorp, Michiel de Jong, Yury Zemlyanskiy, Fedorico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXie preprinl arXiv:2305. 13245, 2023.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws. arXiv preprint Xiv:244 05405, 2024.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Ao- han Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arxiv preprint ar Xiv:22308. 14508, 2023.
Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew Gormley. Unlimiformer: Long-range transformers with unlimited length input. Advances in Newral Information Processing Systems, 36, 2024.
Guanzheng Chen, Xin Li, Zaiqiao Meng, Shangsong Liang, and Lidong Bing. Clex: Continuous length extrapolation for large language models, 2024. URL https://arxiv: org/ abs/abss///// 2310. 16450
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXim:2366.158: 15595, 2023a.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. ArXiu, abs/2309.1222 2023b.
Tri Dao. Fltstotte::ion-2:: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint xrx:::1810.04800 2018.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, and et al. The Iaam 3 herd of models. ArXiv, abs/2407.211 21783, 2024.
Lizhe Fang, Yifci Wang, Zhaoyang Liu, Chenheng Zhang, Stefanic Jegelka, Jinyang Gao, Bolin Ding, and Yisen Wang. What is wrong with perplexity for long-context language modeling?, 2024. URL https:
//arxiv.org/ ab////as//2100 23771
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hanna Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context. ArXiv, abs/ abs/2402.10011, 2024.
Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How to train long-context language models (effectively). arXiv preprint arXiv:2410. 02660, 2024.
Suyu Ge, Xihui Lin, Yunan Zhang, Jiawei Han, and Hao Peng. A little goes long way: Efficient long context training and inference with partial contexts. arXiv preprint arXiv:2410. 01485, 2024.
Xiangming Gu, Tianyu Pang, Chao Du, Qian Liu, Fengzhuo Zhang, Cunxiao Du, Ye Wang, and Min Lin. When attention sink emerges in language models: An empirical view. arXiv preprint arXiv:2/10. 10781, 2024.
17

---

Published in Transactions on Machine Learning Research (02/2025)
Tianyu Guo, Druv Pai, Yu Bai, Jiantao Jiao, Michael I Jordan, and Song Mei. Active-dormant attention heads: Mechanistically demystifying extreme-token phenomena in llss. aiv preprint or orXim:24 10. 13835, 2024a.
Zhiyu Guo, Hidetaka Kamigaito, and Taro Watanabe. Attention score is not all you need for token impor- lance indicator in kv cache reduction: Value also matters, 2024b. URL /arxiv wwterii.orggabs//// abgg/
12335
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint wwXiv:2308.16199 2023.
Dan Hendrycks, Collin Burns, Steven Basart, Andrew Critch, Jerry Li, Dawn Song, and Jacob Steinhardt. Aligning with shared human values. Proceedings of the International Conference on Learning Repre- sentations (ICLR), 2021.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint XrXee:244.00665, 2024.
Yutong Hu, Quzhe Huang, Mingxu Tao, Chen Zhang, and Yansong Feng. Can perplexity reflect large language model's ability in long text understanding? arXiv preprinl arXiv:2405.06105 2024.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajb- handari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models, 2023. URL httt:s://aaxxioor////////2..... 14509
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. Xiv preprint Xiv: 238,0 2023.
Carlos Jimenez, John Yang, Alexander Weltig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik R Narasimhan. SWE-bench: Can language models resolve real-world github issues? In The Twelfth In- ternational Conference on Learning Representations, 2024. URL https://openrevien.nncm wit.errrrieeeeeetttoounmiiii VIFSyNQM66
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia- Yuan Chang, Huiyuan Chen, and Xia Hu. Lim maybe longlm: Self-extend Ilm context window without tuning. arXiv preprint 2rXiv:2401.01322 2024.
Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharma Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang, Hector Yuen, et al. A study of bfloat16 for deep learning training. arXiv preprint arxiv:1905. 12322, 2019.
Marzena Karpinska, Katherine Thai, Kyle Lo, Tanya Goyal, and Mohit lyyer. One thousand and one pairs: A' novel" challenge for long-context language models. arXiv preprint orXiv:2406.1626, 2024.
Guolin Ke, Di He, and Tie- Yan Liu. Rethinking positional encoding Conference on Learning Representations, 2021. http:pppppppoooooeeeeee/m//m/mmp
encoddin in language pre-training, In International
Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective. arXiv preprint arXiv:110511110, 2021.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context, 2023a. URL https://arxiv..
httpaaaaww..oorraa/aaa/aaaaaaaea wwtp/////g/..oo...//////// 0188
Xiaoran Liu, Hang Yan, Shuo Zhang, Chenxin An, Xipeng Qiu, and Dahua Lin. Scaling laws of rope- rose-basod extrapolation. ar preprint 2rXim:2310.05200 2023b.
18

---

Published in Transactions on Machine Learning Research (02/2025)
LocalLLaMA.
Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuming and minimal perplexity degration, 2023. URL hittps://e httpss//wwr.redddt.coom////
Commaenna//LLl2/55//ttaaaarrrrr waapaaaree_aaaedd_ooeeeeeeee allows aeaaalamaalaaaaaaa mww.eooellll.llm to nave/
Chao Lou, Zixia Jia, Zilong Zheng, and Kewei Tu. Sparser is faster and less is more: Efficient sparse attention for long-range transformers. orXie preprinl arXiv:2406. 16747, 2024.
Sahisnu Mazumder and Bing Liu. Lifelong and continual learning dialogue systems. ar preprint arXiv:2211.0688 16.65533 2022.
Xin Men, Mingyu Xu, Bingning Wang, Qingyu Zhang, Hongyu Lin, Xianpei Han, and Weipeng Chen. Base of rope bounds context length, 2024. URL
wwtp:rgggbba.ssoa/a////// 14591
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXim:2309.0000 00000 2023.
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Benmani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint Xrx:22305. 16843, 2023.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. Xixiv preprint arXiv: 1001.180000..... 2018.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R. Steeves, Joel Hestness, and Nolan Dey. SlimPa- jama: A 627B token cleaned and deduplicated version of RedPajama. http::s///eeebbas.ca.
sl htt.aa/mm..aaaeeooooooeeeeeeeaaa
ofeeiiojjaaa-eeaaaaaaaaaaaaaaa 2023. URL https:
am2677
Jianlin Su, Yu Lu, Shongfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. ArXiv, abs/2100.09864, 2021.
Mingjie Sun, Xinlei Chen, J. Zico Kolter, and Zhuang Liu. Massive activations in large language models, 2024. URL https::/arxxw.org/////a
17762
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini rescarch and technology. arXiv preprint or arXiv:208.0082 08209, 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, and et al. Llana 2: Open foundation and fine-tuned chat models, 2023.
Ashish Vaswani, Noam Shazoer, Niki Parmar, Jakob Uszkorcit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Atlention is all you need. Advances in neural information processing systems, 30, 2017.
Minzheng Wang, Longze Chen, Cheng Fu, Shengyi Liao, Xinghua Zhang, Bingli Wu, Haiyang Yu, Nan Xu, Lei Zhang, Run Luo, al. Leave no document behind: Beuchmarking long-context llms with extended multi-doc qa. arXiv preprint arXiv:2406. 17419, 2024a.
Shibo Wang and Pankaj Kanwar.
Bfloat16:
The secret to high performance
cloud tpus.
https:///clou.googg.. http:///ow...olloeeoooooo//////a http:cpooouucaeeoeo wwtp:aaiiaaieeeeaaaeaeaaaeeaeeee learning/ htwpapgh---p.raooomaomaoooooaaaa
bf ort.a//eeeeeccc.oooooooooooooono August 23 2019. Accessed: August 31, 2024.
Suyuchen Wang, Ivan Kobyzev, Peng Lu, Mehdi Rezagholizadeh, and Bang Liu. Resonance rope: Improving context length generalization of large language models. ar preprint arXiv:2403.00071, 2024b.
19

---

Published in Transactions on Machine Learning Research (02/2025)
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pier- ric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: Statc-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empiricul Methods in Natural Language Processing: System Demonstrations, pp. 38-45, Online, October 2020. Association for Computational Linguistics. URL https: / /www//wwwooooo aclweb. hntg/anthologg/2000 wttp:///220020mmmmmmmmmmeeoo emnlp-demos.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Infllm: Unveiling the intrinsic capacity of Ilms for understanding extremely long sequences with training-free memory. arXiv preprint 0r2.v:2402.04000
2024.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint Xiv:20009. 1 17458, 2023.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajiwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective long-context scaling of foundation models. In North American Chapler of the Association for Computational Linguistics, 2023.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, and et al. Qwen2 technical report, 2024. URL org. / 200024411000
Howard Yen, Tianyu Gao, Minmin Hou, Ke Ding, Daniel Fleischer, Peter Izasak, Moshe Wasserblat, and Danqi Chen. Helmet: How to evaluate long-context language models effectively and thoroughly. arXiv preprint (rXit:2410.0261) 2024.
Qingyu Yin, Xuzheng He, Xiang Zhuang, Yu Zhao, Jianhua Yao, Xiaoyu Shen, and Qiang Zhang. Stablemask: Refining causal masking in decoder-only transformer. arXiv preprint arXiv:2402. 04779, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint Xrx: 07830, 2019.
J. Zhang, Easycontext:
Efficient and extensible long context language models. https://github. com/ 2023. Accessed: 2024-09-16.
Jiajie Zhang, Yushi Bai, Xin Lv, Wanjun Gu, Danqing Liu, Minhao Zou, Shulin Cao, Lei Hou, Yuxiao Dong, Ling Feng, and Juanzi Li. Longeite: Enabling Ilms to generate fine-grained citations in long-context qa. arXiv preprint 240: 028977 2024a.
Xiao Zhang, Miao Li, and Ji Wu. Conditional 2rXiv:2406.01974 2024b.
language learning with context. arXiv preprinl
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et al. inftybench: Extending long context evaluation beyond 100k tokens. arXiv preprint Xiv: 13718, 2024c.
Liang Zhao, Tianwen Wei, Liang Zeng, Cheng Cheng, Liu Yang, Peng Cheng, Lijie Wang, Chenxia Li, Xuejic Wu, Bo Zhu, et al. Longskywork: A training recipe for efficiently extending context length in large language models. arXiv preprinl ar Xim:24400 000:2400000605, 2024a.
Yu Zhao, Yuanbin Qu, Konrad Staniszewski, Szymon Tworkowski, Wei Liu, Piotr Milos, Yuxiang Wu, and Pasquale Minervini. Analysing the impact of sequence composition on language model pre-training. anXiv preprint ar Xio:24002. 18991, 2024b.
20

---

Published in Transactions on Machine Learning Research (02/2025)
A Rotary Positional Embedding (RoPE) and Relative Positional Encoding
Rotary Positional Embedding (RoPE) is a method employed to infuse positional information into transformer models by modifying the query (g) and key (k) vectors within the attention mechanism. RoPE enhances the model's ability to caplure positional relationships through the following steps:
A.1 Mechanism of RoPE
Splitting into 2-Dimensional Chunks:
The query and key vectors are partitioned into 2- dimensional segments. This segmentation allows each pair of dimensions to undergo independent rotational transformations.
Applying Rotational Transformations: Each 2-dimensional chunk rotated by an angle deter- mined by a frequency parameter 0. Specifically, for the i-th chunk, a rotation matrix RA is applied. This introduces positional information based on the token's position in the sequence.
The rotation matrix Ri for the i-th chunk is defined as block diagonal matrix composed of individual 2-dimensional rotation matrices, where the d denotes the number of hidden dimension of each head, such as 128 for LLaMA-2-7B:
Feoa(tfi))
(50 cont
(x0g))
cont xil stn
wan Li con 220
(6)
CONE 2/2-1?
a/2-1)
8/2-
A.2 Incorporating RoPE into Attention Mechanism
In the attention mechanism, logits are computed based on the similarity between queries and keys. With RoPE, the rotated queries and keys are utilized in this computation as follows:
A(.) (Ri. (:) (Rg,0 q. Ryokk
A.3 Encoding Relative Positional Information
To explain how RoPE encodes relative positional information, consider the product of the rotation matrices R.g.R Using trigonometric identities for sine and cosine of angle differences:
sin(a 3) = sin a cos 3 - cos a sin 8
cos(a: 3) cos cos 3 + sin sin /
Focusing on single 2-dimensional block, the product is computed
cos(iio) sin(i60)
(cos(joo) -sin(joo)] cos( (i0)) sin
sin( (i00)
R.
cos(j0o)
cos ((i - j)0o)
sin ((i 10-jj)b j)oo
sin ((i j)00)
cos ((i j)0o)
R
Here, m j represents the relative position between the query and key. This derivation demonstrates that the product R,o Rg depends solely on the relative position (i (d-)), rather than the absolute positions of i and individually.
21

---

Published in Transactions on Machine Learning Research (02/2025)
Attention Difference per Sample
Attention Difference per Shift
35.0
Fretrained (BF16), Sample
Pretrained (BF16), Sample
Protrained (BF16),
Pretrained Sample
Attention Score Difference
Attention Score Difference
27.5
20
15
2
20.0
A,-0, Az=0
Ag-0, Ag
17.5
A,-0, 42
A--0, Az
Ag=0,
10:2241517122222
300
1000
0 1012141018202
50 100
200
500
1000
2000
Shift Index
Figure 7: Detailed visualization of attention score differences under BF for individual samples.
B Detailed Analysis of Attention Discrepancies with RoPE under BFloat16
In Section we primarily measured the attention differences averaged over multiple sequences. To provide a more detailed perspective, we present visualizations for individual samples. Figure (left) shows the attention score differences for five samples, offering a detailed view compared to the averaged version represented by the blue line Figure] (left).
As illustrated in Figure 7 (left), the attention score differences are highly consistent across different samples. The curves nearly identical, exhibiting only slight biases for each sample. Given the starting differences (A) 0 and A2 = 16), we can accurately predict the attention differences for various A values.
Furthermore, while Section 2 focuses on fixing A 16 and varying A to study attention differences, we extend our analysis by fixing A = 0 and varying A2. The results are summarized in Figure 7 (right). We observe that the attention difference curve shifts along the x-axis as A₂ changes. These findings demonstrate that the attention score differences are highly predictable.
previous observation that the first token's positional encoding deviates significantly, became curious about its role in establishing positional references for subsequent tokens. Empirical findings led us to hypothesize that the position ID of the first token functions as an absolute position. A more detailed discussion of this hypothesis will be explored in future work.
Building
C Data Statistics
We present the mixture ratios and token contributions from each domain in dataset. The Mizture Ratio is calculated by dividing the number of sequences from each domain by the total number of sequences. The Token Ratio is determined by dividing the token count of each domain by the total token count of the dataset.
For reference, the original SlimPajama token mixture ratios are as follows: the dataset consists of 82% web data (67% from CommonCrawi and 15% from C4), 4.5% code (GitHub), 4.5% Wikipedia, 4.5% books, 2.5% Ar Xie, and 2.0% StackExchange Since this dataset closely mirrors that used to pretrain the L.LMMA models.
22

---

Published in Transactions on Machine Learning Research (02/2025)
Table 9: Domain and Token Distributions
C4
Arxiv Github
StackExchange
Common
Wikipedia
Books
Up-sampled Data Minture
61K 128K 128K]
Mixture Ratio
52.34%
1.1%
3.68%
4.56%
33.40%
4.79% 3.51% %
0.21%
Token Ratio
19.53%
5.86%
6.61%
1.64%
58.14%
4.69%
Original SlimPajama
Mixture Ratio
55.32%
0.30%
3.65%
5,06%
31.01%
4.59%
0.06%
Token Ratio
26.50%
4.64%
5.05%
3.18%
53.42%
3.34%
3.88%
Mixture Ratio 55.05%
0.40%
3.66%
4.97%
31.23%
4.58%
0.10%
Token Ratio 25.43%
5.22%
5.05%
2.95%
54.24%
3.24%
3.86%
there is less concern about distribution shift during continual pretraining: therefore, many recent works have utilized it.
Comparing the data we used for training, as shown in Table we observe several key differences:
C4 Datasct: Our up-sampled data mixture has a slightly lower Mixture Ratio for C4 (52.34% vs. 55.32% in SlimPajama) but a more pronounced decrease in Token Ratio (19.53% vs. 26.50%). This indicates that while the number of sequences from C4 is comparable, they are shorter on average in our datasct.
CommonCrawl: We increased the Mixture Ratio for CommonCrawl (33.40% vs. 31.01%) and observed higher Token Ratio (58.14% vs. 53.42%). This suggests that CommonCrawl sequences in our dataset are only more numerous but also longer, contributing significantly to the total loken count.
The representation of ArXiv and Books enhanced our up-sampled mixture. ArXiv's Mixture Ratio increased from 0.30% 1.01%, and its Token Ratio from 4.64% to 5.86%. Similarly, Books saw an increase in Mixture Ratio from 0.06% 0.21% and in Token Ratio from 3.88% to 4.69%. These increases aim to enrich the dataset with more scholarly and literary content.
ArXiv and Books:
StackExchange and Wikipedia: The proportions of Etacanxc and Wikipedia remain rela- tively consistent between the two datasets, cnsuring stable reproscntation of community-driver and encyclopedic knowledge.
The Mixture Ratios for GitHub similar (3.68% in our dataset 3.65% in SlimPajama), but the Token Ratio is slightly higher in our dataset (6.61% vs. 5.05%), indicating longer code sequences that could benefit code understanding tasks.
GitHub (Code Data):
Overall, up-sampled data. mixture places greater emphasis longer sequences from CommonCrawi and increases the diversity of content by up-sampling underrepresented domains like ArXiv and Books. This rebalancing is designed to enhance the model's ability to generalize across various content types and improve its performance on tasks requiring knowledge from specific domains.
2000
2000
200r
1990
Figure 8: Training Data Sequence Length Distribution
23

---

Published in Transactions on Machine Learning Research (02/2025)
NIAH Multikey 1
NIAH Multikey 2
NIAH Multivalue
NIAH Multiquery
NIAH Multiey
NIAH Single
NIAH Single
NIAH Single
CWE
FWE
Q4 7
QA 2
VT
Model
Llama2
43.0
Chat
95.2
100.0
99.8
95.2
90.0
70.2
95.8
88.7
58.4
34,26
85.91
64.8
39.4
+
Yarn 64K
73.0
24.4
8.0
18.0
5.8
0.8
5.9
6.35
54.22
18.16
67.8
38
27.6
Yarn
34.75
82.07
25.0
Table 10: Results different models across various tasks on 4, 000 context length.
D LLaMA-2-7B RULER Performance
In Table 10 we use the base template for all models. Compared to the results reported in (Hsieh et al. 2024
observe that Llama 2 7B's performance the Common Word Extraction (CWE) task is unstable. When the context length is set to 4,000 (less than 4,096), slightly different context example is provided (refer to the generate_ mewoorae_input_oon input_output function in RULER 3 As shown in Table 11 when the context length is set to 4, 096, we can reproduce the reported results. However, the inclusion of a slightly different example disrupts the Llama 7B model's ability to accurately count words, highlighting the instability of the CWE task for evaluating Llama 7B. Consequently, we omit this task from the long-context evaluation.
compo//ccmonncc/coooco/o hommmnn_woddss.s wwt.eeeetettttteeeeeeeeeeeeeeeee py in
Furthermore, based on hypothesis that continued training of large language models primarily improves their ability to handle longer contexts rather than introducing new capabilities, we have excluded the NIAH Multikey 3 and CWE tasks from our evaluation. This decision is supported by the observation that even within the pretraining context length, Llama 2 7B is unable to effectively solve these tasks.
4,000
4,096
LLaMA-2-78
24.46
76.8
Table 11: Performance of LLaMA-2-7H 2-7B on Common Word Extraction (CWE) with different context lengths.
E
Longbench full result
24

---

Published in Transactions on Machine Learning Research (02/2025)
Code Completion TCL Multi-Doc QA Single-Doc QΛ Summarization Synthetic
SlimPajaman 64K
Full Attention
60.52
62.51
9.68
17.34
16.09
2.87
Cross-Dos Attention
62.95
62.79
9.51
16.82
16.73
2.94
62.76
63.76
9.30
16.40
14.61
3.74
AnthorAttention
62.04
65.38
18.60
17.56
4.24
tag
63.53
66.02
9.51
18.28
15.30
5.24
SlimPajama- 128K
Full Attention
54.17
50.72
6.36
16.43
13.30
2.04
Cross- Doc Attention
54.59
51.22
6.42
15.59
13.92
3.63
reset
52.51
50.07
6.30
16.64
14.45
4.18
AnchorAttention
54.14
51.85
6.32
17.74
12.67
3.89
lng
55.81
51.89
5.93
17.67
12.43
3.41
UysSampleddii-122K
Full Attention
53.13
48.96
6.12
14.66
12.77
4.13
Cross-Doc Attention
54.16
49.51
5.72
14.62
14.38
2.57
reset
54.29
50.18
5.37
14.30
15.23
2.55
AnchorAtt
53.90
50.17
6.30
18.29
13.78
6.13
tog
55.13
49.70
5.65
16.90
15.53
4.20
Table 12: Performance Metrics across Different Attention Mechanisms and Datasets.
25


Reviews:
[{"review_id": "Dm6rp94YfH", "reviewer": "Reviewer_CR2M", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n1.The analysis of the impact of bf16 on RoPE at various positions is both novel and meaningful, yielding intriguing results.\n\n2.The paper provides extensive experiments, including both long and short benchmarks, and the conclusions align well with the observed phenomenon that the difference in attention increases with context length.\n\n\n# Weaknesses\n\nMy main concern is about the motivation of AnchorAttention. The authors assert that AnchorAttention enhances the model by avoiding the breakdown of RoPE in bfloat16, but the experiments provided are insufficient to substantiate this claim. I am quite intrigued by Section 2 and look forward to seeing more experiments that validate the impact of Attention differences. However, I became quite confused while reading Sections 3, 4, and 5, as there is no mention of bfloat16. Although I am convinced that AnchorAttention is indeed beneficial for long-context modeling, there are several problems: First, the Attention Difference does not necessarily correlate with differences in final performance. Secondly, it is doubtful whether the functionality of AnchorAttention is related to precision or any other factors (might be the impact of the first token itself?). The argument would be significantly more compelling with additional experiments demonstrating that the superiority of AnchorAttention is evident only in bf16 training and inference. Also, if the analyses of Attention Difference are accurate, then it should be feasible for AnchorAttention to also benefit full attention, too.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper starts with an interesting prospective of the impact of float16 to rotary position embedding. The thesis introduces AnchorAttention, a method that effectively addresses the issue where the first token in RoPE plays a dominant role in attention differences when using BF16. It includes a wealth of experiments that confirm the superiority of AnchorAttention in long-context modeling.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.\tProvide experiments showing that the superiority only arises in bfloat16 training/inference (e.g., by comparing fp32 training & inference). Also it might help to show that AnchorAttention works for both intra-document and full attention.\n2.\tI am curious why the first token has the largest Attention Difference? This looks not straight forward for me，It would be beneficial to include some fundamental theoretical analyses.Also, in Figure 1, there is a abrupt pitfall of Attention Difference at shift index of 15, why this happens? I would hope to see some theoretical analysis on this issue.", "broader_impact_concerns": "No specific concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ok3UtuBKaB", "reviewer": "Reviewer_y1GF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\n- BFloat16 breaking the relative property of RoPE is an important observation that is highly relevant to other researchers.\n\n- AnchorAttention is simple to implement and shows improvement over full attention and intra-document attention on the RULER benchmark across multiple LLM architectures. \n\n- The codebase provided is easily adaptable by the research community. \n\n\nWeakness:\n\n- I would have liked more intuition behind the observations extracted from the experiments in Section 2. Why is the discrepancy so large at the first token? Why is there a small uptick in the difference in attention at the end of the sequence (figure 1 middle)?\n\n- Section 3.2 investigates the effect of the RoPE + BFloat16 discrepancy on long-context tasks, but the authors did not show the performance on the same benchmark using float32, which I think would be important to understand the effect of BFloat16. Is this drop in performance as the context length increases the result of BFloat16 or something else?\n\n- The connection between AnchorAttention and the limitation of BFloat16 can be further clarified. Is AnchorAttention’s motivation to reduce the impact of the discrepancy introduced by BFloat16 instead of addressing the discrepancy itself?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "RoPE has become the standard positional encoding method for training LLMs with long context windows as it allows for out-of-distribution generalization of relative positions. However, the paper shows that when RoPE is implemented in combination with brain floating point (bfloat16), commonly used to reduce memory requirement by reducing precision, the relative property of RoPE is compromised especially for long-context scenarios. \n\nThe paper further shows empirically that the discrepancy in attention calculation is mainly due to the first token in a sequence. To mitigate the effect of the discrepancy, the paper introduces AnchorAttention, which builds on intra-doc attention by including a special Anchor Token shared by all documents. AnchorAttention shows significant improvement on the RULER benchmark over full attention and intra-doc attention.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Summary of weaknesses to address:\n\n- Since AnchorAttention is motivated by the observation that the first token of a sequence contributes the most to the deviation of RoPE’s relative properties, can you show more experiments that support this claim? Figure 1 (middle) only shows this for a shift of 0 and 16, does this same drastic deviation happen with different deltas?\n\n- Can you show the performance of float32 on the same benchmark in section 3.2 and figure 3? If not, can you justify why the current experiments are sufficient?\n\nMinor comments and questions:\n\n- I think the x-axis of Figure 1 (left) could be more intuitive if it shows the difference in the two deltas instead of the value of delta 1. I also think omitting delta 1 = 16 adds to some confusion, and it would be good to show that when the two deltas are equal the difference goes to zero as a sanity check.\n\n- Are all experiment results in Section 4.3 conducted with bfloat16?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "cHGywKMokF", "reviewer": "Reviewer_QoHJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The introduction of AnchorAttention is timely and addresses a critical issue in long-context model training with BFloat16 precision.\n- The paper is easy to follow and well written.\n- Authors provide extensive experimental results to demonstrate the effectiveness of AnchorAttention\n- Authors release a useful package of AnchorAttention based on flashAttnetion, which can be useful for the community.\n\nWeaknesses:\n- (Minor) Though the piece of work provides very strong empirical results. More theoretical understanding can help readers to understand why the proposed method works.\n- Question: What are the trade-offs of resetting the position IDs? Will it influence multi-document understanding (e.g. QA based on multiple documents)?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces AnchorAttention, a novel attention mechanism designed to mitigate the limitations of RoPE when used with BFloat16 precision in long-context language models. The primary contributions are:\n1. The authors identify that RoPE's relative positional encoding properties break down under BFloat16 precision, especially in long-context scenarios, due to cumulative numerical errors.\n2. AnchorAttention introduces resetting position IDs for intra-document attention and a shared anchor token to alleviate numerical issues. This approach reduces the computational cost and improves performance on long context tasks.\n3. Extensive experiments demonstrate that AnchorAttention reduces training time by over 50% while maintaining or improving performance on benchmarks like RULER and LongBench.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the weaknesses part.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a valuable contribution to the field, and the requested revisions will further strengthen its impact. The identified weaknesses are addressable and do not detract from the overall significance of the work. The suggested certifications will further enhance the paper's credibility and contribution to the research community.

The following points can be taken into consideration for the minor revision:

1. While the empirical results are strong, the paper lacks a deep theoretical understanding of why the RoPE breakdown occurs under BFloat16 and how exactly AnchorAttention mitigates this issue. Reviewers CR2M and QoHJ specifically request more theoretical analysis.

2. Reviewer CR2M raised concerns about the motivation of AnchorAttention, suggesting a need for clearer articulation of how it relates to the BFloat16 precision issue. The authors' response clarifies this, but the revised manuscript should reflect this more explicitly.

3. Reviewer CR2M's question about directly starting from position 2 in intra-document attention warrants a brief discussion in the paper, even if it's just to explain why it's suboptimal.

4.  Incorporate other minor clarifications and suggestions from the reviewers as needed.
Decision: Accept (poster)