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