Contextual Thompson Sampling via Generation of Missing Data

Thank you so much for your thoughtful comments and questions!

(1) Natarajan dimension

We 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:

1. Policy class versus imputation model. Our Proposition 1 upper bounds H(\\mathbf{\\pi}^\\ast(X_{1:T}) \\mid Z\_\\tau) by $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.

2. 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.

3. 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.

(2) Unknown horizon $T$

In 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.

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.

(3) Computational efficiency

When 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:

1. 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).
2. 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].
3. 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.

(4) Typos

• Thank you for pointing out the overlapped citations on Lines 467 and 468. We will correct this!
• 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.

References

• Czarnecki, Pascanu, Osindero, Jayakumar, Swirszcz, & Jaderberg. Distilling policy distillation. AISTATS 2019
• Gelman, Carlin, Stern, Dunson, Vehtari, Rubin. Bayesian data analysis. Chapman and Hall/CRC 1995
• Riquelme, Tucker, & Snoek. Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Sampling. ICLR 2018
• Yi, Wang, He, Chandrasekaran, Wu, Heldt, Hong, Chen, Chi. Online Matching: A Real-time Bandit System for Large-scale Recommendations. RecSys 2023
• Yu & Oh. Optimal and Practical Batched Linear Bandit Algorithm. ICML 2025

Thank you for your questions and engagement with our paper!

As far as I understand, Theorem 1 provides a general result that yields an explicit regret bound once an imputation model $p^\*$ is specified. As the authors noted, when this model is chosen as a linear multiclass predictor, the Natarajan dimension becomes smaller than the dimension of the linear model, which (almost) recovers previously known results. However, as mentioned in Appendix A.5, the resulting regret bound is looser than the prior one.

Based on your comment, we think there may still be some confusion. The complexity of the sequence model $p^\*$ does not impact the Theorem 1 bound and we do not ever discuss the Natarajan dimension of $p^\*$ in our work. We only mention the Natarajan complexity of the policy class $\Pi$, which is used in our Proposition 1 bound.

To clarify the distinction between the sequence model and the policy class, and how the Natarajan dimension fits in, recall:

• After we use the sequence model $p^\*$ to impute missing outcomes, we fit a policy $\pi \in \Pi$ that is used to select an action
• Our Theorem 1 (when sequence model $p^\*$ is used) bounds the per-timestep regret with $\sqrt{\frac{|\mathcal{A}|}{2T} H(\mathbf{\pi}^\*(X_{1:T}) \mid Z_\tau)}$. Our Proposition 1 bounds the term $H(\mathbf{\pi}^\*(X_{1:T}) \mid Z_\tau)$ by $O(b \cdot \log(T \cdot|\mathcal{A}|))$ where $b$ is the Natarajan dimension of the policy class $\Pi$. Note, this is not the sequence model $p^\*$.

Q1. Do you think these additional terms can be reduced, are they unavoidable due to the generality of Theorem 1, or are they inherent to the specific algorithm being analyzed?

Thanks for your question. For clarity, we'll use the linear (non-contextual) bandits example from Appendix A.5, where the expected reward is $\mathbb{E}[ R(Y) | A_t = a] = \varphi(a)^\top \beta$ for some $\beta \in \mathbb{R}^d$.

Combining our Theorem 1, Proposition 1, and the Natarajan dimension of linear classifiers gives us a per-timestep regret bound of $\sqrt{\frac{d \cdot |\mathcal{A}|}{2T} \log(T \cdot|\mathcal{A}|)}$. The reviewer correctly points out this bound is looser than the known result from Russo and Van Roy: $\sqrt{\frac{d }{2T} \log(|\mathcal{A}|)}$.

The additional terms in our bound—the $\log(T)$ and $|\mathcal{A}|$ factors—we believe are not artifacts of our specific algorithm, but rather are a consequence of the generality of our analysis.

• The $\log(T)$ term comes from our use of the Natarajan dimension, a generalization of the VC dimension. This term is common in bandit regret bounds that rely on VC dimension-based analysis, as seen in other work (e.g., Beygelzimer et al). It appears to be an unavoidable consequence of this type of generalized bound.
• The $|\mathcal{A}|$ term is a consequence of the generality of our analysis, which does not utilize a shared parameterization across actions. The Russo and Van Roy bound for linear bandits is tighter because it leverages the linear structure, where $\mathbb{E}[ R(Y) | A_t = a] = \varphi(a)^\top \beta$. In this setting, the parameter $\beta$ 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 a 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 a broader class of problems, but also looser for specific settings like linear bandits with shared parameters across actions.

Ultimately, the primary contribution of our work is not to provide the tightest possible regret bound for a specific parametric model. Instead, our goal is to present a general and robust theoretical framework that characterizes the performance of Thompson Sampling variants that use modern generative sequence models and general policy classes. The generality of our approach means our bounds hold even with misspecified Bayesian models, a scenario where standard Thompson Sampling bounds do not apply.

References

• Beygelzimer, Langford, Li, Reyzin, & Schapire. Contextual Bandit Algorithms with Supervised Learning Guarantees. AISTATS 2011
• Russo and Van Roy. An Information-Theoretic Analysis of Thompson Sampling. JMLR 2016

Q2. While Theorem 1 offers a general formulation, as I noted earlier, the Natarajan dimension of certain neural networks can be very large. This suggests that using more complex models could lead to looser regret bounds, even though the empirical performance may improve, an issue frequently observed in deep learning. In the bandit setting, how should we interpret this phenomenon? Roughly speaking, can we view it as a form of overfitting to limited data, which results in a larger worst-case regret?

Thank you for this question. It gets to the heart of a crucial challenge: the potential trade-off between improved empirical performance and looser theoretical guarantees. You're right that a more complex model can improve performance but may also lead to looser worst-case regret bounds. Our framework addresses this by considering how model complexity impacts two distinct parts of our algorithm.

(a) Complexity of the sequence model $p_\theta$.

The sequence model $p_\theta$ is pre-trained on offline data. A key advantage of our Theorem 2 is that its regret bound depends only on the offline prediction loss of $p_\theta$ compared to the optimal sequence model $p^\*$, i.e., $\ell(p_\theta)-\ell(p^\*)$. The bound does not depend on the intrinsic complexity of $p_\theta$ itself. This implies that using a more complex neural network is not an issue, and in fact is encouraged, as long as it helps to minimize this offline loss.

(b) Complexity on the policy class $\Pi$.

Your question about the Natarajan dimension of complex models is most relevant here, as it directly relates to the policy class $\Pi$. We assume the algorithm designer defines a policy class $\Pi$ and is able to fit a "best-in-class" oracle policy $\pi^* \in \Pi$ given a complete dataset. Our theory bounds the average per-timestep regret between this oracle policy and our algorithm's decisions.

A more complex policy class $\Pi$ in the worst case, it can be much harder for our algorithm to learn $\pi^\*$. This difficulty is captured by the $H(\mathbf{\pi}^\*(X_{1:T}) \mid Z_\tau)$ term in our regret bounds (Theorems 1 and 2). Our Proposition 1 bounds this term by $O(b \cdot \log(T \cdot ∣\mathcal{A}∣))$, where $b$ is the Natarajan dimension of $\Pi$. You are correct that for a complex neural network policy class, a standard characterization of the Natarajan dimension $b$ can be very large and may not be the tightest possible characterization for a specific neural network architecture.

Our goal was to provide a widely applicable theoretical framework for bandit algorithms that can easily take advantage of tighter characterizations of supervised learning models as they are developed in the future. Even though existing Natarajan dimension bounds are loose for certain models, our regret bound results will improve alongside any tighter bounds that researchers develop for supervised learning in the future.

Thank you for your helpful comments and questions!

(1) Experiments comparing to standard Thompson Sampling

The experiments do not compare to regular Thompson sampling (TS) that learn a posterior distribution over $\theta$.

We 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).

(2) Relationship to Thompson Sampling

Can you prove that the proposed algorithm is equivalent to Thompson sampling?

Thank 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.

(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$):

Thanks for your questions and comments! Here's our response to Q1:

Clarifying "Regular TS".

The algorithm you've outlined, which selects the action with the highest sampled outcome ($Y_A$), is a common misunderstanding of standard TS. To clarify, we rely on the standard definition of Thompson Sampling as established in Section 4 of "A Tutorial on Thompson Sampling" by Russo et al. (2018).

Standard TS algorithm:

Specify the prior $q(\mu)$. At the start of the bandit task, receive $Z$. Then, for each decision time $t = 1, 2, ... T$:

1. Receive $X_t$
2. Sample $\tilde{\mu} \sim q(\mu \mid H_t)$
3. Select action $A_t = \mathrm{argmax}\_{a \in \mathcal{A}} \mathbb{E}\_{\tilde{\mu}}[ Y \mid Z, X_t, A = a]$
4. Observe outcome $Y_t$ for action $A_t$
5. Update the history $H_{t+1} \gets H \cup (X_t, A_t, Y_t)$ and the posterior $q(\mu \mid H_{t+1})$

For example, for Beta-Bernoulli multi-armed bandits (no $X$'s or $Z$'s),

• The unknown parameters $\mu = (\mu_a)_{a \in \mathcal{A}}$ are the mean rewards associated with each action
• Specifying the prior: $\mu_a \sim \mathrm{Beta}(\alpha_{1,a}, \beta_{1,a})$ for each $a \in \mathcal{A}$
• Step 2: Form $\tilde{\mu} = (\tilde{\mu}\_a)_{a \in \mathcal{A}}$ by sampling from the posterior $\tilde{\mu}\_a \sim \mathrm{Beta}(\alpha\_{t,a}, \beta\_{t,a})$ for each $a \in \mathcal{A}$
• Step 3: Select action $\mathrm{argmax}\_{a \in \mathcal{A}} ~ \tilde{\mu}_a$, i.e., $\mathbb{E}\_{\tilde{\mu}}[ Y \mid Z, X\_t, A = a] = \tilde{\mu}\_a$
• Step 5: Update posterior of Beta distribution: $\alpha_{t+1,a} \gets \alpha_{t,a} + Y_t$ and $\beta_{t+1,a} \gets \beta_{t,a} + (1-Y_t)$

Your proposed algorithm would replace step 3 above with

• Sample $\tilde{Y}_a \sim \mathrm{Bernoulli}(\tilde\mu_a)$ and select $A_t \gets \mathrm{argmax} ~ \tilde{Y}_a$.

Note that your proposed algorithm would likely perform poorly in practice because $\tilde{Y}_a$ are binary and there would be many ties. Even if the algorithm knows the true mean reward, $\mu_a$, for each action, it would not select the best action with probability $1$.

Comparing our algorithm to "regular TS".

We understand your concern about wanting to compare to a "stronger" regular TS baseline beyond a linear TS algorithm, by using a neural network approach with "approximate Bayesian learning". The literature has been quite stuck on what this "approximate Bayesian learning" should be and it's especially unclear how to learn across tasks and use very rich $Z$'s and $X$'s. We discuss several approaches one might take:

1. Adapt our pretrained sequence model $p_\theta$ to use with standard TS. The algorithm you proposed suggested that we might be able to use our $p_\theta$ sequence model with standard Thompson Sampling. However, our $p_\theta$ model is not a standard reward model that models $\mathbb{E}[Y \mid X, A, Z]$, but rather a sequence model that generates $Y$'s conditional on the history $H_t$. Furthermore, we do not put priors on the weights $\theta$ as we do not treat them as a latent parameter. Our pretraining procedure for learning $p_\theta$ is therefore not appropriate for learning a prior on the weights $\theta$, so it is not directly compatible with standard TS where one samples the latent parameters $\theta$ from the posterior distribution. To reiterate, the sampling in TS-Gen is not from sampling a latent parameter, but rather unknown outcomes in $\tau$, and $\theta$ is not a latent parameter in our setting to sample at all.
2. Use a standard Bayesian Neural Network approach. Our work considers a meta-bandit problem, where the goal is to use data from past tasks to perform well on new tasks. The Thompson sampling and Bayesian neural network literature do not have many practical methods for learning an informative prior over the weights of a neural network for meta-learning problems. We could compare to Bayesian neural network approaches use an uninformative prior, but a Bayesian neural network with a completely uninformative prior that ignores previous task data and uses approximate posterior inference seems like a complex and poor alternative algorithm.
3. Use other approximate Bayesian methods with TS proposed in the literature. When we looked into the literature for TS algorithms with neural networks, the most prominent and comprehensive work comparing methods empirically was "Bayesian Bandit Showdown" by Riquelme et al (ICLR 2018). Surprisingly, they found that simple methods like neural linear (linear TS with the last layer from a neural network as context) was the among the most promising methods and generally outperformed more complex methods like variational inference/Bayes-by-backprop, expectation-propagation, bootstrap, direct noise injection, dropout, Bayesian non-parametrics, etc. This is what led us to believe that a linear based TS algorithm would be a reasonable baseline.

Please let us know if you have further comments on this point.

Thank you for your comments and questions!

(1) Relationship to Thompson Sampling (TS)

In 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.

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.

We 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.

Proposition 2. Algorithm 1 with imputation model $p^\*$ implements Thompson Sampling (probability matching), i.e., the following holds with probability $1$:

Many thanks for the rebuttal! I am sorry but I have been trying to understand the paper and the rebuttal. However, I still do not think this work is up to the standard of the acceptance:

1. Problem with learning objective. The definition of the "best-fitting" policy makes the objective a bit ambiguous. Commonly in the bandit literature, the objective is to learn an optimal policy, but in this work, the policy given the history comes from the oracle. Since the policy is not directly learnable, my understanding is that the main goal is to learn $p^*$, i.e., item (i) in Line 73. However, nowhere in the algorithm design permits the update of the imputation model $p$.
2. Contributions and novelty. The authors claim the following technical novelties: (i) Entropy-based regret bounds, which is interesting but not surprising given the intuition that the regret of the algorithm should depend on the uncertainty of the policy; (ii) Complex reward model, which is possible because the a general best-fitting policy is assumed and no further assumptions are needed to obtain (approximations of) good policies; (iii) flexible policy class & dependence on pretraining loss, which is rendered by the simplicity of analysis. In other words, my understanding is that the work is able to study very general settings because the analysis does not pose too much restrictions to the policy class, imputation model class, etc.
3. Definition of the regret. I made the point that the regret bound was too good to be true, and it turned out that the definition of the regret is defined as the average of the difference between the best-fitting policy and the reward of the action at step $t$. The regret bound becomes unsurprising because the the action at step $t$ is drawn from the best-fitting policy, so there should be no such problem as the constant estimation error in Question 2. I find this definition problematic because this definition disables the discussions about the error between the learner's estimated policy and the best-fitting policy (which I believe is the main source of error), and reduces the regret to some metric of "variance error" or "uncertainty". For lines 100-103, I do not think the regret is the same as the previous works because in this work, the learner has access to the best-fitting policy but not in previous works.

In conclusion, I would like to maintain my score for now. I have also read all other reviews, and find the weaknesses raised by other reviewers mostly minor. I would ask the AC to carefully consider this message, and would like to hear the clarification of the authors.

Thanks for your engagement with our work.

(i) What does the policy fitting oracle output during the process of the algorithm, when the learning agent does not have to the entire dataset?

You're hitting on a key point here. During the online decision-making process, the algorithm does not have a complete dataset of each arm's rewards for all contexts in the current task. Our method proceeds by probabilistically imputing/sampling the unknown rewards with an autoregressive sequence model and then computing the best-fitting policy under that imputed dataset. This is done very carefully to preserve appropriate randomness in the imputation/sampling of the unknown rewards, and that randomness drives careful exploration that provably controls regret.

(ii) How to define the rewards from other tasks and arms if the reward is not queried?

Rewards associated with arms that may or may not be selected are defined as potential outcomes. This is a standard framework in causal inference and bandit literature. Under this framework, all the rewards from a task (from all contexts and arms) are sampled by the environment ahead of time, prior to the algorithm starting to make decisions. Then when the algorithm makes a decision, the environment reveals rewards to the algorithm associated with the selected action.

Many thanks for the reply!

About best-fitting policy: My current understanding is that the best-fitting policy is the best policy if all data, including data from other tasks and arms, were given. If this understanding is correct, then I have two questions: (i) What does the policy fitting oracle output during the process of the algorithm, when the learning agent does not have to the entire dataset? (ii) How to define the rewards from other tasks and arms if the reward is not queried?

These two questions are currently my major confusions.

Thank you for your helpful comments and questions!

(1) "Design based" perspective

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.

Thanks for the comment. There are two relevant points, which we will make clearer in the revision.

(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^\*$: $\\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].$

(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.

(2) Effect of choice of imputation model on performance

We can interpret your question about how "generic" imputation methods perform in our setting in two ways. We provide an answer for each interpretation.

(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.

(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.

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$.

In the table above, observe that less expressive sequence models have higher (worse) sequence loss and worse regret.

References

• Osband, Wen, Asghari, Dwaracherla, Ibrahimi, Lu, & Van Roy. Approximate Thompson Sampling via Epistemic Neural Networks. UAI 2023
• Riquelme, Tucker, & Snoek. Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Sampling. ICLR 2018
• Tran, Snoek, & Lakshminarayanan. Practical uncertainty estimation and out-of-distribution robustness in deep learning. NeurIPS Tutorial 2020

Thank you for your positive comments and for taking the time to engage with our work! We're glad that our rebuttal and the additional experiments helped with understanding our method.