MANATEE: Inference-Time Lightweight Diffusion Based Safety Defense for LLMs

Chun Yan Ryan Kan; Tommy Tran; Vedant Yadav; Ava Cai; Kevin Zhu; Ruizhe Li; Maheep Chaudhary

cs.CRFeb 21, 2026

MANATEE: Inference-Time Lightweight Diffusion Based Safety Defense for LLMs

Authors: Chun Yan Ryan Kan, Tommy Tran, Vedant Yadav, Ava Cai, Kevin Zhu, Ruizhe Li, Maheep Chaudhary

Abstract

Defending LLMs against adversarial jailbreak attacks remains an open challenge. Existing defenses rely on binary classifiers that fail when adversarial input falls outside the learned decision boundary, and repeated fine-tuning is computationally expensive while potentially degrading model capabilities. We propose MANATEE, an inference-time defense that uses density estimation over a benign representation manifold. MANATEE learns the score function of benign hidden states and uses diffusion to project anomalous representations toward safe regions--requiring no harmful training data and no architectural modifications. Experiments across Mistral-7B-Instruct, Llama-3.1-8B-Instruct, and Gemma-2-9B-it demonstrate that MANATEE reduce Attack Success Rate by up to 100% on certain datasets, while preserving model utility on benign inputs.

Figures & tables

Figure 1: Attack Success Rate (ASR) reduction across various models on JBB dataset, before and after diffusion models are applied on finetuned harmful models.

Figure 2: Overview of MANATEE’s inference-time detection and steering pipeline. Given a user query, the base LLM produces the final-layer hidden state $h_{t}$ . We standardize this representation and apply forward diffusion to obtain a noisy state at a timestep $t_{\mathrm{check}}$ . The denoiser’s predicted noise magnitude defines an anomaly score $s(h)$ that measures how far $h_{t}$ lies from the benign manifold. States with scores above the threshold $\tau$ trigger automatic refusal; otherwise, MANATEE applies a DDPM-based purification step by injecting noise and denoising in latent space to produce a purified hidden state that is passed through the LM head to generate a benign, conditionally steered response.

Dataset	Model	Base ASR	+MANATEE
			ASR	$\Delta$ ASR $\downarrow$
MAD	Mistral	55	11	-44
	Gemma	100	40	-60
	Llama	100	28	-72
JBB	Mistral	98	0	-98
	Gemma	46	0	-46

Table 1: ASR (%, lower is better) on the backdoor dataset. MAD = Mechanistic Anomaly Detection Llama3-deployment-backdoor, JBB = JailbreakBench, ASA = Anthropic Sleeper Agent.

Jailbreak Defense

Jailbreak defense studies how to stop adversarial prompts from eliciting harmful outputs while preserving helpful behavior.

See all 60 papers

Black-box, Adaptive, Efficient, Transferable, Harmful, Applicable... Attacks Are All You Need to Break LLMs
Jun 2, 2026Vincent Limbach, Jonas Dornbusch, David Lüdke +2LLM Jailbreak Attacks LLM Safety Evaluation
Department of Computer Science, Technical University of Munich, Germany · Munich Data Science Institute · Munich Center for Machine Learning +1
Automated jailbreak attack targeting multiple defense strategies
Jun 15, 2026Qi Wang, Chengcheng Wan, Weijia He +4Adversarial Prompt Generation LLM Jailbreak Attacks
East China Normal University Shanghai, China · East China Normal University, Shanghai Innovation Institute Shanghai, China · University of Southampton Southampton, England +2
When LLM Defenses Backfire: Characterizing Safety, Performance, and Cost Trade-offs
Jul 27, 2026Tong Zhang, Zexin Li, Simin Chen +1LLM Inference Efficiency Language Model Safety Evaluation
Fudan University · University of California, Riverside · Columbia University

Jun 2, 2026J/K move · Enter open · S save

Black-box, Adaptive, Efficient, Transferable, Harmful, Applicable... Attacks Are All You Need to Break LLMs

Vincent Limbach, Jonas Dornbusch, David Lüdke, Stephan Günnemann+1

Department of Computer Science, Technical University of Munich, Germany · Munich Data Science Institute · Munich Center for Machine Learning+1

LLM Jailbreak Attacks LLM Safety Evaluation Transfer-Based Adversarial Attacks Black-Box Adversarial Attacks Jailbreak Attack Evaluation

Open paper AskPDF arXiv

Accurately evaluating adversarial robustness is a longstanding challenge. A flawed attack design can inflate robustness estimates, making deployment risk assessment and defense comparison unreliable. Historically, standardized attacks such as AutoAttack have largely resolved this for image classifiers, providing a reliable evaluation baseline for systematic comparison across defenses. However, no equivalent exists for LLM jailbreak evaluation yet, where designing such an attack is considerably more difficult. A reliable attack must, among other things, be black-box compatible, applicable to arbitrary defense pipelines, and efficient, which no existing method jointly satisfies. We introduce Indirect Harm Optimization (IHO), a masked diffusion language model attacker trained via iterative preference optimization against a harmfulness judge, requiring only black-box access to the target. The same method can be used without modification as a strong adaptive attack on individual behaviors, or as an efficient amortized policy that transfers to held-out behaviors and unseen target models without fine-tuning. Even against layered defenses, such as a Circuit Breaker-trained model combined with an auxiliary detector, IHO improves attack success considerably over state-of-the-art approaches, without any defense-specific adaptation. Our results position IHO as a practical step toward the kind of standardized jailbreak evaluation that has improved reliability in the past. Code and models are available on GitHub and Hugging Face.

Abstract

Jailbreak Defense

Explore similar work

Black-box, Adaptive, Efficient, Transferable, Harmful, Applicable... Attacks Are All You Need to Break LLMs

Automated jailbreak attack targeting multiple defense strategies

When LLM Defenses Backfire: Characterizing Safety, Performance, and Cost Trade-offs