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

We thank the reviewer for the positive evaluation and the valuable feedback. We have answered all questions to the best of our ability. We are glad to address any further questions you might have.

1. Computing resources and running time.

Our experiments are conducted on a single 48G NVIDIA A6000 GPU, taking about 15 hours to align each model. For comparison, the constrained RLHF [R4] indicates 20 hours for each run on a more powerful NVIDIA A100 GPU. The computing resources and running time are not reported in safe-RLHF [12]. However, since safe-RLHF uses PPO for policy updates, like constrained RLHF, on a much larger dataset, we expect its running time to scale 2-3x (i.e., 40-60 hours) proportional to constrained RLHF. Constrained DPO [22] uses a single 40G NVIDIA A100 GPU without reporting the running time. Therefore, our methods reduce running time by at least 25% compared to the aforementioned methods while using a much cheaper GPU. We will include this discussion in future revisions.

2. Monotonically tuning the dual value.

Although intuitively tuning the dual variable simulates gradient descent/ascent updates to some extent, its efficacy highly depends on the unknown optimization landscape. To our understanding, the suggested strategy implicitly involves three undefined factors: (i) the optimal dual variable $\lambda^\star$; (ii) positive or negative constraint satisfaction; and (iii) annealing step size. Adjusting the dual variable depends on whether the primal variable (i.e., parameters of an LM) violates the constraint. However, the LLM policy, updated via sophisticated optimizers (e.g., DPO), does not necessarily ensure that the updated LLM policy always satisfies (or violates) the constraint. Hence, gradually reducing or increasing the dual variable, even starting from a sufficiently large or small initial value, may not converge to the optimum.

Moreover, small step sizes slow down the convergence speed of the dual variable, while large step sizes often result in severe oscillation of policy iterates. Therefore, iterative primal-dual algorithms often suffer from high computational burden and instability issues, as reported in Figure D.2 and the conclusion section of [R1], and Figure 2 of [22].

We would also like to note that the heuristic tuning strategy does not scale well to multi-constrained cases as the search space grows exponentially. In contrast, our method provides a principled and guaranteed methodology for multi-constraint cases.

3. Insufficiency of model-based evaluation.

We agree on the limitations of model-based evaluation, which is why we have the GPT-based evaluation. While human evaluation can be ideal, it is costly and subject to biases, as you mentioned. It seems there is no gold-standard evaluation for language models. We believe that human and AI evaluations should be complementary [R1].

4. Choices of safety margins

We chose the set of safety margins to achieve a diverse range of safety improvement levels. These margins were made primarily to better visualize the trend of safety improvement versus the dual variable in Figure 2 (left).

5. Code Availability

According to the rebuttal policy, we have shared our source code with the AC in an anonymized link. Please feel free to contact the AC for access to our code. We plan to officially release our code to the public after cleaning it up and adding detailed instructions.

6. Details on sample Responses

Thank you for your great interest in our experiments. We provide the GPT-evaluated safety scores of sample outputs upon the malicious red-teaming prompts on Pages 29--31 of the submitted manuscript as follows. We observe noticeable improvements brought by MoCAN- and PeCAN-alignment. However, we would like to remark that these sample prompts are handcrafted by literature [12] to mainly test safety performance. The helpfulness score evaluated on these prompts may not fully indicate the ground-truth helpfulness of each LM.

Table 1. GPT-evaluated safety levels for sample responses in the submitted manuscript.

Table 1. GPT-evaluated help levels for sample responses in the submitted manuscript.

7. Typos and Other Minor Suggestions

Thank you for your careful reading of our paper and for catching typos. We will fix these typos and double-check the paper's writing in revisions.

8. Social impact.

Thank you for bringing our attention to the broader scope of social impact and for pointing out an excellent reference on the dual use of AI. We will acknowledge the dual use of constrained alignment methods and remark on potential applications that could be negative to our society, such as dialogue systems with gender biases [R3].

References

[R1] Confronting Reward Model Overoptimization with Constrained RLHF

[R2] Complementarity in Human-AI Collaboration: Concept, Sources, and Evidence

[R3] GenderAlign: An Alignment Dataset for Mitigating Gender Bias in Large Language Models

We thank the reviewer for the very positive evaluation and the valuable feedback. We have answered all questions to the best of our ability, and we are glad to address any further questions you might have.

1. Experiments are limited to a single safety constraint.

As stated in the limitation section, our experiments use a single safety constraint due to the lack of suitable datasets. We are eager to test our method with multi-safety datasets if they become available. To the best of our knowledge, all existing works in safety alignment, such as [12, 22, 35], involve only one constraint, even in theory. In contrast, our algorithms and theoretical guarantees have advanced in dealing with multiple constraints.

2. Restriction of the Bradley-Terry preference setup.

We agree that the Bradley-Terry model may not reflect all ground-truth human preferences. However, we want to remark that our two-stage strategy is orthogonal to the preference setup. In fact, our two presented algorithms readily adapt to more general preference setups by generalizing the DPO optimization backbone to the more generic $\Phi$-PO (see Equation (6) of [R1]). Specifically, we can consider $r(x,y) = \mathbb{E}\_{y'\sim\mu}[\Phi(\mathcal{P}\_r(y\succ y'\vert x))]$ and $g(x,y) = \mathbb{E}\_{y'\sim\mu}[\Phi(\mathcal{P}\_g(y\succ y'\vert x))]$, where $\mathcal{P}\_r$ is the preference for helpfulness and $\mathcal{P}\_g$ is the preference for safety, $\Phi$ is a preference-based utility function, and $\mu$ is an underlying behavior policy. When $\Phi(p) = \log (p/(1-p))$, it recovers the standard Bradley-Terry preference setup (e.g., Proposition 1 of [R1]). We have remarked on the general preference setup in the conclusion section and left the extension for future work.

3. Potential sensitivity to pre-trained reward and safety models.

In practice, pre-training reward and safety models can only be estimated up to some error, as characterized by Definition 1. Theorem 3 shows that our strategy can find a nearly optimal LLM policy up to some estimation errors. Therefore, in theory, MoCAN enjoys stability against perturbations in reward/safety models.

Experimentally, we supplemented new experiments of MoCAN using the beaver-7b-v3.0-reward/cost models (see Table 3 in the PDF attached to the global response). It is observed that MoCAN-aligned LMs attain reasonably good performance under the new reward/cost models.

4. PECAN slightly underperforms MOCAN.

Please refer to the first point in our global response.

5. Potential negative effects or failure cases.

Potential negative effects can arise from misusing our constrained alignment method in applications that harm society, such as promoting gender biases [R2]. Our method assumes IID preference datasets and potential failures can be caused by out-of-distribution datasets [R3]. We will discuss potential negative effects and failures in future versions.

6. The comparison of computational complexity.

Our alignment methods turn to solve a one-shot unconstrained problem, while iterative primal-dual algorithms [12, R4, 22] must solve an unconstrained optimization problem for each round of dual update. Moreover, these algorithms need to generate a large batch of on-policy responses for evaluating the update of the dual variable, which is computationally expensive.

In practice, our experiments are conducted on a single 48G NVIDIA A6000 GPU, taking about 15 hours to align each model. For comparison, the constrained RLHF [R4] indicates 20 hours for each run on a more powerful NVIDIA A100 GPU. The computing resources and running time are not reported in safe-RLHF [12]. However, since safe-RLHF uses PPO for policy updates, like constrained RLHF, on a much larger dataset, we expect its running time to scale 2-3x (i.e., 40-60 hours) proportional to constrained RLHF. Constrained DPO [22] uses a single 40G NVIDIA A100 GPU without reporting the running time. Therefore, our methods reduce running time by at least 25% compared to the aforementioned methods while using a much cheaper GPU. We will include this discussion in future revisions.

6. Scaling with larger language models.

Due to resource and compute limits, we cannot experiment with our alignment methods on larger language models. To our knowledge, 7B models are commonly used to evaluate generation performance in recent NeurIPS papers (e.g., [R5, R6]). We believe it is acceptable to experiment with 7B models for fair comparison.

7. Exploration in other constrained ML problems.

Our dualization approach also applies to MinMax-RLHF [10] and alignment with $f$-divergence [R7]. Please check Appendix I and Appendix A for more details.

8. Resistance to adversarial attacks.

We thank the reviewer for highlighting this interesting direction. We are not aware of any common attackers that target safety constraints. However, one could naturally consider formulating the improvement of adversarial robustness into new constraints. Please refer to the paragraphs in our global response.

9. Suggestions for improvement.

Thank you for carefully reviewing our paper and suggesting several important points for improvement. We will discuss them point-by-point in revisions.

References

[R1] A General Theoretical Paradigm to Understand Learning from Human Preferences. Azar et al., 2023.

[R2] GenderAlign: An Alignment Dataset for Mitigating Gender Bias in Large Language Models. Zhang et al., 2024.

[R3] Align Your Prompts: Test-Time Prompting with Distribution Alignment for Zero-Shot Generalization. Hassan et al., 2023.

[R4] Confronting Reward Model Overoptimization with Constrained RLHF. Moskovitz et al., 2023.

[R5] Training language models to follow instructions with human feedback. Ouyang et al., 2022.

[R6] RRHF: Rank Responses to Align Language Models with Human Feedback. Yuan et al., 2023.

[R7] Aligning Language Models with Preferences through $f$-divergence Minimization. Go et al., 2023

We thank the reviewer for the positive evaluation and the valuable feedback. All questions have been answered as best as we could. We are glad to address any further questions you might have.

Additional benchmarks.

We have conducted additional experiments on the benchmark datasets TruthfulQA and AdvBench. Please refer to Tables 1 and 2 in the PDF attached to the global response for detailed results. These tables show that our aligned models achieve higher safety scores with increased $\lambda$ values, even when evaluated on these out-of-distribution prompts.

Dear Reviewer EtW9,

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

I would like to thank the authors for addressing my concerns, and I raised the acceptance rate from "weak accept" to "accept".

We thank the reviewer for the very positive evaluation and the valuable feedback. We have answered all questions as best we could, and we are glad to address any further questions you might have.

1. Incorporating adversarial robustness.

Please refer to the second point in our global response.

2. Ways to improve pre-alignment.

There are several ways to improve pre-alignment and, consequently, PeCAN's performance. As clarified in the global response, models can always be pre-aligned over larger datasets or using stronger reference LMs. With the same experimental resources, properly regularizing the pre-alignment process, as done in the training of reward/cost models in reference [12] (see Equations (13) and (14)), could enhance the reliability of the pre-aligned pseudo-preference labelers. Additionally, integrating more safeguard measures into the vanilla log-probability computation, as demonstrated in our supplementary experiments, could significantly improve the quality of pseudo-preference labeling and PeCAN's overall performance.

We thank the reviewer for the valuable feedback. We have addressed all your comments/questions to the best of our ability; see our detailed response below. We hope that our response would add your openness of re-evaluating our paper. We will be happy to address any further questions you might have.

1. The empirical gap between MoCan and PeCan.

Please see to the first point in our global response all the reviewers above.

2. No clear trade-off in MoCAN.

We note that the trade-off in MoCAN is evident in the model-based evaluation in Figure 3 (left), which is consistent with our theoretical findings. The less clear trade-off in the pairwise win-rate comparison in Figure 3 (middle) is somewhat understandable, as the win-rate performance does not directly correspond to the associated model-based performance. The model-based evaluation and the MoCAN algorithm are associated with proxy reward/safety models, while the pairwise win-rate evaluation is more closely tied to the ground-truth reward/safety and the specific baseline LM (i.e., the reference policy). Additionally, the GPT-based evaluation uses handcrafted prompts from [12] for the safety study and prompts from the Alpaca-eval dataset [20] associated with the "helpful_base" category, following the practice in [22]. The model-based evaluation uses prompts from the SafeRLHF-30K test set. These prompts come from different sources, leading to expected distributional shifts. In particular, the GPT-based evaluation result can be viewed as out-of-distribution performance, and thus, it may differ from the in-distribution performance seen in model-based evaluation. We also note that in related literature, such as [35], the helpfulness-safety trade-off under pairwise win-rate evaluation is unclear.

3. Infeasibility of straight-forward policy update.

There seems to be a misunderstanding regarding the alignment of language models. Sampling responses from a fixed and known policy is generally infeasible, even if the policy can be expressed in a closed form. In alignment, we aim to fine-tune language models, which contain billions of weights/parameters. Any change in the sampling policy for an LM must be implemented by updating its internal weights. Moreover, the input space (prompt tokens) and the output space (response tokens) of an LM policy can be viewed as infinitely large. Sampling in a large discrete space is notoriously difficult [R1, R2]. There is no direct way to force a language model to generate responses by exactly following a mathematically known policy. This practical impossibility is ubiquitous in LM alignment, even in unconstrained problems, departing from the common belief in classic reinforcement learning.

Regarding the reviewer's follow-up comment that "preference-based approaches are sparser than the value of $h$", we do not have a clear understanding. We kindly ask the reviewer for more clarification on this comment so that we can respond accurately.

4. Model-based evaluation of DPOs and RLHF.

We have supplemented the model-based evaluation results of DPOs and Safe-RLHF (see Figure 1 in the PDF attached to the global response). Our findings indicate that the model-based scores of DPO LMs align closely with those in Figure 3 (middle). Notably, Safe-RLHF exhibits better helpfulness vs. safety tradeoff, supporting our claim that there can occasionally be a discrepancy between model-based and GPT-based evaluations. The surprisingly good safety level of Safe-RLHF under in-distribution model-based evaluation may be attributed to the large dataset used in training (nearly one million data points, see https://huggingface.co/datasets/PKU-Alignment/PKU-SafeRLHF).

References

[R1] Oops I Took A Gradient: Scalable Sampling for Discrete Distributions. Grathwohl et al., 2021.

[R2] A Langevin-like Sampler for Discrete Distributions. Zhang et al., 2022.

We thank the reviewer for responding to our rebuttal and for being willing to raise the score. We would like a few more clarifications on your concern.

First, we would like to remark on the relation between $\pi_{\lambda}$, $\pi_{\theta_r}$, and $\pi_{\theta_g}$: $\ln\left(\frac{\pi\_{\lambda}(y_1 | x)}{\pi\_{\rm ref}(y_1 | x)}\right) - \ln\left(\frac{\pi\_{\lambda}(y_0 | x)}{\pi\_{\rm ref}(y_0 | x)}\right) = \ln\left(\frac{\pi\_{\theta_r}(y_1 | x)}{\pi\_{\rm ref}(y_1 | x)}\right) - \ln\left(\frac{\pi\_{\theta_r}(y_0 | x)}{\pi\_{\rm ref}(y_0 | x)}\right) + \lambda \left(\ln\left(\frac{\pi\_{\theta_g}(y_1 | x)}{\pi\_{\rm ref}(y_1 | x)}\right) - \ln\left(\frac{\pi\_{\theta_g}(y_0 | x)}{\pi\_{\rm ref}(y_0 | x)}\right) \right).$ However, this relation does not give the explicit formula of $\pi\_{\lambda}(y | x)$ due to the lack of normalization factors $Z_{r}(x)$, $Z_{g}(x)$ of $\pi_{r}(\cdot \mid x)$ and $\pi_{g}(\cdot \mid x)$, where $Z_{r}(x) =\int_{y}\pi_{\rm ref}(y \mid x)\exp(r(x, y)/\beta){\rm d} y$ and $Z_g(x)$ is defined similarly. Moreover, such normalization constants are intractable to compute due to the high-dimensional sampling space and input space (i.e., response tokens and prompt tokens).

Second, even assuming the probability oracle $\pi_{\lambda}(\cdot \mid x)$ is available, sampling from the oracle $\pi_{\lambda}(\cdot \mid x)$ remains way more challenging. The reason is multi-fold, including the high-dimensional sampling space and input space and the unmeasurable complex landscape of the probability distribution $\pi_{\lambda}(\cdot \mid x)$, unlike sampling from a common distribution like normal or uniform distributions.

We hope this response can clarify the reviewer's questions. We look forward to the follow-up discussion and are happy to address any further comments or questions.