Transformer Decoder Implementation using PyTorch

A modular implementation of a Transformer Decoder developed using PyTorch to understand the fundamental architecture behind modern decoder-only Large Language Models (LLMs).

---

Project Overview

This project focuses on implementing the core components of a Transformer Decoder using PyTorch in a clean and modular way. The implementation is designed to provide a practical understanding of how decoder-only architectures process sequential data and generate contextual representations.

Instead of relying on high-level transformer libraries, each major component of the decoder architecture is implemented as an independent module. This approach helps in understanding how information flows through the model, from token embeddings to the final output predictions.

The project demonstrates the key building blocks used in modern Large Language Models (LLMs), including Self-Attention, Multi-Head Attention, Positional Encoding, Feed Forward Networks, Residual Connections, and Layer Normalization. These concepts form the foundation of decoder-based models such as Llama, Mistral, Gemma, Qwen, Phi, DeepSeek, and GPT-style architectures.

The codebase follows a modular software engineering approach, where every component is organized into a separate Python module. This improves readability, maintainability, testing, and future scalability while providing a clear understanding of the complete decoder pipeline.

Although this project is intentionally lightweight for learning purposes, its overall architecture closely reflects the design principles used in modern transformer-based language models and provides a strong foundation for exploring advanced LLM concepts in future projects.

This implementation is intended for educational purposes and serves as a foundation for understanding the internal architecture of Transformer-based language models before working with full-scale LLM frameworks.

---

Project Objectives

The primary objectives of this project are:

• Understand the internal architecture of a Transformer Decoder.
• Implement Token Embeddings and Positional Encoding.
• Build Multi-Head Self-Attention with Causal Masking.
• Implement the Feed Forward Network (FFN).
• Apply Residual Connections and Layer Normalization.
• Stack multiple Transformer Decoder Blocks to build a mini decoder-only language model.
• Verify the implementation using a simple inference test.

---

Project Structure

Transformer-Decoder/
│
├── config.py
├── embeddings.py
├── attention.py
├── feed_forward.py
├── decoder_block.py
├── model.py
├── test_model.py
│
├── requirements.txt
├── README.md
├── LICENSE
└── .gitignore

---

Features

• Modular Transformer Decoder implementation
• Token Embeddings
• Sinusoidal Positional Encoding
• Multi-Head Self-Attention
• Causal (Masked) Attention
• Feed Forward Network (FFN)
• Residual Connections
• Layer Normalization
• Dropout Regularization
• Stackable Transformer Decoder Blocks
• Mini Decoder-Only Language Model
• Unit Testing

---

Architecture

Input Tokens
      │
      ▼
Token Embedding
      │
      ▼
Positional Encoding
      │
      ▼
──────────────────────────
Transformer Decoder Block
──────────────────────────
│
├── Multi-Head Self Attention
│
├── Residual Connection
│
├── Layer Normalization
│
├── Feed Forward Network
│
├── Residual Connection
│
└── Layer Normalization
      │
      ▼
Multiple Decoder Blocks
      │
      ▼
Final Layer Normalization
      │
      ▼
Linear Projection
      │
      ▼
Vocabulary Logits

---

Modules

• config.py – Stores all model configuration parameters.
• embeddings.py – Implements token embeddings and positional encoding.
• attention.py – Implements Multi-Head Self-Attention with causal masking.
• feed_forward.py – Implements the Feed Forward Network (FFN).
• decoder_block.py – Combines attention, FFN, residual connections, and layer normalization into a Transformer Decoder Block.
• model.py – Builds a mini decoder-only language model by stacking multiple decoder blocks.
• test_model.py – Verifies that the model runs successfully and produces the expected output.

---

Concepts Covered

This implementation includes the following core Transformer concepts:

• Self Attention
• Multi-Head Attention
• Positional Encoding
• Token Embeddings
• Decoder Architecture
• Feed Forward Network (FFN)
• Residual Connections
• Layer Normalization
• Causal Masking

---

Technologies Used

• Python 3.14.0
• PyTorch 2.13.0
• NumPy
• Object-Oriented Programming (OOP)
• Transformer Architecture
• Deep Learning Fundamentals

---

Installation

This project requires Python 3.14.0 and PyTorch 2.13.0.

1. Clone the repository.

2. (Optional) Create and activate a virtual environment.

python -m venv .venv

Windows

.venv\Scripts\activate

3. Install the required dependencies.

pip install -r requirements.txt

---

Run the Project

After installing the required dependencies, run the following command to verify that the Transformer Decoder has been implemented correctly.

python test_model.py

The script performs a simple sanity test by creating a dummy batch of input tokens, passing them through the complete Transformer Decoder model, and verifying that the output tensor has the expected dimensions.

During execution, the following components are tested together:

• Token Embedding
• Positional Encoding
• Multi-Head Self-Attention
• Feed Forward Network (FFN)
• Transformer Decoder Block
• Stacked Decoder Layers
• Final Output Projection

If the implementation is correct, the script displays the input shape, output shape, and a success message confirming that the model is functioning as expected.

---

Expected Output

==================================================
Mini Transformer Decoder Test
==================================================
Input Shape : torch.Size([2, 10])
Output Shape: torch.Size([2, 10, 1000])

✅ Test Passed Successfully!
The Transformer Decoder model is working correctly.

---

Future Improvements

Possible future enhancements include:

• KV Cache
• Hugging Face Integration
• Quantization
• GPU Optimization
• Memory Profiling
• Flash Attention
• Rotary Positional Embeddings (RoPE)

---

Learning Outcome

After completing this project, you will gain a practical understanding of how a Transformer Decoder works internally and how its individual components interact to build the foundation of modern decoder-only Large Language Models.

---

Author

Ifrah Uzair

Software Engineering Student

This project was developed as part of my learning journey to understand the internal architecture of Transformer-based Large Language Models using PyTorch and strengthen my understanding of modern LLM architectures.