Chapter 9: Rejection sampling

Introduction

The introduction defines rejection sampling for preference fine-tuning and relates it to the statistical idea of sampling from a simpler distribution and filtering the results.

• Rejection sampling training workflow: Describe how rejection sampling curates candidate completions, filters them with a trained reward model, and fine-tunes the original model on only the top completions.
  Why it matters: This is the basic workflow that the chapter breaks down step by step.
• Statistical sampling analogy: Explain how the statistical idea behind rejection sampling maps to language models: the current model supplies samples, a reward model filters them, and high-quality prompt completions are the target.
  Why it matters: The analogy explains why the method generates and filters completions instead of drawing ideal answers directly.

Key terms: rejection sampling, candidate completions, reward model (RM), instruction tuning, target distribution, sampling distribution, current model

9.1 Training process, step by step

This section outlines the stages of rejection sampling training, from choosing prompts and a reward model to fine-tuning on selected completions.

• Rejection sampling training stages: Trace the rejection sampling process: select prompts and a trained reward model, generate completions with the starting checkpoint, rank them with the reward model, and instruction fine-tune the starting checkpoint on the selected completions.
  Why it matters: These stages show how reward-model rankings become training data for the starting model.
• Prompt selection tradeoff: Explain why reusing every prompt from the initial supervised or instruction fine-tuning stage is a simple option but may be undesirable.
  Why it matters: The easiest source of prompts is not necessarily the best choice for another training stage.

Key terms: rejection sampling, RM, starting checkpoint, completions, SFT/IFT, prompts, supervised fine-tuning, instruction fine-tuning, overfitting

9.1.1 Generating completions

This subsection defines how multiple candidate completions are generated and organized for a set of prompts.

• Prompt and completion indexing: Describe how M prompts and N generated completions per prompt are represented, including what x_i and y_{i,j} identify.
  Why it matters: The indexing makes clear which candidates belong to each prompt when they are scored later.

Key terms: prompts, completions, candidate completions, x_i, y_{i,j}

9.1.2 Scoring completions

This section shows how to score generated prompt–completion pairs and select completions either per prompt or across the entire batch.

• Reward matrix for completions: Explain how the reward matrix is formed from prompt–completion pairs and what an entry represents.
  Why it matters: The matrix supplies the scores used by both selection methods.
• Best completion per prompt: Describe how per-prompt selection uses the reward matrix to choose training completions.
  Why it matters: This method retains a highest-scoring completion for every prompt.
• Top pairs across prompts: Explain how top-overall selection chooses K pairs and maps their flattened reward indices back to prompts and completions. Distinguish its result from per-prompt selection.
  Why it matters: Selecting across all pairs can keep multiple completions for one prompt instead of guaranteeing one for each prompt.

Key terms: prompt–completion pairs, reward model, reward matrix, RM, top per prompt, max reward per prompt, argmax, top overall pairs, R_flat, top K

9.1.3 Fine-tuning

9.2 Implementation details

This section identifies practical rejection-sampling choices and a batching technique for efficient reward-model inference.

• Completion sampling and quantity: Explain how sampling settings and the number of completions per prompt affect rejection-sampling training.
  Why it matters: The quality and quantity of candidates determine what rejection sampling has available to select.
• Reward model influence: State how the reward model affects the outcome of rejection sampling.
  Why it matters: Selecting candidates with a reward model makes the result sensitive to that model.
• Length-sorted inference batches: Explain why tokenized completions can be sorted by length before batch reward-model inference.
  Why it matters: Length-based batching avoids spending reward-model computation on unnecessary padding.

Key terms: RS, temperature, top-p, top-k, completions per prompt, RM, batch RM inference, tokenized completions, padding tokens, throughput

9.3 Related: Best-of-N sampling

This section explains best-of-N sampling as inference-time selection without fine-tuning, its use as a comparison baseline, and the equivalence of two selection rules when K=1.

• Inference-Time Completion Selection: Distinguish best-of-N sampling from rejection sampling: both generate and score completions, but best-of-N selects an answer at inference time without fine-tuning the model.
  Why it matters: This distinction separates choosing a better answer now from training a model to produce better answers later.
• Sampling Baseline Comparisons: Explain why best-of-N can be compared with RLHF training methods despite not modifying the model, including that its KL distance relative to another policy can still be measured.
  Why it matters: Not changing model weights does not prevent comparison of the answers produced by different methods.
• Single-Prompt Selection Equivalence: Explain why, for one prompt with N scored completions, selecting the reward argmax is equivalent to top-K selection when K=1.
  Why it matters: The two selection rules are not different choices when only the single best completion is kept.

Key terms: Best-of-N (BoN), rejection sampling (RS), inference time, baseline, Kullback–Leibler (KL) distance, online training methods, reward vector, argmax, top-K, K=1

9.4 Summary