ESRL: Efficient Sampling-based Reinforcement Learning for Sequence Generation

TL;DR

ESRL introduces two-stage and dynamic sampling methods to enhance the efficiency of reinforcement learning in sequence generation tasks, significantly reducing computational costs while maintaining or improving performance.

Contribution

The paper proposes novel sampling strategies that improve RL training efficiency for sequence models, applicable to tasks like translation and summarization, including RLHF.

Findings

ESRL outperforms baseline methods in training efficiency and memory use.

ESRL achieves better or comparable performance on translation and summarization tasks.

The approach is effective in RLHF for large language models.

Abstract

Applying Reinforcement Learning (RL) to sequence generation models enables the direct optimization of long-term rewards (\textit{e.g.,} BLEU and human feedback), but typically requires large-scale sampling over a space of action sequences. This is a computational challenge as presented by the practice of sequence generation problems, such as machine translation, where we often deal with a large action space (\textit{e.g.,} a vocabulary) and a long action sequence (\textit{e.g.,} a translation). In this work, we introduce two-stage sampling and dynamic sampling approaches to improve the sampling efficiency during training sequence generation models via RL. We experiment with our approaches on the traditional sequence generation tasks, including machine translation and abstractive summarization. Furthermore, we evaluate our approaches in RL from human feedback (RLHF) through training a large language model using the reward model. Experimental results show that the efficient sampling-based RL, referred to as ESRL, can outperform all baselines in terms of both training efficiency and memory consumption. Notably, ESRL yields consistent performance gains over the strong REINFORCE, minimum risk training, and proximal policy optimization methods.