[Paper Review] Self-Distillation for Further Pre-training of Transformers

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

Multimodal Machine Learning Applications4 citations

TL;DR

This paper proposes self-distillation as a regularization method for further pre-training of vision and language transformers on domain-specific unlabeled data. By using a further pre-trained model as a teacher to guide a copy of the initial pre-trained model (student) via representation distillation during masked autoencoding, the method improves generalization on downstream classification tasks, outperforming baselines and reducing overfitting in Vision Transformers.

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.

Motivation & Objective

• Address the challenge of domain shift between pre-training and fine-tuning data, especially in Vision Transformers where standard further pre-training leads to overfitting.
• Overcome the limitations of existing regularization methods that hinder adaptation by penalizing weight updates from pre-trained initialization.
• Develop a method that enhances generalization during further pre-training without relying on data augmentations.
• Enable effective adaptation of pre-trained transformers to low-resource downstream tasks with limited labeled data.
• Theoretical and empirical validation of self-distillation’s role in reducing generalization error and improving fine-tuning performance.

Proposed method

• First, further pre-train the initial pre-trained model (e.g., ViT or RoBERTa) on target-domain unlabeled data using masked autoencoding (MAE) to produce a teacher model.
• Use the encoder of the further pre-trained model as a teacher to provide target representations for knowledge distillation.
• Take a copy of the original pre-trained model as a student, and train it to match the teacher’s hidden representations during MAE pre-training.
• Optimize the student model with a combined objective: masked autoencoding reconstruction loss and a distillation loss enforcing agreement between student and teacher representations.
• After self-distillation, fine-tune the student model on the labeled downstream dataset.
• The method is applied to both vision and language models, with ablation studies on the number of distillation rounds and the effect on weight distance from initialization.

Experimental results

Research questions

• RQ1 Can self-distillation effectively regularize further pre-training in Vision Transformers, where standard further pre-training leads to overfitting?
• RQ2 How does self-distillation compare to standard fine-tuning and further pre-training in terms of downstream generalization performance?
• RQ3 What is the role of the first round of self-distillation in reducing the L2 distance between the initial pre-trained weights and the final fine-tuned weights?
• RQ4 Does self-distillation reduce the generalization gap in the downstream fine-tuning stage, as suggested by theoretical analysis?
• RQ5 How does the method perform across diverse image and text classification benchmarks with limited labeled data?

Key findings

• Self-distillation significantly improves generalization performance on image and text classification benchmarks compared to standard fine-tuning and further pre-training baselines.
• On the CIFAR-100 and ImageNet-1K datasets, the method achieves higher accuracy than baselines, especially in low-data regimes.
• The first round of self-distillation has the most significant regularization effect, reducing the L2 distance between the initial and final model weights more than subsequent rounds.
• Empirical results confirm that reduced weight distance from initialization correlates with improved generalization, supporting the theoretical link between weight proximity and generalization bound.
• Theoretical analysis with a simplified model shows that self-distillation strictly decreases the generalization bound on the supervised loss in the fine-tuning stage.
• The method is effective across both vision and language models, including ViT and RoBERTa, demonstrating broad applicability.