This repository demonstrates how to train a Language Learning Model (LLM) from scratch using the GPT-2 architecture. The model is trained on numerical sequences to learn and predict patterns.
📌 Overview
This project implements a full machine learning pipeline:
📊 Synthetic dataset generation (number sequences)
🔤 Custom tokenizer training
🧠 Model training using GPT-2
🤖 Inference capabilities
🚧 Progress So Far. We have trained a 6.4 million parameter model that:
Uses base-16 (hexadecimal) conversion for tokenization.
..And we are expanding on generalizing this method - i.e output token at the earliest opportunity so we can have 2 or more predicted in one go.
📊 Attention Visualization
Visualizing attention patterns reveals how the model processes arithmetic operations. Below is an example showing attention patterns for the addition problem: 101 + 1002 = 1103 (represented in reversed form as +101+2001+3011+).
Layer 1 Attention Patterns
Layer 1 Attention Visualization
In this visualization:
Bright vertical bars at positions 1, 5, and 10 show how the model focuses on unit digits from both inputs and the output
The model learns to align corresponding digit positions (units with units, tens with tens, etc.)
Attention patterns reveal how information flows during the addition process, including carry operations
This confirms our block position ID approach helps the model understand the commutative nature of addition and properly align digits for arithmetic operations.
The visualization demonstrates how the model has learned to focus on relevant digits when performing calculations, similar to how humans process arithmetic problems.
🎯 Performance Results
We've rigorously tested our model's arithmetic capabilities with impressive results:
Addition Performance Test
Test Set: 10,000 random pairs of 4-digit numbers
Accuracy: 100%
Consistency: Maintained perfect accuracy across multiple test runs
This perfect accuracy demonstrates that our approach successfully teaches the model to perform addition operations with complete reliability, even on previously unseen number combinations. The combination of our specialized tokenization strategy, position encoding, and multi-token prediction enables the model to generalize arithmetic rules effectively.
These results validate our architectural choices and confirm that transformer-based models can master fundamental arithmetic operations when properly designed.
🚀 Next Steps
Multi-token Generation:
We've proved the model can output more than 1 token at a time
Test if model can generate all tokens in one-go (greedy generation)
Scale Up:
Increase the length/number of digits in operations
