技术总结：高维线性模型中经验贝叶斯的梯度流

问题陈述

本文探讨了在高维线性模型中进行经验贝叶斯（Empirical Bayes, EB）估计所面临的挑战。在这种设定下，潜在参数通过复杂的回归设计进行相互作用，这使得理论和算法研究滞后于发展成熟的序列模型（例如 $y_i = \theta_i + \epsilon_i$）。具体而言，作者考虑了如下线性模型：
$$y = X\theta + \epsilon, \quad \theta_j \overset{iid}{\sim} g^*, \quad \epsilon_i \overset{iid}{\sim} N(0, \sigma^2)$$
其中 $n$（样本量）和 $p$（维度）均很大。目标是从数据 $(X, y)$ 中估计未知的先验分布 $g^*$。

主要的难点在于：

1. 统计学方面： 在一般回归设计 $X$（即 $X$ 不为单位矩阵时）下，建立非参数最大似然估计量（NPMLE）的一致性（当 $n, p \to \infty$ 时）是一个难题。虽然近期的研究 [MSS23] 在强假设（$n \ge p$ 且 $X$ 条件良好）下建立了一致性，但针对一般设计（包括 $n < p$ 和随机设计）的条件仍是一个开放问题。
2. 计算学方面： 边缘对数似拟合函数 $\bar{F}_n(g)$ 涉及对高维 $\theta$ 的难解后验积分，使得直接优化变得困难。现有的近似方法（如平均场变分推断、MCEM）在非凸设定或 $n < p$ 时缺乏完整的理论收敛保证。

方法论

作者提出了一种结合概率测度上的**梯度流（gradient flows）与自适应朗之万动力学（adaptive Langevin dynamics）**的新颖框架。

1. 重参数化与变分表示

为了解决 $\theta$ 的后验分布可能变得离散而缺乏正则性的问题，作者引入了一个平滑的潜在变量 $\phi = \theta + z$，其中 $z \sim N(0, \tau^2 I)$。这使得模型转化为：
$$y = X\phi + \tilde{\epsilon}, \quad \phi_j \overset{iid}{\sim} N_\tau * g^*$$
其中 $N_\tau$ 是高斯核。这种重参数化确保了 $\phi$ 的后验分布具有光滑且正的密度，从而便于采样。

边缘对数似拟合函数可以通过 Gibbs 变分表示来表达：
$$\bar{F}_n(g) = \min_{q \in \mathcal{P}^*(\mathbb{R}^p)} F_n(q, g)$$
其中 $F_n(q, g)$ 是 Gibbs 自由能。优化过程是在后验密度 $q$ 和先验 $g$ 上共同进行的。

2. 双变量梯度流系统

作者提出了一个同时最小化 $F_n(q, g)$ 的同步梯度流方程组：
$$\begin{cases} \frac{d}{dt} q_t(\phi) = -p \cdot \text{grad}_{W_2}^q F_n(q_t, g_t)[\phi] \\ \frac{d}{dt} g_t(\theta) = -\alpha \cdot \text{grad}_{FR}^g F_n(q_t, g_t)[\theta] \end{cases}$$

• $q$-流（Wasserstein-2）： $q_t$ 的演化遵循 Fokker-Planck 方程，对应于一个由当前先验 $g_t$ 决定的随时间变化的漂移项的朗之万扩散过程。
• $g$-流（Fisher-Rao）： $g_t$ 的演化遵循 Fisher-Rao 梯度流，这自然地保持了先验的正定性和有界支撑约束。

3. EBflow 算法

由于精确求解梯度流是难以实现的，作者推导出了一个离散时间算法——EBflow：

• E-步（近似）： 单步非调整朗之万算法（ULA）根据当前先验 $g_t$ 更新样本 $\phi_t$。
• M-步（近似）： 使用单步 Fisher-Rao 梯度下降来更新先验 $g_t$。至关重要的是，高维密度 $q_t$ 通过单个朗之万轨迹 $\phi_t$ 的坐标经验分布进行近似，从而避免了需要多个并行链的情况。

核心贡献

1. 近 NPMLE 的统计一致性

作者在设计矩阵 $X$ 的确定性条件下，确立了任何“近 NPMLE”（即边缘对数似拟合函数的近似优化器）的一致性。

• 假设： 设计矩阵必须允许存在测试向量 $\{z_j\}$，这些向量能够将 $X$ 的列相对于其他列进行“近似正交化”（类似于去偏 Lasso）。这扩展了近期 [MSS23] 的结果，后者要求 $n \ge p$ 且 $X$ 条件良好。
• 结果： 当 $n, p \to \infty$ 时，任何满足 $\bar{F}_n(\hat{g}_n) \le \bar{F}_n(g^*) + o(1)$ 的先验序列 $\hat{g}_n$ 都会弱收敛于真实先验 $g^*$。即使在欠采样机制（$n < p$）以及行向量为亚高斯分布的随机设计下，该结论依然成立。

2. 梯度流的收敛性

论文对连续时间梯度流系统进行了理论分析：

• 单变量流：
  • 对于固定的 $q$，用于 $g$ 的 Fisher-Rao 流以 $O(1/t)$ 的速率全局收敛至序列模型的 NPMLE。
  • 对于固定的 $g$，用于 $q$ 的 Wasserstein 流以 KL 散度指数级快速收敛至后验 $P_g(\phi|y)$。
• 联合流： 在后验分布满足无维度 Log-Sobolev 不等式（LSI）（对于足够大的噪声方差 $\sigma^2$ 已建立）的情况下，联合流收敛至一个近 NPMLE。
• 时间复杂度： 如果边缘对数似拟合函数在由初始化定义的下水平集上是凸的，则算法在时间上与 $(n, p)$ 成线性关系即可达到近 NPMLE。

3. 算法实现

论文介绍了 EBflow，这是一个实用的算法，它：

• 使用单个朗之维链来估计平均边缘分布，利用了高维空间中的“单链粒子混沌传播”（single-chain propagation-of-chaos）直觉。
• 通过 $\phi$ 进行重参数化，以确保即使在真实先验 $g^*$ 是离散的情况下，也能满足朗之万采样的正则性要求。
• 在模拟实验中展示了优于 MCEM 和变分推断基准方法的竞争性能。

结果

理论保证

• 一致性： 定理 3.2 证明，在满足假设 3.1（涵盖了随机设计和 $n < p$）的情况下，边缘似拟合函数的近优化器是 $g^*$ 的一致估计量。这放宽了先前研究 [MSS23] 中更强的条件。
• 收敛率： 定理 3.10 确立了联合梯度流在时间范围 $T = O(\epsilon^{-3.01}(n+p))$ 内达到满足平均一阶最优性条件的密度的性质。
• LSI： 定理 3.7 证明了当 $\sigma^2 > M^2 \|X\|_{op}^2$ 时，平滑变量 $\phi$ 的后验分布具有无维度的 LSI，这一条件与一致性机制是兼容的。

经验表现

在与 Langevin-MCEM、Gibbs-MCEM 以及坐标上升变分推断（CAVI）进行对比的数值实验中：

• 准确性： 与 CAVI 相比，EBflow 在估计先元密度方面通常能获得更低的全变差（TV）误差，并且在面对具有块相关或基因型设计的场景时，其表现与 MCEM 方法相当甚至更优。
• 效率： EBflow 的运行速度比 MCEM 方法更快，因为后者在每个 M 步都需要进行昂贵的凸优化（NPMLE）。
• 鲁棒性： 该方法在各种先验形状（高斯、偏态、双峰、混合尺度）和设计矩阵（单位矩阵、i.i.d.、块相关、遗传设计）下均表现良好。

意义与主张

本文声称填补了高维线性模型中经验贝叶斯理论的空白，其理由如下：

1. 扩展了 NPMLE 的一致性保证，使其适用于一般的线性模型（超越了序列模型），采用了涵盖随机设计和 $n < p$ 的确定性设计假设，放宽了近期研究 [MSS23] 中较强的条件（$n \ge p$ 且 $X$ 条件良好）。
2. 开发了一种新的计算范式，使用双变量梯度流和自适应朗之万动力学，为变分推断和 MCEM 提供了一个具有理论依据的替代方案。
3. 通过为高维后验律建立新的 Log-Sobolev 不等式，确立了无维度的收敛特性，这对贝叶斯采样本身也具有独立的研究价值。

作者指出，虽然收敛保证是针对连续时间流给出的，但他们也为离散的 EBflow 算法提供了见解。他们承认，目标函数的全局景观可能是非凸的，因此他们的理论保证适用于特定初始化的局部收敛，而全局几何性质仍是一个开放性问题。