奇思妙想：多头RAG

大模型技术论文不断，每个月总会新增上千篇。本专栏精选论文重点解读，主题还是围绕着行业实践和工程量产。若在某个环节出现卡点，可以回到大模型必备腔调或者LLM背后的基础模型新阅读。而最新科技（Mamba,xLSTM,KAN）则提供了大模型领域最新技术跟踪。若对于具身智能感兴趣的请移步具身智能专栏。技术宅麻烦死磕AI架构设计。

RAG

检索增强生成 （RAG） 是一种通过集成文档检索系统来增强大型语言模型 （LLMs） 功能的方法。这种集成允许LLMs从外部来源获取相关信息，从而提高生成的响应的准确性和相关性。这种方法解决了传统LLMs方法的局限性，例如需要广泛的培训以及提供过时或不正确信息的风险。RAG的主要优势在于它能够将模型的输出建立在可靠的来源中，从而减少幻觉并确保最新的知识，而无需昂贵的持续培训。

RAG 面临的一个重大挑战是处理需要具有不同内容的多个文档的查询。此类查询在各个行业中都很常见，但会带来困难，因为所需的文档可能具有截然不同的嵌入，因此很难准确检索所有相关信息。这个问题需要一个能够有效地从多个来源获取和组合信息的解决方案。在复杂的场景中，如化工厂事故，从与设备维护、天气状况和工人管理等各个方面相关的文档中检索数据对于提供全面的答案至关重要。

现有的 RAG 解决方案通常使用 Transformer 模型的最后一层解码器模块的嵌入来检索文档。但是，此方法需要充分处理多方面查询，因为它难以检索涵盖明显不同内容方面的文档。当前的一些技术包括 RAPTOR、Self-RAG 和 Chain-of-Note，它们专注于提高检索准确性，但无法有效处理复杂的多方面查询。这些方法旨在优化检索数据的相关性，但需要帮助来处理多方面查询所需的文档内容的多样性。

MRAG

来自苏黎世联邦理工学院、Cledar、巴斯夫欧洲公司和华沙理工大学的研究人员推出了多头 RAG （MRAG） 来解决多方面查询问题。这种新颖的方案利用了 Transformer 模型的多头注意力层的激活，而不是最后一层解码器的激活。研究团队设计了MRAG，利用不同的注意力头来捕获各种数据方面，从而提高复杂查询的检索准确性。<若读者对于多头注意力机制有所忘记的话，可以移步“深入注意力机制”。> 通过利用多头注意力机制，MRAG创建了代表数据在各个维度的嵌入，增强系统跨不同内容区域获取相关信息的能力。

MRAG的关键创新是使用来自多个注意力头的激活来创建嵌入。Transformer 模型中的每个注意力头都可以学习捕获不同的数据方面，从而产生表示数据项和查询各个方面的嵌入。与标准RAG相比，此方法使 MRAG能够更有效地处理多方面查询，而不会增加空间要求。MRAG在数据准备阶段使用来自多头注意力层的激活来构建嵌入。在查询执行期间，这些多方面嵌入允许从不同的嵌入空间中检索相关的文本块，从而解决多方面查询的复杂性。

MRAG背后的关键思想很简单，利用最后一个注意层为最后的Token生成的H个独立激活向量，而不是利用最后一个解码器的前馈层为最后的Token生成的单个激活向量。

对于输入的最后一个Token(𝑥𝑛)，在还没有传送到前馈层前，利用多个注意力头的激活向量构造嵌入：S={ek∀k}，其中ek = headk(𝑥𝑛)，S是输入的最后一个Token(𝑥𝑛)上注意头的所有输出集合。

由于使用多个头进行处理不会改变输出向量的大小，因此S具有与标准 RAG 相同的空间要求。但是，由于在将单独的嵌入与Wo混合之前捕获它们，推测这些向量可以提供有关输入不同部分的关注内容信息，从而有助于捕获更多维的信息。如下图所示，传统的RAG将每个文档块放在一个空间，而MRAG则是将原来的空间切分为更小的子空间。

既然将文本的索引切分到多个子空间，而检索策略需要升级，进而在给定用户查询下从数据库中选择最接近的文本块。通常，MRAG检索策略包括三个步骤。第一，在数据准备期间，为所有h个嵌入空间分配重要性分数。直观地说，这些分数反映了不同空间可能与所用数据的重要性。

头部h的得分 si 由两部分组成，ai 和 bi。向量空间i指的是被刚才算法切分形成的h个子空间的某一个。ai 是向量空间i中所有嵌入的 L2 范数的平均值；它表示给定头部的重要性：范数越大，对该注意力头部的关注就越多。bi 是向量空间i中所有嵌入之间的余弦距离的平均值。直观地讲，bi 是测量向量空间 i 的“分散”。将 si导出为 ai · bi 的乘积可确保奖励具有高平均注意力和高平均分散率的子空间，但同时惩罚相反的子空间。

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

本次的实验将MTEB排行榜上的两个嵌入模型视为候选模型。SFR-embedding-Model和e5-mistral-7b-instruct均基于Mistral 7B架构，每个多头注意力有32 个解码器块和32个注意力头。

上面是一个示例，最上面为查询的Prompt模板。第二行左一为知识库，左二和左三分别为传统的RAG和MRAG。针对两种算法，实验的结果标记在PROMPT中。

对比结果会发现，MRAG 显著提高了检索相关性，在获取多方面文档方面，其性能比标准 RAG 基线高出 20%。该评估使用了合成数据集和真实世界的用例，证明了 MRAG 在不同场景中的有效性。例如，在涉及多方面维基百科条目的测试中，MRAG 的相关性比标准 RAG 基线提高了 20%。此外，MRAG在法律文件合成和化工厂事故分析等实际任务中的表现也展示了其实际优势。在法律文件综合任务中，MRAG从各种法律框架中检索上下文相关文件的能力尤其值得称赞。

此外，MRAG的优势不仅限于检索精度。该方法经济高效且节能，不需要额外的LLM查询、多个模型实例、增加的存储或对嵌入模型进行多次推理传递。这种效率与增强的检索准确性相结合，使MRAG 成为RAG系统领域的宝贵进步LLMs。MRAG 可以与现有的RAG框架和基准测试工具无缝集成，为复杂的文档检索需求提供多功能且可扩展的解决方案。