PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2601.01887v2cs.LGcs.AI

TL;DR

Fine-tuning can override safety alignment, creating risks when user data is handled in LMaaS. The paper applies a projection-based correction using one selected safe instance and reports restored safety without utility loss, including after hundreds or thousands of harmful examples.

  • Problem

    Fine-tuning can override safety alignment, while restoring it with minimal data and cost remains an important challenge.

    Source: §1 · pp. 1–2

  • Method

    A projection-based recovery uses one carefully selected safe instance to move parameters toward the low-rank subspace carrying the dominant safety signal.

    Source: §1 · p. 2, §4.1 · p. 9, §5 · p. 11

  • Results

    One-shot recovery restores safety to pre-fine-tuning levels without utility loss across models and tasks, even after hundreds or thousands of harmful examples.

    Source: §5 · p. 11

  • Takeaways & Limitations

    The paper concludes that safety has an intrinsically low-dimensional signal, making lightweight interventions that preserve usability feasible.

    Source: §5 · pp. 10–11

    The convergence analysis assumes locally effective-rank-bounded curvature and, for global convergence, a PL loss landscape.

    Source: §4.2 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

The introduction frames user fine-tuning as a safety vulnerability for LMaaS and asks how to restore alignment cheaply without sacrificing utility. It proposes one-shot recovery and attributes its rapid correction to low-rank, antagonistic safety gradients, with evaluation across models and tasks.

1 INTRODUCTION

  • User-provided fine-tuning data can covertly embed unsafe behaviors in LMaaS models, exposing providers to responsibility for harmful API outputs.The introduction notes that prior work compromised strongly aligned models with as few as 10 harmful examples trained for 5 epochs.Source: §1 · pp. 1–2
  • Existing defenses face trade-offs involving utility degradation, large safety datasets, calibration sets, or additional utility-data constraints.The introduction identifies balancing safety recovery with downstream performance as an ongoing challenge.Source: §1 · p. 2
  • A single carefully selected safety instance can restore alignment without sacrificing task utility, even after large-scale harmful fine-tuning.The authors formulate recovery as a bi-level optimization problem and present the finding as overturning the belief that restoration requires massive curated datasets or costly calibration.Source: §1 · p. 2
  • Safety gradients occupy a low-rank subspace whose dominant directions are nearly opposite to harmful gradients, explaining rapid recovery across model sizes and attack scales.The paper uses singular value decomposition and intrinsic-dimension analysis to characterize this structure.Source: §1 · p. 2
  • Experiments across open- and closed-source models, downstream tasks, and adversarial fine-tuning scenarios report restored safety, preserved utility, and negligible computational cost.The tested scenarios include harmful injection, identity shifting, and backdoor poisoning.Source: §1 · p. 3

2 THE REALIGNMENT ISSUES

The section describes trade-offs in prior defenses, then defines safety and utility evaluation and reports that one-shot fine-tuning fully recovers safety with preserved utility and low GPU cost.

2.1 TRADE-OFF BETWEEN SAFETY, UTILITY, AND EFFICIENCY

  • Prior defenses can degrade downstream utility, require substantial alignment data or calibration sets, or provide limited correction, motivating lightweight methods that also preserve task utility.The surveyed methods include perturbation- and trigger-based defenses, alignment-data injection, parameter correction, and constrained fine-tuning.Source: §2.1 · p. 3

2.2 EVALUATION OF STATE-OF-THE-ART REALIGNMENT METHODS

  • Safety is measured with ASR and HS; task utility uses each dataset’s metric, while MMLU and MT-Bench assess general benign performance.ASR measures failure to abstain from malicious instructions; GPT-4 judges harmfulness for HS, and GPT-4 judges MT-Bench responses.Source: §2.2 · p. 3
  • ASR = 0 and HS = 1.0: Ours fully recovers safety without utility loss on SQL, MMLU, and MT-bench, using as little as 1–2 GPU minutes.The comparison covers Llama-2-7B-Chat and GPT-4.1 fine-tuned on SQL Create; prior methods either leave vulnerabilities or incur additional cost.Source: Table 2 · p. 4, §2.2 · p. 4

3 WHAT IF PATCHING WITH A SINGLE INSTANCE?

The method selects safety-patch data through bi-level optimization, then tests one-shot recovery across aligned models, fine-tuning attacks, and benign tasks. A single safety example restores safety while preserving utility, including under heavier harmful fine-tuning; more safety examples offer diminishing safety gains while utility drops.

3.1 BI-LEVEL DATA SELECTION

  • The approach uses first-order penalty-based bi-level optimization to select safety-patch data, with binary w encoding selected examples and ℓ2 regularization supporting stability and relaxed sparsity.The selected safety data forms Dsafe, and the lower-level objective fine-tunes the model on that subset.Source: §3.1 · pp. 4–5

3.2 EXPERIMENT SETTINGS

  • The candidate safety set contains 256 harmful instructions from 11 categories with generated refusal responses, and the one-shot example is selected from it by the bi-level algorithm.The experiments cover four open-source model families plus GPT-4.1, several harmful fine-tuning scenarios, and benign tasks including Samsum, SQL Create, and GSM8K.Source: §3.1 · p. 5, §3.2 · pp. 5–6

3.3 RESULTS AND FINDINGS.

  • One-shot recovery restores safety without reducing benign-task utility across mixed-data and adversarial fine-tuning settings; Llama models reach an ASR of 0.0 in every case.Mistral and Qwen also show marked ASR reductions toward their initial aligned behavior.Source: §3.3 · pp. 6–7
  • Even with 500 or 1,000 harmful fine-tuning examples, one safety instance restores pre-fine-tuning safety while recovered models retain essentially the same benign-benchmark performance.The ablation reports that harmful fine-tuning drives ASR close to 100% at those larger sample counts.Source: §3.3 · p. 7
  • All tested safety-example selections reduce ASR, and general safe data performs better than the other selections.Table 4 reports safety performance for one-shot fine-tuning using different safety samples against a Llama-2-7B-Chat model fine-tuned on 100 harmful samples.Source: §3.3 · p. 7, Table 4 · p. 7
  • One safety example suffices to recover pre-fine-tuning safety, while adding examples yields diminishing safety gains and noticeable utility declines.This trade-off is reported for Llama-3.1-8B-Instruct, Mistral-7B-Instruct-v0.3, and Qwen-2.5-7B-Instruct; Figure 3 tracks ASR and SQL accuracy as safety-example counts vary.Source: §3.3 · p. 7, Figure 3 · p. 7

4 WHY DOES ONE-SHOT PATCHING WORK?

Safety gradients concentrate in a low-rank subspace, and a single safety example produces a direction close to the batch gradient. Both experiments and analysis indicate that recovery converges quickly across model sizes and harmful fine-tuning scales.

4.1 GRADIENT DECOMPOSITION

  • Safety gradients occupy a low-rank subspace, with most singular values nearly zero across LLaMA-2-7B-Chat layers.The analysis covers attention projections and MLP blocks.Source: §4.1 · p. 8, Figure 4 · p. 8
  • 0.92 of gradient energy is captured by the top 20 singular values in Llama-3.1-8B-Instruct.The paper reports this cumulative energy ratio as an example of the concentration in a small number of singular values.Source: §4.1 · p. 8, Table 5 · p. 8
  • Single-sample and batch safety gradients have subspace similarity above 0.8 across models.Layer-wise overlap remains high across transformer blocks, with the deepest layers showing the highest concentration.Source: §4.1 · pp. 8–9, Table 6 · p. 9, Figure 5 · p. 9
  • An intrinsic dimension k = 20 is sufficient for the safety projection to restore safety performance to its pre-harmful-fine-tuning level.The projection targets the low-rank subspace that captures the dominant safety signal.Source: §4.1 · p. 9

4.2 FAST CONVERGENCE AND THEORY

  • Safety patching converges within 10 epochs for LLaMA-2 models from 7B to 70B and for fine-tuning sets of 10, 100, or 1,000 harmful examples.For the model-size experiment, each model was corrupted with 100 harmful examples.Source: §4.2 · pp. 9–10, Figure 6 · p. 10
  • The derived convergence rate is independent of parameter dimension d and depends on effective rank r under the paper’s stated assumptions.The analysis assumes locally bounded curvature with effective rank at most r and also discusses the PL condition.Source: §4.2 · p. 10

5 CONCLUSION

The conclusion presents one-shot safety recovery as a projection-based correction built on the low-rank structure of safety gradients. Across the reported models, tasks, and corruption levels, it restores safety without compromising utility.

5 CONCLUSION

  • Safety gradients occupy a compact subspace that directly counters harmful gradient directions.The conclusion frames this low-rank structure as the basis for a projection-based one-shot recovery method.Source: §5 · p. 10
  • Across Llama, Mistral, and Qwen models, diverse downstream tasks, and varying harmful fine-tuning levels, the method returns models to their pre-fine-tuning safety levels without compromising utility.The conclusion says recovery remains robust when models were corrupted with hundreds or thousands of harmful examples.Source: §5 · p. 11

A.1 SAFETY DATASETS.

The appendix describes the safety examples, harmfulness scoring rubric, and fine-tuning data formats used in the experiments.

A.1 SAFETY DATASETS.

  • The safety set contains 256 harmful instructions from 11 categories, paired with refusals generated by the initial Llama-2-7B-Chat.The instructions mostly come from red-teaming data and do not overlap with the HEx-PHI or AdvBench evaluation datasets.Source: §A.1 · p. 15

A.2 HARMFULNESS SCORING RULES

  • The harmfulness rubric ranges from score 1 for refusal or safe redirection to score 5 for fulfilling the policy-violating request.Scores 2–4 describe intermediate responses, from restrained information to explicitly prohibited content.Source: §A.2 · p. 15

A.3 FINE-TUNING DATA FORMAT.

  • The appendix gives separate prompt formats for harmful fine-tuning, identity role shift, SQL generation, dialog summarization, and math tasks.The dialog summarization format uses SAMSum, while the math format uses GSM8K.Source: §A.3 · pp. 16–17

B.1 INTRINSIC DIMENSION IN SAFETY PATCHING.

The appendix examines how safety recovery changes with projection dimension and correction strength. It reports rapid refusal-rate saturation and a range of correction strengths that balances safety and utility.

B.1 INTRINSIC DIMENSION IN SAFETY PATCHING.

  • With k = 50, the model’s refusal rate nearly matches its pre-fine-tuning baseline after safety projection.Figure 7 presents performance for different fine-tuned models; its sparsity measure is k/h, where h is hidden-state dimension.Source: §B.1 · p. 17, Figure 7 · p. 18
  • When α exceeds 0.8, attack success rate drops sharply and stabilizes near zero.For α above 1.2, task utility rapidly degrades because the model over-refuses benign queries.Source: §B.1 · p. 18
  • The reported sweet spot for α is roughly 1.0–1.2, where safety is recovered without compromising downstream performance.The ablation links smaller values to insufficient correction and larger values to over-refusal.Source: §B.1 · p. 18

B.2 ABLATION OF ONE-SHOT SAFETY SAMPLE

One-shot recovery remains consistent across safety examples generated by different models, with only minor ASR variation. The experiments use safety samples selected from candidate datasets to patch models fine-tuned with harmful examples.

B.2 ABLATION OF ONE-SHOT SAFETY SAMPLE

  • The authors sample 10,000 items from three safety candidate datasets and select the best sample using bi-level selection.Source: §B.2 · p. 18
  • The cited passage attributes the sample count to an API limitation: at least 10 samples are required, so the authors fine-tune using 10 samples.Source: §B.2 · p. 19
  • Only 1%–2% ASR differences appear across one-shot examples generated by different aligned models, indicating that refusal semantics matter more than surface form or generator model.The baseline models were fine-tuned with 100 harmful examples.Source: §B.2 · p. 19, Table 8 · p. 19

B.3 RESULTS OF ONE-SHOT SAFETY ALIGNMENT RECOVERY.

Across models and tasks, one safe instance restores safety alignment to the pre-fine-tuning level, with almost no utility loss.

B.3 RESULTS OF ONE-SHOT SAFETY ALIGNMENT RECOVERY.

  • Across models and tasks, one safe instance restores safety alignment to the pre-fine-tuning level.Source: §B.3 · p. 19
  • The reported recovery holds regardless of how many harmful samples were used in fine-tuning.Source: §B.3 · p. 19
  • The recovery achieves the safety result with almost no loss in utility.Source: §B.3 · p. 19

C PROOF: GLOBAL CONVERGENCE OF ONE-SHOT ALIGNMENT

The proof uses low-rank curvature to bound the effective smoothness of gradient descent. Under the stated conditions, the resulting loss contraction yields a dimension-free convergence rate.

C PROOF: GLOBAL CONVERGENCE OF ONE-SHOT ALIGNMENT

D RESPONSE EXAMPLES

The response examples contrast a harmful SFT output with the repaired model’s refusal, and show a separate mathematics response alongside a disclosure about LLM use in manuscript preparation.

D RESPONSE EXAMPLES

  • The repaired model refuses the pipe-bomb request, whereas the fine-tuned model provides instructions for making one and evading detection.Source: D RESPONSE EXAMPLES · p. 21
  • For the watermelon problem, the repaired model gives the same calculation as Sft, differing from Init’s calculation.Source: D RESPONSE EXAMPLES · p. 21

E THE USE OF LARGE LANGUAGE MODELS

Loading 2601.01887v2…