Do LLMs Build World Representations? Probing Through the Lens of State Abstraction

Authors: zichao  li
yanshuai  cao
jackie c k cheung
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:05:17

Abstract: How do large language models (LLMs) encode the state of the world, including the status of entities and their relations, as described by a text? While existing work directly probes for a complete state of the world, our research explores whether and how LLMs abstract this world state in their internal representations. We propose a new framework for probing for world representations through the lens of state abstraction theory from reinforcement learning, which emphasizes different levels of abstraction, distinguishing between general abstractions that facilitate predicting future states and goal-oriented abstractions that guide the subsequent actions to accomplish tasks. To instantiate this framework, we design a text-based planning task, where an LLM acts as an agent in an environment and interacts with objects in containers to achieve a specified goal state. Our experiments reveal that fine-tuning as well as advanced pre-training strengthens LLM-built representations' tendency of maintaining goal-oriented abstractions during decoding, prioritizing task completion over recovery of the world's state and dynamics.

Main paper: 
 introduction: 
 Drawing inspiration from human mental models [9, 12, 22] , AI researchers have introduced the concept of world models for sample-efficient and robust machine learning systems [13, 25] . Specifically, a world model is defined to fulfill a dual role: (1) it estimates information about the world state that may not be directly observable from the input signals, and (2) it distills essential information to predict future states, thereby informing subsequent actions based on these predictions.Recently, there has been growing interest in investigating whether pre-trained Transformer models [36] , especially large language models (LLMs) [32] , construct implicit world models. These investigations aim to determine whether the state of the world as described in the text, either implicitly or explicitly, can be recovered from the internal representations built by Transformer models. However, the field has produced studies with conflicting conclusions when examining different tasks. For example, [27] shows that it is possible to accurately probe the status of entities and their semantic relations encoded in the representations of LLMs within a discourse. Conversely, [24] report negative results when attempting a similar yet more challenging setting. In addition, [28] successfully extracts the board state of a partially played Othello game from the internal representation of a small-scale GPT model [7] trained to complete game scripts. A general abstraction (termed as world-irrelevant abstraction later) enables one to answer a wide range of questions. On the other hand, a count-oriented abstraction is only applicable to the counting task.We argue that these contradictions stem from the fact that recovering the complete world state is not always necessary for solving every tasks, which previous studies have not systematically controlled for. Rather, an abstract representation of the world state may sometimes be sufficient, and the necessary level of abstraction can also vary. Some tasks may require the complete recovery of world dynamics in order to predict future states, while others may get by on one that omits this information as it is unnecessary for task completion. For instance, consider Figure 1 . When parsing a discourse into a representation of the world state, one can either record the location of each object or simply note the number of objects in each box. The former offers the possibility to answer diverse questions, including future state predictions after operations, e.g., Move the key to Box C. The latter, however, is restricted to a counting task, unable to foresee future states. Neglecting this nuance could lead to a mismatch between evidence and conclusion, causing undue pessimism about LLMs lacking awareness of the world or excessive optimism regarding their ability to develop general world representations. For example, fine-tuning on all types of questions in Figure 1 might push LLMs to capture world dynamics, this does not guarantee the same outcome when fine-tuning solely on counting tasks.We formalize this intuition through the application of state abstraction theory, originally proposed in reinforcement learning (RL) [29, 3] to simplify the state space by aggregating similar states to abstract ones without modifying the core aspects of the task or underlying world. The level of abstraction is a spectrum, from more general world-irrelevant abstraction, that allows recovery of world dynamics to more goal-oriented abstractions, guiding task completion while giving up on predicting future states. These include Q * -irrelevant and π * -irrelevant abstraction; the former preserves the long-term impact of actions, while the latter preserves the optimal policy. Through this lens of state abstraction, we propose a new framework for examining the world representations constructed by LLMs. This framework investigates the various types of abstractions that may be encoded by LLM-built representations. To demonstrate its utility, we present a concrete application. First, we design a text-based planning task, REPLACE, which requires altering the state of a simplistic world with a collection of containers and objects. We intentionally craft the task's state space to be highly structured and modular, enabling a precise yet simple derivation of world state abstractions at different levels. Despite its simplicity, the abstract states at different levels are distinct and identifiable. Subsequently, we prompt pre-trained and fine-tuned LLMs to complete this planning task and extract their representation during decoding. Finally, we probe different abstract states within the representations. We conduct experiments on a wide range of Transformer models and LLMs, namely Pythia [6] , Llama2 [35] , Llama3 [1] , Mistral [20] and Phi3 [2] . Our experiments show that LLMs achieving reasonable performance on REPLACE, whether through fine-tuning or advanced pre-training, tend to maintain goal-oriented abstractions rather than more general world representations during decoding. Additionally, pre-trained models with near-random performance fail to efficiently preserve any type of abstractions.Our contributions are as follows: 1) We propose a new framework to probe world abstraction from LLM-built representations that can be adapted for other NLP tasks. 2) We release a new synthetic task, REPLACE, and accompanying datasets that are modular and extendable. 3) Experiments using our framework and task yield novel findings: LLM representations prioritize goal-oriented abstractions that preserve the effect of actions in terms of task completion while abstracting out the world state and dynamics during decoding. 4) Our findings also reconcile conflicting conclusions in prior work. For instance, [21, 28] successfully probe state variables like disc color and obstacle position, which pertains to goal-oriented abstraction for the task that the models are optimized to solve, while [24] struggles to recover entity status, which lies outside these abstractions. 
world models: 
 The concept of world models in machine learning has deep connections with the concept of human internal models [9] in cognitive science, most often referred to as mental models [12, 22] . As originally defined, mental models can build abstract and symbolic representations of entities and their relations in the real world or environment around them [12] . One of the most important characteristics of such mental models is their optimal balance between representation complexity and utility for accomplishing a specific task [18] . In a similar vein, AI researchers have developed symbolic [33] or neural world models [13, 14] to compress vast amounts of input information and extract a simplified and essential representation to predict future states. Our work is related to both domains and applies state abstraction theory [29, 3] from reinforcement learning to assess the representations of the underlying world and the task, if any, encoded by LLMs.There is a growing interest in probing for interpretable features in LLMs, most of which are linguistic in nature, such as morphology [30] , syntax [17] and word-level semantics [5] . Recently, researchers have moved beyond the exploration of shallow linguistic features, investigating whether and how LLMs' hidden representations on the fly encode entities and their relations, as described in the text. More specifically, [27] trains a shallow neural classifier on top of LLMs' representations of discourse to recover the ground-truth situations as depicted in the text. This approach is further modified and formulated as a next-sentence prediction task by [24] . Similarly, [28, 21, 37] probes the internal representation of world state in Transformer models reading (semi-)structured input, such as game scripts and embodied sequences.However, existing work primarily focuses on probing for a comprehensive description of the underlying world, defined as a set of state variables. This approach has two main limitations. First, it fails to distinguish the function of a specific state variable: Is it intended to maintain a general representation of the world, or is it crucial only for specific tasks, or both? Without this distinction, we cannot accurately interpret the positive outcomes of probing. Second, existing studies do not account for the potential abstraction of the world that LLM might build in its representation. Our work overcomes these limitations by introducing a general framework that examines the world abstraction encoded within LLM representations.In contrast to [28, 21, 37] , the input in our task is text data instead of a game script or environment layout. Therefore, the conclusions drawn from our experiments are more applicable to LLMs and Transformer models for NLP tasks. 
framework: 
 In this section, we introduce a new probing framework to investigate the world state abstraction encoded by large language models (LLMs) when prompted to perform a decision-making task. We start with a reinforcement learning (RL) formulation, which is general enough to encompass a wide range of NLP tasks and other tasks adopted by previous work, e.g. completing an Othello game script [28] . An RL problem is characterized by a 4-tuple: (S, A, T, R), where s ∈ S represents the world state, a ∈ A denotes possible actions the model can take, transition function T (s ′ |s, a) measures the probability of transitioning to state s ′ induced by a from s, and R(s, a) is the reward.Our framework builds upon a rigorous definition of world abstraction. In particular, we follow RL literature [29] , deriving state abstraction function ϕ : S → S ϕ that maps each state s into an abstract state ϕ(s) (|S ϕ | ≤ |S|). In addition to the raw state S, we consider three types of abstraction:World-irrelevant abstraction ϕ w2 ensures that ∀s 1 , s 2 where ϕ w (s 1 ) = ϕ w (s 2 ) implies that ∀a, x ′ ∈ S ϕw , R(s 1 , a) = R(s 2 , a), ands ′ ∈ϕ -1 w (x ′ ) T (s ′ |s 1 , a) = s ′ ∈ϕ -1 w (x ′ ) T (s ′ |s 2 , a). Intuitively speaking, it preserves the transition dynamic of the underlying world and the reward function, thereby enabling the prediction of future states induced by subsequent actions. Recalling the definition of world models in Section 1, this type of abstraction enables precise recovery of the world model.Q * -irrelevant abstraction ϕ Q ensures that ∀s 1 , s 2 , if ϕ Q (s 1 ) = ϕ Q (s 2 ), then ∀a, Q * (s 1 , a) = Q * (s 2 , a), where Q * (s, a) = max π Q π (s, a). This form of abstraction preserves the effect of all actions a ∈ A in terms of their optimal Q-value, which is the maximal expected future rewards. However, it discards the world dynamics, making future state predictions infeasible. LLMs can leverage this abstraction to be cost-sensitive, minimizing action counts and avoiding penalties for violating world constraints.π * -irrelevant abstraction ϕ π guarantees that the optimal action, and hence the optimal policy, can be recovered. Formally, ∀s 1 , s 2 ∈ S , ϕ π (s 1 ) = ϕ π (s 2 ) implies that π * (s 1 ) = π * (s 2 ). Therefore, this is the coarsest abstraction that one can recover the optimal policy. Therefore, we expect that at least this type of abstraction can be accurately probed from an LLM that excels in the given task.In practice, one may use learning [23, 4] or heuristics [10, 11] to derive the abstractions at each level.We aim to assess which types of state abstraction are encoded in LLM representations. To do so, we first prompt a pre-trained or fine-tuned LLM with τ (s 0 , A 1:t-1 , z), which is a textual description of the initial state s 0 , previously executed actions A 1:t-1 , and optionally, feedback z received from the world. The feedback could be the opponent's moves in an Othello game or users' utterances in a dialogue system. Next, we extract the hidden states H (m) from the m-th layer of the LLM, considering m as a hyperparameter. Following [28] , we select the last hidden states h (m) t3 , which is used by the LLM to predict subsequent actions. After collecting (s t , h t ) pairs, where s t is the world state induced by s 0 , A 1:t-1 and z, we detect the existence of abstraction of s t in LLM representations by assessing if ϕ(s t ) can be probed from h t with accuracy surpassing a random baseline, potentially achieving near-perfect performance. Previous work [16, 17] on probing presents a challenge in differentiating whether a representation encodes a linguistic property or if the probe itself learns the task. While these works focus on low-level syntax features [16, 30] , our focus is on the recoverability of world states and their abstractions from LLM representations, questioning whether these are maintained or discarded during decoding. Therefore, we train probes to classify raw and abstract states, following previous work [27, 28] , which essentially estimates the mutual information between abstractions and representations [34] . As we will show in Section 6, probing either raw or abstract states from LLM representations is not always successful.Remark: To draw faithful conclusions from probing experiments, it is crucial to carefully design or select the tasks in a way that ensures the spaces of abstract states at different levels differ (S ϕw ̸ = S ϕ Q ̸ = S ϕπ ). Otherwise, one cannot determine whether the success of probing stems from the LLM's preference for learning a general world model or from the necessity to recover the world state while learning the optimal policy. For instance, in the game of Othello, it is feasible to recover most of the current board state from all possible legal moves. As such, the raw state space is almost identical to the coarsest π * -irrelevant abstraction for predicting legal moves. Therefore, a plausible interpretation for the success of probing in [28] could be that the Transformer model learns the π * -irrelevant abstraction rather than deliberately learning a general world model.In the next section, we design a planning task within an RL framework and synthesize a dataset accordingly. In Section 5, we prompt LLMs to perform this task and probe different types of abstractions from their internal representations. 
replace: a text-based planning task: 
 To instantiate our framework, we draw inspiration from recent work on probing discourse representation in LLMs [27, 24] and design a text-based planning task named REPLACE. In REPLACE, an LLM acts as an agent within a simple world consisting of a set of containers and objects. The objective of the LLM is to predict actions that alter the situation of the objects to match a described target situation. We design this task for three reasons. First, in its symbolic form, the task possesses a simple and modular state structure, making it easier to analyze and derive abstract state spaces. Second, it is closely related to the gripper problem [31] and the entity tracking task [24] , which are widely adopted in planning and NLP interpretability research. Third, as we will demonstrate, each type of abstraction has a unique state space, enabling the assessment of whether LLM-built representations encode specific abstractions.In this section, we start with formulating the planning task within the RL framework, then introduce the textual realization of the RL elements, and finally the method of curating prompts.Formally, the specific design of the elements within the RL framework is described as follows: • put(o): Put an object from the agent to the nearby container. 3. Transition function T : Following previous work [27] , we consider a static environment. That means T (s ′ |s, a) ∈ {0, 1}. Moreover, the transition impacts only the current situation u t in s t by taking the effect of a. For example, if agentLoc(2) is True in u t-1 , executing the action move(left) results in agentLoc(1) becoming True in u t .1. State s = [u,Reward function R(s, a) checks whether the new situation of o, induced by executing a, matches the target situation: R(s, a) = r if P o (u ′ ) = P o (g); or -(r + r ′ ) if a / ∈ C(s); or -r otherwise. Here u ′ denotes the new situation, P o (u ′ ) the set of assignments on predicates involving objects o (store and held) in the new situation u ′ , and C : S → A, the constraint function mapping states to permissible actions. We incorporate three spatial constraints: (1) the agent cannot move left and right at the leftmost and rightmost locations, respectively; (2) grab is only possible if the target object o is in a nearby container; (3) put is allowed only for objects currently held by the agent. Given the LLM's insensitivity to exact numerical values in prompts, we do not set concrete values for r and r ′ but simply assume r > 0 and r ′ > 0. Notably, both Q * and π * , as well as the abstract states, are invariant to the actual magnitudes of r and r ′ .We map the situation (u or g), previous actions (A 1:t = {a 1 , .., a t }) to their respective natural language descriptions τ s (u), τ s (g), τ a (A 1:t ), and using a set of textual templates: τ s (•) and τ a (•). As such, the input text to the model comprises two main components: (1) a general task instruction (see Appendix C) specifying the actions the agent can take, {τ a (a)|a ∈ A}, and the constraints it must adhere to, τ (C), and (2) textual descriptions of the world's initial and target situations, as well as previous actions [τ s (u); τ s (g); τ a (A 1:t )]. An example is presented in Figure 2 . Conditioned on this, the LLM generates textual descriptions Ŷ of subsequent actions, which are parsed back into their symbolic form Ât+1:T by a rule-based parser τ -1 a ( Ŷ ).We synthesize two English datasets for REPLACE: GRIPPER and COOK. For each instance in GRIPPER, we first sample a set of o from a list of container names and b from a list of frequent nouns in British National Corpus (BNC) [26] . Subsequently, we generate pairs of initial and target situations, [u 0 , g], by randomly assigning values to the predicates in u 0 and g, such that it requires the execution of 2 to 6 actions to transition from u 0 to g. Towards a more realistic setting, we introduce the following dataset variants: (1) Lexical variants, which includes rare nouns collected from BNC for b, adding color and size modifiers to describe the objects, and use diverse textual templates for translating predicates; (2) Partial τ (g), which omits explicit information about held in the target situation, requiring logical inference by LLMs; (3) Partial [τ (u t ), τ (g)], an extension of (2) that includes previous operations performed by the agent, thus requiring inferring u t . Appendix C provides concrete examples for each variant. The final dataset is a uniform mixture of these four variants and segmented into training, validation, and test sets, with [u 0 , g] splits, allocating 50k, 1k, and 1k instances to each set, respectively. COOK generalizes the setting of GRIPPER to a grid world of containers and is generated with TextWorld [8] . Similarly, we sample and transform the initial and goal situations into textual descriptions. The main differences between COOK and GRIPPER are: (1) the constants of direction d are different: the agent in COOK can move north, south, east or west; (2) COOK uses a different set of textual templates adopted from the cook domain in TextWorld. See Appendix D for more details about COOK. Notably, both datasets share the same world abstractions at all levels. 
probing world abstractions in llm representations using replace: 
 For each LLM prompted to solve REPLACE, we assess which types of abstractions of state s t are preserved in its representation h t during decoding after the collection of {(s t , h t )} pairs as described in Section 3.The modular state structure in REPLACE allows us to analytically derive a set of predicates for each type of world state abstraction. To begin with, we introduce a set of abstract predicates crucial for devising the optimal plan for task completion:• nearby(o): the object o is located within the container nearby to the agent. Similarly, we have distant(l, o), where l is the relative distance to the agent's current location, e.g., distant(-1, o) being True indicates that the container one step to the left contains o.• subgoal(l, j): the agent needs to visit a series of containers to manipulate objects, and the j-th container to be visited is at the relative distance l to the agent's current location. • nextObjDirect(v): the relative direction of the next object to be manipulated, where v could be left or right to the agent, in a nearby container, or with the agent. • nextObj(o): the object o to be manipulated at the next step, if applicable.Appendix E provides a detailed explanation of the definition and derivation of these predicates. Finally, we can infer the optimal action for the next step by incorporating nextObj and nextObjDirect. Figure 3 demonstrates how to derive abstract predicates and actions based on the raw predicates. Figure 4 provides concrete examples of predicates. Now, we derive each type of world abstraction, using both raw and abstract predicates:Raw state The complete world state S includes store, held, agentLoc in both the current and target situations, u and g, and boxName.World-irrelevant abstraction To preserve the world dynamics, we have to track all raw predicates of the current situation, excluding boxName, as actions do not affect it. In addition, we include all predicates in target situation g to preserve the reward function. Therefore, S ϕw includes the predicates store, held and agentLoc in both u and g. Q * -irrelevant abstraction To construct S ϕ Q , we include agentLoc, held u and nearby u to distinguish between the effects of legal and illegal execution of actions. Also, we include subgoal, where subgoal(l, j = 1) differentiates the effect between optimal and non-optimal actions, and subgoal(l, j > 1) quantifies the effect of executing each action. Further, we include nextObj to distinguish the effect between optimal and non-optimal put or grab actions. π * -irrelevant abstraction S ϕπ is factorized by predicates nextObjDirect and nextObj. nextObjDirect specifies the action type (move or manipulate objects) and moving direction (left or right) at the next step. nextObj identifies the particular o for grab or put operations.Appendix G provides a more formal proof.We probe the predicates in LLM representations by training two-layer neural models, following [28] . For single-variable predicates (e.g., nearby(o)), we train an individual probe for each. For two-variable predicates where one variable denotes a positional index (e.g., boxName(i, b)), we train a separate probe for each i-th position. For some predicates, e.g., store(i, o), involving variables o with domains that differ across the dataset, the input to probes is a concatenation of [h t , φ(o)], where φ(o) is the embedding of o. Following the conditional probing principle [15] , we treat the embedding method for each predicate as a hyper-parameter, selecting the method with the largest performance margin between using [h t , φ(o)] and using only φ(o) as the input. Specifically, we explore two embedding methods: (1) averaging the word embedding for all the tokens in o and (2) selecting and averaging the hidden states H across all mentions of o. For other predicates, such as agentLoc and subgoal, probes take as input h t . Appendix H provides more details about the formulation of probing tasks as well as the design (e.g. layer number) and optimization of probes.Evaluation metric We evaluate the performance of probing models for each predicate using the F1-score. To account for variations in candidate numbers across different predicates, we normalize the scores x p for predicate p by max(xp-βp,0) 1-βp × 100%, where β p is the score of a RANDOM baseline. It predicts outcomes by proportional sampling from label candidates of p based on their frequencies. We refer to this normalized score as the recovery rate. We experiment with two groups of LLMs to investigate how fine-tuning and pretraining impact the world abstractions encoded in their representations, respectively. The details follow below.Effect of fine-tuning We adapt Llama2-7b/13b [35] , Mistral-7b [20] , and LLama3-8b [1] on both datasets, using in-context learning (ICL) and supervised fine-tuning (SFT) with LoRA adapters [19] . Appendix F provides more details.Effect of pre-training Since the LLMs above perform near-random without fine-tuning, we also experiment with a more recent state-of-the-art LLM, Phi3-17b [2] , which achieves much higher performance than other pre-trained LLMs on GRIPPER 4 . For comparative analysis, we also include Phi3-3.8b, Pythia-70m [6] and a fine-tuned Phi3-17b.In addition, we train a 6-layer decoder-only Transformer from scratch, probing for abstractions from its internal representations as a baseline. Before the probing experiments, we first verify that the fine-tuned LLMs and Phi3-17b ICL can achieve reasonable performance on REPLACE, and thus, we can draw informative conclusions. We adopt three evaluation metrics for REPLACE: legal rate (the proportion of predictions that comply with constraints), success rate (the proportion of predictions that achieve the target situations), and optimal rate (the proportion of cases where the target situation are achieved with the minimum number of actions). The Llama2-13b and Mistral results are reported in Table 1 , with comprehensive results for all LLMs provided in Appendix I. As shown in Table 1 , both LLMs initially fail in nearly all cases when relying solely on in-context demonstrations. Once fine-tuned, however, they can accomplish the task with reasonable success and optimal rate, and most of the predicted actions adhere to the constraints. Next, we conduct probing experiments on each variant of the LLMs on the same datasets.  
probing results and analysis: 
 We measure probes for all the LLM variants adapted to REPLACE (ICL or SFT) as well as the train-from-scratch Transformer. All LLMs SFT and Phi3-17b ICL , which outperform the random baseline by a substantial margin, primarily preserve Q * -and π * -irrelevant abstractions. Comparing LLMs SFT and LLMs ICL in Figure 5 and Figure 7 (a), Q * -and π * -irrelevant abstractions are probed from LLMs SFT with drastically higher RR than raw and world-irrelevant abstraction. Focusing on the RR of individual predicates, Figure 6 (a) and Figure 7 (b) reveal that the predicates for π * -irrelevant abstraction, nextObj and nextObjDirect have been mostly recovered from LLMs SFT . This result aligns with the high task success rate achieved by these LLMs. Similarly, the predicates for Q * -irrelevant abstraction are effectively probed from the LLMs SFT , although with larger variance across different predicates. This suggests that LLM representations do not merely encode the most coarsest world abstraction for deriving the very next action. Instead, they preserve sufficient information to estimate each action's long-term effect with respect to altering the world situation to match the target, facilitating efficient planning. The successful recovery of nearby, held u and agentLoc indicates that LLM representations capture the world constraints to inform legal actions, as evidenced by the near-perfect legal rate of LLMs SFT .Nevertheless, Figure 5 shows that although the RR margin between goal-oriented abstractions and more general ones is smaller for Phi3-17b ICL , it remains significant. Figure 6 (b) explains this by revealing that the RR for nextObj, nextObjDirect, and subgoal from Phi3-17b ICL is much lower than that of LLMs SFT , which is unsurprising as Phi3-17b ICL achieve a much lower success rate on REPLACE than LLMs SFT .Conversely, more general world abstractions are mostly absent in all LLM representations. Predicates uniquely tied to the raw state and the world-irrelevant abstraction cannot be accurately recovered. The low RR for boxName suggests that LLMs discard world details that are not pertinent to the task's completion. Similarly, the minimal RR on store t , a key predicate for reconstructing the world model, implies that the information to derive world dynamics is mostly omitted during decoding.We create two other variants of GRIPPER to further consolidate this finding. The first one is a counterfactual setting where a predicate originally irrelevant to the goal becomes pertinent to goaloriented abstraction. Our results show that it is accordingly encoded in the LLM SFT representations. Under the other setting, the LLMs are fine-tuned on sub-optimal action sequences, and the results suggest that random exploration does not necessarily improve the encoding of world dynamics. Details of these experiment can be found in Appendix K and Appendix L, respectively.Finding 2. Supervised fine-tuning and advanced pre-training mainly enhance goal-oriented world abstractions.Figure 6 and Figure 7 (b) demonstrate significant differences in the recoverability of goal-oriented abstractions from representations of LLMs, comparing models with and without supervised finetuning. Specifically, the performance gap in terms of predicates nextObjDirect, agentLoc, and subgoal is notable, approximately 40 -70 recovery rate. In contrast, for boxName and store, the difference is much more negligible, less than 5. This disparity implies that when fine-tuned with teacher forcing, LLMs evolve to develop a more goal-oriented world abstraction during decoding. This finding also provides a novel perspective on [21, 28] that successfully probe the world state in Transformer models: the state information is an integral part of the goal-oriented abstraction in their task, which is enhanced as the models are optimized to solve the task.Interestingly, advanced pre-training has a similar effect. Figure 5 and Figure 7 (a) show that the key distinction between train-from-scratch Transformer and LLMs SFT lies in their retention of Q *and π * -irrelevant abstractions. Nonetheless, as LLMs increase in scale and capability (Pythia<Phi3-3.8b<Phi3-17b), they are more likely to maintain goal-oriented abstractions over a more general one. This is apparent from Figure 6 (b), where the RR of predicates does not increase with Phi3-17b unless the predicates pertains to the goal-oriented abstractions. In contrast, the RR of store and held g , essential for world-irrelevant abstraction, are almost identical across all pre-trained LLMs. It implies that advanced pre-training does not necessarily lead to better encoding of the world dynamics. Finding 3. Raw state and world-irrelevant abstractions are mostly suppressed during decoding. Motivated by our initial two findings, we explore whether and how boxName and store are encoded within the contextualized representation by Llama2 13b SFT . To do so, we use the contextualized embedding Ctxt(•) of the label candidate, comparing the use of only the embedding Ctxt(e) against a concatenation [h t ; Ctxt(e)], inspired by [15] . We also examine nearby for comparative analysis. The results in Table 2 show that the objects' location can be probed from the contextualized representation with near-perfect performance while incorporating h t may decrease the performance instead. As for more goal-oriented predicates, nearby cannot be probed from the contextualized representation with high accuracy, and the concatenation with h t boosts the performance vastly, suggesting that LLMs discard the information of store and integrate it with agentLoc to derive nearby. Finding 4. LLMs are limited in building world representations, whether general or goal-oriented.While LLMs tend to preserve a goal-oriented abstraction, they do not totally discard boxName (15-20 RR from LLMs SFT and even 43 from Phi3-17b), which is irrelevant to either task completion or transition dynamics. Also, the probing performance for some goal-oriented predicates, e.g., subgoal, has ample room to improve. Moreover, the information for world dynamics, e.g., store, is mostly discarded. Remark. The interpretation of probing performance for a predicate should depend on the specific type of abstraction it belongs to.Focusing solely on raw world states might lead to a superficial conclusion that the LLM only partially captures the world. In contrast, our framework offers a more thorough and systematic interpretation, demonstrating that well-performing LLMs maintain a goal-oriented world representation during decoding. This underlines the core rationale of our probing framework: the necessity of probing different world abstractions rather than just the raw state. 
conclusion: 
 We propose a new framework to probe the abstract world state in LLM representations through the lens of state abstraction. Experiments with a synthesized task using our framework demonstrate that LLMs tend to preserve a goal-oriented world abstraction instead of a more general one during decoding, abstracting away the world's transition dynamics. Overall, our experiment findings highlight the importance of probing different types of abstraction encoded in LLM representations to draw comprehensive and nuanced conclusions. Nevertheless, our work is not without limitations, which are discussed in Appendix A. 
 
 Appendix: Limitations of our framework Our framework primarily investigates the presence of different abstraction types in LLMs, without delving into the mechanisms through which LLMs generate or link these specific abstractions. We leave these topics for future research. Additionally, as with many probing methods [34, 16] , we cannot ensure that LLMs accurately and effectively make use of the identified abstraction. Our research instead focuses on whether it is possible to recover certain types of abstractions from LLM representations.Limitations of our designed task While the task we designed in this work allows for precise and analytical derivation of world abstractions and involves fundamental skills essential for other NLP and reasoning tasks, it is important to note that it remains relatively simple and artificial compared to real-world NLP tasks. Consequently, future work can extend this task to more realistic and complex settings.Limitations of our datasets Following previous works in this line of research [24, 28] , we synthesize datasets for our task. Although synthetic datasets offer precise control over their properties and the characteristics of the underlying world, there may be discrepancies between our datasets and those used in real-world NLP applications. Unfortunately, realistic NLP datasets paired with a "world simulator"-vital for tracking state updates triggered by predicted actions in probing experiments-are lacking. We leave the development of such datasets for future work.In this work, we have run experiments under our proposed framework with two datasets and four state-of-the-art open-sourced LLMs. Despite the extensive experiments we have done, the findings with our experiments and framework are still specific to the LLMs and the particular task we have investigated, and should not considered universally applicable across other tasks and models. Nevertheless, the findings are important since they are new and complementary to existing literature. Further investigations could also examine how factors like fine-tuning methods, model scale, linguistic characteristics and training data distribution influence the abstractions in LLM representations.Our work introduces a probing framework designed to identify the types of essential information about the underlying world that can be retrieved from LLM representations. Our experiments show that without this detailed analysis, discrepancies can arise between evidence and conclusions. Such discrepancies may lead to undue pessimism about LLMs' lack of world awareness or excessive optimism about their ability to develop general world representations. Consequently, we expect that our work could be potentially helpful to other researchers and LLM users, enabling them to draw more nuanced and comprehensive conclusions from analyzing LLMs regarding their ability to construct world representations. This in-depth understanding is crucial for developing responsible AI systems.Table 9 shows examples of each dataset variant in GRIPPER. Table 10 provides the full task instruction included in the prompt.We generate the COOK dataset with TextWorld [8] . Concretely, we adopt the cook domain in TextWorld, where the agent navigates in a grid world with multiple rooms and manipulate the food ingredients to make a meal. Whereas TextWorld primarily focuses on evaluating RL with respect to react to the feedback from the environment, where the updated state information is included in the feedback, it makes the task of world modeling somehow trivial at a surface level. Therefore, we modify its setting to the one similar to GRIPPER, where only the initial and desired state are given in the prompt, and the models have to infer the intermediate world state induced by predicted or ground-truth actions. Nonetheless, there are two main differences between COOK and GRIPPER. First, the containers (rooms) in COOK are placed in a 2 × 3 grid world, where the agent can move south, north, east or west, and therefore have a larger action space. In contrast, the agent move either left or right in GRIPPER. Second, COOK uses a different set of textual templates adopted from TextWorld, which are listed in Table 3 . Table 11 shows the instruction for COOK and Table 12 We derive following abstract predicates using raw predicates, namely boxName, store, agentLoc and held:• nearby(o): the object o is located within the container nearby to the agent. It is derived by incorporating agentLoc with store. Similarly, distant(l, o), where l denotes the relative distance to the agent's current location, can be derived. That means l ∈[-L + 1, L -1],where L denotes the number of containers. For instance, distant(-1, o) being True indicates that the container one step to the left of the agent contains o.• subgoal(l, j): the agent needs to visit a series of containers to manipulate objects, and the j-th container to be visited is at the relative position l to the agent's current location. Furthermore, for subgoal(j, l), the label for l is in the set {-L + 1, . . . , L -1} ∪ {put, stop}, where put indicates the j-th subgoal is to put an object at the current nearby container, which is particularly useful for the first subgoal, as it allows for deterministic derivation of nextObjDirect. Also, stop implies that there are fewer than j subgoals. subgoal is derived by comparing distant, nearby and held across the current and target situation.• nextObjDirect(v): the relative direction of the object to be grabbed or put at the next step, where v could be left or right to the agent, in a nearby container, or with the agent. This is derived from the value of first subgoal, i.e., subgoal(•, 1).• nextObj(o): the object o to be manipulated at the next step, if applicable. This is derived by comparing nearby and held between u and g.We use Llama2-7b5 , Llama2-13b6 , Mistral-7b7 , Llama3-8b 8 , Phi3-17b 9 , Phi3-3.8b 10 , and Pythia-70m 11 as the base models for experiments. Additionally, we train another Pythia-70m model from scratch on REPLACE datasets as a Transformer baseline.For ICL, we randomly sample two demonstrations from the training set and append them to the prompt. For SFT, we train LoRA adapters [19] with 8-bit quantization on top of LLMs with the learning rate of 1e -4, the batch size of 8, the scaling factor of α = 64, the dropout rate of 0.05, the rank of 32 for 1 epoch. To avoid significantly altering the LLM's representational traits, we use only a subset of 10k training samples for SFT. We use 8-bit quantization for decoding. During decoding, we parse the generated textual description Ŷ of subsequent actions to its symbolic form with a rule-based parser. In our preliminary experiments, we have found using hard-crafted regular expression rules for the parser is good enough, since the the in-context demonstrations and the fine-tuning enforce well LLMs to generate in the same format as the data in Gripper and Cook.Recall that the raw state S in REPLACE is factorized by boxName, store, agentLoc and held.Lemma G.1. S ϕw is a world-irrelevant abstraction of S if it is factorized by store, agentLoc and held, excluding boxName.Proof. Assume ϕ w (s 1 ) = ϕ w (s 2 ). This implies that ∀a ∈ A, T ϕw (s ′ |s 1 , a) = T ϕw (s ′ |s 2 , a), given that boxName in S always remains unchanged after the execution of any action, we have:∀x ′ ∈ S ϕw , s ′ ∈ϕ -1 w (x ′ ) T (s ′ |s 1 , a) = s ′ ∈ϕ -1 w (x ′ ) T ([ϕ w (s ′ ); b s ′ ]|[ϕ w (s 1 ); b s1 ], a) = s ′ ∈ϕ -1 w (x ′ )T ϕw (ϕ w (s ′ )|ϕ w (s 1 ), a)EQUATIONwhere b s ′ denotes the assignment over boxName at s ′ . Similarly, it follows that ϕ w (s1 ) = ϕ w (s 2 ) ⇒ R(s 1 , a) = R(s 2 , a), since R(•, •) is invariant with respect to changes in boxName within a state s.Lemma G.2. S ϕπ is an π * -irrelevant abstraction of S if S ϕπ is factorized by nextObj and nextObjDirect.Proof. To demonstrate that S ϕπ is a π * -irrelevant abstraction, we show that the optimal action a * in any state can be uniquely identified using nextObj and nextObjDirect. This is formalized as follows:a * =        grab(O * ), if nOD = container put(O * ), if nOD = with-agent move(left), if nOD = left move(right), if nOD = right (2), where nOD denotes the value returned by nextObjDirect and O * denotes the object identified by nextObj. This ensures that for any two states s 1 and s 2 where ϕ π (s 1 ) = ϕ π (s 2 ), it follows that π * (a|s 1 ) = π * (a|s 2 ),Lemma G.3. S ϕ Q is an Q * -irrelevant abstraction of S if S ϕ Q is factorizedby agentLoc, held u , nearby, subgoal and nextObj.Proof. We start by deterministically extracting a sequence of containers, to be visited, represented by their relative distances to the agent's current location g. This sequence is denoted as g = [g 1 , ..., g j ],where for each i, subgoal(i, g j ) = True. nextObjDirect is inferred from g 1 . Consider the optimal action a * derived from nextObj and nextObjDirect, as discussed in Theorem G.2. The value of Q * (s, a * ) is given by:-r • j i=1 |g i+1 -g i | -r • j + r,where the first term represents the cost of traveling between containers, the second term accounts for the cost of executing actions, and the third term is a positive reward for achieving the desired situation. It follows that Q * (s, a) for all actions can be measured as below:EQUATION. where C(•) is the constraint function defined in Section 4, dependent on nearby u , held u and agentLoc.Consequently, Q * (s, a) can be precisely recovered conditioned on the predicates agentLoc, held u , nearby, subgoal, and nextObj. This implies thatQ * (ϕ Q (s 1 ), a) = Q * (ϕ Q (s 2 ), a) ⇒ Q * (s 1 , a) = Q * (s 2 , a).We detail the formulation and evaluation of the probing for each predicate as follows:boxName: Given an LLM-built representation h t , and a set of candidate container names b, as well as the embedding φ(b) for each b ∈ b, the probing model M i θ for i-position learns to select the ground-truth container name b * by maximizing the probability of boxName(i, b * ) being True:J θ = -log exp M i θ ([h t ; φ(b * )]) b∈b exp M i θ ([h t ; φ(b)]).During testing, we include the container names in the same data instance as candidates and measure the F1-score using a macro average.agentLoc and nextObjDir: Both of these predicates are formulated as multi-class classification tasks. Given h t , the probing model learns to maximize the ground-truth label, i.e. i * for agentLoc(i) and v * for nextObj(v):J θ = -log exp (P M θ (i * ; h t )) i exp (P M θ (i; h t )).We adopt the F1-score as the evaluation metric.store, held and nearby: These predicates may have multiple labels. For instance, a container may hold a number of objects. Therefore, we decompose a predicate into multiple independent decisions. For instance, we let the i-th probing model M θ for store(i, •) learns to perform pairwise binary classification for store(i, o), where o ∈ O:J θ = q(i, o)• log P M i θ ([h t ; φ(o)]) + (1 -q(i, o)) • log 1 -P M i θ ([h t ; φ(o)]), where q(i, o) = store(i, o) . During testing, we combine all the predictions for store(i, o) for all candidates of o. We use the exact-match score method to evaluate the predictions of store(i, •), held(•), nearby(•), and we calculate the average score as the final metric.We formulate and train the model in the same way as boxName(i, •). Again, we adopt the F1-score as the evaluation metric.subgoal: As mentioned above, we train a separate probing model for the j-th subgoal. Each j-th subgoal is formulated, and the probing model is accordingly trained in the same way as nextObjDirect and agentLoc. We use the F1-score for evaluation.When we compare probing performance across different predicates, we normalize all the metric score to the recovery rate, as described in Section 5.We create probing data with GRIPPER and COOK as well as all variants of LLMs. Following [28] , probing models are built upon two-layer fully-connected neural networks. Consistent with their observation, we have found that the probing models with only one layer have worse performance on some predicates. For instance, when tested with Llama2-13b , the probing performance for nearby and subgoal drops to 80.91 and 55, respectively, whereas performance for other predicates remains nearly unchanged (with variances of less than 1 score). On the other hand, using more layers (e.g. three layers) does not lead to improvement for any predicate. We explore two types of embedding methods φ(•) and select the one with the largest performance margin between using [h t , φ(e)] and using only φ(e) on the validation set. As a result, we use word embedding for boxName, store, held, nextObj, and contextualized embedding for nearby. We explore multiple variants of m for extracting h m t : the last layer, the last 6th layer, and the last 12th layer. Among these, the last 6th layer has been found to lead to the best overall performance on the validation set. We also explore different dimensions of neural probing models: 6600, 7600, 8600, and 9600, selecting 8600 based on performance on the validation set.For the probing models takes as input the concatenation [h t , φ(e)], we uses a trick that transforms the it to [h t + φ(e); h t -φ(e); h t • φ(e)] as the input of neural networks. In our preliminary experiment, we have found it leads to better probing performance.All probing models were trained with the learning rate of 1e -3 and the batch size of 64 for 30 epochs. The main probing experiments were conducted three times, while the others were conducted once. All experiments were run using a single Nvidia A100 GPU, and each of them was completed within 24 hours.Table 4 presents the planning performance of both variants of all LLMs for REPLACE.Table 5 and Table 6 present the probing performance for different predicates across all variants of LLMs on GRIPPER and COOK, respectively. We conduct the probing experiment for each predicate with each LLM variant three times using different random seeds and report the mean value. We also report σ, the mean variance across different runs. Specifically, we average the variance across different LLM variants. The low variance observed indicates that our experimental findings are robust. In the original setting of GRIPPER, the color c of object o, denoted by color(o, c), plays the role of distracting information that is totally irrelevant to task completion. We create another setting, COLORGRIPPER, where the color information is necessary to accomplish the task, and we compare the probing performance with respect to color information under these two settings. Specifically, COLORGRIPPER includes a new predicate color g , the color of each object in the desired situations, and an action paint(o, c), which paints the object o in a nearby container to color c and c ∈ {red, blue, green, yellow}. Also, we incorporate a new constraint that the new color must be different from the original one. As such, there is a new abstract state nearNotColor(o, c), which classifies that the color of a nearby object o is NOT c and it pertains to Q * -irrelevant abstraction. We fine-tune two Llama2-13b models on GRIPPER (Llama2 org ) and COLORGRIPPER (Llama2 color ) respectively, and probe the nearNotColor and color from both fine-tuned LLMs. The results in Table 7 show that the probes fail to recover color, which do not belongs to goal-oriented abstraction under both settings, from Llama2 org and Llama2 color . However, nearNotColor can be probed from Llama2 color with reasonable recovery rate. Again, this confirm our finding 1 that LLM SFT tends to maintain a goal-oriented world abstraction during decoding.Can the LLM representations become more general if we enforce greater random exploration within the environment? To answer this question, we synthesize another variant of GRIPPER, where groundtruth action sequences are sub-optimal. Concretely, they contain random legal moves that do not lead to the goal. We compare the average recovery rates (RR) for each abstraction from Llama3-8b models, fine-tuned on both the original and new datasets. The results are reported in Table 8 . With the new dataset, the RR of goal-oriented abstractions decreases, which is expected as the model imitates a sub-optimal policy. However, the margin between goal-oriented and general abstractions is still substantial. More importantly, the RR of world-irrelevant abstraction does not increase at all, which implies that random exploration does not yield a more general world abstraction.You are an intelligent agent in a fictional environment with a sequence of containers arranged in a straight line.The environment involves movement, interaction with objects, and strategic planning to reach a particular state. It is originally referred to as model-irrelevant abstraction[29]. Here, we rename it to avoid confusion between world models and large language models.We will omit the layer index in the following sections for simplicity.We do not conduct probing experiment with the pre-trained Phi3-17 on COOK as it does not outperform the random baseline significantly.https://huggingface.co/meta-llama/Llama-2-7b-chat-hf (We accept their Meta license before requesting the model.)6 https://huggingface.co/meta-llama/Llama-2-13b-chat-hf (We accept their Meta license before requesting themodel.)7 https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2(licensed by 


Decision: This paper emphasizes the importance of considering different scopes of state abstraction when analyzing internal representations of language models to determine whether they contain world models. The paper demonstrates these results by analyzing representations of several fine-tuned LMs on two synthetic planning domains, and finds that their representations are goal-oriented. After discussion the reviewers generally agree that the revised version of the paper will merit acceptance. There were some remaining concerns about the generality of the result, including the focus on synthetic tasks, the range of models considered, and most centrally the need to fine-tune which will alter the representation content. These concerns limit the paper from addressing the broadest form of the questions posed, but the reviewers do agree that the revised version will addresses some aspects of the questions and makes a useful contribution in doing so.

Reviews: review: paper_summary: paper_summary: In this work, the authors attempt to examine whether large language models (LLMs) possess representations that can work as the world model. To this end, they target different state abstraction levels, world-irrelevant abstraction, $Q^*$-irrelevant abstraction, and $pi^*$-irrelevant abstraction, which are based on state abstraction in reinforcement learning. Their objective is to figure out how much of each state abstraction can be recovered from the last hidden representations of LLMs, and the authors propose a text planning task named *RePlace*. Given the objective, RePlace is designed to have distinct state abstractions for the three abstraction levels as well as the raw state. The authors empirically confirm that both pre-trained LLMs and LLMs fine-tuned on this task fail to recover the world-irrelevant abstraction, whereas the $Q^*$-irrelevant and $pi^*$-irrelevant abstractions exhibit higher recovery rates.

**Post-Rebuttal/Discussion**: I find the author response fair and am raising my score to 7.
strength_weakness: strength_weakness: - The finding that the fine-tuned language models do not preserve the features needed for recovering the world dynamics may not be entirely surprising, as fine-tuning enforces the bias needed for the specific problem to the models and much of irrelevant aspects often end up being ignored.
- In addition to my first point, I believe more interesting findings can be derived with pre-trained LLMs, as there is a possibility that the pre-training provides enough signals to the LLMs to catch the world dynamics information. While the authors do analyze pre-trained LLMs with in-context learning in the proposed benchmark, according to Table 1, their performance is quite low in the first place, which makes it difficult to examine this hypothesis.
questions: questions: Please take a look at the Weaknesses section above.
limitations: limitations: The authors provide a fair list of limitations in Appendix A.
paper_summary: paper_summary: This work makes the observation that prior work in probing LLM for planning tasks comes to different conclusions on whether there are internal state abstractions in LLM hidden layers. This work hypothesizes that the disagreement comes from that these works are probing LLM with different tasks, which may necessitate LLM to learn abstraction at different granularities best suited for next-token prediction. Hence probing without controlling the abstraction level might give rise to different recovery rates of abstractions from hidden layers for different tasks. To support this claim, the paper borrows from RL literature the concept of world-/Q*-/$\pi^*$-irrelevant abstractions to designate different abstraction levels and aims to show whether they can recover these abstractions from hidden layers of Llama and Mistral models (both via in-context prompting and LoRA fine-tuning)
strength_weakness: strength_weakness: The biggest weakness I find is that the experiment may not be sufficient to support their central hypothesis. While it is necessary to show high recovery rate of Q*-/$\pi^*$-irrelevant abstractions for RePlace tasks, it is not sufficient to substantiate the claim that LLM will learn different abstractions given different tasks. RePlace tasks (both Gripper and Cook) is just one category of tasks whose solution requires forming abstractions such as Q value. Necessary additional experiments should show figure 6 on other tasks where 1. world-irrelevant features are important and 2. either Q abstractions or $\pi$ abstractions are important but not both to solve the task. If figure 6 shows different recovery rates for these other categories of tasks, there is sufficient evidence to support your claim.
questions: questions: A lot of the authors' findings (e.g. finding 2 "Supervised fine-tuning mainly enhances a goal-oriented world abstraction") are not necessarily true because they only fine-tune LLM on RePlace tasks that are inherently goal-oriented. Finetuning with other categories of tasks that are not goal-oriented may not lead to goal-oriented world abstractions. For example, if you modify your RePlace domains to generate a dataset of non-optimal plans (potential back-and-forth legal moves but no shortest paths to goals), fine-tuning with this dataset might not lead to purely goal-oriented abstractions.
limitations: limitations: The authors make the remark that "one cannot determine whether the success of probing stems from the LLM’s preference for learning a general world model or from the necessity to recover the world state while learning the optimal policy" in the critique of prior work using the Othello game. However, the proposed probing method might not prevent supervised fine-tuning on a particular task from washing out LLM's innate preference either. In other words, the probing results only reflect abstractions suited to solve the tasks they fine-tune on. Perhaps one can only use ICL results to probe LLM's innate preference for abstraction learning, and to that end, one might need to come up with tasks that LLM are better at solving zero-shot.
paper_summary: paper_summary: This paper investigates whether different levels of world abstractions can be decoded from LLM representations. The study is performed using a synthetic dataset of simple planning problems involving moving objects between containers. The study takes inspiration from RL to define different levels of abstraction, from specific goal-oriented states to general world-related states. The LLMs (Llama and Mistral models) are initially adapted to the task through in-context learning and fine-tuning. The problem is encoded in natural language and the hidden states are extracted and used to train probing classifiers to determine whether each level of abstraction can be recovered. The results show that the recovery of goal-related states is much more successful than general world abstractions.
strength_weakness: strength_weakness: Major comments:

1. Given that the models are first fine-tuned to the task, I wonder whether the results are not a self-fulfilling prophecy of the setup. If full problems and solutions were presented during adaptation, then it seems expected that the model would produce representations that are relevant to this task - the paper result would be, in a sense, a measure of the LLM having adapted to the fine-tuning, as opposed to reflecting the general inner workings of an LLM faced with a realistic or novel problem. A more convincing experiment would involve no task adaptation.

2. It is not clear how the hidden states were selected for probing. Only the "last hidden states" were selected (section 3.2). This again seems to invite the paper's results - it seems like economically preserving only relevant representations in the last layers is a measure of the LLM adaptation to the task. How about early layers of the LLM - can general abstractions be recovered there? Additionally, can any abstractions be recovered at all before adaptation?

Minor comments:

3. Some of the figures lack clarity. Figures 5-6 have unreadably small fonts, tables have barely readable fonts. "Best viewed in color" is a useful warning, but better colour selection would be ideal.

4. Typos: L326 "the most coarsest". Additionally, expressions like L318 "LLMs prefer maintaining a goal-oriented world abstraction over a more general one" are excessive LLM personification.
questions: questions: See above.
limitations: limitations: Adequately discussed, but only in the Appendix, not in the main text.
paper_summary: paper_summary: This paper proposes a new framework for studying world state abstractions from LLM representations. The framework, based on state abstraction theory, focuses not on assessing whether a model has a single world representation but assessing different levels of possible abstractions. These are each roughly functions of states: world-irrelevant abstractions, q*-irrelevant abstractions, and pi*-irrelevant abstractions. A probe is used to assess whether an LLM's representation is encoding each abstraction. Specifically, the authors design a planning task called REPLACE (where each abstraction is known to the researcher but not to the model), and train a probe on a transformer's representation to predict the abstraction when prompted with a textual description corresponding to the abstraction.

In experiments, the authors prompt LLMs to perform the REPLACE task and assess which abstractions each LLM is encoding. They have four findings:
1. LLM representations reflect goal-oriented abstractions during decoding  
2. Supervised-fine tuning increases the level of this goal-oriented abstraction
3. We don't see the other abstractions present during decoding
4. LLMs do not appear to build complete world representations.
strength_weakness: strength_weakness: While the framework proposed in the paper is interesting, the main weakness of the paper is that the experimental results are limited: REPLACE consists of two related planning task based on a limited set of containers and objects. Since the paper is focused on assessing world representations generally, there needs to be more evaluation settings and datasets. For example, the paper frequently refers to Othello as a testbed for assessing LLM world representations, but doesn't include any experiments involving Othello. The findings in the last section of the paper are based on insufficient evidence since they're only using the two related planning tasks that make up REPLACE. Do the same trends hold for game datasets? At the very least, more settings for REPLACE should be considered to provide more ablations of where each abstraction is recovered. 

Similarly, the paper could benefit from more extensive model comparisons to make sure the takeaways are robust. While the authors do cover multiple LLMs, they're all from the Llama family or Mistral. Using models from different families (or at least the larger sizes of the Llama/Mistral models) would help make the findings more compelling.

Less importantly, there are a couple of writing and presentation suggestions that would improve the paper. The text in Figure 5 is difficult to read because it is small. Moreover, although Q* has a standard definition for anyone familiar with RL, it is not defined in the paper even though RL problems are defined.
questions: questions: See questions above relating to robustness.
limitations: limitations: Yes
paper_summary: paper_summary: This work investigates what kind of abstractions LLMs use to encode the world, distinguishing goal-oriented abstractions (discarding world dynamics that are not necessary for achieving the goal) from world-general abstractions (including dynamics irrelevant for the goal). The authors note that prior work looking at LLM world models don't make this distinction, leading to conflicting results. For a text-based planning task, the authors probe LLMs doing the task in-context as well as fine-tuned, finding that from the latter goal-oriented abstractions can be recovered with a relatively higher accuracy than from the former.
strength_weakness: strength_weakness: **LLMs can only do this task after task-specific fine-tuning, in which case it is unsurprising that the representations are goal-oriented. This doesn't mean that LLMs doing tasks out of the box (or in-context) well won't use world-general representations**

The main result is that LLMs fine-tuned for a task start forming more useful goal-oriented state abstraction representations. This is unsurprising, as when fine-tuning the LLM for a task you're essentially training it to discard irrelevant information and encode relevant information for the goal. I would expect world-general abstractions to only be present in a general model that can do a task well out of the box without task-specific fine-tuning. However, the result is presented in the paper like LLMs do not use world-general representations (e.g. from the abstrcct: "Our experiments reveal that LLM-built representations tend to favor goal-oriented abstractions during decoding,
prioritizing task completion over recovery of the world’s state and dynamics."). For this kind of claim to be sufficiently substantiated I would like to see the framework applied to a task that LLMs can do out of the box, or this task using LLMs that can do it in-context (e.g. larger models or better models). Alternatively, I'd like to see the text indicate that the claims are about task-specific (i..e. fine-tuned) LLMs.

**Missing baselines**

It's difficult to interpret the results without a simple baseline that learns a probe on top of a non-LLM trained to do this task. Now, we don't know whether the accuracy of around 20% for world-general predicates is because the LLM has picked up general knowledge from pretraining or because you can learn to represent these things from this task (either in-context or through fine-tuning.)
questions: questions: **Questions**
- For section 6.1, it would be useful to know what percentage of random moves would be legal. It seems like the ICL performance of llama2 and mistral are very low / close to random?
- I would expect more world-general abstractions to be represented in earlier layers, and goal-oriented abstractions in the later layers. Did you look at the earlier layers for world-general predicate accuracy as well (first half)?

**Suggestions**
- I would caveat the abstract (and results mentioned elsewhere) sentence: "Our experiments reveal that LLM-built representations tend to favor goal-oriented abstractions during decoding, prioritizing task completion over recovery of the world’s state and dynamics.". You only find this for fine-tuned LLMs, and it might be entirely different for LLMs that can do a task out-of-the box.
- Should the predicate be subgoal(l=1, j) and subgoal(l>1, j) in line 270 - 271? It says in line 252 that l is the relative distance and j the j-th container to be visited.
limitations: limitations: Discussed and addressed in the appendix.

Comments: title: Further Clarifications (3/3)
Response to Rebuttal
Further Clarifications (2/3)
Further Clarifications (1/3)
Thank you for your review; we're ready to answer any further questions
Thank you for your review; we're ready to answer any further questions
Thank you for your review; we're ready to answer any further questions
Thank you for your review; we're ready to answer any further questions
Thank you for your review; we're ready to answer any further questions
Footnotes for rebuttal
Footnotes and Reference List for rebuttal
Footnotes and Reference list for Rebuttal
Footnotes and Reference list for Global Rebuttal
comment: Dear Reviewer,

We're thrilled to hear that most of your concerns have been addressed. We're truly grateful for your open-mindedness in considering our points and taking the time to review our responses thoroughly. To be honest, it's not often we encounter this level of engagement during the rebuttal process and we really enjoy our fruitful exchanges! We agree that assessing whether and how LLMs encode general world abstractions will be a long-term problem that demands continuous investigation. While our work may not provide a final answer, we hope it serves as a solid starting point. Rest assured, we'll make sure the promised changes are carefully implemented in the revised version!
> "we can agree that SFT is an important type of NTP training, and our new experiments probing pre-trained LLMs address this concern."

I do agree, and the main thing I was trying to get at was the way the results are presented in the paper. Indeed, with your additional experiments probing the pre-trained LLM the concerns are further alleviated, and the only thing that is left is the presentation.

> "You might consider LLMs-icl results uninformative due to their poor performance on the task" and "Fine-tuning NTP doesn't result in a world model" and "pre-training on NTP doesn't inherently produce internal world models"

Indeed, your results show that FT doesn't result in an internal world model for this task, and indeed the results on ICL are uninformative due to poor performance. Of course one can't such much about other tasks, and general claims about LLMs not having internal world models cannot be mode. This paper gives additional evidence to the full picture, but LLMs might very well be learning internal world models for other tasks. That is by itself not at all a weakness, such evidence is important, and I think your paper is a strong addition to the literature together with your supplemented results during the response phase. That's the reason why I increased my score, but not yet further than 5 because of the substantial rewrites I believe necessary.

> "These results directly go against the claim from previous work that Transformers/LLMs trained on NTP develop internal world models. Moreover, we want to emphasize a meta-result that diverges from previous work: probing without considering world abstractions leads to biased evaluation and unnecessary conflicts, as reiterated throughout the paper and rebuttals."

I agree, you're right, please excuse me for misrepresenting your results on this particular aspect in my previous response.

>  "Incorporating the new results into the paper"

This is my main remaining concern, and your details on how this will be done adresses my concerns. 

## Summary

Thanks for bearing with me, and thanks for responding to everything so thoughtfully. I am going to increase further to a 7. The way the authors engage in respond in this rebuttal phase make me relatively confident the proposed changes will actually be implemented. Both my weaknesses, about missing baselines (probe on top of non-LLM) and investigating pre-trained LLMs, are adequately addressed. That means the paper makes an important and interesting contribution of distinguishing goal-oriented and world-general abstractions when investigating world models, and present interesting findings that LLMs (both fine-tuned and general) do not encode world-general abstractions for their task. I do believe that to really show LLMs "in the wild" do not encode world-general abstractions it would need to be shown on more different and realistic types of tasks. Of course, there is always the risk of contamination, and that is a good reason to also show results on synthetic tasks, but for more general conclusions more tasks need to be considered. Nonetheless, the above contributions are sufficient for acceptance to me.
Dear Reviewer,

We appreciate the effort you've put into reviewing our paper and offering really helpful feedback. We're wondering if you've had a chance to go through our responses. As the discussion period is nearing its close, we want to ensure that any remaining questions or concerns are fully addressed. If you feel that the main concerns (suprisal of findings & the selection of hidden states) have been satisfactorily resolved, could you please consider increasing the score to reflect that? We would be truly grateful.
Dear Reviewer,

We appreciate the effort you've put into reviewing our paper and offering really helpful feedback. We're wondering if you've had a chance to go through our responses. As the discussion period is nearing its close, we want to ensure that any remaining questions or concerns are fully addressed. If you feel that the main concerns (surprisal of findings and probing with pre-trained LLMs) have been satisfactorily resolved, could you please consider increasing the score to reflect that? We would be truly grateful.
Dear reviewer,

We've just saw your most updated response. We truly appreciate your active participation in the discussion phase, especially during what might also be a busy week for you. We're grateful to hear that we've addressed some of your concerns and that you find our findings interesting. This feedback boosts our confidence and excitement to share our work with a broader community. Thank you once again for your invaluable feedback; it has greatly enhanced the quality of our work! We will be sure to incorporate the promised changes in the revised version.
### **Making the scope of findings clear in abstract/introduction**


> I disagree that it's a good reason to not include the fact that your results are only on fine-tuned LLMs in the abstract due to space constraints.

To clarify the findings in our experiment section (**Finding 1**): (1) goal-oriented abstractions are much more effectively recovered from LLMs-sft representations than world-irrelevant abstractions, and (2) world-irrelevant abstractions are largely absent in both LLMs-icl and LLMs-sft representations. Thanks to your thoughtful reminder, we realize that we should have summarized the findings for LLM-sft/icl separately in the abstract and introduction. **As stated in our rebuttal, we will ensure this is clearly clarified in these two sections**. We are also open to addressing any remaining confusion.

---

### **Incorporating the new results into the paper**

> I believe the paper requires quite a substantial rewrite to include the right baselines (non-LLMs) and the experiments using a pretrained LLM from the general rebuttal


In our humble opinion, these would be easy adjustments. First, the new results support our original findings, and second, incorporating them involves simply: (1) adding the new Plots A and B from the Global Rebuttal, presented in the same format as Figures 5 and 6 in our paper, along with the corresponding descriptions (points 1-3 in the Global Rebuttal); and (2) including the results with non-LLM baselines in Plot A and Figure 5, represented by one additional bar each.



We sincerely hope the clarifications above encourage you to reassess the value of our work. We're eager to address any further concerns you may have. Thank you again for your time and dedication in the review process.

---

### **Footnotes**

$\color{red}{^1}$ We are not intentionally downplaying the validity of [3]. We know it's one of the earliest works in this space, raising an important question and inspiring many follow-up works, including ours.

$\color{red}{^2}$ We agree that it's important to distinguish the effects of pretraining from fine-tuning when drawing conclusions. Therefore, we compare probing results from LLMs-sft and LLMs-icl (L342-350, **Finding 2**) to investigate how SFT impacts the encoding of world abstractions. The new results (again, thanks to your suggestion) using pre-rained LLMs and non-LLM models help clarify the effects of pre-training, as discussed in the Global Rebuttal.

$\color{red}{^3}$ To clarify, we're just explaining that this assumption is valid and inspired by existing work, and hence findings based on LLMs-sft could be informative and insightful. We're not suggesting that experimenting with LLMs-sft is sufficient for drawing conclusions for pre-trained LLMs.



### **Reference**

[1] Kenneth Li, "Do Large Language Models learn world models or just surface statistics?", The Gradient, 2023.

[2] Li, Kenneth, et al. "Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task." ICLR, 2023.

[3] Li, Belinda Z., Maxwell Nye, and Jacob Andreas. "Implicit Representations of Meaning in Neural Language Models." ACL, 2021.

[4] Hazineh, Dean S., Zechen Zhang, and Jeffery Chiu. "Linear Latent World Models in Simple Transformers: A Case Study on Othello-GPT." arXiv preprint arXiv:2310.07582 (2023).

[5] Kim, Najoung, and Sebastian Schuster. "Entity Tracking in Language Models." ACL, 2023.

[6] Prakash, Nikhil, et al. "Fine-Tuning Enhances Existing Mechanisms: A Case Study on Entity Tracking." ICLR, 2024.

[7] Panigrahi, Abhishek, et al. "Task-specific skill localization in fine-tuned language models." ICML, 2023.

[8] Zhou, Chunting, et al. "Lima: Less is more for alignment." NeurIPS, 2024.

[9] Ha, David, and Jürgen Schmidhuber. "Recurrent world models facilitate policy evolution." NeurIPS, 2018.
Dear Reviewer,

We're truly heartened and thrilled to know that our responses are able to address your concerns. This is incredibly important to us! Your interesting question has inspired us to uncover more insights. We will definitely incorporate them in the revised version.
Thanks for the really well-done rebuttal! Your response clarified a major misunderstanding on my end and thus cleared many of my resulting concerns about your experiments and results. I especially appreciate the added experiment on SFT with sub-optional data following my question. Therefore, I am happy to raise my scores.
### **Clarification about the conclusions from our experiments**

>So even though for your task it might be the case that world-general representations can be useful, the results seemingly can't say much about whether or not pre-trained LLMs or next-word predictors learn world models. 

With great respect, we have to point out that this is NOT true. 

1. In our original findings, world-irrelevant abstractions are mostly missing in both LLMs-sft and LLMs-icl, which are fine-tuned and pre-trained on NTP, respectively. While you might consider the LLMs-icl results uninformative due to their poor performance on the task, the LLMs-sft results clearly demonstrate that **fine-tuning on NTP doesn't result in an internal world model**.
2. Further experiments, prompted by your suggestion, on pre-trained LLMs and non-LLM transformers suggest that **pre-training on NTP doesn't inherently produce internal world models**—instead, more advanced pre-training leads to a higher priority for encoding goal-oriented abstractions.

---


>  I remain of the opinion that **the results** … **do not necessarily go against prior findings**. For example, you also **find world-general abstractions in LLMs**, just more goal-oriented ones.

This is NOT true, either. 

1. Figure 6 in our paper and Plot B in the Global Rebuttal clearly show that **the predicate cannot be reliably recovered unless it pertains to the goal-oriented abstractions**. 
2. Conversely, **the predicates uniquely tied to world-irrelevant abstraction**, namely `store_u`, `store_g`, and `held_g`, **are probed with a recovery rate of less than 12.5% across all LLM variants** (even lower than the one of `boxName`, which is unrelated to both task completion and world dynamics). Both types of NTP training—pre-training and fine-tuning—make little difference in this regard. It suggests that **LLMs tend to discard predicates that serve solely for preserving world dynamics**, the most critical ingredient of world models [9]. To be more concrete, LLM representations do not indicate how the world will change when an agent manipulates a certain object. 
3. **These results directly go against the claim from previous work that Transformers/LLMs trained on NTP develop internal world models**. 
4. Moreover, we want to emphasize a meta-result that diverges from previous work: probing without considering world abstractions leads to biased evaluation and unnecessary conflicts, as reiterated throughout the paper and rebuttals.

---

### **Why we don't work on a simpler task (such that LLMs can do well in context)**

> Of course, I am not suggesting you do an exhaustive search over all possible tasks, as I don't even know what that would mean in practice, but it seems highly unlikely to me that there exists no task open source LLMs can perform in-context that affords both types of abstractions.

We agree it's reasonable to expect a task that open-source LLMs can perform well in context, consisting of distinguishable world abstractions at various levels. However, finding such a task can be exhausting and tricky. Our task is already simple enough, as elaborated in our rebuttal and paper, so we didn't tweak the task setting until the LLM could achieve a reasonable performance. We were not implying that you suggested searching through all possible tasks. Apologize if there is any confusion. 

–--

> Finally, data contamination can be equally an issue in your fine-tuning setup, so I don't fully understand that argument.

Sorry for the confusion. By data contamination, we're referring to the possibility that the LLM might have been pre-trained on very similar or almost identical tasks. However, we're not suggesting that the LLM has encountered the exact same test samples during pre-training—this is a general concern across all setups. This possibility makes it unclear whether experimental results reflect SFT or pre-training effects. It's not an issue for fine-tuning setup, as we're fully aware that the LLMs are fine-tuned specifically for this task. If this still doesn’t seem like an issue to you, we can temporarily set it aside, as it doesn't affect the validity of our findings or how we've addressed your concern.


Simply put, these three points in our rebuttal were to explain why we initially didn't conduct experiments on pre-trained LLMs. However, these were relatively minor, as we later identified a better model during rebuttal, Phi3 (released after the submission deadline), which achieves acceptable performance without fine-tuning. Therefore, the additional experiments in the Global Rebuttal primarily address your concerns regarding pre-trained LLMs.
Thank you for your detailed feedback. We truly appreciate your engagement in further discussions and your acknowledgment that the facts we clarified in our rebuttal (the recovery of $Q^*$-abstractions and the relevance of world dynamics for the task) make our findings more surprising. It is also gratifying to know that our additional experiments echoing our original findings, provide deeper insights into world abstractions in both pre-trained and fine-tuned LLMs.



We address every single one of your remaining concerns below.

---



### **Clarification about the finding based on SFT and its relation to a common belief shared by the community**

In the Global Rebuttal, we stated that *"(Our $\color{blue}{\text{finding}}$ that SFT mainly enhances goal-oriented abstractions) undermines an increasingly widespread belief that **an implicit world model can emerge from next-token prediction (NTP) training** [1]"* and we cited two representative works [2,3] that encourage this belief. However, you raised concerns, arguing that our finding is insufficient because **[3]** *"explicitly controls for the effect on representations of pretraining and fine-tuning"* and they found *"While both fine-tuning and pretraining contribute to the final probe accuracy, **pretraining appears to play a much larger role**"*.
 
Indeed, it doesn't negate the fact that our $\color{blue}{\text{finding}}$ presents a sufficient challenge to the belief discussed above, as we will elaborate below.

(1) The next-token prediction (**NTP**) training mentioned in the belief is **not limited to pre-training alone**; the findings that encourage this belief go beyond just pre-training, including conventional supervised training for particular tasks (i.e., SFT or training from scratch). As you noted, even **[3]** has claimed that fine-tuning enhances world modeling. Moreover, the work [2] in this area that attracted wide attention, along with more recent follow-up work, e.g. [4], has probed from Transformer models trained on synthetic tasks with NTP. Therefore, our $\color{blue}{\text{finding}}$ directly addresses the belief since SFT on NTP is the de facto standard approach to adapt LLMs to a particular task.


(2) Furthermore, the results you quoted from **[3]**, *"semantic state can be recovered well from models with no in-domain fine-tuning"*, don't undermine our finding that pre-trained LLMs struggle to maintain a general world abstraction (in RePlace), as a reanalysis of it [5] has already revealed that the high accuracy in **[3]** was based on trivial cases, and the entity status indeed cannot be recovered from pre-trained LMs with reasonable performance (please refer to section 2 in [5])$\color{red}{^1}$$\color{red}{^2}$. 

(3) Although we don't make such a claim in the paper, it's reasonable that one may expect some insights into the world modeling capabilities of pre-trained LLMs from the findings on LLMs-sft, as several recent studies suggest that it's very rare that fine-tuning alters the pre-trained capabilities [6,7,8]; rather, it tends to accentuate them. This may indicate that SFT could potentially enhance LLMs's prioritization of specific types of word abstractions that emerge during pre-training$\color{red}{^3}$. The findings from our new experiment with pre-trained LLMs, combined with the original ones, seem to align with this general hypothesis.

With all that being said, we agree with you that **the $\color{blue}{\text{finding}}$ based on LLMs-sft should not be considered universally applicable to all LLM variants** (as we consciously warned in L495-498), including pre-trained LLMs. In response to your suggestion, we have conducted additional experiments, as detailed in the Global Rebuttal, to provide a more comprehensive assessment of the belief in question. Even if the findings from LLMs-sft do not fully address the world abstractions in pre-trained LLMs, hopefully, we can agree that **SFT is an important type of NTP training**, and **our new experiments probing pre-trained LLMs address this concern**.
Thanks for the detailed rebuttal, both personal and the general one. I will respond below.

## Using FT-ed models instead of base models

The fact that world dynamics are still relevant for the goal and the fact that you recovered the $Q^*$ abstraction from LLM-sft representations is more convincing to me and alleviates some of my concerns. However, in the general rebuttal you claim your findings are more relevant in light of related work claiming that your finding:

 *"undermines an increasingly widespread belief that an implicit world model can emerge from next-token prediction (NTP) training [1]. This belief is encouraged by recent studies demonstrating that world states can be recovered from Transformer representations [2,3]."*

However, [3] explicitly controls for the effect on representations of pretraining and fine-tuning, from their paper section 4.2:

*"To determine whether the advantage is conferred by LM pretraining or fine-tuning, we ablate either open-domain pretraining, in a -pretrain,+fine-tune ablation, or in-domain finetuning, in a +pretrain,-fine-tune ablation. [...] While both fine-tuning and pretraining contribute to the final probe accuracy, pretraining appears to play a much larger role: semantic state can be recovered well from models with no in-domain fine-tuning."*

So even though for your task it might be the case that world-general representations can be useful, the results seemingly can't say much about whether or not pre-trained LLMs or next-word predictors learn world models. 

In your rebuttal, you say: *""Given the space of possible tasks is extraordinarily huge, it's intractable to do an exhaustive search to identify a task that simultaneously satisfies (1) complex enough to differentiate the spaces of various abstractions, and (2) existing open-source LLMs can perfectly solve without fine-tuning. Even though we find such a task, it is hard to rule out the possibility of data contamination, especially considering that it still remains unknown if LLMs can solve a novel task that involves modeling the underlying world described by the text.*

Of course, I am not suggesting you do an exhaustive search over all possible tasks, as I don't even know what that would mean in practice, but it seems highly unlikely to me that there exists no task open source LLMs can perform in-context that affords both types of abstractions. Additionally, it's not necessary for the model to perform the task perfectly without fine-tuning. Finally, data contamination can be equally an issue in your fine-tuning setup, so I don't fully understand that argument.

Finally, respectfully, I disagree that it's a good reason to not include the fact that your results are only on fine-tuned LLMs in the abstract due to space constraints. The line that is misleading is the following:

"*"Our experiments reveal that LLM-built representations tend to favor goal-oriented abstractions during decoding, prioritizing task completion over recovery of the world’s state and dynamics.".*"

Which can be changed to

"Our experiments reveal that LLM-built representations when fine-tuned tend to favor goal-oriented abstractions during decoding, prioritizing task completion over recovery of the world’s state and dynamics.".

## Learning a probe

Thanks for these additional insights. I believe it's important that these, together with the results on a pretrained model you have in the general rebuttal, are a part of the paper, as without it it remains hard to say much about which part contributes to what (pretraining vs finetuning). 

## Summary

To summarise, some of my concerns are alleviated, and I am still of the opinion that making the distinction between world-general and goal-oriented abstractions is a really cool contribution and makes for an interesting paper. The additional baselines and results on pretrained LLMs are interesting. This makes that I'll raise to score to a 5. The reason I did not raise it further just yet is as follows:
I remain of the opinion that the results are somewhat unsurprising and do not necessarily go against prior findings. For example, you also find world-general abstractions in LLMs, just more goal-oriented ones. All in all I believe the paper requires quite a substantial rewrite to include the right baselines (non-LLMs), the experiments using a pretrained LLM from the general rebuttal, and to change the language about what exactly LLMs (when finetuned or not) do learn and how that relates to prior findings.
Thank you for your feedback. We regret that your concern about experiments on other tasks remains, although we believe we have thoroughly addressed the need for more extensive model comparisons in our rebuttal. If possible, we would like to respectfully reiterate several points for your kind consideration. 

 (1) Existing testbeds are NOT available for feasible and faithful probing;

 (2) Due to (1), we develop a new one, creating TWO datasets in one effort;

 (3) We have conducted two ablations with different task variants, fully addressing your request for "**more settings for ablations**";

 (4) We aren't aware of any work analyzing LLMs/neural models that proposes two or more tasks simultaneously.
Thank you for the rebuttal. While I appreciate the additional experiments and clarifications, I still have my original concerns about whether these conclusions can be made from just REPLACE. Therefore I will keep my original score.
Dear reviewer,

Thank you once again for your review! We are truly grateful for your appreciation. We hope our rebuttal clearly explains why our findings are surprising and important and that our additional experiments with pre-trained LLMs yield some other surprising and interesting findings and address all of your concerns. Please don't hesitate to ask any further questions. We are always ready to respond, and we're very looking forward to an engaged and productive discussion!
Dear reviewer,

Thank you once again for your review! We really appreciate your recognition of all aspects of our work, including our addressed problem, framework, presentation, and the design of our task and experiments. We believe the major concern might be due to some misunderstandings, and hopefully, our rebuttal has already cleared it up. Meanwhile, we hope that our additional experiments in response to your concerns and suggestions will make you more convinced about the soundness of our findings. Please don't hesitate to let us know if you have any concerns that are unaddressed or if we need to clarify something further. We're always on standby to respond, and we're very looking forward to an engaged and productive discussion!
Dear reviewer,

Thank you once again for your review! We are truly grateful for your appreciation. We hope our rebuttal clearly explains why our findings are surprising and important, and that our additional experiments with pre-trained LLMs, along with our clarification about the selection of hidden states, address all of your concerns. Please don't hesitate to ask any further questions. We are always ready to respond, and we're very looking forward to an engaged and productive discussion!
Dear reviewer,

Thank you again for your review! We really appreciate your recognition of our methodologies, new resources, and presentation. We hope our rebuttal clearly explains why our new task representative and the existing testbeds are not ideal for faithful probing. Along with the two ablation studies and additional experiments on new families of LLMs, we hope that your suggestion for more extensive experiments has been adequately addressed. If there are any specific setups you believe might challenge our claims/findings that we haven't covered, please let us know. We'd be glad to accommodate them if feasible. We're very looking forward to an engaged and productive discussion!
Dear reviewer,

Thank you again for your review! We really appreciate your recognition of our methodologies, new resources, and presentation. Hopefully, our additional experiments in response to your comments and in-depth elaboration on the surprisal of our findings and the broader significance of our framework can address most of your concerns. If not, please don't hesitate to ask us any further questions. We are always ready to respond. We're very looking forward to an engaged and productive discussion!
$\color{red}{^1}$ To clarify, we don't need to use different tasks to probe different levels of abstraction. Instead, we derive the abstraction functions, such as world-irrelevant abstraction and $Q^*$-irrelevant abstraction, for the same task and then probe the abstract state at various levels from LLM representations. All we need is to carefully select a task where the spaces of abstract states at different levels differ (highlighted in L162-166). This generic framework is introduced in Section 3, and we demonstrate how to apply it to our new task, RePlace, in Section 5.

$\color{red}{^2}$ In particular, there are three necessary criteria for faithful and feasible probing (explained in L179-184, L488-493): (1) processing text data; (2) the abstractions across various levels are distinguishable; (3) the world state can be easily manipulated and tracked. 

$\color{red}{^3}$ However, our findings can explain and resolve conflicts among existing works. Specifically, the contradiction may stem from various degrees of alignment between the raw states and goal-oriented abstractions across tasks. For instance, [1] successfully probes the Othello board state, which is already the goal-oriented abstraction for predicting legal moves (the task Transformers are trained on), whereas [2] struggles to probe the status of entities, as they have little overlap with goal-oriented abstractions of the task that LLMs are fine-tuned on.

$\color{red}{^4}$ As explained earlier, however, it's entirely possible for a predicate to pertain to the goal-oriented abstractions of Task A yet remain unrelated to Task B (i.e., Task A and B share the same state space but have very different reward functions). Our ablation study confirms that under this scenario, this predicate can be probed from LLMs adapted to Task A but not Task B. Concretely, we have created a new variant of GRIPPER where an originally goal-unrelated predicate (e.g., information of objects' color) pertains to $Q^*$ abstraction. As such, the recovery rate of color information has significantly increased from zero. This further consolidates our conclusion that SFT mainly enhances goal-oriented abstractions. Due to the limited space, we summarize this finding in the main body of the paper (L337-L340) and provide details in Appendix K.
### **Footnotes**

$\color{red}{^1}$ Gripper problem is a classic AI/planning task featured in the first ICAPS's competition [6].

$\color{red}{^2}$ In case you can't find it quickly, we've copied the details here. We synthesize another variant of GRIPPER, where ground-truth action sequences are sub-optimal, including some random legal moves that do not lead to the goal. We compare the average recovery rates for each abstraction from Llama3 models, fine-tuned on both the original and new datasets:

|    Models    | Raw | World | $Q^*$ | $\pi^*$|
| --- | --- |--- |--- |--- |
| Llama3-sft (original, reported in paper)  | 28.5 | 30.6  | 73.0  |  87.8  |
| Llama3-sft (new)  | 25.5 |27.4 | 63.3| 72.1 |


The results show that the recovery rate of goal-oriented abstractions decreases, which is expected as the models imitate a sub-optimal policy. However, the margin between goal-oriented and general abstractions is still substantial. More importantly, the recovery rate of world-irrelevant abstraction doesn't increase at all, implying that training with non-optimal ground truths doesn't lead to a more general world abstraction. These new results echo our main findings, and we'll add them to the revised version. 

$\color{red}{^3}$ While existing work focuses on **one or two** predicates, our probing experiments involve **ten** different predicates, as our framework requires a comprehensive assessment of various types of abstractions. Also, the main results reported in the paper are averaged across multiple runs. In total, **we completed \~500 runs of training and testing probes across all LLMs and predicates**.  


### **Reference**

[1] Li, Kenneth, et al. "Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task." ICLR, 2023.

[2] Hazineh, Dean S., Zechen Zhang, and Jeffery Chiu. "Linear Latent World Models in Simple Transformers: A Case Study on Othello-GPT." arXiv preprint arXiv:2310.07582 (2023).

[3] Li, Belinda Z., Maxwell Nye, and Jacob Andreas. "Implicit Representations of Meaning in Neural Language Models." ACL, 2021.

[4] Kim, Najoung, and Sebastian Schuster. "Entity Tracking in Language Models." ACL, 2023.

[5] Prakash, Nikhil, et al. "Fine-Tuning Enhances Existing Mechanisms: A Case Study on Entity Tracking." ICLR, 2024.

[6] McDermott, D.M., 2000. The 1998 AI planning systems competition. AI magazine, 21(2), pp.35-35.
### **Footnotes**
$\color{red}{^1}$ We mainly focus on reporting the findings for LLMs-sft since it is a common interest [1, 2, 3, 4] in this field that if Transformers/LLMs trained on next-token predictions develop internal world models .

$\color{red}{^2}$ While this result isn't very satisfactory, it is acceptable given that it's the best ICL performance we've seen with open LLMs. We didn't conduct probing experiments on COOK, because Phi3-17 still falls short of this dataset without fine-tuning (\~3% success rate). It's expected as COOK has a more complicated structure of containers, and larger space of actions.

### **Reference**
[1] Li, Kenneth, et al. "Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task." ICLR, 2023.

[2] Li, Belinda Z., Maxwell Nye, and Jacob Andreas. "Implicit Representations of Meaning in Neural Language Models." ACL, 2021.

[3] Hazineh, Dean S., Zechen Zhang, and Jeffery Chiu. "Linear Latent World Models in Simple Transformers: A Case Study on Othello-GPT." arXiv preprint arXiv:2310.07582 (2023).

[4] Kim, Najoung, and Sebastian Schuster. "Entity Tracking in Language Models." ACL, 2023.
### **Footnotes**
$\color{red}{^1}$ Although R4 (nz33) didn't straightforwardly suggest conducting experiments with pre-trained LLMs, their recommendation of conducting experiment with another LLM family is also addressed here.

$\color{red}{^2}$ Due to the space limits, we didn't thoroughly discuss this aspect in the paper. We'll add this discussion in the revised version.

$\color{red}{^3}$ **While some studies continue to produce "surprising" findings suggesting that Transformers/LLMs trained or fine-tuned on next-token prediction develop internal world models** [2,3,4,7], **other research provides counter-evidence, cautioning that the capabilities of world modeling are overestimated** [5,8,9].

$\color{red}{^4}$ Phi3-3.8b and Pythia-70 achieve success rates of approximately 4% and 1%, respectively, on GRIPPER .

### **Reference**

[1] Kenneth Li, "Do Large Language Models learn world models or just surface statistics?", The Gradient, 2023.

[2] Li, Kenneth, et al. "Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task." ICLR, 2023.

[3] Li, Belinda Z., Maxwell Nye, and Jacob Andreas. "Implicit Representations of Meaning in Neural Language Models." ACL, 2021.

[4] Hazineh, Dean S., Zechen Zhang, and Jeffery Chiu. "Linear Latent World Models in Simple Transformers: A Case Study on Othello-GPT." arXiv preprint arXiv:2310.07582 (2023).

[5] Kim, Najoung, and Sebastian Schuster. "Entity Tracking in Language Models." ACL, 2023.

[6] Abdin, Marah, et al. "Phi-3 technical report: A highly capable language model locally on your phone." arXiv preprint arXiv:2404.14219 (2024).

[7] Gurnee, Wes, and Max Tegmark. "Language Models Represent Space and Time." ICLR, 2024.

[8] Altmeyer, Patrick, et al. "Position Paper: Against Spurious Sparks-Dovelating Inflated AI Claims." ICML, 2024.

[9] Wang, Ruoyao, et al. "Can Language Models Serve as Text-Based World Simulators?." ACL, 2024.