PyTorch下4阶因子分解机替代嵌套循环的张量求和优化方法

数据老陈

2026-10-1

PyTorch 4阶因子分解机四阶项高效实现方案

四重Python循环是效率最低的实现方式，PyTorch提供了多个经过底层优化的张量运算接口，可大幅提升计算速度，同时原生支持自动微分和GPU加速。

方案1：`torch.einsum`实现（最推荐，可读性最高）

爱因斯坦求和接口可以直接匹配四阶项的求和公式，代码和数学表达式一一对应，不易出错：

import torch

# 输入张量定义
# x: 形状为 (n,) 的输入特征向量
# t: 形状为 (n, n, n, n) 的4阶权重张量
# u: 形状为 (n,) 的一阶项权重向量
# w: 形状为 (n, n) 的二阶项上三角权重矩阵

# 完整计算输出y
y = u @ x + x @ w @ x.T + torch.einsum('i,j,k,l,ijkl->', x, x, x, x, t)

注：你提供的二阶项代码中变量名存在笔误，描述里二阶权重为n×n的上三角矩阵w，代码中写为v，上述代码已修正对应关系。

该方案优势：

代码简洁，完全贴合数学公式，维护成本低
PyTorch内部会对einsum做运算融合优化，速度远高于Python层面的四重循环

天然支持批处理：如果输入是批次数据（x形状为(batch_size, n)），仅需调整einsum索引规则即可一次性计算全批次结果：

# 批处理版本计算，输出y形状为 (batch_size,)
y = x @ u + torch.einsum('bi,ij,bj->b', x, w, x) + torch.einsum('bi,bj,bk,bl,ijkl->b', x, x, x, x, t)

方案2：广播+逐元素乘+求和实现

如果对einsum语法不熟悉，也可以用张量广播机制实现，性能和einsum基本持平：

# 四阶项单独计算
x_expand = x[:, None, None, None] * x[None, :, None, None] * x[None, None, :, None] * x[None, None, None, :]
y_4 = (t * x_expand).sum()
# 完整y计算
y = u @ x + x @ w @ x.T + y_4

该方案适合需要查看四阶交叉项中间计算结果的场景。