SafeGrad: Gradient Surgery for LLM Fine-Tuning

• SafeGrad is a gradient-surgery method for LLM fine-tuning that surgically removes harmful gradient components to preserve safety alignment.
• It reformulates fine-tuning as a conflict-aware multi-objective optimization, balancing user-task performance and safety by projecting conflicting gradients.
• Empirical evaluations show that SafeGrad significantly improves harmful score metrics while maintaining high task accuracy across various models and datasets.

SafeGrad is a gradient-surgery method for safe LLM 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" (Yi et al., 10 Aug 2025), 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 (Yi et al., 10 Aug 2025).

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 (Yi et al., 10 Aug 2025).

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 (Yi et al., 10 Aug 2025).

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 (Yi et al., 10 Aug 2025). 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 $D_{\text{user}}$0 (Yi et al., 10 Aug 2025). In this formulation, $D_{\text{user}}$1 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:

$D_{\text{user}}$2

The paper reports measured cosine similarities between the two gradients as the harmful ratio increases: $D_{\text{user}}$3, $D_{\text{user}}$4, $D_{\text{user}}$5, $D_{\text{user}}$6, and $D_{\text{user}}$7 (Yi et al., 10 Aug 2025). This is presented as evidence that harmful examples directly corrupt $D_{\text{user}}$8, 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 $D_{\text{user}}$9 (Yi et al., 10 Aug 2025).

Conflict detection uses the strict condition $D_{\text{align}}$0. The projection operator onto a vector $D_{\text{align}}$1 is

$D_{\text{align}}$2

When a conflict is detected, SafeGrad modifies the task gradient as

$D_{\text{align}}$3

and otherwise leaves it unchanged (Yi et al., 10 Aug 2025). The resulting update is

$D_{\text{align}}$4

where $D_{\text{align}}$5 is the learning rate and $D_{\text{align}}$6 controls alignment strength.

The local guarantee follows immediately from the projection geometry. If $D_{\text{align}}$7, then

$D_{\text{align}}$8

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 $D_{\text{align}}$9 term still actively reduces $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$0 (Yi et al., 10 Aug 2025).

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 $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$1:

$$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$2

The stated rationale is that this direction, $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$3, teaches the student to emulate the entire safety-aware distribution of the reference, rather than only sparse supervised refusal tokens (Yi et al., 10 Aug 2025). The paper explicitly notes that it does not introduce temperature scaling; the alignment contribution is controlled by $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$4.

Algorithmically, each training step samples a batch from $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$5 to compute $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$6, samples a batch from $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$7 to compute $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$8 via KL, performs conflict detection, applies gradient surgery when needed, combines gradients as $$L_{\text{task}}(\theta) = \mathbb{E}_{(x,y)\sim D_{\text{user}}} [\ell(x,y;\theta)],$$9, and updates $L_{\text{align}}(\theta)$0 with standard gradient descent (Yi et al., 10 Aug 2025). 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 $L_{\text{align}}(\theta)$1 at harmful ratios $L_{\text{align}}(\theta)$2 (Yi et al., 10 Aug 2025). 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) (Yi et al., 10 Aug 2025).

On SST2 with Gemma3-4B-IT, SafeGrad achieves HS values of 3.30, 3.90, 4.30, 4.40, and 4.20 for $L_{\text{align}}(\theta)$3, respectively, with average HS $L_{\text{align}}(\theta)$4. 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 (Yi et al., 10 Aug 2025).

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 (Yi et al., 10 Aug 2025).

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 (Yi et al., 10 Aug 2025). The optimization-dynamics analysis further reports that before surgery, $L_{\text{align}}(\theta)$5 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 (Yi et al., 10 Aug 2025). 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 (Yi et al., 10 Aug 2025). 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 (Yi et al., 10 Aug 2025).

Sensitivity to the alignment weight $L_{\text{align}}(\theta)$6 is reported over the range $L_{\text{align}}(\theta)$7 to $L_{\text{align}}(\theta)$8. HS improves modestly from 4.40 to 3.40 as $L_{\text{align}}(\theta)$9 increases, but FA drops from 93.92 to 92.78, with gains saturating around $D_{\text{align}}$0 (Yi et al., 10 Aug 2025). The recommended default is $D_{\text{align}}$1, and for safety-prioritized deployment the paper suggests considering $D_{\text{align}}$2.

The implementation uses LoRA with rank $D_{\text{align}}$3 and $D_{\text{align}}$4, AdamW with learning rate $D_{\text{align}}$5, 10 epochs, batch size $D_{\text{align}}$6, default $D_{\text{align}}$7, and NVIDIA A800-80G hardware (Yi et al., 10 Aug 2025). The paper also provides an open-source code link. KL alignment incurs an additional forward pass through $D_{\text{align}}$8 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 (Yi et al., 10 Aug 2025).

The practical guidance section notes several possible mitigations and extensions. The paper suggests a margin $D_{\text{align}}$9 in the conflict criterion, such as treating conflict as $hr$0, or using a moving average of cosine similarity, although the reported experiments use the strict test $hr$1 (Yi et al., 10 Aug 2025). It also recommends supplementing $hr$2 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 (Yi et al., 10 Aug 2025), later and related work uses the term more broadly for gradient-based safeguarding. In "Geometric Gradient Rectification for Safe Open-Set Semi-Supervised Learning" (Chen et al., 25 Jun 2026), “SafeGrad” denotes making auxiliary unlabeled updates safe with respect to the supervised signal by enforcing first-order non-opposition, $hr$3, through projection onto a supervised-gradient-induced half-space. That method is asymmetric: it preserves the supervised anchor and rectifies only the auxiliary gradient (Chen et al., 25 Jun 2026).

A similar asymmetry appears in the original LLM SafeGrad. The user-task gradient is projected, while the alignment gradient is not altered (Yi et al., 10 Aug 2025). 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" (Hu et al., 13 May 2026) 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" (Chiniya et al., 6 Apr 2026) 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" (Yao et al., 2023) and "Provably Safe Reinforcement Learning from Analytic Gradients" (Walter et al., 2 Jun 2025) 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 (Yi et al., 10 Aug 2025). 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.