r/MachineLearning · · 3 min read

Looking for feedback on a small test SLM I built completely from scratch [P]

Mirrored from r/MachineLearning for archival readability. Support the source by reading on the original site.

Architecture:

- Parameter count: 216.5M

- Layers: 10

- Attention / no attention:** Attention — 12-head multi-head self-attention, RoPE positional

encoding, SDPA. Decoder-only, pre-norm, RMSNorm + SwiGLU, tied input/output embeddings.

(hidden 1032, head_dim 86, FFN 4416)

- Tokenizer:** Custom 36k SentencePiece unigram, case-preserving, byte-fallback, with atomic

chat/role + memory special tokens (`<|user|>`, `<|assistant|>`, `<|end|>`, …)

- Context length: 768

- Dataset size/source: ~551M tokens. Pretraining on public English text (Wikipedia,

TinyStories, OpenWebText2); SFT on public instruction/chat sets (SmolTalk, UltraChat, Magpie,

AM-DeepSeek-R1, Orca-Math) + a small custom greetings/assistant set.

- Training time: ~15 hours

- Hardware: single NVIDIA RTX 3080 (10GB), bf16

- Training setup: AdamW (β 0.9/0.95), LR 3e-4, 1k warmup, weight decay 0.1, grad-clip 1.0,

effective batch 16,384 tokens/step (micro 4 × grad-accum 8 × seq 512), 33,650 steps.

Loss curve

Random init → 1.77 smoothed (1.27 last step). Started at exactly ln(36000)=10.49 (true from-scratch).

You can see SFT kick in — the sharp step-down at ~410M tokens.

(attach tinybrainbot_loss_curve.png)

Sample outputs

(F16 GGUF, temp 0.7 / top-p 0.9 / repeat-penalty 1.15 / freq 0.5 / presence 0.3)

> **User:** Hello

> **Bot:** Hello! How can I help you today?

> **User:** Tell me a story about a robot who learns to paint.

> **Bot:** In the heart of the city, there lived an assistant named Max. He was a kind and gentle

> soul with eyes that shone like diamonds, and he had always known that art would bring people

> together...

> **User:** Write a short poem about the rain.

> **Bot:** The rain was crisp and warm, casting dappled shadows on the walls of your town square...

> Your voice was so cool and gentle, so bright and calming all that you could feel at home.

> **User:** Explain what a computer is.

> **Bot:** ...the development of a computer can be categorized as either an operating system or a

> popular language, such as JavaScript or C++...

(^ that last one is the honest failure mode — confidently wrong on facts. It nails conversational form far better than it knows anything.)

What shocked me

- The tokenizer was way more important that I thought. I lost quite a few full runs on a broken char-level (386 vocab) fallback before building a real 36k subword tokenizer — and that fixed it and was the biggest quality jump than any architecture or data tweak.

- A GGUF export issue: llama.cpp has no unigram-tokenizer runtime on the default path, so a SentencePiece unigram*model gets silently re-segmented (Hello → Hell+o) and produces garbage — the weights were fine. Fix was exporting the GGUF vocab as UGM (t5 type) + the SP charsmap so it runs Viterbi. Cost me a day of pain.

What I want feedback on

- Token budget: 551M tokens for 216M params is way under Chinchilla (~4.3B). Would you push pretraining much longer before SFT, or is more/cleaner data the better spend at this scale for a quick demo?

- Data mix / timing: the SFT step-down in the loss is dramatic — is it real learning, or mostly the model latching onto chat format? How would you tell them apart?

- Cutting factual drift at ~200M params: more pretraining, longer context, or distillation froma bigger teacher — where's the best ratio for noticeably better factual responses?

- Arch nits: head_dim 86 is ugly (1032/12). Would 1024 hidden / 16 heads / 64 head_dim be ameaningfully better shape at this size (also allows quantization)?

Happy to answer anything about the pipeline (tokenizer training, the from-scratch loop, GGUF export).

Here's the model link as well https://huggingface.co/nkthebass/TinyBrainBot-demo-216.5m

submitted by /u/nkthebass
[link] [comments]

Discussion (0)

Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.

Sign in →

No comments yet. Sign in and be the first to say something.

More from r/MachineLearning