Scaling Laws of RoPE-based Extrapolation

Dear reviewer, thank you for your review. When it comes to long-range real-world tasks, we compared some tasks often used in long-context research [1-2]. On one hand, concating multi-shot short prompts into a long context is included, such as 25-shot ARC-challenge [3], 10-shot Hellaswag [4], and 16-shot GSM8K [5] in LEval [2]. On the other hand, QA and extraction tasks of long documents are also involved, such as the TOEFL-QA [6] and TopicRetrieval [7] in L-Eval [2]. 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.

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.

[1] Peng, Bowen, et al. "Yarn: Efficient context window extension of large language models." arXiv preprint arXiv:2309.00071 (2023).

[2] An, Chenxin, et al. "L-Eval: Instituting Standardized Evaluation for Long Context Language Models." arXiv preprint arXiv:2307.11088 (2023).

[3] Clark, Peter, et al. "Think you have solved question answering? try arc, the ai2 reasoning challenge." arXiv preprint arXiv:1803.05457 (2018).

[4] Zellers, Rowan, et al. "Hellaswag: Can a machine really finish your sentence?." arXiv preprint arXiv:1905.07830 (2019).

[5] Cobbe, Karl, et al. "Training verifiers to solve math word problems." arXiv preprint arXiv:2110.14168 (2021).

[6] Tseng, Bo-Hsiang, et al. "Towards machine comprehension of spoken content: Initial TOEFL listening comprehension test by machine." arXiv preprint arXiv:1608.06378 (2016).

[7] Li, Dacheng, et al. How long can open-source llms truly promise on context length?, 2023. URL https://lmsys.org/blog/2023-06-29-longchat.

[8] Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. "Reformer: The efficient transformer." arXiv preprint arXiv:2001.04451 (2020).

[9] Beltagy, Iz, Matthew E. Peters, and Arman Cohan. "Longformer: The long-document transformer." arXiv preprint arXiv:2004.05150 (2020).

[10] Zaheer, Manzil, et al. "Big bird: Transformers for longer sequences." Advances in neural information processing systems 33 (2020): 17283-17297.

[11] Wang, Sinong, et al. "Linformer: Self-attention with linear complexity." arXiv preprint arXiv:2006.04768 (2020).

[12] Choromanski, Krzysztof, et al. "Rethinking attention with performers." arXiv preprint arXiv:2009.14794 (2020).

[13] Dao, Tri, et al. "Flashattention: Fast and memory-efficient exact attention with io-awareness." Advances in Neural Information Processing Systems 35 (2022): 16344-16359.

[14] Dao, Tri. "Flashattention-2: Faster attention with better parallelism and work partitioning." arXiv preprint arXiv:2307.08691 (2023).

[15] Li, Shenggui, et al. "Sequence parallelism: Long sequence training from system perspective." arXiv preprint arXiv:2105.13120 (2021).

[16] Liu, Hao, Matei Zaharia, and Pieter Abbeel. "Ring Attention with Blockwise Transformers for Near-Infinite Context." arXiv preprint arXiv:2310.01889 (2023).

[17] Li, Dacheng, et al. "LightSeq: Sequence Level Parallelism for Distributed Training of Long Context Transformers." arXiv preprint arXiv:2310.03294 (2023).

[18] Korthikanti, Vijay Anand, et al. "Reducing activation recomputation in large transformer models." Proceedings of Machine Learning and Systems 5 (2023).

[19] Jacobs, Sam Ade, et al. "DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models." arXiv preprint arXiv:2309.14509 (2023).

[20] Bavarian, Mohammad, et al. "Efficient training of language models to fill in the middle." arXiv preprint arXiv:2207.14255 (2022).

[21] Liu, Nelson F., et al. "Lost in the middle: How language models use long contexts." arXiv preprint arXiv:2307.03172 (2023).

[22] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks." arXiv preprint arXiv:2309.17453 (2023).

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.

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.

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.

[1] Peng, Bowen, et al. "Yarn: Efficient context window extension of large language models." arXiv preprint arXiv:2309.00071 (2023).

[2] Chen, Shouyuan, et al. "Extending context window of large language models via positional interpolation." arXiv preprint arXiv:2306.15595 (2023).

[3] Roziere, Baptiste, et al. "Code llama: Open foundation models for code." arXiv preprint arXiv:2308.12950 (2023).

[4] Xiong, Wenhan, et al. "Effective long-context scaling of foundation models." arXiv preprint arXiv:2309.16039 (2023).

[5] Kwiatkowski, Tom, et al. "Natural questions: a benchmark for question answering research." Transactions of the Association for Computational Linguistics 7 (2019): 453-466.

[6] Joshi, Mandar, et al. "Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension." arXiv preprint arXiv:1705.03551 (2017).

[7] Wang, Alex, et al. "Superglue: A stickier benchmark for general-purpose language understanding systems." Advances in neural information processing systems 32 (2019).

[8] Hendrycks, Dan, et al. "Measuring massive multitask language understanding." arXiv preprint arXiv:2009.03300 (2020).

[9] Paperno, Denis, et al. "The LAMBADA dataset: Word prediction requiring a broad discourse context." arXiv preprint arXiv:1606.06031 (2016).

[10] Clark, Peter, et al. "Think you have solved question answering? try arc, the ai2 reasoning challenge." arXiv preprint arXiv:1803.05457 (2018).

[11] Zellers, Rowan, et al. "Hellaswag: Can a machine really finish your sentence?." arXiv preprint arXiv:1905.07830 (2019).

[12] Lin, Stephanie, Jacob Hilton, and Owain Evans. "Truthfulqa: Measuring how models mimic human falsehoods." arXiv preprint arXiv:2109.07958 (2021).

[13] Mihaylov, Todor, et al. "Can a suit of armor conduct electricity? a new dataset for open book question answering." arXiv preprint arXiv:1809.02789 (2018).

[14] Bisk, Yonatan, et al. "Piqa: Reasoning about physical commonsense in natural language." Proceedings of the AAAI conference on artificial intelligence. Vol. 34. No. 05. 2020.

[15] Sakaguchi, Keisuke, et al. "Winogrande: An adversarial winograd schema challenge at scale." Communications of the ACM 64.9 (2021): 99-106.

[16] An, Chenxin, et al. "L-Eval: Instituting Standardized Evaluation for Long Context Language Models." arXiv preprint arXiv:2307.11088 (2023).

[17] Cobbe, Karl, et al. "Training verifiers to solve math word problems." arXiv preprint arXiv:2110.14168 (2021).

[18] Tseng, Bo-Hsiang, et al. "Towards machine comprehension of spoken content: Initial TOEFL listening comprehension test by machine." arXiv preprint arXiv:1608.06378 (2016).

[19] Li, Dacheng, et al. How long can open-source llms truly promise on context length?, 2023. URL https://lmsys.org/blog/2023-06-29-longchat.

[20] Su, Jianlin. Rectified rotary position embeddings., 2023. https://github.com/bojone/rerope.

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.

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.

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].

Settings：

model_size = 110M

n_layers = 12

embed_dim = 768

train_max_length = 1024

text_max_length = 4096

epochs = 5

datasets = wikitext-103

others as same as gpt2-base

Result (best in 5 epochs):

base=100 PPL=79.72

base=1000 PPL=32.87

base=10000 PPL=159.79

base=100000 PPL=270.94

base=1000000 PPL=320.32

Thank you very much for the sincere reply!

The author's response can address some of my questions, and I hope the author can incorporate these experiments into the next version of the manuscript.

I agree that this work has discovered some rope-based patterns for long context extrapolation, which will be helpful to research in this direction.

However, I believe there's potential for enhancement. For instance, the term "scaling law" is employed, yet the largest evaluation model merely reaches 13 billion. Furthermore, numerous studies have focused on optimizing Llama in long context scenarios, as cited in the authors' response. However, there's a scarcity of evidence demonstrating their tangible impact on practical applications.

Therefore, I decided to improve my score to 6.

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.