RobustKV: Defending Large Language Models against Jailbreak Attacks via KV Eviction

Authors: Tanqiu Jiang, Zian Wang, Jiacheng Liang, Changjiang Li, Yuhui Wang, Ting Wang
Primary area: foundation or frontier models, including LLMs
Keywords: Jailbreak Attack, Large Language Model, KV cache optimization

Abstract:
Jailbreak attacks circumvent LLMs' built-in safeguards by concealing harmful queries within adversarial prompts. While most existing defenses attempt to mitigate the effects of adversarial prompts, they often prove inadequate as adversarial prompts can take arbitrary, adaptive forms. This paper introduces RobustKV, a novel jailbreak defense that takes a fundamentally different approach by selectively removing critical tokens of harmful queries from key-value (KV) caches. Intuitively, for an adversarial prompt to be effective, its tokens must achieve sufficient `importance' (measured by attention scores), which consequently lowers the importance of tokens in the concealed harmful query. Therefore, by carefully evicting the KVs of low-ranked tokens, RobustKV minimizes the harmful query's presence in the KV cache, thus preventing the LLM from generating informative responses. Extensive evaluation using benchmark datasets and models demonstrates that RobustKV effectively counters state-of-the-art jailbreak attacks while maintaining the LLM's performance on benign queries. Notably, RobustKV creates an interesting effectiveness-evasiveness dilemma for the adversary, leading to its robustness against adaptive attacks.{(Warning: This paper contains potentially harmful content generated by LLMs.)}

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

Reviews and discussion:
<thread>
creation date:
1730343493865

thread_type:
official_review

<thread_object>
creator: reviewer_mthc

<content>
summary:
The paper  introduces a novel defense mechanism called RobustKV. This method tackles jailbreak attacks on large language models (LLMs) by selectively evicting key-value (KV) pairs of less important tokens from the model's KV cache. RobustKV identifies and removes tokens related to harmful queries by analyzing their attention scores, thus preventing the LLM from generating malicious responses.

strengths:
1. Unlike existing defenses, which focus on mitigating jailbreak prompts, RobustKV directly addresses harmful queries by manipulating the KV cache, offering a fresh perspective on LLM defenses.
2. RobustKV does not increase computational overhead or response time, making it practical for real-world deployment.

weaknesses:
1. In my opinion, since your method needs to access the LLM’s key-value (KV) caches, it may only be applicable to open-source LLMs. I believe this could be a significant limitation for your work.
2. In the example you provided in line 379, 'how to steal someone' is removed, so the attack will certainly not succeed. If the core content of the jailbreak issue has been removed and the LLM can still be attacked, that would be truly strange！In my opinion, this paper provides a meaningless method, and there's no scenario can the method be used.

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

</thread>
<thread>
creation date:
1730500538076

thread_type:
official_review

<thread_object>
creator: reviewer_k7do

<content>
summary:
This paper presents a novel approach to defending large language models (LLMs) from jailbreak attacks. Based on the idea that tokens associated with jailbreak prompts tend to hold higher importance than the actual harmful query, RobustKV evicts critical tokens of the harmful query from the KV cache. Thus it results in a non-informative response to the harmful query to improve jailbreak defense.  Benchmark evaluations under state-of-the-art jailbreak attacks demonstrate this method's effectiveness.

strengths:
1. This work combines LLM inference efficiency and safety, addressing two critical areas for improving LLM performance.

2. Observation and motivation are both clear, I enjoy the insight this paper provides for the adaptive attack design.

3. Algorithm is the improvement of SnapKV and can be applied to any pre-trained LLMs without finetuning.

4. This paper provides comprehensive experiments to support the stability and efficacy of the algorithm.

weaknesses:
1. For the Observation Window, did you try different setting or different size?
2. I notice that in Figure 3, this method achieve a trade-off between eviction ratio and defense effectiveness (clean accuracy and jailbreak defense). For larger LLMs like Llama2-13B, KV cache's performance may be different, would it still be as effective?
3. I'm curious about the choice to use token importance across all layers in the algorithm. What would the impact be if importance were calculated across all heads within a single layer instead?
4. Could you provide why choose these three different models for your algorithm? As I understand, Mistral uses GQA, how do the other two models compare in this experiment setting? Do these selections provide sufficient evaluation of your algorithm’s generalization?

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

</thread>
<thread>
creation date:
1730694582818

thread_type:
official_review

<thread_object>
creator: reviewer_wnh1

<content>
summary:
Safety-tuned LLMs are still jailbroken by advanced techniques. This paper implements a complimentary method to safety-tuning for defending against jailbreaks. 

The paper makes the following observations:
- Jailbreak tokens must have high attention scores to bypass safety
- High jailbreak token attention necessarily reduces harmful query token attention
- This creates a consistent pattern: harmful query tokens rank lower than jailbreak tokens

They implement RobustKV, which evicts the harmful components of the jailbreaking input. This forces attackers into a dilemma:
- Make harmful query important = jailbreak fails to bypass safety
- Make jailbreak important = harmful query gets evicted
- Can't optimize both simultaneously

strengths:
- This work leverages safety-tuning. Unlike the jailbreak defense pipelines mentioned in [1] which constitute defining properties of harmful outputs explicitly, this approach preserves the learnt safety-tuning policy, and does not seek to redefine safe and unsafe behavior. 
- This defense points out that many jailbreak techniques leverage the fact that safety-tuning does not catch when malicious queries are obfuscated. Given the possibility of RobustKV, potentially this attack model will no longer be as interesting to study.


[1] Testing the Limits of Jailbreaking Defenses with the Purple Problem. Kim 2024.

weaknesses:
- This paper needs to better justify that this sort of defense is useful:
  -  [1] ran an extensive comparison across jailbreaking techniques, and found that [2] and [3] are the strongest attack techniques when comparing to all the techniques mentioned in this paper. Either a larger scope of attacks should be defended against using RobustKV, or proper analysis of patterns within the studied jailbreaks should be provided. Given the evolving offense-defense balance, defense techniques cannot be measured against a select set of attacks without a proper argument for it.
  - [4]  relies on in-context learning, which may not be defended against RobustKV. In security, in order to provide adequate evidence of your hypothesis, please seek out counterexamples that might be most challenging. 

- Generating uninformative responses may be a significant challenge to any model developers implementing such methods. 
- The related works section only cites as one black-box jailbreaking technique, please refer to [1] for a list.

[1] A STRONGREJECT for Empty Jailbreak. Souly 2024.
[2] How Johnny Can Persuade LLMs to Jailbreak Them: Rethinking Persuasion to Challenge AI Safety by Humanizing LLMs. Zeng 2024.
[3] Jailbreaking Black Box Large Language Models in Twenty Queries. Chao 2024.
[4] Many-shot jailbreaking. Anil 2024.

questions:
- The LLM Utility section only has data for Llama, what about the others? 
- If utility was not decreased by percentage scores, did you also also analyze the change in distribution?
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors propose RobustKV, a novel defense mechanism aimed at defending large language models (LLMs) from jailbreak attacks, which are attempts to manipulate the model into producing harmful or undesired outputs. The key innovation of RobustKV is the eviction of critical tokens from the key-value (KV) caches during the inference process, which prevents the model from generating specific responses to adversarial prompts. The paper demonstrates through experiments that this approach is effective in mitigating jailbreak attacks without significantly harming the model's overall performance.