Dense Reward for Free in Reinforcement Learning from Human Feedback

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Chan, Alex J., Sun, Hao, Holt, Samuel, van der Schaar, Mihaela
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866916112260661248
author Chan, Alex J.
Sun, Hao
Holt, Samuel
van der Schaar, Mihaela
author_facet Chan, Alex J.
Sun, Hao
Holt, Samuel
van der Schaar, Mihaela
contents Reinforcement Learning from Human Feedback (RLHF) has been credited as the key advance that has allowed Large Language Models (LLMs) to effectively follow instructions and produce useful assistance. Classically, this involves generating completions from the LLM in response to a query before using a separate reward model to assign a score to the full completion. As an auto-regressive process, the LLM has to take many "actions" (selecting individual tokens) and only receives a single, sparse reward at the end of an episode, a setup that is known to be difficult to optimise in traditional reinforcement learning. In this work we leverage the fact that the reward model contains more information than just its scalar output, in particular, it calculates an attention map over tokens as part of the transformer architecture. We use these attention weights to redistribute the reward along the whole completion, effectively densifying the signal and highlighting the most important tokens, all without incurring extra computational cost or requiring any additional modelling. We demonstrate that, theoretically, this approach is equivalent to potential-based reward shaping, ensuring that the optimal policy remains unchanged. Empirically, we show that it stabilises training, accelerates the rate of learning, and, in practical cases, may lead to better local optima.
format Preprint
id arxiv_https___arxiv_org_abs_2402_00782
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Dense Reward for Free in Reinforcement Learning from Human Feedback
Chan, Alex J.
Sun, Hao
Holt, Samuel
van der Schaar, Mihaela
Machine Learning
Reinforcement Learning from Human Feedback (RLHF) has been credited as the key advance that has allowed Large Language Models (LLMs) to effectively follow instructions and produce useful assistance. Classically, this involves generating completions from the LLM in response to a query before using a separate reward model to assign a score to the full completion. As an auto-regressive process, the LLM has to take many "actions" (selecting individual tokens) and only receives a single, sparse reward at the end of an episode, a setup that is known to be difficult to optimise in traditional reinforcement learning. In this work we leverage the fact that the reward model contains more information than just its scalar output, in particular, it calculates an attention map over tokens as part of the transformer architecture. We use these attention weights to redistribute the reward along the whole completion, effectively densifying the signal and highlighting the most important tokens, all without incurring extra computational cost or requiring any additional modelling. We demonstrate that, theoretically, this approach is equivalent to potential-based reward shaping, ensuring that the optimal policy remains unchanged. Empirically, we show that it stabilises training, accelerates the rate of learning, and, in practical cases, may lead to better local optima.
title Dense Reward for Free in Reinforcement Learning from Human Feedback
topic Machine Learning
url https://arxiv.org/abs/2402.00782