Rejection Sampling Fine-Tuning: How Filtering Model Outputs by Reward Optimizes Alignment Without Policy Gradient Instability

Post-training alignment has become a defining phase in modern large language model development. While supervised fine-tuning (SFT) teaches a model to follow instructions and adopt structured formats, aligning model behavior with human preferences, safety criteria, and domain accuracy requires optimizing against reward signals.

Historically, this optimization has been approached through reinforcement learning from human feedback (RLHF) using policy gradient algorithms such as Proximal Policy Optimization (PPO). However, online policy gradient methods introduce severe systems complexity: actor-critic memory contention, generalized advantage estimation (GAE) variance, policy-value synchronization bottlenecks, and hyperparameter instability.

Rejection Sampling Fine-Tuning (often referred to as RFT, RAFT, or Reinforced Self-Training) provides a mathematically grounded and computationally simpler alternative. Instead of running policy gradient updates in the critical path of generation, Rejection Sampling Fine-Tuning decouples trajectory exploration from policy optimization. By sampling multiple candidates per prompt, scoring them with a reward model or automated verifier, and updating the model via standard supervised loss on the top-performing completions, rejection sampling directly approximates the theoretical optimal policy under KL regularization.

Mathematical Foundations: The Optimal KL-Constrained Policy

To understand why rejection sampling works, consider the standard reinforcement learning objective used in post-training alignment. Given a prompt distribution D, a reference policy $\pi_{\text{ref}}$ , a reward function $r(x, y)$ , and a temperature parameter $\tau > 0$ that scales the Kullback-Leibler (KL) divergence penalty, the goal is to find a parameterized policy $\pi_\theta(y|x)$ that maximizes expected reward while remaining close to the reference model:

$\max_{\pi} \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi(\cdot|x)} [r(x, y)] - \tau D_{\text{KL}}(\pi(\cdot|x) \parallel \pi_{\text{ref}}(\cdot|x))$

Expanding the KL divergence term reveals the objective in unconstrained variational form:

$\max_{\pi} \sum_{y} \pi(y|x) r(x, y) - \tau \sum_{y} \pi(y|x) \log \left( \frac{\pi(y|x)}{\pi_{\text{ref}}(y|x)} \right)$

Subject to the probability simplex constraint $\sum_y \pi(y|x) = 1$ . Solving this optimization problem analytically using Lagrange multipliers yields a closed-form expression for the optimal policy $\pi^*(y|x)$ :

$\pi^*(y|x) = \frac{1}{Z(x)} \pi_{\text{ref}}(y|x) \exp\left( \frac{r(x, y)}{\tau} \right)$

where $Z(x)$ is the sequence-level partition function normalizing the probability distribution:

$Z(x) = \sum_{y} \pi_{\text{ref}}(y|x) \exp\left( \frac{r(x, y)}{\tau} \right)$

This analytical result demonstrates that the optimal aligned policy is a Gibbs (or Boltzmann) distribution that re-weights the base reference policy by an exponential factor proportional to the reward score.

Directly computing $\pi^*(y|x)$ is intractable for modern vocabularies and multi-token sequences because calculating the partition function $Z(x)$ requires summing over an exponentially large sequence space. Rejection sampling and Best-of-N selection serve as non-parametric Monte Carlo approximations to draw samples directly from this target distribution.

Best-of-N Sampling as Monte Carlo Approximation

In practice, sampling from $\pi^*(y|x)$ can be approximated through Best-of-N (BoN) rejection sampling. For a given input prompt $x$ :

Generate $N$ independent candidate completions from the reference or current policy:

$y_1, y_2, \dots, y_N \sim \pi_{\text{ref}}(\cdot|x)$

Evaluate each candidate trajectory using the learned reward model or automated oracle:

$s_i = r(x, y_i) \quad \text{for } i \in \{1, \dots, N\}$

Select the candidate completion that achieves the maximum reward:

$y^* = \arg\max_{i \in \{1, \dots, N\}} r(x, y_i)$

As analyzed by Gao et al. (2023) and Liu et al. (2023), the distribution of the selected best candidate shifts probability mass toward higher-reward regions. The expected KL divergence between the Best-of-N distribution and the sampling policy scales logarithmically with the number of generated candidates:

$D_{\text{KL}}(\pi_{\text{BoN}} \parallel \pi_{\text{ref}}) \approx \log N - \frac{N-1}{N}$

For small to moderate values of $N$ , Best-of-N sampling produces high-quality completions without updating model parameters. However, evaluating $N$ rollouts at inference time multiplies serving latency and compute cost by a factor of $N$ . Rejection Sampling Fine-Tuning resolves this inference bottleneck by distilling the capabilities of the Best-of-N sampling procedure back into the model weights.

The Rejection Sampling Fine-Tuning Workflow

Rejection Sampling Fine-Tuning converts the sample selection process into an offline or iterative supervised training pipeline. Formally introduced across frameworks like Reward rAnked FineTuning (RAFT) by Dong et al. (2023) and Reinforced Self-Training (ReST) by Gulcehre et al. (2023), the workflow proceeds through four phases:

[Prompt Dataset D]
        │
        ▼
[Policy Sampling: K rollouts per prompt via π_θ]
        │
        ▼
[Reward Evaluation: r_φ(x, y_k) or Verification Oracle]
        │
        ▼
[Filtering & Ranking: Top-1 or Top-P Selection]
        │
        ▼
[Curated Dataset D_RS = {(x, y*)}]
        │
        ▼
[Supervised Fine-Tuning: ∇_θ L_SFT(θ)]
        │
        ▼
[Updated Aligned Policy π_θ+]

1. Generation Phase

For each prompt $x$ in a training dataset D, the generator policy $\pi_\theta$ produces $K$ distinct completions $\{y_1, y_2, \dots, y_K\}$ using stochastic decoding (e.g., temperature $T \in [0.7, 1.0]$ or top- $p$ nucleus sampling). The number of candidate rollouts $K$ typically ranges from 4 to 64 depending on prompt complexity and compute budget.

2. Reward Scoring and Verification

Each completion is evaluated by a scoring function. Depending on the task domain, this can take several forms:

Learned Reward Models: A scalar preference model trained on human comparison pairs via Bradley-Terry loss.
Rule-Based or Programmatic Oracles: Exact unit test suites, compiler outputs, or formal theorem verifiers used in coding and mathematical reasoning tasks.
Process Reward Models (PRMs): Step-level verification models that score intermediate reasoning steps.

3. Filtering and Selection Strategy

Rather than selecting only the single top-scoring response, modern RFT frameworks employ flexible selection mechanisms:

Top-M Selection: Retaining the top $M$ completions ( $1 \le M < K$ ) for each prompt to increase training sample diversity.
Reward Thresholding: Retaining only samples whose score exceeds a minimum baseline, discarding prompts where all generated rollouts failed.
Importance Re-weighting: Assigning sample weights proportional to $\exp(r(x, y_i)/\tau)$ during gradient updates.

4. Supervised Optimization

The filtered pairs are collected into a curated training set $\mathcal{D}_{\text{RS}} = \{(x, y^*)\}$ . The policy parameters $\theta$ are updated using standard autoregressive cross-entropy loss:

$\mathcal{L}_{\text{RFT}}(\theta) = -\mathbb{E}_{(x, y^<em>) \sim \mathcal{D}_{\text{RS}}} \left[ \sum_{t=1}^{|y^</em>|} \log \pi_\theta(y_t^* \mid x, y_{<t}^*) \right]$

Because the loss function is identical to standard supervised fine-tuning, training leverages standard SFT infrastructure, avoiding the memory-heavy multi-model scheduling required by PPO.

Iterative Rejection Sampling

A single pass of rejection sampling is bounded by the initial policy's generation capabilities. If the base model assigns near-zero probability to the optimal trajectory, sampling $K$ candidates will rarely surface high-reward responses.

To overcome this limitation, models like Llama 2 Chat employ iterative rejection sampling fine-tuning. The pipeline is executed in discrete successive rounds:

Round t: Policy $\pi_{\theta_t}$ generates candidates over dataset D.
Filtering: The reward model scores all candidates; top samples form dataset $\mathcal{D}_t$ .
Training: Policy $\pi_{\theta_{t+1}}$ is trained on $\mathcal{D}_t \cup \mathcal{D}_{\text{gold}}$ .
Iteration: The updated policy $\pi_{\theta_{t+1}}$ serves as the generator for Round $t+1$ .

π_θ0 ──> Sample K ──> Score & Filter ──> SFT ──> π_θ1 ──> Sample K ──> Score & Filter ──> SFT ──> π_θ2

Across multiple iterations, the model's base probability distribution shifts toward high-reward regions. As the generator improves, the quality of its $K$ sampled candidates rises, allowing the reward model to surface increasingly refined trajectories. In the Llama 2 alignment pipeline, Touvron et al. (2023) demonstrated that iterative rejection sampling prior to PPO significantly boosted helpfulness and safety win-rates compared to direct RLHF alone.

Statistical Rejection Sampling Optimization (RSO)

While standard RFT trains purely on high-reward positive samples via supervised loss, preference optimization methods like Direct Preference Optimization (DPO) benefit from contrasting chosen and rejected responses.

To bridge this gap, Liu et al. (2023) introduced Statistical Rejection Sampling Optimization (RSO). Standard DPO assumes that the preference dataset was sampled from the reference policy $\pi_{\text{ref}}$ . However, static human preference datasets are often collected from third-party models or earlier checkpoints, creating out-of-distribution support mismatch.

RSO resolves this discrepancy by using rejection sampling directly against the target distribution $\pi^*(y|x)$ :

Draw candidate responses from the reference policy $\pi_{\text{ref}}(y|x)$ .
Use the reward model $r_\phi(x, y)$ to construct pairwise preference tuples $(y_w, y_l)$ sampled proportionally to the target Boltzmann distribution:

$p(y_w \succ y_l \mid x) = \sigma\left(\frac{r_\phi(x, y_w) - r_\phi(x, y_l)}{\tau}\right)$

Optimize the policy using the exact DPO objective evaluated on the rejection-sampled pairs:

$\mathcal{L}_{\text{RSO}}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}_{\text{RSO}}} \left[ \log \sigma \left( \beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)} \right) \right]$

By aligning the support of the preference distribution with the actual reference policy, RSO eliminates distribution mismatch and improves alignment over standard DPO.

Architectural Comparison: RFT vs. PPO vs. DPO vs. GRPO

Rejection Sampling Fine-Tuning (RFT / RAFT): Uses standard supervised cross-entropy on top-filtered samples. Requires only a single actor model during training with baseline SFT memory footprint. Needs no critic or value network. Offers high numerical stability and allows arbitrary offline verifiers, but requires generating candidate pools offline.
Proximal Policy Optimization (PPO): Uses online policy gradients with clipped surrogate objectives. Requires coordinating four models simultaneously (Actor, Critic, Reference, and Reward), demanding 3x to 4x baseline VRAM. High tuning complexity and sensitive to reward non-stationarity.
Direct Preference Optimization (DPO): Optimizes closed-form implicit preference losses on static pairwise datasets. Requires two models (Actor and Reference). Avoids online generation and critics, but can suffer from reference model drift and out-of-distribution dataset mismatch.
Group Relative Policy Optimization (GRPO): Employs online group-normalized policy gradients without a value model. Evaluates groups of completions synchronously during online rollouts, balancing policy gradient exploration with lower VRAM overhead than PPO.

Practical Engineering Advantages

1. Decoupled Compute Architecture

In online RL (PPO), generation, reward scoring, advantage estimation, and gradient updates must execute in a tight synchronous loop. If a reward model is large, or if verification involves executing code inside secure sandboxes, GPUs remain idle waiting for evaluation.

Rejection sampling decouples these phases. Generation and verification can run as asynchronous batch jobs across low-cost spot instances or CPU clusters. Once the dataset is materialized on disk, model training proceeds as standard high-throughput supervised fine-tuning.

2. Complex and Ensemble Oracles

Because evaluation happens offline during dataset preparation, practitioners can deploy sophisticated reward mechanisms that would be intractable inside an online RL training step:

Multi-model judge ensembles (such as majority voting across multiple independent reward models).
Multi-turn execution environments (running automated unit test suites, linters, static analyzers, and symbolic provers).
Human-in-the-loop review queues for borderline confidence scores.

3. Training Stability

Online policy gradient methods often suffer from high gradient variance, leading to catastrophic forgetting or degenerate repetitive text. Rejection Sampling Fine-Tuning optimizes exact log-likelihoods over valid token sequences. The resulting gradient pushes probability mass directly onto high-reward trajectories without the numerical instability of online importance sampling ratios.

Limitations and Failure Modes

Despite its simplicity and empirical success, Rejection Sampling Fine-Tuning exhibits several distinct failure modes that require careful mitigation:

1. Reward Model Overoptimization (Goodhart's Law)

As demonstrated by Gao et al. (2023), maximizing a proxy reward model causes true output quality to improve initially, reach a peak, and subsequently degrade as the policy exploits flaws in the reward model:

True Quality
     │       Peak Alignment
     │          ▲
     │         / \
     │        /   \   Reward Hacking Regime
     │       /     \  (Length bias, stylistic filler)
     │      /       \
     └─────/─────────\───────────► KL Divergence / Sample Budget (N)

Because Best-of-N selection aggressively filters for maximum proxy reward, large candidate pools (such as $K > 128$ ) frequently select responses that exploit length bias, verbosity hacks, or adversarial token sequences. Maintaining moderate sampling budgets and incorporating rule-based length penalties prevents proxy gaming.

2. Mode Collapse and Diversity Loss

Training repeatedly on top-1 outputs can cause the policy to collapse into narrow linguistic patterns. If all filtered responses for a given prompt share identical phrasing, the model loses entropy and diversity across downstream tasks. Selecting top- $M$ completions ( $M \ge 2$ ) or applying nucleus sampling during generation preserves output entropy.

3. Exploration Ceilings

Rejection sampling cannot generate solutions that lie entirely outside the base policy's support. If the generator has a 0% baseline pass rate on a complex reasoning task, drawing $K=16$ or $K=64$ samples will yield zero correct rollouts, leaving the dataset empty for that prompt. Solving hard exploration regimes requires combining rejection sampling with prompt mutation, tree search, or bootstrapping from stronger teacher models.

Implementation Summary

Rejection Sampling Fine-Tuning bridges the gap between pure supervised imitation learning and complex reinforcement learning. By framing alignment as the distillation of a reward-filtered Boltzmann distribution, RFT allows engineering teams to achieve frontier alignment performance using standard SFT infrastructure.

When designing a production post-training pipeline:

Generate between 8 and 32 candidate rollouts per prompt using moderate decoding temperature ( $T \approx 0.8$ ).
Evaluate candidates with strict outcome verifiers or well-calibrated reward models, discarding prompts where no response meets quality thresholds.
Apply standard autoregressive supervised fine-tuning on the filtered pool.
Iterate across 2 to 4 rounds to progressively elevate the policy's generation capabilities.

Sources

Dong, H., Xiong, W., Goyal, D., Pan, R., Diao, S., Zhang, J., Shum, K., & Zhang, T. (2023). RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment. Transactions on Machine Learning Research. https://arxiv.org/abs/2304.06767
Liu, T., Zhao, Y., Joshi, R., Khalman, M., Saleh, M., Liu, P. J., & Liu, J. (2023). Statistical Rejection Sampling Improves Preference Optimization. International Conference on Learning Representations (ICLR 2024). https://arxiv.org/abs/2309.06657
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. Meta AI. https://arxiv.org/abs/2307.09288
Gulcehre, C., Le Paine, T., Srinivasan, S., Krivokhizha, A., Sheppard, B., Ahern, M., et al. (2023). Reinforced Self-Training (ReST) for Language Modeling. DeepMind. https://arxiv.org/abs/2308.08998
Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C. D., & Finn, C. (2023). Direct Preference Optimization: Your Language Model is Secretly a Reward Model. NeurIPS 2023. https://arxiv.org/abs/2305.18290
Gao, L., Schulman, J., & Hilton, J. (2023). Scaling Laws for Reward Model Overoptimization. International Conference on Machine Learning (ICML 2023). https://arxiv.org/abs/2210.10760
Nakano, R., Hilton, J., Balaji, S., Wu, J., Ouyang, L., Kim, C., et al. (2021). WebGPT: Browser-assisted question-answering with human feedback. OpenAI. https://arxiv.org/abs/2112.09332

Rejection Sampling Fine-Tuning in Large Language Models: How Best-of-N Filtering, Reward Oracles, and Distillation Align Neural Policies