現(xiàn)有的RAG方案沒(méi)有專注于需要檢索多個(gè)具有顯著不同內(nèi)容的文檔的查詢，這類查詢經(jīng)常出現(xiàn)，但挑戰(zhàn)在于這些文檔的嵌入可能在嵌入空間中距離很遠(yuǎn)，難以全部檢索。

因此，提出一種多頭檢索增強(qiáng)生成方法（MRAG：Multi-Head RAG ），它的關(guān)鍵思想是利用Transformer解碼器的多頭注意力層的激活，而不僅僅是前饋層，來(lái)生成嵌入。不同的注意力頭可以捕獲數(shù)據(jù)的不同方面，因此使用這些激活產(chǎn)生的嵌入能更好地表示文檔和查詢的各個(gè)方面。

MRAG使用最后一個(gè)標(biāo)記的每個(gè)注意力頭的輸出作為一組“單方面”嵌入。這些嵌入的完整集合形成了“多方面”嵌入。重要的是，這與僅使用最終解碼器層輸出的標(biāo)準(zhǔn)RAG相比，不需要額外的空間。

解碼器架構(gòu)的概述，以及標(biāo)準(zhǔn)RAG和多頭RAG嵌入生成方式的比較。

MRAG（多頭RAG）流程的概述，由兩部分組成：數(shù)據(jù)準(zhǔn)備A和查詢執(zhí)行B。嵌入模型C和數(shù)據(jù)存儲(chǔ)D被這兩部分使用。數(shù)據(jù)存儲(chǔ)D包含指向反映三種不同方面（青色、品紅色、黃色）的文本塊的文本嵌入。帶有星號(hào)標(biāo)記的塊是這項(xiàng)工作的新特點(diǎn)。

多頭檢索增強(qiáng)生成（Multi-Head RAG, MRAG）流程的概述：

1. 數(shù)據(jù)準(zhǔn)備（Data Preparation）：在數(shù)據(jù)準(zhǔn)備階段，MRAG通過(guò)使用選定的基于解碼器的嵌入模型（decoder-based embedding model）來(lái)創(chuàng)建文本塊（text chunks）的多方面嵌入（multi-aspect embeddings）。這些嵌入隨后與相應(yīng)的文檔或文本塊一起存儲(chǔ)在數(shù)據(jù)存儲(chǔ)（data store）中。MRAG支持不同類型的數(shù)據(jù)嵌入，包括文檔分塊、整文檔或其他類型的數(shù)據(jù)。
2. 查詢執(zhí)行（Query Execution）：在查詢執(zhí)行階段，首先使用選定的嵌入模型為輸入查詢生成多方面嵌入。然后，使用特殊的多方面檢索策略在數(shù)據(jù)存儲(chǔ)中找到最接近的多方面嵌入及其對(duì)應(yīng)的文本塊。最后，檢索到的數(shù)據(jù)可以選用新的度量標(biāo)準(zhǔn)進(jìn)行評(píng)估，以確定其與多方面需求的對(duì)應(yīng)程度。
3. 多方面嵌入的存儲(chǔ)：與標(biāo)準(zhǔn)RAG不同，MRAG為每個(gè)多方面嵌入包含h個(gè)單方面嵌入，每個(gè)嵌入指向原始文本塊。數(shù)據(jù)存儲(chǔ)包含h個(gè)嵌入空間，每個(gè)空間捕獲文本的不同方面。這使得MRAG能夠在多個(gè)嵌入空間中比較查詢和文本塊，這些空間捕獲數(shù)據(jù)的多個(gè)方面。
4. 多方面檢索策略：MRAG的檢索策略包括三個(gè)步驟：在數(shù)據(jù)準(zhǔn)備期間為所有h個(gè)嵌入空間分配重要性得分；在查詢執(zhí)行期間，MRAG首先為每個(gè)嵌入空間單獨(dú)應(yīng)用傳統(tǒng)RAG檢索，返回每個(gè)嵌入空間的c個(gè)最近文本塊的列表；然后使用預(yù)計(jì)算的重要性得分通過(guò)投票策略從所有hc個(gè)塊中選擇整體前k個(gè)。
5. 與數(shù)據(jù)存儲(chǔ)的集成：MRAG可以與不同類型的數(shù)據(jù)存儲(chǔ)和最近鄰（Nearest Neighbor, NN）搜索方法無(wú)縫使用。它可以結(jié)合精確和近似的最近鄰搜索來(lái)找到匹配的（嵌入，塊）對(duì)。?

用于評(píng)估不同RAG策略的示例查詢。文本中要獲取的文檔，然后評(píng)估不同RAG策略在找到這些文檔及其類別方面的成功率。標(biāo)記了精確文檔匹配、類別匹配、多次匹配類別的文檔，以及沒(méi)有匹配文檔的文本段。最后，展示了每種策略的加權(quán)成功率，采用2:1的加權(quán)（優(yōu)先考慮精確文章匹配）。

提供了評(píng)估方法和指標(biāo)、合成數(shù)據(jù)集和現(xiàn)實(shí)世界用例來(lái)展示MRAG的有效性，顯示出在相關(guān)性方面比標(biāo)準(zhǔn)RAG基線提高了高達(dá)20%。

MRAG（多頭RAG）和標(biāo)準(zhǔn)RAG在25個(gè)查詢中的檢索成功率比較，每個(gè)查詢包含10個(gè)不同的方面。上半部分展示了精確文檔匹配，而下半部分展示了僅類別匹配（我們?cè)诘?節(jié)中解釋了使用的度量標(biāo)準(zhǔn)）。為了展示25個(gè)查詢中詳細(xì)的分布情況，為一個(gè)特定樣本呈現(xiàn)了直方圖（每個(gè)查詢獲取的文檔數(shù)量為30）。MRAG在精確文檔匹配的檢索成功率上一貫優(yōu)于標(biāo)準(zhǔn)RAG（平均提高了超過(guò)10%）。在類別匹配上的檢索性能提升更為顯著（平均提高了超過(guò)25%的檢索成功率）。

MRAG相對(duì)于標(biāo)準(zhǔn)RAG在不同數(shù)量方面和不同嵌入模型的查詢中的相對(duì)檢索改進(jìn)10-20%（左側(cè)為SFR，右側(cè)為e5）。

在多方面數(shù)據(jù)集、法律數(shù)據(jù)集和事故數(shù)據(jù)集上的實(shí)驗(yàn)表明，MRAG在檢索成功率上優(yōu)于標(biāo)準(zhǔn)RAG和Split RAG。

在構(gòu)建法律文件（左）和發(fā)現(xiàn)工業(yè)事故原因（右）這兩個(gè)真實(shí)世界工作負(fù)載中，MRAG和Split RAG相對(duì)于標(biāo)準(zhǔn)RAG在檢索成功率上的平均提升。

通過(guò)對(duì)不同解碼器塊的嵌入和不同投票策略的分析，證明了MRAG設(shè)計(jì)和選擇的合理性。

對(duì)MRAG和Split RAG的不同投票策略的評(píng)估

https://arxiv.org/html/2406.05085v1
https://github.com/spcl/MRAG
Multi-Head RAG: Solving Multi-Aspect Problems with LLMs

本文轉(zhuǎn)載自??PaperAgent??