[论文解读] Simplified TinyBERT: Knowledge Distillation for Document Retrieval

Xuanang Chen, Ben He+3 more|arXiv (Cornell University)|Sep 16, 2020

Topic Modeling参考文献 21被引用 6

一句话总结

该论文提出Simplified TinyBERT，一种知识蒸馏框架，通过将TinyBERT的两阶段训练合并为一个阶段，并在损失函数中引入硬标签，从而提升文档检索性能。其推理速度比BERT-Base快15倍，同时在MS MARCO和TREC 2019 DL基准测试中显著优于BERT-Base和标准TinyBERT。

ABSTRACT

Despite the effectiveness of utilizing the BERT model for document ranking, the high computational cost of such approaches limits their uses. To this end, this paper first empirically investigates the effectiveness of two knowledge distillation models on the document ranking task. In addition, on top of the recently proposed TinyBERT model, two simplifications are proposed. Evaluations on two different and widely-used benchmarks demonstrate that Simplified TinyBERT with the proposed simplifications not only boosts TinyBERT, but also significantly outperforms BERT-Base when providing 15$ imes$ speedup.

研究动机与目标

• 研究知识蒸馏在文档排序任务中的有效性，该任务因BERT的高推理成本而受限于部署。
• 通过简化训练流程并增强蒸馏在排序任务中的效果，改进TinyBERT模型。
• 实现一个更小、更快的学生模型，在保持高效率的同时，其排序性能超越BERT-Base。

提出的方法

• 使用组合损失函数，将TinyBERT的任务特定蒸馏阶段的两个步骤合并为单一步骤。
• 提出统一的损失函数，结合注意力、隐藏状态、嵌入和软标签蒸馏损失：$\mathcal{L} = \mathcal{L}_{attn} + \mathcal{L}_{hidn} + \mathcal{L}_{emb} + \mathcal{L}_{soft}$。
• 将硬标签引入损失函数，以增强模型对相关与非相关文档的判别能力。
• 采用单阶段蒸馏流程，以BERT-Base的前k层初始化学生模型，无需单独的通用蒸馏阶段。
• 在固定大小的学生模型上，对温度$T \in \{1,5,10\}$和知识蒸馏权重$\alpha \in \{0.2,0.5,0.7\}$进行超参数搜索。
• 使用MRR@10和NDCG@10等标准排名指标，在MS MARCO和TREC 2019 DL Track基准上评估模型性能。

实验结果

研究问题

• RQ1 标准知识蒸馏和TinyBERT能否有效应用于文档检索任务，该任务对推理成本有严格限制？
• RQ2 将TinyBERT的任务特定蒸馏的两个步骤合并为一步，是否能提升性能并提高训练效率？
• RQ3 在蒸馏损失中加入硬标签，是否能增强模型区分相关与非相关文档的能力？
• RQ4 Simplified TinyBERT的简化训练流程是否能在实现15倍推理加速的同时，使排名性能超越BERT-Base？
• RQ5 所提出的简化方法是否能在不同模型尺寸和重排序深度下保持或提升性能？

主要发现

• 在15倍推理加速条件下，Simplified TinyBERT的L6_H768配置在MS MARCO开发集上显著优于BERT-Base和标准TinyBERT。
• 在TREC 2019 DL测试集上，使用小型L3_H384学生模型时，Simplified TinyBERT在NDCG@10指标上优于BERT-Base，而标准TinyBERT在相同设置下表现显著逊于BERT-Base。
• 在所有重排序深度（从top-10到top-100）下，Simplified TinyBERT始终优于BERT-Base和TinyBERT，且在MS MARCO开发集top-100深度下，MRR@10指标提升达3.5%。
• 与标准TinyBERT相比，单步训练简化使第二阶段训练时间减少42–52%：L6_H768模型的训练时间从29.95小时降至14.37小时。
• 消融实验表明，单步训练和硬标签集成两种简化方法均独立且协同地促进了性能提升。
• 该模型在蒸馏模型中实现了当前最优的文档检索性能，其中L3_H384变体在实现15倍加速的同时，排名效果超越BERT-Base。