爱因斯坦求和约定Einsum

爱因斯坦求和约定(einsum)

爱因斯坦求和约定由爱因斯坦提出，是一种简洁的表示张量运算的方式。它的核心思想是：当一个指标在一个表达式中出现两次时，意味着对该指标进行求和。

我们对于两个矩阵\(A_{ij}\)与\(B_{jk}\)的矩阵乘法的公式一般如下：

\[C_{ik} = \sum_{j} A_{ij} B_{jk} \]

有下面的例子

\[\begin{bmatrix} 1\quad2\quad3 \\ 4\quad5\quad6 \\ \end{bmatrix} \begin{bmatrix} 1 \\ 2 \\ 3 \end{bmatrix}= \begin{bmatrix} 14 \\ 32 \end{bmatrix} \\ 其中，C_{ik} = C_{21},则其计算过程为： \\ c_{11} = A的第一行 @ B的第一列 \\ c_{21} = A的第二行 @ B的第一列 \]

在求和约定下，重复的指标\(j\)表示对\(j\)进行求和，我们省略求和符号吗，简化为：

\[C_{ik} = A_{ij} B_{jk} \]

torch中的torch.einsum

在pytorch中，torch.einsum函数提供了一个强大的工具来执行爱因斯坦求和约定的操作。它允许我们使用字符串表示张量的索引，从而简化复杂的张量运算。

上述的矩阵乘法，使用torch.einsum可以表示为：

import torch
A = torch.tensor([[1, 2, 3], [4, 5, 6]]) #(2, 3)
B = torch.tensor([[1], [2], [3]]) #(3, 1)

C = torch.einsum('ij,jk->ik', A, B)
print(C)
'''
C = tensor([[14],
            [32]]) #(2, 1)
'''

即对于该函数，其参数为

torch.einsum(equation, *operands)
'''
- equation: 一个字符串，表示张量的索引和操作。箭头左侧为输入的索引，若有两个输入用逗号隔开；箭头右侧为输出的索引。

- operands: 一个或多个张量，表示参与运算的张量。
'''

不难看出，该函数所做的事情很简单：

1. 我们自己找到矩阵乘法的爱因斯坦求和约定公式：\(C_{ik} = A_{ij} B_{jk}\)
2. 表达式左侧为输入的下表，上面的例子下是ij,jk，表示输入的两个张量的下标
3. 表达式右侧为输出的下标，上面的例子下是ik，表示输出的张量的下标

换一种直观的理解即为：对两个矩阵，沿着输出没有的那个方向去做那个方向上的向量内积，并沿着这个方向相加，不断遍历所有的i与j所得出的输出即为输出

下面给出一些表达式的例子：

表达式	数学含义	运算
ij->ji	\(y_{ji}=x_{ij}\)	转置
ij->i	\(y_i=\sum_j x_{ij}\)	每行求和
ij->	\(y=\sum_{i,j}x_{ij}\)	所有元素求和，得到标量
i,i->	\(y=\sum_i a_i b_i\)	向量点积
i,j->ij	\(y_{ij}=a_i b_j\)	向量外积
ij,ij->ij	\(y_{ij}=a_{ij}b_{ij}\)	逐元素相乘
ik,kj->ij	\(y_{ij}=\sum_k a_{ik}b_{kj}\)	矩阵乘法
bik,bkj->bij	\(y_{bij}=\sum_k a_{bik}b_{bkj}\)	批量矩阵乘法

示例

1. 转置

\(y_{ji} = x_{ij}\)

A = torch.tensor([[1, 2, 3], 
                  [4, 5, 6]]) #(2, 3)
torch.einsum('ij->ji', A)
print(A)
'''
A = tensor([1, 4],
           [2, 5],
           [3, 6]]) #(3, 2) 
'''

2. 每行求和

\(y_i = \sum_j x_{ij}\)

A = torch.tensor([[1, 2, 3], 
                  [4, 5, 6]]) #(2, 3)
torch.einsum('ij->i', A) # j这个维度在输出消失了，因此对每个i，将所有j的值相加
print(A)
'''
A = tensor([6],
           [15]) #(2,)
'''

3. 所有元素求和，得到标量

\(y = \sum_{i,j} x_{ij}\)

A = tensor([[1, 2, 3], 
            [4, 5, 6]]) #(2, 3)
torch.einsum('ij->', A) # i,j这两个维度在输出消失了，因此对所有i,j的值相加
print(A)
'''
A = tensor(21) #标量
'''

4. ...的使用

在torch的einsum函数中，可以使用...来表示任意数量的前导维度，这对于处理批量数据或高维张量非常有用。...可以匹配任意数量的维度，使得我们可以在不明确指定所有维度的情况下进行操作。

比如对于一个高维度的张量

\[W_{bsik} \]

我们当前想要做的操作是按照最后一维即k去求和，即
\(Out_{bsi}=\sum_k W_{bsik}\)

# 使用randn随机初始化一个维度为(2, 3, 4, 5)的张量W
W = torch.randn(2, 3, 4, 5)  # (2, 3, 4, 5)
Out = torch.einsum('...k->...', W)  # 对最后一维k求和, ...为省略的前置维度
# 等价于 Out = torch.einsum('bsik->bsi', W)
print(Out.shape)  # 输出形状为 (2, 3, 4)