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

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

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

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

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

DSI 體系架構(gòu)與 DE 比較：

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

可微搜索索引

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

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

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

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

索引策略

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

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

雙向：其在同一個聯(lián)合訓(xùn)練設(shè)置中訓(xùn)練 Inputs2Targets 和 targets2input。附加一個 prefix token 以允許模型知道任務(wù)正在哪個方向執(zhí)行。

表示 Docids 用于檢索

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

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

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

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

實驗結(jié)果

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

下表 3 給出了經(jīng)過微調(diào)的 NQ10K、NQ100K 和 NQ320K 的檢索結(jié)果，表 4 給出了零樣本檢索結(jié)果。對于零樣本檢索，模型僅針對索引任務(wù)而不是檢索任務(wù)進行訓(xùn)練，因此模型看不到標記查詢 → docid 數(shù)據(jù)點。

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

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

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