Paper: Scaling Laws of RoPE-based Extrapolation (ICLR-2024)

Reviewer comment (Presentation):
There is no related work section. It is "okay" for expert readers, but it would be very helpful to make a broader impact if authors can write one. I think authors should at least include:
- scaling seq len with position embedding. The authors discussed some of these approaches, but it would be better to summarize them in a thread.
- scaling seq len with efficient implementation, such as Flash attention 1&2 [1]
- scaling seq len with a distributed system, such as Sequence Parallelism [2][3]

Author response:
When it comes to the related work, if the paper gets accepted, we will summarize scaling seq len with position embedding into one subsection and add a subsection respectively to scaling seq len with efficient attention and scaling seq len with a distributed system. The added content is as follows.

Long sequence modeling faces the following three challenges: the O(N2) computational complexity, the difficulty of parallelism for sequence length, and the collapse of the performance beyond the training length. First of all, due to the O(N2) complexity of the self-attention mechanism in Transformers, long sequence modeling requires a large computational overhead. Early efficient Transformers reduced computational overhead through methods such as recalculation like ReFormer [8] and sparse attention such as LongFormer [9] and BigBird [10]. After that, LinFormer [11] linearizes the self-attention complexity through low-rank approximation, and Performer [12] expands the softmax calculation through a kernel function. Compared with these efficient attention works based on approximation, FlashAttention [13-14] reduces HBM reading and writing and turns self-attention into a form that can be iteratively calculated in a loop, which greatly improves the utilization of computing resources, reduces the cost of long sequence calculations. and has become the standard for long-sequence modeling.

The second challenge is parallelism for sequence length, that is, Sequence Parallelism (SP) [15]. Since the memory and calculation of a single device are limited, ColossalAI-SP [15] first proposes the segmentation and parallelism for the sequence dimension in addition to tensor parallelism for the hidden dimension and pipeline parallelism for model depth. On this basis, Ring Attention [16] uses blockwise self-attention to split long sequences into different devices and overlap the communication of key-value blocks. Besides, LightSeq [17] further improves the efficiency of long sequence modeling through load balancing for causal language modelings and a re-materialization-aware checkpointing strategy. The advantage of the above sequence parallelism is that it can break through architectural limitations and achieve infinitely long sequence modeling. In contrast, another type of sequence parallelism emphasizes maintaining the combination with existing efficient self-attention mechanisms such as FlashAttention [13-14] to achieve input with almost infinity. For example, Megatron-SP [18] only uses sequence parallelism during Dropout and Layernorm operations, thereby reducing more activation redundancy. In addition, DeepSpeed-Ulysses [19] uses an alternative all-to-all collective communication gathering weight for attention computation when segmenting the sequence, avoiding communication overhead that originally increases with length.

Reviewer comment (Experiments):
It would also be very interesting to study the filling-in-the-middle pattern after applying the proposed trick. Is the filling-in-the-middle related to the position embedding? This is just a bonus question. It is an okay paper without this study, but I suggest authors to study this to deliver more insights.

Author response:
When it comes to the filling-in-the-middle pattern [20], filling-in-the-middle [20] is a training paradigm that we believe is orthogonal to RoPE-based extrapolation. If you talk about lost-in-the-middle [21], the relationship between lost-in-the-middle and RoPE-based extrapolation is an interesting research topic. Preliminarily, we find that tuning with a larger base is essentially enlarging the interval lost span in the middle. In essence, the RoPE-based LLMs still only pay attention to the information of the adjacent and initial positions, as discussed in StreamingLLM [22]. We are now conducting more experimental verifications on this.

Reviewer comment (Novelty):
In fact, the analysis results of this article are basically available in the YaRN[4] paper, which had pointed out that certain dimensions have high frequency and enough turns are made during training, therefore, the training of these dimensions is sufficient and should remain unchanged during extrapolation; for those dimensions that are less than one circle, only interpolation (PI) can be used to ensure numerical stability. This is obviously the equivalent expression of the proposed concept "critical dimension" in this work.

Author response:
Dear reviewer, thank you for your question. When it comes to the comparison with YaRN [1], we think our work is our concurrent work, because YaRN is published only one month before our submission. Compared with YaRN [1], our work gives a quantitative analysis of the critical dimension. Based on this analysis, we propose and validate the relationship between RoPE extrapolation, period, and attention score, and give the upper bound of model extrapolation under different fine-tuning bases. These can not be replaced and treated equally with simply saying that low-frequency dimensions can extrapolate directly and high-frequency dimensions need to be interpolated.

Reviewer comment (Evaluation):
Most experimental comparisons only report ppl, which does not represent real long context ability. The authors should conduct experiments on some practical benchmarks and tasks. Eg. PASSKEY RETRIEVAL[2], the Hugging Face Open LLM Leaderboard[3] (ARC, HellaSwag, TruthfulQA, etc).

Author response:
When it comes to perplexity, perplexity is measured in all extrapolation work, such as Linear PI [2], CodeLLaMA [3] and LLaMA2Long [4]. Our work focuses on the curve of perplexity w.r.t. context length, which is commonly reported in related works [1-4]. 

When it comes to the downstream tasks, our work actually provides the results of 0-shot NQ [5], TriviaQA [6], SuperGLUE [7] and 5-shot MMLU [8] in Appendix B.2. In order to present a comprehensive evaluation, we included new results on 0-shot Lambada [9], ARC-easy/challenge [10], Hellaswag [11], TruthfulQA [12], OpenbookQA [13], PIQA [14] and Winogrande [15]. We also show the average score over the above benchmark and updated 0-shot SuperGLUE [7] as well as 5-shot MMLU [8]. Particularly, based on the evaluation results, we find that the effect of tuning on the original context length as emphasized in our work is almost consistently better than tuning on the longer context length with the same base value. As a result, we achieve a balance between long context and short tasks by tuning with a larger base on the original length, which is a common issue raised in the work of expanding the context during the fine-tuning phase [1-4].

Reviewer comment (Evaluation):
The RoPE extrapolation methods discussed in the paper may have certain limitations in practical applications, it is just verified over perplexity, but not enough practical application scene.

Author response:
We also compared some tasks with 16K context length used in long context research [1,16]. On one hand, concating multi-shot short prompts into a long context [1,4] is included, such as 25-shot ARC-challenge [1,10], 10-shot Hellaswag [1,11], and 16-shot GSM8K [17] in LEval [16]. On the other hand, QA and extraction tasks of long documents are also involved, such as TOEFL-QA [18] and TopicRetrieval [19] in L-Eval [16]. Our method is still capable of these tasks, as shown in the table below. Particularly, even for RoPE-based LLMs fine-tuned on pre-training context, namely 4k tokens, they still have certain effects on these tasks with 16K context size. This also proves one of the claims of our work, that tuning with the original pre-training context can also boost the length extrapolation of RoPE-based LLMs. We will publish more evaluation results for long document tasks in the later version.

|  |  | 25-shot ARC-c | 10-shot hellaswag | 10-shot GSM | TOEFL-QA | topic_retrieval |
| --- | --- | --- | --- | --- | --- | --- |
| llama2_7B | pretrained | 48.14  | 72.78  | 0.00  | 34.94  | 0.00  |
|  | ntk_fixed_8 | 46.78  | 68.94  | 5.00  | 34.20  | 31.33  |
|  | ntk_dynamic | 48.14  | 72.73  | 9.00  | 37.17  | 24.67  |
|  | base_10000 | 47.80  | 72.49  | 0.00  | 18.59  | 0.00  |
|  | base_80000 | 47.12  | 72.66  | 11.00  | 36.43  | 29.33  |
|  | base_160000 | 45.42  | 72.60  | 10.00  | 40.15  | 38.00  |
|  | base_400000 | 47.12  | 72.63  | 8.00  | 41.64  | 36.67  |
|  | base_600000 | 46.10  | 72.54  | 12.00  | 38.66  | 34.00  |
|  | base_1000000 | 46.78  | 72.71  | 9.00  | 31.97  | 34.00  |
| llama2_13B | pretrained | 54.58  | 73.89  | 0.00  | 59.48  | 0.00  |
|  | ntk_fixed_8 | 50.85  | 66.88  | 8.00  | 58.74  | 36.00  |
|  | ntk_dynamic | 54.58  | 73.93  | 22.00  | 59.85  | 30.00  |
|  | base_10000 | 53.56  | 75.53  | 0.00  | 56.13  | 0.00  |
|  | base_80000 | 53.56  | 75.69  | 20.00  | 63.94  | 30.00  |
|  | base_160000 | 53.56  | 75.50  | 23.00  | 62.83  | 39.33  |
|  | base_400000 | 54.24  | 75.56  | 17.00  | 59.48  | 34.67  |
|  | base_600000 | 53.90  | 75.29  | 14.00  | 63.94  | 36.00  |
|  | base_1000000 | 52.54  | 75.18  | 10.00  | 64.68  | 28.67  |

Reviewer comment (Evaluation):
Compared to baselines, it lacks some comparison results over computational resources and memory usage.

Author response:
When it comes to the comparison of computational resources and memory usage, since we only modify the rotation angle base of RoPE, no additional computation and memory are added. At the same time, since we conduct the training with an original context length, 4K, to achieve 100K context length, compared with Linear PI [2], CodeLLaMA [3], YaRN [1], and LLaMA2Long [4], our work also has certain advantages in the fine-tuning speed and memory cost. If the paper gets accepted, we will add detailed data about these contents. Thanks to the reviewer for the reminder.

Reviewer comment (Experiments):
Missing reference. The work should compare results with ReRoPE[1].

Author response:
Finally, when it comes to the reference of ReRoPE [20] and the verification of larger (such as 70B) and smaller (such as 330M, 1B) RoPE-based LM, we are conducting the verification of the 70B model at present. Since the computational cost is huge, we will add the results to our paper in the later version.

Reviewer comment (Experiments):
Only 7B and 13B models are verified. Have you experimented with the proposed methods on larger or smaller models? So that the "scaling law" saying can be more convincing.

Author response:
Thank you for your interest in our work and your comments. Regarding the phenomenon you report, I hope you can further report on some of your training configurations. Is it fine-tuning with a different base value as discussed in this article, or directly changing the base value in the inference stage or the pre-training stage? I am still not clear about the changing process of perplexity w.r.t. the length of language modeling. In addition, regarding the laws we propose, there may indeed be a range of applicable model sizes, and even those scaling laws that are generally accepted still have such a phenomenon. Thank you for reporting these results on the 110M size model. It has some reference significance for us. We will do the same in the future. We are also doing some experiments on the 70B size model.

Reviewer comment (Theory):
A bit more clarification around the calculation of critical beta, and why extrapolation capabilities exist even when beta values are large could be added since we will see OOD embeddings as soon as T is larger than T_{train} during inference.

Author response:
Dear reviewer, thank you very much for your recognition and comments on this paper. Concerning the clarification of the critical base and the extrapolation of a larger base beyond T_{train}, we will add the following discussion if the paper gets accepted. I don’t know whether the following discussion can make you feel OK. 

The critical base refers to a base value that exists when extending the further training context. This base value is just enough to make the period of the original critical dimension equal to the length of the further training context length. In this way, RoPE-based LLM only uses the dimensions within the original critical dimension to perceive the position information that happens to be within the continued training length. If the base is smaller, then more dimensions can sense position information, and the perplexity curve will be flatter; if the base is larger, then longer position information can be represented by dimensions within the critical dimension, thereby expanding the supported context window.

When it comes to the extrapolation beyond T_{train}, this is still related to the critical dimension. For dimensions within the critical dimension, the inputs of sin and cos representing the positional embedding can traverse 0 to 2$\pi$, and any value can be expressed normally; for the dimensions beyond the critical dimension, the inputs of sin and cos have not even completed a round. After enlarging the base, the rotation angle $\theta$ will shrink. For the same sin and cos values, the allowed input position index will become larger. 
For the dimensions beyond the critical dimension, when the value of the position index * $\theta$ exceeds the value seen during training, then ood embedding will occur, resulting in extrapolation problems; while the value of the input subscript * $\theta$ still falls within the range seen at the pre-training stage, there will be no extrapolation problem. In other words, the increase in base leads the rotation angle to become smaller, which increases the range of allowed values for the index in the position encoding, thus enabling extrapolation.