如何使用einsum批量计算矩阵B中每个行向量的xAx^T？

我们需要对堆叠在矩阵B（维度(m,n)）中的每个行向量x，计算二次型xAx^T（其中A是(n,n)矩阵），最终得到维度为(m,1)的结果向量。

实现代码 #

直接用EinSum可以一步完成计算，无需循环：

import torch
m = 30
n = 4
B = torch.randn(m, n)
A = torch.randn(n, n)

# EinSum实现，直接得到(m,1)结果
einsum_result = torch.einsum('bi,ij,bj->b1', B, A, B)

# 原循环方法用于验证
result = torch.zeros(m,1)
for i in range(m):
    x = B[i].unsqueeze(0)
    result[i] = torch.matmul(x, torch.matmul(A, x.T))

# 验证结果一致性
print(torch.allclose(einsum_result, result))  # 输出 True

表达式说明 #

EinSum的维度标记'bi,ij,bj->b1'对应逻辑：

• bi：矩阵B的维度，b代表第b个行向量（共m个），i代表特征维度（共n个）
• ij：矩阵A的维度，连接B的i维度和自身的j维度
• bj：再次引用矩阵B，用j维度和A的j维度对应
• ->b1：指定输出为(m,1)维度，直接匹配需求的结果形状

如果不需要直接输出(m,1)，也可以写成'bi,ij,bj->b'得到(m,)的向量，再通过unsqueeze(1)转换维度。

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