LongAttn: Selecting Long-context Training Data via Token-level Attention
Authors: Longyun Wu, Dawei Zhu, Guangxiang Zhao, Zhuocheng Yu, Junfeng Ran, Xiangyu Wong, Lin Sun, Sujian Li
Accepted at: acl25

LONGATTN: Selecting Long-context Training Data via Token-level Attention

Longyun Wu*♥♦  Dawei Zhu*♣♦  Guangxiang Zhao◊  Zhuocheng Yu♣♣
Jungfeng Ran♥♦  Xiangyu Wong♣  Lin Sun†◊  Sujian Li‡♣♣
♥ School of Software and Microelectronics, Peking University
♣ School of Computer Science, Peking University
♦ Key Laboratory of Computational Linguistics, MOE, Peking University
◊ Qiyuan Tech
wulongyun@stu.pku.edu.cn  sunlin1@360.cn  lisujian@pku.edu.cn

Abstract

With the development of large language models (LLMs), there has been an increasing need for significant advancements in handling long contexts. To enhance long-context capabilities, constructing high-quality training data with long-range dependencies is crucial. Existing methods to select long-context data often rely on sentence-level analysis, which can be greatly optimized in both performance and efficiency. In this paper, we propose a novel token-level framework, LongAttn, which leverages the self-attention mechanism of LLMs to measure the long-range dependencies for the data. By calculating token-level dependency strength and distribution uniformity of token scores, LongAttn effectively quantifies long-range dependencies, enabling more accurate and efficient data selection. We filter LongABC-32K from open-source long-context datasets (ArXiv, Book, and Code). Through our comprehensive experiments, LongAttn has demonstrated its excellent effectiveness, scalability, and efficiency. We have released our code and the high-quality long-context training data LongABC-32K at https://github.com/Lyun0912-wu/LongAttn.

1 Introduction

Large language models (LLMs) have achieved impressive performance across a broad spectrum of traditional natural language processing tasks (Touvron et al., 2023). To effectively address real-world applications, these models further require enhanced capabilities in handling longer contexts, particularly in key areas such as in-context learning (Brown et al., 2020), real-world question-answering based on lengthy documents (Wang et al., 2024b), long-context dialogue with historical context (Packer et al., 2023), and comprehensive document summarization (Koh et al., 2022).

To enhance LLMs’ long-context processing capabilities, data engineering remains fundamental. Simple methods to construct long-context datasets are through naïve methods like concatenating short texts or randomly sampling existing sources (e.g., CommonCrawl, GitHub). However, studies by de Vries (2023) and Chen et al. (2024a) emphasize that data obtained through such approaches fail to effectively improve long-context capabilities of LLMs because the data lack meaningful long-range dependencies. Inspired by this, a line of studies focus on exploring the identification and selection of high-quality long-context with consideration of relations between text segments were proposed. ProLong (Chen et al., 2024a) measures long-range dependencies between segments based on the relative perplexity and relative distance. Lv et al. (2024) develop a set of metrics including complexity, coherence, and cohesion based on various kinds of text segments (i.e., sliding windows, sentences, paragraphs) to measure the quality of long texts. However, these methods have two main drawbacks: (1) Linguistic metrics do not fully align with the underlying mechanisms of LLMs, as they often fail to capture fine-grained token-level relationships. (2) They are computationally expensive and inefficient. For example, ProLong reports that the speed of a 7B parameter model is roughly 1/16 of that of a 350M parameter model, making such methods challenging to scale for LLMs.

Attention mechanisms have been proven to effectively model context understanding (Beltagy et al., 2020; Zaheer et al., 2020). Some studies focusing on attention mechanisms and positional encoding have shown that they can significantly improve a model’s long-context ability (Peng and Quesnelle, 2023; Peng et al., 2023). Motivated by this, we propose to address the limitations of sentence-level selection methods leveraging the rich information provided in the attention mechanism. Specifically, we propose LongAttn, a simple yet effective frame-
Figure 1: (a) How to measure long-range dependencies at the token level by using the self-attention mechanism. \( D_{ST} \) indicates that the tokens in this data have strong long-distance dependencies, while \( D_{UT} \) prevents negative impacts from individual tokens’ high scores. (b) The comparison of long-context retrieval capabilities of models trained with different scales of tokens selected randomly, with sentence-level ProLong, and with LongAttn (ours).

work that leverages the attention patterns of LLMs to analyse token-level dependency for long-context data selection.

LongAttn utilizes the long-range dependency indicator, \( LSD_T \), to measure the strength of dependencies between tokens separated by a distance of at least \( k \), which we define as the minimum token distance. We break down the indicator into two scores: dependency strength(\( D_{ST} \)) and distribution uniformity(\( D_{UT} \)). As shown in Figure 1a, \( D_{ST} \) measures the strength of dependencies between tokens separated by a distance of at least \( k \), and \( D_{UT} \) serves as a correction term, ensuring a consistent distribution of token scores and preventing individual tokens with excessively high attention scores from skewing the overall dependency assessment. To enhance computational efficiency and avoid the Attention Sink (Xiao et al., 2023), we use the attention score calculated by the first decoder layer of LLaMA (Dubey et al., 2024). To better integrate \( D_{ST} \) and \( D_{UT} \), we normalize them to the same value range and then multiply the distribution uniformity (\( D_{UT} \)) by a correction factor \( \alpha \). In this way, our framework effectively quantifies the degree of contextual information aggregation at the token level, providing a reliable criterion for selecting high-quality long-context data.

Through comprehensive experiments, the LongAttn framework has demonstrated significant advantages. We selected Arxiv, Book, and Code as the long-context datasets to be studied. After pre-processing, we used the LongAttn framework to make selections, and the resulting data is referred to LongABC-32K. Datasets selected using the ProLong framework (Chen et al., 2024a) and random selection mechanism are designated as ProLong-32K and Random-32K, respectively. As shown in Figure 1b, we compare the long-context retrieval abilities of models trained on these datasets across different token scales. The experimental results demonstrate that models trained on LongABC-32K consistently perform the best, even surpassing those trained on 20B tokens from the randomly selected dataset, despite using only 5B tokens. Through further experiments, we found that, in addition to its effectiveness (As seen in 5), LongAttn exhibits excellent scalability (Performs better with attention map from larger models, as seen in 6.2) and efficiency (As seen in 6.3). Our contributions are summarized as follows:
• We propose LongAttn, a framework which is the first to analyze long-range dependencies at the token level by using self-attention mechanisms.
• We will release LongABC-32K, a high-quality long-context dataset with strong long-range dependencies.
• Through comprehensive experiments, we have demonstrated LongAttn’s excellent effectiveness, scalability, and efficiency.

2 Related Work

Long-context LLMs The ability to process extensive contextual information is a crucial aspect of language models, with context length serving
as a key determinant of their processing capacity. To enable models to accommodate longer inputs and outputs, numerous methodologies have been developed to modify the model’s architecture (Liu et al., 2025). Some approaches focus on altering positional encoding techniques. For instance, Chen et al. (2023a) introduced Positional Interpolation, while PoSE (Zhu et al., 2023) simulates longer texts by modifying positional encodings during training. Similar methods include YaRN (Peng et al., 2023) and LongRoPE (Ding et al., 2024). Other strategies involve modifying the attention mechanism. For example, STRING (An et al., 2024) shifts well-trained positions to overwrite originally ineffective positions during inference, and SelfExtend (Jin et al., 2024) extends the context window of Large Language Models (LLMs) by constructing bi-level attention information. Besides the methods used, the training data is also crucial. Below are related works on data.

Pre-training data Training data that exhibits long-range dependency patterns is crucial for enhancing the model’s ability to handle extended contextual information. For post-training data, numerous methodologies have been explored to generate synthetic long-context data (Wang et al., 2024a; Chen et al., 2024b; Bai et al., 2024; Wu et al., 2024). Conversely, for pre-training data, the predominant approach involves the curation and selection of relevant text from existing corpora,

which is exemplified by prominent models including Qwen (Bai et al., 2023a) and LLaMA (Touvron et al., 2023). While scaling laws suggest that a model’s capabilities improve with more data (Kaplan et al., 2020), large volumes of data bring about high resource demands. Therefore, optimizing data utilization more effectively should become a key area of research. ProLong (Chen et al., 2024a) proposes a framework for calculating long-distance dependencies of data at the sentence level. Long-Wanjuan (Lv et al., 2024) also designed metrics and filtered data based at the sentence level. However, Xiong et al. (2023) assert that the key factor affecting the long-context ability of LLMs is the positional encoding’s capacity to aggregate information from distant tokens. Our method focuses on token-level long-distance dependencies to select high-quality long-context data.

3 Methodology

As shown in Figure 2, our proposed method can be divided into three steps. Firstly, we gather and preprocess the data to a predetermined length. Subsequently, we employ the self-attention mechanism of a LLM to compute the long-distance dependency score for each data instance. Finally, we filter the data based on the score and utilize the refined dataset for continued pre-training of the model.

3.1 Data Collection and Preprocessing

To ensure the training data is suitable for long-context modeling, we carefully curate and preprocess our dataset. We choose books, code, and Arxiv papers as our primary sources of long-context data, drawing from open-source pre-training datasets such as RedPajama (Weber et al., 2024) and Dolma (Soldaini et al., 2024). These sources are known for their rich content and long sequences, which are essential for training models with extended context windows.

Given that the computational complexity of self-attention layers grows quadratically with sequence length, we set the context length to 32k tokens in this work. This length strikes a balance between capturing long-range dependencies and maintaining reasonable computational complexity. To segment/divide the data into 32k-token chunks/segments, we employ a sliding-window approach, which is more effective than naive truncation in preserving the integrity of the information. Let the total number of tokens in a text be \( n \). The sliding-window strategy is as follows:
• If \( 32768(32k) < n \leq 65536(64k) \), take both the front and back windows.
• If \( 65536(64k) < n \leq 98304(96k) \), take the front, back, and middle windows.
• If \( n > 98304(96k) \), iteratively take the front and back windows until one of the two conditions above is met.
The detailed algorithm is presented in Appendix A. After preprocessing, we obtain the long-context pre-training dataset LongABC-32K-Raw, which we denote as \( \mathcal{D} \).

3.2 Assess Long-distance Dependency via Token-level Attention

To effectively select high-quality long-context data, we need to accurately measure the long-range dependencies within the data. In this section, we
Figure 2: LongAttn Framework: After preprocessing the data, the long-distance dependency strength at the token-level is analyzed using the self-attention mechanism of an LLM. This analysis serves as the basis for filtering the data, which is then used for continual pre-training of a base model that initially lacks long-context capabilities, resulting in our LongAttn model

detail the process of assessing long-distance dependencies in the data using token-level attention mechanisms.

3.2.1 Token-level Dependency Strength
Given a data instance \( s \in \mathcal{D} \), we input it into an LLM and extract the masked self-attention matrix \( M \) from the first transformer decoder layer to quantify the long-range dependencies within the data. The choice of using the first layer is driven by two primary reasons: (1) It is computationally efficient, requiring approximately 1/32 of the inference time; (2) Due to the Attention Sink phenomenon (Xiao et al., 2023), deeper layers of the model tend to disproportionately focus on the initial tokens, irrespective of their semantic relevance to the language modeling task. Consequently, leveraging the shallow layers of the model’s decoder is more optimal for capturing the contextual dependencies among tokens in the data.

Define \( A_{m,n} \) as the cumulative attention score assigned by \( n \) to the first \( m \) tokens (i.e., tokens from position 1 to \( m \)):

\[
A_{m,n} = \sum_{i=1}^{m} M_{i,n}
\]

where \( M_{i,n} \) represents the attention score assigned by the \( n \)-th token to the \( i \)-th token. Since the self-attention matrix \( M \) has been normalized by the softmax function, it follows that \( A_{n,n} = 1 \). For the \( n \)-th token in the data, where \( n > k \), \( A_{n-k,n} \) represents the sum of attention scores of all tokens located at least \( k \) positions ahead of it. We define the contextual dependency strength of the \( n \)-th token as:

\[
DS_T^n = \frac{A_{n-k,n}}{A_{n,n}} = A_{n-k,n}
\]

which quantifies the proportion of attention scores assigned to tokens at least \( k \) positions prior to the \( n \)-th token, relative to the total attention scores. For cases where \( n \leq k \), we define \( DS_T^n = 0 \) to account for insufficient context. Finally, the token-level contextual dependency strength of the entire data instance is defined as the average of \( DS_T^n \) over all tokens:

\[
DS_T = \frac{1}{L} \sum_{i=1}^{L} DS_T^i
\]
\[
= \frac{1}{L} \sum_{i=k+1}^{L} DS_T^i
\]
\[
= \frac{1}{L} \sum_{t} M_t
\]
where \( L \) is the total number of tokens in the data and \( M_t \) represents the lower triangular matrix in the bottom left corner of matrix \( M \):

\[
M_t = \begin{pmatrix}
M_{k+1,1} & 0 & \cdots & 0 \\
M_{k+2,1} & M_{k+2,2} & \cdots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
M_{L,1} & M_{L,2} & \cdots & M_{L,L-k}
\end{pmatrix}
\] (6)

3.2.2 Distribution Uniformity of Token Scores
While \( DS_T \) provides a measure of dependency strength, it is important to ensure that individual tokens with high scores do not disproportionately influence the overall dependency assessment. For example, In the previously mentioned Attention Sink phenomenon, the first token’s scores very high in deeper decoder self-attention layers, which can have a significantly negative impact. Instead, the scores across the entire data segment should be consistently high. To achieve this, we introduce the distribution uniformity of token scores \( DU_T \) to measure the uniformity of the score distribution:

\[
DU_T = -Variance(M_t)
\] (7)

This correction term helps to prevent individual tokens with excessively high attention scores from skewing the overall dependency assessment.

3.2.3 Collaborative Ensemble
To obtain a comprehensive measure of long-range dependencies, we combine the dependency strength \( DS_T \) and the distribution uniformity \( DU_T \). Due to the differences in the magnitudes of \( DU_T \) and \( DS_T \), as shown as Appendix E, we compute \( DU_T \) and \( DS_T \) for all data and then standardize them to independent normal distributions. We then use the following formula to calculate the final long-distance dependency score:

\[
LDS_T = Std(DS_T) + \alpha \cdot Std(DU_T)
\] (8)

where \( \alpha \) is a correction factor that balances the contributions of \( DS_T \) and \( DU_T \), and \( Std \) represents Z-Score standardization.

4 Experimental Setup

4.1 LongAttn Setup and Training Details
In the process of filtering data using LongAttn, we utilize the first transformer decoder layer of LLaMA-3.1 to calculate long-distance dependency score. The length of each data segment \( L \) is 32768 and the minimum token distance \( k \) is set to \( L/4 \) (i.e., 8192). We set the correction factor \( \alpha \) in the Eq.8 to 0.5.

We adopt Adjusted Base Frequency (ABF) (Xiong et al., 2023) to continual pretrain LLaMA-3, extending the context window size to 32,768 by adjusting the RoPE theta parameter. The continued pre-training is based on the Megatron training framework (Shoeybi et al., 2019), utilizing 8x8 H800 GPUs. Detailed parameters can be found in the Appendix B.1.

4.2 Continual Pre-trained Datasets
We form the following datasets by combining short-context data with selections made through random sampling, the ProLong framework, LongAttn based on LLaMA-3.1-8B, and LongAttn based on LLaMA-3.1-70B: \( \mathcal{D}_{Rx}(x \in \{1, 3, 5, 10, 20\}) \), \( \mathcal{D}_{Px}(x \in \{1, 3, 5, 10\}) \), \( \mathcal{D}_{Ax,8B}(x \in \{1, 3, 5, 10\}) \), and \( \mathcal{D}_{Ax,70B}(x \in \{1, 3, 5, 10\}) \), with \( x \) representing the data size in Billions.

To ensure the diversity of the filtered data, we apply the filtering process within each category of datasets separately. For detailed data composition, please refer to the Appendix C.

4.3 Baselines
Data-Scale Comparison To demonstrate the effectiveness of LongAttn, We conduct a data-scale comparison of the long-context retrieval capabilities of models continued pre-trained on \( \mathcal{D}_{Rx}(x \in \{1, 3, 5, 10, 20\}) \), \( \mathcal{D}_{Px}(x \in \{1, 3, 5, 10\}) \), \( \mathcal{D}_{Ax,8B}(x \in \{1, 3, 5, 10\}) \), and \( \mathcal{D}_{Ax,70B}(x \in \{1, 3, 5, 10\}) \).

Fixed-Scale Method Comparison To demonstrate the superiority of LongAttn, we conduct fixed-data method comparison of the models trained on \( \mathcal{D}_{Rx}(x \in \{5, 10, 20\}) \), \( \mathcal{D}_{P5} \), \( \mathcal{D}_{A5,8B} \), and \( \mathcal{D}_{A5,70B} \). Additionally, we compare them with similarly sized models that have excellent long-context capabilities. Details of the baselines can be found in Appendix D.

4.4 Evaluation Tasks
We assess the capability of the base model, continually pre-trained within the current window length, based on the following long-context and short-context criteria: (1) The best reflection of the base model’s long-context capabilities is its long-context retrieval ability, followed by its performance on other long-context tasks. (2) No degradation in
<table>
  <tr>
    <th rowspan="2">Method</th>
    <th rowspan="2">Tokens</th>
    <th colspan="3">Niah-Single</th>
    <th colspan="3">Niah-Multickey</th>
    <th colspan="2">Niah-Mult-</th>
    <th rowspan="2">Avg.</th>
  </tr>
  <tr>
    <th>Sigle-1</th>
    <th>Sigle-2</th>
    <th>Sigle-3</th>
    <th>MK-1</th>
    <th>MK-2</th>
    <th>MK-3</th>
    <th>Value</th>
    <th>Query</th>
  </tr>
  <tr>
    <td>Random</td>
    <td></td>
    <td>99.8</td>
    <td><b>100.0</b></td>
    <td>93.4</td>
    <td><b>91.0</b></td>
    <td>11.6</td>
    <td>11.4</td>
    <td><b>91.7</b></td>
    <td>93.2</td>
    <td>74.0</td>
  </tr>
  <tr>
    <td>ProLong</td>
    <td></td>
    <td>99.4</td>
    <td>99.8</td>
    <td>92.4</td>
    <td>89.2</td>
    <td>10.8</td>
    <td>24.0</td>
    <td>91.6</td>
    <td><b>93.6</b></td>
    <td>75.1</td>
  </tr>
  <tr>
    <td>LongAttn-8</td>
    <td>1 B</td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td>91.4</td>
    <td>88.6</td>
    <td>16.2</td>
    <td>19.2</td>
    <td>90.3</td>
    <td>93.4</td>
    <td>74.9</td>
  </tr>
  <tr>
    <td>LongAttn-70</td>
    <td></td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td><b>95.4</b></td>
    <td>88.0</td>
    <td><b>29.0</b></td>
    <td><b>35.0</b></td>
    <td>90.4</td>
    <td>92.4</td>
    <td><b>78.8</b></td>
  </tr>
  <tr>
    <td>Random</td>
    <td></td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td>86.2</td>
    <td>92.8</td>
    <td>62</td>
    <td>8.6</td>
    <td>70.0</td>
    <td>95.9</td>
    <td>76.9</td>
  </tr>
  <tr>
    <td>ProLong</td>
    <td></td>
    <td><b>100.0</b></td>
    <td>99.8</td>
    <td>79.6</td>
    <td><b>93.8</b></td>
    <td><b>60.4</b></td>
    <td><b>32.0</b></td>
    <td>85.9</td>
    <td>95.0</td>
    <td>80.8</td>
  </tr>
  <tr>
    <td>LongAttn-8</td>
    <td>3 B</td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td>88.8</td>
    <td>92.2</td>
    <td>60.0</td>
    <td>31.2</td>
    <td>79.7</td>
    <td>94.7</td>
    <td>80.8</td>
  </tr>
  <tr>
    <td>LongAttn-70</td>
    <td></td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td><b>91.6</b></td>
    <td>93.6</td>
    <td>57.4</td>
    <td>19.8</td>
    <td><b>88.8</b></td>
    <td><b>96.2</b></td>
    <td><b>80.9</b></td>
  </tr>
  <tr>
    <td>Random</td>
    <td></td>
    <td><b>100.0</b></td>
    <td>99.8</td>
    <td>81.8</td>
    <td><b>94.8</b></td>
    <td>56.4</td>
    <td>11.8</td>
    <td>84.4</td>
    <td>96.5</td>
    <td>78.2</td>
  </tr>
  <tr>
    <td>ProLong</td>
    <td></td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td>78.0</td>
    <td>92.8</td>
    <td>64.8</td>
    <td>40.4</td>
    <td>77.8</td>
    <td>95.9</td>
    <td>81.2</td>
  </tr>
  <tr>
    <td>LongAttn-8</td>
    <td>5 B</td>
    <td><b>100.0</b></td>
    <td>99.8</td>
    <td>81.6</td>
    <td>92.4</td>
    <td>62.6</td>
    <td>37.2</td>
    <td><b>87.6</b></td>
    <td><b>97.3</b></td>
    <td>82.3</td>
  </tr>
  <tr>
    <td>LongAttn-70</td>
    <td></td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td><b>83.8</b></td>
    <td>92.8</td>
    <td><b>84.8</b></td>
    <td><b>46.8</b></td>
    <td>78.8</td>
    <td>95.2</td>
    <td><b>85.2</b></td>
  </tr>
  <tr>
    <td>Random</td>
    <td></td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td>84.0</td>
    <td>92.6</td>
    <td>58.2</td>
    <td>14.2</td>
    <td>90.9</td>
    <td><b>96.9</b></td>
    <td>79.6</td>
  </tr>
  <tr>
    <td>ProLong</td>
    <td></td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td>83.4</td>
    <td>92.8</td>
    <td>74.4</td>
    <td>32.2</td>
    <td>88.7</td>
    <td>95.5</td>
    <td>83.4</td>
  </tr>
  <tr>
    <td>LongAttn-8</td>
    <td>10 B</td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td><b>87.4</b></td>
    <td><b>93.0</b></td>
    <td>72.2</td>
    <td>23.0</td>
    <td><b>93.1</b></td>
    <td>96.8</td>
    <td>83.2</td>
  </tr>
  <tr>
    <td>LongAttn-70</td>
    <td></td>
    <td><b>100.0</b></td>
    <td><b>100.0</b></td>
    <td>86.8</td>
    <td>92.4</td>
    <td><b>80.6</b></td>
    <td><b>34.4</b></td>
    <td>92.0</td>
    <td>96.5</td>
    <td><b>85.3</b></td>
  </tr>
  <tr>
    <td>Random</td>
    <td>20 B</td>
    <td>100.0</td>
    <td>100.0</td>
    <td>84.6</td>
    <td>91.0</td>
    <td>66.2</td>
    <td>22.4</td>
    <td>93.3</td>
    <td>96.5</td>
    <td>81.8</td>
  </tr>
</table>

Table 1: Models trained with different methods for selecting varying scales of tokens were evaluated on complex NIAH tasks. Random, ProLong, LongAttn-8, and LongAttn-70 represent random selection, selection based on the ProLong framework, selection based on LongAttn with LLaMA-3.1-8B, and selection based on LongAttn with LLaMA-3.1-70B, respectively. And bold number is used to highlight the better-performing models within each data size category.

<table>
  <tr>
    <th rowspan="2">Model</th>
    <th rowspan="2">Trained Dataset</th>
    <th colspan="4">Short-Context Task</th>
    <th rowspan="2">Avg.</th>
  </tr>
  <tr>
    <th>MMLU</th>
    <th>HS</th>
    <th>HE</th>
    <th>OBQA</th>
  </tr>
  <tr>
    <td rowspan="5">LLaMA-3-Base</td>
    <td>†</td>
    <td>65.9</td>
    <td>49.9</td>
    <td>25.0</td>
    <td>72.0</td>
    <td>53.2</td>
  </tr>
  <tr>
    <td>\( D_{R5} \)</td>
    <td>61.8</td>
    <td>52.4</td>
    <td>19.5</td>
    <td>81.8</td>
    <td>53.9</td>
  </tr>
  <tr>
    <td>\( D_{P5} \)</td>
    <td>61.0</td>
    <td>38.3</td>
    <td>23.2</td>
    <td>79.4</td>
    <td>50.5</td>
  </tr>
  <tr>
    <td>\( D_{A5,8B} \)</td>
    <td>61.6</td>
    <td>47.1</td>
    <td>25.6</td>
    <td><b>82.6</b></td>
    <td>54.2</td>
  </tr>
  <tr>
    <td>\( D_{A5,70B} \)</td>
    <td>61.0</td>
    <td><b>52.8</b></td>
    <td><b>28.1</b></td>
    <td>80.4</td>
    <td><b>55.6</b></td>
  </tr>
</table>

Table 2: The short-context fundamental capabilities of our continued pre-trained models and LLaMA-3-base. † indicates no training. MMLU, HS, HE, and OBQA stand for the MMLU, HellaSwag, HumanEval, and OpenBookQA tasks, respectively.

short-context performance. The evaluation tasks can be divided into the following parts:

Long-context Retrieval  Retrieval ability is the most crucial and best reflects the model’s long-context ability before post-training. The ‘Needle In A Haystack’ task analysis in-context retrieval ability of long-context LLMs. The original ‘needle in a haystack’ task was relatively simple. RULER (Hsieh et al., 2024) introduced a more detailed and complex ‘needle in a haystack’ task, and we use RULER with a length of 32k to comprehensively evaluate long-context retrieval ability.

Long-context Benchmark  In addition to retrieval ability, we also want to evaluate the model’s performance on formal long-context tasks. Long-Bench (Bai et al., 2023b) is the first proposed bilingual long-context benchmark, which includes a total of 21 tasks categorized into 6 main types, with task lengths ranging from about 0 to 20k. RULER provides longer, variable-length evaluations across 13 complex tasks. Here, we will evaluate the tasks at the 32k length to assess changes in the model’s long-context capabilities.

Fundamental Abilities of LLMs.  We use HumanEval (Chen et al., 2021) to assess code evaluation capability and OpenBookQA (Mihaylov et al., 2018) to assess book knowledge extraction ability. Additionally, we use Hellaswag(Zellers et al., 2019) and MMLU (Hendrycks et al., 2020) to assess its broader short-context fundamental capabilities.

5 Experimental Results

We validate the effectiveness, scalability, and high efficiency of LongAttn through a series of data-scale and Lateral comprehensive experiments.

5.1 Performance on Retrieval Ability

We evaluate the retrieval capabilities of models trained with LongAttn-selected data and com-
Figure 3: (a) and (b) show the performance of other long-context LLMs and LongAttn-trained models on the RULER and complex NIAH tasks. (c) and (d) show the performance of models trained with different methods on the same tasks. Toge. and LLORA represent Together and LongLLORA, respectively. 5B-LA and 10B-LA represent models trained on 5B and 10B tokens selected by LongAttn. LA-8 and LA-70 represent LongAttn based on 8B and 70B models, respectively.

pare them with models trained on randomly selected data and ProLong-selected data. The results are shown in Table 1. The models trained with LongAttn-selected data consistently outperform those trained on randomly selected or ProLong-selected data across all data scales, demonstrating the effectiveness of LongAttn in improving data quality for long-context modeling. Notably, models trained on a smaller amount of data filtered using our method even outperform those trained on a larger amount of randomly selected data in retrieval tasks. For example, the model trained on just 5B tokens filtered by LongAttn outperforms models trained on 10B or even 20B randomly selected tokens. This indicates that LongAttn can significantly enhance the efficiency of data usage for long-context pre-training.

5.2 Performance on Fundamental Abilities
The results in Table 2 indicate that data selected by LongAttn not only maintains the model’s short-context capabilities but enhances them in specific domains. For example, the LongABC-32K-Raw dataset includes book and code data, and our model performs well on short-context tasks such as OpenBookQA (Mihaylov et al., 2018) and HumanEval (Chen et al., 2021).

However, there is a slight decline in performance on MMLU (Hendrycks et al., 2020). This is expected, as we do not include such data during continual pre-training, so the base model experienced some forgetting in these areas.

5.3 Performance on Long-context Benchmark
As shown in Figure 3a and 3c, models trained on data filtered by LongAttn outperform those trained on equivalent amounts of data selected randomly or by ProLong. LongAttn’s performance is also comparable to models trained on larger data volumes. Additionally, on the RULER-32K benchmark, LongAttn outperforms all other long-context models of similar parameter sizes. The specific experimental results can be found in Appendix F.

As shown in Table 3, we compare model performance on LongBench, which consists of 21 evaluation tasks. We calculate the average score for each of the six categories to represent overall performance. The results show that LongAttn outperforms models trained on equivalent data selected randomly or by ProLong in almost all tasks and even surpasses models trained on larger amounts of randomly selected data. However, while 5B data selected by LongAttn-70 outperforms 10B randomly selected data, it does not perform as well as 5B data selected by LongAttn-8. We speculate this is because the average context length in LongBench is far below 32k, thus not effectively showcasing the advantage of 5B data selected by LongAttn-70.

6 Analysis

6.1 Ablation Study
To investigate the impact of the constraint factor \( \alpha \) and the correction term \( DU_T \) on regulating \( LDS_T \), we conduct ablation experiments on the \( \mathcal{D}_{43} \) and \( \mathcal{D}_{45} \) datasets using retrieval tasks. The default setting of the constraint factor \( \alpha \) is 0.5.

As shown in the table 4, we can see that the correction term \( DU_T \) plays a positive role in the data selection results. In addition, the constraint on the dependency strength \( DS_T \) by \( DU_T \) should not be too large, which suggests that the constraint on \( DS_T \) by \( DU_T \) should be moderate to avoid over-correction.
<table>
  <tr>
    <th>Method</th>
    <th>Number of Tokens Trained</th>
    <th>Single-Doc QA</th>
    <th>Multi-Doc QA</th>
    <th>Summarization</th>
    <th>Few-shot Learning</th>
    <th>Synthetic Tasks</th>
    <th>Code Completion</th>
    <th>Avg. Score</th>
  </tr>
  <tr>
    <th colspan="9">Trained on 5B Tokens from Different Methods</th>
  </tr>
  <tr>
    <td>Random</td>
    <td>5B</td>
    <td>10.11</td>
    <td>6.57</td>
    <td>13.72</td>
    <td>64.10</td>
    <td>1.83</td>
    <td>65.05</td>
    <td>24.46</td>
  </tr>
  <tr>
    <td>ProLong</td>
    <td>5B</td>
    <td>11.95</td>
    <td><b>12.59</b></td>
    <td>17.87</td>
    <td>63.33</td>
    <td>4.15</td>
    <td>65.01</td>
    <td>26.93</td>
  </tr>
  <tr>
    <td>LongAttn-8</td>
    <td>5B</td>
    <td><b>13.01</b></td>
    <td>11.20</td>
    <td>18.96</td>
    <td><b>64.62</b></td>
    <td><b>5.12</b></td>
    <td><b>65.06</b></td>
    <td><b>27.46</b></td>
  </tr>
  <tr>
    <td>LongAttn-70</td>
    <td>5B</td>
    <td>12.39</td>
    <td>9.33</td>
    <td><b>19.72</b></td>
    <td>64.1</td>
    <td>3.42</td>
    <td>65.03</td>
    <td>26.78</td>
  </tr>
  <tr>
    <th colspan="9">Trained on over 5B Tokens Selected Randomly</th>
  </tr>
  <tr>
    <td>Random</td>
    <td>10B</td>
    <td>9.41</td>
    <td>8.93</td>
    <td>19.30</td>
    <td>63.89</td>
    <td>4.83</td>
    <td>65.57</td>
    <td>26.27</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>20B</td>
    <td>11.45</td>
    <td>11.72</td>
    <td>20.41</td>
    <td>64.13</td>
    <td>9.67</td>
    <td>66.51</td>
    <td>28.23</td>
  </tr>
</table>

Table 3: The performance of models continued pre-trained using data filtered by different methods on LongBench. Random, ProLong, LongAttn-8, and LongAttn-70 represent data selected randomly, data selected using the ProLong framework, data selected by the LongAttn framework with LLaMA-3.1-8B, and data selected by the LongAttn framework with LLaMA-3.1-70B, respectively.

<table>
  <tr>
    <th>Model</th>
    <th>RULER-NIAH-32K</th>
  </tr>
  <tr>
    <td><i>LongAttn_{D_{A^3}}</i></td>
    <td><b>80.83</b></td>
  </tr>
  <tr>
    <td>w/ \( \alpha = 1 \)</td>
    <td>79.49(-1.34)</td>
  </tr>
  <tr>
    <td>w/o \( D_{U_T} \)</td>
    <td>78.28(-2.55)</td>
  </tr>
  <tr>
    <td><i>LongAttn_{D_{A^5}}</i></td>
    <td><b>82.30</b></td>
  </tr>
  <tr>
    <td>w/ \( \alpha = 1 \)</td>
    <td>81.05(-1.25)</td>
  </tr>
  <tr>
    <td>w/o \( D_{U_T} \)</td>
    <td>82.11(0.19)</td>
  </tr>
</table>

Table 4: Ablation experiments on the constraint factor \( \alpha \) and the correction term \( D_{U_T} \) were conducted on the RULER-NIAH-32K task.

<table>
  <tr>
    <th>Method</th>
    <th>Model</th>
    <th>GPU Hours</th>
  </tr>
  <tr>
    <td>ProLong</td>
    <td>OPT-350M</td>
    <td>30</td>
  </tr>
  <tr>
    <td>ProLong</td>
    <td>LLaMA-3.1-8B</td>
    <td>600</td>
  </tr>
  <tr>
    <td>LongAttn</td>
    <td>LLaMA-3.1-8b</td>
    <td>50</td>
  </tr>
  <tr>
    <td>LongAttn</td>
    <td>LLaMA-3.1-70b</td>
    <td>100</td>
  </tr>
</table>

Table 5: Compared the GPU hours used by different methods on LongABC-32K-Raw, using H800 GPUs. For implementation simplicity, we used the traditional attention computation method in LongAttn. If efficient methods like Flash-attn were adopted, the speed would further improve.

6.2 The Scalability of LongAttn

Figures 3c and 3d show that LongAttn significantly improves performance when using stronger models. This indicates that more powerful models can better analyze the dependencies between long-context tokens. It can be envisioned that using LongAttn with larger models could yield even stronger performance.

However, in works like ProLong, computational efficiency is constrained by the approach, making it unfeasible to use larger models. This unique advantage of LongAttn highlights its tremendous growth potential.

6.3 The Efficiency of LongAttn

Compared to sentence-level methods like ProLong, LongAttn is significantly more efficient. ProLong divides the data into sentence segments and calculates the relative perplexity and distance between each segment, which is computationally expensive, especially for LLMs. As a result, only smaller models are used in their work. In contrast, LongAttn only requires a single inference pass to obtain relative scores between all tokens, using just the first layer of the LLM’s decoder. This approach is far more efficient and scalable.

Table 5 compares the GPU hours consumed by the two methods using models of different sizes on the LongABC-32K-Raw dataset. LongAttn, even with the traditional attention computation method, is much faster than ProLong. If more efficient methods like Flash-attention were adopted, the speed of LongAttn could be further improved.

7 Conclusion

In this paper, we introduce LongAttn, a framework evaluates long-range dependencies at the token level. LongAttn is effective as the self-attention mechanism captures relationships between all token contexts during inference. This approach to measuring long-range dependencies aligns better with the underlying operating principles of LLMs. We validate the effectiveness, scalability, and high efficiency of LongAttn through a series of comprehensive experiments. Additionally, our research contributes to the previously limited study of high-quality long-context training data. This finding
suggests promising directions for future research, and we anticipate further advancements in this domain through subsequent investigations.

Limitations

Although LongAttn has demonstrated satisfactory performance, there is still room for improvement. Specifically, we used the traditional attention map calculation method, which is inefficient. While its efficiency is satisfactory, there is still significant potential for enhancement. In future work, we hope to overcome the shortcomings, refine our method further, and advance the development of long-context capabilities in LLMs.

Ethics Statement

This work fully complies with the ACL Ethics Policy. We declare that there are no ethical issues in this paper, to the best of our knowledge.

Acknowledgement

We thank the anonymous reviewers for their helpful comments on this paper. This work was partially supported by National Natural Science Foundation of China (No. 92470205).

References

Chenxin An, Jun Zhang, Ming Zhong, Lei Li, Shansan Gong, Yao Luo, Jingjing Xu, and Lingpeng Kong. 2024. Why does the effective context length of llms fall short? arXiv preprint arXiv:2410.18745.

Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. 2023a. Qwen technical report. arXiv preprint arXiv:2309.16609.

Yushi Bai, Xin Lv, Jiajie Zhang, Yuze He, Ji Qi, Lei Hou, Jie Tang, Yuxiao Dong, and Juanzi Li. 2024. LongAlign: A recipe for long context alignment of large language models. In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 1376–1395, Miami, Florida, USA. Association for Computational Linguistics.

Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. 2023b. Longbench: A bilingual multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508.

Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901.

Longze Chen, Ziqiang Liu, Wanwei He, Yunshui Li, Run Luo, and Min Yang. 2024a. Long context is not long at all: A prospector of long-dependency data for large language models. arXiv preprint arXiv:2405.17915.

Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. 2021. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374.

Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. 2023a. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595.

Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. 2023b. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309.12307.

Zhi Chen, Qiguang Chen, Libo Qin, Qipeng Guo, Haijun Lv, Yicheng Zou, Wanxiang Che, Hang Yan, Kai Chen, and Dahua Lin. 2024b. What are the essential factors in crafting effective long context multi-hop instruction datasets? insights and best practices. ArXiv, abs/2409.01893.

Colin B. Clement, Matthew Bierbaum, Kevin P. O’Keeffe, and Alexander A. Alemi. 2019. On the use of arxiv as a dataset. Preprint, arXiv:1905.00075.

Harm de Vries. 2023. In the long (context) run. https://www.harmdevries.com/post/context-length/.

Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. 2024. Longrope: Extending llm context window beyond 2 million tokens. In Forty-first International Conference on Machine Learning.

Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783.

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2020. The pile: An 800gb dataset of diverse text for language modeling. Preprint, arXiv:2101.00027.

Shousheng Jia Haosheng Zou, Xiaowei Lv and Xiangzheng Zhang. 2024. 360-llama-factory.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.

Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. 2024. Ruler: What’s the real context size of your long-context language models? arXiv preprint arXiv:2404.06654.

Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia-Yuan Chang, Huiyuan Chen, and Xia Hu. 2024. Llm maybe longlm: Selfextend llm context window without tuning. In Forty-first International Conference on Machine Learning.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361.

Huan Yee Koh, Jiaxin Ju, Ming Liu, and Shirui Pan. 2022. An empirical survey on long document summarization: Datasets, models, and metrics. ACM computing surveys, 55(8):1–35.

Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. 2024. World model on million-length video and language with blockwise ringattention. CoRR.

Jiaheng Liu, Dawei Zhu, Zhiqi Bai, Yancheng He, Huanxuan Liao, Haoran Que, Zekun Wang, Chenchen Zhang, Ge Zhang, Jiebin Zhang, et al. 2025. A comprehensive survey on long context language modeling. arXiv preprint arXiv:2503.17407.

Kai Lv, Xiaoran Liu, Qipeng Guo, Hang Yan, Conghui He, Xipeng Qiu, and Dahua Lin. 2024. Longwanjuan: Towards systematic measurement for long text quality. arXiv preprint arXiv:2402.13583.

Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. 2018. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02789.

Charles Packer, Sarah Wooders, Kevin Lin, Vivian Fang, Shishir G Patil, Ion Stoica, and Joseph E Gonzalez. 2023. Memgpt: Towards llms as operating systems. arXiv preprint arXiv:2310.08560.

Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. 2023. The refinedweb dataset for falcon llm: Outperforming curated corpora with web data, and web data only. Preprint, arXiv:2306.01116.

Bowen Peng and Jeffrey Quesnelle. 2023. Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation. https://www.reddit.com/r/LocalLlama/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have.

Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shipppole. 2023. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071.

Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053.

Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Author, Ben Bogin, Khayathi Chandu, Jennifer Dumas, Yanai Elazar, et al. 2024. Dolma: An open corpus of three trillion tokens for language model pretraining research. arXiv preprint arXiv:2402.00159.

Together.Ai. 2023. Preparing for the era of 32k context: Early learnings and explorations, 2023a. https://www.together.ai/blog/llama-2-7b-32k.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288.

Liang Wang, Nan Yang, Xingxing Zhang, Xiaolong Huang, and Furu Wei. 2024a. Bootstrap your own context length. arXiv preprint arXiv:2412.18860.

Minzheng Wang, Longze Chen, Fu Cheng, Shengyi Liao, Xinghua Zhang, Bingli Wu, Haiyang Yu, Nan Xu, Lei Zhang, Run Luo, et al. 2024b. Leave no document behind: Benchmarking long-context llms with extended multi-doc qa. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 5627–5646.

Maurice Weber, Daniel Fu, Quentin Anthony, Yonatan Oren, Shane Adams, Anton Alexandrov, Xiaozhong Lyu, Huu Nguyen, Xiaozhe Yao, Virginia Adams, et al. 2024. Redpajama: an open dataset for training large language models. arXiv preprint arXiv:2411.12372.

Wenhao Wu, Yizhong Wang, Yao Fu, Xiang Yue, Dawei Zhu, and Sujian Li. 2024. Long context alignment with short instructions and synthesized positions. arXiv preprint arXiv:2405.03939.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2023. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453.

Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. 2023. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.16039.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283–17297.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830.

Dawei Zhu, Nan Yang, Liang Wang, Yifan Song, Wen-hao Wu, Furu Wei, and Sujian Li. 2023. Pose: Efficient context window extension of llms via positional skip-wise training. arXiv preprint arXiv:2309.10400.
A Algorithm for Pre-process

Algorithm 1 Sliding Window Sample Algorithm

1: function SLIDINGWINDOW(data, W)
2:   if len(data) < W then
3:     return ∅
4:   end if
5:   l ← 0
6:   r ← len(data)
7:   S ← ∅
8:   while r − l > 3W do
9:     S ← S ∪ {data[l : l + W]}
10:    l ← l + W
11:    S ← S ∪ {data[r − W : r]}
12:    r ← r − W
13:   end while
14:   Δ ← r − l
15:   if W < Δ ≤ 2W then
16:     S ← S ∪ {data[l : l + W], data[r − W : r]}
17:   else if 2W < Δ ≤ 3W then
18:     m ← l + ⌊(Δ − W)/2⌋
19:     S ← S ∪ {data[l : l + W], data[m : m + W], data[r − W : r]}
20:   end if
21:   return S
22: end function

The Algorithm 1 demonstrates how we perform sliding window pre-processing on the data. The length of the data processed using this method will remain consistent with the window size, and compared to the truncation method, this algorithm better preserves the completeness of the original information. Some of our code is based on the (Haosheng Zou and Zhang, 2024).

B Training Details

B.1 Training Parameters

The specific experimental parameters for continual pre-training using Megatron (Shoeybi et al., 2019) are shown in Table 6.

<table>
  <tr>
    <th rowspan="2">Params</th>
    <th colspan="3">Methods</th>
  </tr>
  <tr>
    <th>Random</th>
    <th>ProLong</th>
    <th>LongAttn</th>
  </tr>
  <tr>
    <td>learning rate(Ir)</td>
    <td>1 × 10^{-5}</td>
    <td>1 × 10^{-5}</td>
    <td>1 × 10^{-5}</td>
  </tr>
  <tr>
    <td>lr decay style</td>
    <td>cosine</td>
    <td>cosine</td>
    <td>cosine</td>
  </tr>
  <tr>
    <td>Gpus (H800)</td>
    <td>64</td>
    <td>64</td>
    <td>64</td>
  </tr>
  <tr>
    <td>mbs</td>
    <td>1</td>
    <td>1</td>
    <td>1</td>
  </tr>
  <tr>
    <td>gas</td>
    <td>8</td>
    <td>8</td>
    <td>8</td>
  </tr>
  <tr>
    <td>tp size</td>
    <td>8</td>
    <td>8</td>
    <td>8</td>
  </tr>
  <tr>
    <td>pp size</td>
    <td>1</td>
    <td>1</td>
    <td>1</td>
  </tr>
  <tr>
    <td>dropout</td>
    <td>0.1</td>
    <td>0.1</td>
    <td>0.1</td>
  </tr>
  <tr>
    <td>seq length</td>
    <td>32768</td>
    <td>32768</td>
    <td>32768</td>
  </tr>
</table>

Table 6: Parameter settings for continual pre-training by different methods based on the Megatron framework.

B.2 Training Dataset

When continuing pre-training, we use the data ratios shown as Table 7, where ArXiv, Book, and Code data refer to the data selected through different methods (random selecting, based on the ProLong (Chen et al., 2024a) framework, or based on the LongAttn framework).

<table>
  <tr>
    <th>Types</th>
    <th>length</th>
    <th>Source</th>
    <th>Ratio</th>
  </tr>
  <tr>
    <td>Wiki</td>
    <td>Short</td>
    <td>Dolma (Soldaini et al., 2024)</td>
    <td>3%</td>
  </tr>
  <tr>
    <td>Github</td>
    <td>Short</td>
    <td>Pile (Gao et al., 2020)</td>
    <td>3%</td>
  </tr>
  <tr>
    <td>Web</td>
    <td>Short</td>
    <td>Refinedweb (Penedo et al., 2023)</td>
    <td>4%</td>
  </tr>
  <tr>
    <td>ArXiv</td>
    <td>Long</td>
    <td>LongABC-Arxiv</td>
    <td>30%</td>
  </tr>
  <tr>
    <td>Book</td>
    <td>Long</td>
    <td>LongABC-Book</td>
    <td>30%</td>
  </tr>
  <tr>
    <td>Code</td>
    <td>Long</td>
    <td>LongABC-Code</td>
    <td>30%</td>
  </tr>
</table>

Table 7: The types of data and their proportions used during the continuation of pre-training. LongABC-Arxiv, LongABC-Book, and LongABC-Code refer to the types of data selected using different methods from LongABC-32K-Raw.

C Details of Continual Pre-train Dataset

As shown as Figure 8, LongABC-32K-Raw is a dataset obtained by sampling long-context data and then preprocessing it as mentioned in 3.1.

LongABC-32K-Raw serves as the data source. We filter it using different methods, including random selecting, selecting based on the ProLong framework, and selecting based on the LongAttn framework. The filtered data is then combined with quantified short-context data to form our pre-training dataset, as shown in Table 7.

<table>
  <tr>
    <th>Category</th>
    <th>Source</th>
    <th>Scale</th>
  </tr>
  <tr>
    <td>ArXiv</td>
    <td>ArXiv (Clement et al., 2019)</td>
    <td>12B Tokens</td>
  </tr>
  <tr>
    <td>Book</td>
    <td>Dolma (Soldaini et al., 2024), RedPajama (Weber et al., 2024)</td>
    <td>12B Tokens</td>
  </tr>
  <tr>
    <td>Code</td>
    <td>Dolma (Soldaini et al., 2024)</td>
    <td>12B Tokens</td>
  </tr>
</table>

Table 8: Data source of LongABC-32K-Raw and composition of its various parts.

D Baselines

Table 9 details the models and baselines for our data-scale and fixed-data method comparison experiments.
<table>
  <tr>
    <th>Comparison Method</th>
    <th>Base Model</th>
    <th>Trained Dataset</th>
    <th>Selected Method</th>
    <th>Tokens</th>
  </tr>
  <tr>
    <td rowspan="4">Data-Scale</td>
    <td rowspan="4">LLaMA-3</td>
    <td>\( \mathcal{D}_{R_x} \)</td>
    <td>Selected Randomly</td>
    <td>\( x \in \{1B, 3B, 5B, 10B, 20B\} \)</td>
  </tr>
  <tr>
    <td>\( \mathcal{D}_{P_x} \)</td>
    <td>ProLong</td>
    <td>\( x \in \{1B, 3B, 5B, 10B\} \)</td>
  </tr>
  <tr>
    <td>\( \mathcal{D}_{Ax,8B} \)</td>
    <td>LongAttn-8</td>
    <td>\( x \in \{1B, 3B, 5B, 10B\} \)</td>
  </tr>
  <tr>
    <td>\( \mathcal{D}_{Ax,70B} \)</td>
    <td>LongAttn-70</td>
    <td>\( x \in \{1B, 3B, 5B, 10B\} \)</td>
  </tr>
  <tr>
    <td rowspan="6">Fixed-Scale Method</td>
    <td rowspan="6">LLaMA-3</td>
    <td>\( \mathcal{D}_{R_x} \)</td>
    <td>Selected Randomly</td>
    <td>\( x \in \{5B, 10B, 20B\} \)</td>
  </tr>
  <tr>
    <td>\( \mathcal{D}_{P_x} \)</td>
    <td>ProLong</td>
    <td>\( x \in \{5B, 10B\} \)</td>
  </tr>
  <tr>
    <td>\( \mathcal{D}_{Ax,8B} \)</td>
    <td>LongAttn-8</td>
    <td>\( x \in \{5B, 10B\} \)</td>
  </tr>
  <tr>
    <td>\( \mathcal{D}_{Ax,70B} \)</td>
    <td>LongAttn-70</td>
    <td>\( x \in \{5B, 10B\} \)</td>
  </tr>
  <tr>
    <td>Yarn<br>(Peng et al., 2023)</td>
    <td>†</td>
    <td>†</td>
    <td>†</td>
  </tr>
  <tr>
    <td>LWM<br>(Liu et al., 2024)</td>
    <td>†</td>
    <td>†</td>
    <td>†</td>
  </tr>
  <tr>
    <td>Together<br>(Together.AI, 2023)</td>
    <td>†</td>
    <td>†</td>
    <td>†</td>
  </tr>
  <tr>
    <td>LongLORA<br>(Chen et al., 2023b)</td>
    <td>†</td>
    <td>†</td>
    <td>†</td>
  </tr>
</table>

Table 9: The experiments compared different models and baselines. Selected Method indicates the method used to filter the current training set, and Tokens represents the number of tokens used for training. † indicates the absence of a given option. ProLong, LongAttn-8, and LongAttn-70 represent the ProLong framework, LongAttn based on LLaMA-3.1-8B, and LongAttn based on LLaMA-3.1-70B, respectively.

E Distribution of \( D_{S_T} \) and \( D_{U_T} \)

The distribution \( D_{S_T} \) and \( D_{U_T} \) measured by LongAttn based on LLaMA-3.1-70B is shown in Table 10. They are distributed across different value ranges.

<table>
  <tr>
    <th rowspan="2">Statistical Indicators</th>
    <th colspan="2">Arxiv</th>
    <th colspan="2">Book</th>
    <th colspan="2">Code</th>
  </tr>
  <tr>
    <th>D\( S_T \)</th>
    <th>D\( U_T \)</th>
    <th>D\( S_T \)</th>
    <th>D\( U_T \)</th>
    <th>D\( S_T \)</th>
    <th>D\( U_T \)</th>
  </tr>
  <tr>
    <td>Min Val.</td>
    <td>0.25</td>
    <td>2.2\times10^{-7}</td>
    <td>0.21</td>
    <td>1.6\times10^{-7}</td>
    <td>0.18</td>
    <td>9.7\times10^{-8}</td>
  </tr>
  <tr>
    <td>Max Val.</td>
    <td>0.50</td>
    <td>1.8\times10^{-6}</td>
    <td>0.59</td>
    <td>4.9\times10^{-6}</td>
    <td>0.54</td>
    <td>2.4\times10^{-6}</td>
  </tr>
  <tr>
    <td>Mean</td>
    <td>0.43</td>
    <td>8.5\times10^{-7}</td>
    <td>0.40</td>
    <td>4.8\times10^{-7}</td>
    <td>0.39</td>
    <td>6.1\times10^{-7}</td>
  </tr>
</table>

Table 10: Statistical indicators of \( D_{S_T} \) and \( D_{U_T} \) after evaluating LongABC-32K-Raw using the LongAttn framework based on LLaMA-3.1-70B

F Other Experimental Results

The evaluation results on RULER for models trained with data selected from LongABC-32K-Raw using different methods are shown in Table 11. RULER includes 13 tasks, categorized into four major types: retrieval ability, multi-hop tracking ability, information aggregation ability, and question answering ability. The retrieval ability has been thoroughly evaluated earlier, so only the average score is presented here.
<table>
  <tr>
    <th rowspan="2">Method</th>
    <th rowspan="2">Tokens</th>
    <th colspan="3">Retrival</th>
    <th colspan="3">Aggregation</th>
    <th colspan="3">QA</th>
    <th rowspan="2">Avg. Score</th>
  </tr>
  <tr>
    <th>Retrival Avg.</th>
    <th>VT</th>
    <th>CWE</th>
    <th>FWE</th>
    <th>Avg.</th>
    <th>QA1</th>
    <th>QA2</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td colspan="11"><i>Trained on 5B Tokens from Different Methods</i></td>
  </tr>
  <tr>
    <td>Random</td>
    <td></td>
    <td>78.2</td>
    <td>40.6</td>
    <td>31.4</td>
    <td>66.7</td>
    <td><b>49.0</b></td>
    <td>55.2</td>
    <td>43.8</td>
    <td>49.5</td>
    <td>66.4</td>
  </tr>
  <tr>
    <td>ProLong</td>
    <td>5B</td>
    <td>81.2</td>
    <td><b>51.8</b></td>
    <td>13.0</td>
    <td>65.4</td>
    <td>39.2</td>
    <td>57.2</td>
    <td>43.4</td>
    <td><b>50.3</b></td>
    <td>67.7</td>
  </tr>
  <tr>
    <td>LongAttn-8</td>
    <td></td>
    <td>82.3</td>
    <td>50.3</td>
    <td>19.8</td>
    <td>71.0</td>
    <td>45.4</td>
    <td>53.4</td>
    <td>44.0</td>
    <td>48.7</td>
    <td>69.0</td>
  </tr>
  <tr>
    <td>LongAttn-70</td>
    <td></td>
    <td><b>85.2</b></td>
    <td>43.4</td>
    <td>16.8</td>
    <td>68.5</td>
    <td>42.7</td>
    <td>55.6</td>
    <td>43.0</td>
    <td>49.3</td>
    <td><b>69.9</b></td>
  </tr>
  <tr>
    <td colspan="11"><i>Trained on 10B Tokens from Different Methods</i></td>
  </tr>
  <tr>
    <td>Random</td>
    <td></td>
    <td>79.6</td>
    <td>48.8</td>
    <td>53.3</td>
    <td>74.2</td>
    <td><b>63.7</b></td>
    <td>55.4</td>
    <td>43.6</td>
    <td>49.5</td>
    <td>70.2</td>
  </tr>
  <tr>
    <td>ProLong</td>
    <td>10B</td>
    <td>83.4</td>
    <td>55.1</td>
    <td>19.4</td>
    <td>76.8</td>
    <td>48.1</td>
    <td>54.6</td>
    <td>44.6</td>
    <td>49.6</td>
    <td>70.6</td>
  </tr>
  <tr>
    <td>LongAttn-8</td>
    <td></td>
    <td>83.2</td>
    <td>52.1</td>
    <td>21.8</td>
    <td>77.9</td>
    <td>49.9</td>
    <td>54.6</td>
    <td>43.8</td>
    <td>49.2</td>
    <td>70.4</td>
  </tr>
  <tr>
    <td>LongAttn-70</td>
    <td></td>
    <td><b>85.3</b></td>
    <td><b>55.6</b></td>
    <td>31.9</td>
    <td>67.4</td>
    <td>49.7</td>
    <td>55.4</td>
    <td>44.0</td>
    <td><b>49.7</b></td>
    <td><b>72.1</b></td>
  </tr>
  <tr>
    <td colspan="11"><i>Trained on 20B Tokens Selected Randomly</i></td>
  </tr>
  <tr>
    <td>Random</td>
    <td>20B</td>
    <td>81.8</td>
    <td>47.4</td>
    <td>51.9</td>
    <td>87.9</td>
    <td>69.9</td>
    <td>51.9</td>
    <td>56.0</td>
    <td>46.4</td>
    <td>73.0</td>
  </tr>
</table>

Table 11: The performance of models continued pre-trained using data filtered by different methods on RULER. Random, ProLong, LongAttn-8, and LongAttn-70 represent data selected randomly, data selected using the ProLong framework, data selected by the LongAttn framework with LLaMA-3.1-8B, and data selected by the LongAttn framework with LLaMA-3.1-70B, respectively.

Reviewer paper summaries: This paper proposes LongAttn, a framework for measuring long-range dependency
through the self-attention mechanism. This paper then discusses how this framework
can be utilized to create long context training data for LLMs, and showed the effect
through a couple experiments with long context task including NiaH and its variance.

This paper introduces a framework named LongAttn, which quantifies long-range dependencies at the token level by analyzing the self-attention mechanism of LLMs to calculate dependency strength. Experiments demonstrate that the LongAttn is superior to existing long text data filtering work in terms of effect and efficiency.
Summary of strengths: - This paper is globally well written, with a clear thesis and progression from
LongAttn to LongABC-32K.
- This paper introduces a novel framework that helps in producing long text training
data. The result reported in this paper sounds promising as there might be
opportunity of extending this to a larger base of unlabeled data.

1. LongAttn is the first to measure long-range dependencies at the token level using the self-attention mechanism. This approach aligns more closely with the internal workings of LLMs and can capture long-range dependencies more precisely.
2. LongAttn can be combined with LLM to filter data. This scalability offers great potential for future advancements in long-context modeling using more powerful models.
3. By using the self-attention mechanism of LLMs, LongAttn significantly reduces computational costs and improves efficiency compared to traditional sentence-level methods like ProLong, which are computationally expensive and less scalable.
Summary of weaknesses: - 32 k window is not competitive compared to the current Long LLM works.

- This paper produced labelled data via LongAttn through LLaMA-3.1 8B and 70B, and
the data is used to train LLaMA-3. This raises a slight concern that the work could
be perceived as a fancy way of distilling knowledge from a later model to a prior
one in the same family. It might be worthwhile to also conduct experiments on LLaMA-3.1
models to compare the "teacher model" performance. In addition, conducting the
experiment across other open source LLMs might also help clearing the mist.

1. Although LongAttn performs well in long-context retrieval and benchmark tests, its performance improvements in specific long-context tasks (e.g., long-context generation, long-context question answering) may be less pronounced. This is because these tasks are influenced by other factors beyond long-range dependencies (e.g., semantic understanding, context coherence).
2. LongAttn assumes that long-range dependencies are the key factor in assessing the quality of long-context data. However, data quality may also be influenced by other factors such as coherence and consistency. Relying solely on long-range dependencies for data selection may not fully capture all high-quality long-context data.
3. Lack of more detailed analysis, such as case studies or comparison of the effects or detailed comparisons of the effects of self-attention screening using different layers.
Comments, suggestions and typos: - The authors should consider reporting some additional metrics on evaluating long
context capabilities, such as perplexity or other benchmarks that are not necessarily 
direct offspring of NiaH. 
- The authors should consider designing additional experiments that would
accommodate more than just the LLaMA family, but an additional open-source large
language model(Qwen, Mistral etc). This applies to both the LLM used in LongAttn and the LLM trained
with LongABC.

1. Consider incorporating more efficient attention calculation methods (e.g., Flash-Attention) in future implementations to further enhance the efficiency of LongAttn, enabling it to handle larger datasets.
2. In addition to long-range dependencies, introduce other metrics (e.g., coherence, consistency) to comprehensively evaluate the quality of long-context data, thereby more effectively selecting high-quality data.
3. Investigate the application of the LongAttn framework to contexts longer than 32k tokens(e.g., 64k or 128k tokens).
Overall assessment scores: 2.5

3.5