愛因斯坦求和約定

前言

今天在看一個算法的代碼中，出現了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/