Research map / Large language models / Language, safety and interpretability
Sparse Autoencoders Reveal Temporal Difference Learning in Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows Llama 3 70B performs in-context RL on three tasks and that sparse autoencoder latents in its residual stream track TD errors, Q-values and successor-representation learning, with lesion/clamp interventions causally affecting behavior.
Problem
How LLMs learn to solve reinforcement learning problems in-context, and whether their internal representations implement TD-learning-like computations despite being trained only on next-token prediction.
Approach
Train sparse autoencoders on residual-stream activations of Llama 3 70B during in-context RL prompts, find latents correlated with TD errors and Q-values, and causally test them via lesioning and clamping that replace residual activations with SAE reconstructions.
Key results
- Llama 3 70B's behavior on the Two-Step Task is best fit by a Q-learning model (NLL 2729) vs. myopic (NLL 2864) and repetition (NLL 5745) models.
- A block-34 SAE latent correlates r=0.58 with the Q-learning TD error; deactivating it significantly worsens Two-Step Task performance.
- In the graph task, SAE latents correlate more with the successor representation (max r=0.62) than with the transition matrix (max r=0.49), and lesioning a block-64 TD latent drops next-state accuracy from 24.3% to 18.9%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a Q-learning model rather than the myopic model | Behavioral model of Llama 3 70B · Two-Step Task | negative log-likelihood: 2864 -> 2729 | Helped |
| Use a Q-learning model rather than the repetition model | Behavioral model of Llama 3 70B · Two-Step Task | negative log-likelihood: 5745 -> 2729 | Helped |
| Deactivate the TD latent in block 34 (ablation) | Llama 3 70B · Two-Step Task | task performance | Hurt |
| Deactivate a low-TD-correlation SAE feature in block 34 (ablation) | Llama 3 70B · Two-Step Task | behavioral fit to Q-learning | No clear effect |
| Clamp the TD latent to -10 (ablation) | Llama 3 70B · Two-Step Task | correlation of later representations with Q-values and TD errors | Hurt |
| Provide correct rewards rather than randomly assigned rewards | Llama 3 70B · Grid World | action prediction accuracy | Helped |
| Deactivate four TD latents across blocks (ablation) | Llama 3 70B · Grid World | action prediction accuracy | Hurt |
| Lesion low-correlation control latents rather than TD latents (ablation) | Llama 3 70B · Grid World | action prediction accuracy | No clear effect |
| Deactivate or clamp the identified TD latents (ablation) | Llama 3 70B · Grid World | correlation of later representations with Q-values and TD errors | Hurt |
| Deactivate the block 64 TD latent (ablation) | Llama 3 70B · Graph Learning Task | next-state prediction accuracy: 24.3% -> 18.9% | Hurt |
| Deactivate the block 64 TD latent (ablation) | Llama 3 70B · Graph Learning Task | correlation of later SAE representations with the successor representation | Hurt |
Surprising findings
- The 8B Llama model performs roughly at chance on the Two-Step Task, suggesting in-context RL ability emerges with scale.
- TD-like representations emerge in an LLM trained only on next-token prediction, without an RL objective.
Limitations
- SAEs are task-specific and cannot identify RL variables for arbitrary tasks.
- Circuit-level mechanisms producing the TD representations remain unclear.
- Llama's behavior and TD-like representations are not perfectly aligned with the Q-learning model (e.g., repetition tendency).
Experimental setup
Datasets: THINGS database
Benchmarks: Two-Step Task 5x5 Grid World Graph Learning Task
Metrics: return negative log-likelihood action prediction accuracy Pearson correlation with TD error and Q-values next-state prediction accuracy CKA
Baselines: Q-learning model myopic Q-learning model repetition model transition matrix learning model Llama 3 8B
Builds on: Q-learning sparse autoencoders successor representation centered kernel alignment
Closest papers by meaning
- Reinforcement Learning Finetunes Small Subnetworks in Large Language Models (NeurIPS 2025 · similarity 0.89)
- LongRLVR: Long-Context Reinforcement Learning Requires Verifiable Context Rewards (ICLR 2026 · similarity 0.89)
- Supervised Knowledge Makes Large Language Models Better In-context Learners (ICLR 2024 · similarity 0.88)
- Can large language models explore in-context? (NeurIPS 2024 · similarity 0.88)
- In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering (ICML 2024 · similarity 0.88)
- TTRL: Test-Time Reinforcement Learning (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.