Self-Distillation for Continual Learning in Masked Language Models
Abstract
Masked language modeling is a widely used pretraining objective for learning reusable representation backbones. In practical settings, models pretrained with masked language modeling often need to be adapted sequentially to new supervised tasks without losing performance on previously learned ones. We introduce SDFT-MLM, a self-distillation fine-tuning method for continual adaptation of such pretrained models. SDFT-MLM combines a supervised task loss with a token-level distillation penalty from an exponential-moving-average teacher. As a case study, we apply SDFT-MLM to protein language models, where new experimental assays often become available over time and models must learn new protein property predictors while retaining performance on earlier tasks. We evaluate continual protein property regression across four tasks using ESM2 as the encoder. To assess robustness to task order, we fine-tune the model under three dataset orderings. Across all orderings, SDFT-MLM improves final test Spearman correlation on earlier tasks, which are most susceptible to forgetting, while maintaining performance on later tasks comparable to standard supervised fine-tuning. These results suggest that token-level self-distillation provides a simple and effective mechanism for improving retention during sequential adaptation of models pretrained with masked language modeling.