Source-linked AI summary

Reinforced Self-Training (ReST) for Language Modeling

Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, Wolfgang Macherey, Arnaud Doucet, Orhan Firat, Nando de Freitas

arXiv:2308.08998v2cs.CLcs.LG

TL;DR

RLHF seeks to align language-model outputs with human preferences, while existing online and offline approaches have efficiency, reward-hacking, or dataset-quality limitations. ReST alternates policy-generated dataset growth with repeated offline policy improvement. In machine translation, it improves reward-model scores and human-rated translation quality, while its scope depends on robust reward models and scalable sample generation.

  • Problem

    RLHF aims to align LLM outputs with human preferences, but online training is costly and reward hacking is possible, whereas offline RL depends on fixed-dataset quality.

  • Method

    ReST alternates an outer Grow step that samples outputs from the latest policy with an inner Improve step that filters samples and fine-tunes using offline RL.

  • Results

    ReST significantly improves reward-model scores and produces higher-quality translations than a supervised-learning baseline according to human raters.

  • Takeaways & Limitations

    ReST is a general and efficient growing-batch RL approach for generative tasks when robust human-preference reward models and scalable sampling are available.

  • Takeaways & Limitations

    Reward-model scores do not necessarily reflect human preferences, and reward models can generalize worse as the policy moves away from the behavior model.

Abstract

from arXiv · show

Reinforcement learning from human feedback (RLHF) can improve the quality of large language model's (LLM) outputs by aligning them with human preferences. We propose a simple algorithm for aligning LLMs with human preferences inspired by growing batch reinforcement learning (RL), which we call Reinforced Self-Training (ReST). Given an initial LLM policy, ReST produces a dataset by generating samples from the policy, which are then used to improve the LLM policy using offline RL algorithms. ReST is more efficient than typical online RLHF methods because the training dataset is produced offline, which allows data reuse. While ReST is a general approach applicable to all generative learning settings, we focus on its application to machine translation. Our results show that ReST can substantially improve translation quality, as measured by automated metrics and human evaluation on machine translation benchmarks in a compute and sample-efficient manner.

1. Introduction

RLHF aligns language models with human preferences, but online and offline approaches face efficiency, reward-hacking, or dataset-quality constraints. ReST addresses these issues by alternating dataset growth from the latest policy with repeated offline policy improvement, and improves translation evaluation outcomes.

  • Motivation: Online RLHF repeatedly samples and scores updated-policy outputs, making continual data generation computationally costly as policy and reward networks grow.Online methods are also described as prone to reward hacking.
  • Motivation: Offline RL is more computationally efficient and less prone to reward hacking, but its policy quality depends on the fixed dataset and its curation.Without carefully curated data, gains over supervised learning may be limited.
  • ReST approach: ReST frames language-model alignment as growing-batch RL with an outer Grow loop and an inner Improve loop.The latest policy generates new samples, while policy improvement uses a fixed dataset.
  • ReST approach: During Grow, the policy generates multiple outputs per context; during Improve, samples are ranked and filtered before offline-RL fine-tuning.The Improve step can be repeated with increasing filtering thresholds, and the resulting policy feeds the next Grow step.
  • Advantages: ReST reuses Grow outputs across several Improve steps, reduces computational burden, and decouples data inspection from policy improvement.The paper also describes the method as simple, stable, and requiring few hyperparameters.
  • Experiments: ReST significantly improves reward-model scores and produces higher-quality translations than a supervised-learning baseline according to human raters.The experiments compare offline RL algorithms on IWSLT 2014, WMT 2020, and internal Web Domain benchmarks.

2. Preliminaries

The preliminaries formulate conditional language generation as producing an output sequence from a source context and define the supervised-learning baseline used in the paper. The policy is autoregressive, and behavioral cloning denotes training with negative log likelihood.

  • Conditional language modeling: A conditional language model produces an output sequence y = (y_1, y_2, ..., y_T) given a context x = (x_1, x_2, ..., x_L).The input and output tokens belong to a chosen vocabulary.
  • Conditional language modeling: The language-generation policy is an autoregressive conditional probability distribution parameterized by θ.Previous output tokens are represented using the convention y_1:0 = ∅ and y_1:t−1 = (y_1, ..., y_t−1).
  • Data distribution: The data distribution is written as p(x, y) = p(x)p(y|x), and dataset D consists of samples from this distribution.This separates the context distribution from the conditional output distribution.
  • Supervised baseline: The supervised policy is trained by minimizing the negative log likelihood (NLL) loss.The paper calls the resulting model behavioral cloning (BC), following reinforcement-learning terminology.

3. Reinforced Self-Training (ReST)

ReST aligns language-model outputs with human preferences by alternating offline policy improvement on filtered data with dataset growth from the latest policy. The approach reuses generated data across Improve steps while progressively selecting higher-reward samples.

  • Grow: Grow samples multiple outputs for contexts from the current policy, scores them with a reward function, and adds the resulting trajectories to the dataset.In the described setting, contexts are sampled from the original dataset and outputs from the language policy.
  • Improve: Improve filters generated data to retain samples whose rewards exceed a threshold, then fine-tunes the current policy on the filtered data.The method can use supervised NLL or offline RL losses, including V-MPO and offline actor-critic.
  • Grow–Improve schedule: Only the final improved policy generates the next dataset, amortizing expensive sampling across several Improve steps.The algorithm therefore reuses one generated dataset during multiple offline updates before performing another Grow step.
  • Scope and limitation: Threshold filtering may learn suboptimal high-variance behaviors in stochastic environments, whereas the paper formulates language modeling and translation as deterministic reinforcement-learning problems.This limitation is explicitly scoped to environments with stochastic dynamics.
  • Improve: Increasing thresholds τ1 < · · · < τN produces smaller, higher-quality subsets, and consecutive fine-tuning starts each policy from its predecessor with a lower learning rate.This design addresses rapid overfitting to small datasets and supports policy improvement while the dataset remains fixed.

4. Experiments and analysis

ReST is evaluated on machine-translation benchmarks across multiple datasets, language pairs, offline RL losses, training iterations, inference sampling, and human ratings. Across these analyses, ReST improves reward-based and human-evaluated performance over supervised training, while additional Improve and Grow steps and loss choice affect outcomes.

  • Experimental setup: ReST is evaluated on IWSLT 2014, WMT 2020, and Web Domain datasets using separate validation and test sets with unseen source sentences.The experiments use different language pairs across datasets to test generality, and Metric X supplies reference-free reward scores.
  • Improve steps: Each successive Improve step significantly increases average reward across all three translation datasets, with the suitable number of steps depending on the dataset.The experiments used five Improve steps on IWSLT 2014, four on WMT-2020, and two on Web Domain.
  • Grow steps: A second Grow step with subsequent Improve steps adds 5.3 points on IWSLT 2014 and 0.8 points on Web Domain over the first Grow step.The additional Grow step was evaluated specifically for its effect on reward-model performance.
  • Comparison with supervised learning: Different ReST variants significantly outperform supervised learning after only one Grow step, consistently across the tested datasets and language pairs.The supervised baseline is BC (G=0, I=0), while ReST variants use different losses and Grow/Improve counts.
  • Offline RL losses: BC loss outperforms the other tested offline RL losses in a single ReST step, while reward filtering incorporates the reward for Improve steps with I≥1.Figure 5 compares OAC, BVMPO, GOLD, and BC-based variants; BC loss also performs better than GOLD across increasing Improve steps.
  • Online RL and human evaluation: With comparable training data, multiple Improve steps outperform online PPO; PPO’s validation BLEU drops by nearly 8 points, whereas ReST’s BLEU does not change.All ReST variants also outperform BC in human ratings, although reward-model and human-score rankings differ.

5. Related works

ReST is compared with supervised learning, self-training, online RL, and offline RL as a method combining compute efficiency with exploration data and rewards. Related approaches differ in their use of planning, iteration, agents, RL losses, and filtering schedules.

  • ReST is described as the only compared approach that is compute efficient while leveraging exploration data and rewards.
  • Expert Iteration (EI): Expert Iteration also uses policy-generated data, but ReST needs no planning mechanism and reuses data through iterative Improve steps.
  • Reasoning with language models: Impossible Distillation corresponds to ReST with a single Grow and Improve step, whereas ReST supports varied offline RL losses, planning choices, and filtering mechanisms.
  • Iterated Learning (IL): Iterated Learning differs from ReST because it uses a multi-agent setting and does not use reinforcement learning.
  • Self Imitation Learning (SIL): ReST is closely related to Self Imitation Learning through filtering unsuccessful trajectories, but remains agnostic to the underlying RL algorithm.
  • Reward ranked Fine-Tuning (RAFT): RAFT is characterized as one Improve step per Grow step with a fixed reward-quantile threshold, while ReST uses multiple Improve steps with increasing thresholds for further gains.

6. Discussion

The discussion presents ReST as a simple, flexible, and efficient alignment method, while identifying reward-model overfitting and exploration limitations. Its application requires a robust human-preference reward model and scalable sample generation.

  • ReST has minimal hyperparameters and supports many designs of Grow and Improve steps.
  • Behavior cloning performed best among tested offline RL losses for improving reward-model scores, but reward gains did not necessarily reflect human preferences.
  • Human evaluation favored one Grow step, even though reward-model scores continued increasing with additional Grow steps.
  • Repeated Grow iterations increase the risk of overfitting to the reward model, motivating future reward-model fine-tuning on human-annotated generated data.
  • ReST could benefit from stronger Grow-step exploration, while sparse rewards, credit assignment, hyperparameter sensitivity, and limited exploration challenge offline RL losses.
  • ReST applies broadly when a robust human-preference reward model exists and samples can be generated from the model at scale.

A.1. RLHF for conditional language modeling as MDP

Conditional language modeling can be represented as a sequence-to-sequence mapping and as a Markov Decision Process. In this formulation, token generation is action selection and human-preference rewards score completed outputs.

  • Conditional language modeling maps a source sequence x to a target sequence y, with machine translation as a classic sequence-to-sequence example.
  • An offline MDP consists of states, actions, transitions, rewards, an initial-state distribution, and a dataset of state-action-reward sequences.
  • In language modeling, each state contains the input and generated prefix, while each action is the next token produced by the policy.
  • The reward can be learned from human preferences and is assigned to the full generated sequence after the end-of-sequence token.
  • This conditional-language-modeling RLHF formulation can also be viewed as a contextual bandit with a very large action space.

A.2. Negative results with offline RL

Experiments with several offline RL alternatives did not improve over the supervised baseline, with Q-function methods performing worse. The authors relate these results to the difficulty of offline Q-learning in large action spaces and sparse-reward settings.

  • Q-learning and reward-conditioned approaches did not obtain notable improvements over the supervised baseline.
  • Q-function-based methods performed worse than supervised learning even when initialized from a supervised checkpoint.
  • The authors connect these results to offline Q-learning difficulty in vocabularies of 32 000 tokens and the tendency of return-conditioned methods to learn sub-optimal policies with sparse continuous rewards.

A.3. Data and Model details

ReST experiments use Transformer-based translation models across three datasets, with large generated candidate sets and reward-based iterative improvement. Training and evaluation report dataset sizes and validation reward scores.

  • Model configuration: ReST uses a modified Transformer architecture with a 32,000-token vocabulary and maximum decoding length of 128.These settings apply across the experiments.
  • Datasets: The IWSLT 2014 De-En experiments generate 100 candidate translations per source sentence, yielding |D_g| = 16,000,000.
  • Datasets: The WMT 2020 Zh-En experiments generate 25 candidates per source sentence, yielding |D_g| = 890,000,000.
  • Datasets: The Web Domain En-Zh experiments use an in-house corpus and a modified Transformer-big encoder-decoder.The model has dimension 1024, feedforward dimension 8192, 16 attention heads, and six encoder and decoder layers.
  • Evaluation: Unless stated otherwise, experiments report average reward scores on the validation set.

A.4. Reward model

The experiments use learned translation reward models that score complete translations, with reference-free and reference-based variants. The reference-free model is more flexible but remains imperfect despite unit testing.

  • Reward models: The reward models assign a score to the whole translation.
  • Reward-model types: Reference-free reward models score source-candidate pairs, whereas reference-based models additionally use the human reference translation.
  • Reward-model types: Reference-free reward models are more flexible because they do not require reference translations.
  • Validation and limitation: Unit tests on permutations and repetitions supported reward-model development, but the reference-free model sometimes increased scores when translations repeated sentences.

A.5. Alignment between human evaluation scores and the reward model scores

ReST improves test-set rewards over supervised and value-based offline RL baselines, while performance depends on the loss, filtering strategy, and number of Grow and Improve steps. Reward-model agreement is weaker for low-quality translations.

  • Human–reward alignment: ReST reward-model scores have high variance for samples with low human preference scores, likely because the training data overrepresented high-quality translations.The authors suggest incremental reward-model retraining as a possible remedy.
  • Test-set results: All ReST variations outperform the supervised learning baseline on IWSLT 2014 De-En, and test-set reward increases with more Grow and Improve steps.
  • Test-set results: ReST clearly outperforms baselines on WMT 2020 Zh-En, with test-set reward increasing as Improve steps increase.
  • Loss ablations: BC loss outperforms GOLD and other tested offline RL losses across increasing Improve steps on WMT 2020 Zh-En and IWSLT 2014 De-En.
  • Filtering ablations: Filtering by higher source-specific reward percentiles improves performance, with gains saturating after p = 90.
  • Threshold ablations: The number of Improve steps and threshold selection strongly influence final performance, including thresholds interpolated between source-specific mean and maximum rewards.
Loading 2308.08998v2…