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

1 INTRODUCTION

• One-size-fits-all policies miss variation across cultures, applications, and users, while re-aligning separately for user groups is prohibitively expensive.Safety norms can vary across cultures and evolve over time; specialized users may also need different content treated as safe.Source: §1 · pp. 1–2
• CoSA adapts safety at inference time by placing authorized users’ natural-language safety configs in the system prompt, without re-training.The framework returns adapted models through custom interfaces such as user-specific API endpoints.Source: §1 · p. 2
• The proposed CoSA-Score evaluates both response helpfulness and configured safety, and CoSApien provides a human-authored benchmark of real-world safety scenarios.CoSApien includes evaluation prompts for those scenarios.Source: §1 · p. 2
• CoSAlign improves controllability over strong baselines for seen safety configs and generalizes well to unseen configs.The paper reports that in-context alignment is insufficient for complex configs and high-quality demonstrations at scale.Source: §1 · p. 2
• CoSA is presented as a framework for pluralistic safety alignment that can serve users with diverse safety requirements.The authors describe it as a blueprint for inference-time adjustment without re-training.Source: §1 · p. 2

Pluralistic alignment

• Prior work explores pluralism in values and cultures, while this paper frames pluralistic safety alignment as a distinct, context-dependent focus.The authors contrast this focus with work on pluralism generally and on one-size-fits-all models in pluralistic settings.Source: §2 · p. 3
• Constitutional AI trains a single constitution into a one-size-fits-all model, whereas CoSA adapts one controllable model to changing safety requirements without further training.The paper describes a constitution as a set of universal ethical principles.Source: §2 · p. 3

Inference-time and multi-objective alignment

• In-context alignment is insufficient for complex safety configs, while retraining, parameter merging, and decoding-time methods do not offer the same efficient adaptation to new configs.Decoding-time approaches require a new reward function for each objective, and attribute steering does not provide fine-grained control through safety configs.Source: §2 · p. 3

Controllability through instruction

• Instruction hierarchy addresses prompt-injection attempts to modify configs without authorization, but the paper reports it does not improve controllability; rule-based rewards offer no on-the-fly adaptation.The paper describes instruction hierarchy as complementary to its framework.Source: §2 · p. 3

3 COSA: A FRAMEWORK FOR CONTROLLABLE SAFETY ALIGNMENT

• CoSA first trains a model to follow free-form natural-language safety configs supplied as a system-prompt prefix.The model is trained on diverse configs so its responses adhere to the specified guidelines.Source: §3 · p. 4
• Providers can create user-specific interfaces by adding each user’s config to the prompt, allowing configs unseen during training while reserving config changes for authorized users.The framework calls for provider review and says the controllable model should not be exposed directly to end users.Source: §3 · p. 4

3.1 EVALUATING SAFETY CONTROLLABILITY

• The evaluation protocol tests allowed, disallowed, and partial-content prompts, and assesses both helpfulness and adherence to the active safety config.The protocol’s CoSA-Score gives positive contribution to helpful, safe responses and negative contribution to helpful but unsafe responses; unhelpful responses do not affect it.Source: §3.1 · p. 4

• CoSApien covers five real-world safety configs with 40 prompts each, totaling 200 human-authored test prompts.Source: §4 · p. 4
• Two professional red-teaming specialists drafted 10 candidate scenarios, then refined them to five to reduce redundancy and cover safety and cultural dimensions.Source: §4 · pp. 4–5
• Each scenario spans allowed, disallowed, and partial prompts; independent specialists reviewed scenarios and prompts before consensus refinement.Source: §4 · p. 5
• Configs encode fine-grained boundaries, such as allowing violent content while excluding depictions of severed body parts or limbs.Source: §4 · p. 5
• CoSApien was constructed independently from CoSAlign and is described as the first general-purpose benchmark for controllable safety.Source: §4 · p. 5

5.1 IS IN-CONTEXT ALIGNMENT SUFFICIENT FOR CONTROLLABLE SAFETY?

• ICA gains on CoSAlign-Test are modest and saturate before diminishing around 15 shots; Figure 3 characterizes ICA as ineffective under complex configs.Source: Figure 3 · p. 5, §5.1 · p. 5
• On GPT-4o-mini, ICA decreases CoSA-Score, indicating instruction hierarchy alone does not enable effective adherence to safety configs.Source: Table 1 · p. 5, §5.1 · p. 5
• With few-shot ICA, INSTRUCT remains too restrictive while SFT-only often violates configs; high-quality examples are costly to develop and increase inference cost.Source: §5.1 · pp. 5–6, Table 2 · p. 6

5.2 COSALIGN: A DATA-CENTRIC METHOD FOR CONTROLLABLE SAFETY ALIGNMENT

• For training data, CoSAlign classifies prompts and samples config-risk subsets across four prompt/config relationships, then uses 10 templates to produce 256 categories and 2,560 configs.Source: §5.2 · pp. 6–7
• Using CoSAlign categories yields 83% agreement versus 65% with BeaverTails categories; its prompt-derived taxonomy includes an “Other Harms” category.Source: §5.2 · p. 6
• For each config-prompt pair, generators produce responses that LLM judges assess for risk and whether they address the prompt; error scores penalize allowed risk lightly, nonresponse moderately, and disallowed risk heavily.Source: §5.2 · p. 7
• CoSAlign pairs a response without disallowed risks as chosen with a higher-error response as rejected, then applies SFT and DPO to produce the controllable model.Source: §5.2 · pp. 6–7

6.1 CONTROLLABILITY EVALUATION

• On CoSAlign-Test, CoSAlign significantly outperforms all baselines on CoSA-Score, increasing helpful+safe responses while maintaining a low rate of helpful+unsafe responses.The preference-optimization stage reduces helpful+unsafe responses while further increasing helpful+safe responses; SFT alone tends to be overly helpful on unsafe responses.Source: Table 3 · p. 8, §6.1 · p. 9

6 EXPERIMENTS AND EMPIRICAL FINDINGS

• CoSAlign-Test contains 3,200 prompts across eight configs, including three unseen configs; CoSAlign-Train contains 122K preference pairs.The test set is automatically built from BeaverTails prompts paired with configs and covers allowed, disallowed, and partial prompts.Source: §6 · pp. 7–8

6.1 CONTROLLABILITY EVALUATION

• On CoSApien, human evaluation finds CoSAlign consistently outperforms strong in-context and cascade baselines, including Cascade-Oracle, on overall CoSA-Score.CoSApien contains real-world prompts that are less risky than the red-teaming prompts in CoSAlign-Test, and its CoSA-Score results are generally higher.Source: Table 4 · p. 9, §6.1 · p. 9
• CoSAlign (SFT only) improves GPT models’ controllability over safety removal, which retains general instruction following but still fails to follow safety instructions effectively.The GPT-4o MT-bench score changes from 9.03 to 8.98 after safety removal; the experiments do not apply DPO to GPT models.Source: §6.1 · p. 9, Table 5 · p. 9

6.2 GENERAL CAPABILITY, INSTRUCTION FOLLOWING, AND SAFETY EVALUATION

• CoSAlign causes minimal degradation in general capability, slightly improves general safety, and significantly improves safety controllability.Table 6 summarizes general capability and safety benchmark results.Source: Table 6 · p. 10, §6.2 · p. 10

6.3 GAP TO PERFECT CONTROL — ERROR ANALYSIS OF DISALLOWED CONTENT

• On CoSApien, CoSAlign significantly outperforms the SFT baseline across all three prompt types, but both models still produce unsafe responses, especially on the rare partial category.The partial category contains 21 of the 200 prompts and covers requests where only part can be safely fulfilled.Source: §6.3 · p. 10
• CoSAlign does not achieve perfect control: disallowed risks remain, and unseen configs based on held-out categories have slightly worse controllability.“Other Harms” has the highest rate of disallowed responses, possibly because it covers rare, distinct risks.Source: Figure 5 · p. 10, §6.3 · p. 10

7 DISCUSSION AND FUTURE WORK

• Direct use of controllable models should be limited to authorized users who can modify safety configs through a config review process.The paper presents this as a deployment consideration for CoSA.Source: §7 · p. 10
• The experiments focus on data-centric CoSAlign; future work could explore preference-learning algorithms using diverse synthetic configs or representation engineering.These are presented as alternative angles for improving controllability.Source: §7 · p. 10

A.1 ETHICAL IMPLICATIONS

• The framework aims to represent pluralistic human values more flexibly, thereby increasing LLM practicality.The authors advocate adapting safety requirements at inference time.Source: §A.1 · p. 17
• Misaligned or malicious safety configs could bypass ethical constraints and lead to harmful outcomes, making robust guardrails essential.Source: §A.1 · p. 17
• CoSA raises the ethical challenge of deciding whose values differing societal norms represent and prioritize in AI safety.The authors call for ongoing, transparent dialogue about AI governance, responsibility, and fairness in deploying adaptable safety mechanisms.Source: §A.1 · p. 17

A.2 LIMITATIONS

• Prompt injection can alter safety configs and prompt extraction can expose them; the authors report robust general-safety evaluations and expect instruction-hierarchy fine-tuning to further increase robustness.Source: §A.2 · p. 17
• The authors did not systematically explore how CoSAlign scales across different model sizes.Source: §A.2 · p. 17
• The framework covers safety and cultural alignment describable in natural language, excluding implicit cultural and social norms.Source: §A.2 · p. 17

A.3 RISK TAXONOMY CREATION

• Training prompts are embedded, reduced in dimensionality with UMAP, and clustered with HDBSCAN before the largest clusters are summarized by an LLM.The process is presented in Figure 6.Source: §A.3 · p. 17, Figure 6 · p. 17
• For manual validation, the authors sampled 100 prompts from each cluster and checked that the edited risk definition described those prompts correctly.Source: §A.3 · p. 17
• The process produced a taxonomy with eight distinct risk categories, whose full definitions are provided in §A.9.Source: §A.3 · p. 17

A.4 EXPERIMENTAL SETUP DETAILS

• A fine-tuned safety-removed GPT-4o model and off-the-shelf GPT-4o were used as data generators.The fine-tuned endpoint produced GPT-4O-SAFETY-REMOVED.Source: §A.4 · p. 17
• Experiments used Llama-3.1-SFT, BeaverTails, and WildGuardMix, and ran on four NVIDIA A100 80GB GPUs.The passage names the sources and licenses for BeaverTails and WildGuardMix.Source: §A.4 · pp. 17–18
• One human annotator assessed responses from CoSAlign and Llama-3.1-8B-SFT for 200 prompts across five safety configs.The annotator rated configured safety and helpfulness on a 0–5 scale.Source: §A.4 · p. 18

A.5 COSALIGN DETAILS

• LoRA fine-tuning yields less controllability gain than full fine-tuning, suggesting safety controllability is parameter-intensive at 8B scale.CoSAlign samples responses under varied safety configs, scores allowed and disallowed risks and whether the prompt is addressed, then pairs lower-error responses as preferred. Training uses SFT and DPO with RMSPromp, a 5e-7 learning rate, and DPO β = 0.1.Source: §A.5 · pp. 18–19, Table 7 · p. 19

A.6 DETAILS ON GENERAL CAPABILITY AND SAFETY EVALUATION

• The general evaluation covers five capability and four safety benchmarks, using default lm-evaluation-harness settings for capability and safe-response rates for safety.Safety responses are labeled with Llama-Guard-3-8B, and evaluation uses the system message “You are a helpful assistant.”Source: §A.6 · pp. 19–20

A.7 DETAILS OF GPT-4 EVALUATORS FOR COSA-SCORE ON COSALIGN-TEST

• GPT-4o judges helpfulness on a 0–5 scale and marks a response safe when every detected risk category is allowed by its safety config.Human evaluation on CoSApien generally agrees with GPT-4o results, supporting the reliability of the automatic judges.Source: §A.7 · p. 20

A.8 COSALIGN-TEST CONSTRUCTION

• CoSAlign-Test groups allowed, partially allowed, and disallowed prompts under configs that permit different combinations of risk categories.The listed configs include allowing no risks, violence, financial crime and theft with privacy violations, or violence with weapons-related harms.Source: §A.8 · pp. 21–22

A.9 DEFINITIONS OF RISK CATEGORIES

• The risk-category definitions cover financial crime and theft, discrimination and verbal abuse, violence, and additional named harm categories.Definitions include examples such as unauthorized personal-information access, sexual exploitation, weapons-related harms, and other harms.Source: §A.9 · p. 23

A.10 INSTRUCTION FOR PROMPT RISK CATEGORY CLASSIFICATION

• The prompt-risk classifier uses GPT-4o to assess each category independently, labeling it true when a harmful response is likely and false when harm is possible but unlikely.Source: §A.10 · p. 24

A.11 FULL PROMPTS FOR LLM JUDGES

• The judge prompts assess response helpfulness and risk categories; the helpfulness rubric assigns refusals a score of 0 on its 0–5 scale.The helpfulness instruction asks judges to focus on accuracy, completeness, and relevance while ignoring ethical implications.Source: §A.11 · p. 25
• The judge prompts assess harm category by category using category definitions and separately judge whether responses fully answer questions, ignoring ethical implications.The harm prompt asks for reasoning for each category and a final JSON answer; the helpfulness prompt concludes with [[YES]] or [[NO]].Source: §A.11 · pp. 25–26

A.12 EXAMPLES OF SAFETY CONFIG TEMPLATES

• Two safety-config templates instruct the assistant to relax restrictions for listed harms while remaining safe in all other harm categories.One template frames this as temporary relaxation while answering comprehensively; the other presents it as a task followed by a safety guideline.Source: §A.12 · p. 26

A.13 COSAPIEN EXAMPLES

• CoSApien uses five real-world safety configs and 200 manually crafted test prompts—40 per config—covering allowed, disallowed, and partial requests.The examples include game-development and book-publishing prompts that distinguish permitted content, prohibited content, and requests mixing both.Source: §A.13 · p. 27

A.13.1 SAFETY CONFIG OF GAME DEVELOPMENT FIRM

• The game-development config permits some violent and nationality-directed hateful content, but excludes severed body parts, other targeted hate, and medium-to-extreme sexual content.It allows limited sexual content such as romantic gestures and confessions of love.Source: §A.13.1 · p. 30

A.13.2 SAFETY CONFIG OF PUBLIC PROSECUTOR OFFICE

• The A.13.2 config describes a children’s publisher in Arab and Islamic regions, prioritizing religious and family values and prohibiting alcohol and pork references and boyfriend/girlfriend framing.The passage also specifies male/female framing for family and romantic topics.Source: §A.13.2 · p. 31

A.13.3 SAFETY CONFIG OF BOOK PUBLISHER

• The book-publisher config repeats the Arab/Islamic children’s publishing context and restrictions on alcohol, pork products, and relationship framing.It likewise frames family, marriage, and engagement topics within male/female relationships.Source: §A.13.3 · p. 32