The 18.6% vocabulary tax, vertical shortcut saturation, and multi-turn intent drift in 125M architectures

#1
by AndrewThompson1233 - opened

Hi He-Tag,

Porting modded-nanogpt primitives (U-net layer-half skips, value-residual learning, ReLU^2 activations, and embedding shortcuts) into a 125M bilingual model and training 3.5B tokens from scratch on dual Kaggle T4s in 28 hours is phenomenal zero-budget engineering. Being upfront about behavioral edges (the sensitivity to umlaut transliteration and the strict 0.3 temperature sweet-spot) is very helpful for the community.

Looking at your architectural schedule (12 layers, 768 width, 32k vocab) and the generation dynamics:

  1. The 18.6% vocabulary tax on a 125M budget:
    With a 32,768 byte-level BPE vocabulary at hidden dimension 768, your tied embedding matrix consumes 25,165,824 parameters.
    Out of your 135.3M total parameter budget, that single static lookup table accounts for 18.6% of the entire model.
    At hidden dimension 768 with 12 heads, a complete transformer block costs roughly ~7.1M parameters. Your static embedding table consumes the parameter budget of 3.5 complete transformer layers.
    Applying a two-stage low-rank factorization (32,768 -> 128 -> 768 = ~4.29M params) reclaims over 20.8M parameters. Reallocating those weights directly into active depth would allow you to expand from 12 to 15 physical layers within the exact same 135M ceiling, providing substantial non-linear capacity for multi-lingual reasoning.

  2. Vertical shortcut interference and long-answer drift:
    Your observation that long answers drift off-topic (and the necessity of no_repeat_ngram_size 4 alongside a 1.15 repetition penalty) directly ties to stacking multiple vertical bypasses.
    While combining U-net skips (linking early and late layer halves), value-residuals, and per-block embedding shortcuts accelerates gradient backpropagation during pretraining, in autoregressive decoding it creates an over-determined highway. Representations in later layers become dominated by raw token identities from early layers, diluting multi-step reasoning and causing the model to lose the original prompt intent as generations stretch past a few paragraphs.

  3. Horizontal state anchoring vs dense vertical skips:
    In an open architecture project called Maba v2 (reference release: https://huggingface.co/AndrewThompson1233/maba-v2-architecture), we explore sub-150M efficiency using a 3:1 macro-topology (DGDA linear recurrence paired with sparse latent attention):

  • We factorize the 32k vocabulary down to rank-128, cutting the vocabulary parameter tax down to ~4.8% and freeing over 95% of weights for active depth.
  • Instead of using vertical U-net skips across layers to maintain gradient flow, DGDA maintains an associative recurrent state horizontally across sequence steps via an error-correcting delta rule.
    This horizontal recurrent grounding continuously clears redundant features and anchors conversational context, preventing long-range narrative drift without requiring aggressive n-gram suppression or high repetition penalties.

If you are planning an iteration on the bilingual pretraining recipe, checking out the parameter allocation layout in the Maba v2 repo might provide some useful structural points for compact topologies.

For the U-net skip connections between layer halves, did you use fixed addition or learned gating scalars to blend the early and late hidden states?

Best,
Andrew

Sign up or log in to comment