近日，來自小紅書搜索算法團隊的論文《Generative Dense Retrieval: Memory Can Be a Burden》被自然語言處理領域國際會議 EACL 2024 接收為 Oral，接受率為 11.32%（144/1271）。

他們在論文中提出了一種新穎的信息檢索范式——生成式密集檢索（Generative Dense Retrieval，GDR）。該范式能夠很好地解決傳統(tǒng)生成式檢索（Generative Retrieval，GR）在處理大規(guī)模數(shù)據(jù)集時所面臨的挑戰(zhàn)。它是從記憶機制得到的靈感。

在過往的實踐中，GR 憑借其獨特的記憶機制，實現(xiàn)了查詢與文檔庫間的深度交互。然而，這種依賴于語言模型自回歸解碼的方法，在處理大規(guī)模數(shù)據(jù)時存在著明顯的局限性，包括細粒度文檔特征模糊、文檔庫規(guī)模受限、索引更新困難等。

小紅書提出的 GDR 采用由粗到細的兩階段檢索思想，首先利用語言模型有限的記憶容量，實現(xiàn)查詢到文檔簇的映射，然后通過向量匹配機制完成文檔簇到文檔的精細映射。GDR 通過引入密集檢索的向量匹配機制，有效緩解了 GR 的固有弊端。

此外，團隊還設計了「記憶友好的文檔簇標識符構建策略」與「文檔簇自適應的負采樣策略」，分別提升兩階段的檢索性能。在 Natural Questions 數(shù)據(jù)集的多個設定下，GDR 不僅展現(xiàn)了 SOTA 的 Recall@k 表現(xiàn)，更在保留深度交互優(yōu)勢的同時實現(xiàn)了良好的可擴展性，為信息檢索的未來研究開辟了新的可能性。

1.背景

文本檢索具有重要的研究與應用價值。傳統(tǒng)檢索范式，如基于字詞匹配度的稀疏檢索（sparse retrieval，SR）和基于語義向量匹配度的密集檢索（dense retrieval，DR），雖然各有千秋，但隨著預訓練語言模型的興起，基于此的生成式檢索范式開始嶄露頭角。

在訓練過程中，模型以給定查詢作為上下文，自回歸地生成相關文檔的標識符。這一過程事實上完成了模型對于候選語料庫的記憶。查詢進入模型后與模型參數(shù)交互并自回歸解碼的過程，隱式地發(fā)生了查詢與候選語料庫的深度交互，而這種深度交互正是 SR 和 DR 所欠缺的。因此，當模型能夠準確記憶候選文檔時，GR 能夠表現(xiàn)出優(yōu)異的檢索性能。

然而，GR 的記憶機制并非無懈可擊。我們通過經典 DR 模型（AR2）與 GR 模型（NCI）之間的對比實驗，證實了記憶機制至少會帶來三大挑戰(zhàn)：

1）細粒度文檔特征模糊：

我們分別計算了 NCI 與 AR2 在由粗到細解碼文檔標識符的每一位時發(fā)生錯誤的概率。對于 AR2，我們通過向量匹配找到給定查詢最相關的文檔對應的標識符，再統(tǒng)計標識符的首次出錯步數(shù)，得到 AR2 對應的分步解碼錯誤率統(tǒng)計。如表 1 所示，NCI 在解碼的前半段中表現(xiàn)良好，而后半段中錯誤率則較高，AR2 與之相反。這說明 NCI 通過整體記憶候選文檔庫，能較好地完成查詢到候選文檔語義空間的粗粒度映射。但是由于訓練過程中候選文檔的特征是由查詢來刻畫的，因此其細粒度的特征難以被準確記憶，故而在細粒度映射時表現(xiàn)不佳。

2）文檔庫規(guī)模受限：

如表 2 所示，我們分別以 334K 的候選文檔庫大小（第一行）和 1M 的候選文檔大?。ǖ诙校┯柧毩?NCI 模型并以 R@k 指標進行測試。結果表明 NCI 在 R@100 上下降了 11 point，對比之下 AR2 只下降了 2.8 point。為了探究候選文檔庫規(guī)模擴大使 NCI 性能顯著下降的原因，我們進一步測試了在 1M 文檔庫上訓練的 NCI 模型在以 334K 為候選文檔庫時的測試結果（第三行）。與第一行相比，NCI 記憶更多文檔的負擔導致了其召回性能的顯著下降，這說明模型有限的記憶容量限制了其記憶大規(guī)模的候選文檔庫。

3）索引更新困難：

當新文檔需要加入候選庫時，需要更新文檔標識符，并且需要重新訓練模型以重新記憶所有文檔。否則，過時的映射關系（查詢到文檔標識符和文檔標識符到文檔）將顯著降低檢索性能。

上述問題阻礙了 GR 在真實場景下的應用。為此，我們在分析后認為 DR 的匹配機制與記憶機制有著互補的關系，因此考慮將其引入 GR，在保留記憶機制的同時抑制其帶來的弊端。我們提出了生成式密集檢索新范式（Generative Dense Retrieval，GDR）：

• 我們整體設計了由粗到細的兩階段檢索框架，利用記憶機制實現(xiàn)簇間匹配（查詢到文檔簇的映射），通過向量匹配機制完成簇內匹配（文檔簇到文檔的映射）。
• 為了協(xié)助模型記憶候選文檔庫，我們構建了記憶友好的文檔簇標識符構建策略，以模型記憶容量為基準控制文檔簇的劃分粒度，增益簇間匹配效果。
• 在訓練階段，我們依據(jù)兩階段檢索的特點提出文檔簇自適應的負采樣策略，增強簇內負樣本的權重，增益簇內匹配效果。

2.1 基于記憶機制的簇間匹配

以查詢作為輸入，我們利用語言模型記憶候選文檔庫，并自回歸生成 k 個相關文檔簇（CID），完成如下映射：

在這一過程中，CID 的生成概率為：

其中 

 是編碼器產生的所有查詢嵌入，

 是編碼器產生的一維查詢表征。該概率同時作為簇間匹配分數(shù)被存儲，參與后續(xù)運算。基于此，我們采用標準交叉熵損失訓練模型：

2.2 基于向量匹配機制的簇內匹配

我們進一步從候選文檔簇內檢索候選文檔，完成簇內匹配：

我們引入一個文檔編碼器提取候選文檔的表征，這一過程會離線完成。以此為基礎，計算簇內文檔與查詢間的相似度，作為簇內匹配分數(shù)：

在這一過程中，NLL loss 被用來訓練模型：

最后，我們計算文檔的簇間匹配分數(shù)與簇內匹配分數(shù)的加權值并進行排序，選出其中的 Top K 作為檢索出的相關文檔：

其中 beta 在我們的實驗中設定為 1。

2.3 記憶友好的文檔簇標識符構建策略

為了充分利用模型有限的記憶容量實現(xiàn)查詢與候選文檔庫之間的深度交互，我們提出記憶友好的文檔簇標識符構建策略。該策略首先以模型記憶容量為基準，計算簇內文檔數(shù)上限：

在此基礎上，進一步通過 K-means 算法構建文檔簇標識符，保障模型的記憶負擔不超過其記憶容量：

2.4 文檔簇自適應的負采樣策略

GDR 兩階段的檢索框架決定了在簇內匹配過程中簇內的負樣本所占比重更大。為此，我們在第二階段訓練過程中以文檔簇劃分為基準，顯式增強了簇內負樣本的權重，從而獲得更好的簇內匹配效果：

3.實驗

實驗中使用的數(shù)據(jù)集為 Natural Questions (NQ)，它包含 58K 個訓練對（查詢和相關文檔）以及 6K 個驗證對，伴隨著 21M 個候選文檔庫。每個查詢多個相關文檔，這對模型的召回性能提出了更高的要求。為了評估 GDR 在不同規(guī)模文檔庫上的性能，我們構建了 NQ334K、NQ1M、NQ2M 和 NQ4M 等不同設置，通過向 NQ334K 添加來自完整 21M 語料庫的其余段落來實現(xiàn)。GDR 在每個數(shù)據(jù)集上分別生成 CIDs，以防止更大候選文檔庫的語義信息泄露到較小的語料庫中。我們采用 BM25（Anserini 實現(xiàn)）作為 SR 基線，DPR 和 AR2 作為 DR 基線，NCI 作為 GR 的基線。評價指標包括 R@k 和 Acc@k。

3.1 主實驗結果

在 NQ 數(shù)據(jù)集上，GDR 在 R@k 指標上平均提高了 3.0，而在 Acc@k 指標上排名第二。這表明 GDR 通過粗到細的檢索過程，最大化了記憶機制在深度交互和匹配機制在細粒度特征辨別中的優(yōu)勢。

3.2 擴展到更大的語料庫

我們注意到當候選語料庫擴展到更大的規(guī)模時，SR 和 DR 的 R@100 下降率保持在 4.06% 以下，而 GR 在所有三個擴展方向上的下降率超過了 15.25%。相比之下，GDR 通過將記憶內容集中在固定體量的語料庫粗粒度特征上，實現(xiàn)了平均 3.50% 的 R@100 下降率，與 SR 和 DR 相近。

3.3 消融實驗

表 3 中 GDR-bert 與 GDR-ours 分別代表了傳統(tǒng)和我們的 CID 構建策略下對應的模型表現(xiàn)，實驗證明使用記憶友好的文檔簇標識符構建策略，可以顯著減輕記憶負擔，從而帶來更好的檢索性能。此外，表 4 表明 GDR 訓練時采用的文檔簇自適應的負采樣策略，通過提供更多的文檔簇內辨別信號，增強了細粒度匹配能力。

3.4 新文檔加入

當有新文檔加入候選文檔庫時，GDR 將新文檔加入距離最近的文檔簇聚類中心，并賦予相應標識符，同時通過文檔編碼器提取向量表征更新向量索引，從而完成對新文檔的快速擴展。如表 6 所示，在添加新文檔到候選語料庫的設定下，NCI 的 R@100 下降了 18.3 個百分點，而 GDR 的性能僅下降了 1.9 個百分點。這表明 GDR 通過引入匹配機制緩解記憶機制的難以擴展性，在無需重新訓練模型的情況下保持了良好的召回效果。

3.5 局限性

受限于語言模型自回歸生成的特點，盡管 GDR 在第二階段引入了向量匹配機制，相比于 GR 實現(xiàn)了顯著的檢索效率提升，但相比于 DR 與 SR 仍有較大的提升空間。我們期待未來有更多的研究幫助緩解記憶機制引入檢索框架時帶來的時延問題。

4.結語

本項研究中，我們深入探討了記憶機制在信息檢索中的雙刃劍效應：一方面這一機制實現(xiàn)了查詢與候選文檔庫的深度交互，彌補了密集檢索的不足；另一方面模型有限的記憶容量與更新索引的復雜性，它在面對大規(guī)模和動態(tài)變化候選文檔庫時顯得捉襟見肘。為了解決這一難題，我們創(chuàng)新性地將記憶機制與向量匹配機制進行層次化結合，實現(xiàn)兩者揚長避短、相得益彰的效果。

我們提出了一個全新的文本檢索范式，生成式密集檢索（GDR）。GDR 該范式對于給定查詢進行由粗到細的兩階段檢索，先由記憶機制自回歸地生成文檔簇標識符實現(xiàn)查詢到文檔簇的映射，再由向量匹配機制計算查詢與文檔間相似度完成文檔簇到文檔的映射。

記憶友好的文檔簇標識符構建策略保障了模型的記憶負擔不超過其記憶容量，增益簇間匹配效果。文檔簇自適應的負采樣策略增強了區(qū)分簇內負樣本的訓練信號，增益簇內匹配效果。大量實驗證明，GDR 在大規(guī)模候選文檔庫上能夠取得優(yōu)異的檢索性能，同時能夠高效應對文檔庫更新。

作為一次對傳統(tǒng)檢索方法進行優(yōu)勢整合的成功嘗試，生成式密集檢索范式具有召回性能好、可擴展性強、在海量候選文檔庫場景下表現(xiàn)穩(wěn)健等優(yōu)點。隨著大語言模型在理解與生成能力上的不斷進步，生成式密集檢索的性能也將進一步提升，為信息檢索開辟更加廣闊的天地。

論文地址：https://arxiv.org/abs/2401.10487

5.作者簡介

• 袁沛文
  現(xiàn)博士就讀于北京理工大學，小紅書社區(qū)搜索組實習生，在 NeurIPS、ICLR、AAAI 、EACL 等發(fā)表多篇一作論文。主要研究方向為大語言模型推理與評測、信息檢索。
• 王星霖
  現(xiàn)博士就讀于北京理工大學，小紅書社區(qū)搜索組實習生，在 EACL、NeurIPS、ICLR 等發(fā)表數(shù)篇論文，在國際對話技術挑戰(zhàn)賽 DSTC11 上獲得測評賽道第二名。主要研究方向為大語言模型推理與測評、信息檢索。
• 馮少雄
  負責小紅書社區(qū)搜索向量召回。博士畢業(yè)于北京理工大學，在 ICLR、AAAI、ACL、EMNLP、NAACL、EACL、KBS 等機器學習、自然語言處理領域頂級會議/期刊上發(fā)表數(shù)篇論文。主要研究方向為大語言模型測評推理蒸餾、生成式檢索、開放域對話生成等。
• 道玄
  小紅書交易搜索團隊負責人。博士畢業(yè)于浙江大學，在 NeurIPS、ICML 等機器學習領域頂級會議上發(fā)表數(shù)篇一作論文，長期作為多個頂級會議/期刊審稿人。主要業(yè)務覆蓋內容搜索、電商搜索、直播搜索等。
• 曾書
  碩士畢業(yè)于清華大學電子系，在互聯(lián)網領域從事自然語言處理、推薦、搜索等相關方向的算法工作，目前在小紅書社區(qū)搜索負責召回和垂類搜索等技術方向。