Technical Summary: PFN-TS

Problem Statement

Contextual bandits address sequential decision-making under uncertainty, where a learner observes a context, selects an action, and receives a stochastic reward. The goal is to minimize cumulative regret over time. Thompson Sampling (TS) is a standard strategy that maintains a Bayesian posterior over reward functions and samples actions greedily from this posterior. However, the efficacy of TS is contingent on the quality of the underlying reward model.

In tabular applications (e.g., mobile health, recommendations), linear reward models are often misspecified because they fail to capture complex human behaviors like interactions, threshold effects, and heterogeneity. While richer alternatives exist (kernel methods, neural networks, tree ensembles), they suffer from scalability issues, high computational overhead, or rigid structural assumptions. Furthermore, these models typically require difficult-to-tune hyperparameters in online settings.

Prior-Data Fitted Networks (PFNs), such as TabPFN and TabICL, offer a promising alternative. These are transformer-based models trained to approximate Bayesian posterior predictive distributions (PPDs) across a broad prior of data-generating processes. They operate via in-context learning (ICL), requiring no parameter updates and no hyperparameter tuning. However, a critical gap exists: PFNs output PPDs (uncertainty over noisy future rewards), whereas TS requires uncertainty over the latent mean reward function. Directly sampling from the PPD does not equate to sampling from the posterior of the mean function, leading to suboptimal exploration.

Methodology: PFN-TS

The authors propose PFN-TS, a Thompson Sampling algorithm that bridges this gap by converting PFN posterior predictives into samples of the latent mean reward function. The methodology consists of three core components:

1. Subsampled Predictive Central Limit Theorem (SubCLT)

To recover the posterior distribution of the latent mean $f_0(x)$ from the PFN's predictive sequence, the authors leverage a predictive Central Limit Theorem (CLT).

• Theoretical Basis: Under specific assumptions (asymptotic variance, quasi-martingale conditions, and bounded moments), the tail of the predictive mean sequence converges to a Gaussian distribution centered at the latent mean.
• Efficiency Innovation: Traditional approaches to estimating the posterior variance require computing the predictive mean over all $n$ prefixes of the dataset, which is computationally prohibitive for online bandits ($O(n)$ forward passes). PFN-TS introduces a subsampled estimator that computes the variance using only $O(\log n)$ geometrically spaced prefixes ($t_0, t_1, \dots, t_J$).
• Variance Estimation: The variance is estimated using block increments $D_j = m_{t_j}(x) - m_{t_{j-1}}(x)$ weighted harmonically ($w_j$). The authors prove that this subsampled estimator is consistent ($\hat{V}_{s_n} \xrightarrow{p} V(x)$) and satisfies the CLT, allowing for a tractable Gaussian approximation $N(m_{s_n}(x), \hat{V}_{s_n}(x)/s_n)$ for Thompson sampling.

2. Caching and Computational Optimization

To further reduce the per-round cost, PFN-TS exploits the architecture of TabICL (the specific PFN used):

• KV-Cache Reuse: TabICL caches key-value attention matrices for a fixed training set. PFN-TS reuses these cached representations across rounds.
• Refresh Points: The algorithm only performs new forward passes at the geometric "refresh points" ($t_j$) defined by the SubCLT grid. Between refresh points, it uses the snapshot Gaussian approximation derived from the latest cached state. This reduces the complexity from $O(n)$ to $O(\log n)$ forward passes per arm per round.

3. Adaptive Encoding Strategy

The paper addresses the design choice of how to represent arm identity in the input:

• Disjoint Encoding: Maintains independent models per arm.
• Joint Encoding: Appends arm identity as a one-hot vector to the context, allowing information sharing across arms.
• Adaptive Selection: Since neither encoding universally dominates, PFN-TS runs both in parallel during a "dual-encoding" phase. It selects the active encoding based on the Continuous Ranked Probability Score (CRPS), a strictly proper scoring rule that measures calibration. The encoding with lower cumulative CRPS is retained, and the challenger is discarded after a pre-specified switch time.

Key Contributions

1. Algorithmic Bridge: PFN-TS is the first method to successfully adapt PFNs (which output PPDs) for Thompson Sampling by deriving a consistent estimator for the latent mean posterior variance via a subsampled CLT.
2. Theoretical Guarantees:
   • Consistency: Proves the consistency of the subsampled variance estimator under mild assumptions.
   • Regret Bound: Derives a Bayesian regret bound for PFN-TS. The bound decomposes into the regret of exact Thompson Sampling under the PFN's prior plus an error term dependent on the approximation quality of the PFN and the subsampling process.
3. Efficiency: Demonstrates that $O(\log n)$ forward passes are sufficient for consistent variance estimation, making PFN-based bandits feasible for online deployment.
4. Empirical Performance:
   • Synthetic Benchmarks: PFN-TS achieves the best average rank across nonlinear synthetic tasks (Friedman variants) and remains competitive on linear tasks.
   • OpenML Benchmarks: Outperforms baselines (LinTS, NeuralTS, BFTS, XGBoostTS) on six of eight classification-to-bandit datasets.
   • Real-World Application: In an offline evaluation of the Drink Less mobile health trial, PFN-TS attains the highest estimated policy value (SNIPS) at the final horizon, outperforming BFTS and other baselines.

Significance and Claims

The paper claims that PFN-TS demonstrates that tabular in-context learning can serve as an effective reward model for sequential decision-making once efficient uncertainty extraction and adaptive encoding are implemented.

• Model Flexibility: Unlike traditional models, PFN-TS requires no hyperparameter tuning and automatically adapts to the data structure (sparsity, smoothness, nonlinearity) inherent in the PFN's broad prior.
• Practical Viability: By reducing the computational bottleneck of uncertainty estimation from linear to logarithmic in the dataset size, the method makes the deployment of foundation models in online bandit settings practical.
• Robustness: The method is shown to be robust across diverse data-generating processes, performing well on both synthetic nonlinear tasks and real-world mobile health data.

The authors remain modest regarding limitations, noting that runtime is still a constraint due to repeated model evaluations, and that finite-sample calibration can be poor when arm histories are short or weakly informative. They suggest future work in regularized variance estimation and sharper regret bounds for the specific TabPFN/TabICL priors.