Lookup table allocation across value embeds and epistemic space
Hi team,
Interesting concept embedding token-level epistemic calibration directly into the architecture.
Looking at the tensor breakdown, lookup tables account for nearly two-thirds of the parameter budget:
One 32,768 x 1536 table takes 50.33M parameters.
Across wte (1), lm_head (1), value_embeds (12), and epistemic_table (12), the model holds 26 full-rank tables.
That totals 1.31B parameters in static lookups (65.8% of the model), leaving only 679.6M parameters for active attention and MLP blocks across the 24 layers.
At 1536 hidden width, holding 12 distinct full-rank value tables (604M) and 12 distinct epistemic tables (604M) introduces extreme parameter redundancy:
- Value embeddings: Tokens share underlying semantic subspaces. Routing through a shared low-rank base (Embedding(32768, 128) = 4.19M) with layer-specific up-projections (12 x Linear(128, 1536) = 2.36M) shrinks the 604M value footprint down to 6.55M parameters.
- Epistemic space: Token familiarity and uncertainty are largely driven by frequency, polysemy, and domain specificity. A 1536-dimensional manifold per token per layer is massively over-parameterized for uncertainty tracking. Projecting this into rank 16-64 would retain calibration signals while saving another 500M+ weights.
Factoring these auxiliary spaces frees over 1.1B parameters within the same 2B budget, which could triple the depth or width of the active sequence-modeling backbone.
Did you test low-rank factorized projections for the value and epistemic tables before training on the full 1536 rank?
Best,
Andrew
Hi Andrew — thanks for the careful breakdown; we agree that we could've made this model much more efficient. This was more a proof-of-concept release on the usefulness of the epistemic space and we have not taken the time to optimize it.
We used full-rank value and epistemic tables because they were the simplest baseline for validating the epistemic readout and we did not run a controlled low-rank factorization ablation for this checkpoint. We agree the epistemic table looks compressible, and shared low-rank value tables are a natural next experiment. We’ll add notes if/when we test those variants. We'd like to also try to scale up width and layer depth of the standard embedding space, utilizing the epistemic space as a signal for skip training. We've seen some pretty promising results on the potential for only training on novelty. Keeping everything 1:1 keeps the skip logic simpler so there is an advantage there. Thanks for pushing on this. We appreciate the feedback!
— QA team