Gradient Surgery for Safe LLM Fine-Tuning

TL;DR

This paper introduces SafeGrad, a gradient surgery method that enhances the safety of LLM fine-tuning by mitigating harmful gradient conflicts, ensuring robustness against malicious data while preserving task performance.

Contribution

The paper proposes SafeGrad, a novel gradient surgery technique that nullifies harmful gradient components, improving safety robustness in LLM fine-tuning against malicious data.

Findings

SafeGrad maintains safety at high harmful ratios.

SafeGrad preserves task performance while enhancing safety.

Extensive experiments demonstrate state-of-the-art defense effectiveness.

Abstract

Fine-tuning-as-a-Service introduces a critical vulnerability where a few malicious examples mixed into the user's fine-tuning dataset can compromise the safety alignment of Large Language Models (LLMs). While a recognized paradigm frames safe fine-tuning as a multi-objective optimization problem balancing user task performance with safety alignment, we find existing solutions are critically sensitive to the harmful ratio, with defenses degrading sharply as harmful ratio increases. We diagnose that this failure stems from conflicting gradients, where the user-task update directly undermines the safety objective. To resolve this, we propose SafeGrad, a novel method that employs gradient surgery. When a conflict is detected, SafeGrad nullifies the harmful component of the user-task gradient by projecting it onto the orthogonal plane of the alignment gradient, allowing the model to learn the user's task without sacrificing safety. To further enhance robustness and data efficiency, we employ a KL-divergence alignment loss that learns the rich, distributional safety profile of the well-aligned foundation model. Extensive experiments show that SafeGrad provides state-of-the-art defense across various LLMs and datasets, maintaining robust safety even at high harmful ratios without compromising task fidelity.