💬 NLP

Configurable Reward Model for Balanced Safety Alignment

The paper introduces the Configurable Safety Reward Model (CSRM), a novel approach that leverages configuration-targeted data augmentation to create a reward model capable of adapting to heterogeneous and evolving safety requirements, thereby achieving state-of-the-art performance on configurable safety benchmarks and significantly improving the helpfulness-safety tradeoff in aligned large language models.

Original authors: Zhengping Jiang, Mehran Khodabandeh, Akash Bharadwaj, Manik Bhandari, Mayur Srungarapu, Anqi Liu, Benjamin Van Durme, Li Chen

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhengping Jiang, Mehran Khodabandeh, Akash Bharadwaj, Manik Bhandari, Mayur Srungarapu, Anqi Liu, Benjamin Van Durme, Li Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: One Size Does Not Fit All

Imagine you have a very smart robot assistant (a Large Language Model or LLM). You want this robot to be helpful, but you also want it to be safe.

The problem is that "safety" looks different depending on where you are and what you are doing:

In a creative writing class: A story about a villain stealing a bank might be exciting and safe.
In a bank: That same story is a violation of security protocols.
In a hospital: A story about a patient might need to be strictly confidential.

Currently, most AI safety systems are like frozen statues. Once they are built, they can't change. If a company needs to update its safety rules (like adding a new rule about "no talking about money"), the engineers have to stop everything, gather new human teachers, retrain the robot from scratch, and hope it learns the new rule. This is slow, expensive, and often leads to the robot becoming too scared to answer any questions (a problem called "over-refusal").

The Solution: The "Configurable Safety Reward Model" (CSRM)

The authors introduce a new system called CSRM. Think of CSRM not as a statue, but as a smart, adjustable thermostat.

Instead of having one fixed setting for "Safety," CSRM allows you to plug in a specific "Safety Manual" (written in plain English) for every conversation.

The Input: You give the robot the conversation plus a specific set of rules (e.g., "For this movie script, violence is okay, but hate speech is not").
The Output: The CSRM doesn't just say "Yes/No." It gives a score (like a grade from 0 to 100) telling the robot exactly how safe or unsafe the response is based on those specific rules.

How It Works: The "Chef's Kitchen" Analogy

To understand how they trained this model, imagine a Chef (the AI) who needs to learn how to cook for different dietary restrictions.

The Old Way (Static Training): You teach the Chef, "Never cook with pork." The Chef learns this rule forever. If you later ask for a "Pork-free but spicy" dish, the Chef might refuse to cook anything at all because they are confused or too scared to make a mistake.
The CSRM Way (Configurable Training):
- The "Menu" Augmentation: The researchers created a special training method. They took real conversations and asked a smart AI to invent new "Menu Rules."
  - Scenario A: "Add a rule that says 'No pork'." (The Chef learns to avoid pork).
  - Scenario B: "Add a rule that says 'Pork is okay, but no alcohol'." (The Chef learns to cook pork without wine).
- The "Severity" Augmentation: They also taught the Chef the difference between a small mistake and a big one.
  - Scenario: "Accidentally saying 'pork' once is a minor slip (low penalty)."
  - Scenario: "Serving a whole pig to a strict vegetarian is a major disaster (high penalty)."
- The Result: The Chef learns to read the specific menu for the day and adjust the cooking instantly, without needing to go back to culinary school.

Why This is Better Than Other Systems

The paper compares CSRM to two other types of safety systems:

The "Bouncer" (Standard Classifiers): These stand at the door and just say "In" or "Out." They are fast, but they can't tell the difference between a "slightly risky" joke and a "dangerous" threat. They are too blunt.
The "Judge" (Reasoning Models): These are like lawyers who write long essays explaining why something is bad. They are accurate but very slow and hard to use for training the AI.

CSRM is the "Scorekeeper": It gives a precise number (a reward score) that tells the AI exactly how well it did. This allows the AI to learn gradually, improving its behavior step-by-step, rather than just being told "Wrong!" or "Right!"

The Results: A Better Balance

The researchers tested CSRM on various safety benchmarks and found:

It adapts instantly: It can handle new safety rules it has never seen before without retraining.
It stops "Over-Refusal": Because it understands the nuance of the rules, it doesn't just say "No" to everything. It finds the middle ground where the AI is helpful and safe.
It creates a "Pareto Frontier": This is a fancy way of saying it achieves the best possible balance. You get more helpfulness without losing safety, or more safety without losing helpfulness. It pushes the limits of what is possible.

Summary

The paper presents a new tool that lets AI safety be flexible. Instead of hard-coding safety rules that break when the world changes, CSRM acts like a dynamic translator that reads the specific safety rules for a situation and guides the AI to behave perfectly within those boundaries. It makes AI safer, smarter, and less likely to be annoyingly cautious.

Technical Summary: Configurable Reward Model for Balanced Safety Alignment

Problem Statement

The alignment of Large Language Models (LLMs) to safety requirements faces a critical challenge: safety is not a universal constant but a context-dependent variable shaped by cultural norms, legal jurisdictions, and organizational policies. Current alignment paradigms, particularly Reinforcement Learning from Human Feedback (RLHF), rely on static Reward Models (RMs) where safety knowledge is implicitly encoded and frozen at deployment. This rigidity necessitates a costly full retrain-and-deploy cycle whenever safety policies evolve (e.g., new hate speech regulations or domain-specific compliance).

Conversely, existing "configurable" approaches, such as standalone safety classifiers (e.g., Llama Guard, ShieldGemma) or prompt-conditioned reasoning judges (e.g., DynaGuard), adapt to new policies at inference time. However, these models function as discriminative classifiers or reasoning-based judges. Their outputs often exhibit reward geometries poorly suited for reinforcement learning: probabilities are either sharply peaked (binary decisions) or excessively flat (deliberative judgments). Consequently, they provide sparse, poorly calibrated, and effectively non-differentiable signals for policy optimization, frequently leading to over-refusal (rejecting benign requests) and degrading utility.

There is a missing component in current pipelines: a reward model that is simultaneously configurable at inference time and capable of producing a dense, calibrated scalar reward suitable for inner-loop policy optimization.

Methodology

The authors introduce the Configurable Safety Reward Model (CSRM), a framework designed to condition on natural-language safety configurations at inference time while outputting a scalar reward for RLHF. The methodology relies on three core components:

1. Terminology and Input Structure

A safety configuration $p$ is defined as a set of rules with natural-language guidelines delineating acceptable and unacceptable content for specific categories. The model takes a tuple $(x, r, p)$ as input, where $x$ is the conversation history, $r$ is the agent's response, and $p$ is the safety configuration. It outputs a label $y \in \{\text{safe, unsafe}\}$ and a reward value $c \in [0, 1]$ .

2. Targeted Data Augmentation

To enable generalization to unseen configurations without additional human annotation, the authors propose two augmentation strategies:

Configurable Safety Configuration Augmentation: For a given conversation and configuration, a reasoning model proposes a "positive" category (not in $p$ ) that would render the conversation unsafe if added, and a "negative" category that would render it safe. The configuration is augmented by randomly dropping existing categories and optionally inserting these new ones. This two-sided augmentation increases reward spread and improves pairwise modeling.
Strictness Augmentation: To shape the reward geometry and reflect violation severity, the authors construct partially ordered sets of guideline descriptions. They identify subcategories within a safety category, sort them by estimated severity, and rewrite guidelines to selectively allow or disallow prefixes of this ordered list. To ensure reliability, they employ a confidence-qualified dominance test using the Clopper–Pearson lower bound. Only guideline pairs $(d, d')$ where $d$ is statistically significantly stricter than $d'$ are retained for training. This ensures the model learns a calibrated severity spectrum rather than arbitrary binary distinctions.

3. Joint Discriminative–Generative Training

CSRM unifies safety classification and reward modeling within a single generative framework.

Classification Loss: The model is trained to predict safety labels using verbalized tokens (e.g., "safe", "unsafe").
Reward Loss: Using the strictness-augmented pairs, the model is trained to maximize the margin between the log-probabilities of safe vs. unsafe tokens under strict versus lenient configurations. This is formulated as a Bradley–Terry style preference loss, encouraging the model to assign higher rewards to responses that are safer under stricter guidelines.
Calibration: Rather than using a separate calibration loss, calibration is induced during training by constructing diverse, statistically validated severity pairs. This leverages the empirical link between ranking quality and calibration.

Key Contributions

Configurable, Calibrated Safety Reward Modeling: The CSRM conditions directly on natural-language safety configurations at inference time, producing a dense, calibrated scalar reward rather than a binary judgment. This enables fine-grained control over safety behavior without retraining and remains compatible with gradient-based policy optimization.
Joint Training with Targeted Augmentation: The work unifies safety classification and reward modeling, introducing configuration-targeted and strictness augmentations. This allows the model to distinguish between borderline and severe violations and generalize to unseen configurations without new human annotations.
Improved Safety–Helpfulness Trade-offs: CSRM provides a more informative reward signal for RL, yielding policies that avoid over-refusal while maintaining strong safety guarantees, effectively expanding the Pareto frontier between safety and utility.

Experimental Results

The authors evaluated CSRM on multiple benchmarks and downstream alignment tasks:

Safety Classification: On benchmarks including BeaverTails, WildGuardMix, CoSApien, and DynaBench, CSRM achieved state-of-the-art performance. Notably, it reached 94.6% F1 on CoSApien and 75.8% F1 on DynaBench, outperforming baselines like LlamaGuard-3-8B and ShieldGemma-9B. It also demonstrated superior calibration (lower smECE) compared to reasoning-based judges.
Reward Modeling: In pairwise reward modeling tasks (BeaverTails-Aug, WildGuardMix-Aug), CSRM achieved the highest accuracy in ordering responses by violation severity, significantly outperforming baselines.
Over-Refusal: On over-refusal benchmarks (XSTest, OR-Bench), CSRM attained the highest F1, indicating better calibration and reduced tendency to reject benign requests compared to instruction-tuned baselines.
Downstream RL Alignment: When used as a reward signal for aligning Mistral-7B-Instruct-v0.1 across five domains (e.g., Arab Publisher, Film Production) using Reward Distillation and REINFORCE++, CSRM consistently achieved the best CoSA Score (the dot product of safety and helpfulness rewards). In several domains, it achieved higher helpfulness at comparable safety levels, demonstrating a superior trade-off. For instance, in the Film Production domain under REINFORCE++, CSRM achieved a CoSA score of 3.981, outperforming the LlamaGuard baseline (3.391) and the instruction-tuned baseline (3.113).

Significance and Claims

The paper claims that CSRM addresses a fundamental limitation in current safety alignment: the inability to adapt to evolving safety requirements without costly retraining while maintaining a signal suitable for RL. By functioning as a dense, configurable reward signal within the inner loop of RLHF, CSRM enables the training of inherently safer models rather than merely policing their outputs at inference time.

The authors emphasize that CSRM achieves these results without requiring additional human annotation, relying instead on synthetic augmentation of existing datasets. The work suggests a practical path toward reward-level configurability, allowing safety behavior to be adapted by updating specifications at inference time while retaining an optimization-compatible signal for training.

Limitations noted by the authors:

The augmentation pipeline relies on LLM-generated proposals, which may inherit biases from the proposer model, though this is mitigated by statistical filtering.
The study did not perform a targeted reward-hacking analysis; strategic exploitation under long-horizon online optimization remains future work.
The main results utilize a single base architecture (Llama-3.1-8B), though trends were replicated on Qwen3-2B. Broader backbone and multilingual evaluations are left for future work.