Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining
Authors: Licong Lin
Yu Bai
Song Mei
Venue: ICLR 2024 poster (2024)
Keywords: transformers
in-context learning
reinforcement learning
learning theory
Primary area: learning theory
TL;DR: We prove that transformers can provably implement various reinforcement learning algorithms in context, and learn them through supervised pretraining.

Abstract: Large transformer models pretrained on offline reinforcement learning datasets have demonstrated remarkable in-context reinforcement learning (ICRL) capabilities, where they can make good decisions when prompted with interaction trajectories from unseen environments. However, when and how transformers can be trained to perform ICRL have not been theoretically well-understood. In particular, it is unclear which reinforcement-learning algorithms transformers can perform in context, and how distribution mismatch in offline training data affects the learned algorithms. 

This paper provides a theoretical framework that analyzes supervised pretraining for ICRL. This includes two recently proposed training methods --- algorithm distillation and decision-pretrained transformers. First, assuming model realizability, we prove the supervised-pretrained transformer will imitate the conditional expectation of the expert algorithm given the observed trajectory. The generalization error will scale with model capacity and a distribution divergence factor between the expert and offline algorithms. Second, we show transformers with ReLU attention can efficiently approximate near-optimal online reinforcement learning algorithms like LinUCB and Thompson sampling for stochastic linear bandits, and UCB-VI for tabular Markov decision processes. This provides the first quantitative analysis of the ICRL capabilities of transformers pretrained from offline trajectories.

Published as a conference paper at ICLR 2024
TRANSFORMERS AS DECISION MAKERS: PROVABLE IN-CONTEXT REINFORCEMENT LEARNING VIA SU- PERVISED PRETRAINING
Licong Lin
Yu Bai
University of California, Berkeley edu
Salesforce Al Research yu.bai@salesforee. com
Song Mei
University of California, Berkeley songmei @berkeley.ede
ABSTRACT
Large transformer models pretrained on offline reinforcement learning datasets have demonstrated remarkable in-context reinforcement learning (ICRL) capabil- ities, where they can make good decisions when prompted with interaction tra- jectories from unseen environments. However, when and how transformers can trained to perform ICRL have not been theoretically well-understood. In par- ticular, it is unclear which cerfforreemnt-learnnn algorithms transformers perform in context, and how distribution mismatch in offline training data affects the learned algorithms. This paper provides theoretical framework that analyzes supervised pretraining for ICRL. This includes two recently proposed training methods algorithm distillation and decision-pretraineed transformers. First, as- suming model realizability, we prove the supervised supervised-pretrained -peetraind transformer will imitate the conditional expectation of the expert algorithm given the observed trajectory. The generalization eттоr will scale with model capacity and distri- bution divergence factor between the expert and offline algorithms. Second, we show transformers with ReI.U attention can efficiently approximate near-optimal online reinforcement learning algorithms like LinUCB and Thompson sampling for stochastic linear bandits, and UCB-VI for tabular Markov decision processes. This provides the first quantitative analysis of the ICRL capabilities of transform- pretrained from offline trajectories.
INTRODUCTION
The transformer architecture (Vaswani et al., 2017) for sequence modeling has become a key weapon for modern artificial intelligence, achieving success language (Devlin et al., 2018; Brown al., 2020; OpenAI, 2023) and vision (Dosovitskiy et al., 2020). Motivated by these advances, the search community has actively explored how to best harness transformers for reinforcement learning (RL) (Chen al., 2021; Janner et al., 2021; Lee et al., 2022; Reed et al., 2022: Laskin et al., 2022; Lee et al., 2023; Yang et al., 2023). While promising empirical results have been demonstrated, the theoretical understanding of transformers for RL remains limited.
This paper provides theoretical insights into in-context reinforcement learning (ICRL)-an emerg- ing approach that utilizes learning in newly encountered environments. In ICRL, the model takes as input the current state and past interaction history with the environment (the context), and outputs an action. The key hy- pothesis in ICRL is that pretrained transformers can act RL algorithms, progressively improving their policy based on past observations. Approaches such as Algorithm Distillation (Laskin et al., 2022) and
models like transformers to perform reinforcement
Transformers (Lee et al., 2023) have demonstrated early successes, finding that supervised pretraining can produce good ICRL performance. However, many concrete theoretical questions remain open about the ICRL capabilities of transformers, including but not limited to (1) what RL algorithms can transformers implement in-context; (2) what performance guarantees (e.g. regret bounds) can such transformers achieve when used iteratively online RI.

---

Published as a conference paper at ICLR 2024
algorithm; and (3) when can supervised pretraining find such good transformer. Specifically, this paper investigates the following open question:
How can supervised pretraining on Transformers learn in-context reinforcement learning?
In this paper, we initiate a theoretical study of the ICRL capability of transformers under supervised pretraining to address the open questions outlined above. We show that (1) Transformers can im- plement prevalent RI. algorithms, including LinUCB and Thompson sampling for stochastic linear bandits, and UCB-VI for tabular Markov decision processes; (2) The algorithms learned by trans- formers achieve near-optimal regret bounds in their respective settings; (3) Supervised pretraining find such algorithms long the sample size scales with the covering number of transformer class and distribution ratio between expert and offline algorithms.
Summary of contributions and paper outline
We propose a general framework for supervised pretraining approaches to meta-reiinforeement learning (Section 2). This framework encompasses existing methods like Algorithm Distillation (Laskin et al., 2022), where the expert and context algorithms are identical, as well as Decision- Pretrained Transformers (Lee et al., 2023), where the expert generates optimal actions for the MDP. also includes approximate DPT variants where the expert estimates optimal actions from full interaction trajectories.
pretrained We prove that the puetraised-peetrannee transformer will imitate the conditional expectation of the expert algorithm given the observed trajectory (Section 3). The generalization error scales with both model capacity and distribution ratio measuring divergence between the expert algo- rithm and the algorithm that generated offline trajectories.
We demonstrate that transformers can effectively approximate several near-optimal reinforce- ment learning algorithms by taking observed trajectories as context inputs (Section 4). Specif- ically, we show transformers can approximate LinUCB (Section 4.1) and Thompson sampling algorithms (Section 4.2) for stochastic linear bandit problems, and UCB-VI (Section 4.3) for tabular Markov decision processes. Combined with the generalization error bound from su- pervised pretraining and regret bounds of these RL algorithms, this provides regret bounds for pnpervisedd-reerainee transformers.
Preliminary experiments validate that transformers can perform ICRL in our setup (Section 5).
Technically, we prove efficient approximation of LinUCB by showing transformers can imple- ment accelerated gradient descent for solving ridge regression (Appendíx E.4), enabling fewer attention layers than the vanilla gradient descent approach in Bai et al. (2023). To enable efficient Thompson sampling implementation, we prove transformers can compute matrix square roots through the Pade decomposition (Appendix F.3). These approximation results are interesting in their own right.
Our work is intimately related the lines of work on neta-reinforcement learning, in-context learning, transformers for decision-makee and the approximation theory of transform- ers. Due to limited space, we discuss these related works Appendix A.
Related work
2 FRAMEWORK FOR IN-CONTEXT REINFORCEMENT LEARNING
Let M be the space of malssinnmmaingg environments, where each environment M M shares the same number of rounds T and state-action-rewaard spaces {S,,A, R₂ Reere. Each M (TT JR heem has its own transition model T/ S, x A, + A(S.+)) (with So. Ao {0} so (+) E ((S)) gives the initial state distribution) and reward functions R S, A, A(R,)
We equip M with a distribution A E A(M), the environment prior. While this setting is general, later give concrete examples taking M as T' rounds of bandits K episodes of H-step MDPs with T = K.H.
Distributions of offline trajectories We denote a partial interaction trajectory, consisting of ob- served state-action-rewarl tuples, by D. {(0), (i,rr),...... (st, des r:)) T, (S, x A, x R.) and write D = Dr for short. A standard algorithm Alg maps a partial trajectory D. E T and state 8; S, to distribution over the actions A 8t) ((A)). Given an environment
2

---

Published as a conference paper at ICLR 2024
M and algorithm Alg. the distribution over full trajectory Dr is fully specified:
PM (Dr) = II' (4 9ii (relss.))
In supervised pretraining, we use a context algorithm Algo (which we also refer to as the offline algorithm) collect the offline trajectories Dr. For each trajectory Dr, also assume access to expert actions a = (a, E AtkeT ~ Alg Dr, r M), sampled from an expert algorithm Alg Tr M mm ((A)). This expert could omnisciently observe the full trajectory Dr and environment M to recommend actions. Let Dr Dr U {a} have
be the augmented trajectory. Then we
AlA. AlRE (Dr) = P/ (DT) (DT) Alg'((a:) Dr, M). M
We denote PA Alss as the joint distribution of (M, Dr) where M and Dy Al and AAlgo the joint distribution of (M, Dr) Dr) where M ~ A and Dr Algo
Three special cases of expert algorithms We consider three special cases of the expert algorithm Algn, corresponding to three supervised pretraining setups:
(a) Algorithm distillation (Laskin al., 2022). The algorithm depends only the partial trajectory
D.-1 and current state sy: ( M) = Alg (-,, St). For example, Algg could bandit algorithm like the Upper Confidence Bound (UCB).
(b) Decision pretrained transformer (DPT) (Lee et al., 2023). The algorithm depends on the environment M and the current state Sy: Al M) (-|s: M). For example, Alga could output the optimal action of in state for environment M.
(c) Approzimate DPT. The algorithm depends on the full trajectory Dr but not the environment M: Algg(-|Dr, M) A- For example, Algg could estimate the optimal action äi from the entire trajectory Dr.
For any expert algorithm Algg, define its reduced algorithm where the (-th step is
Alg D 81) AAA [Al''s (- Dr, M) D, 82]-
The expectation the right is over A (Dr, M 1. 8x) A(M) P (Dr) /TPM Also (Dt-1, 8t). Note that the reduced expert algorithm Algg generally depends the context algorithm Alga- How- for cases (a) and (b), Alg is independent of the context algorithm Algo- Furthermore, in case (a), we have Alg = Alg
Transformer architecture written as an input matrix H [h...... . hv] I8 DxN, where each h, is column H (also token). Throughout this paper, we define (t)) = ReLU(t) = max {t, 0} as the standard relu activation function.
We consider sequence of N input vectors [h:), R., compactly
Definition 1 (Masked attention layer). A masked attention layer with M heads is denoted as (-)no(-) with parameters = ((Vm, m. Qm, Km)) me[M]] C DRD On any input sequence H IR DxN, have H Attnn [hi.... hNl E RDN, where
= [Attmo (H)) h; MM 1 a((Qm Kmh,)) V ERD
We remark that the use of ReLU attention layers for technical reasons. In practice, both ReLU attention and softmax attention layers should perform well. Indeed, several studies have shown that ReLU transformers achieve comparable performance to softmax transformers across a variety of tasks (Wortsman et al., 2023; Shen et al., 2023; Bai et al., 2023).
Definition 2 (MLP layer). An MLP layer with hidden dimension D' is denoted as MLPo(.) with parameters 0 = (W,, W2) E DD'xD x DDDD On any input sequence H E RDxN we have H = MLP e H) = [i,..., hv] E DxN where
h h; W2 (WW,)))) ERD

---

Published as a conference paper at ICLR 2024
We next define L-layer decoder-based transformers. Each layer consists of masked attention layer (see Definition 1) followed by an MLP layer (see Definition 2) and a clip operation.
Definition 3 (Decoder-basee Transformer). An L-layer decoder-based transformer, denoted as (-).), is composition of L masked attention layers, each followed by MLP layer and clip operation: TF
H &R DDNN, where H is defined iteratively by taking H (0) clipp (H) E IR DxN and for EE [
H clipR (MPP
(Attn (Hil-11))) E EDDNN clipR =
(0l1:t) ((-L)) consists of o = {(V!) Qm K! Jmees C
Above, the parameter
RDDD and 0 (W(1) w E RD''x x DDD' We define the parameter class of trans-
formers L,,,,M, B. == 10- (0 0 0 |o1 S B}, where the norm of transformer
TE denoted as
M
|0| :- max max [E[L] MEM
{ 1 Ilop:
|
lop
+
Σ
| op + [| llop + l lop) (I)
(0)
We introduced clipped operations in transformers for technical reasons. For brevity, we will write TF TF when there is no ambiguity. We will set the clipping value R to be sufficiently large so that the clip operator does not take effect in any of our approximation results.
Algorithm induced by Transtormers We equip the transformer with embedding mapping h: UrerryS (UETT x R.) R. This assigns any state & E St a D-dimensional embedding vector h(s;) E R², and any action-reward pair (at, (1),rt) E A, x R2 D-dimensional embedding n(a,,rz) R.D The embedding function h should encode the time step of the state, action, and re- ward. With abuse of notation, we denote h( (,-1,8 - We define a concatenation operator cat R - R.xx* that concatenates its inputs cat(h,..... ,,) [h,..... hx] in most examples, but it could also insert special tokens at certain positions (in MDPs we add an additional token at the end of each episode). For a partial trajec- tory and current state (t), we input H cat (h(51), h( r1),............ ......... (( RDx into the transformer. This produces output H ( same shape as H. To extract a distribution over the action space A with [Adl = A actions, we assume fixed linear extraction mapping
[b(81), (,(s)), h(a,,,1).... "ia,,r,)),,,,.,__1 rt-1). 1), h(s))]
E h1 h2 .... h22 h-1] with the
Ax The induced algorithm is then defined
Alga Dz-1, st) softmax(A A h)1). The overall algorithm induced by the transformer is:
A- D 8e) soft A (
(2)
We will always choose a proper concatenation operator cat in examples, so that in the pretraining phase, all the algorithm outputs ( ( along the trajectory be computed in single forward propagation.
3 STATISTICAL ANALYSIS OF SUPERVISED PRETRAINING
In supervised pretraining, we are given n i.idd offline trajectories (Di = (si ai, ri... s ar rf));_ rid FAlE from the interaction of M ~idd A with offline algorithm Algo Given an expert algorithm Algg. we augment each trajectory Dj, by (बद mnd ( Supervised pretraining maximizes the logelikelihood over the algorithm class {Algo Jeee
0
arg max DEE
log Algo ((
n
This section discusses the statistical properties of the algorithm learned via supervised pretraining.
3.1
MAIN RESULT
Our main result demonstrates that the algorithm maximizing the supervised pretraining loss will imitate (- of the expert algorithm Alg, given the observed trajectory. The imitation error bound will scale with the covering number of the algorithm class and distribution ratio factor, defined follows.
EM~A.Dy~Alg, (-Ir. D st). the conditional expectation
4

---

Published as a conference paper at ICLR 2024
Definition 4 (Covering number). For class of algorithus {Alga. )) say Og S θ is p-cover of è if u is a finite set such that for any 8 E θ, there exists 8g E Ou such that
,,, S for all Di-1, St, E [T]. The covering number Np) is the minimal cardinality of Og such that Og is a p-cover of O
log Alge, (-) sx) log
Definition (Distribution ratio). We define the distribution ratio of two algorithms Alg, by
Alg Do-1, = =1+ x² : II
RAIE; Als,
Algz D 8,)
Our main result requires the realizability assumption of algorithm class {Algolaee with respect to the conditional expectation of the expert algorithm.
Assumption A (Approximate realizability). There exists 8" E and Ereal > 0 such that for all
t t [T].
Alg (Te 1, Ereal-
log E
M Drow Also Moz Maz
Alge- (atl 1,)),88)
4
We aim to bound the performance gap between Algg and Algg in terms of expected cumulative rewards, where the expected cumulative reward is defined as
( EMM~A [RM.Aig(T]]
Hi (T) E. D Et= rz].
An intermediate step of the result is controlling the expected Hellinger distance between algo- rithms, where for distributions p. q. have (p)
/((z)) )d..
Theorem 6 (Performance gap between expected cumulative rewards). Let Assumption A hold and ler be solution to Eq. (3). Take R AA as defined in Definition 5, and
Ne NNo((nT) -2) as defined Definition 4. Then for some universal constant c > 0, with probability at least - 8, we have
Σ
T
Du
Algal si)1, & CIVR(
(
log [Ne T/5)
+ V VEreal
Dr PARE
(5)
Further assume that Iral S 1 almost surely. Then with probability at least 1 - 6, the difference of the expected cumulative rewards between Alga and Alg satisfies
log [ee T/8
RA,A (T) RA,NEE (T) S eT
+ V Ereal
(6)
n
The proof Theorem 6 contained in Section D.1.
We remark that when the expectation on the left-hand-side of (5) is with respect to the measure Also, standard MIE analysis will provide a bound without the distribution ratio factor R. RR Algo, in the right-hand side. The distribution ratio factor arises from distribution shift between trajectories generated by the expert algorithm Algg versus the context algorithm Alg, In addition, it should be noted that the result in Theorem 6 holds generally provided Assumption A is satisfied, which does not require that the algorithm class is induced by transformers.
3.2 IMPLICATIONS IN SPECIAL CASES
Algorithm Distillation When we set Algg = Algo- the supervised pretraining approach corre- sponds to the Algorithm Distillation method introduced in Laskin al. (2022). In this case, suffices to set ह' a' for every pretraining trajectory, eliminating the need to sample additional expert actions. The conditional expectation of the expert algorithm is given by Algg Algo- and the distribution ratio RAKE Alg, = 1. Under these conditions, Theorem 6 ensures that Alg, imitates Alg, with reward difference bounded by
(T)) S 2T (1 log (No T/
A RR
+ /Ereal
the context algorithm Alg, does not perform well, we cannot expect the learned algorithm Alga to have good performance, regardless of the number of offline trajectories.
5

---

Published as a conference paper at ICLR 2024
Decision Pretrained Transformer When set Alg': Alg (st, M) = az to be the optimal ac- tion at time t, the supervised pretraining approach corresponds to (DPT) proposed in Lee et al. (2023). In this case. the conditional expectation of the expert algorithm (- st) = [lAlg (-1s, D st] Alg (-1, st) is the Thompson sampling al- gorithm (Lee et al., 2023, Theorem 1), which samples from the posterior distribution of the optimal action of given by P(a (M) st)
Transformers
x A(M) PM (Di-1. st). This implies that learning from optimal actions effectively learns to imitate Thompson sampling. Furthermore, the context algo- rithm is not required to perform well for the learned algorithm to be consistent with Thompson sampling. However, a high-quality context algorithm Algo may help reduce the distribution ratio R, thereby learning Thompson sampling with fewer samples.
Approximate DPT practical scenarios, the learner may not have access to the optimal action aj of the environment M during pretraining. Instead, they might rely on estimated optimal action 2 (- Dr), derived from the entire trajectory Dr. We can offer guarantee analogous to Theorem 6, provided the distribution of the estimated action closely aligns with its posterior distribution:
E DimeE (L(All' (rs,(()Drr) S Fapprox: Vt [T].
Here, Prs,a(-1Dr) represents the posterior distribution of the optimal action o; = of (M) at time t, given the observation Dr, where (M, Dr) Ale..
Let Assumption A hold and let ê be the solution to Eq. (3). Take R = RAIBTS as defined in Definition 5, and No = (N((TT))-) ) as defined in Definition 4. Assume that for each trajectory, an estimated optimal action is provided āj ~ Alg'g-|/y at each time t E [7] satisfying Eq. (7). Assume that the rewards [rel S almost surely. Then for some universal constant O, with probability at least 1 - 8. the difference of the expected cumulative rewards between Algg and Algrs satisfies
Proposition 7.
log (Ne .T/8 + V + Eapprrx (10
|R (T) - Ri A (T) S CVR .T
The proof Proposition 7 is contained in Appendix D.2.
APPROXIMATION BY TRANSFORMERS
In this section, we demonstrate the capability of transformers to implement prevalent reinforcement learning algorithms that produce near-optimal regret bounds. Specifically, we illustrate the imple- mentation of LinUCB for stochastic linear bandits in Section 4.1, Thompson sampling for stochastic linear bandits in Section 4.2, and UCB-V VI for tabular Markov decision process in Section 4.3.
4.1 LINUCB FOR LINEAR BANDITS
A1, A stochastic linear bandit environment is defined by M Ai,., Ar). For each time step t E [T]. the learner chooses an action E Rd from a set of actions A, {ar......, ar.A). which consists of A actions and may vary over time. Upon this action selection, the learner receives reward r. (at, + Here,fe;} ~ad E are zero-mean noise variables, and w' E Rd represents an unknown parameter vector. Stochastic linear bandit can be cast into our general framework by setting A, and adopting a deterministic transition where 8; transits to 8141 deterministically regardless of the chosen action.
We assume the context algorithm Alg, is the soft LinUCB (Chu et al., 2011). Specifically, for each time step t E [T], the learner estimates the parameter w using linear ridge regression Wridge, == arg min Ej=(rr - (a,, w))2 + Al|w// Subsequently, the learner calculates the upper confidence bounds for the reward of each action vik (ar, Wiidge.x) + a (aT& (AI& + 2-11 -
Finally, the learner selects an action at according to proba- bility {pt bJe(A) = softmax( / (TbEA)) for some sufficiently small r > O. Note that the soft LinUCB AlgaLinUCB(()) recovers LinUCB as 0.
We further assume the existence of constants o, bu, Ba, B,, such that the following conditions hold: | < a. b, < ||ar, 5 B,, and ||w*|12 5 BN,, for all t E [2], E [A]. Given these, the con- vAB,, (og T) + dllo((d) TB3) (dd))
fidence parameter is defined

---

Published as a conference paper at ICLR 2024
((vd). The following result shows that the soft LinUCB algorithm can be efficiently approximated by transformers, for which the proof is contained in Appendix E.4.
Theorem 8 (Approximating the soft LinUCB). Consider the embedding mapping h, extraction mapping A, and concatenation operator cat as in E.l. For any small E, > 0, there exists transformer (-) with log R = Ó(1),
D O L Ö(VT), M S 4A, D'= A (d)(T))) |0| = O
such that laking Alga as defined in Eg. (2), we have
log ( (ar, - log
E [A].
Here 0(.) hides some absolute constant, and (()) additionally hides polynomial terms in (a, b Ba, Bm, ) and poly-logarithmic terms in (T, A,
A,d,1/5 1/8, 1/T).
A key component in proving Theorem 8 is demonstrating that the transformer can approximate the accelerated gradient descent algorithm for solving linear ridge regression (Lemma 21), result of independent interest. Leveraging Theorem 8, we can derive the following regret bound for the algorithm obtained via Algorithm Distillation, with the proof provided in Appendix E.S.
DD.LMMDD'BB B be the class of transformers satisfying Eq. (8) with e = 1/T and = 1/ log(4'' ('' AB. (B., + 2a/ / = Ó(T and choose the clip value log R = (1). Let both the context algorithu Algo and the expert algorithm Algg coincide with the soft LinUCB algorithm AlgaLinUCB(s) with parameter t during supervised pretraining. Then with probability at least the learned algorithm Alga solution Eq. (3), entails the regret bound
Theorem 9 (Regret of LinUCB and ICRL). Ler
EM-A Σ max (ar,;; w") Aag (T) s o dv dVTTog(T) + TV log(Ne T/o) k
n
where log S D(L'D(MD+ S d. (49logg)) log Here O hides polynomial terms in (a, b Ba, Bw, and Ô additionally hides poly-logarithmic terms in (T, A, d, 1/T).
4.2 THOMPSON SAMPLING FOR LINEAR BANDIT
We continue to examine the stochastic linear bandit framework of Section 4.1, assuming Gaussian prior w ~ (O, ALa) and Gaussian noises he mid (N,, r). Additionally, we assume existence of (b.. B.) such that b, S ||ar,1|l2 S Ba- In this model, Thompson sampling also utilizes linear ridge regression. Subsequently, we establish that transformers trained under the DPT methodology can learn Thompson sampling algorithms. We state the informal theorem in Theorem 10 below, where its formal statement and proof are contained in Appendix F.
Theorem 10 (Approximating the Thompson sampling, Informal), Consider the embedding mapping h, extraction mapping A, and concatenation operator cat in E. Under Assumption B, C, for sufficiently small e, there exists a transformer (-) with log R = Ö(1)
D = (4A) L=V(V)), = O
(9)
|||| Ö A4 + D'=O(AT
such that taking Alga as defined in Eq. (2), with probability at least 1 - do over (M, Dr) ~ PAAl for any Alg, we huve
log (ars
log (D-1,)
A T
1.1,Ba), Ba), and poly-logarithmic terms in Here, O(-) hides polynomial terms in (Mo, Co, , b-1
(T, A, d, 1/e, 1/66), where (Mo, Co) are parameters in Assumption B and
Central to proving Theorem 10 is establishing that the transformer roots via Pade decomposition (Appendix F.3), a result of independent interest. Theorem 10 thereby implies the subsequent regret bound for transformers trained under DPT.
approximate matrix square

---

Published as a conference paper at ICLR 2024
Theorem 11 (Regret of Thompson sampling and ICRL). Follow the assumptions of Theorem 10. Let = L..LLDD be the class of transformers satisfying Eq. (9) with E = 1/(RT))) 6n = 6/(2),,
and choose the clip value log R = (1). Assume the trajectories are collected by some context algorithm Algu, and we choose the expert algorithm Alg M) aj arg maxask, (a, w") to be the optimal action of the bandit instance M for each trajectory. Then with probability at least 1 6, the learned algorithm Algg. a solution to Eq. (3), entails regret bound
T
max (ackk w") - RM.AlE (T) & so so <o(dVTlogl (ddTloo((T)) + VR-
log (NgT/o)/)
EM~A
where R
-e and lo O(T.D(MD D`) logn) S O A + AvTd) logn). Here O hides polynomial terms in (X=1, r=1 , BB), and Ô additionally hides logy-lggari terms in (Ma. Co. T.A, d, 1/2, 1/
4.3 UCB-VI FOR TABULAR MDPS
A finite-horizon tabular MDP is specified by M (S,A, (PahnetA) {R (R(M)))))))) with H being the time horizon, S the state space of size 5, AA the action space of size A, and μ1 A(S) defining the initial state distribution. At each time step h E [H, Pa. SS x A + ASS) denotes the state transition dynamics and R2 S A [0, 1] gives the reward function. A policy क == {TA. (S x A x A
- ((A) ]IE maps history and state to a distribution over actions. The value of policy æ interacting with environment M is defined as the expected cumulative reward VM(T)= = EM, 12, h=1 RA (5), an)]. A policy π² is said to be optimal if π arg max =EA(T)(T)) VI((T).
We let the context algorithm Algo interact with an MDP instance M to generate K episodes, each consisting of H horizon sequences (sk, OR.b: [k.A) [ These can be reindexed into single trajectory Dr = (s2, 0z, re()) with /. = (kk - 1) + h and T KH. The Bayes regret AM (T))-
of any algorithm Alg gives BM~A[KVM(TT
Near minimax-optimal regret for tabular MDPs can be attained through the UCB- VI algorithm (Azar et al., 2017). We demonstrate that transformers are capable of approximating the soft UCB-VI algorithm AlgauCBvi(r)) slight modification of UCB-VI formalized in Appendix G.I.
Theorem 12 (Approximating the soft UCB-VI). Consider the embedding mapping h, extraction mapping A, and concatenation operator cat as in Appendix G.1. There exists a transformer TF)(.)
with log R = ((1)
D = O(HS A), L= 2H +8, M O A),
(10)
D' HS* |0|| O **
such that AlBsUCBvI(-) (a Dr- 1,) st) Algo (al Da),,,,, for all E TT, a E A. Here O(-) hides universal constants and ((-) hides
terms in (H, K, S, A,17)
Leveraging Theorem 12, we can derive the following regret bound for the algorithm obtained via Algorithm Distillation.
Theorem 13 (Regret of UCB-VI and ICRL). Let θ = OD,L,M be the class of transformers satisfying Eq. (10) with 5 = 1/K, and choose the clip value log R = Õ(). Let both the context al- gorithm Algo and the expert algorithm Algg coincide with the soft UCB-VI algorithm Alg.uCBvI(r) during supervised pretraining. Then with probability least 6, the learned algorithm Algg- solution to Eq. (3), entails regret bound
EM-A ( (KVM(TT) RRM.AlE (T)) S ó 6 HRVSAK + H T
Tog(NeT//8)
n
where No < D(I + logn) = A 2 + HSS HSS l)gon, and Ō(.) hides poly-logarithmic terms in (H, K, S,
8

---

Published as a conference paper at ICLR 2024
Linear bandit
Bernoulli bandit
Emp LinUCE
Emp
50
10-
UCB
Rngret
40
TF
Regret
30
20-
10
0-
50
100
150
200
50
100
150 200
Time
Time
Figure 1: Regrets of transformer (TF), empirical average (Emp), Thompson sampling (TS) and
LinUCB or UCB (LinUCB reduces to UCB for Bernoulli bandits). Left: linear bandit with = 5, A 10, a L.5. Algo = Algg LinUCB. Right: Bernoulli bandit with d 5, Algo (Algunt + Algrs)/2 and Algg = a". The simulation is repeated 500 times, Shading displays the standard deviation of the regret estimates.
5 EXPERIMENTS
In this section, we perform preliminary simulations to demonstrate the ICRL capabilities of trans- formers and validate our theoretical findings. We remark that while similar experiments have been conducted in existing works (Laskin et al., 2022; Lee et al., 2023), our setting differs in several aspects such as imitating the entire interaction trajectory in our pretrain loss (3) as opposed to on the last (query) state only as in Lee et al. (2023). The code is available https://github.cm http:iitiib.o//mioonn... lin/in-context-rl.
We compare pretrained transformers against empirical average, LinUCB (or UCB), and Thompson sampling. We use GPT-2 model Garg et al. (2022); Lee et al. (2023) with L = 8 layers, M = heads, and embedding dimension D 32. We utilize ReLU attention layers, aligning with our theoretical construction. We pretrain the transformer with two setups: (1) Both context algorithm Algo and expert algorithm Alg use LinUCB (the Algorithm Distillation approach); (2) Context algorithms Alg, mixes uniform policy and Thompson sampling, while expert Alg = až provides optimal actions (DPT). See Appendix C for further experimental details.
In the first setup, we consider stochastic linear bandits with d 5 and A = 10. At each & E [200]. the agent chooses an action at and receives reward r, (a,, w*) + St. where Et No, 1.5²). The parameter w* is from Uni 1)(). The action set As A is fixed over time with actions We generate 100K trajectories using Algo = Alg = LinUCB and
i..d. from ( train transformer TF (()) via Eq. (3). Figure 1 (left) shows regrets of the transformer (TF), empirical average (Emp), LinUCB, and Thompson sampling (TS). The transformer outperforms Thompson sampling and empirical average, and is comparable to LinUCB, agreeing with Theorem 9. The small regret gap between TF and LinUCB may stem from the limited capacity of the GPT2 model.
In the second setup, we consider multi-armed Bernoulli bandits with d = The parameter w" from (I0, 1] The fixed action set A A contains one-hot vectors {c; (multi- bandits). each E [200], the agent selects a, receives reward r, Bern w w)) Let Alg
be the uniform policy. We use Algunif and Algts context algorithms to generate 50K trajectories each. The expert fixed as Algg a² We train transformer TF via Eq. (3). Figure (right) shows regrets for the pretrained transformer (TF), empirical average (Emp). UCB, and Thompson sampling (TS). The transformer aligns with Thompson sampling, validating Theorem 11. However, TT underperforms UCB for Bernoulli bandits, as shown.
CONCLUSIONS
This paper theoretically investigates the ICRL capability of supervised-pretrained transformers. We demonstrate how transformers can efficiently implement prevalent RL. algorithms including Lin- UCB, Thompson sampling, and UCB-VI, achieving near-optimal regrets respective settings, We also provide sample complexity guarantees for the supervised pretraining approach to learning these algorithms. The generalization error scales with the covering number of the transformer class as well as the distribution ratio between the expert and offline algorithms. Simulations validate theoret- ical findings. Finally, we discuss the limitations of our results and provide additional discussions in Appendix A.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The authors would like thank Peter Bartlett for the valuable discussions, and Jonathan Lee for the valuable discussions regarding version of its implementation. This work is supported by NSF CCF-2315725, DMS-2210827, NSF Career award DMS-2339904, and an Amazon Research Award.
Transformers as well as providing an early
REFERENCES
Alekh Agarwal, Nan Jiang. Sham M Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms. CS Dept., UW Seattle, Seattle, WA, USA, Tech. Rep, 32, 2019.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent in-context learning. arXiv preprint arXiv:2306.00029) 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661
2022.
Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforce- ment learning. In International Conference on Machine Learning, pp. 263-272. PMLR, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint 2016.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Prov- able in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.0466 2023.
Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning a synaptic learning rule. Citeseer, 1990.
Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.0922 2020.
David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Systems, 35:1542-1553, 2022.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06817, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learnin via sequence modeling. Advances in neural information processing systems, 34:1 15844-15097, 2021.
Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandits with linear payoff func- tions. In Proceedings of the Fourteenth International Conference Artificial Intelligence and Statistics, pp. 208-214. JMLR Workshop and Conference Proceedings, 2011.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: arXiv:1810.04805. 2018.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning-identffabbliy challenges and effective data collection strategies. Advances in Neural Information Processing Systems, 34:4607-4618, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 1929,, 2020.
10

---

Published as a conference paper at ICLR 2024
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI²: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 2016.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. International conference on machine learning, pp. 1126-1 135. PMLR, 2017.
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487, 2021.
Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. arXiv preprint arXiv:2307.11333: 2023.
Justin Fu, Sergey Levine, and Picter Abbeel. One-shot learning of manipulation skills with online dynamics adaptation and neural network priors. In 2016 IEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp. 4019 4026. IEEE, 2016.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes, Advances in Neural Information Processing Systems, 1533058-30598., 2022.
Sara A Geer. Empirical Processes M-estimation, volume 6. Cambridge university press, 2000.
Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning structured exploration strategies. Advances neural information pro- cessing systems, 31, 2018.
Sepp Hochreiter, A Steven Younger, and Peter Conwell. Learning learn using gradient descent. In Artificial Neural Peetoors-ICAAN 2001: International Conference Vienna, Austria, August 21-25, 2001 Proceedings 11, pp. 87-94. Springer, 2001.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nugp and ntk for deep attention networks. In International Conference Machine Learning, pp. 4376- 4386. PMLR, 2020.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Tch, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv:1905 1905.0095564 2019.
Shin Ishii, Wako Yoshida, and Junichiro Yoshimoto. Control of applottation-explratttee meta- parameter in reinforcement learning. Neural networks, 15(4-6):665666) 2002.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273-1286, 2021.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.144: 14215, 2022.
Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selec- tion. Annals of statistics, pp. 1302-1338, 2000.
Jonathan N Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. arXiv preprint arXiv:2306. /4892, 2023.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Sherry Yang. Lisa Lee, Daniel Freeman, Sergio Guadar- rama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi- game decision trans- formers. Advances in Neural Information Processing Systems, 35:27921
2022.
Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reiinfrreement learning via distance metric learning and behavior regularization. arXiv preprint arXiv:20/00011 /2, 2020.

---

Published as a conference paper at ICLR 2024
Yingcong Li, M Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers al- gorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067 2023.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint 1XXV:22210.1049,, 2022.
Ya Yan Lu. A padé approximation method for square roots of symmetric positive definite matrices. SIAM journal matrix analysis and applications, 19(3):833-88) 845, 1998.
Song Mei and Yuchen Wu. Deep networks denoising algorithms: Sample-efficiens learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309. 11420, 2023.
Eric Mitchell, Rafael Rafailov, Xue Bin Peng. Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pp. 7780-7791. PMLR, 2021.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Picter Abbeel, Sergey Levine, and Chelsea Finn. Learning adapt in dynamic, real-world environments through meta- reinforcement learning. arXiv preprint arXiv: 1347, 2018.
Devang K Naik and Richard Mammone. Meta-neural networks that learn by learning. In [Proceed- ings 1992] IJCNN International Joint Conference on Neural Networks, volume 1, pp. 437-442 IEEE, 1992.
Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2003.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint 1803:08000009999 2018.
OpenAl. Gpt-4 technical report. arXiv preprint arXiv:2303.08744 2023.
Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You can't count on luck: Why decision transformers and rvs fail in stochastic environments. Advances in Neural Information Processing Systems, 35: 38966-38979, 2022.
Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. arXiv preprint arXiv:1901.03420 2019
Vitchyr H Pong, Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline meta- reinforcement learning with online self-supervision. In International Conference on Machine PMLR, 2022.
Learning, pp.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Nived Rajaraman, Lin Yang, Jiantao Jiao, and Kannan Ramehandran. Toward the fundamental limits of imitation learning. Advances in Neural Information Processing Systems, 33:291 24:2444-2924, 2020.
Nived Rajaraman, Yanjun Han, Lin F Yang, Kannan Ramchandran, and Jiantao Jiao. Provably breaking the quadratic error compounding barrier in imitation learning, optimally. arXiv preprint 12Xi::2102.12948. 2021.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforremment learning via probabilistic context variables. In International conference on machine learning, pp. 5331-53400 PMLR, 2019.
Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline rein- forcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems, 34:1 7722-17766 2021.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in cn-context learning for regression. arXiv preprint arXiv:2306. 15063, 2023.
12

---

Published as a conference paper at ICLR 2024
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXlv:2205.06115, 2022.
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings the thirteenth international conference on artificial intelligence and statistics, pp. 661-668 JMLR Workshop and Conference Proceedings, 2010.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the fourteenth international con- ference artificial intelligence and statistics, pp. 627-635. JMLR Workshop and Conference Proceedings, 2011.
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta-policy search. arXiv preprint arXiv:/810.06675 2018.
Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, Zheng Wen, al. A tutorial thompson sampling. Foundations and Trends9 in Machine Learning, 11(1):1-96, 2018.
Tom Schaul and Jürgen Schmidhuber. Metalearning. Scholarpedia, 5(6):4650, 2010.
Jürgen Schmidhuber. Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meia-.- hook. PhD thesis, Technische Universität München, 1987.
Jürgen Schmidhuber. Learning to control fast- memories: An alteraatie to dynamic recurrent networks. Neural Computation, 4(1):131-139 1992.
Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning i modes with one stone. Advances in neural information processing sys- tems, 35:22955 22968, 2022.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study on relu and softmax in transformer. arXiv preprint arXiv:2302.06661. 2023.
Miroslav Strupl, Francesco Faccio, Dylan R Ashley, Jürgen Schmidhuber, and Rupesh Kumar Srivastava. Upside-down reinforcement learning can diverge stochastic environments with episodic resets. arXiv preprint arXiv:2205.06500 2022.
Sebastian Thrun and Lorien Pratt. Learning to learn. Springer Science & Business Media, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp. 35151-35174, PMLR, 2023.
Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cam- bridge University Press, 2019.
Jane Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn arXiv preprint 1rXiv:1611.0500 2016.
Lingxiao Wang, Qi Cai, Zhuoran Yang, and Zhaoran Wang. On the global optimality of model- agnostic meta-learnin. In International conference machine learning, pp. 9837-98446 PMLR, 2020.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35: 10011120883 2022.
Mitchell Wortsman. Jachoon Lee, Justin Gilmer, and Simon Korblith.. Replacing softmax with relu in vision transformers. arXiv preprint arXivv2309.08588 2023.
13

---

Published as a conference paper at ICLR 2024
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021.
Mengjiao Yang. Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Dichotomy of control: Sepa- rating what you can control from what you cannot. arXiv preprint arXiv:2210.13433 2022.
Sherry Yang, Ofrr Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foun- dation models for decision making: Problems, methods, and opportunities. arXiv preprint arXiv:2303.041299 2023.
Shunyu Yao, Binghui Peng. Christos Papadimitriou, and Karthik Narasimhan. Self-attention net- works process bounded hierarchical languages. arXiv preprint arXiv:2 2021.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequenee-to-sequence functions? arXiv preprint arXiv:1912.10077. 2019.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2300.009927 2023.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task. arXiv preprint arXiv:2206.04500 2022.
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta- learning. arXiv preprint arXhv:1910.08343. 2019.
14

---

Published as a conference paper at ICLR 2024
CONTENTS
1 Introduction
2 Framework for In-Context Reinforcement Learning
3 Statistical analysis of supervised pretraining
3.1
Main result
3.2
Implications in special cases
Approximation by transformers
4.1
LinUCB for linear bandits
4.2
Thompson sampling for linear bandit
4.3
UCB-VE VI for Tabular MDPs
5 Experiments
6 Conclusions
A Related work
15
A.I Limitation and discussion
16
Technical preliminaries
17
Experimental details
19
C.1 Implementation details
20
C.2 Additional experiments and plots
20
C,3 The effect of distribution ratio
21
D Proofs in Section 3
21
D.1
Proof of Theorem
22
D.2 Proof of Proposition
23
D.3
An auxiliary lemma
Soft LinUCB for linear stochastic bandit
25
E.
Embedding and extraction mappings
26
E.2
LinUCB and soft LinUCB
26
E.3
Approximation of the ridge estimator
27
E,4
Proof of Theorem
29
E.5
Proof of Theorem 9
F
Thompson sampling for stochastic linear bandit
35
F.1
Thompson sampling algorithm
F.2
Definitions and assumptions
................................
F3
Proof of Theorem 23 (and hence Theorem 10)
F.4
Proof of Theorem 11
................................
44
F.5
Proof of Lemma 24
46
G Learning in-context RL in markoy decision processes
48
G.I Embedding and extraction mappings
................................
48
G.2 UCB-VI and soft UCB-VI
................................
49
G.3 Proof of Theorem 12
49
G.4 Proof of Theorem 13
57
A RELATED WORK
Meta-learning and meta-reinforcemmnt learning In-context reinforcement learning can be cast into the framework of meta-leat meta-learning and renforeeifore learning (Schmidhuber, 1987: 1992; Bengio et al., 1990; Naik & Mammone, 1992: Ishii et al., 2002; Schaul & Schmidhuber, 2010; Thrun & Pratt, 2012). More recently, line of work focuses meta-learn certain shared structures such
15

---

Published as a conference paper at ICLR 2024
as the dynamics of the shared tasks (Fu et al., 2016; Nagabandi et al., 2018), task context identifier (Rakelly et al., 2019; Humplik et al., 2019; Zintgraf et al., 2019), exploration strategies (Gupta et al., 2018), or the initialization of the network policy (Finn et al., 2017; Hochreiter et al., 2001; Nichol et al., 2018; Rothfuss et al., 2018). Theories for this last approach of model-agnostie meta-learning have been explored by Wang al. (2020).
Our work focuses on more agnostic approach to learning the learning algorithm itself (Wang et al., 2016; Duan et al., 2016;
Dorfman et al., 2021; Mitchell al., 2021; Li al., 2020; Pong al., 2022; Laskin et al., 2022; Lee et al., 2023). Among these works. Wang et al. (2016); Duan et al. (2016) focus on the online meta-RL setting with the training objective be the total reward. Furthermore, Dorfman et al. (2021); Mitchell et al. (2021); Li et al. (2020); Pong et al. (2022) focus on offline meta-RI., but their training objectives differ from the cross entropy loss used here, requiring explicit handling of distribution shift. The supervised pretraining approach we consider is most similar to the algorithm distillation methods of Laskin et al. (2022) and the decision-prec ppresiion--rtraaeedd transformers of Lee et al. (2023). We provide quantitative sample complexity guarantees and transformer constructions absent from previous work.
In-context learning The in-context learning (ICL) capability of pretrained transformers has gained significant attention since being demonstrated on GPT-3 Brown et al. (2020). Recent work investigates why and how pretrained transformers perform ICL (Garg et al., 2022; Li et al., 2023; Von Oswald et al., 2023; Akyürek et al., 2022; Xie et al., 2021; Bai et al., 2023; Zhang et al., 2023; Ahn et al., 2023; Raventós et al., 2023). In particular, Xie et al. (2021) propose a Bayesian frame- work explaining how ICL works. Garg et al. (2022) show transformers can be trained from scratch to perform ICL of simple function classes. Von Oswald et al. (2023); Akyürek et al. (2022); Bai al. (2023) demonstrate transformers can implement in-context learning algorithms via in-context gradient descent, with Bai et al. (2023) showing transformers can perform in-context algorithm se- lection. Zhang et al. (2023) studied training dynamics of single attention layer for in-context learning of linear functions. Our work focuses on the related but distinct capability of in-context decision-making for pretrained transformers.
Transformers for decision making Besides the ICRL approach, recent work has proposed goal- conditioned supervised learning (GCSL) for using transformers to make decisions (Chen et al., 2021;
Janner et al., 2021: Lee et al., 2022: Reed et al., 2022: Brohan et al.. 2022: Shafiullah et al., 2022: Yang et al., 2023). In particular, Decision Transformer (DT) (Chen et al., 2021; Janner et al., 2021) uses transformers to autoregressively model action sequences from offline data, conditioned on the achieved return. During inference, one queries the model with desired high return. Limitations and modifications of GCSL have been studied in Yang et al. (2022); Paster et al. (2022); Strupl et al. (2022); Brandfonbrener et al. (2022). A key distinction between GCSL and ICRL is that GCSL treats the transformer as a policy, whereas ICRL treats as an algorithm for improving the policy based on observed trajectories.
Expressivity of transformers The transformer architecture, introduced by Vaswani et al. (2017), has revolutionized natural language processing and is used in most recently developed large lan- guage models like BERT and GPT (Devlin et al., 2018; Brown et al., 2020). The expressivity of transformers has been extensively studied (Yun et al., 2019; Pérez et al., 2019; Hron et al., 2020; Yao et al., 2021; Bhattanishra et al., 2020; Zhang et al., 2022; Liu et al., 2022: Wei et al., 2022; Fu et al., 2023; Bai et al., 2023; Akyürek et al., 2022; Von Oswald et al., 2023). Deep neural networks such as ResNets and transformers have been shown to efficiently approximate various algorithms, including automata (Liu et al., 2022), Turing machines (Wei et al., 2022), variational inference (Mei & Wu, 2023), and gradient descent (Bai et al., 2023; Akyürek et al., 2022; Von Oswald al., 2023). Our work provides efficient transformer constructions that implement accelerated gradient descent and matrix square root algorithms, complementing existing expressivity results.
Statistical theories of imitation learning Our generalization error analysis adapts classical anal- ysis of maximum-likclihodd estimator (Geer, 2000). The error compounding analysis for imitation learning appeared in early works (Ross al., 2011; Ross & Bagnell, 2010). More recent theoreti- cal analyses of imitation learning also appear in Rajaraman et al. (2020; 2021); Rashidinejad et al. (2021).
A.1 LIMITATION AND DISCUSSION
In this section, we discuss some limitations of our work and some potential future directions.
16

---

Published as a conference paper at ICLR 2024
Distribution ratio In Theorem 6, our regret bound of the learned algorithms Alga depends on the distribution ratio & Alg,, glgg. While in cases like algorithm distillation (Laskin et al., 2022) the dis- tribution ratio equals one since the offline algorithm matches the expert algorithm, in the worst case, the ratio can exponentially depend on T or even become arbitrarily large. To control the distribution ratio in practice, one approach is augment the offline trajectory dataset with portion of trajecto- ries generated by an expert algorithm or no-regret algorithms resembling the expert algorithm. On the other hand, further research could investigate structural assumptions of decision-makinn makinn prob- lems that avoid pessimistic dependence on the distribution ratio in regret bounds.
Guarantee of pretrained transformer Our statistical result (Theorem 6) only guarantees that the pretrained transformer learns an "algorithm" matching the expert algorithm under the pre-training distribution, even though our approximation results (Theorem 8, 10, 12) show the existence of a transformer approximating the expert algorithm over the entire input space. In our early experi- ments, we noticed the learned transformers do not generalize well on out-of-distributin instances, such as with shifted reward distributions or increased number of runs T. Similar phenomena occur in other in- in-context learning problems (e.g. Garg et al. (2022)). Understanding the actual algorithm im- plemented by the pretrained transformer through theoretical and empirical analysis is an interesting question for future work.
Alternative pretraining methods Our theoretical results study pretraining the transformer by maximizing the log-likelihood of i.i.d. offline trajectories as Eq. (3). This aligns with standard supervised pretraining of large language models. However, alternative pretraining methods may also be effective. For instance, one could replace the log- prgg-abbbbbitty in Eq. (3) with an & loss for contin- uous action spaces, consider other objectives like cumulative reward (Duan al., 2016), explore goal-conditioned reinforcement learning (Chen et al., 2021) for in-context RI. While our work fo- cuses on oigliiili pretraining, thcoretical investigation of alternative methods is an interesting direction for future work.
Possibility of surpassing the expert algorithm by online training Our work considers offline pretraining by imitating the expert algorithm (i.e., Algg), which can only learn transformer match- ing the expert's performance at best. However, through online training, where the transformer in- teracts with the environment, the learned transformer may surpass existing experts by training to improve itself rather than imitating a specific algorithm. Investigating whether online training en- ables surpassing expert algorithms is an interesting direction for future work.
Implications for practice While the focus of our work is primarily theoretical, our results lead to several practical implications for in-context reinforcement learning. One key implication is the importance of training labels (i.e., expert actions a). When the expert algorithm depends solely on past observations, we can learn Alg (see Theorem 9). In contrast, when Alg is the optimal action a" (involving knowledge of the underlying MDP), we can learn the posterior average of this algorithm given past observations. This corresponds to the Thompson sampling algorithm, in Decision- Pretraindd Transformers (see Theorem 11).
Furthermore, as discussed previously, the distribution ratio between the offline and expert algorithms may impact the generalization of the learned algorithm. Both our theory (see Theorem 8) and simulations (see Figure 4) show that a small distribution ratio between the offline algorithm Alg, and the expert algorithm Algg is essential, otherwise the online performance of the learned algorithm may substantially degrade. This suggests that incorporating trajectories generated purely from the expert ("on-policy ICRL") into the offline dataset is advantageous, when feasible.
B TECHNICAL PRELIMINARIES
In this work, we will apply the following standard concentration inequality (see e.g. Lemma A.4 in Foster et al. (2021)).
Lemma 14. For any sequence of random variables (X1)eST adapted to a filtration (FIT have with probability at least 1 8 that
X.
log El(ex((X.) [F-1]
+ Log(1 1
for all LE (T).
17

---

Published as a conference paper at ICLR 2024
Lemma 15. Adopt the notations in Definition 4. Then for any 0 E O there exists 8g E O such that ||Aggo, (-IDi-1, St) - Algo( D1 S 2p.
Proof of Lemmaa 15. For any 0 θ, let Go Oo be such that log Alga, (1) log Algo( (;-1 st)ll S p for all Di-1, St and t E [T]. Then
||Ago (- ( 81) Algo( ( 8)]l/
Σ
Algo, Sc) Algo(a
oeAt
Σ
max AlBon
Alao
aEA.
log Alga, (-1D--1, Algo
maxxflog flog AlBa (-1De-1.0e)),ogg
S p
aEAA.
Σ (Alge, (D)-1,88) + Alga (- 2p:
S p
aEAt
where the second line uses a Taylor expansion of e", the fourth line uses the assumption on 6o.
the last line uses max S ex + ea and the fact that Alga, ( (- are probability functions.
We have the following upper bound on the covering number of the transformer class {TF 0 e D.L.LM.D',B) )
Lemma 16. For the space of transformers {TF B
0 DDD,,M.D' B} with
O DDLL,MDDIB == {o (O.lt.. olip)) max ( S M, max ( 101 BB. FE(T.]
BE[T]
where (()), ()(()) denote the number of heads and hidden neurons in the t-th layer respectively, the covering number of the set of induced algorithms {Alge: 0 E &DDLLM,D'.R (c.f. Eq. 2) satisfies
22 max R]
log e
Dp..,M. os (0) S cT.
for some universal constant
Remark of Lemma 16. Note that the transformer classes OD.LLM.D'.B. DD..MM.D.E. have the same expressivity one can augment any TFo iD.L.M,D..B such that the resulting TF FFg P.L.M, a by adding heads or hidden neurons with fixed zero weights. Therefore, the same bound in Lemma 16 follows for D. n', and throughout the paper we do not distinguish
D.L.M.D and D.L.M, B and use them interchangeably. We also use ((0), D=(()) to represent the number of heads and hidden neurons in the é-th layer of transformers, respectively.
Proof of Lemma 16. We start with introducing Proposition J.1 Bai et (2023).
Proposition 17 (Proposition J.1 in Bai al. (2023)). The function TE is (LB); Be) Bipsciitt w.E.t. OD.LMM.' for any fixed input H. Namely, for any 01, 02 E DDLLMM.'..B we have
||T,R, (H) TFF (H)ll,,0 S LLB
02ll.
where ||A := sup.emm e| for any matrix A kxT, and Bg = BR(1 BR +
B²R2), BH == (1+ (1 (1+ B²R3).
As in the Proof of Theorem 20 in Bai et al. (2023), (2019) that the 5-covering number
verify using Example 5.8 in Wainwright
log (ō; B (r), < L(3MD² + 2DDDD log(1 2r/5),
(11)
where B1-1(7 (T) denotes any ball of radius r under the norm ||-1|. Moreover, we have the following continuity result the log-softmax function
18

---

Published as a conference paper at ICLR 2024
Lemma 18 (Continuity of softmax),,. For any u, E R we have
e log
log
[eall
T
We defer the proof of Lemma 18 to the end of this section.
Note that (- -. st) corresponds to A entries in one column of H applied through the softmax function. Therefore, combining Proposition 17, Lemma 18 and Eq. (11), we conclude that for any r > O, there exists subset Og D.LL,M D'. ,B with size L(3MD² + 2DD') log(1 + 2r/6) such that for any @D,L,M,D',R there exists 6g E Op with
|log Algo( (- (1))). - log Alga, (D) S 22LB Bed
for all Dr. Substituting N... (p) in Lemma 16.
B and letting p/(2LB yields the upper bound on
Proof of Lemma 18. Define w = u-v. Then
log ||e"|| ) - log
eu
ev
|1e^
< ||u - vl1oo + |log ||e"|1 - log ||e" l
||u vlloo + /''
extiw
ww) dt
|ev
$ ||u vl + (-)
c v--w
|| dt
||eettw
2 vl
where the third line uses the Newton-Leibniz formula.
We present the following standard results on the convergence of GD and AGD, We refer the reader to Nesterov (2003) for the proof of these results.
Proposition 19 (Convergence guarantee of GD and AGD). Suppose L(w) a-strongly convex and B-soooth function on Rd. Denote the condition number == 3/a and w' == arg min L(w). = wOO 1/VI(w with stepsize n = 1/3 and
(a). The gradient descent iterates w t+1 GD
initial point w&D = O₂ satisfies
|W -w" 11 < exp( e w'113
3
- -xp(-=))lwG
I(WED) - T.(w")
< 2 expl--
w'1l3
(b). The accelerated gradient descent (AGD, Nesterov (2003)) iterates w t+1
AGD = v
x+11
YGD
w++1
AGG
JR+1
JK-1 (wi+1
AGD WAGG) with w v
AGD 0g
satisfies
|WADD - w0 < (1+ k:) exp(
}| W w'1l3
Vk
L(When) L(w") < 2
+ 8
expl
1| whGD w' Vk
t
C EXPERIMENTAL DETAILS
This section provides implementation details of experiments and some additional simulations.
19

---

Published as a conference paper at ICLR 2024
C.I IMPLEMENTATION DETAILS
Model and embedding Our experiments use GPT-2 model (Radford et al., 2019) with ReLU activation layers. The model has L = 8 attention layers, M attention heads, and embedding dimension D = 32. Following standard implementations in Vaswani et al. (2017), we add Layer Normalization (Ba et al., 2016) after each attention and MLP layer to facilitate optimization. We consider the embedding and extraction mappings as described in Appendix 1, and train trans- former TF (.) via maximizing Eq. (3).
Online algorithms We compare the regret of the algorithm induced by the transformer with em- pirical average, Thompson sampling, and LinUCB (or UCB for Bernoulli bandits).
(Emp) Empirical average. For time t Σ A, the agent selects each action once. For time f. > A, the agent computes the average of the historical rewards for each action and selects the action with the maximal averaged historical rewards.
(TS) Thompson sampling. For linear bandits with Gaussian noises, we consider Thompson sampling introduced in Appendix F.1 with = a 1.5 and À = 1 (note that in this case TS does not correspond posterior sampling as we assume w follows the uniform distri- bution [0, 1) For Bernoulli bandits, we consider the standard TS sampling procedure (see, for example, Algorithm 3.2 in Russo et al. (2018)).
(LinUCB) Linear UCB and UCB. For linear bandits, we use LinUCB (Appendix E.2) with À and a = 2. For multi-armed Bernoulli bandits, LinUCB reduces to UCB, which selects at = arg max.e, Af + vI/NN(a))} where fr,a is the average reward for action a up to time and (a) is the number of times action was selected up to time
C.2 ADDITIONAL EXPERIMENTS AND PLOTS
We provide additional experiments and plots in this section. In all experiments, we choose the number of samples n = 100K.
a, w*) periments in Section 5. In both cases, the transformer is able to imitate the expected expert policy Alg as its suboptimality closely matches Alg, (LinUCB and TS for the left and right panel, re- spectively). While the empirical average (Emp) has lower suboptimality early on, its gap does converge In contrast, both LinUCB and Thompson sampling are near-optimal up ((1) factors in terms of their (long-term) regret.
Additional plots of suboptimality -a,
time are shown in Figure for the two
Linear bandit
Berouuli bandit
1.2
Emp
Emp
Suboptimality
1.0
TS
LinUCB
Suboptimality
0.3
TS
UCB
0.8
TF
0.2
0.6
0.4
0.1
0.2
0.0
0.0
50
100
150
200
50
100
150
200
Time
Time
Figure 2: Suboptimalities of transformer (TF), empirical average (Emp), Thompson sampling (TS), and LinUCB (or UCB). Left: linear bandit with d = 5, A = 10, 1.5. Alg, Alg = LCB.... Right: Bernoulli bandit with d = 5. Algo = (Alguni! + Alg 1/2, and Algg a; The simulation is repeated 500 times. Shading displays the standard deviation of the sub-optimality estimates.
Additional simulations were run with Alg, Alg UCB for Bernoulli bandits, which has fewer actions (A = 5) than linear bandits (A = 10). Figure 3 shows the regret and suboptimality of UCB and the transformer overlap perfectly, with both algorithms exhibiting optimal behavior. This suggests the minor gaps between LinUCB and transformer in the left panel of Figure and are likely due to limited model capacity.
20

---

== 0 Published as a conference paper at ICLR 2024
Bernoulli bandit
Bernoulli bandit
Emp
Emp
12.5
UCB
0.3
UCB
TS
Suboptimality
10.0
Regret
TF
0.2
7.5
5.0
0.1
2.5
0.0
0.0
50
100 Time
150
200
50
100
150
200
Time
Figure 3: Regrets and suboptimalities of transformer (TF), empirical average (Emp). Thompson sampling (TS), and UCB. Settings: Bernoulli bandit with d = 5, and Algo Alg LinUCB. The simulation is repeated 500 times. Shading displays the standard deviation of the estimates.
C.3 THE EFFECT OF DISTRIBUTION RATIO
We evaluate the effect of the distribution ratio R = g (Definition 5) transformer per- formance. We consider the Bemoulli bandit setting from Section 5 with expert Algg a* giving optimal actions. The context algorithm is
Alga aAlgrs (1 Algunif
mixing uniform policy Algunif and Thompson sampling Algrs. for a E {0,0.1, 0.5,, The case 0 corresponds the context algorithm being the i.i.d. uniform policy, and corresponds to the context algorithm being Thompson sampling. Note that the distribution ratio R may scale as ((1/a)) A (O(T) in the worst case.
Figure 4 evaluates the learned transformers against Thompson sampling for varying context algo- rithms. The left plot shows cumulative regret for all algorithms. The right plot shows the regret difference between transformers and Thompson sampling. The results indicate that increased distribution ratio impairs transformer regret, as expected. Moreover, it is observed that the trans- former, even with the uniform policy (i.e., a = 0), capable of imitating Thompson sampling in the early stages (Time < 30), exceeding theoretical predictions. This suggests the transformer can learn Thompson sampling even when the context algorithm differs significantly from the expert algorithm.
Bernoulli bandit
Bernoulli bandit
Difference Regrets
== 0
0.1
4
0.1
== 0.5
0.5
Regret
a= 1
=
2
50
100
150
200
50
100
150
200
Time
Time
Figure Regrets and difference of regrets between transformers and Thompson sampling, for different context algorithms. Settings: Bernoulii bandit with d = 5. Alg ot and Alg, aAlgss (1 ax Alg with a (0,0.1, 1]. The simulation repeated 500 times. Shading displays the standard deviation of the estimates.
D PROOFS IN SECTION 3
In this section, denotes universal constants that may differ across equations.
21

---

Published as a conference paper at ICLR 2024
D.1
PROOF OF THEOREM 6
Proof of Eq. (5) Note that we have
M-A.a V D2 (AIBE(-)D(--,, Alga 8())
Algg (ug (asg(as)aa-l:11 Da-11
) (1 Algyp (
Algo (ag Do-1,
xm
A D2 1 )
Also
Algo (a, 0.1 s.)) (2)
EA,, (AlgE AlgA (
Al Dx- 8z) II
Alao
BA,AEo D2 St) sc)); Σ
Algo (a, D,-1.88)
where the second line follows from change of distribution argument, the third line follows from Cauchy-Schwand inequality, and the fourth line uses the fact that
Q1 (x)
Q1(2) du(x,y) v)
(2))((2)))))
P1
/ Q1 (x)²
(
Q3 (wl=))
(1)))
Q1(2)
du(ulx) du(x)
du(x) = E
(x P2
Pi(x)
P1(z)
for any probability densities (Qu, 1. with respect some base measure jr.
Continuing the calculation of the above lines of bounds, we have
T
E M M-A.ann- i V DH (-)-,()))))))))) (DE)))))
1=1
Alg (Bg DD- 8s)
VT
E MMM.Aerr--iAAl
II
Algo (as Da- 1, 8x)
EM-.A.ane---Al.. oz Di (Ag (- (-1De-1, Algp( (Ie-(-)
Alg (a, Do-1,
VT
E M -1 N
II Algo (a. D
EM-.A.nna----ll..- D Alg (1 APE(--))
log Ne (1((L((())) + (Tg(T/8)
cT A
Ereal
7
log (1gTT)(// + V
I cT R
where the first incquality follows from the Sauchy-Schwand inequality, the first equality is due to change of distribution argument, the second inequality uses Lemma 20. This completes the proof of Eq. (5).
Proof of Eq. (6) For any bounded function f such that (f(D)) < F for some F have
EM-A.aNEE [S(DT) EM--A.-Ale (l, [((D))
22

---

Published as a conference paper at ICLR 2024
T
A [f(DT) -MMA. Aag. - [
< 2F
B M--AALa--ATEer Drv (Agg (- Dr 8:));
where the first equality uses the performance difference lemma, the last line follows from the varia- tional representation of the total variation distance
(Tv(P.Q) Q) = sup -
and
Drv((i(r)Pa() s)) P1)) x))P((z
y)) 1 Erw,, Drv (P2(y x)) Pa(y [=)
(12)
for probability densities 11,3,,,,,,344 with respect
base measure μa. Since Drv (P, Q) s
(P, Q) for any distributions P, Q, follows from Eq. (5) that
E EMLAONTE [f(Dr) - EM-A, [f(DT)]
log We (NG(1((T))))/// + Ereal
Sel RAEES T
with probability at least 1 - 8 for some universal constant noting that [f(rr] s T concludes the proof of Theorem
0. Letting f(Dr) Et= and
D.2 PROOF OF PROPOSITION 7
By the jointly convexity of KL(P ||
Q) with respect to (P,, Q) and the faet that (P, Q) S
(L || Q), we have
D2 (ABgE ( D1)1,
E
NEa KL D D.-1, 8t) PTs( 1,
PAo KL P....(( (r)) S Eapprox-
Therefore, applying Lemma 20 gives
Aleo E MrA.
2E 2.
Dj (- AlE + (AAee((1)) Alg (- (+IDe-1,
<c Togg (2((/(nT)2))/// + T(Ereal + Fapprox))
with probability at least 1 - 8. Proposition 7 follows from similar arguments in the proof of Theorem with Greal replaced by Ereal + Sapprox-
D.3 AN AUXILIARY LEMMA
Lemma 20 (General guarantee for supervised pretraining). Suppose Assumption A holds. Then the solution Eq. (3) achieves
T
Tlog [Ne(L(n2TT3)T//] + Tereal-
( Alg (
n
with probability at least for some universal constant e 0.
23

---

Published as a conference paper at ICLR 2024
Proof of Lemma 20.
Define
(0)
log Algo Da st), and Lne(expert): Σ log Alga (c)) D
and let (0) = Et=1 1=1 (in((), fa(expert) = Ei=l (experr) We claim that with probability
least 1
ED ( Alg (1
Ln (expert) (.(()) + Neog - Tlog(T/6) 4
(13)
n
n
n n
for all 0 E θ, E [T], where Dr follows distribution PM (-), M ~ A. For now, we assume this claim holds. Moreover, it follows from Lemma 14 and the fact L. (0) 2 Ln (0*) that
L (experrt) (0)
(expppert) (0)
(expex)) Lnr
n
n
n
T Tlo((T/o)
T
(1) Alg (Tg
log E Dr
n
Algg- (Te D1-1.
T Tlo((T/0)
+ Tereal
(14)
n
with probability at least
Choosing @ in Eq. (13) and combining with Eq. (14) and a union bound, we obtain
T
Σ EDr
(Algg (. st) (1B
1=1
Tlog + 27 log(2//5) +
S Treal + 2
T'rrrll + cf log (ee log(T/))
with probability at least - ö for some universal constant > 0. This completes the proof.
Proof of Eq. (13) Let O, 1/onT)=-coveriig set with covering number Yloov = |O,]. For
l E [Nenw], E [T), E [m)), define
Alg (a s()
f. log
Alga, (क) )
where (Di, (Di,a') a²) are the trajectory and expert actions collected in the i-th instance. Using Lemma 14
with X, -be and union bound over (k:, 1), conditioned on the trajectories (D)...., D.), we have
f
Σ t + log(nnow Nroow T/8)
>
log E exp
Act
2
2
1=1
for all & E [Meox], E [7] with probability at least
8. Note that
exp (-
tit 2
Di s6) = EEp
Alga, (a DJ 18))
Di 5/
V Alg (ā: D st)
Σ
/Algo, (alDi 3/
Alg DJ s
A.
24

---

Published as a conference paper at ICLR 2024
where the last inequality uses the assumption that the actions क' are generated using the expert A ( (: s)). Therefore, for any 0 E θ covered by Os, we have
- -log exp
R!
2
2 Epi Σ
Algo, (alDi (
aeAL
Epi
Σ
V Algoo (
aEAA
D
V (I
Algo, (a) (a)D)_ Algo(a)Di si)
aEA.
1
D
D AAg (-1 Algo (-D((-D(-21
2
Alga (1Di-1; - v /Algo, (-D)-1, )
ED Σ
aEA
1 D
(ATg (-1 st) Algo( (-1 s{))
|AAgo (-1Di
Algo, (1D(-
2
1 Ep 2
Dj (Alg (-IDi_ s) Alga (-) D
$))]
V2
nT
for all E [n], E [7'] where the first inequality uses - log I, the second inequality follows
from Sauchy-Schward. inequality, the third inequality uses (Vz - /y)² < |x yl for I, y 2 0, the last inequality uses the fact that 9 is covered by 0 and Lemma 15. Since any 0 θ is covered by Og for some k E [Neow]. and for this & summing over / E [T] gives
n
hh (expert)) (0s) S Lu(expert) (0) nT S (expert)) - (n) +
Therefore, with probability at least - 8, we have
C. (expert) Ln(() + 1) + Tlog V2
T
71
EDT
(1((Alo(- AlBexxxer (- st))
2
1=1
for all 0 E θ, where Dj follows PAE. Multiplying both sides by 2/n and letting neov Ne(1/(nT)?) yields Eq. (13).
E SOFT LINUCB FOR LINEAR STOCHASTIC BANDIT
Throughout this section, we use c > to denote universal constants whose values may vary from line to line. Moreover, for notational simplicity, use O(-) to hide universal constants, O(-) hide polynomial terms in the problem parameters (a, Ba. Bw. X and Ö(-) to hide both poly-logarithmic terms in (T, A, d, 1/e, 1/T) and polynomial terms in (a, b- B.s Bw. X+1), We also use the bold font a, E Rd denote the selected action vector a; at time E [T].
This section is organized follows. Section E. discusses the embedding and extraction formats of transformers for the stochastic linear bandit environment. Section E.2 describes the LinUCB and the soft LinUCB algorithms. Section E.3 introduces and proves lemma approximating the linear ridge regression estimator, which is important for proving Theorem We prove Theorem in Section E.4 and prove Theorem 9 in Section E.5.
25

---

Published as a conference paper at ICLR 2024
E.1 EMBEDDING AND EXTRACTION MAPPINGS
Consider the embedding in which for each E [2", we have two tokens hgg-1, h2: E RDD
such that
a,
0a+1
Aag
h
F't
h
h
Odd
h
h 1
0.1
h2r
0
h
0
Ö
h2
pos
p0822
where 12211 A. [aT,
TT[A]' denotes the action set at time t, hng [a] ri]
denotes the action and the observed reward at time t, h is used to store the (unnormalized) policy at time 1, O in ha denotes an additional zero vector with dimension OddA), and pos; == (i, 12, 1) for i [27] is the positional embedding. Note that the token dimension D O(dA). In addition,
we define the token matrix H, := [h,...., h2g E R Dx22
for all t E [7]]
Offline pretraining During pretraining, the transformer TFa takes in H Hy the input token matrix and generates H TF, as the output. For each step [T]. we define 1
E EA whose i-th entry is the probability of the induced policy Al = (x 21 Ill
selecting action ar, given (D.-1,8)). We then find the transformer ① E θ by solving Eq. (3). Due to the decoder structure of transformer TF,, the 2/ 1-th token only has access to the first 2/ tokens. Therefore the induced policy is determined by the historical data (Dr-1. St) and does not depend on future observations.
Rollout At each time t E [T', given the action set rr (i.e., current state st) and the previous data Di-1, we first construct the token matrix H = -H,-1, h2: E D21-11)) The trans- former then takes H roll,t as the input and generates HPost [HPa h hpost TF T (an..) Next, the agent selects an action a E As according to the induced policy Algo (-1 Di-1, == oxp(hprat))
exp(hgr (onpheee Illa E A and observes the reward Tt.
Embedding and extraction mappings To integrate the above construction into the framework described in Section 2, we have the embedding vectors h(ss) := h2,,,, r1) := h2, the con- catenation operator cat(hı..... hx) == [h.,.., hr]. the input token matrix
H rall.r = cat(h(s1), h(a1, b(a,..,,,((a,11, 1 n(-), E Dx
the output token matrix 22-1
Host roll,e and the linear extraction map A satisfies A h = A h
E.2 LINUCB AND SOFT LINUCB
Let T be the total time and X; a >> O be some prespecified values. At each time E (T). LinUCB consists of the following steps:
1. Computes the ridge estimator w A arg min 28
A 2t w E (aj, w)) +
2. For each action A E [A], computes vix : ( (a w whidpo.) av / ac,.
" where
A, = ALa + a,a..
3. Selects the action a., with j := arg max Uis.
Unless stated otherwise, in step above we choose a = a(6) with 1/(2B,B.T) and
(6) = VAB,, /. 2og + dl log (d) (TAA+TBB))/ce TB3) = O( T) = Ö(Vī).
In this work, to facilitate the analysis of supervised pretraining, we consider soft LinUCB (denoted by SLinUCB(7)). which replaces step 3 in LinUCB with
3' Selects the action a.j with probability exp(e7,/r) [exp(o3/>>]1
for j E [4].
Note that soft LinUCB recovers the standard LinUCB as 7 0.
26

---

Published as a conference paper at ICLR 2024
E.3 APPROXIMATION THE RIDGE ESTIMATOR
In this section, we present a lemma on how transformers can approximately implement the ridge regression estimator in-context.
Throughout the proof. for t E [2T]. we let h, denote the i-th token in the output token matrix obtained after passing through an I.-layer transformer. We also define read : RRP + IR be the operator that gives the values of d coordinates in the token vector that are used store the estimation of the ridge estimate.
Lemma 21 (Approximation of the ridge estimator), For any small & > 0, there exists an attention- only (i.e., no MLP layers) transformer TF with
4T + A) A (B, o)/(E)) = =Ō(T),
X+2 2
max M 3, |0| S v2+ eE[r]
= O(1)
B +
(4) w rid S for all E [T].
such that rread (h
Moreover, there exists a transformer (-)()) with
B log
(2T(B2 + A) + X)T C..BBBB.... ) = OO(TT)
= 2v 2V 2T
max M S 4, eE[r]
X2e
max (r) RE(L]
4d. ||0| S 10
2+2 0(1)
B2 +
such that read
(h22 (L)
w ridge,Allz S for all t E [T].
Results similar to Lemma 21 have been shown in Bai et al. (2023) under different scenario. How- ever, we remark that the second part of Lemma has a weaker requirement on the number of layers we prove that transformers implement accelerated gradient descent (AGD, Nesterov (2003)) in-context.
Proof of Lemma 21. Note that AIa 3 A, (TB2 + A)La. Therefore the optimization problem
1
A
w Widgee. arg min L(w): = arg min weRd 2(21
(ra (aj, w)) + 2(21 |lw12
weRd
1)
is A/(2t - 1)-strongly convex and (Ba + A) smooth and the condition number S BB? + A)/2. Moreover, by the definition of w ridge,A we have
1-
1-1
| Alla = [1(AT +
Σ a
alr < (1 +
Σ
112
Σ
agryll2
==1
1=1
j=1
==1
TB, (B, +
X
for all E [T'].
Proof of part By Proposition 19, we see that -= [BT +X) (., + +A)/(A)))/AT
steps of gradient descent with stepsize 7) = 1/(B2 + A) starting from w. = O2 finds w such
that 1/ W'idge, ll S
Now we prove that one attention-only layer can implement one step of gradient descent
2+1 GD
Σ (aj, woD) - -r,)a; 2t 1=1
W
2t-1
We encode the algorithm using the last token (i.c., the 2t 1-th token). Denote the first entries of hg 1 by W and define read (h2--1) w. Starting from T O4,, for each layer E [T.],
27

---

Published as a conference paper at ICLR 2024
let the number of heads M(4) = 3 and choose 2,,,3 1,2.3 V1.2. such that for even tokens haf
K (6)
with; j S ( - and odd tokens h25-1 with j S L
0
00 (6-1)
21
1
K
(4) (t-1)
h2
[r]
aj
V1 h2j
(t-1)
=
aj
K
D2
(e-1)
0,
D2j-1
(6-1
0
0 (0) -0Q10
(0)
K!
K
v(0
-V
1
1
0
(0) (1-1)
(21 1)
Kg
22
(e-1)
Kg
10 (8-1)
22
2/-1
1
V
2011
(f-1)
=nnd
w
1
[2j
-
0
Summing up the three heads and noting that t = o(t) (-t)), we see that the w part of h (i.c., readw (b2r-1) 1) has the update
W w/
11
Σ
(a((a,,w/'))
0(r (aj. : "-1aa
j=l
111
Σ
1-1
(a(1 + 2j - 2t h2j-1 + a(1 + 2j - 2t 1)V -
(0)
((-1))) )) + V b
h2)
21-
j=1
ww111
1)
Kay 1 -r,lay
y1
h
(8-1)
2t j=1
2t
17
Kay, W 1 - ryla,
n
21
2t
which is one step of gradient descent with stepsize 17. Moreover, is easy to see that one can
choose the marices such that maxme]3] ! lop maxme[s]
(+)
||K| (8) llop v2 and | 1lp =
[ llop = 7, || ( llop = An. Therefore the norm of the transformer |0| S v22 + (A+)) A).
Proof of part 2 Similarly, Proposition 19 shows L 12 2T2 (B2 +)/ + B. BB.
B. + +))/(A)))] steps of accelerated gradient descent gives || [ whic AGD Whidge.> alla Se.
Again, we encode the algorithm using the last token (i.e., the 2t 1-th token). Denote the first d, d + 1 ~~ 2d, 2d + 1 ~ 3d entries of h by Wa: Wo, V respectively. Starting from w w Do = O2, AGD updates the parameters as follows:
we w6-1 + (De-1 w - -
(15a)
V -- 1 + W + VK-
Vk + 1 (wh w/-1) - ----11
(15b)
W / + (w. w/
(15c)
We show that one attention layer and one MLP layer can implement one step of AGD as above. Namely, Eq. (15a) can be obtained using the same attention layer we constructed for gradient descent with v replacing W, and an extra head with
2t - 1
Q hh
- (21 - 1)²
K K(
(C))
1
for i < 2t 1 that gives vP-1 we-1 Denote the output tokens of the attention layer by h. Eq. (15b), (15c) can be implemented using one layer of MLP. Concretely, we choose w(O w( such that
wf
VR-
We 1 -I-1
WK+1
0
W 2
w
/K-
7/-1 ( wb ( )
O (W()))))))
W
vx+1
wf w 8-
e-1
0
w + Wb
28

---

Published as a conference paper at ICLR 2024
Since & o(t) o(-t) for t R, it is readily verified that one choose the linear maps such that |/ lop - 4v/2 || op = V2. Combining this with the attention layer for Eq. (15a) and noting that [ llop v2, we verify that the transformer constructed has norm |@|| S 10 + (A 2)/ A). This completes the proof of Lemma 21.
E,4 PROOF OF THEOREM
We construct a transformer that implements the following steps each time [T] starting with h h for E (a, c.,,,
fa. hag-1
h (o,b.c 22-1
Wrdde
W Wrddee
A, A, a,
hpre (a.b.)
h
h post,b
h
pre,ó
step
Wridge
step 2
A - ar,1
step 3
20-1 posticc
hz-1
A, ac.A
h
pre.d
0 pos
1el /T
2t-1 h 22--1
A Ar,a
:
0 pos
DEA/T 0 pos
(16)
where pos : 1,, Wridge is an approximation to the ridge estimator w Widgo, A, A approximations A, a Uex are approximations to Uyg == (Wridge: ar.)) +ap (ares Ar Faan which are also approximations to
a are
(w A ar.))) + avy (ar,ss Ar,))
for k E [A]. After passing through the transformer, we obtain the policy
Algo (- (- := exp (h post.c,
EA4
|| eex (h post. 11
We claim the following results which we will prove later.
Step For any
O, there exists transformer TFo(.) with
L. 2/2T
B2
(2T(B² + A) + A)T (B, By + + =O(T)),
log
X
max (D) RE[L]
S 4,
max (11)) SELL]
< 4d,
|0| S 10 +
13 + 0(11)
2+2
that implements step in (16) with Wridge wt Widggo.ll2 SE.
Step For any > O, there exists transformer TFe(.) with
- 2v 2T
B + X
log (2T(B2 + A) + A)B.
)1 -
max M 4A, EE(T.
X2g
max (( < 4dA, ||0|| < 10 A( [E[T.]
A 3
+ +V2 = -O(A)
B +
that implements step 2 in (16) with || 1 a
A, - for t: E [A].
29

---

Published as a conference paper at ICLR 2024
Step Suppose that the approximation error in Step 2 satisfies E2 b2/12(B3 + + For & > 0, there exists a one-layer transformer TFo(-) with
Forany any
L 2, max MD ) I 4A, max D S r.
|01 O(4-T(a/(7s)))4
FE
that implements step 3 in (16) with jūxx/T - vx/TT S for [A].
Denote the errors appear in each step by 53, respectively, Define for all & E [A] that
(w
Ai A al.) a (ar, A, al.)).
which are the actual values used to compare across different actions in LinUCB. Then for all [A], we have the approximation eтror
Uik
Uek
1k
Utke
Usk
Urk
T
T
T
( Vidgo,
1 Wridge: atkk) 1
a av (r.,k A, A,Ta(s) (aak)-aa -ay (ar.e. Ar aaaxs)) +
7
Bas1
aBaE2
+
(V(a A, (ark) A
V A, Aaxix))
2T min
BaE1 TBB2 N)aBBE2 +
baT
where the last line Eq. (17). For targeted approximation emror choosing 61 eT/(12,,), E2 min Te T(B + A)aBa), by (2/2 + A)] and E3 e/12, we obtain |0///T - uu/Tl S s/2 for all k E [A].
From the proof of each step, we can verify that the token dimension D can be chosen to be of order Odd)). Moreover, due to the convergence guarantee for each iteration of AGD in Proposition 19, it can be verified that there exists some sufficiently large value R > O with log R = Ö(1) such that we have [b! (a) |l2 S R for all layer E [L] and all token i [27] in our TF construction. Therefore, TF and TFe generate the same output for all the token matrices we consider, and w.l.o.g. we may assume in the proof of each step that the transformers consider are those without truncation (i.e.,
TF = T
Finally, combining Step 1-3 with a = Ō(vd) and applying Lemma 18 completes the proof of Theorem 8.
Proof of Step 1 We use the first d entries of h to represent Wridge and the d + ~ 3d entries (denoted by *) to record intermediate results for computing Wridge- Step follows immediately from the second part of Lemma 21.
Proof of Step 2 Note that
A, a = arg min XER4 2(2t - 1)
T A,x (2t-1) (x, ar,k) =: arg min L3 (x)
xER
is the global minimizer of a A/(2t 1) -ttrongyyy convex and (B² + A)-smooth quadratic function with the condition number < 2V2 X)/ X, Moreover, we have
|A a 1 < |A, 1l |ar B./A. It follows from Proposition 19 that L /2 27' + A)/Alog((1 + x) erated gradient descent finds A, at, with ||A, a A, A S..
steps of accel-
Similar to the proof of Lemma 21, we can construct a transformer such that each (self- attention+MLP) layer implements one step of the accelerated gradient descent (AGD) for all k E [A]. Denote the (k+2)d++1~ ~ (k+3),, (A+1+2k)d
[A+2+2k)d, (A+2+2k)d+1
(A + 3 + 2 entries of h2.-1 by Wa,kk Wo.k, Ve for [A]. Note that in the input vector pre.d have w.
attk: Wo v2 0g-
30

---

Published as a conference paper at ICLR 2024
For each layer E E [Z] and & [A]. we choose (2)
63,,42, K (0)
11, A1, such that for
V
even tokens hay with j S t 1 and odd tokens h2--1 1 with j S
0
((-1)
(6-1)
aj
k1 2j = 0, v ((-1)) = -
((-1)
v (L-1) 0
21-1
k1
2j
0
aj
A1 2j 1j =
(0
(8)
K
(0)
K k1
V(
(f)
42
k1
22
k2
k1
R3 21-
(0) (E-1)
1
2t
K
k3 22j
(0) (1-1)
2j
1
K
(0) (1-1) 2j-1
2j - 1 0
1
22
n)
aj,k
0
2f- 1
Q (t-1)
- (2t 1)²
K
K00
h2
K 1:4
(t-1)
k4 21-1
1
1:4
=
(2))²
2--1
(2j 1)2
k4 21-1
o,
0
0
where 17 = 1/ + A) and the values V/2 (E-1), ktt h2j v1O ((-1)) t= 2, 4 are supported on the
23-1
entries corresponding to Wa.k- Summing up the M = 4.4 heads and noting that t = (tt) o we see that the Wa,k part of h; has the update
We 0,k =
wh-1
7
(a((a 1/4 1/ - al- (aj, v "))aa
n
0.k
2f-
2F -
21-1
8--1
11
(aj. VE 1)
nn
v
(8-1)
V n
0.k
21 - 1
2t
k3
D2c-1
21-1
17
(2/, VR
11-1) aj
7/-1
Vk
+
a
21
21
2t
//6 n LL
which is one step of gradient descent with step size 17 (c.f. Eq. 15a). Moreover, can be verified that
one can choose the matrices such that ma) m=[4] Q llop = max ||
(t)
km
llop
I
22 and maxkelA
[ kl llop = max4e[A] [ #2 loop = 7) maxkela |
43 llop
S
(A +
1)n, max [ k4 lop S V2. Therefore, the norm of the attention layer
| s v2(A+ + A(A+ 3)/ A).
Following the construction as in the proof of Lemma 21, we can choose ((()) Eq. (15b), (15b) for all l E [A] simultancously and we also have || llop S 4v2, w! = v2
(E) w/O that implement
with (()) 4dA. It follows from combining the bounds for the weight matrices that
A 3
A A+3
|| < (2
B + + V + v2 + +422 < 10 A(
+ +V2) = O(A).
B +
Proof of Step 3 Denote the i-th token of the output of step 2 (i.e., the input of step 3) by h
We use the (3A + 3)d + 1 ~ (3A + 3)d + A entries of h² to record , D and
use the (3A + 3)d + A+1
(3A + 3)d + 2A entries to store additional information (de- noted by fa.,l,..., CatA) for computing Tn Pu/T,.... SeA/T. Concretely, for all & E [A], we choose
(()) Q Q k4r K(O A3, k2,
such that for even tokens h2j with j S t and odd tokens
h2j-1 with <t
Wridge 21 1
ay.8 - B
0,8
(1) (0)
h
1
K
kl
(1) (0)
h2 =
B(2j - 1)
,
K
R1
(1)
bj
2Bj
-B
0
0
31

---

Published as a conference paper at ICLR 2024
0
0
V 25
(0)
V (0)
2j -
2j
2j
0
0
ridge
( (0)
2
21
K
(1)
K K(
v
-V(1)
12
1
122
k1
A1
0
A 1
at
a
0
(1)
21
(0)
B
K
(1) (0) k3 2j
B
13 21-1
1
k3
2-
B(2j - 1)
2Bj
0
0
0
V (1) (0)
h2
2j
V
(1) (0) 2
2j
A3
0
0
A
a
h
(0)
2t
1
KM
K K3
v
v
k4
64
R3
where = (B.. + o)/A+ + 2B2/>: /x; v!ll
he
(0)
v
v@ (0) (c((c
he (c = 2j - 1,2)) are supported
on the [3.4 + 3)d k:]-th entry of hg, V w he
(0)
(0)
(c 2j 1,) supported on the
[33A 3)d + A + it]-th entry of hg-
Since (Wridge: a²,x) I W Widge l2lla,,0la & B, it follows that
(1) 0 kl 21 k1
(1)
(1)
23 12
= (Wridges ay,e) + (2j- - 11 (2t - 1)) S
for j < i. Likewise
k1 221
(1)
k1 h2
s 0 for < i. Since we assume the error E2 S S Bd and IL = A I (B.T + A)Ia, it follows that
(2 + +a) in Step 2, b. S ||a
(a.,k A, AfJank) 2 (ar,k. A, araak))
||ar, A at.k A a1,1112
b2
b²
11.
(17)
2BB2T + A)
B + 2)
2B X
a1,,k Ar A, aakk) $ a1.kk A, A-- aa.i)) + ||ar, l22lll1222 A
ar. A,' a1,1l12 S
l2- (18)
Therefore, (Q (1) (0)
11
K
(1)
21
/ ))
(0)
ar, a5. (2j - -
(21 1-1
1))B > 0 iff j = i.
Likewise (0 (0)
K
h2j < 0 for j i. Similar results hold for the k2, k4-th heads. By k3
(1)
(0)
21-11
some basic algebra and noting that t a (1) (-t) for t E TR, we see that the attention layer
updates the position for B to with the values (Wridge at,x) (a A, a.,k) for all A [A], respectively. Moreover, can be verified that one can choose the matrices such that
110
K (1) km llop B.
max mE(A]
max kE[A,,E(44
max kE[A),E[44]
Now, compute the value of S in step 3 in (16), what remains is to approximately compute a Va,tk, add the result to the position for U and multiplied it by 1/r.
12), to approximately compute Since Na,tk = (a. AT -_aa,,)) [1/T,
(ar,k, A, an,s) A ). it suf-
fices to approximate fxx = VI E [1/T, 12). For any level of approximation error Sapprox > O. let (I1, xN) E [T/T, [z] satisfy
51 = h/T, NN l
0 VoJti Vay E Eapprox; for j [N 1).
32

---

Published as a conference paper at ICLR 2024
Define the function
N-1
f(x) = 21 +
Σ 3
o - 2z)
V + vay
Note that f(x) is a piecewise linear function on /T,T1 with f(a)) = ve, for E [N]. By some basic algebra, it can the shown that for Eapprox 4/T, the difference between f(x) and f(x)
max (f(c) f(c)] = max Vag + cE|=),15+11
Vas + M+ xz)) S Capprox
tE
V rj +++ vay
when Fi++ Vij cv Eapproxl1 T for universal constant and all j E [N- 11.
Therefore, there exists function f(x) with N = O( /O(VT/Eapppeee that satisfies
max 1f(x)) - [(x) Eapprox. [1/T/2]]
As a consequence, we verify that one can implement f(Co,()) for all [A] simultancously by constructing a two-layer MLP with
Io s O(VN), ll S <O(VTN), D' AN.
Choose Eapprox = Te/a. Substituting the expressions for Enpprox into the upper bounds on the norms, we obtain
[[ llop s ((e)))
[|
1lop s (aT")( D' D' (TA(
Lastly, we can construct another two-layer MLP with weights ((( w such that implements the summation and multiplication updates
- v w (W (2) (2 ~ ~ 11- [ Va.ti.....
VattA
with 1T// U/a for all 7e E [A]. We verify that the weight matrices can be chosen with
|| 2 lop < O(),
[ llop < 0(a/T)
and D' ((AA).
Therefore the norm of the transformer that implements step satisfies
|0| < O(B + + 4.1 + 7 (a + +a/r) T(a +a/T).
This conclude the proof of Step 3.
E.5 PROOF OF THEOREM 9
By Theorem 6 and Theorem 8 with Ereal = 1/7 it suffices to show soft LinUCB with param- eter T has the regret guarantee
Σ max (ar,k."") RM.Als.L..CC (T)) s (OdVVllo(())
EM
This follows directly from a regret analysis similar to that for LinUCB (see c.g. Chu et al. (2011) or Theorem 19.2 in Lattimore & Szepesvári (2020)). Concretely, note that vi (ridge, ar.k + (w
/ (arkk A, a the solution to the optimization problem
maximize (w,ar,e)
subject to w W C, [wl(w- wiidgo.)) Ae(w - whidge,)) S a²),
33

---

Published as a conference paper at ICLR 2024
where we recall a(6o) with dg B,)) and
= a == v\B 200llol (dd (
(19)
Moreover, standard analysis as in the proof of Theorem 19.2 Lattimore & Szepesvári (2020) shows with probability over 1 1/ we have w E C, for all / E [T]. Denote this event by Ep. Moreover, let Pt, denote the probability of soft LinUCB selecting the action ar. at time for all & E [A]. For any e > 0, let (c) == { E [A] Vix max Vig S)).
Therefore, on the event Eg at time we have
max oi
Pi
Σ
(max
vig v((e) + Σ
Pr, vij - o..))
1=1
(EB)
(E))
exp (-
ED (max oij vike
S 5 +
(E))
$ 80 + 2A 1-2Axx
exp
E0 - 7
Ba Bw + 20/ 2a/VA).
where the second line uses
ED
-max Pt,k exp 5D
PLA exp
and the last line follows from that |o/z|
(B + 20 ( on the event Eg. Choosing ég
:= 1//2:11/V4T and noting that = Eu/ so/log((4 (T AT, ( + 2a/VA)) 0 (T)/T), we obtain
A
max 12
P E1
k=1
Now, the event Eg, we have
A + 51 =
A
max (w" ar,)) S max oi - JE[A]
Σ PL
Σ
Pi,b (Wk, a(,s) + E
JE[A] k=
for some Wa E C. [A]. Therefore, on Eo for each E [T]
max (w", (w",ar.) and) Σ PaN (w", ac SE Σ Pi,ke (wx- w,, ac, JE[A
E +
Pt. WK
A S 61 + Br ||ar,ell,. w" w"llA, ||a Iaa,
Moreover, note that maxje[A] awy) - a,,)) S 2B, B, and || I B./VA. Sum- ming over t [T'] and using the tower property of martingales, we obtain
T
EM-A
max (ar,k, w") - RMM AlM.Aeet.n (T)
1=1
E max (w", ar.s)
(
jE(Al
k=1
T [2
a /.s.ppllattlllaa +ET + B.T [sp)]
T'
Ba
< |
A |ac,alla;))
+ + 2.r.B. 16881
m=1
T
< 2B
Bu
+1)
(a (-- + 2B B.TP(E
VX
t=1
34

---

Published as a conference paper at ICLR 2024
1)2Ta (0C + +EI +1,
where the fourth line uses the fact that
BN A a S
Ba VX + (1 [|
|a,,l 1)
and Cauchy-Schhee Schwatz inequality, the last line follows from Lemma 19.4 of Lattimore & Szepesvári (2020). Plugging in E1 = =1/VT and Eq. (19) gives the upper bound on expected regret
T
EM-A Σ
max (ar,k: w") )- AMM.Agatinuun (T) s O(dvTTog(TT)
for soft LinUCB with parameter 7.
Moreover, the second part of Theorem 9 (i.e., the upper bound on log.Ve) follows directly from Lemma 16 and Eq. (8).
F THOMPSON SAMPLING FOR STOCHASTIC LINEAR BANDIT
Throughout this section, we use c > denote universal constants whose values may vary from line to line. Moreover, for notational simplicity, we use (--) to hide universal constants, ((-) to hide polynomial terms in the problem parameters ( b Ba), and ( to hide both poly- logarithmic terms (Mo: Co, T, A, d, 1/e, 1/66) and polynomial terms in (X+1, y= bb Ba). We
also use the bold font letter E d to denote the selected action ag at time t E [7'.
This section organized as follows. Section F.I describes Thompson sampling algorithm for stochastic linear bandits. Section F.2 introduces some additional definitions, assumptions, and the formal version of Theorem 10 as Theorem 23. We prove Theorem 23 in Section F.3 and prove Theorem 11 in Section F.4. Lastly, the proof of Lemma 24 used in the proof of Theorem 23 is provided in Section F.5.
F.1 THOMPSON SAMPLING ALGORITHM
Consider the stochastic linear bandit setup as in Section 4.1, but instead we assume a Gaussian prior distribution w (0, AL2) and Gaussian noises {Et)120 wii N(O,,) r). Furthermore, we assume there exist (b,; B₂) such that b, S ||a,,el|2 & Ba. each time E (T), Thompson sampling
consists of the following steps:
1. Computes
L-1 Pe (T Eajaf
E-1 = "I + Σ a,a
L-1
Σ
E
j=1
5=1
j=1
2. Selects the action a a with probability
P..NN(aEEE, (k= arg max (ar.j. Wi)
jE[A]
Note that Thompson sampling is equivalent to the posterior sampling procedure in our stochastic linear bandit setup, i.e., select an action with probability that equals the posterior probability of the action being optimal. We allow À to be either some constant independent of T, d, or has the form X Ao/d for some constant Ag O. The latter case is considered so that the bandit parameter vector w" has & norm of order Ó11) with high probability. In this case, we O(-) to hide polynomial terms in the problem parameters (XF+) p+1 by Ba), and Ö(-) to hide both poly-logarithmic terms in (Mo, Co, T, A, d, 1/e, 1/86) and polynomial terms in (Xi pi1 b-11 BB).
F.2 DEFINITIONS AND ASSUMPTIONS
For any actions a,1,.... a Rd we define
log P W. (k: (R= arg max (ar,s, wi) JE(A]
fa a (4, TE, rE,1)
35

---

Published as a conference paper at ICLR 2024
For any lo E [A], X1, ,1,...,,, XA E Rd, 11, Vi.....99. . R, we introduce
gk (x.,...., X. M1..... ;))) = log PanN(o..)) (xx: Xx,, 2) + yk - ys 2 0, for all j E [A]).
It can be verified that
fa (a)........ MA 1) r,,
(vrVV 1/2 ar,1.
1/2 a, (v,1) ..... ......u Kue ac,)))
For any Vi E [0, 1]. we also define the truncated pogabrobabilities
fk.m (a.1...... at,i: Me: r., log [P arg max (ar.s. wi) V vm]. jE(A]
Ikmm (XL, XA V1. VA) = log [P( (xx - X3. z) + yv yo 2 0, for all j E [A]) v vm]:
Define in addition the region D, :- {x1,..., X,, M1,...... . Me : ||x xy ll 7, for all i + j}. We verify that on the set Dyp the function Ikm is Lipschitz continuous in any of its arguments (see Lemma 24 for more).
We adopt the following definition in Bai al. (2023).
Definition 22 (Approximability by sum of relus). A function 9 F - R is (c, R. Mo. Co)- approximable by sum of relus, if there exists a "(Mo. Co)-sum of relus" function
Mu Σ
Mo
fMa.Co (z)
Cm (w [z 1])
with
Gml S Co, max || S Wm E Rd+1 Cm E R, [M]M
such that SuP2EI- n (z(z) fMo. co (z)l S.
Assumption B (Approximation of log- boog--ooterior probability). There exist Mo. Co 0 depending (1/e, 1/m,
1/12, Rs, A) such that for any e 0, (0.1), > 0,8 E (0, and k E [A], Jn.m (Xi,...s AA .1........ VA) is (E, Rs, Mo: Co) approxxmiieeeee by sum of relus on D,p with Rş := 2B, VX(1 /log(22/)) + vd) ((VAd).
Assumption B states that the (truncated) log-policy of Thompson sampling can be approximated via a two-layer MLP on a compact set with Ö(v da-riduus when À (1)) (or with ( when do/d = (1/d)).
Assumption C (Difference between the actions). There exists some >> 0 such for all instances M and any time [T). we have ||ar,j at, 2 for all S
A.
With the definitions and assumptions hand, now present the formal statement of Theorem 10 as in Theorem 23.
Theorem 23 (Approximating the Thompson sampling, Formal statement of Theorem 10). For any O < do 1/2, consider the same embedding mapping h and extraction mapping A as for soft LinUCB in E.I, and consider the standard concatenation operator cat. Under Assumption B, for
8 < (1) A 1)/4, there exists transformer B with log R = Ō(1).
D O(TT/ Ad), L L= LVV(T,, M = Ō(ATT//)) D' -O(A(T)/ Mo))
|0| Ô(T 1/4 / MoA + Co),
(20)
such that with probability least 5u over (M, Dr) PAK for any Alg, we have
log Algrs 5t) log Algo (a D.-1, 8t) S E, for all t [T], [A].
Here Mo. Co are the values defined in Assumption B with Tn = e/(4A), = 6
Su. and Õ(.) hides polynomial terms in (Xt1,r b- 1 Ba) and poly-logarithmic terms in
(Mo. Co. T, A, d. 1/80, 1/2).
36

---

Published as a conference paper at ICLR 2024
F.3 PROOF OF THEOREM 23 (AND HENCE THEOREM 10)
We construct transformer implementing the following steps at each time & E [7]] starting with h h for {a, b, c.
pre, (a.b.c)
h2
h preeta,
21
v Vr
21,1
ba
20
1
Σ a4,1
h2:-1
h 2r- pin.c
h h
step
Pr
*
step
stop
Vr 1/2 a
(21)
(Tte a(.1)
0
1/2 a1,A
h pre,d
21
pos
(pe, ar
pos
pos
(a..
21
post,
2
step
1
post.c h2r-11 - h post.e
ALA
hg
where pos == [t, pt, E, are the mean and covariance of the distribution we sample w from; Utx are approximations to u log P(j arg max (a w.)). In addition, we use h', x E {a,b, d) denote the corresponding parts of token vector h. After passing through the transformer, we obtain the policy
exp(her eepphhpet 1010) Ill
exp(V,) (dell A4.
Alga (- =
In step 3-3 of (21), we use transformer to approximately generate the arguments
(vrΣ," 1/2 ar....... VrE," 1/2 ar.A (He, ar.1) ...., (Ue, ar,A))
of the function 9k (or 9k.m) and in step of (21), we use transformer approximate the truncated log-probability Ik.. for some 7n E (0, 1) by exploiting Assumption B, C.
For any 0 50 < 1/2, define B vā(vd 2 2o1o((2/0) and the event
Esp
{max |exl < / 2r 1oT 2T /80) fllw" ||2 $ Bw). BE|T]
Then by a standard tail bound for gaussian variables ( a union bound over E [T], and Eq. (4.3) in Laurent & Massart (2000), we have
P(Ea,)) 2 1 do.
We claim the following results which we will prove later.
Step
Under the high probability event Esg, for any E > 0, there exists a transformer TFe(-) with
B + X
(2T(B2 X) A.TB.BB.R / 2r log =
L 2 2/2T
log À
X3E
max (D) [E(T.]
<4
max (1) < 4d, | < 10 + (A+2) X), FE(T.]
where Ā === r/> that implements step 1 in (21) with 7i mell2
37

---

Published as a conference paper at ICLR 2024
Step For any O, there exists transformer TFo(-) with
L= OvVT, max ( ((TT))) EE[T.]
D O(T4/4/4A Ad), | = O(T+AT)//)
max
RE
that implements step 2 in (21) such that |||/ at.k / a 2l22 S E for all & E [A].
Step3 Under the high probability event Eng. for any e 0, assume Step above are im- plemented with the approximation error less than E/ =\/ / respectively, then there exists transformer (Fo(-) with
I. = = v 2T {)/Alog((1 +1I)A/AA/V T(B2 -
max M 4A, EE[L]
max D I 4Ad, |0| - Ö(T + A)
REIL]
that implements step 3 in (21) with VT 1/2 ar, VrZ," 1/2 a/, ar,k112 S e, (Tte ac,s) (Me, ax,x) S for all k [A].
Step Under Assumption B, C and the high probability event Ess, suppose the approximation error Eg in Step satisfies Eg S Rs,/2 2 = O(VAd), and suppose the vector
V = (VFD, 1/2 a
Vrz," 1/, a (ee, ar,1) ... (Tt, a),
lies in D 2/ for any & > 0 there exists an MLP-only transformer TFa(-) with
1 D' = MoA, |Willop S VMo4, [Wellop S Co
that implements step 4 in (21) such that Po 9k.m (v)) s & forall & E [A] amd mn CE/A for some universal constant e > 0.'
To complete the proof, we in addition present the following lemma.
Lemma 24. For any n E (0, 1), > gkm (x, X.s U1.... MA) 1/2 1/2-Holder continuous in its arguments on DR, namely,
19hm (XI..... Xy.
XA, .i....., ,AA) 9k.m (XI..... Xz... XA. .X....A..
......,)A)
2A
2||x, x/ll2 2||x, x5112
< (V
111
190m (XI, (X xA, 1/1 ....... vU)) - mamm (XI, XA, .... w/.............. -...)..)))) ..A1
2.1ly; 2/21
1771
for any
(x1:...; Xj., XA.PIo....... ... H4), (X1, X'...,,, XA.DI.............. MAA E D
(x1...., XA. V1....................... Uj. ........ Ma), (XI...., .AA. ./....., ,/. ./..... .HA) Dn
for all j,j E [A].
See the proof in Section F.5.
Now, we complete the proof by combining Step 1- 4 and using Lemma 24. Let E2, Es, 44 denote the approximation errors E appearing in Step 1, 2, 3, 4, respectively, W.l.ogg, we assume E 62, E3, 64 < 1/4 m/4. Define the vector
(VTΣ, 1/2
a 1/2,,
V 1/2 a, As (Mes ar,)) (ee, a)))).
Note Mo, Co in the formula implicitly depend 1/e.
38

---

Published as a conference paper at ICLR 2024
By Assumption C and triangular inequality, we have v, v Da/2. By the Lipschitz continuity of f(x)) = exp(x) on (-o0, 1.5], we have
[exp(Tx)) A (ar,a/ st)l
- exp(((xx) AlgTs(an. P1-1, V 1711 711
S c2/22 -I (v)l + 19 (v) - mk. (v }] + 1/1,
2A² (1
228
253
2A
3/2 Es
1 (V 1 7
+
+ 1/1 =:
17711
where the second inequality uses
Ik. (v*) = log [assAAgss(ee (ar, D V v/l.
and the third inequality uses Lemma 24 and Step 4. Therefore.
exp(Or,x) 1| S Aes-
and the constructed transformer TFa satisfies (assume AEs 1) log Algrs Dr- 8,) - log Algo(ar, (D
I (a Dr-1: s1) V Vml - -(,,)) + log
())))))))
Am1
Fit Ik.m (v)l + 19k. (v) Ok.m (v*)l + Ass
(A+ 1)83:
where third line uses log(1 + x < x. Finally, for the prespecified E > 0, choosing E2, C3, E4: such that E5 I {2AA gives
log Algrs | Dr-1, - log Algo (D-1,,))) Se.
This can be done via choosing E1 = c15 E2 o2s", E3 csEt, CEE, m e/(4A), where c (i 1, 3, 4) hide values that could depend polynomially on (A, 1/n). such that Es S E4().
Combining the construction in Step 1- yields Theorem 23.
Similar to the proof of Theorem 8, from the proof of each step, we can verify that the token dimen- sion D can be chosen be of order Ó(T Ad) (see the proof of Step 2b for details). Moreover, due to the convergence guarantee for each iteration of AGD in Proposition 19, we can be verified that
there exists some sufficiently large value R > 0 with log R Ó(1) such that have || (4) 1l2 S R for all layer& E [L]] and all token i E [2T) in our construction. Therefore, TF and TF2 yield identical outputs for all token matrices considered, and hence we do not distinguish them in the proof of each step.
Proof of Step 1
Note that is a ridge estimator of w, with parameter À r/A and the noise sup, |eal I /2r log(T/6o). Step 1 follows immediately from the second part of Lemma 21.
Proof of Step 2 By the boundedness assumption of the actions, we have
S mmin (Ex)= S Fmax Σx) S T(B2 + 2).
Define the condition number Ā T(B² + X)ʃ5 and / TRXB2 + X). Using the Pade decom- position for the square root function in Theorem 3.1 and the discussion afterward in Lu (1998), have
vH( 11 E, pla)) 1/2
x//2 = (ĀLa + Σ
a
1//2
J=1
(m) (E - mlaa
(m) (n)(E) pdd)) -1 -1a,
m
= -v Ia +
Ia
+ Em
==1
39

---

Published as a conference paper at ICLR 2024
for any 0, where
(m)
sin 2
bmm) = cos?
ju
2m +
1
2m + 1
2m + 1
and the error term Em satisfies
Em llop
2m++1
1+ V T(B: + +X)+ VB
2m++1
I max 2 T +X)
2V
VI +
T(B2 + - vp/
VII
R1/4 + 22m 10mmt1111,10 22 VV( 1//1 +
max 2 T(B2 + 1)[1 +
)2m 11-1
R1/1 -
2:11 -
2m {
VT[BB +))(1
2
22m-1
Vā[c1
2
jam+1
21/1
#1/4
Since (1 + ///2//42''//20 /2+1 > c, it follows that choosing
(R) +1 max {{
1//4
[log log
2 T + X)
) log
2V 1)]} Ö(T2/4).
gives Em lop S for any 0 <<<<> 1.
Thus, using Pade decomposition, we can write
aa/aa ai,
V a
m
-1 (m)) (E, xIa)aak Lla + 16Th) (E2 (Ee- --l-aaa) aj
E
j=1
with E 2
S for all E [A] and some m Ö(T2/4). Next, we show that there exists transformer that can implement the following intermediate steps that give Step 2.
(a.be.
21-1 1
Pe
h (abbe 22
O
(,
مL a
Pe
step 2a
( a a
pos
0 pos
h (ab.c)
26-1 1
Pt
pre, (a,b.c)
h
044
E pLa) a
/ a1.1
E, ula)aa, -1
step 2b
(ula - ((m) (E, (E, - ulla)
[m) (E; - -ula)arr1
step 20
(22)
a1
Σ) ar,A
(ula + (mm) (2, (01a)
(m) (Ei (x
ula))aa
0
pos
O
pos
where denotes additional terms in hd that are not of concern analysis.
40

---

Published as a conference paper at ICLR 2024
Step 2a There exists an iitention-only transformer TFe(-) with
= max M 33,, ||0|| - + O(T)
that implements step 2a (22).
Step 2b Denote (E. - ar by qu and (ula + bim (E: (Ez- mLla) by Mm. For any 0, there
exists transformer (Fa(-) with
L 22=2//7 2T A)/Alog(() + +R)
max M = 4Am = 0( A), max FE(!]
1 < O(Adm) = ((T)//A)) |0| S O(Am) & Ö(T!/AA)
RE[T]
approximately implements step 2b in (22) such that the output component Mm q satisfies
ll Mm qk
- -aa M 9x I e for all j E [m]] and E [A].
Step 2c There exists an MLPP-oly transformer TFo(-) with
L 1, = 2Ad(m+1) = Ö(T Ad), ||illopp v2, || S vM((+m) = Ö(T3/4)
that implements step 2c in (22).
Combining the intermediate steps with the approximation error in Step 2b chosen as e/m gives Step 2 desired.
((1) Proof of Step 2a For all l E [A], we choose Ql k2, 41,42, token indices 2ÿ with j S t - and odd token indices with j S f
(
41,42, such that for even
0 a 0
(1) k1
26-1
a 0
ag
0
61
=
K k1
h = [0],
v
k1
h
2j =
[0]
Q (1)
(1)
K
K
(1)
(
(
k2
12
k1
k2
k
1
R3
26-1
(0)
(2t 1)
K
13
(1) (0)
22j
2j - 0
1
K
k3
(1)
2j
(0)
2j
1
V(1)
h3
(0) D
(2
p
at.k
0
where for each k E [A], v
v h2 are supported on the same d entries of hd. It readily verified that summing over the attention heads and k E [A] gives the updates
25
13
0g
28 1
for all l E [A]. We assume the updated vectors are supported on some Ad coordinates of hg. h2 Moreover, one can choose the matrices such that [1Q 1 1, | aa, <
1, only transformer
( ( kalppp max [1, [ =1 s 1+ p.. Therefore the norm of the first layer of the attention- ||e((] < 2 +
(1)
The second layer used multiply the updated vectors by factor of 2t namely, to perform the map
Σ, μLa)ark
(E. - ula)ar,
21
for all &s E [A], where the output vectors are supported on coordinates different from the input vectors (therefore need 2Ad coordinates for embedding in step 2a). This can be achieved by choosing Q (lop T. [ Ilop T, (2 [lop I 1 such that
(21 1)2
1
Q
1 (1)
h2-11
T(2 - 1)2
2 (1) K
h
11
T(2j 1)2
A3
(1) h22
T(2))²
0
0
0
41

---

Published as a conference paper at ICLR 2024
[ 1a--y1 L m
(2i
(2) (1)
2
[ (
(2(-1)
0
Therefore ( T 2 J..
S T + 1 and hence the two layer transformer we constructed has norm |0| s
Proof of step 2b The construction is similar to the construction in Step 2 of the proof of Theo- rem Hence we only provide a sketch of proof here. Note that
(mn) Mj qk = arg min
xER" 2(2t - 1)" x M,x 21 (x.0) (m)
qk
=:
arg min La. (x)
xER
is the global minimizer of X/(21 -)-onveex and (B² + X)-smooth function with the conditional number & s (B2 2). Since
||a]
qull2 S 10, (m)) IIM llplllllllll... < BoT(B + +
Therefore by Proposition 19 we have L 12 2 A)/Alog (11 + R)B,T + ;))(Ae))
M, la," qk 'xll2 I E. Now, it remains to construct a transformer that can implement the accel- erated gradient descent steps. Here we only provide the construction of the gradient (x) f-1 at the I-th iteration x = k.j E Fd which belongs to the output after é - 1 transformer lay- ers. The full construction of AGD steps follows from similar techniques as in Step 2 of the proof of Theorem 8. Concretely, for each layer é E [4] and k E [A], [m]] we choose
steps of accelerated gradient descent with step size rj = 1/(B2 + X) gives
J,T qxll2
o Skjl.kj2,433 K! kjj,k,, kj1,, such that for even token indices 2j with s S t - 1 and odd token indices with 8 Stt
(i
Q Q h2-1
(0) (e-1)
0
(
2s
(e-1)
a,
0
K
2 kil
h (e-1)
2--1
0,
0
v (8) (6-1) 1)
ky1 12
-1
((m)
a
(
kj1
h (1-1)
20
0,
0
k22
Q kyj'
K
kj2
(8)
=
K
(0)
kJl
kf2 - (t))
k51
1
1
1 (e-1)
- (2t - 1)
K
9
1
2kj
1
kj3
(2s 1)
0
0
v(O
kg3 21-1 =
17
ag
(m)
qk [(1
b
m5m))
p + f6m)
[ 8-1
Similarly, it can be verified that the constructed attention layer generates
(m)
" -17 VLR, = - [11 b/th) III_ + bmm) X)
b (5m))
Σ a
1 k.j
8-1
+
q na
1 1
Therefore, construction similar Step 2 of the proof of Theorem 8 yields Step 2b. Moreover, note that for the construction to exist we need the embedding dimension D O(Adm) O(TL/1 and the number of hidden neurons D' O(Adm) Ö(T Ad).
Proof of Step 2c Note that step 2c is a linear transformation from k aj M yk. k E [A].] E
[m] to Vala + Σ== (aIa + b//²) [,, pla)) -1a,"
-1o,) (ml) E, ( and we have the fact x =
(2) - (-z). One can thus choose W1- = [ l)d - 1)d 0 with D' 2A4m + 1)d
and W2 with |alloopp I vp(1 + m) that implements the linear map.
42

---

Published as a conference paper at ICLR 2024
Proof of Step
Similar the proof of Step 2b, given ///² ac, we can apprxoimate 1/2 at using accelerated gradient descent. Concretely, note that
_ E 1/2
a E
VrE E a arg min
xER 2(2t - 1) Σxx
1
1 x, Vr
ar,k
arg argggim Lg (x) XER
2t 2t
the global minimizer \/(21 1) conneex and (Ba X)-smooth function with the conditional number & S ( 2). Since
MV a a S V ' || /// a anall2 S 2/7V ( A)Ba:
where the last inequality uses the assumption in Step 3. Therefore for any > 0, it follows from Proposition 19 that [2 2T A)/Alog((1 R)2 T(B2 + \)oo/ steps of
accelerated gradient descent with step size 1j 1/ 2) gives
IIVEE,
1/2 ar,
vrE, E 2e, S 50.
Following the construction in Step 2 of the proof of Theorem 8 it can be verified that there exists transformer (-) with
L 12 (B + A)/Alog((1 +)2 (L))((LLL))// =
max M 4A,
<E[4]
max D (1) 4Ad, |0| s O(A) EE[L]
that implements the AGD steps. Therefore, the approximation error
11V12, 1/2 a/,
ac.k Vr2,
-1/2 ar,
11VTE, a VrE, E/ ac,bll + 1IVTE, "
////2ar aa. -VE2, " a
80 + vFlE. |lop E1/2 at. 21,111
< + X
| /2
ar, - / a 12 S + E/2,
where the last inequality uses the assumption on the approximation error in Step 2. Letting Eo E/2 1/2
yields l| -, a a1,A /r VrE,1/ 1/2
1/2 a1, S..
In addition to the calculation of VrI, 13-11 1/2 ai, we construct two-layer attention-only transformer
that computes (īa, a(,)). Namely, we choose 141
(11)
V(1)
such that
k1,22 k1,k2
14
0
aj
0
21-1
(0)
(2t 1)
B
K(O)
h2 R1
2j
B
K
(1) (0) R1
2j-1
2 -
B
(()
h
(0)
C
0
0
0
(2t -4
0
a B
0
Q
02 2t-1
(1) (0)
1)
K (1) (0)
k2 2j
2j
B
K (1) (0)
(
2j-1
2j -
(1 (0)
k2
12-1
e
0
0
0
Õ(T) is an upper bound of (ft ar,k) for all k E [A] under the event Es,, and e₁ = (0. ..... .,......,,,.. 1.0.....0 0) E R is the one-hot vector supported on the k-th entry. Summing up the attention heads gives the update
where = 2T + log
(Tc, a
0
43

---

Published as a conference paper at ICLR 2024
Note that can choose the matrices such that
(1) k1 llop - B,
| (2)
ki,k2lpp
1 B,
|
ki,kelop 1.
Thus the norm of the attention layer ||00||| || B
Finally, in the proof of Step 2a we construct single-layer single-head attention-only trans- former with | | ||
s T + that performs the multiplication
(The arkk) + (Pe, ar.))
21
To estimate the approximation error, note that ||at, I Bo and l - pell2 S e/B, by our assumption in Step 3, it follows immediately that (ft, ae,k) - (Me, ar,k) S e for all k e [A]. Combining the construction of the transformer layers above gives Step 3.
Proof Step By Assumption B, 9k.m (v) are (e, Rog, Mo: apppooimxiea by sum of relus for some Mo. Co depend polynomially on (1/e, 1/m, 1/1, 1/60, A). Since
LIVE2, 1/2 al Vi 1/2 122aa/a, S vABa. Mer, ac S ||el|l ||lla|a,, Ba
and Rso = B.B... + ABB), it follows from the assumption E3 S Ra,/ /2 and triangular inequal-
ity that ||vllcc S Rg. Therefore, using Assumption and stacking up the approximation functions for each coordinate [A] we construct a two-layer MLP with [| s VAMo, [[illop I Co. D' 4 such that
0
Va
W20 W W1h
:
V.A
0,
where (Vei...., VEA) is supported on h and Uuz
Ikmm (v)l S & for all & E [4].
F,4 PROOF OF THEOREM 11
Denote the transformer constructed in Theorem 23 by TFg. From the proof of Theorem 23, we have
Algts Dr-1, Se Algo Th 81)
log
under the event
E% = (max [xx| S / log log(2T/66)) U [[|w S Br.} for all t E [T] (E(7]
with probability at least 1 - 8g. where B, == vā(vd 2 logg/(//)))) Note that due to the un- boundedness of the noise and parameter vector w Assumption A may not be satisfied. However, setting d = 6/m2n) and applying a union bound gives
AlgTs (t se, for E [T] [n].
log
(23)
Algo (i 88)
with probability least 1 5/2. From the proof of Theorem 6 that Assumption A is only used Eq. (14) in the proof of Lemma 20. Moreover, it can be verified that the same result holds with Assumption replaced by the condition in Eq. (23). Therefore, have
log No T/6 + V
RA,Ales (7) RA,AlE (T)
< cI CT2V
n
log (Ne T/8) VT
2
S cT VR
n
44

---

Published as a conference paper at ICLR 2024
with probability at least - as in Theorem 6, where the second inequality follows in our setting Ereal = E = R)) Now, it suffices to show Thompson sampling has the expected regret with
T
EM
max (ack, w") AMRT (T) = O(ddTToog(Td).
&
The proof follows similar arguments in Theorem 36.4 in Lattimore & Szepesvári (2020). Define
A : B B= (rrdV//dloog//////g + log(2/6rs) + dl log(1 + TB2/Ad)
where êTS will be specified later, and recall Σ; = AL + Σ asa]. Since ||w*|l2 1|2 I /2 log( (4d/8rs) with probability at least 1 - 0Ts/2 by union bound, it follows from Theo- rem 20.5 in Lattimore & Szepesvárí (2020) that
P[[[ 2 2 log( Ad/ōrs), and wwll,, 3, for some i (T)) & ōTS,
where the probability is taken over the both randomness of the noise and of the bandit instance M.
Let be the event where l/,e - w w' B for all i E [T', and & be the event where {|l {[|w"l) "Il S
v /2 2dd log( 4d/ d/d/rs)) Then PPEE n Eo) 2 1 8TS and the expected regret
EM-A
Σ
max (ar,ks w") MiMMMAllS (T)
T
E
max (w", a(.k) - (w*,ai)) JE[A]
EE
maxx (w" ac,x) - aa)"a) [E(A)
+ E
Σ max (ww", aw)" t=
jE(A)
EE
(max (w" ac,x) aa SE(A)
+ +E2BBTT w (12
E
max (w*, ar,s) - aa)" JE(A]
+ 22,,TE[(w"(|aa(eennnn)
Since
E w ((ne)))) S (P(EEEEE)))) 22dlog(4d/ōis)
V +
P([w"" ll2 2 t)dt
(2Adlog(ad//rs)
s Badlog(4d/oTs66l + d3/2
P([ 2 t)dt
/ 2x log(1d/ōrs)
2Xdlog(4d/0rs)Orr + 2v 2d 1/2
exp(- dt
hog(Ad/6ra)
< 2Xdlog(4d/brs)0m + 22
log(4d/ōrs) 21/2 exp( -1)dt
< 2 2od (0llog(4////)0000
where the second line follows from union bound over id]. and the third line uses properties of subgaussian variables. Therefore, choosing ôTS = V/ gives
EM-A
Σ
max (ar,k, w*) RMMAAETT (T) &
T
E
(max (w",ar.)) - (r,,a)))nnn + +BB, Alog(AdPT).
(24)
/E[A]
45

---

Published as a conference paper at ICLR 2024
Now define the event &, := {lli- w 1x, - 8), then we have & E Fi-1 and nT= E. Also, we define the upper confidence bound Ut(a) := (ue, a) + B||a|| x which does not depend on the true parameter w'. Let (2o) denote the filtration generated by the data collected up time and the random parameter vector w".
Let at denote the optimal action at time L. Due to the construction of Thompson sampling, we have the distribution of a and a are the same conditioned on F: Therefore, E[Ur(at) and
(a;)(a))) [F-]]
E ((w", (w", aa) [F-1]
s B
(ww a Ui(a,) + Uc(as) - (w" an) ) 1e, Fc-1
1 E
(Us(ae) - (w", 13, Fo-1
< E |||اا|ممر 11 - wx,, Bl|a, 18s Fr-1]
s 201E arl Rcll=;11 F(-1].
Moreover, we have
|a | S Ba/Vă.
Combining the last two displays, we obtain
E
max ark) (w*, ax) Leneo jE[A]
T
<
E 28
A (Bo/V
1=1
< <BV (Ba/ X)) E Σ (1 x A 1)]
<2(BV (Ba/ 2)) VT
(|acl| 1/2 2 x.
A 11]
s 2B3 (B./ X)) VT log(1 + (TB2/(Ad)
-OdVTloogTTM
where the fourth line uses Caucch-SShwarre inequality and the fifth line follows from Lemma 19.4 in Lattimore & Szepesvári (2020). Combining the last display with Eq. (24) completes the proof of first part of Theorem 11. Moreover, the second part of Theorem 11 (i.e., the upper bound Nog) follows directly from Lemma 16 and Eq. (9).
F.5 PROOF OF LEMMA 24
For any j # k, by definition of ga.m
19k.m (xI,.... Xf...., XA. V1, yA) 9k.m (XI:..., Xf... XA: XA.V1.-. VV....,.Al .9A)
1
p((x Xi, 2) + M/ yi 2 0, for all i E [A])
1/1
- -P((XX - z) + yk y 2 0, for all a j. ( [xk - X5, z) + yb - y, 20)
(P((Xe - x),2) + 2 - yj z (X Xj, z) + y/k - yg)
+ P((x x/, 2) + yk yj
(xi - Xj, 2) + yk yz))
(P((x)) -x/, z) 2 (XL- x 2) + y) yj 2 0)
7/1
+ P((x 2) + w/ yj < 0 < (x - xj, 2) + 3/ - Us)
46

---

Published as a conference paper at ICLR 2024
- (P((2))-) - xx, (x++ -x), y1 Vj 2 0) +(((x) - -x, S (xt - -x,,2) + yv yj s =0)
1/1
Note that conditioned on Xj, x we have
P(I (x, - x), 2) | S ||3 - 2|1l2221og(2/11) 2 1 51
for any 81 O. Therefore we further have
19k.m (X1..... XYj,.... XA. 12.....,y44. - 9k.m (x1..... X X1. Vi......... yAAl
P [5 X,, + Us Uy E [-1|x, xx| 212 2 1og(2/))) ), ||; 2/ll2//1og((/6))) +51
71
P((x+ - xj. 2) E A) + +d1
sup
AEF ((A)=211xx
x/11a (2g((2/)1)
2 2||x, x 21 2 1og 2/)1)
2||x; - -Xyll2 51)
51) 1
22m
for any 51 > 0. where the last inequality follows from the fact that standard Gaussian has probability
density less than 1/v2x everywhere, ||x; xyll2 I 17 and 1og(2/61) I 4/81². Choosing 81
2 2 -x, gives
1A
19hm (X),.... ,,,,,,,,,... Xj.... XA. XA.B1.... :1........ MA) - 9kmm (x1,...,, xy.... .A, 01,..., :.)]
2/|x, x²) 12 + 2|1x/ x312 12
S V
17
Similarly, for Xk o X we have
19h.m (X1...., Xk. XA. 81. ......... MA) 9k.m (x1...., Xg. .................... XA, Vi...... :B)]
(((x Xi, 2) + y/ Vi 2 0 > (x4 Xi 2) + ys Vi, for some i E [A])
m
+ P((x - x., 2) + yh y S 0 3 (x) - xx, 2) + y/ yi for some i E [A])
1 (P((X)= X;, 2) + yk yı
(X} X1, 2) + yk - ya)
+ P((x x,, 2) + Wk y, < 0 < (x/ x,, z) + U/ w/.
[p((x X 2) (x X z) + yi 2 0) + P((x X/e 2) <(x; - Xx. + yu S 0)
- max
1/1
Following the same argument, we have
Igk.m (X1,..., Xk. XA. 1/. Vl........ .AA) - gh.m. (Xi,..., xk. XA. .......A .A..I......2 MAA1
2AA 2 2||xx
xt,l2
2||xx
x²y l122).
7/1
1)
11
Likewise, for any j + we have
19.m (XI, (XI..... XA. Mj. .UA) gh.m (XI, XA, ........... V/y.... , Ma))
|(((x - -xy; 2) E (min(ye - 1b, - ViJh max(vx Uo Dk - 1/ }1)
1 sup
P((x xx, 2) A)
1T1 ACF AEF.)(A)=22))
-1/1
1 2ly; 1/1 <
2/y, 1/21
111
v 2m
7/71
47

---

Published as a conference paper at ICLR 2024
and
Ik.m (xi,..., X.A 01. 11....... :.s............................. .......99))) 9h.m (xI...., X.A yz. ................................ p6..... ........................ 9.)l
[P((x Xz. E [mintt/ M/J uh Vi max(uh M1 y/2 2/ b1
A
sup
P((x= xx,,2 A)
?パ1
AEF.)(AA=2(u))
A 2/y;
2. Uj V
<
?
V2TT
1/71
G LEARNING IN-CONTEXT RL IN MARKOV DECISION PROCESSES
Throughout this section, c > 0 denote universal constants whose values may vary from line line. Moreover, for notational simplicity, use O(-) to hide universal constants, O(-) to hide poly-logarithmic terms in (H, K, S, A,
This section is organized as follows. Section G.1 discusses the embedding and extraction formats of transformers for Markov decision processes. Section G.2 describes the UCB-VI and the soft UCB-VI algorithms. We prove Theorem 12 in Section G.3 and prove Theorem 13 in Section G.4.
G.1 EMBEDDING AND EXTRACTION MAPPINGS
To embed MDP problems into transformers, we consider an embedding similar to that for linear bandits. For each episode & E [K] we construct 2H + tokens. Concretely, for each t E [T] in the k-th episode, we write H (k: 1) + h and construct two tokens
ag.h
0,+1
h ((t-1)+
hgg 1 he
Fk.b
8k.
h ()++)
0s
(2((-1)+*
0
h (1+
h2-144
0
1+
1+
po - 1)++
pos 1+
where A 0 are represented using one-hot embedding (we let 8k 0s).
h used
to store the (unnormalized) policy at time given current state 8k,6, 0 in hd denotes an additional zero vector, At the end of each episode k. we add an empty token
b = h emp == [o (05
to store intermediate calculations. We also include in the tokens the positional embedding pos; (k, h, u, i, i², 1) for i E [2T' + K]. where v := 1(b;=0) denote the tokens that do not embed ac- tions and rewards. In addition, we define the token matrix H, -h.....,b2- for all E (7).
IR Dx(--1+4))
Offline pretraining Similar to the bandit setting, during pretraining the transformer TF, takes in HeTe Hy as the input token matrix, and generates H H.ost =
TF (H) as the output. For each exp(hpant phstt.
time 1 E [T]. we define the induced policy (- == Texp(h postre
-)-
1o-t--144 1plill E A whose
i-th entry is the probability of selecting the i-th action (denoted by the one-hot vector e;) given (D) We then find the transformer i θ by solving Eq. (3).
Rollout At each time E [T]. given the current state 8; and previous data Di-1, we first construct the token matrix Here E D that consists of tokens up the first token for time t. The rfll.t
transformer then takes H as the input and generates Hpost = TF. (Here (olle))
Next, the agent
exp(hpout. 111k m11
selects an action E A following the induced policy Algo(: D = and observes the reward and next state 8111 (8t+1 ~ M41 if is the last time step in episode).
Texp( Posph 1) 2ll E EA4
48

---

Published as a conference paper at ICLR 2024
Embedding and extraction mappings To integrate the above construction into our general frame- work in Section 2, for t = (h H] + h, we have the embedding vectors
h(s,) = h2(e- 1)+ h(ag,r)) = h2--144
For N 2 1, write
T(N + 1)/2] = (Ai 1)H + hN
for some hy E [H), and define the concatenation operator
cat(h hx) [h...... h2n, h hmmp hhhh+11..... had, h2mp, hhhh+1,........... hy] E R R-N+EN-1
where we insert an empty token h (i.e., a token with h (a,b,c) = 0) at the end of each episode k. In this case, we have the input token matrix
H
H Home rll.. cat(h(s1), ((( "apri.,................... ((1-,,,
E (x((C-1))))
the output token matrix H Hpost and the linear extraction map A satisfies
A h A (2(1-1)+4 = Post,c 2(t- 1)+
post
G.2 UCB-W VI AND SOFT UCB-VI
We show that transformers with the embedding in Section G.I can approximately implement the UCB-VI VI algorithm in Azar et al. (2017). Namely, UCB-VI VI implements the following steps:
for each episode E [K] and each step H....,
Compute the estimated transition matrix (s'lss a) == Ni(s,a,s))
(a.avvi where (s, a, s') de-
notes the number of times the state-action-nexxtstaae tuple (8,a, a, s') has been visited in the first k episodes, and N₂. (s, a) = Eg N₁ (s, a, (we assume Nar(s, a, 8') 0 and let NH a) be the number of times (s, a) is visited at timestep H).
2. Calculate the estimated Q-function
(n a) min (s) an(
En(ss') ())))))((5)))
where the bonus bn (s. 2H
(Sg(SAT/8)
VH+1(s) 0 for all S and Vi((s) Nu(s.a]VI
maxaeA (s, a).
Throughout this section, we choose the small probability 8 1/(KH).
During policy execution, each step h [H], UCB-VI takes the greedy action an arg max, (s). a) and observes the reward and next state (rn $1+1). To facilitate pretraining, in this work we consider a soft version of UCB-VI, which takes action following the softmax policy
exp(Qn a)/T)
nh (alss)
eep (sh, a)/T)
using the estimated Q-function for some sufficiently small r > O. Note that soft UCB-V VI recovers UCB-VI +
G.3 PROOF OF THEOREM 12
Throughout the proof, we abuse the notations h for E {a. c, d} to denote the corresponding positions in the token vector h,. For any t' [T], let ke(t),, hi(t)) be the non-negative integers such that f' = (kk(t') - 1) + h(t') and h(t') E [H. For the current time t, we use the shorthands = k(t), h = h(t). For a token index e [[2H + 1]. let ((t). h(i) be the episode and time step the i-th token corresponds (for the empty tokens set h H 1). Given the input token
49

---

Published as a conference paper at ICLR 2024
matrix roll,s construct transformer that implements the following steps on the last token. 22(--1)42 1)42 h pre, 1)+=1)+R for * E {a, b, c, d}
ba..ee 1)++ (1,,,,
h h tab.c 20+-1)+8
(8,
h pr..ce
NH(s,a,8) (1(5,))
2(+-1)+ (s
Th preca
h pre, -01-1++
Qu(s,a,
2(1- + h ((t-1)++ pre,d 2(1-1)+Ak
step
step 2
step 3
Vi(s)
NH(s.a,8)) (1(8())r((s.a)
0
(s)
Nn 8 (s,
2t 1)+
*
0
pos 1)+k
0
(1)
(r,,(
((--1)+
h post,a
Qn(sr.a))
2(t- 2(t-1)+4B
4
2((-1))+
(25)
hpact,c
(2((-1)+
ax(
h post,d
2(t- 1)+
hd +k
where Na(s, 8'), (s 8'), Qn(s, 8.) E R
Nn a) E RSSA Vh((s) E RR for all h E Given the current state St, the transformer
[1], and denote additional quantities in h2 1)+k* TFo(-) generates the policy
exp(h (x 2(1-1)+k)) EA4
Algo (- ==
|e post,c 2(++))+))
We claim the following results which we will prove later.
Step There exists an attention-only transformer TFo(-) with
4, max ( < O(HS2A), |||| s O(HK + HSAA) EE[L]
that implements step in (21).
Step 2 There exists one-layer transformer TFo(:) with = 1, M < O A), D' HS),, |0] = Ö(HS A+ K3 KKH) that implements step in (21).
Step 3 There exists a transformer (-) with
L 2H1, max ( 2SA, max (0 < 3SA, |0| < 0(1 + RE[L]
[E[L]
that implements step 3 (i.e., value iteration) in (21).
Step There exists attention-only transformer TFo(.) with
- max ( O(HA), ||0|| < O(H(K A) 1/7)
EE[L]
that implements step in (21).
From the construction of Step 1-4, we verify that one choose the constructed transformer to have the embedding dimension = O(HS"A) A). Moreover, due to the boundedness of reward function, Q-function and the fact that the bonus b(s, a) S Ö(H). we verify that there exists [L] and all token i
some R > 0 with log R Ó(1) such that ||h]d] ||2 < R for all layer [K2 + 1)]. Therefore, similar to what we do in the proof of Theorem 8, 10, we may w.l.o.g. consider transformers without truncation (i.e., R oo) in construction of step 1-4 in (25).
50

---

Published as a conference paper at ICLR 2024
Proof of Step two steps:
We prove this step by constructing transformer that implements the following
Step Ia For each t' f with t' (k' 1)H + h for some h [H], add sk (aw N rk. from h 12(1-11++ and h hg h2
Step 1b Compute Nn(s, a,ss,,,,),,, 8'), (s, a) for h [M] and assign them the current token h 1(((-1)+
For step la, we can construct a two-layer attention-only transformer with 21,2,3 that for all < 2(t 1) + k
(())
V1,2,3 such
1,2,3
Te(i)
Q h (0)
h
B
K
h
(0)
k(i)
V
h2g
(0)
0
1
1+3
2(
1)
i
-1
0
agech!
v 20 1+*
TW Os
0,
where we choose B and (0) are supported on some entries in h (0),d, Moreover, we choose Q [1]
Q
Q V
v(
v(2) and K.! K(
such that
B
B
K. hh (0)
T(ii)
K
h
Te(1)
i+2
i i+1
We verify that QQ 1
||
lpp = 4, | Ilop
for * E [3]. Summing up the heads, we
obtain the following update on subset of coordinates in h (0).d
22'
3
2/'
0 Os+A+1 - OS+A+1
Σ
(0)
+
h (0)
V,n, (0)
(0)
2V
(0)
+ 3V
2t' + 1/
21
-(V(")
(0)
v
(1)
h
(0)
(0)
2t' + Rv'
+
[ag.h
21' + k'
[8kr
i - k (i.e., all tokens that embed the action and reward) since O, it follows that no update happens on the tokens in which we embed the action and reward (i.e., the corresponding part of hd remains zero). Moreover, it should be noted that no update happens on tokens with h 1.
Note that (Q (0) K (0))
21'
We then use another attention layer to multiply the updated vectors by a factor of 21' + R', namely, perform the map
1
Tak.k" Tkeck" [Skr
Top.k' Pkrth" LSM,
21'
+
51

---

Published as a conference paper at ICLR 2024
where the output vector is supported on coordinates different from the input vectors. This can be (2H 1)
achieved by choosing |Q1
(2)
llop
| llop I (2H + 1)K, V llop
such that
1 11) - (2H + 1)K? Q Q 1
O ah
K n
11)
1
(2H + 1)
2
21' + kv
TW " SKIM 0 (26)
0
(2) and noting that (Q)²) 1000' h Qi n when j = and otherwise 0.
For step Ib, we show that it can be implemented using a two-layer attention-only transformer.
To compute (s, a, s'), in the first layer we construct M 10H.S² heads with the query, key, value matrices {Q (1) 110 [$].j [4], h E [H]
(1) sahh (Viih,,) 1r0 such that for all s 2(1 - 1) + and i, E ijkh,eJe=1
ijkkh s=1
U 1)7
Be
(T()).((()-
Bej
((()),(())-1
(1) (0) Gijkh,
Be
K
( jkh,
h
(T(().(())
(1)
iihh
h
jjk
N
1
-3B
I
h(1)
h
where we choose B = 2H and eix denotes the one-hot vector supported on the (i.j. k)-entyy in (s, a, s'). We similarly construct
B( (v 1)7
Be:
ST()),((1))
Bej
((())T(A)))
0,10
h Qijkh,2
Bex
K (1)
jjkh.
h
(0)
*K((),(()
V
ijhh
h
(0)
Gijk
N
1
-3B
1
-((1)
h
BBu 1)7
Bei
$T())
Bej
(E()),T())-1
ijjkha h
(1)
(0)
Bo
K (1) (0)
ijkh,3 h,
(T(().(())
(1)
iihh,3 h
h
(0)
ijk
N;
1
-3B
1
TO(i) -1
-h
B(w - 1)7
Be
$6(0)
Bej
(E(),T())-1
0,00) h iijkh,4
(1) (0)
Be
K (1)
ijhh,4
h,o0)
$T(1),T())
(1)
ijhh,
h
00)
Gijk
N
1
-3B
1
Ti(i) -
-h
BB(; 1)7
Be
$E())
iikh. ((0)
Bej
K ijkh.5
(0)
(T())
( (0)
ijkh.5 h
eijk
Be
*k(6),T())
-3B
52

---

Published as a conference paper at ICLR 2024
Summing up the first five heads, verify that such attention updates the token with h 0 and has the form
0-0+ Na(i,,, kjej
on hg. where (n k) denote the number of visits to the state-action-nexttstaa tuple (i.j. k) at time
step before token For * E [5]. we choose v((1) such that
v.(1)
ijkh, and ijkh, K (1)
ijkh,
jkh,++
be
h
((0)
K(1)
ijkh h
(0)
(1) +5
h
(0)
ijkh.
B k(i)
K
+khh
(1)
h
(0)==
k(i)
B
which adds positional embedding about the current episode k(i). We verify that summing up the sixth to the tenth heads gives the update
0 (N.(i...k (i.j.. k))e
on h for s 2(t 1) + k with h O. Therefore, combining all the heads together we have the update
- 0+ Nn e) for all i, E [S].] E [A], E [H]
on he for S 2(1 1) + k with he 0, particular when = 2( - 1) + Moreover, notice
that the matrices {Q (1) 120 (K (1) 2 can be constructed with the operator norm less than
@oijhh,JJs=11
10B 10H, and ( (1) )mm with the operator norm equals 1.
jjkh,)
Following similar construction, can also compute (s, a), Nn a),, a) for all h, 8, a, s' on different supports of coordinates in h via adding additional M O(HSA) heads to the attention- only layer.
Next, we construct the second attention layer to multiply the token vector by the index number i as in the proof of Step la. The construction is similar that in Eq. (26) and we omit it here. Moreover, note that Step lb can be implemented with the embedding dimension D < O(H A) as we need (1) dimensions for each quadruple (i.j. k, h). Combining Step la, lb concludes the proof of Step 1.
Proof of Step 2 After Step I, for the current token i 2( 1) + k, we have (s, a, s'), a')s), store the number of visits and rewards, note that
rn,, a), Nı(s,a), a), Nn(s, a)ra(s,a) lie in h for all h E [H]. Given these vectors that
rn(8,a) an(s,a)n Na(8, a) VT
when N, a) 2
bn (s) 2H log(SAT/A)
Nn a) v1'
Nn(s, an(s,a)VI V
Rr(s,a,s')
Therefore, we may compute Pr. bn via using a transformer layer to implement the functions (a, = a f2(y) 2H (og(SAT/8)
ff(
VI H1 for z, E {0) [K].
vV1
We demonstrate the computation of Ph(s, a, s') (i.c.. the computation of fı y)) here. We start with constructing an attention layer with M O(H52A) A) heads such that implements + x for = ahss,, ()) For Nn(s, a, s'), this can be done by choosing || 110 llop S
K, K Wiillop < K, (1) ijkh llop = such that
K
0
( hi
-i
K
K
K [Nn(e, ej, ex)
j
NN ej, ew)
0
53

---

Published as a conference paper at ICLR 2024
where C₁, ek e₁ denote the i,j-th states and egj denotes the k-th action. Similarly, construct HSA additional heads to compute Ns(s,a) a)² for all possible a.
Next, we compute the exact values of P(s, a', using MLP layer. Namely, construct
w = 12 11 w
w
(1 such that for all h, 8, a, s', on the corresponding 23 22 21
vector component have
1 an(s,a,s))
Nh(s,a,s')
(Niis, a, 5) 5
= Nn(s,a,s') + K 2K 2KNn(s,a,s)
an²
Nn(s,a)²
(Nn(s, - K)²
a) + K2² 2k
1 Nn -
W h
1 1 (Nh(s,a, x)² (Nh,s, a) yy)
(Ni,s, - KY² (N,, a) K}
where x,y E {0} U [K].
Moreover, we construct w so that on the entries corresponding to h,s, a, s implements
/ o (w (") - (11 (Nh(s, (N 4)²) MV
(x.aa,,s)) Mr(sa)vi
can be verified that we can find such ((1) w with
| lop S // Iop W op S (0(3) ((K)=))(()==((K))
[W ( op I O(K), and the number of hidden neurons D' O(K Simlarly, we can compute f2(-) (or fa(-)) exactly following the same construction but with a different w that records all possible values of f2(-) (or fa(-)). Combining the upper bounds on the operator norm of the weight matrices, we further have ||o]| < Ó(H A K + KH).
Proof of Step 3 Given VH+1 VH41 O, show the there exists transformer with
L = 2, max ( < 2SA,
max (1) < <3SA ||o|| < O(H + SA) EE(T.)
KE(T.)
that implements one step of value iteration
(1,, @n(3, a) m (()
Σ Ph (m)))), S'ES
V,(s)) max (8, a) aEA
for some h [H]. Namely, start with constructing an-attention layer with M 2SA and (Q (9 12 [ (I)
1
(1) 1 such that for all i < 2(t 1) + k
B
i
-
K h..
B
Vi. h
(0)
ic
(i+z(.)
(
(1) (0)
-i
K!I)
K
(1)
V(I (1
ij.2
j22
1,2
Vo+101
54

---

Published as a conference paper at ICLR 2024
where B 3H and
Ij Qh R&A is vector supported some coordinates in h reserved for Qn-
Moreover, have ||Q ijj. llops
(1) K (1)
jhhallop
B, I uh.allop 1. Since
( (Vi+z() (+), Ph+1 8, s,)))
- [
S H
as (8) [0, H| and Pn+ (((s,0) a))
1. it follows that summing up two heads gives the
update for i < 2(t 1) k
- 0 |o(
(0) knn a 74 K(
(Q
Denote the resulting token vector by h (2) Moreover, we can construct a two-layer MLP with
/ lop O(H,, [ 3, D' = 3SA
such that for any state-action pair (s.u) 8 A on the corresponding coordinates
bn Eves Prlo'' (s))) w( = a)( + an( + Eves Pa(s' ( H an(s,a) (s) a) Eves Pa(s' ((a
and
w a (W(("(())) - o (s) + Σ Ph(s' (s)))
-((n((, ())
Σ Prll'' (s)()) H)
s'ES
+ + bn
Σ a'ES
Ph 8, (())(())()
max(min(H.nmss.c) by (s,a)+
Σ
Pi(s' s, (s)))) (n((s))
s'ES
Denote the resulting token vector by h,,2), Next, we construct second MLP layer with
( 2, | < VA, D' = AS
such that for any S on the corresponding coordinates have
an)) @n(s, Qn(s,a))
@n(s,as) (n(s,
where aj denotes the j-th th action, and
! (w (Qn(s, a1)) + Σ - - -n(s,
j=2
max (s) = (s). a
Using the upper bounds on the operator norm of the weight matrices, we further have ||] < O(SA H). Combining the steps concludes the construction in Step 3.
55

---

Published as a conference paper at ICLR 2024
(1)
Proof of Step 4 start with constructing an-attention layer with M 2HA and {K] [1)
Jh. [V/ jh.sle=1 such that for all the current token 2(
1) + & and s i
F(0) i
Q aj)
0
(1) jh,1 h
(1)
(0))
B
V
1h, 4
(0)
iejh
jh.1
B
j
0,
S1 (())
Q
(
jh.2
h
(0)
K(1)
Jh.2
=
K(1)
Jh.1
v
jh,2
v
1h,1
B
where we choose B 2H and V h
(0) is one-hot vector supported on some entry of hy. We verify that summing up the heads gives the update
0 Qnk ej)
for all h E [H].j E [A]. Moreover, have 1
lop S 2H, |K Th S 2H, 1IV
Jh.s ol S 1 for
1, Through this attention-only layer, we extract the values (n(s, az) for all [M] from the Q-function.
Similar the proof Step 1b, we construct a second attention-only layer with attention heads (2)2
(O jh 1 (K jh.s
{ (V(2) ) that jh,
1
2 (2) (1)
jh.1 h
Fi(i)
h
(L)
1
V2
h
(1) =-
(
=
jh.1
B
jh,1
B
j
T
17
(222)
K K(2k
h
1
V
V2
B
jh.2
jh.2"
j
where b() update
are supported on some entry of he for 1, Summing up the heads gives the
jh.1
Σ
(s)h.)))
s=h()+1
Similarly, we can construct attention heads (0(2))14 14
(0 1s (K
(v
1 that implements
7())
Σ (.)...))))
Moreover, we construct
Q
1h.5
K
h.,5
(2)
Jh.5
with
2 2
h
K
(2) (1)
h
B
V
h
Jh.5
Jh.1
B
Lj
O,
that implements
H
Σ
(s
Therefore, summing up the M = 5HA heads we obtain the update
1 (n
0.1
56

---

Published as a conference paper at ICLR 2024
Note that 1 S 4H, Kinhlllo 4H, II jhisllop 1 for E (5].
Finally, we apply an attention-only layer to implement the multiplication by factor of i/T using similar construction as in Eq. (26) with | llop OHH)), |||] llop O(H)), |V llop = O(1/T). and assign the resulting vector ((sk,.))) to hf. Combining the three attention-only layers completes Step 4.
G.4 PROOF OF THEOREM 13
By Theorem and 12, it suffices to show the regret soft UCB-VI satisfies
E ([ (ecc (T)) < ÓHH' + H²
for MDP instances M, where 1/K and Ó(-) hides logarithmic dependencies on (H, K, A). Throughout the proof, we may drop the dependence on M for notational simplicity when there is no confusion. For each episode l [K]. let N P. Qh Vn bf, denote the corresponding quantities Nh. Pn. Qh, Vn, bn introduced in UCB-V VI (see Section G.2).
For policy and time step h E [ we define the Q-function Q and the value function Vi"
H
(T(
= E
(82, az) |8h =8,00 aa
t=h
H
(s) = E
(s,, a)
8h = 8, T..
(= sm.h Hm to denote the policies given by UCB-VI and soft UCB-VI in the k-th episode, respectively. Note that we have VM( Eg (()((s) (s) and cumulative the regret
We use x ( T
K E[KVM(TT)) - A.M.AN..ucevv( (T)] = VT (SK.1) V1 "I (SK.1))
km
where the expectation taken over the collected data
Dr {( dk Tk. (1,,)) PaUCBVI(T))
For any function [8,,), a). abuse the notation f(s, x()) Enwa (f(s, a)). Lastly, define
Q T (.)) @h (s. (.),)
Esfuaxx
max
kE|K kE
We claim the following which we will prove later
Exfmax I Ar.
(27)
The proof follows from similar arguments as in the proof of Theorem in Azar et al. (2017) (see also Theorem 7.6 in Agarwal et (2019)). Hence we only provide sketch of proof here. First, from the proof of Theorem 7.6 in Agarwal et al. (2019), it can be shown that
(x) Vi" (s)
for any k, h, with probability at least - 6. Thus with probability at least 1 - 8 for all h E [M]] k e [K]
ViT (SK.N) - Vn -/-
$ Vi (sa.h) - -iTTon
(SK,n Th)()) 2. (sk T (-)) + Q (Sk.h. (-), 0 (sk.h, x.)
(s T 2. (sk.b, mem. (.)) Estiaxx
57

---

Published as a conference paper at ICLR 2024
= (sk,h: (k 2n (sk,, ak.h)) + -MD k.h + 6sfmax:
where the first equality uses V/2 (sx.h) = arg max, 0% (sk,h aknn) ) = 2f (-)) and in the last line
MD
(1 k.h = 10% h mmm, Qh hs (mm,(())) - (sk, axhh) Qi (St,hh bx,
Note that for any fixed h [H], (MD (1) K ing the proof of Theorem 7.6 in Agarwal et al. (2019), we further have
k= is bounded martingale difference sequence. Follow-
Vr (8k.b) 1T- (sk.h)
2 (k,,) 0T (Ek,)) + MD
k.h + Esfmax
(1 H [1th (nttneshh+1) 1
V 26b (st,, akhh)
+ CoLNH + MD
(2) k.h
+ MD k,h + Estmaxx
(sp.h. (ggh)
with probability at least - o for some universal constant Cp > is some universal constant and
0, where Lo = logg(AAK//)),
2D (2) k.h Pa Sk.hi (k.h)) (V - vntf)) (V Vit (sk,bbt))
a bounded martingale difference sequence for any fixed h E [HT. Using the recursive formula and the fact that (1 + 1/H)" e, we obtain
K
E
(sk,1) (8x
k=
< cE
K
H
28/(skh. akhh) +
Co LoL
+ +2D (2)
k.h + MD
(1) k.h
+ E Σ
H-1
(1+ )h Estinax
1
(St,,, an,h)
H
k= h=
h=0
coLoH!?S
<E
(5 (s.ss)
+MD + MDD k.h
ck HAT
(sk,h: ok
O(H + S² AAT
s ŌCH'SAK + 5²)
where 0 is some universal constant, Ó(-) hides logarithmic dependencies on (H, K, S, A), and the last line follows again from the proof of Theorem 7.6 Agarwal et al. (2019), and the assumption that r 1/K. We omit the detailed derivations here they similar to those in Azar et al. (2017); Agarwal et al. (2019). Therefore, we conclude the proof of the first part of Theorem 13. Moreover, the second part of Theorem 13 (i.e., the upper bound log No) follows immediately from Lemma 16 and Eq. (10).
Proof of Eq. (27) By definition of Q and o, ssm.h have
expl ())
(h(,, (+) - (s, m.m. (-)) max (h(
En (s,
exp(Q)(s, 4)/T)
mmax (5( ( (s. a)]
E. (s.(()(ss,)/7)
5 exp( (s, a)/T)
[max (5 a) - (s,
exp(max, (5
< A. (suptexp( (oopp(t/77)1 t/T) < Ar.
120
58


Reviews:
[{"review_id": "EPTl7BxcOi", "reviewer": "Reviewer_ycZf", "summary": "This paper addresses the theory of supervised pretraining for in-context reinforcement learning (ICRL) via transformers. \n- It first proves that supervised pre-training will imitate the expert algorithm, with a generalization error that scales with the model capacity and a distribution ratio between the expert algorithm and the context algorithm that is used to generate offline trajectories. \n- It then transformers can approximate RL algorithms including LinUCB (w ridge regression), Thompson sampling (w ridge regression), and UCB-VI.\n\n**Update after rebuttal and reviewer-AC discussion**\n\nAfter reading the authors' rebuttal and discussing it with the other reviewers/AC, I am partly convinced about the connection to Transformer, though a stronger connection is more desirable. I still had reservations about how useful the bound is, but this is okay as the first step towards the theory and inspiring the community to continue working on it. I am maintaining my score, but I think this is generally a decent work and it would be okay to see it getting accepted.", "questions": "Please see my questions as in the above.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The theory of supervised pretraining for in-context reinforcement learning via transformers is an interesting and important problem. This paper is among the first to explore that area. \n\n2. The paper shows a comprehensive understanding of the existing literature via the related work. \n\n3. The paper shows technical depth in the proofs of Theorems. \n\n4. The paper is technically sound and well-written.", "weaknesses": "1. My main concern is that the two main sets of results are not as dependent/relevant to Transformers. \n- For Theorem 6 (supervised pretraining can imitate the expert algorithm) -- this seems to be a general result that is not specific to Transformers. In fact, any supervised pretraining models that satisfy the approximate realizability assumption should suffice Theorem 6. Is this true? \n- For the theorems in Section 4, the paper mostly constructs a specific Transformer structure (e.g., 3-head attention + MLP with a specified QVK) and says that this certain Transformer can mimic the gradient update in the no-regret RL/bandit algorithms. My understanding is that the Transformer is an overkill, and this can generally apply to NNs that have similar structures. Is this true?\n\n\n2. The idea of proofs of Theorems in Section 4 seems essentially similar to that of Bai et al. 2023, though with nuances. \n\n\n3. If I understand correctly, the distribution ratio $\\mathcal{R}$ in Equation (6) can be arbitrarily large. This significantly limits the usefulness of the bound. In fact, in many cases (except the Algorithm distillation case), this ratio should be sufficiently large as Alg_E and Alg_0 are usually different, and the difference is accumulated over the entire trajectory. \n\n\n4. Minor issues:\n- $M$ is used as both an MDP environment and the number of heads\n- Not sure if I missed this -- $n$ Equation (5) and onwards has not been formally introduced. Is this the number of trajectories/samples?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ty7PLt2JRJ", "reviewer": "Reviewer_HUY7", "summary": "This paper provides a theoretical framework that analyzes supervised pretraining for ICRL. It provides two proofs: 1. the supervised-pretrained transformer will imitate the conditional expectation of the expert algorithm given the observed trajectory; 2. transformers with ReLU attention can efficiently approximate near-optimal online reinforcement learning algorithms.", "questions": "1.Would different pre-trained method lead to different results?\n\n2.Can learning-from-scratch DT figure out similar conclusion? What’s the difference?\n\n3.Will your conclusions help to inspire the empirical study? Such as figuring out better DT-based offline RL algorithms when using D4RL datasets?", "limitations": "", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Their original theory proved the feasibility of using a supervised pretraining transformer on in-context reinforcement learning and provides a quantitative analysis framework.\n2. The theoretical analysis demonstrates the important role of supervised pretraining Transformers in in-context RL.", "weaknesses": "1. Since my research does not focus on the theory of RL. Therefore, I give some general comments. \n2. I am not sure the conclusion is a little bit narrow. The conclusion is limited in a supervised paradigm, using pretraining, and in an ICRL setting. It seems that these conditions restrict the conclusions to be general.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k5Zk7sKJt0", "reviewer": "Reviewer_fr59", "summary": "**Update after rebuttal:** I consider all issues raised by all reviewers sufficiently addressed by the authors. I stand by my original recommendation to accept the paper and have raised my confidence accordingly.\n\nThe paper theoretically investigates the conditions under which Transformers, pre-trained with offline RL datasets, can be trained to perform In-Context RL (ICRL). To this end, the paper distinguishes between an offline algorithm, used to collect trajectories, and an expert algorithm, which is evaluated on the offline trajectories and the goal of the transformer is to learn to imitate the expert algorithm. The main result of the paper is a generalization bound that depends on model capacity (somewhat expected) and on the distribution ratio between the offline algorithm and the expert algorithm to imitate; implying that ICRL is possible if the model capacity is large enough and the online algorithm performs “close enough” to the expert algorithm. The paper then theoretically analyzes three RL algorithms in detail (LinUCB, ThompsonSampling, and UCB-VI), concluding that Transformers are able to learn these and then perform them in-context. The theoretical analysis is complemented with simulations that support the theory.", "questions": "1) (How) Could ICRL be used to surpass the performance of the expert (e.g. going from human chess as the expert to superhuman chess)? Most RL algorithms have fairly simple parameters to increase performance (train longer, lower discount, deeper backups, more unrolls, …) - how could this look like for ICRL. Or, alternatively, is it possible make a theoretical statement showing that ICRL performance is capped at exactly the level of the expert (even if we used ICRL to generate a new dataset of supervised trajectories where it slightly surpasses expert performance)?\n\n2) After Eq (2): “all the algorithm outputs [...] along the trajectory can be computed in a single forward propagation.” Does that also imply having masked (causal) attention?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is quite technical, but well written; leaning towards formalism and clean definitions over extensive prose (which I consider a plus in this case, if the paper length were not limited by 9 pages I would recommend adding additional explanation / interpretation of the formal statements and results). The analysis is timely and original, looks very plausible to me (though I did not check the proofs in detail), and sheds some much needed light on the properties of Transformers that are often portrayed as borderline mystical. I think the work lays some much needed groundwork and provides a theoretical framework for analysis of further RL algorithms; I thus have no concerns regarding the paper's significance to the ICLR community.  \n\n**Pros:**\n * Timely problem, approached from a solid theoretical angle. The result is a sophisticated theoretical framework for analyzing RL algorithms and their learnability by Transformers via offline supervised pre-training.\n * Generalization error bound regarding the imitation performance of the Transformer w.r.t. the expert algorithm. Additionally, when assuming realizability, the supervised pre-trained Transformer (log-loss minimization on offline trajectories) is proven to be able to imitate the expert algorithm, given that the offline policy used to collect data “is close” to the expert algorithm (quantified in terms of a distribution ratio).\n * Analysis of 3 concrete RL algorithms, showing theoretically that Transformers can learn to imitate them in-context (with some extra conditions).\n * Simple simulations to confirm the theoretical results for the 3 algorithms.", "weaknesses": "The paper is interesting and results look good - my main criticism is that the paper stops somewhat short of situating the results within a wider agenda and clearly stating open issues and limitations of the current results. Overall I currently think that the paper is ready for publication and interesting to a significant part of the ICLR audience. I therefore suggest acceptance at the moment, and will update my verdict based on the other reviews and the authors’ response (it is possible that I missed or overlooked some critical details in the theory and required assumptions for it to hold; therefore I have lowered my confidence accordingly).\n\n**Cons:**\n * The paper currently lacks a strong discussion of limitations. The theory requires some assumptions, and while having the generalization bound is very good it remains somewhat unclear how hard/easy it is to make the bound tight in practice (a lot comes down of course to the mismatch between offline and expert algorithm; how easy is it in practice to keep this mismatch low? What can even be easily said about the distribution ratio in practice?). How do the theoretical statements in the paper lead to non-vacuous practical take-aways (after all requiring low distribution ratio between offline and expert algorithm, and sufficient model capacity are not the most surprising requirements)? Or do additional problems need to be solved to make such statements?\n * Another issue that is missing from the discussion (and perhaps even the analysis) is partial observability. If the expert algorithm depends on unobservables, I believe it would be possible to create settings that lead to a failure to imitate the expert; essentially leading to “Self-Delusions” - see [1]. Taking a wild guess here, but this would probably imply additional assumptions around identifiability? This point becomes even more important when considering humans as “expert policies”.\n * This is perhaps more a question than a weakness with actionable improvement (so consider this point addressed if it can easily be clarified/answered). I am wondering about the relationship between the transformer learning an **amortized algorithm** (i.e. the trained net behaves as if it were running the algorithm under the hood, but only on-distribution because it is not actually running the algorithm but has learned a function that produces the same outputs) **vs. actually implementing something close to the expert algorithm**. It seems to me that the paper suggests the latter (which is why it is important that Transformers can implement accelerated gradient descent and matrix square roots algorithmically). Naively I would have guessed that the transformer learns an amortized solution rather than the actual algorithm. On-distribution the two are indistinguishable, but typically the amortized solution generalizes less well. There are some experiments that could tease apart the two possibilities (but I’d also be interested in a general comment by the authors; most importantly whether anything in the paper would break with an amortized solution). Sketches for experiments (I do not necessarily expect to see them performed, but they might help clarify the discussion):\n    * Let the offline algorithm not fully cover the state space. At test time input these parts of the state-space (e.g. a coordinate input in a gridworld that is much larger than anything the Transformer has ever seen; but same dynamics, same reward function, etc). An amortized solution should struggle in this case (and potentially break down catastrophically), whereas learning the correct algorithm should not.\n    * Change the reward function at test-time (e.g. picking up bananas in a gridworld has positive reward and apples have negative reward during pre training, but at test-time the rewards are reversed). The amortized solution should struggle in this case.\n\n[1] Shaking the foundations: delusions in sequence models for interaction and control, Ortega et al. 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DAy7QoBJFr", "reviewer": "Reviewer_E6sG", "summary": "The paper examines the theoretical basis of in-context reinforcement learning (ICRL), specifically focusing on the supervised pretraining approach. This analysis builds upon previous studies [1, 2] and sheds light on the specific capabilities of algorithms trained through supervised pretraining. Theoretical findings demonstrate that transformers can approximate well-known algorithms such as LinUCB for bandit problems and UCB-VI for tabular MDPs under certain conditions. Empirical simulations further confirm the authors' findings that trained transformers effectively imitate bandit algorithms.\n\n[1] Laskin M, Wang L, Oh J, et al. In-context reinforcement learning with algorithm distillation[J]. arXiv preprint arXiv:2210.14215, 2022.\n\n[2] Lee J N, Xie A, Pacchiano A, et al. Supervised Pretraining Can Learn In-Context Reinforcement Learning[J]. arXiv preprint arXiv:2306.14892, 2023.", "questions": "1. Can you provide further explanation on the importance of calculating loss on the entire trajectory in Equation 3?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This work offers a timely investigation into the theoretical understanding of ICRL. Based on previous findings, the proposed framework advances this line of research by providing valuable analysis with guarantees on sample complexity.\n2. The analysis is overall well-executed.  It is good to have a clear and unified overview of the previous work.\n3. The paper is easy to follow.", "weaknesses": "1. Although this work theoretically covers tabular MDPs, it lacks the necessary experimental results to support the analysis.\n2. I am unsure if the proposed framework can offer insights on improving pretraining approaches for ICRL. I recommend that the authors address this in their manuscript.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)