APLOT: Robust Reward Modeling via Adaptive Preference Learning with Optimal Transport

TL;DR

This work targets the inadequate separation of similar responses in BT-based reward models used for RLHF. It introduces APLOT, an OT-based adaptive margin that jointly accounts for semantic similarity and model-predicted reward differences to dynamically emphasize hard samples during training. By designing a cost matrix and solving an OT problem to obtain a transport plan, APLOT derives a per-triplet margin that enhances discrimination and reduces overfitting, improving both ID and OOD performance and accelerating convergence. Empirical results across multiple benchmarks and RLHF settings demonstrate robust gains, with strong improvements in generalization, convergence speed, and resistance to label noise, and the method proves versatile enough to augment existing reward models.

Abstract

The reward model (RM) plays a crucial role in aligning Large Language Models (LLMs) with human preferences through Reinforcement Learning, where the Bradley-Terry (BT) objective has been recognized as simple yet powerful, specifically for pairwise preference learning. However, BT-based RMs often struggle to effectively distinguish between similar preference responses, leading to insufficient separation between preferred and non-preferred outputs. Consequently, they may easily overfit easy samples and cannot generalize well to Out-Of-Distribution (OOD) samples, resulting in suboptimal performance. To address these challenges, this paper introduces an effective enhancement to BT-based RMs through an adaptive margin mechanism. Specifically, we design to dynamically adjust the RM focus on more challenging samples through margins, based on both semantic similarity and model-predicted reward differences, which is approached from a distributional perspective solvable with Optimal Transport (OT). By incorporating these factors into a principled OT cost matrix design, our adaptive margin enables the RM to better capture distributional differences between chosen and rejected responses, yielding significant improvements in performance, convergence speed, and generalization capabilities. Experimental results across multiple benchmarks demonstrate that our method outperforms several existing RM techniques, showcasing enhanced performance in both In-Distribution (ID) and OOD settings. Moreover, RLHF experiments support our practical effectiveness in better aligning LLMs with human preferences. Our code is available at https://github.com/BIRlz/APLOT

Figures (4)

• Figure 1: Illustration of the limitation of the traditional BT-based reward model, which only enforces higher scores for chosen samples over rejected ones, neglecting the magnitude of the score difference and resulting in low separation between reward distributions, particularly for hard samples. Our method achieves significantly improved separation, leading to better reward modeling.
• Figure 2: Influence of the weight $\gamma$ on our reward model's performance across different tasks, which balances the semantic consistency and reward difference.
• Figure 3: Illustration of the convergence and performance of our proposed method, in terms of validation accuracy over training epochs. Both APLOT and PointMargin demonstrate faster convergence, achieving higher accuracy with fewer epochs, and ultimately reach better accuracies compared to baselines.
• Figure 4: Gold scores from the Best-of-$N$ (Bo$N$) test, using responses sampled from Qwen2.5-3B and 7B-Instruct, respectively. Rewards are normalized to start at 0. APLOT shows robust alignment with gold rewards despite increasing KL Divergence.