Understanding the Training and Generalization of Pretrained Transformer for Sequential Decision Making

Authors: Hanzhao Wang, Yu Pan, Fupeng Sun, Shang Liu, KALYAN TEJA TALLURI, Guanting Chen, Xiaocheng Li
Primary area: generative models
Keywords: pretrained transformer, in-context learning, bandits, dynamic pricing, sequential decision making

Abstract:
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.

Decision: reject
Reviewer scores: 6, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1730137040335

thread_type:
official_review

<thread_object>
creator: reviewer_hvzs

<content>
summary:
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.

strengths:
By showing that transformers can outperform traditional algorithms in some sequential decision-making contexts, this work contributes to advancing the understanding of transformers.

weaknesses:
**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.

questions:
See Weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730580192343

thread_type:
official_review

<thread_object>
creator: reviewer_uvas

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
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, ....
</content>
</thread_object>

</thread>
<thread>
creation date:
1730754552153

thread_type:
official_review

<thread_object>
creator: reviewer_svgy

<content>
summary:
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.

strengths:
- 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.

weaknesses:
- 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.

questions:
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730822810449

thread_type:
official_review

<thread_object>
creator: reviewer_8xdh

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732659972707

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary to AC

comment:
We thank the reviewers for their time. To save the AC’s time, below we provide a brief summary of the discussion period.

# Our contribution

We developed a uniform pre-training approach that is applicable to a wide variety of decision-making problems. We provide theoretical and empirical results for stable training and effective performance. We believe this is foundational work that provides a comprehensive analysis for applying Transformers to decision-making problems.

To be more specific.
- **Unified Approach.** This approach can solve a wide variety of decision-making problems (using the same pre-training approach for bandit, pricing, and inventory problems) under different environments (with different types of reward functions, such as linear, contextual, and stochastic). By sharing a uniform pre-training structure, this method eliminates the need for users to design instance-specific (e.g., bandit vs. pricing) or environment-specific (e.g., stochastic bandit vs. linear bandit) algorithms, and provides general results for theoretical regrets and similar empirical performance.

- **Power of Transformers to Adapt.** The inherent power of Transformers enables the pretraining of decision-making agents in simulated environments. The Transformer architecture is capable of generalizing and extracting specific patterns from the trajectories it has encountered. Consequently, as long as the simulator is sufficiently large to encompass trajectories similar to those of the current environment, the Transformer can effectively identify the environment and adapt to the corresponding optimal algorithms. This approach introduces a new in-context algorithm selection mechanism that is both powerful and flexible in practice.


# Summary of the reviewers comments

In summary, almost all reviewers (except 8XDh) have had misunderstandings regarding the foundational setup. After the first round of discussions, the main issues remain focused on the contribution of our work to the general decision-making community.

The reviewers' concerns mainly focus on the limitations of the scope of this work, as it primarily focus on simpler environments (bandit, pricing, inventory problems) without random transition dynamics, which is common in the general RL setting.

# Our response regarding the reviewers major concerns

- **Merit of (currently missing) fundamental work.** First, we believe that the effort to focus on a simpler environment should not be downplayed, given that this area lacks a solid foundation. Indeed, every foundational RL or sequential decision-making algorithm needs to be based on foundations for simple environments like bandit problems. For example, consider the advancement from UCB in bandit problems to UCBVI in RL—UCBVI cannot be developed from nowhere without the insights gained from bandit problems. 

- **Comprehensive results.** Second, we believe we have shown enough new results to contribute to this community. To the best of our knowledge, in this simple bandit setting, we have not seen theoretical developments comparable to our setting. Our paper contains 40 pages of work with both theoretical and empirical developments, and we think it is reasonable to defer the investigation of RL to future work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.

The main concern is on the claimed advantage over standard baselines such as UCB or Thompson sampling: (1) UCB and Thompson sampling can also leverage prior information, and (2) no formal theoretical separation between the proposed approach and UCB/TS. Therefore, the AC recommends rejection.