大语言模型上下文长度扩展技术的内涵、动机与方法——详细综述
Saurav Pawar1 , S.M Towhidul Islam Tonmoy2 , S M Mehedi Zaman2 , Vinija Jain3,4∗,
Aman Chadha3,4∗, Amitava Das5
1 Technology Innovation Institute, UAE
2 Islamic University of Technology, Bangladesh
3 Stanford University, USA, 4 Amazon GenAI, USA
5 AI Institute, University of South Carolina, USA
arXiv:2401.07872v1 [cs.CL] 15 Jan 2024
摘要
大语言模型(Large Language Models, LLMs)的出现标志着自然语言处理(Natural Language Processing, NLP)领域的重要突破,在文本理解与生成两方面均取得了实质性进展。然而,在这些进展之中,值得注意的是,大语言模型(LLM)往往在上下文长度外推(context length extrapolation)方面存在局限。理解并扩展大语言模型(LLM)的上下文长度,对于提升其在各类 NLP 应用中的表现至关重要。在本综述中,我们从多个维度探讨为何需要上下文长度扩展,以及更优技术可能为 NLP 应用带来的变革。我们研究扩展上下文长度所固有的挑战,并对研究者已采用的现有策略给出结构化概览。此外,我们讨论评估上下文扩展技术的复杂问题,并指出该领域研究者面临的开放挑战。我们还探究研究社区在评估标准上是否已形成共识,并识别仍需进一步达成一致的方面。本综述旨在为研究者提供有价值的参考资源,帮助其把握上下文长度扩展技术的细微之处,并推动对该演进领域未来进展的讨论。
∗ 工作内容与其在 Amazon 的职位无关。
1 引言
“对我而言,上下文是关键——由此而来对一切的理解。”——Kenneth Noland
大语言模型(LLMs)的成功故事无处不在:现代大语言模型(LLM)的出现显著推进了众多自然语言处理(NLP)挑战,并达到了前所未有的高度。科学探索的自然进路是不断开拓新的、更具挑战性的前沿。在诸多雄心勃勃的举措中,一项引人注目的努力是扩展大语言模型(LLM)对极长上下文的理解能力。OpenAI 提出了约 128 页上下文理解能力的概念,而 Anthropic 近来则提出了超过 200 页的更长上下文。然而,这些商业发布与公告中明显缺乏科学严谨性。在此背景下,若干问题随之而来:(a)哪些应用需要理解如此扩展的上下文?(b)当大语言模型(LLM)能够理解更长的上下文时,我们如何有效度量应用性能的提升?(c)尽管注意力(attention)机制在 NLP 中已得到充分研究,是否仍有必要设计专门面向更长上下文的新型注意力形式?
旨在处理长上下文(long context)的先进技术之整合,有望重塑语言模型的格局。改进长上下文管理方法可提升模型性能,带来更准确、更细腻的语言理解。此类进展有望增强模型捕获长程依赖的能力,从而提升其在各类语言任务上的整体有效性,例如:
文档摘要(Document Summarization): 改进的长上下文处理有助于生成更连贯、更简洁的文档摘要,在扩展文本片段中捕捉关键信息并提升生成摘要的质量。对整篇文档的透彻理解,以及对关键词与主题的识别,都要求对广泛上下文范围加以娴熟管理。在此情境下使用较短窗口会限制生成能力,并可能导致遗漏关键细节。此外,采用更长的上下文窗口有助于缓解歧义——因为若缺乏对文档细节的充分把握,便难以利用细微信息。这反过来使大语言模型(LLM)能够以更高的辨别力与准确性推进摘要过程。
问答(Question Answering): 考量长上下文的能力可增强模型对复杂问答关系的理解,从而产生更准确、更具上下文相关性的回答。此外,大语言模型(LLM)在处理 QA 任务时表现出更强的能力,因为共指代词的消解与上下文实体紧密相关。在面对多轮对话时,扩展上下文窗口有助于在连续对话中连贯地追踪会话主题。
语言翻译(Language Translation): 在更大文本片段上更好地保持上下文,可增强模型提供准确翻译的能力,尤其是在上下文细微差别起关键作用的情形中。多义词汇在翻译领域构成重大障碍(Falkum and Vicente, 2015),而增大的上下文窗口对这类词项的语境化具有明显助益。此外,在面对专业术语时,若赋予更长的输入范围——尤其是在容纳领域特定上下文细微差别方面——大语言模型(LLM)可展现更高效能。
指代消解(Anaphora Resolution): 对长上下文的高级处理有助于在扩展文本跨度上消解对实体的指称,从而促成更准确的指代消解。指代消解过程需要在代词与其先行词之间建立联系。大语言模型(LLM)上下文窗口的扩展有助于更全面地评估信息,并通过纳入远距离指称与上下文相关细节,协助实现精确的代词消解。
对话式 AI(Conversational AI): 由长上下文模型促成的更好追踪与理解扩展对话,可使对话式 AI 系统产生更具上下文适当性的回应。扩展的上下文窗口在将幽默、讽刺或细腻表达置于会话环境中对大语言模型(LLM)至关重要。这对于生成符合当前对话预期语气与风格细微差别的回应不可或缺。
尽管研究持续推进,仍缺乏对上下文长度外推全部技术范围的全面概览。此外,大语言模型(LLM)的持续演进引入了外推上下文长度的创新面向,对现有扩展方法构成挑战,并凸显了需要充分、多样化外推途径的紧迫性。
本文是对大语言模型(LLM)上下文长度扩展技术的首篇全面综述。如图 1 所示,我们深入考察可在微调(fine-tuning)阶段实现的上下文长度外推方面的既有工作。随后,我们探讨大语言模型(LLM)上下文长度外推中潜在的未来挑战。
图 1:大语言模型(LLM)中上下文长度扩展技术的分类体系。该图将技术区分为插值(interpolation)与外推(extrapolation),并进一步划分为零样本(zero-shot)与微调(fine-tuned)分支。在该领域中,基于位置编码(positional encoding)、检索(retrieval)、注意力(attention)与 RoPE 的技术得到了最多探索。(原文图见 PDF 第 2 页)
分类体系概要(见图 1):
- 外推 Extrapolation §4
- 零样本 Zero-shot §4.1
- 位置编码 Positional encoding §4.1.1:ALiBi (Press et al., 2022);RoPE (Su et al., 2024, 2021);Randomized Positional Encoding (Ruoss et al., 2023a)
- 专用注意力机制 Specialized attention mechanism §4.1.2:xPoS (Sun et al., 2023b);LongNet (Ding et al., 2023b)
- 基于窗口的方法 Window based approaches §4.1.3:GrowLength (Jin et al., 2023)
- 微调 Fine-tuned §4.2
- 记忆/检索增强方法 Memory/Retrieval augmented approaches §4.2.1:Landmark Attention (Mohtashami and Jaggi, 2023a);Augmented Language model with long term memory (Wang et al., 2023);Think-in-Memory (Liu et al., 2023a);Focused Transformer (Tworkowski et al., 2023);MemGPT (Packer et al., 2023)
- 零样本 Zero-shot §4.1
- 插值 Interpolation §5
- 零样本 Zero-shot §5.1
- 专用注意力机制 Specialized attention mechanism §5.1.1:LM-Infinite (Han et al., 2023a);LongLoRA (Chen et al., 2023b);LongQLoRA (Yang, 2023)
- 基于提示压缩的方法 Prompt compression based approaches §5.1.2:LongLLMLingua (Jiang et al., 2023b)
- 微调 Fine-tuned §5.2
- 基于 RoPE 的方法 RoPE based approaches §5.2.1:Linear Positional Interpolation (Chen et al., 2023a);YaRN (Peng et al., 2023);PoSE (Zhu et al., 2023)
- 零样本 Zero-shot §5.1
2 当代技术
已有若干方法被提出以增强大语言模型(LLM)的上下文能力。为系统归类并提升清晰度,我们提出如图 1 所示的分类体系(taxonomy)。该分类体系划定两大主类:插值(Interpolation)与外推(Extrapolation)技术。插值涵盖将来自不同来源或上下文的信息加以融合,以精炼预测准确性;该技术适用于混合来自不同文本片段或具有不同上下文长度的不同模型的信息。相对地,外推涉及对观测数据范围之外的值进行预测,旨在将模型的理解拓展至其规定训练上下文长度之外。此外,还有零样本(Rashid et al., 2021)与微调技术用于进一步归类。分类体系中其余子节将在后续各节讨论。
3 位置技术
有别于绝对位置嵌入,相对位置嵌入基于键(keys)与查询(queries)之间的差异加以构造(Shaw et al., 2018)。相对位置嵌入的一种流行变体由 Transformer-XL 引入(Dai et al., 2019b; Yang et al., 2019)。键与查询之间注意力分数的计算被改变,以融入对应于相对位置的可训练嵌入。与绝对位置嵌入相比,配备相对位置嵌入的 Transformer 展现出泛化到超过训练中所遇长度序列的能力,体现出外推方面的优势(Press et al., 2021b)。与位置编码相关的一个反复出现的制约是:无法扩展到训练期间所观测的上下文窗口之外。已有若干工作试图克服此类局限。
旋转位置编码(Rotary Position Embedding, RoPE)(Su et al., 2021)依据每个 token 的绝对位置采用不同的旋转矩阵。它利用相对位置信息计算键与查询之间的分数,从而在近期大语言模型(LLM)(如 PaLM (Chowdhery et al., 2022) 与 LLaMA (Touvron et al., 2023a))中贡献了出色性能与长程衰减特性。
带线性偏置的注意力(Attention with Linear Biases, ALiBi)(Press et al., 2021b)与 T5 的相对偏置颇为相似,通过按键与查询之间的距离对注意力分数施加惩罚性偏置。与 T5(Raffel et al., 2020)等相对位置嵌入技术不同,ALiBi 分配预定义的惩罚分数,且不含任何可训练参数。经验结果(Press et al., 2021b)表明,ALiBi 在长于训练所遇序列上展现出更优的外推性能,超越多种流行的位置嵌入方法。此外,ALiBi 已被证明能够提升 BLOOM(Scao et al., 2022)的训练稳定性。
T5 的相对偏置最初将位于位置 $i$ 与 $j$ 的 token 的相对间隔 $i-j$ 与标量偏置值 $b=f(i-j)$ 相关联,其中函数 $f$ 对应于查找表。随后,将学得的相对偏置 $b$ 加入查询与键的点积,从而并入自注意力机制。该查找表被设计为将超过特定阈值的距离等同处理,以确保对未见距离的适应性。
位置插值(Position Interpolation, PI)由(Chen et al., 2023a)提出,是一种在预训练语言模型中扩展上下文窗口的有效方法,尤其聚焦于 LLaMA 模型。该方法的要点包括:因微调适应缓慢而产生对位置插值的动机;在预训练期间缩小位置索引的基本思想;展示稳定性的理论验证;证明效率的经验结果;涉及注意力分数修改的替代途径;以及具有稳健适应能力的微调过程。全面探索表明,位置插值在扩展上下文窗口方面行之有效,可得到在多样化语言任务上表现良好的模型。性能基准测试显示,困惑度(perplexity)有所改善,并在口令检索(passkey retrieval)与长文档摘要上取得具有竞争力的分数。结论强调,位置插值是一种以最小微调显著扩展上下文窗口的方法,可提供适用于多种应用的通用语言模型。
长度可外推 Transformer(Length-Extrapolatable Transformer, LEX Transformer)由(Sun et al., 2023b)提出,旨在应对传统 Transformer 的局限。它强调顺序敏感性、平移不变性与长度外推,并利用可外推位置嵌入(Extrapolatable Position Embedding, XPOS)实现具有注意力分辨率的通用设计。同时引入块级因果注意力(block-wise causal attention)以改进长度处理。经验评估表明,XPOS 在不同长度上的困惑度下降方面具有一致优势,而块级因果注意力增强了对更长序列的效能。实验凸显了注意力分辨率在为不同输入长度设计有效 Transformer 中的关键作用。
可外推位置嵌入(Extrapolatable Position Embedding, xPos)(Sun et al., 2022)提升了 Transformer 对平移变化的抗性及其外推上下文长度的能力。在旋转角度向量的每个维度上,xPos 引入独特的指数衰减,其幅度随旋转角度增大而减小。该特性有助于缓解训练过程中——尤其是在距离增大时——的不稳定性。
在扩展上下文长度的探索中,(Chen et al., 2023a)与(Kaiokendev, 2023)不约而同地提出一种方法:通过对 RoPE 做轻微修改以实施位置插值(PI),随后在有限数据集上微调。作为替代途径,(bloc97, 2023b)提出了“NTK-aware”插值,该方法考虑了高频信息的损失。“NTK-aware”插值方法的后续发展产生了两项显著改进,各有侧重:“Dynamic NTK”(emozilla, 2023)插值面向无需微调的预训练模型;而“NTK-by-parts”(bloc97, 2023a)插值在使用包含更长上下文信息的小规模数据集进行微调时表现最优。
YaRN(Peng et al., 2023)不同于线性插值与 NTK 插值,它通过实现斜坡函数(ramp function),在不同维度上对线性插值与 NTK 插值的组合加以变化。此外,它引入温度因子以抵消由冗长输入引起的注意力矩阵分布偏移。
GrowLength(Jin et al., 2023)提出一种在预训练阶段逐步延长训练长度的方法,从而减轻计算开销并提升整体效率。本质上,效率增益来自使用较短序列进行训练并优化资源利用。
随机化位置编码(Randomized Positional Encodings)(Ruoss et al., 2023b)对多项算法推理任务开展大规模经验评估,展示其方法相对先前途径的优越性。其做法是:通过随机选择与序列长度对齐的有序子集,将更长序列的位置纳入其中。
PoSE(Zhu et al., 2023)提出位置跳跃式(Positional Skip-wisE, PoSE)训练,巧妙地用固定上下文窗口模拟长输入。实验结果表明,相较于全长微调,PoSE 显著降低内存与时间开销,同时对性能影响极小。利用这一优势,PoSE 已成功将 LLaMA 模型扩展至 128k token,而其训练上下文窗口仅为 2k。
LongQLoRA(Yang, 2023)提出 LongQLoRA,这是一种高效且稳健的技术,用于扩展基于 RoPE 的大语言模型(LLM)的上下文长度。其在移位短注意力(shift short attention)与标准全局注意力之间的兼容性,确保可与现有推理框架无缝集成。借助 LongQLoRA,使用单块 32GB 内存的 V100 GPU,即可将 LLaMA2 7B 与 13B 等模型的上下文长度扩展至 8192 或 12k。
地标注意力(Landmark Attention)由(Mohtashami and Jaggi, 2023a)提出,是一种通过将更早的输入块直接并入注意力机制来应对 Transformer 上下文长度限制的创新方法。借助地标 token,模型在推理期间高效检索并整合先前块,从而允许处理任意上下文长度。实验结果表明计算成本与内存占用降低,展现了该方法在训练与微调大语言模型(LLM)上的有效性。该方法增强了可解释性,使信息检索过程可被清晰理解。包括英文书籍与数学论文任务在内的语言建模实验显示,困惑度有所改善,且模型能够在更长上下文中有效运行。使用地标 token 进行微调可扩展模型的上下文长度,并在口令恢复(passphrase recovery)上优于基线模型。
Think-in-Memory(Liu et al., 2023b)提出 TiM,这是一种模仿人类记忆的新型长期记忆机制,使大语言模型(LLM)能够记住并有选择地回忆想法。TiM 允许大语言模型(LLM)在记忆中思考,从而消除对长期历史进行冗余推理的需要。
4 外推
在本探索中,我们将两大总体策略加以归类并深入讨论:外推(Extrapolation)与插值(Interpolation)。外推技术旨在将模型的理解扩展到最初所观测长度之外的序列,采用创新策略以捕获更长范围上的依赖。另一方面,插值技术侧重于精炼模型在观测范围内平滑扩展其上下文理解的能力,从而提升在最初所遇上下文长度内的序列上的性能。以下各节分别阐述各类别中的技术,为应对大语言模型(LLM)中上下文长度动态特性的多样化途径提供洞见。
4.1 零样本外推
在大语言模型(LLM)领域,零样本上下文长度外推指模型固有的能力:对长度大于其原始训练所遇输入序列的内容加以理解与生成。这一独特专长无需针对更长序列进行显式微调或额外训练即可显现,展现了模型在给定任务中对扩展上下文长度的适应性。该能力在输入文本长度可变的实际应用中尤为重要。模型通过展示无需任务特定调整即可处理更广上下文范围的能力,凸显了其在面对超出训练暴露范围的上下文时仍能做出有意义预测并生成连贯文本的通用性。这一内在能力增强了模型在输入长度可能变化的多样化现实场景中的效用,有助于其在一系列上下文复杂度上处理与生成内容。
4.1.1 位置编码
位置编码在此语境中成为关键组件,为模型提供关于输入序列顺序结构的洞见。通过注入关于 token 位置的信息,这些技术在增强模型将理解外推至更长序列——而无需针对特定微调——的能力方面发挥基础性作用。本节探讨在零样本外推总括下采用的各种位置编码技术,阐明它们对模型适应更长上下文的贡献,以及对要求细腻理解序列依赖的下游任务的影响。
带线性偏置的注意力(Attention with Linear Biases, ALiBi)
尽管 RoPE 有效扩展了上下文长度,但 ALiBi(Press et al., 2021b)研究论文揭示了其在零样本上下文长度外推方面的局限。
对超出训练所经历上下文长度的考察表明,RoPE 的有效性迅速恶化。ALiBi 论文提出了一种替代技术,在其性能指标上凸显了更优的外推能力。然而,ALiBi 亦有其不足:
它使用基本线性函数来调控跨距离的注意力分数,这限制了其刻画如 RoPE 的傅里叶基所具有的复杂距离–注意力函数的能力。
此外,ALiBi 每个头仅使用单一函数,降低了其表达能力。这或许可以解释为何采用 ALiBi 的模型在诸如 MMLU(Hendrycks et al., 2020)以及评估人类偏好的 LM-Sys arena(Zheng et al., 2023)等评测上表现逊于基于 RoPE 的模型。
图 2:ALiBi(Press et al., 2021b)的实现。在神经网络中计算注意力时,该方法在对每个注意力分数应用 softmax 函数之前加入一个固定偏置。该偏置对某一特定头中的所有注意力分数相同。计算的其余部分保持不变。变量 $m$ 是每个注意力头特有的常量,在训练前设定且训练期间不调整。该方法在不同类型的文本、各种模型以及不同计算资源上均表现良好。(原文图见 PDF 第 5 页)
ALiBi 的工作原理。 对于长度为 $L$ 的输入序列,原始注意力层(Vaswani et al., 2017)在每个头中基于前 $i$ 个键 $K\in\mathbb{R}^{1\times d}$,为第 $i$ 个查询 $q_i\in\mathbb{R}^{1\times d}$(其中 $1\leq i\leq L$)计算注意力分数,其中 $d$ 表示头维度。这些分数乘以缩放因子 $1/\sqrt{d_k}$,随后经过 softmax 函数。所得注意力分数再与值向量相乘,以产生注意力层的输出。在使用 ALiBi 时,网络中任何位置都不并入位置嵌入。唯一的调整发生在查询–键点积之后:加入一个静态的、非学习的偏置 $m$。图 3 给出了示意性解释。
$$\mathrm{softmax}\bigl(q_i K^\top + m[-(i-1),\ldots,-2,-1,0]\bigr) \tag{1}$$
此处,$m$ 是头特定的斜率,在训练前预先确定。该值之所以必要,是因为查询与键之间的点积结果可能迅速增大,因而 $m$ 对其进行归一化,将其维持在 $[0,1]$ 范围内[^1]。
实验。 该研究在更大模型上探索 ALiBi 的有效性,该模型以更大计算预算与更大数据集(CC100+RoBERTa 语料)进行训练。ALiBi 展现出与正弦基线相当的稳健性能,同时使用更短的子序列并显著降低内存占用。数据集结合了 RoBERTa(Liu et al., 2019)的训练语料与 CC-100(Conneau et al., 2019)语料的英语部分(461 GB)。具有 25 个 transformer 层、16 个头、维度为 2048 的模型,在保持具有竞争力的困惑度的同时,训练速度比正弦模型快 7%,内存占用减少 1.6 GB。即便在长度仅为基线一半的序列上训练,ALiBi 仍保持更优的困惑度。进一步比较表明,ALiBi 在更长序列上具有竞争力的性能,展现了改进扩展上下文处理的潜力。该研究讨论了 ALiBi 在内存使用方面的效率,并暗示了增加更多层的可能性。
优势。 ALiBi 已在近期大语言模型(LLM)中获得广泛采用,例如 MPT-30B(Team, 2023)、Bloom(Scao et al., 2022)与 BloombergGPT(Wu et al., 2023a),原因如下:
传统位置嵌入在特定 NLP 应用中表现出若干弊端。例如,当词语与其上下文环境呈现非线性关系时,位置嵌入可能无法准确捕获这些联系。此外,位置嵌入需要额外计算开销;又因其在训练期间被学习,可能需要额外时间进行优化。相比之下,ALiBi 提供了更直接、更快速的途径,更易于实现且所需计算资源更少。而且,ALiBi 无需优化额外参数,因为头特定的标量偏置预先确定,并不经过学习。
相关工作。 与 ALiBi 研究同期,(Wennberg and Henter, 2021)中 Wennberg 等人提出了一种涉及相对定位的方法。与 ALiBi 的方法学颇为相似,其技术根据键与查询元素的邻近度向注意力分数引入偏置。不同的是,其方法整合了具有若干可训练参数的径向基函数。此外,其实验聚焦于文本分类而非语言建模,且未探索上下文长度外推。
Transformer-XL(Dai et al., 2019a)因其带有缓存机制的语言模型而受到关注:通过扩展缓存长度,将推理 token 容量扩展到训练限制之外。然而,所呈结果局限于输出长度遵守 $L$(训练长度)的情形,且其所用相对定位方法较为迟缓(Press et al., 2021a)。另一方面,Longformer(Beltagy et al., 2020)将最初在较短序列上训练的模型适配到文档级任务。然而,这种适配需要对更长序列进行部分训练。相较之下,ALiBi 方法在无需对更长序列进行额外训练的情况下即可实现外推。
旋转位置编码(Rotary Position Embedding, RoPE)
现有预训练语言模型(PLMs)中的自注意力机制源自 Transformer 架构(Vaswani et al., 2017),在运行时并不考虑位置细微差别(Yun et al., 2019)。因此,人们一直在寻求将位置信息整合到学习过程中的各种方法。一种方法是纳入源自预定函数的绝对位置编码(Vaswani et al., 2017),从而丰富上下文表示。
[^1]: ALiBi 偏置不与缩放因子相乘。
相反,另一种策略则采用可适配的绝对位置编码(positional encoding)(Gehring et al., 2017; Devlin et al., 2018; Lan et al., 2019; Clark et al., 2020; Radford et al., 2019)。另一类研究(Parikh et al., 2016; Huang et al., 2020; Shaw et al., 2018; He et al., 2020; Dai et al., 2019b; Raffel et al., 2020; Yang et al., 2019; Ke et al., 2020; Huang et al., 2018)聚焦于相对位置编码,在注意力(attention)机制中嵌入相对位置的具体信息。尽管这些方法行之有效,但它们有一共同特征:都是增强上下文表示,这与序列式自注意力的组织方式有所不同。(Su et al., 2024)提出了一种创新技术——旋转位置嵌入(Rotary Position Embedding, RoPE),以将位置信息无缝注入预训练语言模型(PLM)的学习范式。
RoPE 通过旋转矩阵捕获精确的绝对位置细节,刻画序列中各 token 彼此之间的相对位置,从而实现其功能。该过程以不同速度旋转 query 与 key 投影矩阵的片段,即便对共享同一编码的 token 也能保证旋转唯一。因此,所得点积发生变化,进而影响注意力分数。旋转带来的错位会削弱点积与注意力分数,而对齐则会提高分数。RoPE 对嵌入维度中 query 与 key 的全部 2 维切片精细管理这些旋转,从而在不同距离上建立细致的注意力分数函数。图 3 给出了示意说明。RoPE 的一个关键优势在于它仅依赖 query 与 key 之间的相对距离,无需绝对位置。这一创新方法增强了模型对 token 关系的理解,从而有助于在自注意力形式中做出更准确的预测。
实验。 该研究在多种 NLP 任务上评估 RoFormer 的表现,涵盖机器翻译、基于 BERT 的预训练、使用 GLUE 基准测试(benchmark)(Wang et al., 2018)的下游评估,以及 RoPE 与 PerFormer(Choromanski et al., 2020)线性注意力相结合的实验。所有实验均在两台配备 4× V100 GPU 的云服务器上完成。
在机器翻译任务中,使用约 450 万句对的 WMT 2014 英德数据集。对比对象为(Vaswani et al., 2017)的基于 Transformer 的基线。对该基线的修改使学习过程中能够集成 RoPE。英译德实验采用 37k 词表,并使用(Sennrich et al., 2015)的联合源端与目标端字节对编码(BPE)。评估采用(Papineni et al., 2002)的 BLEU 分数,结果一致表明 RoFormer 优于基线 Transformer。基于 PyTorch 与 fairseq 工具包(Ott et al., 2019)的实现使用 Adam 优化器、标签平滑(0.1)以及线性升温再衰减的学习率。最终指标报告为对最后 5 个检查点取平均的单一模型,并采用束搜索(束宽 4,长度惩罚 0.6)。
预训练实验在预训练阶段用 RoPE 替换 BERT 原有的正弦位置编码,语料来自 Huggingface Datasets 库中的 BookCorpus(Books, 2015)与 Wikipedia Corpus(Foundation, 2021)。语料按 8:2 划分为训练集与验证集。评估指标采用掩码语言建模(MLM)损失值,以 BERT(Devlin et al., 2018)为基线模型。实现方面,RoPE 被集成到 RoFormer 的自注意力块中。训练批次大小为 64,最大序列长度为 512,共 100k 步,优化器为 AdamW(Loshchilov and Hutter, 2017),学习率为 1e-5。结果表明,相较原始 BERT,RoFormer 在预训练阶段的 MLM 损失收敛更快。
跨多种 GLUE 任务的微调(fine-tuning)评估使用 MRPC(Dolan and Brockett, 2005)、SST-2(Socher et al., 2013)、QNLI(Rajpurkar et al., 2016)、STS-B(Al-Natsheh et al., 2017)、QQP(DataCanary et al., 2017)与 MNLI(Williams et al., 2017)等数据集,指标包括 F1 分数、Spearman 相关与准确率。实现基于 Huggingface Transformers 库,各任务微调 3 个 epoch,序列长度 512,批次大小 32,学习率取 2、3、4、5e-5。结果表明,RoFormer 在六个数据集中的三个(MRPC、STS-B、QQP)上显著优于 BERT。
在 Performer(Choromanski et al., 2020)中实现 RoPE 被证明有效,可应对与输入序列长度相关的二次计算成本。在 Enwik8(Mahoney, 2006)数据集(英文 Wikipedia)上的测试显示,在 12 层、768 维、12 头的基于字符的 Performer 中,收敛更快且损失更低。在一致设置下(例如学习率 1e-4、批次大小 128、最大序列长度 1024)对比有无 RoPE 的预训练损失曲线,凸显了带 RoPE 的 Performer 的优势。该实现在保持线性复杂度的同时提升了性能。
图 3:RoPE 的可视化(Su et al., 2024)。 该方法采用旋转矩阵以捕获 token 序列中精确的绝对位置信息。通过对 query 与 key 投影矩阵片段以不同速度旋转,RoPE 保证旋转唯一,并影响注意力分数。该图直观解释这一创新方法,强调 RoPE 依赖相对距离以改进自注意力模型中对 token 关系的理解。(原文图见 PDF 第 8 页)
优势。 RoPE 已在近期大语言模型(LLM)中得到广泛采用,如 PaLM(Chowdhery et al., 2022)、LLaMA(Touvron et al., 2023a)、LLaMA-2(Touvron et al., 2023b)、GPT-NeoX(Black et al., 2022)与 Falcon(Almazrouei et al., 2023),原因在于如下优势:
- 旋转嵌入的一个显著优势在于能够适配不同序列长度,为上下文长度(context length)外推(extrapolation)提供灵活性。与受限于特定序列长度的传统位置嵌入不同,RoPE 可调整以适应多样序列,从而成为处理不同文本长度的 NLP 模型的有价值工具。
- 随着相对距离增大,它们减弱 token 间依赖,使每个 token 对其他 token 的影响随间隔扩大而减小。这对长序列至关重要,有助于在保持准确预测的同时精简计算需求。
相关工作。 近来涌现出众多 RoPE 缩放技术,以克服 RoPE 的外推局限,并使其可应用于更长序列:
- 线性缩放 / 位置插值(Positional Interpolation) kaiokendev(Kaiokendev, 2023)与 Chen 等人(Chen et al., 2023a)各自独立提出了一种简洁而高效的上下文长度扩展(context length extension)方法。该技术将位置向量除以合适的缩放因子,使输入落在原模型的上下文窗口(context window)之内。其直觉是利用语言模型的插值(interpolation)能力,而非依赖外推。
- ReRoPE ReRoPE(Su, 2023)通过改变注意力机制来扩展上下文,因而不仅是一种嵌入插值技术。然而,它目前与 FlashAttention-2(Dao, 2023)不兼容,且在推理时需要两次注意力前向计算。
- NTK-aware RoPE 缩放 在(bloc97, 2023b)中,Reddit 用户 “bloc97” 提出了考虑高频损失的 “NTK-aware” 插值方法。随后又提出了两种增强版 “NTK-aware” 插值,各自侧重不同方面:
- Dynamic NTK:该技术(emozilla, 2023)可用于预训练语言模型,而无需微调。
- NTK-by-parts:该技术(bloc97, 2023a)在仅用有限长上下文数据微调时表现突出。
上述 NTK-aware RoPE 缩放技术已纳入开源模型,例如 Code-LLaMA(使用 NTK-by-parts 插值)(Roziere et al., 2023)与 Qwen-7B(使用 Dynamic NTK 插值)(Bai et al., 2023a)。
- 截断基(Truncated basis) 在该方法(Pal et al., 2023)中,对 RoPE 基础的改动涉及使用两个截断值与一个固定数。思路是保留重要元素,并将较不重要的元素置为 0,从而帮助模型更好地理解更长上下文。固定频率亦有助于模型区分训练数据中的不同距离。
随机化位置编码(Randomized Positional Encodings)
在其工作中,(Ruoss et al., 2023b)说明,上述局限与位置编码在更长序列上变得分布外(out-of-distribution)相关。他们提出一种新颖的位置编码方法,称为随机化位置编码方案。该方案通过随机选取与序列长度匹配的有序子集,来模拟更长序列的位置。
其所提方法在保持域内泛化性能的同时,相较直接在更长序列上训练 Transformer 的朴素做法显著提升了效率。这一新的位置编码方案族明显增强了 Transformer 的长度泛化能力,而不影响其域内泛化表现。他们在多种算法推理任务上的大规模经验评估表明,该方法优于先前做法。仅依赖顺序信息的随机化编码方案,对长度 (M)(其中 (N < M \leq L))的序列展现出显著性能提升,并允许可配置超参数 (L),其中 (N) 表示最大训练序列长度。其方法论旨在以独立于最大训练长度 (N) 的方式保留相对编码的有利性质,从而在测试时能够泛化到更长序列。具体而言,应用随机化位置编码方案时,他们按批次对扩展位置进行一次子采样,而非对每条序列分别采样。
工作原理。 当配备标准位置编码、并在最大长度为 (N) 的序列课程上训练的 Transformer,遇到长度 (M > N) 的测试序列时,会导致位置编码重新分布,偏离训练中所见情形。随着 (M) 增大,这一偏移愈发明显。
为解决该问题,作者提出仅依赖顺序信息的随机化编码方案。该方案预期可将适用性扩展到长度 (M)((N < M \leq L))的序列,并使用可配置超参数 (L)。
作者假设在每一步训练中,过程旨在最小化固定大小数据批次上的损失。他们将 (U(S)) 定义为集合 (S) 上的离散均匀分布,并将 (P_k) 定义为 ({S \subseteq {1, \ldots, L} \mid |S| = k})。
在其方法中,对每一步训练,首先随机选取长度 (n \sim U({1, \ldots, N}))(遵循 Delétang et al., 2023),再随机选取指标集 (I \sim U(P_n))。将这些指标按升序排序,形成 (I = {i_1, \ldots, i_n}),满足 (i_1 < i_2 < \ldots < i_n),并确保无重复采样。对每个 token (1 \leq j \leq N),随机化位置编码计算为 (RPE(j, \cdot) := PE(i_j, \cdot))。
在测试时,当处理长于 (N) 的序列(例如 (M > N))时,他们对所有 token 位置 (1 \leq j \leq M) 采用类似程序。该方法旨在保持相对编码的有利特性,且独立于最大训练长度 (N) 运行,从而便于在测试时处理更长序列。图 4 给出了示意说明。
优势。 相较先前方法,所提方法在长度泛化上表现更优,同时相较在更长序列上训练模型的传统做法也展现出更高的计算效率。该方法可在较短序列上训练,测试准确率超过 90%,且速度约比在更长序列上训练快 35.4 倍。值得注意的是,随机化相对编码能够解决先前方法认为困难的任务,例如反转字符串或缺失重复项等问题。
实验。 研究者在多种算法推理任务上评估该方法,包括模运算、字符串操作(例如反转/复制)、二元运算与桶排序。他们使用(Delétang et al., 2022)的基准测试进行评估,揭示了 Transformer 在此类任务上泛化的局限。
在其研究中,他们采用原始 seq-to-seq Transformer(Vaswani et al., 2017)中的仅编码器模型。对于需要多 token 输出序列的任务(例如字符串复制),他们用 (|y|) 个空 token 填充输入序列,并基于该填充序列计算整个 Transformer 输出。
模型在从 (U(1, N)) 均匀采样的序列上训练,其中 (N = 40),并在长度集合 ({N+1, \ldots, M}) 上测试,取 (M = 500)。最大位置配置为 (L = 2048)。所报告的准确率为对所有未见序列长度((N+1, \ldots, M))取平均,来自在 10 个不同参数初始化种子以及三种不同学习率 (1 \times 10^{-4})、(3 \times 10^{-4})、(5 \times 10^{-4}) 中表现最佳的模型。
图 4:随机化位置编码的实现(Ruoss et al., 2023b)。 当用更长输入序列测试模型时,通常添加位置信息的方式可能导致训练中未见的取值。该图的解决方案是:为每个训练样例分配覆盖测试时全部可能位置范围的随机(或有序)位置编码向量。(原文图见 PDF 第 10 页)
相关工作。 关于 Transformer 位置编码的研究已显著扩展。早期方法将位置信息简单加到输入序列嵌入上,例如缩放正弦(Vaswani et al., 2017)或可学习嵌入(Gehring et al., 2017)。(Dai et al., 2019a)证明了在每一层纳入 key 与 query 向量之间的相对距离,有助于增强对长期跨上下文依赖的建模。
在类似研究中,(Su et al., 2021)建议根据相对距离旋转 key–query 乘积以注入位置信息。此外,(Press et al., 2021b)通过对每个 key–query 注意力分数引入恒定偏置,增强了 NLP 任务中的长度泛化。
然而,这些方法在算法推理任务上的长度泛化仍有困难。
局限。 主要局限在于需要事先知道最大测试序列长度 (M),才能选择合适的 (L)。评估仅限于合成算法推理任务,可能缺乏真实应用的复杂性与多样性。此外,还引入了新的超参数——最大序列位置 (L)。它仅针对 Transformer 在合成数据上长度泛化的一种失效模式,而忽略了其他因素,例如注意力在更长序列上变得更不集中。
4.1.2 专用注意力机制(Specialized attention mechanism)
注意力机制是关键工具,精细编排对输入序列不同片段重要性的分配。这些技术使模型能够动态聚焦于输入中的特定区域,并适配上下文信息显著性的变化。通过对序列不同部分赋予不同程度的注意力,这些机制增强了模型辨别并捕获相关上下文的能力,这对于需要理解跨越多样且扩展上下文的依赖关系的任务至关重要。本节探讨利用注意力机制的方法,阐明模型如何有选择地权衡并优先处理输入序列中的信息,以改进上下文长度外推。
Length-Extrapolatable Transformer
Transformer 常面临一个重要局限:它们通常被设计为处理特定分布规模内的输入,因而难以针对一切可能的输入长度进行训练。为此,开发长度可外推的 Transformer 对更广泛的适用性变得必要。在序列建模中,位置信息对于构建准确表示并理解潜在含义至关重要(Hochreiter and Schmidhuber, 1997)。鉴于多种策略各自聚焦位置特征的不同方面,需要系统方法以全面指导 Transformer 设计。(Sun et al., 2023b)所提出的 Transformer 对顺序敏感,从而避免退化为混淆整体语义的词袋模型。此外,有效的位置平移——尤其是与适当的注意力掩码操作相结合——对于保持表示完整性至关重要。另外,稳健的序列模型必须能够容纳变化的输入长度,这对 Transformer 是独特挑战。该方案引入可外推位置嵌入(Extrapolatable Position Embedding, XPOS),作为 Transformer 的通用且可靠设计,并利用 ROPE 设计的优势。注意力分辨率(attention resolution)被引入作为度量,其数学形式在旋转矩阵中纳入指数衰减,以增强位置单调性度量。XPOS 保持 ROPE 的稳定性,同时在处理长期依赖方面表现一致。纳入块级因果注意力(block-wise causal attention)进一步提升注意力分辨率,改进语言建模中的长度外推表现。从头训练多种 Transformer 时,LEX Transformer 在预训练语料验证集上达到最低困惑度(perplexity),验证了所提设计的有效性。
设计架构。 在探索 Transformer 模型增强时,强调三个关键方面。第一,强调模型对顺序变化的敏感性,以高效捕获长期依赖。位置信息对有效序列建模至关重要,并与多种位置建模策略相一致。第二,引入平移不变性(translation invariance)概念,确保序列表示对位置平移具有稳健性。该性质与先前工作类似,强调相对位置相对于绝对位置的重要性。最后,讨论 Transformer 模型对长度外推能力的需求。可学习的绝对位置嵌入缺乏该能力,且多种策略在不同长度上表现显著下降。虽然 ALiBi 等方法可应对此问题,但在长期依赖处理上存在权衡。所提出的 XPOS 被呈现为 Transformer 的通用且稳健设计,纳入注意力分辨率度量与块级因果注意力,以改进长度外推表现。总体而言,系统设计考量与注意力机制——尤其是涉及相对位置的部分——对于应对 Transformer 模型中的顺序变化、平移不变性与长度外推挑战至关重要。
为增强 Transformer 的长度外推能力,注意力分辨率被认定为关键度量。所提出的 LEX Transformer 引入两种策略以最大化注意力分辨率。第一,一种专门为此设计的相对位置编码方法。第二,推理时的块级因果掩码,以进一步改进分辨率。在语言模型中表示距离的关键因素被认定为注意力分数的单调性。对不同距离上注意力分数的期望对于评估注意力分辨率至关重要。为提升注意力分辨率,还探索了窗口注意力策略。在推理阶段,为自注意力提出块级掩码,尤其是块级因果注意力。这涉及在预训练时将 query 划分为块,从而便于对更长输入编码时改进分辨率。训练采用标准注意力,而推理阶段直接纳入块级因果注意力,有助于增强位置识别,尤其是对更长序列。图 5 给出了示意说明。这些策略共同构成 Length-Extrapolatable Transformer 的框架,为改进 Transformer 模型中的长度外推提供了全面途径。
实验与结果。 在实验阶段,Transformer 从头预训练,参数规模类似中等大小的 GPT-3 模型,使用多样训练语料与 TorchScale 框架,在 16×V100 GPU 上进行。为记忆(memory)效率,预训练最大长度为 1024。在 arXiv 上进行的语言建模评估,聚焦于评估模型处理长依赖场景的能力。XPOS 在长度至 1024 时于困惑度下降上持续展现稳定优势,而块级因果注意力(BCA)的应用进一步增强了 XPOS 在长度 2048 与 4096 上的效力。通过测量分辨率的经验评估,确认了分辨率在有效 Transformer 设计中的重要性。
图 5:块级因果注意力的实现。 该方法与常规 Transformer 类似,在短文本上使用因果掩码训练。对测试时的更长序列,采用块级因果注意力,可高效复用重叠部分,如 key 与 value 向量。(Sun et al., 2022)(原文图见 PDF 第 12 页)
具体而言,对不同方法的评估——包括位置嵌入(ALiBi、ROPE 与 XPOS)在有无块级因果注意力时的表现——凸显了 BCA 的效力:它防止 ROPE 中的困惑度爆炸,并增强 XPOS 有效处理更长输入序列的能力。该实验强调了在设计注意力机制以有效处理长上下文任务时,分辨率的重要性。
相关工作。 面向扩展序列的 Transformer 需应对双重挑战:对长序列处理或记忆利用效率不足,以及效能与资源使用之间固有的权衡。线性注意力等技术(Wang et al., 2020; Katharopoulos et al., 2020; Choromanski et al., 2020)利用基于核或低维近似,优先考虑资源效率,但在典型长度场景中往往表现欠佳。稀疏注意力(Child et al., 2019b; Beltagy et al., 2020; Zaheer et al., 2020; Xiong et al., 2021)借助结构化稀疏性提供计算削减策略。此外,采用循环风格架构(Dai et al., 2019b; Hutchins et al., 2022; Ma et al., 2022)进行因果序列建模的设计,亦是应对这些挑战的竞争方案。在此背景下,重点在于解决语言建模中的外推问题——在短文本上训练、在扩展文本上评估(Press et al., 2021b)。训练方法与常规 Transformer 一致,包括在短序列上以集中注意力计算进行训练。其优势在于能够在不牺牲训练效率的前提下,无缝释放推理时对长序列建模的潜力。该方法保证在典型长度上保留最优性能,相较早期方法消除了与长序列建模相关的权衡。
LongNet
LongNet 引入扩张注意力(Dilated Attention)概念:将输入 ((Q, K, V)) 划分为长度为 (w) 的片段,记为 ((\tilde{Q}_i, \tilde{K}_i, \tilde{V}i){i=1}^{N/w})。这些片段沿序列维度以行选择间隔 (r) 进行稀疏化。注意力计算对稀疏化片段并行执行,随后经散射并拼接形成输出 (O)。通过聚集与散射操作,扩张注意力实现可无缝转化为稠密注意力,从而复用为标准注意力设计的优化,例如 flash attention(Dao et al., 2022)。该转化使计算成本相较标准注意力显著降低 (\frac{N}{wr}) 倍。实践中,片段大小 (w) 在注意力的全局性与效率之间取得平衡,而大小为 (r) 的扩张则通过近似注意力矩阵来最小化计算成本。
为高效同时捕获长程与短程信息,LongNet 使用具有多种片段大小与扩张率 ({r_i, w_i}_k) 的扩张注意力混合。由注意力 softmax 分母计算得到的动态权重优于固定可学习权重,实验表明其表现更优。扩张注意力的混合涉及并行计算,利用它们之间缺乏计算依赖这一特点。该方法对每种注意力模式逐步增大片段大小((w_i))与扩张率((r_i)),直至达到最大长度 (N) 或预定义的注意力模式数 (k),从而提供指数级的注意力感受野。图 6 给出了示意说明。
实验。 在语言建模实验中,LongNet 部署于 MAGNETO(Wang et al., 2022a)架构,并采用 xPOS 相对位置编码(Sun et al., 2022)。扩张注意力替换标准注意力机制,同时保持 MAGNETO base-size 配置,其特征为隐层维度 / 嵌入维度为 768、12 个注意力头、12 个解码器层。
图 6:扩张注意力(Dilated Attention)的实现。 LongNet(Ding et al., 2023a)这一神经网络中扩张注意力的核心组件,是用于同时捕获短程与长程依赖的注意力模式。网络可根据输入序列长度调整注意力模式的数量。(原文图见 PDF 第 13 页)
在 The Stack 数据集(Kocetkov et al., 2022)上进行预训练,该数据集汇集了超过 300 种编程语言的源代码。数据预处理采用 tiktoken 分词器与 cl100k_base 编码,模型以 0.5M token 的批次规模训练 300K 步。将 LongNet 与标准 Transformer 及稀疏 Transformer 进行比较,实验序列长度范围为 2K 至 32K。LongNet 的分段长度定义为 (w = {2048, 4096, 8192, 16384, 32768}),扩张比率记为 (r = {1, 2, 4, 6, 12})。稀疏注意力遵循固定模式(Child et al., 2019b),调整比率以匹配与 LongNet 相当的计算浮点运算量。由于计算成本更高,标准 Transformer 中的稠密注意力受限于 32K 序列长度。注意力变体源自 FlashAttention3 以保证训练效率,并为稀疏注意力与扩张注意力纳入定制的 flash attention 内核。
对于超出模型支持范围的序列,在语言模型推理中实现块级因果注意力(block-wise causal attention, BCA)(Sun et al., 2022),并移除绝对位置编码。结果表明,在训练中增大序列长度通常能提升语言模型表现。然而,当推理阶段的序列长度显著超出模型支持范围时,序列长度外推(extrapolation)会遇到局限。LongNet 持续优于基线模型,证实了其在语言建模中的有效性。
优势。 LongNet 具有显著优势,包括线性计算复杂度,以及序列中任意两个 token 之间的对数级依赖。其适用性还可扩展为超长序列的分布式训练器,并将其扩张注意力无缝集成,作为现有基于 Transformer 优化框架中标准注意力的即插即用替代。LongNet 的线性复杂度便于跨节点并行化训练,通过分布式算法克服计算与记忆(memory)约束。这种可扩展性使得在近乎恒定的运行时下,能够高效地在长达 10 亿 token 的序列上训练,显著改进了标准 Transformer 所面临的二次复杂度限制。LongNet 的线性计算复杂度被用于序列维度的分布式训练。
4.1.3 基于窗口的方法
在推进大语言模型(LLM)的领域中,一类技术通过在预训练阶段逐步延长训练长度,以凸显计算效率。这一策略性调整成为应对模型复杂度与计算成本之间持续权衡的通用方案。通过系统地增加训练长度,这些技术提升了效率,使模型能够把握更长的上下文细微差别,而不会带来过度的计算开销。下文所探讨的这一细致方法,展示了优化各类大语言模型(LLM)实用性与有效性的有前景途径。
GrowLength
大语言模型(LLM)的持续进步带来了显著进展,同时也加剧了对计算资源的需求与高昂成本。为应对这些挑战,该论文(Jin et al., 2023)提出了一种创新、简洁且高效的方法,称为“GrowLength”,以加速大语言模型(LLM)的预训练过程。所提方法在预训练阶段逐步延长训练长度,从而减轻计算开销并提升整体效率。该策略使模型能够在受限时间范围内处理更多 token,有可能提升整体性能。本质上,效率提升来自以较短序列进行训练、优化资源利用。针对多种前沿大语言模型(LLM)的大量实验表明,使用“GrowLength”方法训练的模型不仅收敛更快,而且在性能指标上优于使用现有方法训练的模型。
算法 1 GrowLength 的伪代码
1 | # loader_list : 具有不同长度的数据加载器。 |
工作原理。 GrowLength 的核心思想围绕使用较短序列加速大语言模型(LLM)的预训练展开,与较长序列相比可显著缩短训练时间。此外,从较短序列过渡到较长序列不会导致性能下降,并保持一致的退化趋势。该方法以较短序列启动预训练,并在训练过程中逐步延长序列长度,有望提供一种高效且平滑的途径。
GrowLength 将上下文窗口(context window)扩展技术加以扩展并融入预训练阶段,旨在最小化总体预训练时间,同时与现有加速方法保持兼容。算法 1 给出了 GrowLength 技术的伪代码。优势见下文。
优势。
- 该方法表明,以较短序列训练大语言模型(LLM)比以较长序列训练快得多。
- 在消耗相同 GPU 记忆的情况下,以较短序列训练可使用更大的批次规模。
- 对于较小的序列长度,模型可同时处理更多 token,从而充分利用 GPU 的全部可用记忆。
相关工作。 近期研究中,提升大语言模型(LLM)预训练效率的兴趣日益增长。诸如(Kim et al., 2023)等研究者致力于优化 CUDA 内核以减少记忆访问,从而在训练与推理速度上取得显著改进。另一些工作,如(Dao et al., 2022; Choi et al., 2022; Kwon et al., 2023),探索了流水线并行与张量并行,以在多个 GPU 间有效分配工作负载,从而增强大语言模型(LLM)推理的可扩展性。此外,由(Wu et al., 2023b; Dettmers et al., 2022; Frantar et al., 2022)所研究的量化等策略,旨在压缩大语言模型(LLM)参数,以优化整体推理效率。这些进展在应对新增大语言模型(LLM)开发中的计算成本与时间约束方面发挥关键作用,为现有方法提供了有价值的补充。在大语言模型(LLM)中的位置编码(positional encoding)方面,Transformer 架构经历了从绝对位置嵌入到更现代方法的演进。早期方法采用可学习的绝对位置嵌入(Devlin et al., 2018),提供精确的位置信息。随后引入正弦与固定位置嵌入以编码 token 的位置信息(Vaswani et al., 2017)。更近的创新,如相对位置编码,将重点转向利用 token 之间的距离信息。(Press et al., 2021b)提出了固定的线性注意力偏置,而(Su et al., 2021)引入了旋转位置嵌入(RoPE)这一新颖概念。通过 XPos(Sun et al., 2022)在外推能力方面取得了进一步进展。这一系列方法共同推动了大语言模型(LLM)背景下位置编码格局的演进。
4.1.4 记忆/检索增强方法
记忆增强架构成为一类关键范式,引入创新策略以赋予模型扩展的上下文理解能力。这些方法精巧地纳入外部记忆模块或机制,使模型能够在更广的上下文中存储与检索(retrieval)信息。通过赋予模型一种外部记忆形式,这些架构致力于增强对超出即时上下文窗口的信息的保留与利用。本节探讨记忆增强架构中的多样化方法,阐明外部记忆增强如何助力模型在理解与生成超出其训练阶段所遇长度的序列内容时提升适应性。
地标注意力(Landmark Attention)
(Mohtashami and Jaggi, 2023a)引入了一种创新方法,通过将更早的输入块直接纳入注意力机制来克服上下文长度限制。输入被划分为固定长度的块,每块以一个地标(landmark) token 标记,作为注意力的门控。该方法在保持随机访问灵活性的同时,为循环记忆方法提供了替代方案。在推理时,对地标的注意力分数允许检索并整合先前的块,从而能够处理任意上下文长度。该方法显著降低了计算成本与记忆用量。实验结果证明了其在从头训练模型或微调(fine-tuning)预训练模型方面的有效性,展示了从超过 32k token 的上下文中检索信息的能力。模型在无需额外训练的情况下进行文档检索的潜力亦得到强调。
方法。 本文聚焦于 Transformer 中的因果语言建模,其中每个 token 只能关注其之前的 token。对非因果情形的扩展亦有简要讨论。尽管处理长输入的理想情形是每个 token 关注所有先前 token,但随着输入长度增加,这在计算上变得不可行。为解决这一问题,所提方法将长输入划分为连续的 token 块,并利用注意力检索相关块。为每个块分配代表性向量,可基于注意力分数直接进行块检索。地标 token 便于训练,专门的注意力机制控制检索,提供基于语义的灵活性。文中讨论了训练细节、推理过程、位置编码与计算优势,展示了该方法在处理 Transformer 中大上下文时的效率。
在实验中,研究者着重评估在语言建模任务中检索更早块的有效性,具体聚焦于英语书籍(PG-19)与来自 arXiv 的数学论文。这些任务涉及长程 token 交互,适合评估所提方法。所用数据集分别包含 37 亿 token 的英语书籍与 56 亿 token 的 arXiv 数学论文。使用地标 token 训练的模型展示了检索相关块的能力,在降低浮点运算量的同时达到与 Transformer-XL 相当的困惑度(perplexity)。值得注意的是,该方法增强了信息检索中的可解释性,使人们能够清晰理解用于生成特定答案的被恢复文本部分。结果还表明,采用该推理机制的模型可在显著长于训练所用的上下文中有效运行。
研究者采用类似 GPT-2 的 12 层 Transformer 架构,每层 8 个注意力头,嵌入维度为 1024,前馈隐层大小为 4096。训练使用 AdamW 优化器,(\beta_1 = 0.9)、(\beta_2 = 0.95),权重衰减为 0.001,以及带预热、最小学习率为 0.0004 的余弦学习率调度器。使用 GPT-2 的分词器,在不改变批处理的情况下向数据集添加地标 token。在最多 4 块 Nvidia A100 GPU 上应用 bfloat16 混合精度训练,通过梯度累积与数据并行保持有效批次规模为 128。模型在每个数据集上训练 240K 步,上下文长度 (l_{seq} = 512)。作为对比,Transformer-XL 使用窗口大小 256(有效上下文 512),在 2048 token 的分段上训练 60K 步以观察相同数量的 token。图 7 给出了标准注意力与带地标的注意力之间的比较。
图 7:标准注意力与带地标的注意力的对比示意。 以块大小 (l_{block} = 2) 为例,展示当前 token 对先前 token 的注意力如何受到与键向量以及与各块对应地标向量相似度的共同影响。这解释了为何同一 token 尽管最初表示相同,在不同块中却可具有不同的注意力分数。地标 token 最初共享表示,但通过网络各层后以不同方式演化,从而影响注意力行为。(Mohtashami and Jaggi, 2023a)
结果。 结果评估了模型在多种推理设置下的表现,尤其是上下文长度与块检索粒度。通过将验证数据划分为称为评估长度的等长片段来评估困惑度。各片段单独送入模型,并进一步划分为块。
值得注意的是,局部上下文长度为 250 token、检索 top (k = 2) 个最相关块的设置优于 512 token,相当于关注 360 个 token(250 局部上下文、10 个地标、100 个检索得到的)。与标准的 360-token 推理相比,启用地标的检索更为有效,表明智能地恢复相关块可在保持性能的同时显著减少所关注的 token 数量。
此外,地标使模型能够在长于训练所见的上下文中有效运行,困惑度的改进表明所检索的块贡献显著,结果可与长度为 2048 的 Transformer-XL 相当。与 Transformer-XL 的循环机制不同,所提方法能够关注任意过去的 token,便于保留细粒度细节与可解释性。
在调整检索块与存储块数量时评估了性能。在长度为 2048 与 4096 的上下文中仅检索 2 个块时,模型即优于基线。仅在记忆中保留最近 40 个块可带来更好的 4096 长度性能,提示学到了类似 Transformer-XL 的循环行为。
此外,还探索了缓存块检索粒度。尽管降低灵活性会明显影响性能,模型仍优于基线。跨注意力头检索相同块(各头间有所变化)时,困惑度仅有最小幅度上升。
研究者还展示了使用地标对大语言模型进行微调以扩展其上下文长度。对 LLaMA 7B 微调 15,000 步,并通过恢复文本中隐藏的通行短语进行评估,表明即便在远长于基座模型的上下文中,也能更优地生成正确通行短语。在评估极大输入时,采用额外技术以降低记忆用量:将除地标外的 KV 缓存卸载到 CPU。
用长期记忆增强语言模型
(Wang et al., 2023)引入了长期记忆增强的语言模型(Language Models Augmented with Long-Term Memory, LONGMEM)框架,旨在解决语言模型中的记忆陈旧问题。LONGMEM 允许模型将大量先前上下文缓存到不可微的记忆库中,并采用解耦的记忆模块。该方法引入一种新颖的残差 SideNet,将先前输入向记忆的编码与检索及融合过程分离,从而有效缓解记忆陈旧与灾难性遗忘。通过在高效的记忆增强适配过程中冻结骨干大语言模型(LLM),LONGMEM 在避免计算低效的同时利用预训练知识。
LONGMEM 架构在将多样化长文与知识按下游任务纳入记忆库方面展现出通用性。在语言建模与记忆增强的上下文学习场景中的评估一致表明 LONGMEM 优于强基线。它显著提升了长上下文语言建模能力,在 ChapterBreak 等具有挑战性的基准上达到最先进性能。此外,在记忆中放入 2k 个示范样例时,LONGMEM 在自然语言理解(NLU)任务上展现出显著的上下文学习改进,凸显其在多种上下文与学习场景中增强语言模型的有效性。
设计架构。 在方法部分,作者引入 LONGMEM 框架,以增强大语言模型(LLM)从过去长上下文中获取相关信息的能力。他们提出用解耦记忆模块增强冻结的骨干大语言模型(LLM),并采用轻量残差 SideNet 以进行高效训练。该架构包含三个关键组件:冻结的骨干大语言模型(LLM)、SideNet 与缓存记忆库(Cache Memory Bank)。先前输入与当前输入由冻结的骨干大语言模型(LLM)以不同方式编码,SideNet 模块作为高效适配模型,融合当前输入上下文,并将先前上下文缓存到解耦记忆中。
在残差 SideNet 小节中,作者详述了 SideNet 的架构与初始化过程,强调通过预训练参数实现其效率。引入跨网络残差连接,将骨干大语言模型(LLM)的表示融合到 SideNet 中,确保从预训练参数进行知识迁移。在记忆检索与融合小节中,作者描述了 LONGMEM 通过用于检索与融合的记忆增强模块所实现的长期记忆能力。突出了聚焦 n-gram 结构的 token-to-chunk 记忆检索,以及在特殊记忆增强层内的记忆融合。该方法概述了一种高效且创新的途径,以应对记忆陈旧并增强大语言模型(LLM)在各类下游任务中的能力。
实验。 在实验阶段,作者评估所提 LONGMEM 模型在多种需要记忆内长上下文的任务上的表现。评估涵盖通过将过去长上下文加载到缓存记忆中进行的长文语言建模与语言理解,以及通过将大量示范样例加载到缓存记忆中实现的无限长度上下文学习。
训练设置详述了训练语料的分批过程,强调需在片段层面维持全局因果性。作者从 The Pile 中采样子集作为训练语料,复现 GPT-2(407M 参数)作为预训练骨干大语言模型(LLM),并引入 SideNet 与缓存记忆库组件。训练在 26B token 上迭代,并采用特定的记忆增强适配超参数。记忆检索细节涉及为效率构建与更新记忆检索模块、利用 token-to-chunk 检索,并引入 GPT-2 与 Memorizing Transformer(MemTRM)等基线。随后小节聚焦长上下文语言建模,凸显增强的解耦记忆在提供重要背景与上下文信息方面的益处。给出了 Project Gutenberg 2020–2022、arXiv 与 ChapterBreak 数据集的评估设置,所选指标为困惑度与后缀识别准确率。方法与评估设计提供了对 LONGMEM 能力及其在处理需广泛上下文利用的多样化任务时与相关基线比较的全面理解。
图 8:LongMem(Wang et al., 2023)架构概览。 通过向现有模型添加独立记忆模块,增强语言模型以有效利用来自过去长上下文的信息。同时引入轻量 SideNet 以高效整合记忆上下文信息。图中展示了语言建模问题与 SideNet,并刻画了为更好的语言建模而对过去记忆进行编码、存储、召回与整合的过程。
结果。 所提 LONGMEM 模型在长文语言建模领域相对所有所考虑基线均表现出显著优势。在 PG-22 数据集的多种长度划分上困惑度改进介于 −1.38 至 −1.62,在 arXiv 数据集上困惑度改进为 −1.0,展示了其利用缓存记忆中存储的过去长上下文以增强语言建模的有效性。此外,LONGMEM 在 ChapterBreakAO3 后缀识别基准上以 40.5% 准确率达到最先进性能,超越了强长上下文 Transformer 以及参数量大 313 倍的最新大语言模型(LLM) GPT-3。这些显著提升凸显了 LONGMEM 善于利用缓存记忆、并敏锐理解未来输入以完成语言建模任务的能力。
转向记忆增强的上下文学习领域,LONGMEM 扩展了大语言模型(LLM)在这方面的能力。传统上下文学习受输入上下文长度约束,限制了其从充足示范样例中吸收监督信号的有效性。LONGMEM 通过引入无限长度记忆增强来应对这一局限,使其能够通过将整个训练集加载到缓存记忆中而关注整个训练集。这一创新方法超越了传统的少样本上下文学习,实现了带有数千个辅助示范样例的记忆增强上下文学习。在 SST-2、MPQA、MR、Subj 与 SST-5 等多种自然语言理解(NLU)数据集上的评估结果表明,在 20-shot 与 4-shot 场景下均有显著改进。在 20-shot 设置中,LONGMEM 相对预训练 GPT-2* 与 MemTRM 平均分数提升 +8.0,强调其利用辅助上下文示范进行更优上下文学习的能力。此外,该模型在开放式生成任务中亦展现潜力,在 SQuAD 上 EM 分数提升 +4.5,展示了其利用缓存记忆改进上下文学习的通用性。结果证实了 LONGMEM 在长上下文建模、理解与多样本上下文学习方面的有效性与优越性,使其成为语言模型领域中一种有力方法。消融研究进一步探索了块大小与记忆大小等超参数的影响,为其对任务性能的作用提供了洞见。
相关工作。 GPT-2(Radford et al., 2019)、GPT-3(Brown et al., 2020)、OPT(Zhang et al., 2022)与 BLOOM(Workshop et al., 2022)等杰出语言模型已深刻重塑 NLP 研究,提升了语言理解、生成(Wang et al., 2022c)与视觉–语言(Wang et al., 2022b)任务上的性能基准。这些统称为大语言模型(LLM)的模型,通过扩展参数规模展现出少样本上下文学习与多步推理(Wei et al., 2022)等开创性能力。为应对处理更长上下文的挑战,已提出一类称为“x-formers”的 Transformer 模型。Transformer-XL(Dai et al., 2019b)开创了对过去片段注意力键与值的缓存机制,而 LinFormer(Wang et al., 2020)、LongFormer(Beltagy et al., 2020)与 Routing Transformer(Roy et al., 2021)等近期创新则利用稀疏注意力机制以缓解二次复杂度问题。尽管效率有所提升,这些模型在处理书籍长度序列时仍面临局限。BigBird(Zaheer et al., 2020)扩展了序列长度,但仍受限于 16k token。在任务特定调优领域,Side-Tuning 方法(Zhang et al., 2020)涉及训练一个轻量侧网络,并通过求和与预训练网络整合。相比之下,LONGMEM 引入解耦记忆以增强长期输入记忆,而无需任务特定调优。其独特的跨网络残差连接使其有别于 Side Tuning 中的传统求和方式。
4.2 微调外推
大语言模型(LLM)背景下的微调外推代表了 NLP 领域的一次精巧演进。该过程涉及专门地优化模型的既有能力,使其不仅能够理解、还能准确生成超出其初始训练数据参数范围的文本。与零样本(zero-shot)学习——模型在无需进一步调整的情况下利用其预训练知识——不同,微调外推聚焦于通过额外的、有针对性的训练来增强模型能力。这对于需要高精度生成上下文丰富且细腻文本的应用尤为关键。
经过微调(fine-tuning)后,大语言模型(LLM)变得善于处理复杂且冗长的输入,展现出适应新内容类型与结构的显著灵活性。这一增强能力确保模型能够生成更加连贯、语境恰当且精致的回应,从而大幅提升其在从高级对话界面到综合内容创作等诸多场景中的适用性。微调外推(fine-tuned extrapolation)的出现,标志着迈向更智能、更灵敏、更通用的语言模型的关键一步——这些模型能够以前所未有的精细度把握人类语言的复杂之处。
4.2.1 记忆/检索增强方法(Memory/Retrieval augmented approaches)
两种值得关注的方法——TiM(Think-in-Memory)与 Focused Transformer(FOT)——已出现,用以应对扩展大语言模型(LLM)有效上下文长度(context length)的挑战。TiM 引入动态记忆机制,通过消除重复推理并改进历史想法的组织,促进长期交互中的性能提升。另一方面,FOT 采用受对比学习启发的训练过程,在可访问外部记忆的注意力(attention)层中有效扩展(键,值)空间。FOT 通过对大规模模型进行微调展示其效力,并在需要更长上下文的任务上呈现增强性能。两种方法都对克服与有效上下文长度相关的局限作出重要贡献,为在现实应用中优化大语言模型(LLM)提供了通用解决方案。
Think-in Memory
(Liu et al., 2023b)提出了 TiM,一种模仿人类记忆的新型长期记忆机制,使大语言模型(LLM)能够记住并有选择地回忆想法。TiM 允许大语言模型(LLM)在记忆中思考,从而无需对长期历史进行冗余推理。
工作原理。 所提出的 TiM 使智能体能够进行长期对话,并在多次交互中保留有价值的历史信息。
TiM 由若干相互关联的组件构成,旨在增强延展对话中的连贯性与准确性:第一个组件是 Agent,即针对动态对话定制的预训练大语言模型(LLM)。第二个组件是 Memory Cache,一个不断扩展的哈希表,存储表示各个想法的键值对。最后,Hash-based Mapping 引入局部敏感哈希,以实现相关想法的高效存储与检索。
TiM 框架分两个阶段运作:阶段一:回忆与生成——当用户提问时,大语言模型(LLM)智能体从记忆中检索相关想法,从而在无需对原始对话文本进行冗余推理的情况下给出准确回应。阶段二:后思考与更新——在给出回应之后,大语言模型(LLM)智能体对问答对(Q-R pair)进行后思考,并将新生成的推理想法整合进记忆缓存。
优势。 作者在多轮对话数据集上进行了大量实验,揭示了大语言模型(LLM)性能在多个维度上的显著提升:
- 它能够容纳从开放域到特定域的多样化主题。
- 它支持双语,涵盖中文与英文。
- 它显著提升回应的正确性与连贯性。
实验。 三个不同的数据集——包括 KdConv(Yang et al., 2023)、Generated Virtual Dataset(GVD)(Zhong et al., 2023)以及人工整理的 Real-world Medical Dataset(RMD)——用于展示所提方法的有效性。
为突出 TiM 机制的效力,作者整合了两个稳健的大语言模型(LLM)(Zeng et al., 2022; Yang et al., 2023)。他们采用三项指标——检索准确率、回应正确性与上下文连贯性——来评估该方法。为确保评估公平,所有大语言模型(LLM)的预测结果在人工评估前均被随机化。
在 GVD 数据集(Zhong et al., 2023)的英文与中文测试集上的评估表明,该方法在所有指标上均优于 SiliconFriend(Zhong et al., 2023),尤其在上下文连贯性方面表现突出,表明 TiM 机制跨语言均有效。
在 KdConv 数据集的不同主题(电影、音乐与旅行)上测试时,该方法在所有主题上均展现出优越性能。值得注意的是,它实现了较高的检索准确率,缓解了无记忆机制的大语言模型(LLM)中观察到的较低回应正确性,并显著增强了回应中的上下文连贯性。
在 RMD 数据集上,该方法显著改善了 ChatGLM 与 Baichuan2 在长期医疗对话中的回应正确性与上下文连贯性。该方法更贴近人类记忆工作流,使大语言模型(LLM)能够生成更类人的回应。
作者引入了面向医患对话的医疗智能体 TiM-LLM,将 ChatGLM 与 TiM 相结合。TiM-LLM 通过准确回忆症状并全面理解患者疾病,协助临床医生提供精确诊断与治疗方案。
相关工作。 已有大量策略被探索以增强大语言模型(LLM)的记忆能力。记忆增强网络(Memory-augmented networks, MANNs),例如 Neural Turing Machines(NTMs)(Graves et al., 2014)以及 Meng 等人提出的其他变体(Meng and Huang, 2018),利用外部记忆缓存来处理对话中的大量上下文信息。这些 MANN 对数据进行操作与存储,通过记忆交互促进需要长期上下文的任务。
若干近期研究专门深入探讨了长期对话(Xu et al., 2021, 2022; Zhong et al., 2023; Liang et al., 2023)。例如,(Xu et al., 2021)引入了一个新的英文数据集,汇集多会话的人人众包工作者聊天,以应对长期对话流的细微之处。在并行工作中,(Zhong et al., 2023)提出了 MemoryBank 机制,其灵感来自 Ebbinghaus 遗忘曲线理论。然而,这些方法在为大语言模型(LLM)建立稳健且可适应的长期记忆框架方面面临重大挑战。它们主要关注存储原始对话文本,从而迫使大语言模型(LLM)智能体对相同的历史数据反复推理。此外,这些方法涉及计算密集的成对相似度计算以召回相关信息,在长时间交互中被证明耗时。
Focused Transformer
(Tworkowski et al., 2023)的研究指出了上下文增强中的一个主要挑战:随着文档数量增加,相关与无关 token 的比率下降,导致与无关值及相关值相关联的键之间出现重叠。这一挑战被称为分心问题(distraction issue),会阻碍模型区分它们的能力。为解决该问题,研究者提出 Focused Transformer(FOT),一种专门用于对抗分心问题的技术。FOT 使注意力层能够通过 k 近邻(k-nearest neighbors, kNN)算法访问外部(键,值)对记忆,从而有效扩展总上下文长度。训练过程受对比学习启发,在训练期间使记忆注意力层同时接触相关与无关键,增强其区分语义多样值的能力。研究者引入经 FOT 微调的 OpenLLaMA 模型,展示了其对现有模型的适用性以及在需要长上下文建模的任务上的显著改进。
FOT 被呈现为 transformer 模型的即插即用扩展,既可用于训练新模型,也可用于以更长上下文微调现有大模型。借助记忆注意力层与跨批次(crossbatch)训练过程,FOT 使模型能够在推理期间从外部记忆检索信息,从而有效扩展上下文。跨批次训练过程引导模型学习有利于记忆注意力层使用的(键,值)表示。记忆注意力层在推理期间访问外部记忆数据库,并使用 kNN 搜索算法对键进行排序与检索。该训练过程引入了一种改进(键,值)空间结构的新方法,受对比学习启发,以可微方式使注意力层接触相关与无关键。
研究讨论了分心问题,强调在标准训练中,模型没有被激励去区分来自不同文档的键,从而导致注意力质量在相关与无关文档上均匀分散。通过跨批次训练过程缓解了分心问题,带来聚焦的注意力与改进的性能。所提方法具有通用性,允许使用外部记忆而无需在训练期间依赖它。该研究仅引入极少额外超参数,并展示了 FOT 的有效性,这尤其体现在 LONGLLAMAs 在需要长上下文建模的任务上的显著改进,例如 256k 上下文长度的口令检索(passkey retrieval)任务。图 9 与图 10 分别对 Focused Transformer 在训练与推理期间进行了示意说明。
图 9:Focused Transformer(Tworkowski et al., 2023)训练期间概览。FOT 纳入记忆注意力层并采用跨批次训练方法。记忆注意力层使模型能够在推理期间从额外上下文中获取信息,从而有效扩展上下文。(原文图见 PDF 第 21 页)
图 10:Focused Transformer(Tworkowski et al., 2023)推理期间概览。在推理期间,FOT 中的记忆注意力层促进从扩展上下文中检索信息,增强模型的理解。这得益于模型在训练阶段在跨批次训练过程引导下学习到的(键,值)表示。该过程鼓励模型获取与记忆注意力层特别兼容的表示,从而优化其利用更长上下文信息的性能。(原文图见 PDF 第 21 页)
实验。 研究者展示了 FOT 对微调现有大模型的适用性,具体针对 OpenLLaMA-3B 与 OpenLLaMA-7B 模型。所得模型称为 LONGLLAMAs,展现出外推到超出其训练上下文长度、最高达 256K 的能力,同时在短上下文任务上保持性能。实验设置使用 L = 6, 12, 18(针对 3B)以及 L = 8, 16, 24(针对 7B)作为记忆层,在基于 RedPajama 的数据集混合上以 8k 上下文长度、分别用 10B(针对 3B)与 3B(针对 7B)个 token 进行 FOT 微调。值得注意的修改包括保留位置编码、使用密集注意力而非 kNN 检索,以及调整跨批次训练过程以获得更多控制。
通过口令检索任务评估 LONGLLAMAs 的有效上下文长度,展示了模型解决超出其训练上下文长度任务的能力。随后的评估聚焦于在下游任务上衡量长上下文能力,具体为 TREC 问题分类与 WebQS 问答,展示了更长上下文带来的显著准确率提升。与标准长上下文微调的比较揭示了 FOT 在准确率改进上的优越性能,尤其是在超出训练长度进行评估时。
重要的是,研究者强调,使用 FOT 为更长上下文进行微调不会损害短上下文任务上的性能,从而确保兼容性,并支持将 LONGLLAMAs 作为原始 LLaMA 模型的即插替换使用。该研究为 FOT 在扩展上下文长度并提升模型在各类任务上性能方面的有效性与通用性提供了有价值的洞见。
分析。 在本节中,研究者在较小模型上进行全面实验,以进一步审视并验证其方法。该调查回答关键问题:(1)在推理时扩展上下文长度时 FOT 的性能;(2)FOT 扩展现有预训练模型上下文长度的能力;以及(3)其处理分心问题的有效性及其对长上下文语言建模任务性能的影响。此外,还给出了消融研究与进一步分析。
实验设置涉及具有 12 层与 184M 参数的仅解码器 Transformer 模型,使用 l = 8 作为记忆注意力层,并将 kNN 检索的 top 键数调为 k = 128。区分两种评估设置:单文档与多文档。所评估的数据集包括 PG-19(英文书籍)、arXiv(数学论文)、GitHub(代码)以及 Isabelle(形式化证明)。
在将上下文长度扩展到 16M 的分析中,采用了合成字典查找任务。将 FOT 与基线 transformer 模型进行比较,展示了 FOT 在利用大记忆实现扩展上下文长度方面的有效性。
探索了 FOT 的微调能力与上下文长度外推。研究了各类数据集上的困惑度(perplexity)改进,展示了 FOT 即使在超出训练上下文长度时也能提升性能。与基线的比较揭示了稳定的困惑度收益,强调了 FOT 的优势。
本节深入讨论语言建模任务中的分心问题,尤其是在多文档设置中。使用 PG-19 数据集,研究者测量了不同多文档记忆大小下的困惑度变化。结果表明,跨批次维度取更高值会带来改进的困惑度,这与此前关于缓解分心问题的观察一致。
在单文档设置中探索了上下文长度外推,揭示 FOT 有助于外推到更长上下文,甚至超出训练上下文长度。分析引入了额外参数 w,显示随着上下文增长而改进。
消融研究聚焦于两个关键性质:可微性与负例的纳入。将可微键与值与 Memorizing Transformer 进行比较,确认了 FOT 的益处。负例的重要性得到强调,表明其在实现更好模型性能中的显著作用。
本节最后讨论了与 Memorizing Transformer 的关系,强调训练协议与记忆集成方法的影响。突出了 FOT 在记忆集成上的简洁性,以及结合训练协议的概念验证实验。该研究因其便于微调现有模型以及在训练协议方面的潜在益处而推荐 FOT 的方法。
局限与可能的补救。 当前研究不仅提供了关于 Focused Transformer 发展中挑战与进展的洞见,还识别了若干未来探索与潜在改进的方向。以下概述未来研究途径并承认现有局限:
- 记忆规模: 未来研究的一个关键方向是扩大记忆容量。克服工程挑战以存储超过 1600 万个(键,值)对,将需要实现分布式多节点系统。虽然实验使用了可扩展性有限的精确 kNN 搜索,未来工作可能探索近似 kNN 搜索方法,这需要对模型性能影响进行细致评估。
- 跨批次规模: 研究揭示了增大跨批次维度(d)的益处。当前实验采用 d = 64 或 d = 128,这是可装入单台 TPUv3/TPUv2 机器内存的最大值。未来工作旨在进一步提高 d、探索更大内存设备,或采用多节点训练设置,以增强跨批次维度的可扩展性。
- 对比学习: FOT 训练从基本对比学习(contrastive learning, CL)技术中汲取灵感,有助于改进键结构并缓解分心问题。未来研究可能深入其他 CL 方法,例如难负例挖掘(hard negative mining),以便在训练期间有效利用更大记忆。
- 与其他方法协作: 鉴于长上下文方法的动态格局,该研究认识到通过将 FOT 与其他新兴技术结合可产生协同效应。未来研究工作可探索将 FOT 与互补方法集成,促进互利互动并推进长上下文建模。
相关工作。 该调查探索了多种扩展 transformer 上下文范围的方法。例如,Transformer-XL(Dai et al., 2019b)缓存先前上下文以实现线性扩展,而 Longformer(Beltagy et al., 2020)采用稀疏注意力以促进 token 与远端对应物的交互,从而降低计算复杂度。BigBird(Zaheer et al., 2020)与 LongT5(Guo et al., 2021)等其他模型同样采用稀疏注意力处理延展序列。层次化 transformer(Nawrot et al., 2021)采用激活下采样,而 COLT5(Ainslie et al., 2023)引入条件计算以容纳更大上下文。Memorizing Transformer(Wu et al., 2022)利用 kNN 查找,旨在满足更长注意力上下文需求并增强长上下文能力。此外,该文深入探讨为大语言模型(LLM)进行扩展检索的微调,介绍了 RETRO(Borgeaud et al., 2022)与 Memorizing Transformer(Wu et al., 2022)等方法。所提方法在单阶段中扩展模型上下文,有别于以检索为中心的方法。其他研究,如 Landmark attention(Mohtashami and Jaggi, 2023b)与位置插值(Position Interpolation)(Chen et al., 2023a; Kaiokendev, 2023),聚焦于扩展 LLaMA 的上下文长度。值得注意的是,所提方法避免依赖位置编码,从而使外推到理论上无限的上下文长度成为可能。该研究还探索了零样本(zero-shot)方法,以区别于 KNN-LM(Khandelwal et al., 2019)与 Parallel Context Windows(Ratner et al., 2023)。此处的方法涉及微调模型,允许所有 token 在部分层内关注先前 token。最后,研究深入对比学习,将其与 CLIP(Radford et al., 2021)、SimCLR(Chen et al., 2020)、TRIME(Zhong et al., 2022)以及 ContraCLM(Jain et al., 2022)区分开来。所提方法将负例集成到注意力层中,专注于为扩展上下文训练注意力机制。它引入了专门为处理延长上下文而定制的、受对比学习启发的技术。
Memory-GPT (MemGPT)
工作原理。 MemGPT(Packer et al., 2023)引入一种多级记忆架构,使大语言模型(LLM)能够自主管理记忆以实现无界上下文。该架构区分两类主要记忆:主上下文(main context)与外部上下文(external context)。主上下文类似于计算机的 RAM,代表推理期间大语言模型(LLM)可用的固定上下文窗口(context window)。它由三个组件构成:提供基础大语言模型(LLM)指令的只读系统指令、存储近期对话历史的只读 FIFO 队列,以及用于临时信息的可写草稿区。这些合在一起遵守处理器的最大上下文大小。
外部上下文类似于计算机的磁盘存储,保存大语言模型(LLM)上下文窗口之外的信息。这些超上下文数据可通过显式函数调用引入主上下文。外部上下文存储可针对特定任务配置,例如为对话智能体保留完整聊天日志,或为分析保留大型文档集合。
MemGPT 的一项关键创新是大语言模型(LLM)能够自主管理其记忆。预提示提供关于记忆层次与工具的详细指令,以及用于访问或修改记忆的函数模式。在每个推理周期中,大语言模型(LLM)在执行前解析并验证包含记忆函数调用的输出字符串。这一自我导向机制由反馈环路促成,使系统能够从其行动中学习。对 token 约束的意识对于有效的自我编辑至关重要。MemGPT 以关于 token 限制的警告提示大语言模型(LLM)处理器,以引导记忆决策。MemGPT 中的控制流由事件触发,用户消息、警报、交互或定时事件可启动推理。函数链式调用允许顺序执行多个函数,增强实际任务处理能力。函数可在完成后立即返回控制权,将输出加入上下文以便继续处理而无需暂停。图 11 给出了 MemGPT 组件的概览。
图 11:MemGPT(Packer et al., 2023)的组件。在 MemGPT 中,固定上下文语言模型被增强以分层记忆系统。处理器管理其记忆,使用函数在主上下文与外部上下文之间传输数据。它通过解析器生成文本,产出结果或进行函数调用,并预先请求控制权以链式调用函数。在产出期间,处理器暂停,直至下一外部事件。(原文图见 PDF 第 23 页)
实验。 该研究在两个领域探索 MemGPT 的性能:对话智能体与文档分析。在对话智能体领域,研究扩展了 Multi-Session Chat 数据集(Xu et al., 2021),引入任务以评估智能体在长对话中的知识保留与参与度。MemGPT 在一致性与参与度标准上得到评估,展示了其利用记忆改进对话连贯性与个性化回应的能力。研究引入深度记忆检索任务,并将 MemGPT 与固定记忆基线进行比较,展示了 MemGPT 在保持连贯性方面的优越性能。此外,在对话开场任务中,MemGPT 展现出通过汲取先验知识来撰写引人入胜开场的能力。在文档分析领域,研究应对 transformer 模型中有限上下文窗口带来的挑战。MemGPT 在多文档问答任务(Liu et al., 2023c)中与固定上下文基线进行基准比较,展示了其有效扩展到更大上下文长度并处理跨文档推理的能力。研究还引入嵌套键值检索任务(Liu et al., 2023c),其中 MemGPT 通过访问存储在记忆中的键值对优于 GPT-3.5 与 GPT-4,展示了其在多跳查找方面的熟练程度。研究结果突出了 MemGPT 在对话智能体与文档分析任务中的有效性。
相关工作。 近期工作聚焦于提升大语言模型(LLM)处理更长上下文长度的能力。这一能力对需要连贯对话的对话智能体、以及对需要组合多源信息的问答任务中的大语言模型(LLM)尤为有用。诸如递归摘要(Wu et al., 2021)等方法已被探索,通过在滑动窗口上生成简洁表示来应对固定长度上下文限制。然而,该过程存在无意丢失相关细节的风险。
鉴于许多大语言模型(LLM)应用的上下文长度限制,增强大语言模型(LLM)处理更长序列的能力日益受到关注(Press et al., 2021b; Dong et al., 2023; Beltagy et al., 2020)。MemGPT 可以利用并受益于扩展的上下文长度,因为它能在记忆中存储更多信息。搜索与检索机制,尤其是在检索增强生成(Retrieval-Augmented Generation)范式内,已被集成到用于文档问答、客户支持与聊天机器人的对话智能体中。多项工作(Lin et al., 2023; Ram et al., 2023; Borgeaud et al., 2022; Karpukhin et al., 2020; Lin et al., 2023; Guu et al., 2020)分别优化了检索器或大语言模型(LLM),而 MemGPT 对具体检索方法保持无关。近期研究还聚焦于为大语言模型(LLM)增加作为交互式智能体的额外能力。例子包括为规划增加记忆(Park et al., 2023)、在 Web 环境中使用分页控制上下文大小(Park et al., 2023),以及探索交错推理(Nakano et al., 2021)。MemGPT 专门解决为智能体配备用户输入长期记忆的问题。
5 插值(Interpolation)
在上下文长度外推语境下,插值(interpolation)技术聚焦于微调或优化模型,使其有效处理落在训练期间所遇上下文长度范围内的序列。重点在于改进模型在已观察范围内平滑扩展其对上下文理解的能力,从而提升其在最初所遇上下文长度内的序列上的性能。这些技术有助于在训练限度内对上下文形成更细致且改进的理解,确保模型在训练期间所暴露的上下文长度内达到最优表现。
5.1 零样本外推(Zero-shot extrapolation)
针对插值技术的零样本外推,涉及在不对这些长度进行显式微调或优化的情况下,扩展模型处理落在训练期间所观察上下文长度之外的序列的能力。换言之,模型被期望良好地泛化到超出其训练范围的上下文长度,并依赖从已观察长度中获得的知识。
对于插值,模型通常在已观察的上下文长度内被微调或优化。在此语境下,零样本外推评估模型在没有任何针对那些长度的特定适应时,在更长序列上的表现如何。这需要评估模型对未明确包含在其训练数据中的上下文长度的零样本泛化。
5.1.1 专门化注意力机制(Specialized attention mechanism)
本节深入探讨旨在应对大语言模型(LLM)在面对更长上下文时出现的长度泛化失败的专门化注意力(attention)机制。以下论文对此有所贡献:LM-Infinite (Han et al., 2023b),提出一种 Λ 形注意力掩码与距离限制以实现即时长度泛化的方案;LongQLoRA (Yang, 2023),一种将位置插值(Position Interpolation) (Chen et al., 2023a)、QLoRA (Dettmers et al., 2023) 与移位短注意力(Shift Short Attention) (Chen et al., 2023b) 相结合、以极少训练资源扩展上下文长度的高效方法;以及 LongLoRA (Chen et al., 2023b),一种在保持与现有技术兼容的同时高效扩展上下文规模的微调(fine-tuning)方法。这些工作共同推动了专门化注意力机制的发展,以缓解大语言模型(LLM)零样本上下文长度外推(extrapolation)所面临的挑战。
LM-Infinite
工作原理。 在大语言模型(LLM)领域,提出了名为 LM-Infinite 的新方法,以解决配备相对位置编码的基于 Transformer 的大语言模型(LLM)中的长度泛化问题。LM-Infinite 提出可适用于多种大语言模型(LLM)的总体原则。LM-Infinite 由 Λ 形注意力掩码与距离限制构成。注意力掩码包含全局分支与局部分支,使 token 能够关注由因子 (n_{\text{global}}) 控制的预定数目前置 token。同时施加距离限制,将“有效距离”约束在训练长度上限((L_{\text{pretrain}}))之内。LM-Infinite 确保超出该限制的 token 在注意力中被排除,从而避免在预训练阶段接触未见过的距离。所提出的原则在三个当代开源大语言模型(LLM)系列上进行了评估:LLaMA 系列(LLaMA 与 Llama-2)、MPT-7B 系列,以及采用 RoPE、Alibi 等多种相对位置编码方法的 GPT-J 系列。对于涉及按位置旋转键向量与查询向量的 RoPE (Su et al., 2021),通过引入带有未旋转键向量与旋转查询向量的全局分支,可无缝实现 LM-Infinite。对于在 token 之间对注意力 logits 施加偏移的 Alibi (Press et al., 2021b) 编码,通过对偏移矩阵进行裁剪即可平滑集成 LM-Infinite。这一创新方案为克服大语言模型(LLM)长度泛化相关挑战提供了有前景的策略,增强了其在推理阶段对扩展上下文的适应性。图 11 给出了 LM-Infinite 与一个概念模型的概览。
实验。 评估在 Pile 数据集 (Gao et al., 2020) 的 arXiv 与 OpenWebText2 子集上对 LM-Infinite 进行了检验,前者包含 arXiv 预印本,后者包含 Reddit 帖子。流畅性通过 arXiv 数据上的困惑度(perplexity)评估,结果表明 LM-Infinite 成功将远超训练长度的曲线展平。在长序列上观察到一致的流畅性,最先进的困惑度分数证实了无需参数更新即可取得有效性。值得注意的是,MPT-7B+LM-Infinite 的分数略逊于经微调的 MPT-7B-Storywriter,展示了其作为资源高效替代方案的效率。
生成性能在 arXiv 与 OpenWebText2 上使用 BLEU (Papineni et al., 2002) 与 ROUGE (Lin, 2004) 指标进行评估。LM-Infinite 将质量扩展到超过训练的更长长度,效果类似于无需更新的微调。分析揭示了不同大语言模型(LLM)上的差异化影响:LLaMA 与 GPT-J-6B 在更远位置更好地保持质量,而 Llama-2 在较近位置表现更好。效率评估显示,在 32k 长度下编码加速 3.16 倍、解码加速 2.72 倍。一个 8k 上下文示例展示了成功的生成。
诊断大语言模型(LLM)中的分布外(OOD)问题。 作者通过理论分析与实验,探究影响大语言模型(LLM)长度泛化挑战的分布外(out-of-distribution, OOD)因素。其假设是:虽然带有相对位置编码的预训练大语言模型(LLM)能够处理相对位置,但更长的序列会使注意力权重与隐藏状态变得“陌生”,偏离训练分布。
一个关键的 OOD 因素是未见过的距离。相对位置编码依赖于注意力权重,当距离增长超出预期量级时可能出现困难。在 LLaMA 实验的支持下,作者给出一个定理,表明注意力 logits 可能随长度增加而爆炸。
第二个因素是 token 数量。随着文本变长,注意力权重熵可能增加——除非 logits 爆炸,从而在这些因素之间形成权衡。该困境在理论上与经验上均得到验证。
尽管缺乏绝对位置编码,第三个因素表明:带有相对位置编码的 Transformer 中的注意力仍可隐式编码绝对位置。一个定理与 PCA 投影表明初始 token 具有截然不同的子空间。
图 12:LM-Infinite (Han et al., 2023b) 是一种易于用于不同大语言模型(LLM)的增强方法,在注意力中涉及 Λ 形掩码与距离约束。此外,还描绘了解释相对位置编码如何工作的概念模型。(原文图见 PDF 第 26 页)
优势。 LM-Infinite 通过纳入两项独特设计提出开创性方法:Λ 形注意力掩码,以及在注意力过程中整合距离边界。这些创新元素显著提升了其吸引力:LM-Infinite 无需对预训练大语言模型(LLM)进行参数更新,并以 (O(n)) 的复杂度展现出卓越的计算效率。此外,LM-Infinite 通过切实优势体现其实用性,在编码过程中实现 3.16 倍加速,在解码效率上提升 2.72 倍。
相关工作。 Transformer (Vaswani et al., 2017) 及其变体因有效性与并行训练能力而在现代大语言模型(LLM)中得到广泛应用。位置编码对这些模型至关重要,可分为绝对位置编码(使用正弦位置嵌入或可学习位置嵌入等向量提供绝对位置)与相对位置编码(利用 token 之间的距离信息)。示例包括 T5 (Raffel et al., 2020) 与 Transformer-XL (Dai et al., 2019b) 中的可学习注意力 logit 偏置、Alibi (Press et al., 2021b) 中的线性注意力衰减,以及 RoPE (Su et al., 2021)、CAPE (Likhomanenko et al., 2021) 与 XPos (Sun et al., 2022) 等技术。
在对更长文本进行微调的背景下,现有方案包括插值位置编码 (Chen et al., 2023a)、使用对比学习 (Tworkowski et al., 2023),以及采用填充 (Tao et al., 2023) 或移位 (Kiyono et al., 2021)。然而,这些方法仅提供暂时缓解,且需要大量训练资源。本文通过识别并应对影响长度泛化的分布外(OOD)因素,提供了一种即时方案。
此外,已有多项工作致力于长上下文大语言模型(LLM)。RecurrentGPT (Zhou et al., 2023) 循环生成文本,读取近期上下文与更长历史的摘要。一些工作使用特殊标记 token (Bueno et al., 2022) 或路标(landmarks) (Mohtashami and Jaggi, 2023b) 以访问信息丰富的子集,另一些则提出提示策略 (Anil et al., 2022) 或基于检索的记忆 (Wu et al., 2022; Guu et al., 2020; Borgeaud et al., 2022; Khandelwal et al., 2019; Kaiser et al., 2017; Yogatama et al., 2021)。这些设计往往需要显式微调,且缺乏与最先进大语言模型(LLM)的兼容性。本文关注动态地将现有大语言模型(LLM)扩展到更长文本,发挥其稳健的泛化能力。
LongLoRA
在处理极长序列时,典型的自注意力机制 (Vaswani et al., 2017) 的计算开销会急剧上升,导致训练减速并增加对额外 GPU 内存的需求。标准自注意力的计算复杂度为 (O(n^2)),会带来更高的 GPU 内存成本。相比之下,移位短注意力 (Chen et al., 2023b) 将输入 token 划分为若干簇,并仅在每个簇内部独立计算注意力。为增强相邻簇之间的信息交互,它还会计算邻近簇之间的注意力。借助稀疏局部注意力机制,移位短注意力可大幅节省 GPU 内存。假设输入 token 被划分为 (g) 个簇,计算复杂度可由 (O(n^2)) 降至 (O((n/g)^2))。
LongLoRA 基于如下前提:尽管推理时需要稠密全局注意力,但微调可通过稀疏局部注意力最优地实现。LongLoRA 的关键创新在于在微调期间扩展上下文长度,在高性能与计算效率之间保持平衡。这通过增强版的低秩自适应(Low-Rank Adaptation, LoRA) (Hu et al., 2021) 方法实现——该方法是精简 Transformer 模型微调的成熟技术。LoRA 的独特之处在于在单独矩阵中训练并封装额外的权重调整,同时保持预训练模型权重的完整性。该方法简化并提升了微调过程的效率,使 LoRA 区别于其他方法。
工作原理。 LongLoRA (Chen et al., 2023b) 通过引入两个关键方面缓解计算成本挑战:
- 移位短注意力(Shift short attention, S(^2)-Attn)
- 参数高效微调
在微调阶段,S(^2)-Attn 采用稀疏局部注意力而非稠密全局注意力。本质上,这意味着将输入文档划分为不同组,并在每组内独立应用注意力机制。由于组间信息交换受限,这种分割会提高困惑度。为解决该问题,S(^2)-Attn 引入按组大小一半的 token 移位,从而促进相邻组之间的无缝信息交换。在此过程中,输出被协调合并,构成多头自注意力层的输出,并利用预训练的自注意力权重。
LongLoRA 的运行效率还通过对上下文扩展微调方法的重新审视得到进一步提升。研究表明,将通常用于注意力层的 LoRA 与允许嵌入层及归一化层在训练阶段学习相结合,具有显著效果。图 13 与图 14 对 LongLoRA 与移位短注意力的架构给出了示意说明。
实验。 实验设置包括将预训练的 7B、13B 与 70B LLaMA2 (Touvron et al., 2023b) 模型分别扩展至最大扩展上下文窗口(context window)规模达 100k、65536 与 32768,并采用位置插值 (Chen et al., 2023a)。训练参数遵循位置插值,并针对单台 8×A100 GPU 机器进行调整。训练使用 Redpajama 数据集 (Computer, 2023),评估在 PG19 (Rae et al., 2019) 与 arXiv Math proof-pile (Zhangir Azerbayev, 2022) 数据集上进行。为解决聊天能力限制,创建了 LongQA (Chen et al., 2023b) 数据集。主要结果表明,更长的上下文规模带来更低的困惑度。LongLoRA 在极大规模设置上取得有前景的结果;基于检索的主题检索任务评估显示,其性能与 LongChat-13B (Dacheng et al., 2023) 相当,并在 16k 评估中优于后者。
优势。 LongLoRA 具有以下优势:
保持原有架构结构: 经 S(^2)-Attn 微调的模型在推理时保持原有注意力架构,从而能够与既有优化技术与基础设施无缝集成。
与现有技术与工具集成: LongLoRA 在训练与推理中均可无缝集成 FlashAttention-2 (Dao, 2023) 及其他优化技术,便于纳入现有工作流。
实现简单: LongLoRA 实现直接,训练仅需极少代码,并可选用配置以在推理时保留原始标准自注意力。
图 13:LongLoRA (Chen et al., 2023b) 设计概览。移位稀疏注意力(Shifted Sparse Attention, S(^2)-Attn)在微调期间引入,而训练后的模型在推理时保持原始标准自注意力。除线性层中的 LoRA 权重外,LongLoRA 还通过使嵌入层与归一化层可训练来扩展训练。该扩展对扩大上下文至关重要,且仅引入极少额外可训练参数。(原文图见 PDF 第 28 页)
图 14:移位短注意力。包含三个步骤。特征沿头维度被划分为两个块。其中一个块中的 token 按组大小的一半移位,随后对 token 进行分组与重塑。注意力在各组内计算,组间信息通过移位流动。(Chen et al., 2023b)(原文图见 PDF 第 28 页)
相关工作。 大量研究致力于扩展 Transformer 的上下文长度。若干基于检索的方法 (Karpukhin et al., 2020; Guu et al., 2020; Izacard et al., 2022) 通过将相关文档纳入上下文来增强语言模型。本工作 (Chen et al., 2023b) 与此类方法一致,在推理时保持注意力机制不变。多种技术 (Zaheer et al., 2020; Kitaev et al., 2020; Qiu et al., 2020; Bulatov et al., 2022; Beltagy et al., 2020; Wang et al., 2020) 近似多头注意力以缓解自注意力计算的二次复杂度。值得注意的是,Longformer (Beltagy et al., 2020) 采用稀疏注意力处理扩展序列。另一些工作利用记忆机制作为对过去输入的压缩,以访问相关 token。这些技术的一个显著局限是压缩与全注意力之间存在可感知差距,从而阻碍对预训练大语言模型(LLM)的微调。尽管涉及对注意力机制的近似,本工作 (Chen et al., 2023b) 仍保持与标准注意力相当的形态与较小差距。这使得能够对预训练大语言模型(LLM)进行微调,同时在推理时保留全注意力。
LongQLoRA
在该工作中,(Yang, 2023) 提出 LongQLoRA,一种内存高效且有效的方法,用于扩展 LLaMA 系列模型的上下文长度。借助 LongQLoRA,作者将 LLaMA2 的上下文长度从 4,096 扩展到 8,192,甚至在单块 32GB 内存的 V100 上扩展到 12k。LongQLoRA 结合了位置插值、QLoRA 以及 LongLoRA 的移位短注意力的优势。
工作原理。 LongQLoRA 结合了位置插值 (Chen et al., 2023a)、QLoRA (Dettmers et al., 2023) 与 LongLoRA (Chen et al., 2023b) 的移位短注意力的优势。首先,它使用位置插值将 LLaMA2 (Touvron et al., 2023b) 的上下文长度从 4,096 扩展到目标规模。为节省更多 GPU 内存,在微调期间使用 QLoRA 将基础模型权重量化为 4 位。为进一步节省 GPU 内存,微调中还使用移位短注意力,组大小为目标上下文长度的 1/4。
为恢复不精确量化所损失的性能,它在所有层上添加 LoRA (Hu et al., 2021) 适配器,LoRA 秩为 64。研究发现,使用标准全局注意力可获得更好的推理性能。
优势。 借助单块 32GB V100 GPU,LongQLoRA 可在 1,000 步微调内将 LLaMA2 7B 与 13B 的上下文长度从 4,096 扩展到 8,192,甚至到 12k。LongQLoRA 在 PG19 (Rae et al., 2019) 与 Proof-pile (Azerbayev et al., 2022) 数据集上取得有竞争力的困惑度表现。该模型亦优于 LongLoRA,并在评估上下文长度 8,192 内与 MPT-7B-8K (Team et al., 2023) 非常接近。
实验。 研究主要在 7B 与 13B 模型上运行实验,整个实验过程使用单块 32GB 内存的 V100 GPU。他们将 LLaMA2-7B 与 Vicuna-13B 模型的上下文长度从 4096 扩展到 8192。
首先,采用位置插值技术将上下文长度从 4096 增加到 8192。关于 QLoRA,它将基础模型权重量化为 4 位 Normal Float (Dettmers et al., 2023),将 LoRA 秩设为 64,并将 LoRA 适配器集成到所有层。
在微调 LLaMA2-7B 时,他们实施下一 token 预测任务;在微调 Vicuna-13B 时,仅在目标部分计算交叉熵损失。
他们在微调中采用组大小等于模型最大上下文长度 1/4 的移位短注意力,并在推理时使用标准全局注意力。
关于数据集,作者从 Redpajama 数据集 (Computer, 2023) 中提取约 54k 条长文本样本以微调预训练模型,token 长度跨度从 4096 到 32768。此外,他们使用 PG19 (Rae et al., 2019) 验证集与 Proof-pile (Azerbayev et al., 2022) 测试集对预训练模型进行困惑度评估。
相关工作。 LLaMA 系列模型,如 LLaMA 与 LLaMA2 (Touvron et al., 2023b),以预定上下文长度训练——LLaMA 为 2,048,LLaMA2 为 4,096。其位置编码 RoPE (Su et al., 2021) 的外推能力有限。一旦输入长度超过这些预设上下文长度,模型的困惑度急剧上升,导致需要更长上下文的任务性能下降。
通过进一步预训练扩展上下文长度需要大量资源且收敛缓慢。为解决该问题,已提出位置插值(PI) (Chen et al., 2023a)、聚焦 Transformer(FOT) (Tworkowski et al., 2023) 与 LongLoRA (Chen et al., 2023b) 等技术。然而,这些方法仍需大量计算资源,往往令许多研究者难以获得。
PI (Chen et al., 2023a) 在 32 块 A100 GPU 上以 1,000 步微调 LLaMA,将上下文长度从 2,048 扩展到 8,192。FOT (Tworkowski et al., 2023) 提出在 128 块 TPU 上训练、上下文长度为 256k 的 LongLLaMA。LongLoRA (Chen et al., 2023b) 结合 PI 与 LoRA (Hu et al., 2021),在 8 块 A100 GPU 上将 LLaMA2 的上下文长度从 4,096 扩展到 100k。然而,PI 与 FOT 计算昂贵,且 LongLoRA 仍需 8 块 A100 GPU。
QLoRA (Dettmers et al., 2023) 通过先将模型量化到 4 位再添加低秩适配器,实现更高效的微调,从而降低内存需求。这使得甚至可在单块 48GB GPU 上微调 65B 参数的 LLaMA,同时匹配 16 位微调的性能。
5.1.2 基于提示压缩的方法(Prompt compression-based approaches)
提示压缩(prompt compression)技术是大语言模型(LLM)上下文长度外推领域中的关键探索方向。随着大语言模型(LLM)力求处理更长输入序列或生成更长输出,高效处理庞大提示的挑战日益凸显。提示压缩技术聚焦于从冗长提示中提炼关键信息、同时保持输入完整性与相关性的策略。这些方法旨在使大语言模型(LLM)能够有效管理扩展上下文,而不牺牲计算效率。在此背景下,以下段落概述大语言模型(LLM)中采用的各种提示压缩技术,阐明其在增强模型对不同输入长度适应性方面的作用。
LongLLMLingua
工作原理。 如 (Jiang et al., 2023a) 所阐明,LLMLingua 框架使用小型语言模型 (M_s) 评估初始提示中每个 token 的困惑度,随后消除困惑度较低的 token。该方法的依据在于:困惑度较低的 token 对语言模型总体熵增益的贡献最小,因此移除它们对大语言模型(LLM)的理解影响可忽略。LLMLingua 包含预算控制器、迭代式 token 级提示压缩算法,以及分布对齐机制。LongLLMLingua 是面向长上下文场景的扩展,旨在应对增强大语言模型(LLM)对与提示问题相关的关键信息感知的挑战。LongLLMLingua 深入四个方面:提高关键信息密度、减少中间位置的信息损失、在压缩过程中实现自适应粒度控制,以及增强关键信息的完整性。
为提高关键信息密度,LongLLMLingua 引入问题感知的粗粒度压缩与问题感知的细粒度压缩。在粗粒度压缩中,它使用度量 (r_k) 评估每篇文档的重要性,旨在保留重要性分数较高的文档。相比之下,细粒度压缩评估指令、问题与保留文档中每个 token 的重要性,使用对比困惑度表示 token 与问题之间的关联。该方法旨在确保压缩结果包含更多与问题相关的关键信息,最终提升召回率。
为应对中间位置信息损失的挑战,LongLLMLingua 根据粗粒度压缩得到的重要性分数对文档重新排序。这一策略性重排旨在优化大语言模型(LLM)在上下文不同位置上的信息感知差异。
为实现压缩过程中的自适应粒度控制,LongLLMLingua 根据粗粒度压缩的重要性分数动态分配压缩预算。这种动态分配确保更相关的文档获得更低的压缩比,从而能够基于其与提示问题的相关性对信息进行更细腻的处理。
为增强关键信息的完整性,LongLLMLingua 提出子序列恢复方法。该方法通过迭代选择出现在压缩提示中的最长子串并将其映射回原始提示,从大语言模型(LLM)的响应中恢复原始内容。该子序列恢复机制旨在纠正压缩过程中关键信息丢失可能导致的问题,确保提供给用户的信息准确可靠。图 15 给出了 LongLLMLingua 框架的概览。
实验。 在实验部分,研究评估了 LongLLMLingua 的有效性与效率。所选实验用大语言模型(LLM)为 GPT-3.5-Turbo-0613 与 LongChat-13B-16k,分别来自 OpenAI 与 HuggingFace。实现基于 PyTorch 1.13.1 与 HuggingFace Transformers,并通过贪心解码与温度设为 0 以注重稳定性与可复现性。为保证比较基础一致,压缩过程中的小型语言模型采用 LLaMA-2-7B-Chat。评估所选数据集包括 NaturalQuestions (Liu et al., 2023c)、LongBench (Bai et al., 2023c) 与 ZeroSCROLLS (Shaham et al., 2023),各自服务于评估 LongLLMLingua 性能的不同目的。
对于模拟商业搜索与问答场景中检索增强生成设置的 NaturalQuestions 数据集,准确率作为主要评估指标。LongBench 覆盖多样任务,包括单文档问答、多文档问答、摘要、少样本学习、代码补全与合成任务。评估使用该基准提供的指标与脚本进行全面检验。另一方面,ZeroSCROLLS 涵盖十个数据集上的摘要、问答、情感分类与重排任务。
为建立比较基线,采用了 BM25、Gzip (Jiang et al., 2023c)、SentenceBERT (Reimers and Gurevych, 2019)、OpenAI Embedding 与 (r_k) 度量等基于检索的方法,以及 Selective Context (Li, 2023) 与 LLMLingua (Jiang et al., 2023a) 等基于压缩的方法。LongLLMLingua 在一系列任务与压缩比上持续优于这些基线,尤其在无关信息充斥的场景中展现有效性。所提出的文档重排策略成为一项有价值的增强。
在 V100-32G GPU 上进行了专门的时延分析,聚焦平均 token 数约 10K 的 LongBench 数据集。API 调用中响应长度设为 200 token。结果表明,LongLLMLingua 不仅促进了提示压缩,还加速了整体推理过程。随着压缩率提高,加速效果更加明显,表明其在 API 成本时间较长的场景中具有潜在意义。
相关工作。 近期工作探索通过分阶段预训练 (Nijkamp et al., 2023)、修改位置嵌入 (Chen et al., 2023a; Peng et al., 2023; Han et al., 2023b)、实现线性或稀疏注意力(Ding et al., 2023a; Sun et al., 2023a),以及引入外部记忆(Bertsch et al., 2023)。然而,它们对下游任务的影响仍有待探索。
图 15: LLMLingua(Jiang et al., 2023a)提示压缩示意:引入预算控制器以实现动态压缩分配,在演示(demonstration)层面施加粗粒度压缩,详述用于知识保留的迭代提示算法,并引入对齐以弥合紧凑模型与黑盒模型之间的分布差距。(原文图见 PDF 第 31 页)
经验研究表明,大语言模型(LLM)的性能会随提示信息有效性下降而减弱(Bai et al., 2023c; Dacheng et al., 2023; Shi et al., 2023),并且依赖于相关信息的位置——理解位于中间的信息比位于两端更具挑战(Wu et al., 2023c; Liu et al., 2023c)。
检索方法包括使用潜在向量的稠密方法,如 SentenceBERT(Reimers and Gurevych, 2019),以及基于 n-gram 的稀疏方法,如 BM25。近期,(Jiang et al., 2023c)提出了一种结合压缩与 kNN 的无监督稠密方法。
提示压缩方法包括:(1)token 剪枝/合并(Goyal et al., 2020; Kim and Cho, 2020; Modarressi et al., 2022; Bolya et al., 2022);(2)软提示调优,如 GIST(Mu et al., 2023)与 AutoCompressor(Chevalier et al., 2023),需要大语言模型(LLM)微调;(3)信息熵类技术,如 LLMLingua,通过移除困惑度较低的 token 实现压缩。
5.2 微调外推(Fine-tuned extrapolation)
面向插值技术的微调外推(fine-tuned extrapolation),是指使预训练语言模型适应处理其初始训练阶段未曾见过的更长输入序列。模型首先主要在指定长度范围内的序列上训练(插值),随后再通过微调过程提升其在更长序列上的表现。这一适应过程精炼模型对扩展上下文的泛化能力,从而确保对原先观测长度与外推输入长度均能无缝处理。
5.2.1 基于 RoPE 的方法
线性位置插值(Linear positional interpolation)
基于经验证据可以观察到:在已有 Transformer 上对扩展上下文窗口进行微调的模型,其对更长上下文窗口的适应往往较为迟缓。相比其他方法,位置插值(Position Interpolation)在使特定预训练大语言模型(LLM)(如 LLaMA)实现上下文窗口扩展方面表现更优(Chen et al., 2023a)。其基本思想是放弃外推,转而直接缩小位置索引的尺度。该缩放确保最大位置索引与预训练阶段先前的上下文窗口上限对齐。更简单地说,为了容纳更多输入 token,该方法在相邻整数位置之间对位置编码进行插值。这与在已训练位置之外进行外推不同——后者可能导致灾难性的数值。对该方法的理论验证表明,插值后的注意力分数上界显著更小(在 LLaMA 7B 设定下约小 600 倍),因而更为稳定。因此,模型更容易适应插值后的位置编码。经验上,位置插值的有效性与效率得到证实:仅需短暂微调,模型即可无缝适配大幅扩展的上下文窗口。实验结果展示了通过应用位置插值,将 7B 至 65B 的 LLaMA 模型的上下文窗口从最初的 2048 扩展至 32768。
图 16 对位置插值给出了示意性说明。
实验与结果。 在其全面探索中,研究者展示了位置插值可将上下文窗口扩展至原始大小的至多 32 倍,且仅需数百个训练步即可达成。所得模型在多种语言任务上表现稳健,在长序列语言建模、口令检索(passkey retrieval)以及长文档摘要等方面表现突出。7B、13B、33B 与 65B 的 LLaMA 模型变体被扩展到至多 32768 的多种上下文窗口大小,分别采用直接微调或位置插值方法。值得注意的是,除了用位置插值对位置索引重新缩放外,LLaMA 模型架构未做其他修改。训练过程使用下一 token 预测目标进行微调,采用 AdamW、线性学习率预热,并根据模型与上下文窗口大小调整批大小。扩展后的模型在长序列语言建模上展现出更优的困惑度,且位置插值优于直接微调。此外,模型在口令检索任务中成功扩展了有效上下文窗口大小。在原始上下文窗口大小上的基准测试表明性能相当,对更长上下文窗口仅有轻微回退。在长文档摘要中,经位置插值扩展的模型取得了有竞争力的 ROUGE-1 分数,凸显了其在极少超参数调优下处理该复杂任务的有效性。
优势。 借助最少的微调,位置插值成为一种有效显著扩展 LLaMA 模型上下文窗口的手段。所得扩展模型在扩大后的上下文窗口内的一系列任务上表现熟练,同时保留了其在最初定义模型范围内任务上的原有能力。这种通用性使其成为适用于冗长与精炼输入提示的通用语言模型的可行选择。此外,经位置插值扩展的模型可复用既有基础设施与优化,从而增强该方法在各类应用中的实际吸引力。
相关工作。 该领域研究涵盖多种方法,包括检索增强大语言模型(LLM)。这涉及为 LLM 扩展检索模块,以将相关文档纳入输入上下文。该方向的代表性工作包括(Karpukhin et al., 2020; Guu et al., 2020; Izacard et al., 2022; Jiang et al., 2022; Khattab et al., 2021; Santhanam et al., 2021)。作为对这些工作的补充,该研究引入扩展上下文窗口,以在长文档摘要、少样本学习等多种任务上具备通用性。另一焦点是为 Transformer 集成记忆能力,增强其处理长序列的能力(Bulatov et al., 2022; Wu et al., 2020; Dai et al., 2019b; Wu et al., 2022; Martins et al., 2022; Mu et al., 2023)。该工作允许对所有先前 token 进行注意力计算,在不压缩的情况下保留细节。与 landmark attention(Mohtashami and Jaggi, 2023b; Chen et al., 2023a)不同,它通过未修改的注意力实现对整个输入的完全访问。该研究还与近似多头注意力(Approximated Multi-head Attention)方法兼容,通过近似与稀疏化降低多头注意力机制的内存与计算复杂度(Child et al., 2019a; Zaheer et al., 2020; Beltagy et al., 2020; Wang et al., 2020; Choromanski et al., 2020; Kitaev et al., 2020; Ren et al., 2021)。在长度外推方面,近期研究(Press et al., 2021b; Sun et al., 2022; Haviv et al., 2022)致力于在短序列上训练 Transformer,以便随后在更长序列上推理。然而,这些方法尚未在 LLaMA(Touvron et al., 2023a)或 OPT(Zhang et al., 2022)等主流语言模型中落地,限制了其对成熟预训练模型的适用性。(Chen et al., 2023a)通过扩展既有大语言模型(LLM)填补了这一空白,提供了一种兼顾成本效益的长度外推方案,同时保持原模型质量,即便在上下文窗口适中的任务中亦然。(Dosovitskiy et al., 2020)在 Vision Transformer 中引入对学习位置嵌入的线性插值,以在微调时支持更高分辨率。相比之下,当前工作对位置索引进行插值,专为类 RoPE 位置编码设计,可能需要更少训练且无需额外可训练参数。成功将上下文窗口扩展至 32 倍,(Chen et al., 2023a)超越了(Dosovitskiy et al., 2020)对至多 4 倍扩展的探索,证实了位置插值在扩展语言模型上下文窗口方面的有效性。
图 16: 位置插值(Chen et al., 2023a)方法示意。对于预训练上下文窗口为 2048 个位置的 Llama 模型,左上为标准用法。在长度外推(右上)中,模型处理未见过的位置(红点)直至 4096。位置插值(左下)将位置索引(蓝点与绿点)从 [0, 4096] 缩小到 [0, 2048],使其保持在预训练范围内。
Yet another RoPE extensioN(YaRN)
该论文(Peng et al., 2023)针对位置编码的一个长期局限:无法泛化到训练时所见上下文窗口之外。尽管 ALiBi 等方法表现出有限泛化能力,尚无方法能将序列显著扩展到远超其预训练长度。先前工作提出了位置插值(Chen et al., 2023a)与“NTK-aware”插值(bloc97, 2023b)等方案,并应用于 Code Llama(Roziere et al., 2023)与 Qwen 7B(Bai et al., 2023b)等开源模型。本文提出 YaRN——一种针对 LLaMA、GPT NeoX 与 PaLM 等模型中 RoPE 的增强扩展方法。YaRN 在扩展上下文窗口方面达到最先进性能,微调所需数据不足原始预训练数据的 0.1%。此外,Dynamic-YaRN 结合 Dynamic Scaling 推理时技术,无需微调即可实现超过 2 倍的上下文窗口扩展。
YaRN 的结构。 该论文研究了位置插值在刻画 RoPE 与大语言模型(LLM)动力学之间复杂关系方面的不足。“NTK-aware”插值被引入以应对高频细节损失,但外推效果仍非最优。
为克服这些局限,论文提出“NTK-by-parts”插值,针对特定 RoPE 维度以提升性能。它还探索了动态缩放,特别是“Dynamic NTK”插值,能够优雅处理重复的前向传播,并在未微调模型上展现有效性。
在 YaRN 方法中,提出通过对 logits 进行温度重参数化来修改注意力机制,而无需改动代码。YaRN 在上下文扩展方面高效,训练与推理阶段均零额外开销。YaRN 与“NTK-by-parts”插值相结合可进一步增强所提方法。
在训练阶段,模型在不改变架构的情况下被扩展。学习过程涉及以不同取值重新计算某些量。训练分步进行,调整学习率等参数并采用特定技术。该过程分别针对 7B 与 13B 两种模型规模执行,并应用特定策略以优化其性能。
在外推与迁移学习阶段,模型被评估其将所学应用于新场景的能力。这涉及使用具有特定上下文长度的数据集,并对模型进行微调以增强能力。评估包括探索模型对新的、更长上下文的适应程度——这些上下文超出了原始训练所见范围。
结果表明,更大的模型成功适应了扩展的上下文长度,展现出从不同尺度学习并应用知识的能力。这一适应高效完成,表明模型能将其所学信息迁移到新上下文,而无需大量重新学习。
结果。 评估首先考察模型在不断扩大的上下文窗口上的表现。为此,从 Proof-pile 中选取十个随机样本,每个至少包含 128k 个 token。对这些样本在不同序列长度(从 2k 到 128k 个 token)上评估困惑度,并以 2k 为步长截断。值得注意的是,PI 与“NTK-aware”模型的训练方法遵循(Chen et al., 2023a),而 YaRN 采用类似方案,但训练步数与数据量约为其 2.5 分之一。
结果表明,在目标上下文大小范围内性能稳健,YaRN 插值尤为显著地将 Llama-2 的有效上下文大小扩展至 128k。特别值得注意的是 YaRN(s=32)模型:尽管微调仅限于 64k token 长度,其困惑度在直至 128k 的范围内持续下降。这证明了模型对未见上下文长度的泛化能力。
口令检索任务由(Mohtashami and Jaggi, 2023a)定义,用于评估模型从大量其他无意义文本中检索简单口令的能力。评估对口令检索任务进行十次迭代,口令随机置于不同大小评估上下文窗口(从 8k 到 128k)中均匀分布的位置。经 YaRN 在 128k 上下文大小上微调的 7b 与 13b 模型,在整个上下文窗口大小范围内均展现出非常高的准确率(>99%)。
在 YaRN 模型与各自 Llama-2 基线之间仅观察到极小的性能下降。平均而言,YaRN(s=16)与 YaRN(s=32)模型之间的分数仅下降 0.49%。这表明从 64k 迭代扩展至 128k 所带来的性能损失可忽略不计。
总之,YaRN 改进了现有 RoPE 插值方法,可无缝替代位置插值,无明显缺点且实现代价极低。微调后的模型在多个基准测试上保持原有能力,同时有效关注显著更大的上下文规模。此外,YaRN 通过对更短数据集的微调实现高效外推,并借助迁移学习加快收敛,从而应对计算受限场景中的关键要素。
相关工作。 ReRoPE(Su, 2023)旨在扩展基于 RoPE 的预训练模型的上下文规模,声称无需微调即可获得“无限”上下文长度,并以 Llama 2 13B 模型上直至 16k 单调递减的损失为支撑。与嵌入插值方法不同,ReRoPE 通过修改注意力机制实现上下文扩展。然而,它目前与 Flash Attention 2(Dao, 2023)不兼容,且推理时需要两次注意力计算。同期提出的 LM-Infinite 与 YaRN 有相似理念,但强调对未微调模型的“即时(on-the-fly)”长度泛化。与 ReRoPE 类似,LM-Infinite 修改注意力机制,因而与 Flash Attention 2 不兼容,也不是直接的嵌入插值方法。
Positional Skip-wisE(PoSE)
(Zhu et al., 2023)提出了一种称为 Positional Skip-wisE(PoSE)训练的方法,以有效扩展 Transformer 的上下文窗口。他们将固定长度的上下文窗口切分为多个块(chunk),并为每个块引入专用的位置偏置项。通过按样本变化偏置值与块长度,模型学习到覆盖比实际上下文规模更长范围的序列位置表示。由此,PoSE 训练可在仅使用固定窗口的情况下模拟更大的上下文窗口,使模型积累长程上下文知识。
工作原理。 PoSE 技术通过在固定上下文窗口内操纵位置索引来模拟更长输入。它将原始窗口划分为多个块,并通过为每个块添加不同的跳跃偏置项来调整位置索引。这些偏置项以及块长度对每个训练样本均有变化,使模型在微调期间适应目标上下文窗口内的所有位置。通过在每个块内保持连续的位置索引,PoSE 与预训练高度相似,从而保留模型的语言建模与理解能力。
举例说明,该方法将原始上下文窗口 (L_c) 划分为 (N) 个块 (c_0, c_1, \ldots, c_{N-1}),长度分别为 (l_0, l_1, \ldots, l_{N-1}),其中 (\sum_{i=0}^{N-1} l_i = L_c)。它为每个块 (c_i) 引入起始索引 (st_i),从而将位置索引表述如下:
$$\mathrm{Pos}(c_i) = {st_i, st_i + 1, \ldots, st_i + l_i - 1}, \quad (2)$$
其中,
$$st_i = \sum_{j=0}^{i-1} l_j \quad (3)$$
随后,它使用离散均匀分布 (U(S)) 为每个块 (c_i) 采样跳跃偏置项 (u_i \sim U({u_{i-1}, \ldots, L_t - L_c}))。
该偏置项将原始位置索引变换为:
$$\mathrm{PoSE}(c_i) = {u_i + st_i, u_i + st_i + 1, \ldots, u_i + st_i + l_i - 1} \quad (4)$$
注意约束 (u_i \geq u_{i-1}) 可防止块之间的位置索引重叠。
对于每个块内的文本,采用类似过程从输入文本 (x = {x_0, x_1, \ldots, x_{L_x}}) 中选取连续的 token 跨度。
在确定每个块的位置索引与内容后,再进行位置插值以稳定微调。
图 17 对全长微调与 PoSE 给出了示意性对比。
优势。 PoSE 的优势有三方面:
内存与时间效率: 由于微调仅需原始上下文规模,PoSE 避免了微调阶段计算复杂度相对目标长度的二次增长。这显著降低了内存与时间开销。
极长上下文的潜力: PoSE 架构成功将 LLaMA(Touvron et al., 2023a)的上下文窗口从原始的 2,048 个 token 扩展至 131,072 个 token——实现 64 倍增益。尽管上下文长度大幅扩展,PoSE 仍保持了 LLaMA 的语言建模与理解能力。
与基于 RoPE 的大语言模型(LLM)及 PI 策略的兼容性: PoSE 的有效性已在若干代表性基于 RoPE 的大语言模型(LLM)上经验验证,包括 LLaMA、LLaMA2(Touvron et al., 2023a)、GPT-J(Wang and Komatsuzaki, 2021)以及 Baichuan(Yang et al., 2023)。此外,PoSE 已展示与多种位置插值方法的兼容性,包括 Linear(Chen et al., 2023a)、NTK(bloc97, 2023b)与 YaRN(Peng et al., 2023)插值。
实验。 作者在两项主要任务上考察长文本建模效能:语言建模与口令检索。语言建模作为衡量模型处理长篇文本总体能力的基本指标,而口令检索则衡量推理阶段所考虑的最大 token 距离。语言建模在 GovReport 数据集(Huang et al., 2021)与 Proof-pile 数据集(Azerbayev et al., 2022)上评估。对于口令检索,他们采用(Mohtashami and Jaggi, 2023a)所述方法生成合成提示以供评估。
实验通过全长(Full-length)训练将上下文扩展至 16k 与 32k,并对比两种技术:RandPos(Ruoss et al., 2023b)与 PoSE。报告各扩展模型(含未微调的 LLaMA 模型,记为 None)在不同评估上下文窗口大小(从 2k 到 32k)上的困惑度分数。评估采用(Press et al., 2021a)提出的滑动窗口方法,为效率起见将窗口步长设为 1,024。
观察包括:经 PoSE 扩展的 16k 与 32k 模型整体呈现困惑度下降趋势。尽管微调时上下文长度显著更短,PoSE 仍取得与 Full-length 相当的结果,证实了其有效性。值得注意的是,PoSE 优于 RandPos。
在(Mohtashami and Jaggi, 2023a)提出的口令检索测试中,模型需从冗长文档中恢复隐藏的随机口令。未微调的 LLaMA 模型(None)在提示长度超过 2k 时检索准确率迅速降至 0。相比之下,两个经 PoSE 扩展的模型在各自目标上下文窗口内均保持高检索准确率(≥ 90%)。这表明经 PoSE 训练的模型确实具备关注扩展上下文窗口内所有 token 的能力。
图 17: 全长微调与 PoSE(Zhu et al., 2023)微调对比示意,将上下文从 2,048 扩展至 8,192 个 token。全长方法直接使用全部 8,192 个 token,而 PoSE 以唯一的跳跃偏置项调整 2,048 个 token 的位置索引。这使模型在微调期间适应不同的相对位置。
相关工作。 若干模型——包括(Press et al., 2021b; Sun et al., 2022)与(Haviv et al., 2022)所提——旨在确保即使推理时输入 token 数超过模型训练的上下文窗口大小,性能仍保持一致。另一方面,一些工作(Chen et al., 2023a; bloc97, 2023b; Peng et al., 2023; Chen et al., 2023b)聚焦于以更长上下文窗口微调大语言模型(LLM)。然而,所有这些方法都需要全长微调,计算成本随目标上下文规模增长。(Ruoss et al., 2023a)也尝试在训练期间模拟更长序列以应对分布外长度。他们引入随机位置编码(RandPos),从更长序列中随机选择有序的位置索引子集。PoSE 与 RandPos 显著不同:RandPos 主要在预训练期间增强仅编码器模型的长度泛化能力;相比之下,PoSE 高效扩展预训练大语言模型(LLM)——尤其是仅解码器架构——的上下文窗口。此外,RandPos 在相邻 token 之间缺乏连续位置索引,而 PoSE 有意在每个块内保持这种连续性。该连续性与预训练阶段高度对齐,从而将对已学语言建模与理解能力的干扰降至最低。
管理极长输入序列通常涉及记忆机制。两条突出的研究路线利用记忆:基于循环的方法((Dai et al., 2019b; Bulatov et al., 2022))与基于检索的策略((Wu et al., 2022; Wang et al., 2023; Tworkowski et al., 2023))。基于循环的方法切分冗长输入,将先前片段的隐藏状态复用为当前片段的记忆。然而,它们存在信息损失与有限随机访问能力的问题。相比之下,基于检索的范式将先前序列编码为(key, value)对,使用记忆检索器与阅读器提取编码信息。其缺点是离散记忆片段之间缺乏交互。
近期,(Mohtashami and Jaggi, 2023a)提出 landmark attention,通过 landmark token 实现对输入块的随机访问。然而,PoSE 方法在不改变注意力机制的情况下实现对整个输入的完全访问。
6 结论
总之,本文全面综述了扩展大语言模型(LLM)上下文长度的多样技术与方法。所提出的分类法将这些方法归为两大策略——外推与插值。外推技术旨在扩展模型处理超出其初始训练上下文长度之序列的能力。这包括利用位置编码、注意力机制与记忆增强等专用组件实现即时泛化的零样本方法。亦探索了微调策略,以使模型适应预训练阶段未见过的更长上下文。插值技术聚焦于优化模型,使其在观测到的训练长度内平滑扩展上下文理解。专用注意力机制与提示压缩有助于高效处理冗长上下文。微调插值使模型在序列开始超出训练长度时能够优雅过渡。本综述洞察了涵盖提示工程、注意力机制、位置编码与记忆增强等多种技术的通用性,并凸显了为应对上下文长度限制而定制的模型架构与训练方法创新。广泛的经验分析在基准测试与下游任务上证实了这些多样技术的有效性。通过提供结构化分类法并对既有文献加以综述,本文有助于更清晰地理解大语言模型(LLM)上下文长度扩展这一不断演进的图景。讨论部分指明了有前景的研究方向,强调持续开发擅长处理广泛上下文信息之模型的重要性。随着对长篇文本生成与大规模语料推理兴趣的增长,增强的上下文处理在未来数年仍将是活跃的研究领域。
7 讨论
本全面综述凸显了在开发扩展大语言模型(LLM)上下文能力的多样方法方面所取得的显著进展。然而,若干开放问题与挑战仍然存在,有待研究社区进一步探究。未来工作的一个关键方向是探索本文所综述技术的协同组合。例如,将记忆增强策略与专用注意力机制相结合,有望产生擅长处理显著更长上下文的模型。利用不同技术互补优势的混合方法值得深入探索。另一需要关注的关键方面是开发合适的评估基准测试与指标,以准确评估上下文扩展技术。尽管已有初步基准测试被提出,标准化套件可促进方法间更严格的比较。提供超越简单困惑度分数、对模型上下文能力有细致洞察的指标将很有价值。上下文扩展技术的可解释性亦是一个探索不足的领域。增强模型如何利用扩展上下文之可解释性的方法,可解锁对其内部机制的更深洞察。这种可解释性对调试、分析以及大语言模型(LLM)的负责任部署至关关键。训练效率以及开发扩展上下文模型的高昂资源成本是重大挑战。能够比肩原生短上下文训练效率的技术有望加速进展。多阶段训练流程与面向上下文扩展的迁移学习是有前景的方向。最后,研究长上下文对大语言模型(LLM)涌现能力的影响带来引人入胜的机遇。例如,对文档而非句子进行推理如何改变模型对复杂概念的理解?通过精心设计的评估与实验探究这些高阶效应,仍是未来工作的开放途径。总之,本综述提供了一个结构化基础,总结进展并勾勒关键开放挑战。持续研究若能借助这一既有文献的综合,将进一步推动展现对长程上下文精细感知的大语言模型(LLM)发展。凭借更高的上下文精密度,这些模型有望最终达致更接近人类的语言理解。
参考文献
(参考文献列表共约 8 页,请参见原文 PDF 第 37–45 页,此处不再翻译。)
附录:表 1 上下文长度扩展技术总览
表 1:与上下文长度扩展技术相关的全部工作总览。 此处按以下因素对各工作进行划分:1. 技术;2. 训练长度;3. 评估长度;4. 指标;5. 模型;6. 任务;7. 基准。
(表中具体条目为英文项目名与标签,完整版请参见原文 PDF 第 46–47 页。)
外推:零样本(Extrapolation: Zero-shot)
位置编码(Positional encoding)
| 标题 | 训练长度 | 评估长度 | 指标 | 模型 | 任务 | 基准 |
|---|---|---|---|---|---|---|
| Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (Press et al., 2021b) | 64, 128, 256, 512, 1024, 1536, 2048, 3072 | 64, 128, 256, 512, 1024, 1536, 2048, 3072 | Perplexity, Words per seconds, Memory | Customized Transformer Model | Language modeling, text generation | WikiText-103, Toronto Book Corpus, CC100+RoBerta Corpus |
| RoFormer: Enhanced transformer with Rotary Position Embedding (Su et al., 2021) | 128, 256, 512, 1536 | Details not provided | BLEU, GLEU, Accuracy | Customized Transformer Model (RoFormer), BERT, WoBERT, NEZHA | Machine translation, Semantic text matching | WMT 2014 English-German dataset, Chinese dataset, CAIL2019-SCM, GLUE (MRPC, SST-2, QNLI, STS-B, QQP, MNLI), Wikipedia Corpus Foundation, BookCorpus |
| Randomized Positional Encodings Boost Length Generalization of Transformers (Ruoss et al., 2023b) | 1024, 2048, 4096, 8192 | Details not provided | Accuracy | Encoder-only Customized Transformer Model | Various algorithmic reasoning tasks | Manual |
专用注意力机制(Specialized attention mechanism)
| 标题 | 训练长度 | 评估长度 | 指标 | 模型 | 任务 | 基准 |
|---|---|---|---|---|---|---|
| A Length-Extrapolatable Transformer (Sun et al., 2023b) | 1024 | 256, 512, 1024, 2048, 4096 | Perplexity | Customized Transformer Model (LeX) | Details not provided | arXiv dataset, Pile, Books3, OpenWebText2, Stack Exchange, PubMed Abstracts, Wikipedia, PG-19, BookCorpus2, NIH Exporter, Pile-CC |
| LongNet: Scaling Transformers to 1,000,000,000 Tokens (Ding et al., 2023b) | 8k, 16k, 32k | 2k, 8k, 32k | Perplexity | Customized Transformer Model (LongNet) | Long-sequence modeling | Stack dataset |
基于窗口的方法(Window based approaches)
| 标题 | 训练长度 | 评估长度 | 指标 | 模型 | 任务 | 基准 |
|---|---|---|---|---|---|---|
| GrowLength: Accelerating LLMs Pretraining by Progressively Growing Training Length (Jin et al., 2023) | 128, 256, 512, 1024 | Details not provided | Training loss | 70M, 160M and 410M LLM (specific model not mentioned) | Details not provided | Neural networks and the chomsky hierarchy |
记忆/检索增强方法(Memory/Retrieval augmented approaches)
| 标题 | 训练长度 | 评估长度 | 指标 | 模型 | 任务 | 基准 |
|---|---|---|---|---|---|---|
| Landmark Attention: Random-Access Infinite context length for Transformers (Mohtashami and Jaggi, 2023a) | 250, 256, 300, 360, 512 | 512, 2048, 4096 | Perplexity | Transformer-XL, LLaMA-7B | Language modeling, Next word prediction, Information retrieval over long contexts | PG-19, arXiv |
| Augmenting Language Models with Long-Term Memory (Wang et al., 2023) | 1k | 65k | Perplexity, accuracy, F1 score | GPT-2-407M | Long-context language modeling, long-context understanding, memory-augmented in-context learning | Gutenberg-2022 (PG-22), arXiv, ChapterBreak, SST-2, MPQA, MR, Subj, SST-5 |
外推:微调(Extrapolation: Fine-tuned)
记忆/检索增强方法(Memory/Retrieval augmented approaches)
| 标题 | 训练长度 | 评估长度 | 指标 | 模型 | 任务 | 基准 |
|---|---|---|---|---|---|---|
| Think-in-Memory: Recalling and Post-thinking Enable LLMs with Long-Term Memory (Liu et al., 2023a) | Details not provided | Details not provided | Retrieval Accuracy, Response Correctness, Contextual Coherence | ChatGLM-6B, Baichuan2-13B | Response generation for long-term conversations | KdConv, GVD, RMD |
| Focused Transformer: Contrastive Training for Context Scaling (Tworkowski et al., 2023) | 512, 1024, 2048, 4096 | 2k, 4k, 8k, 16k, 64k, 128k | Perplexity | Decoder-only Customized Transformer Model (LongLLaMA), OpenLLaMA 3B, OpenLLaMA 7B | Passkey retrieval, QA, Long-context language modeling | PG-19, arXiv, Github code, Isabelle, TREC, WebQS |
| MemGPT: Towards LLMs as Operating Systems (Packer et al., 2023) | Details not provided | Details not provided | ROUGE-L, Accuracy, CSIM similarity scores | GPT-4 | Deep memory retrieval task, Conversation opener task, Multi-document QA, Nested key-value retrieval requiring multi-hop lookups | MSC, NQ-Open Wikipedia |
插值:零样本(Interpolation: Zero-shot)
专用注意力机制(Specialized attention mechanism)
| 标题 | 训练长度 | 评估长度 | 指标 | 模型 | 任务 | 基准 |
|---|---|---|---|---|---|---|
| LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models (Han et al., 2023a) | Details not provided | 2k, 4k, 8k, 16k, 32k, 128k | Perplexity | MPT-7B, LLaMA, GPT-J-6B, LLaMA-2 | Long text generation | arXiv, OpenWebText2 |
| LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models (Chen et al., 2023b) | 2048, 4096 | 8192, 16384, 32768, 65536, 100,000 | Perplexity, Passkey retrieval accuracy, Win-rate in topic retrieval | LLaMA2-7B, 13B, and 70B | Long-sequence language modeling, Topic retrieval | RedPajama |
| LongQLoRA: Efficient and Effective Method to Extend context length of Large Language Models (Yang, 2023) | 4k, 8k | 1024, 2048, 4096, 8192 | Perplexity | LLaMA-2-7B, MPT-7B, LongLoRA-7B, LongQLoRA | Long-sequence language modeling, Passkey retrieval, Topic retrieval in long conversations | PG-19, Proof-Pile |
基于提示压缩的方法(Prompt compression based approaches)
| 标题 | 训练长度 | 评估长度 | 指标 | 模型 | 任务 | 基准 |
|---|---|---|---|---|---|---|
| LongLLMLingua: Accelerating and Enhancing LLMs in Long Context Scenarios via Prompt Compression (Jiang et al., 2023b) | 3k-10k | Details not provided | Accuracy | GPT-3.5-Turbo, LongChat-13B | Single and Multi-Doc QA, Summarization, Code completion, Sentiment classification, Information reordering | NQ-multi-document QA, LongBench, ZeroSCROLLS |
插值:微调(Interpolation: Fine-tuned)
基于 RoPE 的方法(RoPE based approaches)
| 标题 | 训练长度 | 评估长度 | 指标 | 模型 | 任务 | 基准 |
|---|---|---|---|---|---|---|
| Extending Context Window of Large Language Models via Positional Interpolation (Chen et al., 2023a) | 8192, 16384, 32768 | 2048, 4096, 16384, 32768 | Perplexity | LLaMA-7B, 13B, 33B, and 65B | Language modeling, Passkey retrieval, Long document summarization | Pile, PG-19, RedPajama, arXiv math |
| YaRN: Efficient Context Window Extension of Large Language Models (Peng et al., 2023) | 4k, 32k, 64k, 100k, 128k | 8192, 32768, 65536, 98304, 131072 | Perplexity | LLaMA-2 7B, LLaMA-2 13B, GPT-NeoX and Mistral 7B v0.1, MistralLite 7B, PaLM | Passkey retrieval | PG-19, Proof-pile, Hugging Face open LLM benchmark suite |
| PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training (Zhu et al., 2023) | 2k, 16k | 2k, 4k, 8k, 16k, 32k, 64k, 96k, 128k | Perplexity | LLaMA-7B | Language modeling, Passkey retrieval | Proof-Pile, GovReport, Gutenberg (PG-19), Books3 |