Blame for README.md

---
language:
- en
library_name: transformers
pipeline_tag: text-generation
tags:
- math
- arithmetic
- reasoning
- grpo
- small-language-model
- from-scratch
datasets:
- arpjw/microreason-math
- HuggingFaceTB/cosmopedia
---

# MicroReason-43M

Small decoder-only language models (42.7 M parameters) trained **from random initialization** to
do exact arithmetic and linear algebra with step-by-step work. They come from a research project
that asked whether a staged pipeline can teach a tiny model to reason within a small compute
budget. The pipeline is base pretraining → math mid-training → SFT → RL. The total GPU spend was
about $23.

All three checkpoints share one architecture and tokenizer:

| Folder | Model | Held-out-shape test | In-distribution test | 4–5 digit extrapolation |
| --- | --- | ---: | ---: | ---: |
| [`sft-long/`](./sft-long) | SFT model (3.1 M supervised examples) | 81.4% | 87.1% | 10.2% |
| [`grpo/`](./grpo) | `sft-long` + GRPO (150 on-policy steps, outcome reward) | 82.0% | 87.8% | 10.1% |
| [`extra-sft/`](./extra-sft) | `sft-long` + continued SFT for the same GPU time as GRPO | 83.1% | 88.8% | 7.8% |

**How these numbers were measured:**
- **Decoding:** greedy, show-work format, exact final answer.
- **Suite sizes:** 6,000 / 7,000 / 7,000 problems.
- **Sealing:** the suites were frozen before training and opened once.

**Which to use:**
- **`extra-sft`** is the most accurate on in-range problems.
- **`grpo`** keeps its ability on numbers larger than any seen in training. Continued SFT erodes
  that ability, and the result replicated on a fresh test set.

## What the models do

**Expressions:** expressions with 1–5 operators over integers up to 999 (`+ − × ÷`, parentheses),
with exact rational results.

**Equations:** linear equations such as `3(x + 4) − 2 = 5x − 10`.

**Output format:** each model writes numbered reduction steps, then an answer:

```
[steps]
1. Subtract: 12 - 5 = 7
2. Multiply: 7 * 3 = 21
3. Divide: 4 / 2 = 2
4. Add: 21 + 2 = 23
[/steps]
[answer]23[/answer][eos]
```

**Accuracy falls with depth:** about 98% with 1 operator, 76% with 3 and 60% with 5. Linear
equations are about 99%.

## Usage

The architecture is custom: RMSNorm, RoPE, SwiGLU, and 12 layers × 512 dimensions. Load it with
`trust_remote_code=True`; the code is in `modeling_microreason.py` and needs only `torch` and
`transformers`.

```python
import torch
from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast
from huggingface_hub import snapshot_download

path = snapshot_download("arpjw/microreason-43m", allow_patterns=["grpo/*"]) + "/grpo"
tok = PreTrainedTokenizerFast(tokenizer_file=f"{path}/tokenizer.json")
model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True).eval()

prompt = "[bos][user]Compute: (12 - 5) * 3 + 4 / 2 Show your work.[eos][assistant]"
ids = torch.tensor([tok(prompt, add_special_tokens=False)["input_ids"]])
eos = tok.convert_tokens_to_ids("[eos]")
with torch.no_grad():
    for _ in range(320):
        nxt = model(ids).logits[0, -1].argmax().item()
        ids = torch.cat([ids, torch.tensor([[nxt]])], dim=1)
        if nxt == eos:
            break
print(tok.decode(ids[0]))
```

**Prompt formats** (the special markers are single tokens):
- **Show work:** `[bos][user]{problem} Show your work.[eos][assistant]`
- **Answer only:** `[bos][user]{problem} Give only the answer.[eos][assistant]` (much less
  accurate: about 35%)

**Problem phrasing** follows the training data, for example `Compute: …` or
`Solve for x: …`.

## Training

| Stage | Data | Notes |
| --- | --- | --- |
| Base pretraining | 1.0 B tokens: 70% general text (Cosmopedia), 30% synthetic math | from random init |
| Math mid-training | 200 M tokens, 85% math with an easy → hard curriculum | |
| SFT | 100 k, then 3 M more verified worked solutions (75% show-work, 25% answer-only) | `sft-long` |
| GRPO | 150 steps × 1,024 completions, group 16, binary exact-answer reward, pass-rate filter (2–14 of 16 correct), LR 3e-6, KL 1e-3, one update per batch | `grpo` |
| Continued SFT | about 2.2 M fresh worked solutions in the same 46 min of GPU time | `extra-sft` |

**Tokenizer:** an 8,192-token byte-level BPE that splits every number into single digits.
Splitting digits, and generating problems from random expression trees instead of fixed
templates, raised held-out accuracy from 0.2% to 70%.

**Synthetic math:** generated by a seeded expression-tree grammar. Every training problem is
checked against the frozen evaluation suites to exclude overlap. The evaluation suites and the
100k SFT set are published as
[arpjw/microreason-math](https://huggingface.co/datasets/arpjw/microreason-math).

## Findings

- **Math mid-training** improves accuracy (+5.3 points on test).
- **SFT** is what gives the answer format.
- **GRPO** gives a small but consistent gain over its SFT parent: +0.5 points, positive on all
  three seeds. Continued SFT on unlimited correct traces gains more at equal compute (+1.5).
- **Extrapolation:** GRPO preserves accuracy on 4–5 digit numbers, which continued SFT erodes.
  GRPO beats it there by +2.2 points, replicated on an independent frozen suite.

## Limitations

- Narrow domain: arithmetic and linear equations only, in the training phrasing. These are not
  general chat or general-purpose math models.
- Accuracy on numbers beyond 3 digits is low (about 10%).
- Small models trained on a narrow distribution; outputs should not be relied on without
  checking.