向量化高斯过程中的 ARD（自动相关性确定）内核实现

Question

我正在尝试使用 GPML 书中给出的 NumPy 实现 ARD 内核（公式 5.2 中的 M3）。

我正在为 NxM 核计算对这个方程进行矢量化而苦苦挣扎。我尝试了以下非矢量化版本。有人可以帮助在 NumPy/PyTorch 中对此进行矢量化吗？

import numpy as np
N = 30 # Number of data points in X1
M = 40 # Number of data points in X2
D = 6  # Number of features (ARD dimensions)

X1 = np.random.rand(N, D)
X2 = np.random.rand(M, D)
Lambda = np.random.rand(D, 1)
L_inv = np.diag(np.random.rand(D))
sigma_f = np.random.rand()

K = np.empty((N, M))
for n in range(N):
  for m in range(M):
    M3 = Lambda@Lambda.T + L_inv**2
    d = (X1[n,:] - X2[m,:]).reshape(-1,1)
    K[n, m] = sigma_f**2 * np.exp(-0.5 * d.T@M3@d)

Answer 1

我们可以使用broadcasting and the neat NumPy function einsum的规则来向量化数组运算。简而言之，广播允许我们通过向结果数组添加新维度来对单行数组进行操作，而 einsum 允许我们通过明确地使用索引符号（而不是矩阵）来对多个数组执行操作.

幸运的是，计算内核不需要循环。请看下面的矢量化解决方案 ARD_kernel 函数，它在我的机器上比原来的 loopy 版本快大约 30 倍。现在，einsum 通常是最快的，但是可能有更快的方法，我没有检查任何其他东西（例如通常的 @ 运算符而不是 einsum）。

此外，代码中缺少一个术语（Kronecker delta），我不知道它是否是故意省略的（如果您在实现它时遇到问题，请告诉我，我会编辑答案） .

import numpy as np
N = 300 # Number of data points in X1
M = 400 # Number of data points in X2
D = 6  # Number of features (ARD dimensions)

np.random.seed(1)  # Fix random seed for reproducibility
X1 = np.random.rand(N, D)
X2 = np.random.rand(M, D)
Lambda = np.random.rand(D, 1)
L_inv = np.diag(np.random.rand(D))
sigma_f = np.random.rand()

# Loopy function
def ARD_kernel_loops(X1, X2, Lambda, L_inv, sigma_f):
    K = np.empty((N, M))
    M3 = Lambda@Lambda.T + L_inv**2
    for n in range(N):
        for m in range(M):
            d = (X1[n,:] - X2[m,:]).reshape(-1,1)
            K[n, m] = np.exp(-0.5 * d.T@M3@d)
    return K * sigma_f**2
    
# Vectorized function
def ARD_kernel(X1, X2, Lambda, L_inv, sigma_f):
    M3 = Lambda.squeeze()*Lambda + L_inv**2  # Use broadcasting to avoid transpose
    d = X1[:,None] - X2[None,...]            #  Use broadcasting to avoid loops
    
    # order=F for memory layout (as your arrays are (N,M,D) instead of (D,N,M))
    return sigma_f**2 * np.exp(-0.5 * np.einsum("ijk,kl,ijl->ij", d, M3, d, order = 'F'))

Answer 2

可能还有一个额外的优化。给出的M个矩阵的例子都是正定的。这意味着可以应用 Cholesky 分解，我们可以找到上三角 U 以便

M = U'*U

重点是如果我们将 U 应用于 xs，那么

y[p] = U*x[p] p=1..

然后

(x[p]-x[q])'*M*(x[p]-x[q]) = (y[p]-y[q])'*(y[p]-y[q])

因此，如果有 N 个向量 x，每个维度为 d，我们将 LHS 上的 N 平方 O(d 平方) 操作转换为 RHS 上的 N 平方 O(d) 操作

这需要额外的 choleski 分解 (O(d cubed)) 和 N O( d 平方) U 对 xs 的应用。

向量化高斯过程中的 ARD（自动相关性确定）内核实现

Vectorizing ARD (Automatic Relevance Determination) kernel implementation in Gaussian processes

numpy

pytorch

vectorization

matrix-multiplication

gaussian-process