GPT-2 Replication from Scratch

Task: Natural Language Processing  |  Architecture: Generative Pre-trained Transformer (GPT-2)

Parameter	Value
Parameters	162,419,712
Model size	619.58 MB
Context length	256
Embedding dimension	768
Attention heads	12
Transformer layers	12

---

Abstract

This project replicates the GPT-2 language model [1] from scratch using PyTorch, following the architectural specification of the original paper. All core components — tokenization, embedding, multi-head causal self-attention, feed-forward layers, and layer normalization — are implemented as standalone modules. The model is trained on the Harry Potter corpus [8] using a standard language modeling objective. Code reference follows Raschka [9].

---

1. Architecture

The model follows the standard decoder-only Transformer pipeline. For an input sequence of token indices, the forward pass proceeds as:

$$\mathbf{x} = \text{Dropout}\left( \mathbf{E}_{\text{tok}}(\text{idx}) + \mathbf{E}_{\text{pos}}(\text{arange}(T)) \right)$$

$$\mathbf{x} = \text{TransformerBlock}^{(L)}(\cdots \text{TransformerBlock}^{(1)}(\mathbf{x}))$$

$$\text{logits} = \mathbf{W}_{\text{out}} \cdot \text{LayerNorm}(\mathbf{x})$$

where $T$ is the sequence length, $L = 12$ is the number of stacked blocks, and $\mathbf{W}_{\text{out}} \in \mathbb{R}^{d \times V}$ projects to vocabulary logits. Both embeddings share dimension $d = 768$:

$$\mathbf{E}_{\text{tok}} \in \mathbb{R}^{V \times d}, \qquad \mathbf{E}_{\text{pos}} \in \mathbb{R}^{T \times d}$$

---

2. Components

2.1 Tokenization

Text is tokenized using OpenAI's Byte Pair Encoding (BPE) scheme [7] via tiktoken with the gpt2 vocabulary ($V = 50{,}257$). BPE iteratively merges the most frequent adjacent byte pairs, yielding a vocabulary of subword units that balances coverage and sequence length.

A simpler whitespace-and-punctuation tokenizer (CustomTokenizer) is also implemented for reference, operating over the corpus vocabulary only.

2.2 Input Embedding

The input representation is formed by summing token and positional embeddings [3]:

$$\mathbf{h}_i = \mathbf{e}_{\text{tok}}(w_i) + \mathbf{e}_{\text{pos}}(i), \quad i = 1, \ldots, T$$

where both are learned lookup tables:

$$\mathbf{e}_{\text{tok}}: \mathbb{Z} \to \mathbb{R}^d, \qquad \mathbf{e}_{\text{pos}}: \mathbb{Z} \to \mathbb{R}^d$$

The positional embedding is learned rather than fixed, consistent with the GPT-2 specification.

2.3 Causal Multi-Head Attention

Scaled dot-product attention is computed as [2]:

$$\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}} + \mathbf{M}\right)\mathbf{V}$$

where $\mathbf{M}$ is the causal mask, defined as:

$$M_{ij} = \begin{cases} 0 & \text{if } i \geq j \ -\infty & \text{if } i < j \end{cases}$$

Setting future positions to $-\infty$ before softmax ensures they contribute zero weight, preventing the model from attending to tokens it has not yet generated.

For multi-head attention with $h = 12$ heads and $d = 768$, the projection dimension per head is $d_k = d / h = 64$. The full sequence of operations is:

$$\mathbf{Q}, \mathbf{K}, \mathbf{V} = \mathbf{x}\mathbf{W}_Q,\ \mathbf{x}\mathbf{W}_K,\ \mathbf{x}\mathbf{W}_V \quad \in \mathbb{R}^{B \times T \times d}$$

$$\text{head}_i = \text{Attention}(\mathbf{Q}_i, \mathbf{K}_i, \mathbf{V}_i) \quad \in \mathbb{R}^{B \times T \times d_k}$$

$$\text{MultiHead}(\mathbf{x}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h),\mathbf{W}_O$$

In the implementation, the head split is performed via .view(B, T, h, d_k).transpose(1, 2) rather than running $h$ separate projections, and the output projection $\mathbf{W}_O \in \mathbb{R}^{d \times d}$ recombines the heads.

2.4 Activation Function

The feed-forward sublayer uses the tanh approximation of the Gaussian Error Linear Unit (GELU) [6]:

$$\text{GELU}(x) = 0.5 \cdot x \cdot \left(1 + \tanh!\left(\sqrt{\tfrac{2}{\pi}},(x + 0.044715,x^3)\right)\right)$$

This is the exact approximation used in the original GPT-2 codebase, and is reproduced verbatim in Activation.py.

2.5 Layer Normalization

Each sublayer applies pre-norm layer normalization [5] with learnable parameters $\gamma$ (scale) and $\beta$ (shift):

$$\text{LayerNorm}(\mathbf{x}) = \gamma \cdot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \varepsilon}} + \beta$$

where:

$$\mu = \frac{1}{d}\sum_j x_j, \qquad \sigma^2 = \frac{1}{d}\sum_j (x_j - \mu)^2$$

(biased estimator, unbiased=False), and $\varepsilon = 10^{-5}$. Pre-norm placement — normalizing the input before each sublayer rather than after — is the key architectural difference from the original Transformer [2].

2.6 Feed-Forward Sublayer

The position-wise feed-forward network expands the embedding dimension by a factor of 4 before projecting back:

$$\text{FFN}(\mathbf{x}) = \mathbf{W}_2,\text{GELU}(\mathbf{W}_1 \mathbf{x} + \mathbf{b}_1) + \mathbf{b}_2$$

where $\mathbf{W}_1 \in \mathbb{R}^{4d \times d}$ expands the dimension and $\mathbf{W}_2 \in \mathbb{R}^{d \times 4d}$ projects back, giving an intermediate dimension of $4 \times 768 = 3072$.

2.7 Transformer Block

Each of the $L = 12$ transformer blocks applies attention and feed-forward as residual branches over a pre-normalized input:

$$\mathbf{x} \leftarrow \mathbf{x} + \text{Dropout}!\left(\text{MHA}(\text{LayerNorm}(\mathbf{x}))\right)$$

$$\mathbf{x} \leftarrow \mathbf{x} + \text{Dropout}!\left(\text{FFN}(\text{LayerNorm}(\mathbf{x}))\right)$$

---

3. Data Pipeline

The corpus is chunked into training samples using a sliding window over the full token sequence. Given a tokenized corpus of length $N$, context length $T$, and stride $s$, the $i$-th sample is constructed as:

$$\text{input}^{(i)} = (w_{is},\ w_{is+1},\ \ldots,\ w_{is+T-1})$$

$$\text{target}^{(i)} = (w_{is+1},\ w_{is+2},\ \ldots,\ w_{is+T})$$

Each target is the input shifted by one position, encoding the next-token prediction objective directly in the dataset construction. When $s &lt; T$, consecutive windows overlap, increasing the number of training samples at the cost of data redundancy. In this project, $s = T = 256$ (no overlap) for training and validation respectively.

The corpus is split 90/10 into train and validation sets prior to windowing, and fed to the model in batches of 32.

---

4. Training

The model is trained with a standard language modeling objective — cross-entropy loss over next-token prediction. For a sequence of targets $y_1, \ldots, y_T$, the loss is:

$$\mathcal{L} = -\frac{1}{T} \sum_{t=1}^{T} \log P(y_t \mid y_{1:t-1};, \theta)$$

Perplexity, reported as a secondary metric, is the exponentiated average negative log-likelihood:

$$\text{PPL} = \exp(\mathcal{L})$$

The optimizer is AdamW with learning rate $4 \times 10^{-4}$ and weight decay $0.1$. The model checkpoint is saved to ./model/ after each training run and automatically reloaded on subsequent runs.

---

5. Module Dependency Tree

0_main.py
│
├── Typewriter.py
├── Tokenizer.py
├── DataLoader.py
└── GPT.py
    │
    ├── Tokenizer.py
    ├── Embedding.py
    │   ├── Tokenizer.py
    │   └── DataLoader.py
    ├── MaskDropout.py
    │   └── Softmax.py
    ├── LayerNorm.py
    └── TransformerBlock.py
        ├── MultiHeadAttention.py
        │   └── MaskDropout.py
        │       └── Softmax.py
        ├── Activation.py
        ├── LayerNorm.py
        └── MaskDropout.py
            └── Softmax.py

1_chat.py

Note: SelfAttention.py is not part of the active pipeline; it is retained as a pedagogical progression from simplified → trainable → batched self-attention, illustrating the derivation of the full MultiHeadAttention implementation.

---

6. Usage

Training

python 0_main.py

Trains the model for 100 epochs on harrypotter.txt, evaluating every 5 steps and printing a sample generation after each epoch. Saves the final checkpoint to ./model/GPT2-epoch-100.pt. If a checkpoint already exists in ./model/, it is loaded automatically before training resumes.

Inference

python 1_chat.py

Loads the latest checkpoint and opens an interactive prompt. Input is prefixed as Voldemort: <input>. Harry: and the model generates a continuation of up to 100 tokens using greedy decoding (argmax). Type stop to exit.

Note: Inference uses a context window of 50 tokens, smaller than the training context length of 256. Only the most recent 50 tokens are visible to the model during generation.

---

References

[1] Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., & Sutskever, I. (2019). Language Models are Unsupervised Multitask Learners. OpenAI Blog. https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf

[2] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. https://arxiv.org/abs/1706.03762

[3] Mikolov, T., Sutskever, I., Chen, K., Corrado, G., & Dean, J. (2013). Distributed Representations of Words and Phrases and their Compositionality. NeurIPS 2013. https://arxiv.org/abs/1310.4546

[4] Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., & Salakhutdinov, R. (2014). Dropout: A Simple Way to Prevent Neural Networks from Overfitting. JMLR, 15(1), 1929–1958. https://www.cs.toronto.edu/~rsalakhu/papers/srivastava14a.pdf

[5] Ba, J. L., Kiros, J. R., & Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450. https://arxiv.org/abs/1607.06450

[6] Hendrycks, D., & Gimpel, K. (2016). Gaussian Error Linear Units (GELUs). arXiv:1606.08415. https://arxiv.org/abs/1606.08415

[7] Gage, P. (1994). A new algorithm for data compression. C Users Journal, 12(2). https://dl.acm.org/doi/abs/10.5555/177910.177914

[8] Kudła, M. Harry Potter all books (preprocessed). Kaggle. https://www.kaggle.com/datasets/moxxis/harry-potter-lstm

[9] Raschka, S. Build a Large Language Model (From Scratch). GitHub. https://github.com/rasbt/LLMs-from-scratch