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