[论文解读] TwinBERT: Distilling Knowledge to Twin-Structured BERT Models for Efficient Retrieval

Wenhao Lu, Jian Jiao+1 more|arXiv (Cornell University)|Feb 14, 2020

Topic Modeling参考文献 37被引用 28

一句话总结

TwinBERT 提出了一种双结构 BERT 模型，将查询和文档编码解耦，使文档嵌入可以预先计算并缓存，从而将 CPU 上的推理时间减少至约 20ms。通过知识蒸馏和高效的网络设计，TwinBERT 在推理速度上比 BERT-Base 和 BERT-Large 快 77–663 倍，达到 BERT-Base 水平的性能，实现了生产环境中低延迟检索系统的高效部署。

ABSTRACT

Pre-trained language models like BERT have achieved great success in a wide variety of NLP tasks, while the superior performance comes with high demand in computational resources, which hinders the application in low-latency IR systems. We present TwinBERT model for effective and efficient retrieval, which has twin-structured BERT-like encoders to represent query and document respectively and a crossing layer to combine the embeddings and produce a similarity score. Different from BERT, where the two input sentences are concatenated and encoded together, TwinBERT decouples them during encoding and produces the embeddings for query and document independently, which allows document embeddings to be pre-computed offline and cached in memory. Thereupon, the computation left for run-time is from the query encoding and query-document crossing only. This single change can save large amount of computation time and resources, and therefore significantly improve serving efficiency. Moreover, a few well-designed network layers and training strategies are proposed to further reduce computational cost while at the same time keep the performance as remarkable as BERT model. Lastly, we develop two versions of TwinBERT for retrieval and relevance tasks correspondingly, and both of them achieve close or on-par performance to BERT-Base model. The model was trained following the teacher-student framework and evaluated with data from one of the major search engines. Experimental results showed that the inference time was significantly reduced and was firstly controlled around 20ms on CPUs while at the same time the performance gain from fine-tuned BERT-Base model was mostly retained. Integration of the models into production systems also demonstrated remarkable improvements on relevance metrics with negligible influence on latency.

研究动机与目标

• 为解决 BERT 在实时信息检索（IR）系统中推理延迟过高的问题。
• 实现在低延迟环境（如赞助搜索）中高效在线服务深度神经网络模型。
• 在大幅降低计算成本的同时，保持高检索和相关性性能。
• 探索能够保留性能但提升推理效率的知识蒸馏技术。
• 实现在不牺牲相关性质量的前提下，将密集语义模型部署在 CPU 上。

提出的方法

• TwinBERT 使用两个独立的类似 BERT 的编码器，分别独立处理查询和文档，将输入编码过程与 BERT 的标准拼接方式解耦。
• 文档嵌入在离线阶段预先计算并缓存在内存中，从而在推理阶段消除文档的运行时编码开销。
• 一个交叉层通过余弦相似度或残差网络结合查询和文档嵌入，计算相关性得分。
• 采用知识蒸馏技术，以 BERT-Base 模型作为教师模型训练 TwinBERT，从而在降低模型复杂度的同时保留性能。
• 通过 ONNX Runtime 对模型进行 CPU 推理优化，最大限度减少生产系统中的服务开销。
• 设计了高效的网络组件和训练策略，在不显著降低性能的前提下减少计算成本。

实验结果

研究问题

• RQ1 在 BERT 中解耦查询和文档编码，是否能在保持高检索性能的同时降低推理延迟？
• RQ2 知识蒸馏在多大程度上能将 BERT 级性能保留在更小、更快的模型中，用于 IR 任务？
• RQ3 预先计算并缓存文档嵌入在多大程度上能减少检索系统中的运行时计算？
• RQ4 TwinBERT 是否能在 CPU 上实现低于 20ms 的推理延迟，同时在相关性排序上匹配 BERT-Base 的性能？
• RQ5 TwinBERT 对生产搜索系统在延迟、准确率和部署可行性方面有何影响？

主要发现

• 在每查询评分 100 个文档时，TwinBERT 在 CPU 上的平均推理时间约为 20ms，相比 BERT 显著降低了延迟。
• 通过预计算文档嵌入，TwinBERT 在推理速度上分别比 BERT-Base 和 BERT-Large 快 77 倍和 422 倍。
• 在生产 A/B 测试中，TwinBERT 保留了微调后的 BERT-12 模型 90% 以上的增量性能增益，同时将不良广告展示次数减少了 10% 以上。
• TwinBERT 模型已成功部署于一个主要的赞助搜索系统中，表现出可忽略的延迟影响和高质量的相关性。
• TwinBERT 的余弦相似度版本在相同条件下，相比 BERT-3 和 BERT-12 分别实现了 121 倍和 663 倍的加速。
• 即使在运行时重新计算查询嵌入，TwinBERT 仍比 BERT-3 更快，凸显了其在效率上的优势。