Transformers Learn Transition Dynamics when Trained to Predict M arkov Decision Processes (2024)

Abstract: Language models have displayed a wide array of capabilities, but the reason for their performance remains a topic of heated debate and investigation. Do these models simply recite the observed training data, or are they able to abstract away surface statistics and learn the underlying processes from which the data was generated? To investigate this question, we explore the capabilities of a GPT model in the context of Markov Decision Processes (MDPs), where the underlying transition dynamics and policies are not directly observed. The model is trained to predict the next state or action without any initial knowledge of the MDPs or the players’ policies. Despite this, we present evidence that the model develops emergent representations of the underlying parameters governing the MDPs.

## introduction recently, large language models (llms) have gained significant popularity and attention due to their versatility and performance, including in writing code, engaging in meaningful conversations, and much more. many of these models, trained on the simple principle of "predicting the next word," go on to become vastly capable polymaths. yet the reason behind how language models come to obtain this performance remains a subject of continuous debate and research. many have suggested, based on the extensive number of parameters of these language models, that their performance may result from merely memorizing "surface statistics," or external correlations that do not necessarily reflect the underlying data generation process. such issues can arise, for instance, when the pre-training corpora contains frequently co-occurring words, which can be preferred over the right answer @xcite . another instance in which a language model has been shown to learn causal statistical dependencies is due to dataset selection bias @xcite . it has also been suggested that language models can construct world models-interpretable and internal characterizations of the environment from which the data generating process is derived @xcite . recent works have shown that llms are able to develop internal representations of concepts such as color @xcite and direction @xcite . a standard way to evaluate the emergence of internal representations of the world state in these models is to assess them in a relatively wellbehaved, self-contained environment in which the rules are clearly stated and understood. to illustrate, @xcite have explored how such models, trained on sequences of chess moves, are able to predict valid chess moves with high accuracy. the authors also suggest that the model keeps track of the current board state for the prediction step. @xcite extended this idea by exploring the internal representations of a gpt-2 variant trained on the game of othello. however, previous works have only investigated how these models are able to internally identify the current state and stop short of demonstrating whether they are able to identify parameters governing the underlying data generation process. the main goal of this paper is to take a step towards filling in this gap in the context of markov decision processes (mdps), where the sequence of states and actions are generated by hidden, parameterized policies and transition dynamics. specifically, we consider the synthetic and wellunderstood game of connectfour for our investigation. first, we generate data in the form of game transcripts where the both players follow a policy guided by either deep q-learning (dql) or monte-carlo tree search (mcts). then, we train 3 transformer models each when the game transcript is represented using only the states (coordinates of the played pieces) or actions (which column the piece is placed in), hence totaling 12 transformers. next, we investigate whether the transformer models trained on the game transcripts contain an internal representation of the parameters governing the transition dynamics, which takes the form of either the players' deep q-values or mcts values. we verify whether the model is able to identify a salient representation when predicting the next state or action conditioned on the partial transcript seen thus far via probing-training classifiers to predict the deep q-values or mcts values of the current game state using the network's internal activations as input @xcite . using this probing technique, we find ample evidence of these models being able to internally represent the generative process despite changing the transition dynamics and representation of the input data. in summary, our contributions are twofold: 1) we show evidence that transformer models contain internal representations of the underlying transition dynamics governing markov decision processes after trained to predict the next tokenized state or action 2) we show that our result is robust to how this process is represented to the transformer model as input data and how the policy of the mdp is defined. ## dataset generation and language modeling we focus on investigating internal representations of language models in a well-understood, selfcontained synthetic game setting. this is motivated by the observation from past works that the language models learn to predict valid game moves by simply being trained to extend game transcripts @xcite . specifically, we select connectfour, a turn-based, twoplayer, board-completion game in which the goal is to connect four pieces of a player's own color. the connectfour environment is shown in figure 1 . in connectfour, the game is played on a 6 × 7 board where two agents place alternating pieces of red or yellow discs on the board, which fall down to the bottom-most unoccupied row of the column chosen by the agent. the objective for both agents is to connect four pieces of the same color before the opponent, whether horizontally, vertically, or diagonally. the agent who wins the game secures a terminal reward. we choose this environment for two reasons: first, connectfour has a game tree that is exponentially large, hence making it infeasible for any transformer model to brute-force "memorize" or "recite" the optimal game-play strategy for all outcomes; second, training deep reinforcement learning agents or growing monte-carlo search trees on connectfour to approximate optimal playing strategies have been shown to enjoy good performance @xcite . ## exploring internal representations as mentioned above, to see if our language model effectively captures the underlying transition dynamics of the markov decision processes, we use a standard tool called "probing," which is the process of training a classifier or regressor using the internal activations of a transformer model as input features to predict labels or values of interest. if we can train probes in all four settings (whether the policy is governed by mcts or deep q-learning and whether the mdp is represented using states or actions), then we can conclude that the transformer models effectively internally characterize information about the parameters governing the mdps. ## conclusion in summary, our study provides compelling evidence that transformer-based models, when trained on data generated from a markov decision processes, are able to develop internal representations of the underlying parameters governing these processes. our investigation, primarily focused on the game of connectfour, shows that these models are able to capture information about the players' policies and hence transition dynamics of the mdps, whether they are guided by deep q-learning or monte carlo tree search, and is robust to how the data is being fed as input to the transformer model. specifically, we show that 1) probes trained on the internal activations of our transformer models always outperform random probes in predicting the deep q-values or mcts values, which suggests that the model encode meaningful information about the mdp parameters 2) the superior performance of non-linear probes suggest that the internal representation of mdp dynamics may have a non-linear structure within the transformer models 3) the probes' ability to identify the best move using the embeddings further support this hypothesis that they capture salient features of the mdps 4) the robustness of these findings across different input representations and types of policy underscores the generality of our result. we hope this study contributes to the ongoing debate about the capabilities of language models, providing evidence that they can develop rich internal representations of underlying data-generating processes. as technological advancements continue to push the boundaries of what these models can achieve, understanding their internal mecha-nisms becomes increasingly crucial. we also wish to extend this work in the future to where there is even greater variability within the generative process, and consider alternative evaluation metrics to provide more insight along this line of research.