奇思妙想：多头RAG

现实场景中，很多查询需要从不同维度的文档中综合信息——比如调查一起化工厂事故，可能同时涉及设备维护记录、天气预报、工人排班等多个方面。这对检索增强生成系统来说是个不小的挑战。那么，RAG 到底是怎么运作的，又能如何应对这类“多方面查询”？

检索增强生成（RAG）是一种通过集成文档检索系统来增强大语言模型功能的方法。它让模型能从外部来源获取相关信息，从而提升生成响应的准确性和相关性。传统做法要么需要大规模训练，要么面临信息过时或错误的风险。RAG 的核心优势在于：把输出锚定在可靠来源上，减少幻觉，同时维持知识的新鲜度，还不用花费成本做持续训练。

RAG 的一大痛点就是处理需要来自多个内容差异明显的文档的查询。这类查询在各行各业都很常见，但麻烦在于，所需文档的嵌入可能截然不同，很难一次性准确检索出所有相关信息。例如在化工厂事故场景中，答案需要综合设备、天气、管理等多个方面的文档——这正是 RAG 力不从心的地方。

现有 RAG 方案通常用 Transformer 最后一层解码器的嵌入来检索文档。但这种方法面对多方面查询时捉襟见肘，因为不同的内容维度在嵌入空间里可能并不靠近。RAPTOR、Self-RAG、Chain-of-Note 等改进方案聚焦于提升检索准确性，但都无法有效处理多方面查询带来的内容多样性问题。

来自苏黎世联邦理工学院、Cledar、巴斯夫欧洲公司和华沙理工大学的研究人员推出了多头 RAG（MRAG），专门解决多方面查询问题。这一方案另辟蹊径：不是用最后一层解码器的激活，而是用 Transformer 多头注意力层的激活来构造嵌入。研究团队巧妙利用了不同注意力头会关注不同数据维度的特性，从而提升复杂查询的检索准确性。

MRAG 的关键创新在于：使用来自多个注意力头的激活来创建嵌入。Transformer 中的每个注意力头都可能学到不同的数据方面，产生的嵌入自然反映了数据项和查询的多个维度。相比标准 RAG，这种方法能更有效地处理多方面查询，而且没有额外增加空间需求。具体来说，在数据准备阶段，MRAG 利用多头注意力层的激活构建嵌入；查询时，这些多方面嵌入允许系统从不同的嵌入空间中检索相关文本块，从而应对多方面查询的复杂性。

核心思路很直接：取最后一个注意力层为最后一个 token 生成的 H 个独立激活向量，而不是取最后一个解码器前馈层的单个激活向量。对于输入的最后一个 token（尚未送到前馈层之前），用多个注意力头的激活向量构造嵌入集合 S = { e^k ∀ k }，其中 e^k = head^k()，S 就是所有注意头输出的集合。

由于多个头处理不会改变输出向量的大小，S 的空间需求与标准 RAG 相同。但在与输出权重矩阵混合之前捕捉这些向量，推测它们能提供输入不同部分的关注内容信息，从而捕获更多维的信息。如图所示，传统 RAG 把每个文档块放在单一空间，而 MRAG 则将原来的空间切分为更小的子空间。

既然文本的索引被切分到多个子空间，检索策略也得升级。MRAG 的检索通常分三步走。第一步：在数据准备阶段，为所有 h 个嵌入空间分配重要性分数。直观上说，这些分数反映了不同空间对数据的重要性。对于头部 h，得分 s 由两部分构成：a 是向量空间 i 中所有嵌入 L2 范数的平均值，表征给定头部的重要性——范数越大，对该注意力头的关注越多；b 是向量空间 i 中所有嵌入之间余弦距离的平均值，度量空间的“分散程度”。s = a × b 的乘积既能奖励高平均注意力和高分散的子空间，也能惩罚相反的情况。

第二步：查询执行时，MRAG 对每个嵌入空间分别应用传统 RAG 检索，返回每个空间的 c 个最接近文本块（总共 h × c 块）。这里需要投票策略，用预先计算的重要性分数从所有块中挑出总体排名前 k 个。

实验选用了 MTEB 排行榜上的两个嵌入模型：SFR-embedding-Model 和 e5-mistral-7b-instruct，它们都基于 Mistral 7B 架构，包含 32 个解码器块和 32 个注意力头。

上图展示了一个示例：最上面是查询的 prompt 模板，第二行左一是知识库，左二和左三分别对应传统 RAG 和 MRAG。实验结果直接在 prompt 中标记。

对比结果很明显：MRAG 在检索相关性上比标准 RAG 基线提升了 20%，尤其在获取多方面文档时效果突出。评估涵盖了合成数据集和真实世界用例，例如涉及多方面维基百科条目的测试中，MRAG 的相关性比标准 RAG 基线高出 20%。在法律文件合成和化工厂事故分析等实际任务中，MRAG 也展现了实际优势——它从不同法律框架中检索上下文相关文件的能力尤其值得称赞。

而且 MRAG 的优势不止于检索精度。它经济高效又节能，不需要额外的大模型查询、多个模型实例、增加的存储，或者对嵌入模型进行多次推理传递。这种效率加上增强的检索准确性，让 MRAG 成为 RAG 系统领域一个很有价值的进步。它可以与现有 RAG 框架和基准测试工具无缝集成，为复杂的文档检索需求提供灵活可扩展的方案。