知乎搜索文本相关性与知识蒸馏

导读：大家好，我是申站，知乎搜索团队的算法工程师。今天给大家分享下知乎搜索中文本相关性和知识蒸馏的工作实践，主要内容包括：

• 知乎搜索文本相关性的演进
• BERT 在知乎搜索的应用和问题
• 知识蒸馏及常见方案
• 知乎搜索在 BERT 蒸馏上的实践

知乎搜索文本相关性的演进

1. 文本相关性的演进

我们首先来介绍下知乎搜索中的文本相关性。在搜索场景中，文本相关性可以定义为⽤户搜索 query 的意图与召回 doc 内容的相关程度。我们需要通过不同模型来对这种相关程度进行建模。整体而言，文本的相关性一般可以分为两个维度，字面匹配和语义相关。知乎搜索中文本相关性模型的演进也是从这两个方面出发并有所侧重和发展。在知乎搜索的整个架构中，文本相关性模型主要定位于为二轮精排模型提供更高维/抽象的特征，同时也兼顾了一部分召回相关的工作。

2. Before NN

知乎搜索中的文本相关性整体演进可以分为三个阶段。在引入深度语义匹配模型前，知乎搜索的文本相关性主要是基于 TF-IDF/BM25 的词袋模型，下图右边是 BM25 的公式。词袋模型通常来说是一个系统的工程，除了需要人工设计公式外，在统计词的权重、词频的基础上，还需要覆盖率、扩展同义词，紧密度等各种模块的协同配合，才能达到一个较好的效果。知乎搜索相关性的一个比较早期的版本就是在这个基础上迭代的。右下部分为在基于词袋模型的基础上，可以参考使用的一些具体特征。

3. Before BERT

基于 BM25 的词袋模型不管如何设计，主要还是只解决文本相关性中的字面匹配这部分问题。第二阶段引入的深度语义匹配模型则聚焦于解决语义相关的问题，主要分为两部分：双塔表示模型和底层交互模型。微软的 DSSM（左下）是双塔模型的典型代表。双塔模型通过两个不同的 encoder 来分别获取 query 和 doc 的低维语义句向量表示，然后针对两个语义向量来设计相关性函数（比如 cosine）。DSSM 摆脱了词袋模型复杂的特征工程和子模块设计，但也存在固有的缺陷：query 和 doc 的语义表示是通过两个完全独立的 encoder 来获取的，两个固定的向量无法动态的拟合 doc 在不同 query 的不同表示。这个反应到最后的精度上，肯定会有部分的损失。

底层交互模型一定程度上解决了这个问题。这个交互主要体现在 query 和 doc term/char 交互矩阵（中）的设计上，交互矩阵使模型能够在靠近输入层就能获取 query 和 doc 的相关信息。在这个基础上，后续通过不同的神经网络设计来实现特征提取得到 query-doc pair 的整体表示，最后通过全连接层来计算最终相关性得分。Match-Pyramid（右下）、KNRM（右上）是交互模型中比较有代表性的设计，我们在这两个模型的基础上做了一些探索和改进，相比于传统的 BM25 词袋模型取得了很大的提升。

4. BERT

BERT 模型得益于 transformer 结构拥有非常强大的文本表示能力。第三阶段我们引入了 BERT 希望能够进一笔提高知乎搜索中文本相关性的表型。BERT 的应用也分为表示模型和交互模型。

对于交互模型来说，如下左图，query 和 doc 分别为 sentence1 和 sentence2 直接输入到 BERT 模型中，通过 BERT 做一个整体的 encoder 去得到 sentence pair 的向量表示，再通过全连接层得到相似性打分，因为每个 doc 都是依赖 query 的，每个 query-doc pair 都需要线上实时计算，对 GPU 机器资源的消耗非常大，对整体的排序服务性能有比较大的影响。

基于上述原因，我们也做了类似于 DSSM 形式的表示模型，将 BERT 作为 encoder，训练数据的中的每个 query 和 doc 在输入层没有区分，都是做为不同的句子输入，得到每个句向量表示，之后再对两个表示向量做点乘，得到得到相关度打分。通过大量的实验，我们最终采用了 BERT 输出 token 序列向量的 average 作为句向量的表示。从交互模型到表示模型的妥协本质是空间换时间，因为 doc 是可以全量离线计算存储的，在线只需要实时计算比较短的 query ，然后 doc 直接通过查表，节省了大量的线上计算。相比于交互模型，精度有一部分损失。

BERT 在知乎搜索的应用和问题

1. 搜索业务架构中的 BERT

在下图中我们可以看到，BERT 在知乎搜索业务的召回和排序阶段都扮演了比较重要的角色。交互模型的主要服务于二轮精排模型，依赖于线上实时的计算 query 和 doc，为精排模块提供相关性特征。表示模型又分为在线和离线两块，在线表示模型实时的为用户输入的 query 提供句向量表示，离线表示模型为库中的 doc 进行批量句向量计算。一方面，doc 向量通过 TableStore/TiDB 和 Redis 的两级存储设计，为线上排序做查询服务；另一方面，使用 faiss 对批量 doc 向量构建语义索引，在传统的 term 召回基础上补充向量语义召回。

2. BERT 表示模型语义召回

下面详细介绍下我们的语义召回模型。首先看个例子，对于「玛莎拉蒂 ghlib」这个 case，用户真正想搜的是「玛莎拉蒂 Ghibli」这款车，但用户一般很难记住完整的名称，可能会输错。在输错的情况下，基于传统的 term 匹配方式（Google 搜索的例子）只能召回“玛莎拉蒂”相关的 doc，而无法进行这辆车型的召回，这种场景下就需要进行语义召回。更通用的来说，语义召回可以理解为增加了字面不匹配但是语义相关的 doc 的召回。

语义召回模型整体是 BERT 相关性任务中双塔表示模型的一个应用。BERT 做为 encoder 来对 query 和 doc 进行向量的表示，基于 faiss 对全量 doc 向量构建语义索引，线上实时的用 query 向量进行召回。这个策略上线后，线上 top20 doc 中语义召回 doc 数量占总召回 doc 数量的比例能到达 5%+。

3. BERT 带来的问题

BEER 模型上线后，为不同的模块都取得了不错收益的同时，也给整个系统带来了不少问题。这些问题整体可以归结为线上实时计算、离线存储、模型迭代三个方面。具体的见上图。

4. 蒸馏前的尝试

针对上述性能或存储的问题，在对 BERT 蒸馏之前，我们也进行了很多不同的尝试。

BERT 交互模型的部署放弃了使用原生 TF serving，而是在 cuda 的基础上用 c++ 重写了模型的加载和 serving，加上混合精度的使用。在我们的业务规模上，线上实时性能提高到原来的约 1.5 倍，使 BERT 交互模型满足了的最低的可上线要求。在这个基础上，对线上的 BERT 表示模型增加 cache，减少约 60% 的请求，有效减少了 GPU 机器资源的消耗。

另一个思路是尝试给 BERT 在横向和纵向维度上瘦身。横向上，一方面可以减小 serving 时 max_seq_length 长度，减少计算量；另一方面可以对表示向量进行维度压缩来降低存储开销。这两种尝试在离线和在线指标上都有不同程度的损失，因此被放弃。纵向上，主要是减少模型的深度，即减少 transformer 层数。这对于显存和计算量都能得到显著的优化。前期尝试过直接训练小模型，以及使用 BERT-base 若干层在下游的相关性任务上进行 fine-tune。这两种方案，在离线指标上的表现就没法达到要求，因此也没有上线。

针对 doc 数量过大，存储开销过大和语义索引构建慢的问题。在这方面做了一个妥协的方案：通过 wilson score 等规则过滤掉大部分低质量的 doc，只对约 1/3 的 doc 存储表示向量和构建语义索引。该方案会导致部分文档的相关性特征存在缺失。对于表示模型存在的低交互问题，尝试 Poly-encoder（Facebook 方案）将固定的 768 维表示向量转为多个 head 的形式，用多个 head 做 attention 的计算，保证性能在部分下降的前提得到部分精度的提升。

智知识蒸馏及常见方案

1. 知识蒸馏

下面简单介绍下知识蒸馏。从下图中看，我们可以把知识蒸馏的整体形式简化为：大模型不考虑性能问题尽量学习更多的知识（数据），小模型通过适量的数据去高效地学习大模型的输出，达到一个知识迁移的效果。实际 serving 使用的是小模型。

知识蒸馏为什么能有效？关键点在于 soft target 和 temperature。soft target 对应的是 teacher 模型的输出，类似于概率分布，知识蒸馏从 hard target 转为 soft target 的学习有利于模型更好的去拟合标签，引入 temperature 则是为了进一步平滑标签，让模型去学习到类别和类别中的知识。这里需要注意的是，temperature 的选取不宜过大，太大的 temperature 会导致不同类别之间的差异被完全平滑掉。

2. BERT 蒸馏方案

对与 BERT 的蒸馏我们做了大量的调研，并对目前主流的蒸馏方案做了归纳分类。基于任务维度来说，主要对应于现在的 pretrain + fine-tune 的两段式训练。在预训练阶段和下游任务阶段都有不少的方案涉及。技巧层面来分的话，主要包括不同的迁移知识和模型结构的设计两方面。后面我会选两个典型的模型简单介绍一下。

3. 蒸馏-MiniLM

MiniLM 是基于预训练任务的蒸馏，其是一种通用的面向 Transformer-based 预训练模型压缩算法。主要改进点有三个，一是蒸馏 teacher 模型最后一层 Transformer 的自注意力模块，二是在自注意模块中引入 values-values 点乘矩阵的知识迁移，三是使⽤了 assistant ⽹络来辅助蒸馏。

4. 蒸馏-BERT to Simple NN

BERT to Simple NN 更多的是做了一些 loss 形式的设计，使其训练方式更高效。

知乎搜索再 BERT 蒸馏上的实践

1. BERT 蒸馏上的实践和收益

前面的介绍中我有提到，在做 BERT 蒸馏前其实已经做了很多尝试，但是多少都会有精度的损失。因此，我们做蒸馏的第一目标是离线模型对⽐线上 BERT 精度⽆损。但对 BERT-base 直接进行蒸馏，无论如何都没办法避免精度的损失，所以我们尝试用更大的模型（比如 BERT-large/Robert-large/XLNET）来作为 teacher 进行蒸馏。这些多层的模型均在我们知乎全量语料先做 pretrain，再做 fine-tune，得到微调后的模型再做蒸馏。

2. 蒸馏-Patient KD

我们对交互模型和表示模型都做了蒸馏，主要采用了 Patient KD 模型的结构设计，Student 模型基于 BERT-base 的若干层运用不同的策略进行参数的初始化，去学习 Robert-large 大模型的方案。

其中知识迁移主要有三部分：student 的预测与真实标签的交叉熵、student 与 teacher 的预测的交叉熵和中间隐层的向量之间的 normalized MSE。

3. BERT 交互模型蒸馏

对于我们选的 teacher 模型 Robert-large，单纯预训练模型其 nDCG 指标为 0.914，线上之前使用的 BERT-base 是 0.907，若对 BERT-base 的若干 6 层直接去做 fine-tune 能达到的最高指标是 0.903，对比于 BERT-base 精度会损失很多。

我们这块做了一些尝试，基于 Robert-large 从 24 层蒸馏到 6 层的话能到 0.911，能超过线上 BERT-base 的效果。

训练数据方面，我们经历了点击日志数据挖掘到逐渐建立起完善的标注数据集。目前，相关性任务训练和蒸馏主要均基于标注数据集。标注数据分为 title 和 content 两部分，Query 数量达到 10w+ 的规模，标注 doc 在 300w ~ 400w 之间。

4. BERT 表示模型蒸馏

在 BERT 表示模型上，蒸馏时我们希望对向量维度和模型层数同时进行压缩，但蒸馏后得到的 student 模型表现不及预期。所以最后上线的方案中，表示模型层数还是维持了 12 层。在蒸馏时，为了提高精度，选取交互模型作为 teacher 进行蒸馏。因为交互模型是 query 和 doc 之间的打分，交互模型得到的 logits 与表示模型点乘后的打分在数量值会有较大差值，所以用 pairwise 形式通过 teacher 差值拟合来进行 loss 的计算。

在维度压缩方面我们做了对比实验，BERT 模型输出做 average pooling 后接全连接层分别压缩至 8 维到 768 维。如图所示，128 维和 64 维的表现跟 768 维差别不大，在上线时选择维度为 64 和 128 进行尝试，两者在线上表现没有太明显的差异，最终选择了 64 维的方案，把模型的维度压缩了 12 倍，存储消耗更低。

 5. 蒸馏的收益

蒸馏的收益主要分为在线和离线两部分。

在线方面：

交互模型的层数从 12 层压缩到 6 层，排序相关性特征 P95 减少为原本的 1/2，整体搜索入口下降 40ms，模型部署所需的 GPU 机器数也减少了一半，降低了资源消耗。

表示模型语义索引存储规模 title 减为 1/4，content 维度从 768 维压缩至 64 维，虽然维度减少了 12 倍，但增加了倒排索引 doc 的数量，所以 content 最终减为 1/6，语义索引召回也有比较大的提升，title 减少为 1/3，content 减少为 1/2。精排模块需要线上实时查询离线计算好的向量，所以查询服务也有提升。

离线方面：

表示模型语义索引的构建时间减少为 1/4，底层知乎自研的 TableStore/TIDB 存储减为原来的 1/6，LTR 训练数据和训练时间都有很大的提升，粗排早期用的是 BM25 等基础特征，后来引入了 32 维的 BERT 向量，提升了精排精度。

嘉宾介绍：

申站 硕士，知乎社区业务中心搜索算法工程师。目前主要负责知乎搜索系统中文本相关性探索和落地优化，在 query 理解和召回上也有长期的经验。

本文转载自：DataFunTalk（ID：datafuntalk）

原文链接：知乎搜索文本相关性与知识蒸馏