Part 1
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.