Test-Time Immunization: A Universal Defense Framework Against Jailbreaks for (Multimodal) Large Language Models

TL;DR

TIM tackles jailbreak vulnerabilities in both text-only and multimodal large language models by introducing Test-time Immunization, a universal defense that evolves during inference. It detects jailbreak attempts with a lightweight gist token and a binary classifier, and when detected, performs safety fine-tuning via a LoRA module using the identified jailbreak instructions, while decoupling detection from defense updates to prevent interference. The framework leverages online test-time adaptation and memory to continually improve both detection and defense, reducing jailbreak success rates with minimal impact on normal queries. Extensive experiments on LLMs and MLLMs show TIM achieves near-zero attack success, low overhead, and robust performance against diverse jailbreak strategies, highlighting its practical potential as a flexible safety mechanism.

Abstract

While (multimodal) large language models (LLMs) have attracted widespread attention due to their exceptional capabilities, they remain vulnerable to jailbreak attacks. Various defense methods are proposed to defend against jailbreak attacks, however, they are often tailored to specific types of jailbreak attacks, limiting their effectiveness against diverse adversarial strategies. For instance, rephrasing-based defenses are effective against text adversarial jailbreaks but fail to counteract image-based attacks. To overcome these limitations, we propose a universal defense framework, termed Test-time IMmunization (TIM), which can adaptively defend against various jailbreak attacks in a self-evolving way. Specifically, TIM initially trains a gist token for efficient detection, which it subsequently applies to detect jailbreak activities during inference. When jailbreak attempts are identified, TIM implements safety fine-tuning using the detected jailbreak instructions paired with refusal answers. Furthermore, to mitigate potential performance degradation in the detector caused by parameter updates during safety fine-tuning, we decouple the fine-tuning process from the detection module. Extensive experiments on both LLMs and multimodal LLMs demonstrate the efficacy of TIM.

Figures (8)

• Figure 1: The overview of test-time immunization. (1): The LLMs with pre-guarded strategy can defend against some jailbreak attacks successfully, but can't defend against all potential types of jailbreak attacks in advance. (2) We resort to adaptively leveraging test jailbreak data during testing to enhance the defense capabilities of LLMs. When a jailbreak attack successfully hacks our model, we learn the distribution of the jailbreak attack and gradually become immune to it.
• Figure 2: Detailed workflow of test-time immunization. (1) We insert a trainable gist token at the sequence's end and utilize the hidden states from intermediate layers along with a classifier $\mathcal{C}_d$ to perform detection. In a real-world application, we can employ the KV Cache and the gist token to perform efficient detection. (2) Upon detecting jailbreak activity during detection, we append the data to jailbreak memory and incorporate detection data into detection memory for training. Then we utilize jailbreak memory $\mathcal{M}_j$ to train the LLM's defense LoRA module by supervised fine-tuning and employ detection memory $\mathcal{M}_d$ to further train the detector (i.e., TTA) by \ref{['equ:detection']}. Additionally, we employ a question-answering dataset $\mathcal{D}_{qa}$ and a detection dataset $\mathcal{D}_d$ for regularization.
• Figure 3: Performance of different variants of the proposed method. All metrics are normalized, and the methods with larger areas have better performance.
• Figure 4: The detection performance under I-FSJ attack.
• Figure 4: Results under hybrid jailbreak attack. We randomly selected 300 jailbreak samples from MM-SafetyBench and 300 from Figstep, combining them into a new jailbreak dataset.