Online Posterior Sampling with a Diffusion Prior
Authors: Branislav Kveton
Boris N. Oreshkin
Youngsuk Park
Aniket Anand Deshmukh
Rui Song
Venue: NeurIPS 2024 poster (2024)
Keywords: posterior sampling
diffusion models
online learning
contextual bandits
Primary area: bandits
TL;DR: We propose Laplace posterior sampling approximations for linear models and GLMs with a diffusion model prior, and apply them to contextual bandits.

Abstract: Posterior sampling in contextual bandits with a Gaussian prior can be implemented exactly or approximately using the Laplace approximation. The Gaussian prior is computationally efficient but it cannot describe complex distributions. In this work, we propose approximate posterior sampling algorithms for contextual bandits with a diffusion model prior. The key idea is to sample from a chain of approximate conditional posteriors, one for each stage of the reverse diffusion process, which are obtained by the Laplace approximation. Our approximations are motivated by posterior sampling with a Gaussian prior, and inherit its simplicity and efficiency. They are asymptotically consistent and perform well empirically on a variety of contextual bandit problems.

Online Posterior Sampling with a Diffusion Prior
Branislav Kveton
Boris N. Oreshkin Amazon
Youngsuk Park AWS AI Labs
Aniket Deshmukh AWS AT Labs
Rui Song Amazon
Adobe Research"
Abstract
Posterior sampling in contextual bandits with Gaussian prior can be implemented exactly approximately using the Laplace approximation. The Gaussian prior is computationally efficient but it cannot describe complex distributions. we propose approximate posterior sampling algorithms for contextual bandits with a diffusion model prior. The key idea to sample from a chain of approximate conditional posteriors, one for each stage of the reverse diffusion process, which are obtained by the Laplace approximation. Our approximations are motivated by posterior sampling with Gaussian prior, and inherit its simplicity and efficiency. They are asymptotically consistent and perform well empirically on a variety contextual bandit problems.
this work,
Introduction
multi-armell bandit (27, 6, 30] is online learning problem where agent sequentially interacts with an environment over n rounds with the goal of maximizing its rewards. In each round, it takes an action and receives its stochastic reward. The mean rewards of the actions are unknown a priori and must be learned. This leads to the exploration-exploitaaaai about them or exploir the action with the highest estimated reward. Bandits have been successfully applied to problems where uncertainty modeling and adaptation are beneficial, such recommender systems [32, 54, 25, 35] and
dilemma: explore actions to learn
optimization [34].
Contextual bandits [29, 32] with linear [13, 1] and generalized linear models (GLMs) [17, 33, 2, 26] have become popular due the their flexibility and efficiency. The features in these models can be hand-crafted or learned from historic data [40], and the models can be also updated incrementally [1, 24]. While the original algorithms for linear and GLM bandits were based on upper confidence bounds (UCBs) [13, 1, 17]. Thompson sampling (TS) is more popular in practice [11, 3, 42, 44J. The key idea in TS is explore by sampling from the posterior distribution of model parameter 0.. TS uses the prior knowledge about 6, to speed up exploration [11, 40, 36, 9, 21, 20, 5]. When the prior a multivariate Gaussian, the posterior of 6_ can updated and sampled from efficiently [11]. This prior has limited expressive power, because it cannot even represent multimodal distributions. To address this, we study posterior sampling with diffusion prior. The main benefit of such priors that they can represent complex distributions and be learned from data.
We make the following contributions. First, propose novel posterior sampling approximations for linear models and GLMs with a diff'usion model prior. The key idea is to sample from chain of approximate conditional posteriors, one for each stage of the reverse process, which are estimated in a closed form. In linear models, each conditional is a product of two Gaussians, representing prior knowledge and diffused evidence (Theorem 2). In GLMs, each conditional is obtained by Laplace approximation, which mixes prior knowledge and evidence (Theorem 4). Our approximations are motivated by posterior sampling with Gaussian priors, and inherit its simplicity and efficiency. Prior works (Section 7) sampled from the posterior using the likelihood score, and their approximations become unstable when the score is high. We combine the likelihood with conditional priors, each
'The work done at AWS AI Labs.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

stage of the diffusion model, using the Laplace approximation. The resulting posterior concentrates single point and can be sampled from efficiently even if the likelihood score high. We prove that this approximation is asymptotically consistent.
Our second contribution is in theory. We properly derive our posterior approximations (Theorems and 4) and show their asymptotic consistency (Theorem 3). The key idea in the proof of Theorem that the conditional posteriors concentrate of observations increases. While this claim is asymptotic, it is an expected property of a posterior distribution. Many prior works, such as Chung et al. [12], do not propose asymptotically consistent approximations. All of our main results rely on a novel approximation of clean samples by scaled diffused samples (Section 4.3). The most challenging part of the analysis is Theorem 3, where we analyze an asymptotic behavior of a chain of T dependent random vectors.
scaled unknown model parameter as the number
Our last contribution is empirical evaluation contextual bandits. We focus bandits because the ability to represent all levels of uncertainty precisely is critical for exploration. Our experiments show that score-based method fails to do (Section 6.2). Note that our posterior approximations are general and not restricted to bandits.
2 Setting
We start with introducing our notation. Random variables capitalized, except for Greek letters like 9. We denote the marginal and conditional probabilities under probability measure p by p(X = x) and p(X x Y = =y). respectively. When the random variables are clear from context, we write p(x) and p( y). We denote by Xnm and Trom collection of random variables and their values, respectively. For positive integer n, define [n] fl,...,,yy ,,). The i-th entry of vector e is Us. If the vector is already indexed, such as uy, we write Uz We denote the maximum and minimum eigenvalues of matrix M E dd by A1 and Xa(M), respectively.
The indicator function is 1 [-)-
The posterior sampling problem can be formalized follows. Let 0g E θ be an unknown model parameter and θ C IRd be the space of model parameters. Let h = {(de, Me) be) be the history of N noisy observations of 0., where de E TRd is the feature vector for y E R. We assume that
Y/ = 9 0.) + t
where g: - R is the mean function and Eş is an independent gero-mean -ussuSGaannia for a > O. Let p(h 0-) be the likelihood of observations in history h under model parameter 6, and p(0-) be its prior probability. By Bayes' rule, the posterior distribution of 0, given h is
Gaussian noise
p(0, h) x p(h| 0.) p(O.).
We want to sample from p(- h) efficiently when the prior distribution is represented by a diffusion model. As stepping stone, review existing posterior formulas for multivariate Gaussian priors. This motivates our solution for diffusion model priors.
2.1 Linear Model
The posterior of 0, in linear models can be derived as follows.
Assumption Let g in (1) be an identity and Eq ~ (0(0 o²). Then the likelihood of under model parameter 0., is p(/ 0.) x expl- Em (T..)/((20)))
Let p(O.) = (0:; 0, Mo) be the prior distribution of 0.. where 0g E Rd and Eo E Rdxd are the prior mean and covariance, respectively. Then p(0. h)) x WO00; 0, Σ) [10], where
0 - & 1 1o +
EN deve)
S = (2 11-1 - o - (107 +
the posterior mean and covariance, respectively. In this work, write them equivalently as = ا 0o + ビ-10),
Σ = 1-1--
1= drwn and = --1= - thee? are the empirical mean and inverse of its covariance, respectively. Therefore, the posterior of 0, is product of two multivariate Gaussians: (Oo,, Eo) representing prior knowledge about 0., and (/0 Σ) representing empirical evidence.
where 0 0= o -

---

Algorithm IRLS:
Iteratively reweighted least squares.
Input: Prior parameters 0n and 2o. history of observations {(oe, 3e Ne)
2: Initialize 0 E Rdd
3: repeat
4: for stage E ,...., NN,NNo do
5:
2t - 0 + (yt g 0)) 0))/0(000
6:
Eg Ee=1 5602 0 spp
7:
+ Eg Ou + Ei=1 1(0) 0)
8: until 0 converges
9: Output: Posterior mean 0 and covariance Σ
2.2 Generalized Linear Model
Generalized linear models (GLMs) [37] extend linear models (Section 2.1) to non-linear monotone mean functions g in (1). For instance, in logistic regression, g(u) = 1/(1 + exp[-ul) is a sigmoid. The likelihood of observations in GLMs has the following form [26].
Assumption 2. Let h {(or, RE(NE be a history of N observations under mean function g and the corresponding noise. Then log 0.) Ee=1 yeol 0. bl 0.) + clye), where is real function and b a function whose derivative is the mean functtin, b = g.
The posterior distribution of 0. in GLMs does not have closed form in general [10]. Therefore, it is often approximated by the Laplace approximation. Let the prior distribution of the model parameter be p(O.) (0,; to Eo), as in Section 2.1. Then the Laplace approximation is (@ Σ). where is the maximum a posteriori (MAP) estimate of 0, and Σ is the corresponding covariance, Note that the Laplace approximation can be applied to non-Gaussian priors.
The MAP estimate i can be obtained by iteratively reweighted least squares (IRLS) [52], which we present in Algorithm 1. IRLS is a Newton-type algorithm that computes Ô iteratively (lines 6 and 7). It converges to the optimal solution due to the strong convexity of the problem. The solution has similar structure to (3). That is, ((,. Σ) is product of two multivariate Gaussians, representing prior knowledge about 0 and empirical evidence. The new quantities in GLMs are the derivative of the mean function and pseudo-observation zς (line 5), which play the role of observations yE in Section 2.
2.3 Towards Diffusion Model Priors
The assumption that p(0.) = N(e:: 0y, Eo) is limiting, for instance because it precludes multimodal priors. We relax it by representing p(f.) by diffusion model, which we call diffusion model prior. We propose efficient posterior sampling approximations for this prior, where the prior and empirical evidence mixed similarly (3) and IRLS. We review diffusion models next.
Diffusion Models
Diffusion models [46, 19] are generative models trained by diffusing samples from unknown and hard to represent distributions. They be viewed in multiple ways [49]. We adopt the probabilistic formulation and presentation of Ho et al. [19]. A diffussio model is a graphical model with stages indexed by E [T]. Each stage t is associated with a latent variable S, E Rdd. A sample from the model is represented by observed variable Sy Rd We visualize a diffusion model in Figure 1. In the forward process, a clean sample 8p is diffused through a sequence of variables S Sr- This process is used to learn the reverse process, where the clean sample 80 is generated through sequence of variables ST.... Sp. To sample 89 from the posterior (Section 4), we add random variable H that represents partial information about so. We introduce forward and reverse diffusion

---

Forward process (probability measure q) S1 + Su
Reverse process (probability measure p) ST - ST - + - So - H
ST + Sr
Figure I: Graphical models the forward and reverse processes in the diftusion model. The variable H represents partial information about So-
processes next. Learning of the reverse process is described in Appendix B. While this is a critical component of diffusion models, it is not necessary introduce our posterior approximations.
Forward process. In the forward process, a clean sample so is diffused through a chain of latent variables S1. (Figure 1). We denote the probability measure under this process by q and define its joint probability distribution as
q($1:T so) IT= ((s st-1), e [ g(s; (t-1) (s Voist-1. Bela),
(4)
where g(st | 8t-1) the conditional density mapping less diffused st-1 to a more diffused 8t- The diffusion rate is set by parameters o, E (0, 1) and 3r = 1 - og. The forward process is sampled from as follows. First, a clean sample 8o is chosen. Then S q(: 8t-1) are sampled, from to T. T.
Reverse process. In the reverse process, clean sample so generated through a chain of variables Sr..... So (Figure 1). We denote the probability measure under this process by and define its joint probability distribution as
p(su:r) = p(8))
II 1=1
p(
(5)
p(ST)= = (sr: I2), vt E :(8 8t) 3 (;
where p(s;-1 8,) is the conditional density of mapping more diffused 8, to less diffused 8;-1- The function je predicts the mean of S s and is learned (Appendix B). As in Ho et al. [19], keep the covariance fixed at E, = la. where 3. 1- Br and a; = IT dg. This known as stable diffusion. We make this assumption only to simplify exposition. All our derivations in Section 4 hold when Σ, is learned, for instance as in Bao et al. [8].
1-
This process is called reverse because it is learned by reversing the forward process. The reverse process is sampled from as follows. First, an initial diffused sample ST ~ p is sampled. After that, 8,-1 p- 82) sampled, from / = T to 1.
4 Posterior Sampling
This section is organized as follows. In Section 4.1, we show how to sample from a chain random variables conditioned observations. In Sections 4.2 and 4.4, specialize this to the observation models in Section
4.1
Chain Model Posterior
Let h = ((ée, beeNeN denote history of N observations (Section 2) and H be the corresponding random variable. In this section, we assume that h is fixed. The Markovian structure of the reverse process (Figure 1) implies that the joint probability distribution conditioned on factors as
p(Sorr 1h) =p(sT 1)) P(s-11 (8,
Therefore, p(so:T h) can be sampled from efficiently by first sampling from p(st h) and then from T' conditional distributions p(s:-1 | st. h). We derive these next.
Lemma 1. Let p be a probability measure over the reverse process (Figure 1). Then
p(sT h) xx Sog p(h | 80) p(50 ST) ds0 p(TT),
Vt E [2] (1) p(5
st, ht x S'p p(h so) p(so st-1) da p(s 9t),
p(so h) x p(h 50) p(so $1).

---

Algorithm 2 LaplaceDPS: Laplace posterior sampling with a diffusion model prior.
Input: Diffusion model parameters (Me, Et)ee history of observations h
2: Initial sample ST (jr++ (1),
3: for stage 1. = T, ......1 do
4: St-1 ~ N (S,, (+))
5: Output: Posterior sample So
Proof. The claim is proved in Appendix A.I.
4.2 Linear Model Posterior
Now we specialize Lemma to the diffusion model prior (Section 3) and linear models (Section 2.1).
The prior distribution is the reverse process in (5),
p(sT)= = (sr: 0g, Ia), A [T] p(s.11 st) = -i E.).
The term p(h so) the likelihood of observations in Assumption 1. The main challenge in using the lemma is that the conditional densities of clean samples p(so ST) and p(sp st) are complex [12]. To get around this, we make an additional assumption, which is discussed in Section 4.3.
Theorem 2. Let p be probability measure over the reverse process (Figure 1). Let 0 and E-1 be defined as in (3). Suppose that
Seg p(h so) p(so st) dsg x p(h stt)
(6)
holds for all E (T). Then p(sT
h) x N(T: (1T (h), 1 where
(h) = Tx)1 ( + E 10/VoT) (r+1 =
1a
Prior
Evidence
Prior
Evidence
For t E [T], we have p(st-1 | St, h) x (N (s h), Ex(h),, where
h) = (E) (a(()) + -- -E
Var-))
Ex(h)
(E
Σ 1/111111)-
(8)
Prior
Evidence
Prior
Evidence
Proof. The proof in Appendix A.2. has four steps. First, we fix stage and apply approximation
(6). Second, we rewrite the likelihood as in (3). Third, we reparameterize it as function of st. At the end, we combine the likelihood with the Gaussian prior using Lemma 6 in Appendix A.5. O
The algorithm that samples from the posterior distribution in Theorem is presented Algorithm 2. We call it Laplace diffusion posterior sampling (LaplaceDPS) because its generalization to GLMs uses the Laplace approximation. LaplaceDPS samples from a chain of products of two distributions: one distribution represents the pre-trained diffusion model and does not depend on history h, and the other represents the history h. The sampling implemented as follows. The initial variable Sr is sampled conditioned on h (line 2) from the distribution in (7). This distribution is product of the h-independent prior (0,,,, Ia) and the h-dependent distribution of the diffused evidence up to stage T. (Var0 E). Then, for any / [T), S. is sampled conditioned on 8, and evidence h (line 4) from the distribution (8). This distribution is product of the h--independent conditional prior N ((u(s)) ),), from the pre-trained model, and the h-dependent distribution of the diffused evidence up to stage t 1, AV( ((va,110 E). The last variable So is the clean sample. When compared to Section 2, the prior and evidence mixed conditionally in T'sstgee chain. This increases the computational cost T times, as discussed in Section 8.
4.3 Key Approximation in Theorem
Now we motivate our assumption in (6). Simply put, we assume that 8q = =8// s./Von. where is clean sample and 8, the corresponding diffused sample in stage f. This motivated by the forward process, which relates and so st Vā,so + vI azEt, where e; NN,,,, Ia) is standard
5

---

Algorithm 3 Contextual Thompson sampling.
1: for round & = 1, I,..., do
2: Sample 8g p(- hx),, where p(- ha) is the posterior distribution in (2)
3: Take action aą arg max aeA r(rk,, a; Og) and observe reward yk
Gaussian noise [19]. After rearranging, we get So (s: - -I- Vāg, and therefore can be viewed as a random variable with mean 8// yā, The consequence of (6) that the likelihood becomes a function of 8t. which yields a closed form when multiplied by the conditional prior, also function of 8,- Our approximation can be also viewed as the Tweedie's formula Chung et [12] where the score component is neglected.
Our approximation has several notable properties. First, (1)aa)/ā + 0 as 1. Therefore, becomes more precise in later stages of the reverse process. Second, in the absence of evidence h, the approximation vanishes, and all posterior distributions in Theorem 2 reduce to the priors in (5). Finally, as the number of observations increases, sampling from the posterior in Theorem is asymptotically consistent.
Theorem 3. Fix 0, E Rd. Let Ō + (ap, E Ee)erry h), where h = ((de, be) is history of N observations. Suppose that >a(E-1) + co as N + co, where is defined in (3). Then [imm 110 0x112 = 0)
Proof. The proof is in Appendix A.3. The key idea that the conditional posteriors (7) and (8) concentrate at scaled unknown model parameter 0, as the number of observations increases, which we formalize as Aa(E-1) + 0o.
The bound in Theorem 3 can be interpreted as follows. The sampled parameter 0 approaches the true unknown parameter 9, as the number of observations increases. To guarantee that the posterior shrinks uniformly in all directions, we assume that the number of observations in all directions grows linearly with N. This is akin assuming that { (E-1) )) = S2(N). This lower bound can be attained in linear models by getting observations according to the D-optimal design [39].
4.4 GLM Posterior
The Laplace approximation in GLMs (Section 2.2) naturally generalizes the exact posterior distribu- tion in linear models (Section 2.1). We generalize Theorem to GLMs along the same lines.
Theorem 4. Let p be probubility measure over the reverse process (Figure 1). Suppose that (6)
holds for all E (T). Then p(sT) h) x XN(ss: (T Er-i(h)), where
(h)+ Värdrti, (h) or Ext1 0rt1,E IRLS(Oa, Ta/ārr,
For 1 E [T], we have p(or-1 88, h) x N (sr h), E,((i)), where
h))st, = var-id,, (1(1)== = o, E, c IRLS(U Von-11 E1/01-1, It)
Proof. The proof in Appendix A.4. has four steps. First, fix stage and apply approximation (6). Second, we reparameterize the prior, from a function of 8, to a function of s Vor. Third, we combine the likelihood with the prior using the Laplace approximation. Finally, repameterize the posterior, from a function of s 1//va, Vo, to a function of 8t-
Similarly to Theorem 2, the distributions in Theorem 4 mix evidence with the diffusion model prior. However, this is done implicitly in IRLS. The posterior can be sampled from using LaplaceDPP where the mean and covariances would taken from Theorem 4. Note that Theorem a special case of Theorem where the mean function is identity.

---

5 Application to Contextual Bandits
Now we apply our posterior sampling approximations (Section 4) contextual bandits. contextual bandit [29, 32] is a classic model for sequential decision making under uncertainty where the agent takes actions conditioned on context. We denote the action set by A and the context set by X. The mean reward for taking action E A in context I E X is r(z, a; B.), where r X' x A x θ + R is a model purameter (Section 2). The agent interacts with the bandit for n rounds indexed by k E [n)] In round k; it observes context FA E X. takes an action dg A, and observes its stochastic reward yı = r( (g;0.) 0.) + ex with independent noise E We assume that the noise is zero-mean oL-sub-Gaussian for a >> 0. The objective of the agent is to maximize its cumulative reward in n rounds, or equivalently to minimize its cumulative regret. We define the n-round regret
a reward function and 0.
R(n) - Ei= [r(IR, dR.0; 0-) - r( aw: 0+)]
where ag,= arg max aEA r (T a; 0.) is the optimal action in round k.
Arguably the most popular method for solving contextual bandit problems is Thompson sampling [50, 11,3] The key idea in TS is to use the posterior distribution of 0, to explore. This done as follows. In round k the model parameter is drawn from the posterior in (2), 0g p(- hy). where hg
is the history of all interactions up to round k. After that, the agent takes the action with the highest mean reward under 0z. The pseudo-code of this algorithm is given in Algorithm 3.
A linear bandit [13] has linear reward function r(z, a; 0.) Φ(2, a) (T)Te, is a feature extractor. The feature extractor can be non-linear in x and a. Therefore, linear bandits can be applied to non-linear functions of x and The feature extractor hand-crafted. We denote the feature vector of the action in round k by Φb 6(IK, ax).. Therefore, the history of interactions up to round i is h4 {(oe, ye)) re)) When the prior distribution is Gaussian, p(0.) = N No.;; lo. Mo), the posterior in round is a Gaussian in (3) for h = fug. When the prior is diffusion model, propose sampling from the posterior using
0., where 6 X' A Rd
be either learned [40] or
0 LaplaceDPS((He: Ee)eery hx),
(10)
where júg and ES in LaplaceDPS computed as in Theorem 2. We call this algorithm DiffTS.
generalized linear bandit [17, 24, 33, 26] is an extension of linear bandits to generalized linear models (Section 2.2). When p(0.) = _N(B,; fo. Eo). the Laplace approximation to the posterior is Gaussian (Section 2.2). When the prior diffusion model, we propose posterior sampling using (10), where iú and E. in LaplaceDPS are computed as Theorem 4.
Experiments
We conduct three experiments: synthetic problems in 2 dimensions (Section 6.2 and Appendix C.I), a recommender system (Section 6.3), and a classification problem (Appendix C.2). In addition, we conduct an ablation study in Appendix C.3, where we vary the number of training samples for the diffusion prior and the number of diffusion stages T.
6.1 Experimental Setup
We have four baselines. Three baselines are variants of contextual Thompson sampling [11, 3]: with an uninformative Gaussian prior (TS), a learned Gaussian prior (TunedTS), and a learned Gaussian mixture prior (MixTS) [22]. The last bascline is diffusion posterior sampling (DPS) of Chung et al. [12]. We implement all TS baselines as described in Section 5. The uninformative prior is (O,,,, Ia). MixTS is used only in linear bandit experiments because the logistic regression variant does not exist. The TS baselines are chosen to cover various levels of prior information. Our implementation of DPS is described in Appendix D. We also experimented with frequentist baselines, such LinUCB [1] and the e-greedy policy. They performed worse than TS and thus we do not report them.
Each experiment is set up as follows. First, the prior distribution of 0, is specified: it can be synthetic or estimated from real-world data. Second, we learn this distribution from 10 000 samples from In DiffTS and DPS, we follow Appendix B. The number of stages is T 100 and the diffusion factor
7

---

Problem cross
Problem
Problem triangles
4
True and diffusiion prtoes
Thue and alffusion arfors
diffusion
-1
-1
-1
Problem cross
Problem rays
Problem triangles
IT
TS
TS
20
TunedTS
TunedTS MixTS
TunedTs
MIXTS
MixTS
Regret
15
DPS
Regret
DPS
Rogret
DPS
DIATS
DaFTs
PILTS
10
10
10
5
5
50
100
150
200
100
150
200
50
100
150
200
Round
Round
Round
Figure 2: Evaluation of DiffTS on three synthetic problems. The first row shows samples from the true (blue) and diffusion model (red) priors. The second row shows the regret DiffTS and the baselines as function of round
ar 0.97. Since 0.97 2s 0.05, most of the information in the training samples is diffused. The regressor in Appendix B is 2-layer neural network with ReLU activations. In TunedTS, we fit the mean and covariance using maximum likelihood estimation. MixTS, we fit the Gaussian mixture using SCIKIT-LEARN, All algorithms are evaluated on 0, sampled from the true prior. The regret is computed as defined in (9). All error bars are standard errors of the estimates.
6.2 Synthetic Experiment
The first experiment is on three synthetic problems. Each problem is linear bandit (Section 5) with K 100 actions in d 2 dimensions. The reward noise = 1. The feature vectors of actions sampled uniformly at random from a unit ball. The prior distributions of 0, are shown in Figure 2. The first is mixture of two Gaussians and the last can be approximated well by mixture of two Gaussians. We implement MixTS with two mixture components. Therefore, it can represent the first prior exactly and approximate the last one well.
Our results are reported in Figure We observe two main trends. First, samples from the diffusion prior closely resemble those from the true prior. In such cases, DiffTS is expected perform well and even outperforms MixTS, because has a better representation all problems. Second, DPS diverges the number of rounds increases. This is because DPS uses approximation based on the likelihood score (Section 7), which is unstable when the score is high. best efforts to tune DPS (Appendix D). We report results another three
the prior. We observe this in
This happens despite synthetic problems in Appendix C.I.
DiffTS should be T times more computationally costly than TS with Gaussian prior (Section 4.2). observe this empirically, As an example, the average cost of 100 runs DiffTS on any problem in Figure 2 is 12 seconds. The average cost of TS is 0.1 seconds. The computation and accuracy can be traded off, and we investigate this in Appendix C.3. In the cross problem, we vary the number of diffusion stages from T = 1 to T 300. We observe that the computational cost is linear in T and the regret drops quickly from 26 at T to 15 at 50.
6.3 MovieLens Experiment
In the second experiment, we learn recommend an item to randomly arriving users. The problem simulated using the Movielens 1M dataset [28], with one million ratings for 706 movies from 6 040 users. We subtract the mean rating from all ratings and complete the sparse rating matrix M by

---

[a] Problem MovieLens
(b) MovieLens bandit
(c) MovieLens logistic bandit
12
16
10
14
155
12
Regret
6
10
T5
5
TunedTS
MiXTS
TuncdTS
DIfTTS
DirTS
0
100
150
200
50
Round
Figure 3: Evaluation of DiffTS on the MovieLens dataset: (a) shows samples from the true (blue) and diffusion model (red) priors, (b) shows regret in the linear bandit, and (c) shows regret in the logistic bandit.
alternating least squares [14] with rank d = 5. The learned factorization is = UVT The i-th row of U, denoted by U, represents user i. The j-th row of V, denoted by V,, represents movie j. We use movie embeddings Vj as model parameters and user embeddings U, as features of the actions. The movies are items.
We experiment with both linear and logistic bandits. In both, an item is initially chosen randomly from Vg and K = 10 actions are chosen randomly from U; in each round. In the linear bandit, the mean reward of item j for user is U. Vj. The reward noise is a 0.75, and we estimate it from data. In the logistic bandit, the mean reward is g(U V₃), where g is a sigmoid.
Our MovieLens results are reported in Figure 3 and we observe similar trends Section 6.2. First, samples from the diffusion prior closely resemble those from the true prior (Figure 3a). Since the problem is higher dimensional, we visualize the overlap using UMAP [45]. Second, DiffTS has lower regret than all baselines, in both linear (Figure 3b) and logistic (Figure 3c) bandits. Finally, MixTS barely outperforms TunedTs. We observe this trend consistently in higher dimensions, and this motivated our work on online learning with more complex priors.
7 Related Work
We start with reviewing related works on bandits with diffusion models. Hsieh al. [23] proposed Thompson sampling with a diffusion model prior for K K-armed bandits. There are multiple technical differences from work. First, the diffusion model in Hsieh et al. [23] is scalars representing individual arms. Our model is over vectors representing model parameters, and thus can be applied to contextual bandits. Second, the approximations are different. In stage I, Hsieh et al. [23] sample from two distributions: the conditional prior and the distribution of the diffused empirical mean up to stage t. Then they take a weighted sum of the samples. We sample only once, from the posterior distribution that combines the conditional prior in stage t and likelihood. Therefore, the method of Hsieh et al. [23] can be viewed as a non-contextual variant of our method, where posterior sampling is done by weighting samples from the prior and empirical distributions. Finally, Hsieh et al. [23] do not analyze their approximation.
Aouali [4] proposed and analyzed contextual bandits with a linear diffusion model prior: (s:)) in (5) is linear in 8, and ((5)) is Gaussian. Therefore, this model is linear Gaussian model and not general diffusion model, as in our work.
The closest related work posterior sampling in diffusion models is DPS of Chung al. [12]. The key idea in DPS is to sample from the posterior distribution using the likelihood score V log p(h 0). where p(h 0) is the likelihood (Assumptions and 2). Note that V log p(h 0) grows linearly in N because the history in p(h| 0) involves N terms. Therefore, DPS becomes unstable as N + co. We show it empirically in Section 6.2 and discuss the implementation of DPS in Appendix D, which was tuned to improve its stability.
Many other posterior sampling methods for diffusion models have been proposed recently: sequen- tial Monte Carlo approximation for the conditional reverse process [53], variant of DPS with an uninformative prior [38], pseudo-inverse approximation to the likelihood of evidence [48], and

---

posterior sampling in latent diffusion models [41]. All of these methods rely on the likelihood score v loo p(h 0) and thus become unstable as the number of observations N increases. Our posterior approximations do not have this issue because they are based on the product of prior and evidence distributions (Theorems and 4), and thus gradient-free. They work well across different levels of uncertainty (Section 6) and do not require tuning.
We note that posterior sampling a special form of snference-time guidance in diffusion models. Other approaches are conditional pre-training [15], a constraint in the reverse process [18], refining the null- null-space content [51], solving an optimization problem that pushes the reverse process towards evidence [47], and aligning the reverse process with the prompt [7].
8 Conclusions
We propose posterior sampling approximations for diffusion models priors. These approximations are contextual, and can be implemented efficiently in linear models and GLMs. We analyze them and evaluate them empirically on contextual bandit problems. Our method has two main limitations.
Computational cost. The cost of posterior sampling in LaplaceDPS with T stages is about times higher than that of posterior sampling with a Gaussian prior (Section 2). We validate it empirically in Section 6.2. We plot the sampling time as a function of T in Figure 6c (Appendix C.3).
Learning cost and hyper-parameter tuning. In all experiments, the number of diffusion stages is T 100 and the diffusion rate is set such that most of the signal diffuses. The regressor is 2-layer neural network and we learn it from 10 000 samples from the prior. These settings resulted in stable performance in all our experiments (Section 6). However, they clearly impact the performance. We plot the regret as a function of the number of training samples in Figure 6a and as a function of T in Figure 6b. When T or the number of training samples small, DiffTS performs very similarly to posterior sampling with Gaussian prior. In summary, there is no benefit in these cases.
Future work. We develop novel posterior approximations rather than bounding their regret. This is because the existing approximations are unstable and may diverge in the online setting (Sections 6.2 and 7). We believe that a proper regret analysis of DiffTS is possible and would require bounding two errors. The first error arises because the reverse process does not reverse the forward process exactly (Appendix B). The second error arises because our posterior distributions are approximate (Section 4.3). One possibility is to start with prior works that already showed the utility of complex priors. For instance, Russo and Van Roy [43] proved a O(VTH((A )n) regret bound for linear bandit, where I' is the maximum ratio of regret to information gain and H (A.) is the entropy of the distribution of the optimal action under the prior. This bound holds for any prior and says that a lower entropy H (A.),, which corresponds to more informative priors, yields a lower regret.
We also believe that our ideas can be extended beyond GLMs. The key idea in Section 4.4 is to use the Laplace approximation of the likelihood. This approximation can be computed exactly in GLMs. More generally though, it is a good approximation whenever the likelihood can be approximated well by a single Gaussian distribution. By the central limit theorem, under appropriate assumptions, this is expected for any observation model when the number of observations is large.
References
[1] Yasin Abbasi-Yadkorf David Pal, and Csaba Szepesvari. Improved algorithms for linear stochastic bandits. In Advances in Neural Information Processing Systems 24, pages 2312-2320, 2011.
[2] Marc 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] Imad Aouali. Linear diffusion models meet contextual bandits with large action spaces. In NewrlPS 2023 Workshop on Foundation Models for Decision Making Workshop. 2023.
[5] Imad Aouali, Branislav Kveton, and Sumeet Katariya. Mixed-effect Thompson sampling. In Proceedings of the 26th International Conference on Artificial Intelligence and Statistics, 2023.
10

---

[6] Peter Auer, Nicolo Cesa-Biand Bianchin and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine Learning, 47:235 2002.
[7] Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the 12th International Conference on Learning Representations, 2024.
[8] Fan Bao, Chongxuan Li, Jiacheng Sun, Jun Zhu, and Bo Zhang. Estimating the optimal covariance with imperfect mean in diffusion probabilistic models. In Proceedings of the 39th International Conference on Machine Learning, 2022.
[9] 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.
[10] Christopher Bishop. Pattern Recognition and Machine Learning. Springer, New York, NY, 2006.
[1]] Olivier Chapelle and Lihong Li. An empirical evaluation of Thompson sampling. In Advances in Neural Information Processing Systems 24, pages 2249-2257 2011.
[12] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In Proceedings of the I1th International Conference on Learning Representations, 2023.
[13] Varsha Dani, Thomas Hayes, and Sham Kakade. Stochastic linear optimization under bandit feedback. In Proceedings of the 21st Annual Conference on Learning Theory, pages 355-366, 2008.
[14] Mark Davenport and Justin Romberg. An overview of low-rank matrix recovery from incomplete observations. IEEE Journal of Selected Topics in Signal Processing, 10(4):608-63) 622, 2016.
[15] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems 34, 2021.
[16] Arnaud Doucet, Nando de Freitas, and Neil Gordon. Sequential Monte Carlo Methods in Practice. Springer, New York, NY, 2001.
[17] 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, 2010.
[18] Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-pin -pay priors. In Advances in Neural Information Processing Systems 35, 2022.
[19] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems 33, 2020.
[20] Joey Hong, Branislav Kveton, Sumeet Katariya, Manzil Zaheer, and Mohammad Ghavamzadeh. Deep hierarchy in bandits. In Proceedings of the 39th International Conference on Machine Learning, 2022.
[21] 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.
[22] Joey Hong, Branislav Kveton, Manzil Zaheer, Mohammad Ghavamzadeh, and Craig Boutilier. Thompson sampling with mixture prior. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics, 2022.
[23] Yu-Guan Hsieh, Shiva Kasiviswanathan, Branislav Kveton, and Patrick Blobaum. Thompson sampling with diffusion generative prior. In Proceedings of the 40th International Conference on Machine Learning, 2023.
[24] Kwang-Sung Jun, Aniruddha Bhargava, Robert Nowak, and Rebecca Willett. Scalable gen- eralized linear bandits: Online computation and hashing, In Advances in Neural Information Processing Systems 30, pages 98-108, 2017.
[25] Jaya Kawale, Hung Bui, Branislav Kveton, Long Tran-Thanh, and Sanjay Chawla. Efficient Thompson sampling for online
recommendation. In Advances Neural Information Processing Systems 28, pages 1297-1305, 2015.

---

[26] Branislav Kveton, Manzil Zaheer, Csaba Szepesvari, Lihong Li, Mohammad Ghavamzadeh, and Craig Boutilier. Randomized exploration in generalized linear bandits. In Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics, 2020.
[27] Tze Leung Lai and Herbert Robbins. Asymptotically efficient adaptive allocation rules. Ad- vances in Applied Mathematics, 6(1 6(12222 1985.
[28] Shyong Lam and Jon Herlocker. MovieLens Dataset. hhtp://gouplees.orr/dttttssssss 2016.
[29] John Langford and Tong Zhang. The epoch-greedy algorithm for contextual multi-armed bandits. In Advances in Neural Information Processing Systems 20, pages 817-824, 2008
[30] Tor Lattimore and Csaba Szepesvari. Bandit Algorithms. Cambridge University Press, 2019.
[31] Yann LeCun, Corinna Cortes, and Christopher Burges. MNIST Handwritten Digit Database.
http: / yann. lecun. com/ excb bxcbb/ma /mnist, 2010.
[32] 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.
[33] 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.
[34] Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hy- perband: A novel bandit-based approach to hyperparameter optimization. Journal of Machine Learning Research, 18(185):1-52 2018.
[35] Shuai Li, Alexandros Karatzoglou, and Claudio Gentile. Collaborative filtering bandits. In Proceedings of the 39th Annual International ACM SIGIR Conference, 2016.
[36] Xiuyuan Lu and Benjamin Van Roy. Information-theorrtiicc confidence bounds for reinforcement learning. In Advances in Neural Information Processing Systems 32, 2019.
[37] P. McCullagh and J. A. Nelder. Generalized Linear Models. Chapman & Hall, 1989.
[38] Xiangming Meng and Yoshiyuki Kabashima. Diffusion model based posterior sampling for noisy linear inverse problems. CoRR, abs/221 2023. URL https: //axiv.. org/ abs 22211221110 12343.
[39] Friedrich Pukelsheim. Optimal Design of Experiments. Society for Industrial and Applied Mathematics, 2006.
[40] Carlos Riquelme, George Tucker, and Jasper Snock. Deep Bayesian bandits showdown: An empirical comparison of Bayesian deep networks for Thompson sampling. In Proceedings of the 6th International Conference on Learning Representations, 2018.
[41] Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alex Dimakis, and Sanjay Shakkottai. Solving linear inverse problems provably via posterior sampling with latent diffusion models. In Advances in Neural Information Processing Systems 36, 2023.
[42] Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Mathematics of Operations Research, 39(4): 1211 2014.
[43] Daniel Russo and Benjamin Van Roy. An information--heoreet analysis of Thompson sampling. Journal of Machine Learning Research, 17(68):1-30, 2016.
[44] Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, and Zheng Wen. A tutorial on Thompson sampling. Foundations and Trends in Machine Learning, 11(1):1-96, 2018
[45] Tim Sainburg, Leland McInnes, and Timothy Gentner. Parametric UMAP embeddings for representation and semisupervised learning. Neural Computation, 33(1 12:81112722277 2021.
[46] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamies. In Pruceedings of the 32nd International Conference on Machine Learning, pages 2256-2265, 2015.
[47] Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, and Liyue Shen. Solving inverse problems with latent diffusion models via hard data consistency. In Proceedings of the 12th International Conference Learning Representations, 2024.
12

---

[48] Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guiee diffusion models for inverse problems. In Proceedings of the 11th International Conference on Learning Representations, 2023.
[49] Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Proceedings of the 9th International Conference on Learning Representations, 2021.
[50] William R. Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-299 1933.
[51] Yinhuai Wang. Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising dif- fusion null-space model. In Proceedings of the 11th International Conference on Learning Representations, 2023.
[52] R. Wolke and H. Schwetlick. Iteratively reweighted least squares: Algorithms, convergence analysis, and numerical comparisons. SIAM Journal on Scientific and Statistical Computing, (5):907-921, 1988.
[53] Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. In Advances in Neural Information Processing Systems 36, 2023.
[54] Xiaoxue Zhao, Weinan Zhang, and Jun Wang. Interactive collaborative filtering. In Proceedings of the 22nd ACM International Conference on Information and Knowledge Management, pages 1411-1420, 2013.
13

---

A Proofs and Supporting Lemmas
This section contains proofs of our main claims and supporting lemmas.
A.1 Proof of Lemma
All derivations are based on basic rules of probability and the chain structure in Figure 1, and are exact. From Figure 1, the joint probability distribution conditioned on H h factors as
T
p(so:T ht) p(ST h) II
p(s SUT, h)) p(ST h)
II t=1
p(s st, h).
t=1
We use that p(s SE.T, h) p(s st, h)) in the last equality. We consider two cases.
Derivation of p(s St h). By Bayes' rule, get
p(s 81, h) p(h 81-11, 86) p(st-1 1 8t)
p(h 8t-1) p(s:-1 88))
p(h| (s)
In the last step, we use that p(h (t) is a constant, since 8; and are fixed, and that p(h $1-1, 8t) = p(h 8t-1). Note that the last term p(8t-1 8x) is the conditional prior distribution. When > the first term can be expressed as
( p(h, so (s-1) dso ff p(h so, 81-1) p(so s2-1) dsu [
p(h s1-1)
p(h 80) p(so 88-11) d
In the last equality, use that our graphical model is chain (Figure 1), and thus p(h so, 8t-1) p(h 80). Finally, we chain all identities and get that
p(s h) [
p(h so) p(so sc-1) dso p(s:-1 st).
(11)
Derivation of p(ST h). By Bayes' rule, we get
p(sT) h)) p(h ST) p(sT) x p(h 18)) p(sT).
p(h)
In the last step, use that p(h) is constant, since h fixed. The first term can be rewritten
p(h ST) = / p(h, s0 88) dso = p(h I s0, ST) p(so 88) d8a
p(h I 8o) p(so ST) doo-
Finally, we chain all identities and get that
p(sr h) x / p(h sp) p(so ST) dsg p(sT).
(12)
This completes the derivations.
A.2 Proof of Theorem
This proof has two parts.
Derivation of p(sr-1 18,5)) h). From (6) and Assumption 1, follows that
/ p(h (0) 18-1)) doo & p(h
Va.-) x ((5--1/ x N(
E.) 1
14

---

The last step treats 0z-1 and as constants, because the forward process and evidence h are fixed. Now we apply Lemma distributions
p(s:-1 $t) = ((8)-11;((s)) Ex), N( Vār-10, az-
and get that
p( st,h) h)) x (i h), (k))
where (8t, h) and (h) are defined in the claim. This is a product of two Gaussians: the prior with mean (se) and covariance Σt, and the evidence with mean vāu-10 and covariance 1
Derivation of p(sT h). Analogously to the derivation of p(s, 8t, h), we establish that
f p(h so) p(so ST) dag o (N v V0r ay.
Then apply Lemma 6 to distributions
p(ST)= Ld) N(S:: vārO, Er
and get that
p(T h) x (ST; (h), ET+1(1))
where ĀT+ (h) and ΣT (h) defined in the claim. This is product of Gaussians: the prior with mean Ou and covariance Id. and the evidence with mean Vār and covariance āzy E.
A.3 Proof of Theorem
We start with the triangle inequality
||0 0.112 110 - + 0 00112 0 - 0-112.
where we introduce from Section 2. Now we bound each term on the right-hand side.
Upper bound on 0 0||2- This part of the proof is based on analyzing the asymptotic behavior of the conditional densities in Theorem 2.
As first step, note that Sp (Npr+ (h), (h)()),,, where
juT+1 (h) = ( 02 + (///))))) (1+1(1)= = (-rr)-
Since Aa(E + Do, we get
(1) 5, ( Voro.
Moreover, >a(E Aa(E-1) oo implies A1(2) + O. and thus lim || v 1a
The same argument can be applied inductively to later stages of the reverse process. Specifically, for any E [T], S ~ (S,, h). (h)). where
(.,S,, h) E,( p4 (S,) Σ V). E,(h) = (E, + E-1/01111-
Since Xx(E-1) and S, v by induction, get
(h) 1. A (S., h1) vār-10.
Moreover, >a(551)~ - oo implies > (5.) - 0, and thus lim ||St-1 - /00-10000 = for any t E [T]. In the last stage, / 1, āo 1, and So 0. Therefore,
lim S - Va 1 lim ||0 0|1 = 0.
N N
N
Upper bound on ||0 0. 1|2- This part of the proof uses the definition of in Section 2.1 and that Ee (O, o²) is independent noise. By definition,
0- 0, = -
Beei 0, = - Σ
sel0 0, + tx) 0, - -
grer
im1
15

---

Since Et is independent zero-mean Gaussian noise with variance o², 0 02 is a Gaussian random variable with mean Og and covariance
Σ var [Eel 08 5 --- -Z=1 0
Cov
o E
DEEt
E E.
r²
Since Aa(E-1) - oo implies A1 (E) - O, we get
lim 10 1|2 O.
N -00
This completes the proof.
A.4 Proof of Theorem
This proof has two parts.
Derivation of p(s st, ht. From (6), we have
[ p(h so) p(so st dsg x p(h v011/(1-11.
Since p(8 8x) is Gaussian, we have
p(s st) (s)), Σ,) N(19)-1571((t)))
yory= = 1/V0a-1. Then by the Laplace approximation,
p(h ((78 Ex) x (N7s Ex) (NooiNN((sss (,,,;
where 8. E, + (R))
Derivation of p(sr h). Analogously to the derivation of p(s st, h), we establish that
[ p(h so) p(sa! ST) dso x p(h ST sT/Vär))
Then by the Laplace approximation for y Vōr, get phh (3T) (2, & ((ST: ET41/72), where 07+1, ET+1 70,
A.5
Supporting Lemmas
We state and prove our supplementary lemmas next.
Lemma 5. Let p(x) NN;;; pI, Σn) and q(x) NN u2, E2), where u1,12 E Rd and Σ1, Ea Rdxd Then
1 (w2 - (41 T (w2 - p1) tr(2 En) log det(E1) d) )) d(p.q) 2
det(
Moreover. when Σ1 Eg.
1 [ - 11) 2 (#2 m1) d(p, q) -1:55-10
Proof. The proof follows from the definitions of KL divergence and multivariate Gaussians.
Lemma 6. Fix #1 E Rd, Σj ž 0, E Rdd and Σg Z 0. Then
Ei Ez) x Nx;
where
>== 4n ES p2), == == ΣS 1
16

---

Proof. This is a classic result, which is proved as
Nax; p1, [x) (2: Σg) x exp
2 ( - 141) /1)"/1 (x M1) (a 12) I (x - p2))
x exp
(x Ej x - 2x ΣΙ
+ x Eng x - 2x T E u2)
2 [ x - 2x Σ I in + Enj EnnN²2)
exp
2 (a n)) "
- (x - u)) & N(G,,
x exp
The neglected factors depend on constants μ4,, M2, E,, and Σz. This completes the proof.
B Learning the Reverse Process
One property of our model is that ((8T) x MsT: Oa: Ta) when T is sufficiently large [19]. Since ST has the same distribution in the reverse process p. can be learned from the forward process q by simply reversing it. This is done as follows. Using the definition of the forward process in (4), Ho al. [19] showed that
(86-1 18,, 18 so), BeTa)
(13)
holds for any and ss. where
(e((s:: s0) = 1 āe
80 + vail -
0;-1 3r: 0. II ag (14)
3, = 1 ā St
1 C₁
Therefore, the latent variable in stage - 1, S 1, is easy to sample when 8 and sy are known. To estimate So, which is unknown when sampling from the reverse process, again. particular, (4) implies that 8; = V + vI - ōeEt, where Eg ~ NO,,, Ia) is a standard Gaussian noise. This identity can be rearranged as
the forward process
80
(0t - v1 a.st).
Vat
To obtain Es, which is unknown when sampling from p. learn regress it from [19].
The regressor is learned as follows. Letex(;;) v) be a regressor of e parameterized by y and {8o} be a dataset of training examples. We sample so uniformly at random from D and then solve
1 = arg min Eg [|lee Ee(S 1)
(15)
per stage. The expectation approximated by sampled 8y- Note that we slightly depart from Ho et al. [19]. Since each regressor has its own parameters, the original optimization problem over T stages decomposes into T subproblems.
C Additional Experiments
This section contains four additional experiments.
C.1 Additional Synthetic Problems
In Section 6.2, we show results for three hand-selected problems out of six. We report results on the other three problems in Figure 4, We observe the same trends in Section 6.2.
C.2 MNIST Experiment
The next experiment on the MNIST dataset [31]. We start with learning an MI.P-bsse multi-way classifier for digits and extract their d dimensional embeddings. These used features in
17

---

Problem swirl
Problem
Problem coreers
4
Thue and difiusiion prtors
Thue and alffusion prfors
diffusion
-1
-1
Problem swirl
Problem
Problem corners
TS
TunedTS MIXTS
TunedTS
TunedTs
MixTS
MixTS
Regret
15
DPS
Regret
DPS
Rogret
OP
DIFES
DIFTS
10
10
100
5
5
5
50
100
150
200
100
150
200
50
100
150
200
Round
Round
Round
Figure 4: Evaluation of DiffTS on another three synthetic problems. The first row shows samples from the true (blue) and diffusion model (red) priors. The second row shows the regret of DiffTS and the baselines function of round
(a) Prablom MNIST
(b) MNIST lineor bandit
(c) MNIST logistic bandit
10
25
20
T5
15
20
10
MixTS
10
TunedTs
6
10
100
200
300
400
500
100
200
300
400
500
Round
Round n
Figure 5: Evaluation of DiffTS on the MNIST dataset: (a) shows samples from the true (blue) and diffusion model (red) priors, (b) shows regret in the linear bandit, and (c) shows regret in the logistic bandit.
our experiment. We generate a distribution over model parameters 0, as follows: (1) we choose random positive label, assign it reward 1, and assign reward -1 to all other labels; (2) we subsample a random dataset of size 20, with 50% positive and 50% negative labels; (3) we train linear model, which gives us single 0_. We repeat this 10000 times and get distribution over 0..
We consider both linear and logistic bandits. both, the model parameter 0, initially sampled from the prior. In each round, K 10 random actions are chosen randomly from all digits. In the linear bandit, the mean reward for a digit with embedding x is x 0-, and the reward noise is 1. In the logistic bandit, the mean reward is g( 0_), where g is sigmoid.
Our MNIST results reported in Figure 5. We observe again that DiffTS has lower regret than all baselines, because the learned prior captures the underlying distribution of 0, well. We note that both the prior and diffusion prior distributions exhibit strong cluster structure (Figure 5a), where each cluster represents one label.
C.3 Ablation Studies
We conduct three ablation studies on the problem in Figure 2.
In experiments, the number of samples for training diffusion priors was 000. Figure 6a, we vary it from 100 to 10 000. We observe that the regret decreases the number of samples increases.
18

---

24
13
0.0
100.0
750
350
Thatiing Samplie
DiPlution Stepe.
09%uddon Stecs,
Figure 6: An ablation study of DiffTS on the cross problem: (a) we vary the number samples for training the diffusion prior and report regret, (b) we vary the number of diffusion stages 7 and report regret, and (c) we vary the number of diffusion stages T and report computation time.
Problem
Problem
Problem
True and diftusion priors
Thue and diftusion priors
Thue and diruuson proors
-2
-2
-2
-3
2
-3
-2
3
-3
-2
-1
Problem
Problem
Problem
1.0
1.0
1.0
0.8
TunedTS
0.8
TuneaTS
0.8
TunedTs
MixTS
MiXTS
MixTS
0.6
0.6
0.6
DIfTTS
DIfTS
DiffTS
0.4
SMC
SMC
0.4
SMG
0.2
0.2
0.0
0.0
10"
10a
107
10a
10a
10a
10'
10a
10a
10a
10'
Sample
Sample
Figure 7: Evaluation Gaussian mixture variants of the synthetic problems in Figure 2. The first row shows samples from the true (blue) and diffusion model (red) priors. The second row shows the earth mover's distance of DiffTS and baseline posteriors from the true posterior as a function of sample size n.
due to learning a better prior approximation. The trend stabilizes around 3 000 training samples, We conclude that the quality of the learned prior approximation has a major impact on regret.
In all experiments, the number diffusion stages was T = 100. In Figure 6b, we vary from 1 to 300 and observe its impact regret. While the regret at T increases. It stabilizes around T 100, which we used in our experiments. In Figure 6c, we vary T from 1 to 300 and observe its effect the computation time of posterior sampling. The time is linear in T, as suggested in Section 4.2. The main contributor the neural network regressor.
is high, decreases quickly as T
C.4 Non-Bandit Evaluation
We use Gaussian mixture variants of the synthetic problems in Figure 2 for our non- bandit evaluation. The action in round i is chosen uniformly random (not adaptively). Since the priors are Gaussian mixtures, the true posterior distribution can be computed in a closed form using MixTS and we can measure the distance posterior approximations from We use the earth mover's distance (EMD) between posterior samples from the true posterior and its approximation. We also considered the KL. divergence. However, we could not apply because the posteriors of DiffTS and DPS do have analytical forms.
19

---

Algorithm 4 DPS of Chung et al. [12].
1: Input: Model parameters o₂ and St
2: Initial sample ST (Oa. Ta)
3: for stage T..... , do
4: $ - 1-
(S.ivl)))
5:
$o + (St ( + (1 - (1-(a)))
6: Z NNa.. La)
7:
St-1 + vār- Su + vaill-āe-)) I-
L St + -ō &V Ei=(ie (ye - 0, So)²
l-ar
8: Output: Posterior sample Sp
We evaluate all methods from Figure 2. In addition, we implement a sequential Monte Carlo (SMC) sampler [16]. The initial particles are sampled uniformly at random from the prior. At each round, the particles are perturbed by a Gaussian noise. The standard deviation of the noise is initialized as a fraction of the observation noise and decays are weighted according to the likelihood of the observation in the round. Finally, we use normalized likelihood weights to resample the particles. The number of particles is 3 000 and we tune SMC to get good posterior approximations. The computational cost of SMC is comparable to DiffTS.
time, the posterior concentrates. The particles
Our results are reported in Figure 7. We observe that DiffTS approximations are comparable to MixTS, which has an exact posterior this setting. The second best performing method is SMC. Its approximations worsen as the sample size n increases. DPS approximations also get worse increases, which caused instability in Figure
D Implementation of Chung et al. [12]
In our experiments, compare diffusion posterior sampling (DPS) with Gaussian observation noise (Algorithm in Chung et al. [12]). Our implementation is presented in Algorithm 4. The score is = (S; //) āt), where (S,; the) is regression estimate of the forward process noise e, in Appendix B. We set ō, / 3. which is the same amount of noise in our reverse process (Section 3). The term
N UE o, S0)
the gradient of the negative log-likelihood with respect to S..
As discussed in Appendices C.2 and D.1 of Chung et al. [12]. Ç. in DPS needs to be tuned for good performance. This is because V S (yee - of $o)² grows with the number of observations, which f-
causes instability. We also observed this in our experiments (Section 6.2). To make DPS work well, we follow Chung et al. [12] and set
1
S1
E N (ye 0f So)²
While this significantly improves the performance of DPS, does not prevent failures. The funda- mental problem is that maadient-based optimization is sensitive the step size, especially when the optimized function is steep. Note that LaplaceDPS does not have any such hyper-parameter.
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state all contributions. The introduction also points to where those contributions are made,
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The increase in computational cost is discussed in Section 4.2 and shown empirically in Section 6.2. We also conduct an ablation study in Appendix C.3, where we show how the regret of DiffTS scales with the number of samples used for pre-training the prior and the number of stages in the diffusion model prior.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: The main claims are stated and discussed in Section 4. Their proofs are in Appendix A.
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 it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We also include code to reproduce the synthetic results in Figures and
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?
Answer: [Yes]
Justification: We include code to reproduce the synthetic results in Figures and 4.
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]
The experiments are described to sufficient level to be reproducible. To make sure, we include code to reproduce the synthetic results in Figures and 4.
Justification:
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All plots in the paper have error bars.
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 reproduce the experiments?
Question:
21

---

Answer: [No]
Justification: Our experiments are not large scale.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: htt.rpeurpppppp/ httppccccc/sttttcccccoe.........
Answer: [Yes]
Justification: We checked the link and comply.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is algorithmic and not tied to particular application that would have immediate negative impact.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper does not pose such risk.
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 the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All used assets are stated and cited.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: crowdsourcing or research with human subjects.
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No crowdsourcing or research with human subjects.
22


Reviews:
[{"review_id": "UwLm8jncup", "reviewer": "Reviewer_Vag9", "summary": "The paper introduces new posterior sampling approximations for contextual bandits with diffusion model priors, allowing the capability to handle complex distributions beyond traditional Gaussian priors. It contributes by developing efficient sampling algorithms, proving their asymptotic consistency, and validating their effectiveness on synthetic and empirical contextual bandit problems.", "questions": "As mentioned in the \"Weakness\" section, a somewhat detailed summary that clarifies the differences between this work and related works, particularly the work of Hsieh et al. [22], would help understand the unique contributions of this paper.\n\nClarity would be enhanced by more explanations on the roles of the reverse process and forward process in posterior sampling.", "limitations": "I recommend that the authors include a more detailed discussion of the limitations of this work.", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is overall well-written, providing a clear background and enhancing accessibility for a broad audience. It advances prior work by extending Thompson sampling with a diffusion model prior from K-armed bandits to a more general setting of contextual bandits, which broadens the application scope. The theoretical claims, such as asymptotic consistency, are supported by sound arguments. The experiment are also comprehensive. Furthermore, the authors also mentioned the potential in extensions beyond GLM, suggesting broader applications.", "weaknesses": "The paper briefly mentions extending the work of Hsieh et al. [22], who proposed Thompson sampling with a diffusion model prior for K-armed bandits, to contextual bandits. However, the discussion lacks depth and specificity. A more detailed comparison is needed to highlight the unique contributions of this paper. For example, the difference between the posterior sampling approximations, and the difference between the theoretical analysis. \n\nThe experiment sections (Figures 4 and 5) lack explanations on the metrics, which are crucial for understanding the variability and reliability of the results. It should be clear whether the error bar is the standard deviation or the standard error of the mean.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for detailed and positive feedback. Our rebuttal is below. If you have any additional concerns, please reach out to us to discuss them.\n\n**Differences from Hsieh et al. (2023)**\n\nThere are multiple differences:\n\n* The posterior approximation in Hsieh et al. (2023) is for scalars (individual arms). Our approximation is for vectors (model parameters).\n* The approximations are different. In stage $t$, Hsieh et al. (2023) sample from the conditional prior and the diffused empirical mean distribution in stage $t$. Then they take a weighted average of the samples. We sample only once, from the posterior distribution obtained by combining the conditional prior in stage $t$ and likelihood. Based on this, Hsieh et al. (2023) can be viewed as a non-contextual variant of our method, where posterior sampling is done by weighting samples from the prior and empirical distributions.\n* Hsieh et al. (2023) do not analyze their approximation.\n\n**Metrics in Figures 4 and 5**\n\nThe regret is defined as in Figures 2 and 3 (Section 6). The mathematical definition is given in line 213 (Section 5). All error bars are standard errors.\n\n**Limitations of our method**\n\nPlease see the **common rebuttal**.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "agCxtKg820", "reviewer": "Reviewer_r4r4", "summary": "The authors propose an algorithm for sampling from a generalised linear model posterior where the prior is defined through a diffusion model. This is achieved by utilizing the Laplace approximation, and is shown to be asymptotically consistent. This model and inference scheme is then applied to contextual bandits, where their performance is demonstrated on a variety of synthetic and real world applications.", "questions": "1)\tIn Thm 2 equation 6 reads $p(h | s_t) = E_{p(s_0 | s_t}[p(h|s_0)] \\approx p(h | s_t / sqrt(\\alpha_t))$. Which seems like the assumption is $s_t \\approx  s_t / sqrt(\\alpha_t)$, why do you want to make this assumption and why is justified?\n\n2)\tWhy does $\\nabla log p(h|\\theta)$ grow linearly in N? \n\n3)\tDoes relying on the score of likelihood only impact the application for contextual bandits or does this also impact these methods generally?\n\n4)\tIs there any error caused/limitation in using the Laplace approximation?", "limitations": "They discuss that they have performed regret analysis. However they do discuss limitations of the proposed sampling algorithm.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "I am not familiar with Bandits or Diffusion models but the paper is interesting and the experiments in the context of contextual bandits seem convincing.", "weaknesses": "1)\tThis paper reads like a paper of two halfs. The first half is proposing a new sampling algorithm, and the second is applying this to contextual bandits. I would have like to have seen more evidence / discussion on how well this algorithm works in general for diffusion models.  For example, do you know how well your method works on the experimental setup of Chung et al [12] ?\n\n2)\tThe math at times can be a bit difficult to parse (however this could be due to not being familiar with this field). Minor things like using bold symbols for vectors/matrices may help here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for valuable feedback. Our rebuttal is below. If you have any additional concerns, please reach out to us to discuss them.\n\n**Experimental setup of Chung et al. (2023)**\n\nChung et al. (2023) experiment with computer vision problems. Their algorithm is unstable in these problems without tuning, and they discuss it in their Appendices C.2 and D.1. We focused on online learning as the first step because a suboptimal model of uncertainty, even after tuning, is unlikely to perform well. See the performance of DPS in Figures 2 and 4, and our discussion in Appendix D. We plan to experiment with vision and video models in our future work.\n\n**Q1: Approximation (6) in Theorem 2**\n\nWe assume that $s_0 = s_t / \\sqrt{\\bar{\\alpha}_t}$, where $s_0$ is a clean sample and $s_t$ is the corresponding diffused sample in stage $t$. This approximation is motivated by the observation that under the forward process, $s_t = \\sqrt{\\bar{\\alpha}_t} s_0 + \\sqrt{1 - \\bar{\\alpha}_t} \\tilde{\\varepsilon}_t$ for any $s_0$, where $\\tilde{\\varepsilon}_t \\sim \\mathcal{N}(\\mathbf{0}_d, I_d)$ is a standard Gaussian noise. See lines 166-170 in Section 4.3. The result of our approximation is that the likelihood becomes a function of scaled $s_t$, and can be easily combined with the conditional prior distribution, which is also a function of $s_t$.\n\n**Q2: Why does $\\nabla \\log p(h | \\theta)$ grow linearly in $N$?**\n\nBecause the history $h$ involves $N$ observations. As an example, in Assumption 1, the likelihood is $p(h \\mid \\theta) \\propto \\exp[- \\sigma^{-2} \\sum_{\\ell = 1}^N (y_\\ell - \\phi_\\ell^T \\theta)^2]$.\n\n**Q3: Is the score issue general?**\n\nYes. This is a general issue of relying on $\\nabla \\log p(h | \\theta)$ when $N$ is large. Chung et al. (2023) tune their gradient step to mitigate this. See our Appendix D for more details.\n\n**Q4: Does the Laplace approximation have an error?**\n\nYes. We discuss this in lines 166-175 (Section 4.3). The good news is that the error vanishes as the number of observations increases. This is what we prove in Theorem 3.\n\n**Limitations of our method**\n\nPlease see the **common rebuttal**.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "e06a3YX3PM", "reviewer": "Reviewer_YMJY", "summary": "The paper presents approximate posterior sampling methods for contextual bandits with a diffusion prior. A key weakness of existing methods designed to work on noisy data is that rely on using the score function is that it becomes unstable as the number of observations grows. Instead, the authors propose to use closed-form analytic updates to merge prior and evidence distributions in the diffusion chain. This works for linear models and can be extended to generalized linear models with the Laplace approximation. This framework is applied in the context of contextual bandits for synthetic and real datasets and performs better than other diffusion model alternatives.", "questions": "In Algorithm 2, the mean and covariance functions for all S_{0:T} are functions of the history which can be of variable size, how does the network take this aspect into account?", "limitations": "The idea seems quite general but was only considered in a contextual bandits setting, is there a reason why you couldn't compare the samples in a non-bandit setting to other methods e.g. VI, SMC, MCMC?", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The problem of Thompson sampling with complex multi-modal distribution is a significant one with a lot of impact. The use of identifying (or approximating with) Gaussians to enable closed-form updating in a diffusion setting is very compelling and results in a simple solution that seems to work well in contextual bandits. The Laplace approximation is a natural choice to move beyond simple linear model settings. The writing and development of the method and theorems was mostly well executed.", "weaknesses": "A few notes about correctness/presentation:\n(1) it is not true that the posterior of \\theta_* is a product of two multivariate Gaussians (as stated below Eq 3) it is instead proportional to that product.\n(2) there seems to be a non-sequitur on L82-83 when Laplace approximation is introduced: Laplace does not require the prior to Gaussian, it approximates the posterior as a Gaussian.\n\nThe introduction of IRLS with corresponding algorithm in Sec 2.2. seemed a bit out of place and I wonder if it could not be delayed until a little bit later in the method section (in part because it is not explained why one couldn't just use autodiff to find the MAP solution as is now standard). \n\nIt was not clear which of the proofs are standard results and which are novel (e.g. Lemma 1 looks like Bayes theorem).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for detailed feedback, and praising our contributions and execution. Our rebuttal is below. If you have any additional concerns, please reach out to us to discuss them.\n\n**Some claims are imprecise**\n\nThe reviewer is right that the posterior in (3) is only proportional to the product of the prior and likelihood. They are also right that the Laplace approximation does not require a Gaussian prior, although this is a very common setting.\n\n**Novelty in proofs**\n\nPlease see the **common rebuttal**.\n\n**How do the mean and covariance in Algorithm 2 change with history?**\n\nThey are weighted sums of the prior quantities, which are represented by a neural network and independent of history $h$, and empirical quantities. As an example, $\\hat{\\Sigma}_t(h)$ in line 3 is computed in (8) in Theorem 2. It is the inverse of a weighted sum of the prior $\\Sigma_t^{-1}$ and empirical $\\bar{\\Sigma}^{-1}$ precisions. The latter is defined right after (3) in Section 2.1.\n\n**Non-bandit evaluation**\n\nWe conduct an empirical evaluation in the non-bandit setting in the pdf attached to the **common rebuttal**. Please see it for more details.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u90GqTW8D7", "reviewer": "Reviewer_1tRW", "summary": "This paper introduces novel posterior sampling approximations tailored for diffusion model priors, specifically designed for use in contextual bandits and applicable to a broader range of online learning problems. The methods are developed for linear models and generalized linear models (GLMs), emphasizing enhanced stability and efficiency in environments where previous algorithms may exhibit instability and divergence. The paper provides the asymptotic consistency of these approximations. Empirically, the performance of these approximations is evaluated on contextual bandit problems, showcasing their capability to manage uncertainties inherent in these settings effectively.", "questions": "Can the author explain why the diffusion model prior solves the instability and divergence issues?\n\nWhat other benefits could the diffusion model prior bring?\n\nWhat other approaches could be employed beyond Laplace approximation?\n\nWhat are the technical difficulties in the paper, especially the analysis?", "limitations": "Discussion on limitations seems to be missing or need to be more explicit.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "This paper proposes approximate posterior sampling algorithms for contextual bandits with a diffusion model prior. Unlike previous research using a Gaussian prior, it addresses the instability and divergence issues. This paper has a well-structured presentation and elucidates background knowledge such as linear models, GLM, and diffusion models. Claims are supported by either theoretical proofs and numerical studies. The proofs are technically sound based on the reviewer's examination.", "weaknesses": "It seems to the reviewer the key idea in this paper is to use the diffusion model instead of a Gaussian distribution as a prior. The proofs seem to be standard and similar to those based on a Gaussian prior. The main weaknesses of the paper are its lack of novelty and the absence of significant technical challenges.\n\nSome format issues: p2 l49, p4 Fig. 1, l121, and l118.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for valuable feedback and praising our execution. Our rebuttal is below. If you have any additional concerns, please reach out to us to discuss them.\n\n**Q1: Score issue in posterior sampling**\n\nThere may be a misunderstanding. The diffusion prior does not solve any instability. We address previous issues with posterior sampling and a diffusion model prior.\n\nPrior works on posterior sampling in diffusion models sampled using the score of the posterior probability (Section 7). The score depends on the gradient of the log-likelihood $\\nabla \\log p(h | \\theta)$, which grows linearly with the number of observations $N$ and causes instability. We combine $p(h | \\theta)$ with the conditional prior, in each stage of the diffusion model, using the Laplace approximation. The resulting conditional posterior concentrates at a single point as $N \\to \\infty$ and is easy to sample from, although its gradient goes to infinity. Our main technical contribution is an efficient and asymptotically consistent implementation of this solution, using a stage-wise Laplace approximation with diffused evidence.\n\n**Q2: Benefit of diffusion model priors**\n\nThe main benefit of diffusion models is that they can learn complex prior distributions from data. These distributions then serve as representations of prior knowledge.\n\n**Q3: Other approaches to posterior sampling**\n\nThe most popular approach is DPS of Chung et al. (2023). The key idea in DPS is to sample from a diffusion model posterior by adding the score of the likelihood of observations to the diffusion model prior. We describe this approach in Appendix D and compare to it empirically in Section 6. We also mention several other approaches in Section 7. All of them rely on the score of the likelihood $\\nabla \\log p(h \\mid \\theta)$ and thus become unstable as the number of observations $N$ increases.\n\n**Q4: Technical challenges**\n\nWe do not just replace a Gaussian prior with a diffusion model prior in classic posterior sampling. We propose posterior sampling with a diffusion model prior that can be implemented efficiently and asymptotically consistently using a stage-wise Laplace approximation with diffused evidence. Please see the **common rebuttal** for more details on technical challenges.\n\n**Limitations of our method**\n\nPlease see the **common rebuttal**.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "I85fftkZ9W", "reviewer": "Reviewer_AZGY", "summary": "The paper studies online learning and proposes to approximate the updating prior with a diffusion model, rather than the more simple, less expressive Gaussian approximation. The main application is contextual bandits with a linear or GLM model. For the latter case, the authors derive a version of the IRLS algorithm (based on a Laplace approximation) that works with the diffusion prior. The authors also demonstrate the consistency of their approximating algorithm. The paper examines a variety of examples, and compares the proposed model to existing methods for contextual bandit.", "questions": "If I understand correctly, computing the regret requires knowing the optimal $a_*$ and $\\theta_*$. Are there evaluations which can be performed without oracle knowledge? What about diagnostics to assess the quality of the underlying Laplace approximation?", "limitations": "The authors address some of the limitations. See my comments above for additional limitations that should be addressed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The idea to extend the IRLS algorithm used for an updating Gaussian prior to a more sophisticated prior is well motivated and seems useful. The paper presents an actionable algorithm and the theoretical results are convincing. There are also some empirical results---both on toy examples and a benchmark data set---, which demonstrate the utility of the diffusion prior.\n\nI also found the paper to be well written: it provides extensive context for existing approaches, notably the Laplace approximation and the IRLS algorithm. The proposed DiffTS then seems like a natural next step.", "weaknesses": "I picked up a few points, but I believe most of them can be addressed in the rebuttal.\n\nI'm not entirely comfortable with the statements of Theorem 2 and 4, because they involve \"$\\approx$\", which is not formally defined. For instance, a reader might wonder if the approximation in (6) is \"as good\" as the one on line 148... I understand what the authors mean, but formal statements need to be more precise. I recommend, as a fix, writing \"Suppose that\" and then equation (6) with an equality sign. The authors can then replace all the \"$\\approx$\" with an \"=\", and state that the initial assumption does not hold in general; assuming it does is where the approximation comes into play.\n\nIn the experiments, the authors report the regret against round $n$. One thing that's not accounted for is that DiffTS requires more computation than its benchmarks. How much of a concern is this in practice? Does the added computation slow down model training or is the wait time dominated by generating a new round? Some discussion around this, potentially with an additional figure that reports regret vs time, could be helpful.\n\nIt was also not clear to me what where the tuning parameters of DiffTS. The authors claim it inherits the simplicity and efficiency of the Gaussian prior, however, determining T and picking a neural network architecture strike me as complications that limit the off-the-shelf use of the algorithms. I believe the authors could be more upfront about this; the paper would also be stronger with examination of those tuning choices across examples.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for positive feedback, which recognizes multiple contributions of our work. Our rebuttal is below. If you have any additional concerns, please reach out to us to discuss them.\n\n**Avoid $\\approx$ in Theorems 2 and 4**\n\nA great comment. We agree that it is cleaner to state (6) as an assumption and then replace all $\\approx$ with $=$.\n\n**Computation time and regret trade-offs of DiffTS**\n\nIn our experiments, posterior sampling in DiffTS with $T$ stages is about $T$ times more computationally costly than posterior sampling with a Gaussian prior. We discuss this in lines 267-272 (Section 6.2). We plot the regret and sampling time as a function of $T$ in Figures 6b and 6c (Appendix C.3), respectively.\n\n**Diffusion model tuning**\n\nWe have not done much tuning. In all experiments, the number of diffusion stages is $T = 100$ and the diffusion rate is set such that most of the signal diffuses. The regressor is a $2$-layer neural network and we learn it from $10000$ samples from the prior. These settings resulted in stable performance across all experiments. We plot the regret as a function of the number of training samples and $T$ in Figures 6a and 6b (Appendix C.3), respectively. When $T$ or the number of training samples is small, DiffTS performs similarly to posterior sampling with a Gaussian prior.\n\n**Non-bandit evaluation**\n\nWe conduct an empirical evaluation in the non-bandit setting in the pdf attached to the **common rebuttal**. Please see it for more details.\n\n**Limitations of our method**\n\nPlease see the **common rebuttal**.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We wanted to thank all reviewers for positive reviews and recognizing our contributions. There were three common comments that we want to address jointly: limitations of the work, technical challenges, and a non-bandit evaluation.

**Limitations of our method**

We point out limitations of our work throughout the paper. However, the reviewers could not always find them. To address this issue, we plan to have a dedicated paragraph for limitations in Conclusions. The main limitations of our approach are:

* **Computational cost:** Posterior sampling in DiffTS with $T$ stages is about $T$ times more computationally costly than posterior sampling with a Gaussian prior. We say this in lines 162-164 (Section 4.2) and validate it experimentally in lines 267-272 (Section 6.2). We plot the sampling time as a function of $T$ in Figure 6c (Appendix C.3).
* **Diffusion model prior tuning**: In all experiments, the number of diffusion stages is $T = 100$ and the diffusion rate is set such that most of the signal diffuses. The regressor is a $2$-layer neural network and we learn it from $10000$ samples from the prior. These settings resulted in stable performance across all experiments without tuning. We plot the regret as a function of the number of training samples and $T$ in Figures 6a and 6b (Appendix C.3), respectively. When $T$ or the number of training samples is small, DiffTS performs similarly to posterior sampling with a Gaussian prior.

**Technical challenges**

Prior works on posterior sampling in diffusion models sampled using the score of the posterior probability (Section 7). The score depends on the gradient of the log-likelihood $\nabla \log p(h | \theta)$, which grows linearly with the number of observations $N$ and causes instability. We combine $p(h | \theta)$ with the conditional prior, in each stage of the diffusion model, using the Laplace approximation. The resulting conditional posterior concentrates at a single point as $N \to \infty$ and is easy to sample from, although its gradient goes to infinity. Our main technical contribution is an efficient and asymptotically consistent implementation of this solution, using a stage-wise Laplace approximation with diffused evidence.

Lemma 1 can be derived using basic rules of probability and we state it in Appendix A.1. All other claims, including efficient posterior derivations (Theorems 2 and 4) and asymptotic consistency (Theorem 3), rely on our proposed approximation of clean samples by scaled diffused samples. The most challening part of the analysis is Theorem 3, which analyzes an asymptotic behavior of a chain of $T$ random vectors, which depend on each other.

**Non-bandit evaluation**

We use Gaussian mixture variants of the synthetic problems in Figure 2 for our non-bandit evaluation. The action in round $k$ is chosen uniformly at random (not adaptively). Since the priors are Gaussian mixtures, the true posterior distribution can be computed in a closed form using MixTS and we can measure the distance of posterior approximations from it. We use the *earth mover's distance (EMD)* between posterior samples from the true posterior and its approximation. We also considered KL divergence, but this one required analytical forms of posterior approximations, which are not available in DiffTS and DPS.

We compare all methods from Figure 2. In addition, we implemented a sequential Monte Carlo (SMC) sampler. The initial particles are chosen uniformly at random from the prior. At each next round, the particles are perturbed by a Gaussian noise. The standard deviation of the noise is initialized as a fraction of the observation noise and decays over time, as the posterior concentrates. The particles are weighted according to the likelihood of the observation in the round. Finally, normalized likelihood weights are used to resample the particles. We tune SMC to get the best approximations. We use $3000$ particles. For this setting, the posterior sampling times of SMC and DiffTS are comparable.

Results of our experiments are reported in Figure 7 in the attached pdf. We observe that the quality of DiffTS approximations is similar to MixTS, which has an exact posterior in this setting. The second best performing method is SMC. The quality of its approximations worsens as the sample size $n$ increases. The quality of DPS approximations also worsens as $n$ increases, which caused instability in Figure 2.

Meta-review: This paper has borderline scores. The reviewers all emphasize good-quality writing. The main weakness noted by the lowest-scoring review is that the techniques used are relatively standard, especially for a paper which is more computational / empirical / similar in nature, compared to other papers in this area that are doing pure theorem-proof-based regret analysis. Still, even that reviewer writes that the work is technically sound in terms of the numerics, which are arguably one of the more important aspects in a paper of this kind. On basis of this thinking, and discussion, I am inclined to accept.
Decision: Accept (poster)