Nucleus (Top-P) Sampling

mediumPROGeneration & OptimizationLLM 3: Applications & Evaluation

Implement nucleus (top-P) sampling.

Top-P sampling selects the smallest set of tokens whose cumulative probability exceeds P:

Convert logits to probabilities (softmax)
Sort by probability descending
Find the smallest set where cumulative probability >= P
Zero out all other positions and renormalize

Input:

Line 1: P (nucleus threshold)
Line 2: space-separated logits

Output: The filtered probability distribution, rounded to 4 decimal places.

Example

Input

0.9
3.0 1.0 4.0 0.5

Output

0.2583 0.0000 0.7417 0.0000

Reasoning

First, we convert the logits to probabilities using the softmax function: $p_i = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}$ , resulting in probabilities for the given logits: $p_1 = \frac{e^{3.0}}{e^{3.0} + e^{1.0} + e^{4.0} + e^{0.5}}$ , $p_2 = \frac{e^{1.0}}{e^{3.0} + e^{1.0} + e^{4.0} + e^{0.5}}$ , $p_3 = \frac{e^{4.0}}{e^{3.0} + e^{1.0} + e^{4.0} + e^{0.5}}$ , $p_4 = \frac{e^{0.5}}{e^{3.0} + e^{1.0} + e^{4.0} + e^{0.5}}$ .
Then, we sort these probabilities in descending order and calculate their cumulative sum until it exceeds the given threshold $P = 0.9$ .
Next, we select the smallest set of tokens whose cumulative probability exceeds $P$ , which in this case are the first and third tokens ( $p_1$ and $p_3$ ), and zero out the other positions ( $p_2$ and $p_4$ ).
Finally, we renormalize the selected probabilities to ensure they sum up to 1, resulting in the filtered probability distribution: $p_1 = \frac{0.2583}{0.2583 + 0.7417} = 0.2583$ , $p_2 = 0$ , $p_3 = \frac{0.7417}{0.2583 + 0.7417} = 0.7417$ , $p_4 = 0$ .

Constraints

0 < P <= 1
Include the token that causes cumulative probability to exceed P
Renormalize selected probabilities to sum to 1
Round to 4 decimal places

More from LLM 3: Applications & Evaluation

Beam Search Decoderhard INT8 Quantization Simulatormedium Top-K Samplingeasy Temperature Samplingeasy BLEU-4 Scorehard

Concept primerPREVIEW

Background Knowledge

The problem revolves around nucleus (top-P) sampling, a technique used in natural language processing and machine learning, particularly in the context of language models. This method is designed to improve the quality of generated text by filtering out unlikely candidates. The core idea is to select a subset of tokens (words or characters) based on their probabilities, such that the cumulative probability of this subset exceeds a given threshold $P$ . This approach helps in avoiding low-probability tokens that might lead to unrealistic or low-quality generated text.

To understand nucleus sampling, one needs to be familiar with softmax functions, which are used to convert logits (unnormalized scores) into probabilities. The softmax function maps the input values to a probability distribution, ensuring that all output values are between 0 and 1 and sum up to 1. This is crucial because it allows us to interpret the outputs as probabilities. Additionally, understanding how to sort and accumulate probabilities is essential for implementing the nucleus sampling algorithm.

The concept of cumulative probability is also key. It refers to the sum of probabilities of all outcomes up to a certain point. In the context of nucleus sampling, we're interested in finding the smallest set of tokens whose cumulative probability is greater than or equal to the threshold $P$ . This involves sorting the tokens by their probabilities in descending order and then iteratively adding tokens to the set until the cumulative probability exceeds $P$ .

Algorithm/Approach

The general approach to solving this problem involves a combination of probability theory and algorithmic sorting. The algorithm can be broadly categorized into two main parts:

Probability Calculation and Sorting: This involves converting logits to probabilities using the softmax function and then sorting these probabilities in descending order.
Cumulative Probability Calculation and Thresholding: After sorting, the algorithm calculates the cumulative probability and identifies the smallest set of tokens that meets the threshold $P$ .

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