A linear output layer projects the final normalized transformer representation into vocabulary-sized logits for next-token prediction. In the GPT-2 configuration, it produces 50,257 values per sequence position. section 4.6
Data reaches the output head after token and positional embeddings, repeated transformer blocks, and final layer normalization. The implementation uses a bias-free linear layer mapping the embedding dimension to the tokenizer vocabulary size. section 4.6
Implementation details:
- Define the head as nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False). section 4.6
- Return logits whose final dimension equals the vocabulary size; convert these unnormalized values into token probabilities afterward. section 4.6
Shape check
For a batch of 2 sequences with 4 tokens each, the model returns torch.Size([2, 4, 50257]). The output-head weights have shape torch.Size([50257, 768]) in the 124-million-parameter configuration. section 4.6
Figure 4.15 An overview of the GPT model architecture showing the flow of data through the GPT model. Starting from the bottom, tokenized text is first converted into token embeddings, which are then augmented with positional embeddings. This combined information forms a tensor that is passed through a series of transformer blocks shown in the center (each containing multi-head attention and feed forward neural network layers with dropout and layer normalization), which are stacked on top of each other and repeated 12 times.
Follow the architecture diagram from embeddings through stacked transformer blocks and final normalization to the vocabulary projection. section 4.6
