Postprocessing Logits into Token IDs

This article begins with code that is already running rather than with a list of definitions. A 124-million-parameter DummyGPTModel is created and fed a batch of tokenized text. The result that comes back is called the model output, and in this context it is commonly referred to as logits. Your job is to watch what the code produces and then work out why the result has the shape it has. Once you can read that shape, the rest of the printed output becomes much easier to interpret.

The run you will trace feeds the model two text samples, and each sample contains four tokens. Everything in this article is built around one question: when the model finishes, what exactly is the structure of the result it hands back?

Two things are printed when the run is shown: the shape, and the values that follow it. The shape tells you how the output is organized. The values are the numbers the model produced, and they only make sense once you know what organization they belong to. So read the shape first.

The shape has three numbers, and they are not interchangeable. Each one answers a different question about the output. Before you can explain any value in the printed tensor, you need to know which question each number answers, so the next section takes the three axes one at a time.

A tensor is easiest to read one axis at a time. The output of this model has three axes, and each one answers a different question. The first axis asks which text sample a result belongs to. The second asks which token position inside that sample it belongs to. The third asks which vocabulary entry that particular number is scored against.

Read in order, the first axis groups the results by text sample, so the 2 tells you there are two rows, one per sample. The second axis, with size 4, gives the token position inside a sample, so each sample has four positions. The third axis, with size 50257, is the vocabulary dimension. The source describes the result as a grid of samples and token positions in which every grid cell holds a whole vocabulary-sized vector rather than a single value.

The reason the vector at each position has 50,257 entries is that each entry corresponds to one unique token in the vocabulary. This is the key to reading the output. The model is not handing back a single guessed word. For each position, it gives a set of numbers spread across the entire vocabulary, and the position of a number inside that vector is what links it to a specific token.

The diagram shows a single token position. Its output is one row of 50,257 slots, and each slot is tied to one token in the vocabulary. The vocabulary therefore sets the width of every vector, and every position in the vector is a slot for one specific token. The same two-way correspondence governs how tokens are represented in general: each token is a high-dimensional vector whose size matches the tokenizer's vocabulary, so a token can be thought of as a position in a space with one dimension for every token the tokenizer knows.

The source also explains what happens after the model returns. The 50,257-dimensional vectors are converted back into token IDs, and those IDs can then be decoded into words. The flow below shows the order of steps from input token IDs, through the model, to the logits tensor, then through postprocessing, and finally to words.

The logits tensor is the point at which the shape is fully visible, and the postprocessing step is where it is turned into token IDs. The source places the placeholder model between the input and the output you studied here. Layer normalization is the next placeholder the source plans to replace, and it sits inside the model in the same position. Its details are outside the scope of this article.

Now test the reading you have built. Suppose the input changes from two samples of four tokens to three samples of six tokens. This example is generated for practice and is not a result taken from the source. It uses the same reasoning the source applies: the first axis corresponds to text samples, the second to tokens, and the last to vocabulary entries.

The before and after columns show the rule. The sample count and the token count are the two dimensions that grow when you change the input, while the vocabulary dimension remains 50257 in both runs. This is the stable part of the shape, and it is the part you can rely on when you read any logits output from this model.

The model output for this run is called logits, and its shape torch.Size([2, 4, 50257]) encodes three facts: two text samples, four tokens per sample, and one slot for each of the 50,257 vocabulary entries. Each slot is tied to one specific token, and the logits are converted into token IDs in postprocessing before being decoded into words later.