AxiomLLM

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:

$$c_t^{KV}=W_{DKV}{h}_t\in {\mathbb{R}}^{r_{kv}},k_{t,i}^C=W_{UK,i}{c}_t^{KV},v_{t,i}=W_{UV,i}{c}_t^{KV}c\_t^\{KV\}\; =\; W\_\{DKV\}\; h\_t\; \backslash in\; \backslash mathbb\{R\}^\{r\_\{kv\}\},\; \backslash qquad\; k\_\{t,i\}^\{C\}\; =\; W\_\{UK,i\}\backslash ,\; c\_t^\{KV\},\; \backslash qquad\; v\_\{t,i\}\; =\; W\_\{UV,i\}\backslash ,\; c\_t^\{KV\}$$

Only $c_t^{KV}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:

$$q_{t,i}=[q_{t,i}^C;\mathrm{R}\mathrm{o}\mathrm{P}\mathrm{E}(q_{t,i}^R)],k_{t,i}=[k_{t,i}^C;\mathrm{R}\mathrm{o}\mathrm{P}\mathrm{E}(k_t^R)]q\_\{t,i\}\; =\; [\backslash ,q\_\{t,i\}^\{C\}\backslash ,;\backslash ,\backslash mathrm\{RoPE\}(q\_\{t,i\}^\{R\})\backslash ,],\; \backslash qquad\; k\_\{t,i\}\; =\; [\backslash ,k\_\{t,i\}^\{C\}\backslash ,;\backslash ,\backslash mathrm\{RoPE\}(k\_t^\{R\})\backslash ,]$$

Feed-forward routes each token to 2 of 8 SwiGLU experts via a softmax router:

$$p_t=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(W_r{h}_t),{\mathcal{T}}_t={\mathrm{t}\mathrm{o}\mathrm{p}}_2(p_t),y_t=\sum _{j\in {\mathcal{T}}_t}{w}_{t,j}\cdot {\mathrm{S}\mathrm{w}\mathrm{i}\mathrm{G}\mathrm{L}\mathrm{U}}_j(h_t)p\_t\; =\; \backslash mathrm\{softmax\}(W\_r\; h\_t),\; \backslash quad\; \backslash mathcal\{T\}\_t\; =\; \backslash mathrm\{top\}\_2(p\_t),\; \backslash quad\; y\_t\; =\; \backslash sum\_\{j\; \backslash in\; \backslash mathcal\{T\}\_t\}\; w\_\{t,j\}\; \backslash cdot\; \backslash mathrm\{SwiGLU\}\_j(h\_t)$$

with a load-balancing auxiliary loss (\(\mathcal{L}_{aux} = N\sum_i f_i P_i\), where $f_{i}f\_i$ is actual routing fraction and $P_{i}P\_i$ is mean router probability for expert $ii$) added at weight 0.01 to the cross-entropy loss, to keep the router from collapsing onto 1–2 experts.

Measured results

	Standard MHA	This model
KV cache per token (floats)	1536	64 (24x smaller)

	Millions of params
Total (8 experts stored/layer)	540
Active per token (top-2 routed)	200 (2.7x less compute)

Model summary

Total parameters	540M (200M active per token)
Layers / heads / embed dim	12 / 12 / 768
MLA: kv_lora_rank / qk_nope_head_dim / qk_rope_head_dim / v_head_dim / q_lora_rank	32 / 32 / 32 / 64 / 96
MoE: experts / top-k	8 / 2
Vocabulary	50,257 (custom BPE)
Context length	1,024 tokens

Training details

Dataset	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.

import torch
ckpt = torch.load("checkpoints/axiomllm_epoch_3.pt", map_location="cpu")
print(f"epoch={ckpt['epoch']}, loss={ckpt['loss']}")

How to use

Not transformers-compatible — no config.json/AutoModel integration. Clone the code repo first:

git clone https://github.com/itsraunak-work/axiomllm.git
cd axiomllm && pip install -r requirements.txt huggingface_hub

import torch
from huggingface_hub import hf_hub_download
from src.model import AxiomLLM
from src.tokenizer import AxiomTokenizer
from src.config import load_config

REPO_ID = "itsraunak-work/axiomllm"
ckpt_path = hf_hub_download(REPO_ID, "axiomllm_epoch_3.pt")
tokenizer_path = hf_hub_download(REPO_ID, "axiom_tokenizer.json")

cfg = load_config("configs/default.yaml")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AxiomTokenizer(vocab_size=cfg.model.vocab_size, save_path=tokenizer_path)
tokenizer.load()

model = AxiomLLM(cfg.model).to(device)
checkpoint = torch.load(ckpt_path, map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

prompt = "Once upon a time"
input_ids = tokenizer.encode(prompt, add_eos=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long, device=device)

with torch.no_grad():
    for _ in range(50):
        logits, _ = model(input_tensor)
        next_logits = logits[0, -1, :] / 0.8  # temperature
        probs = torch.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        input_tensor = torch.cat([input_tensor, next_token.unsqueeze(0)], dim=1)
        if next_token.item() == tokenizer._tokenizer.token_to_id("<eos>"):
            break

print(tokenizer.decode(input_tensor[0].tolist()))

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.

License

MIT — see LICENSE in the source repo.