LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models

We are glad that the reviewer found the paper “creative” and “reasonably well-written”, and the addressed problem of on-the-fly length generalization “an important challenge for LLMs”. To respond to the comments and questions one by one:

W1 & 4 & 5 & 7, Q2, 3 & 4 (Clarification of Figure 1 OOD Analysis)

Thanks for the suggestions on clarifying our OOD factor analysis. First, generally, why do LLMs fail faster than logit & entropy explosion? LLMs are usually built on very large models that can have steep local landscapes, subtle perturbations in the features space can lead LLMs to collapse. Instead, our analysis of OOD factor bounds provides a “best case” scenario where LLMs have to fail in the end (a “macro” view). In this way, dealing with these OOD factors constitutes the minimum requirement for LLM length generalization.

Figure 1a Clarification: As stated in Section 3.1 and we will emphasize more in the revision, different colors in Figure 1a are the 0-th attention head in each Transformer layer.

Factor 2 Attention Entropy: entropy introduced by too many tokens provides a potentially hazardous factor that might harm model performance. LLMs are trained on and familiar with the entropy range of fewer than 4k tokens, while the entropy level larger than that constitutes an out-of-distribution zone for them.

Factor 1c: We thank the reviewer for raising suggestions on more direct evidence of OOD factor 3 analysis. We add a plot with PCA analysis of sequence features with higher layers in Appendix G. The results suggest that, at a higher layer than 1, which was plotted in Figure 1c, the initial few tokens are implicitly encoded with a strong signal. Therefore, abandoning them might cause an OOD factor for the attention mechanism.

W2, Q5 (Pseudo-dimension Analysis)

Pseudo-dimension is a common assumption for describing the expressiveness of a function family. It holds for most common functions that humans develop, including the relative positional encoding attention functions studied in this paper, with more details elaborated in Appendix C.

W3 (fundamentally new insights in OOD factor 1)

This factor serves as a motivation to inspire our method proposal. Beyond empirical studies in prior work, here we provide a theoretical explanation of why it has to happen in various position encoding techniques, enhancing the universality of the insight.

W6, Q1 (dataset preprocessing and method hyperparameter configurations)

Thanks for the suggestion on configuration descriptions. The datasets are tokenized with conventional word tokenizers that are provided along with language models. LM-Infinite is an unparameterized model that relies on very few hyperparameters. The local branch width is selected as the pre-trained length which is the maximum inter-token distance LLMs are applied to. We find that the performance is not sensitive to the global branch width, so the range [10, 100] generally works well.

Q6 (Proposition 1 - Term B)

$B$ is the value bound for attention logits. Alongside Theorem 1, they provide an OOD dilemma for LLMs: either the logits are bounded and the attention entropy will explode, or the logits themselves will explode.

Q7 (LM-Infinite vs. Vanilla Transformer)

The difference between vanilla Transformers and those with LM-Infinite is that some long data cannot be passed to vanilla Transformers without CUDA OOM or NaN values. LM-Infinite has an $O(n)$ space complexity and can be evaluated on all data, all with valid values.

We thank the reviewer for appreciating the “fluency and generation quality over longer sequences” of the proposed method and the advantage that LM-Infinite does not need fine-tuning. Regarding the comments and questions:

W1 & 2 (long-term dependency)

We would like to emphasize that this work is the first of its kind for efficiently generating long sequences for existing LLMs without parameter updates. Considering that LLMs tend to get “lost in the Middle,” as the reviewer mentioned, long-term dependency is hard even for LLMs within their pre-training length, let alone an on-the-fly method. So, we leave investigations in this direction for future work, which is out of the scope of the current paper.

W3 (OOD factor 3)

This factor explains why conventional windowed attention will make LLMs fail, as stated in Section 3.3. The intuitive reason that initial tokens matter is that attention can be imagined as a weighted average (according to attention weights) of the value vectors $v_i$. Windowed attention, however, will reduce the number of seen initial tokens or even discard them as length increases. This will deviate the attention output to “unfamiliar” regions to pre-training scenarios.

W4 (Claim about infinite)

The reason that curves exhibit more noticeable fluctuations over 80k length is the fewer number of datum that exceeds that length, resulting in a larger variance.

Q1 (Positional embedding)

They are at positions (0, 1, 2), and the initial tokens share position 0.

Q2 (Approach design)

Thanks for raising this interesting question. We also discuss this intriguing phenomenon in Section 4.1, which reflects the tolerance of the initial tokens by the LLMs. Because they are already very far from the token being generated, their mere existence matters more as long as the number of tokens is not too large (100 << 4096).

Q3 (task solving)

Thanks for referring to this experiment that we forgot to list in the appendix. We also evaluated the task of passkey retrieval, which requires LLMs to answer a passkey buried in a long context. We follow the setting in Mohtashami et al. (https://arxiv.org/abs/2305.16300). It buries a passkey at a random position in a long distraction text and, in the end, asks what the passkey is. We update the results in Appendix F, which show that LM-Infinite allows LLMs to keep slower decaying accuracy on lengths longer than training, compared to vanilla models, which fail immediately.

Q4 (encoding and decoding)

Typically, LLMs generate from a context or prefix. Encoding refers to passing the context or prefix sequence and storing initial key-value caches, which is an expensive process so that the later decoding of each token can reuse the cached key-value sequences. As encoding the context and decoding a token has totally different computation costs, we separate the two cases for comparison.

Q5 (Transformer-XL)

This work focuses on the on-the-fly generation of long sequences for existing LLMs. As Transformer-XL is not compatible with an LLM of different architecture, and forcing it onto another model will make its generation collapse, making a fair comparison is difficult in our setting.

Q6 (Writing)

Thank you for pointing out the typo. We revised it in the PDF.

First of all, thank you very much for appreciating the “adequate theoretical analysis and experimental verification” in the paper and the “plug-and-play” nature of the proposed method. Regarding your comments and questions:

W1 (Lack of originality and significance):

In this paper, we focus on understanding and applying these techniques in a totally different direction: on-the-fly adaptation of LLMs to longer generations. In contrast to our method, most previous papers focus on training from scratch or fine-tuning longer texts for long generations. This demand for such fine-tuning will be endless if the length of inputs continues to increase. For this reason, our papers are under different settings. More importantly, our paper provides more insights into what position representation was learned during pre-training, which is of larger value for the NLP community.

W2 (Few experiments and ablation studies):

Thank you for this constructive suggestion. In the paper, we evaluate generation quality and decoding speed compared with baselines. We agree that ablation studies can give more insights into why alternative techniques might fail, so we compare with the baselines of vanilla Transformer, using windowed attention, using only a $\Lambda$-shaped attention mask, and only bounding the attention distance value. The result is added as Figure 5 in the revision. We see that all these baselines fail to maintain stable loss when length increases. This validates that it is essential to combine the two proposed components for LLM length generalization.

W3 (Dataset used):

Thank you for your interest in future exploration in this direction. We would like to point out that the current work focuses on the problem of “length representation generalization” of large language models (LLMs). On the text generation task, the proposed method demonstrates dramatic improvement by allowing LLMs to generate at 32k (much longer than pre-training length) with stable quality. To evaluate the model’s ability to solve tasks, we follow Mohtashami et al. (https://arxiv.org/abs/2305.16300) and use the passkey retrieval task. It buries a passkey at a random position in a long distraction text and, in the end, asks what the passkey is. We update the results in Appendix F, which show that LM-Infinite allows LLMs to keep slower decaying accuracy on lengths longer than training, compared to vanilla models, which fail immediately.

So far I didn't see a clear tendency to provide a rebuttal or update the work. Considering the current version, I tend to reject this work.

We are very happy to see that the reviewer finds our work “simple” while “inspired by theoretical & empirical analysis”, and the proposed method “potentially impactful”. In the following, we address your comments and questions one by one.

W1 & 3 (evaluation on downstream tasks, accuracy, and models of larger model sizes)

Our work is focused on extending the generation of language models continuously to longer lengths, which is helpful for tasks like long dialogues and long document generation. In the revised Appendix F, we evaluate the model’s ability to solve tasks by following Mohtashami et al. (https://arxiv.org/abs/2305.16300) and using the passkey retrieval task. It buries a passkey at a random position in a long distraction text, and, in the end, asks what the passkey is. Results show that LM-Infinite allows LLMs to keep slower decaying accuracy on lengths longer than training, compared to vanilla models, which fail immediately. About other model architectures, constrained by academic lab resources, we are only capable of evaluating as large as the size of Llama-2 7B and GPT-J models.

W2 (The connection between the theorems and the empirical results):

Thanks for the suggestions on a clarification of the theoretical part. Figure 1a is hard to verify due to the inaccessibility of constants. However, we are able to add a “bound” curve in Appendix 1b to visualize how the trend is bounded by the theoretical prediction in Proposition 1.

Q1 (Explaining results in Introduction)

The speedup is calculated on the Llama-2-7B model on 100 sequences of 32k length in the ArXiv dataset, with more details elaborated in Appendix D. The comparison baseline is a vanilla Llama-2-7B model. Their performance is the same comparison in Section 4.

Q2 (proof of theorem 1)

Thanks for suggesting clarifications on the proof details. The contradiction is described in the Appendix A. We elaborate it as follows: we first assert a bound on the attention logits within the range [-a, a]. Without loss of generality, we can offset all function values up to [0, 2a] to suit lemma 3 but without affecting other parts of the math. Then for the attention functions on all distances $\mathcal{H}$, lemma 3 tells us that there is a constant bound $\mathcal{N}_P(\epsilon, \mathcal{H}, \mu)$ on the number of “different enough” bins of functions. Here, “different enough” is defined as having an expected distance larger than $\epsilon$. This contrasts with our other assumption that the number of “different enough” bins of functions $\alpha(n)$ grows to let LLMs distinguish newer distance groups. Thus, if $\alpha(n)$ grows, we will need an increasing logit bound $a$ stated in the theorem. We add these explanations in the paper revision.

Q3 ($\ln(n)$ and entropy) If we understand correctly, the reviewer refers to our Proposition 1 instead of Theorem 2. Indeed, when n goes to infinity, the ln(n) also goes to infinity, which is an intuition we want to convey in the theorem. Also, the entropy of attention will indeed rise to infinity but very slowly, which accords with our observation in Figure (b). However, we add an ablation study in the paper revision in Figure 5. It shows that this factor leads to a milder OOD hazard compared with other factors, which can be attributed to the slow (but still harmful) growth of attention.