用Transformer解决泊松分布下的经验贝叶斯估计问题。
Solving Empirical Bayes via Transformers
- 用Transformer通过上下文学习适应未知先验分布
- 100k参数小模型在真实数据上优于经典算法
- 内部机制与传统方法不同，适合高维稀疏数据
本文将现代AI工具（Transformer）应用于一个经典的统计问题：在经验贝叶斯框架下对泊松分布的高维均值向量θ进行估计（泊松-EB）。观测值X服从参数为θ的泊松分布，θ的各分量独立同分布于未知先验π。我们使用合成数据预训练Transformer，使其具备在上下文学习（ICL）中自适应未知π的能力。理论上，足够宽的Transformer在维度趋于无穷时可实现与已知π的最优估计器相当的渐近无悔性。实践中，仅10万参数的小模型在运行时间与验证损失上均优于最佳传统算法（非参数最大似然估计，NPMLE），且在跨分布合成数据及真实数据集（NHL冰球、MLB棒球、BookCorpusOpen）上表现优异。通过线性探针分析，发现Transformer的贝叶斯估计器内部运作方式不同于NPMLE或Robbins估计器。
原文摘要 · Abstract (English)
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 $θ$ (with iid coordinates sampled from an unknown prior $π$) is estimated on the basis of $X=\mathrm{Poisson}(θ)$. A transformer model is pre-trained on a set of synthetically generated pairs $(X,θ)$ and learns to do in-context learning (ICL) by adapting to unknown $π$. Theoretically, we show that a sufficiently wide transformer can achieve vanishing regret with respect to an oracle estimator who knows $π$ 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。