coderian

/

StoryGPT

Like 1

Hi Coderian,

Pretraining a 29M causal transformer from scratch on TinyStories with a clean pre-LN implementation and standalone Hugging Face pipeline integration is a great educational project.

Looking at your exact architectural accounting (4 layers, 256 hidden size, 50,257 vocab) and the 128-token training block:

1. The vocabulary domain mismatch on TinyStories:
   The TinyStories dataset was deliberately synthesized using the basic vocabulary of 3 to 4-year-old children, typically spanning only 3,000 to 4,000 distinct words.
   Using the generic 50,257-token GPT-2 tokenizer means over 90% of the token rows in your embedding table are likely never updated across training.
   At hidden dimension 256, that static lookup table consumes 12,865,792 parameters. Out of your 29.1M total budget, that single table locks up 44.2% of the entire model (or ~88% if input and output matrices are untied).

2. The 4-layer depth bottleneck for plot coherence:
   At hidden dimension 256, a standard transformer block costs roughly 0.8M to 1.0M parameters.
   That static, mostly empty embedding dictionary literally costs the parameter budget of 12 to 14 full transformer layers.
   Because the computational backbone is left with only 4 shallow layers, the network struggles to track multi-character interactions, leading directly to the repetitions and logic lapses mentioned in your limitations.
   Pruning the tokenizer to 4,096 tokens or using two-stage factorized input projections (50,257 -> 64 -> 256 = ~3.2M params) frees up roughly 9.6M parameters. Reallocating those weights into depth would expand the model from 4 layers to 14 or 16 physical layers within the exact same 29M budget.

3. Context and state continuity:
   Stories naturally unfold across sequential narrative beats. While the model card lists 512 context, training on 128-token blocks truncates multi-paragraph story arcs early.
   In an open architecture project called Maba v2 (reference release: https://huggingface.co/AndrewThompson1233/maba-v2-architecture), we explore extreme parameter efficiency for compact models:
   We keep the vocabulary overhead under 5% through factorized projections, and handle long narrative dependencies using a hybrid of linear recurrence (DGDA) and specialized sparse attention (MABA-SA).
   This provides a persistent associative state across long sequences without quadratic cache bloat, keeping generation coherent across extended scenes.

If you are planning to iterate on StoryGPT, exploring the parameter allocation layout and factorized embedding design in the Maba v2 repo could offer practical ways to gain substantial depth on a 29M budget.

Did you keep the 50k GPT-2 tokenizer primarily for direct transformers library compatibility, or did you explore custom BPE vocab sizes for TinyStories?

Best,
Andrew