A small GPT-style language model built from scratch in PyTorch, designed to study how a language model learns arithmetic word problems and how external tools can improve its reliability.
This project intentionally uses a small architecture that can be trained on a CPU, making the complete training pipeline accessible without large-scale GPU infrastructure.
Large language models can perform arithmetic surprisingly well, but language models are fundamentally probabilistic sequence predictors rather than deterministic calculators.
This project explores a simple question:
Can a small language model become more reliable at arithmetic by delegating computation to an external tool?
We approach the problem in two stages:
Direct Model
Word Problem
↓
Tiny GPT
↓
Answer
versus:
Tool-Augmented Model
Word Problem
↓
Tiny GPT
↓
Arithmetic Expression
↓
Calculator
↓
Answer
The project also investigates how well the model generalizes when the wording and numerical distribution change.
The model is a small GPT-style Transformer implemented using PyTorch.
Input Tokens
│
▼
Token Embeddings
│
▼
Positional Information
│
▼
Transformer Block
│
├── LayerNorm
├── Multi-Head Causal Self-Attention
├── Residual Connection
│
├── LayerNorm
├── Feed-Forward Network
└── Residual Connection
│
▼
Transformer Block
│
▼
Language Model Head
│
▼
Token Probabilities
The model is autoregressive and uses causal masking so that each token can only attend to previous tokens.
A lightweight character-level tokenizer was implemented from scratch.
The vocabulary contains:
- Uppercase and lowercase letters
- Digits
- Common punctuation
- Arithmetic operators
- Special control tokens
Special tokens include:
<PAD>
<BOS>
<EOS>
<QUESTION>
<ANSWER>
<TOOL_CALL>
</TOOL_CALL>
<TOOL_RESULT>
The vocabulary is fixed rather than being derived from the training corpus, preventing inference failures caused by previously unseen characters.
The attention mechanism implements:
Q = XWq
K = XWk
V = XWv
followed by scaled dot-product attention:
Attention(Q,K,V)
=
softmax(QKᵀ / √dₖ)V
Causal masking prevents the model from accessing future tokens during training.
Each Transformer block contains:
LayerNorm
↓
Multi-Head Self-Attention
↓
Residual Connection
↓
LayerNorm
↓
Feed-Forward Network
↓
Residual Connection
This provides the core architecture required for autoregressive language modeling.
The initial dataset consists of synthetically generated arithmetic word problems involving:
- Addition
- Subtraction
- Multiplication
- Division
Example:
A box contains 1 pencils.
Another 47 pencils are placed inside.
What is the total number of pencils?
Associated representation:
Expression: 1 + 47
Answer: 48
The synthetic dataset allows controlled experimentation with:
- Number ranges
- Operations
- Sentence templates
- Training/validation/test splits
- Difficulty
- Generalization
The model is trained using next-token prediction.
Given:
X = [token₁, token₂, token₃, ...]
the model learns to predict:
Y = [token₂, token₃, token₄, ...]
using cross-entropy loss.
One of the key experiments in V1 was changing the training objective.
Initially, the model was penalized for predicting every token in the sequence.
For example:
<BOS> <QUESTION> ... <ANSWER> 75 <EOS>
↑ ↑ ↑ ↑
loss loss loss loss
This resulted in:
31.30% exact-answer accuracy
The training objective was then changed so that loss was concentrated on the answer portion of the sequence.
Conceptually:
<BOS> <QUESTION> ... <ANSWER> 75 <EOS>
X X ↑ ↑ ↑
ignored ignored loss loss loss
After this change:
75.40% exact-answer accuracy
This demonstrated the impact of aligning the training objective with the actual task.
The direct-answer model was evaluated on 2,000 held-out synthetic arithmetic problems.
| Metric | Result |
|---|---|
| Exact Accuracy | 75.40% |
| Mean Absolute Error | 1.26 |
| MAE on Incorrect Answers | 5.13 |
| Operation | Accuracy |
|---|---|
| Addition | 74.08% |
| Subtraction | 62.73% |
| Multiplication | 88.16% |
| Division | 78.70% |
| Answer Length | Accuracy |
|---|---|
| 1 digit | 80.96% |
| 2 digits | 72.41% |
| 3 digits | 78.16% |
The results show that the model learned useful arithmetic patterns but still produced approximate rather than deterministic arithmetic.
Example:
Expected: 26 - 18 = 8
Predicted: 10
This motivated the tool-use experiment.
Instead of requiring the language model to perform the arithmetic itself, the model was trained to generate an arithmetic expression.
For example:
Question:
Ravi has 26 apples.
He gives 18 apples to his friend.
How many apples remain?
The model generates:
<TOOL_CALL>26 - 18</TOOL_CALL>
The expression is then passed to a deterministic calculator:
26 - 18
↓
8
The complete pipeline is:
┌─────────────┐
│ Word Problem│
└──────┬──────┘
│
▼
┌─────────────┐
│ Tiny GPT │
└──────┬──────┘
│
▼
26 - 18
│
▼
┌─────────────┐
│ Calculator │
└──────┬──────┘
│
▼
8
On the same 2,000-example test set:
| Metric | Result |
|---|---|
| Expression Accuracy | 100.00% |
| Tool Success Rate | 100.00% |
| End-to-End Accuracy | 100.00% |
By operation:
| Operation | Accuracy |
|---|---|
| Addition | 100.00% |
| Subtraction | 100.00% |
| Multiplication | 100.00% |
| Division | 100.00% |
This represents an improvement from:
Direct-answer model
75.40%
to:
Tool-assisted model
100.00%
on the in-distribution test set.
The 100% tool-use result led to a controlled generalization experiment.
Three new test sets were generated with:
- Slightly different wording
- Moderately different wording
- Significantly different wording
The model's performance dropped dramatically.
For the first reworded test set:
Expression Accuracy: 2.20%
End-to-end Accuracy: 2.20%
The model frequently generated malformed or incorrect expressions, for example:
Expected:
98 + 98
Predicted:
9 + - 98
or:
Expected:
228 / 12
Predicted:
21
This exposed a major limitation of the V1 system:
The model had learned the training distribution and templates rather than developing robust language generalization.
This was an important result rather than simply a failure.
It motivated the development of V2 with:
- A fixed vocabulary
- More diverse language
- Larger training data
- More varied numerical distributions
- Multi-step problems
- Distractor information
- Increased model capacity
The initial model was capable of learning the task, but its performance was heavily dependent on the training distribution.
Changing the loss from full-sequence prediction to answer-focused prediction increased accuracy substantially.
The direct model frequently produced near-correct arithmetic results rather than exact results.
The tool model only needed to identify the correct operands and operation, while the calculator performed deterministic arithmetic.
The 100% tool-use result initially looked impressive, but testing on reworded examples revealed severe overfitting.
This reinforced the importance of held-out distribution testing rather than relying solely on random train/test splits.
tiny-llm/
│
├── data/
│ ├── tiny.txt
│ ├── arithmetic_train.jsonl
│ ├── arithmetic_val.jsonl
│ ├── arithmetic_test.jsonl
│ └── ...
│
├── tokenizer/
│ └── tokenizer.py
│
├── model/
│ ├── attention.py
│ ├── feedforward.py
│ ├── transformer.py
│ └── gpt.py
│
├── training/
│ └── loss.py
│
├── evaluation/
│ ├── evaluate_arith.py
│ ├── evaluate_tooluse.py
│ └── ...
│
├── inference/
│ └── generate.py
│
├── tools/
│ └── calculator.py
│
├── train.py
│
├── model.pt
├── tool_model.pt
└── README.md
Create and activate a virtual environment:
python -m venv venvWindows:
venv\Scripts\activateInstall dependencies:
pip install torchTrain the model:
python train.pyRun inference:
python -m inference.generateEvaluate the arithmetic model:
python -m evaluation.evaluate_arithEvaluate the tool-use pipeline:
python -m evaluation.evaluate_tooluseThe project was intentionally designed to run on consumer hardware.
The V1 experiments were performed on an Intel i9 CPU without requiring a dedicated GPU.
The goal is not to compete with large language models in scale, but to understand and implement the underlying components of a modern autoregressive Transformer.
The next version of the project will focus on improving generalization and multi-step reasoning.
Planned improvements:
50K–100K+ examples
with:
- Diverse linguistic templates
- Larger number ranges
- Distractor information
- Multiple contexts
- Two-step problems
- Three-step problems
- Edge cases
Increase model capacity moderately while remaining CPU-trainable.
Move from:
Question
↓
Expression
↓
Calculator
↓
Answer
to:
Question
↓
Tool Call
↓
Tool Result
↓
Tool Call
↓
Tool Result
↓
Answer
Eventually evaluate the system on datasets such as GSM8K and other mathematical reasoning benchmarks.
Completed baseline
V1 demonstrates:
- A Transformer implemented from scratch
- End-to-end training
- Validation and checkpointing
- Arithmetic reasoning
- Task-specific loss masking
- Tool-call generation
- Deterministic tool execution
- Controlled generalization testing
- Identification of overfitting
V2 will focus on making the model generalize beyond the synthetic training distribution and perform multi-step reasoning through repeated tool use.
Salil Phanse
Built as a hands-on exploration of:
Transformers
+
Language Modeling
+
Training
+
Evaluation
+
Tool Use
+
Reasoning