YAML Metadata Warning:empty or missing yaml metadata in repo card

Check out the documentation for more information.

πŸš€ 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.

Downloads last month
4
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support