Rejection sampling is presented as a practical but underdocumented technique in preference fine-tuning. The core idea is simple: generate several candidate completions for a prompt, score them with a reward model, keep the best ones, and then fine-tune the original model on those selected outputs using standard instruction-tuning loss. The chapter notes that this method appears in many influential RLHF systems because it can be inserted at different points in a training pipeline, including after supervised fine-tuning or after reinforcement learning, even though its exact role and best practices are not always clearly specified.
Conceptually, the method is compared to rejection sampling in statistics, where samples from an easier-to-draw distribution are filtered to approximate a harder target distribution. In the language-model setting, the current model produces candidate completions, the reward model acts as the acceptance filter, and the goal is to concentrate training on higher-quality prompt-completion pairs. The chapter emphasizes that the procedure can be implemented in different ways, such as choosing the best completion for each prompt or selecting the top overall prompt-completion pairs across the full batch, with the exact choice depending on the training objective and experimental design.
The chapter also gives implementation guidance and related context. It explains the need to choose prompts, generate multiple completions per prompt, score them with a reward model, and then retrain on the selected outputs. Practical considerations include using enough sampled completions, tuning sampling settings like temperature and top-p, possibly mixing generations from different models, and sorting sequences by length for efficient reward evaluation. Finally, it distinguishes rejection sampling from best-of-N sampling: both use the same generate-and-rank idea, but best-of-N is only an inference-time selection method and does not update the model, whereas rejection sampling uses the filtered completions for further training.
Rejection sampling overview.
Summary
Rejection sampling is a simple alternative to RL for preference fine-tuning. Generate multiple completions per prompt, score them with a reward model, and instruction fine-tune the model on only the top-scoring completions.
Two main selection strategies exist: picking the best completion per prompt (preserving prompt diversity) or picking the top \(K\) completions globally (maximizing overall quality), with the choice depending on the training goals.
Best-of-N (BoN) sampling follows the same generate-and-score procedure but does not fine-tune the model, serving instead as an inference-time technique and a common baseline for comparing against RL training methods.
FAQ
What is rejection sampling in the context of RLHF?Rejection sampling is a generate-and-filter training method where a model produces multiple candidate completions for each prompt, a reward model scores them, and only the highest-scoring completions are kept for further fine-tuning. In this chapter, the selected completions are then used with the same supervised fine-tuning loss as instruction tuning.Why is rejection sampling used instead of reinforcement learning?It can be used as a simpler alternative to RL because it avoids the complexity of online policy optimization while still improving model outputs using a reward model. It is often easier to implement and can be inserted at several stages of the RLHF pipeline.What are the main stages of the rejection sampling process?The chapter describes four stages: choose prompts and a reward model, generate multiple completions from a starting checkpoint, score all completions with the reward model, and fine-tune the model on the selected top completions.How are prompts and completions represented mathematically in rejection sampling?Prompts are represented as a vector \(X = [x_1, x_2, ..., x_M]\), and for each prompt, multiple completions are stored in a matrix \(Y\), where \(y_{i,j}\) is the \(j\)-th completion for the \(i\)-th prompt.How does the reward model score completions?Each prompt-completion pair is passed through a reward model \(\mathcal{R}\), producing a reward value \(r_{i,j} = \mathcal{R}(y_{i,j} \mid x_i)\). These scores are arranged into a reward matrix \(R\).What is the difference between selecting the top completion per prompt and selecting the top overall pairs?Top per prompt chooses the single best completion within each prompt row of the reward matrix, while top overall pairs flatten all prompt-completion rewards into one vector and select the highest-scoring \(K\) completions across all prompts.How do you map a flattened reward index back to the original prompt and completion?For a zero-indexed flattened index \(k\), the original prompt and completion can be recovered with \(i = \lfloor k / N \rfloor + 1\) and \(j = (k \bmod N) + 1\), where \(N\) is the number of completions per prompt.What are typical implementation settings for rejection sampling?Common settings include sampling with temperature above zero, often around 0.7 to 1.0, and sometimes using top-p or top-k sampling. Successful implementations often generate 10 to 30 or more completions per prompt.How is best-of-N sampling related to rejection sampling?Best-of-N uses the same generate-and-score idea, but it does not fine-tune the model on the selected outputs. Instead, it chooses the best completion at inference time, making it a sampling method rather than a training method.Why might rejection sampling be considered underdocumented or hard to standardize?The chapter notes that key choices such as which prompts to reuse, how to select the reward model, how many completions to sample, and how to sequence the process are not well documented in the literature. As a result, many practical details depend on empirical experimentation rather than a canonical recipe.
pro $24.99 per month
access to all Manning books, MEAPs, liveVideos, liveProjects, and audiobooks!
Reinforcement Learning from Human Feedback ebook for free
