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.

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

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

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

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.