MIT三人团队：用Transformer解决经验贝叶斯问题，比经典方法快100倍

Transformer 很成功，更一般而言，我们甚至可以将（仅编码器）Transformer 视为学习可交换数据的通用引擎。由于大多数经典的统计学任务都是基于独立同分布（iid）采用假设构建的，因此很自然可以尝试将 Transformer 用于它们。

针对经典统计问题训练 Transformer 的好处有两个：

• 可以得到更好的估计器；

• 可以在一个有比 NLP 更加容易和更好理解的统计结构的领域中阐释 Transformer 的工作机制。

近日，MIT 的三位研究者 Anzo Teh、Mark Jabbour 和 Yury Polyanskiy 宣称找到了一个可以满足这种需求 「可能存在的最简单的这类统计任务」，即 empirical Bayes (EB) mean estimation（经验贝叶斯均值估计）。

• 论文标题：Solving Empirical Bayes via Transformers

• 论文地址：https://arxiv.org/pdf/2502.09844

该团队表示：「我们认为 Transformer 适用于 EB，因为 EB 估计器会自然表现出收缩效应（即让均值估计偏向先验的最近模式），而 Transformer 也是如此，注意力机制会倾向于关注聚类 token。」对注意力机制的相关研究可参阅论文《The emergence of clusters in self-attention dynamics》。

此外，该团队还发现，EB 均值估计问题具有置换不变性，无需位置编码。

另一方面，人们非常需要这一问题的估计器，但麻烦的是最好的经典估计器（非参数最大似然 / NPMLE）也存在收敛速度缓慢的问题。

MIT 这个三人团队的研究表明 Transformer 不仅性能表现胜过 NPMLE，同时还能以其近 100 倍的速度运行！

总之，本文证明了即使对于经典的统计问题，Transformer 也提供了一种优秀的替代方案（在运行时间和性能方面）。对于简单的 1D 泊松 - EB 任务，本文还发现，即使是参数规模非常小的 Transformer（< 10 万参数）也能表现出色。

定义 EB 任务

泊松 - EB 任务：通过一个两步式过程以独立同分布（iid）方式生成 n 个样本 X_1, . . . , X_n.

第一步，从某个位于实数域 ℝ 的未知先验 π 采样 θ_1, . . . , θ_n。这里的 π 的作用是作为一个未曾见过的（非参数）隐变量，并且对其不做任何假设（设置没有连续性和平滑性假设）。

第二步，给定 θ_i，通过 X_i ∼ Poi (θ_i) 以 iid 方式有条件地对 X_i 进行采样。

这里的目标是根据看到的 X_1, . . . , X_n，通过 估计 θ_1, . . . , θ_n，以最小化期望的均方误差（MSE） 。如果 π 是已知的，则这个最小化该 MSE 的贝叶斯估计器便是 θ 的后验均值，其形式如下：

其中 是 x 的后验密度。由于 π 是未知的，于是估计器 π 只能近似 。这里该团队的做法是将估计器的质量量化为后悔值，定义成了 多于 的 MSE：

通过 Transformer 求解泊松 - EB

简单来说，该团队求解泊松 - EB 的方式如下：首先，生成合成数据并使用这些数据训练 Transformer；然后，冻结它们的权重并提供要估计的新数据。

该团队表示，这应该是首个使用神经网络模型来估计经验贝叶斯的研究工作。

理解 Transformer 是如何工作的

论文第四章试图解释 Transformer 是如何工作的，并从两个角度来实现这一目标。首先，他们建立了关于 Transformer 在解决经验贝叶斯任务中的表达能力的理论结果。其次，他们使用线性探针来研究 Transformer 的预测机制。

本文从 clipped Robbins 估计器开始，其定义如下：

得出：transformer 可以学习到任意精度的 clipped Robbins 估计器。即：

类似地，本文证明了 transformer 还可以近似 NPMLE。即：

完整的证明过程在附录 B 中，论文正文只提供了一个大致的概述。

接下来，研究者探讨了 Transformer 模型是如何学习的。他们通过线性探针（linear probe）技术来研究 Transformer 学习机制。

这项研究的目的是要了解 Transformer 模型是否像 Robbins 估计或 NPMLE 那样工作。图 1 中的结果显示，Transformer 模型不仅仅是学习这些特征，而是在学习贝叶斯估计器 是什么。

总结而言，本章证明了 Transformer 可以近似 Robbins 估计器和 NPMLE（非参数最大似然估计器）。

此外，本文还使用线性探针（linear probes）来证明，经过预训练的 Transformer 的工作方式与上述两种估计器不同。

合成数据实验与真实数据实验

表 1 为模型参数设置，本文选取了两个模型，并根据层数将它们命名为 T18 和 T24，两个模型都大约有 25.6k 个参数。此外，本文还定义了 T18r 和 T24r 两个模型。

在这个实验中，本文评估了 Transformer 适应不同序列长度的能力。图 2 报告了 4096 个先验的平均后悔值。

图 6 显示 transformer 的运行时间与 ERM 的运行时间相当。

合成实验的一个重要意义在于，Transformer 展示了长度泛化能力：即使在未见过的先验分布上，当测试序列长度达到训练长度的 4 倍时，它们仍能实现更低的后悔值。这一点尤为重要，因为多项研究表明 Transformer 在长度泛化方面的表现参差不齐 [ZAC+24, WJW+24, KPNR+24, AWA+22]。

最后，本文还在真实数据集上对这些 Transformer 模型进行了评估，以完成类似的预测任务，结果表明它们通常优于经典基线方法，并且在速度方面大幅领先。

从表 3 可以看出，在大多数数据集中，Transformer 比传统方法有显著的改进。

总之，本文证明了 Transformer 能够通过上下文学习（in-context learning）掌握 EB - 泊松问题。实验过程中，作者展示了随着序列长度的增加，Transformer 能够实现后悔值的下降。在真实数据集上，本文证明了这些预训练的 Transformer 在大多数情况下能够超越经典基线方法。