作者：卓立、日涉、谨持

一、背景

大规模信息检索一直是搜推广领域的核心问题之一，而基于任意复杂模型的检索方案无疑是业界重要的迭代方向之一。近年来，阿里妈妈展示广告Match团队与预测引擎团队专注于从算法与工程角度推动工业级大规模检索技术的研发，我们在基于任意复杂模型的检索方向上积累了一定经验并取得了不错的业务效果，现整理发布NANN（Neural Approximate Nearest Neighbor，以下简称NANN）并对外开源，希望通过社区的协同创造力，共同推进该领域的发展。

本文介绍的NANN源自阿里妈妈展示广告Match团队研发的二向箔算法体系，该方案在保留复杂模型召回能力的同时，将索引学习和模型训练解耦，提供了轻量化的任意复杂模型召回解决方案。NANN基于Tensorflow，提供了性能benchmarking工具以及完整的由模型训练至在线deployment的demo。

该方案由阿里妈妈技术团队自研，已在阿里巴巴集团内部其他业务进行推广上线，在典型的搜索、推荐、广告场景均取得了显著的业务收益。

本文将围绕NANN核心功能及算法体系更新进行简要介绍，欢迎大家试用和交流讨论。

二、核心功能介绍

NANN代码库基于原生TensorFlow框架实现了高性能的二向箔算法在线检索服务。同时对于NANN算法的训练、索引结构构建、离线效果及在线性能评测、部署等流程提供了详细的示范。

该开源代码库的主要优点可以分为模型训练、性能优化和用户友好性三方面：

2.1 模型训练

任意复杂模型：模型训练与索引构建解耦，因此对模型结构几乎没有任何限制，也避免了索引和模型训练绑定带来的高额训练负担。
对抗训练：我们采用对抗训练来保证复杂模型下的优越检索性能。

此外，我们还提供了预训练模型和格式转化后的测试数据集，以方便大家试用及验证。

2.2 性能优化

高效检索：我们使用TensorFlow Custom Ops实现了HNSW检索过程。就在线检索而言，重写后的HNSW检索比Faiss原生版本更加高效。
运行时优化：我们支持GPU Multi-Streaming with Multi-Contexts，这极大地增强了并行性。
编译优化：我们支持Dynamic Shape下的XLA加速能力，并利用缓存等手段提升了just-in-time（JIT ）编译效率，最终使得NANN整个检索过程可以充分利用XLA加速。
图级优化：我们针对推荐、搜索和广告领域中常用的一些常见的模型结构，基于TensorFlow Grappler提供了一些图级优化。

2.3 用户友好性

原生TensorFlow：NANN的后端服务和前端实现完全基于TensorFlow生态系统。
模型推理和检索解耦：训练-检索解耦使得用户能专注优化深度模型，无需额外考虑检索流程。
性能测试：我们提供了一个简单的基准测试工具，可用于分析延迟、吞吐量等推理性能。

环境搭建、模型训练、部署等更多细节详见开源库。

https://github.com/alibaba/nann

三、算法体系更新

我们去年在《TDM到二向箔：阿里妈妈展示广告Match底层技术架构演进》这篇文章里介绍了初版的二向箔算法体系。在文章末尾的展望部分，我们也提到图索引构建中采用的 L2 距离度量与在线检索采用的模型打分之间存在不一致，这会导致检索精度会随着模型结构变复杂而大幅降低。为此，我们在训练过程中加入了一个对抗训练辅助任务，大幅减轻了相似性度量之间异质性的负面影响，有效提高了检索精度。

具体来说，在内部数据集（Industry）和开源数据集（User Behavior）两个数据集上，随着User侧和Target侧的交互越来越多，虽然检索recall （recall-retrieval）提升明显，但同时与全量打分recall（recall-all）之间的差距（）也越来越大。

我们认为该现象主要是因为模型打分的非线性随模型复杂度的增加而大幅提升。在双塔模型（two-sided）下，模型打分函数与target embedding 为线性关系，因此L2距离近的target embedding对应的模型打分也相近，最终表现为检索损失很低，在Industry和User Behavior两个数据集上分别为0.35%和0.74%；而在结构最复杂的attention模型（DNN w/attention）下，一旦给target embedding 加上微小的扰动，模型打分可能会发生剧烈的变化，因此检索精度损失大幅升高。为此，我们在模型训练中加入了如下辅助函数来控制模型打分的平滑性：

其中为用户，为该用户的正、负样本集，表示用户和target 的模型打分，为模型打分在target embedding 处对抗梯度方向的微小扰动。

加入对抗训练可以在不损失全量打分recall的情况下大幅降低复杂模型的检索损失，从而大幅提升检索性能。

四、总结展望

从模型角度看，加入对抗训练的二向箔算法体系真正做到了任意复杂模型下的高精度、高性能大规模KNN检索。我们将该方案整理开源，为社区提供了一套轻量级的支持任意复杂模型的大规模召回解决方案，该项工作相关工作论文已发表在CIKM2022。期待大家试用和体验，如遇问题，也欢迎与我们讨论。

五、关于我们

阿里妈妈展示广告Match团队多年来致力于大规模检索技术。经过团队同学多年的努力与沉淀，发展并迭代了如TDM、JTM、BSAT以及NANN等多项核心检索技术。近年来，阿里妈妈展示广告Match与阿里妈妈预测引擎团队深度合作，逐步从算法单一视角迭代演变为算法与工程双视角迭代。在阿里妈妈预测引擎团队异构计算技术能力的加持下，进一步释放模型迭代的天花板。同时，从业务及算法视角反哺阿里妈妈预测引擎团队，进一步扩展了妈妈预测引擎团队技术能力的通用性。

代码开源！阿里妈妈展示广告Match底层技术架构最新进展