Do LLMs Build World Representations? Probing Through the Lens of State Abstraction
Authors: Zichao Li
Yanshuai Cao
Jackie CK Cheung
Venue: NeurIPS 2024 poster (2024)
Keywords: Large Language Models
World Models
World Representation
Probing
Reinforcement Learning
State Abstraction
Primary area: natural_language_processing
TL;DR: We introduce a new framework for probing world abstraction within LLM-built representations, and our experiments with a text-based planning task demonstrate LLMs prefer maintaining goal-oriented abstractions during decoding.

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.

Do LLMs Build World Representations? Probing Through the Lens of State Abstraction
Zichao Li
Yanshuai Cao Borealis AI
Mila, McGill University zichao. 1i 1i@mail .ww.ill.ca
Jackie C.K. Cheung Mila, McGill University jackie. cheunggmmcgill..cm
Abstract
How large language models (LLMs) encode the state of the world, including the status of entities and their relations, described by 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 accomplish tasks. To instantiate this framework, we design a text-based planning as agent in environment and interacts with objects
task, where an LLM 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.
Introduction
Drawing inspiration from human mental models [9, 12, 22], Al researchers have introduced the con- cept 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 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 mod- els [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 implic- itly 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 discourse. Conversely, [24] report negative results when attempting a similar yet more challenging setting. In addition, [28] successfully extracts the board state of partially played Othello game from the internal representation of small-scale GPT model [7] trained to complete game scripts.
Work was done during a Mitacs internship at Borealis Al. Code and dataset: https://github.
htteitttuuuooooooeeee wwm//oomaollleeee Bon//oraalss//Ilmm waoliiiwwwliwwwwl--------o-----e world-abs
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 stud- ies 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.
Discourse: Emima enters a dimly lit room and finds three boxes aligned: Box A holds brass key, Box B a porcelain cup, and Box C both a vibrant lotus and an enigmatic clock.
General abstraction
Count-oriented abstraction
What does Box 8 contain?
Some tasks may require the complete re- covery of world dynamics in order to pre- dict future states, while others may get by on one that omits this information as it is unnecessary for task completion. For in- stance, consider Figure 1. When parsing a discourse into 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 op- erations, 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 mismatch between evidence and conclusion, causing undue pessimism about LLMS lacking awareness of the world or excessive optimism regarding their ability develop general world representations. For example, fine-tuning on all types of questions in Figure I might push LLMs to capture world dynamics, this does not guarantee the same outcome when fine-tuning solely on counting tasks.
Move the objects in Box A to Box C. What does Box C contain now?
Move the objects Box A to Box C.
How many objects does Box C have?
Figure 1: A discourse and two possible abstractions of the world state described by it (top half of figure). A general abstraction (termed as world-irrelevant ab- straction later) enables one to answer a wide range of questions. On the other hand, criet-oriented abstrac- tion is only applicable to the counting task.
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
abstraction, that allows recovery of world dynamics to more goal-orienter abstractions, guiding task completion while giving up on predicting future states. These include Q" -"rirrelevant and nr -iiirrelevant 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 buill 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 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 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 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 new framework to probe world abstraction from LLM-built representations that can be adapted for other NLP tasks. 2) We release 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
2

---

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.
2
Related Work
2.1 World models
The concept of world models in machine learning has deep connections with the concept of human internal models [9] 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 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 specific task [18]. In similar vein, Al 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, any, encoded by LLMs.
2.2 Probing LLM Representations
There a growing interest in probing for interpretable features in LLMs, most which are linguistic in nature, such 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 I.LM'' representations of discourse recover the ground-truth situations depicted in the text. This approach is further modified and formulated as 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 probing for comprehensive description of the underly- ing world, defined as a set of state variables. This approach has two main limitations. First, it fails to distinguish the function of specific state variable: Is it intended to maintain general representation of the world, or 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 general framework that examines the world abstraction encoded within LLM representations.
In contrast [28, 21, 37], the input in our task is text data instead of game script 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 perform decision-making task. We start with a reinforcement learning (RL) formulation, which is general enough to encompass 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 4-tuple: (S. AS T, J R). where E S represents the world state, A denotes possible actions the model can take, transition function T T(s'|ss, a) measures the probability of transitioning to state s' induced by from and R(s, a) is the reward.
3.1 Definition and Derivation of World State Abstraction
Our framework builds upon rigorous definition of world abstraction. In particular, we follow RL literature [29], deriving state abstraction function & - Sz that maps each state into an abstract state ((s) (|Sp) < (S)). In addition to the raw state S, consider three types of abstraction:
3

---

World-irrelevant abstraction ow ensures that V8s, S2 where (s)) = (s() implies that Va, E SS... R(s)) a) R(s2, a), and Eveont(rr (ar) [(s"[s), = Ev (ar) T(s' [22 Intu- itively 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* or-irrelevant abstraction &Q ensures that Vs1, 82, if (Q( Q* a), where Q" a) max; Q. (s, a). This form of abstraction preserves the effect of all actions E A in terms of their optimal Q-value, which is the maximal expected future rewards. However, it discards the world dynamies, making future state predictions infeasible. L.LMs can leverage this abstraction to be cost-sensitive, minimizing action counts and avoiding penalties for violating world constraints.
(s2), then Va, Q*((1, a)
T"-irrelevant abstraction OR guarantees that the optimal action, and hence the optimal policy, can be recovered. Formally, Vsi, $2 E 5 ox (s1) = (52))) implies that x¹ (s))) = x" (s2). 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.
3.2 Probing World Abstraction from LLM-built Representations
We aim to assess which types of state abstraction are encoded LLM representations. To do so, we first prompt pre- mre-trained or fine-tuned LLM with r(su, A1:-1, 2). which is textual description of the initial state su. previously executed actions A1 and optionally, feedback 2 received from the world. The feedback could be the opponent's 's moves in an Othello game or users' utterances in a dialogue system. Next, we extract the hidden states H (mm) from the m-th layer of the LLM, considering m as a hyperparameter. Following [28]. we select the last hidden states h; ((m)), which is used by the LLM to predict subsequent actions. After collecting (st, he) pairs, where 8, is the world state induced by 8p, A1:-1 1 and 2, we detect the existence of abstraction of & in L.M representations by assessing if ((s²) can be probed from he with accuracy surpassing a random baseline, potentially achieving near-ped near-perfect performance. Previous work [16, 17] on probing presents challenge in differentiating whether 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 I.LM representations is not always successful.
Remark: To draw faithful conclusions from probing experiments, is crucial to carefully design or select the tasks in a way that ensures the spaces of abstract states at different levels differ (Sa.,. + Saa f 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 π -irellevann abstraction for predicting legal moves. Therefore, a plausible interpretation for the success of probing in [28] could be that the Transformer model learns the o -rrrreelevant abstraction rather than deliberately learning a general world model.
In the next section, we design a planning task within an RL framework and synthesize dataset accordingly. In Section 5, we prompt LLMs to perform this task and probe different types of abstractions from their internal representations.
4 REPLACE: A Text-based Planning Task
To instantiate our framework, we draw inspiration from recent work on probing discourse repre- sentation in I.LMs [27, 24] and design a text-based planning task named REPLACE. In REPLACE,
It originally referred to as model-irrelevan abstraction [29]. Here, we rename to avoid confusion between world models and large language models.
We will omit the layer index in the following sections for simplicity.

---

([Task Instruction]]
A sequence of containers are ordered from left to right as follows: Bucket A, Basket A, Basket B, Crate A, Basket C. You are at the Basket B.
Initial situation: Bucket A is occupied by a car, a key. Basket A contains a creature, a letter. Basket B has beer. Crate A holdy cout. Basket C holds crown.
Desired situation: Buckel A contains a car. Baskel is occupied by a creature. Basket B is empty. Crate A holds a beer, a coat. Basket C has crown. Operations applied: You Move left.
What are the operations to achieve the desired state?
Figure 2: An example of task input in REPLACE.
boxName
[Bucket A, Basket A, Basket B.
Crate A, Basket C]
Rass stote
nextObjDirect
action
agentLoc
2
World abstrect state
storeu
[earr key). (ereature, letter). (beer), (coar), [crown]]
Q' sbstract state
subgoal
nextObj
stores
Hear), (ereature). O. {coat, beer), (crown]
di stantt
Hear, key) [baer)-1, (eoat) (crown)+al
distant,
distant,
nearby,
nearby,
held,
held,
distant,
[[err] -1. (1-,,
fcoat, beer) 12, [crown] s
nearby,
(creature, letter)
storey
store,
agentLox
boxName
nearby,
(craature)
Figure Derivation of abstract predicates and action from raw predicates, Better viewed in color. Predicates represented by multicol- ored squares indicate association with multi- ple abstraction levels. Grey squares represent intermediate predicates, not associated with any abstraction, used to derive abstract predi- cates.
he held, heldy
{ (key, letter)
subqoal nextCo)
[0. 1, 1, +2]
letter
nextob-jiirect nearby-container
Figure 4: The raw and abstract predicates of the state described in Figure 2. u, g: current and target situa- tions.
an LLM acts agent within simple world consisting of 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 the gripper problem [31] and the entity tracking task [24], which are widely adopted in planning and NL.P 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.
4.1 Specification within RL Formulation
Formally, the specific design of the elements within the RL framework is described follows:
State s = [u, g] E S: the current situation u and target situation g of the world, each of which is factorized as the Cartesian product of set of assignments over predicates applied to entities including objects o and containers b. The predicates are:
st stcre(i, o): the object is stored in the i-th container from left to right.
he)d(o): the agent has the object o.
LgentLoc(i): the agent is at the i-th container.
boxName (i, 6): the name of the i-th container from left to right is b.
2. The possible actions E A include:
move(d): Move along a direction d, l or ri ght, for one step.
grab(o): Grab object from the nearby container.
5

---

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,)) E {0, 1}. Moreover, the transition impacts only the current situation un in St by taking the effect of For example, if ag is True in ut-1, executing the action move(left) results in agentLoc(1] becoming True in ut-
4. Reward function R(s, a) checks whether the new situation of o, induced by executing a, matches the target situation: R(s, a) r (u')) = Pa(g); or -(r + r') if a & C(s); or -r otherwise. Here v' denotes the new situation, F.(u') ) the set of assignments on predicates involving objects (store and deld) in the new situation n', and C: S + A. the constraint function mapping states to permissible actions. We incorporate three spatial constraints: (I) 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 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 and r' but simply assume 0 and r' > O. 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 9). previous actions (A1.z {a1,..,..) (:)) to their respective natural language descriptions T, (u), Tg(g). (Ar:)) and using a set of textual templates: T.(+) and a. As such, the input text to the model comprises two main components: (1) general task instruction (see Appendix C) specifying the actions the agent can take, (ra(a)a E A), and the constraints it must adhere to, T(C), and (2) textual descriptions of the world's initial and target situations, as well as previous actions TT Ta (A..)) An example is presented in Figure 2. Conditioned on this, the L.M generates textual descriptions Ý of subsequent actions, which are parsed back into their symbolic form A by a rule-based parser To (Y))
4.2 Datasets
We synthesize two English datasets for REPLACE: GRIPPER and COOK. For each instance in GRIPPER, we first sample set of from 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. [uo. g]. by randomly assigning values the predicates in ug and g. such that it requires the execution of 2 to 6 actions to transition from un to Towards a more realistic setting, we introduce the following dataset variants: (I) 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 r(g), which omits explicit information about he in the target situation, requiring logical inference by LLMs; (3) Partial [T(u). (g)). an extension of (2) that includes previous operations performed by the agent, thus requiring inferring 44 Appendix provides concrete examples for each variant. 'he final dataset is uniform mixture of these four variants and segmented into training, validation, and test sets, with [wo: g) splits, allocating 50k, 1k, and Ik instances to each set, respectively.
COOK generalizes the setting of GRIPPER to grid world of containers and is generated with Werld [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.
5 Probing World Abstractions in LLM Representations using REPLACE
For each LLM prompted to solve REPLACE, we assess which types of abstractions of state s, are preserved in its representation h4 during decoding after the collection of {(s,, ht)) 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 set of abstract predicates crucial for devising the optimal plan for task completion:
nea rby(o): the object is located within the container nearby to the agent. Similarly, we have dí stant (1, o), where is the relative distance to the agent's is current location, e.g.. distant( (oentt(L)) (-1, o) being True indicates that the container one step to the left contains o.

---

subgoal(1,,) 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 I the agent's current location.
the relative direction of the next object to be manipulated, where could left or right to the agent, in a nearby container, or with the agent.
nextObjDirrett
(o): the object 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 next Obj and next Ob jobijDirttt Figure demonstrates how to derive abstract predicates and actions based on the raw predicates. Figure 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 d, agent both the current and target situations, and 9. and boxName.
World-irrelevant abstraction To preserve the world dynamics, we have to track all raw predicates of the current situation, excluding boxName, actions do not affect it. In addition, we include all predicates in target situation g to preserve the reward function. Therefore, Sb_ includes the predicates store, he and agent Loc in both and g.
Q* gr-irrelevant abstraction To construct Sgos we include agent held, and nearby, to distinguish between the effects of legal and illegal execution of actions. Also, we include subgoal, where subgoa 1 j = 1) differentiates the effect between optimal and non-optimal actions, and subgoa 1 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.
π -iireeeevant abstraction Sa, is factorized by predicates next Ob joiDirees Diirect and next nextOb jDireect specifies the action type (move manipulate objects) and moving direction (left or right) at the next step. nextObj identifies the particular for grab or put operations.
Appendix G provides more formal proof.
We probe the predicates in LLM representations by training two-layer neural models, following [28]. For single-variable predicates (e.g., nea (by( we train an individual probe for each. For two-variable predicates where one variable denotes a positional index (e.g., boxName b)), we train 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,, p(o)], where cp(o) is the embedding of Following the conditional probing principle [15], we treat the embedding method for each predicate as a tlper-parameter selecting the method with the largest performance margin between using [he: soo)] and using only (()) 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,, Appendix H provides more details about the formulation of probing tasks 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 tp for predicate p by (0,,,,.,,.0000 100%, where 3, is the score of a RANDOM baseline. predicts outcomes by proportional sampling from label candidates of p based on their frequencies. We refer to this normalized score as the recovery rate.
1-3p
Experiments
6.1 Target Models for Probing Experiments
We experiment with two groups of LLMs to investigate how fine-tuning and pre-
GRIPPER
COok
Models
%llegal
WSuce
3:0ptim
$T.egal
5pppim
training impact the world abstractions coded in their representations, respectively. The details follow below.
Llama2.
62.41
3.07
0.19
35.41
1.70
0.36
Mistralenn
44.41
Llama215t
96.54
88,30
84.02
95.71
85.11
69.17
Mistraler
92.15
87.36
95.04
Table I: Planning performance of Llama2- and Mis-
Effeec of fine-tuning We adapt Llama2-
tral on GRIPPER and COOk.
1b/3bb [35], Mistral-7b [20], and LLama3-
7

---

8b [1] on both datasets, using in in-context learning (ICL) and supervised fine-tuning (SFT) with LoRA adapters [19]. Appendix provides more details.
Effect of pre-training Since the LLMss above perform near-random without fine-tuning, we also experiment with more recent state-of-the-art LLM,, Phi3-17b [2], which achieves much higher performance than other pre-trained LLMs on GRIPPER". For comparative analysis, we also include Phi3-3.8b, Pythia-70m [6] and fine-tuned Phi3-17b.
In addition, we train a 6-layer decoder-only Transformer from scratch, probing for abstractions from its internal representations baseline.
20-
Age Figure 5: Average recovery rate of each world abstraction across different LLMs GRIPPER.
1:00
92
Agestee
Karw
World
Morty
(a) From Llama2 and Mistral.
(b) From Pythia and Phi3.
Figure 6: Recovery rate of all predicates from different LLMs on GRIPPER. Predicates are grouped according to the coarsest abstractions they pertain to. Better viewed in color. The color indicates all abstraction levels the predicates are associated with.
Before the probing experiments, we first verify that the fine-tuned LLMs and Phi3-17bicr. can achieve reasonable performance on REPLACE, and thus, we can draw informative conclusions. We adopt three evaluation metries 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-133 and Mistral results are reported in Table 1, with comprehensive results for all LLMS provided in Appendix As shown in Table I, both LLMs initially fail 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 the constraints. Next, we conduct probing experiments on each variant of the LLMs on the same datasets.
WWe do not conduct probing experiment with the pre-trained Phi3-17 on COOK as it does not outperform the random baseline significantly.
8

---

World
(b) Recovery rate of all predicates from Ilama2- 1-mba211bbba and Mistral.
(a) Average recovery rate of each world abstraction.
Figure 7: Recovery rate of predicates within different world abstractions across different LLMs on Cook.
6.2 Probing Results and Analysis
We measure probes for all the LLM variants adapted to REPLACE (ICL or SFT) as well as the train-from-seae rrom-scratch Transformer. Figure and Figure 7(a) report the average recovery rate (RR) of predicates within each type of world abstractions from different LLMs. Figure 6 and Figure 7(b) take a closer look the RR of each predicate from seven LLM variants to analyze the effect of fine-tuning and pre-training. The complete results are reported in Appendix J. Next, we present four key findings from the results.
Finding 1. Reasonably performing LLMs tend to maintain goal-oriented world abstractions rather than a more general one during decoding.
All LLMSsFT and Phi3-17bicı, which outperform the random baseline by substantial margin, primarily preserve Q"- and T -rrreeelevant abstractions. Comparing LLMSaar and LLMs in Figure 5 and Figure 7(a), Q and π- irreeeelevaant abstractions are probed from LLMssar with drastically higher RR than
and world-irreleval abstraction. Focusing the RR of individual predicates, Figure 6(a) and Figure 7(b) reveal that the predicates for π iirelcvaan abstraction, nextObj and nextOb jDitobjDiiee rect have been mostly recovered from LLMS, This result aligns with the high task success rate achieved by these LLMs. Similarly, the predicates for Q" -i-irrelevaan abstraction effectively probed from the LLMSsFT. 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 's long-term effect with respect to altering the world situation to match the target, facilitating efficient planning. The successful recovery of nearby, he held, and agentLoc indicates that L.M representations capture the world constraints to inform legal actions, as evidenced by the near-perfect legal rate of LLMssFr.
Nevertheless, Figure 5 shows that although the RR margin between goal-oriented abstractions and more general ones is smaller for Phi3-17bicz it remains significant. Figure 6(b) explains this by revealing that the RR for bextObj. nextOB jDi rect, and subgoal from Phi3-17bicu is much lower than that of LLMSsFT, which is unsurprising Phi3-17bicı. achieve much lower success rate on REPLACE than LLMS
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 accurately recovered. The low RR for boxName suggests that LLMs discard world details that not pertinent to the task's s completion. Similarly, the minimal RR on store,, 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 predicate originally irrelevant to the goal becomes pertinent to goal- oriented abstraction. Our results show that it is accordingly encoded the LLM representations. Under the other setting, the L.Mss 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 fine- tuning. Specifically, the performance gap in terms of predicates next jDi rect, agent 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 implics that when fine-tuned with teacher forcing, L.LMs evolve develop 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- fot---rrmmm-aacchh Transformer and LLMSsFr lies in their retention of Q" and --" -irellevannt abstractions. Nonetheless, as LLMs increase in scale and capability (Pythia<Phii3 3.8b<Phi3-177 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 stors store and he 1 dg. essential for world-irrelevant abstraction, are almost identical across all pre-trate 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-irrelever abstractions are mostly suppressed during decoding.
Motivated by our initial two findings, explore whether and how boxName and store are en- coded within the contextualized representation SPT by [J To do so, we use the contex-
Predicates
Abstraction Type
hy Cra(e)
Cixt(c)
boxName
Raw
100.00
100.00
store,
Raw, World-irrelevant
93.32
95.71
store,
Raw, World-irreleaan
93.01
95.18
nearby.
Q" qr-irrelevant
92.11
42.46
tualized embedding Cixt(-) of the label candi- date, comparing the use of only the embedding Cixt(e) against a concatenation [h,; Ctxt(e)], in- spired 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- near-perfect performance while incorporating h2 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 ht boosts the performance vastly, suggesting that LLMs discard the information of store and integrate it with agentLoc to derive nearby.
Table 2: Recovery rate with different encoding methods of label candidates.
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 LLMss 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 raw world states might lead to superficial conclusion that the LLM only partially captures the world. In contrast, our framework offers a more thorough and systematic interpretation, demonstrating that well well-performing LLMs maintain 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.
7 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 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.
10

---

Acknowledgements
This work is supported by the Mitacs Accelerate Program (Project ID: IT27067) and partially enabled by Mila's s computing resources (mila.quebee), We thank Raquel Aoki from Borealis AI for discussions, and Michael Rizvi, Maxime Darrin, and Saba Ahmadi from Mila for their manuscript feedback during paper swap. Finally, we appreciate the anonymous reviewers for their comments on clarity and suggestions for additional experiments.
References
[1] Introducing meta Ilama 3: The most capable openly available IIm to date. https://aa meta. www.bommoloooooo meta- 11ama 3/. Accessed: 2024-04-18.
[2] M. Abdin, A. Jacobs. A. A. Awan, J. Aneja. Awadallah, Awadalla, N. Bach, A. Bahree,
A. Bakhtiari, H. Behl, et al. Phi-3 technical report: A highly capable language model locally on your phone. arXiv preprint arXiv:2404. /4219, 2024.
[3] D. Abel. A theory of abstraction in reinforcement learning. arXiv preprint urXiv:2203.00397, 2022.
[4] D. Abel, D. Arumugam, K. Asadi, Y. Jinnai, M. L. Littman, and L. L. Wong. State abstraction compression in apprenticeship learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pages 3134 1134-114,,2 2019.
[5] M. Beloucif and C. Biemann. Probing pre-trained language models for semantic attributes and their values. In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 2554-2559, 2021.
[6] Biderman, H. Schoelkopf, Q. G. Anthony, H. Bradley, O'Brien, Hallahan, M. A. Khan, S. Purohit, U. S. Prashanth, E. Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. International Conference on Machine Learning, pages 2397-2430. PMLR, 2023.
[7] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances neural information processing systems, 19:1,
[8] M.-A. Côté, A. Kádár, X. Yuan, B. Kybartas, T. Barnes, E. Fine, J. Moore, M. Hausknecht, L. El Asri, M. Adada, al. Textworld: A learning environment for text-based games. In Computer Games: 7th Workshop, CGW 2018, Held in Conjunction with the 27th International Conference on Artificial Intelligence, IJCAI 2018, Stockholm, Sweden, July 13, 2018, Revised Selected Papers 7, pages 41-75. Springer. 2019.
[9] K. J. W. Craik. The nature of explanation, volume 445. CUP Archive, 1967.
[10] T. Dean and R. Givan, Model minimization markov decision processes. In AAAI/AAL, pages 106-111, 111, 1997.
[11) N. Ferns, P. Panangaden, and D. Precup. Metrics for finite markov decision processes. In Proceedings of the 20th conference on Uncertainty artificial intelligence, pages 162-169, 2004.
[12] J. W. Forrester. Counterintuitive behavior of social systems. Theory and decision, 2(2):109-140, 1971.
[13] D. Ha and J. Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems, 31, 2018.
[14] D. Hafner, T. P. Lillicrap, M. Norouzi, and Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2020.
[15] J. Hewitt, K. Ethayarajh, P. Liang, and C. D. Manning. Conditional probing: measuring usable information beyond baseline. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 1626-1639, 2021.

---

[16] J. Hewitt and P. Liang. Designing and interpreting probes with control tasks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNIP-IJCNIPP) pages 2733-2743, 2019.
[17] Hewitt and C. D. Manning. A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers),
pages 4129-44 4138, 2019.
[18] M. K. Ho, D. Abel, C. G. Correa, M. L. Littman, J. D. Cohen, and T. L. Griffiths. People construct simplified mental representations to plan. Nature, 606(7912):129-133 2022.
[19] E. J. Hu, P. Wallis, Z. Allen-Zhu, Li, S. Wang, L. Wang, W. Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2021.
[20] A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, Chaplot, d. I. Casas, Bressand, G. Lengyel, G. Lample, L. Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06625, 2023.
[21] C. Jin and M. Rinard. Emergent representations of program semanties in language models trained on programs. Forty-first International Conference on Machine Leurning.
[22] D. H. Jonassen and P. Henning. Mental models: Knowledge in the head and knowledge in the world. Educational technology, pages 37-42, 1999.
[23] N. K. Jong and P. Stone. State abstraction discovery from irrelevant state variables. In IJCAI, volume 8, pages 752-757, 2005.
[24] N. Kim and S. Schuster. Entity tracking in language models. arXiv preprint arXiv:2305.02363, 2023.
[25] Y. LeCun, A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. 2022..
[26] G. Leech, P. Rayson, et al. Word frequencies in written and spoken English: Based on the British National Corpus. Routledge, 2014.
[27] B. Z. Li, M. Nye, and J. Andreas. Implicit representations of meaning in neural language models. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers), pages 1813-1827, 2021.
[28] K. Li, A. K. Hopkins, D. Bau, F. Viégas. H. Pfister, and M. Wattenberg. Emergent world representations: Exploring a sequence model trained on a synthetic task. The Eleventh International Conference on Learning Representutions, 2022.
[29] L. Li, J. Walsh, and M. L. Littman. Towards a unified theory of state abstraction for mdps. In AI&M, 2006.
[30] X. I.. Li and J. Eisner. Specializing word embeddings (for parsing) by information bottleneck. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and 9th International Joint Conference on Natural Language Processing, 2019.
[31] D. M. McDermott. The 1998 ai planning systems competition. Al magazine, 21(3):35-35, 2000.
[32] B. Min, H. Ross, E. Sulem, A. P. B. Veyseh, T. H. Nguyen, O. Sainz, E. Agirre, I. Heintz, and D. Roth. Recent advances in natural language processing via large pre-trained language models: A survey, ACM Computing Surveys, 56(2):1-40) 2023.
[33] M. Minsky, A framework for representing knowledge, 1974.
[34] T. Pimentel, J. Valvoda, R. H. Maudslay, R. Zmigrod, A. Williams, and R. Cotterell. Information- theoretic probing for linguistic structure. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4609-4622, 2020.
12

---

[35] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, Bhargava, S. Bhosale, al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
[36] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
[37] T. Yun, Z. Zeng, K. Handa, A. Thapliyal, B. Pang, E. Pavlick, and C. Sun. Emergence of abstract state representations in embodied sequence modeling. In Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pages 12190-12205, 2023.
13

---

A Limitations
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 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 important note that 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 synthe- size 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 sim- "latorr-vitas for tracking state updates triggered by predicted actions in probing experiments are lacking. We leave the development of such datasets for future work.
Limitations of our experiments In this work, we have run experiments under our proposed framework with two datasets and four state-of-the-arr open-sourced L.MMs. 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 LLM representations.
B Broader Impact
Our work introduces 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 discrepancies may lead to undue pessimism about I.LMs' 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.
arise between evidence and conclusions. Such
C Task Instructions and Input Examples in GRIPPER
Table 9 shows examples of each dataset variant in GRIPPER. Table 10 provides the full task instruction included the prompt.
D Details and Examples of COOK Dataset
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 evaluating RL with respect to react to the feedback from the environment, where the updated state information included in the feedback, it makes the task of world modeling somehow trivial at 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
14

---

ground-truth actions. Nonetheless, there are two main differences between COOK and GRIPPER. First, the containers (rooms) COOK are placed south, north, east west, and therefore have larger action space. In contrast, the agent move either left or git GRIPPER. Second, COOK uses different set of textual templates adopted from TextWorld, which are listed in Table 3. Table 11 shows the instruction for COOK and Table 12 provides an input example. COOK has the same size with GRIPPER.
grid world, where the agent can move
GRIPPER
COOk
store
container contains object, ner holds sojeet, container has ab object, a ine occupied by object container emply
containe contains ol jbjeet, In conta' you see object, pont a iine has ab jajet in conta ner, you pont iner has nothing
ob ject,
boxName
A sequence of containers are ordered from left to right follows: Rox Rox B.
Room in the northwest connects east Room
and south
Room. A links further easi to Room and south
further
Room
Room. connects south Rocs.
Room leads east to Raom B.
Room connects east
Grab the object from nearby container
Take ab abject from the
Put objject in the nearby container.
drep ob jeject in the room.
move
dí raction
walk to 6 rect inon
Examples of o (objects)
book, car, paper, letter, game...
onion, tomato, potato, mushroom, carrot...
Table 3: Textual templates used in GRIPPER and COOK.
E Detailed Introduction of Abstract Predicates
We derive following abstract predicates using raw predicates, namely boxName, store, agentLoc and he held:
nearby(o): the object o is located within the container nearby to the agent. It is derived by incorporating Loe with store. Similarly, distant o), where I denotes the relative distance to the agent's current location, can be derived. That means 1 e [- 1-11] 1). where I. denotes the number of containers. For instance, di stant(-1,o) being True indicates that the container one step to the left of the agent contains o.
subgoal(L.j): the agent needs visit a series of containers to manipulate objects, and the j-th container to be visited is at the relative position I to the agent's s current location. Furthermore, for subgoal(j, 1), the label for is in the set {-I. + 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. 1, as it allows for deterministic derivation of nextOb jolijDiirttte Also, stop implies that there are fewer than subgoals. subgoal derived by comparing distant, nearby and hele le across the current and target situation.
1}
nextOb rect step, where v could be 1 eft ri right the agent, in nearby container, with the agent. This is derived from the value of first subgoal, i.c., subgoal(-,11
the relative direction of the object to be grabbed or put at the next
nextObi(o)) the object o be manipulated at the next step, if applicable, This is derived by comparing nearby and held between and g.
15

---

Adapting LLMs to REPLACE
We use Llama2-7b', Llama2-13b', Mistral-7b', Llama3-8b*, Phi3-17b', Phi3-3.8b and Pythia- 70m as the base models for experiments. Additionally, 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 LLMSs with the learning rate of le - 4, the batch size of 8, the scaling factor of a 64, the dropout rate of 0.05, the rank of 32 for epoch. To avoid significantly altering the I.L.'s ' representational traits, we use only subset of 10k training samples for SFT. We use 8-bit quantization for decoding, During decoding, parse the generated textual description Y of subsequent actions to its symbolic form with 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 tinnitnning enforce well L.MMs to generate the same format as the data in Gr pper and Cook.
G Proof of World State Abstraction for REPLACE
Recall that the raw state S in REPLACE is factorized by boxName, store, agentloc and held. Lemma G.1. Se., is world-irreleaan abstraction of S is factorize by store, agentLoc and held, d, excluding boxName.
Proof. Assume (81) = ow (82). This implies that Va A, To (s'(s,,,) Tow (5 a), given that boxName in S always remains unchanged after the execution of any action, we have:
Va' E Spu,
Σ
T(b1b),,))
s x' (x')
Σ s'Ew"((x))
T T(l (s'); by b= u.)
Σ
Tow ( (s))
s'Enn (x')
(1)
Σ
Tow (o (sw (s))))))
s s'E (x')
Σ
b,, bx
b.
s'E (z)
Σ
(s)
S'E (z')
where by denotes the assignment over boxName s'. Similarly, it follows that (81) = (x2)
R(s1.a)= R(s2, a). since R(.,.) is invariant with respect to changes in boxName within a state
(We accept their Meta li-
cense before requesting the model.)
"https: / htth//ugggngggaccccc co 2
(We accept their Meta li-
before requesting the model.)
(licensed by
Apache-2.0)
'https: http:///ia..........aoaaaaaaaaaa
(We accept their
Meta license before requesting the model.)
/https://huggingffae co/m /mi http://rrrrott///aaaameaaao--ae-
(licensed by MIT)
#https://hugginggfaa co/m http:/cct//ihii--miniiiii
(licensed by MIT)
"nttps::/huggingtaan co httt:s//wwggggnoooooooottttttat/ (licensed by Apache-2.0)
16

---

Lemma G.2. Sp. is an π -""irrelevvant abstraction of S if Sø, is factorized by next Obj and nextObjpirect.
Proof. To demonstrate that Sø, a a iireelevnn abstraction, we show that the optimal action a" in any state can be uniquely identified using nextObj and nextOb This is formalized as follows:
grab(o*),
if nOD container
put(o*), a"
nOD = with-agent
move(left),
if nOD = left
move(right),
if nOD = right
where nOD denotes the value returned by nextob rect and O* denotes the object identified
by nextObj. This ensures that for any two states 81 and 82 where (5())) = (8())) follows that x (a(s)) = TT" (als2).
L
Lemma G.3. Sag is an Q" -irieeeeeaant abstraction of S if Sãg is factorized by egentLoc, held.. 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 = [gi,..., gi. where for each i, subgoal(( (1, 8i) True. nextobjDirect is inferred from gi- Consider the optimal action a" derived from nextObj and rextObjDirect, as discussed in Theorem G.2.
The value of Q" a²) is given by:
8:+1
Bil -
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. follows that Q" (s, a) for all actions can measured as below:
Q²(s, (*( ('(aaa))-) - - if $ ((s) clse
a = a"
(3)
Q*(s,a")-)
where C()) is the constraint function defined in Section 4, dependent on nearby.. held, and agentLoc.
Consequently, (s a) can be precisely recovered conditioned on the predicates agentLoc, heldw, nearby, subgoal, and next This implies that Q" (éq((1),, Q (dQ a) (*( = (s2 a).
H
Details about Design of Probing Tasks and Training of Probing Models
We detail the formulation and evaluation of the probing for each predicate as follows:
boxName: Given an LLM-built representation h, and a set of candidate container names b, as well the embedding (bb) for each b E b, the probing model M₁ for i-position learns to select the ground- -trut container name b* by maximizing the probability of boxName 5*) being True:
exp (M) ((b"))
To = llog EbeD exp (M3 ((b)))
During testing, we include the container names in the same data instance as candidates and measure the F1-score using a macro average.
Loe and next Ob jDir: Both these predicates are formulated as multi-class classification tasks, Given h. the probing model learns to maximize the ground-truth label, i.e. i" for agentLoc(i) and 6" for (ex)
exp (PMo (i";ha)
To - log E,exp exp (PMa (i: h4))
17

---

We adopt the Fl-score as the evaluation metric.
store, he d and nearby: These predicates may have multiple labels. For instance, 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 Mo for store(i,) learns to perform pairwise binary classification for st store o), where o O O:
To = 0(, log PMù (bhe; (())))
+ (1 - (i, log PMg, (Ie ((o)))
where q(i, o) = [|storeii, o)]l During testing, we combine all the predictions for st store 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.
next We formulate and train the model in the same way as boxName(i,, -). Again, we adopt the Fi-score 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 nextobjiii rect. and agentLoc. We use the F1-score for evaluation.
When compare probing performance across different predicates, normalize all the metric score to the recovery rate, as described in Section
We create probing data with GRIPPER and COOK 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 1 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 (p() and select the one with the largest performance margin between using [he, (e)) and using only sp(e) the validation set. As result, we use word embedding for boxName, store, held, next jbj, and contextualized embedding for near aby. We explore multiple variants of m for extracting ht": 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+, p(e)], we uses a trick that transforms the it to [he + ((e); h4 <pe); h4 (ee) 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 le 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 single Nvidia A100 GPU, and each of them was completed within 24 hours.
I Complete Planning Experiment Results
Table 4 presents the planning performance of both variants of all LLMs for REPLACE.
J Complete Probing Experiment Results
Table 5 and Table 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.
18

---

GRIPPER
CooK
Models
<TLegal
SSucc
Sgee
SSSccc
Uptim
Llama2?
34,89
3.83
0.00
30.29
2.27
0.76
Llama2g2
3.07
Mistrlla
41.12
2.12
0.65
44.41
1.95
0.22
Lluma
62.19
10.36
2.17
39.11
2.26
0.09
Llama22n
95.40
87.20
82.18
89.13
64.86
35.59
Llama2iar
Mistraler
97.07
92.15
87.36
95 04
85.91
61.47
LLama3ar
96.25
89.14
74.32
95.36
86.61
61.21
52.88
18.78
8.33
Table 4: Planning performance of LLMs on GRIPPER and COOK.
05
Hold
$3.13
1726
0157
8.10
0185
1334
4.53
653
6.30
025
Ruk. Audiieskeo Saviiiua
11.35
28.5
21-11
26.0
2106
11.00
31400
2148
28
20-12
20.00
18.30
23.8
0465
2708
1.00
13:09
20.13
12.05
3001
11.04
101
01:11
Table 5: Probing performance (recovery rate) of all predicates on GRIPPER.
Pritese melicates
Alstracton type
Lianz
Limrs-::
Larren
Lisest::
Kiokeler
Listt-gn
Listtee
Litefonner
12.58
3.50
11,02
Swild
Raw. Wodda-indiesatt
3.72
0.00
3.45
0.00
3.26
1.00
3.45
3,00
3.6
0.55
Raw, Wattiindtesatt deeeeaa
39.30
15005
35.13
10.33
30.00
34.72
39.23
14.54
10:33
0.35
Seld,
Raw, Wodll-neeevaat, Q" cincevet
42:09
40.16
37.33
34.34
15.47
41.38
40:90
40.77 T3
36.55
8:21
g" inndeaar
next:2a
Q" undezani -iimloosai
77.48
0.00
94.51
0.00
50.30
0.00
76.59
0.00
0.56
0226
rex rect
e" Inmboeee
87.52
6.77
88.39
5.
91.87
18.00
90.31
12.23
57.40
0.6
Table 6: Probing performance (recovery rate) of all predicates on COOK.
K
Probing Experiments with COLORGRIPPER
In the original setting of GRIPPER, the color 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, GRIPPE,, 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 colorg. the color of each object in the desired situations, and an action paint c), which paints the object o in a nearby container to color and e E fred, blue, green, yellow). Also, incorporate the new color must be different from the original one. As such, there is a new abstract state hearNotColor (o, c). which classifies that the color of nearby object is NOT and it pertains to Q" -irrelevant abstraction. We fine-tune two Llama2-13b models on GRIPPER (Llama2mng) and COLORGRRPPER (Lllam22olarr) respectively, and probe the nea rNotCol and color from both fine-tuned L.MMs. The results Table show that the probes fail to recover co color, which do not belongs goal-oriented abstraction under both settings, from Llama2 and Llama2 However, nearNot Colootoolo can be probed from Llama2cokrr with reasonable recovery rate. Again, this confirm our finding that LLM tends maintain goal-oriented world abstraction during decoding.
constraint that
L Probing from LLMs Fine-tuned with Suboptimal Action Sequences
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 ground- truth 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 With the new dataset, the RR of goal-oriented abstractions decreases, which is expected the model imitates a sub-optimal policy. However, the margin between goal-oriented and general abstractions 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.
19

---

Predicates
Absraction Type
Ilama
Llamaa.okrr
color
Raw
0.00
1.92
nearNotColoz Q or-irrelevant
0.00
37.70
Table 7: Comparative experiments of probing color information under GRIPPER and COLORGRIPPER.
Dataset
Raw
World-imelevant
-iiinelevant
a --iirclevant
Original
28.5
30.6
73.0
Suhoptimal
25.5
27.4
63.3
72.1
Table Comparative probing experiments with Llama3-8b fine-tuned on original and suboptimal actions on GRIPPER.
Base
A sequence containers ordered from left right follows: Crate A, Bucket A, Crate B, Basket A, Box A. You are at the Basket A. Initial situation: Crate A contains painting. shoe. Bucket A contains a block, a disk. Crute contains glass. Basket onything.
rose. contains paper, camera. You don't i hold
Desired situation: Crute A contains painting, a shoe. Bucket A contains a block, disk. Crate contains glass. Basket A empty. Box contains camera. You hold rose, paper.
Lexical variants
sequence of contuiners are ordered from left right as follows: Box A, Basket Box B, Box Basket B. You
the Box
Initial situation: Box A holds small green petroleum, big yellow gastropod. Basket A has nothing. Box B holds a small yellow sabre, a small red tricycle. Box has tiny green eggnog, big yellow nametag, a small blue lantern, a big yellow biscotti. Basket B holds a tiny blue anaconda. You don't hold anything.
Desired simution: Box A is occupied by petroleum, gastropod. Basket A has nothing. Box holds sabre. Box C has a eggnog, nametag, lantern, a biscotti. Basket B has a anaconda, a tricycle. hold anything.
Partial ru(g)
sequence of containers are ordered from left to right as follows: Box Box B, Box C, Basket Box D. You
Basket A.
Initial situation: Box A contains jacket. Box B contains a boat. paper, a ticket. Box C contains bag. Basket A contains a engine. Box D contains a phone, a coffee, a ticket.
Desired situttion: Box contains a jacket.
contains boat, a paper, ticket. Box contains bag. Basket contains engine. Box D contains phone.
Partial (r.(w)s...g)
ordered from left right as follows: Bucket Basket A, Bucket Crate A, Bucket C. You are at the Crate A.
sequence of containers
Initial situation: Bucket A contains bomb, a ring, medicine. Basket A contains boat, a map, bottle, a fish. Bucket B contains a dish, a fish. Crate A contains a plane, stone. Bucket contains radio, book, map, fish.
Desired situation: Bucket A contains bomb, ring, Basket A contains boat, bottle, fish, medicine. Bucket contains a dish, a fish. Crate A contains a plane, a stone. Bucket C contuins radio, a book a map, a fish.
Operations applied: You Move left for 2 steps.
Table 9: Examples of GRIPPER variants.
20

---

You are an intelligent agent a fictional environment with sequence of containers arranged in straight line.
The environment involves movement, interuction with objects, and strategic planning to reuch a purticular state.
Structure of the Environment: There are containers arranged in a line from left to right. You are at one of them in the beginning.
Goals Strategies:
You start with an initial configuration of objects in containers. Your job is to reach the desired state, differrnt configuration, through sequence of operations. If there is any object not mentioned in the desired state, that implies you should hold them in the end.
The legal operations include moving and manipulation grab or put) of objects
1. Moving: You can only walk left or right. Each time you take a step, you 'll be right next to the container on your left or right side.
2. manipulation: Inside some containers, there are objects. To grab an object or put one inside, you need be at that container. If you see an object in container far away, you can't I grab it. You need to walk over to that container first. Constraints: For each step, you can choose only one operation: either move or manipulate one object. You can hold multiple objects at once.
Table 10: The general task instruction in GRIPPER.
In Tex: You are an intelligent agent tasked with arranging food ingredients in grid of rooms to match specific desired configuration.
Structure of the Tex:World: TextWorld consists of a grid-based based map with multiple rooms. You start in one of these rooms. Goals & Strategies: You start with an initial configuration of ingredients in rooms.
Your job is to reach the desired state, different configuration, through sequence of operations. If there any ingredient not mentioned in the desired state, that implies you should hold it in the end. If you think the desired state is reached, terminate the process.
The legal operations include moving and manipulation( take or drop) of ingredients:
1. Moving: You can go north, south, west east to enter adjoining rooms.
manipulation: Interuct with ingredients inside the rooms by either taking or dropping them. You must be physically present in room interact with its contents. Constraints: For each step. you can choose only one operation: either move or manipulate one ingredient. You can hold multiple ingredients at once.
Table 11: The general task instruction in COOK.
You open the map of TextWorld. bedroom A in the northwest connects east to garden A and south to pantry A. garden A links further east backyard A and south to garden B. backyard A connects south to bedroom B. pantry A leads east to garden B. garden B connects east to bedroom B. You are in backyard A.
Initial state: In bedroom A, you can see roasted sliced pecan, fried chopped mustard seed. garden A has grilled chopped nectarine, grilled chopped mango, raw chopped beef in it. backyard A has fried diced pork in it. pantry A contains raw sliced polenta, grilled diced rice, fried uncut tempeh. garden B has raw uncut mussels, grilled diced ham, roasted diced squid in it. bedroom B has grilled chopped nectarine, fried chopped mustard seed, roasted chopped walnut in it.
Desired state: bedroom A has nothing. In garden A. you can find nectarine, mango, beef, mustard seed. backyard A contains nothing. pantry A contains polenta, rice, tempeh. garden B contains mussels, ham, squid. In bedroom B, you can see nectarine, mustard seed, walnut. Operations applied: You take pork from the room. You go west.
Table 12: An input example in COOK.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We list the main contributions at the end of Section
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discuss thoroughly the limitations of our framework, task, experiments in Appendix A.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We have properly cited state abstraction theory, on which our work relies. Due to the limit of space, we provide the formal proof of the derivation of abstractions for our REPLACE task Appendix G.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide the implementation details our probing framework, including the task formulation, probing model design, hyper-parameter search and setting in Appendix F and Appendix H. To further ensure reproducibility, the reported results are averaged across multiple runs, and we have found the variance small (detailed in Appendix J). We will release the synthesized dataset, the fine-tuned LLMs checkpoints and all of the code for experiments.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We will release the synthesized dataset, the code for generate datasets, fine- tunining LLMs, and probing. The instruction our experiments are detailed Appendix and Appendix H.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Data splits is introduced in Section Section 4.2. The setting and impact hyper-parameters as well as other training details are elaborated in Appendix and Appendix H.
Justification:
7. Experiment Statistical Significance
22

---

Question: Does the paper report enror bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the standard deviation for the main probing experiments in Table and Table 5. As we discuss Appendix J, the variance across different runs is very small and it cannot contribute meaningfully to the plot. Therefore, the error bar is omitted from the plot.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide details of compute resources in Appendix H.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: httpa//ppp...ccpll/aa//ccccccooo
Answer:[Yee
Justification: We conform the NeurIPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Due the limit of space and do any potential negative societal impact of our work, we discuss the positive societal impacts of our work in Appendix B.
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We do not see any significant risk that our dataset and probing models could misused.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We conduct probing experiments on top of existing LLMs released by other researchers and organizations. We have cited the resource, such the paper and website, introducing the models and include the license information in Appendix F.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We will release all the code and datasets, which are licensed under CC BY 4.0.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
23

---

Justification: Our work does not involve crowdsourcing nor research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does involve crowdsourcing research with human subjects.
24


Reviews:
[{"review_id": "M5gg2cNXAh", "reviewer": "Reviewer_ZobB", "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.\n\n**Post-Rebuttal/Discussion**: I find the author response fair and am raising my score to 7.", "questions": "Please take a look at the Weaknesses section above.", "limitations": "The authors provide a fair list of limitations in Appendix A.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Given the extensive use of transformer models these days, establishing sound findings in this regard (whether LLMs can provide representations useful for recovering world dynamics information) can be an important problem.\n- The proposed task, RePlace, is carefully designed with different levels of state abstractions in mind. It allows the authors to conduct the empirical analysis that shows the differences in the recovery rates for different state abstractions.", "weaknesses": "- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and positive feedback. We address all of your concerns below.\n\n### **Weakness 1: finding not entirely surprising**\n\n>(**Unsurprising findings**) The finding that the fine-tuned language models do not preserve the features needed for recovering the world dynamics may not be entirely surprising…\n\n* **It's indeed surprising when viewed in the context of existing work**. We provide detailed elaboration in the [Global Rebuttal](https://openreview.net/forum?id=lzfzjYuWgY&noteId=i9JnaUoNpz), with other surprising findings and broader significance of our framework. We sincerely invite you to review it.\n\n> (**It's unsurprising because**) fine-tuning enforces the bias … and much of irrelevant aspects often end up being ignored\n\n* In fact, the **world dynamics is useful for the task**, albeit in a less goal-oriented way. As illustrated in Figure 3, the optimal plan is based on predicting future states. Generally speaking, planning is searching for the shortest action sequences induced by which the predicted future state matches the goal.\n* Nevertheless, **Although $\\pi^\\*$ abstraction is sufficient for predicting the optimal actions, we still have probed $Q^\\*$ abstraction with a high recovery rate** (Figure 5&6, discussed in L324-330). It reveals that LLMs-sft preserves the impact of each possible action, despite being fine-tuned only on next-token prediction. To the best of our knowledge, this interesting finding is novel and can inspire future research to conduct more in-depth analysis.\n\n### **Weakness 2: Probing with pre-trained LLMs**\n\n> (**Experiments with pre-trained LLMs**) more interesting findings can be derived with pre-trained LLMs … (but) their performance is quite low in the first place, which makes it difficult to examine this hypothesis.\n\n* You are correct that we didn't conduct probing experiments using pre-trained LLMs, given their near-random performance without fine-tuning. Also, it's a common practice [1, 2, 3, 4] that probing from Transformers/LLMs that are trained or fine-tuned on next-token prediction to see if implicit world models emerge, as this is a common interest in the field. \n* To completely address your concern, we conduct new probing experiments on top of Phi3-17b, a brand-new SOTA LLM that achieves 19.94% success rate on GRIPPER without any fine-tuning. From its representations, goal-oriented abstractions can be probed with a significantly higher recovery rate than the world-irrelevant abstraction. We further include two smaller and weaker pre-trained LMs for comparison. The results suggest that more powerful pre-training leads to a higher tendency to prioritize goal-oriented abstractions over a more general one. Please refer to the [Global Rebuttal](https://openreview.net/forum?id=lzfzjYuWgY&noteId=i9JnaUoNpz) for detailed results and in-depth analysis. \n\n[1] Li, Kenneth, et al. \"Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task.\" ICLR, 2023.\n\n[2] Li, Belinda Z., Maxwell Nye, and Jacob Andreas. \"Implicit Representations of Meaning in Neural Language Models.\" ACL, 2021.\n\n[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).\n\n[4] Kim, Najoung, and Sebastian Schuster. \"Entity Tracking in Language Models.\" ACL, 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mezp74nbku", "reviewer": "Reviewer_4mUW", "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)", "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": "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.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper is really well-motivated and written. The problem is a varying interesting one and the observation is timely for the field. \n2. It's quite novel to break down abstractions through the lens of q value function and policy to capture aspects that are relevant for planning. \n3. The design of the abstraction together with the RePlace experiment is rigorous and well-thought-out.\n4. Figures are super helpful to ground analysis in concrete examples.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and review. We appreciate the opportunity to clarify two critical aspects that may have been misinterpreted: our claims/hypothesis and the concept of goal-oriented abstractions. \n\n### **Weakness**\n\n> 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 (from **summary**) …  the claim that LLM will learn different abstractions given different tasks (from **weakness**) …\n\n**Indeed, these are NOT our hypotheses and claims, and NONE of our contributions (listed in L74-83) & findings (Sec 6.2) is about LLMs learning different abstractions for different tasks.**\n\nWe clarify our main claims & proofs below.\n\n* The disagreement among existing works is likely due to **(1)** they probe **only the raw world state** on different tasks (L30-31), and **(2)** the raw world state might equal to or overlap with varying levels of abstraction depending on the task (L34-L38).  It's, therefore, plausible that LLMs consistently preserve certain types of abstractions, but previous studies fail to capture this pattern. \n\n* **More critically**, probing WITHOUT the notion of world abstractions leads to biased evaluations, which may underestimate LLMs' world modeling when there is no overlap between the raw state and any abstraction or overestimate when the raw state aligns with goal-oriented abstractions (illustrated in L49-54, L164-170). Our framework directly addresses this limitation by probing different levels of abstractions$\\color{red}{^1}$.\n\n* Within our framework, we primarily investigate **which types of world abstractions are encoded in LLM representations** (mentioned throughout the paper,  L71-73, 77-79, 145, 243-245). Sec 6.2 lists all findings. Moreover, **the results prove that probing only the raw state causes biased evaluations and unnecessary disagreement** (emphasized in L374-379, also elaborated in the Global Rebuttal).\n \nDue to the **lack of faithful testbeds** (explained in L162-170)$\\color{red}{^2}$, we CANNOT further test if LLMs preserve different types of abstractions for other tasks$\\color{red}{^3}$ .\n \n\n> **RePlace tasks … 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**.\n\nThis question doesn't hold up because \n(1) **The world-irrelevant abstraction is indeed important in RePlace** as it's necessary for predicting future states, from which the optimal plan is derived (illustrated in Fig.3).\n(2) **A goal-oriented abstraction unimportant for a task doesn't exist**$\\color{red}{^4}$. State abstractions are task-specific and derived from the task's characteristics (i.e., transition dynamics and reward function). A state abstraction is a function that maps raw states to a smaller space, preserves the interested statistics (such as $Q^*$ and $\\pi^*$), and omits other irrelevant information. Therefore, **if a predicate is not important** (no impact on the $Q^*$ or $\\pi^*$) for a task, **it should be excluded from the goal-oriented abstractions**. Please refer to Sec 3.1 for formal definition and explanation of state abstractions as well as pointers to relevant literature.\n \n### **Questions**\n\n>  Finetuning with other categories of **tasks that are not goal-oriented** may not lead to goal-oriented world abstractions. \n\nTo clarify, we interpret \"tasks that are not goal-oriented\" as fine-tuning with sub-optimal ground truths, given that any RL task has a goal that maximizes the accumulated reward. We address this concern below.\n\n> 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.\n\nThis is an interesting question! To answer it, we synthesize another variant of GRIPPER, where ground-truth action sequences are sub-optimal (as you suggested, with some random legal moves that do not lead to the goal). We compare the average recovery rates (RR) for each abstraction from Llama3 models, fine-tuned on both the original and new datasets:\n\n|    Dataset    | Raw | World | $Q^*$ | $\\pi^*$|\n| --- | --- |--- |--- |--- |\n| original (reported in paper)  | 28.5 | 30.6  | 73.0  |  87.8  |\n| new  | 25.5 |27.4 | 63.3| 72.1 |\n\n\nOn the new dataset, the RR of goal-oriented abstractions decreases, which is expected as the model imitates sub-optimal actions. However, the margin between goal-oriented and general abstractions is still substantial. More importantly, the RR of world-irrelevant abstraction doesn't increase at all, which proves 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. \n\n### **Limitations**\n\n> However, the proposed probing method might not prevent supervised fine-tuning on a particular task from washing out LLM's innate preference either…\n\nIt's completely flexible for our framework to verify this assumption by directly probing pre-trained LLMs. To demonstrate that, we conduct the same probing experiments with three different pre-trained LMs. The results suggest that more advanced pre-training increases the tendency to prioritize goal-oriented abstractions over more general ones, and that SFT further enhances this trend. Please refer to the [Global Rebuttal](https://openreview.net/forum?id=lzfzjYuWgY&noteId=i9JnaUoNpz) for detailed results and in-depth analysis. \n\nMoreover, we thoroughly discuss all aspects of our work's limitations in Appendix A, which includes the scope of our experimental findings.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q19TKSxPO0", "reviewer": "Reviewer_9vdN", "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.", "questions": "See above.", "limitations": "Adequately discussed, but only in the Appendix, not in the main text.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The presentation is rigorous and comprehensive.\n\n2. The synthetic dataset is very good - simple but to the point. Upon release, it could be helpful for various other research papers.\n\n3. The experiment performed using the dataset is clear and well-designed.\n\n4. The question of whether LLMs build state abstractions is very interesting and currently an open question. Progress in this area would lead to more interpretable models.", "weaknesses": "Major comments:\n\n1. 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.\n\n2. 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?\n\nMinor comments:\n\n3. 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.\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful and encouraging feedback. We address all of your concerns below.\n\n### **Major Weakness 1: Findings are expected**\n\n> **(Findings unsurprising)** 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\n \n* World dynamics are also relevant to the task as the optimal plan is derived based on the prediction of future states (illustrated in Figure 3), but they are almost absent from LLM representations. \n* On the other hand, LLMs-sft could have omitted $Q^*$ abstraction since the $\\pi^*$ abstraction alone is sufficient for predicting optimal actions. However, we successfully recovered $Q^*$ abstraction from LLM-sft representations. It reveals that LLMs-sft preserve the long-term impact of each possible action, despite being fine-tuned on next-token prediction only. To the best of our knowledge, this exciting finding is novel and can inspire future research to conduct more in-depth analysis.\n* In the [Global Rebuttal](https://openreview.net/forum?id=lzfzjYuWgY&noteId=i9JnaUoNpz), we further elaborate on the broader significance of our framework and explain why our findings are surprising. We sincerely invite you to review it.\n\n\n> **(Experiments with pre-trained LLMs)** A more convincing experiment would involve no task adaptation.\n\nTo completely address your concern, we conduct new probing experiments on top of Phi3-17b, a brand-new SOTA LLM that achieves 19.94% success rate on GRIPPER without any fine-tuning. From its representations, goal-oriented abstractions can be probed with a much higher recovery rate than the world-irrelevant abstraction. We further include two smaller and weaker pre-trained LMs for comparison. The results suggest that more advanced pre-training leads to a higher tendency to prioritize goal-oriented abstractions over a more general one. Please refer to the [Global Rebuttal](https://openreview.net/forum?id=lzfzjYuWgY&noteId=i9JnaUoNpz) for detailed results and in-depth analysis. \n\n### **Major Weakness 2: Clarity issue**\n\n> **not clear how the hidden states were selected** … How about early layers of the LLM…\n\n* To clarify, the *last hidden states* refers to the hidden states at the last step before decoding, following the common practice of existing work. We treat the layer index as a hyperparameter (detailed in L620-L622) and use the last 6th layer for all experiments. \n* We have also experimented with using early layers. For instance, the average recovery rate (RR) of different abstractions from the 4th layer of Llama2-13b-sft is as follows:\n|    Layer    | Raw | World | $Q^*$ | $\\pi^*$|\n| --- | --- |--- |--- |--- |\n| 4th     | 13.7 |  13.0 | 14.7 | 7.02 |\n| Last 6th | 28.2  |30.2 | 71.1| 85.6 |\n\nThe RR for both the general- and goal-oriented abstractions significantly decreases. Therefore, the world abstraction cannot be recovered using earlier layers.\n\n\n### **Minor Weaknesses**\n\n* **small fonts in the table and better color selection**: Thank you for your feedback! We will enlarge the font in the figure and use deepened color in the revised version.\n* **Typos and expression**: Thank you for your careful reading! We will fix the typos and grammar in the revised version. They would be very easy fixes.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8rudgNZajm", "reviewer": "Reviewer_nz33", "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.\n\nIn experiments, the authors prompt LLMs to perform the REPLACE task and assess which abstractions each LLM is encoding. They have four findings:\n1. LLM representations reflect goal-oriented abstractions during decoding  \n2. Supervised-fine tuning increases the level of this goal-oriented abstraction\n3. We don't see the other abstractions present during decoding\n4. LLMs do not appear to build complete world representations.", "questions": "See questions above relating to robustness.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the main strengths of the paper are as follows:\n\n1. World abstractions: The idea of probing for different types of state abstractions is interesting. This is a new perspective for studying emergent world representations in LLMs, which have typically focused on all-or-nothing notions of world representations.  \n\n2. Clarity: The paper is mostly well-written and clearly structured. The authors do a good job of explaining their framework and describing the experimental methodology.  \n\n3. The planning task proposed in the paper (REPLACE) is an interesting task, and may be useful as a testbed for future research on analyzing world representations in language models.", "weaknesses": "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. \n\nSimilarly, 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.\n\nLess 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and detailed review. We address all of your concerns below.\n\n### **Major Weakness: Running more experiments with more datasets/LLMs**\n\n> REPLACE … planning task based on **a limited set of containers and objects**.\n\n* We acknowledge that RePlace has predefined domains of containers and objects due to the synthetic nature of the datasets. However, **this synthetic setup is crucial for faithful probing experiments**.\n**(1)**  It allows us to manipulate and track the underlying world state for probing experiments.\n**(2)** We can deliberately balance the dataset based on various attributes, such as the order of containers and the initial location of agents and objects, to eliminate potential biases. However, it's hard to do the same thing for real-world datasets.\n\n* Due to the above reasons, almost **all existing work on probing world representations uses synthetic datasets**. [1,2] use game script datasets synthesized with the rule of the Othello game. [3,4] also employ synthesized text datasets based on limited sets of containers and objects, inspiring our datasets' design (mentioned in L175-176). However, we additionally have an agent exploring the environment and manipulating objects with some physical constraints to accomplish a goal, which makes the task more challenging and complex enough to differentiate the spaces of various abstractions. Furthermore, we have introduced different dataset variants to create a more diverse and realistic task setup (detailed in L227-234).\n\n* In addition, **our new task spans a broad domain** (elaborated in L179-182). It is closely related to the Gripper problem$\\color{red}{^1}$ and the EntityTracking task [4,5], which is widely adopted in planning and NLP interpretability research. This task also involves tracking entities, aggregating multi-hop information, and recognizing textual entailment, essential skills for many real-world NLP tasks.\n\n> Since the paper is focused on assessing world representations generally, there needs to be **more evaluation settings and datasets**.\n\n* Our framework, the central contribution of this work, is generic and can be seamlessly applied to other LLMs and tasks wherever the abstract state spaces are distinguishable. However, we are not aiming to reach a conclusion that applies universally to all LLMs and tasks (discussed in L494-498). This is also impossible, given that there is even a lack of available testbeds.\n* We didn't use other datasets employed by existing work as **they don't meet the 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. Our new task/datasets fill this gap.\n\n> … doesn't include any experiments involving Othello\n\n* In the paper (L166-170), we explain **why the Othello dataset is not a faithful testbed**: it's not complex enough to differentiate abstract states at various levels. \n* Nonetheless, the Othello dataset was originally used to probe a small-scale Transformer trained on the game script data. But our work focuses on LLMs, and game script is very different from natural text that LLMs are trained on. \n\n> … more settings for REPLACE should be considered to provide more **ablations** …\n\n* **We have done such ablation**. Due to the limit of space, we highlight the conclusion in the main body of the paper (L337-L340) and provide more details in Appendix K. To sum up, we have created a new setting where an originally goal-unrelated predicate (e.g. color info) belongs to goal-oriented abstractions. As such, the recovery rate of color information has significantly increased from zero. This further consolidates our finding that SFT mainly enhances goal-oriented abstractions. \n* During the rebuttal period, we have done another ablation to answer Reviewer 4mUW's question, which shows that LLMs fine-tuned with sub-optimal demonstrations still prioritize goal-oriented abstractions, echoing our main finding. For the sake of not repeating, we sincerely invite you to refer to the rebuttal to Reviewer 4mUW for more details and results of this ablation study$\\color{red}{^2}$.\n\n> … could benefit from **more extensive model comparisons** … all from the Llama family or Mistral…\n\n* We have already adopted four LLMs for experiments, each of which is experimented with both ICL and SFT, while existing work like [1,2] uses only one and [3] uses two. Even though some of the experimented LLMs are from the same family, they are at different scales. Llama2 and Llama3 have different tokenizers trained on different data and post-training methods (DPO, etc.). \n* To address your concern fully, we adopt another SOTA LLM, phi3-17b, and conduct new probing experiments. The results show that this SOTA LLM, whether or not it has been fine-tuned, tends to maintain a more goal-oriented abstraction of the underlying world. It is totally consistent with our original findings. Please refer to the [Global Rebuttal](https://openreview.net/forum?id=lzfzjYuWgY&noteId=i9JnaUoNpz) for more details and analysis.\n* Consequently, we end up experimenting with **12** LLM variants (5 SFT + 4 ICL + 3 pre-trained), which, to the best of our knowledge, is **the most extensive comparison**$\\color{red}{^3}$ so far in the area of probing world representations.\n\n### **Minor Weakness: text in Figure 5 too small & formal definition of $Q^\\*$**\n\nThank you for your suggestions! We will increase the font size in Fig.5 and quote the formal definition of $Q^*$ in the revised version. They would be very easy fixes.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oBEndNhVys", "reviewer": "Reviewer_CXua", "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.", "questions": "**Questions**\n- 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?\n- 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)?\n\n**Suggestions**\n- 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.\n- 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": "Discussed and addressed in the appendix.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It's a neat idea and good contribution to formalise the types of abstractions LLMs can use to represent the world through state abstraction theory. I believe using this framework will make future work around these questions more interesting and grounded.\n- The paper is well-written and easy to understand\n- The authors design a synthetic task that is both easy to understand, has modular and distinct state abstractions, and is somewhat complex for LLMs to perform (requiring planning to perform optimally).", "weaknesses": "**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**\n\nThe 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,\nprioritizing 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.\n\n**Missing baselines**\n\nIt'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.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and constructive feedback. We address all of your concerns below.\n\n### **Weakness 1: Results are unsurprising**\n\n> **(unsurprising finding)**  LLMs fine-tuned for a task start forming more useful goal-oriented state abstraction representations…  is unsurprising … \n\n* It's indeed surprising when viewed in the context of existing work$\\color{red}{^1}$. We provide detailed elaboration in the [Global Rebuttal](https://openreview.net/forum?id=lzfzjYuWgY&noteId=i9JnaUoNpz), with other surprising findings and broader significance of our framework. We sincerely invite you to review it.\n\n> **(it's unsurprising because)** you're training it to discard irrelevant information and encode relevant information for the goal…\n\n* World dynamics are relevant to the goal since the optimal plan is derived based on predicting future states (illustrated in Fig.3). \n* On the other hand, LLMs-sft could have been abstracted away $Q^*$ abstraction since the $\\pi^*$ abstraction alone is sufficient for predicting optimal actions. However, we successfully recovered $Q^*$ abstraction from LLM-sft representations. It reveals that LLMs-sft preserve the long-term impact of each possible action, despite being fine-tuned on next-token prediction only. To the best of our knowledge, this interesting finding is novel, and we believe it can inspire future research to conduct a more in-depth investigation.\n\n> **(experiments on a task that LLMs do well in context)** 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…\n\n* 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.\n* Due to these challenges, we instead start with a simple enough task. As explained in L179-184, the RePlace task closely resembles the [gripper problem](https://shorturl.at/j0wp5), one of the most basic planning tasks, where the transition is deterministic, the state can be derived without uncertainty, the goal must be achievable and the optimal plan can be searched in O(n). Therefore, we cannot make the task simpler while still having distinctive state space for different abstractions. \n\n> **(experiments with pre-trained LLMs)** …using LLMs that can do it in-context (e.g. larger models or better models)\n\nTo fully address your concern, we conduct the same set of probing experiments with Phi3-17b, a brand-new SOTA LLM that achieves 19.94% success rate on GRIPPER with ICL only$\\color{red}{^2}$. In response to your question that whether better LLMs are more likely to use general abstractions, we employ two smaller and weaker pre-trained LMs for comparative study. The new results indicate that (1) **goal-oriented abstractions are probed from pre-trained Phi3-17b with significantly higher recovery rate**, and (2) **more advanced pre-training leads to higher priority of encoding goal-oriented abstractions over a more general one**. Please refer to the Global Rebuttal for detailed results and in-depth analysis.\n\n> **(clarification needed)** I'd like to see the text indicate that the claims are about task-specific (i..e. fine-tuned) LLMs.\n\nThank you for your suggestion! Due to space constraints, we summarize the findings for LLMs-sft/icl in one sentence in the abstract/Intro. In the experiments section (Sec 6.2), however, we have reported and analyzed the findings for LLMs-sft/icl separately (L320-321, L331). We will do the same thing for the abstract/Intro in the revised version to make it clear, which will be flexible given more space. \n\n### **Weakness 2: learns a probe on top of a non-LLM**\n\nThank you for your suggestion! To address your concern, we train two 6-layer decoder-only Transformers on our datasets, which achieves around 30% success rates. The average recovery rate of each abstraction is: \n\n|Datasets| Raw | World | $Q^*$ | $\\pi^*$|\n| --- | --- |--- |--- |--- |\n|GRIPPER| 27.7 | 29.9  | 52.5  | 28.3| \n|COOK| 23.9 | 23.8  | 32.5 |28.7| \n\nInterestingly, the raw state and world-irrelevant abstractions can be probed with a much similar recovery rates as from LLMs-sft (\\~30% for both). The main difference lies in goal-oriented abstractions, i.e. $Q^*$-irrelevant abstraction (\\~73% from LLMs-sft) and $\\pi^*$-irrelevant abstraction (\\~83%). It suggests that pre-training mainly improves the encoding of goal-oriented abstractions.\n\nWe will include this result in Fig. 5 in our paper. Thanks again for your valuable suggestion!\n\n### **Questions**\n* Legal rate of random move: The legal rate of a random move is 21.65%. All LLMs-icl have significantly higher legal rates than this baseline. \n* Probing earlier layers: We treat the layer index as a hyperparameter (detailed in L620-L622) and use the last 6th layer based on the overall performance on the validation set. We have also experimented with using early layers. For instance, the average recovery rate (RR) of each abstraction from the 4th layer of Llama2-13b-sft is:\n\n|    Layer    | Raw | World | $Q^*$ | $\\pi^*$|\n| --- | --- |--- |--- |--- |\n| 4th     | 13.7 |  13.0 | 14.7 | 7.02 |\n\nThe RR of both the general- and goal-oriented abstractions are much lower.\n\n### **Suggestions**\n* Thank you for the suggestion! We will summarize the findings for LLMs-icl/sft separately in the abstract/Intro to avoid confusion, which would be an easy fix.\n* Thank you for your question! `subgoal(l, j)` denotes that the container to be visited at $j$-th step is at relative distance $l$. We will clarify it in Line 270.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their time and feedback. We are encouraged that all aspects of our work are widely recognized. The reviewers found our addressed problem to be new and important (R1, R2, R3), our framework novel and neat (R2, R4, R5), our task/datasets thoughtfully designed and thus useful for future research (R1, R2, R3, R4, R5), experiments well-designed (R3), findings sound (R1), and our paper well-written (R2, R3, R4, R5). Most of the concerns regard two interrelated aspects, namely the findings' surprisal (R1, R3, R5) and the lack of probing experiments on pre-trained LLMs without ICL/SFT (all reviewers$\color{red}{^1}$). We address both of them below.
### **1. Surprisal of findings**

One primary concern is that the finding, *SFT mainly enhances goal-oriented abstraction*, is not entirely surprising$\color{red}{^2}$. Below, we elaborate on this and other surprising findings, as well as the broader significance of our framework.
* **This finding is surprising when viewed in the context of existing work.** While not entirely unexpected in hindsight, it 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]. When we use a testbed that can differentiate abstraction at various levels, however, we find that LLMs tend to discard the transition dynamics, essential ingredient of world models, when possible. This finding also explains and hence resolves the conflicting conclusions in existing works. For instance, [2,4] successfully probes the Othello board state, which is already the goal-oriented abstraction for predicting legal moves (the task Transformers are trained on), whereas [5] struggles to probe the entity status, which falls outside the goal-oriented abstractions for which LLMs are optimized. 
* **It’s more surprising how probing different sets of predicates WITHOUT the notion of world abstractions can lead to directly opposing conclusions**. Imagine two researchers independently assessing an LLM's world representation on RePlace, using current methodologies. As demonstrated in our experiments (Sec 6.2, illustrated in L374-379), Researcher A selecting `agentLoc`, `nearby`, and `held_u`, might infer that SFT enforces LLMs to develop an implicit world model, while Researcher B, focusing on `store` and `boxName`, comes to a totally different conclusion, arguing that LLMs-sft have little awareness of the underlying world. This disparity reflects the current state of the field$\color{red}{^3}$. Thus, our framework provides a more grounded basis for future research, avoiding blind spot conflict and enabling others to dissect and critically evaluate claims in a principled way.

* It's noteworthy that **LLM-sft representations encode the $Q^\*$ abstraction, which could have been abstracted away** as $\pi^\*$ abstraction is sufficient for deriving optimal actions. It reveals that LLMs-sft preserve the long-term impact of actions even when fine-tuned with NTP only. This finding is novel in this area, and it leaves many interesting questions open for future work: How does this abstraction emerge from SFT? Could different training methods enhance or diminish it? Does it affect only the immediate prediction, or do LLMs continuously revisit previous abstractions during decoding? 

### **2. Probing experiments on pre-trained LMs**
Similarly, another concern is that *the pre-trained LLMs, especially those that can perform the task well without SFT, may tend to maintain general abstractions, but SFT fundamentally alters this pattern*. To best address it, we employ Phi3-17b [6], a brand-new open LLM that hits a 19.94% success rate on GRIPPER without SFT. To further understand the impact of pre-training on the encoding of world abstractions, we employ two weaker and smaller LLMs, Phi3-3.8b and Pythia-70m, which reach near-random success rates$\color{red}{^4}$, for comparative analysis. We additionally fine-tune another Phi3-17b. Plot A in the **attached pdf** reports the average recovery rate ($\color{red}{\text{RR}}$) of predicates within each type of world abstractions from different LLMs, and Plot B reports the RR of each predicate across all LLMs. We highlight three key findings below.

1) **Pre-trained Phi3-17b prioritizes maintaining goal-oriented abstractions over a more general one**. Plot A clearly shows that. And SFT substantially strengthens this tendency. 
2) **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**, thereby widening the gap between their RR. This is apparent from Plot B, where the RR variance among different LLMs is mainly found in predicates that pertain to $\pi^*$- and $Q^*$-irrelevant abstractions. In particular, the predicates that uniquely pertain to goal-oriented abstractions, such as `nearby` and `nextObjDirect`, are probed with remarkably higher recovery rates from Phi3-17b than from Pythia. In contrast, the RR of `store` and `held_g`, essential for world-irrelevant abstraction, are almost identical across all LLMs variants. It suggests that **more advanced pre-training doesn't necessarily enhance the encoding of world dynamics**. 
3) **Advanced pre-training is NOT sufficient for efficient world modeling**. Interestingly, probing from Phi3-17b has much higher RR of `boxName`, which is neither relevant to task completion nor to the transition dynamics. This verifies our Finding 4 (L366-371) that LLMs are limited in building world representations.

We thank the reviewers for their suggestion of adding this experiment. It turns out to be a good chance to gain new insights about SOTA LLM's world representation which supplement and consolidate our original findings. We will include these results in the revised version, which would be an easy fix.

Meta-review: 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.
Decision: Accept (poster)