Functional Interpolation for Relative Positions improves Long Context Transformers

Regarding Weakness 3. Thanks for the references! We would like to first highlight that all these 3 references appeared only after May 28, 2023, with the YaRN paper being uploaded on August 31st, just a month before the ICLR deadline. Hence we did not have comparisons in our submission.

According to ICLR policy, authors are not required to compare their own work to papers published on or after May 28, 2023. As all the references have appeared after that and have not been peer-reviewed (we are not even able to find paper drafts for [3] and [5]), they should be considered concurrent to our work.

Among these references, YaRN [4] seems to be the only one with a formal preprint paper and a complete codebase. We follow your suggestion and additionally compare FIRE with YaRN. The results are shown below:

• Language modeling on C4

• Language modeling on arXiv

• Language modeling on github

• SCROLLS benchmark

The performance of YaRN is worse than FIRE on most datasets and sequence lengths, though it improves over RoPE PI. In [4], YaRN is claimed to be better than NTK-RoPE, so FIRE should also be better than NTK-RoPE. Besides, our experiments show that Alibi-based models consistently underperform RoPE/FIRE-based models. Therefore, we believe NTK-Alibi may be less promising to compare with.

[1] Qin, Zhen, et al. "Toeplitz Neural Network for Sequence Modeling." arXiv preprint arXiv:2305.04749 (2023).

[2] Fu, Daniel Y., et al. "Simple hardware-efficient long convolutions for sequence modeling." arXiv preprint arXiv:2302.06646 (2023).

[3] https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/

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

[5] https://github.com/keezen/ntk_alibi/blob/main/readme_en.md

We sincerely hope that our responses address your concerns and you reevaluate our work based on the responses. We are also willing to discuss with you if you have any further questions.

Thank you for your careful review! Your comments point out very meaningful avenues for further improvement of our work. We respond to your concerns as below.

Regarding Weakness 1. Thanks for the reference! We note that [1] and [2] propose new deep learning architectures for sequence modeling, while our work focuses on the Transformer variants. Besides, The idea of "simple" baseline $b(i,j)=f(i-j)$ is exactly T5's RPE, which has been extensively compared in our experiments. This "simple" baseline, along with many existing methods, are mainly evaluated using perplexity-based measures in previous works. Our systematic experiments show that they suffer significant performance drops with longer contexts across diverse evaluation tasks. Those methods typically have hard-coded inductive biases (with fixed RPE patterns or the same encoding for distant tokens), which prevent them from being universally effective. Our method addresses these drawbacks by using a learnable approach with interpolation to truly enable length generalization.

To reiterate, progressive interpolation is proposed to map relative distances to numbers in $[0,1]$ and avoid out-of-distribution issues during inferences. The $\psi$ transform and thresholding technique further enhance modeling local positions and short sequences. We emphasize that these transformations do not limit the expressive power of FIRE to learn arbitrary biases, which is theoretically shown in Theorem 3.1. Moreover, we carefully ablate these design choices through experiments in Appendix B. Thus, we believe the design of FIRE is both practically effective and theoretically principled.

Regarding Weakness 2. Please note that the performance gap for long context settings between FIRE and baselines is so large ($>2$ perplexity points and $>1$ average SCROLLS score), shifting speed-accuracy pareto curve significantly. In fact, even base-sized FIRE is better than or on par with many large-sized baselines on SCROLLS. Further, in Sec. 4.6, we show that FIRE can be accelerated by layerwise sharing the positional encoding (FIRE-S) and provides the inference speed evaluations. We also follow your suggestion to additionally compare the speed and accuracy of RoPE, Alibi, and FIRE below:

It's clear that our method obtains much stronger performance with decent training speed. We will include a full version of the table in the final draft.

Thank you for supporting our work! We address your concerns as below.

Regarding the thresholding technique. Thank you for your careful reading! We would like to first note that even FIRE without thresholding outperforms all the baselines (including RoPE, T5's RPE, etc) on all the sequence lengths. The thresholding technique is proposed to further enhance FIRE which turns out to be highly effective as demonstrated in our experiments. For your convenience, we reiterate the performance comparison in the table below. These results are in Sec. 4.1 & Appendix B.1 in the paper.

We follow your suggestion and additionally conduct GLUE/SuperGLUE finetuning for the variant based on Eq. (21). The averaged accuracy is 69.06. We will include this result in Table 3 for a more complete comparison.

Regarding the learned thresholding parameter $L$. This is a good question. For the base-sized model with pretraining sequence length 2048, the learned parameter $L$ is between 1200 to 1600 across different layers.

We agree that setting $L$ to a fixed value is another viable option. In our preliminary exploration, FIRE with a fixed/learnable $L$ achieved 24.65/25.64 average SCROLLS scores respectively. Both versions outperform all the baselines, while the learnable variant leads to better performances. Moreover, the fixed variant introduces one more hyper-parameter and may require more tuning. Thus, FIRE uses learnable $L$ as the default choice.

Regarding the short sequence performance. We agree with the reviewer that FIRE does not lead to performance gains on GLUE/SuperGLUE, but is on par with existing approaches. We emphasize that the primary focus of FIRE is length generalization on long sequences, and our experiments on GLUE/SuperGLUE only serve as an additional sanity check of performances on shorter sequences. RoPE shows good performances on GLUE/SuperGLUE (normal sequence length), but it significantly underperforms FIRE on long sequences (Sec. 4.1-4.3). Besides, Sec. 4.6 shows that FIRE-S (with layer-wise sharing) is faster than RoPE, and the GLUE/SuperGLUE performance is comparable to RoPE (FIRE-S 71.04 v.s. RoPE 71.15). We believe it is a merit of our work to conduct more comprehensive evaluations rather than cherry-picking experimental settings.

We sincerely hope that our responses address your concerns and you reevaluate our work based on the responses. We are also willing to discuss with you if you have any further questions.

Thank you for supporting our work! We address your concerns below.

Regarding the performance of Alibi. We would like first note that our evaluation is much more comprehensive (as the reviewer noted earlier) including zero shot and finetuning on downstream tasks, whereas [1,2] mainly rely on language modeling performance for their evaluation. Second, for the language modeling task, we use standard evaluation to compute perplexity where we do a single forward pass on the entire sentence, whereas Alibi uses sliding window to evaluate perplexity. These differences could account for the discrepancies observed in Alibi's performance between our study and those reported in [1,2].

[1] Press, Ofir, Noah Smith, and Mike Lewis. "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation." International Conference on Learning Representations. 2021.

[2] Chi, Ta-Chung, et al. "Kerple: Kernelized relative positional embedding for length extrapolation." Advances in Neural Information Processing Systems 35 (2022): 8386-8399.

Regarding the simple baseline. Thanks for this suggestion. In the ablation study of our submission, we have tested a FIRE variant $f_{\theta}\left(\frac{\psi(i-j)}{\psi(\max\\{i, L\\})}\right)$ where $f_{\theta}$ is a linear function instead of a neural network. The results, presented in Table 4, show that using a linear $f_{\theta}$ leads to a 0.26-point drop in log perplexity on sequence length 8192, though is still better than competing baseline methods. This finding indicates using a neural-network-parametrization for $f_{\theta}$ is also important for the final performance. Besides, although the visualization shows the learned positional encoding is not very complex, the patterns still exhibit some non-linearity which may require a neural network to model.

We sincerely hope that our responses address your concerns. Let us know if you have any further questions.