Addressing the 66M embedding tax and RAG copy precision via block recycling
Hi Jaehyeon,
Scratch-pretraining a dedicated Korean SLM specifically calibrated for RAG document grounding is a very clean initiative, and explicitly separating the 107M backbone from the 66M embedding tables in your model card is great transparency.
Looking at your architecture trade-offs and the limitations on exact entity/digit copying:
The 66M lookup parameter tax:
With a 32,001 vocabulary at 1,024 hidden width and untied weights, your token embeddings and lm_head consume 65.5M parameters (38.1% of your entire 172M model).
Given that one transformer block (1,024 hidden, 2,048 SwiGLU, 4 KV heads) costs roughly 8.9M parameters, static lookup tables cost more than 7 full layers of compute.
Decoupling the input embedding via low-rank projection (32,001 -> 128 -> 1,024 = 4.23M params) while keeping the output head full-rank frees over 28.5M parameters. That reclaimed budget is enough to expand your physical backbone from 12 to 15 layers without increasing total model size.Entity extraction drift across 12 layers:
Your observation regarding swapped syllables in company names and altered digits during RAG extraction highlights an induction depth constraint. Verbatim copying from a [자료] prompt block requires multi-step attention routing: attending to source context, binding adjacent tokens, and copying exact spans to logits. A 12-layer depth struggles with this resolution, leading to approximate rather than exact extraction.
In an open architecture project called Maba (101M reference model: https://huggingface.co/AndrewThompson1233/maba-v1-architecture), we handle this using deterministic 2-pass block recycling:
Passing hidden states through your 12 physical blocks twice with Split RMSNorm (distinct norm vectors for pass 0 and pass 1) expands depth to 24 effective layers at zero additional parameter cost.
Pass 0 aligns the query against retrieved evidence, while pass 1 executes faithful verbatim token extraction. This extra compositional depth significantly reduces digit and proper noun hallucinations in constrained RAG workflows.
Did the memory footprint of the 66M embedding tables constrain you to 12 layers during the initial scratch pretraining run?
Best,
Andrew
Hi Andrew,
First of all, thank you for taking the time to discuss the architecture with me and for sharing such detailed feedback. I really learned a lot from your analysis and suggestions.
Regarding your question, yes, GPU memory was one of the constraints. I was training the model on a GPU with limited VRAM, so I had to be fairly conservative with the model size, batch size, and overall training configuration.
The amount of pretraining data was also relatively limited, so I was trying to balance model capacity, available data, training compute, and training time. The 12-layer configuration was therefore more of an overall engineering trade-off rather than being specifically constrained by the 66M embedding parameters.
I hadn't previously considered reallocating some of the embedding parameters toward increasing Transformer depth in this way. Your explanation of the low-rank embedding and especially the 2-pass block recycling approach gave me some very interesting ideas for future experiments.
In particular, I'd like to test whether increasing the effective depth, either through additional physical layers or block recycling, can improve exact entity and digit copying in my RAG setting.
Thanks again for the thoughtful discussion and feedback. It was a very valuable learning experience for me.
Best,
Jaehyeon
Hi Jaehyeon,
Glad the parameter reallocation breakdown and recycling concepts were helpful!
Balancing tight VRAM margins and limited training tokens is always tough, and keeping the initial baseline conservative makes complete sense from an engineering perspective.
If you decide to experiment with 2-pass block recycling on your next run, one practical tip: keep pass 0 and pass 1 normalization weights decoupled (distinct RMSNorm scale vectors per pass). It costs almost zero extra parameters, but provides the numerical headroom needed to prevent activations from drifting on the second pass.
Really looking forward to seeing how your future RAG iterations turn out. Best of luck with the experiments!
Best,
Andrew