File size: 7,366 Bytes
3020d73 3e12c54 3020d73 3e12c54 3020d73 | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | ---
license: apache-2.0
language:
- en
library_name: transformers
pipeline_tag: text-generation
tags:
- tinybrainbot
- small-language-model
- from-scratch
- recurrent-depth
- looped-transformer
- arithmetic
- custom_code
- base
datasets:
- HuggingFaceFW/fineweb-edu
- HuggingFaceTB/cosmopedia
- roneneldan/TinyStories
- wikimedia/wikipedia
model-index:
- name: tinybrainbot-3m-loop
results:
- task:
type: text-generation
dataset:
type: allenai/ai2_arc
name: ARC-Easy
config: ARC-Easy
split: test
metrics:
- type: acc_norm
value: 30.81
name: acc_norm (0-shot)
- task:
type: text-generation
dataset:
type: allenai/ai2_arc
name: ARC-Challenge
config: ARC-Challenge
split: test
metrics:
- type: acc_norm
value: 22.95
name: acc_norm (0-shot)
- task:
type: text-generation
dataset:
type: Rowan/hellaswag
name: HellaSwag
split: validation
metrics:
- type: acc_norm
value: 27.36
name: acc_norm (0-shot)
- task:
type: text-generation
dataset:
type: ybisk/piqa
name: PIQA
split: validation
metrics:
- type: acc_norm
value: 56.04
name: acc_norm (0-shot)
- task:
type: text-generation
dataset:
type: AxiomicLabs/Arithmark-3.0
name: ArithMark-3.0
split: train
metrics:
- type: acc_norm
value: 89.30
name: acc_norm (0-shot)
---
# TinyBrainBot-3M-Loop
**A 2.96M-parameter 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 5.5B tokens.
| Benchmark (0-shot, full set) | acc_norm |
|---|--:|
| **Open SLM Leaderboard Intelligence Index** | **20.12** |
| *(raw acc)* | *20.58* |
| HellaSwag | 27.36 |
| *(raw acc)* | *26.94* |
| ARC-Easy | 30.81 |
| *(raw acc)* | *33.67* |
| ARC-Challenge | 22.95 |
| *(raw acc)* | *18.00* |
| PIQA | 56.04 |
| *(raw acc)* | *57.67* |
| ArithMark-3.0 | 89.3 |
| *(raw acc)* | *89.7* |
lm-eval 0.4.13 in float32 for HellaSwag, ARC and PIQA; the official `bencharithmark-3.py` (float32) for ArithMark-3.0.
## Arithmetic on wording it has never seen
ArithMark-3.0 is easy to overfit: a tiny model can learn one benchmark's sentence patterns instead of arithmetic. So we tested on held-out sets in the same continuation format (about 1,600 items over 29 grade-school skills). Their stories were written by models that wrote **none** of our training data (Gemma-3-27B and Llama-3.3-70B), Python set every number and answer, and a separate model re-solved every item. Scoring is the same as ArithMark-3.0 (4 endings, length-normalized).
The *traps* versions use the same stories, but the wrong endings are typical mistakes instead of nearby numbers: a number copied from the story, the result of the wrong operation, and the answer plus or minus 1 or 10. A model that echoes or pattern-matches numbers falls below chance (25%) there.
| Model | ArithMark-3.0 (our run) | Llama-written | Llama, traps | Gemma-written | Gemma, traps |
|---|--:|--:|--:|--:|--:|
| **TinyBrainBot-3M-Loop** (2.96M) | **89.3** | **92.8** | **90.9** | **69.3** | **58.1** |
| Qwen2.5-0.5B | 60.9 | 68.3 | 64.5 | 63.8 | 57.9 |
| SmolLM2-360M | 58.9 | 59.6 | 52.4 | 57.4 | 43.8 |
| Pascaline-3M | 69.2 | 47.1 | 22.6 | 39.0 | 11.2 |
| SmolLM2-135M | 38.8 | 33.2 | | 36.2 | |
| GPT-S3-3M | 38.9 | 33.0 | 9.1 | 36.2 | 9.9 |
*AM3 numbers in this table are all our own runs of the official `bencharithmark-3.py` (float32) on the current ArithMark-3.0 files, so every model is scored the same way. They can differ slightly from the values listed on the leaderboard (e.g. SmolLM2-360M: 58.9 here, 59.6 listed).*
It keeps its accuracy when the wording changes and when the wrong answers are plausible mistakes. Its weak spot is length: the Gemma-written stories are longer and chattier (about 200 characters against 120), and it handles plain, compact word problems much better than long ones with extra detail.
## 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
5.5B tokens, constant learning rate then a decay over the last 15%.
| 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 |
The decay used more arithmetic (45%) and WikiHow (20%). No benchmark training or test data was used, and none of the arithmetic generators copy ArithMark's sentences.
## Use
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "nkthebass/tinybrainbot-3m-loop"
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. 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,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 --dtype float32
```
## Limitations
- A 3M base model: it writes simple English and makes up facts. It is not a chat model.
- Arithmetic is strongest in short word problems; long, chatty setups lower its accuracy.
- 1,024-token context.
|