检索增强生成 (RAG) 是为大语言模型注入特定领域知识的有效方法，但其推理效率受限于输入提示长度的增加。本文介绍了一种名为块注意力 (Block-Attention) 的新机制，它通过分块独立编码检索到的文档，避免重复编码计算，从而显著提升 RAG 模型的推理效率。研究表明，Block-Attention 能够让 RAG 模型的响应速度与不使用 RAG 的模型几乎相同，甚至还能略微提升模型准确率。该方法通过将输入序列分成多个块，独立计算每个块的键值状态，并仅让最后一个块关注其他块，从而实现高效推理。实验结果表明，Block-Attention 在处理长文本时优势显著，能够大幅降低推理延迟和计算量，为 RAG 模型的实际应用带来了新的可能性。

🤔**RAG 提升 LLM 效率的挑战：**检索增强生成 (RAG) 通过检索外部知识库中的相关文档来增强 LLM 的知识和回复质量，然而，每次查询都需要检索多个文档并将其整合到输入提示中，导致输入序列长度增加，从而显著降低推理效率，特别是推理延迟（TTFT）大幅提高。例如，RAG 模型的推理延迟通常远高于非 RAG 模型。 传统的自回归注意力机制使得每个文档的键值状态 (KV states) 与上下文相关，这意味着模型每次遇到新查询时都需要重新编码 KV states，导致效率低下。这种重复计算限制了 RAG 模型在实际应用中的性能。

🚀**块注意力机制 (Block-Attention) 的设计与原理：**为了解决 RAG 模型推理效率低下的问题，论文提出了块注意力 (Block-Attention) 机制。该机制将整个输入序列划分为多个独立的块 (block)，每个块独立计算其 KV States。只有最后一个块能够关注其他块，在 RAG 场景中，最后一个块通常是用户的输入。这种分块独立编码的方式避免了重复编码已经处理过的文档，从而显著提升了推理效率。 Block-Attention 的实现相对简单：首先独立编码除最后一个块以外的所有块；然后为每个块重新计算位置编码；最后将所有块拼接在一起，并计算最后一个块的 KV State。然而，直接将模型从自注意力机制切换到块注意力机制会导致模型性能下降，因此需要进行微调以适应这种新的编码方式。

📊**实验结果与应用前景：**实验结果显示，Block-Attention 机制能够有效提升 RAG 模型的推理效率，同时保持甚至略微提升模型的准确率。具体来说，在用户输入长度为 50，prompt 总长度为 32K 的极端情况下，Block-Attention 模型的 TTFT 和 FLOPs-TFT 分别能降低至 self-attention 模型的 1.3% 和 0.2%，与无 RAG 模型的效率基本持平。 实验还表明，随着输入序列长度的增加，Block-Attention 模型的优势更加明显，加速效果显著。例如，当输入序列长度为 512 时，使用 Block-Attention 可以将 TTFT 减少 48%，将 FLOPs-TFT 减少 90.1%。作者们认为，文本越长，Block-Attention 越重要。此外，Block-Attention 的应用并不局限于 RAG 场景，它在其他领域也可能发挥重要作用，但由于保密原因，作者们暂未透露具体应用。

💡**模型微调与位置重新编码：**直接将模型从自注意力机制切换到块注意力机制会导致模型性能下降，因为模型在训练阶段从未见过这种编码方式。为了解决这个问题，作者们对模型进行了进一步的微调。实验结果表明，在 100-1000 步微调之后，模型能够快速适应块注意力机制，并在 RAG 数据集上的平均准确率恢复到与自注意力模型相当的水平，甚至略有提升。 此外，位置重新编码操作对于 Block-Attention 模型至关重要。去除位置重新编码会导致模型性能显著下降。因此，在使用 Block-Attention 时，需要确保正确地重新计算每个块的位置编码。

2024-11-03 11:03 北京

RAG 被公认是最有效的为 LLM 注入特定领域知识的方式。

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

在工业场景中，往往会利用检索技术来为大语言模型添加一些来自外部数据库的知识文档，从而增强大语言模型的回复可信度。一般来说，RAG 被公认是最有效的为 LLM 注入特定领域知识的方式。

然而，RAG 也有其不足之处。通常来说，在实际应用中，为确保能召回包含正确知识的文档，对于每个用户的查询，会检索多个文档（一般在 5 到 30 个之间），并把这些文档整合到输入提示中供大语言模型处理。这样一来，输入提示的序列长度增加，使得推理效率大幅降低。具体来讲，以首次生成标记的时间（TTFT）来衡量，RAG 大语言模型的推理延迟比非 RAG 大语言模型高很多。

由于数据库中同一文档经常会被不同 query 召回，大家很自然的会想到：是否能够把已经算好的文档表示（KV states）存在缓存中，以供二次使用？很遗憾， 由于自回归注意力机制的限制，大语言模型中每个文档的 KV States 都与上下文相关，所以遇到新的 query 时，模型必须重新编码 KV states 才能确保准确预测。

最近，论文《Block-Attention for Efficient RAG》为检索增强 (RAG) 场景实现了一种块状注意力机制，Block-Attention，通过分块独立编码检索到的文档，使得模型无需重复编码计算已经在其他 query 中已经见过的文档，从而实现线上推理效率的有效提升。在实验中，该方法能够让使用 RAG 技术的模型与不使用 RAG 的模型有几乎一样的响应速度。同时，该方法甚至还能略微提升在 RAG 场景下的模型准确率。

论文标题：Block-Attention for Efficient RAG

论文地址：https://arxiv.org/pdf/2409.15355

如下图所示，该方法把整个输入序列分成若干个 block，每个 block 独立计算其 KV States，只有最后一个 block 能够关注其他 blocks（在 RAG 场景中，最后一个 block 即用户的输入）。在 RAG 场景中，block-attention 让模型不再需要重复计算已经在其他 query 中见过的文档。

Block-Attention 的实现并不复杂：1）独立编码除最后一个 block 以外的所有 blocks；2）为每个 blocks 重新计算位置编码；3）将所有 blocks 拼接在一起，并计算最后一个 block 的 KV State。然而直接把模型不加任何修改的从 self-attention 切换到 block-attention 会导致大语言模型懵圈，毕竟模型在训练阶段从来没见过 block-attention 方式编码的输入。一个量化的对比是，直接切换为 block-attention 会让 Llama3-8B 在四个 RAG 数据集上的平均准确率由 67.9% 下降至 48.0%。

为了让模型适应 block-attention，作者们对模型进行了进一步微调，作者们发现在 100-1000 步微调之后，模型就能快速适应 block-attention，在四个 RAG 数据集上的平均准确率恢复至 68.4%。另外，block-attention 方式的模型在 KV cache 技术的帮助下，能达到与无 RAG 模型相似的效率。在用户输入长度为 50 而 prompt 总长度为 32K 的极端情况下，block-attention model 的首字延时（Time To First Token, TTFT）和首字浮点运算数（FLOPs To Frist Token, (FLOPs-TFT）分别能降低至 self-attention model 的 1.3% 和 0.2%，与无 RAG 模型的效率基本持平。

推理流程

关于 block-attention 的实现和详细推导，读者们请移步原文，这里主要介绍 block-attention 模型的推理流程。如下图所示，首先从缓存中查询并提取前 K 个 block 的 KV states。然后，根据每个 block 在输入序列中的位置，作者们对每个 block 的位置编码进行了重新计算。具体的操作过程详见论文的公式 3。最后，根据前 k-1 个 KV States 计算最后一个数据块的键值状态以及模型的输出。

实验结果

在实验中，作者们主要想探究两个问题的答案：1）在 RAG 场景中，block-attention 模型能否达到与自 self-attention 相同的准确率？2）block-attention 对效率的提升有多大？

对于问题一，上图给出了答案。作者们根据实验结果给出了三个结论：

1. 直接从 self-attention 切换到 block-attention 是不可取的，因为这会导致准确率急剧下降。例如，对于 Llama3-8B 和 Mistral-7B 模型，去除微调过程会导致在所有四个基准上平均绝对性能下降 21.99%。

2. 然而，如果作者们在微调阶段使用块注意力机制，那么得到的模型与自注意力模型的性能几乎相同，甚至在某些数据集上略好。例如，Mistral-7B-block-ft 在四个基准上的性能优于自回归方式训练的模型，平均准确率由 59.6% 上升至 62.3%。

3. 位置重新编码操作对于 block-attention 模型至关重要。去除它会导致性能显著下降 —— 在四个数据集上准确率平均下降 4%。

对于效率的提升，作者们也通过另一组实验进行了验证。他们将用户的问题长度固定在 50 个 token，然后逐渐增加被召回文档的数量，让输入序列总长度从 50 一直增加到 32K。模型在不同 prompt 长度下的首字延时（Time To First Token, TTFT）和首字浮点运算数（FLOPs To Frist Token, (FLOPs-TFT）如下图所示。显然，加速效果令人满意：当输入序列的长度为 512 时，使用 block-attention 可以将 TTFT 减少 48%，将 FLOPs-TFT 减少 90.1%。随着总长度的增加，block-attention 模型的 TTFT 和 FLOPs-TTF 保持基本不变的趋势。当总长度达到 32K 时，加速效果可以达到惊人的 98.7%，FLOPs-TFT 的消耗甚至减少了 99.8%。作者们将此实验结果总结为：“文本越长，block-attention 越重要”。

作者们最后还指出，block-attention 在很多场景中都有着重要作用，并不局限于 RAG。由于一些保密原因，作者们暂时无法透露在其他工业应用中是如何使用它的。作者们期待社区的研究人员能够进一步探索 block-attention 的潜力，并将其应用于合适的场景。

© THE END 

转载请联系本公众号获得授权

投稿或寻求报道：liyazhou@jiqizhixin.com

跳转微信打开