🤖 machine learning

Convergence and Emergence of In-Context Reinforcement Learning with Chain of Thought

This paper provides the first theoretical framework demonstrating that Chain-of-Thought generation in linear Transformers is equivalent to temporal difference learning, proving that it drives geometric convergence of policy evaluation error while emerging as a global minimizer of pretraining loss.

Original authors: Zixuan Xie, Xinyu Liu, Rohan Chandra, Shangtong Zhang

Published 2026-05-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zixuan Xie, Xinyu Liu, Rohan Chandra, Shangtong Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart, but slightly rigid, robot assistant. Usually, to teach this robot a new game, you have to stop it, reprogram its brain (update its parameters), and start over. This is slow and expensive.

In-Context Reinforcement Learning (ICRL) is a different approach. Instead of reprogramming the robot, you simply give it a "cheat sheet" or a history of what happened in the game right now. The robot looks at this history and figures out what to do next without changing its brain. It learns "on the fly."

This paper investigates a specific trick called Chain-of-Thought (CoT). You know how when humans solve a hard math problem, we don't just spit out the answer? We write down intermediate steps: "First I do this, then I check that, then I adjust..." CoT is the AI doing the same thing. It generates a sequence of "thoughts" before giving the final answer.

The authors of this paper wanted to understand why this "thinking out loud" makes the robot so much better at learning new tasks instantly. They used a simplified version of a modern AI (called a "Linear Transformer") to prove three main things:

1. The "Thinking" is Actually "Learning"

The authors discovered that when the AI generates a Chain of Thought, it isn't just rambling. It is mathematically equivalent to running a classic learning algorithm called Temporal Difference (TD) learning over and over again.

The Analogy: Imagine you are trying to guess the average temperature of a room.
- Without CoT: You take one quick guess and shout it out.
- With CoT: You take a guess, then you say, "Wait, that felt too high. Let me adjust based on the draft near the window." Then you say, "Okay, that's better, but the sun is coming in, so I'll adjust again."
- The paper proves that every time the AI writes a new "thought" (a new step in the chain), it is mathematically performing one precise update step of a learning algorithm. The more "thoughts" it generates, the more it refines its answer, just like a student correcting their work step-by-step.

2. The More You Think, The Closer You Get (But There's a Limit)

The paper shows that as the AI generates more thoughts (longer CoT), its answer gets closer and closer to the perfect solution. The error drops quickly at first, like a ball rolling down a steep hill.

The Analogy: Imagine you are trying to hit a bullseye on a dartboard while blindfolded, but someone is whispering corrections to you based on where your last throw landed.
- With 1 correction, you are still far off.
- With 10 corrections, you are very close.
- With 100 corrections, you are almost on the bullseye.
The Catch: The paper notes there is a "statistical floor." Even if you think forever, you can't get perfectly perfect if the information you were given (the context) was noisy or incomplete. It's like trying to guess the exact temperature of a room with a broken thermometer; no amount of thinking will fix the broken tool. The limit is determined by how much data you started with.

3. The AI "Discovers" the Right Way to Think

Perhaps the most surprising finding is about how the AI learns to do this. The researchers showed that if you train the AI on a bunch of different tasks (pretraining), it naturally "discovers" the specific internal settings (parameters) that make this "thinking = learning" process work.

The Analogy: Imagine a chef who has cooked thousands of different meals. They haven't been taught a specific recipe for "thinking." However, through experience, they naturally develop a specific way of tasting and adjusting the food that is mathematically the most efficient way to get a perfect dish. The paper proves that the "perfect way" the AI uses to think is actually the "best possible way" to minimize errors, and the AI finds this naturally during training.

Summary

In simple terms, this paper provides a mathematical proof that when an AI "thinks out loud" (Chain of Thought) to solve a new problem, it is secretly running a powerful learning algorithm inside its head. The more it thinks, the better it gets, up to the limit of the information it was given. Furthermore, the AI doesn't need to be explicitly told to do this; it learns to do it automatically because it's the most efficient way to solve the problems it was trained on.

The authors tested this on a simple, controlled environment (a mathematical puzzle called "Boyan's chain") and saw the AI naturally develop the exact internal structure needed to perform these learning steps, confirming their theory.

Technical Summary: Convergence and Emergence of In-Context Reinforcement Learning with Chain of Thought

Problem Statement

In-Context Reinforcement Learning (ICRL) enables agents to adapt to new tasks at inference time without parameter updates by conditioning on additional context, such as interaction trajectories. While empirical studies have demonstrated that Chain-of-Thought (CoT) generation can amplify ICRL capabilities, particularly in Large Language Models (LLMs), a theoretical understanding of how CoT interacts with ICRL remains limited. Existing theoretical works on ICRL primarily focus on supervised pretraining or idealized settings with known dynamics, leaving a gap in understanding the mechanics of CoT-driven policy evaluation under reinforcement pretraining and unknown dynamics.

Methodology

The authors analyze ICRL within a policy evaluation setup using a linear Transformer architecture. The core methodology involves:

Modeling CoT as Iterative Updates: The paper casts policy evaluation as an in-context weight learning task. Given a finite trajectory $\tau_n$ , the Transformer is required to approximate the Temporal Difference (TD) fixed point $w^*$ . The authors formalize CoT generation as a recursive process where the model generates a sequence of intermediate weight iterates $\{w_k\}$ , with each step conditioning on the growing prompt.
Linear Attention Analysis: The study employs linear self-attention (replacing softmax with the identity map) to facilitate theoretical tractability. The authors construct specific parameter matrices $(P, Q)$ for the linear attention layer.
Prompt Construction: A structured prompt matrix $Z_0$ is defined to separate the fixed context (trajectory data) from the iterate column. The CoT process appends new iterate columns to this prompt, allowing the model to perform iterative computations over the fixed context.
Convergence Analysis: The authors analyze the convergence of the generated iterates in two regimes:
- Known Dynamics: Where the transition matrix is available, and the context is constructed using expected successor features.
- Unknown Dynamics: Where the model relies on a single finite Markovian trajectory, introducing estimation errors and dependence issues.
Emergence Analysis: The paper investigates whether the specific parameters required for CoT to implement TD updates are learnable. They define an empirical pretraining objective based on the Empirical Update Norm (EUN) and analyze the global optimality of the constructed parameters.

Key Contributions

1. Equivalence of CoT and TD Updates

The paper provides a constructive proof that under specific parameter configurations $(P, Q)$ , a single linear Transformer layer with CoT generation is mathematically equivalent to repeatedly executing batch Temporal Difference (TD) learning updates on the context.

Theorem 1: Demonstrates that for a structured prompt, the output of the linear attention layer exactly produces the next batch TD iterate $w_{k+1}$ from $w_k$ . This establishes that CoT generation effectively simulates an iterative optimization process within the forward pass.

2. Inference-Time Convergence Guarantees

The authors establish convergence rates for the policy evaluation error (measured by Mean-Squared Projected Bellman Error, MSPBE) in two settings:

Known Dynamics: In an idealized setting, the error decreases geometrically with the number of CoT steps $k$ , requiring $O(\log(1/\epsilon))$ steps to reach precision $\epsilon$ .
Unknown Dynamics: In a more realistic setting with a single trajectory, the error still decreases geometrically but saturates at a statistical floor. This floor is determined by the effective sample size of the context length $n$ and scales as $\tilde{O}(\sqrt{d/n})$ . This result extends prior work by handling unknown dynamics and linear function approximation.

3. Emergence via Reinforcement Pretraining

The paper proves that the specific parameters $(P, Q)$ required to implement the TD updates are global minimizers of a TD-based pretraining loss.

Theorem 3: Shows that as the number of CoT steps $k \to \infty$ , the constructed parameters $\theta^*$ attain the global minimum of the empirical pretraining objective (EUN loss) over a dataset of trajectories.
This provides a theoretical justification for the empirical observation that reinforcement pretraining leads to the emergence of CoT capabilities that mimic RL algorithms.

Results

Theoretical Proofs: The paper rigorously proves that linear Transformers with CoT can implement iterated TD updates and converge to the TD fixed point, with convergence rates dependent on context length and step size.
Empirical Validation: Experiments conducted on Boyan's chain environment demonstrate the "emergence" of the predicted block-sparse structure in the learned parameters $(P, Q)$ . When trained via multi-task reinforcement pretraining, the model learns parameters that closely match the theoretically derived sparse matrices, confirming that the desired algorithmic behavior is learnable.

Significance

This paper offers the first theoretical characterization of how autoregressive Chain-of-Thought generation affects in-context policy evaluation in ICRL. By bridging the gap between CoT generation and iterative TD learning, the work explains why CoT amplifies ICRL capabilities: it allows the model to perform iterative internal computations that refine value estimates based on context. Furthermore, by proving that these algorithmic parameters are global minimizers of a standard pretraining loss, the paper provides a theoretical foundation for the emergence of in-context learning behaviors in reinforcement-pretrained agents, moving beyond the assumption that such behaviors require explicit supervised imitation of RL algorithms.