Contextual Thompson Sampling via Generation of Missing Data
Authors: Kelly W. Zhang
Tiffany Cai
Hongseok Namkoong
Daniel Russo
Venue: NeurIPS 2025 poster ()
Keywords: contextual bandits
Thompson sampling
sequence modeling
Bayesian modeling
uncertainty quantification
decision-making
Primary area: theory
TL;DR: A generative model imputes missing outcomes to drive Thompson Sampling decisions, yielding a flexible algorithm with regret tied to offline prediction quality.

Abstract: We introduce a framework for Thompson sampling (TS) contextual bandit algorithms, in which the algorithm's ability to quantify uncertainty and make decisions depends on the quality of a generative model that is learned offline. Instead of viewing uncertainty in the environment as arising from unobservable latent parameters, our algorithm treats uncertainty as stemming from missing, but potentially observable outcomes (including both future and counterfactual outcomes). If these outcomes were all observed, one could simply make decisions using an "oracle" policy fit on the complete dataset. Inspired by this conceptualization, at each decision-time, our algorithm uses a generative model to probabilistically impute missing outcomes, fits a policy using the imputed complete dataset, and uses that policy to select the next action. We formally show that this algorithm is a generative formulation of TS and establish a state-of-the-art regret bound. Notably, our regret bound depends on the generative model only through the quality of its offline prediction loss, and applies to any method of fitting the "oracle" policy.

Contextual Thompson Sampling via Generation of Missing Data
Kelly W. Zhang
Tiffany (Tianhui) Cai Department of Statistics Columbia University
Department of Mathematics Imperial College London
kelly. zhang@imperial.......mmmm.. ac.uk
ti tiffany. cai@collllloccccooo cnboalmbi.... edu
Hongseok Namkoong
Daniel Russo
Decision, Risk, and Operations Columbia Business School namkoong@gsb, columbia edu
Decision, Risk, and Operations Columbia Business School
djx21748columbia edu
Abstract
We introduce a framework for Thompson sampling (TS) contextual bandit algo- rithms, in which the algorithm's s ability quantify uncertainty and make decisions depends on the quality viewing uncertainty in the environment as arising from unobservable latent param- eters, our algorithm treats uncertainty as stemming from missing, but potentially observable outcomes (including both future and counterfactual outcomes). If these outcomes were all observed, one could simply make decisions using an "oracle" policy fit the complete dataset. Inspired by this conceptualization, decision-timer time,, our algorithm uses a generative model to probabilistically impute missing outcomes, fits policy using the imputed complete dataset, and uses that policy to select the next action. We formally show that this algorithm is generative formulation of 'T and establish state-of-the-art regret bound. Notably, regret bound depends on the generative model only through the quality of its offline prediction loss, and applies to any method of fitting the "oracle" policy.
generative model that is learned offline. Instead of
each
Introduction
Recent advances in machine learning have transformed our ability to develop high quality predictive and generative models for complex data. This work introduces framework for developing decision- making algorithms, specifically for contextual bandit problems, that can take advantage of these machine learning advances. By design, we assume the algorithm developer is able to apply these techniques (e.g., minimize a loss via gradient descent) and employ these methods as subroutines in our decision-making algorithm. Moreover, our theory formally connects the quality of effective (self-) supervidd learning via loss minimization to the quality of decision-making.
Classically, Thompson sampling (TS) algorithms form a parametric model of the environment and consider the decision-make maker's uncertainty arising from unknown latent parameters of that model [Thompson. 1933, Russo al., 2020]. The primitive operations used by TS include i) specifying an informative prior for the latent parameter using domain knowledge ii) sampling from the posterior distribution of the latent parameter, and iii) updating the posterior distribution as more data is collected. Unfortunately. it is well known that all three of these operations are non-trivial to perform with neural networks [Tran et al., 2020, Goan and Fookes, 2020]. this work, we view missing, but potentially observable, counterfactual outcomes the source of the decision-makert uncertainty. This perspective allows us to replace the primitive operations required the classical view with new
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

ones that are more compatible with neural networks, namely the ability to i) effectively minimize an offline prediction loss, ii) autoregressively generate from a learned sequence model, and iii fit desired policy given access to complete dataset (outcomes from all actions and decision-times).
In the missing data view of uncertainty, if we had a complete dataset, there is no uncertainty because we could simply use the entire dataset to fit a desired "oracle" policy to use to make optimal decisions for that task. Inspired by this idea, each decision time our algorithm imputes missing outcomes using a pretrained generative model, fits a desired policy using the imputed complete dataset, and selects the best action according to the fitted policy. We show that this algorithm is a generative implementation of TS. We demonstrate empirically how to learn a generative model to impute missing outcomes using standard machine learning tools in meta reta-bandit settings, where the algorithm learns from data from previous tasks to perform well on new task from the same distribution.
We prove a state-of-the-art regret bound for generative TS with three key properties, which each have significant practical implications. First, the generative model used to impute missing outcomes only affects our bound through the offline prediction loss of the model. This means that our theory is applicable to any imputation model architecture, and that the quality of the generative model can be easily optimized for and evaluated via offline training and validation. Second, our bound is unique in that it applies to any procedure for fitting a desired "oracle" policy. This allows one to easily adapt TS to decision- -maiing problems with constraints, e.g., for fairness or balancing. Finally, our proof approach makes important improvements to previous information theoretic analyses, which may be broadly applicable: i) we accommodate infinite policy classes directly without discretization, and ii) our bound quantifies the benefit of prior task information, such as side information on the actions. Our results hold quite generally and do not require restrictions on generative model or policy class. We demonstrate practical implementation of our framework in Sections 4 and 6.
2 Problem formulation
Meta-contextual bandit problem. Let bandit tasks T be sampled from an unknown distribution p":
T ~ p". where T= [Z,, (x(()) ..... y(a)) PaEA.h:
where each bandit task r consists of prior task infor- mation Zy, action space A,, context vectors X1T
X
X2
X3
X4
XT
(X Xr and potential outcomes (X JacA, = JOEA, [Rubin, 2005]; see Figure for
20)
(1)
(1) (1)
Y (1)
x
(Y/(), ) depiction. We omit subscripting X, and y,(0) with to reduce clutter. Note, in contrast to the design-based based in- ference literature [Neyman, 1992], which conditions on the potential outcomes and treats them as non-random, we assume the potential outcomes T are drawn from a task distribution p". Informally, the agent's 's objective to se- lect actions to maximize the total expected reward for each encountered task. At the start of task, the agent observes prior task information Z, For each timestep t E [1: T]. the agent observes context
(2)
y12
(2)
V (2)
2
(3)
y(33
(3)
(3)
(3)
Vz
Figure 1: Potential outcomes table for a
task T.
X,, selects action A, E Ar, observes outcome Y: y((A)) and computes reward R(YY). for a fixed, known function R in [o, 1]. The history, H, = {Z,, AX, Yi). (X,-1. Ar-1, 1. Yr-1), 11),, X,). includes the current context Xt. In contrast to much of the Bayesian contextual bandit literature [Lattimore and Szepesvári, 2019, Russo et al., 2020], we do not make parametric assumptions about the distribution of outcomes Y conditional
contexts X and prior task information Z.
The agent is able to learn both online within a single task (i.e., over the T total decision times), as well as meta-learn across different lasks (e.g.. learning how task prior information Z, may inform the } J..A)) The algorithm has access to training data collected from previous tasks, sampled from (1). These previous bandit tasks can be used by the algorithm to meta-learn across tasks, i.e., learn about the distribution p" itself to improve decision-makin quality. Our algorithm's decision-making quality depends on how accurately the agent is able to model the task distribution, as well as the policy fitting procedure the algorithm designer chooses. Rather than relying on strong assumptions on the environment structure, we put the onus on the algorithm designer to i) learn a generative model that accurately captures the environment structure of the meta-bandit task at hand, and ii) choose meaningful method for fitting desired "oracle" policy, assuming access
distribution of tx (1))
2

---

to complete dataset. Since generative models learned offline routinely perform much better than expected according existing theory, our theory focuses on formal reductions of decision-making making quality to offline learning.
New articles released
Interact with User
Interact with User T
+
13 Observe context features
Observe context features
2) Action: recommend
2) Action: recommend
Click
Click Like
Obscrvation:
Lise
3) Observation:
Share
Foundation model reads them
Reward:
4) Reward:
Repeat another recommendation task lomorrow
Figure 2: News recommendation meta contextual bandit problem.
Motivating example: News recommendation. As depicted in Figure 2, motivating meta-contextual bandit problem is cold-start news recommendations. Each day, a new set of articles A, is released, which the agent recommends users who arrive throughout the day. In contrast Li al. [2010], our algorithm meta-learns across news recommendation tasks and uses the article text to improve cold- start decisions. We use Z, (ZZ"))uE.A to denote the task-specific prior information; for example, for article a A, - could be the news article text or other article meta-data (category, style, etc.). The context variables X, consist of user-specific features, and Y, are recommendation outcomes observed following the t decision. The modern challenge in this setting is that incorporating news article text Z, can greatly improve the recommendation system's decisions, but a foundation model is needed to process this high dimensional text and inform decision-makes maaiiinn. This motivates us to i) make very minimal structural assumptions the relationship between prior information Z,, context features X,, and outcomes Y,, and ii) develop an algorithm that can leverage foundation models.
Policy fitting. The algorithm designer specifies procedure for fitting a desired "oracle" pol- icy given access to a complete bandit task dataset т. This fitting procedure outputs policies in a function class II where each
II defines mapping from contexts X, to action E A, that does not vary over time. For notational simplicity, the policies in II are assumed to be non-stochastic. Note that do not require that this policy class is "correct". For particu- lar task т, we use x² (-: to denote a "best- "best-fitting" policy X² E II, where the fitting criterion is defined by the algorithm designer. For example, consider simple least squares criterion: argminzem + E/= E (T(()(((X)))) - maxae R)
One should think of (-:)) the policy one would implement if abundant task data, т, were available. This could involve fitting a model, adding prompt tokens to condition a language model, or maximizing hindsight performance. This policy fitting can also incorporate constraints the policy, e.g., to ensure fairness. We aim to match this policy's performance via efficient interactive learning.
Regret. We consider best-in-class style regret objective, which is common in the contextual bandit literature [Foster et al., 2020, 2019, Langford and Zhang, 2007, Agarwal et al., 2017]. The objective of the agent A is to make decisions minimize the per-period regret against the best-in- hiindsht policy n (c:T):
(AA) i
Σ (T (Xx77)) -R(y())))) R(x(4d)))) }
(2)
The "best-fitting" or bbet-in-hinnsiggh policy x <(:TT) T) is well-defined and is well-established concept in the bandit literature, representing generalization of the optimal policy. Refer, for example, to Section 2 of Beygelzimer et al. [201]] I] and Chapter 4 of Bubeck et al. [2012] for very analogous objectives. We emphasize that our algorithm does not have access to the best-fitting policy, means to compute the best-fitting policy if
which would trivialize the problem. They have access they had v, i.e., observed rewards of every arm every context.
The expectation in (2) averages over tasks 7 ~ p" and any randomness in how the algorithm selects actions. A(A) the long-run per-period regret if the algorithm was deployed across many tasks. Note, increasing the complexity of the policy class II increases the average reward under the best- fitting policy, E [ Σ R(Y((T (Xx:)))) However, this increased complexity also means that large sample sizes are required to learn (::7) accurately and will worsen our regret bound (Section 3.2).
3

---

3 Generative Thompson Sampling: General algorithm and regret bounds
Posterior sampling via imputing missing data. In this work, we view missing data as the source of the decision-maker's uncertainty. This contrasts the classical approach of considering unknown model parameters as the source of uncertainty. As we will explore in the following sections, the missing data viewpoint is very amenable to modern deep learning methods, which can be used to train models that are able to impute missing data probabilistically in calibrated fashion. First, consider an idealized setting in which we have the true meta task distribution p². Using p" we can form exact posteriors sample for task outcomes
(Zr, XI.T (119) }} given the history Hz: ft p" (T 6 6- H.)).
(3)
Above, we probabilistically generate values in T that have not yet been observed in the history Hr This consists of future contexts, future outcomes, and outcomes from previous timesteps for actions that were not selected. We discuss how to practically implement such sampling in Section 4. Note, even when p* is known, fg is simply calibrated posterior sample and is not equivalent to the true
With this exact posterior sample, it. we can form posterior samples of any statistic computed using 4. In particular, we are interested in sampling from the posterior distribution of the fitted policy x x'(.:7) which can be computed by finding the fitted policy for the sampled task dataset fr, i.e., x² (-; ft). Pos- terior sampling of best-fitting policy is common subroutine used in Bayesian decision-male making algo- rithms [Kaufmann et al., 2012, Russo and Van Roy, 2018, Ryzhov et al., 2012]. Thus, our posterior sampling approach can easily integrate with these existing Bayesian algorithms.
Observed imputed outcomes
X
X2
X3
X
x2
933
9:3
nt Fitted policy
Z7
x
9
p
Figure 3: The agent imputes missing outcomes and uses the imputed dataset to fit policy.
In this work, we focus on Thompson sampling [Russo and Van Roy, 2016, Thompson, 1933], i.e., probability matching, which selects actions according to the posterior probability that they are optimal. Thompson Sampling (TS) can be implemented with single posterior sample per decision time. In our generative implementation
TS (Algorithm 1) at decision time i, after sampling f; as in (3), TS fits the policy x² (.; ft). and selects the action A, + T² (Xt; ft). See Figure 3 for a depiction. Our algorithm generalizes TS by replacing the true wwraadma-miizinnggg policy with a best--itting policy under given policy class II. A more "standard" TS algorithm is recovered when the best-fitting policy is correctly specified. See the discussion below display (2) for more on the best- best-fitting policy.
Algorithm 1 Generative Thompson Sampling
Require: Imputation model p. actions A,, task input Zr.
1: for / E (I....,T) . do
2: Observe context X, and append it to H.
( 3: Generate / sample Ft ~ p(T E H.))
4: Fit the policy x² (-:
5: Select the action A, + x ft)
6: Observe outcome Y, < (Ar)
7: Update history H4+1 + He U f(X, YY}
8:
end for
Under our generative TS Algorithm 1, the polices in II that are best-in-class optimal under some likely generation of 4t have a chance of being selected. Once no plausible sample of missing outcome ft could result in an action being optimal, it is essentially written off. We formalize that our generative algorithm aligns with the abstract definition of Thompson Sampling (probability matching) in Proposition below when using the correct model p". See Chapter 36.5 of Lattimore and Szepesvári [2020] for further discussion of the probability matching definition of Thompson Sampling.
Proposition 1 (Algorithm Implements Thompson Sampling). Algorithm I with imputation model p" implements Thompson Sampling (probability matching), i.e., the following holds almost surely:
PPA H.) P( ((m'((Xii) H.).

---

A key to proving Proposition 1 is showing that which holds when iz is sampled from the true meta task distribution p" as in (3). See Appendix A.2 for our proof.
= a H4,) = =X; = Ht).
3.1 Regret when using a perfectly calibrated imputation model p'.
We develop a novel analysis of contextual TS, which is applicable infinite policy classes II with finite VC dimension. Our VC dimension bound resembles those from adversarial bandits, but for the first time, we show we can derive this using an information theoretic analysis. We first present regret bound for Algorithm with a perfectly calibrated imputation model, p" from (1), and extend to approximate imputation models in Section 3.2. Note that assuming p" is known is akin to assuming the prior and likelihood of Bayesian model are known, which is standard in Bayesian regret analyses.
Notation, Let x² (XIr) == ( T) be the best fitting policy evaluated at contexts X XIr. Let H(Y X) denote the conditional entropy of Y (discrete) given X; note H(Y X) = - -E P(Y = X) log F(Y
Xdyy is a constant. Let I(Y: X Z) be the mutual information between Y and X conditional on Z: note :Y: Z) marginalizes Z and is also a constant.
Theorem (Regret bound for Generative TS with perfectly calibrated imputation model p*). For Algorithm I with imputation model p", ATT--Gen(p" (p").
Arl 2T
A A(ATS-GeulP") S
H(T (X1:)) Z+).
Moreover, (AATS-Gen (p*)) S V 11 (XI:T) Zr), where I' bounds the information ratio [Russo and Van Roy, 2016], Le., I max; Fe a.s. for re := (Xxe))"(())))))-""""")")"))))))) -R R( y
1(x (X (A/)H)))
Note I can be smaller than |/, when feedback from action informs learning about other actions (Appendix A.6). The entropy. H (Xir)) Zr), quantifies the benefit of using prior information Z. Our bound automatically applies to infinite policy classes since it only depends on the entropy of the optimal policy evaluated at a finite number of contexts, π (X1:))-
Upper bounding the condition entropy using VC dimension. We can construct coarse upper bound for the entropy H (X1.)) 2r) using the VC dimension of the policy class 11. The VC dimension is worst-case quantity that has to with the total number of possible assignments of actions given contexts. In contrast, entropy reflects uncertainty based on the task distribution (learned from past tasks) and the information Z (e.g., article texts), many assignments may be extremely unlikely to be optimal. Since VC dimension is only defined for binary functions, we use the multiclass generalization Nataranjan dimension [Natarajan, 1989] when |Arl > 2.
Proposition 2 (Complexity bound on entropy). For policy class 11 over action space A, with Nataranjan dimension d (equivalent to VC dimension when [A,]= [ol=30 2).
H(T" (Xi:r) Z+) S H(TT = O(d log(T' (A,])).
Note, our bound above depends on the Natarajan dimension of the policy class II, not the Natarajan dimension of the generative sequence model p". Furthermore, the Natarajan dimension of II does not change with T for stationary policies. feature our result is that our bound, when combined with Theorem 1, can be used to derive regret bounds for wide range of policy classes II.
Using Proposition 2, our regret bound (Theorem 1) resembles adversarial regret bounds that depend on VC dimension [Beygelzimer et al., 2011], showing for the first time how such a result can be established through information thcoretic arguments.
Benefits of our approach and relationship to related work. Regret bounds for contextual TS bandits with infinite policy classes have been of great interest in the literature. The predominant approach to generalizing
analyses for TS beyond multi-armed bandits requires discretizing a latent parameter space [Dong and Van Roy, 2018, Gouverneur et al., 2024, Neu et al., 2022, Min and Russo, 2023] and uses cover-number arguments; our proof approach notably does not require any discretization. Furthermore, our bound can be applied broadly, while existing approaches like Neu et al. [2022], Min and Russo [2023] depends on the entropy of latent environment parameter, which only applicable to parametric bandits. By Proposition 2, our result can directly be applied to infinite policy classes by leveraging existing VC dimension bounds, e.g., for decision
infommaiion-thoreee
5

---

trees [Asian et al., 2009]. In parametric, stationary bandit settings, our result approximately matches (up to log factors) existing Bayesian regret bounds for linear logistic bandits [Neu et al., 2022] and matches up to a factor of vd and log factors bounds for linear non-contextual bandits [Russo and Van Roy, 2018, Dong and Van Roy, 2018 8] (Appendix A.6). Finally, though we do not explore it much in this work, since we make minimal assumptions on p', Theorem I applies to nonstationary bandit environments. While the oracle policy n² cannot be time-varying, x² can effectively vary over time by including the timestep / as a context feature in Xt.
3.2 Regret when using an approximate imputation model po-
We now present a regret bound for generative TS with an approximate generative model po- The result is notable because py only affects the regret bound through its offline prediction loss, which means the result can be applied to any model class. Specifically, our regret bound will depend on the following population-level loss (the expectation below averages over the task distribution y*):
f(pe) = -
-E log pe (XurT,
(x (1 k.)].
In Section 4, we discuss training and sampling from learned generative imputation models in practice. Theorem 2 (Regret bound for Generative TS with an approximate imputation model). For Algorithm I with imputation model pa. AAT--Gen (pe).
Ar 2T
A (PTT))
((TT(X
| Z,) + 22(((p))
(5)
Penalty for sub-optimal prediction
Rogret bound for Thompson sampling
What is particularly novel about Theorem 2 is that the analysis holds even when the imputation model Po misspecified and does not correspond to proper Bayesian inference in any way. Comparing Theorem 2 to Theorem from earlier, we can interpret the "cost" of using an approximate model po as v (((p)) ((p*), )]: This penalty depends on how well pe approximates p".
Scaling of loss penalty. While tight theoretical bounds for the penalty term E(po) {(p)) currently do not exist for complex models like neural networks, we can draw intuition from simpler settings. Consider a stationary, stochastic, Bayesian bandit problem. In this setting, for parametric Bayesian models, where pe and p' are exchangeable, posterior predictive distributions [Fortini and Petrone, 2023], classic results by Clarke and Barron [1990] show that the gap {(p)) L(p') scales like log T, under mild regularity conditions. This sublinear growth occurs because in this stationary setting, the Bayesian model po is better able to approximate the next outcome as it observes more data (a phenomenon closely related to Bayesian consistency [Kleijn and Van der Vaart, 2012] and how the effect of the prior eventually washes out). The difference {(pa) - E(p") also scales with the amount of data used to learn pe: This is closely linked to empirical Bayes methods, i.e., approaches to meta-learn a prior distribution from data. When pe and p' correspond to posterior predictive distributions of Bayesian models with correctly specified likelihoods, pa and p" differ only in their initial prior distributions. Existing works bounding the regret of TS with misspecified priors are not directly comparable, as Simchowitz et al. [2021] analyzes modified version of TS that requires multiple posterior samples per decision time, and Liu and Li [2016] bounds the frequentist regret.
Related work on generative TS algorithms. Wen et al. [2021] consider a non-contextual, multi- armed TS algorithm that incorporates a generative outcome model. However, they require modeling latent environment parameters, and their bound requires a histoey-dependent KL. divergence term to be small, which differs from our prediction loss penalty. Cai et al. [2024] proves a regret bound with similar prediction loss penalty generative TS algorithm with misspecified models for much simpler multi-armed non-contextual setting. They do not introduce the concept of a general "oracle" policy fitting procedure, and their result does not apply infinite policy classes. Moreover, not able to directly build on their proof approach because they critically rely on the fact that under p". unobserved outcomes Y are exchangcable given the history. In contrast, our result does not require exchangeability at all and technically applies even if p" is not exchangeable (e.g., nonstationary).
were
Flexibility and advantages of Generative TS. Generative TS requires the algorithm designer to choose an imputation model pe and a policy class IT. The modularity of these two components allows one to easily extend TS to more complex, less standard decision-making problems, e.g., (i) Nonstationarity can be accommodated with Po that models trends over time (see discussion before Section 3.2); (ii) Correlated outcomes can be modeled using a pe that captures dependencies between

---

outcomes across actions or over time; (iii) Constrained decision-makinn can be done by choosing policy class II satisfying such constraints, e-g., to ensure fairness one can use standard constrained optimization approaches to learning decision rules [Corbett-Davies et al., 2017] (Appendix B.9).
Practically implementing generative Thompson Sampling
We now introduce an example of how to learn Pe and implement generative TS. Our overall frame- work is depicted Figure 4: step I, we use of- fline data from previous tasks to learn a pe model; Then in step 2, we use the learned po model to implement generative TS. Here, Po sequence model that we meta-learn by pretraining on his- torical data from previous tasks. As theory accommodates any architecture, our approach can take advantage recent advances in genera- tive sequence models [Vaswani et al., 2017].
Gather offline data
Deploy policy online
Train sequence model offiine 'Ten inpur online policy
Task
Task
Learn
Task
Offline Date
Model py Policy x(w)
Task
Ze, {Y 14
Task
Task
Task
Task
Figure 4: Offline meta-learning and online
decision-making across multiple tasks.
4.1 Step 1: Offiine learning for generative model ps.
We now describe learning a generative, sequence model po from historical data. Our goal is to minimize the loss ((p)) from (4). First note that by rules conditional probabilities. E(pe)
(X. XI--1, (X(()) (Y1=-IacA.)) + log po ((V. JaEA- (r,, JaEAA X (Yia (Yi=-IheA)))
E
To make learning po more practical, the model make variety of simplifying approxima- tions. For example, po could model contexts as evolving independently of past outcomes, i.e., Po(X, Zr, XI--1, (Y ( JaE.,) (X( Z, X1.-1), model contexts as i.id. over time, ie., (X(, ZTaX1= 11 (1EA)) (X,).. Additionally, Pe could model outcomes in- dependently across actions, i.e., (o JaEA, Xia (1142 = IIEEA, (o((, |
z(a) X1. 1 where Z, = )
JaEA, for action-specific task features Z(a).
Under the chosen simplifying modeling approximations, to minimize an empirical loss. For example, in our experiments, our Pe makes several simplifying assumptions, and we minimize the following empirical loss to approximately minimize E(po):
use gradient descent to optimize pe
Σ
Σ
log (X.) + Σ log po ((Y lEEA 2 X1a-1, (lo ) (6)
Dollinee
TEDilline t=m
aEA.
Above, Doflline ideally consists of bandit tasks p¹ as described in (1). In practice, may not have "complete" task datasets t XIT, (K(a) Y JoEA, J. but instead have some partial datasets, e.g., {Zr, (X1, A₁, Yı),..., (Xr, Ar, Yr),. collected by behavior policy. In our experiments. we use several heuristics to construct approximate complete tasks 7 from the partial datasets. We use these approximate task datasets form Datliiee
{ 22, ......, To form 7, we make a simplifying modeling assumption that the tuples (X1, Y )..... (X,, (x1,Y0 r are exchangeable over time. We then bootstrap sampling form approximate complete task datasets F; see Appendix B.2.2. In this appendix, we also formalize all the simplifying modeling assumptions make and show how they match standard stochastic contextual bandits with independent actions.
Algorithm 2 Offline training of a sequence model
Require: Training data Dollliee model class {po
1: while not converged do
2: Sample mini-batch of tasks Dmini-batch C Doffline
3: Compute loss in (6) using tasks E DDmini-batch
4: Backpropagate and take gradient step to update pa
5: end while

---

Algorithm 3 Posterior sampling via autoregressive generation
Require: Sequence model pe. actions Ar, current timestep L, current task history Hr.
1: For each E A,, define MMa) as the set of times E [1: '] where Y:(") not observed in H.
2: For each E A,, define the ordering <a so that all observed outcomes precede unobserved ones
3: Set X1 +- X1:t and sample XI+1,..., Xy from po
4: for AA, do
5: for E (1, (1....,T) in order of <a do
6:
if & M(a) then
Y(() f Y,
7:
8: else
9: Sample p((0) pel. 1 {X,, (/000 D
10: end if
11: end for
12:
end for
Return: +r 12, Cosr:(1ril Jar
(F(e))
4.2
Step 2: Online deeision-making using the learned generative model po-
After the sequence model is trained offline, it is deployed and used for online eecision-making No additional training Pe is needed. Instead, the sequence model learns from recent online observations "in-context" by conditioning [Brown et al., 2020]. Specifically, to implement the generative step of Generative TS (line 3 of Algorithm 1), we use pa to sample future contexts X and missing outcomes Y to form fr. We refer to this procedure as posterior sampling via autoregressive generation; this is depicted in Figure and formalized Algorithm below.
In Algorithm 3, we use M, C [I,...,T) T} to denote the timesteps for which y:(') has not been observed. When generating outcomes in for a, permute pairs of contexts and outcomes (X, Y) so that observed outcomes always precede missing ones; this way. we always condition on all observed outcomes (and corresponding contexts), matching Figure 5. We use 2a to denote this ordering for an action a E Ay; we use -a j whenever either (a) < j or (b) i Ma but j E Ma.
Part 2: For each row with missing entries, autoregressively generate missing outcomes
Observed and missing outcomes
Xr
X1
X2
X
Step
Part Sample future contexts X
Sample from pacy [Zy, X P(a
y(
2x
p2
X
X,
Step 2
X
X
X
Yx22
02
12
Sample from Paly I2y, Lu
x
Figure 5: Posterior sampling via autoregressive generation (Algorithm 3).
5
Related work
Decision-making with generative models. Many recent methods use generative models in decision- making that involve imitation learning, i.c., from demonstrations learn to mimic an expert's actions [Chen al., 2021, Janner al., 2021, Hussein et al., 2017]. Lee et al. [2023] discuss how these approaches can be used even without access to expert demonstrations, as long as one is able to fit an approximate "oracle" policy from offline bandit environments. Our work differs significantly from Lee et al. [2023] and other imitation learning based works because our sequence models are used to sample future oulcomes, instead of predicting optimal actions. Several recent works also use generative models to model future rewards [Mukherjee et al., 2024, Nguyen and Grover, 2022, Müller et al., 2022x, Garnelo et al., 2018, Liu and Li, 2016]. Most previous work on decision- making with sequence models that predict future rewards does not use autoregressive generation

---

to quantify uncertainty [Mukherjee et al., 2024, Nguyen and Grover, 2022, Müller et al., 2022x, Garnelo et al., 2018]; Instead, their algorithms only consider uncertainty in the single next timestep's reward under each action, e.g., using softmax sampling [Mukherjee et al., 2024]. We find empirically that alternative (non- non-autoregressive) ways of sampling from the sequence model can lead to inferior decision-making performance (Figure 6).
(Approximate) TS with neural networks (NN). Implementing TS with NN has been longstanding challenge. Riquelme et al. [2018] investigated TS with a variety of Bayesian uncertainty quantification techniques for NN; they found that linear TS with the last layer of NN as context outperformed many more complex methods. While some TS algorithms directly model uncertainty in NN weights [Zhang et al., 2020, Wang and Zhou, 2020], the foremost approach in the literature implement TS with deep ensembles [Qin et al., 2022, Lu and Van Roy. 2017, Dwaracherla et al., 2020, Osband et al., 2023, Oshand and Van Roy, 2015, Oshand et al., 2023, Li et al., 2024]. Our generative TS algorithm is critically different from ensembling because a) through offline meta-training we are able to learn informed priors from complex task-specific information Z (like text) with benefits that are explicitly reflected in our bound, and b) our approach allows the generative model to learn in-context avoiding retraining online using gradient updates on sub-sampled data, which is sensitive to learning rates.
Meta-bandits, In the bandit literature, many algorithms have been proposed for meta-learning settings. Many prior works focus on different setup, where bandit tasks are encountered sequentially and leveraged for learning across tasks [Lazaric et al., 2013, Basu et al., 2021, Kveton et al., 2021, Wan et al., 2021, Moradipari et al., 2022]. In contrast, our approach uses in-context learning, where a single algorithm adapts to a variety of new task it could, it encounters (see Figure 4). Also, unlike much of the meta-bandit theory
cerrature-wiie focuses on simple models, e.g., linear [Cella et al., 2020, Cella and Pontil, 2021, Moradipari et al., 2022] or TS with parametric Bayesian priors, including mixture models [Wan et al., 2021, Kveton et al., 2021, Hong et al., 2022]-our method accommodates complex sequence models Pe with low loss and any policy class with finite VC dimension. A notable exception is Boutilier et al. [2020], which directly optimizes a non-contextual bandit policy from historical data via gradient descent, but their approach only works for learning differentiable, soft-max based soft-max based algorithms.
6 Experiments
Problem setting. Throughout, = 500, |A| = 10,' outcomes Y are binary, RRy) = y. and Z has separate components z(a) E I²² for each action. Our SYNTHETIC setting uses a Bayesian logistic regression data-generating process with contexts X E IR$. Our SEMI-SYNTHETIC setting mimics a cold-start, news recommendation setting using the MIcrosoft News Dataset [Wu et al., 2020]; Z consists of article headline text, contexts X E TR55 are user features, and Y E (0, 13 represents whether user click on recommendation. See Appendix B.1 for details.
Bandit algorithms. We use Generative TS (TS-Gen) as described in Section 4. For pe. we use a simple recurrent neural network which takes in prior information Z, history H and current context X, and outputs a distribution over Y. In the VEMIETC setting, Pe embeds the article headline Z using DistilBERT [Sanh et al., 2019]. We use logistic regression- based policy class II for the SYNTHETIC setting and a multi-layer perceptron (MLP) policy class for the SOMI-SYNTHETIC setting. For baselines, three algorithms use the same Pg model as TS-Gen, but select actions differently: 1) GREEDY deterministically selects the action predicted by pe to have the greatest next reward. 2) EPSILON-GREEDY employs GREEDY with probability 0.9 and otherwise selects an action uniformly at random. 3) -S-NEURAIL--LNEE the context for a linear TS algorithm with multivariate Gaussian prior; we consider variants with an uninformative prior and a prior fit using historical data. We also compare to a standard linear TS [Agrawal and Goyal, 2013], where X, is used the context, well as LinUCB [Li et al., 2010].
which uses the output of the last layer of the pe model as
Results. As seen in Figure 6, TS-Gen outperforms other algorithms in both the SYNTHETIC and REMI-SYNTHETIC settings. TS-Gen's 's superior performance compared to other algorithms that use the same pe model (GREEDY, EPSILON-GREES our generative approach to uncertainty quantification and ceaision-mmking... We conjecture TS-Gen 's advantage compared to LinUCB and Li-en is attributable to our pretraining procedure and the
NEUURAL-LNEAR validates the benefit of
'Recommendation options often from pre-filtered set [Davidson et al., 2010, Covington al., 2016).

---

Regret Over Timesteps: Synthetic
Regret Over Timesteps: Semisynnhetic
120
100
125
TS Gen
80
Linear
Regret
TS-Linear (fitted prior)
75
T5-Neural-Linear
40
50
TS Neural Lineer (fitted prior)
Greedy
20
25
Epsilon-Greedy
LinUCB
100
200
300
400
500
100
200
300
400
500
Dacision Times
Decision Times
Figure 6: Cumulative regret averaged over 500 bandit tasks. Regret is against the best fitting policy II (logistic for synthetic and ML.P-based for semisynthetic). TS-Ger outperforms methods that use the same pe model (Greedy, e-Greedy, IS-Neural-Linear), Error bars (barely visible) denote +1 s.e.
better use of prior information Z. We also found, suggested by Theorem 2, the lower the offline prediction loss of ps, the lower the regret of TS-Gen; see Appendix B.4.1.
Computational costs. For our semi-synthetic experiments, the generation and policy fitting times per decision were 4.2 and 2.2 seconds, respectively, on CPU (Appendix B.8). Various approaches could be investigated to speed up the algorithm. Distillation: Policy distillation, transferring knowledge from one policy to another, is commonly used to speed up computation. These approaches could distill TS-Gen into policy that maps the current context X, and recent task history Hi to distribution over actions [Czarnecki et al., 2019]. Generation: Generation could be sped up by truncating or reducing the number of outcomes generated per timestep. For sequence models more broadly, there is great interest in speeding up inference time through architecture changes [Tay et al., 2022] and optimizing around hardware constraints [Aminabadi et al., 2022, Dao et al., 2022]. Policy fitting: Policy fitting could be done incrementally instead of being refitted from scratch at each decision time.
7 Discussion
We introduce generative TS algorithm for contextual bandits that is compatible with any generative model with low offline prediction loss and policy fitting procedure with low VC dimension. We prove regret bound for our algorithm that allows for misspecification of the generative model, and provides insights into information theoretic analyses for contextual bandits that may be of independent interest. Open directions include i) developing methods to guide how to choose appropriate policy class II [Foster et al., 2020], ii) quantifying how much offline data is needed to train high quality generative model (including settings where offline data is collected by behavior policy), iii) exploring if the generative approach to modeling uncertainty can be extended to more difficult
settings, like Markov decision processes, and iv) investigating methods to reduce computational cost.
Limitations. We evaluate our generative TS algorithm in only two experimental settings. As result, our experiments are primarily proof-of-concep for the viability of the generative TS approach. Additionally, in practice, complete task datasets, but in practice one may not have access to complete task datasets. We describe heuristic approaches
approach requires training generative model pe approximate
use approximate complete task datasets from partial task datasets in Section 4.1. Further work is needed to assess practical feasibility in more complex settings and to formalize how well our heuristic approaches perform theoretically. Finally, our generative TS algorithm may also be computationally costly, especially when implemented with complex generative models. We discuss potential approaches improve computation cost at the end of Section 6.
Acknowledgments
This work was partially supported by the Al Agents Initiative at the Columbia Business School. Tiffany Cai was partially supported by the National Science Foundation Graduate Research Fellowship under Grant No. DGE-2036197.
References
Alekh Agarwal, Haipeng Luo, Behnam Neyshabur, and Robert E Schapire. Corralling a band of bandit algorithms. In Conference on Learning Theory, pages 12-38. PMLR, 2017.
10

---

Shipra Agrawal and Navin Goyal. Thompson sampling for contextual bandits with linear payoffs. In International conference on machine learning, pages 127-135. PMLR, 2013.
Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng. Olatunji Ruwase, Shaden Smith, Minjia Zhang, Jeff' Rasley, et al. Deepspecd-inferencee enabling efficient inference of transformer models unprecedented scale. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15, IEEE, 2022.
Ozlem Asian, Olcay Taner Yildiz, and Ethem Alpaydin. Calculating the vc-dimension of decision trees. In 2009 24th International Symposium on Computer and Information Sciences, pages 193-198. IEEE, 2009.
Nikolay Babakov, David Dale, Ilya Gusev, Irina Krotova, and Alexander Panchenko. Don't 't lose the message while paraphrasing: A study on content preserving style transfer. In Natural Language Processing and Information Systems, pages 47-61, Cham, 2023. Springer Nature Switzerland. ISBN 378-3-031-353200
Soumya Basu, Branislav Kveton, Manzil Zaheer, and Csaba Szepesvári. No regrets for learning the prior in bandits. Advances in neural information processing systems, 34:28029-28041, 2021.
Alina Beygelzimer, John Langford, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandit algorithms with supervised learning guarantees. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 19- 19-26, 2011.
Craig Boutilier, Chih- Wei Hsu, Branislav Kveton, Martin Mladenov, Csaha Szepesvari, and Manzil Zaheer. Differentiable meta-learning of bandit policies. Advances in Neural Information Processing Systems, 33:2122-2134, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 2020.
Sébastien Bubeck, Nicolo Cesa-Bianchi, et al. Regret analysis of stochastic and nonstochastic multi-armed bandit problems. Foundations and Trends@ in Machine Learning, 5(1): 2012.
Tiffany (Tianhui) Cai, Hongseok Namkoong, Daniel Russo, and Kelly W Zhang. Active exploration via autoregressive generation of missing data. arXiv preprint arXiv:2405. 19466, 2024.
Leonardo Cella and Massimiliano Pontil. Multi-task and meta-learning with sparse linear bandits. In Cassio de Campos and Marloes H. Maathuis, editors, Proceedings of the Thirly-Seventh Conference on Uncertainty in Artificial Intelligence, volume 161 of Proceedings of Machine Learning Research, pages 1692-1702. PMLR, 27-30 Jul 2021. URL https: http://eoooec.....c//e cella21a. html.
www.eeeeeeinnnm....ee...ee mlr. www..prrssvieee
Leonardo Cella, Alessandro Lazaric, and Massimiliano Pontil. Meta-learning with stochastic linear bandits. In International Conference on Machine Learning. PMLR, 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, 2021.
Bertrand S. Clarke and Andrew R. Barron. eeformation-theoredt asymptotics of bayes methods. IEEE Transactions on Information Theory, 36(3):453-471, 1990.
Sam Corbett-Davies, Emma Pierson, Avi Feller, Sharad Goel, and Aziz Huq. Algorithmic decision making and the cost of faimess. In Proceedings of the 23rd aem sigkdd international conference on knowledge discovery and data mining, pages 797-806, 2017.
Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommendations. In Proceedings of the 10th ACM conference recommender systems, pages 191-198, 2016.
Wojciech M Czamecki, Razvan Pascanu, Simon Osindero, Siddhant Jayakumar, Grzegorz Swirszez, and Max Jaderberg. Distilling policy distillation. In The 22nd international conference on artificial intelligence and statistics, pages 1331-1340. PMLR, 2019.
11

---

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness, Advances in neural information processing systems, 35: 16344-16359, 2022.
James Davidson, Benjamin Liebald, Junning Liu, Palash Nandy, Taylor Van Vleet, Ullas Gargi, Sujoy Gupta, Yu He, Mike Lambert, Blake Livingston, et al. The youtube video recommendation system. In Proceedings of the fourth ACM conference Recommender systems, pages 293-296, 2010.
Shi Dong and Benjamin Van Roy. An information-theorriicc analysis for thompson sampling with many actions. Advances in Neural Information Processing Systems, 31, 2018.
Vikranth Dwaracherla, Xiuyuan Lu, Morteza Ibrahimi, Ian Osband, Zheng Wen, and Benjamin Van Roy. Hypermodels for exploration. arXiv preprint wrXiv:2006.07466 2020.
Sandra Fortini and Sonia Petrone. Prediction-based uncertainty quantification for exchangeable sequences. Philosophical Transactions of the Royal Sociery A, 381(2247)20222142, 2023.
Dylan Foster, Akshay Krishnamurthy, and Haipeng Luo. Model selection for contextual bandits.
Advances in Neural Information Processing Systems, 32, 2019.
Dylan J Foster, Akshay Krishnamurthy, and Haipeng Luo. Open problem: Model selection for contextual bandits. In Conference on Learning Theory, pages 1842-3846. PMLR, 2020.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Tch, Danilo Rezende, and SM Ali Eslami. Conditional neural processes. In Proceedings of the 35th International Conference Machine Learning, pages 1704-1713. PMLR, 2018.
Studies in Applied Bayesian Data Science: CIRM Jean-Morlet Chair, Fall 2018, 2020. Ethan Goan and Clinton Fookes. Bayesian neural networks: An introduction and survey. Case
Amaury Gouverneur, Borja Rodríguez Gálvez, Tobias Oechtering, and Mikael Skoglund. An -eformationn-hheortiie analysis of thompson sampling for logistic bandits. In NewrlPS Workshop Bayesian Decision-makiin and Uncertainty, 2024.
David Haussler and Philip M Long. A generalization of sauer's lemma. Journal of Combinatorial Theory, Series A, 71(2):219-22) 240, 1995.
Joey Hong, Branislav Kveton, Manzil Zaheer, Mohammad Ghavamzadeh, and Craig Boutilier. Thompson sampling with mixture prior. In International Conference on Artificial Intelligence and Statistics, pages
PMLR, 2022.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50(2):1-35, 2017.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 14:11273-286 2021.
Emilie Kaufmann, Olivier Cappé, and Aurélien Garivier. On bayesian upper confidence bounds for bandit problems. In Artificial intelligence and statistics, pages 592-600. PMLR, 2012.
B.J.K. Kleijn and A.W. Van der Vaart. The bernstein-von-mite theorem under misspecification. Electronic Journal of Statistics, 6:354-381, 2012.
Branislav Kveton, Mikhail Konobeev, Manzil Zaheer, Chih-wei Hsu, Martin Mladenov, Craig Boutilier, and Csaba Szepesvari. Meta-thompson sampling. In International Conference on Machine Learning, pages 5884-5893, PMLR, 2021.
John Langford and Tong Zhang. The epoch-greedy algorithm for contextual multi-armed bandits. Advances in neural information processing systems, 20(1):6-11 2007.
Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge, 2019.
Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
12

---

Alessandro Lazaric, Emma Brunskill, et al. Sequential transfer in multi- armed bandit with finite set of models. Advances in Neural Information Processing Systems, 26, 2013.
Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. In-context decision-making from supervised pretraining. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023.
Lihong Li, Wei Chu, John Langford, and Robert E Schapire. A contextual-bandit approach to personalized news article recommendation. Proceedings of the 19th international conference on World wide web, pages 661-670, 2010.
Yingru Li, Jiawei Xu, Baoxiang Wang, and Zhi-Quan Luo. Scalable exploration via ensemble++. Preprint. An early version "Adaptive Foundation Models for Online Decisions: HyperAgent with Fast Incremental Uncertainty Estimation" presented ICML 2024 Workshops: (1) "Aligning Reinforcement Learning Experimentalists and Theorists"; (2) "Automated Reinforcement Learning: Exploring Meta-I Learnng,,,,ng, AutoMI., and I.MMs".
Yingru Li, Jiawei Xu, Lei Han, and Zhi-Quan Luo. Q-Star Meets Scalable Posterior Sampling: Bridging Theory and Practice via HyperAgent. In The 41 International Conference on Machine Learning (ICML), Proceedings Machine Learning Research, 2024.
Che-Yo Yu Liu and Lihong Li. On the prior sensitivity of thompson sampling. In International Conference on Algorithmic Learning Theory, pages 321-336. Springer, 2016.
Xiuyuan Lu and Benjamin Van Roy. Ensemble sampling, Advances in neural information processing systems, 30, 2017.
Mohammad Mehrabi and Stefan Wager. Off-policy evaluation in markov decision processes under weak distributional overlap. arXiv:2402.08201 [star.ML], 2024.
Seungki Min and Daniel Russo. An information-theoreet analysis nonstationary bandit learning. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, 2023.
Shira Mitchell, Eric Potash, Solon Barocas, Alexander D'mour, and Kristian Lum. Algorithmic fairness: Choices, assumptions, and definitions. Annual review of statistics and its application, (1):141-163, 2021.
Ahmadreza Moradipari. Mohammad Ghavamzadeh, Taha Rajabzadeh, Christos Thrampoulidis, and Mahnoosh Alizadeh. Multi-environment meta-learning in stochastic linear bandits. In 2022 IEEE International Symposium Information Theory (ISIT), pages 1659-1664, IEEE, 2022.
Subhojyoti Mukherjee, Josiah Hanna, Qiaomin Xie, and Robert Nowak. Pretraining decision transformers with reward prediction for in-context multi-task structured bandit learning. arXiv preprint arXiv:2406. 05064, 2024.
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. In Proceedings of the Tenth International Conference on Learning Representations, 2022x.
Balas K Natarajan. On learning sets and functions. Machine Learning, 4:67-97, 1989.
Gergely Neu, Iuliia Olkhovskaia, Matteo Papini. and Ludovic Schwartz. Lifting the information ratio: An information-thooree analysis of thompson sampling for contextual bandits. Advances in Neural Information Processing Systems, 9498, 2022.
Jerzy Neyman. On the two different aspects of the representative method: sampling and the method of purposive selection. In Breakthroughs in statistics: Methodology and distribution, pages 123-150. Springer, 1992.
the method of stratified
Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty- meta learning sequence modeling. Proceedings of the 39th International Conference on Machine Learning, 2022.
13

---

Ian Osband and Benjamin Van Roy. Bootstrapped thompson sampling and deep exploration. arXiv preprint vrXiv:1507.00::0 507.00000, 2015.
Ian Osband, Zheng Wen, Seyed Mohammad Asghari, Vikranth Dwaracherla, Morteza Ibrahimi, Xiuyuan Lu, and Benjamin Van Roy. Approximate thompson sampling via epistemic neural networks. In Uncertainty in Artificial Intelligence. PMLR, 2023.
Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel. M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Courapeaa, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 2011.
12:2825-2833
Chao Qin, Zheng Wen, Xiuyuan Lu, and Benjamin Van Roy. An analysis of ensemble sampling. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, 2022.
Carlos Riquelme, George Tucker, and Jasper Snock. Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling. In International Conference on Learning Representations, 2018.
Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American statistical Association, 100(469):322 -311, 2005.
Daniel Russo and Benjamin Van Roy. An information-theorriie analysis of thompson sampling. Journal of Machine Learning Research, 17(68):1 2016.
Daniel Russo and Benjamin Van Roy. Learning to optimize via information--ireecte sampling. Operations Research, 66(1):230- 252, 2018.
Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, and Zheng Wen. A tutorial on thompson sampling. 2020.
Ilya O Ryzhov, Warren B Powell, and Peter Frazier. The knowledge gradient algorithm for general class of online learning problems. Operations Research, 60(1): 1001)99919 2012.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108 2019.
Norbert Sauer. On the density of families of sets. Journal of Combinatorial Theory, Series A, 13(1): 145-147, 1972.
Bhadresh Savani. destller-baceuucccccdeeenneeees 2022. URL https::/hugiingaaceee.
Shai Shalev-Shwartz and Shai Ben-David. Understanding Machine Learning: From Theory to Algorithms. Cambridge University Press, 2014.
Saharon Shelah. A combinatorial problem; stability and order for models and theories in infinitary languages. Pacific Journal of Mathematics, 41(1):247-261, 1972.
Max Simchowitz, Christopher Tosh, Akshay Krishnamurthy, Daniel J Hsu, Thodoris Lykouris, Miro Dudik, and Robert E Schapire. Bayesian decision-making under misspecified priors with applications to meta-learning. Advances in Neural Information Processing Systems, 2021.
Yi Tay, Mostafa Dehghani. Dara Bahri, and Donald Metzler. Efficient transformers: survey. ACM Computing Surveys, 55(6):1-28, 2022.
William R Thompson. On the likelihood that one unknown probability exceeds another view of the evidence of two samples. Biometrika, 25(3-4):285-22 1933.
Dustin Tran, Jasper Snoek, and Balaji Lakshninarayanan. Practical uncertainty estimation and out-of-distribution robustness in deep learning, NewrIPS Tutorial, Google Brain, 2020.
14

---

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Sahil Verma and Julia Rubin. Fairness definitions explained. In Proceedings of the international workshop on software faimess, pages 1-7,20 2018.
Runzhe Wan, Lin Ge, and Rui Song. Metadata- based multi-task bandits with bayesian hierarchical models. Advances in Neural Information Processing Systems, 34:29655- -96668, 2021.
Zhendong Wang and Mingyuan Zhou. Thompson sampling via local uncertainty. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 10115-10125, 2020.
Zheng Wen, Ian Osband, Chao Qin, Xiuyuan Lu, Morteza Ibrahimi, Vikranth Dwaracherla, Moham- mad Asghari, and Benjamin Van Roy. From predictions to decisions: The importance of joint predictive distributions. arXiv preprint arXiv:2107 092224, 2021.
Fangzhao Wu, Ying Qiao, Jiun-Hung Chen, Chuhan Wu, Tao Qi, Jianxun Lian, Danyang Liu, Xing Xie, Jianfeng Gao, Winnie Wu, et al. Mind: A large-scale dataset for news recommendation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 3597-3606, 2020.
Weitong Zhang, Dongruo Zhou, Lihong Li, and Quanquan Gu. Neural thompson sampling. arXiv preprint arXiv:2010.00082,, 2020.
15

---

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:
The claims in the abstract and introduction do accurately reflect the paper's
contributions and scope. We propose a generative version of Thompson sampling for contextual bandits, provide regret bounds, and include empirical demonstrations.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Limitations are discussed in the Discussion section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a 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 (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect 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 on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be 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?
16

---

Answer: [Yes]
Justification:
Yes, all assumptions are stated, and proofs are in Appendix A.
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 theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Yes. We describe our experiments thoroughly in Section 6 and Appendix B. We also provide code the supplementary materials and at https: /topp:://ithull com/ namkoong- largg-labtts--eo gen...
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to 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 dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution 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 accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors 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 or verifying the results.
5. Open access to data and code
17

---

Question: 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?
Answer: [Yes]
Justification: We include code for experiments in the supplementary materials and also at attps://githuu www. naw//ongooooaaaaa laaigg-lbbttttgeee tgn...
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
www.ssssssllliiieoo onPolicyl for more details.
While encourage the release of code and data, 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 (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / wttpppccppppppappppppptppppppppp cooollccggidess.
ceabbbmmsssoonPe ssiiinolicyyyy 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 proposed method and baselines. 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
Question: 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 to understand the results?
Answer: [Yes]
Justification:
We include some details in Section 6 and remaining details 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 to 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
Does the report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We provide standard error bars for our results that represent + s.e. in our experiments, with standard errors calculated assuming normality, and we describe over what population they are averaged (500 bandit environments, drawn IID).
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

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 run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is 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 are out range (e.g. negative error rates).
If error bars 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 on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We discuss the compute resources we used in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 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 or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: / http///eeripp.eeee wttppppccccpppppppppppppppppppct ctw.ae/eeeeiaaieeeeeeeeeennee
Answer: [Yes]
Justification:
We have reviewed the NeurIPS code cthics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. g.. if there is special consid- eration due to laws or regulations their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The primary contribution our paper is a very general methodology for contextual bandits, accompanied by similarly general theoretical analysis, and empirical demonstration. Because of the nature of our contribution is unlikely that there are significant societal impacts.
Guidelines:
19

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 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 or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility 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: [NA]
Justification: The paper poses no such risks. a general method for contextual bandit settings.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere 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
Question: 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?
Answer: [Yes]
Justification: The only assets we use are the MIND dataset [Wu et al., 2020], DistilBERT [Sanh et al., 2019], and two pre-trained text classifiers from huggingface. For all of these we ensure we follow their license agreements (see Appendix B for more information on their URLs, usage, and licenses).
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 used and, possible, include URL.
20

---

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.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is 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 is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide code along with our submission in the supplementary material and in https://githu: cwm/nammeeeeeeeeeccooe http//mmmooooog--ooo lab/ts- gen. There are other new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was 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. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do not use human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, 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 describe 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: [NA]
Justification:
We do not use human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
21

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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 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, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs are not used as a core part of our methods.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy httpp:://neuripass for what should or should not be described.
22

---

A Theory
A.1 Notation
Throughout, we use E, to denote expectations conditional on He, i.e., use
E2 [-] E 1H)-
We use H(Y to denote the entropy of a discrete random variable Y, i.e., H(Y) = P(Y v) log P-Y = y)dy. We also use H,(Y) - H(Y He) to denote the entropy of conditional H4; Note that standard in information theory, HY(Y) not random variable, rather, marginalizes over H4:
H.YY) = HY 74) E E P(Y = H4) P(Y H4)dy
Above, the outer expectation marginalizes over the history Hc.
We also use [(Z Y) to denote the mutual information between some random variables and Y = y) log P(Z=s.Y=v) F[2=;)RY =y) dzdy. We further use (,; Y) to denote the mutual information between Z and Y conditional on H, (which we then marginalize over H4), i.e.
Y, i.e., Y) = fg P(Z z,
(Z = I(Z;Y) 74))
/ P(Z = 2, Y (4) log
P(Z = 2, = = Hc) P(Z= HY) = Hc)
= E
drdy
(8)
Above, the outer expectation marginalizes over the history He.
Finally, we use DRr. (p(Z (Z X)) to denote the KL divergence, i.e.,
p(Z X)
DKL (p(Z X) p (Z x))
p(Z X) p/ X)]
Above, the outer expectation marginalizes over X.
A.2 Showing Algorithm implements Thompson Sampling (Probability Matching)
Proposition 1 (Algorithm Implements Thompson Sampling). Algorithm I with imputation model p" implements Thompson Sampling (probability matching), i.e., the following holds almost surely:
PPA, H4) PXx = 74).
Proof. Recall that Algorithm selects actions as follows:
PPA, H4)) = P( (Xx = H4).
Since (T E H4) and from Eq (1) T ~ p", the distributions of and are equal given Hy.
Hence, with probability for any j:
P(# j 1 H.) = P(7 = He))
The above implies that
P(m'(X;; H.) =
Ht.)
Combining the above statements gives the result.
A.3 Bounding the conditional entropy by VC dimension
Proposition (Complexity bound on entropy). For policy class 11 over action space A, with Nataranjan dimension d (equivalent to VC dimension when [A,] 2).
H(T" (Xi.) Z,) < H(m (Xxx)) O(d log(T |A, l)).
23

---

Proof. The first inequality ( (X(rr) 2.) S (TTTT(((X1) holds by the chain rule for entropy.
Note that π (XI:T) ) is a random vector of dimension I' where each dimension can take |AAy| different values. By generalization of the Sauer-Shelah lemma [Sauer, 1972, Shelah, 1972], specifically Theorem 2 and Corollary 3 in [Haussler and Long, 1995], if multi-class function that can take on [A| different values has Nataranjan dimension d, then that function class can produce at most Ee ()) (1Ar - 1) = O(T4 (d) different labelings of any points. Thus, since a coarse upper bound on the entropy of random variable is the log of the number of unique values that variable can take, get that H (T" (X1:)) log Ei=o (14r 1)² O(d O(d-log(T) (AAr))(AA))
A.4 Regret bound for Generative TS with an approximate imputation model
A.4.1 Lemma To minimize loss po needs to approximate p".
The next lemma is a standard result connecting the excess expected loss sequence model ps to its KI. divergence from the true sequence model p". The expected loss of a sequence model po denoted {(p)); See (4). To minimize loss, pa. the learner needs to closely approximate the truc sequence model p².
Lemma (Decomposing loss under pe). For the loss l as defined in (4),
t(po) = E(p") ) + DKL (p" (Xi:r, (V11) JaEA, | (+) 11 X1:r (1 JacA, JaEA, k+:)
Proof. By the definition of the expected loss in (4),
E(po) E(p") -E [logpo (XIr, (V 1Y1a- [+)] + E [loge" (Xur, Dre (x))
DKL (p' (v
1.7 X1T 2+) Pe 1:7 X1.r 2,)
Above, the final equality holds by the definition of the KI. divergence.
A.4.2 Lemma 2: Action selection under perfect vs. imperfect imputation models.
Lemma (KL Divergence in next action distribution). For any t,
DKL. (P, (TT (Xi+) -) || IP, (Ae = )) ((po) Elp').
Proof. Note the following:
DKL. (Pr (I (XGT) = (A,
DKL. (XI:T, (. JaEA,
HHi) IPp (Xur, I,, (K (Hl))
(a)
DKL (pp- (XITS
(X hael Z,) Ppo (XIr, Dr JaeA, (+)
E(po) - E(p²).
(6)
(e)
Inequality (a) holds because n² (Xrit) and A, are both are derived by applying the same function to the contexts XIT and outcomes {x JoEA.
Inequality (b) holds because by the chain rule for KL divergence,
DKL (Pp- (XI.T, { (x(a)) J JaEAA H4) PP (X1.T, { (x-(a) I aEA (84)))
= DKi. (Pyr (XErT (x11) JaEA-
27) || Pp (XI:T, {1 1:7 JoEAr Z+))
+ DKI. (P, (He, X | 27) || Pr, (He, X 2r)),
and the KI. divergence non-negative.
Inequality (c) holds by Lemma (Decomposing loss under po).
24

---

A.4.3 Lemma 3: Mutual information equivalency.
Lemma 3 (Mutual information equivalency).
.A))
E Σ - (TT a). Drr. (P(()) 1 (X(:))
a) (IP:(r!
a,
Proof. Note that
( (XUT): (Y((A),
(A)) Ai)) = (X;::) AA)
(a)
AA.= Po(A: = ( (XiT)); (xx)(("(xxx)));))"" (( A = a)
P=(A;= - a)li(= (X
E3
(6)
a
(c)
uEAr
Σ A,
Pa(A: = a) Σ P((T DKL (P:(Y (X(T)) @) (P,(
(d)
aEA,
Above, equality (a) holds since x T)) and A, are independent conditional on Hr. Equality (b) holds by the definition of conditional mutual information. Equality (c) holds because y,(a) and x (X+;T) are independent of A, conditional Hz. Equality (d) holds by the KL divergence form of mutual information.
A.4.4 Lemma 4: Mutual information bound for policies.
Lemma 4 (Mutual information bound for policies).
(TTT(( (A)) (.)) s (XT)() 2+)
1=1
Proof.
((TT((x)))( At))
Σ
Te TT( (Y(4),
f=1
11 ( (XI:T): (X, (1), Ab)
(18)
Hi (IT (XIT)) H1 ( (T" (XI:))
A1)
AA))
(19)
Hi (X1r)) $ H(T" (Xır) (+))
((r)
(v)
For inequality (i), note that for any random variables X1, X2, Y (where X1, X2 are discrete), by properties of mutual information and entropy,
(2); H(X,, - 7((x1X2213
-H(X)) - H(X; 1Y) + H(X, X1) H(x, YY
+ =I(X;;Y)+ I(X,; |X,)
The above implies that I((X1,X2); > I(X1;Y) since I(X;;Y | X1) > 0. Recall that
(XI.T) : (X;; Thus, since - (X E T (X we have that
(XT:)) (X, (A.) .A)) 2 1 (Xi:)) (A) (A))
Equality (ii) uses the chain rule for mutual information.
25

---

Equality (iii) holds by the relationship between mutual information and entropy.
Inequality (iv) holds since entropy is always nonnegative.
Inquality (v) uses that H1 (T (Xi:r) = (Xi:r) ) Z,) H(TT((I'11) | Z,), where the first equality holds by the definition of H, and the final inequality holds by the chain rule for entropy.
A.4.5 Proof of Theorem 2
Theorem 2 (Regret bound for Generative TS with an approximate imputation model). For Algorithm I with imputation model po, Gen (po).
Ar 2T
A A(A1S-Gen(Po)) S
H(TT( Z,)
(2(((p)) {(p')}
Ponalty for sub-optimal prediction
Regret bound for Thompson sampling
Proof. Note that by the law of iterated expectations,
A(ATS-Gea) = E
1 T
(Xxt) -((y()))))
ROY
= E
1 T
E
R(Y (a (Xist)l) R(((tr(Xett)))) - R(Y((A))))
Consider the following for any t E [1: T):
ROY (T' (X0))) - (()) Eg
Σ Po(t a) a) E. [R())(")) nt (Xi:t)= = =a] Σ
Pr(A: = a) -E, [x(,(r( | =
aeA,
aEA,
Σ
Pi(T = a) E. (((Y -" (Xx:t) =
Σ
Pr(A, = a) EE [x))
dEA,
aεd,
Σ
(XiT) -E [E,( -E,
=
aEAr
Σ (VPT(T(XET) = a) (Pr a))
dEAr
(VP((T E. [R(Y (Xii)) = a] + VP.(A. a)E, [x(
Σ (PT(T Ec ()))(())) . (Xi:)) aj - [.(((
(i) aEA.
IP a) Pi(Ar = =a)l
aedt
Σ PA(A- = a) Σ P((TT(((T)) =)) DKL (P((Y:!a) (XxT)) a) (Y
2
(10)
aeA,
ae.A,
(X:7) a) a) - -P((A = =a]
a
26

---

Above, equality (i) holds since conditional on Hi, the action A, and the outcome Yy(a) are independent. Inequality (ii) uses that R takes values in [0, 1] in the second term. Inequality (ii) above holds because:
Σ aEA,
I IT (Xe:)) (X:T) (
= (E,
[R,
Arl Σ aEA,
(X = = a) (Ee (E(( | a] - E. ()(
(a)
Arl Σ aEAr
Pr(At a) Σ
P.(TT((TT)) a) (E.[R((..") 1 1 (X(ir) = a] E. [x((()(
aEAr
|A- 2
Σ Po(Ar=a) a) Σ (X(T) = a) DKL (P((Y((") TT (Xxx)) a) || (Y. aEA,
GEA,
Incquality (a) uses Cauchy-Schwand inequality. Inequality (b) uses an elementary equality of summation. Inequality (c) uses Fact 9 of Russo and Van Roy [2016] (which uses Pinsker's inequality).
Using the above result, averaging over t and taking an expectation, we get
T
A (AATS-Gca)
=
E
(X
B(Y;
- R(X [1)]
T
E
A
Σ
PA(A- a) E (((TT(X(TT) a) DKL (Pa(Y(!a) ( [**(x) ( IP((Y!!0))
2
a
Pi (XGr) = -PPP))PPP- = الlه
+ E
T
aEA,
Σ
a) Σ
(X)T)(XTTT)= = ū) DKL (Y, (X(:t) = a) || (x,
2
aEA,
REA.
E
DKL (XiT) = -) || || |P -))
T
Ar
1
T
(X(T)) 10100,A) +
T
2 DKL (Po(Tt(Xa:)) = .) ||P((Ae= (A)=))) .))
2
T
T
(6)
1=1
|11
1
I (T(T)) (Y,(A),))))) +
1
2. DKL. (Pr =) 1I :))=)) .)
2
T
T
(A0)
Arl H(T (Xir)
Zr)
+ -(p(()
2T
(10)
Incquality (i) uses Jensen's 's inequality on the first term and Fact 9 of Russo and Van Roy [2016]
(which uses Pinsker's 's inequality) on the second term
Equality (ii) uses Lemma (Mutual information equivalency).
Inequality (iii) uses Jensen's inequality.
The first term in inequality (iv) uses Lemma 4 (Mutual information bound for policies) and the second term uses Lemma 2 (KL Divergence of next action distribution).
27

---

A.5 Regret bound for Generative TS with perfectly calibrated imputation model p*
Theorem (Regret bound for Generative TS with perfectly calibrated imputation model p"). For Algorithm I with imputation model p", Gen
Arl
A Gen (1")) S
H(T (X1:) 2z).
2T
Moreover, ((ATS- Gen Gen""") S V D. T ((xx((X1:1) ?r), where I bounds the information ratio [Russo and Van Roy, 2016], i.e., r max, r a.s. for D
E[Rry("" - 1( (X A..)6.))
Proof. The first result that A(ATS-Cen(P")) Gen()")) S lary of Theorem 2 by setting po = p".
AA 27
(X(TT(
Zr), holds as direct corol-
We now show the second result that S(ATTS-Cen(P) S V
F H(T<(X1:1) Zr). It holds by
very similar argument as Proposition of [Russo and Van Roy, 2016].
A ((ATS.Cien)
R(Y ((y(r"( -
[n(y -
B,
[[Y,
T
T
(5)
T rr (XTT)) .Ar)
T
E
S <V E
(((T) AA)
S V [ H( (XIT) 24,) (so)
r.
E
T ( Ad)
(815)
Equality (i) holds by the law of iterated expectations. Equality (ii) holds by the definition of Ir- Inequality (iii) holds by Cauchy-Shwarts Inequality (iv) holds by Lemma 4 (Mutual information bound for policies).
A.6 Comparison to existing regret bounds
Lemma 5 (Bounding information ratio for linear, non-contextual bandits), Suppose E [R(Y:) A.
((A,) 0r for some 0 Rd. Let the policy class II be such that for any 11, (a) ((A)) T for some 0 E Rd. Then, I, < 2 a.s.
Proof. This result follows by Proposition 5 of Russo and Van Roy [2016].
By our Theorem 1, we have that the per round average Bayesian regret bounded by V 27 (Xi:r) | Z7). Note that by Theorem 29.7 in Shalev-Shwarzz and Ben-David [2014] linear multiclass predictor of the form argmax 0 (p(x, a)) for 0 E R has Nataranjan dimension less than equal to d. Thus, by applying Proposition 2, we have that =7(x=(XX11) 2+) < d log(T [Arl), so the per round average Bayesian regret is bounded by / d/A.1 log(T' (A+I).
Generative TS regret bound for linear and logistic reward settings. |A-
27"
Alternatively we can use the second result Theorem conclude that the per round average Bayesian regret is bounded by / f [TT* (x( Z7). By Lemma we can choose 5. so by applying the same Proposition 2 argument as above, we have that the per round average Bayesian regret is bounded by
2T log IAr
Thus, by combining the above two results, we have that
< d min(d, (Ar)) log(T |Arl)
A(ATT.Gen(P))
28

---

Linear logistic bandits. We compare to Theorem 4 of Neu et al. [2022]. We only provide a brief overview of their result here; Please see the paper for additional details. Additionally note that their result applies to adversarial contextual bandits, whereas our result only applies for stochastic contextual bandits.
In their problem setup, the rewards are generated using a logistic model where 0 E R
R(Y)) Xe; A Bernoulli (expit (x,x, Aa)T0))
They show that cumulative Bayesian regret of Thompson sampling (with a correctly specified Bayesian model) bounded by /2|A, (odgloo(2SCC + 1) + 1]. where ||0*1| S and C' is related Lipschitz smoothness of the logistic function. This means that the per round average Bayesian regret
bounded by / [d][A=1 (log(2SCT + 1). Our result from (9) matches up log factors.
Linear non-contextual bandits. We now compare to the result in Section 6.5 of Russo and Van Roy [2018]. Again, only provide brief overview of their result here; Please the paper for additional details.
In their non-contextual problem setup, the rewards are generated using linear model where 0" E Rdd
E (R(Y.) A. = a) = (
They show that cumulative Bayesian regret of Thompson sampling (with a correctly specified Bayesian model) is bounded by / log([A; DdT,. This means that the per round average Bayesian
regret is bounded by V 27 log(IAAr)). Our result from (9) differs by factor /min(, [Ar]) and (T) term.
The additional |A,| and log( factors, we believe, are not artifacts of our specific algorithm, but rather are consequence of the generality of our analysis.
The log(T') term comes from use of the Natarajan dimension, generalization of the VC dimension. This term is common in bandit regret bounds that rely on VC dimension- based analysis, seen in other work (e.g., Beygelzimer al. [201]] I]). appears to be an unavoidable consequence this type of generalized bound.
The IA| term is a consequence of the generality of our analysis, which does not utilize a shared parameterization across actions. The Russo and Van Roy [2018] bound for linear bandits is tighter because it leverages the linear structure, where E[R(Y)) A a] = p(a) B. In this setting, the parameter 3 is common to all actions, meaning information gained from observing an action can be used to inform beliefs about the rewards of all other actions. Our analysis, however, does not assume or utilize such shared structure. Instead, our regret-bound scales with the number of actions, similar to bounds for multi-armed bandits where the reward distribution for each arm is learned independently. This makes our bound applicable to broader class problems, but also looser for specific settings like linear bandits with shared parameters across actions.
While our result does not provide the tightest possible regret bound for specific parametric model, we present general and robust theoretical framework that characterizes the performance of Thompson Sampling variants that use modern generative sequence models and general policy classes
B Experiment details
B.1 Data generating environment
B.1.1 Synthetic bandit setting.
We form samples tasks r = (Z, (X,, (x()) bandit task consist one feature per action, i.e. = {2(a) [aEAA, JAEA, where only Z(a) R²a We sample task features 7(a) N ((2, I2) independently across |Ar| = 10 actions and contexts X, (0,, Is) independently across time. We let R(y) y and use the following generative model
(aEA))( as follows. The task features Z for a given
29

---

for y((a)
y((0) W(()) (er
(10)
where
UT(a W(() U + U/() 2(4) + U Xx + ((2) cross Z(n)
for (ww) == (1 + exp( =x) -w) Above we use X1,1:2 to denote the first two dimensions of X, The latent variables multivariate Gaussian: Uta N(0,)). U (22,, 0.25"), U
N (1s, J5 0.25²), and Uerdes is diagonal matrix where the diagonal entries are drawn independently from N 0.25%).
B.1.2 Semi-synthetic setting.
We form samples of tasks r = (z, (X,, JoEA as follows. We consider a semi-synthetic news recommendation setting in which we use text headlines Z(a) for action a. We let R(y) = y and use the following generative model for (()))
y-(0) Wh() Bernoulli(((W!11)))
(11)
where
Wh(0) Ula) + UZ Oz U ox (X.) + Φx xx (X:) cross Oz (Z(4)).
Above, Øx (X,) E R and (Z()) E R2 are complex nonlinear function of X,, Z("), which increases the difficulty of the learning task; We describe these functions in detail below. Note, xx (X,) denotes the first two dimensions ox (X,) E R?. The latent variables are multivariate
Gaussian: U(a) N(0, 1), U N(1,, I2 0.25²). U(e) NN4,,, 0.252...... Ta), and the matrix t:(a) is diagonal with diagonal entries drawn independently from (1,05255)
Contexts and ox. The contexts X, ~ N (0s; As) independently over time. We use
Φx (X,) X1,1:4 sign(X..s),
i.., Φxx multiplies the first four dimensions of X, by the sign of the fifth dimension. Above, Xt,1:4 denotes the first dimensions of X,,
Tasks features and Φz. To form task, sample |A+| = 10 headlines (()) uniformly from the MIND large dataset [Wu et al., 2020]. oz (z))) R2 where the each dimension is the output of a pre-trained binary classifier evaluated on the news article. The first dimension is the output of probability output of pre-trained sentiment classifier [Savani, 2022] and the second dimension the probability output of a pre-trained formality classifier [Babakov et al., 2023]: The outputs based on their distribution in the training set. Both classifier models were obtained from huggingface.com.
normalized to have mean 0 and variance 1
B.2
Offline pretraining
B.2.1
Sequence model architecture
Synthetic setting. This architecture is described by Figure 7 except the X MLP head and Distil- BERT head should be replaced by identity mappings. In the synthetic setting Pe is simple recurrent neural network where the MI.P takes as input 7(a), current context X,, as well as summary statisties of the history Hr (discussed below). Before being fed into the MIP head, the summary statistics are then repeated 100 times and concatenated into a single vector. The (()), the current context X,, and the repeated summary statisties of the history are fed into the final MLP head, which has 3 hidden layers, each with width 100. Note that the MLP consists of linear layer taking the input to the first hidden layer, the hidden linear layers, and finally a linear layer taking hidden layer to the output before the sigmoid, which is total of 5 linear layers. The output of the MI.P head is fed through sigmoid function obtain a prediction for the probability that the next outcome is 1 (rather than 0).
output from the last
30

---

Distil BERT
Z
Recent Y 1==1
Recent history summory statistics
Final MLP head
sigmoid
Po .)
Recent X 1--1
MLP
x
head
Figure 7: Diagram of model architecture for pe, for semisynthetic settings. In synthetic settings, the model architecture is the same, except that does not include the DistilBERT [Sanh et al., 2019] encoder to process text, or the MLP encoder process contexts X,.
The summary statistic of H4 only contains information about action a, i.e., {(X,, Y,): << A, a}. For these summary statistics, row is one element in {X, :8 < t, A, = a}. We do the same for {Y, < A, = a} to construct vector Y. The X, and Y, appear in X and Y in order according to timestep The summary statisties are [X X + 1)-1 and X Y.
aggregate the context vectors X, into matrix X, where each
Semisynthetic setting. This architecture is described by Figure 7. In the semisynthetic setting, Pa is implemented to take as input action-specine task feature Z(a), current context X,. well as summary statistics of the history Hy (discussed below). As displayed in Figure 7, the model architecture is as follows. We concatenate DistilBert [Sanh et al., 2019] embedding of headline Z(a) with Xt, and a summary statistics of the history (desribed below) that is repated 100 times. Then, this concatenated vector is fed into the final MLP head (3 hidden layers, width 100). Finally, the output of the MLP is fed through sigmoid function to obtain prediction for the probability that the next outcome is I (rather than 0).
The summary statistic of 74 only contains information about action i.e., {(X,, Y,) :s :s A,
a). For these summary statistics, aggregate learnable MLP embedding øx (of depth 2 and width 100, labeled "X MLP Head" in Figure of the context vectors Đx (X,) into a matrix Φx (X), where each row is one element in (Xx) s < t, A. = a). We do the same for (Y, < t, A, = a} to construct vector Y. The Φx (X&) and Y, appear in 6x (X) and Y in order according to timestep
B.2.2 Forming approximate complete task datasets from partial datasets
As described in Section 4, Doffline ideally consists of bandit tasks p² as described in (1).
In practice, one may not have "complete" task datasets 7 (Zr, XI.7, (Y((0) "i....... YY,,,,,,)) Yx(o)) a }.
but instead have some partial datasets, e.g., {2r, (1, Yi),... ...,, (XT, Ar, Yr),. collected by behavior policy. In our experiments we use several heuristics to construct approximate complete tasks 7 from the the partial datasets. We use these approximate task datasets to form Dollliee (+1. 72, fs.....
The bootstrapping procedure we use makes several modeling simplifying assumptions, which are all common in the bandit literature:
Stationarity over time. We model the X₂ 's's as being drawn i.i.d. from an unknown distribution. )ae:: T]
Additonally, we model the (X,, (o)) as exchangeable over time, i.c., (Xa, (())))))))
D
(Xo(t), y((a) ) T].
Independence across actions. For given task we model the outcomes 1:T i.i.d. con- ditional on XI.T and Z. This means that the outcomes y(a) are not correlated with those from other actions, given contexts and task features.
31

---

Due the independence across actions assumption, instead of generating f, we instead impute rows F(a) I.T Y Y } for individual actions We use a bootstrapping procedure construct j(a),
described in Algorithm 4 below.
Algorithm 4 Bootstrapping historical data to form F(4)
Require: Historical data from action a, denoted 5(a) f {(X,, A. a}
1: Sample (with replacement) 7 tuples from S(a);
(X1, (10) (Xr, f, ((a) i.dd
Σ
((r.))
[S(0]
(Z.0)
2: return 7(0) = {(8,, p(a)), (10),, (Xx, }
B.2.3 Additional sequence model training details
Synthetic setting. For offline training of po, we sample 20k independent "task action" datasets {Z("), (1,Noo x ) according the data generating process from Appendix B.1.1; Specifically we use N(a) = 1000 for all a. This dataset is split into training and validation sets where 10k actions in each set. The training set is used for training pe via gradient descent for 100 epochs, with loss from display (6); Note for approximating the distribution of Xt, we use the empirical distribution of 1000 contexts X's ss from the training set (no gradient descent training). In each training batch, we use bootstrap resampling, specifically, Algorithm 4. The validation set is for choosing best hyperparameters and training epoch. We optimize weights in po with the Adam W optimizer. We try learning rates {0.1, 0.01, 0.001} } and choose the learning rate with the lowest validation loss, which 0.01. We set weight decay to 0.01. The batch size is 500 actions per batch.
For offline training of po. we sample independent "task action" datasets {Z(a), 1ItN(os, y²(a) 1=N For Z(a)'s use 104k headlines from the MIND dataset [Wu al., 2020]; 20k are used for the training set, 10k are used for validation, and 74k are used for bandit evaluation. The outcomes X and Y are generated according to the process described in Appendix 1.1.2; Specif-
Semi-synthetic setting. setting
ically use (a)) 1000 for all The training set is used for training pe via gradient descent for 40 epochs, with loss from display (6); Note for approximating the distribution of Xt, we use the empirical distribution of 1000 contexts X's I's from the training set (no gradient descent training). In each training batch, we use bootstrap resampling, specifically. Algorithm 4. We optimize weights in Pe with the AdamW optimizer. We try learning rates {0.1, 0.01, 0.001} and choose the learning rate and also the training epoch with the lowest validation loss; the learning rate chosen is 0.01. We set weight decay to 0.01. The batch size is 500. We do not fine-tune the DistilBERT encoder, i.e., its weights are frozen.
B.3 Online learning
Bandit datasets are constructed as described in Appendix B.1. In the semisynthetic setting, the headlines used are as described in Appendix B.2.3.
B.3.1 TS-Gen policy-fitting details
Here we describe additional details used to fit T' (-; 11 given an imputed task dataset f. Using ft. for each action
A,, we fit an action-specific model predict (binary) outcome Y given context X; We use f(a) (X; to denote this fitted action-specifie model. Note that these models do not incorporate task features Z(a). Then,
x" <"x;;t = argmax f(a) fA (t)
In our experiments we choose f to be either logistic regression function or an MLP.
For logistic f, we use the default logistic regression implementation from scikit-lean (Pe- dregosa et al., 2011].
32

---

For MLP-based policies, we use the default MLP classifier implementation (including hyperpa- rameters), also from scikit-learn [Pedregosa et al., 2011]. This is an MLP with one hidden layer of width 100, with ReLU activation, trained with Adam optimizer, with initial learning rate 0.001, and batch size 200. There is no early stopping or additional validation split.
B.3.2 Baseline bandit methods
The first three (Greedy, Apsilon-Greedy, and Softmax) are alternative ways to make decisions using an existing pre-trained sequence model po. The others (Linear Thompson Sampling, LinUCB) are contextual bandit methods that do not use pg-
Greedy. We use the samed trained sequence model pa as used by TS-Gen In the online step, at time i, we feed the history Hr (which includes the current context X,) into the model pa- We look at the predicted mean reward I. [R(Y,) Hr, A, = a) for each action a according to po and select the action with the largest predicted mean reward.
Epsilon-Greedy
This algorithm also uses py and at each decision time follows the Greedy policy with probability 1 - and selects an action uniformly at random from Ae with probability c. We use 0.1.
Softmax sampling. Softmax sampling also uses the sequence model Po to select actions. Just like the Greedy algorithm, at time t, we feed the history Hi (which includes the current context X,) into the model pe. We look at the predicted mean reward E [R(Y; He, Ar - a] for each action according to pa and put these values through softmax function with temperature b > 0. We then sample the action A, according to the softmax probabilities. Note that softmax sampling is also called Boltzmann sampling and is also called PreDeToR-7 Mukherjee et al. [2024]. Following Mukherjee et al. [2024], set b 0.05.
For lack of space, this is omitted in the main text but we compare PreDeToR-T with Greedy and Apsilon-Grredd later in this Appendix.
Linear Thompson Sampling (Isotropic Gaussian prior). We use Linear TS [Agrawal and Goyal, 2013] with the following Bayesian model with non-informnn -oniioformeeive prior. For each arm a E A, and time t, outcomes are modeled as linear function of X,.
Y (a) x g(a) + c,
where
gla) N(H, Σ) and (a) ~ ~N(000²))
where ef is modeled as Gaussian with mean 0 and variance 1/4 (since the maximum variance of Bernoulli is 1/4). Note that unlike TS TS-Gen, linear Thompson sampling does not learn rich and flexible prior based on task features Z,.
Linear Thompson Sampling (Fitted prior). We use Linear TS [Agrawal and Goyal, 2013] with the following Bayesian model with prior using historical data Dollinee We use Aollliee to denote all actions across all tasks in Dolline We fit the following Bayesian linear regression model for each action a E Aoffline
y:a y = X, g(a) (0)
where g(a) N(,, E) and (2) N()
where g(a) are drawn iid across a, and E (a) drawn iid across a, t, that
((()) X ~ N( T XG, 02 + x, EX,).
For fitting μ1 E, 2 we do the following:
For each action a E Aoffline in the available historical data (see Appendix B.2.3), we fit the action-specific least squares model:
B(a) = argming
E(r -x,3
IET
where Ti denotes the first 80% of timesteps in [1, 2, ...... T7.
33

---

Then we set = [ EeEAutia 3(a) and Σ to be the sample covariance of the 3( across a E A We set ô² to the sample variance of the residuals, iee. the sample variance of y-u) -X, across and f, where E _ and & T2, and where Tz denotes the final 20% timesteps in [1.2....,
Linear Thompson Sampling Using Learned Features (Isotropic Gaussian prior) Here, we propose variant of Linear Thompson Sampling above, but using features extracted from the learned sequence model po- Let do (Z()) X²) denote the last-layer feature emhedding (using the output of the last hidden layer) in the MLP head in the sequence model po used for TS-Gen (see Section B.2.1) evaluated for the current context X, and action feature Z(#); note we do not feed any history into the sequence model pa when forming on (Z()),
We the following Bayesian linear regression model, which is linear in ()( Xx):
y = 60 ,o())),,,) x,) 3(a) + c (a) where
B(a) (0,, Ia) and 4 (O. 1/4).
(1) are drawn iid across a, Above, the noise variance is set to 1/4, the maximum variance of Bernoulli random variable. Note that while this version of linear Thompson sampling does use Po to form the context (Z()))) X,), it does not utilize a fitted prior.
where 3(a] are drawn iid across a, and et
Linear Thompson Sampling Using Learned Features (Fitted prior) Here, we propose a variant of the Linear Thompson Sampling Using Learned Features method above, but fit the prior using historical data Doffline We use Aottlnne to denote all actions across all tasks in Dettline. We fit the following Bayesian linear regression model for each action a E Auflline
Y:()) ()(())) X g(a)
G
(a)
where
g(a) N(,, Σ) and do) ~ ~N0,00
where g(a) are drawn iid across a, and e are drawn iid across a, I. For fitting μ, Σ, 2 do the following:
For each action E Aoliline in the available historical data (see Appendix B.2.3), fit the action-specific least squares ridge-regression model using the corresponding historical data from Doflinee
g(a) argmin, { |all812 12
(Y Σ(r!") xx)
TET
where Ti denotes the first 80% of timesteps in [1,2, ...,. T7. We add the ridge penalty term because ((()) X,) is 100-dimensional and we found that the adding the ridge penalty leads to more stable coefficient estimates.
We set the ridge parameter & = 0.1.
En 3(a) and Σ to be the sample covariance of the 3(a) across a E AAullliee to ensure the covariance matrix is well-conditioned (to avoid numerical issues when computing posteriors), we add 10 L, the sample covariance. We set 8² to the sample variance of the residuals, i.e. the sample variance of y((") Φo (Z()) ,X)) g(a) across and I, where a E Aoffllnee and / E T2, and where Tz denotes the final 20% of timesteps in [1,2,.... T).
Then we set
We implement LinUCB-disjonie in [Li et al., 2010], contexts X,. We set 0.1 it performs well in comparison to a small set other values tried ((0.1,12))). Note that unlike TS-Gen, LinUCB does not learn a rich and flexible prior based on task features Z,.
LinUCB.
B.4
Additional simulation results
B.4.1 Sequence loss vs. regret under TS-Gen (Figure 8)
We examine the relationship between sequence model loss ((po) and regret of TS- using Pe in the SYNTHETIC setting. Our Theorem 2 suggests that the lower the loss of sequence model po the lower the regret of TS-Gen using that sequence model pa- We examine this by varying the amount of training tasks used to learn po and thus obtain sequence models with different losses. Indeed, in Figure 8, models trained on more data tend to have lower sequence loss, which tend to have lower regret.
34

---

Prediction by Timestep
Cumulative regret timesteps
0.433
10000 Training Rows
3000 Training Rows
25
Training
Cross-Entropy Prediction
300 Training Rows
2,043
2.333
10000 Training rows
Training 1000 Training rows
2.430
Training rows
0.07%
Timestep
Decision Times
Figure 8: Sequence loss bandit regret: and regret for TS-Gen by pre- pre-training our sequence models offline on varying dataset sizes in the semisynthetic setting. As training dataset sizes are smaller, sequence loss (left) is higher (worse), and bandit regret (right) is higher (worse). "Training rows" refers to the number of actions used in the pool of actions to select from to form tasks (Appendix B.2.3). (Left):: Prediction loss timestep. We plot an empirical estimate of the per-timestep (nonccumulative) loss from (4) by evaluating our held-out validation set. Error bars represent =1 s.e. (Right): Cumulative
We demonstrate the relationship between sequence loss
sequence models regret for TS-Gen using the corresponding sequence models, with logistic policy class, and relative to the logistic "oracle". Error bars represent 1 averaged over 500 re-drawn bandit environments.
B.4.2 Policy class for TS-Gen (Figure 9)
The choice of policy class II affects both the reward achieved by TS-Gen, and the "oracle"; see Figure 9. In the semisynthetic setting, TS-Gen has moderately greater reward using MLP-based policy than logistic policy. In contrast, the "oracle" using an MLP-based policy is much better than the "oracle" using a logistic policy.
Regret Timesteps: Semisynthetic
Regret Timesteps: Semisynthetic
70
80
TS-Gen, logistic policy
H5-Gien, MUSbased policy
30
TS-Gen, logistic policy
H5-Gie, MUPbased policy
60
"Oracle", MLP-besed pollcy
logistic policy
"Gracle", logistic policy
20
15
30
10-
10
100
200
300
400
500
100
200
300
400
500
Decision
Times
Decision Times
Figure 9: Varying policy classes in the semisynthetic setting. The same experimental results are plotted on the left and the right. The plot on the right calculates regret relative to the logistic "oracle", while the left calculates regret relative the MLP-based "oracle" Error bars are #1 s.e. across 500 bandit environments.
B.5 TS-Gen with truncated imputation horizon
TS-Gen imputes missing outcomes up to the horizon T. In practice, one may want to truncate the number of imputed timesteps in Algorithm 3 to smaller number than T in order to reduce computation cost in the meaiiin step for TS-Gen, or to run TS-Gen when the total number of timesteps 7 is unknown. Fortunately, regret does not degrade quickly when the imputation horizon is truncated, which we observe in Figure 10.
B.6 TS-Gen with simpler sequence models
To understand how the regret for TS-Gen depends on the complexity of the sequence model po we compare TS-Gen with simpler sequence models. Specifically, we compare the regret results in Section 6 with their counterparts where the final MLP head of the sequence model (Figure 7) is
35

---

Regret Over Timesteps: Synthetic
Regret Over Timesteps: Semisynthetic
25
70
60
20-
15-
40
10
TS-Gen imagine 500 (default
30
TS-Gen imagine 500 (defaulti
TS-Gen imegine 300
TS-Gen imagine 300
Gen imagnne
10
Sen imagine
TS Gen imagine 100
TS Gen imagine 100
100
200
100
200
500
Decision Times
Decision Times
Figure 10: Regret for TS-Gen with truncated imputation horizon in the synthetic (left) and semisyn- thetic (right) settings. Performance degrades slowly and smoothly with reduced number of imputation steps. Error bars 11 s.e. across 500 Monte Carlo repetitions.
replaced with an MLP with fewer layers. Recall that the usual TS-Gen has an input layer, 3 hidden layers, and an output layer (Section B.2.1), adding up to 5 total layers, We compare regret for such variants of TS-Gen in Figure 11.
Regret Over Timesteps: Synthetic
Regret Over Timesteps: Semisynthetic
layer
layer (cieaaltt)
T5-Gen, 4 layer model
TS-Gen. layer model
40
TS- Cen, 3 layer model
-1-
T5 Gen. 3 layer model
model
100
30
TS-Gen, layer model Hlgistic
layer model Dogistic)
80
20-
60
40-
20-
100
200
300
100
200
400
500
Decision Times
Figure 11: Regret for TS-Gen with simpler sequence models in the synthetic (left) and semisyn- thetic (right) settings. Error bars are #1 s.e. across 500 bandit environments.
B.7 Softmax Sampling vs. Greedy
Here we compare Softmax Sampling as described Appendix B.3.2, with Greedy, and e-Greedy in Figure 12. Softmax Sampling is another bandit algorithm that uses Like e- --reedy, it "explores" while using Ps. but it does adequately handle uncertainty as TS-Gen does, as evidenced by the difference in regret.
Regret Over Timesteps: Synthetic
Regret Over Timesteps: Semisynthetic
80
TS-Gen
120
Greedy
60
Epsilon Greedy
100
Regret
Softmax
Regret
80
40
40
20
10
0
0
100
200
300
400
500
0
100
200
300
400
500
Decision Times
Decision
Times
Figure 12: Regret for Softmax Sampling vs Greedy vs e-Greedy in the synthetic (left) and semisyn- thetic (right) settings. Error bars are +1 s.e. across 500 bandit environments.
36

---

B.8
Compute resources
Offline pretraining. Pretraining a single pa for the semisynthetic setting took at most 12 hours; We use CPU cluster at Columbia GSB and request most 50GB of memory per job. The semisynthetic data generating process also involves evaluating two pre-trained text classifiers, and then caching their outputs and/or embeddings (DistilBERT embeddings + text classifier outputs in the semisynthetic setting); this was done once on a single GPU at negligible time cost (several minutes).
Online decision-making For online decision-making, we also GSB and for each job we request at most 10GB of memory. Below is a sample of decision-making time per timestep, across 20 sampled semisynthetic bandit tasks (10, 000 decisions total). Note that we cache the DistilBERT embedding representing the news article text so this is not included in the computation. each sampled bandit task, we compute the average per-timestep time in seconds for generation vs policy fitting; then, we report mean and variance of these quantities across the sampled bandit tasks, We write these times below as mean = standard deviation across the 10, 000 decisions.
CPU cluster Columbia
TS-Gen, using logistic policies: 3.1 + 0.5 seconds for generating f.. 0.01 + 0.02 seconds for policy fitting, 3.1 0.5 seconds total
TS-Gen, using MLP-based policies: 4.2 + 0.5 seconds for generating f,, 2.2 0.03 seconds for policy fitting, 6.4 = 0.5 seconds total
Neural Linear Thompson Sampling: 1.9 + 0.2 seconds total
B.9 Constrained policy classes
Algorithmic fairness is topic of general interest [Mehrabi and Wager, 2024, Mitchell et al., 2021], and fairness can be thought of as modeling constraint [Corbett-Davies et al., 2017]. Because our proposed method takes a policy class as an input, results can be immediately adapted to settings that require specific kinds of constraints, such as fairness balancing constraints.
As a simple example, we could enforce the constraint that at any given timestep f. a fitted policy r (t) must satisfy the condition that it would give specific treatment to approximately the same proportion of user contexts X, across two pre-specified groups. For example, these groups can be two sets of specific individuals, representatively drawn from the population, where each group selects individuals from a different geographic region, and where the groups are not related to contexts drawn in f.. This kind of fairness constraint essentially the notion of predictive parity [Verma and Rubin, 2018].
To implement such policy class, we would modify the policy fitting procedure in Line 4 in Algorithm 1 as follows: Letting fs be the imputed table, and letting G1 = .2,
be these predefined sets of user con- texts X, we would be solving n² = argmax, R Σ/= y, (Xeii)))))) subject to the constraint that Ni N 1fx((),)) = a} No EN (xx()2)) = a} < c for some chosen e > 0.
X1,,,,
and G2
B.10 Licenses
MIND news dataset We use the MIND news dataset [Wu et al., 2020]. Is under Microsoft Re- search License at https//waagiigggee...oodiiiiilet coddiisllbenddiittheet---eeaaaan It has an apache-2.0 license, with license and terms of use https://choeeellllcne.cc///ll/ns
www.ta/www.....ee/oooooooaeeeeee siceee Lieeneeeee___ee aatteioeeeinneeeeaaaaaaaaaa which we comply with. The terms of use are at http://wwww.iiicottttt..//oo//es
DistiBBERT Our semisynthetic sequence models use DistilBERT [Sanh al., 2019] from
Text classifiers for semisynthetic setting in the semisynthetic experiment setting. We use a sentiment classifier [Savani, 2022], accessed at https://huggingfice.oobbbbb classifier [Babakov et al., 2023), accessed https:/hugginggaee http:aannnnnaacccool//oiieee co/s-nip/roberta-baa caatbeerta-ass-ffmmllllee more-formalitt- ranker. Both models were obtained from huggingfacc.com. The sentiment classifier is not associated
We use text classifiers for the data generating process
and formality
37

---

with a paper and is under an Apache 2.0 license
2.0/,
which we comply with. The formality classifier associated with a paper, as cited, and is under cc-by-ne-sa-4.0 license https//ppb.ooooollenesscCCCcCcCC which we also comply with.
38


Reviews:
[{"review_id": "mPKZQsO20y", "reviewer": "Reviewer_ik19", "summary": "The authors propose a generative approach for contextual Thompson sampling by leveraging an offline pretrained model for imputing trajectories. The authors provide a meta-framework with (i) a task-dependent variable $Z$ and (ii) observed trajectories $(X_t^{a}, Y_t^{a})_{a \\in \\mathcal{A}}$ for imputing missing potential outcomes. Using the task-dependent vector $Z$ with observed histories to impute missing observations (e.g. potential outcomes, future observations, etc.), one samples tasks from a posterior sample and fit a policy $\\pi \\in \\Pi$ that solves the task on the imputed data. The authors provide methods of imputing data through autoregression and theoretical results relating regret with the quality of the imputed data. Synthetic experiments are provided to demonstrate the efficacy of this approach.", "questions": "* It still seems somewhat unclear why the design-based approach (viewing observations as fixed, but unobserved) helps with the interpretation of the approach provided by the authors. It would be nice to have a section discussing the distinction between this \"fixed-sample\" approach vs. distributional approach, and how this perspective results in the development and superior performance of the authors' approach. \n* Does the performance of this approach depend heavily on the author's choice of imputation architecture? It would be of interest to see how these approaches perform with generic imputation methods beyond the neural network proposed by the authors - using a neural network for the low-dimensional examples provided may be an \"easy\" task relative to having either (i) a more complex DGP and/or (ii) a less powerful model. While there is theory for the regret bounds based on the quality of imputations, it would be beneficial to see how empirical performance degrades (as bandit regret bounds often provide little guidance for the actual empirical performance of standard approaches.)", "limitations": "The authors adequately address their limitations (small synthetic experiments), and discuss how the computation cost may pose issues. They provide guidance for the latter issue at the end of Section 6.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n* The authors provide theoretical guarantees for their approach based on the quality of imputation, which provides nontrivial regret bounds. The authors provide context for these regret bounds relative to existing works, and show that they are competitive with existing optimal bounds.\n* The authors provide clear motivations and examples throughout their work. Even without familiarity with the meta-bandit literature, the example in Figure 1 provides clear exposition for applications of this work.\n* The authors provide practical guidelines for their method, including implementation details with modern imputation approaches (e.g. neural network models).  \n\n**Weaknesses**\n* As the authors note, the empirical validation of this approach is too limited to draw significant conclusions, even with the clear outperformance of their approach compared to other methods. It would be interesting to see how these approaches perform when $Z$ is less informative and/or the complexity of the synthetic data-generating process is increased.", "was_revised": true, "final_justification": "I maintain my positive score for the reasons discussed above in the review. The authors have answered the major questions posed.", "rebuttal": "Thank you for your helpful comments and questions!\n\n## (1) \"Design based\" perspective\n\n*It still seems somewhat unclear why the design-based approach (viewing observations as fixed, but unobserved) helps with the interpretation of the approach provided by the authors. It would be nice to have a section discussing the distinction between this \"fixed-sample\" approach vs. distributional approach, and how this perspective results in the development and superior performance of the authors' approach.*\n\nThanks for the comment. There are two relevant points, which we will make clearer in the revision.\n\n**(1a) We don't view the outcomes in the table $\\tau$ as *fixed*, but rather as drawn from a task distribution.** We first want to clarify that in contrast to design-based inference in which potential outcomes $\\tau$ are fixed, we view them as a random variable drawn from a \"meta-task\" distribution: $\\tau \\sim p^\\*$ (from Eq (1) in the draft). Our decision-making objective is to minimize the regret on average across sampled tasks $\\tau \\sim p^\\*$:\n$$\\\\mathbb{E} \\\\left[ \\\\frac{1}{T} \\\\sum\\_{t=1}^T \\\\left\\\\{ R(Y\\_t^{(\\\\pi^\\*(X\\_t; \\\\tau))}) - R(Y\\_t^{(A\\_t)}) \\\\right\\\\} \\\\right].$$\n\n**(1b) Our approach is distributional, and models potentially observable outcomes $\\tau$ instead of latent variables.** Standard Thompson Sampling explicitly specifies *latent* environment parameters, places a prior on them, and performs posterior sampling. Performing standard Thompson sampling with neural networks, requires approximate inference methods and can be computationally very challenging [Riquelme et al, Tran et al, Osband et al]. In contrast, our generative Thompson Sampling algorithm maintains a posterior distribution over the *missing outcomes* in the potential outcomes table $\\tau$. Practically, a motivation for our approach of directly modeling potentially observable outcomes $\\tau$ is that it easily accommodates autoregressive sequence modeling; in contrast, we cannot as easily train a sequence model to predict unobserved latent variables. Our theory reflects the benefit of directly modeling observable outcomes, since the pretraining sequence prediction loss function directly controls the regret.\n\n\n## (2) Effect of choice of imputation model on performance\n\nWe can interpret your question about how \"generic\" imputation methods perform in our setting in two ways. We provide an answer for each interpretation.\n\n**(2a) \"Generic imputation methods\" refers to standard methods for imputation used in the missing data literature.** There are a wide variety of imputation methods in the missing data literature, including mean imputation, nearest neighbor imputation, model-based imputation, and Bayesian multiple imputation. Many of the standard imputation methods in the machine learning and statistics literature are not directly appropriate. Our generative Thompson sampling algorithm requires a sequence model be used for imputation that (i) has low offline sequence prediction loss, and (ii) can effectively use high-dimensional task features $Z_\\tau$. So, while imputation is an old topic, the feasibility and promise of our approach hinges on recent advances in sequence modeling.\n\n**(2b) \"Generic imputation methods\" refers to sequence models for imputation that are simpler models than the neural network-based ones we use in our simulations.** Under this interpretation the main concern is understanding how our simulation results may change if we use simpler sequence models for imputation. To address this, we ran additional simulations in our synthetic setting and varied the complexity of the sequence model $p_\\theta$ used for imputation. In our original experiments, we used a $p_\\theta$ parameterized by a multi-layer perceptron (MLP) with three hidden layers; this model takes as input a vector summarizing the history $\\mathcal{H}_t$ and the current context $X_t$, and outputs a distribution over $Y_t$. We ran additional simulation results for an MLP with one hidden layer and a simple logistic regression model (MLP with 0 hidden layers). The results are summarized below, and we will add full results to the draft. \n\n**Table 1: Sequence loss and cumulative regret of TS-Gen in synthetic setting for different sequence imputation models $p_\\theta$.** All sequence models $p_\\theta$ are pretrained using the same set of previous task data and hyperparameter tuning procedure. The 3 layer MLP results match those already presented in the draft (Figure 5). Sequence loss $\\ell(p_{\\theta})$ corresponds to Eq (6) in the draft computed on the validation set; we report $\\frac{1}{T} \\ell(p_{\\theta})$, where we divide by $T=500$.\n\n| | Sequence loss $\\frac{1}{T} \\ell(p_{\\theta})$ | Regret t=100 | Regret t=200 | Regret t=300 | Regret t=400 | Regret t=500 |\n|:-----------------------|:------------------|------|------|------|------|------|\n| 3 layer MLP (original) | 0.437             | 8.204| 13.652| 17.978| 21.434| 24.578|\n| 1 layer MLP            | 0.438             | 8.580| 14.474| 19.192| 23.028| 26.502|\n| Logistic               | 0.482             | 8.666| 16.746| 23.776| 29.564| 34.710|\n\n\nIn the table above, observe that less expressive sequence models have higher (worse) sequence loss and worse regret. \n\n## References\n- Osband, Wen, Asghari, Dwaracherla, Ibrahimi, Lu, & Van Roy. *Approximate Thompson Sampling via Epistemic Neural Networks.* UAI 2023\n- Riquelme, Tucker, & Snoek. *Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Sampling.* ICLR 2018\n- Tran, Snoek, & Lakshminarayanan. *Practical uncertainty estimation and out-of-distribution robustness in deep learning.* NeurIPS Tutorial 2020"}, {"review_id": "ElbS6PsNcO", "reviewer": "Reviewer_iwTT", "summary": "This paper studies the meta contextual bandit problem where the learner aims to minimize the regret in learning a number of tasks (with ground truth distribution $p\\_*$). The learner first learns the distribution of the tasks and then performs online decision making.", "questions": "1. In Theorem 2, it does not appear very clear to me why $\\ell(p_\\theta)\\ge \\ell(p_*)$. Could the authors explain it?\n2. Theorem 1 seems to good to be true to me. The \"best-fitting\" policy is not quite restrictive in this paper, but if the policy incurs a constant gap compared with the optimal policy, then there should be a constant regret, instead of the $1/\\sqrt{T}$ regret in Theorem 1. Moreover, if we restrict the setting to a single task of a stochastic linear bandit, then $p_*$ is trivial, and the regret is $\\tilde O(d\\sqrt T)$ (combining Theorem 1 and Proposition 1), which does not seem to be able to be achieved by the simple algorithm in this submission.\n3. Does it make sense if $p_\\theta$ is learned online? If it does, can the constant term $\\sqrt{\\ell(p_\\theta)-\\ell(p_*)}$ be replaced by some sub-constant term?", "limitations": "I don't find any undiscussed limitations or potential negative societal impact of this work.", "rating": 2, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n1. Although I am unfamiliar with the literature, the problem formulation section is very informative, and the example of news recommendation system makes the problem setting very clear.\n2. The theoretical findings of this work are solid and discussed with sufficient details.\n\n## Weaknesses\n\n1. I would doubt whether the the algorithm is essentially Thompson sampling. Classical Thompson sampling typically involves constructing a distribution over the **reward function**, and the action is selected to maximize the sampled reward function. However, the action selection of the algorithm in this paper seems to depend on the policy fitting. I would agree with the discussion in the section of related works that the algorithm framework in this work is completely different from previous works, but I would hesitate to refer to this method as Thompson sampling.\n2. The algorithm design looks quite straightforward, and I cannot evaluate the novelty of the algorithm because the authors only compare it against the TS literature but not the meta contextual bandit literature. Although there is a comparison in the section of previous works, the comparison is not technical enough to show the technical contributions. It may seem to me that this paper is some simple techniques disguised in a complex framework.\n\nI find the technical findings of this paper quite solid, but currently score this submission as \"reject\" mostly due to my unfamiliarly of the literature and inability to evaluate the contributions of this paper. As the proof is quite short, I am unsure about the contributions of the proof techniques either. I would be happy to discuss the contributions of this submission with other reviewers and the AC.", "was_revised": true, "final_justification": "I would still hesitate to recommend accepting this work, after the discussion with other reviewers. My major concerns are listed in the [official comment](https://openreview.net/forum?id=Fqsl9IfbfJ&noteId=P9uHPy29kT).\n\nFor point 1, I think the setting is reasonable but probably not very practical. For point 2, I think the lack of a ground truth policy is a major issue.", "rebuttal": "Thank you for your comments and questions!\n\n## (1) Relationship to Thompson Sampling (TS)\nIn the revision, we’ll clarify the following points and include Proposition 2 and its proof, which shows our algorithm performs probability matching, i.e., Thompson Sampling.\n\n**Thompson Sampling (TS) is probability matching and our algorithm uses probability matching.** TS selects actions according to the posterior probability that they are optimal. In the bandit/RL literature, the *\"probability matching\"* property is commonly taken to be the abstract definition of TS. When generalizing TS to MDPs [Osband et al], the algorithm samples an MDP from the posterior, picks the optimal policy for that sampled MDP, and applies that policy to the states encountered. Our algorithm fits this framework: we sample a dataset $\\hat{\\tau}$ from the posterior, compute the best-fitting policy $\\pi^*(\\cdot; \\hat{\\tau})$, and act according to that policy.\n\nWe show in Proposition 2 (below) that, when using the correct model $p^\\*$, our algorithm satisfies: $\\mathbb{P}(A_t = a \\mid \\mathcal{H}_t) = \\mathbb{P}(\\pi^\\*(X_t; \\tau) = a \\mid \\mathcal{H}_t).$ Above, $\\mathbb{P}(\\pi^\\*(X_t; \\tau) = a \\mid \\mathcal{H}_t)$ which is probability matching with respect to the oracle policy. This aligns with the abstract definition of TS.\n\n**Proposition 2.** *Algorithm 1 with imputation model $p^\\*$ implements Thompson Sampling (probability matching), i.e., the following holds with probability $1$:*\n$$\\\\mathbb{P}(A\\_t = a \\\\mid \\\\mathcal{H}\\_t) \n= \\\\mathbb{P}(\\\\pi^*(X\\_t; \\\\tau) = a \\\\mid \\\\mathcal{H}\\_t).$$\n\n*Proof:* Recall that Algorithm 1 selects actions as follows: \n$\\\\mathbb{P}(A\\_t = a \\\\mid \\\\mathcal{H}\\_t) = \\\\mathbb{P}(\\\\pi^*(X\\_t; \\\\hat{\\\\tau}) = a \\\\mid \\\\mathcal{H}\\_t).$\n\nSince $\\\\hat{\\\\tau} \\sim p^*( \\tau \\in \\cdot \\mid \\mathcal{H}_t)$ and from Eq (1) $\\\\tau \\\\sim p^\\*$, the distributions of $\\tau$ and $\\hat{\\tau}$ are equal given $\\mathcal{H}_t$. Hence, with probability $1$ for any $j$:\n$\\\\mathbb{P}(\\\\hat{\\tau} = j \\\\mid \\\\mathcal{H}\\_t) = \\\\mathbb{P}(\\\\tau = j \\\\mid \\\\mathcal{H}\\_t).$\n\nThe above implies that $\\\\mathbb{P}(\\\\pi^\\*(X\\_t; \\\\hat\\\\tau) = a \\\\mid \\\\mathcal{H}\\_t) = \\\\mathbb{P}(\\\\pi^\\*(X\\_t; \\\\tau) = a \\\\mid \\\\mathcal{H}\\_t).$ Combining the above statements gives the result. $\\square$\n    \n**Our algorithm generalizes TS** by replacing the true reward-maximizing policy with a best-fitting policy under a given policy class $\\Pi$. A more \"standard\" TS algorithm is recovered when the best-fitting policy is correctly specified. We will clarify these points in the revision.\n\n\n## (2) Novelty and Relationship to Meta-Bandits\n*The algorithm design looks quite straightforward, and I cannot evaluate the novelty...*\n\nThe algorithm we present is \"straightforward\" by design. The real novelty of our work is that we are able to provide theoretical guarantees, even when our algorithm takes advantage of neural network sequence models.\n\n**(2a) Relationship to meta-bandit literature.** Our work differs from much of the meta-bandit literature in two key ways (see Related Works line 306 for further details). To summarize: \n- **Model complexity.** Much of the meta-bandit literature focuses on algorithms that use simple models, e.g., linear or parametric models [Cella et al, Kveton et al]. Our approach uses autoregressive sequence models to capture complex dependencies across task-features $Z_\\tau$, contexts $X$, and actions $A$, enabling richer reward modeling.\n- **Regret guarantees.** Unlike typical meta-bandit algorithms, our regret bounds hold under misspecification of both the Bayesian model ($p_\\theta$) and the policy class $\\Pi$.\n\nWe also include comparisons to a meta-TS baseline in experiments (TS-Linear fitted with history, Section 6 and Appendix B.3.2).\n\n**(2b) Novelty.** Even beyond the meta-bandit literature, our theoretical guarantees have no precedent in the literature (see Section 3.1, especially line 172). To clarify further:\n\n- **Entropy-based regret bound.** Our regret bound depends on $H (\\mathbf{\\pi}^*(X_{1:T}) \\mid Z_\\tau)$ the entropy of oracle policy decisions, rather than the entropy of latent parameters [Neu et al, Dong et al]. We boud this term by the VC/Natarajan dimension of the policy class $\\Pi$, which is new in the information-theoretic stochastic bandit literature.\n- **Accommodates complex reward models.** There is great interest in algorithms that generalize TS beyond parametric models and still have theoretical guarantees. Our work does this, as we allow our outcome (imputation) model to be highly expressive and provide a regret bound (Theorem 1) that does not directly depend on the complexity of this imputation model. We also characterize the regret when the imputation model $p_\\theta$ is misspecified (Theorem 2).\n- **Flexible policy classes.** Many bandit works that assume the true policy class $\\Pi$ is known (e.g. linear bandits). Our work takes a flexible approach: (i) if the true policy class is known, one can choose that to be the \"oracle\" policy, and (ii) if the true policy class is unknown or misspecified, our regret bound provides a best-in-(policy)-class guarantee.\n- **Dependence on pretraining loss.** It is common in TS regret bounds to assume the true Bayesian model (prior and likelihood) is known by the algorithm. For our algorithm, when an approximate $p_\\theta$ model is used, the regret of our algorithm is bounded in terms of the pretraining sequence loss $\\ell(p_\\theta)$, which can be optimized offline. This connects to recent work on analyzing TS under misspecified priors for parametric bandits  [Simchowitz et al, Liu et al].\n\n\n\n## (3) In Theorem 2 questions\n### **(3a) Loss inequality** \n*In Theorem 2, it does not appear very clear to me why $\\ell(p_\\theta) \\geq \\ell(p^\\*)$.*\n\nThe inequality $\\ell(p_\\theta) \\geq \\ell(p^\\*)$ follows directly from Lemma 1 (Appendix A.3.1), which shows:\n$$\\ell(p\\_\\theta) = \\ell(p^\\*) + D\\_{\\text{KL}} \\left( p^\\* \\big(X\\_{1:T}, \\{ Y\\_{1:T}^{(a)} \\}\\_{a \\in \\mathcal{A}\\_\\tau} \\mid Z\\_\\tau \\big) \\| ~p_\\theta \\big(X\\_{1:T}, \\{ Y\\_{1:T}^{(a)} \\}\\_{a \\in \\mathcal{A}\\_\\tau} \\mid Z\\_\\tau \\big) \\right).$$ Since the KL divergence is non-negative, the the result follows. Intuitively, this inequality is saying that $p^\\*$ minimizes the sequence prediction loss, and any misspecification in $p_\\theta$ results in higher loss.\n\n### **(3b) Choice of policy class**\n*The \"best-fitting\" policy is not quite restrictive in this paper, but if the policy incurs a constant gap compared with the optimal policy, then there should be a constant regret...*\n    \nIt is common in the bandit literature to assume that the correct policy class is known (e.g., the expected reward is linear in the context). **Our theoretical results cover this case in which the correct policy class is known to the algorithm design a priori, and also provide theoretical results when this policy class is misspecified.** As discussed in point (1) above, we make a substantial generalization of standard TS in this work by replacing the \"optimal policy\" with an abstracted notion of a \"best-fitting oracle policy\" $\\pi^\\*(X_t; \\tau)$, which is the policy you would want to want to fit if you had the entire dataset $\\tau$. It's appropriate to view our method as a \"generalized\" TS that targets learning a \"best-fitting oracle policy\", which may or may not be the one that exactly maximizes the reward in each context, depending on the choice of the oracle policy fitting procedure $\\pi^\\*(\\cdot)$. A more \"standard\" TS algorithm is recovered when the best-fitting policy is correctly specified.\n\n### **(3c) Regret achievable by our algorithm** \n*Theorem 1 seems to be good and true to me... If we restrict the setting to a single task of a stochastic linear bandit, then $p^\\*$ is trivial, and the regret is $O(d \\sqrt{T})$ (combining Theorem 1 and Proposition 1), which does not seem to be able to be achieved by the simple algorithm in this submission.*\n\nWe did not fully understand your comment. As you state, *\"Theorem 1 seems to be good and true to me\"* and the regret bound follows from Theorem 1 and Proposition 1. We clarify below:\n- Theorem 1 proves the following upper bound the *cumulative* regret: $\\sqrt{\\frac{1}{2} |\\mathcal{A}| T \\cdot H(\\mathbf{\\pi}^*(X_{1:T}) \\mid Z_\\tau) }$.\n- Proposition 1, shows that $H(\\mathbf{\\pi}^*(X_{1:T}) \\mid Z_\\tau) = O(d \\cdot \\log(T \\cdot |\\mathcal{A}|))$, where $d$ is the Natarajan dimension of the policy class $\\Pi$. \n- In stochastic linear bandits with policies $\\text{argmax}_a ~\\beta^\\top \\varphi(X_t, A_t)$ for some $\\beta \\in \\mathbb{R}^d$ have Natarajan dimension $d$ (Appendix A.5).\n- The resulting bound is $O \\sqrt{ |\\mathcal{A}| T d \\log(T |\\mathcal{A}|) }$\n\n\n## (4) Learning $p_\\theta$ online\n*Does it make sense if $p_\\theta$ is learned online?*\n\nWe assume access to past bandit tasks (offline meta-bandit), as in prior work [Cella et al, Hong et al]. To extend to a sequential task setting, one could retrain $p_\\theta$ online (as new tasks arrive), which we expect would decrease the gap $\\ell(p_\\theta) - \\ell(p^*)$, and thus the regret bound. Future work could explore this retraining more formally.\n\n## References\n- Cella et al. *Meta-learning with stochastic linear bandits.* ICML 2020\n- Dong & Van Roy. *An information-theoretic analysis for TS with many actions.* NeurIPS 2018\n- Hong et al. *TS with a mixture prior.* AISTATS 2022\n- Kveton et al. *Meta-TS.* ICML 2021\n- Liu & Li. *On the prior sensitivity of thompson sampling.* ALT 2016\n- Neu et al. *Lifting the information ratio: An information-theoretic analysis of TS for contextual bandits.* NeurIPS 2018\n- Osband, Russo, & Van Roy. *(More) Efficient Reinforcement Learning via Posterior Sampling.* NeurIPS 2013\n- Simchowitz et al. *Bayesian decision-making under misspecified priors with applications to meta-learning.* NeurIPS 2021"}, {"review_id": "egY7E4gvhu", "reviewer": "Reviewer_y7h5", "summary": "The paper describes a new bandit algorithm that addresses the cold start problem by meta learning a generative model with offline data regarding a multitude of tasks. The pre-trained generative model is used to sample complete task data based on which a policy is fitted.  The paper claims that this approach is equivalent to contextual Thompson Sampling.  Some regret bounds are provided.  Experiments show that the proposed generative thompson sampling technique has lower regret than various baselines in a proof of concept that involve one synthetic problem or one semi-synthetic problem.", "questions": "* Can you prove that the proposed algorithm is equivalent to Thompson sampling?\n* What are the advantages of proposed algorithm over Thompson sampling?\n* In the new recommendation example, why does Z correspond to a news article?  It seems more natural to think of each action as a news article.  If the actions are not news articles, then what are the actions?", "limitations": "The paper acknowledges that the empirical evaluation is limited to a synthetic problem and a semi-synthetic problem.  It also acknowledges that the proposed algorithm will be computationally intensive in complex tasks.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n* The proposed approach is novel and interesting.\n* Regret bounds are provided, though I did not verify them.\n* The paper is well written and generally clear.\n\nWeaknesses:\n* Claim that the proposed algorithm is equivalent to Thompson Sampling is not substantiated.  Please provide a proof.\n* It is not clear what are the advantages of the proposed algorithm over regular Thompson Sampling.  Let $p_\\theta(Y|X,A)$ be a distribution parameterized by $\\theta$ over the reward $Y$ when executing $A$ in context $X$.  Regular Thompson Sampling would learn a posterior distribution over $\\theta$, then sample $\\theta$ based on which a reward $Y$ is sampled for each action $A$ and the action with the highest reward is selected.  Since Bayesian learning is rarely tractable, estimating a posterior over $\\theta$ will need to be approximated.  In contrast, the proposed approach learns $p_\\theta(\\tau)$ by maximum likelihood.  If the proposed approach is indeed equivalent to Thompson sampling, then it must learn to implicitly approximate Bayesian learning based on many tasks.  This seems harder than regular Thompson sampling in the sense that it will need a lot of data just to learn to approximate Bayesian learning.  \n* The experiments do not compare to regular Thompson sampling that learn a posterior distribution over $\\theta$.\n* The empirical evaluation is limited to one synthetic problem and one semi-synthetic problem", "was_revised": true, "final_justification": "I read several times the author rebuttal.  I really appreciate all the effort put into responding to my questions.  Nevertheless, my concerns remain.  First the paper claims that the proposed algorithm is equivalent to Thompson sampling, but a formal proof is lacking.  The authors explained a good direction for constructing a proof.  I appreciate the definition of Thompson sampling as probability matching.  The authors provide informal arguments to suggest that their algorithm provides a form of probability matching.  But a proper proof should include a line by line derivation.  In particular, Proposition 2 lacks a proof that the statement holds with probability 1. Also, the argument that it is okay to replace $\\tau$ by $\\hat{\\tau}$ because they are sampled from the same distribution lacks a mathematical derivation.\n\nI do not agree with the claim that TS-Gen generalizes Thompson sampling.  The fact that the optimal policy is replaced by a best fitting policy simply means that an approximation is introduced.  All Thompson sampling algorithms introduce approximations because estimating the posterior is intractable.  If we follow the logic that an approximation leads to a generalization then all previous Thompson sampling techniques should also be considered as different forms of generalized Thompson sampling, but this is misleading.\n\nMy concern about the fairness of the empirical evaluation remains. The paper proposes TS-Gen which does two things: it provides a new way of estimating the posterior and it uses a transformer based seq2seq model.  The paper claims that their new posterior estimation is the key for low regret, but the transformer model could also be a major factor.  This is why I asked the authors to provide a fair comparison with other techniques that use the same transformer model instead of a linear model.  In their last response below, the authors generated new results that use the same transformer model with a neural linear estimation technique. Technically, this is not doing probability matching because the posterior is only computed with respect to the last layer (instead of all parameters), but we can view this as an approximation and therefore I am fine with this.  The main issue is that this baseline is only fed with a single $(X_t, Z)$ while their approach is also fed with the history.  Again, this is unfair. \n\nOverall, the work is promising, but it is not ready for publication.", "rebuttal": "Thank you for your helpful comments and questions!\n\n## (1) Experiments comparing to standard Thompson Sampling\n*The experiments do not compare to regular Thompson sampling (TS) that learn a posterior distribution over $\\theta$.*\n\nWe **do** compare to standard linear Thompson Sampling (TS) in our simulations. See *\"TS-Linear\"* and *\"TS-Linear (fit w history)\"* in Figure 5 and Section 6. These use Gaussian priors: one non-informative and one fit using historical data (Appendix B.3.2).\n\n## (2) Relationship to Thompson Sampling\n*Can you prove that the proposed algorithm is equivalent to Thompson sampling?*\n\nThank you for the question. In the revision we will add the Proposition 2 (see below) and associated proof. We first provide some relevant context on Thompson Sampling and our work, which we will clarify in our revision.\n\n**(2a) Thompson Sampling is probability matching.** \"Standard\" Thompson sampling (TS) uses *probability matching* to select actions, which means *TS selects an action $a$ according to the posterior probability that action $a$ is optimal.* For example a linear contextual bandit TS algorithm selects actions according to the posterior probability that action $a$ is optimal with context $X_t$ (note $X_t \\\\in \\\\mathcal{H}_t$):\n$$\\\\mathbb{P}(A\\_t = a \\\\mid \\\\mathcal{H}\\_t) \n= \\\\mathbb{P}( \\\\text{argmax}\\_{a'} ~\\\\beta^{\\\\top} \\\\varphi(X\\_t, a') = a \\\\mid \\\\mathcal{H}\\_t).$$  Above, the probability on the right hand side averages over the posterior distribution of the latent parameter $\\\\beta$ given $\\\\mathcal{H}_t$.\n    \nIn the bandit and RL literature, the *\"probability matching\"* property is commonly taken to be the abstract definition of TS. When generalizing TS to MDPs [Osband et al], the algorithm samples an MDP from the posterior, picks the optimal policy for that sampled MDP, and applies that policy to the states encountered. Our algorithm fits this framework: we sample a dataset $\\hat{\\tau}$ from the posterior, compute the best-fitting policy $\\pi^*(\\cdot; \\hat{\\tau})$, and act according to that policy. We formalize the probability matching property of our algorithm next.\n\n**(2b) Our algorithm uses probability matching to select actions.** Our TS-Gen algorithm uses probability matching to select actions:\n$$\\\\mathbb{P}(A\\_t = a \\\\mid \\\\mathcal{H}\\_t) = \\\\mathbb{P}(\\\\pi^*(X\\_t; \\\\tau) = a \\\\mid \\\\mathcal{H}\\_t)$$ \n\nAbove, $\\\\mathbb{P}(\\\\pi^*(X\\_t; \\\\tau) = a \\\\mid \\\\mathcal{H}\\_t)$ denotes the posterior probability that under the best-fitting policy on the true dataset $\\tau$ that $a$ is the best action; the probability averages over randomness over unobserved outcomes in $\\tau$. The above property holds because through our sampling procedure, the conditional distributions of the true $\\tau$ and imputed $\\hat{\\tau}$ datasets are the same conditional on the history $\\mathcal{H}_t$; see Proposition 2 below for the proof.\n\nWith approximate models $p_\\theta$, our algorithm performs probability matching under the $p_\\theta$-defined distribution; Theorem 2 quantifies the regret due to this misspecification.\n\n**Proposition 2.** *Algorithm 1 with imputation model $p^\\*$ implements Thompson Sampling (probability matching), i.e., the following holds with probability $1$:*\n$$\\\\mathbb{P}(A\\_t = a \\\\mid \\\\mathcal{H}\\_t) \n= \\\\mathbb{P}(\\\\pi^*(X\\_t; \\\\tau) = a \\\\mid \\\\mathcal{H}\\_t).$$\n\n*Proof:* Recall that Algorithm 1 selects actions as follows: \n$\\\\mathbb{P}(A\\_t = a \\\\mid \\\\mathcal{H}\\_t) = \\\\mathbb{P}(\\\\pi^*(X\\_t; \\\\hat{\\\\tau}) = a \\\\mid \\\\mathcal{H}\\_t).$\n\nSince $\\\\hat{\\\\tau} \\sim p^*( \\tau \\in \\cdot \\mid \\mathcal{H}_t)$ and from Eq (1) $\\\\tau \\\\sim p^\\*$, the distributions of $\\tau$ and $\\hat{\\tau}$ are equal given $\\mathcal{H}_t$. Hence, with probability $1$ for any $j$:\n$\\\\mathbb{P}(\\\\hat{\\tau} = j \\\\mid \\\\mathcal{H}\\_t) = \\\\mathbb{P}(\\\\tau = j \\\\mid \\\\mathcal{H}\\_t).$\n\nThe above implies that $\\\\mathbb{P}(\\\\pi^\\*(X\\_t; \\\\hat\\\\tau) = a \\\\mid \\\\mathcal{H}\\_t) = \\\\mathbb{P}(\\\\pi^\\*(X\\_t; \\\\tau) = a \\\\mid \\\\mathcal{H}\\_t).$ Combining the above statements gives the result. $\\square$\n    \n**(2c) Our work is a generalization of Thompson Sampling (TS).** We make a substantial generalization of standard TS in this work by replacing the \"optimal policy\" with an abstracted notion of a \"best-fitting oracle policy\" $\\pi^\\*(X_t; \\tau)$, which is the policy you would want to use if you had the entire dataset $\\tau$. It is appropriate to view our method as a \"generalized\" TS that targets learning a \"best-fitting oracle policy\". A more \"standard\" TS algorithm is recovered when the best-fitting policy is correctly specified. We will clarify these points in the revision.\n\n\n## (3) Advantages of our algorithm\n*What are the advantages of the proposed algorithm over Thompson sampling?*\n\nWe note several advantages and discuss what revisions we plan to make below.\n\n**(3a) Learning the Bayesian model from data.** Recall that in our news recommendation problem setting, we need learn across tasks how to form prior beliefs conditioned on article text, i.e., fit an appropriate Bayesian model given high-dimensional text features. A strength of our approach is that an appropriate Bayesian model can be learned through standard offline sequence prediction and gradient descent (see empirical Bayes example below); furthermore, our regret bound (Theorem 2) shows that minimizing this loss provably controls regret. This benefit of being able to leverage existing methods for training offline sequence models relies on the fact taht we directly modeling potentially observable outcomes $\\tau$, rather than unobserved latent variables (as is standard in Thompson Sampling). Below we discuss an example of how our pretraining procedure relates to standard methods for fitting Bayesian models from data, i.e., empirical Bayes (Type-II maximum likelihood). We mention this connection to empirical Bayes in the draft briefly in line 209, but will expand on this point in the revision.\n    \n**Empirical Bayes connection:**  *As a simple example, in a Beta-Bernoulli model:*\n$$\\mu_a \\sim \\text{Beta}(\\alpha^\\*, \\beta^\\*), \\qquad Y_1^{(a)}, Y_2^{(a)}, \\dots \\big| \\, \\mu_a \\overset{iid}{\\sim} \\text{Bernoulli}(\\mu_a).$$ *The above model is associated the following class of sequence models (posterior predictives):*\n$$p_{\\theta} \\big( Y_{t+1}^{(a)} = 1 \\mid Y_{1:t}^{(a)} \\big) =  \\frac{ \\alpha  + \\sum_{i=1}^{t} Y_{i}^{(a)} }{ \\alpha + \\beta + t } \\quad \\text{where} \\quad \\theta = (\\alpha, \\beta).$$ *The true sequence model $p^\\*$ is $p_\\theta$ evaluated at $\\theta = (\\alpha^\\*, \\beta^\\*)$. Minimizing our negative log likelihood sequence loss from (6) is equivalent to maximizing the marginal likelihood of the data, the criterion that is used in Empirical Bayes to fit prior distributions to observed data.*\n\n**(3b) Agnostic approach that accommodates complex outcome/reward models.** There has been a lot of interest in being able to effectively generalize Thompson Sampling beyond parametric models to settings that require more complex outcome/reward, and at the same time still provide theoretical guarantees. Our work does this, as we allow our outcome (imputation) model, $p^*$, to be exceptionally complex and provide a regret bound (Theorem 1) that does not directly depend on the complexity of this imputation model (but instead on the complexity of the chosen \"oracle\" policy discussed next). Moreover, our algorithmic procedure (which involves fitting and sampling from a sequence model), is computationally more straightforward than many methods for Thompson Sampling with Bayesian neural networks [Riquelme et al, Tran et al].\n\n**(3c) Tolerance for general \"oracle\" policy construction.** Many works on bandit algorithms assume that the true policy class is known (e.g., the expected reward is linear in the context, so the true policy class is based on a linear model). Our work is a generalization of this standard approach in that (i) if the true policy class is known, one can choose that to be the \"oracle\" policy, i.e., the definition of $\\pi^*$ from line 137, and (ii) if the true policy class is unknown or the policy class is misspecified, our regret bound provides a best-in-(policy)-class guarantee.\n\n\n## (4) Z in the news recommendation setting\n*In the new recommendation example, why does Z correspond to a news article? It seems more natural to think of each action as a news article. If the actions are not news articles, then what are the actions?*\n\nWe use $\\\\mathcal{A}\\_\\\\tau$ to denote the action space and $Z\\_\\\\tau = \\\\{ Z\\_\\\\tau^{(a)} \\\\}_{a \\\\in \\\\mathcal{A}\\_\\\\tau}$ to denote *action features*. In the news recommendation setting an action $a \\\\in \\\\mathcal{A}\\_\\tau$ refers to a particular news article and $Z\\_\\\\tau^{(a)}$ refers to features associated with that news article; in general, features could include the news outlet, article type, article text, and article photos. We separate actions $a$ and their associated features $Z\\_\\\\tau^{(a)}$ since for some choices of features (e.g. news outlet) two distinct actions may have the same action feature (e.g. two articles are distinct but come from the same news outlet).  Note though that even with action features $Z\\_\\\\tau^{(a)}$ there is still residual uncertainty in how well that action will perform (e.g., two articles with similar features $Z\\_\\\\tau^{(a)}$ could perform differently), which requires exploration to learn. We will make these points clearer in the revision.\n\n## References\n- Tran, Snoek, & Lakshminarayanan. *Practical uncertainty estimation and out-of-distribution robustness in deep learning.* NeurIPS Tutorial 2020.\n- Osband, Russo, & Van Roy. *(More) Efficient Reinforcement Learning via Posterior Sampling.* NeurIPS 2013\n- Riquelme, Tucker, & Snoek. *Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Sampling.* ICLR 2018"}, {"review_id": "xLX9xtPBwj", "reviewer": "Reviewer_BmbW", "summary": "This paper proposes a novel framework for Thompson Sampling (TS) in stochastic contextual bandit settings, where a generative model is used to impute all missing data points (with future context vectors), rather than sampling from the posterior distribution (classical approach). Notably, the quality of the generative model affects the regret bound only as an additive term, rather than multiplicatively influencing the leading term. Leveraging this generative approach, the authors also explore the potential for incorporating meta-learning, which could offer practical benefits in real-world applications. Experimental results demonstrate strong empirical performance. However, a significant limitation is the computational cost of TS-Gen, which is approximately 1000 times higher than that of simpler baselines, posing a serious bottleneck for deployment in real-time or resource-constrained online settings.", "questions": "### Questions\n\n1. Although Proposition 1 provides an upper bound on the entropy in terms of the Natarajan dimension $d$ and $\\log ( T  |A_\\tau|))$, to the best of my knowledge, the Natarajan dimension can be quite large when the function class is parameterized by neural networks.\nFor example, Jin (2023) shows that the Natarajan dimension of a neural network class with ReLU activations is upper bounded by  $O(|A_\\tau| \\cdot p^2)$, where $p$ denotes the number of parameters of neural networks. This suggests that the function class may need to be sufficiently restricted to ensure its Natarajan dimension remains below $O(\\sqrt{T})$, in order for the resulting bound to be meaningful. Since my understanding might be wrong, I would appreciate the authors' opinion on this issue.\n\n2. In the current formulation, the algorithm requires the time horizon $T$ as input, since the generative model is used to impute all future outcomes in advance. If $T$ is unknown, a natural alternative would be to generate a fixed number of future outcomes or to adopt a doubling strategy for the number of imputations. Could the authors clarify whether such modifications would affect the theoretical guarantees or empirical performance of the proposed method?\n\n---\nJin, Ying. \"Upper bounds on the natarajan dimensions of some function classes.\" 2023 IEEE International Symposium on Information Theory (ISIT). IEEE, 2023.\n\n### Comments\n\n1. two overlapped citations: see Lines 467 and 468.\n2. In appendix A.3., Line 982: Lemma 2 -> Jensen's inequality?", "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": "### Strength\n\n* This paper is well-written and figures are very helpful to capture the overall idea.\n* A general framework that can be used for single/multiple tasks, several imputation models including neural networks, and various settings such as non-stationary settings.\n* The regret bound clearly separates the quality of generative model from the algorithmic performance and it is confirmed in experiments. The additive dependence on model quality is notable since it would scale as $\\log T$ under mild regularity conditions as the authors mentioned.\n* By relying on supervised learning rather than probabilistic modeling, TS-Gen broadens the applicability of TS to settings where approximate posteriors are difficult to compute or unavailable.\n\n### Weakness\n\nAs acknowledged by the authors, TS-Gen requires significantly more computation than baseline methods (up to 1000× slower in some cases). This scalability issue is a serious bottleneck in online decision-making applications where response time is critical. One possible remedy could be to use an (adaptive or increasing) batch size strategy, where the generative model re-samples the missing outcomes only when a new batch threshold is reached, incorporating the currently observed outcomes and contexts.", "was_revised": true, "final_justification": "I have also read the discussion with the other reviewers and identified several points that could help improve the paper’s quality. While I consider the paper borderline, some of the authors’ responses (to other reviewers) seem reasonable to me, so I think it is acceptable. Therefore, I will keep my score unchanged but lower my confidence.", "rebuttal": "Thank you so much for your thoughtful comments and questions!\n\n## (1) Natarajan dimension\nWe agree that known bounds on the Natarajan dimension of general neural networks can indeed be large, as highlighted by Jin (2023). This critique seems to apply to nearly all learning theory with neural networks. We have several additional comments related to this:\n\n1. **Policy class versus imputation model.** Our Proposition 1 upper bounds $H(\\\\mathbf{\\\\pi}^\\\\ast(X_{1:T}) \\\\mid Z\\_\\\\tau)$ by \n$O(d \\\\cdot \\\\log(T \\\\cdot |\\\\mathcal{A}|\\_\\\\tau))$, where $d$ is the Natarajan dimension of the policy class $\\\\Pi$. Recall that the policy class $\\\\Pi$ is the class of functions that take in $X$ and output an action in $\\\\mathcal{A}$. We first want to clarify that the *policy class* $\\\\Pi$ is distinct from the sequence model, $p^\\\\ast$ or $p_\\\\theta$, used for imputation. After we use the sequence model to impute missing outcomes, we fit a policy $\\\\pi \\\\in \\\\Pi$ that is used to select actions. In summary, the complexity of the sequence model does not impact the bound in Theorem 1 at all, only the complexity of the policy class $\\\\Pi$ does.\n\n2. **Policy class flexibility.** While the Natarajan dimension can be large for arbitrary neural networks, our framework offers the flexibility to choose different policy classes $\\Pi$ depending on how complex of a model is needed to effectively model $R(Y) \\mid X, A$. If the policy class $\\Pi$ is simpler, the Natarajan dimension would be significantly smaller. For example, in our synthetic simulation results, one version of our algorithm uses a logistic regression-based policy class.\n\n3. **Scaling with $T$.** Each policy $\\pi \\in \\Pi$ is a function that maps any context $X$ to an action in $\\mathcal{A}$. Generally, *we do not expect the Natarajan dimension of the policy class $\\Pi$ to depend on $T$.* For example, linear policy classes of the form $\\text{argmax}_{a} \\theta^\\top \\varphi(x, a)$ for $\\theta \\in \\mathbb{R}^p$ have Natarajan dimension $p$ (see Appendix A.5 for additional details and the derivation of the Natarajan dimension for other special case policy classes). In summary, for reasonable policy classes the Natarajan dimension does not grow with $T$, making the overall regret bound meaningful. We will make this point more overt in the revision.\n\n## (2) Unknown horizon $T$\nIn this work, we consider a finite horizon setting in which $T$ is known and used by the algorithm. A simple way to modify our algorithm when $T$ is not known is to have the sequence model only impute past missing outcomes and a fixed number of timesteps, $m$, into the future. Although our theoretical results do not apply to this version of algorithm, we find empirically that the performance of the algorithm for sufficiently large $m$ is often like that of our original algorithm, which imputes outcomes over the entire horizon $T$ (see results table below); performance degrades smoothly as $m$ is decreased. Furthermore, this modified approach increases the computational efficiency of our method in the online decision-making phase. We will include these additional results and discussion in our revision.\n\n**Table 1: Cumulative regret of TS-Gen in synthetic setting, when imputing up to m=100, 200, 300, 500 timesteps steps into the future.** Note that when $m=500$, this matches our original TS-Gen algorithm and synthetic setting results in Figure 5 in the draft.\n\n|  | t=100 | t=200 | t=300 | t=400 | t=500 |\n|:--------------------------|--------|--------|--------|--------|--------|\n| TS-Gen, m=500=T (original) | 8.20   | 13.65  | 17.98  | 21.43  | 24.58  |\n| TS-Gen, m=300              | 8.43   | 13.96  | 18.47  | 21.90  | 25.14  |\n| TS-Gen, m=200              | 8.60   | 14.31  | 18.63  | 22.27  | 25.48  |\n| TS-Gen, m=100              | 9.21   | 14.97  | 19.44  | 23.04  | 26.28  |\n\n \n## (3) Computational efficiency\nWhen implementing Thompson Sampling with neural network models, a variety of methods for approximate posterior inference methods have been proposed. Many of these existing methods (e.g., MCMC and variational inference [Riquelme et al]) are computationally costly. Our approach, as presented in Algorithm 1, is also computationally costly but we believe it can be extended to be more computationally efficient in practice:\n1. **Truncated generation.** As discussed above, one can truncate the generation length to at most a fixed number of $m$ timesteps into the future. For sufficiently large $m$, we find that the decision-making performance is comparable to doing the full imputation (see Table 1 above).\n2. **Batch updating.** Following your comment, a well-known practical strategy to increase the computational efficiency of bandit algorithms in practice is to update after taking a batch of action, rather than updating the algorithm after every action taken [Yi et al, Yu et al].\n3. **Policy distillation.** In the draft we also mention that one could explore policy distillation methods [Czarnecki et al] to increase computational efficiency, as these methods are designed to transfer knowledge from one policy to another, typically to speed up computation.\n\n## (4) Typos\n- Thank you for pointing out the overlapped citations on Lines 467 and 468. We will correct this!\n- We appreciate the careful reading of the appendix. You are correct that in Appendix A.3, Line 982, the inequality should indeed refer to Jensen's inequality, not Lemma 2. We will fix this in the revision.\n\n## References\n- Czarnecki, Pascanu, Osindero, Jayakumar, Swirszcz, & Jaderberg. *Distilling policy distillation.* AISTATS 2019\n- Gelman, Carlin, Stern, Dunson, Vehtari, Rubin. *Bayesian data analysis.* Chapman and Hall/CRC 1995\n- Riquelme, Tucker, & Snoek. *Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Sampling.* ICLR 2018\n- Yi, Wang, He, Chandrasekaran, Wu, Heldt, Hong, Chen, Chi. *Online Matching: A Real-time Bandit System for Large-scale Recommendations.* RecSys 2023\n- Yu & Oh. *Optimal and Practical Batched Linear Bandit Algorithm.* ICML 2025"}]

Author rebuttal:

Meta-review: A generative TS algorithm is proposed, which meta-learns the distribution of tasks and uses the learned model to generate both future and counterfactual outcomes, which are used for optimizing the action-selection policy. The generation model is updated every time a new observation is acquired, and the process of optimizing the action-selection policy, observing rewards, and updating the generation model is repeated all along.
Authors introduce a novel analysis framework in which uncertainty is viewed as arising from missing data rather than from unknown parameters, and the oracle policy to which the algorithm is compared is the one that best fits the complete dataset without any missing data.
The advantage of this new analysis makes the algorithm compatible with complex task environments and complex policies (with finite VC dimension). 

Overall, reviewers agreed that the proposed framework is novel, and that the paper is written with clarity. Having read the paper myself, I also find that the newly proposed framework, which views uncertainty as arising from missing data rather than unknown parameters, will provide fresh insights into the analysis of other bandit algorithms as well, thereby making the contribution of this paper significant.

One of the main issues was whether the proposed algorithm can truly be regarded as Thompson sampling, to which the authors responded affirmatively and with clarity, noting that the algorithm implements ‘probability matching’. Authors are encouraged to clarify this point once more in their final version. Also during reviewer-AC discussion period, some reviewers questioned the appropriateness of the definition of the best-fitting policy and the fairness of policy evaluation, with these issues clarified through conversations between the reviewers and the AC. Author are encouraged to provide additional details on the best-fitting policy and to include their new TS baseline that they proposed during the rebuttal (which uses $p_{\theta}$ as feature extractor) in their final version of the paper. 

Since the main questions raised by the reviewers were resolved as mentioned above, I recommend acceptance of the paper despite its borderline score.
Decision: Accept (poster)