On-Policy Distillation for LLM Safety: A Routing Approach to Template-Robust Realignment

2026-07-29

AI Summary

This work addresses the vulnerability of existing safety alignment methods for large language models to adversarial fine-tuning, which often leads to catastrophic forgetting of domain expertise, failure against unseen attack templates, and susceptibility to jailbreaking via system prompt switching. To overcome these limitations, we propose a Routing-based Online Policy Distillation (ROPD) framework that introduces dynamic routing into safety distillation for the first time. Instead of fitting specific attack patterns, ROPD robustly re-aligns models by modeling the distributional divergence between aligned and compromised outputs. Extensive experiments across three base models and three datasets demonstrate that ROPD significantly outperforms four state-of-the-art baselines, maintaining strong defense efficacy and task performance under both template variations and prompt-switching attacks, with negligible performance degradation.

Abstract

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.

Problem

• LLM safety
• realignment
• template robustness
• catastrophic forgetting
• jailbreaking

Innovation

• On-Policy Distillation
• LLM Safety
• Template-Robust Realignment
• Routing Framework
• Catastrophic Forgetting Mitigation