[论文解读] LM-Infinite: Zero-Shot Extreme Length Generalization for Large Language Models

Chi Han, Qifan Wang+5 more|arXiv (Cornell University)|Aug 30, 2023

Topic Modeling被引用 5

一句话总结

本文提出 LM-Infinite，一种简单、无需参数的零样本极端长度泛化方法，适用于大型语言模型。通过在推理过程中应用 $λ$-形注意力掩码和距离约束，该方法使当前最先进（SOTA）的 LLM 能够在无需微调的情况下生成长达 128k token 的流畅、连贯文本，实现 2.72 倍的解码加速，并在长上下文基准测试中达到与微调模型相当的性能。

ABSTRACT

Today's large language models (LLMs) typically train on short text segments (e.g., <4K tokens) due to the quadratic complexity of their Transformer architectures. As a result, their performance suffers drastically on inputs longer than those encountered during training, substantially limiting their applications in real-world tasks involving long contexts such as encoding scientific articles, code repositories, or long dialogues. Through theoretical analysis and empirical investigation, this work identifies three major factors contributing to this length generalization failure. Our theoretical analysis further reveals that commonly used techniques like truncating the attention window or relative positional encodings are inadequate to address them. Answering these challenges, we propose LM-Infinite, a simple and effective method for enhancing LLMs' capabilities of handling long contexts. LM-Infinite is highly flexible and can be used with most modern LLMs off-the-shelf. Without any parameter updates, it allows LLMs pre-trained with 2K or 4K-long segments to generalize to up to 200M length inputs while retaining perplexity. It also improves performance on downstream tasks such as Passkey Retrieval and Qasper in the zero-shot setting. LM-Infinite brings substantial efficiency improvements: it achieves 2.7x decoding speed up and 7.5x memory saving over the original model. Our codes are released at \url{https://github.com/Glaciohound/LM-Infinite}.

研究动机与目标

• 解决 LLM 在生成长度超过其训练长度时出现的长度泛化失败这一关键挑战。
• 识别并诊断导致长序列性能下降的分布外（OOD）因素——未见过的距离、过多被注意的 token 以及隐式编码的位置信息。
• 开发一种即插即用的解决方案，使现有预训练 LLM 能够在无需任何参数更新或微调的情况下，流畅泛化至极端长度。
• 在多种 LLM 架构和长上下文数据集上，实现高计算效率和一致的生成质量。

提出的方法

• 引入 $λ$-形注意力掩码，限制每个查询所注意的 token 数量，降低计算负载并缓解 OOD 注意力模式。
• 应用距离约束，限制相对位置编码中考虑的最大相对距离，防止推理过程中出现未见过的距离。
• 仅在推理阶段应用该方法，无需对预训练模型进行重新训练、微调或参数更新。
• 通过避免对长序列进行完整的自注意力计算，保持 $O(n)$ 的时间与空间复杂度，实现高效的解码。
• 可无缝集成至现有的相对位置编码方案（如 RoPE 和 Alibi），确保在各类 LLM 中的广泛兼容性。
• 在解码过程中动态应用该方法，无需架构修改即可处理长上下文。

Figure 1: Diagnosis of three OOD factors in LLMs.

实验结果

研究问题

• RQ1 LLM 在生成长度超过其训练长度时，导致长度泛化失败的主要分布外（OOD）因素是什么？
• RQ2 一种简单、无需参数的方法是否能在不进行微调或模型重训练的情况下，提升长上下文生成的流畅性与连贯性？
• RQ3 $λ$-形注意力掩码与距离约束在长达 128k token 的序列上，能在多大程度上保持生成质量？
• RQ4 与截断基线和微调模型相比，该方法在长上下文任务中的效率与性能表现如何？
• RQ5 该方法是否能跨多种 LLM 架构（如 LLaMA、Llama-2、MPT、GPT-J）通用，且适用于使用相对位置编码的模型？

主要发现

• LM-Infinite 使 LLaMA 和 Llama-2 在 ArXiv 的 32k token 序列上仍能保持流畅连贯的生成，BLEU 分别达到 19.7 和 18.4，与微调模型相当。
• 在 OpenWebText2 上，Llama-2 + LM-Infinite 在 16k token 时 BLEU 达到 9.6，而基线 Llama-2 模型则生成无意义文本（BLEU = 0.0）。
• 该方法在 32k 长度序列上实现了 2.72 倍的解码加速和 3.16 倍的编码加速，显著提升了推理效率。
• 与截断基线相比，LM-Infinite 在计算成本相近时实现约 5 的 BLEU 分数优势，或在仅 25% 计算量下达到相同 BLEU 分数。
• 该方法在 ArXiv 和 OpenWebText2 上均能将生成质量稳定保持至 128k token，展现出强大的零样本泛化能力。
• 该方法可普遍适用于多种 SOTA LLM（如 LLaMA、Llama-2、MPT-7B、GPT-J），且无需针对特定模型进行适配。

Figure 2: (a) \model is a plug-and-play solution for various LLMs, consisting of a $\Lambda$ -shaped mask and a distance constraint during attention. (b) We also provide a conceptual model for understanding how relative position encoding works.