Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

T-REG: Preference Optimization with Token-Level Reward Regularization

Wenxuan Zhou, Shujian Zhang, Lingxiao Zhao, Tao Meng

TL;DR

T-REG addresses token-level credit assignment in RLHF by introducing token-level reward regularization that combines sequence-level preferences with self-generated token rewards derived from contrastive prompting. It derives token-level rewards without external annotations and uses them as a regularization term to steer token-level credit assignment during preference optimization. The approach yields consistent improvements on instruction-following benchmarks such as Alpaca Eval 2 and Arena-Hard, and demonstrates compatibility with different preference optimization frameworks, while providing qualitative evidence of more precise token-level reward signals. This work advances practical alignment by enabling finer-grained credit assignment without additional labeling, potentially improving robustness and fidelity of LLM behavior in user-facing tasks.

Abstract

Reinforcement learning from human feedback (RLHF) has been crucial in aligning large language models (LLMs) with human values. Traditionally, RLHF involves generating responses to a query and using a reward model to assign a reward to the entire response. However, this approach faces challenges due to its reliance on a single, sparse reward, which makes it challenging for the model to identify which parts of the sequence contribute most significantly to the final reward. Recent methods have attempted to address this limitation by introducing token-level rewards. However, these methods often rely on either a trained credit assignment model or AI annotators, raising concerns about the quality and reliability of the rewards. In this paper, we propose token-level reward regularization (T-REG), a novel approach that leverages both sequence-level and token-level rewards for preference optimization. Harnessing the self-refinement capabilities of LLMs, our method uses contrastive prompting to enable LLMs to self-generate token-level rewards. These self-generated rewards then act as reward regularization, guiding the model to more effectively distribute sequence-level rewards across tokens. This facilitates better token-level credit assignment and enhances alignment performance. Experiments on the instruction following benchmarks, including Alpaca Eval 2 and Arena-Hard, show that our method consistently outperforms baseline methods by up to 3.8% and 4.4%, respectively. We will release the code and models at https://github.com/wzhouad/T-REG.

T-REG: Preference Optimization with Token-Level Reward Regularization

TL;DR

T-REG addresses token-level credit assignment in RLHF by introducing token-level reward regularization that combines sequence-level preferences with self-generated token rewards derived from contrastive prompting. It derives token-level rewards without external annotations and uses them as a regularization term to steer token-level credit assignment during preference optimization. The approach yields consistent improvements on instruction-following benchmarks such as Alpaca Eval 2 and Arena-Hard, and demonstrates compatibility with different preference optimization frameworks, while providing qualitative evidence of more precise token-level reward signals. This work advances practical alignment by enabling finer-grained credit assignment without additional labeling, potentially improving robustness and fidelity of LLM behavior in user-facing tasks.

Abstract

Reinforcement learning from human feedback (RLHF) has been crucial in aligning large language models (LLMs) with human values. Traditionally, RLHF involves generating responses to a query and using a reward model to assign a reward to the entire response. However, this approach faces challenges due to its reliance on a single, sparse reward, which makes it challenging for the model to identify which parts of the sequence contribute most significantly to the final reward. Recent methods have attempted to address this limitation by introducing token-level rewards. However, these methods often rely on either a trained credit assignment model or AI annotators, raising concerns about the quality and reliability of the rewards. In this paper, we propose token-level reward regularization (T-REG), a novel approach that leverages both sequence-level and token-level rewards for preference optimization. Harnessing the self-refinement capabilities of LLMs, our method uses contrastive prompting to enable LLMs to self-generate token-level rewards. These self-generated rewards then act as reward regularization, guiding the model to more effectively distribute sequence-level rewards across tokens. This facilitates better token-level credit assignment and enhances alignment performance. Experiments on the instruction following benchmarks, including Alpaca Eval 2 and Arena-Hard, show that our method consistently outperforms baseline methods by up to 3.8% and 4.4%, respectively. We will release the code and models at https://github.com/wzhouad/T-REG.

Paper Structure

This paper contains 12 sections, 13 equations, 3 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Reinforcement learning from human feedback.
  4. Method
  5. Preliminaries
  6. Regularized Token-level Preference Optimization
  7. Self-generated Token-level Rewards
  8. Experiments
  9. Experimental Settings
  10. Main Results and Ablation
  11. Case Study
  12. Conclusion

Figures (3)

  • Figure 1: Overview of the Token-Level Reward Regularization (T-REG). Existing DPO directly optimizes the sequence-level rewards to align with user preferences. T-REG prompts LLMs to generate the token-level reward and utilize it as regularization during the preference optimization.
  • Figure 2: Prompt template for revising the output to either a better or worse one. To make the output better, we set use "helpful, correct, coherent, concise" as the description. To make the output worse, we use "unhelpful, incorrect, incoherent, verbose" as the description.
  • Figure 3: Case study on the token-level rewards learned by DPO and DPO-REG, where red means positive reward and blue means negative reward. We use Llama-3-Instruct (8B) as the base model.