Factorized embeddings and depth scaling on ~100M parameter budgets
Hi Pratham,
Really neat release - training a clean ~100M chat model completely from scratch is great work, and being upfront about the reasoning and context limitations is refreshing to see.
Looking at your architecture breakdown (12 layers, hidden 768, 32k vocab, 109.5M total), I noticed a bottleneck that almost every sub-150M Llama-style model runs into:
Vocab parameter tax (~23%): With a 32,768 vocabulary and 768 hidden dimension, the embedding table takes 32,768 * 768 ≈ 25.17M parameters. Nearly a quarter of your total 109.5M budget is spent on static token lookups rather than sequence modeling.
The depth bottleneck (12 layers): Because the embeddings lock down ~25M weights, the parameter budget forces you to stop at 12 layers, which directly caps multi-step reasoning capacity.
I have been working on this exact budget constraint in an open architecture called Maba (101M reference model):
https://huggingface.co/AndrewThompson1233/maba-v1-architecture
(Open weights and benchmarks: https://huggingface.co/AndrewThompson1233/maba-101m)
A couple of architectural observations that might be useful if you do an exploratory Particle 3.0 run:
Factorized embeddings: Projecting vocabulary through a low-rank bottleneck (e.g. 32k -> 128 -> 768) drops the vocab tax down to 4.3% (4.5M params). In our empirical runs on TinyStories, shrinking the embedding rank caused zero degradation in validation perplexity, but reclaimed ~20M parameters that can go straight into expanding depth from 12 to 18-20 layers.
2-pass physical block recycling: Passing tokens through physical blocks twice with layer conditioning effectively doubles your depth (e.g. 24-40 effective layers from a compact block footprint) without adding extra parameter weights.
Hybrid recurrence (75% GDN-2 / 25% GQA): Replaces quadratic attention in 3 out of 4 layers with linear state updates (O(1) memory), making KV-cache scaling past 2k tokens practically negligible on single-card inference.
Curious: during the Particle 1.0/2.0 architectural search, did you test embedding factorization or weight tying to free up budget for more layers?
Best,
Andrew
Thanks Andrew - appreciate the careful look at the 100M budget.
Particle 1.0/2.0 already uses tied input/output embeddings, so the 32k x 768 table is paid once (~24.6M, about 22% of 109.5M), not twice. We did not run embedding factorization or a depth-vs-width search; the 12 x 768 Llama-style block was an intentional vanilla baseline so the training stack (FineWeb-Edu, SFT mix, Spot) was the thing being tested.
Factorizing 32k -> 128 -> 768 would free ~20M params. At this width that is closer to +3 layers (12 -> 15) than 18-20, unless we also narrow hidden size or share blocks. TinyStories is a much easier signal than FineWeb-Edu, so I would want a matched FineWeb ablate before treating rank-128 as free.
2-pass recycling and GDN-2/GQA are interesting, but they change the compute graph and the "clean Llama from scratch" claim. If we do an exploratory 3.0 architecture run, factorization vs extra depth at fixed FLOPs is the first A/B I would actually run. KV cache past 2k is not the current bottleneck - capacity and data are.
Thanks for writing it up.
Hi Pratham,
Spot-on check on the layer count: in a Llama-style block with SwiGLU (3 FFN projections), ~20M parameters translates to roughly +3 full physical layers (12 -> 15). The 18-20 figure was factoring in 2-pass physical recycling.
Your point on FineWeb-Edu vs synthetic data is completely fair. Higher lexical diversity on web data puts more stress on token representation than TinyStories does. If rank-128 feels tight for technical/educational web text, testing rank-128 vs rank-192/256 is a quick way to find the inflection point where reconstruction error matches full-rank while still saving 16-18M parameters.
Fixed-FLOPs factorization vs depth is the cleanest A/B comparison to run for 3.0.
Best of luck with the training stack evals!
Best,
Andrew