--- license: apache-2.0 language: - en library_name: transformers pipeline_tag: text-generation base_model: nkthebass/tinybrainbot-320mV2-base tags: - tinybrainbot - small-language-model - from-scratch - gqa - instruct - chat --- # TinyBrainBot 320M V2 — Instruct A **~326M-parameter** decoder-only language model, trained from scratch on **~10B tokens** and then **supervised-fine-tuned** for chat / instruction following. This is the **instruct** model. - Base model: **`tinybrainbot-320mV2-base`** (full pretraining details there). - Successor to the 303M V2 instruct. **TL;DR:** A compact general-purpose assistant trained from scratch on ~10B tokens, with an added **conversational + in-context-recall** fine-tune (robust multi-turn chat, diverse instruction phrasings, remembers facts stated earlier in the conversation). It matches/beats **Pythia-410M** on general benchmarks (on far fewer training tokens) and sweeps **GPT-2-124M**, follows instructions, and has **2–3-digit arithmetic far stronger than its GSM8K reasoning score would suggest** (94–99% on 2–3-digit addition) — though it hits a clear length-generalization wall beyond 3 digits. It 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 (67% English / 20% code / 13% math), with reserved ``/`` special tokens | | Precision | trained in fp16 with an fp32 master copy (autocast) | ## Usage Prompt with the chat format: ``` <|user|> {user message} <|end|> <|assistant|> {assistant reply} <|end|> ``` Example (greedy): ``` <|user|> What is the capital of France? <|end|> <|assistant|> Paris. <|end|> ``` The model gives concise direct answers and shows worked steps for arithmetic. ### Conversational + recall update This release adds a **conversational + in-context-recall** fine-tuning pass on top of the base instruct SFT. It: - answers open-ended imperatives robustly — `list all the planets in the solar system` → *Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune.*; - handles **multi-turn chat and in-context recall** (state a fact early, ask about it later); - stays consistent across the F16 GGUF (LM Studio / Ollama / llama.cpp) and the fp16 `transformers` weights. *(An earlier build gave terse/empty answers to some imperative phrasings through the GGUF export. That was a real tokenization mismatch in the export — not a capacity limit. This release fixes it two ways: the conversational pass makes the model robust to it, **and** the F16 GGUF now sets `tokenizer.ggml.add_space_prefix=false` + ships a leading-space chat template so llama.cpp tokenizes the chat format token-for-token identically to the native tokenizer, per llama.cpp [#23840](https://github.com/ggml-org/llama.cpp/issues/23840).)* --- ## Training procedure ![training loss curve](loss-curve-instruct.png) *Pretraining (blue, ~10B tokens) → instruct SFT (orange, from step 51k) → conversational + recall SFT (green, steps 52k–55k, final loss ~0.5).* ### Pretraining (base) Pretrained from scratch on **10.03B tokens** (51,000 steps, WSD schedule, peak LR 6e-4, two-phase broad → quality-anneal, final loss **1.436**) across **2× Tesla V100-PCIE-16GB** with PyTorch DDP (gloo), fp16 autocast + fused AdamW. Full pretraining details and the 13-source data mix are documented on the **base** model card. ### Supervised fine-tuning (this model) | Hyperparameter | Value | |---|---| | Steps | ~1,433 (resumed from base @ 51,000) | | Tokens | ~94M | | LR | 1.5e-5, constant | | Warmup | 100 steps | | Global batch | 4 micro × 8 grad-accum × 2 GPUs × 1024 seq | | Loss masking | assistant-only | | Best checkpoint | step 52,000 (selected on eval, not last) | **SFT data mix** (sampling weights; sum = 9.5; math ≈ 15.8% of the mixture): | Source | Weight | |---|---:| | longdef-sft (~16k long / multi-step answers) | 2.0 | | smoltalk | 2.0 | | **math-v2** (~70k verified worked examples) | 1.5 | | qa-distill | 1.0 | | lamini-instructions | 1.0 | | soda-dialogues | 1.0 | | greetings | 0.5 | | reasoning-distill | 0.5 | *`math-v2` is ~70k programmatically-generated, tolerance-verified arithmetic worked examples (multiplication / division / decimals, with the final answer checked). It appears in both pretraining and SFT — see the note below.* > **On where the arithmetic comes from:** the base and instruct models score almost identically on the GPT-3 Arithmetic suite (**aggregate 31.4% vs 31.2%**). That base/instruct parity suggests **most arithmetic computation was acquired during pretraining**, while SFT mainly shaped instruction-following and response format (the one sub-task where SFT clearly helps is composite / order-of-operations, 7.0% vs 3.7%). Pretraining also delivered far more math *tokens* in absolute terms — on the order of ~215M vs ~15M — despite SFT's higher math *proportion*. ### Conversational + recall fine-tuning (final stage — this release) A further SFT pass continued from the instruct checkpoint (**step 52,000 → 55,000**) to make the model a robust multi-turn conversationalist, teach in-context recall, and fix brittle behavior on open-ended instruction phrasings. | Hyperparameter | Value | |---|---| | Steps | 3,000 (resumed from instruct @ 52,000) | | LR | 1.5e-5, constant (WSD, decay-fraction 0) | | Global batch | 2 micro × 16 grad-accum × 2 GPUs × 1024 seq | | Loss masking | assistant-only | | Final loss | ~0.5 | **Data mix** (sampling weights): | Source | Weight | |---|---:| | **convo-recall** (6k multi-turn in-context-recall dialogues) | 2.5 | | **convo-core** (4k instruction-phrasing + small-talk) | 2.5 | | smoltalk | 1.5 | | math-v3 | 0.75 | | longdef-sft | 0.75 | | greetings | 0.5 | | qa-distill | 0.5 | `convo-recall` and `convo-core` are programmatically generated: `convo-recall` teaches **in-context memory** (the user states facts early — name / city / pet / counts — then asks about them later), and `convo-core` covers **diverse instruction verbs** ("list all / name / give me / what are the …") over closed sets with correct complete answers, plus natural small-talk. Alongside the conversational gains, this pass makes the model **robust to the SPM normalization that the GGUF/HF exports drop** — so it now behaves consistently in llama.cpp / LM Studio / Ollama *and* in `transformers`, instead of degrading on some phrasings through the export. --- ## 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 instruct (full test sets) | Benchmark | n | **320M V2 instruct** | 303M instruct | |---|:--:|:--:|:--:| | HellaSwag | 10042 | **34.5** | 30.7 | | ARC-Easy (acc_norm) | 2376 | **49.3** | 47.6 | | ARC-Easy (raw acc) | 2376 | **57.0** | 51.0 | | ARC-Challenge | 1172 | 27.6 | 27.6 | | OpenBookQA | 500 | **31.8** | 29.0 | | WinoGrande | 1267 | **53.3** | 52.2 | | MMLU | 14042 | **28.0** | 27.1 | → **5 wins, 1 tie, 0 losses** over the previous generation. ### vs reference models (headline metric) | Benchmark | **320M V2 instruct** | GPT-2-124M | Pythia-410M | SmolLM-360M | |---|:--:|:--:|:--:|:--:| | HellaSwag | 34.5 | 31 | 34 | 54 | | ARC-Easy | 49.3 (57 raw) | 44 | 52 | 70 | | ARC-Challenge | 27.6 | 22 | 24 | 37 | | OpenBookQA | 31.8 | 29 | 30 | 42 | | WinoGrande | 53.3 | 52 | 53 | 57 | | MMLU | 28.0 | 26 | 25 | 34 | → Sweeps **GPT-2-124M**; **~5 wins + 1 draw vs Pythia-410M**. SmolLM-360M (trained on ~600B aggressively-filtered tokens) remains the 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 therefore best read as a **token-efficiency** result (curated data + quality anneal) rather than a scale win. ### Math — computation vs reasoning The model was trained on arithmetic **computation**, not word-problem **reasoning** — the two benchmarks below show that split clearly. **GSM8K** (grade-school word problems, full 1319-problem test, zero-shot chain-of-thought): | Model | GSM8K | |---|:--:| | GPT-2-124M | ~0% | | **320M V2 instruct** | **0.53%** | | Pythia-410M | ~1–2% | | SmolLM2-360M-Instruct | ~3–5% | → At the floor **for a general-purpose model of this size and training mix**; stronger sub-1B *math-specialized* models can score substantially higher. GSM8K rewards multi-step semantic reasoning, which this recipe did not target. **GPT-3 Arithmetic** (Brown et al. 2020 protocol, exact-match, n=300/sub-task): | Sub-task | Accuracy | |---|:--:| | 2-digit addition | 99.0% | | 3-digit addition | 94.0% | | 2-digit subtraction | 49.3% | | 3-digit subtraction | 42.7% | | 4-digit addition / subtraction | 0.3% / 0.3% | | 5-digit addition / subtraction | 0.0% / 0.0% | | 2-digit multiplication | 21.7% | | single-digit composite (order of ops) | 7.0% | | **Aggregate (all 10 sub-tasks)** | **31.4%** | → **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 and end with a stated answer). Subtraction sits ~42–50% — it handles `a−b` when `a>b` but drops the sign on **negative results**. 2-digit multiplication ~22%; single-digit composite (order of operations) ~7%. > **How to read this:** the GPT-3 Arithmetic suite mainly probes **exact symbolic computation** and short-range algorithmic generalization; it should **not** be interpreted as evidence of strong mathematical *reasoning* (see GSM8K above). The two results together are the point: strong at computing, weak at reasoning. --- ## Intended use & limitations **Intended use:** a capable general chat assistant at ~326M scale, on-device / low-resource deployment, research on small-model SFT, and arithmetic computation. **Limitations:** - **Math reasoning** (word problems, GSM8K/MATH) 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 RLHF/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.