Jailbreak Antidote: Runtime Safety-Utility Balance via Sparse Representation Adjustment in Large Language Models

Authors: Guobin Shen, Dongcheng Zhao, Yiting Dong, Xiang He, Yi Zeng
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Large Language Models, Jailbreak Defense, Safety-Utility Balance, Internal State Manipulation, Sparse Representation Adjustment

Abstract:
As large language models (LLMs) become integral to various applications, ensuring both their safety and utility is paramount. Jailbreak attacks, which manipulate LLMs into generating harmful content, pose significant challenges to this balance. Existing defenses, such as prompt engineering and safety fine-tuning, often introduce computational overhead, increase inference latency, and lack runtime flexibility. Moreover, overly restrictive safety measures can degrade model utility by causing refusals of benign queries. In this paper, we introduce *Jailbreak Antidote*, a method that enables real-time adjustment of LLM safety preferences by manipulating a sparse subset of the model's internal states during inference. By shifting the model's hidden representations along a safety direction with varying strengths, we achieve flexible control over the safety-utility balance without additional token overhead or inference delays. Our analysis reveals that safety-related information in LLMs is sparsely distributed; adjusting approximately *5\%* of the internal state is as effective as modifying the entire state. Extensive experiments on nine LLMs (ranging from 2 billion to 72 billion parameters), evaluated against ten jailbreak attack methods and compared with six defense strategies, validate the effectiveness and efficiency of our approach. By directly manipulating internal states during reasoning, *Jailbreak Antidote* offers a lightweight, scalable solution that enhances LLM safety while preserving utility, opening new possibilities for real-time safety mechanisms in widely-deployed AI systems.

Decision: accept (poster)
Reviewer scores: 5, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1729998734496

thread_type:
official_review

<thread_object>
creator: reviewer_fzr8

<content>
summary:
This paper proposes a defense method against LLM jailbreaking attacks, which operates at model inference-time by manipulating a sparse subset of the model's representations (internal states). Specifically, the internal states are moved towards pre-computed safe directions, which make the model's responses safer against harmful requests. This method enables better model safety without leading to significant utility drop.

strengths:
* This proposed methodology to defend against jailbreaking attacks is neat. The paper is well structured and written.
* The finding of the sparsity (5\%) of safety representations is intriguing.
* The results are promising -- largely diminishing the success rate of jailbreaking attack while maintaining high utility, which surpass the compared baseline defenses.

weaknesses:
* My major concern is the incremental contribution of this work to [1] and [2]. For example, in [1], the authors already provided extensive and thorough study regarding representation engineering, and discussed methods to adjust model behaviors against LLM jailbreak attacks. These existing work largely diminish the contribution of this paper.

* In line 78-79, you argued that safety alignment methods "lacks real-time flexibility." But on the other hand, you also did not highlight the benefit of "real-time flexibility" on your method. Then this point should not be emphasized too much.

* Should include more defense baseline comparison. I appreciate that the authors considered 6 existing jailbreak defense baselines which operate mostly at inference time, but I think it would be beneficial to consider safety alignment defenses as well, e.g., [3] and [4].

  In Line 79, you mentioned that these methods "*may* reduce model utility," but you didn't compare them with your method throughout the paper. Do they?

* In "Inference Efficiency Analysis" (page 8), the authors used "the number of defense tokens required" as a metric to compare method efficiency. This is clearly a bad metric for efficiency study -- compared to additional output tokens, additional input tokens induce less significant overhead. Thus, it's not clear whether the proposed method is indeed more efficient than the baseline methods. I suggest the authors use actual GPU time instead of number of defense tokens as the metric here.

  Also, could you elaborate on the (pre-inference) time cost to obtain the safety representations and masks?

* There should be some adaptive analysis against the proposed defense. That is, you should consider if an adversary is already aware that the victim model is deployed with your proposed defense. Then, if the adversary explicitly optimize for an adversarial prompt $S'$ against your safety representations, whehter / how much can this adaptive attack be defended by Jailbreak Antidote? This is an necessary study that should come alongside the proposed defense.

[1] Representation Engineering: A Top-Down Approach to AI Transparency, Arxiv 2023

[2] Steering Language Models With Activation Engineering, Arxiv 2024

[3] Improving Alignment and Robustness with Short Circuiting, NeurIPS 2024

[4] Safety Alignment Should Be Made More Than Just a Few Tokens Deep, Arxiv 2024

questions:
* I think there need to be more specifications and analysis on Fig 2(a) and Sec 4.2. For example, how is this figure plotted? And how did you come to the conclusion that "the safety representation in LLMs is sparse"?

  From my understanding, most entries in each $\textbf{d}^l_{safe}$ are almost 0, which indicates *sparsity*. Then, what's the principle behind masking these safety representations (Eq 4)? Since most entries are already 0, in Eq 5, it doesn't make a difference whether you mask them or not, right?

* Are the DSP values reported in Table 1 averaged over all jailbreaking methods?

* Why the method works better for larger models? I see there are some discussions on this. But could you also perform some more in-depth studies (e.g., analyzing the representation distribution difference between small and large models) to provide more understanding about this?

* Could you also discuss the relationship of your observation of sparsity in safety representations, with [1]'s observation of sparsity in model weights that account for safety?

[1] Assessing the brittleness of safety alignment via pruning and low-rank modifications, ICML 2024
</content>
</thread_object>

</thread>
<thread>
creation date:
1730005227908

thread_type:
official_review

<thread_object>
creator: reviewer_gjnv

<content>
summary:
The paper propose a new defense against jailbreak attacks, called Jailbreak Antidote. It adjusts internal hidden state of LLMs during inference, with parameters controlling the tradeoff between safety and utility. Extensive experiments prove the performance of the method.

strengths:
* The paper is looking at a well-defined hot security topic, the jailbreak problem.
* The paper is well written and easy to follow.
* The paper has a comprehensive evaluation, covering various state-of-the-art models, attack methods, mitigation baselines, etc., as well as ablation studies.

weaknesses:
* The evaluation of false positive rate of safety blocking is missing in evaluation. See questions.
* I did not find implementation details of the PCA step.

questions:
The paper is in a good shape and I enjoyed reading it. I have several questions on evaluation results.

The evaluation uses defense success rate against jailbreak attacks and win rate on benign prompts to demonstrate the balance between safety and utility. However, the defense success rate only represent the true positive rate of safety blocking, while false positive rate is missing. In other words, the statement in Section 4.3 "A higher \alpha emphasizes safety, making the model more conservative in its responses but potentially affecting utility by **increasing the refusal of borderline benign prompts**." is not validated by experiments. I would admit that the decent win rate on AlpacaEval has demonstrated that such borderline benign prompts is not common, but specific analysis on such failing cases is also meaningful. For instance, how much of the drop on win rate is because of false safety blocking?

Though the modification of internal state at inference stage should be very lightweight, I am unsure about the difficulty of the PCA step generating $d_{safe}$. How much data is used for PCA? How fast is the step especially for 70B large models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730272438541

thread_type:
official_review

<thread_object>
creator: reviewer_eqzp

<content>
summary:
This paper studies an efficient defense, named Jailbreak Antidote, against jailbreak attacks targeting LLMs. The goals of defense design include balancing safety and utility, as well as reducing overhead during inference time. The paper is based on an observation drawn from the representations of harmful and safe prompts. Given the distance between these two representations, the authors modify the hidden state by add the masked distance vector to it to improve safety performance. Evaluations over nine LLMs are presented to demonstrate the results.

strengths:
The tradeoff between LLM safety and utility is important. The paper is well-written and easy to follow. The authors showed some promising experimental results.

weaknesses:
The paper misses comparison or discussion to related work. Additional experimental evaluations on more LLMs and specific downstream tasks should be performed to better support the claims in this paper. See Questions for more detail.

questions:
I have some major comments on scope of related work and experiment design. 
- The design of Jailbreak Antidote reads very similar to SafeDecoding [a], with one modifying the internal state while the other modifying the token distribution. However, I could not find any comparison or discussion with it. It would be great to see the experimental comparisons between these two defenses. 
- Are the observations in A.3 consistent across all models? It is not clear which model was used to generate A.3.
- I wonder how would Jailbreak Antidote perform on uncensored models such as Dolphin-llama.
- The authors should include more attack methods such as [b] to evaluate the effectiveness of Jailbreak Antidote.
- Since utility is an important design goal, the authors should evaluate the model performance on more downstream tasks in addition to Win Rate, such as using WildBench, MMLU, and HellaSwag.
- Jailbreak Antidote seems to be sensitive to the choice of hyper-parameters. Is there any systematic way to choose these parameters, especially $\alpha$, for different models?

Minor comments:
- Although $d_{safe}^l$ and $m^l$ can be computed offline, it would be good to see the associated computational cost.
- Defense Tokens seem undefined.
- Is length controlled win rate also evaluated?
- In the notation of $h_S^l$, the subscript is not clear. The paper sometimes uses prompt as subscript, sometimes uses T as subscript.

[a] SafeDecoding: Defending against Jailbreak Attacks via Safety-Aware Decoding
[b] DeepInception: Hypnotize Large Language Model to Be Jailbreaker
</content>
</thread_object>

</thread>
<thread>
creation date:
1730596191433

thread_type:
official_review

<thread_object>
creator: reviewer_1gvh

<content>
summary:
This paper proposed a new defense to mitigate jailbreak attacks. The key idea of the proposed defense is to modify the internal state of an LLM during inference. The idea of the method is to extract a sparse safety feature vector and add it to the representation of a prompt. As a result, the model is more likely to refuse the malicious prompts.

strengths:
1. How to defend against jailbreak attacks is still an open challenge in the community.

2. Multiple LLMs are evaluated.

3. The proposed attack does not rely on the attacks. Moreover, the method is efficient as it directly manipulates the internal states of an LLM during its inference.

weaknesses:
1. In general, the compared baselines are very weak. The authors may consider doing a comprehensive survey on the defenses against jailbreak attacks and add a discussion on these defenses. For instance, the authors could check this survey papers: https://arxiv.org/pdf/2407.04295 

2. The proposed method requires harmful prompts to identify safety direction. In the Appendix, it is mentioned that the dataset in Phan 2023 is used. Could authors evaluate the difference between this dataset and the dataset used for evaluation? For instance, whether there are many similar questions in the two datasets.

Additionally, a straightforward baseline is to use preference learning to finetune an LLM on this dataset. The authors may also consider comparing with this baseline in some settings (e.g., for small LLMs). 

3. Adaptive attacks are not considered. If an attacker knows the safety vector (or a surrogate one optimized by the attacker), could an attacker perform adaptive attacks to evade the defense?

4. Only one benchmark dataset and 100 prompts are used in the evaluation. The authors may consider evaluating the effectiveness of the defense on more benchmark datasets.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732264494423

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal - Part I/II

comment:
We sincerely thank all the reviewers for their thoughtful feedback and constructive comments on our manuscript titled "Jailbreak Antidote: Mitigating Jailbreak Attacks on LLMs via Sparse Runtime Adjustment of Internal States." We are grateful for the recognition of the importance of our work and the valuable suggestions provided to enhance its quality and clarity. We have carefully considered all comments and have revised the manuscript accordingly to address the concerns raised.


Several common themes emerged from the reviews, which we have carefully addressed in our revisions:

1. **Comparison with Additional Baselines and Related Work**: Reviewers suggested including comparisons with additional defense methods, particularly safety alignment defenses and preference-based fine-tuning approaches. In response, we have incorporated experimental comparisons with methods such as preference-based fine-tuning [7, 9, 3], demonstrating that our approach achieves superior defense success rates while maintaining better utility.

2. **Evaluation Against Adaptive Attacks**: Reviewers requested an assessment of the robustness of our method against adaptive attacks. To address this, we conducted experiments against adaptive attacks, including GCG [8], PAIR [4], and AutoDAN [6]. The results show that our method remains robust even when attackers are aware of our defense mechanism.

3. **Clarification on Dataset Similarity**: Concerns were raised about the potential overlap between the dataset used to generate safety directions and the evaluation dataset. We performed a thorough similarity analysis and confirmed that no significant overlap exists, ensuring unbiased evaluation.

4. **Additional Experiments and Analysis**: Reviewers recommended expanding experimental evaluations, covering more downstream tasks, and analyzing hyperparameter effects. In response, we included additional tasks such as MMLU and HellaSwag to assess utility, provided detailed analyses on hyperparameter impacts, and examined trade-offs between safety and utility, as well as false positive rates and efficiency.

5. **Clarification and Detail Enhancements**: Suggestions were made to improve clarity and provide more implementation details. We enhanced key sections for better readability, added implementation details to facilitate reproducibility, and addressed all minor issues, including notation and definitions.



---

**References**

1. Yi, Sibo, et al. "Jailbreak attacks and defenses against large language models: A survey." *arXiv preprint arXiv:2407.04295 (2024)*.

2. Wei, Boyi, et al. "Assessing the Brittleness of Safety Alignment via Pruning and Low-Rank Modifications." *Forty-first International Conference on Machine Learning*.

3. Xu, Zhangchen, et al. "SafeDecoding: Defending against Jailbreak Attacks via Safety-Aware Decoding." *ICLR 2024 Workshop on Secure and Trustworthy Large Language Models*.

4. Chao, Patrick, et al. "Jailbreaking Black Box Large Language Models in Twenty Queries." *R0-FoMo: Robustness of Few-shot and Zero-shot Learning in Large Foundation Models*.

5. Li, Xuan, et al. "Deepinception: Hypnotize large language model to be jailbreaker." *arXiv preprint arXiv:2311.03191 (2023)*.

6. Liu, Xiaogeng, et al. "AutoDAN: Generating Stealthy Jailbreak Prompts on Aligned Large Language Models." *The Twelfth International Conference on Learning Representations*.

7. Qi, X., et al. (2024). "Safety Alignment Should Be Made More Than Just a Few Tokens Deep." *arXiv preprint arXiv:2405.00159*.

8. Zou, A., et al. (2023). "Universal and Transferable Adversarial Attacks on Aligned Language Models." *arXiv preprint arXiv:2406.05946*.

9. Zou, Andy, et al. "Improving Alignment and Robustness with Short Circuiting." *arXiv preprint arXiv:2406.04313 (2024)*.

10. Chen, Yangyi, et al. "Why Should Adversarial Perturbations be Imperceptible? Rethink the Research Paradigm in Adversarial NLP." *Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing. 2022*.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission received the ratings of four reviewers, which recommended 5, 6, 8 and 6, averaging 6.25. Given the plenty of competitive submissions in ICLR, this stands at a score above the borderline. The reviewers' concerns focus on the insufficient baselines, unclear technical parts or experimental setup, and the lack of analysis. After the substantial rebuttal by authors, most of reviewers considered the concerns well addressed, and one reviewer still maintain no response. After carefully checking the reviewer comments and the authors' feedback, the AC agreed with the authors' effort to address the remaining concerns. Therefore, I tend to recommend acceptance towards the current submission, and hope all advice are well incorporated into the final submission.