[ACL 2022] BERT Learns to Teach: Knowledge Distillation with Meta Learning

[ACL 2022] BERT Learns to Teach: Knowledge Distillation with Meta Learning

Motivation

之前的工作通常训练一个大模型作为teacher，然后保持teacher不动训练student模型来完成teacher模型知识的迁移。然而这种做法有以下缺点：1）teacher不了解student的能力。近来有一些工作通过联合训练teacher和student来引入student-aware的蒸馏，但是这种方法也有提升的空间，因为2）teacher不是为了蒸馏优化的。之前的工作中，teacher通常都是为了自己的推理性能而优化的。然而，teacher不能感知到将其知识迁移给student的需求，这样不是最优的。为了解决这两个问题，我们提出了MetaDistil，一个新的使用meta-learning的蒸馏框架来利用学习过程中student的反馈来提升teacher的知识迁移能力。

Approach (Knowledge Distillation with Meta Learning)

Background

Knowledge Distillation

给定teacher $T$ 和student $S$，以及有标注的训练数据集$\mathcal{D}=\{(x_1,y_1),...,(x_N,y_N)\}$，student的损失函数可以写成下面形式：
$$\begin{gathered} \begin{array}{c}{\mathcal{L}}_S(\mathcal{D};{\theta }_S;{\theta }_T)=\frac{1}{N}\sum _{i=1}^N[\alpha {\mathcal{L}}_{\tau }(y_i,S(x_i;{\theta }_S)) \\ +(1-\alpha ){\mathcal{L}}_{KD}(T(x_i;{\theta }_T),S(x_i;{\theta }_S))]\end{array} \end{gathered}$$
${\mathcal{L}}_{\tau }$是task-specific的loss，${\mathcal{L}}_{KD}$是衡量teacher和student之间相似度的蒸馏损失，如teacher和student 概率分布的KL散度、logits的MSE等等。（此前有研究发现MSE比KL散度更稳定且效果略优，因此本文采用MSE）

Meta Learning

在涉及bi-level优化的meta learning算法中，存在inner-learner $f_i$和meta-learner $f_m$。inner-learner在meta-learner的帮助下被训练来完成任务$\tau$。$f_i$ 在 $f_m$ 的帮助下进行训练的过程被叫做inner-loop，在inner-loop后更新的inner-learner记为$f_i^{\prime }(f_m)$。而meta-leaner通过meta目标优化，通常是最大化inner-loop之后inner-learner的预期性能。这个学习过程被叫做meta-loop，通常通过更新后的inner-learner在一些验证集上的损失$\mathcal{L}(f_i^{\prime }(f_m))$的导数进行梯度下降完成。

Methodology

Pilot Updata

本来的meta learning目的是学习一个好的meta-learner，能够泛化到不同的inner-learner来解决不同的任务。而在MetaDistil中，我们只关注inner-learner的表现。为了同步meta-和inner-learner的学习进度，我们设计了一个pilot update机制。
对一个batch的训练数据$x$，我们复制了inner-learner $f_i$，然后在$x$上更新复制的$f_i^{\prime }$和meta-learner $f_m$。然后。我们丢弃$f_i^{\prime }$并使用更新了的$f_m$在同一份数据$x$上再次更新$f_i$。这个机制可以使用数据$x$来同时影响$f_i$和$f_m$，因此对齐了训练过程。

Learning to Teach

在MetaDistil架构中，teacher的优化目标是：teacher蒸馏之后student的性能。student网络 ${\theta }_S$ 是inner-learner，teacher网络${\theta }_T$ 是meta-learner。在每个训练step中，我们首先将${\theta }_S$ 复制一份成 ${\theta }_S^{\prime }$。然后给一个batch数据以及学习率$\lambda$，${\theta }_S^{\prime }$ 通过传统的蒸馏算法更新：
$$\begin{array}{c}{\theta }_S^{\prime }({\theta }_T)={\theta }_S-\lambda {\nabla }_{{\theta }_S}{\mathcal{L}}_S(x;{\theta }_S;{\theta }_T)\end{array}$$
然后我们观察实验的student参数 ${\theta }_S^{\prime }$ 和 student 在从quiz set上采样的一个batch数据的quiz loss $l_q={\mathcal{L}}_{\tau }(q,{\theta }_S^{\prime }({\theta }_T))$，它是一个关于teacher参数 ${\theta }_T$ 的函数。因此，我们可以对${\theta }_T$ 进行梯度下降来优化$l_q$：
$$\begin{array}{c}{\theta }_T\leftarrow {\theta }_T-\mu {\nabla }_{{\theta }_T}{\mathcal{L}}_{\tau }(q;{\theta }_S^{\prime }({\theta }_T))\end{array}$$
在我们的实验中，student没有在quiz set上训练，teacher只在quiz set上进行meta-update。我们也没有使用动态quiz set 策略，否则student将会在quiz set上训练过，这样loss就不会包含足够的信息了。在mate-update teacher之后，我们使用公式(2) 更新本来的student。使用公式(3) 优化teacher参数${\theta }_T$是最大化被teacher蒸馏后的student期望表现。meta-objective 允许teacher模型调整参数来更好地将知识蒸馏给student。

Experiments

Settings

目标：BERT-base蒸馏6层的BERT，使用模型logits间的MSE loss作为蒸馏目标。

Analysis

Why Does MetaDistill Work?

对于87%的更新，在进行真实更新后student的dev loss要小于没有进行pilot update的。这就说明了pilot update机制能更好地匹配student和teacher模型。
而且，我们发现前一半更新中，91%的更新让teacher与student更相似了（从logits的kl散度来看），这就表明teacher学着去适应低性能的student。在后一半更新中，这个百分比降到了63%。我们认为后面的训练过程中，teacher需要去超越student去自我进化来引导student进一步提升。

Hyper-parameter Sensitivity

Limitation

尽管MetaDistil达到最好性能的时间更久，但是它达到与PKD最好性能持平的时间基本相同。