当前位置：首页 > article >正文

低精度加速高斯过程

article 2025/9/6 23:09:06

0. 摘要

低精度算术对神经网络的训练产生了变革性的影响，减少了计算、内存和能量需求。然而，尽管它有希望，但低精度算术在高斯过程 (GP) 训练中几乎没有受到关注，这主要是因为 GP 需要复杂的线性代数例程，这些例程在低精度下不稳定。我们研究了以半精度训练 GP 时可能发生的不同故障模式。为了规避这些故障模式，我们提出了一种多方面的方法，包括具有重新正交化、紧凑内核和预处理器的共轭梯度。我们的方法在各种设置下显着提高了低精度共轭梯度的数值稳定性和实际性能，并将单个 GPU 上 180 万个数据点的运行时间减少到 10 小时，而无需任何稀疏近似。

1. 介绍

低精度训练可以显着提高训练时间、内存需求和能量消耗。此外，随着硬件加速器的不断发展，低精度算法的可扩展性增益将在未来继续升值。然而，低精度训练会增加舍入误差的幅度并降低数值表示的支持范围，从而导致梯度质量降低并最终牺牲性能，或者在最坏的情况下会牺牲训练过程的收敛性。最近的工作已经开发出高效且稳定的方法，用于在低精度下训练深度神经网络并取得很好的效果，通常是通过开发新的浮点表示 [Intel，2018，Wang 和 Kanwar，2019]，混合精度表示 [Micikevicius 等人。 , 2017 年Das et al., 2018, Yang et al., 2019]，或替代求和技术 [NVIDIA, 2017, Micikevicius et al., 2017, Zamirai et al., 2021]。

高斯过程 (GP) 比神经网络要昂贵得多，因此可能会从低精度计算中获得更多收益。对于 $n$ 个数据点，GP 通常需要 $O(n^3)$ 计算和 $O(n^2)$ 内存。因此，从浮点到半精度已经将矩阵向量乘法的内存需求和运行时间减半。然而，GPs 通常需要复杂的代数运算，例如用于求解线性系统的 Cholesky 分解，这些运算严重受到舍入误差的影响，并且非常不适合半精度。 为低精度计算开发这些方法是一个开放的研究领域 [Higham and Pranesh, 2021]。

相反，我们专注于高斯过程的迭代技术 [Gibbs and MacKay, 1997, Cutajar et al., 2016, Gardner et al., 2018]，它使用共轭梯度来计算线性系统的解，同时利用极其有效的基于 MapReduce 的矩阵向量 使用 KeOps [Charlier et al., 2021, Feydy et al., 2020] 进行乘法运算。这些方法在低精度计算的背景下特别有吸引力：（1）与 Cholesky 分解相比，它们遭受的舍入误差更少 [Gardner et al., 2018]； (2) 它们很容易在 GPU 上并行化，GPU 是专门为加速低精度计算而设计的； (3) 内存消耗是主要的计算瓶颈，因为更多的内存意味着更少的计算昂贵的矩阵分区 [Wang et al., 2019]。

尽管有许多挑战需要克服，但我们证明了以半精度训练 GP 可以稳定、高效且广泛实用。特别是，我们的贡献包括：

我们展示了半精度核矩阵与更高精度表示相比具有质量不同的特征谱，并且表示协方差的距离比浮点或双精度要短得多。
我们研究了内核选择、CG 容差、预处理器等级和紧凑支持对半精度稳定性的影响。
基于这些见解，我们提出了一种新的数值稳定共轭梯度(CG) 算法，该算法在半精度下快速收敛，特别是由于使用了重新正交化、混合精度和logsumexp 技巧。
我们提供了一个半精度的高斯过程的强大实用实现，可将单个 GPU 上 180 万个数据点数据集的训练时间减少到不到 10 小时，这是对最近 8 日 3 天内 100 万个数据点的 CG 里程碑的重大进步 GPU [Wang et al., 2019]。

2. 预备知识

本节介绍不同浮点精度表示的基本信息、训练高斯过程以及如何将优化转化为线性求解，最后介绍作为加速线性求解器的共轭梯度。

2.1 数值精度

现代计算机体系结构以浮点精度表示数字，通常具有 32 位（浮点）或 64 位（双精度）精度 [Kahan，1996]。浮点精度使用 1 位作为数字 (s) 的符号，8 位用于指数 (e)，23 位用于尾数（或有效位，m）。那么，一个数的值表示为
在这里插入图片描述

2.2 高斯过程

观测数据： $D=\{(x_i,y_i)\}_{i=1}^N$ ,高斯过程模型：
$f (.) : g p (m (.), k (., .))$
$y:f(x_i)+\epsilon_i,\epsilon_i:N(0,\sigma^2)$
$m (.)$ 设置为零而不失一般性.
我们通过最小化负边际似然来训练内核超参数 $θ$ ：
在这里插入图片描述
其中 $\tilde{K} ∈ R^{N×N}$ 是所有具有对角观测噪声的数据点的 Gram 矩阵

将评估的协方差矩阵表示为 K。为了执行 θ 的超参数估计，我们需要计算损失函数的梯度，由下式给出

在这里插入图片描述
第一行来自矩阵导数 [Rasmussen and Williams, 2008]，第二行来自 Hutchinson 的迹估计量对迹项的应用 [Gibbs and MacKay, 1997, Gardner et al., 2018]。限制时间复杂度是计算 $=\tilde{K}^{-1} y$ ，使用 Cholesky 分解时需要 $O(N^3)$ 次操作和空间 [Rasmussen and Williams, 2008]。相反，Gibbs 和 MacKay [1997]、Gardner 等人。 [2018] 提出使用共轭梯度 (CG) 来计算这些线性系统，将共轭梯度的 $J$ 步的时间复杂度降低到 $O(JN^2)$ 。

2.3 共轭梯度

我们使用共轭梯度 (CG) 来求解线性系统 $A x = b$ 对于 $x$ ，因此 $x = A^{−1}b$ [Hestenes and Stiefel, 1952, Nocedal and Wright, 2006, Demmel, 1997]。为了估计 $x$ ，CG 使用三项递归迭代地最小化 CG ，其中每个新项只需要与 A 的矩阵向量乘法。形式上，每次 CG 迭代计算以下总和的新项： $A^{−1}b = \sum_{i=1}^N γ_id_i$ ，其中，为简单起见，我们假设算法在 $x_0 = 0$ 处初始化， $γ_i$