Alignment-Enhanced Decoding:Defending via Token-Level Adaptive Refining of Probability Distributions

TL;DR

This work tackles jailbreak vulnerabilities in large language models by addressing the root cause: competing objective pressures during generation. It introduces Alignment-Enhanced Decoding (AED), which uses a defined Competitive Index $I = \frac{S}{S_t}$ and the model’s self-evaluation to produce post-alignment logits $\mathbf{L}_{post}$; these are adaptively blended with the original logits $\mathbf{L}_{model}$ to yield a safe yet helpful token distribution. The method requires no additional training and is validated across five open-source models against four jailbreak attacks, showing strong defense performance while preserving normal functionality. Practically, AED offers an efficient decoding-time defense that can be deployed with minimal overhead to improve safety in real-world deployments.

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 AED 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. Code is available at https://github.com/GIGABaozi/AED.git.

Figures (12)

• Figure 1: Overview of AED: This diagram illustrates the impact of AED on the token probability distribution. The distribution for harmless queries remains unchanged (left), whereas the distribution for malicious queries undergoes correction (right).
• Figure 2: Pipeline of the decoding process depicted with and without AED intervention, addressing the same harmful query: the top sequence demonstrates standard decoding, while the bottom sequence illustrates the AED process: Step 1 involves obtaining the probability distribution of the next token; Step 2 computes the Competitive Index, which reflects the degree of competitions; and Step 3 realigns the distribution to ensure a safe and ethical response.
• Figure 3: Probability density distributions of the Competitive Index for the Vicuna-7B across five datasets. Harmless datasets are represented in green, while the jailbreaks are represented in orange. The threshold $I_t$ is set at 1. For clarity, data are preprocessed by capping indices exceeding twice the threshold at this upper limit.
• Figure 4: These figures display the probability density distributions of the Competitive Index $I$ for three harmless datasets and two jailbreaks across various models. The charts highlight the differences in Competitive Index between harmless and jailbreak inputs. For clarity, we preprocess the data by capping all indices exceeding twice the threshold at this upper limit. Further details are illustrated in Appendix \ref{['appendix: index']}.
• Figure 5: This graph illustrates the probability density distributions of the Competitive Index $I$with and without system prompts across five models. The inclusion of system prompts leads to a noticeable shift of the Index toward zero, indicating a decrease in the degree of competition.