PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhaorui Yang, Tianyu Pang, Haozhe Feng, Han Wang, Wei Chen, Minfeng Zhu, Qian Liu

arXiv:2402.13669v2cs.CL

TL;DR

Fine-tuning for specific tasks can threaten general instruction-following abilities through catastrophic forgetting. SDFT uses a distilled dataset of self-generated responses to reduce the distribution gap. The authors report diminished forgetting and comparable or superior targeted-task performance to vanilla fine-tuning.

  • Problem

    Task-specific fine-tuning can improve downstream performance while risking loss of general instruction-following abilities through catastrophic forgetting.

    Source: §1 · p. 1

  • Method

    SDFT uses semantically equivalent responses generated by the seed model as surrogate fine-tuning targets to reduce the distribution gap.

    Source: §1 · p. 2

  • Results

    SDFT effectively diminishes forgetting and achieves comparable or superior performance to vanilla fine-tuning on targeted tasks.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The authors present SDFT as a plug-and-play strategy for improving target-task performance while maintaining LMs’ broad capabilities.

    Source: §6 · p. 8

    Most experiments use Llama-2-7b-chat with LoRA, and safety evaluation covers only Advbench with fixed adversarial suffixes.

    Source: §6 · p. 9

Abstract

from arXiv · show

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

The paper frames downstream fine-tuning as a trade-off: task performance may improve while general instruction-following abilities and safety alignment are lost. It proposes SDFT to address this by using self-generated responses to reduce the distribution gap, and reports better results than vanilla fine-tuning across benchmarks.

1 Introduction

  • Vanilla fine-tuning can improve downstream performance but risks catastrophic forgetting of general instruction-following abilities and safety alignment.The authors cite studies finding that even benign-data fine-tuning can compromise seed-model safety.Source: §1 · pp. 1–2
  • SDFT hypothesizes that forgetting stems from the distribution gap, then uses semantically equivalent self-generated responses as surrogate targets to preserve the seed model’s distribution.The paper describes systematic evaluations across downstream tasks, helpfulness, and safety alignment.Source: §1 · p. 2
  • Across benchmarks, SDFT outperforms vanilla fine-tuning, mitigating forgetting and improving HumanEval pass@1 from 9.8 to 15.2 after vanilla fine-tuning had reduced it from 13.4.The authors also report that increasing the distilled-dataset proportion decreases catastrophic forgetting.Source: §1 · p. 2

2 Related Work

The related-work discussion covers fine-tuning, continual learning, alignment, and prompting-based methods. It distinguishes SDFT by using self-generated responses to bridge the distribution gap and address catastrophic forgetting during fine-tuning.

Fine-Tuning

  • Self-Play Fine-tuning uses the same LLM as generator and discriminator to steer preference for annotated responses, but the authors say it does not alleviate forgetting as the model distribution converges to training data.The related-work passage cites Chen et al. (2024) for Self-Play Fine-tuning.Source: §2 · p. 2

Continual Learning

  • Continual learning seeks to acquire new knowledge while mitigating forgetting, but LLM-scale parameter and task spaces and unavailable training data limit established approaches.The passage describes replay, parameter-importance computation, and task-specific neurons as traditional approaches.Source: §2 · p. 2

Alignment

  • Prior alignment methods include instruction tuning, reinforcement learning from human feedback, and self-alignment, while benign-data fine-tuning can still compromise safety.The authors state that SDFT can mitigate such safety degradation.Source: §2 · pp. 2–3

Prompting Based Learning

  • Unlike prompting methods that use generated responses for supervised tuning or iterative refinement, this work uses them to bridge the distribution gap and address forgetting.The passage names Self-Instruct (Wang et al., 2022), WizardLM, Self-Refine (Madaan et al., 2024), and Self-Reward (Yuan et al., 2024).Source: §2 · p. 3

3 Method

The method section outlines fine-tuning before presenting SDFT and its implementation. SDFT rewrites task responses into the seed model’s distribution and uses the resulting responses for fine-tuning.

3.1 Fine-tuning LLMs

  • Fine-tuning updates the seed LM using task examples to minimize negative log-likelihood, aligning the fine-tuned model’s response distribution with the task dataset.The seed LM is a model that needs further fine-tuning and is denoted f with parameters θ.Source: §3.1 · p. 3

3.2 Self-Distillation Fine-Tuning

  • SDFT prompts the seed LM to rewrite each task response into its own distribution, then uses the distilled response as the fine-tuning target instead of the original.For math reasoning, a heuristic compares the distilled and original final answers; otherwise, the original response is kept.Source: §3.2 · pp. 3–4

3.3 Distillation Template

  • The task-independent distillation template labels the original response as a “reference answer” and guides the model to generate its own response.The template is slightly adjusted for math-reasoning datasets.Source: §3.3 · p. 4, Figure 3 · p. 4

4 Experiments

The experiments compare SDFT with vanilla fine-tuning across downstream tasks, safety and helpfulness, and general knowledge. Results report mitigated forgetting, comparable or superior downstream performance, and maintained alignment.

4.1 Experimental Setup

  • Most experiments use Llama-2-7b-chat with LoRA for both methods; datasets over 10,000 examples are generally subsampled to 2,000, versus 20,000 for OpenHermes.The seed model is used in most experiments, except where stated otherwise.Source: §4.1 · p. 4

4.2 Datasets for Fine-tuning and Evaluation

  • Evaluation spans single-task GSM8K, OpenFunctions, and MagiCoder/HumanEval, multi-task instruction datasets, safety and jailbreak rates, AlpacaEval helpfulness, and five general-knowledge benchmarks.The multi-task datasets include Alpaca, Dolly, LIMA, and OpenHermes.Source: §4.2 · pp. 4–5

4.3 SDFT Achieves Better Results on Downstream Tasks

  • SDFT mitigates forgetting and achieves comparable or superior performance across downstream tasks.OpenFunctions vanilla fine-tuning lowers coding from 13.41 to 9.76; SDFT scores 15.24 on HumanEval versus the seed model’s 13.41, and target-task accuracy is 36.61 versus vanilla fine-tuning’s 34.82.Source: Table 1 · p. 5, §4.3 · pp. 5–6

4.4 SDFT Preserves Alignment

  • SDFT mitigates vanilla fine-tuning’s safety and helpfulness declines across single-task and multi-task settings, including OpenHermes jailbreak safety.On OpenHermes, SDFT raises jailbreak safe rate from 61.54 to 87.50.Source: Table 2 · p. 5, §4.4 · p. 6

4.5 General Knowledge Remains Intact

  • General-knowledge performance remains relatively unaffected by fine-tuning; after OpenFunctions fine-tuning, the gap from the seed model is under 1.The passage reports the same pattern after SDFT.Source: §4.5 · p. 6

5 Analysis

Across the tested settings, SDFT outperforms vanilla fine-tuning on target tasks while reducing forgetting. The analysis associates greater distribution shift with benchmark decline and finds performance robust across distillation templates.

5 Analysis

  • The analysis measures distribution shift using seed-model response comparisons, BLEU-4, ROUGE-L, embedding cosine similarity, and parameter distance.Lower response-similarity scores and greater parameter distance indicate greater distribution shift.Source: §5 · pp. 6–7
  • Fine-tuning has a marginal effect on general-knowledge benchmark performance, both for vanilla fine-tuning and SDFT.For OpenFunctions, the fine-tuned model's performance disparity from the seed LM is less than 1.Source: Figure 4 · p. 7

5.1 Distribution Shift Correlates with Catastrophic Forgetting

  • As fine-tuning sample size grows, BLEU-4, ROUGE-L, and embedding similarity fall while parameter shift rises; benchmark performance also declines, suggesting intensified catastrophic forgetting.The cited benchmarks include GSM8K, MultiArith, Advbench, and AlpacaEval.Source: §5.1 · p. 7, Figure 6 · p. 8, Figure 5 · p. 8
  • As mix ratio increases, BLEU-4, ROUGE-L, and embedding similarity rise while parameter shift falls, indicating reduced distribution shift; benchmark performance improves, and SDFT has higher fine-tuned-to-seed similarity than vanilla fine-tuning.Source: §5.1 · p. 7

5.2 Robustness among Distillation Templates

  • Performance across diverse benchmarks remains consistent with both distillation templates, supporting SDFT's robustness to template wording.The templates are labeled “Using” and “Refer.”Source: §5.2 · pp. 7–8

5.3 Efficacy of SDFT Across Model Scales and Architectures

  • On OpenFunctions, Table 5 reports that in all scenarios SDFT consistently outperforms vanilla fine-tuning on the target task and reduces forgetting across all other tasks.Source: §5.3 · p. 8

6 Conclusions and Limitations

The paper concludes that SDFT reduces catastrophic forgetting while achieving comparable or better target-task performance than vanilla fine-tuning. The authors qualify these findings by noting that most experiments use Llama-2-7b-chat with LoRA and that safety tests cover only Advbench and fixed adversarial suffixes.

6 Conclusions and Limitations

  • SDFT reduces catastrophic forgetting and achieves comparable or superior performance to vanilla fine-tuning on targeted tasks.The conclusion identifies degradation in general task capabilities, safety alignment, and helpfulness as concerns associated with distribution shift during fine-tuning.Source: §6 · pp. 8–9
  • Most experiments use Llama-2-7b-chat with LoRA; larger models and full fine-tuning remain to be explored.Source: §6 · p. 9
  • Safety evaluations are limited to Advbench and fixed adversarial suffixes, leaving robustness against other jailbreaking strategies for future work.Source: §6 · p. 9

Ethics Statement

The ethics statement reports no additional risks from SDFT and acknowledges potential biases in the training datasets and seed model.

Ethics Statement

  • The authors state that SDFT mitigates forgetting, including safety-alignment degradation, and therefore does not entail additional risks.Source: Ethics Statement · p. 9
  • The authors acknowledge possible biases in the open-source English datasets and the Llama-2-chat seed model.Source: Ethics Statement · p. 9

A Experimental Details

Experiments primarily use LoRA on Llama-2-7b-chat, with dataset-specific sample sizes and training epochs. Evaluation uses separate frameworks for helpfulness, leaderboard benchmarks, and coding.

A Experimental Details

  • Most experiments tune the LoRA query and value matrices at rank 8, with a learning rate starting at 1 × 10−4 and batch size 8.The learning rate decays to zero with cosine annealing; the model follows the default Llama-2 configuration.Source: A Experimental Details · p. 13
  • Alpaca, Dolly, and MagiCoder use 2,000 sampled examples for 2 epochs, while OpenHermes uses 20,000 examples for 2 epochs.Source: A Experimental Details · p. 13
  • GSM8K and OpenFunctions use their full training sets for 5 epochs, while LIMA uses its full set for 2 epochs.Source: A Experimental Details · p. 13
  • Helpfulness is evaluated with AlpacaEval and GPT-4, leaderboard benchmarks with lm-evaluation-harness, and coding with HumanEval and bigcode-evaluation-harness.Source: A Experimental Details · p. 13

B Templates and Examples

This section presents the experiment templates and illustrative distilled-data examples. Most experiments use the standard Alpaca template, while reasoning datasets receive modified templates for training and distillation and a specified answer format at evaluation.

Loading 2402.13669v2…