[ICML 2024] Assessing the Brittleness of Safety Alignment via Pruning and Low-Rank Modifications
Assessing the Brittleness of Safety Alignment via Pruning and Low-Rank Modifications
This repository provides an original implementation of Assessing the Brittleness of Safety Alignment via Pruning and Low-Rank Modifications by Boyi Wei, Kaixuan Huang, Yangsibo Huang*, Tinghao Xie, Xiangyu Qi, Mengzhou Xia, Prateek Mittal, Mengdi Wang and Peter Henderson.
1. Setup
You can use the following instruction to create conda environment
conda env create -f environment.ymlenvironment.yml
Besides, you need to manually install a hacked version of lmeval to support evaluating the pruned model. See wanda.
There are known issues with the transformers library on loading the LLaMA tokenizer correctly. Please follow the mentioned suggestions to resolve this issue.
Before running experiments, make sure you have specified the path pointing to the model stored in your locations.
2. Neuron Level Usage
2.1 Top-down Pruning
The main function is
main.py. When using Top-down pruning, we need to add --neg_prune in the command line.
Important parameters are:
--prune_method: To specify the prune method. Available options are wanda, wandg (SNIP in the paper), random. --prunedata: To specify datasets used for pruning. When doing top-down pruning safety-critical neurons, we can use align(safety-full in the paper) and alignshort (safety-short in the paper) as our dataset. --sparsity_ratio: Specify the prune sparsity. --evalzeroshot: Whether to evaluate the model's zero-shot-accuracy after pruning --eval_attack : Whether to evaluate the model's ASR after pruning. --save: Specify the save location --model: Specify the model. Currently we only support llama2-7b-chat-hf and llama2-13b-chat-hf
Example: Using
llama2-7b-chat-hf to prune 0.5 part of weights, using safety-full dataset. <pre><code class="lang-bash">model="llama2-7b-chat-hf" method="wanda" type="unstructured" suffix="weightonly" savedir="out/$model/$type/${method}${suffix}/align/"
python main.py \ --model $model \ --prune_method $method \ --prune_data align \ --sparsity_ratio 0.5 \ --sparsity_type $type \ --neg_prune --save $save_dir \ --evalzeroshot \ --eval_attack \ --saveattackres</code></pre>
2.2 Pruning the least safety-critical neurons.
Simply remove
--negprune will reverse the order of pruning. We recommend using alignshort (safety-short in our paper) when pruning the least safety-critical neurons to get more obvious results.
2.3 Pruning based on Set Difference
Select option
--prunemethod as wandgsetdifference (SNIP with set difference in our paper). Add option --p, which corresponds to top-p scored entries in alpacanosafety-based wandg score; Add option --q, which corresponds to top-q scored entries in aligned-based wandg score. Please notice that you have to specify a non-zero value of --sparsityratio. For the dataset to compute the utility importance score, we usealpacacleanednosafety by default, --prunedata here is used to specify the dataset to compute the safety importance score. Available options are align (safety-full in our paper) and align_short (safety-short in our paper). Please note, for set difference methods, you need to first dump the Wanda / SNIP score before pruning! (See https://github.com/boyiwei/alignment-attribution-code/edit/main/README.md#24-wandasnip-score-dumping)
Example: Pruning the set difference between top-10% utility-critical neurons (Use alpacacleanedno_safety dataset to identify) and top-10% safety-critical (Use safety-full to identify) safety neurons.
<pre><code class="lang-bash">model="llama2-7b-chat-hf" method="wandgsetdifference" type="unstructured" suffix="weightonly" savedir="out/$model/$type/wandgsetdifference{$suffix}"
python main.py \ --model $model \ --prune_method $method \ --sparsity_ratio 0.5 \ --prune_data align --p 0.1\ --q 0.1\ --sparsity_type $type \ --save $save_dir \ --evalzeroshot \ --eval_attack \ --saveattackres</code></pre>
2.4 Wanda/SNIP Score dumping
Simply add option
--dumpwandascore into the command. If you want to dump Wanda score, use --prunemethod wanda and --dumpwanda_score. If you want to dump SNIP score, use --prunemethod wandg and --dumpwanda_score.
Example: Safety-first pruning with align_llama2-7b-chat dataset:
<pre><code class="lang-bash">model="llama2-7b-chat-hf" method="wanda" type="unstructured" suffix="weightonly" savedir="out/$model/$type/${method}${suffix}/align/"
python main.py \ --model $model \ --prune_method $method \ --prune_data align \ --sparsity_ratio 0.5 \ --sparsity_type $type \ --save $save_dir \ --dumpwandascore</code></pre>
3. Rank Level Usage
3.1 Remove the most safety-critical rank
The main function of this pipeline is
mainlowrank.py. Most of the parameters are similar to the prune neurons situation.
Important parameters are:
--prunemethod: To specify the pruning method, in this case we choose lowrank, which corresponds to ActSVD in our paper.
--prunedata: To specify the dataset used to identify the safety/utility projection matrix. Available options are align (safety-full), alignshort (safety-short), alpacacleanednosafety (filtered alpacacleaned dataset)
--rank: To determine how many ranks needed to be removed .
--top_remove : To determine whether to remove the top-critical ranks or the least-critical ranks. If true, remove the top critical ranks
Example: Prune the top-10 safety-critical rank based on the safety-full(align in the code) dataset.
<pre><code class="lang-bash">model="llama2-7b-chat-hf" method="low_rank" type="unstructured" suffix="weightonly" savedir="out/$model/$type/${method}${suffix}/align/"
python mainlowrank.py \ --model $model \ --prune_method $method \ --prune_data align \ --rank 10 \ --top_remove \ --save $save_dir \ --evalzeroshot \ --eval_attack \ --saveattackres</code></pre>
3.2 Remove the least safety-critical ranks
Similar to 3.1, but here we don't need to add --top_remove in the command line.
Example: Remove the bottom-1000 safety-critical rank based on the safety-short(
align_short in the code) dataset.
<pre><code class="lang-bash">model="llama2-7b-chat-hf" method="low_rank" type="unstructured" savedir="out/$model/$type/${method}/alignshort/"
python mainlowrank.py \ --model $model \ --prune_method $method \ --prunedata alignshort \ --rank 1000 \ --top_remove \ --save $save_dir \ --evalzeroshot \ --eval_attack \ --saveattackres</code></pre>
3.3 Remove rank with orthogonal projection
The main function of this program is mainlowrank_diff.py.
Important parameters are:
--prunemethod: To specify the method of rank removal, here we use lowrank_diff, which corresponds to the (ActSVD with orthogonal projection in the paper)
--rank_pos: Specify the $r^u$ in the paper.
--rank_neg: Specify the $r^s$ in the paper.
--prunedatapos: The data to determine the utility projection matrix, we use alpacacleanedno_safety.
--pruneddataneg: The data to determine the safety projection matrix, we recommend to use align.
Example: Prune based on rank-3000 utility projection matrix and rank-4000 safety projection matrix on alpacacleanednosafety (filtered alpacacleaned dataset without safety-related prompt-response pairs) and safety-full on llama2-7b-chat-hf`.
model="llama2-7b-chat-hf"
type="unstructured"
ru=3000
rs=4000
method="lowrankdiff"
save_dir="out/$model/$type/${method}/align/"
python mainlowrank_diff.py \ --model $model \ --rank_pos $ru \ --rank_neg $rs \ --prunedatapos "alpacacleanedno_safety" \ --prunedataneg "align" \ --save $save_dir \ --evalzeroshot \ --eval_attack \
4. BibTeX
If you find our code and paper helpful, please consider citing our work:@InProceedings{pmlr-v235-wei24f,
title = {Assessing the Brittleness of Safety Alignment via Pruning and Low-Rank Modifications},
author = {Wei, Boyi and Huang, Kaixuan and Huang, Yangsibo and Xie, Tinghao and Qi, Xiangyu and Xia, Mengzhou and Mittal, Prateek and Wang, Mengdi and Henderson, Peter},
booktitle = {Proceedings of the 41st International Conference on Machine Learning},
pages = {52588--52610},
year = {2024},
editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix},
volume = {235},
series = {Proceedings of Machine Learning Research},
month = {21--27 Jul},
publisher = {PMLR},
pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/wei24f/wei24f.pdf},
url = {https://proceedings.mlr.press/v235/wei24f.html},
abstract = {Large language models (LLMs) show inherent brittleness in their safety mechanisms, as evidenced by their susceptibility to jailbreaking and even non-malicious fine-tuning. This study explores this brittleness of safety alignment by leveraging pruning and low-rank modifications. We develop methods to identify critical regions that are vital for safety guardrails, and that are disentangled from utility-relevant regions at both the neuron and rank levels. Surprisingly, the isolated regions we find are sparse, comprising about $3$ % at the parameter level and $2.5$ % at the rank level. Removing these regions compromises safety without significantly impacting utility, corroborating the inherent brittleness of the modelโs safety mechanisms. Moreover, we show that LLMs remain vulnerable to low-cost fine-tuning attacks even when modifications to the safety-critical regions are restricted. These findings underscore the urgent need for more robust safety strategies in LLMs.}
}
