PyTorch einsum最优路径复杂度及自定义层反向传播选型

问题背景 #

我需要执行如下张量收缩操作：

np.einsum('ijk,kl,jlm', x, y, z, optimize = 'optimal')

通过Numpy测试发现，针对我的数据，最优路径几乎始终是先收缩k维度，再收缩j维度，最后收缩l维度。我希望让PyTorch自动寻找最优路径，但查阅文档看到：

当输入张量至少为3个时会触发该优化，否则顺序不影响。需注意，寻找最优路径是NP难问题，因此opt_einsum依赖不同启发式算法得到近似最优结果

关于最优路径的疑惑 #

我尚未完全理解最优路径的本质，有以下疑问：

• 为什么寻找张量收缩的最优路径是NP难问题？它本质上等价于排序算法吗？
• 是哪些因素（比如张量数量、收缩次数、总索引数）导致其复杂度呈非多项式缩放？

反向传播的实现选择 #

此外，我要将该结构整合到自定义PyTorch层中，可学习参数存储在张量z中。为便于反向传播，我是否应该将乘积拆分（即PyTorch中等价于如下代码的操作）：

np.einsum('ijk,jkl', np.einsum('ijk,kl', x, y, optimize = 'optimal'), z, optimize = 'optimal')

还是应该坚持使用自动寻找的最优路径？

内容的提问来源于stack exchange，提问作者DPurple