Mitigating Forgetting in LLM Fine-Tuning via Low-Perplexity Token Learning

TL;DR

The paper tackles catastrophic forgetting during domain-specific fine-tuning of instruction-following LLMs and examines why LLM-generated data can boost target-task performance while preserving cross-domain capabilities. It reveals that low-perplexity token sequences in generated data correlate with reduced non-target degradation and proposes Selective Token Masking (STM) to emulate this effect by masking high-perplexity ground-truth tokens during training. Across multiple model families and fine-tuning strategies, STM achieves target-task gains comparable to Self-Output while substantially preserving non-target performance and requiring far less computation. The work provides a practical, loss- and perplexity-based lens on fine-tuning robustness and suggests a scalable path toward more robust domain adaptation of LLMs.

Abstract

Maintaining consistent model performance across domains is a fundamental challenge in machine learning. While recent work has explored using LLM-generated data for fine-tuning, its impact on cross-domain generalization remains poorly understood. This paper presents a systematic analysis revealing that fine-tuning with LLM-generated data not only improves target task performance but also reduces non-target task degradation compared to fine-tuning with ground truth data. Through analyzing the data sequence in tasks of various domains, we demonstrate that this enhancement of non-target task robustness stems from the reduction of high perplexity tokens found in LLM-generated sequences. Following our findings, we showed that masking high perplexity tokens in ground truth training data achieves similar non-target task performance preservation, comparable to using LLM-generated data. Extensive experiments across different model families and scales, including Gemma 2 IT 2B, Llama 3 8B Instruct, and three additional models, agree with our findings. To the best of our knowledge, this is the first work to provide an empirical explanation based on token perplexity reduction to mitigate catastrophic forgetting in LLMs after fine-tuning, offering valuable insights for developing more robust fine-tuning strategies.

Figures (14)

• Figure 1: An example MATH problem showing more high perplexity tokens (highlighted in red, perplexity $\geq 2.5$) in ground truth than Self-Output responses (Llama 3 self-generated responses).
• Figure 2: Fine-tuned Llama 3 8B Instruct on MBPP labels generated by various LLMs. Backward Transfer (BWT) measures performance drop on non-target tasks, showing that Low-perplexity labels training generally leads to less degradation (upper left) than ground truth (bottom right).
• Figure 3: Comparison of token-level perplexity (PPL) distributions between human-annotated Ground Truth (top) and Llama 3 8B Instruct generations for Rephrase and Self-Output sequence (middle and bottom), where PPL of Self-Output data is low and with fewer spikes.
• Figure 4: MBPP target task testing accuracy, validation and training loss of baseline finetuning with ground truth data, finetuning with Self-Output data, and STM strategy of perplexity filtering threshold=2.5 for Llama 3 8B Instruct. STM and self-output training yield better performances with much lower training and validation loss because of low-perplexity training.
• Figure 5: MATH SFT using STM method on Llama 3 8B Instruct on different levels of token filtering levels. The best in domain performance also matches peak performance on out of domain tasks : GSM8k, ARC, MBPP as well.