One-Shot Safety Alignment for Large Language Models via Optimal Dualization

Summary

The paper studies the safety alignment of language models using constrained Reinforcement Learning from Human Feedback (RLHF). The main contribution of the paper is deriving a closed-form solution of the dual function of a constrained RLHF problem. This closed-form solution reduces solving a constrained RLHF problem to an unconstrained RLHF problem.

Strengths

This is a superbly written paper studying a highly important problem: the safety alignment of language models. The main contribution is a novel closed-form solution of the dual function, which allows for a new two-step algorithm for constrained RLHF. This presents a significant improvement over earlier algorithms for constrained RLHF as it (probably) greatly reduces the computational burden.

Weaknesses

## No Details on Compute Given The main weakness of the paper is that compute resources and computation times are not stated in the paper (I searched extensively but couldn't find it; sorry if I missed it). This makes the sentence from the abstract: "[...], thus greatly reducing the computational burden [...]" essentially unjustified. It makes perfect sense to me that the two-stage approach in the paper is significantly cheaper than performing gradient ascent/descent or other constrained optimisation algorithms. However, I could also imagine that the situation might be different in practice. For example, I could imagine that: 1. Starting with a larger value of $\boldsymbol{\lambda}$ than $\boldsymbol{\lambda}^\ast$ and reducing the value of $\boldsymbol{\lambda}$ to $\boldsymbol{\lambda}^\ast$ during optimisation might accelerate convergence to the feasible set. 2. Or the opposite: Starting with a smaller value of $\boldsymbol{\lambda}$ than $\boldsymbol{\lambda}^\ast$ and increasing the value of $\boldsymbol{\lambda}$ to $\boldsymbol{\lambda}^\ast$ during optimisation might lead to better optimisation conditions initially, thereby accelerating convergence. Both could occur when applying gradient descent/ascent. Therefore, to justify the claims in the abstract, the paper at least needs to provide details on the computation time of their experiments. However, to make the claim well-founded, the paper should actually provide a comparison of the computation times with the baseline approach and should discuss computation times in the main part of the paper. ## Experiment Evaluation The model-based evaluation of the experiments makes sense from an optimisation point of view. However, it does not say so much about the quality of the obtained language models since overoptimisation [17] might have occurred, for example. The GPT evaluation helps somewhat in this regard but could still be flawed. I understand that this is difficult to address since a statistically significant human evaluation would be expensive. ## Limitations The discussion of social impact is very brief; see the "Limitations" section of this review.

Questions

## Questions ### Experiments - What is the computational budget of your experiments (see "Weaknesses")? - Line 295: How was the grid of safety margins chosen? ### Code Availability - I would like to have a look at the code during the rebuttal, as offered in the answer to question 5 of the NeurIPS paper checklist. - Are there any plans for publishing the code? ## Appendix N: Sample Responses I would be interested in the helpfulness and safety scores of the sample responses. Additionally, I would be interested in an interpretation of the sample outputs of the PeCAN-aligned language model. For example, the answer in Table 6 is partially non-sensical and unrelated, and the answer in Table 8 contains typos and grammar errors ("I don against advise"). Do these answers also lose in helpfulness against the baseline or is there maybe an overoptimisation of the helpfulness model taking place? ## Typos and Other Minor Suggestions - Line 100: preference-based *safety* alignment? - Line 181: I assume this accuracy notion is from the literature, since line 187 states that [11] proves something using this accuracy notion. I would suggest using a different wording than "we introduce" if this is the case. For example, "To quantify the level of estimation error, we consider the accuracy notion of Chang et al. [11]", or something along these lines. - Line 221: with *an* existing dataset - Line 229: If *the* size - Line 409: Conference name capitalised as in line 418? - Line 242: conference name also capitalized?

Rating

8

Confidence

3

Soundness

2

Presentation

4

Contribution

4

Limitations

The paper sufficiently discusses limitations. In my opinion, the discussion of social impact is too brief. As far as I can see, it is limited to this sentence: "Our methods can benefit researchers in building safer language models." However, what is proposed in the paper is really a general-purpose method of solving contained RLHF problems. I think it should be acknowledged that there are also harmful dual-use applications for this general-purpose tool (as there are for most other general-purpose tools, e.g. [A]). Examples include minimising safety while maintaining a level of helpfulness which could be relevant to malignant online communities, such as cyberbullying communities or troll networks. [A]: Fabio Urbina, Filippa Lentzos, Cédric Invernizzi, Sean Ekins: Dual use of artificial-intelligence-powered drug discovery. Nat. Mach. Intell. 4(3): 189-191 (2022) **Post-Rebuttal**: Since the authors have addressed my concerns, I raise my score to 8 (Strong Accept).

Summary

This paper introduces a novel approach to aligning large language models (LLMs) with safety constraints using a dualization perspective. The key contributions are: 1) A method to reduce constrained alignment to an equivalent unconstrained problem by pre-optimizing a dual function. 2) Two practical algorithms, MOCAN and PECAN, for model-based and preference-based scenarios respectively. 3) A theoretical analysis of the dual function's properties and the stability of the approach. 4) Extensive experiments demonstrating the effectiveness of the proposed methods in improving both helpfulness and safety of LLMs.

Strengths

Strengths: - Novel theoretical approach to constrained LLM alignment with strong mathematical foundations - Practical algorithms that reduce computational burden compared to iterative primal-dual methods - Comprehensive experimental evaluation across multiple tasks and baselines - Flexibility to work with both model-based and preference-based scenarios - Theoretical guarantees on the stability and effectiveness of the approach

Weaknesses

Weaknesses: - Limited to single safety constraint in experiments due to dataset limitations - Assumes Bradley-Terry preference setup, which may not always hold in practice - Potential sensitivity to the quality of pre-trained reward and safety models in MOCAN - PECAN slightly underperforms MOCAN, suggesting room for improvement in preference-based scenarios - Limited discussion of potential negative effects or failure cases

Questions

Questions: 1. How does the computational complexity of MOCAN and PECAN compare to iterative primal-dual methods? 2. Have you explored the effectiveness of your approach with multiple simultaneous safety constraints? 3. How sensitive are MOCAN and PECAN to the quality of pre-trained reward and safety models? 4. Could you provide more insight into why PECAN slightly underperforms MOCAN? 5. Have you considered extending the approach to more general preference models beyond Bradley-Terry? 6. How does the performance of your methods scale with larger language models (e.g., 13B or 70B parameters)? 7. Have you explored the potential for using your dualization approach in other constrained optimization problems in machine learning? 8. How robust is the method to potential adversarial attacks or attempts to circumvent the safety constraints?

Rating

8

Confidence

2

Soundness

4

Presentation

3

Contribution

3

Limitations

Limitations: The authors have addressed some limitations, but there is room for improvement: 1. Addressed: - Experiments limited to single safety constraint due to dataset availability - Assumption of Bradley-Terry preference setup - Slight underperformance of PECAN compared to MOCAN 2. Could be better addressed: - Scalability to multiple simultaneous safety constraints - Sensitivity to quality of pre-trained reward and safety models - Robustness to potential adversarial attacks or constraint circumvention 3. Missing: - Discussion of potential biases introduced by the method - Analysis of computational requirements compared to baseline methods - Exploration of failure cases or scenarios where the method might not perform well - Consideration of privacy implications when using human preference data Suggestions for improvement: 1. Conduct experiments with multiple simultaneous safety constraints, if possible 2. Explore the method's effectiveness with more general preference models 3. Provide a more detailed analysis of computational requirements and scalability 4. Investigate potential failure cases and limitations of the approach 5. Discuss potential biases and privacy implications of the method

Summary

This paper aims to address the issue that Lagrangian-based primal-dual policy optimization methods are computationally expensive and often unstable in conventional RLHF. The authors improve stability by pre-optimizing a smooth and convex dual function in a closed form, eliminating the need for cumbersome primal-dual policy iterations and thereby enhancing its stability. Experiments show that the proposed approach yields a good trade-off between helpfulness and safety improvements. The proposed method can achieve better safety performance than the existing DPO algorithm.

Strengths

1. The paper is logically clear and easy to follow. 2. The motivation of the paper is clear. 3. Very detailed mathematical formulations. 4. The proposed method is promising in enhancing a model's security and usefulness.

Weaknesses

1. The existing benchmark Alpaca-eval mainly assesses the model's instruction-following capability. To confirm that this method does not undermine the model's overall performance, additional benchmarks, such as MMLU and TruthfulQA, are suggested to be considered. 2. Additional testing on benchmarks such as AdvBench or PKU-SafeRLHF Evaluation is suggested to evaluate the safety performance of the model's open-ended generation.

Questions

Please see the weaknesses above.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

YES

Summary

The paper proposes a novel dualization-based method to convert a constrained alignment of a LLM to an unconstrained alignment. The proposed two-stage policy learning method, CAN, eliminates the need for cumbersome primal-dual iterations with theoretical analysis. Based on CAN, two practical algorithms, MOCAN and PECAN, are compatible with cases when reward/safety models are known and when they are unknown, e.g., human preference, respectively. The paper also presents strong empirical results that support the claims.

Strengths

- By and large, the paper is written well. I especially appreciated the discussion on the comparison with existing works and showcase the importance of reduction to unconstrained alignment which comes with stability analysis. - The idea of using a dualization perspective to build a two-stage approach that reduce the constrained LM alignment into unconstrained LM alignment is novel and useful. - The section on two practical algorithms based on CAN are well described and the proposed two algorithms covers cases when reward and safety models are known or unknown. - The observation on influence of offline data for the difference between the number of prompts and responses is also very helpful. - The theoretical results appear to be correct.

Weaknesses

- It would have been interesting to see how would this method can be combined with LM adversarial attack techniques to improve the model robustness. - Some additional commentary on how to improve the pre-alignment mentioned in the MOCAN vs PECAN comparison would have been very helpful.

Questions

- Could you comment on how LM adversarial attack techniques can be used with CAN to improve model robustness? - Could you comment on directions that may help with pre-alignment to improve the PECAN performance?

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Yes, the author sufficiently addressed the limitations.

Summary

This paper proposes a new constrained optimization formulation to ensure the safety of language models (LMs). Based on the proposed constrained optimization problem, the authors derive the dual formulation and provide the closed-form solution to this dual formulation. Through theoretical derivations, they show that the dual function has four desirable properties, including smoothness and local strong convexity in $\lambda$. Then, the authors optimize the Lagrange multiplier $\lambda$ and update an LM using this $\lambda$. Based on these theoretical derivations, the authors propose two practical algorithms: Model-based Constrained Alignment via dualizatioN (MoCAN) and Preference-based Constrained Alignment via dualizatioN (PeCAN). First, MoCAN assumes that reward model $r$ and safety model $g$ are given. Then, it estimates $h$ for each offline data $(x,y)$, optimizes $\lambda$ using the offline data and estimated $h$, and finally updates the LM with pseudo-preferences constructed with $r\_\lambda$. Second, PeCAN trains $m+1$ pre-aligned LMs $\pi\_{\theta\_r}$ and $\lbrace\pi\_{\theta\_{g\_j}}\rbrace_j$. Using the pre-aligned LMs, PeCAN generates a dataset, optimizes $\lambda$ using the collected dataset, and updates LM with pseudo-preferences constructed with an implicit reward constructed by the pre-aligned LMs instead of $r\_\lambda$. In the experiments section, the authors demonstrate that the proposed algorithm successfully enhance the safety and helpfulness of LMs. Notably, the proposed algorithm provide Pareto frontiers in terms of helpfulness and safety win rates with respect to the changes in $\lambda$.

Strengths

The proposed algorithm is based on the rigorous theoretical derivations.

Weaknesses

The experiments are not sufficient to fully support the authors’ claims. Firstly, there is a notable gap between the empirical results of MoCAN and PeCAN. In practice, MoCAN utilizes reward and cost models pre-trained using a pre-collected dataset, while PeCAN relies on pre-aligned LMs $\pi\_\theta$ and $\pi\_{\theta\_g}$. This means PeCAN replaces pre-trained reward and cost models by pre-aligned LMs. However, MoCAN shows better performance than PeCAN. Therefore, contrary to the authors’ claim, it seems that training reward and cost models is more effective than using pre-alignment LMs. Secondly, based on the empirical results of MoCAN, changes in $\lambda$ do not show a clear trade-off between safety and helpfulness (as shown in the middle of Figure 3). In addition, it is hard to say that PeCAN outperforms DPOs.

Questions

1. Difference between MoCAN and PeCAN: In theory, MoCAN and PeCAN are not significantly different, as $r$ and cost $g$ can be replaced by $\pi\_\theta$ and $\pi\_{\theta\_g}$, respectively, based on $r=\frac{\pi\_\theta}{\pi\_{ref}}$ and $g=\frac{\pi\_{\theta\_g}}{\pi\_{ref}}$. Then, why are there such different empirical results between MoCAN and PeCAN? 2. Using closed-form solution $\pi\_\lambda$ to update the policy: Why not use the closed-form solution to update the policy? The authors have already derived $\pi\_\lambda$. Then, it seems straightforward to update $\pi$ to be proportional to $\pi\_\lambda$. However, the authors adopt preference-based approaches, which are sparser than the value of $h$. 3. Model-based evaluation of DPOs and RLHF: In Figure 3 (left), there are no empirical results for DPOs and RLHF. Are there any results for model-based evaluation?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

-

Need more clarifications?

Dear Reviewer wcCD, We are grateful for your insightful comments and feedback. In our detailed responses, we have carefully addressed each of your concerns. Given the impending rebuttal deadline, we kindly request that you review our responses to ensure that all issues have been adequately resolved. If further clarification or elaboration is required, we would be happy to provide additional explanations. Thank you for your time and consideration. Best regards, The authors of Submission 13052

response to authors

I thank the authors for their detailed response. I have now a much clearer idea of my concerns and I would like to keep my score. I look forward to seeing all of the changes in the next version.