TinyBrainBot-61M-A10M-v4.1-MoE-Base

The first mixture-of-experts TinyBrainBot and the first model of the v4.1 series. Every layer holds 16 experts and each token is routed to 2 of them, so the model stores 60.6M parameters but runs only 10.0M per token (non-embedding; plus a 3.6M tied embedding).

It was trained on exactly the same 11B tokens, in the same order, with the same schedule and tokenizer as the v4 scaling study's dense models, so the comparison below isolates one change: dense feed-forward layers vs experts. At the compute of the dense 10M it scores higher on 10 of 13 benchmarks and averages 37.5 against the dense 10M's 36.5, and it beats every model here, including the dense 25M, on MMLU, PIQA, MathQA and ARC-Challenge.

EleutherAI lm-eval-harness v0.4.13, 0-shot, up to 2,000 examples per task. Each task's first row is acc_norm where the task reports it (acc otherwise); the "(raw acc)" row under it is acc. The average uses the first rows. Every model was run by us on the same harness and settings.

Quick start

It is a standard Mixtral model, so it needs no custom code:

from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "nkthebass/tinybrainbot-61m-a10m-v4.1-moe-base"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo)

ids = tok("Once upon a time, a little fox", return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=80)       # defaults from generation_config.json
print(tok.decode(out[0], skip_special_tokens=True))

GGUF (llama.cpp, LM Studio, Ollama): tinybrainbot-61m-a10m-v4.1-moe-base-F16.gguf (129 MB). It tokenizes and generates identically to the transformers version.

llama-completion -m tinybrainbot-61m-a10m-v4.1-moe-base-F16.gguf -no-cnv -p "Once upon a time, a little fox" -n 80 --temp 0.7 --repeat-penalty 1.15

Default sampling: temperature 0.7, top_p 0.95, top_k 40, repetition_penalty 1.15.

Samples

Unedited, at the default settings.

Once upon a time, a little fox named Sammy had an idea. He wanted to create an amazing treehouse made of pebbles and stones so that it could be easily accessed by everyone.

Photosynthesis is the process by which plants absorb light energy and convert it into chemical energy.

The best way to learn a new language is to try it out. To learn a new language, you need to know the words and phrases that make up the language.

Model

Attention is identical to the dense 10M; only the feed-forward layer changes. Its width (672) is split into experts of half that width, and each token uses two, so a token costs exactly what it costs in the dense model while the layer stores 8x the feed-forward parameters. Routing is Mixtral's (softmax over the experts, keep the top 2, renormalise their weights), trained with a 0.01 load-balancing loss, which keeps all 16 experts in use: in a typical training batch the least-used expert still gets about half of its fair share of tokens.

Training

11.0B tokens on 2x Tesla V100 16GB, about 25.5 hours, fp16, AdamW (0.9, 0.95), weight decay 0.1, peak LR 3e-3, 131k tokens per step, warmup-stable-decay - the v4 scaling study's exact data and schedule. The expert layers run on the sigma-MoE Triton kernel (Csordas et al.), which fuses token routing into the matrix multiply.

Then 500M tokens of the same mix plus 7.8% generated everyday-commonsense, procedural, worked arithmetic and short-fact writing (no benchmark data), and an 803M-token fact-dense anneal while the learning rate decays to zero.

Reproduce the benchmarks

lm_eval --model hf --model_args pretrained=nkthebass/tinybrainbot-61m-a10m-v4.1-moe-base,dtype=float32 \
  --tasks arc_easy,arc_challenge,sciq,openbookqa,mmlu,hellaswag,piqa,winogrande,social_iqa,commonsense_qa,lambada_openai,boolq,mathqa \
  --num_fewshot 0 --limit 2000 --batch_size 8 --seed 1234

TinyBrainBot v4.1 (mixture of experts)

• TinyBrainBot-61M-A10M-v4.1-MoE-Base (this model)

Dense comparison models: TinyBrainBot-25M-v4-Base and the v4 scaling study (500K to 25M).