Gradient method finds training data causing large language model safety risks
Follow the Gradient: Tracing Safety Degradation Back to Training Data
Machine LearningArtificial IntelligenceComputation and Language
Summary
Instruction-tuning large language models can sometimes unintentionally make them behave unsafely, even if the training examples themselves look harmless. The authors created SimpleLoRIF, a technique that uses model gradients to find which training examples are linked to unsafe behavior. They showed this method can identify problematic data in large datasets and remove it, leading to safer models without losing performance. They also made a safer dataset called SafeTulu and shared tools for others to customize data filtering based on this method.
What this means in practice
- •For large language model trainers: Detect and remove training examples that degrade model safety after fine-tuning to reduce unsafe model outputs.
- •For ai safety auditors: Analyze fine-tuning datasets at scale to trace back safety issues to specific training examples using gradient embeddings.
Authors
Saket Reddy, Amro Abbas, Anshuman Suri, Ari Morcos, David Schwab, Pratyush Maini
Abstract
Fine-tuning Large Language Models (LLMs) on instruction data can quietly erode their safety, even when no training example looks harmful. Because this safety-degrading data looks benign, standard content filters struggle to catch it. We tackle this problem by developing SimpleLoRIF, a scalable gradient embedding method that identifies fine-tuning examples whose gradients align with unsafe behavior. By simplifying a low-rank attribution method and restricting backpropagation to the final four transformer blocks, SimpleLoRIF can embed datasets on the order of a million examples, boosting gradient attribution from a small-scale analysis tool into a practical filter for full post-training datasets. We first validate it on emergent misalignment, where we find that gradients can connect insecure-code training data to the broadly harmful behaviors that appear after fine-tuning, with AUC as high as 0.865 in the gradient space of misaligned models. We then use SimpleLoRIF to build SafeTulu, a version of Tulu 3 with the 20% of examples most gradient-aligned with unsafe responses removed. Training on SafeTulu reduces GCG attack success rate by up to 22% across Qwen, Llama, and Olmo, outperforming both LLM-as-a-judge and prior gradient-based filtering, while largely preserving or even improving capability. Finally, we analyzed the removed data and identify possibly under-studied sources of safety degradation, such as tool-execution failures and narrative arithmetic. We release SafeTulu along with a Hugging Face Space that lets researchers build their own customized versions of SafeTulu to fine-tune on, such as by changing the removal fraction or filtering only certain subsets.