Leo Pekelis：如何将 Llama 3 扩展至百万 token 上下文

Matthew Berman 介绍 Gradient 首席科学家 Leo Pekelis，与他讨论团队将 Llama 3 上下文窗口扩展到百万 token 的工作。在 2024 年 6 月的访谈时，原版 Llama 3 支持 8K 上下文。Berman 问什么是上下文窗口，以及为什么扩展它很重要。

Pekelis 解释，自回归语言模型一次生成一个 token。指令、文档和对话历史构成序列的初始内容，引导模型接下来生成什么。上下文窗口就是模型在这个序列中可用的文本范围；Berman 指出，输入和输出共同计入总量。

Pekelis 将上下文比作工作记忆，将训练所得的知识比作记忆深处的内容。更长的窗口让模型可以关注更多提供的材料，可能减少在模型外部预处理、筛选或总结的需求。这个类比有助于说明区别，并非声称模型真的拥有人类记忆。

Berman 指出，更长的窗口可以容纳此前必须分块处理的书籍或多模态材料。Pekelis 说，好处不只是方便：当信息都能放入窗口时，模型可以尝试自行选择哪些部分重要，而不必完全依赖用户事先筛选。实际质量和速度仍取决于模型与任务。

谈到企业编程场景时，Pekelis 将补全一个小函数与实现一个需要整个代码库或多个代码库信息的功能或工作流作比较。他说，智能体可以逐步获取片段，但这种方式可能脆弱、缓慢且昂贵。更大的窗口使一次请求容纳更多材料成为可能，也可能减少调用次数。

Berman 描述了主要只能看到当前打开文件或部分项目结构的助手的局限。Pekelis 说，Gradient 曾尝试将不同代码库放进上下文，询问如何把它们集成起来。这是团队实验的例子，并不保证模型能理解或安全修改每个代码库。

谈到自己的工作流程，Pekelis 说他使用 VS Code，也试过 Cursor。他还尝试 Gradient 自己的长上下文模型，但这些模型与编辑器的集成较少，需要更多来回操作。团队在内部集群上运行它们，并在研究最佳服务方式的同时接受一定等待时间。

团队正在尝试缓存查询的开头部分，例如在多次提问之间基本不变的代码库。这样，后续请求就可以复用与共同前缀相关的计算。Pekelis 说这只是一个近似解释：实际实现比简单保存所有两两注意力结果更复杂。

Pekelis 从注意力机制解释计算瓶颈：在直接实现中，token 之间的交互大致随序列长度的平方增长。复用前缀时，缓存可以避免重复部分工作，但处理长输入仍消耗大量资源。因此，高效提供服务是独立于训练模型处理长上下文的另一项挑战。

Berman 问 Gradient 如何扩展 Meta 的基础 Llama 3。Pekelis 指出两个挑战：处理长序列所需的计算，以及教会为较短上下文训练的模型处理长得多的上下文。仅仅输入更多 token，并不能证明模型能可靠地利用它们。

Pekelis 说，这种训练更接近继续预训练，而非简单的常规微调，位置编码在其中起重要作用。模型需要一种方式来表示可能相隔数百甚至更多 token 的相对位置。他认为位置编码是扩展模型可用范围的核心。

Pekelis 说，长上下文训练成本很高，Gradient 在 Llama 3 出现之前就已着手提高效率。他感谢开放研究和资源对团队技术栈的影响。与 Berkeley 已发表的 Large World Model 论文相比，他估计其计算时间效率约高 30 倍，样本效率或许高近 100 倍；这是团队的比较，而非独立确立的普适比例。

Berman 问扩展后的模型对终端用户硬件的要求是否更高。Pekelis 说，仅仅因为模型支持更大的上限，使用较短的 8K 上下文并不会更加耗费资源；但输入百万 token 就需要多得多的计算。这里的区别在于模型的最大容量和实际使用的长度。

Pekelis 补充说，Gradient 在扩展范围的同时努力保留 Llama 3 的短上下文表现。对团队而言，如果长上下文严重损害基础模型的性能，就不算值得的改进。他的说法描述了设计目标和团队监测，并不能证明所有短上下文能力都没有变化。

为解释评测，Pekelis 描述了大海捞针测试：在大量文本中插入一条信息，再要求模型将其找出来。Gradient 有时使用《战争与和平》作为背景材料的一部分，他估计其长度约为 75 万 token。他认为关联回忆是有用的基础测试，也是更丰富的上下文学习的前提。

Berman 问，当应用需要关联整个代码库中的信息时，找回一个孤立事实是否足够。Pekelis 同意，找到单个信息与对整个项目进行推理之间还有很大差距。这项测试只是第一步，访谈时许多长上下文模型在该测试上已经接近满分。

Berman 说 Gradient 报告的检索热图在测试位置上全部呈绿色；Pekelis 说 Gemini 的结果也全绿。这说明模型在特定检索测试中取得成功，并不证明它能够在百万 token 的整个范围内进行各种推理。要区分都擅长简单任务的系统，还需要更难的评测。

Pekelis 提到 NVIDIA 的 RULER 基准，它将大海捞针扩展为 13 种任务。包括多个目标、不同事实使用相同键，以及从众多干扰项中选出相关信息。变量追踪任务要求模型沿着 X2 = X1、X3 = X2 等引用关系找回一个值。

Berman 将变量追踪比作寻宝游戏，一条线索引向下一条。Pekelis 说，Gradient 的模型当时在 RULER 排行榜上位列第四，排在 Gemini、GPT-4 和 Cohere Command R+ 之后。Berman 对最大上下文长度不同的模型之间如何比较提出疑问。

Pekelis 承认关键限制：由于任务更复杂，他所说的 RULER 对比只运行到 128K token。因此，排行榜名次并不能证明模型在百万 token 时也有同样表现。Gradient 当时正在研究如何对超过 128K 的情况进行更详细的测试。

展望未来，Pekelis 对以内存和计算消耗更少的方式提供长上下文模型服务感兴趣。当时，用户在很大程度上依赖 Google 或 Gradient 等供应商托管这些资源需求高的模型。他提出可以压缩部分记忆，并在需要时有选择地展开；这借用了人们回忆相关章节、而非将整本教科书放在脑海里的类比。

Pekelis 邀请研究人员和用户通过社交渠道及 Discord 联系 Gradient，并说团队受益于与开源社区的合作。Berman 最后表示，他想继续测试百万 token 的 Llama 3 模型，并希望团队最终达到 200 万 token。Pekelis 将其视为未来目标，而非发布公告。