Title: Diffusion Models Meet Contextual Bandits
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion models

Bayesian contextual bandits
Abstract: Efficient online decision-making in contextual bandits is challenging, as methods without informative priors often suffer from computational or statistical inefficiencies. In this work, we leverage pre-trained diffusion models as expressive priors to capture complex action dependencies and develop a practical algorithm that efficiently approximates posteriors under such priors, enabling both fast updates and sampling. Empirical results demonstrate the effectiveness and versatility of our approach across diverse contextual bandit settings.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors consider contextual Thompson sampling for a diffusion reward model. The reward model consists of a (diffusion) hierarchy of latent parameters, each a multivariate normal conditioned on the previous latent with mean given through a (non-linear) link function of the previous latent, followed by a head that independently maps to an action vector for each of the arms of the bandit. The reward is assumed to be a function of the alignment between a provided context vector and the unknown action vectors of each arm, following a generalized linear model. \n\nThe main contribution of this work is a tractable approximation of the posterior action vectors (and latents) in case the link function is non-linear (otherwise the diffusion model marginalizes to a Gaussian). The authors adopt the Laplace approximation for the generalized linear model and approximate the link functions as linear for the posterior update.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The paper conducts both theoretical and experimental analysis for their method. The theoretical analysis is based on assuming linear link functions, where the adopted approximation becomes exact. There they inherit the cumulative regret bounds known for Thompson sampling. \n- The experimental analysis is based on appropriate statistical metrics, including significance. On the Swiss-roll dataset and MovieLens, the authors demonstrate benefits of using a non-uniform prior through pre-training of the diffusion model. \n- Moreover, the authors provide insights on a bias-variance trade off in the number of layers: shallow diffusion models are not expressive enough for the task at hand, whereas deep diffusion models result in excessive capacity and slow exploration.\n\nWeaknesses:\n- From the main text it is not entirely clear whether the theoretical analysis just amounts to computing the covariance matrix of the action space and plugging into Thompson sampling bounds from the standard literature, or whether a more involved analysis takes place. \n\n- Across the paper, the adopted notation is questionable, making the paper unnecessarily difficult to read. For instance, random variables are denoted with subscript asterisks all over the place, and likelihood functions are introduced with specific notation never to be used again. See questions for suggestions on the notation.\n\n- The experimental analysis is limited though promising. The results from Section 5.1 are entirely expected with the benchmark designed according to their method, and as far as I can tell mainly argue that in a calibrated setting adopting dTS is better than assuming independent arms. Also, the paper may benefit significantly from additional experiments in the style of MovieLens, mainly to convince practitioners that pre-training Diffusion models prior to playing bandits can give significantly better cumulative regret in the online setting.\n\n- Since the paper introduces an efficient albeit biased approximation to posteriors in diffusion models, a significant missing piece in the experimental analysis is comparison against unbiased Monte Carlo integration. In particular, a plot with compute resources on the x-axis and regret on the y-axis for different integration budgets of Monte Carlo as well as dTS is absolutely vital for selling the method, assuming dTS is indeed pareto-optimal. This is a crucial missing point, vitally affecting the usefulness of the paper.\n\nQUESTIONS:\n- Concerning notation, I would suggest to only add an asterisk for the optimal action selection / policy, and a special font for random variables (or upper-case lower-case notation). Moreover, the paper would be much easier to parse if instead of defining quantities like $\\mathcal L_l(\\psi_l) := \\mathbb P(H_t | \\psi_{*,l} = \\psi_l)$ the authors could just stick to writing the right-hand side expression in the formulas. Same with $Q_{t,l}$. It is barely any longer, easily fits in the space provided, and results in immediately recognizable posterior update formulas. Also, $P$ is used to denote two distinct probability distributions.\n\n- The authors adopt a latent parameter $\\psi_{1}$ through which correlation of the action vectors $\\theta_{I}$ is defined. It could greatly help the reader if upon introduction of the latent parameters it was explicitly stated that the action vectors are independent conditioned on the latent.\n\n- The histories $H_t$ are taken as a tuple with index $k \\in [t-1]$ which is an unfortunate and confusing choice, because $K$ already denotes the number of arms of the contextual bandit. Instead of using $I \\in [K]$ and $k \\in [t-1]$, I would suggest the authors to consistently use lower and upper case, respectively, to denote matching ordinal and cardinal quantities. \n\n- $\\mathcal R^{ACT}(n)$ and $\\mathcal R^{LAT}_l$ are discussed but nowhere defined in the Bayesian regret bound.\n\n- The Hessian in Equation (3) has a typo, $g$ should have a “double-dot” to indicate the second order derivative.\n\n- Under the linearization assumption in the theoretical analysis dTS becomes exact Thompson sampling and the prior action vectors simply become Gaussian. Doesn’t this suffice to immediately deduce regret bounds by borrowing from the extended analysis done on GP-TS? It would be great to outline the strategy for deriving the regret bounds, even outside the Appendix.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces an approach to contextual bandits by leveraging pre-trained diffusion models as priors to capture complex correlations between actions, which is particularly useful when dealing with large action spaces. The authors develop diffusion Thompson sampling (dTS), a computationally efficient algorithm that approximates posterior distributions under diffusion priors through a hierarchical sampling scheme, enabling better exploration by sharing information across related actions. They provide theoretical analysis showing improved regret bounds in the linear case and demonstrate through experiments that dTS significantly outperforms standard baselines like LinUCB and LinTS, especially as the number of actions increases, while maintaining computational efficiency even with model misspecification.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- The authors develop a two-stage approximation that first handles the nonlinear reward distribution by approximating the likelihood with a Gaussian centered at the maximum likelihood estimate, then addresses nonlinear link functions by borrowing the covariance update structure from the exact linear case - this design ensures the approximate posteriors reduce to exact ones in linear settings while remaining computationally tractable for general nonlinear diffusion models.\n- The algorithm exploits the Markov structure of the diffusion model to avoid maintaining joint distributions over all action parameters, instead decomposing the problem into separate posterior updates for each latent layer and action - this reduces the computational burden from cubic in the total number of action parameters to linear in the number of actions plus layers.\n\nWeaknesses:\n\n- The theoretical analysis only covers linear diffusion models with linear-Gaussian rewards, yet the main algorithmic contribution involves two layers of approximations for nonlinear cases.\n\n- The approach relies on pre-training diffusion models on accurate offline estimates of action parameters, but obtaining such estimates for thousands of actions is precisely the challenging problem that online learning aims to solve - the paper doesn't address how to handle settings where offline data is biased, sparse, or completely unavailable for certain actions.\n\n- The experimental section shows performance degrades significantly when the true generative process doesn't match the diffusion prior, yet there's no theoretical characterization of robustness - the approximation scheme could amplify model misspecification errors through the hierarchical structure, but the paper doesn't analyze worst-case performance or provide guidance on detecting when the diffusion prior is misleading the algorithm.\n\nQUESTIONS:\n- Could you provide either theoretical analysis or of your two-layer approximation scheme? Specifically, I'd like to understand: (a) under what conditions the Gaussian approximation of the GLM likelihood is accurate, (b) how errors propagate through the hierarchical structure when using nonlinear link functions?\n\n- The paper claims computational advantages, but doesn't provide wall-clock time comparisons. Could you benchmark the actual runtime of dTS versus LinTS on problems with varying K and L? \n\n- The algorithm seems to require good offline parameter estimates to pre-train the diffusion model. How does dTS perform in a \"cold start\" scenario with minimal or no offline data?\n\nLIMITATIONS:\n- The significant gap between theoretical guarantees (linear case only) and practical algorithm (nonlinear approximations) should be prominently discussed.\n\n- The computational overhead for real-time decision-making systems needs acknowledgment.\n\n- The lack of formal guarantees on approximation quality and failure modes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n* This paper proposes a new contextual bandit algorithm that extends the literature on hierarchical Bayesian bandits to a more complex structure where the prior distribution is represented as a pre-trained diffusion model.\n\n* This paper considers the contextual bandit setting with prior represented as a diffusion model, and adopts Bayes regret to measure the performance. Then, it proposes a novel algorithm - Diffusion contextual Thompson sampling (DTS).\n    * The first key novelty is deriving an approximation for the posterior update. This approximation has good properties: (1) they match the prior with no data, and the prior's effect diminishes as data accumulates. (2) The approximate posterior is also a diffusion model with updated means and covariances.\n    * The second key contribution is to theoretically analyze the Bayes regret bound in the simplified setting where the known diffusion model's ODE is time-varying linear–Gaussian. The challenge of analysis the recursive update in the diffusion model. The analysis demonstrates the computational and statistical benefits of the proposed method, in contrast to LinTS [54]. The key insight is that the proposed method leverages the correlation between actions through the pre-trained diffusion model, while LinTS [54] has to learn each action independently.\n\n* Empirical analysis contains two experiments:\n  * When the prior matches the generative process, the paper compares the proposed method against LinUCB [1], LinTS [3], and HierTS [31] for the linear-reward setting, and against UCB-GLM [44] and GLM-TS [16] for the nonlinear-reward setting. The result shows that the proposed method achieves lower regret than the baselines. It demonstrates the power of a well-specified prior.\n  * When the prior does not match the generative process, such as Swiss-roll data and MovieLens data, the proposed method also outperformed LinTS.\n\nSTRENGTHS AND WEAKNESSES:\n* This paper solves a well-motivated problem and is well-written.\n* The algorithm is novel and sound. Both theoretical and empirical result supports the method.\n* One concern is that it seems that the result section focuses on comparing the proposed method vs LinTS (a baseline that does not exploit the correlation between action features). This is very effective in showing that the proposed method can exploit the action correlation, and it is actually effective and useful. However, it would be even stronger to also demonstrate that the proposed method is better than other hierarchical Bayesian bandit methods, to demonstrate that representing the prior as a diffusion model is a useful, not an overkill, in real-world scenarios. In this way, the two key takeaway messages are proved: (1) exploiting action correlation is useful (already somewhat demonstrated in prior literature), (2) representing prior distribution as a diffusion model is useful (very novel).\n    * To show this, I think there are a couple of ways. For example, for the real-world dataset - movielens, it would be helpful to benchmark the HierTS [31] method, to see how much modeling the prior as diffusion is beneficial in the real world. Alternatively, the author could take the setting of HierTS [31] and show that the diffusion-model prior allows better learning performance because the diffusion model is more expressive.\n\nQUESTIONS:\n* In Sec.4.1, the paper compares the proposed method vs LinTS [3]. I am not very familiar with [3], but it seems that the paper [3] (https://arxiv.org/pdf/1209.3352) 's Sec.2.1 assumes that there is a common reward parameter µ, which seems different from this paper's setting where each action `i' has its own parameter `θ_{*,i}`. I am curious whether you think it would still be appropriate to compare the proposed method with [3]?\n* Another question I have is that Fig. 2's nonlinear diffusion linear reward case, why not also benchmark the HierTS [31] method?\n* Out of curiosity, this paper focuses on correlations between actions. What about also exploiting correlations between contexts? E.g., similar users would have similar preferences about movies. Would the proposed method also extend to handle that?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper leverages diffusion models to tackle the challenge of decision-making in contextual bandits with large and complex action spaces. The authors propose a framework that uses pre-trained diffusion models as expressive priors within a Thompson Sampling (TS)-style bandit algorithm. Their proposed method enables efficient posterior approximation and sampling by combining offline-trained diffusion models with online updates during bandit online interactions. The authors develop practical implementations for posterior inference under this setup, and conduct experiments to demonstrate strong empirical performance across a range of contextual bandit scenarios. The authors provide theoretical analysis for the linear case, where the regret of the proposed method is bounded.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. This paper studies an important problem of contextual bandits. To the best of my knowledge, this is the first work to leverage diffusion models to enhance contextual bandits. The paper is well-motivated.\n2. The paper is well-written and easy to follow.\n2. The authors provide theoretical guarantees of the proposed method.\n3. Extensive experiments are conducted to show the outperformance of the proposed methods.\n\nWeaknesses:\n1. This paper considers the linear case, which is a slightly restrictive assumption in practice, but it is reasonable in the bandit literature.\n2. There are many works on diffusion models for reinforcement learning, though mostly from the empirical perspective. It would be better to discuss these works and make some comparisons. For example, see the related works in the survey \"Diffusion Models for Reinforcement Learning: A Survey\".\n3. The authors conducted extensive simulations, but they are only on synthetic and MovieLens datasets. It would be better to conduct experiments on more datasets to show the consistent improvement of the proposed methods over existing ones.\n\nQUESTIONS:\n1. How can the proposed method extend to more general nonlinear settings?\n2. How can the proposed method extend to the reinforcement learning settings?\n3. Your work proposes a TS-type algorithm. Do you think the idea of leveraging diffusion models to enhance contextual bandits can also apply to the UCB-type algorithmic approach?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}