Simple LLM Training with GPT-2 Architecture

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.
Can add up to 4-digit numbers with 100% accuracy.
Is publicly available on Github : 🔗 Rajesh-Nair/simple-llm

🏗️ Dataset Generator

Synthetic number sequences are generated based on parameters defined in data_config.yaml.

Example Configuration:

Number range: 0 - 9999
Number of sequences: 100,000
Output path: ../simple-llm-data/sequences.txt
Delimiters: | (columns), \n (rows)

🔧 To Generate the Dataset:

Update data_config.yaml with your desired parameters.
Run the generator:
```
python3 data_generator.py
```

🎯 Training

Step 1: Train the Tokenizer

python3 tokenizer.py

Step 2: Train the Model

python3 trainer.py

Training configurations are managed in train_config.yaml, including:

🔧 Model architecture (layers, heads, embedding size)
⚙️ Training hyperparameters (batch size, learning rate)
💾 Checkpointing and saving
☁️ Hugging Face Hub integration

🔢 Position Embeddings

📐 Learnable vs. Sinusoidal Embeddings

Learnable Embeddings: Adapt to numeric patterns.
Sinusoidal Embeddings: Provide a mathematical basis for position understanding.

🧮 Block Position IDs (Abacus Embedding)

Inspired by the Abacus Embedding paper, we use block position IDs.

Example:

Input: +1342+879+2221+
Block IDs: 012340123012340

🔍 Why Block Position IDs?

✅ Commutative Support: a + b = b + a — block IDs reinforce this.
🧠 Digit Alignment: Helps align units, tens, hundreds positions for easier digit-wise processing.

🔄 Digit Reversal

As part of preprocessing:

5672 → 2765 (reversed)
Output is reversed back during evaluation.

📈 Benefits of Reversal

🧒 Human-like learning: Mimics the left-to-right addition humans use.
🎯 Causal attention compatibility: Enables better carryover handling.
📚 Research-backed approach: Digit reversal has been successfully used in several papers including:
- Transformers Can Do Arithmetic with the Right Embeddings (which also introduces Abacus embedding)
- Transformers Can Achieve Length Generalization But Not Robustly

🧩 Tokenization Strategy

Tokenization is critical for arithmetic modeling. Our approach:

📏 Shortens sequences: Optimizes context window usage.
🧬 Boosts generalization: Learns across number patterns.
🔄 Uses base conversion (e.g., decimal → hexadecimal) for compact, arithmetic-aware tokens.
🧠 Preserves arithmetic logic: Even in higher bases, rules still apply.

We're experimenting with different bases to improve efficiency further.

🔁 Multi-token Prediction

Predicting multiple tokens at once increases efficiency. This is possible since we have reversed all the numbers.

Example: To predict two token at a time, we see output 99 to appear in the first iteration

Input (reversed):     +12+873+PPPPPPPP      (P = padding tokens)
Output (reversed):    PPPPPP99PPPPPPPP      (P = padding tokens)
Position IDs:         0120123000000000

We're currently supporting 2-token prediction and it works well 🔗 mirajnair/simple-llm-gpt2-v2.0

..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

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
- Scale up model size for more complex operations
Length Generalization:
- Implement and test length generalization techniques as described in Transformers Can Achieve Length Generalization But Not Robustly
- Explore methods to improve model's ability to handle varying sequence lengths
Add batch prediction:
- Implement parallel processing of multiple arithmetic operations
- Optimize throughput by processing multiple sequences simultaneously
- Reduce overall inference time for bulk operations
KV cache:
- Implement key-value caching to reduce redundant computations
- Optimize memory usage during autoregressive generation
- Speed up sequential token generation by reusing previous computations