mlsanityData to 7 Oct 2026

TIDE: Every Layer Knows the Token Beneath the Context

arXiv 2605.06216 ↗First posted 7 May 2026cs.CL cs.AI cs.LGcited by 0

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

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

Built on by

Not cited yet.