在pytorch（或Numpy）中实现此方程的更有效方法

Question

我正在实现此函数的解析形式

其中k（x，y）是RBF内核k(x,y) = exp(-||x-y||^2 / (2h))

我的函数原型是

def A(X, Y, grad_log_px，Kxy):
   pass

和X，Y是NxD矩阵，其中N是批处理大小，D是维。因此，X是上述方程x中大小为N的一批grad_log_px，是我使用NxD计算的一些autograd矩阵。

Kxy是NxN矩阵，其中每个条目(i,j)是RBF内核K(X[i],Y[j])

这里的挑战是，在上式中，y只是维度为D的向量。我有点想传递一批y。 （因此，传递尺寸为Y的矩阵NxD）

使用循环遍历批处理大小可以很好地解决问题，但是我很难以更简洁的方式实现

这是我尝试的循环解决方案：

def A(X, Y, grad_log_px，Kxy):
   res = []
   for i in range(Y.shape[0]):
       temp = 0
       for j in range(X.shape[0]):
           # first term of equation
           temp += grad_log_px[j].reshape(D,1)@(Kxy[j,i] * (X[i] - Y[j]) / h).reshape(1,D)
           temp += Kxy[j,i] * np.identity(D) - ((X[i] - Y[j]) / h).reshape(D,1)@(Kxy[j,i] * (X[i] - Y[j]) / h).reshape(1,D) # second term of equation
       temp /= X.shape[0]

        res.append(temp)
    return np.asarray(res) # return NxDxD array 

在等式中：grad_{x}和grad_{y}均为维度D

Answer 1

鉴于我正确推断了各个术语的所有维度，因此这是一种解决方法。但首先要总结一下尺寸（屏幕截图，因为使用数学类型设置更容易解释；请验证尺寸是否正确）：

还要注意第二项的双导数，得出：

其中下标表示样本，上标表示特征。

因此，我们可以使用np.einsum（类似torch.einsum）和array broadcasting创建两个术语：

grad_y_K = (X[:, None, :] - Y) / h * K[:, :, None]  # Shape: N_x, N_y, D
term_1 = np.einsum('ij,ikl->ikjl', grad_log_px, grad_y_K)  # Shape: N_x, N_y, D_x, D_y
term_2_h = np.einsum('ij,kl->ijkl', K, np.eye(D)) / h  # Shape: N_x, N_y, D_x, D_y
term_2_h2_xy = np.einsum('ijk,ijl->ijkl', grad_y_K, grad_y_K)  # Shape: N_x, N_y, D_x, D_y
term_2_h2 = K[:, :, None, None] * term_2_h2_xy / h**2  # Shape: N_x, N_y, D_x, D_y
term_2 = term_2_h - term_2_h2  # Shape: N_x, N_y, D_x, D_y

然后将结果给出：

(term_1 + term_2).sum(axis=0) / N  # Shape: N_y, D_x, D_y