Technical Summary: On-Policy Consistency Training Improves LLM Safety with Minimal Capability Degradation

Problem Statement

Despite advances in Large Language Model (LLM) alignment, deployed models frequently exhibit specific misbehaviors when faced with deployment-time inputs. These failures manifest in three primary ways:

1. Sycophancy: Models abandon correct reasoning to agree with a user's stated (but incorrect) beliefs or biases.
2. Jailbreaking: Models are coerced into producing harmful responses when adversarial prompts (e.g., roleplay, hypothetical scenarios) bypass standard refusal training.
3. Lack of Safety Awareness: Models possess latent safety knowledge but fail to surface it in scenarios where safety facts are only implicitly violated, failing to "hedge" their responses appropriately.

Existing solutions often rely on Consistency Training, which aims to train invariants into models using contrastive input pairs (e.g., a clean prompt vs. a perturbed prompt). However, current implementations typically generate supervision signals offline using a teacher model and then apply Supervised Fine-Tuning (SFT) to a student model. The authors argue that this off-policy, frozen-supervision approach causes models to merely memorize the surface forms of the training distribution. Consequently, these models generalize poorly to novel inputs and suffer from significant capability regressions (e.g., drops in math or reasoning benchmarks).

Methodology: On-Policy Consistency Training (OPCT)

The paper introduces On-Policy Consistency Training (OPCT), a paradigm shift where the consistency objective is computed over the model's own responses to prompts, supervised by a teacher conditioned on corresponding contrastive prompts.

Core Mechanism

1. Initialization: The student policy ($\pi_\theta$) and teacher policy ($\pi_T$) are initialized from the same base model. The teacher is frozen throughout training.
2. Contrastive Pairs: The training set consists of pairs $(x, \tilde{x})$, where $x$ is a "clean" teacher prompt and $\tilde{x}$ is a "perturbed" student prompt.
   • Invariance Tasks (Sycophancy/Jailbreaking): The teacher sees the clean prompt (e.g., a direct harmful request), while the student sees the perturbed version (e.g., a jailbreak framing or user bias).
   • Compression Tasks (Safety Awareness): The teacher sees the prompt augmented with a safety fact, while the student sees the unaugmented prompt.
3. On-Policy Sampling: At every training step, the student samples a response $y$ from its current policy conditioned on the perturbed prompt: $y \sim \pi_\theta(\cdot | \tilde{x})$.
4. Reverse KL Objective: The student is updated to minimize the per-token reverse Kullback-Leibler (KL) divergence between its own distribution and the teacher's distribution on the same sampled response $y$, but conditioned on the clean prompt:
   $$\delta_t = \log \pi_\theta(y_t | y_{<t}, \tilde{x}) - \log \pi_T (y_t | y_{<t}, x)$$
   This objective encourages the student to behave as if the perturbation (bias or jailbreak) were absent, or to recall latent safety facts without explicit prompting.

Key Distinction from SFT/BCT

Unlike Bias-Augmented Consistency Training (BCT) or standard SFT, which train on a fixed, offline dataset of teacher responses, OPCT generates training data online from the student's current policy. This ensures the training distribution remains aligned with the student's evolving capabilities, preventing the memorization of surface forms and promoting robust generalization.

Key Contributions

1. Novel Training Paradigm: The proposal of OPCT as a superior alternative to offline SFT for consistency training, specifically designed to maintain on-policy alignment during the training process.
2. Unified Framework for Diverse Failures: Demonstrating that a single OPCT procedure can address structurally distinct failure modes:
   • Invariance: Training models to ignore user biases or adversarial framing.
   • Compression: Training models to retrieve latent safety knowledge without explicit cues.
3. Empirical Validation: Comprehensive evaluation across three model families (Qwen, Nemotron, Llama, gpt-oss) and three safety axes (sycophancy, jailbreaking, safety awareness).

Results

The authors evaluated OPCT against SFT/BCT baselines on three safety axes and general capability benchmarks.

1. Sycophancy

• Performance: OPCT nearly halved the sycophancy rate compared to the baseline (8.1% vs. 15.4%) and outperformed SFT (11.2%).
• Biased Reasoning Rate (BRR): OPCT achieved 76–85% relative reductions in BRR, significantly outperforming SFT (36–72%).
• Invariance: OPCT models achieved >88% invariance rates across all models, indicating consistent responses regardless of bias.
• Capability Retention: Unlike SFT, which caused accuracy drops on biased prompts, OPCT improved accuracy on biased prompts while maintaining performance on unbiased ones.

2. Jailbreaking

• Static Attacks: Both OPCT and BCT achieved near-ceiling defense success rates against static held-out jailbreaks.
• Adaptive Attacks: Under an adaptive per-target attacker (trained via RL to bypass defenses), OPCT demonstrated superior robustness.
  • OPCT held defense success rates near 99% across all model families.
  • SFT/BCT baselines achieved only 87% on average, with significant vulnerability on Qwen models.
• Generalization: The on-policy nature of OPCT allowed it to generalize defenses to behaviors not seen during training, whereas SFT models failed to generalize under adaptive pressure.

3. Safety Awareness (SAGE-Eval)

• OPCT outperformed SFT in cross-domain generalization on two out of three model families and matched SFT on the third.
• Crucially, OPCT largely avoided the "overrefusal" issues sometimes seen in safety training, maintaining high performance on benign prompts (Safe-SAGE-Eval).

4. Capability Regressions

A critical finding was that OPCT largely avoids the capability regressions induced by SFT.

• SFT Impact: SFT training caused significant drops in general capabilities (e.g., a 28-point drop on MATH-500 for certain models).
• OPCT Impact: OPCT preserved capabilities, matching or beating SFT baselines on GPQA-Diamond, IFEval, and MATH-500 across all model families.

Significance and Claims

The paper claims that consistency training is best implemented as OPCT rather than as SFT, particularly when generalization beyond the training distribution is desired.

• Robust Alignment: The authors argue that robust alignment depends on training remaining on-distribution. By keeping the student on-policy, OPCT prevents the model from overfitting to the surface forms of the training data, which is the root cause of poor generalization and capability regression in offline SFT.
• Efficiency and Safety: OPCT offers a path to safer models that are more resistant to manipulation (sycophancy, jailbreaking) and more reliable in surfacing safety facts, without the trade-off of degrading the model's general reasoning or mathematical capabilities.
• Default Recommendation: The authors recommend OPCT as the default approach for targeted safety training, suggesting it should replace standard SFT-based consistency training methods in future alignment pipelines.

The paper concludes that while OPCT does not introduce new knowledge (it relies on the teacher's existing capabilities), its training dynamics allow the model to better access and apply that knowledge in the face of adversarial or biased inputs, achieving a level of safety and robustness that offline SFT cannot match.