PaperlayerView exact revision on arXiv

Source-linked AI summary

Assessing the Brittleness of Safety Alignment via Pruning and Low-Rank Modifications

Boyi Wei, Kaixuan Huang, Yangsibo Huang, Tinghao Xie, Xiangyu Qi, Mengzhou Xia, Prateek Mittal, Mengdi Wang, Peter Henderson

arXiv:2402.05162v4cs.LGcs.AIcs.CL

TL;DR

The paper asks why aligned LLM safety mechanisms remain vulnerable to jailbreaks and non-malicious fine-tuning. It uses pruning and low-rank modifications to disentangle safety-critical regions from utility regions at neuron and rank levels, finding that sparse safety regions can be removed with safety degradation but little utility loss. The study also finds that restricting modifications to those regions does not eliminate fine-tuning vulnerability.

  • Problem

    Aligned LLMs remain susceptible to jailbreaks and can lose safety after non-malicious fine-tuning, motivating analysis of the weight regions underlying this fragility.

  • Method

    The study identifies and isolates safety-critical neurons and ranks using pruning, ActSVD, set difference, and orthogonal projection while measuring safety and utility changes.

  • Results

    3% of weights and about 2.5% of ranks are safety-critical; removing them raises attack success from 0 to over 90% while largely preserving utility.

  • Takeaways & Limitations

    Safety alignment can depend on sparse regions, and freezing those regions still leaves LLMs vulnerable to fine-tuning attacks.

  • Takeaways & Limitations

    Limited publicly accessible strong safety-aligned models constrain the experiments to Llama2-chat models, and other aligned models may behave differently.

Abstract

from arXiv · show

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.

1. Introduction

The study investigates why aligned LLM safety mechanisms are fragile by disentangling safety-critical weight regions from utility-relevant regions at neuron and rank levels. These regions are sparse, and removing them substantially weakens safety while largely preserving utility.

  • Jailbreaks and non-malicious fine-tuning can weaken safeguards in both open-access and closed-access aligned models.
  • The central challenge is separating safety mechanisms from overlapping general capabilities while locating the smallest safety-critical weight regions.
  • The proposed analysis disentangles safety and utility at both individual-neuron and weight-matrix-rank levels.
  • 3% of weights and about 2.5% of ranks comprise the isolated safety-critical regions in Llama2-chat models.
  • Removing identified safety-critical neurons or ranks raises attack success from 0 to over 90% while leaving overall utility largely unaffected.
  • Freezing safety-critical regions does not prevent low-cost fine-tuning attacks, which may introduce pathways bypassing the original safety mechanisms.
  • The authors propose safety-critical-region sparsity as a model-intrinsic brittleness metric that can complement red teaming.

2. Methodology

The methodology identifies behavior-specific neurons and ranks using calibration data, then isolates safety-critical regions from utility-relevant regions. Weight modifications are evaluated by their effects on layer outputs or final loss.

  • The pipeline modifies weights through neuron removal or rank removal and observes resulting changes in safety and utility behavior.
  • Identifying Important Neurons and Ranks: SNIP scores weight entries using first-order loss sensitivity, while Wanda scores neurons by minimizing output change under pruning.
  • The study excludes rank attribution combined with loss change because differentiating through an orthogonal projection requires additional technical analysis.
  • Identifying Important Neurons and Ranks: ActSVD uses response activations to identify important low-rank components of each weight matrix through singular-value decomposition.
  • Isolating Safety-Critical Neurons and Ranks: Safety-critical regions are selected to have low influence on general-language behavior but high influence on harmful-request responses that should decline.
  • Isolating Safety-Critical Neurons and Ranks: Neuron attribution isolates safety-critical neurons by taking the set difference between highly safety-important and utility-important neurons.
  • Isolating Safety-Critical Neurons and Ranks: Rank attribution uses orthogonal projection to remove safety-important ranks that are orthogonal to important utility ranks.

3. Experimental Setup

Experiments use Llama2-chat models, safety and utility datasets, three ASR settings, averaged zero-shot accuracy, and neuron- and rank-attribution methods.

  • Models: The experiments use Llama2-7B-chat and Llama2-13B-chat because their weights are publicly accessible and extensively safety-tuned.
  • Datasets: Safety data comprise harmful AdvBench instructions, while utility data are filtered Alpaca-Cleaned prompt-response pairs excluding safety-related examples.Safety responses include full-response and judgment-segment variants.
  • Evaluation Metrics: Utility is measured by averaged zero-shot accuracy across six EleutherAI LM Harness tasks.The tasks are BoolQ, RTE, HellaSwag, WinoGrande, ARC Challenge, and OpenbookQA.
  • Evaluation Metrics: Safety is measured as attack success rate on harmful instructions under standard use, adversarial suffixes, and adversarial decoding.ASR counts responses lacking key patterns indicating instruction rejection.
  • Attribution Methods: Neuron attribution evaluates SNIP and Wanda importance scores, including SNIP set difference, while rank attribution evaluates ActSVD and orthogonal projection.The set-difference and projection methods compare safety-related regions with utility-related regions.

4. Experimental Results

The results show that disentangling safety from utility isolates sparse neurons and ranks whose removal sharply weakens safety while preserving utility, whereas removing low-safety-relevance regions can improve robustness. Safety and utility are less overlapping in MLP layers, but freezing identified safety neurons does not generally stop fine-tuning attacks.

  • Method Comparisons: SNIP set difference and rank orthogonal projection outperform alternative methods in isolating safety-critical regions while retaining utility.Attention-head probing predicts harmful versus harmless instructions well but does not isolate neurons solely responsible for safety responses.
  • Safety-Critical Regions: Removing less than 3% of neurons drives ASR in all three scenarios close to 1 while maintaining average zero-shot accuracy above 0.5.For ranks, removing 2.5%—fewer than 100 of 4096—also increases ASR while preserving zero-shot accuracy.
  • Safety-Critical Regions: Less than 1% neuron pruning can significantly compromise adversarial safety while keeping accuracy above 0.53.
  • Safety-Critical Regions: Removing top safety-scored neurons or ranks severely compromises utility, supporting the need to disentangle safety-related regions from utility-related regions.The model’s average accuracy drops to about 0.35 from an original 0.58 when top safety neurons are removed.
  • Safety Enhancement: Pruning neurons with the lowest safety scores slightly enhances safety while keeping accuracy largely stable above 0.5.
  • Safety Enhancement: Removing higher-order ranks decreases ASRAdv-Decoding, while ASRAdv-Suffix varies considerably because adversarial suffixes may not transfer to modified models.
  • Layerwise Overlap: MLP layers show lower neuron overlap and lower rank subspace similarity between safety and utility than attention layers, indicating more differentiated behaviors.

5. Limitations & Future Work

The study identifies limitations in model coverage and localization, while proposing directions for improving safety robustness and attribution.

  • Limitations: Experiments are constrained to Llama2-chat because few strong safety-aligned models are publicly accessible.Other aligned models trained with different datasets and strategies might behave differently under this methodology.
  • Limitations: Standard attention head probing does not effectively localize safety-critical neurons.The findings suggest that MLP layers may offer better localization for safety knowledge.
  • Future Work: Pruning regions least important for safety, or potentially explicitly harmful, could improve safety robustness.Making safety-critical regions difficult to isolate is another proposed direction for building inherently safer models.

6. Conclusion

The study introduces a pipeline that separates safety-critical from utility-relevant regions in LLM weights. Experiments on Llama2-chat find these regions are sparse yet crucial to safety, motivating intrinsic brittleness assessment and further robustness research.

  • Conclusion: The pipeline disentangles regions critical for safety from those vital for utility in LLM weights.
  • Conclusion: 3% at the weight level and 2.5% at the rank level account for safety-critical regions in Llama2-chat models.
  • Conclusion: Removing safety-critical regions destroys safety while retaining utility.
  • Conclusion: The sparsity of safety-critical regions may explain brittleness and could serve as a model-intrinsic metric alongside red teaming.

Contribution Statement

The project combined contributions across idea development, literature review, evaluation, implementation, and manuscript preparation.

  • Idea formulation and preliminary exploration: Core ideas linked safe behaviors to model regions and isolated safety from utility.
  • Literature survey: Team members surveyed alignment, jailbreaks, task attribution, low-rank modifications, and pruning techniques.
  • Evaluation: The team prepared attribution data, established the evaluation pipeline, created visuals, and investigated overlap and freezing effects.
  • Writing: The initial manuscript structure and primary drafting were led by Yangsibo, Boyi, and Kaixuan, with others editing and providing feedback.

Impact Statement

The work aims to improve AI safety by studying sparse and removable safety guardrails, while acknowledging dual-use risks and limits in safety definitions and model coverage.

  • Impact: The research aims to encourage safety mechanisms that are not sparse, easy to isolate, or easy to remove.
  • Dual-use Risk: Safety and security research carries a risk that adversaries could use the work to remove safety guardrails.
  • Potential Benefits: The authors argue that studying guardrail brittleness can encourage stronger guardrails that are harder for attackers to isolate and remove.
  • Impact: The work is intended to improve AI safety, particularly in open models, through analysis and information.
  • Safety and harm definitions: Standard safety and harm benchmarks may not cover all definitions of safety and harm.
  • Alignment and Jailbreak: Aligned LLMs can follow human values and refuse harmful requests, but jailbreaks can bypass or remove those guardrails.

A.2. Identifying Task-Specific Regions in Large Language Models

Weight attribution methods identify task-specific regions in transformer models, including neurons, attention heads, and low-rank components. This paper adapts pruning and low-rank techniques to locate regions associated with safety while distinguishing them from utility.

  • Weight attribution: Weight attribution is a classic explainable-machine-learning problem for linking model behavior to specific parameters.Prior methods include perturbation-based and gradient-based importance scores, increasingly applied to transformer models.
  • Probing: Probing feeds model representations or parameters to a classifier that predicts linguistic properties or reasoning abilities.The paper uses probing as a baseline for localizing safety-related representations.
  • Related localization methods: Recent localization methods target linguistic properties, general capabilities, fine-tuning, and prompt tuning in transformer models.
  • Alignment and localized behavior: Prior work suggests alignment may bypass toxicity-eliciting regions rather than eliminate the ability to generate toxic outputs.
  • Paper's focus: The paper applies pruning-style attribution to identify critical regions responsible for safety rather than only compressing models.

B. Experimental Details

The experiments evaluate neuron- and rank-level attribution and pruning on Llama-family models using safety, utility, adversarial, and downstream-task data. Procedures include block-wise and output-wise pruning, probing baselines, repeated trials, and specified attack configurations.

  • Compute: Experiments use four AMD EPYC CPUs and one NVIDIA A100-80G GPU, with vLLM for faster decoding.Typical GPU-hour requirements are summarized in Table 4.
  • Pruning procedure: Block-wise pruning processes Llama sequentially, pruning seven linear layers in each Transformer block before recomputing that block's outputs.
  • Pruning procedure: Neuron attribution uses output-wise pruning, setting a chosen percentage of each weight matrix row to zero independently.
  • Attribution data: Importance scores and projection matrices are computed from 128 sampled prompt-response pairs, using the safety and utility datasets listed in Table 5.
  • Reproducibility: Experiments in Section 4.2 are repeated three times with different random seeds, reporting means and standard-deviation bands.
  • Evaluation: The probing baseline trains on 420 harmful and 420 harmless instructions, while adversarial evaluations use GCG suffixes and stochastic decoding attacks.Downstream utility includes ARC-Challenge, HellaSWAG, OpenBookQA, WiNoGrande, BoolQ, and RTE.

C.1. More results in Llama2-13B-chat

Additional Llama2-13B-chat results confirm that disentangling safety from utility isolates regions whose removal damages safety while preserving utility. They also show model- and attack-dependent effects when removing regions considered least safety-critical.

  • Safety-critical regions: Set difference for neurons and orthogonal projection for ranks destroy safety while preserving utility in Llama2-13B-chat.These methods outperform alternatives in isolating safety-critical regions.
  • Safety-critical regions: Removing top safety neurons or ranks severely hurts utility, unlike removing disentangled safety-critical regions.
  • Method comparison: SNIP set difference consistently outperforms attention-head probing as an isolation method.
  • Least-critical pruning: Llama2-13B-chat retains nearly zero adversarial-suffix ASR after removing less than 45% of neurons or less than 750 least safety-critical ranks.The text suggests transfer failure from suffixes optimized on 7B models as one possible explanation.
  • Attribution data: Safety-full and safety-short attribution produce similar ASR-versus-accuracy trends, while safety-short changes adversarial metrics more during least-critical pruning.
  • Disentanglement results: For set difference, actual sparsity between 1% and 3% drives all three ASR variants close to 1 while average accuracy remains above 0.5.
  • Disentanglement results: Orthogonal projection is especially effective when the safety rank target approaches the total rank and the utility rank target approaches it.

D. Proof of the Optimality of ActSVD

The ActSVD proof uses the SVD of activation-weight outputs to construct a rank-constrained approximation. Projecting the weight matrix onto the top left singular vectors yields the optimal solution.

  • Problem formulation: ActSVD stores input activations before a layer and seeks a low-rank weight matrix minimizing output change in Frobenius norm.
  • SVD argument: The SVD of WX_in supplies the best rank-r approximation to the layer outputs by the Eckart–Young–Mirsky theorem.
  • Optimality: Setting cW = UU^T W produces a matrix with rank at most r and solves the rank-constrained minimization problem.The projection UU^T uses the top-r left singular vectors.
Loading 2402.05162v4…