Transformers Learn Transition Dynamics when Trained to Predict Markov Decision Processes

Authors: yuxi  chen
suwei  ma
tony  dear
xu  chen
Venue: EMNLP_2024_Workshop_BlackBoxNLP
Published: 2024-08-15 20:02:31

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.

Main paper: 
 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 (Kang and Choi, 2023) .Another instance in which a language model has been shown to learn causal statistical dependencies is due to dataset selection bias (McMilin, 2022) .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 (Goldstein and Levinstein, 2024) . Recent works have shown that LLMs are able to develop internal representations of concepts such as color (Abdou et al., 2021) and direction (Patel and Pavlick, 2021) .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, Toshniwal et al. (2021) 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. Li et al. (2022) 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 (Alain and Bengio, 2016; Tenney et al., 2019) . 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 (Toshniwal et al., 2021; Li et al., 2022) . 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 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 (Alderton et al., 2019; Sheoran et al., 2022) .We describe below how game transcripts are generated using when players are trained on deep Qlearning or guided by MCTS in the ConnectFour environment to be used to autoregressively train our transformer models. Then, we give a highlevel overview of MDPs and its connection to our setting.To start, we use deep Q-learning (DQL) to train both players in the ConnectFour environment since traditional Q-learning often struggles to converge in when the space of outcomes is combinatorially large (Mnih et al., 2013) . We define a neural network parameterized by weights θ, which takes the current state s, and action a and outputs a scalar value Q θ (s, a). The state space S consists of all possible configurations of the 6 × 7 board, while the action space A is defined as placing a disc in the i-th column, where i ∈ {1, 2, . . . , 7}.Our training process is based on a variant of the original deep Q-learning algorithm. Specifically, the architecture of our network Q θ consists of one convolutional layer followed by two linear layers. The network predicts the Q-values for placing a disc in each of the seven columns. See algorithm 1 in Appendix A for further training details. We train nine pairs of RL agents, each pair competing against each other for one million games. For each game, we record the action (the column into which the piece is placed), the state (represented henceforth as coordinate of the played piece), and the deep Q-values of all the feasible moves at each step (if a move is infeasible, then the value is set to 0). Then, we combine the last 111K games played for every pair totaling one million game transcripts.We also generate data using MCTS. MCTS is a heuristic search algorithm that has shown remarkable success in classic board-games, modernboard games, and video games. MCTS combines depth-first search and stochastic simulation to build and use a game tree of possible outcomes based on selection, expansion, simulation, and backpropagation (Chaslot et al., 2008) . In our setting, we implement the standard MCTS algorithm where at each move decision, we run 100 rollouts, and the action with the highest MCTS value (win rate) is selected. Similar to the above, we generate one million game transcripts by running 1 million independent ConnectFour games where both agents play according to the MCTS heuristic. We also record the action, state, and corresponding MCTS values at each step (the value is likewise 0 if the move is infeasible).A Markov Decision Process (Ghavamzadeh et al., 2015) M is a tuple ⟨S, A, P, P 0 , R⟩ where S is the set of states, A is the set of actions, P (•|s, a) ∈ P(S) is the probability distribution over next states, conditioned on action a being taken in state s, P 0 ∈ P(S) is the probability distribution according to which the initial state is selected, and R(s, a) ∈ P(R) is a random variable representing the reward obtained when action a is taken in state s. A policy-a mapping from past observations to a distribution over the set of actions-is a rule for choosing actions at any given state. Policies can be characterized as 1. Markov if the distribution is only dependent on the last state of the observation sequence.2. Stationary if the distribution does not change over time.3. Deterministic if the probability distribution concentrates on a single action for all possible histories of states and actions.A Markov Decision Process is called first-order if the state transition probability P (s t+1 |H t = s 1 , a 1 , . . . , s t , a t ) = P (s t+1 |s t , a t ) depends only on the latest state and action and n-th order if P (s t+1 |H t ) = P (s t+1 |s 1 , a 1 , . . . , s t-n , a t-n ) depends on the last n states and actions.In the ConnectFour setting, we can define the components of the MDP as • S: All possible board configurations.• A: Valid column placements out of the 7 columns on the ConnectFour board.• P : Deterministic transitions based on player actions.• P 0 : The initial empty ConnectFour board state.• R: Reward based on game outcome, which can either be a win, loss, or draw.Having outlined the above, we see that players guided by deep Q-learning in ConnectFour follow a policy that is • Markov: The neural network only considers the current board state as input.• Non-stationary: The network's parameters are continuously updated during training, which means that their game-playing strategy can also evolve.• Non-deterministic: This is due to ϵ-greedy exploration, mini-batch sampling, and other sources of randomness during training.2 The resulting MDP is first-order, as the next state depends only on the current state and action.Players guided by MCTS, on the other hand, can be viewed in two ways. If we consider the Monte-Carlo Search Tree as part of the state, then the policy is• Markov: The current game tree and board state completely determine the distribution over the next action.• Non-stationary: The search tree is able to grow over time and output different moves.• Non-deterministic: This is due to the inherent stochastic nature of MCTS simulations.In this worldview, the MDP under MCTS remains first-order. However, if we regard the search tree as external to the state, then the policy and MDP become n-th order. This is because the search tree's evolution stochastically depends on all previous moves and simulations.In ConnectFour, given a chosen action, the next board state and terminal reward is deterministic. Therefore, the stochasticity in the MDP formulations of both deep Q-learning and MCTS are attributed only to randomness in the policy parameters. This means that a transformer model that internally characterizes deep Q-values or MCTS values with accuracy effectively captures the transition dynamics P (•|s, a). From here, we can conclude that such a model has a internal representation of the underlying parameters governing the MDP that generates the observed data. This insight guides our later experiments.For both settings, each state and action are tokenized as input. We supply no further auxiliary information during training, as our goal is to study how much they can infer the underlying transition dynamics from only information of the observed histories. Each history is treated as a sentence tokenized with a predefined vocabulary (for states, this corresponds to 42 possible coordinates of the discs; for actions, this corresponds to the 7 column placements; an extra padding vocabulary is included for both).For each setting (deep Q-Learning and MCTS) and each representation of the history (using state or action), we train three separate 8-layer GPT models (Radford et al., 2018) with an 8-head attention mechanism and a 512-dimensional hidden space. When we represent the history using only actions, we let the transformer predict state s t conditioned on the history {s 1 , . . . , s t-1 }. In the actionexclusive setting, we let the transformer predict action a t conditioned on {a 1 , . . . , a t-1 }. 3 The models' weights are initialized randomly, including the layer for word embeddings.Training is next performed autoregressively: for each tokenized partial history where each element is either a state or action, the forward process converts the input via the trainable word embedding into {c 0 t } T -1 t=1 , where c i t is the intermediate feature for the t-th token after the i-th layer to be sequentially processed by 8 multi-head attention layers. Using a causal mask, we ensure that only c i-1 ≤t are visible to c i t during training; that is, the prediction step only involves features in the preceding layer and earlier time steps. c 8 T -1 is lastly fed through a linear layer to predict logits for the ground-truth state or action. We use cross-entropy loss between the predicted logits and the ground-truth state or action as the objective during training. The parameters of the network are optimized by gradient descent, and we use the model weights corresponding to the epoch with the lowest validation loss to explore internal representations.4  
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.To train all probes, we first randomly sample one time stamp t in each game to obtain partial histories H t-1 = {s 1 , . . . , s t-1 } or H t-1 = {a 1 , . . . , a t-1 }. Then, we retrieve the corresponding internal embedding E t i of the network that is used to predict s t or a t when the input is H t-1 after the i-th layer of the network. We repeat this process of retrieving the embedding after every layer of the network, and obtain {E t i } 8 i=1 for each sampled partial history H t-1 . We repeat this process 679K times for each probe and split the dataset into training, validation, and testing data according to a 8-1-1 split. The embeddings after each layer are used to train separate probes, that is, we use {E t 1 } t to train the probe who uses embedding information output by the transformer after the first layer, {E t 2 } t to train the probe that uses the em-beddings after the second layer, totaling 8 probes for any particular combination of policy and data representation. We also repeat the probe training process 3 times for any setting corresponding to the 3 transformer models trained in each setting.To train the probes, we use the embedding E t i as input to regress against the true corresponding deep Q-values or MCTS values underlying the MDP at the time-step t -1 given the partial history H t-1 . For example, suppose at time t -1 that a player's MCTS values used to make the decision at time t are m t-1 = (0.2, 0.4, 0.7, 0.9, 0.1, 0.1, 0.1) corresponding to columns 1 through 7. The player would have chosen action a t = 4 or s t = (4, 2) 5 since the MCTS value corresponding to column 4 is highest. Then, we extract the embedding E t i associated with predicting a t or s t and use it to regress the 7-dimensional vector m t-1 . The parameters of each probe is optimized by gradient descent, and we select the model weights with the lowest validation loss to explore our hypothesis.Inspired by Li et al. (2022) , we also explore if the performances of linear and non-linear probes have significantly different accuracies, which may suggest how the parameters of the MDP are represented in the transformer model. In both settings, we compare probe performance trained on internal activations after each layer against a probe trained and validated on randomly generated embeddings. 6 It is clear that probes trained even on randomized embeddings may perform better than blindly "guessing" a random real-valued vector. 7 This approach allows us to see whether a random probe can encode information about the parameters of the MDP without any additional data as good as a properly trained probe. If the test loss between the two types of probes are indistinguishable, then this suggest that the transformers' internal activations do not contain any effective information of the MDP parameters.For linear probes, the prediction of the deep Qvalues or MCTS values parameterized by weights θ is given by W E t i where θ = {W ∈ R D×d }, D =5 Here, we suppose there already exists a disc beneath it played before, hence the current y-coordinate is 2.6 Each entry in the embedding is drawn independently from a normal distribution with mean of 0 and standard deviation of 5. We refer to these probes as "random probes" hereinafter for concision.7 Since even a network with random valued vectors as input can encode the empirical mean of the observed data. Then, if the distribution of the training and testing data are the same, we should expect to see that this network performs better than blindly guessing.512 is the number of dimensions of the internal embedding E t i and d = 7 is the dimension of the output space. For nonlinear probes, the prediction can be written asW 1 ReLU (W 2 E t i ), where θ = {W 1 ∈ R D×d , W 2 ∈ R D×D }.We verify the performance of our probes on 2 different metrics: 1) mean squared error between the predicted and ground truth deep Q-values or MCTS values of the moves 2) whether the best move predicted by a probe matches the ground-truth best move. We do not normalize any of the groundtruth or predicted values prior to evaluation.We first show the test MSE loss when the trained probes regress against the target deep Q-values or MCTS values of the player. The first column denotes the layer of the transformer model after which the embeddings are used to train the probe, where "R" stands for randomly generated embeddings. We report the mean and standard deviation of the test losses obtained from the three probes in each scenario in Tables 1 2 3 4 . We see that trained probes have a significantly lower test loss compared to random probes across all settings, which strongly suggests that the internal activations do contain representations of the MDP parameters. We also see that non-linear probes consistently yield lower losses than linear ones, which suggests that the MDP parameters may admit a non-linear representation in the transformer models. In addition, the difference in the scale between the DQL and MCTS settings can be easily explained: while MCTS values are bounded between 0 and 1, it is known that conventional Deep Q-learning is impacted by an overestimation bias (Hessel et al., 2017) . Nevertheless, our conclusion remains valid since all the probes in each setting are trained to regress against values generated from the same space.In terms of robustness to data representation, we see how the losses of the probes when the data is being represented using only states or using only actions do not differ significantly: the non-linear layers when data is represented using states performs slightly better than that using actions. This intuitively makes sense since encoding using states inherently provide more explicitly information (since they include the y-coordinate of the played discs) compared to actions. 3: MSE | MCTS | H t-1 = {s 1 , . . . , s t-1 }Non-Linear 1 0.0419 ± 1.0 e-7 0.0411 ± 4.0 e-6 2 0.0323 ± 1.7 e-4 0.0206 ± 2.1 e-5 3 0.0290 ± 1.9 e-4 0.0196 ± 5.6 e-5 4 0.0273 ± 1.6 e-4 0.0191 ± 8.8 e-5 5 0.0270 ± 1.6 e-4 0.0191 ± 1.0 e-4 6 0.0269 ± 1.6 e-4 0.0189 ± 4.7 e-5 7 0.0270 ± 1.2 e-4 0.0190 ± 5.9 e-5 8 0.0272 ± 1.3 e-4 0.0191 ± 3.7 e-5 R 1.4103 ± 1.0 e-7 1.4259 ± 1.1 e-4 Non-Linear 1 0.0420 ± 3.0 e-6 0.0416 ± 1.0 e-6 2 0.0341 ± 7.6 e-5 0.0269 ± 7.9 e-5 3 0.0309 ± 2.0 e-4 0.0224 ± 2.7 e-4 4 0.0297 ± 1.7 e-4 0.0205 ± 3.8 e-5 5 0.0285 ± 1.1 e-4 0.0202 ± 1.3 e-4 6 0.0275 ± 1.8 e-4 0.0196 ± 1.0 e-4 7 0.0267 ± 2.3 e-4 0.0195 ± 8.2 e-5 8 0.0261 ± 2.3 e-4 0.0195 ± 9.3 e-5 R 1.4103 ± 1.0 e-7 1.4259 ± 1.1 e-4Here, we would like to investigate whether the best move predicted by the probe matches the best ground-truth move. We define the loss function as1[Best Predicted Move ̸ = True Best Move]In other words, we want to see whether arg max i ṽi ̸ = arg max i v i for i ∈ {1, 2, . . . , 7} where ṽ, v ∈ R 7 are our predicted and ground-truth target deep Q-values or MCTS values respectively. We report the mean and standard deviation of the test losses across different settings in Tables 5 6 7 8 . Here, we observe that the performance of the trained probes significantly excel that of the random probes, meaning that the embeddings also contain internal information on how to make the best moves. 8In terms of robustness to data representation, we see how the data encoded using only actions yield a lower loss compared to that of states. This may be explained by how encoding the input data using actions is more directed towards identifying the best move, since the dimensionality of the space of actions and the space of best moves are identical and their structure may hence share greater similarity. Nevertheless, both ways of representing the input data to the transformer exceeds the performance of random probes.It should also be noted that mean-squared error and correctly identifying the best move are nots=1 p s-1 |{τ i } s i=1 ∩ {τ i } s i=1 | sThe output is bounded between 0 and 1 and captures how well the values of our predicted moves match those of the ground truth with regards to their orderings. It is clear that this metric, while not evaluating the numerical differences at each index, is able to preserve some notion of structural similarity between the predicted and ground-truth values. By varying p from close to 0 to close to 1, one is able to interpolate between putting emphasis on only the best move to virtually all the moves.We also remark that the choice of the evaluation metric may be highly problem-specific. For instance, one may resort to evaluation using the Kullback-Leibler (KL) divergence when the outputs are or can be normalized to probability distributions. We defer investigating alternative choices of metrics and their properties for future works. 
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. 
 
 Appendix: As mentioned above, to train our players using deep Q-Learning, we use Algorithm 1 shown below. The DQN architecture consists of a single convolutional layer followed by two fully connected layers and an output layer, designed to map 2D input states to Q-values for 7 possible actions. In our implementation of the network, we use a replay buffer of size 1000, and a mini-batch size of 32. We select actions using an epsilon-greedy policy, with ϵ = 0.1.We trained our transformers models on datasets consisting of 1 million games, with each game represented in four different forms: (1) sequences of states generated by deep Q-Learning, (2) sequences of states generated by MCTS, (3) sequences of actions generated by deep Q-Learning, and (4) sequences of actions generated by MCTS. Every dataset is split into 80% for training, 10% for validation, and 10% for testing.The transformer model we use have a block size of 42, an embedding dimension of 512, and 8 attention heads for a total of 8 layers with a predefined vocabulary size. The dropout rates are 0.1 for embedding dropout, 0.1 for residual dropout, and 0.1 for attention dropout. The model consists of an embedding layer, followed by a series of transformer blocks, each containing a causal self-attention mechanism and a feed-forward neural network. The final layers include layer normalization and a linear projection to the vocabulary size.For each of the four dataset representations, three transformers with identical architectures were trained to account for variability and potential error. This resulted in a total of 12 trained transformers. The models were optimized using Adam with a learning rate of 0.001, and training was conducted for 15 epochs with a batch size of 32. The loss function used is cross-entropy, calculated between the predicted logits and the true next tokens in the sequence.All training was performed on instances equipped with 8 NVIDIA RTX 4090 GPUs. x 0 ← empty board 5:for t = 0 : T do 1: Sample (x j , a j , r j , x j+1 ) from D 2: y j ← r j for terminal x j+1 r j + γ max a ′ Q(x j+1 , a ′ ; θ) for non-terminal x j+1 3: Gradient descent on loss (y j -Q(x j , a j ; θ)) 2However, if all these random factors are controlled, the policy becomes deterministic.As mentioned above, we want to explore both representations to see if the transformers' model's learning of the parameters of the MDP, if successful, is robust to how the input data is represented during training.See Appendix B for more training details.In addition, our probes are trained to minimize the MSE between the predicted and target values, not cross-entropy loss of the predicted and actual best move. This also implies that minimizing MSE can help partially achieve this functionality. 


Decision: Reviewers agree that the research question and experiments are clear and well-motivated. I encourage the authors to incorporate the reviewers' comments regarding the choice of metrics and extending their discussion in the camera-ready version.

Reviews: review: paper_summary: paper_summary: The paper explores the ability of language models to abstract beyond surface-level statistics and learn underlying processes by training a transformer on ConnectFour game transcripts, with the goal of investigating whether the model can encode the parameters governing the underlying process used by agents playing the game (e.g. Q-values in the case of an agent guided by deep Q-Learning or MCTS values in the case of an agent governed by Monte-Carlo tree search). The authors present an experimental setup where the model is trained to predict the next state in the context of Markov Decision Processes (MDPs), without direct observation of transition dynamics or player policies. The paper provides evidence suggesting that the models develop emergent representations of Q-values and MCTS values by showing lower mean squared error when prompted to predict these values compared to a randomly initialized model.
strength_weakness: strength_weakness: - MSE as a metric lacks interpretability and comparing against a random baseline is limited: While MSE is used as the primary evaluation metric, it may not be the most appropriate measure. MSE quantifies how far the model's predictions are from the ground truth but does not provide insight into whether the model captures the structure of decision-making or the policy. Also, the comparison to a random baseline is useful for establishing a lower bound, showing that the model is not making purely random predictions. However, this baseline alone is insufficient evidence to conclude that the model has learned Q-values. A random baseline merely demonstrates that the model has learned something, but it doesn’t reveal the depth of what has been learned. Additional comparisons against more meaningful baselines would provide stronger evidence of the model's ability to learn the underlying decision-making processes. 
Since MSE is sensitive to outliers, focusing on predicting only the best move (the action with the highest Q-value for instance) could artificially lower the MSE. This would allow the model to achieve a lower error compared to a random baseline, even if the other Q-values for suboptimal moves are incorrect or ignored. Essentially, the model could appear to perform well simply by predicting the next optimal move, without truly capturing the entire distribution of Q-values across all possible actions. However, it's important to note that this demonstrates the model has captured some level of information, which is interesting in itself.
questions: questions: -
paper_summary: paper_summary: This paper shows that Transformers learn representations that encode information about the transition dynamics and policies of Markov Decision Processes just by being trained to predict the next state or action. This is demonstrated by training autoregressive Transformers to play Connect Four and studying their internal representations with probing classifiers.
strength_weakness: strength_weakness: This paper was written very clearly, so I have no glaring concerns. However, I do wonder how big an MSE of e.g. 467 is for the action prediction (or how big an MSE of 0.0272 is for the state prediction). Are there more intuitive measures that can be presented (or alternatively, are there ways to make the MSE range more intuitive to readers)?
questions: questions: - Line 54, 59: You can use \citet for narrative citation
- Line 107: MDp $\rightarrow$ MDP
- Line 245: Remove whitespace before footnote
- Footnote 1: The full stop is missing
- Line 262: becomes $\rightarrow$ because (?)
- Section 2.3: Maybe some tokenized examples of the training data would be informative for readers
- Line 337: $\mathcal{H}_{t-1}$ for the sequence of actions
- Line 469: "probe strained" $\rightarrow$ "probes trained"
paper_summary: paper_summary: The paper trains deep Q-learners to play Connect Four, then trains a small GPT-style Transformer on the sequences of moves of such players, and then shows that MDP state and next action can be decoded linearly or nonlinearly with decent accuracy from the Transformer. The paper argues that this shows meaningful generalization by the Transformer, going beyond "surface statistics" in training data.
strength_weakness: strength_weakness: While the work seems to be solid, it does not meaningfully advance our knowledge or opinion about neural network capabilities. Essentially these same studies (albeit training networks to mimic human players, not trained MDPs as far as I know) have been done for example with Othello (cited by the authors, but there is a lot more), and indeed some knowledge of variables like game state can indeed be decoded from networks in previous work. That work triggered a number of interesting debates about the validity of using probes to determine whether a neural network "represents" something---see for example Melanie Mitchell's claims that the apparent ability to decode Othello states was in fact a property of the decoder rather than the underlying neural network. These debates are not engaged with here, nor does the study provide new information touching on those debates.
questions: questions: probe strained -> probes trained