File size: 3,747 Bytes
a080586 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | ---
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.
|