Token Highlighter: Inspecting and Mitigating Jailbreak Prompts for Large Language Models

TL;DR

Token Highlighter addresses jailbreak threats to aligned LLMs by locating jailbreak-critical tokens through the Affirmation Loss gradient and attenuating their influence with Soft Removal, yielding a highlighted user query that guides the model toward safe refusals. The method demonstrates strong defense across six jailbreak attacks on two protected models (LLaMA-2-7B-Chat and Vicuna-7B-V1.5) while preserving benign performance as measured by AlpacaEval, and it remains computationally efficient by requiring only a single forward/backward pass for token selection. It further analyzes the trade-offs between defense strength and utility, evaluates running-time advantages over smoothing/detector baselines, and tests resilience to adaptive attacks, showing robust performance and interpretability. The work also shows compatibility with prompt-engineering defenses, such as Self Reminder, enabling additive defense strategies for practical deployment.

Abstract

Large Language Models (LLMs) are increasingly being integrated into services such as ChatGPT to provide responses to user queries. To mitigate potential harm and prevent misuse, there have been concerted efforts to align the LLMs with human values and legal compliance by incorporating various techniques, such as Reinforcement Learning from Human Feedback (RLHF), into the training of the LLMs. However, recent research has exposed that even aligned LLMs are susceptible to adversarial manipulations known as Jailbreak Attacks. To address this challenge, this paper proposes a method called Token Highlighter to inspect and mitigate the potential jailbreak threats in the user query. Token Highlighter introduced a concept called Affirmation Loss to measure the LLM's willingness to answer the user query. It then uses the gradient of Affirmation Loss for each token in the user query to locate the jailbreak-critical tokens. Further, Token Highlighter exploits our proposed Soft Removal technique to mitigate the jailbreak effects of critical tokens via shrinking their token embeddings. Experimental results on two aligned LLMs (LLaMA-2 and Vicuna-V1.5) demonstrate that the proposed method can effectively defend against a variety of Jailbreak Attacks while maintaining competent performance on benign questions of the AlpacaEval benchmark. In addition, Token Highlighter is a cost-effective and interpretable defense because it only needs to query the protected LLM once to compute the Affirmation Loss and can highlight the critical tokens upon refusal.

Figures (5)

• Figure 1: Overview of Token Highlighter. (a) The top panel illustrates the concept of LLM jailbreaks by presenting examples of two types of jailbreak prompts (token-level jailbreak by GCG gcg and sentence-level jailbreak by TAP tap. (b) The bottom left panel explains how Token Highlighter finds the jailbreak-critical tokens and mitigates the potential jailbreak effects. We define a loss function called $\mathtt{Affirmation\text{~}Loss}$ to measure the model's willingness to generate affirmative responses to the user query. In step 1, our method selects a set of tokens in the user query that have a large influence on generating the affirmation. In step 2, our method applies Soft Removal on these tokens by shrinking the embeddings of these tokens. We call the user query modified by Soft Removal the Highlighted User Query. The bottom right panel demonstrates that Token Highlighter can inspect suspicious tokens and help the LLM to correctly refuse malicious user queries.
• Figure 2: Performance evaluation on Vicuna-7B-V1.5 (a) and LLaMA2-7B-Chat (b). The horizon axis represents the Attack Success Rate (ASR) averaged over 6 jailbreak attacks, and the vertical axis shows the Win Rate on Alpaca Eval of the protected LLM when the corresponding defense is deployed. Complete results can be found in Appendix \ref{['subapp:complete']}.
• Figure 3: Trade-off between Win Rate and Attack Success Rate by adjusting the values of $\alpha$ and $\beta$.
• Figure 4: Running time analysis. We select 25 queries from AlpacaEval dataset to evaluate the wall clock time of all defenses on Vicuna-7B-V1.5. The printed value for each marker is the running time averaged across the 25 samples. Larger size of a marker means lower running time cost.
• Figure 5: Token Highlighter against adaptive attacks.