Decoder-Only Transformer Architecture Implementation in Pytorch (Reference work)

A complete implementation of a Decoder-Only Transformer (GPT-style) built from the ground up using PyTorch, without relying on high-level abstractions. This implementation includes all core components: token embeddings, positional embeddings, multi-head self-attention, feedforward networks, causal masking, and output logits generation.

Table of Contents

• Overview
• Architecture Components
• Data Flow: Tokens to Softmax Probabilities
• Installation & Setup
• Usage Examples
• Detailed Component Explanations
• References & Resources

---

Overview

A Transformer is a neural network architecture that processes sequences of tokens through layers of self-attention and feedforward networks. This implementation focuses on the Decoder-Only variant, which can be used for:

• Language modeling: Predicting the next token given previous tokens
• Text generation: Autoregressive sequence generation
• Instruction following: Like GPT, ChatGPT, and similar models

Key Features

✅ Pure PyTorch implementation from scratch
✅ Multi-head self-attention with causal masking
✅ Layer normalization and residual connections
✅ GELU activation function
✅ Learnable token and positional embeddings
✅ Configurable model dimensions and layers

---

Architecture Components

1. Token Embedding

Converts discrete token IDs (integers 0 to vocab_size-1) into continuous dense vectors of dimension d_model.

Input: Token indices [0, 1, 5, 2] 
Output: Dense vectors of shape (batch_size, seq_len, d_model)

2. Positional Embedding

Adds sequential position information to preserve word order, since self-attention is position-agnostic.

Token Embedding + Positional Embedding = Input Representation

3. Multi-Head Self-Attention

Allows the model to attend to different positions, computing relationships between tokens.

Attention(Q, K, V) = softmax(Q·K^T / √d_k)·V

4. Causal Masking

Prevents the model from attending to future tokens during training/inference (autoregressive constraint).

Mask prevents: token at position t from seeing positions > t

5. Feedforward Network (FFN)

Two linear transformations with GELU activation for non-linearity.

FFN(x) = Linear(GELU(Linear(x)))

6. Layer Normalization & Residual Connections

Stabilizes training and enables deeper networks.

Output = LayerNorm(Input) → AttentionOrFFN → Add(Input)

7. Output Layer (LM Head)

Projects the final hidden states to vocabulary logits for probability distribution over next tokens.

Logits = Linear(d_model → vocab_size)

---

Data Flow: Tokens to Softmax Probabilities

Here's the complete journey of data through the model:

[STEP 1] INPUT TOKENS
─────────────────────
Raw token IDs: [0, 5, 12, 3]  (integers from 0 to vocab_size-1)
Shape: (batch_size=1, seq_len=4)

       ↓

[STEP 2] TOKEN EMBEDDING
──────────────────────
Token embedding maps each ID to a d_model-dimensional vector
ID 0 → [-0.21, 0.54, -0.12, ..., 0.33]  (d_model dimensions)
ID 5 → [0.15, -0.34, 0.67, ..., -0.21]
...
Output shape: (1, 4, 512) with d_model=512

       ↓

[STEP 3] POSITIONAL EMBEDDING
──────────────────────────────
Add position information to preserve sequence order
Position 0 → [0.0, 1.0, 0.0, 1.0, ..., 0.0, 1.0]  (sinusoidal patterns)
Position 1 → [0.841, 0.540, 0.998, 0.0, ..., 0.001]
Position 2 → [0.909, -0.416, 0.995, -0.100, ...]
Position 3 → [0.141, -0.990, 0.959, -0.283, ...]
Output shape: (1, 4, 512)

Token Embedding + Positional Embedding = (1, 4, 512)

       ↓

[STEP 4] DECODER BLOCKS (×N layers)
───────────────────────────────────

For each of N decoder blocks:

  ┌─────────────────────────────────────────┐
  │                                         │
  │  [STEP 4a] MULTI-HEAD SELF-ATTENTION  │
  │  ─────────────────────────────────     │
  │  • Compute Q, K, V from input          │
  │  • Split into n_heads (e.g., 8)        │
  │  • Each head: Attention(Q,K,V)         │
  │  • Apply CAUSAL MASK (prevent future)  │
  │  • Concatenate heads                   │
  │  • Output projection                   │
  │  Residual connection: x + Attention(x)│
  │                                        │
  │  [STEP 4b] LAYER NORM                  │
  │  ─────────────────────                 │
  │  Normalize across feature dimension    │
  │                                        │
  │  [STEP 4c] FEEDFORWARD NETWORK         │
  │  ──────────────────────────────        │
  │  • Linear(512 → 2048)                  │
  │  • GELU activation                     │
  │  • Linear(2048 → 512)                  │
  │  Residual connection: x + FFN(x)      │
  │                                        │
  │  Output: (1, 4, 512)                   │
  └─────────────────────────────────────────┘
       ↓ (repeat for N layers)

[STEP 5] FINAL LAYER NORMALIZATION
──────────────────────────────────
Normalize the final hidden states
Output shape: (1, 4, 512)

       ↓

[STEP 6] LM HEAD (Output Projection)
────────────────────────────────────
Linear transformation from d_model to vocab_size
Output shape: (1, 4, vocab_size)  [logits]
Example: (1, 4, 50000)

       ↓

[STEP 7] SOFTMAX PROBABILITIES
───────────────────────────────
Apply softmax to convert logits to probabilities
P(next_token | input) = softmax(logits)
Output shape: (1, 4, 50000)  [probabilities]
Each value in [0, 1], sum across vocabulary = 1

For position 4, probabilities:
  Token "the": 0.0342
  Token "cat": 0.0157
  Token "is": 0.1893  ← highest probability
  Token "dog": 0.0012
  ... (50000 total values)

SAMPLING / GREEDY DECODING:
argmax(probabilities) → predicted next token ID
Sample from distribution → next token ID (more diverse)

---

Installation & Setup

Prerequisites

• Python 3.8+
• PyTorch 1.9+

Installation Steps

1. Clone or download the repository

   git clone https://github.com/shaheennabi/Transformer-Architecture-Implementation-from-Scratch-in-PyTorch.git
   cd Transformer-Architecture-Implementation-from-Scratch-in-PyTorch

2. Create a virtual environment (recommended)

   python -m venv venv

3. Activate virtual environment

   • On Windows:

     venv\Scripts\activate

   • On macOS/Linux:

     source venv/bin/activate

4. Install PyTorch

   Visit pytorch.org and install the appropriate version for your system:

   # CPU only
   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
   
   # GPU (CUDA 11.8)
   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   
   # GPU (CUDA 12.1)
   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

5. Verify installation

   python -c "import torch; print(torch.__version__)"

---

Detailed Component Explanations

Multi-Head Self-Attention

Purpose: Allow tokens to attend to all previous positions in parallel.

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

Steps:

1. Project input into Query (Q), Key (K), Value (V)
2. Split into multiple heads: d_model / n_heads per head
3. Compute attention scores for each head
4. Apply softmax to get attention weights
5. Multiply by values and concatenate heads
6. Project back to d_model

Code:

scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.head_dim)  # scaling
scores = scores.masked_fill(mask == 0, float('-inf'))          # causal mask
attn = torch.softmax(scores, dim=-1)                           # attention weights
out = attn @ V                                                  # apply attention

Causal Masking

Purpose: Enforce autoregressive constraint (token can only attend to previous tokens).

Mechanism:

• Create lower triangular matrix of 1s (allowed positions)
• Fill upper triangle with 0s (forbidden positions)
• Set forbidden attention scores to -inf before softmax
• After softmax, these become 0 probability

Example (3-token sequence):

Mask:
[[1, 0, 0],      After softmax:  position 0 → can see [0]
 [1, 1, 0],                      position 1 → can see [0, 1]
 [1, 1, 1]]                      position 2 → can see [0, 1, 2]

Feedforward Network (FFN)

Purpose: Add non-linearity and increase model capacity.

Architecture:

Linear(d_model → d_ff) → GELU → Linear(d_ff → d_model)

Typical sizes:

• d_model = 768 → d_ff = 3072 (4× expansion)
• d_model = 512 → d_ff = 2048 (4× expansion)

Layer Normalization & Residual Connections

Layer Norm: Normalizes inputs to have mean=0, std=1 across feature dimension.

Residual Connection:

Output = LayerNorm(Input) → Process → Add(Original Input)

Benefits:

• Stabilizes training
• Prevents vanishing gradients
• Enables deeper networks

---

References & Resources

Papers

• "Attention Is All You Need" (Vaswani et al., 2017)
  arXiv:1706.03762

• "Language Models are Unsupervised Multitask Learners" (Radford et al., 2019) - GPT-2
  OpenAI Blog

• "Language Models are Few-Shot Learners" (Brown et al., 2020) - GPT-3
  arXiv:2005.14165

Educational Resources

• The Illustrated Transformer by Jay Alammar
  https://jalammar.github.io/illustrated-transformer/

• A Gentle Introduction to Positional Encoding in Transformers
  https://machinelearningmastery.com/positional-encoding-transformer-neural-networks/

• Understanding Multi-Head Attention
  https://towardsdatascience.com/multi-head-attention-f3ca05bbf6ef

Official Documentation

• PyTorch Official Docs: https://pytorch.org/docs/stable/index.html
• PyTorch Attention Mechanism: https://pytorch.org/docs/stable/nn.html#attention-layers

Implementation References

• Hugging Face Transformers: https://huggingface.co/transformers/
• nanoGPT by Andrej Karpathy: https://github.com/karpathy/nanoGPT

---

Key Configuration Parameters

Parameter	Description	Typical Value
vocab_size	Number of tokens in vocabulary	50000
d_model	Hidden dimension size	512, 768, 1024
n_heads	Number of attention heads	8, 12, 16
n_layers	Number of decoder blocks	6, 12, 24
d_ff	Feedforward hidden dimension	4 × d_model
max_seq_len	Maximum sequence length	256, 512, 1024, 2048

---

License

This project is provided as-is for educational purposes. See LICENSE for details.

---

Contributing

Contributions are welcome! Feel free to submit issues and pull requests.

---

Happy transforming! 🚀