A decoder-only transformer built from scratch in PyTorch around two design decisions: attention that caches a compressed latent instead of full-size keys and values, and a feed-forward layer that routes each token to a small subset of experts instead of running one dense MLP on every token. Full implementation and training code: github.com/itsraunak-work/axiomllm.
Research/educational checkpoint, not a general-purpose chatbot. Trained 3 epochs on a 5,000-story subset of TinyStories on a single free-tier Colab T4. Expect short, simple, TinyStories-flavored English — see Limitations.
Architecture
Attention compresses K/V through a shared low-rank latent instead of projecting full-size keys/values per head:
Only \(c_t^{KV}\) needs to be cached during generation — its size doesn't depend on the number of heads. Since RoPE doesn't commute with this compression, position is carried by a small decoupled slice instead, rotated once and shared across all heads:
with a load-balancing auxiliary loss (\(\mathcal{L}_{aux} = N\sum_i f_i P_i\), where \(f_i\) is actual routing fraction and \(P_i\) is mean router probability for expert \(i\)) added at weight 0.01 to the cross-entropy loss, to keep the router from collapsing onto 1–2 experts.
roneneldan/TinyStories, train split, first 5,000 samples
Epochs
3
Batch size
2, grad accumulation 8 (effective 16)
Learning rate
3e-4, AdamW, weight decay 0.1
Precision
bf16 autocast
Hardware
1x NVIDIA T4 (Colab free tier)
Seed
42
Final training loss
[fill in — see snippet below]
T4 is Turing-generation and lacks native BF16 tensor-core support (that arrived with Ampere); training likely ran without the acceleration bf16 is meant to provide. fp16 + GradScaler (already supported in train.py) would typically be faster on this hardware.
Or run the included REPL after downloading both files into checkpoints/ and assets/:
python scripts/chat.py
Limitations
5,000 of TinyStories' ~2.1M stories, 3 epochs — a small fraction of data relative to a 540M-parameter model. Expect underfitting and repetition, not fluent long-form generation.
TinyStories-only vocabulary and style — simple, GPT-generated children's stories. Not suited to factual questions, code, reasoning, or anything outside that register.
No instruction tuning, no RLHF, no safety alignment. Raw next-token-prediction base model.
Router balance not separately audited for this checkpoint — the auxiliary loss is implemented and included in training, but per-expert utilization at this specific checkpoint hasn't been measured post-hoc.
Temperature-only sampling in the reference generation loop — no top-k/top-p filtering.
Intended use
A working, from-scratch implementation of latent-attention KV compression and load-balanced MoE routing, trained end-to-end — useful for reading the code, studying the training loop, or as a starting checkpoint for further training on more data. Not intended for downstream deployment as-is.
