# Correlation and Navigation in the Vocabulary Key Representation Space of Language Models

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 42 references that are in the map.
- First posted on arXiv 3 Oct 2024.

## Builds on (most cited first, 28 of 28)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Language Models as Knowledge Bases?](https://mlsanity.com/paper/1909.01066/) (2019)
- [Are NLP Models really able to Solve Simple Math Word Problems?](https://mlsanity.com/paper/2103.07191/) (2021)
- [Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval](https://mlsanity.com/paper/2007.00808/) (2020)
- [ZeroGen: Efficient Zero-shot Learning via Dataset Generation](https://mlsanity.com/paper/2202.07922/) (2022)
- [Chain-of-Thought Reasoning Without Prompting](https://mlsanity.com/paper/2402.10200/) (2024)
- [A Review on Language Models as Knowledge Bases](https://mlsanity.com/paper/2204.06031/) (2022)
- [From Decoding to Meta-Generation: Inference-time Algorithms for Large Language Models](https://mlsanity.com/paper/2406.16838/) (2024)
- [Word Embeddings Are Steers for Language Models](https://mlsanity.com/paper/2305.12798/) (2023)
- [ProGen: Progressive Zero-shot Dataset Generation via In-context Feedback](https://mlsanity.com/paper/2210.12329/) (2022)
- [Answer is All You Need: Instruction-following Text Embedding via Answering the Question](https://mlsanity.com/paper/2402.09642/) (2024)
- [Empower Entity Set Expansion via Language Model Probing](https://mlsanity.com/paper/2004.13897/) (2020)
- [BertNet: Harvesting Knowledge Graphs with Arbitrary Relations from Pretrained Language Models](https://mlsanity.com/paper/2206.14268/) (2022)
- [Crafting Interpretable Embeddings by Asking LLMs Questions](https://mlsanity.com/paper/2405.16714/) (2024)
- [Retrieval is Accurate Generation](https://mlsanity.com/paper/2402.17532/) (2024)
- [Incubating Text Classifiers Following User Instruction with Nothing but LLM](https://mlsanity.com/paper/2404.10877/) (2024)

## Built on by (most cited first, 1 of 1)

- [Linguistic Nepotism: Trading-off Quality for Language Preference in Multilingual RAG](https://mlsanity.com/paper/2509.13930/) (2025)

## Abstract (arXiv)

Language model (LM) decoding is based on the next-token prediction (NTP) probability distribution. For neural LMs (e.g., Transformer-based), NTP distribution is essentially a softmax-regularized dot product between an encoded input context (query) and fixed vocabulary representations (keys). In this paper, we study the effect of the key distribution on the NTP distribution, with a focus on whether the similarity between keys will trigger spurious correlations in NTP. Through knowledge-probing tasks, we show that in the NTP distribution, the few top-ranked tokens are typically accurate. However, the middle-ranked prediction is highly biased towards the tokens that are distributionally (not necessarily semantically) similar to these top ones. For instance, if "P" is predicted as the top-1 token, "A"-"Z" will all be ranked high in NTP, no matter whether they can lead to correct decoding results. This hurts the sampling diversity and makes the sampling of correct, long-tail results hopeless and noisy. We attempt to alleviate this issue via a novel in-context method that iteratively pushes the query representation away from explored regions. Specifically, we include the explored decoding results in the context and prompt the LM to generate something else, which encourages the LM to produce a query representation that has small dot products with explored keys. Experiments on knowledge-probing tasks show that our method leads to efficient navigation away from explored keys to correct new keys. We further extend our method to open-ended and chain-of-thought (for reasoning) generation. Experiment results show that ICN contributes to better generation diversity and improved self-consistency voting performance. Finally, we discuss potential training issues caused by the fixed key space together with the challenges and possible ways to address them in future research.