Transformer作为贝叶斯上下文实验者：平滑自适应的高效ATE估计

# 1 引言 来源：https://arxiv.org/html/2606.31184 \\RRHSecondLine\\LRHSecondLine \\RUNTITLE Transformer作为贝叶斯上下文实验者\\TITLETransformer作为贝叶斯上下文实验者：平滑自适应的高效ATE估计 \\RUNAUTHOR Li 和 Simchi-Levi\\ARTICLEAUTHORS\\AUTHORJiachun Li\\AFF麻省理工学院信息与决策系统实验室，jiach334@mit\.edu (https://arxiv.org/html/2606.31184v1/[email protected])\\AUTHORDavid Simchi-Levi\\AFF麻省理工学院信息与决策系统实验室，dslevi@mit\.edu (https://arxiv.org/html/2606.31184v1/[email protected]) ###### 摘要 平均处理效应(ATE)的自适应实验需要在有效推断与统计效率之间取得平衡的随机化分配。Oracle设计是一种依赖于协变量的Neyman规则，其由未知的手臂条件结果方差决定。我们研究这个序贯方差估计与分配过程能否通过上下文学习实现摊销。我们引入贝叶斯上下文实验者：一种经过训练以模仿贝叶斯后验Neyman教师的Transformer策略。该教师利用实验历史更新非参数信念，并据此分配后验Neyman处理概率。此设计收敛于Oracle规则，支持高效的ATE推断。Transformer通过基于注意力的充分统计量和投影梯度下降，以构造性方式实现这一映射，模拟具有高斯级数先验的贝叶斯更新。为解决未知结果平滑度的问题，我们通过专家混合Transformer (Mixture-of-Experts) 结合了按平滑度索引的实验者。其门控充当平滑度类上的分层后验，集中于接近Oracle的专家。通过限制Transformer类的复杂度，我们证明可以利用监督预训练通过经验风险最小化学习这一摊销策略。实验证实了准确的教师模仿、自适应分配以及相对于基线的ATE精度提升。自适应实验广泛应用于临床试验和在线平台，以提高平均处理效应(ATE)估计的精度 (Dai et al., 2023 (https://arxiv.org/html/2606.31184#bib.bib8); Hu and Rosenberger, 2006 (https://arxiv.org/html/2606.31184#bib.bib12))。贝叶斯实验设计在这些场景中尤为自然，因为它将先验领域知识和历史数据纳入序贯决策 (Chaloner and Verdinelli, 1995 (https://arxiv.org/html/2606.31184#bib.bib6); Berry et al., 2010 (https://arxiv.org/html/2606.31184#bib.bib4))。对于ATE推断，最优规则是由未知的手臂条件残差方差确定的Oracle协变量依赖性Neyman分配 (Li et al., 2024 (https://arxiv.org/html/2606.31184#bib.bib17); Zhao, 2023 (https://arxiv.org/html/2606.31184#bib.bib30); Armstrong, 2022 (https://arxiv.org/html/2606.31184#bib.bib1); Kato et al., 2020 (https://arxiv.org/html/2606.31184#bib.bib15))。一个正确设定的贝叶斯实验者可以通过更新关于潜在结果的后验信念来恢复此规则：结果矩的后验收缩驱使序贯分配概率趋近于Oracle Neyman设计 (Dai et al., 2023 (https://arxiv.org/html/2606.31184#bib.bib8); Cook et al., 2024 (https://arxiv.org/html/2606.31184#bib.bib7); Li et al., 2024 (https://arxiv.org/html/2606.31184#bib.bib17))。实现这一贝叶斯更新与设计过程需要准确估计手臂条件结果矩，以及将后验不确定性高效转换为处理分配概率。在高维协变量空间中，主要瓶颈在于指定潜在结果模型的平滑度 (van der Vaart and van Zanten, 2009 (https://arxiv.org/html/2606.31184#bib.bib27); van Waaij and van Zanten, 2017 (https://arxiv.org/html/2606.31184#bib.bib28))。过度平滑的设定会导致偏差和错误的分配，而过度粗糙的设定则收敛缓慢。因此，适应未知平滑度对于后验准确性和实验效率至关重要 (Szabó et al., 2015 (https://arxiv.org/html/2606.31184#bib.bib21))。在部署层面，贝叶斯自适应设计需要一个自动化系统，能将传入的实验历史直接映射为统计上有原则的分配概率，而无须针对每个实验进行工程调整 (Tian et al., 2021 (https://arxiv.org/html/2606.31184#bib.bib22))。最近的研究表明，预训练的Transformer可以作为上下文贝叶斯学习器，从历史中推断潜在结构而无需更新参数 (Bai et al., 2023 (https://arxiv.org/html/2606.31184#bib.bib2); Lin et al., 2023 (https://arxiv.org/html/2606.31184#bib.bib18); Lee et al., 2023 (https://arxiv.org/html/2606.31184#bib.bib16))。我们将这一思想扩展到因果实验设计，提出Transformer不是作为处理效应估计器，而是作为摊销的贝叶斯实验者。通过平衡偏差与方差，所需的非参数贝叶斯更新可以通过动态有效维度掩码进行截断，将难以处理的推断问题简化为有限秩掩码岭回归。该更新通过基于注意力的充分统计量和投影梯度下降块以构造性方式实现。为处理模型选择，我们采用专家混合(MoE)架构，其门控近似于平滑度类上的分层贝叶斯后验，并随着数据累积集中于适当的复杂度尺度。这些机制共同使Transformer能够学习从历史到倾向得分的映射，并适应潜在结果的未知复杂度。尽管我们的构造表明Transformer可以代表贝叶斯实验者，但实际模型是通过监督预训练而非手动工程设计来学习的。因此，我们分析了有限预训练轨迹上的监督经验风险最小化(ERM) (Huang et al., 2025 (https://arxiv.org/html/2606.31184#bib.bib14); Hu et al., 2025 (https://arxiv.org/html/2606.31184#bib.bib13); Lin et al., 2023 (https://arxiv.org/html/2606.31184#bib.bib18))。我们的理论给出了一个算法模仿保证：训练后的Transformer模仿了教师的从历史到倾向得分的映射行为，误差在近似误差和估计误差之内。重要的是，学习到的模型不需要恢复精确的后验状态或显式执行Neyman分配；预训练使其可以直接从实验历史中学习设计映射。我们的贡献如下：贝叶斯设计与平滑度适应。我们定义并分析了一个贝叶斯后验Neyman教师，用于面向高效ATE的自适应实验。该教师维护手臂条件均值和二阶矩的非参数贝叶斯信念，将后验残差方差估计转换为已知的分配概率，以实现有效推断。我们进一步引入了平滑度类上的分层贝叶斯扩展，并证明由此产生的设计适应未知结果复杂度，实现了接近Oracle的收缩率和最优效率。Transformer的表达能力与泛化性。我们证明了贝叶斯更新与设计映射可以由Transformer架构表示和学习。构造上，基于注意力的充分统计量和投影梯度下降块实现了有限秩掩码岭回归更新，而MoE门控通过分层后验实现了平滑度适应。然后我们建立了一个ERM泛化理论，表明在有限轨迹上的监督预训练使得学习到的Transformer能够模仿教师的从历史到倾向得分的映射。实证验证。在合成自适应轨迹上，一个没有显式MoE头的直接估计Transformer在七个未见过的平滑度水平上实现了平滑自适应极小化最优速率。训练后的设计Transformer再现了贝叶斯-Neyman分配行为，包括在线部署时与贝叶斯教师匹配的步骤级波动模式，而无需显式学习结果矩。端到端地，它将AIPW ATE估计的MSE相对于均匀随机化降低，并接近Oracle Neyman分配。 ## 2 问题设定与贝叶斯统计保证 ### 2\.1 序贯实验设定 我们考虑一个有序贯协变量的双手臂实验。在每一轮 $t=1,\ldots,n$，一个带有协变量的单元到达 $X_t \sim P_X$ 在 $\mathcal{X} \subset [0,1]^d$ 上，具有有界、连续的密度。基于历史 $\mathcal{H}_{t-1} = \{(X_s, W_s, Y_s): s < t\}$，实验者选择处理分配概率 $e_t(x) = \mathbb{P}(W_t = 1 \mid X_t = x, \mathcal{H}_{t-1})$，其中 $e_t(x) \in [\eta, 1-\eta]$ 用于某些固定的 $\eta \in (0, 1/2)$。剪辑确保在推断中处理和控制组都有重叠。然后观察到结果 $Y_t$。对于每个 $w \in \{0,1\}$，令 $\mu_w^\star(x) = \mathbb{E}[Y_t(w) \mid X_t = x]$ 和 $q_w^\star(x) = \mathbb{E}[Y_t(w)^2 \mid X_t = x]$ 为条件矩函数。残差方差为 $v_w^\star(x) = q_w^\star(x) - \mu_w^\star(x)^2$。我们假设 $v_w^\star$ 是上下有界的：$c_V \leq v_w^\star(x) \leq C_V$，对于所有 $w$ 和 $x$，其中 $c_V > 0$。 ### 2\.2 先验与似然设定 在实验开始前，教师在一个固定的光滑度超参数 $\beta$ 上放置一个先验。对于一个候选的 $\beta>0$，它假设每个 $\mu_w^\star$ 和 $q_w^\star$ 是 $\beta$ 次平滑函数。更精确地说，我们采用高斯级数先验。令 $\{\psi_j\}_{j=1}^\infty$ 为 $\mathcal{X}$ 上关于 $L_2(P_X)$ 的标准正交基，例如三角多项式张量积。**假设 2** (Sobolev 平滑度)。对于 $w\in\{0,1\}$，存在 $0<\beta_{w,\mu}^\star, \beta_{w,q}^\star<\infty$，使得系数 $\mu_{w}^\star$, $q_{w}^\star$ 在上述基中以速率 $j^{-\beta_{w,\mu}^\star/d}$, $j^{-\beta_{w,q}^\star/d}$ 衰减。噪声 $\xi_{t,\mu,w} = Y_t - \mu_w^\star(X_t)$, $\xi_{t,q,w} = Y_t^2 - q_w^\star(X_t)$ 条件于子指数，均匀地关于 $t,w$。对于固定的 $\beta>0$，在 $\mu_w$ 和 $q_w$ 上放置独立的高斯级数先验。例如，
$$\mu_w(x) = \sum_{j=1}^\infty a_{w,\mu,j} \psi_j(x), \quad a_{w,\mu,j} \overset{\mathrm{ind}}{\sim} N(0, \rho_j(\beta)),$$
并且类似地对于 $q_w$，其中 $\rho_j(\beta) \asymp j^{-1-2\beta/d}$。改变 $\beta$ 会改变分配给高频坐标的先验方差。较小的 $\beta$ 允许更粗糙的函数，而较大的 $\beta$ 会更激进地压缩高频坐标。给定手臂 $w$ 的观测值，教师通过简化的高斯模型更新其均值信念：
$$Y_t \mid X_t, W_t = w, \mu_w \overset{\mathrm{model}}{\sim} N(\mu_w(X_t), \bar{\sigma}_\mu^2),$$
其中 $\bar{\sigma}_\mu^2>0$ 是用于更新的用户选择的调整常数，而非真实的噪声方差。这给出：
$$d\Pi_{\mu,w,t}^{(\beta)}(\mu) \propto \exp\left\{-\frac{1}{2\bar{\sigma}_\mu^2} \sum_{s \leq t: W_s = w} \bigl(Y_s - \mu(X_s)\bigr)^2\right\} d\Pi_{\mu,w,0}^{(\beta)}(\mu).$$
二阶矩信念的更新类似，只需将 $Y_s$, $\mu$, $\bar{\sigma}_\mu^2$ 替换为 $Y_s^2$, $q$, $\bar{\sigma}_q^2$。虽然这些高斯模型可能被错误设定，但关键点在于其平方损失总体目标仍然是正确的条件矩 $\mu_w^\star$ 和 $q_w^\star$；在假设 2 下，由此产生的信念会以通常的非参数速率收缩到这些目标。 ### 2\.3 后验收缩 令 $N_{t,w}$ 为截至第 $t$ 期的手臂特定样本量。剪辑设计意味着 $N_{t,w} \asymp t$ 以高概率成立。由于 $e_t$ 是可预测的，手臂 $w$ 的观测值形成了一个自适应回归样本，具有鞅差噪声和远离零的重叠。 **命题 1** (矩后验收缩)。假设假设 1 和 2 成立。考虑任意具有 $e_t(x) \in [\eta, 1-\eta]$ 的可预测自适应设计。如果 (2.2) 中的先验平滑度与相应的真实平滑度匹配，那么对于足够大的常数 $M$，
$$\Pi_{\mu,w,t}^{(\beta)}\left(\|\mu_w - \mu_w^\star\|_{L_2(P_X)} > M \epsilon_{\mu,w,t} \,\middle|\, \mathcal{H}_t\right) \to 0, \quad \Pi_{q,w,t}^{(\beta)}\left(\|q_w - q_w^\star\|_{L_2(P_X)} > M \epsilon_{q,w,t} \,\middle|\, \mathcal{H}_t\right) \to 0,$$
其中 $\epsilon_{\mu,w,t} \asymp N_{t,w}^{-\beta_{w,\mu}^\star / (2\beta_{w,\mu}^\star + d)}$, $\epsilon_{q,w,t} \asymp N_{t,w}^{-\beta_{w,q}^\star / (2\beta_{w,q}^\star + d)}$ 是非参数收敛速率。换句话说，后验信念将集中在局部球 $\{\mu_\omega: \|\mu_\omega - \mu_\omega^*\|_{L_2} \lesssim \epsilon_{\mu,\omega,t}\}$ 内。证明见附录 8。它基于可预测自适应采样下平方损失贝叶斯信念的标准高斯级数后验收缩论证。非高斯和异方差噪声不影响收缩速率，只要子指数尾部条件成立。 ### 2\.4 后验 Neyman 分配与半参数效率 我们现在将后验矩一致性转化为自适应设计保证和高效的 ATE 推断保证。给定后验抽取 $(\mu_w, q_w)$，定义诱导的残差方差抽取为：
$$v_w(x) = \Pi_{[c_V, C_V]} \left\{ q_w(x) - \mu_w(x)^2 \right\},$$
其中 $\Pi_{[c_V, C_V]}$ 是逐点投影到 $[c_V, C_V]$ 上。令 $\bar{v}_{w,t}(x) = \mathbb{E}_{\Pi_t}[v_w(x) \mid \mathcal{H}_t]$ 为后验方差信念。贝叶斯教师根据后验 Neyman 倾向得分分配下一个单元：
$$e_t^{\mathrm{B}}(x) = \mathrm{clip}_{[\eta, 1-\eta]} \left\{ \frac{\sqrt{\bar{v}_{1,t-1}(x)}}{\sqrt{\bar{v}_{0,t-1}(x)} + \sqrt{\bar{v}_{1,t-1}(x)}} \right\}.$$
其 Oracle 对应项为：
$$e^\star(x) = \frac{\sqrt{v_1^\star(x)}}{\sqrt{v_0^\star(x)} + \sqrt{v_1^\star(x)}}.$$
由于方差有上下界，Oracle 倾向得分不需要剪辑。 **命题 2** (设计收敛)。在命题 1 的条件下，
$$\|\bar{v}_{w,t} - v_w^\star\|_{L_2(P_X)} \to 0, \qquad w \in \{0,1\}.$$
因此，后验倾向得分收敛于 Neyman 分配：
$$\|e_t^{\mathrm{B}} - e^\star\|_{L_2(P_X)} \to 0.$$
自适应设计在每个分配时刻产生倾向得分 $e_t(X_t)$。实验结束后，我们使用标准的增强逆概率加权 (AIPW) 估计量进行推断。给定后验均值估计 $\widehat{\mu}_0, \widehat{\mu}_1$，定义：
$$\widehat{\tau}_{\rm AIPW} = \frac{1}{n} \sum_{t=1}^n \left[ \widehat{\mu}_1(X_t) - \widehat{\mu}_0(X_t) + \frac{W_t\{Y_t - \widehat{\mu}_1(X_t)\}}{e_t(X_t)} - \frac{(1-W_t)\{Y_t - \widehat{\mu}_0(X_t)\}}{1-e_t(X_t)} \right].$$
使用已知的被剪辑倾向得分确保了有效的逆概率校正，同时矩信念的后验收缩使得