Sparse Autoencoders Reveal Temporal Difference Learning in Large Language Models

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:

Step 1: Keep a track of the maximum points for each episode in the context.

Step 2: Predict the actions from the episode that has the maximum points.

We would like to thank the reviewer for engaging with our paper and for the constructive feedback. The reviewer raised an important point about alternative explanations for the results of the paper. We agree that the proposed mechanism could in principle account for some of our results and therefore conducted new experiments and analyses to test this alternative account. We find that TD learning predicts our results much more accurately than the alternative account proposed by the reviewer (which we called Look-up). Details as well as answers to the other questions are presented below.

In the Two-Step Task, we implemented Look-up with a softmax decision rule (in state $s$, it assigns an action a logit of 1 if the action was performed in s in the maximally rewarding episode), where we fitted a temperature parameter to the softmax to minimize the negative log-likelihood of Llama’s actions. We found that the Look-up model does not predict Llama’s behavior as well as a $Q$-learning agent (Page 19, Figure 11).

We also implemented the Look-up model for the Grid World task. Fitting it to the actions produced by the $Q$-learning agent, we see that its ability to predict actions matches Llama’s ability to predict actions (page 20, Figure 12 left). To further investigate this we tested Llama on a larger Grid World with 49 states (page 20, Figure 12 right), almost doubling the state space from the small Grid World. We also randomly initialized the starting position every episode, meaning that Llama had to generalize to predict the $Q$-learning agent’s actions. Here Llama’s action predictions diverged strongly from the Look-up model, indicating that it relied on a different mechanism than pure look-up to predict actions, which is likely TD learning.

Lastly, we showed that the Grid World SAEs showed much stronger correlations with $Q$ values compared to the Look-up’s estimates (Page 20, Fig. 13 left). Notably, these differences were further amplified in the new Grid World (Page 20, Fig. 13 right), paralleling our behavioral findings. Taken together, both our behavioral and representational results show stronger evidence for Llama using TD learning than Look-up.

(...) is the correlation computed with the value/error of the action predicted by the LLM at the given state?

We compute the correlations between the Q/TD of each action and the SAE features separately. What we plot is the average of the correlations across the different actions.

if yes, then it would be valuable to check whether there are features that correlate with value/error of non-optimal actions.

Please see Figure 10 in the Appendix, where we plot an example $Q$-value latent and the $Q$ value estimated by the RL model. We observe dips both in the SAE feature and the $Q$-value at certain points that are late in the task. If Llama does TD-learning, this is to be expected when Llama ends up in a state where a particular action is suboptimal, just like the $Q$-value. This cannot be explained by the Look-up model.

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?

Thanks for pointing out this missing detail. Indeed, we do a softmax over the $Q$-values to compute the log probabilities. This is now added to the appendix on page 16.

Are you using any discount rate for the Grid World Task? If yes, please provide it.

Yes, we use $\gamma = .99$ for all tasks. We have made this clearer in the appendix Line 822 now.

Thank you for your response and the new suggestion. We have modified the Look-up model so that it conditions on the initial state when predicting actions, just like the reviewer proposed. We see that this model’s action predictions do not match Llama’s action predictions in the 7x7 Grid World either (see the updated Figure 12, page 20). Since the state-space is combinatorially large, the Look-up mechanism cannot predict the $Q$-learning agent’s actions nearly as well as Llama, suggesting that Llama maintains representations of cached $Q$-values. Furthermore, we still see that $Q$-values correlate much more strongly with SAE features than the Look-up action-values (see updated Figure 13, page 20).

We also followed the reviewer’s suggestion to analyze the models’ predictions in states that don't occur at the start of the trajectories, but in the middle of other trajectories. Here the Look-up model predicts actions at random. Llama’s action probabilities are not only considerably higher, but also more closely match the action probabilities of the $Q$-learning agent with an $\epsilon$ -greedy policy (see Figure 14, page 21).

We hope this addresses the reviewer’s concern.

We thank the reviewer for their thoughtful review. We are glad the reviewer found our paper well-written and our methodology sound. The reviewer raised some important points, particularly, if our results generalize to other LLMs outside the Llama family. To address this, we repeated the majority of our analyses for all three tasks (two-step task, grid world, and graph learning) on two new models, Gemma-2-27b and Qwen2.5-72B. Notably, 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.

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)

Thanks for this suggestion. We have repeated all three experiments using two new models (Gemma-2-27b and Qwen2.5-72B). We place these findings in the Appendix, and here is a summary of the results:

• Two-Step Task (Page 23 Figures 15 and 16): For both Qwen and Gemma, we find SAE latents that correlate highly with TD errors and values. However, neither model can do the task as well as Llama. Consequently, the RL models predict behavior not as well for these models.
• Grid World (Page 24, Figures 17 and 18): The results are qualitatively the same. $Q$-learning values and error signals are more strongly correlated with SAE features than the value/error signals of the myopic model. We found both models can learn to predict actions from rewards.
• Graph Learning (Page 25, Figures 19 and 20): The results are qualitatively the same. Strong SR and TD correlations are found in both of these models, and the community structures emerge over the transformer blocks.

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

We tested Llama on a new grid world with two important differences from the grid world we initially tested: The new environment had $49$ states, almost double the size of the grid world we initially tested, which had $25$ states. In each episode, we randomized the initial location of Llama, requiring strong generalization.

The results are shown on page 20, Figure 12. We found that Llama can learn to predict the Q learning agent’s actions here as well, though performance is a bit weaker in this more challenging task. Importantly, since the starting position is randomly initialized in each episode, Llama cannot solely rely on looking up what the $Q$-learning agent did in the past to predict what it will do in the future. Furthermore, we observed strong correlations between TD/$Q$-values and SAE features, as shown on page 20 Fig. 13. These findings provide further evidence that Llama uses TD errors and $Q$-values in larger environments as well.

One minor concern regards the extent of the novelty of the work (...) 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

Past works have indeed shown that Transformers can implement RL algorithms when trained either to solve RL tasks directly or to predict action sequences produced by agents trained with an explicit RL objective (e.g. behavioral cloning) [1, 2]. However, training LLMs differs significantly from these training setups: LLMs are multi-billion parameter models trained on internet-scale text data using next-token prediction as its objective, whereas past work has investigated small-scale Transformers trained directly on state-action-reward triplets. We find LLMs doing TD learning surprising, given the major differences in their training compared to the other transformer models discussed. Moreover, past works have not investigated the mechanistic basis of the learning algorithms these Transformers implement after training. Our paper fills this gap by examining the mechanisms of in-context RL in LLMs. We added these points and references to the Related Work section.

Furthermore, the methodology used for the mechanistic analysis is also already well established in the mechanistic interpretability literature.

Indeed, SAEs and targeted interventions have seen wide applications in recent years in terms of uncovering the inner workings of LLMs, which is why we chose these methods. An important novelty of our work is that we use SAEs to identify learning algorithms implemented in-context, in contrast to identifying static concepts. This marks a novel use case of these models and demonstrates how they can be used to understand in-context learning better.

We would like to thank the reviewer for engaging with our work and for their encouraging feedback. We are happy that the reviewer thought our paper was excellent and easy to follow. The reviewer raised some minor points which we have addressed below.

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.

Thanks for raising this point. We have added the following clarification to our paper on line 119:

“For subsequent analyses, we rely on $Q$-learning, which is a variant of TD learning that learns a value function for the optimal policy.”

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?

In a control lesion, we set the latent unit that has the lowest correlation with Q/TD from a given layer to $0$. We have added the following clarification on line 330:

“We also conducted control lesions, where we set the activity of a latent unit from the same block with the lowest TD correlation to $0$.”

Lesioning means that we set the activation of a particular unit to 0. We have added the following clarification:

“Lesioning refers to setting the activations of specific units to $0$. This is also commonly referred to as zero ablation in the literature [1].”

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?

Thank you for raising this point. By clamping, we refer to multiplying the activity with a scalar. We replaced this term with scaling in the main text to improve clarity. The purpose of the negative scaling analyses was to show degraded Q/TD representations in subsequent blocks, as shown in Fig. 2G and Fig. 2H

Line 458, mangled sentence "our study is, we have explored".

Thanks, we corrected the sentence:

“While specific static concepts have been identified using SAEs, we have explored the characteristics of in-context learning algorithms using this technique.”

[1] Heimersheim, S., & Nanda, N. (2024). How to use and interpret activation patching. arXiv. https://doi.org/10.48550/arxiv.2404.15255