在不确定性量化方面,高斯过程 (GPs) 堪称机器学习领域的“瑞士军刀”。它们提供了一种强大且有原则的方法,不仅能进行预测,还能评估我们对这些预测的置信度。这使得它们在药物研发、机器人技术和自动化科学探索等高风险应用中具有不可估量的价值,因为在这些领域,了解模型不知道什么与它预测什么同样重要。
然而,高斯过程有一个众所周知的“阿喀琉斯之踵”:** 计算复杂度**。精确的高斯过程推断的计算复杂度与数据点数量呈三次方关系,即 \(O(N^3)\)。这意味着,虽然高斯过程在处理几千个观测数据时可能表现优异,但当面对现代数据集中常见的数百万个数据点时,它就会陷入停滞。
几十年来,研究人员开发了各种巧妙的方法来解决这个扩展性问题。大多数解决方案可归为两类:
- 诱导点方法: 用一小组伪输入来概括数据集。
- 迭代求解器,如共轭梯度法 (CG) ,它能更高效地逼近精确的高斯过程解。
最近的一篇论文《使用随机梯度下降从高斯过程后验进行采样》提出了一种引人入胜的新方法:** 随机梯度下降 (SGD)** 。这听起来可能有些反直觉: SGD 是深度学习的主力军——简单、常被认为是“朴素”的——而大多数人不会选择它来解决高斯过程核心的那些精确、结构化的线性代数问题。
然而,作者们表明,SGD 不仅可以用于高斯过程推断,而且还能取得顶尖水平的结果——尤其是在处理大型或病态 (ill-conditioned) 数据集时。其秘诀在于一个他们称为**良性非收敛 **(benign non-convergence) 的反直觉现象,即 SGD 不完全收敛的倾向实际上可以成为一种优势。在这篇文章中,我们将剖析他们的方法,看看他们如何调整 SGD 用于高斯过程后验采样,并理解其惊人效果背后的优美理论。
高斯过程简述
高斯过程是函数的分布。与拟合多项式中的权重等参数不同,高斯过程定义了一个函数上的先验,由一个均值函数 \(\mu(\cdot)\) (通常设为零) 和一个协方差函数 (核函数) \(k(\cdot, \cdot')\) 指定,后者用于衡量输入之间的相似性。
给定训练输入 \(\mathbf{x}\) 和带噪声的观测值 \(\mathbf{y}\),我们使用贝叶斯规则将这个先验更新为后验高斯过程。后验均值给我们预测值;后验协方差则量化了不确定性。
在测试点 \((\cdot)\) 处的高斯过程后验均值和协方差为:
\[ \mu_{f|\boldsymbol{y}}(\cdot) = \mathbf{K}_{(\cdot)\boldsymbol{x}} (\mathbf{K}_{\boldsymbol{x}\boldsymbol{x}} + \boldsymbol{\Sigma})^{-1} \boldsymbol{y} \]\[ k_{f|\boldsymbol{y}}(\cdot, \cdot') = \mathbf{K}_{(\cdot,\cdot')} - \mathbf{K}_{(\cdot)\boldsymbol{x}}(\mathbf{K}_{\boldsymbol{x}\boldsymbol{x}} + \boldsymbol{\Sigma})^{-1} \mathbf{K}_{\boldsymbol{x}(\cdot')} \]计算瓶颈在于 \((\mathbf{K}_{\boldsymbol{x}\boldsymbol{x}} + \boldsymbol{\Sigma})^{-1}\)——对这个稠密的 \(N\times N\) 矩阵求逆的成本为 \(O(N^3)\)。
路径式条件化: 高斯过程后验的另一种视角
另一种方法是使用**路径式条件化 **(pathwise conditioning) ,它直接从后验中构造函数样本:
\[ (f \mid \boldsymbol{y})(\cdot) = f(\cdot) + \mathbf{K}_{(\cdot)\boldsymbol{x}}(\mathbf{K}_{\boldsymbol{x}\boldsymbol{x}} + \boldsymbol{\Sigma})^{-1}(\boldsymbol{y} - f(\boldsymbol{x}) - \boldsymbol{\varepsilon}), \quad \boldsymbol{\varepsilon} \sim \mathcal{N}(\mathbf{0}, \boldsymbol{\Sigma}), \quad f \sim \mathrm{GP}(0, k) \]这表示: 取一个先验函数 \(f(\cdot)\),用一个修正项将其“推向”数据点。但问题在于,这个修正项仍然需要求解大型线性系统。
为了使其可行,我们需要:
- 一种从先验高效采样 \(f(\cdot)\) 的方法;
- 一种高效的线性系统求解方法。
对于 (1) ,作者使用了随机傅里叶特征 (Random Fourier Features, RFFs) : 一种通过低维特征映射近似核函数的方法:
\[ k(x, x') \approx \langle \Phi(x), \Phi(x') \rangle, \quad f(\cdot) \approx \tilde{f}(\cdot) = \boldsymbol{\theta}^\top \Phi(\cdot), \quad \boldsymbol{\theta} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \]这样采样先验成本很低。对于 (2) ,他们转向了 SGD。
从二次优化到 SGD
求解高斯过程后验均值可表述为最小化一个二次损失函数:
\[ \boldsymbol{v}^* = \arg\min_{\boldsymbol{v} \in \mathbb{R}^N} \sum_{i=1}^N \frac{\big(y_i - \mathbf{K}_{x_i\boldsymbol{x}} \boldsymbol{v}\big)^2}{\Sigma_{ii}} + \|\boldsymbol{v}\|_{\mathbf{K}_{\boldsymbol{x}\boldsymbol{x}}}^2 \]\[ \mu_{f| \mathbf{y}}(\cdot) = \mathbf{K}_{(\cdot)\boldsymbol{x}} \boldsymbol{v}^* \]通常会使用共轭梯度法 (CG) ,它利用二次结构比普通梯度下降收敛更快。然而在病态问题中,CG 仍可能收敛缓慢。
这引出了论文的核心问题: 不利用二次结构的普通 SGD 能否竞争?
使用 SGD 进行后验采样
作者为后验均值和后验样本设计了可用小批量 (minibatch) SGD优化的随机目标函数。
图1: 在病态条件下,当 CG 无法在规定时间内收敛时,SGD 表现出色。在良态条件下,两者都能恢复精确解。
1. 用于后验均值的 SGD
通过以下方式使二次均值目标函数适用于 SGD:
- 对损失项进行小批量处理,使用 \(D \ll N\) 个数据点;
- 使用 RFFs 进行随机正则化估计。
得到的目标函数为:
\[ \mathcal{L}(\boldsymbol{v}) = \frac{N}{D} \sum_{i=1}^D \frac{(y_i - \mathbf{K}_{x_i\boldsymbol{x}}\boldsymbol{v})^2}{\Sigma_{ii}} + \sum_{\ell=1}^L \left(\boldsymbol{v}^\top \phi_\ell(\boldsymbol{x})\right)^2 \]这个目标函数每步的计算成本为 \(O(N)\)。
2. 用于后验样本的低方差目标函数
从路径式形式来看,“不确定性减少项”涉及另一个二次问题,但目标是带噪声的数据 \(f(x_i)+\varepsilon_i\),朴素的小批量会产生高噪声梯度。
通过将噪声移到正则化项可以降低方差:
\[ \min_{\boldsymbol{\alpha} \in \mathbb{R}^N} \sum_{i=1}^N \frac{(f(x_i) - \mathbf{K}_{x_i x} \boldsymbol{\alpha})^2}{\Sigma_{ii}} + \|\boldsymbol{\alpha} - \boldsymbol{\delta}\|_{\mathbf{K}_{\boldsymbol{x}\boldsymbol{x}}}^2, \quad \boldsymbol{\delta} \sim \mathcal{N}(\mathbf{0}, \boldsymbol{\Sigma}^{-1}) \]这个目标与带噪声版本的最小化器相同,但梯度方差大幅降低。
图2: 低方差重构降低了 SGD 更新中的噪声。诱导点变体使得能高效处理海量数据集成为可能。
进一步扩展: 诱导点
即使每步 \(O(N)\) 的开销,千万级数据仍压力巨大。诱导点方法使用 \(M\) 个伪输入 (\(M \ll N\)) 概括数据集,只需学习 \(M\) 个参数。
在作者提出的诱导点 SGD中,均值和样本目标都用 \(M\) 维表示权重表达,将每步开销降为 \(O(SM)\) (\(S\)为样本数) 。这样可以使用数十万个诱导点,既能紧密逼近完整 GP,又保持高速。
良性非收敛的魔力
SGD 的收敛模式与 CG 截然不同。
图3: 尽管在参数空间中收敛缓慢,SGD 的预测准确率却快速达到最优。
为何如此?作者将输入空间分为:
1. 先验区域
远离数据时,核相似度 \(k(x_i, x_{\text{test}}) \approx 0\)。精确后验与 SGD 近似后验都会回归先验: 误差为零。
2. 插值区域
在数据附近,预测很关键。对核矩阵做谱分解得到基函数:
\[ u^{(i)}(\cdot) = \sum_{j=1}^N \frac{U_{ji}}{\sqrt{\lambda_i}} k(x_j, \cdot) \]大特征值基函数集中在数据密集区。理论 (命题1) 表明,SGD 在大 \(\lambda_i\) 对应方向收敛很快:
\[ \|\operatorname{proj}_{u^{(i)}} \mu_{f|y} - \operatorname{proj}_{u^{(i)}} \mu_{\text{SGD}}\|_{H_k} \le \frac{1}{\sqrt{\lambda_i}} (\text{error terms shrinking as } t^{-1}) \]因此,SGD 能迅速捕捉这些关键方向。
3. 外推区域
在先验区与插值区之间,小 \(\lambda_i\) 基函数占主导。SGD 在这些方向收敛缓慢,导致了图3中较大的 RKHS 误差。但因数据稀少,这些误差对预测影响甚微——这就是良性非收敛。
图4: SGD 的高精度源于插值区域;外推区域的误差是无害的。
实验证据
大规模回归
作者在九个数据集 (最多200万点) 上测试,核函数与超参数固定。
表1: 在大型或病态数据集上,SGD 在 RMSE 和 NLL 上优于 CG 与 SVGP。
图5: 在大规模场景下,SGD 在获得良好预测所需的时间方面优于 CG。
高风险决策: 并行汤普森采样
作者进行大规模贝叶斯优化,使用批量汤普森采样 (每步1000次采集) 。这要求快速获得大量后验样本。
图6: 在紧张计算预算下,SGD 在更短时间内得到更优的目标函数值。
核心要点
此工作将高斯过程后验采样重新表述为随机优化问题,实现了使用 SGD 的线性成本采样及利用诱导点的亚线性成本采样。
关键在于:
- SGD 沿高特征值方向的早期收敛带来快速的预测精度;
- 病态条件主要影响小特征值方向,而 SGD 会自然降低这些方向的优先级,使其稳健;
- 良性非收敛是一种“因祸得福”。
作者证明了一个“不完美”的优化器可以与高斯过程后验的几何结构高度契合,挑战了可扩展机器学习中关于收敛与最优性的传统认知。这为将高斯过程应用于真正海量数据集与复杂序列决策问题打开了大门——在这些场景中,不确定性估计必须快速、准确且可信。