TIDE: Every Layer Knows the Token Beneath the Context

Abstract: We revisit a universally accepted but under-examined design choice in every modern LLM: a token index is looked up once at the input embedding layer and then permanently discarded. This single-injection assumption induces two structural failures: (i) the Rare Token Problem, where a Zipf-type distribution of vocabulary causes rare-token embeddings are chronically under-trained due to receiving a fraction of the cumulative gradient signal compared to common tokens; and (ii) the Contextual Collapse Problem, where limited parameters models map distributionally similar tokens to indistinguishable hidden states. As an attempt to address both, we propose TIDE, which augments the standard transformer with EmbeddingMemory: an ensemble of K independent MemoryBlocks that map token indices to context-free semantic vectors, computed once and injected into every layer through a depth-conditioned softmax router with a learnable null bank. We theoretically and empirically establish the benefits of TIDE in addressing the issues associated with single-token identity injection as well as improve performance across multiple language modeling and downstream tasks.

• The paper introduces TIDE, which supplements transformers with a persistent token identity pathway to address gradient starvation and contextual collapse.
• It employs multiple independent MemoryBlocks that amplify gradients for rare tokens and maintain semantic distinctions across layers.
• Empirical results show significant reductions in cross-entropy loss for rare tokens and improved performance with negligible VRAM and computational cost.

TIDE: Persistent Token Identity Across All Transformer Layers

Motivation and Problem Analysis

The paper rigorously interrogates the standard single-injection assumption in transformer-based LLMs, namely that token identity (via input embeddings) is consulted once at the input and never revisited throughout the network's depth. Two key structural deficiencies are formalized: (1) the rare token problem, where Zipf-distributed vocabularies result in chronic gradient starvation for rare token embeddings, and (2) contextual collapse, where semantically distinct tokens in similar contexts yield indistinguishable hidden states across the residual stream.

The rare token issue is empirically demonstrated by tracking the ℓ₂-norms of token embeddings throughout training. Mean norms for rare tokens display stagnation or even decline, while common token norms exhibit sustained growth, directly illustrating under-training and increased noise for rare tokens.

Figure 1: Mean embedding $l_2$-norms of a pretrained LLaMa-Base-1B checkpoint highlighting persistent under-trained rare token embeddings.

Contextual collapse is independently validated using sentence templates with distinct token insertions. Measurements of mean ℓ₂-distances between token-pair hidden states across layers reveal that, except for the final few layers, representations of distinct tokens rapidly converge—a manifestation of the FFN’s inability to differentiate in the absence of explicit token identity conditioning.

Figure 2: Heatmap showing the mean $\ell_2$-distance across layers between hidden states for distinct token pairs in identical contexts, confirming the prevalence of contextual collapse.

The TIDE Architecture

To resolve these pathologies, TIDE (Token Identity Delivered Everywhere) augments the transformer forward pass with a parallel, globally shared EmbeddingMemory: an ensemble of $K$ independent MemoryBlocks, each being its own vocabulary-wide embedding table. At every layer, a depth-conditioned softmax router linearly blends the $K$ memory vectors for each token and injects this context-free, token-identity signal into the residual stream independently of the FFN or attention pathways. A dedicated NULL bank lets the network learn, per layer and per token, how much to rely on this additive identity.

Figure 3: TIDE augments the transformer with a parallel EmbeddingMemory module composed of $K$ independent MemoryBlocks, with per-layer routed fusion into the residual stream.

The memory pathway is computationally efficient, imposing negligible VRAM or compute overhead, as the memory block lookups are cheap and the tables can be aggressively quantized or offloaded to SSD.

Figure 4: VRAM and SSD storage breakdowns for TIDE-1B versus LLaMA-Base-1B across varying $K$; VRAM remains nearly constant while SSD footprint scales linearly with $K$.

Theoretical Insights and Empirical Results

Rare Token Gradient Amplification

TIDE provides a $K$-fold amplification in the cumulative gradient signal available to each token, as MemoryBlocks are independent and each receives gradients whenever a token is present. This fundamentally alters the statistical bottleneck for rare token learning: gradient starvation is no longer dominated by base embedding update frequency but now receives multiplicative signal through multiple pathways.

Empirically, TIDE delivers strict and monotonic reductions in held-out cross-entropy loss for every token frequency decile compared to LLaMA-Base-1B, with especially pronounced improvements for rare tokens (up to a 4.8× gain differential between rarest and most frequent deciles).

Figure 5: Mean validation cross-entropy loss per frequency decile; TIDE-8E-1B outperforms baseline at all deciles, especially for rare tokens.

Resolving Contextual Collapse

The ability of TIDE to inject a token-discriminative signal at each layer enables persistent separation for semantically distinct tokens, independent of their contextual indistinguishability in the main residual path. Theoretical results confirm that token pairs collapsed under the FFN-Lipschitz bottleneck can be separated arbitrarily within EmbeddingMemory, as block embeddings are indexed strictly by discrete token identity.

Empirically, measuring ℓ₂-separation across layers for canonical collapse case studies (homophones, numerics, rare synonyms) shows that TIDE restores substantial token-wise separation, particularly in deeper layers where contextual collapse is most severe for the base model.

Figure 6: Layer-wise $\ell_2$ separation between token-pair hidden states; TIDE maintains greater separation, mitigating collapse across layers.

Scaling, Ablations, and Analysis

Language modeling and downstream evaluations across model scales (750M–3B parameters) and $K$ values (number of MemoryBlocks) demonstrate that TIDE confers consistent performance gains. Increased $\ell_2$0 monotonically improves validation perplexity and rare-token loss without evident saturation, and even modest $\ell_2$1 values (2–4) recover a majority of the rare-token gain.

Figure 7: Cross-entropy loss for rare, mid, and common tokens decreases with larger $\ell_2$2, highlighting greater benefit for rare tokens.

Figure 8: Validation PPL on Wikitext-2, PubMed, DCLM as a function of training tokens; TIDE's performance improves monotonically with $\ell_2$3.

Router and block specialization analyses indicate MemoryBlocks do not collapse to the baseline embedding subspace. Cosine distance statistics show that MemoryBlocks encode complementary, non-redundant information, and routing weights skew toward activating the memory pathway far more for rare tokens.

Figure 9: Mean cosine distance between baseline embedding and each of 8 MemoryBlocks under TIDE-8E-1B; high distance confirms non-collapsing, diverse signal.

Figure 10: Per-bin routing weights for MemoryBlocks (left) and the NULL bank (right); router preferentially allocates memory pathway weight to rare tokens.

Nearest neighbor analyses further reveal that, for rare tokens, MemoryBlock embedding neighborhoods are highly disjoint from those of the base embedding, demonstrating specialization and information enrichment.

Figure 11: NNS agreement between base and MemoryBlock embeddings: rare tokens exhibit substantially less overlap, confirming complementary information.

Practical Implications and Future Directions

By integrating token identity at every layer via EmbeddingMemory, TIDE corrects the undertraining of rare tokens and restores model capacity to preserve semantic distinctions otherwise lost due to contextualization. TIDE achieves these gains with minimal memory and compute penalty and is amenable to low-rank approximation and quantization, enabling practical deployment even at scale.

The theoretical perspective motivating TIDE extends beyond language modeling, suggesting opportunities in other domains where context-independent identity or static feature information must be persistently available to deep hierarchical architectures. Further, the block specialization observed in TIDE recommends future investigation into adaptive or interpretable memory-routing architectures. Potential integration with retrieval-augmented or MoE transformer variants could combine the strengths of persistent identity and external knowledge injection.

Conclusion

TIDE represents a rigorous architectural re-examination of how transformers handle token identity information. By rejecting the single-injection convention and injecting token-identity signals at each layer through independent EmbeddingMemory blocks, TIDE robustly addresses gradient starvation in rare tokens and contextual collapse in FFNs. The approach delivers consistent empirical improvements in both language modeling and downstream tasks, especially for rare and semantically ambiguous tokens, with practical and scalable implementation characteristics. The framework opens several pathways for future theoretical and engineering advancements in persistent memory architectures for deep sequence models.