Towards Cost-Effective Reward Guided Text Generation

TL;DR

This paper introduces a novel reward model architecture for reward-guided text generation that enables faster inference by scoring all candidate tokens simultaneously with a single call, improving efficiency and performance.

Contribution

The paper proposes a new reward model trained with a Bradley-Terry loss for efficient, step-wise preference scoring during text generation, reducing inference overhead.

Findings

Faster inference compared to existing RGTG methods

Requires fewer calls to the reward model during generation

Performs competitively with previous RGTG and RLHF approaches

Abstract

Reward-guided text generation (RGTG) has emerged as a viable alternative to offline reinforcement learning from human feedback (RLHF). RGTG methods can align baseline language models to human preferences without further training like in standard RLHF methods. However, they rely on a reward model to score each candidate token generated by the language model at inference, incurring significant test-time overhead. Additionally, the reward model is usually only trained to score full sequences, which can lead to sub-optimal choices for partial sequences. In this work, we present a novel reward model architecture that is trained, using a Bradley-Terry loss, to prefer the optimal expansion of a sequence with just a \emph{single call} to the reward model at each step of the generation process. That is, a score for all possible candidate tokens is generated simultaneously, leading to efficient inference. We theoretically analyze various RGTG reward models and demonstrate that prior techniques prefer sub-optimal sequences compared to our method during inference. Empirically, our reward model leads to significantly faster inference than other RGTG methods. It requires fewer calls to the reward model and performs competitively compared to previous RGTG and offline RLHF methods.