Constrained Language Model Policy Optimization via Risk-aware Stepwise Alignment
Lijun Zhang, Lin Li, Wei Wei, Yajie Qi, Huizhong Song, Jun Wang, Yaodong Yang, Jiye Liang
TL;DR
This work addresses the challenge of balancing usefulness and safety in large language model alignment by introducing Risk-aware Stepwise Alignment (RSA), a token-level constrained policy optimization framework that incorporates nested risk measures into a risk-sensitive CMDP. It develops a two-step, monotonic policy-improvement scheme that maps risk-aware state-action values to reward- and cost-aware policies, with theoretical guarantees under mild assumptions. A practical RSA(P) variant avoids unstable dual updates by mixing a reward-focused policy with a safety-adjusted counterpart, improving stability and deployment practicality. Empirically, RSA achieves superior helpfulness while maintaining strong safety and suppressing tail risks across text generation and multi-turn conversations, outperforming Safe RLHF, SACPO, DPO, and Ra-DPO. Overall, RSA provides a principled, scalable approach to safer LLM deployment with explicit tail-risk control and robustness across diverse risk scenarios.
Abstract
When fine-tuning pre-trained Language Models (LMs) to exhibit desired behaviors, maintaining control over risk is critical for ensuring both safety and trustworthiness. Most existing safety alignment methods, such as Safe RLHF and SACPO, typically operate under a risk-neutral paradigm that is insufficient to address the risks arising from deviations from the reference policy and offers limited robustness against rare but potentially catastrophic harmful behaviors. To address this limitation, we propose Risk-aware Stepwise Alignment (RSA), a novel alignment method that explicitly incorporates risk awareness into the policy optimization process by leveraging a class of nested risk measures. Specifically, RSA formulates safety alignment as a token-level risk-aware constrained policy optimization problem and solves it through a stepwise alignment procedure that yields token-level policy updates derived from the nested risk measures. This design offers two key benefits: (1) it mitigates risks induced by excessive model shift away from a reference policy, and (2) it explicitly suppresses low-probability yet high-impact harmful behaviors. Moreover, we provide theoretical analysis on policy optimality under mild assumptions. Experimental results demonstrate that our method achieves high levels of helpfulness while ensuring strong safety and significantly suppresses tail risks, namely low-probability yet high-impact unsafe responses.