MRAG：一種多頭RAG，提升具有顯著不同內容的多個文檔復雜查詢的檢索準確性
現有的RAG方案沒有專注于需要檢索多個具有顯著不同內容的文檔的查詢，這類查詢經常出現，但挑戰在于這些文檔的嵌入可能在嵌入空間中距離很遠，難以全部檢索。
因此，提出一種多頭檢索增強生成方法（MRAG：Multi-Head RAG ），它的關鍵思想是利用Transformer解碼器的多頭注意力層的激活，而不僅僅是前饋層，來生成嵌入。不同的注意力頭可以捕獲數據的不同方面，因此使用這些激活產生的嵌入能更好地表示文檔和查詢的各個方面。
MRAG使用最后一個標記的每個注意力頭的輸出作為一組“單方面”嵌入。這些嵌入的完整集合形成了“多方面”嵌入。重要的是，這與僅使用最終解碼器層輸出的標準RAG相比，不需要額外的空間。
解碼器架構的概述，以及標準RAG和多頭RAG嵌入生成方式的比較。
MRAG（多頭RAG）流程的概述，由兩部分組成：數據準備A和查詢執行B。嵌入模型C和數據存儲D被這兩部分使用。數據存儲D包含指向反映三種不同方面（青色、品紅色、黃色）的文本塊的文本嵌入。帶有星號標記的塊是這項工作的新特點。
多頭檢索增強生成（Multi-Head RAG, MRAG）流程的概述：
- 數據準備（Data Preparation）：在數據準備階段，MRAG通過使用選定的基于解碼器的嵌入模型（decoder-based embedding model）來創建文本塊（text chunks）的多方面嵌入（multi-aspect embeddings）。這些嵌入隨后與相應的文檔或文本塊一起存儲在數據存儲（data store）中。MRAG支持不同類型的數據嵌入，包括文檔分塊、整文檔或其他類型的數據。
- 查詢執行（Query Execution）：在查詢執行階段，首先使用選定的嵌入模型為輸入查詢生成多方面嵌入。然后，使用特殊的多方面檢索策略在數據存儲中找到最接近的多方面嵌入及其對應的文本塊。最后，檢索到的數據可以選用新的度量標準進行評估，以確定其與多方面需求的對應程度。
- 多方面嵌入的存儲：與標準RAG不同，MRAG為每個多方面嵌入包含h個單方面嵌入，每個嵌入指向原始文本塊。數據存儲包含h個嵌入空間，每個空間捕獲文本的不同方面。這使得MRAG能夠在多個嵌入空間中比較查詢和文本塊，這些空間捕獲數據的多個方面。
- 多方面檢索策略：MRAG的檢索策略包括三個步驟：在數據準備期間為所有h個嵌入空間分配重要性得分；在查詢執行期間，MRAG首先為每個嵌入空間單獨應用傳統RAG檢索，返回每個嵌入空間的c個最近文本塊的列表；然后使用預計算的重要性得分通過投票策略從所有hc個塊中選擇整體前k個。
- 與數據存儲的集成：MRAG可以與不同類型的數據存儲和最近鄰（Nearest Neighbor, NN）搜索方法無縫使用。它可以結合精確和近似的最近鄰搜索來找到匹配的（嵌入，塊）對。?
用于評估不同RAG策略的示例查詢。文本中要獲取的文檔，然后評估不同RAG策略在找到這些文檔及其類別方面的成功率。標記了精確文檔匹配、類別匹配、多次匹配類別的文檔，以及沒有匹配文檔的文本段。最后，展示了每種策略的加權成功率，采用2:1的加權（優先考慮精確文章匹配）。
提供了評估方法和指標、合成數據集和現實世界用例來展示MRAG的有效性，顯示出在相關性方面比標準RAG基線提高了高達20%。
MRAG（多頭RAG）和標準RAG在25個查詢中的檢索成功率比較，每個查詢包含10個不同的方面。上半部分展示了精確文檔匹配，而下半部分展示了僅類別匹配（我們在第3節中解釋了使用的度量標準）。為了展示25個查詢中詳細的分布情況，為一個特定樣本呈現了直方圖（每個查詢獲取的文檔數量為30）。MRAG在精確文檔匹配的檢索成功率上一貫優于標準RAG（平均提高了超過10%）。在類別匹配上的檢索性能提升更為顯著（平均提高了超過25%的檢索成功率）。
MRAG相對于標準RAG在不同數量方面和不同嵌入模型的查詢中的相對檢索改進10-20%（左側為SFR，右側為e5）。
在多方面數據集、法律數據集和事故數據集上的實驗表明，MRAG在檢索成功率上優于標準RAG和Split RAG。
在構建法律文件（左）和發現工業事故原因（右）這兩個真實世界工作負載中，MRAG和Split RAG相對于標準RAG在檢索成功率上的平均提升。
通過對不同解碼器塊的嵌入和不同投票策略的分析，證明了MRAG設計和選擇的合理性。
對MRAG和Split RAG的不同投票策略的評估
https://arxiv.org/html/2406.05085v1
https://github.com/spcl/MRAG
Multi-Head RAG: Solving Multi-Aspect Problems with LLMs本文轉載自??PaperAgent??
