stat.MLMay 11, 2026

PFN-TS: Thompson Sampling for Contextual Bandits via Prior-Data Fitted Networks

Authors: Yan Shuo Tan, Kenyon Ng, Ruizhe Deng, Sumetha Loganathan, Qiong Zhang, Bibhas Chakraborty

Organizations: National University of Singapore · Monash University · Renmin University of China · Duke-NUS Medical School

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)O(\log n) dataset prefixes rather than the full O(n)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/.

Explore similar work

Jun 22, 2026eess.SY

Flow-Corrected Thompson Sampling for Non-Stationary Contextual Bandits

We study non-stationary linear contextual bandits where the reward model drifts over time, rendering classical contextual bandit algorithms brittle because historical data becomes systematically biased. We propose Flow-Corrected Thompson Sampling (fcTS), a Bayesian method that reuses experience by transporting past rewards to the present using an explicit drift model and incorporating each transported observation with a confidence weight that reflects transport reliability. This yields a unified template that specializes in (i) linear parameter drift via online slope estimation and reward correction, (ii) periodic variation via phase-aware reuse across cycles, and (iii) recurring regime switches via changepoint detection and regime-specific posterior memory. The resulting posterior updates remain closed-form under a linear Gaussian model and can be implemented efficiently with truncated, incrementally updated sufficient statistics. Across five controlled case studies and a semi-synthetic portfolio-selection benchmark with multiple overlapping non-stationarities, fcTS outperforms standard forgetting-based baselines (discounting, sliding windows, and periodic restarts), with the largest gains in settings exhibiting recurring temporal structure. These results demonstrate that when non-stationarity is structured, correcting and reweighting historical observations can be substantially more sample-efficient than uniformly discarding them.
AmirHossein Naghdi, Ali Baheri
Sep 12, 2023stat.ML

Generalized Regret Analysis of Thompson Sampling using Fractional Posteriors

Thompson sampling (TS) is one of the most popular and earliest algorithms to solve stochastic multi-armed bandit problems. We consider a variant of TS, named αα-TS, where we use a fractional or αα-posterior (α∈(0,1)α\in(0,1)) instead of the standard posterior distribution. To compute an αα-posterior, the likelihood in the definition of the standard posterior is tempered with a factor αα. For αα-TS we obtain both instance-dependent O(∑k≠i∗Δk(log⁡(T)C(α)Δk2+12))\mathcal{O}\left(\sum_{k \neq i^*} Δ_k\left(\frac{\log(T)}{C(α)Δ_k^2} + \frac{1}{2} \right)\right) and instance-independent O(KTlog⁡K)\mathcal{O}(\sqrt{KT\log K}) frequentist regret bounds under very mild conditions on the prior and reward distributions, where ΔkΔ_k is the gap between the true mean rewards of the kthk^{th} and the best arms, and C(α)C(α) is a known constant. Both the sub-Gaussian and exponential family models satisfy our general conditions on the reward distribution. Our conditions on the prior distribution can be easily satisfied by a density that is positive, continuous, and bounded. We also establish another instance-dependent regret upper bound that matches (up to constants) to that of improved UCB [Auer and Ortner, 2010]. Our regret analysis carefully adapts and combines recent theoretical developments in the non-asymptotic concentration analysis and Bernstein-von Mises type results for the αα-posterior distribution. Moreover, our analysis does not require additional structural properties such as closed-form posteriors or conjugate priors.
Prateek Jaiswal, Debdeep Pati, Anirban Bhattacharya +1
Jun 19, 2026cs.LG

AdaPrivate-TS: Private Thompson Sampling for Contextual Bandits with Privacy Amplification

We present AdaPrivate-TS, a differentially private contextual bandit algorithm that combines Thompson Sampling with batched zCDP composition. Our key insight is that differential privacy noise inflates the posterior covariance in a structured way: adding Gaussian noise N(0,σ2I)N(0,σ^2 I) to bb yields sampling covariance v2A−1+σ2A−2v^2 A^{-1} + σ^2 A^{-2}, which Thompson Sampling interprets as increased uncertainty rather than pure corruption. Under event-level privacy (protecting individual interactions) with stochastic contexts, we prove that the privacy cost is only O(d log⁡T/ρ)O(\sqrt{d}\,\log T/\sqrtρ), logarithmic in TT, because parallel composition amortizes noise across batches. Additionally, we explore privacy amplification via Poisson subsampling, which can reduce effective noise at stringent privacy budgets. Experiments on synthetic and real-world datasets demonstrate: (1) AdaPrivate-TS achieves 93-99% of non-private performance at ε∈[0.5,5]\varepsilon \in [0.5, 5], outperforming UCB by 0.5-3.7% and up to 18% with tuned adaptive exploration at extreme ε\varepsilon; (2) privacy amplification provides additional 2-5% gains at low ε\varepsilon; (3) on MovieLens and Jester, AdaPrivate-TS achieves the best overall performance among event-level baselines, dominating at ε≥2\varepsilon \geq 2; (4) under DP-SVD private features, TS's advantage over UCB grows to +11%, confirming noise-as-uncertainty is not limited to reward privacy. We provide rigorous proofs for privacy guarantees under interactive zCDP composition and comprehensive evaluation including convergence curves, 12-seed CIs, and DP-SVD feature ablation.
Mohammadreza Riyazat, Eranga Ukwatta