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

TL;DR

The paper tackles catastrophic forgetting and distribution shift that occur when fine-tuning LLMs for downstream tasks. It introduces Self-Distillation Fine-Tuning (SDFT), which rewrites task targets into the seed LM's distribution via a distillation template and trains on these distilled targets to preserve broad capabilities. Across experiments with Llama-2-7b-chat and LoRA, SDFT matches or surpasses vanilla fine-tuning on target tasks while maintaining safety alignment and general knowledge. This plug-and-play approach offers a practical strategy to adapt LLMs to specific tasks without sacrificing their general instruction-following abilities.

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.

Figures (18)

• Figure 1: Unlike vanilla fine-tuning, which may compromise seed LMs, our proposed self-distillation fine-tuning (SDFT) approach enhances seed LMs with improved downstream task performance while largely maintaining broad capabilities already learned.
• Figure 2: Left: An illustration of a generated distilled response that demonstrates a reduced distribution shift relative to the seed LLM. Right: The diminished distribution shift contributes to a moderate parameter shift, thereby alleviating the issue of catastrophic forgetting.
• Figure 3: The distillation template used in most of our experiments. It designates the original response as "reference answer" and prompts the model to generate a response using the reference answer as a guide.
• Figure 4: Performance comparisons of models on general knowledge benchmarks after fine-tuning on each dataset, as reported in the OpenLLM Leaderboard. Fine-tuning on these datasets demonstrates a marginal effect on the models' general knowledge.
• Figure 5: With increasing data for fine-tuning, there is a decrease in models' performance across various benchmarks, including math, safety alignment and instruction following capability.