C heckers GPT : 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.

## introduction though the large language models (llms) have shown a remarkable ability to perform well on a wide range of tasks @xcite , 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' @xcite and 'causal parrots' @xcite . 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' @xcite in the learned representations has attracted significant research interest in recent years @xcite @xcite @xcite . 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 @xcite . recent works in this line have shown these representations encoding complex properties present in the language @xcite @xcite @xcite . more recently, researchers have explored the learning behavior using a simulated setting coming from a board game like othello @xcite and chess @xcite . we extend this setting to the game of checkers (also see app. a) by training a gpt-style architecture @xcite 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- ## related work the study of representations as a proxy for world models has been an active area of research interest in recent times. @xcite 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 @xcite @xcite , following previous works that study the linguistic features learned by sentence representations @xcite ## methodology 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. ## results and discussion 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. 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. ## discussion the notion of emergence argues that as the complexity of a system increases, new properties start to appear @xcite . 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' @xcite @xcite . though there has been some evidence of these emergent abilities, @xcite 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 @xcite , 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. ## conclusion 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.