Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

The Power of Active Multi-Task Learning in Reinforcement Learning from Human Feedback

Ruitao Chen, Liwei Wang

TL;DR

This work tackles the data-inefficiency problem in reinforcement learning from human feedback (RLHF) by casting RLHF as a contextual dueling bandit problem with a shared, low-dimensional linear representation. It introduces active multi-task learning, which allocates samples across $M$ source tasks according to a task relevance vector $\nu$, and develops algorithms to estimate $\nu$ from limited data when unknown. Theoretical guarantees show that the source-task sample complexity scales with $C^*$ and $\| u\|_1^2$, while the target-task sample complexity is linear in the latent dimension $k$, thanks to representation learning. The approach supports both known and unknown task relevance settings and provides a practical two-stage procedure (pre-sampling and Lasso-based relevance estimation) to achieve $\varepsilon$-optimal policies with provable guarantees, offering a principled path to leveraging diverse human feedback efficiently in RLHF.

Abstract

Reinforcement learning from human feedback (RLHF) has contributed to performance improvements in large language models. To tackle its reliance on substantial amounts of human-labeled data, a successful approach is multi-task representation learning, which involves learning a high-quality, low-dimensional representation from a wide range of source tasks. In this paper, we formulate RLHF as the contextual dueling bandit problem and assume a common linear representation. We demonstrate that the sample complexity of source tasks in multi-task RLHF can be reduced by considering task relevance and allocating different sample sizes to source tasks with varying task relevance. We further propose an algorithm to estimate task relevance by a small number of additional data and then learn a policy. We prove that to achieve $\varepsilon-$optimal, the sample complexity of the source tasks can be significantly reduced compared to uniform sampling. Additionally, the sample complexity of the target task is only linear in the dimension of the latent space, thanks to representation learning.

The Power of Active Multi-Task Learning in Reinforcement Learning from Human Feedback

TL;DR

This work tackles the data-inefficiency problem in reinforcement learning from human feedback (RLHF) by casting RLHF as a contextual dueling bandit problem with a shared, low-dimensional linear representation. It introduces active multi-task learning, which allocates samples across source tasks according to a task relevance vector , and develops algorithms to estimate from limited data when unknown. Theoretical guarantees show that the source-task sample complexity scales with and , while the target-task sample complexity is linear in the latent dimension , thanks to representation learning. The approach supports both known and unknown task relevance settings and provides a practical two-stage procedure (pre-sampling and Lasso-based relevance estimation) to achieve -optimal policies with provable guarantees, offering a principled path to leveraging diverse human feedback efficiently in RLHF.

Abstract

Reinforcement learning from human feedback (RLHF) has contributed to performance improvements in large language models. To tackle its reliance on substantial amounts of human-labeled data, a successful approach is multi-task representation learning, which involves learning a high-quality, low-dimensional representation from a wide range of source tasks. In this paper, we formulate RLHF as the contextual dueling bandit problem and assume a common linear representation. We demonstrate that the sample complexity of source tasks in multi-task RLHF can be reduced by considering task relevance and allocating different sample sizes to source tasks with varying task relevance. We further propose an algorithm to estimate task relevance by a small number of additional data and then learn a policy. We prove that to achieve optimal, the sample complexity of the source tasks can be significantly reduced compared to uniform sampling. Additionally, the sample complexity of the target task is only linear in the dimension of the latent space, thanks to representation learning.
Paper Structure (26 sections, 22 theorems, 131 equations, 2 algorithms)

This paper contains 26 sections, 22 theorems, 131 equations, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Contexual Dueling Bandits.
  4. Multi-task Learning.
  5. Notations
  6. Preliminaries
  7. Active Multi-task Learning.
  8. Fisher Information Assumption.
  9. Learning Objective.
  10. Algorithm and Theory
  11. Results for the Known Task Relevance Setting
  12. Results for the Unknown Task Relevance Setting
  13. Estimating $\nu^*$.
  14. Estimating $\theta^*$.
  15. Conclusion
  16. ...and 11 more sections

Key Result

Theorem 3.1

Suppose Assumptions smp:lowdim, smp:bound, and smp:var hold. Then for any sufficiently small $\varepsilon > 0$, Algorithm alg:dueling outputs an $\varepsilon$-optimal policy with probability at least $1-\delta$ when where the coverage coefficeint $C^*$ is defined as

Theorems & Definitions (27)

  • Theorem 3.1
  • Remark 3.2
  • Remark 3.3
  • Theorem 3.4: Informal
  • Remark 3.5
  • Theorem B.1: Formal Version of Theorem \ref{['thm:known']}
  • Lemma B.2: Approximation of log-likelihood
  • Lemma B.3: Hessian estimation
  • Lemma B.4: Concentration of single task
  • Lemma B.5: Concentration of the norm of $g_m$
  • ...and 17 more