Product-based Neural Network (PNN) 介绍与源码浅析_product-based neural network(pnn) tensorflow-CSDN博客

本文链接：https://blog.csdn.net/Eric_1993/article/details/107971640

Product-based Neural Network (PNN) 介绍与源码浅析

前言

继续介绍论文~ 本文初看的时候有些懵逼, 多看几次总算有些 Get 到了, 总结一下.

广而告之

可以在微信中搜索 “珍妮的算法之路” 或者 “world4458” 关注我的微信公众号；另外可以看看知乎专栏 PoorMemory-机器学习, 以后文章也会发在知乎专栏中；

Product-based Neural Network (PNN)

文章信息

论文标题: Product-based Neural Networks for User Response Prediction
论文地址: https://arxiv.org/abs/1611.00144
代码地址: https://github.com/Atomu2014/product-nets, 这是作者提供的代码; 另外在 https://github.com/princewen/tensorflow_practice/tree/master/recommendation/Basic-PNN-Demo 我看到了另外的精彩的实现, 是完全按照论文中介绍的公式进行编码的;论文作者提供的代码和论文中的介绍稍有差别.
发表时间: ICDM 2016
论文作者: Yanru Qu, Han Cai, Kan Ren, Weinan Zhang, Yong Yu, Ying Wen, Jun Wang
作者单位: Shanghai Jiao Tong University

请注意

后面在博文中介绍的代码主要以 https://github.com/princewen/tensorflow_practice/tree/master/recommendation/Basic-PNN-Demo 中的为主, 其实现是完全按照论文的思路, 本文原作者的代码实现再其后介绍, 原作者的代码和论文介绍的内容稍有差别, 但是写的很有意思, 也是不能错过的精彩.

核心观点

本文主要介绍了 Product Layer 用于捕获类别特征 (Categorical Features) 的二阶交互特性 (Inter-field Feature Interaction). 就具体实现来说, 主要实现了基于向量 Inner-Product 的 IPNN 网络以及基于向量 Outer-Product 的 OPNN 网络. 这两个网络处理在 Product Layer 中使用的向量交互方式不同之外, 其他结构完全相同.

核心观点介绍

下面来看具体的网络结构:

将网络分为三个部分来看, 分别是:

Embedding Layer: 将 [Field 1, ..., Field N] 等 N 个输入的类别特征转换为稠密向量, 分别为 [Feature 1, ..., Feature N];
Product Layer: 这里面实际上要分为两个部分, 其中 z 表示的内容实际上仍然是一阶特征, 并未进行特征交互, 而 p 表示的内容由本文提出的 Product Layer 生成, 以获取特征间的相关性;
MLP: 上一层得到的 z 和 p 进行 concat 后 (注意在论文的公式中 z 和 p 的结果是相加起来), 再输入到一个 MLP 中学习高阶特征交互, 并输出最终的预估结果.

符号介绍

在做进一步介绍之前, 先明确一下之后会用到的符号, 以方便讨论:

$N$ : 每个样本类别特征的个数, 也就是常说的 Field Size
$M$ : Embedding Size 的大小, 也就是说, Embedding Layer 会将每个类别特征映射为 $M$ 维的稠密向量
$\bm{f}_i$ : 第 $i$ 个类别特征所对应的稠密向量, 对每个样本来说, Embedding Layer 将类别特征分别映射为 $(\bm{f}_1, \bm{f}_2, \ldots, \bm{f}_N)\in\mathbb{R}^{N\times M}$
操作符 $\odot$ : 表示两个 tensor 的 inner product, 定义如下:

$\boldsymbol{A} \odot \boldsymbol{B} \triangleq \sum_{i, j} \boldsymbol{A}_{i, j} \boldsymbol{B}_{i, j}$

上式为对 $\boldsymbol{A}$ 和 $\boldsymbol{B}$ 进行 element-wise 的乘法.

(突然发现有的符号其实一眼看过去就能明白啥函数, 算了就不描述了 🤣)

下面着重介绍 Product Layer 的操作. 分为 $z$ 和 $p$ 部分的计算.

网络结构图中 z 部分的计算

网络结构图中, Product Layer 的输出主要包含两个部分, 左边部分 (即 $z$ 部分), 假设输出结果用 $\bm{l}_z$ 表示, 那么

$\bm{l}_{z}=\left(l_{z}^{1}, l_{z}^{2}, \ldots, l_{z}^{n}, \ldots, l_{z}^{D_{1}}\right), \quad l_{z}^{n}=\boldsymbol{W}_{z}^{n} \odot \boldsymbol{z}$

其中 $\bm{z} = (\bm{f}_1, \bm{f}_2, \ldots, \bm{f}_N)$ , 那么:

$\boldsymbol{W}_{z}^{n} \odot \boldsymbol{z} = \sum_{i=1}^{N} \sum_{j=1}^{M} \left(\boldsymbol{W}_{z}^{n}\right)_{i, j} \boldsymbol{z}_{i, j}$

注意到 $\boldsymbol{W}_{z}^{n}\in\mathbb{R}^{N\times M}$ , 那么 $\boldsymbol{W}_{z} = \left(\boldsymbol{W}_{z}^{1}, \boldsymbol{W}_{z}^{2}, \ldots, \boldsymbol{W}_{z}^{n}, \ldots \boldsymbol{W}_{z}^{D_1}\right)\in\mathbb{R}^{D_1\times N\times M}$

从公式来看, 会有些抽象, 我们转化成代码来看:

# B * N
feat_index = tf.placeholder(tf.int32,
                          shape=[None,None],
                          name='feat_index')

# B * N * M
embeddings = tf.nn.embedding_lookup(weights['feature_embeddings'], feat_index)

## embeddings 的大小为 B * N * M
## weights['product-linear'] 的大小为 D1 * N * M
linear_output = []
for i in range(D1):
    linear_output.append(
    	tf.reshape(
        	tf.reduce_sum(
        		tf.multiply(embeddings,
        		        weights['product-linear'][i]), # B * N * M
				axis=[1, 2]), # B
           shape=(-1, 1)
        ) # B * 1
    )

lz = tf.concat(linear_output, axis=1) # B * D1

网络结构图中 p 部分的计算

这部分的计算是本文的核心. 在网络结构图中, 右边部分 (即 $p$ 部分) 的输出结果使用 $\bm{l}_p$ 表示:

$\bm{l}_{p}=\left(l_{p}^{1}, l_{p}^{2}, \ldots, l_{p}^{n}, \ldots, l_{p}^{D_{1}}\right), \quad l_{p}^{n}=\boldsymbol{W}_{p}^{n} \odot \boldsymbol{p}$

其中 $\boldsymbol{p} = \{\bm{p}_{ij}\}, i = 1 \ldots N, j = 1 \ldots N$ . 注意 $\bm{p}\in\mathbb{R}^{N\times N}$ 是一个矩阵, 里面的元素 $\bm{p}_{ij} = g(\bm{f}_i, \bm{f}_j)$ 定义了两个特征 pair 的相关性, 或者说交叉特性.

于是现在的重点是如何对 $g(\cdot)$ 进行设计. 当我们得到 $\bm{p}$ 后, 就可以通过下式计算出输出结果:

$\boldsymbol{W}_{p}^{n} \odot \boldsymbol{p} = \sum_{i=1}^{N} \sum_{j=1}^{N} \left(\boldsymbol{W}_{z}^{n}\right)_{i, j} \boldsymbol{p}_{i, j}$

针对 $g(\cdot)$ 的设计, 本文考虑了两种方法, 分别是计算两个特征的内积和外积.

Inner Product

基于内积的 $g(\cdot)$ 定义为:

$g(\bm{f}_i, \bm{f}_j) = \left\langle \bm{f}_{i}, \bm{f}_{j}\right\rangle$

之后计算输出结果 $\bm{l}_{p}$ 时, 为了减少时间和空间复杂度, 作者借鉴了 FM 的思路, 利用矩阵分解, 设计 $\boldsymbol{W}_{p}^{n} = \boldsymbol{\theta}^{n} \boldsymbol{\theta}^{n T}$ , 其中 $\boldsymbol{\theta}^{n}\in\mathbb{R}^N$ , 此时可以简化 $\bm{l}_{p}$ 的计算如下:

$\bm{l}_{p}^{n} = \boldsymbol{W}_{p}^{n} \odot \boldsymbol{p}=\sum_{i=1}^{N} \sum_{j=1}^{N} \theta_{i}^{n} \theta_{j}^{n}\left\langle\boldsymbol{f}_{i}, \boldsymbol{f}_{j}\right\rangle=\left\langle\sum_{i=1}^{N} \boldsymbol{\delta}_{i}^{n}, \sum_{i=1}^{N} \boldsymbol{\delta}_{i}^{n}\right\rangle$

其中 $\boldsymbol{\delta}^{n}_{i} = \theta_{i}^{n} \boldsymbol{f}_{i}\in\mathbb{R}^{M}$ , 它表示每个特征向量 $\boldsymbol{f}_{i}$ 都有一个对应的参数 $\theta_{i}^{n}$ 进行加权, 此外还可以得到 $\boldsymbol{\delta}^{n}=\left(\boldsymbol{\delta}_{1}^{n}, \boldsymbol{\delta}_{2}^{n}, \ldots, \boldsymbol{\delta}_{i}^{n}, \ldots, \boldsymbol{\delta}_{N}^{n}\right) \in \mathbb{R}^{N \times M}$ .

由于 $\bm{l}_{p}^{n} = \left\langle\sum\limits_{i=1}^{N} \boldsymbol{\delta}_{i}^{n}, \sum\limits_{i=1}^{N} \boldsymbol{\delta}_{i}^{n}\right\rangle = \left\|\sum\limits_{i} \boldsymbol{\delta}_{i}^{n}\right\|$ , 此时我们可以得到 $\bm{l}_{p}$ 最后的结果:

$\boldsymbol{l}_{p}=\left(\left\|\sum_{i} \boldsymbol{\delta}_{i}^{1}\right\|, \ldots,\left\|\sum_{i} \boldsymbol{\delta}_{i}^{n}\right\|, \ldots,\left\|\sum_{i} \boldsymbol{\delta}_{i}^{D_{1}}\right\|\right)$

下面继续来看看代码:

## embeddings 的大小为 B * N * M
## weights['product-quadratic-inner'] 的大小为 D1 * N
quadratic_output = []
for i in range(D1):
	theta = tf.multiply(embeddings,  # B * N * M
				tf.reshape(weights['product-quadratic-inner'][i],
							(1, -1, 1)) # 1 * N * 1
			) # N * N * M 
	quadratic_output.append(
		tf.reshape(
			tf.norm(
				tf.reduce_sum(theta, axis=1), # B * M, 注意这里是对 N 个特征向量求和得到 theta
				axis=1), # B, 这里是对 theta 求范数
			shape=(-1, 1)
		) # B * 1
	)  
lp = tf.concat(quadratic_output, axis=1) # B * D1

Outer Product

基于外积的 $g(\cdot)$ 定义为:

$g\left(\boldsymbol{f}_{i}, \boldsymbol{f}_{j}\right)=\boldsymbol{f}_{i} \boldsymbol{f}_{j}^{T}$

注意注意, 此时 $\bm{p}_{ij} = g\left(\boldsymbol{f}_{i}, \boldsymbol{f}_{j}\right)=\boldsymbol{f}_{i} \boldsymbol{f}_{j}^{T}\in\mathbb{R}^{M\times M}$ 是一个矩阵而不是一个数值了. 而之后计算 $\bm{p} = \{\bm{p}_{ij}\}, i = 1 \ldots N, j = 1 \ldots N$ 时, 为了减少时间复杂度和空间复杂度, 作者采用了 superposition (重叠) 的思路, 即将 $N\times N$ 个大小为 $M\times M$ 的矩阵给累加起来 (要知道, 如果按照前面内积的思路, 把这里的 $\bm{p}_{ij}$ 作为矩阵中的一个元素, 那么整个矩阵的大小就是 $(N\times M)\times (N\times M)$ , 显然是不可接受的, 因此使用重叠的思路降低复杂度).

用公式体现如下:

$\boldsymbol{p}=\sum_{i=1}^{N} \sum_{j=1}^{N} \bm{f}_{i} \boldsymbol{f}_{j}^{T}=\boldsymbol{f}_{\Sigma}\left(\boldsymbol{f}_{\Sigma}\right)^{T}, \quad \boldsymbol{f}_{\Sigma}=\sum_{i=1}^{N} \boldsymbol{f}_{i}$

这里得到的 $\bm{p}\in\mathbb{R}^{M\times M}$ , 此时 $\boldsymbol{W}_{p}\in\mathbb{R}^{D_1\times M\times M}$ .

下面来看看具体的代码实现:

## embeddings 的大小为 B * N * M
## weights['product-quadratic-outer'] 的大小为 D1 * M * M
quadratic_output = [] 
embedding_sum = tf.reduce_sum(embeddings, axis=1) ## B * M, 即公式中对 f 进行累加
p = tf.matmul(tf.expand_dims(embedding_sum, 2), # B * M * 1
			  tf.expand_dims(embedding_sum, 1) # B * 1 * M
			  ) # B * M * M, 做外积
for i in range(D1):
	theta = tf.multiply(p, # B * M * M
				tf.expand_dims(weights['product-quadratic-outer'][i], 0) # 1 * M * M
				) # B * M * M
    quadratic_output.append(
    			tf.reshape(tf.reduce_sum(theta, axis=[1, 2]), # B
    					  shape=(-1, 1)) # B * 1
    			)
    			
lp = tf.concat(quadratic_output, axis=1) # B * D1

Product Layer 的输出

按照论文中的介绍, Product Layer 最终的输出为:

$\bm{l}_{1}=\operatorname{relu}\left(\bm{l}_{z}+\bm{l}_{p}+\bm{b} \right)$

其中 $\bm{b}$ 为 bias, 而 $\bm{l}_z$ 和 $\bm{l}_p$ 上文已经介绍过了.

代码实现如下:

## lz 和 lp 大小均为 B * D1
## weights['product-bias'] 大小为 D1
y_deep = tf.nn.relu(tf.add(tf.add(lz, lp), 
						   weights['product-bias']
						   ))

至此, PNN 的核心思路已经介绍好了~ 鼓掌~~👏👏👏

原作者的代码介绍

下面再简单分析下原作者提供的代码: https://github.com/Atomu2014/product-nets, 其实现和论文中的公式还是有很多区别的, 但是实践过程中, 这些实现还是有很大的参考意义的. 另外作者的实现看起来特别精炼, 代码很值得学习. 作者总共实现了 LR, FM, DeepFM, FNN, CCPM, PNN1, PNN2 等 7 个算法. 注意, PNN1 和 PNN2 不是分别代表 inner-product 和 outer-product, 其中 PNN2 中完整实现了 inner-product 和 outer-product 方法, 可以先看看 PNN1 进行学习理解, 再来看 PNN2 会轻松一些.

下面主要介绍 PNN2 的实现.

(2020-08-14 TODO: 放心, 不是 Flag, 先去搬砖. 🤣🤣🤣 )

(2020-08-14 夜, 来填坑啦, 😁😁😁)

PNN2 的实现在 https://github.com/Atomu2014/product-nets/blob/master/python/models.py 中:

初始化部分代码如下:

class PNN2(Model):
    def __init__(self, field_sizes=None, embed_size=10, layer_sizes=None, layer_acts=None, drop_out=None,
                 embed_l2=None, layer_l2=None, init_path=None, opt_algo='gd', learning_rate=1e-2, random_seed=None,
                 layer_norm=True, kernel_type='mat'):
        Model.__init__(self)
        init_vars = []
        num_inputs = len(field_sizes)
        for i in range(num_inputs):
            init_vars.append(('embed_%d' % i, [field_sizes[i], embed_size], 'xavier', dtype))
        num_pairs = int(num_inputs * (num_inputs - 1) / 2)
        node_in = num_inputs * embed_size + num_pairs
        if kernel_type == 'mat':
            init_vars.append(('kernel', [embed_size, num_pairs, embed_size], 'xavier', dtype))
        elif kernel_type == 'vec':
            init_vars.append(('kernel', [num_pairs, embed_size], 'xavier', dtype))
        elif kernel_type == 'num':
            init_vars.append(('kernel', [num_pairs, 1], 'xavier', dtype))
        for i in range(len(layer_sizes)):
            init_vars.append(('w%d' % i, [node_in, layer_sizes[i]], 'xavier', dtype))
            init_vars.append(('b%d' % i, [layer_sizes[i]], 'zero',  dtype))
            node_in = layer_sizes[i]

其中: (下面描述中用到的数学符号会完全和论文中的一致, 以便对比)

init_vars: 用来保存所有的权重参数的配置, 之后会调用 utils.init_var_map 函数, 使用 init_vars 来构建权重参数; 其中名字为 embed_%d 的参数是特征对应的 embedding layer 的参数, 名字为 kernel 的参数是 Product 层的参数, 名字为 w%d 的参数是 MLP 层的权重参数, 名字为 b%d 的参数是 MLP 层的 bias;
field_sizes 是一个列表, 里面存储着各个 Field 下 unique 的特征个数, 对一个具体的样本来说, 它拥有 $N$ 个特征, 那么 $N$ 就等于 len(field_sizes), 那么 num_inputs 的大小就是 $N$ ;
num_inputs: 就是 Field 的个数 $N$ ;
num_pairs: $N$ 个特征两两组合, 总共可以组合出 $\frac{(N - 1)\times N}{2}$ 个结果; 同时 num_pairs 大小等于论文中的 $D_1$ , 即 Product Layer 的输出节点个数.
node_in: 它的大小为 MLP 的输入节点个数. 包含两个部分: num_inputs * embed_size, 用论文中的符号就是 $N\times M$ , 提前说明一下, 这相当于将 $N$ 个特征向量进行 concat 得到的结果; 另一个部分是 num_pairs, 即 $D_1$ , 为 Product Layer 的输出结果; 从这里可以看出, 作者的实现并未完全按照论文中来, 论文在设计输入到 MLP 的结果时, 使用的是:

$\bm{l}_{1}=\operatorname{relu}\left(\bm{l}_{z}+\bm{l}_{p}+\bm{b} \right)$

其中 $\bm{l}_{z}$ 和 $\bm{l}_{p}$ 是相加; 而这个代码实现中, 使用的是 concat([lz, lp]), 此外, $\bm{l}_z$ 的输出大小也不是 $D_1$ , 而是 $N\times M$ , 相当于将 $(\bm{f}_1, \bm{f}_2, \ldots, \bm{f}_N)\in\mathbb{R}^{N\times M}$ 进行 concat.

kernel_type: 这是一个很重要的参数, 如果 kernel_type == 'mat', 那么将采用 Outer Product Layer, 此外权重的大小被设置为 [embed_size, num_pairs, embed_size], 翻译为论文中的符号就是 $M\times D_1\times M$ ; 这里其实会问这样一个问题, 为啥不设置为 $D_1\times M\times M$ 呢? 这样不就刚好和论文中介绍的 $W_p$ 契合了? 实际上, 看完作者的代码后, 我认为设置为 $D_1\times M\times M$ 也是没有问题的, 具体原因后面分析代码时我再介绍. 如果 kernel_type == 'vec' 或者 num 的时候, 此时将采用 Inner Product Layer, 具体实现时和论文有点出入, 分析代码时再说.

继续分析:

with self.graph.as_default():
     if random_seed is not None:
         tf.set_random_seed(random_seed)
     self.X = [tf.sparse_placeholder(dtype) for i in range(num_inputs)]
     self.y = tf.placeholder(dtype)
     self.keep_prob_train = 1 - np.array(drop_out)
     self.keep_prob_test = np.ones_like(drop_out)
     self.layer_keeps = tf.placeholder(dtype)
     self.vars = utils.init_var_map(init_vars, init_path)
     w0 = [self.vars['embed_%d' % i] for i in range(num_inputs)]
     xw = tf.concat([tf.sparse_tensor_dense_matmul(self.X[i], w0[i]) for i in range(num_inputs)], 1)
     xw3d = tf.reshape(xw, [-1, num_inputs, embed_size])

self.X: 接受 $N$ 个 Field 的结果, 使用 SparseTensor 来表示;
self.vars: 使用前面的 init_vars 作为参数, 生成各种权重参数;
w0: 保存每个 Field 所对应的 Embedding Layer
xw: 首先, tf.sparse_tensor_dense_matmul(self.X[i], w0[i]), 相当于获取每个特征所对应的 embedding, 如果设 Batch 的大小为 $B$ , 那么其结果为 $B\times M$ , 这样的话, xw 的 shape 为 $B\times (N\times M)$
xw3d: 对 xw 进行 reshape, 得到的结果大小为 $B\times N\times M$

下面高能来了:

(注释来自作者, 其中 batch 相当于我们这的 $B$ ; pair 相当于我们这的 num_pairs, 即 $D_1$ ; 而 $k$ 为 embedding 的大小, 相当于我们这的 $M$ . 不要被几个符号给搞混了, 要清楚了解其含义, 这才是核心和重点)

row = []
col = []
for i in range(num_inputs - 1):
    for j in range(i + 1, num_inputs):
        row.append(i)
        col.append(j)
# batch * pair * k
p = tf.transpose(
    # pair * batch * k
    tf.gather(
        # num * batch * k
        tf.transpose(
            xw3d, [1, 0, 2]),
        row),
    [1, 0, 2])
# batch * pair * k
q = tf.transpose(
    tf.gather(
        tf.transpose(
            xw3d, [1, 0, 2]),
        col),
    [1, 0, 2])

# b * p * k
p = tf.reshape(p, [-1, num_pairs, embed_size])
# b * p * k
q = tf.reshape(q, [-1, num_pairs, embed_size])

row 和 col 分别保存特征 pair $(\bm{f}_i, \bm{f}_j)$ 的下标 $i$ 和 $j$ , 由于特征两两组合总共有 $\frac{(N - 1)\times N}{2}$ 种, 所以它们俩的大小都是 $\frac{(N - 1)\times N}{2}$ ;
p 的操作相当于将 $\bm{f}_i$ 给取出来
q 的操作相当于将 $\bm{f}_j$ 给取出来

下面接着高能:

k = self.vars['kernel']
if kernel_type == 'mat':
   # batch * 1 * pair * k
   p = tf.expand_dims(p, 1)
   # batch * pair
   kp = tf.reduce_sum(
       # batch * pair * k
       tf.multiply(
           # batch * pair * k
           tf.transpose(
               # batch * k * pair
               tf.reduce_sum(
                   # batch * k * pair * k
                   tf.multiply(
                       p, k),
                   -1),
               [0, 2, 1]),
           q),
       -1)
else:
   # 1 * pair * (k or 1)
   k = tf.expand_dims(k, 0)
   # batch * pair
   kp = tf.reduce_sum(p * q * k, -1)

kernel_type == 'mat' 部分是 Outer Product Layer 的结果, 而 else 的部分是 Inner Product Layer 的结果. 上面的代码需要体会, 咱们从最简单的方式入手.

首先, 对于特征对 $(\bm{f}_i, \bm{f}_j)$ , 它们的 Outer Product 可以使用 $\bm{f}_i\bm{f}_j^T$ 进行计算, 假设 $\bm{f}_i = [p_1, p_2]$ , $\bm{f}_j = [q_1, q_2]$ , 那么 Outer Product 的结果是:

$\bm{f}_i\bm{f}_j^T = \left[\begin{matrix} p_1 \\ p_2 \end{matrix}\right]\cdot\left[\begin{matrix}q_1, q_2\end{matrix}\right] = \left[\begin{matrix} p_1q_1 & p_1q_2 \\ p_2q_1 & p_2q_2 \end{matrix}\right]$

我们要注意到的是, 在计算 $l_p^n$ 时, 需要将 $\bm{W}_p^n$ 也考虑上, 同时还需要考虑:

$l_p^n = \bm{W}_p^n\odot\bm{p}$

$\odot$ 的结果无非是加权求和, 因此上面的结果可能表示为:

$\begin{aligned} p_1q_1w_{11} &+ p_1q_2 w_{12} + \\ p_2q_1w_{21} &+ p_2q_2w_{22} \end{aligned}$

可以进一步表示为:

$(p_1w_{11} + p_2w_{21}) \times q_1 + (p_1w_{12} + p_2w_{22}) \times q_2$

从上式可以看出, 可以先用 $\bm{f}_i$ 和权重 $\bm{W}_p^n$ 进行相乘, 然后再来和 $\bm{f}_j$ 进行相乘. 这就是作者代码需要做的事情. 总之可以理解为, 最后的效果就是先计算出了 $\bm{f}_i\bm{f}_j^T$ , 然后再乘上 $\bm{W}_p^n$ .

此外, 可以注意到作者代码中, 权重 $\bm{W}_p$ 大小设置为 $M\times D_1\times M$ 了, 但实际上, 写成 $D_1\times M\times M$ 也是没有问题的, 只是在计算时, 需要做些简单的变换; 为了证实我这一观点, 用 PyTorch 写了个简单的程序验证了一下:

import torch
import torch.nn as nn

## p 的大小为 P * K, 这个大写的 P 相当于作者代码中的 pair, 也是论文中的 D1 参数
p = torch.arange(4).view(2, 2).float() # P * K
q = torch.arange(1, 5).view(2, 2).float() # P * K
w = torch.arange(8).view(2, 2, 2).float() # P * K * K

"""求解方式一"""
## 如果 w 权重为 P * K * K, 那么按照下面的方式进行求解
p1 = torch.unsqueeze(p, 2)
d = torch.sum(p1 * w, dim=1) * q # P * K
print(d)

"""求解方式二"""
## 如果 w 权重为 K * P * K, 那么按照下面的方式求解, 这个求解顺序和作者代码中的顺序相同
w1 = w.permute(2, 0, 1) # K * P * K
p1 = torch.unsqueeze(p, 0) # 1 * P * K
d = torch.sum(p1 * w1, dim=-1).permute(1, 0) * q
print(d)

"""结果"""
"""均为"""
tensor([[  2.,   6.],
        [ 78., 124.]])