PaperlayerView exact revision on arXiv

Source-linked AI summary

THINKSAFE: Self-Generated Safety Alignment for Reasoning Models

Seanie Lee, Sangwoo Park, Yumin Choi, Gyeongman Kim, Minki Kang, Jihun Yun, Dongmin Park, Jongho Park, Sung Ju Hwang

arXiv:2601.23143v4cs.AI

TL;DR

Reasoning-intensive post-training can weaken safety, while external teacher supervision may distort the student’s native distribution. ThinkSafe formalizes safety realignment as a KL projection and uses refusal steering to self-generate filtered safety traces. Across Qwen3 and DeepSeek-R1-Distill, it reports the strongest safety-reasoning balance with substantially lower compute than online alternatives.

  • Problem

    Reasoning optimization can degrade safety, while external teacher distillation introduces distributional discrepancy that may reduce native reasoning.

  • Method

    ThinkSafe uses refusal steering to elicit the student’s latent safety knowledge, filters generated traces, and fine-tunes on a static self-generated dataset.

  • Results

    ThinkSafe consistently achieves the most favorable safety-reasoning trade-off across Qwen3 and DeepSeek-R1-Distill, outperforming teacher and online-RL alternatives.

  • Takeaways & Limitations

    Self-generated, refusal-steered data can restore safety while preserving reasoning performance without external teachers.

  • Takeaways & Limitations

    ThinkSafe assumes retained latent safety knowledge, depends on an external safety classifier, and is evaluated only with LoRA on single-turn prompts.

Abstract

from arXiv · show

Large reasoning models (LRMs) achieve remarkable performance by leveraging reinforcement learning (RL) on reasoning tasks to generate long chain-of-thought (CoT) reasoning. However, this over-optimization often prioritizes compliance, making models vulnerable to harmful prompts. To mitigate this safety degradation, recent approaches rely on external teacher distillation, yet this introduces a distributional discrepancy that degrades native reasoning. We formalize safety realignment as a KL projection onto the safe simplex and prove that the student's own safety-filtered distribution is the unique KL-optimal target, while any external teacher incurs an irreducible excess KL penalty. Guided by this analysis, we propose ThinkSafe, a self-generated alignment framework that restores safety without external teachers. Our key insight is that while compliance suppresses safety mechanisms, models often retain latent knowledge to identify harm. ThinkSafe unlocks this via lightweight refusal steering, which preserves the KL-optimal target while increasing the acceptance rate. Experiments on DeepSeek-R1-Distill and Qwen3 show ThinkSafe significantly improves safety while preserving reasoning proficiency, and achieves superior safety and comparable reasoning to GRPO with roughly an order of magnitude less compute. Code, models, and datasets are available at https://github.com/seanie12/ThinkSafe and https://huggingface.co/Seanie-lee/collections.

1 Introduction

Reasoning-oriented post-training can weaken safety, while external teacher distillation introduces distributional mismatch. ThinkSafe uses refusal steering to recover latent safety knowledge while preserving the student’s native reasoning distribution.

  • 1 Introduction: Reasoning optimization can degrade safety alignment, creating a trade-off between harmfulness reduction and retained problem-solving capability.Prior work describes this as a negative correlation between reasoning and safety and as a safety tax.
  • 1 Introduction: External teacher distillation can deviate from the student’s internal distribution, causing safety internalization failures or degraded native reasoning.The paper formalizes this deviation as an irreducible KL penalty.
  • 1 Introduction: Naive self-distillation preserves in-distribution data but can produce unsafe responses because compliance priors suppress the model’s safety mechanisms.Online student sampling avoids teacher shift but incurs prohibitive computational cost.
  • 1 Introduction: ThinkSafe prepends refusal-oriented instructions to harmful prompts, eliciting the student’s latent safety reasoning while sampling benign prompts directly.The resulting data preserves native helpfulness for benign instructions.
  • 1 Introduction: ThinkSafe achieves the most favorable safety-reasoning balance across Qwen3 and DeepSeek-R1-Distill, with superior safety and comparable reasoning to GRPO at roughly an order of magnitude less compute.The contribution summary reports consistent gains across diverse benchmarks and comparisons with teacher distillation and online RL.

2 Related Works

Related work seeks to preserve reasoning while restoring safety, but teacher-based and naive self-distillation approaches face distributional or data-quality limitations. ThinkSafe addresses these limitations through self-generated safety alignment.

  • 2 Related Works: Reasoning-oriented RL can compromise safety alignment, while reacquiring safety can reduce reasoning capability.The related-work discussion describes this two-sided trade-off as a negative correlation and a safety tax.
  • 2 Related Works: Safety alignment methods increasingly preserve reasoning by generating structured safety traces rather than bypassing reasoning entirely.SafeChain uses structured reasoning steps, while STAR-1 uses policy-guided traces from a larger teacher.
  • 2 Related Works: Self-distillation can improve generalization and mitigate fine-tuning distribution gaps, but harmful queries lack safe ground-truth references for rewriting.Its application to safety alignment is therefore challenging.

3 Method

ThinkSafe generates filtered safety data from a frozen student reference model, using refusal steering for harmful prompts and direct sampling for benign prompts. Fine-tuning on this static dataset approximates the safety objective while retaining native distributions.

  • 3 Method: The setup assumes an initially safety-aligned model whose reasoning-oriented post-training degraded guardrails but preserved latent safety knowledge.The goal is robustness against harmful prompts while retaining general reasoning capabilities.
  • 3 Method: Teacher distillation samples and filters responses from an external model, introducing a distributional gap that can degrade native reasoning even when teacher and student sizes match.The paper formalizes this gap as an irreducible excess KL.
  • 3 Method: Refusal steering is presented as a way to elicit latent safety reasoning while avoiding the unsafe outputs of naive student sampling.This provides a middle ground between teacher-induced distribution shift and computationally expensive online learning.
  • 3 Method: For harmful prompts, ThinkSafe adds a refusal-oriented instruction to elicit safe responses and explicit safety reasoning traces.The instruction shifts probability mass from compliant-unsafe responses toward safety-aligned reasoning paths.
  • 3 Method: Benign prompts are sampled without additional instruction so benign targets remain on the student’s native distribution.The benign split is crucial for preserving general instruction-following.
  • 3 Method: ThinkSafe filters harmful and benign generations with a safety guard and fine-tunes the student on the resulting valid traces.The practical implementation merges filtered prompt-response pairs into a static dataset instead of performing online updates.

4 Theoretical Analysis

The analysis identifies the student's safety-filtered distribution as the unique KL-optimal safe target and shows that refusal steering can reach it while improving acceptance rates. External teacher sources incur irreducible distributional drift, whereas steered self-generation preserves the target and reduces sampling cost.

  • KL-optimal target: The student's safety-filtered distribution p+ref is the unique KL-optimal target for safety realignment.Filtering converts the frozen student's distribution into a safe target while minimizing forward-KL drift from that student.
  • THINKSAFE objective: THINKSAFE combines zero excess KL relative to p+ref with a tractable acceptance rate through offline sampling.This combines preservation of the KL-optimal target with reduced generation cost.
  • Teacher-induced drift: Any external teacher source incurs a strictly positive, irreducible excess KL penalty when it differs from the frozen student's distribution.The penalty cannot be removed through additional filtering or data and is identified as teacher-induced distribution shift.
  • Safe projection: Filtering is the most conservative source modification that places all probability mass on safe outputs.The result applies to student, teacher, and steered sources and is characterized as the smallest χ2-ball step achieving perfect safety reward.
  • Refusal steering: Under refusal-tilt steering, the filtered steered source remains p+ref while its acceptance rate increases, strictly when 0 < αref(xh) < 1.Steering reweights safe outputs relative to unsafe outputs while preserving within-group relative probabilities.
  • Refusal steering: For small αref(xh), refusal steering's sampling speedup approaches ω(xh), while lower-acceptance prompts receive larger multiplicative boosts.The steered source achieves the same KL minimum as benign self-generation and reduces acceptance-rate imbalance across harmful prompts.

5 Experiment

ThinkSafe is evaluated across Qwen3 and DeepSeek-R1-Distill models using safety, over-refusal, and reasoning benchmarks. It consistently improves the safety-reasoning trade-off, while ablations show that refusal steering and explicit safety reasoning are central to its effectiveness.

  • Experimental Setup: ThinkSafe evaluates Qwen3 and DeepSeek-R1-Distill models on harmfulness, benign refusal rate, and reasoning pass@1 across multiple benchmarks.Safety uses HarmBench, StrongReject, and WildJailbreak; reasoning uses GSM8K, MATH500, AIME24, and GPQA.
  • Main Results: ThinkSafe consistently achieves the most favorable safety-reasoning trade-off across Qwen3 and DeepSeek-R1-Distill families.On Qwen3-4B, HarmBench harmfulness falls from 38.21 to 9.63 while average reasoning rises from 74.47 to 77.18.
  • Main Results: External-teacher baselines frequently degrade reasoning, including SafeChain drops from 44.95 to 39.86 on Qwen3-0.6B and from 64.87 to 60.93 on Qwen3-1.7B.The paper attributes this pattern to an irreducible excess KL caused by student-teacher distribution mismatch.
  • Online Comparisons: ThinkSafe retains an efficiency advantage over online baselines: GRPO requires over 21 hours and is approximately 8 times slower, while OPD requires over 88 hours.ThinkSafe’s reported time includes data generation, yet its additional cost remains marginal.
  • Ablation Studies: Removing safety reasoning increases harmful responses from 29.5 to 44.4 at 7B and from 19.1 to 33.7 at 8B, while also degrading reasoning.Naive rejection sampling similarly fails without refusal steering: Qwen3-8B safety changes only from 19.6% to 21.3%, versus 4.5% for ThinkSafe.

6 Conclusion

THINKSAFE addresses the distributional discrepancy of external teacher supervision by using refusal steering to unlock latent safety knowledge and generate student-aligned reasoning traces. It achieves a favorable safety–reasoning trade-off, while remaining subject to assumptions about retained safety knowledge, filtering quality, training regime, and evaluation scope.

  • THINKSAFE uses lightweight refusal steering to unlock latent safety knowledge and generate self-aligned reasoning traces without disrupting native problem-solving.
  • THINKSAFE achieves the most favorable safety–reasoning trade-off across Qwen3 and DeepSeek-R1-Distill families.
  • The method assumes that students retain latent safety knowledge from prior alignment, limiting applicability to base models without safety training.
  • Performance depends on the external safety classifier, while static offline data becomes increasingly off-policy during fine-tuning.
  • Evaluation is limited to LoRA fine-tuning on single-turn prompts, excluding larger scales, full fine-tuning, and multi-turn or agentic settings.

A Proofs

The proofs establish that safety realignment is a projection problem: among safe response distributions, the student’s safety-filtered distribution uniquely minimizes KL divergence to the frozen student. Refusal steering preserves this target while reducing the sampling cost of collecting accepted traces.

  • The unique safe distribution minimizing KL divergence to the frozen student is r*=p+.
  • The accepted conditional distribution π+ is the unique optimizer among distributions supported on safe responses.
  • The source expansion is algebraically correct, but its second term is not sign-definite, unlike the projection-form decomposition.
  • Under refusal tilt, safe and unsafe outputs retain their within-group probabilities while only the odds between groups change.
  • Refusal steering acts as a cost-reducing proposal distribution that leaves the accepted conditional unchanged while lowering the sampling cost for accepted traces.
  • If steering raises acceptance above 0.5 with ω≥19, the variance bound becomes 2m, a roughly 190-fold reduction from 380m.

B.2 Statistics

THINKSAFE response statistics show that benign prompts receive longer responses than harmful prompts, and that response lengths become longer and more stable as model size increases. Figures 9 and 10 summarize these statistics for Qwen3 and DeepSeek-R1-Distill, excluding the longest 1% of outliers.

  • Nh and Nb denote harmful and benign prompt counts, while µh and µb denote their average response lengths in tokens.
  • Benign responses consistently have longer generation lengths than harmful responses across Qwen3 and DeepSeek-R1-Distill model series.The difference reflects more detailed reasoning traces for benign queries.
  • As model size increases, harmful and benign responses tend to become longer and more stable in distribution.
  • Figures 9 and 10 report THINKSAFE statistics for Qwen3 and DeepSeek-R1-Distill after excluding top-1% length outliers.

C Experimental Details

The experiments evaluate reasoning with pass@1 across four benchmarks and examine THINKSAFE through reasoning ablations, cross-model teachers, alternative refusal prompts, and different safety classifiers. These studies show that explicit reasoning, self-generation, simple refusal steering, and classifier substitution each materially affect the safety–reasoning balance.

  • Experimental setup: Reasoning is evaluated on AIME 2024, GSM8K, MATH500, and GPQA using average pass@1 over eight sampled responses.
  • Reasoning ablation: Removing refusal-response reasoning traces degrades both safety and reasoning; Qwen3-4B Avg Pass@1 falls from 77.2% to 57.8%.
  • Cross-model distillation: Cross-model teachers can improve safety, including a 25.3% safety gain from Qwen3-8B data for DeepSeek-R1-Distill, but consistently degrade reasoning.
  • Refusal steering strategies: Suffix refusal steering matches the default prefix strategy, whereas Risk and Intent prompts produce higher harmful response ratios.
  • Filtering ablation: WildGuard filtering yields nearly identical safety to baseline while preserving superior reasoning, supporting robustness to the choice of safety classifier.

E.5 Experiments with Larger Models

THINKSAFE maintains a favorable safety–reasoning trade-off as model size increases. In a controlled Qwen3-8B comparison, it also outperforms SafePath under matched data scale.

  • Larger-model experiments: 5.7× reduction in average harmfulness, from 17.44 to 3.04, while reasoning changes only from 82.56 to 82.12 on Qwen3-14B.The method remains effective beyond the 8B regime and preserves reasoning within 0.44 points of the initial model.
  • Larger-model experiments: THINKSAFE achieves the most favorable safety–reasoning trade-off at both Qwen3-14B and Qwen3-32B scales.These results support empirical realizability of the KL-optimal target at larger scales.
  • Controlled comparison: 4.50% vs 11.50% average harm rate, with THINKSAFE preserving reasoning while more than halving harm relative to controlled SafePath.The comparison uses the same prompt set and matched data scale.

F.4 Experimental Setup

The GRPO and On-Policy Distillation baselines use online rollout generation, with implementation choices designed to support comparable training conditions. Both setups incur substantial generation or scoring requirements.

  • GRPO: GRPO generates G = 8 rollouts per prompt, using β = 0.04 for the KL penalty and ϵ = 0.2 for clipping.The implementation uses TRL with vLLM and keeps optimizer, batch size, and hardware identical to the primary setup.
  • On-Policy Distillation: On-Policy Distillation trains on student-generated rollouts with a full-vocabulary backward KL loss against teacher distributions.The trainer is built on TRL and uses vLLM for online generation.
  • On-Policy Distillation: A 4096-token maximum output length is used because repeated rollout generation and teacher scoring introduce substantial latency.The budget is described as sufficiently large for the implementation.

G Limitations

The paper identifies boundaries involving latent safety knowledge, evaluation scale, fine-tuning scope, deployment scenarios, data staleness, and operational safety risks. It also reports safeguards and low over-refusal rates, while emphasizing that adversarial robustness is not established.

  • Method dependencies: THINKSAFE may fail on base models without prior safety training because refusal steering depends on latent safety knowledge.External teacher supervision or safety-reward RL may remain necessary as an initial step in that setting.
  • Evaluation scope: The evaluation covers 0.6B–8B models across two families, leaving behavior at 70B+ and on closed frontier models unvalidated.Larger-scale effectiveness is suggested by filtering ratios but requires further study.
  • Training scope: LoRA fine-tuning limits conclusions about distributional discrepancy because full fine-tuning may exhibit different trade-offs.A comprehensive comparison is left for future work.
  • Deployment scope: The study focuses on single-turn harmful prompts and does not explicitly address multi-turn jailbreaks, long-context adversarial prefixes, or agentic tool-call settings.Extending refusal steering to these scenarios is identified as future work.
  • Operational limitations: The static offline dataset becomes increasingly off-policy as the student shifts during fine-tuning, while iterative self-training would add computational cost.The paper also warns that THINKSAFE reduces but does not eliminate harmful responses and should not be treated as adversarially robust.
Loading 2601.23143v4…