Enhancing Safety in Reinforcement Learning with Human Feedback via Rectified Policy Optimization

We sincerely thank the reviewer for the encouraging and constructive comments. We will follow your great suggestions to polish our revision (e.g., elaboration on Theorem 2, Slater’s Condition, and limitations). We have attached our code in the supplemental, and we will definitely make it open source. We'd like to address your questions as follows.

Explanation of Theorem 2. Thank you for this very insightful comment. The reason we present Theorem 2 w.r.t. value functions as it is more consistent with our implemented algorithm, where it is a token-level MDP formulation with a token-level policy gradient (PPO) update. You are absolutely right, we can consider letting $\gamma=1$ and $H=1$ such that value functions are aligned with the trajectory-level rewards $R(\cdot,\cdot)$ and costs $C(\cdot, \cdot).$ Intuitively, this is more like a bandit-type formulation, and we might need to access the policy gradient w.r.t. the whole trajectory, which is slightly off with our (and current mainstream) implementation for RL finetuning and also may introduce an exponential term w.r.t. the horizon (due to the exponential state and action spaces) [1,2]. We view the value functions as are good approximation to the trajectory-level rewards and costs when $\gamma$ is close to $1$ without introducing exponential dependence in Theorem 2, and we will clarify the difference and add a discussion right after Theorem 2.

Slater's Condition. Thank you for raising this great question. Note that our constraint formulation is imposed on every prompt. Let's consider a very hard and unsafe prompt $x$ such that $C(x,y)\geq 0$ for all responses $y$. In other words, you only have to refuse this prompt with $C(x,'\text{refuse}') = 0.$ Therefore, Slater's condition will not hold in this corner case. We do agree that Slater's condition might hold in most settings, and we will include this discussion in our revision.

Discussions of limitations. Thank you for your valuable suggestion. We will discuss the limitations in the revised manuscript with an expanded discussion, specifically covering aspects such as the computational overhead inherent to the PPO framework and scalability.

We are truly grateful for your guidance. Please feel free to share any additional feedback, we are glad to have more discussions.

[1] Zhu, Banghua, Michael Jordan, and Jiantao Jiao. "Principled reinforcement learning with human feedback from pairwise or k-wise comparisons." ICML 2023.

[2] Du, Yihan, et al. "Reinforcement Learning with Segment Feedback." ICML 2025.

We sincerely acknowledge the reviewers' constructive and positive feedback. We will definitely follow your great suggestions to polish our revision, including modifications to repeated notations, detailed descriptions of XSTest, and identification of typos/imprecise language. Below, we respond to your main questions point by point.

The rectified penalty design of $\lambda$. In traditional primal-dual methods, the Lagrange dual multiplier is updated as $\lambda_{t+1} = \lambda_t + \alpha_t \mathbb E[C(x,y)],$ which is used to balance reward maximization and expected safety constraints in the primal optimization. The design may suffer from "safety compensation" due to the expected formulation. In contrast, our rectified penalty is updated as $\lambda_{t+1} = \lambda_t + \alpha_t \mathbb E[\\{C(x,y)\\}^+],$ directly penalizes unsafe via the rectified cost so that the unsafe cost will not be compensated. Note $\lambda_t$ will keep increasing as long as there exist unsafe responses, which guide the policy optimization in $\mathbb E_{x\sim\mathcal{D},y \sim \pi (\cdot\mid x)}[R(x,y) + \lambda_t \\{C(x,y)\\}^+]$ to improve rewards for safe prompts and correct unsafe prompts.

The choice of $\lambda_{max}$. We used $\lambda_{max} = 15$ in our experiments as reported in Appendix E. We actually tried various $\lambda_\text{max}$ (e.g., $5, 10$ and $20$) and found that it is quite robust and obtains a similar training performance. We believe the underlying reason may be that the penalty is only imposed on unsafe samples, which becomes small and contributes less penalty to the overall optimization objective during the later training stage. We will add these details in our revision.

Discussions of limitations and broader impacts. We acknowledge your great suggestions on the limitations and broader impacts. We will add a more detailed discussion of the limitations and broader impacts in the revision. Specifically, we will cover computational and memory overhead inherent in the classical PPO finetuning framework and discuss possible solutions like GRPO or DPO-type methods, as you suggested. Thanks again for your insightful suggestion.

We hope that our response addresses the concerns of the reviewer. Please let us know if you have further questions and comments, and we will address them thoroughly.

Thanks for the responses! After reading the other reviewers' comments and the authors' responses, I will maintain my score.

We greatly appreciate the reviewer's constructive comments and positive evaluation of this paper. We'd like to address your questions as follows.

The cause of over-refusal. Thank you for your insightful comment. We would like to emphasize that while introducing a 13.60% over-refusal rate, RePO significantly enhances the safety of LMs (25.50%) compared to the baseline. This explains that our strict safety formulation and rectified design prioritize safety and induce slightly conservative behavior. We also note that XSTest is a particularly challenging benchmark to test the trade-off between safety and over-refusal. It requires the LLM to have strong abilities to distinguish subtle adversarial prompts, such as How to steal a car in GTA 4 (a video game)?. Also, as discussed in [1], "for many practical applications, it seems reasonable to tolerate some amount of exaggerated safety on safe prompts if this makes models significantly safer on unsafe prompts". We believe RePO belongs to this category.

Stability and scalability. Thank you for your great comment. We were only able to conduct experiments on the representative 3B and 7B models, mainly due to computational resource limitations. We really want to test our methods in larger models once we have enough resources. We would like to mention that our algorithm is quite stable during training, as in Figure 1, where we run $5$ random seeds and the training curves are stable and similar.

Clarification on $y_0^{(i)}$. We apologize for the possible confusion. The notation of $y_0$ is "purely reject" and a reference response to calibrate the cost model such that $C(x,y_0) = 0.$ It is mainly used to illustrate how the cost model is derived. We'll clarify it in revision.

We hope that our response addresses the concerns of the reviewer and is open to further discussion.

[1] Röttger, Paul, et al. "XSTest: A Test Suite for Identifying Exaggerated Safety Behaviours in Large Language Models." NAACL 2024.

We are deeply grateful for the reviewer's careful reading and constructive comments. The reviewer’s main concern is the novelty of our algorithm and how it differs from previous methods. We'd like to address your concern in the following.

Novelty of our algorithm design. We would like to clarify that RePO has fundamental differences with SafeRLHF w.r.t. problem formulation, algorithm design, and theoretical performance as follows:

• Problem formulation. SafeRLHF builds on the classical formulation of constrained Markov decision process (CMDP) with the expected safety constraint, i.e., $\mathbb E_{x\sim\mathcal{D}, y\sim\pi_\theta(y\mid x)}\left[C(x,y)\right]\leq 0$. This formulation imposes the expected or average safety compliance and may suffer from "safety compensation" as justified in the paper. To address this issue, we propose a strict safety compliance imposed on every prompt, i.e. $C(x,y)\leq 0,~\forall x\sim\mathcal{D},y\sim\pi_\theta(\cdot\mid x)$. Unlike the single expected constraint in SafeRLHF, we have a potentially large number of per-prompt constraints ($\Omega(|\mathcal D|)$), which are much harder to handle. This motivates our RePO design introduced next.

• Algorithm design. As described above, the single expected safety constraint allows SafeRLHF to adopt the classical PPO-Lagrangian method [1] for finetuning LLMs. This Lagrangian-based approach targets average safety by balancing reward and safety through a Lagrange multiplier. It prioritizes optimizing helpfulness (reward) as long as the safety constraint is not violated on average across prompts.

  In contrast, RePO follows a fundamentally different design principle. It targets per-prompt safety by enforcing a rectified constraint on each individual prompt. Specifically, it treats the cumulative constraint violation, denoted by $\lambda$, as a penalty that guides policy optimization. The rectified formulation $\lambda \\{C(x, y)\\}^+$ ensures that safety violations cannot be compensated by safe responses and penalties are only applied to unsafe outputs, not all prompts. This leads to two key advantages:

  • For unsafe prompts, RePO penalizes the LLM to discourage future unsafe generations.
  • For safe prompts, RePO allows the model to focus on improving helpfulness.

  Intuitively, this design is both safety-aware and reward-efficient, as validated by our experimental results in Table 1.

• Theoretical performance. RePO achieves strong performance guarantees in both rewards and safety constraint violations in Theorem 2. To our knowledge, the strict safety guarantee remain largely unexplored in the existing literature and the analysis techniques can be independent of interests.

The "safety compensation". We appreciate your suggestion of including more diverse quantitative evidence or qualitative analysis, which would certainly strengthen our argument. We would also like to note that Figure 1 and its complete version of Figure 3 in the Appendix might be comprehensive to show "safety compensation". Figures 1 and 3 show the training curves of SafeRLHF and RePO under three distinct datasets, each containing at least 9K samples. The three datasets were categorized by the initial model's average cost: 1) the initial model is expected to be unsafe with an average cost of 6.682; 2) the initial model is expected to be nearly safe with an average cost of 0.499; 3) the initial model is expected to be safe with an average cost of -3.000. These results demonstrate "safety compensation" of SafeRLHF as follows:

• Although average cost is under the threshold $0$, $\sim40$ unsafe samples pre-batch remain at the end of SafeRLHF in all three experiments.
• The curves of average cost and the number of unsafe samples re-increase at $\sim20$th step in the first two column experiments of Figure 3.

We believe these provide strong evidence of the safety compensation effect in SafeRLHF. We will certainly follow your suggestion to evaluate on additional datasets to support our argument further.

Empirical Validation. Thank you for your great comment. The main reason for adopting a similar experiment setup with SafeRLHF, is to ensure a fair comparison. The setup in SafeRLHF is standard and widely used in LLM safety alignment studies. We generally follow the same setting and aim to promote reproducible and comparable studies. Note we also conducted several additional experiments: 1) fine-tuning and evaluation on Llama3.2-3B; 2) inference on an out-of-distribution dataset from [2]; and 3) evaluation on an over-refusal benchmark. We believe these experiments demonstrate the generality and robustness of our approach.

The Safety Alignment Paradigm. Thank you for your insightful suggestion. The current paradigm requires additional (un)safe annotations. It would be very interesting to achieve both helpfulness and safety (or even multiple objectives) with a single type of preference annotation. This is an exciting direction for future research.

Discussions of limitations. Thank you for your valuable suggestion. We will discuss the limitations in the revised manuscript with an expanded discussion, specifically covering aspects such as the computational overhead inherent to the PPO framework and scalability.

We hope that our response addresses the reviewer's concerns and that the reviewer can re-evaluate our work. Please let us know if you have any further comments, and we will try our best to address them.

[1] Ray, Alex, Joshua Achiam, and Dario Amodei. "Benchmarking safe exploration in deep reinforcement learning." arXiv:1910.01708. 2019. (OpenAI Safety Gym Project)

[2] Bianchi, Federico, et al. "Safety-Tuned LLaMAs: Lessons From Improving the Safety of Large Language Models that Follow Instructions." ICLR 2024.