tf.einsum—爱因斯坦求和约定

1. einsum记法

如果你像我一样，发现记住PyTorch/TensorFlow中那些计算点积、外积、转置、矩阵-向量乘法、矩阵-矩阵乘法的函数名字和签名很费劲，那么einsum记法就是我们的救星。einsum记法是一个表达以上这些运算，包括复杂张量运算在内的优雅方式，基本上，可以把einsum看成一种领域特定语言。一旦你理解并能利用einsum，除了不用记忆和频繁查找特定库函数这个好处以外，你还能够更迅速地编写更加紧凑、高效的代码。而不使用einsum的时候，容易出现引入不必要的张量变形或转置运算，以及可以省略的中间张量的现象。此外，einsum这样的领域特定语言有时可以编译到高性能代码，事实上，PyTorch最近引入的能够自动生成GPU代码并为特定输入尺寸自动调整代码的张量理解（Tensor Comprehensions）就基于类似einsum的领域特定语言。此外，可以使用opt einsum和tf einsum opt这样的项目优化einsum表达式的构造顺序。

比方说，我们想要将两个矩阵AϵRI×KA\epsilon\mathbb{R}^{I\times K}AϵRI×K和BϵRK×JB\epsilon\mathbb{R}^{K \times J}BϵRK×J相乘，接着计算每列的和，最终得到向量cϵRJc \epsilon\mathbb{R}^JcϵRJ。使用爱因斯坦求和约定，这可以表达为：
cj=∑i∑kAikBkj=AikBkjc_j=\sum_i\sum_kA_{ik}B_{kj}=A_{ik}B_{kj} cj=i∑k∑AikBkj=AikBkj

这一表达式指明了ccc中的每个元素cic_ici是如何计算的，列向量Ai:A_{i:}Ai:乘以行向量B:jB_{:j}B:j，然后求和。注意，在爱因斯坦求和约定中，我们省略了求和符号∑\sum∑，因为我们隐式地累加重复的下标（这里是k）和输出中未指明的下标（这里是i）。

在深度学习中，我经常碰到的一个问题是，变换高阶张量到向量。例如，我可能有一个张量，其中包含一个batch中的N个训练样本，每个样本是一个长度为T的K维词向量序列，我想把词向量投影到一个不同的维度Q。如果将这个张量记作TϵRN×T×KT\epsilon\mathbb{R}^{N\times T\times K}TϵRN×T×K，将投影矩阵记作WϵRK×QW\epsilon\mathbb{R}^{K\times Q}WϵRK×Q，那么所需计算可以用einsumeinsumeinsum表达为：
Cntq=∑kTntkWkq=TntkWkqC_{ntq}=\sum_kT_{ntk}W_{kq}=T_{ntk}W_{kq}Cntq=k∑TntkWkq=TntkWkq

最后一个例子，比方说有一个四阶张量TϵRN×T×K×MT\epsilon\mathbb{R}^{N\times T\times K\times M}TϵRN×T×K×M，我们想要使用之前的投影矩阵将第三维投影至QQQ维，并累加第二维，然后转置结果中的第一维和最后一维，最终得到张量CϵRM×Q×NC\epsilon \mathbb{R}^{M\times Q\times N}CϵRM×Q×N。einsum可以非常简洁地表达这一切：
Cmqn=∑t∑kTntkmWkq=TntkmWkqC_{mqn}=\sum_t\sum_kT_{ntkm}W_{kq}=T_{ntkm}W_{kq}Cmqn=t∑k∑TntkmWkq=TntkmWkq
注意，我们通过交换下标n和m（Cmqn而不是Cnqm），转置了张量构造结果。

2. Numpy、PyTorch、TensorFlow中的einsum

einsum在numpy中实现为np.einsum，在PyTorch中实现为torch.einsum，在TensorFlow中实现为tf.einsum，均使用一致的签名einsum(equation, operands)，其中equation是表示爱因斯坦求和约定的字符串，而operands则是张量序列（在numpy和TensorFlow中是变长参数列表，而在PyTorch中是列表）。

例如，我们的第一个例子，cj=∑i∑kAikBkj=AikBkjc_j=\sum_i\sum_kA_{ik}B_{kj}=A_{ik}B_{kj} cj=i∑k∑AikBkj=AikBkj写成equation字符串就是ik,kj -> j。注意这里(i, j, k)的命名是任意的，但需要一致。

PyTorch和TensorFlow像numpy支持einsum的好处之一是einsum可以用于神经网络架构的任意计算图，并且可以反向传播。典型的einsum调用格式如下：

上式中◻是占位符，表示张量维度。上面的例子中，arg1和arg3是矩阵，arg2是二阶张量，这一einsum运算的结果（result）是矩阵。注意einsum处理的是可变数量的输入。在上面的例子中，einsum指定了三个参数之上的操作，但它同样可以用在牵涉一个参数、两个参数、三个以上参数的操作上。学习einsum的最佳途径是通过学习一些例子，所以下面我们将展示一下，在许多深度学习模型中常用的库函数，用einsum该如何表达（以PyTorch为例）。

1 矩阵转置
Bji=AijB_{ji}=A_{ij}Bji=Aij

import torch
a = torch.arange(6).reshape(2, 3)
torch.einsum('ij->ji', [a])
tensor([[ 0.,  3.],[ 1.,  4.],[ 2.,  5.]])

2 求和
b=∑i∑jAij=Aijb=\sum_i\sum_jA_{ij}=A_{ij}b=i∑j∑Aij=Aij

a = torch.arange(6).reshape(2, 3)
torch.einsum('ij->', [a])
tensor(15.)

3 列求和
bj=∑iAij=Aijb_j=\sum_iA_{ij}=A_{ij}bj=i∑Aij=Aij

a = torch.arange(6).reshape(2, 3)
torch.einsum('ij->j', [a])
tensor([ 3.,  5.,  7.])

4 行求和
bj=∑jAij=Aijb_j=\sum_jA_{ij}=A_{ij}bj=j∑Aij=Aij

a = torch.arange(6).reshape(2, 3)
torch.einsum('ij->i', [a])
tensor([  3.,  12.])

5 矩阵-向量相乘
ci=∑kAikbk=Aikbkc_i=\sum_kA_{ik}b_k=A_{ik}b_kci=k∑Aikbk=Aikbk

a = torch.arange(6).reshape(2, 3)
b = torch.arange(3)
torch.einsum('ik,k->i', [a, b])
tensor([  5.,  14.])

6 矩阵-矩阵相乘
Cij=∑kAikBkj=AikBkjC_{ij}=\sum_kA_{ik}B_{kj}=A_{ik}B_{kj}Cij=k∑AikBkj=AikBkj

a = torch.arange(6).reshape(2, 3)
b = torch.arange(15).reshape(3, 5)
torch.einsum('ik,kj->ij', [a, b])
tensor([[  25.,   28.,   31.,   34.,   37.],[  70.,   82.,   94.,  106.,  118.]])

7 点积
c=∑iaibi=aibic=\sum_ia_ib_i=a_ib_ic=i∑aibi=aibi

a = torch.arange(3)
b = torch.arange(3,6)  # [3, 4, 5]
torch.einsum('i,i->', [a, b])
tensor(14.)

8 哈达玛积
Cij=AijBijC_{ij}=A_{ij}B_{ij}Cij=AijBij

a = torch.arange(6).reshape(2, 3)
b = torch.arange(6,12).reshape(2, 3)
torch.einsum('ij,ij->ij', [a, b])
tensor([[  0.,   7.,  16.],[ 27.,  40.,  55.]])

9 外积
Cij=aibjC_{ij}=a_ib_jCij=aibj

a = torch.arange(3)
b = torch.arange(3,7)
torch.einsum('i,j->ij', [a, b])
tensor([[  0.,   0.,   0.,   0.],[  3.,   4.,   5.,   6.],[  6.,   8.,  10.,  12.]])

10 batch矩阵相乘
Cijl=∑kAijkBikl=AijkBiklC_{ijl}=\sum_kA_{ijk}B_{ikl}=A_{ijk}B_{ikl}Cijl=k∑AijkBikl=AijkBikl

a = torch.randn(3,2,5)b = torch.randn(3,5,3)torch.einsum('ijk,ikl->ijl', [a, b])tensor([[[ 1.0886,  0.0214,  1.0690],[ 2.0626,  3.2655, -0.1465]],[[-6.9294,  0.7499,  1.2976],[ 4.2226, -4.5774, -4.8947]],[[-2.4289, -0.7804,  5.1385],[ 0.8003,  2.9425,  1.7338]]])

11 张量缩约
batch矩阵相乘是张量缩约的一个特例。比方说，我们有两个张量，一个n阶张量A ∈ ℝI1 × ⋯ × In，一个m阶张量B ∈ ℝJ1 × ⋯ × Jm。举例来说，我们取n = 4，m = 5，并假定I2 = J3且I3 = J5。我们可以将这两个张量在这两个维度上相乘（A张量的第2、3维度，B张量的3、5维度），最终得到一个新张量C ∈ ℝI1 × I4 × J1 × J2 × J4，如下所示：
Cpstuv=∑q∑rApqrsBtuqvr=ApqrsBtuqvrC_{pstuv}=\sum_q\sum_rA_{pqrs}B_{tuqvr}=A_{pqrs}B_{tuqvr}Cpstuv=q∑r∑ApqrsBtuqvr=ApqrsBtuqvr

a = torch.randn(2,3,5,7)b = torch.randn(11,13,3,17,5)torch.einsum('pqrs,tuqvr->pstuv', [a, b]).shapetorch.Size([2, 7, 11, 13, 17])

12 双线性变换
Dij=∑k∑lAikBjklCil=AikBjklCilD_{ij}=\sum_k\sum_lA_{ik}B_{jkl}C_{il}=A_{ik}B_{jkl}C_{il}Dij=k∑l∑AikBjklCil=AikBjklCil

a = torch.randn(2,3)
b = torch.randn(5,3,7)
c = torch.randn(2,7)
torch.einsum('ik,jkl,il->ij', [a, b, c])
tensor([[ 3.8471,  4.7059, -3.0674, -3.2075, -5.2435],[-3.5961, -5.2622, -4.1195,  5.5899,  0.4632]])

3.总结

einsum是一个函数走天下，是处理各种张量操作的瑞士军刀。话虽如此，“einsum满足你一切需要”显然夸大其词了。从上面的真实用例可以看到，我们仍然需要在einsum之外应用非线性和构造额外维度（unsqueeze）。类似地，分割、连接、索引张量仍然需要应用其他库函数。

使用einsum的麻烦之处是你需要手动实例化参数，操心它们的初始化，并在模型中注册这些参数。不过我仍然强烈建议你在实现模型时，考虑下有哪些情况适合使用einsum.

from