【论文】 Simplified TinyBERT文档检索的知识蒸馏

本文链接：https://blog.youkuaiyun.com/elf1110/article/details/124647909

Simplified TinyBERT: Knowledge Distillation for Document Retrieval

ECIR2021（European Conference on Information Retrieval）

论文链接：https://arxiv.org/abs/2009.07531v2

code：GitHub - cxa-unique/Simplified-TinyBERT: Simplified TinyBERT: Knowledge Distillation for Document Retrieval

Abstract

因为BERT的高计算成本限制了它的使用。

本文首先实证研究了两种知识蒸馏模型在文档排序任务中的有效性。此外，在最近提出的 TinyBERT 模型之上，还提出了两个简化。

Simplified TinyBERT 不仅提升了 TinyBERT，也优于 BERT-Base15 倍速。

Introduction

BERT 中的大量参数使其应用不广泛。知识蒸馏（KD）可用于学习较小的 BERT 模型，而不会过多影响有效性。

TinyBERT 被提出以两阶段蒸馏方法在预测层和中间层上进行蒸馏，并在多个 NLP 任务上取得了接近教师模型的有效性。

本文首先研究标准知识蒸馏模型和最近的 TinyBERT在文档排序任务中的用途。此外，对 TinyBERT 提出了两个简化，希望进一步提高蒸馏排序模型的有效性。

数据集MS MARCO 和 TREC 2019 DL Track

Background

Passage-level BERT-based Document Re-ranking

文档被拆分为重叠的段落，通过BERT得到与查询的一个相关性分数，之后产生文档的分数（如利用BERT-MaxP）对文档重排序。

Knowledge Distillation (KD)

一些对 BERT 的 KD 方法，例如 DistilBERT 、BERT-PKD 、TinyBERT 和 MiniLM 。早期的KD方法依赖于教师模型的软标签，设计了一个损失函数，使学生模型直接模拟教师模型的输出。同时，实际的注释也被考虑在损失函数中，两种交叉熵损失分别被称为软损失 $L_{soft}$ 和硬损失 $L_{hard}$ 。

TinyBERT中，预训练和微调知识分别从一般阶段和特定任务阶段的目标任务上的预训练 BERT 和微调 BERT 中提炼出来。采用三个 MSE 损失，使学生模型从教师模型的三种内部权重中学习：即注意力权重、隐藏状态权重和不同层的嵌入权重 $L_{attn}$ 、 $L_{hidn}$ 和 $L_{emb}$ 。中间层两个阶段都使用 $L_{attn}$ 、 $L_{hidn}$ 和 $L_{emb}$ 。而预测层仅在特定任务阶段使用 $L_{soft}$ 进行蒸馏。

Simplified TinyBERT for Ranking

对 TinyBERT 模型提出了两个简化，希望在文档排序任务上取得更好的性能。

Method

Merge two steps in the task-specific stage into one step

TinyBERT 涉及两个阶段，第二阶段有两个步骤，其中训练过程非常耗时。然而，本文的实证研究发现这两个步骤可以通过简单地一次优化所有损失来合并为一个步骤，如公式1

这种简化不仅减少了训练时间，而且提高了排序性能，如表 1 所示。这意味着学生模型可以更有效地学习预测层和中间层。

此外还发现可以通过将两个阶段合并为一个来进一步简化 TinyBERT 蒸馏过程，即，如果可用，则使用预训练的 BERT 模型，并使用其前 k 层来初始化学生模型来代替一般的蒸馏阶段。例如，表 1 中被称为 L6 H768 的学生模型也可以通过使用 BERT-Base 的前六层初始化学生模型，只用一个阶段进行提炼，而不会影响性能。将在以后的工作中对这部分做进一步的研究。

推测硬标签可以帮助更好地区分相关和非相关文档。因此，通过将蒸馏过程中的硬损失添加到等式（1）中，最终得到等式（2）