爱因斯坦求和约定

前言

今天在看一个算法的代码中，出现了tf.einsum()这个函数，之前没见过，所以查了下，居然是一块自己缺失的知识——爱因斯坦求和约定，赶紧恶补一下。知乎上有一个提问说——爱因斯坦求和约定除了增加歧义有任何好处吗，看来有些人对这个用法有不少疑惑，问题答案中很多答主们都在为有这么一个方便的标记法而庆幸，是他们开发深度学习模型中最喜欢用的函数。

einsum记法释义

在深度学习中，我们会碰到诸如点积、外积、转置、矩阵-向量乘法、矩阵-矩阵乘法等各种计算，当然这些在numpy、keras、tensorflow或者pytorch中都可以很简单地计算，但是有没有想过有一个函数或者方法可以同时做这些事情，是不是很优雅呢。没错，einsum记法是一个表达以上这些运算，包括复杂张量运算在内的优雅方式，基本上，可以把einsum看成一种领域特定语言。
在einsum约定中，省略了求和符号 \(\sum\)，因为它隐式地累加重复的下标和输出中未指名的下标。

比如我们要将两个矩阵\(A_{ik}\)和\(B_{kj}\)相乘，然后按列求和得到向量\(c\)。常用表达可表示为

\[\sum_{i}\sum_{k}A_{ik}B_{kj} \]

用einsum标记可写为：

\[c_{j}=A_{ik}B_{kj} \]

注意看上面的写法，等号右边有重复下标k，所以计算时它会隐式地累加重复下标k，等号左边表示输出结果，因为它只写了下标j，也即缺少下标i，所以它也会隐式地累加下标i，即输出中未指名的下标。

einsum在深度学习框架或numpy中的使用

einsum在numpy中实现为np.einsum，在PyTorch中实现为torch.einsum，在TensorFlow中实现为tf.einsum。我们以tf.einsum为例。

tf.einsum(equation, *inputs, **kwargs)

其中equation是表示爱因斯坦求和约定的字符串，而inputs则是张量序列。比如上面举例中，equation可以表示为ik,kj->j。这里(i, j, k)的命名是任意的，但需要一致。假设张量A和B分别用a, b来表示，则tensorflow中可以讲上述举例用einsum表示为

tf.einsum('ik,kj->ij', a, b)

应用举例

我们以tensorflow来举例

矩阵转置

\[B_{ij}=A_{ji} \]

a = tf.reshape(tf.range(6), shape=(2,3))
print(a)
tf.einsum('ij->ji', a)

output

tf.Tensor(
[[0 1 2]
 [3 4 5]], shape=(2, 3), dtype=int32)
<tf.Tensor: shape=(3, 2), dtype=int32, numpy=
array([[0, 3],
       [1, 4],
       [2, 5]])>

求和

\[b=\sum_{i}\sum_{j}A_{ij}=A_{ij} \]

a = tf.reshape(tf.range(6), shape=(2,3))
print(a)
tf.einsum('ij->', a)

output

tf.Tensor(
[[0 1 2]
 [3 4 5]], shape=(2, 3), dtype=int32)
<tf.Tensor: shape=(), dtype=int32, numpy=15>

列求和

\[ b_j = \sum_{i}A_{ij}=A_{ij} \]

a = tf.reshape(tf.range(6), shape=(2,3))
print(a)
tf.einsum('ij->j', a)

output

tf.Tensor(
[[0 1 2]
 [3 4 5]], shape=(2, 3), dtype=int32)
<tf.Tensor: shape=(3,), dtype=int32, numpy=array([3, 5, 7])>

行求和

\[ b_i = \sum_{j}A_{ij}=A_{ij} \]

a = tf.reshape(tf.range(6), shape=(2,3))
print(a)
tf.einsum('ij->i', a)

output

tf.Tensor(
[[0 1 2]
 [3 4 5]], shape=(2, 3), dtype=int32)
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([ 3, 12])>

矩阵和向量相乘

\[ c_{i}=\sum_{k}A_{ik}b_{k}=A_{ik}b_{k} \]

a = tf.reshape(tf.range(6), shape=(2,3))
print(a)
b = tf.range(3)
print(b)
tf.einsum('ij,j->i', a, b)

output

tf.Tensor(
[[0 1 2]
 [3 4 5]], shape=(2, 3), dtype=int32)
tf.Tensor([0 1 2], shape=(3,), dtype=int32)
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([ 5, 14])>

矩阵和矩阵相乘

\[ C_{ij}=\sum_{k}A_{ik}B_{kj}=A_{ik}B_{kj} \]

a = tf.reshape(tf.range(6), shape=(2,3))
print(a)
b = tf.reshape(tf.range(15), shape=(3,5))
print(b)
tf.einsum('ik,kj->ij', a, b)

output

tf.Tensor(
[[0 1 2]
 [3 4 5]], shape=(2, 3), dtype=int32)
tf.Tensor(
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]], shape=(3, 5), dtype=int32)
<tf.Tensor: shape=(2, 5), dtype=int32, numpy=
array([[ 25,  28,  31,  34,  37],
       [ 70,  82,  94, 106, 118]])>

点积

\[ c=\sum_{i}a_{i}b_{i}=a_{i}b_{i} \]

a = tf.range(3)
print(a)
b = tf.range(3, 6)
print(b)
tf.einsum('i,i->', a, b)

output

tf.Tensor([0 1 2], shape=(3,), dtype=int32)
tf.Tensor([3 4 5], shape=(3,), dtype=int32)
<tf.Tensor: shape=(), dtype=int32, numpy=14>

哈达玛积

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

a = tf.reshape(tf.range(6), (2,3))
print(a)
b = tf.reshape(tf.range(6, 12), (2,3))
print(b)

output

tf.Tensor(
[[0 1 2]
 [3 4 5]], shape=(2, 3), dtype=int32)
tf.Tensor(
[[ 6  7  8]
 [ 9 10 11]], shape=(2, 3), dtype=int32)
<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
array([[ 0,  7, 16],
       [27, 40, 55]])>

外积

\[C_{ij}=a_{i}b_{j}=a_{i}b_{j} \]

a = tf.range(3)
print(a)
b = tf.range(3, 6)
print(b)
tf.einsum('i,j->ij', a, b)

output

tf.Tensor([0 1 2], shape=(3,), dtype=int32)
tf.Tensor([3 4 5], shape=(3,), dtype=int32)
<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[ 0,  0,  0],
       [ 3,  4,  5],
       [ 6,  8, 10]])>

batch矩阵相乘

\[ C_{ijl}=\sum_{k}A_{ijk}B_{ikl}=A_{ijk}B_{ikl} \]

a = tf.random.normal([3,2,5])
b = tf.random.normal([3,5,4])
c = tf.einsum('ijk,ikl->ijl', a, b)
c.shape

output

TensorShape([3, 2, 4])

张量缩约

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，如下所示：

\[C_{pstuv}=\sum_{q}\sum_{r}A_{pqrs}B_{tuqvr}=A_{pqrs}B_{tuqvr} \]

a = tf.random.normal([2,3,4,5])
b = tf.random.normal([7,8,3,6,4])
c = tf.einsum('pqrs,tuqvr->pstuv', a, b)
c.shape

output

TensorShape([2, 5, 7, 8, 6])

文章参考：

1、https://blog.csdn.net/zzq060143/article/details/89107567

2、https://ajcr.net/Basic-guide-to-einsum/

3、https://rockt.github.io/2018/04/30/einsum

4、https://obilaniu6266h16.wordpress.com/2016/02/04/einstein-summation-in-numpy/