1 Introduction↩︎

Developing language models capable of handling long context lengths poses several challenges. First, as the context length increases, an effective modeling of extended input sequences becomes increasingly critical. This often requires advancements in positional encoding [1], extrapolation techniques [5], or architectural innovations [6], [7]. Second, training long context large language models with billions of parameters demands significant computational resources. Overcoming this challenge requires scalable algorithms, high-quality datasets, and robust infrastructure. Lastly, deploying these models in real-world applications demands low latency and low memory usage, which requires meticulous optimization of both the model architecture and the serving infrastructure.

On the modeling front, two components of the transformer architecture are particularly crucial for long context capabilities: the attention mechanism and positional embeddings. Recent research has proposed various methods to enhance these components. For instance, Landmark Attention [8] trains attention modules to select relevant blocks using a representative token, referred to as a “landmark token", for efficient retrieval within extended text corpora. Similarly, Focused Transformer [6] adopts a contrastive training approach to prioritize attending the most relevant portions of the input sequence, allowing the model to focus on smaller, contextually significant subsets of tokens. Although these approaches improve long context modeling, stabilizing training remains a key challenge for extending transformer capabilities to longer sequences. Query-Key Normalization (QK-Norm) [9], [10] has been introduced to address the stability issue, which normalizes the query-key vectors along the head dimension before computing attention. Although QK-Norm mitigates numerical instability during training and is widely used [11]–[13], it may impair long context capabilities.

In addition to the chosen attention mechanism, positional embeddings play a crucial role in long context modeling. Various approaches have been proposed to improve their effectiveness. Popular methods include Absolute Position Embedding (APE) [14], Relative Position Embedding [15], ALiBi [16], and Rotary Position Embedding (RoPE) [1]. Among these, RoPE has gained significant adoption in large language models (LLMs) [17]–[19] due to its simplicity and effectiveness. In particular, it has the ability to extrapolate context lengths by adjusting RoPE \(\theta\) values during training [20]–[22]. Other techniques, such as relative bias [15], [16], [23] and contextualized position embeddings [24], introduce distance-based bias terms or condition the position information on input semantics. These methods often affect attention distribution by incorporating auxiliary information, such a positional indices or explicit recency bias. However, whether certain information or biases are beneficial to long context modeling or overall performance remains less explored. Additionally, the concept of No Positional Embedding (NoPE) has been explored by [25], suggesting that removing explicit positional embeddings and relying solely on implicit positional information derived from the causal mask can enhance long context performance.

Despite advances in long context modeling, training and serving such models remain challenging due to the quadratic complexity of standard attention. Techniques like Sliding Window Attention (SWA) [26], [27] mitigate this by restricting each token’s attention to a local window, reducing compute while maintaining local coherence. Sparse attention methods [28]–[30] extend this by introducing structured sparsity, including random [31] and dilated patterns [32]. More recent approaches further compress attention, such as activation beacon [33], which sequentially summarizes local keys and values, and attention sink [34], which stabilizes long-sequence training by preserving early tokens with a sliding window. On the serving side, KV cache trimming methods [35]–[37] selectively discard cached states based on heuristics to reduce memory usage and boost inference efficiency. However, these gains often come with trade-offs in model quality, emphasizing the need for careful design.

In this paper, we begin analyzing attention patterns of different attention mechanisms, RoPE, NoPE, and QK-Norm and its impacts on long context performance trained up to 750 billion tokens. Building on these insights, we propose a novel hybrid attention architecture and extensively pretrain up to 5 trillion tokens, followed by supervised fine-tuning on a diverse set of datasets tailored for long context. We show that this architecture surpasses existing state-of-the-art extrapolation-based RoPE models [3] by a large margin, striking a balance between efficiency and performance.

2 Observation↩︎

In this section, we assess three models with different attention components on needles-in-a-haystack [38] (NIAH) and analyze the attention patterns to understand how these variants affect performance. Analysis in this section guides our architectural design choices throughout this work.

2.1 Experimental Setup↩︎

All model variants share a common configuration consisting of 8 billion parameters (including the token embedding parameters), with detailed architectural specifications provided in Table 1. For this set of experiments, the model is trained in two stages: a pretraining stage followed by a supervised fine-tuning (SFT) stage. Previous research shows that the SFT stage is necessary for long context evaluations, as it can reduce variance in long context tasks and enables the emergence of long context capabilities that may not manifest in models trained solely through pretraining [39].

We pretrain the model with a batch size of 4 million tokens. We use AdamW with a peak learning rate of \(7e^{-3}\), a linear warmup of 2000 steps and a cosine learning rate schedule decaying to \(3.5e^{-4}\) over 179,000 steps for a total of 750 billion tokens. For the SFT stage, we adopt an interleaved training strategy: we combine short- and long context data in a 3:1 ratio, with context lengths of 8192 and 65536 tokens, respectively. We use a batch size of 0.5 million tokens.

The 3 model variants we test are:

1. RoPE Model: For this variant, we employ Rotary Position Embedding (RoPE) to encode positional information. During the pretraining stage, the RoPE parameter \(\theta\) is set to 10,000. In the subsequent supervised fine-tuning (SFT) stage, \(\theta\) is increased to 2 million to account for the increased context length. This variant serves as the baseline model configuration, maintaining an architecture similar to that of most existing models [17], [26], [40].

2. QK-Norm Model: Layer normalization [41] is applied to both the query and the key vectors before performing the angular rotation used in RoPE. All other hyperparameters, including the \(\theta\) values and training methodology, remain identical to those of the RoPE variant.

3. NoPE Model: Previous research [42], [43] has demonstrated that transformer variants trained without positional embeddings (NoPE) can perform effectively on long context tasks. However, these models often exhibit inferior performance in terms of perplexity and downstream task evaluations within the trained sequence length [43]. In our study, the NoPE variant does not have QK-Norm. This variant is trained using the same methodology as the other two variants.

2.2 Evaluation and Attention Analysis↩︎

2.2.1 Evaluations↩︎

We evaluate the variants on a set of core evaluation benchmarks, including MMLU [44], HellaSwag [45], CommonsenseQA [46], ARC [47] for core model capabilities and NIAH benchmark [38] for long context capability. NIAH evaluates a model’s ability to retrieve information accurately from a specific sentence (the “needle”) embedded within a lengthy document (the “haystack”). The needle is randomly placed at varying depths within the sequence to examine performance across different context lengths. To improve robustness, we modify the original NIAH benchmark, where each context-depth combination is tested 16 times with different random seeds, creating diverse context compositions for comprehensive evaluation. The results of all standard benchmark evaluations and results with 65536 context length needles are presented in Table 2. Although prior research has emphasized the limitations of NIAH [48] for evaluating deeper and more general context understanding, our focus is solely on testing basic long context capabilities and gaining insights on model architecture design, for which this benchmark is sufficient.

Table 2 shows that the RoPE and QK-Norm variants exhibit comparable performance on standard benchmarks, while the NoPE variant lags behind. This finding is consistent with previous studies [25], [42]. For long context evaluations, QK-Norm performs the worst among the three variants, despite its decent performance in other capabilities. This observation is consistent with the results from the comparisons between Command R and Command R+, where Command R, despite being a significantly smaller model, outperforms Command R+ overall on longer context benchmarks [49]. Although the NoPE variant has slightly lower needles score compared to the RoPE variant, it is decent given that its base capabilities is relatively low.

2.2.2 Attention Pattern Analysis↩︎

To better understand the impacts of different architectures, we also analyze the attention patterns within each model. This approach is inspired by previous studies [50], [51] where attention scores assigned to different parts of the context are closely examined.

We still utilize the NIAH task by first dividing the context into four segments:

• Begin: The first 10 tokens. This part of the context is also often referred to as the “attention sink” [34], where a disproportionately large amount of attention is typically allocated in a trained transformer model.

• Needle: The tokens of the inserted needle sentence. Ideally, a well-trained model should assign a relatively high amount of attention to this part of the context.

• End: The query and completion tokens, which can represent recency bias.

• Context: The remainder of the context, typically consisting of noise or irrelevant content.

We position the needles at approximately 50% depth within the context to increase the complexity of the task, as most models suffer from the lost-in-the-middle problem, as highlighted in previous works [52], [53]. For each model, we first calculate the attention scores between the query tokens of “End” segment and the key tokens of all four segments across all heads and layers. The attention scores are summed within each segment and then aggregated across all heads and layers to obtain the average attention weight for each segment. These scores are further averaged across multiple samples at sequence lengths of 8,000 tokens, 32,000 tokens, and 128,000 tokens. We refer to this metric as “attention mass” in the following sections. The results of each variant are visualized in Figure 1.

We begin by comparing attention distributions across different sequence lengths from Figure 1. We observe a consistent decrease in attention mass allocated to the “Needle” segment and a corresponding increase in attention mass on the “Context” segment as sequence length increases. This trend indicates that retrieving relevant information becomes increasingly difficult as the context grows longer. When comparing across model variants, the NoPE variant consistently allocates the highest attention mass to the Needle” segment, followed by the RoPE variant, while the QK-Norm variant assigns the least attention to this segment. Furthermore, the QK-Norm variant exhibits markedly lower attention mass on the “Begin” segment and substantially higher attention mass on the “Context” segment relative to other variants. This indicates that models trained with the QK-Norm component exhibit a weak attention sink and are more susceptible to interference from noisy content. These patterns are consistent with QK-Norm’s poor performance on the NIAH task. We argue that QK-Norm has this effect because the normalization operation mitigates magnitude information from the dot product of Query and Key vectors which tends to result in attention logits being closer in terms of magnitude and flatter in terms of distribution. A more detailed analysis of why QK-Norm is detrimental to long-context modeling is provided in Appendix 8.

2.2.3 Hybrid Model↩︎

Building on the findings above, we are inspired to combine RoPE and NoPE layer-wise to leverage the strengths of both variants. NoPE provides an effective attention mechanism for retrieving information based on vector similarity, while RoPE explicitly models positional information and introduces a recency bias. By combining them, we aim to enhance overall performance. We propose a new variant that alternates between NoPE and Rotary Position Embedding (RoPE) layers. Specifically, the two position-embedding strategies are interleaved, with NoPE applied in one layer and RoPE in the next. To ensure consistency and enable meaningful comparisons, the RoPE parameter \(\theta\) is initially set to 10,000 during pre-training. The pre-training procedure follows the setup described in Section 2.1 in terms of data and training protocol. Subsequently, we perform multiple fine-tuning runs with varying \(\theta\) values—10,000, 100,000, 2 million, and 4 million—to evaluate performance across different configurations, using the same training steps and data as in Section 2.1. We refer to these models as RNoPE-10k, RNoPE-100k, RNoPE-2M, and RNoPE-4M with the specific RoPE \(\theta\) value used for each during supervised fine-tuning (SFT).

For the RoPE baseline model and all RNoPE variants, we report the needles evaluation score of all SFT models at a sequence length of 128,000 in Table 3. We also display the attention mass of all variants in Figure 2. The attention mass is aggregated separately for all RoPE and NoPE layers. For simplicity, we present results based on samples with 32,000 sequence lengths, with the complete table for all sequence lengths provided in Appendix 7.

The results in Table 3 reveal that increasing the RoPE parameter \(\theta\) during fine-tuning of the RNoPE variants—where NoPE and RoPE layers are interleaved—leads to a decline in the model’s long-context capability. In contrast, previous studies on pure RoPE architectures [20], [54] have shown that using a larger RoPE \(\theta\) during pre-training or fine-tuning enhances long-context performance and expands the effective attention receptive field. To investigate this discrepancy, we compare the attention mass across different model variants from Figure 2.

First, when comparing the NoPE and RoPE layers across all RNoPE variants, we observe a clear divergence in their behavior. The NoPE layers exhibit strong retrieval capabilities, characterized by a pronounced spike in attention mass on the “Needle” segment and a phenomenon of attention sink [34] on the “Begin” segment. Additionally, these layers show a significantly weaker recency bias indicated by the low attention mass on the “End” segment. In contrast, the RoPE layers within the RNoPE variants demonstrate extremely weak retrieval performance, evidenced by near-zero attention mass on the “Needle” segment and only modest attention on the “Begin” segment—indicating an absence of an attention sink. However, these RoPE layers exhibit a much stronger recency bias compared to the pure RoPE baseline. In summary, interleaving RoPE and NoPE layers induces a spontaneous “division of labor” phenomenon, where RoPE layers focus on local information aggregation and NoPE layers specialize in long-range retrieval. Remarkably, this functional specialization emerged naturally during training, without any explicit training constraints, data augmentation strategies, or specific loss function designs.

Next, we examine the RNoPE variants fine-tuned with different \(\theta\) values. As \(\theta\) increases, we observe a consistent decline in the recency bias of the RoPE layers. Specifically, the attention mass on the “End” segment drops from 0.314 in RNoPE-10k to 0.223 in RNoPE-4M, indicating a flatter attention distribution that reaches deeper into the context. This aligns with prior findings and theoretical analyses suggesting that larger \(\theta\) values expand the effective receptive field of the attention mechanism [54]. However, our empirical results indicate that this expanded receptive field in the RoPE layers introduces noise into the overall architecture, which disrupts the NoPE layers’ ability to compute similarities and perform retrieval effectively. This degradation is reflected in both attention mass and task performance: the needle attention mass in the NoPE layers drops from 0.0765 to 0.0369, and the needle evaluation score decreases from 8.036 to 6.203 as \(\theta\) increases from 10,000 to 4 million. These findings further underscore that the model spontaneously develops a “division of labor” mechanism during training, with distinct roles emerging between RoPE and NoPE layers.

From these observations, we draw the following insights:

1. Division of Labour: Combining NoPE and RoPE layers yields complementary strengths, with each type naturally assuming specialized roles after training. NoPE layers are adapted to information retrieval, while RoPE layers become effective at modeling local context due to their inherent recency bias.

2. Potential Efficiency Gains: Restricting the effective attention span of the RoPE layers in RNoPE models can reduce noise in the attention distribution and reinforce the functional specialization of each layer. Additionally, it lowers the computational cost (FLOPs) during training—particularly for longer context lengths—thereby improving training efficiency while maintaining or even enhancing performance.

Guided by these insights, we fine-tune a new variant, RNoPE-10k-swa, where “swa” denotes sliding window attention. This modification imposes a hard limit on the attention span of RoPE layers, operationalizing the second insight above. Specifically, the sliding window size for RoPE layers is set to 8,192, while the NoPE layers retain full attention to support long-context retrieval. All other training configurations remain identical to the RNoPE-10k variant, including the use of \(\theta = 10{,}000\). Evaluation results, presented in Table 3, show a marked improvement. The RNoPE-10k-swa variant achieves a needles-128k score of 9.562, surpassing both the baseline and the original RNoPE-10k model. Moreover, its NoPE layers exhibit a well-structured attention pattern with high attention mass on both the “Begin” and “End” segments from 2, reflecting strong long-context retrieval capabilities.

3 Model Architecture↩︎

Based on the analysis above, we make the following design choices on top of the Command R+ architecture [40]. First, we remove the QK-Norm component due to its poorly shaped attention patterns, which adversely impacted long context performance. Second, NoPE layers with a full attention span are introduced to enhance the model’s retrieval capabilities. Third, a sliding window size of 4,096 is applied to RoPE embeddings, leveraging RoPE’s inherent recency bias to improve performance on short-to-medium context ranges. In particular, the sliding-window approach has been employed in several prior works [26], [27], [55]. Regarding the number of layers, we perform an ablation study on the interleaving ratio of full attention and sliding window layers, testing the configurations of 1:1, 1:3, and 1:7. The results show that a 1:3 ratio strikes an optimal balance between computational efficiency and performance. We position the full attention layers at the end of each interleaving group, preceded by three sliding window layers. All other hyperparameters for the model architecture remain consistent with those outlined in Table 1. A visualization is shown in Figure 3.

For additional context, the resulting design shares similarities with other well-explored long-sequence architectures, such as Mega [56] and state-space models (SSMs) [57], [58]. For example, [56] and [59] introduced a multi-dimensional damped version of an exponential moving average component in conjunction with the gated attention unit (GAU) architecture [60], aiming to balance local and long-term dependencies, a common challenge in time-series modeling. Similarly, previous studies have proposed a diagonal variant of the S4 architecture [57], incorporating an exponentially decaying measure to enhance the model’s ability to capture long-range dependencies [61].

Earlier transformer variants, such as GPT-J [62] and GPT-NeoX [63], explored hybridizing RoPE and NoPE by applying rotational embeddings to a subset of the head dimensions. The work of p-RoPE [64] further advanced this line of research by analyzing RoPE in the frequency domain, revealing that low-frequency components capture semantic relationships while high-frequency components encode positional information. Similar to prior partial-RoPE approaches, p-RoPE removes selected low-frequency components along the head dimensions. These findings align with our observation that NoPE primarily facilitates retrieval (semantic focus), whereas RoPE exhibits a stronger recency bias (positional focus). However, prior works did not identify the cross-layer patterns underlying the “division of labor” phenomenon described in Section 2.2.3, which enables our approach to achieve both improved performance and practical engineering benefits.

4 Experiments↩︎

In this section, we detail the experiments conducted on the model architectures, covering the stages of pretraining, cooldown, supervised fine-tuning, and evaluations. Alongside long context evaluations, we provide a comprehensive assessment of short-context benchmarks, a dimension often disregarded in other long context studies. We train two models: one with the RoPE architecture as a baseline and another employing the architecture introduced in Section 3.

Pretraining and Cooldown. The models are pretrained for 5 trillion tokens of diverse data with batch size of 8 million tokens using FP8 precision format [65]. We use a cosine learning rate schedule of 5e-3 peak learning rate and 5% end learning rate with 8,000 linear warmup steps. From the pre-trained model, we linearly anneal the learning rate from 2.5e-4 to 1e-6 for 50,000 steps in BF16 precision. The context length was initially maintained at 8k for the first 35,000 steps, then extended to 32k and 128k for 10,000 steps and 5,000 steps, respectively. For the baseline model, the RoPE \(\theta\) values were increased to 1,000,000 for 32k and 8,000,000 for 128k contexts during the length extrapolation phase, while remaining constant for the RNope-SWA variant. Both models utilized the same interleaved training strategy outlined in Section 2.1.

Supervised Finetuning. We supervise fine-tune on top of the pretrained models. As the primary focus of this study is to evaluate the impact of architectural design on downstream tasks, preference training is deferred to future work. To preserve the long context capabilities of the model, the fine-tuning process utilizes interleaved datasets containing 8k and 128k prompt-response pairs. The long context SFT data at 128k includes Retrieval-Augmented Generation (RAG) [66] datasets, multilingual translation datasets with extended passages, long code instruction datasets, and long context retrieval datasets. Training was performed for two epochs across the entire dataset.

5 Experimental Results↩︎

Our evaluation contains a comprehensive analysis over standard benchmarks below 8k context length such as MMLU [44], HellaSwag [45], ARC [47], SAT [67], GSM8k [68], Winograde [69] and MBPP [70], as well as popular long context benchmarks with needles-in-a-haystack [38] and the retrieval and QA portion of Ruler [49]. We test the context lengths up to 256k so we can examine the impact of these choices in the extrapolation capability of the model. We denote the baseline model trained with RoPE scaling as “Baseline" and the architecture with interleaved attention span and position embeddings as”RNope-SWA".

5.1 Standard Benchmarks↩︎

In this set of evaluations, we evaluate baseline and RNope-SWA on a standard LLM benchmark covering various language, math and code capabilities. The results are shown in Table 4. We observe that the model is better or on-par on most of the metrics compared to the baseline and has some gains over the baseline numbers on certain benchmarks (+2.0% on MMLU and +1.8% on GSM8k). This set of results also indicates that although RNope-SWA explicitly removed position embeddings from 25% of all its layers, positional information is retained by the interleaving attention span and captured by RoPE from previous layers. RNope-SWA does not have the performance loss due to the removal of explicit position embeddings, as previous works have shown [25], [42].

5.2 Long Context Benchmarks↩︎

To evaluate the long context performance of these models, we use NIAH and the retrieval and QA components of Ruler [49]. To better understand how architectural choices affect long context performance, we also evaluate with context lengths extending beyond training sequence length. This allows us to assess how well these models can interpolate as well as extrapolate to unseen context lengths and how specific architectural decisions influence these capabilities.

5.2.1 NIAH Evaluations↩︎

Following the settings of section 2.2, we run NIAH test 256k context lengths. The scores are reported in Table 5 . The figure indicates that although both models are able to get close to perfect scores below the context length seen during training, RNope-SWA has better extrapolation capabilities and achieves almost no loss up to 256k context length while the baseline fails to extrapolate well – despite using a high RoPE \(\theta\) value of 8 million. We attach the graphical visualization with depth and length dimension expanded in Appendix 9.

5.2.2 Ruler Evaluations↩︎

First introduced in [49], the Ruler benchmark aims to provide a set of more difficult tasks than NIAH. It covers a wider range of retrieval under a Multi-Query/Key/Value settings, more realistic tasks with a long-context Question-Answering format and more. Although our modification of NIAH introduced more context variants and proves to be more difficult than the vanilla version, it still cannot test the limits of the model. Therefore, we evaluate our models on the retrieval and QA portion of the Ruler so we can better separate their performance. From the results, we can observe that the baseline model with RoPE \(\theta\) scaling approach suffers from a sharp drop in the longer context range, especially 64k and longer. Comparing the difference between the scores obtained at 8k and 256k, the baseline model dropped from 96.6 to 57.1 (about 41%) on retrieval and from 53.5 to 30.0 (about 44%) on QA, whereas the RNope-SWA model dropped from 96.1 to 74.8 (about 22.1%) on retrieval and from 55.5 to 42.5 (about 23.4%) on QA respectively. From the original Ruler Paper [49], models that adopt similar RoPE scaling approaches have shown a similar degradation over longer context lengths [17], [18], [71] as the baseline.

5.3 Impacts on Training and Inference↩︎

We also report the differences in training and inference speed, as well as memory requirements, of RNope-SWA compared to the baseline model. Let \(S\) denote the sliding window size and \(L\) represent the full training context length. During training, 75% of all layers now operate with a computational complexity of \(\mathcal{O}(\textit{S} \textit{L})\), rather than the quadratic complexity of \(\mathcal{O}(L^2)\). This results in the model being approximately 50% faster than the baseline at a 64k context length and nearly 2x faster at 128k in terms of training throughput, using flash attention [72]–[74] and a sequence-parallel scheme similar to [75], [76]. With alternative implementations, such as Ring Attention [77], [78] or its variants [79], sliding window adoption can reduce key-value block communication overhead if carefully sharded along the sequence dimension, potentially improving speed. In theory, KV cache memory and time complexity of RNoPE-SWA can be reduced by up to 75%. Empirically, we observed a 44% latency reduction at 132k input and 96 output tokens, and nearly 70% at 990k input and 8 output tokens—approaching the theoretical limit as sequence length grows. Increasing the ratio of sliding window to full attention layers can further improve speed and memory efficiency.

6 Discussions and Future Work↩︎

In this paper, we introduced RNope-SWA, an architecture that interleaves NoPE and RoPE position embeddings with varying attention spans (RNope-SWA). RNope-SWA is able to strike a balance between effective attention modeling and computational efficiency, achieving a nearly 4x reduction in KV cache size and significantly boosting both training and inference speeds without compromising performance. The integration of NoPE layers with full attention spans enhances long context capabilities, eliminating the need for RoPE scaling. This simplification improves the stability of training and delivers excellent long context performance.

Our findings align with recent work, such as YoCo [80], Jamba-1.5 [81], and MiniMax-01 [82], which demonstrate that hybrid attention mechanisms generally outperform full attention mechanisms in handling long contexts. However, the underlying reasons behind this seemingly counterintuitive observation remain largely unexplored. This opens an intriguing area of study, particularly as models push towards multi-million-token context lengths. Re-visiting and re-thinking the foundational components of transformer architectures, such as attention mechanisms, may become essential to accommodating these extreme requirements. Recent works [50], [51], [83] have begun to explore this direction by focusing on reducing attention noise across large context windows, a promising approach to refine the performance of attention modules.

7 Attention Distribution of All Lengths↩︎

This table contains the attention distribution of RoPE and RNoPE variants from Section 2.2.2 over 8k, 32k and 128k sequence lengths.

8 Attention Distribution of RoPE and QK-Norm variants↩︎

In this section, we further investigate the suboptimal performance of the QK-Norm variant. We present three plots comparing the attention distribution between the RoPE and QK-Norm variants across sequence lengths of 8k, 32k, and 128k on needle samples, following the setup outlined in Section 2.2.2. Additionally, we provide the aggregated attention entropy for each variant to quantitatively support the arguments.

To enhance the clarity of the distribution plots, we preprocess the attention distribution array by removing the first 10 tokens and the last 3% of tokens from each sequence. This preprocessing step mitigates the disproportionate attention mass resulting from the attention sink effect and the recency bias observed in RoPE, thereby making the attention patterns more interpretable. We then compute a moving average with a window size of 100 tokens and average the results across all samples and layers to generate the final distributions.

From Figure 4, we observe that the QK-Norm variant exhibits a lower spike on needle tokens but distributes more attention mass across context tokens. However, it also demonstrates a stronger recency bias compared to the RoPE variant. This characteristic results in a lower signal-to-noise ratio for the QK-Norm variant, which hampers its ability to effectively retrieve relevant information from long contexts. To further quantify this observation, we calculate the entropy values of the attention distributions for both variants, averaging across samples and layers at each sequence length. The results, listed in Table 9, show that the QK-Norm variant has significantly higher entropy values than the RoPE variant. This aligns with its weaker performance in long context retrieval tasks, as higher entropy reflects a more dispersed and less focused attention distribution.

9 Needles Score at 256k↩︎

References↩︎