Transformer Pretrained From Scratch

A ~57M-non-embedding-parameter decoder-only transformer, pretrained from scratch in PyTorch: RMSNorm, rotary position embeddings, causal self-attention, SwiGLU MLP, and a weight-tied embedding/output head, all authored as raw nn.Modules. No AutoModel, no fine-tune of an existing checkpoint. Trained on ~1B tokens of FineWeb-Edu with a custom byte-level BPE tokenizer, also trained from scratch rather than reusing GPT-2's. The value here is depth, not scale: every claim below is measured, and every negative result is reported as one, not hidden.

0.057bits-per-byte, held-out (vs. Pythia-70M's 1.135)

0.0071raw token-level calibration ECE, near-perfect

0clear induction heads found, causally verified

0Mnon-embedding parameters, ~1B training tokens

Evaluation

Held-out loss, perplexity, and bits-per-byte (BPB), measured on ~497K held-out tokens never seen during training. BPB, not raw perplexity, is the metric that is actually comparable across models with different tokenizers, since it normalizes by the byte length of the original text rather than by token count.

This model scores lower BPB than Pythia-70M on this held-out set. The likely reason is domain match, not general capability: it was trained exclusively on FineWeb-Edu, a narrower and more predictable distribution than Pythia's Pile-trained generalist scope, and the two models are close in parameter count.

Calibration, and a genuinely surprising result

The raw confidence is already close to perfectly calibrated: the fitted temperature is barely different from 1, and temperature scaling only marginally improves an already-low ECE. That is a real, honest finding, not a failure to find something more dramatic. Models trained end to end with softmax cross-entropy on next-token prediction are documented to calibrate more naturally than classifiers trained on one-hot labels, since the training objective directly targets the true conditional distribution rather than a decision boundary.

This also connects to the Cost-Aware LLM Router's calibrator, which measured 0.1671 ECE. That is not a fair "worse" comparison: it is a different task entirely, task-level response correctness, judged, via a separately-trained logistic regression on hand-engineered features, included here for reference, not as a competing number. What is comparable is calibration behavior itself: whether raw confidence over- or under-shoots accuracy, and whether a standard post-hoc fix narrows the gap.

Interpretability: a probe, and the causal test that complicates it

Causal ablation: held-out loss delta when each head is zeroed, full sweep.

A full causal ablation sweep, every head in every layer, zeroed one at a time, measured by the resulting held-out loss increase, is what turns the induction-head numbers into an actual causal claim rather than a picture. The two heads with the highest induction score are not the most important by ablation: they rank only 8th and 11th of 12 heads in their own layer. The heads that matter most causally show no elevated induction score at all. The likely explanation: the synthetic probe targets exact token repeats, a pattern rarely seen verbatim in natural text, so a head specialized for it may matter less for real-text loss than a head doing more general work the probe was never designed to detect. A real limitation of attention-pattern-only probing, reported as a finding rather than smoothed into a cleaner story.

An honest negative result: the Triton kernel

A hand-written Triton kernel (forward and backward) for RMSNorm, correctness and gradient-verified against the PyTorch reference, then benchmarked against it at real model scale.

The Triton kernel measured 6.6% slower, not faster. torch.compile already lowers the plain PyTorch RMSNorm into a fused, autotuned Triton kernel automatically as part of compiling the surrounding graph, and wrapping a hand-written kernel in a torch.autograd.Function forces a compile graph break around it, adding dispatch overhead the compiler's own fusion does not pay. Not every hand-written kernel beats a modern compiler baseline, and that is a real, useful finding about when custom kernels are and are not worth writing, not a result to hide.

Try it

The live demo runs the model entirely client-side via ONNX Runtime Web: no server compute, so it stays free to host as a static page rather than needing a paid always-on backend.

Stack