miniKimiK3 — 50M

A 50M parameter language model written from scratch in plain PyTorch, combining Kimi Delta Attention, Multi-head Latent Attention, a latent Mixture-of-Experts with auxiliary-loss-free load balancing, hyper-connection block residuals, and a Muon + AdamW hybrid optimizer. Trained on TinyStories on a single Apple Silicon GPU.

No Triton, no CUDA kernels, no transformers — every component is implemented directly so the math is readable end to end.

Code, full architecture notes and derivations: https://github.com/Shiveshrane/MiniKimiK3

Results

parameters	50.17M total, 36.71M active per token
validation loss	1.6429 (perplexity ≈ 5.17)
trained on	82M tokens (10,000 steps × 8,192 tokens), 1.22 epochs of 67M
hardware	one Apple Silicon GPU via MPS, 16.1 hours
throughput	~2,300 tokens/s training, ~13 tokens/s single-stream decode

Samples

Prompt: Once upon a time there was a little girl who, temperature 0.8, top-k 50.

Once upon a time there was a little girl who loved to go to the park. One day she was walking with her mom and she saw a big tree. She wanted to climb the tree, but her mom said no.

The little girl was sad, but she kept walking. Suddenly, she saw a man walking towards the tree. He was very big and he had a bag in his hand.

The little girl was scared, so she ran away. But then she heard a voice in the tree. It was the voice of someone yelling.

Once upon a time there was a little girl who liked to take a bath. Her mommy would take her in the bathtub and filled it with warm water. The little girl liked to splash around in the warm water.

One day, the little girl was feeling very miserable. She had no water, so she lay down in the bathtub and cried out loud.

The mommy said, "What's wrong?"

The little girl said, "I'm so miserable. I don't want to get out of the bathtub."

Architecture

13 layers, 4n+1 rule    9 KDA + 4 MLA (layer 12 forced to MLA)
hidden 512              8 heads × 64
KV / MoE latent 128     4× compression
MoE per layer           16 routed (top-4) + 2 shared, StiGLU experts
vocab 8192              byte-level BPE, tied input/output embedding

• Kimi Delta Attention — linear attention whose state update is one step of gradient descent on $\lVert S^\top k - v\rVert^2$, giving an error-correcting write instead of blind accumulation, with per-channel gated decay. Trains through a chunked parallel form; decodes as a recurrence with a fixed $64\times64$ state per head, so cost per token is constant in sequence length.
• Multi-head Latent Attention every fourth layer, for exact long-range lookups a fixed-size state cannot hold.
• Latent MoE — routed experts operate inside a 128-wide latent rather than the 512-wide model dimension, which is what makes 16 experts per layer affordable at this scale.
• Quantile balancing — per-expert routing bias solved as a fixed point, so expert load is equalized with no auxiliary loss term and no loss weight to tune.
• Muon for the 2-D hidden matrices (Newton-Schulz orthogonalized momentum, per-head on q/k/v projections), AdamW for embeddings, norms and biases.

Files

file
best.pt	checkpoint: weights, both optimizers' state, step, val loss
tokenizer.json	the byte-level BPE this checkpoint was trained with — ids are meaningless without it

Usage

git clone https://github.com/Shiveshrane/MiniKimiK3
cd MiniKimiK3

hf download ItsProtesilaus/MiniKimiK3 best.pt --local-dir checkpoints
hf download ItsProtesilaus/MiniKimiK3 tokenizer.json --local-dir data

python3 tests/generate.py "Once upon a time" --ckpt checkpoints/best.pt

The architecture is reconstructed from the checkpoint's stored config, so no flags need to match by hand.

Limitations

• Trained only on TinyStories — simple-vocabulary children's stories. It has no world knowledge, no instruction following, no chat behaviour, and will not perform on any general benchmark.
• 82M training tokens for 50M parameters is well under a compute-optimal budget (~20 tokens/parameter would be 1B), so it is undertrained.
• No positional encoding of any kind: order information reaches the model only through KDA's decay and short convolutions plus the causal mask.
• generate is single-sequence; batched sampling needs a rework.
• Research and educational use. Not evaluated for safety, bias, or factuality.

Training details

corpus	300k TinyStories → 66.99M train / 0.38M val tokens, uint16
batch	8 × 4 accumulation × 256 tokens = 8,192 tokens/step
optimizer	Muon (lr 0.02) + AdamW (lr 3e-3), weight decay 0.1, grad clip 1.0
schedule	300 warmup steps, cosine to 10% over 10,000 steps
init	$\mathcal{N}(0, 0.02)$ — required with tied embeddings, which would otherwise start at loss ≈ 494 instead of $\ln 8192 \approx 9.0$