Enhancing Reinforcement Learning with Dense Rewards from Language Model Critic (2024)

Abstract: Reinforcement learning (RL) can align language models with non-differentiable reward signals, such as human preferences. However, a major challenge arises from the sparsity of these reward signals - typically, there is only a single reward for an entire output. This sparsity of rewards can lead to inefficient and unstable learning. To address this challenge, our paper introduces an novel framework that utilizes the critique capability of Large Language Models (LLMs) to produce intermediate-step rewards during RL training. Our method involves coupling a policy model with a critic language model, which is responsible for providing comprehensive feedback of each part of the output. This feedback is then translated into token or span-level rewards that can be used to guide the RL training process. We investigate this approach under two different settings: one where the policy model is smaller and is paired with a more powerful critic model, and another where a single language model fulfills both roles. We assess our approach on three text generation tasks: sentiment control, language model detoxification, and summarization. Experimental results show that incorporating artificial intrinsic rewards significantly improve both sample efficiency and the overall performance of the policy model, supported by both automatic and human evaluation.

## introduction large language models (llms) have seen a rapid advancement in recent years, demonstrating a remarkable ability to understand and generate natural language @xcite @xcite @xcite . in the meanwhile, † work done during an internship at google. ‡ corresponding author(s): leimeng@google.com, meng.cao@mail.mcgill.ca reinforcement learning has emerged as a complementary tool for further refining the capabilities of lms. rl allows for the optimization of lms towards any non-differentiable reward signal. for example, techniques like reinforcement learning from human feedback (rlhf) @xcite have been used to steer language models to better align with human preferences. however, the reward signals received from the environment are usually sparse, a fundamental bottleneck that restricts the efficiency of learning @xcite . typically, in text generation tasks, a single scalar reward is obtained after a sentence or paragraph has been fully generated. this single reward signal introduces a temporal credit assignment problem, making it difficult for the model to learn which tokens were responsible for the received reward. previous attempts to circumvent the sparsity of rewards in rl have included reward shaping @xcite @xcite , curiosity-driven exploration @xcite @xcite , and hierarchical rl @xcite . however, these methods either require handcrafted features or do not translate straightforwardly into the domain of text generation. a direct solution is to refine the environment's holistic reward model with one that offers dense rewards. @xcite and @xcite have explored employing human annotators to provide detailed feedback at each intermediate step of model's generation. these annotations can then be used to train a fine-grained reward model. however, this method incurs high costs, and the resulting reward models tend to be highly taskspecific, limiting their applicability across different tasks. in light of these limitations, we introduce relc (rewards from language model critique), an novel framework that leverages the critique capability of llms @xcite @xcite to provide artificial reward signals for intermediate steps during rl training. as illustrated in figure 1 , we explicitly define an rl agent as the integration of 1) a policy model responsible for output generation, and 2) a critic model that tasked with assessing the quality of the outputs produced by the policy model. the critic lm, informed by the question, the policy model's output, and the single reward signal provided by the environment, generates verbal evaluation of each segment of the policy model's output. these evaluations are then converted into reward signals. we label the rewards generated by the critic model as "intrinsic rewards" to differentiate them from the reward signals provided by the environment. the critic model can be seamlessly integrated into rl algorithms such as ppo @xcite , requiring no or very little modification to the algorithms themselves. our evaluation of the proposed method is carried out in two distinct settings: one employs a smaller policy model (gpt-2 large) coupled with a more advanced critic model (gpt-3.5), and the other, a more challenging "self-critique" setting, where a single model (llama 2) fulfills both roles. we evaluate the effectiveness of our method through three text generation tasks: sentiment control, lm detoxification, and abstractive text summarization. the experimental results show that the use of llmgenerated intrinsic rewards significantly enhances sample efficiency across three tasks, with our approach outperforming established baseline methods according to both automated and human evaluation. despite the additional inference cost incurred by incorporating the critic model, our approach is shown to be more computationally efficient, achieving superior performance to the baseline within the same computational budget. ## method the basic idea behind our method is to leverage a llm to generate dense intrinsic reward signal r in and provide it to an rl agent, which will optimize a combination of the intrinsic and extrinsic rewards. in this section, we first establish the markov decision process (mdp) for text generation. then, we discuss the policy gradient-based rl method widely used for text generation tasks. finally, we detail the process of incorporating llm-generated intrinsic rewards into rl training. ## experiments in this section, we demonstrate that our method outperforms the ppo baseline in three text generation tasks: sentiment control, lm detoxification, and text summarization. ## conclusion in this work, we introduced a novel framework that incorporates a critic lm to generate dense intrinsic reward signals, addressing the issues of reward sparsity and credit assignment in rl fine-tuning. the critic model evaluates segments of the policy model's output, providing token-or span-level feedback that is converted into numeric rewards. these intrinsic rewards are then combined with extrinsic rewards in rl training. evaluated on sentiment control, detoxification, and summarization tasks, our method not only significantly improve the sample efficiency of the ppo algorithm but also outperformed baseline methods using automatic and human evaluation. furthermore, we demonstrate the effectiveness of the "self-critique" approach across all three tasks, where the same model acts as both the policy and the critic.