高性能高维向量的KNN搜索方案

hogwarts · 2022 年5 月 18 日 08:39

背景

机器学习以及深度学习中经常会使用向量表示原始数据（例如图片、视频、自然语言中的embedding等），而且通过查询词（例如图片搜索、视频搜索、词向量空间近似计算等）计算KNN是强需求。
在此业务需求背景下产生了高性能高维向量KNN搜索方案。

KNN

KNN算法（K近邻法）是简单常用的监督学习方法，模型简单无需训练，但是应用范围广而且变形多，近些年深度学习广泛应用于推荐系统，但是对比由KNN等基线算法，效果受到很大挑战

KNN算法通过多数投票的方式进行预测，可以应用于分类与回归，利用训练数据将特征向量空间进行划分。

由三个基本要素组成：

1.K值的选择（该超参数直接影响效果）

2.距离度量的方法（例如闵可夫斯基距离、马氏距离、相关系数、余弦夹角等，不同的距离衡量方法，导致结果不同）

3.决策规则（回归经常采用均值，分类是多数投票）

结构化数据可以直接应用KNN方法，非结构化数据需要先转成抽象含义的向量，然后应用KNN方法。

为了实现快速的KNN计算，暴力搜索的复杂度是N，基本不会被采用。常用的是KDTree、BallTree等树状结构算法，将复杂度降低到对数量级。

高性能搜索方案

除了KDTree等算法之外，近些年工业界产生了大量的高性能方案，这些方案的主要特点是毫秒级延迟、可处理亿级数据、利用GPU分布式计算、可接受精度损失。

其中速度最快属于nmslib，faiss应用最广，因此后续介绍下faiss如何使用

FAISS

FAISS（Facebook AI Similarity Search）使用C++实现，封装了python调用接口，支持CPU和GPU的异构计算方案，GPU上的运行速度相比CPU大致提升了一个量级。

如果使用GPU，需要安装cuda（支持最新版10.1）

向量是32维，百万数据集上进行K=4搜索：

用时29.438秒，单个数据耗时都在微秒量级。

参考文献

阿里妈妈推荐系统开发的深度树匹配模型可以参考下：