Alignment-Enhanced Decoding: Defending Jailbreaks via Token-Level Adaptive Refining of Probability Distributions (2024)

Abstract: Large language models are susceptible to jailbreak attacks, which can result in the generation of harmful content. While prior defenses mitigate these risks by perturbing or inspecting inputs, they ignore competing objectives, the underlying cause of alignment failures. In this paper, we propose Alignment-Enhanced Decoding (AED), a novel defense that employs adaptive decoding to address the root causes of jailbreak issues. We first define the Competitive Index to quantify alignment failures and utilize feedback from self-evaluation to compute post-alignment logits. Then, AED adaptively combines Competitive Index and post-alignment logits with the original logits to obtain harmless and helpful distributions. Consequently, our method enhances safety alignment while maintaining helpfulness. We conduct experiments across five models and four common jailbreaks, with the results validating the effectiveness of our approach.

## introduction large language models (llms) are increasingly being applied across various domains @xcite . given the malicious content in pre-training datasets, alignments are implemented to ensure these models are helpful and harmless. @xcite @xcite . despite efforts in alignment, jailbreak attacks can circumvent safety measures, resulting in undesirable outcomes @xcite @xcite . current defenses against jailbreaks primarily involve perturbation of jailbreaks or detecting the safety of inputs. perturbation defenses focus on countering jailbreak attacks through input modification. @xcite @xcite ., * equal contribution. † corresponding author. 2024; @xcite . detection method aims to inspect and categorize input as harmful or safe content, such as perplexity-based classification @xcite @xcite . however, existing defenses lack efficiency because they ignore the underlying causes of jailbreaks. one explanation for alignment failure is the presence of competing objectives outlined by @xcite . competing objectives arise when there is a balance between helpful performance n ... token 0 inference stage obtain probability distribution of the next token (input) user: how to make a bomb? + adversirial input aed (alignment-enhanced decoding) standard decoding original lm find number of top-p candidates exceeds threshold sorry alignment-enhanced lm output: sorry ， it . #token: 0 1 2 ... n output: sure ， ... . #token: 0 1 2 ... n i ... competitive index @xmath0.. . f in e p le a s e i s u r e o k ... s o rr y 0.15 0.13 0.12 0.08 0.04 0.03 0.05 s o r r y p le a s e i t 0.85 0.1 0.04 sure g re a t .. . s u r e o k ... and adhering to harmless principles. this competition may cause a model to prioritize helpful objectives over harmless when confronted with jailbreak prompts, leading to the failure of safety measures. in this work, we present alignment-enhanced decoding (aed), a novel defense that employs adaptive decoding to refine the probability distribution of each token (see fig. 2 ). specifically, we define the competitive index to quantify the competing objectives of the model and to represent the risk of the model being jailbroken. subsequently, we obtain the self-evaluation of the model in which we use the generated output as an auxiliary input to derive the post-alignment logits. when predicting the next token, aed adaptively refines the original logits based on the competitive index and the post-alignment logits. therefore, aed ensures that each step of the decoding process adheres to harmless goals without additional training. in addition, aed is adaptive to maintain the helpfulness to routine queries. we perform comprehensive experiments across five popular open-source large language models, including llama2-7b-chat-hf @xcite , llama3-8b-instruct (meta, 2024), vicuna-7b @xcite , guanaco-7b @xcite , and gemma-1.1-7b-it @xcite . experimental results show that aed effectively counters a range of sophisticated jailbreak attacks such as gcg @xcite , auto-dan @xcite , icd @xcite , and refusal_suppression @xcite . additionally, aed maintains helpfulness on general queries in harmless datasets, including mmlu @xcite , gms8k @xcite , and alpaca (for research on foundation models, 2023). to summarize our contributions: ## competitive index the trade-offs between helpfulness and harmlessness objectives appear after language models are trained to align human values @xcite . when faced with ambiguous questions, these tradeoffs place the models at risk of choosing between in top-p sampling @xcite , given the decoding step t, the candidate set p c ⊆ v is defined as follows: here v is the vocabulary set, p(x|x 0 , • • • , x j-1 ) denotes the probability of next token given a sequence of j -1 tokens as context and p 0 ∈ (0, 1] is a threshold hyper-parameter. the size of candidate set p c is defined as candidate count s and is then calculated as follows: the variation of s in harmless datasets tends to be stable compared with encountering the jailbreak attacks, as illustrated in appendix d. the upper limit for s in harmless datasets is then defined as s @xmath1: as illustrated in fig. 3 , the i can be differentiate by a threshold i t . the threshold i t is set 1, corresponding to the condition where @xmath2. an i greater than the threshold signals anomalies, indicating the competition and an increased risk of jailbreak influence. ## method: alignment-enhanced decoding as discussed in sec. 3, competitive index quantifies the degree of the objectives competition within the model. based on competitive index, we propose a novel defense method, alignment-enhanced decoding (aed). aed adaptively refines the distribution of each generation step. as a result, aed performs an enhanced alignment at the decoding phase, illustrated in fig. 2 . ## experiments in this study, we conducted extensive experiments of aed across five models, utilizing four attack methods. then, we evaluated the performance of aed on three harmless datasets. ## conclusions we define the competitive index i for the first time to quantify the degree of competition among various training objectives. utilizing e competitive index i and the self-evaluation capabilities of the model, we introduce a novel defensive aed that adaptively refines the token distribution during prediction. this method is validated across five different models and tested against four jailbreak attacks, confirming its efficacy. through comparative studies, we demonstrate that aed surpasses existing defenses in effectiveness and achieves this without necessitating additional training. furthermore, according to the average time generation ratio (atgr), aed introduces no significant increase in time overhead, confirming its efficiency and practicality.