CheckersGPT: Learning World Models through Language Modeling (2024)

Abstract: Although Large Language Models (LLMs) have been trained using just the next token prediction objective, these have shown impressive performance on various tasks. Consequently, it has attracted research interests in this regard. While one line of work in the past has suggested that LLMs learn surface-level statistics from the dataset, another line of work emphasizes that the learned representations are effective for simulating the underlying world model, considering the causal relationship for the next token prediction. This phenomenon is often referred to as the emergence of a world model in sequence prediction tasks. Recent work has demonstrated this phenomenon in a simulated setting of board games like Othello and Chess. In this paper, we analyze the game of Checkers to find out the emergence of a world model in a language model. By training a GPT-style autoregressive language model using only the next character prediction objective, we find that the model does show a hint of learning a world model representation of the board positions. We perform our analysis on two datasets: 1) synthetic dataset, which comes from the checkers game tree, and 2) human gameplay dataset. With multiple models trained with different layer sizes, we find that increasing the parameter size does help learn better world model representation decoded by linear probes.

heading: 1 Introduction
text: Though the Large Language Models (LLMs) have shown a remarkable ability to perform well on a wide range of tasks (Radford et al., 2019; Brown et al., 2020), the underlying process behind predicting the desired next token remains a black box. Since these language models are trained on a huge corpus of human-written text, it remains challenging to validate whether the network models the causal relationships or relies on spurious correlations occurring in the large corpus. Some initial
* Equal Contributions
findings suggest LLMs rely on surface-level statistics, stating LLMs are ‘stochastic parrots’ (Bender et al., 2021) and ‘causal parrots’ (Zečević et al., 2023). On the other hand, some of the recent work highlights LLMs learn feature representations that help create a proxy for an internal representation of the world. This property of the emergence of the ‘world model’ (Ha and Schmidhuber, 2018) in the learned representations has attracted significant research interest in recent years (Li et al., 2021; Toshniwal et al., 2022; Li et al., 2023; Nanda et al., 2023; Karvonen, 2024). The former line of work deals with model interpretations from an end-to-end perspective, whereas the latter heavily relies on understanding the complex properties via learned hidden representations. In this work, we focus on the latter part and analyze the learning process of GPT-style autoregressive models (Radford et al., 2018). Recent works in this line have shown these representations encoding complex properties present in the language (Bau et al., 2020; Goh et al., 2021; Geva et al., 2021; Burns et al., 2023; Elhage et al., 2022; Gurnee et al., 2023). More recently, researchers have explored the learning behavior using a simulated setting coming from a board game like Othello (Li et al., 2023; Nanda et al., 2023) and Chess (Toshniwal et al., 2022; Karvonen, 2024). We extend this setting to the game of Checkers (also see App. A) by training a GPT-style architecture (Radford et al., 2018) over the task of the next character prediction. We consider gameplay trajectories (sequence of moves in the game of checkers) coming from human gameplay as well as create a synthetic dataset using a random legal move in the game, i.e., considering the random sequences obtained from the game tree of checkers. A noteworthy feature of this training setup is that no information is provided to the network about the game of checkers, i.e., no legal moves feedback, no game design, and no information about the quality of the trajectory (more strategic or less strategic). Inter-
576
estingly, the training setup we follow does not even consider the next move prediction as the learning objective but rather the next character prediction (multiple characters form a move in the game), and the network only considers the sequence of characters as an instream to learn about the game. The closed setup of the game with definite board states helps facilitate simulating the underlying ‘world model’, i.e., the rules of the game decide the legal moves. We explore if the next character prediction tasks help learn useful representation forming the proxy for the board game, i.e., ‘world model’ for checkers. We train multiple networks with different layer sizes and observe the learning behavior across the training. Further, for analyzing the emergent behavior of the trained model, we consider using probes (Alain and Bengio, 2018; Belinkov, 2022) trained over latent representations learned by the model. Figure 1 provides a brief overview of the experimental setup. We observe that the deeper layers of the generative models tend to predict the game board state better, showcasing a piece of evidence for emergent behavior. In a nutshell, we make the following contributions:
• We create a simulated setting of decisionmaking tasks using the game of Checkers. Following previous works on OthelloGPT (Li et al., 2023; Nanda et al., 2023) and ChessGPT (Karvonen, 2024), this adds a new setting, fa-
cilitating the interpretability research. • We curate a human gameplay dataset and a
synthetic version of the dataset, containing 22,607 and 16 million gameplay trajectories, respectively. • We analyze the learned representations in detail and provide empirical evidence to support the emergent behavior of GPT-style autoregressive models. We release the codebase and the dataset for the experiments https://github.com/Exploration-Lab/ CheckersGPT

heading: 2 Related Work
text: The study of representations as a proxy for world models has been an active area of research interest in recent times. Li et al. (2021) study sequence generation models over synthetic tasks and report finding latent features encoding a proxy of world models. Some of the other works on similar lines highlight language models encoding ground concepts (Abdou et al., 2021; Patel and Pavlick, 2022; Burns et al., 2023), following previous works that study the linguistic features learned by sentence representations (Conneau et al., 2018; Tenney et al., 2019). Moreover, other recent works like McGrath et al. (2022); Lovering et al. (2022) analyze the representation learned by AlphaZero (Silver et al.,
2017) encoding chess concepts. Some initial works in the simulated setting of board games study the game of Chess for GPT-style architectures (Toshniwal et al., 2022), which was extended for the Othello game in Li et al. (2023); Nanda et al. (2023); Hazineh et al. (2023). More recently, Karvonen (2024) provides a detailed analysis by training a GPT-syle architecture on the Chess gameplay sequences.

heading: 3 Methodology
text: To study the learning behavior of GPT-style architecture trained on a next token prediction task, we create an experimental setup by considering the sequence of gameplay moves in checkers (see Figure 1). We consider a character as a token and train the model to predict the next character in the gameplay sequence strings. Note that in the gameplay string, multiple sets of characters define a move in the gameplay. We intentionally train the models in such a fashion so that it doesn’t provide any prior knowledge about the game of checkers (no game rules or board states). We further validate if such a training objective can learn a world model for the game of checkers while optimizing to predict the next token based on the previous tokens.

heading: 3.1 Dataset
text: To obtain the set of sequences coming from a Checkers game, we considered two sources: 1) Synthetic dataset, and 2) human gameplay dataset. The synthetic sequences are generated via uniformly sampling next moves from a Checkers game tree, whereas the Human gameplays are obtained from an online Checkers platform.1 Figure 2 shows the distribution of token lengths in the obtained human gameplay dataset. To have a similar distribution for the synthetic version, we clip the sequence in the range of 100 to 400. Previous approaches (Li
1https://www.fierz.ch/download.php
et al., 2023; Nanda et al., 2023; Karvonen, 2024) used a similar strategy to create synthetic and human gameplay versions of the datasets for board games like Othello and Chess. Table 1 compares the created Checkers game sequence resource with the existing works. The gameplay sequences in the created datasets are represented as a Portable Draught Notation (PDN) string. App. Figure 5 provides a detailed explanation of the gameplay sequence format.

heading: 3.2 Model Training
text: We focus on studying the learning behavior of the model trained on the next token prediction task. Rather than providing the task as the next move prediction in a sequence, we reduce the inductive bias by training the architecture on the next character prediction instead. Moreover, the model has no idea about the underlying mechanism from which these strings are generated. We use a vocabulary size of 17 characters to encapsulate the gameplay PDN string, which is “-./0123456789;x” white space “ ” and new line “\n.” Note that the squares (positions in the Checkers board) are also encoded in the string, reducing the inductive bias of encoding a single board position as a token. We consider a GPT-style architecture (Radford et al., 2018) for all of our experiments and train them in an autoregressive fashion to predict the next character in the sequence. For a sequence of N tokens t1, t2, . . . tN defining the gameplay trajectory. The model is trained to predict the token ti using all the previous tokens t1, t2, . . . ti−1. We employ the autoregressive style training using a causal mask and minimize the cross-entropy loss between the actual sequence token and the predicted logits. App. Table 2 provides the architecture-specific details for
the trained models. To notice the effect of parameter size, we use the same set of hyperparameters for training different variants of the same architecture, i.e., 1-layer, 2-layer, 4-layer, 8-layer. App. Figure 5 provides some qualitative examples of the input format provided to the model. For both datasets, we used a 99/1 train/test split to train the model in an autoregressive fashion.

heading: 3.3 Evaluation
text: We monitor the learning behavior of the models using various objectives. 1) Loss: We monitor the training/test losses across the training to validate if the model is able to minimize the learning objective. 2) Valid Moves Accuracy: Another metric that helps capture the learning behavior is the percentage of valid moves generated by the model. We consider the Checkers game engine to explore the game tree and validate the trajectory generated by the trained models. Note that the model predicts the game sequence at the character level, whereas the game tree validates the sequence of legal game moves. We sample N sequences from the trained model and check if the sequence contains strings processable by the game engine, i.e., all the generated sets of characters signify a gameplay sequence with all legal moves. We consider the generated string invalid even if one of the predicted characters is invalid. We report the percentage of valid strings from the sampled K strings. Probing Internal Model Representations: To gain a deeper insight into the latent representations learned by the model, we make use of the probing literature (Alain and Bengio, 2018). Probing provides a way to quantify the quality of learned representations (Belinkov, 2022). We consider the underlying world model representation of the game board state and train linear probes over the layer activations to predict the current board state. The linear probe for square s at layer l of the model is formulated as follows:
Ps,l = Softmax(Ws,l ·Al),
where Pi,l is the probability distribution over the 3 classes for square s, as predicted by the probe at layer l, Wi,l is the weight matrix for the linear probe associated with square i at layer l, Al is the 512-dimensional activation vector from layer l of the model. We trained separate linear probes for each model layer as well as each board square.
Each probe at layer l takes the input as the activation from the model’s lth layer.
We hypothesize that if a linear probe with little capacity can learn to predict the state of the board (which is not a linear function of the input), it demonstrates the capability of the model to transform the input into a linear representation of the board states in the latent activations. Essentially, the higher the accuracy of linear probes, the better the quality of learned representations, highlighting the learning of the underlying world model. Though the board contains 64 squares, only 32 are active board places where a piece can exist. We only consider the active 32 squares for reporting the probing accuracy. For each square, the probe is designed to classify it into one of the three possible states (Black Piece, White Piece, and Empty square). Recently, Nanda et al. (2023) reported the OthelloGPT representing (Mine, Yours, Empty) rather than (Black, White, Empty) in the form of linear representations. We follow a similar format for our experiments, where the probes predict the current game state relative to the current player at each turn, i.e., for odd turns, the model considers Black pieces as Mine and White pieces as Yours, and vice versa for even turns. Moreover, in our setup, since the model is trained on a next character prediction task rather than the next move prediction, multiple characters form a single move. For computing the probing accuracy, we consider the features corresponding to the last token before the next move. For white’s move, we consider the representations corresponding to ‘.’, i.e., the white move starts after <turn#>, and for black’s move, we consider features corresponding to <space> token before black’s moves (see App. Figure 5 for a reference PDN string). We perform a 50/50 train/test split for probing experiments and report the test accuracy.

heading: 4 Results and Discussion
text: We perform all the experiments by considering different layer-sized models over the synthetic and human gameplay datasets. Overall, we found the training done on human trajectories to be much more stable than the ones done on randomly sampled trajectories from the game tree, with the models trained on the synthetic version of the dataset having low valid moves accuracy. App. Table 3 and Table 4 show the number and percentage of generated legal moves, respectively. We observe that
the models with 4 and 8 layers learn better than the ones with a smaller number of layers and are able to generate sequences with 90% and 99% legal moves, respectively. However, we found that the model trained over the synthetic sequences generates a very low number of valid moves (in ranges of 100 even after 600K training steps), when compared to models trained on human gameplay sequences.
Probing Experiments: Figure 3 provides probing accuracy for various models with different layer sizes. We observe that for trained models, the probe accuracy for predicting board states using latent activations increases significantly compared to models with randomly initialized parameters. Moreover, we observe that the board states become more predictable for deeper layers in the larger network (8-layer), highlighting the quality of learned representations forming the underlying world model. App. Table 5 shows the probe accuracy for different layers of the models. Surprisingly, we observe that though the model trained on a synthetic dataset generates a very low number of valid moves, the latent representations learned by the model are good enough to predict the board state with a comparable probing accuracy. Note that, for validation of the model’s learning, we considered the metric of valid moves percentage, where all the trajectories with a single wrong character are ignored, making it a rigid evaluation scheme. Previous works like Karvonen (2024), have considered top-k% to report the accuracy of the valid move; however, in Checkers specifically, multiple sets of moves are possible with different numbers of characters, i.e., when taking a capture x move, the number of characters in a move change (in contrast to chess where the number of characters is fixed). Given the setting, it becomes difficult to process the validation
metric in such a fashion. App. Figure 6 shows the predictions obtained for the best model with 8 layers along with the corresponding ground truth board states.

heading: 5 Discussion
text: The notion of emergence argues that as the complexity of a system increases, new properties start to appear (Anderson, 1972). This idea of emergence has recently gained attention due to the increasing sizes of datasets as well as a number of parameters in LLMs, i.e., exhibiting ‘emergent abilities’ (Ganguli et al., 2022; Srivastava et al., 2023; Wei et al., 2022). Though there has been some evidence of these emergent abilities, Schaeffer et al. (2023) argues that the design of evaluation metrics plays a crucial role in observing the emergent behavior when increasing parameter size. In this work, we studied another variant of emergence, more specifically, if a simple objective of next character prediction can lead to learning the underlying ‘world models’, i.e., does the model learn the governing dynamics of the actual system (Checkers in our case) to satisfy the next token prediction objective. Though our findings point towards models learning the underlying ‘world models,’ it is to be noted that the evaluation measure we use may not be a good measure and only provides a weak signal. As recently highlighted by Vafa et al. (2024), a better evaluation metric is required for justifying the presence of implicit world model representations. Moreover, the settings we used had huge constraints; it would be interesting to consider a more open-ended setting encapsulated using a written text in the future.

heading: 6 Conclusion
text: In this work, we developed and studied a simulated setting, considering the board game of Checkers, and asked if a next token prediction objective in a GPT-style architecture helps learn the underlying world model for a task. With a detailed set of experiments over a varying range of models, we find that the deeper models (with about 25M parameters) show evidence of learning representations that encapsulate a proxy for the underlying world model. We believe this study will help facilitate future research on the emergent behavior of autoregressive models, by providing a firm playground for understanding representations learned during training over Checkers game sequences.
Limitations
We only consider the probing accuracy to get a hint of model learning ‘world models.’ For making the claim more concrete, previous works (Li et al., 2023; Nanda et al., 2023; Karvonen, 2024) have also proposed various intervention schemes applied over the trained probes, where they intervene over the learned representation to observe the changes in model’s predictions, highlighting the formation of world models. In the future, it would be interesting to try similar strategies for the proposed setting of checkers. Another major limitation of this work is the limited setting for the number of experiments; we only considered one architecture with different layer sizes to gain deeper insights into the training behavior of the models. In the future, it would be interesting to validate the same for a range of models with high parameter sizes. Due to the limited number of game trajectories available online as well as limited computing, we performed all the experiments in a small dataset/parameter size regime. It would be good to study the training process with a larger number of parameters, specifically for the synthetic dataset where we could not get the model to generate valid game trajectories.
Ethical Considerations
In this work, we focused on a simulated setting of the game of Checkers. To the best of our knowledge, the proposed setting and the model does not have any negative consequences on the society at large. Moreover, the proposed approach and model are restricted to research only; we do not advocate the usage of the model in real-life online Checker gameplay.

heading: A The Game of Checkers
text: Checkers2 is a board game played by two opponents on opposite sides of the game board. One player has dark pieces; the other has light pieces. A move consists of moving a piece forward to an adjacent unoccupied square. If the adjacent square contains an opponent’s piece, and the square immediately beyond it is vacant, the piece may be captured (and removed from the game) by jumping over it. Only the dark squares of the checkerboard are used, and a piece can only move forward into an unoccupied square. When capturing an opponent’s piece is possible, capturing is mandatory in most official rules. In almost all variants, a player with no valid move remaining loses. This occurs if the player has no pieces left or if all the player’s pieces are obstructed from moving by opponent pieces. The gameplay sequence is captured in the PDN format3, which is of the form “1. 10-15 22-18 2. 15x22 . . .”. Figure 5 describes the structure of the gameplay sequences in detail. The moves can be broadly classified into a normal move represented with "-" where a piece moves from one position to another, a capture move represented with "x" where an opponent piece is captured, and a multi-capture move "x x" where multiple pieces are captured. Figure 4 shows an example on the checker’s board.

heading: B Hyperparameters
text: The full set of hyperparameters is given in the Table 2. We trained the model on 4 NVIDIA A40 with a total training time of close to 14 hours. We used a 50/50 split to train the linear probes and a 99/1 split to train the model. We used the AdamW optimizer (Loshchilov and Hutter, 2019), with a maximum learning rate of 3e-4 and a minimum learning rate of 3e-5. We used the block size of 399 for our autoregressive task. We chose the model dimension of the transformer to be 512 with 8 heads and 8/4/2/1 layers. We trained the model for 600,000 iterations.
2https://en.wikipedia.org/wiki/Checkers 3https://en.wikipedia.org/wiki/Portable_
Draughts_Notation

heading: Future Work
text: We believe this study will help facilitate future research on the emergent behavior of autoregressive models, by providing a firm playground for understanding representations learned during training over Checkers game sequences. Limitations
We only consider the probing accuracy to get a hint of model learning ‘world models. ’ For making the claim more concrete, previous works (Li et al. , 2023; Nanda et al. , 2023; Karvonen, 2024) have also proposed various intervention schemes applied over the trained probes, where they intervene over the learned representation to observe the changes in model’s predictions, highlighting the formation of world models. In the future, it would be interesting to try similar strategies for the proposed setting of checkers. Another major limitation of this work is the limited setting for the number of experiments; we only considered one architecture with different layer sizes to gain deeper insights into the training behavior of the models. In the future, it would be interesting to validate the same for a range of models with high parameter sizes. Due to the limited number of game trajectories available online as well as limited computing, we performed all the experiments in a small dataset/parameter size regime. It would be good to study the training process with a larger number of parameters, specifically for the synthetic dataset where we could not get the model to generate valid game trajectories. Ethical Considerations
In this work, we focused on a simulated setting of the game of Checkers. To the best of our knowledge, the proposed setting and the model does not have any negative consequences on the society at large. Moreover, the proposed approach and model are restricted to research only; we do not advocate the usage of the model in real-life online Checker gameplay