Self-Distillation Bridges Distribution Gap in Language Model Fine-Tuning — part 1/2

Cleaned-text code points [0, 15661). Bibliography omitted here; source.md includes it.

# Self-Distillation Bridges Distribution Gap in Language Model Fine-Tuning

###### 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 .

## 1 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 ( Brown et al., 2020 ) and PaLM ( Chowdhery et al., 2023 ) 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) ( Ouyang et al., 2022b ; Chung et al., 2022 ) has further propelled the capabilities of LLMs, particularly in enhancing their instruction-following abilities.

Interestingly, even when starting with the same base LLM ( Touvron et al., 2023 ; Bai et al., 2023 ) , minor variations in the supervised dataset can lead to significant differences in model performance ( Zhou et al., 2023 ; Wang et al., 2023 ) . Consequently, the open-source community has witnessed rapid growth in the diversity of LLM variants, incorporating various SFT datasets and techniques, thereby enhancing their usefulness and accessibility.

However, SFT typically prioritizes improving general instruction-following abilities, suggesting that LLMs with SFT might face challenges in specific downstream tasks. As a result, repurposing these models as Seed Language Models (seed LMs) for subsequent fine-tuning tailored to specific downstream tasks has emerged as an appealing approach. While the approach seems promising, our preliminary study reveals the challenge of simultaneously enhancing task-specific performance and preserving general instruction-following abilities through vanilla fine-tuning, primarily due to the issue of catastrophic forgetting. Echoing our findings, recent studies have highlighted that fine-tuning, even with benign datasets, can compromise the safety of seed LMs ( Qi et al., 2024 ; Yang et al., 2023 ; Zhan et al., 2023 ; Pelrine et al., 2023 ) . As evidenced, fine-tuning methods aimed at mitigating catastrophic forgetting are still absent.

In this paper, we propose a novel fine-tuning method, Self-Distillation Fine-Tuning (SDFT), to mitigate catastrophic forgetting during fine-tuning. We hypothesize that catastrophic forgetting stems from the distribution gap between the task dataset and the seed LMs. To address the issue, as shown in Figure 1 , SDFT first prompts the seed LM to generate responses that uphold semantic equivalence with the original responses present in the task dataset, resulting in the distilled dataset. A representative example of rewriting is depicted in Figure 2 . After rewriting, the self-generated responses serve as surrogate targets during subsequent fine-tuning. Through the approach, SDFT inherently maintains the original distribution, avoiding distribution shift and thereby preserving capabilities.

We systematically evaluate SDFT by comparing its performance against that of vanilla fine-tuning and the seed LM across a variety of benchmarks. These benchmarks encompass: (1) diverse downstream tasks, including mathematical reasoning , tool using and code generation ; (2) assessments of general helpfulness and safety alignment. Results on all benchmarks demonstrate the superiority of SDFT compared to vanilla fine-tuning. For instance, vanilla fine-tuning on the OpenFunctions dataset ( Patil et al., 2023 ) leads to a significant decrease in pass@1 on the HumanEval benchmark ( Chen et al., 2021 ) from 13.4 13.4 to 9.8 9.8 , constituting a decline of 27 % 27\% . In contrast, SDFT not only mitigates this degradation, but also marginally enhances the accuracy to 15.2 15.2 . The in-depth analysis of our method indicates that increasing the proportion of distilled dataset for fine-tuning leads to a decrease in catastrophic forgetting, thereby confirming that SDFT mitigates catastrophic forgetting by bridging the distribution gap.

## 2 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 ( Roziere et al., 2023 ; Luo et al., 2024 ) , arithmetic ( Luo et al., 2023a ) , healthcare ( Jin et al., 2023 ) and finance ( Wu et al., 2023 ) . Vanilla fine-tuning directly maximizes the log-likelihood of target responses. Similar to our work, Self-Play Fine-tuning ( Chen et al., 2024 ) 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 ( French, 1999 ) . A related domain is continual learning ( Kirkpatrick et al., 2017 ; Lopez-Paz and Ranzato, 2017 ) , 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 ( Scialom et al., 2022 ; Luo et al., 2023b ) , the computation of parameter importance ( Kirkpatrick et al., 2017 ; Aljundi et al., 2018 ) , or the assignment of distinct neurons to different tasks ( Mallya and Lazebnik, 2018 ) . 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 ( Kirkpatrick et al., 2017 ; Lopez-Paz and Ranzato, 2017 ; Scialom et al., 2022 ) . Although recent research ( Luo et al., 2023b ; Scialom et al., 2022 ) 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 ( Perez et al., 2022 ; Ganguli et al., 2022 ) . 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 ( Ouyang et al., 2022a ; Touvron et al., 2023 ) , reinforcement learning from human feedback ( Ouyang et al., 2022a ; Bai et al., 2022 ) , and self-alignment techniques ( Sun et al., 2023 ) . Employing these alignment techniques, LLMs strike a dedicate tradeoff between utility and safety ( Bianchi et al., 2023 ; Qi et al., 2024 ) . 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-tuning with benign data can lead to compromised safety ( Qi et al., 2024 ; Yang et al., 2023 ; Zhan et al., 2023 ; Pelrine et al., 2023 ) . 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 Wang et al. (2022) and WizardLM Xu et al. (2024) utilize the generated responses for supervised fine-tuning, with the latter employing GPT-4 as the generator. Other methods, such as Self-Refine Madaan et al. (2024) and Self-Reward Yuan et al. (2024) , 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.

## 3 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.

### 3.1 Fine-tuning LLMs

While LLMs demonstrate remarkable proficiency across various tasks, they often encounter limitations when it comes to downstream tasks that necessitate fine-tuning. Specifically, we refer to a LM in need of further fine-tuning as seed LM, denoted as f f and parameterized by θ \theta . The seed LM typically undergoes general SFT, indicating its capacity to map any natural language instruction x ∈ X x\in X contextualized by the task description c ∈ C c\in C , to its corresponding output y ∈ Y y\in Y . f θ : C × X → Y . f_{\theta}:C\times X\rightarrow Y. (1)

The fine-tuning process of the seed LM can be outlined as follows: for the target task t t with context c t c^{t} , each task example ( x t , y t ) (x^{t},y^{t}) is utilized to update the model parameters. This update aims at minimizing the disparity between the data distribution and the LM distribution, as expressed below:

L FT ​ ( θ ) = − log ⁡ f θ ​ ( y t ∣ c t , x t ) , L_{\text{FT}}(\theta)=-\log f_{\theta}(y^{t}\mid c^{t},x^{t}), (2)

which seeks to minimize the negative log likelihood of the target output y t y^{t} given the context c t c^{t} and input x t x^{t} , with respect to the model parameters θ \theta . L FT L_{\text{FT}} converges when the generated response y ^ \hat{y} matches y t y^{t} , i.e., the distribution of fine-tuned LM aligns with the task dataset distribution.

### 3.2 Self-Distillation Fine-Tuning

As the distribution of the seed LM converges towards that of the task dataset, it naturally enhances performance on target tasks. However, vanilla fine-tuning is susceptible to catastrophic forgetting in general instruction-following capabilities and safety alignment.

To address this issue, we propose S elf- D istillation F ine- T uning (SDFT) to better align the distribution of the task dataset with that of the seed LM.

As depicted in Figure 2 , the initial step of SDFT involves prompting the seed LM to rewrite the original response y t y^{t} into y ~ \tilde{y} : y ~ ∼ f θ ​ ( y ∣ c t , x t , y t ) . \tilde{y}\sim f_{\theta}(y\mid c^{t},x^{t},y^{t}). (3) This step marks the primary distinction between our method and vanilla fine-tuning, as it involves mapping the original response into a response within the seed LM’s distribution. To accomplish the rewriting, we utilize a self-distillation template, which imposes minimal requirements on the seed LM, simply requiring it to adhere to our directive for rewriting responses. The exact specifications of this prompt are elaborated later.

Next, to ensure the quality of the distilled responses, we employ simple heuristics to evaluate the distilled response. For instance, in math reasoning problems, we extract the final answer from the distilled response y ~ \tilde{y} and compare it with the one from the original response y t y^{t} . Otherwise, we keep the original response. We formalize this conditional selection process as: y ~ ′ = { y ~ if Extract ​ ( y ~ ) = y t , y t otherwise. \tilde{y}^{\prime}=\begin{cases}\tilde{y}&\text{if Extract}(\tilde{y})=y^{t},\\ y^{t}&\text{otherwise.}\end{cases} (4) Finally, the distilled response is used as a replacement for the original response y t y^{t} for fine-tuning, i.e., the loss becomes: L SDFT ​ ( θ ) = − log ⁡ f θ ​ ( y ~ ′ ∣ c t , x t ) . L_{\text{SDFT}}(\theta)=-\log f_{\theta}(\tilde{y}^{\prime}\mid c_{t},x_{t}). (5) Hence, the distribution gap is mitigated by utilizing the distilled dataset instead of the task dataset, as depicted on the right side of Figure 2 .

### 3.3 Distillation Template

In our work, the distillation template plays a crucial role. Designed to be task-independent, it can be applied seamlessly across various tasks without requiring modification. Within this framework, the template designates the original response within the task dataset as the “reference answer” and guides the model to generate a response accordingly. The template employed in the majority of our experiments is illustrated in Figure 3 . When dealing with datasets involving math reasoning, we slightly adjust the template to better accommodate the reasoning process. Further details about these templates can be found in Appendix B .

## 4 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.

### 4.1 Experimental Setup

We utilize the Llama-2-7b-chat model ( Touvron et al., 2023 ) as the seed LM in most of our experiments, except where explicitly stated otherwise. Due to limited computation resources, we utilize the Low Rank Adaptation (LoRA) technique ( Hu et al., 2022 ) during both vanilla fine-tuning and our proposed SDFT.

To ensure fair comparison, we maintain consistency in nearly all hyperparameters for both methods. For datasets comprising more than 10,000 10,000 examples, we randomly select 2,000 2,000 examples for fine-tuning to ensure comparability in size across most datasets. For the OpenHermes dataset, we randomly select 20,000 20,000 examples to validate the effect of SDFT with larger, mixed dataset. More experimental details can be found in Appendix A .

### 4.2 Datasets for Fine-tuning and Evaluation

We fine-tune the seed LM on a variety of datasets, including those for both single-task and multi-task scenarios. We then evaluate the performance of both the seed model and the fine-tuned models across diverse tasks. The datasets for fine-tuning and evaluation are categorized as follows:

Single-task datasets. For single-task datasets, we explore boosting the mathematical reasoning, tool using, and code generation capabilities of LMs during fine-tuning. The mathematical reasoning capabilities are improved using the GSM8K dataset ( Cobbe et al., 2021 ) , which comprises 8.8k high-quality arithmetic word problems designed at grade school level. The tool using proficiency is assessed by leveraging function-calling datasets such as the Gorilla OpenFunctions dataset ( Patil et al., 2023 ) . Additionally, code generation skills are boosted using the MagiCoder dataset ( Wei et al., 2023 ) , while evaluation is conducted using the HumanEval dataset ( Chen et al., 2021 ) .