Understanding the Training and Generalization of Pretrained Transformer for Sequential Decision Making
pretrained transformer;in-context learning;bandits;dynamic pricing;sequential decision making
We thank the reviewer for their quick feedback. While we understand the concern, we respectfully disagree with the perspective on the necessity of adopting a transformer-based decision-making agent. We hope the following clarification provides a clearer picture and kindly ask the reviewer to reconsider after reviewing our explanation.

# General Considerations

The reviewer may have overlooked two important factors.

1. **Applicability to Complex Decision-Making Problems**

We would like to highlight that our pipeline is capable of handling more complex yet practical problems beyond simple bandits. In real-world decision-making scenarios, such as dynamic pricing and newsvendor problems, it is often necessary to switch between different models depending on the environment. To make an analogy to the bandit setting, an environment may alternate between a linear bandit problem (requiring the LinUCB algorithm to perform well [1]) and a stochastic bandit problem (requiring the UCB algorithm). A unified framework like our transformer-based approach effectively handles these environments by identifying the environment and selecting the right algorithm when seeing the current trajectory, as demonstrated in Figure 12. This versatility extends the applicability of our method beyond simple bandit problems to more complex decision-making tasks. We provide a detailed discussion of the pricing settings at the end of this response.

2. **Foundational Step Toward General AI**

This approach is an essential step toward artificial general intelligence (AGI), which holds significant implications for the future of AI. If we do not analyze it within the bandit setting as a foundation, how can we extend it to more general RL or practical applications? For instance, the well-known online RL algorithm UCBVI [2] builds on the development of UCB in bandit problems. Therefore, we believe that studying the transformer framework in the bandit setting is both meaningful and necessary.

# Practical Relevance of Pretraining in Decision-Making Problems

Next, we elaborate on why the transformer-based decision-making agent is practical for our decision-making problems.

In the dynamic pricing example (a widely studied subject, see [3]), imagine a company (e.g., Amazon) that sells products to customers. At time $t$, the company needs to set the price $p_t$, which corresponds to a demand function $d(p_t)$, with the goal of dynamically maximizing revenue $\sum_{t=1}^T p_t \cdot d(p_t)$. In practice, the demand pattern changes from week to week—sometimes $d(p)$ is a linear function, and other times it is polynomial.

To achieve maximal revenue, two issues need to be addressed:

* Identify the environment and the right algorithm: this is important because using an algorithm designed for linear demand will perform poorly for an environment with polynomial demand.
* Quickly adapt: Often, it is impractical to start learning from scratch each time since the demand pattern changes frequently. The algorithm must quickly find the near-optimal price.

Therefore, it is necessary to pretrain the decision-transformer on many simulated environments. As shown in Figure 1, the larger the simulator environment $\mathcal{S}$, the more likely the transformer will have encountered patterns observed in the real environment. Once the transformer identifies these patterns, it can quickly adapt and find the near-optimal price. The transformer we developed is highly effective in addressing this setting (as shown in Figures 4(c) and 4(d), and Figure 12 in the appendix).





[1] Azar, Mohammad Gheshlaghi, Ian Osband, and Rémi Munos. "Minimax regret bounds for reinforcement learning." International conference on machine learning. PMLR, 2017. 

[2] Li, Lihong, et al. "A contextual-bandit approach to personalized news article recommendation." Proceedings of the 19th international conference on World wide web. 2010.

[3] Den Boer, Arnoud V. "Dynamic pricing and learning: historical origins, current research, and new directions." Surveys in operations research and management science 20.1 (2015): 1-18.
2nd Round Response
Thank you for the clarification.

My concern about computational and memory overhead (or other potential drawbacks) was specifically in comparison to running TS/UCB in practice, rather than to other generative language models. I remain somewhat unconvinced that, in practice, the more complex pre-trained transformer approach offers a clear advantage or necessity over simpler structured methods like TS/UCB in bandit problems, especially given the costs associated with pre-training on a large number of environments and trajectories. In reinforcement learning, I understand that it’s often challenging to estimate every state-action pair due to the complexities introduced by the transition matrix; therefore, having access to a good simulator can make pre-training models offline to narrow down the scope beneficial. However, in the simpler bandit setting, TS/UCB-type algorithms can already learn optimally from scratch, so I’m not sure why pre-training a more complicated model would be preferable.

Again, I do see the conceptual value of the framework itself and its connection to performative prediction, but the concern above raises questions about the practical significance of the contribution. I’m not an expert in the field, but with my best understanding, I’ll maintain my score.
Response to author
We appreciate the reviewer’s questions and suggestions. It seems there may be some misunderstandings about the specific setup and other details of our paper. We hope our response clarifies these points and kindly invite the reviewer to re-evaluate our contributions in light of this clarification.
``` 
Weakness 1: My biggest concern is the lack of discussion ...
```
We do not see significant issues with respect to the computational and memory overhead. In terms of computational overhead, our model scales linearly with respect to the context length, as is the case with many generative language models, which is totally acceptable. In terms of memory, our transformer adopts the GPT-2 architecture, which requires less than 1 gigabyte of VRAM and can run on almost all GPUs compatible with CUDA. Therefore, we don’t think there is a significant issue in terms of inference speed or memory requirements.
```
Weakness 2: The comparison between the proposed method and UCB/TS seems unfair, as the transformer benefits from...
```
We believe there may be a misunderstanding of our setup. Our transformer is pretrained on a vast number of environments, with each environment generating hundreds of thousands of trajectories. In the evaluation or online learning phase, if the current environment and its sampled trajectories were known, we could leverage that data in combination with the UCB/TS algorithms. However, since the environment is unknown beforehand, using the large volume of offline data indiscriminately could potentially degrade performance. While some algorithms may leverage massive pretraining data from heterogeneous environments online, fully understanding their analytical and computational properties would require a dedicated study, which is beyond the scope of this paper. For this reason, we adhere to the original UCB/TS algorithms in our comparisons.
```
Weakness 3: The claimed short-term advantage of the proposed method also seems unconvincing. While it’s true that...
```
To address this question, we would like to reiterate our main contribution: a stable and unified pretraining pipeline for transformers in the context of decision-making, with both theoretical guarantees and strong empirical performance. This pipeline is designed to be broadly applicable across various settings, including stochastic bandits, linear bandits, dynamic pricing, and newsvendor problems.

We believe it would be inappropriate to compare our method with problem-specific algorithms, as those algorithms are tailored to individual problem instances rather than offering a unified framework. The strength of our approach lies in its generality and versatility across different problem domains, which is distinct from optimizing performance in specific scenarios.
```
Weakness 4: I am not very familiar with related work on applying deep learning...
```
We thank the reviewer for suggesting this, and we will put this in the updated version.
```
Weakness 5: The clarity and organization of the paper could be improved...
```
We thank the reviewer for suggesting this, and we will put this in the updated version.
```
Question 1: What is the difference between the setting described in the paper...
```
The setting in our paper encompasses both the stochastic bandit and contextual bandit settings. In Appendix B, we provide detailed examples that fall within this framework. We believe this is the most general formulation, as it includes these specific settings. A contextual bandit framework, by contrast, would be considerably less general.
```
Question 2: What are the disadvantages of using pre-trained transformers compared to structured algorithms like UCB and Thompson Sampling...
```
The computational cost grows linearly (similar to the UCB and TS algorithms) with respect to the context length, as it corresponds to the number of tokens that need to be evaluated by the large language model, which we believe is entirely acceptable in practice. In terms of memory, our transformer is based on the GPT-2 architecture, requiring less than 1 GB of VRAM and running on most CUDA-compatible GPUs. Therefore, we do not anticipate significant practical issues with inference speed or memory requirements.
```
Question 3: If UCB/TS had access to the same amount of pretraining data, how would...
```
As mentioned earlier, for UCB/TS to effectively leverage the pretraining data, they would need to know the current environment and its corresponding trajectories. Without this knowledge, it is straightforward to construct counterexamples where trajectories from unrelated environments could adversely impact UCB/TS performance by distorting the confidence bounds or posterior estimates.
Rebuttal to Reviewer Svgy
Yes
No ethics review needed.
1. How do the authors choose the environments for pre-training, and do variations in the gamma values during training impact the model’s ability to generalize? It would be interesting to know if a big difference between training and testing gamma values affects performance or creates bias.

2. Have the authors thought about using curriculum learning, where self-generated actions gradually increase over time, instead of the two-phase switch? I think it would enjoy less theoretical property, but possibly perform better in numerical experiments.
1. This paper applies transformers in the domain of sequential decision-making, typically dominated by traditional RL approaches. The authors ground their framework with rigorous theoretical proofs, establishing conditions under which the mixed pre-training algorithm minimizes cumulative regret. They also prove that the learned decision function asymptotically matches the Bayes-optimal policy as the data approaches infinity.

2. I like how the authors propose the two-phase training strategy to effectively mitigate out-of-distribution problems by including the model's self-generated actions in the second training phase. This solution is simple yet effective and also maintains theoretical rigor.
This paper presents a supervised training framework for transformers applied to sequential decision-making problems—a subset of reinforcement learning (RL) tasks that lack an explicit transition matrix. In this framework, transformers are pre-trained on the target task class by simulating multiple environments to generate trajectories consisting of states, actions, and rewards. However, a discrepancy exists between action selection during pre-training and at test time, which leads to a gap between the empirical loss observed in training and the expected loss during testing. This gap can cause out-of-distribution issues, as studied by previous literature, causing the regret bound of the model’s performance to increase exponentially over time. To address this, the authors propose a mixed pre-training algorithm, in which the models learn not only from simulated data but also from trajectories generated by their own actions. This hybrid approach reduces the discrepancy between training and testing distributions. The proposed algorithm is theoretically validated and supported by experimental results.
1. I think the setup in Section 2 goes on a bit too long and could be more concise. A detail: the way f is used in Equation (4) isn’t fully explained until after Equation (5), which makes it hard to follow. Plus, describing f as "a prespecified decision function used to generate the data" is too vague.

2. The framework leans heavily on simulated environments ($gamma_i$ and $f$) for pre-training, which might limit scalability to real-world cases. Generating high-quality simulations can be both costly and challenging.
Yes
No ethics review needed.
- What is the difference between the setting described in the paper (reinforcement learning/sequential decision-making without transition probabilities) and the stochastic bandit or contextual bandit settings? If it aligns with an existing setting, that should be clarified. The use of reinforcement learning terminology here seems to add more confusion than clarity.

- What are the disadvantages of using pre-trained transformers compared to structured algorithms like UCB and Thompson Sampling? One notable concern could be the computational and memory cost; can this be quantified? For instance, how would the algorithms compare if the x-axis were aligned with computational metrics? What is the computational or memory overhead compared to simpler methods? Given that the proposed method is intended to be practical but lacks theoretical rigour, it is important to assess whether using these algorithms makes sense in practice.

- If UCB/TS had access to the same amount of pretraining data, how would their performance compare to the proposed method?
- The paper introduces notation and a framework for studying pre-trained transformers in decision-making tasks.
- It proposes a new algorithm and demonstrates its practical performance.
- The connection to performative prediction is intriguing.
This paper investigates the training and generalization of pre-trained transformers for sequential decision-making tasks without transition probabilities. The authors propose an algorithm that incorporates transformer-generated action sequences during pretraining, establishes a connection to performative prediction, and addresses the challenge of limited exploration. They highlight three advantages of pre-trained transformers over structured algorithms like UCB and TS: improved utilization of pretraining data, robustness to model misspecification, and enhanced short-term performance.
- My biggest concern is the lack of discussion regarding the disadvantages of using transformers for tasks like bandits. The paper does not sufficiently address the computational and memory overhead involved.

- The comparison between the proposed method and UCB/TS seems unfair, as the transformer benefits from extensive prior data while UCB/TS is evaluated in a cold-start setting. Claiming that UCB/TS is less effective at utilizing prior knowledge is misleading since these algorithms aren’t provided with any pre-training data. In fact, UCB/TS algorithms do incorporate observed data as part of their operation, in the form of empirical mean rewards and reward uncertainty. 

- The claimed short-term advantage of the proposed method also seems unconvincing. While it’s true that UCB/TS in their original forms is designed for asymptotic optimality rather than short-term performance, there are likely modifications or more state-of-the-art structured algorithms that can improve short-term performance. For instance, restricting the variance (exploration) in these algorithms could make them more greedy. I’m not convinced that this is an inherent limitation of structured algorithms.

- I am not very familiar with related work on applying deep learning or transformer methods to bandits, so a more comprehensive related work section would be helpful. For example, some studies use generated data as priors for algorithms like TS.

- The clarity and organization of the paper could be improved. For instance, while the paper emphasizes reinforcement learning and sequential decision-making settings, I believe the studied problem is more aligned with the contextual bandit setting. I believe some of the notations on pages 3-5 could be streamlined, and a discussion in the main text about how these settings differ and why the setting studied in the paper is more general would be beneficial.
Yes
No ethics review needed.
1.	Do the benchmarks used need pre-training? If not, then what would the authors argue is the benefit of the pre-trained method given the increased computation that is required for the pre-training?
2.     In line 167, it is stated that "...reinforcement learning algorithms and are usually hard to combine with prior knowledge such as pre-training data." However, I believe there are many options to combine prior knowledge in previously collected data: initializing replay buffer, offline learning, transfer learning, ....
This paper offers an analytical analysis of using pre-trained transformers on bandit-style problems and demonstrates the benefits exceeding the typical asymtotic optimal regret. Along the way, the paper provides interesting insights, and the resulting algorithm is well-motivated and explained clearly.
The authors explore the implications of using pre-trained transformers on bandit-style decision-making problems. They perform a scientific analysis that yields interesting insights, and provide empirical case studies that show that the use of pre-trained transformers yields competitive results with existing methods.
The main weakness can be the narrow scope of impact as the sequential decision making problem in the paper is very narrowly defined. 

The text does very little to motivate the need for the Learned Decision Function (LDF) described in Section 4. In particular, the relationship between the LDF and the pre-trained methods described in this paper is unclear on a first read. The authors need to better explain why the LDF is necessary. It seems unlikely that there is not an acceptable baseline that can be used from the literature. The confusion is further amplified in the experiments section, where the authors use different baselines to evaluate the pre-trained methods.

Additionally, while the insights generated are interesting, the paper lacks a clear and compelling rationale for why the pre-training is worth the extra computation compared to simpler methods. In particular, the use of a complex transformer architecture on a bandit-style problem would, on the surface, appear to be ‘overkill’, and the paper is not able to clearly articulate a compelling reason that justifies the use of such a complex solution method.

Ovearll, math notation is confusing and seems even incorrect/inconsistent. 

Some minor comments:
•	Figures 1 and 2 are not referenced anywhere in the main text.
Yes
No ethics review needed.
See Weakness.
By showing that transformers can outperform traditional algorithms in some sequential decision-making contexts, this work contributes to advancing the understanding of transformers.
This paper looks at how supervised pre-trained transformers can be used in sequential decision-making tasks, which is a subfield of reinforcement learning that doesn't involve transition probability matrices. The authors believe that transformers can perform well in tasks like bandits, dynamic pricing, and newsvendor problems if you leverage closed-form or easily computable optimal actions during the pre-training phase. They look at an issue that arises when there are differences between the training and test phases, and suggest an algorithm that lets transformer-generated actions be part of the training, which helps with generalisation. The authors also give some theoretical insights by looking at the transformer model as a near-Bayesian optimal predictor and suggest some numerical experiments comparing its performance against structured algorithms like UCB and Thompson sampling, which show where transformers work well.
**Confusing claims and demonstration:**

- The notations abuse. For example, decision function $f$ generate $a$ or $a^*$. In Eq. 5, the authors claim to generate $a^*$  but there is no $a^*$ in Eq. 5. I can only infer that $f$ generate $a^*$. And then below Eq. 7 it becomes $a_{\tau}$.

- In the paragraph, below Eq. 9, the authors claim that OOD in this setting is because "when generating the training data, there is no way we can know the final parameter. This is not an OOD problem to me. Training on some tasks and testing on other tasks, or changing test-time environmental parameters are more likely to be OOD problem.

- Above Eq. 9, pertraining -> pretraining, and this paragraph is unclearly written. Overly referring to the equation make me hard to follow the idea and intuition.

   

**Lack of Novelty** 

- The theory is comprehensively covered, while simply adding a mixed training phase and testing only on simple experimental setup doesn't seem novel enough to me, given the existing works [1-3]. 

- In addition, given that Transformer is a powerful sequence model, isn't studying pretraining on state-free bandit problems less meaningful? The authors claim that bandit is a special case, but it's to me more like out-of-scope because sequential problems are more related.

  

Since the motivations and clarity are less satisfying, I didn't look into the proof in details. While it's good to see that mixing online and offline data can preserve a theoretical guarantee.



[1] Lee, Jonathan, et al. "Supervised pretraining can learn in-context reinforcement learning." *Advances in Neural Information Processing Systems* 36 (2024).

[2] Laskin, Michael, et al. "In-context reinforcement learning with algorithm distillation." *arXiv preprint arXiv:2210.14215* (2022).

[3] Sodhani, Shagun, Amy Zhang, and Joelle Pineau. "Multi-task reinforcement learning with context-based representations." *International Conference on Machine Learning*. PMLR, 2021.
@inproceedings{
anonymous2024understanding,
title={Understanding the Training and Generalization of Pretrained Transformer for Sequential Decision Making},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=CiiLchbRe3},
note={under review}
}
In this paper, we consider the supervised pre-trained transformer for a class of sequential decision-making problems. The class of considered problems is a subset of the general formulation of reinforcement learning in that there is no transition probability matrix; though seemingly restrictive, the subset class of problems covers bandits, dynamic pricing, and newsvendor problems as special cases. Such a structure enables the use of optimal actions/decisions in the pre-training phase, and the usage also provides new insights for the training and generalization of the pre-trained transformer. We first note the training of the transformer model can be viewed as a performative prediction problem, and the existing methods and theories largely ignore or cannot resolve an out-of-distribution issue. We propose a natural solution that includes the transformer-generated action sequences in the training procedure, and it enjoys better properties both numerically and theoretically. The availability of the optimal actions in the considered tasks also allows us to analyze the properties of the pre-trained transformer as an algorithm and explains why it may lack exploration and how this can be automatically resolved. Numerically, we categorize the advantages of pre-trained transformers over the structured algorithms such as UCB and Thompson sampling into three cases: (i) it better utilizes the prior knowledge in the pre-training data; (ii) it can elegantly handle the misspecification issue suffered by the structured algorithms; (iii) for short time horizon such as $T\le50$, it behaves more greedy and enjoys much better regret than the structured algorithms designed for asymptotic optimality.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
pretrained transformer in-context learning bandits dynamic pricing sequential decision making
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/ba073469b6a1eee183ce7f5c868adb0eeff29211.pdf
generative models
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/10d6cf76cc86da016e5acafe79640e6041109561.zip
Understanding the Training and Generalization of Pretrained Transformer for Sequential Decision Making
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission