亚马逊云创新「神经稀疏检索」：仅需要文本匹配就能实现语义搜索

AIxiv专栏是本站发布学术、技术内容的栏目。过去数年，本站AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本文作者是来自 OpenSearch 中国研发团队的机器学习负责人杨扬博士以及机器学习工程师耿志超和管聪。OpenSearch 是一个由亚马逊云科技发起的纯开源搜索和实时分析引擎项目。目前软件超过 5 亿下载量，社区在全球拥有 70 个以上的企业合作伙伴。

• 相关性表现在不同查询上的稳定性：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

                               ^{稀疏编码与稠密编码对比}

                             ^{*整个系统在只运行 OpenSearch 时的内存，包括 JVM 的堆内和堆外内存}

                              ^{部分数据集上几个方法的相关性表现比较}

                               ^{两阶段式搜索速度对比}

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": "","status": "CREATED"}

用 task_id 来得到详细的注册信息：

GET /_plugins/_ml/tasks/

在 API 返回中，我们可以拿到具体的 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}

在索引之前，每个文档需要被编码的文本字段需要被转变成稀疏向量。在 OpenSearch 中，这一过程是通过预处理器来自动实现的。你可以使用以下 API 来创建离线索引时的处理器管线：

PUT /_ingest/pipeline/neural-sparse-pipeline{  "description": "An example neural sparse encoding pipeline",  "processors" : [    {      "sparse_encoding": {        "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"      }    }  }}

PUT /my-neural-sparse-index/_doc/{   "passage_text": "Hello world"}

在索引中进行稀疏语义搜索的接口如下，将 替换为第二步中注册的 model_id：

GET my-neural-sparse-index/_search{  "query":{    "neural_sparse":{      "passage_embedding":{        "query_text": "Hi world",        "model_id":       }    }  }}