Correlation and Navigation in the Vocabulary Key Representation Space of Language Models
Authors: Letian Peng
Chenyang An
Jingbo Shang
Venue: ICLR 2025 Poster (2025)
Keywords: Language Modeling
Next Token Prediction
Spurious Correlation
Generation Diversity
Primary area: foundation or frontier models, including LLMs

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.

Published as a conference paper at ICLR 2025
CORRELATION AND NAVIGATION IN THE VOCABULARY KEY REPRESENTATION SPACE OF LANGUAGE MODELS
Letian Peng, Chenyang An, Jingho Shang (lepeng, c5an, jshang) @ucsd. edu University of California, San Diego
ABSTRACT
Language model (LM) decoding is based on the next-token prediction (NTP) prob- ability distribution. For neural LMs (c.g.. pransformer-based) NTP distribution essentially softmax-regularizeed dot product between encoded input context (query) and fixed vocabulary representations (keys). In this paper, we study the effect of the key distribution the NTP distribution, with a focus the similarity between keys will trigger spurious correlations NTP. Through snowledge-probingg tasks, show that in the NTP distribution, the few top-ranked tokens are typically accurate. However, the biased towards the tokens that are distributionally (not necessarily semantically) similar to these top ones. For instance, "p" predicted as the top-1 token, "A"-"Z" will all be ranked high in NTP, matter whether they 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 the context and prompt the LM to generate something else, which encourages the LM to produce query representation that has small dot products with explored keys. Experiments on shooweddgeprobiiggg 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. Experi- ment results show that ICN contributes to better generation diversity and improved self- 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. Code:
whether
prediction is highly
Show some black and white animals
Capitalized Character Cluster
Rank
First-token
Answen gredy path
The dbstabbutiona similartty the embedding space loads to quaiouus comrition in next token prodiction
Panda
2
Zebra
Show me black and white animals "Panda"
3
C
Cow
"Zebra"
Quokka
Sk
"Cow"
X
X enopus
4.
Rac
42
Rac
Racoon
We append the explored answers to the contest, which finds new query dissimilar prubed lokens but similer
43
Sk
Skunk
leading answers.
Next Token Prediction
Embedding (Vocabulary Key) Space
In-context Navigation (ICN)
Figure 1: We showcase how next token prediction happens in the key space of the vocabulary embeddings. In the key space, there 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.

---

Published as a conference paper at ICLR 2025
INTRODUCTION
Since the era of statistical language models (LMs) (Brown et al. 1992). LMs have long been decoded by next token prediction (NTP) probability based on statistics from real-world texts. Neural LMs (Bengio et al.) 2000 Sutskever et al., 2014; Brown et al. 2020: Dubey et al.) 2024) are particularly successful in NTP. Specifically, given input context, neural LM first encodes it into vector (query), and then calculates a softmax-regularin
dot product between the query and fixed vocabulary representations (keys), leading the predicted NTP distribution.
However, previous works have shown NTP distributions are highly restricted by the arrangement of key embeddings in a context-agnostic manner, such as probability bounds for less frequent words (Demeter et al. 2020) (unpredictable in extreme situations (Grivas et al.) 2022)). ). With the dis- covery of context- aggosticc embedding's ' limitation modeling multi- word distributions (Chang & McCallum 2022), we are motivated to explore and mitigate the potential spurious correlation in the embedding space. Different from the major focus on probability bounds in previous works, we emphasize more on distributional precision is NTP rank a good reference for diverse sampling?
Our investigation begins with the knowledge probing task (Petroni et al.) 2019 AlKhamissi et al., 2022 Hao et al. (022) that extracts factual knowledge from language models. Knowledge-probinn requests are typically like "Show me some black and white animals." " where the answers are not unique. Based on the request, we decode multiple first-tokens of the potential answer and approximate their correctness by whether they can lead to correct answers by greedy decoding or sampling. After probing the top first-tokens predicted by NTP, we visualize the distribution of their key representations. The results show the existence of many incorrect next tokens with high similarity to top (correct) next tokens, which might be introduced by spurious correlation in the key vocabulary space. Thus, cluster vocabulary embeddings in the key space, which shows many tokens without semantic relation lie in the same cluster. One example is the cluster of capitalized characters, "A"-'Z", which are mutually unrelated when they are subwords in generation. When probing "black and white animals", "P" is predicted as the top-1 token as "p" "Panda", which leads "Q" also to be high-ranked by its similarity to "P"". However, "Q" can hardly lead to correct answer for "black and white animals"' which showcases how spurious correlation in the key space affects the quality of NTP.
To quantify the spurious correlation, we compare the correctness of tokens with and without key space similarity to top-ranked tokens. For 60 probing prompts expanded from CGExpan (Zhang et al.) 2020), we group the next tokens ranked among 1 11-100 by whether they fall in the same clusters as the top-10 tokens tokens (in-top-cluster). Our result shows tokens in different clusters from the top-10 tokens
(out-of-top-cluster have higher accuracy than those cn-top-cluster tokens, but are ranked lower by NTP. This indicates that some in in-top-cluster tokens are overestimated by spurious correlation in the key space rather than really leading to correct decoding paths.
To alleviate such spurious correlation, we propose a simple yet effective method, in-context navigation (ICN), to efficiently push the query away from explored keys. Specifically, we explicitly append answers starting with probed first-tokens to the context and instruct the LM to generate different answers. Following the instruction, NTP will eliminate the probability of explored tokens, resulting in a low similarity between the new query and the explored keys. Consequently, this simple modification pushes the new query representation away from probed clusters to explore new ones containing the correct first-tokens. In comparison with simple rephrasing the prompt, ICN produces queries dissimilar from explored keys, which reflects a strong pushing-an awway ability of ICN.
We further benchmark the precision of ICN- based knowledge probing. By iteratively producing new queries away from explored spaces, we discover significant precision improvement in knowledge probing. We continue to extend this method to explore the potential first-token generation for open-ended generation ( (Ye et al.] 2022c a) and thahngtf generation (Wei et al. 2022) for self-consistency (Wang et al.) 2022). The results show higher generation diversity and reasoning accuracy than simply exploring top first-tokens.
Finally, we discuss training risks that might be caused by the fixed key space. By comparing the key space of large LMs before and after the large-scale fine-tuning, we observe the key space is almost unchanged, indicating the key representations have converged at the very early stage. This indicates that fine-tuming the language model is only learning the query encoder to push the query toward certain key, which unfortunately also pushes the query closer to incorrect keys similar in the
2

---

Published as a conference paper at ICLR 2025
key space. Our quantitative experiment shows when fine-tuning correct token, the knowledge is generalized to tokens in the same cluster (by increasing their probability) rather than other correct tokens. Thus, we also propose potential refinement methods for future works to further address the spurious correlation in inference and training in this paper. Our contributions are presented as follows,
We unveil the spurious correlation in the vocabulary key space for NTP, which introduces incorrect tokens into the prediction by their context-agnostic similarity in the key space.
We propose simple method, in-context navigation (ICN), to mitigate the spurious correlation using explicit context to search for new queries away from explored keys. ICN is beneficial to knowledge probing, open-ended generation, and seain-of-thougeee generation. We extend the discussion to large-scale model training, revealing that the space remains unchanged even during extensive fine-tuning. This also posts a question about the generalization ability of NTP learning to similar keys or to correct keys?
key
2
RELATED WORK
2.1 LANGUAGE MODELING
Language modeling has been an essential problem in long history of natural language processing, which has attracted more and more attention from both industry and academia since the boom of
large language models (Brown et al. 2020 Achiam et al.
[023] Touvron et al., 2023 Team et al.
2024). Similarly to their statistical ancestors Brown et al. (1992). these large neural language models are generally decoded from probabilistic view, such as random sampling, greedy decoding, beam
search, nucleus decoding (Holtzman et al. 2019 Welleck et al. 2024) and contrastive decoding (Li
cet al. 2023) that exploits predictions from weaker LMs.
However, there is also evidence indicating the limitation of the probabilistic view. The effectiveness of self-consistency (Wang et al. 2022) shows that the higher probability decoding path is not generally more correct. [Wang & Zhou also presents that probing the first-token even with a very low probability might lead to correct answer. Unfortunately, LMs might not be able to model such diverse decoding paths because of the softmax bottleneck (Yang et al.) 2018) as they are inclined to monomodal words (Chang & McCallum, 2022). The probabilities of less frequent words also likely to be bounded when falling inside convex hulls in the distribution (Demeter et al.] (020),, causing some words to be unpredictable (Grivas et al. (2022). Based on previous efforts, our work further illustrates and mitigates the potential spurious correlation phenomenon in NTP.
2.2 VOCABULARY KEY SPACE OF LANGUAGE MODEL
In language models, the next tokens are represented by fixed vectors to calculate their similar- ity with encoded context, which is similar dense retrieval in information retrieval. Dense re- trieval (Karpukhin et al.] 2020 Xiong et al. 2020) is an essential technique to retrieve relevant documents according to user's ' request, which encodes requests&dooumments into query&key rep- resentations. When user request comes, it is first transformed into query representation whose dot product values with key representations are ranked determine the most relevant documents. Next token prediction in neural language models can also be discussed from this view that the context is encoded into query representation, whose dot product values with the vocabulary's key representations determine the most possible next tokens (Cao et al. 2024). There several trials to improve language modeling by modifying the key space, Adapting key representations by vector transformation is proposed to steer generation (Han et al. (2044). Phrase-augmented language models (Lan et al. 2023 Cao et al. 2024) whose vocabulary is augmented by encoded phrases have been proposed for a more diverse open-ended generation. Our work is inspired by the success of language modeling as retrieval and delves into how vocabulary embeddings are distributed.
2.3 DECODING PATH PROBING
There are generally multiple decoding paths that lead to correct responses (Wang et al. 2022). Probing these paths is important for harvesting knowledge, sampling diverse passages, and exploring multiple ways of reasoning from language models. Knowledge-prote (Petroni et al.) 2019; AlKhamissi
et al. 2022 Hao et al. 2022) prompts language models with factual questions and collects the top
3

---

Published as a conference paper at ICLR 2025
answers as knowledge. Sampling diverse passages is important for data synthesis like training text classifier (Ye et al. 2022a]b). Voting with diverse chain-of of Chain-of-thouuht reasoning paths is also promising in improving the reasoning accuracy (Wang al., 2022; Wang & Zhou, 2024p. There are also trials to encode the feature of contexts with the possible answer paths according to questions (Peng
et al.; [2024] [enara et al. 2024). Our work explores how path probing is affected by the potential
correlation in the key space and proposes corresponding methods to address the issue.
3 PRELIMINARY AND MOTIVATION
Language models LM( are trained to make NTP which produces the probability distribution of the next token to generate. Given an input context X = [11, 22, .NN.. where n = |X represents the number of input tokens. The language model LM(X) outputs a probabilistic distribution P(
of the vocabulary V with probabilities of all vocabularies the next token.
For neural language models, the NTP probability essentially produced by representation calculation. Inside the language model, the encoder E(.) encodes X into a D-dimension representation Rx, which is the same as the vocabulary representations Ry R The dot product between Rx and Ry produces logit L E R RVI, which is regularized by softmax function to finally output the probabilistic distribution P(5n+1) (X).
The fixed key representations might introduce unexpected spurious correlation into NTP because the vocabularies are assigned with context-agnostie similarity. For instance, "p" and "Q" generally show high similarity in the key spaces of different LMs because they are both capitalized characters. However, this should be viewed as a spurious correlation in many contexts (e.g. generating subwords). In the following sections, we will use experiments to visualize and quantify the severity of this issue.
4 SPURIOUS KEY CORRELATION
4.1 EXPERIMENT SETUP
Knowledge probing (Petroni 2019 AlKhamissi et al. 2022: Hao al. 2022) is task that aims to extract as much as possible knowledge from LMs (e.g. "Show me some computer scientists." "). We select the knowledge probing task to investigate the spurious correlation issue because there are many first-tokens that can lead to correct decoded answers and the answers are easy to validate. Our probing experiment aims to unveil how the correct next tokens are distributed in the key space and how the distribution affects the NTP result. We illustrate the result on strong open-source L.M,
1lama - 3--3--8-- 8b inst ruct Dubey et al. 2024), and include the result on other LMs in Appendix H
We follow Zhang et al. 2020) for the probing tar- gets, and extend its coverage into 12 categories, such "Scientist", "Astronomical Object", and "Sports League". To further challenge LLM, we broaden the comprehensiveness of the probing by adding extra sub-categories for each cat- egory (12
Scientist
Food
City
Computer Scientist
Chinese Food
Canadian City
Social Scientist
Spicy Food
Asian City
Female Scientist
Veggic Food
Coastal City
German Scientist
Dessert
Capital City
12 60 probing categories in total). For instance, "Computer Scientist" is included as an extra sub-category probing exper- iment for "Scientist" probing. In the experiment, we refer to the expanded knowledge probing set as ProbeSet. Some categories and sub-categories involved in our experiments are presented in Table We explore the NTP of the first-token in the answer and use greedily decoded answers (Monte-Care sampling results also included Appendix EE) to approximate the decoding space for experiment efficiency like in Wang & Zhou (024)). For validation, we prompt the state-of- the-htt-trtt LM, GPT 4o, whose accuracy is verified to be around 95% by annotated reference in Zhang et al. (2020) and human consistency evaluation shown in Appendix The prompts used can be found in Appendix K
Table 1:
Examples of categories and sub- categories. Full list can be found in Appendix C
4.2 CASE VISUALIZATION
The first stage in our experiment is to visualize the distribution of predicted first-tokens in the key space. Thus, we select three cases, "Scientists", "Astronomical Objectives", "Scientists", to visualize
4

---

Published as a conference paper at ICLR 2025
how correct/incorrect next tokens are distributed. We validate the top-100 predicted next tokens and apply t-distributed Stochastic Neighbor Embedding (t-SNE) (Van der Maaten & Hinton 2008) to reduce the dimensions of corresponding key representations together with the query representation (encoded context). We select E-SNE because of its ability to maintain the cosine similarity relationship in the high dimension space.
Sports Leagues
Astronomical Obiects
Scientists
Query (Carme)
(henneer)
Figure 2: Visualization of the relationship between key representations of first-tokens with their probing correctness in knowledge probing cases.
we illustrate the visualization of different probing cases. We can observe the query to be encoded near several correct key representations leading to decoding paths to correct answers, which correspond to the top next tokens in NTP. However, there are also incorrect next tokens with high similarity with these top tokens, which are consequently closer to the query representation than other correct tokens. This supports the hypothesis of spurious correlation in NTP, will further verify the issue existence via metric quantification.
In Figure 2
4.3 ISSUE QUANTIFICATION
We first propose metric to descrihe how NTP is impacted by the spurious correlation in the
CID In-Cluster subwords
key space. The ultimate goal of our metric is to depict the difference between the next tokens
896
A, B.
200
0, 0.1.2.345.......0... 20. 10. 201....
640
Gwonder, Gbeautiful. Gamazing, Gexcellent. ...
that are spuriously correlated to top tokens and those aren't. The spurious correlation is inde-
996
urn. Gecho, printf, echo, Gprintf, decho,
100 elf, cld, elp, els, EI,, elt, elves, ael, EI, elay,
pendent on the query and only dependent on how keys are similar to one other. Thus, we run a clustering algorithm on the key represen- tations to divide the vocabularies into 1024 clus- ters. Specifically, we select the K-means (Lloyd. 1982) algorithm because it outputs clusters in the same size, which indicates the same amount of spurious correlation. We showcase some clusters in Table 2 with corresponding in-cluster subwords. The clusters include capitalized characters from "A" to "Z" (CID=896), and numbers from "O" to "9" (CID=896). There are highly explainable clusters like positive adjectives (CID=640) but there are also clusters without valid reason for similarity. especially for subwords (CID=100).
350 loh, Glunch, Glens, Glip, Globby, Glaptop,
Table 2: Examples of the next token clusters in key space clustering of LLaMA-3. CID: The cluster identifier for reference in the paper. "G" and "e" are special tokens, which are decoded into blanks.
Based on the clusters, we design the metric to quantitatively analyze the spurious correlation. Our focus is on the middle-ranked next tokens that are ranked high but not top in NTP, specifically, from all tokens in top-N (N 100) but not in top- K (K 10). The top- next tokens are viewed as the top tokens that might inject the spurious correlation. The middle-ranked tokens falling in the same cluster as the top tokens are supposed to be affected by the spurious correlation. The middle-ranked tokens are thus divided into two groups: in-top-cluste cluster (InTop) and out-of-top -oo--tutee (OutTop).
InTop (, = C(o,) ^j s K) (K < (i
Out = {w (v; & InTop) A (K <is
where C(-) returns the cluster of the i-th (ranked by NTP) token 02- We accuracy and average rank to compare them in performance and distribution. Besides, we illustrate the proportion of the two groups to show the results are not based on very few data points.
5

---

Published as a conference paper at ICLR 2025
Calegory
Sports
League OutTop
Astronomical
Scientist
Landmark
Country
City
Group
InTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
Accuracy
50.12
60.36
49.21
71.80
50.62
64.17
70.64
76.65
68.61
65.50
69.87
Rank
54.01
56.97
54.42
56.65
52.70
58.51
54.79
56.10
54.69
55.71
52.00
59.
Proportion
50.89
49.11
54.00
46.00
53.11
46.89
48.67
51.33
41.78
58.22
51.78
48.22
Category
TV Channel
Restaurant
Company
Creature
Discase
Food
Group
InTop
OurTop
InTop
OutTop
InTop
OurTop
InTop
OutTop
InTop
OurTop
InTop
OutTop
Accuracy
60.33
62.97
71.15
75.97
76.41
81.00
72:87
77.66
65.55
71.44
68.00
73.66
Rank
54.59
56.96
52.03
59.39
52.04
59.82
50.54
61.69
50.50
59.64
50.07
59.87
Proportion
46.89
53.11
51.11
48.89
54.00
46.00
36.41
43.56
46.22
53.78
46.00
54.00
Table 3: Knowledge probing results ProbeSet. The shown metrics calculated by an average over main category + 4 sub-categories. We also provide experiments in probing "words starting with given characters", whose result is independent of LLM evaluator in Appendix(]
We present the knowledge probing result in Table across the 12 categories, which shows that the OutTop first-tokens not only have hetter accuracy but are ranked lower as well. Thus, the division verifies the existence of spurious correlation, which leads group of less accurate tokens to be ranked higher. The balanced proportion indicates this not to be an issue stemming from just a few isolated cases, but a general phenomenon that hinders the NTP. Following the observation above, will discuss potential ways to mitigate the spurious correlation in the following sections.
IN-CONTEXT NAVIGATION
5.1 METHODOLOGY
From previous experiments, the middle-ranked tokens are shown to be the victims of spurious correlation. Thus, our strategy is to decode the LM multiple times, each time only for the relatively accurate top tokens. A simple method to rephrase the prompt, which perturbs the query in the representation space search for new top tokens. Our proposed method, in-context navigation (ICN), inherits this idea and steps further to navigate the query represen- tation away from the explored key representations with an explicit instruction in the context. For instance, when we have explored "Alan Turing" for "Computer Scientist", a prompt with ICN can be "A computer scientist other than Alan Turing is", which discourages the LM to generate "Alan Turing" and consequently climinates the similarity between the query representation and the key representation of "Alan". In our experiments, we use the format shown Figure to handle long list of explored keys.
Prompt Format
Show me some computer scientists.
nswer:
1. "Alan Turing"
Explored
2. "Charles Babbage"
Figure 3: The prompt format used for ICN in knowledge probing experiments.
Query Similarity with Vanilla Query
Query-Key Similarity with Top Keys
Vasilla
XCN
Replaaa:
Density
6/5
Cosine Similarity
Cosine Similarity
Figure 4: Exploration of the navigation ability of ICN. Left: Query similarity with the original query Query-key similarity with the top key representations corresponding to the original query. Rephrasing: Some rephrasing cases shown in Appendix K
representation. Right:
We first conduct experiments to certificate the navigation ability of ICN based on ProbeSet. For each instruction, we decode the top-10 first-tokens and append the decoded result to the context, which is encoded to the new query representation. We evaluate two types of similarity: 1) the similarity with the original query representation and 2) the average similarity with the top-10 key representations corresponding to the original query. We include simply rephrasing the probing prompt for comparison. The results are plotted in Figure by the distribution curve. The left subfigure illustrates the query similarity, which shows the query is successfully navigated away from the
bascline

---

Published as a conference paper at ICLR 2025
original one by ICN while simple rephrasing still leads the new query to position near the original one (with cosine similarity close to 1.0). The right subfigure shows ICN to be better at navigating the query away from explored keys in comparison with the simple rephrasing. Thus, we certificate the ability of ICN to navigate the query to a different location away from probed keys.
The next step is to verify the aceuracy of the navigated query, i.e. the correctness of new top first- tokens. Specifically, we compare the probing accuracy between direct decoding and decoding with ICN. To utilize ICN, we introduce the iterative ICN traverse through key representations. In each iteration, we will decode the path for top first-tokens and append them to the context, which encoded new query for the next iteration. Also, the probed first-tokens are skipped in future iterations as each first-token should be probed only once for comparison.
For iterative ICN, the two key parameters are the number of encoded queries (#Query) and the number of probed top keys (#Key) in each iteration, whose multiplication (#Query #Key) should be equal to the number of probed paths. When #Query 1, the procedure equals direct decoding, which is taken as the baseline result. When #Key 1, the procedure equals to prompting the LM to list the answers.
Sports Leagues
A Astronomical Objects
Scientists
23382
RK.ey
94.27
24.12
52.23
78,02
70 44
12
18
#Query
#Query
#Query
Figure 5: Exploring the impact of ICN frequency on knowledge probing.
We depict probing result of different (#Query, #Key) pair for iterative ICN in Figure 5 We set the number of probing path to be 50 and present the MAP@ 50 (Everingham et al. 2010) value as the metric, which considers the rank of correct predictions. For comparison fairness, all probed first-tokens by navigated queries are appended back to the initial prompt. Thus, the generation spaces are kept the same for different experiment configurations.
The experiment results show the probing performance with ICN to outperform direct decoding, which indicates ICN to be able to navigate the query close to correct keys. On the other hand, navigating queries too frequently (like listing the answers when #Key= 1) also does not lead to the best performance, which indicates the importance probing multiple top tokens. In general, the best performance is achieved with a balanced value between #Key and #Query, for instance, #Query = 10 and #Key = 5. In conclusion, navigating queries away from explored keys by ICN is beneficial to mitigate the spurious correlation.
5.2 MAIN COMPARISON
Based on the verification of ICN's effect, we further systematically use ProbeSet to benchmark the probing performance of different methods against the spurious correlation in the key vocabulary space. We continue using the ProbeSet for benchmarking. We still probe 50 decoding paths by using MAP@50 and precision as the metric. Based on the conclusion in the probing experiment, we set #Query to 10 and #Key to 5. Similar with previous experiments, the first-tokens probed by ICN or other methods are appended to the initial probing prompt to eliminate the influence of extra context towards a fair comparison.
Baselines We include the vanilla method, which simply probes decoding paths of top-50 first-tokens. The rephrasing method mentioned in the navigation ability evaluation replaces the appending explored paths in ICN by rephrasing the probing instruction to navigate to different first-tokens. We rephrase each instruction for 10 times and probe top-5 (not probed) keys. Another reranking baseline penalizes tokens by the number of probed tokens in the same cluster. Specifically, the rank of each token will be added by its rank in its cluster. For instance, when the first 5 tokens in cluster 1, 1,
7

---

Published as a conference paper at ICLR 2025
1, 2, 2, their in-cluster ranks should be 2. 2. Consequently, the added ranks will be 2, 4, 6, 5, 7. which switch the rank between the initial third and fourth tokens. When the added ranks are equal, the initial rank is considered as the prior one in our implementation. The reranking method explores whether we can eliminate the spurious correlation without searching for multiple queries.
Category Metric
Sports League PREC
Astronomical MAP
Scientist
Landmark
Country
City
MAP
PREC
MAP
PREC
MAP
PREC
MAP
PREC
MAP
PREC
Vanilla
73.82
60.40
75.98
72.80
69.46
62.40
84.48
79.60
81.45
74.40
91.57
85,20
Rephrase
73.78
59.60
76.69
74.00
72.15
62.80
84.31
79.60
81.94
74.80
92.05
85.60
Rerank
74.82
61.00
77.57
74.00
70.79
63.20
84.82
80.40
81.81
75.40
92.42
85.00
ICN
82.23
71.60
78.02
68.40
79.44
76.40
89.44
84.00
92.57
82.80
96.38
93.20
Category
TV Channel
Restaurant
Company
Creature
Discase
Food
Metrie
MAP
PREC
MAP
PREC
MAP
PREC
MAP
PREC
MAP
PRFC
MAP
PRFC
Vanilla
79.56
71.20
88.80
81.60
87.63
82.80
95.73
91.00
85.55
77.20
81.43
76.40
Rephrase
80.87
72.00
89.35
81.60
86.83
81.60
89,59
84.80
85.48
76.80
80.90
76.80
Rerank
79.86
71.20
89.60
83.40
88.23
84.60
96.19
90.80
85.96
78.60
82.16
78.00
ICN
87.70
82.80
93.70
84.00
97.60
94.40
97.60
92.00
88.37
78.80
93.53
91.20
Table 4: Knowledge probing results with different methods to mitigate spurious correlation.
The experiment result in Thole that ICN significantly outperforms the vanilla probing strategy. which is consistent with the conclusion in the probing experiment. Rephrasing generally does not show much difference with the vanilla method (sometimes higher and sometimes lower), suggesting it not able to navigate the query to new correct keys. This is consistent with the indication from the similarity experiment, which shows a high similarity between the initial and the rephrased query. Finally, the reranking method achieves small yet consistent improvement across all experiments, which again verifies the existence of spurious correlation. However, only reranking significantly underperforms ICN, which emphasizes the importance of multiple queries. Still, the efficiency advantage from single time encoding maintains the usage of reranking for efficiency.
5.3 OPEN-ENDED GENERATION
We further explore the usage of ICN beyond knowledge probing. The first task is open-ended generation, which differs from knowledge probing by generating sentences instead of entities. Thus, the framework of ICN is kept with the instruction changed to generate sentences. The evaluation concentrates on the diversity and usage of the generated sentences. For diversity, we select unique n-gram (UNG) (Buck et al. 2014) as the metric, which counts the proportion of unique n- n-grams, averaged over n. = ~ 4. For usage, we evaluate the classifier trained on the generated texts on test sets annotated by humans. This scenario has been proposed as ZeroGen (Ye et al., 2022b), so we name the metric as ZeroGen accuracy (ZGN), As more diverse datasets train better classifiers (Peng & Shang 2024), the result simultaneously reflects the semantic diversity of text generation.
We select 3 datasets for evaluation, SST-2 (positive, negative) (Socher et al. 2013), AG-News (World, Sports, Business, SciTech) (Zhang et al., 2015). Emotion (Sadness, Joy, Anger, Fear, Love, Surprise) (Saravia et al. 2018). Baselines include repetitively sampling sentences from the prompt, which is generally applied in existing text generation scenarios. Another baseline probed top first- tokens, which corresponds to the vanilla method in the knowledge probing experiments. All methods generate 100 sentences for diversity evaluation and classifier training (selected as RoBERTa-Large, hyperparameters listed in Appendix K). For ICN, we set #Query 10 and #Key 10.
The results are presented in Table 5 we can observe the ICN achieving consistent improvement on text generation with bet- ter UNG diversity and ZGN accuracy. In comparison with repetitive sampling, prob- ing different first-tokens shows better perfor- mance, which indicates the influence of the first-tokens in generation even for sequences (sentences) longer than entities. The advantage of ICN over simple top probing be similarly explained as the diverse yet correct first-tokens explored by ICN leading to more diverse sequence generation.
SST-2
AG-News
Emotion
Method
UNG
ZGN
UNG
ZGN
UNG
ZGN
Sample
58.95
83.20
48.17
38.37
41.39
38.37
Top ICN
61.29
84.90
50.97
40.63
50.85
40.63
64.74
86,33
53,37
42.19
53.34
42.19
Table 5: Open-ended generation results. "Top" ex- plores decoding paths with top predicted first tokens.
8

---

Published as a conference paper at ICLR 2025
5.4 CHAIN-OF-THOUGHT GENERATION
2022) refers to the reasoning chains in complex tasks such as math problem solving. Similar to open-endeee generation, diversity influences the success rate of large model reasoning (Naik et al.
Chain-of- Chain-of-Thought (CoT) (Wei ct al.
2023) when answers from different CoTss are merged for self-consistency (Wang et al. 2022). When multiple CoT's are generally sampled by multiple times from large languages, Wang & Zhou propose better way to probe the reasoning path starting from different first-tokens. Their strategy probes the top tokens, similar to the baseline setup in our knowledge probing experiments. Consequently, we apply ICN to this CoT generation framework by appending probed CoTs in the context to probe new first-tokens. (Remind that explored CoTs will be removed when generating the new CoT, which prevents copying the answer in explored CoTss)
For benchmarking, we use math problem solving datasets from the initial CoT experiments (Wei et al. 2022). which are GSM8K (Cobbe al. 2021), SVAMP (Patel et al., 2021). and AQuA (Deep Mind 2017). GSM8K and SVAMP directly ask for numeric answers while AQuA contains multiple- choice math questions, which select answers from 5 candidates. For self-consistency, we the number of CoTs to 4. The hyperparameters for ICN are set to #Query = and #Key = 1. For self-consistency, the most voted answer is selected the final answer.
In Table 6 we illustrate the reasoning performance of celf-consistency with different strategies. Consistent with results on previous tasks, our ICN contributes more to ooffcconsistecyy than sampling and probing only top tokens by proposing diverse and accurate CoTIs, which verified by 3 datasets. While the first generated tokens might be considered to have limited impact on the whole CoT quality, our result (together with Wang & Zhou (2024) suggests the benefit in probing them. On the other hand, the benefit from ICN to CoT generation is not as significant as knowledge probing, which indicates lengthy generation might weaken the benefit from ICN.
Method
GSM8K
SVAMP
AQuA
w/o SC
74.15
79.00
53.54
Sample Top ICN
76.50
83.10
55.51
77.56
84.20
57.09
78.09
85.10
58.66
Table 6: Reasoning benchmark results. SC:
Self- Seff-Consistency
6 LM TRAINING RISKS FROM FIXED KEY SPACE
Our previous contents mainly concentrate on the impact key space during the inference time. In this section, we will dive deeper into the potential influence of the query-key matching procedure during training neural language models.
We first illustrate an important property of the key space, its convergence after the large scale pre-training. Specifically, we compare the key spaces between 1lama-3-88 and 1ama-3--b-inst
is based on 11ama 3-8b with further supervised fine-tuning (SFT) and reinforcement learning with human feedback (RLHF). We evaluate the similarity three scenarios, 1) Token Similarity. calculating the cosine similarity between key representations of the same vocabulary (before or after SFT&RLHF), which is applied to all vocabularies, 2) Pair Similarity Difference, calculating the difference cosine similarity between the same pair of vocabularies (before or after SFT&RLHF), which is applied to 100,000 randomly sampled pairs. 3) Similarity Rank Difference, calculating the similarity between a token with all vocabulary tokens, then calculating the spearman correlation between the similarity distribution of the same vocabularies (before or after SFT&RLHF), which is applied to 10,000 randomly sampled vocabularies.
11ama-3 -ew---iinnsttreee
Token Similarity
Pair Similarity Difference
Similarity Rank Difference
Density
0.994
0.994
0.098
Cosine Similarity
ACosine Similarity
Spearman Correlation
Figure 6: Similarity between 1lama 3- 8b and 1lama-2 3- 3-8b 1e------inttt iinstruct in key spaces.

---

Published as a conference paper at ICLR 2025
As the 3 evaluation results presented in Figure 6 observe the key spaces before and after SFT&RLHF have very high Token Similarity and Similarity Rank Difference, together with almost zero Pair Similarity Difference. This indicates the key space hardly changed after SFT&RLHF, even though these stages also include numerous training data. A highly possible explanation is the shallow network (only embedding layer) to encode vocabularies can only capture some spurious correlation between them as shown in Table 2 ignoring the complex interaction between queries and keys in language modeling. Given the early-converged key space and the high performance difference between the models before and after SFT&RLHF, we conclude the fine-tuning stages are mainly learned to encode the context to queries but hardly adjust the key space.
Based on the conclusion above, we would like to further point out the potential vulnerability in fine-tuning large models. As only queries are effectively adjusted, the ability of language models to store multiple knowledge is questionable, especially when two correct NTP answers are in different vocabulary clusters. We quantify this question as "When correct next token is used for fine fine-tuning, it is generalizing to (increasing the probability of) other correct next tokens generalizing to other next tokens in the same cluster?" To answer this question, we go back to the knowledge-probing task and fine-tund (learning rate is set to 10 10-6) the large model on the correct top- 10 next tokens. For each optimization step, we calculate the probability sum difference between groups of tokens. For comparison, the first group of tokens is those in the same cluster the one used for fine-tuning, and the second one is just the correct next tokens. For experiment efficiency, we only include the top-100 tokens before fine-tuning into the groups.
We illustrate the probability differences by NTP among instructions in ProbeSet in Figure 7 which shows in-cluster tokens to benefit more from fine-tuning than correct tokens. In fact, while the correct next tokens suffer from -0.99% drop in probability sum on average, the in-cluster tokens reversely get +0.99% lift in probability sum. This indicates learning cor- rect knowledge does not naturally generalize to other correct knowledge, but with a high poten- tial to generalize by spurious correlation in the key space. This discovery challenges the ability of language models reflect the real world as they might generalize to hallucination injected by the spurious correlation in the key space.
Probability Sum Difference of Token Groups
0.16-
brapellate Tikew -a80
Cersa (1
0.13-
0.10-
016-
am--
0.0
ore
-10
15
AProbability (%)
Figure 7: The distribution of probability switching of different next token groups in fine-tuning the language model.
Finally, we would like to propose some potential ways further address the spurious correlation by editing the language modeling framework.
Adding a reranking stage, which decodes top next tokens and uses reranker rescore them based on the context. Reranker is a commonly applied module in information retrieval and NTP can be viewed as a retrieval stage (query-key matching) from the view of information retrieval. The reranker scores with the predicted token in the context, allowing the next tokens to interact with the context to produce less biased NTP. possible challenge is the ability of the reranker to recognize whether rather non-informative tokens (like subwords) can lead to correct decoding path.
Adding a contextualization layer for vocabularies, which adjusts the distribution of key vocab- ulary representations based on the context as an input. This strategy has potential as the query representations are well contextualized by the Transformer architecture, which can be extended to contextualize the key vocabularies. A potential challenge is the cost to contextualize the large scale vocabularies, which requires multiple times of interactions between them and the input context.
7 CONCLUSION AND FUTURE WORK
In this paper, we unveil the potential spurious correlation in the key vocabulary spaces of neural language models for next token prediction. We use knowledge probing experiment to verify the potential issue and correspondingly propose in-context navigation for better token probing. We show in-context navigation can be extended to benefit open-ended and -fain-of-thougee generation. Finally, we discuss the further impact of the spurious correlation on language models and propose potential ways to address issues for future works.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Badr AlKhamissi, Millicent Li, Asli Celikyilmaz, Mona Diab, and Marjan Ghazvininejad. A review on language models as knowledge bases, 2022. URL https: org/
wtwp/aa/////.....ae abs hbtg://b//22000
06031
Vinamra Benara, Chandan Singh, John X. Morris, Richard Antonello, Ion Stoica, Alexander G. Huth, and Jianfeng Gao. Crafting interpretable embeddings by asking Ilms questions. CoRR, abs/2405. 167144 2024. doi: 10448550/ABXIV.2205..6777 URL
48550/arXiv 2405. 16714
Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. A neural probabilistic language model. Advances in neural information processing systems, 13, 2000.
Peter F Brown, Vincent J Della Pietra, Peter V Desouza, Jennifer C Lai, and Robert L Mercer. Class-based n-gram models of natural language. Computational linguistics, 18(4):467-480 1992.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
2020.
Christian Buck, Kenneth Heafield, and Bas van Ooyen. N-gram counts and language models from the common crawl. In Nicoletta Calzolari, Khalid Choukri, Thierry Declerek, Hrafn Loftsson, Bente Maegaard, Joseph Mariani, Asunción Moreno, Jan Odijk, and Stelios Piperidis (eds.), Pro- ceedings of the Ninth International Conference on Language Resources and Evaluation, LREC 2014, Reykjavik, Iceland, May 26-31, 2014, pp. 3579-3584. European Language Resources Associ- ation (ELRA), 2014. URL http:/ttp/..... /wwp: www. rec conf wttp:/oowogooo.eeenn//////////// es 1097.htm httll
Bowen Cao, Deng Cai, Leyang Cui, Xuxin Cheng, Wei Bi. Yuexian Zou, and Shuming Shi. Retrieval is accurate generation. arXiv preprint 2rXiv:2402.15532. 2024. URL https: //tp:iv.//xxovvvg//
2aa//240.11500
Haw-Shiuan Chang and Andrew McCallum. Softmax bottleneck makes language models unable to represent multi-mode word distributions. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings the 60th Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pp. 8048-8073. LingiuttiisnLiLLiiiiisiiiiiii Association for Computational Linguistics, 2022. doi:
URL
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers solve math word problems. CoRR, abs/2110.14168, 2021. URL
https: Settp://arriiz owapiii.orrraree abs wwt.oo/ggggbbaoooo
14168
Google DeepMind. Aqua-rat (algebra question answering with rationales) dataset. https://
github. httpbb/ccooogoooleeeeeeeeeee AAQuAAA 2017.
David Demeter, Gregory Kimmel, and Doug Downey, Stolen probability: A structural weakness of
neural language models. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault (eds.),
Proceedings of the 58th Annual Meeting the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pp. 19912197. Association for Computational Linguistics, 2020. doi: 12200//11/22.AAAAAAAAAAA 198. URL wttph:ttp//// /tt::///o...oo wrt. 10. 11t.1..1111533//222222
ac mal-maine n. 1e..an1119. 198
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv preprint 207.22240..1178. 2024.
11

---

Published as a conference paper at ICLR 2025
Mark Everingham, Lue Van Gool, Christopher K. I. Williams, John M. Winn, and Andrew Zisserman. The pascal visual object classes (VOC) challenge. Int. J. Comput. Vis., 88(2): 303-338, 2010. doi: 10.1007/11223300000 026.099002 URL https:/ /tt://doi. www..ooooo...000000/000
s s11263 009-0275-----
Andreas Grivas, Nikolay Bogoychev, and Adam Lopez. Low-rank softmax can have unargmax- able classes theory but rarely in practice. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume Long Pupers), ACI. 2022, Dublin, Ireland, May 22-27, 2022, pp. 6738-6758. URL
Association for Computational Linguistics, 2022. doi: 1211666//1120222ALLLLLL
https://doi /ttt:://bo..00331011166333333333
long 465
Dirk Groeneveld, Iz Beltagy, Evan Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, Shane Arora, David Atkinson, Russell Authur, Khyathi Raghavi Chandu, Arman Cohan, Jennifer Dumas, Yanai Elazar, Yuling Gu, Jack Hessel, Tushar Khot, William Merrill, Jacob Morrison, Niklas Muennighoft, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Valentina Pyatkin, Abhilasha Ravichander, Dustin Schwenk, Saurabh Shah, Will Smith, Emma Strubell, Nishant Subramani, Mitchell Wortsman, Pradeep Dasigi, Nathan Lambert, Kyle Richardson, Luke Zettlemoyer, Jesse Dodge, Kyle Lo, Luca Soldaini, Noah Smith, and Hannaneh Hajishirzi. Olmo: Accelerating the science of language models. In Lun-We Ku, Andre Martins, and Vivek Srikumar (eds.). Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), ACI. 2024, Bangkok, Thailand, August |1- 1-16, 2024, pp. 15789-1 15809. Association for Computational Linguistics, 2024. doi: 10.18633W120044A4L--LLOO
URL /dttps /tt.///10..o
18653/ v0.33v11220 202. acl long 841
Chi Han, Jialiang Xu, Manling Li, Yi Fung, Chenkai Sun, Nan Jiang, Tarek F. Abdelzaher, and Heng Ji. Word embeddings are steers for language models. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational
Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16,
2024, pp. 16410-16430, Association for Computational Linguistics, 2024, doi: 10. 116538//////224, ACL-LONG.866
htLpstts:
(tlp://1...153vvv11 http://////.//2222.1111111111 acl- acl-longl 864
Shibo Hao, Bowen Tan, Kaiwen Tang, Bin Ni, Xiyan Shao, Hengzhe Zhang, Eric P Xing, and Zhiting Hu. Bertnet: Harvesting knowledge graphs with arbitrary relations from pretrained language models. arXiv preprint arXiv:2206. 2022.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint arXv:: 2019.
Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXiv:2004.04900 049900 2020.
Tian Lan, Deng Cai, Yan Wang, Heyan Huang, and Xian-Ling Mao. Copy is all you need. ArXiv, 0iss/307.00666. 2023. URL httpps//ppa/ semanticscho ar. orgg/orpussr 259298789
Xiang Lisa Li, Ari Holtzman, Daniel Fried, Percy Liang, Jason Eisner, Tatsunori Hashimoto, Luke Zettlemoyer, and Mike Lewis. Contrastive decoding: Open-ended text generation as optimization. In Anna Rogers, Jordan L. Boyd-Graber, and Naoaki Okazaki (eds.). Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023, Toronto, Canada, July 9-1 2023, pp. 12286-12312 Association for Computational Linguistics, 2023. doi: 10. 10,18653/ 18653/v1/203.
URL htttp:://coi.rg////
687
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1990..1169 2019. URL http: /attt://arxiv.... wtw.iiiiiiv./.... org/ abss hatp:///b///1/11 /107/ 11692
Stuart P. Lloyd. Least squares quantization in pem. JEEE Trans. Inf. Theory, 28: 139- 1982. URL https://epp semanticscho ar. www.ooor/oooruuroo ID: 0833328
12

---

Published as a conference paper at ICLR 2025
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Ranjita Naik, Varun Chandrasekaran, Mert Yüksekgönüll, Hamid Palangi, and Besmira Nushi. Di- versity of thought improves reasoning abilities of large language models. CoRR, abs/2310.07088, 2023. doi: 10.48550ARXXV.231100000
URL https://doi org/ /tlgg//1..855///000 48550/arXiv
2310. 07088
Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are NLP models really able solve simple math word problems? In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilck Hakkani- -ürr, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou (eds.). Proceed- ings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL- 2021, Online, June 6-11, 2021, pp. 2080- 2094. Association for Computational Linguisties, 2021. doi: 168. URL httthtts: http:::///ooo.rg/////// wwg.ooggg/ooo.o....o 18 653 /v1 2001 naacl main. in. 168
Letian Peng and Jingbo Shang. Incubating text classifiers following user instruction with nothing but LLM. CoRR, abs/2404.10877, 2024. doi: 10.4850/ARXIV.240..0000
URL https
//doi. org.10.rg1.... 1111111..85550/1 (18855//axiiii 2404. 10877.
Letian Peng, Yuwei Zhang, Zilong Wang. Jayanth Srinivasa, Gaowen Liu, Zihan Wang, and Jingho Shang. Answer is all you need: Instruction-follee following text embedding via answering the question. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Lonng Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, pp. 459-477, Association for Computational Linguisties. 2024. doi: 10.18653/V 10202,,11222AAALLLOLOOOC 2024.ac1-1es 27
Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models as knowledge bases? arXiv preprint arXiv: 1909.990666 2019.
Elvis Saravia, Hsien-Chi Toby Liu, Yan-Hao Huang, Junlin Wu, and Yi-Shin Chen. CARER: Contex- tualized affect representations for emotion recognition. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 3687-3697, Brussels, Belgium, October- November 2018 Association for Computational Linguistics. doi: 10.18653/v /D18-1140 URL 0
https: / www. a web.
Richard Socher. Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Y. Ng. and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, EMNLP 2013. 18-21 October 2013, Grand Hyatt Seattle, Seattle, Washington, USA, A meeting of SIGDAT, a Special Interest Group of the ACL, pp. 1631-1642. ACL, 2013. URL org http////ooooooooo//1111000 D13////111117/
Ilya Sutskever, Oriol Vinyals, and Quoc V Le. Sequence to sequence learning with neural networks. Advances in neural information processing systems, 27, 2014.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Rohert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint wxXiv:2403.08295 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09280 2023.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008
Xuezhi Wang and Denny Zhou. Chain-of tf-thuught reasoning without prompting. arXiv preprint arXiv:2402. 10200, 2024.
13

---

Published as a conference paper at ICLR 2025
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang. Aakanksha Chowdh- ery, and Denny Zhou. consstenny improves chain of thought reasoning in language models. arXiv preprint arXhv:2203.11171, 2022.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 24837,, 2022.
Sean Welleck, Amanda Bertsch, Matthew Finlayson, Hailey Schoelkopf, Alex Xie, Graham Neubig, Ilia Kulikov, and Zaid Harchaoui. From decoding to veta-geneeattin: Inference-time algorithms for large language models. arXiv preprint urXiv:2406. 16838, 2024.
Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Amold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. arXiv preprint arXiv:2007.00800 2020.
Zhilin Yang, Zihang Dai, Ruslan Salakhutdinov, and William W. Cohen. Breaking the softmax bottleneck: A high-rank RNN language model. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL http::::///ppeeeviieeeeee net forum?id-
AkwZsG-CO CZ
Jiacheng Ye, Jiahui Gao, Jiangtao Feng, Zhiyong Wu, Tao Yu, and Lingpeng Kong. Progen: Progressive zero-shot dataset generation via in-context feedback. arXiv preprint arXiv:2210.12322 2022a.
Jiacheng Ye, Jiahui Gao, Qintong Li, Hang Xu, Jiangtao Feng, Zhiyong Wu, Tao Yu, and Ling- peng Kong. Zerogen: Efficient zero-shot learning via dataset generation. arXiv preprint arXiv:2202.07922. 079222 2022b.
Jiacheng Ye, Jiahui Gao, Qintong Li, Hang Xu, Jiangtao Feng, Zhiyong Wu, Tao Yu, and Ling- peng Kong. Zerogen: Efficient zero-shot learning via dataset generation. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.). Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pp. 11653-11660 1669. Association for Computational Linguistics, 2022c doi: 1:0::///1/202..EMMPPMMMMMM
URL https: //ttss://doi.ooo
10
emn
ma in. m0....801 801
Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-1 2015, Montreal, Quebec, Canada, pp. 2015. URL htRp:tttsp///pcccciiigs/ neurips. cc/ paper/ 21///2115hhhhsee Mbs .htll
Abst
Yunyi Zhang, Jiaming Shen, Jingbo Shang, and Jiawei Han. Empower entity set expansion via language model probing. arXiv preprini 2rX4v:2004.1339 2020.
14

---

Published as a conference paper at ICLR 2025
A LIMITATION
The main limitations of our content in this paper are in the in-context navigation (ICN) query searching method, which is established on the ideal performance of large language models. Thus, it might not be applicable to weaker models, especially for language models that are not trained by supervised fine-tuning (e.g. GPT-2). The navigation performance of ICN might also be dependent on the generation abilitty of the language model itself as it might append incorrect results to the context, misleading the space for navigation. Finally, ICN requíres encoding contexts to queries for multiple times, which will reduce the generation efficiency. In summary. ICN should be considered as compromising method in the condition that the language model frozen and thus can be fine-tuned. Some more fundamental ways to address the spurious correlation can be the ones we have discussed in training impact or changing the architecture of the language model to produce multiple queries. Besides further refinement in methods, our experiment scope can also be extended other language modeling scenarios such as masked language modeling, which can be a potential future work.
B LLM DISCRIMINATOR ACCURACY VALIDATION
We validate the discriminative ability of GPT-4o by testing it on the dataset from CG- Expan (Zhang et al. (020) with the prompts in
Category
Positive
Negative
Country
United Kingdom
London
Appendix|K K The dataset include 10 categories, each with positive and negative examples. We showcase some examples in Table 7 GPT-4o achieves 72% accuracy on the test set, which established it as a competent discriminator. As the entities generated from LLaMA might be dif- ferent from the CGExpan test set, we manually check 10% (600 in all) of the discrimination result, which shows 94.67% accuracy. We find the accuracy to be higher because some entities generated from LLaMA makes sense because of bad first-tokens. Thus, we conclude the discrimination of knowledge probing to be a easy task for GPT-4o to make trustful predictions.
Disease
Lymphoma
Chemotherapy
Party
Democrats
Episcopalians
Table 7: Examples of data from CGExpan.
C KNOWLEDGE PROBING CATEGORY LIST
Sports League
Astronomical Objects
Scientist
Landmark
Baskethall Sports League
Planet
Computer Scientist
European Landmark
Baseball Sports League
Nehulae
Social Scientist
Modern Landmark
USA Sports League
Fixed Stars
Female Scientist
Tower Landmark
European Sports League Hot Astronomical Object
German Scientist
Natural Landmark
Country
City
TV Channel
Restaurant
Developing Country
Canadian City
USA TV Channel
Fast Food Restaurant
African Country
Coastal City
Entertainment TV Channel
French Restaurant
Small Country
Capital City
News TV Channel
Pizza Restaurant
Island Country
Asian City
Premium TV Channel
Expensive Restaurant
Company
Creature
Discase
Food
USA Company
Four-leg Creature
Infective Disease
Chinese Food
Japanese Company
Flying Creature
Stomach Disease
Spicy Food
Technology Company
Mammal
Childhood Disease
Veggie Food
Medical Company
Bacteria
Fatal Disease
Dessert
Table 8: The full list of categories and sub-categories.
15

---

Published as a conference paper at ICLR 2025
D
PROBING RESULT ON STARTING CHARACTER
100
80
60
40
20
A
B
C
D
E
F
G
H
1
K
L
M
100
InTop Accuracy
Out Accuracy
80
20
N
O
P
R
S
Y
Figure 8: Probing result on generating words starting with given characters.
In Figure 8 showcase much simpler probing target than the main content generating words starting with given characters, for instance, "P". We can easily evaluate the probing correctness by checking the first character from generation without any LLM. The observed result is consistent with the main content (except for "W" and "X") that tokens with high similarity to top tokens do not necessarily start with the same token, which introduces spurious correlation. This evaluation independent from LLM evaluator further validate the existence of the spurious correlation in NTP by key space similarity.
16

---

Published as a conference paper at ICLR 2025
E RESULT WITH MULTIPLE SAMPLING AS APPROXIMATION
Category
Sports League Out
Astronomical InTop
Scientist
Landmark
Country
City
Group
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OufTop
Accuracy
11.95
54.42
40.75
63.27
38.80
46.91
51.31
68.06
41.07
61.90
55.62
67.45
Rank
54.01
56.97
54.42
52.70
58.51
54.79
56.10
55.71
59.29
Proportion
50.89
49.11
54.00
46.00
53.11
46.89
48.67
51.33
41.78
58.22
51.78
48.22
Cutegory
TV Channel
Restaurant
Company
Creature
Discase
Food
Group
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
Accuracy
50.97
59.07
62.50
74.70
68.75
74.05
58.36
74.48
52.51
68.50
68.96
Rank
54.59
56.96
52.03
59.39
52.04
59.82
50.54
61.69
50.50
59.64
50.07
59.87
Proportion
46.89
53.11
51.11
48.89
54.00
46,00
56.44
43.56
46.22
53.78
46,00
54.00
Table 9: Knowledge probing result with multiple sampling approximation.
Sports Leagues
Astronomical Objects
Scientists
028
58.8
63.
56.4
#Query
+Query
#Query
Figure 9: ICN results with multiple sampling as approximation.
In Table 9 and Figure 9 represent knowledge probing result approximated by Monte-Carlo sampling (5 times each first-token), which shows a consistent result with the main content. For Figure as MAP is inapplicable for multiple answers from first-token, report precision the metric.
F EFFECT OF APPENDED EXAMPLES
Sports Leagues
Astronomical Objects
Scientists
73.22
75.98
69.45
1-00
a0-0000
33.42
7.57
$2.55
20-
3-11
78.88
$3.80
20.56
0.3 3.9
66.11
4.3.5
65.71
23.5
003.1
#Query
#Query
#Quary
Figure 10: ICN results when probing after first-tokens with appended examples.
In the main content, the first-tokens are appended back the initial prompt to avoid the influence of in-context examples on the generation space. In Figure I0] we show what if we append them in the prompt for probing. The result shows that in-context examples overall improve the probing precision. However, it also sometimes results in significant drop when wrong answers appended, which undermines the quality of in-context examples.
17

---

Published as a conference paper at ICLR 2025
G RESULT ON OTHER LANGUAGE MODEL
Category
Sports League Out
Astronomical InTop
Scientist
Landmark
Country
City
Group
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OufTop
Accuracy
39.92
47.72
42.18
56.50
23.21
35.64
47.69
47.99
42.95
43,53
64.12
61.14
Rank
59.67
53.34
53.79
53.25
61.64
54.90
56.74
58.27
60.10
Proportion
42.00
58.00
60.14
39.56
70.22
29.78
63.11
36.89
61.56
38.41
56.22
43.78
Cutegory
TV Channel
Restaurant
Company
Creature
Discase
Food
Group
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
Accuracy
50.28
55.17
63.31
58.15
56.17
$7.25
44.52
62.68
48 04
58.11
66.14
67.97
Rank
54.09
57:21
54.01
57.14
56.16
56.44
51.33
62.38
52.41
59.10
52.18
61.46
Proportion
46.00
54.00
61.33
38.67
47.56
52.44
62.44
37.56
55.11
44.89
70.22
29.78
Table 10: Knowledge probing results olmo--b--nntt
Sports Leagues
Astronomical Objects
Scientists
20.79
15.0
71.65
59.78
37.16
00.02
#Query
#Query
#Query
Figure 11: ICN results on almo 7b- inw-inistuutt-io -hf.
We include experiments on another open-source language model, olmo fommo-7b-insteee aob-instruct.. -hf (Groeneveld et al. 2024), which uses different vocabulary dictionary from LLaMA. Our experiment
shows result consistent to the main content in Table I0 and Figure [1] Thus, our discovery is more widely supported as common issue among different language models.
H RESULT ON LARGER LANGUAGE MODEL
Category
Sports League OutTop
Astronomical InToo
Scientist InTop
Landmark
Country
City
Group
InToop
OutTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
Accuracy
37.18
67.38
36.99
59.70
38.38
59.83
80.78
53.93
46.68
51.44
Rank
51.83
56.62
51.17
59.07
51.54
57.55
51.65
37.31
55.27
56.40
51.69
56.82
Proportion
64:22
35.75
69.11
30.89
71.11
28.89
62.89
37.11
58.22
41.78
58.89
41.11
Category
TV Channel
Restaurant
Company
Creature
Disease
Food
Group
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
Accuracy
39.39
58.02
51.96
69.51
69.73
75.52
56.45
67.55
50.43
76.61
48.33
81.08
Rank
57.17
53.51
53.30
60.35
53.7
59.72
53.94 94
59.03
54.96
56.09
53.91
58. 68.66
Proportion
61.78
38.22
70.14
29.56
66.44
33.56
68.0
32.0
63.11
36.89
67.11
32.89
Table 11: Knowledge probing results on 1lama 3-700-
we illustrate the spurious correlation and ICN performance on the larger 70B LLaMA-3 model. We find consistent conclusion with the discovery the main content. Thus, the conclusion in our work is further scaled-up to even larger LMs.
In Table II and Figure 12
aallmmi001.Mo--7--7241-sss
18

---

Published as a conference paper at ICLR 2025
Sports Leagues
Astronomical Objects
Scientists
RK.ey
60.13
21.31
359
2141
18
#Query
#Query
#Query
Figure 12: ICN results on 1lama 10lma----0000 ---00------------
I RESULT ON DIFFERENT CLUSTERING RESULT
Category
Sports League Out
Astronomical OutTop
Scientist
Landmark
Country
City
Group
InTop
InTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
Accuracy
57.08
55.78
51.95
75.82
53.13
65.52
70.78
76.68
61.11
70.84
61.37
76.86
Rank
52.76
57.61
54.27
58.13
54.57
56.90
55.79
54.46
54.39
55.61
52.77
61.25
Proportion
42.00
58.00
60:11
39.56
70.22
29.78
63.11
36.89
61.56
38.41
56.22
43.78
Category
TV Channel
Restaurant
Company
Creature
Discuse
Food
Group
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
InTop
OutTop
Accuracy
61.03
67.77
68. 53
82.24
74.78
80.72
71.32
78.25
65.84
72.94
72.24
74.91
Rank
54.69
54.02
58.3
58.30
Proportion
56.22
43.78
65.11
34.89
56.89
43.11
70.22
29.78
42.67
57.33
52.67
47.33
Table 12: Knowledge probing results on different clustering result.
As the clustering algorithm might output different clusters by different initialization, we get another version of the clusters and rerun the experiment in Table 3 Table 12 presents our result, which is quite consistent with Table 3. Thus, our conclusion on the existence of spurious correlation is further solidified.
19

---

Published as a conference paper at ICLR 2025
J CASE STUDY
Vasilla (PQuery-1, MKey5
ICN (1Q (Kay=5)
Infections discrses, Malaria, Ebola, HIVAIDS, Tubercalois
Infectious discsses Malaria, Elola, HIWAIDS, Tubercsloii
The plagur Mexdles, Typhoid Fever AIDS, Robies
Pertinsis, Menles, Cholem, SARS-CoV-2, Sypbilis
Penussis, SARS Cov TB, Smallpos, Bubonic Plague Influenza, Plague. [Ounlength), Cood Jwian Influenza
Lassa Fever, Dengue Fever, Typhodd Fever, Rabies, Gonerrhes Herpes simples, Salmonelia, Trichomnnasiss Cordd 19, Stmp Threat MRSA, Whacoping Cough, Ringoora, Shigelkoiss, Croptococooiis Plague Scahios Beruliom
Ringwoem, What is... Choden, Desgue Feues, MRSA
Mad Cow, Chickeapos Tuberculoii, Swite
Threat
Epidenic, Chickeupos, Cuunyllbucter cepospimoiis Rubonic Plagne,
COVID-19, Foidemic, Fla. Fovrr, Whesping Touophammsis Common Cold,
Clestidium diffcclle, Tocaplaempii, Listria moteeeylogeeoo
Legiannires discaso
Ganorbees You're going... Worns, have... you're.. re... [OntLength), (Ouilength) Monkeypos, Spanish Coronavinos
Rocky Moomain, Prinacadis, (Outlengih), Aspengillosis Histrplasmosis, Feverr AIDS, Nipah virus,
Table 13:
Case study of ICN in knowledge probing ("Infective Disease").
In Table 1 we showcase the effect of ICN on knowledge probing, which illustrates ICN to success- fully navigate the query to correct first-tokens away from explored ones. Also, we can observe some nonsense generation introduced by spurious correlation in simply probing the top first-tokens.
Vanilla (#Query=l, #Key=20)
ICN (#Query=4, #Key=5)
The company announced a significant incrense in profits....
The company announced a significant increase in profits,...
Apple's stock price surged 5% after the company...
Apple's stock price surged 5% after the company... ...
Amazon's 's stock price surged 10% after the company...
Amazon's 's stock price surged 10% after the company...
According to the latest report, the company's 's profits...
According to the latest report. the company's 's profits...
Tech giant Apple has announced new line...
Tech giant Apple announced new line...
Microsoft's 's stock price surged 5% after the compuny...
Investors optimistic about the company' 's product...
Coca Cola announced a major restructuring plan to cut... E-commerce giant Amazon announced a 20% increase in... Google's parent company, Alphabet, reported 13% increase... Walmart has unnounced major expansion of its...
In a major deal, tech giant Apple Apple has... A recent report by the International Monetary Fund...
Microsoft's stock price surged 5% after the company...
Google's 's parent company, Alphabet, reported a 13% increase...
Alibubu's 's stock price surged 10% after the company... Stock prices plummeted
Facebook's 's stock price plummeted 10% after the company...
Street yestenday, with... Global economic growth is expected to slow down... Facebook's sss stock price plummeted 10% after the company... Technology giant Apple has announced significant increase...
Tesla's 's stock price surged 10%
company...
Coca-Coe Cola announced major restructuring plan to cut...
China's 's cconomy grew at its slowest pace in...
Ford Moter Company announced a significant increase in...
Shares XYZ plummeted after the...
General Motors announced a major restructuring plan to...
Tesla's stock price surged 10% after the company...
Walmart has announced a major expansion of its...
Wall Street experienced significant decline stocks...
Goldman Sachs reports record profits for the quarter,...
After successful IPO. the tech company's 's stock...
Procter & Gamble's quarterly profits rose 12% duc...
Dow Jones Industrial Average (DJIA) closed at a...
Netffix's stock price surged 10% after the company...
Table 14:
Case study of ICN in open-ended generation ("Business News").
In Table 14 we showcase the effect of ICN on open-ended generation, which shows the first-token
explored by ICN has better diversity, which also leads more diverse sentence structure in generation.
Question
Rita hand-picks Junebugs off of her plants every summer. On Monday, she removed 39 Junebugs. On both Tuesday and Wednesday, she removed twice as many Junebugs as she did on Monday.
Thursday she removed 48 and on Friday she removed 57.
What is the average number of Junebugs that she removes per day?
Answer
On both Tuesday and Wednesday, she removed twice as many as the 39 on Monday so she removed 2*39 78 Junebugs. All total she has removed 39+78+78+48+57 300 Junebugs She has removed 300 Junebugs over 5 days so she has removed, on average 300/5 = 60 Junebugs per day
Table 15: The case used for chain-of-thouge generation in Table 16
20

---

Published as a conference paper at ICLR 2025
Vanilla (#Query-1l 1, #Key-4)
ICN (#Quury tKey-1)
* SSepe 1:
e Step 1:49
Find the total number of Junebugs removed.
Find the total number of Junebugs removed.
39 x 89 48 + 57 39 78 + 48 + 57 222
39 2 39 48 + 57 - 39 78 48 + 57 222
enStep 2:00 Find the average number of Junchugs remaved per day. To find the average. we divide the total numbbr of Junebags removed number days. 222 5 4444
etSSep 2:44
Find the avenge number of Junchugs removed per day. To fiod the average, we divide the total number of Junchugs removed by the number of days. 222 5 44.4
Step
Step
Calculate the total number Junchugs removed:
List the number of Junchugs removed each day:
2 39 +57 39 78 48 57 222
Monday; 39
seStep
Tuesday: 2.x. 39
Calculate the overage number of Junchugs removed per day: 222
Wednesday: 2 39
78
Thursday:
Friday: 57
stSpp 2:
Add number Junebugs removed each day:
39 + 78 78 48 + 57 300
e Step 3:40
Find the average number of Junebugs removed per day: 300 5 60
Step 1,50
**SS 1:
We can add up all the Junebugs that Rita removed:
Use the information that Ria removed twice as many Junebuga Tuesday and Wednesday
S"Step 339 +-48 57 39 78 48 57
did Monday.
This means that she removed x 39 = Junebugs on Tuesday
We can find the average dividing the total number of Junchags number days:
and 78 Junchugs on Wednesday.
*SStop 2:44
222/ 5 44.4.
Add up the number of Junehugs she removed each day:
39 78 78 + 57 300.
*Step
To find the average, divide the total number
of Junchups by the number of 300 60.
SSupp 1:
SSeop 1:*4
Calculate the total number Junehugs removed:
Rit removed Junebugs Monday.
39 + x 39 48 57 78 + 48 57 222
*-Step 2:40
**Step 2:40
On Tuesday and Wednesday, she removed twice as many Junehugs she did Moncayy So removed 78 Junebogs on Tuesday and 78 Junehugs on Wednesday.
Calculate average number Junchugs removed day: 222 44.4
a SSepp 3:*a
Thursday, she sSSScp 4:4x
48 Junehugs.
Friday, she removed Junebugs.
s*SStop 5:
To find the average number of Junehugs she removed per day, Juncbugs removed divide by
add up
days.
sSepp 6:*9
The total number Junebugs she removed 39 + 78 48 57 300. *Sep 7:
The average number of Junehugs she removed per day is 300 = 60.
Table 16: Case study of ICN in chain-of-thought generation.
In T 6 we showcase how ICN helps the generation of diverse Chain-of-thoughee which conse- quently improves the performance in reasoning.
21

---

Published as a conference paper at ICLR 2025
K
PROMPTS, CASES AND HYPERPARAMETERS
In Table 17
we present the prompts used in our experiment for reproduction.
Function
Prompt
"Please show me some [CATEGORY]"
Knowledge Probing
1. [ANSWER
"Please show [CLASS] [DOMAIN]."
Open-ended Generation
[SENTENCE IJ
"[QUESTION] Please show me some different ways to solve problem." " #*#Way I: [CoT I]
Chain-of-Thought
#*# 2:
Rephrasing
Reprahse this request: [REQUEST]
Probing Evaluation
"Please show me some [CATEGORY]. Is '[ANSWER]' considered as among correct answers? Answer only "Yes" or 'No'."
Get Answer (Numeric)
[CoT] F*Final Answer (Only Number):
Get Answer (Choice)
[CoT] **Final Answer (A, C. D, E):**
Table 17: The prompts used in experiments.
In Table 18, we showcase the rephrased knowledge probing requests generated by GPT-4o in the "Rephrasing" method.
Original
Rephrased
"Please show me some sports leagues."
"Could you please provide me with list of sports leagues?"
"Plcase show me some USA companies."
"Could you provide me with
companies based the USA?"
"Please me some pizza restaurants."
"Could you provide information on some pizza restaurants" illnesses that affect children?"
"Please show
childhood discases."
"Could you provide information
"Pleuse show
veggic foods."
"Could you please offer selection vegetarian dishes?"
Table 18: Cases in the rephrasing of knowledge probing requests.
For ZeroGen implementation, we fine-tune RoBERTa Liu et (2019) PoBBRRaa-aagge the classifier, which is optimized by AdamW W Loshchilov Hutter (2019). The learning rate is initialized 10-4 The classifier is fine-tuned by 10 epochs with batch size 16. For the result, we report the averaged performance over 5 different runs.
22


Reviews:
[{"review_id": "cOVaYov4Pi", "reviewer": "Reviewer_ciXY", "summary": "This paper investigates the impact of word-embeddings parameters in the last layer (also referred to as un-embedding parameters, referred to as \"keys\" in this work) on spurious correlations during decoding via language models. Specifically, this paper finds that top-k items in response to a prompt tend to be close together (low cosine distance) and hence a lot of plausible and correct responses which might have a higher cosine distance from the top-k items are ranked lower in LM predictions. This impact is shown on knowledge probing tasks which shows that a significant number of correct answers lie outside of the cluster formed by top few items. To ameliorate this issue, the paper proposes a simple prompt-based fix called ICN, which appends the top-few items to the query in a listed form to encourage the language model to suggest \"under-explored\" items. This approach is shown to improve accuracy on knowledge probing. This approach is also used for open-ended generation to stimulate diversity among the responses. When this approach is used for CoT consistency decoding, it tends to improve upon standard CoT sampling baseline slightly. Finally, this paper also provides evidence that fine-tuning doesn't change the unembedding layer much and is still susceptible to the same spurious correlations as pre-fintuned models.", "questions": "Please address the concerns above.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "-- This paper is well-motivated and identifies a focused problem in natural language generation. However, as I note in the weaknesses section, it misses out on discussion about other work investigating the inadequacy of softmax for similar reasons this paper explores.\n\n-- The effects shown are significant. The knowledge probing experiments show concerning effects of the spurious correlations introduced due to the unembedding parameters.\n\n-- The proposed fix, while simple seems to be effective at ameliorating the issue.\n\n-- The proposed approach also seems to improve diversity of generated text.", "weaknesses": "--  The paper should improve the presentation. I am unsure about the exact prompt behind ICN while I understand it at a high level. Moreover, I it is also difficult to understand how this prompting scheme was explicitly adapted for open-ended generation and CoT. Clear examples would help. Also, the baselines should be clearly explained in greater detail.\n\n-- The proposed ICN scheme also seems expensive to implement as a general method. The results indicate that one needs a large number of queries and a large number of keys to work.\n\n-- My biggest concern is that this work ignores a lot of prior work that precisely characterizes the problem of spurious correlations due to the unembedding parameters. [1] introduces the popular concept of softmax bottleneck which shows rank-deficient behavior of the softmax function during decoding which is directly linked to the unembedding parameters. [2] investigates exactly the problem this paper studies. [3] relatedly shows that some words might never be chosen as argmax even if correct because of the manifold of the unembedding parameters.\n\n[1] Breaking the softmax bottleneck: A high-rank RNN language model. Yang et al. 2017\n\n[2] Stolen probability: A structural weakness of neural language models. Demeter et al. 2020\n\n[3] Low-Rank Softmax Can Have Unargmaxable Classes in Theory but Rarely in Practice. Grivas et al. 2022\n\nAll of these works and more are closely related to the proposed work and should be discussed in detail.\n\n-- It is not discussed/analyzed in detail what makes ICN work. The prompt as designed could have worked against the desires as well in which it explores the nearest items more eagerly. I don't come away with a good understanding of why this method of prompting is effective over other potential prompts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L0upnsDU5H", "reviewer": "Reviewer_2odz", "summary": "This paper present an interesting fact that the vocabulary representations (keys) has spurious correlations with the generation process. More specifically, when some token gets a higher generation probability, the tokens that are closer to this token will also have higher generation probability, even though they may not lead to correct generation results.\n\nThe paper also present a simple solution, which includes the explored decoding results in the context and prompt the LLM to generate differently. Experiments on knowledge-probing tasks show that the method leads to exploration away from the explored keys.", "questions": "Since the key representations are learned during large scale pretraining, it is also possible that the generation probability behind these keys still matches the knowledge in the world. As shown in Figure 2, there are quite a lot of correct key around the query. It is not possible, in this case, to have all the correct answer appears around the query anyway.\n\nI am wondering what prompt is used for rephrasing and reranking, respectively. I am wondering why rephrasing does not bring to much difference.\n\nI am also curious about how the method related to the contrastive decoding approach.\n\n[Contrastive Decoding: Open-ended Text Generation as Optimization. ACL 2023]", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper use clustering and visualization to show that the similarities in keys did resulting in generation bias.\n\nThe proposed method, ICN, could increase the generation probability for tokens that are different from previous explored keys.", "weaknesses": "I dont think it is reasonable to ask the LLM to be correct for the top 100 candidate. At least it has already shown strong performance on the top 10 or maybe to the top 20 candidate.\n\nI think ask the LLMs to generate different tokens is an practical approach, but it will also decreases the generation probability for correct answers that shows the same tokens with those explored ones.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rRsvlDx3le", "reviewer": "Reviewer_7cDM", "summary": "The paper discover a fundamental issue that the next-token prediction is affected by correlations within the vocabulary key space. By analyzing key representations in transformer LMs, the authors observe that context-agnostic similarities between certain keys can lead to spurious correlations, such as the capital alphabets are predicted with high probabilities together without regarding their real relevance. These correlations cause certain “middle-ranked” tokens, often unrelated semantically to the top-ranked token, to be underestimated in ranking. This issue introduces bias and reduces sampling diversity in open-ended or exploratory text generation. The author further proposed In-Context Navigation (ICN) to incorporate the previous generated outputs into the input context to iteratively guide the LM away from explored key regions.", "questions": "- I wonder why a regular decoding process cannot achieve the same effects of ICN did? In regular decoding, the previous generated outputs are natively placed in the input context. If we put always put the generated outputs in the context, we should be able to achieve the same thing as ICN did, right?\n- Is there any analysis of how ICN can increase the context length and inference cost?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- Novel discovery on the the issue of very fundamental design of transformer LMs. It's valuable to facilitate the understanding of LM decoding biases\n- Novelty in proposed approach ICN to mitigate such bias.\n- The paper shows strong empirical evidence through diverse tasks, including knowledge probing and chain-of-thought generation, to validate ICN’s effectiveness.\n- The author proposed some potential direction to solve this issue in the LM pretraining stage, which can be positive contributing to the community.", "weaknesses": "- Lack of exploration on larger models: the paper only discuss LLaMA-3-8B in the experiments. However, I expect that a larger model with higher-dimensional embedding space can more or less mitigate such issue. When the dimension is large enough, the query embedding can be similar to different key embeddings in many different sub-spaces, so as to reduce the problem of similar key vectors always get activated together. It will be good to show more results in the larger models like LLaMA-3-70B.\n- ICN increase the context length. When making the model generating longer responses (e.g. CoT) while considering previous generated outputs, ICN can largely increase the context length and inference cost.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RBCesI3E2L", "reviewer": "Reviewer_Re3A", "summary": "Following the terminology in the paper, let us denote the encoded input hidden sate representation at some position (the one that is multiplied with the vocabulary matrix to get the probability distribution at that position) as \"query\", and let us denote the vectors in the vocabulary matrix (each correspondent to some token in the vocabulary) as \"keys\". \n\nThe paper explores the confounding factor that is brought about by the input-context-insensitive similarity of keys in their vector space in next token predictions (NTPs). Particularly, in the context where there are multiple good answers for an input prompt, many of the good answers are not as highly ranked, because input-irrelevant vocabulary keys can get highly ranked because they are close (in the non-contextual vocabulary embedding space) to some of the top ranked keys that are actually relevant to the input. Essentially, even if the top ranked keys are appropriate, they being scored high leads to a side effect of inappropriate (unrelated to input) keys to be also ranked moderately high only because they are close to the top keys in the vocabulary embedding space (which is insensitive to input context by design).\n\n1. The paper provides empirical analysis of this phenomenon by probing LLMs. \n2. The paper explores several ways to mitigate this issue - example by reranking/rephrasing and the main proposed method ICN (in-context navigation). ICN - employees iterative prompting to sequentially generate multiple good answer candidates (instead of just multi-sampling), where the model is prompted explicitly to generate something other than prior generated answers -- thus allowing to focus on different vocabulary key clusters unrelated to the prior answer-clusters. \n3. The paper presents the utility of ICN as a strategy in different contexts - like incorporating it with self-consistency for mathematical reasoning or investigating open-ended generation.", "questions": "1. Line 256-257 is phrased in a confusing manner (and probably wrongly). Instead of \"specifically, from\ntop-N (N = 100) to top-(K + 1) (K = 10)\", you should probably say something like \"specifically, from\nall tokens in top-N (N = 100) but not in top-K (K = 10)\".\n\n1. I think the equation in 261 defining InTop is wrong. It should have $j \\leq K$ not $K \\leq j$. j is supposed to be one of the top-k tokens (thus $\\leq K$) based on the text description.  \n\n1. Related to the last weakness: I think it would be good to more explicitly show what prompt template is used for ICN in different contexts and maybe some concrete examples about how iteratively the input prompt is being changed for different cases (mathematical reasoning, open ended generation, knowledge probing). I saw the case studies in the appendix but they seem to focus only on the output, or just the original input and output (rather than the processed prompt input with all the templates and such). And Table 16 doesn't seem to show anything related to ICN.\n\n1. Still somewhat related to the last weakness: I am a bit confused by your definition of COT prompt in Table 16. Normally, COT prompt is supposed to be add something like  \"Let's think step by step\" to the beginning of the assistant response. Your COT prompt seems to be just about generating multiple answers. Moreover, self-consistency normally used distinct samples, whereas this COT prompt seems to be designed to generate multiple answer procedures in a single sample.\n\n1. What exactly is Top in Table 6 or 5? Do you mean greedy decoding? If so I am not sure how that can be used to generate different sequences for diversity measures or self-consistency voting.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Provides a reasonable analysis of an interesting phenomenon in the context of LLMs. Devises reasonable ways to quantify it. \n1. Proposes a simple method (ICN) to mitigate the identified issue, and shows its utility empirically in a wide set of experimental contexts. \n1. In the end, suggestions some interesting ideas to potentially solve this issue in a more principled manner (contextual layer for vocabulary or re-ranking). \n1. The paper is well written (lower presentation score because of weakness in literature review)", "weaknesses": "1. Similar analyses have been already done. If I understand correctly, this issue is related to the well-known softmax-bottleneck phenomenon. And there are a bunch of related works on that - especially see [1]. While I think the current paper has some virtues over some of those prior works (like investigation in terms of LLMs, exploration of a tuning-free method ICN to mitigate it, and potentially better empirical measure by cluster divisions and such of this issue) - the past related works should be still acknowledged and used for contextualization.  \n\n1. If I am not missing anything, most of the analyses and investigation is based on a single LLM (Llama 3 8B instruct). While olmo results seem to be presented in the appendix, but ICN-related experiments seem to be missing for olmo. So the universality of the utility of ICN is still a question.\n\n1. There could be more discussion about how much does this issue have anything to do with NTP in particular. Seems like a general issue -- that can apply to other language modeling settings (like Masked language modeling).\n\n1. The technical details hyperparameters and the prompt settings are not crystal clear to me. \n\n \n\n\n\n\n[1] Softmax Bottleneck Makes Language Models Unable to Represent Multi-mode Word Distributions - Chang et al. ACL 2022\n https://aclanthology.org/2022.acl-long.554/", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)