[论文解读] Self-Distillation for Further Pre-training of Transformers

Seanie Lee, Minki Kang+3 more|arXiv (Cornell University)|Sep 30, 2022

Multimodal Machine Learning Applications被引用 4

一句话总结

本文提出一种自蒸馏正则化方法，用于在领域特定的无标签数据上对视觉与语言Transformer模型进行进一步预训练。通过在掩码自编码过程中，使用进一步预训练的模型作为教师，指导初始预训练模型的副本（学生）进行表示蒸馏，该方法提升了下游分类任务的泛化能力，优于基线方法，并减少了视觉Transformer模型的过拟合现象。

ABSTRACT

Pre-training a large transformer model on a massive amount of unlabeled data and fine-tuning it on labeled datasets for diverse downstream tasks has proven to be a successful strategy, for a variety of vision and natural language processing tasks. However, direct fine-tuning of the pre-trained model may be suboptimal if there exist large discrepancies across data domains for pre-training and fine-tuning. To tackle this issue, several previous studies have proposed further pre-training strategies, where we continue to pre-train the model on the target unlabeled dataset before fine-tuning. However, all of them solely focus on language models and we empirically find that a Vision Transformer is vulnerable to overfitting as we continue to pretrain the model on target unlabeled data. In order to tackle this limitation, we propose self-distillation as a regularization for a further pre-training stage. Specifically, we first further pre-train the initial pre-trained model on the target unlabeled data and then consider it as a teacher for self-distillation. Then we take the same initial pre-trained model as a student and enforce its hidden representations to be close to those of the teacher while optimizing the student with a masked auto-encoding objective. We empirically validate the efficacy of self-distillation on a variety of benchmark datasets for image and text classification tasks. Experimentally, we show that our proposed method outperforms all the relevant baselines. Theoretically, we analyze the proposed method with a simplified model to understand how self-distillation for further pre-training can potentially help improve the performance of the downstream tasks.

研究动机与目标

• 解决预训练与微调数据之间的领域偏移问题，尤其在视觉Transformer中，标准的进一步预训练会导致过拟合。
• 克服现有正则化方法的局限性，这些方法通过惩罚从预训练初始化开始的权重更新，阻碍了模型适应能力。
• 开发一种在不依赖数据增强的前提下，提升进一步预训练过程中泛化能力的方法。
• 实现预训练Transformer模型在低资源下游任务中的有效适应，尤其在标注数据有限的情况下。
• 对自蒸馏在减少泛化误差和提升微调性能方面的作用进行理论与实证验证。

提出的方法

• 首先，使用目标领域无标签数据，通过掩码自编码（MAE）对初始预训练模型（如ViT或RoBERTa）进行进一步预训练，生成教师模型。
• 将进一步预训练模型的编码器作为教师，为知识蒸馏提供目标表示。
• 将原始预训练模型的副本作为学生，训练其在MAE预训练过程中匹配教师的隐藏表示。
• 通过联合优化目标函数来训练学生模型：包括掩码自编码重建损失，以及强制学生与教师表示一致的蒸馏损失。
• 完成自蒸馏后，在标注的下游数据集上对学生模型进行微调。
• 该方法同时应用于视觉与语言模型，并通过消融实验研究了蒸馏轮数的影响以及对权重与初始化距离的影响。

实验结果

研究问题

• RQ1 自蒸馏是否能有效正则化视觉Transformer的进一步预训练，从而缓解标准进一步预训练导致的过拟合？
• RQ2 与标准微调和进一步预训练相比，自蒸馏在下游泛化性能方面表现如何？
• RQ3 第一轮自蒸馏在减少初始预训练权重与最终微调权重之间L2距离方面起到什么作用？
• RQ4 如理论分析所建议，自蒸馏是否能减少下游微调阶段的泛化差距？
• RQ5 该方法在标注数据有限的多样化图像与文本分类基准测试中表现如何？

主要发现

• 与标准微调和进一步预训练基线相比，自蒸馏在图像与文本分类基准测试中显著提升了泛化性能。
• 在CIFAR-100与ImageNet-1K数据集上，该方法的准确率高于基线，尤其在低数据场景下表现更优。
• 第一轮自蒸馏具有最显著的正则化效果，相比后续轮次，能更显著地减少初始权重与最终模型权重之间的L2距离。
• 实证结果证实，权重与初始化距离的减小与泛化能力的提升密切相关，支持了权重接近性与泛化界之间理论联系的合理性。
• 通过简化模型的理论分析表明，自蒸馏可严格降低微调阶段监督损失的泛化界。
• 该方法在视觉与语言模型（包括ViT与RoBERTa）中均表现有效，展现出广泛的适用性。