Solving Empirical Bayes via Transformers

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.

Overview of Solving Empirical Bayes via Transformers

The paper "Solving Empirical Bayes via Transformers" by Anzo Teh, Mark Jabbour, and Yury Polyanskiy presents a novel application of transformer models to the classical statistical problem of estimating Poisson means in the empirical Bayes framework (Poisson-EB). This paper addresses the challenge of leveraging modern deep learning models, particularly transformers, to improve upon traditional empirical Bayes estimators in terms of both computational efficiency and predictive performance.

Summary of Approach

The authors focus on the Poisson-EB problem where the objective is to estimate a high-dimensional mean vector $\theta$ from observed data $X \sim \text{Poisson}(\theta)$. The entries of $\theta$ are sampled i.i.d. from an unknown prior $\pi$. The proposed approach involves pre-training a transformer model on synthetic data pairs $(X, \theta)$, enabling the model to perform in-context learning (ICL) by adapting to different unknown priors $\pi$.

Theoretical Insights

The paper provides theoretical evidence showing that a sufficiently wide transformer can achieve vanishing regret compared to an oracle estimator with precise knowledge of $\pi$ as the problem dimension increases. The vanishing regret indicates that the transformer’s estimates approach the optimal Bayes estimator’s performance as dataset size grows.

Empirical Results

Empirically, small transformer models with approximately 100k parameters demonstrate the ability to outperform the best classical algorithm, the non-parametric maximum likelihood estimator (NPMLE), in both runtime efficiency and validation loss. The transformer-based approach is validated on out-of-distribution synthetic data and real-world datasets, including NHL hockey, MLB baseball, and BookCorpusOpen. Notably, the models perform exceptionally well in terms of runtime, showing near 100x speed improvements over NPMLE.

The paper further employs linear probes to test how the transformer's internal estimation process compares to traditional empirical Bayes estimators like Robbins and NPMLE. Findings suggest that the transformer implements a distinct mechanism, neither emulating Robbins nor NPMLE.

Implications and Future Directions

The development of this transformer-based approach for solving Poisson-EB has both practical and theoretical implications. Practically, the method offers a new tool for statisticians and data scientists dealing with empirical Bayes problems, providing faster and potentially more accurate estimates. Theoretically, it adds to the understanding of how deep learning models, specifically transformers, can be utilized in traditional statistical frameworks.

In terms of future directions, the paper hints at several areas for potential development. One critical area is extending the approach to handle multi-dimensional inputs, which would significantly enhance the applicability of the method across different domains. Furthermore, understanding the limitations and capabilities of transformers in approximating sophisticated function classes in statistical tasks continues to be an area ripe for exploration.

In conclusion, this paper contributes a compelling case for integrating transformers into empirical Bayes methods, demonstrating significant advancements in both efficiency and efficacy. The insights gleaned here pave the way for broader applications of transformers in statistical learning and beyond, offering a promising intersection of modern AI and classical statistical inference.