Sparse Autoencoders Reveal Temporal Difference Learning in Large Language Models
Authors: Can Demircan
Tankred Saanum
Akshay Kumar Jagadish
Marcel Binz
Eric Schulz
Venue: ICLR 2025 Poster (2025)
Keywords: reinforcement learning
in-context learning
representation learning
sparse autoencoders (SAEs)
large language models (LLMs)
Primary area: interpretability and explainable AI
TL;DR: SAEs uncover a temporal-difference learning algorithm used by Llama for in-context reinforcement learning.

Abstract: In-context learning, the ability to adapt based on a few examples in the input prompt, is a ubiquitous feature of large language models (LLMs). However, as LLMs' in-context learning abilities continue to improve, understanding this phenomenon mechanistically becomes increasingly important. In particular, it is not well-understood how LLMs learn to solve specific classes of problems, such as reinforcement learning (RL) problems, in-context. Through three different tasks, we first show that Llama $3$ $70$B can solve simple RL problems in-context. We then analyze the residual stream of Llama using Sparse Autoencoders (SAEs) and find representations that closely match temporal difference (TD) errors. Notably, these representations emerge despite the model only being trained to predict the next token. We verify that these representations are indeed causally involved in the computation of TD errors and $Q$-values by performing carefully designed interventions on them. Taken together, our work establishes a methodology for studying and manipulating in-context learning with SAEs, paving the way for a more mechanistic understanding.

Published as a conference paper at ICLR 2025
SPARSE AUTOENCODERS REVEAL TEMPORAL DIF- FERENCE LEARNING IN LARGE LANGUAGE MODELS
Can Demircan #,1
Tankred Saanum
can roanghe Imoo 2--mun .he
tankred, saanum@t wewinggneeennnnnn..mmmm.. wareiiinnnn........ gen.mpg......e
Akshay K. Jagadish 1.2
Marcel Binz²
Eric Schulz
ABSTRACT
In-context learning, the ability adapt based few examples in the input prompt, is a ubiquitous feature of large language models (LLMs). However, as LLMs' in-context learning abilities continue to improve, understanding this phe- nomenon mechanistically becomes increasingly important. In particular, it is not well-understood how LLMs learn solve specific classes of problems, such as reinforcement learning (RL) problems, in-context. Through three different tasks, we first show that Llama 3 70B can solve simple RL problems in-context. We then analyze the residual stream of Llama using Sparse Autoencoders (SAFs) and find representations that closely match temporal difference (TD) errors. Notably, these representations emerge despite the model only being trained to predict the next token. We verify that these representations are indeed causally involved the computation of TD errors and Q-values by performing carefully designed in- terventions on them. Taken together, our work establishes a methodology for studying and manipulating in-context learning with SAEs, paving the way for more mechanistic understanding.
INTRODUCTION
Large language models (LLMs) pretrained on large-scale text corpora are proficient in-context learn- ers. They can predict and learn novel rules and functions conditioned only on a text prompt with a few examples. This phenomenon has been demonstrated under many different paradigms, such as translation (Brown et al., 2020; Wei et al., 2022). function learning and sequence prediction (Garg al., 2022; Coda-Forno al., 2023), and even reinforcement learning (Binz & Schulz, 2023; Shinn et al., 2024; Brooks et al., 2022; Schubert et al., 2024; Hayes et al., 2024). Consequently, under- standing how in-context learning is implemented mechanistically becomes increasingly important. In this paper, we investigate the mechanisms underlying in-context reinforcement learning (RL) in LLMs. specifically how they can learn to generate actions that maximize future discounted rewards through trial and error, given only scalar reward signal.
In RI., the central learning signal the temporal difference (TD) error (Sutton, 1988). The TD error is the difference between the agent's 's belief about state-action pair's 's value, and a target value, which is constructed as the immediate reward plus the discounted value of the successor state. For instance, when the agent receives an unexpected reward, unexpectedly transitions to high-value successor state, the corresponding TD error can be used to update beliefs about the value of the current state. This learning rule forms the foundation of many RL algorithms, such Q-learning (Mnih et al., 2015; Watkins & Dayan, 1992) and Critic (Konda & Tsitsiklis, 1999; Sutton et al., 1999; Haarnoja et al., 2018). TD learning has also played an important role in neuroscience, where several experiments have shown links between TD-like computations and dopamine in multiple species
* Equal contribution.
Instittte for Human-Centered AI, Helmholtz Computational Health Center, Munich, Germany Max Planck Institute for Biological Cyberneties, Tübingen, Germany

---

Published as a conference paper at ICLR 2025
B
8
KING
Sparse Autcencoder (SAE)
Compare Identily
Prompt
SOE Acttllty
You eccions, _ofi and Right. Your goel is to get as many points as possible with an epieode Here gX%
TO Enor
Episode book Left. brange You pixk Right. You get 5 points. Episade over.
T3 Lesion
Reward
Esisode 1. You Left, You ai orange. piok get Episode is over.
Episode
We study the mechanisms of in-context learning in three different tasks: (A) The Two-Step Task, (B) the Grid World task, and (C) the graph prediction task. (D) Example pipeline for the Two- Step Task. We prompt Llama as shown on the left. As selects its actions, we record the internal representation for the tokens that precede the actions, which are highlighted in orange. We train SAES on these representations (middle) and correlate with the learned latents of the SAEs against the TD еттоr signals and other variables of interest we obtain from reinforcement learning agents (right). After identifying such latents, we lesion them and replace Llama's internal representations with reconstructions following the lesions (middle). We then test whether the lesion created the expected effects in the behavior (right).
Figure 1:
(Montague et al., 1996; Schultz et al., 1997; Fiorillo et al., 2003; Flagel et al., 2011; Niv et al., 2015).
TD learning is a fundamental algorithm that has been successful modeling both machine and an- imal learning. Could LLMs learn implement TD learning in-context simply from being trained to perform next-token prediction? We use sparse autoencoders (SAEs) to analyze the representa- tions supporting in-context learning RL settings. These models have successfully been used to build mechanistic understanding of neural networks and their representations (Bricken et al., 2023; Templeton al., 2024; Gao al., 2024; Lieberum et al., 2024). Across several experiments, we establish methodology for systematically studying and manipulating in-context RL in Llama 3 70B! (Dubey al., 2024) with SAES (see Figure 1). Using our methodology, we uncover represen- tations similar TD eπrоrs and Q-values in multiple tasks. Moreover, find that manipulating these features changes Llama's behavior and representations in predictable ways. We believe our study paves the way for mechanistic understanding of in-context learning, setting precedent for investigating how LLMss solve other types of learning problems in context.
From here onwards, we refer to this model as Illaaa
2I "Lesioning refers setting the activations of specific units to O. This is also commonly referred to zero ablation in the literature (Heimersheim & Nanda, 2024).
2

---

Published as a conference paper at ICLR 2025
2 METHODS
2.1 REINFORCEMENT LEARNING
We investigate in-context learning in the setting of Markov Decision Processes (MDPs). An MDP consists of state space S, an action space A. and transition dynamics 8t+1 T( a,), defining the probability distribution of successor states given the current state and action, well as a reward function r(8,, (t) that maps state-action pairs to a scalar reward term. The goal of the agent is to learn policy (s() that maximizes future discounted rewards, e.g. Q-values
Qxo (st, a1) = Ems Ei_o 7'r(s;, at) where is discount factor.
canonical algorithm for learning the Q-value function for fixed policy ma, is TD learning. From an initial estimate of the Q-values, TD learning bootstraps sequence of estimates that converge to the true values. This is done by minimizing TD errors, the difference between the agent's previous estimated Q-values, and the immediate reward plus the discounted Q-value of the successor state. Assuming the policy always picks the action that maximizes Q-values, the temporal difference error and subsequent Q-value update can be written as follows:
&TD = r(st, + max Q(s++1, - Q(s;,
(1)
Q(s,, a.) - Q( a.) + a ôTD
(2)
For subsequent analyses, we rely Q-learning which is variant of TD learning that learns value function for the optimal policy. See Appendix A.1 for details surrounding the implementation the Q-learning model.
2.2 SPARSE AUTOENCODERS
A natural place to look for the traces of TD errors and Q-values is in the residual streams of the Llama's transformer blocks (Templeton al., 2024; Bricken et al., 2023; Gurnee et al., 2023). The residual stream of a transformer block carries information from the input and all preceding blocks, making it viable candidate for initial exploration.
Directly searching for features that correspond to Q-values or TD errors in the residual stream, while possible, is often impractical. In Llama, each token is represented vector in the residual stream. Due to the high number of statistical comparisons such an approach would necessitate, analyses showing connections between activations and TD errors will suffer from reduced power. Furthermore, it has been shown that LLMs represent concepts in an entangled and distributed manner. This phenomenon, known as polysemanticity (Elhage et al., 2022), occurs when single semantic concepts are represented across multiple neurons. Using tools from the field of mechanistic interpretability (Olah, 2023), we seek to learn a disentangled, monosemantic, and potentially lower dimensional set of features from the residual stream. We hypothesize that features corresponding to TD errors or action values can be found in this oow-dimensional disentangled latent space.
8192-dimensional
One popular model for learning such features is the Sparse Autoencoder (SAE) architecture. SAEs trained with gradient descent reconstruct the potentially entangled features learned by models like transformers through a linear combination of sparse latent features. SAEs have been shown to learn more easily interpretable and monosemantie features (Cunningham al., 2023; Bricken al., 2023). We used SAEs with a single dense encoder layer (f), followed by ReLU non-linearity (a) and single dense decoder layer (g). We train SAEs reconstruct vectors from the residual stream h while minimizing the L1 norm of the encoded features:
CSAE = |hh h113 + (f())
(3)
where h = gf)(). We trained separate SAEs for each of the 80 blocks of Llama and each experiment. Further training details are provided under Appendix A.2.
3

---

Published as a conference paper at ICLR 2025
Values
D
Errors
7
3000
Max Correlation
1.00
Correlatoon
1.00
70B
Myopic
Myopic
Return
6
-logpla)
2750
0.75
Olcarring
0.75
Q-leaming
2250
0.25
0.25
3
2000
0.00
0.00
10
20
Myopic
0
25
50
75
25
50
Episode
learning
Block
Block
E
G
Values
H
Errors
3000
Max Correlation
1.00
Correlation
1.00
No lesion
No losion
No lcsion
Return
6
Lesion
pla)
2750
0.75
TD lesion
0.75
2500
0.50
0.50
2250
0.25
0.25
3
2000
0.00
0.00
10
20
30
TD No Control lesion lesion lesion
0
25
50
75
0
25
50
75
Episode
Block
Block
Figure 2: Llama relies on TD-like features solve RL tasks in-context. (A) Llama 70B often learns the optimal policy in the Two-Step Task through trial and error, whereas the smaller 8B counterpart does not improve beyond chance level. Shaded regions show standard error of the mean. (B) Llama's behavior is best described by a Q-learning algorithm. (C & D) SAE features with significant correlations to both reward estimates (myopic values) and Q-value estimates, as well as temporal difference errors, appear gradually through the transformer blocks. (E and F) Deactivating single TD feature Llama is sufficient to impair performance and make behavior less consistent with Q-learning. (G& & H) Negatively scaling the TD feature decreases subsequent representations' similarity to Q-values and TD errors.
3 LLAMA USES TD FEATURES TO LEARN POLICIES IN-CONTEXT
We first sought assess Llama's ability solve RL problem purely in-context. We designed a simple MDP inspired by the Two-Step Task (Daw et al., 2011; Kool et al., 2016). In this task, Llama first had choose to go Left or Right to enter either the Apple state the Orange state, respectively. This first state-transition was always awarded with 0 rewards. From there, Llama again had to choose Left or Right enter the terminal state and receive a reward. The reward and task structure are visualized in Figure 1A. Rewards and transition dynamics were deterministic. We instructed Llama to maximize reward and encouraged experiment (see Appendix A.3 for the exact prompt). When entering a new state, we sampled an action from Llama's 'ss predictive distribution of the possible action tokens with a temperature of 1. Llama completed 100 independent experiments initialized with unique seeds, each consisting of 30 episodes. We sampled actions from a random policy in the first 7 episodes to ease the exploration problem.
explore more at the beginning of the
We first observe that Llama achieved significantly higher returns than chance in the Two-Step Task, and often learned the optimal policy (see Figure 2A). The smaller 8B parameter version of Llama performed roughly at chance level, suggesting that the ability to learn policies through RI. emerges with scale. As a result, focus the 70B parameter version for subsequent analyses. Next, fitted behavioral models to the sequences of actions performed by Llama. We fitted Q-learning model that updated Q-values using TD errors, a myopic Q-leaning model that only cared about immediate rewards?, which we implemented by setting the discount parameter ? of the Q-lennning model to 0, as well as a repetition model that computed choice probabilities as a function of how often an action was picked in given state in the past. We see that Llama's 's behavior is best captured by a Q-learning model, attaining a negative log-likehood (NLL) score of 2729, outperforming the myopic model with an NLL of 2864. This indicates that it integrates information about future
In the animal learning literature, this is referred to as the Kascorla-Wagnee model (Rescorla & Wagner, 1972).
4

---

Published as a conference paper at ICLR 2025
Run
Run
Run
Run
Run
Run 6
Standard
M
0.5
0.0
1.0
A
Moni
NW
Reward disruption
0.5
1.0
Transition disruption
30
60
60
60
Step
TD SAE fealurre
Figure 3: Qualitative comparison of the TD error and the best matching SAE feature from block 34 for three separate runs in all three variations of the Two-Step Task. The SAE feature shows similar jumps TD when Llama encounters surprising events, such transitioning unexpected states receiving an unexpected reward. The dashed line indicates the onset of the change in either transition dynamics or reward function.
discounted rewards to select actions (see Figure 2B). The repetition model, which predicted that Llama simply repeated action patterns in the prompt, fits behavior worse than chance with an NLL of 5745.
Next, we trained SAEs on the residual stream outputs from all of Llama's transformer blocks. To create diverse training set for the SAEs. collected representations from Llama on two additional variations of the Two-Step Task where either the reward function or the transition function changed in the middle of the experiment. Since unexpected changes either transition dynamics or reward function will be followed by large positive or negative TD errors, we expected it to be useful for detecting possible TD errors in Llama's activations. We trained each SAE using regularization strength 3 le 05 for 30 epochs 18000 residual stream representations. In the resulting SAEs, we find features gradually developing over transformer blocks that track both the reward function (captured in the myopic values) well Q-values (see Figure 2C).
Furthermore, we observe features that show significant correlation with TD errors (see Figure 2D). In block 34, we find the SAE feature with the highest correlation (r 0.58) with the TD error of our Q-learnin model. We visualize the normalized SAE feature overlaid with the normalized TD error in Figure for six different runs and the three variations of the Two-Step Task. The feature shows the same qualitative change as the TD error when the reward function and transition dynamics change. We refer these types of latent features that resemble TD errors TD latents. To establish whether this TD latent has a causal connection to RL behavior, we tested Llama on the Two-Step Task with the TD latent either deactivated or negatively scaled. Before each choice, we replaced the residual stream activations in block 34 with reconstructions from the trained SAE in which the TD latent was set
scaled. We see that simply deactivating the TD latent led to significantly worse performance in the task (see Figure 2E), producing behavior that differed more from Q-learning agent (Figure 2F). In contrast, deactivating the SAE feature block 34 with the lowest correlation to TD did not produce behavior that was less likely according to the Q-learning model. Finally, we analyzed correlations between SAE features and Q-values and TD errors after we seale the TD latent by value of -10, similar to Templeton et al. (2024). Here we the correlation strengths for both Q-values and TD errors (see Figure 2G H). In sum, we found an SAE feature showing the characteristics of TD error across three different variations of the Two-Step
substantial decrease in
5

---

Published as a conference paper at ICLR 2025
Values
Errors
1.0
1.0
1.0
O-learning
Tue rewards
Q-learming
0.8
Random rewards
Correlation
0.8
Myopic
Max Correlation
Myopic
Random rewards
0.6
0.6
0.6
plad
0.2
0.2
0.2
0.0
0.0
0.0
100
200
300
25
50
75
25
50
75
Step
Block
Block
D
Values
Errors
1.0
1.0
rewards
lesion
0.8
Control lesion
lesion
Correlation
0.8
TD lesion
Correlation
0.8
TD lesion
plad)
0.6
M
Wyw
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
100
200
25
Step
Block
Block
Figure 4: Llama can predict the actions Q-learning agent and keeps track of variables similar to Q-values and TD errors. (A) Llama predicts action sequences better when given correct infor- mation about rewards. (B & C) Q-values and a reward-tracking variable, as well as accompanying signals, are significantly correlated with SAE features. Max correlations shown after Gaussian smoothing with a 0.5. (D) Lesioning the TD latents impacts action prediction accuracy, whereas lesioning other features barely affects action predictions. (E & F) Lesioning TD latents also impacts subsequent Q-value and TD error representations.
Task. Intervening on this single SAE feature was sufficient to significantly alter Llama's 's ability to perform RL in-context and impacted subsequent representations.
4 Q- Q-VALUES AND TD SUPPORT OPTIMAL ACTION PREDICTION
In the previous task, showed that Llama's behavior and internal representations in simple RL task could be predicted by a Q-learning model. Next, we investigated whether we could also identify TD errors and Q-values in Llama when prompted by sequence of observations and actions generated by a Q-learning agent. This setup allowed us to analyze Llama's representations in a more complex MDP. We trained 50 separate Q-learning agents navigate from fixed starting location to fixed goal location in a 5 x 5 grid using 4 different actions (UP, DOWN, LEFT, and RIGHT) (Figure IB). Entering the goal state led +1 reward and terminated the episode, while entering any other state yielded -1 reward. Following Kaplan et al. (2020), we quantified Llama's in- in-context reinforcement learning ability by how well it could predict the Q-learning agent's actions over time. Further details about the task, including the prompt structure, are provided in Appendix A.4.
First, we observed that Llama improved at predicting the actions of the Q-learning agent as it re- ceived more observations. While this indicates in-context learning, we cannot yet attribute this ability to reinforcement learning, as Llama could simply learn the action sequence without attend- ing to the reward. To control for this, we provided Llama with second prompt, where the action sequences were identical, but the rewards were randomly set to +1 or -1, independently of whether the agent reached the goal. As shown in Figure 4A, Llama is better at predicting the same action sequences when given the correct rewards, indicating that integrated reward information to predict actions. See Appendix A.4 for additional control analyses.

---

Published as a conference paper at ICLR 2025
We then trained SAEs on Llama's residual stream representations for this task and examined whether the learned latents corresponded to the Q-values and the TD errors of the agent that generated the observations. The training done for 15 epochs with regularization strength of 3 1e - 05. As in the previous task, we correlated each latent representation learned by the SAEs with the Q-values, as well as the myopic values. As an additional control, included estimates from Q-learning agent trained on the random reward sequence from the control analysis presented above. Through- out Llama's 's transformer blocks, we observed the highest correlations against the Q-values that were used to generate the observations (Figure 4B). Similarly, most blocks showed the highest correla- tions for the TD errors of the same model. We also found high correlations with the myopic error signal, indicating that Llama learned rewards as well (Figure 4C). These findings further support the hypothesis that Llama's internal representations encode reinforcement learning-like computations. even in more complex environments with larger state and action spaces.
Lastly, we validated the causal role of the identified TD latents by deactivating and scaling them. We selected 4 different latents across 4 different blocks and set their activity to O. All of these latents were selected from block 40 onwards and had r > 0.75 with the TD error. We selected more features to deactivate in this experiment to elicit stronger effects on Llama's behavior. We also conducted control lesions, where we set the activity of latent units from the same blocks that has the lowest correlation with TD to O. We find that lesioning the TD latents significantly degraded Llama's ss ability to predict actions, whereas the control lesions led to very small changes in Llama's action predictions (Figure 4D). Deactivating or scaling the TD latents, we observed a decrease in both the Q-value correlations (Figure 4E) and the TD error correlations (Figure 4F). These results provide evidence for the causal role of the identified TD latents in Llama's reinforcement learning capabilities.
LEARNING GRAPH STRUCTURES WITHOUT REWARDS
So far, we have studied TD learning in the context of RL. However, TD learning can also be used to learn about statistical structures in the absence of rewards. For example, the Successor Repre- sentation (SR; Dayan, 1993) encodes an MDP in terms of future discounted state occupancies given the agent's 's policy, and can be learned using TD learning. The SR is represented as a state by state matrix M. where M(s s')
is the future discounted expected occupancy of state when the agent
(1((s= = s')[s0 s where is the indicator function. The TD error computation and the update defined in equation 1 & equation 2 can be generalized to learn this representation as follows:
at state s, c.g. M(s, E Σ
ÔTD = 1s, + - M(s.:) M(s,,)) - M(s..) + a 8TD
(4)
(5)
Here 1,, is a one-hot vector at state index 8, The SR has seen wide application not only in RL. research (Kulkarni et al., 2016; Barreto al., 2017: Machado et al., 2017), but has also been used as a model of hippocampal function in neuroscience (Garvert et al., 2023; Gershman, 2018; Stachenfeld al., 2017; Gardner et al., 2018). For example, it has been shown that when humans are presented with a sequence of stimuli that are generated from latent graph, the representations in the medial temporal lobe resemble the diffusion properties of the graph (Schapiro et al., 2013; 2016; Garvert et al., 2017), which is captured by the SR.
If Llama uses TD-like computations to represent Q-values, does it similarly rely on a temporal difference error signal to learn something akin to the SR? To test this hypothesis, we prompted Llama with a sequence of observations generated from a random walk on a latent community graph (Schapiro et al., 2013) (Figure IC). This graph had three communities, each community with five fully connected nodes. Each community also had two bottleneck nodes that connected to other com- munities. The model was prompted to predict the next observation at each time point. There were 401 observations, and we repeated this procedure 20 independent times, using randomly sampled node names.
7

---

Published as a conference paper at ICLR 2025
Liama Block
Liama Block 40
Liama Block 79
Successar Representation
A
B
Values
Errors
1.00
1.00
Decoding Accuracy
0.7
SR
SR
0.4
Bottleneck State
Max Correlatin
0.75
Transition Matrix
Max Correlation
0.75
Transition Matrix
Accuracy
0.6
0.50
0.50
0.2
0.5
0.25
0.25
mmmw
0.00
200
400
20
40
80
0
40
80
Step
Block
Block
Block
Figure 5: Llama learns graph structures through TD-learning, representing them similarly to the successor representation (SR). (4) Llama's state representations projected in 2D space, using multi- dimensional scaling, shows the emergence of latent graph structure across transformer blocks. (B) Llama quickly achieves high accuracy in predicting the next state. Accuracy is averaged over 100 runs. (C) Bottleneck states can be linearly decoded from middle blocks onward. (D & E) Latent representations of SAEs trained Llama's 's representations strongly correlate with the SR and asso- ciated TD learning signals, outperforming model-based alternatives. Shaded regions in B-E indicate 95% confidence intervals. Dashed horizontal lines indicate chance level performance, and the solid horizontal line represents the ceiling.
First, we saw that Llama learns the local structure of the graph and can predict the next state with ceiling accuracy after around 100 observed state transitions (Figure 5B). Does Llama also learn the global connectivity structure of the graph? To test this, we performed multidimensional scaling (Torgerson, 1952) on all 80 residual streams' representation of each node in the graph after learning- We saw that the 2D projections of these representations gradually grow more and more similar to the 2D projection of the SR of the graph (Figure 5A), indicating that Llama represents the global geometry of the latent graph.
To further probe Llama's 's understanding of the global graph structure, we tested whether it explicitly represents bottleneck states. These are special states that allow transitions between communities and are particularly useful for planning in hierarchically structured environments. Although such representations are not needed to predict the next tokens, we found that they can be linearly decoded starting around block 20 onward (Figure 5C), with the decoding accuracy peaking around halfway through the model.
These findings show that Llama builds a representation of the graph's global geometric properties over blocks. However, this could be achieved without TD learning, such as through a model that learns the environment's transition matrix by counting experienced transitions. To test whether Llama acquires structural knowledge through TD learning, we trained SAES on its representations and compared them to agents learning either the SR using TD errors, or simply the graph's 's transition matrix, respectively (see Appendix A.5 for training details).
We identified latents that were maximally correlated with each set of representations. We found stronger correlations with the SR (max r = 0.62) than with the transition matrix (max r = 0.49) throughout the model, as shown in Figure 5D. These representations build gradually and peak late in the model. Consequently, we found that the TD errors (max r 0.60) that are used to learn the SR had much stronger correlations than a surprise signal corresponding to the log- likeelhodoo of the next state under the learned transition matrix (max findings using more traditional methods such as representational similarity analysis (Kriegeskorte et al., 2008; Kornblith et al., 2019) in Appendix A.5.
0.18; Figure 5E). We additionally replicated these
8

---

Published as a conference paper at ICLR 2025
Block 65 Following
Block 65 Following
Values
Errors
No Lesion
TD Lesion
A
E
1.00
TD Lesion
Max Correlation
0.2
0.75
0.50
0.1
0.0
No Lesion
TD
20
40 Block
60
80
0
20 40 Block
60
80
Lesion
Figure 6: Lesioning TD SAE latents impair behavior and representations. (A) Following a lesion in block 64, Llama's accuracy in predicting the next state drops. (B) The SAE representations following the lesioning have reduced correlations with the SR, despite the recovery of the TD error following lesioning (C). While the community structure is reflected in the original representations in Block 65 (D), these representations are disrupted as a result of lesioning in earlier block (E).
Lastly, we deactivated the TD latent that had the strongest correlation with the TD error, which was in block 64. Following this intervention, Llama's prediction accuracy for the next state decreased from 24.3% to 18.9% (Figure 6A). This also reduced the subsequent correlations between the SAE latents and the SR (Figure 6B), despite rapid recovery of the TD correlations after the lesion (Figure C). The disruption of the learned structure can be seen in block 65, where the community structure of the graph is distorted (Figures 6D and 6E). Taken together, behavioral and representational analyses show that Llama's graph knowledge is represented similarly to the SR, and the TD latents identified support learning these representations.
RELATED WORK
Understanding in-context learning in transformers has received significant attention (Garg et al., 2022: von Oswald al., 2023; Dai et al., 2023; Ahn et al., 2023). Previous work has shown that transformers can discover known and existing algorithms without explicit guidance towards these algorithms. Akyürek et al. (2023) have trained transformers on linear regression problems and shown that the model can implement ridge regression, gradient descent, or ordinary least squares in- context. Focusing on RL, Wang et al. (2024) have shown that linear transformers trained explicitly to solve reinforcement learning problems discover TD learning methods. On larger scale, Laskin et al. (2022) has shown that in-context RL can be distilled to smaller models. We build on this by consequence of language pre- mre-training,
showing that in-context RL emerge in LLMs simply even if the model was never explicitly trained to solve RL problems.
SAEs have been successfully used decompose LLM activations into interpretable features. This has been demonstrated both toy models (Cunningham et al., 2023; Bricken et al., 2023) and state- of-the-art language models (Gao al., 2024). SAEs can identify features ranging from concrete objects to abstract concepts (Templeton et al., 2024). While specific static concepts have been identified using SAEs, we explore the characteristics of in-context learning algorithms using this technique.
In the broader scope RL and LLMs, several papers have investigated how the two methods can be integrated to improve LLMs. Le et al. (2022) coupled LLMs with deep RL in an Actor-Criti Critic (Konda & Tsitsiklis, 1999) setup for program synthesis, where the deep RL model served as the critic for the code generated by the language models. Shinn et al. (2024) built framework where an LLM received verbal RL feedback, allowing it to perform better in decision-making tasks, coding, and reasoning problems. Lastly, Brooks et al. (2022) developed method for improving in-context policy iteration LLMs for RL tasks. In contrast to these approaches where RL was used to improve LLM behavior, focused understanding the mechanisms through which LLM can do RL. Given the demonstrated usefulness RL in improving LLMs' abilities, is essential to understand these mechanisms.

---

Published as a conference paper at ICLR 2025
7 DISCUSSION
TD learning fundamental algorithm artificial intelligence research (Sutton & Barto, 1987; 2018; Mnih et al., 2015). It offers a simple yet efficient solution the problem of distilling tem- porally distant consequences of actions into an immediate value signal. TD learning is general: it can be used not only to learn future discounted rewards but also state occupancies (Dayan, 1993), uncertainties (Guez et al., 2012) and prediction errors (Burda et al., 2018; Saanum et al., 2024). It therefore perhaps not surprising that neuroscience research has found traces of TD learning in striatal neural populations in several species. In this paper, we offer evidence that similar substrates of prediction and reward are implemented in the eircuits of a large transformer network pretrained on internet-scale text corpora. This is surprising because the LLM was not trained using RL objective but simply to perform next-token prediction. One hypothesis is that TD learning could facilitate next-token prediction by allowing the LLM build more general models of the task. In the Graph Learning Task, the TD error can be used to build the SR which offers global geometric information about the graph structure. Representing features that span longer temporal horizons could be important for building rich representations of the world, which TD learning affords.
7.1 LIMITATIONS & FUTURE WORK
There are some important limitations and extensions of our work. First, our SAEs are task-specific and are not suitable for identifying RL-related variables for arbitrary tasks. Future work should attempt to build more generic SAEs and repeat our analyses with them. Perhaps, this can be ac- complished through training SAEs representations coming from large and diverse corpus of RL tasks, or by fine-tune sone-tuning pretrained SAEs (c.g.. Lieberum et al. (2024)) on RL data. Similarly, to establish the generality of our results, future work could aim replicate findings using different LLMs.
Second, while mapped out TD errors and Q-values across the residual blocks of Llama, circuit- level understanding of how these representations come about remains unclear. Can induction heads (Olsson al., 2022) give rise these representations? Using methods such as path patching (Hanna et al., 2024; Wang et al., 2022), can in-context TD circuits be identified? These remain important questions for complete understanding of how LLMs do in-context RL.
Lastly, while we can predict Llama's behavior using Q-learning agents and find TD-like represen- tations in the model, their alignment is perfect. Future work should identify which additional components are needed to fully account for the model's behavior. For example, adding a prefer- ence for repetition to the Q-learning model can be a natural extension to capture the tendency of pretrained LLMs to repeat themselves.
7.2 CONCLUSION
In this work, we have shown evidence that LLMs implement TD learning to solve RL problems in-context. Our work only demonstrates that SAES disentangle features used during in- context learning but also that these discovered features can be used to manipulate LLM behavior and representations in systematic fashion. Our finding paves the way for studying other types of in-context learning abilities using SAEs. Lastly, our work establishes line of convergence between the learning mechanisms of LLMs and biological agents, who have been shown to implement TD computations in similar settings.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, 2023. URL https: /arppsiv//axx...///////sp org/ abs/2306. 00297.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm in-context learning? investigations with linear models, 2023. URL https:// arxiv. org/ abs /2211 15661.
André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado P Hasselt, and David Silver. Successor features for transfer in reinforcement learning. Advances in neural information processing systems, 30, 2017.
Marcel Binz and Eric Schulz. Using cognitive psychology to understand GPT-3. Proceed- ings of the National Academy of Sciences, 12066):22118523120 February 2023. doi: 10.
1073/pnas.221 2215233100, URL https:/ /www. pnas org/ / /do///aabbb///////// /10. 1073/pnas.
2218523120. Publisher: Proceedings of the National Academy of Sciences.
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards Monosemanticity: Decomposing Language Models With Dictionary Learning. Transformer Circuits Thread, 2023.
Ethan Brooks, Logan A Walls, Richard Lewis, and Satinder Singh. In-context policy iteration. In NewrlPS 2022 Foundation Models for Decision Making Workshop, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810.122 12894, 2018
Julian Coda-Forno, Marcel Binz, Zeynep Akata, Matt Botvinick, Jane Wang. and Eric Schulz. Meta-in-context learning large language models. In A. Oh. T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 65189-65201. Curran Associates, Inc., 2023. URL https::/procest: www.:eeeddnnnnnennnnneeeeeeeeeee neurips...... cttppp/eeepeeeeeeee////eeeeeeeee fiwee/ileepeeee Mttees/ppper/0000 fttle//ccoccraaaaaaaaafffaeeeeee
pdf.. -ddf.
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoen- coders find highly interpretable features in language models, 2023. URL https://arxin: org/abs/ 2309..... 08600.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers, 2023. org/ abs 2211 2. 10559.
URL
Nathaniel D. Daw, Samuel J. Gershman, Ben Seymour, Peter Dayan, and Raymond J. Dolan. Model- Based Influences on Humans' Choices and Striatal Prediction Errors. Neuron, 1244:1201- March 2011. ISSN 0896-6273. doi: 10. 1016///.eeuron.200 1. 1.02.027. URL https: /www. cell
6273 (11)00125-5,
Publisher: Elsevier.
Peter Dayan. Improving Generalization for Temporal Difference Learning: The Successor Repre- sentation. Neural Computation, 5(4):613 July 1993. ISSN 0899-7667. doi: 10.1162/neco. 1993.5.4.613. URL https:/s:/// /tt.///...oo///o/// /10 10/.10111111 62. necoo 1993. 5999.5...... 5.4.61......
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Leuman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv preprint axXiv:2407.21788 2024.
11

---

Published as a conference paper at ICLR 2025
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Satfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models of superpo- sition. Transformer Circuits Thread, 2022.
Christopher D Fiorillo, Philippe N Tobler, and Wolfram Schultz. Discrete coding of reward proba- bility and uncertainty by dopamine neurons. Science, 299(5614):1898-9909 2003.
Jaden Fiotto-Kaffmaa Alexander R Loftus, Eric Todd, Jannik Brinkmann, Caden Juang. Koyena Pal, Can Rager, Aaron Mueller, Samuel Marks, Arnab Sen Sharma, Francesca Lucchetti, Michael Ripa, Adam Belfki, Nikhil Prakash, Sumeet Multani, Carla Brodley, Arjun Guha, Jonathan Bell, Byron Wallace, and David Bau. Nnsight and ndif: Democratizing access to foundation model internals. 2024. URL https: / aaaiiiiiv. wwg/aos//e....ooo.o............e 24407 .ET.407.156... 14561.
Shelly B Flagel, Jeremy Clark, Terry E Robinson, Leah Mayo, Alayna Czuj. Ingo Willuhn, Christina A Akers, Sarah M Clinton, Paul EM Phillips, and Huda Akil. A selective role for dopamine in rewardd learning. Nature, 469(7328):53 2011.
Leo Gao, Tom Dupré la Tour, Henk Tillman, Gabriel Goh, Rajan Troll, Alec Radford, Ilya Sutskever, Jan Leike, and Jeffrey Wu. Scaling and evaluating sparse autoencoders, 2024. URL https:
//arxiv. org abaaabas/0000 2406. 04093.
Matthew P H Gardner, Geoffrey Schoenbaum, and Samuel J Gershman. Rethinking dopamine as November 2018.
generalized prediction error. Proc. Biol. Sci.,
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30599 2022.
Mona M Garvert, Raymond J Dolan, and Timothy EJ Behrens. A map of abstract relational knowl- edge in the human hippocampal-entoriee entorhinal cortex. eLife, 6:e17086, April 2017. ISSN 2050- 084X. doi: 10.7554/eLife. 7086. URL https::/dols http///o.ooggr//////// Publisher: eLife Sciences Publications, Ltd.
1//5/////fe..// 17086.
Mona M Garvert, Tankred Saanum, Eric Schulz, Nicolas Schuck, and Christian F Doeller. Hip- pocampal ppatio-predictive cognitive maps adaptively guide reward generalization. Nature Neu- roscience, 26(4):615-626, 2023.
Samuel Gershman. The successor representation: its computational logic and neural substrates. Journal of Neuroscience, 3803))77338 2018.
Arthur Guez, David Silver, and Peter Dayan. Efficient bayes-adaptive reinforcement learning using sample-based search. Advances neural information processing systems, 25, 2012.
Wes Gurnee, Neel Nanda, Matthew Pauly, Katherine Harvey, Dmitrii Troitskii, and Dimitris Bert- simas. Finding neurons in haystack: Case studies with sparse probing. arXiv preprint arXiv:2305.01610, 2023.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine Soft actor-critie: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018
Michael Hanna, Ollie Liu, and Alexandre Variengien. How does gpt-2 compute greater-than?: Inter- preting mathematical abilities in a pre-trained language model. Advances in Neural Information Processing Systems, 36, 2024.
William M. Hayes, Nicolas Yax, and Stefano Palminteri. Large language models are biased rein- forcement learners, 2024. URL pettts://arxtaa org/ abs/2405 11102...11......
Martin N. Hebart, Adam H. Dickter, Alexis Kidder, Wan Y. Kwok, Anna Corriveau, Caitlin Van Wicklin, and Chris Baker. THINGS: A database of 1,854 object concepts and more than 26,000 naturalistic object images. PLOS ONE, 14(10):0223792, October 2019. ISSN 1932-6203. doi: 10. 10.1377/ouummal.pon.00222
URL https: atttps://jounnlle journals plcs. org/plosone/
article?i=111 1371/journal pone. 0223792. Publisher: Public Library of Science.
12

---

Published as a conference paper at ICLR 2025
Stefan Heimersheim and Neel Nanda. How to and interpret activation patching, 2024. URL alss http://p/p//.....///// 2404.
attps::/arxiv...
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. URL https: / arxiv. wrg. abss/2000 08361.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL
https: http:v/axvvvvvooo .rg. wttp://rrgr//ooo/ooo/o/oooss//// 1412. 6980.
Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems, 12, 1999.
Wouter Kool, Fiery A. Cushman, and Samuel J. Gershman. When Does Model-Based Control Pay Off? PLOS Compulational Biology, 1208)=1000000, August 2016. ISSN 1553-7358. doi: poi URL /ttpss//jouuraale paw.srannnsnsssl.a....oooooooole plos......... org/ httg///oscoppiiin
1ttiic/erddt11.13///oorrrree 1tt...1111o//aaa//////aa pobi wtt.i...000000 1005090. Publisher: Public Library of Science.
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International conference machine learning, pp. 3519- 3529. PMLR, 2019.
Nikolaus Kriegeskorte, Marieke Mur, and Peter A Bandettini. Representational similarity analysis- connecting the branches systems neuroscience. Frontiers in systems neuroscience, 2:249, 2008.
Tejas D Kulkarni, Ardavan Saeedi, Simanta Gautam, and Samuel J Gershman. Deep successor reinforcement learning. arXiv preprint (to.:.26009966 2016.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, al. In-context reinforcement learning with algorithm distillation. arXiv preprint
2022.
Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven Hoi. Coderl: Mastering code generation through pretrained models and deep reinforcement learning, 2022.
URL https::///// /warxivvvvv.////// org/ aba/ 01780.
Tom Lieberum, Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Nicolas Sonnerat, Vikrant Varma, János Kramár, Anca Dragan, Rohin Shah, and Neel Nanda. Gemma scope: Open sparse autocncoders everywhere all at once gemma 2024. URL https::/aaxxi..rrrrrese// 2408. 05147.
Marlos C Machado, Clemens Rosenbaum, Xiaoxiao Guo, Miao Liu, Gerald Tesauro, and Murray Campbell. Eigenoption discovery through the deep successor representation. arXiv preprint 17X.:11710.11089 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529-5333 2015.
P Read Montague, Peter Dayan, and Terrence Sejnowski. A framework for mesencephalic dopamine systems based on predictive hebbian learning. Journal of neuroscience, 16(5):1936- 1947, 1996.
Yael Niv, Reka Daniel, Andra Geana, Samuel J. Gershman, Yuan Chang Leong, Angela Radulescu, and Robert C. Wilson. Reinforcement learning in multidimennional environments relies atten- tion mechanisms. Journal of Neuroscience, 35(21):8145-81)) 8157, 2015. ISSN 0270-6474. doi: 10. 10.153/NEEROSSC1.2977 14.2015. URL https: www. http://wwwwwwwwuuwweceeeo www.oooi..oog/ooooooooeeeeeeeoee content// 35/21/8144
Chris Olah. Mechanistic Interpretability, Variables, and the Importance of Interpretable Bases, 2023. httpa:://raassoonneeee www.rrruuustttu.aeeeeee ppbp/022/ http:///cccnineepeeeeeeeeeee nndex. html.
13

---

Published as a conference paper at ICLR 2025
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfiel--oodss Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kermion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack
Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. https:/ttansforrerccicuitts
Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit-learn: Machine learning in python. the Journal of machine Learning research, 12:2825-2883 2011.
Qwen Team. 5wen2.5: A party of foundation models, September 2024. URL httht::: /qwenlm.. github. rithub.io//ble io/blg// hwwp://////gg////m2mmmmmm qwen2.5/
RA Rescorla and Allan Wagner. A theory of Pavlovian conditioning: Variations in the effectiveness of reinforcement and nonreinforcement, volume Vol. 2. 01 1972.
Tankred Saanum, Noémi Éltetō, Peter Dayan, Marcel Binz, and Eric Schulz. Reinforcement learning with simple sequence priors. Advances in Neural Information Processing Systems, 36, 2024.
Anna C. Schapiro, Timothy T. Rogers, Natalia I. Cordova, Nicholas B. Matthew M. Botvinick. Neural representations events arise from temporal community struc- ture. Nature neuroscience, 16(4))486-49) 492,448664442-- April 2013. ISSN 1097-6256. doi: 10. 10.10/nn.3331.. URL wttptttpp///cc http://ww/////.mm ncbi nlm. nih wtt.ggg////ccctgcggg
and
Anna C. Schapiro, Nicholas B. Broo-Browne. Kenneth A. Norman, and Matthew M. Botvinick. Statistical learning of temporal community structure in the hippocampus. Hippocampus, 26(1): 3-8, January 2016. ISSN 1098-1063. doi: 10. 00.1002/hipo.2555 222523.
Johannes A. Schubert, Akshay Kumar Jagadish, Marcel Binz, and Eric Schulz. In-context learn- ing agents are asymmetric belief updaters. In Proceedings of the 41st International Con- ference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 43928 43946. PMLR, 21-27 Jul 2024. URL https: httt:s://procceeingg mlr. press/ v235 2.35/schubert2244 html.
W. Schultz, Dayan, and P. Montague. A neural substrate of prediction and reward. Science (New York, N.Y.), 275(5306) 1599, March 1997. ISSN 0036-8075. doi: 10.1126/science 275.5306.1593.
Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Kimberly L. Stachenfeld, Matthew M Botvinick, and Samuel Gershman. The hippocampus predictive map. Nature neuroscience, 20411)164331165, 2017.
Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine Learning, 3(1):9-44, August 1988. ISSN 1573-0565. doi: 10.1107//F001 URL https://dol org/ /tgg/10.o 100.0000001000000111111011111111 /tt001001010000
Richard Sutton and Andrew G Barto. A temporal-difference model of classical conditioning. In Proceedings of the ninth annual conference of the cognitive science society, pp. 355-378. Seattle, WA, 1987.
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: An introduction, 2nd ed. Rein- forcement learning: An introduction, 2nd ed. The MIT Press, Cambridge, MA, US, 2018. ISBN 978-0-262-03924-6. Pages: xxii, 526.
Richard Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
14

---

Published as a conference paper at ICLR 2025
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhu- patiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma 2: Improving open language models at a practical size. arXiv e-prints, pp. arXiv-2408, 2024.
Adly Templeton, Tom Conerly, Jonathan Marcus, Jack Lindsey, Trenton Bricken, Brian Chen, Adam Pearce, Craig Citro, Emmanuel Ameisen, Andy Jones, Hoagy Cunningham, Nicholas Turner, Callum McDougall, Monte MacDiarmid, C. Daniel Freeman. Theodore R. Sumers, Edward Rees, Joshua Batson, Adam Jemmyn, Shan Carter. Chris Olah, and Tom Henighan. Scaling monosemanticity: Extracting interpretable features from claude sonnet. Trans- former Circuits Thread, 2024. URL https: wow.ransooommerrrrcciooooo pubp//222
scaling s mww.senanticits/ hineeitt/iiii.i..eee html.
Warren S. Torgerson. Multidimensional scaling: I. theory and method. Psychometrika, 17:401-419, 1952. URL https: //pi semant nt cschola wrg. (NOrpusTD:12084990
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2023. URL https:://///....c. http///axiv..../ www.v.org/ab.....a..oo httgg//bs//2222/ /tlpo//221..000 07677.
Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, and Shangtong Zhang. Transformers learn temporal difference methods for in-context reinforcement learning, 2024.
abs/2405. 13861.
Kevin Wang, Alexandre Variengien, Arthur Commy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593, 2022.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8:279-299 292, 1992.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo- gatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent Abilities of Large Language Models, October 2022. URL [cs].
hts//bb//20555 07682. arXiv:2206.07682
Safoora Yousefi, Leo Betthauser, Hosein Hasanbeig, Raphaël Millière, and Ida Momennejad. De- coding in-context learning: Neuroscience-inspires analysis of representations in large language models, 2024.
15

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 Q-LEARNING
In most RL settings we compute the TD error over a batch of observations. In offline off-policy RL (Haarnoja al., 2018), we calculate the TD error over batch state transitions experienced in the past, potentially collected using different policy. To model TD learning in Llama, whose transformer architecture allows to casily integrate information over multiple episodes the past, calculate the TD error average over the k last state transitions, where & is hyperparameter. In the Two-Step Task, we use a = 4, as this gave a good behavioral and representational fit. In the entire history of previous observations to compute the TD error, slightly
the other tasks, smoothing out the error relative to k 4. See Figure 7 for a qualitative comparison of the TD error with these settings. Other hyperparameters used train the Q-learning model include the discount parameter 7 = 0.99 across all tasks. The learning rate a was 0.1 in the Two-Step Task and the Grid World, and 0.05 in the Graph Learning Task. In the Two-Step task and the Grid World, we passed the Q-value estimates through softmax to obtain action probabilities, which were used to compute -log p(a).
In the Grid World task, we train the Q-learning agent using an c-greedy exploration strategy, where we linearly decay for the first episodes, leaving at 0 for the remaining episodes of the training run.
n=1
n=4
n=10
n=full
GTD
50
25
25
50
Step
Step
Step
Step
Figure 7: As the replay window & grows, the TD signal is smoothed.
A.2 SAE TRAINING
For each task and each block of Llama, we train separate SAEs. For SAEs, batch size of 256, learning rate of Te 04, and B = 1e 05 were used. We used the Adam optimizer (Kingma Ba, 2017) with the default parameters and shuffled the training data. The latent space for the Two-Step and the Grid World tasks was doubled in size. For the Graph Learning Task, we used an 5192-dimensional latent space, equal to the input space size. We scaled our input data e [gaxd xd as follows before passing it through the SAEs: H' Vd.
I I hh When we replaced Llama's
activations with the reconstructed activity H' from the SAEs, we applied the inverse transformation. All the interventions were performed using the nns ght (Fiotto- Kaufman al., 2024) library.
A.2.1 SAE Lo NORMS
Training SAEs residual stream activations for the three experiments, that the original -192-dimensional representational space can be reduced to low-dimensional feature spaces with few dozen active features. See Figure 8.
16

---

Published as a conference paper at ICLR 2025
Two-Step Task
Grid World
Community Graph
30
/ lur phague Mam Malp
25
20
15
M Myr
10
N
25
50
75
0
25
50
75
25 50
75
Block
Block
Block
Figure 8: The Lo norm of SAE representations (e.g. the number of features with non-zero variance) tends peak in the middle blocks of Llama, except for in the Graph Learning Task, where it in- creases again towards the last blocks. This is presumably because Llama needs to predict the next token out of larger set of token candidates in that task.
A.3 THE TWO-STEP TASK
Prompt
You are an optimal reinforcement learning agent. You are playing little game. You have two actions, Left and Right. Your goal is to get as many points as possible within an episode. Here go.
Episode O: You pick right. You see an orange. You pick left. You get 6 points. Episode is over.
[...]
Episode 19: You pick left. You apple. You pick right. You get points. Episode is over.
We used the prompt shown above for the Two-Step Task. Llama's internal representations were recorded at the pick tokens.
Model comparison. To test whether the negative log- likelihkooods between the Myopic and the Q- learning were significant, we averaged negative-log likelihoods over each run and conducted paired t-test. The Q-learning model was significantly better than the Myopic model (t(99) -3.40,
p = 0.001).
A.4 REINFORCEMENT LEARNING A GRID WORLD
Prompt
What is the next object in the following sequence? You are living on a two-dimensional grid-world. You have actions, left, right, up and down. Your goal is get the goal in few steps as possible.
Episode 0: You
1. You go left. You get -1 points. You 0, 1. You go down.
You get 1 points. You are at 0, O. You down [...]
Episode 1: You are at 1. You go right. You get -1 points. [...] You go right. You get +1 points. You finished with -3 points.
17

---

Published as a conference paper at ICLR 2025
For the Grid World task, we prompt Llama as shown above. The go tokens are the recording points for internal representations, and the logits here were used to predict the action probabilities. This task was repeated 50 independent times, and we analyzed the behavior and the representations from the first 320 observations.
Original No Reward
Original Swapped Reward
0.06
0.05
0.04
Aplay)
0.03
Malle
0.02
0.00
-0.01
100
250
Step
Step
Figure 9: Additional control analyses for the Grid World. (A) The change in Llama's prediction accuracy of the Q-learning agent's actions from control prompt with no rewards to the original prompt. (B) Same as A, except the control prompt uses sign-swapped rewards from the original prompt. Higher values in both plots indicate the effeet of reinforcement during learning.
Run
Run
Run
Run
Run
Run
Standard
Reward aisruption
Transition disruption
0.5
0.0
30
60
30
6.0
60
60
0
30
Step
Q-value
SAE to fecture
Figure 10: Examples of SAE features and Q-value for the Left action in the Two-Step task. There dips both in the SAE feature and the Q-value certain points that late in the task. If Llama does TD-learning, this is to be expected when Llama ends up in a state where particular action is suboptimal, just like the Q-value.
Additional control analyses. In addition to randomizing the rewards as described in the main text, conducted two other control analyses. In the first (Figure 9A), removed the reward
18

---

Published as a conference paper at ICLR 2025
information altogether from the prompt. During the exploration phase, this did not change how well Llama predicts actions, but we observed better predictions for the original prompt in later trials. In another control analysis, instead of randomizing the rewards, swapped the + and the signs (Figure 9B). This led to a smaller effect as removing the reward, but the shift was smaller in magnitude.
Testing a simpler alternative to TD learning. A reviewer raised the point that Llama's behavior could potentially be explained by a simpler mechanism than TD learning, for instance a model that keeps track of the rewards obtained at the end of each episode: Suppose in state 8, it assigns an action logit of if the action was performed in s the maximally rewarding episode. We then get the action probabilities by applying a softmax with temperature parameter to the logits, where we fit the temperature parameter minimize the negative log likelihood on I. Llama's 's behavior. We refer to this as the Look-up model. As can be seen in Fig. 11, the Q-learning model predicts Llama's choices better than the look up model.
3400
3200
3000
2800
2600
2400
2200
2000
Myopic
Q-learning
Look-up
Figure 11: Q-learnng fits Llama's behavior better than the Look-up model.
A more difficult grid world. We additionally tested Llama different variant of the grid world. In this setup, we used a by grid, increasing the state space from 25 to 49. Additionally, each episode, we randomized the initial location of Llama, requiring strong generalization. The results are shown in Fig. 12. We found that Llama can learn to predict the Q-learnin agent's 's actions in this task well.
Using the larger grid world also lets us dissociate Llama's behavior from the Look up model better. In the original grid world, observed highly similar predictions between the two models. However, when we tested Llama on a larger grid world and initialized its starting position randomly in each episode, observed that their predictions severely diverged, where Look-up's predictions stayed around chance level. In the former environment, the sequences of actions become more deterministic over time, which is why Look-up is also good model for that environment. However, changing the initial location of the model in each episode also changes the sequence of optimal actions. Llama can generalize adapt to this, whereas Look-up cannot. The discrepancy in the new environment further suggests Llama is not implementing just Look-up but is indeed doing TD learning.
We also tested whether Llama's internal representations correlate more strongly with Q-values or Look-up values. As displayed in Fig. 13, in both environments, SAE features across blocks correlate more strongly with Q-values than with Look-up values. Notably, the differences in correlations are larger in the more difficult 7 by Grid World (Fig. 13B), paralleling our behavioral findings. Taken together, both our behavioral and representational results show stronger evidence for Llama using TD learning than Look-up.
Lastly, we analyzed the models' predictions in states that don't 't occur at the start of the trajectories, but in the middle of other trajectories. Here the Look-up model predicts actions random. Llama's ' action probabilities are not only considerably higher, but also match the action probabilities of the Q-learning agent with e-greedy policy (Fig. 14).
19

---

Published as a conference paper at ICLR 2025
5 x 5 Grid
7 7 Grid
1.0
1.0
0.8
0.8
0.6
0.6
plad)
play)
0.2
0.2
0.0
0.0
50
100
200
250
50
100
150
250
300
Step
Step
Llama 70B
Look-up
Figure 12: Llama learns to predict actions in larger grid-world environments. While Llama's pre- dictions match those of the Look-up model in the small grid- morld, this is not the case with larger grids which requires the agent to cache state-action values.
5 x 5 Grid
x 7 Grid
1.0
1.0
0.8
Max correlation
h M M 1014
Max correlation
0.6
0.6
0.4
0.4
0.0
0.0
20
40
60
80
0
20
40
60
80
Block
Block
Q-values
Look-up
Figure 13: In both Grid World environments, SAEs show stronger correlations with Q-values than Look-up values.
A.5 GRAPH LEARNING
Prompt
What is the next object in the following sequence?
cake
horse
gift
horse
woman
wall
sock
wall
20

---

Published as a conference paper at ICLR 2025
0.8
0.6
p(ar)
0.4
0.2
0.0
Look-up
Llama Q
(E-greedy)
Figure 14: We analyze the models' predictions states that have been visited before in the past, but from different starting states. Here the Look-up model predicts actions at random. Llama's 's action probabilities are not only considerably higher, but also match the action probabilities of the Q-learning agent with an e-greedy policy.
For the Graph Learning Task, prompt Llama shown above. Colors indicate how the text is tokenized. The representations corresponding to the token are recorded, and we obtain Llama's predictions for the tokens that follow this. Recording representations always at the same token eliminates any differences in representations that may arise from using different tokens.
The latent community graph consisted of 3 communities, with nodes in each community. Every node had 4 adjacent nodes. The node names sampled from the category labels in the THINGS database (Hebart et al., 2019). We only sampled from labels represented as single token in isola- tion, and the sampling
done with replacement across the 20 runs.
Behavior. On average, Llama had an accuracy of 24.3% in predicting the next token, and the ceiling accuracy expected in the environment
25%.
Multidimennsonal scaling (MDS). We used cosine distance to calculate the pairwise dissimilarity of representations both for Llama and the SR. This done separately for different blocks of the transformer and the SR. For this analysis, we only considered the last encounter of both the Llama and the SR for each state. The dissimilarity matrices were used project the representations onto 2D space using metric MDS as implemented in scikit -laann (Pedregosa et al., 2011 I). What is plotted Figure 5A an example from one of the runs.
Decoding bottleneck states. We trained a linear support vector machine on Llama's 's state represen- tations to predict whether the corresponding state bottleneck state. We did this in leave-one- run-out fashion, where the data from each run served as the test data in one fold. This procedure also ensured that the state names could not be leveraged aid classification, as they are randomly assigned across runs.
Learning models. We introduced some new learning models in this task that learn representations of the graph. First is the Successor Representation (Daan, 1993), which is defined in the main text. We also consider an agent learning the transition matrix of the graph. This agent learns a representation T. where T(s, s') is the estimated immediate transition probability from state to s', which can be written as (ss.8 s') [[[[(,+1 = 8')}s, s]. We learn this by counting the times the agent transitions from to s'. divided by the total number of transitions from
21

---

Published as a conference paper at ICLR 2025
To accompany the transition probabilities, also compute the surprise of a transition, measure of the error the learned model.
log 8'), log - T(st, s')), otherwise
if = 8t-1
surprise st. s')
(6)
This yields vector-valued signal surprise(s,) for a given state s.
SAE analyses. The SAEs were trained for 20 epochs with an 8192 dimensional latent space, which was identical the imput space size. After training the SAEs for each transformer block, we corre- lated the activity of the SAE latents with non-zero variance against estimates of our learning models. For each model, we stack representations row-wise and obtain step-by-state matrix. We then cor- relate every column of this matrix with every SAE unit. Then, for each model and each state, we take the maximum correlation and report the average of this across runs.
8
Transition Matrix
TD
Surprise
0.8
CKA
0.4
0.0
20
40
20
40
60
Layer
Layer
Figure 15: Replication of Graph Learning SAE results using CKA. SR (TD) is similar to Llama than the learned transition matrix (surprise).
Replication using representational similarity analysis (RSA). Comparing vector-valued model representations across two different spaces has traditionally been done using RSA both in neuro- science (Kriegeskorte al., 2008) and machine learning (Yousefi et al., 2024). Here, we replicate our SAE findings using RSA. Specifically, we used variation of RSA called Centered Kernel Align- ment (Kornblith et al., 2019) (CKA) that is commonly used in machine learning, where similarity is calculated between representations X and Y as follows:
T CKA(X, Y) |XTX YTY
where || ||r denotes the Frobenius norm. CKA is bounded between 0 and 1, and higher CKA indicates stronger similarity. We stacked Llama's and the learning models's 's representations during the task row-wise and calculated the CKA between them for each run. As shown below, arrive at the same conclusion using this method: Llama represents the environment like the SR and learns this representation through TD learning.
B APPENDIX: REPLICATIONS WITH OTHER LLMS
To test whether our results generalize to other LLMs outside the Llama family, we repeated the majority of our analyses for all three tasks (Two-Step, Grid World, and Graph Learning) on two
22

---

Published as a conference paper at ICLR 2025
models, Gemma-2-27B (Team et al., 2024) and Qwen-2.5-72B (Qwen Team, 2024). Remarkably, almost all analyses yielded qualitatively matching results, showing a high degree of similarity in how these models solve RL tasks in-context. We discuss each individual point raised in the review in more detail below.
We observed that neither model can do the Two-Step task well Llama. Consequently, the RL models predict behavior not as accurately for these models. However, the SAE correlations remain similar the original findings. The results visualized in Fig. 16 and Fig. 17 respectively. For the Grid World, the replications closely match the original findings. We visualize them in Fig. 18 and 19. Both models can learn to predict the actions in the sequence, and we observe stronger correlations in the SAE features with TD driven value and error signals compared to alternatives. Similarly, in the Graph Learning task, both models have strong SR and the community structures emerge over the transformer blocks. These findings are visualized in Fig. 20 and Fig. 21.
23

---

Published as a conference paper at ICLR 2025
Gemma-2-27B
Errors
6000
Correlation
1.0
Max Correlation
1.0
Myopic O-leariing
Myopic
-logpla)
5000
Qlleaming
0.5
4000
Myopic
25 Black
Episade
learning
Block
Figure 16: For Gemma-2-27B, there are SAE features that correlate highly with TD errors and values. The model can not do the task well Llama. Consequently, the RL models predict behavior not as well for these models
Qwen-2.5-72B
Values
Errors
B
6000
Correlation
1.0
Max Correlation
1.0
Qwen 728
Myopic Relarboos
Myopic
-logp(a)
5000
O-leaming
0.5
0.5
3000
0.0
0.0
20
Myopic
25
50
75
0
25
50
75
Episode
learning
Block
Block
Figure 17: For Qwen-2.5-72B. there are SAE features that correlate highly with TD errors and values. The model can not do the task as well as Llama. Consequently, the RL models predict behavior not well for these models
24

---

Published as a conference paper at ICLR 2025
Gemma-2-27B
A
Values
Errors
1.0
1.0
Q-learning
rewards
O-leaming
0.8
Random rewards
Correlation
0.8
Myopic
Max Correlation
Myopic
Random rewards
Random rewards
plad
0.6
0.6
0.6
0.4
MW Woshoy
0.2
0.2
0.2
0.0
0.0
0.0
100
200
20
40
Step
Block
Block
Figure 18: Behavioral and SAE correlation analyses reported in Fig. 4 replicate with Gemma-2-277
Qwen-2.5-72B
A
Values
Errors
1.0
1.0
1.0
Q-leerning
Tue rewards
O-leamiin
0.8
Random rewards
Correlation
0.8
Myopic
Correlation
Myopic
Random rewards
Random
plac)
0.6
0.6
MW
0.4
0.4
0.4
0.2
0.2
0.0
0.0
100
200
300
25
50
75
0
25
50
75
Step
Block
Block
Figure 19: Behavioral and SAE correlation analyses reported in Fig. 4 replicate with
25

---

Published as a conference paper at ICLR 2025
Gemma-2-27B
Germa Block
Gemma Block 20
Gemma Block
Successor Representation
Values
Errors
1..00
Decoding Accuracy
0.7
SR
SR
0.4
Bottleneck State
Correlation
0.75
Transition Matrix
Max Correlation
Transition Matrix
Accuraci
0.6
0.50
0.2
0.25
0.25
0.0
0.4
0.00
0.00
200
400
10
20
30
40
10
20
30
40
0
10
20
30
40
Step
Block
Block
Figure 20: Behavioral, decoding, and SAE correlation analyses reported in Fig. 5 replicate with Gemma-2-27B.
wwen-2.5-722
Owen Block
Qwen Block
Qwen Block
Successor Representation
Values
Errors
1..00
Decoding Decodlig&ccuureer Accuracy
0.7
SR
SR
0.4
Bottleneck State
Max Correlation
0.75
Transition Matrix
Max Correlation
0.75
Transition Matrix
Accuracy
0.6
0.2
0.50
0.25
0.25
0.0
0.00
0.00
200
400
20
40
60
80
20
40
60
80
0
20
40
60
80
Step
Block
Block
Figure 21: Behavioral, decoding and SAE correlation analyses reported in Fig. replicate with Qwen-2.5-72B.
26


Reviews:
[{"review_id": "XfHLBIt8ZA", "reviewer": "Reviewer_BhvT", "summary": "The paper investigates whether Llama 3 70B has internal representations that support temporal difference learning. First, it demonstrates that Llama can solve RL tasks significantly better than chance. Next, it trains a sparse autoencoder (SAE) and finds features correlated with TD error. Finally, it causally intervenes on these features to show that in-context RL performance degrades without those specific TD features.", "questions": "Could you please provide clarification re: weaknesses 2 & 3?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This is an excellent paper. It asks a very interesting question and provides compelling evidence for the conclusion that Llama represents TD error and uses it to solve RL problems in-context. The section on successor representations was a welcome surprise in section 5, and offered more evidence for TD learning, even absent any rewards. The paper was also quite easy to follow and laid out the argument in a very natural way. I don't have any major complaints.", "weaknesses": "Only minor weaknesses.\n\n1. In the background section on RL, TD is presented for a fixed policy, and then the paper switches to Q-learning, assuming the policy chooses \\argmax_a Q(s,a). But this will change the policy as the Q function is updated, so it's not technically the same setting.\n2. It was a bit unclear what \"control lesion\" referred to in Fig. 2F. And more generally, I was not familiar with the \"lesion\" terminology, so a brief definition would be welcome. I assume it's a form of activation patching?\n3. I would have liked slightly more explanation regarding \"clamping\" the activations. I assume this means setting them to a specific value, but how is that different from deactivating them (i.e. clamping them to zero)? Is the purpose of clamping the activations to show degraded, unchanged, or improved performance?\n4. Line 458, mangled sentence \"our study is, we have explored\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wXtaU7NEmY", "reviewer": "Reviewer_H3vx", "summary": "The paper presents a mechanistic analysis of internal activations of the Llama 3 70b model during three different in-context reinforcement learning tasks. In particular, the authors use Sparse Auto-Encoders (SAEs) to generate latent space representations of the residual streams and show how they correlate with TD Errors and Q-values of Q-learning agents trained to solve the same tasks. Furthermore, they show that relationship between the latent representations and the TD Errors is causal by intervening on the latent representations, which causes a degrading in the performance of the model.", "questions": "Some small additional comments and questions I had:\n- In the definition of the Q function in Section 2 (Methods, at page 2), shouldn't there be a conditioning on the initial state and action inside the expectation? Also, shouldn't the sum start from $t=0$ instead of $t=1$?\n- In Section 3, you claim that Llama 3 most likely implements \"classic\" Q-Learning rather than myopic Q-learning based on the negative log-likelihood. However, in Figure 2, looking at the correlations, it seems that the myopic Q-learning has in general comparable if not higher correlations to the latent representations. Couldn't this suggest that the model is implementing the myopic algorithm instead? Furthermore, is the difference in negative log-likelihood statistically significant?\n- In Figure 5, what do the horizontal lines in subplots B & C represent?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the paper is well written and the setting and the experimental details are generally well explained. The contributions are also clearly stated. Furthermore, as far as I can tell, the presented experimental methodology is also sound. Although it is a known fact that Transformers can potentially perform In-Context RL, especially if trained for it, it is the first time, to the best of my knowledge, that a mechanistic analysis is conducted on a model which was pre-trained on next token prediction. In addition, even if the methods used (e.g. SAEs) are already well established in the mechanistic interpretability literature, it is insightful to see how they can be successfully used also to better understand how LLMs solve In-Context RL. Hence, even if the problem of In-Context RL is well studied in the literature and the interpretability methods used are also well established, overall I think the presented results shed more light on the inner workings of how LLMs can solve RL tasks in-context, which can be significant and insightful for the community.", "weaknesses": "- The main weakness of the paper is that being an experimental work, I find the number of experiments conducted to be a bit limited. I think that more experiments should be conducted to further support the contributions of the paper (I saw that the authors mention this in future works/limitations, but I think the current paper would benefit from more ablation to make the evidence stronger). In particular, I suggest that the authors (as they also mention) should try to repeat the experiments they present with different models (at least one more) to prove that their results hold in general for \"big enough\" models. This would be really insightful since it would tell us that different models, even if trained differently, learn similar representations or make use of similar strategies to solve tasks. Furthermore, I think it would be insightful to conduct experiments on larger environments to better understand both to what extent these models are capable of performing In-Context RL and to analyze if, even at larger scale, these models still make use of TD Erros and Q-Values to solve the task\n- One minor concern regards the extent of the novelty of the work: as I mentioned above, although I agree with the authors that it is the first time (to the best of my knowledge) that it was shown that models trained on next-token prediction perform In-Context RL exploiting TD Errors, there are already quite some works exploring TD Learning in Transformers (both at a theoretical and experimental level). Furthermore, the methodology used for the mechanistic analysis is also already well established in the mechanistic interpretability literature.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4yC7Chgh9x", "reviewer": "Reviewer_fLyU", "summary": "This paper is way out of my expertise and hence I cannot provide a meaningful review.", "questions": ".", "limitations": "", "rating": 5, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": ".", "weaknesses": ".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m6Y2mlFgl1", "reviewer": "Reviewer_VvRq", "summary": "The paper looks for evidence of whether Llama3-70B model can simulate the TD-learning algorithm for solving RL tasks. The authors evaluate this using simple toy RL environments. They train different SAEs for the different tasks and find features that correlate highly with TD-errors and Q-values. They confirm that these features are causal in predicting actions by performing interventions on such features. Based on these evidence, the authors conclude that the LLM is simulating the TD-learning algorithm.", "questions": "1. In the plots on max correlation with values/errors (eg fig 2c, 2d, 3, 4b, 4c, etc.), is the correlation computed with the value/error of the action predicted by the LLM at the given state? If yes, then it would be valuable to check whether there are features that correlate with value/error of non-optimal actions. This could help in distinguishing whether the LLM is actually implementing TD-learning or the max-point episode algorithm provided above.\n2. Can you provide how the NLL score computed? I couldn't find it in the appendix either. Particularly, are you computing the log probabilities of Q-learning agent by doing a softmax using the Q-values over actions?\n3. Are you using any discount rate for the Grid World Task? If yes, please provide it.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "1. The study evaluates their hypothesis through a series of tasks to substantiate their empirical claims.\n2. Intervention experiment with the features to confirm their causal roles.\n3. The writing is clear and easy to understand. However, some details are missing. See in questions.", "weaknesses": "My main objection with the paper is that there is a simpler alternative hypothesis that could equally explain all of the results. Given the simplicity of the task, the LLM could be implementing the following algorithm to solve the tasks:\n\n\nStep 1: Keep a track of the maximum points for each episode in the context.\n\nStep 2: Predict the actions from the episode that has the maximum points.\n\n\nThis algorithm is simple to implement for the LLM given previous works on LLMs implementing greater than circuit [1] and induction heads [2]. Also, for the Two-Step Task, the first 7 episodes are provided by using a random policy, which should cover all the 4 different trajectories possible in the task.\n\nThe features that the authors find using SAEs could be features that are tracking the maximum points across episodes. These features will have high correlation with Q-values, and are also causal so interventions on them should show similar results as shown in the paper.\n\nI recommend that the authors conduct experiments designed to refute this hypothesis. See questions for some suggestions on experiments that can be performed.\n\nReferences:\n\n[1] How does GPT-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model. https://arxiv.org/abs/2305.00586\n\n[2] In-context Learning and Induction Heads. https://arxiv.org/abs/2209.11895", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)