Diffusion Models Meet Contextual Bandits

## Abstract

Abstract Efficient exploration is a key challenge in contextual bandits due to the large size of their action space, where uninformed exploration can result in computational and statistical inefficiencies. Fortunately, the rewards of actions are often correlated and this can be leveraged to explore them efficiently. In this work, we capture such correlations using pre-trained diffusion models; upon which we design diffusion Thompson sampling ( dTS ). Both theoretical and algorithmic foundations are developed for dTS , and empirical evaluation also shows its favorable performance.

## 1 Introduction

A *contextual bandit* is a popular and practical framework for online learning under uncertainty (Li et al., ). 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., ) or *Thompson sampling (TS)* (Scott, ).

The action space in contextual bandits is often large, resulting in less-than-optimal performance with standard exploration strategies. Luckily, actions often 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 prior distributions (Hong et al., ) that can capture these correlations. Inspired by the achievements of diffusion models (Sohl-Dickstein et al., ; Ho et al., ), which effectively approximate complex distributions (Dhariwal & Nichol, ; Rombach et al., ), this work captures action correlations by employing diffusion models as priors in contextual Thompson sampling.

We illustrate the idea using a video streaming scenario. The objective is to optimize watch time for a user $j$ by selecting a video $i$ from a catalog of $K$ videos. Users $j$ and videos $i$ are associated with context vectors $x_{j}$ and unknown video parameters $\theta_{i}$, respectively. User $j$’s expected watch time for video $i$ is linear as $x_{j}^{\top}\theta_{i}$. Then, a natural strategy would be to independently learn video parameters $\theta_{i}$ using LinTS or LinUCB (Agrawal & Goyal, ; Abbasi-Yadkori et al., ), but this proves statistically inefficient for larger $K$. Luckily, videos exhibit correlations and can provide informative insights into one another. To capture this, we leverage offline estimates of video parameters denoted by $\hat{\theta}_{i}$ and build a diffusion model on them. This diffusion model approximates the video parameter distribution, capturing their dependencies. This model enriches contextual Thompson sampling as a prior, effectively capturing complex video dependencies while ensuring computational efficiency.

Formally, we present a unified contextual bandit framework represented by diffusion models. On this basis, we design a computationally and statistically efficient diffusion Thompson sampling (dTS). We then specialize dTS on linear instances, for which we provide closed-form solutions and establish an upper bound for its Bayes regret. The regret bound reflects the structure of the problem and the quality of the priors, demonstrating the benefits of using diffusion models as priors (dTS) over the standard methods. Beyond enabling theoretical analysis, these linear instances motivate an efficient approximation for general non-linear ones. Finally, our empirical evaluations validate our theory and demonstrate the strong performance of dTS with closed-form solutions as well as dTS with approximations.

Diffusion models have been used for offline decision-making (Ajay et al., ; Janner et al., ; Wang et al., ). However, their use in online learning was only recently explored by Hsieh et al. (), who focused on *multi-armed bandits without theoretical guarantees*. Our work extends Hsieh et al. () in two ways. First, we extend the idea to the broader contextual bandit framework. This allows us to consider problems where the rewards depend on the context, which is more realistic. Second, we show that when the diffusion model is parametrized by linear score functions, we can derive recursive closed-form posteriors without the need for approximations. Closed-form posteriors are particularly interesting because they facilitate theoretical analysis and motivate more computationally efficient approximations for the case with non-linear score functions; making dTS more scalable. Finally, we provide a theoretical analysis of dTS, which effectively captures the benefits of employing linear diffusion models as priors within contextual Thompson sampling. An extended comparison to related works is provided in [Appendix A](#A1), where we discuss the broader topics of diffusion models and decision-making, hierarchical, structured and low-rank bandits, approximate posteriors, etc.

## 2 Setting

The agent interacts with a *contextual bandit* over $n$ rounds. In round $t\in[n]$, the agent observes a *context* $X_{t}\in\mathcal{X}$, where $\mathcal{X}\subseteq\mathbb{R}^{d}$ is a $d$-dimensional *context space*, it takes an *action* $A_{t}\in[K]$, and then receives a stochastic reward $Y_{t}\in\mathbb{R}$ that depends on both the context $X_{t}$ and the taken action $A_{t}$. Each action $i\in[K]$ is associated with an *unknown action parameter* $\theta_{*,i}\in^{d}$, so that the reward received in round $t$ is $Y_{t}\sim P(\cdot\mid X_{t};\theta_{*,A_{t}})$, where $P(\cdot\mid x;\theta_{*,a})$ is the reward distribution of action $a$ in context $x$. We consider the *Bayesian* bandit setting (Russo & Van Roy, ; Hong et al., ), where the action parameters $\theta_{*,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 $\theta_{*,i}$ are captured through a diffusion model, where they share a set of $L$ consecutive *unknown latent parameters* $\psi_{*,\ell}$ for $\ell\in[L]$. Precisely, the action parameter $\theta_{*,i}$ depends on the $L$-th latent parameter $\psi_{*,L}$ as $\theta_{*,i}\mid\psi_{*,1}\sim\mathcal{N}(f_{1}(\psi_{*,1}),\Sigma_{1})$, where the *score function* $f_{1}:^{d}\rightarrow^{d}$ is *known*. Also, the $\ell-1$-th latent parameter $\psi_{*,\ell-1}$ depends on the $\ell$-th latent parameter $\psi_{*,\ell}$ as $\psi_{*,\ell-1}\mid\psi_{*,\ell}\sim\mathcal{N}(f_{\ell}(\psi_{*,\ell}),\Sigma_{\ell})$, where the score function $f_{\ell}:^{d}\rightarrow^{d}$ is known. Finally, the $L$-th latent parameter $\psi_{*,L}$ is sampled as $\psi_{*,L}\sim\mathcal{N}(0,\Sigma_{L+1})$. We summarize this model in ([1](#S2.E1)) and its graph in [Figure 1](#S2.F1).

$$ $\displaystyle\psi_{*,L}$ $\displaystyle\sim\mathcal{N}(0,\Sigma_{L+1})\,,$ (1) $\displaystyle\psi_{*,\ell-1}\mid\psi_{*,\ell}$ $\displaystyle\sim\mathcal{N}(f_{\ell}(\psi_{*,\ell}),\Sigma_{\ell})\,,$ $\displaystyle\forall\ell\in[L]/\{1\}\,,$ $\displaystyle\theta_{*,i}\mid\psi_{*,1}$ $\displaystyle\sim\mathcal{N}(f_{1}(\psi_{*,1}),\Sigma_{1})\,,$ $\displaystyle\forall i\in[K]\,,$ $\displaystyle Y_{t}\mid X_{t},\theta_{*,A_{t}}$ $\displaystyle\sim P(\cdot\mid X_{t};\theta_{*,A_{t}})\,,$ $\displaystyle\forall t\in[n]\,.$ $$

The model in ([1](#S2.E1)) represents a Bayesian bandit. The agent interacts with a bandit instance defined by the action parameters $\theta_{*,i}$ over $n$ rounds (4-th line in ([1](#S2.E1))). These action parameters $\theta_{*,i}$ are drawn from the generative process in the first 3 lines of ([1](#S2.E1)). Note that ([1](#S2.E1)) can be built by pre-training a diffusion model on offline estimates of the action parameters $\theta_{*,i}$ (Hsieh et al., ).

Figure: Figure 1: Graphical model induced by ([1](#S2.E1)).
Refer to caption: /html/2402.10028/assets/x1.png

A natural goal for the agent in this Bayesian framework is to minimize its *Bayes regret* (Russo & Van Roy, ) that measures the expected performance across multiple bandit instances $\theta_{*}=(\theta_{*,i})_{i\in[K]}$,

$$ $\displaystyle\mathcal{B}\mathcal{R}(n)=\mathbb{E}\left[\sum_{t=1}^{n}r(X_{t},A_{t,*};\theta_{*})-r(X_{t},A_{t};\theta_{*})\right]\,,$ (2) $$

where the expectation in ([2](#S2.E2)) is taken over all random variables in ([1](#S2.E1)). Here $r(x,i;\theta_{*})=\mathbb{E}_{Y\sim P(\cdot\mid x;\theta_{*,i})}\left[Y\right]$ is the expected reward of action $i$ in context $x$ and $A_{t,*}=\operatorname*{arg\,max}_{i\in[K]}r(X_{t},i;\theta_{*})$ 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 Algorithm

We design Thompson sampling that samples the latent and action parameters hierarchically (Lindley & Smith, ). Precisely, let $H_{t}=(X_{k},A_{k},Y_{k})_{k\in[t-1]}$ be the history of all interactions up to round $t$ and let $H_{t,i}=(X_{k},A_{k},Y_{k})_{\{k\in[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 $\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,H_{t}\right)$ recursively as

$$ $\textstyle\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,H_{t}\right)=$ $\textstyle\int_{\psi_{1:L}}Q_{t,L}(\psi_{L})\prod_{\ell=2}^{L}Q_{t,\ell-1}(\psi_{\ell-1}\mid\psi_{\ell})$ $\textstyle\qquad\,\,P_{t,i}(\theta\mid\psi_{1})\mathop{}\!\mathrm{d}\psi_{1:L}\,,\quad\text{where}$ (3) $$

$Q_{t,L}(\psi_{L})=\mathbb{P}\left(\psi_{*,L}=\psi_{L}\,\middle|\,H_{t}\right)$ is the *latent-posterior* density of $\psi_{*,L}\mid H_{t}$. Moreover, for any $\ell\in[L]/\{1\}$, $Q_{t,\ell-1}(\psi_{\ell-1}\mid\psi_{\ell})=\mathbb{P}\left(\psi_{*,\ell-1}=\psi_{\ell-1}\,\middle|\,H_{t},\psi_{*,\ell}=\psi_{\ell}\right)$ is the *conditional latent-posterior* density of $\psi_{*,\ell-1}\mid H_{t},\psi_{*,\ell}=\psi_{\ell}$. Finally, for any action $i\in[K]$, $P_{t,i}(\theta\mid\psi_{1})=\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,H_{t,i},\psi_{*,1}=\psi_{1}\right)$ is the *conditional action-posterior* density of $\theta_{*,i}\mid H_{t,i},\psi_{*,1}=\psi_{1}$.

The decomposition in ([3](#S3.Ex4)) inspires hierarchical sampling. In round $t$, we initially sample the $L$-th latent parameter as $\psi_{t,L}\sim Q_{t,L}(\cdot)$. Then, for $\ell\in[L]/\{1\}$, we sample the $\ell-1$-th latent parameter given that $\psi_{*,\ell}=\psi_{t,\ell}$, as $\psi_{t,\ell-1}\sim Q_{t,\ell-1}(\cdot\mid\psi_{t,\ell})$. Lastly, given that $\psi_{*,1}=\psi_{t,1}$, each action parameter is sampled *individually* as $\theta_{t,i}\sim P_{t,i}(\theta\mid\psi_{t,1})$. This is possible because action parameters $\theta_{*,i}$ are conditionally independent given $\psi_{*,1}$. This leads to [Algorithm 1](#alg1), named diffusion Thompson Sampling (dTS). dTS requires sampling from the $K+L$ posteriors $P_{t,i}$ and $Q_{t,\ell}$. 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}(\cdot\mid\psi_{1})$ writes

$$ $\textstyle P_{t,i}(\theta\mid\psi_{1})\propto\prod_{k\in S_{t,i}}P(Y_{k}\mid X_{k};\theta)\mathcal{N}(\theta;f_{1}(\psi_{1}),\Sigma_{1})\,,$ $$

where $S_{t,i}=\{\ell\in[t-1],A_{\ell}=i\}$ are the rounds where the agent takes action $i$ up to round $t$. Moreover, let $\mathcal{L}_{\ell}(\psi_{\ell})=\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)$ be the likelihood of all observations up to round $t$ given that $\psi_{*,\ell}=\psi_{\ell}$. Then, for any $\ell\in[L]/\{1\}$, the $\ell-1$-th conditional latent-posterior $Q_{t,\ell-1}(\cdot\mid\psi_{\ell})$ writes

$$ $\textstyle Q_{t,\ell-1}(\psi_{\ell-1}\mid\psi_{\ell})$ $\textstyle\propto\mathcal{L}_{\ell-1}(\psi_{\ell-1})\mathcal{N}(\psi_{\ell-1},f_{\ell}(\psi_{\ell}),\Sigma_{\ell})\,,$ $$

and $Q_{t,L}(\psi_{L})\propto\mathcal{L}_{L}(\psi_{L})\mathcal{N}(\psi_{L},0,\Sigma_{L+1})$. All the terms above are known, except the likelihoods $\mathcal{L}_{\ell}(\psi_{\ell})$ for $\ell\in[L]$. These are computed recursively as follows. First, the basis of the recursion writes

$$ $\displaystyle\mathcal{L}_{1}(\psi_{1})=\prod_{i=1}^{K}\int_{\theta_{i}}\prod_{k\in S_{t,i}}P(Y_{k}\mid X_{k};\theta_{i})\mathcal{N}(\theta_{i};f_{1}(\psi_{1}),\Sigma_{1})\mathop{}\!\mathrm{d}\theta_{i}.$ $$

Then for any $\ell\in[L]/\{1\}$, the recursive step follows as

$$ $\displaystyle\mathcal{L}_{\ell}(\psi_{\ell})=\int_{\psi_{\ell-1}}\mathcal{L}_{\ell-1}(\psi_{\ell-1})\mathcal{N}(\psi_{\ell-1};f_{\ell}(\psi_{\ell}),\Sigma_{\ell})\mathop{}\!\mathrm{d}\psi_{\ell-1}\,.$ $$

Figure: Algorithm 1 dTS: diffusion Thompson Sampling

All posterior expressions above use known quantities $(f_{\ell}$, $\Sigma_{\ell}$ and $P(y\mid x;\theta))$. However, these expressions either lead to closed-form solutions or should be approximated, depending on the form of score functions $f_{\ell}$ and reward distribution $P(\cdot\mid x;\theta)$. Specifically, for linear score functions $f_{\ell}$ ([Section 3.1](#S3.SS1)), we can directly derive closed-form posteriors when rewards are also linear ([Section 3.1.1](#S3.SS1.SSS1)). When dealing with non-linear rewards, we can use an efficient approximation based on the closed-form expressions derived for linear rewards ([Section 3.1.2](#S3.SS1.SSS2)). These posterior derivations are beneficial both for theoretical understanding and for creating efficient approximations in situations with non-linear score functions $f_{\ell}$ as we show in [Section 3.2](#S3.SS2).

### 3.1 Linear Diffusion Models

Assume the score functions $f_{\ell}$ are linear such as $f_{\ell}(\psi_{*,\ell})=\mathrm{W}_{\ell}\psi_{*,\ell}$ for $\ell\in[L]$, where $\mathrm{W}_{\ell}\in^{d\times d}$ are the *mixing matrices* and they are known. Then, the model ([1](#S2.E1)) becomes a linear Gaussian system (LGS) (Bishop, ) that writes

$$ $\textstyle\psi_{*,L}$ $\textstyle\sim\mathcal{N}(0,\Sigma_{L+1})\,,$ (4) $\textstyle\psi_{*,\ell-1}\mid\psi_{*,\ell}$ $\textstyle\sim\mathcal{N}(\mathrm{W}_{\ell}\psi_{*,\ell},\Sigma_{\ell})\,,$ $\textstyle\forall\ell\in[L]/\{1\}\,,$ $\textstyle\theta_{*,i}\mid\psi_{*,1}$ $\textstyle\sim\mathcal{N}(\mathrm{W}_{1}\psi_{*,1},\Sigma_{1})\,,$ $\textstyle\forall i\in[K]\,,$ $\textstyle Y_{t}\mid X_{t},\theta_{*,A_{t}}$ $\textstyle\sim P(\cdot\mid X_{t};\theta_{*,A_{t}})\,,$ $\textstyle\forall t\in[n]\,.$ $$

This model is important, both in theory and practice. For theory, ([4](#S3.E4)) leads to closed-form posteriors that allow bounding the regret of dTS. For practice, these closed-form expressions could be used to motivate efficient approximations for the general case in ([1](#S2.E1)) as we show in [Section 3.2](#S3.SS2).

#### 3.1.1 Linear Diffusion and rewards

We start with linear-Gaussian reward distributions $P(\cdot\mid x;\theta_{*,a})=\mathcal{N}(\cdot;x^{\top}\theta_{*,a},\sigma^{2})$ where $\sigma>0$ is the observation noise. Then, we can derive closed-form posteriors for this model(^1^11The proofs are provided in [Appendix B](#A2) and they are based on standard derivations of linear Gaussian systems (Bishop, ).), and they are computed recursively. First, let

$$ $\textstyle\hat{G}_{t,i}=\sigma^{-2}\sum_{k\in S_{t,i}}X_{k}X_{k}^{\top}\,,$ $\textstyle\hat{B}_{t,i}=\sigma^{-2}\sum_{k\in S_{t,i}}Y_{k}X_{k}\,,$ $$

where $S_{t,i}=\{\ell\in[t-1]:A_{\ell}=i\}$ are the rounds where the agent takes action $i$ up to round $t$. Then, the conditional action-posterior reads $P_{t,i}(\cdot\mid\psi_{1})=\mathcal{N}(\cdot;\hat{\mu}_{t,i},\hat{\Sigma}_{t,i})$, with

$$ $\displaystyle\hat{\Sigma}_{t,i}^{-1}=\Sigma_{1}^{-1}+\hat{G}_{t,i},\,\,\hat{\mu}_{t,i}=\hat{\Sigma}_{t,i}\big{(}\Sigma_{1}^{-1}\mathrm{W}_{1}\psi_{1}+\hat{B}_{t,i}\big{)}.$ (5) $$

For $\ell\in[L]/\{1\}$, the $\ell-1$-th conditional latent-posterior writes $Q_{t,\ell-1}(\cdot\mid\psi_{\ell})=\mathcal{N}(\bar{\mu}_{t,\ell-1},\bar{\Sigma}_{t,\ell-1})$, where

$$ $\displaystyle\bar{\Sigma}_{t,\ell-1}^{-1}$ $\displaystyle=\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1}\,,$ $\displaystyle\bar{\mu}_{t,\ell-1}$ $\displaystyle=\bar{\Sigma}_{t,\ell-1}\big{(}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\psi_{\ell}+\bar{B}_{t,\ell-1}\big{)}\,,$ (6) $$

and the $L$-th latent-posterior is $Q_{t,L}(\cdot)=\mathcal{N}(\bar{\mu}_{t,L},\bar{\Sigma}_{t,L})$,

$$ $\textstyle\bar{\Sigma}_{t,L}^{-1}=\Sigma_{L+1}^{-1}+\bar{G}_{t,L}\,,$ $\textstyle\bar{\mu}_{t,L}=\bar{\Sigma}_{t,L}\bar{B}_{t,L}\,.$ (7) $$

Finally, $\bar{G}_{t,\ell}$ and $\bar{B}_{t,\ell}$ for $\ell\in[L]$ are computed recursively. The basis of the recursion are $\bar{G}_{t,1}$ and $\bar{B}_{t,1}$, which read

$$ $\textstyle\bar{G}_{t,1}$ $\textstyle=\mathrm{W}_{1}^{\top}\sum_{i=1}^{K}\big{(}\Sigma_{1}^{-1}-\Sigma_{1}^{-1}\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\big{)}\mathrm{W}_{1}\,,$ $\textstyle\bar{B}_{t,1}$ $\textstyle=\mathrm{W}_{1}^{\top}\Sigma_{1}^{-1}\sum_{i=1}^{K}\hat{\Sigma}_{t,i}\hat{B}_{t,i}\,.$ (8) $$

Then, the recursive step for $\ell\in[L]/\{1\}$ is,

$$ $\displaystyle\bar{G}_{t,\ell}$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}\,,$ $\displaystyle\bar{B}_{t,\ell}$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\bar{B}_{t,\ell-1}\,.$ (9) $$

#### 3.1.2 Linear diffusion and non-linear rewards

After presenting the closed-form posteriors when both the score functions and rewards are linear, we now consider the case where the score functions are linear but the rewards are non-linear. Precisely, we assume that the reward distribution is parametrized as a generalized linear model (GLM) (McCullagh & Nelder, ). That is, for any $x\in\mathcal{X}$, $P(\cdot\mid x;\theta)$ is an exponential-family
distribution with mean $g(x^{\top}\theta)$, where $g$ is the mean function. For example, we recover logistic bandits (Filippi et al., ) if we let $g(u)=(1+\exp(-u))^{-1}$ and $P(\cdot\mid x;\theta)=\mathrm{Ber}(g(x^{\top}\theta))$, where $\mathrm{Ber}(p)$ be the Bernoulli distribution with mean $p$. Despite the linearity in score functions in ([4](#S3.E4)), we cannot derive closed-form posteriors due to the non-linearity of the rewards. However, since non-linearity only comes from the reward distribution, we will approximate it by a Gaussian distribution which will enable us to use the closed-form solutions in [Section 3.1.1](#S3.SS1.SSS1). This is achieved as follows.

We approximate the likelihoods $\mathbb{P}\left(H_{t,i}\,\middle|\,\theta_{*,i}=\theta\right)$ by multivariate Gaussian densities as follows. Precisely, the reward distribution $P(\cdot\mid x;\theta)$ is an exponential-family distribution. Therefore, the log-likelihoods write $\log\mathbb{P}\left(H_{t,i}\,\middle|\,\theta_{*,i}=\theta\right)=\sum_{k\in S_{t,i}}Y_{k}X_{k}^{\top}\theta-A(X_{k}^{\top}\theta)+C(Y_{k})$, where $C$ is a real function, and $A$ is a twice continuously differentiable function whose derivative is the mean function, $\dot{A}=g$. Now we let $\hat{\theta}_{t,i}^{\textsc{GLM}}$ and $\hat{G}_{t,i}^{\textsc{GLM}}$ be the maximum likelihood estimate (MLE) and the Hessian of the negative log-likelihood, respectively, defined as

$$ $\textstyle\hat{\theta}_{t,i}^{\textsc{GLM}}$ $\textstyle=\operatorname*{arg\,max}_{\theta\in^{d}}\log\mathbb{P}\left(H_{t,i}\,\middle|\,\theta_{*,i}=\theta\right)\,,$ (10) $\textstyle\hat{G}_{t,i}^{\textsc{GLM}}$ $\textstyle=\sum_{k\in S_{t,i}}\dot{g}\big{(}X_{k}^{\top}\hat{\theta}_{t,i}^{\textsc{GLM}}\big{)}X_{k}X_{k}^{\top}\,.$ (11) $$

Then we let $\mathbb{P}\left(H_{t,i}\,\middle|\,\theta_{*,i}=\theta\right)\approx\mathcal{N}\big{(}\theta;\hat{\theta}_{t,i}^{\textsc{GLM}},\big{(}\hat{G}_{t,i}^{\textsc{GLM}}\big{)}^{-1}\big{)}$. This approximation allows us to recover the posterior expressions in [Section 3.1.1](#S3.SS1.SSS1), except that we replace $\hat{B}_{t,i}$ and $\hat{G}_{t,i}$ in [Section 3.1.1](#S3.SS1.SSS1) by $\hat{G}_{t,i}^{\textsc{GLM}}\hat{\theta}_{t,i}^{\textsc{GLM}}$ and $\hat{G}_{t,i}^{\textsc{GLM}}$, respectively.
A question that may arise is why this *Laplace approximation* is applied to the likelihoods $\mathbb{P}\left(H_{t,i}\,\middle|\,\theta_{*,i}=\theta\right)$ instead of the posteriors $\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,H_{t,i}\right)$, as is common in Bayesian inference. The reason is that it allows us to use the expressions for linear rewards ([Section 3.1.1](#S3.SS1.SSS1)) with slight adaptations.

### 3.2 Non-Linear Diffusion Models

After deriving the posteriors for linear score functions $f_{\ell}$, we now get back to the general case in ([1](#S2.E1)), where the score functions are potentially non-linear. We start with the case of non-linear score functions and linear rewards.

#### 3.2.1 Non-linear diffusion and linear rewards

Assume that the reward distribution is linear-Gaussian and writes $P(\cdot\mid x;\theta_{*,a})=\mathcal{N}(\cdot;x^{\top}\theta_{*,a},\sigma^{2})$ where $\sigma>0$ is the observation noise. Then, despite the non-linearity in functions $f_{\ell}$, the conditional action-posteriors $P_{t,i}$ can be computed in closed form thanks to linearity in rewards. Precisely, it has the form in ([5](#S3.E5)) where the term $\mathrm{W}_{1}\psi_{1}$ is replaced by $f_{1}(\psi_{1})$. However, approximation is needed for the latent-posteriors $Q_{t,\ell}$ since the functions $f_{\ell}$ are non-linear. To avoid any computational challenges, we use a simple and intuitive approximation, where all posteriors $Q_{t,\ell}$ are approximated by the Gaussian distributions in [Section 3.1.1](#S3.SS1.SSS1), with few changes. First, the terms $\mathrm{W}_{\ell}\psi_{\ell}$ in ([3.1.1](#S3.Ex13)) are replaced by $f_{\ell}(\psi_{\ell})$. This accounts for the fact that the prior mean is now $f_{\ell}(\psi_{\ell})$ rather than $\mathrm{W}_{\ell}\psi_{\ell}$. Second, the matrix multiplications that involve the matrices $\mathrm{W}_{\ell}$ in ([3.1.1](#S3.Ex14)) and ([3.1.1](#S3.Ex15)) are simply removed.

This approximation, despite being simple, 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](#S2.E1)), and as more data is accumulated, the influence of the prior diminishes.

#### 3.2.2 Non-linear diffusion and rewards

After providing approximate posteriors when the score functions are non-linear but the rewards are linear, we now consider the more general case where both the score functions and rewards are non-linear. Precisely, we assume that the reward distribution is parametrized as a generalized linear model (GLM) as in [Section 3.1.2](#S3.SS1.SSS2). Thus approximation is needed for both $P_{t,i}$ and $Q_{t,\ell}$. To achieve this, we combine the approximation ideas in [Section 3.1.2](#S3.SS1.SSS2) and [Section 3.2.1](#S3.SS2.SSS1). From [Section 3.1.2](#S3.SS1.SSS2), we use the Laplace approximation that allowed us to use the posteriors in [Section 3.1.1](#S3.SS1.SSS1) where we only replace $\hat{B}_{t,i}$ and $\hat{G}_{t,i}$ in [Section 3.1.1](#S3.SS1.SSS1) by $\hat{G}_{t,i}^{\textsc{GLM}}\hat{\theta}_{t,i}^{\textsc{GLM}}$ and $\hat{G}_{t,i}^{\textsc{GLM}}$ in ([10](#S3.E10)), respectively. Then, from [Section 3.2.1](#S3.SS2.SSS1), we borrow the idea of replacing the terms $\mathrm{W}_{\ell}\psi_{\ell}$ in ([3.1.1](#S3.Ex13)) by $f_{\ell}(\psi_{\ell})$, and removing the matrix multiplications that involve $\mathrm{W}_{\ell}$ in ([3.1.1](#S3.Ex14)) and ([3.1.1](#S3.Ex15)). To summarize, the posteriors are approximated here by using the expressions in [Section 3.1.1](#S3.SS1.SSS1), with three changes. First, $\hat{B}_{t,i}$ and $\hat{G}_{t,i}$ are replaced by $\hat{G}_{t,i}^{\textsc{GLM}}\hat{\theta}_{t,i}^{\textsc{GLM}}$ and $\hat{G}_{t,i}^{\textsc{GLM}}$. Second, $\mathrm{W}_{\ell}\psi_{\ell}$ are replaced by $f_{\ell}(\psi_{\ell})$. Third, the matrices $\mathrm{W}_{\ell}$ in ([3.1.1](#S3.Ex14)) and ([3.1.1](#S3.Ex15)) are removed.

## 4 Analysis

We will analyze dTS under the linear diffusion model in ([4](#S3.E4)) with linear rewards ([Section 3.1.1](#S3.SS1.SSS1)). Although our result holds for milder assumptions, we make some simplifications for the sake of clarity and interpretability. We assume that (A1) Contexts satisfy $\|X_{t}\|_{2}^{2}=1$ for any $t\in[n]$. (A2) Mixing matrices and covariances satisfy $\lambda_{1}(\mathrm{W}_{\ell}^{\top}\mathrm{W}_{\ell})=1$ for any $\ell\in[L]$ and $\Sigma_{\ell}=\sigma_{\ell}^{2}I_{d}$ for any $\ell\in[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 $\Sigma_{\ell}$ and arbitrary mixing matrices $\mathrm{W}_{\ell}$. In this section, we write $\tilde{\mathcal{O}}$ for the big-O notation up to polylogarithmic factors. We start by stating our bound for dTS.

###### Theorem 4.1 .

For any $\delta\in(0,1)$, the Bayes regret of *dTS* under ([4](#S3.E4)) with linear rewards, *(A1)* and *(A2)* is bounded as

$$ $\textstyle\mathcal{B}\mathcal{R}(n)\leq\sqrt{2n\big{(}\mathcal{R}^{\textsc{act}}(n)+\sum_{\ell=1}^{L}\mathcal{R}_{\ell}^{\textsc{lat}}\big{)}\log(1/\delta)\Big{)}}+cn\delta\,,$ with $c>0$ is constant, $\sigma_{\textsc{max}}^{2}=\max_{\ell\in[L+1]}1+\frac{\sigma_{\ell}^{2}}{\sigma^{2}}$ , (12) $\textstyle\mathcal{R}^{\textsc{act}}(n)=c_{0}dK\log\big{(}1+\frac{n\sigma_{1}^{2}}{d}\big{)},\ c_{0}=\frac{\sigma_{1}^{2}}{\log\left(1+\sigma_{1}^{2}\right)}\,,$ $\textstyle\mathcal{R}_{\ell}^{\textsc{lat}}=c_{\ell}d\log\big{(}1+\frac{\sigma_{\ell+1}^{2}}{\sigma_{\ell}^{2}}\big{)},c_{\ell}=\frac{\sigma_{\ell+1}^{2}\sigma_{\textsc{max}}^{2\ell}}{\log\left(1+\sigma_{\ell+1}^{2}\right)},$ $$

([4.1](#S4.Ex16)) holds for any $\delta\in(0,1)$. In particular, the term $cn\delta$ is constant when $\delta=1/n$. Then, the bound is $\tilde{\mathcal{O}}(\sqrt{n})$, and this dependence on the horizon $n$ aligns with prior Bayes regret bounds. The bound comprises $L+1$ main terms, $\mathcal{R}^{\textsc{act}}(n)$ and $\mathcal{R}_{\ell}^{\textsc{lat}}$ for $\ell\in[L]$. First, $\mathcal{R}^{\textsc{act}}(n)$ relates to action parameters learning, conforming to a standard form (Lu & Van Roy, ). Similarly, $\mathcal{R}_{\ell}^{\textsc{lat}}$ is associated with learning the $\ell$-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 $\theta_{*,i}\mid H_{t}$ is not explicitly defined. The first contribution is deriving $\theta_{*,i}\mid H_{t}$ using the total covariance decomposition combined with an induction proof, as our posteriors in [Section 3.1](#S3.SS1) were derived recursively. Unlike standard analyses where the posterior distribution of $\theta_{*,i}\mid 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\in[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](#S3.SS1) that connect the posterior covariance of each latent parameter $\psi_{*,\ell}$ with the covariance of the posterior action parameters $\theta_{*,i}$. This allows us to propagate the information gain associated with the action taken in round $A_{t}$ to all latent parameters $\psi_{*,\ell},$ for $\ell\in[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 C](#A3).

To include more structure, we propose the *sparsity* assumption (A3) $\mathrm{W}_{\ell}=(\bar{\mathrm{W}}_{\ell},0_{d,d-d_{\ell}})$, where $\bar{\mathrm{W}}_{\ell}\in^{d\times d_{\ell}}$ for any $\ell\in[L]$. Note that (A3) is not an assumption when $d_{\ell}=d$ for any $\ell\in[L]$. Notably, (A3) incorporates a plausible structural characteristic that could be captured by a diffusion model. Next we present [Theorem 4.1](#S4.Thmtheorem1) under (A3).

###### Proposition 4.2 (Sparsity) .

For any $\delta\in(0,1)$, the Bayes regret of *dTS* under ([4](#S3.E4)) with linear rewards, *(A1)*, *(A2)* and *(A3)* is bounded as

$$ $\textstyle\mathcal{B}\mathcal{R}(n)\leq\sqrt{2n\big{(}\mathcal{R}^{\textsc{act}}(n)+\sum_{\ell=1}^{L}\tilde{\mathcal{R}}_{\ell}^{\textsc{lat}}\big{)}\log(1/\delta)\Big{)}}+cn\delta\,,$ with $c>0$ is constant, $\sigma_{\textsc{max}}^{2}=\max_{\ell\in[L+1]}1+\frac{\sigma_{\ell}^{2}}{\sigma^{2}}$ , (13) $\textstyle\mathcal{R}^{\textsc{act}}(n)=c_{0}dK\log\big{(}1+\frac{n\sigma_{1}^{2}}{d}\big{)},c_{0}=\frac{\sigma_{1}^{2}}{\log\left(1+\sigma_{1}^{2}\right)}\,,$ $\textstyle\tilde{\mathcal{R}}_{\ell}^{\textsc{lat}}=c_{\ell}d_{\ell}\log\big{(}1+\frac{\sigma_{\ell+1}^{2}}{\sigma_{\ell}^{2}}\big{)},c_{\ell}=\frac{\sigma_{\ell+1}^{2}\sigma_{\textsc{max}}^{2\ell}}{\log\left(1+\sigma_{\ell+1}^{2}\right)}.$ $$

From [Proposition 4.2](#S4.Thmtheorem2), our bounds writes

$$ $\textstyle\mathcal{B}\mathcal{R}(n)=\tilde{\mathcal{O}}\Big{(}\sqrt{n(dK\sigma_{1}^{2}+\sum_{\ell=1}^{L}d_{\ell}\sigma_{\ell+1}^{2}\sigma_{\textsc{max}}^{2\ell})}\Big{)}\,,$ (14) $$

since $\mathcal{R}^{\textsc{act}}(n)=\tilde{\mathcal{O}}(dK\sigma_{1}^{2})$ and $\mathcal{R}^{\textsc{lat}}_{\ell}=\tilde{\mathcal{O}}(d_{\ell}\sigma_{\ell+1}^{2}\sigma_{\textsc{max}}^{2\ell})$. Then, smaller values of $K$, $L$, $d$ or $d_{\ell}$ translate to fewer parameters to learn, leading to lower regret. The regret also decreases when the initial variances $\sigma_{\ell}^{2}$ decrease. These dependencies are common in Bayesian analysis, and empirical results match them. They arise from the assumption that true parameters are sampled from a known distribution that matches our prior. When the prior is more informative (such as low variance), the problem is easier, resulting in lower Bayes regret. The reader might question the dependence of our bound on both $L$ and $K$, wondering why $K$ is present. This arises due to our conditional learning of $\theta_{*,i}$ given $\psi_{*,1}$. Rather than assuming deterministic linearity, $\theta_{*,i}=\mathrm{W}_{1}\psi_{*,1}$, we account for stochasticity by modeling $\theta_{*,i}\sim\mathcal{N}(\mathrm{W}_{1}\psi_{*,1},\sigma_{1}^{2}I_{d})$. This makes dTS robust to misspecification scenarios where $\theta_{*,i}$ is not perfectly linear with respect to $\psi_{*,1}$, at the cost of additional learning of $\theta_{*,i}\mid\psi_{*,1}$. If we were to assume deterministic linearity ($\sigma_{1}=0$), our regret bound would scale with $L$ only.

###### Theorem 4.1 .

###### Proposition 4.2 (Sparsity) .

### 4.1 Discussion

#### 4.1.1 Computational benefits

Action correlations prompt an intuitive approach: marginalize all latent parameters and maintain a joint posterior of $(\theta_{*,i})_{i\in[K]}\mid H_{t}$. Unfortunately, this is computationally inefficient for large action spaces. To illustrate, suppose that all posteriors are multivariate Gaussians ([Section 3.1](#S3.SS1)). Then maintaining the joint posterior $(\theta_{*,i})_{i\in[K]}\mid H_{t}$ necessitates converting and storing its $dK\times dK$-dimensional covariance matrix. Then the time and space complexities are $\mathcal{O}(K^{3}d^{3})$ and $\mathcal{O}(K^{2}d^{2})$. In contrast, the time and space complexities of dTS are $\mathcal{O}\big{(}\big{(}L+K\big{)}d^{3}\big{)}$ and $\mathcal{O}\big{(}\big{(}L+K\big{)}d^{2}\big{)}$. This is because dTS requires converting and storing $L+K$ covariance matrices, each being $d\times d$-dimensional. The improvement is huge when $K\gg 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 $\theta_{*,i}\in\mathbb{R}^{d}$ (LinTS). This improves time and space complexity to $\mathcal{O}\big{(}Kd^{3}\big{)}$ and $\mathcal{O}\big{(}Kd^{2}\big{)}$, respectively. However, LinTS maintains independent posteriors and fails to capture the correlations among actions; it only models $\theta_{*,i}\mid H_{t,i}$ rather than $\theta_{*,i}\mid 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.

#### 4.1.2 Statistical benefits

The linear diffusion model in ([4](#S3.E4)) can be transformed into a Bayesian linear model (LinTS) by marginalizing out the latent parameters; in which case the prior on action parameters becomes $\theta_{*,i}\sim\mathcal{N}(0,\,\Sigma)$, with the $\theta_{*,i}$ being not necessarily independent, and $\Sigma$ is the marginal initial covariance of action parameters and it writes $\Sigma=\sigma_{1}^{2}I_{d}+\sum_{\ell=1}^{L}\sigma_{\ell+1}^{2}\mathrm{B}_{\ell}\mathrm{B}_{\ell}^{\top}$ with $\mathrm{B}_{\ell}=\prod_{k=1}^{\ell}\mathrm{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 $\Sigma$ despite integrating them. This causes the *marginal* action uncertainty $\Sigma$ to be much higher than the *conditional* action uncertainty $\sigma_{1}^{2}I_{d}$ in ([4](#S3.E4)), since we have $\Sigma=\sigma_{1}^{2}I_{d}+\sum_{\ell=1}^{L}\sigma_{\ell+1}^{2}\mathrm{B}_{\ell}\mathrm{B}_{\ell}^{\top}\succcurlyeq\sigma_{1}^{2}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 $\Sigma$. This is also reflected by our regret bound. To simply comparisons, suppose that $\sigma\geq\max_{\ell\in[L+1]}\sigma_{\ell}$ so that $\sigma_{\textsc{max}}^{2}\leq 2$. Then the regret bounds of dTS (where we bound $\sigma_{\textsc{max}}^{2\ell}$ by $2^{\ell}$) and LinTS read

$$ $\textstyle\texttt{dTS}:\tilde{\mathcal{O}}\big{(}\sqrt{n(dK\sigma_{1}^{2}+\sum_{\ell=1}^{L}d_{\ell}\sigma_{\ell+1}^{2}2^{\ell})}\big{)}\,,$ $\textstyle\texttt{LinTS}:\tilde{\mathcal{O}}\big{(}\sqrt{ndK(\sigma_{1}^{2}+\sum_{\ell=1}^{L}\sigma_{\ell+1}^{2})}\big{)}\,.$ $$

Then regret improvements are captured by the variances $\sigma_{\ell}$ and the sparsity dimensions $d_{\ell}$, and we proceed to illustrate this through the following scenarios.

(I) Decreasing variances. Assume that $\sigma_{\ell}=2^{\ell}$ for any $\ell\in[L+1]$. Then, the regrets become

$$ $\textstyle\texttt{dTS}:\tilde{\mathcal{O}}\big{(}\sqrt{n(dK+\sum_{\ell=1}^{L}d_{\ell}4^{\ell}))}\big{)}\,,$ $\textstyle\texttt{LinTS}:\tilde{\mathcal{O}}\big{(}\sqrt{ndK2^{L})}\big{)}$ $$

Now to see the order of gain, assume the problem is high-dimensional ($d\gg 1$), and set $L=\log_{2}(d)$ and $d_{\ell}=\lfloor\frac{d}{2^{\ell}}\rfloor$. Then the regret of dTS becomes $\tilde{\mathcal{O}}\big{(}\sqrt{nd(K+L))}\big{)}$, and hence the multiplicative factor $2^{L}$ in LinTS is removed and replaced with a smaller additive factor $L$.

(II) Constant variances. Assume that $\sigma_{\ell}=1$ for any $\ell\in[L+1]$. Then, the regrets become

$$ $\textstyle\texttt{dTS}:\tilde{\mathcal{O}}\big{(}\sqrt{n(dK+\sum_{\ell=1}^{L}d_{\ell}2^{\ell}))}\big{)}\,,$ $\textstyle\texttt{LinTS}:\tilde{\mathcal{O}}\big{(}\sqrt{ndKL)}\big{)}$ $$

Similarly, let $L=\log_{2}(d)$, and $d_{\ell}=\lfloor\frac{d}{2^{\ell}}\rfloor$. Then dTS’s regret is $\tilde{\mathcal{O}}\big{(}\sqrt{nd(K+L)}\big{)}$. 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., ) 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](#S2.E1)), the latent parameters cannot be analytically marginalized, making LinTS with exact marginalization inapplicable.

Large action space aspect. dTS’s regret bound scales with $K\sigma_{1}^{2}$ instead of $K\sum_{\ell}\sigma_{\ell}^{2}$, particularly beneficial when $\sigma_{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., ; Xu & Zeevi, ; Zhu et al., ) proposed bandit algorithms that do not scale with $K$. However, our setting differs significantly from theirs, explaining our inherent dependency on $K$ when $\sigma_{1}>0$. Precisely, they assume a reward function of $r(x,i;\theta_{*})=\phi(x,i)^{\top}\theta_{*}$, with a shared $\theta_{*}\in\mathbb{R}^{d}$ and a known mapping $\phi$. In contrast, we consider $r(x,i;\theta_{*})=x^{\top}\theta_{*,i}$, with $\theta_{*}=(\theta_{*,i})_{i\in[K]}\in^{dK}$, requiring the learning of $K$ separate $d$-dimensional action parameters. In their setting, with the availability of $\phi$, the regret of dTS would similarly be independent of $K$. However, obtaining such a mapping $\phi$ 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.

#### 4.1.3 Link to two-level hierarchies

The linear diffusion ([4](#S3.E4)) can be marginalized into a 2-level hierarchy using two different strategies. The first one yields,

$$ $\displaystyle\psi_{*,L}$ $\displaystyle\sim\mathcal{N}(0,\sigma_{L+1}^{2}\mathrm{B}_{L}\mathrm{B}_{L}^{\top})\,,$ (15) $\displaystyle\theta_{*,i}\mid\psi_{*,L}$ $\displaystyle\sim\mathcal{N}(\psi_{*,L},\,\Omega_{1})\,,$ $\displaystyle\forall i\in[K]\,,$ $$

with $\Omega_{1}=\sigma_{1}^{2}I_{d}+\sum_{\ell=1}^{L-1}\sigma_{\ell+1}^{2}\mathrm{B}_{\ell}\mathrm{B}_{\ell}^{\top}$ and $\mathrm{B}_{\ell}=\prod_{k=1}^{\ell}\mathrm{W}_{k}$.

The second strategy yields,

$$ $\displaystyle\psi_{*,1}$ $\displaystyle\sim\mathcal{N}(0,\Omega_{2})\,,$ (16) $\displaystyle\theta_{*,i}\mid\psi_{*,1}$ $\displaystyle\sim\mathcal{N}(\psi_{*,1},\,\sigma_{1}^{2}I_{d})\,,$ $\displaystyle\forall i\in[K]\,,$ $$

where $\Omega_{2}=\sum_{\ell=1}^{L}\sigma_{\ell+1}^{2}\mathrm{B}_{\ell}\mathrm{B}_{\ell}^{\top}$. Recently, HierTS (Hong et al., ) was developed for such two-level graphical models, and we call HierTS under ([15](#S4.E15)) by HierTS-1 and HierTS under ([16](#S4.E16)) by HierTS-2. Then, we start by highlighting the differences between these two variants of HierTS. First, their regret bounds scale as

$$ $\textstyle\texttt{HierTS-1}:\tilde{\mathcal{O}}\big{(}\sqrt{nd(K\sum_{\ell=1}^{L}\sigma_{\ell}^{2}+L\sigma_{L+1}^{2}}\big{)}\,,$ $\textstyle\texttt{HierTS-2}:\tilde{\mathcal{O}}\big{(}\sqrt{nd(K\sigma_{1}^{2}+\sum_{\ell=1}^{L}\sigma_{\ell+1}^{2})}\big{)}\,.$ $$

When $K\approx 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 $\psi_{*,1}$, rather than $K$ individual $d$-dimensional action parameters $\theta_{*,i}$. More importantly, HierTS-1 implicitly assumes that action parameters $\theta_{*,i}$ are conditionally independent given $\psi_{*,L}$, which is not true. Consequently, HierTS-2 outperforms HierTS-1. Note that, under the linear diffusion model ([4](#S3.E4)), 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 $\sigma_{1}^{2}$, and both enjoy improved performance compared to HierTS-1. That said, note that [Theorems 4.1](#S4.Thmtheorem1) and [4.2](#S4.Thmtheorem2) provide an understanding of how dTS’s regret scales under linear score functions $f_{\ell}$, and do not say that using dTS is better than using HierTS when the score functions $f_{\ell}$ 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 $\mathrm{W}_{\ell}$ 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 $\psi_{*,1}$ coordinates given $\psi_{*,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., ), 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 [Theorems 4.1](#S4.Thmtheorem1) and [4.2](#S4.Thmtheorem2). More importantly, HierTS is inapplicable in the general case in ([1](#S2.E1)) with non-linear score functions since the latent parameters cannot be analytically marginalized.

## 5 Experiments

We evaluate dTS using synthetic data, to validate our theory and test dTS in large action spaces. We omit simulations derived from classification data (Riquelme et al., ) since they typically result in smaller action spaces. This choice is further justified by the fact that Hsieh et al. () has already demonstrated the advantages of diffusion models in multi-armed bandit using such data, without theoretical guarantees. In all experiments, we run 50 random simulations and plot the average regret with its standard error.

Linear diffusion. We consider the settings in [Sections 3.1.1](#S3.SS1.SSS1) and [3.1.2](#S3.SS1.SSS2). The linear rewards ([Section 3.1.1](#S3.SS1.SSS1)) are generated as $\mathcal{N}(x^{\top}\theta,\sigma^{2})$ with $\sigma=1$, and the non-linear rewards ([Section 3.1.2](#S3.SS1.SSS2)) are binary and generated as $\mathrm{Ber}(g(x^{\top}\theta)))$, where $g$ is the sigmoid function. For both settings, the covariances are $\Sigma_{\ell}=I_{d}$, and the context $X_{t}$ is uniformly drawn from $[-1,1]^{d}$. We set $d=20$, $L=4$ and $K=10^{4}$. The score functions are linear as $f_{\ell}(\psi)=\mathrm{W}_{\ell}\psi$ where $\mathrm{W}_{\ell}$ are uniformly drawn from $[-1,1]^{d\times d}$. To introduce sparsity, we zero out the last $d_{\ell}$ columns of $\mathrm{W}_{\ell}$, resulting in $\mathrm{W}_{\ell}=(\bar{\mathrm{W}}_{\ell},0_{d,d-d_{\ell}})$, where $(d_{1},d_{2},d_{3},d_{4})=(20,10,,2)$.

Non-linear diffusion. We use the same setting above, except that the score functions $f_{\ell}$ are now parametrized by two-layer neural networks with random weights in $[-1,1]$, ReLU activation, and a hidden layer dimension of 60. Similarly, we consider both linear rewards ([Section 3.2.1](#S3.SS2.SSS1)) and non-linear ones ([Section 3.2.2](#S3.SS2.SSS2)).

Baselines. When the rewards are linear, we use LinUCB (Abbasi-Yadkori et al., ), LinTS (Agrawal & Goyal, ), and HierTS (Hong et al., ) that marginalizes out all latent parameters except $\psi_{*,L}$. This corresponds to HierTS-1 in [Section 4.1](#S4.SS1). When the rewards are non-linear, we include UCB-GLM (Li et al., ), and GLM-TS (Chapelle & Li, ). GLM-UCB (Filippi et al., ) induced high regret while HierTS was designed for the Gaussian case only and thus both are not included. Finally, we name dTS for each setting as dTS-dr, where the suffix d represents the type of diffusion; L for linear and N for non-linear. The suffix r indicates the nature of rewards; L for linear and N for non-linear rewards. For instance, dTS-LL signifies dTS in linear diffusion with linear rewards ([Section 3.1.1](#S3.SS1.SSS1)).

Figure: Figure 2: Regret of dTS-LL ([Section 3.1.1](#S3.SS1.SSS1)) and dTS-LN ([Section 3.1.2](#S3.SS1.SSS2)) when $d=20$, $K=10^{4}$ and $n=5000$.
Refer to caption: /html/2402.10028/assets/x2.png

Results for linear diffusion ([Figure 2](#S5.F2)). We begin by examining the case of linear rewards in [Figure 2](#S5.F2). First, dTS-LL consistently outperforms all baselines that either disregard the latent structure (LinTS and LinUCB) or incorporate it only partially (HierTS). Specifically, the baselines that disregard the structure (LinTS and LinUCB) fail to converge in $n=5000$ rounds since our problem involves a high-dimensional setting with $d=20$ and a large action space of $K=10^{4}$. These baselines appear to require a more extended horizon $n$ for convergence. Meanwhile, HierTS, which partially leverages the structure, manages to converge within the considered horizon but still incurs a significantly higher regret compared to dTS-LL. In contrast, dTS-LL enjoys rapid convergence and maintains low regret.

In the case of non-linear rewards in [Figure 2](#S5.F2), dTS-LN also surpasses all baselines. Moreover, dTS-LN outperforms dTS-LL in this setting, underscoring the advantages of our approximation in [Section 3.1.2](#S3.SS1.SSS2). While dTS-LN lacks theoretical guarantees, the results suggest that it is preferable to use dTS-LN when dealing with non-linear rewards, emphasizing the generality and adaptability of dTS and the general posterior derivations in [Section 3](#S3). However, despite the misspecification of the reward model in dTS-LL, it still outperforms models that use the correct reward model but neglect the latent structure, such as GLM-TS and UCB-GLM. This highlights the importance of accounting for the latent structure, which can outweigh the correctness of the reward model itself in some cases. Finally, we also conduct an additional experiment to verify the impact of the number of actions $K$, the context dimension $d$, and the diffusion depth $L$ on the regret of dTS. The results are in [Appendix D](#A4) and they match [Theorem 4.1](#S4.Thmtheorem1).

Figure: Figure 3: Regret of dTS-LL and LinTS with varying $K$.
Refer to caption: /html/2402.10028/assets/x3.png

Large action spaces ([Figure 3](#S5.F3)). 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 ([Section 3.1.1](#S3.SS1.SSS1)). [Figure 3](#S5.F3) reports the final cumulative regret for varying values of $K$ for both dTS-LL ([Section 3.1.1](#S3.SS1.SSS1)) and LinTS, observing that the gap in the performance becomes larger as $K$ increases.

Figure: Figure 4: Regret of dTS-NL ([Section 3.2.1](#S3.SS2.SSS1)) and dTS-NN ([Section 3.2.2](#S3.SS2.SSS2)) when $d=20$ and $K=10^{4}$.
Refer to caption: /html/2402.10028/assets/x4.png

Results for non-linear diffusion ([Figure 4](#S5.F4)). We also evaluate our proposed approximation for non-linear diffusion models, with both linear rewards ([Section 3.2.1](#S3.SS2.SSS1)) and non-linear ones ([Section 3.2.2](#S3.SS2.SSS2)). In the linear reward setting,
dTS-NL demonstrates significant performance gains compared to both LinTS and LinUCB baselines. With non-linear rewards, dTS-NN surpasses both GLM-TS and UCB-GLM. Furthermore, dTS-NN outperforms dTS-NL in this setting due to its combined use of the correct reward model and the correct latent structure. Similarly to the linear diffusion case previously discussed, dTS-NL surpasses algorithms that employ the correct reward model but neglect the latent structure (GLM-TS and UCB-GLM). The performance gap between dTS-NL and these algorithms is even more pronounced in this non-linear diffusion setting compared to the linear diffusion setting, potentially due to the increased complexity of the latent structure overshadowing the impact of the reward model itself.

## 6 Conclusion

Grappling with large action spaces in contextual bandits is challenging. Recognizing this, we focused on structured problems where action parameters are sampled from a diffusion model; upon which we built diffusion Thompson sampling (dTS). We developed both theoretical and algorithmic foundations for dTS in numerous practical settings.

We identified several directions for future work. Notably, exploring other approximations for non-linear diffusion models, both empirically and theoretically. From a theoretical perspective, future research could explore the advantages of non-linear diffusion models by deriving their Bayes regret bounds, akin to our analysis in [Section 4](#S4). Empirically, investigating other approximations and assessing their performance in complex tasks would be interesting. Additionally, exploring the extension of this work to offline (or off-policy) learning in contextual bandits represents a promising avenue for future research.

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

## Supplementary Materials

Notation. For any positive integer $n$, we define $[n]=\{1,2,...,n\}$. Let $v_{1},\ldots,v_{n}\in^{d}$ be $n$ vectors, $(v_{i})_{i\in[n]}\in^{nd}$ is the $nd$-dimensional vector obtained by concatenating $v_{1},\ldots,v_{n}$. For any matrix $\mathrm{A}\in\mathbb{R}^{d\times d}$, $\lambda_{1}(\mathrm{A})$ and $\lambda_{d}(\mathrm{A})$ denote the maximum and minimum eigenvalues of $\mathrm{A}$, respectively. Finally, we write $\tilde{\mathcal{O}}$ for the big-O notation up to polylogarithmic factors.

## Appendix 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 & Van Roy, ; Chapelle & Li, ). 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., ; Agrawal & Goyal, ), Gaussian-Gaussian model (Agrawal & Goyal, ), and in the exponential family using Jeffrey’s prior (Korda et al., ). For linear bandits, TS is nearly-optimal (Russo & Van Roy, ; Agrawal & Goyal, ; Abeille & Lazaric, ). In this work, we build TS upon complex diffusion priors and analyze the resulting Bayes regret (Russo & Van Roy, ) in the linear contextual bandit setting.

Decision-making with diffusion models gained attention recently, especially in offline learning (Ajay et al., ; Janner et al., ; Wang et al., ). However, their application in online learning was only examined by Hsieh et al. (), which focused on meta-learning in multi-armed bandits without theoretical guarantees. In this work, we expand the scope of Hsieh et al. () 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., ; Kveton et al., ; Basu et al., ; Simchowitz et al., ; Wan et al., ; Hong et al., ; Peleg et al., ; Wan et al., ; Aouali et al., ) 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 meta- and multi-task learning for multi-armed bandits, except in cases such as Aouali et al. (); Hong et al. () that consider the contextual bandit setting. Precisely, Aouali et al. () 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. () 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_{\ell}$. 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, ). Prior works (Riquelme et al., ; Chapelle & Li, ; Kveton et al., ) highlight the favorable empirical performance of approximate Thompson sampling. Particularly, (Kveton et al., ) 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 & Mannor, ; Hong et al., ), a single latent variable indexes multiple candidate models. Within structured finite-armed bandits (Lattimore & Munos, ; Gupta et al., ), 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., ; Gopalan et al., ). 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., ; Gentile et al., ; Deshmukh et al., ; Cella et al., ). 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., ; Cella et al., ; Yang et al., ) 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 $\theta_{*,i}=\mathrm{W}_{1}\psi_{*,1}$, whereas we incorporate additional uncertainty in the covariance $\Sigma_{1}$ to account for possible misspecification as $\theta_{*,i}=\mathcal{N}(\mathrm{W}_{1}\psi_{*,1},\Sigma_{1})$. Consequently, these algorithms might suffer linear regret due to model misalignment. Second, we assume that the mixing matrix $\mathrm{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\sigma_{1}^{2}$ rather than $K\sum_{\ell}\sigma_{\ell}^{2}$. This is particularly beneficial when $\sigma_{1}$ is small, a common scenario in diffusion models with decreasing variances. A notable case is when $\sigma_{1}=0$, where the regret becomes independent of $K$. Also, our analysis ([Section 4.1](#S4.SS1)) 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., ; Xu & Zeevi, ; Zhu et al., ) proposed bandit algorithms that do not scale with $K$. However, our setting differs significantly from theirs, explaining our inherent dependency on $K$ when $\sigma_{1}>0$. Precisely, they assume a reward function of $r(x,i)=\phi(x,i)^{\top}\theta_{*}$, with a shared $\theta_{*}\in\mathbb{R}^{d}$ across actions and a known mapping $\phi$. In contrast, we consider $r(x,i)=x^{\top}\theta_{*,i}$, requiring the learning of $K$ separate $d$-dimensional action parameters. In their setting, with the availability of $\phi$, the regret of dTS would similarly be independent of $K$. However, obtaining such a mapping $\phi$ 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](#S4.SS1)). This becomes particularly notable in large action spaces. Our empirical results in [Figure 2](#S5.F2), 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.

Bayesian analyses in contextual bandits. Traditionally, contextual bandit algorithms are evaluated using frequentist regret, which measures the performance against a fixed but unknown environment parameter $\theta_{*}$. This approach assumes no prior knowledge about $\theta_{*}$. In contrast, Bayesian regret (Russo & Van Roy, ) assumes that $\theta_{*}$ is drawn from a known prior distribution. The Bayesian regret then calculates the average regret across all possible values of $\theta_{*}$ sampled from the prior. This work employs Bayesian regret because it allows us to leverage informative priors (Hong et al., ; Aouali et al., ). Interestingly, PAC-Bayesian theory also offers a Bayesian perspective and has seen success in capturing the benefits of prior information, as it was demonstrated in offline contextual bandits (London & Sandler, ; Sakhi et al., ; Aouali et al., ). However, this work does not employ PAC-Bayesian analysis, and a connection between PAC-Bayesian theory and Bayesian regret remains unclear.

## Appendix B Derivation of Closed-Form Posteriors for Linear Diffusion Models

In this section, we derive the $K+L$ posteriors $P_{t,i}$ and $Q_{t,\ell}$, for which we provide the expressions in [Section 3.1](#S3.SS1). In our proofs, $p(x)\propto f(x)$ means that the probability density $p$ satisfies $p(x)=\frac{f(x)}{Z}$ for any $x\in^{d}$, where $Z$ is a normalization constant. In particular, we extensively use that if $p(x)\propto\exp[-\frac{1}{2}x^{\top}\Lambda x+x^{\top}m]$, where $\Lambda$ is positive definite. Then $p$ is the multivariate Gaussian density with covariance $\Sigma=\Lambda^{-1}$ and mean $\mu=\Sigma m$. These are standard notations and techniques to manipulate Gaussian distributions (Koller & Friedman, , Chapter 7). We note that the posterior derivations in this section are similar to those used in linear Gaussian systems (Bishop, ).

### B.1 Derivation of Action-Posteriors for Linear Diffusion Models

###### Proposition B.1 .

Consider the following model, which corresponds to the last two layers in ([4](#S3.E4))

$$ $\displaystyle\theta_{*,i}\mid\psi_{*,1}$ $\displaystyle\sim\mathcal{N}\left(\mathrm{W}_{1}\psi_{*,1},\Sigma_{1}\right)\,,$ $\displaystyle Y_{t}\mid X_{t},\theta_{*,A_{t}}$ $\displaystyle\sim\mathcal{N}\left(X_{t}^{\top}\theta_{*,A_{t}},\sigma^{2}\right)\,,$ $\displaystyle\forall t\in[n]\,.$ $$

Then we have that for any $t\in[n]$ and $i\in[K]$, $P_{t,i}(\theta\mid\psi_{1})=\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,\psi_{*,1}=\psi_{1},H_{t,i}\right)=\mathcal{N}(\theta;\hat{\mu}_{t,i},\hat{\Sigma}_{t,i})$,
where

$$ $\displaystyle\hat{\Sigma}^{-1}_{t,i}=\hat{G}_{t,i}+\Sigma_{1}^{-1}\,,$ $\displaystyle\hat{\mu}_{t,i}$ $\displaystyle=\hat{\Sigma}_{t,i}\left(\hat{B}_{t,i}+\Sigma_{1}^{-1}\mathrm{W}_{1}\psi_{1}\right)\,.$ $$

###### Proof.

Let $v=\sigma^{-2}\,,\quad\Lambda_{1}=\Sigma_{1}^{-1}\,.$ Then the action-posterior decomposes as

$$ $\displaystyle P_{t,i}(\theta\mid\psi_{1})$ $\displaystyle=\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,\psi_{*,1}=\psi_{1},H_{t,i}\right)\,,$ $\displaystyle\propto\mathbb{P}\left(H_{t,i}\,\middle|\,\psi_{*,1}=\psi_{1},\theta_{*,i}=\theta\right)\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,\psi_{*,1}=\psi_{1}\right)\,,\quad(\text{Bayes rule})$ $\displaystyle=\mathbb{P}\left(H_{t,i}\,\middle|\,\theta_{*,i}=\theta\right)\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,\psi_{*,1}=\psi_{1}\right)\,,(\text{given $\theta_{*,i}$, $H_{t,i}$ is independent of $\psi_{*,1}$})$ $\displaystyle=\prod_{k\in S_{t,i}}\mathcal{N}(Y_{k};X_{k}^{\top}\theta,\sigma^{2})\mathcal{N}(\theta;\mathrm{W}_{1}\psi_{1},\Sigma_{1})\,,$ $\displaystyle=\exp\Big{[}-\frac{1}{2}\Big{(}v\sum_{k\in S_{t,i}}(Y_{k}^{2}-2Y_{k}X_{k}^{\top}\theta+(X_{k}^{\top}\theta)^{2})+\theta^{\top}\Lambda_{1}\theta-2\theta^{\top}\Lambda_{1}\mathrm{W}_{1}\psi_{1}$ $\displaystyle\qquad\qquad+\big{(}\mathrm{W}_{1}\psi_{1}\big{)}^{\top}\Lambda_{1}\big{(}\mathrm{W}_{1}\psi_{1}\big{)}\Big{)}\Big{]}\,,$ $\displaystyle\propto\exp\Big{[}-\frac{1}{2}\Big{(}\theta^{\top}(v\sum_{k\in S_{t,i}}X_{k}X_{k}^{\top}+\Lambda_{1})\theta-2\theta^{\top}\Big{(}v\sum_{k\in S_{t,i}}X_{k}Y_{k}+\Lambda_{1}\mathrm{W}_{1}\psi_{1}\Big{)}\Big{)}\Big{]}\,,$ $\displaystyle\propto\mathcal{N}\big{(}\theta;\hat{\mu}_{t,i},\hat{\Lambda}_{t,i}^{-1}\big{)}\,,$ $$

where $\hat{\Lambda}_{t,i}=v\sum_{k\in S_{t,i}}X_{k}X_{k}^{\top}+\Lambda_{1}\,,$ and $\hat{\Lambda}_{t,i}\hat{\mu}_{t,i}=v\sum_{k\in S_{t,i}}X_{k}Y_{k}+\Lambda_{1}\mathrm{W}_{1}\psi_{1}$. Using that $\hat{B}_{t,i}=v\sum_{k\in S_{t,i}}X_{k}Y_{k}$ and $\hat{G}_{t,i}=v\sum_{k\in S_{t,i}}X_{k}X_{k}^{\top}$ concludes the proof.
∎

###### Proposition B.1 .

###### Proof.

### B.2 Derivation of Recursive Latent-Posteriors for Linear Diffusion Models

###### Proposition B.2 .

For any $\ell\in[L]/\{1\}$, the $\ell-1$-th conditional latent-posterior reads $Q_{t,\ell-1}(\cdot\mid\psi_{\ell})=\mathcal{N}(\bar{\mu}_{t,\ell-1},\bar{\Sigma}_{t,\ell-1})$, with

$$ $\displaystyle\bar{\Sigma}_{t,\ell-1}^{-1}=\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1}\,,$ $\displaystyle\bar{\mu}_{t,\ell-1}=\bar{\Sigma}_{t,\ell-1}\big{(}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\psi_{\ell}+\bar{B}_{t,\ell-1}\big{)}\,,$ (17) $$

and the $L$-th latent-posterior reads $Q_{t,L}(\cdot)=\mathcal{N}(\bar{\mu}_{t,L},\bar{\Sigma}_{t,L})$, with

$$ $\displaystyle\bar{\Sigma}_{t,L}^{-1}=\Sigma_{L+1}^{-1}+\bar{G}_{t,L}\,,$ $\displaystyle\bar{\mu}_{t,L}=\bar{\Sigma}_{t,L}\bar{B}_{t,L}\,.$ (18) $$

###### Proof.

Let $\ell\in[L]/\{1\}$. Then, Bayes rule yields that

$$ $\displaystyle Q_{t,\ell-1}(\psi_{\ell-1}\mid\psi_{\ell})$ $\displaystyle\propto\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell-1}=\psi_{\ell-1}\right)\mathcal{N}(\psi_{\ell-1},\mathrm{W}_{\ell}\psi_{\ell},\Sigma_{\ell})\,,$ $$

But from [Lemma B.3](#A2.Thmtheorem3), we know that

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell-1}=\psi_{\ell-1}\right)\propto\exp\Big{[}-\frac{1}{2}\psi_{\ell-1}^{\top}\bar{G}_{t,\ell-1}\psi_{\ell-1}+\psi_{\ell-1}^{\top}\bar{B}_{t,\ell-1}\Big{]}\,.$ $$

Therefore,

$$ $\displaystyle Q_{t,\ell-1}(\psi_{\ell-1}\mid\psi_{\ell})$ $\displaystyle\propto\exp\Big{[}-\frac{1}{2}\psi_{\ell-1}^{\top}\bar{G}_{t,\ell-1}\psi_{\ell-1}+\psi_{\ell-1}^{\top}\bar{B}_{t,\ell-1}\Big{]}\mathcal{N}(\psi_{\ell-1},\mathrm{W}_{\ell}\psi_{\ell},\Sigma_{\ell})\,,$ $\displaystyle\propto\exp\Big{[}-\frac{1}{2}\psi_{\ell-1}^{\top}\bar{G}_{t,\ell-1}\psi_{\ell-1}+\psi_{\ell-1}^{\top}\bar{B}_{t,\ell-1}$ $\displaystyle\qquad\qquad-\frac{1}{2}(\psi_{\ell-1}-\mathrm{W}_{\ell}\psi_{\ell})^{\top}\Sigma_{\ell}^{-1}(\psi_{\ell-1}-\mathrm{W}_{\ell}\psi_{\ell}))\Big{]}\,,$ $\displaystyle\stackrel{{\scriptstyle(i)}}{{\propto}}\exp\Big{[}-\frac{1}{2}\psi_{\ell-1}^{\top}(\bar{G}_{t,\ell-1}+\Sigma_{\ell}^{-1})\psi_{\ell-1}+\psi_{\ell-1}^{\top}(\bar{B}_{t,\ell-1}+\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\psi_{\ell})\Big{]}\,,$ $\displaystyle\stackrel{{\scriptstyle(ii)}}{{\propto}}\mathcal{N}(\psi_{\ell-1};\bar{\mu}_{t,\ell-1},\bar{\Sigma}_{t,\ell-1})\,,$ $$

with $\bar{\Sigma}_{t,\ell-1}^{-1}=\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1}$ and $\bar{\mu}_{t,\ell-1}=\bar{\Sigma}_{t,\ell-1}\big{(}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\psi_{\ell}+\bar{B}_{t,\ell-1}\big{)}$. In $(i)$, we omit terms that are constant in $\psi_{\ell-1}$. In $(ii)$, we complete the square. This concludes the proof for $\ell\in[L]/\{1\}$. For $Q_{t,L}$, we use Bayes rule to get

$$ $\displaystyle Q_{t,L}(\psi_{L})$ $\displaystyle\propto\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,L}=\psi_{L}\right)\mathcal{N}(\psi_{L},0,\Sigma_{L+1})\,.$ $$

Then from [Lemma B.3](#A2.Thmtheorem3), we know that

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,L}=\psi_{L}\right)\propto\exp\Big{[}-\frac{1}{2}\psi_{L}^{\top}\bar{G}_{t,L}\psi_{L}+\psi_{L}^{\top}\bar{B}_{t,L}\Big{]}\,,$ $$

We then use the same derivations above to compute the product $\exp\Big{[}-\frac{1}{2}\psi_{L}^{\top}\bar{G}_{t,L}\psi_{L}+\psi_{L}^{\top}\bar{B}_{t,L}\Big{]}\times\mathcal{N}(\psi_{L},0,\Sigma_{L+1})$, which concludes the proof.
∎

###### Lemma B.3 .

The following holds for any $t\in[n]$ and $\ell\in[L]$,

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)\propto\exp\Big{[}-\frac{1}{2}\psi_{\ell}^{\top}\bar{G}_{t,\ell}\psi_{\ell}+\psi_{\ell}^{\top}\bar{B}_{t,\ell}\Big{]}\,,$ $$

where $\bar{G}_{t,\ell}$ and $\bar{B}_{t,\ell}$ are defined by recursion in [Section 3.1](#S3.SS1).

###### Proof.

We prove this result by induction. To reduce clutter, we let $v=\sigma^{-2}$, and $\Lambda_{1}=\Sigma_{1}^{-1}$. We start with the base case of the induction when $\ell=1$.

(I) Base case. Here we want to show that $\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,1}=\psi_{1}\right)\propto\exp\Big{[}-\frac{1}{2}\psi_{1}^{\top}\bar{G}_{t,1}\psi_{1}+\psi_{1}^{\top}\bar{B}_{t,1}\big{)}\Big{]}$, where $\bar{G}_{t,1}$ and $\bar{B}_{t,1}$ are given in ([3.1.1](#S3.Ex14)). First, we have that

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,1}=\psi_{1}\right)$ $\displaystyle\stackrel{{\scriptstyle(i)}}{{=}}\prod_{i\in[K]}\mathbb{P}\left(H_{t,i}\,\middle|\,\psi_{*,1}=\psi_{1}\right)=\prod_{i\in[K]}\int_{\theta}\mathbb{P}\left(H_{t,i},\theta_{*,i}=\theta\,\middle|\,\psi_{*,1}=\psi_{1}\right)\mathop{}\!\mathrm{d}\theta\,,$ $\displaystyle=\prod_{i\in[K]}\int_{\theta}\mathbb{P}\left(H_{t,i}\,\middle|\,\theta_{*,i}=\theta\right)\mathcal{N}\left(\theta;\mathrm{W}_{1}\psi_{1},\Sigma_{1}\right)\mathop{}\!\mathrm{d}\theta\,,$ $\displaystyle=\prod_{i\in[K]}\underbrace{\int_{\theta}\Big{(}\prod_{k\in S_{t,i}}\mathcal{N}(Y_{k};X_{k}^{\top}\theta,\sigma^{2})\Big{)}\mathcal{N}\left(\theta;\mathrm{W}_{1}\psi_{1},\Sigma_{1}\right)\mathop{}\!\mathrm{d}\theta}_{h_{i}(\psi_{1})}\,,$ $\displaystyle=\prod_{i\in[K]}h_{i}(\psi_{1})\,,$ (19) $$

where $(i)$ follows from the fact that $\theta_{*,i}$ for $i\in[K]$ are conditionally independent given $\psi_{*,1}=\psi_{1}$ and that given $\theta_{*,i}$, $H_{t,i}$ is independent of $\psi_{*,1}$. Now we compute $h_{i}(\psi_{1})=\int_{\theta}\left(\prod_{k\in S_{t,i}}\mathcal{N}(Y_{k};X_{k}^{\top}\theta,\sigma^{2})\right)\mathcal{N}\left(\theta;\mathrm{W}_{1}\psi_{1},\Sigma_{1}\right)\mathop{}\!\mathrm{d}\theta$ as

$$ $\displaystyle h_{i}(\psi_{1})=\int_{\theta}\Big{(}\prod_{k\in S_{t,i}}\mathcal{N}(Y_{k};X_{k}^{\top}\theta,\sigma^{2})\Big{)}\mathcal{N}(\theta;\mathrm{W}_{1}\psi_{1},\Sigma_{1})\mathop{}\!\mathrm{d}\theta\,,$ $\displaystyle\propto\int_{\theta}\exp\Big{[}-\frac{1}{2}v\sum_{k\in S_{t,i}}(Y_{k}-X_{k}^{\top}\theta)^{2}-\frac{1}{2}(\theta-\mathrm{W}_{1}\psi_{1})^{\top}\Lambda_{1}(\theta-\mathrm{W}_{1}\psi_{1})\Big{]}\mathop{}\!\mathrm{d}\theta\,,$ $\displaystyle=\int_{\theta}\exp\Big{[}-\frac{1}{2}\Big{(}v\sum_{k\in S_{t,i}}(Y_{k}^{2}-2Y_{k}\theta^{\top}X_{k}+(\theta^{\top}X_{k})^{2})+\theta^{\top}\Lambda_{1}\theta-2\theta^{\top}\Lambda_{1}\mathrm{W}_{1}\psi_{1}$ $\displaystyle\qquad\qquad+(\mathrm{W}_{1}\psi_{1})^{\top}\Lambda_{1}(\mathrm{W}_{1}\psi_{1})\Big{)}\Big{]}\mathop{}\!\mathrm{d}\theta\,,$ $\displaystyle\propto\int_{\theta}\exp\Big{[}-\frac{1}{2}\Big{(}\theta^{\top}\Big{(}v\sum_{k\in S_{t,i}}X_{k}X_{k}^{\top}+\Lambda_{1}\Big{)}\theta-2\theta^{\top}\Big{(}v\sum_{k\in S_{t,i}}Y_{k}X_{k}$ $\displaystyle\qquad\qquad+\Lambda_{1}\mathrm{W}_{1}\psi_{1}\Big{)}+(\mathrm{W}_{1}\psi_{1})^{\top}\Lambda_{1}(\mathrm{W}_{1}\psi_{1})\Big{)}\Big{]}\mathop{}\!\mathrm{d}\theta\,.$ $$

But we know that $\hat{G}_{t,i}=v\sum_{k\in S_{t,i}}X_{k}X_{k}^{\top}$, and $\hat{B}_{t,i}=v\sum_{k\in S_{t,i}}Y_{k}X_{k}$. To further simplify expressions, we also let

$$ $\displaystyle V=\big{(}\hat{G}_{t,i}+\Lambda_{1}\big{)}^{-1}\,,\quad U=V^{-1}\,,\quad\beta=V\big{(}\hat{B}_{t,i}+\Lambda_{1}\mathrm{W}_{1}\psi_{1}\big{)}\,.$ $$

We have that $UV=VU=I_{d}\,,$ and thus

$$ $\displaystyle h_{i}(\psi_{1})$ $\displaystyle\propto\int_{\theta}\exp\left[-\frac{1}{2}\left(\theta^{\top}U\theta-2\theta^{\top}UV\left(\hat{B}_{t,i}+\Lambda_{1}\mathrm{W}_{1}\psi_{1}\right)+(\mathrm{W}_{1}\psi_{1})^{\top}\Lambda_{1}(\mathrm{W}_{1}\psi_{1})\right)\right]\mathop{}\!\mathrm{d}\theta\,,$ $\displaystyle=\int_{\theta}\exp\left[-\frac{1}{2}\left(\theta^{\top}U\theta-2\theta^{\top}U\beta+(\mathrm{W}_{1}\psi_{1})^{\top}\Lambda_{1}(\mathrm{W}_{1}\psi_{1})\right)\right]\mathop{}\!\mathrm{d}\theta\,,$ $\displaystyle=\int_{\theta}\exp\left[-\frac{1}{2}\left((\theta-\beta)^{\top}U(\theta-\beta)-\beta^{\top}U\beta+(\mathrm{W}_{1}\psi_{1})^{\top}\Lambda_{1}(\mathrm{W}_{1}\psi_{1})\right)\right]\mathop{}\!\mathrm{d}\theta\,,$ $\displaystyle\propto\exp\left[-\frac{1}{2}\left(-\beta^{\top}U\beta+(\mathrm{W}_{1}\psi_{1})^{\top}\Lambda_{1}(\mathrm{W}_{1}\psi_{1})\right)\right]\,,$ $\displaystyle=\exp\left[-\frac{1}{2}\left(-\left(\hat{B}_{t,i}+\Lambda_{1}\mathrm{W}_{1}\psi_{1}\right)^{\top}V\left(\hat{B}_{t,i}+\Lambda_{1}\mathrm{W}_{1}\psi_{1}\right)+(\mathrm{W}_{1}\psi_{1})^{\top}\Lambda_{1}(\mathrm{W}_{1}\psi_{1})\right)\right]\,,$ $\displaystyle\propto\exp\left[-\frac{1}{2}\left(\psi_{1}^{\top}\mathrm{W}_{1}^{\top}\left(\Lambda_{1}-\Lambda_{1}V\Lambda_{1}\right)\mathrm{W}_{1}\psi_{1}-2\psi_{1}^{\top}\left(\mathrm{W}_{1}^{\top}\Lambda_{1}V\hat{B}_{t,i}\right)\right)\right]\,,$ $\displaystyle=\exp\left[-\frac{1}{2}\psi_{1}^{\top}\Omega_{i}\psi_{1}+\psi_{1}^{\top}m_{i}\right]\,,$ $$

where

$$ $\displaystyle\Omega_{i}$ $\displaystyle=\mathrm{W}_{1}^{\top}\left(\Lambda_{1}-\Lambda_{1}V\Lambda_{1}\right)\mathrm{W}_{1}=\mathrm{W}_{1}^{\top}\left(\Lambda_{1}-\Lambda_{1}(\hat{G}_{t,i}+\Lambda_{1})^{-1}\Lambda_{1}\right)\mathrm{W}_{1}\,,$ $\displaystyle m_{i}$ $\displaystyle=\mathrm{W}_{1}^{\top}\Lambda_{1}V\hat{B}_{t,i}=\mathrm{W}_{1}^{\top}\Lambda_{1}(\hat{G}_{t,i}+\Lambda_{1})^{-1}\hat{B}_{t,i}\,.$ (20) $$

But notice that $V=(\hat{G}_{t,i}+\Lambda_{1})^{-1}=\hat{\Sigma}_{t,i}$ and thus

$$ $\displaystyle\Omega_{i}=\mathrm{W}_{1}^{\top}\big{(}\Lambda_{1}-\Lambda_{1}\hat{\Sigma}_{t,i}\Lambda_{1}\big{)}\mathrm{W}_{1}\,,$ $\displaystyle m_{i}$ $\displaystyle=\mathrm{W}_{1}^{\top}\Lambda_{1}\hat{\Sigma}_{t,i}\hat{B}_{t,i}\,.$ (21) $$

Finally, we plug this result in ([B.2](#A2.Ex53)) to get

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,1}=\psi_{1}\right)=\prod_{i\in[K]}h_{i}(\psi_{1})$ $\displaystyle\propto\prod_{i\in[K]}\exp\left[-\frac{1}{2}\psi_{1}^{\top}\Omega_{i}\psi_{1}+\psi_{1}^{\top}m_{i}\right]\,,$ $\displaystyle=\exp\left[-\frac{1}{2}\psi_{1}^{\top}\sum_{i\in[K]}\Omega_{i}\psi_{1}+\psi_{1}^{\top}\sum_{i\in[K]}m_{i}\right]\,,$ $\displaystyle=\exp\left[-\frac{1}{2}\psi_{1}^{\top}\bar{G}_{t,1}\psi_{1}+\psi_{1}^{\top}\bar{B}_{t,1}\right]\,,$ $$

where

$$ $\displaystyle\bar{G}_{t,1}$ $\displaystyle=\sum_{i=1}^{K}\Omega_{i}=\sum_{i=1}^{K}\mathrm{W}_{1}^{\top}\big{(}\Lambda_{1}-\Lambda_{1}\hat{\Sigma}_{t,i}\Lambda_{1}\big{)}\mathrm{W}_{1}=\mathrm{W}_{1}^{\top}\sum_{i=1}^{K}\big{(}\Sigma_{1}^{-1}-\Sigma_{1}^{-1}\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\big{)}\mathrm{W}_{1}\,,$ $\displaystyle\bar{B}_{t,1}$ $\displaystyle=\sum_{i=1}^{K}m_{i}=\sum_{i=1}^{K}\hat{\Sigma}_{t,i}\hat{B}_{t,i}=\mathrm{W}_{1}^{\top}\Sigma_{1}^{-1}\sum_{i=1}^{K}\hat{\Sigma}_{t,i}\hat{B}_{t,i}\,.$ $$

This concludes the proof of the base case.

(II) Induction step. Let $\ell\in[L]/\{1\}$. Suppose that

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell-1}=\psi_{\ell-1}\right)\propto\exp\left[-\frac{1}{2}\psi_{\ell-1}^{\top}\bar{G}_{t,\ell-1}\psi_{\ell-1}+\psi_{\ell-1}^{\top}\bar{B}_{t,\ell-1}\right]\,.$ (22) $$

Then we want to show that

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)\propto\exp\left[-\frac{1}{2}\psi_{\ell}^{\top}\bar{G}_{t,\ell}\psi_{\ell}+\psi_{\ell}^{\top}\bar{B}_{t,\ell}\right]\,,$ $$

where

$$ $\displaystyle\bar{G}_{t,\ell}$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}=\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}(\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1})^{-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}\,,$ $\displaystyle\bar{B}_{t,\ell}$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\bar{B}_{t,\ell-1}=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}(\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1})^{-1}\bar{B}_{t,\ell-1}\,.$ $$

To achieve this, we start by expressing $\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)$ in terms of $\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell-1}=\psi_{\ell-1}\right)$ as

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)=\int_{\psi_{\ell-1}}\mathbb{P}\left(H_{t},\psi_{*,\ell-1}=\psi_{\ell-1}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)\mathop{}\!\mathrm{d}\psi_{\ell-1}\,,$ $\displaystyle=\int_{\psi_{\ell-1}}\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell-1}=\psi_{\ell-1},\psi_{*,\ell}=\psi_{\ell}\right)\mathcal{N}(\psi_{\ell-1};\mathrm{W}_{\ell}\psi_{\ell},\Sigma_{\ell})\mathop{}\!\mathrm{d}\psi_{\ell-1}\,,$ $\displaystyle=\int_{\psi_{\ell-1}}\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell-1}=\psi_{\ell-1}\right)\mathcal{N}(\psi_{\ell-1};\mathrm{W}_{\ell}\psi_{\ell},\Sigma_{\ell})\mathop{}\!\mathrm{d}\psi_{\ell-1}\,,$ $\displaystyle\propto\int_{\psi_{\ell-1}}\exp\Big{[}-\frac{1}{2}\psi_{\ell-1}^{\top}\bar{G}_{t,\ell-1}\psi_{\ell-1}+\psi_{\ell-1}^{\top}\bar{B}_{t,\ell-1}\Big{]}\mathcal{N}(\psi_{\ell-1};\mathrm{W}_{\ell}\psi_{\ell},\Sigma_{\ell})\mathop{}\!\mathrm{d}\psi_{\ell-1}\,,$ $\displaystyle\propto\int_{\psi_{\ell-1}}\exp\Big{[}-\frac{1}{2}\psi_{\ell-1}^{\top}\bar{G}_{t,\ell-1}\psi_{\ell-1}+\psi_{\ell-1}^{\top}\bar{B}_{t,\ell-1}$ $\displaystyle\qquad\qquad+(\psi_{\ell-1}-\mathrm{W}_{\ell}\psi_{\ell})^{\top}\Lambda_{\ell}(\psi_{\ell-1}-\mathrm{W}_{\ell}\psi_{\ell})\Big{)}\Big{]}\mathop{}\!\mathrm{d}\psi_{\ell-1}\,.$ $$

Now let $S=\bar{G}_{t,\ell-1}+\Lambda_{\ell}$ and $V=\bar{B}_{t,\ell-1}+\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}$. Then we have that,

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)$ $\displaystyle\propto\int_{\psi_{\ell-1}}\exp\Big{[}-\frac{1}{2}\psi_{\ell-1}^{\top}\bar{G}_{t,\ell-1}\psi_{\ell-1}+\psi_{\ell-1}^{\top}\bar{B}_{t,\ell-1}$ $\displaystyle\qquad\qquad+(\psi_{\ell-1}-\mathrm{W}_{\ell}\psi_{\ell})^{\top}\Lambda_{\ell}(\psi_{\ell-1}-\mathrm{W}_{\ell}\psi_{\ell})\Big{)}\Big{]}\mathop{}\!\mathrm{d}\psi_{\ell-1}\,,$ $\displaystyle\propto\int_{\psi_{\ell-1}}\exp\Big{[}-\frac{1}{2}\Big{(}\psi_{\ell-1}^{\top}S\psi_{\ell-1}-2\psi_{\ell-1}^{\top}\left(\bar{B}_{t,\ell-1}+\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}\right)+\psi_{\ell}^{\top}\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}\Big{)}\Big{]}\mathop{}\!\mathrm{d}\psi_{\ell-1}\,,$ $\displaystyle=\int_{\psi_{\ell-1}}\exp\Big{[}-\frac{1}{2}\Big{(}\psi_{\ell-1}^{\top}S(\psi_{\ell-1}-2S^{-1}V)+\psi_{\ell}^{\top}\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}\Big{)}\Big{]}\mathop{}\!\mathrm{d}\psi_{\ell-1}\,,$ $\displaystyle=\int_{\psi_{\ell-1}}\exp\Big{[}-\frac{1}{2}\Big{(}(\psi_{\ell-1}-S^{-1}V)^{\top}S(\psi_{\ell-1}-S^{-1}V)$ $\displaystyle\qquad\qquad+\psi_{\ell}^{\top}\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}-V^{\top}S^{-1}V\Big{)}\Big{]}\mathop{}\!\mathrm{d}\psi_{\ell-1}.$ $$

In the second step, we omit constants in $\psi_{\ell}$ and $\psi_{\ell-1}$. Thus

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)$ $\displaystyle\propto\int_{\psi_{\ell-1}}\exp\left[-\frac{1}{2}\left((\psi_{\ell-1}-S^{-1}V)^{\top}S(\psi_{\ell-1}-S^{-1}V)+\psi_{\ell}^{\top}\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}-V^{\top}S^{-1}V\right)\right]\mathop{}\!\mathrm{d}\psi_{\ell-1},$ $\displaystyle\propto\exp\left[-\frac{1}{2}\left(\psi_{\ell}^{\top}\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}-V^{\top}S^{-1}V\right)\right]\,.$ $$

It follows that

$$ $\displaystyle\mathbb{P}\left(H_{t}\,\middle|\,\psi_{*,\ell}=\psi_{\ell}\right)$ $\displaystyle\propto\exp\left[-\frac{1}{2}\left(\psi_{\ell}^{\top}\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}-V^{\top}S^{-1}V\right)\right]\,,$ $\displaystyle=\exp\left[-\frac{1}{2}\left(\psi_{\ell}^{\top}\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}-\left(\bar{B}_{t,\ell-1}+\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}\right)^{\top}S^{-1}\left(\bar{B}_{t,\ell-1}+\Lambda_{\ell}\mathrm{W}_{\ell}\psi_{\ell}\right)\right)\right]$ $\displaystyle\propto\exp\left[-\frac{1}{2}\left(\psi_{\ell}^{\top}\left(\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}\mathrm{W}_{\ell}-\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}S^{-1}\Lambda_{\ell}\mathrm{W}_{\ell}\right)\psi_{\ell}-2\psi_{\ell}^{\top}\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}S^{-1}\bar{B}_{t,\ell-1}\right)\right]\,,$ $\displaystyle=\exp\left[-\frac{1}{2}\psi_{\ell}^{\top}\bar{G}_{t,\ell}\psi_{\ell}+\psi_{\ell}^{\top}\bar{B}_{t,\ell}\right]\,.$ $$

In the last step, we omit constants in $\psi_{\ell}$ and we set

$$ $\displaystyle\bar{G}_{t,\ell}$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\left(\Lambda_{\ell}-\Lambda_{\ell}S^{-1}\Lambda_{\ell}\right)\mathrm{W}_{\ell}=\mathrm{W}_{\ell}^{\top}\left(\Lambda_{\ell}-\Lambda_{\ell}(\Lambda_{\ell}+\bar{G}_{t,\ell-1})^{-1}\Sigma_{\ell}^{-1}\Lambda_{\ell}\right)\mathrm{W}_{\ell}\,,$ $\displaystyle\bar{B}_{t,\ell}$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}S^{-1}\bar{B}_{t,\ell-1}=\mathrm{W}_{\ell}^{\top}\Lambda_{\ell}(\Lambda_{\ell}+\bar{G}_{t,\ell-1})^{-1}\bar{B}_{t,\ell-1}\,.$ $$

This completes the proof.
∎

###### Proposition B.2 .

###### Proof.

###### Lemma B.3 .

###### Proof.

## Appendix C Regret Proof and Discussion

### C.1 Sketch of the Proof

We start with the following standard lemma upon which we build our analysis (Hong et al., ; Aouali et al., ).

###### Lemma C.1 .

Assume that $\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,H_{t}\right)=\mathcal{N}(\theta;\check{\mu}_{t,i},\check{\Sigma}_{t,i})$ for any $i\in[K]$, then for any $\delta\in(0,1)$,

$$ $\textstyle\mathcal{B}\mathcal{R}(n)$ $\textstyle\leq\sqrt{2n\log(1/\delta)}\sqrt{\mathbb{E}\left[\sum_{t=1}^{n}\|X_{t}\|_{\check{\Sigma}_{t,A_{t}}}^{2}\right]}+cn\delta\,,\qquad\text{where $c>0$ is a constant}\,.$ (23) $$

Applying [Lemma C.1](#A3.Thmtheorem1) requires proving that the *marginal* action-posteriors $\mathbb{P}\left(\theta_{*,i}=\theta\,\middle|\,H_{t}\right)$ in ([3](#S3.Ex4)) are Gaussian and computing their covariances, while we only know the *conditional* action-posteriors $P_{t,i}$ and latent-posteriors $Q_{t,\ell}$. This is achieved by leveraging the preservation properties of the family of Gaussian distributions (Koller & Friedman, ) and the total covariance decomposition (Weiss, ) which leads to the next lemma.

###### Lemma C.2 .

Let $t\in[n]$ and $i\in[K]$, then the marginal covariance matrix $\check{\Sigma}_{t,i}$ reads

$$ $\textstyle\check{\Sigma}_{t,i}=\hat{\Sigma}_{t,i}+\sum_{\ell\in[L]}\mathrm{P}_{i,\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{i,\ell}^{\top}\,,$ where $\,\mathrm{P}_{i,\ell}=\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\prod_{k=1}^{\ell-1}\bar{\Sigma}_{t,k}\Sigma_{k+1}^{-1}\mathrm{W}_{k+1}$ . (24) $$

The marginal covariance matrix $\check{\Sigma}_{t,i}$ in ([24](#A3.E24)) decomposes into $L+1$ terms. The first term corresponds to the posterior uncertainty of $\theta_{*,i}\mid\psi_{*,1}$. The remaining $L$ terms capture the posterior uncertainties of $\psi_{*,L}$ and $\psi_{*,\ell-1}\mid\psi_{*,\ell}$ for $\ell\in[L]/\{1\}$. These are then used to quantify the posterior information gain of latent parameters after one round as follows.

###### Lemma C.3 (Posterior information gain) .

Let $t\in[n]$ and $\ell\in[L]$, then

$$ $\textstyle\bar{\Sigma}_{t+1,\ell}^{-1}-\bar{\Sigma}_{t,\ell}^{-1}\succeq\sigma^{-2}\sigma_{\textsc{max}}^{-2\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\,,$ $\textstyle\text{where }\,\sigma_{\textsc{max}}^{2}=\max_{\ell\in[L+1]}1+\frac{\sigma_{\ell}^{2}}{\sigma^{2}}\,.$ (25) $$

Finally, [Lemma C.2](#A3.Thmtheorem2) is used to decompose $\|X_{t}\|_{\check{\Sigma}_{t,A_{t}}}^{2}$ in [Equation 23](#A3.E23) into $L+1$ terms. Each term is bounded thanks to [Lemma C.3](#A3.Thmtheorem3). This results in the Bayes regret bound in [Theorem 4.1](#S4.Thmtheorem1).

###### Lemma C.1 .

###### Lemma C.2 .

###### Lemma C.3 (Posterior information gain) .

### C.2 Technical Contributions

Our main technical contributions are the following.

[Lemma C.2](#A3.Thmtheorem2). In dTS, sampling is done hierarchically, meaning the marginal posterior distribution of $\theta_{*,i}|H_{t}$ is not explicitly defined. Instead, we use the conditional posterior distribution of $\theta_{*,i}|H_{t},\psi_{*,1}$. The first contribution was deriving $\theta_{*,i}|H_{t}$ using the total covariance decomposition combined with an induction proof, as our posteriors in [Section 3.1](#S3.SS1) were derived recursively. Unlike in Bayes regret analysis for standard Thompson sampling, where the posterior distribution of $\theta_{*,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 single-level hierarchy, where $\theta_{*,i}|\psi_{*,1}\sim\mathcal{N}(\psi_{*,1},\Sigma_{1})$ with $L=1$. In contrast, our model proposes a multi-level hierarchy, where the first level is $\theta_{*,i}|\psi_{*,1}\sim\mathcal{N}(W_{1}\psi_{*,1},\Sigma_{1})$. This also introduces a new aspect to our approach – the use of a linear function $W_{1}\psi_{*,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_{\ell}\psi_{*,\ell}$, instead of a single-level hierarchy under a single identity function $\psi_{*,1}$.

[Lemma C.3](#A3.Thmtheorem3). 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\in[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](#S3.SS1) that connect the posterior covariance of each latent parameter $\psi_{*,\ell}$ with the covariance of the posterior action parameters $\theta_{*,i}$. This allows us to propagate the information gain associated with the action taken in round $A_{t}$ to all latent parameters $\psi_{*,\ell},$ for $\ell\in[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](#S4.Thmtheorem2). Building upon the insights of [Theorem 4.1](#S4.Thmtheorem1), we introduce the sparsity assumption (A3). Under this assumption, we demonstrate that the Bayes regret outlined in [Theorem 4.1](#S4.Thmtheorem1) can be significantly refined. Specifically, the regret becomes contingent on dimensions $d_{\ell}\leq 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](#S4.Thmtheorem1) and modify a crucial equality. This adjustment results in a more precise representation by partitioning the covariance matrix of each latent parameter $\psi_{*,\ell}$ into blocks. These blocks comprise a $d_{\ell}\times d_{\ell}$ segment corresponding to the learnable $d_{\ell}$ parameters of $\psi_{*,\ell}$, and another block of size $(d-d_{\ell})\times(d-d_{\ell})$ that does not necessitate learning. This decomposition allows us to conclude that the final regret is solely dependent on $d_{\ell}$, marking a significant refinement from the original theorem.

### C.3 Proof of Lemma C.2

In this proof, we heavily rely on the total covariance decomposition (Weiss, ). Also, refer to (Hong et al., , Section 5.2) for a brief introduction to this decomposition. Now, from ([5](#S3.E5)), we have that

$$ $\displaystyle\mathrm{cov}\left[\theta_{*,i}\,\middle|\,H_{t},\psi_{*,1}\right]$ $\displaystyle=\hat{\Sigma}_{t,i}=\left(\hat{G}_{t,i}+\Sigma_{1}^{-1}\right)^{-1}\,,$ $\displaystyle\mathbb{E}\left[\theta_{*,i}\,\middle|\,H_{t},\psi_{*,1}\right]$ $\displaystyle=\hat{\mu}_{t,i}=\hat{\Sigma}_{t,i}\left(\hat{B}_{t,i}+\Sigma_{1}^{-1}\mathrm{W}_{1}\psi_{*,1}\right)\,.$ $$

First, given $H_{t}$, $\mathrm{cov}\left[\theta_{*,i}\,\middle|\,H_{t},\psi_{*,1}\right]=\left(\hat{G}_{t,i}+\Sigma_{1}^{-1}\right)^{-1}$ is constant. Thus

$$ $\displaystyle\mathbb{E}\left[\mathrm{cov}\left[\theta_{*,i}\,\middle|\,H_{t},\psi_{*,1}\right]\,\middle|\,H_{t}\right]=\mathrm{cov}\left[\theta_{*,i}\,\middle|\,H_{t},\psi_{*,1}\right]=\left(\hat{G}_{t,i}+\Sigma_{1}^{-1}\right)^{-1}=\hat{\Sigma}_{t,i}\,.$ $$

In addition, given $H_{t}$, both $\hat{\Sigma}_{t,i}$ and $\hat{B}_{t,i}$ are constant. Thus

$$ $\displaystyle\mathrm{cov}\left[\mathbb{E}\left[\theta_{*,i}\,\middle|\,H_{t},\psi_{*,1}\right]\,\middle|\,H_{t}\right]$ $\displaystyle=\mathrm{cov}\left[\hat{\Sigma}_{t,i}\left(\hat{B}_{t,i}+\Sigma_{1}^{-1}\mathrm{W}_{1}\psi_{*,1}\right)\,\middle|\,H_{t}\right]\,,$ $\displaystyle=\mathrm{cov}\left[\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\psi_{*,1}\,\middle|\,H_{t}\right]\,,$ $\displaystyle=\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t}\right]\mathrm{W}_{1}^{\top}\Sigma_{1}^{-1}\hat{\Sigma}_{t,i}\,,$ $\displaystyle=\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\bar{\bar{\Sigma}}_{t,1}\mathrm{W}_{1}^{\top}\Sigma_{1}^{-1}\hat{\Sigma}_{t,i}\,,$ $$

where $\bar{\bar{\Sigma}}_{t,1}=\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t}\right]$ is the marginal posterior covariance of $\psi_{*,1}$. Finally, the total covariance decomposition (Weiss, ; Hong et al., ) yields that

$$ $\displaystyle\check{\Sigma}_{t,i}=\mathrm{cov}\left[\theta_{*,i}\,\middle|\,H_{t}\right]$ $\displaystyle=\mathbb{E}\left[\mathrm{cov}\left[\theta_{*,i}\,\middle|\,H_{t},\psi_{*,1}\right]\,\middle|\,H_{t}\right]+\mathrm{cov}\left[\mathbb{E}\left[\theta_{*,i}\,\middle|\,H_{t},\psi_{*,1}\right]\,\middle|\,H_{t}\right]\,,$ $\displaystyle=\hat{\Sigma}_{t,i}+\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\bar{\bar{\Sigma}}_{t,1}\mathrm{W}_{1}^{\top}\Sigma_{1}^{-1}\hat{\Sigma}_{t,i}\,,$ (26) $$

However, $\bar{\bar{\Sigma}}_{t,1}=\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t}\right]$ is different from $\bar{\Sigma}_{t,1}=\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]$ that we already derived in ([3.1.1](#S3.Ex13)). Thus we do not know the expression of $\bar{\bar{\Sigma}}_{t,1}$. But we can use the same total covariance decomposition trick to find it. Precisely, let $\bar{\bar{\Sigma}}_{t,\ell}=\mathrm{cov}\left[\psi_{*,\ell}\,\middle|\,H_{t}\right]$ for any $\ell\in[L]$. Then we have that

$$ $\displaystyle\bar{\Sigma}_{t,1}=\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]$ $\displaystyle=\big{(}\Sigma_{2}^{-1}+\bar{G}_{t,1}\big{)}^{-1}\,,$ $\displaystyle\bar{\mu}_{t,1}=\mathbb{E}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]$ $\displaystyle=\bar{\Sigma}_{t,1}\Big{(}\Sigma_{2}^{-1}\mathrm{W}_{2}\psi_{*,2}+\bar{B}_{t,1}\Big{)}\,.$ $$

First, given $H_{t}$, $\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]=\big{(}\Sigma_{2}^{-1}+\bar{G}_{t,1}\big{)}^{-1}$ is constant. Thus

$$ $\displaystyle\mathbb{E}\left[\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]\,\middle|\,H_{t}\right]=\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]=\bar{\Sigma}_{t,1}\,.$ $$

In addition, given $H_{t}$, $\bar{\Sigma}_{t,1}$, $\tilde{\Sigma}_{t,1}$ and $\bar{B}_{t,1}$ are constant. Thus

$$ $\displaystyle\mathrm{cov}\left[\mathbb{E}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]\,\middle|\,H_{t}\right]$ $\displaystyle=\mathrm{cov}\left[\bar{\Sigma}_{t,1}\Big{(}\Sigma_{2}^{-1}\mathrm{W}_{2}\psi_{*,2}+\bar{B}_{t,1}\Big{)}\,\middle|\,H_{t}\right]\,,$ $\displaystyle=\mathrm{cov}\left[\bar{\Sigma}_{t,1}\Sigma_{2}^{-1}\mathrm{W}_{2}\psi_{*,2}\,\middle|\,H_{t}\right]\,,$ $\displaystyle=\bar{\Sigma}_{t,1}\Sigma_{2}^{-1}\mathrm{W}_{2}\mathrm{cov}\left[\psi_{*,2}\,\middle|\,H_{t}\right]\mathrm{W}_{2}^{\top}\Sigma_{2}^{-1}\bar{\Sigma}_{t,1}\,,$ $\displaystyle=\bar{\Sigma}_{t,1}\Sigma_{2}^{-1}\mathrm{W}_{2}\bar{\bar{\Sigma}}_{t,2}\mathrm{W}_{2}^{\top}\Sigma_{2}^{-1}\bar{\Sigma}_{t,1}\,.$ $$

Finally, total covariance decomposition (Weiss, ; Hong et al., ) leads to

$$ $\displaystyle\bar{\bar{\Sigma}}_{t,1}=\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t}\right]$ $\displaystyle=\mathbb{E}\left[\mathrm{cov}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]\,\middle|\,H_{t}\right]+\mathrm{cov}\left[\mathbb{E}\left[\psi_{*,1}\,\middle|\,H_{t},\psi_{*,2}\right]\,\middle|\,H_{t}\right]\,,$ $\displaystyle=\bar{\Sigma}_{t,1}+\bar{\Sigma}_{t,1}\Sigma_{2}^{-1}\mathrm{W}_{2}\bar{\bar{\Sigma}}_{t,2}\mathrm{W}_{2}^{\top}\Sigma_{2}^{-1}\bar{\Sigma}_{t,1}\,.$ $$

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

$$ $\displaystyle\bar{\bar{\Sigma}}_{t,\ell}=\bar{\Sigma}_{t,\ell}+\bar{\Sigma}_{t,\ell}\Sigma_{\ell+1}^{-1}\mathrm{W}_{\ell+1}\bar{\bar{\Sigma}}_{t,\ell+1}\mathrm{W}_{\ell+1}^{\top}\Sigma_{\ell+1}^{-1}\bar{\Sigma}_{t,\ell}\,,$ $\displaystyle\forall\ell\in[L-1]\,.$ $$

Then, by induction, we get that

$$ $\displaystyle\bar{\bar{\Sigma}}_{t,1}=\sum_{\ell\in[L]}\bar{\mathrm{P}}_{\ell}\bar{\Sigma}_{t,\ell}\bar{\mathrm{P}}_{\ell}^{\top}\,,$ $\displaystyle\forall\ell\in[L-1]\,,$ $$

where we use that by definition $\bar{\bar{\Sigma}}_{t,L}=\mathrm{cov}\left[\psi_{*,L}\,\middle|\,H_{t}\right]=\bar{\Sigma}_{t,L}$ and set $\bar{\mathrm{P}}_{1}=I_{d}$ and $\bar{\mathrm{P}}_{\ell}=\prod_{k=1}^{\ell-1}\bar{\Sigma}_{t,k}\Sigma_{k+1}^{-1}\mathrm{W}_{k+1}$ for any $\ell\in[L]/\{1\}$. Plugging this in ([C.3](#A3.Ex109)) leads to

$$ $\displaystyle\check{\Sigma}_{t,i}$ $\displaystyle=\hat{\Sigma}_{t,i}+\sum_{\ell\in[L]}\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\bar{\mathrm{P}}_{\ell}\bar{\Sigma}_{t,\ell}\bar{\mathrm{P}}_{\ell}^{\top}\mathrm{W}_{1}^{\top}\Sigma_{1}^{-1}\hat{\Sigma}_{t,i}\,,$ $\displaystyle=\hat{\Sigma}_{t,i}+\sum_{\ell\in[L]}\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\bar{\mathrm{P}}_{\ell}\bar{\Sigma}_{t,\ell}(\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1})^{\top}\,,$ $\displaystyle=\hat{\Sigma}_{t,i}+\sum_{\ell\in[L]}\mathrm{P}_{i,\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{i,\ell}^{\top}\,,$ $$

where $\mathrm{P}_{i,\ell}=\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\bar{\mathrm{P}}_{\ell}=\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\prod_{k=1}^{\ell-1}\bar{\Sigma}_{t,k}\Sigma_{k+1}^{-1}\mathrm{W}_{k+1}$.

### C.4 Proof of Lemma C.3

We prove this result by induction. We start with the base case when $\ell=1$.

(I) Base case. Let $u=\sigma^{-1}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}X_{t}$ From the expression of $\bar{\Sigma}_{t,1}$ in ([3.1.1](#S3.Ex13)), we have that

$$ $\displaystyle\bar{\Sigma}_{t+1,1}^{-1}-\bar{\Sigma}_{t,1}^{-1}$ $\displaystyle=\mathrm{W}_{1}^{\top}\left(\Sigma_{1}^{-1}-\Sigma_{1}^{-1}(\hat{\Sigma}_{t,A_{t}}^{-1}+\sigma^{-2}X_{t}X_{t}^{\top})^{-1}\Sigma_{1}^{-1}-(\Sigma_{1}^{-1}-\Sigma_{1}^{-1}\hat{\Sigma}_{t,A_{t}}\Sigma_{1}^{-1})\right)\mathrm{W}_{1}\,,$ $\displaystyle=\mathrm{W}_{1}^{\top}\left(\Sigma_{1}^{-1}(\hat{\Sigma}_{t,A_{t}}-(\hat{\Sigma}_{t,A_{t}}^{-1}+\sigma^{-2}X_{t}X_{t}^{\top})^{-1})\Sigma_{1}^{-1}\right)\mathrm{W}_{1}\,,$ $\displaystyle=\mathrm{W}_{1}^{\top}\left(\Sigma_{1}^{-1}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}(I_{d}-(I_{d}+\sigma^{-2}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}X_{t}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}})^{-1})\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}\Sigma_{1}^{-1}\right)\mathrm{W}_{1}\,,$ $\displaystyle=\mathrm{W}_{1}^{\top}\left(\Sigma_{1}^{-1}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}(I_{d}-(I_{d}+uu^{\top})^{-1})\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}\Sigma_{1}^{-1}\right)\mathrm{W}_{1}\,,$ $\displaystyle\stackrel{{\scriptstyle(i)}}{{=}}\mathrm{W}_{1}^{\top}\left(\Sigma_{1}^{-1}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}\frac{uu^{\top}}{1+u^{\top}u}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}\Sigma_{1}^{-1}\right)\mathrm{W}_{1}\,,$ $\displaystyle\stackrel{{\scriptstyle(ii)}}{{=}}\sigma^{-2}\mathrm{W}_{1}^{\top}\Sigma_{1}^{-1}\hat{\Sigma}_{t,A_{t}}\frac{X_{t}X_{t}^{\top}}{1+u^{\top}u}\hat{\Sigma}_{t,A_{t}}\Sigma_{1}^{-1}\mathrm{W}_{1}\,.$ (27) $$

In $(i)$ we use the Sherman-Morrison formula. Note that $(ii)$ says that $\bar{\Sigma}_{t+1,1}^{-1}-\bar{\Sigma}_{t,1}^{-1}$ is one-rank which we will also need in induction step. Now, we have that $\|X_{t}\|^{2}=1$. Therefore,

$$ $\displaystyle 1+u^{\top}u=1+\sigma^{-2}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}X_{t}\leq 1+\sigma^{-2}\lambda_{1}(\Sigma_{1})\|X_{t}\|^{2}=1+\sigma^{-2}\sigma_{1}^{2}\leq\sigma_{\textsc{max}}^{2}\,,$ $$

where we use that by definition of $\sigma_{\textsc{max}}^{2}$ in [Lemma C.3](#A3.Thmtheorem3), we have that $\sigma_{\textsc{max}}^{2}\geq 1+\sigma^{-2}\sigma_{1}^{2}$. Therefore, by taking the inverse, we get that $\frac{1}{1+u^{\top}u}\geq\sigma_{\textsc{max}}^{-2}$. Combining this with ([27](#A3.E27)) leads to

$$ $\displaystyle\bar{\Sigma}_{t+1,1}^{-1}-\bar{\Sigma}_{t,1}^{-1}\succeq\sigma^{-2}\sigma_{\textsc{max}}^{-2}\mathrm{W}_{1}^{\top}\Sigma_{1}^{-1}\hat{\Sigma}_{t,A_{t}}X_{t}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}\Sigma_{1}^{-1}\mathrm{W}_{1}$ $$

Noticing that $\mathrm{P}_{A_{t},1}=\hat{\Sigma}_{t,A_{t}}\Sigma_{1}^{-1}\mathrm{W}_{1}$ concludes the proof of the base case when $\ell=1$.

(II) Induction step. Let $\ell\in[L]/\{1\}$ and suppose that $\bar{\Sigma}_{t+1,\ell-1}^{-1}-\bar{\Sigma}_{t,\ell-1}^{-1}$ is one-rank and that it holds for $\ell-1$ that

$$ $\displaystyle\bar{\Sigma}_{t+1,\ell-1}^{-1}-\bar{\Sigma}_{t,\ell-1}^{-1}\succeq\sigma^{-2}\sigma_{\textsc{max}}^{-2(\ell-1)}\mathrm{P}_{A_{t},\ell-1}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell-1}\,,$ $\displaystyle\text{where }\,\sigma_{\textsc{max}}^{-2}=\max_{\ell\in[L]}1+\sigma^{-2}\sigma_{\ell}^{2}.$ $$

Then, we want to show that $\bar{\Sigma}_{t+1,\ell}^{-1}-\bar{\Sigma}_{t,\ell}^{-1}$ is also one-rank and that it holds that

$$ $\displaystyle\bar{\Sigma}_{t+1,\ell}^{-1}-\bar{\Sigma}_{t,\ell}^{-1}\succeq\sigma^{-2}\sigma_{\textsc{max}}^{-2\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\,,$ $\displaystyle\text{where }\,\sigma_{\textsc{max}}^{-2}=\max_{\ell\in[L]}1+\sigma^{-2}\sigma_{\ell}^{2}.$ $$

This is achieved as follows. First, we notice that by the induction hypothesis, we have that $\tilde{\Sigma}_{t+1,\ell-1}^{-1}-\bar{G}_{t,\ell-1}=\bar{\Sigma}_{t+1,\ell-1}^{-1}-\bar{\Sigma}_{t,\ell-1}^{-1}$ is one-rank. In addition, the matrix is positive semi-definite. Thus we can write it as $\tilde{\Sigma}_{t+1,\ell-1}^{-1}-\bar{G}_{t,\ell-1}=uu^{\top}$ where $u\in^{d}$. Then, similarly to the base case, we have

$$ $\displaystyle\bar{\Sigma}_{t+1,\ell}^{-1}-$ $\displaystyle\bar{\Sigma}_{t,\ell}^{-1}=\tilde{\Sigma}_{t+1,\ell}^{-1}-\tilde{\Sigma}_{t,\ell}^{-1}\,,$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}+\tilde{\Sigma}_{t+1,\ell-1}\big{)}^{-1}\mathrm{W}_{\ell}-\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}+\tilde{\Sigma}_{t,\ell-1}\big{)}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Big{[}\big{(}\Sigma_{\ell}+\tilde{\Sigma}_{t+1,\ell-1}\big{)}^{-1}-\big{(}\Sigma_{\ell}+\tilde{\Sigma}_{t,\ell-1}\big{)}^{-1}\Big{]}\mathrm{W}_{\ell}\,,$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\Big{[}\big{(}\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1}\big{)}^{-1}-\big{(}\Sigma_{\ell}^{-1}+\tilde{\Sigma}_{t+1,\ell-1}^{-1}\big{)}^{-1}\Big{]}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\Big{[}\big{(}\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1}\big{)}^{-1}-\big{(}\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1}+\tilde{\Sigma}_{t+1,\ell-1}^{-1}-\bar{G}_{t,\ell-1}\big{)}^{-1}\Big{]}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\Big{[}\big{(}\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1}\big{)}^{-1}-\big{(}\Sigma_{\ell}^{-1}+\bar{G}_{t,\ell-1}+uu^{\top}\big{)}^{-1}\Big{]}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\Big{[}\bar{\Sigma}_{t,\ell-1}-\big{(}\bar{\Sigma}_{t,\ell-1}^{-1}+uu^{\top}\big{)}^{-1}\Big{]}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\Big{[}\bar{\Sigma}_{t,\ell-1}\frac{uu^{\top}}{1+u^{\top}\bar{\Sigma}_{t,\ell-1}u}\bar{\Sigma}_{t,\ell-1}\Big{]}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\frac{uu^{\top}}{1+u^{\top}\bar{\Sigma}_{t,\ell-1}u}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}$ $$

However, we it follows from the induction hypothesis that $uu^{\top}=\tilde{\Sigma}_{t+1,\ell-1}^{-1}-\bar{G}_{t,\ell-1}=\bar{\Sigma}_{t+1,\ell-1}^{-1}-\bar{\Sigma}_{t,\ell-1}^{-1}\succeq\sigma^{-2}\sigma_{\textsc{max}}^{-2(\ell-1)}\mathrm{P}_{A_{t},\ell-1}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell-1}$. Therefore,

$$ $\displaystyle\bar{\Sigma}_{t+1,\ell}^{-1}-\bar{\Sigma}_{t,\ell}^{-1}$ $\displaystyle=\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\frac{uu^{\top}}{1+u^{\top}\bar{\Sigma}_{t,\ell-1}u}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle\succeq\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\frac{\sigma^{-2}\sigma_{\textsc{max}}^{-2(\ell-1)}\mathrm{P}_{A_{t},\ell-1}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell-1}}{1+u^{\top}\bar{\Sigma}_{t,\ell-1}u}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle=\frac{\sigma^{-2}\sigma_{\textsc{max}}^{-2(\ell-1)}}{1+u^{\top}\bar{\Sigma}_{t,\ell-1}u}\mathrm{W}_{\ell}^{\top}\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\mathrm{P}_{A_{t},\ell-1}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\mathrm{W}_{\ell}\,,$ $\displaystyle=\frac{\sigma^{-2}\sigma_{\textsc{max}}^{-2(\ell-1)}}{1+u^{\top}\bar{\Sigma}_{t,\ell-1}u}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\,.$ $$

Finally, we use that $1+u^{\top}\bar{\Sigma}_{t,\ell-1}u\leq 1+\|u\|_{2}\lambda_{1}(\bar{\Sigma}_{t,\ell-1})\leq 1+\sigma^{-2}\sigma_{\ell}^{2}$. Here we use that $\|u\|_{2}\leq\sigma^{-2}$, which can also be proven by induction, and that $\lambda_{1}(\bar{\Sigma}_{t,\ell-1})\leq\sigma_{\ell}^{2}$, which follows from the expression of $\bar{\Sigma}_{t,\ell-1}$ in [Section 3.1](#S3.SS1). Therefore, we have that

$$ $\displaystyle\bar{\Sigma}_{t+1,\ell}^{-1}-\bar{\Sigma}_{t,\ell}^{-1}$ $\displaystyle\succeq\frac{\sigma^{-2}\sigma_{\textsc{max}}^{-2(\ell-1)}}{1+u^{\top}\bar{\Sigma}_{t,\ell-1}u}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\,,$ $\displaystyle\succeq\frac{\sigma^{-2}\sigma_{\textsc{max}}^{-2(\ell-1)}}{1+\sigma^{-2}\sigma_{\ell}^{2}}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\,,$ $\displaystyle\succeq\sigma^{-2}\sigma_{\textsc{max}}^{-2\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\,,$ $$

where the last inequality follows from the definition of $\sigma_{\textsc{max}}^{2}=\max_{\ell\in[L]}1+\sigma^{-2}\sigma_{\ell}^{2}$. This concludes the proof.

### C.5 Proof of Theorem 4.1

We start with the following standard result which we borrow from (Hong et al., ; Aouali et al., ),

$$ $\displaystyle\mathcal{B}\mathcal{R}(n)$ $\displaystyle\leq\sqrt{2n\log(1/\delta)}\sqrt{\mathbb{E}\left[\sum_{t=1}^{n}\|X_{t}\|_{\check{\Sigma}_{t,A_{t}}}^{2}\right]}+cn\delta\,,\qquad\text{where $c>0$ is a constant}\,.$ (28) $$

Then we use [Lemma C.2](#A3.Thmtheorem2) and express the marginal covariance $\check{\Sigma}_{t,A_{t}}$ as

$$ $\displaystyle\check{\Sigma}_{t,i}=\hat{\Sigma}_{t,i}+\sum_{\ell\in[L]}\mathrm{P}_{i,\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{i,\ell}^{\top}\,,$ where $\,\mathrm{P}_{i,\ell}=\hat{\Sigma}_{t,i}\Sigma_{1}^{-1}\mathrm{W}_{1}\prod_{k=1}^{\ell-1}\bar{\Sigma}_{t,k}\Sigma_{k+1}^{-1}\mathrm{W}_{k+1}$ . (29) $$

Therefore, we can decompose $\|X_{t}\|_{\check{\Sigma}_{t,A_{t}}}^{2}$ as

$$ $\displaystyle\|X_{t}\|_{\check{\Sigma}_{t,A_{t}}}^{2}=\sigma^{2}\frac{X_{t}^{\top}\check{\Sigma}_{t,A_{t}}X_{t}}{\sigma^{2}}\stackrel{{\scriptstyle(i)}}{{=}}\sigma^{2}\big{(}\sigma^{-2}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}X_{t}+\sigma^{-2}\sum_{\ell\in[L]}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}\big{)}\,,$ $\displaystyle\stackrel{{\scriptstyle(ii)}}{{\leq}}c_{0}\log(1+\sigma^{-2}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}X_{t})+\sum_{\ell\in[L]}c_{\ell}\log(1+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})\,,$ (30) $$

where $(i)$ follows from ([29](#A3.E29)), and we use the following inequality in $(ii)$

$$ $\displaystyle x=\frac{x}{\log(1+x)}\log(1+x)\leq\left(\max_{x\in[0,u]}\frac{x}{\log(1+x)}\right)\log(1+x)=\frac{u}{\log(1+u)}\log(1+x)\,,$ $$

which holds for any $x\in[0,u]$, where constants $c_{0}$ and $c_{\ell}$ are derived as

$$ $\displaystyle c_{0}=\frac{\sigma_{1}^{2}}{\log(1+\frac{\sigma_{1}^{2}}{\sigma^{2}})}\,,\quad c_{\ell}=\frac{\sigma_{\ell+1}^{2}}{\log(1+\frac{\sigma_{\ell+1}^{2}}{\sigma^{2}})}\,,\text{with the convention that }\sigma_{L+1}=1\,.$ $$

The derivation of $c_{0}$ uses that

$$ $\displaystyle X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}X_{t}\leq\lambda_{1}(\hat{\Sigma}_{t,A_{t}})\|X_{t}\|^{2}\leq\lambda_{d}^{-1}(\Sigma_{1}^{-1}+G_{t,A_{t}})\leq\lambda_{d}^{-1}(\Sigma_{1}^{-1})=\lambda_{1}(\Sigma_{1})=\sigma_{1}^{2}\,.$ $$

The derivation of $c_{\ell}$ follows from

$$ $\displaystyle X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}\leq\lambda_{1}(\mathrm{P}_{A_{t},\ell}\mathrm{P}_{A_{t},\ell}^{\top})\lambda_{1}(\bar{\Sigma}_{t,\ell})\|X_{t}\|^{2}$ $\displaystyle\leq\sigma_{\ell+1}^{2}\,.$ $$

Therefore, from ([C.5](#A3.Ex149)) and ([28](#A3.E28)), we get that

$$ $\displaystyle\mathcal{B}\mathcal{R}(n)$ $\displaystyle\leq\sqrt{2n\log(1/\delta)}\Big{(}\mathbb{E}\Big{[}c_{0}\sum_{t=1}^{n}\log(1+\sigma^{-2}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}X_{t})$ $\displaystyle+\sum_{\ell\in[L]}c_{\ell}\sum_{t=1}^{n}\log(1+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})\Big{]}\Big{)}^{\frac{1}{2}}+cn\delta$ (31) $$

Now we focus on bounding the logarithmic terms in ([C.5](#A3.Ex154)).

(I) First term in ([C.5](#A3.Ex154)) We first rewrite this term as

$$ $\displaystyle\log(1$ $\displaystyle+\sigma^{-2}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}X_{t})\stackrel{{\scriptstyle(i)}}{{=}}\log\det(I_{d}+\sigma^{-2}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}X_{t}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}})\,,$ $\displaystyle=\log\det(\hat{\Sigma}_{t,A_{t}}^{-1}+\sigma^{-2}X_{t}X_{t}^{\top})-\log\det(\hat{\Sigma}_{t,A_{t}}^{-1})=\log\det(\hat{\Sigma}_{t+1,A_{t}}^{-1})-\log\det(\hat{\Sigma}_{t,A_{t}}^{-1})\,,$ $$

where $(i)$ follows from the Weinstein–Aronszajn identity. Then we sum over all rounds $t\in[n]$, and get a telescoping

$$ $\displaystyle\sum_{t=1}^{n}$ $\displaystyle\log\det(I_{d}+\sigma^{-2}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}X_{t}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}})=\sum_{t=1}^{n}\log\det(\hat{\Sigma}_{t+1,A_{t}}^{-1})-\log\det(\hat{\Sigma}_{t,A_{t}}^{-1})\,,$ $\displaystyle=\sum_{t=1}^{n}\sum_{i=1}^{K}\log\det(\hat{\Sigma}_{t+1,i}^{-1})-\log\det(\hat{\Sigma}_{t,i}^{-1})=\sum_{i=1}^{K}\sum_{t=1}^{n}\log\det(\hat{\Sigma}_{t+1,i}^{-1})-\log\det(\hat{\Sigma}_{t,i}^{-1})\,,$ $\displaystyle=\sum_{i=1}^{K}\log\det(\hat{\Sigma}_{n+1,i}^{-1})-\log\det(\hat{\Sigma}_{1,i}^{-1})\stackrel{{\scriptstyle(i)}}{{=}}\sum_{i=1}^{K}\log\det(\Sigma_{1}^{\frac{1}{2}}\hat{\Sigma}_{n+1,i}^{-1}\Sigma_{1}^{\frac{1}{2}})\,,$ $$

where $(i)$ follows from the fact that $\hat{\Sigma}_{1,i}=\Sigma_{1}$. Now we use the inequality of arithmetic and geometric means and get

$$ $\displaystyle\sum_{t=1}^{n}\log\det(I_{d}+\sigma^{-2}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}}X_{t}X_{t}^{\top}\hat{\Sigma}_{t,A_{t}}^{\frac{1}{2}})$ $\displaystyle=\sum_{i=1}^{K}\log\det(\Sigma_{1}^{\frac{1}{2}}\hat{\Sigma}_{n+1,i}^{-1}\Sigma_{1}^{\frac{1}{2}})\,,$ $\displaystyle\leq\sum_{i=1}^{K}d\log\left(\frac{1}{d}\operatorname{Tr}(\Sigma_{1}^{\frac{1}{2}}\hat{\Sigma}_{n+1,i}^{-1}\Sigma_{1}^{\frac{1}{2}})\right)\,,$ (32) $\displaystyle\leq\sum_{i=1}^{K}d\log\left(1+\frac{n}{d}\frac{\sigma_{1}^{2}}{\sigma^{2}}\right)=Kd\log\left(1+\frac{n}{d}\frac{\sigma_{1}^{2}}{\sigma^{2}}\right)\,.$ $$

(II) Remaining terms in ([C.5](#A3.Ex154)) Let $\ell\in[L]$. Then we have that

$$ $\displaystyle\log(1$ $\displaystyle+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})=\sigma_{\textsc{max}}^{2\ell}\sigma_{\textsc{max}}^{-2\ell}\log(1+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})\,,$ $\displaystyle\leq\sigma_{\textsc{max}}^{2\ell}\log(1+\sigma^{-2}\sigma_{\textsc{max}}^{-2\ell}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})\,,$ $\displaystyle\stackrel{{\scriptstyle(i)}}{{=}}\sigma_{\textsc{max}}^{2\ell}\log\det(I_{d}+\sigma^{-2}\sigma_{\textsc{max}}^{-2\ell}\bar{\Sigma}_{t,\ell}^{\frac{1}{2}}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}^{\frac{1}{2}})\,,$ $\displaystyle=\sigma_{\textsc{max}}^{2\ell}\Big{(}\log\det(\bar{\Sigma}_{t,\ell}^{-1}+\sigma^{-2}\sigma_{\textsc{max}}^{-2\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell})-\log\det(\bar{\Sigma}_{t,\ell}^{-1})\Big{)}\,,$ $$

where we use the Weinstein–Aronszajn identity in $(i)$. Now we know from [Lemma C.3](#A3.Thmtheorem3) that the following inequality holds $\sigma^{-2}\sigma_{\textsc{max}}^{-2\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\preceq\bar{\Sigma}_{t+1,\ell}^{-1}-\bar{\Sigma}_{t,\ell}^{-1}$. As a result, we get that $\bar{\Sigma}_{t,\ell}^{-1}+\sigma^{-2}\sigma_{\textsc{max}}^{-2\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\preceq\bar{\Sigma}_{t+1,\ell}^{-1}$. Thus,

$$ $\displaystyle\log(1+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})$ $\displaystyle\leq\sigma_{\textsc{max}}^{2\ell}\Big{(}\log\det(\bar{\Sigma}_{t+1,\ell}^{-1})-\log\det(\bar{\Sigma}_{t,\ell}^{-1})\Big{)}\,,$ $$

Then we sum over all rounds $t\in[n]$, and get a telescoping

$$ $\displaystyle\sum_{t=1}^{n}\log(1+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})$ $\displaystyle\leq\sigma_{\textsc{max}}^{2\ell}\sum_{t=1}^{n}\log\det(\bar{\Sigma}_{t+1,\ell}^{-1})-\log\det(\bar{\Sigma}_{t,\ell}^{-1})\,,$ $\displaystyle=\sigma_{\textsc{max}}^{2\ell}\Big{(}\log\det(\bar{\Sigma}_{n+1,\ell}^{-1})-\log\det(\bar{\Sigma}_{1,\ell}^{-1})\Big{)}\,,$ $\displaystyle\stackrel{{\scriptstyle(i)}}{{=}}\sigma_{\textsc{max}}^{2\ell}\Big{(}\log\det(\bar{\Sigma}_{n+1,\ell}^{-1})-\log\det(\Sigma_{\ell+1}^{-1})\Big{)}\,,$ $\displaystyle=\sigma_{\textsc{max}}^{2\ell}\Big{(}\log\det(\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}})\Big{)}\,,$ $$

where we use that $\bar{\Sigma}_{1,\ell}=\Sigma_{\ell+1}$ in $(i)$. Finally, we use the inequality of arithmetic and geometric means and get that

$$ $\displaystyle\sum_{t=1}^{n}\log(1+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})$ $\displaystyle\leq\sigma_{\textsc{max}}^{2\ell}\Big{(}\log\det(\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}})\Big{)}\,,$ $\displaystyle\leq d\sigma_{\textsc{max}}^{2\ell}\log\left(\frac{1}{d}\operatorname{Tr}(\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}})\right)\,,$ (33) $\displaystyle\leq d\sigma_{\textsc{max}}^{2\ell}\log\left(1+\frac{\sigma_{\ell+1}^{2}}{\sigma_{\ell}^{2}}\right)\,,$ $$

The last inequality follows from the expression of $\bar{\Sigma}_{n+1,\ell}^{-1}$ in ([3.1.1](#S3.Ex13)) that leads to

$$ $\displaystyle\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}}$ $\displaystyle=I_{d}+\Sigma_{\ell+1}^{\frac{1}{2}}\bar{G}_{t,\ell}\Sigma_{\ell+1}^{\frac{1}{2}}\,,$ $\displaystyle=I_{d}+\Sigma_{\ell+1}^{\frac{1}{2}}\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}\Sigma_{\ell+1}^{\frac{1}{2}}\,,$ (34) $$

since $\bar{G}_{t,\ell}=\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}$. This allows us to bound $\frac{1}{d}\operatorname{Tr}(\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}})$ as

$$ $\displaystyle\frac{1}{d}\operatorname{Tr}(\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}})$ $\displaystyle=\frac{1}{d}\operatorname{Tr}(I_{d}+\Sigma_{\ell+1}^{\frac{1}{2}}\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}\Sigma_{\ell+1}^{\frac{1}{2}})\,,$ $\displaystyle=\frac{1}{d}(d+\operatorname{Tr}(\Sigma_{\ell+1}^{\frac{1}{2}}\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}\Sigma_{\ell+1}^{\frac{1}{2}})\,,$ $\displaystyle\leq 1+\frac{1}{d}\sum_{k=1}^{d}\lambda_{1}(\Sigma_{\ell+1}^{\frac{1}{2}}\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}\Sigma_{\ell+1}^{\frac{1}{2}}\,,$ $\displaystyle\leq 1+\frac{1}{d}\sum_{k=1}^{d}\lambda_{1}(\Sigma_{\ell+1})\lambda_{1}(\mathrm{W}_{\ell}^{\top}\mathrm{W}_{\ell})\lambda_{1}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\,,$ $\displaystyle\leq 1+\frac{1}{d}\sum_{k=1}^{d}\lambda_{1}(\Sigma_{\ell+1})\lambda_{1}(\mathrm{W}_{\ell}^{\top}\mathrm{W}_{\ell})\lambda_{1}\big{(}\Sigma_{\ell}^{-1}\big{)}\,,$ $\displaystyle\leq 1+\frac{1}{d}\sum_{k=1}^{d}\frac{\sigma_{\ell+1}^{2}}{\sigma_{\ell}^{2}}=1+\frac{\sigma_{\ell+1}^{2}}{\sigma_{\ell}^{2}}\,,$ (35) $$

where we use the assumption that $\lambda_{1}(\mathrm{W}_{\ell}^{\top}\mathrm{W}_{\ell})=1$ (A2) and that $\lambda_{1}(\Sigma_{\ell+1})=\sigma_{\ell+1}^{2}$ and $\lambda_{1}(\Sigma_{\ell}^{-1})=1/\sigma_{\ell}^{2}$. This is because $\Sigma_{\ell}=\sigma_{\ell}^{2}I_{d}$ for any $\ell\in[L+1]$. Finally, plugging ([C.5](#A3.Ex160)) and ([C.5](#A3.Ex172)) in ([C.5](#A3.Ex154)) concludes the proof.

### C.6 Proof of Proposition 4.2

We use exactly the same proof in [Section C.5](#A3.SS5), with one change to account for the sparsity assumption (A3). The change corresponds to ([C.5](#A3.Ex172)). First, recall that ([C.5](#A3.Ex172)) writes

$$ $\displaystyle\sum_{t=1}^{n}\log(1+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})$ $\displaystyle\leq\sigma_{\textsc{max}}^{2\ell}\Big{(}\log\det(\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}})\Big{)}\,,$ $$

where

$$ $\displaystyle\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}}$ $\displaystyle=I_{d}+\Sigma_{\ell+1}^{\frac{1}{2}}\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}\Sigma_{\ell+1}^{\frac{1}{2}}\,,$ $\displaystyle=I_{d}+\sigma_{\ell+1}^{2}\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}\,,$ (36) $$

where the second equality follows from the assumption that $\Sigma_{\ell+1}=\sigma_{\ell+1}^{2}I_{d}$. But notice that in our assumption, (A3), we assume that $\mathrm{W}_{\ell}=(\bar{\mathrm{W}}_{\ell},0_{d,d-d_{\ell}})$, where $\bar{\mathrm{W}}_{\ell}\in^{d\times d_{\ell}}$ for any $\ell\in[L]$. Therefore, we have that for any $d\times d$ matrix $\mathrm{B}\in^{dd\times d}$, the following holds, $\mathrm{W}_{\ell}^{\top}\mathrm{B}\mathrm{W}_{\ell}={\begin{pmatrix}\bar{\mathrm{W}}_{\ell}^{\top}\mathrm{B}\bar{\mathrm{W}}_{\ell}&0_{d_{\ell},d-d_{\ell}}\\
0_{d-d_{\ell},d_{\ell}}&0_{d-d_{\ell},d-d_{\ell}}\end{pmatrix}}$. In particular, we have that

$$ $\displaystyle\mathrm{W}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\mathrm{W}_{\ell}={\begin{pmatrix}\bar{\mathrm{W}}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\bar{\mathrm{W}}_{\ell}&0_{d_{\ell},d-d_{\ell}}\\ 0_{d-d_{\ell},d_{\ell}}&0_{d-d_{\ell},d-d_{\ell}}\end{pmatrix}}\,.$ (37) $$

Therefore, plugging this in ([C.6](#A3.Ex181)) yields that

$$ $\displaystyle\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}}={\begin{pmatrix}I_{d_{\ell}}+\sigma_{\ell+1}^{2}\bar{\mathrm{W}}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\bar{\mathrm{W}}_{\ell}&0_{d_{\ell},d-d_{\ell}}\\ 0_{d-d_{\ell},d_{\ell}}&I_{d-d_{\ell}}\end{pmatrix}}\,.$ (38) $$

As a result, $\det(\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}})=\det(I_{d_{\ell}}+\sigma_{\ell+1}^{2}\bar{\mathrm{W}}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\bar{\mathrm{W}}_{\ell})$. This allows us to move the problem from a $d$-dimensional one to a $d_{\ell}$-dimensional one. Then we use the inequality of arithmetic and geometric means and get that

$$ $\displaystyle\sum_{t=1}^{n}\log(1$ $\displaystyle+\sigma^{-2}X_{t}^{\top}\mathrm{P}_{A_{t},\ell}\bar{\Sigma}_{t,\ell}\mathrm{P}_{A_{t},\ell}^{\top}X_{t})\leq\sigma_{\textsc{max}}^{2\ell}\Big{(}\log\det(\Sigma_{\ell+1}^{\frac{1}{2}}\bar{\Sigma}_{n+1,\ell}^{-1}\Sigma_{\ell+1}^{\frac{1}{2}})\Big{)}\,,$ $\displaystyle=\sigma_{\textsc{max}}^{2\ell}\log\det(I_{d_{\ell}}+\sigma_{\ell+1}^{2}\bar{\mathrm{W}}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\bar{\mathrm{W}}_{\ell})\,,$ $\displaystyle\leq d_{\ell}\sigma_{\textsc{max}}^{2\ell}\log\left(\frac{1}{d_{\ell}}\operatorname{Tr}(I_{d_{\ell}}+\sigma_{\ell+1}^{2}\bar{\mathrm{W}}_{\ell}^{\top}\big{(}\Sigma_{\ell}^{-1}-\Sigma_{\ell}^{-1}\bar{\Sigma}_{t,\ell-1}\Sigma_{\ell}^{-1}\big{)}\bar{\mathrm{W}}_{\ell})\right)\,,$ $\displaystyle\leq d_{\ell}\sigma_{\textsc{max}}^{2\ell}\log\left(1+\frac{\sigma_{\ell+1}^{2}}{\sigma_{\ell}^{2}}\right)\,.$ (39) $$

To get the last inequality, we use derivations similar to the ones we used in ([C.5](#A3.Ex175)). Finally, the desired result in obtained by replacing ([C.5](#A3.Ex172)) by ([C.6](#A3.Ex182)) in the previous proof in [Section C.5](#A3.SS5).

## Appendix D Additional Experiment

Regret scaling with $K$, $L$, $d$. In this experiment, we aim to empirically verify the relationships outlined in [Theorem 4.1](#S4.Thmtheorem1) between the regret of dTS and several key factors: the number of actions $K$, the context dimension $d$, and the diffusion depth $L$. We maintain the same experimental setup with linear rewards, for which we have derived a Bayes regret. In [Figure 5](#A4.F5), we plot the regret of dTS-LL across varying values of these parameters: $K\in\left\{10,100,500,1000\right\}$, $d\in\left\{5,10,15,20\right\}$, and $L\in\left\{2,4,5,6\right\}$. 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. Interestingly, the empirical regret of dTS-LL increases as the number of actions increases, which is consistent with the regret bound outlined in [Theorem 4.1](#S4.Thmtheorem1). This observation may appear counterintuitive, as one might expect the regret to depend solely on the diffusion depth $L$ in our setting. However, as discussed in [Section 4](#S4), this behavior is caused by the fact that action parameters are not deterministic given the latent parameters.

Figure: Figure 5: Regret of dTS-LL w.r.t. $K,d,L$.
Refer to caption: /html/2402.10028/assets/x5.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How does the choice of the link function $ f_\ell$ impact the posterior approximation and computational efficiency in non-linear scenarios?

2. How would the regret bounds change if the action parameters $\theta^*$ were updated using a non-linear transformation?

3. How does the choice of the number of layers $L$ in the diffusion model affect the convergence rate of the posterior approximations, especially as the action space $K$ increases?

4. What impact does the sparsity assumption on the mixing matrices $W_\ell$ have on the model's computational efficiency and the quality of the posterior approximations?


5. As I have mentioned in (3) of Weaknesses, can you discuss alternative approximation techniques (e.g. Monte Carlo methods [1,2]) that could potentially offer more flexible or accurate approximations, especially for non-linear reward distributions?

References:

[1] Xu, Pan, et al. "Langevin monte carlo for contextual bandits." International Conference on Machine Learning. PMLR, 2022.

[2] Karbasi, Amin, et al. "Langevin thompson sampling with logarithmic communication: bandits and reinforcement learning." International Conference on Machine Learning. PMLR, 2023.
soundness: 3
strengths: 1. The paper thoroughly explains the posterior approximation process for both linear and non-linear diffusion models.

2. The recursive hierarchical sampling in dTS is well-defined, which simplifies the complex diffusion model and supports efficient computational sampling.
summary: This paper introduces Diffusion Thompson Sampling (dTS), a novel algorithm leveraging pre-trained diffusion model priors to optimize exploration in contextual bandits. By capturing correlations among actions, dTS enhances both computational and statistical efficiency. It offers theoretical insights into posterior approximations and Bayes regret bounds, providing a structured and computationally manageable solution for large action spaces in contextual bandit problems.
weaknesses: 1. The method’s performance is heavily reliant on the accuracy of the pre-trained diffusion model. If the model's prior assumptions are incorrect or misspecified, the effectiveness of the posterior approximations and subsequent regret bounds could be compromised.

2. There is limited discussion on alternative approximation techniques (e.g., variational inference or Monte Carlo methods) that could potentially offer more flexible or accurate approximations, especially for non-linear reward distributions.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Here is a list of suggestions for the authors. 
- In section 4.1, Statistical benefits, the authors mention "The only Bayesian lower bound that we know of is $\Omega(\log^2(n))$. The authors could have mentioned the minimax results from Dani et al. (2008) in $\Omega(d\sqrt{n})$ for the $d$-dimensional linear bandit setting.
- On line 354, the authors mention "we can show that dTS’s regret is independent of K in their setting, assuming the availability of $\phi$". It would be interesting to add this proof in the Appendix. 
- As mentioned before, it would be interesting to compare the performance of dTS against fairer baselines for the experiments on MovieLens problem such as the Vits: Variational Inference Thompson Sampling" from Clavier et al. (2023). 
- It would also be interesting to compare the performance of dTS against the DiffTS proposed by Hsieh et al. (2023) on bandit problems and verify if dTS can recover the same performance while presenting computational advantages.
- In the Appendix "Extended related work", the authors are pointed to two papers that studied the performance of TS for contextual bandits and derived Bayesian regret bounds: Neu et al. "Lifting the information ratio: An information-theoretic analysis of thompson sampling for contextual bandits" (2022) and Gouverneur et al. "Thompson sampling regret bounds for contextual bandits with sub-Gaussian rewards." (2023).
- For the sake of completeness, the authors are suggested to include in the Appendix the pseudocode of the LinTS algorithm used for the experiments.
- On line 239, the authors mention that their "bound is $\tilde{O}(\sqrt{n})$. A suggestion is to include the non-logarithmic dependencies on $d$ and $K$.
- On line 127, a suggestion to the authors to change "We design Thompson sampling that samples" to "We design a Thompson sampling algorithm that samples".
- On lines 122-123, the authors wrote "The Bayes regret is known to capture the benefits of using informative priors, and hence it is suitable for our problem". The authors are suggested to provide references to support their claim.
soundness: 3
strengths: The main strength of the paper is its overall clarity and coherence. The authors introduce the interesting idea of extending the work of Hsieh et al. (2023) to contextual bandit problems, show how to efficiently compute an approximation of the posterior to sample from, derive regret bounds under some specific assumptions, and perform experiments to demonstrate the performance of the proposed method. The different ideas are explained clearly, the code for the experiments is provided and user-friendly, the notations used are rigorous, and the proof techniques are thorough and explicit.
summary: This paper studies the performance of Thompson Sampling for contextual bandit problems with generalized linear model reward distribution. Starting from a prior distribution over the "action-parameter", Thompson Sampling algorithm works by randomly selecting actions according to their posterior probability of being optimal. More specifically, at each time step, it samples an "action-parameter" estimate from the posterior distribution conditioned on the history and selects the action that is optimal for the sampled parameter estimate given the context.

In this work, the authors study the case where the prior is (or can be approximated by) a diffusion model. This idea is inspired by the work of Hsieh et al. (2023), where they showed they could use diffusion model to model the prior together with Thompson Sampling algorithm to solve bandit problems. This paper builds on Hsieh results and extends their idea to contextual bandit problems. They also show how to efficiently compute an approximation of posterior and how to efficiently sample from this approximation, and call the resulting algorithm diffusion Thompson sampling. Under the assumption that the true prior can be written as a linear Gaussian system, their approximate posterior matches the exact posterior and the authors could derive Bayesian regret bounds in $\tilde{O}(\sqrt{n (d K \sigma_{1}^2 + \sum_{l=1}^L d \sigma_{l+1}^2 \sigma_{\text{MAX}}^{2l}})$ where $n$ is the number of time steps, $d$ is the dimension of the problem, $K$ is the number of possible actions, $\sigma$ is the variance of the rewards, $\sigma_1$ is the isotropic variance of the action parameter given the first latent variable, $L$ is the number of latent parameters and $\sigma_{l+1}$ is the isotropic variance of the latent parameter $l$ conditioned on the latent parameter $l+1$ (by construction the $L+1$ latent parameter is zero) and $\sigma^2_{\text{MAX}} = \max_{l\in[L+1]} 1+\frac{\sigma_l^2}{\sigma^2}$. 

The authors perform two experiments to demonstrate the performance of the proposed method. The first experiment is performed on synthetic problems where the true prior is a diffusion model. They compare their method with several baselines (LinTS, LinUCB, HierTS, GLM-TS, UCB-GLM) and show that the diffusion Thompson sampling performs better. The second experiment has a prior distribution that is not a diffusion model. In this case, the author first pre-train a diffusion model to approximate the prior distribution before running the algorithm. They show that their method performs better than LinTS.
weaknesses: Although the paper's main ideas are interesting, the experiment section presents some weaknesses. The first experiment, Figure 2, compares the proposed algorithm dTS against several baselines HierTS, LinTS, GLM-TS, UCB-GLM for problems where the true prior is a diffusion model. Unsurprisingly, this setting perfectly fits the proposed algorithm, and it outperforms the baseline algorithms. This first experiments can be understood as a sanity check test that dTS passed. A more interesting experiment, Figure 4, tests the performance of dTS on problems where the true prior is not a diffusion model. However, the performance of dTS in this setting is only compared to LinTS, which is by design not suited to this contextual bandit problem as it cannot capture the correlations among actions. It is, therefore, not surprising that dTS improves the performance of LinTS for this problem, and it is difficult to appreciate its performance. A more fair comparison would have been against algorithms suited to the setting, such as "Vits: Variational Inference Thompson Sampling". Another interesting experiment would have been to compare the performance of dTS against the DiffTS proposed by Hsieh et al. (2023) on bandit problems and verify if dTS can recover the same performance while presenting computational advantages.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Theorem 4.1 presents a **$O(\sqrt{d(K+L)T})$ Bayes regret** bound. Compared to Russo and Van Roy’s (2014) **$O(d\sqrt{T})$ Bayes regret** for linear Thompson sampling, I am curious about how your bound handles the **$O(\sqrt{d})$** factor. While I understand that the additional **K** factor stems from a different setting (with shared parameters across arms in Russo and Van Roy’s work), I would appreciate further elaboration on how your method results in a **$O(\sqrt{d})$** regret bound. Could you explain this in more detail?
soundness: 3
strengths: - The paper leverages *pre-trained diffusion models* to capture correlations in large action spaces, which provides a more structured and efficient approach to exploration compared to traditional methods in contextual bandits.

- The paper provides a Bayes regret bound for diffusion Thompson sampling (dTS) in the linear case, offering some theoretical backing to the proposed algorithm, though the assumptions could be more thoroughly justified (see the Weakness section).

- Comprehensive Empirical Results: The paper presents empirical evaluations that show dTS outperforming standard methods in several settings, particularly in large-scale action spaces, which lends support to the algorithm’s practical effectiveness.
summary: The paper introduces Diffusion Thompson Sampling (dTS), an algorithm designed to handle the exploration-exploitation trade-off in contextual bandits with large action spaces. The key idea is to leverage pre-trained diffusion models to capture correlations between actions, which can guide more efficient exploration.

The authors propose using diffusion models as priors in a Thompson sampling framework, allowing the algorithm to explore the action space more strategically by exploiting the underlying structure of the action set. The paper provides a theoretical analysis of the algorithm, including an upper bound on the Bayes regret for linear diffusion and reward structures.

In addition to the theoretical results, the paper presents empirical evaluations that demonstrate the performance of the proposed method across various settings, showing that it outperforms traditional approaches in contexts with both linear and non-linear reward and diffusion models.
weaknesses: - The unclear presentation of the assumptions, along with the lack of thorough discussion or justification, makes it difficult to fully evaluate the soundness of the work (see more details below).
- In terms of numbering, the assumptions should begin from **1** instead of 0. For example, (A0), (A1), (A2) should be renamed to (A1), (A2), (A3), etc., for consistency and clarity.
- Assumption (A0) should be separated into two distinct parts:  
  (i) The assumption that **$W_\ell$ is known** for all $\ell \in [L]$, and  
  (ii) The assumption of **Gaussian reward noise**, which is stronger than the sub-Gaussian noise typically assumed in the contextual bandit literature.
- Regarding **Assumption (A2)**, I encourage the authors to justify why such an assumption is necessary (rather than simply referring to "milder assumptions"). This assumption is non-standard compared to existing literature, and further elaboration is needed to position the analysis on firm ground. Providing a more explicit explanation would allow future work to build upon these assumptions, even if they seem restrictive.
- A fundamental concern lies in how the **$O(\sqrt{d(K+L)T})$ Bayes regret** in Theorem 4.1 compares with the well-established **$O(\sqrt{KT})$ Bayes regret** for the non-contextual multi-armed bandit (as in Russo and Van Roy, 2014). The bound of $O(\sqrt{d(K+L)T})$ in Theorem 4.1 is worse than both the non-contextual result in Bayes regret and the **worst-case regret bound of $O(\sqrt{KT})$** from Agrawal and Goyal (2023b) without prior. Why should practitioners consider this method in the linear case when the theory suggests it’s more effective to treat arms independently and use the non-contextual MAB version of Thompson sampling?
- The section “**Why the bound increases with L?**” (Line 268) is confusing. The theoretical result in the paper shows that increasing $L$ worsens the upper bound, implying that in the linear case, one should use **$L=1$**. The current explanation seems contrary to the theoretical findings for the linear case, and the authors should clarify this discrepancy.
- A **table of notations** is necessary, at least at the beginning of the appendix, to improve clarity and assist readers in following the paper’s mathematical presentation.


### **Minor Comments**:

- **Line 1151-1152**: "This sparsity assumption is both a novel and a key technical contribution to our work."  
  It is debatable whether adding a structural assumption qualifies as novel unless it offers new insights or introduces milder assumptions compared to existing literature.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Please refer to the Weaknesses section above.
soundness: 2
strengths: 1. Combining bandit algorithms with diffusion models is an innovative approach.

2. The paper introduces an efficient approximation for posterior updates in diffusion models, which I found interesting. The provided experiments offer valuable insights into the performance of this approximation.
summary: This paper studies contextual bandits where the feature vectors of the actions are hidden and generated by a known diffusion model. The authors propose an algorithm based on posterior sampling and show that the posterior update can be performed efficiently, using an approximation inspired by the linear case. The authors also theoretically demonstrate that the algorithm achieves a sublinear regret bound when the link functions are linear. Lastly, they present experiments to illustrate the practical performance of the algorithm.
weaknesses: 1. I am unsure if the setting makes sense. Specifically, since $\theta_{\*,i}$ are fixed across episodes, one can interpret their concatenation as the hidden vector. This reduces the problem to a generalized linear contextual bandit with dimension $dK$, where the feature vector for each action $i \in [K]$ is $[0, \dots, 0, X_t, 0, \dots]$, activating only the entries corresponding to $\theta_{\*,i}$. According to my calculations, this should result in an algorithm with regret $\tilde{O}(\sqrt{dKT})$. Therefore, I question the advantage of introducing the posterior sampling over diffusion model in this context.

2. The theoretical analysis assumes linear link functions. However, since a combination of linear functions remains linear, the diffusion model may effectively reduce to a single linear function, potentially trivializing the problem.

3. Certain parts of the paper are unclear. For instance, it is not specified whether the covariance matrix $\Sigma_l$ or the link function $g$ are revealed to the agent.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: I would refer to the weaknesses section. Any clarification or illustration would be very helpful.
soundness: 3
strengths: 1. The paper proposes a new algorithm for contextual bandits with dependent arms, using diffusion models.
2. The paper is relatively clearly written and easy to follow.
3. The paper includes a variety of interpretations of the results, both theoretical and numerical.
4. The theoretical statements and numerical experiments are solid, as they are clearly presented and explained. The regret bound is nearly optimal in that it is approximately of order $\sqrt{T}$, when the $\log{T}$ term is negligible.
summary: The paper considers the classical contextual bandit problem from a Bayesian perspective, particularly through the use of a diffusion model. The arms are correlated via a shared distribution for their parameters. The paper proposes a new Thompson sampling algorithm that incorporates a diffusion model prior. The estimation of the posterior distribution and efficient sampling method are applicable to other Thompson sampling problems. The effectiveness of the algorithm is demonstrated by the regret bound in a linear case, where the regret is of order $\sqrt{KT(\log{T})^2}$. Additionally, the paper presents a numerical study to validate the algorithm, showing improvement over existing benchmarks in cases where the true prior is a diffusion model as well as in cases where the true prior is not a diffusion model.
weaknesses: 1. The actual role of using a diffusion model remains unclear. If the prior is not a diffusion model, will the algorithm still perform effectively in this scenario?

2. The theoretical statements provided are only valid for the linear case. It is also unclear whether there would be a significant gap in performance between linear and non-linear cases, which could impact the applicability of the algorithm in more complex settings.

3. The paper addresses both the contextual bandit and bandits with dependent arms. However, this combined focus makes it challenging to isolate the specific impact of using a diffusion model. It remains unclear whether the diffusion model could be effectively applied in the traditional contextual bandit problem without dependent arms, as dependent arms generally allow for more informative decision-making.

4. Lastly, I am curious whether a diffusion model is truly necessary in this setting, especially if it is already pre-trained, to handle the contextual bandit scenario with dependent arms.

Meta-review: This paper introduces Diffusion Thompson Sampling (dTS), an algorithm that addresses the exploration-exploitation trade-off in contextual bandits with large action spaces. By integrating pre-trained diffusion models, dTS identifies correlations between actions, facilitating more efficient exploration. The algorithm incorporates diffusion models as priors within a Thompson sampling framework, leveraging the underlying structure of the action set to explore the action space more effectively. The theoretical analysis provides upper bounds on the Bayes regret for linear diffusion and reward structures.

The paper is generally well-written, and the proposed algorithm is novel, supported by both solid theoretical results and empirical evaluations. However, during the discussion phase, reviewers identified some weaknesses in the method and its regret analysis. Specifically, the approach seems fundamentally limited when applied to the linear diffusion case. Reviewer KvQk pointed out that linear diffusions reduce to a single Gaussian, making the motivation for using the proposed model in the linear diffusion setting unjustified. As such, the linear assumption appears unreasonable, raising questions about the necessity of a diffusion model in this context. Additionally, the paper does not provide regret analysis for nonlinear diffusions, leaving a significant gap between the theoretical framework and the method’s potential practical applicability. Given these concerns, the reviewers have recommended rejecting the paper.

Decision: Reject