Wavelet-based Positional Representation for Long Context

Authors: Yui Oka, Taku Hasegawa, Kyosuke Nishida, Kuniko Saito
Primary area: foundation or frontier models, including LLMs
Keywords: Positional Encoding, Extrapolation, Wavelet Transform, Transformers, RoPE, ALiBi, NLP

Abstract:
In the realm of large-scale language models, a significant challenge arises when extrapolating sequences beyond the maximum allowable length. 
This is because the model's position embedding mechanisms are limited to positions encountered during training, thus preventing effective representation of positions in longer sequences.
We analyzed conventional position encoding methods for long contexts and found the following characteristics.
(1) When the representation dimension is regarded as the time axis, Rotary Position Embedding (RoPE) can be interpreted as a restricted wavelet transform using Haar-like wavelets. 
However, because it uses only a fixed scale parameter, it does not fully exploit the advantages of wavelet transforms, which capture the fine movements of non-stationary signals using multiple scales (window sizes). 
This limitation could explain why RoPE performs poorly in extrapolation.
(2)
Previous research as well as our own analysis indicates that Attention with Linear Biases (ALiBi) functions similarly to windowed attention, using windows of varying sizes.
However, it has limitations in capturing deep dependencies because it restricts the receptive field of the model.
From these insights, we propose a new position representation method that captures multiple scales (i.e., window sizes) by leveraging wavelet transforms without limiting the model's attention field.
Experimental results show that this new method improves the performance of the model in both short and long contexts. 
In particular, our method allows extrapolation of position information without limiting the model's attention field.

Decision: accept (poster)
Reviewer scores: 5, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730508155071

thread_type:
official_review

<thread_object>
creator: reviewer_kkux

<content>
summary:
This paper conducts a unified analysis of positional encoding methods in large language models and proposes a novel wavelet-based approach. The work begins by demonstrating that RoPE can be interpreted as a restricted wavelet transform using Haar-like filters operating at a fixed scale. The authors then analyze ALiBi, revealing that it functions similarly to windowed attention with varying window sizes but is limited by constraints on the attention mechanism's receptive field. Building on these insights, the paper introduces a new wavelet-based positional representation method that leverages multiple scales through wavelet transforms. The method is designed to capture both local and long-range dependencies without restricting the model's receptive field. The authors validate their approach through experiments on short and long contexts.

strengths:
1. The motivation and connection are clear as the authors provide a unified analysis of the positional encoding methods.
2. The paper introduces a new wavelet-based positional representation method that leverages multiple scales through wavelet functions.
3. The proposed wavelet positional embedding shows improvements across tasks.

weaknesses:
1. Novelty Concerns: 
- The core idea of wavelet-based positional encoding has been explored in the previous work GMT (Ngo et al., 2023a,b), which uses graph wavelets to generate node positional representations that can capture the structural information of a center node on the graph at different resolutions, though in a different domain. However, there is insufficient discussion of how this approach differs fundamentally from or improves upon prior works.
2. Mathematical Rigor:
- The unified comparison for different positional encodings is not rigorous as the wavelet property needs stronger mathematical justification. The proposed "Haar-like wavelets" in Equation (7) require verification of wavelet admissibility conditions beyond just square integrability.
3. Experimental Design: 
- Limited exploration of wavelet families (only four tested).
- Insufficient analysis of scale parameter choices (only two tested).
4. Minor issues:
- The paper lacks of clear notation definition and problem setup affects readability.
- Inconsistent terminology between abstract ("simple Haar wavelets" in Line 18) and technical content ("Haar-like wavelets" in Sec 3.2).
- The phrase 'single window size' (in Line 19) should be replaced with 'fixed scale parameter' to align with standard wavelet terminology, as it specifically refers to the dilation factor in wavelet transforms.
- The space definition L⊭(R) in Line 135 appears to have a typographical error - it should be L²(R), the space of square-integrable functions.
- Some mathematical notation is not explicitly defined. For example, the superscript T in equation (3) seems to represent the length of a discrete sequence, but this isn't explicitly defined.

questions:
1. How does your approach fundamentally differ from MGT's wavelet position encoding?
2. Can you provide formal verification of the wavelet admissibility conditions for the proposed "Haar-like wavelets"? Or what else condition you might need to add to make it a wavelet?
3. Why were these specific wavelet families chosen? Have you explored other wavelet families, particularly discrete wavelets like Daubechies or biorthogonal wavelets?
4. What guided your choice of scale parameter ranges? Have you conducted sensitivity analyses on these choices?
5. What is the theoretical justification for removing the $\frac{1}{\sqrt{a}}$ amplitude term?

[1] Ngo, Nhat Khang, Truong Son Hy, and Risi Kondor. "Multiresolution graph transformers and wavelet positional encoding for learning long-range and hierarchical structures." The Journal of Chemical Physics 159.3 (2023).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730593922793

thread_type:
official_review

<thread_object>
creator: reviewer_dfcs

<content>
summary:
The paper proposes a wavelet transform-based positional representation for transformer models. It begins with highlighting the properties of existing position representation techniques such as relative position bias, RoPE and ALiBi. The authors then show how RoPE can be interpreted as a single scale wavelet transform with a Haar-like wavelet. Then the authors present a position embedding based on wavelet transform. This multi-scale embedding can be viewed as a generalization of RoPE and possess attractive "multi-window" properties of ALiBi. Experiment results have been reported on short and longer-context scenarios showing improved perplexity over existing methods and better extrapolation properties.

strengths:
- The paper is well-written barring some typographical errors. It motivates the problem well, starts with a well-defined goal, describes existing methods clearly, and presents the proposed method in a manner which is easy to appreciate.
- The paper tackles a critical problem of context length extrapolation which often arises in practical settings. The method holds significance, not just for the language modeling community, but also other domains such as time series forecasting where context length extrapolation may improve the performance of models on high-frequency data (see [1]). 
- The proposed method is novel to the best of my knowledge. It creatively relates RoPE to time-frequency analysis and proposes a promising method based on wavelet transforms. 
- The method outperforms alternatives and scales gracefully with extrapolated context lengths. 

[1] Ansari, Abdul Fatir, et al. "Chronos: Learning the language of time series." arXiv preprint arXiv:2403.07815 (2024).

weaknesses:
- While the discussion of related methods is generally well done, discussion of RoPE scaling techniques (linear, NTK-aware) is missing. A discussion and comparison with these techniques would significantly improve the positioning of this work.
- The results reported in sections 6 and 7 are excellent proofs of concept but they lack comprehensiveness. Particularly, in section 7, evaluations beyond the CodeParrot dataset would be needed to thoroughly appreciate the proposed method. Furthermore, the paper only evaluates the model in term of the perplexity. For a stronger evaluation, this needs to be augmented with task-based evaluations. Apart from common language tasks, you might also want to consider associative recall tasks [2] and the long range arena [3] since the focus of this work is context length extrapolation. To further demonstrate the robustness of the proposed method, domains beyond natural language may also be considered, for example, DNA modeling [4], audio generation [4] and time series forecasting [1]. 

I am willing to raise my score if the evaluation is strengthened. 

[1] Ansari, Abdul Fatir, et al. "Chronos: Learning the language of time series." arXiv preprint arXiv:2403.07815 (2024).    
[2] Arora, Simran, et al. "Zoology: Measuring and improving recall in efficient language models." arXiv preprint arXiv:2312.04927 (2023).    
[3] Tay, Yi, et al. "Long range arena: A benchmark for efficient transformers." arXiv preprint arXiv:2011.04006 (2020).    
[4] Gu, Albert, and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces." arXiv preprint arXiv:2312.00752 (2023).

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730686463207

thread_type:
official_review

<thread_object>
creator: reviewer_fca9

<content>
summary:
the paper presents a novel positional encoding method leveraging wavelet transforms to address the challenges of limited receptive field and  extrapolation in LLMs. 

The proposed wavelet-based approach aims to overcome these limitations by introducing a multi-scale analysis that captures the fluidity of natural language and enhances the model's ability to extrapolate beyond its training context length.

In extrapolation experiments, model that used Ricker-based wavelet positional embedding had the lowest PPL trained on the WikiText-103 dataset, compared with RoPE, ALibi, Transformer-XL.

strengths:
- The authors provide a solid theoretical foundation by drawing parallels between RoPE and wavelet transforms, and by extending this analogy to propose their method.
- The proposed method has the potential to be widely applicable to various transformer-based models
- The paper is easy to read and positions itself clearly with respect to related work.

weaknesses:
- The paper primarily evaluates the method on language modeling tasks. It would be valuable to see how the approach generalizes to other NLP tasks such as question answering or text summarization.
- The paper does not provide sufficient experimental evidence to support the authors' claim that Wavelet Transform can capture the dynamic changes in a sequence over positions.（L84-85）
- See my questions/suggestions below

questions:
- How did the “shift and scale parameters” been decided? Is there any ablation study or results for this particular setting?
- Did you compare the length extrapolation results with other RoPE-based methods? (eg. relative works you mentioned in L47-L48)
- Both RoPE and Wavelet methods aim to model distances at different scales by placing them in different dimensions [1]. In Figure 3, the diagonal lines in the RoPE plot appear because the model can focus on the features of specific token distances. However, Wavelet functions do not have this stable frequency inductive bias, which is why they do not produce diagonal patterns in the attention scores. Do you think this is a desirable property?
- In the experiment represented by Figure 3, what kind of text is input into the model? Why is the initial word particularly important? Which specific tokens does the model attend to? Is this phenomenon common?

[1] Hong, Xiangyu, et al. "On the token distance modeling ability of higher RoPE attention dimension." arXiv preprint arXiv:2410.08703 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1731255109802

thread_type:
official_review

<thread_object>
creator: reviewer_qwsg

<content>
summary:
In this work, the authors proposed wavelet-based positional encoding for the length extrapolation problem of language models. The motivation is based on an observation that the widely used Rope is related to wavelet transform via simple Haar wavelet functions with a fixed scale. By further investigating the properties of ALiBi, the authors modified the form of relative positional encoding to use wavelet transform based approaches. Several experiments are conducted to demonstrate the performance of the proposed positional encoding.

strengths:
1. The paper is easy-to-follow. 
2. The length extrapolation problem is important for language models.

weaknesses:
1. **There exist gaps from the motivation to the proposed approach**. In Section 3 and 4, the authors provide analysis on the relationship between RoPE and wavelet transform, and the properties of ALiBi like positional encodings. The ability of ALiBi to accommodate multiple window sizes is concluded as the key point for better length extrapolation performance, while RoPE is claimed to be worse. Based on these statements, a natural question is, what is the advantage of using wavelet transform? The authors do not provide enough supportive quantitative evidence for it, but directly modify the relative positional encoding by using wavelet transform. If it has superiorities, why don't we follow the form of RoPE to use wavelet transform instead of RPE? These points should be better clarified to make the motivation of the proposed approach more convincing.

2. **The design choice for the proposed approach should be better supported**. In Section 5, the authors propose to use Ricker Wavelet as the base function of wavelet transform, and set the shift and amplititude parameters with predefined strategies. However, the reasons to choose these designs are not well clarified (except lines 294-303, which are rather superficial).

3. **The empirical improvement is marginal**. From Table 2 and 3, we can see that the proposed PE does not show significant improvement compared to compared approaches such as ALiBi, which is proposed in 2022 and is not the state-of-the-art nowadays. The evaluations are also limited in task types and scales. Overall, the empirical results are not convincing to show the superiority of the proposed approach.

Overall, I hope the authors can well address the above concerns, which I think are important for the quality of this work.

questions:
See the Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732538924536

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Official Comment by Authors

comment:
Dear reviewers,

Thank you for taking the time to review this paper. We have added experiments and notes in response to some of the comments, and have reported them. As the discussion period is now coming to an end, we would be grateful if you could discuss them.

We are also currently conducting experiments on state-of-the-art RoPE improvement methods and other tasks. We will report the results of these experiments as soon as they are ready.
We look forward to discussing them with you!
</content>
</thread_object>

</thread>
<thread>
creation date:
1732861634883

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

Thank you very much for taking the time to review our manuscript. We deeply appreciate your thoughtful comments and suggestions.
In response to the valuable feedback we received, we have revised our manuscript accordingly and conducted additional experiments to address the points raised. We have carefully documented these changes and believe that your insights have significantly improved the quality of our work. We are truly grateful for your constructive input.

If you have any additional comments or further suggestions, we would greatly appreciate hearing from you. We understand that you are very busy, and we sincerely apologize for any inconvenience caused. However, we would be most grateful if you could kindly provide any feedback at your earliest convenience to help facilitate the review process.
Thank you again for your time and effort. We greatly appreciate your guidance and look forward to hearing from you.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a wavelet transform-based positional representation for transformer models. This multi-scale embedding can be viewed as a generalization of RoPE and possess attractive "multi-window" properties of ALiBi. Several experiments are conducted in long and short context scenarios to demonstrate the performance of the proposed positional encoding. They show improved perplexity, better length extrapolation properties and  sometimes even improved question answering performance compared to RoPE.

Strengths: The paper tackles a critical problem of context length extrapolation which often arises in practical settings. The proposed method is novel, and the authors provide a solid theoretical foundation by drawing parallels between RoPE and wavelet transform. Their unified analysis of positional encoding methods sheds more  light on this important topic. 

Weaknesses: Reviewers raised a few concerns regarding novelty, mathematical rigor and experimental evaluation.  In my estimation they have been adequately addressed by the authors during the discussion phase.

Overall I think this is a good paper which adds theoretical understanding, as well as a practical method with potential significance beyond the language learning community. I recommend accepting it.