TLCR : Token-Level Continuous Reward for Fine-grained Reinforcement Learning from Human Feedback (2024)

Abstract: Reinforcement Learning from Human Feedback (RLHF) leverages human preference data to train language models to align more closely with human essence. These human preference data, however, are labeled at the sequence level, creating a mismatch between sequence-level preference labels and tokens, which are autoregressively generated from the language model. Although several recent approaches have tried to provide token-level (i.e., dense) rewards for each individual token, these typically rely on predefined discrete reward values (e.g., positive: +1, negative: -1, neutral: 0), failing to account for varying degrees of preference inherent to each token. To address this limitation, we introduce TLCR (Token-Level Continuous Reward) for RLHF, which incorporates a discriminator trained to distinguish positive and negative tokens, and the confidence of the discriminator is used to assign continuous rewards to each token considering the context. Extensive experiments show that our proposed TLCR leads to consistent performance improvements over previous sequence-level or token-level discrete rewards on open-ended generation benchmarks.

## introduction with the emergent properties of large language models (llms) in relation to the scale of both the model and data, there has been a notable performance improvement, particularly in various language generation tasks including summarization @xcite , translation @xcite , and question-answering @xcite . central to these rapid performance gains in language generation is the reinforcement learning from human feedback (rlhf), which employs reinforcement learning techniques to fine-tune llms in alignment with human preferences @xcite @xcite . ## preliminary in this work, we focus on aligning the language model to human preference in terms of helpfulness and harmfulness of the generated texts through reinforcement learning from human feedback (rlhf). language generation as an mdp we consider the language generation procedure as a markov decision process (mdp) as detailed in ramamurthy et al. @xcite , which is defined by the tuple (s, a, p, r, γ, t max ) using finite predefined vocabulary tokens v. an episode in the mdp starts with a human instruction or a history of human-assistant conversation pairs followed by a human instruction @xmath0, ..., x l ). this is used as our initial state s 0 = (x 0 , ..., x l ), where s 0 ∈ s and s represents every possible states with x @xmath1. at time step t, an action in the environment a @xmath2. for @xmath3, the transition function p : s × @xmath4, ...x l , a 0 , ..., a t-1 ). this process continues until the current t exceeds t max or the end of sentence (<@xmath5, yielding the generated sequence @xmath6, ..., a t ). a reward function r : s × @xmath7. sequence-level reward in rlhf typically, in rlhf for language models, a sequence-level (i.e., sparse) reward model is used, which maps the concatenation of the input x and the generated output y to a single scalar value. this scalar value is solely assigned to the last generated token or the end of sentence (<@xmath8. formally, this process of sparse reward modeling used in previous rlhf frameworks can be defined as: token-level discrete reward in rlhf in contrast to the sparse reward approach, token-level discrete rewards in rlhf assign a predefined discrete value at each time step @xmath9, t ], reflecting the preference for each token, formalized as: ## tlcr: token-level continuous reward in this section, we introduce an rlhf framework that utilizes a discriminator-based token-level re-ward model that allows continuous-scale dense reward signals for rlhf. section 4.1 details the procedure for obtaining token-level preference labels. in section 4.2, we elaborate on the training procedure of the token-level preference discriminator. subsequently, section 4.3 outlines how the trained discriminator is utilized for continuous dense reward signals during the rl phase. ## experimental settings this section presents the datasets, models, compared baselines, training details, and the evaluation metrics used in the paper. dataset throughout the experiments, we use the full-hh-rlhf dataset @xcite which is centered on improving the helpfulness and harmlessness of the language model generation. this dataset includes 112k training instances and 12.5k instances for evaluation. every instance in the dataset features a prompt, along with a chosen response considered preferable and a rejected response, offering a clear basis for performing model alignment for better helpfulness and harmlessness. following the setting in @xcite , we randomly divide the dataset into three parts: 20% for supervised fine-tuning (sft), 40% for reward model learning, and 40% for reinforcement learning for reward maximization. model during the supervised fine-tuning (sft) phase we use llama-2-7b @xcite as our base model. similarly, for the reward model training, we initialize our discriminator d ϕ with llama-2-7b. during the rl phase, we initialize our policy model with the trained sft model. for the external language model lm ext , we use gpt-4-0125-preview 1 . training details we use deepspeed-chat @xcite framework for performing rlhf, where ppo is the default algorithm. within this framework, our proposed tlcr is integrated to provide token-level rewards to enhance the ppo algorithm. experimental details are presented in appendix b. all experiments were conducted 3 times with different random seeds, and the average value is reported. the experiments were conducted using 8 x nvidia a100 80gb pcie. evaluation automatic evaluation of instructionfollowing language models are becoming a primary metric for preference evaluation with leaderboards foot_0 for ranking. in light of this trend, , we report mt-bench @xcite and alpacaeval (al-paca_eval_gpt4) @xcite for our evaluations. also, we perform human evaluation on 100 random samples from the full-hh-rlhf test set. ## conclusions in this study, we introduced the token-level continuous reward (tlcr), a novel reward model aimed at providing detailed, token-based continuous rewards for reinforcement learning from human feedback (rlhf). by utilizing the prediction confidences from a discriminator to assign rewards at the token level, our approach offers a sophisticated reward system that adjusts to the varying degrees of relevance each token holds within its context. the superiority of tlcr over traditional sequence-level and token-level discrete reward mechanisms has been established through detailed experiments and analyses.