微软最新突破：如何让大型语言模型在超长文本中聊得又快又好？

近日，由微软研究院的孙宇涛、叶天竹、董力、夏玉清、陈健、高亦昭、曹世杰，以及清华大学的王建勇和香港大学的研究团队共同发表了一篇名为《Rectified Sparse Attention》（矫正稀疏注意力机制）的研究论文。这项研究针对大型语言模型处理超长文本时的效率问题提出了一种全新的解决方案，团队的研究成果可在https://aka.ms/ReSA-LM获取。

随着ChatGPT、Claude等大型语言模型的普及，人们越来越希望AI能够处理更长的文本内容。想象一下，你希望AI能阅读并分析一本完整的书籍，或者处理成百上千页的法律文件。这样的长文本处理需求正在迅速增长，一些最新的模型甚至可以处理数百万个词的文本。但是，处理如此长的文本带来了一个严峻的挑战：效率问题。

为什么处理长文本会如此耗时呢？这里需要简单了解一下大型语言模型的工作方式。当模型生成文本时，它需要"记住"之前的所有内容，这些记忆被存储在所谓的"KV缓存"（Key-Value缓存）中。想象这个缓存就像模型的短期记忆，存储着它刚刚处理过的信息。在标准的处理方式中，模型每生成一个新词，都需要查看所有之前的内容，这就像是一个人在写作时，每写一个词都要回头完整阅读之前写的所有内容。当文本变得很长时，这个过程会变得异常缓慢。

有研究者提出了"稀疏注意力"的解决方法，就像是让模型只关注最相关的部分内容，而不是所有内容。这就好比写作时只回顾最重要的几段，而不是整篇文章。这种方法确实能提高速度，但随着生成文本的增长，质量会逐渐下降。为什么会这样呢？因为在稀疏处理过程中，微小的计算误差会不断积累，最终导致模型的"记忆"与训练时的状态越来越不一致，就像是复述一个故事时，每次复述都会有小的偏差，最终导致故事变得面目全非。

微软研究团队提出的"矫正稀疏注意力"（ReSA）方法巧妙地解决了这个问题。这个方法就像是定期"校准"模型的记忆，确保它不会偏离轨道太远。具体来说，ReSA结合了两个关键步骤：首先，它使用"分组块稀疏注意力"让模型只关注最相关的内容，加快处理速度；其次，它会定期进行"密集矫正"，即完整地检查一遍最近生成的内容，修正可能积累的误差。这就像是一个人在讲述长故事时，大多数时候只回忆关键部分，但每隔一段时间会完整地回顾一下最近的内容，确保故事不会偏离原意。

研究团队通过多项实验证明，这种方法不仅保持了生成内容的高质量，还显著提高了处理速度。在数学推理、语言建模和信息检索等任务上，ReSA的表现几乎与完整的密集注意力机制相当，但速度快了许多。特别是在处理25.6万词的超长文本时，整体速度提升了2.42倍，这是一个相当显著的改进。

一、分组块稀疏注意力：智能筛选重要信息

为了理解ReSA的工作原理，我们需要先了解其核心组件之一：分组块稀疏注意力。想象你正在阅读一本千页的百科全书，当你阅读到第800页时，你不可能记住之前所有页面的每一个细节。相反，你的大脑会选择性地记住一些重要的信息块，比如关键概念、重要事实或与当前阅读内容相关的段落。

分组块稀疏注意力机制模拟了这种人类的选择性注意力。它将之前的内容划分为多个"记忆块"，每个块包含一组连续的词。当模型需要生成新内容时，它不会检查所有之前的内容，而是智能地选择最相关的几个记忆块进行参考。

那么，模型如何知道哪些块是重要的呢？每当模型需要生成新内容时，它会进行一个"相关性评分"的过程。对于每个记忆块，模型会计算它与当前生成需求的相关程度。这就像是你在写作时，会下意识地想到与当前段落最相关的之前内容。

为了提高效率，研究团队采用了一种巧妙的表示方法。对于每个记忆块，模型不会存储其中的所有详细信息，而是存储两个关键"摘要"：块中最小值和最大值。这就像是给每个内容块做了一个"范围标签"，说明这个块中包含的内容大致在什么范围内。这种简化的表示方法大大减少了计算量，同时保持了足够的信息来判断相关性。

此外，研究团队还引入了"共享分组"技术，允许多个注意力头共享同一个稀疏模式。这就像是团队成员共享同一份阅读笔记，而不是每个人都单独做笔记，进一步提高了效率。

在实际运行中，模型会动态确定需要关注的记忆块数量。它总是会保留最近的几个块（称为"局部窗口"），确保短期连贯性；同时，它还会根据内容长度，选择一定比例的历史块进行关注。例如，当处理非常长的文本时，模型可能只关注10%的历史内容，但这些内容是经过精心筛选的最相关部分。

二、密集矫正：定期校准确保质量不下滑

虽然稀疏注意力机制大大提高了处理速度，但正如前面提到的，它会导致误差累积问题。这就好比是导航系统中的"漂移"现象：如果你只依靠估算行走距离，而不参考GPS信号进行校正，那么误差会随着时间不断累积，最终导致你的位置估计与实际位置相差甚远。

ReSA通过"密集矫正"巧妙地解决了这个问题。密集矫正的核心思想非常直观：每生成一定数量的新内容后（比如每32个词），模型会暂停稀疏处理，对最近生成的这批内容进行一次完整的密集注意力计算。这就像是导航系统定期接收GPS信号来校准位置，确保累积误差不会太大。

这个过程在技术上是如何实现的呢？当达到预设的矫正频率（默认是32个词）时，模型会将最近生成的这批词一次性输入，使用完整的密集注意力机制重新计算它们的表示，并用这些精确的表示更新KV缓存。这样，即使之前的稀疏计算引入了一些误差，这些误差也会被定期"重置"，不会无限制地累积。

值得注意的是，这种矫正是批量进行的，而不是逐词进行。这意味着虽然矫正过程使用了计算成本较高的密集注意力，但由于是一次性处理多个词，平均到每个词的额外成本其实很小。就像是你可能每周花一个小时整理一次笔记，而不是每写一段就停下来整理，这样总体效率会更高。

研究者们发现，通过适当选择矫正频率，可以在保持生成质量和提高计算效率之间取得良好的平衡。当矫正频率设为32（即每生成32个词进行一次矫正）时，模型的表现接近于完全使用密集注意力的上限，同时仍然保持了显著的速度优势。

三、全面评估：ReSA在各种任务中的表现

为了验证ReSA的有效性，研究团队在多种任务上进行了广泛的实验评估，包括数学推理、语言建模和信息检索。

在数学推理任务上，研究者使用了多个具有挑战性的数据集，包括Minerva Math、高考2023英文数学题、奥林匹克数学题集、AIME24和AMC23。这些任务特别适合测试长文本处理能力，因为解决复杂数学问题通常需要模型维持长期的推理链。结果显示，使用纯稀疏注意力的模型表现明显不如密集注意力，特别是在生成较长答案时。而使用ReSA的模型几乎能达到与密集注意力相同的准确率，同时保持了更高的效率。

在语言建模任务中，研究者设计了一个巧妙的实验，将输入序列分为两部分：前缀使用密集注意力处理，后缀使用稀疏注意力或ReSA处理。这样可以模拟不同长度生成过程中的性能变化。结果再次证明，随着生成长度的增加，纯稀疏注意力的性能明显下降，而ReSA能够保持稳定的高性能。特别是当矫正频率设为32时，ReSA的表现几乎达到了理论上限。

研究者还尝试了不同的稀疏率设置。当模型只关注10%的历史内容时（即稀疏率为0.9），ReSA仍能保持与密集注意力相当的性能，这证明了该方法在高稀疏率下的有效性。

在RULER信息检索基准测试中，由于输出序列相对较短，矫正的作用不太明显，但结果仍然表明适当的稀疏率设置（p=0.9）能够在保持性能的同时提高效率。

四、效率分析：ReSA带来的速度提升

除了保持生成质量外，ReSA的另一个主要优势是显著提高了处理效率。研究团队在标准GPU硬件（NVIDIA A100-80G）上进行了详细的效率测试。

从理论上分析，ReSA的内存访问成本可以表示为：mem(KV缓存)×(1/b + p + 1/f)，其中b是块大小，p是稀疏率，f是矫正频率。相比于密集注意力需要访问整个KV缓存，ReSA只需要访问其中的一小部分，从而大大减少了内存带宽需求。

在实际测试中，研究者比较了在不同序列长度下的延迟情况。结果显示，当序列长度增加时，密集注意力的延迟急剧增加，而ReSA的延迟增长则平缓得多。特别是在超长序列（25.6万词）下，ReSA的速度优势最为明显。

延迟分析还显示，稀疏估计和注意力计算占用了相当的时间，这是因为它们的内存访问模式类似（一个与缓存大小/块大小成比例，另一个与缓存大小×稀疏率成比例）。而矫正过程的额外开销相对较小，在25.6万词长度下仅占总延迟的32.7%。

在端到端的吞吐量测试中，ReSA在FP16精度下提供了最高2.28倍的加速，在INT4量化精度下提供了最高2.44倍的加速。这种加速效果在序列越长时越明显，充分证明了ReSA在实际应用中的价值。

五、参数选择与权衡

为了帮助用户选择最适合自己需求的参数设置，研究团队还进行了详细的消融实验，分析了不同矫正频率和稀疏率对性能的影响。

实验结果显示，在大多数任务上，矫正频率f=32或f=64能够在保持高质量和提高效率之间取得最佳平衡。虽然更频繁的矫正（f=16）可能带来轻微的性能提升，但矫正开销也会相应增加，降低了整体效率。另一方面，即使较低频率的矫正（f=128）也能保持显著的性能提升，这表明矫正机制对不频繁更新具有很强的鲁棒性。

对于稀疏率，研究者发现p=0.9（即关注10%的历史内容）通常能提供最佳的效率与性能平衡。虽然更低的稀疏率（如p=0.8，关注20%的内容）可能在某些情况下提供略好的性能，但考虑到较低稀疏率会导致更多的内存访问和计算成本，p=0.9成为了推荐的默认设置。

值得一提的是，ReSA的设计是完全无需训练的，它可以直接应用于现有的预训练模型，无需任何微调或重新训练。这使得它特别适合在实际部署中使用，因为用户可以立即享受到效率提升，而不需要耗费大量资源进行模型再训练。

这项研究的成功之处在于，它巧妙地结合了稀疏注意力的高效率和密集注意力的高质量，为长文本生成找到了一条实用的中间道路。正如研究者所展示的，通过适当的参数设置，ReSA可以在几乎不损失性能的情况下，将超长文本处理的速度提高2倍以上，这对于实际应用具有重要意义。

当人工智能技术日益融入我们的生活时，处理长文本的能力变得越来越重要。无论是分析法律文件、医学文献，还是与用户进行深入的多轮对话，都需要模型能够高效地处理长文本。微软研究团队提出的ReSA方法为这一挑战提供了一个简单而有效的解决方案，它不需要复杂的模型修改或重新训练，就能显著提高处理效率。

这项研究的代码已在https://aka.ms/ReSA-LM公开，感兴趣的读者可以直接尝试将这一技术应用到自己的项目中。随着这类高效处理技术的发展，我们可以期待人工智能在处理复杂长文本任务方面的能力将继续提升，为用户提供更加流畅、自然的交互体验。