Best-of-N Sampling

hardPROReward ModelingLLM 2: Training & Alignment

Implement Best-of-N (rejection) sampling for LLM generation.

Given N candidate responses with their reward scores, select the best one. Also compute the KL cost of this selection strategy.

The expected KL divergence cost of Best-of-N sampling is approximately: $\text{KL}_{\text{BoN}} \approx \log(N) - \frac{N-1}{N}$

Input:

Line 1: N (number of candidates)
Next N lines: reward_score response_text

Output:

Line 1: The best response text (highest reward)
Line 2: The best reward score, rounded to 4 decimal places
Line 3: The KL cost approximation, rounded to 4 decimal places

Example

Input

3
1.5 Hello world
2.3 Hi there
0.8 Hey

Output

Hi there
2.3000
0.4319

Reasoning

We read the input values: $N = 3$ , and the reward scores with their corresponding responses: $(1.5, \text{Hello world})$ , $(2.3, \text{Hi there})$ , and $(0.8, \text{Hey})$ .
We select the best response by finding the highest reward score, which is $2.3$ for the response $\text{Hi there}$ .
We calculate the KL cost approximation using the formula: $\text{KL}_{\text{BoN}} \approx \log(N) - \frac{N-1}{N} = \log(3) - \frac{3-1}{3} \approx 0.4319$ .
The final output is the best response text, the best reward score rounded to 4 decimal places ( $2.3000$ ), and the KL cost approximation rounded to 4 decimal places ( $0.4319$ ).

Constraints

1 <= N <= 100
If tied, pick the first one encountered
Round to 4 decimal places

Learn this in

Probability & StatisticsFoundations

More from LLM 2: Training & Alignment

RLHF Reward Shapinghard Implicit Reward from DPOmedium PPO Clipped Objectivemedium DPO Lossmedium KL Divergence Penaltyeasy

Concept primerPREVIEW

Background Knowledge

The Best-of-N sampling strategy is a technique used in reward modeling for selecting the best response from a set of candidate responses generated by a Large Language Model (LLM). This approach involves evaluating each candidate response based on a reward score, which represents the quality or relevance of the response. The goal is to choose the response with the highest reward score, as it is likely to be the most suitable or accurate response.

In the context of LLM training and alignment, Best-of-N sampling is used to improve the model's performance by selecting the best responses and using them as targets for further training. The KL divergence cost is a measure of the difference between the selected response distribution and the original distribution of candidate responses. It is used to evaluate the effectiveness of the selection strategy and to guide the optimization of the model's parameters. The KL divergence cost approximation provided in the problem description is a simplified formula that can be used to estimate the cost of the Best-of-N sampling strategy.

The KL divergence is a fundamental concept in information theory and machine learning, which measures the difference between two probability distributions. In the context of Best-of-N sampling, the KL divergence cost represents the loss of information that occurs when selecting the best response from a set of candidates. Understanding the KL divergence and its role in reward modeling is essential for developing effective selection strategies and optimizing LLM performance.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Reading the input data, including the number of candidates and their corresponding reward scores and response texts
Implementing the Best-of-N sampling strategy to select the best response based on the reward scores
Computing the KL cost approximation using the provided formula
Outputting the selected response text, the best reward score, and the KL cost approximation

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results