Four-Layer Tiny Transformer Training Run
June 22, 2024

A complete tiny Transformer training run kept together with the mechanism that produced it.

This is a small generative Transformer run: 4 layers, 16 attention heads, 128-dimensional embeddings, 128-token context, a 361-token custom vocabulary, and about 835k logged parameters, trained on CPU past 50,000 iterations until it produced recognizable small-story text.

The useful part is not benchmark quality. The useful part is that the model shape, tokenizer size, context length, command, loss history, PR, and generated samples stay together.

Training runThis page delivers a small public training run, including the source path and sample output.

Included: 4 layers, 16 attention heads, 128-dimensional embeddings, 128-token context, 361-token vocabulary, about 835k logged parameters, CPU training past 50,000 iterations, loss logs, command, PR, and generated samples.

Article focus: one complete small-model training run: model shape, vocabulary, command, loss history, generated samples, and source path.

The Artifact

The public artifact is a GitHub pull request that adds a single training log/script file. It records the model shape, training command, inference command, sample output, and enough loss-log history to show the run moving from initialization into a working tiny-story generator.

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

Training Run

The training 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, warmup over 2,500 iterations, and a target maximum of 100,000 iterations.

The log starts at train loss 13.0864 and validation loss 13.0707. The artifact includes excerpts through step 53,130, where the logged loss is already in the rough 0.75 to 0.86 range.

The inference command in the artifact uses checkpoint ./outminimagic23/ckpt.pt, tokenizer ./data/tok361.model, --max_new_tokens=10240, --top_k=200, and --temperature=0.8.

The Tweak

The PR body notes one LLaMA tweak: the embedding matrix has a computed pseudo-inverse used as the unembedding, and that unembedding is backpropagated. That note should be treated as an implementation detail of the artifact. A later article can explain the linear-algebra and training implications if the supporting notes are published.

Generated Text

The sample output is not clean prose. It repeats names, loses references, and breaks grammar. But it also has recognizable story structure: characters, actions, dialogue, simple causal turns, and moral-like endings. That is why this tiny run is worth preserving as a public artifact. The setup is small enough to keep the mechanism and the result on the same page, yet large enough to show emergent generative behavior.

Run Boundary

A run this small is not a replacement for a useful production language model. It is useful because the architecture, tokenizer size, context length, parameter tensors, training command, loss log, and sample behavior all fit into one small public run record.

Next Reproduction Step

The next useful addition would be the exact data snapshot, environment versions, and a fresh rerun command that reproduces the same loss range and sample behavior from a clean checkout.