Self-Distillation Bridges Distribution Gap in Language Model Fine-Tuning (2024)

Abstract: The surge in Large Language Models (LLMs) has revolutionized natural language processing, but fine-tuning them for specific tasks often encounters challenges in balancing performance and preserving general instruction-following abilities. In this paper, we posit that the distribution gap between task datasets and the LLMs serves as the primary underlying cause. To address the problem, we introduce Self-Distillation Fine-Tuning (SDFT), a novel approach that bridges the distribution gap by guiding fine-tuning with a distilled dataset generated by the model itself to match its original distribution. Experimental results on the Llama-2-chat model across various benchmarks demonstrate that SDFT effectively mitigates catastrophic forgetting while achieving comparable or superior performance on downstream tasks compared to the vanilla fine-tuning. Moreover, SDFT demonstrates the potential to maintain the helpfulness and safety alignment of LLMs. Our code is available at https://github.com/sail-sg/sdft.

## introduction in recent years, the development of large language models (llms) has emerged as one of the most groundbreaking advancements in natural language processing (nlp). llms such as gpt-3 @xcite and palm @xcite have revolutionized the field by leveraging massive textual corpora during pre-training, enabling them to achieve remarkable few-shot performance across a wide range of tasks. the introduction of supervised fine-tuning (sft) @xcite has further propelled the capabilities of llms, particularly in enhancing their instruction-following abilities. interestingly, even when starting with the same base llm @xcite ), * corresponding authors vanilla fine-tuning self-distillation fine-tuning (ours) seed language models task dataset task dataset compromised language models enhanced language models distilled dataset ## related work fine-tuning fine-tuning is a prevalent strategy for improving the performance of models on downstream tasks, as demonstrated in domains including coding @xcite , arithmetic @xcite , healthcare @xcite and finance @xcite . vanilla finetuning directly maximizes the log-likelihood of target responses. similar to our work, self-play finetuning @xcite employs the identical llm as both generator and discriminator, steering the model to prefer annotated response over generated outputs. as the llm's distribution ultimately converges with that of the training data, the method does not alleviate forgetting during fine-tuning. continual learning fine-tuning enables models to adapt to new data distributions, improving their efficacy on downstream tasks. however, this process can lead to the loss of previously acquired knowledge, an issue known as catastrophic forgetting @xcite . a related domain is continual learning @xcite , which seeks to enable models to acquire new knowledge while mitigating such forgetting. traditional methods often depend on the preservation of historical data for replay @xcite , the computation of parameter importance @xcite , or the assignment of distinct neurons to different tasks @xcite . however, fine-tuning llms is particularly challenging due to their extensive parameter and task space, compounded by the frequent unavailability of original training datasets, which diminishes the feasibility of these established techniques @xcite @xcite . although recent research @xcite highlights the significance of continual learning for lms, there are scant feasible solutions for llms. in this paper, we conduct a comprehensive evaluation of the catastrophic forgetting issue during the fine-tuning of llms and propose a simple yet effective strategy specifically designed for llms. alignment as the capabilities of large language models (llms) expand, so does the potential for generating toxic content, engendering significant safety concerns @xcite . in response, various strategies have been proposed to align llms with human ethical standards and prevent the generation of toxic content. prevalent methods including instruction tuning @xcite , reinforcement learning from human feedback @xcite , and self-alignment techniques @xcite . employing these alignment techniques, llms strike a dedicate tradeoff between utility and safety @xcite . while these methods have demonstrated efficacy in safety alignment, they do not cover further risks that arise from fine-tuning. recent research reveals that even fine- seed lm distribution task dataset distribution distilled dataset distribution instruction: name three types of machine learning algorithms. response: three types of machine learning algorithms are supervised learning, unsupervised learning, and reinforcement learning. self-distillation via response rewriting task dataset instruction: name three types of machine learning algorithms. response: i can name three types of machine learning algorithms as follows:1. supervised learning: this type of algorithm … distilled dataset seed language models pr oje cti on tuning with benign data can lead to compromised safety @xcite @xcite . our proposed strategy can effectively mitigate such safety degradation. prompting based learning recently, the use of prompting in llms to generate responses for model training has garnered significant interest. approaches like self-instruct @xcite and wizardlm @xcite utilize the generated responses for supervised fine-tuning, with the latter employing gpt-4 as the generator. other methods, such as self-refine @xcite and self-reward @xcite , use the responses as feedback to iteratively refine the model's outputs. in contrast, our work introduces a novel perspective by leveraging the responses to bridge the distribution gap and address the catastrophic forgetting issue during the fine-tuning process. ## method in this section, we begin by outlining the process of fine-tuning, followed by the introduction of our proposed self-distillation fine-tuning method and its implementation details. ## experiments in this section, we begin by presenting the dataset employed for fine-tuning and evaluation purposes. following that, we conduct a comparative analysis of the experimental results obtained from vanilla fine-tuning and our proposed sdft approach across various tasks, encompassing mathematical reasoning, code generation, and tool using. finally, we assess the impact of both methods on safety, general knowledge, and helpfulness. ## analysis in this section, we conduct a detailed analysis to understand the impact of distribution shift on catastrophic forgetting. in addition to the evaluation metrics outlined in section 4, we incorporate four supplementary metrics to assess the degree of distribution shift. we utilize both the seed model and fine-tuned models to generate responses on the advbench @xcite dataset and engage in a comparative analysis of these responses. in particular, we calculate the bleu-4 and rouge-l scores for the fine-tuned models, using the outputs from the seed model as references to evaluate the extent of distribution shift. we also utilize sentence-bert @xcite to derive sentence embeddings and use the cosine similarity between these embeddings following @xcite . lastly, we quantify the extent of parameter shift by comparing the updated parameters with those of the seed model, considering their distance as a measure of the parameter shift magnitude. the lower the bleu-4, rouge-l, and embedding similarity scores, the greater the ## conclusions and limitations in this paper, we perform a systematic evaluation of catastrophic forgetting during the fine-tuning of language models for downstream tasks. gate catastrophic forgetting. extensive experiments show that sdft effectively diminishes forgetting and delivers comparable or superior performance to vanilla fine-tuning on targeted tasks. our study is subject to certain limitations. owing to constraints in computational resources, most of our experiments are based on the llama-2-7bchat model with lora. further investigations involving larger models and full fine-tuning remain to be explored. furthermore, our safety evaluations are limited to the advbench dataset and fixed adversarial suffixes, leaving the robustness against other jailbreaking strategies for future work.