PFN-TS: Thompson Sampling for Contextual Bandits via Prior-Data Fitted Networks
Not on the map. Why?
Abstract (arXiv)
Thompson sampling is a widely used strategy for contextual bandits: at each round, it samples a reward function from a Bayesian posterior and acts greedily under that sample. Prior-data fitted networks (PFNs), such as TabPFN v2+ and TabICL v2, are attractive candidates for this purpose because they approximate Bayesian posterior predictive distributions in a single forward pass. However, PFNs predict noisy future rewards, while Thompson sampling requires uncertainty over the latent mean reward function. We propose PFN-TS, a Thompson sampling algorithm that converts PFN posterior predictives into mean-reward samples using a subsampled predictive central limit theorem. The method estimates posterior variance from a geometric grid of O(log n) dataset prefixes rather than the full O(n) predictive sequence used in previous predictive-sequence approaches, and reuses TabICL's cached representations across rounds. We prove consistency of the subsampled variance estimator and give a Bayesian regret bound that decomposes PFN-TS regret into exact posterior-sampling regret under the PFN prior plus approximation terms. Empirically, PFN-TS achieves the best average rank across nonlinear synthetic and OpenML classification-to-bandit benchmarks, remains competitive on linear and BART-generated rewards, and attains the highest estimated policy value in an offline mobile-health evaluation. Code is available at https://anonymous.4open.science/r/PFNTS-36ED/.
Key facts. Cited by 2 papers in the map (data through 14 Oct 2026). Builds on 15 of its 30 references that are in the map. First posted on arXiv 11 May 2026.
Builds on
15 of 30 references.
References · 15
- 2022TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Secondcited by 674
- 2014OpenML: networked science in machine learningcited by 526
- 2017A Tutorial on Thompson Samplingcited by 433
- 2021Transformers Can Do Bayesian Inferencecited by 360
- 2017On Kernelized Multi-armed Banditscited by 352
- 2014An Information-Theoretic Analysis of Thompson Samplingcited by 178
- 2018Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Samplingcited by 169
- 2020On Information Gain and Regret Bounds in Gaussian Process Banditscited by 158
- 2026TabICLv2: A better, faster, scalable, and open tabular foundation modelcited by 144
- 2025TabPFN: One Model to Rule Them All?cited by 14
- 2025TabMGP: Martingale Posterior with TabPFNcited by 9
- 2025Contextual Thompson Sampling via Generation of Missing Datacited by 6
3 more
Built on by
Citing papers per year: 2026: 2.
2 papers.