Optimal Transport for LLM Reward Modeling from Noisy Preference
Proposes SelectivelyRM, an optimal transport-based framework for robust reward modeling from noisy preferences, outperforming SOTA by 5-10% in error metrics.
Key Findings
Methodology
The approach formulates reward modeling as a distribution alignment problem using joint consistency discrepancy, combining semantic distance and preference differences. It employs partial optimal transport to dynamically exclude high-cost noisy samples, effectively filtering out preference errors. Theoretically, it minimizes a tighter upper bound on the clean risk, providing generalization guarantees. During training, the model iteratively computes cost matrices based on fixed semantic embeddings, solves partial transport plans, and updates rewards to focus on consistent supervision, thus enhancing robustness against instance-dependent noise.
Key Results
- On HelpSteer, UltraFeedback, and PKU-SafeRLHF datasets, SelectivelyRM reduces MSE by 5-10%, e.g., from 0.070 to 0.063 on HelpSteer, outperforming all baselines. It maintains stable performance across noise ratios from 10% to 50%, with noise exclusion rates above 80%. Ablation shows Mass Relaxation and Semantic Alignment contribute 40% and 35% respectively to improvements.
- The method demonstrates robustness to varying model sizes (7B-72B parameters), with consistent error reductions. It effectively filters noisy samples, leading to more accurate preference recovery. Comparative analysis confirms superior performance over statistical correction and heuristic filtering methods.
- Ablation studies reveal that removing partial transport or semantic alignment significantly degrades performance, confirming their critical roles. The approach also scales well with large datasets, maintaining efficiency and stability.
Significance
This work addresses a fundamental challenge in reward modeling—noise in preference data—by introducing a theoretically grounded, noise-aware distribution alignment method. It enhances the robustness and reliability of RLHF, facilitating safer and more aligned AI systems. The framework’s theoretical guarantees and empirical success set a new standard for noise-resilient preference learning, with broad implications for multi-modal and multi-task settings. It paves the way for future research into adaptive, scalable noise filtering techniques that can be integrated into large-scale language model training pipelines, ultimately contributing to AI systems that better understand and respect human values.
Technical Contribution
The key technical innovation is the integration of partial optimal transport into reward modeling, allowing the model to autonomously exclude high-cost, noisy preference samples. The joint consistency discrepancy combines semantic and preference information, providing a more sensitive alignment metric. Theoretical analysis proves that minimizing this discrepancy tightens the upper bound on the true risk, offering rigorous generalization guarantees. Practically, the method employs an efficient iterative algorithm that adapts to different model architectures and noise levels, significantly improving robustness without sacrificing scalability.
Novelty
This is the first work to incorporate partial optimal transport into reward modeling for preference learning, specifically targeting instance-dependent noise. Unlike previous methods relying on full transport or static filtering, it dynamically filters noisy samples during training, guided by a theoretically justified discrepancy measure. The combination of semantic and preference alignment in a partial transport framework represents a novel approach, bridging distribution matching theory with practical noise filtering in large-scale language models.
Limitations
- The method relies on fixed semantic embeddings, which may be less effective if the embeddings do not capture nuanced semantics or contain biases. In highly noisy scenarios (>50%), filtering effectiveness diminishes, requiring further refinement.
- Computational overhead increases with dataset size and model complexity, especially during iterative transport plan computation. Optimization for large-scale deployment remains an open challenge.
- The approach assumes the availability of pre-trained semantic encoders and may be sensitive to hyperparameter choices like the mass quota κ. Further research is needed to automate hyperparameter tuning and extend robustness.
Future Work
Future directions include integrating adaptive hyperparameter tuning, exploring multi-modal preference data, and extending the framework to multi-task settings. Developing more efficient algorithms for large-scale partial transport, and investigating causal mechanisms behind preference noise, could further improve robustness. Additionally, applying this approach to real-world applications such as dialogue safety, content moderation, and multi-modal alignment will be valuable for deploying safer, more aligned AI systems.
AI Executive Summary
In the rapidly advancing field of large language models (LLMs), aligning AI behavior with human preferences remains a central challenge. Reinforcement Learning from Human Feedback (RLHF) has emerged as a promising approach, but the quality of preference data is often compromised by noise—either from human annotators or automated systems. Traditional denoising techniques, such as statistical correction or heuristic filtering, struggle to handle instance-dependent noise, which varies with input complexity and ambiguity.
This paper introduces SelectivelyRM, a novel framework grounded in optimal transport theory, to robustly model preferences despite noisy annotations. The core idea involves formulating reward modeling as a distribution alignment problem, where the goal is to match the model's predicted distribution with the observed noisy preferences. To achieve this, the authors define a joint consistency discrepancy that combines semantic similarity and preference divergence, serving as a sensitive measure of distributional mismatch. Recognizing that strict mass conservation in traditional optimal transport can force the model to fit outliers, they incorporate a partial transport mechanism, allowing the model to selectively exclude high-cost, noisy samples during training.
Theoretically, the authors demonstrate that this partial transport-based discrepancy tightens the upper bound on the true, noise-free risk, providing rigorous generalization guarantees. Empirically, extensive experiments on multiple benchmark datasets show that SelectivelyRM outperforms existing methods by reducing error metrics by 5-10%, with noise exclusion rates exceeding 80%. Ablation studies confirm the importance of the mass relaxation and semantic alignment components. The approach scales well across different model sizes and noise levels, offering a practical, scalable solution for real-world noisy preference data.
This work significantly advances the robustness of reward models in RLHF, addressing a critical bottleneck in deploying safe and aligned AI systems. Its theoretical insights and empirical successes open new avenues for scalable, noise-resilient preference learning, with broad implications for AI safety, multi-modal alignment, and human-AI interaction.
Deep Analysis
Background
The evolution of large language models (LLMs) has revolutionized NLP, with techniques like RLHF enabling models to better align with human values. Early works such as InstructGPT demonstrated the effectiveness of preference-based fine-tuning. However, as datasets scale, annotation noise—stemming from human inconsistency, fatigue, or automated biases—becomes inevitable. Existing solutions like loss correction methods (e.g., F-correction) or sample filtering (e.g., Co-Teaching) partially address noise but struggle with instance-dependent noise, which varies with input ambiguity. The challenge remains to develop methods that can adaptively filter noisy preferences, especially in high-noise scenarios, without sacrificing data efficiency or theoretical guarantees.
Core Problem
The core issue is that noisy preference data, especially with instance-dependent errors, causes reward models to memorize erroneous signals, leading to poor generalization and potential reward hacking. Traditional training approaches treat all data equally, which results in overfitting to outliers. This is particularly problematic in RLHF, where the reward model directly influences policy optimization. The difficulty lies in designing a filtering mechanism that can dynamically identify and exclude noisy samples during training, ensuring the model learns the true underlying preferences. Achieving this while maintaining theoretical guarantees and computational efficiency remains an open problem.
Innovation
The main innovations include: 1) introducing a joint consistency discrepancy that combines semantic and preference information for sensitive distribution alignment; 2) employing partial optimal transport to relax the mass conservation constraint, enabling the model to autonomously filter out high-cost, noisy samples; 3) providing rigorous theoretical bounds showing that minimizing this discrepancy tightens the upper bound on the clean risk. Unlike prior methods relying on static thresholds or full transport, this approach dynamically adapts to data quality, improving robustness against instance-dependent noise. The framework is also compatible with various model architectures and scalable to large datasets.
Methodology
- �� Encode prompt-response pairs into fixed semantic vectors using a pre-trained encoder, ensuring semantic consistency.
- �� Construct a cost matrix combining Euclidean semantic distance and preference divergence.
- �� Define the joint consistency discrepancy as the optimal transport cost with mass conservation constraints.
- �� Introduce partial transport with a mass quota κ, allowing the model to selectively exclude high-cost, noisy samples.
- �� During training, compute the cost matrix, solve for the optimal partial coupling, and update the reward model by minimizing the weighted loss based on the coupling.
- �� Iteratively refine the model, dynamically filtering out noisy preferences, and improving alignment with true preferences.
Experiments
Experiments used datasets like HelpSteer, UltraFeedback, and PKU-SafeRLHF, injecting synthetic noise at varying levels (10%-50%). Baselines included statistical correction methods and heuristic filtering. Metrics such as MSE, MAE, and R2 evaluated alignment accuracy. Hyperparameters like learning rate and batch size were tuned via grid search. The training employed a fixed semantic encoder and iterative partial transport optimization. Ablation studies isolated the contributions of mass relaxation and semantic alignment, confirming their importance. Results demonstrated consistent error reduction and robustness across model sizes and noise levels.
Results
SelectivelyRM achieved 5-10% error reduction over baselines, with MSE dropping from 0.070 to 0.063 on HelpSteer. Noise exclusion rates exceeded 80% at 50% noise ratio, effectively filtering out noisy samples. Ablation showed that removing partial transport or semantic alignment decreased performance by 30-40%. The method maintained stable performance across models from 7B to 72B parameters, confirming scalability. These results validate the approach’s robustness and practical utility in noisy real-world scenarios.
Applications
Applicable in any preference-based RLHF system, such as chatbot safety, content moderation, and recommendation systems. By improving noise robustness, it enhances model safety and user trust. The framework can be extended to multi-modal preferences, aiding in tasks like visual or multi-sensory alignment. Its scalability makes it suitable for industrial deployment, where noisy feedback is common. Future integration with active learning could further improve data efficiency and model reliability.
Limitations & Outlook
Dependence on pre-trained semantic encoders may limit performance if embeddings are biased or insufficiently expressive. Effectiveness diminishes with extremely high noise levels (>50%). Computational complexity increases with dataset size, requiring further optimization. Hyperparameter sensitivity, especially the mass quota κ, necessitates careful tuning. Future work should focus on adaptive hyperparameters, more efficient algorithms, and broader application scenarios.
Plain Language Accessible to non-experts
Imagine you’re organizing a big school event, and you ask students to rate different activities. Some students give wrong or inconsistent ratings because they’re tired or confused—like giving high scores to boring games or low scores to fun ones. If you treat all ratings equally, you might end up planning an event based on these mistakes, making it less enjoyable.
Now, think of a smart organizer who can look at all these ratings and figure out which ones are trustworthy. This organizer uses a special method to compare the ratings with the actual quality of activities, focusing on the most reliable feedback. It also ignores or filters out the obviously wrong ratings, so the final plan is based on the best information. This way, the event turns out great because it’s built on accurate feedback, not the noisy or misleading ones. The paper’s method is like that organizer, using a clever mathematical way to identify and exclude bad ratings, ensuring the final decision is as close to the truth as possible.
ELI14 Explained like you're 14
Imagine you’re playing a game with your friends, and everyone gives you scores about how fun each level is. Sometimes, friends might give weird scores—like giving a boring level a 10 out of 10 just to be funny, or giving a super fun level a 1 because they’re in a bad mood. If you listen to all their scores equally, you might end up thinking the boring level is the best, which isn’t true.
So, you decide to be smarter. You look at all the scores and try to find out which ones are real and which ones are just jokes or mistakes. You focus on the scores that match what you see and feel about the game, and ignore the weird ones. This way, you can figure out which levels are actually fun and plan your next game better.
That’s what the paper’s method does for AI models. It helps the AI learn what people really like, even if some feedback is noisy or wrong. It uses a clever trick to filter out the bad feedback, so the AI can understand human preferences more accurately. Just like you, the AI becomes better at understanding what’s truly good or bad, making it safer and more helpful.
Abstract
Reward models are fundamental to Reinforcement Learning from Human Feedback (RLHF), yet real-world datasets are inevitably corrupted by noisy preference. Conventional training objectives tend to overfit these errors, while existing denoising approaches often rely on homogeneous noise assumptions that fail to capture the complexity of linguistic preferences. To handle these challenges, we propose SelectiveRM, a framework grounded in optimal transport. We first devise a Joint Consistency Discrepancy to align the distribution of model predictions with preference data. Furthermore, to address the limitation of strict mass conservation which compels the model to fit outliers, we incorporate a Mass Relaxation mechanism via partial transport. This enables the autonomous exclusion of samples with noisy preference that contradict semantic consistency. Theoretically, we demonstrate that SelectiveRM optimizes a tighter upper bound on the unobserved clean risk. Extensive experiments validate that our approach significantly outperforms state-of-the-art baselines across diverse benchmarks.