Four-Layer Transformer Trained On A CPU

June 22, 2024

An 834,644-parameter Transformer learned recognizable story structure on a CPU in 53,130 recorded iterations.

Four layers, 16 attention heads, 128-dimensional embeddings, a 128-token context, and a custom 361-token vocabulary keep the complete generative mechanism compact enough to inspect from initialization through output.

The model shape, training and inference commands, loss history, source change, and generated samples appear together, turning a machine-learning result into a reproducible engineering narrative.

Open The Complete Run

A GitHub pull request carries the training log and script in one file. It records the architecture, training command, inference command, loss progression, and generated output from the same run.

GitHub PR: 0.8M parameters model training

Raw training script and log

Model Shape

Layers
4
Attention heads
16
Key/value heads
16
Embedding dimension
128
Context length
128 tokens
Vocabulary
361 tokens
Dropout
0.15
Logged parameters
834,644 total from 832,128 decayed and 2,516 non-decayed parameters
Training device
CPU
Tokens per iteration
8,192

From Initialization To Generative Behavior

The command initializes a new model from scratch with a custom 361-token vocabulary, batch size 64, 128-token sequence length, learning rate 3e-4, weight decay 0.1, a 2,500-iteration warmup, and a target maximum of 100,000 iterations.

Training starts at loss 13.0864 with validation loss 13.0707. By step 53,130, logged values occupy roughly the 0.75 to 0.86 range as the model begins producing recognizable story structure.

The inference command loads ./outminimagic23/ckpt.pt and ./data/tok361.model, then runs with --max_new_tokens=10240, --top_k=200, and --temperature=0.8.

The Pseudo-Inverse Unembedding

The implementation computes a pseudo-inverse of the embedding matrix, uses it for unembedding, and propagates gradients through that transformation. This change makes the output projection depend directly on the learned embedding geometry.

Story Structure Emerges

The samples form characters, actions, simple causal turns, and moral-like endings. Repetition and unstable references reveal the pressure of the compact vocabulary and model shape, while the recognizable narrative structure shows what 834,644 trained parameters can organize.

Why The Run Matters

The entire path—from tensor dimensions and tokenizer size through optimization settings, checkpoint selection, loss progression, and output—fits on one engineering surface. That density makes the run useful for teaching, adaptation, and direct experimentation with architecture changes.

Use the source to reproduce the run, change the vocabulary or model geometry, and watch each decision alter the learning curve and generated structure.