Diffusion Models Meet Contextual Bandits
Authors: Imad Aouali
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion models
Bayesian contextual bandits
Primary area: reinforcement_learning

Abstract: Efficient online decision-making in contextual bandits is challenging, as methods without informative priors often suffer from computational or statistical inefficiencies. In this work, we leverage pre-trained diffusion models as expressive priors to capture complex action dependencies and develop a practical algorithm that efficiently approximates posteriors under such priors, enabling both fast updates and sampling. Empirical results demonstrate the effectiveness and versatility of our approach across diverse contextual bandit settings.

Diffusion Models Meet Contextual Bandits
Imad Aouali
Criteo AT Lab CREST, ENSAE, IP Paris deouali@rriin com
Abstract
Efficient online decision-makii in contextual bandits is challenging, as methods without informative priors often suffer from computational or statistical inefficien- cies. In this work, leverage pre-trained diffusion models as expressive priors to capture complex action dependencies and develop a practical algorithm that efficiently approximates posteriors under such priors, enabling both fast updates and sampling. Empirical results demonstrate the effectiveness and versatility of our approach across diverse contextual bandit settings.
1 Introduction
contextual bandit models online decision-making under uncertainty [49]. At each round, an agent observes context, selects action, and receives reward, aiming maximize cumulative reward by balancing exploitation high-reward actions and exploration of uncertain ones. However, in large- scale settings (e.g, the number of actions K is large), standard exploration strategies (e.g., LinUCB [11] or LinTS [60]) often become computationally expensive or statistically inefficient. Fortunately, actions in many real-world problems exhibit correlations, enabling more efficient exploration since observing one action can inform the agent about others. Thompson sampling is particularly well- suited for this, as naturally incorporates informative priors [33] that capture complex action dependencies. Inspired by the success of diffusion models [62, 30], which excel at approximating complex digensinnnii distributions [23, 58], this work leverages pre-trained diffusion models as priors in contextual Thompson sampling.
Precisely, introduce framework for contextual bandits with diffusion-derivel prior, and develop diffusion Thompson sampling (dTS) that is both computationally and statistically efficient. dTS achieves fast posterior updates and sampling through efficient approximation that becomes exact when the diffusion prior and the likelihood are linear. A key contribution, beyond applying pre-trained diffusion models in contextual bandits, the efficient computation and sampling of the posterior distribution of d-dimensional parameter 0 DD with D representing the data, when using a pre-trained diffusion model prior on This is relevant not only to bandits and RL but also broader range of applications [19]. Our approximations are motivated by exact closed-form solutions obtained in cases where both the pre-trained diffusion model and the likelihood are linear. These solutions form the basis for our approximations for the non-linear case, demonstrating both strong empirical performance and computational efficiency. Our approach avoids the computational burden of heavy approximate sampling algorithms required for each latent parameter.
Diffusion models have been applied to offline decision-making [6, 36, 65], but their in online learning has only recently been explored by Hsich et al. [34] who studied the multi-armed bandit setting, and Kveton al. [44] who explored similar direction, with their preprint appearing shortly after the first version of this work. An carlier, deliberately simplified version of our approach, restricted to a linear diffusion model prior, was introduced in Aouali [7]. Here, we extend that formulation to the more realistic and expressive non-linear case. A detailed discussion of related work is provided in Appendix A.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

2 Setting
The agent interacts with contextual bandit over rounds. In round t [ni], the agent observes context X, E X, where X' C IR is context space, takes an action A, E [K], and then receives stochastic reward Y, E IR that depends on both the context X₁ and the taken action Ac.
We focus on the per-action (disjoint) setting, where each action a E [K] parameter vector 0₂ E R so that the reward received in round t is Y, ~ P(- 0 where P( x;; 0a) is the reward distribution of action in context The reward distribution is parametrized as a generalized linear model (GLM) [54]. That is, P(. x;; (a) is an exponential-famin distribution with mean g(x² T 0n), where g is the mean function. For example, we recover linear bandits when P(1 2; (2)) -V(: 0a. o²) where > O, and logistic bandits [24] with g(u) + (xpp(-))) and P(- æ; 0a) = Ber(g( (.)). where Ber(p) denotes the Bernoulli distribution with mean p. All derivations and algorithms extend naturally to the parrmee- case described in Remark 2.1.
] is represented by an unknown
We consider the Bayesian bandit setting [59, 33, 55, 28], where the true action parameters O. are assumed to be drawn from a known prior distribution. As both the true parameters and the model parameters are sampled from this same prior, notation. We proceed to define this prior distribution using a diffusion model. The correlations between the action parameters O,, captured through diffusion model, where they share set of I. consecutive unknown latens parameters the E Rd for l E [7.]. Precisely, the action parameter 0₂ depends the I--th latent parameter 1/1. as
them interchangeably
a slight abuse of
1/1 0a 1/1 ~ ((f(((1)))
where the link function fi and covariance Σ1 are known. In particular, the action parameters 9g are conditionally independent given 1/1 Also, the 1-th latent parameter the-1 depends on the f-th latent parameter u/e as
wz-1 the N (ffe(t)),
where fe and Eg are known. Finally, the L-th latent parameter d!z is sampled as th ~ (O, EL+1), where 11.11 is known. We summarize this model in Eq. (1) below
th ~ ((0,Er+)),
(1)
V1-1 1/2 N(Je((be) Ex),
V E [1]/(1)
0g th ~ ((J(((i), En), Y, 12, 0A, ~ P(. Xx;
Va E [K]
14 [n]
Eq. (1) represents a Bayesian bandit, where the agent interacts with a bandit instance defined by 0 over n rounds (4th line in Eq. (1)). These action parameters 0g are drawn from the generative process in the first three lines of Eq. (1). practice, Eq. (1) can be built by pre-training a diffusion model on offline estimates of the action parameters 0..
The goal of the agent is minimize its Bayes regret [59], which measures the expected performance across multiple bandit instances 0 that are sampled from the prior,
BR(n) = E Σ
where the expectation is taken over all random variables in Eq. (1). Here r(æ, a; By [Y] is the expected reward of action in context x, and Ar.- = arg max r(X,, a; 0) is the optimal action in round t. The Bayes regret captures the benefits of using informative priors [33, 32, 8), and hence it is suitable for our problem.
Remark 2.1 (Single shared action parameter). Our algorithm and analysis also apply to the case where all actions share a single unknown parameter E Rdd. Let yp : X [K] - Rdd be a known feature map, and assume the reward distribution mean is g g a) 0). Then, the diffusion prior in Eq. (1) specializes by replacing the per-action parameters (0a)aE(K) with single shared parameter 0:
VL NO, Ez+1),
(2)
d the ~ N(fe(1b)) Ee), 0 01 (N1 E1), Y | 4,,0 ~ P(- (X,,) TO),
AL [L) f1).
(11).
Vt

---

This formulation is useful when a shared feature map is available. In that case, the diffusion model can be pre-trained on parameters {0,}--1 from previous tasks, and dTS can then be applied to new task S+1 using the pre-trained prior. To avoid clutter, our main exposition focuses on the model in Eq. (1), but all theoretical results and algorithmic components extend naturally this Shared-paramete case, which we also include in some experiments (explicitly noted when applicable).
Diffusion contextual Thompson sampling
3.1 Algorithm
We design Thompson sampling algorithm that samples the latent and action parameters hierar- chically [51]. Let H, = (X;, A;, Y-) denote the history of all interactions up to round 4, and let Ht, = (X,, A,, =)) be the history of interactions with action a up to round t. To motivate our algorithm, we decompose the posterior density p(O,. He) recursively as
p(O, Ho)
pLUL 11) II P(ve tee H.)p(O, (r,a) dii.L J
Hierarchical sampling. This decomposition induces the following sampling procedure in round I. First, draw sample t according the posterior density p(ıy Hz). Then, for each & [L] {1}, draw 1- from the conditional posterior p(ve-1 tr,,, Hz). Finally, given 1hr,1, draw each action parameter independently from p(0, vhr,,, He,a) (the 0, are conditionally independent given $1))
This defines Algorithm 1, diffusion Thompson Sampling (dTs).
Posterior components via recursion. To implement dTS, we provide an efficient recursive scheme express the required posteriors using known quantities. These expressions may always admit closed forms and can require approximation. The conditional action action-posterior can be written
p(0a | 1/ Hr,a)
II P(Y X,;0) f1(1)), Ex), vESt..
(4)
where St,, (e E - 1]: A₂ a} is the of rounds in which action was selected. Moreover, let p( H; 4/2) denote the likelihood of the observations up to round t given the. For any & E [Z]] X (1). the conditional posteroost is
p( ( 1/, H.) x p(H, ((1- Ex).
and the top-layer posterior is
p H.) x p(H.) Wr.) 0, Er.1).
All terms above are known except the likelihoods p(H. tir), which are computed recursively. The recursion starts with
K II a=1
p(H,
II P(Y X1i0a) (((1),) de.. LieS.,a
(5)
and for é [4]] (1), proceeds as
p(Hel whe) = / p(H, ((--)) Ex) dib-11
(6)
All posterior expressions above use known quantities (fé. E,, P x;; B)). However, these expres- sions typically need to be approximated, except when the link functions fe are linear and the reward distribution P(- xn; 0) is Sheea-Gaussian where closed-form solutions can be obtained with careful derivations. These approximations are not trivial, and prior studies often rely computationally intensive approximate sampling algorithms. the following sections, we explain how we derive our efficient approximations which are motivated by the closed-fone solutions of linear instances.
3

---

Algorithm dTS: diffusion Thompson Sampling
Input: Prior components {fe, Ee
L+t1 and reward model P.
for f I...., do
Draw L according to the posterior density p(tL H³)
for L L..... 2 do
L Draw 1r,c-1 according to p( < 1 eereer Hz) do L Draw or, according to p(0, 11,1, He,a)
for a aara=,.....
Select action A, arg m[ r(Xx e; a; 0z), where 0. = (Or,a)ae(k]
Observe reward Y ~ P(- Xx: AA,) and update the posteriors.
3.2 Posterior approximation
The reward distribution is parameterized as generalized linear model (GLM) [54], which allows for non-linear rewards. In addition, the diffusion model itself highly non-linear due to the link functions fe. These two sources of non-linearity make the posterior intractable, so apply two layers of approximation: (i) a likelihood approximation to linearize the reward model, and (ii) a diffusion approximation to handle the non-linear hierarchy induced by the diffusion model prior.
(i) Likelihood approximation. We use an approach similar the Laplace approximation, but instead of approximating the entire posterior, approximate only the likelihood by a Gaussian. Precisely, the reward distribution P(.
xz; (a) belongs to the exponential family with mean function g. Thus
II iES...
P(Y, 0a) = N(Oa: Br,as GT.a))
where Br,a is the maximum likelihood estimate and Gr,a is the Hessian of the negative log-likeli lieeiike
Bra arg max
Σ S
log P(Y, Xx; Gr.0 = - (
(8)
0RER
IESt.a
and St.a = (e A, =, a} is the set of rounds in which action was selected. Unlike Laplace, which fits a global Gaussian to the full posterior, this step only linearizes the local likelihood, allowing the hierarchical diffusion structure of the prior to remain intact and expressive.
(ii) Diffusion approximation. Plugging the Gaussian likelihood approximation (7) into the posterior expressions p(O, | 11, He,a) and p( - the Hz) removes the non-linearity of the reward model. However, the diffusion hierarchy remains non-linear through fe. To handle this, we build on the closed-form posteriors of the linear diffusion case (where fe(we = Wayfe see Appendix B) and generalize them by replacing the linear terms W Wee with their non-linear counterparts fe($)). This substitution yields a posterior diffusion model that retains the same hierarchical form as the prior but with data-dependent means and covariances. Details on how we transition from the linear the general non-linear setting are provided in Appendices B and C. The resulting approximate posteriors (both action and latent) admit the following closed-form expressions.
Approximate action posterior. We approximate the conditional action posterior as p(0a v1, Hr,)) ss M(Oaifrr, Er,))
where
$b-1
G data precision
j4, Et, prior contribution
Σ (1
Graa CraBBrg Brag )
prior precision
data contribution
This posterior update has clear interpretation. The posterior precision En is the sum of the prior precision and the data precision. The posterior mean Ωr,a is the precision- average of the prior mean and the MLE Bra. As more data observed, the covariance shrinks and the mean moves from the prior mean fi (1(1) toward the MLEE Br,a- When no data are available (Gt,a = 0), the posterior reduces to the prior N(Ji (, E,): in the limit of infinite data (G.,a -0), the posterior collapses to the MLE Br with fr,a + Br,a and Er +

---

Approximate latent posteriors. For each & E [L + 1] (1}, we approximate the latent posterior as
p(ee-1 We, H1) = -e
with
Σ-1 26-1
+
Gr.c- Pr,e-1 = 11,0-1( Ee fe(ee)
Bi-e-1 )
(10)
prior precision
data precision
prior contribution
data contribation
where, by convention, fI-1 = since the top layer 1 has no parent 11+11 The quantities G and B computed recursively. The base recursion
K G.,1
K
- - "T,
-
B₁ = E
Σ
Cra
(1I)
a=1
a=1
and for each l e [L]] {1)}
E -, -1,6-122 E-
B,, = = E B,,
(12)
The latent posterior update in Eq. (10) has the same structure as the action posterior. The posterior precision 1 is the sum of the prior and data precisions and the posterior mean is their precision- weighted combination. The data terms G and B are computed recursively (Eqs. (11)
and (12)), so information collected at the action level propagates upward through the hierarchy.
Interpretation. The resulting approximate posterior remains diffusion model whose conditional Gaussians have updated, data-dependent means and covariances. The latent-posterior means can be viewed refined link functions:
fee((b))) = fu, = Er
Se(('e) + (E( Br.--1)
and Σr,, represents their updated uncertainty. Both are updated with data: covariances contract as uncertainty decreases, and means move from the prior toward the MLE. Unlike a full Laplace approximation, this formulation preserves the expressiveness of the posterior rather than replacing it globally with single Gaussian, while also avoiding the heavy computation required by other approximate inference methods.
3.3 Extension single shared action parameter
For the shared-parameter model in Remark 2.1, dTS's posterior approximations are similar. The action posterior is p(0 1/, Hy) ((z,, Ex,, where
--1= = ET + G.
f4 = (i + GGBB.)
(13)
where
B, arg max BER4
log P(Y | A.))) G. = Σ Ar) B:) ((X,, (X,, T.)
i<t
ist
Similarly, for l [L + 1]} {1}, the latent posterior is p(th-1 v/, Hi) v ( Et,e-1), where
-= = Ex + GG,--1 P = E fo( + B,.--i..
(14)
where, by convention, f = and the quantitics Gi,e and B are computed recursively as
Base case:
-. E. -
B Σ7 1
(15)
Recursive case:
E
(16)
Again, this shared-parameter variant of dTS is presented for completeness and to illustrate generality of our posterior derivations; the main focus of the paper remains on the per-action formulation Eq. (1). Unless stated otherwise, all theoretical results and experiments use the main version of dTS described in Algorithm I.
5

---

4 Informal theoretical insights
In this section, we present an informal Bayes regret analysis of dTS to build intuition around its statistical efficiency. We assume simplified the reward distribution is Gauss-Gn and each link function fe((r) Weip is a known linear mixing matrix. These assumptions induce a hierarchy of L linear Gaussian layers from the latent root to the action parameters. In this case, our posterior approximation becomes exact which enables an analysis reminiscent of linear contextual bandits [3]. However, our recursive hierarchical structure introduces technical differences: the posteriors must be derived inductively using total covariance decompositions, and regret bounds require tracking information flow across all latent layers. We emphasize that this regret bound does not hold in the general nonlinear case studied experiments and on which we focus in this paper, and is only included here to provide theoretical intuition under simplifying assumptions. Formal statements and derivations are provided in Appendices E and F.
setting to make the analysis tractable:
Informal Bayes regret bound. The bound of dTS in this case is
v + dE 22 o
where AAx max 1 _.. This dependence on the horizon n aligns with prior Bayes regret bounds scaling with n. However, the bound comprises L + 1 main terms. First, one relates action parameters learning, conforming to a standard form [52], while the remaining terms are associated with learning each of the latent parameters.
Sparsity refinement. If each mixing matrix exhibits column sparsity, that, We = -We, Od,d-d;) with dy < d active columns, then the bound becomes
BR(n) = Ō mndKof + Eie d
Hence, informative, spurse priors can cut the cost of learning deep latent chains down from d to de- This Bayes regret bound has a clear interpretation: if the true environment parameters are drawn from the prior, then the expected regret of an algorithm stays below that bound. Consequently, less informative prior (such as high variance) leads to a more challenging problem and thus a higher bound. Then, smaller values of K, L, d, de translate to fewer parameters to learn, Icading to lower regret. The regret also decreases when the initial variances of decrease. These dependencies are common in Bayesian analysis, and empirical results match them.
The reader might question the dependence of our bound on both L and K. Details can be found in Appendix D.2, but in summary, we model the relationship between 9, and 1 stochastically as N( 1h,, of Ia) to account for potential nonlinearity. This choice makes the model robust to model misspecification but introduces extra uncertainty and requires learning both the O2 and the thy. This results in a regret bound that depends on both K and However, thanks to the use of informative priors, bound has significantly smaller constants compared to both the Bayesian regret for LinTS and its frequentist counterpart, as demonstrated empirically in Appendix G.5 where it is much tighter than both and in Section 4.1 where we theoretically compare our Bayes regret bound to that of LinTs.
Technical contributions. dTS uses hierarchical sampling. Thus the marginal posterior distribution of 0, H; is not explicitly defined. The first contribution is deriving 0, H, using the total covariance decomposition combined with an induction proof, as our posteriors were derived recursively. Unlike standard analyses where the posterior distribution of 0, H, is predetermined due to the absence of latent parameters, our method necessitates this recursive total covariance decomposition. Moreover, in standard proofs, we need to quantify the increase in posterior precision for the action taken A, in each round E [n]. However, in dTS, our analysis extends beyond this. We not only quantify the posterior information gain for the taken action but also for every latent parameter, since they also learned. To elaborate, we use our recursive posteriors that connect the posterior covariance of each latent parameter 2/2 with the covariance of the posterior action parameters 0,, This allows us to propagate the information gain associated with the action taken in round A, to all latent parameters d/,, for E [L] by induction. Details are given in Appendix F.
Limitations. We identified several limitations that should be addressed in future work. First, our Bayes regret analysis is established only for the collapses to hierarchy of Gaussian distributions and dTS becomes exact Thompson Sampling. While this setting does not require diffusion model, it validates posterior approximation (exact in this
case, where the diffusion prior

---

limit) and clarifies how prior structure and diffusion depth L affect regret. Extending the theory to nonlinear diffusion or non- Gaussian rewards remains open. Second, for general nonlinear cases, dTs employs (i) a Laplace approximation for the reward likelihood and (ii) layer-wise linearization of diffusion links. A full analysis should account for errors coming from both approximations. We leave formal guarantees for future work. Third, dTS relies on pre-trained diffusion prior. With scarce or biased offline data, the prior may be reder--egularized. Empirically, performance degrades gracefully: dTS still outperforms LinTS and HierTS with as little as 1-5% pretraining data. Overall, dTS is advantageous when actions exhibit structured correlations and some offline data exist. In unstructured or purely online regimes, simpler methods such as LinTS may suffice.
4.1 Discussion
Computational benefits. Action correlations prompt an intuitive approach: marginalize all latent parameters and maintain a joint posterior of (0a)ae(K) H, Unfortunately. this is computationally inefficient for large action spaces. To illustrate, suppose that all posteriors are multivariate Gaussians. Then maintaining the joint posterior (0o)aE[K] Ht, necessitates converting and storing its dK x dK- dimensional covariance matrix, leading to d() and OKK2d²) time and space complexities. In contrast, the time and space complexities of dTS are ( + K) K)d") and (( K) This because dTS requires converting and storing L K covariance matrices, each being d x d- dimensional. The improvement is huge when K > L. which is common in practice. Certainly, more straightforward way to enhance computational efficiency is discard latent parameters and maintain K individual posteriors, each relating to an action parameter 0g E Rd (LinTS). This improves time and space complexity to O( dd)) and O (Kd).. However, LinTS maintains independent posteriors and fails to capture the correlations among actions; it only models 0a Hr rather than 0a H2 as done by dTS. Consequently, LinTS incurs higher regret due to the information loss caused by unused interactions of similar actions. Our regret bound and empirical results reflect
this aspect.
Statistical benefits. We do not provide matching lower bound. The only Bayesian lower bound that we know of is (n)) for much simpler K -armed bandit [45, Theorem 3]. All seminal works on Bayesian bandits do not match and providing such lower bounds on Bayes regret still relatively unexplored (even in standard settings) compared to the frequentist one. Also, a min-max lower bound of ((dyn) was given by Dani et al. [21]. In this work, we argue that our bound reflects the overall structure of the problem by comparing dTS to algorithms that only partially use the structure or do not use it at all as follows.
When the link functions are linear, we can transform the diffusion prior into Bayesian linear model (LinTS) by marginalizing out the latent parameters; in which case the prior on action parameters becomes 02 ~ (O, Σ), with the 0,, being not necessarily independent, and Σ is the marginal initial covariance of action parameters and it writes Σ of Σí 2 B, with B, -II W;, Then, it is tempting to directly apply LinTS to solve our problem. This approach will induce higher regret because the additional uncertainty of the latent parameters despite integrating them. This causes the marginal action uncertainty Σ to be much higher than the conditional action uncertainty ofLa, since we have Σ ofla + Σ/_ BrB discrepancy leads to higher regret, especially when K is large. This is due to LinTS needing to learn K independent d-dimensional parameters, each with considerably higher initial covariance Σ. This also reflected by our regret bound. To simply comparisons, suppose that o 2 max 0z so that o < 2. Then the regret bounds of dTS (where we bound ofAx by 2¹) and LinTS read
accounted for in Σ
of This
dTS Ö (n Et d
LinTS Ö (yndk + -b=)),4))
Then regret improvements captured by the variances or and the sparsity dimensions de, and proceed to illustrate this through the following scenarios.
(I) Decreasing variances. Assume that oe = 2 for any E [L + 1]. Then, the regrets become
dTS ((Vn(dda + Σi_ dd4%)).
LinTS ((Vndk2²)
Now to see the order of gain, assume the problem is high-dimensional (d >>> 1), and set T. = log, (d) and de = d Then the regret of dTS becomes Ö / /d([ + L))), and hence the multiplicative factor 2 in LinTS is removed and replaced with a smaller additive factor L.
7

---

(II) Constant variances. Assume that de = 1 for any E [L + 1]. Then, the regrets become
dTs: O nldk Ei=1 dr2%)))
LinTS: O (VndkL)
/d Similarly, let L = (d), and de = J Then dTS's regret is ( K + L)). Thus the multiplicative factor LinTS is removed and replaced with the additive factor By comparing this to (I), the gain with decreasing variances is greater than with constant ones. In general, diffusion models use decreasing variances [30] and hence we expect great gains in practice. All observed improvements in this section could become even more pronounced when employing non-linear diffusion models. In our theory, we used linear diffusion models, and yet we can already discern substantial differences. Moreover, under non-linear diffusion Eq. (1), the latent parameters cannot be analytically marginalized, making LinTS with exact marginalization inapplicable.
Regret independent of K? dTS's regret bound scales with K Koî rather than K E. which is particularly advantageous when 01 is small, as is often the case with diffusion model priors. Both our theoretical bound and empirical results show that dTS's advantage over LinTS increases as the action space grows. Nevertheless, dTS's 's regret still depends on K; this dependence arises from the problem setting rather than from the algorithm itself. Prior works [25, 67, 70] have proposed bandit algorithms whose regret does not scale with K. This difference stems from the setting considered: we study the disjoiní (per-action) case r(z, 0) x: Oa. where 0 (0a)aE[K) E RaK requiring the learning of Kd parameters and thus introducing an inherent dependence on when σn O. In contrast, setting described in Remark 1, where r(x, a; 8) = (0x, a) T f and only a single d-dimensional parameter must be learned. However, this formulation requires access to the feature map yp. Fortunately, dTS is also compatible with this setting (Section 3.3), in which case its regret would indeed be independent K.
K K-independent regret results are obtained in the shared-paramend
Experiments
Experimental setup. We evaluate dTS using both synthetic and MovieLens problems. In our experiments, we run 50 random simulations and plot the average regret with standard error. Our main contribution is to demonstrate that pretraining diffusion model offline enables the construction of ex- pressive and informative priors that substantially improve exploration efficiency in contextual bandits. We first evaluate dTS in setting where the prior matches the true generative process (Section 5.1 to isolate the benefit of informative priors), and then consider a misspecified regime (Section 5.2 and Appendix G) where the prior is either trained These experiments show that even when the prior is imperfect, dTS maintains strong performance: highlighting its robustness and practical relevance. Code can be found in this Gut repository.
-attrf-disttibution data intentionally perturbed.
5.1 True prior is diffusion model
Synthetic bandit problems are generated from the diffusion model in Eq. (1) with both linear and non-linear rewards. Linear rewards follow P(- | x;; 0a) NNxx' 0n, 1), while non-linear rewards are binary from P(- x;; 0a) Ber(g(x 0a)). with as the sigmoid function. Covariances are Dg = la. and contexts X, are uniformly drawn from [-1, We vary d E (5, 20}. L E (2, 4}, 5000, considering both linear and non-linear models.
K {102, 1043, and set the horizon
Linear diffusion. We consider Eq. (1) with fe(()) = Wew. where We uniformly drawn from [-1, dxd Sparsity is introduced by zeroing the last d, columns of We as W,- = (We, O.,d-d)) For d ==5 5 and T, = 2, d2) (5, 2); for d = 20 and I. = 4, (di, d2, dg. d4) (20, 10, 10,5,2) 2).
Non-linear diffuusion. We consider Eq. (1) where fe are 2-layer neural networks with random weights in [1,, ReLU activation, and hidden layers of size h = 20 for d = 5, and /t = 60 for d = 20.
Baselines. For linear rewards, we use LinUCB [1], LinTS [3], and HierTs [33]. marginalizing out all latent parameters except v/z, which corresponds to HierTs-1 in Appendix D. For non- lonelr rewards, include UCB-GLM [50] and GLM-TS [18]. We exclude GLM-UCB [24] due to high regret and HierTS as it's designed for linear rewards. We name dTS as dTS-dr, where d refers to diffusion type (L for linear, N for non-linear) and r indicates reward type (L for linear, N for non-linear). For example, dTS-LL signifies dTS in linear diffusion with linear rewards.
Results and interpretations. Results are shown in Fig. I and make the following observations:
8

---

Linear diffusion, linear reward Km100, Lm2, dinb
Linear diffusion, nonlineer reward
Nonlinar diffusion,, lineer reward
Nonlinear diffusion, nonlineer reward
1.8
K= L-2, dmn
104
K=100, LmZ, dins
1e3
Km 100, LmZ, disb
1.5
HierTS
14
dTS-LL
LinT5
dTS-NIL
Lint'
GLM-IS
UnDCC
GUM-18
Linuce
0.8
0.6
GLM
0.4
GLM
0.2
0.7
1000 2000 3000 4000 5000
2000 3000 4000 5000
1000 2000 3000 5000
1000
2000 3000 4000 5000
K-10000, LH4,, d-20
2c3 K=00000 Le4, d=20
1e7 K=10000, L-44, Ci=20
1e2 K- 10000 L.4, L=4, d-20
4.0
3.5
4TS-LL
2.5
dIT--IL
2.5
2.0
3.0
LinTs
HeerTs
2.0
OTS-LL
GLM-TS
LinTS
LinUce
dTS NL GLMITS UCB-GLM
1.5
Linugg
1.0
UCB
1.0
0.5
0.5
1000 2000 3000 4000 5000
0
1000 2000 3000 4000 5000
1000 2000 3000 4000 5000
1000 2000 3000 4000 5000
Hound hn
Hound Jn
Hound [c(
Round to
Figure I: Regret of dTS with varying diffusion and reward models and varying parameters d, K, L.
1) dTS demonstrates superior performance (Fig. 1). dTS consistently outperforms the baselines across all settings, including the four combinations of linear/non linear diffusion and reward (columns in Fig. 1) and both bandit settings with varying K, I., and d (rows in Fig. 1).
2) Latent diffusion structure may be more important than the reward distribution. When rewards are non-linear (second and fourth columns in Fig. 1), we include variants of dTS that use the correct diffusion prior but the wrong reward distribution, applying linear-Gaussian instead of BentsilllB (dTS-LL in the second column and dTS-NL in the fourth). Despite the reward misspecification, these variants outperform models using the correct reward distribution but ignoring the latent diffusion structure, such as GLM-TS and UCB-GLM. This highlights the importance of accounting for latent structure, which can be more critical than an accurate reward distribution.
3) Performance gap between dTS and LinTS widens as K increases (Fig, 2a). To show dTS's improved scalability, evaluate its performance with varying values of K E [10, 5 x 10*), in the linear diffusion and rewards setting. Fig. 2a shows the final cumulative regret for varying K values for both dTS-LL and LinTS, revealing widening performance gap K increases.
4) Regret scaling with K, and L matches our theory (Fig. 2b). We assess the effect of the number of actions K, context dimension d, and diffusion depth L on dTS's regret. Using the linear diffusion and rewards setting, for which we have derived a Bayes regret upper bound, we plot dTS -VSS-LL's 's regret across varying values of K {10, 100, 500, 1000), d (5, 10, 15, 20}. and L {2, 4, 5, 6) in Fig. 2b. As predicted by our theory, the empirical regret increases with larger values of K, d, or L, as these make the learning problem more challenging. leading to higher regret.
5) Diffusion prior misspecification (Fig. 2c). Here, dTS's diffusion prior parameters differ from the true diffusion prior. In the linear diffusion and reward setting, we replace the true parameters We and E, with misspecified ones, W, + é and Σ, + E2, where é and é2 are uniformly sampled from [v. + 0.5]dxd, with controlling the misspecification level. We vary E {0.5, 1, 1.5) and assess dTS's performance, comparing it to the well-specified dTS-LL and the strongest baseline in this fully-linear setting, HierTs. As shown in Fig. 2c, dTS's performance decreases with increasing misspecificatioon but remains superior to the baseline, except at v = 1.5, where their performances are comparable. Additional misspecification experiments are presented in Section 5.2, where the bandit environment is not sampled from a diffusion model.
5.2
True prior is not a diffusion model
Swiss roll data. Unlike previous experiments, the true action parameters are now sampled from the Swiss roll distribution (see Fig. 4 in Appendix G.1), rather than from diffusion model. The diffusion model used by dTS is pre-trained on samples from this distribution, with the offline pre-training procedure described in Appendix G.2. Fig. 3a shows that larger sample sizes increase the performance gap between dTS and LinTS. More samples improve the estimation of the diffusion prior (see Fig. 4 in Appendix G. I), leading to better dTS performance. Notably, comparable performance was achieved with few as 10 samples, and dTS outperformed LinTS by factor of 1.5 with just 50 samples.

---

1.8
1ee
function
regres Lindfs
Effect of prlor misspecifiatioo
1.6
dTS LL
3000
2500
1.2
1.0
Lin75
2000
2500
Regret
2000
2500
(v=0.5)
LindTS (vm1)
Bagrez
0.6
increasing
1000
UnddS (v-1.5)
500
Increasing
500
Hierts
0.0
5000
50000
0.0
0.5
1.0
1.3
2.0
2.5 3.0
0
Number actions
Increasing
Round re
(a) Perf. gap increases with K.
(b) Regret scaling with K, d. L.
(c) Diffusion prior misspecification.
Figure 2: Effect of various factors on dTS's performance.
3.5
Cumulattv regrot
3.0
2.5
Lint dTS
3.0
UnT5 reges
1.5
1.0
2.0
dTS
LinTS
10
50
100
2500 5000 1000250000
10
70
100
20
40
60
80
100
Number
Difusion death
Round
(a) Ratio of LinTS/dTS cumulative regret in last round varying
(b) Ratio of LinTS/dTS cumulative
(c) Regret of dTS in MovieLens.
regret in the last round
vary-
diffusion model with L =
pre-training sample size in [10, 10%. Higher values mean bigger performance gap.
ing diffusion depth T. in [2, 100].
pre-trained on embeddings ohtained
Higher values mean a bigger per-
by low-rank factorization of Movic- Lens rating matrix.
formance gap.
Figure 3: (a) and (b): Impact of pre-training sample size and diffusion depth for the Swiss roll data. (c): Regret of dTS in MovieLens.
While more samples may be required for more complex problems, LinTS would also struggle in such cases. Therefore, we expect these gains be even more sigmificant in more challenging settings.
We studied the effect of the pre-trained diffusion model depth and found that ~ 40 yields the best performance, with drop beyond that point (Fig. 3b). While our theory doesn't 't apply directly here, as it assumes a linear diffusion model, it still offers some intuition on the decreased performance for L > 40. The theorem shows dTS's regret bound increases with L when the true distribution is diffusion model. For small L, the pre-trained model doesn't It fully capture the true distribution, making the theorem inapplicable, but at L 40, the distribution is nearly captured, and further increases in lead to higher regret, consistent with our theory.
MovieLens data. We also evaluate dTS using the standard MovieLens [46] setting. In this semi- synthetic experiment, user is sampled from the rating matrix in each interaction round, and the reward is the rating the user gives movie (see Clavier al. [20, Section 5] for details about this setting). Here, the true distribution of action parameters is unknown and not a diffusion model. The diffusion model is pre- pre-trained offline estimates of action parameters obtained through low- rank factorization of the rating matrix. Fig. 3c demonstrates that dTS outperforms LinTS in this setting, Additional CIFAR-10 ablation studies are provided in Appendix G.4 where similar strong improvements are observed.
Conclusion
We use pre-trained diffusion model as a strong and flexible prior for dTS. Diffusion pre-training leverages abundant offline data, which is then fine-tuned through online interactions via our tractable posterior approximation. This approximation enables efficient posterior sampling and updates while maintaining strong empirical performance. Moreover, dTS admits simple Bayesian regret bound in the Gaassinn setting. Broader impact and computational considerations are discussed in Appendices I and J, and directions for future work are provided in Appendix H.
10

---

References
[1] Yasin Abbasi-Yann Yadkori, David Pal, and Csaba Szepesvari. Improved algorithms for linear stochastic bandits. In Advances in Neural Information Processing Systems 24, pages 2312-2320, 2011.
[2] Mare Abeille and Alessandro Lazaric. Linear Thompson sampling revisited. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, 2017.
[3] Shipra Agrawal and Navin Goyal. Thompson sampling for contextual bandits with linear payoffs. In Proceedings of the 30th International Conference on Machine Learning, pages 127-135, 2013.
[4] Shipra Agrawal and Navin Goyal. Further optimal regret bounds for thompson sampling. In Proceedings of the Sixteenth International Conference Artificial Intelligence and Statistics, pages 99-107, 2013.
[5] Shipra Agrawal and Navin Goyal. Near-optimal regret bounds for thompson sampling, Journal of the ACM (JACM), 64(5):1-24, 2017.
[6] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making arXiv preprint arXiv:2211.1500 2022.
[7] Imad Aouali. Linear diffusion models meet contextual bandits with large action spaces. In NeurlPS 2023 Workshop on Foundation Models for Decision Making, 2023.
[8] Imad Aouali, Branislav Kveton, and Sumeet Katariya. Mixed-effect thompson sampling. In International Conference on Artificial Intelligence and Statistics, pages 2087-2115. PMLR, 2023.
(9] Imad Aouali, Victor-Emmanuel Brunel, David Rohde, and Anna Korba. Unified pac-bayesian study of pessimism for offline policy learning with regularized importance sampling. In Uncertainty in Artificial Intelligence, pages 88-109. PMLR, 2024.
[10] Imad Aouali, fictor-Emmanne Brunel, David Rohde, and Anna Korba. Bayesian off-policy evaluation and learning for large action spaces. In International Conference on Artificial Intelligence and Statistics, pages 136- 144. PMLR, 2025.
[11] Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine Learning, 47:235-256, 2002.
[12] Mohammad Gheshlaghi Azar, Alessandro Lazaric, and Emma Brunskill. Sequential transfer in multi-armed bandit with finite set of models. In Advances in Neural Information Processing Systems 26, pages 2220-2228, 2013.
[13] Hamsa Bastani, David Simchi-Levi and Ruihao Zhu. Meta dynamic pricing: Transfer learning across experiments. Management Science, 1865)-1865-1885 2022.
[14] Soumya Basu, Branislav Kveton, Manzil Zaheer, and Csaba Szepesvari. No regrets for learning the prior bandits. In Advances in Neural Information Processing Systems 34, 2021.
[15] Christopher M Bishop. Pattern Recognition and Machine Learning, volume 4 of Information science and statistics. Springer, 2006.
[16] Leonardo Cella, Alessandro Lazaric, and Massimiliano Pontil. Meta-learning with stochastic linear bandits. In Proceedings the 37th International Conference Machine Learning, 2020.
[17] Leonardo Cella, Karim Lounici, and Massimiliano Pontil. Multi-task representation learning with stochastic linear bandits. arXiv preprint 20X1v:2202.10066, 2022.
[18] Olivier Chapelle and Lihong Li. An empirical evaluation Thompson sampling. In Advances in Neural Information Processing Systems 24, pages 2249-2257, 2012.

---

[19] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Mare L Klasky, and Jong Chul Ye. Diffu- sion posterior sampling for general noisy inverse problems. arXiv preprint MxXiv:2209.14687 2022.
[20] Pierre Clavier, Tom Huix, and Alain Durmus. Vits: Variational inference thomson sampling for contextual bandits. arXiv preprint arXiv:2307. 2023.
[21] Varsha Dani, Thomas Hayes, and Sham Kakade. The price of bandit information for online optimization. In Advances in Neural Information Processing Systems 20, pages 345-352, 2008.
[22] Aniket Anand Deshmukh, Urun Dogan, and Clayton Scott. Multi-task learning for contextual bandits. In Advances in Neural Information Processing Systems 30, pages 4848 4856, 2017.
[23] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 14:8780-8794 2021.
[24] Sarah Filippi, Olivier Cappe, Aurelien Garivier, and Csaba Szepesvari. Parametric bandits: The generalized linear case. In Advances in Neural Information Processing Systems 23, pages 586-594, 594, 2010.
[25] Dylan J Foster, Claudio Gentile, Mehryar Mohri, and Julian Zimmert. Adapting mis- specification in contextual bandits. Advances in Neural Information Processing Systems, 33: 11478-11489, 2020.
[26] Claudio Gentile, Shuai Li, and Giovanni Zappella. Online clustering of bandits. In Proceedings of the 31st International Conference on Machine Learning, pages 757-765, 2014.
[27] Aditya Gopalan, Shie Mannor, and Yishay Mansour. Thompson sampling for complex online problems. In Proceedings of the 31st International Conference on Machine Learning, pages 100-108, 2014.
[28] Amaury Gouverneur Borja Saarfguez--Gälver Tobias Oechtering, and Mikael Skoglund. Thompson sampling regret bounds for contextual bandits with sub-gaussian rewards. In 2023 IEEE International Symposium on Information Theory (ISIT), pages 1306-1311. IEEE, 2023.
[29] Samarth Gupta, Shreyas Chaudhari, Subhojyoti Mukherjee, Gauri Joshi, and Osman Yagan. A unified approach to translate classical bandit algorithms to the structured bandit setting. IEEE Journal on Selected Areas in Information Theory, 1(3):840 14380,33,, 2020.
[30] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-668 68000068811100 2020.
[31] Joey Hong, Branislav Kveton, Manzil Zaheer, Yinlam Chow, Amr Ahmed, and Craig Boutilier. Latent bandits revisited. In Advances in Neural Information Processing Systems 33, 2020.
[32] Joey Hong, Branislav Kveton, Sumeet Katariya, Manzil Zaheer, and Mohammad Ghavamzadeh. Deep hierarchy in bandits. In International Conference on Machine Learning, pages 8833-8851. PMLR, 2022.
[33] Joey Hong, Branislav Kveton, Manzil Zaheer, and Mohammad Ghavamzadeh. Hierarchical Bayesian bandits. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics, 2022.
[34] Yu-Guan Hsich, Shiva Prasad Kasiviswanathan, Branislav Kveton, and Patrick Blöbaum. Thompson sampling with diffusion generative prior. arXiv preprint wxXiv:2301.05188 2023.
[35] Jiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li, and Liwei Wang. Near-optimal representation learning for linear bandits and linear rl. In International Conference on Machine Learning, pages 4349 4358. PMLR, 2021.
[36] Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint 2rXtv:2205.09991, 2022.
12

---

[37] Emilie Kaufmann, Nathaniel Korda, and Rémi Munos. Thompson sampling: An asymptotically optimal anite-time analysis. In International conference on algorithmic learning theory, pages 199-213. Springer, 2012.
[38] Daphne Koller and Nir Friedman. Probabilisticc Graphical Models: Principles and Techniques. MIT Press, Cambridge, MA, 2009.
[39] Nathaniel Korda, Emilie Kaufmann, and Remi Munos. Thompson sampling for I-dimensional exponential family bandits. Advances in neural information processing systems, 26, 2013.
[40] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[41] John Kruschke. Bayesian data analysis. Wiley Interdisciplinary Reviews: Cognitive Science, 1(5):658 116666666,.--6 2010.
[42] Branislav Kveton, Manzil Zaheer, Csaba Szepesvari, Lihong Li, Mohammad Ghavamzadeh, and Craig Boutilier. Randomized exploration in generalized linear bandits. In International Conference on Artificial Intelligence and Statistics, pages 2066-2076. PMLR, 2020.
[43] Branislav Kveton, Mikhail Konobeev, Manzil Zaheer, Chih-Wei Hsu, Martin Mladenov, Craig Boutilier, and Csaba Szepesvari. Meta-Thomppon sampling. In Proceedings of the 38th International Conference on Machine Learning, 2021.
[44] Branislav Kveton, Boris Oreshkin. Youngsuk Park. Aniket Anand Deshmukh, and Rui Song. Online posterior sampling with diffusion prior. Advances in Neural Information Processing Systems,
130484, 2024.
[45] Tze Leung Lai. Adaptive treatment allocation and the multi-armed bandit problem. The Annals of Statistics, 15(3)) 1091-1144, 1987.
[46] Shyong Lam and Jon Herlocker. MovieLens Dataset. http://gooooeee...ryccttttttseer 2016.
[47] Tor Lattimore and Remi Munos. Bounded regret for finite-armed structured bandits. In Advances in Neural Information Processing Systems 27, pages 550-558, 2014.
[48] Yann LeCun, Corinna Cortes, and Christopher Burges. MNIST Handwritten Digit Database. http:// /ttt://yann... lecun. com/ wtt.////////c/cieeeeeee exdb/mnist, 2010.
[49] Lihong Li, Wei Chu, John Langford, and Robert Schapire. A contextual- approach to personalized news article recommendation. In Proceedings the 19th International Conference on World Wide Web, 2010.
[50] Lihong Li, Yu Lu, and Dengyong Zhou, Provably optimal algorithms for generalized linear contextual bandits. In Proceedings of the 34th International Conference on Machine Learning, pages 2071-2080) 2017.
[51] Dennis Lindley and Adrian Smith. Bayes estimates for the linear model. Journal of the Royal Statistical Society: Series B (Methodological), 34(1):1-18, 1972.
[52] Xiuyuan Lu and Benjamin Van Roy. ceeemattio--thorettte confidence bounds for reinforcement learning. In Advances in Neural Information Processing Systems 32, 2019.
[53] Ambrri Maillard and Shie Mannor. Latent bandits. In Proceedings of the 31st International Conference on Machine Learning. pages 136-144, 2014.
[54] McCullagh and J. A. Nelder. Generalized Linear Models. Chapman & Hall, 1989.
[55] Gergely Neu, Iuliia Olkhovskaia, Matteo Papini, and Ludovic Schwartz. Lifting the information ratio: An information-theoreeie analysis of thompson sampling for contextual bandits. Advances in Neural Information Processing Systems, 35:9486 0498699988. 2022.
[56] Amit Peleg, Naama Pearl, and Ron Meirr. Metalearning linear bandits by prior update. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics, 2022.
13

---

[57] Carlos Riquelme, George Tucker, and Jasper Snoek. Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling. arXiv preprint 1802.09 2018.
[58] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 10684-10695, 2022.
[59] Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Mathematics of Operations Research, 39(4): 221-)223,,, 2014.
[60] Steven Scott. A modern bayesian look at the multi-armed bandit. Applied Stochastic Models in Business and Industry, 26:639 658, 2010.
[61] Max Simchowitz, Christopher Tosh, Akshay Krishnamurthy, Daniel Hsu, Thodoris I Lykouris, Miro Dudik, and Robert Schapire. Bayesian decision-making under misspecified priors with applications to meta-learning.
Advances in Neural Information Processing Systems 34, 2021.
[62] Jascha Sohl-Dicksttin, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference machine learning, pages 2256-2265. PMLR, 2015.
[63] Runzhe Wan, Lin Ge, and Rui Song. Metadata-based multi-task bandits with Bayesian hierar- chical models. In Advances in Neural Information Processing Systems 34, 2021.
[64] Runzhe Wan, Lin Ge, and Rui Song. Towards scalable and robust structured bandits: meta- learning framework. In International Conference on Artificial Intelligence and Statistics, pages 1144- 1174. PMLR, 2023.
[65] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint 2rXiv:2208.06193, 2022.
[66] Neil Weiss. A Course in Probability. Addison- Addison-Wesley. 2005.
[67] Yunbei Xu and Assaf Zeevi. Upper counterfactual confidence bounds: new optimism principle for contextual bandits. arXiv preprini urXiv:2007.07876 2020.
[68] Jiaqi Yang, Wei Hu, Jason D Lee, and Simon S Du. Impact of representation learning in linear bandits. arXiv preprint XXx::201000653/ 2020.
[69] Tong Yu, Branislav Kveton, Zheng Wen, Ruiyi Zhang, and Ole Mengshoel. Graphical mod- els meet bandits: A variational Thompson sampling approach. In Proceedings of the 37th International Conference on Machine Learning, 2020.
[70] Yinglun Zhu, Dylan Foster, John Langford, and Paul Mineiro. Contextual bandits with large action spaces: Made practical. In International Conference Machine Learning, pages 27428-27453. PMLR, 2022.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Section 4 and Appendix E and Section 5
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Main text
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

Justification:
Section 4 and Appendix E
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section and Appendix G
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

Answer: [Yes]
Justification: Code provided in supplementary materials
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/ Belliiigguides/Cooee
wwwiiissssopllooooooooooo clycc for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ //iips.cooo wttpccccppppcceeeecececccccccccc httpa///ggiieee/eeeuuuubiiiiiiia /Codeeeooiissseeeeee cetpuubmmsssoonni ssisnnoll onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section and Appendix G Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Section 5 and Appendix G
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
17

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //wuurrrrsp.../ cc
Answer: [Yes]
Justification:
The authors confirm reading and adhering the NeurIPS Code of Ethics
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Appendix
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
18

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: This is bandit paper whose algorithms and data are not pretrained language models, image generators, or scraped datasets.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The Authors coded themselves all the baselines and algorithms. MovieLens data was used and cited properly.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
19

---

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
Supplementary materials contains the well-documented code
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You either create an anonymized URL or include an anonymized zip file.
Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
No crowdsourcing was involved
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects were involved
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact core methodology, scientific rigorousness, originality of the research, declaration is not required.
20

---

Answer: [NA]
Justification: LLMs were only used in writing polishing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: htmp::///neee //tpur/eessp./ee/e /2//Connerencce/22255///// for what should or should not be described.
21

---

Supplementary materials
Notation, For any positive integer ni, define [n] (1, Let 01: U1:-..;Up E TR be vectors, (Us)ke(n) E Rnd is the nd-dimensional vector obtained by concatenating ",,.... un For any matrix E Rdxd X1 (A) and (A) denote the maximum and minimum eigenvalues of A, respectively. Finally, we write O for the big-O notation up to polylogarithmic factors.
Table of notations.
Table 1: Notation.
Symbol
Definition
n X
Learning horizon
Context space Number of actions
K
[K]
Set of actions
d
Dimension of contexts and action parameters d d-dimensional parameter of action a E [K] Reward distribution of context x and action Reward function of context and action Bayes regret after n interactions
On
p( 0z)
r(z,a; 0.)
BR(n)
N
Multivariate Gaussian distribution of parameters μ and Σ Multivariate Gaussian density of parameters and Σ Diffusion model depth ê-th d-dimensional latent parameter
N(:AE))
L
ft
Link functions of the diffusion model
Et
Covariances of the link function History of interactions
Hi
A Extended related work
Thompson sampling (TS) operates within the Bayesian framework and involves specifying prior/likelihood model. In each round, the agent samples unknown model parameters from the current posterior distribution. The chosen action the one that maximizes the resulting reward. TS is naturally randomized, particularly simple to implement, and has highly competitive empirical performance both simulated and real-world problems [59, 18). Regret guarantees for heuristic remained open for decades even for simple models. Recently, however, significant progress has been made. For standard multi-armes bandits, TS optimal in the Beta-Bernoulli model [37, 4), Gaussian-Grr Gaussian model [4], and in the exponential family using Jeffrey's 's prior [39]. For linear bandits, TS is nearly-optimal [59, 5, 2]. In this work, we build TS upon complex diffusion priors and analyze the resulting Bayes regret [59, 55, 28] in the linear contextual bandit setting.
Decision-making with diffusion models gained attention recently, especially in offline learning [6, 36, 65]. However, their application in online learning was only examined by Hsieh et al. [34], which focused on meta-learning we expand the scope of Hsieh et al. [34] to encompass the broader contextual bandit framework. In particular, we provide theoretical analysis for linear instances, effectively capturing the advantages of using diffusion models as priors in contextual Thompson sampling. These linear cases are particularly captivating due to closed-fome form posteriors, enabling both theoretical analysis and computational efficiency; an important practical consideration.
multi-armed bandits without theoretical guarantees. In this work.
Hierarchical Bayesian bandits [13, 43, 14, 61, 63, 33, 56, 64, 8] applied TS simple graphical models, wherein action parameters are generally sampled from a Gaussian distribution centered at single latent parameter. These works mostly span meta- and multi-task learning for multi-armes bandits, except in cases such as Aouali et al. [8], Hong et al. [32] that consider the contextual bandit setting. Precisely, Aouali et [8] assume that action parameters are sampled from Gaussian distribution centered at linear mixture of multiple latent parameters. On the other hand, Hong et al. [32] applied TS to graphical model represented by tree. Our work can seen as an extension of all these works to much more complex graphical models, for which both theoretical and algorithmic
22

---

foundations are developed. Note that the settings in most of these works can be recovered with specific choices of the diffusion depth L and functions ft. This attests the modeling power of dTS.
Approximate Thompson sampling is a major problem in the Bayesian inference literature. This is because most posterior distributions intractable, and thus practitioners must resort to sophisticated computational techniques such as Markov chain Monte Carlo [41]. Prior works [57, 18, 42] highlight the favorable empirical performance of approximate Thompson sampling. Particularly, [42] provide theoretical guarantees for Thompson sampling when using the Laplace approximation in generalized linear bandits (GLB). In our context, we incorporate approximate sampling when the reward exhibits non-linearity. While our approximation does not come with formal guarantees, it enjoys strong practical performance. An in-depth analysis of this approximation is left a direction for future works. Similarly, approximating the posterior distribution when the diffusion model is non non-linear as well as analyzing is an interesting direction of future works.
Bandits with underlying structure also align with our work, where we assume a structured rela- tionship among actions, captured by diffusion model. In latent bandits [53, 31], single latent variable indexes multiple candidate models. Within structured finite-armed bandits [47, 29], each action is linked to a known mean function parameterized by a common latent parameter. This latent parameter is learned. TS was also applied to complex structures [69, 27J. However, simultaneous computational and statistical efficiencies aren't 't guaranteed. Meta- and multi-tast learning with upper confidence bound (UCB) approaches have a long history in bandits [12, 26, 22. 16]. These, however, often adopt frequentist perspective, analyze stronger form of regret, and sometimes result in conservative algorithms. In contrast, our approach is Bayesian, with analysis centered on Bayes regret. Remarkably, our algorithm, dTS, performs well as analyzed without necessitating additional tuning. Finally, Low-rank bandits [35, 17, 68] also relate to our linear diffusion model when I, = 1. Broadly, there exist two key distinctions between these prior works and the special case of our model (linear diffusion model with L 1). First, they assume 02- W 13, whereas we incorporate additional uncertainty in the covariance Σ₁ to account for possible misspecification as 0,, = (W 1&,, E.). Consequently, these algorithms might suffer linear regret due to model misalignment. Second, we assume that the mixing matrix W₁ is available and pre-learned offline, whereas they learn it online. While this is more general, it leads computationally expensive methods that are difficult to employ in a real-world online setting.
The regret bound of dTS scales with Kof rather than K KE,o of. which is particularly advantageous when d1 is small: a common case in diffusion models with decreasing layer variances. In the limiting case d O, the regret becomes independent of K. Our theoretical analysis (Section 4.1) and empirical results both show that the performance gap between dTS and LinTs widens as the number of actions increases, highlighting dTS's suitability for large action spaces.
Large action spaces.
Some prior works [25, 67, 70] achieve regret bounds that do not scale with K. This discrepancy stems from the problem setting rather than the algorithm itself. Specifically, those works adopt a shared- parredd-rramtter model r(z, a) = p(x, a) 8 with a single parameter 0 Ied and a known feature map s2, whereas we study the disjoint case r(x, a) x: TTO, 0. with K separate d-dimensional parameters. In the shared-parameter setting (see Remark 2.1 and Section 3.3), dTS would similarly achieve regret independent of K.
In summary, the dependence on K arises from the modeling choice rather than a limitation of dTS. When & is available. dTS scales only with d; otherwise, in the per-action setting, it remains both computationally and statistically efficient (Section 4.1). Empirically, even for very large action spaces (c.g., K = 10t), dTS substantially outperforms existing baselines, with the performance gap increasing as K grows: highlighting its scalability to large action spaces.
B Posterior derivations for linear diffusion models
Our posterior approximation builds on the simplified setting where the diffusion model is fully linear, i.e., each link function ft is linear in ehz. This linear case, studied in our earlier work [7], serves as the analytical foundation for our posterior approximation used in the general non-linear case. In Appendix C, we show how the exact posteriors derived in this linear setting inspire our efficient approximation, which extends naturally to practical diffusion models that are typically highly non-linear.
23

---

B.1 Linear diffusion model
Here, we assume the link functions fe are linear such as fe(tz) Were for & [L], where
We E IRddd are known mixing matrices. Then, Eq. (1) becomes linear Gaussian system (LGS) [15] and can be summarized as follows
UL NiO, (0,,),+)),
(17)
F Ve NWWer Ee),
AL E [1/ Va [K]
0g 1/1 N(WIV Σ1), Y X,, OA. ~ P(. 0A)),
1A E [m].
This model is important, both in theory and practice. For theory, it leads to closed-form posteriors when the reward distribution is linear-Gaussian as P(. x; 0a) N(: 0n, )²). This allows bounding the Bayes regret of dTS. For practice, the posterior expressions are used to motivate efficient approximations for the general case in Eq. (1) as we show in Section 3.2. These derivations can be proven following standard techniques [15], and the reader may refer to Aouali [7, Appendix B] for example of how these posteriors can be derived in the case of contextual bandits.
B.2 Posterior derivation in the linear diffusion case
We now consider the linear link function case, where fe((e) = Wabe for E [L] (the setting above in Appendix B. Recall that the reward distribution is modeled as a generalized linear model (GLM) [54], allowing for non-linear rewards even when the diffusion links are linear. This non-linearity in the reward distribution prevents closed-form posteriors. However, since the non-linearity arises only through the reward likelihood, we approximate by a Gaussian, leading efficient posterior updates that are exact whenever the reward model itself is Gaussian; special case of the GLM framework.
Specifically, let P(- x;; 0.) be associated with action a is
aaponential-famill distribution. The log-likelihood of the data
lo, 0a) Σ iES..
[X,, 0g 1 (()
where a real-valued function and A is twice continuously differentiable, with derivative Á representing the mean function. Let Br,a and G.,a denote the maximum likelihood estimate (MLE) and the Hessian of the negative log- ligllikelihood, respectively:
B = arg max log pl He,a 0a),
GID= Σ B..)
(18)
0, ERE
iESt.
where Staa (e E [t 1): A a} is the set of rounds in which action was taken up round t.
We approximate the likelihood as
HHH (a) xe No Br. GT,d);
(19)
which renders all subsequent posteriors Gaussian. Once this approximation is done, all other derivations of the action posterior and latent posteriors are exact.
Action posterior. The conditional action posterior becomes
p(0a W1, H.a) N(Oa:fraa Er,),
with parameters
Erd DT- + Gi,,,
tro = Eg,a (IT W + GraBiaa.
(20)
Latent posteriors. For each E [Z] (1), the conditional latent posterior is
p(vt-1 Uee H.) z M
where
5-1 = 2x1 + Gre-1
P=,-1= = E 1 Br.c1).
(21)
24

---

The top-layer posterior is
p(VL Hi) = MELL:E.LEE..)))
with
E = E441 + Gt.L.
PL, = Et,
(22)
Recursive updates. recursion is
The matrices Gz and Bed for l [Z] defined recursively. The base
G1,1 = WT Σ
(ET--E) E
B E Σ E Er Gra
(23)
Then, for l E [4] (1}, the recursive step is
Gi - -., We
Bi,== -
(24)
Discussion. This completes the derivation of the linear posterior approximation. All posteriors Gaussian and exact whenever the reward distribution follows a linear-Gaussian model, i.e.
P( x; 0a) ~(; 0a:0²)
In this case, the above posterior updates coincide with the exact Bayesian updates, while for general GLMs they serve as efficient and accurate approximations.
C Posterior derivations for non-linear diffusion models
The general diffusion model (Eq. (1), which is our case of interest) involves two sources of non- linearity: (i) the reward distribution P(. x; x;0), which may follow a non-linear generalized linear model (GLM), and (ii) the diffusion links fe(Y)), which can be arbitrary non-linear functions. Both sources make the posterior intractable, and therefore two approximations are needed.
First approximation (likelihood), We first approximate the reward likelihood by Gaussian density (as we did above in Eq. (19)). After this substitution, the model becomes conditionally Gaussian given the latent variables. This step is exact when the reward model is Gaear-Ga and approximate otherwise.
Second approximation (diffusion hierarchy). Even after the likelihood is approximated, the diffusion hierarchy remains non-linear because of the non-linear mappings To handle this, reuse the exact Gaussian posteriors derived for the linear diffusion case (Appendix B.2) and generalize them as follows:
Replace each linear mapping Weth by its non-linear counterpart fe(Uz), which represents the mean of the diffusion prior at layer L.
Remove matrix multiplications involving W, in the recursive updates.
This step can be viewed as extending the linear-Gaussian posterior formulas to a general non-linear setting, without performing explicit lincarization or optimization.
Resulting approximation. The two steps above yield posterior where each conditional factor p(0, V1, .l..a) and p t/,, H,) remains Gaussian with updated means and covariances, while the overall model retains the hierarchical diffusion structure. The approximation satisfies two desirable properties: it exactly recovers the diffusion prior when data is available, and more data is observed, the likelihood terms dominate and the prior influence fades naturally.
This approach is computationally efficient, avoids costly posterior sampling or variational optimiza- tion, and remains expressive since the overall posterior still diffusion model with non-linear link functions and covariances that are now data-dependent.
25

---

D Additional discussions
D.1 Additional discussion: link two-level hierarchies
The linear diffusion Eq. (17) can be marginalized into a 2-1evel hierarchy using two different strategies. The first one yields,
1 (0,,
(25)
a 4D ~ ( 21),
Va E [K],
with Q2 = LiLa- + 1 oz BB and B₁ = -I W.. The second strategy yields,
1 N0,
(26)
0a 4/ NW,,, of
Va E [K]
where f22 = -/ o2 B,B Recently, HierTS [33] was developed for such two-level graphical models, and call HierTs under Eq. (25) by HierTs-1 and HierTs under Eq. (26) by HierTs-2 2. Then, we start by highlighting the differences between these two variants of HierTS. First, their regret bounds scale
HierTS-1: : (v nd( El og Lof HierTs-2: ( + Ei (211)
When Rs L., the regret bounds of HierTS-1 and HierTS-2 are similar. However, when K I. HierTS-2 2 outperforms HierTS-1. This is because HierTS-2 puts more uncertainty on a single d- dimensional latent parameter 1h1, rather than K individual d-dimenniional action parameters 0u. More importantly, HierTS-1 implicitly assumes that action parameters 0, are conditionally independent given vr, which is not true. Consequently, HierTS-2 outperforms HierTS-1. Note that, under the linear diffuusion model Eq. (17), dTS and HierTS-2 have roughly similar regret bounds. Specifically, their regret bounds dependency on K is identical, where both methods involve multiplying K by of and both enjoy improved performance compared to HierTs-1. That said, note that Theorem E.1 and Proposition E.2 provide an understanding of how dTS's ' regret scales under linear link functions fe, and do not say that using dTS is better than using HierTS when the link functions fe are linear since the latter can be obtained by a proper marginalization of latent parameters (i.e., HierTs-2 instead of HierTS-1). While such comparison not the goal of this work, we still provide it for completeness next.
When the mixing matrices We are dense (i.e., assumption (A3) is not applicable), dTS and HierTS-2 have comparable regret bounds and computational efficiency. However, under the sparsity assumption (A3) and with mixing matrices that allow for conditional independence of 1 coordinates given v/2, dTS enjoys a computational advantage over HierTs-22 2. This advantage explains why works focusing on multi-level hierarchies typically benchmark their algorithms against two-level structures akin to HierTs-1. 1. rather than the more competitive HierTs-2 2. This is also consistent with prior works in Bayesian bandits using multi-level hierarchies, such Tree-based priors [32], which compared their method to HierTs-1. In line with this, we also compared dTS with HierTS-1 in our experiments. But this is only given for completeness as this is not the aim of Theorem 1 and Proposition E.2. More importantly, HierTS is inapplicable in the general case in Eq. (1) with non- non-linear link functions since the latent parameters cannot be analytically marginalized.
D.2 Additional discussion: why regret bound depends on K and L
Why the bound increases with K? This arises due to our conditional learning of 9g given th Rather than assuming deterministic linearity, 0, W1 account for stochasticity by modeling 0₂- ~ WWW aa). This makes dTS robust to misspecification scenarios where 0, is not perfeetly linear with respect to d1, at the cost of additional learning of 0., V/. If we were to assume deterministic linearity (01 = 0), our regret bound would scale with only.
Why the bound increases with L? This is because increasing the number of layers L adds more initial uncertainty due to the additional covariance introduced by the extra layers. However, this does not imply that we should always use L = (the minimum possible L). Precisely, the theoretical results predict that regret increases with L. when the true prior distribution matches a diffusion model of depth L, as increasing L reflects a more complex action parameter distribution and hence more
26

---

complex bandit problem. However, in practice, when L is small, the pre-trained diffusion model may be too simple to capture the true prior distribution. violating the assumptions of our theory. Increasing L. improves the pre-trained model's quality, reducing regret. Once I. is large enough and the pre-trained model adequately captures the true prior distribution, the theoretical assumptions hold, and regret begins to increase with I,, as predicted. This is validated empirically in Fig. 3b.
E Formal theory
We analyze dTS assuming that: (A1) The rewards are linear P(. 2; 0a) = =N(x2 (2, 02). (A2) The link functions ft are linear such as fe(()) Weeh for E [L], where We E Rddd are known mixing matrices. This leads to a structure with L layers of linear Gaussian relationships detailed in Appendix B. In particular, this leads to closed-form posteriors given in Appendix B.2 that inspired our approximation and enable theory similar to linear bandits [3]. However, proofs are not the same, and technical challenges remain (explained in Appendix F).
Although our result holds for milder assumptions, we make additional simplifications for clarity and interpretability. We assume that (A3) Contexts satisfy ||x, l1 = 1 for any & E [n]. Note that (A3) can be relaxed to any contexts X, with bounded norms || |Xe|2 2- (A4) Mixing matrices and covariances satisfy A (W/ Wx) for any l [L] and Σ. = o La for any l E [L
+ 1]. (A4) can be relaxed to
positive definite covariances Eg and arbitrary mixing matrices We. In particular, this is satisfied once
we use diffusion model parametrized with linear functions. In this section, we write O for the big-O notation up to polylogarithmie factors. We start by stating our bound for dTS.
Theorem E.1. Let aAx = max + d.. There exists a constant e > 0 such that for any ô E (0, 1), the Bayes regret of dTS under (A1), (A2), (A3) and (A4) is bounded as
BR(n) < / 2n 2n((RCC(1) +
Σ R( log(1/0) + cno
not
of Go log (1 + of)
(n) = cidcK log (1+
2241
0 MAx
RT lod (1+
c log + +
(27)
Eq. (27) holds for any 8 E (0, 1). In particular, the term onô is constant when 8 1/n. Then, the bound is ō Vn(dKo + E'e 2 ( ofxx) and this dependence on the horizon n aligns with prior Bayes regret bounds. The bound comprises L + main terms, (ACT(M) and R for E [L]. First, RACT (n) relates to action parameters learning, conforming to standard form [52]. Similarly, R.AA is associated with learning the f-th latent parameter.
To include more structure, we propose the sparsity assumption (A5) Wt = (We, Od,d-a;), where W. E dx for any E [L]. Note that (A5) is not assumption when de d for any E E [L]. Notably, (A5) incorporates a plausible structural characteristic that a diffusion model could capture.
Proposition E.2 (Sparsity). Let 0a maxte[L+1]] + g.. There exists a constant > 0 such that for any å E (0, 1), the Bayes regret of dTS under (A1), (A2), (A3), (A4) and (A5) is bounded as
BR(n) < 2n (2((RACT()) +
Σ R( log(1/8) cnd.
(=1
no?
of log (1 + of)
R() =codK log (1 +
10
d
2241)
2 Max
28 MAX
RINT c,d, log (1+
of
c
log (1+ 1
(28)
27

---

From Proposition E.2, our bounds scales as
BR(n) - Ö6 mlakot ( +
+E deo 1 )
(29)
The Bayes regret bound has clear interpretation: if the true environment parameters are drawn from the prior. then the expected regret of an algorithm stays below that bound. Consequently, less informative prior (such as high variance) leads more challenging problem and thus higher bound. Then, smaller values of K, L, d or de translate to fewer parameters to learn, leading to lower regret. The regret also decreases when the initial variances of decrease. These dependencies are common in Bayesian analysis, and empirical results match them.
The reader might question the dependence of bound both L and K. Details be found in Appendix G.5, but in summary, we model the relationship between 0g and s/ stochastically as N(W141 L4) to account for potential nonlinearity. This choice makes the model robust to model misspecification but introduces extra uncertainty and requires learning both the 6g and the the. This results in a regret bound that depends on both K and L. However, thanks to the use of informative priors, our bound has significantly smaller constants compared to both the Bayesian regret for LinTS and its frequentist counterpart, as demonstrated empirically Appendix G.5 where it much tighter than both and in Section 4.1 where we theoretically compare our Bayes regret bound to that of LinTS.
Technical contributions. dTS uses hierarchical sampling. Thus the marginal posterior distribution of O, H2 is not explicitly defined. The first contribution is deriving 0, H, using the total covariance decomposition combined with induction proof, our posteriors were derived recursively. Unlike standard analyses where the posterior distribution of 0g H, is predetermined due to the absence of latent parameters, our method necessitates this recursive total covariance decomposition. Moreover, in standard proof', we need to quantify the increase in posterior precision for the action taken A, in each round E [n). However, dTS, our analysis extends beyond this. We not only quantify the posterior information gain for the taken action but also for every latent parameter, since they are also learned. To elaborate, we use our recursive posteriors that connect the posterior covariance of each latent parameter ahe with the covariance of the posterior action parameters 6.. This allows us to propagate the information gain associated with the action taken in round A, to all latent parameters the, for & E [L] by induction. Details are given in Appendix F.
F Regret proof
F.1 Sketch of the proof
We start with the following standard lemma upon which we build our analysis [8].
Lemma F.I. Assume that p(0a! H.) (Aaa ir,a, Eg,a) for any [K], then for any 5 (0,1), 1),
BR(n)= S 22n log( (omllo/1//) 1/0)
E EL=
+ end where > 0 is a constant. (30)
Applying Lemma F.1 requires proving that the marginal action-posteni posterior densities of 8, H, in Eq. (3) are Gaussian and computing their covariances, while we only know the conditional action- posteriors p(0, vI, H,) and putent-postrrios p(v:-1 vr,, H,). This is achieved by leveraging the preservation properties of the family of Gaussian distributions [38] and the total covariance decomposition [66] which leads to the next lemma.
Lemma F.2. Lett E [n] and E [K], then the marginal covariance matrix En, reads
En = Σq,a + EeELE) P...I..PTT. aet where Paε I I E E W (31)
The marginal covariance matrix Er, in Eq. (31) decomposes into + terms. The first term corresponds to the posterior uncertainty of 0, 1h1. The remaining L, terms capture the posterior uncertainties of 4/z and vir-1 1 the for & [1]/(1}. These are then used to quantify the posterior information gain of latent parameters after one round as follows.
28

---

Lemma F.3 (Posterior information gain), Lett E [n]] and l E [L], then
E - -Enp = o - - XcX, where d = maxxE[L+1)
1 +
(32)
Finally, Lemma F.2 used to decompose || L..A, in Eq. (30) into L + terms. Each term is
bounded thanks to Lemma E.3. This results in the Bayes regret bound in Theorem E.1.
F.2 Technical contributions
Our main technical contributions are the following.
Lemma F.2. In dTS, sampling is done hierarchically, meaning the marginal posterior distribution of 0a H, is not explicitly defined. Instead, we use the conditional posterior distribution of 9, HH, th1. The first contribution was deriving 0& HH, using the total covariance decomposition combined with an induction proof, our posteriors in Appendix B.2 were derived recursively. Unlike in Bayes regret analysis for standard Thompson sampling, where the posterior distribution of 0a H2 is predetermined due the absence of latent parameters, decomposition, marking a first difference from the standard Bayesian proofs of Thompson sampling. Note that Hi erTS, which is developed for multi-task linear bandits, also employs total covariance decomposition, but it does so under the assumption of a single latent parameter; on which action parameters are centered. Our extension significantly differs as it tailored for contextual bandits with multiple, successive levels of latent parameters, moving away from HierTS's assumption of a 1-level structure. Roughly speaking, HierTS when applied to contextual would consider single-level hierarchy, where 0./01 ~ (1 21) with L = 1. In contrast, our model proposes multi-level hierarchy, where the first level is 0z 61 N(W 11,, E1). This also introduces a new aspect to our approach the use of linear function W1 V1, as opposed to HierTS's ' assumption where action parameters are centered directly on the latent parameter. Thus, while HierTS also uses the total covariance decomposition, our generalize it to multi-level hierarchies under L linear functions Weeh. instead of single-level hierarchy under a single identity function 4/1.
method necessitates this recursive total covariance
Lemma F.3. In Bayes regret proofs for standard Thompson sampling, we often quantify the posterior information gain. This is achieved by monitoring the increase in posterior precision for the action taken A, in each round E [n). However, in dTS, our analysis extends beyond this. We not only quantify the posterior information gain for the taken action but also for every latent parameter, since they are also learned. This lemma addresses this aspect. To claborate, we use the recursive formulas in Appendix B.2 that connect the posterior covariance of each latent parameter the with the covariance of the posterior action parameters 0g. This allows us to propagate the information gain associated with the action taken in round A, all latent parameters the, for E [L] by induction. This a novel contribution, as it is not feature of Bayes regret analyses in standard Thompson sampling.
Proposition E.2. Building upon the insights of Theorem 1, introduce the sparsity assumption (A3). Under this assumption, we demonstrate that the Bayes regret outlined in Theorem E. can be significantly refined. Specifically, the regret becomes contingent on dimensions de S d, as opposed to relying on the entire dimension d. The underlying principle of this sparsity assumption is straightforward: the Bayes regret is influenced by the quantity of parameters that require learning. With the sparsity assumption, this number is reduced to less than d for each latent parameter. To substantiate this claim, we revisit the proof of Theorem E.1 and modify a crucial equality. This adjustment results in more precise representation by partitioning the covariance matrix of each latent parameter the into blocks. These blocks comprise d, x de segment corresponding to the learnable de parameters of tú,, and another block of size (d d,) (d - de) that does not necessitate learning. This decomposition allows us to conclude that the final regret is solely dependent on de, marking significant refinement from the original theorem.
29

---

F.3 Proof of lemma E.2
In this proof, we heavily rely on the total covariance decomposition [66]. Also, refer to [33, Section 5.2] for brief introduction to this decomposition. Now, from Eq. (20), we have that
cov 10 H2, vi] Er, = (Gra + +
E 10 H1, 4/1 = fie.a = Eea (Gr, + +i WT'W
First, given H4, cov 10, He, 41] (Graa + +2 is constant. Thus
E ccvv 10x Hr, wh1] H,] cov 10,1 He, tbi1 (Gr.a 2TT)² -1 In addition, given He, Σe,a. Gr.a and Bt,a constant. Thus
Em
cov E 10 He, 1/1] Hi] = cov Eun (Gr,Be. + +Ei +I''Wi W
Hg
= cov D 1itaaii-i110 W1th1 H1
Ej W Wicov (#1 Hj) WT Ej Eea
Et 1 W W - Er
where Σ., cov [t/1 H.) is the marginal posterior covariance of v/1 Finally, the total covariance decomposition [66, 33] yields that
Σn,a cov [0. Hy] E [cov [0, He, vh1 H., cov [0. H ril H..
= Σt, + Er Ej 1 W; E, W Σ E 't,a.
(33)
However, Σ,, = cov [ Hj] is different from Σg, = cov [eh He, e/2] that we already derived in Eq. (21). Thus we do not know the expression of E., But we can use the same total covariance decomposition trick find it. Precisely, let E.,e cov [vbe Hc] for any E [L]. Then have that
Σ.1 cov [4/1 He, e/2a (Σg + Graa)
F41 E [4/2 He, w/2a Σ,, Eg Wer BC.1)
First, given H,, cov [th He, thy = (Eng + G.,1) constant. Thus
E ccov [ He, 2/21 H.] cov (bb H,, 2/2a Er,
In addition, given H. E, Et,1 and B are constant. Thus
cov [E [ He, d22] Hj] cov [x,1 Σ Wer' + BB.1) Hi
cov [Er, Eg W22t2 E He]
Σy Wycov (e/2
Hy] W En Ex,1
Σx, W2 W.J 2n E.,
Finally, total covariance decomposition [66, 33] leads
Σ.1 cov [ H,} = E [cov [ He, 1221 Hj] + cov
E [14 He,
vi) Hy)
Σ., Σn, ΣS W2 S. W Σー Σ- E
Now using the techniques, this can be generalized using the same technique as above to
Σre Ee.t + Ee F+1 Wea Ee,411 W M Er,e
V8 E [L - 1].
30

---

Then, by induction, we get that
E
Σ EE(L)
E E [1 -011] 1];
where we use that by definition E.,,L cav fo H.) E,, and set P1 I E Wi+1 for any E [L]/[1}. Plugging this in Eq. (33) leads to
La and P
Ena -Er + Σ
W P.E, w E x E
1,a
rE(L]
- +
Σ
En 1
W P
(E[L]
Ei, +
Σ
P aft Pa,
EE[L]
where Pa.e Eq W,, En W
W1
Ii=
En
Wi+1-
F.4 Proof of lemma E.3
We prove this result by induction. We start with the base case when L = 1.
(I) Base case. Let a 18,2 A.A, X, From the expression of Σ,, in Eq. (21), we have that
2141,1 - - WT ET 1 (Eid, to t xx,x,) 'T (ET² Ex
EJ ES A, ET
WI,
=W7 [Σ (Ét.A. (L to 2x(x,) ")ET|) W1,
-W7 (xy 1, (Is - (Ta + - X,X, E
[ 1 ti xz WT,
wT (x- 1,A, (Id - (Ia + uu
)> 1.A, ΣT W1.
( wT
Σ-
uu T
E-1 W1.
A.A
TTu "L,,
X.X,
W
Eو A 1 u S. -
(34)
In (i) use the Sherman-Morrison formula. Note that (ii) says that Σ-1 Egp is one-rank which we will also need induction step. Now, we have that ||aell2 = 1. Therefore,
+ u Tu
2x, Ex,,, X, < 1 o - >1 XI, o -00 -
2 MAX
where we use that by definition of oA in Lemma F.3, we have that o 2 1 o 2o7 Therefore, by taking the inverse, we get that Th
-2 Combining this with Eq, (34) leads to
1T E 2
WT I SAA,, E. ET W1
MAX
Noticing that PA..I D. concludes the proof of the base case when f=1. 1. (II) Induction step. Let E [L](11} and suppose that 5-1 holds for - that
24116-1 Eiy is one-rank and that it
14116-1 -E 2 -0-20- 2 1(0-1)) A X
X.X, C...P.AA.... PA.C-1, where T
max 1 + -20 SE[L]
Then, we want to show that E441, Σ is also one-rank and that it holds that
211, - End 2 -0--0000 -0, TT.x X,X, X. PAL,
where a max
1 + 0 -00 LE
31

---

This is achieved as follows. First, we notice that by the induction hypothesis, we have that Σ-1 Gt,e-1 -1111-11 2,,-11 is one-rank. In addition, the matrix is positive semi-definite Thus we
141,6-1
write as Σ-1
2+1,6-1 G1,e-1=ua uu where Rd. Then, similarly the base case, we have
244,. 21, -Eit1, ---11
We (Ex + Ec+1,-11) - - W07 (Es + E1,-1)
W/ (// + E1+111111 -1 (Eg + c-,e-1) -1 We,
W, _
(=) + G1,--1) -1 (Ee" + E-1 E(111-1) - Wee
W; x
(ΣT)) + Gr,--)) -1 - (E + G E Gr D
W; 2, 1
(27 + G(,--1) -1 -1 - + GG +eu uu)) " We.
W, Σ
E ( su )u))
Eg'We.
uu
W, Σイ¹
[1,0-1 1 u E
E 1/ Wee
uu T
= We Et,e-1 1+ u 2
1 We
However, we follows from the induction hypothesis that wu = - -
- Gi,-1 24118-1
E = 0 -0 -(((-1) 22p--()PTTTT
I PA Therefore.
uu
(+1, Expt = W 23 1
+ u 1 E1 We,
2(18 /PT X P PA..-1
MAX
W/ 23 E1.011
111-1210 We. 1 Ee,e-1u
σ 20 MAX
1)
W 2g Dr PT X P. 1- We,
-We,
+ u Σ1,e-14
1(
MAX
PT P. XxX, PA.E
U Eq,e-14
Finally, use that 1 + u Er,c-iu < + (ull2>((Ereei))) & 1 0 207. Here use that ||ul| 5 o² which can also be proven by induction, and that X (Ve,e-1)) S oz. which follows from the expression of Σ in Appendix B.2. Therefore, we have that
20 2(--1)
E
Σ-1
+ u E1,--1u PA
PT
X,X,
2(6-1)
MAX
PT
PA X.X,XXX,PA PALe
202
2
TT
X,X, PALE
where the last inequality follows from the definition of o = maxte(z] 1 1 This concludes the proof.
F.5 Proof of theorem E.1
We start with the following standard result which we borrow from [32, 8].
(R(n) < 22n /22log(1/0),
Xx
53 Se
+ eno,
where is constant,
(35)
32

---

Then we use Lemma F.2 and express the marginal covariance Σr,A,
2-1
Er, = Dr,a + Σ
PartsePPo RE|L]
where Pa.e = -7
E Wi+t. (36)
Therefore, we can decompose
|X
Se,
2
as
|Xe E.At
=0
xx
D LL,,,
(i)
2 LY X1 to PA
o²
EEIL]
co (x)
X1)
Σ
cy lo 2 Pf A.)))).
(37)
FE(L]
where (i) follows from Eq. (36), and use the following inequality in (it)
x
log(1 + +2) <
max [, log(1 + x:) log(1 + 2) log(1 11 + u) log(1 + x))
log(1 + x)
which holds for any E [0, w/. where constants and are derived as
co of
2 24 log(1 +
cr
with the convention that 0441
log(1 + )
0²
The derivation of co uses that
X, E, S (E. S A7 ( + GLAY) S S =1 of
The derivation of eg follows from
X, P Ee X, Xt S (1 P. (A (A S 5o
Therefore, from Eq. (37) and Eq. (35), we get that
(R(n)=
2n log( 1/8)
co
log(1 + to (x.
o - P E ex, 1 + end
log(1
(38)
EEIL
t=1
Now we focus on bounding the logarithmic terms in Eq. (38).
(I) First term in Eq. (38) We first rewrite this term
log(1 + - Xx. E, X.) log det (Ta + o 2 L.
X,X, Σ3 17,,),
let ( o - xx,,]) - log det ( log det ($-1 (E log det (f (
where (i) follows from the Weinstein-Aronsene identity. Then get a telescoping
over all rounds t [mi]. and
n log det [5 (ETH.A.)) det
log det La + -5 r.A,
T
X.X,
$ t.A.
K
K
log ded (eg log det (5rr.)
log det dedd(((( (2tTt11)))) - log det (D),)),
1=1 a=1
a
K
Σ
log det ( (,,1, - log det(
(
A-1 log det det
n11,a E/).
33

---

where (i) follows from the fact that Σi, E1. geometric means and get
Now we use the inequality of arithmetic and
K
log det (12 to -0 XxX,
X
log det
det((tt
S S-1
Ei),
1 1--1 (
d
dlog
Tr22 d
(39)
K
o
n of [+
dlog
1
Kdlog
d
o²
d d²
(II) Remaining terms in Eq. (38) Let E [L]. Then we have that
log(1 10 2x, PA Er, .P A,, (x))) = 0 2
O 22 MAX log(1 to Xx P Prt A.x)t)
-2,
S MAx lo 1 +
2 MAX X, PA. E P P X(),
log (ea + 30
28 5
P Xix, P (.
MAX
MAX i
o 26
dot dot +
PP X PA.,e) log dot(E,))) MAX
where we use the Beinstein-Aronspajj identity in (i). Now we know from Lemma F33 that the following inequality holds o -2 2PP
22 PT
A,, XxX, PAde I 241, - As a result, we get that 5-11 2411. Thus,
E +o a 2 Raxxx A., XxX, P ALd
log(1 t P PA Su P X1) S oMix (logdet - log det det (5T,)))
Then we sum over all rounds t E [n], and get a telescoping
log + to PA A Xi) S o MAX
log deedeetTTTT log del (ET)),
log det doggeetSnnt.d - log det (2ij))
σ MAX
28 MAX
(tt,,)) log det (Ez
5-1 log S sxt1ee=211))
MAX
where we use that Σn, = Eet1 in (1). Finally, we use the inequality of arithmetic and geometric means and get that
Σ log(1 to -2, PA (.,.P.)))))) s 22 (log E-- ( [
L < do max log
2 2 (40)
o2 oz
s do MAx log (+
The last inequality follows from the expression of Σ,1 1,,,e in Eq. (21) that leads
Σ3 1 La + + 14
= +a++ - W (2, - Ei 1- "E "
(41)
34

---

since Ge WT (5rg² - Ex 1-11,,,-1 - We, This allows us to bound Tr(E 5-1 1 as و
ä T المر الشل = d
Tr(la + 2 WT - - Wo
(d Th W (E) E E WW
s1+
1 w (2, -
111-121 W
d
S1+
X1
(WW - 20
S1+
(
d
2+1
o2 of
$ +
1
(42)
d
of
im
where the assumption that A1 (W/ We) 1 (A2) and that X1 (Ee+1) o and Y1 (E,
1/03- This is because E, of la for any l [T+ 1). Finally, plugging Eqs. (39) and (40) in Eq. (38)
concludes the proof.
F.6 Proof of proposition E.2
We use exactly the same proof in Appendix F.5, with one change account for the sparsity assumption (A3). The change corresponds to Eq. (40). First, recall that Eq. (40) writes
log(1 + t xx, PA. P S xax (logdet(22+)))) Smt1.
where
24 D n
La + 2 WT (2, - We Ia + o W (Σp² -E Eイ 1)
2011
(43)
where the second equality follows from the assumption that Ee+1 = - Id- But notice that in assumption, (A3), assume that We (We, Od,d-d;), where We dx for any e [4]. Therefore, have that for any d x matrix B Radxd, the following holds, W.BWe
W. W.TNV BW, Od ) In particular, we have that Od-de,de Od-dedd-de)
WT (E+²
ET Wt
Ode )
(E,1 --' E,
(44)
Os-di
od-de
Therefore, plugging this in Eq. (43) yields that
x =
(Th
Id, + t WT (5-1 --0-------0 - --1
Od Ta-de
(45)
Od-de:de
As result, del dEt(E
5-1
x nnt1e ( = del(/a, + + W (Σr -
We). This allows
us to move the problem from a d-dimensional one dyddimensional one. Then we the inequality
35

---

of arithmetic and geometric means and get that
log(1 + 2x, PA ee e., (x)) MAX det(( Dna ((tz)))
Max log det(Ld, + 0241 We -2 We),
S dio log de
Tr(la, + +o W. ( = ST'E1 E1,6-110 We))
l+41 of
S dio log
1+
(46)
To get the last inequality, use derivations similar the ones we used in Eq. (42). Finally, the desired result in obtained by replacing Eq. (40) by Eq. (46) in the previous proof in Appendix F.5.
G Additional experiments
G.1
Swiss roll data
Fig. 4 shows samples from the Swiss roll data and samples from generated by the pre-trained diffusion model for different pre-training sample sizes.
2
*** bola Diffusion model
distribution 2050 Diffusion model
distribution X/x Diffusion model
-2
-2
-3
(a) Diffusson traitre 50 sam- ples from the Swiss roll dataset.
(b) Diffusion pre-trained on 1 (c) Diffusion pre-trained on 10 samples from the Swiss roll dataset. samples from the Swiss roll dataset.
Figure 4: True distribution of action parameters (blue) vs. distribution of pre-trained diffusion model (red).
G.2 Diffuusion models pre-training
We used JAX for diffusion model pre-training, summarized as follows:
Functions fo are parameterized with fully connected 2-layer neural network (NN) with ReLU activation. The step is provided as input to capture the current sampling stage. Covariances are fixed (not learned) as Σg ofIa with σ increasing with É.
Parameterization:
Offline data samples are progressively noised over steps f E [L]. creating increasingly noisy versions of the data following predefined noise schedule [30]. The NN is trained to reverse this noise (i.e., denoise) by predicting the noise added at each step. The loss function measures the L2 norm difference between the predicted and actual noise at each step, as explained Ho al. [30].
Loss:
Optinization: Adam optimizer with 10-3 learning rate was used. The NN was trained for 20,000 epochs with a batch size min pre-training sample size). We used CPUs for pre-training, which was efficient enough to conduct multiple ablation studies.
A After pre-training: The pre-trained diffusion model is used as prior dTS and compared to LinTS as the reference baseline. In our ablation study, we plot the cumulative regret of LinTS in the last round divided by that of aTS. ratio greater than indicates that dTS outperforms LinTS, with higher values representing larger performance gap.
36

---

G.3 Quality of our posterior approximation
To assess the quality of our posterior approximation, we consider the scenario where the true distribution of action parameters is N (0a. Id) with d = 2 and rewards are linear. We pre-train diffusion model using samples drawn from N (0g, Ia). We then consider two priors: the true prior N (Oa, Ia) and the pre-trained diffusion model prior. This yields two posteriors:
Pi Uses N (0d, Ia) as the prior. P1 is exact posterior since the prior is Gaussian and rewards are Gaussian.
P2 Uses the pre-trained diffusion model as the prior. Pg our approximate posterior.
The learned diffusion model prior matches the true Gaussian prior (as seen in Fig. 5a). Thus, if our approximation accurate, their posteriors Pi and P2 should also similar. This is observed in Fig. 5b where the approximate posterior P2 nearly matches the exact posterior Pt.
-0.30
Gaussian Prior
-0.32
Exact Posterior (P))
xXx
Diffusion model Prior
xxx
Approximate Posterior (Pt)
0.36
0.38
0.40
-2
0.42
-3
-0.44
-6
-2
6
0.1
0.0
0.2
(a) Gaussian distribution vs. diffusion model pre- trained 10ah samples drawn from it.
Exact posterior Pi approximate posterior after 100 rounds of interactions.
Figure 5: Assessing the quality of our posterior approximation.
Empirical posterior validation with MCMC. Above, we assessed posterior approximation in tractable Ainar-Gaussian case, showing that the diffusion-basee posterior closely matches the exact posterior. To further validate approximation in more complex, non-linear settings, compare our diffusion Thompson sampling (dTS) posterior samples to those obtained via two MCMC variants on the non linear MovieLens benchmark:
MCMC-Fast: Uses fewer sampling steps for efficiency.
MCMC-Slow: Uses more sampling steps for higher accuracy.
As shown Table even the high-compute MCMC variant yields higher regret than dTS, motivating efficient approximation for online bandits.
Table 2: Comparison between dTS and MCMC-based posteriors on MovieLens.
Baseline
Regret Improvement (%)
Time Speed-Up (%)
dTS MCMC-Fast dTS vs. MCMC-Slov
50.6 %
47.6 %
12.7
80.5
G.4 CIFAR-I0 ablation study
CIFAR-10. In Fig. 3a in Section 5.2, we showed that with only 10 pre-training samples, dTS outperforms LinTS on Swiss-roll benchmark. We now extend this analysis to the vision dataset CIFAR-10 10 [40] (similar results were obtained on MNIST [48]). Our setting is similar to that in Hong et al. [32] and we use dTS's variant that uses single shared parameter R& (Remark 2.1 and Section 3.3) because is more suited for this setting. These additional ablations CIFAR-10
37

---

confirm that dTS consistently benefits from offline pre-training, even when the true prior is not diffusion model. Specifically, we vary the percentage of offline data used train the prior and compare against both HierTS and LinTS.
Table Regret improvement (%) of dTS CIFAR-1O.
Offiin Data (%)
HierTS
LinTS
1%
69.11%
87.74%
5%
79.56%
92.18%
25%
80.65%
92.48%
50%
81.67%
92.88%
G.5 Bound comparison
Here, compare our bound in Theorem E.1 to bounds of LinTS from the literature.
Linear Bandit: d-5, k-10000, L-5
Linear Bandit: d-5, K- 10000 L-5
600000
90000
Bayes Bound for dTS (Ours) Frequentist Bound LinTS
80000
Bayes Bound drS (Gurs)
500000
Bayes Bound
400000
60000
Regret
50000
300000
40000
200000
30000
100000
10000
200
400
600
800
1000
200
400
600
800
1000
Horizon
Horizon
(a) Comparing our bound the frequentist bound of LinTS in Abeille and Lazaric [2].
(b) Comparing our bound the standard Bayesian bound of LinTS.
Figure 6: Comparing our bound to the frequentist and Bayesian bounds of LinTS.
H Extensions
Theory beyond iinar-Gaussian Extending analysis nonlinear settings is nontrivial. A promising direction is an information-theoretiic analysis of Bayesian regret with structured priors [59, 52] or PAC-Bayesian treatment similar to that used in offline contextual bandits [9]. Closing the gap to lower bounds remains open: the only known Bayesian lower bound applies to K-armed bandits [45], while minimax frequentist bounds scale as ((dv/)) [21].
K K-independent regret. In the setting of Remark 2.1, where r(ı, a; 0) = p(æ, a) and is shared across actions, our proof techniques imply K-independent regret once & known or accurately estimated. This connects to structured large-action-spart results where regret does not scale with K [25, 67, 70]. Exploring further the use of diffusion models such setting is promising.
Robustness and misspecification. dTS may face misspecification at both the prior and likelihood levels. When the diffusion prior is biased or trained limited data, it remains empirically stable but lacks robustness guarantees. Moreover, dTS assumes a generalized linear reward model; deviations from this assumption leads to model misspecification.
Beyond contextual bandits. The posterior derivations dTS extend naturally to other settings. For instance, in off-policy learning, since dTS defines tractable posterior over reward parameters, it can be combined with off- off-policy estimators and policy improvement methods in structured offline contextual bandits environments [10].
Online fine-tuning and offline RL. Pre-training diffusion model on offline data and refining online via dTS amounts to diffusion fine-tuning from implicit bandit feedback. Extending this to sequential decision-min makiin with dynamics aligns with recent liffusiion-fordeciinee work. concrete next step is to use dTS for fine-tuning pre-trained diffusion models collected reward data.
38

---

I
Broader impact
This work contributes the development and analysis of practical algorithms for online learning to act under uncertainty. While our generic setting and algorithms have broad potential applications, the specific downstream social impacts are inherently dependent on the chosen application domain. Nevertheless, we acknowledge the crucial need to consider potential biases that may be present in pre-trained diffusion models, given that our method relies on them.
Amount of computation required
Our experiments conducted internal machines with 30 CPUs and thus they required moder- ate amount of computation. These experiments are also reproducible with minimal computational resources.
39


Reviews:
[{"review_id": "ykJwgKPWAW", "reviewer": "Reviewer_taUg", "summary": "This paper leverages diffusion models to tackle the challenge of decision-making in contextual bandits with large and complex action spaces. The authors propose a framework that uses pre-trained diffusion models as expressive priors within a Thompson Sampling (TS)-style bandit algorithm. Their proposed method enables efficient posterior approximation and sampling by combining offline-trained diffusion models with online updates during bandit online interactions. The authors develop practical implementations for posterior inference under this setup, and conduct experiments to demonstrate strong empirical performance across a range of contextual bandit scenarios. The authors provide theoretical analysis for the linear case, where the regret of the proposed method is bounded.", "questions": "1. How can the proposed method extend to more general nonlinear settings?\n2. How can the proposed method extend to the reinforcement learning settings?\n3. Your work proposes a TS-type algorithm. Do you think the idea of leveraging diffusion models to enhance contextual bandits can also apply to the UCB-type algorithmic approach?", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. This paper studies an important problem of contextual bandits. To the best of my knowledge, this is the first work to leverage diffusion models to enhance contextual bandits. The paper is well-motivated.\n2. The paper is well-written and easy to follow.\n2. The authors provide theoretical guarantees of the proposed method.\n3. Extensive experiments are conducted to show the outperformance of the proposed methods.\n\nWeaknesses:\n1. This paper considers the linear case, which is a slightly restrictive assumption in practice, but it is reasonable in the bandit literature.\n2. There are many works on diffusion models for reinforcement learning, though mostly from the empirical perspective. It would be better to discuss these works and make some comparisons. For example, see the related works in the survey \"Diffusion Models for Reinforcement Learning: A Survey\".\n3. The authors conducted extensive simulations, but they are only on synthetic and MovieLens datasets. It would be better to conduct experiments on more datasets to show the consistent improvement of the proposed methods over existing ones.", "was_revised": true, "final_justification": "The authors have addressed my concerns, and I believe this is a good work with solid contributions", "rebuttal": "We thank Reviewer **taUg** for your thoughtful review and your positive remarks on our motivation, clarity, theoretical contributions, and empirical results. Below, we summarize the main updates made (in this rebuttal only, as external links are disallowed) and respond to your comments point-by-point.\n\n---\n\n## Summary of Key Revisions (Version 2)\n\n1. **Expanded Related Work**\n   We added discussion contrasting our *online contextual bandit* setting with recent *diffusion-RL* works and cite the 2024 survey *Diffusion Models for Reinforcement Learning* along with key references therein.\n\n2. **New Experiments on MNIST & CIFAR**\n   Added standard supervised-to-bandit benchmarks. With just **5 %** offline data for pretraining, dTS yields:\n\n   | Dataset | vs. HierTS | vs. LinTS |\n   | ------: | ---------: | --------: |\n   |   MNIST |    42.17 % |   82.30 % |\n   |   CIFAR |    79.56 % |   92.18 % |\n\n3. **Cold-Start Ablations**\n   We include ablations across 1–50 % offline data. dTS consistently outperforms all baselines:\n\n   | Offline Data (%) | vs. HierTS | vs. LinTS |\n   | ---------------: | ---------: | --------: |\n   |              1 % |    69.11 % |   87.74 % |\n   |              5 % |    79.56 % |   92.18 % |\n   |             25 % |    80.65 % |   92.48 % |\n   |             50 % |    81.67 % |   92.88 % |\n\n4. **Clarified Nonlinear Reward Support**\n   dTS supports any twice-differentiable reward function via Laplace approximation. The linearity assumption is used solely for deriving theoretical regret bounds.\n\n5. **Limitations Moved from Appendix to Main Text**\n   The limitations section has been elevated from the appendix for improved visibility.\n\n---\n\n## Point-by-Point Responses\n\n### 1. Linear-Case Theory\n\n> **“The linear assumption is slightly restrictive.”**\n\nOur theoretical analysis is limited to the linear setting (common in bandit theory for tractability), but the algorithm itself supports non-linear rewards and non-linear diffusion models via our approximation. Our empirical results on high-dimensional tasks (e.g., CIFAR) confirm strong performance.\n\n---\n\n### 2. Diffusion-RL Discussion\n\n> **“There are many works on diffusion models for RL.”**\n\nWe cite relevant papers including the 2024 survey. Our setting is distinct in that we focus on *exploration in online contextual bandit* via posterior sampling, not offline policy optimization.\n\n---\n\n### 3. Dataset Diversity\n\n> **“It would be better to conduct experiments on more datasets.”**\n\nWe agree and added experiments on **MNIST** and **CIFAR**, which are common benchmarks in contextual bandits. These show large and consistent improvements over LinTS and HierTS, even with minimal offline data. See results in the Summary of Key Revisions section above.\n\n---\n\n### 4. Q1 – Extension to Nonlinear Settings\n\n> **“How can the method extend to more general non-linear settings?”**\n\nThis is already supported: our approximation handles any non-linear twice-differentiable reward (e.g., logistic) and any non-linear diffusion model. We clarified this and expanded examples in App C.\n\n---\n\n## Extensions\n\n__RL Extension.__\n\nWe agree this is a promising direction. A natural extension is to use dTS for **model-based RL**, where policy or transition dynamics parameters are drawn from a diffusion prior and updated online via posterior sampling. We look forward to exploring this idea further.\n\n__UCB-Type Extension.__\n\nIn Appendix J, we added a sketch for **BayesUCB-style variant**, where the selected action maximizes the posterior mean plus a bonus term based on posterior uncertainty (computed using the posterior covariance). We believe this could yield promising performance and plan to explore this direction in future work.\n\n---\n\nWe hope these additions, more datasets, deeper context, and clarification of practical extensions, address your concerns and strengthen your confidence in our submission. We would greatly appreciate your consideration for raising your support of our paper."}, {"review_id": "Ih5K8JAkv3", "reviewer": "Reviewer_oesN", "summary": "* This paper proposes a new contextual bandit algorithm that extends the literature on hierarchical Bayesian bandits to a more complex structure where the prior distribution is represented as a pre-trained diffusion model.\n\n* This paper considers the contextual bandit setting with prior represented as a diffusion model, and adopts Bayes regret to measure the performance. Then, it proposes a novel algorithm - Diffusion contextual Thompson sampling (DTS).\n    * The first key novelty is deriving an approximation for the posterior update. This approximation has good properties: (1) they match the prior with no data, and the prior's effect diminishes as data accumulates. (2) The approximate posterior is also a diffusion model with updated means and covariances.\n    * The second key contribution is to theoretically analyze the Bayes regret bound in the simplified setting where the known diffusion model's ODE is time-varying linear–Gaussian. The challenge of analysis the recursive update in the diffusion model. The analysis demonstrates the computational and statistical benefits of the proposed method, in contrast to LinTS [54]. The key insight is that the proposed method leverages the correlation between actions through the pre-trained diffusion model, while LinTS [54] has to learn each action independently.\n\n* Empirical analysis contains two experiments:\n  * When the prior matches the generative process, the paper compares the proposed method against LinUCB [1], LinTS [3], and HierTS [31] for the linear-reward setting, and against UCB-GLM [44] and GLM-TS [16] for the nonlinear-reward setting. The result shows that the proposed method achieves lower regret than the baselines. It demonstrates the power of a well-specified prior.\n  * When the prior does not match the generative process, such as Swiss-roll data and MovieLens data, the proposed method also outperformed LinTS.", "questions": "* In Sec.4.1, the paper compares the proposed method vs LinTS [3]. I am not very familiar with [3], but it seems that the paper [3] (https://arxiv.org/pdf/1209.3352) 's Sec.2.1 assumes that there is a common reward parameter µ, which seems different from this paper's setting where each action `i' has its own parameter `θ_{*,i}`. I am curious whether you think it would still be appropriate to compare the proposed method with [3]?\n* Another question I have is that Fig. 2's nonlinear diffusion linear reward case, why not also benchmark the HierTS [31] method?\n* Out of curiosity, this paper focuses on correlations between actions. What about also exploiting correlations between contexts? E.g., similar users would have similar preferences about movies. Would the proposed method also extend to handle that?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "* This paper solves a well-motivated problem and is well-written.\n* The algorithm is novel and sound. Both theoretical and empirical result supports the method.\n* One concern is that it seems that the result section focuses on comparing the proposed method vs LinTS (a baseline that does not exploit the correlation between action features). This is very effective in showing that the proposed method can exploit the action correlation, and it is actually effective and useful. However, it would be even stronger to also demonstrate that the proposed method is better than other hierarchical Bayesian bandit methods, to demonstrate that representing the prior as a diffusion model is a useful, not an overkill, in real-world scenarios. In this way, the two key takeaway messages are proved: (1) exploiting action correlation is useful (already somewhat demonstrated in prior literature), (2) representing prior distribution as a diffusion model is useful (very novel).\n    * To show this, I think there are a couple of ways. For example, for the real-world dataset - movielens, it would be helpful to benchmark the HierTS [31] method, to see how much modeling the prior as diffusion is beneficial in the real world. Alternatively, the author could take the setting of HierTS [31] and show that the diffusion-model prior allows better learning performance because the diffusion model is more expressive.", "was_revised": true, "final_justification": "This paper is well motivated, with theoretical and promising empirical results. I think it extends contextual bandit to be able to bake richer information into the action prior. The rebuttal clarified my questions. I would keep my score.", "rebuttal": "We thank Reviewer **oesN** for the thorough review and for highlighting our paper’s clear motivation, sound algorithmic design, and strong theoretical and empirical results. Since NeurIPS disallows updated PDFs or external links, all revisions are detailed below.\n\n---\n\n## Summary of Key Revisions\n\n1. **Added HierTS \\[31] Baseline.**\n   We now benchmark HierTS on MovieLens and Swiss-roll. On MovieLens, HierTS incurs ≈ 45.5 % higher cumulative regret than dTS, with dTS > HierTS > LinTS in terms of performance.\n\n2. **Expanded Nonlinear-Reward Evaluation.**\n   Fig. 2’s nonlinear-diffusion, linear-reward experiments now include HierTS, where dTS still outperforms it.\n\n3. **Clarified LinTS \\[3] Implementation.** We run LinTS per action (one Gaussian posterior per \\$\\theta\\_{\\*,i}\\$), matching our setting. We cite \\[3, Sec. 2.1] and per-arm TS variants. Joint-parameter variants would require modeling the full \\$Kd\\$-dimensional posterior, which is computationally prohibitive.\n\n4. **Context-Correlation as Future Work.**\n   We agree that modeling correlations between contexts (e.g., similar users) is valuable. While this paper focuses on action correlations, extending diffusion priors to contexts is an exciting direction for future research.\n\n---\n\n## Point-by-Point Responses\n\n**1. Comparison to other hierarchical Bayesian bandits (HierTS \\[31])**\nReviewer: *How does dTS compare to HierTS in MovieLens?*\nDespite using the same offline data, HierTS’s Gaussian hyperprior yields up to **45.5 %** higher regret than dTS’s diffusion prior in MovieLens, demonstrating that diffusion models capture richer action correlations. We also conducted additional MNIST and CIFAR benchmarks. With just **5 %** offline data, dTS outperforms both HierTS and LinTS:\n\n| Dataset | vs. HierTS | vs. LinTS |\n| :-----: | ---------: | --------: |\n|  MNIST  |     42.2 % |    82.3 % |\n|  CIFAR  |     79.6 % |    92.2 % |\n\n**2. Exploiting correlations between contexts**\nReviewer: *Can dTS handle context correlations (e.g., user similarity)?*\nWe agree this is an important extension. Prior work (e.g., clustering-based Thompson sampling) has explored simple context correlations, but diffusion-based context priors remain unexplored. We plan to investigate this direction in future work.\n\n---\n\nWe believe these updates address your concerns and further demonstrate the practical advantages of diffusion-model priors. Thank you again for your valuable feedback; we welcome any additional comments or suggestions."}, {"review_id": "mDjQrOwwYC", "reviewer": "Reviewer_uoU2", "summary": "This paper introduces an approach to contextual bandits by leveraging pre-trained diffusion models as priors to capture complex correlations between actions, which is particularly useful when dealing with large action spaces. The authors develop diffusion Thompson sampling (dTS), a computationally efficient algorithm that approximates posterior distributions under diffusion priors through a hierarchical sampling scheme, enabling better exploration by sharing information across related actions. They provide theoretical analysis showing improved regret bounds in the linear case and demonstrate through experiments that dTS significantly outperforms standard baselines like LinUCB and LinTS, especially as the number of actions increases, while maintaining computational efficiency even with model misspecification.", "questions": "- Could you provide either theoretical analysis or of your two-layer approximation scheme? Specifically, I'd like to understand: (a) under what conditions the Gaussian approximation of the GLM likelihood is accurate, (b) how errors propagate through the hierarchical structure when using nonlinear link functions?\n\n- The paper claims computational advantages, but doesn't provide wall-clock time comparisons. Could you benchmark the actual runtime of dTS versus LinTS on problems with varying K and L? \n\n- The algorithm seems to require good offline parameter estimates to pre-train the diffusion model. How does dTS perform in a \"cold start\" scenario with minimal or no offline data?", "limitations": "- The significant gap between theoretical guarantees (linear case only) and practical algorithm (nonlinear approximations) should be prominently discussed.\n\n- The computational overhead for real-time decision-making systems needs acknowledgment.\n\n- The lack of formal guarantees on approximation quality and failure modes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The authors develop a two-stage approximation that first handles the nonlinear reward distribution by approximating the likelihood with a Gaussian centered at the maximum likelihood estimate, then addresses nonlinear link functions by borrowing the covariance update structure from the exact linear case - this design ensures the approximate posteriors reduce to exact ones in linear settings while remaining computationally tractable for general nonlinear diffusion models.\n- The algorithm exploits the Markov structure of the diffusion model to avoid maintaining joint distributions over all action parameters, instead decomposing the problem into separate posterior updates for each latent layer and action - this reduces the computational burden from cubic in the total number of action parameters to linear in the number of actions plus layers.\n\nWeaknesses:\n\n- The theoretical analysis only covers linear diffusion models with linear-Gaussian rewards, yet the main algorithmic contribution involves two layers of approximations for nonlinear cases.\n\n- The approach relies on pre-training diffusion models on accurate offline estimates of action parameters, but obtaining such estimates for thousands of actions is precisely the challenging problem that online learning aims to solve - the paper doesn't address how to handle settings where offline data is biased, sparse, or completely unavailable for certain actions.\n\n- The experimental section shows performance degrades significantly when the true generative process doesn't match the diffusion prior, yet there's no theoretical characterization of robustness - the approximation scheme could amplify model misspecification errors through the hierarchical structure, but the paper doesn't analyze worst-case performance or provide guidance on detecting when the diffusion prior is misleading the algorithm.", "was_revised": true, "final_justification": "The author has addressed my questions; therefore, I increase my score accordingly.", "rebuttal": "We thank Reviewer **uoU2** for the careful review and for emphasizing our work’s strengths in novelty, clarity, theoretical insights, and empirical performance. Since NeurIPS does not support updated PDFs or external links, all changes are detailed directly below.\n\n---\n\n## Summary of Key Revisions\n\n1. **Wall-Clock Runtime Benchmarks.**\n   We added end-to-end timing comparisons demonstrating that dTS scales linearly in both K (number of actions) and L (diffusion depth), yielding orders-of-magnitude speedups over joint-parameter LinTS.\n\n2. **Cold-Start Robustness Experiments.**\n   We introduced new ablations on CIFAR (and other benchmarks), varying the fraction of offline data used to pre-train the diffusion prior from 1 % to 50 %, and show that even with just **1 %** pre-training samples, dTS consistently outperforms LinTS and HierTS.\n\n3. **Clarified Theory–Practice Connection.**\n   We expanded our discussion of how dTS’s posterior approximation is directly inspired by the fully linear case, and we show strong empirical evidence supporting this approximation across diverse settings.\n\n---\n\n## Detailed Responses**\n\n### 1. Theoretical Analysis of the Two-Layer Approximation\n\n> **“Could you provide theoretical analysis of the Gaussian approximation and error propagation?”**\n\nFormal proofs for nonlinear link functions exceed this paper’s scope. Instead, we support our approximation with:\n\n* **Empirical Posterior Validation.** In Appendix F, we compare dTS samples to the exact posterior in a tractable linear-Gaussian case. Results show that dTS's approximation closely matches the exact posterior. This is the only case where we have access to an exact posterior to which we can compare our approximation. To futher support our approximation, we compared to two MCMC variants in the non-linear MovieLens experiment:\n\n  |          Baseline | Regret Improvement (%) | Time Speed-Up (%) |\n  | ----------------: | ---------------------: | ----------------: |\n  | dTS vs. MCMC-Fast |                 50.6 % |            47.6 % |\n  | dTS vs. MCMC-Slow |                 12.7 % |            80.5 % |\n\n  Even the high-compute MCMC variant yields higher regret than dTS, motivating our efficient approximation for online bandits.\n\n* **Strong Empirical Performance.** We added MNIST and CIFAR benchmarks (common practice datasets). With just **5 %** offline data, dTS outperforms HierTS and LinTS:\n\n  | Dataset | vs. HierTS | vs. LinTS |\n  | ------: | ---------: | --------: |\n  |   MNIST |     42.2 % |    82.3 % |\n  |   CIFAR |     79.6 % |    92.2 % |\n\n* **Design-Principle Alignment.** Our approximation reduces to the exact posterior when the diffusion model is linear and the reward model is Gaussian. As more data accrue, the prior’s influence naturally diminishes; this behavior matches the exact posterior distribution behavior, and hence we believe it is a strong point of our approximation.\n\n### 2. Role of Linear-Case Guarantees\n\n> **“How do linear-case guarantees inform the nonlinear algorithm?”**\n\nOur linear analysis inspired our approximation for two reasons:\n\n1. **Exact-Case Posterior Approximation.** dTS matches the linear-case posterior when models are linear. Thus, the linear-case posterior fully inspired our posterior approximation in the non-linear case.\n2. **Regret-Scaling Alignment.** Once the prior captures the true distribution, dTS’s regret scaling with L matches the linear-case scaling predicted by our theory (Figure 4 (b)).\n\n### 3. Wall-Clock Runtime Comparisons\n\n> **“Please benchmark actual runtimes of dTS vs. LinTS.”**\n\nWe measured average runtimes over 100 GPU runs:\n\n* **Varying K (L=2, d=5):**\n\n  |        Method |   K=200 |     400 |      800 |      1600 |\n  | ------------: | ------: | ------: | -------: | --------: |\n  |       **dTS** |  1.95 s |  3.78 s |   7.39 s |   14.82 s |\n  | LinTS (joint) | 17.13 s | 98.88 s | 611.61 s | 3834.61 s |\n\n* **Varying L (K=200, d=5):**\n\n  |        Method |     L=4 |       8 |      16 |      32 |\n  | ------------: | ------: | ------: | ------: | ------: |\n  |       **dTS** |  3.73 s |  3.78 s |  3.92 s |  4.05 s |\n  | LinTS (joint) | 95.05 s | 95.73 s | 96.08 s | 95.84 s |\n\nThese results confirm dTS’s **linear** scaling versus LinTS (joint) **cubic** cost, making dTS more practical at large scales.\n\n### 4. Cold-Start Performance with Minimal Offline Data\n\n> **“How does dTS perform with sparse or no offline data?”**\n\nIn Sec 5.4 and Fig 4(a), we show that with only **10 pre-training samples**, dTS outperforms LinTS on the Swiss-roll benchmark. Our additional CIFAR ablations yield:\n\n| Offline Data (%) | vs. HierTS | vs. LinTS |\n| ---------------: | ---------: | --------: |\n|              1 % |     69.1 % |    87.7 % |\n|              5 % |     79.6 % |    92.2 % |\n|             25 % |     80.7 % |    92.5 % |\n|             50 % |     81.7 % |    92.9 % |\n\nThese results confirm that dTS **effectively leverages limited prior data** to capture action correlations and achieve substantial regret reductions.\n\n---\n\nWe believe these revisions significantly strengthen our theoretical grounding, empirical validation, and practical utility. If these changes address your concerns, we kindly request your renewed support for our submission."}, {"review_id": "phnKaWWCRv", "reviewer": "Reviewer_8Fbo", "summary": "The authors consider contextual Thompson sampling for a diffusion reward model. The reward model consists of a (diffusion) hierarchy of latent parameters, each a multivariate normal conditioned on the previous latent with mean given through a (non-linear) link function of the previous latent, followed by a head that independently maps to an action vector for each of the arms of the bandit. The reward is assumed to be a function of the alignment between a provided context vector and the unknown action vectors of each arm, following a generalized linear model. \n\nThe main contribution of this work is a tractable approximation of the posterior action vectors (and latents) in case the link function is non-linear (otherwise the diffusion model marginalizes to a Gaussian). The authors adopt the Laplace approximation for the generalized linear model and approximate the link functions as linear for the posterior update.", "questions": "- Concerning notation, I would suggest to only add an asterisk for the optimal action selection / policy, and a special font for random variables (or upper-case lower-case notation). Moreover, the paper would be much easier to parse if instead of defining quantities like $\\mathcal L_l(\\psi_l) := \\mathbb P(H_t | \\psi_{*,l} = \\psi_l)$ the authors could just stick to writing the right-hand side expression in the formulas. Same with $Q_{t,l}$. It is barely any longer, easily fits in the space provided, and results in immediately recognizable posterior update formulas. Also, $P$ is used to denote two distinct probability distributions.\n\n- The authors adopt a latent parameter $\\psi_{1}$ through which correlation of the action vectors $\\theta_{I}$ is defined. It could greatly help the reader if upon introduction of the latent parameters it was explicitly stated that the action vectors are independent conditioned on the latent.\n\n- The histories $H_t$ are taken as a tuple with index $k \\in [t-1]$ which is an unfortunate and confusing choice, because $K$ already denotes the number of arms of the contextual bandit. Instead of using $I \\in [K]$ and $k \\in [t-1]$, I would suggest the authors to consistently use lower and upper case, respectively, to denote matching ordinal and cardinal quantities. \n\n- $\\mathcal R^{ACT}(n)$ and $\\mathcal R^{LAT}_l$ are discussed but nowhere defined in the Bayesian regret bound.\n\n- The Hessian in Equation (3) has a typo, $g$ should have a “double-dot” to indicate the second order derivative.\n\n- Under the linearization assumption in the theoretical analysis dTS becomes exact Thompson sampling and the prior action vectors simply become Gaussian. Doesn’t this suffice to immediately deduce regret bounds by borrowing from the extended analysis done on GP-TS? It would be great to outline the strategy for deriving the regret bounds, even outside the Appendix.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The paper conducts both theoretical and experimental analysis for their method. The theoretical analysis is based on assuming linear link functions, where the adopted approximation becomes exact. There they inherit the cumulative regret bounds known for Thompson sampling. \n- The experimental analysis is based on appropriate statistical metrics, including significance. On the Swiss-roll dataset and MovieLens, the authors demonstrate benefits of using a non-uniform prior through pre-training of the diffusion model. \n- Moreover, the authors provide insights on a bias-variance trade off in the number of layers: shallow diffusion models are not expressive enough for the task at hand, whereas deep diffusion models result in excessive capacity and slow exploration.\n\nWeaknesses:\n- From the main text it is not entirely clear whether the theoretical analysis just amounts to computing the covariance matrix of the action space and plugging into Thompson sampling bounds from the standard literature, or whether a more involved analysis takes place. \n\n- Across the paper, the adopted notation is questionable, making the paper unnecessarily difficult to read. For instance, random variables are denoted with subscript asterisks all over the place, and likelihood functions are introduced with specific notation never to be used again. See questions for suggestions on the notation.\n\n- The experimental analysis is limited though promising. The results from Section 5.1 are entirely expected with the benchmark designed according to their method, and as far as I can tell mainly argue that in a calibrated setting adopting dTS is better than assuming independent arms. Also, the paper may benefit significantly from additional experiments in the style of MovieLens, mainly to convince practitioners that pre-training Diffusion models prior to playing bandits can give significantly better cumulative regret in the online setting.\n\n- Since the paper introduces an efficient albeit biased approximation to posteriors in diffusion models, a significant missing piece in the experimental analysis is comparison against unbiased Monte Carlo integration. In particular, a plot with compute resources on the x-axis and regret on the y-axis for different integration budgets of Monte Carlo as well as dTS is absolutely vital for selling the method, assuming dTS is indeed pareto-optimal. This is a crucial missing point, vitally affecting the usefulness of the paper.", "was_revised": true, "final_justification": "The authors addressed the initially unclear notation and added further experiments, in particular, an absolutely essential accuracy-compute comparison against MCMC, which demonstrates that their method is pareto-optimal. Thus I improved my clarity assessment and the final score.", "rebuttal": "We thank Reviewer **8Fbo** for their detailed and constructive feedback. We appreciate the recognition of our theoretical and empirical contributions and your suggestions for improving clarity and notation. Below, we summarize our main updates in **Version 2** (updated PDFs or external links are not supported in NeurIPS rebuttals) and address each comment in turn.\n\n---\n\n## Main Revisions in Version 2\n\n1. **Expanded Experiments.**\n\n   * **MCMC vs. dTS:** Added direct comparisons showing that dTS achieves substantially lower regret and faster convergence than both “fast” (low-compute) and “slow” (high-compute) MCMC baselines.\n   * **Additional Datasets:** Beyond MovieLens and Swiss-roll, we now include MNIST and CIFAR. With only 5 % offline data, dTS yields up to 92.18 % regret improvement over LinTS and 79.56 % over HierTS on CIFAR, with similar gains on MNIST.\n\n2. **Clarified Theoretical Analysis (Sec 3).**\n   We now explicitly outline our regret-bound derivation strategy, clarifying that it is **not** a simple covariance plug-in. Instead, we exploit the diffusion model’s **hierarchical structure** and track **prior-influence decay** via the diffusion ODE to decompose regret into **L+1 interpretable terms**, each corresponding to a latent layer of the diffusion model.\n\n3. **Notation & Typos Fixed.**\n   We revised our notation. Thank you for these suggestions.\n\n---\n\n## Point-by-Point Responses\n\n**1. Missing MCMC comparison.**\nWe agree that an MCMC baseline is essential. We compared dTS to two MCMC variants:\n\n| Baseline      | Regret Improvement (%) | Time Speed-up (%) |\n| ------------- | ---------------------: | ----------------: |\n| dTS-MCMC-Fast |                 50.6 % |            47.6 % |\n| dTS-MCMC-Slow |                 12.7 % |            80.5 % |\nwhere:\n* **Regret Improvement** = 100 × (Regret\\_baseline – Regret\\_dTS) / Regret\\_baseline\n* **Time Speed-up** = 100 × (Time\\_baseline – Time\\_dTS) / Time\\_baseline\n\nEven the high-compute “Slow” MCMC variant has higher regret than dTS. Applying MCMC with diffusion priors in an **online bandit** setting proved highly challenging, this was actually our first attempt and motivated our efficient approximation, which we believe is more practical.\n\n---\n\n**2. Additional real-world datasets**\n\n> *“Beneficial to include more experiments like MovieLens to convince practitioners.”*\n\n* **Response:** We add MNIST and CIFAR benchmarks (very common benchmarks in practice). With priors trained on just **5 %** of offline data, dTS outperforms both HierTS and LinTS:\n\n| Dataset | vs. HierTS | vs. LinTS |\n| :-----: | ---------: | --------: |\n|  MNIST  |    42.17 % |   82.30 % |\n|  CIFAR  |    79.56 % |   92.18 % |\n\nWe also include **cold-start ablations** where we vary the amount of offline data available to pre-train the diffusion model prior:\n\n| Offline Data (%) | vs. HierTS | vs. LinTS |\n| :--------------: | ---------: | --------: |\n|        1 %       |    69.11 % |   87.74 % |\n|        5 %       |    79.56 % |   92.18 % |\n|       25 %       |    80.65 % |   92.48 % |\n|       50 %       |    81.67 % |   92.88 % |\n\nThese results demonstrate the **practical value** of pretraining diffusion priors with minimal data.\n\n---\n\n**3. Depth of theoretical analysis**\n\n> *“Unclear if regret bound is mere covariance plug-in or more involved.”*\n\n* **Response:** * **Response:** We clarified in the revised version that while the diffusion model collapses to a Gaussian when everything is linear-Gaussian, our analysis does **not** simply compute a prior covariance and plug it into standard TS bounds. Instead, we:\n* Derive bounds from scratch,\n* Decompose the regret into **L+1 interpretable terms**,\n* Track **prior influence decay** across layers, exploiting the **ODE structure**.\n\nThis yields bounds that capture **how diffusion model depth and structure** affect regret, beyond what a covariance plug-in would reveal.\n\n---\n\n**4. Notation concerns**\n\n> *“Notation is confusing: random-variable fonts, duplicate symbols, tuple indices, undefined terms.”*\n\n* **Response:** We improved notation per your suggestions.\n\n---\n\nThank you again for your insightful comments. We believe that we fully addresses your concerns on empirical validation, theoretical clarity, and notation.\n\nIf you feel your main concerns have been resolved, we would greatly appreciate your renewed support. Please let us know if anything remains unclear."}]

Author rebuttal:

Meta-review: The paper proposes a new contextual bandit algorithm that extends the literature on hierarchical Bayesian bandits to a more complex setting, where the prior distribution is represented by a pre-trained diffusion model. The authors analyze the Bayesian regret when the reward functions are linear and the prior is Gaussian. Numerical experiments are presented to illustrate the efficiency of the proposed method.

All reviewers acknowledged the novelty of the approach and the promising experimental results. However, several shortcomings were noted that should be addressed in a revised version of the paper:

•	The theoretical guarantees are derived only for the linear-Gaussian case, which essentially means that the prior is Gaussian and that the diffusion model is not actually necessary. The authors should comment on this limitation and/or explain how the analysis might be extended.

•	The use of a pre-trained diffusion model can be problematic if offline data is limited. While the authors addressed this concern in the rebuttal, the discussion should be further developed and explicitly included in the paper.

•	The experiments should be expanded and completed as suggested by Reviewer 8Fbo.

Overall, the paper introduces a novel approach to contextual linear bandits. Unfortunately, the theoretical results only support the Gaussian prior setting, essentially reducing the role of the diffusion model. Therefore, to strengthen the contribution, the paper should include extended numerical experiments that both motivate the use of diffusion priors and demonstrate their advantages in practice.
Decision: Accept (poster)