使用AWS Opensearch KNN插件实现向量检索

需求场景

业界有很多对商品，内容进行分类检索的业务场景，比如图片搜索、UGC视频版权保护、人脸识别、搜索&推荐等诸多应用层面。随着最近几年向量检索作为研究热点的不断普及，数据对象更多的通过一个抽象的向量形式来表达，通过最近邻的查找，查找最符合条件的检索结果，比如：

图像&视频向量相似度检索：对于内容提供商而言，线上有几十万用戶，每个用戶保存有数千图像的特征，例如⻋辆特征，人脸特征等。当有新的某个用戶的视频发送到算法云服务时，分析提取出该视频中的⻋辆特征，将该⻋辆特征发送到数据库中，从该用戶已保存的⻋辆特征中，实时查询到最相似的特征id，返回给算法云服务。

广告&推荐系统embedding向量检索：DSP及个性化推荐系统中，经常对通过NLP，Node2Vec，Object2Vec等算法模型对商品内容标签或者用户画像标签做embedding向量化，以便用于基于向量相似度的推荐结果召回并进行TopN排序。对于海量用户及商品标签（如电商，游戏，数字媒体内容…etc)，其向量检索和相似度计算会耗费大量资源，时效性差。

解决方案

对于上述向量检索场景，业界通常采用kNN搜索算法。KNN（K-Nearest Neighbor）算法是机器学习算法中最基础、最简单的算法之一。它既能用于分类，也能用于回归。KNN通过测量不同特征值之间的距离来进行分类。KNN算法的思想非常简单：对于任意n维输入向量，分别对应于特征空间中的一个点，输出为该特征向量所对应的类别标签或预测值。

传统KNN算法是全量计算向量相似度，在海量向量数据规模的情况下，训练时间复杂度为O(n)，计算空间复杂性较高，开销巨大。另一方面随着数据的爆发式增长，对检索的时延和准确率又提出了更高的要求，因此目前业界考虑是否能将传统KNN算法思想延展到分布式系统上，利用并行计算和，降低时间复杂度，提升检索效率，节省成本。

对于分布式检索场景，目前业界最有益的选择就是OpenSearch KNN插件, Opensearch是一种分布式，由社区驱动并取得 Apache 2.0 许可的 100% 开源搜索和分析套件，它提供了一个高度可扩展的系统，通过集成的可视化工具 OpenSearch 控制面板为大量数据提供快速访问和响应，使用户可以轻松地探索他们的数据。

Amazon OpenSearch Service (AOS)是托管的opensearch服务，提供与其他 AWS 服务的集成并且您可以选择开源引擎，包括 OpenSearch 和开源Elasticsearch版本，Amazon OpenSearch在机器学习方面有深层次的融合，如RCF算法进行异常检查，以及本文将向大家介绍的KNN算法插件进行向量检索。

KNN在Amazon Opensearch上，其向量的ingestion，存储及检索都是分布式的，Amazon OpenSearch集群的每个shard负责一部分向量数据，检索的时候会在每个shard上计算num_candidates个候选向量和待检索向量的topN个近邻，每个shard的结果再最后merge成最终的topN，其架构如下所示：

Amazon Opensearch 的KNN插件支持faiss，nsmlib等多种向量相似度计算算法，通过对向量index的参数配置，可以方便的调用底层算法进行向量的检索调试和优化，并通过标准化的restful API进行调用，而不需要开发一行代码。

AOS KNN使用details

以下章节详细介绍在Amazon Opensearch托管服务上使用KNN插件进行向量检索的具体实施步骤。

AWS Opensearch 支持两种kNN搜索的方法：

近似kNN检索

多数情况下，我们只需要使用近似kNN。近似kNN提供低延迟是以较慢的索引速度和不完美的准确性为代价，它使用几种算法之一将近似k最近邻返回到查询向量。通常，这些算法牺牲索引速度和搜索精度，以换取性能优势，如更低的延迟、更小的内存占用和更具伸缩性的搜索

为了使用近似kNN搜索，需要使用kNN搜索api来搜索已经被索引的dense_vector字段明确映射一个或多个dense_vector字段，近似kNN搜索要求下列mapping选项设置index：

PUT my-knn-index-1
{
 "settings" : {
    "number_of_shards" : 9,
    "number_of_replicas" : 0
  },
  "mappings": {
    "properties": {
      "my_vector": {
        "type": "knn_vector",
        "dimension": 300,
        "method": {
        "name": "hnsw",
        "space_type": "l2",
        "engine": "faiss",
        "parameters": {
          "ef_construction": 128,
          "m": 24
         }
       }
      }
    }
  }
}

里面最重要的参数是method，该定义是指要使用的近似k-NN算法的基本配置。method定义用于创建knn_vector字段或在训练期间创建模型，然后可以使用该模型创建knn_向量字段

method定义将始终包含算法名称、为其构建方法的space_type、要使用的引擎（库）和参数映射等，是在大量向量数据时调优方向

以faiss引擎的method为例，其算法可以是hnsw和ivf，其支持的space_type为l2, innerproduct，而如果是hnsw的enginer引擎，其method配置下的space_type为l2, innerproduct, cosinesimil, l1, linf。

详细可以参考Opensearch官方KNN插件配置参数

近似KNN检索语法见如下所示：

GET my-knn-index-2/_search
{
  "size": 2,
  "query": {
    "knn": {
      "my_vector": {
        "vector": [0.17412458703561207,
            0.8247188588762435,
            0.8579923977127106,
                         ……
                        0.6437698111182779],
        "k": 10
      }
    }
  }
}

K指定每个shards分片上返回的candidate的最近邻向量相似的结果，size指定最终merge成的result的数量，如上标识从每个shard选择10个candidate候选最相似向量结果，再merge后返回最符合检索条件的2个document记录

精确kNN检索

精确KNN检索扩展了OpenSearch的脚本script功能，对“knn_vector”字段或可以表示二进制对象的字段执行强力、精确的k-NN搜索。使用这种方法，可以对索引中的向量子集运行k-NN搜索（有时称为预过滤搜索）

精确KNN检索方法适合搜索较小或需要预筛选的向量集合。在大型索引上使用此方法可能会导致比近似KNN检索更高延迟。

精确KNN检索与上文近似检索语法类似，但增加了script_score从而可以添加各种detail的过滤条件

GET my-knn-index-2/_search
{
 "size": 1,
 "query": {
   "script_score": {
     "query": {
       "match_all": {}
     },
     "script": {
       "source": "knn_score",
       "lang": "knn",
       "params": {
         "field": "my_vector",
         "query_value": [0.17412454545561207,
            0.8247188588762435,
                                …….
                         0.6437698111182779],
         "space_type": "cosinesimil"
       }
     }
   }
 }
}

注意lang需要设置为knn，source需要设置为knn_score

性能测试

我们使用opensearch-py 客户端的一个demo程序构造1百万记录，300维向量的dummy数据，用于测试Amazon Opensearch的向量化检索性能

from opensearchpy import OpenSearch, RequestsHttpConnection
import boto3
import random 
from requests_aws4auth import AWS4Auth
service = 'es'
host='search-osdemo02-fwvq3527gsykkz5tjp2iz7exam.ap-southeast-1.es.amazonaws.com'
region='ap-southeast-1'
credentials = boto3.Session().get_credentials()
awsauth = AWS4Auth(credentials.access_key, credentials.secret_key, region, service, session_token=credentials.token)

color=('red','organge','green','blue','yellow','black','pickacl','spice')
search = OpenSearch(
     hosts = [{'host': host, 'port': 443}],
     http_auth = awsauth ,
     use_ssl = True,
     verify_certs = True,
     connection_class = RequestsHttpConnection
 )
for iter in range(1000000):
    document = { "color": color[int(random.uniform(1,8))], "my_vector": [float(random.uniform(0.0, 1.0)) for _ in range(300)]} 
    search.index(index="my-knn-index-2", body=document)

对插入的100w记录向量数据做KNN检索

GET my-knn-index-2/_search
{
  "size": 2,
  "query": {
    "knn": {
      "my_vector": {
        "vector": [0.17412458703561207,
            0.8247188588762435,
            0.8579923977127106,
                         ……
                        0.6437698111182779],
        "k": 10
      }
    }
  }
}

可以看到百万级数据量下Amazon Opensearch KNN向量检索性能，在128shards，9 nodes r6g.4xlarge的集群规模下 ,检索响应时间都在秒级（这里测试的是纯计算和检索的性能，因为随机生成的dummy数据，没有考虑业务层面数据的标签的准确性）

小结

Amazon Opensearch的KNN插件可以方便的进行向量的ingestion，storage和TopN 检索，并很好的利用了Opensearch的分布式特性，对于海量向量数据的检索可以实现并行计算，性能能随集群size扩缩而线性增长
多数情况下，我们只需要使用近似kNN。近似kNN提供低延迟是以较慢的索引速度和不完美的准确性为代价。
精确的暴力kNN保证精确结果但是不能很好地适应大型数据集。通过这种方式，script_score查询必须扫描每个匹配的文档节点来计算向量函数，这将导师搜索非常慢。然而我们可以通过限制传给向量函数的文档数量来降低延迟。如果我们能筛选出一个很小的文档子集，那么我们就可以通过这个方法获取较好的性能。
我们可以通过调大num_candidates的值来更精确的获取结果，代价是搜索的速度会变慢。如果使用一个比较大的num_candidates值会从每个分片获取更多的候选者。这就要花费更多的时间，也就有更大的可能找到真正的k个邻近结果。
类似的，我们可以降低num_candidates的值来获取更快的搜索，同时要接受潜在不太精确的结果。

参考资料

Opensearch 社区

KNN算法

Amazon Opensearch 托管服务

Opensearch KNN插件

亚马逊AWS官方博客