axiomllm

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}}, \qquad k_{t,i}^{C} = W_{UK,i}, c_t^{KV}, \qquad v_{t,i} = W_{UV,i}, c_t^{KV}$$

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:

$$q_{t,i} = [,q_{t,i}^{C},;,\mathrm{RoPE}(q_{t,i}^{R}),], \qquad k_{t,i} = [,k_{t,i}^{C},;,\mathrm{RoPE}(k_t^{R}),]$$

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

$$p_t = \mathrm{softmax}(W_r h_t), \quad \mathcal{T}_t = \mathrm{top}2(p_t), \quad y_t = \sum{j \in \mathcal{T}t} w{t,j} \cdot \mathrm{SwiGLU}_j(h_t)$$

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.

Measured results

Model summary

Training details

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.