Correlation and Navigation in the Vocabulary Key Representation Space of Language Models
Keywords: Language Modeling, Next Token Prediction, Spurious Correlation, Generation Diversity
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.
Primary Area: foundation or frontier models, including LLMs
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 246
Loading