大型语言模型（LLMs）在通过上下文学习能力转变了许多机器学习任务的同时，也面临着训练成本高、数据泄露和幻觉等问题。为了解决这些问题，提出了检索增强生成（RAG）方法，通过集成检索系统来提高LLMs的生成能力和可信度。现有的RAG解决方案主要关注于单一方面查询的文档检索，而忽略了需要同时检索
多个不同内容文档的多方面查询。为解决这类问题，本文提出的Multi-Head RAG（MRAG）方案通过利用Transformer解码器的多头注意力层的激活，而不仅仅是前馈层，来生成嵌入。不同的注意力头可以捕获数据的不同方面，因此使用这些激活产生的嵌入能更好地表示文档和查询的各个方面。类似相关工作有：RAPTOR、Self-RAG、Chain-of-Note等不同的RAG设计方案，以及RAGAS（用于RAG评估的基准和数据集），但是这些工作并未明确针对具有高度多方面性的问题进行优化。
点关注
不迷路
MRAG研究方法
多头RAG（MRAG）以解决现有RAG方法在处理多面性查询时的不足。
具体来说，研究方法包括以下几点：
-
解码器架构：论文介绍了解码器的架构，忽略了一些不必要的细节如层归一化。输入是一个包含n个token的文本块，每个注意力头h对第i个token 的输出定义为：
其中，
这些输出组合形成第i个多头注意块的输出：
其中，矩阵是将所有注意力头的输出组合起来的线性层。
-
标准RAG公式：假设一个包含n个token的文本块作为输入文本块，该块的嵌入是通过最后一个解码器块的馈向层生成的激活向量，即：
-
多头RAG公式：MRAG的关键思想是使用最后一个注意力层的激活向量而不是最后一个馈向层的激活向量。具体公式为：
即最后一个token的所有注意力头的输出。
-
MRAG管道概述：
MRAG管道包括两个主要部分：数据准备和查询执行。
数据准备阶段，使用选定的基于解码器的嵌入模型创建多面MRAG文本嵌入及其对应的文档或文本块。查询执行阶段，首先生成输入查询的多面嵌入，然后在数据存储中找到最近的多面嵌入及其对应的文本块。
实验设计
-
数据集生成：选择概念上不同的文档类别，主要是公开可用的维基百科文章，共25个类别，每类50篇文档。每个文档的概览部分用于嵌入。 
-
查询生成：生成涉及多个方面的查询，涵盖1到20个方面。 -
指标设计：设计了一套新的指标来评估RAG方案支持多面性的能力，主要包括检索成功率、类别检索成功率和加权检索成功率。
结果分析
-
绝对检索性能：MRAG在检索成功率方面显著优于标准RAG，平均精确文档匹配的成功率提高了超过10%，类别匹配的成功率提高了超过25%。
-
相对性能提升：MRAG在不同数量的方面和不同的嵌入模型上的相对加权性能提升表明，它在所有情况下都优于标准RAG，平均提升10-20%。
-
与融合RAG的结合：MRAG与融合RAG结合后，性能进一步提升，但计算成本和延迟也相应增加。
-
真实工作场景负载：在合成法律文件和分析化工厂事故原因的真实世界工作负载中，MRAG表现出优于其他方案的优势。
论文读后感
优势
-
多头注意力机制的应用：MRAG利用Transformer解码器块的多头注意力层的激活作为检索键，提高了复杂查询的检索准确性。 -
多维度嵌入的生成：通过多头注意力机制生成的多个嵌入，使得MRAG能够更好地表示数据项和查询的多维度特性，从而提高检索的准确性和相关性。 -
评估方法和指标：提出了一套全面的评估方法学和指标，包括合成数据集和真实世界用例，展示了MRAG的有效性。 -
与现有框架的兼容性：MRAG可以无缝集成到现有的RAG框架和不同类型的数据存储系统中。 -
成本效益和能效：MRAG不需要额外的LLM查询、多个模型实例、增加的存储或多次嵌入模型的推理传递，具有高效性和成本效益。
不足
-
单维度任务的性能：尽管MRAG在多维度查询任务中表现出色，但在仅需单维度检索的任务中，其性能并未显著下降。未来可进一步探讨MRAG在单维度任务中的表现。 -
不同嵌入层的选择：实验表明，MRAG在最后一个多头注意力层提取的多维度嵌入效果最佳，但未详细探讨其他嵌入层的效果。 -
投票策略的优化：虽然采用了指数降低重要性的投票策略，但其他策略可能更有效。未来可以探索更多有效的投票策略以提高MRAG的性能。
结论
本文提出的MRAG方案通过利用Transformer解码器块的多头注意力层的激活，显著提高了复杂多方面查询的检索准确性。实验结果表明，MRAG在检索成功率和类别检索成功率方面均优于现有的RAG基线。此外，MRAG还具有成本低、计算效率高和易于集成等优点。本文的研究为进一步改进LLMs和RAG系统提供了新的思路和方法。
引用
-
📖 https://arxiv.org/pdf/2406.05085v1 -
🔍 https://github.com/spcl/MRAG
关于社区
『社区持续纳新中』主要关注LLM、RAG、智能文档等技术方向，并持续提供AI行业热点资讯、精选论文解读。欢迎加入。
-
加入方式1：点击➡️加入·技术交流群 -
加入方式2：关注 公众号，在菜单栏中点击【会员社区】->【技术社区群】
特此声明：本文未经许可不得转载、私自发布在其他平台。后台私信获取转发、转载授权。
往期推荐
原文始发于微信公众号（老贾探AI）：『MRAG』Multi-Head RAG：基于多注意力头检索多文档内容的解决方案
