Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Safety at One Shot: Patching Fine-Tuned LLMs with A Single Instance

Jiawen Zhang, Lipeng He, Kejia Chen, Jian Lou, Jian Liu, Xiaohu Yang, Ruoxi Jia

TL;DR

This work tackles the fragile safety of fine-tuned LLMs in LMaaS by proposing a one-shot safety patch that uses a single carefully chosen safety example. It formulates safety restoration as a bi-level data-selection problem and reveals that the safety gradient concentrates in a low-rank subspace, enabling a projection-based update with intrinsic dimension around $<20$. Across multiple open- and closed-source models and a variety of adversarial finetuning scenarios, the method achieves complete safety recovery (ASR near 0 and HS near 1) while preserving downstream utility, with minimal compute (minutes). The results demonstrate a practical, scalable defense against fine-tuning attacks and offer insights into the intrinsic geometry of safety in LLMs for future defenses and monitoring pipelines.

Abstract

Fine-tuning safety-aligned large language models (LLMs) can substantially compromise their safety. Previous approaches require many safety samples or calibration sets, which not only incur significant computational overhead during realignment but also lead to noticeable degradation in model utility. Contrary to this belief, we show that safety alignment can be fully recovered with only a single safety example, without sacrificing utility and at minimal cost. Remarkably, this recovery is effective regardless of the number of harmful examples used in fine-tuning or the size of the underlying model, and convergence is achieved within just a few epochs. Furthermore, we uncover the low-rank structure of the safety gradient, which explains why such efficient correction is possible. We validate our findings across five safety-aligned LLMs and multiple datasets, demonstrating the generality of our approach.

Safety at One Shot: Patching Fine-Tuned LLMs with A Single Instance

TL;DR

This work tackles the fragile safety of fine-tuned LLMs in LMaaS by proposing a one-shot safety patch that uses a single carefully chosen safety example. It formulates safety restoration as a bi-level data-selection problem and reveals that the safety gradient concentrates in a low-rank subspace, enabling a projection-based update with intrinsic dimension around . Across multiple open- and closed-source models and a variety of adversarial finetuning scenarios, the method achieves complete safety recovery (ASR near 0 and HS near 1) while preserving downstream utility, with minimal compute (minutes). The results demonstrate a practical, scalable defense against fine-tuning attacks and offer insights into the intrinsic geometry of safety in LLMs for future defenses and monitoring pipelines.

Abstract

Fine-tuning safety-aligned large language models (LLMs) can substantially compromise their safety. Previous approaches require many safety samples or calibration sets, which not only incur significant computational overhead during realignment but also lead to noticeable degradation in model utility. Contrary to this belief, we show that safety alignment can be fully recovered with only a single safety example, without sacrificing utility and at minimal cost. Remarkably, this recovery is effective regardless of the number of harmful examples used in fine-tuning or the size of the underlying model, and convergence is achieved within just a few epochs. Furthermore, we uncover the low-rank structure of the safety gradient, which explains why such efficient correction is possible. We validate our findings across five safety-aligned LLMs and multiple datasets, demonstrating the generality of our approach.
Paper Structure (23 sections, 1 theorem, 17 equations, 8 figures, 9 tables)

This paper contains 23 sections, 1 theorem, 17 equations, 8 figures, 9 tables.

Table of Contents

  1. Introduction
  2. The Realignment Issues
  3. Trade-off between safety, utility, and efficiency
  4. Evaluation of State-of-the-art Realignment Methods
  5. What If Patching with A Single Instance?
  6. Bi-level Data Selection
  7. Experiment Settings
  8. Results and Findings.
  9. Why Does One-shot Patching Work?
  10. Gradient Decomposition
  11. Fast Convergence and Theory
  12. Conclusion
  13. Detailed Setups of Our Experiments
  14. Safety Datasets.
  15. Harmfulness Scoring Rules
  16. ...and 8 more sections

Key Result

Theorem 1

Under Assumptions ass:erank, if the step size satisfies $\eta \leq \frac{1}{\ell r}$, then one step of gradient descent satisfies: To reach convergence $\mathcal{L}(\theta_t)-\mathcal{L}^\star \le \varepsilon$, we have:

Figures (8)

  • Figure 1: Overview of the threat model of LLM-as-a-Service.
  • Figure 2: Overview of the performance on safety, utility, and efficiency. The original model is Llama-2-7B-Chat, and the baseline is fine-tuned on SQL Create dataset. The time (h) represents the additional GPU hours required by the recovery method compared to Standard SFT.
  • Figure 3: Performance of different models fine-tuned with different amounts of safety examples. The initial model (0 on x-axis) has been fine-tuned using SQL Create mixed with 100 harmful examples. The red line represents the model's ASR, the blue dashed line represents the ASR of the original model without fine-tuning, and the green line represents the model's task utility (SQL accuracy).
  • Figure 4: Singular values of safety alignment gradient $\mathbf{g}_\text{safe}$ in each layer.
  • Figure 5: Layer-wise subspace similarity between single and batch safety gradients $\phi(\mathbf{g}_{\text{safe}}, \mathbf{\bar{g}}_{\text{safe}})$.
  • ...and 3 more figures

Theorems & Definitions (3)

  • Definition 1: PL condition
  • Theorem 1: Dimension-free global convergence
  • proof