Online Posterior Sampling with a Diffusion Prior

Authors: branislav  kveton
adobe  research
boris  oreshkin
youngsuk  park
aniket  deshmukh
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 22:36:47

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.

Main paper: 
 introduction: 
 A multi-armed bandit [26, 6, 29] is an online learning problem where an 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-exploitation dilemma: explore actions to learn about them or exploit the action with the highest estimated reward. Bandits have been successfully applied to problems where uncertainty modeling and adaptation to it are beneficial, such as in recommender systems [31, 53, 24, 34] and hyper-parameter optimization [33] .Contextual bandits [28, 31] with linear [13, 1] and generalized linear models (GLMs) [16, 32, 2, 25] have become popular due to the their flexibility and efficiency. The features in these models can be hand-crafted or learned from historic data [39] , and the models can be also updated incrementally [1, 23] . While the original algorithms for linear and GLM bandits were based on upper confidence bounds (UCBs) [13, 1, 16] , Thompson sampling (TS) is more popular in practice [11, 3, 41, 43] . The key idea in TS is to explore by sampling from the posterior distribution of model parameter θ * . TS uses the prior knowledge about θ * to speed up exploration [11, 39, 35, 9, 20, 19, 5] . When the prior is a multivariate Gaussian, the posterior of θ * can be updated and sampled from efficiently [11] . This prior has a limited expressive power, because it cannot even represent multimodal distributions. To address this, we study posterior sampling with a diffusion prior. The main benefit of such priors is that they can represent complex distributions and be learned from data.We make the following contributions. First, we propose novel posterior sampling approximations for linear models and GLMs 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 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 a Laplace approximation, which mixes prior knowledge and evidence (Theorem 4). Our approximations are motivated by posterior sampling with a Gaussian prior, and inherit its simplicity and efficiency. In prior works (Section 7), posterior sampling is implemented using the likelihood score, which tends to infinity as the number of observations increases; and therefore causes instability. We combine the likelihood with the conditional prior, in each stage of the diffusion model, using the Laplace approximation. The resulting posterior concentrates at a single point and can be easily sampled from, even if the likelihood score tends to infinity. We propose an efficient and asymptotically consistent implementation of this idea, using a stage-wise Laplace approximation with diffused evidence.Our second contribution is in theory. We properly derive our posterior approximations (Theorems 2 and 4) and show their asymptotic consistency (Theorem 3). The key idea in the proof of Theorem 3 is that the conditional posteriors concentrate at a scaled unknown model parameter as the number 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 have this property. All of our main results rely on our 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.Our last contribution is an empirical evaluation on contextual bandits. We focus on bandits because the ability to represent all levels of uncertainty precisely is critical for exploration. Our experiments show that a score-based method fails to do so (Section 6.2). Note that our posterior approximations are general and not restricted to bandits. 
setting: 
 We start with introducing our notation. Random variables are capitalized, except for Greek letters like θ. 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(x | y). We denote by X n:m and x n:m a collection of random variables and their values, respectively. For a positive integer n, we define [n] = {1, . . . , n}. The indicator function is 1{•}. The i-th entry of vector v is v i . If the vector is already indexed, such as v j , we write v j,i . We denote the maximum and minimum eigenvalues of matrix M ∈ R d×d by λ 1 (M ) and λ d (M ), respectively.The posterior sampling problem can be formalized as follows. Let θ * ∈ Θ be an unknown model parameter and Θ ⊆ R d be the space of model parameters. Let h = {(ϕ ℓ , y ℓ )} ℓ∈[N ] be the history of N noisy observations of θ * , where ϕ ℓ ∈ R d is the feature vector for y ℓ ∈ R. We assume thatEQUATIONwhere g : R → R is the mean function and ε ℓ is an independent zero-mean σ 2 -sub-Gaussian noise for σ > 0. Let p(h | θ * ) be the likelihood of observations in history h under model parameter θ * and p(θ * ) be its prior probability. By Bayes' rule, the posterior distribution of θ * given h isEQUATIONWe want to sample from p(• | h) efficiently when the prior distribution is represented by a diffusion model. As a stepping stone, we review existing posterior formulas for multivariate Gaussian priors. This motivates our solution for diffusion model priors.The posterior of θ * in linear models can be derived as follows.Assumption 1. Let g in (1) be an identity and ε ℓ ∼ N (0, σ 2 ). Then the likelihood of h under modelparameter θ * is p(h | θ * ) ∝ exp[- N ℓ=1 (y ℓ -ϕ ⊤ ℓ θ * ) 2 /(2σ 2 )].Let p(θ * ) = N (θ * ; θ 0 , Σ 0 ) be the prior distribution, where θ 0 ∈ R d and Σ 0 ∈ R d×d are the prior mean and covariance, respectively. Then p(θ * | h) ∝ N (θ * ; θ, Σ) [10] , whereEQUATIONθ = σ -2 Σ N ℓ=1 ϕ ℓ y ℓ , and Σ-1 = σ -2 N ℓ=1 ϕ ℓ ϕ ⊤ ℓ .Therefore, the posterior of θ * is a product of two multivariate Gaussians: N (θ 0 , Σ 0 ) represents prior knowledge about θ * and N ( θ, Σ) represents empirical evidence. for stage ℓ = 1, . . . , N do 5:z ℓ ← ϕ ⊤ ℓ θ + (y ℓ -g(ϕ ⊤ ℓ θ))/ ġ(ϕ ⊤ ℓ θ) 6: Σ ← Σ -1 0 + N ℓ=1 ġ(ϕ ⊤ ℓ θ)ϕ ℓ ϕ ⊤ ℓ -1 7: θ ← Σ Σ -1 0 θ 0 + N ℓ=1 ġ(ϕ ⊤ ℓ θ)ϕ ℓ z ℓ 8: until θ converges 9:Output: Posterior mean θ and covariance ΣGeneralized linear models (GLMs) [36] 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[-u]) is a sigmoid.The likelihood of observations in GLMs has the following form [25] .Assumption 2. Let h = {(ϕ ℓ , y ℓ )} ℓ∈[N ]be a history of N observations under mean function g and the corresponding noise. Then log p(h| θ * ) ∝ N ℓ=1 y ℓ ϕ ⊤ ℓ θ * -b(ϕ ⊤ ℓ θ * ) + c(y ℓ ), where c is a real function and b is a function whose derivative is the mean function, ḃ = g.The posterior distribution of θ * in GLMs does not have a closed form in general [10] . Therefore, it is often approximated by the Laplace approximation. Let the prior distribution of the model parameter be p(θ * ) = N (θ * ; θ 0 , Σ 0 ), as in Section 2.1. Then the Laplace approximation is N ( θ, Σ), where θ is the maximum a posteriori (MAP) estimate of θ * and Σ is the corresponding covariance. Note that the Laplace approximation can be applied to non-Gaussian priors.The MAP estimate θ can be obtained by iteratively reweighted least squares (IRLS) [51] , which we present in Algorithm 1. IRLS is a Newton-type algorithm that computes θ iteratively (line 6). The convergence rate is fast due to the strong convexity of the problem. In the updates, ġ is the derivative of the mean function and the pseudo-observation z ℓ (line 4) plays the role of the observation y ℓ in Section 2.1. The IRLS solution has a similar structure to (3). Specifically, N ( θ, Σ) is a product of two multivariate Gaussians, representing prior knowledge about θ * and empirical evidence.The assumption that p(θ * ) = N (θ * ; θ 0 , Σ 0 ) is limiting, for instance because it precludes multimodal priors. We relax it by representing p(θ * ) by a diffusion model, which we call a diffusion model prior.We propose an efficient posterior sampling approximation for this prior, where the prior and empirical evidence are mixed similarly to (3) and IRLS. We review diffusion models next. 
diffusion models: 
 Diffusion models [45, 18] are generative models trained by diffusing samples from unknown and hard to represent distributions. They can be viewed in multiple ways [48] . We adopt the probabilistic formulation and presentation of Ho et al. [18] . A diffusion model is a graphical model with T stages indexed by t ∈ [T ]. Each stage t is associated with a latent variable S t ∈ R d . A sample from the model is represented by an observed variable S 0 ∈ R d . We visualize a diffusion model in Figure 1 . In the forward process, a clean sample s 0 is diffused through a sequence of variables S 1 , . . . , S T . This process is used to learn the reverse process, where the clean sample s 0 is generated through a sequence of variables S T , . . . , S 0 . To sample s 0 from the posterior (Section 4), we add a random variable H that represents partial information about s 0 . We introduce forward and reverse diffusion 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 to introduce our posterior approximations. Forward process (probability measure q) Reverse process (probability measure p) Forward process. In the forward process, a clean sample s 0 is diffused through a chain of latent variables S 1 , . . . S T (Figure 1 ). We denote the probability measure under this process by q and define its joint probability distribution asS T ← S T -1 ← • • • ← S 1 ← S 0 → H S T → S T -1 → • • • → S 1 → S 0 → HEQUATIONwhere q(s t | s t-1 ) is the conditional density of mapping a less diffused s t-1 to a more diffused s t . The diffusion rate is set by parameters α t ∈ (0, 1) and β t = 1 -α t . We also define ᾱt = t ℓ=1 α ℓ . The forward process is sampled as follows. First, a clean sample s 0 is chosen. Then S t ∼ q(• | s t-1 ) are sampled, starting from t = 1 to t = T .Reverse process. In the reverse process, a clean sample s 0 is generated through a chain of variables S T , . . . , S 0 (Figure 1 ). We denote the probability measure under this process by p and define its joint probability distribution asEQUATIONp(s T ) = N (s T ; 0 d , I d ) , ∀t ∈ [T ] : p(s t-1 | s t ) = N (s t-1 ; µ t (s t ), Σ t ) ,where p(s t-1 | s t ) is the conditional density of mapping a more diffused s t to a less diffused s t-1 .The function µ t predicts the mean of S t-1 | s t and is learned (Appendix B). As in Ho et al. [18] , we keep the covariance fixed at Σ t = βt I d , where βt is defined in (14) in Appendix B. This is also known as a stable diffusion. We make this assumption only to simplify exposition. All our derivations in Section 4 hold when Σ t is learned, for instance as in Bao et al. [8] .This process is called reverse because it is learned by reversing the forward process. The reverse process is sampled as follows. First, a diffused sample S T ∼ p is chosen. After that, S t-1 ∼ p(• | s t ) are sampled, starting from t = T to t = 1. 
posterior sampling: 
 This section is organized as follows. In Section 4.1, we show how to sample from a chain of random variables conditioned on observations. In Sections 4.2 and 4.4, we specialize this to the observation models in Section 2.Let h = {(ϕ ℓ , y ℓ )} ℓ∈[N ] denote a 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 h factors asp(s 0:T | h) = p(s T | h) T t=1 p(s t-1 | s t , h) .Therefore, p(s 0:T | h) can be sampled from efficiently by first sampling from p(s T | h) and then from T conditional distributions p(s t-1 | s t , h). We derive these next. Lemma 1. Let p be a probability measure over the reverse process (Figure 1 ). Thenp(s T | h) ∝ s0 p(h | s 0 ) p(s 0 | s T ) ds 0 p(s T ) , ∀t ∈ [T ] \ {1} : p(s t-1 | s t , h) ∝ s0 p(h | s 0 ) p(s 0 | s t-1 ) ds 0 p(s t-1 | s t ) , p(s 0 | s 1 , h) ∝ p(h | s 0 ) p(s 0 | s 1 ) .Proof. The claim is proved in Appendix A.1.Algorithm 2 LaplaceDPS: Laplace posterior sampling with a diffusion model prior. S t-1 ∼ N (μ t (S t , h), Σt (h))5: Output: Posterior sample S 0Now we specialize Lemma 1 to the diffusion model prior (Section 3) and linear models (Section 2.1). The prior distribution is the reverse process in (5),p(s T ) = N (s T ; 0 d , I d ) , p(s t-1 | s t ) = N (s t-1 ; µ t (s t ), Σ t ) .The term p(h | s 0 ) is the likelihood of observations in Assumption 1. The main challenge in using the lemma are potentially complex conditional densities of clean samples p(s 0 | S T ) and p(s 0 | s t ).To get around this issue, we make an additional assumption and then discuss it in Section 4.3. Theorem 2. Let p be a probability measure over the reverse process (Figure 1 ). Let θ and Σ-1 be defined as in (3) . Suppose thatEQUATIONEQUATIONEQUATIONΣt (h) = (Σ -1 t + Σ-1 /ᾱ t-1 ) -1 .Proof. The proof is in Appendix A.2. It has four steps. First, we fix stage t and apply approximation (6) . Second, we rewrite the likelihood as in (3) . Third, we reparameterize it as a function of s t . At the end, we combine the likelihood with the Gaussian prior using Lemma 6 in Appendix A.5.The algorithm that samples from the posterior distribution in Theorem 2 is presented in 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 is implemented as follows. The initial variable S T is sampled conditioned on h (line 2) from the distribution in (7) . This distribution is a product of the h-independent initial prior N (0 d , I d ) and the h-dependent distribution of the diffused evidence up to stage T , N ( √ ᾱT θ, ᾱT Σ). All remaining variables, S t-1 for t ∈ [T ], are sampled conditioned on s t and evidence h (line 3) from the distribution in (8) . This is again a product of the h-independent conditional prior N (µ t (s t ), Σ t ), from the pre-trained model, and the h-dependent distribution of the diffused evidence up to stage t -1, N ( √ ᾱt-1 θ, ᾱt-1 Σ). The final variable S 0 represents a clean sample. When compared to Section 2, the prior and evidence are mixed conditionally in a T -stage chain. This increases computational cost, as discussed in Section 8.Now we motivate our assumption in (6) . Simply put, we assume that s 0 = s t / √ ᾱt , where s 0 is a clean sample and s t is the corresponding diffused sample in stage t. This is motivated by the forward process, which relates s t and s 0 as [18] . After rearranging, we get s 0 = (s t -√ 1 -ᾱt ε t )/ √ ᾱt , and therefore s 0 can be viewed as a random variable with mean s t / √ ᾱt . The consequence of ( 6) is that the likelihood Take action a k ← arg max a∈A r(x k , a; θk ) and observe reward y k becomes a function of scaled s t , which yields a closed form when multiplied by the conditional prior, which is also a function of s t . Our approximation can be also viewed as the Tweedie's formula used in Chung et al. [12] where the score component is neglected.s t = √ ᾱt s 0 + √ 1 -ᾱt ε t , where ε t ∼ N (0 d , I d ) is a standard Gaussian noiseOur approximation has several notable properties. First, (1 -ᾱt )/ᾱ t → 0 as t → 1. Therefore, it 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 2 is asymptotically consistent.Theorem 3. Fix θ * ∈ R d . Let θ ← LaplaceDPS((µ t , Σ t ) t∈[T ] , h), where h = {(ϕ ℓ , y ℓ )} ℓ∈[N ] is a history of N observations. Suppose that λ d ( Σ-1 ) → ∞ as N → ∞, where Σ is defined in (3). Then P lim N →∞ ∥ θ -θ * ∥ 2 = 0 = 1.Proof. The proof is in Appendix A.3. The key idea is that the conditional posteriors in ( 7) and ( 8) concentrate at a scaled unknown model parameter θ * as the number of observations increases, which we formalize asλ d ( Σ-1 ) → ∞.The bound in Theorem 3 can be interpreted as follows. The sampled parameter θ approaches the true unknown parameter θ * as the number of observations N 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 to assuming that λ d ( Σ-1 ) = Ω(N ). This lower bound can be attained in linear models by getting observations according to the D-optimal design [38] . Since our proof is asymptotic, we can neglect some finite-time errors and this simplifies the argument.The Laplace approximation in GLMs (Section 2.2) naturally generalizes the exact posterior distribution in linear models (Section 2.1). We generalize Theorem 2 to GLMs along the same lines.Theorem 4. Let p be a probability measure over the reverse process (Figure 1 ). Suppose that (6)holds for all t ∈ [T ]. Then p(s T | h) ∝ N (s T ; μT +1 (h), ΣT +1 (h)), where μT +1 (h) = √ ᾱT θT +1 , ΣT +1 (h) = ᾱT ΣT +1 , θT +1 , ΣT +1 ← IRLS(0 d , I d /ᾱ T , h) . For t ∈ [T ], we have p(s t-1 | s t , h) ∝ N (s t-1 ; μt (s t , h), Σt (h)), where μt (s t , h) = √ ᾱt-1 θt , Σt (h) = ᾱt-1 Σt , θt , Σt ← IRLS(µ t (s t )/ √ ᾱt-1 , Σ t /ᾱ t-1 , h) .Proof. The proof is in Appendix A.4. It has four steps. First, we fix stage t and apply approximation (6) . Second, we reparameterize the prior, from a function of s t to a function of s t / √ ᾱt . Third, we combine the likelihood with the prior using the Laplace approximation. Finally, we repameterize the posterior, from a function of s t / √ ᾱt to a function of s t .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 LaplaceDPS, where the mean and covariances would be taken from Theorem 4. Note that Theorem 2 is a special case of Theorem 4 where the mean function g is an identity. 
application to contextual bandits: 
 Now we apply our posterior sampling approximations (Section 4) to contextual bandits. A contextual bandit [28, 31] 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 a ∈ A in context x ∈ X is r(x, a; θ * ), where r : X × A × Θ → R denotes a reward function and θ * ∈ Θ is a model parameter (Section 2). The agent interacts with the bandit for n rounds indexed by k ∈ [n]. In round k, the agent observes a context x k ∈ X , takes an action a k ∈ A, and observes a stochastic reward y k = r(x k , a k ; θ * ) + ε k , where ε k is independent zero-mean σ 2 -sub-Gaussian noise for σ > 0. The goal 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 asEQUATIONwhere a k, * = arg max a∈A r(x k , a; θ * ) is the optimal action in round k.Arguably the most popular method for solving contextual bandit problems is Thompson sampling [49, 11, 3] . The key idea in TS is to use the posterior distribution of θ * to explore. This is done as follows. In round k, the model parameter is drawn from the posterior in (2), θk ∼ p(• | h k ), where h k is the history of all interactions up to round k. After that, the agent takes the action with the highest mean reward under θk . The pseudo-code of this algorithm is given in Algorithm 3.A linear bandit [13, 1] is a contextual bandit with a linear reward function r(x, a; θ * ) = ϕ(x, a) ⊤ θ * , where ϕ : X × A → R d 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 in x and a. The feature extractor can be hand-designed or learned [39] . To simplify notation, we let ϕ ℓ = ϕ(x ℓ , a ℓ ) be the feature vector of the action in round k. So the history of interactions up to round k ish k = {(ϕ ℓ , y ℓ )} ℓ∈[k-1] . When the prior distribution is a Gaussian, p(θ * ) = N (θ * ; θ 0 , Σ 0 ), the posterior in round k is a Gaussian in (3) for h = h k .When the prior is a diffusion model, we propose sampling asEQUATIONwhere μt and Σt in LaplaceDPS are set according to Theorem 2. We call this algorithm DiffTS.A generalized linear bandit [16, 23, 32, 25] is an extension of linear bandits to generalized linear models (Section 2.2). When p(θ * ) = N (θ * ; θ 0 , Σ 0 ), the Laplace approximation to the posterior is a Gaussian (Section 2.2). When the prior is a diffusion model, we propose sampling from (10) , where μt and Σt in LaplaceDPS are set according to Theorem 4. 
experiments: 
 We conduct three experiments: synthetic problems in 2 dimensions (Section 6.2 and Appendix C.1), 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 .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) [21] . The last baseline is diffusion posterior sampling (DPS) of Chung et al. [12] . We implement all TS baselines as described in Section 5. The uninformative prior is N (0 d , I d ).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 as LinUCB [1] and the ε-greedy policy. They performed worse than TS, and therefore we do not report them here.Each experiment is set up as follows. First, the prior distribution of θ * is specified: it can be synthetic or estimated from real-world data. Second, we learn this distribution from 10 000 samples from it. In DiffTS and DPS, we follow Appendix B. The number of stages is T = 100 and the diffusion factor is α t = 0.97. Since 0.97 100 ≈ 0.05, most of the information in the training samples is diffused. The regressor in Appendix B is a 2-layer neural network with ReLU activations. In TunedTS, we fit the mean and covariance using maximum likelihood estimation. In MixTS, we fit the Gaussian mixture using SCIKIT-LEARN. All algorithms are evaluated on θ * sampled from the true prior. The regret is computed as defined in (9) . All error bars are standard errors of the estimates.The first experiment is on three synthetic problems. Each problem is a linear bandit (Section 5) with K = 100 actions in d = 2 dimensions. The reward noise is σ = 1. The feature vectors of actions are sampled uniformly at random from a unit ball. The prior distributions of θ * are shown in Figure 2 .The first is a mixture of two Gaussians and the last can be approximated well by a 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 2 . We observe two main trends. First, samples from the diffusion prior closely resemble those from the true prior. In such cases, DiffTS is expected to perform well and even outperforms MixTS, because it has a better representation of the prior. We observe this in all problems. Second, DPS diverges as the number of rounds increases. This is because DPS relies on the likelihood score (Section 7), which tends to infinity and causes instability. This happens despite tuning (Appendix D). We report results on additional synthetic problems in Appendix C.1.DiffTS should be T times more computationally costly than TS with a Gaussian prior (Section 4.2). We observe this empirically. As an example, the average cost of 100 runs of 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 around 85 at T = 1 to 50 at T = 25.In the second experiment, we learn to recommend an item to randomly arriving users. The problem is simulated using the MovieLens 1M dataset [27] , with one million ratings for 3 706 movies from 6 040 users. We subtract the mean rating from all ratings and complete the sparse rating matrix M by alternating least squares [14] with rank d = 5. The learned factorization is M = U V ⊤ . The i-th row of U , denoted by U i , represents user i. The j-th row of V , denoted by V j , represents movie j. We use movie embeddings V j as model parameters and user embeddings U i 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 V j and K = 10 actions are chosen randomly from U i in each round. In the linear bandit, the mean reward of item j for user i is U ⊤ i V j . The reward noise is σ = 0.75, and we estimate it from data. In the logistic bandit, the mean reward is g(U ⊤ i V j ), where g is a sigmoid. Our MovieLens results are reported in Figure 3 and we observe similar trends to 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 a UMAP projection [44] . Second, DiffTS has a significantly 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, which motivated our work on online learning with more complex priors. 
related work: 
 We start with reviewing related works on bandits with diffusion models. Hsieh et al. [22] proposed Thompson sampling with a diffusion model prior for K-armed bandits. There are multiple technical differences from our work. First, the diffusion model in Hsieh et al. [22] is over 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 t, Hsieh et al. [22] 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 that combines the conditional prior in stage t and likelihood. Because of this, Hsieh et al. [22] 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. [22] do not analyze their approximation.Aouali [4] proposed and analyzed contextual bandits with a linear diffusion model prior: µ t (s t ) in ( 5) is linear in s t and q(s 0 ) is a Gaussian. Because of this, their model is a linear Gaussian model and not a general diffusion model, as in our work.The closest related work on posterior sampling in diffusion models is DPS of Chung et al. [12] . The key idea in DPS is to sample from the posterior distribution using the likelihood score ∇ log p(h | θ), where p(h | θ) is the likelihood (Assumptions 1 and 2). Note that ∇ log p(h | θ) grows linearly in N because the history h in p(h | θ) involves N terms. Therefore, DPS becomes unstable as N → ∞. See our empirical results in Section 6.2 and 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: a sequential Monte Carlo approximation for the conditional reverse process [52] , a variant of DPS with an uninformative prior [37] , a pseudo-inverse approximation to the likelihood of evidence [47] , and posterior sampling in latent diffusion models [40] . All of these methods rely on the likelihood score ∇ log p(h | θ) 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 2 and 4), and thus gradient-free. They work well across different levels of uncertainty (Section 6) and do not require tuning.We also wanted to note that posterior sampling is a special from of guiding generation in diffusion models. Other approaches include conditional pre-training [15] , a constraint in the reverse process [17] , refining the null-space content [50] , solving an optimization problem that pushes the reverse process towards evidence [46] , and aligning the reverse process with the prompt [7] . 
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 T 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 a 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 is small, DiffTS performs very similarly to posterior sampling with a 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 reverses 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 [42] proved a O( ΓH(A * )n) regret bound for a linear bandit, where Γ 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 reflects how concentrated the prior is, 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 conditions, this is expected for any observation model when the number of observations is large. 
 
 Appendix: This section contains proofs of our main claims and supporting lemmas.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 asp(s 0:T | h) = p(s T | h) T t=1 p(s t-1 | s t:T , h) = p(s T | h) T t=1 p(s t-1 | s t , h) .We use that p(s t-1 | s t:T , h) = p(s t-1 | s t , h) in the last equality. We consider two cases.Derivation of p(s t-1 | s t , h). By Bayes' rule, we getp(s t-1 | s t , h) = p(h | s t-1 , s t ) p(s t-1 | s t ) p(h | s t ) ∝ p(h | s t-1 ) p(s t-1 | s t ) .In the last step, we use that p(h | s t ) is a constant, since s t and h are fixed, and that p(h| s t-1 , s t ) = p(h | s t-1). Note that the last term p(s t-1 | s t ) is the conditional prior distribution. Let t > 1. Then we rewrite the first term asp(h | s t-1 ) = s0 p(h, s 0 | s t-1 ) ds 0 = s0 p(h | s 0 , s t-1 ) p(s 0 | s t-1 ) ds 0 = s0 p(h | s 0 ) p(s 0 | s t-1 ) ds 0 .In the last equality, we use that our graphical model is a chain (Figure 1 ), and thus p(h | s 0 , s t-1 ) = p(h | s 0 ). Finally, we chain all identities and get thatEQUATIONDerivation of p(s T | h). By Bayes' rule, we getp(s T | h) = p(h | s T ) p(s T ) p(h) ∝ p(h | s T ) p(s T ) .In the last step, we use that p(h) is a constant, since h is fixed. The first term can be rewritten asp(h | s T ) = s0 p(h, s 0 | s T ) ds 0 = s0 p(h | s 0 , s T ) p(s 0 | s T ) ds 0 = s0 p(h | s 0 ) p(s 0 | s T ) ds 0 .Finally, we chain all identities and get thatEQUATIONThis completes the derivations.This proof has two parts.Derivation of p(s t-1 | s t , h). From p(h | s 0 ) ∝ N (s 0 ; θ, Σ) and ( 6), it immediately follows thats0 p(h | s 0 ) p(s 0 | s t-1 ) ds 0 ∝ N (s t-1 / √ ᾱt-1 ; θ, Σ) ∝ N (s t-1 ; √ ᾱt-1 θ, ᾱt-1 Σ) .The last step treats ᾱt-1 and Σ as constants, because the forward process, t, and evidence are fixed. Now we apply Lemma 6 to distributionsp(s t-1 | s t ) = N (s t-1 ; µ t (s t ), Σ t ) , N (s t-1 ; √ ᾱt-1 θ, ᾱt-1 Σ) ,and get thatp(s t-1 | s t , h) ∝ N (s t-1 ; μt (s t , h), Σt (h)) ,where μt (s t , h) and Σt (h) are defined in the claim. This is a product of two Gaussians: the prior with mean µ t (s t ) and covariance Σ t , and the evidence with mean θ and covariance Σ.Derivation of p(s T | h). Analogously to the derivation of p(s t-1 | s t , h), we establish thats0 p(h | s 0 ) p(s 0 | s T ) ds 0 ∝ N (s T ; √ ᾱT θ, ᾱT Σ) .Then we apply Lemma 6 to distributionsp(s T ) = N (s T ; 0 d , I d ) , N (s T ; √ ᾱT θ, ᾱT Σ) ,and get thatp(s T | h) ∝ N (s T ; μT +1 (h), ΣT +1 (h)) ,where μT +1 (h) and ΣT +1 (h) are defined in the claim. This is a product of two Gaussians: the prior with mean 0 d and covariance I d , and the evidence with mean θ and covariance Σ.We start with the triangle inequality∥ θ -θ * ∥ 2 = ∥ θ -θ + θ -θ * ∥ 2 ≤ ∥ θ -θ∥ 2 + ∥ θ -θ * ∥ 2 ,where we introduce θ from Section 2.1. Now we bound each term on the right-hand side.Upper bound on ∥ θ -θ∥ 2 . This part of the proof is based on analyzing the asymptotic behavior of the conditional densities in Theorem 2.As a first step, note that S T ∼ N (μ T +1 (h), ΣT +1 (h)), whereμT +1 (h) = ΣT +1 (h)(I d 0 d + Σ-1 θ/ √ ᾱT ) , ΣT +1 (h) = (I d + Σ-1 /ᾱ T ) -1 .Since λ d ( Σ-1 ) → ∞, we getΣT +1 (h) → ᾱT Σ , μT +1 (h) → √ ᾱT θ .Moreover, λ d ( Σ-1 ) → ∞ implies λ 1 ( Σ) → 0, and thus lim N →∞ ∥S T -√ ᾱT θ∥ 2 = 0.The same argument can be applied inductively to later stages of the reverse process. Specifically, for any t ∈ [T ], S t-1 ∼ N (μ t (S t , h), Σt (h)), whereμt (S t , h) = Σt (h)(Σ -1 t µ t (S t ) + Σ-1 θ/ √ ᾱt-1 ) , Σt (h) = (Σ -1 t + Σ-1 /ᾱ t-1 ) -1 . Since λ d ( Σ-1 ) → ∞ and S t → √ ᾱt θ by induction, we get Σt (h) → ᾱt-1 Σ , μt (S t , h) → √ ᾱt-1 θ .Moreover, λ d ( Σ-1 ) → ∞ implies λ 1 ( Σ) → 0, and thus lim N →∞ ∥S t-1 -√ ᾱt-1 θ∥ 2 = 0.In the last stage, t = 1, ᾱ0 = 1, and S 0 = θ, which implies thatlim N →∞ ∥ θ -θ∥ 2 → 0 .Upper bound on ∥ θ -θ * ∥ 2 . This part of the proof uses the definition of θ in Section 2.1 and that ε ℓ ∼ N (0, σ 2 ) is independent noise. By definition,θ -θ * = σ -2 Σ N ℓ=1 ϕ ℓ y ℓ -θ * = σ -2 Σ N ℓ=1 ϕ ℓ (ϕ ⊤ ℓ θ * + ε ℓ ) -θ * = σ -2 Σ N ℓ=1 ϕ ℓ ε ℓ .Since ε ℓ is independent zero-mean Gaussian noise with variance σ 2 , θ -θ * is a Gaussian random variable with mean 0 d and covariancecov σ -2 Σ N ℓ=1 ϕ ℓ ε ℓ = σ -4 Σ N ℓ=1 ϕ ℓ var [ε ℓ ] ϕ ⊤ ℓ Σ = Σ N ℓ=1 ϕ ℓ ϕ ⊤ ℓ σ 2 Σ = Σ . Since λ d ( Σ-1 ) → ∞ implies λ 1 ( Σ) → 0, we get lim N →∞ ∥ θ -θ * ∥ 2 = 0 .This completes the proof.This proof has two parts. We state and prove our supplementary lemmas next. Lemma 5. Let p(x) = N (x; µ 1 , Σ 1 ) and q(x) = N (x; µ 2 , Σ 2 ), whereµ 1 , µ 2 ∈ R d and Σ 1 , Σ 2 ∈ R d×d . Then d(p, q) = 1 2 (µ 2 -µ 1 ) ⊤ Σ -1 2 (µ 2 -µ 1 ) + tr(Σ -1 2 Σ 1 ) -log det(Σ 1 ) det(Σ 2 ) -d .Moreover, whenΣ 1 = Σ 2 , d(p, q) = 1 2 (µ 2 -µ 1 ) ⊤ Σ -1 2 (µ 2 -µ 1 ) .Proof. The proof follows from the definitions of KL divergence and multivariate Gaussians.Lemma 6. Fix µ 1 ∈ R d , Σ 1 ⪰ 0, µ 2 ∈ R d , and Σ 2 ⪰ 0. Then N (x; µ 1 , Σ 1 ) N (x; µ 2 , Σ 2 ) ∝ N (x; µ, Σ) ,where embeddings of all digits in the dataset, which are used as feature vectors in our experiment. We generate a distribution over model parameters θ * as follows: (1) we choose a 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 a linear model, which gives us a single θ * . We repeat this 10 000 times and get a distribution over θ * .µ = Σ(Σ -1 1 µ 1 + Σ -1 2 µ 2 ) , Σ = (Σ -1 1 + Σ -1 2 ) -1 .We consider both linear and logistic bandits. In both, the model parameter θ * is 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 ⊤ θ * and the reward noise is σ = 1.In the logistic bandit, the mean reward is g(x ⊤ θ * ), where g is a sigmoid.Our MNIST results are reported in Figure 5 . We observe again that DiffTS has a lower regret than all baselines, because the learned prior captures the underlying distribution of θ * well. We note that both the prior and diffusion prior distributions exhibit a strong cluster structure (Figure 5a ), where each cluster represents one label.We conduct three ablation studies on the cross problem in Figure 2 .In all experiments, the number of samples for training diffusion priors was 10 000. In Figure 6a , we vary it from 100 to 10 000. We observe that the regret decreases as the number of samples increases, 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 of diffusion stages was T = 100. In Figure 6b , we vary it from 1 to 300 and observe its impact on regret. While the regret at T = 1 is high, it decreases quickly as 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 on the computation time of posterior sampling. The time is linear in T , as suggested in Section 4.2. The main contributor to it is the neural network regressor.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 evaluate all methods from Figure 2 . In addition, we implement a sequential Monte Carlo (SMC) sampler. The initial particles are selected 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 over time, as the posterior concentrates. The particles are weighted according to the likelihood of the observation in the round. Finally, we use normalized likelihood weights to resample the particles. We tune SMC to get good posterior approximations. We use 3 000 particles. For this setting, the computational costs of posterior sampling in SMC and DiffTS are comparable.Our results are reported in Figure 7 . We observe that DiffTS approximations are comparable to MixTS, which has an exact posterior in this setting. The second best performing method is SMC. Its approximations worsen as the sample size n increases. DPS approximations also get worse as n increases, which caused instability in Figure 2 .In our experiments, we compare to diffusion posterior sampling (DPS) with a Gaussian observation noise (Algorithm 1 in Chung et al. [12] ). Our implementation is presented in Algorithm 4. The score is Ŝ = ε t (S t ; ψ t )/ √ 1 -ᾱt , where ε t (S t ; ψ t ) is a regression estimate of the forward process noise ε t in Appendix B. We set σt = βt , which is the same amount of noise as in our reverse process (Section 3). The term∇ N ℓ=1 (y ℓ -ϕ ⊤ ℓ Ŝ0 ) 2is the gradient of the negative log-likelihood with respect to S t .As noted in Appendices C.2 and D.1 of Chung et al. [12] , ζ t in DPS needs to be tuned for good performance. We also observed this in our experiments (Section 6.2). To make DPS work well, we follow Chung et al. [12] and setζ t = 1 N ℓ=1 (y ℓ -ϕ ⊤ ℓ Ŝ0 ) 2.While this significantly improves the performance of DPS, it does not prevent failures. The fundamental problem is that gradient-based optimization is sensitive to the step size, especially when the optimized function is steep. Note that LaplaceDPS does not have any such hyper-parameter. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

The 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.
questions: 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.

Clarity would be enhanced by more explanations on the roles of the reverse process and forward process in posterior sampling.
limitations: limitations: I recommend that the authors include a more detailed discussion of the limitations of this work.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1)	This 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] ?

2)	The 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.
questions: questions: 1)	In 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?

2)	Why does $\nabla log p(h|\theta)$ grow linearly in N? 

3)	Does relying on the score of likelihood only impact the application for contextual bandits or does this also impact these methods generally?

4)	Is there any error caused/limitation in using the Laplace approximation?
limitations: limitations: They discuss that they have performed regret analysis. However they do discuss limitations of the proposed sampling algorithm.
paper_summary: paper_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.
strength_weakness: strength_weakness: A few notes about correctness/presentation:
(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.
(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.

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

It was not clear which of the proofs are standard results and which are novel (e.g. Lemma 1 looks like Bayes theorem).
questions: 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: 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?
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

Some format issues: p2 l49, p4 Fig. 1, l121, and l118.
questions: questions: Can the author explain why the diffusion model prior solves the instability and divergence issues?

What other benefits could the diffusion model prior bring?

What other approaches could be employed beyond Laplace approximation?

What are the technical difficulties in the paper, especially the analysis?
limitations: limitations: Discussion on limitations seems to be missing or need to be more explicit.
paper_summary: paper_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.
strength_weakness: strength_weakness: I picked up a few points, but I believe most of them can be addressed in the rebuttal.

I'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.

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

It 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.
questions: 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: limitations: The authors address some of the limitations. See my comments above for additional limitations that should be addressed.

Comments: title: Reply to authors
thank you
Response to authors' rebuttals
comment: Thanks for the authors' responses. I'll keep my score.
Thank you for the rebuttal including the common responses. I maintain my current score of weak accept.
Thanks for the authors' responses. The discussions on the limitations and the prior work now look more comprehensive to me. I'm happy to keep my positive score.
Thank you for your response. I will maintain my already positive score.
Dear reviewers,

Thank you for the reviews and taking our rebuttal into account when evaluating the paper. The reviewer-author discussion period will end soon (August 13 EoD AoE). If you have any additional concerns about the paper, we would be happy to address them in the next 2 days.

Sincerely,

Authors
I've read the authors' rebuttals.

The authors have addressed my comments and I'm happy to maintain my good score.