Decode autoregressively

Learn it

Words to know

autoregressive

Producing output one token at a time, with each new token added to the input for the next step.

greedy decoding

Always choosing the token with the highest score.

end-of-sequence token

A special token that means the text is finished. It is often written eos.

bigram model

The simplest language model: it predicts the next token from the current token alone.

How it works

1. Start with the prefix. Here the model is a bigram table: the scores for the next token are the row of logits at the last token's ID.
2. Pick the position of the largest score in that row. On a tie, the smallest position wins. Append it to the sequence.
3. Stop after max_new new tokens.
4. Also stop as soon as the sequence ends with eos, whether that token was in the prefix or was added a moment ago.

Worked example

Input
prefix = [0]
logits = [[0, 3, 1], [0, 0, 4], [0, 0, 1]]
max_new = 5
eos = 2

Output
[0, 1, 2]

Five new tokens were allowed, but only two were written. The end token stopped the loop early.

Why it matters

The workshop's GPT uses this loop with two changes. It has no end token, so it always writes a fixed number of tokens. And it draws each token at random, likelier tokens more often, instead of taking the top one.

Your turn

Extend a sequence of token IDs by greedy decoding from a bigram score table.

prefix

the token IDs to start from; at least one

logits

a square table: row i holds the scores for the token after token i

max_new

the largest number of tokens to add; 0 or more

eos

the ID of the end token

Returns

The whole sequence: the prefix followed by the new tokens. Nothing is added once the sequence ends with eos.

Input
prefix = [0]
logits = [[0, 3, 1], [0, 0, 4], [0, 0, 1]]
max_new = 5
eos = 2

Output
[0, 1, 2]

Input
prefix = [2]
logits = [[0, 1, 2], [0, 1, 2], [0, 1, 2]]
max_new = 4
eos = 2

Output
[2]

Input
prefix = [0]
logits = [[1, 1], [0, 1]]
max_new = 2
eos = 1

Output
[0, 0, 0]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   Copy the prefix into a new list, then loop at most max_new times. Each pass adds one token or stops.

2. Hint 2: The key step

   At the top of each pass, check whether the last token is eos. If not, take the row logits[last] and find the index of its largest value.

3. Hint 3: The outline

   Repeat max_new times: if result[-1] == eos, break. Otherwise set row = logits[result[-1]] and append row.index(max(row)): max(row) is the largest score and index gives the first position that holds it. Return result.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. Work on a copy, so the caller's prefix is not changed. list(prefix) makes a new list. Appending to it leaves the original as it was.
2. The end token means the text is finished. Checking at the top of each pass covers both cases: an end token already in the prefix, and one appended on the pass before.
3. The scores for what follows the last token. A bigram model looks only at the last token. A GPT would run the whole sequence through the network here.
4. Greedy: take the position of the largest score. max(row) is the largest score, and row.index returns the first position that holds it. On a tie that is the smallest ID.

The same numbered steps are the comments in the solution.

def solve(prefix, logits, max_new, eos):
    # 1. Work on a copy, so the caller's prefix is not changed
    result = list(prefix)
    for _ in range(max_new):
        # 2. The end token means the text is finished
        if result[-1] == eos:
            break
        # 3. The scores for what follows the last token
        row = logits[result[-1]]
        # 4. Greedy: take the position of the largest score
        best = row.index(max(row))
        result.append(best)
    return result

Watch out

• Always cap the number of new tokens. A model that never produces the end token would otherwise loop for ever.
• max(row) gives the largest score, not its position. The token ID is the position.

Check yourself