Exploring Context Window of Large Language Models via Decomposed Positional Vectors

Summary

This paper disentangles positional vectors from the hidden states of a pretrained Transformer language model to facilitate the understanding of length extrapolation. After a series of analyses, this paper proposes two context extending techniques. Experiments show that the proposed methods lower the perplexity on the task of language modeling.

Strengths

It's always good to have a mechanistic interpretability view of the hidden states of language models. The findings presented in this paper might inspire follow-up work along this direction.

Weaknesses

The experiments presented in the current draft are not convincing enough to me. See questions below.

Questions

1. Instead of continue training from the tinyllama model, I think training models from scratch using the 50B tokens budget will make the results in this paper more convincing. This is because you can get rid of the ripple effect of the originally used rope positional embeddings. Maybe your models were trying to unlearn rope during the continue training stage? 2. Apart from testing perplexity scores on the task of language modeling, I highly recommend the authors adding the experiment of needle in a haystack, otherwise I do not know if the models are really using all the tokens. 3. How do you decide the values of alpha and lambda in section 4.1 and 4.2? In addition, the temperature scaling technique was also used in several other places [1, 2] with explanations of how they did temperature selection. [1] YaRN: Efficient Context Window Extension of Large Language Models [2] Attention Alignment and Flexible Positional Embeddings Improve Transformer Length Extrapolation

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

See above.

Summary

This paper proposes a mean-based decomposition technique to analyze the formation and effect of positional encodings in LLMs. It then uses these results to propose methods to extend the context window, resulting in models that generalize better to longer texts.

Strengths

1. This paper is very well-written, and the main findings are properly highlighted. 2. This paper not only explains how positional vectors are formed, but also introduces methods to interpolate them based on the findings. 3. Experiments are performed to show that the new methods result in better perplexity scores beyond the context window.

Weaknesses

I believe this contribution is novel and insightful enough, and there is no apparent weakness.

Questions

1. The legends and graphs in Figure 4 overlap.

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have adequately addressed the limitations of their work.

Summary

This paper dives into the inner workings of how transformer-based language models handle positional information. By decomposing hidden states into semantic and positional vectors, the authors give a series of analysis about how the positional information are encoded and propagated through layers. I believe this work offers valuable insights for understanding the positional information within the transformer architecture.

Strengths

Very detailed and clear analysis about how the positional information is encoded and propagated within the transformer architecture, and to the best of my knowledge, I have not seen similar work before. I particularly enjoyed reading Figure 2 and 3, which shows how positional information is propagated through layers and goes beyond the window size, and shows how the manipulation of the positional embedding causally influence the attention patterns, particularly removing the attention sink.

Weaknesses

There are few points that I would like to suggest here to make the paper even stronger. - Section 4 feels weak and unnecessary. The performance of replacing the positional vector, if my understanding is correct, seems to be much worse than Dynamic NTK. Given the current mainstream approach is modifying the base of Rope (like YaRN), which is much easier than the approach proposed by this work, I do not think this work’s proposed context extension will be accepted by mainstream model builder. - That being said, I think the in-depth analysis of the positional embeddings are strong enough for me to give an acceptance (I learned a lot from it), so **I would strongly suggest removing the content of section 4, and use its space for more experimental analysis of the positional vectors** There are a few important problems that I believe will receive the communities’ attention and worth being addressed: - Although this paper shows the positional information can propagate through layers (Figure 2), in practice, many work found that models with window attention cannot pass the needle in a haystack test, and this is why Mistral 1.5 changed its attention back to full attention. It would be insightful if the authors can discuss the relationships between positional information and needle-in-a-haystack performance (because needle in haystack is what makes long-context models useful), i.e., why window attention cannot pass needle in haystack even it does have the correct positional information? - This paper’s analysis is restricted on TinyLLaMA, but TinyLLaMA is not a widely used open-source model, thus casting the doubt whether this discovery of this paper will hold for other model families, particularly mainstream open-weight models like LLaMA 3, Mistral, QWen or Yi. I would strongly suggest the authors verify the behavior of positional embedding on either LLaMA 3, Mistral, QWen, or Yi. Currently I’m given a borderline accept, and I will consider increasing my scores if the authors could either (1) discuss the relationship between positional vectors v.s. needle-in-a-haystack or (2) verify that the properties of positional vectors hold for LLaMA 3, Mistral, QWen or Yi (any 2 out of the 4).

Questions

see the above weakness section

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

see the above weakness section

Thank you for the rebuttal

Q1 and Q3: Thank you for addressing the concerns. We are good now. Q2: I believe there should be other methods to test whether the model is truly utilizing long-context information. I would assume TinyLLama still has a decent in-context learning ability given its size. If so, one experiment I can suggest is: Take a segment (S) of natural text with length L. Repeat it N times and concatenate them to form an artificial sequence (A) with a length of L*N. Feed A into your length-extended model and observe the output. If the model learns to use long-context information, it should output the first token of the original short sequence (S[0]). Note that S[0] is not a natural continuation of A; the model will only output S[0] if it leverages the long artificial in-context examples provided. You can then experiment with different values of L and N to demonstrate the model's ability to process long-context information. The experiment above is just one example. Feel free to devise other methods that can effectively demonstrate the long-context processing capabilities of your model.