| --- |
| license: mit |
| tags: |
| - text-generation |
| - from-scratch |
| - grpo |
| - reinforcement-learning |
| - arithmetic-reasoning |
| --- |
| |
| # tinyzero-countdown-19m |
|
|
| A ~18.9M-parameter decoder-only transformer, pretrained from scratch and |
| post-trained with GRPO (Group Relative Policy Optimization) to solve |
| Countdown-style arithmetic puzzles: given a set of numbers and a target, |
| find an equation using each number exactly once that reaches the target. |
|
|
| ## Architecture |
| RoPE positional embeddings, RMSNorm, grouped-query attention (via |
| `F.scaled_dot_product_attention`), SwiGLU MLP, tied embeddings. Custom |
| 8192-token BPE vocabulary trained on the pretraining corpus (not GPT-2's |
| tokenizer -- see rationale below). |
|
|
| - Parameters: ~18.88M (verified exactly, not estimated) |
| - Context length: 256 |
| - Vocab size: 8192 (custom-trained BPE) |
| - d_model: 384, layers: 10, heads: 6 (2 KV heads, GQA) |
| |
| ## Training pipeline and what I learned building it |
| |
| **Pretraining**: ~380M tokens on FineWeb-Edu + synthetic arithmetic |
| text, at a ~20:1 token:parameter ratio (Chinchilla-optimal). An earlier |
| attempt at 116M params / 150M tokens (1.3:1 ratio) showed the failure mode |
| directly: healthy train/val loss gap but weak generalization. This version |
| also fixes a subtler issue -- at small model scale, a standard 50k-token |
| vocabulary's embedding table dominates the parameter budget (60-75% of |
| total params); training a small custom vocab instead keeps embedding |
| overhead to ~17%, leaving actual capacity for reasoning. |
| |
| **SFT**: an instruction-format fine-tune initially looked successful by |
| loss (train 1.02->0.34) but generation accuracy was 0% -- a real |
| loss/accuracy divergence caused by a response template that was mostly |
| easy-to-predict boilerplate, diluting the loss signal on the tokens that |
| actually mattered (the numbers/operators). Root-caused to a large, |
| un-bridged distribution shift between the pretraining corpus's format and |
| the instruction phrasing; fixed by skipping the instruction wrapper and |
| running GRPO directly on the pretrained checkpoint's native prompt format |
| instead. |
| |
| **GRPO**: trained directly on the pretrained checkpoint, using the |
| verifier (exact equation checker) as a binary+partial-credit reward, group- |
| relative advantage normalization, PPO-style clipping, and a KL penalty |
| against a frozen reference to prevent collapse. Result: 31.6% -> |
| 34.4% accuracy on a held-out 250-problem set (+2.8pp), with stable |
| KL throughout (no collapse). This is a modest, honestly-reported effect -- |
| run at only ~500 steps on an 18.9M model, not a large or highly significant |
| result, and reported with that caveat intentionally. |
| |
| ## Usage |
| |
| ```python |
| import torch |
| from tokenizers import ByteLevelBPETokenizer |
| # adapt these imports to wherever you place model.py / config.py from this repo |
| from model import TinyTransformer |
| from config import ModelConfig |
| |
| mcfg = ModelConfig() |
| model = TinyTransformer(mcfg) |
| ckpt = torch.load("pytorch_model.pt", map_location="cpu") |
| model.load_state_dict(ckpt["model_state_dict"]) |
| model.eval() |
| |
| tokenizer = ByteLevelBPETokenizer("vocab.json", "merges.txt") |
| prompt = "Numbers: [12, 45, 7, 3], Target: 88, Equation:" |
| ids = tokenizer.encode(prompt).ids |
| x = torch.tensor([ids]) |
| out = model.generate(x, max_new_tokens=40, temperature=1.0, top_k=1) |
| print(tokenizer.decode(out[0, len(ids):].tolist())) |
| ``` |
| |
| ## Limitations |
| - Small model (~19M params) -- general text fluency is weak; this is |
| specialized for the countdown arithmetic task, not general-purpose use. |
| - Only handles the raw prompt format shown above; natural-language |
| instruction phrasing was found to significantly degrade output quality |
| (see training notes above) and was not used for the released checkpoint. |
| - Evaluated on synthetically generated countdown problems only. |
| |
