TIDE: Every Layer Knows the Token Beneath the Context
Not on the map. Why?
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 42 of its 49 references that are in the map. First posted on arXiv 7 May 2026.
Builds on
42 of 49 references.
References · 42
- 2020Language Models are Few-Shot Learnerscited by 32,199
- 2020Retrieval-Augmented Generation for Knowledge-Intensive NLP Taskscited by 9,160
- 2023Toolformer: Language Models Can Teach Themselves to Use Toolscited by 4,072
- 2020Dense Passage Retrieval for Open-Domain Question Answeringcited by 3,779
- 2016Pointer Sentinel Mixture Modelscited by 2,992
- 2015Neural Machine Translation of Rare Words with Subword Unitscited by 2,898
- 2022Locating and Editing Factual Associations in GPTcited by 2,609
- 2020REALM: Retrieval-Augmented Language Model Pre-Trainingcited by 2,126
- 2021WebGPT: Browser-assisted question-answering with human feedbackcited by 1,722
- 2019Root Mean Square Layer Normalizationcited by 1,425
- 2020Transformer Feed-Forward Layers Are Key-Value Memoriescited by 1,307
- 2021Improving language models by retrieving from trillions of tokenscited by 1,278
30 more
- 2019PubMedQA: A Dataset for Biomedical Research Question Answeringcited by 1,271
- 2022Mass-Editing Memory in a Transformercited by 976
- 2022Atlas: Few-shot Learning with Retrieval Augmented Language Modelscited by 940
- 2019Generalization through Memorization: Nearest Neighbor Language Modelscited by 788
- 2017Reading Wikipedia to Answer Open-Domain Questionscited by 777
- 2020How Much Knowledge Can You Pack Into the Parameters of a Language Model?cited by 713
- 2014Neural Turing Machinescited by 630
- 2021Knowledge Neurons in Pretrained Transformerscited by 590
- 2020KILT: a Benchmark for Knowledge Intensive Language Taskscited by 571
- 2022Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Spacecited by 522
- 2018Lipschitz regularity of deep neural networks: analysis and efficient estimationcited by 413
- 2024DataComp-LM: In search of the next generation of training sets for language modelscited by 381
- 2014Memory Networkscited by 354
- 2016Using the Output Embedding to Improve Language Modelscited by 266
- 2023Outlier Weighed Layerwise Sparsity (OWL): A Missing Secret Sauce for Pruning LLMs to High Sparsitycited by 163
- 2022Analyzing Transformers in Embedding Spacecited by 131
- 2019Large Memory Layers with Product Keyscited by 120
- 2023LoSparse: Structured Compression of Large Language Models based on Low-Rank and Sparse Approximationcited by 103
- 2019Augmenting Self-attention with Persistent Memorycited by 92
- 2016Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modelingcited by 87
- 2024Mixture of A Million Expertscited by 70
- 2023Shall We Pretrain Autoregressive Language Models with Retrieval? A Comprehensive Studycited by 63
- 2024Understanding Factual Recall in Transformers via Associative Memoriescited by 45
- 2023The Emergence of Essential Sparsity in Large Pre-trained Models: The Weights that Mattercited by 42
- 2023LORD: Low Rank Decomposition Of Monolingual Code LLMs For One-Shot Compressioncited by 26
- 2023Cuttlefish: Low-Rank Model Training without All the Tuningcited by 25
- 2024FFN-SkipLLM: A Hidden Gem for Autoregressive Decoding with Adaptive Feed Forward Skippingcited by 23
- 2025Mixture of Lookup Expertscited by 11
- 2026STEM: Scaling Transformers with Embedding Modulescited by 10
- 2026MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformerscited by 3
Built on by
Not cited yet.