亞馬遜雲創新「神經稀疏檢索」：僅需要文字匹配就能實現語義搜尋

机器之心發表於2024-07-01
亞馬遜
圖片
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本文作者是來自 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、全文搜尋等。

相關文章