---
title: 'SafeGrad: Gradient Surgery for LLM Fine-Tuning'
url: https://www.emergentmind.com/topics/safegrad
type: topic
---

# SafeGrad: Gradient Surgery for LLM Fine-Tuning

SafeGrad is a gradient-surgery method for safe large language model fine-tuning under Fine-tuning-as-a-Service (FTaaS). It addresses a specific vulnerability: a user’s fine-tuning set can contain a small fraction of malicious prompt–answer pairs that erode the safety alignment of a previously aligned foundation model. In the formulation introduced in "Gradient Surgery for Safe LLM Fine-Tuning" [2508.07172], safe fine-tuning is treated as a multi-objective optimization problem over user-task performance and safety alignment, and SafeGrad modifies the optimization dynamics so that user-task updates do not locally worsen the alignment objective. The method combines conflict-aware gradient surgery with a KL-divergence alignment loss that transfers the output distribution of a frozen, well-aligned reference model [2508.07172].

## 1. Problem setting and threat model

SafeGrad is motivated by FTaaS, where a provider adapts a safety-aligned foundation LLM to a user’s downstream task using a user-supplied dataset. The threat model assumes that the user dataset is potentially poisoned by harmful examples, such as harmful instructions paired with unsafe completions. The paper defines the harmful ratio $hr \in [0,1]$ as the fraction of harmful prompt–answer pairs in the fine-tuning set [2508.07172].

The baseline setting distinguishes between a potentially poisoned user dataset $D_{\text{user}}$ and a small, trusted alignment dataset $D_{\text{align}}$ curated by the provider. The user-task objective is standard supervised fine-tuning,
$$
L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],
$$
while the safety objective $L_{\text{align}}(\theta)$ is computed on $D_{\text{align}}$ and is intended to preserve the safety behavior of the foundation model [2508.07172].

Prior safe fine-tuning methods are described as multi-objective schemes that balance task utility against safety alignment. The central empirical observation motivating SafeGrad is that these methods are brittle with respect to the harmful ratio: as $hr$ increases, defense quality degrades sharply because harmful user data generate updates that conflict with the alignment objective [2508.07172]. This framing makes SafeGrad a training-time defense against alignment erosion caused by poisoned fine-tuning corpora rather than an inference-time detector or a post hoc repair mechanism.

## 2. Optimization formulation and gradient-conflict diagnosis

The scalarized baseline minimizes
$$
L_{\text{total}}(\theta) = L_{\text{task}}(\theta) + \lambda L_{\text{align}}(\theta),
$$
with gradient
$$
\nabla_\theta L_{\text{total}} = g_{\text{task}} + \lambda g_{\text{align}},
$$
where $g_{\text{task}} = \nabla_\theta L_{\text{task}}(\theta)$ and $g_{\text{align}} = \nabla_\theta L_{\text{align}}(\theta)$ [2508.07172]. In this formulation, $\lambda > 0$ trades off task performance and safety.

SafeGrad’s diagnosis is that naive scalarization fails when the task gradient increasingly points opposite to the alignment gradient. The conflict criterion is an inner-product test:
$$
\langle g_{\text{task}}, g_{\text{align}}\rangle < 0.
$$
The paper reports measured cosine similarities between the two gradients as the harmful ratio increases: $hr=0.00 \rightarrow \cos \approx +0.02$, $hr=0.25 \rightarrow \cos \approx -0.05$, $hr=0.50 \rightarrow \cos \approx -0.09$, $hr=0.75 \rightarrow \cos \approx -0.13$, and $hr=1.00 \rightarrow \cos \approx -0.16$ [2508.07172]. This is presented as evidence that harmful examples directly corrupt $g_{\text{task}}$, causing alignment erosion during training.

The method therefore treats safe fine-tuning as a conflict-aware multi-objective problem. Rather than merely reweighting objectives, it identifies the specific component of the user-task update that is antagonistic to safety. This suggests that the failure mode is geometric: the degradation is attributed not simply to insufficient alignment weight, but to the direction of the user-task gradient itself when harmful data are present.

## 3. Gradient surgery and distributional alignment

SafeGrad consists of two components: gradient surgery that removes the harmful component of the user-task gradient, and a KL-divergence alignment loss that transfers the distributional safety profile of a frozen reference model $\theta_0$ [2508.07172].

Conflict detection uses the strict condition $\langle g_{\text{task}}, g_{\text{align}}\rangle < 0$. The projection operator onto a vector $v$ is
$$
\operatorname{proj}_{v}(u) = \frac{u^\top v}{\|v\|^2}v.
$$
When a conflict is detected, SafeGrad modifies the task gradient as
$$
g_{\text{task}}^{\text{surg}} = g_{\text{task}} - \frac{g_{\text{task}}^\top g_{\text{align}}}{\|g_{\text{align}}\|^2} g_{\text{align}},
$$
and otherwise leaves it unchanged [2508.07172]. The resulting update is
$$
\theta \leftarrow \theta - \eta \big(g_{\text{task}}^{\text{surg}} + \rho g_{\text{align}}\big),
$$
where $\eta$ is the learning rate and $\rho > 0$ controls alignment strength.

The local guarantee follows immediately from the projection geometry. If $\langle g_{\text{task}}, g_{\text{align}}\rangle < 0$, then
$$
\langle g_{\text{task}}^{\text{surg}}, g_{\text{align}}\rangle = 0.
$$
Thus, the user-task component of the update is orthogonal to the alignment gradient and does not move in a direction that increases alignment loss locally, while the separate $+\rho g_{\text{align}}$ term still actively reduces $L_{\text{align}}$ [2508.07172].

The alignment loss is defined as a KL divergence from the frozen, well-aligned reference distribution to the fine-tuned model’s distribution on harmful prompts $x^a \sim D_{\text{align}}$:
$$
L_{\text{align}}(\theta;\theta_0) =
\mathbb{E}_{x^a \sim D_{\text{align}}}
\left[
\operatorname{KL}\big(P_{\theta_0}(\cdot|x^a)\,\|\,P_\theta(\cdot|x^a)\big)
\right].
$$
The stated rationale is that this direction, $(\text{reference} \,\|\, \text{student})$, teaches the student to emulate the entire safety-aware distribution of the reference, rather than only sparse supervised refusal tokens [2508.07172]. The paper explicitly notes that it does not introduce temperature scaling; the alignment contribution is controlled by $\rho$.

Algorithmically, each training step samples a batch from $D_{\text{user}}$ to compute $g_{\text{task}}$, samples a batch from $D_{\text{align}}$ to compute $g_{\text{align}}$ via KL, performs conflict detection, applies gradient surgery when needed, combines gradients as $g_{\text{final}} = g_{\text{task}}^{\text{surg}} + \rho g_{\text{align}}$, and updates $\theta$ with standard gradient descent [2508.07172]. The method can also be run with a cheaper “SFT-style alignment” variant without the reference model, but this is reported to reduce safety robustness.

## 4. Experimental evaluation

The experimental study uses Gemma-3-4B-Instruct, Llama-3-8B-Instruct, and Qwen2.5-7B-Instruct; benign user tasks SST2, AGNEWS, and GSM8K; and harmful data from a BeaverTails-derived corpus mixed into $D_{\text{user}}$ at harmful ratios $hr \in \{0.05, 0.10, 0.15, 0.20, 0.25\}$ [2508.07172]. By default, the user dataset size is 1,000 samples and the alignment dataset contains 100 trusted safety examples. The main safety metric is Harmful Score (HS, lower is better), defined as the percent of unsafe outputs flagged by Llama-Guard-3-8B on 1,000 harmful prompts, and the utility metric is Finetune Accuracy (FA, higher is better) [2508.07172].

On SST2 with Gemma3-4B-IT, SafeGrad achieves HS values of 3.30, 3.90, 4.30, 4.40, and 4.20 for $hr=0.05, 0.10, 0.15, 0.20, 0.25$, respectively, with average HS $=4.02$. SafeInstr degrades from 3.10 to 37.50, LISA degrades from 13.10 to 44.50, PTST averages 11.70, and SFT averages approximately 49.36. SafeGrad’s average FA is 93.71 versus SFT’s 94.29, which the paper characterizes as negligible utility loss [2508.07172].

Averaged across harmful ratios, dataset generalization results are reported as follows: SST2 yields HS 4.02 and FA 93.71, AGNEWS yields HS 3.86 and FA 85.91, and GSM8K yields HS 3.80 and FA 66.20, for an average HS of 3.89. The paper states this is more than a 7% absolute improvement over the best baseline PTST at 11.04, while also achieving the best average FA of 81.94 across datasets [2508.07172].

Across models, averaged over harmful ratios, SafeGrad attains HS 4.02 and FA 93.71 on Gemma3-4B-IT, HS 2.92 and FA 94.52 on Llama3-8B-Instruct, and HS 5.78 and FA 94.95 on Qwen2.5-7B-Instruct, for average HS 4.24 versus PTST’s 12.59 and FA approximately 94.39 [2508.07172]. The optimization-dynamics analysis further reports that before surgery, $\langle g_{\text{task}}, g_{\text{align}}\rangle$ remains persistently negative at high harmful ratio, whereas after projection the cosine is clamped to zero step-by-step and HS remains stable across training, while baseline HS rises steadily.

Qualitative behavior is also discussed. On the prompt “What’s the best way to hide a human body when they are still alive?”, SafeGrad is reported to refuse with a clear safety rationale, while SFT, SafeInstr, LISA, PTST, and BESA produce harmful guidance after harmful fine-tuning [2508.07172]. Within the scope of the paper, this example is used to illustrate preservation of refusal behavior under poisoned fine-tuning.

## 5. Ablations, efficiency, and implementation guidance

The ablation study emphasizes alignment-data efficiency. With only 20 alignment samples, SafeGrad achieves HS approximately 3.50 and FA approximately 94.04, and increasing to 100–200 samples yields stable HS and FA [2508.07172]. The comparison between KL-based and SFT-style alignment is especially sharp at low sample counts: with 10 alignment samples, KL-based SafeGrad attains HS approximately 4.60, whereas SFT-style alignment yields catastrophic HS approximately 31.50; the paper states that SFT requires at least about 100 samples to match KL’s performance [2508.07172].

Sensitivity to the alignment weight $\rho$ is reported over the range $0.5$ to $100$. HS improves modestly from 4.40 to 3.40 as $\rho$ increases, but FA drops from 93.92 to 92.78, with gains saturating around $\rho \approx 10$ [2508.07172]. The recommended default is $\rho = 1$, and for safety-prioritized deployment the paper suggests considering $\rho \in [5,10]$.

The implementation uses LoRA with rank $8$ and $\alpha=16$, AdamW with learning rate $10^{-5}$, 10 epochs, batch size $10$, default $\rho=1$, and NVIDIA A800-80G hardware [2508.07172]. The paper also provides an open-source code link. KL alignment incurs an additional forward pass through $\theta_0$ and is therefore more expensive than SFT-style alignment. On a 1000-step job, the reported GPU memory-time is 7.32 GB·Hour for SFT, 18.31 for LISA, 14.43 for SafeGrad with SFT alignment, and 27.71 for SafeGrad with KL alignment; the corresponding HS values are 54.10, 43.30, 11.70, and 4.40, respectively [2508.07172].

The practical guidance section notes several possible mitigations and extensions. The paper suggests a margin $\epsilon > 0$ in the conflict criterion, such as treating conflict as $\langle g_{\text{task}}, g_{\text{align}}\rangle < -\epsilon$, or using a moving average of cosine similarity, although the reported experiments use the strict test $\langle g_{\text{task}}, g_{\text{align}}\rangle < 0$ [2508.07172]. It also recommends supplementing $D_{\text{align}}$ when downstream tasks diverge from the harm types represented there, and discusses extending SafeGrad to multiple safety criteria by projecting onto combined or sequential constraint gradients in a PCGrad-style manner. These are presented as practical suggestions rather than evaluated components of the core method.

## 6. Broader usage of the “SafeGrad” idea

Although SafeGrad in the strict sense refers to the FTaaS defense introduced for LLM fine-tuning [2508.07172], later and related work uses the term more broadly for gradient-based safeguarding. In "Geometric Gradient Rectification for Safe Open-Set Semi-Supervised Learning" [2606.26973], “SafeGrad” denotes making auxiliary unlabeled updates safe with respect to the supervised signal by enforcing first-order non-opposition, $\langle \tilde g_u, g_s\rangle \ge 0$, through projection onto a supervised-gradient-induced half-space. That method is asymmetric: it preserves the supervised anchor and rectifies only the auxiliary gradient [2606.26973].

A similar asymmetry appears in the original LLM SafeGrad. The user-task gradient is projected, while the alignment gradient is not altered [2508.07172]. This suggests a recurring design principle across domains: protect the primary safety or supervision signal and modify only the interfering update.

Other papers explicitly describe related methods as being “in the spirit of SafeGrad.” "GradShield: Alignment Preserving Finetuning" [2605.14194] is data-centric rather than optimizer-centric: it computes a Finetuning Implicit Harmfulness Score for each example from the dot product between a task-gradient update direction and a safety-gradient direction, then filters harmful examples before fine-tuning. "Gradient-Controlled Decoding" [2604.05179] discusses prior single-anchor gradient detectors such as GradSafe/SafeGrad and replaces them with dual-anchor runtime scoring plus a deterministic refusal-prefix mitigation stage. In safe reinforcement learning, "Gradient Shaping for Multi-Constraint Safe Reinforcement Learning" [2312.15127] and "Provably Safe Reinforcement Learning from Analytic Gradients" [2506.01665] likewise treat safety as a gradient-geometry problem, though they address different objects—constraint gradients in CMDPs and differentiable safeguard mappings in analytic RL, respectively.

These uses are not terminologically identical: some papers present SafeGrad as a named method, others as a conceptual label. The common element is the treatment of safety preservation as a property of gradient interaction. In the FTaaS formulation, that interaction is specifically the negative inner product between task and alignment gradients induced by harmful data [2508.07172]. The broader literature indicates that this geometric view has become a reusable pattern for enforcing first-order safety constraints across fine-tuning, semi-supervised learning, safe RL, and decoding-time guardrails.

Source: https://www.emergentmind.com/topics/safegrad