# TIDE: Every Layer Knows the Token Beneath the Context

arXiv: https://arxiv.org/abs/2605.06216 · HTML: https://mlsanity.com/paper/2605.06216/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 42 of its 49 references that are in the map.
- First posted on arXiv 7 May 2026.

## Builds on (most cited first, 30 of 42)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [WebGPT: Browser-assisted question-answering with human feedback](https://mlsanity.com/paper/2112.09332/) (2021)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [PubMedQA: A Dataset for Biomedical Research Question Answering](https://mlsanity.com/paper/1909.06146/) (2019)
- [Mass-Editing Memory in a Transformer](https://mlsanity.com/paper/2210.07229/) (2022)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Reading Wikipedia to Answer Open-Domain Questions](https://mlsanity.com/paper/1704.00051/) (2017)
- [How Much Knowledge Can You Pack Into the Parameters of a Language Model?](https://mlsanity.com/paper/2002.08910/) (2020)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Knowledge Neurons in Pretrained Transformers](https://mlsanity.com/paper/2104.08696/) (2021)
- [KILT: a Benchmark for Knowledge Intensive Language Tasks](https://mlsanity.com/paper/2009.02252/) (2020)
- [Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space](https://mlsanity.com/paper/2203.14680/) (2022)
- [Lipschitz regularity of deep neural networks: analysis and efficient estimation](https://mlsanity.com/paper/1805.10965/) (2018)
- [DataComp-LM: In search of the next generation of training sets for language models](https://mlsanity.com/paper/2406.11794/) (2024)
- [Memory Networks](https://mlsanity.com/paper/1410.3916/) (2014)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Outlier Weighed Layerwise Sparsity (OWL): A Missing Secret Sauce for Pruning LLMs to High Sparsity](https://mlsanity.com/paper/2310.05175/) (2023)
- [Analyzing Transformers in Embedding Space](https://mlsanity.com/paper/2209.02535/) (2022)
- [Large Memory Layers with Product Keys](https://mlsanity.com/paper/1907.05242/) (2019)
- [LoSparse: Structured Compression of Large Language Models based on Low-Rank and Sparse Approximation](https://mlsanity.com/paper/2306.11222/) (2023)

## Abstract (arXiv)

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.