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

Abstract: 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 respond to harmful prompts and propose a novel defense method termed L ayer-specific Ed iting (LED) to enhance the resilience of LLMs against jailbreak attacks. Through LED, we reveal that several critical 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 identified toxic 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 https://github.com/ledllm/ledllm .

## introduction large language models (llms) such as gpt-4 @xcite , llama2 @xcite , vicuna @xcite , and mistral @xcite have demonstrated remarkable capabilities across a wide range of natural language tasks and have been increasingly adopted in many real-world applications. despite extensive efforts @xcite @xcite figure 1 : example responses from normal and pruned llms to harmful and jailbreak prompts. when some crucial layers are removed, llms surprisingly provide harmful responses to unchanged harmful queries. @xcite @xcite to align llms' responses with human values to generate helpful and harmless content, recent studies @xcite @xcite @xcite @xcite @xcite @xcite reveal that these aligned models are still vulnerable to intentionally crafted adversarial prompts, also termed as "jailbreak attacks", which can elicit harmful, biased, or otherwise unintended behaviors from llms, posing significant challenges to their safe deployment. to mitigate the jailbreak attacks on llms, various methods have been proposed. however, existing defense methods primarily focus on two aspects: 1) detecting whether the prompt or response contains harmful or unnatural content via perplexity filter @xcite , input mutation @xcite , or using the llm itself @xcite ; and 2) reducing the probability of generating harmful responses through safe instruction @xcite @xcite or logit processor @xcite . while these methods reduce the attack success rate of adversarial prompts to some extent, their effectiveness may be overcome by adaptive attacks @xcite . our take is that there remains a significant gap in our understanding of how llms handle harmful prompts and whether jailbreak attacks exploit llms in specific ways to produce harmful responses. without diving into the inner workings of llms, existing efforts to improve their safety may only scratch the surface. recent research about llm pruning and layer skipping @xcite @xcite found that removing certain layers does not significantly affect llm performance. additionally, observations made in @xcite suggest that early layers are crucial in defending against adversarial attacks. these studies suggest that not all layers contribute equally when responding to harmful prompts and adversarial prompts. in this work, we take a step toward in terms of understanding the underlying safety mechanisms of llms. we conduct a systematic layer-wise analysis to identify layers that significantly influence llm responses in the presence of harmful and jailbreak prompts. figure 1 illustrates example responses from normal and pruned llms where some selected layers are removed to different prompts. surprisingly, we observe several critical safety layers responsible for handling harmful prompts. once these layers are pruned, llms can be jailbroken by simply inputting the original harmful prompts without any modification. furthermore, our detailed analysis of each layer's hidden states reveals that not all layers contain toxic information that triggers llms to generate harmful responses; some layers maintain a relatively high probability of decoding refusal tokens. based on these insights, we propose a novel jailbreaking defense method termed layer-specific editing (led), which uses targeted model editing to enhance llm defense against adversarial attacks. through led, we reveal that several critical safety layers are located in the early layers of llms. realigning these safety layers and some selected additional layers that merely contribute to defense with the decoded safe response from identified toxic layers can significantly improve the safety alignment of llms, whilst maintaining their performance on benign prompts. our main contributions are summarized as follows. ## led: layer-specific editing for enhancing defense of llms in this work, we propose led, layer-specific editing to enhance the safety alignment of llms. figure 2 illustrates the overview of led's workflow. specifically, led consists of three critical steps: 1) selecting edited layers, which consist of safety layers mostly relevant to the safety alignment for the harmful queries and additional layers which merely contribute to the defense; 2) locating toxic layers, which act as the optimization object for fully eliminating unexpected information; and 3) layer-specific editing to align the edited layers with the decoded safe responses from toxic layers, thereby enhancing the defense effectiveness against jailbreak attacks. given the importance of the initial embedding layer, we start with layer 1 rather than layer 0 in the probing. we limit the pruning to a maximum of half the model's layers to retain meaningful output from the pruned llm as suggested in previous work @xcite . then we investigate the response of f l,n to a set of harmful prompts. if the response is harmful, we stop the pruning process and treat layer l to l + n as safety layer candidates. in practice, we use an array layer_frequency to count the frequency of safety layer candidates over all harmful prompts. the top-k layers that appear most frequently during the pruning process are designated as safety layers s as described in equation 2: to enhance the robustness of the model editing, we also incorporate a selection of additional layers that merely contribute to the defense into s to obtain the edited layers e, aiming to involve more layers in the defense against jailbreak attacks. to automatically detect toxic regions, we introduce a layer-specific toxic score t (h l ) to quantify the number of toxic responses in the decoded output of layer l, defined as equation 3 : in the next section, we first present some interesting findings regarding safety layers and toxic layers. then, we conduct extensive experiments to evaluate the effectiveness of our proposed led. ## a closer look at llms: safety and toxic layers early safety layers dominates defense. to identify the presence of safety layers, we conduct a layer-wise pruning analysis on a set of 100 randomly selected harmful prompts from advbench @xcite as demonstrated in section 3. figure 3( a) presents the results of four different llms responding to naturally harmful prompts before and after pruning safety layers, where a higher attack success rate (asr) indicates lower defense performance. surprisingly, safety layers 26 27 28 29 30 31 layer 0.0 0.2 0.4 0.6 0.8 1.0 safety score llama2-7b safety score on harmful and adversarial promtps series harmful adversarial are widely present in different llms, regardless of their varying structures and training data. pruning these layers significantly increases the asr for naturally harmful prompts compared to the original model. figure 3(b) illustrates the distribution of safety layers for each model. for different llms, their safety layers are mainly concentrated in early layers, indicating that these early layers play a crucial role in discriminating the harmfulness of input queries, consistent with previous findings @xcite . notably, almost all later layers do not participate in the defense against harmful queries. this observation leads to the idea of enhancing the llm's safety by involving more layers in defending against harmful queries. see multiple later layers contain toxic information. figure 4 shows that later layers, rather than only the final layer, generally tend to output the toxic token with high probability which requires realignment. notably, we do not directly edit the parameters in these toxic layers because it is impractical to eliminate all harmful knowledge within them. we cannot identify all adversarial prompts that trigger harmful content in these layers, and even if we could, we still can not ensure that new methods would not retrieve harmful information from these layers. instead, we set these toxic layers as targets and fine-tune the model to generate only safe refusal responses from them. ## conclusion this paper systematically investigates the intrinsic defense mechanism of llms against harmful and adversarial prompts through layer-wise pruning and decoding analysis, finding that some layers play a crucial role in defending harmful queries and many layers of defense capabilities are not fully utilized according to the unbalanced distribution of safety layers. therefore, we introduce led, a layer-specific editing approach to enhance the robustness of llms against adversarial attacks while maintaining the model's helpfulness on benign queries