Vol. XIX · Reinforcement Learning for Reasoning Plate 6 of 12rejection sampling, best-of-n, expert iteration, STaR, self-training, filtering, RFTLE

Keeping the Good Samples.

Sample, keep what the verifier accepts, fine-tune on it, repeat.

Before any policy gradient, there is a simpler loop. Sample n answers per prompt from the current model; keep those the verifier accepts; fine-tune on the kept ones with the ordinary supervised loss; repeat. This is expert iteration, or rejection-sampling fine-tuning; STaR (Zelikman et al., 2022) applied it to reasoning traces, and Llama 2 and later models used rejection sampling against a reward model as a step of post-training. At inference the same idea is best-of-n: sample n answers and return the one the verifier — or a reward model — likes best.

Expert iteration is RL with a particular advantage: +1 for the kept samples, 0 for the rest, applied with supervised training. It never pushes down wrong answers, so it can be slower and less sharp than a policy gradient, but it is stable and simple. The workshop runs it on the RL world for a chosen number of rounds and samples per prompt.

Below: held-out accuracy after each round of expert iteration, and the number of training prompts for which at least one sample was kept.

What you will learn
Workshop
From the fine-tuned model: each round samples the chosen number of answers at temperature 1 for each of the 70 training prompts, keeps the prompts with at least one verified answer, and fine-tunes for 50 steps (batches of 16, step size 0.001, loss on answers). Held-out accuracy: greedy answers on the 30 held-out sums.

Climbing on its own successes

With four samples per prompt, the first round keeps 60 of the 70 training prompts — those where at least one sample was right — and three rounds take held-out accuracy from 53% to 77%. With a single sample per prompt, round one keeps 53 prompts and three rounds reach 80%: on prompts this easy for the model one sample usually finds something to keep, and the difference from four samples is within run-to-run noise. On hard prompts, where success is rare, n decides whether there is anything to learn from at all. Nothing in the loop knew the answers except the verifier; the model improved by imitating its own verified attempts.

The loop’s limits are the policy gradient’s, made plain. It learns only where the model sometimes succeeds — a prompt it never solves in n samples contributes nothing, which is why n matters. It never says which wrong answers to avoid, so mistakes fade only as right answers crowd them out. And it can narrow the model towards whatever it already does well, reducing diversity in the same way as other RL (plate XIX.9). In practice rejection sampling is a robust first stage, and best-of-n with a verifier is the simplest way to spend more compute at inference (plate XIX.10).

Expert iteration on a fixed set of prompts eventually saturates: once every prompt yields a correct sample, more rounds only reinforce the same answers. New, harder prompts are what keep it climbing.

Terms

Rejection sampling
Generating samples and keeping only those that pass a check.
Expert iteration
Alternating between sampling verified solutions and training on them.
Best-of-n
Sampling n answers and returning the highest-scoring one.
STaR
Self-taught reasoner: expert iteration on reasoning traces.
Sample, keep what the verifier accepts, fine-tune, repeat: expert iteration climbs on the model’s own successes, wherever it has some.

Sources

  1. Zelikman, Wu, Mu and Goodman, “STaR: Bootstrapping Reasoning With Reasoning”, NeurIPS 2022.
  2. Anthony, Tian and Barber, “Thinking Fast and Slow with Deep Learning and Tree Search”, NeurIPS 2017 (expert iteration).
  3. Touvron et al., “Llama 2”, arXiv:2307.09288, 2023, section 3.2.3 (rejection sampling).