Self-Distillation for Further Pre-training of Transformers

ICLR2023| Self-Distillation for Further Pre-training of Transformers基于自蒸馏的Transformer下游任务进一步微调改进算法

• 论文链接：https://openreview.net/pdf?id=kj6oK_Hj40
• 源码链接：https://openreview.net/attachment?id=kj6oK_Hj40&name=supplementary_material

简介

大型Transformer模型在大量未标注数据预训练标注数据集后在标注数据集微调的应用已经在各种视觉和自然语言处理下游任务上取得了巨大的成功。然而如果预训练域和微调域存在较大差异时，直接微调的方法可能获得次优性能。一些改进方法提出在微调前在目标未标注数据集上继续预训练的策略。然而这些方法只能在语言模型上使用，在视觉Transformer模型上容易出现过拟合问题。为了解决这个限制，本文提出一种自蒸馏作为正则化的策略。

本文方法

相关工作

问题定义

假定神经网络参数$({\theta }_{init},{\varphi }_{init})$，神经网络参数由大容量未标注数据和掩蔽自编码策略（masked auto-encoding）预训练得到。$f_{{\theta }_{init}}$是用于抽取输入隐藏表示的编码器，$g_{{\varphi }_{init}}$是重建掩蔽输入的解码器。主要目标是在标注数据集${\mathcal{D}}^{tr}=\{({\mathbf{x}}^{(i)},{\mathbf{y}}^{(i)}){\}}_{i=1}^n$上微调带随机初始化特定任务头$h_w$的预训练模型$f_{{\theta }_{init}}$。将预训练算法定义为$\mathcal{A}$。优化目标可以描述为：
$$\underset{\theta ,w}{\min }{L}_{CE}(\theta ,w,{\mathcal{D}}^{tr})$$

进一步微调策略

在少量特定域标注数据条件下微调容易过拟合。Don’t stop pretraining: Adapt language models to domains and tasks展示了在目标未标注数据集 D u = { x ( i ) } i = 1 n \mathcal{D}^{u} = \{\mathbf{x}^{(i)}\}_{i=1}^{n} Du={ x(i)}i=1n​上预训练后在训练集${\mathcal{D}}^{tr}$的方法在提升泛化性能上是有效的。其中${\mathcal{D}}^u$与${\mathcal{D}}^{tr}$相同除了移除了标签。

掩蔽自编码

定义${\mathbf{x}}^{(i)}=(x_1^{(i)},\dots ,x_K^{(i)})$是一系列图像的补丁或固定长度$K$的句子令牌。对于每一个$x_k^{(i)}$独立地根据以$\gamma$概率的伯努利分布得到二进制掩码${\mathbf{z}}^{(i)}=(z_1^{(i)},\dots ,z_K^{(i)})$。如果$z_k^{(i)}=1$，$x_k^{(i)}$被一个特殊的掩码令牌取代，否则在掩蔽输入使用相同的$x_k^{(i)}$。定义${\hat{\mathbf{x}}}^{(i)}=({\hat{x}}_1^{(i)},\dots ,{\hat{x}}_K^{(i)})$，掩蔽自编码的最终目标可以描述为：
$${\mathcal{L}}_{\mathrm{MAE}}\left(\theta ,\varphi ;{\mathcal{D}}^u\right)=\frac{1}{n}\sum _{i=1}^n{\mathbb{E}}_{{\mathbf{z}}^{(i)}\sim p_{\gamma ,T}(\mathbf{z})}\left[-\sum _{k=1}^K\frac{z_k^{(i)}}{Z^{(i)}}\cdot \log p_{\theta ,\varphi }\left(x_k^{(i)}\mid {\hat{\mathbf{x}}}^{(i)}\right)\right],Z^{(i)}=\sum _{k=1}^K{z}_k^{(i)}$$

基于自蒸馏的进一步微调策略

给定预训练模型$g_{{\varphi }_{init}}\circ f_{{\theta }_{init}}$，继续基于目标未标注数据集${\mathcal{D}}^u$和掩蔽自编码策略训练模型获得编码器$f_{{\theta }_0}$和$g_{{\varphi }_0}$。丢弃解码器考虑使用编码器$f_{{\theta }_0}$作为自蒸馏的教师模型。复制预训练初始网络参数$g_{{\varphi }_{init}}\circ f_{{\theta }_{init}}$作为学生模型，进一步使用掩蔽自编码策略预训练学生模型，且强制学生模型$f_{{\theta }_{init}}$编码器的隐藏表示接近教师模型的隐藏表示：
$${\mathcal{L}}_{Distill}(\theta ;{\theta }_0,{\mathcal{D}}^u)=\frac{1}{n}\sum _{i=1}^n||f_{\theta }({\mathbf{x}}^{(i)})-\text{StopGrad}({\mathbf{x}}_0^{(i)})|{|}_2^2$$
本文算法可以重复这个过程执行多轮自蒸馏，其中前一轮自蒸馏的学生称为教师模型。新的学生模型用预训练的权重初始化。经验发现，第一轮的自蒸馏在提升下游任务的最终范围性能方便发挥了最重要的作用。在自蒸馏之后，丢弃解码器$g_{{\theta }_1}$，使用初始化${\theta }_1$权重和随机初始化的特定任务头$h_{\omega }$微调学生编码器$f_{{\theta }_1}$