Research map / Large language models / Language, safety and interpretability

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

ICLR 2025 · Large language models › Language, safety and interpretability · approach: Linguistic encoding in LLM representations

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that next-token prediction in LMs suffers from spurious key-space correlations among vocabulary embeddings, and proposes in-context navigation (ICN), which iteratively appends explored answers to the prompt to push the query away from explored keys, improving knowledge probing, open-ended generation, and chain-of-thought self-consistency.

Problem

Next-token prediction is a softmax over dot products between a context-dependent query and fixed, context-agnostic vocabulary keys; similarity between keys can inject spurious correlations that bias middle-ranked tokens and reduce sampling diversity and correctness.

Approach

Cluster vocabulary key embeddings to quantify spurious in-cluster correlation with top-ranked tokens, then use in-context navigation: append already-explored answers to the context and instruct the LM to generate something else, iteratively moving the query away from explored keys to reach correct new first-tokens.

Key results

• On ProbeSet with llama-3-8b-instruct, ICN raises MAP@50 for Sports League from 73.82 (vanilla) to 82.23 and for Scientist from 69.46 to 79.44.
• ICN improves self-consistency accuracy on GSM8K, SVAMP, and AQuA to 78.09, 85.10, and 58.66 versus 77.56, 84.20, and 57.09 for top-token probing.
• Fine-tuning on a correct token raises the probability sum of same-cluster tokens by +0.99% on average while correct tokens drop by -0.99%.

Measured effects reported in the paper

Surprising findings

• Key embeddings of llama-3-8b and llama-3-8b-instruct are nearly unchanged after SFT and RLHF, suggesting fine-tuning mainly adapts the query encoder.
• Fine-tuning on a correct token tends to increase probability of same-cluster tokens rather than other correct tokens.

Limitations

• ICN relies on the instruction-following and generation ability of the LM and may append incorrect results, misleading navigation.
• ICN requires multiple query encodings, lowering generation efficiency.
• Not applicable to weaker models such as GPT-2 that lack SFT.

Experimental setup

Datasets: ProbeSet (60 categories expanded from CGExpan) SST-2 AG-News Emotion GSM8K SVAMP AQuA CGExpan

Benchmarks: ProbeSet SST-2 AG-News Emotion GSM8K SVAMP AQuA

Metrics: MAP@50 precision accuracy average rank unique n-gram (UNG) ZeroGen accuracy (ZGN) self-consistency accuracy

Baselines: vanilla top-k probing rephrasing reranking repetitive sampling

Builds on: self-consistency ZeroGen CGExpan chain-of-thought prompting

Closest papers by meaning

• Implicit Optimization Bias of Next-token Prediction in Linear Models (NeurIPS 2024 · similarity 0.87)
• On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.87)
• Large Language Models Explore by Latent Distilling (ICML 2026 · similarity 0.87)
• From Tokens to Words: On the Inner Lexicon of LLMs (ICLR 2025 · similarity 0.87)
• L-MTP: Leap Multi-Token Prediction Beyond Adjacent Context for Large Language Models (NeurIPS 2025 · similarity 0.87)
• Rep2Text: Decoding Full Text from a Single LLM Token Representation (NeurIPS 2026 · similarity 0.86)