LYCEUMAI

Documentation

LyceumAI is a neural-network framework written from scratch in Mojo 1.0 — tensor backend, arena autograd, dense/CNN/Transformer layers, optimizers, trainer — plus the models trained with it, including a full reproduction of grokking. This page is the project's official documentation.

9 modules23 gradient checks0 dependencies60k epochs verified

How the project thinks

Three rules that shape every file in mojo-model/ml.

Verify, then trust

Every op's analytic backward is checked against central differences at init time. Gate: max |analytic − numeric| < 1e-6. All 23 checks pass at ≤ 2.4e-8. If a check fails, the process aborts — exit codes are the only honest CI signal (release-mode assert is a no-op in Mojo).

No dependencies, no shortcuts

ml/ imports nothing but the Mojo standard library. Tensors are plain row-major Lists of Float64; matmul is a triple loop. The point is that every line of the learning machinery is readable and owned.

Honest gates over vibes

Fit runs end in explicit thresholds (train acc > 0.9, diag acc == 1.0) with printed actual values. A run either passes its gate or aborts — there is no 'looks about right' path through the codebase.

The framework, module by module

Everything lives under mojo-model/ml/. Import as from ml.<module> import <symbol>; the full API reference is ml/docs/API.md in the repo.

  • ml/tensor.mojo

    Tensor struct (row-major List[Float64]) plus the op vocabulary: matmul, broadcast_add, reductions, reshape, transpose, elementwise math, add_inplace for grad accumulation.

  • ml/autograd.mojo

    The arena. Value(id/data/grad/parents/op) nodes appended to one Graph; ~20 op builders (g_linear, g_softmax_ce, g_conv2d, g_sdpa_h, …); backward() dispatches on the op string and accumulates into node grads in place.

  • ml/optim.mojo

    sgd_step, momentum_step, adam_step, adamw_step (decoupled weight decay), each operating on arena grads with per-parameter state structs.

  • ml/layers.mojo

    Stage A dense (Linear, activations, Dropout, BatchNorm1d), Stage B conv (Conv2D im2col/col2im, MaxPool2D max-index), Stage C transformer (Embedding, LayerNorm, causal SDPA, MHA, TransformerBlock).

  • ml/loss.mojo

    cross_entropy as one fused stable softmax+CE autograd op, mse, binary_cross_entropy with consistent eps clamping in value and grad.

  • ml/data.mojo

    Synthetic tasks (modular arithmetic pairs, parity, polynomial/sin regression) plus Dataset/DataLoader with deterministic LCG shuffling and train_test_split.

  • ml/models.mojo

    MLP, SimpleCNN, MiniGPT structs, the grokking recipe (fit_full_batch on all-but-diagonal pairs), reregister_params lifecycle, and the fit gates.

  • ml/trainer.mojo

    The NNBase trait and generic fit/predict/evaluate that work over ANY conforming model, optimizer dispatch by string, TrainingConfig.

One training epoch

The same four-phase lifecycle runs every model in the repo, from a 114-parameter CNN to the grokking MLP.

  1. 1 · RegisterLayer params are registered in the arena as leaf Values at stable ids 0..P-1. reregister_params snapshots weights, resets the graph, and re-registers so ids never drift between epochs.
  2. 2 · ForwardEvery op appends nodes to the Graph. Parameterized layers re-sync their weights from the arena before each forward; leaf inputs read their own data. Sequences travel as flattened [1, T*D] rows.
  3. 3 · Loss in-graphPer-sample losses are summed INSIDE the graph with g_add chains, so ONE backward(seed) produces exact full-batch gradients — verified equal to per-sample accumulation within 1e-12.
  4. 4 · StepOne optimizer step applies the accumulated grads from graph.nodes[id].grad, then the next epoch resets the arena and starts over. The arena grows ~5 nodes per forward; reset keeps memory flat.

The models

Three architectures, three fit gates, zero pretrained weights. Charts and per-layer verification numbers live on the Mojo page.

MLP — the grokking vehicle

5,518

28→128→14 MLP trained on all 182 off-diagonal (a+b) mod 14 pairs with AdamW lr=1e-3. It memorizes by epoch ~1000; held-out accuracy on the 14 diagonal pairs (a == b, never trained) snaps to 1.0 at epoch 9000 with wd=1.0 and epoch 25000 with wd=0.3 — the classic grokking transition, reproduced honestly at full 30k epochs.

SimpleCNN — 2D structure

114

Conv2D(im2col forward, col2im backward) → ReLU → MaxPool → Flatten → Linear on flattened-NCHW input. Fit gate: stripes-task train accuracy 0.978 > 0.9, held-out 0.80.

MiniGPT — attention that pays rent

26,182

Embedding + learned positional table + 2 TransformerBlocks (post-LN, causal MHA) + vocab head via seq_linear. Fits a deterministic period-8 char grammar over vocab 6: train 97.7%, held-out 94.5% — through real attention, not a lookup shortcut.

Grokking curves, gradient-check table, and fit results — with the actual numbers from the run logs.Open results

Run it yourself

Requires Mojo 1.0 (python3 -m pip install --break-system-packages mojo). Every command exits nonzero via abort() on any failure.

Verification commands
mojo run ml/layers.mojoAll Stage A+B+C layer gradient checks (~2 min)
mojo run ml/models.mojoUnit checks + smoke/grok/CNN/MiniGPT fit gates (~5 min)
mojo run framework_tests.mojoCross-module integration harness
mojo build grok_full.mojo -o /tmp/grok_full && /tmp/grok_fullFull 30k-epoch grokking, both decays (~65 min)

About this site

The hub you are reading is itself part of the repo: Next.js App Router + React 19, Tailwind v4 tokens, GSAP scroll motion, shadcn/ui primitives. No mock registry data anywhere — every model card and metric renders numbers produced by the framework's own verification gates. Repo docs: README.md and ml/docs/API.md.