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

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

Cumulative Connect-4 points to a general fragility in next-token prediction as an evaluation metric that can be understood in the context of the Myhill-Nerode theorem [ 26 , 27 ] , a classic result from language theory. The Myhill-Nerode theorem states that the sets of sequences accepted by a minimal DFA starting at two distinct states are distinct (see Appendix C for a full statement). More formally, for states q 1 ≠ q 2 q_{1}\neq q_{2} , we have L W ​ ( q 1 ) ≠ L W ​ ( q 2 ) L^{W}(q_{1})\neq L^{W}(q_{2}) . However, while distinct, the two sets may exhibit a great deal of overlap. Cumulative Connect-4 exhibits this behavior; any board for which there are less than k k disks in each column will have the same set of valid moves for the next n − k n-k moves. This intuition motivates a pair of definitions:

###### Definition 2.4 .

Given a DFA W W , the Myhill-Nerode interior for the pair q 1 , q 2 ∈ F q_{1},q_{2}\in F is the set of sequences accepted when starting at both states: MNI W ​ ( q 1 , q 2 ) = { s ∈ Σ ∗ | s ∈ L W ​ ( q 1 ) ∩ L W ​ ( q 2 ) } . \text{MNI}^{W}(q_{1},q_{2})=\{s\in\Sigma^{*}\penalty\ |\penalty\ s\in L^{W}(q_{1})\cap L^{W}(q_{2})\}. The Myhill-Nerode boundary is the set of minimal suffixes accepted by a DFA at q 1 q_{1} but not q 2 q_{2} : MNB W ( q 1 , q 2 ) = { s = a 1 a 2 … a k | s ∈ L W ( q 1 ) ∖ L W ( q 2 ) and ∀ j < k : a 1 … a j ∈ MNI W ( q 1 , q 2 ) } . \text{MNB}^{W}(q_{1},q_{2})=\{s=a_{1}a_{2}...a_{k}\penalty\ |\penalty\ s\in L^{W}(q_{1})\setminus L^{W}(q_{2})\text{ and }\forall j<k:a_{1}...a_{j}\in\text{MNI}^{W}(q_{1},q_{2})\}.

Figure 1 depicts an example Myhill-Nerode interior and boundary for cumulative Connect 4. Sequences on the interior are accepted by both states; it is only when we reach the boundary that these states will be distinguishable. Thus, models that pool together states with large interiors will perform well on next-token prediction tests; this is why the simple generative model succeeds in the cumulative Connect-4 example. To properly differentiate states, we must consider sequences that are long enough to be differentiated. In the remainder of the paper, we (i) use the Myhill-Nerode logic to develop new evaluation metrics and (ii) apply these to several applications.

### 2.4 Compression and distinction metrics for evaluating world models

We propose metrics to evaluate a model’s implicit world model by comparing the true Myhill-Nerode boundary to the one implied by the model.

###### Definition 2.5 .

For two sequences s 1 , s 2 s_{1},s_{2} , the Myhill-Nerode boundary implied by model m ⁡ ( ⋅ ) m(\cdot) is MNB m ( s 1 , s 2 ) = { x = x 1 … x k | x ∈ L m ( s 1 ) ∖ L m ( s 2 ) and ∀ j < k : x 1 … x j ∈ L m ( s 1 ) ∩ L m ( s 2 ) } . \text{MNB}^{m}(s_{1},s_{2})=\{x=x_{1}...x_{k}\penalty\ |\penalty\ x\in L^{m}(s_{1})\setminus L^{m}(s_{2})\text{ and }\forall j<k:x_{1}...x_{j}\in L^{m}(s_{1})\cap L^{m}(s_{2})\}. (1)

This is the set of minimal suffixes that are accepted by the model conditioned on s 1 s_{1} but not s 2 s_{2} . Since we now focus on the generative model rather than the DFA, the definition refers to pairs of sequences rather than to pairs of states.

Our evaluation metrics summarize how well a generative model identifies sequences that distinguish a given pair of states. Given a pair of states q 1 q_{1} and q 2 q_{2} , the metric is formed by first sampling sequences that lead to each state, s 1 ∈ S ⁡ ( q 1 ) s_{1}\in S(q_{1}) and s 2 ∈ S ⁡ ( q 2 ) s_{2}\in S(q_{2}) . We then calculate the true Myhill-Nerode boundary between the states and the model’s boundary between the sequences. Our metrics then compare the resulting boundaries using two statistics as building blocks:

###### Definition 2.6 .

The boundary recall of generative model m ⁡ ( ⋅ ) m(\cdot) with respect to a DFA W W is defined as | MNB W ​ ( q 1 , q 2 ) ∩ ( L m ​ ( s 1 ) ∖ L m ​ ( s 2 ) ) | | MNB W ​ ( q 1 , q 2 ) | , \frac{|\text{MNB}^{W}(q_{1},q_{2})\cap(L^{m}(s_{1})\setminus L^{m}(s_{2}))|}{|\text{MNB}^{W}(q_{1},q_{2})|}, (2) and the boundary precision is defined as | MNB m ​ ( s 1 , s 2 ) ∩ ( L W ​ ( q 1 ) ∖ L W ​ ( q 2 ) ) | | MNB m ​ ( s 1 , s 2 ) | . \frac{|\text{MNB}^{m}(s_{1},s_{2})\cap(L^{W}(q_{1})\setminus L^{W}(q_{2}))|}{|\text{MNB}^{m}(s_{1},s_{2})|}. (3)

Notice that boundary recall and boundary precision are not affected by whether the Myhill-Nerode interior is large between the two states. Returning to cumulative Connect-4, the simple generative model that outputs { 1 , … , 7 } \{1,\ldots,7\} with equal probability will perform poorly on these metrics; its recall will be 0 for all pairs of distinct states.

Based on the building blocks of recall and precision, we construct evaluation metrics to summarize whether the generative model correctly compresses sequences that arrive at the same state under the DFA and correctly distinguishes sequences that arrive at different states under the DFA. These two metrics correspond to different methods of sampling state pairs.

Sequence compression metric. To evaluate sequence compression, we sample equal state pairs q 1 = q 2 q_{1}=q_{2} . Since a DFA provides multiple ways to arrive at the same state, this test assesses whether a generative model recognizes that two sequences correspond to the same state. For example, in cumulative Connect-4, there may be multiple sequences that arrive at the same board position. Recall is undefined for equal states because there is no true boundary, so our compression metric only reports precision, averaged over states sampled uniformly at random (we say a generative model’s precision is 1 if its boundary is correctly empty).

Sequence distinction metric. To evaluate sequence distinction, we sample distinct state pairs, i.e. q 1 ≠ q 2 q_{1}\neq q_{2} . Here, there must be a true boundary, so we test how well a generative model recovers it. We report both precision and recall averaged over state pairs sampled uniformly at random.

Both metrics are depicted in Figure 2 . Although we have defined a generative model as accepting all sequences it assigns positive probability to, in practice sequence models are regularized to assign all sequences nonzero probability. Our evaluation metrics therefore depend on an acceptance threshold parameter ϵ > 0 \epsilon>0 . In practice, we explore sensitivity to different values of ϵ \epsilon and other acceptance mechanisms. We present ablations and other details in more depth in Section 3 and Appendix E .

## 3 Illustration: Do Transformers Recover the Street Map of New York City?

To illustrate these metrics, we create a dataset consisting of taxi rides in New York City. We process each ride into sequences of turn-by-turn directions and train transformers to predict the next direction. We show that transformers trained on these sequences have surprising route planning abilities: they not only find valid routes between two intersections but usually find the shortest path.

We then examine the underlying world model of the trained models. Despite the route planning capabilities of these models, our metrics reveal that their underlying world models are incoherent. Using a graph reconstruction technique, we show that each model’s implicit street map of New York City bears little resemblance to the actual map. Finally, we demonstrate that the route planning capabilities of these models break down when detours are introduced, a consequence of their incoherent world models.

### 3.1 Data and models

We base our analysis on a dataset of taxi rides released by the NYC Taxi & Limousine Commission, containing the latitude and longitude of each ride’s pickup and dropoff location in Manhattan. Each taxi ride obeys a true world model: the weighted graph corresponding to the system of intersections and streets in New York City. The graph is defined as G = ( V , E , W ) G=(V,E,W) , where V V is the set of intersections, E E the set of streets, and W : E → ℝ + W:E\to\mathbb{R}^{+} a weighting function containing the distance of each street. 2 2 2 A real-world intersection may be represented as multiple intersections here. For example, if a turn is only valid from one direction, it is represented as two different nodes. Each edge is labeled corresponding to its cardinal direction, represented as a function D : V × V → { □ , N , S , E , W , NE , NW , SE , SW } D:V\times V\rightarrow\{\square,\texttt{N},\texttt{S},\texttt{E},\texttt{W},\texttt{NE},\texttt{NW},\texttt{SE},\texttt{SW}\} with □ \square indicating that the edge does not exist. Each intersection has at most one edge in each direction. The graph has 4580 nodes (i.e. intersections) and 9846 edges (i.e. streets).

A traversal is a sequence of nodes where an edge exists between each consecutive node in the sequence. To study how the construction of traversals affects the resulting generative model, we consider three different approaches. Shortest paths constructs traversals by finding the shortest path between two nodes. Since these may not be reflective of real-world traversals due to traffic conditions, noisy shortest paths constructs multiple shortest paths by perturbing the magnitude of each edge weight in the underlying graph. Finally, random walks samples random traversals instead of approximating shortest paths. See Appendix F for details.

We convert each traversal into a sequence of directions. Each sequence begins with the origin and destination, followed by the cardinal directions in the traversal, and concludes with a special end-of-sequence token. Figure 5 gives an example of a set directions and the corresponding path. Since this language corresponds to a DFA W W with | V | 2 + 1 |V|^{2}+1 accept states, corresponding to all combinations of current intersection/destination intersection pairs and an additional end state, we can apply the evaluation metrics in Section 2.4 .

We randomly split data into train and test splits, ensuring no origin-destination pair is in both train and test sets. We include all sequences containing less than 100 directions. Our training sets consist of 2.9M sequences (120M tokens) for shortest paths; 31M sequences (1.7B tokens) for noisy shortest paths; and 91M sequences (4.7B tokens) for random walks. We train two types of transformers [ 38 ] from scratch using next-token prediction for each dataset: an 89.3M parameter model consisting of 12 layers, 768 hidden dimensions, and 12 heads; and a 1.5B parameter model consisting of 48 layers, 1600 hidden dimensions, and 25 heads. We follow the architecture of GPT-2 for each model [ 29 ] . We train models on 8 A100 GPUs. For each dataset, we analyze the model with the best held-out performance: the 89.3M parameter model for shortest paths, and the 1.5B parameter for noisy shortest paths and random walks.

### 3.2 Evaluating world models

To assess their capabilities, we first assess whether the trained models can recover the shortest paths between unseen (origin, destination) pairs. We prompt each model with (origin, destination) pairs from the test set and use greedy decoding to generate a set of directions. All models consistently generate valid traversals — between 96% and 99%. Impressively, 97% of the sequences generated by the shortest paths model are the true shortest path, and 94% of the sequences generated by the model trained on noisy shortest paths find a shortest path for one of the noisy graphs used to generate data. Figure 5 provides an example of a shortest path traversal.

To assess whether these capabilities correspond to coherent implicit world models, we first consider two existing diagnostics [ 36 , 20 ] . The next-token test assesses whether a model, when conditioned on each subsequence in the test set, predicts a legal turn for its top-1 predicted next-token. In our example, a directional move is legal if a street in the direction exists at the current intersection. Predicting the end token is only legal if the traversal implied by the sequence is at the listed destination. Meanwhile, the current-state probe trains a probe [ 11 ] from a transformer’s representation to predict the current intersection implied by the directions so far. We train a linear probe on a transformer’s last layer representation.

To implement the sequence compression metric, we randomly sample states (i.e., [intersection, destination] pairs) and two distinct traversals (i.e. prefixes) that arrive at each state. We then assess whether a model correctly admits the same suffixes for each prefix. We average over pairs of prefixes to report a score for each state and average over states to report a final score. To implement the sequence distinction metrics, we sample pairs of distinct states and traversals (i.e. prefixes) that arrive at each state, comparing the model’s approximate Myhill-Nerode boundary to the true one. We average over pairs of prefixes to report a score for each pair of states, and average over 1000 randomly sampled state pairs to report a final scores. Both metrics depend on a threshold parameter ϵ \epsilon : a prefix is only sampled or accepted if the model’s assigned probability for each token is above ϵ \epsilon . Here, we consider ϵ = 0.01 \epsilon=0.01 for all models and metrics. We describe implementation details, provide parameter ablations, and consider other acceptance rules (e.g. top-p and top-k) in Appendix E .

Table 1 summarizes our results. As references, we compare each trained transformer to a randomly initialized transformer baseline following Li et al. [20] as well as to the true world model. The three trained transformers perform exceptionally well on existing diagnostics; nearly 100% of next-token predictions are valid and the probe recovers the true intersection for more than 90% of examples. 3 3 3 While the next-token test accuracy is rounded to 100%, no model performs perfectly.

Our evaluation metrics, however, reveal that these existing diagnostics are incomplete. All trained transformers perform poorly on sequence compression, frequently failing to recognize that two prefixes leading to the same state should admit the same continuations. Even the transformer trained on random walks, which sees many distinct types of traversals during training, fails to compress prefixes for half the states. For the sequence distinction metrics, the transformers trained on shortest paths or noisy shortest paths perform poorly. In contrast, the transformer trained on random walks performs well on the sequence distinction metric. Both metrics are therefore valuable for evaluating world models; a model can perform well on one metric and poorly on the other. Here, a model that distinguishes separate states at a high rate fails to recognize that two prefixes that lead to the same state should have the same valid continuations.

### 3.3 Reconstructing implicit maps

Our evaluation metrics point to deficiencies in recovering world models. We now show that these metrics reveal underlying incoherence. In the maps setting, the state structure of the true world model is easy to interpret and visualize: it is defined by the map itself. We attempt to “reconstruct” the map implied by sequences sampled from each generative model.