transformer-pretraining-from-scratch
A ~57M-non-embedding-parameter decoder-only transformer, pretrained from scratch in
PyTorch. Every nn.Module (RMSNorm, rotary position embeddings, causal self-attention,
SwiGLU MLP, weight-tied embedding/output head) is authored in the training repo: 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.
Full training code, tokenizer, evaluation, calibration study, interpretability pass, and writeups: GitHub repo.
Try it live, no install needed: the model runs entirely in your browser via ONNX Runtime Web.
Why this exists
This model's companion project, a cost-aware LLM router, calls a pretrained transformer through a high-level API and asks when to trust its confidence. This project builds the model itself and asks where that confidence actually comes from, mechanistically: attention patterns, calibration behavior, what a causal ablation reveals about which heads matter. It isn't trying to be state of the art, instruction-tuned, or chat-capable.
Usage
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "hari-8/transformer-pretraining-from-scratch"
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(repo)
ids = tokenizer("The history of the Roman Empire", return_tensors="pt").input_ids
out = model.generate_simple(ids, max_new_tokens=40, temperature=0.8, top_k=40)
print(tokenizer.decode(out[0].tolist()))
trust_remote_code=True is required: the architecture (modeling_decoder_transformer.py)
ships in this repo rather than being a stock transformers model class.
model.generate_simple(...) is a minimal temperature/top-k sampling loop, not HF's
KV-cached .generate(). It recomputes the full sequence each step, which is fine for
short demo completions but not optimized for long generations or serving at scale.
Architecture
| Layers | 8 |
| Attention heads | 12 |
| Embedding dim | 768 |
| Context length | 1024 |
| Vocab size | 16,384 (custom byte-level BPE) |
| Non-embedding params | ~56.6M |
| Position encoding | RoPE |
| Normalization | RMSNorm |
| MLP | SwiGLU |
| Attention | causal, F.scaled_dot_product_attention |
| Embedding/output head | weight-tied |
Training
- ~1B tokens, FineWeb-Edu sample
- Single NVIDIA L4 GPU, bf16,
torch.compile, gradient accumulation - 7,630 steps, cosine LR schedule with warmup
- Final training loss: 3.16
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's actually comparable across models with different tokenizers, since it normalizes by the byte length of the original text rather than by token count.
| loss (nats) | perplexity | BPB | |
|---|---|---|---|
| this model | 3.19 | 24.2 | 1.057 |
| Pythia-70M (reference) | 3.68 | 39.7 | 1.135 |
This model scores lower BPB than Pythia-70M on this held-out set. The likely reason is
domain match rather than 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. Full methodology: src/eval/ in the
GitHub repo.
Calibration
Token-level calibration (confidence is the max softmax probability, the label is whether that top-1 prediction matches the actual next token, the standard Guo et al. 2017 definition), measured on ~248K held-out tokens disjoint from the temperature-fitting split.
| ECE (5 bins) | Brier | |
|---|---|---|
| raw (T=1) | 0.0071 | 0.157 |
| temperature-scaled (T=1.014) | 0.0047 | 0.157 |
The model's raw confidence is already close to perfectly calibrated: the fitted
temperature is barely different from 1. Full write-up: results/calibration_report.md in
the GitHub repo.
Interpretability
Induction-head probing (Olsson et al. 2022 methodology) found two clear induction heads
concentrated in layers 6-7, with prefix-matching scores of 0.79 and 0.74, far above every
other head. A full causal ablation sweep across every head in every layer found that these
aren't necessarily the most causally important heads on natural text, which points to a
real limit of attention-pattern-only probing. Full write-up:
results/interpretability_report.md in the GitHub repo.
Limitations
- Not instruction-tuned. This is a raw next-token-prediction language model, not a chat assistant, so expect continuation-style completions rather than question-answering or dialogue.
- Small (~57M params) and trained on ~1B tokens, so short-range text stays coherent but long-form factual accuracy or reasoning isn't reliable.
generate_simple's decoding has no KV cache, so it's O(n²) and slow for long generations.- Custom architecture, so it requires
trust_remote_code=Trueto load.
Citation
If this is useful, please link back to the GitHub repo instead of citing it formally. This is a learning project, not a paper.