README.md

---
license: apache-2.0
language:
  - en
library_name: transformers
pipeline_tag: text-generation
tags:
  - tinybrainbot
  - small-language-model
  - from-scratch
  - gqa
  - pretrained
---

# TinyBrainBot 320M V2 — Base

A **~326M-parameter** decoder-only language model **pretrained from scratch** on **~10B tokens** on 2× NVIDIA Tesla V100 (16GB). This is the **base** (pretrained + quality-annealed) model — a text-completion model with **no instruction tuning**.

- For the chat / instruction-following version, see **`tinybrainbot-320mV2-instruct`**.
- Successor to the 303M V2 base.

**TL;DR:** A compact general-purpose base LM trained from scratch on ~10B tokens. It matches/beats **Pythia-410M** on general benchmarks (on far fewer training tokens) and cleanly sweeps **GPT-2-124M**. Its **2–3-digit arithmetic is notably strong for the parameter count** (94–99% on 2–3-digit addition) — a product of curated numerical pretraining data — though it hits a clear length-generalization wall beyond 3 digits and is not math-specialized.

---

## Model details

| | |
|---|---|
| Parameters | **325,899,264** (~326M) |
| Architecture | Decoder-only transformer, pre-norm, RMSNorm, SwiGLU MLP, RoPE |
| Hidden size | 1024 |
| Layers | 26 |
| Attention heads | 16 (query) / **4 KV heads** (grouped-query attention) |
| FFN size | 2816 |
| Context length | 1024 |
| Vocabulary | 32,000 |
| Tokenizer | `tbb-32k-v2` — 32k BPE trained on the training mix (67% English / 20% code / 13% math), with reserved `<think>`/`</think>` special tokens |
| Precision | trained in fp16 with an fp32 master copy (autocast) |

## Usage

This is a **base completion model** — it continues text and does **not** follow a chat template. Prompt it with text to continue:

```python
prompt = "The capital of France is"
# -> the model continues the text
```

For Q&A / chat behavior, use the **instruct** variant.

> **Format:** provided as **fp16 safetensors** (`AutoModelForCausalLM`). No GGUF is shipped for the base model — it's a raw completion model, and GGUF chat apps expect an instruction-tuned model. For chat / GGUF, use the **instruct** variant.

---

## Training procedure

![training loss curve](loss-curve-base.png)

*Pretraining loss over 51,000 steps / ~10B tokens — the dip near step 40k is the quality-anneal phase.*

Trained with the custom **TinyBrainBot** trainer using PyTorch DDP (gloo backend) across **2× Tesla V100-PCIE-16GB** on Windows, fp16 autocast + `GradScaler`, fused AdamW. One continuous **Warmup–Stable–Decay (WSD)** run.

| Hyperparameter | Value |
|---|---|
| Total steps | 51,000 |
| Tokens | 10.03B |
| Global batch | 4 micro × 24 grad-accum × 2 GPUs × 1024 seq ≈ 196.6k tokens/step |
| Peak LR | 6e-4 |
| Warmup | 1,000 steps |
| Schedule | WSD — stable LR, then cosine decay over the final ~20% (from step 40,800) |
| Optimizer | fused AdamW |
| Final loss | **1.436** |
| Throughput | ~22k tok/s (2 GPU), ~4–5 day wall-clock |

**Two-phase data:** a broad base phase (0 → ~40k steps) followed by a **quality-anneal** phase (~40k → 51k) in which the data mix was swapped to a knowledge-dense blend (Wikipedia leads up, web data halved, math tripled) as the LR decayed. The anneal pulled loss from ~1.62 to 1.436.

**Pretraining data mix** (sampling weights; design principle: *real > synthetic, synthetic ≤ ~35%*):

| Source | Weight | Notes |
|---|---:|---|
| dclm-clean | 3.0 | DCLM web slice, light-cleaned (~3.15B tok) |
| wiki-dense | 2.0 | 5.2M Wikipedia article leads (knowledge-dense, ~1.66B) |
| fineweb-edu-10bt | 1.3 | Educational web (~1B weighted) |
| qa-distill | 0.6 | Synthetic Q&A |
| bookcorpus-clean | 0.5 | Modern novels, boilerplate stripped (~1.68B) |
| code-pyjs | 0.5 | the-stack filtered to Python/JS/TS (~0.48B) |
| facts-distill | 0.4 | Synthetic scoped facts |
| reasoning-distill | 0.3 | Synthetic reasoning |
| tiny-stories-v2 | 0.25 | Simple narrative |
| capability-distill | 0.2 | Synthetic capability data |
| **math-v2** | 0.15 | **~70k** verified mul/div/decimal worked examples (tolerance-checked) |
| writing-distill | 0.1 | Synthetic writing |
| math-scratchpad | 0.05 | Step-by-step arithmetic |

*(weights are sampling proportions; sum = 9.35. Explicit math ≈ 2.1% of the mixture — but ~215M tokens in absolute terms, the bulk of the model's math exposure.)*

---

## Evaluation

Measured on our own log-likelihood MC harness (lm-eval style, fixed seed). Headline metric = `acc_norm` for HellaSwag/ARC/OpenBookQA, `acc` for WinoGrande/MMLU. Reference values are published lm-eval approximations — treat gaps under ~±2 points as ties.

### vs the previous 303M base (n=2000)

| Benchmark | **320M V2 base** | 303M base |
|---|:--:|:--:|
| HellaSwag | **34.0** | 30.1 |
| ARC-Easy (acc_norm) | **50.0** | 48.5 |
| ARC-Easy (raw acc) | **57.0** | 51.8 |
| ARC-Challenge | **27.4** | 26.4 |
| OpenBookQA | **31.6** | 30.8 |
| WinoGrande | **54.3** | 51.6 |
| MMLU | **27.4** | 26.6 |

→ **Clean 6/6 sweep** over the previous generation.

### vs reference models (headline metric)

| Benchmark | **320M V2 base** | GPT-2-124M | Pythia-410M | SmolLM-360M |
|---|:--:|:--:|:--:|:--:|
| HellaSwag | 34.0 | 31 | 34 | 54 |
| ARC-Easy | 50.0 (57 raw) | 44 | 52 | 70 |
| ARC-Challenge | 27.4 | 22 | 24 | 37 |
| OpenBookQA | 31.6 | 29 | 30 | 42 |
| WinoGrande | 54.3 | 52 | 53 | 57 |
| MMLU | 27.4 | 26 | 25 | 34 |

→ Sweeps **GPT-2-124M** 6/6; **~5 wins + 1 draw vs Pythia-410M**. SmolLM-360M (trained on ~600B aggressively-filtered tokens) remains the out-of-reach frontier for this size.

**Training efficiency.** These results come from **~10B pretraining tokens** — roughly an order of magnitude fewer than the Pythia suite's ~300B. The Pythia-410M parity is best read as a **token-efficiency** result (curated data + quality anneal) rather than a scale win.

### Arithmetic (GPT-3 Arithmetic protocol, exact-match, n=300/sub-task)

| Sub-task | Accuracy |
|---|:--:|
| 2-digit addition | 99.3% |
| 3-digit addition | 94.7% |
| 2-digit subtraction | 50.0% |
| 3-digit subtraction | 42.3% |
| 4-digit addition / subtraction | 0.3% / 0.3% |
| 5-digit addition / subtraction | 0.0% / 0.0% |
| 2-digit multiplication | 21.0% |
| single-digit composite (order of ops) | 3.7% |
| **Aggregate (all 10 sub-tasks)** | **31.2%** |

→ **Strong through 3 digits, then a hard wall.** 2–3-digit addition is near-solved (94–99%, with correct carrying), but 4+-digit accuracy collapses to ~0%: the model executes a fixed **~3-column** addition routine and silently drops the higher place values — a **length-generalization limit** tied to the training distribution (`math-v2` operands are ≤3 digits), *not* truncation (generations complete normally). Subtraction sits ~42–50% — it handles `a−b` when `a>b` but drops the sign on **negative results**. 2-digit multiplication ~21%; single-digit composite ~4%.

> **How to read this:** the GPT-3 Arithmetic suite mainly probes **exact symbolic computation** and short-range algorithmic generalization; it should **not** be read as evidence of mathematical *reasoning* (word problems), which is a separate, much harder skill at this scale.

---

## Intended use & limitations

**Intended use:** research on small-model training, base for further fine-tuning, text completion, and arithmetic computation.

**Limitations:**
- Base completion model — **does not follow instructions or a chat format** (use the instruct variant for that).
- **Math reasoning** (word problems) is at the floor — the model computes but does not reason through multi-step problems.
- **Negative-result subtraction** is unreliable (drops the sign).
- **WinoGrande and MMLU** sit near the random floor — consistent with the capacity and data limits of a ~326M model under this training recipe.
- Trained predominantly on English; 1024-token context; no safety tuning — outputs may be incorrect or inappropriate and should not be relied upon unchecked.

## Hardware & framework

2× NVIDIA Tesla V100-PCIE-16GB · Windows · PyTorch DDP (gloo) · fp16 autocast (fp32 master) · fused AdamW · custom TinyBrainBot trainer.