Extending Context Window in Large Language Models with Segmented Base Adjustment for Rotary Position Embeddings
Source abstract
Abstract
In the realm of large language models (LLMs), extending the context window for long text processing is crucial for enhancing performance. This paper introduces SBA-RoPE (Segmented Base Adjustment for Rotary Position Embeddings), a novel approach designed to efficiently extend the context window by segmentally adjusting the base of rotary position embeddings (RoPE). Unlike existing methods, such as Position Interpolation (PI), NTK, and YaRN, SBA-RoPE modifies the base of RoPE across different dimensions, optimizing the encoding of positional information for extended sequences. Through experiments on the Pythia model, we demonstrate the effectiveness of SBA-RoPE in extending context windows, particularly for texts exceeding the original training lengths. We fine-tuned the Pythia-2.8B model on the PG-19 dataset and conducted passkey retrieval and perplexity (PPL) experiments on the Proof-pile dataset to evaluate model performance. Results show that SBA-RoPE maintains or improves model performance when extending the context window, especially on longer text sequences. Compared to other methods, SBA-RoPE exhibits superior or comparable performance across various lengths and tasks, highlighting its potential as an effective technique for context window extension in LLMs.
Source: crossref · Quality: baseline
Research context
OpenAlex topics: Topic Modeling · Natural Language Processing Techniques · Multimodal Machine Learning Applications
SciVX discovery tags are terminology-based. OpenAlex topics are source-provided classifications.
Affiliations
Shenzhen International Graduate School, Tsinghua University, Shenzhen 518071, China
Pengcheng Laboratory, Shenzhen 518055, China
School of Dentistry, Seoul National University, Seoul 03080, Republic of Korea
Funding
National Natural Science Foundation of China · Research Center for Computer Network (Shenzhen) Ministry of Education · Natural Science Foundation of Guangdong Province · Basic Research Fund of Shenzhen City · Major Key Project of PCL for Experiments and Applications · Overseas Cooperation Research Fund of Tsinghua Shenzhen International Graduate School
References (25)
Show source-provided references
1. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., and Polosukhin, I. (2017). Advances in Neural Information Processing Systems, Curran Associates, Inc.
2. Dong, Q., Li, L., Dai, D., Zheng, C., Wu, Z., Chang, B., Sun, X., Xu, J., and Sui, Z. (2022). A survey for in-context learning. arXiv.
3. Recurrent neural networks
4. Long short-term memory10.1162/neco.1997.9.8.1735 ↗
5. Shaw, P., Uszkoreit, J., and Vaswani, A. (2018, January 1–6). Self-Attention with Relative Position Representations. Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), New Orleans, LA, USA.10.18653/v1/N18-2074 ↗
6. Press, O., Smith, N., and Lewis, M. (2022, January 25–29). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. Proceedings of the International Conference on Learning Representations, Virtual.
7. Sun, Y., Dong, L., Patra, B., Ma, S., Huang, S., Benhaim, A., Chaudhary, V., Song, X., and Wei, F. (2022). A Length-Extrapolatable Transformer. arXiv.10.18653/v1/2023.acl-long.816 ↗
8. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., and Azhar, F. (2023). Llama: Open and efficient foundation language models. arXiv.
9. Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L., Golding, L., He, H., Leahy, C., McDonell, K., and Phang, J. (2022). Gpt-neox-20b: An open-source autoregressive language model. arXiv.10.18653/v1/2022.bigscience-1.9 ↗
10. Anil, R., Dai, A.M., Firat, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., and Chen, Z. (2023). Palm 2 technical report. arXiv.
11. Su, J., Lu, Y., Pan, S., Wen, B., and Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv.
12. Chen, S., Wong, S., Chen, L., and Tian, Y. (2023). Extending context window of large language models via positional interpolation. arXiv.
13. bloc97 (2024, March 01). NTK-Aware Scaled RoPE Allows LLaMA Models to Have Extended (8k+) Context Size without any Fine-Tuning and Minimal Perplexity Degradation. Available online: https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have.
14. Peng, B., Quesnelle, J., Fan, H., and Shippole, E. (2023). Yarn: Efficient context window extension of large language models. arXiv.
15. Biderman, S., Schoelkopf, H., Anthony, Q.G., Bradley, H., O’Brien, K., Hallahan, E., Khan, M.A., Purohit, S., Prashanth, U.S., and Raff, E. (2023, January 23–29). Pythia: A suite for analyzing large language models across training and scaling. Proceedings of the International Conference on Machine Learning, Honolulu, HI, USA.
16. Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., and Nabeshima, N. (2020). The pile: An 800gb dataset of diverse text for language modeling. arXiv.
17. Lefaudeux, B., Massa, F., Liskovich, D., Xiong, W., Caggiano, V., Naren, S., Xu, M., Hu, J., Tintore, M., and Zhang, S. (2024, March 01). xFormers: A Modular and Hackable Transformer Modelling Library. Available online: https://github.com/facebookresearch/xformers.
18. Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., and Antiga, L. (2019, January 8–14). PyTorch: An Imperative Style, High-Performance Deep Learning Library. Proceedings of the 33rd International Conference on Neural Information Processing Systems, Vancouver, BC, Canada.
19. A tutorial on the cross-entropy method10.1007/s10479-005-5724-z ↗
20. Loshchilov, I., and Hutter, F. (2017). Decoupled weight decay regularization. arXiv.
21. Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., and Venkatesh, G. (2017). Mixed precision training. arXiv.
22. Rasley, J., Rajbhandari, S., Ruwase, O., and He, Y. (2020, January 6–10). Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, Virtual Event.10.1145/3394486.3406703 ↗
23. Rae, J.W., Potapenko, A., Jayakumar, S.M., Hillier, C., and Lillicrap, T.P. (May, January 26). Compressive Transformers for Long-Range Sequence Modelling. Proceedings of the International Conference on Learning Representations, Virtually.
24. Azerbayev, Z., Ayers, E., and Piotrowski, B. (2024, March 01). Proof-Pile. Available online: https://github.com/zhangir-azerbayev/proof-pile.
25. Mohtashami, A., and Jaggi, M. (2023). Landmark Attention: Random-Access Infinite Context Length for Transformers. arXiv.