Reproduce “Contrastive Reasoning Alignment: Reinforcement Learning from Hidden Representations” on OpenResearch

About this paper

Haozheng Luo, Yimin Wang, Jiahao Yu, Binghui Wang, Yan Chen

A framework called CRAFT significantly enhances the safety alignment of Large Reasoning Models by addressing superficial safety alignment, where internal thought processes might be unsafe despite a safe final refusal. Developed by researchers from Northwestern University, University of Michigan, and Illinois Institute of Technology, CRAFT uses latent contrastive learning to structure the model's internal reasoning space and reinforcement learning to guide it toward safety, resulting in substantial improvements in robustness against jailbreak attacks (up to 38.3% over baselines in final-response safety) while also improving general reasoning performance by 8.0%.

Key insights

• Initial analysis revealed that safe, unsafe, and 'rethink' reasoning traces occupy geometrically distinct regions in the model's latent space, providing a foundation for targeted alignment.
• Directly manipulating and structuring the latent space of reasoning traces, coupled with a reward that ensures consistency between internal latent safety and external textual safety, is effective in preventing superficial safety alignment.
• Integrating contrastive learning for latent space structuring with reinforcement learning for policy optimization allows for a robust, reasoning-level safety alignment that also preserves or enhances model utility.

Summary by alphaXiv.