Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning
Authors: Jiuqi Wang
Ethan Blaser
Hadi Daneshmand
Shangtong Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: in-context reinforcement learning
policy evaluation
temporal difference learning
Primary area: reinforcement learning

Abstract: Traditionally, reinforcement learning (RL) agents learn to solve new tasks by updating their neural network parameters through interactions with the task environment. However, recent works demonstrate that some RL agents, after certain pretraining procedures, can learn to solve unseen new tasks without parameter updates, a phenomenon known as in-context reinforcement learning (ICRL). The empirical success of ICRL is widely attributed to the hypothesis that the forward pass of the pretrained agent neural network implements an RL algorithm. In this paper, we support this hypothesis by showing, both empirically and theoretically, that when a transformer is trained for policy evaluation tasks, it can discover and learn to implement temporal difference learning in its forward pass.

Published as a conference paper at ICLR 2025
TRANSFORMERS CAN LEARN TEMPORAL DIFFERENCE METHODS FOR IN-CONTEXT REINFORCEMENT LEARNING
Jiuqi Wang"
Ethan Blaser"
University of Virginia
University of Virginia
jiuqisemail. www.wiiiiiiiiaii..mm..
blaser virginia. edu
Hadi Daneshmand
Shangtong Zhang
University of Virginia
University of Virginia
dhadi@virginia.ed
shangtong@rirggiiaa.ee
ABSTRACT
Traditionally, reinforcement learning (RL) agents learn to solve new tasks by updat- ing their neural network parameters through interactions with the task environment. However, recent works demonstrate that some RL agents, after certain pretraining procedures, can learn to solve unseen new tasks without parameter updates, a phenomenon known as in-context reinforcement learning (ICRL). The empirical success of ICRL is widely attributed to the hypothesis that the forward pass of the pretrained agent neural network implements an RI. algorithm. In this paper, we support this hypothesis by showing, both empirically and theoretically, that when a transformer is trained for policy evaluation tasks, it can discover and learn to implement temporal difference learning in its forward pass.
INTRODUCTION
In reinforcement learning (RL, Sutton and Barto (2018)), an agent typically learns to solve new tasks by updating its neural network parameters based on interactions with the task environment. For example, the DQN agent (Mnih et al., 2015) incrementally updates the parameters of its Q-network while playing the Atari games (Bellemare et al., 2013). However, recent works (e.g., Duan al. (2016); Wang et al. (2016); Laskin et al. (2022)) demonstrate that RL can also occur without any parameter updates. These works demonstrate that an RL agent with fixed pretrained parumeters can take as input its observation history in the new task (referred to as context) and output good actions for that task. Specifically, let T, = (So, Ag, R St-1, A:-1, 1, R&) be a sequence of state-action- reward triples that agent obtains until time in some new task. This тe is referred to the context. The agent then outputs an action A, based on the context and the current state S, without updating its parameters. Notably, the context can span multiple episodes. As the context length increases, action quality improves, suggesting that this improvement is not due to memorized policies encoded in the fixed parameters. Instead, indicates that reinforcement learning process occurs during the forward pass as the agent processes the context, a phenomenon termed in-context reinforcement learning (ICRL), where RL occurs at inference time within the forward pass. See Moeini al. (2025) for a comprehensive survey of ICRL.
Previous works (e.g., Lin et al. (2023)) understand the ICRL phenomenon from the supervised pretraining perspective, where during the pretraining stage, the RL agent is explicitly tasked to imitate the behavior of some existing RL algorithms. It is therefore not surprising that the pretrained agent neural network implements the corresponding RL algorithm in its forward pass. This paper instead provides the first theoretical analysis of the emergence of ICRL from reinforcement pretraining, where in the pretraining stage, the RL agent is only asked to complete some task, but there is no
"Equal contribution. The order is determined by tossing a fair coin. Work performed while affiliated with MIT I. LIDS//oston University.

---

Published as a conference paper at ICLR 2025
constraint on how it should complete it. ICRL emerges in the sense that the agent neural network itself discovers and implements a certain RL algorithm in the forward pass.
Although most existing ICRL studies focus control tasks (i.e., outputting actions given state and context), to better understand ICRL, in this work, we investigate ICRL for policy evaluation, as it is widely known in the RL community that understanding policy evaluation is often the first step toward understanding control (Sutton and Barto, 2018). Specifically, suppose an agent with fixed pretrained parameters follows some fixed policy ι in a new task. We explore how the agent can estimate the value function Ux (8) for a given state 8 based on its context Tr without parameter updates. We call this in-context policy evaluation and believe that understanding will pave the way for a comprehensive understanding of ICRL. We demonstrate that ICRL can emerge even with simplified neural network architectures, e.g., transformer (Vaswani et al., 2017) with linear attention
1:30
1.25
MSVE 1.00
TFp. v(s)
So Au, Ri,..., A,-1,A,-1 Res
~ Ux (s)
0.50
input: context + query
oulpul
0.25
Contest Tength (s)
Figure 1: transformer capable io-context policy evaluation. This 15-layer transformer TFg. takes the context and a state of interest s as input and outputs TFo, (Tr, s) as the estimation of the state value U, (s). The y-axis is the mean square value error (MSVE) E. d (T) (T,, - u* (s)) with dz (s) being the stationary state distribution. The curves are averaged over 300 randomly generated policy evaluation tasks, with shaded regions being standard errors. The tasks vary in state space, transition function, reward function, and policy. Yet a single 0, is used for all tasks. See Appendix B for more details.
Figure provides a concrete example of a transformer capable of in-context policy evaluation. To our knowledge, this the first empirical demonstration of in-context policy evaluation. Let TF.. denote the transformer used in Figure 1 with parameters 0,. Figure demonstrates that the value approximation eтror of this transformer drops when the context length increases even though 6, remains fixed. Notably, this improvement cannot be attributed to 8, hard-coding the true value function. The approximation error Figure I averaged over wide range of tasks and policies, each with distinct value functions, while only single 0. is used. The only plausible explanation seems to be that the transformer TFg. able to perform some policy evaluation algorithm in the forward pass to process the context and thus predict the value of This immediatcly raises two key questions:
(QI) What is the specific policy evaluation algorithm that TF., is implementing? Q2) What kind of pretraining can generate such a powerful transformer?
This work aims to answer these questions to better understand in-context RL for policy evaluation. To this end, this work makes three contributions.
First, we confirm the existence of such a 0, by construction. We prove that this 0, enables in-context policy evaluation because the crration-by-iteration updates of batch version of temporal difference learning (TD, Sutton (1988)). To summarize, a short answer to (Q1) is "TD." This is the first time that a pretrained agent neural network for ICRL is fully white-boxed. Furthermore, also prove by construction that
forward pass of TFp. is precisely equivalent to the
'ee, of course, also need to provide the discount factor to the agent. We ignore it for now to simplify the presentations.
'Linear attention a widely used transformer variant for simplifying both computation and analysis (Katharopoulos et al., 2020; Wang et al., 2020; Schlag et al., 2021; Choromanski et al., 2020; Mahankali et al., 2023; Ahn et al., 2023; Von Oswald et al., 2023a;b; Wu et al., 2023; Ahn et al., 2024; Gatmiry et al., 2024; Zhang al., 2024; Zheng al., 2024; Sander al., 2024).
2

---

Published as a conference paper at ICLR 2025
transformers are capable of implementing many other policy evaluation algorithms, including TD(A) (Sutton, 1988), residual gradient (Baird, 1995), and average reward TD (Tsitsiklis and Roy, 1999).
Second, we empirically demonstrate that this 0.. naturally emerges after regard TF2 a standard nonlinear function approximator and train it using nonlinear TD on multiple randomly generated policy evaluation tasks (similar training a single DQN agent on multiple Atari games). This empirical finding is surprising because the pretraining only drives TF2 to output good value estimates. There is no explicit mechanism that forces the transformer's weights to implement TD in its forward pass (cf. that the forward pass of DQN's Q-network can be anything as long as it outputs good action value approximations). Despite having the capacity to implement other algorithms like residual gradient, the pretraining process consistently leads the transformer weights to converge to those that implement TD. This observation parallels the historical development of the RI. community itself, where TD became the favored method for policy evaluation after extensive trial-and-error with alternative approaches. Thus, a short answer to Question (Q2) is also "TD.." Naturally, this leads to our third and final question.
(Q3) Why does TD pretraining give rise to in-context TD?
Our third contribution addresses this question by proving that the parameters O. that implement TD in the forward pass lie in an invariant set of the TD pretraining algorithm. is, of course, not complete answer. Similar to Wu et al. (2023); Zhang et al. (2024), we only prove the single-layer case, and we do not prove that the parameters will surely converge to this invariant set. However, we argue that our invariant set analysis and the techniques developed to prove it are significant step toward future work that can fully characterize how in-context reinforcement learning emerges from reinforcement pretraining.
2 RELATED WORKS
Our first question (Q1) is closely related to the expressivity of neural networks (Siegelmann and Sontag, 1992: Graves et al., 2014; Jastrzębski et al., 2017: Hochreiter et al., 2001; Lu et al., 2017). Per the universal approximation theorem (Hornik al., 1989; Cybenko, 1989; Leshno et al., 1993; Bengio et al., 2017), sufficiently wide neural networks can approximate any function arbitrarily well. However, this theorem focuses only on input-output behavior, meaning that given the same input, the network will produce similar outputs as the targer function. It does not say anything about how the forward pass is able to produce the desired outputs. nor how the number of layers affects the approximation error. In the supervised learning community, there are few works that are able to white-box the forward pass of neural networks to some extent (Frosst and Hinton, 2017; Alvarez Melis and Jaakkola, 2018; Chan et al., 2022; Yu et al., 2023; Von Oswald et al., 2023a; Ahn et al., 2024). But in the RL community, this work is, to our knowledge, the first to white-bat box how the forward pass of pretrained transformer can implement RL algorithms. Notably, Lin et al. (2023) also construct some weights of transformer such that the forward pass of the transformer can implement some RL algorithm. However, their constructed transformer is overly complicated and there is no evidence that their weight construction can emerge through any kind of pretraining in practice.
Our second question (Q2) is closely related to the pretraining in ICRL, which can be divided into supervised pretraining and reinforcement pretraining. In supervised pretraining, the agent is explicitly tasked with imitating the behavior of some existing RL algorithms demonstrated in an offline dataset (Xu et al., 2022; Laskin al., 2022; Raparthy et al., 2023; Sinii al., 2023; Zisman al., 2023; Shi et al., 2024; Dai et al., 2024; Huang et al., 2024a;b; Kirsch et al., 2023; Wang et al., 2024; Lee et al., 2024). It is thus less surprising (Krishnamurthy et al.. 2024) that the pretrained agent network does implement some RL algorithm in the forward pass. Supervised pretraining can be understood through the lens of behavior cloning, see Lin et al. (2023) for theoretical analysis of supervised pretraining. In reinforcement pretraining, the agent is only tasked with maximizing the return, and there is no constraint on how the agent network should achieve this in the forward pass (Duan et al., 2016; Wang et al., 2016; Kirsch et al., 2022; Bauer et al., 2023; Grigsby et al., 2023: Lu et al., 2023; Park et al., 2024; Xu et al., 2024; Grigsby et al., 2024; Cook et al., 2024). Reinforcement pretraining also is closely related to algorithm discovery in meta RL. The difference is that the reinforcement pretraining in ICRL discovers and implements novel RL algorithm in the forward pass without parameter updates, while large body of prior works of algorithm discovery in meta RL. (Kirsch et al., 2019; Oh et al.,
3

---

Published as a conference paper at ICLR 2025
2020; Lu et al., 2022) require parameter updates when executing the discovered algorithm. See Beck et al. (2023) for a comprehensive survey of meta RL and see Mocini et al. (2025) for a comprehensive survey of different pretraining methods for ICRL. The reinforcement pretraining method we use is very simple version of multi-task RI. and is very standard in the meta RI. community. We do not claim any novelty in our pretraining method. Instead, the novelty lies in the empirical and theoretical analysis of this simple yet standard pretraining method.
Our third question (Q3) is closely related the dynamics of RL algorithms (Borkar and Meyn, 2000; Bhandari et al., 2018; Cai et al., 2019; Qian et al., 2024; Liu et al., 2025), which is an active research area. In the context of ICRL, the dynamics of supervised pretraining is previously studied in Lin et al. (2023) following the behavior cloning framework. This work is to our knowledge the first theoretical analysis of reinforcement pretraining for ICRI..
ICRL is broadly related to the general in-context learning (ICL) community in machine learning (Garg et al., 2022; Müller et al., 2022; Akyürek et al., 2023; Von Oswald et al., 2023a; Zhao al., 2023; Allen-Zhu and Li, 2023; Mahankali et al., 2023; Ahn et al., 2024; Zhang et al., 2024). While ICL widely studied in the context of large language models (LLMs) (Brown et al., 2020), ICRL and LLM-based ICL represent distinct areas of research. ICRL typically needs RL-based pretraining while LLM's pretraining is usually unsupervised. Additionally, ICRL focuses on RL capabilities during inference, while LLM-based ICL typically examines supervised learning behavior during inference. RL and supervised learning are fundamentally different problems, and similarly, ICRL and in-context supervised learning (ICSL) require different approaches. For example, Ahn al. (2024) prove that ICSL can be viewed as gradient descent in the forward pass. While our work draws inspiration from Ahn et al. (2024), the scenario in ICRL is more complex. TD, which we analyze in this paper, is not equivalent to gradient descent. Our proof that transformers can implement TD in the forward pass is, therefore, more intricate, especially when extending it to TD(A) and average reward TD. Moreover, Ahn et al. (2024) consider a gradient descent- based pretraining paradigm where the transformer trained minimize an in-context regression loss. As a result, they analyze the critical points of the regression loss to understand their pretraining. By contrast, we consider TD-based pretraining, which is not gradient descent. To address this, we introduce a novel invariant set perspective to analyze the behavior of transformers under TD-based pretraining.
3 BACKGROUND
Transformers and Linear Self-Attention. All vectors are column vectors. We denote the identity matrix in R by I, and an m n all-zero matrix by 0mxn- We use ZT to denote the transpose of
Z and use both (zz y) and x to denote the inner product. Given prompt Z E Rdx=, standard single-head self-attention (Vaswani et al., 2017) processes the prompt by Attnw,,w,,.r w. (Z) = W. Z softmax Z W W,Z). where W. E Raxd W E Rmxd and W, E R represent the value, key and query weight matrices. The softmax function is applied to each row. Linear attention is a widely used architecture in transformers (Mahankali et al., 2023; Ahn et al., 2023; Von Oswald et al., 2023a:b; Wu al., 2023; Ahn al., 2024; Gatmiry et al., 2024; Zhang et al., 2024; Zheng et al., 2024; Sander et al., 2024), where the softmax function is replaced by an identity function. Given a prompt 7 E R (m+++ linear self-attention is defined as
(Z; = PZ (Z QZ)),
(1)
where P E (224+))((24+ and Q (224+))((2++1) are parameters and M R(++1)((++1) is fixed mask of the input matrix Z, defined as
T
0 0
M
(2)
01xn
Note that we can view P and Q as reparameterizaaions of the original weight matrices for simplifying presentation. The mask M is introduced for in-context learning (Von Oswald et al., 2023a) to designate the last column of 7 as the query and the first n columns as the context. We use this fixed mask in most of this work. However, the linear self-attention mechanism can be altered using different mask M', when necessary, by defining LinAttn(2 (; P,Q, M') = PZM'(Z' QZ). In an L-layer transformer with parameters ((P), Q ho...r.... the input Zo evolves layer by layer
Z41 = 2, + LLinAtttn (7t) Zt + P.Z. (Z)Q.))))
(3)
4

---

Published as a conference paper at ICLR 2025
Here, is a normalization factor simplifying presentation. We follow the convention in Von Oswald et al. (2023a); Ahn et al. (2024) and use
(Zo {P, Qth Qth - 1, +1
(4)
to denote the output of the L-layer transformer, given input Zo. Note that Zi (d 1, n, 1] is the bottom-right element of Z, Equation (4) establishes the notation convention that we adopt to define the output of L-layer transformer. Specifically, línear attention produces a matrix, but for policy evaluation, we require a scalar output. Following prior works, we define the bottom-right element of the output matrix as this scalar.
Reinforcement Learning, We consider an infinite horizon Markov Decision Process (MDP, Puterman (2014)) with a finite state space S. finite action space A, reward function MMDP : S A E. a transition function PMDP: & x S x A + [0, 1J. discount factor E [0, 1), and an initial distribution Po : + [o, 1]. An initial state So is sampled from po- At time t, an agent at state S, takes an action A, ~ T-[Ss,), where x A x S - [0, 1] is the policy being followed by the agent, receives a reward Re+1- = (M At), and transitions to a successor state St+1 ~ PMMP(-|Ss At). the policy is fixed, the MDP can be simplified to Markov Reward Process (MRP) where transitions and rewards are determined solely by the current state:S(+1 ~ p(-[S,) with R++1 = r(S)). Here, p(s'ls) = E. ((MDs))MDP((' (s. a) and r(s) = =.. MDD(s) (s, a). In this work, we consider the policy evaluation problem where the policy æ is fixed. So, it suffices to consider only an MRP represented by the tuple (Po, (Po.p,r), and trajectories (So, R1, S1, R2.... ...) sampled from it. The value function of this MRP is defined as v(s) - E function v is one of the fundamental tasks in RL. To this end, one can consider a linear architecture. Let Φ : S + Rd be the feature function. The goal is then to find weight vector w E I" such that for each 8, the estimated value i(s; w) = w (08) approximates u(s). TD is a prevalent method for learning this weight vector, which updates w iteratively as
Li=++411 7 R2 SS
-00
Estimating the value
W/++1 =t;; 02 (Re41 + 70 (Sc-1: w.) - 0 (Si; wx)) V (Sc: w.) =t. + 04 (Re+1 + ywu (S.+1) - w. (SS)) ((S)),
(5)
where {a,} } is a sequence of learning rates. Notably, TD is not gradient descent algorithm. It is instead considered as a semi-gradient algorithm because the gradient is only taken with respect to 6 (Sr,w) and does not include the dependence on i (St+1; (Sutton and Barto, 2018). Including this dependency modifies the update
W++1 = We + ar (Re+1 + ywu ((S(+1) - w, ((S)) (o(S)) - -10(0++11)))
(6)
known the (naive version of) residual gradient method (Baird, 1995).3 3 The update in (5) is also called TD(0) - a special case of the TD(A) algorithm (Sutton, 1988). TD(A) employs an eligibility trace that accumulates the gradients as e-1 = 0, e, = - + o(S,) and updates w iteratively as Wi+11 W. + au [Re+1 + y ((S++1) - w, ((S.))e; The hyperparameter À controls the decay rate of the trace. If X = 0, we recover (5). On the other end with = 1, known that TD(A) recovers Monte Carlo (Sutton, 1988). Another important setting in RI. is the average-reward setting (Puterman, 2014; Sutton and Barto, 2018), focusing on the rate of receiving rewards, without using a discount factor 7. The average reward F is defined as F + limy 1 Σ E[R,]. Similar to the value function in the discounted setting, differential value function v(s) is defined for the average-reward setting as t(s) = E |E 1 (R, F)[ s). One can similarly estimate E(s) using a linear architecture with a vector w w Φ(s). Average-reward TD (Tsitsiklis and Roy, 1999) updates iteratively as W/++1 W) + ar [Re+1 Fr+1 + w, ((S++1) !! ((S.)))((i) where R, is the empirical average of the received reward.
Fz Ei= =
4 TRANSFORMERS CAN IMPLEMENT IN-CONTEXT TD(0)
In this section, we reveal the parameters of the transformer used to generate Figure 1 and answer (QI. Namely, we construct that transformer below and prove that implements TD(0) its forward
This is naïve version because the update does not account for the double sampling issue. We refer the reader Chapter 11 of Sutton and Barto (2018) for detailed discussion.
5

---

Published as a conference paper at ICLR 2025
pass. Given trajectory (Su, R1, S1, R2, S3, R4. ................................ Sn) sampled from MRP, using shorthand
6 ((S)). we define for 1 = 0, 1.........
Φo
Øn
O2ax2d 0xx2d
0zax1 1
Gi
CJT
0ax1
Zo = 701
you
0
PID
Qi
Oaxd
Oaxd
0ax1
(7)
R1
R₂
0xxd
0xxd
0
Here, Zn E g(22+1)x(x+)) is the prompt matrix, Ct E Raxd is an arbitrary matrix, and {(P)D, Q11) t0..,,......- the parameters of the L-layer transformer. We then have
Theorem (Forward pass as TD(O)). Consider the L-layer linear transformer followiig (3), using the mask (2), parameterized by {P7D QTD) 11=0,...,-11 in (7). Let yi" (n-1) be the bottom right element of the 1-th layer's output, i.e,, Vi
(n+1) = Z [2d + 1, 1). Then, it holds that vi (n+1) (Ón: wt).
where {wi} is defined as wo = 0 and
w1+1 = WT + C. Eno (Rz+1 + + Φ5+1 - -wi (bj) dj--
The proof is in Appendix A.1 and with numerical verification in Appendix H as sanity check. Notably, Theorem holds for any Ct- In particular, if C. a (this used in the transformer to generate Figure 1), then the update (8) becomes batch version of TD()) in (5). For a general Cr the update (8) can be regarded as preconditioned batch TD(0) (Yao and Liu, 2008). Theorem precisely demonstrates that transformers forward pass. We call this in-context TD. It should be noted that although the construction of Zu in (7) uses 6, as the query state for conceptual clarity, any arbitrary state E S can serve the query state and Theorem still holds. In other words, by replacing on with ((s), the transformer will then estimate o(s). Notably, if the transformer has only one layer, i.e., L 1, there are other parameter contigurations that can also implement in-context TD(0).
expressive enough to implement iterations of TD in its
Corollary 1. Consider the 1-layer linear transformer followingg (3), using the mask (2). Consider the following parameters
- CT Oaxd 01xd
Oaxa Oaxea
0ax 1
POD = O2ax2a
02dx1
2 =
0ax1
01x2d
01xd
Then, it holds that V1
(n+1) - (dn; wi), where wi is defined
wi wo + C = (Ry+1 Tno ΦJ+1 wh Φb)) Øj with wo- 0.
The proof is in Appendix A.2. observant reader may notice that this corollary holds primarily because wh- = O, making it unique result for L = 1. Nevertheless, this special case helps understand a few empirical and theoretical results below.
TRANSFORMERS Do IMPLEMENT IN IN-CONTEXT TD(0)
In this section, we reveal our pretraining method that generates the powerful transformer used in Figure 1, answering (Q2).* 4 We also theoretically analyze this pretraining method, answering (Q3).
Multi-Task Temporal Difference Learning. In existing ICRL works for control, the transformer takes the observation history as input and outputs actions. A behavior cloning loss is used during pretraining ensure that the transformer outputs actions similar to those in the pretraining data. In contrast, our work seeks to understand ICRL through the lens of policy evaluation, where the goal for the transformer to output value estimates rather than actions. To ground the value estimation, we use the most straightforward method in RL: the TD loss. This yields pretraining algorithm (Algorithm 1) where the transformer trained using nonlinear TD on multiple tasks. We call multi-task TD.
Recall that a policy evaluation task is essentially a tuple (po, p, r, 6). In Algorithm 1, we assume that there is task distribution dusk over those tuples. Recall that TF, (Zo: 0) and TF, 0) are
"The implementation is available at https: //gi thub. can wwwrsequentiiiieooooooo a1 Intel 1 gigenceellle lab/ InContextTD

---

Published as a conference paper at ICLR 2025
Algorithm 1: Multi-Task Temporal Difference Learning
1: Input: context length n,, MRP sample length т, number of training tasks k, learning rate o, discount factor 7, transformer parameters 0 = (Pi, Qt
2: for i f to & do
3: Sample (Po- p. r, Φ) from dask
4: Sample (So, R1, S1, R2 .2,....,.5 S,, R7+1 ST+1) from the MRP (po.p,))
5: for O,...,T= - n do
61
Wt+n-1 70
Øt+n+1
01+1
Oten
Pt+n+
6:
Zo -
70++1
4g +
701+2
70++++1
Re+1
Runn
0
Ra+2
Rotn+1
7: 0 - 0 + a [ee+t+2 + (Z);
(Zo (
// TD
8: end for
9: end for
intended to estimate o(S(+n+1) and ((Sen++2) respectively. So, Algorithm essentially applies TD using (St+n+1, +2+ S(+n+2) to train the transformer. Ideally, when new prompt Ziesi is constructed using a trajectory from a new (possibly out-of-distribution) evaluation task (Po.Pt P, r, &)ea. the predicted value TFy 0) with 0 from Algorithm should be close to the value of the query state in Zest- This problem is a multi-task meta-learning problem, a well-explored area with many existing methodologies (Beck et al., 2023). However, the unique and significant aspect of our work is the demonstration that in-context TD emerges in the learned transformer, providing a novel explanation for how the model solves the problem.
Empirical Analysis. We first empirically study Algorithm To this end, we construct djakk based on Boyan's is chain (Boyan, 1999), canonical environment for diagnosing RI algorithms. We keep the structure of Boyan's 's chain but randomly generate initial distributions po, transition probabilities reward functions r, and the feature function ø. Details of this random generation process are provided in Algorithm 2 with Figure 3 visualizing Boyan's 's chain, both in Appendix C.
For the linear transformer specified in (3), we first consider the autoregressive case following (Akyürek et al., 2023; Von Oswald et al., 2023a), where all the transformer layers share the same parameters, i.e., P1- Po and Q: = Qu for (T. (I=3), Importantly, all elements of Po and Qo are equally trainable - we did not force any element of Po or Qo to be O. We then run Algorithm with Boyan's chain-based evaluation tasks (i.e., dast) to train this autoregressive transformer. The dimension of the feature is d = 4 (i.e., d(s) E IR. Other hyperparameters of Algorithm are specified in Appendix D.I.
l......... L- We consider a three-layer transformer
Figure 2a visualizes the final learned Po and Qo by Algorithm I after 4000 MRPs (i.e., k: 4000), which closely match our specifications pTD and QTD in (7) with C La. Figure 2b, we visualize the element- wise learning progress of Po and Qo- We observe that the bottom right element of Po increases (the Po[-1, -1] curve), while the average absolute value of all other elements remain close
(the "Avg Abs Others" curve), closely aligning with PTD up some scaling factor. Furthermore, the trace of the upper left d x d block of Qo approaches d (the tr(Qol: [:: d. d) curve), and the trace of the upper right block (excluding the last column) approaches (the tr tr(Qol: d, d 2d)) curve). Meanwhile, the average absolute value of all the other elements in Qo remain near zero, aligning with QTD using C₁ = Ia up to some scaling factor.
More empirical analysis is provided in the Appendix. In particular, besides showing the parameter- wise convergence in Figure 2, we also use other metries including value difference, implicit weight similarity, and sensitivity similarity. inspired by Von Oswald et al. (2023a); Akyürek et al. (2023), to examine the learned transformer. We also study normal transformers without parameter sharing (Appendix D.3), as well as different choices of hyperparameters in Algorithm 1. Furthermore, we empirically investigate the original softmax-base conducted experiments where we constructed dunk based the Cartpole environment (Brockman et al., 2016) (Appendix F). The overall conclusion is the same in-context TD emerges in the transformers learned by Algorithm I. Notably, Theorem I and Corollary I suggest that for L 1, there are two distinct ways to implement in-context TD (i.e., (7) v.s. (9)). Our empirical results in Appendix D.2 show that Algorithm I ends up with (9) in Corollary 1 for I. 1, aligning well with
transformers (Appendix E). Finally, we also
7

---

Published as a conference paper at ICLR 2025
Final
Final Qq
Fl 3oriic
2.0
PC:
(10k) 4::4)
1.5
Aug Ale Others
Abe Ohhers
10-
40 0.5
0.0
3300
3500
3003
0
1800
KSR0s 500
(a) Learned Po and Qo after 4000 MRPs
(b) Element-wise learning progress of Po and Qu
Figure 2: Visualization of the learned transformers and the learning progress. Both (a) and (b) averaged across 30 seeds and the shaded regions in (b) denotes the standard errors. Since Po and Qo in the same product in (I), the algorithm can rescale both or flip the sign of both, but still end up with exactly the same transformer. Therefore, to make sure the visualization are informative, we rescale Po and Qo properly first before visualization. See Appendix D.1.1 for details.
Theorem 2. For T. 2,3, Algorithm always ends up with (7) in Theorem 1, shown in Figure 4 in Appendix D.2. We also empirically observed that for in-context TD to emerge, the task distribution dask has to be "difficult" enough. For example, if (po: p) or & are always fixed, did not observe the emergence of in-context TD.
Theoretical Analysis. The problem that Algorithm 1 aims solve is highly non-convex and non-linear (the linear transformer is still nonlinear function). We analyze a simplified version of Algorithm and leave the treatment the full version for future work. particular, we study the single-layer case with L = 1, and let 0 = (Po. Qo) be the parameters of the single-layer transformer. We consider expected updates, i.e.,
0k+1 =0x ((,) with (0) = E [R + >TF (Z) 0) TF, (Zo, VTF (Zo, 0)] (10) Here, the expectation integrates both the randomness in sampling (po. p. Ø) from dask and the randomness in constructing (R, Zo, Z/) thereafter. We sample (So. R1, ......, Sn+1. Rx42, Sn42)
following (Po, p)r) and construct using shorthand Φ; = (S:)
Φn
Φn-1
Øn+1
Φ1
On
@nt2]
Zo =
701 R
yOu R₂
1/ =
702
70++1
R = Rn+2-
(11)
R
Rm
The structure of Zo and Z/ is similar to those in Algorithm 1. The main difference is that we do not use the sliding window. We recall that (Po. p, r, o) are random variables with joint distribution dask- Here, @ is essentially random matrix taking value in Rdxls represented Φ [d(s])ees. We - to denote "equal in distribution" and make the following assumptions. Assumption 5.1. The random matrix & is independent of (po.P,")
Assumption 5.2. II & o,Ao & o, where II any d-dimensional permutation matrix and A is any diagonal matrix in Rd where each diagonal element of A can only be -1 1.
Those assumptions are easy to satisfy. For example, as long as the elements of the random matrix are i.i.d. from symmetric distribution centered at zero, e.g., a uniform distribution on [-1, then both assumptions hold. We say a set θ is an invariant set (10) for any k 0g E e = 0x+1 E O. Define
O2axxdd 01x2d
02dx1 1)
cTa Odxa
0ax1
0. (m(,,c,c) = Po
Qo =
c'Ta Odxd
0ax1
0
01xd
0
Theorem 2. Let Assumptions 5.1 and 5.2 hold. For the construction (11) of (R, Zu. Z%). the set O. (0, d)),,, c, d' E [R} is an invariant set of (10).
The proof is in Appendix A.3. Theorem 2 demonstrates that once 0g enters O, some k, it can never leave, i.e., θ, is a candidate set that the update (10) can possibly converge to. Consider a subset O C O, with stricter constraint c' = 0, i.c., O = {#. (n, c, 0) E R]. Corollary 1 then confirms that all parameters in O,, implement in in-context TD. That being said, whether (10) is guaranteed converge to O., further to O', is left for future work.
8

---

Published as a conference paper at ICLR 2025
TRANSFORMERS CAN IMPLEMENT MORE RL ALGORITHMS
In this section, we prove that transformers are expressive enough to implement three additional well- known RL algorithms the forward pass. We warm up with the (naive version of) residual gradient (RG). We then move to the more difficult TD(A). This section culminates with average-reward TD, which requires multi-head linear attention and memory within the prompt. We do note that whether those three RL algorithms will emerge after training is left for future work.
Residual Gradient. The construction RG is an casy extension of Theorem We define
0ax1
PRO PP Qi
Ci
0ax1
E (2224+1)((2+1))
(12)
01xd
Oxxd
Corollary 2 (Forward pass Residual Gradient). Consider the L-layer linear transformer following
(3), using the mask (2), parameterized by {P!c QP()) (2]")__,.... I=.0...,..-1 in (12). Define y) (n+1) = [2(2d
1, + 11. Then, it holds that yi fi=+1 (m+1) = - wi), where {wi} is defined as wi = 0 and
w1+1 wy + C (Rz+1 + +wu ØJ+1 wi Φ5) (Ø) - 700++).
(13)
The proof is in A.4 with numerical verification in Appendix as samity check. Again, if C La then (13) can be regarded batch RG. Notably, Figure empirically demonstrates that Algorithm I eventually ends up with in-context TD instead of in-context RG. This observation aligns with the conventional wisdom in the RL community that TD usually superior to the naïve RG (see, e.g., Zhang et al. (2020) and references therein).
batch version of (6). For general C₁, it is then preconditioned
TD(A). Incorporating eligibility traces is an important extension of TD()). We now demonstrate that by using a different mask, transformers are able to implement in-context TD(A). We define
À
0
0 0
M TD(A)
: E R(x+1)((++1)
(14)
A=-1
X=-2
An-3
A
0
0
0
0
0 0
Notably, if X O, the above mask for TD(A) recovers the mask for TD(0) in (2).
Corollary 3 (Forward pass as TD(A)). Consider the L-layer linear transformer parameterized by specified in (7) with the input mask used in (3) being MTDOO in (14).
{P!" PD) Ji=0...,F.-1
Define yi
(n+1) _ Zi(2d 1,n 1]. Then, it holds that yi (n+1) -ddn wi) where {wi} is defined
with wh 0, en = 0, Ej- = -ee-- + Øj. and
#k+1 = wk + L. EnD (re41 + Two di+1 w. dz)e:.
The proof is in A.5 with numerical verification in Appendix H as sanity check.
Rverd TD. We now demonstrate that transformers are expressive enough implement in-context average-reward TD. Different from TD(0), average-reward TD (Tsitsiklis and Roy, 1999) exhibits additional challenges in that it updates two estimates (i.e., wy and Ft) in parallel. To account for this challenge, we use two additional mechanisms beyond the basic single-head linear transformer.
Namely, we allow additional "memory" in the prompt and consider two-head linear transformers. Given trajectory (So, Ri, S1, R2, S3, R4,... Sx) sampled from an MRP, we construct the prompt matrix Zu as
Φn-1
2n
Zo
Φ1
On
R (xd+)))
Ri
Rn
0
9

---

Published as a conference paper at ICLR 2025
Notably, the last row of zeros is the "memory", which is used by the transformer to store some intermediate quantities during the inference time. We then define the transformer parameters and masks
22x2d
02dx1
02dx1
O2dx2d 01x2d
02dx1 0
02dx1
(1)
(D.
P1
01x2d
1
0
Pt
=
0
(15)
01x2d
0
0
01x2d
0
1
-CJ
C
Odx2
02ax2d
02ax1
(24x(2442) 01x(2)+2)
Ozax
Q
Oaxd
Oaxd
0xx2
Wi
(16)
01x2d
02xd
02xd
02x2
M (2)))
I
Onx1
.M Tb.(
(In+1 - U Un+ Un+idiag([
n+1 (+1l)) (2)
,(17)
0 1xx
0
2
where C Rddd is again an arbitrary matrix, Un41 is the (n + 1) (n + 1) upper triangle matrix where all the nonzero elements are 1, and diag(z) constructs diagonal matrix, with the diagonal entry being x. Here,
{PI (PTB(()) Q1D
are the parameters of the first attention heads, with
1 1-1
the input mask being M (2)
QTD
1=0....L--1 are the parameters of the second attention
heads, with the input mask being M TD.(2) The two heads coincide on some parameters. W2 is the affine transformation that combines the embeddings from the two attention heads. Define the LinAttt( Z; Q, M) LinAunn Z; P', Q' M')
two-head imear-attention as TwoHead P,Q, P','P'''' Q' M", WW) w
The L-layer transformer we are interested in then given by
Z = Zi + bh TwoHead P/ (T)(1) QT M DD.1) P
TD.22
Q TD Wi).
(18)
Theorem 3 (Forward pass as average-reward TD). Consider the L-layer transformer in (18). Let 1i (n++1) be the boltom-right element of the 1-th layer oulput, i.e., b4 (n+1) = Z [2d 2, 1}. Then,
holds that h. 11++1) = -(0n. wi) where {wi} is defined as w/ = 0,
W1+1 = wy + 1/4 C Ej= (Ry - Fg + w Ø3 - w 1-
for = ....., 1, where Fy = Ei= Rs-
The iroof is in A.6 with numerical verification Appendix as sanity check.
7 CONCLUSION
This work makes the first step towards white-boxing the mechanism of ICRL under reinforcement pretraining, focusing specifically on policy evaluation. We provide constructive proof that trans- formers can implement multiple temporal difference algorithms in the forward pass for in-context policy evaluation. Additionally, we theoretically and empirically show that the parameters enabling in-context policy evaluation emerge naturally through multi-task TD pretraining. We find com- pelling that randomly initialized transformer, only trained for simple policy evaluation tasks, can learn to discover and implement TD, a provably capable RL algorithm for policy evaluation.
Admittedly, this work does have a few limitations. First, we focus solely on policy evaluation. Second, to facilitate theoretical analysis, we make a few assumptions (e.g., Assumptions 5.1 & 5.2) and simplifications (e.g., using linear attention instead of softmax attention). Yet those assumptions and simplifications may not hold in many practical scenarios. Third, our pretraining method (Algorithm I) requires the ability to randomly generate policy evaluation tasks, which may not be available in many cases, such as offline training. Fourth, despite the fact that we evaluate Algorithm in both Boyan's 's chain and CartPole, it is not evaluated on large-scale environments such as the Atari games (Bellemare et al., 2013) or DeepMindLab (Beattic et al., 2016). We believe that addressing those limitations would be fruitful directions for future works.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work is supported in part by the US National Science Foundation (NSF) under grants III-2128019 and SLES-2333904 EB acknowledges support from the NSF Graduate Research Fellowship (NSF- GRFP) under award 1842490. HD acknowledges support from the NSF TRIPODS program under award DMS-2022448.
REFERENCES
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. (2024). Transformers learn implement precondi- tioned gradient descent for in-context learning. In Advances in Neural Information Processing Systems.
Ahn, K., Cheng, X., Song, M., Yun, C., Jadbabaie, A., and Sra, S. (2023). Linear attention is (maybe) all you need (to understand transformer optimization). ArXiv Preprint.
Akyürek, E.., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. (2023). What learning algorithm is in-context learning? investigations with linear models. In Proceedings of the International Conference on Learning Representations.
Allen-Zhu, Z. and Li, Y. (2023). Physics of language models: Part 1, context-free grammar. ArXiv Preprint.
Alvarez Melis, D. and Jaakkola, T. (2018). Towards robust interpretability with self-explaining neural networks. In Advances in Neural Information Processing Systems.
Ansel, J., Yang, E.,, He, H., Gimelshein, N., Jain, A., Voznesensky, M., Bao, B., Bell, P., Berard, D., Burovski, E., Chauhan, G., Chourdia, A., Constable, W., Desmaison, A., DeVito, Z., Ellison, E., Feng. W., Gong, J., Gschwind, M., Hirsh, B., Huang, S., Kalambarkar, K., Kirsch, L., Lazos, M., Lezcano, M.. Liang. Y., Liang, J., Lu, Y.. Luk, C.. Maher, B., Pan, Y., Puhrsch, C., Reso, M.. Saroufim, M., Siraichi. M. Y., Suk, H.,, Suo, M.. Tillet, P., Wang. E., Wang, X., Wen. W.,
Zhang, S., Zhao, X., Zhou, K., Zou, R., Mathews, A., Chanan, G., Wu, P., and Chintala, S. (2024). PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In Proceedings of the ACM International Conference on Architectural Support for Programming Languages and Operating Systems.
Baird, L. C. (1995). Residual algorithms: Reinforcement learning with function approximation. In Proceedings of the International Conference on Machine Learning.
Bauer, J., Baumli, K., Behbahani, F., Bhoopchand, A., Bradley- N., Chang, M., Clay, N., Collister, A., Dasagi, V.. Gonzalez, L., et al. (2023). numman-timesaale adaptation in an open-ended task space. In Proceedings of the International Conference on Machine Learning.
Beattie, C., Leibo, J. Z., Teplyashin, D., Ward, T.,, Wainwright, M., Küttler, H., Lefrancq, A., Green, S., Valdés, V., Sadik, A., et al. (2016). Deepmind lab. ArXiv Preprint.
Beck, J., Vuorio, R., Liu, E. Z., Xiong, Z., Zintgraf, L., Finn, C., and Whiteson, S. (2023). A survey of meta-reiinforcement learning, ArXiv Preprint.
Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. (2013). The arcade learning environment:
An evaluation platform for general agents. Journal of Artificial Intelligence Research.
Bengio, Y., Goodfellow, I., and Courville, A. (2017). Deep learning. MIT press Cambridge, MA, USA.
Bhandari, J., Russo, D., and Singal, R. (2018). A finite time analysis of temporal difference learning with linear function approximation. In Proceedings of the Conference on Learning Theory.
Borkar, V. S. and Meyn, S. P. (2000). The ode method for convergence of stochastic approximation and reinforcement learning. SIAM Journal on Control and Optimization.
Boyan, J. A. (1999). Lquates temporal difference learning. In Proceedings of the International Conference Machine Learning.
11

---

Published as a conference paper at ICLR 2025
Brockman, G., Cheung, V., Pettersson, L.,, Schneider, J., Schulman, J., Tang, J., and Zaremba, W. (2016). OpenAl Gym. ArXiv Preprint.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam,
P., Sastry, G., Askell, A., Agarwal, S., Beerbert-Voss A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. (2020). Language models are few-shot learners. In Advances in Neural Information Processing Systems.
Cai, Q., Yang, Z., Lee, D., and Wang. Z. (2019). Neural temporal-difference and q-learning provably converge to global optima. ArXiv Preprint.
Chan, R., Yu, Y., You, C.,, Qi. H., Wright, J., and Ma, (2022). Redunet: A white-box deep network from the principle of maximizing rate reduction. Journal of Muchine Learning Research.
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L.,, et al. (2020). Rethinking attention with performers. ArXiv Preprint.
Cook, J., Lu, C., Hughes, E., Leibo, Z.., and Foerster, J. (2024). Artificial generational intelligence: Cultural accumulation in reinforcement learning. ArXiv Preprint.
Cybenko, G. (1989). Approximation by superpositions of a sigmoidal function. Mathematics of Control, Signals and Systems.
Dai, Z., Tomasi, F., and Ghiassian, S. (2024). In-context exploration-exxloitatonn for reinforcement learning. ArXiv Preprint.
Duan, Y., Schulman, J., Chen, X., Bartlett, P. L., Sutskever, I,, and Abbeel, P. (2016). RI2: nforcement learning via slow reinforcement learning, ArXiv Preprint.
Elman, J. L. (1990). Finding structure in time. Cognitive Science.
Frosst, N. and Hinton, G. (2017). Distilling a neural network into soft decision tree. ArXiv Preprint.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. (2022). What can transformers learn in-context? a case study of simple function classes. In Advances in Neural Information Processing Systems.
Garrett, J. D. (2021). garreti403/SciencePlos..
Gatmiry, K., Saunshi, N., Reddi, S. J., Jegelka. S., and Kumar, S. (2024). Can looped transformers learn to implement multi-step gradient descent for in-context learning? In Proceedings of the International Conference Machine Learning.
Graves, A., Wayne, G., and Danihelka, I. (2014). Neural turing machines. ArXiv Preprint.
Grigsby, J., Fan, L., and Zhu, Y. (2023). Amago: Scalable in-context reinforcement learning for adaptive agents. ArXiv Preprint.
Grigsby, J., Sasek, J., Parajuli, S., Adebi, I. D., Zhang, A., and Zhu, Y. (2024). Amago-2: Breaking the multi-task barrier in meta-reinforcemmnt learning with transformers. In Advances in Neural Information Processing Systems.
Harris, C. R., Millman, K. J., der Walt, S. Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Río, F., Wiebe, M., Peterson, P., Gérard-Marchant, P., Sheppard, K., Reddy, T., Weckesser, W., Abbasi. H., Gohlke, C., and Oliphant, T. E. (2020), Array programming with NumPy. Nature.
Hochreiter, S., Younger, A. S., and Conwell, P. R. (2001). Learning to learn using gradient descent. In Proceedings of the International Conference on Artificial Neural Networks.
Hornik, K., Stinchcombe, M., and White, H. (1989). Multilayer feedforward networks are universal approximators. Neural Networks.
12

---

Published as a conference paper at ICLR 2025
Huang, S., Hu, J., Chen, H., Sun, L., and Yang, B. (2024a). In-context decision transformer: Reinforcement learning via hierarchical chain-of tf-ouuggh.. ArXiv Preprint.
Huang. S., Hu, J., Yang, Z.,, Yang, L., Luo, T., Chen, H., Sun, L.., and Yang, B. (2024b). Decision mamba: Reinforcement learning via hybrid selective sequence modeling. ArXiv Preprint.
Hunter, J. D. (2007). Matplotlib: 2d graphics environment. Computing in Science & Engineering.
Jastrzebski, S., Arpit, D., Ballas, N., Verma, V., Che, T., and Bengio, Y. (2017). Residual connections encourage iterative inference. ArXiv Preprint.
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. (2020). Transformers are rnns: Fast autoregressive transformers with linear attention, In Proceedings of the International Conference on Machine Learning.
Kingma, D. P. and Ba, J. (2015). Adam: A method for stochastic optimization. In Proceedings of the International Conference on Learning Representations.
Kirsch, L., Flennerhag, S., Hasselt, H. v., Friesen, A., Oh, J., and Chen, Y. (2022). Introducing symmetries to black box meta reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence.
Kirsch, L., Harrison, J., Freeman, C., Sohl-Dickstein, J., and Schmidhuber, J. (2023). Towards general-purpose in-context learning agents. In NeurlPS Foundation Models for Decision Making Workshop.
Kirsch, L., van Steenkiste, S., and Schmidhuber, J. (2019). Improving generalization in meta reinforcement learning using learned objectives. ArXiv Preprint.
Krishnamurthy, A., Harris, K., Foster, D. J., Zhang, C., and Slivkins, A. (2024). Can large language models explore in-context? ArXiv Preprint.
Laskin, M., Wang, L., Oh, J., Parisotto, E.,, Spencer, S., Steigerwald, R., Strouse, D., Hansen, S., Filos, A., Brooks, E., et al. (2022). In-context reinforcement learning with algorithm distillation. ArXiv Preprint.
Lee, J., Xie, A., Pacchiano, A., Chandak, Y., Finn, C., Nachum, O., and Brunskill, E. (2024). Super- vised pretraining can learn in-context reinforcement learning. In Advances in Neural Information Processing Systems.
Leshno, M., Lin, V. Y., Pinkus, A., and Schocken, S. (1993). Multilayer feedforward networks with a nonpolynomial activation function can approximate any function. Neural Networks.
Lin, L., Bai, Y., and Mei, S. (2023). Transformers as decision makers: Provable in-context reinforce- ment learning via supervised pretraining. ArXiv Preprint.
Liu, S., Chen, S., and Zhang, S. (2025). The ODE method for stochastic approximation and reinforcement learning with markovian noise. Journal of Machine Learning Research.
Lu, C., Kuba, J., Letcher, A., Metz, L., Schroeder de Witt, C., and Foerster, J. (2022). Discovered policy optimisation. In Advances in Neural Information Processing Systems.
Lu, C., Schroecker, Y., Gu, A., Parisotto, E., Foerster, J., Singh, S., and Behbahani, F. (2023). Structured state space models for in-context reinforcement learning. In Advances in Neural Information Processing Systems.
Lu, Z., Pu, H., Wang, F., Hu, Z., and Wang, L. (2017). The expressive power of neural networks: A view from the width. In Advances Neural Information Processing Systems.
Mahankali, A., Hashimoto, T. B., and Ma, T. (2023). One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. ArXiv Preprint.
Mnih, V., Kavukcuoglu, K., Silver, D.,. Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M. A., Fidjeland, A., Ostrovski, G., Petersen. S., Beattie, C., Sadik, A., Antonoglou, I., King, H., Kumaran, D., Wierstra, D., Legg, S., and Hassabis, D. (2015). Human-leve control through deep reinforcement learning. Nature.
13

---

Published as a conference paper at ICLR 2025
Moeini, A., Wang, J., Beck, J., Blaser, E.., Whiteson, S., Chandra, R., and Zhang, S. (2025). A survey of in-context reinforcement learning. ArXiv Preprint.
Müller, S., Hollmann, N., Arango, S. P., Grabocka, J., and Hutter, F. (2022). Transformers can do bayesian inference. In Proceedings of the International Conference on Learning Representations.
Oh, J., Hessel, M., Czarnecki, W. M., Xu, Z., van Hasselt, H. P., Singh, S., and Silver, D. (2020).
Discovering reinforcement learning algorithms. In Advances in Neural Information Processing Systems.
Park, C., Liu, X., Ozdaglar, A., and Zhang, K. (2024). Do Ilm agents have regret? a case study in online learning and games. ArXiv Preprint.
Puterman, M. L. (2014). Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons.
Qian, X., Xie, Z., Liu, X., and Zhang, S. (2024). Almost sure convergence rates and concentration of stochastic approximation and reinforcement learning with markovian noise. ArXiv Preprint.
Raparthy, S. C.,, Hambro, E., Kirk, R., Henaff, M., and Raileanu, R. (2023). Generalization sequential decision making tasks with in-context learning. ArXiv Preprint.
Sander, M. E., Giryes, R., Suzuki, T., Blondel, M., and Peyré, G. (2024). How do transformers perform in-context autoregressive learning? ArXiv Preprint.
Schlag, I. Irie, K., and Schmidhuber, J. (2021). Linear transformers are secretly fast weight programmers. In Proceedings of the International Conference on Machine Learning.
Shi, L. X., Jiang, Y., Grigsby, J., Fan, L., and Zhu, Y. (2024). Cross-episodic curriculum for transformer agents. In Advances in Neural Information Processing Systems.
Siegelmann, T. and Sontag, E. D. (1992). On the computational power of neural nets. In Proceedings of the Annual Workshop on Computational Learning Theory.
Sinii, V., Nikulin, A., Kurenkov, V., Zisman, I,, and Kolesnikov, S. (2023). In-context reinforcement learning for variable action spaces. ArXiv Preprint.
Sutton, R. S. (1988). Learning to predict by the methods of temporal differences. Machine Learning.
Sutton, R. S. and Barto, A. G. (2018). Reinforcement Learning: An Introduction (2nd Edition). MIT press.
Tsitsiklis, J. N. and Roy, B. V. (1999). Average cost temporal-differene learning. Automatica.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. (2017). Attention all you need. In Advances in Neural Information Processing
Systems.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. (2023a). Transformers learn in-context by gradient descent. In Proceedings of the International Conference on Machine Learning.
Von Oswald, J., Niklasson, E., Schlegel, M., Kobayashi, S., Zucchet, N., Scherrer, N., Miller, N., Sandler, M., Vladymyrov, M., Pascanu, R., al. (2023b). Uncovering mesa-optimization algorithms in transformers. ArXiv Preprint.
Wang, J. X., Kurth-Nelson, Z., Tirumala. D.. Soyer, H.. Leibo, J. Z., Munos, R., Blundell, C.. Kumaran, D., and Botvinick, M. (2016). Learning to reinforcement learn. ArXiv Preprint.
Wang, S., Li. B, Z., Khabsa, M., Fang, H., and Ma, H. (2020), Linformer: Self-attention with linear complexity, ArXiv Preprint.
Wang, Z., Wang, H., and Qi. Y. (2024). Hierarchical prompt decision transformer: Improving few-shot policy generalization with global and adaptive. ArXiv Preprint.
14

---

Published as a conference paper at ICLR 2025
Wu, J., Zou, D., Chen, Z., Braverman, V., Gu, Q., and Bartlett, P. L. (2023). How many pretraining tasks are needed for in-context learning of linear regression? ArXiv Preprint.
Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum, J., and Gan, C. (2022). Prompting decision transformer for few-shot policy generalization. In Proceedings of the International Conference on Machine Learning.
Xu, T., Li, Z., and Ren, Q. (2024). Meta-reinforcement learning robust to distributional shift via performing lifelong in-context learning. Proceedings of the International Conference on Machine Learning.
Yao, H. and Liu, Z.-Q. (2008). Preconditioned temporal difference learning. In Proceedings the International Conference on Machine Learning.
Yu, Y., Buchanan, S., Pai, D., Chu, T., Wu, Z., Tong, S., Haeffele, B., and Ma. Y. (2023). White-box transformers via sparse rate reduction. Advances in Neural Information Processing Systems.
Zhang, R., Frei, S., and Bartlett, L.. (2024). Trained transformers learn linear models in-context. Journal of Machine Learning Research.
Zhang, S., Boehmer, W., and Whiteson, S. (2020). Deep residual reinforcement learning. In Proceedings of the International Conference on Autonomous Agents and Multiagent Systems.
Zhao, H., Panigrahi, A., Ge, R., and Arora, S. (2023). Do transformers parse while predicting the masked word? ArXiv Preprint.
Zheng, C., Huang, W., Wang, R., Wu, G., Zhu, J., and Li, C. (2024). On mesa-optimization in autoregressively trained transformers: Emergence and capability. ArXiv Preprint.
Zisman, I., Kurenkov, V., Nikulin, A., Sinii, V., and Kolesnikov, S. (2023). Emergence of in-context reinforcement learning from noise distillation. ArXiv Preprint.
15

---

Published as a conference paper at ICLR 2025
TABLE OF CONTENTS
1 Introduction
2 Related Works
3 Background
Transformers Can Implement In-Context TD(O)
5
Transformers Do Implement In-Context TD(0)
6 Transformers Can Implement More RL Algorithms
7
Conclusion
10
Proofs
18
A.1 Proof of Theorem
18
A.2 Proof of Corollary
22
A.3 Proof of Theorem 2
24
A.4 Proof of Corollary 2
29
A.5 Proof of Corollary 3
31
A.6 Proof of Theorem 3
36
Experimental Details of Figure
41
Boyan's Chain Evaluation Task Generation
Additional Experiments with Linear Transformers
42
D.I
Experiment Setup
42
D.1.1 Trained Transformer Flement-wise Convergence Metrics
43
D.1.2 Trained Transformer and Batch TD Comparison Metrics
43
D.2 Autoregressive Linear Transformers with L= 1, 2, Layers
44
D.3 Sequential Transformers with =, 4 Layers
46
Nonlinear Attention
47
Experiments with CartPole Environment
48
F.1 CartPole Evaluation Task Generation
48
F.2 Experimental Results Pre-training with CartPole
49
G Investigation of In-Context TD with RNN
50
G.I Theoretical Analysis of Linear RNN
50
G.2 Multi-task TD with Deep RNN
51
16

---

Published as a conference paper at ICLR 2025
H Numerical Verification of Proofs
52
17

---

Published as a conference paper at ICLR 2025
A PROOFS
A. PROOF OF THEOREM 1
Proof. We recall from (3) that the embedding evolves according
Z+1 - PZ QiZ.).
11
We first express Zi using elements of Zo. To this end, it is convenient to give elements of Zr different in the following way 1..
(zf Vi (1)
names, particular, we refer to the elements Ze as
(1)
zj (n)
Ti (a+1)
Zi
yi in-
y1
yi
where recall that Z R(24+1)x(x+1)
21
(1)
2d
y/i RR Sometimes it is more convenient
(6)
to refer to the first half and second half of 2/ separately, by, e.g, 1i E R 5l R i.c.,
(1)
Then we have
10)
())
y [n+1)
Zi =
(n)
(n++1)
(++1)
Vi
Vi
We utilize the shorthands
X1
1
())
x
[n]
E R
2dxn
(1)
Vi
(n)
E
R1xn
Then we have
X1
(n+1)
Zi =
Yi
Vi
(n+1)
For the input Zo, we assume 60 (n-11) = 0. yo
(x+1)
0 but all other entries of Zu are arbitrary. We
recall our definition of M in (2) and {PiTD Q
I- in (7). In particular, we can express
Q! in a more compact way as
M
-
-
Id
Ia
E 2dx2d
Oaxd
Oaxd
Bi
CT
Oaxd
E R 2ddx2d
Odxd
0axd
-CT
CT
A
-
B,Mi
E 2dx22
0axd
0axd
Q
A
Ozax1
E ((22+1)x(2411)
01x2d
We now proceed with the following claims.
Claim 1. X1 = Xo, (n+1)
(n+11) VI.
= =x0
O2ax2d 01x2d
0zax1]
Recall that PTD
(2(24+))x(2011) Let
1
Wi = ZiM (Z Q. (P2
Zi) E (22++)))
18

---

Published as a conference paper at ICLR 2025
The embedding evolution can then be expressed
1 WM
Z41 = 2i + n
By simple matrix arithmetic, we get
02ax(n+1) Wi(2d +1]
PP
where W2d + 1) denotes the (2d + 1)-th row of W,. Therefore, we have X(+1 X1, [4]
(x+1)
(a+1) By induction, we get X₁ Xu and z (n+1) Fo
(n+1) for all I = [O....,L 11.
In light of this, we drop all the subscripts of X,, as well as subscripts of z14) for i 1, 1,.... n
Claim 2.
Y14 Yi - YiX A,X (n+1) 9i41 vi (n+1)
Yix (xx(1+1)
The easier way to show why this claim holds is to factor the embedding evolution into the product of pid Z.M and Z Q1 74. Firstly, we have
02dxx Yi
02ax (n+1)
PID
y1
Applying the mask, we get
O2dxo Yi
02dx1 0
PIDZ,
Then, analyze Zí
Q1 Z₁. Applying the block matrix notations, get
XT
Y.
A
02ax1
X
(n+ (a+1)
2 Q1 =
I(x+1)T
Vi
(n+1)
01x2d
Yi
Vi
XTA
Onxi
X xn(n117
(1++
A
Vi
Vi (n+1)
X AAX AxX
X (yz((++1)
(n++1) AiX
we (n+ 1)" (x1()
Combining the two, get
P. 2))
Ozdxn
02ax1 0
A
X A(z((++1 ((+1)) Acx((++))
Yi
x(++1) A&X
Ozdxn
Ozax1
YXX
A
YiX (cx(
Hence, according to our update rule in (3), we get
V4+1 Yi TiXT A&X
(n+1) / + - - (1+1)
(n+1)
1
Vi41
Claim 3.
1 = 30 + Mi (1)
(6) ())
BJ M2XY,T
1=
19

---

Published as a conference paper at ICLR 2025
Ia
Oaxd Oaxd]
for I,...,,7 n + 1, where M2
Oaxd
Following Claim 2, we can unroll Yi+1
Y41 = + - -Yx A,X
n
Yi Yi-1 + - A
n
Y1 Yo + -Y -oX AnX.
n
We can then compactly express Vi+1 as
Y41 = Yo +
Σ
YyXTAN A,X.
==0
Recall that we define A, ByM1. Then, rewrite Y1+1
Y41 = Yo + n ==
Y, 3TXTMBB,M
The introduction of M2 here does not break the equivalence because B, MyB3. However, it will help make our proof steps casier to comprehend later.
With the identical procedure, we can easily rewrite 3/41 as
(n+1)
(n+1) (n+1) M1+
Σ
10
j=0
In light of this, we define vho = and for O....
J BJ M2XY, e 1224
Σ
(19)
j=0
Then we can write
(1
(8)
+ (MMrd(9.</4 2/11))
9i41
yo
(20)
(n+1)
for 1, ,,...,7e n + 1, which is the claim we made. In particular, since we assume yo have
0, we
(x+1) Mi+1 -M1aa(at)) v111).
Claim 4. The bottom elements of th are always O, i.c., there exists a sequence {w) e Rd) such that we can express 1/7
4/
wt
(21)
0ax1
for all I..... L.
We prove the claim by induction. The base case holds trivially since th 0. Suppose that for some 1, (21) holds. It can be easily verified from the definition of v/++1 in (19) that
+1+1 t + -B, M2XY,T.
(22)
20

---

Published as a conference paper at ICLR 2025
If we let
1 N₂ Ti==MMXY M2XY, E R22xx1
the evolution of 1/2+1 can then be compactly expressed as,
V1+1 = 1 + Bi No.
By matrix arithmetic, we have
C
Odxd Odxd
: Ni(d:
B₁ Nr
Oaxd
C1Ni(1 : dd]
0ax1
where N1 d) E R' and Ni(d 2d) E R represent the first d and second elements of N₁ respectively. Substituting in our inductive hypothesis into (22), we have:
wit C d)
W1+1
0dx1]
Odx1
Wi + C d)
0ax1
we let Wi+1 = Wi + C.N((1 : d), we can see that the property holds for s/141, thereby verifying Claim 4.
Given all the claims above, we can then compute that
1h+ Mix(at1))
(1r,, 112, M2 (zx(1+1))
1 BJ M2XY, Maa(ntt)
(By (22))
n
(1r, M1z((++))
1
Bi M2 (0 M
71
-
1
Σ
BJ M2
((0r, + :0"). M M
(By (20))
1=1
(1r,. M2x(n+t))
+
1
Σ
BJ
(())
[0001-1000 0ax1
((10))
-v(0 + ((()
1) + 160 M Maa((tt)
0axi
12
1=1
-tbr M1 M2z(n+1))
1
n
Cru(i)
0zx1
30
(6)
+ w/ e(i) un vl1) M)a((t1))
(By Claim 4)
72
-(ubr (1x(m+1))
1
Civit)
yo
+ wi g(i) wi vll)
12
0ax1
This means
(wi+), ((tt))(att)) - -wn, sat1))
(Cull) (100 wj wT((()) wi v(i). elttti)
Since the choice of the query ((++1) is arbitrary, we get
Σ a (16) w/ T g(()) - -w7 u1)) 2 21))
Wi+1 = 10 +
In particular, when we construct Zo such that (+) = -b:11 ((1) = 70: and yn (1)
Rs, we get
Wi+1 = wi + Σ (R: + Two Φz w, Φb 00-1
21

---

Published as a conference paper at ICLR 2025
which is the update rule for pre-conditioned TD learning. We also have
(m+1)
( 1/. (1,. M,alnttt) - wi, p((t1))
This concludes proof.
A.2 PROOF OF COROLLARY
Proof. The proof presented here closely mirrors the methodology and notation established in Theorem Since we are only considering I-layer transformer in this Corollary, we can recall the embedding evolution from (3) and write
1
21 = Zo + -
-PozoM QoZo).
We once again refer to the elements in Z1 as {(o{)) ,111
())
in the following way
#41
[n)
(n+1)
Z =
(1)
nn
i--1)
y1
V1
Vi
where recall that Z₁ R
(2d++))
21
(i)
R
y1
(
E
R. We utilize, vf Rd, 5/ R
(i)
refer to the first half and second half of ie., 27
(1)
(
Then have
Vi [n+1)
Vi
7,
(n)
(1+1)
(n+1)
y1
Vi
We further define as shorthands
Xi
(1)
2
(n)
E
R2dxn
Y Y= =
yi
(1)
vi (n)
E RRix.
Then the block-wise structure of Zi can be succinetly expressed as:
(n+1)
X₁
Z1
y1 in+1)
Y
For the input Zo, we assume 60
(n--1) = 0, yo
(n+1)
= but all other entries of Zu are arbitrary. We
recall our definition of M in (2) and {Po, Qo) in (7). In particular, we can express Qu in a more compact way
M1
Ta
Odxd Oaxd
E 2dx2d
Bo =
CT 0axd
Oaxd Oaxd
R 2ddxxd
Oaxd
Ao
Mo =
-C
Oaxd
Rd2x2
R
0dxd
Odxd
Qo
Ao
02dx1
(224+1))(24+1))
01x2d
0
We will proceed with the following claims.
Claim 1. X1 Xo, xx+,( (x+1)
z (n+1)
Because we are considering the special case of L = 1 and because we utilize the same definition of Po as in Theorem 1, the argument proving Claim 1 in Theorem I holds here as well. As result, drop all the subscripts of X1, as well as subscripts of (6) for i = 1,..., n 1.
Claim 2.
Yi = Yo
YoX AgX
n
22

---

Published as a conference paper at ICLR 2025
y1 [++1)
21-0"" - -Yo (gz(1+1),
(n-1)
This claim is special case Claim 2 from the proof Theorem Appendix A.I, where L Our diiee construction of Qu matches that in the proof of Theorem 1. Although our Ag here differs from the specific form of Au the proof of Theorem 1, this specific form is not utilized the proof of Claim 2. Therefore, the proof of Claim 2 in Appendix A.1 applies here, and omit the steps to avoid redundancy.
Claim 3.
yo + M1 - Bd MaXYo
V1
la
Oaxd Odxd]
for = I,..., + 1, where M2 =
Odxd
This claim once again is the L 1 case Claim 3 from the proof of Theorem Appendix A.1. The specific form of M1 is not utilized in the proof of Claim from Appendix 1. so it applies here.
We can then define tho = 0 and.
1 By M2XY, R2d
4/ = n
(23)
Then we can write
vi = uo + (Mia ,b)).
(n+1)
for ,,..., 1, which is the claim we made. In particular, since we assume y0 have
y1 (n+1) = -M ..))
Claim 4. The bottom d elements of th are always O, i.c., there exists wi e Rd such that we can express t/1
Wi
11
0axx
Since our Bn here is identical to that in the proof of Theorem 1 in A.1, Claim 4 holds for the same reason. We therefore omit the proof details to avoid repetition.
Given all the claims above, we can then compute that
( (/. M Mix((+1)
BJ M2XY, Mia(n-1)
(By (23))
Bi M2x ( yo M2
(+t) (1(()) M1x((+1)
BJ
0dx1
()
(v6"). M11o((+1))
(By Claim 4)
0ax1
(CCor(0sh (2)
0dx1
i=1
This means
(1, (((1))
Σ
(Coul y0 wl0+1)
1=1
23

---

Published as a conference paper at ICLR 2025
Since the choice of the query p(++)) is arbitrary, get
W1 Σ Covo Govo2ul (())
In particular, when we construct Zo such that v10) Φ and yo = R4 we get
Colidii1
which is the update rule for a single step of TD(O) with up = O. We also have
Vi (x+1)
( 1/1, M1 (w1,, ((ott))))
This concludes our proof.
A.3 PROOF OF THEOREM 2
Preliminaries Before we present the proof, we first introduce notations convenient for our analysis.
We decompose Po and Qo
Qa E Rddd
Q E Rdxd
Qe E Rd 1 L E Rdx1 de E R
P E 12dx(ad+11 Po = e R1x(2d+1)
Qu = Q' E Rdxd
Q E Rdxd
4a Rixd
To Blxd
One readily check that TF; is independent of Qb V, 9b. le, Les %- these matrices zero. Let 26) be the i-th column of Zo. Indeed, TF; can be written
Thus,
assume that
TF1 (Zo, (Po. Qo1) = Zn 22d + 1,n +1]
(By (4))
1
Σ 20)(())
Qo2
n1
1
(p. (1))
),,(())(00)) 2(1)
Q02(++1)
11
(p))(0)) (0T= + +70 120 + Roo (24)
7
(prid), Di-1) + (P 0i) + p12421 Ri
Po.
oI + T Voonti + Redh+110
Bi(7a.Qo)
We prepare the following gradient computations for future use:
TZo (Po, 2o1) = -
3. Q0)ē.-1
Piid
i=1
n
P[441/20 [2, (Po, 2o1) = 2
Σ
3. Qo)o.
71 1=1
1
Ve. (F, (Po, Q6))
as (Z,, Po) 1 0
(25)
n
Vo: (7,, (Po, (Po, Qo)) = 1
0s os, Po) Po)
i=1
1
Vg 4. (Z,, (Po, Qo1) =
Σ
Rça5 (Zo, Po)
n i=1
24

---

Published as a conference paper at ICLR 2025
We will also reference the following two lemmas in main proof. Lemma A.3.1. Let A be a diagonal matrix whose diagonal elements are i.i.d Rademacher random variables "I.... Sd. For any matrix K IRdda we have that EA [AKKAA diag(K) K))
Proof. First, can write AKA explicitly
[
K11
h1
kid
[[
0
0
... 0
12
R22
K2a
0
C
0
AKA
Sa
kaz
kad
Ca
Using (AKA),, ly to denote the element in the i-th row at column j of AKA, from elementary matrix multiplication we have
(AKA);; CihtefSi-
When + j. R[ck)] = [c, becasue Si and Ç are independent. For i= j. R[(Gs)
E|c?] 1. We can then compute the expectation
kis i=
[X[(KA]];
0
i j.
Consequently,
BA[AKA] = diag((K)
Lemma A.3.2. Let II Raxd be random permutation matrix uniformly distributed over all d x d permutation matrices and E Rdxd be a diagonal matrix. Then, it holds that
En 11II
tr tr(L)
d
Proof. By definition,
d
lg Σ
Ti
We note that each of TT is standard basis. Given the orthogonality of standard bases, get
0
i i i=
ILIII lig Lg.gr
where 9; is the unique index such that 1gg. = If the distribution of II is uniform, then [IIIII Jss Las equal to one of Li....,, Las with the same probability. Thus, the expected value [IIII] Ja is tr(.).
Now, we start with the proof of the theorem statement.
Proof. We recall the definition of set O* as
cIa
Oaxd
Odx1 0ax1
O2ax2d
0zax1
- Une p=
Q
e' Ta
Oaxa
O1x2d
[1xd
Ozxd
0
Suppose 0g E O³, then by (24) and (25), we get
TF 0z)
7
Rs (ctt, 1n 6470, on-1)
(26)
1=1
"A Rademacher random variahle takes values
11, each with equal prohahility of 0.5.
25

---

Published as a conference paper at ICLR 2025
TFi (Z0,
Ret1 (cre Pu+2 + 4e (
n
(2,,
(c.0) Aut1 + C xn+1)
==1
TFi 04) Σ
(c+0)-1 10 + 0n+1)00
P1111:20)
n
7 n
Qq. TF₁ (Zo, 0%)
R 10
V Q. TF₁ (Zu, 0x)
7Tk
R2diOn+1
i=1
V 4. TF (Zo, = 7k
Σ
R
Øn+1
Recall the definition of (0) in (10). With slight abuse of notation, we define A(p1::)) to be the P[::d] component of (0)) i.e.
(p = (R (TF TF, (Zo 8TF (2ri(ZZ.1) 0)
Oppt:dl
Same goes for
and A(qa).
We will prove that
(a) A(pp:4) = A(p)t122)) ((qa) = for A(0));
(b) A(Qa) öIa and A(Q') = 8' d' for some 5, 8' for A(++)
using Assumptions 5.1 and 5.2. We can see that the combination of (a) and (b) are sufficient for proving the theorem. Recall that Zo and Zij are sampled from (po. P. 6). We make the following claims to assist our proof of (a) and (b).
Claim 1. Let be Rademacher random variable. We denote Z4 and 7 as the prompts sampled from (Po. p, r, (Ø). We then have Zu - Zç and Z & ZZ.. To show this true, we notice that for any realization of Ç. denoted as { E {1, -1}, we have
Pr(po. p.r, 6) = Pr(po r) Pr(p) = =Pr(po, Pr(CTat)
(Assumption 5.1)
(Assumption 5.2)
PPr(po.p,,.Co).
(Assumption 5.1)
It then follows that
Pr(po. P,r, 3) = Pr(poo,
Σ
Pr(s= = 5)
1E(1-11)
Σ
Pr(po.p.r, 4) Pr(s = =()
(E(1-11}
Σ
Pr(p..P.r, (0) Pr(s 2)
(E(1-11}
Prr(po.p...Co)
This implies Claim holds.
Claim 2. Define A as the diagonal matrix whose diagonal elements are i.i.d. Rademacher random variables (1,... Siy..ra Ca- We denote ZA and Z', as the prompts sampled from (Po.P,r, AΦ), where Ag
26

---

Published as a conference paper at ICLR 2025
means [Ao())]ES. We then have Zu _ Za and Zg ZzA The proof follows the same procedures as Claim
Let II be random permutation matrix uniformly distributed over all d x d permutation matrices. We denote Zn and Zn1 as the prompts sampled from (Po: p. 7; Ilo), where 1l means [[ We then have Zo - ZIT and Zo - Zi The proof follows the same procedures as Claim
Claim 3.
Proof of (a) using Claim 1 It is easy to check by (26) that
TF, (Zc, 0k) = ?/
Ra 2 + 267 ( Φn-1)
i=1
2
TF1 (Zo, 0k)
=
= =TF 0x).
(27)
Similarly, one can check that TF, (Z 0x) Furthermore,
TF1 (Zy, 0z).
TFi (2c,
2 Co- 1 +.47 2 0n+1 cdi-1
Ck
i=1
=1
5
n
+ coro7 @nt1)di-1
i=1
VV P[L.a]
TF (Zo: 0n).
(28)
Then, from (10), we get
Appp1a))
=E[ + TF, 0x) (Zo,
TF, (Zo, 0z)]
PI:d
=E [ +
+TFF 02) TF, (Zç. (+)) P(.J TF (Zç. 0x)]
(By Claim 1)
=B,g [E[((E [RRtt22 + 7TF; (Z,, 06) TF, (2c, 0x)) PPIN TF, (Zc, I [[]
-E [E[ + TTF; (Z, 02) TF, (Zo, V P!L.a] TF, (Zo, 02) | []]
(By (27), (28))
_Eg/<IB (But2 7TTF, (Z), 0z) TF, (Zo, 0x)) P(n. TF, (Zu, 0x) []]
-Bş [ (Rn+2 1 (Z) 0%)
TF, 1Z.. 0x)) P[L.0] TF (Zu, 0x]]]
EE [ (Ru+2 + TTFF (Z), 02) TF, (Zo, 0z) 0z)) P[a:a] TF (o, 0s)]
=0.
The proof is analogous for A(p(4+1224) = 0, and (ga) = 0.
Proof of (b) using Claims and 3 We first show that A(Q.) is a diagonal matrix. Similar to (a),
have
TF, = - Σ
Ma Coo.111 12 ont1 + ToOT A @nt1
(29)
1=1
-TF; (7o- 0z).
Similarly, get TF; (Z), TF, (Z), 0x). Additionally, we have
- Σ
Qq. TF (ZA, 0x)
T A AVQ. TF 0x)
(30)
By (10) again, get
A(Qa)
27

---

Published as a conference paper at ICLR 2025
[En (00F442+TTFF/FFC 0x) QQ. TFi (Zu 8x]]
= TTF (2), - -TF, (Q. TFi
(By Claim 2)
-E, R
(%)) -
A]]
-EA [B[ R2 + 7/TF/ (Z) 0x)) - TFI Q. QF; TFi (Zu. 0x)A AI]
(By (29), (30))
[AE ( R2+2 + 7/TF (Z), 0x) TF, (Zo, Q. 0z) A]A]
-IEA [AE]( R2+ + 7TF, (Z), 02) TF, Qz) TF (Zo. 0z]]A]
-diag(EE) (EE R2 + 7TF, (Z), 0x) - TF1 (Zo, 8x)) TF1 (Zo, 0x)l)
(By Lemma A.3.1)
(Qa)).
The last equation holds if and only if AQQ.) is diagonal. We have proven this claim.
Now, we prove that (QQ.) = a for some ā E using Claim and Lemma A.3.2. Let IT be random permutation matrix uniformly distributed over all permutation matrices. Recall the definition of Zn and Z1 in Claim 3. We have
n 1
TF: (Zn, 0k)
Σ R, coo. II I
+ ro II I @n+1
TF, (Zo, (31.
Analogously, we get TF1 (2 0x) = TF, (Zi, 0x). Furthermore, we have
Q. TF (Zn, 0k)
m 1 T = Q. (Z, 0x)
(32)
By (10), we ready to show that A(Qa)
=R((Rn+2 + TF, (7), 02) TF, (Zo. TF, (7o, 0k)]
=E[(Rn+2 + TTF (Z 0%) TF, (Zn, (x)) V Q TF1 0x]
(By Claim 3)
[En|E[(Rn++ + TTF; (Zi, 0x) TF, (Zn, VQ. TF, (Zn, II]
=En [B [ + 7/TF (Z), 02)
TF, (Zo. 0x)] Q TF, 01)11 [
(By (31), (32))
=En [IE + 7TF, (Z), 02)
TF, (Zo, VQ) (Z,, (Z),0s) [1j11']
=En
IEE (Rn+2 + 7TF, (Z), 0x)
TF, (Z. Vo. TF, (Z, [x)
=En
=d tr(A(Qu))
(By Lemma A.3.2)
-ô1a-
The proof is analogous for (QQ'.) = 8' La for some E R.
Suppose that A(p(2++1) p E I. now conclude that
88a
Oaxd
0dx1
02ax2d O1xed
02dx1
((kx) A(Po)
A(Qn)
8'La
Oaxd
0ax1
01xd
Oxxd
Therefore, according to (10), we get
Dx+1 -0p ax*(
={
O2axed
Ozax1 M + aup
]
[c + agdls
Oaxd
0ax1
01x2d
dk + a
Oaxa O1xd
0ax1
E O..
Oxxd
28

---

Published as a conference paper at ICLR 2025
A.4 PROOF OF COROLLARY
Proof. We recall from (3) that the embedding evolves according to
2+11 7 1 PZ Q.7i).
(o/
(1)
(
We again refer to the elements in Z₁ as
in the following way
n+1 n+1
11+1)
Z
=
yi (1)
Vi nx)
1--1)
Vi
where we recall that Z E ((2++))((+1))
2y E
(
2d
31
E RR Sometimes, it more convenient
to refer to the first half and second half of 2 separately, by, e.g, 2s 10 Rd R R ie.,
(t)
(2)
(4) Then, we have
(n)
(n+1)
1
Zi =
nn)
+1
1+)
91
y1
yi
We utilize the shorthands
X1
1
(1)
[n)
R 2d
Yi
(1)
y1
E
R
Then we have
Xt
(n+1)]
Zi =
in+1)
Yi
Vi
For the input Zo. we assume c(n-1) recall our definition of M in (2) and { (PRG QRG) more compact way as
(n+1) but all other entries of Zu are arbitrary. We
yo
} in (12). In particular, we can express Q in
M1 =
la
Ia Odxd
E 2ddx2d
0dxd
M₂ _ M1
B. = CT
Odxa Odxd
E
TR. 2d2xx2d
Oaxd
-OT C
CT
A = =M, B,M
-CT
E 2dxxd
Q
A
02ax1
E ((2241)x(2241)
01x2d
O
We then verify the following claims.
Claim 1. X₁ Xo, xi
(n+1) (n+1)
=
Vl.
We note that PRo is the key reason Claim holds and is the same as the TD(0) case. Referring to A.1, we omit the proof of Claim here.
Claim 2.
Y441 Y + - YiXT A,X
(n+1)
= = (1+1) + YixT (n+1))
29

---

Published as a conference paper at ICLR 2025
Since the only difference between the true residual gradient and TD(0) configurations is the internal structure of At, we argue that it's irrelevant to Claim 2. We therefore again refer the readers to A.1 for detailed proof.
Claim 3.
(0) Mit1
(8) +
M1r(1 Σ j=0
BJ M2XY,
for = 1n...,n+1,
By Claim 2, we can unroll Y41 as
Y41 + TixTAR A,X
Y Yı-1 + n
- A
Y1 = Yo + YoX AgX.
We can then compactly express Y(+1 as
Vi+1
1o+ +
Y
A,X.
n j=
Recall that we define A₁ M2 ByM Then, rewrite Yi+1
r M2 B,MiX.
Y441 Yo +
With the identical procedure, we can easily rewrite (n+1)
Vi41
1 Mi+1 Mo + Σ YX M2 (1
(n+1) (n+1)
3=0
In light of this, we define vho = 0 and for O,...
1h+ BJ. M,XY, 2d
T
=1i + n
BJ MrXY,
(33)
Then write
9i41
(1
v0 + (M s
(34)
for i = 1, + 1, which is the claim made. In particular, since assume 30 (n+1) 0, have
(n+1) D141
= (Mia(atl) vit1)
Claim 4. The bottom elements of th are always O, i.e., there exists sequence {w) E [Rd] such that we can express t/ as
W1
0dx1
30

---

Published as a conference paper at ICLR 2025
for all 1 = 0.I....,L L.
Since B, is the key reason Claim holds and is identical to the TD(0) case, we refer the reader to A.1 for detailed proof.
Given all the claims above, we can then compute that
, Mia(at)))
Ve. Mzx(n+1)
+
Bj M2XY, MMz(nnt)
(By (33))
n
1 n
Ur, M(x((t1)
(6) Bi M2x 10 Vi M M
71
(WV, Mz M)z((+1))
1
Σ
B.] M2x! ((0b,, Morl!) + +v6") M M,z((tt))
(By (34))
11
i=1
(/r. Mzz(nt1)
1
Σ
B,
-g(0-
((10)) ((]]
-v(0) {(1)
}) + +b6²))
Myx M2x(utz)
Oax1
0ax1
i=1
([Ci(ul))
- (00(((060 + wi wTTe(s ((()
-(1, U1, (1x(1+1))
Σ
w/ vl1)) M (1+11
0ax1
n
i=1
(By Claim 4)
-(v Miz(utl))
+ Σ
1
Ci((()) - -g(0) 100" + wi c(1) - w/ v(il) MMa((tt)
71 i=1
0dx1
This means
(wit), ((+1)))) = -ww,
1
- -c() + Tw 1/ v()) mlot))
Since the choice of the query ((++1) is arbitrary, we get
1 Σ n
yo (2) (36)
ai
(TT(())
E(() wwT((()) ml)) (v00 ((1).
Wi+1 wi +
In particular, when we construct Zo such that (1() 1. g(+) 7h; and yn
Rrs we get
wi+1 = ti Σ ci rwt o, wi ( Yo;)
which is the update rule for pre-conditioned residual gradient learning. We also have
yi (a+1) ( (22, My -(( (wi, s(x+1)).
This concludes our proof.
A.5 PROOF OF COROLLARY
Proof. The proof presented here closely mirrors the methodology and notation established in the proof of Theorem from Appendix A. We begin by recalling the embedding evolution from (3) as,
1 Zi41 Zi + -PZM TDAA (ZT
31

---

Published as a conference paper at ICLR 2025
where we have substituted the original mask defined in (2) with the TD(A) mask in (14). We once again refer to the elements in Z as
(al) v1())
in the following way
[n)
(nt1)
Zi
y1 (1)
...
21
V1 nn)
yi 1-11
where recall that Z I
((2d+1))((+11)
(6)
IR
:3
(0)
E RR. We utilize, 1/ Rd, 5/ e CR
(8)
(1)
(i)
refer the first half and second half of
ie.,
11
Then we have
v (1)
y (n)
uy in+1)
(1+
Zi
5)
in+1)
Vi
We further define as shorthands,
X1
(1)
2j
[n)
E
R
2dxx
Y1 yi
())
(n)
E
R
xn
Then the blockwise structure of Z can be succinetly expressed as:
(n+1) (n+1)
X
Ti
Zi
Yi
V/
We proceed to the formal arguments by paralleling those in Theorem As in the theorem, we assume that certain initial conditions, such as $6 c(n+1) = and y/ (n-1) = 0, hold, but other entries of Zo are arbitrary. We recall our definition M TD(A) in (14) and {PTTD QT _TP) 1 in (7). In particular,
we can express QT in a more compact way as
M
- IId
Ia
E 2ddx2d
Oaxd
0dxd
B
CT
Oaxd
E R. Rd2x2d
Oaxd
0axd
--T
CT
Ar
-B.M
E 2d2x2d
Oaxd
0axd
Q
A
0zax1
E (2(24+1)((2+1))
01xed
We now proceed with the following claims.
In subsequent steps, sometimes useful to refer to the matrix M TD(A) in block form. Therefore. we will define H E (nx2dd) as the first rows of MTD(A) Z except for the last column, which define (A) E IR".
M TIDO
H
y:(1] %//
((++)))(2++)))
01x2d
Let (9) denote i-th column of H.
We proceed with the following claims.
Claim 1. X, xf =
(n+1)
(n+1) VI.
32

---

Published as a conference paper at ICLR 2025
Because utilize the same definition of PT in Theorem 1, the argument proving Claim 1 in Theorem holds here as well. As result, we drop all the subscripts of X₁, as well as subscripts of
(t) for = I,...,
Claim 2. Let H E IR (2dx),, where the i-th column of H is,
Σ
x (1)
E Rddd
A=1
Then write the updates for Y(+1, and (a+1)
Vi+1 as,
Y441 Y Y!H A,X,
(n+1)
(n+1) Y (a+(+
Vi41
y1
We will show this by factoring the embedding evolution into the product P.D Zi and M TD(A z and Qi*Zz, Z₁ Firstly, we have
T2axx
02dx1 (n+1)
PID Zi
Vi
Next we analyze TD(A) ZT.
From basic matrix algebra we have,
1
0
x(1)T
2000
x12)
((2)
2
X
0
x(3)T
(13)
M TD
A3
X²
A
0
An-1
X
An-3
n=-4
0
x(n)
((m)
0
0
0
0
0
x
(m+1)
0
(1)T
V
x(2)
+ (z)
y1
+
An
E An-in
Σ
A
01x2d
0
T())T
h(2)
(2)
(1)
Vi
+
Av
h(m)T
E
1
V/
(n)
01x2d
0
H
{2(A)]
K
O1x2d
where K(A) E Rd is introduced for notation simplicity.
Then, we analyze MIDAN) 7,7 QT Z,. Applying the block matrix notations, we get
(DT)) 27 ) =
H
[((A]
A
Ozax
X x(++1)
01x2d
0
01x2a
Yi
yi
(a+1)
H Ai
Onx1
[
((+1))
01x2d
Yi yf
(a+1)
H AAX
H 1cz(++1))
01x2d
0
33

---

Published as a conference paper at ICLR 2025
Combining the two, get
O2dxn
02dx1 (n+1)
H A1X 01x2d
11 Acæ(n+1)]
pid Zi (M 2iQ, Qi
Yi
Mi
0
Ozaxn YH A,X
02dx1 YiHT A
Hence, according to our update rule in (3), we get
Y44 Y1 1 YY A,X
(x+1)
(n+ + YYH (1a(1+1),
D141 Vi
Claim 3.
(1)
36 + Mix(1) Σ
B. M2XY,
Vi+1
im0
Ia
Oaxd Odxd]
for I,..., + 1, where M2
Oaxd
Following Claim 2, we can unroll the recursive definition of Y141 and express it compactly as,
Y1+ = Yo Σ
TAX.
A,X.
1=0
Recall that we define A; B,M, Then, we can rewrite Yi41 as
Y14 Yo Yo+ 1 n Σ T M.MM,
The introduction of M2 here does not break the equivalence because B: M2B,, However, it will help make our proof steps easier to comprehend later.
With the identical recursive unrolling procedure, can rewrite V/41 as
(n+1)
(n+1) = Mo + 1 n Wi
(n+1)
Σ Y,H (att))
In light of this, we define v/o = and for O....
1h
B. M,HY, e add
(35)
n
Then write
(0) (8) + (Mazl)), (br42)) 0i41
(36)
for i = 1, 1, which is the claim made. In particular, since assume y0 have
(n+1)
0,
(n+1) Di+1
= (M ve+1).
Claim 4. The bottom elements of th are always O, i.e., there exists sequence {w) E IRd] such that we can express t/ as
W1
0dx1
34

---

Published as a conference paper at ICLR 2025
for all ( 0,1,..,,
Because we utilize the same definition of Bt as in Theorem when defining 11+1, the argument proving Claim 4 in Theorem I holds here as well. We omit the steps to avoid redundancy.
Given all the claims above, we can then compute that
(W4,,, Max(ot))
dh, M2x(1+1)
1
Bi M .Ma((tt)
(By (35))
n2
dr, M1x
1
Bi Myh(o)y) ( M
- (0,,, Mz
Bi Mzh(( ((rr)))
+ =b6"). M)
(By (36))
12
(0s, 01, M1x((tt))
1
Σ
BTT
(Σ (
((v0
[- (
[v(0 +
+ v666))
M (n+(1)
12
0ax1
Oax1
i=1
Ci EE (
- (ors, M2 (1
(10 (t)
+ wi ((()
wp v()) M M2 (m+1)
n2
0ax1
(By Claim 4)
1 n
C (16)) (1)
+ +wi g(t) - wo v()))
(Ein )=
-(bb, (1)(
M (1x()
n2
0ax1
This means
((tt)))( - ((mm,((nt))
1 m
Σ
y0 (1100)
Ci
(
T(())
wi v())
+
+ wi
n
Since the choice of the query w(n+1) is arbitrary, we get
1 Σ (166 + tw W1+1 wt +
-000-w)))) 189-ww) wT vn)))
(
1=1
In particular, when we construct Zn such that (0) = $6-1, = 70; and (6) Rú we get
n
wi+1 = wi +
Σ
(R, + ywu 0; w/ Φi-1)es-1
n
where
e, = Σ x i Øk E Rd
which is the update rule for pre-conditioned TD(A). We also have
(n+1)
( (Vir Myxllttt) (wu,, (Ww,o(0+t)) (((+1))
yi
This concludes our proof.
35

---

Published as a conference paper at ICLR 2025
A.6 PROOF OF THEOREM 3
Proof. We recall from (18) that the embedding evolves according to
741 74
-wooHead Zz; Pi
(D.())
ID
M
(D)
Pl TT
Q (D wi)
n
1 WW
LinAttm Z5; P TD.11) QTD MTD.(1))
Z
n
LinAttn/ Z5; P. (2))
Q MTD.(2))
In this configuration, we refer elements Zi as (off 121 .1(1)) way.
the following
i= #+1 n+1
(n+1)
2
2j
if
Z1
yi in)
Vi in+1)
y
h
hi
in)
h ((+t1)
where we recall that Z e ((2d+2)x(a+))
R. Sometimes, it is more convenient refer the first half and second half of 2 separately, by, e.g..
(
E
IRdd,
Vi
E R and hy
(0)
(t) Rd, ( E Rdd i... 2/
Then we have
(n)
(n+1)
((1)
(nn)
g(
Zi
(n+1)
Vi
11
1
11
(n+1)
We further define as shorthands
X1 = 11
(m)
2ddx
Y = 11
(n)
gxxn
H
(m)
1xn.
Then express
Xi
It (n+1)
24
Vi (nt1)
Hi
11 (n+1)
For the input Zo. we assume $0 (
0 and (1)
0 for i = I...., n 1. All other entries of
Zo are arbitrary. We recall our definition of MDD.(1) (T) (17), { {p (TTB(11)
p.pp)i)ppBBB.. (T QT W
in (15) and (16). We again express
Q!
M1
La
E 2dx2d
Oaxd
0axd
B,
GT
Odxe
22d2d
Oaxd
Oaxd
CCT
CT
A, -
B,M1
E 2dx2d
Oaxe
Oaxd
Q
A
02ax2
(22422)((24+2)
Ozxed
02x2
We now proceed with the following claims that assist in proving main theorem.
36

---

Published as a conference paper at ICLR 2025
Claim 1. X, No, (n+1)
= =10 Yi Yo, Vosuin+1) 100
(n+1)
(n+1) (n+1) VI.
We define
v(() = P. (1) (B.(1)
24 2T
E E
v1() = (. Z,M (2)
Z @ Z
(2244))((H
Then the evolution of the embedding can written as
1 Z1+1 Zi + Wi
(1)
By simple matrix arithmetic, realize Wi is merely summing up the (2d 1)-th row of v,(1) and the (2d + 2)-th row of vy(2) and putting the result on its bottom row. Thus, we have
v2
(2d++1)x(n11) ( (2d + 1i (2d (2d+2)
2o (22042)x((++)
Wi
(2)
where 17 (2d + 1) and v/(2 (2d + 2) respectively indicate the (2d + 1)-th row of Vi and the (2d 2)-th row of v1(0) It clearly holds according the update rule that
741(1 2d + 1) = 711 2d + 1)
X1+1 Xx;
(n+1) = =] =x/ntt) (n+1).
[+1
Y1+1 = Yí;
(n+1) yi (n+1).
vi+1
Then, we can easily arrive at our claim by a simple induction. In light this, we drop the subscripts
of X1, x1, (1) Yi and (9) for all i = 1...., n + and write Z₁ as
x(++1)7
2 =
ylu+1) H h (n+1)
Claim 2.
H141 H1 - (Hi (H,+Y--)X* -Y)X AxX
n
(n+1) = 1/ (x+1) -11, +Y - -YX (xt(+1))
1+1
where g(i) = Σi k= F and Y = [,(1),1()
((2) j(m]] RIx"
We show how this claim holds by investigating the function of each attention head in our formulation. The first attention head, corresponding to vi in claim 1, has the form
(1) TD.))
Z Q %1)
We first analyze P. TD.(1) Z,M TD.(1) It should be clear that (TD,(1) Zi selects out the (2d + 1)-th row of Zi and gives us
Ozdxn Y
O2dx1 y Iin
P.
01xn
The matrix M TD) is essentially computing Y y and filtering out the (n + 1)-th entry when applied to P T)) Zj. We break down the steps here:
p (1) 7,
37

---

Published as a conference paper at ICLR 2025
-P
(D.
Zi (In+11 - Un+idiag ([1
M TD.2)
-P.(D(() Z,M (D.(2)
P.D.(1))
21Un+1diag(
1
b MTM
2
O2axx
02ax1
O2dxx y(2)
02ax1
0
02ax1
Y
0
[ + y())
Σ y0)
n++1
i= (() TD.)
m+1
01xn
0
O2dxn
02dx1
Ozdxn
02dx1
Y
0
Y
0
Oixe
0
01xn
0
0ndxn Y- Y
02dx1
01xn
We then analyze the remaining product Z Q11 Zz.
Zi 1IQ22 Zi
XT
YT
HT
A
02dx1
02ax1
X
x(+1))
a((+1)7
y(x+1)+
hy
(x+1))
x12d
0xx2d
0
0
Hi
Y
h
y(a+1)
(
X
(
((++1)
XTA A
A
Oax1
0
Onx1
0
Y
y
(n+1)
(a+1)
Hi
X A
X (1z(1)1)
(
A,X
((++1) Apx(1+1)
Putting them together, get
P (T) 2, ID,1)) (Z QF
Ozdxn Y-Y
02ax1
X A
X Anx((+t))
0
(((+1)" A,X
x(++1)" App((+11))
01xn
O2dxn
02ax1
(Y - -) A (Y (T-Y)xTAAa
01xn
The second attention head, corresponding to v,(2) in claim 1, has the form
p TT) 2 (D.22)
Z Q! 2i).
It's obvious that P (2) selects out the (2d + 2)-th row of Zi as
0(24+1)x1 00 ((241))x
p (2) Zi
Hi
b
(n+1)
Applying the mask M (2). we get
p (2) Z.M (2,
1)2a+1)xm
0(2441)x1
Hi
The product Z, QT 2, is identical to the first attention head. Hence, we see the computation of the second attention head gives us
P
()) TD(2)
Q.0Z %
Opa4t)yn H
022+11)x1 0
XT
XT A ((+1)) A
(((+1))
A.X
38

---

Published as a conference paper at ICLR 2025
1raa HiX TAX
0124+1)x1
A,X
H2X Axx(+t1))
Lastly, the matrix Wi combines the output from the two heads and gives
P (1) ZMM)
Z 21
0(24+1)x 1)x1 (HI+Y-P) AxX (Hr + - A
0pa+1)xn
Wi
p T2)) Z.M (2.)
2 (TQP2 21
8)X
Hence, we obtain the update rule for Hi and 111
((+t1)_ as
H1+1 H1 _
1
( A,X
71
1x+ = m1 (n+1) 1
(Hi Y -Y)x Asz(a+1)
n
and claim 2 has been verified.
Claim 3.
MIz
1 Σ
B. M2 +Y M -
3=0
ld
Oaxd 0axd
for = 1,....,. nn + 1, where M2 =
Odxd
Following claim unroll H1+1
H41 = Hi+ + - -
H. Hi-1 -
(H1-1 Y P)XT A-1X
n
1
8)X H1 Ho -- - - AnX.
n
We therefore can express H141
H1I1 Ho + - Σ
+
n j=0
Recall that we have defined A₁ = B,M and assumed Ho 0. Then, we have
1 1
H41
Σ -V n ==
Note that the introduction of M2 here does not break the equivalence because Bj = M2Bj. We include it in our expression for the convenience of the main proof later.
With the identical procedure, we can easily rewrite (x++1
(x+1) 141
Σ
H
MrB,Mizl
1=0
In light of this, we define eho 0, and for = O....
11
Σ
BJ M2X + - +Y- E 2d
7 1=0
39

---

Published as a conference paper at ICLR 2025
We then can write
11 = (2 11411
(37)
for 1, .,.... n + 1, which is the claim we made.
Claim 4. The bottom elements of dh are always O, i.c., there exists sequence {w: E Rd such that we can express 1/ as
Wt
d1
0dx1
for all 0,1,...,L.
Since our B, here is identical to the proof of Theorem in A. for j = 1, Claim 4 holds for the same reason. We therefore omit the proof details to avoid repetition.
Given all the claims above, we proceed to prove our main theorem.
U141, M2x(nt1))
Vi, M2z((+t))
1
BJ M2X +Y - P)T, M1x(ut1))
+
72
((r,, M (xx(1+1))
1
Bi Max (0) (1)) + (0) g(0),, M1z(111)
n
(Wr, M1Z((+11)
1
Σ
Bi M2x (1)
(11r, (1x()) + - wit)) -111)11(()
M Maa(nt)))
(By (37))
71
2=1
vi0
-v(() + + + yi - g(o) Moslo-1)
-(rr, M1 M2zdtnt1)
+
Bi
12
0ax1
0ax1
1=1
{|
([Ciol()
-(br, M2 (((())))))))))
(y()) g(0) + w/ ((\) wa vl0)) M (x+11)
+
0dx1
(By Claim 4)
- -(bi, (1r,
1
Σ
{ (Qu(0(()1)) -g(0) wwTT((
wt v(o)]
M M
71
0xx1
This means
(11H)) = -w,, wlntt)
+ (Crili)) (100 -00--00 g(t) (Cw + g(T(()) - w/ vli)) vlatt)).
Σ 1=1
n
Since the choice of the query ((+1)) is arbitrary, we get
G (v())) 7( + wT +w] c(0) 10--w) wi v()) Ju(i.
W1+1 w/ +
1=1
In particular, when we construct Zo such that 11() -1-1,5(0) = di and y(t) Ri, we get
w1+1 wi + n Σ
CiR - Fi+ wT ww d)
which is the update rule for pre-conditioned average reward TD learning. We also have
h (x+1)
=( ( Mys(atl))) -((ut, cf(m+1)).
This concludes proof.
40

---

Published as a conference paper at ICLR 2025
B EXPERIMENTAL DETAILS OF FIGURE
We generate Figure with 300 randomly generated policy evaluation tasks. Each task consists of a randomly generated Markov Decision Process (MDP), randomly generated policy, and randomly generated feature function (See Section 3 for detailed definition). The number of states of the MDP ranges from to 10, while the features are always in R The reward is also randomly generated, but we make sure the true value function is representable (cf. Algorithm 3). This treatment ensures that the minimal possible MSVE for each task is always O. The discount factor is always ? 0.9.
C BOYAN'S CHAIN EVALUATION TASK GENERATION
To generate the evaluation tasks used to meta-train our transformer in Algorithm 1, we utilize Boyan's chain, detailed in Figure 3. Notably, we make some minor adjustments to the original Boyan's 's chain Boyan (1999) make infinite horizon chain.
Recall that an evaluation task is defined by the tuple (Po.P. 0). We consider Boyan's 's chain MRPs with m states. To construct po. we first sample m-ddimensional random vector uniformly [0, 1 and then normalize it to probability distribution. To construct p. we keep the structure of Boyan's chain but randomize the transition probabilities. In particular, the transition function p can be regarded random matrix taking value in m To simplify the presentation, p(s'1s (|s) to denote the probability of transitioning to from s. In particular, for = 1,..., m 2, we p(i, / 1) and p(é, + 2)
both p(s, s') and
e, with sampled uniformly from (0, 1). For the last two states, we have p(m /m 1) = 1 and p(-[m) is random distribution over all states. Each element 1" and the matrix
Rdxm are sampled i.i.d. from a uniform distribution over [-1,]] The overall task generation process is summarized in Algorithm 2. Almost surely, no task will be generated twice. In our experiments the main text, we use Boyan Chain MRPs, which consist of m = 10 states, each with feature dimension d = 4.
of the vector
m-1
m
Figure 3: Boyan's 's Chain of m States
Representable Value Function. With the above sampling procedure, there is no guarantee that the true value function always representable by the features. In other words, there no guarantee that there exists E Rd satisfying (s)) = (w, ((s)) for all s E S. Most of our experiments use this setup. It is, however, also beneficial sometimes to work with evaluation tasks where the true value function is guaranteed to be representable. Algorithm 3 achieves this by randomly generating a first and compute v(s) ((s)). The reward is then analytically computed as (Im - vp). We recall that in the above, we regard matrix in m
41

---

Published as a conference paper at ICLR 2025
Algorithm 2: Boyan Chain MRP and Feature Generation (Non-Representable))
1: Input: state space size m = |S], feature dimension
2: for E S do
3: ((s) ~ Uniform (111,1)] // feature
4: end for
5: Po ~ Uniform [0, 1)" // initial distribution
6: Po po/ >. (s)
7: ~ Uniform (-1,1)" // zeward function
8: p < 0mxm
// transition function
9: for i =1.,.... .....
do
6~ Uniform (0.
11: p(i,++1) t-r
12: p(i,i+ 2) f- 1
13: end for
14: p(m pmm-1.ml t
15: = Uniform [0. 1)m
16: =- 2
p(m, 1 <- z
18: Output: MRP (po. r) and feature map
Algorithm 3: Boyan Chain MRP and Feature Generation (Representable)
1: Input: state space size m = |S], feature dimension d, discount factor y
2: w" Uniform ((-1, 1, 1)d) // ground-truth weight
3: for s E S do
4: ((s) ~ Uniform (-1,1) // feature
5: o(s) + (w* ((s)) // ground-truth value function
6: end for
7: Po Uniform (00. 1) // initial distribution
8: Po Po/ >, Ps())
9: p + 0mxm // transition function
10: for = I,..., 2 do
11: e~ Uniform 1)]
12: p(i,i+ +-
13: p(i,i+ 2) +- 1
14: end for
15: pmm-1, 1.m) 11
2 Uniform
[0,, 1)
17: 2 // 2ss)
p(m, 1::0 t
19: + (Im up)u // reward function
20: Output: MRP (po. p, and feature map
D ADDITIONAL EXPERIMENTS WITH LINEAR TRANSFORMERS
D. EXPERIMENT SETUP
We use Algorithm dan for the experiments in the main text with Boyan's chain of states. In particular, consider context of length 30. feature dimension d = 4. and utilize discount factor 7 = 0.9. In Section 5, we consider a 3-layer transformer (L = 3), but additional analyses on the sensitivity to the number of transformer layers (L) and results from larger scale experiment with d = 8, n = 60, and [S] = 20 are presented in D.2. We also explore non-autoregressive (i.e., "sequential") layer configurations in D.3.
When training our transformer, we utilize an Adam optimizer (Kingma and Ba. 2015) with an initial learning rate of 0.001 and weight decay rate of 1 x 10-6 Po and Qo randomly
42

---

Published as a conference paper at ICLR 2025
initialized using Xavier initialization with gain of 0.1. We trained our transformer on 4000 different evaluation tasks. For each task, we generated a trajectory of length T = 347, resulting in n 2 = 320 transformer parameter updates.
Since the models in these experiments are small (~ 10 KB), we did not use any GPU during our experiments. We trained our transformers standard Intel 19-12900-HK, CPU, and training each transformer took ~ 20 minutes.
For implementation", used NumPy (Harris et al., 2020) to process the data and construct Boyan's chain, PyTorch (Ansel et al., 2024) to define and train our models, and Matplotlib (Hunter, 2007) plus SciencePlots (Garrett, 2021) generate our figures.
D.1.1 TRAINED TRANSFORMER ELEMENT-WISE CONVERGENCE METRICS
To visualize the parameters of the linear transformer trained by Algorithm 1, we report element wise metrics. For Po. we report the value of its bottom-right entry, which, as noted in (7), should approach one if the transformer is learning to implement TD. The other entries of Po should remain close to zero. Additionally, we report the average absolute value of the elements of Po, excluding the bottom-right entry, to check if these elements stay zero during training.
For Qo. we recall from (7) that if the transformer learned to implement normal batch TD, the upper- left d x d block of the matrix should converge some -Le, while the upper-right d x d block (excluding the last column) should converge to Ta. To visualize this, we report the trace of the upper-left d d block and the trace of the upper-right d x d block (excluding the last column). The rest of the elements of Qo should remain close to O, and to verify this, we report the average absolute value of the entries of Qo, excluding the entries that were utilized in computing the traces.
Since, Po and Qo are in the same product in (1) we sometimes observe during training that Po converges -PP and Qu converges QQ simultaneously. When visualizing the matrices, we negate both Po and Qo when this occurs.
It's also worth noting that in Theorem we prove a L-layer transformer parameterized as (7) with Cu = La implements L steps of batch TD exactly with a fixed update rate of one. However, the transformer trained using Algorithm could learn perform TD with an arbitrary learning rate (a in (5)). Therefore, even if the final trained Po and Qo differ from their constructions in (7) by some scaling factor, the resulting algorithm implemented by the trained transformer will still be implementing TD. In light of this, we rescale Po and Qo before visualization. In particular, we divide Po and Qo by the maximum of the absolute values of their entries, respectively, such that they both stay in the range [_ after rescaling.
D.1.2 TRAINED TRANSFORMER AND BATCH TD COMPARISON METRICS
To compare the transformers with batch TD we report several metrics following Von Oswald et al. (2023a); Akyürek et al. (2023). Given context E p((2+1)xxn and a query E Rd,, we construct the prompt
2(o,0)
C
0dx1
We will suppress the context C' in subscript when it does not confuse. We use ( = (((()) as shorthand. We use d, to denote the stationary distribution of the MRP with transition function and assume the context C is constructed based on trajectories sampled from this MRP. Then,
(Z) can define 0i E T where uo (x))) (8) = TF ; () for each & E S. Notably, Up is then the value function estimation induced by the transformer parameterized by 0 = {(P,, Qi)) given the context C.
In the rest of the appendix, we will use θTF as the learned parameter from Algorithm 1. As a result,
UTF = dpr denotes the learned value function.
We define θTD = ((P7D) QP)) (P))) with C = all (see (7)) and
(s) = TF, (Z(+), (TD)
'The code will he made publicly available upon publication.
43

---

Published as a conference paper at ICLR 2025
In light of Theorem 1, UID is then the value function estimation obtained by running the batch TD algorithm (8) on the context for iterations, using a constant learning rate a.
We would like to compare the two functions VTF and VID future examine the behavior of the learned transformers. However, UrD is not well-defined yet because it still has a free parameter a, the learning rate. Von Oswald et al. (2023a) resolve similar issue in the in-context regression setting via using line search to find the (empirically) optimal a. Inspired by Von Oswald et al. (2023a), we also aim find the empirically optimal for VTD. We recall that CTD is essentially the transformer TFL (Z)"), (TD) with only single free parameter a. We then train this transformer with Algorithm I. complete the definition of tTD- We
We observe that quickly converges and use the converged are now ready to present different metrics to compare UTp and UTD We recall that both are dependent on the context C.
Value Difference (VD). First, for a given context C, we compute the Value Difference (VD) to measure the difference between the value function approximated by the trained transformer and the value function learned by batch TD, weighted by the stationary distribution. To this end, we define.
VD(ETF, UTD) = ||TT VTDllá,
We recall that dp E RISI is the stationary distribution of the MRP, and the weighted &z norm is defined as I|ul|a = VE. d(s))
Implicit Weight Similarity (IWS). We recall that uTD a linear function, i.e., ByD(s) (wz, ((s) with wr. defined in Theorem We refer to this wy. as wTD for clarity. The learned value function is, however, not linear even with linear transformer. Following Akyürek et al. (2023), we compute the best linear approximation of uyp. In particular, given context C, we define
WTF arg min Pow
erelld,
Here & R is the feature matrix, each of which is (0s) Such wrr is referred to implicit weight in Akyürek et al. (2023). Following Akyürek et al. (2023), we define
IWS UTD) = doos (WTF, WTD)
to measure the similarity between WyF and wip. Here (. computes the cos similarity between two vectors.
Sensitivity Similarity (SS). Recall that UTF(s) = TF QTF) and VTD(s) (s) = (Z θTD). In other words, given context C,, both uTF(s) and UTD (s) are functions of Φ(s). Following Von Oswald et al. (2023a), we then measure the sensitivity of (Ir(s) and VTD(S) w.r.t. (s)). This similarity is easily captured by gradients. particular, define
SS(ere, UID) = Edg(s)dos (VTTFL(Z())) 0IF)
ToTF (0)) BTD)
d=
x=d(s))
Notably, it trivially holds that
WTD T OTD)
(=0(s)
We note that the element-wise convergence of learned transformer parameters (e.g., Figure 2a) is the most definite evidence for the emergence of in-context TD. The three metries defined in this section are only auxiliary when linear attention is concerned. That being said, the three metrics are important when nonlinear attention concerned.
D.2 AUTOREGRESSIVE LINEAR TRANSFORMERS WITH 1,2,3, LAYERS
In this section, we present the experimental results for autoregressive linear transformers with different numbers of layers. In Figure 4, we present the element- wise convergence metrics for autoregressive transformers with L that for the L = case, Po and Qo converge to the construction Corollary 1, which, as proved, implements TD(0) in the single layer case. For the I, = 2, cases,
4 layers. The plot with L is in Figure 2 in the main text. We can
that Po and Qo converge
44

---

Published as a conference paper at ICLR 2025
Final Po
Final Qq
R Mördes
3eeaacc
Erity (4:4:24)
Ang Ale Othere
0.-
18:01
2,000
30.03
3093
3008
3.0
32300
(a) Learned Po and Qo with L
(b) Element-wise learning progress of Po and Qu
Final Po
Final Qq
Ai Mordoe
Metiiss
0.5
Others
3200
3000
(c) Learned Po and Qo with L
3
(d) Element-wise learning progress of Po and Qo
Final Po
Final Qu
P Miric
Q. Metrice
toQd) 4.44)
Abr Othera
1.5
32303
(e) Learned Po and Qo with L 4
(f) Element-wwis learning progress of Po and Qo
Figure 4: Visualization of the learned autoregressive transformers and the learning progress. eraged across 30 seeds and the shaded region denotes the standard error. See Appendix D.I.1 for details about normalization of Po and Qo before visualization.
45

---

Published as a conference paper at ICLR 2025
the construction in Theorem 1. We also observe that as the number of transformer layers L increases, the learned parameters are more aligned with the construction of Po and Q with Co =
We also present the comparison of the learned transformer with batch TD according to the metrics described in Appendix D.1.2. In Figure 5, we present the value difference, implicit weight similarity, and sensitivity similarity. In Figures 5a 5d, we present the results for different transformer layer numbers L = 2, 3,4. 4. In Figure 5e, we present the metrics for a 3-layer transformer, but we increase 60.
the feature dimension d 8 and also the context length
In all instances, we see a strong similarity between the trained linear transformers and batch TD. We that the cosine similarities of the sensitivities near one, as the implicit weight similarities. Additionally, the value difference approaches zero during training. This further demonstrates that the autoregressive linear transformers trained according Algorithm learn implement TD(0).
Leamed and Barzh TD Comparison
Learned TP and Burch Comparion
Learned TF and Berch TD Comparison
1025
-025
-025
Coate Swilwety
Coale Stillery
Coate Sauilertyy
6.8
(1.00
Upier Dalentan
58
DWS
0.30
0.20
0.00
0.6
0.00
2000
2500
1000
MEPPa
MEP
L
(b) T.
(c)
3
Learned TP and Batch TD Comparison
Learned TF and Batch TD Comparioon
0.30
-.25
2.35
Cuate Saatealy
Volle Difforan::
Coalte Saalleity
0.8
11.00
ZWS
TWS
11:2
1405
5.05
2002
3000
0.00
2000
MRNN
MRR
T.
(e) 3 (d 8, 60)
Figure 5: Value difference (VD), implicit weight similarity (IWS), and sensitivity similarity (SS) between the learned autoregressive transformers and batch TD with different layers. All curves are averaged over 30 seeds and the shaded regions are the standard errors.
D.3 SEQUENTIAL TRANSFORMERS WITH = 2,3, 4 LAYERS
So far, we have been using linear transformers with one parametric attention layer applied repeatedly for L steps to implement an L-layer transformer. Another natural architecture in contrast with the autoregressive transformer is sequential transformer with I. distinct attention layers, where the embedding passes over each layer exactly once during one pass of forward propagation.
In this section, we repeat the same experiments we conduct on the autoregressive transformer with sequential transformers with 2, 3, as their architectures coincide when L 1. We compare the sequential transformers with batch TD(0) and report the three metrics in Figure 6. We observe that the implicit weight similarity and the sensitivity similarity grow drastically to near 1, and the value difference drops considerably after few hundred MRPs for all three layer numbers. It suggests that sequential transformers trained via Algorithm I are functionally close to batch TD.
Figure 7 shows the visualization of the converged (Pi, (1)=..2 of 3-layer sequential linear transformer and their element-wise convergence. Sequential transformers exhibit very special patterns in their learned weights. We see that the input layer converges to a pattern very close to our configuration in Theorem (1). However, the deeper the layer, we observe the more the diagonal of Q d, d + : 2d] fades. The P matrices, on the other hand, follow our configuration closely, especially for the final layer. We speculate this pattern emerges because sequential transformers have
46

---

Published as a conference paper at ICLR 2025
more parametric attention layers and thus assign slightly different role to each layer but together implement batch TD(O) as suggested by the black box functional comparison in Figure
Learndd and Barch TD Compartion
Learned TP and Barch TD Comparison
Learndd TP and Barch TD Comparison
Siibbity
(120
Dattmnzz
Crdto Boiluity
0.22
0.6
0.16
0.15
80.0
0.15
0.4
VD
0.
VD
0.4
VD
0.2
0.2
0.06
0.2-
0.0
10G0
2000
9030
4000
0.00
0.0
1000
5000
3030
4000
0.00
0.0
0
1000
2000
3030
4000
MiiPP
MEPP
(a) L
(b) =
(c)
Figure 6: Value difference (VD), implicit weight similarity (IWS), and sensitivity similarity (SS) between the learned sequential transformers and batch TD with different layers. All curves are averaged over 30 seeds, and the shaded regions are the standard errors.
Final Final Q
Ft Merdic
Melaica
2.0
11 Ang Alea Others
unQuia
1.0
unt ls dedra4)
Arg Abs Otbers
1.-
0.5
0.0
30:37
3
30.03
1003
240
(a) Learned Po and Qu
(b) Element-wise learning progress of Po and Qu
Final P.
Final Q
#. Mordee
Metalcs
2.0
PLE I]
1.2
Aug Ale Orfors
in (a) 4. Aug Ale Orlers
10-
0.5
0.0
1000
3000
2003
2000
0
1000
3000
2000
500
(c) Learned Pi and Q
(d) Element-wise learning progress of Pı and Q
Final P
Final Qg
Py Miedo
Metaicc
Arg Otface
1.5
-6
unQ) (nQ)):-dd::44
Are Ab Othere
0.0
3030
2300
3003
4000
3000
2300
3000
4000
22007
(e) Learned P2 and Q2
(1) Element-wise learning progress of Pa and Q2
Figure 7: Visualization of the learned L 3 sequential transformers and the learning progress. Averaged across 30 seeds and the shaded region denotes the standard errors. See Appendix D.I.1 for details about normalization of Po and Qo before visualization.
E NONLINEAR ATTENTION
Until now, we have focused on only linear attention. In this section, we empirically investigate original transformers with the softmax function. Given matrix Z, recall that self-attention computes its embedding as
Attn Z2 P,)) Mo (7 (QQ)..
47

---

Published as a conference paper at ICLR 2025
Cocapariaco
Learnod TF and Batch TD Comporion
0.00
1.0
DWS
0.25
IWS
Stailatty
Orebe Smiaart
50
0.4
0.2
6.0
1:000
2060
3000
0000
00
3000
0000
0.0
MRRPs
(a) General Value Function
(b) Representable Value Function
Figure 8: Value difference (VD), implicit weight similarity (IWS), and sensitivity similarity (SS) between the learned softmax transformers and linear batch TD. All curves are averaged over 30 seeds, and the shaded regions are the standard errors.
Let Zı E ((2d+1)x(++1) denote the input to the 1-th layer, the output of an I-layer transformer with parameters {(P, Q1)). is then computed as
Z41 Z + 1, Atm 2,; Pt, Q1) = Zi + PZMsoftmax (Z))))))
Analogous the linear transformer, we define
TFL Zo: (Pi, Q1 101-......1)) = 22d +
As shorthand, we use TF, (Zg) (Z)) denote the output of the softmax transformers given prompt Zo- We use the same training procedure (Algorithm 1) to train the softmax transformers. In particular, we consider a 3-layer autoregressive softmax transformer.
Notably, the three metrics in Appendix D.1.2 apply to softmax transformers as well. We still compare the learned softmax transformer with the linear batch TD in (8). In other words, the UTD related quantities are the same, and we only recompute Figure 8a, the value difference remains small, and the implicit weight similarity increases. This suggests that the learned softmax transformer behaves similarly to linear batch TD. The sensitivity similarity, however, drops. This expected. The learned softmax transformer TF, is unlikely to be a linear function w.r.t. to the query while UTD is linear w.r.t. the query. So their gradients w.r.t. additionally consider
related quantities in Appendix D.1.2. As shown
the query are unlikely to match. To further investigate this hypothesis, evaluation tasks where the true value function is guaranteed to be representable (Algorithm 3) and is thus linear function w.r.f. the state feature. This provides more incentives for the learned softmax transformer to behave like linear function. As shown in Figure 8b, the sensitivity similarity now increases.
F EXPERIMENTS WITH CARTPOLE ENVIRONMENT
In this section, we present additional experimental results demonstrating that in-context TD emerges after large-scale pretraining using Algorithm where dask is derived from the CartPole environment (Brockman et al., 2016).
F.1 CARTPOLE EVALUATION TASK GENERATION
Recall that in the main text, as well as Appendix D and E, the transformers are pre-trained with tasks drawn from dusk based on Boyan's Chain (See Appendix C). Here, we extend the analysis by introducing dask based on the CartPole environment. Figure 9 provides introduction to the CartPole environment.
Recall that an evaluation task is defined by the tuple (Po,P,T, Φ). In the canonical CartPole envi- ronment, the states are a vector s E R where the entries are the current position of the cart, the velocity of the cart, the angle of the pole, and the angular velocity of the pole. In experiments, the initial state distribution Po and environment transition dynamics p(s'|s [ss a) are given by the standard CartPole equations (e.g. see OpenAl CartPole Github). These transition dynamics, which we denote
48

---

Published as a conference paper at ICLR 2025
Figure 9: The OpenAl Gym CartPole environment (Brockman et al., 2016) is classic RL control task where the goal is to balance a pole on a cart by applying forces to move the cart left or right. The state consists of the cart's position and velocity and the pole's angle and angular velocity. The episode ends if the cart moves out of bounds or the pole falls beyond a threshold angle.
as PCanPoie(s'|ss (8, a), implicitly depend on the physical parameters VP = (Mean, pole: 9, pole: T; f) representing the mass of the cart and pole, gravitational constant, length of the pole, frame rate, and the force magnitude. We abuse the notation of (s (s, v to highlight the transition dependency on V. The joint distribution over these parameters, denoted by Av. defines the the possi- ble CartPole environments. In our experiments, we sampled Meart. mpple pplle ~ Uniform [0.5, 1.5), g ~ Uniform [7, 12], Uniform [0.01, 0.05), f Uniform [5, 15].
Then, the state transition function (()'ss) which characterizes an MRP is defined using PCanPole a), and a fixed random policy Te (als)) parameterized by Uniform [0, 1)]. Un- der ne (a(s). the probability of moving the cart to the right is and the probability of moving the cart to the left - 6. This means that p(s'ls) = E. (0.1) p(s'|s, a) (T (als) where 0 means going left and means going right. The environment is extended to an infinite horizon. When the pole falls, or the cart moves out of bounds, the state is reset by sampling new initial state from po-
Rather than using the standard CartPole observations and reward structure +1 per time step until failure, provide diverse set of reward functions and features by sampling and & randomly. In CartPole, the state is continuous, resulting in an infinite state space S. To address this, we use tile coding (Sutton and Barto, 2018) with random projection generate feature function 6 :: + Rd for E S. Tile coding with random projection maps to a feature vector sampled from - R, is mapped also sampled from Uniform [anfform-11,1] The joint distribution over random features and reward functions denoted pg.-(d (d). For each CartPole MRP, we sample from Ag.r to obtain the feature and reward functions and r. This approach, detailed in Algorithm 4, enables the transformer to encounter variety of tasks during pre-training.
Uniform [--1, 1)d]. Similarly, for the reward function
reward value,
Algorithm 4: CartPole MRP and Feature Generation
Input: feature dimension d, action space A {0, joint distribution over CartPole parameters As, joint distribution over features and rewards Ag.r
2: V Av // sample CartPole parameter
3: Po < Uniform (--005, 0.05)*] // CartPole initial distribution sample features and rewards
4: Ø, r. +- A p.
5: e Uniform (0,, 1)] // sample random policy parameter
6: p(s'ls) + Eas PCan)o)e a; V) // CartPole state transition
7: Output: MRP (po, r) and feature map
F.2 EXPERIMENTAL RESULTS OF PRE-TRAINING WITH CARTPOLE
In our experiments in Figure 10, we pre-train a 3-layer autoregressive transformer using Algorithm 1, where the task distribution dusk is generated using CartPole MRPs (see Algorithm 4) with a feature
vector of dimension d = 4. We used a significantly larger context window length n = 250. Despite the increased complexity of the transition dynamics in the CartPole environment compared to Boyan's
49

---

Published as a conference paper at ICLR 2025
chain environment used in Figure 2, results demonstrate that Po and Qu still converge to the construction in Theorem (up to some noise), which we proved exactly implements TD(0).
is worth noting that theoretical results (Theorem 2), which prove that the weights implementing TD are in the invariant set of the updates in Algorithm 1, do not depend on any specific properties of the environment Thus, it is unsurprising that TD(0) emerges naturally even after pre-training environments with complicated dynamics.
Final
Final
Po Metrics
Qu Metrics
R1-1,-11
#70kiid +0) *39x
(a) Learned Po and Qo after 10000 MRPs
(b) Element-wise learning progress of Po and Qo
Figure 10: Visualization of the learned transformers and the learning progress after pretraining with the CartPole environment for 10,000 MRPs. Both (a) and (b) are averaged across 30 seeds and the shaded regions in (b) denote the standard errors.
G INVESTIGATION OF IN-CONTEXT TD WITH RNN
We have focused primarily on the transformer's capability to implement TD in context. Before transformers, the canonical architecture to tackle sequence modelling problems is the recurrent neural network (RNN) (Elman, 1990; Bengio et al., 2017). Thus, it's worth investigating the algorithmic capacity of RNN implementing TD in its forward pass. In particular, we try to answer the following two questions in this section:
1. Can RNN implement TD context?
2. Does in-context TD emerge in RNN via multi-task pre-training?
A canonical deep RNN with I, layers is parameterized by Will, w!l)) (0))
-= Let m
denote the dimension of the input tokens and h denote that of the hidden states, respectively. Then, we have W-(0) E haxm, will)
E haxh for / = 1...., L - 1, and wi!h E RAXA (() E R for 1 1 o....,. L L-1 1. Let (1) denote the input token and a; denote the hidden state for layer I at time step Unlike transformers that process the whole sequence at once, an RNN processes one token after another sequentially by updating the hidden states. The hidden state evolves according
(1) 4+41
(w()) (1) W a b(!))
l where f is an activation function. In addition, we have xt = az ((-1) for and | = 1,..., In other words, the input to the next depth is the hidden state from the previous depth except for the first layer. The initial hidden states o!" I O,..., L are selected arbitrarily. Popular options include zero initialization and random normal initialization.
When we apply RNN policy evaluation, are interested in predicting scalar value the end, also known as many-to-one prediction. Suppose the input sequence has tokens one typically passes L-1 the final hidden state at the last recurrent layer, through fully connected output layer W. E R 1xh such that
i = Wo
G.1 THEORETICAL ANALYSIS OF LINEAR RNN
We first investigate Question 1 via a theoretical analysis of RNN the context of TD. Due to the intractable difficulty of nonlinear activations present in deep neural network analysis, we resort to
50

---

Under this representation, it impossible to construct the TD error, not to mention applying the semi-gradient term. Therefore, it is safe for us to claim that linear RNN is incapable of implementing since it is only Published as a conference paper at ICLR 2025
analyzing single-layer linear RNN, i..., L 1 and is the identity mapping. Hence, will drop the superscript indicating the layer index and in our notation to simplify the presentation. We shall also remove the bias term b, because it is a constant independent of the context. Under this formulation, the hidden state evolves according to
04+1 = Wantt Waa44-
we initialize ag = 0, we then have
4g =
d1 Wa + W Va WorDo
42 = Wan + WarT1 = W Wa + Waa W ar
43 = W Wan + War = W 02 + W az + W an To
Assuming a sequence of tokens, the final hidden state On is
On W Waz
Applying a linear output layer Wo E 1xh to the hidden state for value prediction, we then get
i = Woa,
W. WoWn
WanTt
Tt:
(38)
t=
t=
where w, = (W.Wn-t-lW.. aar
T is vector. (38) demonstrates that the predicted value is the sum of the inner product between each token and some vector for linear RNN. Recall that each context token 2y for in-context TD is defined
Ft
7
R
corresponding to column t of the prompt Z. Hence, write
Φ
W
T
R.
TD in its forward pass. This result easily extendable the multi-layer performing linear combinations of the tokens, thus reducible to the format of (38). One important insight gained by comparing the forward pass of RNN and transformer under linear activation is that one at least needs x, QI, where Q is a square matrix to have any hope to compute the TD error, which necessary for TD. Therefore, we speculate that deep RNN equipped with common nonlinear activation such as tanh and ReLU is also unable to implement TD in context. We will leave the investigation to Question 2. For now, we can confidently give negative answer to Question concerning linear RNNs.
G.2 MULTI-TASK TD WITH DEEP RNN
We answer Question 2 via an empirical study with deep RNN. We employ a 3-layer RNN with hidden state dimension of f = 4 and tanh as the activation function and train it via multi-task TD (Algorithm 1) on 4,000 randomly generated Boyan's chain MRPs with a feature dimension Since cannot apply mask M like in the transformer to distinguish the query from the context, we instead append a binary flag to each token for the same purpose. Suppose there are n context columns, the prompt Z has the form
= 4.
Φ
22
on+1
7
n
0
R (2d+2)x(n+1)
R1
R₂
R.
0
0
0
1
51

---

Published as a conference paper at ICLR 2025
The forward pass of the deep RNN processes the tokens sequentially in the prompt update the hidden states. The final hidden state the last layer of the RNN is fed into fully connected layer to output scalar value prediction. Figure 11 shows the learning curve of the RNN throughout the multi-task TD training. The MSVE decreases for the first 1, 1,000 MRPs and stays low for the remainder of the training. Thus, some learning occurs during the training of RNN. However, it is unclear whether it is implementing in-context TD. To clarify, we use the last checkpoint of the model and repeat the same experiment used generate Figure 1. We gradually increase the context length and verify if the MSVE drops as observed in the transformers. We run the experiment on the Loop environment used generate Figure and the Boyan's 's chain environment used for training for 500 instances each to produce Figure 12. The MSVE increases with context length in both environments for the trained RNN, exhibiting a trend opposite to transformer. Furthermore, the standard errors are much higher than in Figure despite having more runs. Therefore, the prediction does not improve with more context data for the RNN, indicating the absence of any in-context policy evaluation algorithms. Consequently, the answer Question 2 is again negative.
30
MSVE
10
1000
2000
3000
4000
Number of MRPs
Figure 11: RNN MSVE against the number MRPs multi-task TD training.
3.0
2.8
MWW mmM
2.6
2.2
MSVE
2.4
2.1
2.2
2.0
2.0
1.9
0
20
40
60
80
1.00
200
40
60
80
100
Context Length (1)
Contecc Length (5)
(a) Loop
(b) Boyan's 's chain
Figure 12: MSVE vs. context length with the trained RNN. The shaded regions are the standard errors.
H NUMERICAL VERIFICATION OF PROOFS
We provide numerical verification for our proofs by construction (Theorem 1, Corollary 2, Corollary 3, and Theorem as a sanity check. In particular, we plot log - wi) v4 x++1 against the number of layers 1, For example, for Theorem
first randomly generate Zg and {Ci). Then y/ (n+1) is computed by unrolling the transformer layer by layer following (3) while w, is computed iteration by iteration following (8). We double-precision floats and run for 30 seeds, each with prompt.
52

---

Published as a conference paper at ICLR 2025
As shown in Figure 13, even after 40 layers/iterations, the difference is still in the order of 10 10-10 It not strictly 0 because of numerical errors. It sometimes increases because the accumulation of numerical errors.
TDJ0)
Reidual Goudiett
TD(A)
Aug Reward
log (0h, m0)
Layers
Figure 13: Differences between transformer output and batch TD output. Curves are averaged over 30 random seeds with the (invisible) shaded region showing the standard errors.
53


Reviews:
[{"review_id": "KBl68283Vh", "reviewer": "Reviewer_77PR", "summary": "The paper shows that linear transformers can execute TD(0) over their context both by constructing the matrices that would do so, but also by showing a pre-training method that gives rise to such properties. Finally they show similar constructions for other Policy Evaluation algorithms like TD-lambda, Average-Reward TD etc. Besides some empirical evidences the paper presents many proofs of its statements.", "questions": "1. Would it be possible for the authors to change the paper's title to something that better reflects what is shown by the paper?\n2. Could the authors better justify the study of In-Context RL, both for Linear Attention-based transformers and otherwise?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 2, "contribution": 2, "strengths": "The paper argues with proofs and empirical evidences for its results and is able to remove any possible doubt of their veracity. Furthermore it's a novel exploration of the capabilities of linear transformers and further helps the understanding of In-Context Reinforcement Learning.", "weaknesses": "The title of the paper is rather strong, not only is the paper only about Linear Attention, as opposed to the commonly used Softmax attention, but also they prove that transformers with a specific kind of training implement TD-Learning, so I believe \"Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning\" would be more appropriate.  Furthermore one must wonder how relevant is studying Reinforcement Learning implemented by linear attention over such short horizons such as 40 steps, as most RL problems involve orders of magnitude longer trajectories, whereas attention tends to become computationally expensive and lose precision as the context length grows.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5e7nClskyv", "reviewer": "Reviewer_5HXw", "summary": "The authors performed a mathematical analysis of the transformer architecture commonly used in the field of in-context reinforcement learning (ICRL). They proved that when a transformer with a linear attention structure is trained to estimate value functions in a multi-task environment, the forward pass within the transformer becomes equivalent to temporal-difference (TD) learning in reinforcement learning. Even when out-of-distribution task contexts are provided as input to the transformer, this TD learning effect is induced internally, enabling it to estimate the correct value function. The authors experimentally demonstrate this in a simple Boyan's chain environment.", "questions": "1. The authors explain that in-context TD learning occurs when the task distribution is sufficiently challenging. I wonder why such a phenomenon occurs. They trained the transformer across thousands of Markov Reward Processes. Also, I wonder if the same TD learning effect would emerge internally within the transformer if it were trained in a single-task setup rather than a multi-task one.\n\n2. In most prior research, transformers are directly applied to perform control tasks. However, in this paper, the transformer is used to estimate the value function. While a generalized approach to value function estimation is indeed important in reinforcement learning, this study does not directly address the control problem in RL. I would be interested to hear the authors' perspective on how their approach could be applied to control tasks in future research.\n\n3. When a transformer estimates the value function, its internal operations align with TD learning. Most transformers, however, serve as policies rather than value function estimators. It would be interesting to know the authors’ perspective on whether, if the transformer were to act as a policy instead of estimating a value function, the internal operations would still resemble TD learning, or if another reinforcement learning algorithm would take place internally.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors mathematically analyze that defining a Transformer as a value function estimator and training it in a multi-task environment makes the attention mechanism within the Transformer equivalent to TD(0) learning. The proof process appears to be concrete, and the resulting application seems straightforward. Particularly, while previous reinforcement learning studies using Transformers have merely mentioned its high performance, this paper analyzes why the Transformer structure works so effectively. In this regard, this paper seems to present a notable novelty. \n\n2. The authors mathematically demonstrated that the internal operations of the transformer can be extended not only to TD(0) learning but also to various RL algorithms, such as TD($\\lambda$), average reward TD, and residual gradient. This suggests that these algorithms can be applied in diverse ways to approximate a generalized value function through a transformer. \n\n3. The proposed theorems and corollaries in the paper are meticulously detailed in the appendix, where each proof is presented with clarity and rigor. This comprehensive approach allows readers to follow the logical reasoning behind each statement. This thoroughness underpins the robustness of the paper and enhances its reproducibility, which will be beneficial for future researchers.", "weaknesses": "1. The experimental setup and evaluation metrics may not effectively demonstrate the algorithm's impact. The authors defined the value function estimation error for new states in out-of-distribution (OOD) tasks as the evaluation metric in Figure 1. However, since the Boyan's chain environment is relatively simple, it's possible that generalization effects could also be achieved using standard architectures like MLPs and RNNs, resulting in graphs similar to those in Figure 1. This raises the question of whether the results in Figure 1 reflect the effects of the transformer's in-context TD learning or if they could be similarly obtained with conventional networks. To address this, it seems the authors should consider conducting comparative experiments.\n\n2. The comparison experiments seem insufficient. If there were experiments comparing value estimation between the authors' proposed value-estimating Transformer algorithm and existing reinforcement learning algorithms that utilize Transformers (other than behavior cloning methods), it would support the claim that the proposed algorithm is more effective. Although this paper focuses more on mathematical analysis, it seems to not have enough experiments on other algorithms or tasks. \n\n3. The paper appears to be written in a way that makes it difficult to read. It seems to borrow equations and results from previous studies, but the notation and equations are challenging to follow at first glance. For instance, it is hard to understand how equations like Equation (4) were derived. A brief explanation of the existing equations would be helpful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MKdfqL8KgF", "reviewer": "Reviewer_zyPc", "summary": "This paper provides theoretical analysis on the equivalence between a forward pass and in-context TD in a simple setting - linear single-layer transformers.", "questions": "1. In Corollary 3, the authors show that the transformer can implement TD(λ) updates with a specific mask matrix, $M_{TD(\\lambda)}$. The matrix structure requires strict alignment of eligibility traces, which might be difficult to achieve in real tasks where temporal dependencies and eligibility traces shift dynamically. \n\n2. Lines 523-524 in the Conclusion are unclear. Are you trying to emphasize that policy iteration or optimization forms the key foundation?\n\n3. In Lines 531-532, the claim seems unsupported; working well in a small-scale setting doesn’t necessarily imply the same for larger, more complex tasks.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "It is a new try and an interesting direction to investigate the equivalence between transformers and in-context TD.", "weaknesses": "1. Many of the proofs rely on strict initial conditions, like specific embeddings and matrix structures. The validity of these proofs depends heavily on these controlled setups, which may not be feasible or sustainable in more complex, real-world settings. As a result, some claims in the paper appear intuitive and unsubstantiated.\n2. While the paper claims the emergence of in-context TD learning through multi-task TD pretraining, the empirical analysis is constrained to quite simple tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)