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

1. Problem Statement

Large Language Models (LLMs) remain vulnerable to adversarial jailbreak attacks (e.g., GCG, priming attacks) and backdoor attacks (e.g., Sleeper Agents), which can bypass safety training and elicit harmful outputs. Existing defense mechanisms suffer from significant limitations:

• Fine-tuning: Repeatedly fine-tuning models to patch vulnerabilities is computationally expensive, degrades model capabilities, and often fails to eliminate deceptive behaviors.
• Binary Classifiers: Traditional classifiers that detect harmful inputs are brittle; adversaries can craft inputs that fall outside the learned decision boundary yet still trigger harmful outputs.
• Generalization: Defenses often fail to generalize across different attack classes or model families without retraining.

The core challenge is to develop a defense that is computationally lightweight, requires no harmful training data, does not modify the base model architecture, and effectively projects anomalous (potentially harmful) representations back into a safe region.

2. Methodology: MANATEE

The authors propose MANATEE, an inference-time defense mechanism that operates entirely in the representation space (specifically the final hidden layer of the LLM). Instead of asking "Is this input harmful?", MANATEE asks "How can this representation be projected onto the manifold of benign representations?"

The approach treats safety as a density estimation problem solvable via diffusion models.

Core Components:

1. Benign Manifold Learning:

   • The system extracts final-layer hidden states from benign model completions (no harmful data is used).
   • These vectors are standardized (zero mean, unit variance).
   • A Denoising Diffusion Probabilistic Model (DDPM) is trained on these benign vectors to learn the score function (gradient of the log-density) of the benign representation manifold. The denoiser is a lightweight 3-layer MLP.

2. Inference-Time Detection & Steering:

   • Anomaly Scoring: When a user query is processed, the model's final hidden state ($h$) is standardized. The diffusion model predicts the noise required to reconstruct this state from a noisy version at a fixed timestep ($t_{check}$).
   • Score Calculation: The anomaly score $s(h)$ is the magnitude of the predicted noise ($\|\epsilon_\theta(\tilde{h}, t_{check})\|_2$).
     • Low Score: The state lies within the benign manifold (safe).
     • High Score: The state is out-of-distribution (potentially harmful/backdoored).
   • Decision Logic:
     • If $s(h) > \tau$ (Threshold): The system triggers an automatic refusal (e.g., "I cannot comply").
     • If $s(h) \le \tau$: The system applies Diffusion Steering. It injects noise into the hidden state and iteratively denoises it using the trained diffusion model. This process "purifies" the representation, projecting it toward high-density regions of the benign manifold before it is passed to the language model head for token generation.

3. Key Contributions

• Inference-Time Defense: MANATEE operates without modifying the base LLM weights or requiring retraining of the main model. It is a "plug-and-play" module.
• No Harmful Data Required: The defense is trained exclusively on benign data, avoiding the ethical and logistical issues of collecting harmful training examples.
• Manifold Projection: Unlike binary classifiers that simply reject inputs, MANATEE actively corrects anomalous representations, steering them back to safe regions while preserving semantic utility for benign inputs.
• Cross-Model Transferability: The method is trained on one model's benign states and can be applied to defend different model families (e.g., Mistral, Llama, Gemma) without retraining the diffusion model.

4. Experimental Results

The authors evaluated MANATEE on three base models (Mistral-7B-Instruct, Llama-3.1-8B-Instruct, Gemma-2-9B-it) across three datasets:

1. Mechanistic Anomaly Detection (MAD): Llama3-deployment-backdoor.
2. JailbreakBench (JBB): Standard jailbreak attacks.
3. Anthropic Sleeper Agent (ASA): Deceptive backdoor attacks.

Key Findings:

• Attack Success Rate (ASR) Reduction:
  • MAD: Average reduction of 58.7%.
  • JBB: Average reduction of 77.6%.
  • ASA: 100% reduction (complete elimination of successful attacks) across all tested models.
  • Specific examples show ASR dropping from 98% to 0% on Mistral for JBB, and 100% to 0% for ASA.
• Preservation of Utility: For benign inputs (where $s(h) \le \tau$), the diffusion steering step is either skipped or has negligible impact, ensuring the model's helpfulness and reasoning capabilities remain intact.
• Qualitative Results: In case studies, backdoored models that previously generated harmful instructions (e.g., how to smuggle drugs, create fake prescriptions, or falsely accuse others) were successfully steered to refuse the request or provide safe, generic responses.

5. Significance and Conclusion

MANATEE represents a paradigm shift in LLM safety from reactive filtering to proactive representation correction. By leveraging diffusion models to define a "safe" latent space, it offers a robust defense against both known and novel adversarial attacks without the computational overhead of fine-tuning.

The work demonstrates that safety and capability need not be at odds; a lightweight, inference-time mechanism can effectively neutralize backdoors and jailbreaks while maintaining the model's utility on benign tasks. This approach provides a scalable, transferable, and ethically sound solution for securing deployed LLMs in critical infrastructure and consumer products. Future work aims to extend this to latent-based reasoning models (e.g., Chain of Thought) and a broader variety of attack vectors.