亿级用户相似度计算：Faiss与Spark的实践探索

在系列前作中，我们已经探讨了如何利用ALS算法为海量用户及其兴趣标签构建向量表征（Embeddings）（详见《机器学习实践：ALS的矩阵分解》），并进一步研究了如何让这些向量模型具备动态适应线上数据更新的能力。

手握这些用户画像向量，一个新的工程挑战摆在了面前：如何在两亿级别的用户海洋中，为每一位用户精准且高效地“捞”出与其兴趣图谱最为契合的Top-K“知己”？这对于个性化推荐系统（例如“猜你认识”或关联内容推荐）而言，无疑是核心一环。

面对如此庞大的用户基数，传统的“两两握手”式比较方法，在计算资源和时间成本上显然力不从心。本文将聚焦于一种结合Faiss与Spark的分布式解决方案，旨在攻克这一大规模用户相似度计算的难题。

漫漫长路：那些年我们踩过的“坑”

在抵达“光明顶”之前，咱们也曾在崎岖小路上摸索过。

暴力破解？难以承受

最直观的想法，自然是把每个用户都和其他所有用户比一遍。假设我们有 Dataset A（待查询用户，2亿）和 Dataset B（目标用户池，1000万），每个用户一个24维向量。计算复杂度是 O(N * M * d)，其中 N 和 M 是用户数，d 是向量维度。这在实际工程中几乎是难以承受之重。

LSH 的尝试：理想很丰满，现实有点“慢”

局部敏感哈希（LSH）是个不错的思路，它能把相似的向量“拍”到同一个“桶”里，这样我们只需要在桶内进行比较。Spark MLlib 也提供了 BucketedRandomProjectionLSH。

理论上听起来很美，但在我们的2亿用户量级下，实践起来却发现：

• 调参地狱：bucketLength、numHashTables 这些参数非常敏感，需要针对反复实验，才能在召回率和效率间找到一个微妙的平衡点，当向量生成过程变化后，参数也需要重新调整。
• 时间都去哪儿了：即便调好了参数，一次完整的计算也动辄需要十几个小时。对于需要快速迭代和验证的推荐场景，这显然太慢了。

Faiss 单机版：快是真快，但“小马拉大车”

后来，我们把目光投向了 Faiss。这是 Facebook AI 开源的向量检索神器，以极致的速度著称。在小数据集上，Faiss 的表现堪称惊艳。

但问题是，它是为单机环境设计的。当我们的目标用户池（Dataset B）达到千万级别，要把所有向量加载到一台机器的内存里构建索引，已经开始捉襟见肘。更别提要为2亿用户（Dataset A）进行查询了，单台机器的内存和计算能力都将不堪重负。

Faiss是否主要用于小数据集

Faiss是否主要用于小数据集？

• 不完全是。Faiss本身设计得非常高效，对于单机能够容纳的数据集（例如，几千万甚至上亿的低维向量，取决于机器内存），它的性能是非常出色的。很多场景下，如果数据量在单机可控范围内，Faiss是首选。
• 博客中强调“小马拉大车”，更多是从“面临的2亿用户查询，以及千万级目标用户池构建索引”这个特定场景出发。对于这种规模，特别是当向量维度较高或需要存储原始数据时，单台标准服务器的内存和计算能力确实会成为瓶颈。

Faiss能否原生支持分布式？

• Faiss本身 并不直接提供像Spark那样的开箱即用的、与YARN等资源管理器深度集成的原生分布式计算框架。它更像一个高效的“引擎”或“库”。

• 但是，Faiss的设计考虑了可扩展性，它提供了 构建分布式索引的组件和策略。例如：

  • 索引分片（Index Sharding）：可以将一个大的索引拆分成多个小的分片，分布在不同的机器上。查询时，可以将查询向量发送到所有分片，然后合并结果。
  • IndexProxy：Faiss提供了一些工具类，可以帮助管理和查询分布式的索引分片。
  • 与MPI等结合：在高性能计算领域，Faiss可以与MPI (Message Passing Interface) 等并行编程模型结合，实现多机并行。

• 所以，说Faiss“不能原生支持分布式”可能不太准确，更准确的说法是：Faiss提供了构建分布式检索系统的基础能力和组件，但它本身不是一个完整的分布式计算框架。

为什么选择Spark + Faiss，而不是Faiss原生的分布式方式？

• 充分利用现有大数据生态：很多公司已经拥有成熟的Spark和YARN（或其他资源调度系统如Kubernetes）集群。基于Spark来封装和调度Faiss任务，可以：

  • 复用现有集群资源，避免额外搭建和维护一套专门的Faiss分布式集群。
  • 利用Spark强大的数据处理和ETL能力，方便地从Hive、HDFS等数据源加载和预处理向量数据。
  • 统一技术栈，降低团队学习和运维成本。

• 部署和管理的便捷性：自己从零开始搭建和管理一个稳定高效的Faiss分布式集群，确实比利用Spark的成熟生态要复杂得多。Spark提供了任务调度、容错、资源管理等一系列成熟机制。

• 灵活性和易用性：通过Spark的RDD/DataFrame API，可以非常灵活地组织Faiss索引的构建和查询流程

解决方案：Faiss与Spark的分布式协同

既然一台机器搞不定，那就“众人拾柴火焰高”！我们决定让 Spark 这位分布式计算大师，带着 Faiss 这位向量检索专家一起干活。这就是我们最终的解决方案（详见 similarity_search.ipynb）。

核心思想很简单：分而治之，各个击破，最后汇总。

1. 分布式构建索引 (针对 Dataset B - 我们的“目标用户池”)
   • Driver 领航：首先，Spark Driver 节点从 Dataset B 中抽取一部分样本数据，用它们训练出一个“全局的聚类中心”（Coarse Quantizer）。这就像是给我们的数据世界画一个大致的地图，标出主要的区域。
   • 蓝图共享：然后，这份“地图”被广播给所有的 Spark Worker 节点。
   • Worker 分区建索引：Dataset B 被切分成很多小块，每个 Worker 节点负责处理一小块数据。它们会根据收到的“地图”，把自己负责区域内的用户向量组织起来，建立一个局部的 Faiss 索引。
   • 合并成最终索引：最后，Driver 节点收集所有 Worker 建好的局部索引，将它们巧妙地合并成一个巨大且完整的全局 Faiss 索引。这个索引就代表了我们整个 Dataset B 用户群体的向量分布。
2. 分布式搜索 (针对 Dataset A - 我们的“待查询用户”)
   • 索引广播：将上一步得到的全局 Faiss 索引（Dataset B 的索引）再次广播给所有 Worker 节点。
   • 查询任务分发：Dataset A 同样被切分成小块，分发给各个 Worker。
   • Worker 并行查找：每个 Worker 节点拿到自己负责的 Dataset A 用户向量后，就在本地的全局索引副本中进行高速查找，找出每个用户的 Top-K 相似邻居。
   • 结果汇总：最后，收集所有 Worker 的查找结果，就是我们想要的最终相似用户列表啦！

图：Spark与Faiss协同工作流程示意

Faiss IndexIVFFlat 工作机制解析：“图书馆奇遇记”

我们选用的Faiss核心索引类型是IndexIVFFlat。其工作原理，不妨把它想象成一个超大型图书馆的高效图书检索系统：

1. 图书分类 (KMeans 聚类 - IVF 部分)
   • 想象图书馆有几千万本书（用户向量），直接一本本翻肯定不行。聪明的图书管理员会先把书按主题分成几千个大区（比如 nlist=1024 个区），例如“历史区”、“科幻区”、“烹饪区”等。这个分类过程，Faiss 用的是 K-Means 聚类算法，把相似的向量聚到同一个“区”（Cell）。每个区都有一个“代表”（Centroid），可以理解为这个区的“核心主题”或“平均品味”。
2. 缩小范围 (查 nprobe 个区)
   • 现在，你带着一本书（查询用户向量）想找类似的书。你不会跑遍所有几千个区，而是先看看目录，或者问问管理员，哪些区最有可能包含你想要的书。Faiss 也是这样，它会先比较你的查询向量和所有区的“代表”，找出最可能相关的 nprobe 个区（比如 nprobe=50，代表你决定先逛这50个最相关的区）。
3. 区内精细查找 (Flat Search)
   • 锁定了这 nprobe 个重点区域后，事情就简单了。Faiss 会在这几个区内，把每一本书都拿出来和你的书仔细比对一遍，找出最相似的那些。这个过程虽然是“暴力”比对（Flat Search），但因为范围已经大大缩小，所以依然非常高效。
4. 用什么标准判断相似？(METRIC_INNER_PRODUCT)
   • 我们用“内积”（Inner Product）来衡量两本书（两个用户向量）有多相似。对于已经归一化处理的向量（就像我们从ALS得到的用户向量），内积越大，说明它们在方向上越一致，代表用户兴趣越相近。

通过这种“先粗筛，再精选”的两步策略，IndexIVFFlat 就能在海量数据中快速定位到相似的向量。

进一步的优化方向：量化存储

本文主要采用的IndexIVFFlat策略，其优势在于精确存储原始向量，保证了检索的准确性。然而，这也意味着当目标用户池（Dataset B）的规模进一步扩大时，每个Worker节点仍需承载部分原始向量数据，可能会带来一定的内存压力。

Faiss同样提供了多种内存占用更优的索引策略，例如

• 乘积量化（Product Quantization, PQ） 的IndexIVFPQ
• 标量量化（Scalar Quantization, SQ） 的IndexIVFScalarQuantizer

这些策略通过对向量进行有损压缩，能够显著降低索引的内存占用，代价通常是检索精度上略有折损。选择这类索引时，需要在内存、速度和精度之间进行权衡。其分布式构建和查询流程与IndexIVFFlat在大体框架上相似，但在量化器的训练、向量的编码解码等具体环节会有所不同，是未来进一步优化系统资源占用的一个重要方向。

性能评估：效率与效果的提升

那么，这套组合拳打下来，效果如何呢？

• 速度起飞！ 原先用 Spark LSH 需要跑十几个小时的任务，换上 Faiss + Spark 的组合后，根据我们的实验（详见 similarity_search.ipynb），最快可以在半小时到一小时内完成！这效率提升了不止一个数量级。
• 轻松应对海量数据！ 通过分布式处理，无论是千万级的索引构建，还是亿级的用户查询，都不再是瓶颈。
• 调参free！ 索引是根据数据自动训练的，不需要人工调参数，当我们切换到不同的向量进行检索时，无需额外操作，减少了维护成本。

这套方案不仅解决了我们最初遇到的性能问题，也为后续更复杂的推荐策略打下了坚实的基础。

承前启后：从 ALS “炼丹”到 Faiss “寻宝”

看到这里，有些朋友可能会问，这和你之前分享的 ALS 矩阵分解有啥关系呢？

关系可大了！咱们在这里用来计算相似度的用户向量 (Embeddings)，完全可以是由之前文章中讨论的 ALS 算法“炼制”出来的。

• 第一篇 ALS 博客 主要讲了如何基于用户行为数据，通过矩阵分解得到用户和物品的向量表示。
• 第二篇 ALS 博客 探讨了如何让 ALS 模型能够适应线上数据动态更新的需求，实现类似 Inductive Learning 的效果。

而今天这篇文章，则是聚焦于如何高效地 应用 这些得来不易的用户向量，在亿万用户中为他们找到“知己”。可以说，这是我们推荐技术探索之旅的又一个重要里程碑。

版权声明

以上文章为本人@迪吉老农原创，文责自负。文中如有引用他人内容的部分（包括文字或图片），均已明文指出，或做出明确的引用标记。如需转载，请联系作者，并取得作者的明示同意。感谢。

Comments