Improving Instruction Following in Language Models through Proxy-Based Uncertainty Estimation

2024-05-10

International Conference on Machine Learning

AI Summary

To address ambiguous response quality assessment and inconsistent scoring in instruction-following evaluation, this paper proposes the Uncertainty-Aware Reward Model (URM), the first approach to incorporate Bayesian approximate uncertainty modeling into instruction-following reward learning. URM jointly optimizes response scoring and uncertainty quantification on preference data. It enables dual-dimensional evaluation—assessing both response quality and confidence—thereby providing a unified framework for high-quality data filtering and robust policy optimization. Empirical results on benchmarks including Vicuna and MT-Bench demonstrate that reinforcement learning guided by URM significantly improves instruction-following capability over state-of-the-art methods. Moreover, URM enhances the model’s ability to discriminate training data quality and improves policy convergence stability.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

Assessing response quality to instructions in language models is vital but challenging due to the complexity of human language across different contexts. This complexity often results in ambiguous or inconsistent interpretations, making accurate assessment difficult. To address this issue, we propose a novel Uncertainty-aware Reward Model (URM) that introduces a robust uncertainty estimation for the quality of paired responses based on Bayesian approximation. Trained with preference datasets, our uncertainty-enabled proxy not only scores rewards for responses but also evaluates their inherent uncertainty. Empirical results demonstrate significant benefits of incorporating the proposed proxy into language model training. Our method boosts the instruction following capability of language models by refining data curation for training and improving policy optimization objectives, thereby surpassing existing methods by a large margin on benchmarks such as Vicuna and MT-bench. These findings highlight that our proposed approach substantially advances language model training and paves a new way of harnessing uncertainty within language models.

Problem

• Language Model Accuracy
• Complexity and Ambiguity in Language
• Response Quality Evaluation

Innovation

• Agent-based Uncertainty Estimation
• Uncertainty-aware Reward Model (URM)
• Enhanced Instruction Understanding and Execution