AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：[email protected]；[email protected]
本文作者是来自 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、全文搜刮等。