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

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

Reconstruction is an open-ended problem: the generative model produces directions between an origin and destination that do not necessarily correspond to a fixed graph over the intersections in Manhattan. To narrow the scope, our goal is to produce a visually interpretable reconstructed map. To that end, we fix the reconstructed graph to have the same set of vertices as the true world model, corresponding to intersections in Manhattan, and ensure that the reconstruction algorithm returns a map consistent with the true model whenever it is run on valid sequences. Further, (a) we enforce each node has at most one outgoing edge of any direction, (b) we limit the maximum degree of each node, and (c) we limit the Euclidean distance spanned by any edge. Altogether, our reconstruction algorithm gives the generative model the benefit of the doubt, attempting to reconstruct edges belonging to the true map until forced to do otherwise in order to map a generated sequence. The algorithm is detailed in Appendix B .

Figure 3 shows three reconstructed maps using sequences generated by the transformer trained on random walks. The sequences underlying each map are generated by randomly sampling 6400 (origin, destination) pairs and then sampling the model’s traversal for each pair ( Appendix G shows similar results for when the distribution of origin/destination pairs follows the sampling distribution used to train each model). On the left is the reconstructed map on only sequences which are valid under the true world model. On the right is the reconstructed map using the transformer’s sequences. The transformer’s underlying world model is incoherent; it recovers streets whose orientations are physically impossible (e.g. labeled NW but facing east) and require flyovers above other streets.

To show that this map is not the product of a model that has the right world model but makes a few transcription errors, we artificially corrupt sequences drawn from the true model. With probability equal to the probability of an error for the random walks transformer, we randomly re-label an edge in a sequence consistent with the world model. The middle panel of Figure 3 shows the reconstructed graph. It is much closer to the true world model than the transformer (which makes errors at the same rate). While these results are for random walks and one setting of graph reconstruction, Appendix G shows maps for the other models and different reconstruction settings. All settings recover incoherent underlying maps.

### 3.4 Implication of failing to recover the world model: detour fragility

Does it matter that the transformer has an incoherent world model? After all, it does very well at the practical task of finding shortest paths. Here we look at a slightly adjacent task and consider a driver facing detours while driving; how well does the model re-route?

Concretely, we feed each transformer an (origin, destination) pair from the test set and greedily decode a traversal. But with probability p p for each token, we add one of two kinds of detours: for “random detours”, the model’s proposed token is replaced with a randomly chosen (true) valid token; for “adversarial detours”, it is replaced with the model’s lowest ranked valid token. We always ensure a valid path to the destination exists (shorter than length 100) after each detour. Table 2 shows the fraction of valid traversals produced. While all models perform well initially, detours erode performance, illustrating how faulty world models can prove problematic. Notably, the transformer trained on random walks is more robust to detours than models trained on (noisy) shortest paths, mirroring its advantage on our proposed evaluation metrics. The similarity between model performance on these evaluation metrics and detour robustness illustrates the effectiveness of our proposed metrics for assessing world model recovery.

## 4 Other Applications: Othello and Logic Puzzles

We apply our evaluation metrics to two other settings: sequence models trained on games of Othello and large language models prompted to solve logic puzzles. In both cases, our framework finds the same type of incoherence that we found in the previous section for maps.

Othello. Li et al. [20] study the question of evaluating world models in the context of Othello, a board game that consists of players placing tokens on an 8x8 board. They train transformers on game transcripts to predict the next move of each game. They show these models perform well on both the next-token test and current-state probe considered in Section 2.4 . Since the true Othello game can be represented as a DFA, we can apply our world model evaluation metrics. The sequence compression metric assesses whether openings that lead to the same board position have the same predicted next moves, while the sequence distinction metrics assess whether the model can differentiate two distinct boards.

We apply our metrics to the two Othello sequence models considered by Li et al. [20] : one trained on real games from Othello championship tournaments and another trained on synthetic games. Table 3 in Appendix D reports the metrics in both settings. The model trained on real games performs poorly on both compression and distinction metrics, failing to group together most pairs of game openings that lead to the same board. In contrast, the model trained on synthetic games performs well on both metrics. This discernment is not captured by the existing metrics, which show both models performing similarly. We validate this discernment by performing a “detours” exercise for Othello in Table 4 in Appendix D ; while the model trained on synthetic data produces near-perfect games regardless of detours, the model trained on championship data fails immediately. Similar to the navigation setting, we again find that models trained on random/synthetic data recover more structure than those trained on real-world data.

Logic puzzles. We consider an additional application involving LLMs. Our metrics require that the ground truth language can be expressed as a DFA, so we consider a “seating arrangement” logic puzzle similar to those in Suzgun et al. [35] . There are n n seats and n n individuals. The vocabulary consists of statements like “Person ‘A’ is sitting in seat 1” and “Person ‘B’ is two seats away from Person ‘C”. A state is the set of seating arrangements that are consistent with all of the statements so far, and a statement is valid if it doesn’t contradict all arrangements in the given state.

We analyze Llama 2 (70B) [ 37 ] , Llama-3 (8B and 70B), Mixtral (8x22B Instruct) [ 15 ] , Qwen 1.5 Chat (72B and 110B) [ 2 ] , GPT-3.5 turbo, and GPT-4. We consider n n =3 individuals. We first assess whether the LLMs solve the logic puzzle task when the seating arrangement is fully specified by the statements. Figure 4 shows that most LLMs perform well at this task; GPT-4 is accurate on all examples. We then apply our metrics, assessing if each LLM compresses correctly (whether two sets of statements that lead to the same state lead to the same assessments) and has high recall for distinction (we do not compute precision for distinction because it is too expensive to approximate each LLM’s Myhill-Nerode boundary). See Appendix E for further discussion.

Figure 4 shows the results averaged over 100 samples. While most LLMs can solve the logic puzzle when it’s fully specified, they perform poorly on the compression and distinction metrics: no model has a compression precision higher than 40%. More than half the time a model is conditioned on two sequences with the same set of viable states, it asserts that different continuations are allowed for each sequence; see Figure 8 for an example. No model has distinction recall higher than 0.60. These results bring up an interesting point: LLMs can perform well at some logic tasks (such as when the seating arrangement is fully specified) without having a coherent world model.

## 5 Conclusion

In order to build high-fidelity algorithms that meaningfully capture the logic of the problems they model, we need ways to measure how close we are to that goal. This paper suggests theoretically grounded metrics for assessing the world models implicit inside generative models. Applications to maps, games, and logic puzzles suggest these metrics are both feasible to implement and insightful. Our results show that generative models can perform impressive tasks with incoherent world models (e.g. provide directions for taxi rides). But this incoherence makes them fragile for other tasks (e.g. providing directions when there are detours). This incoherence is also problematic when we hope to use a generative model to learn something latent about the world in scientific domains.

Our primary limitation is the focus on DFAs. While it is suitable for many applications like games, logic, and state tracking, extending it would be quite valuable, e.g. to situations where the underlying world model is more complicated than a DFA or is unknown. We suspect that the core ideas related to sequence compression and sequence distinction generalize to these richer settings, but leave that to future work.

## Acknowledgements

Keyon Vafa is supported by the Harvard Data Science Initiative. Justin Chen is supported by an NSF Graduate Research Fellowship under Grant No. 174530. Jon Kleinberg is supported in part by a Vannevar Bush Faculty Fellowship, a Simons Collaboration grant, and a grant from the MacArthur Foundation. We thank the Chicago Booth School of Business for generous support. We thank Foundry 4 4 4 https://www.mlfoundry.com/ for providing the compute required to conduct this research. We also thank Sarah Bentley, Juan Carlos Perdomo, and Neekon Vafa for helpful comments and feedback.

## Appendix A Proof

###### Proof of Proposition 2.3 .

We will first prove the forward direction that if a generative model m ⁡ ( ⋅ ) m(\cdot) recovers the world model DFA W W , then m ⁡ ( ⋅ ) m(\cdot) satisfies exact next-token prediction. By assumption, L W ​ ( q ) = L m ​ ( s ) L^{W}(q)=L^{m}(s) . Consider any state q ∈ F q\in F and sequence reaching that state s ∈ S ⁡ ( q ) s\in S(q) . This implies that for any sequence s ​ a sa for any character a ∈ Σ a\in\Sigma , δ ⁡ ( q , a ) ≠ q reject ⇔ m ⁡ ( a | s ) > 0 , \delta(q,a)\neq q_{\text{reject}}\iff m(a\penalty\ |\penalty\ s)>0, which is the definition of next-token prediction.

Now, we will prove the backwards direction that if m ⁡ ( ⋅ ) m(\cdot) achieves exact next-token prediction, it recovers W W . Fix any state q q and sequence s ∈ S ⁡ ( q ) s\in S(q) . Consider any sequence a 1 ​ a 2 . . a k ∈ Σ ∗ a_{1}a_{2}..a_{k}\in\Sigma^{*} and let q ′ = δ ^ ( q , a 1 a 2 . . a k ) q^{\prime}=\hat{\delta}(q,a_{1}a_{2}..a_{k}) be the state reached by following the sequence from q q . Note that by definition of q reject q_{\text{reject}} not having any outgoing transitions, q ′ ≠ q reject ⇔ ( ∀ j < k ) : δ ^ ( q , a 1 a 2 . . a j ) ≠ q reject . q^{\prime}\neq q_{\text{reject}}\iff(\forall j<k):\hat{\delta}(q,a_{1}a_{2}..a_{j})\neq q_{\text{reject}}. By assumption, m ( a k | a 1 a 2 . . a k − 1 ) > 0 ⇔ q ≠ q reject . m(a_{k}\penalty\ |\penalty\ a_{1}a_{2}..a_{k-1})>0\iff q\neq q_{\text{reject}}. If a 1 ​ a 2 ​ … ​ a k ∈ L W ​ ( q ) a_{1}a_{2}...a_{k}\in L^{W}(q) , then q ′ ≠ q reject q^{\prime}\neq q_{\text{reject}} . It then must be the case that m ( a j | a 1 a 2 . . a j − 1 ) > 0 m(a_{j}|a_{1}a_{2}..a_{j-1})>0 for all j < k j<k , implying that a 1 ​ a 2 ​ … ​ a k ∈ L m ​ ( s ) a_{1}a_{2}...a_{k}\in L^{m}(s) .

Conversely, if a 1 ​ a 2 ​ … ​ a k ∉ L W ​ ( q ) a_{1}a_{2}...a_{k}\notin L^{W}(q) , then q ′ = q reject q^{\prime}=q_{\text{reject}} . It follows that m ( a k | a 1 a 2 . . a k − 1 ) = 0 m(a_{k}|a_{1}a_{2}..a_{k-1})=0 , and thus a 1 ​ a 2 ​ … ​ a k ∉ L m ​ ( s ) a_{1}a_{2}...a_{k}\notin L^{m}(s) . ∎

## Appendix B Reconstructed Maps

In this section, we give more details on our reconstruction algorithm and display maps for sequences generated from models trained on shortest paths, noisy shortest paths, and random walks each for several parameter settings.

### B.1 Algorithm

Our reconstruction algorithm in Algorithm 1 takes a set of sequences from a generative model and attempts to reconstruct the underlying map implied by the sequences. The reconstructed map has the same set of vertices as the true world model (i.e. the intersections of Manhattan), and we visualize each intersection by placing them at their real-world latitude/longitude. The reconstruction algorithm thus attempts to recover the edges and edge directions implied by each set of sequences.

For each sequence that is valid under the true map, the algorithm adds the true edges and edge directions implied by the sequence. In a sense, this algorithm gives the generative model the benefit of the doubt. However, sometimes the model errs, i.e. it produces a sequence which is only valid by adding an additional edge to the true map. In this case, our reconstruction algorithm adds a new edge at the first invalid step of the traversal. There are usually multiple possible edges to add that would make the traversal valid; the algorithm adds the edge that maximizes the number of subsequent steps that would be valid. In other words, the algorithm adds new edges only when there is a discrepancy between the transformer and the world model, choosing the edge which greedily maximizes the number of subsequent steps that accord with the true world model. The direction label for each added edge is thus the token at the corresponding invalid step of the traversal.

Our algorithm may fail to reconstruct some of the input sequences as they cannot be made consistent with the partial graph we have reconstructed so far. Note that any reconstruction algorithm that satisfies the constraints in the main text of having a single edge with a given direction coming out of any intersection, a maximum degree, and a maximum edge distance will not be able to reconstruct every set of sequences. For example, two sequences [ v 1 , v 2 , E , end ] [v_{1},v_{2},\texttt{E},\texttt{end}] and [ v 1 , v 3 , E , end ] [v_{1},v_{3},\texttt{E},\texttt{end}] cannot be reconstructed without violating the first constraint.

### B.2 Maps

We include reconstructed maps built from 6400 6400 transformer-generated sequences for transformers trained on shortest paths, noisy shortest paths (modeling traffic), and random walks in Manhattan. Each sequence is generated by randomly sampling an (origin, destination) pair and then sampling the model’s traversal for each pair. Because the transformer never sees sequences of length more than 100 during training, we only sample pairs for which there exists a valid traversal in less than 100 moves. We note that this distribution of (origin, destination) pairs varies from the distributions used to train and evaluate each model, and the percent of traversals that are valid falls from >95% to 65-80%; however, Figure 12 shows similar maps when each map is constructed by sampling (origin, destination) pairs from the training/evaluation distributions. We vary the constrained maximum degree between 4 4 and 8 8 (the true map has maximum degree 4 4 and there are 8 8 possible cardinal directions), and the maximum edge distance between 1 / 2 \nicefrac{{1}}{{2}} and 1 1 mile.