Untied embedding parameter tax and depth scaling in BabyLM
Hi Yusuke,
In your 24.28M parameter budget, untied embeddings consume an oversized share of the model. With a 16,384 vocabulary and 256 hidden dimension, the separate input and output matrices take 16,384 * 256 * 2 = 8.39M parameters. That is 34.55% of your total budget spent on static token projections, leaving only 15.89M parameters across your 7 layers.
Your evaluation highlights a steep drop on Entity Tracking (18.56%). At 7 layers, shallow depth makes tracking state transformations difficult across a context window.
In an open architecture I work on called Maba (101M reference model: https://huggingface.co/AndrewThompson1233/maba-v1-architecture; weights: https://huggingface.co/AndrewThompson1233/maba-101m), we avoid this with two techniques:
Low-rank factorized embeddings: Projecting untied vocabulary through a rank-64 bottleneck (16,384 -> 64 -> 256) drops the combined input/output footprint to roughly 2.13M parameters. Reclaiming that 6.26M budget allows adding 2 to 3 full VISTA layers (expanding from 7 to 9-10 layers) within the exact same 24.3M ceiling.
2-pass physical block recycling: Passing representations through your 7 variance-informed blocks twice with layer conditioning gives 14 effective layers of non-linear transformations without increasing parameter count.
Under the 10M-word data constraint of BabyLM Strict-Small, did you test low-rank embedding bottlenecks to direct more parameters into the layer stack?
Best,
Andrew
Hi Andrew,
Thank you for the detailed suggestion. Your comment prompted me to run a direct series of follow-up experiments.
I kept the tokenizer, corpus order, seed, training schedule, and evaluator fixed, then compared three parameter-matched allocations:
| Model | Embedding design | Layers | Parameters |
|---|---|---|---|
| Original VISTA | Full-rank, untied | 7 | 24.281M |
| Factorized VISTA | Rank 64, untied | 9 | 24.248M |
| Factorized VISTA | Rank 128, untied | 8 | 24.228M |
At the first 10M-word exposure point, their local six-axis Full averages were:
| Model | Score |
|---|---|
| Original VISTA | 45.537 |
| Rank 64 / 9 layers | 45.302 |
| Rank 128 / 8 layers | 45.318 |
I also continued them on repeated exposure to the same 10M-word corpus as a research diagnostic. These later results are not official Strict-Small submission scores, since total training exposure reached 100M words.
Their best checkpoints were:
| Model | Best score | Best checkpoint |
|---|---|---|
| Original VISTA | 48.072 | 80M |
| Rank 128 / 8 layers | 47.704 | 90M |
| Rank 64 / 9 layers | 47.132 | 80M |
The rank-128 version produced an interesting tradeoff. Entity Tracking improved from 18.56 to 19.34, suggesting that the additional depth helped with state transformations. However, EWoK fell from 52.83 to 49.04, and the overall score remained 0.368 points below the original model. Rank 64 appeared too restrictive for the lexical interface under this training regime.
I also tested a learned block-recycling variant. It was not identical to your fixed two-pass, layer-conditioned design: a router decided whether to revisit shared blocks or exit. The model quickly learned to exit after the first pass, so the recycling path received almost no use later in training. Its peak score was 47.760. Entity Tracking rose to 19.28, but EWoK fell to 50.88. I therefore regard that experiment as evidence about the failure mode of adaptive routing, rather than evidence against deterministic two-pass recycling.
As a counter-test, I moved in the opposite direction and converted depth into width: four width-320 blocks with a 1,280-dimensional FFN, matched closely to the original model’s average compute. That model peaked at 46.994. The result strengthened the case that VISTA still needs sequential depth, even though the original embedding matrices consume a large part of its parameter budget.
So your diagnosis was valuable: the parameter allocation really does matter, and greater depth improved Entity Tracking. My results suggest that rank 64 compresses the lexical pathway too aggressively, while rank 128 recovers much of the performance but still loses important world-knowledge capability.
I would be very interested to know how Maba conditions its second pass. In particular, does it use a learned pass embedding, separate normalization parameters, or another signal that prevents the recycled blocks from simply repeating their first-pass computation?
Best,
Yusuke
Hi Yusuke,
Running that full sweep so quickly is incredible. Seeing Entity Tracking jump from 18.56 to 19.34 cleanly isolates the benefit of extra depth for state representation.
The drop on EWoK (52.83 -> 49.04) points straight to the bottleneck on the output head. In strict 10M-word data budgets, compressing the output projection matrix caps logit entropy and hurts factual recall. An asymmetric test (rank-128 input projection with a full-rank 256 output head) keeps the world-knowledge capacity intact while still freeing ~3.1M parameters for depth.
Your router finding also makes complete sense. Learned routing often degenerates into an early-exit shortcut because the optimizer minimizes early loss by skipping pass 2.
To prevent recycled blocks from repeating pass-1 transformations, Maba relies on fixed deterministic passes with two conditioning signals:
- Learned pass embedding: At the start of pass 2, add a learned vector to the residual stream (x = x + pass_embed[1]), identical to how segment embeddings signal sentence identity.
- Independent normalization: Keep the heavy attention and FFN projection matrices shared across both passes, but maintain separate RMSNorm weights for pass 1 and pass 2.
Activation variance shifts significantly after 7 layers of processing. Giving each pass its own scale vector costs negligible parameters (a few hundred floats) but allows the shared weights to operate in the correct statistical regime during pass 2.
Best,
Andrew