How to Make LLMs Safer? Detecting and Editing Key Heads in LLMs

Kuan-Lin Chu, Chung-En Sun, Tsui-Wei Weng
UCSD
NeurIPS 2025 Lock-LLM Workshop
description Paper code Code

Abstract

Ensuring the safety of large language models (LLMs) is crucial. In this work, we investigate the internal attention mechanisms that detect harmful content and refusal behaviors in LLMs. We introduce systematic methods to identify detection heads, which are highly sensitive to harmful prompts, and refusal heads, which contribute to the model’s tendency to reject unsafe requests. Building on these insights, we propose the Detection–Refusal Advanced LLM (DRefA), an enhanced model. Safety is quantified as the proportion of responses judged safe by Llama-Guard-3-8B, referred to as the safety rate. DRefA achieves substantial robustness gains—for instance, the safety rate of LLaMA3 increases from 77% → 99% under GCG attacks and from 15% → 99% under ADV-LLM attacks.

Our Contributions

Issue of Jailbreaking in LLMs

As LLMs become widely used, malicious users often attempt jailbreaking to bypass safety mechanisms. Jailbreaking involves crafting inputs that trick the model into producing harmful, biased, or otherwise unsafe content. Even state-of-the-art LLMs can be vulnerable under adversarial prompts, as shown by their baseline safety rates. Here, the safety rate is measured by Llama-Guard-3-8B, which evaluates whether a model’s response is safe. This highlights the need for interventions to improve robustness and ensure safer outputs.

Figure 2: Safety rates (%) of baseline models under different adversarial attack methods, measured by Llama-Guard-3-8B. All models show vulnerability to certain attacks, highlighting the challenge of LLM jailbreaking.

DRefA: Identifying and Editing Key Safety Heads in LLMs

Figure 1: Overview of Identifying Detection Heads.

Figure 2: Overview of Identifying Refusal Heads.

In transformer-based LLMs, certain attention heads play a disproportionately important role in specific behaviors. Our analysis identifies two types of key heads:

By analyzing harmful-neutral prompt pairs and recording attention and residual contributions, we systematically rank heads and select the top candidates for intervention. This allows us to target safety behaviors while minimally affecting general model performance.

Intervention on Key Heads

After identifying detection and refusal heads, we scale their contributions in the transformer's residual stream:

Figure 3a: Safety improvements from intervention on detection heads. Safety rate measured by Llama-Guard-3-8B.

Figure 3b: Safety improvements from intervention on refusal heads. Safety rate measured by Llama-Guard-3-8B.

Experiment I: Examining Causal Relationship Between Detection and Refusal Heads

Suppressing detection heads in Mistral-7B-Instruct-v0.2 sharply reduces the contributions of refusal heads, showing that detection heads activate refusal heads to enforce safety. This provides causal evidence that detection and refusal heads form a safety circuit in LLMs.

Baseline (scaling factor = 1.0)

Detection heads suppressed (-2.0)

Detection heads suppressed (-4.0)

Figure 4: Refusal head contributions under baseline and detection-head suppression. Reducing detection head activity weakens refusal signals, highlighting their causal relationship.

Experiment II: Reinforcing Detection and Refusal Heads

Based on the causal link from Experiment I, we reinforce the top 3% detection and refusal heads by scaling their contributions (3.0 for detection heads, 2.0 for refusal heads). We evaluate 4 instruction-tuned LLMs using harmful prompts and two attack methods (GCG and ADV-LLM), measuring safety with Llama-Guard-3-8B.

Scaling detection or refusal heads individually improves safety, but the strongest results come from jointly reinforcing both, consistently achieving the highest safety rates across models and attacks.

Figure 5: Safety rates (%) under different attack methods (Pure Harmful Prompt / GCG / ADV-LLM) for four models. Jointly scaling detection and refusal heads achieves the highest safety.

Experiment III: Robustness and Accuracy of DRefA

The reinforced model, DRefA, scales detection heads by 3.0 and refusal heads by 2.0. We evaluate its robustness under regenerated GCG attacks and its impact on general-purpose accuracy.

Robustness under Adaptive Attacks

DRefA significantly improves safety rates even when attackers regenerate GCG prompts specifically against it, confirming that reinforcing detection and refusal heads strengthens the safety circuit.

Figure 6: Safety rates under regenerated GCG attacks. DRefA improves robustness compared to baseline.

Effect on General Accuracy

DRefA incurs only modest drops on standard zero-shot benchmarks, demonstrating that safety improvements come at minimal cost to general-purpose utility.

Figure 7: Comparison of baseline vs. DRefA on HellaSwag, PIQA, and ARC benchmarks. Slight accuracy reductions are observed, but overall performance remains high.

Conclusion

Detection heads drive refusal heads to form safety circuits in LLMs. Leveraging this, DRefA a targeted head-scaling intervention—substantially improves robustness across models and attacks, achieving strong safety gains with minimal impact on accuracy. Future work can optimize scaling to further balance safety and utility.

Cite this work

Kuan-Lin Chu, Chung-En Sun, and Tsui-Wei Weng. "How to Make LLMs Safer? Detecting and Editing Key Heads in LLMs.", NeurIPS 2025 Lock-LLM Workshop,. 


@article{drefa,
    title={How to Make LLMs Safer? Detecting and Editing Key Heads in LLMs},
    author={Chu, Kuan-Lin and Sun, Chung-En and Weng, Tsui-Wei},
    journal={NeurIPS 2025 Lock-LLM Workshop},
    year={2025},
    url={https://github.com/Trustworthy-ML-Lab/Safer_LLM_with_Drefa}
}