Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Data-adaptive Safety Rules for Training Reward Models

Xiaomin Li, Mingye Gao, Zhiwei Zhang, Jingxuan Fan, Weiyu Li

TL;DR

This work introduces a data-adaptive framework for labeling fine-grained safety preferences in RLHF by selecting a small set of the most informative rules per data instance. A max-discrepancy criterion, aided by a relevance term, drives the Rule Adapter to pick top-$5$ rules, maximizing the mutual information with latent ground-truth preferences via Jensen–Shannon divergence. The approach yields RAMO, an 8B reward model that achieves state-of-the-art safety on RewardBench and improves safety when integrated into PPO-aligned LLMs, while preserving general capabilities. By translating rule-discrepancy signals into binary preferences, the method reduces labeling cost and bias compared to fixed-rule or GPT-labeling baselines. The work also demonstrates generalization to human preference data and releases useful resources to advance safety-alignment research.

Abstract

Reinforcement Learning from Human Feedback (RLHF) is commonly employed to tailor models to human preferences, especially to improve the safety of outputs from large language models (LLMs). Traditionally, this method depends on selecting preferred responses from pairs. However, due to the variability in human opinions and the challenges in directly comparing two responses, there is an increasing trend towards fine-grained annotation approaches that evaluate responses using multiple targeted metrics or rules. The challenge lies in efficiently choosing and applying these rules to handle the diverse range of preference data. In this paper, we propose a dynamic method that adaptively selects the most important rules for each response pair. We introduce a mathematical framework that utilizes the maximum discrepancy across paired responses and demonstrate theoretically that this approach maximizes the mutual information between the rule-based annotations and the underlying true preferences. We then train an 8B reward model using this adaptively labeled preference dataset and assess its efficacy using RewardBench. As of January 25, 2025, our model achieved the highest safety performance on the leaderboard, surpassing various larger models.

Data-adaptive Safety Rules for Training Reward Models

TL;DR

This work introduces a data-adaptive framework for labeling fine-grained safety preferences in RLHF by selecting a small set of the most informative rules per data instance. A max-discrepancy criterion, aided by a relevance term, drives the Rule Adapter to pick top- rules, maximizing the mutual information with latent ground-truth preferences via Jensen–Shannon divergence. The approach yields RAMO, an 8B reward model that achieves state-of-the-art safety on RewardBench and improves safety when integrated into PPO-aligned LLMs, while preserving general capabilities. By translating rule-discrepancy signals into binary preferences, the method reduces labeling cost and bias compared to fixed-rule or GPT-labeling baselines. The work also demonstrates generalization to human preference data and releases useful resources to advance safety-alignment research.

Abstract

Reinforcement Learning from Human Feedback (RLHF) is commonly employed to tailor models to human preferences, especially to improve the safety of outputs from large language models (LLMs). Traditionally, this method depends on selecting preferred responses from pairs. However, due to the variability in human opinions and the challenges in directly comparing two responses, there is an increasing trend towards fine-grained annotation approaches that evaluate responses using multiple targeted metrics or rules. The challenge lies in efficiently choosing and applying these rules to handle the diverse range of preference data. In this paper, we propose a dynamic method that adaptively selects the most important rules for each response pair. We introduce a mathematical framework that utilizes the maximum discrepancy across paired responses and demonstrate theoretically that this approach maximizes the mutual information between the rule-based annotations and the underlying true preferences. We then train an 8B reward model using this adaptively labeled preference dataset and assess its efficacy using RewardBench. As of January 25, 2025, our model achieved the highest safety performance on the leaderboard, surpassing various larger models.
Paper Structure (34 sections, 3 theorems, 36 equations, 1 figure, 9 tables)

This paper contains 34 sections, 3 theorems, 36 equations, 1 figure, 9 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Definitions and Notations
  5. Preliminaries
  6. Maximum Discrepancy Selection
  7. Theoretical Analysis
  8. Experiments
  9. Rule Pool
  10. Rule Adapter
  11. Reward Model
  12. Evaluation
  13. Results
  14. Ablation study
  15. Generalization: Relabel Human Preference Data
  16. ...and 19 more sections

Key Result

Theorem 3.4

Given a trio $(x, y_A, y_B) \in \mathcal{D}$, use ${\bm{v}}_A, {\bm{v}}_B$ to denote $(x, y_A)$ and $(x, y_B)$, respectively. Let $H\in\{\pm1\}$ be the hidden ground-truth preference label such that Without loss of generality, assume that the data are balanced so that $H$ is uniformly distributed (i.e. $H \sim \text{Bern}(1/2)$). For each rater $\psi_i$ where $i \in [R]$, let $T_i \in \{\pm1\}$ b

Figures (1)

  • Figure 1: Pipeline of our framework. First, we train a Rule Adapter that learns to identify $r=5$ most critical rules for a given trio. These rules are selected based on their ability to maximize the discrepancy between the two responses and their relevance to the prompt. Both responses are then rated according to the $r$ selected rules, and preferences are labeled based on the aggregated ratings. Then we proceed to train a reward model, which is subsequently integrated into the standard RLHF process.

Theorems & Definitions (13)

  • Definition 3.1: Rule-based Raters
  • Definition 3.2: Preference Labeling
  • Definition 3.3: Mutual Information
  • Theorem 3.4
  • proof
  • Definition A.1: Mutual Information (equivalent definition)
  • Definition A.2: Kullback–Leibler Divergence
  • Definition A.3: Jensen-Shannon Divergence
  • Definition A.4
  • Lemma A.5
  • ...and 3 more