In-context learning refers to the learning ability of a model during inference time without adapting its parameters. The input (i.e., prompt) to the model (e.g., transformers) consists of both a context (i.e., instance-label pairs) and a query instance. The model is then able to output a label for the query instance according to the context during inference. A possible explanation for in-context learning is that the forward pass of (linear) transformers implements iterations of gradient descent on the instance-label pairs in the context. In this paper, we prove by construction that transformers can also implement temporal difference (TD) learning in the forward pass, a phenomenon we refer to as in-context TD. We demonstrate the emergence of in-context TD after training the transformer with a multi-task TD algorithm, accompanied by theoretical analysis. Furthermore, we prove that transformers are expressive enough to implement many other policy evaluation algorithms in the forward pass, including residual gradient, TD with eligibility trace, and average-reward TD.

§ INTRODUCTION

In-context learning
has emerged as one of the most remarkable abilities of large language models [13, 32, 39, 10].
In in-context learning,
the input (i.e., prompt) to the model consists of both a context (i.e., instance-label pairs) and a query instance.
The model then outputs a label for the query instance during inference (i.e., the forward pass).
An example of the model input and output could be

\begin{align}
\label{eq sl example}
\underbrace{5 \to \text{number} ; a \to \text{letter} ; 6 \to}_{\text{input}} \underbrace{\text{number}}_{\text{output}},
% \qty{5 \to \text{number} ; a \to \text{letter} ; 6 \to}^{\texttt{// input}} \qty{\text{number}}^{// \text{output}},
\end{align}

where “$5 \to \text{number} ; a \to \text{letter}$” is the context consisting of two instance-label pairs and “$6$” is the query instance.
Based on the context,
the model (e.g., 48, 49, 1) infers the label “number” for the query “6”.
this entire process occurs during the model's inference time without any adjustment to the model's parameters.
Understanding the mechanism behind in-context learning has recently garnered significant attention [18, 3, 53, 2].

The example in (<ref>) illustrates a supervised learning problem.
In the canonical machine learning framework [9],
this supervised learning problem is typically solved by first training a classifier based on the instance-label pairs in the context using methods such as gradient descent, and then asking the classifier to predict the label for the query instance.
3, 53, 2
show that transformers are able to implement this gradient descent training process in their forward pass without adapting any of their parameters,
providing a possible explanation for in-context learning.

Beyond supervised learning,
intelligence involves sequential decision-making,
where Reinforcement Learning (RL, 47) has emerged as a successful paradigm.

Can transformers preform in-context RL during inference, and how?
To address these questions,

we start with a simple evaluation problem in a Markov Reward Process (MRP, 38).

In an MRP,
an agent transitions from state to state at every time step.
We denote the sequence of states that the agent visits by $(S_0, S_1, S_2, \dots)$.
At each state,
the agent receives a reward.
We denote the sequence of rewards that the agent receives along the way as $(r(S_0), r(S_1), r(S_2), \dots)$.
The evaluation problem is to estimate the value function $v$,
which computes for each state the expected total (discounted) rewards the agent will receive in the future.
An example of the desired input-output could be

\begin{align}
\label{eq rl example}
\underbrace{S_0 \to r(S_0); S_1 \to r(S_1); S_2 \to r(S_2); s \to}_\text{input} \underbrace{v(s)}_\text{output}.
\end{align}

the above task is fundamentally different from supervised learning as the goal is to predict the value $v(s)$ and not the immediate reward $r(s)$.

Moreover, the query state $s$ is arbitrary and does not have to be $S_3$.

Temporal Difference learning (TD, 46) is the

most widely used RL algorithm for solving such evaluation problems in (<ref>).
And it is well known that TD is not gradient descent [47].

In this work,
we make three main contributions.
First, we prove by construction that transformers are expressive enough to implement TD in the forward pass,
a phenomenon we refer to as in-content TD.
In other words,
transformers can solve problem (<ref>) during inference time via in-context TD.
Beyond the most straightforward TD,
transformers can also implement many other policy evaluation algorithms,
residual gradient [7],
TD with eligibility trace [46],
and average-reward TD [50].
In particular,
to implement average-reward TD,
transformers require the use of multi-head attention and over-parameterized prompts, e.g.,

\begin{align}
\label{eq rl example memory}
\underbrace{S_0 \to r(S_0) \, \square; S_1 \to r(S_1)\, \square; S_2 \to r(S_2)\, \square; s \to}_\text{input} \underbrace{v(s)}_\text{output}.
\end{align}

Here, “$\square$” acts as a dummy placeholder that the transformers will use as “memory” during inference.
we empirically demonstrate that by training transformers with TD on multiple randomly generated evaluation problems,
in-context TD emerges.
In other words, the learned transformer parameters closely match our construction in proofs.
We call this training scheme multi-task TD.
we bridge the gap between our theories and empirical results
by showing that for a single layer transformer,
the transformer parameters required in the proof to implement in-context TD is in a subset of the invariant set of the training algorithm multi-task TD.

§ BACKGROUND

Transformers and Linear Self-Attention.
All vectors in this paper are column vectors.
We denote the identity matrix in $\R^n$ by $I_n$ and an $m \times n$ all-zero matrix by $0_{m\times n}$.
We use $Z^\top$ to denote transpose of $Z$ and use both $\indot{x}{y}$ and $x^\top y$ to denote the inner product.
Given a prompt $Z \in \R^{d \times n}$, standard single-head self-attention [51] processes the prompt by
% \begin{align}
\text{Attn}_{W_k, W_q, W_v}(Z) \doteq W_v Z\,  \softmax{Z^\top W_k^\top W_q Z},
% \end{align}
where $W_v \in \R^{d \times d}, W_k \in \R^{m \times d}$, and  $W_q \in \R^{m \times d}$ represent the value, key and query weight matrices, respectively.
The softmax function is applied to each row.
Linear attention has recently drawn more attention [43, 53, 2],
where the softmax function is replaced by an identity function.

Given a prompt $Z\in\R^{(2d+1) \times (n+1)}$,
we follow 2 and define
linear self-attention as

\begin{align}
\label{eq linear attention}
\attn(Z; P, Q) \doteq PZM(Z^\top Q Z),
\end{align}

where $P \in \R^{(2d+1) \times (2d+1)}$ and $Q \in \R^{(2d+1) \times (2d+1)}$ are parameters and $M \in \R^{(n+1) \times (n+1)}$ is a fixed mask of the input matrix $Z$, defined as

\begin{align}
\label{eq td0 mask}
\textstyle M \doteq \mqty[I_{n} & 0_{n \times 1} \\
0_{1 \times n} & 0].
\end{align}

Note that we can view $P$ and $Q$ as reparameterizations of the original weight matrices for simplifying presentation.

The mask $M$ is introduced for in-context learning, following 2, to designate the last column of $Z$ 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 a different mask $M'$, when necessary, by defining $\attn(Z; P, Q, M') = PZM'(Z^\top QZ)$.

In an $L$-layer transformer with parameters $\qty{(P_l, Q_l)}_{l=0,\dots, L-1}$, the input $Z_0$ evolves layer by layer as

\begin{align}
\textstyle Z_{l+1} \doteq Z_l + \frac{1}{n} \attn_{P_l, Q_l}(Z_l) = Z_l + \frac{1}{n} P_l Z_l M(Z_l^\top Q_l Z_l). \label{eq: Z_l update}
\end{align}

Here $\frac{1}{n}$ is a normalization factor simplifying presentation.
We follow the convention in 53, 2 and use

\begin{align}
\textstyle \tf_L(Z_0; \qty{P_l, Q_l}_{l=0,1,\dots L-1} ) \doteq -Z_L[2d+1, n+1] \label{eq: tf output convention}
\end{align}

to denote the output of the $L$-layer transformer,
given an input $Z_0$. Note that $Z_l[2d+1, n+1]$ is the bottom-right element of $Z_l$.

In-Context Supervised Learning as Gradient Descent.
A linear regression task can be represented by an instance distribution $d_\fX$ and a ground truth wieght $w_*$.
A training set
$\{(x^{(i)} \in \R^{2d}, y^{(i)} \in \R)\}_{i=1,\dots, n}$
is usually constructed by sampling $n$ instances $\{x^{(i)}\}$ from $d_\fX$ in an i.i.d. manner and constructing the targets as $y^{(i)} \doteq w_*^\top x^{(i)}$.
For a new instance $x^{(n+1)}$ sampled from $d_\fX$, the goal is to predict the correct target $y^{(n+1)}$.

To demonstrate in-context learning,
one constructs a prompt matrix as

Z_0 \doteq \mqty[x^{(1)} & \dots & x^{(n)} & x^{(n+1)} \\ y^{(1)} & \dots & y^{(n)} & 0],
where the bottom right zero reflects that the target for $x^{(n+1)}$ is unknown.

The $L$-layer transformer is trained via gradient descent to minimize the following in-context loss

\begin{align}
% \fL (\qty{P_l, Q_l}_{l=0,1\,\dots,L-1})
% \doteq
% $
\label{eq ic regression loss}
% \textstyle \E_{Z_0, w_*}\qty[\qty(\tf_L(Z_0; \qty{P_l, Q_l}_{l=0,1,\dots,L-1}) + \indot{w_*}{x^{(n+1)}})^2],
\textstyle \E_{(d_\fX, w_*) \sim d_\text{task}, Z_0 \sim d_\fX}[(\tf_L(Z_0; \qty{P_l, Q_l}_{l=0,1,\dots,L-1}) - w_*^\top x^{(n+1)})^2],
% $
\end{align}

where we have assumed that there is a distribution $d_\text{task}$ over such regression tasks.

When a new regression task $(d_\fX^\text{test}, w_*^\text{test})$ is sampled from $d_\text{task}$ and a new input $Z_0^\text{test}$ is constructed, the trained transformer, using $Z_0^\text{test}$ as input, approximates the target $\indot{x^{(n+1), \text{test}}}{w_*^\text{test}}$.

This is a form of meta-learning [52].

Surprisingly, the transformer's ability to achieve this stems from its implementation of gradient descent within its forward pass.

As proved by [2],
by minimizing the in-context loss in (<ref>),
we may end up with a transformer parameterized by, say $\{(P_l^*, Q_l^*)\}_{l=0,\dots,L-1}$,
that has the following remarkable effect.
Feeding the prompt $Z_0$ into this $L$-layer transformer,

we get $Z_1, \dots, Z_L$ following (<ref>).

We denote the right bottom element of $Z_l$ as $y^{(n+1)}_l$.

[2] then prove that for $l=0,1,\dots, L$,

we have

y_l^{(n+1)} = - w_l^\top x^{(n+1)},

w_{l+1} \doteq  w_l + \frac{1}{n} \sum_{i=1}^n (y^{(i)} - w_l^\top x^{(i)}) x^{(i)}
with $w_0 = 0$.
This sequence $\qty{w_l}$ mirrors that produced by running gradient descent on the demonstrations $\{(x^{(i)}, y^{(i)})\}$ to minimize the squared loss $\frac{1}{n} \sum_{i=1}^n (y^{(i)} - w^\top x^{(i)})^2$.

In other words,

unrolling this transformer layer by layer is equivalent to performing gradient descent iteration by iteration.

Reinforcement Learning.
We consider an infinite horizon Markov Decision Process (MDP, 38) with a finite state space $\fS$,
a finite action space $\fA$,
a reward function $r_\mdp: \fS \times \fA \to \R$,
a transition function $p_\mdp: \fS \times \fS \times \fA \to [0, 1]$,
a discount factor $\gamma \in [0, 1)$,
and an initial distribution $p_0: \fS \to [0, 1]$.
An initial state $S_0$ is sampled from $p_0$.
At a time $t$,
an agent at a state $S_t$
takes an action $A_t \sim \pi(\cdot | S_t)$,
where $\pi: \fA \times \fS \to [0, 1]$ is the policy being followed by the agent,
receives a reward $R_{t+1} \doteq r_\mdp(S_t, A_t)$,
and transitions to a successor state $S_{t+1} \sim p_\mdp(\cdot | S_t, A_t)$.
If the policy $\pi$ is fixed, the MDP can be simplified to a Markov Reward Process (MRP) where transitions and rewards are determined solely by the current state:$S_{t+1} \sim p(\cdot | S_t)$ with $R_{t+1} \doteq r(S_t)$.
Here $p(s' | s) \doteq \sum_a \pi(a | s) p_\mdp(s'|s, a)$ and $r(s) \doteq \sum_a \pi(a|s) r_\mdp(s, a)$.
In this work,
we consider the policy evaluation problem where the policy $\pi$ is fixed.
So it suffices to consider only
an MRP represented by the tuple $(p_0, p, r)$,
and trajectories $(S_0, R_1, S_1, R_2, \dots)$ sampled from it.
The value function of this MRP is defined as

v(s) \doteq \E\left[\sum_{i = t+1}^\infty \gamma^{i - t -1} R_i | S_t = s\right]
Estimating the value function $v$ is one of the fundamental tasks in RL.

To this end, one can consider a linear architecture.
Let $\phi: \fS \to \R^d$ be the feature function.
The goal is then to find a weight vector $w \in \R^d$ such that for each $s$, the estimated value $\hat{v}(s; w) \doteq w^\top \phi(s)$

approximates $v(s)$.

TD is a prevalent method for learning this weight vector,
which updates $w$ iteratively as

\begin{align}
w_{t+1} =& w_t + \alpha_t \left(R_{t+1} + \gamma \hat{v}\left(S_{t+1};w_t\right) - \hat{v}\left(S_{t};w_t\right)\right) \nabla \hat{v}\left(S_{t};w_t\right) \\
=& w_t + \alpha_t \left(R_{t+1} + \gamma w_t^\top \phi(S_{t+1}) - w_t^\top \phi(S_t)\right) \phi(S_t), \label{eq: linear td(0)}
\end{align}

where $\qty{\alpha_t}$ is a sequence of learning rates.

Notably, TD is not a gradient descent algorithm.
It is instead considered as a semi-gradient algorithm because the gradient is only taken with respect to $\hat{v}\left(S_{t};w_t\right)$ and does not include the dependence on $\hat{v}\left(S_{t+1};w_t\right)$ [47].
Including this dependency modifies the update to

\begin{align}
w_{t+1} = w_t + \alpha_t \left(R_{t+1} + \gamma w_t^\top \phi(S_{t+1}) - w_t^\top \phi(S_t)\right) \left(\phi(S_t)- \gamma\phi(S_{t+1}) \right), \label{eq: residual grad td(0)}
\end{align}

known as the (naive version of) residual gradient method [7].[This is a naive version because the update does not account for the double sampling issue. We refer the reader to Chapter 11 of 47 for detailed discussion.]
The update in (<ref>) is also called TD(0) – a special case of the TD($\lambda$) algorithm [46].
TD($\lambda$) employs an eligibility trace that accumulates the gradients as

$e_{-1} \doteq 0, \ e_{t} \doteq \gamma\lambda e_{t-1} + \phi(S_t)$ and updates $w$ iteratively as

\begin{align}
w_{t+1} = w_t + \alpha_t (R_{t+1} + \gamma w_t^\top \phi(S_{t+1}) - w_t^\top \phi(S_t)) e_t.
\end{align}

The hyperparameter $\lambda$ controls the decay rate of the trace.
If $\lambda=0$, we recover (<ref>).

On the other end with $\lambda = 1$,
it is known that TD($\lambda$) recovers Monte Carlo [46].

Another important setting in RL is the average-reward setting [38, 47], focusing on the rate of receiving rewards, without using a discount factor $\gamma$.
The average reward $\bar{r}$ is defined as $\bar{r} \doteq \lim_{T \to \infty} \frac{1}{T} \sum_{t=1}^T \E[R_t]$.
Similar to the value function in the discounted setting,
a differential value function $\bar v(s)$ is defined for the average-reward setting as $\bar v(s) \doteq \E\left[\sum_{i=t+1}^\infty (R_i - \bar{r}) | S_t = s\right]$.
One can similarly estimate $\bar v(s)$ using a linear architecture with a vector $w$ as $w^\top \phi(s)$.
Average-reward TD [50] updates $w$ iteratively as

\begin{align}
w_{t+1} = w_t + \alpha_t \qty(R_{t+1} - \bar{r}_{t+1} + w_t^\top \phi(S_{t+1}) - w_t^\top \phi(S_t)) \phi(S_t),
\end{align}

where $\bar r_t \doteq \frac{1}{t}\sum_{i=1}^t R_i$ is the empirical average of the received reward.

§ TRANSFORMERS CAN IMPLEMENT IN-CONTEXT TD(0)

In this section,
we prove that transformers are expressive enough to implement TD(0) in its forward pass.
Given a trajectory $(S_0, R_1, S_1, R_2, S_3, R_4, \dots, S_n)$ sampled from an MRP,
using as shorthand $\phi_i \doteq \phi(S_i)$,
we define for $l = 0, 1, \dots, L-1$

\begin{align}
\!\!\!\!\!\! Z_0 = \mqty[\phi_0 & \dots & \phi_{n-1} & \phi_n \\ \gamma \phi_1 & \dots & \gamma \phi_n & 0 \\ R_1 & \dots & R_n & 0], \, P_l^{\td} \doteq& \mqty[0_{2d \times 2d} & 0_{2d \times 1} \\ 0_{1 \times 2d} & 1], \,    Q_l^{\td} \doteq \mqty[-C_l^\top & C_l^\top & 0_{d \times 1} \\ 0_{d \times d} & 0_{d \times d} & 0_{d \times 1} \\ 0_{1 \times d} & 0_{1 \times d} & 0].
\label{eq: Z_0 P Q TD define}
\end{align}

Here $Z_0 \in \R^{(2d+1) \times (n+1)}$ is the prompt matrix,
$C_l \in \R^{d \times d}$ is an arbitrary matrix,
and $\qty{(P_l^\td, Q_l^\td)}_{l=0,1,\dots, L-1}$ are the parameters of the $L$-layer transformer.
We then have

Consider the L-layer linear transformer following (<ref>), using the mask (<ref>), parameterized by $\qty{P_l^{\td}, Q_l^{\td}}_{l=0,\dots, L-1}$ in (<ref>).
Let $y_l^{(n+1)}$ be the bottom right element of the $l$-th layer's output, i.e.,  $y_l^{(n+1)} \doteq Z_l[2d+1, n+1]$.

Then, it holds that $y_l^{(n+1)} =  - \indot{\phi_n}{w_l}$,
where $\qty{w_l}$ is defined as $w_0 = 0$ and

\begin{align}
w_{l+1} &\textstyle = w_l + \frac{1}{n} C_l \sum_{j=0}^{n-1}  \left(R_{j+1} + \gamma w_l^\top\phi_{j+1} - w_l^\top\phi_j\right)\phi_j. \label{eq: TD(0) w_update}
\end{align}

The proof is in Appendix <ref> and with numerical verification in Appendix <ref> as a sanity check.
Notably, Theorem <ref> holds for any $C_l$.
In particular,
if $C_l = \alpha_l I$,
then the update (<ref>) becomes a batch version of TD(0) in (<ref>).
For a general $C_l$,
the update (<ref>) can be regarded as preconditioned batch TD(0) [55].
Theorem <ref> precisely demonstrates that transformers are expressive enough to implement iterations of TD in its forward pass.
We call this in-context TD.
It should be noted that although the construction of $Z_0$ in (<ref>) uses $\phi_n$ as the query state for conceptual clarity, any arbitrary state $s \in \mathcal{S}$ can serve as the query state and Theorem <ref> still holds.
In other words, by replacing $\phi_n$ with $\phi(s)$,
the transformer will then estimate $v(s)$.
if the transformer has only one layer, i.e., $L=1$,
there are other parameter configurations that can also implement in-context TD(0).

Consider the 1-layer linear transformer following (<ref>), using the mask (<ref>).
Consider the following parameters

\begin{align}
\label{eq pq one layer}
P_0^{\td} \doteq& \mqty[0_{2d \times 2d} & 0_{2d \times 1} \\ 0_{1 \times 2d} & 1], \,    Q_0^{\td} \doteq \mqty[-C_l^\top & 0_{d\times d} & 0_{d \times 1} \\ 0_{d \times d} & 0_{d \times d} & 0_{d \times 1} \\ 0_{1 \times d} & 0_{1 \times d} & 0]
\end{align}

Then, it holds that $y_1^{(n+1)} =  - \indot{\phi_n}{w_1}$,
where $w_1$ is defined as

\begin{align}
w_{1} &\textstyle = w_0 + \frac{1}{n} C_l \sum_{j=0}^{n-1}  \left(R_{j+1} + \gamma w_0^\top\phi_{j+1} - w_0^\top\phi_j\right)\phi_j \qq{with $w_0 = 0$.}
\end{align}

The proof is in Appendix <ref>. An observant reader may notice that this corollary holds primarily because $w_0 = 0$, making it a unique result for $L=1$.
this special case helps understand a few empirical and theoretical results below.

§ TRANSFORMERS DO IMPLEMENT IN-CONTEXT TD(0)

It has been observed that in-context gradient descent emerges during the minimization of the in-context regression loss (<ref>) via gradient descent.
In this section,
we demonstrate the emergence of in-context TD both theoretically and empirically.

Multi-Task Temporal Difference Learning.
The in-context regression loss essentially trains the transformer with multiple regression tasks.
Inspired by this,
we propose to train the transformer with multiple evaluation tasks from multiple MRPs.
Recall, an MRP is defined by the tuple $(p_0, p, r)$.
For the evaluation problem, the feature function $\phi$ also matters.
We therefore
define an evaluation task to be the tuple $(p_0, p, r, \phi)$.
Assuming a distribution $d_\text{task}$ over these tuples, we sample evaluation tasks from this distribution.
For each sampled task,
we apply TD to train the transformer to solve the corresponding evaluation problem,
as described in the following multi-task TD algorithm (Algorithm <ref>).

Multi-Task Temporal Difference Learning
[1]
context length $n$,
MRP sample length $\tau$,
number of training MRPs $k$,
learning rate $\alpha$,
discount factor $\gamma$,
transformer parameters $\theta \doteq \qty{P_l, Q_l}_{l=0,1,\dots L-1}$
$i \gets 1$ to $k$
Sample $(p_0, p, r, \phi)$ from $d_\text{task}$   // see, e.g., Algorithm <ref> in Appendix <ref>

Sample $(S_0, R_1, S_1, R_2, \dots, S_{\tau}, R_{\tau+1}, S_{\tau+1})$ from the MRP $(p_0, p, r)$
$t = 0, \dots, \tau- n - 1$

$Z_0 \gets \mqty[\phi_t & \cdots & \phi_{t+n-1} & \phi_{t+n+1} \\
\gamma \phi_{t+1} & \cdots & \gamma \phi_{t+n} & 0\\
R_{t+1} & \cdots & R_{t+n} & 0 ]$,
$Z_0' \gets \mqty[\phi_{t+1} & \cdots & \phi_{t+n} & \phi_{t+n+2} \\
\gamma \phi_{t+2} & \cdots & \gamma \phi_{t+n+1} & 0\\
R_{t+2} & \cdots & R_{t+n+1} & 0 ]$
$\theta \gets \theta + \alpha (R_{t+n+2} + \gamma \tf_L(Z_0'; \theta) - \tf_L(Z_0; \theta)) \nabla_\theta \tf_L(Z_0; \theta)$   // TD

Recall that $\tf_L(Z_0; \theta)$ and $\tf_L(Z_0'; \theta)$ are intended to estimate $v(S_{t+n+1})$ and $v(S_{t+n+2})$ respectively.
So Algorithm <ref> essentially applies TD using $(S_{t+n+1}, R_{t+n+2}, S_{t+n+2})$ to train the transformer.
when a new prompt $Z_\text{test}$ is constructed using a trajectory from a new evaluation task $(p_0, p, r, \phi)_\text{test} \sim d_\text{task}(\cdot)$,
we would like the predicted value $\tf_L(Z_\text{test}; \theta)$ with $\theta$ from Algorithm <ref> to be close to the value of the query state in $Z_\text{test}$.
This problem is a multi-task meta-learning problem, a well-explored area with many existing methodologies [8].
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.

Theoretical Analysis.
The problem that Algorithm <ref> aims to solve is highly non-convex and non-linear (the linear transformer is still a nonlinear function).
We analyze a simplified version of Algorithm <ref> and leave the treatment to the full version for future work.
In particular,
we study the single layer case with $L=1$

let $\theta \doteq (P_0, Q_0)$ be the parameters of the single-layer transformer.
We consider expected updates, i.e.,

\begin{align}
\!\!\!\!\!\!\!
\theta_{k+1} =& \theta_k + \alpha_k \Delta(\theta_k) \text{ with }
\Delta(\theta) \doteq \E\left[(R + \gamma \tf_1(Z_0', \theta) - \tf_1(Z_0, \theta)) \nabla \tf_1(Z_0, \theta)\right]. \label{eq: expected theta update}
\end{align}

Here the expectation integrates both the randomness in sampling $(p_0, p, r, \phi)$ from $d_\text{task}$ and the randomness in constructing $\qty(R, Z_0, Z_0')$ thereafter.

We sample $(S_0, R_1, S_1, \dots, S_{n+1}, R_{n+2}, S_{n+2})$ following $(p_0, p, r)$ and construct using shorthand $\phi_i \doteq \phi(S_i)$

\begin{align}
\label{eq: prompt for invariant set markovian}
\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! Z_0 \doteq \mqty[\phi_0 & \dots & \phi_{n-1} & \phi_{n+1} \\
\gamma \phi_1 & \dots & \gamma \phi_{n} & 0 \\
R_1 & \dots & R_n & 0],
Z_0' \doteq \mqty[\phi_1 & \dots & \phi_{n} & \phi_{n+2} \\
\gamma \phi_2 & \dots & \gamma \phi_{n+1} & 0 \\
R_2 & \dots & R_{n+1} & 0],
R \doteq R_{n+2}.
% \tag{Markov}
% \in \R^{(2d + 1) \times (n+1)},
\end{align}

The structure of $Z_0$ and $Z_0'$ is similar to those in Algorithm <ref>. The main difference is that we do not use the sliding window.
We recall that $(p_0, p, r, \phi)$ are random variables with joint distribution $d_\text{task}$. Here, $\phi$ is essentially a random matrix taking value in $\R^{d \times \ns}$, represented as, $\phi = [\phi(s)]_{s \in \fS}$.

We use $\triangleq$ to denote “equal in distribution" and make the following assumptions.

The random matrix $\phi$ is independent of $(p_0, p, r)$.

$\Pi \phi \triangleq \phi, \Lambda \phi \triangleq \phi$,
where $\Pi$ is any $d$-dimensional permutation matrix and $\Lambda$ is any diagonal matrix in $\R^d$ where each diagonal element of $\Lambda$ can only be $-1$ or $1$.

Those assumptions are easy to satisfy.
For example,
as long as the elements of the random matrix $\phi$ are i.i.d. from a symmetric distribution centered at zero, e.g., a uniform distribution on $[-1, 1]$,
then both assumptions hold.

We say a set $\Theta$ is an invariant set of (<ref>) if for any $k$, $\theta_k \in \Theta \implies \theta_{k+1} \in \Theta$.

\begin{align}
\theta_*(\eta, c, c') \doteq&  \textstyle \left(P_0 = \mqty[0_{2d \times 2d} & 0_{2d \times 1} \\ 0_{1 \times 2d} & \eta], Q_0 =  \mqty[c I_d & 0_{d \times d} & 0_{d \times 1} \\ c' I_d  & 0_{d \times d} & 0_{d \times 1} \\ 0_{1 \times d} & 0_{1 \times d} & 0] \right)  \label{eq: optimal single layer}.
% \Theta_* \doteq& \qty{\theta_*(\eta, c, c') | \eta, c, c' \in \R}, \, \Theta_*' \doteq \qty{\theta_*(\eta, c, 0) | \eta, c\in \R}.
\end{align}

Let Assumptions <ref> and <ref> hold.
For the (<ref>) construction of $(R, Z_0, Z_0')$,
then $\Theta_* \doteq \qty{\theta_*(\eta, c, c') | \eta, c, c' \in \R}$ is an invariant set of (<ref>).

The proof is in Appendix <ref>. Theorem <ref> demonstrates that once $\theta_k$ enters $\Theta_*$ at some $k$,
it can never leave,
$\Theta_*$ is a candidate set that the update (<ref>) can possibly converge to.
Consider a subset $\Theta_*' \subset \Theta_*$ with a stricter constraint $c' = 0$,
i.e., $\Theta_*' \doteq \qty{\theta_*(\eta, c, 0) | \eta, c \in \R}$.
Corollary <ref> then confirms that all parameters in $\Theta_*'$ implement in-context TD.

That being said,
whether (<ref>) is guaranteed to converge to $\Theta_*$,
or further to $\Theta_*'$,
is left for future work.

Empirical Analysis.
We now empirically study Algorithm <ref>.
To this end,
we construct $d_\text{task}$ based on Boyan's chain [11],
a canonical environment for diagnosing RL algorithms.

We keep the structure of Boyan's chain but randomly generate initial distributions $p_0$,
transition probabilities $p$, reward functions $r$,
and the feature function $\phi$.
Details of this random generation process are provided in Algorithm <ref> with Figure <ref> visualizing Boyan's chain,
in Appendix <ref>.

For the linear transformer specified in (<ref>), we first consider the autoregressive case following  [3, 53],
where all the transformer layers share the same parameters,
i.e., $P_l \equiv P_0$ and $Q_l \equiv Q_0$ for $l=0,1, \dots, L-1$. We consider a three layer transformer $(L=3)$.
Importantly, all elements of $P_0$ and $Q_0$ are equally trainable –
we did not force any element of $P_0$ and $Q_0$ to be 0.
We then run Algorithm <ref> with Boyan's chain based evaluation tasks (i.e., $d_\text{task}$) to train this autoregressive transformer.
The dimension of the feature is $d = 4$ (i.e., $\phi(s) \in \R^4$).
Other hyperparameters of Algorithm <ref> are specified in Appendix <ref>.

Figure <ref> visualizes the final
learned $P_0$ and $Q_0$ by Algorithm <ref> after 4000 MRPs (i.e., $k=4000$),
closely match our specifications $P^{\td}$ and $Q^{\td}$ in (<ref>) with $C_l = I_d$.
In Figure <ref>,
we visualize the element-wise learning progress of $P_0$ and $Q_0$.

We observe that the bottom right element of $P_0$ increases (the $P_0[-1, -1]$ curve)
while the average absolute value of all other elements remain close to zero (the “Avg Abs Others” curve),
closely aligning with $P^{\td}$ up to some scaling factor.

Furthermore, the trace of the upper left $d \times d$ block of $Q_0$ approaches $-d$ (the $\tr(Q_0[:d,:d])$ curve),
and the trace of the upper right block (excluding the last column) approaches $d$ (the $\tr(Q_0[:d, d:2d])$ curve).
Meanwhile, the average absolute value of all the other elements in $Q_0$ remain near zero,
aligning with $Q^\td$ using $C_l = I_d$ up to some scaling factor.

Learned $P_0$ and $Q_0$ after 4000 MRPs

Element-wise learning progress of $P_0$ and $Q_0$

Visualization of the learned transformers and the learning progress.
Both (a) and (b) are averaged across 30 seeds and the shaded region in (b) denotes the standard errors.
Since $P_0$ and $Q_0$ are in the same product in (<ref>),
the algorithm can rescale both or flip the sign of both, but still end up with exactly the same transformer.
to make sure the visualization are informative,
we rescale $P_0$ and $Q_0$ properly first before visualization.
See Appendix <ref> for details.

More empirical analysis is provided in the Appendix.
In particular,
besides showing the parameter-wise convergence in Figure <ref>,
we also use other metrics including value difference,
implicit weight similarity, and sensitivity similarity,
inspired by 53, 3,
to examine the learned transformer.
We also study normal transformers without parameter sharing (Appendix <ref>),
as well as different choices of hyperparameters in Algorithm <ref>.
Furthermore, we empirically investigate the original softmax-based transformers (Appendix <ref>).
The overall conclusion is the same – in-context TD emerges in the transformers learned by Algorithm <ref>.
Notably, Theorem <ref> and Corollary <ref> suggests that for $L=1$,
there are two distinct ways to implement in-context TD (i.e., (<ref>) v.s. (<ref>)).
Our empirical results in Appendix <ref> show that Algorithm <ref> ends up with (<ref>) in Corollary <ref> for $L=1$,
aligning well with Theorem <ref>.
For $L=2,3,4$,
Algorithm <ref> always ends up with (<ref>) in Theorem <ref>,
as shown in Figure <ref> in Appendix <ref>.
We also empirically observed that for in-context TD to emerge,
the task distribution $d_\text{task}$ has to be “difficult” enough.
For example,
if $(p_0, p)$ or $\phi$ are always fixed,
we did not observe the emergence of in-context TD.

§ TRANSFORMERS CAN IMPLEMENT MORE RL ALGORITHMS

In this section,
we prove that transformers are expressive enough
to implement three additional well-known RL algorithms in the forward pass.
We warm up with the (naive version of) residual gradient (RG).
We then move to the more difficult TD($\lambda$).
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 of RG is an easy extension of Theorem <ref>.
We define

\begin{align}
P_l^{\rg} = P_l^{\td}, Q_l^{\rg}
\doteq \mqty[-C_l^\top & C_l^\top & 0_{d \times 1} \\
C_l^\top  & -C_l^\top & 0_{d \times 1} \\
0_{1 \times d} & 0_{1 \times d} & 0] \in \R^{(2d+1) \times (2d+1)}. \label{eq: P and Q def RG}
\end{align}

Consider the L-layer linear transformer following (<ref>), using the mask (<ref>), parameterized by $\qty{P_l^{\rg}, Q_l^{\rg}}_{l=0,\dots, L-1}$ in (<ref>).

Define $y_l^{(n+1)} \doteq Z_l[2d+1, n+1]$.

Then, it holds that $y_l^{(n+1)} =  - \indot{\phi_n}{w_l}$,
where $\qty{w_l}$ is defined as $w_0 = 0$ and

\begin{align}
w_{l+1} &= \textstyle w_l + \frac{1}{n} C_l \sum_{j=0}^{n-1}  \left(R_{j+1} + \gamma w_l^\top\phi_{j+1} - w_l^\top\phi_j\right)(\phi_j - \gamma \phi_{j+1}). \label{eq: w_update RG}
\end{align}

The proof is in <ref> with numerical verification in Appendix <ref> as a sanity check.
if $C_l \doteq \alpha_l I_d$,
then (<ref>) can be regarded as a batch version of (<ref>).
For a general $C_l$,
it is then preconditioned batch RG.

Notably, Figure <ref> empirically  demonstrates that Algorithm <ref> eventually ends up with in-context TD instead of in-context RG.
This matches the conventional wisdom in the RL community that TD is usually superior to the naive RG (see, e.g., 57 and references therein).

Incorporating eligibility traces is an important extension of TD(0).
We now demonstrate that by using a different mask,
transformers are able to implement in-context TD($\lambda$).
We define

\begin{align}
M^\tdlambda \doteq
\mqty[1 & 0 & 0 & 0 & \cdots & 0 & 0\\
\lambda & 1 & 0 & 0 & \cdots & 0 & 0\\
% \lambda^2 & \lambda & 1 & 0 &\cdots & 0 & 0\\
% \lambda^3 & \lambda^2 & \lambda & 1 & \cdots & 0 & 0\\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots\\
\lambda^{n-1} & \lambda^{n-2} & \lambda^{n-3} & \lambda^{n-4} & \cdots & 1 & 0\\
0 & 0 & 0 & 0 & \cdots & 0 & 0] \in R^{(n+1)\times(n+1)}. \label{eq: td lambda mask}
\end{align}

Notably, if $\lambda = 0$, the above mask for TD($\lambda$) recovers the mask for TD(0) in (<ref>).

Consider the L-layer linear transformer parameterized by $\qty{P_l^{\td}, Q_l^{\td}}_{l=0,\dots, L-1}$ as specified in (<ref>)
with the input mask used in (<ref>) being
$M^\tdlambda$ in (<ref>).

Define $y_l^{(n+1)} \doteq Z_l[2d+1, n+1]$.

Then, it holds that $y_l^{(n+1)} =  - \indot{\phi_n}{w_l}$ where $\qty{w_l}$ is defined with $w_0 = 0, e_0 = 0, \, e_j = \lambda e_{j-1} + \phi_j$, and

\begin{align}
\textstyle w_{k+1} = w_k + \frac{1}{n}C_k\sum_{i=0}^{n-1} \qty(r_{i+1} + \gamma w_k^\top \phi_{i+1} - w_k^\top \phi_i )e_i.
\end{align}

The proof is in <ref> with numerical verification in Appendix <ref> as a sanity check.

Average-Reward TD.
We now demonstrate that transformers are expressive enough to implement in-context average-reward TD.
Different from TD(0),
average-reward TD exhibits additional challenges in that it updates two estimates (i.e., $w_t$ and $\bar r_t$) in parallel.
To account for this challenge,
we use two additional mechanisms beyond the naive single-head linear transformer.
we allow additional “memory” in the prompt and consider two-head linear transformers.

Given a trajectory $(S_0, R_1, S_1, R_2, S_3, R_4, \dots, S_n)$ sampled from an MRP,

we construct the prompt matrix $Z_0$ as

\begin{align}
Z_0 = \mqty[\phi_0 & \dots & \phi_{n-1} & \phi_n \\
\phi_1 & \dots & \phi_n & 0 \\
R_1 & \dots & R_n & 0 \\
0 & \dots & 0 & 0] \in \R^{(2d + 2) \times (n+1)}.
\end{align}

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 as

\begin{align}
P_l^{\bartd, (1)}\doteq & \mqty[0_{2d \times 2d} & 0_{2d \times 1} & 0_{2d \times 1}\\
0_{1 \times 2d} & 1 & 0\\
0_{1\times 2d} & 0 & 0],
P_l^{\bartd, (2)} \doteq \mqty[0_{2d \times 2d} & 0_{2d \times 1} & 0_{2d \times 1}\\
0_{1 \times 2d} & 0 & 0\\
0_{1\times 2d} & 0 & 1] \label{eq: two-head P matrices},\\
Q^{\bartd}_l \doteq & \mqty[-C_l^\top & C_l^\top & 0_{d \times 2} \\
0_{d \times d} & 0_{d \times d} & 0_{d \times 2}\\
0_{2 \times d} & 0_{2 \times d} & 0_{2\times2}],
W_l \doteq \mqty[0_{2d \times 2d} & 0_{2d \times 1} & 0_{2d \times (2d + 2)} & 0_{2d \times 1} \\
0_{1 \times 2d} & 1 & 0_{1 \times (2d+2)} & 1]\label{eq: two-head Q and W}, \\
M^{\bartd, (2)} \doteq & \mqty[I_n & 0_{n \times 1}\\
0_{1 \times n} & 0], \,
M^{\bartd, (1)} \doteq  \qty(I_{n+1} - U_{n+1} \text{diag}\qty(\mqty[1& \frac{1}{2}& \dots& \frac{1}{n+1}]))M^{\bartd, (2)}, \label{eq: two-head mask}
\end{align}

where $C_l \in \R^{d\times d}$ is again an arbitrary matrix,
$U_{n+1}$ is the $(n+1) \times (n+1)$ upper triangle matrix where all the nonzero elements are 1,
and $\text{diag}\qty(x)$ constructs a diagonal matrix with the diagonal entry being $x$.
Here, $\qty{P_l^{\bartd, (1)}, Q_l^{\bartd}}$ are the parameters of the first attention heads,
with the input mask being $M^{\bartd, (1)}$.
$\qty{P_l^{\bartd, (2)}, Q_l^{\bartd}}$ are the parameters of the second attention heads,
with the input mask being $M^{\bartd, (2)}$.
The two heads coincide on some parameters.
$W_l$ is the affine transformation that combines the embeddings from the two attention heads.
Define the two-head linear-attention as

\begin{align}
\twohead(Z; P, Q, M, P', Q', M', W) \doteq W \mqty[\attn(Z; P, Q, M) \\ \attn(Z; P', Q', M')].
\end{align}

The $L$-layer transformer we are interested in is then given by

\begin{align}
\label{eq: Z two-head update}
\textstyle Z_{l+1} \doteq Z_l + \frac{1}{n} \twohead(Z_l; P_l^{\bartd, (1)}, Q_l^{\bartd}, M^{\bartd, (1)}, P_l^{\bartd, (2)}, Q_l^{\bartd}, M^{\bartd, (2)}, W_l).
\end{align}

Consider the $L$-layer transformer in (<ref>).

Let $h_l^{(n+1)}$ be the bottom-right element of the $l$-th layer output, i.e.,  $h_l^{(n+1)} \doteq Z_l[2d+2, n+1]$.
Then, it holds that $h_l^{(n+1)} =  -\indot{\phi_n}{w_l}$ where $\qty{w_l}$ is defined as $w_0 = 0$,

\begin{align}
\textstyle{w_{l+1} = w_l + \frac{1}{n} C_l \sum_{j=1}^n  \qty(R_j - \bar{r}_j + w_l^\top \phi_j - w_l^\top \phi_{j-1})\phi_{j-1}}
\end{align}

for $l = 0, \dots, L-1$, where $\bar{r}_j \doteq \frac{1}{j}\sum_{k=1}^j R_k$.

The proof is in <ref> with numerical verification in Appendix <ref> as a sanity check.

§ RELATED WORKS

In-Context Learning.
Understanding in-context learning empirically and theoretically has recently emerged as an active research area [18, 3, 53, 58, 4, 56, 35, 2], building on prior research demonstrating that neural networks are able to implement algorithms.

[44, 20, 24].
This work advances this line of research by demonstrating how transformers implement in-context TD,
accompanied by a theoretical understanding of its emergence.

In-Context Reinforcement Learning.

Existing research on in-context RL predominantly adopts a policy-based approach, often relying on supervised pre-training
[28, 40, 45, 59, 27].
Transformers are trained to output the action, instead of the value, for the query state.
the prompts used in this setup consist of previous trajectories from an MDP

\begin{align}
\underbrace{S_0 A_0 R_1 S_1 A_2 R_2 \dots S_{t-1} A_{t-1}}_{\text{context}} \underbrace{S_t}_{\text{query}} \to \underbrace{A_t}_{\text{output}}.
\end{align}

The dataset usually consists of multiple such prompt-query-output pairs, where maximum likelihood estimation is essentially used to train the transformers.
the prompt can be generated by following multiple policies.
The prompt can also be offline data containing all trajectories generated during prior RL algorithm training across multiple episodes.
This line of research is closely related to offline policy distillation,
the goal of which is to learn a policy from offline data using transformers [14, 23, 31, 41].

Despite that empirical successes observed in the work above,
theoretical analysis is often missing.
[33] provide theoretical analysis for this policy-based supervised pre-training approach and show that the transformers can approximate a few RL algorithms,
including LinUCB [15] and Thompson sampling [42] for linear bandits [29] and UCB-VI [6] for MDPs.
Specifically, [33] prove the inference process of the learned transformers behaves similarly to those aforementioned RL algorithms in terms of action selection probabilities, regret, and other metrics.
This behavioral similarity is also investigated in 30.
However, the underlying mechanisms within the learned transformers that induce this similarity remains unclear.
In contrast,
we go beyond behavioral similarity and
prove that transformers can exactly implement a few RL algorithms in its forward pass.
we do not use the supervised pre-training paradigm, which is centered on maximum likelihood estimation.
As shown in Algorithm <ref>,
we instead use RL pre-training predicated on TD,
a value-based method.
37 concurrently use a regret-based loss for training transformers in online learning.
12 implement policy iteration, a value-based strategy, with transformers, but perform the required $\arg\max$ operation outside the transformers.
Despite the observed empirical success, 12
also lack a theoretical analysis of their approach.

Meta-Learning of RL algorithms.
Our Algorithm <ref> can be regarded as a meta RL algorithm [8],
where $d_\text{task}$ is the task distribution in the meta RL framework.
The learned transformers can be regarded as a learned algorithm,
which is used to solve new evaluation tasks from the task distribution.
Such meta learning of RL algorithms has been explored in [16, 54, 17, 26, 36, 34].

However, those discovered algorithms lack interpretability – it is not clear how the neural network implements the discovered algorithms.
By contrast,
the discovered transformer from Algorithm <ref> is well explained.

§ CONCLUSION

This work demonstrates that transformers can and do learn to implement temporal difference methods for in-context policy evaluation in the forward pass.
We further provide a theoretical explanation of how in-context TD emerges by characterizing an invariant set of the multi-task TD algorithm used in pre-training,
bridging the gap between “can” and “do”.
there are a few limitations.
First, this work is focused on policy evaluation, with control algorithms deferred to future research.
Second, the analysis is largely theoretical – we leave the large-scale verification of the multi-task TD pre-training paradigm for future work.
Third, the theoretical analysis of the pre-training paradigm is confined to single-layer linear transformers, leaving the exploration of multi-layer softmax transformers for future studies. In conclusion, this research aims to illuminate the mechanisms of in-context learning, and motivate further investigation into in-context value-based RL.

§ ACKNOWLEDGEMENTS

This work is supported in part by the US National Science Foundation (NSF) under grants III-2128019 and SLES-2331904. 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.

[1]
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L.,
Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. (2023).
Gpt-4 technical report.
arXiv preprint arXiv:2303.08774.

[2]
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. (2024).
Transformers learn to implement preconditioned gradient descent for
in-context learning.
Advances in Neural Information Processing Systems, 36.

[3]
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. (2023).
What learning algorithm is in-context learning? investigations with
linear models.
The Eleventh International Conference on Learning

[4]
Allen-Zhu, Z. and Li, Y. (2023).
Physics of language models: Part 1, context-free grammar.
arXiv preprint arXiv:2305.13673.

[5]
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 29th ACM International Conference on Architectural Support
for Programming Languages and Operating Systems, Volume 2 (ASPLOS '24). ACM.

[6]
Azar, M. G., Osband, I., and Munos, R. (2017).
Minimax regret bounds for reinforcement learning.
In International conference on machine learning, pages
263–272. PMLR.

[7]
Baird, L. C. (1995).
Residual algorithms: Reinforcement learning with function
In Proceedings of the International Conference on Machine

[8]
Beck, J., Vuorio, R., Liu, E. Z., Xiong, Z., Zintgraf, L., Finn, C., and
Whiteson, S. (2023).
A survey of meta-reinforcement learning.
arXiv preprint arXiv:2301.08028.

[9]
Bishop, C. M. (2006).
Pattern recognition and machine learning.
Springer google schola, 2:1122–1128.

[10]
Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L., Golding, L., He,
H., Leahy, C., McDonell, K., Phang, J., et al. (2022).
Gpt-neox-20b: An open-source autoregressive language model.
arXiv preprint arXiv:2204.06745.

[11]
Boyan, J. A. (1999).
Least-squares temporal difference learning.
In Proceedings of the International Conference on Machine

[12]
Brooks, E., Walls, L., Lewis, R. L., and Singh, S. (2024).
Large language models can implement policy iteration.
Advances in Neural Information Processing Systems, 36.

[13]
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P.,
Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S.,
Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler,
D., Wu, J., 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 Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H.,
editors, Advances in Neural Information Processing Systems, volume 33,
pages 1877–1901. Curran Associates, Inc.

[14]
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P.,
Srinivas, A., and Mordatch, I. (2021).
Decision transformer: Reinforcement learning via sequence modeling.
Advances in neural information processing systems,

[15]
Chu, W., Li, L., Reyzin, L., and Schapire, R. (2011).
Contextual bandits with linear payoff functions.
In Proceedings of the Fourteenth International Conference on
Artificial Intelligence and Statistics, pages 208–214. JMLR Workshop and
Conference Proceedings.

[16]
Duan, Y., Schulman, J., Chen, X., Bartlett, P. L., Sutskever, I., and Abbeel,
P. (2016).
Rl2: nforcement learning via slow reinforcement learning.
arXiv preprint arXiv:1611.02779.

[17]
Finn, C., Abbeel, P., and Levine, S. (2017).
Model-agnostic meta-learning for fast adaptation of deep networks.
In International conference on machine learning, pages
1126–1135. PMLR.

[18]
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. (2022).
What can transformers learn in-context? a case study of simple
function classes.
Advances in Neural Information Processing Systems,

[19]
Garrett, J. D. (2021).

[20]
Graves, A., Wayne, G., and Danihelka, I. (2014).
Neural turing machines.
arXiv preprint arXiv:1410.5401.

[21]
Harris, C. R., Millman, K. J., van der Walt, S. J., 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, J. 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, 585(7825):357–362.

[22]
Hunter, J. D. (2007).
Matplotlib: A 2d graphics environment.
Computing in Science & Engineering, 9(3):90–95.

[23]
Janner, M., Li, Q., and Levine, S. (2021).
Offline reinforcement learning as one big sequence modeling problem.
Advances in neural information processing systems,

[24]
Jastrzębski, S., Arpit, D., Ballas, N., Verma, V., Che, T., and Bengio,
Y. (2017).
Residual connections encourage iterative inference.
arXiv preprint arXiv:1710.04773.

[25]
Kingma, D. P. and Ba, J. (2015).
Adam: A method for stochastic optimization.
In Proceedings of the International Conference on Learning

[26]
Kirsch, L., van Steenkiste, S., and Schmidhuber, J. (2019).
Improving generalization in meta reinforcement learning using learned
arXiv preprint arXiv:1910.04098.

[27]
Krishnamurthy, A., Harris, K., Foster, D. J., Zhang, C., and Slivkins, A.
Can large language models explore in-context?
arXiv preprint arXiv:2403.15371.

[28]
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 arXiv:2210.14215.

[29]
Lattimore, T. and Szepesvári, C. (2020).
Bandit algorithms.
Cambridge University Press.

[30]
Lee, J., Xie, A., Pacchiano, A., Chandak, Y., Finn, C., Nachum, O., and
Brunskill, E. (2024).
Supervised pretraining can learn in-context reinforcement learning.
Advances in Neural Information Processing Systems, 36.

[31]
Lee, K.-H., Nachum, O., Yang, M. S., Lee, L., Freeman, D., Guadarrama, S.,
Fischer, I., Xu, W., Jang, E., Michalewski, H., et al. (2022).
Multi-game decision transformers.
Advances in Neural Information Processing Systems,

[32]
Lieber, O., Sharir, O., Lenz, B., and Shoham, Y. (2021).
Jurassic-1: Technical details and evaluation.
White Paper. AI21 Labs, 1:9.

[33]
Lin, L., Bai, Y., and Mei, S. (2023).
Transformers as decision makers: Provable in-context reinforcement
learning via supervised pretraining.
arXiv preprint arXiv:2310.08566.

[34]
Lu, C., Kuba, J., Letcher, A., Metz, L., Schroeder de Witt, C., and Foerster,
J. (2022).
Discovered policy optimisation.
Advances in Neural Information Processing Systems,

[35]
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 arXiv:2307.03576.

[36]
Oh, J., Hessel, M., Czarnecki, W. M., Xu, Z., van Hasselt, H. P., Singh, S.,
and Silver, D. (2020).
Discovering reinforcement learning algorithms.
Advances in Neural Information Processing Systems,

[37]
Park, C., Liu, X., Ozdaglar, A., and Zhang, K. (2024).
Do llm agents have regret? a case study in online learning and games.
arXiv preprint arXiv:2403.16843.

[38]
Puterman, M. L. (2014).
Markov decision processes: discrete stochastic dynamic
John Wiley & Sons.

[39]
Rae, J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann, J., Song, F.,
Aslanides, J., Henderson, S., Ring, R., Young, S., et al. (2021).
Scaling language models: Methods, analysis & insights from training
arXiv preprint arXiv:2112.11446.

[40]
Raparthy, S. C., Hambro, E., Kirk, R., Henaff, M., and Raileanu, R. (2023).
Generalization to new sequential decision making tasks with
in-context learning.
arXiv preprint arXiv:2312.03801.

[41]
Reed, S., Zolna, K., Parisotto, E., Colmenarejo, S. G., Novikov, A.,
Barth-Maron, G., Gimenez, M., Sulsky, Y., Kay, J., Springenberg, J. T.,
et al. (2022).
A generalist agent.
arXiv preprint arXiv:2205.06175.

[42]
Russo, D. J., Van Roy, B., Kazerouni, A., Osband, I., Wen, Z., et al. (2018).
A tutorial on thompson sampling.
Foundations and Trends® in Machine Learning,

[43]
Schlag, I., Irie, K., and Schmidhuber, J. (2021).
Linear transformers are secretly fast weight programmers.
In International Conference on Machine Learning, pages
9355–9366. PMLR.

[44]
Siegelmann, H. T. and Sontag, E. D. (1992).
On the computational power of neural nets.
In Proceedings of the fifth annual workshop on Computational
learning theory, pages 440–449.

[45]
Sinii, V., Nikulin, A., Kurenkov, V., Zisman, I., and Kolesnikov, S. (2023).
In-context reinforcement learning for variable action spaces.
arXiv preprint arXiv:2312.13327.

[46]
Sutton, R. S. (1988).
Learning to predict by the methods of temporal differences.
Machine Learning.

[47]
Sutton, R. S. and Barto, A. G. (2018).
Reinforcement Learning: An Introduction (2nd Edition).
MIT press.

[48]
Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, J., Soricut, R.,
Schalkwyk, J., Dai, A. M., Hauth, A., et al. (2023).
Gemini: a family of highly capable multimodal models.
arXiv preprint arXiv:2312.11805.

[49]
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y.,
Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. (2023).
Llama 2: Open foundation and fine-tuned chat models.
arXiv preprint arXiv:2307.09288.

[50]
Tsitsiklis, J. N. and Roy, B. V. (1999).
Average cost temporal-difference learning.

[51]
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N.,
Kaiser, L. u., and Polosukhin, I. (2017).
Attention is all you need.
In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R.,
Vishwanathan, S., and Garnett, R., editors, Advances in Neural
Information Processing Systems, volume 30. Curran Associates, Inc.

[52]
Vilalta, R. and Drissi, Y. (2002).
A perspective view and survey of meta-learning.
Artificial intelligence review, 18:77–95.

[53]
von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A.,
Zhmoginov, A., and Vladymyrov, M. (2023).
Transformers learn in-context by gradient descent.

[54]
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 arXiv:1611.05763.

[55]
Yao, H. and Liu, Z.-Q. (2008).
Preconditioned temporal difference learning.
In Proceedings of the 25th international conference on Machine
learning, pages 1208–1215.

[56]
Zhang, R., Frei, S., and Bartlett, P. L. (2023).
Trained transformers learn linear models in-context.
arXiv preprint arXiv:2306.09927.

[57]
Zhang, S., Boehmer, W., and Whiteson, S. (2020).
Deep residual reinforcement learning.
In Proceedings of the International Conference on Autonomous
Agents and Multiagent Systems.

[58]
Zhao, H., Panigrahi, A., Ge, R., and Arora, S. (2023).
Do transformers parse while predicting the masked word?
arXiv preprint arXiv:2303.08117.

[59]
Zisman, I., Kurenkov, V., Nikulin, A., Sinii, V., and Kolesnikov, S. (2023).
Emergence of in-context reinforcement learning from noise
arXiv preprint arXiv:2312.12275.

§ PROOFS

§.§ Proof of Theorem <ref>

We recall from (<ref>) that the embedding evolves according to

\begin{align}
Z_{l+1} = Z_l + \frac{1}{n} P_l Z_l M(Z_l^\top Q_l Z_l).
\end{align}

We first express $Z_l$ using elements of $Z_0$.
To this end,
it is convenient to give elements of $Z_l$ different names,
in particular,
we refer to the elements in $Z_l$ as $\qty{(x^{(i)}_l, y^{(i)}_l)}_{i=1,\dots, n+1}$ in the following way

\begin{align}
Z_l = \mqty[x_l^{(1)} & \dots & x_l^{(n)} & x_l^{(n+1)} \\ y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)}],
\end{align}

where we recall that $Z_l \in \R^{(2d + 1) \times (n+1)}, x_l^{(i)} \in \R^{2d}, y^{(i)}_l \in \R$.
Sometimes it is more convenient to refer to the first half and second half of $x_l^{(i)}$ separately,
by, e.g., $\nu^{(i)}_l \in \R^d, \xi_l^{(i)} \in \R^d$, i.e., $x_l^{(i)} = \mqty[\nu_l^{(i)} \\ \xi_l^{(i)}]$.
Then we have

\begin{align}
Z_l = \mqty[\nu_l^{(1)} & \dots & \nu_l^{(n)} & \nu_l^{(n+1)} \\
\xi_l^{(1)} & \dots & \xi_l^{(n)} & \xi_l^{(n+1)} \\
y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)}].
\end{align}

We utilize the shorthands

\begin{align}
X_l =& \mqty[x_l^{(1)} & \dots & x_l^{(n)}] \in \R^{2d \times n}, \\
Y_l =&  \mqty[y_l^{(1)} & \dots & y_l^{(n)}] \in \R^{1 \times n}.
\end{align}

Then we have

\begin{align}
Z_l = \mqty[X_l & x_l^{(n+1)} \\ Y_l & y_l^{(n+1)}]. \label{eq: z_l blockwise X Y}
\end{align}

For the input $Z_0$,
we assume $\xi_0^{(n+1)} = 0, y_0^{(n+1)} = 0$ but all other entries of $Z_0$ are arbitrary.
We recall our definition of $M$ in (<ref>) and $\qty{P_l^\td, Q_l^\td}_{l=0,\dots,L-1}$ in  (<ref>).
In particular,
we can express $Q_l^{\td}$ in a more compact way as

\begin{align}
M_1 \doteq& \mqty[-I_{d} & I_{d} \\0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d}, \label{eq:M_1 define}\\
B_l \doteq& \mqty[C_l^\top & 0_{d\times d} \\ 0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
A_l \doteq& B_l M_1  = \mqty[-C_l^\top & C_l^\top \\ 0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
%   A_l^\top =& M_1^\top B_l^\top \in \R^{2d \times 2d},\\
Q_l^{\td} \doteq& \mqty[A_l & 0_{2d \times 1} \\ 0_{1 \times 2d} & 0] \in \R^{(2d + 1) \times (2d+1)}.
\end{align}

We now proceed with the following claims.

Claim 1. $X_l \equiv X_0 , x_l^{(n+1)} \equiv x_0^{(n+1)}, \forall l.$

Recall that $ P_l^{\td} \doteq \mqty[0_{2d \times 2d} & 0_{2d \times 1} \\ 0_{1 \times 2d} & 1]
\in \R^{(2d+1) \times (2d+1)}$.

\begin{align}
W_l \doteq Z_l M \qty(Z_l^\top Q_l^{\td} Z_l) \in \R^{(2d+1) \times (n+1)}.
\end{align}

The embedding evolution can then be expressed as

\begin{align}
Z_{l+1} = Z_l + \frac{1}{n} P_l^{\td} W_l.
\end{align}

By simple matrix arithmetic, we get

\begin{align}
P_l^{\td} W_l = \mqty[0_{2d \times (n+1)}\\
\end{align}

where $W_l(2d+1)$ denotes the $(2d + 1)$-th row of $W_l$.
Therefore, we have $X_{l+1} = X_l, x_{l+1}^{(n+1)} = x_l^{(n+1)}$.
By induction, we get $X_l \equiv X_0$ and $x_l^{(n+1)} \equiv x_0^{(n+1)}$ for all $l = [0,\dots,L-1]$.

In light of this,
we drop all the subscripts of $X_l$,
as well as subscripts of $x_l^{(i)}$ for $i=1,\dots, n+1$.

Claim 2.

\begin{align}
Y_{l+1} &= Y_l + \frac{1}{n} Y_l X^\top A_l X\\
y_{l+1}^{(n+1)} &= y_l^{(n+1)} + \frac{1}{n} Y_l X^\top A_l x^{(n+1)}.
\end{align}

The easier way to show why this claim holds is to factor the embedding evolution into the product of $P_l^{\td} Z_l M$ and $Z_l^\top Q_l^{\td} Z_l$.
Firstly, we have

\begin{align}
P_l^{\td} Z_l = \mqty[0_{2d \times n} & 0_{2d \times 1}\\
Y_l & y_l^{(n+1)}].
\end{align}

Applying the mask, we get

\begin{align}
P_l^{\td} Z_l M = \mqty[0_{2d \times n} & 0_{2d \times 1}\\
Y_l & 0].
\end{align}

Then, we analyze $Z_l^\top Q_l^{\td} Z_l$.
Applying the block matrix notations, we get

\begin{align}
Z_l^\top Q_l^{\td} Z_l &= \mqty[X^\top & Y_l^\top\\
x^{(n+1)^\top} & y_l^{(n+1)}]
\mqty[A_l & 0_{2d \times 1} \\
0_{1 \times 2d} & 0]
\mqty[X & x^{(n+1)} \\
Y_l & y_l^{(n+1)}]\\
&= \mqty[X^\top A_l & 0_{n \times 1}\\
x^{(n+1)^\top} A_l & 0]
\mqty[X & x^{(n+1)} \\
Y_l & y_l^{(n+1)}]\\
&= \mqty[X^\top A_l X & X^\top A_l x^{(n+1)}\\
x^{(n+1)^\top} A_l X & x^{(n+1)^\top} A_l x^{(n+1)}].
\end{align}

Combining the two, we get

\begin{align}
P_l^{\td} Z_l M \qty(Z_l^\top Q_l^{\td} Z_l)
&= \mqty[0_{2d \times n} & 0_{2d \times 1}\\
Y_l & 0]
\mqty[X^\top A_l X & X^\top A_l x^{(n+1)}\\
x^{(n+1)^\top} A_l X & x^{(n+1)^\top} A_l x^{(n+1)}]\\
&= \mqty[0_{2d \times n} & 0_{2d \times 1}\\
Y_l X^\top A_l X & Y_l X^\top A_l x^{(n+1)}].
\end{align}

Hence, according to our update rule in (<ref>), we get

\begin{align}
Y_{l+1} &= Y_l + \frac{1}{n} Y_l X^\top A_l X\\
y_{l+1}^{(n+1)} &= y_l^{(n+1)} + \frac{1}{n} Y_l X^\top A_l x^{(n+1)}.
\end{align}

Claim 3.

\begin{align}
y_{l+1}^{(i)} = y_0^{(i)} + \indot{M_1 x^{(i)}}{\frac{1}{n}\sum_{j=0}^l B_j^\top M_2 X Y_j^\top},
\end{align}

for $i = 1, \dots, n+1$, where $M_2 = \mqty[I_{d} & 0_{d\times d}\\ 0_{d\times d} & 0_{d\times d}]$.

Following Claim 2, we can unroll $Y_{l+1}$ as

\begin{align}
Y_{l+1} &= Y_l + \frac{1}{n} Y_l X^\top A_l X\\
Y_l &= Y_{l-1} + \frac{1}{n} Y_{l-1} X^\top A_{l-1} X\\
Y_1 & = Y_0 + \frac{1}{n} Y_0 X^\top A_0 X.
\end{align}

We can then compactly express $Y_{l+1}$ as

\begin{align}
Y_{l+1} = Y_0 + \frac{1}{n}\sum_{j=0}^l Y_j X^\top A_j X.
\end{align}

Recall that we define $A_j = B_j M_1$.
Then, we can rewrite $Y_{l+1}$ as

\begin{align}
&= Y_0 + \frac{1}{n} \sum_{j=0}^l Y_j X^\top M_2 B_j M_1 X.
\end{align}

The introduction of $M_2$ here does not break the equivalence because $B_j = M_2 B_j$.
However, it will help make our proof steps easier to comprehend later.

With the identical procedure, we can easily rewrite $y_{l+1}^{(n+1)}$ as

\begin{align}
y_{l+1}^{(n+1)} &= y_0^{(n+1)} + \frac{1}{n} \sum_{j=0}^l Y_j X^\top M_2 B_j M_1 x^{(n+1)}.
\end{align}

In light of this,
we define $\psi_0 \doteq 0$ and for $l = 0, \dots$

\begin{align}
\psi_{l+1} \doteq& \frac{1}{n}\sum_{j=0}^l B_j^\top M_2 X Y_j^\top \in \R^{2d}. \label{eq: theta_l define}
\end{align}

Then we can write

\begin{align}
y_{l+1}^{(i)} = y_0^{(i)} + \indot{M_1 x^{(i)}}{\psi_{l+1}}, \label{eq: y_l+1 evol theta}
\end{align}

for $i=1, \dots, n+1$, which is the claim we made.
In particular, since we assume $y_0^{(n+1)} = 0$, we have

\begin{align}
y_{l+1}^{(n+1)} = \indot{M_1 x^{(n+1)}}{\psi_{l+1}}.
\end{align}

Claim 4. The bottom $d$ elements of $\psi_l$ are always 0, i.e., there exists a sequence $\qty{w_l \in \R^d}$ such that we can express $\psi_l$ as

\begin{align}
\psi_l = \mqty[w_l \\ 0_{d \times 1}]. \label{eq: theta_l sparsity}
\end{align}

for all $l = 0,1,\dots, L$.

We prove the claim by induction. The base case holds trivially since $\psi_0 \doteq 0$. Suppose that for some $l$, (<ref>) holds. It can be easily verified from the definition of $\psi_{l+1}$ in (<ref>) that

\begin{align}
\psi_{l+1} = \psi_{l} + \frac{1}{n} B_l^\top M_2 X Y_l^\top. \label{eq: theta evolve recursive}
\end{align}

If we let

\begin{align}
N_l = \frac{1}{n}M_2XY_l^{\top} \in \R^{2d \times 1},
\end{align}

the evolution of $\psi_{l+1}$ can then be compactly expressed as,

\begin{align}
\psi_{l+1} = \psi_l + B_l^\top N_l.
\end{align}

By matrix arithmetic, we have

\begin{align}
B_l^\top N_l &= \mqty[C_l^\top & 0_{d\times d} \\ 0_{d \times d} & 0_{d \times d}]^\top \mqty[N_l(1:d) \\ N_l(d:2d)] \\
&= \mqty[C_l N_l(1:d) \\ 0_{d\times 1}]
\end{align}

where $N_l(1:d)\in \R^d$ and $N_l(d:2d)\in \R^d$ represent the first $d$ and second $d$ elements of $N_l$ respectively. Substituting in our inductive hypothesis into (<ref>), we have:

\begin{align}
\psi_{l+1} &=  \mqty[w_l \\ 0_{d\times 1}] + \mqty[C_l N_l(1:d) \\ 0_{d\times 1}], \\
& = \mqty[w_l + C_l N_l(1:d) \\ 0_{d\times 1}]
\end{align}

if we let $w_{l+1} = w_l + C_l N_l(1:d)$, we can see that the property holds for $\psi_{l+1}$,
thereby verifying Claim 4.

Given all the claims above,
we can then compute that

\begin{align}
&\indot{\psi_{l+1}}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} + \frac{1}{n}\indot{B_l^\top M_2 X Y_l^\top}{M_1 x^{(n+1)}} \explain{By~\eqref{eq: theta evolve recursive}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} + \frac{1}{n}\sum_{i=1}^n \indot{B_l^\top M_2 x^{(i)} y_l^{(i)}}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n \indot{B_l^\top M_2 x^{(i)} \qty(\indot{\psi_l}{M_1 x^{(i)}} + y_0^{(i)})}{M_1 x^{(n+1)}} \explain{By \eqref{eq: y_l+1 evol theta}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{B_l^\top \mqty[\nu^{(i)}\\0_{d\times 1}]
\qty( \indot{\psi_l}{\mqty[-\nu^{(i)} + \xi^{(i)}\\ 0_{d\times 1}]} + y_0^{(i)})}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_l \nu^{(i)}\\0_{d\times 1}]
\qty( y_0^{(i)} +  w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)} )}{M_1 x^{(n+1)}} \explain{By Claim 4}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_l \nu^{(i)} \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})\\0_{d\times 1}]}
{M_1 x^{(n+1)}}\\
\end{align}

This means

\begin{align}
\indot{w_{l+1}}{\nu^{(n+1)}}
= \indot{w_l}{\nu^{(n+1)}}
+ \frac{1}{n} \sum_{i=1}^n
\indot{C_l \nu^{(i)} \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})}{
\nu^{(n+1)}}.
\end{align}

Since the choice of the query $\nu^{(n+1)}$ is arbitrary, we get

\begin{align}
w_{l+1} = w_l + \frac{1}{n}\sum_{i=1}^n C_l \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)}) \nu^{(i)}.
\end{align}

In particular, when we construct $Z_0$ such that $\nu^{(i)} = \phi_{i-1}$, $\xi^{(i)} = \gamma \phi_i$ and $y_0^{(i)} = R_i$,
we get

\begin{align}
w_{l+1} = w_l + \frac{1}{n}\sum_{i=1}^n C_l \qty(R_i + \gamma w_l^\top \phi_i - w_l^\top \phi_{i-1}) \phi_{i-1}
\end{align}

which is the update rule for pre-conditioned TD learning.

We also have

\begin{align}
= \indot{\psi_l}{M_1 x^{(n+1)}}
= -\indot{w_l}{\phi^{(n+1)}}.
\end{align}

This concludes our proof.

§.§ Proof of Corollary <ref>

The proof presented here closely mirrors the methodology and notation established in Theorem <ref>. Since we are only considering a 1-layer transformer in this Corollary, we can recall the embedding evolution from  (<ref>) and write

\begin{align}
Z_{1} = Z_0 + \frac{1}{n} P_0 Z_0 M(Z_0^\top Q_0 Z_0).
\end{align}

We once again refer to the elements in $Z_l$ as $\qty{(x^{(i)}_l, y^{(i)}_l)}_{i=1,\dots, n+1}$ in the following way

\begin{align}
Z_l = \mqty[x_l^{(1)} & \dots & x_l^{(n)} & x_l^{(n+1)} \\ y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)}],
\end{align}

where we recall that $Z_l \in \R^{(2d + 1) \times (n+1)}, x_l^{(i)} \in \R^{2d}, y^{(i)}_l \in \R$.
We utilize, $\nu^{(i)}_l \in \R^d, \xi_l^{(i)} \in \R^d$, to refer to the first half and second half of $x_l^{(i)}$ i.e., $x_l^{(i)} = \mqty[\nu_l^{(i)} \\ \xi_l^{(i)}]$.
Then we have

\begin{align}
Z_l = \mqty[\nu_l^{(1)} & \dots & \nu_l^{(n)} & \nu_l^{(n+1)} \\
\xi_l^{(1)} & \dots & \xi_l^{(n)} & \xi_l^{(n+1)} \\
y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)}].
\end{align}

We further define as shorthands

\begin{align}
X_l =& \mqty[x_l^{(1)} & \dots & x_l^{(n)}] \in \R^{2d \times n}, \ Y_l =  \mqty[y_l^{(1)} & \dots & y_l^{(n)}] \in \R^{1 \times n}.
\end{align}

Then the blockwise structure of $Z_l$ can be succinctly expressed as:

\begin{align}
Z_l = \mqty[X_l & x_l^{(n+1)} \\ Y_l & y_l^{(n+1)}]. \label{eq: z_l blockwise X Y L1}
\end{align}

For the input $Z_0$,
we assume $\xi_0^{(n+1)} = 0, y_0^{(n+1)} = 0$ but all other entries of $Z_0$ are arbitrary.
We recall our definition of $M$ in (<ref>) and $\qty{P_0, Q_0}$ in (<ref>).
In particular,
we can express $Q_0$ in a more compact way as

\begin{align}
M_1 \doteq& \mqty[-I_{d} &  0_{d \times d} \\0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d}, \label{eq: L1 M_1 define} \
B_0 \doteq \mqty[C_0^\top & 0_{d\times d} \\ 0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
A_0 \doteq& B_0M_1  = \mqty[-C_0^\top &  0_{d \times d}\\ 0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
%   A_l^\top =& M_1^\top B_l^\top \in \R^{2d \times 2d},\\
Q_0 \doteq& \mqty[A_0 & 0_{2d \times 1} \\ 0_{1 \times 2d} & 0] \in \R^{(2d + 1) \times (2d+1)}.
\end{align}

We will proceed with the following claims.

Claim 1. $X_1 \equiv X_0 , x_1^{(n+1)} \equiv x_0^{(n+1)}$

Because we are considering the special case of $L=1$ and because we utilize the same definition of $P_0$ as in Theorem <ref>, the argument proving Claim 1 in Theorem <ref> holds here as well. As a result,
we drop all the subscripts of $X_1$,
as well as subscripts of $x_1^{(i)}$ for $i=1,\dots, n+1$.

Claim 2.

\begin{align}
Y_{1} &= Y_0 + \frac{1}{n} Y_0 X^\top A_0 X\\
y_{1}^{(n+1)} &= y_0^{(n+1)} + \frac{1}{n} Y_0 X^\top A_0 x^{(n+1)}.
\end{align}

This claim is a special case of Claim 2 from the proof of Theorem <ref> in Appendix <ref>, where $L = 1 $. Our block-wise construction of $Q_0$ matches that in the proof of Theorem <ref>. Although our $A_0$ here differs from the specific form of $A_0$ in the proof of Theorem <ref>, this specific form is not utilized in the proof of Claim 2. Therefore, the proof of Claim 2 in Appendix <ref> applies here, and we omit the steps to avoid redundancy.

Claim 3.

\begin{align}
y_{1}^{(i)} = y_0^{(i)} + \indot{M_1 x^{(i)}}{ \frac{1}{n} B_0^\top M_2 X Y_0^\top},
\end{align}

for $i = 1, \dots, n+1$, where $M_2 = \mqty[I_{d} & 0_{d\times d}\\ 0_{d\times d} & 0_{d\times d}]$.

This claim once again is the $L=1$ case of Claim 3 from the proof of Theorem <ref> in Appendix <ref>. The specific form of $M_1$ is not utilized in the proof of Claim 3 from Appendix <ref>, so it applies here.

We can then define $\psi_0 \doteq 0$ and,

\begin{align}
\psi_{1} \doteq \frac{1}{n} B_0^\top M_2 X Y_0^\top \in \R^{2d}. \label{eq: L1 theta_l define }
\end{align}

Then we can write

\begin{align}
y_{1}^{(i)} = y_0^{(i)} + \indot{M_1 x^{(i)}}{\psi_{1}}, \label{eq: L1 y_l+1 evol theta}
\end{align}

for $i=1, \dots, n+1$, which is the claim we made.
In particular, since we assume $y_0^{(n+1)} = 0$, we have

\begin{align}
y_{1}^{(n+1)} = \indot{M_1 x^{(n+1)}}{\psi_{1}}.
\end{align}

Claim 4. The bottom $d$ elements of $\psi_1$ are always 0, i.e., there exists $w_1 \in \R^d$ such that we can express $\psi_1$ as

\begin{align}
\psi_1 = \mqty[w_1 \\ 0_{d \times 1}]. \label{eq: L1 theta_l sparsity}
\end{align}

Since our $B_0$ here is identical to that in the proof of Theorem <ref> in <ref>, 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

\begin{align}
\indot{\psi_{1}}{M_1 x^{(n+1)}} =& \frac{1}{n}\indot{B_0^\top M_2 X Y_0^\top}{M_1 x^{(n+1)}} \explain{By~\eqref{eq: L1 theta_l define }}\\
=& \frac{1}{n}\sum_{i=1}^n \indot{B_0^\top M_2 x^{(i)} y_0^{(i)}}{M_1 x^{(n+1)}}\\
=& \frac{1}{n} \sum_{i=1}^n
\indot{B_0^\top \mqty[\nu^{(i)} \\0_{d\times 1}]
\qty(y_0^{(i)})}{M_1 x^{(n+1)}}\\
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_0 \nu^{(i)}\\0_{d\times 1}]
\qty( y_0^{(i)} )}{M_1 x^{(n+1)}} \explain{By Claim 4}\\
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_0 \nu^{(i)} y_0^{(i)}\\0_{d\times 1}]}
{M_1 x^{(n+1)}}\\
\end{align}

This means

\begin{align}
\indot{w_{1}}{\nu^{(n+1)}}
= \frac{1}{n} \sum_{i=1}^n
\indot{C_0 \nu^{(i)} y_0^{(i)}}{
\nu^{(n+1)}}.
\end{align}

Since the choice of the query $\nu^{(n+1)}$ is arbitrary, we get

\begin{align}
w_{1} = \frac{1}{n}\sum_{i=1}^n C_0 y_0^{(i)} \nu^{(i)}.
\end{align}

In particular, when we construct $Z_0$ such that $\nu^{(i)} = \phi_{i-1}$ and $y_0^{(i)} = R_i$,
we get

\begin{align}
w_1 = \frac{1}{n}\sum_{i=1}^n C_0 R_i \phi_{i-1}
\end{align}

which is the update rule for a single step of TD(0) with $w_0 = 0$.

We also have

\begin{align}
= \indot{\psi_1}{M_1 x^{(n+1)}}
= -\indot{w_1}{\phi^{(n+1)}}.
\end{align}

This concludes our proof.

§.§ Proof of Theorem <ref>

Before we present the proof, we first introduce notations convenient for our analysis.
We decompose $P_0$ and $Q_0$ as

\begin{align}
P_0 = \mqty[P \in \R^{2d\times (2d+1)} \\
p \in \R^{1 \times (2d+1)}],
Q_0 = \mqty[Q_a \in \R^{d\times d} &  Q_b \in \R^{d\times d} & q_c \in \R^{d\times 1} \\
Q_a' \in \R^{d\times d} & Q_b' \in \R^{d\times d} & q_c' \in \R^{d\times 1}\\
q_a \in \R^{1\times d} &  q_b \in \R^{1\times d} & q_c'' \in \R].
\end{align}

One can readily check that $\tf_1$ is independent of $P, Q_b, Q_b', q_b, q_c ,q'_c, q''_c$. Thus, we can assume that these matrices are zero.
Let $z^\paren{i}$ be the $i$-th column of $Z_0$.
Indeed, $\tf_1$ can be written as

\begin{align}
\tf_1(Z_0, \qty{P_0, Q_0})
&= -Z_1\qty[2d+1,n+1] \explain{By \eqref{eq: tf output convention}}\\
& = -\frac{1}{n} p^\top \qty(\sum_{i=1}^n z^{(i)} z^{(i)^\top})  Q_0 z^{(n+1)} \\
& = -\frac{1}{n} \sum_{i=1}^n  \indot{p}{z^{(i)}} z^{(i)^\top} Q_0  z^{(n+1)}\\
& = -\frac{1}{n} \sum_{i=1}^n \indot{p}{z^{(i)}}
\qty(\phi_{i-1}^\top Q_a \phi_{n+1} + \gamma\phi_{i}^\top Q_a' \phi_{n+1} + R_{i} \phi_{n+1}^\top q_a) \label{eq: tf1 block matrix vector form}\\
& = -\frac{1}{n} \sum_{i=1}^n
\qty( \underbrace{\indot{p_{[1:d]}}{\phi_{i-1}} + \gamma \indot{p_{[d+1:2d]}}{\phi_i} + p_{[2d+1]} R_{i}}_{\alpha_i(Z_0, P_0)})\\
& \cdot \qty(\underbrace{\phi_{i-1}^\top Q_a \phi_{n+1} + \gamma(\phi_i)^\top Q_a' \phi_{n+1} + R_i \phi_{n+1}^\top q_a}_{\beta_i(Z_0, Q_0)}).
\end{align}

We prepare the following gradient computations for future use:

\begin{equation}
\label{eq: gradient formula}
\begin{aligned}
\nabla_{p_{[1:d]}} \tf_1(Z_0,\qty{P_0, Q_0}) & = -\frac{1}{n}\sum_{i=1}^n \beta_i(Z_0, Q_0) \phi_{i-1} \\
\nabla_{p_{[d+1:2d]}} \tf_1(Z_0, \qty{P_0, Q_0}) & = -\frac{\gamma}{n} \sum_{i=1}^n \beta_i(Z_0, Q_0) \phi_i\\
\nabla_{Q_a} \tf_1(Z_0,\qty{P_0, Q_0}) & = -\frac{1}{n} \sum_{i=1}^n \alpha_i(Z_0, P_0)  \phi_{i-1} \phi_{n+1}^\top \\
\nabla_{Q_a'} \tf_1(Z_0,\qty{P_0, Q_0}) & = -\frac{\gamma}{n} \sum_{i=1}^n \alpha_i(Z_0, P_0)  \phi_i \phi_{n+1}^\top \\
\nabla_{q_a} \tf_1(Z_0,\qty{P_0, Q_0}) & =  -\frac{1}{n} \sum_{i=1}^n R_{i} \alpha_i(Z_0, P_0) \phi_{n+1}.
\end{aligned}
\end{equation}

We will also reference the following two lemmas in our main proof.

Let $\Lambda$ be a diagonal matrix whose diagonal elements are i.i.d Rademacher random variables  [A Rademacher random variable takes values $1$ or $-1$, each with an equal probability of $0.5$.] $\zeta_1, \dots \zeta_d$.
For any matrix $K \in \R^{d\times d}$, we have that $\E_\Lambda[\Lambda K \Lambda] = \diag(K)$.

First, we can write $\Lambda K \Lambda$ explicitly as

\begin{align}
\Lambda K \Lambda = \mqty[\zeta_1 & 0& \dots & 0 \\
0 & \zeta_2 & \dots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
0 & 0& \dots & \zeta_d]
\mqty[k_{11} & k_{12}& \dots & k_{1d} \\
k_{21} & k_{22}& \dots & k_{2d} \\
\vdots & \vdots & \ddots & \vdots \\
k_{d1} & k_{d2}& \dots & k_{dd}]
\mqty[\zeta_1 & 0& \dots & 0 \\
0 & \zeta_2 & \dots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
0 & 0& \dots & \zeta_d].
\end{align}

Using $\qty(\Lambda K \Lambda)_{ij}$ to denote the element in the $i$-th row at column $j$ of $\Lambda K \Lambda$, from elementary matrix multiplication we have

\begin{align}
\qty(\Lambda K \Lambda)_{ij} = \zeta_i k_{ij} \zeta_j.
\end{align}

When $i \neq j$, $\E\qty[\zeta_i \zeta_j] = \E\qty[\zeta_i]\E\qty[\zeta_j] = 0$ becasue $\zeta_i$ and $\zeta_j$ are independent.
For $i=j$, $\E[\zeta_i \zeta_j] = \E[\zeta_i^2] = 1$.
We can then compute the expectation

\begin{align}
\E_\Lambda \qty[\qty(\Lambda K \Lambda)]_{ij} =
\begin{cases}
k_{ij} &  i=j \\
0 & i\neq j.
\end{cases}
\end{align}

\begin{align}
\E_\Lambda \qty[\Lambda K \Lambda] = \diag(K).
\end{align}

Let $\Pi \in \R^{d\times d}$ be a random permutation matrix uniformly distributed over all $d\times d$ permutation matrices and $L\in\R^{d\times d}$ be a diagonal matrix.
Then, it holds that

\begin{align}
\E_\Pi\qty[\Pi L \Pi^\top] = \frac{1}{d}\tr(L)I_d.
\end{align}

By definition,

\begin{align}
[\Pi L \Pi^\top]_{ij} = \sum_{k=1}^d \Pi_{ik} L_{kk} \Pi_{jk}.
\end{align}

We note that each row of $\Pi$ is a standard basis.
Given the orthogonality of standard bases, we get

\begin{align}
[\Pi L \Pi^\top]_{ij} = \begin{cases}
0 & i\neq j \\
L_{q_i q_i}  & i=j
\end{cases},
\end{align}

where $q_i$ is the unique index such that $\Pi_{iq_i} = 1$.
If the distribution of $\Pi$ is uniform, then $[\Pi L \Pi^\top]_{ii}$ is equal to one of $L_{11},\dots, L_{dd}$ with the same probability.
Thus, the expected value  $[\Pi L \Pi^\top]_{ii}$ is $\frac{1}{d} \tr(L)$.

Now, we start with the proof of the theorem statement.

We recall the definition of the set $\Theta^*$ as

\begin{align} \label{eq:tf_simplified}
\Theta^* \doteq
\cup_{\eta, c, c' \in \R}
\qty{P = \mqty[0_{2d \times 2d} & 0_{2d \times 1} \\ 0_{1 \times 2d} & \eta],
Q =  \mqty[c I_d & 0_{d \times d} & 0_{d \times 1} \\c'I_d  & 0_{d \times d} & 0_{d \times 1} \\ 0_{1 \times d} & 0_{1 \times d} & 0]}.
\end{align}

Suppose $\theta_k \in \Theta^*$, then by (<ref>) and (<ref>), we get

\begin{align}
\tf_1(Z_0, \theta_k)
&= -\frac{\eta_k}{n}\sum_{i=1}^n  R_i \qty( c_k \phi_{i-1}^\top \phi_{n+1} + c'_k \gamma \phi_{i}^\top \phi_{n+1}) \label{eq: tf1 theta belongs to Theta}\\
\tf_1(Z'_0, \theta_k)
&= -\frac{\eta_k}{n}\sum_{i=1}^n  R_{i+1} \qty( c_k \phi_{i}^\top \phi_{n+2} + c'_k \gamma \phi_{i+1}^\top \phi_{n+2}) \label{eq: tf1 prime theta belongs to Theta}\\
\nabla_{p_{[1:d]}} \tf_1(Z_0, \theta_k)
& = -\frac{1}{n}\sum_{i=1}^n \qty( c_k \phi_{i-1}^\top \phi_{n+1} + c'_k \gamma \phi_{i}^\top \phi_{n+1})\phi_{i-1}
\label{eq: tf1 p[1:d] grad theta belongs to Theta}\\
\nabla_{p_{[d+1:2d]}} \tf_1(Z_0, \theta_k)
& = -\frac{\gamma}{n}\sum_{i=1}^n \qty( c_k \phi_{i-1}^\top \phi_{n+1} + c'_k \gamma \phi_{i}^\top \phi_{n+1})\phi_i
\label{eq: tf1 p[d+1:2d] grad theta belongs to Theta}\\
\nabla_{Q_a} \tf_1(Z_0, \theta_k)
& = -\frac{\eta_k}{n} \sum_{i=1}^n R_i  \phi_{i-1} \phi_{n+1}^\top
\label{eq: tf1 Q_a grad theta belongs to Theta}\\
\nabla_{Q'_a} \tf_1(Z_0, \theta_k)
& = -\frac{\gamma \eta_k}{n} \sum_{i=1}^n  R_i  \phi_i \phi_{n+1}^\top
\label{eq: tf1 Q'_a grad theta belongs to Theta}\\
\nabla_{q_a} \tf_1(Z_0, \theta_k)
& = -\frac{\eta_k}{n} \sum_{i=1}^n  R_i^2 \phi_{n+1}
\label{eq: tf1 q_a grad theta belongs to Theta}
\end{align}

Recall the definition of $\Delta(\theta)$ in (<ref>). With a slight abuse of notation, we define $\Delta(p_\qty[{1:d}])$ to be the $p_\qty[{1:d}]$ component of $\Delta(\theta)$,

\begin{align}
\Delta(p_\qty[{1:d}]) \doteq \E\left[(R + \gamma \tf_1(Z_0', \theta) - \tf_1(Z_0, \theta)) \pdv{\tf_1(Z_0, \theta)}{p_\qty[1:d]}\right].
\end{align}

Same goes for $\Delta(p_\qty[{d+1:2d}]), \Delta(Q_a), \Delta(Q_a')$, and $\Delta(q_a)$.

We will prove that

(a) $\Delta(p_\qty[{1:d}]) = \Delta(p_\qty[{d+1:2d}]) = \Delta(q_a) = 0$ for $\Delta (\theta_k)$;

(b) $\Delta(Q_a) = \delta I_d$ and $\Delta(Q'_a) = \delta' I_d$ for some $\delta, \delta' \in\R$ for $\Delta(\theta_k)$

using Assumptions <ref> and <ref>.
We can see that the combination of (a) and (b) are sufficient for proving the theorem.
Recall that $Z_0$ and $Z_0'$ are sampled from $(p_0, p, r, \phi)$.
We make the following claims to assist our proof of (a) and (b).

Claim 1.
Let $\zeta$ be a Rademacher random variable.
We denote $Z_\zeta$ and $Z_\zeta'$ as the prompts sampled from $(p_0, p, r, \zeta \phi)$. We then have $Z_0 \triangleq Z_\zeta$ and $Z_0' \triangleq Z_\zeta'$.
To show this is true, we notice that for any realization of $\zeta$, denoted as $\bar \zeta \in \qty{1, -1}$,
we have

\begin{align}
\Pr(p_0, p, r, \phi) &= \Pr(p_0, p, r)\Pr(\phi) \explain{Assumption~\ref{assumption independence}}\\
&= \Pr(p_0, p, r)\Pr(\bar \zeta I_d \phi) \explain{Assumption~\ref{assumption symmetry}}\\
&= \Pr(p_0, p, r, \bar \zeta \phi) \explain{Assumption~\ref{assumption independence}}.
\end{align}

It then follows that

\begin{align}
\Pr(p_0, p, r, \phi) =& \Pr(p_0, p, r, \phi) \sum_{\bar \zeta \in \qty{1, -1}}  \Pr(\zeta = \bar \zeta) \\
=&  \sum_{\bar \zeta \in \qty{1, -1}}  \Pr(p_0, p, r, \phi) \Pr(\zeta = \bar \zeta) \\
=&  \sum_{\bar \zeta \in \qty{1, -1}}  \Pr(p_0, p, r, \bar \zeta \phi) \Pr(\zeta = \bar \zeta) \\
=& \Pr(p_0, p, r, \zeta \phi).
\end{align}

This implies Claim 1 holds.

Claim 2.
Define $\Lambda$ as the diagonal matrix whose diagonal elements are i.i.d. Rademacher random variables $\zeta_1, \dots, \zeta_d$.
We denote $Z_\Lambda$ and $Z'_\Lambda$ as the prompts sampled from $(p_0, p, r, \Lambda\phi)$, where $\Lambda\phi$ means $[\Lambda\phi(s)]_{s\in\fS}$.
We then have $Z_0 \triangleq Z_\Lambda$ and $Z_0' \triangleq Z_\Lambda'$.
The proof follows the same procedures as Claim 1.

Claim 3.
Let $\Pi$ be a random permutation matrix uniformly distributed over all $d\times d$ permutation matrices.
We denote $Z_\Pi$ and $Z'_\Pi$ as the prompts sampled from $(p_0, p, r, \Pi\phi)$, where $\Pi\phi$ means $[\Pi\phi(s)]_{s\in\fS}$.
We then have $Z_0 \triangleq Z_\Pi$ and $Z_0' \triangleq Z_\Pi'$.
The proof follows the same procedures as Claim 1.

Proof of (a) using Claim 1

It is easy to check by (<ref>) that

\begin{align}
\tf_1(Z_\zeta,\theta_k)
&= -\frac{\eta_k}{n}\sum_{i=1}^n  R_i \qty( c_k \zeta^2 \phi_{i-1}^\top \phi_{n+1} + c'_k \gamma \zeta^2 \phi_{i}^\top \phi_{n+1})\\
&= \underbrace{\zeta^2}_{=1} \tf_1(Z_0,\theta_k)\\
&= \tf_1(Z_0,\theta_k). \label{eq: (a) proof TF equivalence zeta}
\end{align}

Similarly, one can check that $\tf_1(Z'_\zeta, \theta_k) = \tf_1(Z'_0, \theta_k)$.

\begin{align}
\nabla_{p_{[1:d]}} \tf_1(Z_\zeta, \theta_k)
=& -\frac{1}{n}\sum_{i=1}^n \qty( c_k \underbrace{\zeta^2}_{=1} \phi_{i-1}^\top \phi_{n+1} + c'_k \gamma \underbrace{\zeta^2}_{=1} \phi_{i}^\top \phi_{n+1})\zeta \phi_{i-1} \\
=& -\frac{\zeta}{n}\sum_{i=1}^n \qty( c_k \phi_{i-1}^\top \phi_{n+1} + c'_k \gamma \phi_{i}^\top \phi_{n+1})\phi_{i-1}\\
=& \zeta \nabla_{p_{[1:d]}} \tf_1(Z_0, \theta_k).\label{eq: (a) proof grad equivalence zeta}
\end{align}

Then, from  (<ref>), we get

\begin{align}
&\Delta (p_{[1:d]}) \\
=& \E\qty[\qty(R_{n+2} + \gamma\tf_1(Z_0', \theta_k) - \tf_1 (Z_0, \theta_k))\nabla_{p_{[1:d]}}\tf_1(Z_0, \theta_k)] \\
=& \E\qty[\qty(R_{n+2} + \gamma\tf_1(Z_\zeta', \theta_k) - \tf_1 (Z_\zeta, \theta_k))\nabla_{p_{[1:d]}}\tf_1(Z_\zeta, \theta_k)] \explain{By Claim 1}\\
=& \E_\zeta\qty[\E\qty[\qty(R_{n+2} + \gamma\tf_1(Z_\zeta', \theta_k) - \tf_1 (Z_\zeta, \theta_k))\nabla_{p_{[1:d]}}\tf_1(Z_\zeta, \theta_k) \mid \zeta ]] \\
=& \E_\zeta\qty[\E\qty[\qty(R_{n+2} + \gamma\tf_1(Z_0', \theta_k) - \tf_1 (Z_0, \theta_k))\zeta\nabla_{p_{[1:d]}}\tf_1(Z_0, \theta_k) \mid \zeta ]] \explain{By~\eqref{eq: (a) proof TF equivalence zeta},~\eqref{eq: (a) proof grad equivalence zeta}}\\
=& \E_\zeta\qty[\zeta\E\qty[\qty(R_{n+2} + \gamma\tf_1(Z_0', \theta_k) - \tf_1 (Z_0, \theta_k)) \nabla_{p_{[1:d]}}\tf_1(Z_0, \theta_k) \mid \zeta ]] \\
=& \E_\zeta\qty[\zeta\E\qty[\qty(R_{n+2} + \gamma\tf_1(Z_0', \theta_k) - \tf_1 (Z_0, \theta_k)) \nabla_{p_{[1:d]}}\tf_1(Z_0, \theta_k)]] \\
=& \E_\zeta\qty[\zeta]\E\qty[\qty(R_{n+2} + \gamma\tf_1(Z_0', \theta_k) - \tf_1 (Z_0, \theta_k)) \nabla_{p_{[1:d]}}\tf_1(Z_0, \theta_k)] \\
=& 0.
\end{align}

The proof is analogous for $\Delta(p_{[d+1:2d]}) = 0$, and $\Delta (q_a) = 0$.

Proof of (b) using Claims 2 and 3
We first show that $\Delta (Q_a)$ is a diagonal matrix.
Similar to (a), we have

\begin{align}
\label{eq: (b) flip tf equivalence}
\tf_1(Z_\Lambda,\theta_k)
&= -\frac{1}{n}\sum_{i=1}^n  \eta_k R_i \qty(c_k  \phi_{i-1}^\top \underbrace{\Lambda^2}_{= I} \phi_{n+1} + c_k' \gamma  \phi_{i}^\top \underbrace{\Lambda^2}_{= I} \phi_{n+1})\\
&= \tf_1(Z_0,\theta_k).
\end{align}

Similarly, we get $\tf_1(Z'_\Lambda,\theta_k) = \tf_1(Z'_0,\theta_k)$.
Additionally, we have

\begin{align}
\label{eq: (b) flip gradient relation}
\nabla_{Q_a} \tf_1(Z_\Lambda,\theta_k)
= -\frac{1}{n}\sum_{i=1}^n \eta_k R_i  \Lambda \phi_{i-1} \phi_{n+1}^\top  \Lambda^\top
= \Lambda \nabla_{Q_a} \tf_1(Z_0, \theta_k) \Lambda.
\end{align}

By (<ref>) again, we get

\begin{align}
& \Delta(Q_a)\\
=& \E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k))\nabla_{Q_a}\tf_1(Z_0, \theta_k)]\\
=& \E\qty[(R_{n+2} + \gamma \tf_1(Z'_\Lambda, \theta_k) - \tf_1(Z_\Lambda, \theta_k))\nabla_{Q_a} \tf_1(Z_\Lambda, \theta_k)]\explain{By Claim 2}\\
=& \E_\Lambda\qty[\E\qty[(R_{n+2} + \gamma \tf_1(Z'_\Lambda, \theta_k) - \tf_1(Z_\Lambda, \theta_k))\nabla_{Q_a}\tf_1(Z_\Lambda, \theta_k) \mid \Lambda]]\\
=& \E_\Lambda\qty[\E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k))\Lambda \nabla_{Q_a}\tf_1(Z_0, \theta_k)\Lambda \mid \Lambda]] \explain{By \eqref{eq: (b) flip tf equivalence}, \eqref{eq: (b) flip gradient relation}}\\
=& \E_\Lambda\qty[\Lambda \E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k)) \nabla_{Q_a}\tf_1(Z_0, \theta_k) \mid \Lambda]\Lambda]\\
=& \E_\Lambda\qty[\Lambda \E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k)) \nabla_{Q_a}\tf_1(Z_0, \theta_k)]\Lambda]\\
=& \diag\qty(\E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k))\nabla_{Q_a}\tf_1(Z_0, \theta_k)])\explain{By Lemma~\ref{lemma: rademacher diagonal}}\\
=& \diag(\Delta(Q_a)) \label{eq: delta Q diag}.
\end{align}

The last equation holds if and only if $\Delta(Q_a)$ is diagonal.
We have proven this claim.

Now, we prove that $\Delta(Q_a) = \delta I_d$ for some $\delta \in \R$ using Claim 3 and Lemma <ref>.
Let $\Pi$ be a random permutation matrix uniformly distributed over all permutation matrices.
Recall the definition of $Z_\Pi$ and $Z'_\Pi$ in Claim 3.
We have

\begin{align}
\label{eq: (b) perm tf equivalence}
\tf_1(Z_\Pi,\theta_k) = -\frac{1}{n}\sum_{i=1}^n  \eta_k R_i \qty(c_k \phi_{i-1}^\top \underbrace{\Pi^\top \Pi}_{= I} \phi_{n+1} + c'_k \gamma \phi_i^\top \underbrace{\Pi^\top \Pi}_{= I} \phi_{n+1})
= \tf_1(Z_0,\theta_k).
\end{align}

Analogously, we get $\tf_1(Z_\Pi', \theta_k) = \tf_1(Z_0', \theta_k)$.
Furthermore, we have

\begin{align}
\label{eq: (b) perm gradient relation}
\nabla_{Q_a} \tf_1(Z_\Pi, \theta_k) = -\frac{1}{n}\sum_{i=1}^n \eta_k R_i  \Pi \phi_{i-1} \phi_{n+1}^\top \Pi^\top
= \Pi \nabla_{Q_a} \tf_1(Z_0, \theta_k) \Pi^\top.
\end{align}

By (<ref>), we are ready to show that

\begin{align}
& \Delta(Q_a)\\
=& \E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k))\nabla_{Q_a}\tf_1(Z_0, \theta_k)]\\
=& \E\qty[(R_{n+2} + \gamma \tf_1(Z'_\Pi, \theta_k) - \tf_1(Z_\Pi, \theta_k))\nabla_{Q_a}\tf_1(Z_\Pi, \theta_k)]\explain{By Claim 3}\\
=& \E_{\Pi} \qty[\E\qty[(R_{n+2} + \gamma \tf_1(Z'_\Pi, \theta_k) - \tf_1(Z_\Pi, \theta_k))\nabla_{Q_a}\tf_1(Z_\Pi, \theta_k) \mid \Pi]]\\
=& \E_{\Pi} \qty[\E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k))\Pi \nabla_{Q_a}\tf_1(Z_0, \theta_k) \Pi^\top\mid \Pi]]\explain{By \eqref{eq: (b) perm tf equivalence}, \eqref{eq: (b) perm gradient relation}}\\
=& \E_{\Pi} \qty[\Pi \E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k))\nabla_{Q_a}\tf_1(Z_0, \theta_k) \mid \Pi]\Pi^\top]\\
=& \E_{\Pi} \qty[\Pi \E\qty[(R_{n+2} + \gamma \tf_1(Z'_0, \theta_k) - \tf_1(Z_0, \theta_k))\nabla_{Q_a}\tf_1(Z_0, \theta_k)]\Pi^\top]\\
=& \E_{\Pi} \qty[\Pi \diag\qty(\Delta(Q_a))\Pi^\top] \\
=& \frac{1}{d}\tr(\Delta(Q_a))I_d\explain{By Lemma~\ref{lemma: permutation diagonal}}\\
=& \delta I_d.
\end{align}

The proof is analogous for $\Delta(Q_a') = \delta'I_d$ for some $\delta' \in \R$.

Suppose that $\Delta(p_{[2d+1]}) = \rho \in \R$,
we now can conclude that

\begin{align}
\Delta(\theta_k)
= \qty{
\Delta(P_0)  = \mqty[0_{2d\times2d} & 0_{2d\times1}\\
0_{1\times2d} & \rho],
\Delta(Q_0)  = \mqty[\delta I_d & 0_{d \times d} & 0_{d \times 1} \\
\delta' I_d & 0_{d \times d} & 0_{d \times 1} \\
0_{1 \times d} & 0_{1 \times d} & 0]}.
\end{align}

Therefore, according to (<ref>), we get

\begin{align}
=&\theta_k  + \alpha_k\Delta(\theta_k)\\
\mqty[0_{2d\times2d} & 0_{2d\times1}\\
0_{1\times2d} & \eta_k + \alpha_k\rho],
\mqty[c_k + \alpha_k\delta I_d & 0_{d \times d} & 0_{d \times 1} \\
c_k' + \alpha_k\delta' I_d & 0_{d \times d} & 0_{d \times 1} \\
0_{1 \times d} & 0_{1 \times d} & 0]
\in \Theta_*.
\end{align}

§.§ Proof of Corollary <ref>

We recall from (<ref>) that the embedding evolves according to

\begin{align}
Z_{l+1} = Z_l + \frac{1}{n} P_l Z_l M(Z_l^\top Q_l Z_l).
\end{align}

We again refer to the elements in $Z_l$ as $\qty{(x^{(i)}_l, y^{(i)}_l)}_{i=1,\dots, n+1}$ in the following way

\begin{align}
Z_l = \mqty[x_l^{(1)} & \dots & x_l^{(n)} & x_l^{(n+1)} \\
y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)}],
\end{align}

where we recall that $Z_l \in \R^{(2d + 1) \times (n+1)}, x_l^{(i)} \in \R^{2d}, y^{(i)}_l \in \R$.
Sometimes, it is more convenient to refer to the first half and second half of $x_l^{(i)}$ separately,
by, e.g., $\nu^{(i)}_l \in \R^d, \xi_l^{(i)} \in \R^d$, i.e., $x_l^{(i)} = \mqty[\nu_l^{(i)} \\ \xi_l^{(i)}]$.
Then, we have

\begin{align}
Z_l = \mqty[\nu_l^{(1)} & \dots & \nu_l^{(n)} & \nu_l^{(n+1)} \\
\xi_l^{(1)} & \dots & \xi_l^{(n)} & \xi_l^{(n+1)} \\
y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)}].
\end{align}

We utilize the shorthands

\begin{align}
X_l =& \mqty[x_l^{(1)} & \dots & x_l^{(n)}] \in \R^{2d \times n}, \\
Y_l =&  \mqty[y_l^{(1)} & \dots & y_l^{(n)}] \in \R^{1 \times n}.
\end{align}

Then we have

\begin{align}
Z_l = \mqty[X_l & x_l^{(n+1)} \\
Y_l & y_l^{(n+1)}].
\end{align}

For the input $Z_0$,
we assume $\xi_0^{(n+1)} = 0, y_0^{(n+1)} = 0$ but all other entries of $Z_0$ are arbitrary.
We recall our definition of $M$ in (<ref>) and $\qty{P_l^{\rg}, Q_l^{\rg}}$ in (<ref>).
In particular,
we can express $Q_l^{\rg}$ in a more compact way as

\begin{align}
M_1 \doteq& \mqty[-I_{d} & I_{d} \\
0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
M_2 \doteq& -M_1\\
B_l \doteq& \mqty[C_l^\top & 0_{d\times d} \\
0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
A_l \doteq& M_2^\top B_l M_1  = \mqty[-C_l^\top & C_l^\top \\
C_l^\top & -C_l^\top] \in \R^{2d \times 2d},\\
Q_l^{\rg} \doteq& \mqty[A_l & 0_{2d \times 1} \\ 0_{1 \times 2d} & 0] \in \R^{(2d + 1) \times (2d+1)}.
\end{align}

We then verify the following claims.

Claim 1. $X_l \equiv X_0 , x_l^{(n+1)} \equiv x_0^{(n+1)}, \forall l.$

We note that $P_l^{\rg}$ is the key reason Claim 1 holds and is the same as the TD(0) case.
Referring to <ref>, we omit the proof of Claim 1 here.

Claim 2.

\begin{align}
Y_{l+1} &= Y_l + \frac{1}{n} Y_l X^\top A_l X\\
y_{l+1}^{(n+1)} &= y_l^{(n+1)} + \frac{1}{n} Y_l X^\top A_l x^{(n+1)}.
\end{align}

Since the only difference between the true residual gradient and TD(0) configurations is the internal structure of $A_l$, we argue that it's irrelevant to Claim 2.
We therefore again refer the readers to <ref> for a detailed proof.

Claim 3.

\begin{align}
y_{l+1}^{(i)} = y_0^{(i)} + \indot{M_1 x^{(i)}}{\frac{1}{n}\sum_{j=0}^l B_j^\top M_2 X Y_j^\top},
\end{align}

for $i = 1, \dots, n+1$.

By Claim 2, we can unroll $Y_{l+1}$ as

\begin{align}
Y_{l+1} &= Y_l + \frac{1}{n} Y_l X^\top A_l X\\
Y_l &= Y_{l-1} + \frac{1}{n} Y_{l-1} X^\top A_{l-1} X\\
Y_1 & = Y_0 + \frac{1}{n} Y_0 X^\top A_0 X.
\end{align}

We can then compactly express $Y_{l+1}$ as

\begin{align}
Y_{l+1} = Y_0 + \frac{1}{n}\sum_{j=0}^l Y_j X^\top A_j X.
\end{align}

Recall that we define $A_j = M_2^\top B_j M_1$.
Then, we can rewrite $Y_{l+1}$ as

\begin{align}
&= Y_0 + \frac{1}{n} \sum_{j=0}^l Y_j X^\top M_2^\top B_j M_1 X.
\end{align}

With the identical procedure, we can easily rewrite $y_{l+1}^{(n+1)}$ as

\begin{align}
y_{l+1}^{(n+1)} &= y_0^{(n+1)} + \frac{1}{n} \sum_{j=0}^l Y_j X^\top M_2^\top B_j M_1 x^{(n+1)}.
\end{align}

In light of this,
we define $\psi_0 \doteq 0$ and for $l = 0, \dots$

\begin{align}
\psi_{l+1}
\doteq & \frac{1}{n}\sum_{j=0}^l B_j^\top M_2 X Y_j^\top \in \R^{2d}\\
= & \psi_l + \frac{1}{n} B_l^\top M_2 X Y_l^\top \label{eq: true RG psi evol}
\end{align}

Then we can write

\begin{align}
y_{l+1}^{(i)} = y_0^{(i)} + \indot{M_1 x^{(i)}}{\psi_{l+1}}, \label{eq: y_l+1 evol psi true RG}
\end{align}

for $i=1, \dots, n+1$, which is the claim we made.
In particular, since we assume $y_0^{(n+1)} = 0$, we have

\begin{align}
y_{l+1}^{(n+1)} = \indot{M_1 x^{(n+1)}}{\psi_{l+1}}.
\end{align}

Claim 4. The bottom $d$ elements of $\psi_l$ are always 0, i.e., there exists a sequence $\qty{w_l \in \R^d}$ such that we can express $\psi_l$ as

\begin{align}
\psi_l = \mqty[w_l \\ 0_{d \times 1}].
\end{align}

for all $l = 0,1,\dots, L$.

Since $B_l$ is the key reason Claim 4 holds and is identical to the TD(0) case, we refer the reader to <ref> for detailed proof.

Given all the claims above,
we can then compute that

\begin{align}
&\indot{\psi_{l+1}}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} + \frac{1}{n}\indot{B_l^\top M_2 X Y_l^\top}{M_1 x^{(n+1)}} \explain{By~\eqref{eq: true RG psi evol}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} + \frac{1}{n}\sum_{i=1}^n \indot{B_l^\top M_2 x^{(i)} y_l^{(i)}}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n \indot{B_l^\top M_2 x^{(i)} \qty(\indot{\psi_l}{M_1 x^{(i)}} + y_0^{(i)})}{M_1 x^{(n+1)}} \explain{By \eqref{eq: y_l+1 evol psi true RG}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{B_l^\top \mqty[\nu^{(i)} - \xi^{(i)} \\
0_{d\times 1}]
\qty( \indot{\psi_l}{\mqty[-\nu^{(i)} + \xi^{(i)}\\ 0_{d\times 1}]} + y_0^{(i)})}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_l\qty(\nu^{(i)}- \xi^{(i)})\\
0_{d\times 1}]
\qty( y_0^{(i)} +  w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)} )}{M_1 x^{(n+1)}} \explain{By Claim 4}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_l \qty(\nu^{(i)}-\xi^{(i)} ) \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})\\0_{d\times 1}]}
{M_1 x^{(n+1)}}\\
\end{align}

This means

\begin{align}
\indot{w_{l+1}}{\nu^{(n+1)}}
= \indot{w_l}{\nu^{(n+1)}}
+ \frac{1}{n} \sum_{i=1}^n
\indot{C_l \qty(\nu^{(i)}-\xi^{(i)}) \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})}{
\nu^{(n+1)}}.
\end{align}

Since the choice of the query $\nu^{(n+1)}$ is arbitrary, we get

\begin{align}
w_{l+1} = w_l + \frac{1}{n}\sum_{i=1}^n C_l \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)}) \qty(\nu^{(i)}- \xi^{(i)}).
\end{align}

In particular, when we construct $Z_0$ such that $\nu^{(i)} = \phi_{i-1}$, $\xi^{(i)} = \gamma \phi_i$ and $y_0^{(i)} = R_i$,
we get

\begin{align}
w_{l+1} = w_l + \frac{1}{n}\sum_{i=1}^n C_l \qty(R_i + \gamma w_l^\top \phi_i - w_l^\top \phi_{i-1})\qty(\phi_{i-1} -\gamma \phi_i)
\end{align}

which is the update rule for pre-conditioned residual gradient learning.

We also have

\begin{align}
= \indot{\psi_l}{M_1 x^{(n+1)}}
= -\indot{w_l}{\phi^{(n+1)}}.
\end{align}

This concludes our proof.

§.§ Proof of Corollary <ref>

The proof presented here closely mirrors the methodology and notation established in the proof of Theorem <ref> from Appendix <ref>. We begin by recalling the embedding evolution from (<ref>) as,

\begin{align}
Z_{l+1} = Z_l + \frac{1}{n} P_l Z_l M^\tdlambda(Z_l^\top Q_l Z_l).
\end{align}

where we have substituted the original mask defined in (<ref>) with the TD($\lambda$) mask in (<ref>).
We once again refer to the elements in $Z_l$ as $\qty{(x^{(i)}_l, y^{(i)}_l)}_{i=1,\dots, n+1}$ in the following way

\begin{align}
Z_l = \mqty[x_l^{(1)} & \dots & x_l^{(n)} & x_l^{(n+1)} \\ y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)}],
\end{align}

where we recall that $Z_l \in \R^{(2d + 1) \times (n+1)}, x_l^{(i)} \in \R^{2d}, y^{(i)}_l \in \R$.
We utilize, $\nu^{(i)}_l \in \R^d, \xi_l^{(i)} \in \R^d$, to refer to the first half and second half of $x_l^{(i)}$ i.e., $x_l^{(i)} = \mqty[\nu_l^{(i)} \\ \xi_l^{(i)}]$.

Then we have

\begin{align}
Z_l = \mqty[\nu_l^{(1)} & \dots & \nu_l^{(n)} & \nu_l^{(n+1)} \\
\xi_l^{(1)} & \dots & \xi_l^{(n)} & \xi_l^{(n+1)} \\
y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)}].
\end{align}

We further define as shorthands,

\begin{align}
X_l =& \mqty[x_l^{(1)} & \dots & x_l^{(n)}] \in \R^{2d \times n}, \\
Y_l =&  \mqty[y_l^{(1)} & \dots & y_l^{(n)}] \in \R^{1 \times n}.
\end{align}

Then the blockwise structure of $Z_l$ can be succinctly expressed as:

\begin{align}
Z_l = \mqty[X_l & x_l^{(n+1)} \\ Y_l & y_l^{(n+1)}]. \label{eq: z_l blockwise X Y lambda}
\end{align}

We proceed to the formal arguments by paralleling those in Theorem <ref>. As in the theorem, we assume that certain initial conditions, such as $\xi_0^{(n+1)} = 0$ and $y_0^{(n+1)} = 0$, hold, but other entries of $Z_0$ are arbitrary.
We recall our definition of $M^\tdlambda$ in (<ref>) and $\qty{P_l^{\td}, Q_l^{\td}}_{l=0,\dots,L-1}$ in (<ref>).
In particular,
we can express $Q_l^{\td}$ in a more compact way as

\begin{align}
M_1 \doteq& \mqty[-I_{d} & I_{d} \\0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d}, \label{eq: lambda M_1 define}\\
B_l \doteq& \mqty[C_l^\top & 0_{d\times d} \\ 0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
A_l \doteq& B_l M_1  = \mqty[-C_l^\top & C_l^\top \\ 0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
%   A_l^\top =& M_1^\top B_l^\top \in \R^{2d \times 2d},\\
Q_l^{\td} \doteq& \mqty[A_l & 0_{2d \times 1} \\ 0_{1 \times 2d} & 0] \in \R^{(2d + 1) \times (2d+1)},
\end{align}

We now proceed with the following claims.

In subsequent steps, it sometimes is useful to refer to the matrix $M^\tdlambda Z^{\top}$ in block form. Therefore, we will define $H^\top \in \R^{\paren{n \times 2d}}$ as the first $n$ rows of $M_{\td(\lambda)}Z^\top$ except for the last column, which we define as $Y_l^{(\lambda)} \in \R^n$.

\begin{align}
M^\tdlambda Z_l^{\top} = \mqty[H^\top & Y_l^{(\lambda)} \\ 0_{1\times 2d} & 0] \in \R^{(n+1) \times (2d+1)} \label{eq: M lambda ZT}
\end{align}

Let $h^{(i)}$ denote $i$-th column of $H$.

We proceed with the following claims.

Claim 1. $X_l \equiv X_0 , x_l^{(n+1)} \equiv x_0^{(n+1)}, \forall l.$

Because we utilize the same definition of $P_l^{\td}$ as in Theorem <ref>, the argument proving Claim 1 in Theorem <ref> holds here as well. As a result,
we drop all the subscripts of $X_l$,
as well as subscripts of $x_l^{(i)}$ for $i=1,\dots, n+1$.

Claim 2.  Let $H \in \R^{\paren{2d \times n}}$, where the $i$-th column of $H$ is,

\begin{align}
h^{(i)} = \sum_{k=1} ^i \lambda^{i-k} x^{(i)} \in \R^{2d}. \label{eq: h^i definition lambda}
\end{align}

Then we can write the updates for $Y_{l+1}$, and $y_{l+1}^{(n+1)}$ as,

\begin{align}
Y_{l+1} &= Y_l + \frac{1}{n} Y_l H^\top A_l X,\\
y_{l+1}^{(n+1)} &= y_l^{(n+1)} + \frac{1}{n} Y_l H^\top A_l x^{(n+1)}.
\end{align}

We will show this by factoring the embedding evolution into the product of $P_l^{\td} Z_l$ and $M^\tdlambda  Z_l^\top$, and  $Q_l^{\td} Z_l$.
Firstly, we have

\begin{align}
P_l^{\td} Z_l = \mqty[0_{2d \times n} & 0_{2d \times 1}\\
Y_l & y_l^{(n+1)}].
\end{align}

Next we analyze $M^\tdlambda  Z_l^{\top}$. From basic matrix algebra we have,

\begin{align}
M^\tdlambda Z^\top  &=
\mqty[1 & 0 & 0 & 0 & \cdots & 0 & 0\\
\lambda & 1 & 0 & 0 & \cdots & 0 & 0\\
\lambda^2 & \lambda & 1 & 0 &\cdots & 0 & 0\\
\lambda^3 & \lambda^2 & \lambda & 1 & \cdots & 0 & 0\\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots\\
\lambda^{n-1} & \lambda^{n-2} & \lambda^{n-3} & \lambda^{n-4} & \cdots & 1 & 0\\
0 & 0 & 0 & 0 & \cdots & 0 & 0] \mqty[x^{(1)^\top} & y^{(1)}\\
x^{(2)^\top} & y^{(2)}\\
x^{(3)^\top} & y^{(3)}\\
\vdots & \vdots\\
x^{(n)^\top} & y^{(n)}\\
x^{(n+1)^\top} & 0
] \\
&= \mqty[ x^{(1)^\top} &  y_l^{(1)}\\
x^{(2)^\top} + \lambda x^{(1)^\top} & y_l^{(2)} + \lambda y_l^{(2)}\\
\vdots & \vdots\\
\sum_{i=1}^{n} \lambda^{n-i} x_i^\top &  \sum_{i=1}^{n} \lambda^{n-i} y_l^{(i)}\\
0_{1\times 2d} & 0
&=  \mqty[ h^{(1)^\top} &  y_l^{(1)}\\
h^{(2)^\top} & y_l^{(2)} + \lambda y_l^{(1)}\\
\vdots & \vdots\\
h^{(n)^\top} &  \sum_{i=1}^{n} \lambda^{n-i} y_l^{(n)}\\
0_{1\times 2d} & 0
&= \mqty[H^\top & K_l^{(\lambda)} \\ 0_{1\times 2d} & 0], \label{eq: block Mtd ZT}
\end{align}

where $K_l^{(\lambda)} \in \R^d$ is introduced for notation simplicity.

Then, we analyze $M^\tdlambda  Z_l^\top Q_l^{\td} Z_l$.
Applying the block matrix notations, we get

\begin{align}
\qty(M^\tdlambda  Z_l^\top) Q_l^{\td} Z_l &= \mqty[H^\top & K_l^{(\lambda)}\\
0_{1\times 2d} & 0]
\mqty[A_l & 0_{2d \times 1} \\
0_{1 \times 2d} & 0]
\mqty[X & x^{(n+1)} \\
Y_l & y_l^{(n+1)}]\\
&= \mqty[H^\top A_l & 0_{n \times 1}\\
0_{1\times 2d} & 0]
\mqty[X & x^{(n+1)} \\
Y_l & y_l^{(n+1)}]\\
&= \mqty[H^\top A_l X & H^\top A_l x^{(n+1)}\\
0_{1\times 2d} & 0].
\end{align}

Combining the two, we get

\begin{align}
P_l^{\td} Z_l \qty(M^\tdlambda  Z_l^\top Q_l^{\td} Z_l)
&= \mqty[0_{2d \times n} & 0_{2d \times 1}\\
Y_l & y_l^{(n+1)}]
\mqty[H^\top A_l X & H^\top A_l x^{(n+1)}\\
0_{1\times 2d} & 0]\\
&= \mqty[0_{2d \times n} & 0_{2d \times 1}\\
Y_l H^\top A_l X & Y_l H^\top A_l x^{(n+1)}].
\end{align}

Hence, according to our update rule in (<ref>), we get

\begin{align}
Y_{l+1} &= Y_l + \frac{1}{n} Y_l H^\top A_l X\\
y_{l+1}^{(n+1)} &= y_l^{(n+1)} + \frac{1}{n} Y_l H^\top A_l x^{(n+1)}.
\end{align}

Claim 3.

\begin{align}
y_{l+1}^{(i)} = y_0^{(i)} + \indot{M_1 x^{(i)}}{\frac{1}{n}\sum_{i=0}^l B_i^\top M_2 X Y_i^\top},
\end{align}

for $i = 1, \dots, n+1$, where $M_2 = \mqty[I_{d} & 0_{d\times d}\\ 0_{d\times d} & 0_{d\times d}]$.

Following Claim 2, we can unroll the recursive definition of $Y_{l+1}$ and express it compactly as,

\begin{align}
Y_{l+1} = Y_0 + \frac{1}{n}\sum_{i=0}^l Y_i H^\top A_i X.
\end{align}

Recall that we define $A_i = B_i M_1$.
Then, we can rewrite $Y_{l+1}$ as

\begin{align}
&= Y_0 + \frac{1}{n} \sum_{i=0}^l Y_i H^\top M_2 B_i M_1 X.
\end{align}

The introduction of $M_2$ here does not break the equivalence because $B_i = M_2 B_i$.
However, it will help make our proof steps easier to comprehend later.

With the identical recursive unrolling procedure, we can rewrite $y_{l+1}^{(n+1)}$ as

\begin{align}
y_{l+1}^{(n+1)} &= y_0^{(n+1)} + \frac{1}{n} \sum_{i=0}^l Y_i H^\top M_2 B_i M_1 x^{(n+1)}.
\end{align}

In light of this,
we define $\psi_0 \doteq 0$ and for $l = 0, \dots$

\begin{align}
\psi_{l+1} \doteq& \frac{1}{n}\sum_{i=0}^l B_i^\top M_2 H Y_i^\top \in \R^{2d}. \label{eq: lambda theta_l define}
\end{align}

Then we can write

\begin{align}
y_{l+1}^{(i)} = y_0^{(i)} + \indot{M_1 x^{(i)}}{\psi_{l+1}}, \label{eq: lambda y_l+1 evol theta}
\end{align}

for $i=1, \dots, n+1$, which is the claim we made.
In particular, since we assume $y_0^{(n+1)} = 0$, we have

\begin{align}
y_{l+1}^{(n+1)} = \indot{M_1 x^{(n+1)}}{\psi_{l+1}}.
\end{align}

Claim 4. The bottom $d$ elements of $\psi_l$ are always 0, i.e., there exists a sequence $\qty{w_l \in \R^d}$ such that we can express $\psi_l$ as

\begin{align}
\psi_l = \mqty[w_l \\ 0_{d \times 1}]. \label{eq: lambda theta_l sparsity}
\end{align}

for all $l = 0,1,\dots, L$.

Because we utilize the same definition of $B_l$ as in Theorem <ref> when defining $\psi_{l+1}$, the argument proving Claim 4 in Theorem <ref> holds here as well. We omit the steps to avoid redundancy.

Given all the claims above,
we can then compute that

\begin{align}
&\indot{\psi_{l+1}}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} + \frac{1}{n}\indot{B_l^\top M_2 H Y_l^\top}{M_1 x^{(n+1)}} \explain{By~\eqref{eq: lambda theta_l define}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} + \frac{1}{n}\sum_{i=1}^n \indot{B_l^\top M_2 h^{(i)} y_l^{(i)}}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n \indot{B_l^\top M_2 h^{(i)} \qty(\indot{\psi_l}{M_1 x^{(i)}} + y_0^{(i)})}{M_1 x^{(n+1)}} \explain{By \eqref{eq: lambda y_l+1 evol theta}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{B_l^\top \mqty[\qty(\sum_{k=1} ^i \lambda^{i-k} \nu^{(i)})\\0_{d\times 1}]
\qty( \indot{\psi_l}{\mqty[-\nu^{(i)} + \xi^{(i)}\\ 0_{d\times 1}]} + y_0^{(i)})}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_l \qty(\sum_{k=1} ^i \lambda^{i-k} \nu^{(i)})\\0_{d\times 1}]
\qty( y_0^{(i)} +  w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)} )}{M_1 x^{(n+1)}} \explain{By Claim 4}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_l \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})\qty(\sum_{k=1} ^i \lambda^{i-k} \nu^{(i)})\\0_{d\times 1}]}
{M_1 x^{(n+1)}}\\
\end{align}

This means

\begin{align}
\indot{w_{l+1}}{\nu^{(n+1)}}
= \indot{w_l}{\nu^{(n+1)}}
+ \frac{1}{n} \sum_{i=1}^n
\indot{C_l \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})\qty(\sum_{k=1} ^i \lambda^{i-k} \nu^{(i)})}{
\nu^{(n+1)}}.
\end{align}

Since the choice of the query $\nu^{(n+1)}$ is arbitrary, we get

\begin{align}
w_{l+1} = w_l + \frac{1}{n}\sum_{i=1}^n C_l \qty(y_0^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})\qty(\sum_{k=1} ^i \lambda^{i-k} \nu^{(i)}).
\end{align}

In particular, when we construct $Z_0$ such that $\nu^{(i)} = \phi_{i-1}$, $\xi^{(i)} = \gamma \phi_i$ and $y_0^{(i)} = R_i$,
we get

\begin{align}
w_{l+1} = w_l + \frac{1}{n}\sum_{i=1}^n C_l \qty(R_i + \gamma w_l^\top \phi_i - w_l^\top \phi_{i-1}) e_{i-1}
\end{align}

\begin{align}
e_i = \sum_{k=1}^i \lambda^{i-k} \phi_k. \in \R^d
\end{align}

which is the update rule for pre-conditioned TD($\lambda$).

We also have

\begin{align}
= \indot{\psi_l}{M_1 x^{(n+1)}}
= -\indot{w_l}{\phi^{(n+1)}}.
\end{align}

This concludes our proof.

§.§ Proof of Theorem <ref>

We recall from (<ref>) that the embedding evolves according to

\begin{align}
&= Z_l + \frac{1}{n} \twohead(Z_l; P_l^{\bartd, (1)}, Q_l^{\bartd}, M^{\bartd, (1)}, P_l^{\bartd, (2)}, Q_l^{\bartd}, M^{\bartd, (2)}, W_l)\\
&= Z_l + \frac{1}{n} W_l \mqty[\attn(Z_l; P_l^{\bartd, (1)}, Q_l^{\bartd}, M^{\bartd, (1)}) \\
\attn(Z_l; P_l^{\bartd, (2)}, Q_l^{\bartd}, M^{\bartd, (2)})]
\end{align}

In this configuration, we refer to the elements in $Z_l$ as $\qty{(x^{(i)}_l, y^{(i)}_l, h^{(i)}_l)}_{i=1,\dots, n+1}$ in the following way,

\begin{align}
Z_l = \mqty[x_l^{(1)} & \dots & x_l^{(n)} & x_l^{(n+1)} \\
y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)} \\
h_l^{(1)} & \dots & h_l^{(n)} & h_l^{(n+1)}],
\end{align}

where we recall that $Z_l \in \R^{(2d + 2) \times (n+1)}, x_l^{(i)} \in \R^{2d}, y^{(i)}_l \in \R$ and $h^{(i)}_l \in \R$.

Sometimes, it is more convenient to refer to the first half and second half of $x_l^{(i)}$ separately,
by, e.g., $\nu^{(i)}_l \in \R^d, \xi_l^{(i)} \in \R^d$, i.e., $x_l^{(i)} = \mqty[\nu_l^{(i)} \\ \xi_l^{(i)}]$.
Then we have

\begin{align}
Z_l = \mqty[\nu_l^{(1)} & \dots & \nu_l^{(n)} & \nu_l^{(n+1)} \\
\xi_l^{(1)} & \dots & \xi_l^{(n)} & \xi_l^{(n+1)} \\
y_l^{(1)} & \dots & y_l^{(n)} & y_l^{(n+1)} \\
h_l^{(1)} & \dots & h_l^{(n)} & h_l^{(n+1)}].
\end{align}

We further define as shorthands

\begin{align}
X_l &\doteq \mqty[x_l^{(1)} & \dots & x_l^{(n)}] \in \R^{2d \times n}, \\
Y_l &\doteq \mqty[y_l^{(1)} & \dots & y_l^{(n)}] \in \R^{1 \times n}, \\
H_l &\doteq \mqty[h_l^{(1)} & \dots & h_l^{(n)}] \in \R^{1 \times n}.
\end{align}

Then we can express $Z_l$ as

\begin{align}
Z_l = \mqty[X_l & x_l^{(n+1)} \\
Y_l & y_l^{(n+1)} \\
H_l & h_l^{(n+1)}].
\end{align}

For the input $Z_0$, we assume $\xi_0^{(n+1)} = 0$ and $h_0^{(i)} = 0$ for $i = 1, \dots, n+1$.
All other entries of $Z_0$ are arbitrary.
We recall our definition of $M^{\bartd, (1)}, M^{\bartd, (2)}$ in (<ref>), $\qty{P_l^{\bartd, (1)}, P_l^{\bartd, (2)}, Q_l^{\bartd}, W_l}$ in (<ref>) and (<ref>).
We again express $Q_l^{\bartd}$ as

\begin{align}
M_1 \doteq& \mqty[-I_{d} & I_{d} \\
0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d}, \label{eq: two-head M_1 define}\\
B_l \doteq& \mqty[C_l^\top & 0_{d\times d} \\
0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
A_l \doteq& B_l M_1  = \mqty[-C_l^\top & C_l^\top \\
0_{d \times d} & 0_{d \times d}] \in \R^{2d \times 2d},\\
Q_l^{\bartd} \doteq& \mqty[A_l & 0_{2d \times 2} \\
0_{2 \times 2d} & 0_{2\times 2}] \in \R^{(2d+2) \times (2d+2)}.
\end{align}

We now proceed with the following claims that assist in proving our main theorem.

Claim 1. $X_l \equiv X_0 , x_l^{(n+1)} \equiv x_0^{(n+1)}, Y_l \equiv Y_0, y_l^{(n+1)} = y_0^{(n+1)}, \forall l.$

We define

\begin{align}
V_l^{(1)} &\doteq P_l^{\bartd, (1)} Z_l M^{\bartd, (1)} \qty(Z_l^\top Q_l^{\bartd} Z_l) \in \R^{(2d+2)\times(n+1)}\\
V_l^{(2)} &\doteq P_l^{\bartd, (2)} Z_l M^{\bartd, (2)} \qty(Z_l^\top Q_l^{\bartd} Z_l) \in \R^{(2d+2)\times(n+1)}.
\end{align}

Then the evolution of the embedding can be written as

\begin{align}
Z_{l+1} = Z_l + \frac{1}{n} W_l \mqty[V_l^{(1)}\\
\end{align}

By simple matrix arithmetic, we realize $W_l$ is merely summing up the $(2d+1)$-th row of $V_l^{(1)}$ and the $(2d+2)$-th row of $V_l^{(2)}$ and putting the result on its bottom row.
Thus, we have

\begin{align}
W_l \mqty[V_l^{(1)}\\
= \mqty[0_{(2d+1)\times(n+1)}\\
V_l^{(1)}(2d+1) + V_l^{(2)}(2d+2)] \in \R^{(2d+2) \times (n+1)},
\end{align}

where $V_l^{(1)}(2d+1)$ and $V_l^{(2)}(2d+2)$ respectively indicate the $(2d+1)$-th row of $V_l^{(1)}$ and the $(2d+2)$-th row of $V_l^{(2)}$.
It clearly holds according to the update rule that

\begin{align}
Z_{l+1}(1:2d+1) &= Z_l (1: 2d+1)\\
\implies X_{l+1} &= X_l;\\
x_{l+1}^{(n+1)} &= x_l^{(n+1)};\\
Y_{l+1} &= Y_l;\\
y_{l+1}^{(n+1)} &= y_l^{(n+1)}.
\end{align}

Then, we can easily arrive at our claim by a simple induction.
In light of this, we drop the subscripts of $X_l, x_l^{(i)}, Y_l$ and $y_l^{(i)}$ for all $i = 1, \dots, n+1$ and write $Z_l$ as

\begin{align}
Z_l = \mqty[X & x^{(n+1)}\\
Y & y^{(n+1)}\\
H_l & h_l^{(n+1)}].
\end{align}

Claim 2.

\begin{align}
H_{l+1} &= H_l + \frac{1}{n} (H_l + Y - \bar{Y}) X^\top A_l X\\
h_{l+1}^{(n+1)} &= h_l^{(n+1)} + \frac{1}{n} (H_l + Y - \bar{Y}) X^\top A_l x^{(n+1)},
\end{align}

where $\bar{y}^{(i)} \doteq \sum_{k=1}^i \frac{y^{(k)}}{i}$ and $\bar{Y} \doteq \mqty[\bar{y}^{(1)}, \bar{y}^{(2)}, \dots, \bar{y}^{(n)}]\in \R^{1 \times n}$.

We show how this claim holds by investigating the function of each attention head in our formulation.
The first attention head, corresponding to $V_l^{(1)}$ in claim 1, has the form

\begin{align}
P_l^{\bartd, (1)} Z_l M^{\bartd, (1)} \qty(Z_l^\top Q_l^{\bartd} Z_l).
\end{align}

We first analyze $P_l^{\bartd, (1)} Z_l M^{\bartd, (1)}$.
It should be clear that $P^{\bartd, (1)} Z_l$ selects out the $(2d+1)$-th row of $Z_l$ and gives us

\begin{align}
P_l^{\bartd, (1)} = \mqty[0_{2d\times n} & 0_{2d \times 1}\\
Y & y^{(n+1)}\\
0_{1\times n} & 0].
\end{align}

The matrix $M^{\bartd, (1)}$ is essentially computing $Y - \bar{Y}$ and filtering out the $(n+1)$-th entry when applied to $P_l^{\bartd, (1)} Z_l$.
We break down the steps here:

\begin{align}
&P_l^{\bartd, (1)} Z_l M^{\bartd, (1)}\\
=& P_l^{\bartd, (1)} Z_l \qty(I_{n+1} - U_{n+1} \diag\qty(\mqty[1& \frac{1}{2}& \dots& \frac{1}{n}]))M^{\bartd, (2)}\\
=& P_l^{\bartd, (1)} Z_l M^{\bartd, (2)} - P_l^{\bartd, (1)} Z_l U_{n+1} \diag\qty(\mqty[1& \frac{1}{2}& \dots& \frac{1}{n}])M^{\bartd, (2)}\\
=& \mqty[0_{2d\times n} & 0_{2d \times 1}\\
Y & 0\\
0_{1\times n} & 0]
- \mqty[0_{2d\times 1} & 0_{2d \times 1} & \cdots & 0_{2d \times 1} & 0_{2d \times 1}\\
y^{(1)} & \frac{1}{2} \qty(y^{(1)} + y^{(2)}) & \cdots & \frac{1}{n}\sum_{i=1}^{n} y^{(i)} & \frac{1}{n+1}\sum_{i=1}^{n+1} y^{(i)} \\
0 & 0 & \cdots & 0 & 0] M^{\bartd, (2)}\\
=& \mqty[0_{2d\times n} & 0_{2d \times 1}\\
Y & 0\\
0_{1\times n} & 0]
- \mqty[0_{2d\times n} & 0_{2d \times 1}\\
\bar{Y} & 0\\
0_{1\times n} & 0]\\
=& \mqty[0_{2d\times n} & 0_{2d \times 1}\\
Y - \bar{Y} & 0\\
0_{1\times n} & 0].
\end{align}

We then analyze the remaining product $Z_l^\top Q_l^{\bartd} Z_l$.

\begin{align}
& Z_l^\top Q_l^{\bartd} Z_l\\
=& \mqty[X^\top & Y^\top & H_l^\top\\
x^{(n+1)^\top} & y^{(n+1)^\top} &  h_l^{(n+1)^\top}]
\mqty[A_l & 0_{2d\times 1} & 0_{2d\times 1}\\
0_{1\times 2d} & 0 & 0\\
0_{1\times 2d} & 0 & 0]
\mqty[X & x^{(n+1)}\\
Y & y^{(n+1)}\\
H_l & h_l^{(n+1)}]\\
=& \mqty[X^\top A_l & 0_{n\times1} & 0_{n\times1}\\
x^{(n+1)^\top} A_l & 0 & 0]
\mqty[X & x^{(n+1)}\\
Y & y^{(n+1)}\\
H_l & h_l^{(n+1)}]\\
=& \mqty[X^\top A_l X & X^\top A_l x^{(n+1)}\\
x^{(n+1)^\top} A_l X & x^{(n+1)^\top} A_l x^{(n+1)}].
\end{align}

Putting them together, we get

\begin{align}
P_l^{\bartd, (1)} Z_l M^{\bartd, (1)} \qty(Z_l^\top Q_l^{\bartd} Z_l)
&= \mqty[0_{2d\times n} & 0_{2d \times 1}\\
Y - \bar{Y} & 0\\
0_{1\times n} & 0]
\mqty[X^\top A_l X & X^\top A_l x^{(n+1)}\\
x^{(n+1)^\top} A_l X & x^{(n+1)^\top} A_l x^{(n+1)}]\\
&= \mqty[0_{2d\times n} & 0_{2d \times 1}\\
\qty(Y - \bar{Y}) X^\top A_l X & \qty(Y - \bar{Y}) X^\top A_l x^{(n+1)}\\
0_{1\times n} & 0].
\end{align}

The second attention head, corresponding to $V_l^{(2)}$ in claim 1, has the form

\begin{align}
P_l^{\bartd, (2)} Z_l M^{\bartd, (2)} \qty(Z_l^\top Q_l^{\bartd} Z_l).
\end{align}

It's obvious that $P_l^{\bartd, (2)}$ selects out the $(2d+2)$-th row of $Z_l$ as

\begin{align}
P_l^{\bartd, (2)} Z_l = \mqty[0_{(2d+1)\times n} & 0_{(2d+1)\times 1}\\
H_l & h_l^{(n+1)}].
\end{align}

Applying the mask $M^{\bartd, (2)}$, we get

\begin{align}
P_l^{\bartd, (2)} Z_l M^{\bartd, (2)} = \mqty[0_{(2d+1)\times n} & 0_{(2d+1)\times 1}\\
H_l & 0].
\end{align}

The product $Z_l^\top Q_l^{\bartd} Z_l$ is identical to the first attention head.
Hence, we see the computation of the second attention head gives us

\begin{align}
&P_l^{\bartd, (2)} Z_l M^{\bartd, (2)} \qty(Z_l^\top Q_l^{\bartd} Z_l)\\
=& \mqty[0_{(2d+1)\times n} & 0_{(2d+1)\times 1}\\
H_l & 0]
\mqty[X^\top A_l X & X^\top A_l x^{(n+1)}\\
x^{(n+1)^\top} A_l X & x^{(n+1)^\top} A_l x^{(n+1)}]\\
=& \mqty[0_{(2d+1)\times n} & 0_{(2d+1)\times 1}\\
H_l X^\top A_l X & H_l X^\top A_l x^{(n+1)}].
\end{align}

Lastly, the matrix $W_l$ combines the output from the two heads and gives us

\begin{align}
W_l \mqty[P_l^{\bartd, (1)} Z_l M^{\bartd, (1)} \qty(Z_l^\top Q_l^{\bartd} Z_l)\\
P_l^{\bartd, (2)} Z_l M^{\bartd, (2)} \qty(Z_l^\top Q_l^{\bartd} Z_l)]
= \mqty[0_{(2d+1)\times n} & 0_{(2d+1)\times 1}\\
\qty(H_l + Y - \bar{Y})X^\top A_l X & \qty(H_l + Y - \bar{Y}) X^\top A_l x^{(n+1)}].
\end{align}

Hence, we obtain the update rule for $H_l$ and $h_l^{(n+1)}$ as

\begin{align}
H_{l+1} &= H_l + \frac{1}{n} (H_l + Y - \bar{Y}) X^\top A_l X\\
h_{l+1}^{(n+1)} &= h_l^{(n+1)} + \frac{1}{n} (H_l + Y - \bar{Y}) X^\top A_l x^{(n+1)}
\end{align}

and claim 2 has been verified.

Claim 3.

\begin{align}
h_{l+1}^{(i)} = \indot{M_1 x^{(i)}}{\frac{1}{n}\sum_{j=0}^l B_i^\top M_2 X (H_j + Y - \bar{Y})^\top},
\end{align}

for $i = 1, \dots, n+1$, where $M_2 = \mqty[I_{d} & 0_{d\times d}\\ 0_{d\times d} & 0_{d\times d}]$.

Following claim 2, we unroll $H_{l+1}$ as

\begin{align}
H_{l+1} &= H_l + \frac{1}{n} (H_l + Y - \bar{Y}) X^\top A_l X\\
H_l &= H_{l-1} + \frac{1}{n} (H_{l-1} + Y - \bar{Y}) X^\top A_{l-1} X\\
H_1 &= H_0 + \frac{1}{n} (H_0 + Y - \bar{Y}) X^\top A_0 X.
\end{align}

We therefore can express $H_{l+1}$ as

\begin{align}
H_{l+1} &= H_0 + \frac{1}{n}\sum_{j=0}^l (H_j + Y - \bar{Y})X^\top A_j X.
\end{align}

Recall that we have defined $A_j \doteq B_j M_1$ and assumed $H_0 = 0$.
Then, we have

\begin{align}
H_{l+1} = \frac{1}{n} \sum_{j=0}^l (H_j + Y - \bar{Y})X^\top M_2 B_j M_1 X.
\end{align}

Note that the introduction of $M_2$ here does not break the equivalence because $B_j = M_2 B_j$.
We include it in our expression for the convenience of the main proof later.

With the identical procedure, we can easily rewrite $h_{l+1}^{(n+1)}$ as

\begin{align}
= \frac{1}{n} \sum_{j=0}^l (H_j + Y - \bar{Y})X^\top M_2 B_j M_1 x^{(n+1)}.
\end{align}

In light of this, we define $\psi_0 \doteq 0$, and for $l = 0, \dots$

\begin{align}
\psi_{l+1} = \frac{1}{n} \sum_{j=0}^l B_j^\top M_2 X (H_j + Y - \bar{Y})^\top \in \R^{2d}.
\end{align}

We then can write

\begin{align}
h_{l+1}^{(i)} = \indot{M_1 x^{(i)}}{\psi_{l+1}} \label{eq: h inner prod}
\end{align}

for $i = 1, \dots, n + 1$, which is the claim we made.

Claim 4. The bottom $d$ elements of $\psi_l$ are always 0, i.e., there exists a sequence $\qty{w_l \in \R^d}$ such that we can express $\psi_l$ as

\begin{align}
\psi_l = \mqty[w_l \\ 0_{d \times 1}]. \label{eq: avg theta_l sparsity}
\end{align}

for all $l = 0,1,\dots, L$.

Since our $B_j$ here is identical to the proof of Theorem <ref> in <ref> for $j = 0, 1, \dots$, 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.

\begin{align}
&\indot{\psi_{l+1}}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} + \frac{1}{n}\indot{B_l^\top M_2 X (H_l + Y - \bar{Y})^\top}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} + \frac{1}{n}\sum_{i=1}^n \indot{B_l^\top M_2 x^{(i)} (h_l^{(i)} + y^{(i)} - \bar{y}^{(i)})}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n \indot{B_l^\top M_2 x^{(i)} \qty(\indot{\psi_l}{M_1 x^{(i)}} + y^{(i)} - \bar{y}^{(i)})}{M_1 x^{(n+1)}} \explain{By \eqref{eq: h inner prod}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{B_l^\top \mqty[\nu^{(i)}\\0_{d\times 1}]
\qty( \indot{\psi_l}{\mqty[-\nu^{(i)} + \xi^{(i)}\\ 0_{d\times 1}]} + y^{(i)} - \bar{y}^{(i)})}{M_1 x^{(n+1)}}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_l \nu^{(i)}\\0_{d\times 1}]
\qty( y^{(i)} - \bar{y}^{(i)} +  w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)} )}{M_1 x^{(n+1)}} \explain{By Claim 4}\\
=&\indot{\psi_l}{M_1 x^{(n+1)}} +
\frac{1}{n} \sum_{i=1}^n
\indot{\mqty[C_l \nu^{(i)} \qty(y^{(i)} - \bar{y}^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})\\0_{d\times 1}]}
{M_1 x^{(n+1)}}\\
\end{align}

This means

\begin{align}
\indot{w_{l+1}}{\nu^{(n+1)}}
= \indot{w_l}{\nu^{(n+1)}}
+ \frac{1}{n} \sum_{i=1}^n
\indot{C_l \nu^{(i)} \qty(y^{(i)} - \bar{y}^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)})}{
\nu^{(n+1)}}.
\end{align}

Since the choice of the query $\nu^{(n+1)}$ is arbitrary, we get

\begin{align}
w_{l+1} = w_l + \frac{1}{n}\sum_{i=1}^n C_l \qty(y^{(i)} - \bar{y}^{(i)} + w_l^\top \xi^{(i)} - w_l^\top \nu^{(i)}) \nu^{(i)}.
\end{align}

In particular, when we construct $Z_0$ such that $\nu^{(i)} = \phi_{i-1}, \xi^{(i)} = \phi_i$ and $y^{(i)} = R_i$, we get

\begin{align}
w_{l+1} = w_l + \frac{1}{n}\sum_{i=1}^n C_l \qty(R_i - \bar{r}_i + w_l^\top \phi_i - w_l^\top \phi_{i-1}) \phi_{i-1}
\end{align}

which is the update rule for pre-conditioned average reward TD learning.
We also have

\begin{align}
= \indot{\psi_l}{M_1 x^{(n+1)}}
= -\indot{w_l}{\phi^{(n+1)}}.
\end{align}

This concludes our proof.

§ EVALUATION TASK GENERATION

To generate the evaluation tasks used to meta-train our transformer in Algorithm <ref>, we utilize Boyan's chain, detailed in Figure <ref>.
we make some minor adjustments to the original Boyan's chain in 11 to make it an infinite horizon chain.

Recall that an evaluation task is defined by the tuple $(p_0, p, r, \phi)$.
We consider Boyan's chain MRPs with $m$ states.
To construct $p_0$,
we first sample a $m$-dimensional random vector uniformly in $[0, 1]^m$ and then normalize it to a 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 as a random matrix taking value in $\R^{m \times m}$.
For simplifying presentation,
we use both $p(s, s')$ and $p(s' | s)$ to denote probability of transitioning to $s'$ from $s$.
In particular, for $i = 1, \dots, m-2$,
we set $p(i, i+1) = \epsilon$ and $p(i, i+2) = 1 - \epsilon$,
with $\epsilon$ sampled uniformly from $(0, 1)$.
For the last two states, we have
$p(m | m-1) = 1$
and $p(\cdot | m)$ is a random distribution over all states.
Each element of the vector $r \in \R^m$ and the matrix $\phi \in \R^{d \times m}$ are sampled i.i.d. from a uniform distribution over $[-1, 1]$.
The overall task generation process is summarized in Algorithm <ref>.
Almost surely,
no task will be generated twice.
In our experiments in the main text,
we use Boyan Chain MRPs which consist of $m = 10$ states each with feature dimension $d = 4$.

Boyan's Chain of $m$ States

Boyan Chain MRP and Feature Generation (Non-Representable)
[1]
Input: state space size $m = \abs{\fS}$, feature dimension $d$
$s \in \fS$
$\phi(s) \sim \uniform{(-1, 1)^d}$   // feature
$p_0 \sim \uniform{(0, 1)^m}$   // initial distribution
$p_0 \gets p_0 /\sum_s p_0(s)$
$r \sim \uniform{(-1, 1)^m}$   // reward function
$p \gets 0_{m \times m}$   // transition function
$i = 1, \dots, m - 2 $
$\epsilon \sim \uniform{(0, 1)}$
$p(i, i+1)$ $\gets$ $\epsilon$
$p(i, i+2)$ $\gets$ $1 - \epsilon$
$p(m-1, m)$ $\gets$ 1
$z$ $\gets$ $\uniform{(0, 1)^m}$
$z$ $\gets$ $z/\sum_s z(s)$
$p(m, 1:m)$ $\gets$ $z$
Output: MRP $(p_0, p, r)$ and feature map $\phi$

Representable Value Function.
With the above sampling procedure,
there is no guarantee that the true value function $v$ is always representable by the features.
In other words,
there is no guarantee that there exists a $w \in \R^d$ satisfying $v(s) = \indot{w}{\phi(s)}$ for all $s \in \fS$.
Most of our experiments use this setup.
It is,
also beneficial sometimes to work with evaluation tasks where the true value function is guaranteed to be representable.
Algorithm <ref> achieves this by randomly generating a $w_*$ first and compute $v(s) \doteq \indot{w_*}{\phi(s)}$.
The reward is then analytically computed as $r \doteq (I_m - \gamma p) v$.
We recall that in the above we regard $p$ as a matrix in $\R^{m \times m}$.

Boyan Chain MRP and Feature Generation (Representable)
[1]
Input: state space size $m = \abs{\fS}$, feature dimension $d$, discount factor $\gamma$
$w^* \sim \uniform{(-1, 1)^d}$   // ground-truth weight
$s \in \fS$
$\phi(s) \sim \uniform{(-1, 1)^d}$
// feature
$v(s) \gets \indot{w^*}{\phi(s)}$   // ground-truth value function
$p_0 \sim \uniform{(0, 1)^m}$   // initial distribution
$p_0 \gets p_0 /\sum_s p_0(s)$
$p \gets 0_{m \times m}$   // transition function
$i = 1, \dots, m - 2 $
$\epsilon \sim \uniform{(0, 1)}$
$p(i, i+1)$ $\gets$ $\epsilon$
$p(i, i+2)$ $\gets$ $1 - \epsilon$
$p(m-1, m)$ $\gets$ 1
$z$ $\gets$ $\uniform{(0, 1)^m}$
$z$ $\gets$ $z/\sum_s z(s)$
$p(m, 1:m)$ $\gets$ $z$
$r \gets \qty(I_m - \gamma p)v$   // reward function
Output: MRP $(p_0, p, r)$ and feature map $\phi$

§ ADDITIONAL EXPERIMENTS WITH LINEAR TRANSFORMERS

§.§ Experiment Setup

We use Algorithm <ref> as $d_\text{task}$ for the experiments in the main text with Boyan's chain of 10 states.
In particular,
we consider a context of length $n=30$, feature dimension $d=4$, and utilize a discount factor $\gamma = 0.9$. In Section <ref>, we consider a 3-layer transformer ($L=3$), but additional analyses on the sensitivity to the number of transformer layers ($L$) and results from a larger scale experiment with $d=8, n=60,$ and $|\mathcal{S}|=20$ are presented in <ref>. We also explore non-autoregressive (i.e., "sequential") layer configurations in <ref>.

When training our transformer, we utilize an Adam optimizer  [25] with an initial learning rate of $\alpha = 0.001$, and weight decay rate of $1\times 10^{-6}$. $P_0$ and $Q_0$ are randomly initialized using Xavier initialization with a gain of $0.1$.
We trained our transformer on $k=4000$ different evaluation tasks.
For each task, we generated a trajectory of length $\tau = 347$, resulting in $\tau - n -2 = 320$ transformer parameter updates.

Since the models in these experiments are small ($\sim 10$ KB), we did not use any GPU's during our experiments. We trained our transformers on a standard Intel i9-12900-HK CPU and training each transformer took $\sim20$ minutes.

For implementation[The code will be made publicly available upon publication.], we used NumPy [21] to process the data and construct Boyan's chain, PyTorch [5] to define and train our models, and Matplotlib [22] plus SciencePlots [19] to generate our figures.

§.§.§ Trained Transformer Element-wise Convergence Metrics

To visualize the parameters of the linear transformer trained by Algorithm <ref>, we report element-wise metrics. For $P_0$, we report the value of its bottom-right entry, which, as noted in (<ref>), should approach one if the transformer is learning to implement TD. The other entries of $P_0$ should remain close to zero. Additionally, we report the average absolute value of the elements of $P_0$, excluding the bottom-right entry, to check if these elements stay near zero during training.

For $Q_0$, we recall from (<ref>) that if the transformer learned to implement normal batch TD, the upper-left $d\times d$ block of the matrix should converge to some $-I_d$, while the upper-right $d\times d$ block (excluding the last column) should converge to $I_d$. To visualize this, we report the trace of the upper-left $d\times d$ block, and the trace of the upper-right $d\times d$ block (excluding the last column). The rest of the elements of $Q_0$ should remain close to 0, and to verify this, we report the average absolute value of the entries of $Q_0$, excluding the entries that were utilized in computing the traces.

Since, $P_0$ and $Q_0$ are in the same product in (<ref>) we sometimes observe during training that $P_0$ converges to $-P_0^{\td}$ and $Q_0$ converges to $-Q_0^{\td}$ simultaneously. When visualizing the matrices,  we negate both $P_0$ and $Q_0$ when this occurs.

It's also worth noting that in Theorem <ref> we prove a $L$-layer transformer parameterized as in (<ref>) with $C_0 = I_d$ implements $L$ steps of batch TD exactly with a fixed update rate of one. However, the transformer trained using Algorithm <ref> could learn to perform $\td$ with an arbitrary learning rate ($\alpha$ in (<ref>)). Therefore, even if the final trained $P_0$ and $Q_0$ differ from their constructions in (<ref>) by some scaling factor, the resulting algorithm implemented by the trained transformer will still be implementing TD.
In light of this,
we rescale $P_0$ and $Q_0$ before visualization.
In particular, we divide $P_0$ and $Q_0$ by the maximum of the absolute values of their entries respectively, such that they both stay in the range $[-1, 1]$ after rescaling.

§.§.§ Trained Transformer and Batch TD Comparison Metrics

To compare the transformers with batch TD we report several metrics following 53, 3.
Given a context $C \in \R^{(2d+1) \times n}$ and a query $\phi \in \R^d$,
we construct the prompt as

\begin{align}
Z^{(\phi, C)} \doteq \mqty[C & \mqty[\phi \\ 0_{d \times 1} \\ 0]].
\end{align}

We will suppress the context $C$ in subscript when it does not confuse.
We use $Z^{(s)} \doteq Z^{(\phi(s))}$ as shorthand.

We use $d_p$ to denote the stationary distribution of the MRP with transition function $p$
and assume the context $C$ is constructed based on trajectories sampled from this MRP.
Then, we can define $v_{\theta} \in \R^{\abs{\fS}}$, where $v_{\theta}(s) \doteq \tf_L(Z_0^{(s)}; \theta)$ for each $s \in \mathcal{S}$.
$v_\theta$ is then the value function estimation induced by the transformer parameterized by $\theta \doteq \qty{(P_l, Q_l)}$ given the context $C$.
In the rest of the appendix,
we will use $\theta_\tf$ as the learned parameter from Algorithm <ref>.
As a result, $v_\tf \doteq v_{\theta_\tf}$ denotes the learned value function.

We define $\theta_\td \doteq \qty{(P_l^\td, Q_l^\td)}_{l=0,\dots, L-1}$ with $C_l = \alpha I$ (see (<ref>)) and

\begin{align}
v_\td(s) \doteq \tf_L(Z_0^{(s)}; \theta_\td).
\end{align}

In light of Theorem <ref>,
$v_\td$ is then the value function estimation obtained by running the batch TD algorithm (<ref>) on the context $C$ for $L$ iterations,
using a constant learning rate $\alpha$.

We would like to compare the two functions $v_\tf$ and $v_\td$ to future examine the behavior of the learned transformers.
$v_\td$ is not well-defined yet because it still has a free parameter $\alpha$,
the learning rate.

[53] resolve a similar issue in the in-content regression setting
via using a line search to find the (empirically) optimal $\alpha$.

Inspired by [53],
we also aim to find the empirically optimal $\alpha$ for $v_\td$.

We recall that $v_\td$ is essentially the transformer $\tf_L(Z_0^{(s)}; \theta_\td)$ with only 1 single free parameter $\alpha$.
We then train this transformer with Algorithm <ref>.
We observe that $\alpha$ quickly converges and use the converged $\alpha$ to complete the definition of $v_\td$.
We are now ready to present different metrics to compare $v_\tf$ and $v_\td$.
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,

\begin{align}
\text{VD}(v_\tf, v_\td) \doteq \norm{v_\tf - v_\td}_{d_p}^2,
\end{align}

We recall that $d_p \in \R^\ns$ is the stationary distribution of the MRP and the weighted $\ell_2$ norm is defined as $\norm{v}_d \doteq \sqrt{\sum_s v(s)^2 d(s)}$.

Implicit Weight Similarity (IWS).
We recall that $v_\td$ is a linear function, i.e., $v_\td(s) = \indot{w_L}{\phi(s)}$ with $w_L$ defined in Theorem <ref>.
We refer to this $w_L$ as $w_\td$ for clarity.
The learned value function $v_\tf$ is, however, not linear even with linear transformer.
Following 3,
we compute the best linear approximation of $v_\tf$.
In particular,
given a context $C$,
we define

\begin{align}
w_\tf \doteq \arg\min_w \norm{\Phi w - v_\tf}_{d_p}.
\end{align}

Here $\Phi \in \R^{\ns \times d}$ is the feature matrix,
each of which is $\phi(s)^\top$.
Such a $w_\tf$ is referred to as implicit weight in 3.
Following 3,
we define

\begin{align}
\text{IWS}(v_\tf, v_\td) \doteq d_\text{cos}(w_\tf, w_\td)
\end{align}

to measure the similarity between $w_\tf$ and $w_\td$.
Here $d_\text{cos}(\cdot, \cdot)$ computes the cos similarity between two vectors.

Sensitivity Similarity (SS). Recall that $v_{\tf}(s) = \tf_L(Z_0^{(s)}; \theta_\tf)$ and $v_{\td}(s) = \tf_L(Z_0^{(s)}; \theta_{\td})$.
In other words, given a context $C$,
both $v_\tf(s)$ and $v_\td(s)$ are functions of $\phi(s)$.
Following 53,
we then measure the sensitivity of $v_\tf(s)$ and $v_\td(s)$ w.r.t. $\phi(s)$.
This similarity is easily captured by gradients.

In particular,
we define

\begin{align}
\text{SS}(v_\tf, v_\td) \doteq \sum_s d_p(s) d_\text{cos}\left(\eval{\nabla_{\phi} \tf_L(Z_0^{(\phi)}; \theta_\tf)}_{\phi = \phi(s)}, \eval{\nabla_{\phi} \tf_L(Z_0^{(\phi)}; \theta_\td)}_{\phi = \phi(s)}\right).
\end{align}

it trivially holds that

\begin{align}
w_\td = \eval{\nabla_{\phi} \tf_L(Z_0^{(\phi)}; \theta_\td)}_{\phi = \phi(s)}.
\end{align}

We note that the element-wise converge of learned transformer parameters (e.g., Figure <ref>) is the most definite evidence for the emergence of in-context TD.
The three metrics defined in this section are only auxiliary when linear attention is concerned.
That being said,
the three metrics are important when nonlinear attention is concerned.

§.§ Autoregressive Linear Transformers with $L=1,2,3,4$ Layers

In this section,
we present the experimental results for autoregressive linear transformers with different numbers of layers.
In Figure <ref>, we present the element-wise convergence metrics for autoregressive transformers with $L=1,2,4$ layers.
The plot with $L = 3$ is in Figure <ref> in the main text.

We can see that for the $L=1$ case, $P_0$ and $Q_0$ converge to the construction in Corollary <ref>,