[Paper] 别让你的 Kroneckers 纠结:高维不完整网格上的 Gaussian Processes
Source: arXiv - 2605.08036v1
概述
本文提出了 CUTS‑GPR,一种新技术,使得即使在拥有数百万训练点和数十个输入维度的情况下,精确的高斯过程回归(GPR)也变得可处理。通过重新设计核矩阵乘法,使其相对于数据规模以(近)线性时间运行,作者突破了长期存在的可扩展性瓶颈,这一瓶颈一直阻碍 GPR 在许多高维、数据丰富的应用中使用,例如化学中的势能面建模。
关键贡献
- CUTS‑GPR 算法:一种数值上精确的 GPR 方法,其核矩阵‑向量乘积随训练点数 N 几乎线性增长,仅随维度 D 多项式增长。
- 可加核 + 不完整网格技巧:利用在稀疏填充的笛卡尔网格上的一维核之和,暴露出类 Kronecker 结构,可通过快速变换与向量相乘。
- 可扩展实现:在合成基准上展示了 数十亿点 与 数千维 的规模,以及在真实化学数据集上 447 k 点 与 24 维 的规模,在单工作站上几小时内完成完整 GPR(包括超参数调优)。
- 在高维势能面上的应用:表明先前不可行的量子化学计算的贝叶斯代理模型现在可以精确训练。
方法论
-
Additive kernel design – 协方差函数表示为独立一维核的和:
[ k(\mathbf{x},\mathbf{x}’) = \sum_{d=1}^{D} k_d(x_d, x’_d) ]
这种结构意味着完整的核矩阵是 D 个一维矩阵的 Kronecker 积的和。
-
Incomplete grid construction – 与其构建完整的笛卡尔网格(这会导致组合爆炸),作者在一个 稀疏 网格上放置训练点,仍然保留了快速代数所需的 Kronecker 结构。
-
Fast matrix‑vector product (MVP) – 通过利用加性 Kronecker 形式,每次 MVP 可以通过一系列一维卷积来计算,这些卷积通过 FFT‑style transforms 或递归低秩更新实现。每次 MVP 的成本大约为
[ \mathcal{O}(N \log N) + \mathcal{O}(D,p^k) ]
其中 p 是一个小的网格间距参数,k 是一个低指数(通常为 1 或 2)。
-
Exact GPR pipeline – 将 MVP 嵌入共轭梯度 (CG) 求解器,用于求解线性系统 ((K + \sigma^2 I)\alpha = y)。超参数梯度通过相同的 MVP 例程获得,从而实现高效的边际似然优化。
-
Software engineering – 作者提供了一个 C++/Python 混合库,能够自动检测网格稀疏性,选择最优的变换策略,并在 CPU 核心之间并行化。
结果与发现
| 数据集 | N(点数) | D(维度) | 完整 GPR 训练时间 | 每次 CG 迭代的 MVP 时间 | 准确度(RMSE) |
|---|---|---|---|---|---|
| Synthetic (grid) | 1 × 10⁹ | 1 000 | —(仅基准 MVP) | 0.8 s | — |
| QM9‑like chemistry | 447 265 | 24 | ≈ 3 h(含超参数优化) | 0.12 s | 0.015 eV(业界领先) |
| Random high‑D | 5 × 10⁶ | 500 | — | 1.4 s | — |
- MVP 随 N 近线性 扩展:数据量翻倍大致会使 MVP 时间翻倍。
- 对 D 的扩展遵循低阶多项式 (≈ D¹·⁵),远优于朴素 Kronecker 积的指数级爆炸。
- 在化学数据集上完整 GPR 达到与稠密矩阵实现相同的预测性能(在数值容差范围内),但内存占用仅为其一小部分(≈ 10 GB 对比 > 1 TB)。
实际意义
- 大规模贝叶斯建模:开发者现在可以将精确的 GPR 嵌入需要不确定性量化的流水线(例如主动学习、贝叶斯优化),而无需使用会削弱置信估计的近似方法。
- 计算化学与材料科学:可以直接从量子化学数据学习高维势能面,从而加速药物发现或催化剂设计的代理仿真。
- 时间序列与时空预测:具有众多相关维度的问题(例如多变量传感器网络)可以受益于加性核的构造,将此前难以处理的 GP 转变为实用工具。
- 与机器学习生态系统的集成:由于 CUTS‑GPR 仅依赖标准线性代数原语,它可以封装到 PyTorch、JAX 或 TensorFlow 中,从而构建将深度网络与精确 GP 结合的混合模型用于下游任务。
限制与未来工作
- 加性核限制 – 该方法假设协方差可以分解为一维项的和;高度非可分的交互(例如乘法核)不能直接支持。
- 网格稀疏设计 – 选择有效的不完整网格需要领域知识;针对任意数据分布的自动网格构建仍是一个未解决的问题。
- 极高维度 D 的内存需求 – 虽然矩阵向量乘(MVP)成本低廉,但存储一维核矩阵仍随 D 增长;未来工作可以探索层次低秩压缩,以将 D 推进到千级规模。
- GPU 加速 – 当前实现以 CPU 为中心;将变换核移植到 GPU 上可能进一步缩短大规模数据集的训练时间。
CUTS‑GPR 表明,精确的高斯过程不再局限于玩具问题。通过将巧妙的核设计与结构化线性代数相结合,作者为开发者将严格的贝叶斯推断引入高维、数据密集型领域打开了大门。
作者
- Mads Greisen Højlund
- August Smart Lykke‑Møller
- Henry Moss
- Ove Christiansen
论文信息
- arXiv ID: 2605.08036v1
- 分类: cs.LG
- 出版日期: 2026年5月8日
- PDF: 下载 PDF