One Head to Rule Them All: Amplifying LVLM Safety through a Single Critical Attention Head¶
🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link
Ask a follow-up
Open an assistant pre-loaded with this paper's context.
TL;DR¶
LVLMs' safety capabilities are not evenly distributed across attention heads — a tiny subset of heads (as few as one) strongly responds to unsafe inputs while leaving safe-input representations intact. The authors exploit this by masking a single "safety-critical" head, measuring the angular shift in final-layer hidden states (deflection angle), and using that signal as a training-free jailbreak detector. Upon detection, a safety prompt with anchored output tokens is appended to steer the model toward refusal.
Problem¶
LVLMs inherit safety-aligned weights from base LLMs but exhibit degraded safety alignment in multimodal contexts. Existing defenses — safety fine-tuning, input purification (e.g., BlueSuffix), and post-processing (e.g., ECSO) — are resource-intensive, introduce high inference latency (minutes of overhead), and still fail against sophisticated jailbreaks. The core gap is that models already possess latent safety capability that goes unused rather than being amplified.
Method¶
Safety head identification. For each attention head (l, i), the Query weight matrix is scaled by ε = 10⁻⁵ (effectively suppressing it), then the deflection angle between the principal hidden-state vectors of a small safe dataset D_safe and unsafe dataset D_unsafe is computed at the final layer. A head is a safety-head candidate if (1) its deflection angle on D_safe is below the 80th percentile across all heads, and (2) its differential deflection angle Δθ = θ_unsafe − θ_safe ranks in the top-k = 5 among qualifying heads. One head is randomly (or via validation) selected from this candidate set.
Input detection. For a new input, the deflection angle at the selected critical head is computed per-sample using the full hidden-state vector. A threshold θ is determined by kernel density estimation on deflection angle distributions of D_safe and D_unsafe; inputs exceeding θ are flagged as unsafe.
Response strategy. Detected unsafe inputs trigger a safety prompt appended to the model context, with the first output tokens anchored to "I cannot", forcing the autoregressive generation onto a refusal trajectory and resisting prompt-injection overrides.
Key Contributions¶
- Systematic head-level analysis showing LVLM safety signals concentrate in a sparse subset of attention heads, identifiable via masking-induced deflection angles.
- Oh Defense: a two-stage, training-free framework (head-based detection + prompt anchoring) requiring no auxiliary models, fine-tuning, or extra parameters.
- Empirical validation across four LVLMs (LLaVA-v1.5-7B, Qwen2-VL-7B, Aya-Vision-8B, Phi-3.5-Vision) and four attack datasets (VLSafe, QR, FigStep, LLM Transfer), demonstrating that one safety head suffices and additional heads degrade performance.
- Ablation showing Query-matrix masking produces larger, cleaner deflection signals than Value-matrix masking, and that final-layer hidden states (beyond layer ~20) are optimal monitoring points.
Results¶
- Detection rates (DR): >99% on VLSafe, QR, FigStep, and LLM Transfer for most models; false positive rate 0–9.62% on Instruct-80k and 0–1.40% on ShareGPT4V.
- ASR (LLaVA-v1.5-7B): Oh Defense average ASR = 0.11% vs. HiddenDetect 3.54%, SAHs 21.17%, BlueSuffix 23.26%, ECSO 21.31%, no defense 53.22%.
- ASR (Qwen2-VL-7B): Oh Defense average ASR = 1.55% vs. SAHs 2.71%, HiddenDetect 6.70%.
- ASR (Aya-Vision-8B): Oh Defense average ASR = 1.71% vs. HiddenDetect 4.01%.
- ASR (Phi-3.5-Vision): Oh Defense average ASR = 0.69% vs. HiddenDetect 1.31%.
- FigStep (typographic jailbreak) is driven to 0% ASR on LLaVA, Qwen2-VL, and Aya-Vision by Oh Defense, whereas SAHs and BlueSuffix leave it at 41.85% and 30.25% respectively on LLaVA.
- Using top-3 or top-5 safety heads (instead of top-1) degrades or does not improve DR, confirming the single-head sufficiency claim.
Limitations¶
- The critical head is model-dependent; identification must be re-run for each distinct LVLM.
- Calibration requires a small labeled safe/unsafe dataset (D_safe, D_unsafe) per model.
- False positive rate reaches up to 9.62% on Instruct-80k for Qwen2-VL-7B, meaning a non-trivial fraction of benign inputs are incorrectly flagged.
- Risks of misuse for overconfident or fully automated content filtering without human oversight are not explored.
- Adaptive white-box attacks that specifically target the identified safety head are not evaluated.
Relevance to Vision-Language Models¶
This paper reveals a mechanistic property of LVLM safety: alignment signals are not diffuse but localized in a handful of attention heads, with a single head often sufficient for robust detection. This mechanistic lens — using deflection angles in hidden states as a safety probe — provides interpretable, lightweight access to the model's internal safety representation without external modules, making it directly relevant to anyone studying LVLM robustness, interpretability, or alignment. The finding that visual input degrades text-LLM safety alignment yet the underlying safety structure persists as a recoverable internal signal has broad implications for how multimodal fine-tuning affects safety geometry in the residual stream. It also establishes a strong inference-time defense baseline that future VLM safety work should compare against.