LightDefense: A Lightweight Uncertainty-Driven Defense against Jailbreaks via Shifted Token Distribution

TL;DR

LightDefense tackles jailbreak defenses for white-box LLMs by shifting the initial token distribution toward a safety-oriented direction, without requiring training or auxiliary models. It identifies a safety direction from the contrast between safe and unsafe token distributions and applies a decoded-time distribution shift, controlled by an uncertainty-based strength parameter. The method achieves strong defense across multiple jailbreak attacks and models while preserving benign utility, evidenced by near-zero attack success rates and high helpfulness benchmarks, and it remains computationally efficient. This approach offers a practical, resource-efficient defense that can be deployed at inference time to bolster AI safety without sacrificing usefulness.

Abstract

Large Language Models (LLMs) face threats from jailbreak prompts. Existing methods for defending against jailbreak attacks are primarily based on auxiliary models. These strategies, however, often require extensive data collection or training. We propose LightDefense, a lightweight defense mechanism targeted at white-box models, which utilizes a safety-oriented direction to adjust the probabilities of tokens in the vocabulary, making safety disclaimers appear among the top tokens after sorting tokens by probability in descending order. We further innovatively leverage LLM's uncertainty about prompts to measure their harmfulness and adaptively adjust defense strength, effectively balancing safety and helpfulness. The effectiveness of LightDefense in defending against 5 attack methods across 2 target LLMs, without compromising helpfulness to benign user queries, highlights its potential as a novel and lightweight defense mechanism, enhancing security of LLMs.

Figures (5)

• Figure 1: (a) is a comparison of defense methods. Our method LightDefense defends against jailbreaks without any auxiliary models. (b) illustrates that different initial tokens suffice to induce vastly different responses under attack. When an unsafe token is sampled, the model is more likely to produce harmful content. Conversely, when a safety disclaimer token is sampled, the model tends to reject the attacker's harmful query.
• Figure 2: Visualization of Qwen3’s generated token representations using 2-dimensional PCA. Left: Safe and unsafe responses can be naturally distinguished, whose boundary (grey dashed line) can be easily fitted by logistic regression using responses’ harmfulness as labels. The difference vector (grey arrow) represents the safety-oriented direction. Right:LightDefense moves responses' representations towards the safety-oriented direction (red arrow for harmful queries and blue arrow for harmless ones). Q represents query and A represents answer.
• Figure 3: This figure illustrates the detail of LightDefense. During Step I, we identify the safety-oriented direction by utilizing the difference vector of token distributions between safe and unsafe responses. During Step II, we adjust token probability distribution by shifting token distribution along this direction to amplify the probabilities of safety disclaimer tokens. Additionally, we introduce LLM's uncertainty for given prompts as defense strength to adjust the shifting weight towards safety.
• Figure 4: (a) shows the negative correlation between uncertainty scores and harmfulness levels of queries, with a Pearson correlation coefficient of $-0.86$. The red dashed vertical line indicates the optimal threshold ($\tau$) for classifying harmless and harmful queries, determined by maximizing F1 score. (b) are predictions from running the uncertainty-based classifier on test data. The classifier is highly effective at separating harmless prompts from harmful attacks.
• Figure 5: The figures above present an ablation analysis of the effects of hyperparameters $\beta$, $m$, $k$, and $\tau$ on Qwen3 using the XSTest benchmark. We observe that LightDefense is insensitive to $\beta$, $m$ and $k$ when $\beta \geq 3$, $m \geq 2$, and $k \geq 3$. However, the selection of $\tau$ is critical for the balance between safety and helpfulness in LightDefense.