Temperature and Token Sampling

A GPT model does not simply announce the next word. It produces a score for every token in its vocabulary, and then a token is chosen from those scores by sampling. The setting that controls how that choice is made is called temperature. This article follows the path from raw token indices to the scores the model produces, and then to the counts that come out of sampling. Work through it in the order given, because each stage consumes the output of the stage before it.

A GPT model is built from three main parts: an embedding stage, a sequence of transformer blocks run with nn.Sequential, and a linear output head with no bias. The embedding stage turns token indices into vectors. The transformer blocks process those vectors. The output head converts the processed vectors into logits, which are the raw scores for each vocabulary entry. The GPTModel class is a PyTorch module, and all three parts are attributes of it. The diagram below shows how they sit inside the whole model.

The embedding stage has two lookups that are later added together. The token embedding represents which tokens appear. The positional embedding represents where in the sequence they appear. The transformer blocks are not a single layer but a stack of identical-shaped blocks, and the source builds that stack in one line. The output head is the last part, and it is the only part that produces the scores used for sampling.

The source builds the sequence of transformer blocks in a single expression. It creates one TransformerBlock for each layer, using the count stored under the n_layers key of the configuration dictionary, and wraps the blocks in nn.Sequential. Wrapping them this way means they run one after another, so the output of the first block becomes the input of the second, and so on. You do not need to call each block by hand inside forward; the sequential container handles the ordering.

The forward method is where an input of token indices becomes a set of output scores called logits. Read it as a pipeline. For each stage, ask three questions: what does this line receive, what does it produce, and which part of the GPT model does it belong to. The diagram below traces the input named in_idx from the embedding stage to the logits.

The phrase with no bias is part of the definition of the output head. The layer is created with bias set to False, so it only applies a weight matrix to its input. It has no added offset term. A learner reading the constructor should note this explicitly, because a bias term would change the arithmetic of every output score. Without a bias, each score is produced purely by the weight matrix acting on the embedding, with nothing added on top.

The output head maps each embedding to a vector whose size is vocab_size. That means each position in the sequence receives one score per vocabulary entry. The final dimension of the logits is therefore a lookup: position number k in that dimension corresponds to one specific token in the vocabulary, and the value stored there is the score for that token. Reading the logits this way lets you treat each entry as a score for one token, which is exactly what sampling needs.

Once the model has produced scores, a token is chosen by sampling. Each time a token is drawn it adds one to that token's count. If you draw N samples in total and a particular token was drawn c times, then c divided by N is an estimate of that token's probability: $\frac{c}{N}$. The estimate improves as N grows, but it is still an estimate. A count from 1000 draws is a sample-based estimate, and it is not the same thing as reading the probability directly from the distribution. Keeping both numbers in view helps you see the difference between what was sampled and what the distribution holds.

Temperature changes the probability distribution from which tokens are sampled. A higher temperature increases how often certain tokens are sampled. The source makes this concrete with the token 'pizza'. At temperature 0 or 0.1, 'pizza' is sampled zero times. When the temperature is scaled up to 5, it is sampled 32 times. The counts come from the print_sampled_tokens function that the source defines in an earlier section, which is the function to use if you want to reproduce these numbers yourself.

The before and after counts show that a token with no samples under a low temperature can accumulate many samples under a higher one. The source does not give the full distribution at each temperature, so the important lesson is the direction of the change: raising temperature shifts sampling toward tokens that were rarely picked before.