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

TL;DR

The paper reveals that neural language models' next-token predictions are biased by a fixed, context-agnostic vocabulary key space, causing spurious correlations that inflate middle-ranked candidates. By combining knowledge probing, clustering of token embeddings, and visualization, the authors show that tokens similar in the key space to top predictions can be mistaken as plausible continuations. They introduce In-context Navigation (ICN), an iterative prompting method that pushes the query away from explored keys, improving knowledge-probing precision, and extending to open-ended generation and chain-of-thought tasks with higher diversity and self-consistency. The work also documents training-time risks from fixed key spaces, showing that fine-tuning mostly adjusts the query encoder and can propagate in-cluster biases, motivating reranking or contextualized vocabularies as potential remedies. Overall, the findings offer a retrieval-inspired lens on NTP and practical strategies to enhance decoding diversity and reliability while highlighting gaps in current fine-tuning regimes.

Abstract

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.

Figures (11)

• Figure 1: We showcase how next token prediction happens in the key space of the vocabulary embeddings. In the key space, there are clusters, like the capitalized character cluster, containing vocabularies similar to each other. This can introduce spurious correlation in next token prediction because key space similarity is (input) context-agnostic. Tokens with high similarity with top predictions are high-ranked by NTP but might not lead to correct decoding paths. In-context Navigation (ICN): We propose a simple in-context method to navigate the query away from the probed keys to efficiently explore the potential decoding paths underestimated by the initial ranking.
• Figure 2: Visualization of the relationship between key representations of first-tokens with their probing correctness in knowledge probing cases.
• Figure 3: The prompt format used for ICN in knowledge probing experiments.
• Figure 4: Exploration of the navigation ability of ICN. Left: Query similarity with the original query representation. Right: Query-key similarity with the top key representations corresponding to the original query.
• Figure 5: Exploring the impact of ICN frequency on knowledge probing.