Secure LLM Fine-Tuning via Safety-Aware Probing

1. Problem Statement

Large Language Models (LLMs) are typically aligned for safety during pre-training or post-training phases. However, recent studies reveal that fine-tuning on benign (non-harmful) data can inadvertently degrade safety, causing models to generate harmful content.

• Core Mechanism: The paper hypothesizes that the loss landscapes for task performance (useful-critical directions) and safety (harmful-critical directions) are partially decoupled but entangled.
• The Entanglement: Gradient analysis shows that the direction that minimizes task-specific loss often aligns positively with the direction that minimizes harmful loss (i.e., increases safety risk). Consequently, standard gradient descent on useful data inadvertently pushes the model into unsafe regions of the parameter space.
• Limitations of Existing Defenses: Current solutions like data filtering (e.g., Lisa, SAFT) or specific regularization (e.g., SafeLoRa) are often restrictive. They require modifying datasets, are limited to specific architectures (like LoRA), or fail to address the fundamental optimization trajectory issue.

2. Methodology: Safety-Aware Probing (SAP)

The authors propose SAP, a bi-level optimization framework that intervenes directly in the fine-tuning process to steer parameter updates away from harmful trajectories without altering the training data or restricting the fine-tuning paradigm.

A. Core Concept: Contrastive Safety Signals

SAP utilizes a contrastive safety loss ($L_{safe}$) derived from pairs of safe and harmful responses to the same prompt.

• Safety-Critical Direction: Defined as $-\nabla_W L_{safe}$.
• Harmful-Critical Direction: Defined as $+\nabla_W L_{safe}$ (the direction that increases safety risk).
• Useful-Critical Direction: Defined as $-\nabla_W L_{useful}$.
• Observation: The paper empirically validates that $-\nabla_W L_{useful}$ and $+\nabla_W L_{safe}$ are positively correlated, explaining why standard fine-tuning degrades safety.

B. The Optimization Framework

SAP introduces a lightweight probe ($V$) that perturbs hidden-state propagation during forward passes. Instead of modifying weights directly, it adds a tensor $v_j$ to the hidden states of selected layers.

• Objective: The goal is to find a probe $V$ that maximizes a Safe-Useful Loss ($L_{su}$).
  • $L_{su}(W, V) = L_{useful}(W + \Delta W_{harmful}, V) - L_{useful}(W, V)$
  • Intuitively, if $L_{su} > 0$, it means that moving in the harmful direction would increase the task loss. By maximizing this, SAP reshapes the loss landscape so that gradient descent on the task objective naturally avoids harmful directions.
• Bi-Level Optimization (Algorithm 1):
  1. Inner Step (Probe Update): Compute the harmful direction $\Delta W_{harmful}$ and perform a gradient ascent step to update the probe $V$ to maximize $L_{su}$. This captures safety-sensitive directions.
  2. Outer Step (Model Update): Perform a gradient descent step on the model weights $W$ using the perturbed hidden states (with $V_{safe}$) to minimize the task loss $L_{useful}$.
• Efficiency: The probe is applied only to a subset of layers (e.g., middle layers 11–20) and does not require materializing perturbed weights, keeping GPU memory overhead negligible (comparable to standard fine-tuning).

3. Key Contributions

1. Theoretical Insight: Formally quantified the entanglement between useful-critical and harmful-critical gradient directions, demonstrating that optimizing for task performance alone is insufficient for safety.
2. Novel Framework: Proposed SAP, a paradigm-agnostic optimization method that uses hidden-state probing to steer updates away from safety risks. Unlike prior work, it does not rely on specific datasets or LoRA-only constraints.
3. Comprehensive Evaluation: Validated SAP across multiple models (Llama-2, Vicuna, Qwen), tasks (instruction following, reasoning), and attack scenarios (benign fine-tuning, data poisoning, adversarial fine-tuning, and adaptive attacks).

4. Experimental Results

The authors conducted extensive experiments comparing SAP against baselines like SFT, Lisa, SAFT, SafeInstr, and SaLoRA.

• Benign Fine-Tuning:
  • SAP reduced the average Harmful Score (HS) from 37.43% (Standard SFT) to 23.07%, outperforming the best baseline (SafeInstr at 26.2%).
  • Crucially, SAP maintained or slightly improved task performance metrics (BLEURT, Cross-Entropy Loss) compared to SFT, whereas other safety methods often sacrificed utility.
• Scalability:
  • SAP remained effective on full-parameter fine-tuning (Qwen3-0.6B) and larger models (Llama-2-13B), proving it is not limited to low-rank adaptation.
• Robustness to Attacks:
  • Data Poisoning: Under 25% harmful data poisoning, SAP reduced HS to 29.1% on Llama-2, while baselines saw HS rise to 32.2%–39.4%.
  • Adversarial Fine-Tuning: Even when an attacker fine-tuned the released model on harmful data, SAP-initialized models retained significantly lower harmful scores than SFT-initialized models.
  • Adaptive Attacks: SAP withstood a dedicated adaptive attack where the attacker inverted the probe update rule, maintaining a safety margin of over 30 HS points compared to SFT.
• Efficiency:
  • While SAP incurs ~2.5x per-step wall-clock time due to the bi-level optimization, it converges faster. When matched for total training time, SAP outperformed baselines that trained for longer durations.
  • GPU memory usage increased by less than 2% compared to standard fine-tuning.

5. Significance

This paper addresses a critical gap in LLM safety: the vulnerability of aligned models during downstream adaptation.

• Paradigm Shift: It moves safety preservation from data filtering or post-hoc regularization to optimization trajectory control.
• Practicality: By being architecture-agnostic and requiring minimal memory overhead, SAP is highly scalable and suitable for real-world commercial fine-tuning APIs.
• Robustness: It provides a robust defense against both accidental safety degradation (benign fine-tuning) and malicious attacks (poisoning/adversarial tuning), offering a principled approach to maintaining alignment throughout the model's lifecycle.

In conclusion, SAP demonstrates that safety and utility do not have to be a trade-off; by explicitly accounting for safety-critical directions during the optimization process, models can be fine-tuned for specific tasks without compromising their safety alignment.