Link the MicroReason Math dataset

97cb057 verified 1 day ago

5.51 kB

---
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.

	---
	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.