LongReD: Mitigating Short-Text Degradation of Long-Context Large Language Models via Restoration Distillation
Authors: zican Dong, Junyi Li, Jinhao Jiang, Mingyu Xu, Xin Zhao, Bingning Wang, Weipeng Chen
Accepted at: acl25

LongReD: Mitigating Short-Text Degradation of Long-Context Large Language Models via Restoration Distillation

Zican Dong1*, Junyi Li2*, Jinhao Jiang1, Mingyu Xu3,
Wayne Xin Zhao1†, Bingning Wang3, Weipeng Chen3
1 Gaoling School of Artificial Intelligence, Renmin University of China
2 Department of Computer Science, National University of Singapore
3 Baichuan Inc.
dongzican@ruc.edu.cn, junyi_cs@nus.edu.sg
batmanfly@gmail.com, daniel@baichuan-inc.com

Abstract

Large language models (LLMs) have gained extended context windows through scaling positional encodings and lightweight continual pre-training. However, this often leads to degraded performance on short-text tasks, while the reasons for this degradation remain insufficiently explored. In this work, we identify two primary factors contributing to this issue: distribution drift in hidden states and attention scores, and catastrophic forgetting during continual pre-training. To address these challenges, we propose Long Context Pre-training with Restoration Distillation (LongReD), a novel approach designed to mitigate short-text performance degradation through minimizing the distribution discrepancy between the extended and original models. Besides training on long texts, LongReD distills the hidden state of selected layers from the original model on short texts. Additionally, LongReD also introduces a short-to-long distillation, aligning the output distribution on short texts with that on long texts by leveraging skipped positional indices. Experiments on common text benchmarks demonstrate that LongReD effectively preserves the model’s short-text performance while maintaining comparable or even better capacity to handle long texts than baselines. Our code is available at https://github.com/RUCAIBox/LongReD.

1 Introduction

Large language models (LLMs) have exhibited remarkable performance across a wide range of text and multimodal tasks (Brown et al., 2020; OpenAI, 2023; Zhao et al., 2023; Touvron et al., 2023; Dubey et al., 2024; Tang et al., 2025c; Wang et al., 2025; Chen et al., 2025). However, their abilities to process long contexts are constrained by the positional encodings and the attention mechanism, which define the context window size based on the length of the pre-training data (Su et al., 2024; Press et al., 2022). When the input text exceeds this context window, the model encounters out-of-distribution positional information, resulting in significant performance degradation (Chen et al., 2023; Peng et al., 2023b; Dong et al., 2024a).

To address the challenges of long context processing, various methods have been proposed to extend the context window of LLMs (bloc97, 2023; Xiao et al., 2023; Xiong et al., 2024; Chen et al., 2023). A prominent approach involves scaling positional encodings combined with lightweight continual pre-training. Leveraging the properties of RoPE (Su et al., 2024), these methods interpolate positional indices or adjust the RoPE base to prevent out-of-distribution rotary angles beyond the original context window (Chen et al., 2023; Ding et al., 2024; Peng et al., 2023b; Xiong et al., 2024). Subsequently, lightweight continual pre-training adapts LLMs to the extended context and modified encodings, enabling context window up to 128K or even 1M tokens (Xiong et al., 2024; Fu et al., 2024; Zeng et al., 2024). However, as shown in Figure 1, these long context window extension techniques come at the cost of degraded performance on short-text tasks (Xiong et al., 2024; Dong et al.,

![Bar chart comparing performances on short-text benchmarks for Original, ABF, and PI models](page_1012_613_377_246.png)

Figure 1: Comparison of the original and long-context models via ABF (Xiong et al., 2024) or PI (Chen et al., 2023) on common short-text benchmarks.
2024b). Yet, the root causes of this performance decline and potential mitigation strategies remain under-explored.

In this study, we aim to demystify the factors contributing to the performance degradation of short-text tasks after context window extension. Our analysis identifies two critical factors: distribution drift and catastrophic forgetting. The key findings of our work include: (1) Continual pre-training seeks to recover the original model’s internal distribution, but the restoration is inherently imperfect. (2) Distribution shift in hidden states potentially leads to performance degradation. (3) During continual pre-training, the performance on short-text tasks initially improves but subsequently declines as training progresses, highlighting the presence of catastrophic forgetting. (4) Replaying short text data is effective in mitigating forgetting and improving performance stability.

Based on our observations, we propose a novel approach called Long Context Pre-training with Restoration Distillation (LongReD) to mitigate the degradation in short-text capabilities of long-context LLMs. The central idea is that the short-text capacities of the extended model can be better preserved if it accurately simulates the original distributions before extension. To achieve this, in addition to typical long-text training, we introduce a short-text distillation objective, which employs the original model as a teacher to distill its hidden states on short texts into the extended model. This training objective minimizes distribution drift and alleviates catastrophic forgetting. Moreover, we propose a short-to-long distillation training objective to bridge the gap between short-text distillation and long-text training. In this setup, the original and extended models are fed with normal positional indices and skipped positional indices, respectively. By applying the distillation on the output distributions of the last layer, the short-text capacities can be effectively transferred and integrated into long-text processing.

To the best of our knowledge, this work represents the first systematic analysis of the reasons behind the performance degradation of long-context LLMs on short-text tasks. Furthermore, we propose a general method to mitigate this short-text degradation. To assess the effectiveness of our method, we extend the context window of Llama-3-8B and Mistral-7B-v0.3 and assess their performance on both short-text and long-text tasks. Experimental results demonstrate that our method preserves the original models’ performance on short-text tasks while maintaining or even improving their long-context modeling capabilities.

2 Background

2.1 Transformer

Owing to strong capabilities and scalability, Transformer decoders have served as the backbone of most LLMs (Vaswani et al., 2017; Touvron et al., 2023; Dubey et al., 2024; Tang et al., 2025b). Given a Transformer decoder with \( L \) layers and \( N \) heads and an input sequence \( x = \{x_1, \ldots, x_T\} \) consisting with \( T \) tokens, the output of the \( l \)-th layer can be denoted as \( \mathbf{H}_l = \{\mathbf{h}_{l,1}, \ldots, \mathbf{h}_{l,T}\} \). Each Transformer layer consists of a multi-head attention (MHA) module and a feed-forward network (FFN) module with residual connections connecting them, as shown by the following formula:

\[
\tilde{\mathbf{H}}_l = \mathrm{MHA}(\mathbf{H}_{l-1}) + \mathbf{H}_{l-1},
\]
\[
\mathbf{H}_l = \mathrm{FFN}(\tilde{\mathbf{H}}_l) + \tilde{\mathbf{H}}_l.
\]

In the \( i \)-th head of the MHA module at the \( l \)-th layer, the hidden states are first projected into query, key, and value matrices, i.e., \( Q_i^l, K_i^l, \) and \( V_i^l \). Positional information is then incorporated into the query and key matrices through the RoPE with the rotation matrix \( \mathbf{R}_\theta \) (\( \theta \) is the RoPE base) (Su et al., 2024). These matrices are subsequently processed through a dot product followed by a softmax operation to compute the attention scores \( \mathbf{A}_i^l \). Finally, the value representations are weighted by the attention scores, and all attention heads are concatenated and projected to produce the attention output of the \( l \)-th layer as follows:

\[
\mathbf{A}_i^l = \mathrm{Softmax}(Q_i^l \mathbf{T} \mathbf{R}_\theta K_i^l / \sqrt{d}),
\]
\[
\mathrm{MHA}(\mathbf{H}_{l-1}) = \mathrm{Concat}(\{\mathbf{A}_i^l V_i^l\}_{i=1}^N) \mathbf{W}^O,
\]

where \( d \) is the dimension of query and key, \( \mathbf{W}^O \) is the projection matrix, and Concat is the concatenation of hidden states.

2.2 Measures of Distributional Discrepancy

After extending the context window, the parameters of LLMs undergo changes, resulting in a shift in the distribution of hidden states. To quantify such distributional discrepancy, we propose two measures, i.e., hidden state similarity and attention Kullback-Leibler (KL) divergence. Specifically, given the hidden states \( \mathbf{H}_l, \tilde{\mathbf{H}}_l \) from the original
and extended models, hidden state similarity measures the average cosine similarity between hidden states at the \( l \)-th layer across all positions. Similarly, attention KL divergence calculates the average KL divergences between the \( i \)-th attention head distributions \( \mathbf{A}_l^i, \hat{\mathbf{A}}_l^i \) from the original and extended models at the \( l \)-th layer. Higher hidden state similarity or lower attention KL divergence indicates less discrepancy between the two models. These metrics are formulated as follows:

\[
\operatorname{Sim}(\mathbf{H}_l, \hat{\mathbf{H}}_l) = \frac{1}{T} \sum_{t=1}^T \frac{\mathbf{h}_{l,t}^\top \hat{\mathbf{h}}_{l,t}}{\|\mathbf{h}_{l,t}\| \|\hat{\mathbf{h}}_{l,t}\|},
\]
(5)

\[
\operatorname{KL}(\mathbf{A}_l^i, \hat{\mathbf{A}}_l^i) = \frac{1}{T} \sum_{t=1}^T \sum_{j=1}^t \frac{a_{l,t,j}^i \log \frac{a_{l,t,j}^i}{\hat{a}_{l,t,j}^i}}{a_{l,t,j}^i}.
\]
(6)

3 Empirical Analysis

Prior work (Xiong et al., 2024; Ding et al., 2024) demonstrated that the short-text capacities will degrade after lightweight continual pre-training for extending the context window of LLMs. However, the secrets behind the performance decline have not been fully explored. In this section, through conducting empirical analysis, we attribute the performance degradation to two factors, i.e., distribution drift and catastrophic forgetting.

3.1 Distribution Drift

Imperfect Distribution Restoration. In previous work (Chen et al., 2023; Ding et al., 2024), RoPE configurations are simply modified to extend the context window, undoubtedly leading to the change of distributions. Although continual pre-training is proposed to adapt LLMs to the extended context, a natural question is whether continual pre-training can eliminate the distributional discrepancy. To verify this, we modify RoPE base to 2e7 and 1e8, and train Llama-3-8B (with context window size of 8192 tokens) on 1B tokens with the sequence length of 32K and 128K, respectively, denoted as Llama-3-8B-32K and Llama-3-8B-128K. Besides, we select Llama-3-8B-Instruct-262K (Pekelis et al., 2024), which is trained on Llama-3-8B-Instruct. To measure the distributional discrepancy, we employ these models before and after continual pre-training for inference on 1000 test samples with the length of 8192 tokens from SlimPajama (Soboleva et al., 2023). The results of hidden state similarity and attention KL divergence are shown in Table 1. As can be seen, as the RoPE base increases, the extended LLMs before continual pre-training exhibit lower hidden state similarity and higher KL divergence from the original models. After training on long-context data, the similarity increases, and the KL divergence declines to some extent. However, despite the LLM striving to restore the inner distribution resembling the original model during continual pre-training, there still exists a certain degree of distributional discrepancy.

<table>
  <tr>
    <th>Model</th>
    <th>Base</th>
    <th colspan="2">Simi ↑</th>
    <th colspan="2">KLD(\( \times 10^{-5} \)) ↓</th>
  </tr>
  <tr>
    <th></th>
    <th></th>
    <th>before</th>
    <th>after</th>
    <th>before</th>
    <th>after</th>
  </tr>
  <tr>
    <td>Llama-3-8B-32K</td>
    <td>2e7</td>
    <td>0.92</td>
    <td>0.95</td>
    <td>3.51</td>
    <td>1.51</td>
  </tr>
  <tr>
    <td>Llama-3-8B-128K</td>
    <td>1e8</td>
    <td>0.83</td>
    <td>0.94</td>
    <td>6.74</td>
    <td>1.68</td>
  </tr>
  <tr>
    <td>Llama-3-8B-Ins-262K</td>
    <td>2.8e8</td>
    <td>0.76</td>
    <td>0.91</td>
    <td>8.75</td>
    <td>2.50</td>
  </tr>
</table>

Table 1: Results of hidden state similarity and KL divergence before and after continual pre-training.

Relationship between Distribution Drift and Performance Degradation. Owing to the distribution drift of the extended model, the inner working mechanisms may also change when processing the original short text data. To answer the question of whether the distributional discrepancy contributes to the performance decline in short text, we train the Llama-3-8B model using various context window extension methods and training strategies (detailed in Appendix D). Additionally, we incorporate existing open-source long-context models, which were extended from Llama-3-8B-Instruct by Gradient AI (Pekelis et al., 2024). We compute the average hidden state similarity across all layers between these extended models and their original counterparts and assess their overall performance on the MMLU benchmark (mainly focused on short-text tasks) (Hendrycks et al., 2021a). Then, we compute the MMLU performance preservation ratio of long-context models compared to the original models w.r.t. the hidden state similarity between them, as shown in Figure 2. Generally, extended models that maintain higher similarity to their original versions preserve more short-text performance. This finding highlights that **distribution shifts in extended models are a significant factor contributing to performance degradation in short text**.

3.2 Catastrophic Forgetting

Catastrophic forgetting is always a critical problem during continual pre-training (Wu et al., 2024b). Owing to the distributional difference of short texts and long texts, adapting for long context window may unavoidably leads to a trade-off on short texts. In this section, we examine the effect of the train-
Figure 2: Relationship between MMLU performance preservation of long-context models w.r.t. the hidden states similarity.

ing steps and training data on the short-text performances to verify the forgetting phenomenons.

Effect of Training Steps. We train Llama-3-8B with the RoPE base of 2e7 on the SlimPajama dataset and obtain checkpoints with different training steps, where the batch size is 64 and the training length is 32K. Then, we evaluate the models on four short-text benchmarks, i.e., MMLU (Hendrycks et al., 2021a), HumanEval (Chen et al., 2021), TriviaQA (Joshi et al., 2017), and PIQA (Bisk et al., 2020). Figure 3 presents the performance change with different training steps. It is clear that the performance is restored at the initial steps (usually less than 32 steps) in spite of fluctuations in some datasets. However, the performance gradually drops with the increase of training steps, which demonstrates that the continual adaptation for long context window will lead to catastrophic forgetting issues on short-text tasks.

Figure 3: Results of models with different training steps.

Effect of Training Data Mixture. Beyond the training steps, we also evaluate the effect of mixtures of training data. We replace half of the long text data with short text data with a length of 8K. We train Llama-3-8B on 1B tokens and evaluate them on different benchmarks. As shown in Table 2, the length of training data plays a critical role in downstream tasks. Introducing short-text data achieves better short-text performances than only training on long-text data. This verifies that long texts are critical for forgetting problems, which can be mitigated by replaying short texts.

<table>
  <tr>
    <th>Length</th>
    <th>MMLU</th>
    <th>HUMANEVAL</th>
    <th>PIQA</th>
    <th>TriviaQA</th>
  </tr>
  <tr>
    <td>Long</td>
    <td>62.0</td>
    <td>14.02</td>
    <td>74.10</td>
    <td>70.67</td>
  </tr>
  <tr>
    <td>Long+Short</td>
    <td>62.5</td>
    <td>16.46</td>
    <td>78.24</td>
    <td>72.82</td>
  </tr>
</table>

Table 2: Results of models trained with data mixtures with different lengths.

4 Approach

Inspired by the above findings, we propose Long Context Pre-training with Restoration Distillation (LongReD), a novel context window extension framework to improve short-text capacities of long-context LLMs through decreasing distribution drift and mitigating catastrophic forgetting. Unlike only training on long texts, our strategy combines three different training objectives at each training step, i.e., long-text training, short-text distillation, and short-to-long distillation. For the three objectives, we employ three datasets \( \mathcal{D}_1, \mathcal{D}_2, \mathcal{D}_3 \) with different lengths \( T_1, T_s (< T) \), and \( T \), respectively, where \( T \) and \( T_i \) are the original and target long context window length, and \( T_s \) is the short text length. The overall architecture is displayed in Figure 4.

4.1 Long Text Training

To extend the context window of an LLM, we follow previous work to first scale positional encodings via ABF (Xiong et al., 2024) or PI (Chen et al., 2023) and directly continually training on the long text dataset. Given any long text \( x \) from the dataset \( \mathcal{D}_1 \), the model \( \Theta_e \) is trained with the language modeling objective to minimize the cross-entropy loss \( \mathcal{L}_{long} \) as follows:

\[
\mathcal{L}_{long} = - \sum_{x \in \mathcal{D}_1} \frac{1}{T_1} \sum_{t=1}^{T_1} \log \Pr(x_t | x_{<t}; \Theta_e).
\]

Through training on long texts, the model can learn to model long-term dependencies and adapt to the extended context window.
Figure 4: Overview of our proposed Long Context Pre-training with Restoration Distillation (LongReD). The method consists of three parts, i.e., long-text training, short-text distillation, and short-to-long distillation.

4.2 Short-Text Distillation

Section 3 demonstrates the negative effect of distribution drift and catastrophic forgetting on short-text capacities while only training on long texts. Though replaying short text data can partly mitigate the issue, the performance restoration is limited. Therefore, we propose short-text distillation, a knowledge distillation method to minimize the distribution discrepancy between the extended model \( \Theta_e \) and the original model \( \Theta_o \).

Distillation Layer Selection. To decrease the distribution discrepancy, we choose to distill the inner hidden states of LLMs. Yet, ensuring all the hidden states are the same as the original model is impossible and may significantly decrease the long text modeling capacities. Thus, we only select \( M (< L) \) layers as the distill layers \( \mathbf{L}_M = \{l_1, \ldots, l_M\}, \forall l_i \in \{0, \ldots, L\} \) and distill the outputs of these layers. Specifically, we propose an attention-based layer selection method where we first compute the KL divergence of attention scores between the extended and original models and then select layers with the largest KL divergences for distillation (Men et al., 2024; Dong et al., 2025). We also select the critical last layer that directly determines the output.

Short-Text Distillation Loss. Given an input sample \( x \) from the short text dataset \( \mathcal{D}_2 \), we input the sample into the two models with its original positional indices \( p = [0, \ldots, T_s - 1] \). Then, we obtain the output hidden states of the selected layers of the two models, i.e., \( \{\mathbf{H}_{l_1:p;\Theta_e}, \ldots, \mathbf{H}_{l_M:p;\Theta_e}\} \) and \( \{\mathbf{H}_{l_1:p;\Theta_o}, \ldots, \mathbf{H}_{l_M:p;\Theta_o}\} \). We compute the cosine similarity between the hidden states of the two models at the same layer and position and sum up them as the short-text distillation loss \( \mathcal{L}_{short} \), formulated as follows:

\[
\mathcal{L}_{short} = - \sum_{x \in \mathcal{D}_2} \sum_{i=1}^M \operatorname{Sim}(\mathbf{H}_{l_i:p;\Theta_e}, \mathbf{H}_{l_i:p;\Theta_o}).
\]

(8)

4.3 Short-to-Long Distillation

The two objectives of long-text training and short-text distillation separately optimize the model on long and short texts, resulting in a challenging performance trade-off between long and short texts. We argue that long and short texts share some inherent linguistic commonalities. Thus, we propose an additional short-to-long distillation to transfer the short-text capabilities to long texts, further bridging the gap between short and long texts.

Skipped Positional Indices. To distill short-text capabilities to long texts, we take the short text with the original context window size \( T \) as input but employ the skipped positional indices method, which can simulate long-text positions and capture longer-term dependencies with short texts (Zhu et al., 2024; Wu et al., 2024a). To be specific, for a sequence \( x \) drawn from \( \mathcal{D}_3 \), we first define a threshold \( T_b \) to split the positional indices into three parts, i.e., \( p_{head}, p_{mid} \), and \( p_{tail} \):

\[
p = p_{head} \cup p_{mid} \cup p_{tail},
\]
(9)
\[
p_{head} = \{0, \ldots, T_b - 1\},
\]
(10)
\[
p_{mid} = \{T_b, \ldots, T - T_b - 1\},
\]
(11)
\[
p_{tail} = \{T - T_b, \ldots, T - 1\}.
\]
(12)
Then, we modify the positional indices of the three segments. We keep \( p_{head} \) unchanged but modify the end positional index of \( p_{tail} \) to be equal to the target long text length \( T_l \). For the middle segment, we employ uniform sampling or CREAM (Wu et al., 2024a) which upsamples middle positions to determine the end positional index of \( p_{mid} \) as \( T_{me} \in \{T - T_b, \ldots, T_l - T_b - 1\} \). And the start positional index of \( p_{mid} \) is equal to \( T_{me} \) minus the length of this middle segment \( T - 2T_b \), i.e., \( T_{me} - T + 2T_b \). The modified positional indices of the input text are formulated as follows:

\[
\hat{p} = \hat{p}_{head} \cup \hat{p}_{mid} \cup \hat{p}_{tail},
\]
(13)
\[
\hat{p}_{head} = \{0, \ldots, T_b - 1\},
\]
(14)
\[
\hat{p}_{mid} = \{T_{me} - T + 2T_b, \ldots, T_{me}\},
\]
(15)
\[
\hat{p}_{tail} = \{T_l - T_b, \ldots, T_l - 1\}.
\]
(16)

Details of the skipped positional indices methods are introduced in Appendix A.

Short-to-Long Distillation Loss. After obtaining skipped positional indices \( \hat{p} \), we input these indices into the extended model \( \Theta_e \) while using the normal positional indices \( p \) for teacher model \( \Theta_o \). Then, we obtain the outputs of the \( L \)-th last layer of both models and compute the cosine similarity between them. Then, we regard the negative cosine similarity as the loss to distill the output distribution into longer positions. The training objective \( \mathcal{L}_{s2l} \) can be represented as follows:

\[
\mathcal{L}_{s2l} = - \sum_{x \in \mathcal{D}_3} \operatorname{Sim}(\mathbf{H}_{L;\hat{p},\Theta_e}, \mathbf{H}_{L;p,\Theta_o}).
\]
(17)

Notably, we only distill the output of the last layer instead of selected layers in short-text distillation to avoid the disturbance of latent positional information in the middle layers.

4.4 Joint Training Objective

Finally, we aggregate losses from these three training objectives. For each batch of training data, we sample data points from the three datasets with a fixed ratio. Then, we independently compute loss through the three different objectives on the corresponding data and sum them up to balance the long and short text capacities. The final loss \( \mathcal{L}_{final} \) can be represented as follows:

\[
\mathcal{L}_{final} = \mathcal{L}_{long} + \alpha_1 \mathcal{L}_{short} + \alpha_2 \mathcal{L}_{s2l},
\]
(18)

where \( \alpha_1 \) and \( \alpha_2 \) are hyper-parameters that control the proportion of the three parts.

5 Experiment

5.1 Experimental Settings

Pre-training Setup. In the experiments, we choose Llama-3-8B (Dubey et al., 2024) and Mistral-7B-v0.3 (Jiang et al., 2023) to evaluate the effectiveness of our methods on context window extension. We sample the long text dataset \( \mathcal{D}_1 \) and short-to-long dataset \( \mathcal{D}_3 \) from SlimPajama following the setting of Fu et al. (2024). Inspired by Gao et al. (2024), we employ higher-quality data as the short text dataset \( \mathcal{D}_2 \) with the length of 1K. Then, we trained our model on these three datasets with a token quantity ratio of 4:3:1 and a total of 1B tokens. The number of distillation layers \( M \) is set to 3 for extending Llama-3-8B to 128K context window size and 6 for other settings for better balancing the short and long text modeling performances. The details of training and parameter configurations and data mixture are displayed in Appendix B.

Evaluation Benchmarks. To thoroughly evaluate the short-text performance of long-context models, we choose 17 benchmarks covering 5 capacities, i.e., general, coding, math, reading comprehension, and commonsense question answering. We also select RULER (Hsieh et al., 2024) to assess the long text processing abilities. The evaluation datasets and details are shown in Appendix C.

Baselines. In our experiments, we choose two baseline methods for comparative analysis: (1) Long+CPT: Only continual pre-training on long text datasets \( \mathcal{D}_1 \); (2) Mix+CPT: Continually pre-training on a mixed-length dataset (\( \mathcal{D}_1 \cup \mathcal{D}_2 \cup \mathcal{D}_3 \)), keeping the training data the same as our method.

5.2 Main Results

Table 3 presents the performance comparison between our proposed methods and various baselines on both short-text and long-text tasks. Detailed results for each dataset can be found in Appendix F.

First, LongReD performs better on short-text tasks and achieves competitive long-text modeling performances compared to baseline methods. Across various scaling methods and target context window sizes, LongReD consistently outperforms most baselines on short-text benchmarks. Specifically, when extending Llama-3-8B to a 32K-token context window using ABF and LongReD-C, the short-text capabilities are retained up to 99.4% of the original model’s performance, compared to
<table>
  <tr>
    <th>Model</th>
    <th>CW</th>
    <th>PE</th>
    <th>Data</th>
    <th>Method</th>
    <th>General</th>
    <th>Coding</th>
    <th>Math</th>
    <th>RC</th>
    <th>Common</th>
    <th>Short Avg.</th>
    <th>RULER</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td rowspan="4">Llama-3-8B</td>
    <td>8K</td>
    <td>-</td>
    <td>-</td>
    <td>-</td>
    <td>68.18</td>
    <td>41.20</td>
    <td>30.12</td>
    <td>63.78</td>
    <td>72.54</td>
    <td>55.16</td>
    <td>-</td>
    <td>-</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>65.39</td>
    <td>29.32</td>
    <td>29.50</td>
    <td>60.98</td>
    <td>69.83</td>
    <td>51.00</td>
    <td>82.80</td>
    <td>56.30</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>66.04</td>
    <td>27.40</td>
    <td>29.32</td>
    <td>34.65</td>
    <td>71.88</td>
    <td>45.86</td>
    <td><b>85.04</b></td>
    <td>52.39</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td><b>67.65</b></td>
    <td><b>39.90</b></td>
    <td><b>31.52</b></td>
    <td><b>62.84</b></td>
    <td>72.35</td>
    <td><b>54.85</b></td>
    <td>84.98</td>
    <td><b>59.87</b></td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>67.05</td>
    <td>39.71</td>
    <td>30.90</td>
    <td>62.51</td>
    <td><b>72.61</b></td>
    <td>54.56</td>
    <td>84.08</td>
    <td>59.48</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Long</td>
    <td>CPT</td>
    <td>63.84</td>
    <td>27.13</td>
    <td>26.18</td>
    <td>61.95</td>
    <td>67.71</td>
    <td>49.36</td>
    <td><b>82.17</b></td>
    <td>54.83</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>64.22</td>
    <td>27.29</td>
    <td>29.16</td>
    <td>36.63</td>
    <td>71.09</td>
    <td>45.68</td>
    <td>81.63</td>
    <td>51.67</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td><b>66.45</b></td>
    <td>38.39</td>
    <td>30.78</td>
    <td><b>62.79</b></td>
    <td>71.30</td>
    <td>53.94</td>
    <td>81.19</td>
    <td>58.48</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>66.34</td>
    <td><b>38.83</b></td>
    <td><b>31.68</b></td>
    <td>62.76</td>
    <td><b>72.36</b></td>
    <td><b>54.39</b></td>
    <td>79.30</td>
    <td><b>58.54</b></td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>64.06</td>
    <td>30.04</td>
    <td>27.72</td>
    <td>60.49</td>
    <td>68.37</td>
    <td>50.14</td>
    <td><b>69.70</b></td>
    <td>53.40</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td><b>66.39</b></td>
    <td>14.03</td>
    <td>29.02</td>
    <td>35.28</td>
    <td><b>71.91</b></td>
    <td>43.33</td>
    <td>69.64</td>
    <td>47.71</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>65.99</td>
    <td><b>40.03</b></td>
    <td><b>30.26</b></td>
    <td>62.39</td>
    <td>71.50</td>
    <td><b>54.03</b></td>
    <td>64.93</td>
    <td>55.85</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>65.61</td>
    <td>39.87</td>
    <td>29.67</td>
    <td><b>62.47</b></td>
    <td>71.64</td>
    <td>53.85</td>
    <td>68.41</td>
    <td><b>56.28</b></td>
  </tr>
  <tr>
    <td>32K</td>
    <td>-</td>
    <td>-</td>
    <td>-</td>
    <td>65.36</td>
    <td>33.20</td>
    <td>26.98</td>
    <td>68.07</td>
    <td>61.86</td>
    <td>51.09</td>
    <td>-</td>
    <td>-</td>
  </tr>
  <tr>
    <td>Mistral-7B-v0.3</td>
    <td>128K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>55.29</td>
    <td>26.37</td>
    <td>15.79</td>
    <td>52.93</td>
    <td>53.07</td>
    <td>40.68</td>
    <td>44.63</td>
    <td>41.35</td>
  </tr>
  <tr>
    <td>Mistral-7B-v0.3</td>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>55.99</td>
    <td>26.36</td>
    <td>15.30</td>
    <td>60.21</td>
    <td>52.74</td>
    <td>40.66</td>
    <td>45.76</td>
    <td>41.51</td>
  </tr>
  <tr>
    <td>Mistral-7B-v0.3</td>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>61.45</td>
    <td>26.66</td>
    <td><b>19.88</b></td>
    <td>65.25</td>
    <td>58.26</td>
    <td>46.30</td>
    <td><b>58.37</b></td>
    <td>48.31</td>
  </tr>
  <tr>
    <td>Mistral-7B-v0.3</td>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td><b>62.42</b></td>
    <td><b>29.64</b></td>
    <td>19.56</td>
    <td><b>66.86</b></td>
    <td><b>59.98</b></td>
    <td><b>47.69</b></td>
    <td>53.60</td>
    <td><b>48.68</b></td>
  </tr>
</table>

Table 3: Comparison of performances of short-text and long-text benchmarks of our methods with other baselines. CW denotes the context window length, PE denotes the scaling method of RoPE, RC denotes reading comprehension, Common denotes commonsense question answering, Short Avg. denotes averaged scores on short benchmarks, and Avg. denotes averaged scores of all benchmarks. LongReD-C and LongReD-U denote our method with different skipped positional indices methods, i.e., CREAM and Uniform Sampling.

92.5% achieved by naive long-text training. Furthermore, our methods demonstrate competitive performance with Llama-3-8B and superior performance with Mistral-7B-v0.3 on the RULER benchmark for long-text processing, when compared to continual pre-training approaches. These compelling results robustly validate the efficacy of our proposed methodology.

Second, the skipped positional indices method is crucial for long-text performances. Comparing LongReD with CREAM and uniform sampling methods, we find that the short-text performances are similar while the long-text performances on RULER are largely different. When the extension ratio is relatively small (e.g., 4 times), CREAM demonstrates superior performance by effectively mitigating the lost-in-the-middle problem. However, on the length of 128K tokens, the model trained with CREAM performs largely worse than that with uniform sampling. We hypothesize that the excessive focus on the middle positions introduced by CREAM results in some positions not being adequately trained.

Finally, our method benefits from the scaling techniques of positional encodings. When scaling LLaMA-3-8B to 32K with either ABF or PI, our LongReD consistently achieves superior short-text performances. However, a comparison between models with different extension methods reveals that PI exhibits inferior performance to ABF on both short-text and long-text tasks when combined with LongReD. We hypothesize that this performance gap arises due to the significant distribution discrepancy introduced by PI, which poses greater challenges for our method to mitigate effectively.

5.3 Detailed Analysis

In this section, we conduct further detailed analysis to investigate our approach.

Ablation Study. In addition to long-text training, our method incorporates both short-text distillation and short-to-long distillation. To evaluate the individual contributions of these two components, we employ half of the training data for long-text training and the other half for short-text distillation or short-to-long distillation. Moreover, we also evaluate the effect of the hyper-parameters \( \alpha_1 \) and \( \alpha_2 \) on both short and long text tasks. The results are compared against the LongReD-C method with \( \alpha_1 = 5, \alpha_2 = 10 \), as summarized in Table 4. The results indicate that excluding short-text distillation significantly degrades the model’s performance on short-text tasks. Conversely, omitting the short-to-long distillation stage results in noticeable declines in long-text performance. In addition, the adapta-
tion of the hyper-parameters is essential. Increasing or decreasing \( \alpha_2 \) will lead to performance degradation on long-text tasks, while slightly decreasing \( \alpha_1 \) will lead to better long-text performances. When either \( \alpha_1 \) or \( \alpha_2 \) is set to an extremely large value, the model’s performance on long-context tasks drops precipitously.

<table>
  <tr>
    <th>\( \alpha_1 \)</th>
    <th>\( \alpha_2 \)</th>
    <th>General</th>
    <th>Code</th>
    <th>Math</th>
    <th>Common</th>
    <th>RC</th>
    <th>RULER</th>
  </tr>
  <tr>
    <td>5</td>
    <td>10</td>
    <td>67.65</td>
    <td>39.90</td>
    <td>31.52</td>
    <td>72.35</td>
    <td>62.84</td>
    <td>84.98</td>
  </tr>
  <tr>
    <td>-</td>
    <td>15</td>
    <td>66.54</td>
    <td>39.93</td>
    <td>31.68</td>
    <td>70.64</td>
    <td>62.47</td>
    <td>85.48</td>
  </tr>
  <tr>
    <td>5</td>
    <td>-</td>
    <td>67.39</td>
    <td>41.25</td>
    <td>30.78</td>
    <td>71.50</td>
    <td>62.25</td>
    <td>83.61</td>
  </tr>
  <tr>
    <td>5</td>
    <td>100</td>
    <td>66.96</td>
    <td>40.47</td>
    <td>33.12</td>
    <td>71.86</td>
    <td>62.47</td>
    <td>80.40</td>
  </tr>
  <tr>
    <td>5</td>
    <td>30</td>
    <td>67.00</td>
    <td>40.14</td>
    <td>31.59</td>
    <td>72.26</td>
    <td>62.29</td>
    <td>83.82</td>
  </tr>
  <tr>
    <td>5</td>
    <td>15</td>
    <td>67.20</td>
    <td>38.20</td>
    <td>31.46</td>
    <td>72.39</td>
    <td>62.55</td>
    <td>84.87</td>
  </tr>
  <tr>
    <td>5</td>
    <td>5</td>
    <td>66.91</td>
    <td>41.84</td>
    <td>32.67</td>
    <td>72.36</td>
    <td>62.42</td>
    <td>83.75</td>
  </tr>
  <tr>
    <td>100</td>
    <td>10</td>
    <td>65.33</td>
    <td>37.81</td>
    <td>29.24</td>
    <td>72.00</td>
    <td>61.85</td>
    <td>80.19</td>
  </tr>
  <tr>
    <td>2</td>
    <td>10</td>
    <td>67.02</td>
    <td>41.10</td>
    <td>31.83</td>
    <td>72.25</td>
    <td>63.04</td>
    <td>86.03</td>
  </tr>
</table>

Table 4: Ablation results on five tasks. “-” denotes the objective is not employed.

Distillation Layers. We explore the impact of distillation layers by varying the number of distilled layers and their selection strategies. In comparison to our method of selecting six layers via KL divergence (denoted as \( KL(6) \)), we design three variants: (1) All: distilling all layers; (2) Last: distilling only the last layer; and (3) Uniform(6): uniformly selecting six layers. The results are summarized in Table 5. We can observe that distilling either all the layers or the last layer leads to worse long-text performance and general capacities than selecting six layers based on KL divergence. Furthermore, uniformly selected layers consistently exhibit inferior performance to those selected via attention KL divergence across the long-text tasks for the same number of distilled layers. These findings underscore the critical role of layer selection methods and distillation granularity in the performances of long-context LLMs.

<table>
  <tr>
    <th>Layers</th>
    <th>General</th>
    <th>Code</th>
    <th>Math</th>
    <th>Common</th>
    <th>RC</th>
    <th>RULER</th>
  </tr>
  <tr>
    <td>KL(6)</td>
    <td>67.65</td>
    <td>39.90</td>
    <td>31.52</td>
    <td>72.35</td>
    <td>62.84</td>
    <td>84.98</td>
  </tr>
  <tr>
    <td>Uniform(6)</td>
    <td>66.86</td>
    <td>40.24</td>
    <td>32.19</td>
    <td>72.33</td>
    <td>60.27</td>
    <td>82.53</td>
  </tr>
  <tr>
    <td>All</td>
    <td>66.89</td>
    <td>39.63</td>
    <td>31.80</td>
    <td>72.33</td>
    <td>62.61</td>
    <td>81.47</td>
  </tr>
  <tr>
    <td>Last</td>
    <td>66.40</td>
    <td>40.49</td>
    <td>30.69</td>
    <td>72.11</td>
    <td>63.19</td>
    <td>82.24</td>
  </tr>
</table>

Table 5: Results with different distillation layers.

Distillation Length. In our experiments, we set the length of short-text distillation \( T_s \) as 1024. To explore the effect of distillation length, we set \( T_s \) as 2048 and 8192, and check the performance changes. As shown in Table 6, increasing the distillation length will harm the model’s ability to handle long texts while the short-text capacities are generally the same. Overfitting to the long distillation length will destroy the implicit positional information in hidden states, which may be the main reason for the performance drop (details are shown in Appendix E).

<table>
  <tr>
    <th>\( T_s \)</th>
    <th>General</th>
    <th>Code</th>
    <th>Math</th>
    <th>Common</th>
    <th>RC</th>
    <th>RULER</th>
  </tr>
  <tr>
    <td>1024</td>
    <td>67.65</td>
    <td>39.90</td>
    <td>31.52</td>
    <td>72.35</td>
    <td>62.84</td>
    <td>84.98</td>
  </tr>
  <tr>
    <td>2048</td>
    <td>66.80</td>
    <td>40.13</td>
    <td>31.01</td>
    <td>72.63</td>
    <td>62.32</td>
    <td>83.20</td>
  </tr>
  <tr>
    <td>8192</td>
    <td>67.43</td>
    <td>39.26</td>
    <td>31.41</td>
    <td>71.75</td>
    <td>63.04</td>
    <td>75.54</td>
  </tr>
</table>

Table 6: Results with varying distillation length \( T_s \).

Comparison with Continual Learning Methods. Beyond simple continual pre-training, we also compare LongReD with two continual learning methods that can effectively mitigate catastrophic forgetting issues: (1) Model Merging: Following common model merging methods (Hu et al., 2024b; Chen et al., 2024b), we average the checkpoints of the original model and the extended model; (2) Parameter-Efficient Tuning: Since modifying RoPE directly affects attention computation without altering the knowledge stored in MLP parameters, we only tune the parameters in the attention blocks (Chen et al., 2024a). We train Llama-3-8B on a 32K context length with these methods and evaluate their performance. As shown in Table 7, though these methods can achieve better performances than naive continual pre-training, our method consistently performs better on most short-context capacities evaluations. This may be because these methods only address catastrophic forgetting. However, our approach not only addresses catastrophic forgetting through continual learning but also leverages knowledge distillation to alleviate distribution drift.

<table>
  <tr>
    <th>Method</th>
    <th>General</th>
    <th>Code</th>
    <th>Math</th>
    <th>Common</th>
    <th>RC</th>
    <th>RULER</th>
  </tr>
  <tr>
    <td>LongReD-C</td>
    <td>67.51</td>
    <td>39.90</td>
    <td>31.52</td>
    <td>72.35</td>
    <td>62.84</td>
    <td>84.98</td>
  </tr>
  <tr>
    <td>Merging</td>
    <td>65.92</td>
    <td>35.04</td>
    <td>32.38</td>
    <td>71.61</td>
    <td>62.65</td>
    <td>83.38</td>
  </tr>
  <tr>
    <td>PET</td>
    <td>66.82</td>
    <td>39.20</td>
    <td>30.16</td>
    <td>70.78</td>
    <td>62.26</td>
    <td>85.86</td>
  </tr>
  <tr>
    <td>CPT</td>
    <td>65.39</td>
    <td>29.32</td>
    <td>29.50</td>
    <td>69.83</td>
    <td>60.98</td>
    <td>82.80</td>
  </tr>
</table>

Table 7: Comparison with continual learning methods, where Merging denotes model merging and PET denotes parameter-efficient tuning.

6 Related Work

Context Window Extension. LLMs typically have constrained context windows owing to positional
encodings and attention mechanisms (Press et al., 2022; Han et al., 2023; Dong et al., 2023). In order to meet the growing demand for processing long texts (Wang et al., 2024; Tang et al., 2025a; Peng et al., 2025), various methods have been proposed to extend the context windows of LLMs based on modification of RoPE (Su et al., 2024). Positional Interpolation down-scales positional indices to prevent out-of-distribution positions. In addition, NTK (bloc97, 2023; emozilla, 2023), ABF (Xiong et al., 2024), Yarn (Peng et al., 2023b), and LongRoPE (Ding et al., 2024) modify the base of RoPE to control the rotation angles of different dimensions. After modifying positional encodings, LLMs are lightweightly continued pre-training on long texts (Fu et al., 2024). Through this pipeline, the context windows of LLMs can be extended to 128K or even 1M tokens (Dubey et al., 2024; Zeng et al., 2024). Typically, they can achieve utilization of long texts within their context windows at the cost of degradation of general capacities on short texts (Dubey et al., 2024). Our work mainly focuses on the reasons for the degradation of short-text capacities and approaches to mitigating the decline.

Knowledge Distillation. Knowledge distillation serves as an approach to transferring knowledge and abilities of teacher models to student models (Hinton et al., 2015; Xu et al., 2024). During the era of LLMs, leveraging black-box and powerful LLMs, e.g., GPT-4 (OpenAI, 2023), to generate instructions and even fine-grained reasoning path to training smaller LLMs is a common and effective method to enhance general and domain-specific capacities (Chiang et al., 2023; Peng et al., 2023a; Mukherjee et al., 2023; Ho et al., 2023). In addition, distilling the inner distribution of white-box LLMs can be another effective way. Divergence-based distillation methods minimize the divergence of the output probability distribution of teacher and student models (Gu et al., 2024; Agarwal et al., 2024; Jiang et al., 2024). Similarity-based distillation methods align the intermediate hidden states, enabling the student models to process input like the teacher models (Liang et al., 2023; Muralidharan et al., 2024). Different from these methods of transferring knowledge of powerful models to small or compressed models, our method aims to restore and preserve the original capacities of LLMs during the context window extension stage.

7 Conclusion

In this paper, we analyzed two reasons for the performance degradation in short-text tasks after context window extension, i.e., distribution drift and catastrophic forgetting. Based on the observations, we proposed an effective training method to better preserve short-text capacities, named Long Context Pre-training with Restoration Distillation (LongReD). Besides continual pre-training on long texts, LongReD introduced two additional training objectives: short-text distillation and short-to-long distillation. Experiments demonstrated that our method could achieve better performances on common short-text benchmarks while achieving comparable long-text modeling abilities on the RULER benchmark compared with continual pre-training. In future work, we will explore how to directly integrate continual training with the distillation of the original model on long texts, rather than applying them separately to texts of different lengths.

Limitations

In this paper, we present a novel perspective on the extension of context windows and the associated decline in the general capabilities of LLMs. We also propose a method to preserve general capabilities while enhancing long-context abilities. However, a notable degradation in short-text performance after lightweight continual pretraining on only several billion tokens. Xiong et al. (2024) and Dubey et al. (2024) have shown that, with training on more than 100 billion tokens, some capabilities on short texts of models may remain largely unaffected or even improve. Furthermore, our proposed method offers a general training strategy that is compatible with different positional encoding extension techniques. A more refined extension method that minimizes disruption to the model’s distribution could further enhance the effectiveness of our approach, a topic we leave for future investigation.

Acknowledgement

This work was partially supported by National Natural Science Foundation of China under Grant No. 92470205 and 62222215, Beijing Municipal Science and Technology Project under Grant No. Z231100010323009, and Beijing Natural Science Foundation under Grant No. L233008. Xin Zhao is the corresponding author.
References

Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. 2024. On-policy distillation of language models: Learning from self-generated mistakes. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net.

Jacob Austin, Augustus Odena, Maxwell I. Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie J. Cai, Michael Terry, Quoc V. Le, and Charles Sutton. 2021. Program synthesis with large language models. CoRR, abs/2108.07732.

Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen Marcus McAleer, Albert Q. Jiang, Jia Deng, Stella Biderman, and Sean Welleck. 2024. Llemma: An open language model for mathematics. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net.

Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. 2020. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pages 7432–7439. AAAI Press.

bloc97. 2023. NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.

Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Pondé de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. 2021. Evaluating large language models trained on code. CoRR, abs/2107.03374.

Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. 2023. Extending context window of large language models via positional interpolation. CoRR, abs/2306.15595.

Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. 2024a. Longlora: Efficient fine-tuning of long-context large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net.

Zhipeng Chen, Yingqian Min, Beichen Zhang, Jie Chen, Jinhao Jiang, Daixuan Cheng, Wayne Xin Zhao, Zheng Liu, Xu Miao, Yang Lu, Lei Fang, Zhongyuan Wang, and Ji-Rong Wen. 2025. An empirical study on eliciting and improving r1-like reasoning models. CoRR, abs/2503.04548.

Zhipeng Chen, Liang Song, Kun Zhou, Wayne Xin Zhao, Bingning Wang, Weipeng Chen, and Ji-Rong Wen. 2024b. Extracting and transferring abilities for building multi-lingual ability-enhanced large language models. CoRR, abs/2410.07825.

Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing. 2023. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality.

Eunsol Choi, He He, Mohit Iyyer, Mark Yatskar, Wentau Yih, Yejin Choi, Percy Liang, and Luke Zettlemoyer. 2018. Quac: Question answering in context. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 - November 4, 2018, pages 2174–2184. Association for Computational Linguistics.

Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. 2019. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pages 2924–2936. Association for Computational Linguistics.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind
Tafjord. 2018. Think you have solved question answering? try arc, the AI2 reasoning challenge. CoRR, abs/1803.05457.

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems. CoRR, abs/2110.14168.

OpenCompass Contributors. 2023. Opencompass: A universal evaluation platform for foundation models. https://github.com/open-compass/opencompass.

Yiran Ding, Li Lyna Zhang, Chengrudong 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, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net.

Zican Dong, Junyi Li, Xin Men, Wayne Xin Zhao, Bingbing Wang, Zhen Tian, Weipeng Chen, and Ji-Rong Wen. 2024a. Exploring context window of large language models via decomposed positional vectors. CoRR, abs/2405.18009.

Zican Dong, Han Peng, Peiyu Liu, Wayne Xin Zhao, Dong Wu, Feng Xiao, and Zhifeng Wang. 2025. Domain-specific pruning of large mixture-of-experts models with few-shot demonstrations. CoRR, abs/2504.06792.

Zican Dong, Tianyi Tang, Junyi Li, and Wayne Xin Zhao. 2023. A survey on long text modeling with transformers. CoRR, abs/2302.14502.

Zican Dong, Tianyi Tang, Junyi Li, Wayne Xin Zhao, and Ji-Rong Wen. 2024b. BAMBOO: A comprehensive benchmark for evaluating long text modeling capacities of large language models. In Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evaluation, LREC/COLING 2024, 20-25 May, 2024, Torino, Italy, pages 2086–2099. ELRA and ICCL.

Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. 2019. DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pages 2368–2378. Association for Computational Linguistics.

Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravanakumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Alonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esibou, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Grégoire Milon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel M. Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasudenh Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, and et al. 2024. The llama 3 herd of models. CoRR, abs/2407.21783.

emozilla. 2023. Dynamically Scaled RoPE further increases performance of long context LLaMA with zero fine-tuning.

Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. 2024. Data engineering for scaling language models to 128k context. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net.

Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. 2024. How to train long-context language models (effectively). CoRR, abs/2410.02660.

Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2024. Minilm: Knowledge distillation of large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net.

Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Simon Wang. 2023. Lm-infinite: Simple on-the-fly length generalization for large language models. CoRR, abs/2308.16137.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2021a. Measuring massive multitask language understanding. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.

Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021b. Measuring mathematical
problem solving with the MATH dataset. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1, NeurIPS Datasets and Benchmarks 2021, December 2021, virtual.

Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015. Distilling the knowledge in a neural network. CoRR, abs/1503.02531.

Namgyu Ho, Laura Schmid, and Se-Young Yun. 2023. Large language models are reasoning teachers. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023, Toronto, Canada, July 9-14, 2023, pages 14852–14882. Association for Computational Linguistics.

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? CoRR, abs/2404.06654.

Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zhen Leng Thai, Kai Zhang, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. 2024a. Minicpm: Unveiling the potential of small language models with scalable training strategies. CoRR, abs/2404.06395.

Zhiyuan Hu, Yuliang Liu, Jinman Zhao, Suyuchen Wang, Yan Wang, Wei Shen, Qing Gu, Anh Tuan Luu, See-Kiong Ng, Zhiwei Jiang, et al. 2024b. Longrecipe: Recipe for efficient long context generalization in large language models. CoRR, abs/2409.00509.

Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressard, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavauud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2023. Mistral 7b. CoRR, abs/2310.06825.

Jinhao Jiang, Junyi Li, Wayne Xin Zhao, Yang Song, Tao Zhang, and Ji-Rong Wen. 2024. Mix-cpt: A domain adaptation framework via decoupling knowledge learning and format alignment. CoRR, abs/2407.10804.

Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. 2017. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics, ACL 2017, Vancouver, Canada, July 30 - August 4, Volume 1: Long Papers, pages 1601–1611. Association for Computational Linguistics.

Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, and Tuo Zhao. 2023. Less is more: Task-aware layer-wise distillation for language model compression. In International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 20852–20867. PMLR.

Hao Liu, Matei Zaharia, and Pieter Abbeel. 2023. Ring attention with blockwise transformers for near-infinite context. CoRR, abs/2310.01889.

Ilya Loshchilov and Frank Hutter. 2019. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net.

Anton Lozhkov, Raymond Li, Loubna Ben Allal, Federico Cassano, Joel Lamy-Poirier, Nouamane Tazi, Ao Tang, Dmytro Pykhart, Jiawei Liu, Yuxiang Wei, Tianyang Liu, Max Tian, Denis Kocetkov, Arthur Zucker, Younes Belkada, Zijian Wang, Qian Liu, Dmitry Abulkhanov, Indraneil Paul, Zhuang Li, Wending Li, Megan Risdal, Jia Li, Jian Zhu, Terry Yue Zhuo, Evgenii Zheltonozhskii, Nii Osae Osae Dade, Wenhao Yu, Lucas Krauß, Naman Jain, Yixuan Su, Xuanli He, Manan Dey, Edorado Abati, Yekun Chai, Niklas Muenninghoff, Xiangru Tang, Muhtasham Obhokulov, Christopher Akiki, Marc Marone, Cheng-hao Mou, Mayank Mishra, Alex Gu, Binyuan Hui, Tri Dao, Armel Zebaze, Olivier Dehaene, Nicolas Patry, Canwen Xu, Julian J. McAuley, Han Hu, Torsten Scholak, Sébastien Paquet, Jennifer Robinson, Carolyn Jane Anderson, Nicolas Chapados, and et al. 2024. Starcoder 2 and the stack v2: The next generation. CoRR, abs/2402.19173.

Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. 2024. Shortgpt: Layers in large language models are more redundant than you expect. CoRR, abs/2403.03853.

Todor Mihailev, Peter Clark, Tushar Khot, and Ashish Sabharwal. 2018. Can a suit of armor conduct electricity? A new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 - November 4, 2018, pages 2381–2391. Association for Computational Linguistics.

Subhabrata Mukherjee, Arindam Mitra, Ganesh Jawahar, Sahaj Agarwal, Hamid Palangi, and Ahmed Awadallah. 2023. Orca: Progressive learning from complex explanation traces of GPT-4. CoRR, abs/2306.02707.

Saurav Muradilaharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochoowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. 2024. Compact language models via pruning and knowledge distillation. CoRR, abs/2407.14679.
OpenAI. 2023. GPT-4 technical report. CoRR, abs/2303.08774.

Denis Paperno, Germán Kruszewski, Angeliki Lazari-dou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, ACL 2016, August 7-12, 2016, Berlin, Germany, Volume I: Long Papers. The Association for Computer Linguistics.

Leonid Pekelis, Michael Feil, Forrest Moret, Mark Huang, and Tiffany Peng. 2024. Llama 3 gradient: A series of long context models.

Guilherme Penedo, Hynek Kydlíček, Loubna Ben Al-lal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro von Werra, and Thomas Wolf. 2024. The fineweb datasets: Decanting the web for the finest text data at scale. CoRR, abs/2406.17557.

Baolin Peng, Chunyuan Li, Pengcheng He, Michel Galley, and Jianfeng Gao. 2023a. Instruction tuning with GPT-4. CoRR, abs/2304.03277.

Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. 2023b. Yarn: Efficient context window extension of large language models. CoRR, abs/2309.00071.

Han Peng, Jinhao Jiang, Zican Dong, Wayne Xin Zhao, and Lei Fang. 2025. Cafe: Retrieval head-based coarse-to-fine information seeking to enhance multi-document qa capability. arXiv preprint arXiv:2505.10063.

Ofir Press, Noah A. Smith, and Mike Lewis. 2022. Train short, test long: Attention with linear biases enables input length extrapolation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net.

Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. 2020. Compressive transformers for long-range sequence modelling. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016. Squad: 100, 000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016, pages 2383–2392. The Association for Computational Linguistics.

Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin Choi. 2019. Social iqa: Commonsense reasoning about social interactions. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019, pages 4462–4472. Association for Computational Linguistics.

Daria Soboleva, Faisal Al-Khatteeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. 2023. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama.

Jianlin Su, Murtadha H. M. Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. 2024. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063.

Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V. Le, Ed H. Chi, Denny Zhou, and Jason Wei. 2023. Challenging big-bench tasks and whether chain-of-thought can solve them. In Findings of the Association for Computational Linguistics: ACL 2023, Toronto, Canada, July 9-14, 2023, pages 13003–13051. Association for Computational Linguistics.

Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. 2019. Commonsenseqa: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pages 4149–4158. Association for Computational Linguistics.

Tianyi Tang, Yiwen Hu, Bingqian Li, Wenyang Luo, Zijing Qin, Haoxiang Sun, Jiapeng Wang, Shiyi Xu, Xiaoxue Cheng, Geyang Guo, Han Peng, Bowen Zheng, Yiru Tang, Yingqian Min, Yushuo Chen, Jie Chen, Yuanqian Zhao, Luran Ding, Yuhao Wang, Zican Dong, Chunxuan Xia, Junyi Li, Kun Zhou, Wayne Xin Zhao, and Ji-Rong Wen. 2024. Llmbox: A comprehensive library for large language models. CoRR, abs/2407.05563.

Xinyu Tang, Xiaolei Wang, Zhihao Lv, Yingqian Min, Wayne Xin Zhao, Binbin Hu, Ziqi Liu, and Zhiqiang Zhang. 2025a. Unlocking general long chain-of-thought reasoning capabilities of large language models via representation engineering. CoRR, abs/2503.11314.

Xinyu Tang, Xiaolei Wang, Wayne Xin Zhao, Siyuan Lu, Yaliang Li, and Ji-Rong Wen. 2025b. Unleashing the potential of large language models as prompt optimizers: Analogical analysis with gradient-based model optimizers. In AAAI-25, Sponsored by the Association for the Advancement of Artificial Intelligence, February 25 - March 4, 2025, Philadelphia, PA, USA, pages 25264–25272. AAAI Press.

Xinyu Tang, Xiaolei Wang, Xin Zhao, and Ji-Rong Wen. 2025c. DAWN-ICL: strategic planning of problem-solving trajectories for zero-shot in-context learning. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for
Computational Linguistics: Human Language Technologies, NAACL 2025 - Volume 1: Long Papers, Albuquerque, New Mexico, USA, April 29 - May 4, 2025, pages 1918–1934. Association for Computational Linguistics.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenjin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yunjing Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Runqta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. 2023. Llama 2: Open foundation and fine-tuned chat models. CoRR, abs/2307.09288.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems, 30.

Yuhao Wang, Ruixiang Ren, Junyi Li, Xin Zhao, Jing Liu, and Ji-Rong Wen. 2024. REAR: A relevance-aware retrieval-augmented framework for open-domain question answering. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, EMNLP 2024, Miami, FL, USA, November 12-16, 2024, pages 5613–5626. Association for Computational Linguistics.

Yuhao Wang, Ruixiang Ren, Yucheng Wang, Wayne Xin Zhao, Jing Liu, Hua Wu, and Haifeng Wang. 2025. Unveiling knowledge utilization mechanisms in llm-based retrieval-augmented generation. CoRR, abs/2505.11995.

Tong Wu, Yanpeng Zhao, and Zilong Zheng. 2024a. An efficient recipe for long context extension via middle-focused positional encoding. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.

Tongtong Wu, Linhao Luo, Yuan-Fang Li, Shirui Pan, Thuy-Trang Vu, and Gholamreza Haffari. 2024b. Continual learning for large language models: A survey. CoRR, abs/2402.01364.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2023. Efficient streaming language models with attention sinks. CoRR, abs/2309.17453.

Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajwal Bhargava, Rui Hou, Louis Martin, Rashi Runqta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. 2024. Effective long-context scaling of foundation models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), NAACL 2024, Mexico City, Mexico, June 16-21, 2024, pages 4643–4663. Association for Computational Linguistics.

Xiaohan Xu, Ming Li, Chongyang Tao, Tao Shen, Reynold Cheng, Jinyang Li, Can Xu, Dacheng Tao, and Tianyi Zhou. 2024. A survey on knowledge distillation of large language models. CoRR, abs/2402.13116.

Aohan Zeng, Bin Xu, Bowen Wang, Chenhui Zhang, Da Yin, Diego Rojas, Guanyu Feng, Hanlin Zhao, Hanyu Lai, Hao Yu, Hongning Wang, Jiadai Sun, Jiajie Zhang, Jiale Cheng, Jiayi Gui, Jie Tang, Jing Zhang, Juanzi Li, Lei Zhao, Lindong Wu, Lucen Zhong, Mingdao Liu, Minlie Huang, Peng Zhang, Qinkai Zheng, Rui Lu, Shuaiqi Duan, Shudan Zhang, Shulin Cao, Shuxun Yang, Weng Lam Tam, Wenyi Zhao, Xiao Liu, Xiao Xia, Xiaohan Zhang, Xiaotao Gu, Xin Lv, Xinghan Liu, Xinyi Liu, Xinyue Yang, Xixuan Song, Xunkai Zhang, Yifan An, Yifan Xu, Yilin Niu, Yuantao Yang, Yueyan Li, Yushi Bai, Yuxiao Dong, Zehan Qi, Zhaoyu Wang, Zhen Yang, Zhengxiao Du, Zhenyu Hou, and Zihan Wang. 2024. Chatglm: A family of large language models from GLM-130B to GLM-4 all tools. CoRR, abs/2406.12793.

Peiyuan Zhang, Kaichen Zhang, Bo Li, Guangtao Zeng, Jingkang Yang, Yuanhan Zhang, Ziyue Wang, Haoran Tan, Chunyuan Li, and Ziwei Liu. 2024. Long context transfer from language to vision. CoRR, abs/2406.16852.

Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, Yifan Du, Chen Yang, Yushuo Chen, Zhipeng Chen, Jinhao Jiang, Ruixiang Ren, Yifan Li, Xinyu Tang, Zikang Liu, Peiyu Liu, Jian-Yun Nie, and Ji-Rong Wen. 2023. A survey of large language models. CoRR, abs/2303.18223.

Dawei Zhu, Nan Yang, Liang Wang, Yifan Song, Wen-hao Wu, Furu Wei, and Sujian Li. 2024. Pose: Efficient context window extension of llms via positional skip-wise training. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net.
A  Skipped Positional Indices

Skipped positional indices are widely employed for effective long-context training. The method simulates long-distance dependency via modifying the positional indices. In this section, we introduce two skipped positional indices employed in our methods.

A.1 CREAM

Continuity-Relativity indExing with gAussian Middle (CREAM) (Wu et al., 2024a) is a method that modifies position indices to simulate long positions. The main idea is to make the model better focused on the middle part of the long sequences.

As described in Section 4.3, the positional indices are first split into three non-overlapped parts, i.e., \( p_{head}, p_{tail}, \) and \( p_{tail} \), where the length of the head and the tail are \( T_b \). There are two segmentation methods and we randomly select one method to determine the number of \( T_b \). To achieve the continuity of the middle part, \( T_b \) is set as a small number, where we set it to four times the context window extension factor. To learn more relative positions, \( T_b \) is set as one-third of the input length \( T/3 \).

To better concentrate on the middle positions, the method proposes a truncated Gaussian function to determine the end position of the middle part \( p_e \). Initially, the method defines a Gaussian distribution \( f(x) \) and its cumulative distribution function \( F(x) \):

\[
f(x) = \frac{1}{\sigma \sqrt{2\pi}} \exp(-\frac{(x - \mu)^2}{2\sigma^2}), \tag{19}
\]
\[
F(x) = \int_{-\infty}^{x} f(t) dt, \tag{20}
\]

where \( \mu = 1 + T_l/T \) and \( \sigma = 3 \). Then, the method samples 1000 points \( x_i \) uniformly from 1 to extension factor \( T_l/T \) and calculates their cumulative distribution function \( F(x_i) \). After that, the method samples \( F(u) \) from the Gaussian distribution and uses interpolation and inverse transformation to obtain the corresponding \( u \). The number \( u \) then rounds to the nearest integer \( \alpha \).

\[
u = x_{i-1} + \frac{(x_i - x_{i-1})(F(u) - F(x_{i-1}))}{F(x_i) - F(x_{i-1})}, \tag{21}
\]
\[
\alpha = \text{round}(u), \tag{22}
\]

where \( F(x_i) \) and \( F(x_{i-1}) \) are the two nearest numbers to \( F(u) \) and \( F(u) \) is sampled uniformly from 0 to 1. Finally, we sample an integer as the end of the middle part:

\[
T_{me} \sim \text{Uniform}(T_b + \alpha(T_l - 2T_b), \alpha T - T_b - 1), \tag{23}
\]
\[
T_{mb} = T_{me} + 2T_b - T. \tag{24}
\]

A.2 Uniform Sampling

In the uniform sampling method, we keep the same as the CREAM method except the sampling method of the positions of the middle part. Instead, we employ a uniform sampling method to determine the end positions of the middle part, as formulated as follows:

\[
T_{me} \sim \text{Uniform}(T - T_b, T_l - T_b - 1), \tag{25}
\]
\[
T_{mb} = T_{me} + 2T_b - T, \tag{26}
\]

where Uniform denotes uniformly sampling from this range.

B  Training Details

B.1 Training Datasets

In our proposed method, we employ three datasets: \( \mathcal{D}_1, \mathcal{D}_2, \) and \( \mathcal{D}_3 \), which are specifically designed for long-text training, short-text distillation, and short-to-long distillation, respectively. Following the sampling strategy of Fu et al. (2024), we sample data in the same proportion as Llama-2 from the SlimPajama datasets (Soboleva et al., 2023), tokenize the samples, and pack them to the target length. To enhance the preservation of short-text capabilities during the short-text distillation phase, we adopt the approach proposed by (Gao et al., 2024), utilizing higher-quality datasets. Specifically, we select Finweb-Edu (Penedo et al., 2024), Stack-v2 (Lozhkov et al., 2024), the mathematical subset of Finweb-Edu (Penedo et al., 2024), Proof-Pile-2 (Azerbayev et al., 2024), PG19 (Rae et al., 2020), Arxiv (Soboleva et al., 2023), and Wikipedia (Soboleva et al., 2023). These datasets are carefully chosen to balance quality and diversity. Additionally, we incorporate instruction-tuned datasets that have been demonstrated to be beneficial for model performance (Hu et al., 2024a), including Tulu-v2-sft-mixture\(^1\), MathInstruct\(^2\), WizardLM-evol-instruct-

\(^1\)https://huggingface.co/datasets/allenai/tulu-v2-sft-mixture
\(^2\)https://huggingface.co/datasets/TIGER-Lab/MathInstruct
V2\textsuperscript{3}, and Magicoder-Evol-Instruct-110K\textsuperscript{4}. The detailed proportions of these datasets utilized in our method are summarized in Table 8.

<table>
  <tr>
    <th>Dataset</th>
    <th>Data Mixture</th>
  </tr>
  <tr>
    <td>\( \mathcal{D}_1 \)</td>
    <td>Wikipedia: 0.034, CommonCrawl: 0.534, StackExchange: 0.032, C4: 0.266, Github: 0.050, ArXiv: 0.043, Book: 0.041</td>
  </tr>
  <tr>
    <td>\( \mathcal{D}_3 \)</td>
    <td>Wikipedia: 0.034, CommonCrawl: 0.534, StackExchange: 0.032, C4: 0.266, Github: 0.050, ArXiv: 0.043, Book: 0.041</td>
  </tr>
  <tr>
    <td>\( \mathcal{D}_2 \)</td>
    <td>Fineweb-Edu: 0.476, Fineweb-Edu-Math: 0.722, Proof-Pile-2: 0.230, Stack-v2: 0.095, PG19: 0.095, Wikipedia: 0.096, ArXiv: 0.095, Instructions: 0.047</td>
  </tr>
</table>

Table 8: Data Mixture of three datasets.

B.2 Evaluated Models

We select Llama-3-8B and Mistral-7B-v0.3 as our evaluated models. For Llama-3-8B, we extend its context window to 32K and 128K, with the ABF and PI methods. For Mistral-7B-v0.3, we extend its context window to 128K with ABF method. The configuration and training length of these models are displayed in Table 9.

<table>
  <tr>
    <th>Model</th>
    <th>CW</th>
    <th>PE</th>
    <th>Ratio</th>
    <th>\( L_{\mathcal{D}_1} \)</th>
    <th>\( L_{\mathcal{D}_2} \)</th>
    <th>\( L_{\mathcal{D}_3} \)</th>
  </tr>
  <tr>
    <td>Llama-3-8B</td>
    <td>32K</td>
    <td>ABF</td>
    <td>2e7</td>
    <td>32000</td>
    <td>1024</td>
    <td>8192</td>
  </tr>
  <tr>
    <td>Llama-3-8B</td>
    <td>32K</td>
    <td>PI</td>
    <td>4</td>
    <td>32000</td>
    <td>1024</td>
    <td>8192</td>
  </tr>
  <tr>
    <td>Llama-3-8B</td>
    <td>128K</td>
    <td>ABF</td>
    <td>1e8</td>
    <td>128000</td>
    <td>1024</td>
    <td>8192</td>
  </tr>
  <tr>
    <td>Mistral-7B-v0.3</td>
    <td>128K</td>
    <td>ABF</td>
    <td>2e7</td>
    <td>128000</td>
    <td>1024</td>
    <td>32768</td>
  </tr>
</table>

Table 9: Configuration of models and length of datasets. CW denotes the context window of target models, PE denotes the scaling methods of RoPE, Ratio denotes the RoPE theta for ABF methods, and the interpolation ratio for PI methods. \( L_{\mathcal{D}_1} \), \( L_{\mathcal{D}_2} \), and \( L_{\mathcal{D}_3} \) denotes the length of the three datasets.

B.3 Training Configurations

We employ ring flash attention (Liu et al., 2023) in EasyContext framework (Zhang et al., 2024) (Apache-2.0 License) to train our models with 8 A800 GPUs. The learning rate is fixed as 2e-5 without warmup and the training tokens of each batch are about 2M tokens. We employ AdamW optimizer (Loshchilov and Hutter, 2019) with the weight decay of 0.1, \( \beta_1 \) of 0.9, and \( \beta_2 \) of 0.95. The hyper-parameters \( \alpha_1 \) and \( \alpha_2 \) are set as \( \alpha_1 = 5, \alpha_2 = 10 \) for 32K, and \( \alpha_1 = 2, \alpha_2 = 15 \) for 128K. All the models are trained with 512 steps with about 1B tokens.

B.4 Training Costs

We report the training costs of our method and other baselines. All the models are trained on 8 A800 GPUs with ring attention. The time costs for training these models are shown in Table 10. Compared with baselines, our methods will cost about 10% extra computation than the only continually pre-training on the same data. Additionally, comparing only training on long texts, LongReD costs similar time in 32K while saving about 20% computations in 128K settings.

<table>
  <tr>
    <th>Model</th>
    <th>CW</th>
    <th>Data</th>
    <th>Method</th>
    <th>Time(h)</th>
  </tr>
  <tr>
    <td rowspan="4">Llama-3-8B</td>
    <td>32K</td>
    <td>Long</td>
    <td>CPT</td>
    <td>22.4</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>19.9</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>Mix</td>
    <td>LongReD</td>
    <td>22.4</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>Long</td>
    <td>CPT</td>
    <td>42.8</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>30.4</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>Mix</td>
    <td>LongReD</td>
    <td>33.4</td>
  </tr>
</table>

Table 10: Training time of Llama-3-8B with different methods and target context windows.

C Evaluation Details

To evaluate the performance of LLMs on short-text and long-text tasks, we employ multiple benchmarks. In this section, we introduce the datasets for evaluating different capacities:

• *General*: MMLU (Hendrycks et al., 2021a), BBH (Suzgun et al., 2023), LAMBADA (Paperno et al., 2016)

• *Math*: MATH (Hendrycks et al., 2021b), GSM8K (Cobbe et al., 2021)

• *Coding*: HumanEval (Chen et al., 2021), MBPP (Austin et al., 2021)

• *Reading comprehension*: Squadv2 (Rajpurkar et al., 2016), Quac (Choi et al., 2018), TriviaQA (Joshi et al., 2017), Drop (Dua et al., 2019), BoolQ (Clark et al., 2019)

• *Commonsense question answering*: Openbookqa (Mihaylov et al., 2018), Commonsenseqa (Talmor et al., 2019), ARC-C (Clark et al., 2018), SIQA (Sap et al., 2019), PIQA (Bisk et al., 2020)),
• Long text processing: RULER (Hsieh et al., 2024)

we also present details of our evaluation configurations, as shown in Table 11.

<table>
  <tr>
    <th>Dataset</th>
    <th>CoT</th>
    <th>#Shots</th>
    <th>Metric</th>
  </tr>
  <tr>
    <td>MMLU</td>
    <td>×</td>
    <td>5</td>
    <td>Probability</td>
  </tr>
  <tr>
    <td>BBH</td>
    <td>✓</td>
    <td>3</td>
    <td>EM</td>
  </tr>
  <tr>
    <td>Lambada</td>
    <td>×</td>
    <td>0</td>
    <td>accuracy</td>
  </tr>
  <tr>
    <td>MBPP</td>
    <td>×</td>
    <td>3</td>
    <td>pass@1</td>
  </tr>
  <tr>
    <td>HumanEval</td>
    <td>×</td>
    <td>0</td>
    <td>pass@1</td>
  </tr>
  <tr>
    <td>GSM8K</td>
    <td>✓</td>
    <td>4</td>
    <td>accuracy</td>
  </tr>
  <tr>
    <td>MATH</td>
    <td>✓</td>
    <td>4</td>
    <td>accuracy</td>
  </tr>
  <tr>
    <td>PIQA</td>
    <td>×</td>
    <td>0</td>
    <td>Probability</td>
  </tr>
  <tr>
    <td>SIQA</td>
    <td>×</td>
    <td>0</td>
    <td>Probability</td>
  </tr>
  <tr>
    <td>OpenBookQA</td>
    <td>×</td>
    <td>0</td>
    <td>Probability</td>
  </tr>
  <tr>
    <td>ARC(C)</td>
    <td>×</td>
    <td>25</td>
    <td>Probability</td>
  </tr>
  <tr>
    <td>CommonSenseQA</td>
    <td>×</td>
    <td>7</td>
    <td>PPL</td>
  </tr>
  <tr>
    <td>Quac</td>
    <td>×</td>
    <td>1</td>
    <td>F1</td>
  </tr>
  <tr>
    <td>TriviaQA</td>
    <td>×</td>
    <td>5</td>
    <td>F1</td>
  </tr>
  <tr>
    <td>DROP</td>
    <td>×</td>
    <td>3</td>
    <td>F1</td>
  </tr>
  <tr>
    <td>BoolQ</td>
    <td>×</td>
    <td>3</td>
    <td>F1</td>
  </tr>
  <tr>
    <td>Squad-v2</td>
    <td>×</td>
    <td>1</td>
    <td>F1</td>
  </tr>
  <tr>
    <td>RULER</td>
    <td>×</td>
    <td>0</td>
    <td>Accuracy</td>
  </tr>
</table>

Table 11: Configurations of evaluated benchmarks. CoT denotes employing the Chain-of-Thought prompt, #Shot denotes the number of shots in prompts, and Metric denotes the evaluation metric for the benchmark.

The details of metrics employed to evaluate the performance are described as follows:

• Probability: For a choice task, all the choices are provided in the prompt and the Probability of each choice is calculated. The choice with the maximum Probability will be compared with the correct choice and the accuracy is reported.

• PPL: For a choice task, the choices are not provided. Instead, the content of each choice is input after the question, and the average perplexity of each choice is calculated. The choice with the lowest PPL is selected as the correct answer.

• Pass@1: For coding tasks, the model is asked to give one solution. Then, we evaluate the solution with the unit testing and count the numbers that pass the testing.

• EM: For the question-answering task, the exact matching metric measures whether the output is the same as the golden answer.

• F1: For the question-answering task, the F1 first splits the golden answer and output of the model into words, and calculates the F1 score of these word lists.

• Accuracy: Accuracy measures the number of correctly answered samples out of all the samples. For RULER, the sample is correct if the answer is in the output or each part of the answer (only for qa1 and qa2).

We evaluate these models using the LLM-Box (Tang et al., 2024) library for most benchmarks, and OpenCompass (Contributors, 2023) for the math-related benchmarks.

D Models for Evaluating Distribution Drift

To evaluate the relationship between the distribution drift and the short performance, we first employ open-sourced three long-context models from Gradient AI (Pekelis et al., 2024), which are continually pre-trained from Llama-3-8B-Instruct with only several hundred million tokens. Then, we train Llama-3-8B with ABF and PI methods with different extension ratios on 256M tokens with a length of 32K. We also follow Pekelis et al. (2024), gradually extending the context window length to twice its original length. For each extension, we train the model with 256M tokens and enlarge the RoPE theta ten times. The full model list and their MMLU scores and hidden states similarity can be seen in Table 12. Typically, a larger extension ratio will cause a greater shift in the model’s distribution, further leading to a decline in the model’s performance on short texts. We provide theoretical proof of the impact of the RoPE bases on the drifts of attention distributions.

D.1 Proof of RoPE Base Impact on Attention Scores

Given a query \( \mathbf{q}^m \) and a key \( \mathbf{k}^n \), their inner product with RoPE encoding can be expressed as:

\[
(\mathbf{R}_{\Theta,m}^d \mathbf{q}^m)^{\top} (\mathbf{R}_{\Theta,n}^d \mathbf{k}^n)
\]
\[
= \mathrm{Re} \left[ \sum_{i=0}^{d/2-1} \mathbf{q}^m_{[2i:2i+1]} \mathbf{k}^n_{[2i:2i+1]} * e^{i(m-n)\theta_i} \right]
\]

where \( [2i : 2i + 1] \) denotes the \( i \)-th subspace. Following the definition from Su et al. (2024), we let \( h_i = \mathbf{q}^m_{[2i:2i+1]} \mathbf{k}^n_{[2i:2i+1]} * \) and \( S_j = \sum_{k=0}^{j-1} e^{i(m-n)\theta_k} \). We extend these definitions with
<table>
  <tr>
    <th>Model</th>
    <th>RoPE</th>
    <th>CW</th>
    <th>Base Model</th>
    <th>Tokens</th>
    <th>MMLU</th>
    <th>MMLU Ratio</th>
    <th>Simi</th>
  </tr>
  <tr>
    <th colspan="8">GradientAI Models</th>
  </tr>
  <tr>
    <td>Llama-3-8B-Ins-G-262k</td>
    <td>ABF(2.8e8)</td>
    <td>262K</td>
    <td>-</td>
    <td>-</td>
    <td>61.9</td>
    <td>0.945</td>
    <td>0.920</td>
  </tr>
  <tr>
    <td>Llama-3-8B-Ins-G-1048k</td>
    <td>ABF(3.6e9)</td>
    <td>1048K</td>
    <td>-</td>
    <td>-</td>
    <td>59.7</td>
    <td>0.911</td>
    <td>0.900</td>
  </tr>
  <tr>
    <td>Llama-3-8B-Ins-G-4194k</td>
    <td>ABF(4.5e10)</td>
    <td>4194K</td>
    <td>-</td>
    <td>-</td>
    <td>56.1</td>
    <td>0.856</td>
    <td>0.872</td>
  </tr>
  <tr>
    <th colspan="8">Our Models (Directly Extension)</th>
  </tr>
  <tr>
    <td>Llama-3-32K-5e6</td>
    <td>ABF(5e6)</td>
    <td>32K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>62.2</td>
    <td>0.960</td>
    <td>0.943</td>
  </tr>
  <tr>
    <td>Llama-3-32K-5e7</td>
    <td>ABF(5e7)</td>
    <td>32K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>62.0</td>
    <td>0.957</td>
    <td>0.934</td>
  </tr>
  <tr>
    <td>Llama-3-32K-5e8</td>
    <td>ABF(5e8)</td>
    <td>32K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>60.9</td>
    <td>0.940</td>
    <td>0.924</td>
  </tr>
  <tr>
    <td>Llama-3-32K-5e9</td>
    <td>ABF(5e9)</td>
    <td>32K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>60.4</td>
    <td>0.932</td>
    <td>0.915</td>
  </tr>
  <tr>
    <td>Llama-3-32K-5e10</td>
    <td>ABF(5e10)</td>
    <td>32K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>59.8</td>
    <td>0.923</td>
    <td>0.908</td>
  </tr>
  <tr>
    <td>Llama-3-32K-5e20</td>
    <td>ABF(5e20)</td>
    <td>32K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>56.5</td>
    <td>0.872</td>
    <td>0.856</td>
  </tr>
  <tr>
    <td>Llama-3-32K-P14</td>
    <td>PI(4)</td>
    <td>32K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>60.9</td>
    <td>0.940</td>
    <td>0.917</td>
  </tr>
  <tr>
    <td>Llama-3-32K-P116</td>
    <td>PI(16)</td>
    <td>32K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>56.1</td>
    <td>0.866</td>
    <td>0.862</td>
  </tr>
  <tr>
    <th colspan="8">Our Models (Gradually Extension)</th>
  </tr>
  <tr>
    <td>Llama-3-16K-5e6</td>
    <td>ABF(5e6)</td>
    <td>16K</td>
    <td>Llama-3-8B</td>
    <td>256M</td>
    <td>62.3</td>
    <td>0.961</td>
    <td>0.946</td>
  </tr>
  <tr>
    <td>Llama-3-32K-5e7</td>
    <td>ABF(5e7)</td>
    <td>32K</td>
    <td>Llama-3-16K-5e6</td>
    <td>256M</td>
    <td>60.7</td>
    <td>0.938</td>
    <td>0.927</td>
  </tr>
  <tr>
    <td>Llama-3-64K-5e8</td>
    <td>ABF(5e8)</td>
    <td>64K</td>
    <td>Llama-3-32K-5e7</td>
    <td>256M</td>
    <td>60.3</td>
    <td>0.930</td>
    <td>0.912</td>
  </tr>
  <tr>
    <td>Llama-3-128K-5e9</td>
    <td>ABF(5e9)</td>
    <td>128K</td>
    <td>Llama-3-64K-5e8</td>
    <td>256M</td>
    <td>57.2</td>
    <td>0.883</td>
    <td>0.901</td>
  </tr>
</table>

Table 12: Details of evaluated models for the analysis of the relationship between distribution drift and short-text performances.

\( h_{d/2} = 0 \) and \( S_0^{\Theta} = 0 \). Applying Abel transformation, we rewrite the summation:

\[
\sum_{i=0}^{d/2-1} h_i e^{i(m-n)\theta_i} = \sum_{i=0}^{d/2-1} h_i (S_{i+1}^{\Theta} - S_i^{\Theta}) \tag{29}
\]
\[
= - \sum_{i=0}^{d/2-1} S_{i+1}^{\Theta} (h_{i+1} - h_i) \tag{30}
\]

**Theorem.** Let \( \Theta = b^{-2i/d} \) and \( \Theta' = b'^{-2i/d} \) be two RoPE bases with \( b > b' > 1 \). For fixed \( q^m \) and \( k^n \), their attention score difference satisfies:
\[
|Attn(b) - Attn(b')| \leq C \cdot \sum_{i=0}^{d/2-1} \left( |S_{i+1}^{\Theta}| + |S_{i+1}^{\Theta'}| \right)
\]
where \( C = \max_{0 \leq i < d/2} |h_{i+1} - h_i| \).

**Proof.** The attention scores under different bases are:
\[
Attn(b) = \operatorname{Re} \left[ - \sum_{i=0}^{d/2-1} S_{i+1}^{\Theta} (h_{i+1} - h_i) \right], \tag{31}
\]
\[
Attn(b') = \operatorname{Re} \left[ - \sum_{i=0}^{d/2-1} S_{i+1}^{\Theta'} (h_{i+1} - h_i) \right] \tag{32}
\]

Their difference satisfies:
\[
|Attn(b) - Attn(b')|
\leq \left| \sum_{i=0}^{d/2-1} \left( S_{i+1}^{\Theta} - S_{i+1}^{\Theta'} \right) (h_{i+1} - h_i) \right| \tag{33}
\]
\[
\leq \sum_{i=0}^{d/2-1} |S_{i+1}^{\Theta} - S_{i+1}^{\Theta'}| \cdot |h_{i+1} - h_i| \tag{35}
\]
\[
\leq C \sum_{i=0}^{d/2-1} \left( |S_{i+1}^{\Theta}| + |S_{i+1}^{\Theta'}| \right) \tag{36}
\]
where the last inequality follows from \( |S_{i+1}^{\Theta} - S_{i+1}^{\Theta'}| \leq |S_{i+1}^{\Theta}| + |S_{i+1}^{\Theta'}| \) by triangle inequality.

To analyze the expected difference across positions, consider:
\[
\mathbb{E} \left[ |Attn(b) - Attn(b')| \right] \tag{37}
\]
\[
\leq \frac{1}{N} \sum_{q,k,t} C^{q,k} \sum_{i=0}^{d/2-1} \left( |S_{i+1}^{\Theta,t}| + |S_{i+1}^{\Theta',t}| \right) \tag{38}
\]
\[
\leq \frac{C}{N} \sum_{t=0}^T (T-t) \sum_{i=0}^{d/2-1} \left( |S_{i+1}^{\Theta,t}| + |S_{i+1}^{\Theta',t}| \right) \tag{39}
\]
where \( T \) is the maximum sequence length and \( t = |m - n| \) is the relative position. Let \( B(\Theta) = \sum_{t=0}^T (T-t) \sum_{i=0}^{d/2-1} |S_{i+1}^{\Theta,t}| \). Previous work has demonstrated that RoPE has long-term decay and
<table>
  <tr>
    <th>Base</th>
    <th>5e5</th>
    <th>5e6</th>
    <th>5e7</th>
    <th>5e8</th>
    <th>5e9</th>
    <th>5e10</th>
    <th>5e11</th>
    <th>5e12</th>
    <th>5e13</th>
  </tr>
  <tr>
    <th>\( B(\Theta) \)</th>
    <td>1.3e6</td>
    <td>1.5e6</td>
    <td>1.9e6</td>
    <td>2.3e6</td>
    <td>3.0e6</td>
    <td>3.3e6</td>
    <td>3.6e6</td>
    <td>3.9e6</td>
    <td>4.1e6</td>
  </tr>
</table>

Table 13: Upper bound of different RoPE bases.

a large base will lead to slow attention decay. For bases \( b_3 > b_2 > b_1 \), we have \( B(\Theta_3) > B(\Theta_2) > B(\Theta_1) \) due to slower attention decay with larger bases. As shown in Table 13, empirical measurements on Llama-3-8B confirm that \( B(\Theta) \) increases with \( b \), proving that larger bases induce greater attention score changes during context window extension.

E Analysis of Distillation Length

As discussed in Section 5.3, increasing the distillation length will harm the long text modeling capacities. To explore the reason for the performance decline, we employ positional vectors (Dong et al., 2024a), a method to extract latent positional information from the hidden states. Specifically, given the hidden states from different samples, we average the hidden states at the same position to obtain vectors that representing positional information, denoted as positional vectors:

\[
p_{l,i} = \frac{1}{N} \sum_{s=1}^{N} h_{l,s}^i
\]

Then, we compute the cosine similarities of positional vectors with different places. All the positional vectors are calculated on samples from SlimPajama (Soboleva et al., 2023) with a length of 32000 tokens. The similarity matrices of the positional vectors of models with distillation lengths of 1024, 2048, and 8192 as well as the models trained with only long texts are shown in Figure 5. We can observe that a large distillation length will cause a significant discontinuity in the implicit positional information inside and outside the original context window.

![Simialrity matrices of positional vectors inside and outside the context window.](page_1012_1012_388_388.png)

Figure 5: Simialrity matrices of positional vectors inside and outside the context window.

F Results Details

In this section, we display the details of performances on these benchmarks. The results of benchmarks evaluating general, coding, and math capacities are shown in Table 14. Table 15 and Table 16 show the performance on commonsense question answering and reading comprehension benchmarks respectively.
<table>
  <tr>
    <th>Model</th>
    <th>CW</th>
    <th>PE</th>
    <th>Data</th>
    <th>Method</th>
    <th>MMLU</th>
    <th>BBH</th>
    <th>LAMBADA</th>
    <th>HUMANEVAL</th>
    <th>MBPP</th>
    <th>GSM8K</th>
    <th>MATH</th>
  </tr>
  <tr>
    <td rowspan="4">Llama-3-8B</td>
    <td>8K</td>
    <td>-</td>
    <td>-</td>
    <td>-</td>
    <td>64.80</td>
    <td>63.97</td>
    <td>75.78</td>
    <td>34.75</td>
    <td>48.28</td>
    <td>50.04</td>
    <td>10.20</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>62.00</td>
    <td>59.90</td>
    <td>74.27</td>
    <td>14.02</td>
    <td>44.62</td>
    <td>49.20</td>
    <td>9.80</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>62.60</td>
    <td>60.54</td>
    <td>74.97</td>
    <td>16.46</td>
    <td>38.34</td>
    <td>49.43</td>
    <td>9.20</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>64.40</td>
    <td>62.78</td>
    <td>75.35</td>
    <td>35.37</td>
    <td>44.42</td>
    <td>52.84</td>
    <td>10.20</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>64.40</td>
    <td>61.75</td>
    <td>75.02</td>
    <td>34.76</td>
    <td>44.66</td>
    <td>52.69</td>
    <td>9.20</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Long</td>
    <td>CPT</td>
    <td>61.40</td>
    <td>56.20</td>
    <td>73.92</td>
    <td>12.20</td>
    <td>42.06</td>
    <td>44.96</td>
    <td>7.40</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>60.50</td>
    <td>57.17</td>
    <td>74.99</td>
    <td>16.46</td>
    <td>38.12</td>
    <td>49.13</td>
    <td>9.20</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>63.30</td>
    <td>60.83</td>
    <td>75.22</td>
    <td>34.76</td>
    <td>42.02</td>
    <td>51.55</td>
    <td>10.00</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>61.40</td>
    <td>56.22</td>
    <td>74.56</td>
    <td>17.07</td>
    <td>43.00</td>
    <td>47.23</td>
    <td>8.20</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>62.40</td>
    <td>61.62</td>
    <td>75.16</td>
    <td>21.95</td>
    <td>6.12</td>
    <td>50.04</td>
    <td>8.00</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>63.60</td>
    <td>60.25</td>
    <td>74.11</td>
    <td>35.37</td>
    <td>44.68</td>
    <td>52.92</td>
    <td>7.60</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>63.60</td>
    <td>59.45</td>
    <td>73.78</td>
    <td>35.96</td>
    <td>43.78</td>
    <td>49.13</td>
    <td>10.20</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>-</td>
    <td>-</td>
    <td>-</td>
    <td>62.30</td>
    <td>58.49</td>
    <td>75.28</td>
    <td>24.39</td>
    <td>42.00</td>
    <td>45.56</td>
    <td>8.40</td>
  </tr>
  <tr>
    <td rowspan="4">Mistral-7B-v0.3</td>
    <td>128K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>51.90</td>
    <td>43.3</td>
    <td>70.68</td>
    <td>21.34</td>
    <td>31.4</td>
    <td>26.38</td>
    <td>5.20</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>54.70</td>
    <td>44.87</td>
    <td>68.39</td>
    <td>19.51</td>
    <td>33.2</td>
    <td>27.60</td>
    <td>3.00</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>58.90</td>
    <td>51.58</td>
    <td>73.86</td>
    <td>20.12</td>
    <td>33.2</td>
    <td>33.36</td>
    <td>6.40</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>60.1</td>
    <td>52.08</td>
    <td>75.1</td>
    <td>23.78</td>
    <td>35.5</td>
    <td>34.12</td>
    <td>5.00</td>
  </tr>
</table>

Table 14: Comparison of performances of short-text and long-text benchmarks of our methods with other baselines. CW denotes the context window length, PE denotes the scaling method of RoPE, RC denotes reading comprehension, and Common denotes commonsense question answering.

<table>
  <tr>
    <th>Model</th>
    <th>CW</th>
    <th>PE</th>
    <th>Data</th>
    <th>Method</th>
    <th>CommonsenseQA</th>
    <th>OpenBookQA</th>
    <th>PIQA</th>
    <th>SIQA</th>
    <th>ARC-C</th>
  </tr>
  <tr>
    <td rowspan="4">Llama-3-8B</td>
    <td>8K</td>
    <td>-</td>
    <td>-</td>
    <td>-</td>
    <td>73.22</td>
    <td>71</td>
    <td>77.97</td>
    <td>61</td>
    <td>79.52</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>71.50</td>
    <td>66.40</td>
    <td>74.1</td>
    <td>59.42</td>
    <td>77.73</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>72.65</td>
    <td>69.60</td>
    <td>77.09</td>
    <td>62.74</td>
    <td>77.30</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>72.56</td>
    <td>71.80</td>
    <td>77.80</td>
    <td>60.34</td>
    <td>79.27</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>72.89</td>
    <td>71.80</td>
    <td>77.91</td>
    <td>61.46</td>
    <td>79.01</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Long</td>
    <td>CPT</td>
    <td>71.25</td>
    <td>62.80</td>
    <td>72.03</td>
    <td>55.58</td>
    <td>76.88</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>70.19</td>
    <td>69.00</td>
    <td>77.04</td>
    <td>62.28</td>
    <td>76.96</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>71.91</td>
    <td>70.80</td>
    <td>75.90</td>
    <td>59.98</td>
    <td>77.90</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>65.19</td>
    <td>64.00</td>
    <td>77.15</td>
    <td>59.83</td>
    <td>75.68</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>72.65</td>
    <td>70.00</td>
    <td>77.53</td>
    <td>61.62</td>
    <td>77.73</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>72.40</td>
    <td>69.20</td>
    <td>76.01</td>
    <td>61.16</td>
    <td>78.75</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>72.40</td>
    <td>69.80</td>
    <td>77.04</td>
    <td>61.00</td>
    <td>77.99</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>-</td>
    <td>-</td>
    <td>-</td>
    <td>71.50</td>
    <td>66.8</td>
    <td>66.65</td>
    <td>56.4</td>
    <td>79.01</td>
  </tr>
  <tr>
    <td rowspan="4">Mistral-7B-v0.3</td>
    <td>128K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>66.99</td>
    <td>34.8</td>
    <td>56.8</td>
    <td>38.89</td>
    <td>67.15</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>66.99</td>
    <td>34.8</td>
    <td>56.8</td>
    <td>38.89</td>
    <td>67.15</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>69.04</td>
    <td>65.4</td>
    <td>62.68</td>
    <td>53.89</td>
    <td>75.26</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>70.84</td>
    <td>65.6</td>
    <td>63.87</td>
    <td>55.48</td>
    <td>78.5</td>
  </tr>
</table>

Table 15: Comparison of performances of short-text and long-text benchmarks of our methods with other baselines. CW denotes the context window length, PE denotes the scaling method of RoPE, RC denotes reading comprehension, and Common denotes commonsense question answering.
<table>
  <tr>
    <th>Model</th>
    <th>CW</th>
    <th>PE</th>
    <th>Data</th>
    <th>Method</th>
    <th>SquadV2</th>
    <th>quacQ</th>
    <th>TriviaQA</th>
    <th>BoolQ</th>
    <th>DROP</th>
  </tr>
  <tr>
    <td rowspan="5">Llama-3-8B</td>
    <td>8K</td>
    <td>-</td>
    <td>-</td>
    <td>-</td>
    <td>72.06</td>
    <td>35.65</td>
    <td>75.42</td>
    <td>82.6</td>
    <td>53.18</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>66.56</td>
    <td>34.03</td>
    <td>70.67</td>
    <td>83.12</td>
    <td>50.52</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>0.89</td>
    <td>16.97</td>
    <td>72.60</td>
    <td>82.60</td>
    <td>0.17</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>74.18</td>
    <td>34.39</td>
    <td>74.24</td>
    <td>82.29</td>
    <td>49.10</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>72.45</td>
    <td>34.50</td>
    <td>74.16</td>
    <td>82.26</td>
    <td>49.18</td>
  </tr>
  <tr>
    <td rowspan="3">Llama-3-8B</td>
    <td>32K</td>
    <td>PI</td>
    <td>Long</td>
    <td>CPT</td>
    <td>73.26</td>
    <td>33.09</td>
    <td>70.95</td>
    <td>81.16</td>
    <td>51.29</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>7.53</td>
    <td>21.43</td>
    <td>72.55</td>
    <td>81.44</td>
    <td>0.21</td>
  </tr>
  <tr>
    <td>32K</td>
    <td>PI</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>71.96</td>
    <td>34.73</td>
    <td>74.87</td>
    <td>82.94</td>
    <td>49.47</td>
  </tr>
  <tr>
    <td rowspan="4">Mistral-7B-v0.3</td>
    <td>128K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>66.01</td>
    <td>32.81</td>
    <td>73.13</td>
    <td>81.16</td>
    <td>49.35</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>1.08</td>
    <td>18.41</td>
    <td>74.06</td>
    <td>82.66</td>
    <td>0.21</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>70.67</td>
    <td>33.97</td>
    <td>75.19</td>
    <td>82.51</td>
    <td>49.63</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>70.26</td>
    <td>34.85</td>
    <td>75.13</td>
    <td>82.84</td>
    <td>49.29</td>
  </tr>
  <tr>
    <td rowspan="5">Mistral-7B-v0.3</td>
    <td>32K</td>
    <td>-</td>
    <td>-</td>
    <td>-</td>
    <td>72.47</td>
    <td>31.84</td>
    <td>76.1</td>
    <td>83.43</td>
    <td>45.46</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Long</td>
    <td>CPT</td>
    <td>55.15</td>
    <td>29.07</td>
    <td>62.9</td>
    <td>80.83</td>
    <td>37.4</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>CPT</td>
    <td>53.92</td>
    <td>31.88</td>
    <td>63.85</td>
    <td>78.32</td>
    <td>35.73</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-C</td>
    <td>70.45</td>
    <td>30.16</td>
    <td>68.16</td>
    <td>80.61</td>
    <td>41.9</td>
  </tr>
  <tr>
    <td>128K</td>
    <td>ABF</td>
    <td>Mix</td>
    <td>LongReD-U</td>
    <td>72.6</td>
    <td>30.82</td>
    <td>71.84</td>
    <td>81.16</td>
    <td>43.5</td>
  </tr>
</table>

Table 16: Comparison of performances of short-text and long-text benchmarks of our methods with other baselines. CW denotes the context window length, PE denotes the scaling method of RoPE, RC denotes reading comprehension, and Common denotes commonsense question answering.

Reviewer paper summaries: This paper identifies two primary factors contributing to the performance degradation of short-text tasks after context window extension: distribution drift and catastrophic forgetting. Then, this paper proposes a Long Context Pre-training with Restoration Distillation (LongReD) to mitigate the degradation in short-text capabilities of long-context LLMs. The LongReD includes long-text training, short-text distillation, and short-to-long distillation. Experimental results show its effectiveness.

The paper introduces LongReD, a method designed to mitigate the degradation of short-text performance in long-context LLMs. It identifies distribution drift and catastrophic forgetting as key causes of performance issues. LongReD addresses this by using restoration distillation to align hidden states and output distributions between the original model and the long-context model. Experiments show that LongReD successfully preserves short-text performance while maintaining long-text processing capability.

The paper examines the impact of scaling context lengths on the performance of short texts, focusing specifically on the LLaMA-3-8B model. It investigates how extending the context length of rotary positional embeddings (RoPE) from 8K to 32K/128K affects model performance. The authors identify that the performance degradation on short texts caused by context size scaling has not been adequately explored.

The study attributes this issue to two key factors: "distribution drift" and "catastrophic forgetting." To address these challenges, the authors propose two novel loss functions designed for use during long-context continual training, incorporating knowledge distillation techniques to preserve short-text performance.

The proposed method, named LongReD, is evaluated on two positional embedding scaling methods: RoPE-PI and RoPE-ABF. Experimental results show that LongReD successfully improve the short-text performance and the performance is much closer to original model while introducing only minor performance degradation on long-context benchmarks.
Summary of strengths: 1. The paper identifies two primary factors contributing to the performance degradation of short-text tasks.
2. LongReD seems to make sense.
3. The paper is well-written.

- The paper tackles the short-text performance degradation in long-context models, an important and underexplored challenge. This issue is critical for real-world applications.
- This paper provides a comprehensive analysis to identify the key reasons for the short-text performance degradation, which is enlightening for further research.
- The paper provides thorough experimental details, including training costs, making it easier for other researchers to reproduce the results and evaluate the approach.

- **Clear Explanation of Methodology:**
The paper clearly explains the proposed method and problem, supported by an informative figure that effectively illustrates how the two novel loss functions operate.

- **Comprehensive Experimental Evaluation:**
The experiments are comprehensive, evaluating the approach on two different methods of context scaling and various short-text evaluation benchmarks, demonstrating the effectiveness of the proposed solution.

- **Robust Ablation Study:**
The ablation study is thorough, addressing concerns about the loss multiplier hyperparameters $\alpha_1$ and $\alpha_2$, which adds to the robustness of the findings.

- **Well-Structured Appendix:**
The appendix is well-structured and provides valuable supplementary information, including details about datasets, training costs, and additional experimental results, enhancing the paper's transparency and reproducibility.
Summary of weaknesses: 1. It would be better if the theoretical explanation of why RoPE base affects the distribution could be provided.

2. Why were the ablation experiments conducted on these specific datasets? This raises doubts about the effectiveness of the components presented in the paper.

3. The performance improvement in this paper doesn't seem as significant as it appears, because the average in Table 3 includes multiple short-text datasets.

4. If it can include results from more LLMs, it would make the article more solid.

- While the LongReD method successfully improves short-text performance, the results show an acceptable but universal degradation in long-context performance. Ideally, the method should improve short-text performance without compromising the model’s ability to handle long-context tasks. It would be beneficial if the approach could achieve both short- and long-context performance improvements simultaneously or at least maintain the long-context performance.
- The paper identifies catastrophic forgetting as a key challenge and uses knowledge distillation to address it. However, knowledge distillation is a widely used method in continual learning, and the paper would benefit from a more in-depth discussion of continual learning techniques. A comparison with other continual learning methods, rather than just data mixture, could provide a more comprehensive view of how LongReD fits within the broader context of continual learning and knowledge retention.

- **Evaluation Limited to a Single Model:**
The proposed method is evaluated only on a single model (LLaMA-3-8B). Expanding the analysis to include additional models would strengthen the generalizability and credibility of the paper’s claims.

- **Limited Scope of Problem Applicability:**
As acknowledged by the authors in the limitations section, the issue of short-context degradation may not arise if the model is exposed to a sufficient number of tokens during long-context continual training. This reduces the broader applicability of the proposed solution.
Comments, suggestions and typos: 1. Why is LongReD-U missing in the 32K setting?

2. Typos: In line 315, forgetting issues on short-text tasks..

- As mentioned in the weaknesses part, it would be better to discuss and compare with other continual learning methods.
- I am curious about the performance with extreme $\alpha$ settings in Table 4.

Beyond the points mentioned in the weaknesses section, I do not have any additional suggestions for improvement.
Overall assessment scores: 3.5

3.5

4
Metareview score: 3
Metareview: This paper introduces LongReD, a technique designed to mitigate performance degradation in short-text tasks within long-context large language models (LLMs). The authors identify two key factors—distribution drift and catastrophic forgetting—as contributors to this issue. To address these, LongReD employs restoration distillation to align hidden states and output distributions between the original and long-context models. The paper is well-motivated, clearly presented, and supported by robust experiments and ablation studies. Overall, the rebuttal period seems to have addressed the reviewers' initial concerns with the paper; reviewers all vote to accept the paper. Combined with the strengths of the paper noted above, I recommend that the paper be accepted.
Summary of reasons to publish: 1. The paper is well-motivated and addresses an important problem of short-text degradation in long-context LLMs. 

2. The proposed approach effectively balances short- and long-context performance, demonstrating strong results across benchmarks. 

3. Comprehensive analysis and extensive experiments are conducted to evaluate that LongReD preserves its short-text performance and extends to long context.   

4. The author experimented on more different LLMs to examine the generalizability of their proposed method.