Transformer from Scratch

A decoder-only, GPT-style Transformer language model built from scratch in PyTorch and trained at the character level on the Tiny Shakespeare dataset. Every part of the architecture is written by hand: self-attention heads, multi-head attention, the feed-forward network, residual Transformer blocks, and the language-model head. PyTorch supplies only the basic layers (nn.Linear, nn.Embedding, nn.LayerNorm).

📝 Read the full write-up on Medium: How I Trained My First Transformer from Scratch

---

Table of Contents

• Highlights
• Repository Structure
• Getting Started
• Implementation Details
  • 1. Dataset
  • 2. Tokenization
  • 3. Train / Validation Split
  • 4. Batching
  • 5. Model Architecture
  • 6. Training
  • 7. Text Generation
• Hyperparameters
• Parameter Count
• Results
• Possible Improvements
• Acknowledgements

---

Highlights

• Architecture: Decoder-only Transformer (GPT-style) with causal masked self-attention
• Tokenization: Character-level, 65-token vocabulary
• Size: ~0.21M parameters (208,193)
• Depth: 4 Transformer blocks × 4 attention heads, 64-dim embeddings
• Training: 10,000 iterations with AdamW
• Result: Loss drops from 4.25 → 2.09 (validation), and the model generates text that looks like Shakespeare

---

Repository Structure

Transformer_from_scratch/
├── Transformer_from_scratch.ipynb   # Full implementation: data, model, training, generation
└── README.md

The dataset (input.txt) is downloaded by the notebook itself. It is not stored in the repo.

---

Getting Started

Requirements

• Python 3
• PyTorch
• Jupyter Notebook / JupyterLab (or Google Colab)

pip install torch

A CUDA GPU is used automatically if one is available. Otherwise the code falls back to CPU, and the model is small enough to train there too.

Run

1. Open Transformer_from_scratch.ipynb.
2. Run all cells from top to bottom. The notebook will:
   • download the Tiny Shakespeare dataset,
   • build the vocabulary and the train/val split,
   • define and train the model,
   • print the train/val loss every 100 steps,
   • generate 2,000 characters of new text.

Note: The dataset is downloaded with !wget. If wget isn't installed on your machine (common on macOS), download the file manually instead:

curl -o input.txt https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

---

Implementation Details

1. Dataset

The model is trained on Tiny Shakespeare, a single 1.1 MB text file (1,115,394 characters) containing a collection of Shakespeare's plays.

2. Tokenization

The model uses a simple character-level tokenizer:

• The vocabulary is every unique character in the dataset, sorted: 65 characters (letters, digits, punctuation, space and newline).
• stoi maps each character to an integer and itos maps it back.
• encode(str) -> list[int] and decode(list[int]) -> str convert between text and token IDs.

The whole dataset is then encoded into a single torch.long tensor.

3. Train / Validation Split

• The first 90% of the tokens are used for training.
• The remaining 10% are used for validation.

The split is sequential (not shuffled), so the validation set is a continuous piece of text the model never trains on.

4. Batching

get_batch(split) samples batch_size random starting positions and returns:

• x: input chunks of length block_size, with shape (B, T)
• y: the same chunks shifted one character to the right, with shape (B, T)

So every chunk gives block_size training examples at once. For each position t, the model sees the context x[:t+1] and has to predict the next character y[t]:

when input is [24]                 the target: 43
when input is [24, 43]             the target: 58
when input is [24, 43, 58]         the target: 5
...

5. Model Architecture

Input token IDs (B, T)
        │
        ├──► Token Embedding    (vocab_size → n_embd)
        ├──► Position Embedding (block_size → n_embd)
        ▼
   tok_emb + pos_emb            (B, T, C)
        │
        ▼
┌───────────────────────────────┐
│  Transformer Block  × n_layer │
│                               │
│  x = x + MultiHeadAttn(LN(x)) │
│  x = x + FeedForward(LN(x))   │
└───────────────────────────────┘
        │
        ▼
   Final LayerNorm (ln_f)
        │
        ▼
   Linear LM Head (n_embd → vocab_size)
        │
        ▼
   Logits (B, T, vocab_size)

Shape notation: B = batch size, T = sequence length (time), C = embedding dimension (n_embd).

Head: one head of causal self-attention

• Three bias-free linear projections turn each token into a key, query, and value of size head_size.

• Attention scores use scaled dot-product attention:

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

  where d_k = head_size. Scaling by 1/√d_k keeps the variance of the scores under control, so the softmax doesn't become too peaked when training starts.

• A causal mask is applied before the softmax: a lower-triangular matrix (tril, stored with register_buffer) sets future positions to -inf, so each token can only attend to itself and earlier tokens.

• Dropout is applied to the attention weights, and then they are used to take a weighted sum of the values.

MultiHeadAttention

• Runs n_head independent Heads in parallel, each with head_size = n_embd // n_head.
• Concatenates their outputs along the last dimension, so the result is back to size n_embd.
• Applies an output linear projection followed by dropout.

FeedFoward: position-wise feed-forward network

• Linear(n_embd → 4·n_embd) → ReLU → Linear(4·n_embd → n_embd) → Dropout
• It is applied to each token on its own. Attention is where tokens communicate; this layer is where each token computes on what it gathered.

Block: Transformer block

• Combines communication (attention) and computation (feed-forward).
• Uses residual connections around both sub-layers: x = x + sublayer(x).
• Uses pre-LayerNorm: LayerNorm is applied before each sub-layer, not after as in the original Attention Is All You Need paper. This makes training more stable.

GPTLanguageModel

• Token embedding and learned positional embedding. The two are added together.
• n_layer Transformer blocks stacked with nn.Sequential.
• A final LayerNorm, then a linear head that projects to vocabulary logits.
• When targets are passed in, the logits are reshaped to (B·T, C) and cross-entropy loss is computed against the next-character targets.

Since this is a decoder-only model (like GPT), there is no encoder and no cross-attention. Only masked self-attention is used.

6. Training

• Optimizer: AdamW, learning rate 1e-3, with no learning-rate schedule.
• Loop: For max_iters = 10,000 steps: sample a training batch → forward pass → zero_grad → backward → optimizer.step().
• Evaluation: Every eval_interval = 100 steps (and at the last step), estimate_loss() averages the loss over eval_iters = 200 batches for both the train and val splits. It runs under @torch.no_grad() and switches the model between eval() and train() modes.

7. Text Generation

generate(idx, max_new_tokens) generates text autoregressively:

1. Crop the context to the last block_size tokens, since that is the longest context the position embeddings support.
2. Run the model and keep only the logits for the last time step.
3. Apply softmax to get a probability distribution over the next character.
4. Sample the next token with torch.multinomial.
5. Append the new token to the sequence and repeat.

Generation starts from a single token with ID 0 (the newline character \n), and the notebook generates 2,000 new characters.

---

Hyperparameters

The first cell defines the hyperparameters. However, batch_size and block_size are set again in the data-exploration cells (to 4 and 8) before the model is built. So the recorded training run actually used the values in the "Used in run" column:

Hyperparameter	Set in first cell	Used in run	Description
batch_size	16	4	Independent sequences processed in parallel
block_size	32	8	Maximum context length (in characters)
max_iters	10,000	10,000	Total training iterations
eval_interval	100	100	Evaluate every N iterations
eval_iters	200	200	Batches averaged per loss estimate
learning_rate	1e-3	1e-3	AdamW learning rate
n_embd	64	64	Embedding dimension
n_head	4	4	Attention heads per block (head_size = 16)
n_layer	4	4	Number of Transformer blocks
dropout	0.0	0.0	Dropout rate

To train with the larger values (batch_size = 16, block_size = 32), remove or change the reassignment in the batching cells so that the first cell's values are the ones used.

---

Parameter Count

Total: 208,193 parameters (~0.21M)

Component	Parameters
Token embedding (65 × 64)	4,160
Position embedding (8 × 64)	512
4 × Transformer block (49,792 each)	199,168
• 4 heads × Q/K/V (64 × 16, no bias)	12,288
• Attention output projection	4,160
• Feed-forward (64 → 256 → 64)	33,088
• 2 × LayerNorm	256
Final LayerNorm	128
LM head (64 → 65)	4,225

---

Results

Loss

Step	Train Loss	Val Loss
0	4.2503	4.2551
1000	2.4142	2.4251
2000	2.2798	2.2797
5000	2.1210	2.1951
9999	2.0162	2.0908

The starting loss of ~4.25 is close to the loss you'd expect from uniform random guessing over 65 characters, which is −ln(1/65) ≈ 4.17. That confirms the model starts out with no knowledge. The train and val losses stay close together for the whole run, so the model is not overfitting badly.

Sample Output

An excerpt from the 2,000 characters the trained model generated:

DETENWICHANAABELLA:
The this my depity:
Augh are plewty.

SIBENSIRY ISABE:
And whom
Is would thake only so what exicks to they, rive weeshirn sto-lite lix thuse known,
Turt f so;
Angreen love,
I know ofter:
Firse?

The model has learned the structure of a play: character names in capitals followed by a colon, line breaks, dialogue-like rhythm and punctuation. It has also learned many real English words. With only 0.21M parameters and an 8-character context, it cannot yet keep meaning coherent. Because generation uses random sampling, your output will be different.

---

Possible Improvements

• Use the larger batch_size / block_size from the first cell, or go bigger, so the model sees more context.
• Scale up n_embd, n_head and n_layer, and turn on dropout (e.g. 0.2) to regularize the larger model.
• Add a learning-rate warmup and cosine decay schedule.
• Switch to a subword tokenizer (e.g. BPE) instead of characters.
• Add temperature and top-k sampling to generate().
• Use PyTorch's fused F.scaled_dot_product_attention for speed.

---

Acknowledgements

• Andrej Karpathy's Let's build GPT: from scratch, in code, spelled out lecture and nanoGPT, which this implementation follows.
• The Tiny Shakespeare dataset from karpathy/char-rnn.
• Vaswani et al., Attention Is All You Need (2017).