IC-062Llama 3 70B implements temporal difference learning in-context for reward-based RL, with causally relevant SAE features in its residual stream, while Llama 3 8B performs at chance
Can Demircan, Tankred Saanum, Akshay Kumar Jagadish, Marcel Binz, Eric Schulz
The paper prompts Llama 3 70B with a two-step MDP and a 5x5 grid world, and shows that its action choices are best explained by a q-learning model (NLL 2729) over a myopic model (NLL 2864) or a repetition model (NLL 5745). SAEs trained on the residual stream reveal features that correlate with TD errors (r = 0.58 in block 34 for the two-step task; r >= 0.75 in blocks 40+ for the grid world) and with q-values. Lesioning a single TD latent in block 34 significantly impairs task performance and makes behavior less consistent with q-learning, while control lesions have no effect. Negatively scaling the TD latent by 10 reduces subsequent correlations with both q-values and TD errors. The 8B variant performs at chance level, indicating the ability emerges with scale.
The SAEs are task-specific and not suitable for identifying RL-related variables for arbitrary tasks; the alignment between q-learning and Llama's behavior is not perfect, and a repetition preference may be needed to fully account for it.