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
[link] [comments]
More from r/MachineLearning
-
For the people who got reviews back from neurips, cvpr, eccv, etc and also tested their paper through an agentic reviewer like the stanford one, how different were the reviews? [D]
Aug 14
-
Building text to ASCII diffusion model , need advice and guidance [P]
Aug 14
-
A collision-entropy floor for watermark/retrieval AI-text detection. Looking for a sanity check before I take this further [D]
Aug 14
-
Are supervised and unsupervised learning still relevant today? [D]
Aug 14
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.