Aadit-032/LiteGPT-16M

Model Architecture

The goal of LiteGPT-16M is not to achieve state-of-the-art performance, but to provide a clean and understandable implementation of a GPT-style language model that can be trained from scratch and extended with modern techniques in future experiments.

Overview

• —Model type: Decoder-only Transformer
• —Parameters: ~16M
• —Context length: 128
• —Vocabulary size: 50,257
• —Attention: Causal Self-Attention
• —Positional Encoding: Learned Position Embeddings

Architecture Diagram

Input Tokens [B, T]
        │
        ▼
┌─────────────────────┐
│  Token Embeddings   │
│  [vocab, d_model]   │
└─────────────────────┘
        │
        ├──────────────┐
        ▼              │
┌─────────────────────┐│
│ Position Embeddings ││
│ [seq_len, d_model]  ││
└─────────────────────┘│
        │              │
        └──────┬───────┘
               ▼
          x = tok + pos
               │
               ▼
╔══════════════════════════════╗
║ Transformer Block × 4        ║
║                              ║
║       LayerNorm              ║
║           │                  ║
║           ▼                  ║
║     Multi-Head Attention     ║
║           │                  ║
║           ▼                  ║
║       Residual Add           ║
║           │                  ║
║           ▼                  ║
║       LayerNorm              ║
║           │                  ║
║           ▼                  ║
║          FFN                 ║
║           │                  ║
║           ▼                  ║
║      Residual Add            ║
╚══════════════════════════════╝
            │
            ▼
┌─────────────────────┐
│   Final LayerNorm   │
└─────────────────────┘
            │
            ▼
┌─────────────────────┐
│      LM Head        │
└─────────────────────┘
            │
            ▼
      Logits [B,T,V]

Configuration

Transformer Block

Attention

• —Multi-Head Self Attention
• —Causal Masking

Feed Forward Network

FFN(x) = W2(GELU(W1(x)))

Expansion ratio: 4×

Residual Connections

x = x + Attention(x)

x = x + FFN(x)

Normalization

• —LayerNorm

Parameter Count

Design Decisions

This model is intentionally kept as close to GPT-2 as possible to build a strong understanding of decoder-only transformers before introducing modern architectural improvements.

GPT-2 Baseline

The model uses:

• —Learned token embeddings
• —Learned positional embeddings
• —Multi-Head Self Attention (MHSA)
• —GELU activations
• —LayerNorm
• —Causal masking

Simplicity Over Performance

Features such as RoPE, GQA, FlashAttention, SwiGLU, RMSNorm, and Mixture-of-Experts are intentionally omitted. While these improve efficiency or performance, they add implementation complexity and make it harder to study the core transformer architecture.

Small Scale Training

The model is designed to train on a single NVIDIA T4 GPU using Google Colab. Model size, context length, and batch size are chosen to fit within limited compute resources.

Dataset

shakespeare.txt
       │
       ▼
┌──────────────────┐
│ GPT-2 Tokenizer  │
│ (tiktoken)       │
└────────┬─────────┘
         │
         ▼
   Token IDs
         │
         ▼
┌──────────────────┐
│ 90/10 Split      │
│ Train / Val      │
└────────┬─────────┘
         │
    ┌────┴────┐
    ▼         ▼
train.bin  val.bin
(uint16)   (uint16)

Tokenizer

• —Type: tiktoken
• —Encoding: gpt2
• —Vocabulary size: 50,257

Storage Format

train.bin

val.bin

dtype = uint16

Training

Hyperparameters

Results

Final Results