面向超长上下文，大语言模型如何优化架构，这篇综述一网打尽了

文章来源：机器之心

作者重点关注了基于 Transformer 的 LLM 模型体系结构在从预训练到推理的所有阶段中优化长上下文能力的进展。

图片来源：由无界 AI生成

ChatGPT 的诞生，让基于 Transformer 的大型语言模型 (LLM) 为通用人工智能（AGI）铺开了一条革命性的道路，并在知识库、人机交互、机器人等多个领域得到应用。然而，目前存在一个普遍的限制：由于资源受限，当前大多 LLM 主要是在较短的文本上进行预训练，导致它们在较长上下文方面的表现较差，而长上下文在现实世界的环境中是更加常见的。

最近的一篇综述论文对此进行了全面的调研，作者重点关注了基于 Transformer 的 LLM 模型体系结构在从预训练到推理的所有阶段中优化长上下文能力的进展。

论文链接：https://arxiv.org/pdf/2311.12351.pdf

论文首先分析了使用当前基于 Transformer 的模型处理长上下文输入和输出的问题。然后，提供了一个全面的分类体系，以指导 Transformer 架构升级的领域，来解决这些问题。作者对长上下文 LLM 广泛使用的评估需求进行了调研，包括数据集、度量标准和基准模型，以及一些令人惊奇的优化工具包，如库、系统和编译器，以增强 LLM 在不同阶段的效率和功效。最后，文章进一步讨论了这一领域未来研究的主要挑战和潜在方向。作者还建立了一个仓库，汇总了相关文献，并提供实时更新 https://github.com/Strivin0311/long-llms-learning。

综述概览

文章从基本的语言建模目标 (第 2.1 节) 开始，内容涵盖从典型的建模阶段到在基于 Transformer 的仅解码 LLM 中找到的关键架构模块，如图 1 (a) 所示。随后，作者对 LLM 在遇到扩展上下文窗口时的架构限制进行了简要分析 (第 2.2 节)。最后提出了一个全面的方法论分类法 (第 2.3 节)，旨在通过架构创新增强 LLM 的长上下文能力 (见图 1 (b))。这个分类法作为文章的第 3、4、5、6、7 节的指南。

图 1：文章核心概述：(a) 现代基于 Transformer 的仅解码 LLMs 的典型架构解剖图，右上角有图例；(b) 用于增强 Transformer 架构模块的方法论分类法（与 (a) 相对应的颜色）：高效注意力（注意力核心的子模块），长期记忆（针对 KV 缓存），外推性 PEs（针对位置嵌入模块），上下文处理（与上下文预 / 后处理有关）和杂项（整个解码器块以及损失模块通用）。

长上下文，目前有哪些难点待突破？

注意力复杂度。在典型情况下 L ≫ d，MHA 的计算复杂性可以简洁总结如下：它涉及 O (L 2d) 的时间复杂度，包括 QKV 投影的 O (Ld2)，计算 P 的 O (L 2d)，通过 softmax 运算获取 A 的 O (L 2 )，A 与 V 相乘的 O (L 2d)，以及输出投影 O 的 O (Ld2)。它还产生 O (L 2) 的空间复杂度，包括 Q、K、V、O 的嵌入的 O (Ld)，以及额外的 O (L 2) 缓冲区用于存储权重 P 和 A。因此，随着序列长度的增加，时间和空间计算成本都呈二次增加，这对于训练和推理可能都是繁重的。

上下文记忆。LLM 缺乏显式的记忆机制，完全依赖 KV 缓存来存储列表中所有先前 token 的表示。这种设计一旦在一个调用中完成查询，Transformer 在后续的调用中不会保留或召回任何先前的状态或序列，除非整个历史记录逐个 token 重新加载到 KV 缓存中。因此，Transformer 在每个调用中仅具有一个上下文工作记忆，而不是像长短时记忆 (LSTM) 这样的内在记忆机制。这种无状态性在并行性方面提供了计算优势，但在聊天机器人应用等需要长期记忆保留的任务中的挑战很明显。

最大长度约束。在训练阶段，工程师通常需要确定一个关键的超参数 max-length，本文中表示为 L_max。这个超参数代表了批次中任何训练样本的序列长度的上限，通常根据可用的计算资源设置为 1k、2k 或 4k，以避免在 GPU 上发生内存溢出 (OOM) 错误。在推理阶段，LLM 的服务提供者还必须限制用户提示的长度或自动截断它们以与预定义的 L_max 对齐，即使推理资源通常比训练阶段更丰富。需要注意的是 Transformer 的任何模块在本质上都不需要这样的限制，因为所有学习的权重仅依赖于维度大小。因此，理论上只要资源足够，Transformer 可以处理任意长度的序列。然而，当前的语言模型在处理超过 L_max 的输入序列时通常表现出明显的性能下降，经常导致重复和不切实际的输出。

改进的新方法

对于上述限制，有多种改进方法可以探索，例如在训练过程中减少注意力复杂性、设计高效的记忆机制，以及增强长度外推的能力，该模型在短序列上进行训练，但在推理过程中对更长的序列进行测试。

因此，论文全面回顾了致力于改进 LLM 长上下文能力的各个阶段的最新方法，并将它们组织成一个统一的分类法，如图 1 (b) 所示。具体而言，这些方法被分为五个主要的类别，如下：

高效注意力 (论文第 3 节)：这些方法侧重于实现具有降低计算要求的高效注意力机制，甚至实现了线性复杂度。通过这样做，它们能够通过直接在预训练阶段增加 L_max 来推进 LLM 在推理期间的有效上下文长度边界。

长期记忆 (论文第 4 节)：为了解决上下文工作记忆的局限性，一些方法旨在设计明确的记忆机制，弥补 LLM 中缺乏高效和有效的长期记忆的不足。

外推性 PEs (论文第 5 节)：最新的研究致力于通过改进现有位置编码方案的外推性能来增强 LLM 的长度泛化能力。

上下文处理 (论文第 6 节)：除了增强特定低级 Transformer 模块的方法外，一些方法涉及对现成的 LLM 与额外的上下文预 / 后处理。这些方法确保每次调用 LLM 时输入始终满足最大长度要求，并通过引入多个调用开销打破上下文窗口限制。

杂项 (论文第 7 节)：探讨了各种一般且有价值的方法，这些方法不容易归入前面四类，为推进 LLM 的长上下文能力提供了更广泛的视角。

未来方向

论文的第 3、4、5、6 节中讨论了该领域取得的显著进展，但仍然存在一些挑战。下面是对一些关键挑战的探讨以及未来在增强基于 Transformer 的 LLM 的长上下文能力方面进行研究和开发的潜在方向，重点关注架构的增强。

注意力 Trade-off。在第 3 节，作者探讨了高效注意方法往往涉及在保持全尺度注意力依赖性（例如局部注意力）或通过近似注意力提高注意力分数精度以减轻标准注意内核的计算需求之间的微妙权衡。然而，随着上下文的延长，话语结构和相互关联的信息变得越来越复杂，需要捕捉全局、长距离的依赖性，同时保持精确的相关性。

解决这一挑战需要在计算效率和尽可能保留注意模式精度之间找到最佳平衡。因此，在长上下文 LLM 领域，这仍然是一个持续追求的目标。最近的创新如 Flash Attention，探索了算法级别之外的 IO 感知解决方案，极大地提高了运行时和记忆开销的效率，而不会丧失注意精度。这是在实际应用中解决这个问题的一个激动人心的潜在途径。此外，可以探索在「即插即用」替代方案中集成先前的高效策略，利用强大的 GPU 内核编程工具 (如 CUDA) 或更轻量级的 Triton。

记忆效果和效率。正如在文章第 2.1、2.2 节中前面讨论的，作者已经概述了由于缺乏明确的记忆机制，仅依赖上下文内工作记忆以及在延长上下文交互期间 KV 缓存记忆消耗显著增加而产生的限制。这些挑战共同强调了在基于 Transformer 的 LLM 领域需要更有效和高效的记忆机制。虽然第 4 节中引入了各种长期记忆机制，但它们受到其复杂启发式设计引入的额外记忆开销的限制，因此随着时间的推移可能导致性能下降。为了解决这一挑战，研究人员可以从最近的进展中汲取灵感，比如 Paged Attention，研发更有效的记忆存储策略，增强读 / 写吞吐量。

长度外推挖掘。在第 5 节中，作者对与基于 Transformer 的模型的长度外推相关的挑战进行了彻底的分析，重点关注了位置嵌入的普遍设计。文章提供了对最近突破的全面概述，特别是应用于 RoPE 的扩展策略，作者相信这在解决外推限制方面具有重要的前景。值得注意的是，这些进步往往依赖于对复杂高维位置嵌入属性的简化观察，并包含简单的启发式调整。作者对使用高维嵌入来建模序列性的理论基础提出质疑，并探索在这些启发式设计的指导下引导具有许多超参数的可学习嵌入的潜在复苏。作者认为未来的研究应该深入探讨这一领域，尤其是在 Transformer 设置下为建模序列性开发健壮的理论框架方面，比如 CLEX 所实现的内容。

特定但通用目标。前文已经为长文本建模量身定制的具体目标做了讨论，但值得注意的是，许多目标仅限于某些类型的任务，或者仅与 MLM 目标兼容，而不是如今更普遍的 CLM 目标。这突显了需要特定但普遍适用的因果语言建模目标，可以在模型训练的早期有效捕捉长距离依赖性。通过与先前提到的目标相一致，这是可能实现的。

可靠的度量需求。在评估度量方面，文章的第 8 节中研究了许多可选项。根据在评估中的先前经验，常用的度量，如 ROUGE 分数，与人类判断分数存在显著差异，后者可以看作是「神谕」。随着 LLM 在现实世界场景中的快速部署，越来越迫切地需要更可靠的度量来评估长上下文能力，特别是在生成性任务中，其中精确的真实性难以捉摸。一个有希望的途径涉及利用最先进的 LLM (如 GPT4) 的鲁棒性作为人类评审的替代，尽管相关的高成本仍然在更广泛地在研究界中采用方面带来挑战。

更多研究细节，可参见原论文。