Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Tanqiu Jiang, Zian Wang, Jiacheng Liang, Changjiang Li, Yuhui Wang, Ting Wang

TL;DR

By strategically evicting the KVs of the lowest-ranked tokens, RobustKV diminishes the presence of the harmful query in the KV cache, thus preventing the LLM from generating malicious responses, and contributes to RobustKV's robustness against adaptive attacks.

Abstract

Jailbreak attacks circumvent LLMs' built-in safeguards by concealing harmful queries within jailbreak prompts. While existing defenses primarily focus on mitigating the effects of jailbreak prompts, they often prove inadequate as jailbreak prompts can take arbitrary, adaptive forms. This paper presents RobustKV, a novel defense that adopts a fundamentally different approach by selectively removing critical tokens of harmful queries from key-value (KV) caches. Intuitively, for a jailbreak prompt to be effective, its tokens must achieve sufficient `importance' (as measured by attention scores), which inevitably lowers the importance of tokens in the concealed harmful query. Thus, by strategically evicting the KVs of the lowest-ranked tokens, RobustKV diminishes the presence of the harmful query in the KV cache, thus preventing the LLM from generating malicious responses. Extensive evaluation using benchmark datasets and models demonstrates that RobustKV effectively counters state-of-the-art jailbreak attacks while maintaining the LLM's general performance on benign queries. Moreover, RobustKV creates an intriguing evasiveness dilemma for adversaries, forcing them to balance between evading RobustKV and bypassing the LLM's built-in safeguards. This trade-off contributes to RobustKV's robustness against adaptive attacks. (warning: this paper contains potentially harmful content generated by LLMs.)

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

TL;DR

By strategically evicting the KVs of the lowest-ranked tokens, RobustKV diminishes the presence of the harmful query in the KV cache, thus preventing the LLM from generating malicious responses, and contributes to RobustKV's robustness against adaptive attacks.

Abstract

Jailbreak attacks circumvent LLMs' built-in safeguards by concealing harmful queries within jailbreak prompts. While existing defenses primarily focus on mitigating the effects of jailbreak prompts, they often prove inadequate as jailbreak prompts can take arbitrary, adaptive forms. This paper presents RobustKV, a novel defense that adopts a fundamentally different approach by selectively removing critical tokens of harmful queries from key-value (KV) caches. Intuitively, for a jailbreak prompt to be effective, its tokens must achieve sufficient `importance' (as measured by attention scores), which inevitably lowers the importance of tokens in the concealed harmful query. Thus, by strategically evicting the KVs of the lowest-ranked tokens, RobustKV diminishes the presence of the harmful query in the KV cache, thus preventing the LLM from generating malicious responses. Extensive evaluation using benchmark datasets and models demonstrates that RobustKV effectively counters state-of-the-art jailbreak attacks while maintaining the LLM's general performance on benign queries. Moreover, RobustKV creates an intriguing evasiveness dilemma for adversaries, forcing them to balance between evading RobustKV and bypassing the LLM's built-in safeguards. This trade-off contributes to RobustKV's robustness against adaptive attacks. (warning: this paper contains potentially harmful content generated by LLMs.)

Paper Structure

This paper contains 22 sections, 1 theorem, 8 equations, 5 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background
  4. Threat Model
  5. Token Importance
  6. Method
  7. Observation
  8. RobustKV
  9. Implementation
  10. Evaluation
  11. Experimental Setting
  12. Q1: Defense Effectiveness
  13. Q2: LLM Utility
  14. Q3: Ablation Study
  15. Q4: Adaptive Attacks
  16. ...and 7 more sections

Key Result

Proposition 1

Let us define: $n$ as the total number of tokens, $k$ as the number of important tokens, $m$ as the number of attention layers. For a given layer, we define $p_\mathrm{t}$ and $p_\mathrm{f}$ as the probability of correctly selecting an important token and mistakenly selecting a non-important token, where $\delta \in (0, 1)$ is a properly chosen constant (e.g., $\delta \leq \frac{p_\mathrm{t} - p_

Figures (5)

  • Figure 1: Illustration of jailbreak attacks and RobustKV.
  • Figure 2: Rankings of tokens in harmful queries and jailbreak prompts (in ascending order of importance scores).
  • Figure 3: Impact of eviction rate and eviction randomness on RobustKV.
  • Figure 4: RobustKV's response to adaptive attacks that duplicate harmful queries.
  • Figure 5: RobustKV's response to adaptive attacks that manipulate token importance.

Theorems & Definitions (1)

  • Proposition 1