Title: Diffusion Models Meet Contextual Bandits with Large Action Spaces
Venue: NeurIPS 2024
Accepted: no
Keywords: Diffusion models

Bayesian bandit

Thompson sampling

Contextual bandit
Abstract: Efficient exploration in contextual bandits is crucial due to their large action space, where uninformed exploration can lead to computational and statistical inefficiencies. However, the rewards of actions are often correlated, which can be leveraged for more efficient exploration. In this work, we use pre-trained diffusion model priors to capture these correlations and develop diffusion Thompson sampling (dTS). We establish both theoretical and algorithmic foundations for dTS. Specifically, we derive efficient posterior approximations (required by dTS) under a diffusion model prior, which are of independent interest beyond bandits and reinforcement learning. We analyze dTS in linear instances and provide a Bayes regret bound highlighting the benefits of using diffusion models as priors. Our experiments validate our theory and demonstrate dTS's favorable performance.

Abstract

Efficient exploration in contextual bandits is crucial due to their large action space, where uninformed exploration can lead to computational and statistical inefficiencies. However, the rewards of actions are often correlated, which can be leveraged for more efficient exploration. In this work, we use pre-trained diffusion model priors to capture these correlations and develop diffusion Thompson sampling ( dTS ). We establish both theoretical and algorithmic foundations for dTS . Specifically, we derive efficient posterior approximations (required by dTS ) under a diffusion model prior, which are of independent interest beyond bandits and reinforcement learning. We analyze dTS in linear instances and provide a Bayes regret bound highlighting the benefits of using diffusion models as priors. Our experiments validate our theory and demonstrate dTS 's favorable performance.

1 Introduction

A contextual bandit is a popular and practical framework for online learning under uncertainty [Li et al., 2010]. In each round, an agent observes a context , takes an action , and receives a reward based on the context and action. The goal is to maximize the expected cumulative reward over n rounds, striking a balance between exploiting actions with high estimated rewards from available data and exploring other actions to improve current estimates. This trade-off is often addressed using either upper confidence bound (UCB) [Auer et al., 2002] or Thompson sampling (TS) [Scott, 2010].

The action space in contextual bandits is often large, resulting in less-than-optimal performance with standard exploration strategies. Luckily, actions usually exhibit correlations, making efficient exploration possible as one action may inform the agent about other actions. In particular, Thompson sampling offers remarkable flexibility, allowing its integration with informative priors [Hong et al., 2022b] that capture these correlations. Inspired by the achievements of diffusion models [SohlDickstein et al., 2015, Ho et al., 2020], which effectively approximate complex distributions [Dhariwal and Nichol, 2021, Rombach et al., 2022], this work captures action correlations by employing diffusion models as priors in contextual Thompson sampling.

We illustrate the idea using video streaming. The objective is to optimize watch time for a user j 27 by selecting a video i from a catalog of K videos. Users j and videos i are associated with context 28 vectors x j and unknown video parameters θ i , respectively. User j 's expected watch time for video i 29 is linear as x ⊤ j θ i . Then, a natural strategy is to independently learn video parameters θ i using LinTS 30 or LinUCB [Agrawal and Goyal, 2013a, Abbasi-Yadkori et al., 2011], but this proves statistically 31 inefficient for larger K . Fortunately, the reward when recommending a movie can provide informative 32 insights into other movies. To capture this, we leverage offline estimates of video parameters denoted 33 by ˆ θ i and build a diffusion model on them. This diffusion model approximates the video parameter 34 distribution, capturing their dependencies. This model enriches contextual Thompson sampling as a 35 prior, effectively capturing complex video dependencies while ensuring computational efficiency. 36

We introduce a framework for contextual bandits with diffusion model priors, upon which we develop 37 diffusion Thompson sampling ( dTS ) that is both computationally and statistically efficient. dTS 38 requires fast updates of the posterior and fast sampling from the posterior , both of which are achieved 39 through our novel efficient posterior approximations. These approximations become exact when 40 both the diffusion model and likelihood are linear. We establish a bound on dTS 's Bayes regret for 41 this specific case, highlighting the advantages of using diffusion models as priors. Our empirical 42 evaluations validate our theory and demonstrate dTS 's strong performance across various settings. 43

44

45

46

47

48

49

50

51

52

Diffusion models were applied in offline decision-making [Ajay et al., 2022, Janner et al., 2022, Wang et al., 2022], but their use in online learning was only recently explored by Hsieh et al. [2023], who focused on multi-armed bandits without theoretical guarantees . Our work extends Hsieh et al. [2023] in two ways. First, we apply the concept to the broader contextual bandit, which is more practical and realistic. Second, we demonstrate that with diffusion models parametrized by linear score functions and linear rewards, we can derive exact closed-form posteriors without approximations. These exact posteriors are valuable as they enable theoretical analysis (unlike Hsieh et al. [2023], who did not provide theoretical guarantees) and motivate efficient approximations for non-linear score functions in contextual bandits, addressing gaps in Hsieh et al. [2023]'s focus on multi-armed bandits.

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

A key contribution, beyond applying diffusion models in contextual bandits, is the efficient computation and sampling of the posterior distribution of a d -dimensional parameter θ | H t , with H t representing the data, when using a diffusion model prior on θ . This is relevant not only to bandits and reinforcement learning but also to a broader range of applications [Chung et al., 2022]. To motivate our approximations, we start with exact closed-form solutions for cases where both the score functions of the diffusion model and the likelihood are linear. These solutions form the basis for our approximations for non-linear score functions, demonstrating both strong empirical performance and computational efficiency. Our approach avoids the computational burden of heavy approximate sampling algorithms required for each latent parameter. For a detailed comparison with existing studies, see Appendix A, where we discuss diffusion models in decision-making, structured bandits, approximate posteriors, and more.

2 Setting

The agent interacts with a contextual bandit over n rounds. In round t ∈ [ n ] , the agent observes a context X t ∈ X , where X ⊆ R d is a context space , it takes an action A t ∈ [ K ] , and then receives a stochastic reward Y t ∈ R that depends on both the context X t and the taken action A t . Each action i ∈ [ K ] is associated with an unknown action parameter θ ∗ ,i ∈ R d , so that the reward received in round t is Y t ∼ P ( · | X t ; θ ∗ ,A t ) , where P ( · | x ; θ ∗ ,i ) is the reward distribution of action i in context x . Throughout the paper, we assume that the reward distribution is parametrized as a generalized linear model (GLM) [McCullagh and Nelder, 1989]. That is, for any x ∈ X , P ( · | x ; θ ∗ ,i ) is an exponential-family distribution with mean g ( x ⊤ θ ∗ ,i ) , where g is the mean function. For example, we recover linear bandits when P ( · | x ; θ ∗ ,i ) = N ( · ; x ⊤ θ ∗ ,i , σ 2 ) where σ > 0 is the observation noise. Similarly, we recover logistic bandits [Filippi et al., 2010] if we let g ( u ) = (1 + exp( -u )) -1 and P ( · | x ; θ ∗ ,i ) = Ber( g ( x ⊤ θ ∗ ,i )) , where Ber( p ) be the Bernoulli distribution with mean p .

We consider the Bayesian bandit setting [Russo and Van Roy, 2014, Hong et al., 2022b], where the action parameters θ ∗ ,i are assumed to be sampled from a known prior distribution. We proceed to define this prior distribution using a diffusion model. The correlations between the action parameters θ ∗ ,i are captured through a diffusion model, where they share a set of L consecutive unknown latent parameters ψ ∗ ,ℓ ∈ R d for ℓ ∈ [ L ] . Precisely, the action parameter θ ∗ ,i depends on the L -th latent parameter ψ ∗ ,L as θ ∗ ,i | ψ ∗ , 1 ∼ N ( f 1 ( ψ ∗ , 1 ) , Σ 1 ) , where the score function f 1 : R d → R d is known . Also, the ℓ -1 -th latent parameter ψ ∗ ,ℓ -1 depends on the ℓ -th latent parameter ψ ∗ ,ℓ as ψ ∗ ,ℓ -1 | ψ ∗ ,ℓ ∼ N ( f ℓ ( ψ ∗ ,ℓ ) , Σ ℓ ) , where the score function f ℓ : R d → R d is known. Finally, the L -th latent parameter ψ ∗ ,L is sampled as ψ ∗ ,L ∼ N (0 , Σ L +1 ) . We summarize this model in (1) and its graph in Fig. 1.

: taken action

in round

The model in (1) represents a Bayesian bandit, where the agent interacts with a bandit instance 86 defined by θ ∗ ,i over n rounds (4-th line in (1)). These action parameters θ ∗ ,i are drawn from the 87 generative process in the first 3 lines of (1). In practice, (1) can be built by pre-training a diffusion 88 model on offline estimates of the action parameters θ ∗ ,i [Hsieh et al., 2023]. 89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

Anatural goal for the agent in this Bayesian framework is to minimize its Bayes regret [Russo and Van Roy, 2014] that measures the expected performance across multiple bandit instances θ ∗ = ( θ ∗ ,i ) i ∈ [ K ] ,

where the expectation in (2) is taken over all random variables in (1). Here r ( x, i ; θ ∗ ) = E Y ∼ P ( ·| x ; θ ∗ ,i ) [ Y ] is the expected reward of action i in context x and A t, ∗ = arg max i ∈ [ K ] r ( X t , i ; θ ∗ ) is the optimal action in round t . The Bayes regret is known to capture the benefits of using informative priors, and hence it is suitable for our problem.

3 Diffusion contextual Thompson sampling

We design Thompson sampling that samples the latent and action parameters hierarchically [Lindley and Smith, 1972]. Precisely, let H t = ( X k , A k , Y k ) k ∈ [ t -1] be the history of all interactions up to round t and let H t,i = ( X k , A k , Y k ) { k ∈ [ t -1]; A k = i } be the history of interactions with action i up to round t . To motivate our algorithm, we decompose the posterior P ( θ ∗ ,i = θ | H t ) recursively as

Q t,L ( ψ L ) = P ( ψ ∗ ,L = ψ L | H t ) is the latent-posterior density of ψ ∗ ,L | H t . Moreover, for any ℓ ∈ [2 : L ] , Q t,ℓ -1 ( ψ ℓ -1 | ψ ℓ ) = P ( ψ ∗ ,ℓ -1 = ψ ℓ -1 | H t , ψ ∗ ,ℓ = ψ ℓ ) is the conditional latentposterior density of ψ ∗ ,ℓ -1 | H t , ψ ∗ ,ℓ = ψ ℓ . Finally, for any action i ∈ [ K ] , P t,i ( θ | ψ 1 ) = P ( θ ∗ ,i = θ | H t,i , ψ ∗ , 1 = ψ 1 ) is the conditional action-posterior density of θ ∗ ,i | H t,i , ψ ∗ , 1 = ψ 1 .

The decomposition in (3) inspires hierarchical sampling. In round t , we initially sample the L -th latent parameter as ψ t,L ∼ Q t,L ( · ) . Then, for ℓ ∈ [ L ] / { 1 } , we sample the ℓ -1 -th latent parameter given that ψ ∗ ,ℓ = ψ t,ℓ , as ψ t,ℓ -1 ∼ Q t,ℓ -1 ( · | ψ t,ℓ ) . Lastly, given that ψ ∗ , 1 = ψ t, 1 , each action parameter is sampled individually as θ t,i ∼ P t,i ( θ | ψ t, 1 ) . This is possible because action parameters θ ∗ ,i are conditionally independent given ψ ∗ , 1 . This leads to Algorithm 1, named d iffusion T hompson S ampling ( dTS ). dTS requires sampling from the K + L posteriors P t,i and Q t,ℓ . Thus we start by providing an efficient recursive scheme to express these posteriors using known quantities. We note that these expressions do not necessarily lead to closed-form posteriors and approximation might be needed. First, the conditional action-posterior P t,i ( · | ψ 1 ) can be written as

where S t,i = { ℓ ∈ [ t -1] , A ℓ = i } are the rounds where the agent takes action i up to round t . Moreover, let L ℓ ( ψ ℓ ) = P ( H t | ψ ∗ ,ℓ = ψ ℓ ) be the likelihood of observations up to round t given that ψ ∗ ,ℓ = ψ ℓ . Then, for any ℓ ∈ [ L ] / { 1 } , the ℓ -1 -th conditional latent-posterior Q t,ℓ -1 ( · | ψ ℓ ) is

and Q t,L ( ψ L ) ∝ L L ( ψ L ) N ( ψ L , 0 , Σ L +1 ) . All the terms above are known, except the likelihoods L ℓ ( ψ ℓ ) for ℓ ∈ [ L ] . These are computed recursively as follows. First, the basis of the recursion is

Then for ℓ ∈ [ L ] / { 1 } , the recursive step is L ℓ ( ψ ℓ ) = ∫ ψ ℓ -1 L ℓ -1 ( ψ ℓ -1 ) N ( ψ ℓ -1 ; f ℓ ( ψ ℓ ) , Σ ℓ ) d ψ ℓ -1 .

All posterior expressions above use known quantities ( f ℓ , Σ ℓ , P ( y | x ; θ )) . However, these expressions typically need to be approximated, except when the score functions f ℓ are linear and the reward distribution P ( · | x ; θ ) is linear-Gaussian, where closed-form solutions can be obtained with careful derivations. These approximations are not trivial, and prior studies often rely on computationally intensive approximate sampling algorithms. In the following sections, we explain how we derive our efficient approximations which are motivated by the closed-form solutions of linear instances.

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

3.1 Linear diffusion model

Assume the score functions f ℓ are linear such as f ℓ ( ψ ∗ ,ℓ ) = W ℓ ψ ∗ ,ℓ for ℓ ∈ [ L ] , where W ℓ ∈ R d × d are known mixing matrices . Then, (1) becomes a linear Gaussian system (LGS) [Bishop, 2006] in this case. This model is important, both in theory and practice. For theory, it leads to closed-form posteriors when the reward distribution is linear-Gaussian as P ( · | x ; θ ∗ ,i ) = N ( · ; x ⊤ θ ∗ ,i , σ 2 ) . This allows bounding the Bayes regret of dTS . For practice, the posterior expressions are used to motivate efficient approximations for the general case in (1) as we show in Section 3.2.

The reward distribution is parameterized as a generalized linear model (GLM) [McCullagh and Nelder, 1989], allowing for non-linear rewards. Thus, we need posterior approximation despite linearity in score functions. Since this non-linearity arises solely from the reward distribution, we approximate it by a Gaussian and propagate this approximation to the latent parameters. This results in efficient posterior approximations that are exact when the reward function is Gaussian (a special case of the GLM model). Specifically, the reward distribution P ( · | x ; θ ) is an exponential family distribution with a mean function denoted by g . Then, we approximate the corresponding likelihood as P ( H t,i | θ ∗ ,i = θ ) ≈ N ( θ ; ˆ B t,i , ˆ G -1 t,i ) , where ˆ B t,i and ˆ G t,i are the maximum likelihood estimate (MLE) and the Hessian of the negative log-likelihood, respectively, and they are defined as

where S t,i = { ℓ ∈ [ t -1] : A ℓ = i } represents the rounds where the agent takes action i up to round t . This simple approximation makes all posteriors Gaussian. Specifically, the conditional action-posterior is Gaussian and is given by P t,i ( · | ψ 1 ) = N ( · ; ˆ µ t,i , ˆ Σ t,i ) , where ˆ µ t,i and ˆ Σ t,i are computed using ˆ B t,i and ˆ G t,i in (7). Moreover, for ℓ ∈ [ L -1] , the ℓ -th conditional latent-posterior is also Gaussian, Q t,ℓ ( · | ψ ℓ +1 ) = N ( · ; ¯ µ t,ℓ , ¯ Σ t,ℓ ) , where ¯ µ t,ℓ and ¯ Σ t,ℓ are computed recursively. The recursion starts with ¯ µ t, 1 and ¯ Σ t, 1 , which are calculated using ˆ B t,i and ˆ G t,i in (7). Full expressions are provided in Appendix B.1. The only approximation made is P ( H t,i | θ ∗ ,i = θ ) ≈ N ( θ ; ˆ B t,i , ˆ G -1 t,i ) , and we propagated it to latent posteriors. Thus, these posterior approximations become exact when the reward distribution follows a linear-Gaussian model, P ( · | x ; θ ∗ ,a ) = N ( · ; x ⊤ θ ∗ ,a , σ 2 ) .

3.2 Non-linear diffusion model

After deriving the posteriors for linear score functions, we return to the general model in (1). Approximation is needed since both the score functions and rewards can be non-linear. To avoid computational challenges, we use a simple and intuitive approximation, where all posteriors P t,i and Q t,ℓ are approximated by Gaussians that are computed recursively. First, the conditional actionposterior is approximated by a Gaussian distribution as P t,i ( · | ψ 1 ) = N ( · ; ˆ µ t,i , ˆ Σ t,i ) , where

In the absence of samples, G t,i = 0 d × d . Thus, the approximate action posterior in (8) matches 157 precisely the term N ( f 1 ( ψ 1 ) , Σ 1 ) in the diffusion prior (1). Moreover, as more data is accumulated, 158 G t,i increases, and the influence of the prior diminishes as ˆ G t,i ˆ B t,i will dominate the prior term 159 Σ -1 1 f 1 ( ψ 1 ) . Similarly, for ℓ ∈ [ L ] / { 1 } , the ℓ -1 -th conditional latent-posterior is approximated by 160

a Gaussian distribution as Q t,ℓ -1 ( · | ψ ℓ ) = N (¯ µ t,ℓ -1 , ¯ Σ t,ℓ -1 ) , where 161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

and the L -th latent-posterior is Q t,L ( · ) = N (¯ µ t,L , ¯ Σ t,L ) ,

Here, ¯ G t,ℓ and ¯ B t,ℓ for ℓ ∈ [ L ] are computed recursively. The basis of the recursion are

Then, the recursive step for ℓ ∈ [ L ] / { 1 } is,

Similarly, in the absence of samples, Q t,ℓ -1 in (9) precisely matches the term N ( f ℓ ( ψ 1 ) , Σ ℓ ) in the diffusion prior (1). As more data is accumulated, the influence of this prior diminishes. Therefore, this approximation retains a key attribute of exact posteriors: they match the prior when there is no data, and the prior's effect diminishes as data accumulates.

4 Analysis

We analyze dTS under the linear diffusion model in Section 3.1 with linear rewards P ( · | x ; θ ∗ ,a ) = N ( · ; x ⊤ θ ∗ ,a , σ 2 ) . This assumption leads to a structure with L layers of linear Gaussian relationships, allowing for theory inspired by linear bandits [Agrawal and Goyal, 2013a, Abbasi-Yadkori et al., 2011]. However, proofs are not the same, and technical challenges remain (explained in Appendix D).

Although our result holds for milder assumptions, we make some simplifications for clarity and interpretability. We assume that (A1) Contexts satisfy ∥ X t ∥ 2 2 = 1 for any t ∈ [ n ] . (A2) Mixing matrices and covariances satisfy λ 1 (W ⊤ ℓ W ℓ ) = 1 for any ℓ ∈ [ L ] and Σ ℓ = σ 2 ℓ I d for any ℓ ∈ [ L +1] . Note that (A1) can be relaxed to any contexts X t with bounded norms ∥ X t ∥ 2 . Also, (A2) can be relaxed to positive definite covariances Σ ℓ and arbitrary mixing matrices W ℓ . In this section, we write ˜ O for the big-O notation up to polylogarithmic factors. We start by stating our bound for dTS .

Theorem 4.1. Let σ 2 MAX = max ℓ ∈ [ L +1] 1 + σ 2 ℓ σ 2 . For any δ ∈ (0 , 1) , the Bayes regret of dTS under Section 3.1 with linear rewards, (A1) and (A2) is bounded as

(13) holds for any δ ∈ (0 , 1) . In particular, the term cnδ is constant when δ = 1 /n . Then, the bound is ˜ O ( √ n ) , and this dependence on the horizon n aligns with prior Bayes regret bounds. The bound comprises L +1 main terms, R ACT ( n ) and R LAT ℓ for ℓ ∈ [ L ] . First, R ACT ( n ) relates to action parameters learning, conforming to a standard form [Lu and Van Roy, 2019]. Similarly, R LAT ℓ is associated with learning the ℓ -th latent parameter. Roughly speaking, our bound captures that our problem can be seen as L +1 sequential linear bandit instances stacked upon each other.

Technical contributions. dTS uses hierarchical sampling. Thus the marginal posterior distribution of θ ∗ ,i | H t is not explicitly defined. The first contribution is deriving θ ∗ ,i | H t using the total covariance decomposition combined with an induction proof, as our posteriors in Section 3.1 were derived recursively. Unlike standard analyses where the posterior distribution of θ ∗ ,i | H t is predetermined due to the absence of latent parameters, our method necessitates this recursive total covariance decomposition. Moreover, in standard proofs, we need to quantify the increase in posterior precision for the action taken A t in each round t ∈ [ n ] . However, in dTS , our analysis extends beyond this. We not only quantify the posterior information gain for the taken action but also for every latent parameter, since they are also learned. To elaborate, we use the recursive formulas in Section 3.1 that connect the posterior covariance of each latent parameter ψ ∗ ,ℓ with the covariance of the posterior action parameters θ ∗ ,i . This allows us to propagate the information gain associated with the action

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

taken in round A t to all latent parameters ψ ∗ ,ℓ , for ℓ ∈ [ L ] by induction. Finally, we carefully bound the resulting terms so that the constants reflect the parameters of the linear diffusion model. More technical details are provided in Appendix D.

To include more structure, we propose the sparsity assumption (A3) W ℓ = ( ¯ W ℓ , 0 d,d -d ℓ ) , where ¯ W ℓ ∈ R d × d ℓ for any ℓ ∈ [ L ] . Note that (A3) is not an assumption when d ℓ = d for any ℓ ∈ [ L ] . Notably, (A3) incorporates a plausible structural characteristic that a diffusion model could capture.

Proposition 4.2 (Sparsity) . Let σ 2 MAX = max ℓ ∈ [ L +1] 1 + σ 2 ℓ σ 2 . For any δ ∈ (0 , 1) , the Bayes regret of dTS under Section 3.1 with linear rewards, (A1) , (A2) and (A3) is bounded as

From Proposition 4.2, our bounds scales as BR ( n ) = O n ( dKσ 2 1 + ∑ ℓ =1 d ℓ σ 2 ℓ +1 σ 2 ℓ MAX ) . The Bayes regret bound has a clear interpretation: if the true environment parameters are drawn from the prior, then the expected regret of an algorithm stays below that bound. Consequently, a less informative prior (such as high variance) leads to a more challenging problem and thus a higher bound. Then, smaller values of K , L , d or d ℓ translate to fewer parameters to learn, leading to lower regret. The regret also decreases when the initial variances σ 2 ℓ decrease. These dependencies are common in Bayesian analysis, and empirical results match them. The reader might question the dependence of our bound on both L and K . We will address this next.

Why the bound increases with K ? This arises due to our conditional learning of θ ∗ ,i given ψ ∗ , 1 . Rather than assuming deterministic linearity, θ ∗ ,i = W 1 ψ ∗ , 1 , we account for stochasticity by modeling θ ∗ ,i ∼ N (W 1 ψ ∗ , 1 , σ 2 1 I d ) . This makes dTS robust to misspecification scenarios where θ ∗ ,i is not perfectly linear with respect to ψ ∗ , 1 , at the cost of additional learning of θ ∗ ,i | ψ ∗ , 1 . If we were to assume deterministic linearity ( σ 1 = 0 ), our regret bound would scale with L only.

Why the bound increases with L ? This is because increasing the number of layers L adds more initial uncertainty due to the additional covariance introduced by the extra layers. However, this does not imply that we should always use L = 1 (the minimum possible L ). While a higher L complicates online learning and increases regret bound, it also enables the capture of a more complex prior distribution through offline pre-training of the diffusion model. Thus, a trade-off exists in practice. A smaller L results in faster computation and easier learning for dTS , but the learned prior might deviate from reality, potentially violating the "true prior assumption" used to derive the regret bound. On the other hand, a larger L allows for better modeling of complex action distributions, producing a prior that more accurately reflects reality and strengthens the validity of the bound.

4.1 Discussion

Computational benefits. Action correlations prompt an intuitive approach: marginalize all latent parameters and maintain a joint posterior of ( θ ∗ ,i ) i ∈ [ K ] | H t . Unfortunately, this is computationally inefficient for large action spaces. To illustrate, suppose that all posteriors are multivariate Gaussians (Section 3.1). Then maintaining the joint posterior ( θ ∗ ,i ) i ∈ [ K ] | H t necessitates converting and storing its dK × dK -dimensional covariance matrix. Then the time and space complexities are O ( K 3 d 3 ) and O ( K 2 d 2 ) . In contrast, the time and space complexities of dTS are O (( L + K ) d 3 ) and O (( L + K ) d 2 ) . This is because dTS requires converting and storing L + K covariance matrices, each being d × d -dimensional. The improvement is huge when K ≫ L , which is common in practice. Certainly, a more straightforward way to enhance computational efficiency is to discard latent parameters and maintain K individual posteriors, each relating to an action parameter θ ∗ ,i ∈ R d ( LinTS ). This improves time and space complexity to O ( Kd 3 ) and O ( Kd 2 ) , respectively. However, LinTS maintains independent posteriors and fails to capture the correlations among actions; it only models θ ∗ ,i | H t,i rather than θ ∗ ,i | H t as done by dTS . Consequently, LinTS incurs higher regret due to the information loss caused by unused interactions of similar actions. Our regret bound and empirical results reflect this aspect.

Statistical benefits. We do not provide a matching lower bound. The only Bayesian lower bound that we know of is Ω(log 2 ( n )) for a much simpler K -armed bandit [Lai, 1987, Theorem 3]. All

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

seminal works on Bayesian bandits do not match it and providing such lower bounds on Bayes regret is still relatively unexplored (even in standard settings) compared to the frequentist one. Therefore, we argue that our bound reflects the overall structure of the problem by comparing dTS to algorithms that only partially use the structure or do not use it at all as follows.

The linear diffusion model in Section 3.1 can be transformed into a Bayesian linear model ( LinTS ) by marginalizing out the latent parameters; in which case the prior on action parameters becomes θ ∗ ,i ∼ N (0 , Σ) , with the θ ∗ ,i being not necessarily independent, and Σ is the marginal initial covariance of action parameters and it writes Σ = σ 2 1 I d + ∑ L ℓ =1 σ 2 ℓ +1 B ℓ B ⊤ ℓ with B ℓ = ∏ ℓ k =1 W k . Then, it is tempting to directly apply LinTS to solve our problem. This approach will induce higher regret because the additional uncertainty of the latent parameters is accounted for in Σ despite integrating them. This causes the marginal action uncertainty Σ to be much higher than the conditional action uncertainty σ 2 1 I d in (3.1), since we have Σ = σ 2 1 I d + ∑ L ℓ =1 σ 2 ℓ +1 B ℓ B ⊤ ℓ ≽ σ 2 1 I d . This discrepancy leads to higher regret, especially when K is large. This is due to LinTS needing to learn K independent d -dimensional parameters, each with a considerably higher initial covariance Σ . This is also reflected by our regret bound. To simply comparisons, suppose that σ ≥ max ℓ ∈ [ L +1] σ ℓ so that σ 2 MAX ≤ 2 . Then the regret bounds of dTS (where we bound σ 2 ℓ MAX by 2 ℓ ) and LinTS read

( √ ( √

Then regret improvements are captured by the variances σ ℓ and the sparsity dimensions d ℓ , and we proceed to illustrate this through the following scenarios.

(I) Decreasing variances. Assume that σ ℓ = 2 ℓ for any ℓ ∈ [ L +1] . Then, the regrets become

Now to see the order of gain, assume the problem is high-dimensional ( d ≫ 1 ), and set L = log 2 ( d ) and d ℓ = ⌊ d 2 ℓ ⌋ . Then the regret of dTS becomes ˜ O (√ nd ( K + L )) ) , and hence the multiplicative factor 2 L in LinTS is removed and replaced with a smaller additive factor L .

(II) Constant variances. Assume that σ ℓ = 1 for any ℓ ∈ [ L +1] . Then, the regrets become

Similarly, let L = log 2 ( d ) , and d ℓ = ⌊ d 2 ℓ ⌋ . Then dTS 's regret is ˜ O (√ nd ( K + L ) ) . Thus the multiplicative factor L in LinTS is removed and replaced with the additive factor L . By comparing this to (I) , the gain with decreasing variances is greater than with constant ones. In general, diffusion models use decreasing variances [Ho et al., 2020] and hence we expect great gains in practice. All observed improvements in this section could become even more pronounced when employing non-linear diffusion models. In our current analysis, we used linear diffusion models, and yet we can already discern substantial differences. Moreover, under non-linear diffusion (1), the latent parameters cannot be analytically marginalized, making LinTS with exact marginalization inapplicable. Finally, Appendix D.7 provide an additional comparison and connection to hierarchies with two levels.

Large action space aspect. dTS 's regret bound scales with Kσ 2 1 instead of K ∑ ℓ σ 2 ℓ , particularly beneficial when σ 1 is small, as often seen in diffusion models. Our regret bound and experiments show that dTS outperforms LinTS more distinctly when the action space becomes larger. Prior studies [Foster et al., 2020, Xu and Zeevi, 2020, Zhu et al., 2022] proposed bandit algorithms that do not scale with K . However, our setting differs significantly from theirs, explaining our inherent dependency on K when σ 1 > 0 . Precisely, they assume a reward function of r ( x, i ; θ ∗ ) = ϕ ( x, i ) ⊤ θ ∗ , with a shared θ ∗ ∈ R d and a known mapping ϕ . In contrast, we consider r ( x, i ; θ ∗ ) = x ⊤ θ ∗ ,i , with θ ∗ = ( θ ∗ ,i ) i ∈ [ K ] ∈ R dK , requiring the learning of K separate d -dimensional action parameters. In their setting, with the availability of ϕ , the regret of dTS would similarly be independent of K . However, obtaining such a mapping ϕ can be challenging as it needs to encapsulate complex context-action dependencies. Notably, our setting reflects a common practical scenario, such as in recommendation systems where each product is often represented by its unique embedding.

5 Experiments

We evaluate dTS using synthetic data, to validate our theory and test dTS in large action spaces. We omit semi-synthetic data [Riquelme et al., 2018] as they often result in small action spaces. This

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

0

10002000300040005000

0

1

2

3

4

5

6

7

8

9

Regret

1e3

Linear diffusion, linear reward

K=100, L=2, d=5

dTS-LL

HierTS

LinTS

LinUCB

0

10002000300040005000

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

1.8

1e3

Linear diffusion, nonlinear reward

K=100, L=2, d=5

dTS-LN

dTS-LL

GLM-TS

UCB-GLM

0

10002000300040005000

0.0

0.2

0.4

0.6

0.8

1.0

1e4

Nonlinear diffusion, linear reward

K=100, L=2, d=5

dTS-NL

LinTS

LinUCB

0

10002000300040005000

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

1e3

Nonlinear diffusion, nonlinear reward

K=100, L=2, d=5

dTS-NN

dTS-NL

GLM-TS

UCB-GLM

0

10002000300040005000

Round

t

2

[

n

]

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Regret

1e5

K=10000, L=4, d=20

dTS-LL

HierTS

LinTS

LinUCB

0

10002000300040005000

Round

t

2

[

n

]

0.0

0.5

1.0

1.5

2.0

2.5

3.0

1e3

K=10000, L=4, d=20

dTS-LN

dTS-LL

GLM-TS

UCB-GLM

0

10002000300040005000

Round

t

2

[

n

]

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

1e7

K=10000, L=4, d=20

dTS-NL

LinTS

LinUCB

0

10002000300040005000

Round

t

2

[

n

]

0

1

2

3

4

5

6

7

8

1e2

K=10000, L=4, d=20

dTS-NN

dTS-NL

GLM-TS

UCB-GLM

choice is further justified by the fact that Hsieh et al. [2023] has already demonstrated the advantages of diffusion models in multi-armed bandits using such data, without theoretical guarantees.

5.1 Settings and baselines

We run 50 random simulations and plot the average regret with its standard error. We consider both linear and non-linear rewards. The distribution of linear rewards is P ( · | x ; θ a ) = N ( x ⊤ θ a , σ 2 ) with σ = 1 . The non-linear rewards are binary and generated from P ( · | x ; θ a ) = Ber( g ( x ⊤ θ a ))) , where g is the sigmoid function. The covariances are Σ ℓ = I d , and the context X t is uniformly drawn from [ -1 , 1] d . We vary d ∈ { 5 , 20 } , L ∈ { 2 , 4 } and K ∈ { 10 2 , 10 4 } . We set the horizon n = 5000 .

Linear diffusion. We consider the linear diffusion model in (3.1) where score functions are linear as f ℓ ( ψ ) = W ℓ ψ where W ℓ are uniformly drawn from [ -1 , 1] d × d . To introduce sparsity, we zero out the last d ℓ columns of W ℓ , resulting in W ℓ = ( ¯ W ℓ , 0 d,d -d ℓ ) , where ( d 1 , d 2 ) = (5 , 2) when d = 5 and L = 2 and ( d 1 , d 2 , d 3 , d 4 ) = (20 , 10 , 5 , 2) when d = 20 and L = 4 .

Non-linear diffusion. We consider the general diffusion model in (1) with score functions f ℓ defined by two-layer neural networks with random weights in [ -1 , 1] , ReLU activation, and a hidden layer dimension of h = 20 when d = 5 and h = 60 when d = 20 .

Baselines. When rewards are linear, we use LinUCB [Abbasi-Yadkori et al., 2011], LinTS [Agrawal and Goyal, 2013a], and HierTS [Hong et al., 2022b] that marginalizes out all latent parameters except ψ ∗ ,L . This corresponds to HierTS-1 in Appendix D.7. When rewards are non-linear, we include UCB-GLM [Li et al., 2017], and GLM-TS [Chapelle and Li, 2012]. GLM-UCB [Filippi et al., 2010] induced high regret while HierTS was designed for linear rewards only and thus both are not included. We name dTS for each setting as dTS-dr , where the suffix d indicates the type of diffusion; L for linear and N for non-linear. The suffix r indicates the type of rewards; L for linear and N for non-linear. For instance, dTS-LL signifies dTS in linear diffusion (Section 3.1) with linear rewards.

5.2 Results and interpretations

Results are shown in Fig. 2 and we make the following observations:

1) dTS has better performance. dTS outperforms the baselines. First, when both the diffusion and rewards are linear, dTS-LL consistently outperforms all baselines that disregard the latent structure ( LinTS and LinUCB ) or incorporate it only partially ( HierTS ). Second, when the diffusion is linear and rewards are non-linear, dTS-LN surpasses all baselines. Third, when the diffusion is non-linear and rewards are linear, dTS-NL demonstrates significant performance gains compared to both LinTS and LinUCB . With non-linear diffusion and rewards, dTS-NN surpasses both GLM-TS and UCB-GLM .

2) Latent diffusion structure may be more important than the reward distribution. When 325 rewards are non-linear (second and fourth columns in Fig. 2), we included variants of dTS that use 326 the correct diffusion prior but the wrong reward distribution, employing linear-Gaussian instead of 327 logistic-Bernoulli ( dTS-LL in the second column and dTS-NL in the fourth column). In both cases, 328 despite the misspecification of the reward distribution, these variants outperform models that use the 329 correct reward distribution but neglect the latent diffusion structure, such as GLM-TS and UCB-GLM . 330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

This underscores the significance of accounting for the latent structure, which can sometimes be more crucial than having an accurate reward distribution. Also, the performance gap between dTS-NL (non-linear diffusion) and GLM-TS and UCB-GLM is even more pronounced compared to the gap between dTS-LL (linear diffusion) and these baselines, possibly due to the increased complexity of the latent structure, in the non-linear diffusion, overshadowing the impact of the reward model itself.

3) Prior misspecification (Fig. 3). We consider a scenario where the prior used by dTS does not match the true prior. To simulate this, we use our setting with linear diffusion and rewards above, but the true parameters W ℓ and Σ ℓ are replaced by misspecified parameters W ℓ + ϵ 1 and Σ ℓ + ϵ 2 . Here, ϵ 1 and ϵ 2 are sampled uniformly from [ v, v +0 . 5] d × d , with v controlling the level of misspecification. The higher the value of v , the greater the misspecification. We vary v ∈ { 0 . 5 , 1 , 1 . 5 } and analyze its impact on dTS 's performance. For comparison, we include the well-specified dTS-LL and the most competitive baseline, HierTS . Results are shown in Fig. 3. As expected, dTS 's performance decreases with increasing misspecification. However, even

0

1000

2000

3000

4000

5000

Round

t

[

n

]

0

500

1000

1500

2000

2500

Regret

Effect of prior misspecification

HierTS

2

with misspecification, dTS outperforms the most competitive baseline, except when v = 1 . 5 , where their performances are comparable. Note that the entries of the true parameters W ℓ and Σ ℓ are smaller than 1, so values of v ∈ { 0 . 5 , 1 , 1 . 5 } can lead to significant parameter misspecification. Yet, the performance of dTS with misspecified prior parameters remains favorable, suggesting that even an imperfect pre-trained diffusion model can be beneficial when used as prior.

4) Regret scaling with K , d and L matches our theory (Fig. 4). We verify the impact of the number of actions K , the context dimension d , and the diffusion depth L on the regret of dTS . We maintain the same experimental setup with linear diffusion and rewards, for which we have derived a Bayes regret upper bound. In Fig. 4, we plot the regret of dTS-LL across varying values of these parameters: K ∈ { 10 , 100 , 500 , 1000 } , d ∈ { 5 , 10 , 15 , 20 } , and L ∈ { 2 , 4 , 5 , 6 } . As anticipated and aligned with our theory, the empirical regret increases as the values of K , d , or L grow. This trend arises because larger values of K , d , or L result in problem instances that are more challenging to learn, consequently leading to higher regret.

5) Performance gap between dTS and LinTS widens as K increases (Fig. 5). To showcase dTS 's improved scalability to larger action spaces, we examine its performance across a range of K values, from 10 to 50 , 000 , in our setting with linear diffusion and rewards. Fig. 5 reports the final cumulative regret for varying values of K for both dTS-LL and LinTS , observing that the gap in the performance becomes larger as K increases.

6 Conclusion

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Increasing values of either

K;d

or

L

0

500

1000

1500

2000

2500

3000

3500

Regret

Effect of

K;d; L

on the regret of LindTS

Increasing

K

Increasing

d

Increasing

L

10

100

500

5000

50000

Number of actions K

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

1.8

Regret in round n

1e4

Regret as a function of K

dTS-LL

LinTS

Grappling with large action spaces in contextual bandits is challenging. Recognizing this, we focused 376 on structured problems where action parameters are sampled from a diffusion model; upon which we 377 built diffusion Thompson sampling ( dTS ). We developed both theoretical and algorithmic foundations 378 for dTS in numerous practical settings. We identified several directions for future work. Exploring 379 other approximations for non-linear diffusion models, both empirically and theoretically. From a 380 theoretical perspective, future research could explore the advantages of non-linear diffusion models 381 by deriving their Bayes regret bounds, akin to our analysis in Section 4. Empirically, investigating 382 our and other approximations in complex tasks would be interesting. Additionally, exploring the 383 extension of this work to offline (or off-policy) learning in contextual bandits [Swaminathan and 384 Joachims, 2015, Aouali et al., 2023a] represents a promising avenue for future research. 385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

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, 2013a.

Shipra Agrawal and Navin Goyal. Further optimal regret bounds for thompson sampling. In Proceedings of the Sixteenth International Conference on Artificial Intelligence and Statistics , pages 99-107, 2013b.

Shipra Agrawal and Navin Goyal. Near-optimal regret bounds for thompson sampling. Journal of the ACM (JACM) , 64(5):1-24, 2017.

Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 , 2022.

Imad Aouali, Victor-Emmanuel Brunel, David Rohde, and Anna Korba. Exponential smoothing for off-policy learning. In International Conference on Machine Learning , pages 984-1017. PMLR, 2023a.

Imad Aouali, Branislav Kveton, and Sumeet Katariya. Mixed-effect thompson sampling. In International Conference on Artificial Intelligence and Statistics , pages 2087-2115. PMLR, 2023b.

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

Mohammad Gheshlaghi Azar, Alessandro Lazaric, and Emma Brunskill. Sequential transfer in multi-armed bandit with finite set of models. In Advances in Neural Information Processing Systems 26 , pages 2220-2228, 2013.

Hamsa Bastani, David Simchi-Levi, and Ruihao Zhu. Meta dynamic pricing: Transfer learning across experiments. CoRR , abs/1902.10918, 2019. URL https://arxiv.org/abs/1902.10918 .

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 M Bishop. Pattern Recognition and Machine Learning , volume 4 of Information science and statistics . Springer, 2006.

Leonardo Cella, Alessandro Lazaric, and Massimiliano Pontil. Meta-learning with stochastic linear bandits. In Proceedings of the 37th International Conference on Machine Learning , 2020.

Leonardo Cella, Karim Lounici, and Massimiliano Pontil. Multi-task representation learning with stochastic linear bandits. arXiv preprint arXiv:2202.10066 , 2022.

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

Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.14687 , 2022.

Aniket Anand Deshmukh, Urun Dogan, and Clayton Scott. Multi-task learning for contextual bandits. In Advances in Neural Information Processing Systems 30 , pages 4848-4856, 2017.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems , 34:8780-8794, 2021.

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.

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

Dylan J Foster, Claudio Gentile, Mehryar Mohri, and Julian Zimmert. Adapting to misspecification in contextual bandits. Advances in Neural Information Processing Systems , 33:11478-11489, 2020.

Claudio Gentile, Shuai Li, and Giovanni Zappella. Online clustering of bandits. In Proceedings of the 31st International Conference on Machine Learning , pages 757-765, 2014.

Aditya Gopalan, Shie Mannor, and Yishay Mansour. Thompson sampling for complex online problems. In Proceedings of the 31st International Conference on Machine Learning , pages 100-108, 2014.

Samarth Gupta, Shreyas Chaudhari, Subhojyoti Mukherjee, Gauri Joshi, and Osman Yagan. A unified approach to translate classical bandit algorithms to the structured bandit setting. CoRR , abs/1810.08164, 2018. URL https://arxiv.org/abs/1810.08164 .

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Joey Hong, Branislav Kveton, Manzil Zaheer, Yinlam Chow, Amr Ahmed, and Craig Boutilier. Latent bandits revisited. In Advances in Neural Information Processing Systems 33 , 2020.

Joey Hong, Branislav Kveton, Sumeet Katariya, Manzil Zaheer, and Mohammad Ghavamzadeh. Deep hierarchy in bandits. In International Conference on Machine Learning , pages 8833-8851. PMLR, 2022a.

Joey Hong, Branislav Kveton, Manzil Zaheer, and Mohammad Ghavamzadeh. Hierarchical Bayesian bandits. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics , 2022b.

Yu-Guan Hsieh, Shiva Prasad Kasiviswanathan, Branislav Kveton, and Patrick Blöbaum. Thompson sampling with diffusion generative prior. arXiv preprint arXiv:2301.05182 , 2023.

Jiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li, and Liwei Wang. Near-optimal representation learning for linear bandits and linear rl. In International Conference on Machine Learning , pages 4349-4358. PMLR, 2021.

Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

Emilie Kaufmann, Nathaniel Korda, and Rémi Munos. Thompson sampling: An asymptotically optimal finite-time analysis. In International conference on algorithmic learning theory , pages 199-213. Springer, 2012.

Daphne Koller and Nir Friedman. Probabilistic Graphical Models: Principles and Techniques . MIT Press, Cambridge, MA, 2009.

Nathaniel Korda, Emilie Kaufmann, and Remi Munos. Thompson sampling for 1-dimensional exponential family bandits. Advances in neural information processing systems , 26, 2013.

John K Kruschke. Bayesian data analysis. Wiley Interdisciplinary Reviews: Cognitive Science , 1(5): 658-676, 2010.

Branislav Kveton, Manzil Zaheer, Csaba Szepesvari, Lihong Li, Mohammad Ghavamzadeh, and Craig Boutilier. Randomized exploration in generalized linear bandits. In International Conference on Artificial Intelligence and Statistics , pages 2066-2076. PMLR, 2020.

Branislav Kveton, Mikhail Konobeev, Manzil Zaheer, Chih-Wei Hsu, Martin Mladenov, Craig Boutilier, and Csaba Szepesvari. Meta-Thompson sampling. In Proceedings of the 38th International Conference on Machine Learning , 2021.

Tze Leung Lai. Adaptive treatment allocation and the multi-armed bandit problem. The Annals of Statistics , 15(3):1091-1114, 1987.

Tor Lattimore and Remi Munos. Bounded regret for finite-armed structured bandits. In Advances in Neural Information Processing Systems 27 , pages 550-558, 2014.

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

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 20712080, 2017.

Dennis Lindley and Adrian Smith. Bayes estimates for the linear model. Journal of the Royal Statistical Society: Series B (Methodological) , 34(1):1-18, 1972.

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

Odalric-Ambrym Maillard and Shie Mannor. Latent bandits. In Proceedings of the 31st International Conference on Machine Learning , pages 136-144, 2014.

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

Amit Peleg, Naama Pearl, and Ron Meirr. Metalearning linear bandits by prior update. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics , 2022.

Carlos Riquelme, George Tucker, and Jasper Snoek. Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling. arXiv preprint arXiv:1802.09127 , 2018.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10684-10695, 2022.

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

Steven Scott. A modern bayesian look at the multi-armed bandit. Applied Stochastic Models in Business and Industry , 26:639 - 658, 2010.

Max Simchowitz, Christopher Tosh, Akshay Krishnamurthy, Daniel Hsu, Thodoris Lykouris, Miro Dudik, and Robert Schapire. Bayesian decision-making under misspecified priors with applications to meta-learning. In Advances in Neural Information Processing Systems 34 , 2021.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning , pages 2256-2265. PMLR, 2015.

Adith Swaminathan and Thorsten Joachims. Counterfactual risk minimization: Learning from logged bandit feedback. In International Conference on Machine Learning , pages 814-823. PMLR, 2015.

Runzhe Wan, Lin Ge, and Rui Song. Metadata-based multi-task bandits with Bayesian hierarchical models. In Advances in Neural Information Processing Systems 34 , 2021.

Runzhe Wan, Lin Ge, and Rui Song. Towards scalable and robust structured bandits: A meta-learning framework. CoRR , abs/2202.13227, 2022. URL https://arxiv.org/abs/2202.13227 .

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06193 , 2022.

Neil Weiss. A Course in Probability . Addison-Wesley, 2005.

Yunbei Xu and Assaf Zeevi. Upper counterfactual confidence bounds: a new optimism principle for contextual bandits. arXiv preprint arXiv:2007.07876 , 2020.

Jiaqi Yang, Wei Hu, Jason D Lee, and Simon S Du. Impact of representation learning in linear bandits. arXiv preprint arXiv:2010.06531 , 2020.

522

523

524

525

526

527

Tong Yu, Branislav Kveton, Zheng Wen, Ruiyi Zhang, and Ole Mengshoel. Graphical models meet bandits: A variational Thompson sampling approach. In Proceedings of the 37th International Conference on Machine Learning , 2020.

Yinglun Zhu, Dylan J Foster, John Langford, and Paul Mineiro. Contextual bandits with large action spaces: Made practical. In International Conference on Machine Learning , pages 27428-27453. PMLR, 2022.

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

570

571

572

573

574

575

576

577

578

579

580

Supplementary materials

Notation. For any positive integer n , we define [ n ] = { 1 , 2 , ..., n } . Let v 1 , . . . , v n ∈ R d be n vectors, ( v i ) i ∈ [ n ] ∈ R nd is the nd -dimensional vector obtained by concatenating v 1 , . . . , v n . For any matrix A ∈ R d × d , λ 1 (A) and λ d (A) denote the maximum and minimum eigenvalues of A , respectively. Finally, we write ˜ O for the big-O notation up to polylogarithmic factors.

A Extended related work

Thompson sampling (TS) operates within the Bayesian framework and it involves specifying a prior/likelihood model. In each round, the agent samples unknown model parameters from the current posterior distribution. The chosen action is the one that maximizes the resulting reward. TS is naturally randomized, particularly simple to implement, and has highly competitive empirical performance in both simulated and real-world problems [Russo and Van Roy, 2014, Chapelle and Li, 2012]. Regret guarantees for the TS heuristic remained open for decades even for simple models. Recently, however, significant progress has been made. For standard multi-armed bandits, TS is optimal in the Beta-Bernoulli model [Kaufmann et al., 2012, Agrawal and Goyal, 2013b], GaussianGaussian model [Agrawal and Goyal, 2013b], and in the exponential family using Jeffrey's prior [Korda et al., 2013]. For linear bandits, TS is nearly-optimal [Russo and Van Roy, 2014, Agrawal and Goyal, 2017, Abeille and Lazaric, 2017]. In this work, we build TS upon complex diffusion priors and analyze the resulting Bayes regret [Russo and Van Roy, 2014] in the linear contextual bandit setting.

Decision-making with diffusion models gained attention recently, especially in offline learning [Ajay et al., 2022, Janner et al., 2022, Wang et al., 2022]. However, their application in online learning was only examined by Hsieh et al. [2023], which focused on meta-learning in multi-armed bandits without theoretical guarantees. In this work, we expand the scope of Hsieh et al. [2023] to encompass the broader contextual bandit framework. In particular, we provide theoretical analysis for linear instances, effectively capturing the advantages of using diffusion models as priors in contextual Thompson sampling. These linear cases are particularly captivating due to closed-form posteriors, enabling both theoretical analysis and computational efficiency; an important practical consideration.

Hierarchical Bayesian bandits [Bastani et al., 2019, Kveton et al., 2021, Basu et al., 2021, Simchowitz et al., 2021, Wan et al., 2021, Hong et al., 2022b, Peleg et al., 2022, Wan et al., 2022, Aouali et al., 2023b] applied TS to simple graphical models, wherein action parameters are generally sampled from a Gaussian distribution centered at a single latent parameter. These works mostly span metaand multi-task learning for multi-armed bandits, except in cases such as Aouali et al. [2023b], Hong et al. [2022a] that consider the contextual bandit setting. Precisely, Aouali et al. [2023b] assume that action parameters are sampled from a Gaussian distribution centered at a linear mixture of multiple latent parameters. On the other hand, Hong et al. [2022a] applied TS to a graphical model represented by a tree. Our work can be seen as an extension of all these works to much more complex graphical models, for which both theoretical and algorithmic foundations are developed. Note that the settings in most of these works can be recovered with specific choices of the diffusion depth L and functions f ℓ . This attests to the modeling power of dTS .

Approximate Thompson sampling is a major problem in the Bayesian inference literature. This is because most posterior distributions are intractable, and thus practitioners must resort to sophisticated computational techniques such as Markov chain Monte Carlo [Kruschke, 2010]. Prior works [Riquelme et al., 2018, Chapelle and Li, 2012, Kveton et al., 2020] highlight the favorable empirical performance of approximate Thompson sampling. Particularly, [Kveton et al., 2020] provide theoretical guarantees for Thompson sampling when using the Laplace approximation in generalized linear bandits (GLB). In our context, we incorporate approximate sampling when the reward exhibits non-linearity. While our approximation does not come with formal guarantees, it enjoys strong practical performance. An in-depth analysis of this approximation is left as a direction for future works. Similarly, approximating the posterior distribution when the diffusion model is non-linear as well as analyzing it is an interesting direction of future works.

Bandits with underlying structure also align with our work, where we assume a structured relationship among actions, captured by a diffusion model. In latent bandits [Maillard and Mannor, 2014, Hong et al., 2020], a single latent variable indexes multiple candidate models. Within structured

581

582

583

584

585

586

587

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

626

627

finite-armed bandits [Lattimore and Munos, 2014, Gupta et al., 2018], each action is linked to a known mean function parameterized by a common latent parameter. This latent parameter is learned. TS was also applied to complex structures [Yu et al., 2020, Gopalan et al., 2014]. However, simultaneous computational and statistical efficiencies aren't guaranteed. Meta- and multi-task learning with upper confidence bound (UCB) approaches have a long history in bandits [Azar et al., 2013, Gentile et al., 2014, Deshmukh et al., 2017, Cella et al., 2020]. These, however, often adopt a frequentist perspective, analyze a stronger form of regret, and sometimes result in conservative algorithms. In contrast, our approach is Bayesian, with analysis centered on Bayes regret. Remarkably, our algorithm, dTS , performs well as analyzed without necessitating additional tuning. Finally, Low-rank bandits [Hu et al., 2021, Cella et al., 2022, Yang et al., 2020] also relate to our linear diffusion model when L = 1 . Broadly, there exist two key distinctions between these prior works and the special case of our model (linear diffusion model with L = 1 ). First, they assume θ ∗ ,i = W 1 ψ ∗ , 1 , whereas we incorporate additional uncertainty in the covariance Σ 1 to account for possible misspecification as θ ∗ ,i = N (W 1 ψ ∗ , 1 , Σ 1 ) . Consequently, these algorithms might suffer linear regret due to model misalignment. Second, we assume that the mixing matrix W 1 is available and pre-learned offline, whereas they learn it online. While this is more general, it leads to computationally expensive methods that are difficult to employ in a real-world online setting.

Large action spaces. Roughly speaking, the regret bound of dTS scales with Kσ 2 1 rather than K ∑ ℓ σ 2 ℓ . This is particularly beneficial when σ 1 is small, a common scenario in diffusion models with decreasing variances. A notable case is when σ 1 = 0 , where the regret becomes independent of K . Also, our analysis (Section 4.1) indicates that the gap in performance between dTS and LinTS becomes more pronounced when the number of action increases, highlighting dTS 's suitability for large action spaces. Note that some prior works [Foster et al., 2020, Xu and Zeevi, 2020, Zhu et al., 2022] proposed bandit algorithms that do not scale with K . However, our setting differs significantly from theirs, explaining our inherent dependency on K when σ 1 > 0 . Precisely, they assume a reward function of r ( x, i ) = ϕ ( x, i ) ⊤ θ ∗ , with a shared θ ∗ ∈ R d across actions and a known mapping ϕ . In contrast, we consider r ( x, i ) = x ⊤ θ ∗ ,i , requiring the learning of K separate d -dimensional action parameters. In their setting, with the availability of ϕ , the regret of dTS would similarly be independent of K . However, obtaining such a mapping ϕ can be challenging as it needs to encapsulate complex context-action dependencies. Notably, our setting reflects a common practical scenario, such as in recommendation systems where each product is often represented by its embedding. In summary, the dependency on K is more related to our setting than the method itself, and dTS would scale with d only in their setting. Note that dTS is both computationally and statistically efficient (Section 4.1). This becomes particularly notable in large action spaces. Our empirical results in Fig. 2, notably with K = 10 4 , demonstrate that dTS significantly outperforms the baselines. More importantly, the performance gap between dTS and these baselines is larger when the number of actions ( K ) increases, highlighting the improved scalability of dTS to large action spaces.

B Posterior derivations for linear diffusion models

Here, we assume the score functions f ℓ are linear such as f ℓ ( ψ ∗ ,ℓ ) = W ℓ ψ ∗ ,ℓ for ℓ ∈ [ L ] , where W ℓ ∈ R d × d are known mixing matrices . Then, (1) becomes a linear Gaussian system (LGS) [Bishop, 2006] and can be summarized as follows

In this section, we derive the K + L posteriors P t,i and Q t,ℓ , for which we provide the full expressions in Appendix B.1. In our proofs, p ( x ) ∝ f ( x ) means that the probability density p satisfies p ( x ) = f ( x ) Z for any x ∈ R d , where Z is a normalization constant. In particular, we extensively use that if p ( x ) ∝ exp[ -1 2 x ⊤ Λ x + x ⊤ m ] , where Λ is positive definite. Then p is the multivariate Gaussian density with covariance Σ = Λ -1 and mean µ = Σ m . These are standard notations and techniques to manipulate Gaussian distributions [Koller and Friedman, 2009, Chapter 7].

628

629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647

648

649

650

651

652

653

654

655

656

657

B.1 Posterior expressions for linear diffusion models

Recall that we posit that the reward distribution is parameterized as a generalized linear model (GLM) [McCullagh and Nelder, 1989], allowing for non-linear rewards. As a result, despite linearity in score functions, the non-linearity in rewards makes it challenging to obtain closed-form posteriors. However, since this non-linearity arises solely from the reward distribution, we approximate it using a Gaussian distribution. This leads to efficient posterior approximations that are exact in cases where the reward function is indeed Gaussian (a special case of the GLM model). Precisely, the reward distribution P ( · | x ; θ ) is an exponential-family distribution. Therefore, the log-likelihoods write log P ( H t,i | θ ∗ ,i = θ ) = ∑ k ∈ S t,i Y k X ⊤ k θ -A ( X ⊤ k θ ) + C ( Y k ) , where C is a real function, and A is a twice continuously differentiable function whose derivative is the mean function, ˙ A = g . Now we let ˆ B t,i and ˆ G t,i be the maximum likelihood estimate (MLE) and the Hessian of the negative log-likelihood, respectively, defined as

where S t,i = { ℓ ∈ [ t -1] : A ℓ = i } are the rounds where the agent takes action i up to round t . Then we approximation the respective likelihood as P ( H t,i | θ ∗ ,i = θ ) ≈ N ( θ ; ˆ B t,i , ˆ G -1 t,i ) . This approximation makes all posteriors Gaussian. First, the conditional action-posterior reads P t,i ( · | ψ 1 ) = N ( · ; ˆ µ t,i , ˆ Σ t,i ) ,

For ℓ ∈ [ L ] / { 1 } , the ℓ -1 -th conditional latent-posterior is Q t,ℓ -1 ( · | ψ ℓ ) = N (¯ µ t,ℓ -1 , ¯ Σ t,ℓ -1 ) ,

and the L -th latent-posterior is Q t,L ( · ) = N (¯ µ t,L , ¯ Σ t,L ) ,

Finally, ¯ G t,ℓ and ¯ B t,ℓ for ℓ ∈ [ L ] are computed recursively. The basis of the recursion are

Then, the recursive step for ℓ ∈ [ L ] / { 1 } is,

This concludes the derivation of our posterior approximation. Note that these approximations are exact when the reward distribution follows a linear-Gaussian model, P ( · | x ; θ ∗ ,a ) = N ( · ; x ⊤ θ ∗ ,a , σ 2 ) .

B.2 Derivation of Action-Posteriors for Linear Diffusion Models

To simplify derivations, we consider the case where the reward distribution is indeed linearGaussian as P ( · | X t ; θ ∗ ,A t ) = N ( X ⊤ t θ ∗ ,A t , σ 2 ) , but the same derivations can be applied when the rewards are non-linear. In this case, the likelihood approximation in (16) becomes exact as we have that P ( H t,i | θ ∗ ,i = θ ) ∝ N ( θ ; ˆ B t,i , ˆ G -1 t,i ) , where ˆ B t,i is the corresponding MLE and ˆ G t,i = σ -2 ∑ k ∈ S t,i X k X ⊤ k in this case. Our derivations rely on the fact that the MLE ˆ B t,i in this linear-Gaussian case satisfies: ˆ G t,i ˆ B t,i = v ∑ k ∈ S t,i X k Y ⊤ k .

Proposition B.1. Consider the following model, which corresponds to the last two layers in Eq. (15)

Then we have that for any t ∈ [ n ] and i ∈ [ K ] , P t,i ( θ | ψ 1 ) = P ( θ ∗ ,i = θ | ψ ∗ , 1 = ψ 1 , H t,i ) = 658 N ( θ ; ˆ µ t,i , ˆ Σ t,i ) , where 659

661

662

663

with ˆ Λ t,i = v ∑ k ∈ S t,i X k X ⊤ k +Λ 1 , ˆ Λ t,i ˆ µ t,i = v ∑ k ∈ S t,i X k Y k +Λ 1 W 1 ψ 1 . Using that, in this linear-Gaussian case, ˆ G t,i = v ∑ k ∈ S t,i X k X ⊤ k and ˆ G t,i ˆ B t,i = v ∑ k ∈ S t,i X k Y k concludes the proof.

664

665

666

667

668

669

670

671

672

673

674

The same proof applies when the reward distribution is not linear-Gaussian, with the approximation P ( H t,i | θ ∗ ,i = θ ) ≈ N ( θ ; ˆ B t,i , ˆ G -1 t,i ) . Using this approximation in the derivations above leads to the same results.

B.3 Derivation of recursive latent-posteriors for linear diffusion models

Again, to simplify derivations, we consider the case where the reward distribution is indeed linearGaussian as P ( · | X t ; θ ∗ ,A t ) = N ( X ⊤ t θ ∗ ,A t , σ 2 ) , but the same derivations can be applied when the rewards are non-linear.

Proposition B.2. For any ℓ ∈ [ L ] / { 1 } , the ℓ -1 -th conditional latent-posterior reads Q t,ℓ -1 ( · | ψ ℓ ) = N (¯ µ t,ℓ -1 , ¯ Σ t,ℓ -1 ) , with

and the L -th latent-posterior reads Q t,L ( · ) = N (¯ µ t,L , ¯ Σ t,L ) , with

Proof. Let ℓ ∈ [ L ] / { 1 } . Then, Bayes rule yields that

675

But from Lemma B.3, we know that

Therefore, 676

with ¯ Σ -1 t,ℓ -1 = Σ -1 ℓ + ¯ G t,ℓ -1 and ¯ µ t,ℓ -1 = ¯ Σ t,ℓ -1 ( Σ -1 ℓ W ℓ ψ ℓ + ¯ B t,ℓ -1 ) . In ( i ) , we omit terms that 677 are constant in ψ ℓ -1 . In ( ii ) , we complete the square. This concludes the proof for ℓ ∈ [ L ] / { 1 } . For 678 Q t,L , we use Bayes rule to get 679

Then from Lemma B.3, we know that 680

We then use the same derivations above to compute the product exp [ -1 2 ψ ⊤ L ¯ G t,L ψ L + ψ ⊤ L ¯ B t,L ] × N ( ψ L , 0 , Σ L +1 ) , which concludes the proof.

Lemma B.3. The following holds for any t ∈ [ n ] and ℓ ∈ [ L ] ,

where ¯ G t,ℓ and ¯ B t,ℓ are defined by recursion in Section 3.1.

Proof. We prove this result by induction. To reduce clutter, we let v = σ -2 , and Λ 1 = Σ -1 1 . We start with the base case of the induction when ℓ = 1 .

(I) Base case. Here we want to show that P ( H t | ψ ∗ , 1 = ψ 1 ) ∝ exp [ -1 2 ψ ⊤ 1 ¯ G t, 1 ψ 1 + ψ ⊤ 1 ¯ B t, 1 ) ] , where ¯ G t, 1 and ¯ B t, 1 are given in Eq. (20). First, we have that

681

682

683

684

685

686

687

688

where ( i ) follows from the fact that θ ∗ ,i for i ∈ [ K ] are conditionally independent given 689 ψ ∗ , 1 = ψ 1 and that given θ ∗ ,i , H t,i is independent of ψ ∗ , 1 . Now we compute h i ( ψ 1 ) = 690 ∫ θ ( ∏ k ∈ S t,i N ( Y k ; X ⊤ k θ, σ 2 ) ) N ( θ ; W 1 ψ 1 , Σ 1 ) d θ as 691

But we know that ˆ G t,i = v ∑ k ∈ S t,i X k X ⊤ k , and ˆ G t,i ˆ B t,i = v ∑ k ∈ S t,i Y k X k (because we assumed 692 linear-Gaussian likelihood). To further simplify expressions, we also let 693

We have that UV = V U = I d , and thus 694

where 695

But notice that V = ( ˆ G t,i +Λ 1 ) -1 = ˆ Σ t,i and thus 696

Finally, we plug this result in Eq. (24) to get 697

where 698

This concludes the proof of the base case. 699

(II) Induction step. Let ℓ ∈ [ L ] / { 1 } . Suppose that 700

Then we want to show that 701

where 702

To achieve this, we start by expressing P ( H t | ψ ∗ ,ℓ = ψ ℓ ) in terms of P ( H t | ψ ∗ ,ℓ -1 = ψ ℓ -1 ) as 703

Now let S = ¯ G t,ℓ -1 +Λ ℓ and V = ¯ B t,ℓ -1 +Λ ℓ W ℓ ψ ℓ . Then we have that, 704

In the second step, we omit constants in ψ ℓ and ψ ℓ -1 . Thus 705

It follows that 706

In the last step, we omit constants in ψ ℓ and we set 707

708

709

710

711

712

713

714

715

716

717

718

719

720

721

722

723

724

725

726

727

728

729

730

731

This completes the proof.

Similarly, this same proof applies when the reward distribution is not linear-Gaussian, with the approximation P ( H t,i | θ ∗ ,i = θ ) ≈ N ( θ ; ˆ B t,i , ˆ G -1 t,i ) . Using this approximation in the derivations above leads to the same results.

C Posterior derivations for non-linear diffusion models

After deriving the posteriors for linear score functions f ℓ , we now get back to the general case in (1), where the score functions are potentially non-linear. Approximation is needed since both the score functions and rewards can be non-linear. To avoid any computational challenges, we use a simple and intuitive approximation, where all posteriors P t,i and Q t,ℓ are approximated by the Gaussian distributions in Appendix B.1, with few changes. First, the terms W ℓ ψ ℓ in (18) are replaced by f ℓ ( ψ ℓ ) . This accounts for the fact that the prior mean is now f ℓ ( ψ ℓ ) rather than W ℓ ψ ℓ , and this is the main difference between the linear diffusion model in (15) and the general, potentially non-linear, diffusion model in (1). Second, the matrix multiplications that involve the matrices W ℓ in (20) and (21) are simply removed. Despite being simple, this approximation is efficient and avoids the computational burden of heavy approximate sampling algorithms required for each latent parameter. This is why deriving the exact posterior for linear score functions was key beyond enabling theoretical analyses. Moreover, this approximation retains some key attributes of exact posteriors. Specifically, in the absence of data, it recovers precisely the prior in (1), and as more data is accumulated, the influence of the prior diminishes.

D Regret proof and additional discussions

D.1 Sketch of the proof

We start with the following standard lemma upon which we build our analysis [Aouali et al., 2023b]. Lemma D.1. Assume that P ( θ ∗ ,i = θ | H t ) = N ( θ ; ˇ µ t,i , ˇ Σ t,i ) for any i ∈ [ K ] , then for any δ ∈ (0 , 1) ,

Applying Lemma D.1 requires proving that the marginal action-posteriors P ( θ ∗ ,i = θ | H t ) in Eq. (3) 732 are Gaussian and computing their covariances, while we only know the conditional action-posteriors 733 P t,i and latent-posteriors Q t,ℓ . This is achieved by leveraging the preservation properties of the 734 family of Gaussian distributions [Koller and Friedman, 2009] and the total covariance decomposition 735 [Weiss, 2005] which leads to the next lemma. 736

Lemma D.2. Let t ∈ [ n ] and i ∈ [ K ] , then the marginal covariance matrix ˇ Σ t,i reads 737

738

739

740

741

742

743

744

745

746

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

763

764

765

766

767

768

769

770

771

772

773

774

775

776

777

778

779

780

781

782

783

784

785

The marginal covariance matrix ˇ Σ t,i in Eq. (29) decomposes into L + 1 terms. The first term corresponds to the posterior uncertainty of θ ∗ ,i | ψ ∗ , 1 . The remaining L terms capture the posterior uncertainties of ψ ∗ ,L and ψ ∗ ,ℓ -1 | ψ ∗ ,ℓ for ℓ ∈ [ L ] / { 1 } . These are then used to quantify the posterior information gain of latent parameters after one round as follows.

Lemma D.3 (Posterior information gain) . Let t ∈ [ n ] and ℓ ∈ [ L ] , then

Finally, Lemma D.2 is used to decompose ∥ X t ∥ 2 ˇ Σ t,A t in Eq. (28) into L +1 terms. Each term is bounded thanks to Lemma D.3. This results in the Bayes regret bound in Theorem 4.1.

D.2 Technical contributions

Our main technical contributions are the following.

Lemma D.2. In dTS , sampling is done hierarchically, meaning the marginal posterior distribution of θ ∗ ,i | H t is not explicitly defined. Instead, we use the conditional posterior distribution of θ ∗ ,i | H t , ψ ∗ , 1 . The first contribution was deriving θ ∗ ,i | H t using the total covariance decomposition combined with an induction proof, as our posteriors in Section 3.1 were derived recursively. Unlike in Bayes regret analysis for standard Thompson sampling, where the posterior distribution of θ ∗ ,i | H t is predetermined due to the absence of latent parameters, our method necessitates this recursive total covariance decomposition, marking a first difference from the standard Bayesian proofs of Thompson sampling. Note that HierTS , which is developed for multi-task linear bandits, also employs total covariance decomposition, but it does so under the assumption of a single latent parameter; on which action parameters are centered. Our extension significantly differs as it is tailored for contextual bandits with multiple, successive levels of latent parameters, moving away from HierTS 's assumption of a 1-level structure. Roughly speaking, HierTS when applied to contextual would consider a singlelevel hierarchy, where θ ∗ ,i | ψ ∗ , 1 ∼ N ( ψ ∗ , 1 , Σ 1 ) with L = 1 . In contrast, our model proposes a multi-level hierarchy, where the first level is θ ∗ ,i | ψ ∗ , 1 ∼ N ( W 1 ψ ∗ , 1 , Σ 1 ) . This also introduces a new aspect to our approach - the use of a linear function W 1 ψ ∗ , 1 , as opposed to HierTS 's assumption where action parameters are centered directly on the latent parameter. Thus, while HierTS also uses the total covariance decomposition, our generalize it to multi-level hierarchies under L linear functions W ℓ ψ ∗ ,ℓ , instead of a single-level hierarchy under a single identity function ψ ∗ , 1 .

Lemma D.3. In Bayes regret proofs for standard Thompson sampling, we often quantify the posterior information gain. This is achieved by monitoring the increase in posterior precision for the action taken A t in each round t ∈ [ n ] . However, in dTS , our analysis extends beyond this. We not only quantify the posterior information gain for the taken action but also for every latent parameter, since they are also learned. This lemma addresses this aspect. To elaborate, we use the recursive formulas in Section 3.1 that connect the posterior covariance of each latent parameter ψ ∗ ,ℓ with the covariance of the posterior action parameters θ ∗ ,i . This allows us to propagate the information gain associated with the action taken in round A t to all latent parameters ψ ∗ ,ℓ , for ℓ ∈ [ L ] by induction. This is a novel contribution, as it is not a feature of Bayes regret analyses in standard Thompson sampling.

Proposition 4.2. Building upon the insights of Theorem 4.1, we introduce the sparsity assumption (A3) . Under this assumption, we demonstrate that the Bayes regret outlined in Theorem 4.1 can be significantly refined. Specifically, the regret becomes contingent on dimensions d ℓ ≤ d , as opposed to relying on the entire dimension d . This sparsity assumption is both a novel and a key technical contribution to our work. Its underlying principle is straightforward: the Bayes regret is influenced by the quantity of parameters that require learning. With the sparsity assumption, this number is reduced to less than d for each latent parameter. To substantiate this claim, we revisit the proof of Theorem 4.1 and modify a crucial equality. This adjustment results in a more precise representation by partitioning the covariance matrix of each latent parameter ψ ∗ ,ℓ into blocks. These blocks comprise a d ℓ × d ℓ segment corresponding to the learnable d ℓ parameters of ψ ∗ ,ℓ , and another block of size ( d -d ℓ ) × ( d -d ℓ ) that does not necessitate learning. This decomposition allows us to conclude that the final regret is solely dependent on d ℓ , marking a significant refinement from the original theorem.

D.3 Proof of lemma D.2 786

In this proof, we heavily rely on the total covariance decomposition [Weiss, 2005]. Also, refer to 787 [Hong et al., 2022b, Section 5.2] for a brief introduction to this decomposition. Now, from Eq. (17), 788 we have that 789

First, given H t , cov [ θ ∗ ,i | H t , ψ ∗ , 1 ] = ( ˆ G t,i +Σ -1 1 ) -1 is constant. Thus 790

In addition, given H t , ˆ Σ t,i , ˆ G t,i and ˆ B t,i are constant. Thus 791

where ¯ ¯ Σ t, 1 = cov [ ψ ∗ , 1 | H t ] is the marginal posterior covariance of ψ ∗ , 1 . Finally, the total covariance 792 decomposition [Weiss, 2005, Hong et al., 2022b] yields that 793

However, ¯ ¯ Σ t, 1 = cov [ ψ ∗ , 1 | H t ] is different from ¯ Σ t, 1 = cov [ ψ ∗ , 1 | H t , ψ ∗ , 2 ] that we already derived 794 in Eq. (18). Thus we do not know the expression of ¯ ¯ Σ t, 1 . But we can use the same total covariance 795 decomposition trick to find it. Precisely, let ¯ ¯ Σ t,ℓ = cov [ ψ ∗ ,ℓ | H t ] for any ℓ ∈ [ L ] . Then we have that 796

First, given H t , cov [ ψ ∗ , 1 | H t , ψ ∗ , 2 ] = ( Σ -1 2 + ¯ G t, 1 ) -1 is constant. Thus 797

In addition, given H t , ¯ Σ t, 1 , ˜ Σ t, 1 and ¯ B t, 1 are constant. Thus 798

Finally, total covariance decomposition [Weiss, 2005, Hong et al., 2022b] leads to 799

Now using the techniques, this can be generalized using the same technique as above to 800

Then, by induction, we get that 801

where we use that by definition ¯ ¯ Σ t,L = cov[ ψ ∗ ,L | H t ] = ¯ Σ t,L and set ¯ P 1 = I d and ¯ P ℓ = 802 ∏ ℓ -1 k =1 ¯ Σ t,k Σ -1 k +1 W k +1 for any ℓ ∈ [ L ] / { 1 } . Plugging this in Eq. (31) leads to 803

where P i,ℓ = ˆ Σ t,i Σ -1 1 W 1 ¯ P ℓ = ˆ Σ t,i Σ -1 1 W 1 ∏ ℓ -1 k =1 ¯ Σ t,k Σ -1 k +1 W k +1 . 804

D.4 Proof of lemma D.3 805

We prove this result by induction. We start with the base case when ℓ = 1 . 806

(I) Base case. Let u = σ -1 ˆ Σ 1 2 t,A t X t From the expression of ¯ Σ t, 1 in Eq. (18), we have that 807

In ( i ) we use the Sherman-Morrison formula. Note that ( ii ) says that ¯ Σ -1 t +1 , 1 -¯ Σ -1 t, 1 is one-rank 808 which we will also need in induction step. Now, we have that ∥ X t ∥ 2 = 1 . Therefore, 809

810

811

812

where we use that by definition of σ 2 MAX in Lemma D.3, we have that σ 2 MAX ≥ 1 + σ -2 σ 2 1 . Therefore, by taking the inverse, we get that 1 1+ u ⊤ u ≥ σ -2 MAX . Combining this with Eq. (32) leads to

Noticing that P A t , 1 = ˆ Σ t,A t Σ -1 1 W 1 concludes the proof of the base case when ℓ = 1 .

(II) Induction step. Let ℓ ∈ [ L ] / { 1 } and suppose that ¯ Σ -1 t +1 ,ℓ -1 -¯ Σ -1 t,ℓ -1 is one-rank and that it 813 holds for ℓ -1 that 814

Then, we want to show that ¯ Σ -1 t +1 ,ℓ -¯ Σ -1 t,ℓ is also one-rank and that it holds that 815

This is achieved as follows. First, we notice that by the induction hypothesis, we have that ˜ Σ -1 t +1 ,ℓ -1 -816 ¯ G t,ℓ -1 = ¯ Σ -1 t +1 ,ℓ -1 -¯ Σ -1 t,ℓ -1 is one-rank. In addition, the matrix is positive semi-definite. Thus we 817 can write it as ˜ Σ -1 t +1 ,ℓ -1 -¯ G t,ℓ -1 = uu ⊤ where u ∈ R d . Then, similarly to the base case, we have 818

However, we it follows from the induction hypothesis that uu ⊤ = ˜ Σ -1 t +1 ,ℓ -1 -¯ G t,ℓ -1 = ¯ Σ -1 t +1 ,ℓ -1 -819 ¯ Σ -1 t,ℓ -1 ⪰ σ -2 σ -2( ℓ -1) MAX P ⊤ A t ,ℓ -1 X t X ⊤ t P A t ,ℓ -1 . Therefore, 820

Finally, we use that 1 + u ⊤ ¯ Σ t,ℓ -1 u ≤ 1 + ∥ u ∥ 2 λ 1 ( ¯ Σ t,ℓ -1 ) ≤ 1 + σ -2 σ 2 ℓ . Here we use that 821 ∥ u ∥ 2 ≤ σ -2 , which can also be proven by induction, and that λ 1 ( ¯ Σ t,ℓ -1 ) ≤ σ 2 ℓ , which follows from 822 the expression of ¯ Σ t,ℓ -1 in Section 3.1. Therefore, we have that 823

where the last inequality follows from the definition of σ 2 MAX = max ℓ ∈ [ L ] 1 + σ -2 σ 2 ℓ . This concludes 824 the proof. 825

826

D.5 Proof of theorem 4.1

We start with the following standard result which we borrow from [Hong et al., 2022a, Aouali et al., 827 2023b], 828

Then we use Lemma D.2 and express the marginal covariance ˇ Σ t,A t as 829

Therefore, we can decompose ∥ X t ∥ 2 ˇ Σ t,A t as 830

where ( i ) follows from Eq. (34), and we use the following inequality in ( ii ) 831

which holds for any x ∈ [0 , u ] , where constants c 0 and c ℓ are derived as 832

The derivation of c 0 uses that 833

The derivation of c ℓ follows from 834

Therefore, from Eq. (35) and Eq. (33), we get that 835

Now we focus on bounding the logarithmic terms in Eq. (36). 836

(I) First term in Eq. (36) We first rewrite this term as 837

where ( i ) follows from the Weinstein-Aronszajn identity. Then we sum over all rounds t ∈ [ n ] , and 838 get a telescoping 839

where ( i ) follows from the fact that ˆ Σ 1 ,i = Σ 1 . Now we use the inequality of arithmetic and 840 geometric means and get 841

(II) Remaining terms in Eq. (36) Let ℓ ∈ [ L ] . Then we have that 842

where we use the Weinstein-Aronszajn identity in ( i ) . Now we know from Lemma D.3 that the 843 following inequality holds σ -2 σ -2 ℓ MAX P ⊤ A t ,ℓ X t X ⊤ t P A t ,ℓ ⪯ ¯ Σ -1 t +1 ,ℓ -¯ Σ -1 t,ℓ . As a result, we get that 844 ¯ Σ -1 t,ℓ + σ -2 σ -2 ℓ MAX P ⊤ A t ,ℓ X t X ⊤ t P A t ,ℓ ⪯ ¯ Σ -1 t +1 ,ℓ . Thus, 845

Then we sum over all rounds t ∈ [ n ] , and get a telescoping 846

where we use that ¯ Σ 1 ,ℓ = Σ ℓ +1 in ( i ) . Finally, we use the inequality of arithmetic and geometric 847 means and get that 848

The last inequality follows from the expression of ¯ Σ -1 n +1 ,ℓ in Eq. (18) that leads to 849

1 1 as 850

where we use the assumption that λ 1 (W ⊤ ℓ W ℓ ) = 1 (A2) and that λ 1 (Σ ℓ +1 ) = σ 2 ℓ +1 and λ 1 (Σ -1 ℓ ) = 851 1 /σ 2 ℓ . This is because Σ ℓ = σ 2 ℓ I d for any ℓ ∈ [ L +1] . Finally, plugging Eqs. (37) and (38) in Eq. (36) 852 concludes the proof. 853

854

855

856

D.6 Proof of proposition 4.2

We use exactly the same proof in Appendix D.5, with one change to account for the sparsity assumption (A3) . The change corresponds to Eq. (38). First, recall that Eq. (38) writes

where 857

where the second equality follows from the assumption that Σ ℓ +1 = σ 2 ℓ +1 I d . But notice that in 858 our assumption, (A3) , we assume that W ℓ = ( ¯ W ℓ , 0 d,d -d ℓ ) , where ¯ W ℓ ∈ R d × d ℓ for any ℓ ∈ [ L ] . 859 Therefore, we have that for any d × d matrix B ∈ R dd × d , the following holds, W ⊤ ℓ BW ℓ = 860 ( ¯ W ⊤ ℓ B ¯ W ℓ 0 d ℓ ,d -d ℓ 0 d -d ℓ ,d ℓ 0 d -d ℓ ,d -d ℓ ) . In particular, we have that 861

Therefore, plugging this in Eq. (41) yields that 862

As a result, det(Σ 1 2 ℓ +1 ¯ Σ -1 n +1 ,ℓ Σ 1 2 ℓ +1 ) = det( I d ℓ + σ 2 ℓ +1 ¯ W ⊤ ℓ ( Σ -1 ℓ -Σ -1 ℓ ¯ Σ t,ℓ -1 Σ -1 ℓ ) ¯ W ℓ ) . This allows 863 us to move the problem from a d -dimensional one to a d ℓ -dimensional one. Then we use the inequality 864

of arithmetic and geometric means and get that 865

To get the last inequality, we use derivations similar to the ones we used in Eq. (40). Finally, the 866 desired result in obtained by replacing Eq. (38) by Eq. (44) in the previous proof in Appendix D.5. 867

868

869

870

871

872

873

874

875

876

877

878

879

880

881

882

883

884

885

886

887

888

889

890

891

892

893

894

895

896

897

898

899

D.7 Additional discussion: link to two-level hierarchies

The linear diffusion (15) can be marginalized into a 2-level hierarchy using two different strategies. The first one yields,

with Ω 1 = σ 2 1 I d + ∑ L -1 ℓ =1 σ 2 ℓ +1 B ℓ B ⊤ ℓ and B ℓ = ∏ ℓ k =1 W k . The second strategy yields,

where Ω 2 = ∑ L ℓ =1 σ 2 ℓ +1 B ℓ B ⊤ ℓ . Recently, HierTS [Hong et al., 2022b] was developed for such two-level graphical models, and we call HierTS under (45) by HierTS-1 and HierTS under (46) by HierTS-2 . Then, we start by highlighting the differences between these two variants of HierTS . First, their regret bounds scale as

When K ≈ L , the regret bounds of HierTS-1 and HierTS-2 are similar. However, when K > L , HierTS-2 outperforms HierTS-1 . This is because HierTS-2 puts more uncertainty on a single d -dimensional latent parameter ψ ∗ , 1 , rather than K individual d -dimensional action parameters θ ∗ ,i . More importantly, HierTS-1 implicitly assumes that action parameters θ ∗ ,i are conditionally independent given ψ ∗ ,L , which is not true. Consequently, HierTS-2 outperforms HierTS-1 . Note that, under the linear diffusion model (15), dTS and HierTS-2 have roughly similar regret bounds. Specifically, their regret bounds dependency on K is identical, where both methods involve multiplying K by σ 2 1 , and both enjoy improved performance compared to HierTS-1 . That said, note that Theorem 4.1 and Proposition 4.2 provide an understanding of how dTS 's regret scales under linear score functions f ℓ , and do not say that using dTS is better than using HierTS when the score functions f ℓ are linear since the latter can be obtained by a proper marginalization of latent parameters (i.e., HierTS-2 instead of HierTS-1 ). While such a comparison is not the goal of this work, we still provide it for completeness next.

When the mixing matrices W ℓ are dense (i.e., assumption (A3) is not applicable), dTS and HierTS-2 have comparable regret bounds and computational efficiency. However, under the sparsity assumption (A3) and with mixing matrices that allow for conditional independence of ψ ∗ , 1 coordinates given ψ ∗ , 2 , dTS enjoys a computational advantage over HierTS-2 . This advantage explains why works focusing on multi-level hierarchies typically benchmark their algorithms against two-level structures akin to HierTS-1 , rather than the more competitive HierTS-2 . This is also consistent with prior works in Bayesian bandits using multi-level hierarchies, such as Tree-based priors [Hong et al., 2022a], which compared their method to HierTS-1 . In line with this, we also compared dTS with HierTS-1 in our experiments. But this is only given for completeness as this is not the aim of Theorem 4.1 and Proposition 4.2. More importantly, HierTS is inapplicable in the general case in (1) with non-linear score functions since the latent parameters cannot be analytically marginalized.

900

901

902

903

904

905

906

907

908

909

910

911

912

913

914

E Broader impact

This work contributes to the development and analysis of practical algorithms for online learning to act under uncertainty. While our generic setting and algorithms have broad potential applications, the specific downstream social impacts are inherently dependent on the chosen application domain. Nevertheless, we acknowledge the crucial need to consider potential biases that may be present in pre-trained diffusion models, given that our method relies on them.

F Limitations

Our work investigated contextual bandits, laying the groundwork for future exploration into reinforcement learning. This exploration can be done from both practical (empirical) and theoretical angles. While our method, which approximates rewards using a Gaussian distribution, worked well for linear rewards and those following a generalized linear model, its effectiveness in real-world, complex scenarios needs further testing. Another interesting direction for future research is pre-training the diffusion model prior. Hsieh et al. [2023] proposed a method for this in multi-armed bandits, but its application to contextual bandits remains unexplored.

G Amount of computation required

Our experiments were conducted on internal machines with 30 CPUs and thus they required a moder915 ate amount of computation. These experiments are also reproducible with minimal computational 916 resources. 917

918

919

920

921

922

923

924

925

926

927

928

929

930

931

932

933

934

935

936

937

938

939

940

941

942

943

944

945

946

947

948

949

950

951

952

953

954

955

956

957

958

959

960

961

962

963

964

965

966

967

968

969

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: All claims are supported by the theory in Section 4 (with proofs provided in the appendix) and experiments in Section 5.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Limitations were discussed in Section 6 and Appendix F.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Assumptions are mentioned in the main text. Complete proofs are provided in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main 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: Information needed to reproduce the main experimental results of the paper is described in Section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

1023

1024

1025

1026

1027

1028

1029

1030

1031

1032

1033

1034

1035

1036

1037

1038

1039

1040

1041

1042

1043

1044

1045

1046

1047

1048

1049

1050

1051

1052

1053

1054

1055

1056

1057

1058

1059

1060

1061

1062

1063

1064

1065

1066

1067

1068

1069

1070

1071

1072

1073

1074

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: The code for the main experiments is shared in the supplementary material. Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

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: All experimental details are described in Section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

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

Answer: [Yes]

Justification: Standard error bars are included in the figures.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

1075

1076

1077

1078

1079

1080

1081

1082

1083

1084

1085

1086

1087

1088

1089

1090

1091

1092

1093

1094

1095

1096

1097

1098

1099

1100

1101

1102

1103

1104

1105

1106

1107

1108

1109

1110

1111

1112

1113

1114

1115

1116

1117

1118

1119

1120

1121

1122

1123

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

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: [Yes]

Justification: As mentioned in Appendix G, our experiments were conducted on internal machines with 30 CPUs and thus they required a moderate amount of computation. These experiments are also reproducible with minimal computational resources.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This work contributes to the development and theoretical analysis of online learning to act under uncertainty and it adheres to the Neurips Code Of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Broader Impacts are discussed in Appendix E.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

1124

1125

1126

1127

1128

1129

1130

1131

1132

1133

1134

1135

1136

1137

1138

1139

1140

1141

1142

1143

1144

1145

1146

1147

1148

1149

1150

1151

1152

1153

1154

1155

1156

1157

1158

1159

1160

1161

1162

1163

1164

1165

1166

1167

1168

1169

1170

1171

1172

1173

1174

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

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

Answer: [NA]

Justification: Our paper is mainly theoretical and the used data is simulated. Thus, we believe that our work poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

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: To the best of our knowledge, all relevant and used papers were cited.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

1175

1176

1177

1178

1179

1180

1181

1182

1183

1184

1185

1186

1187

1188

1189

1190

1191

1192

1193

1194

1195

1196

1197

1198

1199

1200

1201

1202

1203

1204

1205

1206

1207

1208

1209

1210

1211

1212

1213

1214

1215

1216

1217

1218

1219

1220

1221

1222

1223

1224

1225

1226

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We include our code as supplementary material, with all details needed for reproducibility given in Section 5.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

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: This work does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This work does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

1227

1228

1229

1230

1231

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
The authors propose diffusion Thompson sampling, which uses a diffusion model to leverage reward under similar actions for more efficient exploration. The authors derive efficient posterior approximations under a diffusion model prior and prove a regret bound in linear instances. To efficiently compute and sample posterior distribution, the authors provide an approximation that relies on close-form solutions for case where both the score functions of the diffusion model and the likelihood are linear. For nonlinear diffusion model, the authors approximate posteriors by a Gaussian distribution.

STRENGTHS:
The proof of Theorem 4.1 requires novel techniques such as recursive total variance decomposition and refined arguments such as quantifying not only the posterior information gain for the taken action but also for every learnt latent parameter. 

The paper is well-written. The main contributions and key observations from the regret bound are nicely summarized. Experimental results for all four combinations of linear and nonlinear reward, linear and nonlinear diffusion model are provided. In experiments, the authors made a number of insightful observations, accompanied by ablation results.

WEAKNESSES:
The authors discussed how the number of layers L affect the regret bound. A higher L increases regret bound and a smaller L may fail to capture a more complex prior. It would improve the paper to provide a heuristics on choosing an appropriate L along with justifications for the heuristics.

QUESTIONS:
Please see weakness above.

LIMITATIONS:
The authors addressed limitations and societal impact at Appendix E and F.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work presents the use of Diffusion models as priors for Thompson sampling.\n\nNamely, they propose to learn diffusion models (as replacement to other parametric priors) to accommodate more complex correlations between context, action and reward functions than with simple parametric form priors.\n\nGiven that Thompson sampling requires sampling from the posterior of the model, the authors derive a linear-Gaussian posterior approximation (under the proposed diffusion model prior).\n\nThe authors analyze the proposed algorithm for the linear-Gaussian reward case, which enables them to provide a Bayes regret bound.\n\nExperimental results demonstrate some of the benefits of the proposed diffusion-based Thompson sampling: learning the correct latent-structure is beneficial, learning more parameters (as a function of $d$, $K$ and $L$) is a harder problem, hence incurs in higher regret.\n\nSTRENGTHS:\n- The use of diffusion models to learn complex priors for their use within MAB problems is of interest and significant.\n\n- The authors provide a theoretical analysis of their proposed algorithm (only for the linear-Gaussian case), for which they: \n    - use the recursive total covariance decomposition,\n    - showcase the dependency over K ---induced by the hierarchical parameter learning--- and\n    - demonstrate the dependency with $L$, inherent to having more parameters to learn.\n    \n- The theoretical analysis and the experiments showcase the benefits of learning the true hierarchical model (as specified by a diffusion model) in comparison to LinTS.\n\nWEAKNESSES:\n- The proposed diffusion-based algorithm does not learn the diffusion model as it sequentially interacts with the world\n    - Instead, using the diffusion model as a complex prior requires offline learning, so that non-trivial prior distributions can be learned, before it can be used within Thompson sampling.\n    - The cost of learning such a diffusion model is not acknowledged nor discussed.\n\n- The proposed posterior approximation seems to be equivalent to the well known Laplace approximation, i.e., a linear-Gaussian approximation to a (non-linear and non-Gaussian) posterior. See questions below.\n\n- The provided Bayesian regret is limited to the linear-Gaussian case, and in fact is acknowledged to be similar to \"L + 1 sequential linear bandit instances stacked upon each other\".\n\n- The empirical evaluation is executed on synthetic experiments simulated from the assumed model prior, with $L$ latent parameters. Hence, the benefits of learning the true model are somehow expected.\n\nQUESTIONS:\n- The authors present this work as a new bandit algorithm that merges diffusion models with Thompson sampling.\n    - However, one must \"pre-train a diffusion model on offline estimates of the action parameters\", and then use such diffusion model as a prior for Thompson sampling.\n    - Hence, is the novelty of this work on proposing to use a flexible prior, that is learned offline, which is then applied to TS?\n    - **Authors have clarified in their rebuttal how they pre-train the diffusion model offline, and update it online within TS**\n\n- How does the proposed linear-Gaussian approximation to posteriors with non-Gaussian likelihoods in Equation (7) and non-linear functions in Section 3.2 relate to the Laplace approximation?\n   - Such approximation is well known in the nonlinear state-space literature, can the authors please clarify the connection and significance?\n   - **Clarified in the response**\n\n- The authors claim to \"demonstrate that with diffusion models parameterized by linear score functions and linear rewards, we can derive exact closed-form posteriors without approximations\".\n    - How is this result different from the well known result that linear combinations of Gaussian distributions are still Gaussian distributions?\n    - **Addressed in rebuttal**\n\n- Can the authors provide details on what is the offline procedure to learn the diffusion model prior in their synthetic experiments?\n    - How much data did they use to learn these?\n    - How sensitive is the procedure to poorly fitted diffusion priors?\n    - **Addressed in rebuttal, please incorporate details and experiments to final manuscript**\n    \n- Even if the authors provide some results under the misspecified case in Section 5.2.3), can the authors provide more empirical results showcasing how the evaluated models perform when the true underlying model is different?\n    - E.g., when there is a misspecification on the number of latent parameter layers $L$?\n    - E.g., when the actual model does not have latent parameters, i.e., the assumptions made by LinTS?\n       - This would help assess the impact of learning a complex diffusion model, when simpler parametric priors suffice.\n    - **Addressed in rebuttal, please incorporate details and experiments to final manuscript**\n    \n- The authors refer to $f_l$ (the mapping of parameters at layer $l$ to the expected value of the parameters at layer $l-1$) as the \"score function\" . In statistics, the score function is often referred to as the derivative of the log-likelihood function with respect to the parameter. Are there any connections here that motivate the authors to use such naming? Clarifications on this matter would be appreciated.\n    - **Clarified in their rebuttal**\n\nLIMITATIONS:\nThe authors do present general limitations of their work, although the cost associated with learning a diffusion model prior is less clear.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

comments: SUMMARY:
The work provides a great example of diffusion modeling on bandit action parameter for better exploration.

STRENGTHS:
The work provides a comprehensive description on how to employ diffusion modeling on bandit parameters for contextual bandit problems. 

The discussion on linear and non-linear diffusion model is clean and precise for readers with background in Thompson Sampling 

The analysis part also provide comprehensive discussion on how the regret of the  proposed diffusion Thompson Sampling scales with main dimension of contextual bandit problems.

WEAKNESSES:
I am satisfied with current version of the paper.

QUESTIONS:
(1) Could you discuss how to validate assumption (A1)(A2)(A3) in practice and also contrast the assumptions in the literature? 

(2) Could you discuss on some strategy to prove frequentist regret for diffusion Thompson sample?

LIMITATIONS:
Yes, the author clearly states the assumptions to address the limitation of the theoretical analysis.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper considers the problem of contextual bandits in large action spaces.  In this problem, the reward of an arm is a function of the context and an unknown, arm specific parameter vector. To efficiently learn good policies in such large action spaces, the paper places a structured-prior distribution on the unknown arm parameters that can effectively capture the correlations between the arms. The specific form of the prior distribution considered in the paper resembles a diffusion model. The main contribution of the paper is to provide a computationally efficient heuristic for performing Thompson sampling with this prior. Experiments on synthetic data show that the proposed technique is much better at learning optimal policies than other popular baselines such as LinTS, LinUCB.\n\nSTRENGTHS:\nThe problem of handling large action spaces in contextual bandits seems interesting. The empirical evaluation shows promise in the proposed approach\n\nWEAKNESSES:\n- **Related Work:** There are several ways in which large action spaces are typically handled in contextual bandits. One popular approach that is used in practice is to associate a feature vector to each arm (this feature vector is known to the learner ahead of time), and the reward of pulling certain arm for a context is a function of both the context and arm features. In the absence of arm features, the other approach is to impose some structure on the unknown arm parameter vectors. There are several works which do this, and the current paper falls in this line of research. Some of these works assume the arms can be clustered into a small number of groups or can be embedded in a low-dimensional latent space and learn the low-dimensional features during the course of the online learning (https://arxiv.org/pdf/2010.12363, https://arxiv.org/abs/2209.03997, https://arxiv.org/pdf/1810.09401, https://proceedings.neurips.cc/paper_files/paper/2023/file/f334c3375bd3744e98a0ca8eaa2403b0-Paper-Conference.pdf). The diffusion prior used in the current work resembles the low-dimensional embedding assumption. In particular, it is assumed there is a latent vector (psi_1) from which all the arm parameter vectors are generated; this is a form of rank-1 assumption on the arm features. Unfortunately, none of these works were brought up in the paper. It would be great if the authors perform a thorough literature review and better position their work.\n\n- **Linear Setting**: A lot of emphasis has been placed on the linear model in the paper. I understand it is used to derive the heuristic for the non-linear setting. Beyond that, I do not find the regret bounds derived in section 4 to be interesting. In the linear setting, there isn't a need to work with the complex hierarchical diffusion prior. It looks like one could totally remove the latent variables psi_{*, L}, .... psi_{*, 2} and simply place a Gaussian prior on psi_{*,1} and get an equally powerful model. This would also improve the regret bounds, by removing the L factor in the regret. Given this, I'm not sure about the utility of section 4.\n  - Section 4.1 compares the regret bounds obtained in this work with other baselines. But this comparison is only meaningful under the assumption that the diffusion prior is properly specified. This raises the following question: why is this a reasonable prior to use in practice? How do various techniques compare if this prior is misspecified? (There are some experiments section 5.2 on prior misspecification, but the misspecifications considered there seem to be very minor)\n\n- **Quality of Heuristics:** How good is the heuristic used for non-linear diffusion model? There is no discussion on this in the paper (In my opinion, this needs to be thoroughly discussed in the paper, as it is the primary novelty of the work). Some empirical evaluation comparing it with other standard estimation techniques (such as variational techniques, and other posterior estimation techniques) would have been helpful in understanding his question.\n\nQUESTIONS:\n- **Choice of prior**: What is the reason behind choosing this particular choice of prior? \n- **Offline samples for estimating prior**: Line 33 says that offline estimates of parameters are leveraged to build a diffusion model. But this point was never brought up in the rest of the paper. What is the problem setting considered in this work? If the prior is estimated from offline data, then more information needs to be provided on how this is done and how many samples are needed to get a good estimate.\n- **Non-linear diffusion model:** Can something be said about the quality of the approximation in the asymptotic setting where the samples go to infinity?\n\nLIMITATIONS:\nSee my comments above", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-21"}