PaperlayerView exact revision on arXiv

Source-linked AI summary

Defending Large Language Models Against Jailbreak Attacks via Layer-specific Editing

Wei Zhao, Zhe Li, Yige Li, Ye Zhang, Jun Sun

arXiv:2405.18166v2cs.AI

TL;DR

Jailbreak prompts can elicit harmful behavior from aligned LLMs, and defenses based on their internal mechanisms remain underexplored. The paper analyzes layer roles and proposes LED, which edits selected layers toward decoded safe responses. Experiments report reduced attack success across models and attacks while preserving helpfulness on benign benchmarks.

  • Problem

    Aligned LLMs remain vulnerable to deliberately crafted jailbreak prompts, while how their internal mechanisms handle harmful prompts is insufficiently understood.

    Source: §1 · pp. 1–2

  • Method

    LED identifies safety and selected defense-contributing layers, then edits them to align with decoded safe responses from identified toxic layers.

    Source: §1 · p. 2

  • Results

    LED significantly reduces attack success rates across different LLMs and jailbreak attacks while preserving helpfulness on benign benchmarks.

    Source: §6 · p. 8

  • Takeaways & Limitations

    The analysis identifies critical safety layers among early LLM layers, and realigning these layers can improve safety alignment while maintaining benign-prompt performance.

    Source: §1 · p. 2

    The paper calls for further study of LLM component functions to refine defenses and broaden their applicability.

    Source: §6 · p. 8

Abstract

from arXiv · show

Large language models (LLMs) are increasingly being adopted in a wide range of real-world applications. Despite their impressive performance, recent studies have shown that LLMs are vulnerable to deliberately crafted adversarial prompts even when aligned via Reinforcement Learning from Human Feedback or supervised fine-tuning. While existing defense methods focus on either detecting harmful prompts or reducing the likelihood of harmful responses through various means, defending LLMs against jailbreak attacks based on the inner mechanisms of LLMs remains largely unexplored. In this work, we investigate how LLMs response to harmful prompts and propose a novel defense method termed \textbf{L}ayer-specific \textbf{Ed}iting (LED) to enhance the resilience of LLMs against jailbreak attacks. Through LED, we reveal that several critical \textit{safety layers} exist among the early layers of LLMs. We then show that realigning these safety layers (and some selected additional layers) with the decoded safe response from selected target layers can significantly improve the alignment of LLMs against jailbreak attacks. Extensive experiments across various LLMs (e.g., Llama2, Mistral) show the effectiveness of LED, which effectively defends against jailbreak attacks while maintaining performance on benign prompts. Our code is available at \url{https://github.com/ledllm/ledllm}.

1 Introduction

Jailbreak prompts can elicit harmful behavior despite alignment, while existing defenses leave the models’ internal handling of harmful prompts insufficiently understood. The paper identifies critical safety layers and proposes LED to strengthen jailbreak defense while maintaining benign-prompt performance.

1 Introduction

  • Removing selected early safety layers can make LLMs answer unchanged harmful queries harmfully, as if their alignment were undone.The paper also finds that only certain early layers play a crucial role in identifying harmful prompts.Source: §1 · p. 2
  • LED uses targeted editing to realign safety layers and selected defense-contributing layers with decoded safe responses, while maintaining benign-prompt performance.The paper locates several critical safety layers among the early layers of LLMs.Source: §1 · p. 2
  • Some layers retain a high probability of decoding refusal tokens under jailbreak attacks, suggesting attacks may alter final responses without changing every intermediate output.Source: §1 · p. 2
  • Experiments across Llama2, Mistral, and other LLMs report that LED defends against various state-of-the-art adversarial attacks.Source: §1 · p. 2

2 Related Work

Prior work studies both ways to jailbreak LLMs and defenses that detect harmful content or reduce harmful outputs, while knowledge editing offers another route to changing model behavior. The paper distinguishes LED from detoxification approaches by editing safety layers rather than directly editing toxic layers.

2 Related Work

  • Jailbreak methods range from handcrafted and social-media prompts to gradient-based, genetic, random-search, and auxiliary-LLM approaches.Examples include GPTFuzzer and PAIR, which use an LLM to refine templates or select candidate attack prompts.Source: §2 · p. 2
  • Existing defenses mainly detect harmful inputs or outputs, or reduce harmful-response probability through safe instructions and toxic-token logit changes.The passage notes that these methods do not provide a comprehensive understanding of LLMs’ underlying safety mechanisms.Source: §2 · p. 3
  • Knowledge editing includes fine-tuning, meta-learning, and locate-and-edit methods that target stored knowledge or model parameters.Locate-and-edit methods use findings about MLP knowledge storage; ROME and MEMIT use causal mediation analysis to identify relevant hidden-state components.Source: §2 · p. 3
  • Unlike detoxification methods, LED does not directly edit toxic layers; it identifies safety layers and edits them to realign toxic layers toward safe outputs.The paper notes that existing knowledge-editing methods may fail to erase all knowledge, leaving ways to retrieve it.Source: §2 · p. 3

3 LED: Layer-specific Editing for Enhancing Defense of LLMs

LED selects layers involved in safety defense, locates layers that decode toxic content, then edits selected layers to align those decoded outputs with safe responses. Its workflow combines pruning-based safety-layer selection, toxic-layer scoring, and layer-specific editing.

3 LED: Layer-specific Editing for Enhancing Defense of LLMs

  • LED’s three steps are selecting edited layers, locating toxic layers, and aligning edited-layer outputs with decoded safe responses from toxic layers.After optimization, the method updates the weights of each edited layer.Source: §3 · pp. 4–5
  • LED selects safety layers by pruning consecutive layers and identifying those whose removal makes responses to harmful prompts harmful.It counts candidate-layer frequency across harmful prompts and selects the top-k layers; selected additional defense-contributing layers are included in the edited set.Source: §3 · p. 4
  • LED locates toxic layers by decoding hidden states from adversarial prompts and scoring how strongly each layer supports the final model’s toxic token.Layers with an average toxic score above 0.5 are optimization objectives; the method aligns their outputs with safe content rather than trying to eliminate all harmful knowledge.Source: §3 · p. 4
  • Figure 3 presents layer-wise pruning results and the frequency distribution of safety layers, which are mainly found in early layers.The pruning analysis covers four LLMs and 100 randomly selected harmful prompts from AdvBench.Source: Figure 3 · p. 5
  • The editing objective aligns decoded content from toxic layers with safe responses to harmful prompts.LED uses harmful-prompt and desired-safe-response pairs, with the hidden states of toxic layers as targets.Source: §3 · p. 5

4 A Closer Look at LLMs: Safety and Toxic Layers

Layer-wise pruning shows that safety layers occur across different LLMs and are concentrated in early layers. Later layers can also decode toxic tokens, so LED treats them as alignment targets rather than directly editing away their harmful knowledge.

4 A Closer Look at LLMs: Safety and Toxic Layers

  • Pruning safety layers significantly increases attack success rates on naturally harmful prompts across four LLMs.The analysis uses 100 randomly selected harmful prompts from AdvBench; higher ASR indicates lower defense performance.Source: §4 · p. 5
  • Safety layers are mainly concentrated in early layers, which play a crucial role in discriminating harmful queries; almost all later layers do not participate in this defense.The distribution is reported across LLMs with different structures and training data.Source: §4 · p. 5
  • Later layers beyond the final layer can assign high probability to toxic tokens, so LED targets their decoded outputs for safe refusal alignment instead of directly editing their parameters.The paper notes it is impractical to eliminate all harmful knowledge in these layers or identify every adversarial prompt that could trigger it.Source: Figure 4 · p. 5, §4 · p. 6

5 Experiments

LED is evaluated against multiple jailbreak defenses and LoRA settings, with layer-selection and Self-Examination ablations. The experiments report reduced attack success while largely retaining helpfulness, and identify a combination of safety and middle-layer edits as effective.

5.1 Experiment Setup

  • Experiments use AdvBench prompts, five jailbreak attacks, five comparison defenses, and Llama2-7B and Mistral-7B; LED layer choices differ by model.The study also assesses helpfulness with MT-bench and Just-Eval.Source: §5.1 · p. 6

5.2 Effectiveness of LED against Jailbreak Attacks

  • LED reduces average Mistral-7B attack success rate (ASR) to 11.3%, brings Llama2-7B ASR for all attacks to nearly 0%, and reduces helpfulness by 2% and 1%, respectively.The reported comparisons cover multiple jailbreak attacks and defense strategies.Source: §5.2 · p. 6, Table 1 · p. 7, Table 2 · p. 7
  • Fine-tuning only safety layers with LoRA does not significantly improve robustness except against GPTFuzzer; larger LoRA datasets help, but remain behind LED.The comparison examines LoRA settings on Mistral-7B.Source: §5.2 · pp. 6–7, Figure 5 · p. 8

5.3 Ablation Studies

  • Editing a single early layer gives the strongest jailbreak defense among single-layer edits, while later-layer edits have minimal effect.The layer-selection ablation evaluates Mistral-7B.Source: §5.3 · p. 7, Table 3 · p. 8
  • Editing safety layers alone improves defense but reduces helpfulness through sensitivity to benign queries resembling harmful prompts; adding middle layers yields security gains with negligible helpfulness degradation.Adding later layers does not significantly improve helpfulness or robustness.Source: §5.3 · p. 7
  • Using the edited model for Self-Examination improves results across attacks, but DeepInception remains less effective than direct defense with the edited model.DeepInception's imaginative scenarios and character interactions make harmful behavior harder to identify in long outputs.Source: §5.3 · p. 8, Figure 6 · p. 8

6 Conclusion

The paper investigates LLM defenses through layer-wise pruning and hidden-state analysis, then introduces LED to improve jailbreak robustness while retaining helpfulness. It identifies harmful-knowledge erasure and layer-specific evaluation as important boundaries and open questions.

6 Conclusion

  • LED improves robustness across LLMs and jailbreak attacks while preserving helpfulness on benign benchmarks, and the conclusion reports higher performance with fewer datasets than fine-tuning.The paper also states that LED does not require adversarial samples.Source: §6 · p. 8
  • Editing safety layers does not directly erase harmful knowledge; locating that knowledge and determining how to erase it remain open questions.The paper distinguishes the locations of defense mechanisms from the storage of harmful knowledge.Source: §6 · p. 8

A Layer-wise Pruning Jailbreak Examples

B More Analysis On Hidden States

The analysis compares adversarial prompts with minimally edited harmful and benign prompts, then examines how their hidden-state similarities vary across model layers.

B More Analysis On Hidden States

  • The analysis starts with 50 GCG- and AutoDAN-generated adversarial prompts, modifying tokens stepwise to create corresponding harmful and benign prompts with minimal changes.Adversarial suffixes are altered until harmful prompts elicit safe refusals; benign prompts replace harmful-query tokens with benign alternatives.Source: B More Analysis On Hidden States · p. 12, Figure 7 · p. 12
  • Hidden-state similarity is high across prompt types in layers 0–10; in layers 11–31, harmful and adversarial prompts remain similar in middle layers but diverge late.Normal prompts have low similarity to harmful or adversarial prompts in the latter layers.Source: B More Analysis On Hidden States · p. 12, Figure 8 · p. 12
  • The authors interpret the late-layer divergence as evidence that adversarial prompts mainly interfere with later layers.Source: B More Analysis On Hidden States · p. 12

C Detail Experiment Setting

The appendix describes the safety message added to potentially harmful prompts for weak-aligned models.

C.1 Safety Conversation Template

  • For weak-aligned models such as Vicuna and Mistral, the setup adds a safety instruction to each potentially harmful prompt to encourage denial of harmful requests.The template asks for helpful responses while excluding harmful, unethical, dangerous, or illegal content.Source: §C.1 · p. 12

C.2 Attack Setup

The evaluation uses five jailbreak attacks, implemented through EasyJailbreak with its default settings and GPT-4 as the attack model.

  • The evaluation covers PAIR, AutoDAN, GPTFuzzer, GCG, and DeepInception jailbreak attacks.Source: §C.2 · p. 12
  • EasyJailbreak provides the implementation for the attack evaluation.Source: §C.2 · p. 12
  • The setup follows EasyJailbreak defaults and uses GPT-4 to generate jailbreak material, including toxic prefixes, suffixes, or full prompts.Source: §C.2 · p. 12

C.3 Defense Baseline Setup

The defense baselines include input and output checks, prompt transformations, and a no-defense setting, with each method configured as described below.

  • The No Defense baseline uses the FastChat 0.2.36 conversation template without a system prompt.Source: §C.3 · p. 13
  • Self-Reminder applies a safe conversation template to encourage non-harmful responses.Source: §C.3 · p. 13
  • Self-Examination asks the model after each output whether it is harmful and returns a refusal when the answer is yes.Source: §C.3 · p. 13
  • The PPL baseline uses GPT-2 to score input prompts and rejects prompts above the highest perplexity of harmful Advbench queries.Source: §C.3 · p. 13
  • Paraphrase uses GPT-4 to rewrite each prompt before input, aiming to remove unrelated or potential jailbreak material.Source: §C.3 · p. 13
  • SafeDecoding amplifies safety-disclaimer token probabilities and attenuates probabilities of token sequences aligned with jailbreak objectives.The study uses its released fine-tuned Llama-2 and fine-tunes a safer Mistral version with default settings.Source: §C.3 · p. 13

D Detail Algorithm

This section points to Algorithm 1 for LED’s procedure and includes tables on the effects of pruning safety layers and examples from pruned models.

Loading 2405.18166v2…