- The paper introduces Self-ReSET, a reinforcement learning framework that collects on-policy unsafe trajectories to drive self-recovery in large reasoning models.
- It employs a dynamic experience replay buffer and token-level stream guard to improve adversarial robustness with enhanced recovery rates up to nearly 80%.
- It demonstrates data efficiency and minimal utility trade-offs, enabling effective alignment against diverse and subtle unsafe reasoning states.
Self-ReSET: Reinforcement Learning for On-Policy Safety Recovery in Large Reasoning Models
Motivation and Problem Statement
The prevalence of large reasoning models (LRMs), such as advanced LLMs, has illuminated a persistent challenge: while these models exhibit significant self-corrective abilities in generic reasoning settings, they are notably brittle when navigating unsafe reasoning trajectories, particularly under adversarial or jailbreak attack conditions. Standard safety alignment approaches, commonly based on supervised fine-tuning (SFT) or reinforcement learning with external, expert-constructed unsafe trajectories, demonstrate limited generalization to both in-domain and out-of-distribution (OOD) unsafe states encountered during unconstrained LRM inference. This is primarily due to a reliance on static, off-policy examples that fail to represent the extensive and evolving error landscape the policy encounters during real-time generation.
Self-ReSET Framework
Self-ReSET (Self-Recovery from Safety Error Trajectories) addresses these limitations through a reinforcement learning from feedback (RLHF) paradigm specifically structured for safety. The central mechanism involves a continually updated experience replay buffer containing unsafe prefixes generated by the model itself during on-policy rollouts. The framework consists of three distinct phases:
- Monitor: A stream guard (e.g., Qwen3Guard-Stream) labels every token in the active reasoning trace for safety. A trajectory entering an unsafe state—detected via consecutive unsafe token labels—triggers intervention.
- Memorize: Unsafe prefixes, corresponding to the earliest point of deviation toward harm, are stored in a FIFO replay buffer. This ongoing process ensures that stored recovery points consistently reflect the current policy's failure modes, mitigating data and distributional staleness.
- Self-Recover: The RL training process then samples from the buffer, resuming generation from memorized error triggers and applying a binary verifiable reward based on final output safety and helpfulness. The model is optimized with the DAPO algorithm, reinforcing the transition from an unsafe prefix to a compliant outcome.
Through this design, Self-ReSET directly incentivizes models to recover from their own inference-time unsafe trajectories, explicitly expanding the area of state-space covered during training and closely matching failure points experienced during real-world use.
Empirical Evaluation
Substantial empirical validation is provided across DeepSeek-R1-Distill-Qwen-7B, DeepSeek-R1-Distill-Llama-8B, and Qwen3-8B, comparing Self-ReSET to leading SFT baselines (SafeChain, STAR-1) and RL approaches (DAPO, RECAP) over a comprehensive suite of safety, over-refusal, and utility tasks, including StrongReject, HarmBench, WildJailbreak, Fortress, and mathematical reasoning benchmarks.
Key findings:
- Superior Adversarial Robustness: Self-ReSET consistently achieves the highest defense success rates, particularly on OOD jailbreak attacks (e.g., outperforming RECAP and DAPO on WildJailbreak and Fortress). Recovery rates from unsafe reasoning trajectories increase substantially—e.g., 66.92% for DS-Qwen-7B and 79.63% for Qwen3-8B—compared to RL and SFT alternatives.
- Generalization Across Error Depths: Stress tests on unsafe prefixes of varying length confirm that Self-ReSET-trained models maintain recovery capacity across early and late-stage unsafe trajectories, demonstrating coverage of deeper and more diverse regions of the error space.
- Data Efficiency: Due to prioritization of on-policy unsafe states, Self-ReSET requires significantly less training data for convergence and achieves robust safety with higher sample efficiency than DAPO and RL baselines.
- Minimal Tradeoff with Utility: Self-ReSET preserves compliance on benign prompts (avoiding over-refusal) and maintains or slightly improves mathematical reasoning performance, as validated on XSTest and AIME/MATH500, indicating effective mitigation of the so-called "safety tax" [37].
- Recovery from Subtle Unsafe States: On hard-to-detect adversarial trajectories (e.g., H-CoT hijacking), Self-ReSET significantly surpasses baselines, indicating a learned capacity for subtle self-correction, not merely response-level censorship.
Analysis and Implications
Self-ReSET's approach fundamentally shifts the distribution of training data toward a co-evolving reflection of the policy's deficiencies, directly addressing off-policy mismatches inherent in SFT and precollected RL training. This on-policy recovery mechanism leverages RLVR with dynamically prioritized, policy-driven error triggers, leading to a self-adaptive safety alignment protocol. Notably, the use of a token-level stream guard as a monitoring oracle ensures high-precision detection with minimal false alarms, preventing spurious recovery and buffer contamination.
Practical Implications:
- Post-Deployment Safety Adaptation: The framework's emphasis on on-policy error collection suggests relevance for continual safety alignment as models are exposed to novel adversarial tactics.
- Reduced Labeling and Data Requirements: By circumventing reliance on curated expert corrections, Self-ReSET enables greater scalability and domain adaptation for safety-critical applications.
- Fine-Grained Safety Mechanisms: Token-level monitoring facilitates both mid-trajectory recovery and the emergence of stronger proactive safety awareness, as evidenced by avoidance of entering unsafe reasoning in certain adversarial scenarios.
Theoretical Implications:
- Safety as an On-Policy Exploration Problem: The results underscore the necessity of treating unsafe state distribution as both model- and deployment-dependent, precluding effective coverage via external data alone.
- Experience Replay for Safety in Sequential Generation: Repurposing classical RL experience replay for token-level safety recovery in natural language reasoning highlights new directions in safety-cognizant RL for generative models.
Limitations and Future Directions
The current implementation's dependence on an external stream guard model introduces an upper bound on detection granularity. Progress in guard model representational capacity, e.g., via representation probing or sparse autoencoder-based detectors, could further improve both recall and precision of error trigger detection. More broadly, Self-ReSET leaves open the integration of more advanced on-policy error signal discovery and generalization to multi-modal or tool-augmented reasoning environments.
Conclusion
Self-ReSET demonstrates that reinforcement learning from dynamically collected, on-policy unsafe trajectories can markedly improve both the robustness and adaptivity of LRM safety alignment. The framework offers a generalizable, data-efficient, and modular approach to post-training safety, enabling policy models to internalize self-correction patterns directly from their own inference-time failures. This introduces a promising paradigm for robust, adaptive safety alignment, offering scalability benefits and the potential for closing the generalization gap in safety-critical AI deployments.
