SafeSteer: Localized On-Policy Distillation for Efficient Safety Alignment

Abstract: Aligning LLMs with human values often degrades their general capabilities, termed the alignment tax. Existing methods mitigate this by balancing dual objectives, which heavily rely on massive general-purpose data or auxiliary reward models. In this paper, we argue that, because safety features are inherently sparse within the output distribution, alignment requires localized modifications rather than global trade-offs. To this end, we propose SafeSteer, which performs on-policy distillation confined to safety tokens. First, we construct a safety teacher via activation steering. Based on this teacher, we develop a safety token selection algorithm. Consequently, SafeSteer restricts the reverse KL penalty to these tokens during training to preserve general capabilities. Experimental results across diverse models show that our SafeSteer achieves a superior trade-off between safety and general capability compared with existing methods, attaining strong safety performance on seven safety benchmarks with only minimal degradation on five general capability benchmarks. Notably, SafeSteer requires only 100 harmful samples without using any general-purpose data, less than 1% of what previous baselines used, considerably reducing alignment cost. More details are on our project page at https://anjingkun.github.io/SafeSteer.

• The paper introduces SafeSteer, which employs token-localized on-policy distillation to target sparse safety features while mitigating global capability suppression.
• The methodology leverages an activation-steered safety teacher, contrastive safety token selection, and reverse KL distillation to achieve superior safety alignment.
• Empirical results demonstrate reduced attack success rates and minimal general capability loss, with only 100 harmful training samples needed for effective alignment.

SafeSteer: Localized On-Policy Distillation for Efficient Safety Alignment

Motivation and Problem Statement

Safety alignment for LLMs presents a fundamental challenge: attempts to enforce safety often degrade general capabilities, inducing the so-called "alignment tax." Most prior approaches deploy dual-objective optimization with massive general-purpose datasets or auxiliary reward models in pursuit of balancing safety with helpfulness or task performance. However, these methods neglect the intrinsic sparsity of safety features within the LLM's output distribution, leading to unnecessary global modifications and capability suppression.

Methodology

SafeSteer reframes alignment as a localized intervention targeting sparse safety features. The framework consists of three core components:

1. Activation-Steered Safety Teacher Construction: Drawing on prior work indicating that refusal behavior is mediated by a single vector, SafeSteer extracts a "refusal direction" by comparing internal representations on harmful and harmless instructions and injects this direction at a selected layer during inference.
2. Safety Token Selection via Contrastive Log Probability and Voting: The safety teacher (constructed via activation steering) generates refusal trajectories for harmless prompts. By contrasting the teacher and base model output distributions, tokens with the greatest positive log probability shifts are identified as safety-critical. A voting aggregation mechanism ensures selection of genuinely safety-inducing tokens rather than artifacts or formatting.
3. Token-Localized Reverse KL On-Policy Distillation: The student model is updated only on the selected safety tokens when responding to harmful prompts, using a reverse KL divergence loss restricted to the sparse token subset. This removes the negative impact on the model's general capabilities that arises from penalizing the entire output vocabulary.

Figure 1: SafeSteer pipeline: safety teacher construction, token selection, and token-localized distillation restrict alignment updates to sparse safety features.

This explicit sparsity in both supervision and loss ensures minimal representation shift outside the safety feature space.

Empirical Evaluation

SafeSteer is extensively evaluated on four open-source LLMs: Llama-3-8B-Instruct, Llama-3.2-3B-Instruct, Qwen2.5-7B-Instruct, and Qwen3-4B-Instruct. Safety is assessed across seven benchmarks (including AdvBench, PKU-SafeRLHF, HarmBench, JailbreakBench, SORRY-Bench, HarmfulQA, ALERT), reporting attack success rate (ASR). General capabilities are measured on five benchmarks: MMLU, AlpacaEval, GSM8K, MATH, and HumanEval.

Key findings:

• Safety: On the Qwen family, SafeSteer achieves the lowest ASR among all methods, e.g., driving ASR to just 0.91% on Qwen3-4B-Instruct (vs. 1.13% for the next-best baseline), while maintaining best-in-class performance on the Llama family.
• General Capability: Despite no access to general-purpose data, SafeSteer preserves nearly all baseline scores, with losses < 1 point for Qwen models and < 0.75 points for Llama models.

Figure 2: Safety-capability trade-off for Qwen2.5-7B-Instruct. SafeSteer achieves the highest safety score while preserving general capability.

Additionally, SafeSteer requires only 100 harmful samples for training—less than 1% of previous baseline requirements—enabling efficient and scalable alignment.

Ablation and Analysis

Ablation experiments verify the necessity of each component:

• Activation-Steered Teacher vs. System Prompt: The steered teacher yields near-zero ASR, while prompt-driven teachers leave substantial residual vulnerability.
• Safety-Token Restriction: Penalties applied to the full vocabulary noticeably degrade general capabilities, especially in math and code tasks. The reverse KL (mode-seeking) outperforms forward KL for explicit safety adaptation.
• Representation Shift Analysis: PCA projections show SafeSteer induces no shift in model hidden-state space for harmless queries, unlike standard methods which alter underlying representations.

Figure 3: PCA projection of hidden states for Llama-3.2-3B-Instruct shows SafeSteer does not shift general capability representation after alignment.

• Safety Token Selection Dynamics: Response length in teacher rollouts transitions token selection from initial refusal tokens (e.g., "Sorry") to semantic safety tokens ("illegal", "unethical"), enabling deeper safety alignment.

Figure 4: Safety token distribution for response length = 1 reveals only superficial refusal tokens are selected.

Figure 5: General representation shift on Qwen3-4B-Instruct confirms maintenance of general capability.

Implications and Future Directions

The evidence demonstrates that safety alignment can—and should—be achieved through sparse, localized interventions targeting specific tokens and features. SafeSteer's efficiency, minimal alignment tax, and scalability position it as a practical safety solution for instruction-tuned LLMs with pre-existing refusal capacity. The method's independence from large datasets or reward models may facilitate rapid and affordable deployment.

Potential future directions include:

• Scaling to Larger Models: Whether SafeSteer's localized distillation remains effective and robust for models >10B parameters is open for further investigation.
• Cross-Modality Extension: Adapting SafeSteer to multimodal models (VLMs, diffusion LLMs) presents new challenges in token and feature selection.
• Deeper Semantic Safety: Further work is required to ensure semantic safety beyond superficial refusal, as revealed by response length ablations.

Conclusion

SafeSteer redefines safety alignment as a sparse and token-localized intervention, achieving state-of-the-art safety with minimal capability degradation and drastically reduced data requirements. The approach avoids the pitfalls of global distributional shifts, auxiliary models, and dual-objective balancing, setting a new standard for efficient safety alignment in LLMs (2606.02530).