Building a GPT Model Architecture

When you build the 124-million-parameter GPT model in this chapter, a common first reaction is confusion. You set up a model with that name, then add up the sizes of its parameter tensors, and the total comes out noticeably larger. It is easy to assume the setup has a bug or that the count is wrong. Neither is true. The gap comes from a design choice in the original GPT-2 architecture, and understanding that choice is the main purpose of this article.

The model is initialized with the GPT_CONFIG_124M configuration. Its total is computed by adding the element count of every parameter tensor. In Python, the element count of a tensor comes from the numel() method, and the total is produced by the expression below. Before you look at the answer, commit to a number of your own, because the point of the exercise is to notice where your intuition about the name and the actual count disagree.

The resolution to the mismatch is weight tying. In the original GPT-2 architecture, the weights of the token embedding layer are reused in the output layer. Instead of storing two separate matrices, the model keeps one set of values and both layers draw on it. This is the definition to hold onto: weight tying means reusing the token embedding weights in the output layer.

The sizes make the connection concrete. Both layers in this model have a weight tensor of shape 50257 by 768. The 50257 comes from the tokenizer's vocabulary, with one row for each token, and 768 is the embedding dimension, so each token has a 768-value vector. The output layer must map each hidden vector back to one score per vocabulary token, so its matrix has the same vocabulary-by-embedding dimensions. Those matching shapes are what make reuse possible. Reuse is only possible because the two matrices have identical dimensions.

Under weight tying, the output layer does not add its own large block of values. It reads from the same stored weights that the embedding layer uses. Because the output layer alone accounts for 38,597,376 values here, sharing it is a substantial saving. The tied count is therefore not a separate calculation of the model's size. It is the full count minus the parameters that weight tying makes unnecessary.

Parameter counts describe the model's size, but initialization also produces an output you should be able to read. The model is run on a batch of two input texts, each with four tokens. The output tensor therefore has shape [2, 4, 50257]. Each dimension has a clear meaning. The first is the batch size, which is 2 here. The second is the sequence length, which is 4 tokens per text. The last is the vocabulary size of the tokenizer, which is 50257.

The final dimension is the important one to interpret. Each token position receives a vector of 50,257 numbers, one for each token in the vocabulary. These are not yet words. A later step converts each of these 50,257-dimensional vectors back into tokens. At this stage, the goal is only to confirm that the shape is what the design predicts. Reading the shape this way means that for any position in any text in the batch, you know there is one score per vocabulary entry.

Weight tying is a trade-off, not a rule that every model must follow. It reduces the overall memory footprint and the computational complexity of the model. The author of the course text also reports that keeping separate token embedding and output layers led to better training and model performance. For that reason, the GPTModel implementation in this chapter keeps the layers separate. Weight tying will be revisited in chapter 6, when pretrained weights from OpenAI are loaded. Modern LLMs also follow the separate-layer approach.