Summary
This paper proposes a way of modifying RoPE embeddings so that they extrapolate better to sequences longer than the model was originally trained on. The YaRN method leverages a recent (unpublished) method, called "NTK-by-parts" [1], as well as modifying the default temperature value of the attention softmax. The NTK-by-parts method is based on the following observation: Performing position interpolation [2] on the high-frequency (small wavelength) RoPE rotation angles could make it such that neighboring tokens have rotation angles that are very difficult for the model to distinguish, thus potentially confusing the model regarding the relative positions of those tokens. Thus, the NTK-by-parts method functions by determining the amount of position interpolation to perform based on the frequency of the RoPE rotation angles: no interpolation is used for high-frequency angles, while standard interpolation is used for the low-frequency angles (and "middle" frequency angles are handled with some interpolation). Once these updates to the RoPE embeddings (and attention temperature) are made, the whole model is fine-tuned on a small amount of data (~0.1% of the pre-training corpus).
The YaRN method is shown to perform relative well relative to baseline RoPE extension methods across various experiments (e.g., long sequence language modeling, and several standard LLM benchmark tasks from HuggingFace Open LLM leaderboard).
- [1] https://github.com/jquesnelle/yarn/pull/1
- [2] S. Chen, S. Wong, L. Chen, and Y. Tian. Extending context window of large language models via
positional interpolation, 2023. arXiv: 2306.15595.
Strengths
- The proposed method appears to perform well relative to a few baseline RoPE extension methods.
- The proposed method is relatively well-motivated --- it seems like a reasonable idea to not interpolate for high-frequency rotation angles.
- The proposed method has already been used in open-source LLMs in industry.
Weaknesses
- I found the paper relatively difficult to follow. I think the method could be presented in a much simpler and more direct manner. The "NTK-aware interpolation" could likely be moved to the appendix, as it is not part of the YaRN method. The background section could be significantly shortened (currently 1.5 pages).
- I think the experimental results could be much more thorough, and much more clearly presented.
- Including more baselines throughout all the experimental results: Standard RoPE embeddings (e.g., Llama-2), ALiBi, T5 relative position embeddings, absolute position embeddings, Position Interpolation (PI), NTK-aware interpolation, YaRN (with and without fine-tuning, with and without temperature scaling, with static vs. dynamic scaling, different scale factors, etc.).
- Replacing tables with plots (similar to Figure 4 in appendix, but with clearer+more baselines). The tables are more difficult to interpret. And the plots show results for short sequences as well as long sequences, which is helpful for understanding the performance of the method.
- Adding detailed ablations for YaRN, as mentioned above (with and without fine-tuning, with and without temperature scaling, with static vs. dynamic scaling, different scale factors, etc.).
- Choosing better baselines. Why are Together and Code Llama chosen as the main baselines? Why is PI not included in every table?
Overall, given the issues with the presentation of the method and (most importantly) the experimental results, I have currently chosen "marginal reject" (I was torn between "marginal accept" and "marginal reject"). While I think the community could benefit from seeing the proposed idea (in particular, the part about only doing interpolation for low-frequency RoPE angles), I think there are currently too many open questions related to the results, and how this method compares with baselines (and with itself, with ablations), to accept the paper in the current form to a peer-reviewed conference. Open to being swayed though, given the potential of this line of work!
Questions
- In equation (17), should you be using $b'$ or $b$? It seems more natural to use $b$, and would be confusing to use $b'$, which already includes some scaling modifications.
- For equations 19 and 20, isn't this equivalent to $g(m) = (1-\gamma(r(d))) * (1/s) + \gamma(r(d))$, and $h(\theta_d) = \theta_d$? Perhaps this would be simpler, to make this look more like the position interpolation equations (which only modify $g(m)$), as opposed to the "NTK-aware" equations which only modify $h(\theta_d)$.
- Can you show results of YaRN with and without fine-tuning?
- Can you clarify how the sliding window evaluation method works (Press et al., 2022)? I read the reference but was still confused, so adding a simple explanation in the paper would be very helpful. When reporting on performance with evaluation context window size $N$, does this mean you measure perplexity on predicting token N+1? Or you measure the average perplexity across all $N$ tokens (so each token on average sees N/2 context)?
- When you say in your plots "Yarn-Llama-2-7b-64k", this means you used s=16? What fine-tuning data and sequence lengths did you use for both s=16 and s=32?
- Are the only results in your paper that use dynamic scaling those in Figure 5 of appendix B.4?
- Why do you compare with Code-Llama, which is specialized for code?
- Do you have intuition for why temperature scaling is necessary? Does it improve performance with regular RoPE embeddings, for short documents? Or just for long documents when using YaRN? Can you add more explanation regarding the intuition/reasons for this method, and how it performs outside the context of YaRN?
Rating
6: marginally above the acceptance threshold
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.