Evaluating the World Models Used by Pretrained Learners

Thank you for your review. We’re glad that you found our methodology novel and our results "inspiring". We appreciate the positive feedback, and your suggestions have helped improve the paper.

I wonder if authors have plans to scale up the experiments to larger models, or if they could discuss potential limitations in generalizing their findings to LLMs.

This is a great question about LLMs. Our metrics are applicable to any learning algorithm where ground truth state is known to the analys. Much of the prior world model literature has focused on simpler synthetic domains like board games [1-3]. These testbed settings are important for developing general metrics and assessing model capabilities; e.g. what can be said about a transformer's world modeling capabilities if it can't extrapolate in even simple domains? Having developed and validated these metrics on the testbed settings, we're eager to apply them to LLMs in future work.

[1] Toshniwal, Shubham, et al. “Chess as a Testbed for Language Model State Tracking.” AAAI. 2022.
[2] Li, Kenneth, et al. “Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task.” ICLR. 2023.
[3] Vafa, Keyon, et al. Evaluating the World Model Implicit in a Generative Model. NeurIPS. 2024.

Please provide a more explicit explanation of the question "what does it mean to test if a learner has a world model embodied in it?", and how experiment findings relate to this question. And please state the motivation behind this question.

Thank you for this question. Below we provide more background and we've updated the paper to include this point more clearly.

While existing approaches often study whether a fixed model's outputs are consistent with a world model (e.g., checking if language models make physically plausible predictions), this doesn't tell us whether the model actually learns and transfers knowledge using the correct world model. The motivation for this question is both theoretical and practical:

• Theoretical: Having a world model means more than just making valid predictions - it implies having a compressed, generalizable representation of how a domain works that can be applied to new situations. A true world model should manifest in how a system learns and adapts, not just in its static outputs.
• Practical: Many of the benefits of world models (like few-shot learning and transfer) specifically arise from how models learn new tasks. For example, if a language model truly has a physics world model, it should learn new physics-related tasks by using its underlying physics knowledge rather than learning each task from scratch.

Our experimental findings directly address this question by:

• Providing a formal framework that defines what it means for a learner (rather than just a fixed model) to use a world model, based on whether it shows inductive bias toward using state information when learning new tasks
• Demonstrating that models which appear to have world models based on traditional evaluations (e.g., making valid next-token predictions) often fail to show the learning patterns we would expect from a true world model
• Revealing that many models rely on "bundles of heuristics" rather than coherent world models -- they can make valid predictions but don't learn new tasks by leveraging an underlying model of the domain

The orbital mechanics example particularly illustrates this -- while the model can make accurate predictions about planetary motion, our metrics reveal it doesn't actually learn using Newtonian physics as its inductive bias. Instead, it learns different heuristics for different situations, showing it lacks a true physics world model despite superficially valid outputs.

What are the advantages of the proposed measuring procedure over other studies to assess whether pretrained models develop world models (e.g., studying the behavior of fixed models)?

Both fixed models and learners are important to study. In addition to the above, another reason we focus on learners is because one of the promises of foundation models having world models is that they can be adapted to new tasks by using world models. In other words, having reliable world models means algorithms can build on top of each other. The physics example illustrates why failures of world models in learners can be damaging; because the learner doesn't use Newtonian physics, it will be brittle for adapting to other tasks that rely on it.

The title for Table 1 and Table 2 should be more detailed and include result analysis, including key trends that are highlighted in the tables and why the trend exists.

Thank you for your suggestion for the captions for Tables 1 and 2. We've updated them to include more analysis.

Thanks for the authors' response. My concerns have been addressed. And the explanations above are expected to be included in the final version of the paper. One more suggestion is to improve the readability of current version, especially for those who don't have relevant background. I've increased my confidence score.

Thank you for your thoughtful review of our paper and we’re glad you found our results interesting and surprising. We appreciate the positive feedback, and your suggestions have helped improve the paper.

I’m curious whether the measurement of inductive bias is highly dependent on the distribution of the training data. For instance, if the data distribution is more uniform and not covers only a small subspace, would this significantly alter the conclusions? In the case of orbital mechanics, would a model trained on a more uniform distribution still capture universal laws?

This is a great question about the distribution of training data. We empirically study this in Othello and orbital mechanics:

• Othello: We test models on two kinds of Othello datasets -- transcripts from real championship games (championship Othello), and those that come from random gameplay (synthetic Othello). These cover two kinds of training data distributions -- more uniform versus narrower. We don't see a significant difference in inductive bias but we find that state recovery is larger for championship Othello.
• Orbital mechanics: We tried to get as much coverage in the training distribution as possible by, for instance, generating orbits with 10 distinct uniformly-spaced eccentricities and by randomly sampling the masses of the two objects and the initial distance between the two objects. It’s an interesting question whether the strength of inductive bias is dependent on the coverage of the data distribution, but our experiments show that despite the wide coverage, the trained model consistently shows weak inductive bias.

In the cases of Lattice and Othello, do the ground-truth states represent the minimal information needed for next-token prediction? I’m particularly considering the results shown in Figure 2.

This is a great question. The answer is yes: the state structures of both lattice and Othello obey a DFA, so the ground-truth state represents the minimal information needed for next-token prediction. [1] shows that there is an if-and-only-if relationship: perfect next-token prediction is necessary and sufficient for recovering ground truth states. But at the same time, near-perfect next-token prediction does not imply near-perfect state recovery, hence the importance of additional metrics.

[1] Vafa, Keyon, et al. “Evaluating the World Model Implicit in a Generative Model.” NeurIPS. 2024.

I appreciate the 'empirical' evaluation framework, but it seems to depend on various factors, like the IB loss function, reconstruction loss, transfer setup, data distribution, and so on. How sensitive are the conclusions in the paper to these parameters, and how well do they generalize?

This is a good question. We provide ablations across settings in Appendix D of the submission and find that the results are robust against these choices. Additionally, we've included new complementary evaluation metrics in Appendix G. These new evaluation metrics are based on correlation measures, so they don't require making choices about things like model architecture. Importantly, they come to similar conclusions as the original metrics in the main text.

Could you comment more on how these world models connect with the memorization and reasoning concepts in the LLM community if possible?

This is an interesting question. Just as we found models can achieve good performance through "bundles of heuristics" rather than true world models, LLMs can generate valid outputs through memorization and pattern matching without deeper understanding. However, true reasoning capabilities likely require having a world model -- compressed, generalizable representations that can be applied to new situations. This is similar to how we find that models without strong inductive bias toward state struggle with transfer learning, similar to how LLMs that rely primarily on memorization often struggle with novel reasoning tasks. While memorization can enable good performance on tasks similar to training data, our metrics suggest that building AI systems capable of robust reasoning may require developing training approaches that encourage the formation of true world models rather than surface-level patterns.

We greatly appreciate your detailed read of the paper and your thorough review. We’re glad that you found our work to be novel and of broad scientific interest. Your comments were incredibly constructive and we believe we've addressed all of them in the updated paper -- addressing your comments has improved the clarity of the paper and sharpened the main conceptual points.

To summarize, we have:

• Revised the framework section in response to your comments.
• Included additional ways to estimate our metrics that reach the same conclusions (Appendix G).
• Added additional empirical results (e.g. more symbolic regression experiments in Physics in Appendix F)
• Sharpened the writing overall.

We have more details below.

The main contribution of the paper are the metrics defined in Section 2.1. While their motivation is understandable, it is not immediately clear, nor in my opinion sufficiently justified, why these metrics are tracking what the authors propose they are tracking... For example, even when dividing by the random baseline, it is hard to interpret what different scales correspond to... On the gravitational orbits experiments, the authors also turn to symbolic regression... and this seems like a much more straightforwardly interpretable result.

You bring up a good point about scaling, e.g. interpreting whether IB scores of 0.6 are small or large. As you mention in your review, this is one of the reasons we include baselines such as a model trained on state. Nevertheless, in order to provide more clarity, we've included new quantitative measures in Appendix G that correspond to correlations, where the scaling is more intuitive. At a high level, these additional metrics measure whether a learner's extrapolation behavior for two data points in the same state is more correlated than its extrapolations in different states. We've re-implemented the analyses in our submission with these new metrics and find similar results (we've re-implemented the metrics for all models except for Mamba-2, which we did not have time to re-run during the rebuttal period). We include more details and discussion in our updated appendix.

We're also glad you appreciated the interpretability of the symbolic regression example. While symbolic regression is suitable for physics, where the underlying relationships tend to be mathematically structured and well-defined, the nature of the state space structure for other domains like Othello make other analysis methods more suitable.

My understanding is that the calculated metrics only provide potentially interesting information when the states used to compute them are the real states used to generate the data... The settings the authors choose for their experiments make sense, but... we don’t really care that much whether a specific architecture, run on a setting with a specific encoding, does or doesn’t recover a particular functional form. Rather, they serve mostly as exemplifications and explorations of the use of the proposed metrics, which are indeed the main contribution

You bring up an important point: our metrics are intended for domains where there is some notion of true state used to generate data. In doing this we've followed much of the prior literature on assessing the world models of LLMs, which focus on synthetic domains such as games [1, 2], logic puzzles [3], and navigation [4]. These testbeds are important for developing general metrics and assessing model capabilities; e.g. what can be said about a transformer's world model if it can't extrapolate in even simple domains? Settings where state is known also come up in more complex systems (e.g. anything that obeys a DFA), such as search engines, control systems, and genetics [5].

We agree that the specific experimental findings are not the main contribution. Rather, the experiments serve two key purposes: (1) They demonstrate how to implement our metrics in practice and validate that they capture meaningful properties of learners. (2) They reveal interesting phenomena about how different architectures learn from data - e.g. that learners pretrained on next-token prediction can recover enough of state to make accurate next token predictions despite incoherent full states (Figure 2). This also follows the LLM world model literature [1, 2, 3, 4]. These insights, while derived from synthetic settings, help us understand potential limitations of current approaches.

[1] Toshniwal, Shubham, et al. “Chess as a Testbed for Language Model State Tracking.” AAAI. 2022
[2] Li, Kenneth, et al. “Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task.” ICLR. 2023
[3] Li, Belinda Z., et al. “Implicit Representations of Meaning in Neural Language Models.” ACL. 2021
[4] Vafa, Keyon, et al. “Evaluating the World Model Implicit in a Generative Model.” NeurIPS. 2024
[5] Gribkof, Eric. “Applications of Deterministic Finite Automata.” Lecture note, UC Davis EC120, Sprint 2013

Thank you for your review. Your review makes several helpful points that will improve our paper. However, our rebuttal clarifies a potential confusion about differences between the study of world models in the RL and LLM literatures. We've made this more clear in our revision.

The paper defines a world model as a mapping from inputs to “state” (page 1); however, this mapping is clearly insufficient... Given the usage of reinforcement learning terminology, I am surprised by the lack of engagement with the relevant literature... The "partial reconstruction of state" metric, defined in Definition 2.2, seems related to concepts like bisimulation metrics (how can we cluster states that behave similarly?) and MDP homomorphisms (what state mappings preserve the essential dynamics?)... Could the authors please clarify and motivate this instantiation of a world model?

We think there's a potential confusion: while the term "world models" is used in both the RL and LLM literatures, it means different things in each literature and the goals are quite different. While your comments refer to the concept of world models in RL, our paper is contributing to the LLM literature that evaluates whether LLMs build world models that accord with the real world [1, 2, 3, 4]. In RL, world models refer to representations (or even the specific neural network) learned by an agent, and their quality is typically measured by how well they perform policy optimization [5, 6]. In contrast, the LLM literature focuses on evaluation: a learning algorithm is evaluated by how well it can recover an externally defined mapping from a true world model. In other words, it seeks to answer the question, “does a model's behavior reflect understanding of the true world state?” Our work follows this evaluation perspective in the LLM literature, and we contribute to this literature by studying learning algorithms rather than fixed models. Similarly, despite using similar terminology, the state abstraction literature in RL (e.g., bisimulation, MDP homomorphism), which is primarily focused on aggregating and simplifying the state-space of an MDP [7, 8], has a different goal from that of world model evaluation in LLMs. Again, by contrast, the LLM literature focuses on assessing whether generative models recover underlying world models with high representational fidelity.

We agree that an arbitrary lookup table recovered by a model is not an interesting world model. This is why the LLM literature studies whether learning algorithms can recover real-world mappings. What makes mappings from inputs to states interesting in the real world is that we can build on top of them; for example, in the physics example, the mapping from a sequence of planetary movement to a state vector is scientifically useful because those states (e.g. relative velocity and masses) are important concepts for other physical tasks.

We've updated the related work to include a discussion about the differences in world model recovery in LLMs and RL. Please let us know if you have any more questions or if anything is unclear about the LLM world model literature.

[1] Toshniwal, Shubham, et al. “Chess as a Testbed for Language Model State Tracking.” AAAI. 2022.
[2] Li, Kenneth, et al. “Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task.” ICLR. 2023.
[3] Li, Belinda Z., et al. “Implicit Representations of Meaning in Neural Language Models.” ACL. 2021.
[4] Vafa, Keyon, et al. “Evaluating the World Model Implicit in a Generative Model.” NeurIPS. 2024.
[5] Ha, David, and Jürgen Schmidhuber. “World Models.” 2018.
[6] Chen, Chang, et al. “TransDreamer: Reinforcement Learning with Transformer World Models.” arXiv:2202.09481. 2024.
[7] Abel, David, et al. Near Optimal Behavior via Approximate State Abstraction. arXiv:1701.04113. 2017.
[8] Li, Lihong, Thomas J. Walsh, and Michael L. Littman. "Towards a unified theory of state abstraction for MDPs." AI&M1.2 (2006): 3.

I don’t understand the synthetic dataset construction part in Section 2.2. If one constructs a dataset that is generated from some behavior policy, there is no reason that assigning random outputs would map to any meaningful task.

The synthetic datasets are constructed to study the extrapolation behavior of learners. Importantly, they're not completely random outputs. They're constructed to enforce that within a dataset, two points in the same state have the same mapping. It's just the assignments of points that's random. It's exactly because there's no structure beside state that enables us to test whether learners extrapolate based on state. By repeatedly applying a learner to different datasets, we can test whether its extrapolations can be predicted based on extrapolations of other points with the same state. We've made this more clear in the revision.