Learning how to Forget: Fine-tuning for Long-Context Sparse Attention
本文介绍了一种针对 Transformer 模型的微调方法，该方法通过允许模型与各种 KV 缓存策略进行协同适应，使得在适度硬件条件下实现高效的长文本推理成为可能，且其性能通常优于精确注意力机制的方法。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
现代生成类人文本的人工智能系统依赖于一种充当短期记忆的机制，使其能够记住对话或长文档中之前说过的话。这种记忆存储在一个数字缓冲区中，随着系统处理的每一个新词而增长。为了让这些系统在长篇文本中表现良好，这种记忆需要非常庞大，但用于容纳它的计算机硬件既昂贵又有限。当记忆缓冲区填满时，系统必须决定丢弃哪些旧信息以腾出空间给新信息。如果它丢弃了错误的信息，系统就会失去推理或准确回答问题的能力。这产生了一个两难的权衡：保持较小的记忆可以节省成本并允许系统在标准设备上运行，但会面临失去智能所需的上下文信息的风险。
研究人员长期以来一直试图通过教系统学会对保留内容进行选择（这一过程被称为稀疏注意力）来解决这个问题。然而，一项新研究揭示了迄今为止这些系统训练方式中的一个关键缺陷。大多数现有方法使用完美的、无限的记忆来训练 AI，然后试图在稍后强制其在有限的记忆下运行。研究人员发现，这种方法之所以失败，是因为 AI 从未学习过如何在它实际面临的特定约束下运作。通过让模型从一开始就学习有目的地遗忘并适应固定的记忆大小，该团队证明了系统可以比那些使用无限资源的模型表现得更好，即使是在运行于单块中等性能的计算机芯片上时也是如此。
由亚马逊网络服务（AWS）和阿姆斯特丹大学的科学家领导的团队开发了一种微调这些大型语言模型的新方法。他们并没有使用大规模超级计算机来模拟完美记忆，而是教会了模型与特定的记忆管理策略进行协同进化。想象一下一位图书管理员，他接受的是在拥有无限书架的图书馆中整理书籍的训练，结果后来却被告知要在只有一个书架的小房间里工作。在大型图书馆接受训练的管理员在面对小房间时可能会感到吃力。相比之下，本文提出的方法直接在小房间里训练这位管理员，教会他们根据那个特定空间的规则来决定保留哪些书以及丢弃哪些书。这使得模型能够学习自身局限性的节奏，而不是试图去改掉拥有过多空间的习惯。
研究人员在一个拥有 40 亿参数的模型上测试了这种方法，这个规模虽然可观但仍属可控。他们在配备 40 GB 显存的单张图形卡上进行了实验，与之前的方法相比，这种配置对于许多组织来说比需要数十张卡组成的集群更经济实惠。他们将这种新训练技术与标准方法进行了对比，标准方法使用一种称为序列并行（sequence parallelism）的技术，将记忆负载拆分到多个昂贵的设备上。结果显示，使用新方法的模型通常优于标准模型，特别是在任务要求系统生成具体、简洁的答案而非冗长、啰嗦的文本时。在涉及复杂问题和数据提取的几项测试中，标准方法生成的输出往往过于冗长且充满了随机、无意义的数字，而新方法则学会了在正确的时间停止并提供正确的单一数值。
成功的关键部分在于改进了“重击者预言机”（heavy-hitter oracle），这是一种流行的决策哪些信息值得保留的策略。该策略的工作原理是追踪模型在一段时间内投入最多注意力的信息片段。研究人员发现，该策略的原型版本既缓慢又低效。他们重写了底层代码，使其运行得更快，从而使系统能够在不减慢整个过程的情况下计算这些重要性评分。这种优化意味着系统可以实时做出关于遗忘什么的聪明决策，而无需通常伴随此类任务的大规模计算能力。该团队还发布了一个新的开源软件库，旨在向他人提供这些技术，从而降低任何人构建长上下文 AI 系统时的门槛，而不必需要巨额的硬件投入。
这项研究强调，模型的训练方式与运行它的硬件同样重要。当研究人员强迫模型使用与其使用时完全相同的记忆约束进行训练时，它学会了如何有效地应对这些约束。在一次涉及 JSON 数据的特定测试中，标准方法完全失败，无法找到正确的数据点，而新方法成功识别了其中的大约一半。这表明，处理长上下文的能力不仅在于拥有更多内存，还在于教会系统如何管理它已有的内存。研究结果表明，对于许多应用而言，最有效的途径不是建造更大的计算机，而是教会软件如何更高效地利用其现有的资源。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。