Optimal Dynamic Regret by Transformers for Non-Stationary Reinforcement Learning
Authors: Baiyuan Chen
Shinji Ito
Masaaki Imaizumi
Venue: NeurIPS 2025 poster ()
Keywords: transformer
in-context learning
reinforcement learning
bandit problem
non-stationary
Primary area: reinforcement_learning
TL;DR: We prove that transformers can achieve nearly optimal dynamic regret bounds in non-stationary environments.

Abstract: Transformers have demonstrated exceptional performance across a wide range of domains. While their ability to perform reinforcement learning in-context has been established both theoretically and empirically, their behavior in non-stationary environments remains less understood. In this study, we address this gap by showing that transformers can achieve nearly optimal dynamic regret bounds in non-stationary settings. We prove that transformers are capable of approximating strategies used to handle non-stationary environment, and can learn the approximator in the in-context learning setup. Our experiments further show that transformers can match or even outperform existing expert algorithms in such environments.

Optimal Dynamic Regret by Transformers for Non-Stationary Reinforcement Learning
Baiyuan Chen
Shinji Ito
The University of Tokyo chenbaiyuan750g. ecc. u-tokyo. ac. jp
The University of Tokyo RIKEN AIP shinji@mist. w.wkkkyo..........o jp.
Masaaki Imaizumi
The University of Tokyo / RIKEN AIP imaizumi@g.cn ecc. w-toc.u-tkkk.aa ac. jp
Abstract
Transformers have demonstrated exceptional performance across a wide range of domains. While their ability to perform reinforcement learning in-context has been established both theoretically and empirically, their behavior in non- stationary environments remains less understood. In this study, address this gap by showing that transformers can achieve nearly optimal dynamic regret bounds non-stationary settings. We prove that transformers capable ap- proximating strategies used to handle non-stationary environments and can learn the approximator the in-context learning setup. Our experiments further show that transformers can match or even outperform existing expert algorithms in such environments.
Introduction
Transformers have emerged powerful class of sequence models with remarkable expressive capabilities. Originally popularized in the context of natural language processing, they leverage self-attention mechanisms to in-context learn
tasks without any parameter updates (Vaswani et al., 2017; Liu et al., 2021; Dosovitskiy et al., 2021; Yun et al., 2019; Dong et al., 2018). other words, large transformer model can be given prompt consisting of example input-output pairs for an unseen task and subsequently produce correct outputs for new queries of that task, purely by processing the sequence of examples and queries (Lee al., 2022: Laskin al., 2023; Yang al., 2023; Lin et al., 2024). This ability to dynamically adapt via context rather than gradient-based fine- tuning has spurred extensive interest in understanding the theoretical expressivity of transformers and how they might learn algorithms internally during training.
Recent theoretical work has begun to analyze the various aspects of transformers. On the expres- siveness front, transformers have been shown to be universal or near-universa. function approxima- tors under various conditions (Yun et al., 2020). Beyond mere approximation of static functions, researchers have demonstrated that transformers can encode and execute entire computational pro- cedures. For example, by carefully setting transformer's weights, one can make its forward pass simulate iterative algorithms such as gradient descent and many others (Bai et al., 2023; Cheng et al., 2024; Huang et al., 2025).
Studies have shown that transformers can approximate various reinforcement learning (RL) algo- rithms. The question here is whether transformer can be trained to act RL algorithm when presented with an interaction history as context. Lin et al. (2024) shows that a sufficiently large transformer model, after being pre-trained on offline interaction sequences, approximate near-
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Lin et al. (2024) (LinUCB)
Lin et al. (2024) (TS)
Study
Ours
Dynamic
O(T)
O 1/4/////4
O(min(vJT, A A1332/3 vT))
Regret
Table 1: Upper bounds on dynamic regret of transformers the non-stationary setup. T is the time budget, and A and J represent the amount and the number of changes in the environment, respectively. LinUCB suffers from O(T) dynamic regret non-stationary environments. The rate of Thompson Sampling (TS) is from Kim and Tewari (2020), whose algorithm is approximated by transformers through the analysis of Lin et al. (2024). The bound for ours shows the optimal rate derived in Besbes et al. (2014).
optimal online RI. algorithms like LinUCB and Thompson Sampling in a multi-armed bandit, as well as UCB value iteration for tabular Markov decision processes.
Despite these encouraging results, an important open challenge remains: Can transformers adapt to non-stationary environments? Most existing theoretical analyses assume fixed task or reward distribution during the transformer's in-context inference (Mao et al., 2021; Domingues et al., 2021). However, in many real- world scenarios, the environment is non-stationary- the reward- generrting process evolves over time, violating the assumptions of a static underlying model. In multi-armed bandits, for instance, the true reward probabilities of the arms might drift or change abruptly due to seasonality or shifts in user preferences. Classical algorithms for non-stationary bandits-such as sliding-window UCB, periodic resets, change-point detection- known achieve sublinear dynamic regret that closely tracks the moving optimum (Besbes et al., 2014; Garivier and Moulines, 2011; Wei and Luo, 2021; Mao al., 2021; Cheung et al., 2020). However, the behavior of trans- formers in non-stationary environments remains largely unexplored.
Our study. In this paper, we investigate whether transformer, trained in supervised in-context learning setting, can learn in non-stationary environments. Specifically, we introduce a transformer architecture designed to adapt to non-stationarity, and analyze both the approximation error and the generalization error arising from learning it from data. Based on these results, we derive an upper bound for the dynamic regret of the transformer and show its optimality.
Our main contributions are summarized as follows:
1. Regret optimality in non-stationary bandits: We prove that transformer acting in- context bandit learner can achieve cumulative regret matching the minimax optimal rate for non-stationary environments (e.g.. T2/3 for bounded changes) (Besbes et al., 2014). Rates are summarized in Table 1.
2. Generalization analysis under distribution shifts: Extending the offline-to-onlin general- ization theory of Lin al. (2024) (Lin et al., 2024), we derive conditions pretraining data diversity and model size needed to guarantee low regret in new non-stationary envi- ronments.
3. Architectural requirements for implementation: We identify key features- depth, atten- tion heads, and activations- enable transformer forget outdated information and implement restart-style strategies akin to handling non-stationary environments.
4. Novel proof technique via internal algorithm selection: We develop proof approach that treats the transformer's hidden state as maintaining multiple hypothesis policies and uses learned randomness in self-attention to sample among them, yielding tight regret bounds.
1.1 Related Works
In-context learning In-context learning (ICL) was first introduced by Brown et al. (2020), allow- ing large language models (LLMs) learn tasks by using few demonstration examples without fine-tuning or updating model parameters. Since its introduction, researchers have studied the prop- erties and mechanisms of ICL in depth (Mao et al., 2024; Zhao et al., 2024: Mosbach et al., 2023; Singh et al., 2024; Bertsch et al., 2025). For example, Mao et al. (2024) adopts a data generation perspective, showing that skill recognition and skill learning abilities in LLMS can transfer between tasks. Zhao et al. (2024) examines the decision boundaries of L.MMs in classification tasks and finds that, while in-context examples can improve accuracy, the models often display fragmented decision boundaries, meaning small input changes can lead to different outputs. Moreover, like fine-tuning,

---

ICL can be unstable and perform poorly on both in-domain and out-of-domane data, although fine- tuning tends generalize more effectively (Mosbach et al., 2023).
Approximation capability of transformers Transformers have been widely applied tasks such as reinforcement learning, computer vision, graph processing, and speech recognition (Chen et al.. 2021; Lin et al., 2024; Liu et al., 2021; Dosovitskiy al., 2021; Yun et al., 2019; Min et al., 2022: Dong et al.. 2018). For instance. Lin et al. (2024) explores using transformers for in-context rein- forcement learning. Chebotar et al. (2023) applies them to model Q-functions for multi-task policies, and Zhao et al. (2022) combines vision transformers with reinforcement learning for video caption- ing. While transformers have shown great potential, their limitations have also been discussed in the literature (Hahn, 2020; Bhattamishra et al., 2020).
Non-stationary Reinforcement Learning Non-stationary reinforcement learning (RI.) addresses scenarios where rewards and transitions evolve over time, with the degree of difficulty often charac- terized by the frequency and magnitude of these changes (Auer et al., 2019; Chen et al., 2019; Che- ung al., 2020). To tackle this challenge, Wei and Luo (2021) propose a method that periodically schedules instances with a restart mechanism, while Cheung et al. (2020) advocate leveraging recent data and employing wider confidence intervals to adapt to shifting conditions. Meta-reinforcement learning approaches have also been explored in non- non-stationary environments (Bing et al., 2022). Other research efforts address non-stationary settings as well, although many rely on prior knowl- edge, such as the number extent of environmental changes, to achieve effective adaptation (Mao al., 2021; Li and Li, 2019; Domingues et al., 2021).
1.2 Notation
We use the following notations throughout the paper. [s, e] for s, E N with s < e denotes the integer interval {s, s + .....,e), e). A(X') denotes the space of probability distributions over a finite Xx, denote the Cartesian product.
set X'. For sets N'...., . Xn, we write x X₂ := X1 C denotes element- wise multiplication, and + denotes concatenation. || llop denotes the operator norm of matrix. For an event E, we use |g[=] to denote an indicator function on E, which equals 1 if x E, and O otherwise. O-) is Landau's 's Big-O notation which hides some absolute constant, and O()) additionally hides logarithmic terms. [a: denotes that and are stacked vertically.
2 Problem Setup
2.1 Learning Setup
We first define environment of reinforcement learning following Lin et al. (2024). Let N be a number of rounds, and define a tuple of state/action/rewarrd spaces (S, Reherry a transition model T, S, A, 1)), with Sg. Ao (Ø), and the reward function R, S, x A. ARR.). We assume finite action space A [Ar C0.
Next, we define a learning scheme for the environment. Let D. = (3..a,,ry))] C 0 x Aj R, be sequence of observed state//ation/reeward tuples, where r. = Rs(st, az). We define an algorithm alg (0 x A₁ R2) x St A(A.). Then, we obtain distribution function over a trajectory Dr as
(DT)
I| Te- ( (t) t-1
(1)
For training, we assume a base algorithm alg, that provides high-quality actions and serves as a guiding policy for training the transformer-base algorithm.
Given N i.i.d. trajectories Dj- (si af rf...., st,, at,, rf) Page, with offline algorithm algn, we augment each Dj by of m.s.d. alg,; ( (e(r)) with a base algorithm alg and de- note as Dy. The base algorithm can observe the full trajectory and environment to generate actions for the supervised learning of algo-
3

---

Using the reward function defined above, we further define non- non-stationary environments following Wei and Luo (2021).
Definition 1 (Non-stationary measure). A function A [T) measure if satisfies
is defined a non-stationarity
(tt)
max SES,
[R((s, a) Re+1(s,
(2)
for all t. Given any interval I = ss, e). define AT == => A(T) and JI == 1+ [[A(T)= # 0j. We denote A pp,i and = J1.T1 as the amount and number of changes in the environment, respectively.
Our goal is to develop a bound on the dynamic regret of transformers under the non-stationary setup. In preparation, we define the maximal achievable reward r := max Rr a) at round its accumulated version R* = Σ r+ and the expected reward Ralg(T) DD for an algorithm alg. The regret with an algorithm alg is then defined as
Et=
(T) = (T) (T)
(3)
We will show that transformers can effectively solve the non-stationary data problem.
2.2 Transformer
We define the architecture of a transformer. Given an input sequence E Rddx, where is the feature dimension and the sequence length, consider masked attention layer with M heads. Denote the parameters as {(V,, m. Qm km)imeja = (RRxd))x3 The output of the masked attention
layer is denoted as H Attng (H) [hg...., hn] Raxn, where each h, for i e [n]] is given by:
h AAtnn [H)]] +
ReLU (4Qmh, Kmhy)) Vmhy Ed
where ReLU(z)) max(0, æ). With slight abuse of notation, write h, Attno (h, H). Next, the attention output is fed into a feedforward MLP layer MLP (:) with parameter m (W1, W2) Red' xd x Rd xd, where h h; + W2 ReLU(( (h))) Rdd, and D' is the hidden dimension of the MLP layer, transformer with layers is then defined as (H) = H where for e E [L]:
H MIP (2) (Attn,
(H( E Rdxn.
The transformer parameters collectively denoted O m0tt mii rameters are gl5) {(V!) )
0(2:2))). The attention pa-
3em[M]]] Rdxd and the MLP parameters are ((e) mip
(W()) w(( gd xd x idd The norm a transformer TFg is denoted as
1 ( lopp Kl [op) + EM V llop + W lop + /| (0 1op (4)
||||||| == max [t[
max MM
where lpp is the operator norm.
2.3 Algorithm induced by transformers in In-Context Learning
Let s, St and (a,, r.) A. R₂ be embedded by h: Use[z] St Utetre (Ar R2) Rd such that h(st) E IRd and h(a,, rt) E dR.. For the input H [b(s1), ), h(a,,,((aa)) r1), ......B. (((--1171-1) h(s))]
(2x(21-)), By intro- ducing a linear mapping Axd to extract a distribution over Ae with |Ail = A actions. the algorithm induced by the transformer is defined as:
(2--1, the transformer outputs H = TF. (H) = [h,, h2...., .e....,,,--l,.
(- = softmax (A TF. (H)
(5)
For convenience, we denote alge algo ( D1-1,

---

In supervised pretraining, the transformer parameters are learned by maximizing the log rog-likelhood over the algorithm class feee:
N
T
argmax BEG
Σ
log algo [ago(;;)
Di
(6)
N
where θ = (0 (B
is the finite parameter class of transformers.
Main Result
We show that under non-stationary environments, transformers trained with 1 in (6) can approxi- mate the reduced base algorithm algæ aa (- Dr- This is achieved by bounding the dynamic regret of transformers.
3.1 Preparation
We first introduce definitions required for the approximation theorem. In particular, we define the covering number to quantify the complexity of class of algorithms. This concept is standard in learning theory (Anthony and Bartlett, 2009) and has been used in analyses of transformers (Lin al., 2024).
Definition (Covering number). finite subset Ou θ called a p-cover of the algorithm class {algo 0 ) is. for every 0 θ there exists Og On such thut:
log algo, (1D,,1 s,) log algo (-I st)llo < P: &D, [T].
The covering number Ne(p) is the minimal cardinality of any such p-cover Ou.
Next, we define the distribution ratio as measure of discrepancy between two algorithms, which is used in Lin et al. (2024).
Definition (Distribution ratio). The distribution ratio of algorithms alg, and alg, defined as:
II
T
alg, (a.l D 8t)
Ralg, .lg, == =EDr~Pa
alg D 86)
Many reinforcement learning algorithms generate an auxiliary value at each round, for instance the upper confidence bound (UCB) in UCB-based algorithms. Let R,(s,, a,) denote such auxiliary value round f, under state 8, S, and action a, A,- For simplicity, write F, = R a,)- In transformer-based models, this auxiliary value can be interpreted as the output probability dis- tribution that guides action selection. To introduce the approximation theorem, we introduce the following assumptions used in the non-stationary analysis, e.g., Wei and Luo (2021).
Assumption 1 (Non-stationarity). For all O E O, alge outputs an auxiliary quantity 7, E [0, the beginning of each round There exists a -tatsaationarity measure and non-increasing function [T] + R such that rumning algo satisfies that without knowing A[I,) algo ensures with probability at least - 1/T:
1 Σ
F min re[1,[]
Ap. and
( -r-rr) p(t) Ap.
for all EE [T], provided that Ap.q & p(t) and p(t) > 1/Vi. Additionally, we assume the function C(t) = =t((t) is non-decreasing.
In addition, we assume that the base algorithm generating the observation sequence can be approx- imately realized by transformer, which is the following assumption. This assumption has been proved under several traditional RL. algorithms such as LinUCB and Thompson Sampling (Lin al., 2024). For additional discussion of this assumption, see Appendix A.
5

---

Assumption 2 (Realizability). Given a reinforcement learning algoritiu for any > 0, there
exist a parameter 0 such that the following holds: there exist constants Do. Mo. Lg. D/- Co > 0,
which depend such that transformer TFo with D Do: M Mo: I. Lo, D
P/ ||l||
Co, satisfies
log (e-11)))) - log alg (01K 1 kt) <
for all rounds and actions at
3.2
Regret Bound
As our main result, we develop an upper bound on the dynamic regret of the transformer-inducee policy algø- This theorem quantifies how well transformer-based policy performs in non- stationary environment compared to an optimal baseline algorithm. The proof is provided in Ap- pendix D.2.
Theorem 1. Let alga be reinforcement learning algorithm satisfying Assumptions and 2. Sup- pose Assumption holds with C(t) t? for some p [1/2, 1), and let 0 be solution of (6). Assume that Irel & holds almost surely, and define Ne := NN (NT) -2). Then, for any E > O, the dy- namic regret of algg induced by transformer with D & Du + 10, M max 2}, L Lu.
D' = (max Dy T log2 T/s}), [| = Õ + loga T/s) satisfies the following with prob- ability least max 1/T) with R. R Ralgo
Ralsa (T) = TQVR( =TTRR
log (og(NNT/() N
Ve
+ EAT min JT, A 1/3 + VT)
Taprx
Palgo
Toomp
(8)
This bound consists of three main components Teomp: Taprxs and Talgo Toomp captures the effect of pretraining the transformer N supervised samples learn the base distribution. T Taprx represents the approximation error incurred when using the transformer to imitate the base policy in the non- stationary setting. Talgo corresponds the intrinsic regret of the base algorithm in non-stationary environment.
Theorem shows that the transformer-based algorithm performs nearly optimally even in non- stationary environments, with the error diminishing as the training dataset grows and the envi- romment stabilizes. Specifically, when the number of observations N is sufficiently large and the approximation error & is sufficiently small, the regret bound in (8) is dominated by the third term Talgo- To further clarify this, present the following corollary, whose proof is contained in Ap- pendix D.3.
Corollary 2. Consider the setup Theorem 1. There exists a universal constant C such that, for N 3 8 log T, and & T-3, the dynamic regret of algå satisfies
( (imin(V.T. JJ,, A 1/33333 VT}).
Ralg, (T) =Ö min
This result implies that, with sufficiently large dataset and sufficiently expressive transformer (small e). the transformer-based policy achieves the optimal dynamic regret in non-stationary envi- ronments, matching the bounds established by prior works (Wei and Luo, 2021). Intuitively, smaller corresponds to larger transformer architectures, as indicated in Thcorem 1. Hence, with enough data and a suitably large architecture, the transformer is capable of realizing the optimal strategy that minimizes dynamic regret.
While the above bounds could also be achieved with standard transformer, for the sake of simpler proofs we use non-continuous transformer. Structurally, this model is equivalent to regular trans- former except that its inputs are duplicated few times to form an augmented input sequence. The motivation for this construction, along with discussion on how the proofs extend to standard trans- formers, is provided in Appendix E.2. An abstract illustration of the non-continuous transformer shown in Figure 1.

---

Proof Outline
In this section, we provide an overview of the proof of Theorem 1. The most critical component is the analysis of the approximation error Tppxx. To this end, we first introduce common operations and algorithmic structures used handle non-stationary data, and then show how transformer can approximate these operations effectively.
4.1 Operation for Non-Stationary Environments
In non-stationary environments, maintaining strong performance requires limiting the reliance historical data. Existing approaches primarily employ two strategies: using a window scheduler, where only data within sliding window available to the model (Cheung et al., 2022; Trovo al., 2020), and rest-and-restart mechanism, which resets the algorithm when significant changes in the environment detected (Gomes al., 1998; Cheung al., 2020; Cayci al., 2020; Wei and Luo, 2021: Mao et al., 2021). Among these approaches, the MASTER algorithm (Wei and Luo, 2021) stands out. MASTER not only achieves near-optimal performance in non-stationary environments but also does not require prior knowledge of the environment's ' change parameters, such as the number or magnitude of changes.
We briefly deseribe the window scheduler and the restarting operation MASTER, while their formal definitions will be provided in Appendix E.1. The overview of the proof is illustrated in Figure 1.
Stationary tests
Algorithm
Dr
Window Scheduler
alge(a|Dy.>>
Rostart Mcch.
1
Prop
Prop
Thm
Scheduler block
Transformer
Dr
MLP
algo Dy.s)
+
Randomizer block
Figure I: Overview of the proof structure for the approximation analysis (Theorem 5). The blue box represents the algorithmic operations for handling non-stationarity, while the green box represents the transformer approximating the algorithm. Each block or aub-architecture of the transformer corresponds to specific operation of the algorithm.
Window scheduler: Denote n + copies of trajectories Dr as {D) Yi-0- Given window sizes [W o and probabilities (r( J. we define following two modules: (i) a masking operator, and (ii) selection operator.
First, the masking operator 0n maps the input trajectory Dr randomly masked version of the copies (D y using the window sizes and probabilities. Formally,
σ1 (Dr)
m (6) 3 D
where m(( E {0, 137 is a binary mask for copy i defined elementwise as
Bernoulli(p
if (mod 2'), k [t,t+ W 1-0-- 1]
( m
otherwise
This operator probabilistically selects candidate time points of interest within each copy based on the window size Wii) and time position 4.
Second, the selection operator 0g chooses single active trajectory instance among the masked copies:
(6) (() ))=0)) 02 (1(04) o1
[0 =[0;... 0;...; :0; (, a .7 (k) 0:
(k)
(k)

---

.1() * 0 is the index of the first non-zero entry where k min no. (1) ot (order-k). The selected instance is called "active", while the others are "inactive". This ensures that only one instance is active at each round, and it is always the lowest-order scheduled instance.
()
Finally, the window scheduler (WS) is defined as: WS(Dr)
01 (Dr).
The window scheduler stochastically schedules multi-scale windows to restrict the algorithm to recent data (Figure 2).
D
X(3 + 1).
order
Figure 2:
Illustration of WS when 3. Purplish blocks represent instances scheduled by while reddish blocks represent the active instances selected by d2. Reddish blocks connected by a dashed line are concatenated.
Stationary Test and Restart Mechanism: After passing through the window scheduler WS, the base algorithm interacts with the environment, receives a reward, and performs a stationary test, which outputs change-detection signal C {0, 1}. If significant change in the environment is detected (C, 1), the algorithm is restarted. We denote the output of wS as the sequence of active instances ((// 0/2 r())==-- We then define the auxiliary memory A. {(8pp dp rp {p]]p=1r where the auxiliary value Ft is added at each step.
Given the change-detection signal C, {0, 1} and auxiliary memory At, the restart mechanism (RM) is defined as:
[Ao (uc,A D.)
i 1
RM (A(+1)
otherwise
where Ao is an empty/reset auxiliary state (all counts=O, rewards=2), and Ul updates the auxiliary memory (e.g.. appending new rewards). Typically, Ft evaluated to determine when restart is necessary.
It is worth noting that the sliding window and restart mechanisms used in previous works (e.g., (Cheung et al., 2022; Trovo et al., 2020; Cayci et al., 2020)) can be considered special cases of WS and RM. We discuss their equivalence in Appendix E.4. As result, this approximation can be extended to a wide range of algorithms designed to handle non-stationarity.
4.2 Transformers for Approximating the Operations
Our goal is to demonstrate that transformers can approximate the above techniques and the overall algorithm. Given a base algorithm alg,, we have the following statement.
Proposition 3 (Approximating WS). For any small & > O, there exists a transformer TFo(-) with D = O(1), L = O(1), M = O(1/VE), D' = Ol Toog, T/e) such that the algorithm alge- defined in (5), can generate a random number drawn from uniform distribution and satisfies
|log alge ( D,-1,88) 81) log alg ( 8c] < E, VL [T], k [A].
(10)
This proposition shows the existence of transformer block (called the scheduler block) that ap- proximates the window scheduler WS. In the construction, we also include a block that extracts randomness from the data to implement the stochastic scheduling in WS. The proof of Proposition is provided in Appendix E.3.
Furthermore, since RM involves only zeroing out certain values and applying mask [ to detect changes, approximation by transformer follows directly. Details for approximating are included in Appendix E. part of the overall proof.

---

Proposition 4 (Approximating RM), The restart mechanism RM can be implemented by two-layer MLP with D' = O11)
Finally, we consider an expert algorithm algæ. which incorporates both the window scheduler and the restart mechanism to handle the non-stationary environment. The full details of this algorithm are provided in Section C. Then, the following theorem holds:
Theorem 5. For any
there exists transformer TFg with D Do 10, =
(ax 2), Lo. D' O(max( D/. log2 T/e)). ||@||| Ó(Co + /Tlog, T/E), such that the transformer-basee algorithm algo defined in (5) satisfies
( D,- 81) log ( Dt) <,
Wt [T, ke [4].
(11)
With this result, we obtain the approximation error T aprx- The entire proof is shown in Figure 1. The proofs of Theorem 5 is contained in Appendix D.4.
Together with the analyses of the training error Toomp and the algorithmic regret Talgo Theorem shows that transformers can effectively replicate the techniques required for handling non-stationary environments, completing the argument for Theorem 1. Full technical details are provided in the supplementary material.
5 Experiment
In this section, we evaluate transformers and other algorithms in linear bandit setting. The stochas- tic linear bandit framework is given by M (w* E, : At each round t [T]. the learner selects action E Rd from the set At =
dt,A), which may vary over time. (a,, (a,,"*) wn> + E., where et Wiid. d. & and w" e Rd is unknown. The problem generalizes by setting 8; = Ar, with the state transitioning deterministically to 8141 regardless of the action.
The learner then receives reward
We compare transformers against Linear UCB (LinUCB) and Thompson Sampling (TS), as well as MASTER (Wei and Luo, 2021) combined with LinUCB/TS (denoted as expert algorithms) under environments with varying degrees of -tatiotttio,,itt. In our experiments, we set d 32, A 10, ~ NO, 1.5²), and w* ~ Unif((0, 1J We consider two types of environments:
(1) Low -tatioaatonaarit: Models evaluated over 1,000 rounds, with elevated rewards in E [50, 100] u [350, 400] scaled to r E [3, [3. and the remaining rewards in r, E [0, Training data consists of 100,000 samples with normalized rewards rt 1].
(2) High Non-Stationarity: The reward is defined as = (Ka, w*) > + E2) (2bb). For training, we generate 100.000 samples for each b E {0.005, 0.01, 0.015, 0.02). For eval- uation, we test on unseen environments with b E {0.018, 0.025], running 200 rounds per environment.
For transformer models, we use GPT-2 with L = 16 layers and M = 16 attention heads, trained for 200 epochs. The objective is assess generalization non-stationary environments unseen during training.
From Figure 3, we observe that transformers achieve performance comparable to, and sometimes surpassing, both the expert algorithms and their MASTER variants, attaining near-optimal cumula- tive regret. Additional experimental results can be found in Appendix B.
6 Conclusion
In this study, we have shown that transformers can effectively handle non-stationary environments, achieving near-optimal performance by minimizing dynamic regret. By demonstrating that trans- formers can implement strategies commonly used for adapting non-stationarity, we provide a the- oretical guarantee for their dynamic regret bound. Our experiments further show that transformers can match, and in some cases outperform, existing expert algorithms. As a limitation, we evaluated transformers empirically only against two RL algorithms and their variants in the linear bandit set- ting. Future work should explore additional experimental setups and broader algorithm comparisons to further validate findings.

---

Regret Over Time
Regret Over Time
Regret Over
Regret Time
100
Cymáletive Regret
Regret
Thonlp
Cumputsive Ragrat
140
75
90
Thomp
100
Cumulative
80
Thomp
100
TF
25
70
60
0-
0
500
1000
800
900
1000
500
1000
900
1000
Rounds
Rounds
Rounds
Rounds
Regret Over Time
Regret Over Time
Regret Over Time
Regret Over Time
Cumulative Regret
100
MASTER
Cumullative Regret
120
Currulative Regret
100
Thomp
Inomm
Curryjattvc Aurrulstic Regret
120
Thomp
TF
π
11.0
110
1bp
100
TF
100 Episodes
190
200
160
200
Figure 3: Cumulative regret comparison for LinUCB, Thompson Sampling (TS), MAS- TERR+LinCBB/TS and transformer (TF) in linear bandits d 32, A = 10. The first row corre- sponds Low von-Sttatonnaiity environments, while the second row shows High Non-Stationarity environments. Shading indicates the standard deviation of the regret estimates.
Acknowledgements
Shinji Ito was supported by JSPS KAKENHI (JP25K03184). Masaaki Imaizumi supported by ISPS KAKENHI (JP24K02904), JST CREST (JPMJCR21D2), and JST FOREST (JPMJFR2161).
References
Anthony, M. and Bartlett, P. L.. (2009). Neural Network Learning: Theoretical Foundations. cam-
bridge university press.
Auer, P., Gajane, P., and Ortner, R. (2019). Adaptively tracking the best bandit arm with an unknown number of distribution changes. In Conference Learning Theory, pages 138-158. PMLR.
Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. (2023). Transformers as statisticians: Prov- able in-context learning with in-context algorithm selection. Advances in Neural Information Processing systems, 36.
Bertsch, A., Ivgi, M., Xiao, F.,, Alon, U., Berant, J., Gormley, M. R., and Neubig, G. (2025). In-context learning with long-context models: An in-depth exploration. Association for Com- putational Linguistics, pages 121 12149.
Besbes, O.. Gur. Y., and Zeevi, A. (2014). Stochastic serti-ammed--aanae problem with non- stationary rewards. Advances in Neural Information Processing systems, 27.
Bhattamishra, S., Ahuja, K., and Goyal, N. (2020). On the ability and limitations of transform- ers to recognize formal languages. In Conference on Empirical Methods in Natural Language Processing, pages 7096-7116.
Bing, Z., Lerch, D., Huang, K., and Knoll, A. (2022). Meta-reiinooremment learning in non-stationary and dynamic environments. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(3):347633491.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. Advances in Neural Information Processing systems, 33.
Cayci, S., Eryilmaz, A., and Srikant, (2020). Continuous-time multi-armed bandits with con- trolled restarts. arXiv preprint arXiv:2007.00081
10

---

Chebotar, Y., Vuong, Q. Hausman, K., Xia, F., Lu, Y.. Irpan, A., Kumar, A., Yu, T., Herzog, A., Pertsch, K., et al. (2023). Q-transformer: Scalable offline reinforcement learning via autoregres- sive g-functions. In Conference on Robot Learning, pages 99299. PMLR.
Chen, L., Lu, K., Rajeswaran, A., Lee. K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mor- datch, I. (2021). Decision transformer. Reinforcement learning via sequence modeling. Advances in Neural Information Processing systems, 34.
Chen, Y., Lee, C.-W,, Luo, H., and Wei, C.-Y. (2019). A new algorithm for non-stationary contextual bandits: Efficient, optimal and parameter-free. In Conference on Learning Theory, pages 696- 726. PMLR.
Cheng, X., Chen, Y., and Sra, S. (2024). Transformers implement functional gradient descent to learn non-linear functions context. In International Conference on Machine Learning, pages 8002-8037. PMLR.
Cheung, W. C., Simchi-Levi, D., and Zhu, R. (2020). Reinforcement learning for non-stationary markov decision processes: The blessing of (more) optimism. In International Conference on Machine Learning, pages 1843-1854. PMLR.
Cheung, W. C., Simchi-Levi, D., and Zhu, R. (2022). Hedging the drift: Learning to optimize under nonstationarity. Management Science, 6863):1996-1713
Domingues, O. D., Ménard, P., Pirotta, M., Kaufmann, E., and Valko, M. (2021). A kernel-based approach to non-stationary reinforcement learning in metric spaces. In International Conference Artificial Intelligence and Statistics, pages 3538-3546. PMLR.
Dong, L., Xu, S., and Xu, B. (2018). Speech- :reccttrransfommrr:: no-recurence sequence-to-sequenee model for speech recognition. In International Conference on Acoustics, Speech and Signal Pro- cessing, pages 5884 5884-5888. IEEE.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. (2021). An image is worth 16x16 words: Transfomees for image recognition at scale. In International Conference Learning Representations.
Furuya, T., de Hoop, M. V., and Peyré, G. (2025). Transformers are universal in-context learners. In International Conference Learning Representations.
Garivier, A. and Moulines, E. (2011). Upper-confidence bound policies for non-stationary bandit problems. In Algorithmic Learning Theory, pages 174-188.
Gomes, C. P., Selman, B., and Kautz, H. (1998). Boosting combinatorial search through random- ization. In Conference on Artificial Intelligenceehnnovvtiie Applicaao of Artificial Intelligence, pages 431-437.
Hahn, M. (2020). Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171.
Hataya, R. and Imaizumi, M. (2024). Transformers as stochastic optimizers. ICML 2024 Work- shop on In-Context Learning.
Huang, J., Wang, Z., and Lee, J. D. (2025). Transformers learn to implement multi-step gradient descent with chain of thought. In International Conference on Learning Representations.
Kim, B. and Tewari, A. (2020). Randomized exploration for non-stationary stochastic linear bandits. In Conference Uncertainty in Artificial Intelligence, pages 71-80. PMLR.
Laskin, M., Wang, L., Oh, J., Parisotto, E., Spencer, S.. Steigerwald, R., Strouse, D., Hansen, S. S., Filos, A., Brooks, E., et al. (2023). In-context reinforcement learning with algorithm distillation. In International Conference on Learning Representations.
Lee, K.-H., Nachum, O., Yang, M. S., Lee, L., Freeman, D., Guadarrama, S., Fischer, I., Xu, W. Jang, E.,, Michalewski, H., et (2022). Multi-game decision transformers. Advances Neural Information Processing Systems, 35.

---

Li, Y. and Li, N. (2019). Online learning for markov decision processes in nonstationary environ- ments: A dynamic regret analysis. In American Control Conference, pages 2332-133, IEEE.
Lin, L., Bai, Y., and Mei, S. (2024). Transformers as decision makers: Provable in-context rein- forcement learning via supervised pretraining. In International Conference Learning Repre- sentations.
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. (2021). Swin transformer: Hierarchical vision transformer using shifted windows. In International Conference Computer Vision, pages 10012-10022.
Mao, H., Liu, G., Ma, Y., Wang. R., and Tang, J. (2024). A data generation perspective to the mechanism of in-context learning. CoRR.
Mao, W., Zhang, K., Zhu, R., Simchi-Levi, D., and Basar, (2021). Near-optimal model-free reinforcement learning in non-stationary episodic mdps. In International Conference on Machine Learning, pages 7447-7458 PMLR.
Min, E., Chen, R., Bian, Y., Xu, T., Zhao, K., Huang, W., Zhao, P., Huang, J., Ananiadou, S., and Rong, Y. (2022). Transformer for graphs: An overview from architecture perspective. arXiv preprint arXiv:2202.08455.
Mosbach, M., Pimentel, T., Ravfogel, S., Klakow, D., and Elazar, Y. (2023). Few-shot fine-tuning in-context learning: fair comparison and evaluation. CoRR.
Singh, A., Chan, S., Moskovitz, T., Grant, E., Saxe, A., and Hill, F. (2024). The transient nature of emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 36.
Streeter, M. and Golovin, D. (2008). An online algorithm for maximizing submodular functions. Advances in Neural Information Processing Systems, 21.
Trovo, F., Paladino, S., Restelli, M., and Gatti, N. (2020). Sliding-window thompson sampling for non-stationary settings. Journal of Artificial Intelligence Research, 68:311-364.
Vaart, A. W. v. d. and Wellner, J. A. (2023). Empirical processes. In Weak Convergence and Empiri- cal Processes: With Applications Statistics, pages 127-384. Springer International Publishing. Cham.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polo- sukhin, I. (2017). Attention is all you need, Advances in Neural Information Processing systems, 30.
Wei, C.-Y. and Luo, H. (2021). Non-stationary reinforcement learning without prior knowledge: An optimal black-box approach. In Conference on Learning Theory, pages 4300-4354. PMLR.
Yang, S., Nachum, O., Du, Y., Wei, J., Abbeel, P., and Schuurmans, D. (2023). Foundation models for decision making: Problems, methods, and opportunities. arXiv preprint 233hv:2303.04129,
Yun, C., Rao, M., Becker, B., al. (2020). Are transformers universal approximators of sequence- to-sequence functions? In International Conference on Machine Learning, pages 9116-9126.
Yun, S., Jeong, M., Kim, R., Kang J., and Kim, H. J. (2019). Graph transformer networks. Advances in Neural Information Processing systems, 32.
Zhao, H., Chen, Z., Guo, L., and Han, (2022). Video captioning based on vision transformer and reinforcement learning. PeerJ Computer Science, 8:e916.
Zhao, S., Nguyen, T., and Grover, A. (2024). Probing the decision boundaries of in-context learning in large language models. Advances in Neural Information Processing Systems, 37.
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We clearly stated the main claim in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We stated the limitation in the last section for discussion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the ap- proach. For example, a facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Ys]]
13

---

Justification:
We have described the full assumption, statements, and proof in the main body and the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We have described the detailed experimental details in the corresponding section in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models. it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
14

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We clearly stated the source of data. Also, we will open the source code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips..
for more details.
While we encourage the release code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //niss... "/nppsccppppii coc/public/guti wwwpc//ggiiiieegeeeeeeeeeeeeeeee Sed
onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We have clearly described the experimental details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question:
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We have repeated the experimental multiple times and report the standard deviation comes from the replication.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
15

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors). should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK report I-sigma error bars, but one should state The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification: Our experiments are small-scale and implementable by a small laptop. Also, we do not pursue the computational cost in this study, the computational resource is out our focus.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't t make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //teer/is....///eeeee
Answer:
[Yes]
Justification:
We have checked the code.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No. they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: A main focus of this study is fundamental, there is almost no effect on social impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
16

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Justification:
Since this paper is theoretical, the outcome does not have righ risk for
misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We have cited the source of the data which we use.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
17

---

If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswithcode com/aatasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
Answer: [No]
Justification:
We have not created a new asset throughout this study.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Answer: [No]
Justification:
We have not performed the crowdsourcing experiments and others.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper deseribe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [No]
Justification: Since this study is fundamental, there no potential risk on this point.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
18

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [No]
Justification:
We have used LLM only for the formatting purposes.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy peetps://neuriie for what should or should be described.
22055/LLM)
19

---

A Limitation and Discussion
Regret bound In Corollary 2. the regret bound of the learned algorithm alg; is guaranteed when the approximation error is sufficiently small- requires a sufficiently large model size-and when the dataset size N is sufficiently large. Although these requirements may seem inefficient compared to the model sizes and training datasets of traditional RL algorithms, modern LLMs are typically trained with scales that readily meet both conditions.
Proof method To establish the regret bound, we employ a non-continuous transformer, which is structurally equivalent to a standard transformer. This choice is made primarily for proof simplicity. While we also provide a brief description of how the standard transformer can be used for the proof (Appendix E.2), developing more elegant proof directly in that setting is left for future work.
Pretraining of transformer The regret bound holds without assumptions on the non-stationarity of the training data. A natural direction for future work is to examine how training data with different types and degrees of non-stationarity influence the generality and performance of transformers.
Assumption and 2 Our theoretical guarantees rest on Assumptions and 2. The former has been shown to hold for wide range of RL algorithms (Wei and Luo, 2021), whereas the latter may appear somewhat strong. We emphasize, however, that Assumption 2 builds on prior work demonstrating that transformers can emulate standard reinforcement learning algorithms, such as LinUCB and Thompson Sampling, through in-context learning (Lin et al., 2024). In addition, Furuya al. (2025) shows that transformers universal approximators over distributions, which further supports the plausibility of this assumption. We also emphasize that the assumption does not require the transformer exactly replicate the expert policy. Instead, assumes that the transformer can learn an amortized algorithm that approximates the expert's 's output distribution, which is a more relaxed and often more realistic requirement practice.
B Further Experiment Results
In our experiments, we set the confidence scaling parameter a LinUCB to 1, and the noise vari- ance for Thompson Sampling (TS) to 0.3.
The suboptimality of the models in Low Non-stationary environments and in High Non-stationary environments with b = 0.018 (Figure 3) is shown in Figures 4 and 5. We observe that the transformer only outperforms the expert algorithms but also maintains a consistently low suboptimality rate across rounds.
Online Evaluation
Online Evaluation
0.20
MASTER LinUCB
0.15
LinUCB
Suboptimality
Suboptimality
0.15
TF
0.10
wwww
0.10
0.05
MASTER+ Thomp
0.05
Thomp
TF
0.00
0.00
500
1000
500 Rounds
1000
Rounds
Figure 4: Suboptimality Comparisons of LinUCB, Thompson Sampling (TS), MAS- TER+LinUCB/TS. and transformer (TF) in Low Non-stationary environments. Linear bandit with = 32, = 10. Shading indicates the standard deviation of the regret estimates.
Figures 6 and 7 present the results for High Non-stationary environments with b 0.025. We observe that, when the level of non-stationarity becomes too high, the transformer no longer outper- forms the other models. This indicates that the model generalizes well to b 0.018, but struggles
20

---

Online Evaluation
Online Evaluation
MASTER+
0.65
0.7
LinUCB
Suboptimality
TF
Suboptima lita
0.60
0.6
0.55
MASTER+
0.5
Thamp
0.50
TF
100 Rounds
200
100 Rounds
200
Figure 5: Suboptimality Comparisons of LinUCB, Thompson Sampling (TS). MAS-
TER+LinUCB/TS, and transformer (TF) in High Non-stationary environments (b 0.018). Linear bandit with d 32, A 10. Shading indicates the standard deviation of the regret estimates.
with b 0.025, likely due limited coverage of highly non-stationary scenarios in the training data.
Online Evaluation
Regret Over Time
Regret Over Time
0.8
Cumulative Regret
MINTER+LinUCH
Cumulaative Regret
LinUCB
100
Suboptimality
TF
120
0.6 MASTER+LinUCB LinUCB
50
MASTER+LinUCE
110
LinUCB
0.4
100 Episodes
200
Rounds
Rounds
Figure 6: Suboptimality and Cumulative Regret Comparisons LinUCB, MASTER+LinUCB, and
transformer (TF) in High Non-stationary environments (b 0.025). Linear bandit with d 32, = 10. Shading indicates the standard deviation of the regret estimates.
Online Evaluation
Regret Over Time
Regret Over Time
0.70
MASTER+ Thomp
Cumulative Regret
MASTER+ MASTER+Thoma
Cumulafive Regret
MASTER+TThopp
Suboptimality
Thomp
100
Thomp
125
Thomp
0.65
TF
0.60
50
115
0.55
19.0
0.50
100
100 Episodes
200
180
190
Rounds
Rounds
Figure 7: Suboptimality and Cumulative Regret Comparisons of Thompson sampling (TS), MAS- TER+TS, and transformer (TF) in High Non-stationary environments (b 0.025). Linear bandit with d 32, A 10. Shading indicates the standard deviation of the regret estimates.
C MASTER Algorithm
We begin by introducing MALG (Multi-scale AL.Gorithm) (Wei and Luo, 2021), which extends base reinforcement learning algorithm, ALG, by running multiple instances at different time scales.
21

---

Given an integer and a non non-increasiig function p [T] R. MALG operates over time horizon of length 2".
At each time step 7. given an integer m, if r is a multiple of 2m, MALG schedules a new instance of length 2 with probability p((")) These are referred order-m instances. The start and end times of each instance are denoted by alg.s and alg.e, respectively. Among all scheduled instances, only the one with the lowest order remains active any given time, producing an auxiliary value Fz.
After each step, MALG updates the active instance based observed rewards R2 from the environ- ment. In this framework, all reinforcement learning algorithms generate a specific auxiliary value; for instance, in the LinUCB setting, it corresponds to score that combines a reward estimate with an uncertainty term.
Algorithm 1: MALG (Multi-scale (Wei and Luo, 2021)
Input: n. p())
for 0..... 2 do
2
for m = n. -..-..... 0
if is multiple 2 then
With probability p(2") p(()))) schedule a new instance alg of ALG at scales 2";
Run the active instance alg to output 7,, select an action, and update with feedback.
Algorithm 2: MALG with Stationarity TEsts and Restarts (MASTER)(Wei and Luo, 2021)
Input: p()) where p(t) = 6(1og2 T + 1) log(T)) (T: block length)
Initialize & -
for n 0, 0... do
Set f, and initialize an MALG (Algorithm for the block [tn. fn + 1]: while (. < f,, + 2 do
Run MALG to obtain prediction it select action de, and receive reward Rd;
Update MALG with feedback, and set U, = mi) Fr: Perform Test and Test (see below);
Increment < 1+ 1; if either test returns fail then restart from Line 2;
10
Test 1: Ift alg.e e for some order-m alg and I aaige 2 -o- R, U,+ + 9p(2"),, return fail.
Test 2: If 1-bn+1 Lt, (f- r-) 3p(t in 1), return fail.
MASTER (MALG with Stationarity Tests and Restarts) further enhances this process by incorpo- rating repeated checks for stability and performance. uses two key tests: the first ensures that the average reward within any completed instance does not significantly exceed threshold determined by the auxiliary value Fe; the second verifies that the difference between the auxiliary output and observed rewards remains bounded on average. If either test fails. MASTER resets the algorithm, thereby maintaining robustness under changing environmental conditions.
D Proofs in Section 3
D.I Intermediate Statement
We provide the following intermediate statement show Theorem
Theorem 6. Let Assumption I and 2 hold. Then for any small e > O, there exists an algorithm alg introduced transformer with
D Do 10, M max{Mo. 2}, L = Lo.
(12)
D' D'= 0'= 2.0(max(D)) log, T/s)),
|
Õ(C + Vllog, T/E)
22

---

satisfying
log(NeT/0) n
Raaua (T)
Rate (T)
=Ö T
+EAT)
(13)
D.2 Proof of Theorem
As established in Wei and Luo (2021), for any reinforcement learning algorithm algg, the combina- tion of MASTER and alg results in stabilized version: algg- The dynamic regret of algğ can then be expressed as
Halsa
Halsa
(T)
R Magg
(T)
R
aTg g
(14)
=:
can be obtained from Theorem 6, and B bounded as follows:
B = O min
a
2 VJT, (o + c2c 4/a) A 1/372/3 + (a 2
vT)).
where C(t) = 1/2 + 02 with c1, 02 > 0 (Wei and Luo, 2021). Setting c = and C2 & 1, we obtain:
Ō min {JJT, A AV07233 vT))
D.3 Proof of Corollary 2
suffices to prove that
log(NeT/0)
2 V R
+ Ve + +EAT Ö(VT).
algg algn
According to Theorem 5, we have
algø( (a, 12-1, 81)
$ E
algo (44 D,-1,8t) Si)
Thus, the cumulative distribution ratio over T rounds satisfies
R Ra Teg
eT
Since we assume E < T 3, it follows that
R Ralg.iTT Tk. = 0(1).
Consequently, have
T
R ales dlg
=
= O(VT).
Next, following Vaart and Wellner (2023), there exist universal constants C1 and C2 such that
No = -N((NNT 3) & CI
where V(e, θ) denotes the fat-shattering dimension of O. Let V(@) be the VC-dimension of O then, for any e > O, it holds that V(e, θ) V(@) (Vaart and Wellner, 2023). Since O a finite parameter class, there exists constant C3 such that V(E,, θ) & V(@) & C3 for all > 0. Therefore, under the assumption that N > CT log T. we obtain:
(oo(N)
log( (N6T)))
2 R Raggg
< 2
N
N
(1 + 20, (C) log T + 20 log N + log(2
T
N
Ölog T) OT' log T)
< 2
= Ö(VT)
which completes the proof.
23

---

D.4 Proof of Theorem 5
As shown in Section E, a noncontinuous transformer (Definition with
D ((d + (m + 1) + 5), M 2, L 1>0, 0, D' loge T/e, |||l| Ö (VTlog, T/e)
can approximate the MASTER algorithm. Treating any reinforeement learning algorithm that uses MASTER as the expert algorithm algg in Theorem 5, the result follows directly from Assumption
D.5 Proof of Theorem
By Theorem have
algo 01, D <8
log
algg ( (1)
algg(ar,k Dt-1. St)
< .
algg (a,k 81)
(0r, D st)
alg (t ( D st)l (e-
where the last inequality uses the fact that alg(-) < With slight abuse of notation, let rr, denote
the reward of the k-th action at round Since [rx & 1 almost surely for all i e [T], we obtain the policy imitation error between transformers and MASTER as follows:
Policy Imitation Error
(algg (0t, Dr-1, - algg(ar, 8t)) rikk
algo (ur. (D-1, algg(ar,c)Dr-11 Irea
Σ (e 1)
t= k-
(e A)AT.
Finally, since e > 0 small, we use the approximation
= ===+5: O(=2),
and then the policy imitation error becomes e.T..
Next, we leverage the result from Lin et al. (2024):
(1
log(NoT/0) N
Ralg, (T) R (T) = 0
T
R Rala alaa
+ V
(15)
where
alg algns(a Dr-1,
log E Dy Patea .eze
< Eycal
algø (7g Dr-1,
for all / [T]. We extend this result to non-stationary settings to derive bound for [Rale, (T) Rag (T)).
Suppose there nn orders of instances, and alg runs for T rounds. Let k; 2 0 represent the number rounds for order-i instances, so that kip + hn + ..... + kin = T. Define T; as:
T, =
). T k; ie (O.
which is the set of rounds during which order-i instances are active. If h; 0, order-i instances are inactive. If 0, one or more order-i instances are active. Using Lemma 7, we know that running instance during T, yields higher regret bound (as in (15)) than running zero multiple

---

instances in T. Furthermore, Theorem guarantees that the regret for reinforcement learning algorithm at T, bounded by EAk, Thus, we have:
(V (loo((Vokk/00) N
R (T:)
9 alg
(T:)
O
ki Ralss Ralg
+ +.Veevre /
+ EAkss
Summing over all Tú, we get:
Ralap (T) Ragg (T)
alg (Ti)
Halgs (Ti)
Rales (Ti)
RITE (Ii)
log log(NNkk.00 N
o
Σ
12
alg
+ VEreal
+ EAki
2
Rolsa u8g
+ Vercal
+ EAki
N
2
(og(NoT
Σ hi
V Rolso
+ V
+ Σ EAk;
N
(og(NoT/0)
0
1
R Ralgg STe
+ V
+ EAT
N
where the final inequality follows from the Cauchy-Schwarz Schwarz inequality. By Assumption have
alga Te Di-1, st) e
alga āhe Dz-1, St)
Therefore,
alg (hp D1-1, & Ereal &
log E D Jao dls
algg (ate D1-1, 8t)
Combining this with the result of Theorem 6 completes the proof.
D.5.1 An auxiliary lemma
Lemma 7 (Regret bound for multiple instances). Let a, N,,,N, 0 with + b c. The following bound holds:
a² ( Tog(Cra)
(VVos(CLo) +C2 + 6²
log( C1b) + C32
0² e [og(Cic)
+ CO2 (V
(16)
Proof of Lemma 7. Let < u < be such that a = uc and obtain:
(1 u)c. Substituting into (16),
(V Vlo(Crruc) C22 + (1-w)) u)² (V w)) C2 E ( log(Crc)
( C2
N
Denoting G log(Cic)/N, we rewrite the inequality as:
2² (Vlogu G + Cz) +1 -1-u)) u)²
- u) - +)+ & V Cz.
(17)
Since 0 < < 1, we have log < and log( 1 - u) < 0. Substituting these into (17), we obtain:
2 + (1 - -1= & 1,
which follows directly from 0 < u < 1, as 2 + (1 - w)² 1- 1 2u u) < 1.
25

---

E Proofs in Section
E.1 Structural Approximation
To match the order-n instance length as described in Appendix C, we consider single input matrix
H [11. .... h2 pdx2" where [h, ) R. This matrix is then extended H. Rxx2"
defined H H H H* TR. Dx2" (Figure 8) where D := ((d + 5) 1) 5),
and H contains auxiliary entries for approximation purposes. Here, H :- [h ...., ............. hgg and
h (6) [b: 2", 0, 0, O, Ohwheee the four auxiliary entries capture the order information. Details of each entry given in (21) and (22). Once H. is defined, the transformer produces the output H e RDx2", Note that even if the input length T is not exactly 2", the transformer can still replicate the operations of MASTER and generate length-T output by selecting the smallest 2 greater than T and running the algorithm for T rounds.
We then define functions or and 22 replicate the MALG operation: σi stochastically schedules instances for each order, and 02 selects the instance with the lowest order to remain active in each round.
Definition (T1). Given a mann-incrrasiingggnnccooo p [2"] R. we define of Px2 Dx2" as
(H)) ( (H ... :: ;0 (H n, n,
where for each i (O,....n)
a1 (H (H,
h B
(2")
j (1, 1,.,, 1}, mod 2 =1] E Dx2"
p(2t)
Here, B(k) is a Bernoulli random variable with parameter k, meaning B(k) 1 with probability & and B(k) = 0 otherwise.
Using 01, multiple instances can scheduled simultaneously (represented as the bluish blocks in Figure 8). Since only one instance can be active the environment at any given time, define ø2 to select the instance with the lowest order at each time step t.
We denote b, h (0)
h
(16)
R for each t [2"]. so that H [b1.... b2 Then, 02 is
defined as follows.
Definition (a2). We define 02 Dx2" - Dx2" by
02 (H) [bi, .... b/gm
61 - "o 0 0; +;;;,;;;,;** reproduce the uniqueness of instance scheduled every moment in MALG. Here, denote the last five entries of bc.
(n,0) h
(k-1)
O, h
* 0, 0 E
After applying 02, only a single instance remains active at each time t; this active instance is always the one with the lowest order among all scheduled instances. The reddish blocks in Figure 8 illussrate this selection.
To approximate the stationary tests (Test and Test 2) in Algorithm 2, we define the following:
Definition 6 (TEST), TEST represents the stationary tests in MASTER, involving two functions test1 and test2 that map from R RDx2" to R. Given defined in Algorithm 2 and an active instance in be of order where by contains m, I, U2, Fo and Tú, we define:
1, if some order-m alg ends and 2m
Tr U. + 97(2
test1,(b,, H)
T=1- 2m
0, otherwise
26

---

5+5
H
d+5
H
02
5+5
H
d+5
H
H
T= 23
T= 2
Figure 8: Illustration of H when 3. Purplish blocks represent instances scheduled by øı, while reddish blocks represent active instances selected by 02. Reddish blocks connected by a dashed line are concatenated.
1, if 1-++1
Σ r-1
rr) (
< 3p(t fn + 1)
test H)
0, otherwise
TEST is then defined as:
1, if testi(h:) = test2(h:)
Ths H) teati(h: H) test2(h:, H)
0, otherwise
Finally, we define the noncontinuous transformer. This transformer incorporates the stationarity tests and halts processing of the remaining sequence whenever test fails. To implement this, design transformer module that processes each element individually, performs the required tests, and then concatenates the results.
Definition (Noncontinuous Transformer). For t [2 denote
6""
MLP OLZ)
Attn attm
o... (MLP (2))
Attn (bi, H.) E RD"
Astn
where the parameters are g(f) ann
me[[2"] = DDD and ( ((w(9)
m
mlp
) c DDD x DxD' After processing hr, the model interacts with the environment, observes the reward, and inserts into the sequence: byt h, ((21) and (22)). By concatenating [ne[2"), we define the L-layer noncontinuous transformer TF Dx2" Rx22 as
TFo -TF(2")
(18)
TF())
(T) (H) HT( e Dx2" (2))
(19)
where ((T) = bi...., 5r. Tr br.1.... Tr h2n
Dx2"
Here, Tr E RDD is the test matrix
time T:
()++1
Tr == diag
TEST, TEST, TEST, TEST, 1, 1, TEST, TEST, TEST) (20)
d+1
where TEST = Tr H)) Order information E+ -1 Tr. Ei Pt .7t historical data
((k
E'-1 rr E'=1 Fr, U,-1, and rand in (21) and (22) are reset when TEST 0. 7 denotes the auxiliary value generated order i.
27

---

Interaction with the environment ment, we consider the following:
To illustrate the transformer's interaction with the environ
(
X
2
2
- Σ
2 - 1 2
Processed module
h
(21)
Interact with the environment
rand
Σ %
((k
Σ 8
(kk
6)
h (0)
h
h (1)
h
h
2
Processed module
2"
b
(22)
Interact with the environment
5
TT
Up
Here, Ee=1 and r, for & O. The rand entry is random number generated by the transformer. The vectors x!" IRd are used for in-context learning, and x( denotes their processed forms; for example, these vectors can encode action choices in bandit problems (Lin et al., 2024).
Each entry of the vectors captures cumulative information or constants relevant to the model's 's oper- ations up time t. After passing the sequence through the MLP+ Attention module, certain orders are scheduled at each time I, while all other instances are set to 0. The operation ensures that only one instance remains active time by selecting the instance with the lowest order and zeroing out all others. If the active instance is order & n, then p(() is assigned as n at time t, and the update U, min(U,-1, p(()) is applied. Finally, the module interacts with the environment, observes the reward R,, and updates the tensor accordingly.
Notation with regard to the input sequence To avoid potential confusion, the notations defined above are summarized below:
h E RR: The 1-th element of the input sequence.
H = [h....., h E pdx2"
The input sequence as a matrix.
H [H H (0) H
[h] ... h2-] E Dx2"; The extended input sequence, where
(+) denotes the order-i sequence.
6. = b h
h
E RP: The t-th vector of the extended sequence.
6/ 0.... 0; h (K)
;
RP, The t-th vector of the extended input after
being processed by 02. Here, k denotes the order of the active instance selected by a2, i.e., the lowest nonzero order at time
6,: The result of processing h, by the noncontinuous transformer module.
b,: The output of the noncontinuous transformer after incorporating the observed reward
at time
The restart mechanism (line 9 of Algorithm 2) is integrated into the noncontinuous trans- former. If all stationary tests pass, the matrices Tr (T E [2*]) remain identity matrices, and the sequence is processed normally. If a test fails at time T, the matrices TT (T [2"]) modify the states br+1.... b the first D 7 elements-includinn
Restart
xy, T, 2", and 2*-are preserved, while the remaining seven elements are set to zero. This ensures that when test fails, all historical infor-
mation is crased, but essential identifiers and order information retained to start new block.
28

---

Rollout The noncontinuous transformer's architecture is more complex than the classic version, so we outline its input-output flow.
Starting with input sequence H E Dx2", extend it to H. [ H H* Dx2", where submatrices H6S and H() are only different in several entries record order information.
H. is passed through di and d2: di schedules instances for each order, and d2 activates only the instance with the lowest order at any given time t. As result, at each time t, only one instance is active, and the active instance may come from different order each time step.
Next, each element in the sequence is processed by the traditional MLP+ Attention module. The module uses this information to interact with the environment, collects the reward, and inserts into the sequence: hy - b, ((21) and (22)). The updated sequence h2 is evaluated by TEST, which performs stationary tests. If the tests pass, the block remains unchanged; otherwise, the remaining block entries are set to zero, keeping only essential information. This effectively ends the current 2" le-Ieh block and starts new 2" 2""-lngthh block (line 2, Algorithm 2). On restart, variables like R, and ft are reinserted at their appropriate positions, following the same procedure. After processing the entire sequence, the final tensor H is produced. By gathering the rewards from all orders at each t (only one reward is nonzero at given time), the resulting sequence serves as the output of the noncontinuous transformer, containing cumulative rewards up to time 2".
E.2 Motivation of the Non-Continuous Transformer
In this section, we first briefly demonstrate how the proofs can be adapted a regular transformer, then explain the motivation of using the non-continuous transformer.
The core challenge in using regular transformer lies in approximating da (Definition 5), which must select the first non-zero entry among all entries. This operation requires comparing all trajectories simultaneously to identify the first non-zero one, and it is the main reason why the augmented inputs are used.
One alternative approach is to use regular transformer with n + 1 heads to approximate 0i and 02 as follows:
Layer 1 (Compression): Each head uses masked multi-head attention extract per- trajectory representations, i.e., H(O = Whomp MaskAutn, (H) E Ro xT Here, MaskAttn restricts information flows across heads, and the combined output is H' (H). ..... [()] RXT.
Layer 2~ Generate Bernoulli masks, block masks, and select the first non-zero trajectory to approximate dn and 02 (same as the procedure in Appendix E.3.2 and E.3.3).
Final Layer (Decompression): Σ Wideropp.. RRXT
Reconstruct the selected trajectory:
H"
However, this method adds several analytical burdens, such ensuring that Weccoop. Wccop lp for m E {O..., n. to preserve information, and constructing MaskAttn, which make the analysis more cumbersome.
In contrast, our proposed method is simple and transparent preprocessing step where augmented inputs generated via linear transformation h. M, h, This preserves the core semantics while avoiding unnecessary architectural detours. Importantly, the resulting non-continuous trans- former is structurally equivalent standard transformer with transformed inputs.
E.3 Approximation of WS (Proposition 3)
We divide the proof into three parts: approximating random number generation, ai and 02- Throughout this section, denote the ReLU activation function by a(-).
29

---

E.3.1 Approximating random number generation
We start by approximating 1= According to the definition of the ReLU function, we have
0, I & 0
a(kr) okkr -1) =
kar 0 x<
1, x >
(0, x 00
When l - s0, it approximates [[x]]
11, >
We denote 23 := Xk.0 and (2(z) being the cumulative distribution function (CDF) {E1,..., R2n ). Following the proof in Hataya and Imaizumi (2024), we construct 2-head attention layer to ap- proximate the CDF.
From the approximation of [[-]. P=(1) - 1/2" Ex= 1[xE:0 & t] can be approximated by sum of ReLU functions
2
(t(t)
1
Σ (k(t z)) (k(t z) 1)}
2
where k is sufficiently large. Suppose the vector in formula 21 and 22 (left) is one of the input h By selecting (Q1,2 21,, V1.2] such that
kxx,0
ox
kx,,
Q,
Q₂h,
K,h, K2h,
0
1
0
1
Vih;
V2h;
0
0
we have
Σ
Kmb)) Vmh, o (k(x),0 xx,0)) (k((x,,
x1,0) - 1).
m111
Therefore, the output
2"
bn rand
h h;
Σ Σ
2n
1-1 mmml
where rand= (x.,o) can be regarded random variable sampled from 2/(0, 1).
E.3.2 Approximating 01
Step Generating Bernoulli masks. Define Bernoulli Masks as Вi, [[rand, pp(2)) We consider a random number rand,, E R order and time t, and rand p(-) [0. Following the approach in Step 1, we begin by considering the input h This input includes rand,,,, 2 2", and we also assume that p(2*) and p(2") are added to (6) through p(.). this case, two-layer MLP can implement Bernoulli Masks.
First, consider the function fx, 3) z y R. For x. E [0, 1]. dividing the domain [o, [0. 1] into [1/c] segments allows us to express it as:
O.. + Be.ty Nk.t-
This form can approximate f using an MLP. Based on this strategy, the two-layer MLP is constructed follows:
(W b1)
30

---

h2 (p(2%) (01,,rand + 31,1p(2)) 7111)
122 (p(2") - (aj Jz Ia + P1 (1((2)))))))))))))))))))))))(1 (t
li2 (p(2*) randi, (1,10(22) + m1,))) -
k2 ((22)) ) + $ 11 )
where W1 TR2 The output is then processed by another layer:
W2 b (2(2))
rrand, Bx,Ip(2)) (k.1)))
k,==1
-k2(222(22)) Σ
(ag,irand,, Bx,up(2') (R.1)) -1)
k,1=1
(2)(((2)) - rand ((21)) (H2(o(2") - rand p(21)) - 1)
(1, rand,,1 $ p(2")
- 0, rand,,1 p(2))
p(2")
p(2:)
(12 - 00), W 1
The hidden dimensions of W, W are O(1e). If the input is expanded from h to H the hidden dimensions increase O(2"/e) Further expand further H. the hidden dimensions of Wi, W, become (e) OT log, T/s). The operator norms | Wi lop scale as O( T loga T/E). Finally, by taking ky - 0o, the error approaches zero.
Step 2: Generating block masks. By setting rand ... 2 (t = 1 (mod 2)), we have B = = Bi,e+2'-1 (f (mod 2)). Since the output of Bernoulli masks can expressed as H = [b ..... h2nl Dx2", where for / E [2"]:
(1)
(n)
X
hg
0
he
h (n)
h
Bo.t (0)
h
B1,
h
B 1
(0)
2p St
pp)
(n)
choose
(Qm: Km, V V mo....(.....)) such that for tokens has
he
Kmh
V Vnnba
0; 0; h (m)) 0; O...,
(m)
Bm.t
0.5
In this way, with we have
Σ
1s (mm Kmb V. ba = O,
m==
31

---

and with k = t we have
Bo (0)
n Σ m==
(QQmb Kmh:))
Vh.
1 [Bm,t 0.5]
(m)
17
B (n))
Finally, referring the approximation of -= (E.3.1), note that this indicator function be approximated using two ReLU functions. Consequently, the block mask can be approximated by 2(n 1)-head masked attention layer. By setting 1/e 4(n 1)², we have M O(I/VE)
Combining this block mask with the random number generation and the Bernoulli mask completes the proof.
E.3.3 Approximating d2
Based on the output of Ji we can express the input of 02 as H = [b, ... 1 ..-b-
R Dx2", where for
I [
(n)
hg
X
0
be
h
(0)
0
h!"
h.m)) (n)
h
(()))
(p)
(1)
From the definition of σ1, we know that some h values are nonzero, while others are zero. If we (ic)
assume i1 < < is (0 < ho & n) such that h, (11) ..... h,(t) * O while all other (n) values are 0, the goal of 02 is to ensure that ( and the others being
We choose { Qm., Km; Vm3meeo...... such that for tokens he.
Σ mm f(0)
Qmhe =
Km
CD h
bn = 0; O; 0;
(m²) ((m) ..., 0]
60
where > is sufficiently large and eo > 0 is sufficiently small such that - coofs < 0 for any fg 0. In this way, with k < / have
1s (mnba), hx)) V Vnhb
0,
and with 1: = /. we have
0
m7
Σ
120
6o
Co
ff
h
0
32

---

(m)
mn
S =
p-D
0,
otherwise
Since Σp-0 H. 0 equivalent to (10)
h
(p-1)
0, we have for each E [T] that
=0 [Qmbe Kmha>] Whb
h (g)
m=0
0
where h (0) (E.3.1), approximate - using ReLU functions. Therefore, d2 can be approximated by a 2( + 1)-head masked attention layer.
h (--1) and h (m) * 0. Finally, referring the approximation --
E.4
Equivalence of WS and RM with previous works
E.4.1
WS
Sliding windows used in previous works (e.g., (Cheung et al., 2022; Trovo al., 2020)) can be regarded as special cases of wS. Specifically, they are equivalent to WS when WS contains only one instance, all windows are of equal size, and they connected without overlap. Figure shows the equivalence between the stochastic instance scheduler in MASTER and the sliding window schedulers from previous works.
3
Time
order
order
Stochastic window scheduler
order
order
order 3
order
Fixed window scheduler
Figure 9: Equivalence between the stochastic instance scheduler in MASTER and the sliding win- dow schedulers from previous works. Purplis blocks represent scheduled blocks, while reddish blocks represent active ones. Reddish blocks connected by dashed line are concatenated. Here, 8 and the sliding window size is 2.
33

---

E.4.2 RM
The restart strategies in previous works can be broadly categorized into two types: stochastic and deterministic (Gomes al., 1998; Streeter and Golovin, 2008). Stochastic restarts typically involve some form of randomness or a test, whereas deterministic restarts are executed after a fixed number of rounds. Consequently, deterministic restarts can viewed as special case of sliding window strategies. As for stochastic restarts, given that transformers are capable of generating random num- bers and MLPs can implement stationary tests, it follows from the universal approximation theorem (Definition 8) that transformers can also approximate stochastic restart mechanisms.
Definition 8 (Universal Approximation Theorem for MLPs), Let f be continuous function compact subset K R R". Suppose a function applied component-wise.
non-polynomial, continuous activation
Then, for any O, there exists single hidden-lant layer MLP that approximates within Specif- ically, there exist weights A E Rakxn, be R C E IR mx and a sufficiently large number of hidden units such that the MLP
g(x) = C. C x
satisfies
sup esk
1f(x) - g(x)l <c
34


Reviews:
[{"review_id": "ml6GNp2MCU", "reviewer": "Reviewer_kshj", "summary": "This paper investigates the ability of transformer models to achieve optimal dynamic regret in non-stationary RL environments via in-context learning. The authors show that transformers, when trained with supervised learning on trajectories generated by expert algorithms, can implement mechanisms akin to restart-based and sliding-window strategies—common in adaptive RL—to adapt to non-stationarity. The main theoretical result proves that a transformer-based policy can achieve the dynamic regret bound in non-stationary RL.", "questions": "1. The experiments are limited to linear bandit environments. How do the authors anticipate their theoretical findings will extend to more general reinforcement learning settings, particularly those involving large or continuous state spaces and longer temporal dependencies?\n2. Assumption 2 requires that the transformer can approximate the base algorithm arbitrarily well. How realistic is this in practice, especially for complex algorithms or environments with long horizons?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n1. The paper tackles an important and under-explored question: whether transformers trained via in-context learning can effectively handle non-stationary reinforcement learning environments. This is an important extension over prior work, which mostly focuses on stationary settings.\n2. The authors prove that transformers can achieve optimal dynamic regret rates. The regret decomposition into approximation, training, and algorithmic terms is insightful and clearly articulated. Moreover, this paper shows that transformers can approximate key non-stationary strategies like windowing and restarts.\n3. The empirical results show that transformers can match or exceed baseline and expert algorithms, even in previously unseen conditions.\n\n**Weakness:**:\n\n1. The theoretical guarantees hinge on a realizability assumption: that the transformer can approximate the expert policy to arbitrary precision. Although motivated by prior work, this assumption is strong and not easily verifiable in practice.\n2. The architectural implementation of non-stationary mechanisms (e.g., the restart mechanism via 2-layer MLPs) is elegant, but it may oversimplify real-world scenarios. The paper does not explore how robust these designs are to scaling or how sensitive performance is to architecture choices.\n3. While the experiments are well designed, they are limited to linear bandit environments. The results may not generalize to more complex non-stationary MDPs or higher-dimensional, structured RL tasks. This limits the empirical reach of the theoretical claims.", "was_revised": true, "final_justification": "My main concerns have been addressed.", "rebuttal": "We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review.\n\n**Q2/W1:**\n\n> Assumption 2 requires that the transformer can approximate the base algorithm arbitrarily well. How realistic is this in practice, especially for complex algorithms or environments with long horizons?\n\nThank you for raising this important point regarding the realizability assumption. We clarify below its motivation, scope, and practical implications.\n\n**Theoretical grounding:** Assumption 2 is based on prior work showing that transformers can emulate standard reinforcement learning algorithms such as LinUCB and Thompson Sampling via in-context learning [1]. Furthermore, [2] shows that transformers are universal approximators over distributions, which supports the feasibility of such an assumption in principle. It’s also worth emphasizing that the assumption doesn't require the transformer to exactly replicate the expert policy. Rather, it assumes the transformer can learn an *amortized algorithm* that reproduces the expert’s output distribution, which is a more relaxed and realistic target in many cases.\n\n**Practical feasibility:** We emphasize that Assumption 2 does not imply that any complex algorithm or environment is immediately approximable by arbitrary transformers. Rather, it reflects a scaling law—namely, that approximation error $\\varepsilon$ can be made small by increasing model size and training data, as reflected in the regret bound (Theorem 1) and its dependency on architecture parameters (e.g., $D' = O(\\sqrt{T \\log^2 T / \\varepsilon})$). In this sense, the assumption is realistic to the extent that the base algorithm itself is implementable by a recurrent or sequential decision process. For many commonly studied algorithms, including sliding-window UCB and restart-based methods like MASTER, this has been shown to be the case in structured architectures [3].\n\n**Constructive verification:** In Section 4.2 and Appendix D.2, we construct specific transformer blocks that replicate components of non-stationary algorithms—such as window schedulers and restart logic—using standard attention and MLP modules. This shows the assumption is not merely abstract but realizable within standard architectures.\n\n**W2:**\n\n> The architectural implementation of non-stationary mechanisms (e.g., the restart mechanism via 2-layer MLPs) is elegant, but it may oversimplify real-world scenarios. The paper does not explore how robust these designs are to scaling or how sensitive performance is to architecture choices.\n\n**Theory:** Instead of exhaustively exploring design choices, our algorithmic construction is designed to demonstrate the existence of a transformer that can approximate the expert algorithm across the entire input space. For the restart mechanism specifically, many RL algorithms adopt either deterministic strategies (e.g., periodic restarts) or rule-based ones (e.g., resets triggered by specific conditions). Deterministic approaches can be viewed as a form of window scheduling, where the algorithm either discards historical data or resets its internal state at fixed intervals. Rule-based strategies often rely on binary indicators, which can be naturally modeled by a 2-layer MLP. In this sense, our architectural implementation of non-stationary mechanisms is expressive enough to capture a broad class of restart behaviors, at least from a theoretical standpoint.\n\n**Experiment:** While more in-depth experiments on scaling and performance sensitivity of the architecture are left to future work, we conducted experiments using a transformer with 4 heads and 4 layers (compared to the 16 heads and 16 layers described in the paper) under the Low Non-stationarity setting to address the reviewer’s concerns. The \"Regret Over Time\" results are summarized in the table below. It demonstrates that the transformer is capable of achieving performance comparable to that of the expert, even with 4 heads and 4 layers.\n\n| Model\\Round | 50 | 100 | 200 | 300 | 400 | 500 | 600 | 700 | 800 | 900 | 1000 |\n|-------------|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|------|\n| Transformer | 14.5 | 29.5 | 59.7 | 89.8 | 119.9 | 150.1 | 175.3 | 205.6 | 235.5 | 269.2 | 301.2 |\n| LinUCB | 13.3 | 26.1 | 51.7 | 77.3 | 102.8 | 128.3 | 153.7 | 179.2 | 210.3 | 235.7 | 262.3 |\n| LinUCB+MASTER | 14.3 | 29.0 | 57.1 | 83.7 | 111.0 | 138.8 | 165.5 | 193.7 | 221.8 | 250.1 | 278.3 |\n\n**Q1/W3:**\n\n> How do the authors anticipate their theoretical findings will extend to more general reinforcement learning settings, particularly those involving large or continuous state spaces and longer temporal dependencies?\n\nThank you for the question. We believe that the implications from our results are applicable to settings such as continuous state spaces and longer temporal dependencies. As this work is primarily theoretical, we initially focused less on experiments. However, to better address the reviewer’s concerns, we extended our evaluation to a continuous state space environment.\n\nWe ran additional experiments on the MuJoCo Hopper task, which features an 11D continuous state space (e.g., joint angles, velocities) and a 3D continuous action space corresponding to joint torques. The objective is to learn a hopping behavior that balances forward progress, energy efficiency, and stability. Rewards are shaped accordingly to encourage forward velocity while penalizing excessive control effort and contact forces.\n\nWe used Soft Actor-Critic (SAC) [4] as our base algorithm and SAC+MASTER as the expert. For comparison, we trained a 3-layer transformer (with 256D embeddings) on stationary data for 20,000 episodes. During training, the transformer was evaluated on non-stationary data every 100 episodes, by switching to ```model.eval()``` to perform evaluation for 10 episodes, and then switching back to ```model.train()``` for further training. Similarly, SAC was trained on stationary data, and every 10 episodes, it was evaluated in a non-stationary environment.\n\nThis setup follows prior work (e.g., [5, 6]). Although individual state components range over $\\mathbb{R}$, typical ranges are more bounded: the z-axis lies roughly within [0, 2], joint angles within [–1, 1], and velocities within [–10, 10]. We introduce non-stationarity by adding a drift value to the states, defined as $\\sin(0.01 t)$. (new_state = state + $\\sin(0.01 t)$)\n\n**Results:**\n\n**Transformer:**\n\n| Episode | 100 | 200 | 300 | 400 | 500 | 1K | 1K~10K | 10K~20K |\n|---------|-----|-----|-----|-----|-----|-----|--------|----------|\n| Avg. Reward | 113 | 54 | 68 | 52 | 51 | 61 | 42~62 | 42~47 |\n\n**SAC:**\n\n| Episode | 200 | 400 | 1K | 2K | 2K~5K | 5K~500K |\n|---------|-----|-----|----|----|-------|--------|\n| Avg. Reward | 26 | 42 | 40 | 40 | 6~40 | 5~44 |\n\n**SAC+MASTER:**\n\n| Step | 200 | 400 | 1K | 2K | 2K-5K | 5K-50K | 50K-100K | 100K-1M |\n|------|-----|-----|----|----|-------|--------|----------|---------|\n| Avg. Reward | 14 | 12 | 19 | 30 | 10-214 (mostly 10-30) | 9-201 (mostly 10-30) | 7-180 (mostly 10-30) | 6-217 (mostly 10-30) |\n\n** Typically, SAC requires over 1M steps to converge. Due to time constraints, we only trained SAC+MASTER for 1M steps, as its learning process is much faster than SAC's due to the restart mechanism, which maintains smaller active memory pools and results in faster computation.\n\n** SAC+MASTER Avg Reward after 900K steps (converged at this phase): Among all rewards in the evaluation episodes, 21% of the rewards are over 50, 13% over 70, 7% over 100, and 2% over 150.\n\nThe results show that the transformer maintains more stable performance under observation drift and achieves rewards comparable to the expert. While the reported SAC and transformer scores may appear lower than those in earlier work (which often exceed 1,000), we validated our SAC and transformer implementation under standard stationary conditions and confirmed it aligns with previous findings before introducing non-stationarity.\n\nWe acknowledge that this single environment does not fully capture the complexity of general RL tasks, but we view this as a meaningful first step toward evaluating the theoretical framework in richer, more dynamic environments. We plan to explore additional tasks with longer horizons and more complex dynamics in future work.\n\n[1] Lin, L., Bai, Y., & Mei, S. Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining. In The Twelfth International Conference on Learning Representations.\n\n[2] Furuya, T., de Hoop, M. V., & Peyré, G. Transformers are Universal In-context Learners. In The Thirteenth International Conference on Learning Representations.\n\n[3] Wei & Luo (2021). Non-stationary reinforcement learning without prior knowledge: An optimal black-box approach. COLT.\n[4] Haarnoja, T., Zhou, A., Abbeel, P., & Levine, S. (2018, July). Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning (pp. 1861-1870). Pmlr.\n\n[5] Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., ... & Mordatch, I. (2021). Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 15084-15097.\n\n[6] Huang, S., Hu, J., Yang, Z., Yang, L., Luo, T., Chen, H., ... & Yang, B. (2024). Decision mamba: Reinforcement learning via hybrid selective sequence modeling. Advances in Neural Information Processing Systems, 37, 72688-72709."}, {"review_id": "RbncjWAMGn", "reviewer": "Reviewer_NBt8", "summary": "This paper investigates the performance of transformers in non-stationary reinforcement learning environments. Building on prior work that showed transformers can approximate optimal online RL algorithms in stationary settings, this study extends the analysis to dynamic environments where reward distributions and transitions evolve over time. The authors prove that a transformer, trained in a supervised in-context learning setting, can achieve minimax optimal cumulative regret for non-stationary environments for bounded changes. They provide theoretical guarantees on generalization under distribution shifts and identify architectural requirements for transformers to implement restart-style strategies necessary for adapting to non-stationarity. A novel proof technique based on internal algorithm selection is introduced.", "questions": "While your theoretical results are strong for bandit settings, what are the challenges in extending this to more general RL algorithms and potentially more structured models of change?", "limitations": "Yes, the authors have addressed the limitations and potential negative social impact of their work.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n- The problem of understanding transformer behavior for non-stationary reinforcement learning is important\n- The paper establishes a $\\mathcal{\\tilde{O}}(T^{2/3})$ regret bound for an in-context bandit learner matching the information theoretic lower bound\n- The elegant internal algorithm selection approach offers a new method for analyzing in-context learning in transformers by conceptually formalizing how they might choose between different internal strategies\n- Identifying specific architectural features that enable adaptation to non-stationarity provides valuable insights for designing more robust transformer models\n\n\n**Weaknesses:**\n- While the paper sets up a general RL environment, the regret optimality results and much of the discussion are heavily geared towards bandits. While bandits serve as an excellent starting point for theoretical analysis, it is not entirely clear how directly these results translate to more complex general reinforcement learning settings.\n-  The reliance on a supervised in-context learning setup for the transformer raises questions about data efficiency in practical non-stationary RL. Continuously generating or providing optimal input-output pairs for training the transformer to adapt to changing dynamics might be prohibitively expensive or impossible.", "was_revised": true, "final_justification": "Why Accept: the paper addresses an interesting and important question of understanding the behavior of transformers as in-context bandit/RL learners; elegant internal algorithm selection idea; insights for designing more robust transformer architecture; establishes matching upper and lower bounds \n\nWhy not strong accept: concerns on data efficiency remain", "rebuttal": "We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review.\n\n**Q1/W1:**\n\n> While your theoretical results are strong for bandit settings, what are the challenges in extending this to more general RL algorithms and potentially more structured models of change?\n\nThank you for the insightful question. While our experiments are limited to bandit settings, the theoretical framework extends to more general reinforcement learning (RL) problems, as long as the expert policy can be expressed as a distribution over actions given trajectories—a property shared by many episodic or tabular RL algorithms.\n\nA key challenge in general RL is the added complexity from temporal credit assignment and state transitions. Unlike bandits, RL requires modeling long-term dependencies, which may necessitate deeper architectures and longer contexts.\n\nTo address this, our approach does not require the transformer to replicate the expert’s full decision process. Instead, it *approximates the expert’s output distribution (amortized inference)*, which is more tractable and better aligned with in-context learning paradigms. This relaxation is consistent with prior work (e.g., [1]) and backed by universal approximation results.\n\nExtending our analysis to structured forms of non-stationarity (e.g., piecewise-stationary MDPs) is an important next step. We believe our decomposition, especially of restart and windowing mechanisms, provides a promising foundation, and we now discuss this in the conclusion.\n\n**W2:**\n\n> Continuously generating or providing optimal input-output pairs for training the transformer to adapt to changing dynamics might be prohibitively expensive or impossible.\n\nWe appreciate the reviewer’s observation. Indeed, in certain non-stationary reinforcement learning settings, conventional RL algorithms may offer greater reliability or data efficiency than transformer-based approaches. However, the primary goal of our work is to establish a theoretical guarantee for the robustness of transformers under non-stationarity, rather than to optimize for data efficiency or other empirical performance aspects.\n\nCorollary 2 states that the regret bound holds when the number of training samples $N\\geq CT^3 \\log T$. If we take $T=10^3$, which already covers the horizon of many RL tasks, this gives a data requirement on the order of $10^9$. While this is certainly large compared to what's typically needed in traditional RL algorithms, it's still relatively modest compared to the scale of data used to train LLMs.\n\n[1] Furuya, T., de Hoop, M. V., & Peyré, G. Transformers are Universal In-context Learners. In The Thirteenth International Conference on Learning Representations."}, {"review_id": "4U525VEcnX", "reviewer": "Reviewer_JYfA", "summary": "This paper considers transformers and their use in approximating reinforcement learning algorithms. In particular, they consider whether a transformer can be trained to act as an RL algorithm when given an interaction history. Previous work (Lin et al 2024) showed that a sufficiently large transformer model can approximate popular near-optimal online RL algorithms such as linUCB. However, the question of whether transformers can adapt to non-stationary environments has not been answered. In a non-stationary environment, the reward can evolve over time. The paper proves that a transformer can achieve cumulative regret matching the minimax optimal rate for non-stationary environments.", "questions": "My main question is how strong of an assumption is Assumption 2 (realizability)? I am not familiar with this assumption and would like some more context surrounding it.", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths: \n1) The theoretical results seem strong and sound, although I did not fully verify the proofs. \n2) The paper is pretty well-written and understandable. \n3) There are empirical results to back up the theoretical results. \n\nWeaknesses: \n1) Assumption 2 seems a bit strong, and there is little discussion on its importance and whether it is standard/ \n2) There are few minor typos and grammar mistakes.", "was_revised": true, "final_justification": "I am keeping my score. I think the paper is generally well written and I like the theoretical results. I thank the authors for providing more context regarding their assumption.", "rebuttal": "We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review.\n\n**Q1/W1:**\n\n> How strong Assumption 2 is?\n\nWe appreciate the reviewer raising the question. While Assumption 2 is indeed strong and difficult to verify empirically, it builds on solid theoretical foundations. For example, [1] shows that transformers are universal approximators over distributions, which supports the feasibility of such an assumption in principle. \n\nIt’s also worth emphasizing that the assumption doesn't require the transformer to exactly replicate the expert policy. Rather, it assumes the transformer can learn an *amortized algorithm* whose output distribution matches that of the expert, which is a more relaxed and realistic target in many cases.\n\nA similar assumption is also used in [2], and several previous works (e.g., [3, 4, 5, 6]) have proved the expressivity of transformers from other aspects.\n\n**W2:**\n\n> There are few minor typos and grammar mistakes.\n\nThank you for pointing this out. We have carefully revised the manuscript to correct all identified typos and grammatical issues.\n\n[1] Furuya, T., de Hoop, M. V., & Peyré, G. Transformers are Universal In-context Learners. In The Thirteenth International Conference on Learning Representations.\n\n[2] Lin, L., Bai, Y., & Mei, S. Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining. In The Twelfth International Conference on Learning Representations.\n\n[3] Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S., & Kumar, S. Are Transformers universal approximators of sequence-to-sequence functions?. In International Conference on Learning Representations.\n\n[4] Zhang, R., Frei, S., & Bartlett, P. L. (2024). Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49), 1-55.\n\n[5] Bai, Y., Chen, F., Wang, H., Xiong, C., & Mei, S. (2023). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 57125-57211.\n\n[6] Ahn, K., Cheng, X., Daneshmand, H., & Sra, S. (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 45614-45650."}, {"review_id": "txSv0OHf6h", "reviewer": "Reviewer_e7t9", "summary": "This paper analyzes the performance of transformers in the setting of in-context reinforcement learning (ICRL) in non-stationary environments. \n\nThe learning setup is the same as the (stationary) ICRL setup of Lin et al (2024).  The learner is given N iid trajectories of the form $D_T = (s_1, a_1, r_1, … s_T, a_T, r_T)$ generated by an offline algorithm alg0, as well as corresponding optimal actions $\\bar{a_1}, …, \\bar{a_T}$ generated by a base algorithm algB (which produces “high-quality” actions), the transformer is trained to maximize $ \\log  p(\\bar{a}_t | D_{t-1}^i, s_t^i)$. \nThe paper analyzes the dynamic regret of such a pre-trained transformer in a non-stationary environment, with non-stationarity defined similarly as in Wei and Luo (2021). \n\nSimilarly to Wei and Luo (2021), the authors consider the availability of a base algorithm ALG for stationary environments which produces a quantity $\\tilde{r}_t$ that can be used to detect non-stationarity (such as an UCB on the reward). Then they consider the MASTER algorithm of Wei and Luo,  which converts a base algorithm ALG into one that can deal with non-stationarities (by running ALG at randomly-selected multiple time-scales, with periodic resets whenever non-stationarity is detected). The authors break down the dynamic regret of the transformer into three components: (1) regret of MASTER, (2) regret of transformer wrt MASTER, and (3) approximation error in the transformer implementing MASTER. (2) is bounded by extending the work of Lin et al (2024) to the non-stationary setting. \n\nBounding the approximation error is the main technical contribution. The authors show that small transformers / MLPs can be used to implement the randomized window-scheduler and restarts. They subsequently show that the overall MASTER algorithm can be implemented by a \"non-continuous transformer” (only introduced and defined in the appendix). \n\nThe paper also includes some experiments on a linear bandit setting.", "questions": "Why is the definition of the non-stationarity measure different from Wei and Luo? They define it as the maximum change in the expected reward of a policy, $\\max_\\pi |r_{t+1}(\\pi) - r_{t}(\\pi)|$. Your definition has $\\max_{s, a, s’, a’} |r_{t+1}(s’, a’) - r_t(s, a)|$ which could be non-zero for a stationary reward.\n\nThe MASTER algorithm can be implemented by the “non-continous transformer”, which is different from the standard architecture, and operates on augmented inputs. Why not state that in the main text? Does using this architecture complicate adapting the results of Lin et al (2024)?\n\nAssumption 1 assumes that the transformer algorithm “outputs” the auxiliary quantity $\\tilde{r}_t$. But the transformer is only trained to output the actions. Should it be trained to output both? Or should it be given access to some other function computing auxiliary values? Or does it suffice for the transformer to be able to compute these values in some intermediate block?", "limitations": "The discussion of limitations seems inadequate. The only stated limitation is in relation to the experiments. The authors should additionally discuss the discrepancies between the analyzed non-continuous transformer and the standard architecture. The limitations discussed in Lin et al (2024) also apply to this work and could be mentioned.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Quality:\nMost components of the submission seem technically sound. I am worried that the architecture and data processing described in Appendix D, used to prove Theorem 5 (that transformers can approximate the MASTER algorithm), does not match the architecture description in the main text. \nThe experiments are inadequately described and not especially convincing (but this is mostly a theory paper). \n\n\nClarity:\n\nWriting is clear overall, but there is room for improvement in the organization of the paper.  \nIn particular, the non-continuous transformer (architecture that can implement MASTER) is only introduced and defined in the appendix, and the description is difficult to parse. This architecture and data augmentation should be discussed (or at least mentioned) in the main body. \n\nThe proof outline mostly focuses on the approximation error, but it may be useful to include other components as well (e,g. Section C.2, and and outline of how the results of Lin et al are adapted to the non-stationary setting in C.5). \n\nIn the experimental setup, the first paragraph states that the action set is changing over the rounds, but the subsequent text suggests that the actions are fixed and the range of the rewards is changing. In (a) it changes to [3, 4] for some rounds (how? do you just add 3 to the reward?) and in (b) reward is multiplied by a cosine? The Appendix does not clarify anything. \n\nMinor:\n- $alg_{\\mathcal{E}}$ is used to denote the expert algorithm, and $\\varepsilon$ to denote the approximation error, and these epsilons look very similar. \n- The “base” algorithm is sometimes denoted by alg_B and sometimes by alg_E throughout. Assumption 2 has alg_B in the text and alg_E in the equation. I assume they mean the same thing - please pick one and make the text consistent. \n- Appendix C.4 (proof of Thm 5) starts with “consider ${\\bf x}_t^{(i)} \\in \\mathbb{R}^d$ \" but then never references this x again\n\nSignificance:\n\nThe submission advances the understanding of in-context RL by transformers to the setting of non-stationary RL, which is important. It does not provide any new algorithms, and the architecture changes used in the proof are not evaluated empirically (perhaps a bit cumbersome to implement). Therefore this is unlikely to influence practitioners. \n\nOriginality:\n\nThe paper puts together the results of Lin et al (2024) and We and Luo (2021), and mostly looks at the expressivity of transformers when it comes to the MASTER algorithm. While this is not especially original, it is highly non-trivial to execute and seems like a worthy contribution.", "was_revised": true, "final_justification": "Minor issues resolved. Authors have committed to describing the actual architecture used for the proof into the main text.", "rebuttal": "We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review.\n\n> Why is the definition of the non-stationarity measure different from Wei and Luo?\n\nWe thank the reviewer for raising this important point. Our initial intention was to align with the learning setup in Lin et al. (2024), and thus defining the non-stationary measure as $\\Delta(t) = \\max_{s,a} | R_t (s,a) - R_{t+1} (s,a) |$. However, we inadvertently wrote it as $\\Delta(t) = \\max_{s,s',a,a'} | R_t (s,a) - R_{t+1} (s',a') |$. We have now corrected the definition to $\\Delta(t) = \\max_{s,a} | R_t (s,a) - R_{t+1} (s,a) |$, which is the **upper bound** of the one in Wei et al. It is worth noting that using $\\Delta(t) = \\max_{s,a} | R_t (s,a) - R_{t+1} (s,a) |$ instead of the definition in Wei et al. has minimal impact on the overall approximation, while offering better alignment with the setup in Lin et al.\n\n> Does using the “non-continuous” architecture complicate adapting the results of Lin et al (2024)?\n\nNo, using the \"non-continuous\" architecture does not complicate adapting the results of Lin et al. (2024). The non-continuous transformer is essentially equivalent to conventional transformers. In the proofs of Lin et al. (2024), a specific transformer was constructed to approximate each RL algorithm, demonstrating the existence of a transformer that can achieve the error bound. Similarly, we constructed a specific transformer to approximate the MASTER algorithm. While the non-continuous transformer may differ in terms of augmented inputs, the main impact of these inputs is simply increasing the model size. Since our work is focused on the theoretical approximation rather than practical considerations, this does not complicate adapting the results of Lin et al. (2024). \n\nIn fact, our initial proof used a conventional transformer without augmented inputs and relied on multiple heads to compensate. However, due to the complexity of the proof, we later switched to the augmented input strategy for simplification. \n\nWe appreciate the reviewer for highlighting this, and we will include a brief explanation of the non-continuous transformer in the revised version of the main text.\n\n> Should the transformer algorithm be trained to output both the auxiliary quantity and the actions? Or should it be given access to some other function computing auxiliary values? Or does it suffice for the transformer to be able to compute these values in some intermediate block?\n\nThank you for this insightful question. In conventional RL algorithms, auxiliary quantities (e.g., UCB scores, Q-values) are explicitly computed to inform action selection. In our transformer-based framework, we adopt a different but related strategy: the transformer is trained to output a probability distribution over actions, from which the action is sampled. This output serves a role of implicit encoding for the auxiliary information used to guide decision-making.\n\nIn other words, the auxiliary quantity is not a separate supervised target, but rather emerges from the transformer’s internal computation. Our theoretical construction (see Proposition 3 and Theorem 5) shows that the transformer can implement implicit mechanisms, such as window schedulers and restart tests, that rely on internal auxiliary statistics. Thus, it is sufficient for the transformer to internally represent and use such quantities, rather than outputting them explicitly or requiring access to an external function.\n\nWe appreciate the reviewer for raising this important point, and will add a brief explanation of this in the revised version to clarify how these auxiliary quantities are handled in the transformer framework.\n\n> The discussion of limitations seems inadequate. The authors should additionally discuss the discrepancies between the analyzed non-continuous transformer and the standard architecture. The limitations discussed in Lin et al (2024) also apply to this work and could be mentioned.\n\nThank you for your suggestions. We will expand the discussion of limitations in the revised version to address the important points you raised.\n\nIn particular, we will explicitly discuss the discrepancies between our analyzed non-continuous transformer and standard transformer architectures used in practice. Our theoretical construction relies on controlled masking and conditional restarts, which, while implementable in principle, do not directly correspond to standard autoregressive or encoder-decoder implementations. This mismatch may affect practical applicability and will be acknowledged clearly.\n\nWe also agree that the limitations discussed in Lin et al. (2024) such as brittleness under distribution shift, reliance on synthetic pretraining distributions, and sensitivity to context length—are relevant to our setting as well, and we will incorporate them in our limitations section.\n\nWe’ll also address the gap between theory and experiments, as well as the generalizability of our results to broader RL settings, since we only did experiments on bandit settings. \n\n> In the experimental setup, the first paragraph states that the action set is changing over the rounds, but the subsequent text suggests that the actions are fixed and the range of the rewards is changing. In (a) it changes to [3, 4] for some rounds (how? do you just add 3 to the reward?) and in (b) reward is multiplied by a cosine? The Appendix does not clarify anything. \n\nWe apologize for the unclear explanation in the original text. To clarify:\n* In (1), the rewards in the rounds $t\\in[50,100]\\cup[350,400]$ are within the range $[3,4]$, while for the rest of the rounds, the reward $r\\in[0,1]$. To achieve this, we added 3 to the rewards in the specified rounds.\n* In (2), the reward is indeed multiplied by a cosine function, as stated.\nWe will revise the manuscript to make these details clearer and avoid any confusion in the future. Thank you for your understanding.\n\n> Clarity: In particular, the non-continuous transformer is only introduced and defined in the appendix, and the description is difficult to parse. This architecture and data augmentation should be discussed (or at least mentioned) in the main body.\n\nThank you for the suggestion. We have added descriptions and discussion of the relevant parts in the main text to provide greater clarity."}]

Author rebuttal:

Meta-review: This paper demonstrates that transformers can achieve optimal dynamic regret using in-context learning in non-stationary RL environments.  After discussion, the reviewers were uniformly positive about the paper.  There were some concerns about the strength of the assumptions (e.g., the realizability assumption) that I read as having been resolved in discussion.  The initial experimental evaluation was limited to a bandit setting; in response to reviewer comments, the authors have added an evaluation on a continuous MuJoCo environment (the Hopper task).

Overall this seems like a sound contribution that is likely to be influential.
Decision: Accept (poster)