亞馬遜云創新「神經稀疏檢索」：僅需要文本匹配就能實現語義搜索
本文作者是來自 OpenSearch 中國研發團隊的機器學習負責人楊揚博士以及機器學習工程師耿志超和管聰。OpenSearch 是一個由亞馬遜云科技發起的純開源搜索和實時分析引擎項目。目前軟件超過 5 億下載量，社區在全球擁有 70 個以上的企業合作伙伴。
自從大模型爆火以來，語義檢索也逐漸成為一項熱門技術。尤其是在 RAG（retrieval augmented generation）應用中，檢索結果的相關性直接決定了 AI 生成的最終效果。
目前市面上絕大部分的語義檢索實現方案，都是利用語言模型（Language Model）將一串文本編碼為一個高維向量，并利用近似 k - 鄰近搜索（k-NN）進行檢索。面對 VectorDB 和語言模型部署（需要 GPU）高昂的費用，很多人望而卻步。
近日，亞馬遜 OpenSearch 連同亞馬遜上海人工智能研究院，在 OpenSearch NeuralSearch 插件中推出了 Neural Sparse 功能，解決了當前語義檢索正在面臨的以下三個挑戰：
- 相關性表現在不同查詢上的穩定性：zero-shot 語義檢索要求語義編碼模型在不同背景的數據集上都有不錯的相關性表現，即要求語言模型即開即用，無需用戶在自己的數據集上 fine-tune。利用稀疏編碼與詞向量（Term Vector）同源的特性，Neural Sparse 可以在遇到陌生文字表述（行業專有詞、縮寫等等）的時候向文本匹配降級，從而避免離譜的檢索結果。
- 在線搜索的時間效率：低時延對于實時檢索應用的意義是顯而易見的。目前流行的語義檢索方法一般都會包含語義編碼以及索引兩個過程，這兩者的速度決定了一個檢索應用端到端的檢索效率。Neural Sparse 獨特的 doc-only 模式，無需在線編碼，即能在與文本匹配相近的時延情況下，達成與一流語言模型相媲美的語義檢索的精度。
- 索引的存儲資源消耗：商業化的檢索應用對存儲資源的消耗是非常敏感的。在對海量數據進行索引時，搜索引擎的運行成本與存儲資源的消耗強相關。在相關實驗中，索引相同規模的數據，Neural Sparse 僅需要 k-NN 索引的 1/10。同時內存消耗也大大小于 k-NN 索引。
Relevance Demo
- 文檔主頁：https://opensearch.org/docs/latest/search-plugins/neural-sparse-search/
- 項目 Github 地址：https://github.com/opensearch-project/neural-search
技術亮點
稀疏編碼與原生 Lucene 索引結合
當前語義檢索的主要方法來自于稠密編碼（Dense Encoding），待檢索的文檔以及查詢文本都會被語言編碼模型轉換為一個高維空間中的向量。例如 Sentence-BERT 中的 TASB 模型會生成 768 維的向量，All-MiniLM-L6 則會將文本轉化為 384 維的向量。這一類高維向量的索引需要用到特殊的 k-NN 搜索引擎，例如最早基于樹結構的 FLANN、基于哈希的 LSH、還有后來出現基于鄰近圖與跳表的 HNSW 以及最新基于量化的 FAISS 引擎。
而稀疏編碼（Sparse Encoding）則會將文本轉化為一組 token 與權值的組合。這里的 token 是語言編碼模型采用分割器對文本進行切割后產生的文字單元。例如使用 WordPiece 分割器，token 可以在一定程度上理解為「單詞」，但也會出現單詞過長而被切分為兩個 token 的情況。
稀疏編碼與稠密編碼對比
由于稀疏編碼所產生的 token - 權值組合，與傳統文本匹配方法采用的 term-vector 非常類似，所以在 OpenSearch 中可以采用原生的 Lucene 索引去存儲文檔稀疏編碼。相較于 k-NN 搜索引擎，原生的 Luence 引擎會更加輕便，占用的資源也較少。
下表展示了采用 Lucene 進行文本匹配，采用 k-NN 引擎存儲稠密編碼以及采用 Lucene 存儲稀疏編碼的磁盤消耗以及運行時內存（runtime RAM）消耗的比較。
*整個系統在只運行 OpenSearch 時的內存，包括 JVM 的堆內和堆外內存
在陌生數據集上的自適應性
根據 BEIR 文章中提及的，由于目前絕大部分的稠密編碼模型都是基于 MSMARCO 數據集上精調（fine-tune）得到，模型在該數據集上表現非常優越。然而在其他的 BEIR 的數據集上進行 zero-shot 的測試時，稠密編碼模型在大約有 60%～70% 的數據集上的相關性無法超越 BM25。這一點也可以從我們自己復現的對比實驗中間看出（見下表）。
部分數據集上幾個方法的相關性表現比較
我們在實驗中發現稀疏編碼在陌生數據集上的表現要優于稠密編碼。雖然目前還沒有更加詳細的量化數據來印證，但根據在部分樣本上的分析，其優勢主要在兩點：1）稀疏編碼在近義詞的聯想方面更加突出，2）在遇到完全陌生的文本表述，例如一些專業術語，稀疏編碼會更傾向于增強這些術語 token 的權值而弱化聯想出的 token 的權值，使得檢索過程向關鍵詞匹配退化，追求的一個穩定的相關性表現。
在 BEIR 基準上的實驗中我們可以看到，Neural Sparse 的兩個方法相較于稠密編碼模型以及 BM25，相關性得分更高。
極致速度：僅文檔編碼模式
Neural Search 同時提供了一種能夠提供極致線上檢索速度的模式。在這種模式下，僅有待檢索的文檔會進行稀疏編碼。相反，在在線檢索的過程中，查詢文本并不會調用語言編碼模型進行編碼。而僅僅使用分割器（tokenizer）對查詢文本進行分割。由于省去了對深度學習模型的調用過程，不但大大降低了在線檢索的時延，也節省了模型推理所需要的大量計算資源，例如 GPU 算力等。
下表比較了文本匹配檢索方法 BM25、稠密編碼檢索 BERT-TASB 模型、稀疏編碼檢索帶查詢編碼 bi-encoder 方式以及稀疏編碼檢索僅文檔編碼 doc-only 在 MSMARCO v2 一百萬量級數據集上的速度對比。我們可以清楚地看到僅文擋編碼模式具有和 BM25 相近的速度表現，而且從上一節的表格中我們可以看到僅文檔編碼模式的相關性表現，并沒有與帶查詢稀疏編碼的方法差太多。可以說，僅文檔編碼模式是一個非常具有性價比的選擇。
還要更快：使用兩段式搜索進行加速
前文中提到，在稀疏編碼的過程中，文本被轉化為一組 token 與權值的組合。這種轉化產生了大量權值較低的 token，這些 token 雖然在搜索過程中占用了大部分時間，但對最終搜索結果的貢獻并不顯著。
因此，我們提出了一種新的搜索策略，首先在第一次搜索中過濾掉這些低權值 token，僅依賴高權值 token 來定位排名較高的文檔。隨后在這些精選的文檔上，重新引入之前被過濾的低權值 token 進行第二次詳細評分，從而獲取最終得分。
通過這種方法，我們顯著減少了兩部分的延時：首先，在第一階段搜索中，僅通過高權值 token 在倒排索引中進行匹配，大幅減少了不必要的計算時間。其次，在精確的小范圍結果文檔內再次評分時，我們僅對具有潛在相關性的文檔計算低權值 token 的分數，進一步優化了處理時間。
最終，這種改進的方法在僅文檔編碼模式（doc-only）上實現了與 BM25 搜索接近的延時表現，在帶查詢編碼模式 (bi-encoder) 上則加快了 5 到 8 倍，極大地提升了 Neural Search 的延時表現和吞吐量。以下是在四個典型 BEIR 數據集上標準 neural sparse、兩階段 neural sparse、BM25 的時延對比：
兩階段式搜索速度對比
用 5 步在 OpenSearch 中搭建 Neural Sparse 語義檢索應用
1. 設置啟用 Neural Search
首先設置集群配置來使得模型可以在本地集群上運行。
PUT /_cluster/settings
{
"transient" : {
"plugins.ml_commons.allow_registering_model_via_url" : true,
"plugins.ml_commons.only_run_on_ml_node" : false,
"plugins.ml_commons.native_memory_threshold" : 99
}
}2. 部署編碼器
Opensearch 目前開源了 3 個模型。相關注冊信息都可以在官方文檔中獲取。我們以 amazon/neural-sparse/opensearch-neural-sparse-encoding-v1 為例，首先使用 register API 來注冊：
POST /_plugins/_ml/models/_register?deploy=true
{
"name": "amazon/neural-sparse/opensearch-neural-sparse-encoding-v1",
"version": "1.0.1",
"model_format": "TORCH_SCRIPT"
}在集群的返回中，可以看到 task_id
{
"task_id": "<task_id>",
"status": "CREATED"
}用 task_id 來得到詳細的注冊信息：
GET /_plugins/_ml/tasks/在 API 返回中，我們可以拿到具體的 model_id:
{
"model_id": "<model_id>",
"task_type": "REGISTER_MODEL",
"function_name": "SPARSE_TOKENIZE",
"state": "COMPLETED",
"worker_node": [
"wubXZX7xTIC7RW2z8nzhzw"
],
"create_time": 1701390988405,
"last_update_time": 1701390993724,
"is_async": true
}3. 設置預處理管線
在索引之前，每個文檔需要被編碼的文本字段需要被轉變成稀疏向量。在 OpenSearch 中，這一過程是通過預處理器來自動實現的。你可以使用以下 API 來創建離線索引時的處理器管線：
PUT /_ingest/pipeline/neural-sparse-pipeline
{
"description": "An example neural sparse encoding pipeline",
"processors" : [
{
"sparse_encoding": {
"model_id": "<model_id>",
"field_map": {
"passage_text": "passage_embedding"
}
}
}
]
}如果需要開啟兩階段加速功能 (非必需功能)，則需要建立一個兩階段搜索管線，并在索引建立之后設置為默認的搜索管線。
建立一個默認參數的兩階段加速搜索管線方式如下，更詳細的參數設置和意義請參考 2.15 及以后版本的 OpenSearch 官方文檔。
PUT /_search/pipeline/two_phase_search_pipeline
{
"request_processors": [
{
"neural_sparse_two_phase_processor": {
"tag": "neural-sparse",
"description": "This processor is making two-phase processor."
}
}
]
}4. 設置索引
神經稀疏搜索利用 rank_features 字段類型來存儲編碼得到的詞元和相對應的權重。索引將使用上述預處理器來編碼文本。我們可以按以下方式創建索一個包含兩階段搜索加速管線的索引（如果不想開啟此功能，可把 `two_phase_search_pipeline` 替換為 `_none` 或刪除 `settings.search` 這一配置單元）。
PUT /my-neural-sparse-index
{
"settings": {
"ingest":{
"default_pipeline":"neural-sparse-pipeline"
},
"search":{
"default_pipeline":"two_phase_search_pipeline"
}
},
"mappings": {
"properties": {
"passage_embedding": {
"type": "rank_features"
},
"passage_text": {
"type": "text"
}
}
}
}5. 使用預處理器導入文檔并搜索
在設置索引之后，客戶可以提交文檔。客戶提供文本字段，而攝取進程將自動將文本內容轉換為稀疏向量，并根據預處理器中的字段映射 field_map 將其放入 rank_features 字段：
PUT /my-neural-sparse-index/_doc/
{
"passage_text": "Hello world"
}在索引中進行稀疏語義搜索的接口如下，將 <model_id> 替換為第二步中注冊的 model_id：
GET my-neural-sparse-index/_search
{
"query": {
"neural_sparse": {
"passage_embedding": {
"query_text": "Hi world",
"model_id": <model_id>
}
}
}
}關于 OpenSearch
OpenSearch 是一種分布式、由社區驅動并取得 Apache 2.0 許可的 100% 開源搜索和分析套件，可用于一組廣泛的使用案例，如實時應用程序監控、日志分析和網站搜索。OpenSearch 提供了一個高度可擴展的系統，通過集成的可視化工具 OpenSearch 控制面板為大量數據提供快速訪問和響應，使用戶可以輕松地探索他們的數據。
OpenSearch 由 Apache Lucene 搜索庫提供技術支持，它支持一系列搜索及分析功能，如 k - 最近鄰（KNN）搜索、SQL、異常檢測、Machine Learning Commons、Trace Analytics、全文搜索等。
本文轉自 機器之心 ，作者：機器之心
