Title: Online Posterior Sampling with a Diffusion Prior
Venue: NeurIPS 2024
Accepted: yes
Keywords: posterior sampling

diffusion models

online learning

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

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.

1 Introduction

A multi-armed bandit [27, 6, 30] 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 are beneficial, such recommender systems [32, 54, 25, 35] and hyper-parameter optimization [34].

Contextual bandits [29, 32] with linear [13, 1] and generalized linear models (GLMs) [17, 33, 2, 26] have become popular due to 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, 44]. 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, 40, 36, 9, 21, 20, 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 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, in each

Aniket Deshmukh AWS AI Labs

stage of the diffusion model, using the Laplace approximation. The resulting posterior concentrates at a single point and can be sampled from efficiently even if the likelihood score is high. We prove that this approximation is asymptotically consistent.

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

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.

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

where 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 is

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

2.1 Linear Model

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 model parameter θ ∗ is p ( h | θ ∗ ) ∝ exp[ -∑ N ℓ =1 ( y ℓ -ϕ ⊤ ℓ θ ∗ ) 2 / (2 σ 2 )] .

Let p ( θ ∗ ) = N ( θ ∗ ; θ 0 , Σ 0 ) be the prior distribution of θ ∗ , where θ 0 ∈ R d and Σ 0 ∈ R d × d are the prior mean and covariance, respectively. Then p ( θ ∗ | h ) ∝ N ( θ ∗ ; ˆ θ, ˆ Σ) [10], where

are the posterior mean and covariance, respectively. In this work, we write them equivalently as

where ¯ θ = σ -2 ¯ Σ ∑ N ℓ =1 ϕ ℓ y ℓ and ¯ Σ -1 = σ -2 ∑ N ℓ =1 ϕ ℓ ϕ ⊤ ℓ are the empirical mean and inverse of its covariance, respectively. Therefore, the posterior of θ ∗ is a product of two multivariate Gaussians: N ( θ 0 , Σ 0 ) representing prior knowledge about θ ∗ and N ( ¯ θ, ¯ Σ) representing empirical evidence.

Algorithm 1 IRLS : Iteratively reweighted least squares.

1: Input: Prior parameters θ 0 and Σ 0 , history of observations h = { ( ϕ ℓ , y ℓ ) } ℓ ∈ [ N ]

2: Initialize ˆ θ ∈ R d

3: repeat

4: for stage ℓ = 1 , . . . , N do

5: z ℓ ← ϕ ⊤ ℓ ˆ θ +( y ℓ -g ( ϕ ⊤ ℓ ˆ θ )) / ˙ g ( ϕ ⊤ ℓ ˆ θ ) -1

6: ˆ Σ ← ( Σ -1 0 + ∑ N ℓ =1 ˙ g ( ϕ ⊤ ℓ ˆ θ ) ϕ ℓ ϕ ⊤ ℓ )

7: ˆ θ ← ˆ Σ ( Σ -1 0 θ 0 + ∑ N ℓ =1 ˙ g ( ϕ ⊤ ℓ ˆ θ ) ϕ ℓ z ℓ )

8: until ˆ θ converges

9: Output: Posterior mean ˆ θ 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[ -u ]) is a sigmoid. The likelihood of observations in GLMs has the following form [26].

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, ˙ b = 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) [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 a similar structure to (3). That is, N ( ˆ θ, ˆ Σ) is a product of two multivariate Gaussians, representing prior knowledge about θ ∗ and empirical evidence. The new quantities in GLMs are the derivative of the mean function ˙ g and pseudo-observations z ℓ (line 5), which play the role of observations y ℓ in Section 2.1.

2.3 Towards Diffusion Model Priors

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 efficient posterior sampling approximations for this prior, where the prior and empirical evidence are mixed similarly to (3) and IRLS . We review diffusion models next.

3 Diffusion Models

Diffusion models [46, 19] are generative models trained by diffusing samples from unknown and hard to represent distributions. They can be viewed in multiple ways [49]. We adopt the probabilistic formulation and presentation of Ho et al. [19]. 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

Figure 1: Graphical models of the forward and reverse processes in the diffusion model. The variable H represents partial information about S 0 .

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

where 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 . The forward process is sampled from as follows. First, a clean sample s 0 is chosen. Then S t ∼ q ( · | s t -1 ) are sampled, 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 as

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. [19], we keep the covariance fixed at Σ t = ˜ β t I d , where ˜ β t = 1 -¯ α t -1 1 -¯ α t β t and ¯ α t = ∏ t ℓ =1 α ℓ . This is 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 from as follows. First, an initial diffused sample S T ∼ p is sampled. After that, S t -1 ∼ p ( · | s t ) are sampled, from t = T to t = 1 .

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

4.1 Chain Model Posterior

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 as

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

Algorithm 2 LaplaceDPS : Laplace posterior sampling with a diffusion model prior.

1: Input: Diffusion model parameters ( µ t , Σ t ) t ∈ [ T ] , history of observations h

2: Initial sample S T ∼ N (ˆ µ T +1 ( h ) , ˆ Σ T +1 ( h ))

3: for stage t = T, . . . , 1 do

4: S t -1 ∼ N (ˆ µ t ( S t , h ) , ˆ Σ t ( h ))

5: Output: Posterior sample S 0

Proof. The claim is proved in Appendix A.1.

4.2 Linear Model Posterior

Now 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),

The term p ( h | s 0 ) is the likelihood of observations in Assumption 1. The main challenge in using the lemma is that the conditional densities of clean samples p ( s 0 | S T ) and p ( s 0 | s t ) are complex [12]. To get around this, we make an additional assumption, which is discussed 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 that

holds for all t ∈ [ T ] . Then p ( s T | h ) ∝ N ( s T ; ˆ µ T +1 ( h ) , ˆ Σ T +1 ( h )) , where

For t ∈ [ T ] , we have p ( s t -1 | s t , h ) ∝ N ( s t -1 ; ˆ µ t ( s t , h ) , ˆ Σ t ( h )) , where

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 prior N ( 0 d , I d ) and the h -dependent distribution of the diffused evidence up to stage T , N ( √ ¯ α T ¯ θ, ¯ α T ¯ Σ) . Then, for any t ∈ [ T ] , S t -1 is sampled conditioned on s t and evidence h (line 4) from the distribution in (8). This distribution is 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 last variable S 0 is the clean sample. When compared to Section 2, the prior and evidence are mixed conditionally in a T -stage chain. This increases the computational cost T times, as discussed in Section 8.

4.3 Key Approximation in Theorem 2

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 s t = √ ¯ α t s 0 + √ 1 -¯ α t ε t , where ε t ∼ N ( 0 d , I d ) is a standard

Algorithm 3 Contextual Thompson sampling.

1: for round k = 1 , . . . , n do

2: Sample ˜ θ k ∼ p ( · | h k ) , where p ( · | h k ) is the posterior distribution in (2)

3: Take action a k ← arg max a ∈A r ( x k , a ; ˜ θ k ) and observe reward y k

Gaussian noise [19]. 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 becomes a function of s t , which yields a closed form when multiplied by the conditional prior, also a function of s t . Our approximation can be also viewed as the Tweedie's formula in Chung et al. [12] where the score component is neglected.

Our 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 [39].

4.4 GLMPosterior

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

For t ∈ [ T ] , we have p ( s t -1 | s t , h ) ∝ N ( s t -1 ; ˆ µ t ( s t , h ) , ˆ Σ t ( h )) , where

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.

5 Application to Contextual Bandits

Now we apply our posterior sampling approximations (Section 4) to contextual bandits. A 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 a ∈ A in context x ∈ X is r ( x, a ; θ ∗ ) , where r : X × A × Θ → R is 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 , it observes a context x k ∈ X , takes an action a k ∈ A , and observes its stochastic reward y k = r ( x k , a k ; θ ∗ ) + ε k with independent noise ε k . We assume that the noise is zero-mean σ 2 -sub-Gaussian for σ > 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 as

where 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 [50, 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] has 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 of x and a . The feature extractor can be either learned [40] or hand-crafted. We denote the feature vector of the action in round k by ϕ k = ϕ ( x k , a k ) . Therefore, the history of interactions up to round k is h 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 from the posterior using

where ˆ µ t and ˆ Σ t in LaplaceDPS are computed as in Theorem 2. We call this algorithm DiffTS .

A generalized linear bandit [17, 24, 33, 26] 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 posterior sampling using (10), where ˆ µ t and ˆ Σ t in LaplaceDPS are computed as in Theorem 4.

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

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 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 thus we do not report them.

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

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

Problem cross

0

50

100

150

200

Round n

0

5

10

15

20

25

Regret

Problem cross

TS

TunedTS

MixTS

DPS

DiffTS

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

Problem rays

0

50

100

150

200

Round n

0

5

10

15

20

25

Regret

Problem rays

TS

TunedTS

MixTS

DPS

DiffTS

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

Problem triangles

0

50

100

150

200

Round n

0

5

10

15

20

25

Regret

Problem triangles

TS

TunedTS

MixTS

DPS

DiffTS

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.

6.2 Synthetic Experiment

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 uses an approximation based on the likelihood score (Section 7), which is unstable when the score is high. This happens despite our best efforts to tune DPS (Appendix D). We report results on another three 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 26 at T = 1 to 15 at T = 50 .

6.3 MovieLens Experiment

In the second experiment, we learn to recommend an item to randomly arriving users. The problem is simulated using the MovieLens 1M dataset [28], 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

-4

-2

0

2

4

6

8

10

12

0

2

4

6

8

10

12

0

50

100

150

200

Round n

0

5

10

15

20

Regret

(b) MovieLens linear bandit

TS

TunedTS

MixTS

DiffTS

0

50

100

150

200

Round n

0

2

4

6

8

10

12

14

16

Regret

(c) MovieLens logistic bandit

TS

TunedTS

DiffTS

alternating least squares [14] with rank d = 5 . The learned factorization is M = UV ⊤ . 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 UMAP [45]. Second, DiffTS has a 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 et al. [23] 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. [23] 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. [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: µ t ( s t ) in (5) is linear in s t and q ( s 0 ) is a Gaussian. Therefore, this 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 →∞ . 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: a sequential Monte Carlo approximation for the conditional reverse process [53], a variant of DPS with an uninformative prior [38], a 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 ∇ 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 note that posterior sampling is a special form of inference-time guidance in diffusion models. Other approaches are conditional pre-training [15], a constraint in the reverse process [18], refining the 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 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 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 ( √ Γ 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 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

Yasin Abbasi-Yadkori, David Pal, and Csaba Szepesvari. Improved algorithms for linear stochastic bandits. In Advances in Neural Information Processing Systems 24 , pages 2312-2320, 2011.

Marc Abeille and Alessandro Lazaric. Linear Thompson sampling revisited. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics , 2017.

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.

Imad Aouali. Linear diffusion models meet contextual bandits with large action spaces. In NeurIPS 2023 Workshop on Foundation Models for Decision Making Workshop , 2023.

Imad Aouali, Branislav Kveton, and Sumeet Katariya. Mixed-effect Thompson sampling. In Proceedings of the 26th International Conference on Artificial Intelligence and Statistics , 2023.

Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine Learning , 47:235-256, 2002.

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.

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.

Soumya Basu, Branislav Kveton, Manzil Zaheer, and Csaba Szepesvari. No regrets for learning the prior in bandits. In Advances in Neural Information Processing Systems 34 , 2021.

Christopher Bishop. Pattern Recognition and Machine Learning . Springer, New York, NY, 2006.

Olivier Chapelle and Lihong Li. An empirical evaluation of Thompson sampling. In Advances in Neural Information Processing Systems 24 , pages 2249-2257, 2011.

Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In Proceedings of the 11th International Conference on Learning Representations , 2023.

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.

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-622, 2016.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems 34 , 2021.

Arnaud Doucet, Nando de Freitas, and Neil Gordon. Sequential Monte Carlo Methods in Practice . Springer, New York, NY, 2001.

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.

Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. In Advances in Neural Information Processing Systems 35 , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems 33 , 2020.

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.

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.

Joey Hong, Branislav Kveton, Manzil Zaheer, Mohammad Ghavamzadeh, and Craig Boutilier. Thompson sampling with a mixture prior. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics , 2022.

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.

Kwang-Sung Jun, Aniruddha Bhargava, Robert Nowak, and Rebecca Willett. Scalable generalized linear bandits: Online computation and hashing. In Advances in Neural Information Processing Systems 30 , pages 98-108, 2017.

Jaya Kawale, Hung Bui, Branislav Kveton, Long Tran-Thanh, and Sanjay Chawla. Efficient Thompson sampling for online matrix-factorization recommendation. In Advances in Neural Information Processing Systems 28 , pages 1297-1305, 2015.

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.

Tze Leung Lai and Herbert Robbins. Asymptotically efficient adaptive allocation rules. Advances in Applied Mathematics , 6(1):4-22, 1985.

Shyong Lam and Jon Herlocker. MovieLens Dataset. http://grouplens.org/datasets/movielens/, 2016.

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.

Tor Lattimore and Csaba Szepesvari. Bandit Algorithms . Cambridge University Press, 2019.

Yann LeCun, Corinna Cortes, and Christopher Burges. MNIST Handwritten Digit Database. http://yann.lecun.com/exdb/mnist , 2010.

Lihong Li, Wei Chu, John Langford, and Robert Schapire. A contextual-bandit approach to personalized news article recommendation. In Proceedings of the 19th International Conference on World Wide Web , 2010.

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.

Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hyperband: A novel bandit-based approach to hyperparameter optimization. Journal of Machine Learning Research , 18(185):1-52, 2018.

Shuai Li, Alexandros Karatzoglou, and Claudio Gentile. Collaborative filtering bandits. In Proceedings of the 39th Annual International ACM SIGIR Conference , 2016.

Xiuyuan Lu and Benjamin Van Roy. Information-theoretic confidence bounds for reinforcement learning. In Advances in Neural Information Processing Systems 32 , 2019.

P. McCullagh and J. A. Nelder. Generalized Linear Models . Chapman & Hall, 1989.

Xiangming Meng and Yoshiyuki Kabashima. Diffusion model based posterior sampling for noisy linear inverse problems. CoRR , abs/2211.12343, 2023. URL https://arxiv.org/ abs/2211.12343 .

Friedrich Pukelsheim. Optimal Design of Experiments . Society for Industrial and Applied Mathematics, 2006.

Carlos Riquelme, George Tucker, and Jasper Snoek. 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.

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.

Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Mathematics of Operations Research , 39(4):1221-1243, 2014.

Daniel Russo and Benjamin Van Roy. An information-theoretic analysis of Thompson sampling. Journal of Machine Learning Research , 17(68):1-30, 2016.

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.

Tim Sainburg, Leland McInnes, and Timothy Gentner. Parametric UMAP embeddings for representation and semisupervised learning. Neural Computation , 33(11):2881-2907, 2021.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd International Conference on Machine Learning , pages 2256-2265, 2015.

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 on Learning Representations , 2024.

Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guided diffusion models for inverse problems. In Proceedings of the 11th International Conference on Learning Representations , 2023.

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.

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-294, 1933.

Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. In Proceedings of the 11th International Conference on Learning Representations , 2023.

R. Wolke and H. Schwetlick. Iteratively reweighted least squares: Algorithms, convergence analysis, and numerical comparisons. SIAM Journal on Scientific and Statistical Computing , 9 (5):907-921, 1988.

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.

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.

A Proofs and Supporting Lemmas

This section contains proofs of our main claims and supporting lemmas.

A.1 Proof of Lemma 1

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

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 get

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. When t > 1 , the first term can be expressed as

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 that

Derivation of p ( s T | h ) . By Bayes' rule, we get

In the last step, we use that p ( h ) is a constant, since h is fixed. The first term can be rewritten as

Finally, we chain all identities and get that

This completes the derivations.

A.2 Proof of Theorem 2

This proof has two parts.

Derivation of p ( s t -1 | s t , h ) . From (6) and Assumption 1, it follows that

The last step treats ¯ α t -1 and ¯ Σ as constants, because the forward process and evidence h are fixed. Now we apply Lemma 6 to distributions

and get that

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 √ ¯ α t -1 ¯ θ and covariance ¯ α t -1 ¯ Σ .

Derivation of p ( s T | h ) . Analogously to the derivation of p ( s t -1 | s t , h ) , we establish that

Then we apply Lemma 6 to distributions and get that

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 √ ¯ α T ¯ θ and covariance ¯ α T ¯ Σ .

A.3 Proof of Theorem 3

We start with the triangle inequality

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

Since λ d ( ¯ Σ -1 ) →∞ , we get

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

Since λ d ( ¯ Σ -1 ) →∞ and S t → √ ¯ α t ¯ θ by induction, we get

Moreover, λ d ( ¯ Σ -1 ) →∞ implies λ 1 ( ¯ Σ) → 0 , and thus lim N →∞ ∥ S t -1 - √ ¯ α t -1 ¯ θ ∥ 2 = 0 for any t ∈ [ T ] . In the last stage, t = 1 , ¯ α 0 = 1 , and S 0 = ˜ θ . Therefore,

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,

Since ε ℓ is independent zero-mean Gaussian noise with variance σ 2 , ¯ θ -θ ∗ is a Gaussian random variable with mean 0 d and covariance

Since λ d ( ¯ Σ -1 ) →∞ implies λ 1 ( ¯ Σ) → 0 , we get

This completes the proof.

A.4 Proof of Theorem 4

This proof has two parts.

Derivation of p ( s t -1 | s t , h ) . From (6), we have

Since p ( s t -1 | s t ) is a Gaussian, we have

for γ = 1 / √ ¯ α t -1 . Then by the Laplace approximation,

where ˙ θ t , ˙ Σ t ← IRLS ( γµ t ( s t ) , γ 2 Σ t , h ) .

Derivation of p ( s T | h ) . Analogously to the derivation of p ( s t -1 | s t , h ) , we establish that

Then by the Laplace approximation for γ = 1 / √ ¯ α T , we get

where ˙ θ T +1 , ˙ Σ T +1 ← IRLS ( 0 d , γ 2 I d , h ) .

A.5 Supporting Lemmas

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

Moreover, when Σ 1 = Σ 2 , where

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

Proof. This is a classic result, which is proved as

The neglected factors depend on constants µ 1 , µ 2 , Σ 1 , and Σ 2 . This completes the proof.

B Learning the Reverse Process

One property of our model is that q ( s T ) ≈ N ( s T ; 0 d , I d ) when T is sufficiently large [19]. Since S T has the same distribution in the reverse process p , 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 et al. [19] showed that

holds for any s 0 and s t , where

Therefore, the latent variable in stage t -1 , S t -1 , is easy to sample when s t and s 0 are known. To estimate s 0 , which is unknown when sampling from the reverse process, we use the forward process again. In particular, (4) implies that s t = √ ¯ α t s 0 + √ 1 -¯ α t ε t , where ε t ∼ N ( 0 d , I d ) is a standard Gaussian noise. This identity can be rearranged as

To obtain ε t , which is unknown when sampling from p , we learn to regress it from s t [19].

The regressor is learned as follows. Let ε t ( · ; ψ ) be a regressor of ε t parameterized by ψ and D = { s 0 } be a dataset of training examples. We sample s 0 uniformly at random from D and then solve

per stage. The expectation is approximated by sampled s 0 . 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 as in Section 6.2.

C.2 MNIST Experiment

The next experiment is on the MNIST dataset [31]. We start with learning an MLP-based multi-way classifier for digits and extract their d = 8 dimensional embeddings. These are used as features in

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

Problem swirl

0

50

100

150

200

Round n

0

5

10

15

20

25

Regret

Problem swirl

TS

TunedTS

MixTS

DPS

DiffTS

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

Problem H

0

50

100

150

200

Round n

0

5

10

15

20

25

Regret

Problem H

TS

TunedTS

MixTS

DPS

DiffTS

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

Problem corners

0

50

100

150

200

Round n

0

5

10

15

20

25

Regret

Problem corners

TS

TunedTS

MixTS

DPS

DiffTS

-6 -4 -2

0

2

4

6

8

10

12

0

2

4

6

8

10

12

(a) Problem MNIST

0

100

200

300

400

500

Round n

0

10

20

30

40

50

60

70

Regret

(b) MNIST linear bandit

TS

TunedTS

MixTS

DiffTS

0

100

200

300

400

500

Round n

0

5

10

15

20

25

30

35

40

45

Regret

(c) MNIST logistic bandit

TS

TunedTS

DiffTS

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 θ ∗ .

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.

C.3 Ablation Studies

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,

0

2000

4000

6000

8000

10000

Training Samples

14

15

16

17

18

19

Regret

(a) Training samples ablation

0

50

100

150

200

250

300

Diffusion Steps, T

14

16

18

20

22

24

26

28

Regret

(b) Diffusion steps ablation

0

50

100

150

200

250

300

Diffusion Steps, T

0.00

0.05

0.10

0.15

0.20

0.25

Compute time, sec

(c) Compute time per run

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

10 0

10 1

10 2

10 3

Sample size n

0.0

0.2

0.4

0.6

0.8

1.0

EMD

TS

TunedTS

MixTS

DPS

DiffTS

SMC

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

Problem P2

10 0

10 1

10 2

10 3

Sample size n

0.0

0.2

0.4

0.6

0.8

1.0

EMD

Problem P2

TS

TunedTS

MixTS

DPS

DiffTS

SMC

-4 -3 -2 -1

0

1

2

3

4

-4

-3

-2

-1

0

1

2

3

4

True and diffusion priors

Problem P3

10 0

10 1

10 2

10 3

Sample size n

0.0

0.2

0.4

0.6

0.8

1.0

EMD

Problem P3

TS

TunedTS

MixTS

DPS

DiffTS

SMC

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 is the neural network regressor.

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 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 the KL divergence. However, we could not apply it because the posteriors of DiffTS and DPS do not have analytical forms.

Algorithm 4 DPS of Chung et al. [12].

1: Input: Model parameters ˜ σ t and ζ t

2: Initial sample S T ∼ N ( 0 d , I d )

3: for stage t = T, . . . , 1 do

4: ˆ S ←-ε t ( S t ; ψ t ) 1 -¯ α t

5: ˆ S 0 ← 1 √ ¯ α t ( S t +(1 -¯ α t ) ˆ S )

6: Z ∼ N ( 0 d , I d )

7: S t -1 ← √ ¯ α t -1 β t 1 -¯ α t ˆ S 0 + √ α t (1 -¯ α t -1 ) 1 -¯ α t S t + ˜ σ t Z -ζ t ∇ ∑ N ℓ =1 ( y ℓ -ϕ ⊤ ℓ ˆ S 0 ) 2

8: Output: Posterior sample S 0

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

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.

D Implementation of Chung et al. [12]

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 ˆ S = -ε 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

is the gradient of the negative log-likelihood with respect to S t .

As discussed in Appendices C.2 and D.1 of Chung et al. [12], ζ t in DPS needs to be tuned for good performance. This is because ∇ ∑ N ℓ =1 ( y ℓ -ϕ ⊤ ℓ ˆ S 0 ) 2 grows with the number of observations, which 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

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.

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 set 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We also include code to reproduce the synthetic results in Figures 2 and 4.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions 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 2 and 4.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experiments are described to a sufficient level to be reproducible. To make sure, we include code to reproduce the synthetic results in Figures 2 and 4.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All plots in the paper have error bars.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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://neurips.cc/public/EthicsGuidelines ?

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 a particular application that would have immediate negative impact.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not pose such a risk.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

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 to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nIf 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?\n\nLIMITATIONS:\nThe authors address some of the limitations. See my comments above for additional limitations that should be addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nIt 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.\n\nQUESTIONS:\nCan 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?\n\nLIMITATIONS:\nDiscussion on limitations seems to be missing or need to be more explicit.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nA 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).\n\nQUESTIONS:\nIn 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?\n\nLIMITATIONS:\nThe 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?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nI am not familiar with Bandits or Diffusion models but the paper is interesting and the experiments in the context of contextual bandits seem convincing.\n\nWEAKNESSES:\n1)\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.\n\nQUESTIONS:\n1)\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?\n\nLIMITATIONS:\nThey discuss that they have performed regret analysis. However they do discuss limitations of the proposed sampling algorithm.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nAs 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.\n\nLIMITATIONS:\nI recommend that the authors include a more detailed discussion of the limitations of this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-16"}