Temperature Sampling & Generation

The problem

A language model only gives logits for the next token. To generate text you need a loop:

1. turn the logits into probabilities (with a temperature);
2. pick one token at random, according to those probabilities;
3. append it to the text and ask the model again.

You will write one function for each step.

Step 1: softmax with temperature

Divide the logits by the temperature $T$, then apply softmax:

where $z_i$ is the logit of token $i$, $q_i$ its probability, and $T > 0$ the temperature. A temperature below 1 makes the distribution sharper (the favourite gets even more), above 1 flatter.

Example. Logits [2, 1, 0]:

• T = 1: [0.665, 0.245, 0.090], plain softmax;
• T = 0.5: the logits become [4, 2, 0], and the probabilities [0.867, 0.117, 0.016];
• T = 0: dividing by 0 is impossible, so it is defined as greedy decoding: all the probability on the largest logit, [1, 0, 0]. If several logits are tied for the largest, use the first one.

Step 2: drawing a token

To draw a token with the right probabilities, line the tokens up on the interval from 0 to 1, each taking a piece as long as its probability, then throw a random point u on the interval and see where it lands. This is the inverse CDF method: the CDF (cumulative distribution function) is the running total of the probabilities.

1. Compute the running totals of the probabilities: the cumulative sums.
2. Draw one random number u = rng.random(), between 0 and 1.
3. Return the first index whose cumulative sum is greater than u.

Example. probs = [0.2, 0.5, 0.3]. The cumulative sums are [0.2, 0.7, 1.0], so token 0 owns the piece from 0 to 0.2, token 1 from 0.2 to 0.7, token 2 from 0.7 to 1.

• u = 0.1 returns 0 (the first sum above 0.1 is 0.2);
• u = 0.65 returns 1 (the first sum above 0.65 is 0.7);
• u = 0.95 returns 2;
• u = 0.2 exactly returns 1: the sum must be greater than u, not equal.

Token 1 owns half of the interval, so it is drawn half of the time: exactly its probability.

Step 3: generating

next_logits(tokens) stands for the model: it takes the list of tokens so far and returns the logits of the next one. Starting from the prompt, repeat n_new times: get the logits, turn them into probabilities with the temperature, draw a token, append it.

Example. prompt = [7] and n_new = 3: the model is asked with [7], say it draws 4; then with [7, 4], say 4 again; then with [7, 4, 4], say 1. The result is [7, 4, 4, 1], prompt included.

What to implement

• softmax_with_temperature(logits, T): for a 1-D array of logits, return the probabilities of step 1.
• sample_token(probs, rng): draw one index as in step 2.
• generate(next_logits, prompt, n_new, T, rng): return the whole list of step 3, prompt included.

Requirements

• Stable for large logits: subtract the maximum before exp (section 0).
• sample_token must call rng.random() exactly once per token, so the tests can reproduce your samples.
• generate must not modify prompt.

Hints

• np.searchsorted(np.cumsum(probs), u, side="right") finds the first index whose cumulative sum is greater than u.
• Rounding can make the last cumulative sum slightly less than 1: ten probabilities of 0.1 add up to 0.9999999999999999. A u above that would give an index one past the end, so clip the result to len(probs) - 1.