Controllable Safety Alignment: Inference-Time Adaptation to Diverse Safety Requirements

We sincerely appreciate reviewer 6dSX for their thoughtful review. We are glad that you find our method flexible and our evaluation protocol novel. We hope our response addresses your questions and concerns:

Response to Weaknesses

The core contribution of the work is not clear. It is more like multiple small points mixed up.

Summarized in Line 92-96, our core contributions are introducing the Controllable Safety Alignment framework and formulating the task of efficiently adapting models to diverse safety requirements at inference time. Our framework rethinks the current paradigm of safety alignment and enables LLMs to be responsibly deployed in diverse settings.

To construct this comprehensive framework, we propose the CoSA-Score evaluation protocol, a human-curated benchmark CoSApien, and CoSAlign, a method for improving the controllability of LLMs. We believe all of these contributions are crucial for effectively adapting LLMs to diverse use cases. By providing this comprehensive set of artifacts, we enable rich future work on controllable safety alignment as pointed out in our discussion section.

While the method offers flexibility with safety configs, it may face challenges in generalizing across highly diverse or conflicting safety requirements.

We have already conducted experiments in Section 6.1 (discussed in Line 450-464) to evaluate the impact of the distribution shift from simpler safety configs to more complex ones. We refer the reviewer to that section of a detailed discussion. We summarize it here again for your convenience: we train on CoSAlign-Train, which only contain synthetic categorical risk categories, and test on CoSApien, our proposed benchmarks consist of complex, fine-grained real-world configs (see examples in Appendix A.12). Human evaluation (Table 4) shows that CoSAlign still possesses strong controllability in safety configs that are more complex than those the model was trained on, and outperforms all baselines. This shows CoSAlign’s strong generalization ability from simpler training configs to more complex test configs.

Finally, if there is a specific section you found unclear or uneasy to follow, please don’t hesitate to let us know!

Response to questions

Section 3: The evaluation protocol is like the one shown in "Lin, Stephanie, Jacob Hilton, and Owain Evans. "Truthfulqa: Measuring how models mimic human falsehoods." arXiv preprint arXiv:2109.07958 (2021)." What are the advantages of your protocol?

We’d like to clarify that while both our protocol and TruthfulQA use LLM as a judge, CoSA-Score use two separate judge-help and judge-safe to measure helpfulness and configured safety, and summarizing them into a single score that take account of both aspects. We argue this is a novel way to depict the tradeoff between helpfulness and safety aspects that allows distinguishing nuances such as a response that is both unsafe and helpful should be penalized more than a response that is unsafe but not helpful (for example, if the prompt requests criminal advice, a highly helpful response would likely cause greater harm to society than one that is less helpful).

Section 5: The CoSAlign pipeline is more like a rule-based method. I doubt its generalization ability in real-world deployment. Moreover, this part is kind of messy and not easy to follow.

We’d like to clarify that because CoSAlign selects response pairs through the error score derived from LLM judge outcomes, it is not a rule-based method. On the other hand, a recent work on rule-based rewards has recently shown to be effective for controlling nuanced aspects of safety [1], so we agree there is indeed room to explore in the rule-based space.

Related to CoSAlign’s generalization ability in real-world development, we have created the human-authored benchmark, CoSApien, with the help of professional red teaming specialists who are adept at real-world safety scenarios (detailed in Section 4). Thus we believe that CoSApien covers common safety scenarios in practice. In Table 4, Section 5.2, human evaluation shows that CoSAlign not only maintains high controllability but also outperforms all baselines, demonstrating the real-world applicability of our method. Due to the space constraints, we had to move some details of CoSAlign to the appendix, but we will continually iterate to improve the clarify in this part.

We also conducted additional experiments on adversarial robustness, over-refusal, and combining CoSAlign with Cascade methods. Please find more details in the “Response to all reviewers” section. Thank you!

References

[1] Rule Based Rewards for Language Model Safety

Response to questions

Performance on jailbreak method & ASR: Per your suggestion, we have conducted additional experiments on 3 popular jailbreak attacks and shown the results above in the “Response to all reviewers” section. We find that CoSAlign lead to significantly improved adversarial robustness against these attacks.

Oversafety performance: We conduct experiments on XSTest, shown in the “Response to all reviewers” section, and find CoSAlign lead to notably less refusal compared to Llama3.1-8B-Instruct. This indicates the increased controllability can allow model to better determine the safety boundary for refusal.

Thank you for your suggestions of the related work on efficient alignment. Note that we already have a focused discussion on inference-time alignment which focuses on efficient training-free approaches, but we will add these to the manuscript for further completeness.

References

[1] A General Theoretical Paradigm to Understand Learning from Human Preferences

[2] KTO: Model Alignment as Prospect Theoretic Optimization

[3] SimPO: Simple Preference Optimization with a Reference-Free Reward

[4] Llama Guard: LLM-based Input-Output Safeguard for Human-AI Conversations

[5] BeaverTails: Towards Improved Safety Alignment of LLM via a Human-Preference Dataset

[6] WildGuard: Open One-Stop Moderation Tools for Safety Risks, Jailbreaks, and Refusals of LLMs

We sincerely thank reviewer hHMw for their thoughtful review. We are grateful that you find our controllable safety alignment method is “valuable for applications with varied cultural, legal, or organizational safety needs” and “[reduces] manual annotation needs for large-scale applications.” We hope our response below addresses your concerns:

Response to Weaknesses

On error-scoring mechanism assign arbitrary penalties: We’d like to clarify that the purpose of the error penalties is to ensure responses that do not violate the safety configs and maintains helpfulness, are preferred over responses that do not satisfy these two criteria. It is only used in the response pairing stage and only the relative error between two responses matters. For example, suppose there are responses A, B, and C:

• Response A contain one disallowed risk but address the question -> error $\beta$
• Response B contain one allowed risk but does not address the question -> error $\alpha+\gamma$
• Response C contain one allowed risk and address the question -> error $\alpha$

As long as $\alpha < \gamma < \beta$ is satisfied, we will prefer Response C over A, C over B, and B over A no matter what the absolute value of these hyperparameters are. We set $\alpha=0.1, \gamma=1, \beta=3$. Further tuning these hyperparameters might lead to additional gains in rare cases (e.g., when the response contains many types of allowed risks), but we empirically find that CoSAlign is already very effective (see Table 3, 4, 5 in Section 6).

On convergence properties of preference optimization: We respectfully disagree with the reviewer that this should be listed as a weakness of our work. We do not discuss convergence properties or other math details simply because we are not proposing any preference optimization algorithm. Instead, we are just utilizing the existing DPO algorithm, one of the most commonly used and widely adopted algorithms for preference optimization, in our data-centric CoSAlign method. We clarify that our main contribution in this work (summarized already in Line 92-96) is proposing a comprehensive framework (defining the task, setup, evaluation protocol, benchmark, and the CoSAlign method) on controllable safety alignment. While the convergence properties of preference optimization are an active research area, it is not the focus of this work. Moreover, our framework makes no assumptions of the preference optimization technique used. It is straightforward to utilize more advanced algorithms such as IPO [1], KTO [2], and SimPO [3] with CoSAlign.

On input distribution shift for safety configs: Indeed, the distribution of safety configs may change between training and real-world deployment, and maintaining sufficient controllability under config distribution shift is an important issue, and we already conducted experiments in Section 6.1 (discussed in Line 450-464). We summarize it here again for your convenience: we conducted experiments on config distribution shift by training on CoSAlign-Train, which only contain synthetic categorical risk categories, and test on CoSApien, our proposed benchmarks consist of complex, fine-grained real-world configs (see examples in Appendix A.12). Human evaluation (Table 4) shows that CoSAlign still possesses strong controllability in this out-of-distribution setting and outperforms all baselines. This is strong evidence that CoSAlign can generalize from simpler training configs to more complex test configs. Please see Section 6.1 for more details.

It is difficult to theoretically quantify the controllability of CoSAlign in all settings, where safety configs can be arbitrarily complex. Since our focus and main contributions are proposing a comprehensive framework of controllable alignment, we leave the theoretical aspect of quantifying controllability to future work.

On CoSAlign’s risk taxonomy: We’d like to clarify that the risk taxonomy is only used during the training of CoSAlign in order to synthesize large-scale diverse and relevant preference data. Although this risk taxonomy is likely not exhaustive, we show in Table 4 that CoSAlign maintains controllability gain on the real-world CoSApien dataset and generalizes to novel risks that are more fine-grained than the training categories. Moreover, Table 3 and 5 shows CoSAlign is still effective on unseen synthetic configs (configs that contain risk categories held-out from training). These results provide plenty of evidence that our current risk taxonomy is diverse enough to not overfit to training categories. Using a risk taxonomy with a finite set of categories is also a common approach in recent works such as Llama Guard [4], BeaverTails [5], and WildGuard [6]. Therefore, we argue that risk taxonomy is an effective practical approach for large-scale data synthesis in controllable safety alignment.

Dear reviewer hHMw,

We greatly value your feedback and have provided clarifications to your questions and additional experiments on jailbreaking and over-refusal as you requested. To ensure that we have properly addressed your concerns, we would greatly appreciate if you could review our responses and provide any further comments. We look forward to engaging with you before the discussion period ends.

Thank you for your time and consideration.

Dear Authors,

Thank you for your clarifications. Based on the current state of the work, I am inclined to slightly increase my score. However, I fell that the key points I raised should have been addressed in the first version hence not able to increase much.

I recommend that the authors incorporate these discussions and the referenced citations into the appendix in future revisions.

Thanks.

Line 268: Are the training prompts corresponding to or independent of any safety configs used in the training? Can the risk taxonomy made in the paper used in place of the other prior taxonomies, or is it specially tailored to CoSA?

In our current pipeline, the training prompts correspond to the training safety configs because we first derive the risk taxonomy base on the training prompts (Appendix A.2), and then synthesize relevant safety configs (Line 302). But in principle, the training prompts do not need to curate together with the risk taxonomy as long as the taxonomy covers a broad range of risks relevant to the training prompts. Our derived taxonomy is not specifically tailored to CoSA and can be used in place of other prior taxonomies.

Is $C_{i, j} \in R$, in line 313?

Yes, each config risk category is a subset of the taxonomy $R$.

How is judge-help conceptually different from judge-addr?

While judge-help evaluates the helpfulness of a response in detail and gives a score of 0 to 5, judge-addr only considers whether the response is a refusal or not. It does not consider “how well” the model answers the response.

Line 350: Why are allowed risks penalized by $\alpha$? I think they should be rewarded, as the model is adhering to the safety config.

Great question! We argue that models should only use allowed risks as needed in order to achieve better helpfulness. For example, if violent content is allowed in the videogame setting, the model should not be rewarded for producing violent descriptions on prompts where violence is not needed. Therefore, we give a small but positive penalty for allowed risks.

Line 363: Why is the adversarial partition of WildGuardTrain removed from the training set?

We made this design choice to limit the number of training prompts due to our compute resources. Nevertheless, as shown in additional results in “Response to all reviewer” section, CoSAlign achieves increased robustness against adversarial attacks.

Line 407: Does "helpful" mean that the judge-help outputs anything > 0 or equal to 1? Similarly, what values of the helpfulness scores from humans are considered to be "helpful" in Table 4?

We’d like to clarify that (detailed in Appendix A.6) for both the GPT-4 and human judges, raw scores are given on the scale of 0 to 5. A response is considered helpful if judge-help output $\geq$ 1. This ensures consistent results between GPT-4 and human judges.

References

[1] Investigating Cultural Alignment of Large Language Models

[2] Constitutional AI: Harmlessness from AI Feedback

[3] Collective Constitutional AI: Aligning a Language Model with Public Input

[4] A proposal for importing society’s values

[5] We're Afraid Language Models Aren't Modeling Ambiguity

On “The experiments are limited to Llama-3.1 8B and GPT-4o models”: We experiments on four model variants, Llama-3.1-8B-Instruct, Llama-3.1-8B-SFT, GPT-4o, and GPT-4o-mini, within the Llama and GPT model families, and covers both popular open-source and proprietary models. Thank you so much for your suggestions of more models. We believe that the Llama model family is a good representation of open-source models, thus we don’t see an urgent need to experiment with a similar open-source model Mistral, but we’ll take that into consideration.

I would suggest that the authors add some justification for not doing DPO on GPT-4o around line 472.

Thanks for mentioning this point — we actually already have a footnote on page 9 that clarifies only SFT is publicly available for GPT.

The authors should discuss the generally lower (hence better) helpful+unsafe values of the Cascade method over CoSAlign.

We have added additional results of CoSAlign+Cascade methods and included extended discussions. Please see the “Response to all reviewers” section for more details.

Thank you so much for your comments on improving the terminology use & fixing typos. We will make sure to fix these issues in the manuscript.

Response to questions

On ambiguous or contradictory safety configs: While it’s challenging to formally define what configs are ambiguous and modeling ambiguity in general [5], qualitatively we find that when a CoSAlign-tuned model face prompts that are under-defined in config, or low-quality configs in general, the model will “default” to standard safety behavior as achieved by the one-size-fits-all aligned base model. We find this phenomenon very interesting because it shows that the model is protected by regular safety alignment when configs do not apply.

Line 114-115: "However, because of the complexity of safety configs and the difficulty of constructing high-quality demonstrations at scale" - the complexity of safety configs is a problem for this work too. Moreover, in-context learning requires only a handful of demonstrations, hence the argument of absence of demonstrations at scale is void. Hence, I think these arguments against in-context learning are not very strong.

We’d like to clarify that the difficulty for in-context learning lies in the issue of constructing a set of demonstrations that fully covers the desired safety config. Given real-world safety configs as those in CoSApien, we find it very difficult to fully specify the desired safety behavior with examples alone. For example, “allowing violence but excluding descriptions of severed body parts or limbs” requires multiple carefully designed demonstrations but is easily described in natural language.

Line 183: How do you check/verify coverage? Line 405: Manually checking that the test set contains all the 3 kinds of prompts is ok, but what is ensuring that to be the case in general for all the prompts? As far as I understand, it is just a dataset containing some prompts, not necessarily according to some safety configs.

We’d like to clarify that both test sets are test safety configs paired with prompts relevant to the configs (detailed in Line 173 to 183). As detailed in Appendix A.12, during the data curation stage for CoSApien, we produce three types of targeted prompts (allowed, disallowed, partial) for each test config and manually verified all prompts adhere to the desired types. For CoSAlign-Test, we conduct human verification of the automatically produced prompt risk category labels on a subset of 600 prompts and find a high human agreement rate of 89.8%. We then use this category label as proxy to select prompts and ensure each test config is covered. This is detailed in A.7 where we provide the full breakdown by prompt types.

How would GPT-4's own safety alignment affect its responses as judge-safe(.) in line 188?

We kept track of the proportion of judge-safe requests that are refused by GPT-4, and find this number to be very low in all cases (less than 2% of responses). We also qualitatively find that when GPT-4 does not produce a refusal, it gives reasonably good rationales and results. Note that besides GPT-4 evaluation, we also conduct human evaluation on CoSApien and find consistent results on judge-safe.

Can the CoSA-Score be made more useful by assigning -1 to refusals/unhelpful responses?

Thanks for the suggestion! We believe it’s better to give CoSA-Score to refusals/unhelpful response because these responses should be more preferred than responses that are both helpful and unsafe under the current config, which lead to negative scores. A responsible system should never give unsafe responses that are helpful to conducting harmful activities, but if the user query is not answerable given the current safety guidelines, a refusal is a reasonable response.

We genuinely appreciate reviewer pCoa for their insightful review. We’re thankful that you recognize our work contributes to the “practical importance for effective and safe usage of LLMs in diverse applications,” and our paper is well-supported by human evaluations. We address your questions and concerns below:

Response to weaknesses:

On “It may not be feasible to specify all aspects of safety, even by domain experts, in natural language”: We’d like to clarify that safety configs do not necessarily need to specify “all aspects of safety”: just like what we do in our synthetic configs (exemplified in Appendix A.11), the safety config be used as providing exceptions to regular “one-size-fits-all" alignment, indicated by the instruction “On the other hand, you must continue to ensure safety in all other categories.” Therefore, only the difference between the desired safety behavior and a “one-size-fits-all” safety aligned model need to be specified, greatly reducing the overhead of producing safety configs.

In practice, we show that starting CoSAlign from a one-size-fits-all safety aligned model can effectively retain safety on input distributions that are not directly specified in the safety config. In this case, the model will “default” to the one-size-fits-all safety behavior. In our new experiments shown in “Response to all reviewers” section, we find that CoSAlign can significantly increase adversarial robustness against popular jailbreak attacks and reduce over-refusal, while these behaviors are not directly specified in the safety config at all.

On the difference between our work and pluralistic alignment in non-safety settings: Safety is a special setting because while current safety alignment techniques make the model robust to adversarial attacks, they also severely curtail the steerability of models on safety. Instructing models to modify safety standards can be seen as a jailbreak, thus not possible with one-size-fits-all safe models. We experimentally show this is indeed the case due to the ineffectiveness of in-context alignment in Section 5.1. On the other hand, for non-safety settings such as cultural alignment, techniques such as Anthropological Prompting [1] has shown to be an effective baseline. However, similar prompting-based techniques do not apply to the safety settings, motivating our CoSA framework and CoSAlign approach. Additionally, we conduct thorough safety-specific evaluation and propose metrics, benchmarks and data synthesis methods that are tailored to the safety-specific settings. All these contributions are novel and essential for studying pluralistic safety alignment.

On determining disallowed content: Similar to your first point on “feasibility to specify all aspects of safety”, we’d like to clarify that since we start with a one-size-fits-all safety aligned model, the safety guideline only need to specify the difference between the “default” safety behavior. This design makes the model robust to prompts not specified in the config, as shown in results in the “Response to all reviewers” section.

In general, determining what should be the disallowed content is a key sociotechnical problem and ongoing research. Approaches such as Constitutional AI [2], Collective Constitutional AI [3], and Simulated Deliberative Democracy [4] have been proposed to import societal value to determine what should be considered allowed/disallowed. We propose an alternative approach where this boundary can be flexibly and efficiently adapted during test time by adjusting the specification in safety config.

I think it may be advantageous to evaluate CoSApien also with the kinds of content that an actual LLM can generate for the different kinds of prompts, for a real proof of its prompts being effective.

We’d like to clarify that our CoSApien benchmark curated real-world safety configs by professional red teaming specialists who are adept with real-world safety scenarios. We also collected diverse human-written prompts for each config, which are natural and what actual LLMs are often asked about. Please refer to Section 4 for more details.

On details about risk taxonomy construction: due to space constraints, we had to move some of the details to appendix, and please see a detailed description in Appendix A.2. To create this risk taxonomy, we first perform prompt clustering, and produce cluster descriptions with GPT-4o. Next, we manually edit the descriptions of largest clusters and produce the final risk taxonomy. We will also make the reference to the appendix clearer in the main text.

We deeply appreciate reviewer c6ps for their thoughtful feedback. We are thankful for your recognition of the practical relevance of our topic, the advantages of our proposed alignment method, and most interestingly, the framework of obtaining controllability through safety configs. We now address your questions and concerns below.

Response to weaknesses

For example, a negative score can rise without safety being altered at all simply by making the model less helpful (if helpfulness drops to zero, the score becomes zero which is better than negative). Similar things can also happen with positive scores

Thanks for highlighting this point! We argue that instead of a weakness, this is characteristic actually a benefit of CoSA-Score: We believe that unsafe responses that are less helpful should indeed be preferred over unsafe responses that are more helpful. For instance, if the prompt requests criminal advice, a highly helpful response would likely cause greater harm to society than one that is less helpful. By reducing helpfulness in unsafe responses, the potential negative impact is mitigated, aligning with our goal of promoting safer model behavior. Importantly, it is only by considering safety and helpfulness together, as CoSA-Score does, that such nuanced trade-offs can be effectively addressed. For positive scores, CoSA-Score's construction allows safe responses that are more helpful to be preferred over safe responses that are less helpful, as one would intuitively expect. We’d love to provide further clarifications and engage in further discussions on this design if you are interested.

Furthermore, since the helpfulness scores are between {0,1,2,3,4,5}, you may lose a lot of information [...] Thus the other numbers reported in the tables are important as they do remove this ambiguity - by reporting helpful+safe and helpful+unsafe you get a complete picture.

We agree on your point that there might be more than one situation where the aggregated CoSA-Score is the same (e.g., all responses are somewhat helpful v.s. half responses are very helpful & the other half is not helpful at all). This is the tradeoff we are making by summarizing safety and helpfulness into a single score. We believe this design is still very meaningful because it allows us to (1) obtain an aggregated metric across all safety configs and prompt distributions to measure overall controllability (2) take into consideration of the nuanced trade-offs between helpfulness and safety together (related to the first point above). As you mentioned, we also report the helpful+safe and helpful+unsafe responses as focused metrics on positive-scored and negative-scored responses, making it easier to capture the complete picture. Thanks so much for pointing out these important nuances and we will add corresponding discussions to our draft.

Higher helpfulness and lower safety (relative to cascade methods)

We acknowledge that CoSAlign has a drastically higher rate of helpful+unsafe responses at the cost of slightly higher rate of helpful+unsafe responses relative to the Cascade methods. Note that this result is not so surprising, given that Cascade methods specifically focus on improving safety by filtering out unsafe responses deemed by a classifier and replacing them with refusals. Therefore, in this case, the balance between helpfulness and safety is very important because a very safe but not helpful response is useless. As determined by the much higher CoSA-Score, which is specifically designed to capture the trade off between helpfulness and safety, CoSAlign clearly outperforms Cascade method. Also note that Cascade-Oracle is a very strong baseline that uses the evaluator model to conduct filtering, since it’s guaranteed to have 0% unsafe responses. Even in this case, CoSAlign outperforms Cascade-Oracle because there is a limited rate of helpful+safe responses.

Inspired by your suggestions, we have also applied the Cascade methods on top of CoSAlign methods and Cascade-Oracle demonstrated additional gains. Please see “Response to all reviewers” section for details.

On related works: Thank you so much for your suggestion. We agree that these can be a nice addition to our related work discussion. We will add a dedicated paragraph on balancing safety and helpfulness.

Response to questions

Is there some way to combine the cascades method with the CoSAlign method?

Yes! The Cascade method can be naturally incorporated into any model because it adds a post-hoc filtering step, which replaces unsafe responses deemed by the filtering model with refusals. This is similar to what deployed systems usually do (e.g., OpenAI’s content filter on ChatGPT). Although the filtering model is usually smaller than the generator model, we use (1) the generator model itself (2) the evaluator model as filtering model to construct this strong baseline. We have included results of combining Cascade with CoSAlign in the “Response to all reviewers” section.

Dear Reviewer c6ps,

We sincerely appreciate your thoughtful feedback and have provided clarifications to your questions, including conducting new experiments on combining CoSAlign with Cascade methods as per your request. Your insights have been very constructive, and we would greatly value your review of our responses to ensure we have fully addressed your concerns.

Please don’t hesitate to share any additional comments and we are eager to engage with you further before the discussion period ends. Thank you for your time and consideration.

Additional experiments on combining CoSAlign and Cascade methods

In this section, we show that Cascade methods can be effectively incorporated into CoSAlign-tuned models to trade off helpfulness for safety.

Because the Cascade methods use a filtering model to label unsafe responses and replace them with refusals, they can also be applied on top of CoSAlign. Below, we demonstrate the effectiveness of CoSAlign+Cascade on the seen split of CoSAlign-Test (the unseen split follows the same pattern):

Results show that, similar to applying Cascade on the base Llama-3.1-8B model, applying Cascade on CoSAlign-tuned model can also reduce the rate of helpful+unsafe responses. Cascade lowers helpful+unsafe responses at the cost of decreasing helpful+safe responses, trading off helpfulness for better safety. We acknowledge that while CoSAlign+Cascade led to slightly more helpful+unsafe responses than Llama3.1-8B-Instruct+Cascade, the gap is small and CoSAlign lead to significantly improved helpful+safe responses, and thus a much higher CoSA-Score.

In summary, applying Cascade-Oracle on CoSAlign achieves the highest CoSA-Score, the highest helpful+safe responses, and the lowest helpful+unsafe responses across the board, demonstrating the effectiveness of combining both methods.

References

[1] A StrongREJECT for Empty Jailbreaks

[2] GPTFUZZER: Red Teaming Large Language Models with Auto-Generated Jailbreak Prompts

[3] Jailbreaking Black Box Large Language Models in Twenty Queries

[4] Tree of Attacks: Jailbreaking Black-Box LLMs Automatically

[5] HarmBench: A Standardized Evaluation Framework for Automated Red Teaming and Robust Refusal

[6] XSTest: A Test Suite for Identifying Exaggerated Safety Behaviours in Large Language Models