Back to AI Research

AI Research

On-Policy Distillation for LLM Safety: A Routing Ap... | AI Research

Key Takeaways

  • On-Policy Distillation for LLM Safety: A Routing Approach to Template-Robust Realignment addresses the vulnerability of large language models (LLMs) to "misa...
  • We conduct extensive experiments comparing ROPD against four state-of-the-art baselines across three datasets and three base models with varying alignment strengths.
  • Our results demonstrate that when baseline defenses face template mismatches, often accompanied by severe degradation in downstream task performance.
  • In contrast, ROPD substantially mitigates template-mismatch risks, maintaining superior robustness in both defense effectiveness and capability preservation.
  • While our analysis indicates ROPD is not entirely immune to template shifts, its performance degradation is negligible compared to existing methods, establishing a new standard for robust LLM realignment.
Paper AbstractExpand

Fine-tuning is the dominant paradigm for specializing large language models (LLMs), yet it exposes a critical vulnerability: malicious data providers can embed harmful behaviors into downstream corpora, creating models that retain professional skills while violating human values on demand. Existing safety-realignment defenses often fail in practice due to three key limitations: they frequently cause catastrophic forgetting of specialized skills; their effectiveness collapses when the defender cannot observe the attacker's prompt template; and successfully realigned models remain susceptible to re-jailbreaking via simple system prompt switches. To address these challenges, we propose Routing-based On-Policy Distillation (ROPD), a novel realignment framework that models the divergence between aligned and compromised output probability distributions rather than fitting specific prompt templates. We conduct extensive experiments comparing ROPD against four state-of-the-art baselines across three datasets and three base models with varying alignment strengths. Our results demonstrate that when baseline defenses face template mismatches, often accompanied by severe degradation in downstream task performance. In contrast, ROPD substantially mitigates template-mismatch risks, maintaining superior robustness in both defense effectiveness and capability preservation. While our analysis indicates ROPD is not entirely immune to template shifts, its performance degradation is negligible compared to existing methods, establishing a new standard for robust LLM realignment.

On-Policy Distillation for LLM Safety: A Routing Approach to Template-Robust Realignment addresses the vulnerability of large language models (LLMs) to "misalignment through fine-tuning." In this scenario, malicious data providers embed harmful behaviors into a model while keeping it functional for professional tasks, making the compromise difficult to detect. Existing safety-realignment defenses often fail because they damage the model's specialized skills, rely on knowing the attacker's specific prompt template, or remain susceptible to simple system prompt switches.

The Routing-based On-Policy Distillation (ROPD) Framework

ROPD is a realignment framework that shifts the focus from matching specific prompt templates to modeling the divergence between the output probability distributions of an aligned model and a compromised one. The framework uses two frozen "teacher" models to guide a student model during the realignment process:

  • Safety Teacher: The original, pre-attack aligned model, which provides a refusal prior that is largely independent of the prompt template.

  • Task Teacher: The fine-tuned, compromised model, which retains the specialized downstream skill.
    During training, the system routes each token to one of these two teachers based on the source of the example—harmful prompts are routed to the safety teacher, while task-specific examples are routed to the task teacher. The student model is then trained to match the output distribution of the routed teacher using a top-K KL divergence objective. This dual-teacher approach allows the student to inherit safety behaviors from the original model while preserving the task capabilities of the fine-tuned model.

Experimental Results

Researchers evaluated ROPD against four state-of-the-art baselines (SSRD, RESTA, soft-SFT, and rollback) using three base models (Llama-2, Qwen2.5, and Gemma-2) and three downstream tasks (SQL generation, dialogue summarization, and shell synthesis).
The experiments showed that when baseline defenses faced a mismatch between the defender's template and the attacker's template, their effectiveness dropped by more than 30%, often resulting in a total loss of downstream task performance. In contrast, ROPD maintained superior robustness. When using its "self" template, ROPD consistently reduced the attack success rate (ASR) while preserving task performance across all three models, regardless of the template used by the attacker.

Limitations and Considerations

While ROPD significantly reduces the risks associated with template mismatches, the authors note that it is not entirely immune to them. Because the student model is realigned and evaluated under a specific deployment template, a substantial difference between that template and the attacker's original template can still result in residual safety risks.
The authors conclude that by decoupling the safety and task-retention objectives through source-routed distillation, ROPD establishes a more robust standard for realigning models that have been compromised through malicious fine-tuning.

Comments (0)

No comments yet

Be the first to share your thoughts!