TinyBrainBot-3M-Loop-v2

A 2.96M-parameter math-first language model that reuses its layers. Four transformer blocks are run three times in a loop, so every token passes through 12 layers of computation while only 4 layers of weights are stored. Trained from scratch on 6.0B tokens. Compared with TinyBrainBot-3M-Loop, its arithmetic carries over to bare expressions, order of operations and numbers written in words, not only word-problem stories.

Benchmark (0-shot, full set) acc_norm
Open SLM Leaderboard Intelligence Index 20.15
  (raw acc) 19.91
HellaSwag 27.17
  (raw acc) 27.01
ARC-Easy 33.46
  (raw acc) 33.59
ARC-Challenge 22.78
  (raw acc) 18.17
PIQA 54.79
  (raw acc) 55.98
ArithMark-3.0 90.7
  (raw acc) 90.6

lm-eval 0.4.13 in float32 for HellaSwag, ARC and PIQA; the official bencharithmark-3.py (float32) for ArithMark-3.0.

Arithmetic beyond one benchmark

Every column below is a zero-shot, log-likelihood multiple-choice math test, all scored by us with the official bencharithmark-3.py scoring in float32, so every model is measured the same way (numbers can differ slightly from leaderboard listings).

Model Params ArithMark-3 AM3, numbers in words AM3 as bare equations ArithMark-2 ArithMark-1 Held-out stories
TinyBrainBot-3M-Loop-v2 2.96M 90.7 55.0 96.8 83.8 36.1 80.1
TinyBrainBot-3M-Loop 2.96M 89.3 27.3 88.2 58.1 29.4 81.0
Inferno-Exp-1.7 394M 85.9 59.1 86.9 59.9 47.9 74.0
Pascaline-3M 3M 69.2 27.6 99.7 66.5 36.2 43.0
SmolLM2-135M 135M 38.8 34.2 38.6 33.2 33.1 34.7
GPT-S3-3M 2.97M 38.9 25.8 36.2 31.0 31.1 34.6
  • Numbers in words and bare equations are ArithMark-3's own problems with only the surface changed.
  • ArithMark-2 uses its official metric (raw log-likelihood); its official benchmark_arithmark-2.0.py gives the same results.
  • Held-out stories: about 1,600 ArithMark-style problems written by Gemma-3-27B and Llama-3.3-70B, models that wrote none of our training data, with every number and answer set by Python.
  • ArithMark-1: its answer key mostly evaluates expressions strictly left to right (e.g. 14 * -17 + -1 * -18 = 4302), while ArithMark-2 and ordinary arithmetic use operator precedence (that expression is -220). This model follows precedence, so it scores lower on ArithMark-1's longer expressions (62% on its two-operand stage, falling with each extra operator).

How it works

  • Recurrent depth. 4 shared blocks (width 256, 4 attention heads of 64 with 2 KV heads, SwiGLU 560, RoPE, RMSNorm) are applied 3 times. Each pass adds a learned loop embedding and a gated copy of the input embedding. Training sampled 2 to 4 loops per step, and you can change the count at inference (model.model.eval_loops = 4); 3 is best.
  • Factorized embedding. The 4,104-token vocabulary uses 96-dimensional embeddings projected up to 256 (and down again before the tied output layer), which frees about 630k parameters for the 4th block.
  • Numbers. A units-first digit tokenizer writes every number as a length token plus its digits, smallest place first, and a digit-place embedding tells the model which place each digit is in.

Training

Stage 1 (4.46B tokens): the TinyBrainBot-3M-Loop run up to the start of its learning-rate decay.

Part Share What
Q&A and facts 25% our distilled Q&A, short facts, explanations
Arithmetic 20% word problems: about 150k different story wordings written by ling-3.0-flash and Mistral-Small-3.2, each checked by solving it, with every number and answer filled in by Python; plus our own relation (more/fewer than) and story generators
Synthetic textbooks 15% Cosmopedia (Stanford, OpenStax)
Stories 15% TinyStories
Commonsense 10% Cosmopedia WikiHow, our everyday-commonsense prose
Web 10% FineWeb-Edu
Reference 5% Wikipedia

Stage 2 (1.5B tokens, math-first): constant learning rate for 1.0B tokens, then a decay to zero over the last 0.5B.

Part Share What
Story arithmetic 25% the same verified wordings, with a quarter of the problems' numbers written in words
Bare arithmetic 25% our generator: expressions with order of operations, parentheses, negative numbers and 2-5 operands, in about 19 forms ("43 + 19 = 62", "What is ...?", "Compute ...:")
Number words 8% words to digits and back, arithmetic in words, counting, comparisons
Relations and stories 12% our relation and story generators
Q&A and facts 20% as in stage 1
Cosmopedia, FineWeb-Edu 10% as in stage 1

No benchmark training or test data was used, and none of our generators copy any benchmark's problems, code or sentences.

Use

from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "nkthebass/tinybrainbot-3m-loop-v2"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
ids = tok("Mia had 46 stickers. She gave 18 to her brother, so now Mia has", return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=8, do_sample=False)
print(tok.decode(out[0][ids.shape[1]:]))

trust_remote_code=True is required (custom architecture and tokenizer). There is no KV cache: generation recomputes the prefix, which is fast at this size. Padded batches work on either side (the model honors attention_mask). There is no GGUF, because llama.cpp does not support the looped architecture.

To reproduce the scores:

lm_eval --model hf --model_args pretrained=nkthebass/tinybrainbot-3m-loop-v2,dtype=float32,trust_remote_code=True --tasks arc_easy,arc_challenge,hellaswag,piqa --num_fewshot 0
python bencharithmark-3.py --model nkthebass/tinybrainbot-3m-loop-v2 --dtype float32

Limitations

  • A 3M base model: it writes simple English and makes up facts. It is not a chat model.
  • Arithmetic is strongest on short problems; long, chatty setups and long left-to-right chains lower its accuracy.
  • 1,024-token context.
Downloads last month
-
Safetensors
Model size
2.96M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Datasets used to train nkthebass/tinybrainbot-3m-loop-v2

Evaluation results