Scalable Derivative Gaussian Processes via Exact Gradient Reduction¶

会议: NeurIPS 2026
arXiv: 2606.02909
OpenReview: https://openreview.net/forum?id=0mN2CCr29C
代码: https://github.com/hseung88/tera
领域: 优化/理论
关键词: 高斯过程, 导数观测, Vecchia近似, 精确梯度降维, 贝叶斯优化

一句话总结¶

针对导数高斯过程在联合状态下 $\mathcal{O}(n^3 d^3)$ 的立方级计算瓶颈，本文提出 Tera 方法，证明平稳核下与目标差值正交的梯度分量条件独立于目标函数值，从而将局部全梯度精确降维为至多 $m^2$ 个方向导数，并结合 Vecchia 近似将单因子求逆复杂度彻底解耦为与输入维度 $d$ 无关的 $\mathcal{O}(m^6)$。

研究背景与动机¶

在基于自动微分与伴随求解器的现代代理模型任务中，高阶模拟器在返回标量目标函数值的同时往往能以极低边际成本输出全维度输入梯度。将这些高阶导数信息作为观测值注入高斯过程（GP）能够极大地约束函数曲面、提升小样本回归精度并加速贝叶斯优化寻找全局极值的效率。然而，传统导数高斯过程（Derivative GP）通过核函数的各阶偏导将函数值与全量梯度拼成一个规模高达 $n(d+1)$ 的联合高斯状态，其精确推断需要对密集协方差矩阵进行求逆或 Cholesky 分解，面临着 $\mathcal{O}(n^3 d^3)$ 的严苛时间和 $\mathcal{O}(n^2 d^2)$ 的内存墙，导致现有方法在样本量 $n$ 较大或维度 $d$ 超过几十维时彻底瘫痪。

为了缓解这一瓶颈，现有的可扩展导数 GP 方法主要沿着两条路线探索：其一是基于全局核插值（如 D-SKI 与 DSoftKI），通过将导数协方差投影到结构化网格上加速矩阵向量乘法，但其精度和拟合复杂度严重依赖于全局插值网格的设计；其二是基于稀疏变分推断（如 DSVGP 与 DDSVGP），引入伪输入导数或方向导数诱导变量，但诱导点表征会引入额外变分下界松弛与方向近似误差。现有方法普遍将联合函数-梯度状态本身作为不可分割的全局计算对象，忽视了一个关键的数学特性：在大多数实际任务中，推断的终极目标仅仅是预测未知点的标量函数值或评估标量采集函数，我们真正需要的并非全局高维梯度场，而是梯度对标量目标后验分布的边际贡献。

从高斯条件分布几何的角度审视，平稳核函数下观测梯度对标量目标的条件均值修正与方差缩减，本质上只沿着连接目标输入点与参考观测点的差值向量方向传播，而在与这些差值张成子空间正交的方向上，梯度分量在统计上条件独立于目标函数值。核心 idea：利用平稳核的局部几何正交性，将目标周围 $m$ 个邻近观测点的 $md$ 维全梯度精确压缩为至多 $m^2$ 维目标特异性标量投影统计量，并将其作为局部因子嵌入 Vecchia 似然近似框架中，在数学严谨保留全量梯度推断精度的同时，彻底消除高斯条件块求逆对维度 $d$ 的依赖。

方法详解¶

整体框架¶

Tera（Target-specific Exact gradient Reduction for Accelerating derivative GP inference）以标量函数值预测为推断核心目标。在给定 $n$ 个观测点输入 $\mathbf{X} \in \mathbb{R}^{n \times d}$、标量函数观测 $\mathbf{y} \in \mathbb{R}^n$ 及梯度观测 $\mathbf{g} \in \mathbb{R}^{nd}$ 的前提下，Tera 首先利用尺度化最大最小距离（Maximin）排序构建 Vecchia 条件集，将全局后验联合概率近似分解为局部条件因子的连乘；随后在每个局部因子内，通过目标输入的差值投影矩阵将局部的 $md$ 维全梯度精确降维为 $m^2$ 维统计量 $\mathbf{q}_c$，绕过构建超大导数协方差矩阵，直接依据解析几何公式生成紧凑协方差块完成推断与参数学习。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据: 函数观测 y 及高维梯度 g"] --> B["Vecchia 条件集构建<br/>Maximin 排序与 m 个近邻检索"]
    B --> C["目标特异性精确梯度降维<br/>正交分解保留投影统计量 q_c"]
    C --> D["局部协方差直接解析构建<br/>基于 m×m Gram 矩阵绕过 md×md 矩阵"]
    D --> E["条件 Vecchia 似然优化与推断<br/>无偏评分小批量 SGD 训练"]
    E --> F["输出: 目标函数值标量后验均值与方差"]

关键设计¶

1. 目标特异性精确梯度降维：利用平稳核几何消除正交无关梯度
针对传统导数 GP 需同时处理局部参考集所有 $md$ 维梯度坐标的维度灾难，Tera 基于平稳核 $k_{\bm{\theta}}(\mathbf{x}, \mathbf{x}') = \kappa_{\bm{\theta}}(\|\tilde{\mathbf{x}} - \tilde{\mathbf{x}}'\|^2)$（其中 $\tilde{\mathbf{x}} = \bm{\Lambda}_{\bm{\theta}}^{1/2}\mathbf{x}$ 为尺度化输入）证明了精确正交分离定理。对于目标输入 $\mathbf{x}_i$ 及大小为 $m_c$ 的局部参考集 $c = \{j_1, \dots, j_{m_c}\}$，构造尺度化差值矩阵 $\bm{\Delta}_i = [\tilde{\mathbf{x}}_{j_1} - \tilde{\mathbf{x}}_i, \dots, \tilde{\mathbf{x}}_{j_{m_c}} - \tilde{\mathbf{x}}_i] \in \mathbb{R}^{d \times m_c}$。设子空间 $\mathcal{S}_i = \mathrm{col}(\bm{\Delta}_i)$，任何与之正交的梯度分量 $\tilde{\mathbf{g}}_{c, \perp} \in \mathcal{S}_i^\perp$ 与目标标量函数值 $y_i$ 以及参考点函数值 $\mathbf{y}_c$ 的协方差严格为零。因此，定义目标特异性降维梯度统计量： $$ \mathbf{q}c = (\mathbf{I}} \otimes \mathbf{Di^\top)\mathbf{g}c \in \mathbb{R}^{m_c^2}, \quad \text{其中 } \mathbf{D}i^\top \mathbf{g}} = \begin{bmatrix} (\mathbf{x} - \mathbf{xi)^\top \mathbf{g}} \ \vdots \ (\mathbf{x}} - \mathbf{xi)^\top \mathbf{g} $$ 在无噪梯度或匹配噪声 } \end{bmatrix} \in \mathbb{R}^{m_c$\bm{\epsilon}_j^g \sim \mathcal{N}(\mathbf{0}, \sigma_g^2 \bm{\Lambda}_{\bm{\theta}})$ 假设下，严格成立条件概率等价性：$p_{\bm{\theta}}(y_i \mid \mathbf{y}_c, \mathbf{g}_c) = p_{\bm{\theta}}(y_i \mid \mathbf{y}_c, \mathbf{q}_c)$。这意味着无需接触任何正交维度的梯度噪声与冗余特征，仅凭 $m_c^2$ 个未归一化方向导数即可完整复现全导数高斯过程的后验更新。

2. 降维协方差块的直接解析构建：避免实例化稠密局部导数矩阵
若先实例化常规的 $m_c d \times m_c d$ 局部导数协方差矩阵再通过矩阵乘法降维，仍将引入 $\mathcal{O}(m_c^3 d^3)$ 的分解与构造开销。Tera 提出了基于核几何内积的直接构造法：只需计算局部 $m_c \times m_c$ 的目标中心化 Gram 矩阵 $\mathbf{H}_i = \mathbf{D}_i^\top \bm{\Lambda}_{\bm{\theta}} \mathbf{D}_i$。对于任意参考点对 $(j_a, j_b)$，其两两平方距离与核输入由 Gram 矩阵元素直接导出：$r_a = [\mathbf{H}_i]_{aa}$ 以及 $r_{ab} = [\mathbf{H}_i]_{aa} + [\mathbf{H}_i]_{bb} - 2[\mathbf{H}_i]_{ab}$。由此可直接解析写出规模仅为 $m_c^2 \times m_c^2$ 的降维协方差块： $$ \mathbf{K}{\mathbf{q}a \mathbf{q}b} = -2\kappa'(r}})\mathbf{Hi - 4\kappa''(r}})(\mathbf{h}_a - \mathbf{h}_b)(\mathbf{h}_a - \mathbf{hb)^\top + \mathbbm{1}_i $$ 其中 }\sigma_g^2 \mathbf{D}_i^\top \mathbf{D$\mathbf{h}_a = \mathbf{H}_i \mathbf{e}_a$。整个计算过程完全跳过了全局及局部高维导数矩阵的内存分配与分解，使得稠密 Cholesky 分解的对象直接限制在 $m_c^2 \times m_c^2$ 的降维条件协方差上，单步求逆耗时收敛为 $\mathcal{O}(m^6)$，彻底解除了维度 $d$ 的三次指数束缚。

3. 条件 Vecchia 似然与无偏小批量学习机制
为了摆脱全量样本对求逆全局协方差矩阵 $n$ 的三次依赖，Tera 采用 Vecchia 链式法则因式分解，将全局似然转换为单变量条件高斯分布之积：$\tilde{p}_{\bm{\theta}}^{\mathrm{Vec}}(\mathbf{y} \mid \mathbf{g}) = \prod_{i=1}^n p_{\bm{\theta}}(y_i \mid \mathbf{y}_{c(i)}, \mathbf{q}_{c(i)})$。本文进一步从理论上证明了评分函数的无偏性（Score Unbiasedness）：对于真实数据生成参数 $\bm{\theta}_0$，在给定排序与条件集下，$\mathbb{E}_{\bm{\theta}_0}[\nabla_{\bm{\theta}} \tilde{\ell}_n(\bm{\theta}_0)] = \mathbf{0}$。这表明直接最大化条件似然 $\tilde{\ell}_n(\bm{\theta})$ 可获得一致且无偏的参数估计。在训练时，每次迭代仅需均匀采样大小为 $|\mathcal{B}|$ 的局部条件因子子集，构造随机梯度更新超参数，使每步反向传播的复杂度降为 $\mathcal{O}(|\mathcal{B}|(d m^2 + m^6))$，彻底脱离了对总样本数 $n$ 的单步依赖。

损失函数 / 训练策略¶

Tera 的模型超参数包含核参数（振幅、各向同性或 ARD 逆尺度对角阵 $\bm{\Lambda}_{\bm{\theta}}$）以及观测噪声方差 $\sigma_y^2, \sigma_g^2$。训练目标为最大化对数条件 Vecchia 似然： $$ \tilde{\ell}n(\bm{\theta}) = \sum}^n \left[ -\frac{1}{2} \log \left(\sigma^2_{y_i \mid \mathbf{y{c(i)}} - s_i\right) - \frac{\left(y_i - \mu{c(i)}} - t_i\right)^2}{2\left(\sigma^2 $$ 其中 }_{c(i)}} - s_i\right)} \right] + \text{const$t_i$ 和 $s_i$ 分别表示条件均值修正项与方差削减项。在优化过程中，采用 Adam 优化器进行小批量梯度更新；当更新导致尺度参数大幅变动时，阶段性重构近邻排序图。在下游贝叶斯优化（BO）场景中，采用周期性温启动策略（如每 20 次 BO 迭代仅更新 10 步 Adam），在保持局部几何结构稳定的同时逐步吸纳新数据自适应更新超参数。

实验关键数据¶

主实验¶

论文在大型分子物理势能与原子力基准 MD22（包含 6 个真实分子体系，样本量 $n$ 从 4,528 到 62,777，维度 $d$ 从 168 到 1,110）上进行了能量标量回归评估。基线模型包括标准精确高斯过程（Exact GP）、稀疏变分导数高斯过程（DDSVGP）以及导数软核插值方法（DSoftKI）。所有方法基于相同的各向同性平方指数（SE）核与相同的训练/测试划分，评估恢复物理真实尺度的单原子测试均方根误差（Raw Energy RMSE per atom, 单位 eV/atom）。

分子系统 (体系规模 $n$, 维度 $d$)	Exact GP	DDSVGP	DSoftKI	本文方法 (Tera)	相对 DDSVGP 误差降低
DHA ($n=62,777, d=168$)	显存超限中断	$0.155 \pm 0.001$	显存超限中断	0.004 ± 0.001	-97.4%
AT-AT ($n=18,000, d=180$)	$0.054 \pm 0.001$	$0.119 \pm 0.003$	$0.170 \pm 0.004$	0.009 ± 0.000	-92.4%
Stachyose ($n=24,544, d=261$)	$0.056 \pm 0.002$	$0.107 \pm 0.001$	$0.139 \pm 0.005$	0.005 ± 0.000	-95.3%
AT-AT-CG-CG ($n=9,137, d=354$)	$0.042 \pm 0.001$	$0.068 \pm 0.005$	$0.126 \pm 0.003$	0.011 ± 0.000	-83.8%
Buckyball-catcher ($n=5,491, d=444$)	$0.024 \pm 0.000$	$0.056 \pm 0.009$	$0.089 \pm 0.005$	0.007 ± 0.000	-87.5%
Double-walled-nanotube ($n=4,528, d=1,110$)	$0.017 \pm 0.001$	$0.045 \pm 0.003$	$0.048 \pm 0.003$	0.002 ± 0.000	-95.6%

在 200 维高维贝叶斯优化合成基准（Ackley-200D 与 Levy-200D，预算 150 次评估，5 次独立重复）上，Tera 的最终简单遗憾值（Simple Regret）显著优于纯函数值优化基线：在 Levy-200D 上 Tera 取得约 $0.03$ 的低遗憾值（对比 VBO 的 $\approx 1.2$ 和 TuRBO-1 的 $\approx 1.0$）；在 Ackley-200D 上 Tera 收敛至 $\approx 0.18$（对比 VBO $\approx 3.2$、TuRBO-1 $\approx 3.1$），且单步优化耗时仅约 2 秒，GPU 峰值显存始终在 1GB 以下。

消融实验¶

论文在 Buckyball-catcher 分子体系上深入消融了 Vecchia 邻域条件集大小 $m$ 对测试 RMSE、端到端墙上时钟耗时以及峰值显存的影响（各向同性 Matérn-5/2 核）：

局部条件集规模 $m$	单原子测试 RMSE	端到端运行时间 (s)	GPU 峰值显存 (MB)	说明
$m = 5$	0.0152	8.2	145	邻域过小，未能充分捕获局部导数曲率
$m = 10$	0.0094	14.7	210	预测误差明显下降，显存开销平稳
$m = 20$ (默认配置)	0.0071	32.5	380	综合预测性能与运算效率的最优折中点
$m = 30$	0.0070	89.4	1,120	精度趋于完全饱和，但计算耗时剧烈增加

关键发现¶

降维过程严格无损：在控制先验光滑性的模拟高斯过程中，Tera 与全梯度 Vecchia dGP 的边际 KL 散度曲线完全重合，证实其预测误差与精确导数 GP 之间的细微差距全部来源于 Vecchia 局部近邻近似，目标特异性梯度投影本身在数学上没有造成任何信息丢失。
维度扩展性极佳：当输入维度 $d$ 从 10 维扩展至 1,000 维时，Tera 的显存占用与计算耗时曲线均呈现近乎水平的平坦趋势，彻底攻克了传统方法面对高维导数矩阵内存暴涨的顽疾。
单轮训练即可收敛：相较于 DDSVGP 与 DSoftKI 需针对大规模变分诱导参数进行数十轮繁重训练且易陷入局部次优，Tera 的条件似然直接对齐预测目标，只需对数据进行 1 个 epoch 的扫视即可获得 SOTA 精度。

亮点与洞察¶

从推断目标反推有效观测：传统方法盲目维护全局各维度的完整协方差状态，而 Tera 抓住“只有平行于目标差值向量的导数才会影响目标条件高斯更新”这一本质几何特征，从源头上斩断了无关梯度的代数累赘。
代数结构与局部近似的精巧共生：单独进行目标特异性降维只能降低维度维数，无法降低点数；单独使用 Vecchia 近似只能降低点数，无法降低维度。将二者融合使得降维统计量规模被严格约束在 $m^2$，实现了 $n$ 与 $d$ 的双重解耦。
高阶先验在贝叶斯优化中的高效变现：在 200 维超高维空间中，纯函数值代理模型因样本稀疏极易受困于维度诅咒，而 Tera 以极低的计算开销将一阶梯度注入局部曲率建模，将优化效率推向了全新量级。

局限与展望¶

计算复杂度随近邻规模 $m$ 六次方增长：局部降维协方差大小为 $m^2 \times m^2$，稠密 Cholesky 分解包含 $\mathcal{O}(m^6)$ 项，若遇到非平稳或长程依赖强烈的极端函数需扩大 $m$ 时，计算成本将陡增。
非各向同性噪声下的精确性限制：精确降维等价性严格依赖于无噪梯度或与输入缩放矩阵 $\bm{\Lambda}_{\bm{\theta}}$ 匹配的噪声分布；若面对存在复杂相关性或任意各向异性的观测噪声，投影将退化为一种近似。
仅适用于一阶导数降维：当前理论建立在梯度与标量核的一阶交互上，推广至 Hessian 矩阵等高阶导数或多输出 GP 系统仍需进一步发展代数工具。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [严谨证明了平稳核下正交梯度分量的条件独立性，构思巧妙]
实验充分度: ⭐⭐⭐⭐⭐ [包含仿真KL验证、大规模分子实际动力学数据及200维超高维BO]
写作质量: ⭐⭐⭐⭐⭐ [数学推导清晰严密，复杂度与几何洞察剖析透彻]
价值: ⭐⭐⭐⭐⭐ [为长期受限于立方计算瓶颈的导数高斯过程提供了真正可落地的突破性方案]

分子系统 (体系规模 \(n\), 维度 \(d\))	Exact GP	DDSVGP	DSoftKI	本文方法 (Tera)	相对 DDSVGP 误差降低
DHA (\(n=62,777, d=168\))	显存超限中断	\(0.155 \pm 0.001\)	显存超限中断	0.004 ± 0.001	-97.4%
AT-AT (\(n=18,000, d=180\))	\(0.054 \pm 0.001\)	\(0.119 \pm 0.003\)	\(0.170 \pm 0.004\)	0.009 ± 0.000	-92.4%
Stachyose (\(n=24,544, d=261\))	\(0.056 \pm 0.002\)	\(0.107 \pm 0.001\)	\(0.139 \pm 0.005\)	0.005 ± 0.000	-95.3%
AT-AT-CG-CG (\(n=9,137, d=354\))	\(0.042 \pm 0.001\)	\(0.068 \pm 0.005\)	\(0.126 \pm 0.003\)	0.011 ± 0.000	-83.8%
Buckyball-catcher (\(n=5,491, d=444\))	\(0.024 \pm 0.000\)	\(0.056 \pm 0.009\)	\(0.089 \pm 0.005\)	0.007 ± 0.000	-87.5%
Double-walled-nanotube (\(n=4,528, d=1,110\))	\(0.017 \pm 0.001\)	\(0.045 \pm 0.003\)	\(0.048 \pm 0.003\)	0.002 ± 0.000	-95.6%

局部条件集规模 \(m\)	单原子测试 RMSE	端到端运行时间 (s)	GPU 峰值显存 (MB)	说明
\(m = 5\)	0.0152	8.2	145	邻域过小，未能充分捕获局部导数曲率
\(m = 10\)	0.0094	14.7	210	预测误差明显下降，显存开销平稳
\(m = 20\) (默认配置)	0.0071	32.5	380	综合预测性能与运算效率的最优折中点
\(m = 30\)	0.0070	89.4	1,120	精度趋于完全饱和，但计算耗时剧烈增加