| --- |
| library_name: pytorch |
| pipeline_tag: text-generation |
| datasets: |
| - roneneldan/TinyStories |
| language: |
| - en |
| tags: |
| - linear-attention |
| - mixture-of-experts |
| - delta-rule |
| - muon |
| - from-scratch |
| --- |
| |
| # 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 |
|
|
| ```bash |
| 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$ | |
|
|
