信息檢索 (Information Retrieval, IR) 從互聯網誕生之日起，便有著不可撼動的地位。如何從海量數據中找到用戶需要的信息是當前研究的熱點。目前比較流行的 IR 方法是先檢索后排序（retrieve-then-rank）策略。在檢索算法中，比較常用的是基于反向索引或最近鄰搜索，其中基于對比學習的雙編碼器 (dual encoders，DE) 是目前性能最優的模型。

近日，谷歌研究院在論文《Transformer Memory as a Differentiable Search Index》中提出了一種替代架構，研究者采用序列到序列 (seq2seq) 學習系統。該研究證明使用單個 Transformer 即可完成信息檢索，其中有關語料庫的所有信息都編碼在模型的參數中。

該研究引入了可微搜索索引（Differentiable Search Index，DSI），這是一種學習文本到文本新范式。DSI 模型將字符串查詢直接映射到相關文檔；換句話說，DSI 模型只使用自身參數直接回答查詢，極大地簡化了整個檢索過程。

此外，本文還研究了如何表示文檔及其標識符的變化、訓練過程的變化以及模型和語料庫大小之間的相互作用。實驗表明，在適當的設計選擇下，DSI 明顯優于雙編碼器模型等強大基線，并且 DSI 還具有強大的泛化能力，在零樣本設置中優于 BM25 基線。

論文鏈接：https://arxiv.org/pdf/2202.06991.pdf

DSI 體系架構與 DE 比較：

論文一作、谷歌高級研究員 Yi Tay 表示：在這個新范式中，檢索的所有內容都映射到易于理解的 ML 任務上。索引是模型訓練的一種特殊情況，不再依賴外部不可微的 MIPS 操作進行檢索。這使得統一模型更容易。

可微搜索索引

DSI 背后的核心思想是在單個神經模型中完全參數化傳統的多階段先檢索后排序 pipeline。為此，DSI 模型必須支持兩種基本操作模式：

• 索引：DSI 模型應該學會將每個文檔內容 d_j 與其對應的 docid j （ 文檔標識符 ：document identifiers，docid）相關聯。本文采用一種簡單的序列到序列方法，將文檔 token 作為輸入并生成標識符作為輸出；
• 檢索：給定輸入查詢，DSI 模型應返回候選 docid 排序列表。本文是通過自回歸生成實現的。

在這兩個操作之后，DSI 模型可以用來索引文檔語料庫，并對可用的帶標記數據集（查詢和標記文檔）進行微調，然后用于檢索相關文檔 —— 所有這些都在單個、統一的模型中完成。與先檢索后排序方法相反，DSI 模型允許簡單的端到端訓練，并且可以很容易地用作更大、更復雜的神經模型的可微組件。

雙編碼器 (DE) 的概述，它有兩個獨立的步驟，編碼和檢索。

索引策略

Inputs2Target：研究者將其構建為 doc_tokens → docid 的 seq2seq 任務，此方式能夠以直接輸入到目標的方式將 docid 綁定到文檔 token。

Targets2Inputs：從標識符生成文檔 token，即 docid → doc token。直觀來講，這相當于訓練一個以 docid 為條件的自回歸語言模型。

雙向：其在同一個聯合訓練設置中訓練 Inputs2Targets 和 targets2input。附加一個 prefix token 以允許模型知道任務正在哪個方向執行。

表示 Docids 用于檢索

基于 seq2seq 的 DSI 模型中的檢索是通過解碼給定輸入查詢 docid 來完成的。如何有效地解碼很大程度上取決于模型中 docid 的表示方式。在本節中，研究者探討了表示 docid 的多種可能方式以及如何處理解碼。

非結構化原子標識符（Atomic Identifiers）：表示文檔最簡單的方法是為每個文檔分配一個任意的（并且可能是隨機的）唯一整數標識符，該研究將這些標識符稱為非結構化原子標識符。研究者要想使用這些標識符，一個明顯的解碼方式是學習標識符上的概率分布。在這種情況下，模型被訓練為每個唯一的 docid (|Ndocuments|) 發出一個 logit。為了適應這種情況，該研究將標準語言模型的輸出詞匯表擴展如下：

簡單的結構化字符串標識符：該研究還考慮了另一種方法，將非結構化標識符 (即任意唯一整數) 視為可標記的（tokenizable）字符串，將其稱為簡單的結構化標識符。在此標識符下，檢索是通過依次解碼一個 docid 字符串來完成的。解碼時，使用 beam search 來獲得最佳 docid。但是，使用這種策略不容易獲得 top-k 排名。不過，研究者可以徹底梳理整個 docid 空間，并獲得給定查詢的每個 docid 的可能性。

語義結構化標識符：其目標是自動創建滿足以下屬性的標識符：(1) docid 應該捕獲一些語義信息，(2) docid 的結構應該是在每一個解碼 step 之后有效地減少搜索空間。給定一個需要索引的語料庫，所有文檔都聚集成 10 個簇。每個文檔分配有一個標識符，其簇的編號從 0 到 9。下表為這個進程的偽代碼：

實驗結果

所有 DSI 模型均使用標準預訓練 T5 模型配置進行初始化。配置名稱和對應的模型參數數量為：Base (0.2B)、Large (0.8B)、XL (3B) 和 XXL (11B)。該研究用實驗驗證了上述各種策略的效果。

下表 3 給出了經過微調的 NQ10K、NQ100K 和 NQ320K 的檢索結果，表 4 給出了零樣本檢索結果。對于零樣本檢索，模型僅針對索引任務而不是檢索任務進行訓練，因此模型看不到標記查詢 → docid 數據點。

下圖 4 給出了 NQ320K 上的結果。總的來說，研究者發現直接索引方法效果最好，并且由于 docid 反復暴露于不同的 token，因此很難訓練倒排索引（ inverted index）方法。他們還發現，較短的文檔長度似乎在性能大幅下降超過 64 個 token 時效果很好，這表明當存在大量文檔 token 時，可能更難優化或有效記憶。最后，研究者還發現對文檔 token 應用集合處理或停用詞預處理沒有額外的優勢。

下圖 3 繪制了三種方法的縮放表現（以對數尺度計），它們分別是 DE、具有 naive ID 的 DSI 和具有語義 ID 的 DSI。其中，DSI (naive) 可以從 base 到 XXL 的尺度變化中獲益，并且似乎仍有改進的空間。同時，DSI (語義) 在開始時與 DE base 具有同等競爭力，但會隨尺度增加表現得更好。DE 模型在較小的參數化時基本處于穩定狀態。

下圖 5 展示了修改索引比例對檢索樣本的影響。研究者發現，索引和檢索任務之間的相互作用會顯著影響優化過程。r 值設置得過高或過低一般都會導致性能不佳。他們發現，索引比例為 32 時通常表現良好。