Original: Diffusion Prior for Online Decision Making: A Case Study of Thompson Sampling
Revised: Diffusion Prior for Online Decision Making: A Case Study of Thompson Sampling

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: In this work, we investigate the possibility of using denoising diffusion models tolearn priors for online decision making problems.
Revised: In this work, we investigate the possibility of using denoising diffusion models to learn priors for online decision making problems.

Original: Our special focus is on the metalearning for bandit framework, with the goal of learning a strategy that performswell across bandit tasks of a same class.
Revised: Our special focus is on the metalearning for bandit framework, with the goal of learning a strategy that performs well across bandit tasks of a same class.

Original: To this end, we train a diffusion model thatlearns the underlying task distribution and combine Thompson sampling with thelearned prior to deal with new task at test time.
Revised: To this end, we train a diffusion model that learns the underlying task distribution and combine Thompson sampling with the learned prior to deal with new task at test time.

Original: Our posterior sampling algorithmis designed to carefully balance between the learned prior and the noisy observations that come from the learner’s interaction with the environment.
Revised: Our posterior sampling algorithm is designed to carefully balance between the learned prior and the noisy observations that come from the learner’s interaction with the environment.

Original: Preliminaryexperiments clearly demonstrate the potential of the considered approach.
Revised: Preliminary experiments clearly demonstrate the potential of the considered approach.

Original:  Introduction
Revised:  Introduction

Original: Uncertainty quantification is an integral part of online decision making and forms the basis of variousonline algorithms that trade-off exploration against exploitation.
Revised: Uncertainty quantification is an integral part of online decision making and forms the basis of various online algorithms that trade-off exploration against exploitation.

Original: Among these methods, Bayesianapproaches allow us to quantify the uncertainty using probability distributions, with the help of thepowerful tools of Bayesian inference.
Revised: Among these methods, Bayesian approaches allow us to quantify the uncertainty using probability distributions, with the help of the powerful tools of Bayesian inference.

Original: Nonetheless, their performance are known to be sensitive tothe choice of prior.
Revised: Nonetheless, their performance is known to be sensitive to the choice of prior.

Original: For concreteness, let us consider the problem of stochastic multi-armed bandits (MABs) [2, 9], inwhich a learner repeatedly pulls one of the K arms from a given set 
Revised: For concreteness let us consider the problem of stochastic multi-armed bandits (MABs) [ in which a learner repeatedly pulls one of the K arms from a given set A = {

Original: A = { 1 , ..., K } and receivesrewards that depend on the learner’s choices.
Revised: K } and receives rewards that depend on the learner’s choices.

Original: More precisely, when arm a is pulled at round t , thelearner receives reward r t drawn from an arm-dependent distribution P a .
Revised: More precisely, when arm a is pulled at round t the learner receives reward r t drawn from an arm-dependent distribution P a .

Original: The goal of the learner iseither to i ) accumulate the highest possible reward over time (a.k.a.
Revised: The goal of the learner is either to i ) accumulate the highest possible reward over time (a.k.a.

Original: regret-minimization) or to ii ) findthe arm with the highest expected reward within a prescribed number of rounds. (a.k.a.
Revised: regret-minimization) or to ii ) find the arm with the highest expected reward within a prescribed number of rounds (a.k.a.

Original: best-armidentification).
Revised: best-arm identification).

Original: For both purposes, we need to have a reasonable estimate of the arms’ mean rewards µ a =23 E r a ∼P a [ r a ] .
Revised: For both purposes we need to have a reasonable estimate of the arms’ mean rewards µ a = E r a ∼P a [ r a ] .

Original: In general, this would require us to pull each arm a certain number of times, whichbecomes inefficient when K is large.
Revised: In general, this would require us to pull each arm a certain number of times, which becomes inefficient when K is large.

Original: Moreover, unfortunately, the no-free-lunch principle precludes the possibility of improving upon this bottleneck, since an algorithm thatperforms better on one instance necessarily sacrifices its performance on the other. In spite of the said limitation, the bandit instances (referred as tasks hereinafter) that we encounter in most practical situations are far from arbitrary.
Revised: While the no-free-lunch principle prevents us from improving upon this bottleneck in general situations, it is worth noticing that the bandit instances (referred as tasks hereinafter) that we encounter in most practical problems are far from arbitrary.

Original: To name a few examples, in recommendation systems, each task corresponds to a user with certain underlying preferences that affect how much they like each item; in online shortest path routing we operate in real-world networks that feature specific characteristics.
Revised: To name a few examples, in recommendation systems, each task corresponds to a user with certain underlying preferences that affect how much they like each item; in online shortest path routing, we operate in real-world networks that feature specific characteristics.

Original: In this regard, introducing such inductive bias to the learning algorithm would be beneficial.
Revised: In this regard, introducing such inductive bias to the learning algorithm would be beneficial.

Original: In Bayesian models, this can be expressed through the choice of the prior distribution.
Revised: In Bayesian models, this can be expressed through the choice of the prior distribution.

Original: Moreover, as suggested by the mete-learning paradigm the prior itself can also be learned from data which oftenleads to superior performance [6, 12]
Revised: Moreover, as suggested by the mete-learning paradigm, the prior itself can also be learned from data, which often leads to superior performance [7, 13].

Original: Our contributions.
Revised: Our contributions.

Original: This work tackles the problem of meta-learning the priors for bandits [1, 3, 10],using Thompson sampling  as the base algorithm in each task.
Revised: This work tackles the problem of meta-learning a prior for bandits [1, 3, 11], with Thompson sampling used as the base algorithm in each task.

Original: In order to approximate the complexprior that arise in practice, we build upon the powerful tools of deep generative modeling, whoserecent progress have enabled impressive results in various areas ranging from image generation [14]to protein design [20].
Revised: In order to approximate the complex priors that arise in practice, we build upon the powerful tools of deep generative modeling, whose recent progress have enabled impressive results in various areas ranging from image generation [15] to protein design [21].

Original: Specifically, we propose to meta-learn the prior using denoising diffusionmodels [5, 16], and develop an algorithm to perform Thompson sampling under the learned prior.
Revised: Specifically, we propose to meta-learn the prior using denoising diffusion models [6, 17], and develop an algorithm to perform Thompson sampling under the learned prior.

Original: The designed algorithm strikes a delicate balance between the learned prior and bandit observations,bearing in mind the importance of having an accurate uncertainty estimate.
Revised: The designed algorithm strikes a delicate balance between the learned prior and bandit observations, bearing in mind the importance of having an accurate uncertainty estimate.

Original: Through syntheticexperiments, we demonstrate the benefit of the considered approach against several baseline methods.
Revised: Through synthetic experiments, we demonstrate the benefit of the considered approach against several baseline methods.

Original:  Preliminaries and Problem Description
Revised:  Preliminaries and Problem Description

Original: In this section, we briefly review denoising diffusion models and introduce our meta-learningframework over a class of bandit tasks.
Revised: In this section, we briefly review denoising diffusion models and introduce our meta-learning framework over a class of bandit tasks.

Revised: The associated pseudo-codes can be found in Appendix A.

Original: 2.1 Denoising Diffusion Probabilistic Model
Revised: 2.1 Denoising Diffusion Probabilistic Model

Original: First introduced by Sohl-Dickstein et al.
Revised: First introduced by Sohl-Dickstein et al.

Original: [16] and recently popularized by Ho et al.
Revised: [17] and recently popularized by Ho et al.

Original: [5] and Song and
Revised: [6] and Song and

Original: Ermon [17], denoising diffusion models (or the closely related score-based models) have been shownto achieve state-of-the-art performance in various data generation tasks.
Revised: Ermon [18], denoising diffusion models (or the closely related score-based models) have been shown to achieve state-of-the-art performance in various data generation tasks.

Original: Numerous variants of thesemodels have been proposed.
Revised: Numerous variants of these models have been proposed.

Original: Below, we mainly adopt the notations and formulation of Ho et al.
Revised: Below, we mainly adopt the notations and formulation of Ho et al.

Original: [5],with minimal modifications to adapt it to our purpose.
Revised: [6], with minimal modifications to adapt it to our purpose.

Original: Intuitively speaking, diffusion models learn to approximate a distribution Q 0 by training a series ofdenoisers with samples drawn from this distribution.
Revised: Intuitively speaking, diffusion models learn to approximate a distribution Q 0 by training a series of denoisers with samples drawn from this distribution.

Original: Writing q for the probability density function(assume everything is Lebesgue measurable for simplicity) and X 0 for the associated random variable,we define the forward diffusion process with respect to a sequence of scale factors ( α ℓ ) ∈ (0 , 1) L byq ( x 1: L | x 0 ) =
Revised: Writing q for the probability density function (assume everything is Lebesgue measurable for simplicity) and X 0 for the associated random variable, we define the forward diffusion process with respect to a sequence of scale factors ( α ℓ ) ∈ (0 , 1) L by q ( x 1: L | x 0 ) =

Original: L − 1 (cid:89) ℓ = q ( x ℓ +1 | x ℓ) , q ( X ℓ +1 | x ℓ ) =
Revised: L − 1 (cid:89) ℓ = q ( x ℓ +1 | x ℓ

Original: N ( X ℓ +1 ; √ α ℓ +1 x ℓ , (1 − α ℓ +1 ) I ) .
Revised: [Equation]

Original: The first equality suggests that the forward process is Markovian, while the second equality impliesthat the transition kernel is Gaussian.
Revised: The first equality suggests that the forward process is Markovian, while the second equality implies that the transition kernel is Gaussian.

Original: Further denoting the product of the scale factors by ¯ α ℓ =59 (cid:81) ℓi
Revised: Further denoting the product of the scale factors by ¯ α ℓ = (cid:81) ℓi

Original: =

Revised: =

Original: α i , we then have q ( X ℓ | x 0 ) =
Revised: α i , we then have q ( X ℓ | x 0 ) =

Original: N ( X ℓ ; √ ¯ α ℓ x 0 , (1 − ¯ α ℓ ) I ) .
Revised: N ( X ℓ ; √ ¯ α ℓ x 0 , (1 − ¯ α ℓ ) I ) .

Original: The sequence ( α ℓ ) ∈ (0 , 1) L is chosen to be decreasing and such that ¯ α L ≈ 0 .
Revised: The sequence ( α ℓ ) ∈ (0 , 1) L is chosen to be decreasing and such that ¯ α L ≈ 0 .

Original: We thus expectq ( X ℓ ) ≈ N (0 , 1) .
Revised: We thus expect q ( X ℓ ) ≈ N (0 , 1) .

Original: A denosing diffusion model learns to reverse the diffusion process with adistribution P θ over random variables X ′ 

Revised: A denosing diffusion model learns to reverse the diffusion process with a distribution P θ over random variables X ′

Original: L , in the hope that the marginal distribution P θ ( X ′ 0 ) is agood approximation of Q 0 .
Revised: L , in the hope that the marginal distribution P θ ( X ′ 0 ) is a good approximation of Q 0 .

Original: This is achieved by setting p θ ( X ℓ ) =
Revised: This is achieved by setting p θ ( X ℓ ) =

Original: N (0 , 1) , enforcing the learnedreverse process to be Markovian, and modeling p θ ( X ℓ | x ℓ +1 ) as a Gaussian parameterized byp θ ( X [Equation] | [Equation]
Revised: N (0 , 1) , enforcing the learned reverse process to be Markovian, and modeling p θ ( X ℓ | x ℓ +1 ) as a Gaussian parameterized by

Original: are

Original: [Equation]
Revised: [Equation]

Original: In the above h θ is the learned denoiser and h θ ( x ℓ +1 , ℓ + 1) is the predicted clean sample.
Revised: In the above h θ is the learned denoiser and h θ ( x ℓ +1 , ℓ + 1) is the predicted clean sample.

Original: 2.2 Meta-Learning of Bandit Tasks
Revised: 2.2 Meta-Learning of Bandit Tasks

Original: We consider a meta-learning for bandits framework in which the bandit tasks are drawn from anunderlying distribution T .
Revised: We consider a meta-learning for bandits framework in which the bandit tasks are drawn from an underlying distribution T .

Original: Unlike [3, 10], we focus on the a more standard scenario which featuresa meta-train and a meta-test phase.
Revised: Unlike [3, 11], we focus on the a more standard scenario which features a meta-train and a meta-test phase.

Original: The goal is to learn a good prior using the training set suchthat we can perform better once the algorithm is deployed with the learned prior in the test phase.
Revised: The goal is to learn a good prior using the training set such that we can perform better once the algorithm is deployed with the learned prior in the test phase.

Original: Mathematically, we can characterize the performance of the algorithm with the so-called transferregret [3], defined by Reg T ( π ) =
Revised: Mathematically, we can characterize the performance of the algorithm with the so-called transfer regret [3], defined by Reg T ( π ) =

Original: E B ∼T Reg T ( π, B ) where π is the algorithm that uses the learnedprior and Reg T ( π, B ) is the regret incurred by the algorithm within task B (see Appendix B).
Revised: E B ∼T Reg T ( π, B ) where π is the algorithm that uses the learned prior and Reg T ( π, B ) is the regret incurred by the algorithm within task B (see Appendix C).

Original: Throughout the work, we assume that the noise in the rewards are Gaussian with known variance σ 2 .
Revised: Throughout the work, we assume that the noise in the rewards are Gaussian with known variance σ 2 .

Original: The learner thus only needs to learn the vector of the mean rewards µ = ( µ a ) a ∈A .
Revised: The learner thus only needs to learn the vector of the mean rewards µ = ( µ a ) a ∈A .

Original: Accordingly, theprior is defined as a distribution in R K .
Revised: Accordingly, the prior is defined as a distribution in R K .

Original:  Algorithm
Revised:  Algorithm

Original: In this section, we describe our backbone algorithm for Thompson sampling with diffusion prior.
Revised: In this section, we describe our backbone algorithm for Thompson sampling with diffusion prior.

Original: 3.1 Diffusion Model as Prior
Revised: 3.1 Diffusion Model as Prior

Original: Our training data for the diffusion prior can have different forms.
Revised: Our training data for the diffusion prior can have different forms.

Original: For example, these data may becomposed of the interaction history with tasks using some bandit algorithms.
Revised: For example, these data may be composed of the interaction history with tasks using some bandit algorithms.

Original: How to leverage suchconfounded and noisy data to learn a diffusion model is a challenging problem.
Revised: How to leverage such confounded and noisy data to learn a diffusion model is a challenging problem.

Original: In this short paper,we simply assume that the mean vectors of these tasks are known, from which we can apply thestandard diffusion model training process.
Revised: In this short paper, we simply assume that the mean vectors of these tasks are known, from which we can apply the standard diffusion model training process.

Original: Once the diffusion model is learned, we may use it directly as a prior in any downstream tasks.
Revised: Once the diffusion model is learned, we may use it directly as a prior in any downstream tasks.

Original: Nonetheless, the original sampling process of the diffusion model does not take the quality of thelearned model into account, and overly relying on a bad prior does more harm than good.
Revised: Nonetheless, the original sampling process of the diffusion model as described in [6] falls short in providing a good estimate of the variances that reflects the right level of uncertainty.

Original: To addressthis problem, we propose to inject more noise in the sampling process when the model is of poorerquality.
Revised: To address this problem, we propose to further calibrate the variances of the reverse process with a separate calibration set V .

Original: Precisely, we write
Revised: For this, we write

Original: [Equation]
Revised: [Equation]

Original: In the above, p ′ θ ( X 0 | x ℓ +1 ) is a Gaussian distribution centered at ˆ x 

Revised: In the above, p ′ θ ( X 0 | x ℓ +1 ) is a Gaussian distribution centered at ˆ x

Original: = h θ ( x ℓ +1 , ℓ + 1) .
Revised: = h θ ( x [Equation] +1 , [Equation]

Original:  This is91 different from (1) where by analogy we may interpret p ′ θ ( x 0 | x ℓ +1 ) dx 0 as a Dirac.
Revised: + 1) . This is different from (1) where by analogy we may interpret p ′ θ ( x 0 | x ℓ +1 ) dx 0 as a Dirac.

Original: As learning thecovariance ofthe reverse process has been shown to result in poor performance in [5], the covarianceof p ′ θ ( X 0 | x ℓ +1 ) is simply taken as a diagonal matrix diag( τ 2 ℓ +1 ) whose diagonal elements are the(coordinate-wise) mean squared reconstruction errors of the model on a calibration set V .
Revised: The covariance of p ′ θ ( X 0 | x ℓ +1 ) is then taken as a diagonal matrix diag( τ 2 ℓ +1 ) whose diagonal elements are the (coordinate-wise) mean squared reconstruction errors of the model on the calibration set V .

Original: Thatis, we construct ˜ V ℓ containing pairs ( x 0 , x ℓ ) with x 0 ∈ V and x ℓ sampled from X ℓ | x 0 , and setτ ℓ( i )
Revised: That is, we construct ˜ V ℓ containing pairs ( x 0 , x ℓ ) with x 0 ∈ V and x ℓ sampled from X ℓ | x 0 , and set τ aℓ = (cid:113)(cid:80) x 0 ,x ℓ ∈ ˜ V ℓ

Original: [Equation] = [Equation] (˜ V ℓ ) .
Revised: ∥ x a 0 − h aθ ( x ℓ , ℓ ) ∥ 2 / card(˜ V ℓ ) .

Original: Intuitively, thisautomatically adjusts how muchwe rely on the prior in the upcoming tasks by taking the reconstruction error as a proxy for the qualityof the learned model.
Revised: Intuitively, this adjusts how much we rely on the learned model in the upcoming tasks by taking the reconstruction error as a proxy for its quality.

Original: 3.2 Thompson Sampling with Diffusion Prior
Revised: 3.2 Thompson Sampling with Diffusion Prior

Original: Thompson sampling [13, 19] is one of the most popular strategies for tackling stochastic bandits dueto its simplicity and generality.
Revised: Thompson sampling [14, 20] is one of the most popular strategies for tackling stochastic bandits due to its simplicity and generality.

Original: This strategy takes as input a prior distribution for the parameterof interest, and samples a guess of the parameter from the posterior distribution at each round todetermine which arm to pull.
Revised: It takes as input a prior distribution for the parameter of interest, and samples a guess of the parameter from the posterior distribution at each round to determine which arm to pull.

Original: In the MAB setting in Section 1, the parameter of interest is the vectorof mean rewards µ .
Revised: In the MAB setting in Section 1, the parameter of interest is the vector of mean rewards µ .

Original: At each round t , a vector ˜ µ t isthen sampled from the posterior distributiondetermined by the prior and the interaction history ( a s , r s ) s ∈{ 1 ,...,t − 1 } .
Revised: At each round t , a vector ˜ µ t is sampled from the posterior distribution determined by the prior and the interaction history ( a s , r s ) s ∈{ 1 ,...,t − 1 } .

Original: Subsequently, the learner106 pulls an arm a t ∈ arg max a ∈A ˜ µ at .
Revised: Subsequently, the learner pulls an arm a t ∈ arg max a ∈A ˜ µ at .

Original: Our algorithm takes a prior learned by a diffusion model as input. We thus need to design an algorithmthat samples from the posterior distribution when the prior is described in this way.
Revised: In the following we design an algorithm that samples from the posterior distribution when the prior is described by a diffusion model.

Original: While an exactsolution does not exist in general, we aim to find a good approximation that balances well betweenthe meta-learned prior and noisy observations.
Revised: While an exact solution does not exist in general, we aim to find a good approximation that balances well between the meta-learned prior and noisy observations.

Original: In contrast, most existing approaches for posteriorsampling with diffusion models focus on the solution of inverse problems [4, 7, 16, 18], and do nottake the underlying uncertainty into account.
Revised: In contrast, most existing approaches for posterior sampling with diffusion models focus on the solution of inverse problems [5, 8, 17, 19], and do not take the underlying uncertainty into account.

Original: To proceed, the general idea of our algorithm is to guide the sample towards the observation duringthe sampling process.
Revised: To proceed, the general idea of our algorithm is to guide the sample towards the observation during the sampling process.

Original: Formally, let us consider Y 0 as a random observation of X 0 with knownq ( Y 0 | X 0 ) .
Revised: Formally, let us consider Y 0 as a random observation of X 0 with known q ( Y 0 | X 0 ) .

Original: For given y 0 , we are interested in sampling from X 0 | y 0 .
Revised: For given y 0 , we are interested in sampling from X 0 | y 0 .

Original: We achieve this by goingthrough the reverse Markovian process, conditioning on Y 0 = y 0 .
Revised: We achieve this by going through the reverse Markovian process, conditioning on Y 0 = y 0 .

Original: In fact, to sample from X ℓ | y 0 , weonly need to first sample x ℓ +1 from X ℓ +1 | y 

Revised: In fact, to sample from X ℓ | y 0 , we only need to first sample x ℓ +1 from X ℓ +1 | y

Original: and then sample x ℓ from X ℓ | x ℓ +1 , y 0 .
Revised: and then sample x ℓ from X ℓ | x ℓ +1 , y 0 .

Original: Repeating theprocess for ℓ =
Revised: Repeating the process for ℓ =

Original: L − 1 , . .
Revised: L − 1 , . .

Original: . , 0 then gives the desired result.
Revised: . , 0 then gives the desired result.

Original: Below we briefly explain the initializationand the recursive steps of our method.
Revised: Below we briefly explain the initialization and the recursive steps of our method.

Original: Detailed derivation of the algorithm along with alternativeapproaches are discussed in Appendix A.
Revised: Detailed derivation of the algorithm along with alternative approaches are discussed in Appendix B.

Original: Sampling from X L | y 0 .
Revised: Sampling from X L | y 0 .

Original: For this part, we simply ignore y 0 and sample from N (0 , 1) as before.
Revised: For this part, we simply ignore y 0 and sample from N (0 , 1) as before.

Original: Sampling from X ℓ | x ℓ +1 , y
Revised: Sampling from X ℓ | x ℓ +1 , y

Original: For simplicity, we restrict our attention to the case of multi-armedbandit.
Revised: For simplicity, we restrict our attention to the case of multi-armed bandit.

Original: Here, y 0 is the interaction history and x 0 = µ is the vector of mean rewards, leading to 
Revised: Here, y 0 is the interaction history and x 0 = µ is the vector of mean rewards, leading to [Equation] [Equation] a

Original: [Equation]
Revised: [Equation]

Revised: x

Revised: [Equation]

Revised: [Equation]

Original: We distinguish between two situations.
Revised: We distinguish between two situations.

Original: • Arm a has never been pulled in the first t rounds: We sample the corresponding coordinate x aℓ from the Gaussian distribution ˜ q ( x aℓ | x ℓ +1 , y 0 )
Revised: • Arm a has never been pulled in the first t rounds: We sample the corresponding coordinate x aℓ from the Gaussian distribution ˜ q ( x aℓ | x ℓ +1 , y 0 )

Original: =
Revised: =

Original: p θ ( x ℓ | x ℓ +1 ) introduced in (2).
Revised: p θ ( x ℓ | x ℓ +1 ) introduced in (2).

Original: • Arm a has been pulled in the first t rounds: We denote by ˆ µ at = (cid:80) ts =

Revised: • Arm a has been pulled in the first t rounds: We denote by ˆ µ at = (cid:80) ts =

Original: r s 1 { a s = a } /N
Revised: r s 1 { a s = a } /N

Original: at as the empirical mean and σ at = σ/ (cid:112) N at as the adjusted standard deviation, where N at is the numberof times that arm a has been pulled up to time t (included).
Revised: at as the empirical mean and σ at = σ/ (cid:112) N at as the adjusted standard deviation, where N at is the number of times that arm a has been pulled up to time t (included).

Original: We also write ¯ z ℓ +1 for the noise predicted by the denoiser from x ℓ +1 .
Revised: We also write ¯ z ℓ +1 for the noise predicted by the denoiser from x ℓ +1 .

Original: Then, with ρ ℓ = ¯ α ℓ +1 (1 − ¯ α ℓ ) /
Revised: Then, with ρ ℓ = ¯ α ℓ +1 (1 − ¯ α ℓ ) /

Original: (¯ α ℓ (1 − ¯ α ℓ +1 )) , we sample from the Gaussian distribution ˜ q ( x aℓ | x ℓ +1 , y 0 ) satisfying that ˜ q ( x aℓ | x ℓ +1 , y 0 )
Revised: (¯ α ℓ (1 − ¯ α ℓ +1 )) , we sample from the Gaussian distribution ˜ q ( x aℓ | x ℓ +1 , y 0 ) satisfying that ˜ q ( x aℓ | x ℓ +1 , y 0 )

Original: ∝ p θ ( x aℓ | x ℓ +1 ) N (cid:0) x aℓ ; √ ¯ α ℓ ˆ µ at + √ 1 − ¯ α ℓ ¯ z aℓ +1 , ¯ α ℓ (( σ at ) 2 + ρ ℓ ( τ aℓ +1 ) 2 ) (cid:1) .
Revised: ∝ p θ ( x aℓ | x ℓ +1 ) N (cid:0) x aℓ ; √ ¯ α ℓ ˆ µ at + √ 1 − ¯ α ℓ ¯ z aℓ +1 , ¯ α ℓ (( σ at ) 2 + ρ ℓ ( τ aℓ +1 ) 2 ) (cid:1) .

Original: In words, we sample x ℓ from a Gaussian whose mean is a weighted average of the mean ofp θ ( x aℓ | x
Revised: In words, we sample x ℓ from a Gaussian whose mean is a weighted average of the mean of p θ ( x aℓ | x

Original: ℓ
Revised: ℓ

Original: +1 ) and a noisy version of the empirical mean ˆ µ t (only defined for those arms that havebeen pulled).
Revised: +1 ) and a noisy version of the empirical mean ˆ µ t (only defined for those arms that have been pulled).

Original: Importantly, the noisy version of the observation is computed with the predictednoise, which distinguishes our work from existing ones such as [18].
Revised: Importantly, the noisy version of the observation is computed with the predicted noise, which distinguishes our work from existing ones such as [19].

Original:  Numerical Experiments
Revised:  Numerical Experiments

Original: [Figure]
Revised: [Figure]

Original: In this section, we illustrate the benefit of using diffusionprior through two synthetic experiments.
Revised: In this section, we illustrate the benefit of using diffusion prior through two synthetic experiments.

Original: Missing details canbe found in Appendix B.
Revised: Missing details can be found in Appendix C.

Original: Problem construction.
Revised: Problem construction.

Original: For the two problems, we fix σ =0 .
Revised: For the two problems, we fix σ = 0 .

Original: 1 and construct the means as follows
Revised: 1 and construct the means as follows

Original: Labeled Arms Problem.
Revised: Labeled Arms Problem.

Original: Let K = 500 .
Revised: Let K = 500 .

Original: In this problem,each arm is associated to 7 of the 50 labels.
Revised: In this problem, each arm is associated to 7 of the 50 labels.

Original: For eachbandit task, we randomly pick 7 labels; the expectedreward of an arm is positively correlated to the number of labels that fall into the intersection of the arm’s labels and the task’s labels.
Revised: For each bandit task, we randomly pick 7 labels; the expected reward of an arm is positively correlated to the number of labels that fall into the intersection of the arm’s labels and the task’s labels.

Original: This is a simplest model for modeling user preferences over a list of items.
Revised: This is a simplest model for modeling user preferences over a list of items.

Original: Popular and Niche Problem.
Revised: Popular and Niche Problem.

Original: Let K
Revised: Let K

Original: = 200 .
Revised: = 200 .

Original: In this problem, the arms are separated into 20 groups, each of size 10 .
Revised: In this problem, the arms are separated into 40 groups, each of size 5 .

Original: Arms of the same groups have close expected rewards. Moreover, 10 groups of arms tend to have highexpected rewards but these arms are never the optimalones.
Revised: Among these, 20 groups of arms tend to have high expected rewards but these arms are never the optimal ones.

Original: The other 10 groups of arms have lower expected rewards in general but contain a group ofarms that yield much higher rewards.
Revised: The other 20 groups of arms have lower expected rewards in general but contain the optimal arm.

Original: This represents a situation where Gaussian prior can lead to poor performance due to its inability to model multi-modal distributions
Revised: This represents a situation where Gaussian prior can lead to poor performance due to its inability to model multi-modal distributions

Original: Training and baselines.
Revised: Training and baselines.

Original: In terms of algorithms, we compare our method, DiffTS, with UCB and
Revised: In terms of algorithms, we compare our method, DiffTS, with UCB and

Original: Thompson sampling with Gaussian prior using either diagonal or full covariance matrix (GTS-diagand GTS-full).
Revised: Thompson sampling with Gaussian prior using either diagonal or full covariance matrix (GTS-diag and GTS-full).

Original: For each algorithm, we test two configurations: either the algorithm is given the truestandard deviation σ ′ = 
Revised: For each algorithm, we test two configurations: either the algorithm is given the true standard deviation σ ′ = 0 .

Original: 1 or an underspecified standard deviation σ ′ = 0 .
Revised: 1 or an underspecified standard deviation σ ′ = 0 .

Original: To train the diffusion model, we use a training set of 5000 samples and a calibration set of 1000samples.
Revised: To train the diffusion model, we use a training set of 5000 samples and a calibration set ofsamples.

Original: The 6000 samples put together are also used to compute the covariance matrices used bythe Gaussian Thompson sampling algorithms.
Revised: The 6000 samples put together are also used to compute the means and the covariances used by the Gaussian Thompson sampling algorithms.

Original: To test the performance of the algorithms, we thensample another 100 bandit tasks and run the aforementioned algorithms on these tasks.
Revised: To test the performance of the algorithms, we sample another 100 bandit tasks and run the aforementioned algorithms on these tasks.

Original: Results.
Revised: Results.

Original: We report the regrets of the algorithms averaged over the 100 test tasks in Figure 1 (thetop and the bottom figures correspond respectively to the Labeled Arms and the Popular and
Revised: We report the regrets of the algorithms averaged over the 100 test tasks in Figure 1 (the top and the bottom figures correspond respectively to the Labeled Arms and the Popular and

Original: Niche problem).
Revised: Niche problem).

Original: We verify that using a diffusion prior which describes better the task distributionindeed helps achieve smaller regret.
Revised: We verify that using a diffusion prior which describes better the task distribution indeed helps achieve smaller regret.

Original: However, the amount of improvement varies.
Revised: However, the amount of improvement varies.

Original: In the Labeled
Revised: In the Labeled

Original: Arms problem, learning the correlation between arms already lower the regret significantly, and theassumed noise standard deviation σ ′ seems to have a greater impact.
Revised: Arms problem, learning the correlation between arms already lower the regret significantly and the assumed noise standard deviation σ ′ seems to have a greater impact.

Original: In the Popular and Nicheproblem, using a Gaussian prior with full covariance matrix however leads to poor performance, andthe ability of denoising diffusion models to model complex distribution helps greatly here. 
Revised: In the Popular and Niche problem, using a Gaussian prior with full covariance matrix however leads to poor performance and the ability of denoising diffusion models to model complex distribution helps greatly here. A Missing Pseudo Codes 1: Meta Training Deployment

Original: A Posterior Sampling: Algorithms and Derivation
Revised: 1: Input: Learned denoiser (diffusion model) h θ , reconstruction errors ( τ 2 ℓ ) ℓ ∈{ 1 ,...,L } , presumed noise standard deviation σ ′ 2: for t do 3: Sample x L ∼ N ( I ) 4: for ℓ ∈ L −

Revised: 0 do 5:

Original: {+Predict clean sample ˆ x 
+}
Revised: {+Predict clean sample ˆ x 
+}

Revised: =

Revised: h θ x ℓ ℓ and associated noise ¯ z ℓ + 6:

Revised: Compute diffused observation ˜ y aℓ =

Revised: √ ¯ α ℓ ˆ µ at − 1 + √ 1 − ¯ α ℓ ¯ z ℓ + 7: for a ∈ A do 8:

Revised: If N at − 1 = 0 , sample x aℓ ∼ p θ ( X aℓ | x ℓ +1 ) 9:

Revised: If N at − 1 > 0 , sample x aℓ ∼ ˜

Revised: [Equation]

Revised: ( X aℓ

Revised: [Equation]

Revised: x ℓ +1 , y

Revised: [Equation]

Revised: ∝ p

Revised: [Equation]

Revised: ( X aℓ

Revised: [Equation]

Revised: x ℓ +1 ) N (cid:0) X aℓ ; ˜ y aℓ , ¯

Revised: [Equation]

Revised: ℓ (( σ at ) 2 + ρ ℓ ( τ aℓ +1 ) 2 (cid:1) 10: end for 11: end for 12:

Revised: Pull arm a t ∈ arg max a ∈A x a13:

Revised: Update number

Revised: 14: end for

Revised: B Posterior Sampling: Algorithms and Derivation

Original: Below we provide two approximations to sample from X L | x ℓ +1 , y 0 .
Revised: Below we provide two approximations to sample from X L | x ℓ +1 , y 0 .

Original: The first one is easier to derivebut the second one consistently achieves better performance in our experiments.
Revised: The first one is easier to derive but the second one consistently achieves better performance in our experiments.

Original: 3 We thus utilize thesecond approach in Sections 3 and 4.1.
Revised: 4 We thus utilize the second approach in Sections 3 and 4.

Original: Approach 1, acting on x 0 .
Revised: Approach 1, acting on x 0 .

Original: We write
Revised: We write

Original: [Equation]
Revised: [Equation]

Original: As in Section 3.1, we approximate q ( x 0 | x ℓ +1 ) by
Revised: As in Section 3. [Equation] we approximate q ( x 0 | x ℓ +1 ) by

Original: [Equation]
Revised: [Equation]

Original: Then, if q ( Y 0 | x 0 ) is also Gaussian, the integral in (4) with q ( x 0 | x ℓ +1 ) replaced by p ′
Revised: Then, if q ( Y 0 | x 0 ) is also Gaussian, the integral in (4) with q ( x 0 | x ℓ +1 ) replaced by p ′

Original: θ ( x 0 | x ℓ +1 )can be computed in close form.
Revised: θ ( x 0 | x ℓ +1 ) can be computed in close form.

Original: For further illustration, we focus on the case of multi-armed bandits where the relation betweeny 0 the interaction history and x 0 = µ the mean reward vector is given by (3).
Revised: For further illustration, we focus on the case of multi-armed bandits where the relation between y 0 the interaction history and x 0 = µ the mean reward vector is given by (3).

Original: The approximatedistribution of X ℓ | x ℓ +1 , y 0 is thus independent across coordinates (arms) and we can computethe density individually for each arm, which we denote by ˜ q ( x aℓ | x ℓ +1 , y 0 ) .
Revised: The approximate distribution of X ℓ | x ℓ +1 , y 0 is thus independent across coordinates (arms) and we can compute the density individually for each arm, which we denote by ˜ q ( x aℓ | x ℓ +1 , y 0 ) .

Original: It is computed using(4) with q ( x 0 | x ℓ +1 )
Revised: It is computed using (4) with q ( x 0 | x ℓ +1 )

Original: ≈ p
Revised: ≈ p

Original: ′ θ ( x 0 | x ℓ +1 )
Revised: ′ θ ( x 0 | x ℓ +1 )

Original: followed by a normalization.
Revised: followed by a normalization.

Original: When arm a has never been pulled, ˜ q ( x aℓ | x ℓ +1 , y 0 ) is nothing but p θ ( x aℓ | x ℓ +1 ) .
Revised: When arm a has never been pulled, ˜ q ( x aℓ | x ℓ +1 , y 0 ) is nothing but p θ ( x aℓ | x ℓ +1 ) .

Original: Otherwise,˜ q ( x aℓ | x ℓ +1 , y 0 ) is computed first using the fact that
Revised: Otherwise, ˜ q ( x aℓ | x ℓ +1 , y 0 ) is computed first using the fact that

Original: [Equation]
Revised: [Equation]

Original:   .

Original: Next, with
Revised: Next, with

Original: [Equation]
Revised: [Equation]

Original: 
Revised: 

Original: [Equation]
Revised: [Equation]

Original: we see that we just need to replace x a 0 by ˆ x a 0 and augment the variance by π 21 (ˆ σ a ) 2 when samplingx aℓ .
Revised: we see that we just need to replace x a 0 by ˆ x a 0 and augment the variance by π 21 (ˆ σ a ) 2 when sampling x aℓ .

Original: Compared to the case where y 0 is not given, the main difference lies in that we first changethe estimated ˆ x 0 before sampling x ℓ .2.
Revised: Compared to the case where y 0 is not given, the main difference lies in that we first change the estimated ˆ x 0 before sampling x ℓ .

Original: Approach 2, acting on x ℓ .
Revised: Approach 2, acting on x ℓ .

Original: Alternatively, we may write
Revised: Alternatively, we may write

Original: [Equation]
Revised: [Equation]

Original: To begin, we use p θ ( x ℓ | x ℓ +1 ) to approximate q ( x ℓ | x ℓ +1 ) .
Revised: To begin, we use p θ ( x ℓ | x ℓ +1 ) to approximate q ( x ℓ | x ℓ +1 ) .

Original: Next, one natural way to tackle theintegral is to use q ( x 0 | x ℓ , x ℓ +1 ) =
Revised: Next, one natural way to tackle the integral is to use q ( x 0 | x ℓ , x ℓ +1 ) =

Original: q ( x 0 | x ℓ )
Revised: q ( x 0 | x ℓ )

Original: ≈ p
Revised: ≈ p

Original: ′ θ ( x 0 | x ℓ ) .
Revised: ′ θ ( x 0 | x ℓ ) .

Original: In the simplest case q ( y 0 | x 0 ) =
Revised: In the simplest case q ( y 0 | x 0 ) =

Original: N ( y 0 ; x 0 , σ 2 I ) , and we deduce
Revised: N ( y 0 ; x 0 , σ 2 I ) , and we deduce

Original: [Equation]
Revised: [Equation]

Original: Nonetheless, as the denoiser h θ can be arbitrarily complex, this does not lead to a close formexpression to sample x ℓ .
Revised: Nonetheless, as the denoiser h θ can be arbitrarily complex, this does not lead to a close form expression to sample x ℓ .

Original: We may resort to Langevin dynamics sampling, i.e., taking gradient stepson x ℓ to make sure the denoised result is coherent with y 0 .
Revised: We may resort to Langevin dynamics sampling, i.e., taking gradient steps on x ℓ to make sure the denoised result is coherent with y 0 .

Original: This is a general approach that may beof interest for arbitrary conditional distribution Y 0 | X 0 , but we have simpler solution when thenoise is Gaussian.
Revised: This is a general approach that may be of interest for arbitrary conditional distribution Y 0 | X 0 , but we have simpler solution when the noise is Gaussian.

Original: In fact, it holds that
Revised: In fact, it holds that

Original: [Equation]
Revised: [Equation]

Original: where both ¯ Z ℓ and Z ℓ +1 are random variable with distribution N (0 , I ) .
Revised: where both ¯ Z ℓ and Z ℓ +1 are random variable with distribution N (0 , I ) .

Original: This results in
Revised: This results in

Original: [Equation]
Revised: [Equation]

Original: for
Revised: for

Original: [Equation]
Revised: [Equation]

Original: Therefore, we may take ¯ Z ℓ +1 as a reasonable approximation of ¯ Z ℓ , while sampling ¯ Z ℓ +1 isbasically the same as sampling from p ′
Revised: Therefore, we may take ¯ Z ℓ +1 as a reasonable approximation of ¯ Z ℓ , while sampling ¯ Z ℓ +1 is basically the same as sampling from p ′

Original: θ ( X 0 | x ℓ +1 ) .
Revised: θ ( X 0 | x ℓ +1 ) .

Original: To summarize, we write
Revised: To summarize, we write

Original: [Equation]
Revised: [Equation]

Original: | x ℓ , x ℓ [Equation] where ρ ℓ = ¯ α ℓ +1 (1 − ¯ α ℓ ) / (¯ α ℓ (1 − ¯ α ℓ +1 )) and ¯ z ℓ +1 represents the noise predicted by the denoiserfrom x ℓ +1 , that is,
Revised: where ρ ℓ = ¯ α ℓ +1 (1 − ¯ α ℓ ) / (¯ α ℓ (1 − ¯ α ℓ +1 )) and ¯ z ℓ +1 represents the noise predicted by the denoiser from x ℓ +1 , that is,

Original: [Equation]
Revised: [Equation]

Original: Then, focusing on the arm that has been pulled at least once and ignoring the multiplicativeconstant that does not depend on x ℓ, we get
Revised: Then, focusing on the arm that has been pulled at least once and ignoring the multiplicative constant that does not depend on x ℓ

Original: [Equation]
Revised: [Equation]

Original: Subsequently
Revised: we get

Original: [Equation]
Revised: [Equation]

Original:  we can approximate the posterior distribution of x ℓ using (5)
Revised: Subsequently, we can approximate the posterior distribution of x ℓ using (5)

Original: [Equation]
Revised: [Equation]

Original: This results in the algorithm that we present in Section 3. [Equation]
Revised: This results in the algorithm that we present in Section 3.2.

Original: B Missing Experimental Details
Revised: C Missing Experimental Details

Original: In this section, we provide missing experimental details mainly concerning the used diffusion modelsand the construction of the problem instances.
Revised: In this section, we provide missing experimental details mainly concerning the used diffusion models and the construction of the problem instances.

Original: All the simulations are run on an Amazon p3.2xlargeinstance equipped with an NVIDIA Tesla V100 GPU.
Revised: All the simulations are run on an Amazon p3.2xlarge instance equipped with an NVIDIA Tesla V100 GPU.

Original: Diffusion models.
Revised: Diffusion models.

Original: In our experiments, we set the diffusion steps of the diffusion models to L = 100and adopt a linear variance schedule from 1 − α 1 = 10 − 4 to 1 − α L = 0 . 
Revised: In our experiments, we set the diffusion steps of the diffusion models to L =and adopt a linear variance schedule from 1 − α 1 = 10 − 4 to 1 − α L = 0 . 1 .

Revised: Moreover, the models are trained to predict the clean sample x 0 instead of the noise ¯ z ℓ since it is reported in [4] that this leads to better performance when the data are binary, and as explained below and shown in Figuresand 3, the mean rewards of the two considered tasks are nearly binary.

Original: The denoiser itselfis a 1-dimensional U-Net adapted from [8, 11].
Revised: The denoiser itself is a 1-dimensional U-Net adapted from [9, 12].

Original: We use 5 residual blocks, each block containing 6residual channels.
Revised: We use 5 residual blocks, each block containing 6 residual channels.

Original: These numbers are rather arbitrary and do not seem to affect much our results.
Revised: These numbers are rather arbitrary and do not seem to affect much our results.

Original: More importantly, since the patterns that we want to learn do not have the spatial correlation thatconvolutional layers are designed for, we add a fully connected layer at the beginning to map theinput to a vector of size 128 × 6 , before reshaping these vectors into 6 channels and feeding them tothe convolutional layers.
Revised: More importantly, since the patterns that we want to learn do not have the spatial correlation that convolutional layers are designed for, we add a fully connected layer at the beginning to map the input to a vector of size 128 × 6 , before reshaping these vectors into 6 channels and feeding them to the convolutional layers.

Original: In a similar fashion, we also replace the last layer of the architecture bya fully connected layer.
Revised: In a similar fashion, we also replace the last layer of the architecture by a fully connected layer.

Original: We find out that these minimal modification already enable the model toperform well on our problems, but believe a thorough investigation into the architecture design wouldfurther benefit our approach.
Revised: We find out that these minimal modification already enable the model to perform well on our problems, but believe a thorough investigation into the architecture design would further benefit our approach.

Original: These diffusion models are then trained with Adam for 600 epochs, with a learning rate of 1 × 10 − 3 ,5 × 10 − 4 , and 2 × 10 − 4 , for the first, middle, and last 200 epochs. The batch sizeis fixed at 128 .
Revised: These diffusion models are then trained with Adam for 15000 steps, with a learning rate of 5 × 10 − 4 , with a batch size fixed at 128 (which corresponds to 384 epochs).

Original: Again, we believe our synthetic problems are simple enough for diffusion models so that thesehyperparameters should have minimal effect on the results.

Original: Construction of bandit instances.
Revised: Construction of bandit instances.

Original: We now give more details on how the mean reward vectorsare constructed in the two problems.
Revised: We now give more details on how the mean reward vectors are constructed in the two problems.

Original: Some illustrations of the constructed instances and the vectorsgenerated by the trained diffusion models are provided in Figures 2 and 3.1.
Revised: Some illustrations of the constructed instances and the vectors generated by the trained diffusion models are provided in Figures 2 and 3.

Original: Labeled Arms ( K = 500 ).
Revised: Labeled Arms ( K = 500 ).

Original: As described in Section 4, in our first problem we are provided aset of 50 labels L = { 1 , ..., 50 } .
Revised: As described in Section 4, in our first problem we are provided a set of 50 labels L = { 1 , ..., 50 } .

Original: Each arm is associated to a subset L a of these labels with sizecard( L a ) = 7 .
Revised: Each arm is associated to a subset L a of these labels with size card( L a ) = 7 .

Original: To sample a new bandit task B , we randomly draw a set L B ⊆ L again with size7 .
Revised: To sample a new bandit task B , we randomly draw a set L B ⊆ L again with size 7 .

Original: Then for each arm a , we set ¯ µ a = 1 − 1 / 4 card( L a ∩L B ) .
Revised: Then for each arm a , we set ¯ µ a = 1 − 1 / 4 card( L a ∩L B ) .

Original: Finally, to obtain the mean rewardsµ , we perturb the coordinates of ¯ µ by independent Gaussian noises of standard deviation 0 .
Revised: Finally, to obtain the mean rewards µ , we perturb the coordinates of ¯ µ by independent Gaussian noises of standard deviation 0 .

Original: 1 andscale the resulting vector to the range [0 , 1] .2.
Revised: 1 and scale the resulting vector to the range [0 , 1] .

Original: Popular and Niche ( K = 200 )
Revised: Popular and Niche ( K = 200 ).

Original: The arms are split into 20 groups of equal size.
Revised: The arms are split into 40 groups of equal size.

Original: Conceptually,10 of these groups represent the ‘popular’ items while the other 10 represent the ‘niche’ items.
Revised: Conceptually, 20 of these groups represent the ‘popular’ items while the other 20 represent the ‘niche’ items.

Original: For each bandit task, we first construct a vector ¯ µ whose coordinates’ values default to 0 .
Revised: For each bandit task, we first construct a vector ¯ µ whose coordinates’ values default to 0 .

Original: However, we  choose one group of niche items and set their values to 1 . For eachgroup of popularitems, with probability 0 .
Revised: However, we randomly choose 1 to 3 groups of niche items and the value of each of these items is set towith probability 0 .

Original: 8 their values are set to 0 .
Revised: 7 (independently across the selected items).

Original:  Then, to construct the mean reward vectorµ , werespectively perturb the values of the niche items and of the popular items by independent Gaussian noises with standard deviation of 0 . 1and of 0 .
Revised: Similarly, we randomly chooseto 17 groups of popular items and set their values to 0 . Then, to construct the mean reward vector µ , we perturb the values of ¯ µ by independent Gaussian noises with standard deviation of 0 . 1 .

Original: After that, we clip the values of thepopular items to make them smaller than 0 .
Revised: After that, we clip the values of the popular items to make them smaller than 0 .

Original: 95 and finally scale the  vector to the range [0 , 1] .
Revised: 95 and clip the entire vector to the range [0 , 1] .

Original: UCB.
Revised: UCB.

Original: The most standard implementation of the UCB algorithm sets the upper confidence bound to
Revised: The most standard implementation of the UCB algorithm sets the upper confidence bound to

Original: [Equation]
Revised: [Equation]

Original: Instead, in our experiments we use U at = ˆ µ at + σ ′ / (cid:112) N at .
Revised: Instead, in our experiments we use U at = ˆ µ at + σ ′ / (cid:112) N at .

Original: Eq.
Revised: Eq.

Original: (6) is more conservative than ourimplementation.
Revised: (6) is more conservative than our implementation.

Original: However, looking at Figure 1, we figure out that UCB without the log t factor isalready the most conservative algorithm in our experiments.
Revised: However, looking at Figure 1, we figure out that UCB without the log t factor is already the most conservative algorithm in our experiments.

Original: Therefore, while in the long term UCBwith (6) may achieve lower regret, with limited budget Thompson sampling with a wrong but goodenough prior would generally be preferable [15].
Revised: Therefore, while in the long term UCB with (6) may achieve lower regret, with limited budget Thompson sampling with a wrong but good enough prior would generally be preferable [16].

Original: Regret.
Revised: Regret.

Original: For sake of completeness, we state here the definition of the (pseudo-)regret with respectto a sequence of arms ( a t ) t ∈{ 1 ,...,T } in a single bandit task.
Revised: For sake of completeness, we state here the definition of the (pseudo-)regret with respect to a sequence of arms ( a t ) t ∈{ 1 ,...,T } in a single bandit task.

Original: Let a ⋆ ∈ arg max a ∈A µ a be an optimalarm.
Revised: Let a ⋆ ∈ arg max a ∈A µ a be an optimal arm.

Original: It is defined as Reg T= Tµ a ⋆ − (cid:80) Tt =1 µ a t .
Revised: It is defined as Reg T

Original: [Figure] [Figure]
Revised: [Equation] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure][Figure]
Revised: [Figure] [Figure]

Original: [Figure][Figure] [Figure] [Figure]
Revised: [Figure]