Evaluating the World Model Implicit in a Generative Model — part 1/5

Cleaned-text code points [0, 15462). Bibliography omitted here; source.md includes it.

# Evaluating the World Model Implicit in a Generative Model

###### Abstract

Recent work suggests that large language models may implicitly learn world models. How should we assess this possibility? We formalize this question for the case where the underlying reality is governed by a deterministic finite automaton. This includes problems as diverse as simple logical reasoning, geographic navigation, game-playing, and chemistry. We propose new evaluation metrics for world model recovery inspired by the classic Myhill-Nerode theorem from language theory. We illustrate their utility in three domains: game playing, logic puzzles, and navigation. In all domains, the generative models we consider do well on existing diagnostics for assessing world models, but our evaluation metrics reveal their world models to be far less coherent than they appear. Such incoherence creates fragility: using a generative model to solve related but subtly different tasks can lead to failures. Building generative models that meaningfully capture the underlying logic of the domains they model would be immensely valuable; our results suggest new ways to assess how close a given model is to that goal.

## 1 Introduction

Large language models (LLMs) appear to have capacities that far exceed the next-token prediction task they were trained to perform [ 17 , 39 , 35 ] . Recent work suggests a reason: they are implicitly recovering high-fidelity representations of the underlying domains they are trained on [ 1 , 20 ] .

An algorithm that recovers a “world model” from sequence data would be extremely valuable. As an example, consider how one might build a navigation tool today: meticulously map each street and intersection, and then use a search algorithm to provide directions. The success of language models suggests an alternative approach: collect turn-by-turn sequences from trips in a city (e.g. “ East North …”) and then train a sequence model on them. If the sequence model successfully recovers the world model, we would obtain a map of the city without ever mapping it and a routing algorithm simply by predicting the next turn. This example is not far-fetched: it is the reason language models are used in scientific domains such as protein generation, genetics and chemistry [ 7 , 21 , 3 , 14 , 6 ] .

All of this relies on the presumption that the sequence model has recovered the true world model; but how can we test whether it actually has? Answering this question requires first defining what we mean by the true world model. Toshniwal et al. [36] and Li et al. [20] proposed a concrete and influential approach: study whether sequence models trained on board game transcripts (e.g. chess and Othello) recover the underlying game rules. Inspired by this approach, we consider the case where the underlying world can be summarized by a finite collection of states and rules governing transitions between the states; this includes many domains such as logic [ 19 ] , location tracking [ 28 , 9 ] , games [ 36 , 20 ] , and several of the scientific applications described above. As a result, the “world” in these domains can be modeled as a deterministic finite automaton (DFA).

We show the difficulty in evaluating implicit world models. Consider an existing approach: for a given sequence, compare the next tokens outputted by the generative model to the set of valid next tokens for the state implied by that sequence [ 36 , 20 ] . Though intuitive, this approach can fail to diagnose severe problems, and we illustrate this concretely. The classic Myhill-Nerode theorem [ 26 , 27 ] provides intuition: every pair of distinct states can be distinguished by some sequence (admitted by one state but not the other). Unless those minimal distinguishing sequences are of length one, looking at the next single token outputted will not reliably assess whether the generative model has an accurate model of the underlying state.

The logic of Myhill-Nerode suggests two metrics for measuring whether a generative model effectively captures underlying states and transitions. The first metric summarizes sequence compression : under the DFA, sequences that lead to the same state must have the same continuations; so one can test whether the generative model has similar sequences of outputs when started on these two sequences. The second metric summarizes sequence distinction : under the DFA, two sequences that lead to distinct states should have distinct continuations; so one can test whether the generative model matches these distinct outputs when started at these two sequences. We formally define these metrics and provide model-agnostic procedures for calculating them when given query access to the true DFA.

To illustrate these ideas, we first take the stylized mapping example literally. We construct a turn-by-turn sequence dataset of taxi rides in New York City. We then assess to what extent transformers successfully recover the true street map of Manhattan. By the usual metrics, the transformers do very well: their predicted next-direction is a valid turn nearly 100% of the time and their state representations even appear to encode the current location of the ride. Our evaluation methods reveal they are very far from recovering the true street map of New York City. As a visualization, we use graph reconstruction techniques to recover each model’s implicit street map of New York City. The resulting map bears little resemblance to the actual streets of Manhattan, containing streets with impossible physical orientations and flyovers above other streets (see Figure 3 ). Because these transformers fail to recover the true street map of New York City, they are fragile for downstream tasks. While they sometimes have amazing route planning abilities, their performance breaks down when detours are introduced.

These results are not unique to maps and navigation. For both Othello and logic puzzles, we use our evaluation metrics to show language models can perform remarkably well on some tasks despite being far from recovering the true world model. These results demonstrate the importance of using theoretically-grounded evaluation metrics if our goal is to build language models that capture accurate world models of the domains they are trained in. We release our benchmark dataset of taxi rides in New York City along with software implementing our evaluation metrics. 1 1 1 https://github.com/keyonvafa/world-model-evaluation

Related work. Our paper builds on influential work studying whether generative models recover a world model in the context of games. Toshniwal et al. [36] and Li et al. [20] pioneered the study of games as a testbed for world model evaluation, studying tests for chess and Othello, respectively, which were further studied by Hazineh et al. [10] and Kuo et al. [18] . Our evaluation metrics apply to these games because they are DFAs. A common method for assessing whether a trained model has recovered a world model uses probes that assess whether a neural network’s representation can recover some real-world state [ 11 , 19 , 1 , 16 , 20 ] . By contrast, our evaluation metrics are model-agnostic: they’re based only on sequences. While the results from our evaluation metrics sometimes align with those used in existing work, they also reveal incoherence in world models that are not captured by existing diagnostics.

We study whether a language model trained on sequences of directions recovers the true underlying map. This question relates to other state tracking and navigation problems studied in the language modeling literature [ 31 , 32 ] . For example, Patel & Pavlick [28] show that larger LLMs ground spatial concepts like cardinal directions to locations in a grid world and generalize to various grid layouts. Relatedly, Schumann & Riezler [30] demonstrate that transformer-based models can generate navigation instructions in language from underlying graphs. Additionally, Guan et al. [9] use LLMs to perform planning tasks from natural language descriptions. Our results suggest that LLMs can perform some of these tasks well (such as finding shortest paths between two points on a map) without having a coherent world model.

Additionally, our evaluation metrics compare the language accepted by a sequence model to that of an underlying DFA. Existing work studies whether transformers and other sequence models are theoretically capable of recognizing languages in different complexity classes [ 34 , 4 , 22 , 23 , 24 ] . Most relevant to our work, Liu et al. [22] show that low-depth transformers can theoretically represent any finite state automata, and show that transformers trained explicitly to predict their labeled states are capable of doing so. In contrast, our paper doesn’t aim to study whether models are theoretically capable of recovering underlying automata or whether they can do so when given state labels. Instead, we provide metrics for assessing how closely a given model recovers the underlying DFA.

## 2 Framework

In this section, we lay out a framework to interface between generative sequence models and world models represented by deterministic finite automata. Both of these are built on the shared scaffolding of tokens, sequences (a.k.a. strings), and languages.

Tokens and sequences. We consider a finite alphabet Σ \Sigma with tokens a ∈ Σ a\in\Sigma , and sequences s = ( a 1 , a 2 , … ) s=(a_{1},a_{2},\ldots) . Let Σ ∗ \Sigma^{*} denote the collection of sequences on the alphabet.

Generative models. A generative model m ⁡ ( ⋅ ) : Σ ∗ → Δ ⁡ ( Σ ) m(\cdot)\colon\Sigma^{*}\rightarrow\Delta(\Sigma) is a probability distribution over next-tokens given an input sequence. That is, m ⁡ ( s ) ∈ Δ ⁡ ( Σ ) m(s)\in\Delta(\Sigma) , and m ⁡ ( a | s ) m(a\penalty\ |\penalty\ s) is the probability assigned to token a ∈ Σ a\in\Sigma given an input sequence s s . Starting at a sequence s s , the set of non-empty sequences the model can generate with positive probability is: L m ​ ( s ) = { a 1 ​ a 2 ​ … ​ a k : ∀ j ⁡ < k , m ⁡ ( a j + 1 | s ​ a 1 ​ … ​ a j ) > ​ 0 } . L^{m}(s)=\{a_{1}a_{2}...a_{k}:\forall j<k,\ m(a_{j+1}\penalty\ |\penalty\ sa_{1}...a_{j})>0\}. For simplicity, we write the equation above for next-tokens with nonzero probability, but in practice we set a minimum probability ϵ > 0 \epsilon>0 corresponding to next-tokens with non-negligible probability.

Deterministic finite automata (DFA). We use standard notation for a deterministic finite state automaton W = ( Q , Σ , δ , q 0 , F ) W=(Q,\Sigma,\delta,q_{0},F) (see Appendix C for a complete definition). As a simplifying assumption, we consider the case where there is a special state q reject q_{\text{reject}} with no outgoing transitions and F = Q ∖ { q reject } F=Q\setminus\{q_{\text{reject}}\} (i.e., the DFA accepts all valid states). An extended transition function δ ^ \hat{\delta} takes a state and a sequence, and it inductively applies δ \delta to each token of the sequence. A token or a sequence is valid if and only if the output of δ \delta or δ ^ \hat{\delta} respectively starting from q 0 q_{0} is not q reject q_{\text{reject}} .

We define L W ​ ( q ) L^{W}(q) to be the set of valid, non-empty sequences that are accepted by the DFA starting at state q q . We also define q ⁡ ( s ) ∈ F q(s)\in F to be the state that sequence s s leads to in the DFA starting from q 0 q_{0} and S ⁡ ( q ) ⊆ Σ ∗ S(q)\subseteq\Sigma^{*} to be the collection of all sequences that lead from state q 0 q_{0} to state q q in the DFA.

### 2.1 Recovering world models

Throughout this paper we assume that the ground-truth sequences used to train and test a generative model belong to the language of a deterministic finite state automaton W W . This generalizes past work (e.g., on assuming sequences come from legal moves in a game [ 36 , 20 ] ) and allows us to formally define world recovery.

###### Definition 2.1 .

A generative model m ⁡ ( ⋅ ) m(\cdot) recovers the DFA W W if ∀ q ∈ F , ∀ s ∈ S ⁡ ( q ) : L W ​ ( q ) = L m ​ ( s ) . \forall q\in F,\forall s\in S(q)\colon L^{W}(q)=L^{m}(s). That is, recovery requires that a sequence can be generated with positive probability by the model m ⁡ ( ⋅ ) m(\cdot) if and only if the sequence is valid in the DFA W W .

Recovery is defined at the language level. However, generative models are often built and evaluated token-by-token. It turns out that exact next-token prediction is enough for recovery of the language of the world model.

###### Definition 2.2 .

A generative model m ⁡ ( ⋅ ) m(\cdot) satisfies exact next-token prediction under the DFA W W if ∀ q ∈ F , ∀ s ∈ S ⁡ ( q ) , ∀ a ∈ Σ : m ⁡ ( a | s ) > 0 ⇔ δ ⁡ ( q , a ) ≠ q reject . \forall q\in F,\forall s\in S(q),\forall a\in\Sigma\colon m(a\penalty\ |\penalty\ s)>0\iff\delta(q,a)\neq q_{\text{reject}}.

###### Proposition 2.3 .

A generative model m ⁡ ( ⋅ ) m(\cdot) recovers the DFA W W if and only if it satisfies exact next-token prediction under the DFA W W .

Proposition 2.3 (proof given in Appendix A ) suggests a way to evaluate whether a generative model recovers the true DFA: assess the validity of next-token predictions. Existing world model diagnostics are motivated by this intuition; for example, one way that Toshniwal et al. [36] and Li et al. [20] assess world model recovery is by measuring the percent of top next-token predictions that are valid.

### 2.2 Next-token prediction is a fragile metric for recovering structure

Next-token prediction, however, is a limited evaluation metric. While exact next-token prediction implies perfect world model recovery, being very nearly correct on next-token prediction does not mean having very nearly recovered the world model. This can be illustrated by a simple example.

Example: Cumulative Connect-4. Consider a vertical grid with n n rows and 7 columns. Two players take turns dropping a disk in a column, and they can choose any column that contains less than n n disks. When a disk is dropped in a column, it occupies the bottom-most position that isn’t occupied by another disk, and it remains in that position for the full game. The game continues until the entire board is filled, for 7 ​ n 7n moves, regardless of whether a player has achieved four in a row. Games are represented as sequences of moves, where each sequence has 7 ​ n 7n tokens and each token is an integer between 1 and 7 indicating the column the disk is placed in. Here, Σ = { 1 , … , 7 } \Sigma=\{1,\dots,7\} denotes the columns and the state corresponds to the count in each column. A column is a valid move if that column is not already filled.

Consider a generative model that outputs { 1 , … , 7 } \{1,\dots,7\} with uniform probability given any sequence, i.e. m ⁡ ( a | s ) = m ⁡ ( a ′ | s ′ ) = 1 / 7 m(a\penalty\ |\penalty\ s)=m(a^{\prime}\penalty\ |\penalty\ s^{\prime})=1/7 for all a , a ′ ∈ Σ a,a^{\prime}\in\Sigma and s , s ′ ∈ Σ ∗ s,s^{\prime}\in\Sigma^{*} . This model clearly encodes no information about the board. However, for any board where there are no columns filled, this model provides a valid next move (e.g., the right panel of Figure 1 ), and so it will be a near-perfect next-token predictor when n n is large. For example, when n = 1000 n=1000 , it predicts a valid next move for more than 99% of all states. Metrics based on next-token prediction will imply this algorithm is close to recovering a world model.

### 2.3 The Myhill-Nerode interior and boundary