A Critical Look At Tokenwise Reward-Guided Text Generation

TL;DR

This paper critically examines tokenwise reward-guided text generation, revealing limitations of current methods and proposing a new reward model trained on partial sequences that improves generation quality without large-scale fine-tuning.

Contribution

It introduces a novel Bradley-Terry reward model trained on partial sequences and analyzes its properties, outperforming previous heuristic RGTG methods.

Findings

The new reward model is compatible with partial sequence scoring.

The proposed policy is proportional to the ratio of two RLHF policies.

Our method performs similarly to strong offline baselines without large-scale fine-tuning.

Abstract

Large language models (LLMs) can be improved by aligning with human preferences through fine-tuning -- the so-called reinforcement learning from human feedback (RLHF). However, the cost of fine-tuning an LLM is prohibitive for many users. Due to their ability to bypass LLM fine-tuning, prediction-time tokenwise reward-guided text generation (RGTG) methods have recently been proposed. They use a reward model trained on full sequences to score partial sequences during decoding in a bid to steer the generation towards sequences with high rewards. However, these methods have so far been only heuristically motivated and poorly analyzed. In this work, we show that reward models trained on full sequences are not compatible with scoring partial sequences. To alleviate this, we propose to train a Bradley-Terry reward model on partial sequences explicitly, and autoregressively sample from the implied tokenwise policy during decoding. We study the properties of this reward model and the resulting policy: we show that this policy is proportional to the ratio of two distinct RLHF policies. Our simple approach outperforms previous RGTG methods and performs similarly to strong offline baselines without large-scale LLM fine-tuning. Code for our work is available at https://github.com/ahmadrash/PARGS