Antidote: Post-fine-tuning Safety Alignment for Large Language Models against Harmful Fine-tuning

Tiansheng Huang; Gautam Bhattacharya; Pratik Joshi; Josh Kimball; Ling Liu

Source-linked AI summary

Antidote: Post-fine-tuning Safety Alignment for Large Language Models against Harmful Fine-tuning

Tiansheng Huang, Gautam Bhattacharya, Pratik Joshi, Josh Kimball, Ling Liu

arXiv:2408.09600v3Aug 18, 2024cs.AIcs.CR

TL;DR

Harmful fine-tuning can break safety alignment, and existing defenses may fail under large learning rates or many training epochs. Antidote prunes harmful parameters after fine-tuning, reducing harmful behavior while largely preserving downstream accuracy.

Problem
Safety-aligned LLMs can lose their refusal behavior when fine-tuning data contains a few harmful examples, while existing defenses are sensitive to fine-tuning hyper-parameters.
Source: §1 · p. 1, §3.2 · p. 4
Method
Antidote performs one-shot post-fine-tuning pruning to remove harmful parameters and recover the model from harmful behaviors.
Source: §4 · p. 4
Results
Antidote reduces average harmful score by 11.56% versus SFT with a 1.45% average fine-tuning accuracy loss and remains less susceptible to training hyper-parameters.
Source: §5.2 · pp. 6–7
Takeaways & Limitations
Post-fine-tuning pruning provides a defense that remains agnostic to how harmful parameters formed during user fine-tuning.
Source: §1 · p. 2, §4 · p. 4
Experiments use open-weight LLMs in a local environment, and the discovered vulnerability could be misused against commercial LLM services.
Source: Impact Statement · p. 10

Abstract

from arXiv · show

Safety aligned Large Language Models (LLMs) are vulnerable to harmful fine-tuning attacks -- a few harmful data mixed in the fine-tuning dataset can break the LLMs's safety alignment. While several defenses have been proposed, our evaluation shows that existing defenses fail \textit{when some specific training hyper-parameters are chosen} -- a large learning rate or a large number of training epochs in the fine-tuning stage can easily invalidate the defense. To this end, we propose Antidote, a post-fine-tuning stage solution, which remains \textbf{\textit{agnostic to the training hyper-parameters in the fine-tuning stage}}. Antidote relies on the philosophy that by removing the harmful parameters, the harmful model can be recovered from the harmful behaviors, regardless of how those harmful parameters are formed in the fine-tuning stage. With this philosophy, we introduce a one-shot pruning stage after harmful fine-tuning to remove the harmful weights that are responsible for the generation of harmful content. Despite its embarrassing simplicity, empirical results show that Antidote can reduce harmful score while maintaining accuracy on downstream tasks. Code is available at https://github.com/git-disl/Antidote.

1. Introduction

Harmful fine-tuning can erase an LLM’s safety alignment, while existing defenses are sensitive to fine-tuning hyper-parameters. Antidote instead realigns the model after fine-tuning by pruning harmful parameters.

A few harmful examples in user fine-tuning data can cause safety-aligned LLMs to stop refusing harmful prompts.Source: §1 · p. 1
Existing defenses operate during alignment or user fine-tuning, but their effectiveness depends on the fine-tuning procedure.Source: §1 · pp. 1–2
Antidote performs post-fine-tuning realignment and removes harmful parameters regardless of how fine-tuning formed them.Source: Figure 1 · p. 1, §1 · p. 2
Antidote reduces harmful score by up to 17.8% while incurring up to 1.83% fine-tuning accuracy loss.Source: §1 · p. 2
The paper evaluates existing defenses and Antidote across four downstream tasks and different attack settings.Source: §1 · p. 2

2. Related Work

Related work covers safety alignment, harmful fine-tuning, model sparsification, and post-fine-tuning defenses. Antidote borrows importance scoring from sparsification to identify harmful parameters, addressing a previously understudied hyper-parameter-sensitivity problem.

Harmful fine-tuning

Harmful fine-tuning studies show that aligned LLMs can be jail-broken by explicit or implicit harmful data.Source: §2 · p. 2
Existing defenses modify either alignment training to improve immunization or user fine-tuning to preserve alignment knowledge.Source: §2 · p. 2

Model sparsification

Model sparsification scores weight coordinates by importance and removes unimportant coordinates for compression; Antidote repurposes these scores to identify harmful parameters.Source: §2 · pp. 2–3

Post-fine-tuning defenses

Concurrent post-fine-tuning defenses use interpolation, perturbation, projection, model fusion, decoding, neuron correction, activation patching, or optimization.Source: §2 · p. 3

Research gap

Prior work had not systematically studied whether defenses are sensitive to fine-tuning hyper-parameters.Source: §2 · p. 3

3. Preliminaries

The threat model assumes a service provider performs user fine-tuning and retains datasets for safety alignment and post-fine-tuning realignment. Experiments show that existing defenses degrade with larger learning rates or more epochs, despite these settings supporting fine-tuning accuracy.

3.1. Threat Model and Assumptions

3.1. Threat Model and Assumptions: Users upload fine-tuning data, after which the provider fine-tunes and deploys the aligned model as a personalized service.Source: §3.1 · p. 3
3.1. Threat Model and Assumptions: The provider hosts a harmful re-alignment dataset and maintains a safety alignment dataset for model re-alignment.Source: §3.1 · p. 3

3.2. Hyper-parameter Sensitivity Issue

3.2. Hyper-parameter Sensitivity Issue: Vaccine, RepNoise, Lisa, and LDIFS tend to show larger harmful scores when the fine-tuning learning rate increases.Source: §3.2 · p. 3, Figure 2 · p. 3
3.2. Hyper-parameter Sensitivity Issue: Large learning rates can subvert strengthened alignment or drive regularized iterates away from the aligned model.Source: §3.2 · p. 3
3.2. Hyper-parameter Sensitivity Issue: More fine-tuning epochs similarly enlarge harmful scores and break existing defenses.Source: §3.2 · p. 3, Figure 3 · p. 4
3.2. Hyper-parameter Sensitivity Issue: Sufficiently large learning rates are necessary for good fine-tuning accuracy, creating a hyper-parameter sensitivity issue.Source: §3.2 · p. 4

4. Methodology

Antidote is a post-fine-tuning safety-alignment method that identifies harmful parameters and removes them, remaining agnostic to the fine-tuning hyper-parameters. It uses Wanda scores on a re-alignment dataset to construct a harmful mask, then prunes the masked parameters to produce a deployable re-aligned model.

Antidote removes harmful parameters after fine-tuning, making the defense agnostic to the fine-tuning hyper-parameters.The method is designed to deactivate harmful parameters regardless of how they formed during fine-tuning.Source: §4 · p. 4

Antidote algorithm

The procedure takes a safety-alignment-broken fine-tuned model, computes importance scores, and constructs the harmful mask.Algorithm 1 specifies the broken model as input and the importance-score calculation followed by mask construction.Source: §4 · p. 4

Identify harmful parameters

Wanda scores identify important, potentially harmful parameters using the fine-tuned model weights and a harmful re-alignment dataset.The score relates parameter importance to the absolute weight value and its associated input activation.Source: §4 · p. 4
The harmful mask selects the top-k parameter coordinates for removal, with α specifying the ratio of coordinates masked to one.ArgTopKα produces a binary mask whose selected coordinates indicate parameters important for harmful-content generation.Source: §4 · p. 4

Removal of Harmful parameter

Pruning removes the masked harmful parameters from the fine-tuned weights, yielding a re-aligned model ready for deployment.The recovered model is then deployed to serve users’ customized tasks.Source: §4 · p. 4

5. Experiments

Experiments evaluate Antidote across harmful-fine-tuning conditions, datasets, models, system costs, ablations, extensions, and qualitative behavior. Overall, Antidote reduces harmful behavior while preserving downstream accuracy and remains comparatively robust to fine-tuning hyper-parameters.

5.2. Main Results

11.56% average harmful-score reduction versus SFT comes with a 1.45% average finetune-accuracy loss, while Antidote remains robust as the harmful-data ratio increases.Antidote achieves the lowest harmful score in most harmful-ratio settings and avoids the degradation seen in other defenses.Source: §5.2 · pp. 5–6, Table 1 · p. 5
6.56% average harmful-score reduction with a 0.38% average finetune-accuracy drop shows lower learning-rate sensitivity than competing fine-tuning-stage defenses.Lisa and LDIFS obtain larger reductions, 7.72% and 9.50%, but with drastic finetune-accuracy reductions.Source: §5.2 · pp. 6–7, Table 3 · p. 6

5.3. Generalizations on Datasets and Models

Antidote generalizes across fine-tuning datasets, reducing harmful score by 11.75% on average with a 3.08% finetune-accuracy loss.With a stronger safety-alignment dataset, it achieves over 40% harmful-score reduction when p = 0.2.Source: §5.3 · p. 7, Table 6 · p. 7, Table 7 · p. 7

5.4. Statistical and System Evaluation

Antidote maintains small harmful embedding drift after pruning, whereas other defenses’ drift escalates with larger learning rates and training epochs.The post-fine-tuning pruning stage lowers drift relative to the shared pre-pruning process, supporting recovery of preserved alignment knowledge.Source: §5.4 · pp. 7–8, Figure 5 · p. 8
1.02x clock time and unchanged GPU memory make Antidote substantially lighter than defenses requiring over 2x alignment time or over 1.7x GPU memory.Its extra cost comes mainly from calculating Wanda scores and applying the top-k mask.Source: §5.4 · p. 8, Table 10 · p. 8

5.5. Hyper-parameters Analysis and Ablation Study

Larger mask ratios simultaneously decrease harmful score and finetune accuracy, while using a harmful re-alignment dataset is necessary to identify harmful parameters.A benign dataset performs worst because it lacks harmful data needed to reveal harmful parameters.Source: §5.5 · pp. 8–9, Table 11 · p. 8, Table 12 · p. 9

5.6. Extensions

V-S-A lowers average harmful score by 2.46 and raises average finetune accuracy by 0.19 compared with Vanilla Antidote, while other combinations trade off these metrics differently.The visualization shows Antidote refusing a malicious request that other evaluated methods answer harmfully.Source: §5.6 · p. 9, Table 14 · p. 9, §5.7 · p. 9

6. Conclusion

The paper finds that existing harmful fine-tuning defenses fail under large learning rates or long fine-tuning, motivating Antidote as a hyperparameter-agnostic post-fine-tuning defense that preserves downstream accuracy.

Existing alignment-stage and fine-tuning-stage defenses fail when large learning rates or many fine-tuning epochs are used, despite these settings being necessary for downstream accuracy.Source: §6 · p. 9
Antidote removes harmful parameters after fine-tuning, aiming to recover the harmful model from harmful behaviors regardless of how those parameters formed.Source: §6 · p. 9
Antidote achieves strong defense performance while maintaining on-par accuracy on downstream tasks.Source: §6 · p. 9

Impact Statement

The paper studies harmful fine-tuning attacks as a security vulnerability of LLM fine-tuning APIs and evaluates them on open-weight models locally.

The study concerns the security vulnerability known as harmful fine-tuning attack in LLM fine-tuning APIs.Source: Impact Statement · p. 10
All experiments use open-weight LLMs in a local environment and therefore are not presented as posing direct societal risk.Source: Impact Statement · p. 10
The authors acknowledge that the finding could be misused to attack commercial LLM services and may negatively affect society.Source: Impact Statement · p. 10

A.1. Detailed Setup

The detailed setup specifies training hyperparameters, prompt formatting, and task-specific accuracy evaluation for alignment and fine-tuning experiments.

Experiments use the Alpaca prompt template to construct supervised alignment and fine-tuning examples.Source: §A.1 · p. 15
SST2, AGNEWS, and GSM8K accuracy counts correct test answers, while AlpacEval uses ChatGPT-rated win rate against text_Devinci_003.Source: §A.1 · p. 15

A.2. Implementation of Baselines and Their Idea

The baseline implementations cover vanilla SFT, alignment-stage defenses, and fine-tuning-stage defenses, with method-specific objectives and tuned hyperparameters.

The evaluation measures harmful score and fine-tuning accuracy on the finetuned models across five baselines.Source: §A.2 · p. 15
SFT aligns the pretrained model on alignment data and then fine-tunes it on user data without an additional defense mechanism.Source: §A.2 · p. 15
Baseline hyperparameters include Vaccine ρ = 2, RepNoise α = 0.1 and β = 0.001, Lisa ρ = 1, and LDIFS λ = 0.0001.Source: §A.2 · p. 15
Vaccine perturbs hidden embeddings during alignment, while RepNoise degrades harmful-data embeddings toward Gaussian noise to improve robustness.Source: §A.2 · pp. 15–16
Lisa alternates alignment-data and fine-tuning-data optimization with proximal penalties, whereas LDIFS brings current fine-tuning embeddings closer to the original aligned model.Source: §A.2 · p. 16