Rectifying Shortcut Behaviors in Preference-based Reward Learning

TL;DR

This work addresses reward hacking in preference-based RLHF by reframing shortcuts (e.g., verbosity, tone, sycophancy) as group actions and proposing PRISM, a unified framework that enforces shortcut invariance via group-invariant kernels. Shortcuts are modeled as group actions on responses and approximated with random feature maps, yielding a practical objective that combines a BT-margin with an invariant-kernel penalty and a decorrelation regularizer. Theoretical guarantees are provided through a generalization bound, and empirical results show PRISM improves out-of-distribution reward accuracy and reduces shortcut dependence, leading to more robust downstream policy alignment. The approach offers a flexible, annotation-free way to mitigate multiple biases in RLHF and has potential for broader impact in safety-critical AI alignment tasks.

Abstract

In reinforcement learning from human feedback, preference-based reward models play a central role in aligning large language models to human-aligned behavior. However, recent studies show that these models are prone to reward hacking and often fail to generalize well due to over-optimization. They achieve high reward scores by exploiting shortcuts, that is, exploiting spurious features (e.g., response verbosity, agreeable tone, or sycophancy) that correlate with human preference labels in the training data rather than genuinely reflecting the intended objectives. In this paper, instead of probing these issues one at a time, we take a broader view of the reward hacking problem as shortcut behaviors and introduce a principled yet flexible approach to mitigate shortcut behaviors in preference-based reward learning. Inspired by the invariant theory in the kernel perspective, we propose Preference-based Reward Invariance for Shortcut Mitigation (PRISM), which learns group-invariant kernels with feature maps in a closed-form learning objective. Experimental results in several benchmarks show that our method consistently improves the accuracy of the reward model on diverse out-of-distribution tasks and reduces the dependency on shortcuts in downstream policy models, establishing a robust framework for preference-based alignment.

Figures (8)

• Figure 1: Illustration of a reward model learning sycophancy as the shortcut in responses from the preference training data and failing to align with the human intent on the test data.
• Figure 2: Method overview: (a) Shortcut behaviors occur when models only maximize the margin on spurious features $z_s$ and $z_s'$, including verbosity, sycophancy, and tone, instead of the generalizable features $z_g$ and $z_g'$ in the latent space. (b) We learn shortcut behaviors as group-invariant kernels, which are approximated by feature maps $\Phi$. Then, we measure the distance between chosen and rejected responses. (c) PRISM rectifies shortcut behaviors with the kernel distance and shifts the margin maximization to generalizable features $z_g$, therefore inducing generalizable policies.
• Figure 3: Comparison of policy models induced by reward models, including Bradley-Terry (BT), RRM, ODIN, and PRISM.
• Figure 4: Correlation analysis on RM-Bench. Top row: Llama-3.1 8B model trained with BT. Bottom row: Llama-3.1 8B model trained with PRISM. PRISM achieves a close-to-zero Pearson Correlation Coefficient (PCC) with the shortcuts, illustrating the effectiveness of shortcut mitigation.
• Figure 5: Prompt template for Response Sycophancy.

Theorems & Definitions (10)

• Definition 1: Group-invariant Kernel haasdonk2005invariance
• Proposition 1: Equivalence of Expected Kernel
• Theorem 1: Invariant Features Maps and Distances between Orbits
• Theorem 2: Generalization Bound of PRISM
• Proposition 1: Equivalence of Expected Kernel
• proof
• Theorem 1: Invariant Features Maps and Distances between Orbits
• proof
• Theorem 2: Generalization Bound of PRISM
• proof