【推搜】embedding评估 | faiss的topK向量检索

note

• 在双塔模型召回中，正样本即用户点击过的物品，负样本：全体物品中负采样（简单做法）、被排序淘汰的物品（物品）。
• faiss库使用三部曲：构建向量库；构建index，并将向量添加到index中；进行topk的检索。根据具体的场景使用合适的index，有的索引还支持GPU构建。

零、背景介绍

在推荐系统的召回阶段中，很多时候item变化不会很明显，所以一般会将item embedding存入faiss之类的向量数据库（也可以使用HnswLib等，注意要建立对应的向量索引），然后在线计算更新user embedding（可由用户的历史item序列的item embedding的mean pooling操作生成，也可以用其他方法），最后将user embedding和item embedding进行最近邻的topk查找。即典型的【离线存储】和【线上召回】策略。最初的双塔模型就是这么做的。

一、离线指标

hitrate和ndcg等，其中hitrate的指标计算如下：

二、降维可视化

可以通过将embedding进行PCA或者TSNE降维可视化，看同一类别item的embedding分布。

三、通过RS线上指标体现

很多时候更实际的是，embedding用于下游任务，最后看线上指标（如CTR/CVR指标等）。

四、LSH原理及多桶策略

4.1 LSH介绍

Locality Sensitive Hashing——LSH
局部敏感哈希的基本思想是希望让相邻的点落入同一个“桶”，这样在进行最近邻搜索时，我们仅需要在一个桶内，或相邻几个桶内的元素中进行搜索即可。

时间复杂度：如果保持每个桶中的元素个数在一个常数附近，我们就可以把最近邻搜索的时间复杂度降低到常数级别。

先看实验得出的一个结论（上图就是将高维的彩色点映射到低维的abc中）：欧式空间中，将高维空间的点映射到低维空间，原本接近的点在低维空间中肯定依然接近，但原本远离的点则有一定概率变成接近的点。

（1）将高维embedding映射到低维：
由于 Embedding 大量使用内积操作计算相似度，因此我们也可以用内积操作来构建局部敏感哈希桶。假设 v 是高维空间中的 k 维 Embedding 向量，x 是随机生成的 k 维映射向量。那我们利用内积操作可以将 v 映射到一维空间，得到数值 h(v)=v⋅x。

（2）分桶：
一维空间也会部分保存高维空间的近似距离信息。因此，我们可以使用哈希函数 h(v) 进行分桶，公式为：$$h^{x,b}(v)=\lfloor \frac{x\cdot v+b}{w}]$$ 其中， ⌊ ⌋ 是向下取整操作， w 是分桶宽度，b 是 0 到 w 间的一个均匀分布随机变量，避免分桶边界固化。

因为如果总是固定边界，很容易让边界两边非常接近的点总是被分到两个桶里。这是我们不想看到的。所以随机调整b，生成多个hash函数，并且采用【或】的方式组合，就可以一定程度避免这些边界点的问题。

（3）用多个哈希函数同时分桶：
映射的过程会损失部分的距离信息，并且为了防止相近点误判，可以同时用m个哈希函数同时进行分桶操作——如果两个点都同时掉入同一个桶中，则他们是相似的点的概率就很大。就这样得到的候选集再通过遍历得到K邻近。

哈希策略是基于内积操作来制定的，内积相似度也是我们经常使用的相似度度量方法，事实上距离的定义有很多种，比如“曼哈顿距离”、“切比雪夫距离”、“汉明距离”等等。针对不同的距离定义，分桶函数的定义也有所不同，但局部敏感哈希通过分桶方式保留部分距离信息，大规模降低近邻点候选集的本质思想是通用的。

4.2 多桶策略

刚才3.1的使用多个哈希函数进行分桶操作：

（1）分桶栗子

背景：
假设有 A、B、C、D、E 五个点，有 h1和 h2两个分桶函数。
使用 h1来分桶时，A 和 B 掉到了一个桶里，C、D、E 掉到了一个桶里；
使用 h2来分桶时，A、C、D 掉到了一个桶里，B、E 在一个桶。
那么请问如果我们想找点 C 的最近邻点，应该怎么利用两个分桶结果来计算呢？

（1）用“且”（And）操作：
这样处理两个分桶结果之间的关系，则找到与点 C 在 h1函数下同一个桶的点，且在 h2函数下同一个桶的点，作为最近邻候选点。我们可以看到，满足条件的点只有一个，那就是点 D。也就是说，点 D 最有可能是点 C 的最近邻点。

作为多桶策略，可以最大程度地减少候选点数量。但是，由于哈希分桶函数不是一个绝对精确的操作，点 D 也只是最有可能的最近邻点，不是一定的最近邻点，因此，“且”操作其实也增大了漏掉最近邻点的概率。

（2）采用“或”（Or）操作：
找到与点 C 在 h1函数下同一个桶的点，或在 h2函数下同一个桶的点。这个时候，我们可以看到候选集中会有三个点，分别是 A、D、E。这样一来，虽然我们增大了候选集的规模，减少了漏掉最近邻点的可能性，但增大了后续计算的开销。

better：局部敏感哈希的多桶策略还可以更加复杂，比如使用 3 个分桶函数分桶，把同时落入两个桶的点作为最近邻候选点等等。

（2）取值建议：

• 点数越多，我们越应该增加每个分桶函数中桶的个数；相反，点数越少，我们越应该减少桶的个数；
• Embedding 向量的维度越大，我们越应该增加哈希函数的数量，尽量采用且的方式作为多桶策略；相反，Embedding 向量维度越小，我们越应该减少哈希函数的数量，多采用或的方式作为分桶策略。

（3）时间复杂度

局部敏感哈希能在常数时间得到最近邻的结果吗？
答案是可以的，如果我们能够精确地控制每个桶内的点的规模是 C，假设每个 Embedding 的维度是 N，那么找到最近邻点的时间开销将永远在 O(C⋅N) 量级。采用多桶策略之后，假设分桶函数数量是 K，那么时间开销也在 O(K⋅C⋅N) 量级，这仍然是一个常数。

4.3 局部敏感哈希实践

使用 Spark MLlib 完成 LSH 的实现。

在将电影 Embedding 数据转换成 dense Vector 的形式之后，我们使用 Spark MLlib 自带的 LSH 分桶模型 BucketedRandomProjectionLSH（我们简称 LSH 模型）来进行 LSH 分桶。其中最关键的部分是设定 LSH 模型中的 BucketLength 和 NumHashTables 这两个参数：
（1）BucketLength 指的就是分桶公式中的分桶宽度 w；
（2）NumHashTables 指的是多桶策略中的分桶次数。

和其他 Spark MLlib 模型一样，都是先调用 fit 函数训练模型，再调用 transform 函数完成分桶的过程。

def embeddingLSH(spark:SparkSession, movieEmbMap:Map[String, Array[Float]]): Unit ={
  //将电影embedding数据转换成dense Vector的形式，便于之后处理
  val movieEmbSeq = movieEmbMap.toSeq.map(item => (item._1, Vectors.dense(item._2.map(f => f.toDouble))))
  val movieEmbDF = spark.createDataFrame(movieEmbSeq).toDF("movieId", "emb")


  //利用Spark MLlib创建LSH分桶模型
  val bucketProjectionLSH = new BucketedRandomProjectionLSH()
    .setBucketLength(0.1)
    .setNumHashTables(3)
    .setInputCol("emb")
    .setOutputCol("bucketId")
  //训练LSH分桶模型
  val bucketModel = bucketProjectionLSH.fit(movieEmbDF)
  //进行分桶
  val embBucketResult = bucketModel.transform(movieEmbDF)
  
  //打印分桶结果
  println("movieId, emb, bucketId schema:")
  embBucketResult.printSchema()
  println("movieId, emb, bucketId data result:")
  embBucketResult.show(10, truncate = false)
  
  //尝试对一个示例Embedding查找最近邻
  println("Approximately searching for 5 nearest neighbors of the sample embedding:")
  val sampleEmb = Vectors.dense(0.795,0.583,1.120,0.850,0.174,-0.839,-0.0633,0.249,0.673,-0.237)
  bucketModel.approxNearestNeighbors(movieEmbDF, sampleEmb, 5).show(truncate = false)
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

使用 LSH 模型对电影 Embedding 进行分桶得到的五个结果打印：

+-------+-----------------------------+------------------+
|movieId|emb                          |bucketId          |
+-------+-----------------------------+------------------------+
|710    |[0.04211471602320671,..]     |[[-2.0], [14.0], [8.0]] |
|205    |[0.6645985841751099,...]     |[[-4.0], [3.0], [5.0]]  |
|45     |[0.4899883568286896,...]     |[[-6.0], [-1.0], [2.0]] |
|515    |[0.6064003705978394,...]     |[[-3.0], [-1.0], [2.0]] |
|574    |[0.5780771970748901,...]     |[[-5.0], [2.0], [0.0]]  |
+-------+-----------------------------+------------------------+
1
2
3
4
5
6
7
8
9

可以使用多桶策略：BucketId 这一列，因为我们之前设置了 NumHashTables 参数为 3（用了3个分桶/哈希函数），所以每一个 Embedding 对应了 3 个 BucketId。

4.4 LSH的一个实践栗子

场景：现在已知3个用户的特征向量，我们想要找到ABC三个用户之间的相似关系，当然最简单的就是用余弦cos计算，而如果用LSH局部敏感哈希如下pyspark版本的栗子：

import os
import re
import hashlib

from pyspark import SparkContext, SparkConf
from pyspark import Row
from pyspark.sql import SQLContext, SparkSession
from pyspark.sql.functions import *
from pyspark.sql.types import *
from pyspark.sql.functions import udf,collect_list, collect_set
from pyspark.ml.feature import MinHashLSH, BucketedRandomProjectionLSH
from pyspark.ml.linalg import Vectors, VectorUDT

# 控制spark服务启动
spark = SparkSession.builder.appName('app_name').getOrCreate()
spark.stop()
spark = SparkSession.builder.appName('app_name').getOrCreate()

class PySpark(object):
    # 装饰器
    @staticmethod
    def execute(df_input):
        """
        程序入口，需用户重载
        :return:必须返回一个DataFrame类型对象
        """
        # step 1：读入DataFrame
        df_mid = df_input.select('id','name','data','mat')
        print("--------读入的dataframe--------")
        df_mid.show()
        
        # step 2：特征向量预处理
        def mat2vec(mat):
            """
            定义UDF函数，将特征矩阵向量化
            :return:返回相似度计算所需的VectorUDT类型
            """
            arr = [0.0]*len(mat)
            for i in range(len(mat)):
                if mat[i]!='0':
                    arr[i]=1.0
            return Vectors.dense(arr)
        # udf自定义函数，第一个参数为函数，第二个参数为返回值类型
        udf_mat2vec = udf(mat2vec,VectorUDT())
        df_mid = df_mid.withColumn('vec', udf_mat2vec('mat')).select(
                                 'id','name','data','mat','vec')
        print("--------特征向量处理后的的dataframe--------")
        df_mid.show()
        
        # step 3：计算相似度
        ## MinHashLSH，可用EuclideanDistance
        minlsh = MinHashLSH(inputCol="vec", outputCol="hashes", seed=123, numHashTables=3)
        model_minlsh = minlsh.fit(df_mid)
        ## BucketedRandomProjectionLSH
        brplsh = BucketedRandomProjectionLSH(inputCol="vec", 
                                             outputCol="hashes",  
                                             seed=123, 
                                             bucketLength=10.0, 
                                             numHashTables=10)
        model_brplsh = brplsh.fit(df_mid)

        # step 4：计算（忽略自相似，最远距离限制0.8）
        ## model_brplsh类似，可用EuclideanDistance
        df_ret = model_minlsh.approxSimilarityJoin(df_mid, df_mid, 0.8, distCol='JaccardDistance').select(
                col("datasetA.id").alias("id"),
                col("datasetA.name").alias("name"),
                col("datasetA.data").alias("data"),
                col("datasetA.mat").alias("mat"),
                col("JaccardDistance").alias("distance"),
                col("datasetB.id").alias("ref_id"),
                col("datasetB.name").alias("ref_name"),
                col("datasetB.data").alias("ref_data"),
                col("datasetB.mat").alias("ref_mat")
            ).filter("id != ref_id")
        return df_ret

df_in = spark.createDataFrame([
    (1001,"A","xxx","1010001000010000001001101010000"),
    (1002,"B","yyy","1110001000010000000011101010000"),
    (1003,"C","zzz","1101100101010111011101110111101")],
    ['id', 'name', 'data', 'mat'])
df_out = PySpark.execute(df_in)
df_out.show()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

有以下的几个注意点：

• 一开始的三个向量是对用户的特征化描述（如one hot向量一样）
• 可以学习这种装饰器的写法（很常见）；pyspark中调用LSH有两种方法：MinHashLSH, BucketedRandomProjectionLSH两个选择，前者基于Jaccard距离，后者基于欧式距离。
• pyspark中的udf函数参数，第一个参数是传入的函数，第二个参数是该udf的函数返回值；
• 对应的结果如下，可以从第三个表看出，A和B之间的distance为0.27（表示不太相似），而A和B，分别与C之间的distance值都是大约为0.75（表示相似，通过看向量之间的数值也可以发现是符合的）。

--------读入的dataframe--------
+----+----+----+--------------------+
|  id|name|data|                 mat|
+----+----+----+--------------------+
|1001|   A| xxx|10100010000100000...|
|1002|   B| yyy|11100010000100000...|
|1003|   C| zzz|11011001010101110...|
+----+----+----+--------------------+

--------特征向量处理后的的dataframe--------
+----+----+----+--------------------+--------------------+
|  id|name|data|                 mat|                 vec|
+----+----+----+--------------------+--------------------+
|1001|   A| xxx|10100010000100000...|[1.0,0.0,1.0,0.0,...|
|1002|   B| yyy|11100010000100000...|[1.0,1.0,1.0,0.0,...|
|1003|   C| zzz|11011001010101110...|[1.0,1.0,0.0,1.0,...|
+----+----+----+--------------------+--------------------+

+----+----+----+--------------------+------------------+------+--------+--------+--------------------+
|  id|name|data|                 mat|          distance|ref_id|ref_name|ref_data|             ref_mat|
+----+----+----+--------------------+------------------+------+--------+--------+--------------------+
|1001|   A| xxx|10100010000100000...|0.2727272727272727|  1002|       B|     yyy|11100010000100000...|
|1003|   C| zzz|11011001010101110...|              0.75|  1001|       A|     xxx|10100010000100000...|
|1002|   B| yyy|11100010000100000...|              0.76|  1003|       C|     zzz|11011001010101110...|
|1002|   B| yyy|11100010000100000...|0.2727272727272727|  1001|       A|     xxx|10100010000100000...|
|1001|   A| xxx|10100010000100000...|              0.75|  1003|       C|     zzz|11011001010101110...|
|1003|   C| zzz|11011001010101110...|              0.76|  1002|       B|     yyy|11100010000100000...|
+----+----+----+--------------------+------------------+------+--------+--------+--------------------+
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

这里也可以和最常用余弦相似度方法对比一下：

import numpy as np
from numpy import random

# numpy 计算 cos 距离
def cosine_distance(a, b):
    if a.shape != b.shape:
        raise RuntimeError("array {} shape not match {}".format(a.shape, b.shape))
    if a.ndim==1:
        a_norm = np.linalg.norm(a)
        b_norm = np.linalg.norm(b)
    elif a.ndim==2:
        a_norm = np.linalg.norm(a, axis=1, keepdims=True)
        b_norm = np.linalg.norm(b, axis=1, keepdims=True)
    else:
        raise RuntimeError("array dimensions {} not right".format(a.ndim))
    # 计算cos即相似度
    similiarity = np.dot(a, b.T)/(a_norm * b_norm.T)
    return similiarity

# a = random.randint(size(2, 4))
a = np.array([1,2,3,4,5])
b = np.array([5,4,0,2,1])
ans = cosine_distance(a, b)
ans
# 0.5169078021169037
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

如果是用两个相同的向量最后的cos值自然是1了。

五、faiss向量检索使用

5.1 faiss的基本操作

• Faiss：Facebook 开源的向量相似检索库：https://github.com/facebookresearch/faiss，类似的有annoy向量检索库。从多媒体文档中快速搜索出相似的条目——这个场景下的挑战是基于查询的传统搜索引擎无法解决的。Faiss使用的数据流如上图所示。可以通过如下命令下载：

#cpu 版本
conda install faiss-cpu -c pytorch
# GPU 版本
conda install faiss-gpu cudatoolkit=8.0 -c pytorch # For CUDA8
conda install faiss-gpu cudatoolkit=9.0 -c pytorch # For CUDA9
conda install faiss-gpu cudatoolkit=10.0 -c pytorch # For CUDA10
1
2
3
4
5
6

linux上的安装：

# CPU安装
pip install faiss-cpu
# GPU安装
pip install faiss-gpu
1
2
3
4

按照三部曲进行一个简单的栗子（如下），这里使用最基础的索引类型indexFlatL2，即对embedding进行简单的L2距离搜索：

import numpy as np

# 1. 构建向量库
d = 64                                           # 向量维度
nb = 100000                                      # index向量库的数据量
nq = 10000                                       # 待检索query的数目
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')   # index向量库的向量
xq = np.random.random((nq, d)).astype('float32')   # 待检索的query向量

# 2. 构建index, 并将向量添加到index中
import faiss
index = faiss.IndexFlatL2(d)    # 创建索引时必须指定向量的维度d
print(index.is_trained)         # 输出为True，代表该类index不需要训练，只需要add向量进去即可
index.add(xb)                   # 将向量库中的向量加入到index中
print(index.ntotal)             # 输出index中包含的向量总数，为100000

# 3. 进行topk的检索
k = 4                     # topK的K值, 即对于每个需要检索的向量 找出top4个最相似的embedding
D, I = index.search(xq, k)# xq为待检索向量矩阵，返回的I为每个待检索query最相似TopK的索引list，D为其对应的距离
# 显示前5个待检索embedding的top4检索结果
print(I[:5])
print(D[:5])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

最后返回的D为距离，I是返回的每个待检索向量对应的topk列表，如I[: 5]就显示了前五个用户对应的top4个item检索向量：

上面的faiss.IndexFlatIP也可以换成其他的，注意faiss中最重要的是索引index，其他的API介绍见下表：

5.2 基于faiss的向量召回

• 序列推荐：将用户历史item_emb进行pooling后的user_emb（当然也可以不用mean pooling，用注意力机制就变成DIN模型了，可以参考图文解读：推荐算法架构——精排），和item_emb进行内积得到score偏好分数。交叉熵损失函数的参数即socre和标签item_id（该用户交互的item_id）对应的embedding，典型的多分类问题。
• 基于Faiss的向量召回：下面就是如一开始背景说的，一般在faiss向量数据库中离线存储训练好的item embedding，然后在线拼好用户user embedding或者根据用户历史item序列进行mean pooling之类的操作得到user embedding，最后两者进行faiss的topk检索。

# 第一步：我们获取所有Item的Embedding表征，然后将其插入Faiss(向量数据库)中
item_embs = model.output_items().cpu().detach().numpy()
item_embs = normalize(item_embs, norm='l2')
gpu_index = faiss.IndexFlatIP(hidden_size)
gpu_index.add(item_embs)

# 第二步：根据用户的行为序列生产User的向量表征
user_embs = model(item_seq,mask,None,train=False)['user_emb']
user_embs = user_embs.cpu().detach().numpy()

# 第三步：对User的向量表征在所有Item的向量中进行Top-K检索
# Inner Product近邻搜索，D为distance，I是index
D, I = gpu_index.search(user_embs, topN)  
1
2
3
4
5
6
7
8
9
10
11
12
13

5.3 faiss进阶操作

（1）使用GPU加速检索

可以使用gpu加速（注意安装gpu版的faiss），如创建索引时将IndexFlatL2替换为GpuIndexFlatL2。

注意：报错module ‘faiss‘ has no attribute ‘StandardGpuResources‘还没解决，留个坑，查了下说是faiss版本的问题，但是换了几个版本还是。。

（2）倒排索引（inverted index）

为什么需要倒排索引呢，首先看下面这样的表：

如果想搜索还有字符串winter的哪一行，使用sql语句Select * from Quotes where quote_text like '%winter%'语句需要遍历所有行 开销大，一种基础理想的倒排索引操作就是转为如下的结构，这样就能直接找出winter对应的quote_id。

（3）其他索引

附上忘记在哪里看到的一张不同类型的index索引相关总结图：

Reference

[1] Metrics for Evaluating Quality of Embeddings for Ontological
Concepts
[2] 向量检索库Faiss使用指北
[3] https://pytorch.org/docs/stable/generated/torch.nn.GRU.html?highlight=gru#torch.nn.GRU
[4] Faiss官方文档：https://github.com/facebookresearch/faiss
[5] GRU4Rec v2 - Recurrent Neural Networks with Top-k Gains for Session-based Recommendations
[6] Faiss从入门到实战精通
[7] 某乎：深度ctr预估中id到embedding目前工业界主流是端到端直接学习还是预训练
[8] https://github.com/facebookresearch/faiss
[9] 相似度检索Faiss模型
[10] Faiss(3)：基于IndexIVFPQ的demo程序
[11] IndexFlatL2、IndexIVFFlat、IndexIVFPQ三种索引方式示例
[12] facebook. Product quantization for nearest neighbor search（faiss的原论文）
[13] What is an inverted index?(介绍倒排索引)
[14] A brief explanation of the Inverted Index
[15] Indexing 2: inverted index.Victor Lavrenko
[16] 通俗讲解faiss（推荐）
[17] 如何加速faiss python search速度？
[18] 推荐系统(3):倒排索引在召回中的应用