MiniGPT
Next-token char corpus (vocab 6, period-8 pattern)
- Train
- 97.7%
- Held-out
- 94.5%
The LyceumAI framework
A complete neural-network stack written from scratch in Mojo 1.0 — tensor backend, arena autograd, CNN and Transformer layers, optimizers, trainer — verified against numeric gradients, then trained until it generalizes arithmetic it was never shown.
Modules
9
Grad checks
23
Grokking, reproduced honestly
MLP K=14 · hidden=128 · AdamW lr=1e-3 · full-batch
Train on all-but-diagonal pairs of (a+b) mod 14; hold out the 14 pairs where a == b. The model first memorizes (train acc 1.0 by epoch ~1000), then — thousands of epochs later — the held-out accuracy snaps to 1.0. Both weight decays cross, faster with stronger decay, matching the reference ordering.
Grokking — wd=0.3
Held-out accuracy over the 14 never-trained diagonal pairs.
held-out (diagonal)train33 min wall, 30k epochs
Grokking — wd=1
Held-out accuracy over the 14 never-trained diagonal pairs.
held-out (diagonal)train32 min wall, 30k epochs
train pairs
182
held-out pairs
14
parameters
5518
throughput
15.4 ep/s
Every gradient verified numerically
Each op's analytic backward is checked against central differences at init time. Gate: max |analytic − numeric| < 1e-6. Measured worst across all 23 checks: 2.4e-8, most under 2e-10 — about four orders of magnitude inside the bar.
| Layer / op | max |Δ| |
|---|---|
| Linear | ≤ 1.7e−10 |
| ReLU / GELU / Sigmoid / Tanh | ≤ 1.7e−10 |
| Dropout (train path) | ≤ 1.7e−10 |
| BatchNorm1d (train + eval) | ≤ 1.7e−10 |
| Conv2D (im2col fwd / col2im bwd) | ≤ 1.7e−10 |
| MaxPool2D (max-index backward) | ≤ 1.7e−10 |
| Embedding (scatter accumulation) | ≤ 1.2e−10 |
| LayerNorm (gamma/beta/x) | ≤ 1.2e−10 |
| Causal SDPA (Q/K/V) | ≤ 1.7e−10 |
| MultiHeadAttention | ≤ 1.7e−10 |
| TransformerBlock | ≤ 1.7e−10 |
| CrossEntropy / MSE / BCE ops | ≤ 2.4e−8 |
Three architectures, three fit gates
SimpleCNN
Stripes 2D pattern classification
- Train
- 97.8%
- Held-out
- 80.0%
MLP
Grokking memorization phase (182 pairs)
- Train
- 100.0%
- Held-out
- 100.0%
The load-bearing rules
Arena autograd
Every op appends nodes to one Graph; grads live only in graph.nodes[id].grad. One epoch = reset + re-register params at stable ids, sum losses in-graph, ONE backward, ONE optimizer step — exact full-batch gradients (verified to 1e-12).
Sequences travel flattened
No 3D tensors. A sequence is one [1, T*D] row, position-major; attention slices per-position blocks out of projected rows, and Conv2D uses the same trick for flattened-NCHW images.
Every layer carries a numeric gradient check
Analytic vs central-difference < 1e-6 gate with printed actual diffs and abort() on failure — release-mode assert is a no-op in Mojo 1.0, so process exit codes are the only honest CI signal.
| ml/tensor.mojo | Row-major Tensor + matmul/broadcast/reduce ops |
| ml/autograd.mojo | Dynamic-graph arena: Value/Graph, ~20 op builders, backward dispatch |
| ml/optim.mojo | sgd / momentum / adam / adamw (decoupled decay) |
| ml/layers.mojo | Stage A dense, Stage B conv, Stage C transformer |
| ml/loss.mojo | Fused stable softmax+CE, MSE, BCE with grad checks |
| ml/data.mojo | Synthetic tasks + deterministic LCG DataLoader |
| ml/models.mojo | MLP, SimpleCNN, MiniGPT + fit gates |
| ml/trainer.mojo | NNBase trait, generic fit/predict/evaluate |
Verify it yourself
Every suite exits nonzero via abort() on any failure. The full grokking reproduction takes about an hour for both decays.
cd mojo-model
mojo run ml/layers.mojo # all layer grad checks (~2 min)
mojo run ml/models.mojo # fit gates incl. MiniGPT (~5 min)
mojo build grok_full.mojo -o /tmp/grok_full && /tmp/grok_full