plus0pt minus3pt plus0pt minus3pt

1 Introduction↩︎

Foundation models have moved machine learning from modeling concise numerical data to natively handling unstructured data—text, audio, video, sensor streams—that convey richer information and are often more natural to collect. Yet these models seem to lack a capability that has long been studied in statistics: quantifying epistemic uncertainty, the uncertainty that comes from limited data and would resolve if more data were gathered, as opposed to aleatoric uncertainty, which comes from irreducible noise.

We view foundation models through the lens of the sequence-prediction task that they are often trained on: predicting the distribution of the next outcome given previous ones. Even a perfectly accurate next-outcome predictor does not obviously expose epistemic uncertainty on the underlying environment. In this paper, we explore whether this capability for next-outcome prediction is nevertheless enough to support epistemic uncertainty quantification. We use a multi-armed-bandit style sequential decision-making setting as a demanding testbed—one that requires not only recognizing remaining uncertainty, but also acting to resolve it and updating beliefs as new information is incrementally gathered. We aim to power this entire loop using capabilities native to foundation models:

This turns out to be surprisingly subtle. Methods that summarize only the next-step predictive distribution—for instance, by sampling a single next outcome and acting greedily, as proposed in recent work [1]–[4]—can suffer linear regret, eventually badly outperformed by simple algorithms with no foundation model access at all (see Section 7.5, Figure 14). The issue is that good decisions require more than predicting the next outcome; this fact has been reinforced by recent work demonstrating the challenges of using large language models (LLMs) for exploration [5]–[7]. The agent must reason about how future outcomes, once observed, would resolve epistemic uncertainty relevant to later decisions. A one-step predictive distribution marginalizes over both the epistemic and aleatoric uncertainty and therefore cannot drive exploration on its own. How, then, can next-outcome prediction be used to represent the uncertainty that future observations would resolve?

A classical approach to this sequential decision-making problem posits latent parameters governing aspects of the unknown environment, such as parameters in a reward model. In a Bayesian view, these latents are random variables, and the decision-maker tracks a posterior distribution over these latents as new data arrives. This posterior over the latents represents epistemic uncertainty, i.e. uncertainty about the environment that should guide exploration toward observations that would be most informative. But maintaining such a posterior is often a critical computational bottleneck when observations are high-dimensional or unstructured, and explicit posterior distributions are not a primitive that modern foundation models natively support. In the decision-making settings we study, this latent parameter approach is also not strictly necessary. The decision-maker ultimately cares about the outcomes that would result from potential future actions—for example, which action would have the highest reward on average, or which observations would be informative for future decisions—not about the latent parameter itself. If the relevant future outcomes were exactly known under each possible sequence of actions, the decision problem would be resolved. This motivates a missing-data view of uncertainty: rather than represent epistemic uncertainty through an unobserved latent parameter, we represent it through the potentially observable outcomes that have not yet been revealed, in the style of [8].

1.0.0.1 Our approach.

  Operationalizing this missing-data viewpoint yields concrete replacements for the subroutines that classical Bayesian exploration algorithms rely on; see Figure 1. These algorithms typically require i) constructing informed priors over latent parameters, ii) performing probabilistic inference under the current model of uncertainty, and iii) updating beliefs as more data is gathered—each of which becomes difficult with high-dimensional or unstructured data. Our approach instead i) trains a generative model through next-outcome prediction, ii) assesses uncertainty through autoregressive generation of missing outcomes—including those not yet observed, and iii) adapts to new information by extending the sequence model’s context [11] rather than explicitly updating a posterior. Marginal next-outcome prediction thus becomes the basic primitive for epistemic uncertainty quantification and active exploration: the same sequence-prediction machinery that already handles rich, unstructured inputs can drive exploration, without an explicit prior, likelihood, or posterior update over latent environment parameters. This recipe is not tied to any single algorithm—autoregressive generation serves as a drop-in replacement for posterior sampling of latents in Thompson sampling, Bayes UCB, knowledge gradient, information-directed sampling, and related methods [12]–[15].

1.0.0.2 Outline.

  Section 2 develops the missing-data view of uncertainty in a passive-observation setting and shows that autoregressive generation serves as a drop-in replacement for a posterior sampling step that appears in classical Bayesian exploration algorithms such as Thompson sampling, Bayes UCB, knowledge gradient, and information-directed sampling. Section 2.3 itemizes our contributions in full. Sections 3 and 4 formalize the online decision-making problem and the use of autoregressive generation inside exploration algorithms, with a special focus on Thompson sampling [12], [16]; we call our Thompson sampling based algorithm Thompson sampling via autoregressive generation (TSAR). Section 5 proves a regret bound that reduces online decision-making to offline next-outcome prediction. Section 7 validates these ideas in synthetic settings and in a semi-synthetic news recommendation task that requires leveraging article headline text to focus exploration.

2 Key ideas: Missing data view of uncertainty↩︎

While later sections of the paper study active exploration problems, we start by developing key ideas in a setting where observations accrue passively. A few examples help to explain the setting. Consider a news recommendation platform that has just released a new article. From the headline text, the platform can form an initial guess about how engaging the article will be, but substantial uncertainty remains—uncertainty that gets progressively resolved as the article is recommended to one user after another [17]. A similar phenomenon arises in political science, where a previous survey may suggest an increased probability of approval on a certain issue, but considerable uncertainty remains [18]. As more individuals are polled, the overall approval becomes increasingly clear.

In both cases, we face a fundamental inferential question: given limited observations so far, what can we confidently conclude, and how much uncertainty remains? In Section 2.1, we review classical Bayesian inference, which approaches these problems by representing unknowns through latent parameters—the engagement rate of an article or the approval rate of a particular policy, for example. In Section 2.2, we present our alternative approach that views missing data—the user feedback or surveys not yet gathered—as the source of uncertainty. Under this view, the inferential question becomes: given limited observations so far, what can we already infer about these missing but observable quantities, and what remains uncertain? Correct probabilistic imputation of missing data powers this style of uncertainty quantification without any explicit reference to latent variables.

2.0.0.1 A note on prior knowledge.

  Both approaches to uncertainty quantification require substantial probabilistic modeling. In the classical Bayesian setting, the prior and likelihood together specify the joint distribution over latents and observables. In our missing-data-based formulation, we require a sequence model that specifies the distribution of the next observation conditioned on previous observations. We begin by assuming that respective probabilistic models are available and focus on the resulting inference problems in Sections 2.1 and 2.2. In Section 2.3, we return to where these probabilistic models come from and illustrate how the missing data view of uncertainty allows using foundation models to integrate complex prior knowledge (e.g., text and video information).

2.1 Classical Bayesian modeling↩︎

A standard way to formalize uncertainty is to posit a Bayesian model with a latent parameter, say, an underlying quality \(\eta\) that governs observed outcomes \(Y_1,\ldots,Y_T\). Here, one can treat available prior information as shaping initial beliefs about \(\eta\), and then one updates those beliefs as observations arrive. Throughout, we explicitly denote any initial information the modeler may have as \(Z\), to highlight the significance of such information in modern application scenarios (e.g., article text or past clinical notes). Throughout, we use the phrase “prior information” to refer to \(Z\), as it is information about the task prior to observing outcomes \(Y\).

Standard Bayesian approaches model exchangeable observations \(Y_1, Y_2, \ldots, Y_T\) as being drawn independently and identically distributed conditioned on \(\eta\). Given a subset of observations \(Y_{1:t}\), the agent forms the posterior belief

\[\label{eq:posterior95formula} p(\eta \mid Z, Y_1, \ldots, Y_t) \propto \prod_{\tau=1}^{t} p(Y_\tau \mid \eta) \cdot p(\eta \mid Z),\tag{1}\] which summarizes everything that can (and cannot) be inferred about \(\eta\) given the observed information. The prior distribution \(p(\eta \mid Z)\) of \(\eta\) given \(Z\) and likelihood \(p(Y_\tau \mid \eta)\) are assumed to be known.

Typically, quantities of interest can be expressed as a function of \(\eta\). One example is the conditional mean \(f(\eta)=\E[Y_\tau \mid \eta]\) where we assume \(Y_\tau\) is a real-valued random variable. We formalize the problem of performing posterior inference about any quantity of interest \(f(\eta)\) below. Grouping together prior information \(Z\) and observations as \(Y_{1:t} = (Y_1, \ldots, Y_t)\) to infer the distribution of \(f(\eta)\) given \((Z, Y_{1:t})\) we can repeatedly draw \(\widehat{\eta}\) from the posterior distribution 1 , and compute \(f(\widehat{\eta})\).

While elegant for simple settings, this framework faces practical challenges when the data is rich, and complex models are used for the prior and likelihood. Specifying an appropriate prior distribution \(p(\eta \mid Z)\) for \(Z\) with complex, unstructured information (like article text) involves incorporating knowledge from foundation models to inform the distribution of \(\eta\), which is non-trivial. Additionally, when outcomes \(Y\) are themselves high-dimensional or unstructured, the latent \(\eta\) itself may be complex (e.g., the parameters in a neural network), making posterior computation difficult, particularly as data accrues over time \(t\).

2.2 Uncertainty as missing data↩︎

This paper takes an alternative approach that sidesteps latent parameters entirely. Instead of focusing on inferring a latent, intangible quality of the “true” environment, \(\eta\), we consider a quantity of interest that is a function \(f(Y_1, \ldots, Y_T)\) of the full trajectory of outcomes \(Y_1, \ldots, Y_T\) sampled from some joint distribution \(p(\cdot \mid Z)\); in contrast to Section 2.1, here we do not assume that \(Y_{1:T}\) are necessarily exchangeable given \(Z\). For example, practically the quantity of interest \(f(Y_1, \ldots, Y_T)\) could be the average engagement if a news article were shown broadly, i.e., \(f(Y_1, \ldots, Y_T) = \frac{1}{T} \sum_{t=1}^T Y_t\). At any timepoint \(t\), the uncertainty about this quantity of interest arises from missing future outcomes \(Y_{t+1}, \ldots, Y_T\). That is, uncertainty does not stem from a latent environment \(\eta\) parameter, but from outcomes \(Y_{t+1}, \ldots, Y_T\) that are yet to be seen.

Our uncertainty quantification task is as follows: given the prior information \(Z\) and the outcomes observed so far \(Y_1,\ldots,Y_t\), what is the distribution of the quantity of interest? Even among instances with similar prior information \(Z\) and the same early outcomes, there can be substantial variability in the eventual final outcome, so this distribution is genuinely non-degenerate. We can formalize the problem in a manner that parallels the problem of inference under an uncertain latent \(\eta\) to a striking degree.

While so far we have focused on passively observed outcomes, when extended to the decision-making setting, this view of missing future outcomes as the source of uncertainty is particularly aligned with active exploration problems, where gathering missing observations tends to be a concrete choice the agent faces. This viewpoint is also very general. In particular, it does not require that the data is conditionally i.i.d. given some \(\eta\); in fact, this viewpoint does not even require exchangeability of \(Y_1, \dots, Y_T\).

2.2.0.1 Autoregressive generation of missing data.

   Even with a correctly specified model of the joint distribution \(Y_1 \dots, Y_T \mid Z \sim p(\cdot \mid Z)\), quantifying posterior uncertainty over the realization \(f(Y_1, \ldots, Y_T)\) after observing \((Z, Y_{1:t})\) generally is not available in closed form, and so a sampling-based approach to estimation is recommended. As shown in Figure 2, our inferential approach uses the variation in the estimated \(f(Y_1, \ldots, Y_T)\) across probabilistic imputations of the missing outcomes to quantify uncertainty.

In particular, we probabilistically impute the next unobserved outcome \(Y_{t+1}\) given the prior information and observed information so far \(Z, Y_{1:t}\) using \(\widehat{Y}_{t+1} \sim p(Y_{t+1} \mid Z, Y_{1:t})\). We then impute \(Y_{t+2}\) via autoregressively sampling the next conditioned on the same history \((Z, Y_{1:t})\) in addition to the sampled value \(\widehat{Y}_{t+1}\), i.e., \(\widehat{Y}_{t+2} \sim p(Y_{t+2} \mid Z, Y_{1:t}, Y_{t+1} = \widehat{Y}_{t+1})\). We continue this process until all remaining outcomes \(\widehat{Y}_{t+1,T}\) are probabilistically imputed; then, evaluating \(f\) on the imputed trajectory produces one draw from the induced posterior over \(f\).

As an example, if \(Y_1\) and \(Y_2\) have been observed:

• Sample \(\widehat{Y}_3 \sim p(Y_3 \mid Z, Y_1, Y_2)\)

• Sample \(\widehat{Y}_4 \sim p(Y_4 \mid Z, Y_1, Y_2, Y_3 = \widehat{Y}_3)\), conditioning on both real observations and past generations

• Continue until all missing entries have been imputed.

We call the model of conditional distributions of the form \(p(Y_{t+1} \mid Z, Y_1, \ldots Y_{t})\) a sequence model. By autoregressively sampling from the correct conditional distribution \(p\), we can sample from the correct joint distribution over missing outcomes.

The autoregressive imputation approach described above does not require explicitly modeling a latent \(\eta\), but is compatible with such models from Section 2.1: one can simply write the distribution of future outcomes as posterior predictives using the law of total probability. In this special case, our imputation approach matches the sampling methodology used in the multiple imputation literature [19], [20].

2.2.0.2 Why the joint distribution matters, not just the next-step marginal.

The autoregressive procedure does more than predict the next outcome; it posits a joint distribution over the remaining sequence. To see why this matters, consider two distributions over a sequence of coin tosses. The first is a fair coin with i.i.d. tosses of mean \(1/2\); the second is a deterministic coin whose side is unknown but, once observed, is fixed for every subsequent toss. The two have identical marginal distributions over the next toss but very different joint distributions: a single observation resolves all uncertainty in the second world and none in the first; see [21]. Methods that summarize only the one-step distribution—such as generating a single next outcome [1], [2], [22] or averaging several independent next-outcome draws—cannot tell the two apart, and consequently fail to drive exploration correctly when used inside a bandit algorithm. Section 7.5 demonstrates these failure modes empirically.

2.3 Approximate sequence models and our contributions↩︎

This paper applies the ideas developed in Section 2.2 to online learning problems where present actions impact data observed in the future. The autoregressive imputation approach to uncertainty quantification opens up a path to leveraging foundation models for probabilistic uncertainty quantification in ways that appear difficult with the traditional Bayesian modeling of a latent parameter. The central challenge we focus on in this work is the inherently approximate nature of learned sequence models and how uncertainty is used to guide exploration.

We begin by assuming access to an accurate sequence model in Section 4, and provide regret results for Thompson sampling and provide a novel result for the benefit of using action features for screening unpromising actions. In the following sections, we assume access to an approximate sequence model \(\widehat{p}\) for predicting next outcomes to quantify uncertainty—without any additional Bayesian inference machinery. There are several possibilities for forming such a \(\widehat{p}\):

• Prompting. Relying on their strong zero-shot performance, one can simply prompt an LLM to predict the next observation given prior information \(Z\) and past outcomes \(Y_1, \ldots, Y_{t}\).

• In-context learning. One can augment prompting by including in the prompt historical examples of prior information \(Z\) and their corresponding past outcomes \(Y_1,\ldots,Y_{t}\) [23].

• Fine-tuning. When domain-specific outcome data is available, one can fine-tune a pretrained model to predict sequences of outcomes conditioned on prior information \(Z\).

• Training from scratch. For organizations with abundant historical data (e.g., Google [24]), one can train a sequence model entirely from scratch over past outcome sequences. We discuss our procedure for this in Section 7.1, which we use in a semi-synthetic news recommendation experiment.

The rest of the paper develops three contributions that, together, turn the missing-data view into a working algorithm with formal guarantees.

1. Uncertainty quantification alone does not prescribe action; an agent must translate uncertainty about missing outcomes into decisions that balance exploitation of current knowledge against exploration to resolve remaining uncertainties. To study this problem, Section 3 extends the missing data formulation to bandit-style decision-making tasks. Then, in Section 4, we adapt classical Bayesian exploration algorithms by replacing posterior sampling over latent parameters with autoregressive generation of missing outcomes conditioned on prior information \(Z\) and outcomes observed so far in the current task. Section 4 formally interprets this generation process as posterior sampling whenever the sequence model can sample exactly from the relevant conditional distributions, providing a drop-in replacement in algorithms such as Bayes UCB, knowledge gradient, information directed sampling, and Thompson sampling [13]–[15]. We develop and analyze the Thompson sampling version most fully.

2. In online bandit problems, the agent’s actions determine which observations it gathers, errors in an approximate \(\widehat{p}\) can compound: mistaken beliefs may lead an agent to choose actions consistent with those beliefs, causing it to rarely collect data that would correct them. [25] make this vivid by showing that Thompson sampling with a misspecified prior can suffer worst-case cumulative regret that scales quadratically in the horizon \(T\), and a reflexive imitation-learning argument would have offline prediction error enter the regret bound multiplied by a distribution-shift coefficient [26]. Section 5 rules out such multiplicative cascades. We prove that, averaged over the prior that actually generates the bandit task, the agent’s (Bayes) regret is directly controlled by the sequence model’s offline next-outcome prediction error, scaling at worst linearly in \(T\) and recovering the standard \(\sqrt{T}\) Thompson sampling rate as offline predictions get better. This establishes a clean reduction from online decision-making to offline next-outcome prediction.

3. We demonstrate that our theoretical insights bear out in simulations. Our methods not only match oracle Bayesian procedures in synthetic settings, but do so without relying on an explicit latent parameter specification. Moreover, our approach scales to a news article recommendation setting in which the best performance requires leveraging unstructured prior information like article headlines. Even when incorporating and fine-tuning a language model, we find our autoregressive generation procedure implements accurate uncertainty quantification (credible intervals) and drives exploration with low regret.

2.3.0.1 Relationship to a follow-up paper.

  A follow-up paper by the same authors extends one component of this work—the Thompson sampling regret analyses in Sections 4.2.1 and 5.2—to contextual bandits with general policy classes, adding VC-dimension-based entropy arguments specific to that setting [27]. The present paper is broader in scope; specifically it develops the conceptual foundations discuss in this section, develops a novel regret bound on screening actions using high-dimensional action features (Section 4.2.2), general-purpose reductions of regret to offline prediction for multiple adaptive policies (Section 5.1), a formal comparison with previous analyses of Thompson sampling under prior misspecification (Section 5.2.2), and empirical evaluations of epistemic uncertainty quantification (Section 7).

3 Preliminaries: Online decision-making problem↩︎

In this section, we provide a formal definition of the decision-making problem. We consider a bandit-style decision-making task in which an agent sequentially selects actions from a finite set \(\MC{A}\) across time periods \(t\in \{1, 2, \ldots, T\}\), adapting future actions in response to feedback. Formally, we say a bandit task \(\mathcal{B} = (Z, \tau)\) consists of prior task information \(Z\) and a two-dimensional array or “table” of potential outcomes \[\tau = (Y_{1}^{(a)}, \ldots, Y_{T}^{(a)})_{a \in \MC{A}},\] where each \(Y_t^{(a)}\) represents what would be observed if action \(a \in \MC{A}\) were taken at time \(t\). This table is the multi-action analogue of the outcome sequence \(Y_1, \ldots, Y_T\) from the passive setting of Section 2.2: each row represents the potential outcomes for one action, and, exactly as in Section 2.2, what the agent does not yet know about the task is encoded in the entries of \(\tau\) that have not been revealed—not in any latent parameter. Prior information \(Z\) can be made up of rich, unstructured data, and could comprise a set of action-specific information \(\{Z^{(a)}\}_{a\in \mathcal{A}}\). Prior information lets us model problems where rich information about the task is available before any outcomes are observed, and where that information should be used to guide exploration. For example, in a news recommendation problem, \(Z\) could represent news article text.

The pair \(\mathcal{B}=(Z,\tau)\) defines a bandit problem as follows: at each timestep \(t\), the agent selects action \(A_t \in \MC{A}\), observes a vector of outcomes \(Y_t\), and associates with this a reward \(R(Y_t) \in [0,1]\) by applying a fixed, known function \(R(\cdot)\). In particular, \(Y_t = Y_t^{(A_t)}\). We index rows of the table by actions and columns by time, writing \(\tau[a,t] = Y_t^{(a)}\). The agent can make decisions using its own task-specific history from that day: \[\HH_{t-1}\triangleq \big( Z, \, (A_1, Y_1),\ldots, (A_{t-1}, Y_{t-1}) \big), \label{eqn:history}\tag{2}\] which includes all historical observations available at the start of time \(t\) and prior information \(Z\). We call this an informed bandit problem as there is prior information \(Z\) associated with each task. Contrast this with contextual bandit problems, where fresh context arrives at each timestep and describes the current decision-making situation rather than the entire task. While existing works have studied meta-bandits that utilize task features [28], to our knowledge, existing bandit algorithms have not previously studied this kind of informed bandit problem with high-dimensional, unstructured \(Z\) (e.g., text).

The goal, informally, is to select a rule for adaptively selecting actions that maximizes cumulative expected reward. We assume that the bandit task \(\mathcal{B}\) is sampled i.i.d. from the task distribution: \[\begin{align} \label{eq:pstar} \mathcal{B} \sim p_{\TN{tasks}}^*. \end{align}\tag{3}\] This is the direct extension of the model in Section 2.2, where we assumed \(Z\) and the single sequence \(Y_1, \ldots, Y_T\) are drawn jointly from some distribution.

3.0.0.1 Online decision-making objective.

  In the passive setting of Section 2.2, inference centered on a trajectory-level quantity of interest \(f(Y_1, \ldots, Y_T)\). Similarly, in this active multi-action setting, various quantities of interest to the decision-maker can be written as functions of the potential outcomes table \(\tau\). The long-run mean reward of action \(a\) is \[\label{eq:mean95reward95tau} \mu_{a}(\tau)= \frac{1}{T} \sum_{t=1}^{T} R \big( Y_{t}^{(a)} \big).\tag{4}\] This is essentially the “population mean” for a large finite population of size \(T\), a common object of interest in statistics [29]. Similarly, the action with maximal reward is \[\label{eq:optimal95action95tau} A^*(\tau) \in \argmax_{a\in \MC{A}} \, \mu_{a}(\tau).\tag{5}\] Above, we assume ties are broken in some deterministic fashion. A decision-maker who knew \(\tau\) upfront could immediately implement the action \(A^*(\tau)\) with highest average reward throughout all periods, with no need to engage in costly exploration.

The cost of exploration is typically measured through regret: the gap in rewards accrued relative to an omniscient agent who knew from the beginning the best fixed action over the horizon, \(A^*(\tau)\). To define this formally, we let \(\pi\) denote some adaptive policy—i.e. any (possibly randomized) rule that specifies which action to select next, given \(Z\) and the history of past actions and \(Y\)’s. The per-period regret incurred by \(\pi\) on bandit instance \(\mathcal{B}\) is given by¹ \[\Delta(\pi, \mathcal{B}) = \E_{\pi}\left[ \max_{a\in \mathcal{A}} \frac{1}{T} \sum_{t=1}^{T} R(Y_t^{(a)}) - \frac{1}{T} \sum_{t=1}^{T} R(Y_t^{(A_t)}) \, \bigg| \,\mathcal{B} \right].\] The expectation integrates over any internal randomness used in \(\pi\). For deterministic policies, no expectation is needed.

We measure performance through expected regret \[\Delta(\pi) \triangleq \E \left[\Delta(\pi, \mathcal{B}) \right],\] which averages over the draw of the bandit task \(\mathcal{B}\). In our framework, this has the interpretation (by the law of large numbers) of being the long-run regret a fixed policy \(\pi\) would incur on average across many i.i.d. draws of bandit tasks.

4 Active exploration via generative implementations of “posterior sampling”↩︎

We now formally introduce our decision-making algorithm. Throughout Section 4, we assume the algorithm has knowledge of the task distribution \(p_{\TN{tasks}}^*\) defined in 3 , from which \(\mathcal{B}\) is drawn. Specifically, knowledge of \(p_{\TN{tasks}}^*\) means that the agent knows at all times the conditional distribution of any entry in potential outcomes table, written as \[\label{eq:missing95entry95distribution95general} p_{\TN{tasks}}^* \big(\tau[a, i] = y \mid Z, (\tau[a',i'])_{(a',i') \in \mathcal{O}} \big)\tag{6}\] for any tuple \((a, i)\) and any set of observed tuples \(\mathcal{O}\). The above conditional distribution can be thought of as a sequence model, as it is the probability of one outcome, given a previous sequence of outcomes, and task information \(Z\). For now, assume \(p^*_{\rm{tasks}}\) is known; we discuss the case in which we only have an approximate sequence model later in Section 5. In practice, we will make simplifying assumptions on the structure of \(p^*_{\TN{tasks}}\), which will make implementation easier (see Section 4.3).

Note that even if the task distribution \(p_{\TN{tasks}}^*\) is known, the decision-making problem for our particular task \(\MC{B}\) is non-trivial: \(p_{\TN{tasks}}^*\) specifies the distribution of tasks, but not our particular task \(\MC{B}\). For example, we may know that bandit tasks with certain features \(Z\) on average have certain characteristics, but we do not know the realization of the potential outcomes \(\tau\) for the task at hand; this is the missing-data uncertainty introduced in Section 2.2.

4.1 Posterior sampling by imputation when \(p_{\TN{tasks}}^*\) is known↩︎

While the agent knows the conditional distributions of missing outcomes according to \(p_{\TN{tasks}}^*\), the agent does not know the realization of \(\tau\). All decisions are made under uncertainty about missing entries in this potentially observable data table \(\tau\). As decisions are made sequentially, each observed outcome exactly reveals one entry in \(\tau\), resolving uncertainty about that entry while potentially informing predictions about the remaining unobserved entries. The agent may also resolve substantial uncertainty by leveraging prior information \(Z\). One way to reason about this uncertainty is to imagine what the full table could be, drawing a hypothetical realization of \(\tau\) that is consistent with what has already been observed. We call this a posterior sample, which we define formally in Definition 1 below.

An ability to draw posterior samples enables us to approximate the posterior distribution of any function of \(\tau\). For instance, at timestep \(t\), the posterior probability some action \(a\) is optimal on average across the \(T\) rounds could be approximated by Monte-Carlo simulation by repeatedly drawing posterior samples \(\widehat{\tau}_t\), then calculating and recording the fraction of draws under which \(A^*(\widehat{\tau}_t)=a\).

Posterior sampling requires sampling from the joint distribution of the missing entries. This can be implemented by iterating over them sequentially, at each step drawing a value from the conditional distribution 6 given the previously sampled entries, as previewed in Section 2. Algorithm 3 shows pseudocode where we begin with a partially revealed potential outcomes table, then repeatedly select a missing entry (line [line:generic95psar95pick]), sample its value conditioned on both observed entries and previous samples (line [line:generic95psar95sample]), and then record it into the table (line [line:generic95psar95record]). This is an extension of the imputation procedure in Section 2.2, where the procedure from Figure 2 is carried out over the potential outcomes table \(\tau\) rather than a single sequence.

Correctness of Algorithm 3 is an immediate consequence of the chain rule of probability, which factorizes a joint probability distribution into a product of conditional distributions. We state the result and omit the proof.

4.2 A generative implementation of Thompson Sampling↩︎

Here we start by introducing our decision-making algorithm that uses posterior sampling by autoregressive generation (Algorithm 3, “PSAR”) as a subroutine. We will show that the decision-making algorithm implements an analogue of Thompson sampling (i.e., probability matching) that replaces the usual step of posterior sampling of latent parameters \((\eta)\) for a posterior sample of the potential outcomes table \((\tau)\).

The algorithm maintains an internal copy of the potential outcomes table, initialized with all values set to missing (line [line:generic95tspsar95init]). As actions are selected and outcomes are observed, they are recorded in the table. To select an action in any period \(t\), the algorithm generates a posterior sample of the potential outcomes table \(\widehat{\tau}_t\) by copying observed entries, and sequentially imputing missing ones (line [line:generic95tspsar95impute]), calculates the action with maximal average reward under that table and chooses to explore that action (line [line:generic95tspsar95calculate95select]). Note that imputations generated at time \(t\) are used only to select an action for that timestep and are not used thereafter.

We call this procedure Thompson sampling because it satisfies the probability matching property that is often taken as the definition for Thompson sampling [12], [32], [33], which is that

4.2.1 Regret bounds for online decision-making using \(p^*\) (Algorithm 4).↩︎

The next result applies pre-existing analyses to bound the regret of this generative Thompson sampling variant. In particular, per-period regret vanishes if the number of interactions per action, \(T/|\MC{A}|\) is large, which are settings with ample opportunity for exploration. Note that Proposition 1 below does not make a stationarity assumption on the reward distribution over time.

4.2.2 The benefit of prior information: Regret with effective screening↩︎

The value of prior information \(Z\) in our framework is that it lets the agent largely avoid exploring unpromising actions. Probability matching makes this concrete: under Thompson sampling with a correctly specified model, since the probability an action is played equals the model’s posterior probability that it is optimal, an action assigned prior probability below \(1/T\) of being optimal is played less than once in expectation over the whole horizon. In regimes that motivate informed exploration—many candidate actions relative to the horizon—this screening is what makes strong performance possible at all.

The following result shows that this screening effect is captured automatically by our regret bound, without any modification to the algorithm. Although the intuition is natural—Thompson sampling should not waste effort on arms the prior already rules out—we are not aware of a similar result in the literature, and provide our own proof based on the information-ratio technique of Russo and Van Roy [33].

For a threshold \(\epsilon > 0\), define the candidate set \[S(Z) \;=\; \bigl\{a \in \mathcal{A}: \PP(A^*\!=\!a \mid Z) > \epsilon\bigr\},\] the actions that the prior information \(Z\) leaves plausibly optimal. Note that the size of the set \(|S(Z)|\) is a random variable determined by \(Z\); we write \(\E[|S(Z)|]\) for its expectation over the task distribution.

Compared with Proposition 1, which gives \(\Delta(\pi_{\rm TSAR}) \le \sqrt{|\mathcal{A}|\log|\mathcal{A}|/(2T)}\), the bound ?? replaces \(|\mathcal{A}|\) in the leading \(\sqrt{T}\) term with the expected number of actions that survive screening, \(\E[|S(Z)|]\). The arms screened out by \(Z\) contribute only a \(T\)-independent additive cost of \((|\mathcal{A}| - \E[|S(Z)|])/T\) per period. When \(Z\) is highly informative—so that \(\E[|S(Z)|] \ll |\mathcal{A}|\)—this represents a substantial improvement. The proof in Appendix 9.6 utilizes information theoretic techniques to analyze Thompson sampling from [33].

4.2.3 Online decision-making beyond Thompson sampling↩︎

The missing-data view of uncertainty applies more generally to other active exploration algorithms. In particular, our posterior sampling via autoregressive generation approach (Algorithm 3) can be used as a replacement for posterior sampling of latent variables in many common decision-making algorithms. For example, we can implement upper confidence bound algorithms by simulating several imputed average rewards for each action and taking an action that maximizes their quantile; see Algorithm 5 for pseudocode for BayesUCB [13]. Similarly, we can implement look-ahead procedures (e.g., Knowledge Gradients [15]) by imputing an outcome from an action taken this timestep, and appending it to the context of the sequence model to simulate outcomes for future timesteps. Our missing data perspective allows one to extend these existing exploration algorithms to utilize potentially complex prior information through the pretrained sequence model.

4.3 Autoregressive posterior sampling under structural assumptions↩︎

In this section, we present a specific instantiation of our posterior sampling procedure, Algorithm 3, under two specific structural assumptions: independence across actions (Assumption 1) and exchangeability over time (Assumption 2). Additionally, we assume that the prior information \(Z=\{Z^{(a)}\}_{a\in \MC{A}}\) consists of prior information for each action. Since these structural assumptions make implementing our algorithm more straightforward, the remainder of this paper will focus on the setting in which these two structural assumptions hold.

Assumption 1 implies that given \(p_{\TN{tasks}}^*\) is known, the prior information and potential outcomes for action \(a\), \(\big( Z^{(a)}, Y_{1:T}^{(a)}\big)\), do not provide any additional information about \(\big( Z^{(a')}, Y_{1:T}^{(a')}\big)\) for a different action \(a'\). This assumption is reasonable in certain domains (like product recommendations) and simplifies the generative posterior sampling procedure substantially because it allows us to draw inferences independently across actions.

Next, we make an exchangeability over time assumption, which means that the distribution of the potential outcomes table is invariant under permutation of its columns (timesteps). In our setting, where the columns correspond to the time periods in a bandit problem, this amounts to a stationarity assumption that is inherent to most studies of multi-armed bandit problems².

One example of when exchangeability holds is when the outcomes \((Y_1^{(a)}, Y_2^{(a)}, \ldots, Y_T^{(a)})\) are independent and identically distributed (i.i.d.), but assuming outcomes that are i.i.d. is too strong for our purposes. For example, in the classical Bayesian setting described in Section 2.1, the outcomes are not i.i.d.—they are potentially highly correlated due to their common dependence on the latent factor \(\eta\)—but they are exchangeable. Exchangeability encodes a temporal stability that ensures early observations from playing an action are informative about future performance, allowing for a stable “best action” \(A^*\) that an agent could hope to identify and consistently play over time.

Under Assumptions 1 and 2, the procedure in Algorithm 3 reduces to Algorithm 6 (depicted in Figure 7). The main difference between these algorithms is that Assumptions 1 and 2 allow us to generate outcomes \(\widehat{Y}^{(a)}_i\) independently across actions, and outcomes for action \(a\) to depend on the prior information for only that action, \(Z^{(a)}\).

5 Autoregressive “posterior sampling” with approximate sequence models↩︎

In this section, we provide theoretical guarantees for our autoregressive generation approach to forming posterior samples when we do not know the true task distribution \(p_{\TN{tasks}}^*\), but we instead have access to an approximation \(\widehat{p}_{\TN{task}}\). Specifically, we provide two main theoretical results:

1. Forming posterior samples with a sequence model \(\widehat{p}_{\TN{task}}\) approximates posterior sampling up to an error term that depends on the excess prediction loss of \(\widehat{p}_{\TN{task}}\) compared to \(p_{\TN{tasks}}^*\) (Section 5.1). Notably, this guarantee holds even when sampling is conditioned on a history produced by an arbitrary adaptive data-collection policy.

2. For the regret of online decision-making, we first prove a general reduction that holds for any policy: regret against the true environment is controlled by regret in an environment consistent with \(\widehat{p}_{\TN{tasks}}\), plus a penalty for the excess offline prediction loss. Specializing this reduction to Thompson sampling via autoregressive generation (TSAR) then yields a concrete regret bound in terms of the offline prediction loss of \(\widehat{p}_{\TN{tasks}}\) (Section 5.2).

Throughout this section, we work under Assumptions 1 and 2.

5.1 Theory Part I: Autoregressive generation from approximate sequence models approximates posterior sampling↩︎

When can we trust that autoregressive generation will produce faithful posterior samples to guide exploration? While it is natural to expect high-quality samples when the sequence model exactly matches the data-generating distribution, this section establishes a far more general guarantee: the quality of generated posterior samples is controlled by the model’s offline prediction performance, even when sampling is performed conditional on strategically collected observations. This provides a concrete foundation for our algorithm’s core premise—that accurate next-outcome prediction enables reliable posterior sampling during online exploration.

We quantify the quality of a sequence model \(\widehat{p}_{\TN{per-action}}\) by its log-loss, \[\begin{align} \label{eq:pop95loss} \ell(\widehat{p}_{\TN{per-action}}) &\triangleq \E\bigg[ -\sum_{t=1}^{T} \log \, \widehat{p}_{\TN{per-action}} \Big( Y_t^{(a)} \mid Z^{(a)}, Y_{1:t-1}^{(a)} \Big) \bigg], \end{align}\tag{7}\] where the expectation above is taken over \((Z^{(a)},Y^{(a)}_{1:T})\sim p_{\TN{per-action}}^*\). We use this metric since sequence prediction models have had immense success in optimizing this exact quantity. First, note the following identity, which expresses the excess loss of a sequence model, \(\ell(\widehat{p}_{\TN{per-action}})- \ell(p_{\TN{per-action}}^*)\), as the average KL divergence between the prediction of \(\widehat{p}_{\TN{per-action}}\) and those of the true data-generating model \(p_{\TN{per-action}}^*\). \[\begin{gather} \ell(\widehat{p}_{\TN{per-action}}) = \ell(p_{\TN{per-action}}^*) \\ + \E_{Z^{(a)} \sim P_Z} \left[D_{\rm KL}\left( p_{\TN{per-action}}^* \big(Y_1^{(a)}, \ldots, Y_T^{(a)} \mid Z^{(a)} \big) \; \Big\| \; \widehat{p}_{\TN{per-action}} \big( Y_1^{(a)}, \ldots, Y_T^{(a)} \mid Z^{(a)} \big) \right) \right]. \label{eq:log-loss-is-kl} \end{gather}\tag{8}\] For completeness, see Lemma 5 in Appendix 9.2 for the derivation of the above identity. Informally, we say a sequence model has been successfully optimized for offline predictions if \(\ell(p_{\TN{per-action}}^*)-\ell(\widehat{p}_{\TN{per-action}}) \leq \epsilon\), for some small \(\epsilon > 0\).

Our main result formalizes how offline prediction quality controls the quality of generated samples, even under adaptive data collection. Recall that \(\HH_{t-1} \triangleq \big\{ Z, \, (A_1, Y_1),\ldots, (A_{t-1}, Y_{t-1}) \big\}\) denotes the history of observations available at the start of time \(t\), including unstructured prior information \(Z=\{Z^{(a)}\}_{a\in \MC{A}}\). The proof is given in Appendix 9.3.

Note that the expectation in the proposition is taken over the history under \(\pi\). This result relates offline and online performance. When the model achieves near-optimal prediction loss, the bound implies that in an information-theoretic sense it is difficult to distinguish our algorithm’s samples from true posterior samples. The second part of the proposition extends this guarantee to any derived quantity, like mean rewards or optimal actions, that we might compute from these samples.

This result formalizes two deep properties of our algorithm. First, it shows that incremental errors in autoregressive generation do not severely compound across iterations. Second, it shows that a single measure of offline sequence prediction loss directly controls the quality of posterior generations for any given process for generating the history \(\HH_{t-1}\). This may be surprising, as theory in reinforcement learning typically requires certain distributions under the policy used to collect offline data to nearly match those under the online policy [36]–[38].

5.2 Theory Part II: Exploration using approximate sequence models \(\widehat{p}_{\TN{per-action}}\) results in low regret↩︎

This section formally shows how successfully optimizing next-outcome prediction on an offline dataset reliably enables online exploration with low regret. Recall from Section 3, we measure the agent’s performance through expected per-period regret: \[\label{eq:regret} \Delta(\pi; \MC{B})= \E_{\pi}\left[ \max_{a\in \MC{A}} \frac{1}{T} \sum_{t=1}^{T} \left(R(Y_t^{(a)}) - R(Y_t^{(A_t)})\right)\right].\tag{9}\] The expectation above is taken over the joint draw of unstructured prior information and potential outcomes \((Z^{(a)}, Y_{1:T}^{(a)})_{a\in \MC{A}}\) from \(p^*_{\TN{per-action}}\) as well as any additional randomness in the actions selected by policy \(\pi\).

5.2.1 A general reduction: optimized simulators reflect true regret↩︎

Our analysis relies on interpreting an approximate sequence model as inducing a simulator for bandit tasks, as in the definition below. One could use such a simulator to train policies offline using reinforcement learning techniques, or, as we do, to perform online planning.

It is not hard to show that, according to this definition, the simulator induced by \(p_{\TN{per-action}}^*\) perfectly matches the true data-generating process. Let \(\Delta(\pi; \widehat{p}_{\TN{per-action}})\) denote the expected regret policy \(\pi\) incurs under the simulated environment and reserve \(\Delta(\pi; p_{\TN{per-action}}^*)\) for the true regret defined earlier in 9 .

The proof is in Appendix 9.4. Similar statements apply to simulated quantities beyond regret. Informally, if one were training a sequence model \(\hat{p}_{\TN{per-action}}\) from scratch, we expect \(\ell(\widehat{p}_{\TN{per-action}}) - \ell(p_{\TN{per-action}}^*)\) goes to zero as the size of the dataset used to fit \(\widehat{p}_{\TN{per-action}}\) grows, as long as the sequence model class chosen is sufficiently flexible.

As a corollary of the general reduction above, we attain a regret guarantee for our algorithm when deployed with an offline-optimized sequence model.

[39] prove a regret bound that superficially looks similar to ours, however their result only applies to binary outcomes and requires a specific conditional KL divergence to be small, which does not appear to follow from training a model with low validation loss. Moreover, we use a very different proof technique. Our Corollary 2 relies on Proposition 3, which is a result that may be of independent interest.

5.2.2 The price of informed exploration: prior sensitivity of Thompson sampling and the tightness of our regret bound.↩︎

Under the simulator induced by \(\widehat{p}_{\TN{per-action}}\), actions that the model believes are unlikely to be optimal given prior information \(Z\) are therefore effectively screened out and receive little direct exploration. As formalized in Section 4.2.2, in the regimes that motivate informed exploration—many candidate actions relative to the horizon, such as recommendation systems—this screening effect is precisely what can enable strong performance.

The natural question, then, is what this aggressive screening costs when the model’s initial beliefs are wrong. An action that is screened out generates no data that could correct the model’s mistaken beliefs about it, so such errors are never washed out by additional interactions, even over long horizons. The persistent impact of misspecified priors differs markedly from passive learning, where prior effects typically fade with more data. This intuition helps explain why the penalty term in ?? does not decay with \(T\), and why per-period regret need not vanish under misspecification. Indeed, this phenomenon is closely related to the known sensitivity of Thompson sampling to prior misspecification. Classical frequentist regret bounds for Thompson sampling establish sublinear cumulative regret under fixed priors with full support [40], which may suggest that the cost of prior misspecification is limited because the influence of the prior is eventually overcome by data. However, this interpretation can be misleading: worst-case analyses show that the sensitivity of Thompson sampling to prior misspecification can scale quadratically with the horizon [25], [41].³ The apparent tension is resolved by noting that classical regret guarantees fix a problem instance and study performance as the horizon grows, whereas worst-case misspecification analyses allow the hardest problem instance to change with the horizon. These hardest instances can place progressively less prior mass on the truth, increasing the regret incurred before the data can overcome the prior.

To contextualize our result within this literature, and to verify that the non-vanishing penalty in Proposition 3 reflects the problem rather than slack in our analysis, we specialize to the classical setting of Bayesian bandits with a misspecified prior. In this special case, autoregressive generation reduces to sampling from a posterior predictive distribution akin to 1 , and Corollary 2 yields a regret bound for Thompson sampling under prior misspecification. This provides a short information-theoretic route to a prior-sensitivity guarantee. In contrast, [25] study prior sensitivity for algorithms including a “\(k\)-shot” variant of Thompson sampling through an intricate analysis, showing that \(\varepsilon\) misspecification in total variation distance can inflate per-period regret by a term growing linearly with the horizon \(T\). In a few lines, our reduction shows that \(\varepsilon\) misspecification in KL divergence inflates per-period regret by a term independent of \(T\).

The bound ?? allows the per-period expected regret to remain inflated by a non-vanishing KL divergence term. An upper bound alone cannot show that this persistence is real rather than an artifact of the analysis; our lower bound does. In Appendix 9.7, we provide a lower bound example where even with minimal prior misspecification (differing only by a \(\log(T)\) term), the per-period regret remains constant, aligning with known lower bounds in related settings [41]. In this sense the \(T\)-independence of the penalty term in Corollary 2 is tight; we do not claim the bound is tight in all settings.

6 Related work↩︎

6.0.0.1 Distinctive contributions.

  While there has been intense study on sequence modeling, bridging foundation modeling and sequential decision making, and the use of neural networks in multi-armed bandit problems, our work makes distinct contributions that do not seem to have precedent in the literature. First, we focus on a new problem: informed meta-bandit problems where unstructured prior information (like text) must be leveraged to minimize exploration costs, with bandit tasks drawn from a task distribution. Second, in developing principled algorithms for this setting, we show that autoregressive generation of missing outcomes enables principled uncertainty quantification and exploration without explicit parameter inference. Third, our theoretical guarantees establish a powerful reduction, showing that success in offline sequence prediction directly translates to effective online decision-making. A subsequent conference paper by the same authors extends the Thompson sampling component of this work to contextual bandits with general policy classes [27]; see the end of Section 2 for a detailed discussion of the relationship.

6.0.0.2 Core ideas from statistical sequence prediction.

  As discussed in Section 2 and Figure 1, our work builds on a classical vision in statistics where inference is performed solely through potentially observable quantities, rather than through latent parameters. This perspective traces back to de Finetti’s work on exchangeable sequences [42] and appears across multiple literatures [21], [43]–[53]. Among these, our initial efforts on this direction were especially inspired by [53], who uses sequence prediction loss to evaluate the quality of epistemic uncertainty quantification. Several authors recently proposed neural network architectures for learning the joint distribution / posterior predictives [1], [2], [22]. While some of these papers include Bayesian bandit experiments, they do not use autoregressive generation to quantify epistemic uncertainty (see Section 7.5 for further discussion).

6.0.0.3 Sequential decision-making with foundation models.

  Recent work leveraging foundation models for sequential decision-making broadly falls into three categories [54]. The first involves training sequence models on expert demonstrations. Some of these works try to directly imitate the typical expert, while others uses goal-conditioned action selection to improve upon demonstrated behavior [55]–[57]; while effective in some settings, this approach can provably fail when reward signals are noisy [58]–[60]. The second category explores variants of in-context decision making—decision-making where a fixed sequence model adapts by conditioning on its interaction history with the current environment—where pretrained foundation models are carefully prompted without additional fine-tuning [61]–[63]; these include papers recent work that discuss the challenges of using LLMs for in-context exploration in bandit problems [5]–[7]. Another closely related work by [64] leverages TabPFN [3] for exploration; our approach differs from theirs because they quantify uncertainty using a sub-sampled central limit theorem, rather than through our autoregressive generation approach. The third category develops generative models as learned environment simulators [65]–[67]. Section 5.2 interprets our sequence models as simulators, but our work distinctively uses these simulator models for online autoregressive generation to quantify epistemic uncertainty and enable principled exploration.

6.0.0.4 Sequence modeling in bandits.

  A particularly relevant approach in this space is the Decision Pretrained Transformer (DPT) [68], which implements Thompson sampling by pretraining a sequence model to predict the reward-maximizing action. This elegant approach works by observing that optimal in-context prediction of the reward-maximizing action necessarily models its posterior distribution. Our work develops a broader framework for uncertainty quantification through autoregressive generation of missing outcomes, an idea that is applicable well beyond Thompson sampling. While the works share connections in implementing Thompson sampling, their core insights and implementations differ substantially, each with distinct advantages. DPT learns a “generative model of behavior” [54]—predicting the best action based on past histories. In contrast, we learn a “generative model of the world” [54]—predicting reward/outcome sequences for each action individually. We conjecture this is key to our unique theoretical guarantees linking offline sequence modeling to online performance. Given these complementary strengths, it would be interesting to explore hybrid approaches, where our models could serve as offline simulators to distill DPT-style policies, which reduce online computation.

Autoregressive predictive models are also employed by [69], but with very different goals and insights. Their work recognizes that Thompson sampling can over-explore in nonstationary environments, where constant environmental changes introduce fresh uncertainty that may not be worth resolving. They address this through a clever conditioning of imagined reward trajectories in their Bayesian model. While our work shares the use of generating potential rewards, we do so with an entirely different objective—replacing explicit Bayesian inference with a more scalable approach rather than moderating Thompson sampling’s exploration. Later works that approximate predictive sampling with neural networks [70] rely on deep ensembles trained via online stochastic gradient descent, whereas our methods do not train/fine-tune the sequence model at decision-making time; our sequence models are fixed prior to the line decision-making phase and learn in-context.

6.0.0.5 Other approaches to neural network inference in bandit problems.

  Beyond approaches that leverage sequence modeling, there are several approaches to implement Thompson sampling with neural networks. The first places a Bayesian prior on the neural network weights themselves, either through Bayesian linear regression on the final layer [71], [72] or full Bayesian neural networks [73]. A second thread uses ensembles of neural networks to approximate posterior sampling [74]–[76], including more efficient variants like Epinets [53], [77], [78] and HyperModels [79], [80] that reduce computational overhead. Our work offers a different perspective entirely: rather than placing uncertainty on network parameters or using ensembles, we model uncertainty through autoregressive generation of missing outcomes, allowing us to leverage the full power of modern sequence models while maintaining principled uncertainty quantification. Moreover, we emphasize informed meta-problems where leveraging rich priors from unstructured information is crucial. This setting poses fundamental challenges for existing approaches. Ensemble methods typically rely on bootstrapping data to inject variance, but at the start of a decision problem, there is no data to bootstrap, leaving no basis for principled exploration. While randomized prior functions [74] attempt to address this through random initialization, our approach is fundamentally different—extracting meaningful priors implicitly through training and scaling laws, and updating beliefs by through in-context learning rather than incremental gradient descent.

6.0.0.6 Meta-bandit learning.

  Prior work on meta-bandit learning has primarily focused on parametric settings, studying how to fit hyperparameters for Thompson sampling [81]–[83] or linear UCB algorithms [84]. While these approaches demonstrate convergence to oracle performance given sufficient meta-examples, they differ fundamentally from our work which tackles informed exploration with unstructured information and relies on autoregressive generation rather than parameter inference. An alternative approach learns exploration policies from scratch using reinforcement learning, potentially discovering novel exploration strategies but current approaches require more complex training procedures and larger sample sizes than our sequence prediction framework [85]–[88].

7 Experiments↩︎

We evaluate our algorithm in a synthetic setting and in a semi-realistic news recommendation setting. We focus on settings with binary outcomes and rewards (\(R(Y_t)=Y_t \in \{0, 1\}\)) since the news dataset we build on has binary click/no-click outcomes. Both settings satisfy the structural assumptions introduced in Section 4.3, so that the version of TSAR implemented is Algorithm 4 using Algorithm 6 as a subroutine. Additionally, in both settings, the methodology for obtaining \(\widehat p\) is as described in Section 7.1. Additional details and results are in Appendix 10.

In Section 7.2 we use a synthetic setting to validate our approach and see that with our pretraining strategy, our algorithm almost exactly matches the decision-making performance of a classical version of Thompson sampling with a correctly specified Bayesian model. Then, in Section 7.3, we demonstrate our method scales to a news recommendation setting where articles are associated with text headlines, and where \(\widehat p\) incorporates a language model.

7.1 Learning an approximate sequence model \(\widehat{p}_{\TN{per-action}}\) from historical data↩︎

In this section, we present an approach to learning an approximate sequence model \(\widehat{p}_{\TN{per-action}}\) using historical bandit task data under Assumptions 1 (independence across actions) and 2 (exchangeability) as described in Section 4.3. We will use this approach throughout our experiments. Note that this is just one approach among many to form \(\widehat{p}_{\TN{per-action}}\). Some additional ways to obtain \(\widehat{p}_{\text{per-action}}\) are discussed in Section 2.3.

We assume the agent has access to a large amount of historical data \(\MC{D}^{\TN{offline}}\) from past bandit tasks, with data also drawn from \(p_{\TN{per-action}}^*\). This historical data is used to form \(\widehat{p}_{\text{per-action}}\) in an offline learning phase, prior to decision-making. Note that this historical data is distinct from recent history within the current online bandit task, \(\mathcal{H}_{t-1}\), as defined in 2 . Specifically, this historical data consists of \[\begin{align} \MC{D}^{\TN{offline}}\triangleq \big\{ \big( Z^{(a)}, Y_{1:T}^{(a)}\big) \}_{a\in \mathcal{\MC{A}}^{\TN{offline}}} \qquad \TN{where} \qquad \big( Z^{(a)}, Y_{1:T}^{(a)}\big) \iidsim p_{\TN{per-action}}^* \end{align}\] and \(a\in \MC{A}^{\TN{offline}}\) indexes a large historical set of previous actions, their features \(Z^{(a)}\), and their potential outcomes \(Y_{1:T}^{(a)}\), with actions \(a\) sampled across many previous bandit tasks. In practice, to augment the training data, and also to promote exchangeability (Assumption 2), we also train on bootstrapped samples \(\tilde{Y}_{1:T}^{(a)}\) formed as follows: \[\begin{align} \tilde{Y}_1^{(a)}, \dots, \tilde{Y}_T^{(a)} \mid Y_{1:T}^{(a)} \iidsim \frac{1}{T} \sum_{t=1}^T \delta_{Y_t^{(a)}}. \label{eqn:bootstrap} \end{align}\tag{10}\]

Note that by assuming we have access to \(Y_{1:T}^{(a)}\) for each action \(a\), we are assuming we have access to “complete” historical bandit task data, in the sense that we observe all potential outcomes for any given action. If such “complete” historical bandit data is not available, in practice, during training, we can impute a full dataset by bootstrap resampling rewards across timesteps \(t\), where the bootstrap resampling is with replacement and done independently for each action.

Using \(\MC{D}^{\TN{offline}}\), we fit a sequence model \(\widehat p_{\TN{per-action}}\) to approximate \(p_{\TN{per-action}}^*\) by minimizing the empirical version of the log-loss of \(\widehat{p}\) from 7 : \[\begin{align} \label{eq:train95loss} \ell(\widehat p_{\TN{per-action}}; \MC{D}^{\TN{offline}}) \triangleq - \sum_{a \in \MC{A}^{\TN{offline}}} \sum_{t=1}^T \log \widehat p_{\TN{per-action}} \big( Y_t^{(a)} \mid Z^{(a)}, Y_{1:t-1}^{(a)} \big). \end{align}\tag{11}\] Figure 8 visualizes the training objective, and Algorithm 9 describes how to optimize \(\widehat p_{\TN{per-action}}\) via stochastic gradient descent-based approaches.

7.1.1 Sequence model architectures↩︎

In our experiments, we focus on settings where \(Z^{(a)}\) consists of either numeric or text features (e.g., headlines in a news recommendation setting). The choice of architecture for the sequence model \(\widehat p_{\TN{per-action}}\) is flexible. Our theoretical results in Sections 5.1 and 5.2 establish that any architecture achieving strong predictive performance is suitable. Below, we describe two architectural approaches that we use, which offer different tradeoffs; this is by no means an exhaustive list of potential implementations for \(\widehat{p}_{\text{per-action}}\) (for more examples, see Section 2.3).

7.1.1.1 Structured Bayesian models.

  When strong domain knowledge exists about the outcome distribution, we can incorporate it into the model architecture. For binary outcomes (\(Y_{t}^{(a)}\in \{0,1\}\)), we can use a Beta-Bernoulli model where a neural network maps \(Z^{(a)}\) (e.g. text features) to the Beta prior parameters \(\widehat \alpha(Z^{(a)})\) and \(\widehat \beta(Z^{(a)})\): \[\label{eq:ppd} \widehat p_{\TN{per-action}} \big( Y_{t+1}^{(a)} = 1 \mid Z^{(a)}, Y_{1:t}^{(a)} \big) = \frac{\widehat \alpha(Z^{(a)}) + \sum_{j=1}^{t} Y_j^{(a)}}{\widehat \alpha(Z^{(a)}) + \widehat \beta(Z^{(a)}) + t}\tag{12}\] Here, \(\widehat\alpha(Z^{(a)})\) and \(\widehat \beta(Z^{(a)})\) are learned functions that transform text features into Beta prior parameters and 12 is the corresponding “posterior predictive” distribution. This architecture explicitly encodes conjugate Bayesian updating while allowing the prior to be fit from data in the style of empirical Bayes methods [92]. For a more thorough discussion on the connection to Bayesian sequence modeling, see Appendix 11.

7.1.1.2 Simple neural network architecture.

  One approach uses separate encoders for prior information \(Z^{(a)}\) and outcome sequences \(Y_{1:t}^{(a)}\), whose outputs are combined by a final multi-layered perceptron (MLP) to form predictions. The encoder for \(Z^{(a)}\) depends on the input type: a language model (e.g. DistilBERT, [93]) for text, or an MLP for numeric inputs. This encoder and the final MLP are fine-tuned end-to-end. In contrast, the encoder for outcomes is a fixed function extracting sufficient statistics, which is appropriate because the outcome distribution has known sufficient statistics. Figure 10 visualizes this architecture for the case where \(Z^{(a)}\) is text (news experiments, Section 7.3); the synthetic experiments (Section 7.2) in which \(Z^{(a)}\) is a vector use the same architecture except with an MLP as the \(Z^{(a)}\) encoder.

7.2 Synthetic setting: Mixture Beta-Bernoulli↩︎

Despite never defining a Bayesian model with explicit latent parameters, we find that our Thompson sampling via autoregressive generation (TSAR) algorithm matches the performance of an oracle Thompson sampling method that has access to the true prior and likelihood. We show this in a synthetic environment where the true data-generating process is known. The true data-generating process uses a Bernoulli likelihood for outcomes, with a mixture of Beta distributions as the prior, where the Beta parameters are determined by \(Z^{(a)} \in \real^2\). See Appendix 10.1 for more details.

7.2.1 Decision-making algorithms↩︎

7.2.1.1 TSAR variants.

  We implement TSAR with two different types of sequence models as described in Section 7.1.1:

Both models are trained offline using gradient descent on loss 11 and later used to power online decision-making under new bandit tasks (details in Appendix 10.1).

Our primary comparison is between TSAR Flexible NN, which uses the trained Flexible NN for autoregressive generation, and TS Oracle, which performs Thompson sampling with exact Bayesian inference using the true generative model.

7.2.1.2 Truncating generation lengths.

  For large population sizes \(T\), generating all \(T\) missing outcomes during online decision-making (Algorithm 6) is computationally expensive. We therefore modify PSAR to generate only \(m=500\) missing outcomes per action and use their average to estimate the mean reward \(\mu^{(a)}(\tau)\) from 4 . Our experiments show this approximation remains effective when \(m\) is sufficiently large (see Section 7.4).

7.2.1.3 Baselines.

  We consider two additional Thompson Sampling (TS) based baselines, none of which match the performance of the oracle Bayesian procedure.

Finally, we also consider the UCB algorithm from Section 6 of [94]. Full details on baselines are in Appendix 10.3.

7.2.2 Simulation Results↩︎

7.2.2.1 TSAR with a well-fitted sequence model has regret close to the TS Oracle: Figure 11 (Left).

  TS Oracle achieves the lowest regret, with TSAR Flexible NN performing nearly as well. Methods using misspecified unimodal priors—TSAR Beta-Bernoulli NN and TS Beta-Bernoulli (Uniform Prior)—achieve intermediate performance, significantly better than both UCB and Thompson sampling with a misspecified Bayesian model.

7.2.2.2 Our posterior sampling via autogressive generation (PSAR, Algorithm 6) procedure with a well-fitted sequence model accurately quantifies epistemic uncertainty: Figure 11 (Right).

  For each of 1000 unseen articles, we generate 250 posterior samples \(\mu^{(a)}(\widehat{\tau}_1)\) of the mean action reward \(\mu^{(a)}(\tau)\) by autoregressively sampling outcomes conditioned \(Z^{(a)}\). Using quantiles of these samples, we construct credible intervals and measure coverage, i.e., how often the true \(\mu^{(a)}(\tau)\) falls within these intervals. The Flexible NN nearly matches TS Oracle’s theoretically ideal performance in both coverage and interval width. Methods using misspecified Beta priors show substantially worse coverage, even with fitted parameters. While these results reflect inference conditioned only on \(Z^{(a)}\) without observations of \(Y\), similar patterns hold when conditioning on additional observations (see Appendix 10.1).

7.3 News recommendation setting↩︎

Next, we build a semi-realistic news recommendation task using the MIcrosoft News Dataset (MIND) [95]. Our Thompson sampling via autoregressive generation algorithm scales effectively to scenarios where performance depends on unstructured information like article headlines. The method maintains accurate uncertainty quantification and low-regret exploration even when incorporating and finetuning a language model.

For each article \(a\), \(Z^{(a)}\) represents either headline text (in our best methods) or categorical information like “politics” or “sports” (in simpler variants). The dataset contains approximately 11,000 articles after preprocessing, with binary click/no-click rewards.

7.3.1 Decision-making algorithms↩︎

7.3.1.1 TSAR variants.

  We consider TSAR with two different types of sequence models: (i) Flexible NN (Text) and (ii) Beta-Bernoulli NN (Text) are analogous to those from Section 7.2, but we modify them to use article text \(Z^{(a)}\) embedded using DistilBERT [93], which is fine-tuned end-to-end during the offline stage. (iii) Flexible NN (Category) treats \(Z^{(a)}\) as a one-hot encoding of category information (e.g. “politics" or”sports") instead of headline text. As in Section 7.2, we generate \(m=500\) missing outcomes per action and use their average to estimate \(\mu^{(a)}(\tau)\), instead of generating all \(T\) missing outcomes to reduce computation costs.

7.3.1.2 Baselines.

  To our knowledge, this kind of informed bandit problem is new and lacks established baselines. We therefore adapt several popular approaches from related areas. TS Neural Linear implements Thompson sampling with a Gaussian linear regression, where the regression features consist of the output from a neural network trained to predict outcomes given DistilBERT article embeddings as input, adapting ideas from neural contextual bandits [71]. We also implement an ensemble of 50 neural network models following [96] and [97] as baselines for uncertainty quantification at the \(t=0\) timestep. For comparison to methods that ignore prior information \(Z^{(a)}\), we include both the UCB algorithm in Section 6 of [94] and TS Beta-Bernoulli (Uniform Prior), which performs exact Thompson sampling with a misspecified uniform prior.

7.3.2 Simulation Results↩︎

7.3.2.1 PSAR provides accurate uncertainty quantification with language models in the loop: Figure 12 (right).

  A longstanding challenge in ML is implementing uncertainty quantification with neural networks, especially with large foundation models in the loop [71], [98]. Using the same methodology as Section 7.2, we evaluate uncertainty quantification by constructing intervals for the mean action reward across 2,280 articles held out from the offline training data. Notably, all PSAR variants achieve the nominal coverage level out of the box without any post-hoc adjustments to improve calibration. Also note that the text-based models produce slightly narrower intervals compared to ones that do not use text information. See Figure 18 in Appendix 10.2 for uncertainty quantification evaluations for posterior sampling via autoregressive generation.

Many of the baseline approaches have extremely poor coverage, including deep ensembling, which is often considered a state-of-the-art approach for uncertainty quantification with neural networks. We believe ensembling performs poorly in our informed exploration setting where, loosely speaking, the relevant uncertainty is the mean action reward for this specific action, rather than the mean action reward for actions with this action’s prior information \(Z^{(a)}\). Ensembling focuses on uncertainty for the latter quantity. Additional details and analysis are provided in Appendix 10.2.

7.3.2.2 TSAR using headline text has the best regret, but gains are modest: Figure 12 (left).

  TSAR methods that incorporate headline text achieve the lowest regret among all approaches. However, comparing our best TSAR variant against TS Beta-Bernoulli (Uniform Prior) demonstrates modest gains from leveraging offline data and article text. While these results confirm such gains are achievable, this dataset likely understates the potential value of unstructured prior information; in some applications, recent work suggests such information alone can drive effective decision-making [99]. TS Neural Linear perform poorly in terms of regret, showing that simply adding neural-network-based predictions to standard Thompson sampling, for instance through a Gaussian linear model on top of trained neural network outputs, is not enough to ensure good bandit performance.

7.4 Examining the impact of truncating generation length↩︎

7.4.0.1 Empirically comparing full imputation vs. truncated generation (Figure 13 (a)).

  We empirically compare our proposed method (Algorithm 4), i.e., “full imputation”, with the computationally cheaper version of Thompson sampling with autoregressive generation that limits generation to \(m=500\) future timesteps, on the news setting (Section 7.3). We follow the procedure described in Appendix 10.2 in forming the regret plots: we run \(500\) repetitions of each bandit algorithm and in each repetition we draw a new set of \(10\) actions/articles from the validation set to represent a “new task”. We find that the two versions have similar performance.

7.4.0.2 Varying the truncation generation length (Figure 13 (b)).

  Throughout all our previous experiments (aside from Figure 13 (a)), we set \(m=500\). In Figure 13 (b), we examine the impact of varying \(m\) on the regret of the algorithm with the Flexible NN (text) sequence model in the news recommendation setting. As before, we follow the procedure described in Appendix 10.2 in forming the regret plots. We find that increasing \(m\) reduces the regret of the algorithm; however, when \(m\) is sufficiently large, the benefit of increasing \(m\) further is negligible.

7.5 Important subtleties in autoregressive generation↩︎

Our implementation of autoregressive generation in Section 7.1 makes three critical choices:

1. Potential outcomes are computed autoregressively, meaning we condition each newly generated outcome on those previously generated. This captures persistent patterns in performance, where actions that have demonstrated strong results are more likely to continue performing well.

2. Autoregressive generation is performed through random sampling rather than selecting the most likely next outcome. This preserves the model’s implicit understanding of uncertainty, whereas always choosing the most likely outcome does not.

3. At each decision-time, instead of sampling a single reward per action, we consider an aggregate statistic: the average of the imputed rewards over time. This approach better captures uncertainty about an action’s overall performance, which is most relevant for exploration.

These implementation choices are critical to the algorithm’s correctness. Figure 14 demonstrates that two alternative approaches, including one recently proposed in several published papers (to be discussed shortly), lead to poor decision-making by either over- or under-exploring. Each tested variant modifies the steps in our proposed algorithm where we average autoregressively sampled and observed rewards to form a “posterior sample" of the mean reward, while keeping the remainder of the algorithm unchanged. See Appendix 10.1.1 for additional details on the experimental setup. Our analysis reveals two key findings:

1. The importance of averaging out idiosyncratic noise. The method labeled Generate one reward implements a proposal from recent works [1], [2], [22]. These works propose generating only the outcome/reward from the next action draw, and taking the action with the highest sample reward. While this captures overall predictive uncertainty in the next outcome—including both uncertainty in an action’s overall quality (epistemic uncertainty) and uncertainty in the idiosyncratic noise of a single outcome (aleatoric uncertainty)—it fails in settings with substantial idiosyncratic uncertainty.⁴ Figure 14 shows this approach incurs linear regret due to endless randomized exploration.

2. The importance of autoregressive generation. The method labeled Generate rewards i.i.d. averages rewards across many independent (non-autoregressive) draws of the next outcome and takes the action with the highest average sampled reward. With sufficient generations, this average converges to the mean of the predictive distribution, effectively selecting the action currently believed to be best without purposeful exploration. Such “greedy" algorithms perform poorly in the long run because they fail to gather the data needed to resolve uncertainty about action quality.

8 Discussion↩︎

By viewing uncertainty through the lens of missing data rather than latent parameters, our work establishes a principled bridge between modern sequence models and exploration algorithms, demonstrating through algorithms, theory, and experiments that success in offline sequence prediction translates to effective online decision-making and uncertainty quantification. We believe our missing data view of uncertainty extends well beyond the setting studied in this work, including to sequences with predictable time-dependence. We are optimistic that operationalizing this viewpoint beyond the setting studied in this paper will bear fruit.

Several important challenges remain open. While enormous investments in both software and hardware are going toward speeding up autoregressive generation (e.g. for language model inference), we still expect that repeatedly generating posterior samples online can be computationally expensive. Future work could address this computational overhead by distilling exploration policies into neural networks through imitation learning [100]. Second, our experiments assume access to sufficient historical data for offline learning. In practice, available data may be limited and partially confounded by previous action selection policies. While many deployed systems face and overcome similar challenges, carefully evaluating our methods’ robustness to these issues at scale remains important.

8.0.0.1 Overview of Appendices

• Appendix 9: Theoretical results

• Appendix 10: Experimental details

• Appendix 11: Bayesian sequence models

• Appendix 12: Finite vs infinite population formulations and Thompson Sampling variants

9 Theoretical results↩︎

9.1 Proof of Proposition 1: Regret of exact Thompson Sampling by sequential imputation↩︎

9.2 Derivation of Equation 8 : KL divergence representation of the offline loss function↩︎

The next lemma is a standard result connecting the excess expected loss of a sequence model \(n\) to its KL divergence from the true sequence model \(p^*\). In this section, for notational convenience, we use \(p^*\) and \(\widehat{p}\) to denote \(p_{\TN{per-action}}^*\) and \(\widehat{p}_{\TN{per-action}}\), respectively. Recall, the expected loss of a sequence model \(\widehat{p}\) is denoted \(\ell(\widehat{p})\), defined in 7 . To (nearly) minimize loss, \(\widehat{p}\) the learner needs to closely approximate the true sequence model \(p^*\).

9.3 Proof of Proposition 2↩︎

Consider the full process of generating a table \(\widehat{\tau}_t\). This involves first i) the sequential revelation of \(t-1\) (real) potential outcomes determined by selected actions under \(\pi\) and ii) the sequential revelation of \(T\times |\MC{A}|-(t-1)\) (imagined) potential outcomes produced by the autoregressive generation procedure in Algorithm 6 (see also Figure 7.).

We introduce an abstraction which subsumes both kinds of revelation. The revelation of a full potential outcome table proceeds across \(|\MC{A}|\times T\) rounds; we use the letter \(i\) to index these rounds, \(a\)’s to index rows of the potential outcomes table, and \(j\)’s to index columns. Take \((A_i, J_i) \in \MC{A}\times [T]\) to denote the index of the entry selected in round \(i\) and \(\widehat{\tau}[A_i, J_i]\) to denote an entry itself in sampled potential outcomes table. For \(i\leq t-1\), the column is simply \(J_i =i\) and \(A_i\) is determined adaptively according to the policy \(\pi\); For \(i\geq t-1\), the choice of \((A_i, J_i)\) proceeds across missing entries according to the deterministic order of Algorithm 6.

We can represent this procedure abstractly as an adaptive selection rule \(\psi = (\psi_1, \ldots, \psi_{|\MC{A}|T})\). For any \(i\), \(\psi_i\) is a function which selects \(A_i,J_i\) as a function of the revealed/generated outcomes so far and possibly some exogenous noise. Writing it formally, let the (generalized) history \(\HH^+_{i-1} \in \mathcal{Z}^{|\MC{A}|} \times \left( \MC{A}\times [T] \times \mathcal{Y}\right)^{i-1}\) denote the sampled unstructured inputs \((Z's)\) together with a sequence of table indices and entries; take \(\xi_i\) to be a random seed, independent of all else. The procedure described above can be thought of as corresponding to a rule which determines selections \((A_i, J_i)=\psi_i(\HH_{i-1}, \xi_i)\) for each \(i \leq |\MC{A}|T\). While the \(\psi\) we are studying has many special properties, our analysis does not rely on these. Instead, our proof leverages only one property of \(\psi\): that it never selects a previously revealed entry of the table and hence reveals all entries of the table after \(|\MC{A}|T\) revelations. We call this such a selection rule non-repeating.

The proof compares three procedure for sampling outcomes of selected entries conditioned on prior selections and samples.

1. True environment: Conditioned on \(\HH^{+}_{i-1}\) and \((A_i, J_i)\), the revealed outcome \(\widehat{\tau}[A_i, J_i]\) is drawn according to \(p^*\). That is, \(\widehat{\tau}[a,j] \mid \HH^{+}_{i-1}, A_i=a, J_i=j \sim p^*\left( \cdotspace \mid Z^{(a)}, (\widehat{\tau}[A_k, J_k])_{k<i, A_k=a} \right)\). We denote probabilities under \(\psi\) employed in this environment by \(\mathbb{P}_{*,\psi}\).

2. Misspecified environment: Conditioned on \(\HH^{+}_{i-1}\) and \((A_i, J_i)\), the revealed outcome \(\widehat{\tau}[A_i, J_i]\) is drawn according to \(\widehat{p}\). That is, \(\widehat{\tau}[a,j] \mid \HH^{+}_{i-1}, A_i=a, J_i=j \sim \widehat{p}\left( \cdotspace \mid Z^{(a)}, (\widehat{\tau}[A_k, J_k])_{k<i, A_k=a} \right)\). We denote probabilities under \(\psi\) employed in this environment by \(\mathbb{P}_{\theta,\psi}\).

3. Hybrid environment: Conditioned on \(\HH^{+}_{i-1}\) and \((A_i, J_i)\), the revealed outcome \(\widehat{\tau}[A_i, J_i]\) is drawn according to \(p^*\) if \(i<t\) and drawn according to \(\widehat{p}\) otherwise. We denote probabilities under \(\psi\) employed in this environment by \(\mathbb{P}_{h,\psi}\). This represents the procedure used to generate \(\widehat{\tau}_{t}\), where the first \(t-1\) entries are revealed according the true environment, and the remainder are drawn by using Algorithm 6 with \(\widehat{p}\) for generation.

With the abstractions and notations above, Proposition 2 can be rewritten as \[\label{eq:posterior95sample95bound95rewritten} \mathbb{E}_{\psi, *} \left[ D_{\rm KL}\left( \PP_{\psi, *}\left( \widehat{\tau} \in \cdot \mid \HH^+_{t-1} \right) \, \| \, \PP_{\psi, h}\left( \widehat{\tau} \in \cdot \mid \HH^+_{t-1} \right) \right) \right] \leq |\MC{A}| \left( \ell(p^*) - \ell(\widehat{p}) \right).\tag{14}\] We establish this result in three steps. The first bounds the KL divergence on the left hand side of 14 by the KL divergence between full histories.

The next step, stated as a standalone lemma, bounds the KL divergence between the true and hybrid environments by the KL divergence between the true and misspecified environments.

Our next lemma calculates the KL divergence between (generalized) histories under the true and misspecified environments, finding it is exactly equal to the number of actions times the excess offline prediction loss of the misspecified sequence model \(\widehat{p}\).

9.4 Proof of Proposition 3 (Section 5.2): Bounding the deployment regret in terms of regret on a simulator↩︎

9.4.1 Implementing the simulator with apriori randomness↩︎

In this section, for notational convenience, we use \(p^*\) and \(\widehat{p}\) to denote \(p_{\TN{per-action}}^*\) and \(\widehat{p}_{\TN{per-action}}\), respectively. The simulator corresponding to sequence model \(\widehat{p}\), formalized in Definition 2, appears to involve sequential draws of outcomes as actions are selected. An alternative, which is useful in the analysis and potentially in implementation, is to draw all randomness upfront. Doing so in a manner that corresponds to the simulator involves a slightly different potential outcomes table, denoted \(\tilde{\tau}\), in which columns are interpreted as indexing arm play-counts rather than time.

We emphasize that this construction is used solely as proof technique, and is not part of the model of the problem. In particular, it is used for studying simulators corresponding to misspecified non-exchangeable sequence models \(\widehat{p}\). Under exchangeable sequence models, like \(p^*\), order has no bearing.

9.4.2 Completing the proof↩︎

9.5 Proof of Corollary 2↩︎

In light of Proposition 3, the following proof is largely a review of an information-theoretic analysis of Thompson sampling due to [33]. It was observed by [34], [35] that this analysis can be applied without modification to analyze regret with respect to the best fixed action (\(A^*\)) even in nonstationary environments (e.g. non-exchangeable models \(\widehat{p}\) as in Definition 3.) In this section, for notational convenience, we use \(p^*\) and \(\widehat{p}\) to denote \(p_{\TN{per-action}}^*\) and \(\widehat{p}_{\TN{per-action}}\), respectively.

9.6 Proof of Corollary 1↩︎

9.7 The tightness of our regret bound with misspecified priors↩︎

In this section, for notational convenience, we use \(p^*\) and \(\widehat{p}\) to denote \(p_{\TN{per-action}}^*\) and \(\widehat{p}_{\TN{per-action}}\), respectively. In Example 2 we showed the following bound on regret under Thompson sampling applied with a misspecified prior: \[\begin{align} \label{eq:mispecified-prior-regret-repeat} \Delta \big( {\pi_{\TN{TSAR}}}(\widehat{p} ); \, p^* \big) \leq \underbrace{ \sqrt{ \frac{ |\MC{A}| \log (|\MC{A}|) }{2 T} } }_{\TN{Regret bound for Thompson sampling}} + \underbrace{ \sqrt{ \frac{ |\MC{A}| }{2} D_{\rm KL}\big( \nu_{*} \; \big\| \; \nu_{\theta} \big)} }_{\TN{Penalty for mispecified prior}}. \end{align}\tag{20}\] As noted in the discussion following the example, our upper bound suggests that misspecified priors have persistent effects that do not diminish with the time horizon. The per-period expected regret remains inflated by a non-vanishing KL divergence term. Here we show this is unavoidable in general.

10 Experimental details↩︎

• Appendix 10.1 provides details on synthetic experiments

• Appendix 10.2 provides details on news article recommendation experiments

• Appendix 10.3 provides details on bandit baseline algorithms

• Appendix 10.4 discusses modeling best arm rather than arm reward

10.1 Synthetic Experiments: Mixture Beta-Bernoulli↩︎

10.1.0.1 Data Generating Process.

In this setting, we use article attributes \(Z^{(a)} = \big( Z^{(a)}_1, Z^{(a)}_2 \big) \in \real^2\) where \(Z^{(a)}_1, Z^{(a)}_2 \iidsim \TN{Uniform}(0, 0.25)\). We sample \(Y_{1:T}^{(a)}\) by first sampling a click rate \(\mu^{(a)} \in [0, 1]\) from a mixture: \[\begin{align} \mu_{\infty}^{(a)} \mid Z^{(a)} \sim \begin{cases} \TN{Beta} \big( 25 \cdot \frac{1}{4} Z_1^{(a)} + 1, ~ 25 \cdot (1-\frac{1}{4}Z_1^{(a)}) + 1 \big) & \TN{w.p. } 1/2 \\ \TN{Beta} \big( 25\cdot (1-\frac{1}{4}Z_2^{(a)}) + 1, ~ 25\cdot \frac{1}{4}Z_2^{(a)} + 1 \big) & \TN{w.p. } 1/2 \end{cases} \end{align}\] Then, outcomes are sampled as \(Y_1^{(a)}, \dots, Y^{(a)}_T \mid \mu_{\infty}^{(a)}, Z^{(a)} \iidsim \TN{Bernoulli}(\mu_{\infty}^{(a)})\). We use the identity reward mapping \(R(y) = y\).

10.1.0.2 Additional training details.

The training and validation datasets contain 2500 and 1000 actions each, respectively. During offline learning (Section 7.1), we train up to sequence lengths of \(500\). Hyperparameters and early stopping epochs are chosen using the validation dataset.

• Flexible NN. This sequence model \(p_\theta\) is based on a multi-layer perceptron (MLP). This model takes as input the action attribute \(Z^{(a)}\) and a summary statistic that represents the history of outcomes for that action, and outputs a probability in \([0,1]\). The summary statistic we use is simple because outcomes \(Y_t^{(a)}\) are binary. Specifically, it consists of a tuple with the mean of outcomes from action \(a\), and the reciprocal of 1 plus the total number of outcome observations for action \(a\), i.e., \(\big( \frac{1}{N^{(a)}} \sum_{t'=1}^{t-1} Y_{t'} \ind_{A_{t'} = a}, ~ \frac{1}{ 1 + N^{(a)} } \big)\), where \(N^{(a)}\triangleq\sum_{t'=1}^{t-1} \ind_{A_{t'} = a}\).⁵

• Beta-Bernoulli NN. This is a sequential model based on the closed-form posterior predictive for a Beta-Bernoulli model. The prior parameters for the Beta distribution, \(\alpha_\theta(Z^{(a)})\) and \(\beta_\theta(Z^{(a)})\), are each parameterized by separate neural network MLP models that take in \(Z^{(a)}\) and output a non-negative scalar.

The MLPs for both models each have three layers (width \(50\)) and use ReLU activations. Both models are trained for 1000 epochs with learning rate 0.001, batch size 500, and weight decay 0.01 using the AdamW optimizer.

10.1.0.3 Posterior width and calibration plots at additional timesteps (Figure 16).

We display posterior calibration and posterior widths as in Figure 11 (right), but for additional timesteps beyond \(t=0\). To generate posterior samples for Beta-Bernoulli NN, we use the closed-form posterior.

10.1.1 Additional details on Figure 14: Comparing to alternative sampling approaches.↩︎

Figure 14 considers the synthetic decision-making setting described in this section. Figure 14 considers the synthetic decision-making setting described in this section. All decision-making algorithms in Figure 14 use the same pretrained sequence model \(p_\theta\), with the Flexible NN architecture and trained as described in Section 7.2 and Appendix 10.1.

• Greedy. Select \(A_t = \argmax_{a \in \MC{A}} \, \widehat{\mu}_{t,a}(p_\theta)\) where \(\widehat{\mu}_{t,a}(p_\theta) = \int_y R(y) \cdot p_\theta( Y_t^{(a)} = y \mid \HH_{t-1}) dy\). In a setting with binary outcomes and \(R(y)=y\) such as ours, \(\widehat\mu_{t,a}(p_\theta)=p_\theta(Y_t^{(a)}=1\mid \mathcal{H}_{t-1})\).

• Epsilon-greedy. Select the action chosen by Greedy 90% of the time, and a uniformly random action for the remaining 10% of the time.

• Generate one reward. For each action \(a \in \MC{A}\), sample \(\widehat{Y}_t^{(a)} \sim p_\theta( Y_t^{(a)} \in \cdotspace \mid \HH_{t-1})\) and select \(A_t = \argmax_{a \in \MC{A}} \, R (\widehat{Y}_t^{(a)})\).

• Generate rewards i.i.d. For each action \(a \in \MC{A}\), sample
  \(\widehat{Y}_t^{(a)}, \dots, \widehat{Y}_{t+m}^{(a)} \iidsim p_\theta( Y_t^{(a)} \in \cdotspace \mid \HH_{t-1})\) for \(m=500\) and select \(A_t = \argmax_{a \in \MC{A}} \, \frac{1}{m} \sum_{i=0}^m R (\widehat{Y}_{t+m}^{(a)})\).

• Generate rewards autoregressively (TSAR). Probabilistically impute the table of potential outcomes using Algorithm 6, but truncated to 500 imputed values as described in Section 7.4, to form \(\widehat{\tau}_t\). Then select \(A_t=A^*(\widehat{\tau}_t)\).

10.2 News Recommendation Experiment Details↩︎

10.2.0.1 Dataset preprocessing.

  We process the MIcrosoft News Dataset (MIND) as follows:⁶

1. Collect all articles from the MIND “large” dataset (training split only) [95].

2. Remove any article with fewer than 100 total impressions.

3. For each article, we use \(\mu^{(a)}\) to denote its empirical click rate. We normalize the click rates to be centered around 0.5 in a way that preserves their relative rank. We do this transformation to speed up the learning, as it requires more data to accurately learn Bernoulli’s with success probabilities close to 0 or 1.

   We transform as follows: Let \(\mu_0^{(1)},\ldots,\mu_0^{(|\MC{A}|)}\) be the original empirical success probabilities (average click rate). We use \(\MC{A}\) to denote all articles in the MIND large dataset. The new success probabilities are defined as follows for each \(a \in \MC{A}\): \[\begin{align} \mu^{(a)} \gets \begin{cases} \mu_0^{(a)} & \TN{if~} \mu_0^{(a)} \in \{ 0, 1 \} \\ \text{logit}^{-1} \left(\text{logit}(\mu_0^{(a)}) - \bar{\mu}_0 \right) & \TN{otherwise} \end{cases}. \end{align}\] Above, \(\bar{\mu}_0 \triangleq \frac{1}{|\MC{A}|} \sum_{a' \in \MC{A}} \text{logit}(\mu_0^{(a')})\) and \(\TN{logit}(x) \triangleq \log \frac{x}{1-x}\). See Figure 17 for comparison of the success click rates before and after the transformation.

   

4. Randomly select 20% of the remaining articles to be in the validation set; the rest are in the training set. The training and validation datasets contain 9122 and 2280 articles respectively.

10.2.0.2 Additional training details.

  During training, we train to sequence lengths of \(500\). Hyperparameters and early stopping epochs are chosen using the validation dataset.

• Flexible NN (text). This model is very similar to the Flexible NN model in Appendix 10.1, with the exception that in place of a two-dimensional \(Z^{(a)}\), the MLP head of the neural network from before is fed as input a DistilBERT [93] embedding of text data \(Z^{(a)}\). Also, the MLP linear layers have width 100 and the summary statistics are repeated 100 times.

• Beta-Bernoulli NN (text). This is very similar to the Beta-Bernoulli posterior predictive sequence model in Appendix 10.1, with the exception that in place of a two-dimensional \(Z^{(a)}\), the MLP head of the neural network from before is fed as input a DistilBERT [93] embedding of text data \(Z^{(a)}\). On top of the one DistilBERT embedding are two separate MLP heads for \(\alpha(Z^{(a)})\) and \(\beta(Z^{(a)})\), which are trained together. Also, the summary statistics are repeated 100 times and MLP layers have width 50.

• Flexible NN (category). The model and training procedure are the same as for Flexible NN from Appendix 10.1, except it uses a one-hot new category vector for \(Z^{(a)}\).

The models Flexible NN (text) and Beta-Bernoulli NN (text), which both incorporate the DistilBERT model, are each trained for 500 epochs with a learning rate of 1e-5 on MLP weights and 1e-8 on the DistilBERT weights using the AdamW optimizer.

10.2.0.3 Posterior width and calibration plots at additional timesteps.

  In Figure 18, we display posterior calibration and posterior widths as in Figure 12 (right), but for additional timesteps beyond \(t=0\).

10.2.0.4 Ensembles (Figure 12).

  We describe the ensembling approach used in the uncertainty quantification plots in Figure 12 (right). Our ensembles take a base model that is a DistilBERT that is trained to take in article headlines \(Z^{(a)}\), which is fed into an MLP with a sigmoid head that predicts binary outcomes.⁷ To construct our ensembles computationally efficiently, we just ensemble the weights of the MLP head. We freeze the DistilBERT weights and train 50 independent MLP heads from scratch with random initialization and bootstrapped training data to create the ensemble.

10.2.0.5 Additional ensemble results (randomized prior).

  In Figure 19, we include a “randomized prior” variant of ensembling that initializes each neural network model in the ensemble in a particular way to encourage diversity [74]. Specifically, we modify the ensembling procedure described earlier so that each of the 50 MLP heads incorporates their own prior networks, which we denote using \(g_{1},\ldots,g_{50}\), which are MLP heads with the same architecture with randomly initialized weights. Note that the prior network weights are not trained. Let \(f_{\theta_1},\ldots,f_{\theta_{50}}\) denote the trainable MLP heads. The model predicts Bernoulli success probabilities as follows \[\begin{align} \text{sigmoid}(f_{\theta_i}(\varphi(Z^{(a)}))+\sigma \cdot g_{i}(\varphi(Z^{(a)}))). \end{align}\] Above, we use \(\varphi(Z^{(a)})\) to denote the trained DistilBERT embedding of the text headline \(Z^{(a)}\). Also, \(\sigma\) is a prior scaling term hyperparameter.

Although tuning the prior scaling for the randomized prior can achieve approximately calibrated uncertainty (in contrast to regular ensembling), posterior sampling with autoregressive generation (PSAR) is advantageous over ensembling with randomized priors because i) PSAR simply updates its posterior beliefs by conditioning on the history, while ensembling must update models by fine-tuning the models via gradient updates on the new data, and ii) the pretraining procedure for PSAR automatically calibrates uncertainty well through standard sequence loss modeling.

10.3 Baseline Bandit Algorithms↩︎

10.3.0.1 TS Beta Bernoulli (Uniform Prior).

  This algorithm uses Thompson sampling with a Beta-Bernoulli model with a uniform prior [16]. Note that unlike TSAR using the Beta-Bernoulli NN, the prior here does not depend on action attributes \(Z^{(a)}\).

10.3.0.2 TS Neural Linear.

  We implement a simplified version of “neural linear” [71], [72]. Here, we assume that for each arm, data are generated independently as \[Y^{(a)}_t=W^{(a)\top} \beta^{(a)} +\epsilon^{(a)}_t\] where \(\epsilon^{(a)}_t\sim N(0,\sigma^2)\) are drawn iid across arms \(a\) and timesteps \(t\), and \(W^{(a)}:=[g(Z^{(a)}),1]^\top\) where \(g(\cdot)\) is a model trained to predict reward from features \(Z^{(a)}\) by using historical data (details below). Additionally, we assume \[(\sigma^{(a)})^2\sim \text{InverseGamma}(a,b)\quad\text{ and }\beta^{(a)} \mid (\sigma^{(a)})^2\sim N([1,0]^\top, (\sigma^{(a)})^2 \mathbf{I}_2).\] We set \(a=3\) and \(b=0.5\) since outcomes for each arm are Bernoulli and thus can have variance at most 1/4, and an Inverse Gamma with parameters \(a=3\) and \(b=0.5\) have reasonable mass under 1.4 while also being somewhat diffuse so that it does not make too strong an assumption on \(\sigma^{(a)}\).

For our synthetic experiments, we train \(g\) using the same architecture and training procedure as we do for Flexible NN from Appendix 10.1 except it does not take in a sufficient statistic of the history. Similarly, for our news recommendations experiment, we train \(g\) using the same architecture and training procedure as we do for Flexible NN from Appendix 10.2 (which incorporates a DistilBERT model) except it does not take in a sufficient statistic of the history. (We also change the learning rates slightly as we found it led to lower validation set prediction loss: 1e-2 for synthetic setting, and 1e-5 for the news setting.)

10.3.0.3 UCB.

  For UCB we use the multi-arm bandit algorithm described in Section 6 of [94]. We set the failure probability \(\delta = 0.1\) and use sub-Gaussian parameter \(0.5\) (since we have binary rewards).

10.4 Modeling best action rather than modeling action reward↩︎

[68] introduce Decision Pretrained Transformers as an alternative approach of using sequence models to implement Thompson sampling. Their procedure trains a sequence model to predict the best action, mimicking an expert/optimal policy, given current state and recent interactions. This section aims to make a fair empirical comparison between our TSAR algorithm and the approach introduced in [68].

10.4.0.1 Experiment details.

  In order to make the most fair comparisons between modeling best actions (which we call TS Action) and modeling rewards (TSAR), we use nearly identical model architecture and training procedures. The only adjustments we make are i) while the TSAR sequence models output a distribution over the next outcome \(Y\), the TS Action sequence models output a softmax distribution over the set of action candidates, ii) the training target is set to the best action, and iii) we must sample offline trajectories collected by a behavior policy to use for training (“in-context datasets”). For collecting these offline trajectories we follow a procedure identical to that used in [68] for the multi-armed bandit setting.⁸ We train these action-based sequence models for 2500 epochs and select the best learning rate out of \(\{10^{-1}, 10^{-2},10^{-3},10^{-4}\}\).

10.4.0.2 Results.

  When training with \(|\MC{A}^{\TN{offline}}| = 2500\) actions in the offline pretraining dataset, there is little difference between TSAR and TS-Action. In contrast, TSAR is much better than TS Action when the number of offline training arms is drastically reduced, e.g. to 25 (Figure 20 (c)). We hypothesize that this gap in performance is because the best action provides less information than the rewards for each action. For example, one action may have better rewards on average than another, but while TSAR can comprehend how much better the rewards are for action vs the other, TS Action only knows relative ordering. Further understanding the gap in performance is an interesting direction for future work.

11 Bayesian sequence models↩︎

This section connects proposed offline learning procedure (for example, as in Section 7.1) to empirical Bayes methods (Type-II maximum likelihood), which is an approach used to fit a Bayesian model to observed data [92], [101], [102]. This connection is precise in cases where \(p_{\theta}\) is the posterior predictive distribution of some Bayesian model and \(\theta\) consists of hyperparameters of the prior. In other settings, where we fit parameters of a more flexible (and not necessarily exchangeable) sequence model is used, this connection may help readers view our offline sequence model learning procedure as “learning to form prior beliefs given \(Z\), using historical data”.

11.1 Offline learning as empirical Bayes↩︎

Standard Bayesian mixture models define exchangeable sequence models. This is formalized in Example 4 below.

To keep the presentation simple, we illustrate empirical Bayes in a simple case. Consider a Beta-Bernoulli Bayesian model without \(Z\)’s, which was introduced in 12 : \[\begin{align} \label{eqn:betaBernoulli} \mu^{(a)} \sim \TN{Beta}(\alpha^*, \beta^*) ~~\TN{then}~~ Y_1^{(a)}, Y_2^{(a)}, \dots \big| \, \mu^{(a)} \iidsim \TN{Bernoulli}(\mu^{(a)}). \end{align}\tag{21}\] It is well known that the above Bayesian model 21 is associated with the following class of autoregressive sequence models, i.e., the posterior predictive distribution: \[\begin{align} \label{eqn:betaBernoulliPostpred2} 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 \TN{where} \quad \theta = (\alpha, \beta). \end{align}\tag{22}\] Here, the true sequence model \(p^*\) is exactly \(p_\theta\) evaluated at \(\theta = (\alpha^*, \beta^*)\). For the above sequence model class \(p_\theta\), minimizing our negative log likelihood sequence training criterion 11 is equivalent to maximizing the marginal likelihood of the data, the criterion that is used in Empirical Bayes (Type-II maximum likelihood) to fit prior distributions to observed data [92], [101], [102].

11.2 Recovering the true prior via offline learning (Empirical Bayes)↩︎

Here, we numerically demonstrate a setting for which we perform “empirical Bayes” using our offline learning procedure (Section 7.1). We find that we recover the true prior fairly well.

11.2.0.1 Data generation.

  We use a synthetic Beta-Binomial data generating process. We consider one-dimensional action features \(Z^{(a)}\iidsim \TN{Uniform}(0, 1)\). We then sample \(\mu^{(a)}\) from a Beta distribution: \[\begin{align} \label{eqn:empiricalBayesPrior} \mu^{(a)} \mid Z^{(a)} \sim \TN{Beta} \big( Z^{(a)} \cdot 5 + 1, ~ (1-Z^{(a)}) \cdot 5 + 1 \big). \end{align}\tag{23}\] Then, outcomes are drawn as \(Y^{(a)}\mid \mu^{(a)}\sim \TN{Bernoulli}(\mu^{(a)})\). We use \(R(y) \triangleq y\). We use a training dataset of size 25,000 actions and a validation set of size 10,000 actions. We use sequences of outcomes of length \(500\).

11.2.0.2 Autoregressive model.

  We use an architecture for \(p_\theta\) that matches the posterior predictive of a Beta-Bernoulli model described above. Specifically, we use the following sequence model: \[\begin{align} \label{eqn:betaBernoulliPostpred} p_{\theta} \big( Y_{t+1}^{(a)} = 1 \mid Y_{1:t}^{(a)} \big) = \frac{ \alpha_\theta(Z^{(a)}) + \sum_{i=1}^{t} Y_{i}^{(a)} }{ \alpha_\theta(Z^{(a)}) + \beta_\theta(Z^{(a)}) + t }. \end{align}\tag{24}\] We parameterize the prior hyperparameters \(\alpha_\theta(Z^{(a)}), \beta_\theta(Z^{(a)})\) using neural networks. We use the exact procedure from Appendix 10.1 for Beta-Bernoulli NN (except our MLPs have width 100).

11.2.0.3 Recovering the Prior (Figure 21).

  Through our pretraining procedure, we approximately recover the true prior. We show this by comparing means and standard deviations of samples from our learned prior (using \(p_\theta\)) vs. the true prior (according to the data generating process), for different draws of \(Z^{(a)}\). In the scatter plots in Figure 21, each point corresponds to one \(Z^{(a)}\).

Specifically, in these plots we use \(100\) actions sampled uniformly from the validation set. For each of these \(100\) actions we form \(10,000\) samples of \(\mu^{(a)}(\widehat{\tau}_1)\) using Algorithm 6 using our learned \(p_\theta\) model. We also form \(10,000\) samples from the true data generating prior 23 for each of the \(100\) actions. Then for each action, we compute the mean and standard deviations of the samples on the “prior” samples \(\mu^{(a)}(\widehat{\tau}_1)\) from \(p_\theta\); we also compute the mean and standard deviations of the samples from the true prior. We then plot these in a scatter plot; for each action, we have the prior mean according to \(p_\theta\) vs the prior mean according to the data generating process—this forms one point on the scatter plot. A similar procedure is plotted on the right. There, instead of computing the mean of the prior samples, we compute a measure of the spread of the prior samples: let \(\mu^{(a)}(\widehat{\tau}_{1,1}), \mu^{(a)}(\widehat{\tau}_{1,2}) \ldots, \mu^{(a)}(\widehat{\tau}_{1,10000})\) be the prior samples. Let \(\bar{\mu}^{(a)} = \frac{1}{10000}\sum_{i=1}^{10000} \mu^{(a)}(\widehat{\tau}_{1,i})\). Then we compute the mean absolute deviation \(\frac{1}{10000}\sum_{i=1}^{10000} \big| \mu^{(a)}(\widehat{\tau}_{1,i}) - \bar{\mu}^{(a)} \big|\) for this set of prior samples.

12 Finite vs infinite population formulations and Thompson Sampling variants↩︎

This section discusses the intimate connections between (large) finite-population formulations that were discussed in the main body of the paper and infinite-population formulations that are more common in the Bayesian bandit literature. We do this in the special case of the Bayesian mixture model from Example 4.

We emphasize that from our perspective, the main advantages or disadvantages of the finite population view are conceptual. In terms of advantages: (1) the definitions do not require any explicit assumptions around mixture modeling or latent variables, and (2) the finite nature of the problem lets us visualize the procedure as in Figure 7, without abstract reference to limits across infinite sequences.

12.1 Review of Thompson sampling in infinite populations↩︎

Thompson sampling is most often defined for a Bayesian mixture model, e.g., as in Example 4. Following that example, we consider in the subsection the canonical example of exchangeable sequences: a mixture model wherein the outcomes are i.i.d. conditioned on a latent variable \(\eta^{(a)}\). That is, \(p^*(Y_1^{(a)}, \ldots, Y_t^{(a)} \mid Z^{(a)}) = \intop \prod_{t=1}^{T} P(Y_t^{(a)} \mid \eta^{(a)}=u) P(\eta^{(a)} = u\mid Z^{(a)}) du\). The unknown latent variable represents the decision-maker’s uncertainty about an action’s performance.

The literature typically defines the “true arm means” given \(\eta^{(a)}\) as \[\mu_{\infty}^{(a)} = \intop R(y) \cdot P(y \mid \eta^{(a)})\,dy.\] The subscript highlights that this has the interpretation of a long-run average reward across an infinite population of users (or infinite set of rounds). By the law of large numbers (applied conditional on \(\eta^{(a)}\), one has \[\mu_{\infty}^{(a)} = \lim_{T\to \infty} \frac{1}{T} \sum_{t=1}^{T} R( Y_t^{(a)} ).\] The true best arm is defined as \[A^*_\infty \in \argmax_{a\in \MC{A}} ~ \mu_{\infty}^{(a)}\] Randomness in the latent parameters \((\eta^{(a)})\) means \(\mu_{\infty}^{(a)}\) and \(A^*_\infty\) are random variables whose realizations are uncertain even given the history \(\HH_{t-1}.\) Thompson sampling selects an action by probability matching on \(A^*_{\infty}\), defined by the property \[\label{eq:infinite95pop95TS} \PP(A_t=a \mid \HH_{t-1}) = \PP(A^*_{\infty} = a \mid \HH_{t-1}) \quad \text{for all } a \in \MC{A}.\tag{25}\] Per-period Bayesian regret over \(T\) periods is defined as \[\label{eq:infinite95pop95regret} \E\left[ \frac{1}{T}\sum_{t=1}^{T}\left( R \big( Y^{(A_{\infty}^*)}_t \big) - R \big( Y^{(A_t)}_t \big) \right) \right]\tag{26}\]

12.2 Thompson sampling in finite populations↩︎

One can define the true mean of a finite population as \[\mu^{(a)}_T = \frac{1}{T} \sum_{t=1}^{T} R \big( Y_t^{(a)} \big).\] Note in the main paper we use the notation \(\mu^{(a)}_T\) to denote the above mean, but we change the notation here to make the contrast to \(\mu_\infty^{(a)}\) more clear. The true best arm for this finite population is defined as \[A^* \in \argmax_{a\in \MC{A}} ~ \mu^{(a)}_T.\] As in Lemma 3, Thompson sampling selects an action by probability matching on the (finite-population) optimal action \(A^*\), defined by the property \[\label{eq:finite95pop95TS} \PP(A_t=a \mid \HH_{t-1}) = \PP(A^* = a \mid \HH_{t-1}) \quad \text{for all } a \in \MC{A}.\tag{27}\] Per-period Bayesian regret over \(T\) periods is defined as \[\label{eq:finite95pop95regret} \E\left[ \frac{1}{T}\sum_{t=1}^{T}\left( R \big(Y^{(A^*)}_t \big) - R \big( Y^{(A_t)}_t \big) \right) \right]\tag{28}\] It is not hard to show that 28 is a more stringent notion of regret than in 26 , since
\(\frac{1}{T}\sum_{t=1}^{T} R \big( Y^{(A^*)}_t \big) \geq \frac{1}{T}\sum_{t=1}^{T} R \big( Y^{(A^*_\infty)}_t \big)\) by definition of \(A^*\). Both definitions are widely used, with the more stringent finite-population version being more common in the adversarial bandit literature; see [31].

12.3 The gap between finite and infinite population formulations is small↩︎

We analyze the gap between the two formulations in the case of a mixture model. By a sub-Gaussian maximal inequality \[\begin{align} \E\left[ \max_{a\in \MC{A}} \left| \mu_{\infty}^{(a)} - \mu_T^{(a)} \right| \right] = \E\left[ \E\left[ \max_{a\in \MC{A}} \left| \mu_{\infty}^{(a)} - \mu_T^{(a)} \right| \bigg| \, \{ \eta^{(a)} \}_{a\in \MC{A}} \right]\right] \leq \sqrt{ \frac{2 \log (| \MC{A}|)}{T} }. \end{align}\] To justify the last inequality, note that since the function \(R\) takes values in \([0, 1]\), \(R \big( Y_t^{(a)} \big) - \mu^{(a)}_{\infty}\) is sub-Gaussian with variance proxy \(1\), conditional on \(\eta^{(a)}\) (by Hoeffding’s Lemma). Since it is the average of independent sub-Gaussian random variables, \(\mu_{\infty}^{(a)} - \mu_T^{(a)}\) is sub-Gaussian with variance proxy \(\frac{1}{T}\), conditional on \(\{ \eta^{(a)} \}_{a\in \MC{A}}\). The last step follows then from applying the subgaussian maximal inequality, conditional on \(\{ Z^{(a)} \}_{a \in \MC{A}}, \{ \eta^{(a)} \}_{a\in \MC{A}}\).

It follows easily that the infinite population optimum \(A^*_{\infty}\) is near optimal for finite populations: \[0 \leq \E\left[ \max_{a\in \MC{A}} \mu_{T}^{(a)} - \mu^{(A^*_\infty)}_T \right] \leq 2\sqrt{ \frac{2 \log (| \MC{A}|)}{T} }.\] Analogously, the finite population optimum is near-optimal in infinite populations: \[0 \leq \E\left[ \max_{a\in \MC{A}} ~ \mu_{\infty}^{(a)} - \mu^{(A^*)}_\infty \right] \leq 2\sqrt{ \frac{2 \log (| \MC{A}|)}{T} }.\] Supported by this theory, we do not focus on the distinction between \(A^*\) and \(A^*_\infty\).

12.4 Similar insights in empirical results↩︎

Some empirical insight can also be gleaned from Figure 13 (a) in Appendix 7.4. The implementation that performs full imputation can be interpreted as Thompson sampling for a finite population. The implementation that performs forward generation of fixed-length \(m\) does not include past observed rewards in the average. For very large \(m\), it is a direct approximation to infinite-horizon Thompson sampling. See Figure 13 (a) for empirical evidence that these implementations have very similar performance.

References↩︎

---

1. We benchmark against the best action in hindsight, which is very common in adversarial formulations. Less stringent versions of regret that move the maximum outside a certain conditional expectation are sometimes called pseudo-regret [30], [31]. In versions of our formulation where both are well-defined, the distinction between these are insignificant; see Appendix 12.↩︎

2. Somewhat remarkably, Thompson sampling satisfies a regret bound (relative to the best action in hindsight \(A^*\)) even in nonstationary settings where the exchangeability assumption fails (see Proposition 1 and the references therein), but other standard bandit algorithms are only studied in stationary settings.↩︎

3. For example, [41] show that cumulative regret in the bad-prior regime scales as \(O(\sqrt{T/\rho})\), where \(\rho\) is the prior mass assigned to the true model. Thus, if \(\rho=1/T^3\), this dependence yields a \(T^2\) cumulative regret penalty.↩︎

4. In domains like recommender systems, there is substantial “noise" in individual user interactions. The aforementioned papers apply their method to benchmark problems with almost no reward noise.↩︎

5. We found that repeating the summary statistic before being fed into the sequence model improved performance. Thus, the model took as input vectors in \(\real^{22}\) which consisted of a \(2\)-dimensional \(Z^{(a)}\) and \(10\) copies of the summary statistic tuple.↩︎

6. This dataset is free to download for research purposes at https://msnews.github.io/.↩︎

7. MLP width 100, 3 layers, learning rate 1e-5 on the head and 1e-8 on DistilBERT, weight decay 0.01, AdamW optimizer.↩︎

8. For each new bandit task, we load a random batch of arms and their corresponding sequences of rewards. This batch of size 5000 is split into multiple bandit environments, each with 10 arms. Actions are sampled according to a mixture of a Dirichlet and a point mass (see Section 4 of [68] for more details).↩︎