Solving Empirical Bayes via Transformers

TL;DR

This paper introduces a transformer-based approach to solving the classical empirical Bayes problem for Poisson means, demonstrating both theoretical guarantees and practical advantages over traditional methods.

Contribution

It applies transformers to empirical Bayes, providing theoretical analysis and showing small models outperform classical algorithms in real-world and synthetic data.

Findings

Transformers achieve vanishing regret as dimension increases.

Small models outperform classical algorithms like NPMLE.

Transformers internally work differently from traditional estimators.

Abstract

This work applies modern AI tools (transformers) to solving one of the oldest statistical problems: Poisson means under empirical Bayes (Poisson-EB) setting. In Poisson-EB a high-dimensional mean vector $\theta$ (with iid coordinates sampled from an unknown prior $\pi$) is estimated on the basis of $X=\mathrm{Poisson}(\theta)$. A transformer model is pre-trained on a set of synthetically generated pairs $(X,\theta)$ and learns to do in-context learning (ICL) by adapting to unknown $\pi$. Theoretically, we show that a sufficiently wide transformer can achieve vanishing regret with respect to an oracle estimator who knows $\pi$ as dimension grows to infinity. Practically, we discover that already very small models (100k parameters) are able to outperform the best classical algorithm (non-parametric maximum likelihood, or NPMLE) both in runtime and validation loss, which we compute on out-of-distribution synthetic data as well as real-world datasets (NHL hockey, MLB baseball, BookCorpusOpen). Finally, by using linear probes, we confirm that the transformer's EB estimator appears to internally work differently from either NPMLE or Robbins' estimators.