# Contextual Thompson Sampling via Generation of Missing Data

arXiv: https://arxiv.org/abs/2502.07064 · HTML: https://mlsanity.com/paper/2502.07064/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 47 of its 81 references that are in the map.
- First posted on arXiv 10 Feb 2025.

## Builds on (most cited first, 30 of 47)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [A Tutorial on Thompson Sampling](https://mlsanity.com/paper/1707.02038/) (2017)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale](https://mlsanity.com/paper/2207.00032/) (2022)
- [Algorithmic decision making and the cost of fairness](https://mlsanity.com/paper/1701.08230/) (2017)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [An Information-Theoretic Analysis of Thompson Sampling](https://mlsanity.com/paper/1403.5341/) (2014)
- [Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Sampling](https://mlsanity.com/paper/1802.09127/) (2018)
- [Learning to Optimize via Information-Directed Sampling](https://mlsanity.com/paper/1403.5556/) (2014)
- [Bayesian Neural Networks: An Introduction and Survey](https://mlsanity.com/paper/2006.12024/) (2020)
- [Neural Thompson Sampling](https://mlsanity.com/paper/2010.00827/) (2020)
- [Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Modeling](https://mlsanity.com/paper/2207.04179/) (2022)
- [Corralling a Band of Bandit Algorithms](https://mlsanity.com/paper/1612.06246/) (2016)
- [Distilling Policy Distillation](https://mlsanity.com/paper/1902.02186/) (2019)
- [Meta-Thompson Sampling](https://mlsanity.com/paper/2102.06129/) (2021)
- [Ensemble Sampling](https://mlsanity.com/paper/1705.07347/) (2017)
- [Model selection for contextual bandits](https://mlsanity.com/paper/1906.00531/) (2019)
- [Meta-learning with Stochastic Linear Bandits](https://mlsanity.com/paper/2005.08531/) (2020)
- [Bayesian decision-making under misspecified priors with applications to meta-learning](https://mlsanity.com/paper/2107.01509/) (2021)
- [Bootstrapped Thompson Sampling and Deep Exploration](https://mlsanity.com/paper/1507.00300/) (2015)
- [Sequential Transfer in Multi-armed Bandit with Finite Set of Models](https://mlsanity.com/paper/1307.6887/) (2013)
- [Hypermodels for Exploration](https://mlsanity.com/paper/2006.07464/) (2020)
- [An Information-Theoretic Analysis for Thompson Sampling with Many Actions](https://mlsanity.com/paper/1805.11845/) (2018)
- [Approximate Thompson Sampling via Epistemic Neural Networks](https://mlsanity.com/paper/2302.09205/) (2023)

## Built on by (most cited first, 6 of 6)

- [Pre-Trained AI Model Assisted Online Decision-Making under Missing Covariates: A Theoretical Perspective](https://mlsanity.com/paper/2507.07852/) (2025)
- [PFN-TS: Thompson Sampling for Contextual Bandits via Prior-Data Fitted Networks](https://mlsanity.com/paper/2605.10137/) (2026)
- [OLIVIA: Online Learning via Inference-time Action Adaptation for Decision Making in LLM ReAct Agents](https://mlsanity.com/paper/2605.11169/) (2026)
- [Learning with Incomplete Context: Linear Contextual Bandits with Pretrained Imputation](https://mlsanity.com/paper/2510.09908/) (2025)
- [In-Context Learning for Data-Driven Censored Inventory Control](https://mlsanity.com/paper/2605.14840/) (2026)
- [Bootstrap-Conditioned Action Selection with Tabular Foundation Models](https://mlsanity.com/paper/2608.06559/) (2026)

## Abstract (arXiv)

We introduce a framework for Thompson sampling (TS) contextual bandit algorithms, in which the algorithm's ability to quantify uncertainty and make decisions depends on the quality of a generative model that is learned offline. Instead of viewing uncertainty in the environment as arising from unobservable latent parameters, our algorithm treats uncertainty as stemming from missing, but potentially observable outcomes (including both future and counterfactual outcomes). If these outcomes were all observed, one could simply make decisions using an "oracle" policy fit on the complete dataset. Inspired by this conceptualization, at each decision-time, our algorithm uses a generative model to probabilistically impute missing outcomes, fits a policy using the imputed complete dataset, and uses that policy to select the next action. We formally show that this algorithm is a generative formulation of TS and establish a state-of-the-art regret bound. Notably, our regret bound depends on the generative model only through the quality of its offline prediction loss, and applies to any method of fitting the "oracle" policy.