Более эффективный способ реализовать это уравнение в 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) и широковещательная рассылка :

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