The 35% vocabulary tax and associative matrix states for sub-10M recurrent models
Hi Squeal Studio team,
Pretraining a custom 5.8M Russian causal language model from scratch on a Tesla T4 using a clean LSTM backbone, custom Unigram tokenizer, and rigorous deduplication pipelines (MinHash/LSH) is a really neat architectural experiment. Exploring recurrent inductive biases in the era of bloated dense transformers is genuinely refreshing.
Looking at your architectural accounting (2 LSTM layers, hidden dim 512, embed dim 256, 4k vocab) and generation trajectories:
The 35% untied vocabulary footprint:
Even with a compact 4,000-token vocabulary at embedding dimension 256, your input embedding table (1.024M params) and output lm_head (1.024M params) are untied.
Together, static token lookups consume 2.05M parameters, which accounts for exactly 35.0% of your entire 5.86M model budget.1.58M for layer 0 and ~2.10M for layer 1).
The two LSTM layers combined cost roughly 3.68M parameters (
Simply tying the input embeddings to the lm_head (setting self.lm_head.weight = self.embedding.weight) would immediately reclaim 1.02M parameters (17.5% of your total budget) at zero capacity loss. Those reclaimed weights could fund a 3rd full LSTM layer or allow widening the hidden dimension from 512 to 640 within the exact same 5.8M ceiling.
Vector hidden state bottlenecks versus associative matrix memory:
Standard LSTMs compress historical context into 1D hidden and cell vectors (h_t, c_t in R^512). Because scalar cell state updates lack key-value binding geometry, multi-step relational facts decay rapidly across steps, which explains the associative drift visible in the generation samples.
In modern linear recurrence architectures like DGDA (Decoupled Gated Delta Attention), the hidden state is upgraded from a 1D vector to a 2D associative matrix (S_t in R^{d_k x d_v}). By using data-dependent write and erase gates with delta-rule updates, the recurrent state acts as a dynamic fast-weight memory.
Training parallelism:
While standard PyTorch nn.LSTM is strictly sequential in time during training, modern linear recurrent formulations admit fused parallel associative scans.
In an open architecture project called Maba v2 (101M reference release: https://huggingface.co/AndrewThompson1233/maba-v2-architecture), we use DGDA linear recurrence across 75% of the network:
On a Tesla T4, fused Triton DGDA kernels achieve 264k to 375k tokens/second prefill throughput across 4,096 tokens, while maintaining the exact same constant O(1) decode memory and latency advantages of classic RNNs.
Did memory constraints or training step latency on the T4 lead to keeping the network at 2 layers, and have you considered testing tied embeddings for the next checkpoint?
Best,
Andrew