Contextual Position Encoding: Learning to Count What’s Important

We appreciate your review of our paper.

A core issue here is whether the series of problems currently described are truly caused by PE.

We experimentally demonstrated that changing PE brings improvement in various tasks where position in abstract terms are more useful. This is evidence that PE is a major component of these issues. Position information is crucial in locating N-th sentence or paragraph. But, we argue that position alone is not sufficient, thus proposing a contextualized version. Similar findings in failure mechanisms were demonstrated both empirically and theoretically on context-dependent counting toy tasks in earlier papers, such as [1,3,4], and are main motivators to search for alternative architecture, such as [2].

The analysis in Table 1 attributes the problems to the current PE, but the analysis is insufficient and may be due to the inadequate reasoning abilities of LLMs.

Table 1 is meant to be a motivating example highlighting the problem, rather than a quantitative analysis. In Section 5.4, we do careful experimental tests on this word counting issue, and demonstrate that our method helps with it. Given that current LLMs are capable of solving complex math problems, counting words in a sentence is unlikely to be more challenging in terms of reasoning. If only the target sentence is given those LLMs, then they count words easily. This indicates that the problem is inability to isolate and attend to a specific sentence. We do provide extensive analysis of why this might happen in Section 3.1.

One problem with this work is that it requires extensive training; perhaps optimizing during the reasoning phase could resolve these issues (as mentioned in Table 1).

We argue that all PE methods require extensive training. To use RoPE, for example, one needs to add it to the model during both pre-training and post-training. The same is true for CoPE. As a part of the architecture, all PE methods usually are trained together with the model through all stages (pre-training and post-training). While we agree that prompt engineering or finetuning might help to steer the model in the right direction, the goal of our approach is to eliminate the core reason of the issues.

Since this work is general, the experiments on generality (Table 12) should not only compare RoPE with RoPE+CoPE but also display the performance of CoPE alone.

Please see our response to A1r6 for why we made this comparison. We report results with CoPE alone, as well as CoPE + RoPE, on smaller models where specialized hardware-aware attention implementation is not necessary.

The content in lines 32 to 36 of the Introduction is somewhat redundant and could be moved to the related work or background section.

Yes, it can be moved to the background section. We thought some people that are not very familiar with PE may find those parts useful in understanding our motivation.

Why emphasize OOD in the experiment? For PE, does improved OOD indicate better generalizability and extensibility of the method?

OOD generalization is generally a good property in most methods as it means the method is robust. We want LLM to perform well in OOD situations because users might ask questions that are different from what training data covered. PE is part of this, so yes PE with better OOD is desirable for building a robust LLM, specifically when the context gets longer or the number of things in the context change.

Could you provide examples of correct outputs from the experiments in Table 7? This would allow us to better understand the output patterns of LLMs on this task, thereby gaining insights into issues related to PE.

Table 7 is showing a motivational example showing the shortcoming of LLMs. It is obtained by prompting the LLMs as shown in Table7, with those exact words. The correct output should have had the correct answers, which is that Alice nor book is not mentioned in the last sentence.

References:

1. Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Exposing attention glitches with flip-flop language modeling.
2. lbert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces.
3. Michael Hahn. Theoretical Limitations of Self-Attention in Neural Sequence Models.
4. David Chiang and Peter Cholak. Overcoming a Theoretical Limitation of Self-Attention

We thank the reviewer for the extensive and detailed review.

Scalability: The presented method is a parameterized method of encoding positions. The authors mention that to control the number of added parameters one could control the value of p_max (line 227). However, the effect of limiting p_max on performance is unclear, especially for longer sequences (say >10k tokens).

We didn’t see a significant performance drop when reducing p_max. You can see those results in Table 9 where we tried p_max values of 16, 64, and 1024. Also, limiting p_max is an optional feature and CoPE will work without this limit, as we have done in some of our experiments. Without p_max limit (i.e. set to max), CoPE still has linear computational complexity, which is much lighter weight than the attention computations that are quadratic.

Secondly, for CoPE one would need to write ‘qk’ dot product and hence cannot use the sub-quadratic attention methods such as flash attention and ring attention. This is a crucial limiting point because sub-quadratic attention has been a key for scaling transformer LLMs to process longer sequences.

This is not accurate. Exact attention computation requires a QK dot product. Sub quadratic attention methods always have some sort of approximation. Ring attention is not sub-quadratic and computes QK product exactly as it does not do approximation. It makes computation efficient to run on multiple devices (it is distributed through manageable blocks of the input sequence across multiple devices). Note that every operation in CoPE is linear and only quadratic operation is QK and QV multiplications, which are part of normal attention. If there are approximate methods for calculating these in sub-quadratic time, then surely those can be adapted to be used in CoPE as well.

The authors show in Section 5.5 (figure 3) that the model trained on a sequence length of 512 can have reducing PPL values until the context reaches 2.5k tokens. It is then critical to see how long this pattern stays stable.

We demonstrated generalization to 5x longer sequences, which we think provides sufficient evidence of better generalization. Note that when RoPE is extended to longer sequences, it generally requires adjusting base frequency and finetuning on longer sequences. In our experiments, we do not do any finetuning and directly test OOD generalization to longer sequences. This is possible because the same CoPE embeddings appear in different locations in each sample. Thus, when there is a longer sequence, a position embedding that has been seen in shorter sequences can now appear in distant context, allowing the model to generalize and attend to such distances.

Or if the authors could present some analysis regarding the relationship between p_max and extension limit for different model sizes, that’d be helpful for the readers

We provided results with different p_max values in Table 9. We also experimented with diverse model sizes ranging from very small to 1.4B size, covering 7 different tasks. While adding even more experiments could be interesting, the main goal of the paper is to demonstrate the need for contextual position, and we feel the current experimental evidence provides sufficient evidence for that.

It is also crucial to report cost analysis, including the additional time and memory required by the CoPE method compared to RoPE (with and without flash attn), in order to compare the methods fairly. This will give readers and practitioners a clear idea of the required compute investment for better performance.

As you can see in the implementation in Appendix B, none of the operations in CoPE has more complexity in terms of computation and memory than what is already done in normal attention. As we discussed above, there is a difference when hardware-aware efficient attention implementations are involved as there is none for CoPE yet. However, we showed that even adding CoPE to only some layers brings a significant improvement while managing the extra computation due to lack of efficient implementation. We hope future work will bring efficient attention to CoPE as well.

We thank the reviewer for the time and effort.

The mapping from a count (p_{ij}) to an embedding (eq. 5) is fine, but the rounding to the nearest integer count seems a little complicated, and learning different embeddings for every integer count adds a lot of parameters and imposes a maximum count.

Our method of interpolating between integer position embedding is simple to implement (see Appendix B), and performs well in practice, as we have shown with our experiments. We have also released our code to ensure reproducibility of our experiments and further exploration. The number of new parameters added is really small and negligible compared to the model size. This can be seen in Table 9 where replacing CoPE embeddings with ALibi bias terms doesn’t bring a noticeable change in the number of parameters. Also Table 6 shows that CoPE does not increase the parameter count compared to RoPE. But we agree that we only explored two methods (interpolation and MLP) of converting fractional positions to embeddings in this work, and future work could find better methods.

What would be wrong with treating p_{ij} as an angle and converting it to a vector using a sinusoidal encoding? It seems worthwhile to compare this simpler scheme with eq. 5.

This is a good question. We can treat p_ij as an angle between vectors A and B. However, this is not very practical. A normal RoPE can be efficiently implemented because A can be a function of i, A=f(i), and B can be a function of j, B=g(j). This makes it RoPE efficient because we only need to add f and g embeddings once per batch. This is not true for CoPE because A and B cannot be expressed by i and j independently. The angle p_ij changes depending on the context and the query. This means there are no A=f(i) and B=g(j) functions exist, making it impossible to implement efficiently.

Perplexity doesn't seem like the best way to measure language model performance any more; a benchmark like MMLU would have been much more convincing.

Besides perplexity, we report accuracy metrics in many of our experiments, which is more reliable. The problem with benchmarks like MMLU in our setup is that the average context length is relatively short, thus the benefit of measuring position in a context dependent way is less beneficial. In contrast, when we measure perplexity, the model has a very long context that fills up the entire context window, making it beneficial to attend to far away tokens where abstract positions (e.g. number of paragraphs) are more useful.

In Table 5, why is only CoPE+RoPE tested? Does CoPE alone not perform well? Even if so, the results should be reported.

At this scale (1.4B), training using an efficient attention implementation becomes crucial to keep the computational resources at a manageable level. There are specialized very fast attention implementations that exist for RoPE that take advantage of hardware features. We do not have such specialized implementations for CoPE yet, thus we added CoPE to only a few layers while keeping RoPE in all layers to make the pretraining experiment possible. However, we still saw improvements with only a few layers of CoPE, which indicates the usefulness of CoPE.

Thank you for taking the time to review our paper, but we need to clarify several things:

The major concern is the novelty. The learnable vector or contextual position is based on the idea of position interpolation and fractional positional encoding explored in Insert Transformer, particularly in [1]

We respectively disagree. The cited work is significantly different from our work. Our main goal is to measure position in a context dependent way. Note that this problem is not specific to word or sentence counting, but it is well highlighted in the counting tasks we used to demonstrate the new capabilities. In [1] position only depends on token ordering and where tokens are inserted, without any context dependency. It has a completely different goal of computing position embeddings efficiently when tokens are inserted between existing tokens. They achieve this by passing the neighboring position embeddings through a neural network to obtain a new position embedding for the inserted token. This is a completely different mechanism than CoPE where we count positions by adding gated values, then convert them to vectors via interpolation. Thus we argue that the Insert Transformer method won’t be able to resolve failures demonstrated in our paper.

Experimental settings are not transparent and consistent. For example, Table 2 considers Absolute PE and RoPE as baselines, while Tables 3,4 and 5 consider other baselines. I do not see any particular consideration in changing the baseline.

We used two different codebases to make sure there is a less chance of bugs. Those codebases have different position embeddings implemented in them, thus our baselines differ in some experiments. There was no intent to be less transparent.

The authors add CoPE to the model for several layers, while baseline methods are only considered for the input. Is this a factor to affect the final results?

This is not true. All relative position encoding methods including RoPE have to be added at each layer, and that’s how they are commonly implemented. The same is true for our implementations where those baselines are added at each layer.

In Table 7, LLMs can be prompted to explain counting. It is interesting to show a similar result from CoPE. Therefore, we might understand that LLMs with CoPE can "understand" counting tasks instead of guessing a good number.

Yes, that would be an interesting experiment. But such high-level instruction following capabilities emerge only after a very large scale pre-training (not to mention the instruction tuning and RLHF stages), which are extremely compute-expensive and out of scope of this paper.

Word count and word index are not exchangeable in some contexts. It is better to make this more concise. For example, in the abstract "However, current PE methods use token counts to derive position," I suppose this should be "token index" instead of "token count"

Index usually refers to the absolute position of things. However, the current PE methods like RoPE measure word count between the source and target words, this word count is more suited terminology.

Visualizations should be consistent. The style of Figure 5 is different from others.

Sure, we can fix that.

Have you tried your pretrained model for language understanding task?

Yes, we applied our pretrained model to language understanding tasks and presented the result in Table 12.

W1: The major concern is the novelty. The learnable vector or contextual position is based on the idea of position interpolation and fractional positional encoding explored in Insert Transformer, particularly in [1].

Let me give you an example. Suppose we need to add PE to a sequence: a b c d e f g. I will do this step-by-step based on my understanding.

CoPE:

Step 1 (1st layer):

• seq: a b c d e f g
• PE : 1 1/2 2 3 4 5 6 (sigmoid-> integers , interpolation> fractional values)

Step 2 (e.g., 2nd layer):

• seq: a b c d e f g
• PE: 1 1/5 2/5 3/5 4/5 2 3 (refresh)

What insertion Transformer do:

Step 1 (1st decoding):

• seq: a d
• PE: 1 2 (relative or absolute PE)

Step 2 (2nd decoding):

• seq: a b d g
• PE: 1 2 3 4 (refresh integers based on context)

insertion Transformer + [1]:

Step 1 (1st decoding):

• seq: a d
• PE: 1 2

Step 2 (2nd decoding):

• seq: a b c d g
• PE : 1 1/3 2/3 2 3 (do interpolation if the neighboring positions exist. Note that you can also "refresh all integers based on context", similar to insertion Transformer .)

As presented in this toy example, the method significantly overlaps with the insertion Transformer family. The sigmoid value idea is also shared with relative PE. Please point out if I have substantial misunderstandings.

1. To be clear, I want to use my toy example to demonstrate that recomputing PE based on context is not a new angle. As you can see, the context is changed at each step in Insertion Transformers because of newly inserted words.
2. Thanks for clarifying. Now, the method is clearer to me, and I can raise my score.

----------------------------New Review----------------------------

Now, I have some new questions.

Based on my understanding. the p_max is an important component. However, from Table 9, we can see even a small p_max can work well. This is unconvincing or counterintuitive because a small p_max will produce the same PE for many positions. In this case, CoPE can not distinguish positions. Do you have any other experiments or insights for this? (Note that, since the method uses contextualized representations, the anisotropic problem will make the attention scores even in deep layers, meaning gate values are very similar across all positions.)