RAFT: replace PPO's four-model setup with best-of-K sampling plus fine-tuning, and win

RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment

Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, Tong Zhang

cs.LG, cs.AI, cs.CL, cs.CV, stat.ML

2023-04-14

RAFT replaces PPO for alignment: sample K responses per prompt, keep the highest-reward, fine-tune, repeat. LLaMA-7B reward 2.294 beats PPO's 2.077, loading one model instead of four.

What problem this solves

Aligning large models with human feedback (RLHF) is close to an industry standard, but the machinery is heavy. The dominant algorithm, PPO (Proximal Policy Optimization), learns by trial and error: unstable and inefficient. The bigger problem is memory. PPO loads four models at once: the policy being trained, a frozen reference policy, a reward model, and a value critic. Each takes a share, so most of a GPU is gone. On top of that, an imperfect reward model gets gamed for high score (reward hacking), and optimizing reward alone degrades output quality (alignment tax).

The authors want to sidestep all of this. The intuition: best-of-K, sampling K responses per prompt and keeping the highest-reward one, already approaches RLHF quality at inference cost with no training. Can those filtered good samples be turned back into a supervised signal and iterated?

Method

RAFT (Reward rAnked FineTuning) is three steps, repeated until reward converges:

Sample: from the current model, draw K responses per prompt.
Rank: the reward model scores all K; keep the top one per prompt.
Fine-tune: ordinary supervised fine-tuning on this best-of-K subset.

The updated model loops back to step one. The key design choice is decoupling generation from training: ranking uses the reward model, fine-tuning loads only the policy, so peak memory no longer holds four models at once. That is also why the same loop drops onto diffusion models; anything with a scoring reward function works.

Knobs: larger K biases toward high reward (at sampling cost); higher temperature λ means more diverse generation, though too high produces gibberish and narrows usable λ; an optional KL penalty coefficient β reins the model toward the initial distribution to prevent over-optimization.

The authors give a theoretical bound on best-of-K: E[max r] ≤ E[r] + B√(log K / 2), meaning the marginal gain from K shrinks fast. That is why iterated sample-rank-train beats one-shot large K.

Results

Main experiment: LLaMA-7B on Anthropic's HH-RLHF human-preference dialog data.

Model	Reward	Perplexity	Avg length
LLaMA-7B (base)	-0.435	4.781	119.9
LLaMA-7B-SFT	0.772	3.781	145.4
LLaMA-7B + PPO	2.077	4.156	127.8
RAFT-K32-λ1.0	2.294	4.031	156.2

RAFT takes the highest reward, 2.294, beating PPO's 2.077, with perplexity in a sane range. In GPT-4 and human pairwise comparisons, RAFT-K32 mostly beats PPO: against PPO-β0.1, GPT-4 scores 65 wins to 32 losses and humans 66 to 14.

The K ablation: reward climbs from 2.180 at K=8 to 2.329 at K=32, confirming "larger K, higher reward." K=32 converges in about 10-12 iterations versus 15-18 for K=8. On compute, K in {8, 16, 32} takes 5, 6.05, and 7.05 hours respectively, against roughly 8.7 hours for the fastest PPO config. RAFT is faster.

The diffusion result is more lopsided. On Stable Diffusion v1.5 at 256x256:

Metric	Pretrained	DDPO	RAFT
In-domain aesthetic	4.63	6.04	6.14
Out-of-domain aesthetic	4.64	5.76	6.07
Training time (single A40)	N/A	415 min	8.4 min

RAFT matches or beats DDPO on aesthetic score and trains about 50x faster. It also distills: LLaMA-7B as teacher supplies samples, GPT-Neo-2.7B as student lifts reward from -1.23 to 0.739.

Why it matters

RAFT gives alignment a cheaper engineering path. Skip PPO's critic-plus-reference-plus-reward co-residency; the sample-rank-fine-tune loop matches or beats RLHF. For teams short on memory who want fast iteration, it lands directly. For generative models with no off-the-shelf RL training stack, like diffusion, RAFT is close to plug-and-play.

Seen from today, this 2023 paper being re-cited by recent RLVR (reinforcement learning from verifiable rewards) work is no accident. "Best-of-K sampling during training" is the direct ancestor of the drop-the-critic, estimate-baseline-from-samples idea behind methods like GRPO. Reading RAFT clarifies where the current push to simplify RL training comes from.

Limitations

Results are sensitive to K; pick it wrong and reward drops, but the paper gives no clean rule for choosing K.
High temperature λ produces gibberish tokens, narrowing usable λ and pitting diversity against usability.
The KL penalty needs extra forward passes, which the authors concede eats part of the savings.
Reward-model calibration error and classification noise affect RAFT as much as PPO.
Only one dialog dataset (HH-RLHF) and one diffusion task; generalization is narrow, and LLaMA-7B is small by today's standards. Scaling is unverified.
A cold-start risk the paper does not address: best-of-K assumes you can sample a good response. If the current model is far from target and none of the K samples is any good, the loop can stall on bad samples.

Terms

RLHF (Reinforcement Learning from Human Feedback): train a reward model on human pairwise labels, then fine-tune the generator toward it.
PPO (Proximal Policy Optimization): the RL algorithm most RLHF pipelines use; stable but needs policy, reference, reward, and critic models in memory.
best-of-K: sample K candidates and pick the best by some scorer; cheap inference-time improvement.
Alignment tax: the quality or diversity loss from optimizing purely for reward.
Reward hacking: the model exploiting flaws in the reward model to score high without actually being good.
GRPO / RLVR: later methods that drop the critic and estimate baselines from group samples; RAFT is a conceptual forerunner.