Rejection-sampling fine-tuning, STaR and expert iteration, explained

Rejection-sampling fine-tuning (RFT) improves a model by training it on its own best outputs. For each prompt, sample several responses from the current model, score them with a verifier or reward model, keep the ones that pass (or the single best), and fine-tune on the kept set with the ordinary SFT loss. Repeat with the improved model. The same loop appears under several names: STaR (Self-Taught Reasoner) when the outputs are reasoning chains filtered by answer correctness, expert iteration in the reinforcement-learning tradition where a search procedure plays the expert, ReST and ReST^EM in the self-training literature, RAFT and best-of-$K$ fine-tuning when a reward model does the ranking. It sits between SFT and full RL in the post-training stack: it is the simplest way to turn a verifier into training signal, it is often the stage that bootstraps reasoning data before GRPO or RLVR, and in several published pipelines it has been the RL stage.

Where it came from#

The idea of alternating between an improvement operator and imitation is older than LLMs. Anthony, Tian and Barber (2017) named expert iteration (ExIt): tree search acts as an expert that plays better than the current network, the network is trained to imitate the search, and the improved network makes the next search stronger. Their agent, trained from scratch, beat MoHex 1.0 at Hex. Polu et al. (2022) applied expert iteration to formal mathematics in Lean: proof search interleaved with learning on found proofs dramatically outperformed proof search alone at the same compute and discovered a curriculum of harder problems without ground-truth proofs.

For natural-language reasoning, Zelikman, Wu, Mu and Goodman (2022) introduced STaR. Starting from a few-shot prompt with about ten worked rationales, the model generates rationales for many questions, keeps those whose final answer is correct, and fine-tunes on them. For questions it got wrong, it is shown the correct answer as a hint and asked to produce a rationale for it (rationalisation); those rationales are added with the hint removed. STaR on GPT-J performed comparably to fine-tuning a 30× larger model on CommonsenseQA. Yuan et al. (2023) studied the method as data augmentation for maths and named it rejection-sampling fine-tuning: collecting correct reasoning paths from fine-tuned models, with more distinct paths giving larger gains, took LLaMA-7B on GSM8K from 35.9% with SFT to 49.3% when combining samples from several models. In the alignment setting, Llama 2-Chat (Touvron et al., 2023) used best-of-$K$ rejection sampling against a reward model as its only RLHF algorithm until RLHF V4, after which PPO was applied on top of the rejection-sampling checkpoint. Gulcehre et al. (2023) framed the loop as growing-batch offline RL (ReST) for translation, and Singh et al. (2023) simplified it to ReST^EM, an expectation-maximisation procedure with binary rewards, showing on MATH and APPS with PaLM 2 that it scaled with model size and significantly surpassed fine-tuning on human data. Havrilla et al. (2024) then compared expert iteration, PPO and return-conditioned RL for LLM reasoning and found all performed comparably, with expert iteration best in most cases and with sample complexity similar to PPO's.

The objective#

Let ${\pi }_{\theta }$ be the policy, $x$ a prompt from a set $X$, $y$ a response, and $R(x,y)\in \{0,1\}$ a binary verifier reward (answer matches, tests pass). One RFT iteration samples $N$ responses per prompt from a frozen copy ${\pi }_{{\theta }_k}$ of the current policy and minimises

where $\mathcal{K}$ is the kept set. With a real-valued reward model, $\mathcal{K}$ is instead the top-1 (or top fraction) of the $N$ samples per prompt, or all samples above a threshold $\tau$.

Why it works: two derivations

As policy gradient. The RL objective is $J(\theta )={\mathbb{E}}_x{\mathbb{E}}_{y\sim {\pi }_{\theta }}[R(x,y)]$, with gradient $\nabla J=\mathbb{E}[R(x,y){\nabla }_{\theta }\log {\pi }_{\theta }(y\mid x)]$. With binary $R$, the samples with $R=0$ contribute nothing, so this is the gradient of the log-likelihood of the correct samples. The RFT gradient at $\theta ={\theta }_k$ is exactly that, up to normalisation. STaR's authors make this observation: the filtered fine-tuning step approximates a policy gradient with the 0/1 reward. The approximations are that the samples are drawn once from ${\pi }_{{\theta }_k}$ and reused for many gradient steps (so the update becomes off-policy as $\theta$ moves), and that there is no baseline and no negative signal from failures.

As expectation-maximisation. Singh et al. derive the same loop from EM on the probability that a response is "optimal". The E-step forms the posterior $q_{k+1}(y\mid x)\propto R(x,y){\pi }_{{\theta }_k}(y\mid x)$: the current policy restricted to correct responses. The M-step fits the policy to it,

which, estimated with samples, is the RFT loss. The EM view explains the structure of the algorithm: data collection (E-step) and optimisation (M-step) are decoupled, the policy used to sample is fixed for a whole iteration, and each iteration provably does not decrease a lower bound on expected reward under the idealised assumptions. It also makes clear what the policy converges toward: the base model's distribution, renormalised onto correct answers, sharpened iteration by iteration.

Side by side with its neighbours

The difference from GRPO that matters most in practice is the sign. RFT only ever pushes probability toward correct samples; GRPO also pushes it away from incorrect ones, using the group mean as a baseline. Negative signal makes GRPO more sample-efficient on problems the model already sometimes solves, at the cost of an online sampling loop and more moving parts.

Worked example: four samples, two kept#

Take one maths problem and $N=4$ samples with model probabilities (per-sequence, for illustration) of 0.20, 0.05, 0.10 and 0.30, and verifier outcomes 1, 0, 1, 0. The kept set is samples 1 and 3, and their contribution to the loss is $-\frac{1}{2}(\ln 0.20+\ln 0.10)=\frac{1}{2}(1.609+2.303)=1.956$ nats. Samples 2 and 4 contribute nothing, even though sample 4 is the model's most likely response and is wrong; RFT never tells the model to make it less likely. The gradient raises the probability of samples 1 and 3, and because probabilities sum to one, everything else, including sample 4, loses mass indirectly. A GRPO update on the same group would compute advantages $(R-0.5)/0.5=\pm 1$ and push sample 4 down explicitly.

Two consequences of the arithmetic matter at the dataset level. A problem the model solves in 8 of 8 samples contributes 8 kept examples; one it solves in 1 of 8 contributes 1; one it never solves contributes 0. Without intervention, the kept set is dominated by problems the model can already do, and the hardest problems, where improvement is most needed, give no signal at all. This is why STaR added rationalisation, why ReST^EM caps the number of solutions per problem, and why practitioners deduplicate.

A toy run: iterations, caps and diversity#

We simulated RFT with a log-linear policy over 16 candidate solution paths per problem, 400 training and 400 held-out problems. Each problem has between 0 and 6 correct paths, and correct paths share a feature direction (the "skill") that the base policy is slightly biased against, so the model can generalise what it learns to unseen problems. Each iteration samples $N=8$ paths per training problem from the current policy, keeps the correct ones, and fine-tunes from the base parameters (as STaR and ReST^EM do) for 30 epochs of SGD. We compute held-out pass@1 and pass@8 exactly from the policy's probabilities. Seeds are fixed. One run keeps every correct sample; the other caps the kept set at 2 per problem.

Iterating works: held-out pass@1 nearly tripled over five iterations, and the share of training problems with at least one correct sample rose from 54.5% to about 70%. The base model already solved half of the held-out problems within 8 tries (pass@8 of 51.7%), so the first iterations mostly convert "sometimes right" into "usually right". The cap helped despite using less than half as much data: problems with 5 or more correct paths are 25.8% of the training set, but made up 41 to 47% of the uncapped kept set and 35 to 42% of the capped one. Pass@8 in the uncapped run peaked at iteration 3 and then fell while pass@1 kept rising, a diversity-for-reliability trade of the kind Havrilla et al. report during SFT (maj@1 against pass@96); the capped run held pass@8 steady. Finally, 6.8% of training problems have no correct path at all; no number of iterations can learn from them, and in a real model they correspond to problems beyond its current reach, which need hints, a teacher, or a curriculum.

The algorithm, step by step#

base = load(sft_checkpoint)                 # or pretrained model + few-shot prompt (STaR)
policy = base
for k in range(n_iterations):                # 2 to 5 is typical
    kept = []
    for x, answer in prompts:                # thousands to hundreds of thousands
        ys = policy.sample(x, n=N, temperature=0.7..1.0)          # N = 4 to 64
        good = [y for y in ys if verify(y, answer)]                # exact match, unit tests, reward >= tau
        if not good and use_rationalisation:
            good = [y for y in policy.sample(x + hint(answer), n=N2) if verify(y, answer)]
        good = dedup(good)[:cap]                                    # cap per prompt, e.g. 2 to 4
        kept += [(x, y) for y in good]            # hint removed from x
    kept += best_samples_from_previous_iterations   # optional; Llama 2 found this mattered
    policy = sft(start=base, data=kept, epochs=1..3)   # fresh from base each time (STaR, ReST-EM)
    evaluate(policy, held_out)                  # pass@1, pass@k, per-difficulty accuracy

Data needed: prompts with a way to check responses. For maths that is a final answer; for code, unit tests; for instruction following, a reward model or rubric grader. Prompt sets range from a few thousand (GSM8K's 7.5k training problems were the basis of Yuan et al.'s experiments) to hundreds of thousands. Samples per prompt in published work range from a handful to 64; ReST^EM generated 32 per problem for MATH and 64 for APPS, at temperature 0.7 with top-$K$ 40.

Two design choices recur. Restart from the base each iteration: STaR trains from the original pretrained model each time "to avoid overfitting", and ReST^EM found that fine-tuning the base rather than the previous iterate gave comparable task performance and much better transfer to held-out tasks. Keep samples across iterations: Llama 2's authors found that training each iteration only on samples from the previous one led to regressions (RLHF V3 got worse at rhyming poetry) and fixed it by including top samples from all prior iterations.

Hyperparameters that matter#

The most useful single diagnostic is the solve-rate histogram: for each training prompt, the fraction of its $N$ samples that pass. Prompts at 0% give no signal, prompts at 100% give little new information, and the useful middle should be a large share. If most prompts sit at the extremes, change the prompt mix rather than the hyperparameters.

Compute and memory cost#

Each iteration is a large batch of generation followed by an SFT run. Generation costs about $2P$ FLOPs per generated token and usually dominates: 50k prompts × 16 samples × 500 tokens is 400M generated tokens per iteration, while the kept set might be 100k examples and 50M tokens of SFT, about $6P$ FLOPs each. Memory is never more than the larger of an inference deployment and an SFT run, and the two never need to coexist, which is the main practical advantage over online RL: no reference model, no reward model in the training loop, no value network, and generation can run on separate hardware with a standard inference engine. Verification cost depends on the verifier: string matching is free, sandboxed unit tests are not.

Relative to other methods, an RFT iteration costs roughly one SFT run plus one bulk-inference job. Havrilla et al. found expert iteration's sample complexity similar to PPO's (on the order of $10^6$ samples to converge from a pretrained checkpoint on their tasks), so the saving is in engineering and peak memory rather than in samples.

Failure modes and how to detect them#

Right answer, wrong reasoning

With answer-only verification, some kept samples reach the correct answer through flawed reasoning or lucky guesses (multiple-choice questions are the worst case: a random guess passes a quarter of the time). Training on them teaches the flaws. Detect by sampling kept examples for manual or model-graded review; mitigate with stricter verifiers, process supervision (process reward models), or by requiring consistency across samples.

Easy-problem domination

Without caps, the kept set over-represents what the model already does well, as the toy shows. Detect by plotting the kept set's distribution over solve rate or difficulty. Cap and deduplicate.

Plateau on hard problems

Problems the model never solves are never learned. Held-out accuracy stalls while accuracy on the subset with nonzero solve rate keeps rising. Remedies: rationalisation with hints, a stronger model's solutions (which turns RFT into distillation), larger $N$, or a curriculum.

Diversity collapse

Each iteration sharpens the distribution around solutions the model already finds. Pass@1 rises while pass@k flattens or falls and solutions become stereotyped. Track pass@k and the number of distinct correct solutions per problem; cap per prompt and keep temperature up.

Reward hacking with learned rewards

Best-of-$K$ against a reward model selects whatever the reward model over-scores (length, flattery, formatting), and iterating amplifies it. Monitor response length and an independent judge across iterations.

Regression on untrained skills

Repeated narrow fine-tuning erodes unrelated abilities, as Llama 2's poetry regression showed. Train from the base each iteration, keep samples from earlier iterations, and mix in general data.

When to use it, and what to use instead#

Use RFT when you have a verifier or a trustworthy reward model and want the simplest possible way to improve on it: as a first RL-like stage after SFT, to generate reasoning data for a later GRPO run, when you cannot run an online RL loop, or when you want each iteration's training data to be inspectable before it is used. It is also the natural method when the base model solves a reasonable fraction of problems some of the time. It is weaker when the model rarely succeeds (no signal), when negative examples carry important information (it ignores them), and when the reward is a learned model that is easy to exploit.

Minimal implementation#

A complete iteration with an open-source inference engine for sampling and any SFT trainer for the update. The verifier here is a GSM8K-style final-answer check. In POFE, RFT is a recipe that runs these phases as separate jobs, so the kept set can be inspected between them.

import re, random
from vllm import LLM, SamplingParams

def final_answer(text):
    m = re.findall(r"####\s*(-?[\d,\.]+)", text)
    return m[-1].replace(",", "") if m else None

def rft_iteration(model_path, base_path, problems, n=16, cap=2, temperature=0.8):
    llm = LLM(model=model_path)
    params = SamplingParams(n=n, temperature=temperature, top_p=1.0, max_tokens=1024)
    outs = llm.generate([p["prompt"] for p in problems], params)
    kept, solve_rates = [], []
    for p, out in zip(problems, outs):
        texts = [o.text for o in out.outputs]
        good = [t for t in texts if final_answer(t) == p["answer"]]
        solve_rates.append(len(good) / n)
        uniq = list(dict.fromkeys(t.strip() for t in good))    # dedup exact duplicates
        random.shuffle(uniq)
        kept += [{"messages": [{"role": "user", "content": p["prompt"]},
                               {"role": "assistant", "content": t}]} for t in uniq[:cap]]
    del llm
    print(f"kept {len(kept)} | prompts solved {sum(r > 0 for r in solve_rates) / len(problems):.1%}")
    # standard SFT, starting from the base checkpoint each iteration
    new_path = sft_train(start=base_path, data=kept, epochs=2, lr=1e-5)
    return new_path, solve_rates

model = BASE
for k in range(3):
    model, rates = rft_iteration(model, BASE, train_problems)
    print(k, evaluate(model, heldout_problems))

Any SFT implementation works for sft_train; see the SFT article for the masking and normalisation details. For best-of-$K$ with a reward model, replace the answer check with scoring every sample and keeping the argmax per prompt, and log the reward distribution of winners at every iteration.

Further reading#

Anthony, Tian and Barber (2017) for expert iteration in its original setting, and Polu et al. (2022) for its application to theorem proving. Zelikman et al. (2022) for STaR and rationalisation. Yuan et al. (2023) for RFT as data augmentation and the importance of distinct reasoning paths. Gulcehre et al. (2023) and Singh et al. (2023) for ReST and the EM derivation. Dong et al. (2023) for RAFT, and Touvron et al. (2023) for rejection sampling inside a production RLHF pipeline. Havrilla et al. (2024) for the head-to-head with PPO. For the online alternatives, see GRPO and RLVR.

Frequently asked questions#

References#

1. Thomas Anthony, Zheng Tian, David Barber (2017). Thinking Fast and Slow with Deep Learning and Tree Search. arxiv.org/abs/1705.08439
2. Eric Zelikman, Yuhuai Wu, Jesse Mu, Noah D. Goodman (2022). STaR: Bootstrapping Reasoning With Reasoning. arxiv.org/abs/2203.14465
3. Stanislas Polu, Jesse Michael Han, Kunhao Zheng, Mantas Baksys, Igor Babuschkin, Ilya Sutskever (2022). Formal Mathematics Statement Curriculum Learning. arxiv.org/abs/2202.01344
4. Zheng Yuan, Hongyi Yuan, Chengpeng Li, Guanting Dong, et al. (2023). Scaling Relationship on Learning Mathematical Reasoning with Large Language Models. arxiv.org/abs/2308.01825
5. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arxiv.org/abs/2307.09288
6. Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, et al. (2023). Reinforced Self-Training (ReST) for Language Modeling. arxiv.org/abs/2308.08998
7. Avi Singh, John D. Co-Reyes, Rishabh Agarwal, Ankesh Anand, et al. (2023). Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models. arxiv.org/abs/2312.06585
8. Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, et al. (2023). RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment. arxiv.org/abs/2304.06767
9. Alex Havrilla, Yuqing Du, Sharath Chandra Raparthy, Christoforos Nalmpantis, et al. (2024). Teaching Large Language Models to Reason with Reinforcement Learning. arxiv.org/abs/2403.04642
10. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, et al. (2021). Training Verifiers to Solve Math Word Problems. arxiv.org/abs/2110.14168

Related reading