Low-Rank Contextual Reinforcement Learning from Heterogeneous Human Feedback

TL;DR

This paper introduces LoCo-RLHF, a novel framework for reinforcement learning from heterogeneous human feedback that leverages low-rank structures and pessimistic policies to improve personalization and robustness.

Contribution

The paper proposes a low-rank contextual model for RLHF and a new pessimistic policy to handle distributional shifts, advancing personalized alignment methods.

Findings

Outperforms existing methods in personalized RLHF tasks

Demonstrates robustness to distributional shifts in feedback

Achieves tighter sub-optimality bounds theoretically

Abstract

Reinforcement learning from human feedback (RLHF) has become a cornerstone for aligning large language models with human preferences. However, the heterogeneity of human feedback, driven by diverse individual contexts and preferences, poses significant challenges for reward learning. To address this, we propose a Low-rank Contextual RLHF (LoCo-RLHF) framework that integrates contextual information to better model heterogeneous feedback while maintaining computational efficiency. Our approach builds on a contextual preference model, leveraging the intrinsic low-rank structure of the interaction between user contexts and query-answer pairs to mitigate the high dimensionality of feature representations. Furthermore, we address the challenge of distributional shifts in feedback through our Pessimism in Reduced Subspace (PRS) policy, inspired by pessimistic offline reinforcement learning techniques. We theoretically demonstrate that our policy achieves a tighter sub-optimality gap compared to existing methods. Extensive experiments validate the effectiveness of LoCo-RLHF, showcasing its superior performance in personalized RLHF settings and its robustness to distribution shifts.