Paper: Linear diffusion models meet contextual bandits with large action spaces (ICLR-2024)

Reviewer comment (Novelty):
The adoption of the diffusion prior in bandits is interesting and powerful according to previous empirical studies. However, the paper's focus on the linear diffusion prior, essentially a linear Gaussian system, simplifies the problem considerably. Although deriving the exact formula for the posterior mean and covariance, as indicated in equations (5)-(8) of this paper, requires some calculations (see also weakness 2, 3, 4), once these results are established, I don't observe too much additional difficulty compared to prior Thompson sampling papers for linear Gaussian bandits.

Author response:
We are grateful for the reviewer's insights regarding the use of results from linear Gaussian systems (LGS) to streamline our proof exposition. We plan to revise our paper to incorporate these suggestions into more simplified proofs, but we need to make sure that the direct application of these suggestions is possible due to our use of level-specific covariances and mixing matrices $\Sigma_\ell$ and $W_\ell$ (they depend on $\ell$). We will promptly update our paper with these revisions. Additionally, we'd like to note that our current proofs for posterior derivations in the Gaussian case were included for completeness, following the convention in prior research on hierarchical Bayesian bandits where detailed posterior derivations are typically provided.

Reviewer comment (Theory):
As pointed out in the first weakness, the linear diffusion considered in this paper is actually a linear Gaussian system (LGS). Several standard results from LGS theory could be directly applied to this paper to simplify the proof.

Author response:
We are grateful for the reviewer's insights regarding the use of results from linear Gaussian systems (LGS) to streamline our proof exposition. We plan to revise our paper to incorporate these suggestions into more simplified proofs, but we need to make sure that the direct application of these suggestions is possible due to our use of level-specific covariances and mixing matrices $\Sigma_\ell$ and $W_\ell$ (they depend on $\ell$). We will promptly update our paper with these revisions. Additionally, we'd like to note that our current proofs for posterior derivations in the Gaussian case were included for completeness, following the convention in prior research on hierarchical Bayesian bandits where detailed posterior derivations are typically provided.

Reviewer comment (Novelty):
Proposition 2 in Appendix B is straightforward from equations (13.89) and (13.90) in [1] when the parameters $\Gamma = 0, A = I, z_1 = W_1 \psi_{\*,1}, V_0 = \Sigma_1, \Sigma = \sigma^2 I,$ and $C_n = \mathbf{1}(A_n = i) X_{n}^\top$ are substituted into equations (13.78)-(13.83) of [1]. Here, $\Gamma, A, z_1, V_0, C_n$ and $\Sigma$ are notations from [1], while $W_1, \psi_{\*,1}, \Sigma_1, \sigma^2, A_n,$ and $X_n$ come from the paper. It's worth noting that although [1] let $C_n \equiv C$ for the sake of simplicity, generalizing to allow varying $C_n$ is straightforward.

Author response:
We are grateful for the reviewer's insights regarding the use of results from linear Gaussian systems (LGS) to streamline our proof exposition. We plan to revise our paper to incorporate these suggestions into more simplified proofs, but we need to make sure that the direct application of these suggestions is possible due to our use of level-specific covariances and mixing matrices $\Sigma_\ell$ and $W_\ell$ (they depend on $\ell$). We will promptly update our paper with these revisions. Additionally, we'd like to note that our current proofs for posterior derivations in the Gaussian case were included for completeness, following the convention in prior research on hierarchical Bayesian bandits where detailed posterior derivations are typically provided.

Reviewer comment (Novelty):
Proposition 3 in Appendix B can be directly derived by using equations (2.109) and (2.110) from [1] and making particular choices of prior covariance and design matrices. This also makes Lemma 4 redundant since it only acts as an intermediary result to prove Proposition 3.

Author response:
We are grateful for the reviewer's insights regarding the use of results from linear Gaussian systems (LGS) to streamline our proof exposition. We plan to revise our paper to incorporate these suggestions into more simplified proofs, but we need to make sure that the direct application of these suggestions is possible due to our use of level-specific covariances and mixing matrices $\Sigma_\ell$ and $W_\ell$ (they depend on $\ell$). We will promptly update our paper with these revisions. Additionally, we'd like to note that our current proofs for posterior derivations in the Gaussian case were included for completeness, following the convention in prior research on hierarchical Bayesian bandits where detailed posterior derivations are typically provided.

Reviewer comment (Theory):
Given the knowledge of {$Q_{t,\ell}$} and $P_{t,i}$, it appears that the posterior distribution of $\theta_i$ given $H_{t}$ can be computed analytically using the properties of Gaussian distributions. Why do the authors tend to firstly sample {$\psi_{t,\ell}$} instead of directly sampling $\theta_{t,i}$ from this distribution?

Author response:
Regarding the reviewer's question on why we opt for hierarchical sampling over marginalizing the latent parameters, the primary reason is computational efficiency. Marginalization would require calculating the mean and covariance of the marginal $\theta_i | H_t$, which entails several matrix multiplications, thereby increasing computational complexity. Additionally, marginalization becomes infeasible in the context of non-linear diffusions (Section 3), and empirical studies in prior works [1] have also demonstrated that hierarchical sampling offers more efficient approximations in non-linear diffusion models for multi-armed bandits, further justifying our choice in methodology.

Reviewer comment (Novelty):
Novelty. It would be nice if the authors could better illustrate how this work differs from previous works studying contextual bandits with informative priors, especially the adopted techniques. I have come across some works, e.g., [1, 2]. It seems that the essence of this work is to leverage diffusion models are priors, while diffusion models are one particular kind of graphical model. As extending the design of Thompson sampling is rather straightforward in my view, I would love to hear the authors' comments on the technical challenges in dealing with transformers as priors.

Author response:
We would like to thank you very much for your positive feedback and valuable time. We split our response into two parts. Please let us know if you have any additional questions or concerns.

__I) Novetly__

__Algorithmic novelty.__ Hierarchical Thompson sampling has been previously explored in simpler graphical models; our work expands this to diffusion models. While adapting hierarchical sampling to multi-level graphical models is conceptually straightforward, deriving posteriors can be complex. We offer general formulas for approximations applicable in general scenarios (Section 3). Then, specifically, for linear diffusions with linear rewards  (Section 3.1), we achieve closed-form Gaussian posteriors computed recursively in an efficient way. For linear diffusions with non-linear rewards (Section 3.2), we implement an efficient Laplace approximation, leveraging the recursive formulas from Section 3.1. Both approaches demonstrate excellent empirical performance. Notably, the recursive formulas from Section 3.1 also facilitated the development of a regret bound.

__Theoretical novelty.__ We have included a detailed discussion of the technical challenges and contributions in Appendix C.1. Below is a more succinct summary of these points.

- __Lemma 2.__ In dTS, sampling is done hierarchically, and thus the marginal posterior distribution of $\theta_i | H_t$ is not explicitly defined since we only use the conditional posterior distribution of $\theta_i | H_t, \psi_1$. Our first contribution was deriving the covariance of $\theta_i | H_t$ using total covariance decomposition combined with an induction argument, as our posteriors in Section 3.1 were derived recursively. This differs from the standard Bayes regret, where the posterior distribution of $\theta_i | H_t$ is predetermined due to the absence of latent parameters. Note that the total covariance decomposition was originally used in HierTS for multi-task linear bandits with a single latent parameter. Here, we extend it to contextual bandits with multiple, successive levels of latent parameters. Furthermore, our proof diverges from HierTS by incorporating linear functions ($W_\ell\psi_\ell$), as opposed to HierTS's reliance on the more straightforward identity function $\psi_\ell$.
- __Lemma 3.__ Standard Bayes regret analysis quantifies the posterior information gain from actions, particularly by examining the increase in posterior precision of the taken action $A_t$. Our proof, however, requires considering the information gain for each latent parameter, since they are also learned. Then, our contribution lies in using recursive formulas in Section 3.1 to propagate the posterior information gain of the taken action $A_t$ to all latent parameters within the range $\ell \in [L]$. This is not only novel but also technically challenging, as it combines recursive formulas, induction arguments, the Sherman–Morrison formula, and properties of rank-1 matrices. This enables a nuanced quantification of how the information gain from a particular action can impact the learning of latent parameters.
- __Proposition 1.__ Proposition 1 enhances Theorem 1's Bayes regret analysis through the sparsity assumption __(A3)__, which is grounded in the intuitive idea that sparsity leads to fewer parameters to learn, and hence smaller regret. This assumption refines regret to depend on the reduced dimensions $d_\ell \leq d$, instead of $d$. We adapt Theorem 1's proof by partitioning the covariance matrix of each latent parameter $\psi_\ell$ into two blocks: a $d_\ell \times d_\ell$ block for learnable parameters and a $(d-d_\ell) \times (d-d_\ell)$ block for non-learnable ones. This leads to the important finding that the final regret depends only on the learnable $d_\ell$ parameter, leading to tighter regret and highlighting the impact of the sparsity assumption.

Reviewer comment (Theory):
Additionally, while the target is to handle large action space, the obtained regrets still contain $K$, i.e., the number of arms, unless $\theta$'s are fixed given $\psi_{*,1}$. I am wondering whether this meets the design goal of handling large action space.

Author response:
__II) Large Action Spaces__

We explain next how dTS is efficient in large action spaces using both theory and practice.

- __Theory.__ Roughly speaking, the regret bound of dTS scales with $K \sigma_1^2$, contrasting with the standard $K \sum_{\ell}\sigma_\ell^2$ scaling, which proves advantageous especially when $\sigma_1$ is small, as often seen in diffusion models. In particular, the regret becomes independent of $K$ when $\sigma_1=0$. Moreover, analysis in Section 4.1 shows that the performance gap between dTS and LinTS is more pronounced in large action spaces, highlighting its suitability for such scenarios. There are some works [1,2,3] where the regret does not depend on $K$. However, our setting differs significantly from [1,2,3], and the dependency on $K$ is unavoidable in our setting when $\sigma_1>0$. Precisely, while [1,2,3] use a reward function $r(x, i) = \phi(x, i)^\top \theta_*$ with a common $\theta_*$ and a known mapping $\phi$, our model employs $r(x, i) = x^\top \theta_{*, i}$, necessitating the learning of $K$ distinct $d$-dimensional parameters. In the setting of [1,2,3] with an available mapping $\phi$, dTS’s regret would be independent of $K$. However, the practical challenge of obtaining such a mapping, which must capture complex context-action relations, makes our setting relevant. For example, in recommendation systems, a product is often associated with its separate embedding. To summarize, the dependency on $K$ is intrinsic to our specific setting rather than a limitation of dTS.

- __Experiments.__ dTS enjoys notable computational and statistical efficiency, especially in large action spaces (Section 4.1). For instance, our empirical findings, showcased in Figure 2 with $K=10^4$, reveal that dTS substantially surpasses the baselines. More importantly, as $K$ increases, so does the performance gap between dTS and the baselines, underscoring dTS's enhanced scalability and effectiveness in large action spaces.

[1] "Adapting to Misspecification in Contextual Bandits" by Foster et al. NeurIPS 2021

[2] "Upper Counterfactual Confidence Bounds: a New Optimism Principle for Contextual Bandits" by Xu et al. 2021

[3] "Contextual Bandits with Large Action Spaces: Made Practical" by Zhu et al. ICML 2022

Reviewer comment (Novelty):
Despite the excellence of the proposed model, the theoretical contribution may be considered somewhat limited. Both the derivation of closed-form posteriors and the analysis of Bayes regret seem somewhat straightforward, given the existing works.

Author response:
__I) Theoretical Contributions__

We have included a detailed discussion of the technical challenges and contributions in Appendix C.1. Below is a more succinct summary of these points.

- __Lemma 2.__ In dTS, sampling is done hierarchically, and thus the marginal posterior distribution of $\theta_i | H_t$ is not explicitly defined since we only use the conditional posterior distribution of $\theta_i | H_t, \psi_1$. Our first contribution was deriving the covariance of $\theta_i | H_t$ using total covariance decomposition combined with an induction argument, as our posteriors in Section 3.1 were derived recursively. This differs from the standard Bayes regret, where the posterior distribution of $\theta_i | H_t$ is predetermined due to the absence of latent parameters. Note that the total covariance decomposition was originally used in HierTS for multi-task linear bandits with a single latent parameter. Here, we extend it to contextual bandits with multiple, successive levels of latent parameters. Furthermore, our proof diverges from HierTS by incorporating linear functions ($W_\ell\psi_\ell$), as opposed to HierTS's reliance on the more straightforward identity function $\psi_\ell$.

- __Lemma 3.__ Standard Bayes regret analysis quantifies the posterior information gain from actions, particularly by examining the increase in posterior precision of the taken action $A_t$. Our proof, however, requires considering the information gain for each latent parameter, since they are also learned. Then, our contribution lies in using recursive formulas in Section 3.1 to propagate the posterior information gain of the taken action $A_t$ to all latent parameters within the range $\ell \in [L]$. This is not only novel but also technically challenging, as it combines recursive formulas, induction arguments, the Sherman–Morrison formula, and properties of rank-1 matrices. This enables a nuanced quantification of how the information gain from a particular action can impact the learning of latent parameters.

- __Proposition 1.__ Proposition 1 enhances Theorem 1's Bayes regret analysis through the sparsity assumption __(A3)__, which is grounded in the intuitive idea that sparsity leads to fewer parameters to learn, and hence smaller regret. This assumption refines regret to depend on the reduced dimensions $d_\ell \leq d$, instead of $d$. We adapt Theorem 1's proof by partitioning the covariance matrix of each latent parameter $\psi_\ell$ into two blocks: a $d_\ell \times d_\ell$ block for learnable parameters and a $(d-d_\ell) \times (d-d_\ell)$ block for non-learnable ones. This leads to the important finding that the final regret depends only on the learnable $d_\ell$ parameter, leading to tighter regret and highlighting the impact of the sparsity assumption.

Reviewer comment (Theory):
Equation (2): At this point, as the specification to the linear regret model has not been made, the Bayes regret should be defined based on the mean of the reward distribution. Furthermore, in the definition of $A_{t, \*}$, neither the function $r()$ nor $\Theta_\*$ has been defined.

Author response:
__Typos and suggestions.__ We appreciate the reviewer pointing out these typos, and we fixed them in the revised version. To clarify, $\theta_*$ represents the concatenation of all $K$ action parameters. Moreover, $r(x, i; \theta_*)$ denotes the expected reward for action $i$ in context $x$. Also, we added that $\theta_{*, i}$ are not necessarily independent in Eq. (16).

Reviewer comment (Theory):
In the case that $\sigma_1=0$ ($\theta_{\*, i}$ is deterministic given $\psi_{*, 1}$), since all the arm vectors are the same, is it true that the regret is always zero?

Author response:
__Regret when $\sigma_1=0$.__ When $\sigma_1=0$, the regret is not always zero as it is not gap dependent. Our Bayesian regret does not capture this similarly to frequentist problem-independent regret bounds.

__Bayes regret when $\sigma_1=0$.__ We apologize for any confusion; it appears there was a misunderstanding as we were talking about our regret bound in Theorem 1 while your question is regarding the Bayes regret itself in Eq. (2). Yes, the Bayes regret in Eq. (2) indeed becomes zero when $\sigma_1=0, and this is because all action parameters have the same mixing matrix. 

In contrast, if different mixing matrices were applied for each action parameter, such as $\theta_i \sim N(W_{1, i} \psi_1, \Sigma_1)$ with $W_{1, i}$ unique to each action $i$, the Bayes regret in Eq. (2) would not be zero. Our regret bound, as currently formulated, does not reflect this nuance, similar to gap-free bounds that remain non-zero even when expected rewards for all actions are identical.

Reviewer comment (Theory):
In Equation (16), it is beneficial to explicitly state that while $\theta_{*, i}$ for $i\in[K]$ follow the same normal distribution $\mathcal{N}(0, \Sigma)$, they are not independent.

Author response:
__Typos and suggestions.__ We appreciate the reviewer pointing out these typos, and we fixed them in the revised version. To clarify, $\theta_*$ represents the concatenation of all $K$ action parameters. Moreover, $r(x, i; \theta_*)$ denotes the expected reward for action $i$ in context $x$. Also, we added that $\theta_{*, i}$ are not necessarily independent in Eq. (16).

Reviewer comment (Novelty):
Since this paper is mainly theoretical, can the authors highlight the technical contribution of this paper? More specifically, can authors elaborate on the differences in analysis between the standard Thompson Sampling? It seems that the main difference comes from the sampling of the latent parameters \psi; but under linear diffusion model, posteriors of \psi can be easily calculated.

Author response:
We have included a detailed discussion of the technical challenges and contributions in Appendix C.1. Below is a more succinct summary of these points.

- __Lemma 2.__ In dTS, sampling is done hierarchically, and thus the marginal posterior distribution of $\theta_i | H_t$ is not explicitly defined since we only use the conditional posterior distribution of $\theta_i | H_t, \psi_1$. Our first contribution was deriving the covariance of $\theta_i | H_t$ using total covariance decomposition combined with an induction argument, as our posteriors in Section 3.1 were derived recursively. This differs from the standard Bayes regret, where the posterior distribution of $\theta_i | H_t$ is predetermined due to the absence of latent parameters. Note that the total covariance decomposition was originally used in HierTS for multi-task linear bandits with a single latent parameter. Here, we extend it to contextual bandits with multiple, successive levels of latent parameters. Furthermore, our proof diverges from HierTS by incorporating linear functions ($W_\ell\psi_\ell$), as opposed to HierTS's reliance on the more straightforward identity function $\psi_\ell$.

- __Lemma 3.__ Standard Bayes regret analysis quantifies the posterior information gain from actions, particularly by examining the increase in posterior precision of the taken action $A_t$. Our proof, however, requires considering the information gain for each latent parameter, since they are also learned. Then, our contribution lies in using recursive formulas in Section 3.1 to propagate the posterior information gain of the taken action $A_t$ to all latent parameters within the range $\ell \in [L]$. This is not only novel but also technically challenging, as it combines recursive formulas, induction arguments, the Sherman–Morrison formula, and properties of rank-1 matrices. This enables a nuanced quantification of how the information gain from a particular action can impact the learning of latent parameters.

- __Proposition 1.__ Proposition 1 enhances Theorem 1's Bayes regret analysis through the sparsity assumption __(A3)__, which is grounded in the intuitive idea that sparsity leads to fewer parameters to learn, and hence smaller regret. This assumption refines regret to depend on the reduced dimensions $d_\ell \leq d$, instead of $d$. We adapt Theorem 1's proof by partitioning the covariance matrix of each latent parameter $\psi_\ell$ into two blocks: a $d_\ell \times d_\ell$ block for learnable parameters and a $(d-d_\ell) \times (d-d_\ell)$ block for non-learnable ones. This leads to the important finding that the final regret depends only on the learnable $d_\ell$ parameter, leading to tighter regret and highlighting the impact of the sparsity assumption.

Reviewer comment (Theory):
How does the authors address the problem with large action spaces? It seems that both statistical and computational complexities depends on poly(K), which can be exponentially large for linear bandits. Note that recent work [1, 2, 3] have removed the dependency in K, and thus work for large action spaces.

Author response:
We thank the reviewer for the references, they were added to Appendix A. We explain next how dTS is efficient in large action spaces using both theory and practice, but more details can be found in Appendix A.

- __Theory.__ Roughly speaking, the regret bound of dTS scales with $K \sigma_1^2$, contrasting with the standard $K \sum_{\ell}\sigma_\ell^2$ scaling, which proves advantageous especially when $\sigma_1$ is small, as often seen in diffusion models. In particular, the regret becomes independent of $K$ when $\sigma_1=0$. Moreover, analysis in Section 4.1 shows that the performance gap between dTS and LinTS is more pronounced in large action spaces, highlighting its suitability for such scenarios.  Furthermore, our setting differs significantly from [1,2,3], and the dependency on $K$ is unavoidable in our setting when $\sigma_1>0$. Precisely, while [1,2,3] use a reward function $r(x, i) = \phi(x, i)^\top \theta_*$ with a common $\theta_*$ and a known mapping $\phi$, our model employs $r(x, i) = x^\top \theta_{*, i}$, necessitating the learning of $K$ distinct $d$-dimensional parameters. In the setting of [1,2,3] with an available mapping $\phi$, dTS’s regret would be independent of $K$. However, the practical challenge of obtaining such a mapping, which must capture complex context-action relations, makes our setting relevant. For example, recommendation systems, each product is associated with its embedding. To summarize, the dependency on $K$ is intrinsic to our specific setting rather than a limitation of dTS.

- __Experiments.__ dTS enjoys notable computational and statistical efficiency, especially in large action spaces (Section 4.1). For instance, our empirical findings, showcased in Figure 2 with $K=10^4$, reveal that dTS substantially surpasses the baselines. More importantly, as $K$ increases, so does the performance gap between dTS and the baselines, underscoring dTS's enhanced scalability and effectiveness in large action spaces.