🚀 Transformer From Scratch (PyTorch)

A complete implementation of the Transformer architecture from the paper Attention Is All You Need, built entirely with PyTorch.

---

📖 Overview

The Transformer changed Natural Language Processing by replacing recurrent networks with self-attention, allowing models to process entire sequences in parallel.

This repository implements every major component from scratch without using torch.nn.Transformer.

It is designed for:

• 🎓 Students learning Transformers
• 👨‍💻 Deep Learning practitioners
• 🔬 AI researchers
• 💼 Interview preparation
• 🚀 Building custom NLP models

---

✨ Features

• Token Embeddings
• Sinusoidal Positional Encoding
• Multi-Head Self Attention
• Masked Multi-Head Attention
• Encoder–Decoder Attention
• Position-wise Feed Forward Network
• Residual Connections
• Layer Normalization
• Stacked Encoder Layers
• Stacked Decoder Layers
• Final Vocabulary Projection

---

🏗️ Overall Architecture

flowchart TD

A[Source Tokens]
B[Embedding]
C[Positional Encoding]

D["Encoder × N"]

E[Encoder Memory]

F[Target Tokens]
G[Embedding]
H[Positional Encoding]

I["Decoder × N"]

J[Linear Layer]

K[Vocabulary Probabilities]

A --> B --> C --> D --> E

F --> G --> H --> I

E --> I

I --> J --> K

---

🧩 Transformer Components

graph TD

Transformer

Transformer --> Embedding
Transformer --> PositionalEncoding
Transformer --> Encoder
Transformer --> Decoder
Transformer --> Linear

Encoder --> MultiHeadAttention
Encoder --> FeedForward
Encoder --> LayerNorm

Decoder --> MaskedAttention
Decoder --> CrossAttention
Decoder --> FeedForward2
Decoder --> LayerNorm2

---

⚙️ Encoder Block

Each encoder layer consists of:

Input
   │
   ▼
Multi-Head Self Attention
   │
Add & LayerNorm
   │
Feed Forward Network
   │
Add & LayerNorm
   │
Output

---

⚙️ Decoder Block

Each decoder layer consists of:

Input
   │
   ▼
Masked Multi-Head Attention
   │
Add & LayerNorm
   │
Cross Attention
   │
Add & LayerNorm
   │
Feed Forward Network
   │
Add & LayerNorm
   │
Output

---

📂 Project Structure

transformer-from-scratch/

├── model.py
├── encoder.py
├── decoder.py
├── attention.py
├── positional_encoding.py
├── config.py
├── train.py
├── inference.py
├── README.md
│
└── notebooks/

---

⚡ Model Configuration

Hyperparameter	Value
Encoder Layers	6
Decoder Layers	6
Attention Heads	8
Embedding Size	512
Feed Forward Size	2048
Maximum Sequence Length	5000

---

🚀 Quick Start

import torch
from model import Transformer

src = torch.randint(0, 10000, (64, 20))
tgt = torch.randint(0, 12000, (64, 15))

model = Transformer(
    src_vocab_size=10000,
    tgt_vocab_size=12000,
    num_heads=8,
    num_layers=6,
    emb_dim=512,
    nn_dim=2048
)

output = model(src, tgt)

print(output.shape)

Output

torch.Size([64, 15, 12000])

---

🔄 Forward Pass

sequenceDiagram

participant Source
participant Encoder
participant Decoder
participant Output

Source->>Encoder: Source Tokens

Encoder->>Encoder: Self Attention

Encoder-->>Decoder: Encoder Memory

Decoder->>Decoder: Masked Self Attention

Decoder->>Encoder: Cross Attention

Decoder->>Output: Vocabulary Logits

---

🏋️ Training

criterion = nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-4
)

---

📚 What You'll Learn

After studying this repository, you'll understand:

• Self-Attention
• Multi-Head Attention
• Positional Encoding
• Encoder Architecture
• Decoder Architecture
• Residual Connections
• Layer Normalization
• Feed Forward Networks
• Sequence-to-Sequence Modeling
• Machine Translation Pipeline

---

🚧 Future Improvements

• Greedy Decoding
• Beam Search
• Label Smoothing
• Learning Rate Scheduler
• Mixed Precision Training
• Flash Attention
• KV Cache
• Weight Sharing
• Byte Pair Encoding (BPE)
• Hugging Face Checkpoint Support
• ONNX Export

---

📄 Reference Paper

Attention Is All You Need

Ashish Vaswani et al.

NeurIPS 2017

---

⭐ Support

If this project helped you understand Transformers, consider giving it a ⭐ on GitHub.

It helps others discover the project and motivates future improvements.

---

📜 License

Released under the MIT License.