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

Language model (LM) decoding is based on the next-token prediction (NTP) probability distribution. For neural LMs (e.g., Transformer-based), NTP distribution isessentially a softmax-regularized dot product between an encoded input context(query) and fixed vocabulary representations (keys). In this paper, we study theeffect of the key distribution on the NTP distribution, with a focus on whetherthe similarity between keys will trigger spurious correlations in NTP. Throughknowledge-probing tasks, we show that in the NTP distribution, the few top-rankedtokens are typically accurate. However, the middle-ranked prediction is highly biasedtowards the tokens that are distributionally (not necessarily semantically) similar tothese top ones. For instance, if “P” is predicted as the top-1 token, “A”-“Z” will allbe 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-tailresults hopeless and noisy. We attempt to alleviate this issue via a novel in-contextmethod that iteratively pushes the query representation away from explored regions.Specifically, we include the explored decoding results in the context and promptthe LM to generate something else, which encourages the LM to produce a queryrepresentation that has small dot products with explored keys. Experiments onknowledge-probing tasks show that our method leads to efficient navigation awayfrom explored keys to correct new keys. We further extend our method to open-ended and chain-of-thought (for reasoning) generation. Experiment results showthat ICN contributes to better generation diversity and improved self-consistencyvoting performance. Finally, we discuss potential training issues caused by thefixed key space together with the challenges and possible ways to address them infuture research.