Rejection Sampling Fine-Tuning: How Filtering Model Outputs by Reward Replaced RLHF Complexity

Rejection sampling fine-tuning (RAFT) has emerged as the practical workhorse of LLM alignment. While PPO-based RLHF dominated early literature, production systems from Llama 2 to DeepSeek-R1 rely on a simpler loop: generate multiple completions per prompt, score them with a reward model, keep the best, and fine-tune on the filtered data. The technique converts the reinforcement learning problem into supervised learning — stable, memory-efficient, and embarrassingly parallel.

From Computational Statistics to LLM Alignment

The name originates in statistics: when you cannot sample directly from a target distribution, you sample from an easier proposal distribution and accept or reject each sample based on a weighting function. In LLMs, the target distribution is "high-quality completions," the proposal distribution is the current model, and the acceptance criterion is a reward model score.

WebGPT (Nakano et al., 2021), Anthropic's Helpful and Harmless work (Bai et al., 2022), and OpenAI's process reward models (Lightman et al., 2023) all used variants. The method was formalized as RAFT (Reward rAnked FineTuning) by Dong et al. (2023) and independently deployed at scale in Llama 2 (Touvron et al., 2023).

The Three-Step Iteration

At each iteration, RAFT executes:

1. Data collection — Sample K completions per prompt from the current policy using temperature-controlled sampling (typically 0.7–1.0, with K = 10–30). This balances diversity against quality.
2. Reward-based ranking — Pass every (prompt, completion) pair through a trained reward model R. Select the highest-scoring completion per prompt (argmax per row) or the top K pairs globally.
3. Supervised fine-tuning — Run standard SFT on the selected pairs using the same cross-entropy loss as instruction tuning.

Mathematically, given prompts X = [x₁, ..., xₘ] and N completions per prompt forming matrix Y, the reward model produces matrix R where rᵢⱼ = R(yᵢⱼ | xᵢ). The per-prompt selection function S(R) = [argmaxⱼ r₁ⱼ, ..., argmaxⱼ rₘⱼ] yields chosen completions Y_chosen. The model then minimizes −Σ log π_θ(y_chosen | x).

Llama 2 applied this only to their 70B model, then distilled the filtered data to smaller sizes — a pattern RAFT later formalized as a distillation pathway.

Why It Works: Stability Over Sophistication

PPO-based RLHF requires four models in memory simultaneously (policy, reference, critic, reward), careful KL scheduling, and gradient updates that can destabilize. RAFT needs only one model loaded at a time. The learning curves show consistent reward improvement without the spikes and collapses typical of policy gradients.

The "alignment tax" — where RLHF degrades fluency or diversity — is largely absent. RAFT maintains or improves perplexity while increasing reward scores, because the SFT objective directly maximizes likelihood of high-reward tokens rather than indirectly via policy gradients.

Dong et al. demonstrated this on LLaMA-7B with HH-RLHF: RAFT exceeded PPO on both reward model score and human/GPT-4 preference, with 50× less compute on diffusion model alignment (8.4 minutes vs 415 minutes for DDPO on Stable Diffusion 256×256).

Best-of-N Inference vs. Training-Time Filtering

Best-of-N (BoN) uses the same generate-and-score pipeline at inference time: sample N completions, return the highest-reward one. RAFT instead bakes the selection into the weights. BoN spends compute per query; RAFT spends compute once during training. They are complementary — BoN can further boost a RAFT-aligned model at test time.

Practical Hyperparameters

The RLHF Book (Lambert, 2025) documents settings that work in practice:

• Temperature: 0.7–1.0 (higher early in training, lower later)
• Completions per prompt: 10–30+ (fewer makes selection noisy; more yields diminishing returns)
• Selection strategy: Argmax per prompt is standard; top-K overall can concentrate on "easy" prompts
• Deduplication: Often applied before scoring to avoid wasting reward model calls on near-duplicate outputs
• Iterative scheduling: Re-generate from the updated checkpoint each round; Llama 2 accumulated top samples across all prior iterations

Extensions and Variants

KL regularization — Add a penalty β · KL(π_θ || π_ref) to the SFT loss to prevent catastrophic drift from the base model's capabilities.

Knowledge distillation — The decoupled generation/training design lets a small student learn from a large teacher's filtered outputs. Dong et al. showed GPT-Neo-2.7B improving significantly when trained on RAFT-aligned LLaMA-7B samples.

Statistical Rejection Sampling Optimization (RSO) — Liu et al. (2024) provide a principled analysis connecting rejection sampling to DPO and IPO objectives, showing it implicitly optimizes a variant of the pairwise preference loss.

Cross-Modal Generality

RAFT's only modality-specific component is the reward model. The same loop aligns diffusion models: generate K images per prompt, score with a text-image alignment reward (e.g., CLIP or a trained aesthetic predictor), fine-tune on the best. This generality makes it a default alignment tool for any generative foundation model.

When to Use It

• You have a reward model (or verifiable reward function for reasoning tasks)
• You want alignment without PPO's infrastructure burden
• You need to align models larger than your training GPU memory allows for multi-model RLHF
• You want interpretable, inspectable training data — the filtered dataset is the alignment signal

Sources

• Dong et al., "RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment," arXiv:2304.06767 (2023), published in TMLR. https://arxiv.org/abs/2304.06767
• Touvron et al., "Llama 2: Open Foundation and Fine-Tuned Chat Models," arXiv:2307.09288 (2023). https://arxiv.org/abs/2307.09288
• Nakano et al., "WebGPT: Browser-Assisted Question-Answering with Human Feedback," arXiv:2112.09332 (2021). https://arxiv.org/abs/2112.09332
• Bai et al., "Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback," arXiv:2204.05862 (2022). https://arxiv.org/abs/2204.05862
• Lightman et al., "Let's Verify Step by Step," ICLR 2024, arXiv:2305.20050. https://arxiv.org/abs/2305.20050
• Lambert, "Rejection Sampling," Chapter 9 in Reinforcement Learning from Human Feedback, 2025. https://rlhfbook.com/c/09-rejection-sampling
• Liu et al., "Statistical Rejection Sampling Improves Preference Optimization," ICLR 2024. https://arxiv.org/abs/2309.06657
• Yuan et al., "Rejection Sampling Fine-Tuning," referenced in Dong et al. (2023) and subsequent surveys.