Shivam3002/OxMini-v2

OxMini v2

OxMini v2 is a 2.75M-parameter, CPU-first educational language model. It combines a gated delta-rule recurrent attention approximation (KDA-lite), low-rank latent causal attention (MLA-lite), four residual streams with Sinkhorn-constrained mHC-style routing, SwiGLU blocks, and cached decoding.

V2 replaces v1's Shakespeare-only training with a 3M-character TinyStories base run and revision-pinned, response-only Wikipedia post-training with TinyStories replay. It is measurably better at the held-out instruction format, but it is still far too small for coherent or factual long-form writing.

Measured release snapshot

Property	Value
Parameters	2,752,568
Context	320 characters
Vocabulary	117 characters, printable ASCII included
Layers	6 KDA-lite + 2 MLA-lite
Width / heads	144 / 6
Residual streams	4, Sinkhorn-routed
TinyStories test CE / PPL / accuracy	1.2885 / 3.6272 / 61.54%
Held-out SFT CE / PPL / accuracy	2.0018 / 7.4027 / 45.57%
Cached decoding	about 200 chars/s in measured qualitative runs

See V2_RESULTS.md for the base, pure-SFT, and replay-SFT comparison. The release is the replay checkpoint because pure SFT improved instruction loss by slightly more but degraded TinyStories perplexity to 5.4063.

Load and generate

git clone https://huggingface.co/Shivam3002/OxMini-v2
cd OxMini-v2
python -m venv .venv
source .venv/bin/activate
pip install .

from huggingface_hub import snapshot_download
from oxmini import CharTokenizer, OxMiniForCausalLM
from oxmini.generate import generate_text

repo = snapshot_download("Shivam3002/OxMini-v2")
model = OxMiniForCausalLM.from_pretrained(repo)
tokenizer = CharTokenizer.load(f"{repo}/tokenizer.json")

prompt = "User: Write a short story about a helpful robot.\nAssistant:"
print(
    generate_text(
        model,
        tokenizer,
        prompt,
        max_new_tokens=300,
        temperature=0.7,
        top_k=40,
        seed=20260828,
    )
)

The package is plain PyTorch plus safetensors. It does not need trust_remote_code, and no pickle checkpoint is published.

Architecture

• KDA-lite (6 layers): NoPE, gated delta-rule associative memory with a per-head finite-norm bound. This is a transparent approximation, not the optimized Kimi Linear kernel.
• MLA-lite (2 layers): shared low-rank K/V latent, causal global attention, and RoPE on Q/K.
• mHC-lite: four parallel residual streams, pre/post mixing, and Sinkhorn projection toward doubly stochastic routing matrices.
• Decoder cache: persistent KDA associative states and cropped MLA K/V histories. Cached logits match full causal forward logits in tests.

Training and post-training

Base training used 1,200 CPU updates over a deterministic 90/5/5 split of 3,000,000 TinyStories characters. Post-training used 40 source-backed English Wikipedia leads for response-only SFT and 10 held-out topics for selection. Each update also replayed TinyStories with weight 0.5 to reduce catastrophic forgetting. Exact data hashes, revisions, and commands are in V2_RESULTS.md, posttraining_data/manifest.json, and the repository scripts.

The Wikipedia excerpts remain CC BY-SA 4.0 and retain per-row attribution; see POSTTRAINING_DATA_LICENSE.md.

Limitations and safety

OxMini v2 is not a capable assistant. Its unedited outputs contain malformed words, repetition, topic drift, and fabricated statements. It cannot reliably write an essay, answer factual questions, follow complex instructions, or sustain a 1,000-word argument. Character tokenization and the 320-character training context are severe constraints. Evaluation is small, English-only, single-seed, and not a standardized safety or capability benchmark.

Use it for architecture study, CPU training experiments, caching tests, and small ablations only. Do not use it for factual, medical, legal, financial, safety-critical, production, or user-facing decisions.

References

• Eldan and Li, TinyStories.
• Kimi Team et al., Kimi Linear.
• DeepSeek-AI, DeepSeek-V2 and mHC.
• English Wikipedia contributors, revision-pinned sources listed in posttraining_data/manifest.json.