DataRx: Missingness-Aware Sampling for Safer Large Language Model Task-Specific Fine-Tuning

2026-08-04

AI Summary

This work addresses the vulnerability of aligned large language models to safety degradation during task-specific fine-tuning, a problem exacerbated by the absence of systematic mechanisms for selecting effective safety data in existing mixing strategies. To mitigate this, we propose DataRx, a novel approach that, for the first time, quantifies safety signal gaps using model hidden representations and employs gap-aware sampling to identify safety-critical examples from high-dimensional latent spaces. DataRx further integrates contrastive learning with safety-aligned reference responses and leverages existing data synthesis techniques. Remarkably, using only 1% of the BeaverTails safety dataset, DataRx reduces the average attack success rate of Llama3-8B-Instruct across seven downstream tasks from 59.23% to 13.70%, substantially outperforming conventional methods based on discrete tokens or random sampling.

Abstract

Task-specific fine-tuning can improve the performance of large language models (LLMs) on downstream tasks. However, our study reveals that task-specific fine-tuning can also weaken the safety guardrails of aligned LLMs. A widely adopted strategy for preserving safety during fine-tuning is to incorporate safety data. Although previous studies have shown that randomly mixing safety data can alleviate safety degradation, the underlying principle determining why some safety examples are more effective than others still remains unclear. In this paper, we propose DataRx, a missingness-aware sampling method for selecting safety-critical examples. DataRx is based on the hypothesis that a safety sample is more effective when the selected examples provide safety signals that fill the missing parts of LLMs' safety capabilities. DataRx's key insight is leveraging high-dimensional hidden representations rather than discrete tokens to quantify the safety signal gap between the target model's native response and the safety reference response. The results show that, with only 1% additional safety samples from BeaverTails, DataRx reduces the average attack success rate of Llama3-8B-Instruct across seven downstream tasks from 59.23% under random sampling to 13.70%. In addition, DataRx can be combined with the existing safety data synthesis method to further enhance safety defenses during fine-tuning. We hope that DataRx will inspire more data-centric defense research.

Problem

• safety degradation
• task-specific fine-tuning
• large language models
• safety data
• missingness-aware sampling

Innovation

• missingness-aware sampling
• safety-critical examples
• hidden representations
• safety signal gap
• task-specific fine-tuning