SafeSteer: Localized On-Policy Distillation for Efficient Safety Alignment

The evolution of large language models (LLMs) has revolutionized natural language processing, enabling applications from chatbots to automated content generation. Early models like GPT-2 and BERT laid the groundwork, followed by more capable architectures such as LLaMA, PaLM, and Qwen, which demonstrated remarkable zero-shot and few-shot learning abilities. Despite these advances, safety concerns—such as the generation of biased, harmful, or misleading content—have become prominent. Traditional safety measures involve supervised fine-tuning with curated datasets or reinforcement learning with human feedback (RLHF), which help reduce unsafe outputs but often degrade the model’s broader capabilities—a phenomenon termed 'alignment tax.'

Recent research has explored various techniques, including orthogonal projection of safety gradients, sparse activation rerouting, and activation-based control, aiming to improve safety without sacrificing performance. However, these methods typically require large amounts of data, external reward models, or complex engineering, limiting their scalability and practicality. The inherent sparsity of safety signals—refusal behaviors and harmful content indicators—suggests that targeted, sparse adjustments could be more effective. This paper builds on these insights, proposing a localized, sparse approach to safety alignment that leverages the internal representations of models to identify and adjust only the critical safety tokens.

Existing safety alignment techniques often rely on broad, global modifications that inadvertently impair the model’s ability to perform general tasks. This is because safety signals are sparse and often disjoint from tokens used in normal operations. Consequently, global penalties or extensive data augmentation lead to a trade-off: enhancing safety at the cost of capability. Moreover, many methods depend on external reward models or large-scale datasets, increasing complexity and cost. The core challenge is to develop a method that can precisely target safety signals—those sparse, critical tokens—without affecting the model’s overall performance. Achieving this requires identifying these tokens automatically and applying localized adjustments during training, which remains a significant technical hurdle.

The paper introduces a set of key innovations:

1) Activation Steering for Safety Teacher Construction: By comparing the hidden states of the base model on harmful versus harmless prompts, a refusal direction is extracted and injected into the model’s residual stream, creating a stable safety teacher without external models or prompt engineering.

2) Contrastive Log Probability for Safety Token Mining: This algorithm compares the output distributions of the safety teacher and the base model, identifying tokens with the highest contrastive log probability differences as safety-critical. These tokens are sparse and highly relevant for safety adjustments.

3) Localized Reverse KL Divergence Penalty: Instead of applying a global penalty, the method restricts the reverse KL divergence to the identified safety tokens, enabling sparse, targeted fine-tuning that preserves the model’s general capabilities.

This combination of activation-based control, contrastive token selection, and localized divergence regularization constitutes a novel framework that effectively balances safety and performance with minimal data and cost.

• �� Construct a safety teacher model (πt) by comparing the hidden states of the base model (π0) on harmful and harmless prompts, extracting a refusal direction vector, and injecting it into the residual stream at a chosen layer.
• �� Generate refusal trajectories using πt on harmful prompts, and compare these with baseline responses to identify safety-sensitive tokens via contrastive log probabilities.
• �� For each trajectory, concatenate the response with the input and perform forward passes through both πt and π0, computing token-wise conditional probabilities.
• �� Calculate the contrastive log probability difference (Δ) for each token, selecting the top-K tokens with the highest Δ values as safety tokens.
• �� Aggregate votes for each token across multiple trajectories using a contrastive voting scheme, forming a safety token subset S.
• �� During on-policy distillation, generate responses from the student model (πs) on harmful prompts, and compute the reverse KL divergence only over the safety token subset S.
• �� Update the model parameters by minimizing this localized reverse KL loss, ensuring safety behaviors are learned without degrading general capabilities.
• �� Evaluate safety success rate and model performance on multiple benchmarks to validate effectiveness.

The experimental setup involves training models such as Llama-3-8B-Instruct, Llama-3.2-3B, and Qwen-2.5-7B using only 100 harmful samples from PKU-SafeRLHF, significantly less than previous methods. Baselines include MoCAN, BFPO, NSPO, and DPO-Mix, trained on larger datasets. Evaluation metrics encompass safety benchmarks like AdvBench, HarmBench, JailbreakBench, and SORRY-Bench, measuring attack success rate (ASR). General capabilities are assessed via MMLU, HumanEval, and AlpacaEval, using the DeepSeek-V4-Flash and lm-eval frameworks. Hyperparameters include a safety token subset size of |S|, response length H, and a temperature setting of 0 or 1.0. The training involves multiple rollouts (M) with localized reverse KL loss, with ablation studies testing the effects of response length, token selection, and normalization strategies.

Across multiple models, SafeSteer achieves a safety success rate below 2% on several benchmarks, outperforming strong baselines like MoCAN (around 3.75%) and BFPO (around 3.8%). On Qwen-3-4B-Instruct, the harmful response rate drops to 1.13%, with minimal performance loss on general tasks (less than 1.5%). Ablation studies confirm that restricting the reverse KL to safety tokens prevents capability degradation, while full-vocabulary penalties cause significant drops. Visualization of hidden states shows that the internal representations of the model remain nearly identical before and after alignment, indicating that the sparse, localized approach effectively avoids catastrophic forgetting. These results demonstrate that SafeSteer provides a practical, low-cost solution for safe deployment of large models.

This approach is directly applicable to deploying safer AI systems in content moderation, enterprise chatbots, and sensitive data filtering. Its low data requirement and high efficiency make it suitable for rapid iteration and deployment in real-world scenarios. By focusing on sparse safety signals, it can be integrated into existing pipelines with minimal retraining. In the long term, extending this framework to multimodal models and larger architectures could enable safer AI in autonomous vehicles, medical diagnostics, and other critical domains. The ability to dynamically adjust safety policies based on sparse signals also opens avenues for adaptive, context-aware safety mechanisms.

The method assumes that the base model already exhibits some refusal behavior; models lacking this will face challenges in safety token mining. Its effectiveness on models larger than 10B parameters remains unverified, requiring further scaling studies. The current focus is on autoregressive text models, and adaptation to multimodal or non-autoregressive architectures is yet to be explored. Additionally, the safety token selection relies on the quality of contrastive signals, which may vary across datasets and tasks, potentially affecting robustness. Future work should address these limitations by developing more generalizable token mining algorithms and extending the approach to broader AI modalities.