Research map / Bandits and online learning / Contextual bandits and Thompson sampling
Diffusion Models Meet Contextual Bandits
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes diffusion Thompson sampling (DiffTS), using a pre-trained diffusion model as a hierarchical prior over bandit action parameters, with closed-form-inspired Gaussian posterior approximations for contextual bandits.
Problem
Online decision-making in contextual bandits with many actions is computationally expensive or statistically inefficient when no informative prior capturing action correlations is available.
Approach
Use a pre-trained diffusion model as the prior over action parameters and approximate the hierarchical posterior recursively: Laplace-style Gaussian likelihood approximation plus layer-wise replacement of linear mixing terms with nonlinear link functions, yielding fast posterior updates and sampling. Exact in the linear-Gaussian case.
Key results
- DiffTS outperforms LinUCB, LinTS, HierTS, UCB-GLM and GLM-TS across synthetic linear and nonlinear diffusion/reward settings with d in {5,20}, K up to 10^4 and n=5000.
- The performance gap between DiffTS and LinTS widens as the number of actions K grows (evaluated up to K=5x10^4).
- In Swiss roll experiments, DiffTS outperforms LinTS by a factor of 1.5 with just 50 pre-training samples.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use diffusion Thompson sampling instead of the evaluated baselines | Diffusion Thompson sampling · Synthetic contextual bandits, all reported settings | cumulative regret | Helped |
| use the correct diffusion prior but a linear-Gaussian reward model instead of the nonlinear reward model (ablation) | Diffusion Thompson sampling · Synthetic contextual bandits with nonlinear rewards | cumulative regret | Helped |
| use a diffusion prior rather than LinTS | Diffusion Thompson sampling · Synthetic linear diffusion and linear rewards, varying action count K | final cumulative regret | Helped |
| increase the number of actions K (ablation) | Diffusion Thompson sampling · Synthetic linear diffusion and linear rewards | cumulative regret: K ∈ {10, 100, 500, 1000} | Hurt |
| increase context dimension d (ablation) | Diffusion Thompson sampling · Synthetic linear diffusion and linear rewards | cumulative regret: d ∈ {5, 10, 15, 20} | Hurt |
| increase diffusion depth L (ablation) | Diffusion Thompson sampling · Synthetic linear diffusion and linear rewards | cumulative regret: L ∈ {2, 4, 5, 6} | Hurt |
| increase diffusion-prior misspecification (ablation) | Diffusion Thompson sampling · Synthetic linear diffusion and linear rewards | cumulative regret: v ∈ {0.5, 1, 1.5} | Hurt |
| increase diffusion pretraining sample size (ablation) | Diffusion Thompson sampling · Swiss roll data | cumulative regret: 10 samples; 50 samples; 1.5 | Helped |
| increase diffusion depth toward approximately 40, then increase it beyond that (ablation) | Diffusion Thompson sampling · Swiss roll data | cumulative regret: L ≈ 40 | Mixed |
| use diffusion Thompson sampling instead of LinTS | Diffusion Thompson sampling · MovieLens | regret | Helped |
Surprising findings
- Using the correct latent diffusion structure with a misspecified (linear-Gaussian) reward model outperforms models with the correct reward distribution but ignoring latent structure.
Limitations
- Regret theory only holds for the linear-Gaussian case; no formal guarantees for nonlinear diffusion or the Laplace-based likelihood approximation.
- Relies on a pre-trained diffusion prior; with scarce or biased offline data the prior may be poorly fit.
- Regret still depends on K in the per-action setting.
Experimental setup
Datasets: synthetic bandits from diffusion model Swiss roll distribution samples MovieLens CIFAR-10
Benchmarks: synthetic contextual bandits (linear and nonlinear diffusion/reward) MovieLens semi-synthetic bandit Swiss roll prior bandit
Metrics: cumulative regret Bayes regret
Baselines: LinUCB LinTS HierTS UCB-GLM GLM-TS GLM-UCB
Builds on: Thompson sampling LinTS HierTS Laplace approximation GLM bandits
Closest papers by meaning
- Online Posterior Sampling with a Diffusion Prior (NeurIPS 2024 · similarity 0.92)
- Thompson Sampling with Diffusion Generative Prior (ICML 2023 · similarity 0.90)
- Contextual Bandits with Large Action Spaces: Made Practical (ICML 2022 · similarity 0.90)
- Stochastic Bandits with Context Distributions (NeurIPS 2019 · similarity 0.90)
- Active Context Selection Improves Simple Regret in Contextual Bandits (NeurIPS 2026 · similarity 0.89)
- Efficient Contextual Bandits with Uninformed Feedback Graphs (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.