TIDE: Every Layer Knows the Token Beneath the Context
:::info Stub — Full Engineering Breakdown Coming This paper was featured on Hugging Face Daily Papers on 2026-05-07 with 5 upvotes. A full breakdown with production viability rating, implementation notes, and honest limitations is being written. Subscribe to AI Letters → :::
| Authors | Ajay Jaiswal et al. |
| Year | 2026 |
| HF Upvotes | 5 |
| arXiv | 2605.06216 |
| Download | |
| HF Page | View on Hugging Face |
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.
Engineering Breakdown
Plain English
This paper identifies a fundamental design flaw in how all modern LLMs handle token embeddings: tokens are converted to vectors once at the input layer and never revisited, even as context changes. This causes two problems—rare tokens get undertrained because they appear less frequently in training data, and similar tokens get mapped to nearly identical representations regardless of context. The authors propose TIDE, which uses multiple independent embedding memory blocks that can re-encode tokens based on their current context, rather than using static embeddings throughout the entire model.
Key Engineering Insight
Token embeddings are currently static one-shot lookups, but context fundamentally changes what a token means. By allowing multiple embedding layers to re-encode tokens throughout the model based on surrounding context, you can train rare tokens more effectively and maintain finer semantic distinctions between similar words.
Why It Matters for Engineers
Rare token undertraining directly impacts real-world model quality—code tokens, technical terms, and domain-specific vocabulary all suffer from insufficient gradient updates in production models. Contextual collapse wastes model capacity by forcing dissimilar tokens into overlapping representations. Fixing this could improve performance on long-tail vocabularies and specialized domains without scaling up the base model size.
Research Context
Prior work treated token embedding as a solved problem—pick a vocabulary size, initialize embeddings, and freeze the lookup table. This paper challenges that assumption by showing the single-injection design choice has measurable architectural consequences. TIDE extends the standard transformer paradigm to make embedding adaptive, similar to how recent work has made other 'fixed' components learnable, positioning token representation as a context-dependent rather than context-agnostic problem.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::