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 n n 个数据点,GP 通常需要 O ( n 3 ) O(n^3) O(n3) 计算和 O ( n 2 ) O(n^2) O(n2) 内存。 因此,从浮点到半精度已经将矩阵向量乘法的内存需求和运行时间减半。 然而,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
D=\{(x_i,y_i)\}_{i=1}^N
D={(xi,yi)}i=1N,高斯过程模型:
f
(
.
)
:
g
p
(
m
(
.
)
,
k
(
.
,
.
)
)
f(.) :gp(m(.),k(.,.))
f(.):gp(m(.),k(.,.))
y
:
f
(
x
i
)
+
ϵ
i
,
ϵ
i
:
N
(
0
,
σ
2
)
y:f(x_i)+\epsilon_i,\epsilon_i:N(0,\sigma^2)
y:f(xi)+ϵi,ϵi:N(0,σ2)
m
(
.
)
m(.)
m(.) 设置为零而不失一般性.
我们通过最小化负边际似然来训练内核超参数
θ
θ
θ:
其中
K
~
∈
R
N
×
N
\tilde{K} ∈ R^{N×N}
K~∈RN×N 是所有具有对角观测噪声的数据点的 Gram 矩阵
将评估的协方差矩阵表示为 K。为了执行 θ 的超参数估计,我们需要计算损失函数的梯度,由下式给出
第一行来自矩阵导数 [Rasmussen and Williams, 2008],第二行来自 Hutchinson 的迹估计量对迹项的应用 [Gibbs and MacKay, 1997, Gardner et al., 2018]。 限制时间复杂度是计算
v
=
K
~
−
1
y
v =\tilde{K}^{-1} y
v=K~−1y,使用 Cholesky 分解时需要
O
(
N
3
)
O(N^3)
O(N3) 次操作和空间 [Rasmussen and Williams, 2008]。 相反,Gibbs 和 MacKay [1997]、Gardner 等人。 [2018] 提出使用共轭梯度 (CG) 来计算这些线性系统,将共轭梯度的
J
J
J 步的时间复杂度降低到
O
(
J
N
2
)
O(JN^2)
O(JN2)。
2.3 共轭梯度
我们使用共轭梯度 (CG) 来求解线性系统
A
x
=
b
Ax = b
Ax=b 对于
x
x
x,因此
x
=
A
−
1
b
x = A^{−1}b
x=A−1b [Hestenes and Stiefel, 1952, Nocedal and Wright, 2006, Demmel, 1997]。 为了估计
x
x
x,CG 使用三项递归迭代地最小化 CG ,其中每个新项只需要与 A 的矩阵向量乘法。形式上,每次 CG 迭代计算以下总和的新项:
A
−
1
b
=
∑
i
=
1
N
γ
i
d
i
A^{−1}b = \sum_{i=1}^N γ_id_i
A−1b=∑i=1Nγidi,其中,为简单起见,我们假设算法在
x
0
=
0
x_0 = 0
x0=0 处初始化,
γ
i
γ_i
γi 是系数,
d
i
d_i
di 是正交共轭搜索方向 Golub 和 Loan [2018]。 在无限精度表示中,CG 的 N 次迭代产生所有 N 个求和项并恢复精确解。 整体CG算法见附录2。
3. 相关工作
迭代高斯过程:
由于使用 Cholesky 训练高斯过程 (GP) 的三次运行时复杂性,已经对使用替代算法来减少该开销进行了广泛的研究。在这项工作中,我们只关注与迭代 GPs 的比较,而忽略了替代方法的大量文献。 Gibbs 和 MacKay [1997] 首先研究了 GP 的迭代技术,通过使用共轭梯度来求解线性系统并使用迹估计,这将 GP 回归的整体复杂度降低到了 k 步 CG 的
O
(
k
N
2
)
O(kN^2)
O(kN2),而 Cutajar 等人[2016] 重新发现了该技术并应用预处理来加速仅关注双精度的线性系统。加德纳等人 [2018] 提出了我们构建的批量共轭梯度算法,并使用 Lanczos 来估计对数行列式,重点关注浮点精度。王等人 [2019] 扩展了这项工作,并通过数据分区在 3 天内使用 8 个 GPU,使精确的 GP 能够在 100 万个数据点上进行训练。在显式 GP 回归之外,Meanti 等人。 [2020] 还在多个 GPU 上使用 KeOps 进行基于 Nyström 的内核回归,在多达 10 亿个数据点的数据集上取得了结果。
低精度算术:
早在 1990 年 [Hammerstrom, 1990] 就对以低于传统双精度或浮点精度的精度训练神经网络模型感兴趣。在现代,古普塔等人 [2015]和陈等人 [2014] 率先使用较低精度的算法,目标是降低内存成本,以实现大型深度神经网络模型的训练。古普塔等人 [2015] 提出了使用随机舍入来减少从浮点向下移动到一半(例如 16 位算术)精度时因精度损失而导致的误差,wile Chen 等人。 [2014] 使用特殊表示来实现 DNN 的准确训练。这些工作引发了对深度神经网络混合精度训练的一系列研究[Micikevicius et al., 2017, Das et al., 2018, Yang et al., 2019]。在启发我们方法的混合精度训练中,每个 DNN 层的激活和梯度传播一半,但主权重副本以浮点精度存储 [Micikevicius et al., 2017]。这些较低精度方法的成功导致通过神经网络层的量化进一步降低内存开销[Jacob et al., 2018, Das et al., 2018],开发加速专用芯片架构现代 GPU [NVIDIA, 2017] 和 TPU 核心 [Wang and Kanwar, 2019] 等较低精度的算法,以及 bfloat16 [Wang and Kanwar, 2019, Intel, 2018] 等增强混合精度训练的新标准。降低精度的承诺已经在深度学习社区之外蔓延开来。卡塔哈尔等人。 [2016] 发现预处理对于防止共轭梯度 (CG) 来估计数值精度的解是必要的,这一发现得到了 Gardner 等人的回应。 [2018] 以及最近的 Wenger 等人。 [2021]。加德纳等人。 [2018],王等人。 [2019],温格等人。 [2021] 主张使用枢轴 Cholesky 预处理器 [Harbrecht et al., 2012]。在数值线性代数社区中,Gratton 等人 [2021] 主张在使用 CG 作为通用优化算法时使用完全重新正交化。海达尔等人。 [2017, 2018] 替代地提出了一种迭代细化和使用 GMRES(一种直接与 CG 相关的方法)在 GPU 上以较低精度算术求解密集线性系统,而 Abdelfattah 等人。 [2020] 考虑了使用迭代细化的混合精度求解。 Higham [2021]、Higham 和 Pranesh [2021] 主张使用分块算法来减少舍入误差,我们在图 1b 中证明了这一点具有相当大的优势。事实上,较低精度的算术已经直接进入内核近似。然而,这些方法主要是基于结构化核基函数,如 Fastfood [Le et al., 2013, Yang et al., 2015] 或基于使用随机傅立叶特征的量化 [Zhang et al., 2019, Li and Li]。
4 方法
我们的方法结合了两组方法。 首先,它扩展了 Gardner 等人的想法。 [2018],王等人 [2019] 用于将训练 GP 的问题转化为只需要矩阵向量运算的问题。 其次,它通过尺度平移、混合精度聚合和重新正交化来提高数值稳定性。
4.1半精度矩阵向量乘法
按照王等人 [2019],Meanti 等人[2020],我们利用 KeOps Charlier 等[2021] 产生我们的无矩阵方法。 有关 KeOps 的更详细说明,请参阅 Charlier 等人[2021]费迪等人 [2020]。 更具体地说,为了有效地使用共轭梯度,我们需要评估
K
v
Kv
Kv,它可以被构造为
KeOps 通过将矩阵的结果行和列划分为块
B
1
、
⋅
⋅
⋅
B
k
B_1、···B_k
B1、⋅⋅⋅Bk 来执行此乘法的内部循环,然后再减少结果矩阵,并行计算逐行操作。为了利用加速硬件上的并行性,我们将数据分解为块
B
1
、
.
.
.
.
.
,
B
K
B_1、... . . ,B_K
B1、.....,BK 大小
∣
B
k
∣
=
M
<
<
N
|B_k| = M << N
∣Bk∣=M<<N (M 是 192 的 CUDA 块大小) 使得
∪
k
=
1
K
B
k
=
1
,
.
.
.
,
N
∪_{k=1}^KB_k = {1, . . . ,N}
∪k=1KBk=1,...,N。为了计算
K
v
Kv
Kv,我们将运算分解为
K
K
K 个单独的乘积
K
(
k
)
v
(
k
)
K^{(k)}v^{(k)}
K(k)v(k),其中
K
(
k
)
K^{(k)}
K(k) 是由
K
i
,
j
(
k
)
=
a
2
k
(
x
i
,
x
j
)
+
σ
2
δ
i
,
j
K^{(k)}_{i,j} = a^2k (x_i, x_j)+σ^2δ_{i,j}
Ki,j(k)=a2k(xi,xj)+σ2δi,j 组成的 M × M 矩阵对于所有
i
,
j
∈
B
k
i,j ∈ B_k
i,j∈Bk,其中
v
i
(
k
)
=
v
i
v^{(k)}_i = v_i
vi(k)=vi 对于
i
∈
B
k
i ∈ B_k
i∈Bk。然后将每个单独的乘积
K
(
k
)
v
(
k
)
K^{(k)}v^{(k)}
K(k)v(k)计算为正则矩阵向量乘积 (MVM)。我们阐明,无矩阵方法确实涉及块矩阵的创建,但不涉及完整矩阵 K 的创建,但只需要一次在内存中存储 M × M 矩阵。对于精确的矩阵乘积,我们使用更高精度的块求和,这通常用于专用 GPU 硬件 [Zamirai et al., 2021, NVIDIA, 2017]。
为了进一步利用并行性并减少内存开销,我们以半精度计算所有矩阵向量乘积,而不是浮点精度 [Gardner et al., 2018] 或双精度。 与半精度相比,移动到半精度会立即产生大约 2 倍的加速,如图 1a 所示,根据 UCI 存储库 [Dua 和 Graff,2017] 的 180 万个数据点房屋电子数据集上的数据集大小连续变大的数据集进行评估。
图 1:(a) 电上单个矩阵向量乘法的时序。 Half 始终是单精度的两倍左右。 (b) 与各种方案的浮点精度相比的精度。 块求和在半精度内执行每个块,然后以一半或浮点数进行累加,然后将结果转换回一半,而 Kahan 求和则以一半的形式执行所有操作。 使用浮点数和 Kahan 求和的块求和几乎与完全浮点精度的 MVM 一样准确。
此外,如图 1b 所示,当我们以浮点精度执行块求和时,这些矩阵向量乘法非常准确,相对误差平均小于 0.1%。 相比之下,半精度的块求和明显不太准确,而半精度的 Kahan 求和 [Kahan,1965] 几乎与浮点数的块求和一样准确。 如附录所示,所有三种方法都具有相似的速度。 因此,通过移动到半精度,我们立即实现了几乎 2 倍的时间复杂度加速(以及类似的内存减少),而准确度成本可以忽略不计。
但是,范围的限制存在一个小问题,这很容易导致大矩阵向量乘法的数值溢出。 要了解为什么这是有问题的,我们注意到以半精度表示的最大值是
2
13
2^{13}
213 ,而核矩阵向量的输出乘以与核矩阵的大小 N 呈准线性关系,
为了解决这些数值溢出问题,我们改为将每个矩阵乘以
N
−
1
/
2
N^{−1/2}
N−1/2 缩小比例,改为计算
K
~
(
v
/
N
1
/
2
)
\tilde{K}(v/N^{1/2})
K~(v/N1/2),这会产生一个比例为 $N^{1/2} 的上限。 我们可以尝试除以
N
N
N,但是对于最大的数据集甚至不能以半精度存储 N!
4.2 半精度共轭梯度
重新缩放共轭梯度
CG 算法的步长有一个自然的解释:
α
k
α_k
αk 确保我们沿着路径
x
k
+
α
k
p
k
x_k+α_kp_k
xk+αkpk 最小化目标函数,
β
k
β_k
βk 保证搜索方向之间的共轭。 CG 的加速收敛取决于
α
k
α_k
αk 和
β
k
β_k
βk 的正确计算,但是在有限算术中我们不能精确地计算这些量。 更糟糕的是,当我们计算步长
α
k
α_k
αk 和
β
k
β_k
βk 项时,使用半精度计算时舍入误差会被放大
我们需要防止这些术语的舍入和溢出错误,因此我们仅将它们存储在对数刻度中(因为它们在定义上是正数)。 为此,我们利用众所周知的 logsumexp 技巧,该技巧指出对于给定向量
w
w
w 和
z
z
z 使得
w
T
z
>
0
w^T z > 0
wTz>0,为了计算
log
w
T
z
\log w^T z
logwTz 我们使用以下变换
L
Σ
E
(
⋅
)
LΣE (·)
LΣE(⋅) 使得
其中
y
i
=
l
o
g
∣
w
i
∣
+
log
∣
z
i
∣
y_i = log |w_i| + \log |z_i|
yi=log∣wi∣+log∣zi∣ 和
s
i
=
s
i
g
n
(
w
i
z
i
)
s_i = sign(w_iz_i)
si=sign(wizi)。 例如,我们计算并存储
log
α
k
=
L
Σ
E
(
r
k
⊤
z
k
)
−
L
Σ
E
(
d
k
⊤
K
d
k
)
\log α_k = LΣE(r^⊤_k z_k) - LΣE(d^⊤_ k Kd_k)
logαk=LΣE(rk⊤zk)−LΣE(dk⊤Kdk)。
重正交化
在无限精度下,CG 的每个方向向量
d
j
d_j
dj 是 K 正交的,即当
i
≠
j
i\neq j
i=j 时
d
i
T
K
d
j
=
0
d^T_i Kd_j = 0
diTKdj=0,产生相互正交的残差向量
r
k
r_k
rk [Nocedal and Wright, 2006]。 在实践中,由于舍入误差导致残差向量中的收敛速度变慢甚至发散,正交性会丢失 [Gratton et al., 2021, Cutajar et al., 2016]。 为了防止这些舍入误差的累积,我们在 CG 算法中应用了显式的 Gram-Schmidt 重新正交化 [Gratton et al., 2021]。 具体来说,我们重新正交化时间步
k
k
k 的残差向量,将新的残差
r
k
=
K
x
k
−
b
r_k = Kx_k - b
rk=Kxk−b 更新为:
对于
i
=
1
,
.
.
.
,
k
i = 1, . . . , k
i=1,...,k 和其中
u
i
u_i
ui 是从残差
r
i
r_i
ri 创建的前一次迭代的正交向量。 重新正交化会增加内存成本,因为我们必须存储所有先前的残差,为 CG 的
J
J
J 步增加
O
(
J
N
)
O(JN)
O(JN) 的内存复杂度。 然而,我们发现通过重新正交化(和相当高的容差),我们能够在很少的步骤中收敛,通常
J
<
50
J < 50
J<50。
预处理
我们还使用了预处理,特别是使用了一个旋转的 cholesky 预处理器 [Gardner et al., 2018, Harbrecht et al., 2012]。 通过预处理,我们希望找到一个算子,使得
P
−
1
K
≈
I
P^{−1}K ≈ I
P−1K≈I 和
P
−
1
(
v
)
P^{−1}(v)
P−1(v) 相对计算和求解系统
P
−
1
K
v
=
P
−
1
w
P^{−1}Kv = P^{−1}w
P−1Kv=P−1w 而不是
K
v
=
w
Kv = w
Kv=w更加低成本 。
继加德纳等人之后。 [2018],我们使用旋转 Cholesky 分解,它只需要访问矩阵的行(例如
K
i
.
=
K
e
i
K_{i.} = K_{e_i}
Ki.=Kei,其中
e
i
e_i
ei 是零向量,第
i
i
i 个条目等于 1)和近似对角线(在我们的例子中是常数) . 在 K 上运行 k 步旋转 Cholesky 分解会产生近似值
K
~
≈
L
k
L
k
⊤
+
σ
2
I
\tilde K ≈ L_kL^⊤_k + σ^2I
K~≈LkLk⊤+σ2I,然后我们可以使用 Woodbury 矩阵恒等式来构造
P
−
1
P^{-1}
P−1。
内部系统的大小为
k
<
<
N
k << N
k<<N,因此可以使用浮点精度的直接方法(例如 Cholesky 或 QR)求解,因为这些方法不支持半精度的本地 LAPACK。 总而言之,我们修改后的 CG 程序显示在算法 1 中,与 Gardner 等人的 BBMM CG 算法不同。 [2018] 以粗体蓝色字体显示。
4.3 Loss函数
方程1 的计算需要在前向传递中进行昂贵的特征分解或 Lanczos 迭代来计算对数行列式项,这可能非常不稳定 [Gardner et al., 2018]。 相反,我们使用与等式 (2) 具有相同梯度的伪损失。通过同时针对响应
y
y
y 和探测向量
z
i
z_i
zi 求解
K
~
\tilde K
K~。 也就是说,我们使用 算法1 找到线性系统的解决方案。
对于
u
i
.
u_i.
ui. 然后我们分离这些解决方案并简单地计算:
其梯度:
这与等式2,7相同。 然后只需要关于矩阵向量乘积的梯度,这些梯度再次用矩阵向量乘积计算(并且在 KeOps 中本地完成)。等式(7) 不使用 Wenger 等人的方差减少策略。但这可以直接包含在未来。
5.实验
我们的实验分为两部分。 在 5.1 节中,我们经验性地研究了半精度对核矩阵的影响,发现降低精度会增加稀疏性。 在第 5.2 节中,我们研究了降低精度对我们稳定的 CG 求解器的影响,同时在超过 100,000 个数据点的大规模 UCI 基准测试中展示了显着的加速。
5.1 半精度内核矩阵
我们的第一组实验试图理解和量化标准核矩阵不同精度之间的差异。 我们发现较低精度的核矩阵更稀疏,可能会导致加速,但收敛速度较慢的特征值可能会影响泛化。
图 2:双精度(左)、浮点(中)和半精度(右)的 RBF 核矩阵的元素对数。 白色的值表示内核矩阵在这些点处恰好为零,从而产生具有紧凑支持的内核。 紧凑支持的内核产生更有效的内核矩阵 MVM,因为我们可以安全地删除与给定测试点的距离大于数值精度的数据点。
首先,在图 2 中,我们显示了大小为 500 的 RBF 核矩阵的元素对数,因为我们将精度从双精度(左)降低到浮点(中)再降低到半精度。 元素对数使我们可以看到,与高精度内核相比,半精度内核的内核矩阵的值明显更多,表示为零,因此显示为白色。 对于半和浮点内核,我们考虑均匀分布在 [-10, 10] 中的数据,而对于双精度,我们考虑在 [-20, 20] 中的数据。
对于给定的长度尺度,我们可以准确地量化常见固定核矩阵的支持范围。例如,对于长度尺度为
l
l
l 的 Matérn-1/2 核,当
k
(
x
,
x
′
)
=
k
(
d
)
=
exp
{
−
d
/
l
}
<
ϵ
k(x, x′) = k(d) = \exp\{−d/l\} < ϵ
k(x,x′)=k(d)=exp{−d/l}<ϵ (满足
d
>
−
log
ϵ
/
l
d > −\log ϵ/l
d>−logϵ/l) 其中
ϵ
ϵ
ϵ 是在给定精度下的最小可表示值(其他内核的方程在附录中给出。)。 因此,如果我们有
d
=
0.75
d = 0.75
d=0.75(相当于
x
=
−
0.5
x = -0.5
x=−0.5 和
x
′
=
0.25
x' = 0.25
x′=0.25),那么我们不能在对于大于约
9.24
≈
−
log
ϵ
/
0.75
9.24 ≈ −\log ϵ/0.75
9.24≈−logϵ/0.75 的长度尺度,两点之间的半精度。我们在图 3 中显示了几个具有不同长度尺度的内核的双精度、半精度和浮点精度可表示的最大距离,将 5 个单位(代表输入数据的最大距离标准化为均值 0 和方差 1)显示为灰色虚线线供参考。 RBF 核的支持最小,尤其是对短长度尺度,其次是 Matérn 核;有理二次核有很高的支持度。
图 3:Matérn-1/2、RBF 和有理二次 (5) 内核的最大可表示精度之间的距离。 作为参考的灰线显示为 5 个距离单位。 RBF 核具有明显更短的可表示距离,即使对于相当大的长度尺度也是如此,而有理二次核具有迄今为止最高的,尤其是在 double 和 float 中。
这些精确为零的数据点表示产生了一个在实践中具有紧凑支持的内核,尽管理论上具有无限支持。事实上,我们在图 4 中计算 GP 预测均值时展示了这种类型的结构利用的潜力;
图 4:(a)随着数据集大小的增长,完整和截断 MVM 的时间。 截断 MVM 的增长速度明显低于完整 MVM,因为它仅在完整矩阵的四分之一上运行。 (b) 截断 MVM 的误差如预期为零。 为了生成稀疏 MVM,我们截断了任何具有半精度无法表示的内核条目的数据点。
首先,我们在计算完全预测平均值
μ
f
u
l
l
=
K
(
x
∗
,
X
)
v
μ_{full} = K(x^∗,X)v
μfull=K(x∗,X)v 或截断预测平均值之前计算预测平均值缓存
v
=
K
~
−
1
y
v = \tilde K^{−1}y
v=K~−1y ,例如
μ
t
r
u
n
c
.
=
K
(
x
∗
,
X
c
l
o
s
e
)
v
c
l
o
s
e
μ_{trunc.} = K(x∗,X_{close})v_{close}
μtrunc.=K(x∗,Xclose)vclose,在这个例子中大约减少了 3/4 的数据。正如预期的那样,误差为零,表明预测平均值根本没有变化!我们接下来研究跨精度的特征值谱,考虑 Matérn-1/2 内核和从 [-3, 3] 提取的数据,显示图 5a 中的经验谱。半精度内核的特征值在完全下降之前往往会衰减得更慢。然而,最大估计特征值在精度上是相同的(见附录),表明线性系统的条件数在精度上大致相同,因为最小特征值是
σ
2
σ^2
σ2。虽然 CG 收敛的许多界限取决于条件数,但 CG 的收敛有更严格的替代界限,它取决于内部特征值,参见 Nocedal 和 Wright [2006] 的 Thm 5.5,它描述了 CG 收敛的行为从“由外而内”沿着谱(即首先隐式估计最大和最小特征值)[Demmel,1997]。也就是说,半精度核的衰减越慢的特征值谱往往需要更多的 CG 迭代来解决。
图 5:(a) Matérn- 1/2 内核跨精度观察到的特征值谱; 半评估内核的衰减谱比浮点或双精度要慢得多。 (b) 随着矩阵大小的增加,半精度内核的矩阵有效维数明显高于浮点数或半精度。 ED 较高是因为中间特征值高出一半。
此外,内核矩阵在精度上的不同谱,特别是较慢的衰减也可以在方法本身的泛化能力中发挥作用。 这里有一个明确的度量,核矩阵的有效维数,定义为 N e f f ( A , σ 2 ) : = ∑ i = 1 N λ i λ i + σ 2 N_{eff}(A, σ^2 ) := \sum ^N_{i=1} \frac {λ_i} {λ_i+σ^2} Neff(A,σ2):=∑i=1Nλi+σ2λi ,它在核学习的泛化界限中起着重要作用 [Zhang, 2005,卡彭内托和维托,2007]。 我们选择 σ 2 = 0.2 σ^2 = 0.2 σ2=0.2(尽管总体趋势对 σ2 是稳健的)并发现随着 N 的增加,半精度内核的有效维数增加得更快,如图 5b 所示,表明半精度可能具有更差的泛化性能。
5.1.1 直接方法呢?
一个很自然的问题是,看看像 Cholesky 分解这样的直接方法是否可以求解线性系统。 但是,基于 GPU 的线性代数库 (LAPACK) 中不包含对半精度 Cholesky 分解的支持,并且半精度 Cholesky 分解的方法在后端使用迭代技术 [Higham 和 Pranesh,2021]。 我们试图模拟直接方法在求解线性系统中的有效性:通过使用方程。 4 以半精度求解 UCI 蛋白质数据集上的线性系统,直接改变预处理器的大小,结果如图 7 所示。求解产生的残差范数大大大于 1,而预处理器大于 30, 000 失败。
5.2 基准结果 (UCI)
我们使用 RBF ARD 内核和单个 NVIDIA V100 GPU 直接比较使用 KeOps 和 GPyTorch 在来自 UCI 存储库 [Dua 和 Graff,2017 年] 的多达 180 万个数据点的基准上的浮点精度结果,其他内核的结果显示在 附录。 我们发现,仅使用半精度矩阵向量乘法的混合精度策略比等效浮点精度模型产生至少两倍的加速。
残差范数
首先,我们考虑在三个不同大小的数据集上的相对求解误差:elevators、 keggDirected和优化运行结束时收集的超参数的Buzz。我们发现,我们的求解器通常在不到 50 个优化步骤中收敛到小于 0.5 的剩余容差,因为重新正交化和预处理使系统非常容易求解,如图 6a 的电梯和 6b 的 keggDirected 所示,发现运行半精度的 CG 发散得非常快。然而,我们稳定的 CG 实现达到了与浮点精度相似的容差,但需要更多的 MVM。此外,我们展示了无预处理器、等级 50 预处理器和等级 50 预处理器的预处理对 Buzz 的影响,发现预处理对于达到图 6c 中的低残差是必要的。与 Maddox 等人类似。 [2021],我们发现预条件子的大小对求解容差的影响最小,但它是收敛所必需的。
图 6:(a,b)elevators和 keggD 的剩余规范; 标准半精度 CG 失败,而我们的方法在相似数量的迭代中收敛到浮动 CG。 © 在 Buzz 上使用我们的 CG 求解器的残差范数; 使用预处理可以防止由于舍入误差而导致的分歧。 (d) 运行 CG 迭代次数过多会产生舍入误差,这些误差会随着时间的推移而累积,从而产生分歧。 (e) Matérn-1/2、3/2 核往往比 pol 上的 RBF 核收敛得更快。
然而,尽管进行了预处理,但低精度的影响确实会产生随迭代次数增长的舍入误差,从而阻止我们以多次 CG 迭代或低容差(至少在训练时)运行我们的方法。 我们在图 6d 中展示了三个数据集的效果,表明需要做进一步的工作来鼓励收敛。 最后,我们在图 6e 中展示了具有不同 ν 和 RBF 内核的 Matérn 内核的结果,在 pol 上没有 ARD,发现 Matérn-1/2 内核趋向于最快收敛。
优化轨迹
接下来,在图 8 中,我们显示了蛋白质的三个超参数(噪声、输出尺度和长度尺度)的演变,与图 A.5 中的 3d road 显示的图相同。 我们看到求解器产生了足够精确的求解,因此梯度不受影响,并且超参数训练与使用 GPyTorch 求解浮点精度一样稳定。
图 8:标准浮点精度模型和我们的半精度实现的蛋白质优化和参数轨迹。 噪声项(左)和长度尺度(右)在两种精度上都非常接近,而输出尺度(中心)显示出一些细微的差异。
基准测试结果
最后,在表 2 中,我们显示了 5 个种子的浮点和半精度的均方根误差 (RMSE)、负对数似然 (NLL) 和拟合时间,显示了标准偏差。 我们看到,使用半精度需要多达 100、000 个数据点才能加快速度,我们将此归因于 KeOps 中的编译时间增加了一半。 然而,对于歌曲和 houseelectric,我们看到两个最大的数据集在相同数量的优化步骤下显着加速(高达 3 倍)。 有趣的是,我们无法让嗡嗡声以浮点数收敛,而是收敛了一半,这表明较低的精度可能对正则化有所帮助。
6.结论
我们已经证明,由于内存和计算增益的减少,降低精度对于基于 CG 的高斯过程具有相当大的好处,特别是对于超过 100, 000 点的大型数据集。 虽然减少的范围和舍入误差在低精度算术中提出了重大挑战,但我们表明,令人愉悦的简单策略(例如重新正交化和 logsumexp 技巧)可以在很大程度上缓解这些问题。 未来,我们希望通过利用随机舍入误差 [Higham and Mary, 2019] 的特性进一步降低内存开销 [Yang et al., 2021],并将我们的混合精度方法扩展到其他迭代 用于迭代 GP 的方法,例如 Lanczos [Gardner et al., 2018] 和 MINRES [Pleiss et al., 2020]。