DCIS: Efficient Length Extrapolation of LLMs via Divide-and-Conquer Scaling Factor Search

Frequency Scaling. Despite RoPE’s effectiveness in incorporating positional information, the model’s capacity to handle sequences exceeding pre-training lengths remains limited, primarily due to inadequate training of low-frequency dimensions LocalLLaMA (2023) within the conventional context window size $L$. To address this, several scaling methods have been proposed to adjust RoPE’s rotation frequency. We denote $\mathbf{R}^{d}_{\mathbf{\Theta},m}$ succinctly as:

where $\beta_{i}=10000^{2i/d}$ represents the base frequency, and $\lambda_{i}$ are the scaling factors for each frequency dimension.

Scaling Factors Methods. Both NTK-aware approach LocalLLaMA (2023) and YaRN Peng et al. (2023) apply theories from the Neural Tangent Kernel (NTK) and suggest varying the scaling factors $\lambda_{i}$ based on the dimension’s training need, achieving better performance with less fine-tuning data. LongRoPE Ding et al. (2024), on the other hand, utilizes a search-based method to identify better scaling factors $\lambda_{i}$, employing an evolutionary search algorithm for initial short text search (128k/256k) and fine-tuning, and search again at the target length (2M). This method facilitates up to a significant increase in processing length compared to the baseline, highlighting the potential of scaling adjustments in extending model capacities.

Searching Scaling Factors during inference. Recent research Wu et al. (2024) reveal that specific retrieval heads contain long-form textual information, suggesting the intrinsic capacity of LLMs to process extended texts. To exploit this potential within LLMs, our framework first involves searching for optimal scaling factors during the preliminary inference phase. More specifically, our approach encourages the model to autonomously select suitable scaling factors through low-cost inference, thereby enhancing its extrapolation capability with low-cost training.

Fine-Tuning with Searched Factors. Directly applying the searched scaling factors to the original LLMs leads to suboptimal performance, indicating a lack of sufficient adaptation within the LLMs Ding et al. (2024). Following YaRN, our framework further involves a fine-tuning phase with searched scaling factors. Due to our advanced search algorithm (Section 4.2), better scaling factors are searched and initialized, thereby further reducing the number of fine-tuning steps.

Here, we introduce our Divide-and-Conquer Incremental Search (DCIS) algorithm, which integrates the divide-and-conquer strategy to efficiently approximate better scaling factors. Beyond Figure 2, Algorithm 1 presents the detailed procedure of DCIS. We also demonstrate that the search space of DCIS is half that of conventional search methods, in Section 6.1.

Algorithmic Strategy. The scaling factors sequence, designated as $[\lambda_{0},\lambda_{1},\cdots,\lambda_{d/2-1}]$ (as formulated in Eq. 4), is methodically processed using a divide-and-conquer strategy. Specifically, we divide the sequence into a set of segments, and focus on one segment (as highlighted in red in Figure 2) at a time while maintaining the others constant. For each segment, we adopt an incremental search strategy, where a new sequence of scaling factors is generated by adding a predetermined value from a set range to the current sequence. The most effective increment is then selected based on the lowest perplexity (PPL), which is also used to narrow the range for subsequent searches in the divided sub-segments. This iterative process, shown during the first iteration for $d=16$ with three values evaluated at each segment in Figure 2, incrementally refines the scaling factors.

Optimizations Incorporated. Based on empirical insights, several optimizations have been incorporated into the DCIS algorithm:

1. 1.

   Initial Scaling Factors: Drawing from the success of YaRN Peng et al. (2023), its scaling factors are used as starting points from which our search begins.

2. 2.

   Priority to High-Dimensional Scaling Factors: Inspired by the NTK-aware approach LocalLLaMA (2023), which suggests that higher dimensions might require more extensive interpolation, our algorithm prioritizes these dimensions for updates, thereby speeding up the convergence to better scaling factors.

3. 3.

   Discarding Non-Guiding Increments: Increments resulting in a PPL greater than 100 are considered ineffective and are thus excluded from the search to maintain focus on potentially successful modifications.

4. 4.

   Avoiding Local Optima: To prevent falling into local optima, when updating the range for the next layer of values, in addition to using the top 3 PPL incremental values, we also expand the upper and lower bounds outward by one step. Here, one step is defined as the difference between adjacent incremental values.

Adaptive Scaling Factors (ASF). Unlike methods such as YaRN and LongRoPE, which prescribe a strictly increasing order for scaling factors as frequency decreases, our approach does not confine the model to predetermined scaling paths. Considering the complex and often opaque internal mechanisms of models, we posit that different dimensions within each head may require distinct treatments, some might even need interpolation despite being high-frequency components. Thus, our framework allows for flexible scaling factors adjustments, tailored to the specific needs of each dimension.

Model and Evaluation Tasks. We carried out our experiments using the Llama2-7B Touvron et al. (2023) and Mistral-7B-v0.1 Jiang et al. (2023), aiming to extend the model’s context window to 64k tokens. We adopted YaRN’s methodology for perplexity (PPL) evaluation, utilizing ten samples from the Proof-pile dataset Rae et al. (2019) with lengths of at least 128k tokens for assessment. Additionally, we assessed model performance using 50 passkey Mohtashami and Jaggi (2023) tests at each length.

Fine-tuning Parameters. We used Llama2-7B as the base model. Initially, our DCIS algorithm was employed to identify scaling factors at the target length of 64k, setting the initial range between $[-5,5]$ with $C=10$ incremental values per segment. Following this, we segmented the PG19 dataset Gao et al. (2020) into 4k, 16k, and 64k context lengths, and then performed fine-tuning on each segment. The fine-tuning process closely mirrors YaRN’s protocol Peng et al. (2023) with a learning rate of $2\times 10^{-5}$. For context lengths of {4k, 16k, 64k}, we employed total batch sizes of {512, 64, 32}, and all models were fine-tuned for 400 steps.

Baseline. By fine-tuning with the aforementioned hyperparameters, we obtained YaRN-{16k, 64k} Peng et al. (2023), LongRoPE-16k Ding et al. (2024), and our proposed DCIS-{4k, 16k, 64k} models. The CLEX-16k Chen et al. (2023a) used the original model. Additionally, to compare with other types of methods, we employed the InfLLM Xiao et al. (2024), which claims to support infinitely long context windows by chunking and storing text sequences and retrieving the top-k most relevant chunks during inference.

We assessed the PPL of our proposed model and baseline models on the Proof-pile dataset, with results visualized in Figure 1. Figure 3 depicts the performance on the passkey evaluation. Notably, models from other methods that were fine-tuned on 16k-length and subsequently generalized to a 64k context windows experienced a marked increase in PPL at the target length (64k). Furthermore, these models entirely failed the passkey test at the target length. In stark contrast, our approach exhibited consistent performance across various context lengths, including the target length, and even outperformed other methods at shorter fine-tuning lengths (4k). Additionally, when fine-tuned on a 64k-length context, DCIS consistently outperforms YaRN on both PPL and passkey metrics.

All results underscore the significance of scaling factors. DCIS identify superior scaling factors, leading to improved performance across various sequence lengths and demonstrating strong generalization capabilities, thereby reducing the costs associated with model fine-tuning. The superior initial scaling factors contribute to further reduction in the number of fine-tuning steps. For example, in Appendix A.1, we report our observations that our method requires fewer fine-tuning steps to achieve comparable performance to YaRN.

Figure 4 illustrates the outcomes of our experiments wherein inference was conducted solely by adjusting scaling factors without fine-tuning. A comprehensive search for scaling factors was performed at target lengths of 64k and 128k, followed by direct PPL evaluation. The left and middle plots clearly indicate a consistent decline in our model’s PPL values at both target lengths, in stark contrast to the upward trends observed in other methods. To further validate the efficacy of our method, we extended our experiments to the Mistral model, employing the DCIS algorithm for scaling factors search, and evaluated the resulting models on the PG19 test set. The results unequivocally demonstrate the superior performance of our method across diverse model architectures.

In Appendix A.2, we compare the PPL of scaling factors for various methods at shorter lengths, without fine-tuning.

The efficiency of our DCIS algorithm is highlighted by comparing the search space. Specifically, our search space is the product of the total number of processed segments $d-2$ and the number of increments per segment $C$. The search space of evolutionary search utilized by LongRoPE is the product of the number of iterations $T$ and the population size $P$. For example, for the Llama2-7B model with default parameters, our search space equates to $(d-2)\times C=(128-2)\times 10=1260$. In contrast, the search space of evolutionary search calculates as $T\times P=40\times 64=2560$, signifying that our algorithm’s search speed is effectively double that of evolutionary search.

DCIS employs Adaptive Scaling Factors (ASF) (Section 4.2), and we investigated their impact in this section. Since LongRoPE utilizes strictly monotonically increasing scaling factors, we performed ablation studies on it to examine the effects of our proposed ASF. Specifically, we adjusted the evolutionary search algorithm in LongRoPE to a non-strictly increasing one, applied it to fine-tune the model on a 16k-length context, and subsequently generalized it to a 64k context window. Figures 5 and 6 depict the experimental outcomes. The left plot of Figure 5 visualizes the scaling factor distributions employed by various methods. Notably, both DCIS and LongRoPE + ASF exhibited irregular, sawtooth-like scaling factors, while YaRN and the original LongRoPE demonstrated more stable scaling factors. Furthermore, as shown in the right plot of Figure 5 and Figure 6, our ASF can improve LongRoPE in terms of both PPL and passkey scores, suggesting that imposing fewer constraints on scaling factors may enhance model performance.

Recent research may provide an explanation for our findings Barbero et al. (2024). They suggests that RoPE might not exhibit the long-range decay property as originally anticipated, indicating that distant tokens may have a larger influence than previously thought. Figure 2 in Barbero et al. (2024) shows that this decay property does not exist for Gaussian random vectors, further supporting our hypothesis that the complex and opaque internal mechanisms of models are better suited to learning appropriate scaling factors without constraints.

In addition to the previous assessments on PPL and passkey, employing LongBench and Open LLM Leaderboard, we conducted a comprehensive assessment of the models’ general abilities in both long and short context scenarios. The empirical results, as depicted in Tables 1 and 3, indicate that there is no significant difference in performance among the various methods, with different models performing best on different subsets. While models utilizing full attention generally achieved slightly better results, InfLLM, which leverages a retrieval-based approach, demonstrated a notable advantage in terms of memory efficiency. These findings suggest that the optimal choice of model is contingent upon specific application requirements.