PFN-TS: Thompson Sampling for Contextual Bandits via Prior-Data Fitted Networks
Authors: Yan Shuo Tan, Kenyon Ng, Ruizhe Deng et al.
Summary
The authors develop PFN-TS, a Thompson sampling algorithm for contextual bandits that uses prior-data fitted networks (PFNs) like TabPFN and TabICL to approximate Bayesian posteriors in one forward pass. The challenge is that PFNs predict noisy rewards, but Thompson sampling needs uncertainty over the mean reward function. They solve this by estimating posterior variance from a subsampled predictive sequence (a logarithmic grid of dataset prefixes instead of the full sequence), then sampling mean rewards via a central limit theorem. They prove the variance estimator is consistent, bound the regret, and show strong empirical results across synthetic, OpenML, and mobile-health benchmarks.
Main takeaways:
- Converts PFN posterior predictives (which model noisy rewards) into samples of the mean reward function using a subsampled predictive CLT
- Estimates posterior variance from O(log n) dataset prefixes rather than the full O(n) sequence, reusing cached representations for efficiency
- Provides a regret bound decomposing error into exact posterior-sampling regret under the PFN prior plus approximation terms
- Achieves best average rank across nonlinear benchmarks and high estimated policy value in an offline mobile-health task
Relevance
Not directly related to my persona/midtraining work—this is about Thompson sampling for contextual bandits with PFN-based Bayesian approximation. Tangential; only relevant if I needed in-context bandit methods for adaptive persona selection, but that's a stretch.
Abstract
arXiv:2605.10137v1 Announce Type: new Abstract: 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/PFN_TS-36ED/.