如何在PyTorch中通过广播与批次操作计算向量间的外积

你此前使用的torch.einsum('bij, bj->bij', x, y)实现的是逐元素点乘逻辑，而非外积，因为外积需要将两个输入向量的维度分别映射到输出的不同轴上，因此输出维度会比输入多一维。

正确实现方式 #

方法1：使用einsum实现 #

外积需要为两个输入的特征维度单独分配输出轴，正确的einsum表达式如下：

import torch

batch_size = 2
seq_len = 2
dim = 3
x = torch.rand([batch_size, seq_len, dim])
y = torch.rand([batch_size, dim])

out_einsum = torch.einsum('bsd, bd -> bsdd', x, y)

输出out_einsum的形状为[batch_size, seq_len, dim, dim]，完全匹配你示例中循环调用torch.outer的结果。

方法2：使用广播机制实现 #

如果不使用einsum，也可以通过维度扩充后相乘的方式实现，逻辑更直观：

# 给x扩充最后一维：形状变为 [batch_size, seq_len, dim, 1]
x_expand = x.unsqueeze(-1)
# 给y扩充中间两个维度：形状变为 [batch_size, 1, 1, dim]
y_expand = y.unsqueeze(1).unsqueeze(1)
# 广播相乘得到外积结果
out_broadcast = x_expand * y_expand

两种方法得到的结果完全等价，你可以通过如下代码验证和你示例中循环输出的结果一致：

print(torch.allclose(out_einsum[0,0], torch.outer(x[0][0], y[0])))
print(torch.allclose(out_einsum[0,1], torch.outer(x[0][1], y[0])))
print(torch.allclose(out_einsum[1,0], torch.outer(x[1][0], y[1])))
print(torch.allclose(out_einsum[1,1], torch.outer(x[1][1], y[1])))

内容的提问来源于stack exchange，提问作者YQ.Wang