Preventing Catastrophic Forgetting: Behavior-Aware Sampling for Safer Language Model Fine-Tuning

TL;DR

This paper addresses catastrophic forgetting of safety during supervised fine-tuning of large language models and proposes a behavior-aware sampling framework to preserve safety with minimal additional data. It jointly optimizes for refusal-type safety signals and semantic diversity across harm categories using two sampling strategies: Stratified Safety Sampling (SSS) and Prototypical Safety Sampling (PSS), including behavioral variants that prioritize T1 refusals. Empirical results show substantial reductions in harmful outputs (up to 41% with only 0.5% additional data) while preserving helpfulness, with the strongest gains when combining diversity and refusal-focus (SSS-B, Cossim-B). The findings offer practical guidance for scalable safety alignment, demonstrating that targeted data selection can outperform brute-force augmentation in data-efficient fine-tuning of LLMs.

Abstract

Large language models often lose previously aligned safety behaviors when fine-tuned on benign data, a phenomenon known as catastrophic forgetting. Prior work shows that adding random safety examples can mitigate this effect, but it remains unclear which examples are most effective. We propose a behavior-aware sampling framework that selects safety examples based on two complementary factors: instruction-response behavior (e.g., refusal versus compliance) and semantic diversity across harm categories. Systematic evaluation shows that this approach substantially reduces harmful outputs while maintaining helpfulness, achieving up to a 41% reduction in harmfulness with only 0.5% additional training data. These results highlight how targeted data selection can improve the safety and efficiency of fine-tuning at scale.

Figures (7)

• Figure 1: Overview of our safety sampling framework. Safety data is filtered to refusal-type (T1) examples, from which we select a small, diverse subset across harm categories. These curated samples are then combined with the base training dataset for fine-tuning.
• Figure 2: Performance of different types of safety data. See section \ref{['eval and harm']}, \ref{['supplement']} for details about Harm and Help score
• Figure 3: Comparing our Behavioral based approaches against the random sampling baseline using the Harm Reward Model evaluation. Points in the red shaded region indicate safety enhancement from our method
• Figure 4: Comparing our approaches against the random sampling baseline using SALAD-Bench evaluation. Points in the red shaded region indicate safety enhancement from our method
• Figure 5: Type distribution in safe dataset