編者按：FAIR研究科學家Tim Rockt?schel簡要介紹了einsum表示法的概念，并通過真實例子展示了einsum的表達力。

當我和同事聊天的時候，我意識到不是所有人都了解einsum，我開發深度學習模型時最喜歡的函數。本文打算改變這一現狀，讓所有人都了解它！愛因斯坦求和約定（einsum）在numpy和TensorFlow之類的深度學習庫中都有實現，感謝Thomas Viehmann，最近PyTorch也實現了這一函數。關于einsum的背景知識，我推薦閱讀Olexa Bilaniuk的numpy的愛因斯坦求和約定以及Alex Riley的einsum基本指南。這兩篇文章介紹了numpy中的einsum，我的這篇文章則將演示在編寫優雅的PyTorch/TensorFlow模型時，einsum是多么有用（我將使用PyTorch作為例子，不過很容易就可以翻譯到TensorFlow）。

1. einsum記法

如果你像我一樣，發現記住PyTorch/TensorFlow中那些計算點積、外積、轉置、矩陣-向量乘法、矩陣-矩陣乘法的函數名字和簽名很費勁，那么einsum記法就是我們的救星。einsum記法是一個表達以上這些運算，包括復雜張量運算在內的優雅方式，基本上，可以把einsum看成一種領域特定語言。一旦你理解并能利用einsum，除了不用記憶和頻繁查找特定庫函數這個好處以外，你還能夠更迅速地編寫更加緊湊、高效的代碼。而不使用einsum的時候，容易出現引入不必要的張量變形或轉置運算，以及可以省略的中間張量的現象。此外，einsum這樣的領域特定語言有時可以編譯到高性能代碼，事實上，PyTorch最近引入的能夠自動生成GPU代碼并為特定輸入尺寸自動調整代碼的張量理解（Tensor Comprehensions）就基于類似einsum的領域特定語言。此外，可以使用opt einsum和tf einsum opt這樣的項目優化einsum表達式的構造順序。

比方說，我們想要將兩個矩陣A ∈ ?I × K和B ∈ ?K × J相乘，接著計算每列的和，最終得到向量c ∈ ?J。使用愛因斯坦求和約定，這可以表達為：

這一表達式指明了c中的每個元素ci是如何計算的，列向量Ai:乘以行向量B:j，然后求和。注意，在愛因斯坦求和約定中，我們省略了求和符號Sigma，因為我們隱式地累加重復的下標（這里是k）和輸出中未指明的下標（這里是i）。當然，einsum也能表達更基本的運算。比如，計算兩個向量a, b ∈ ?J的點積可以表達為：

在深度學習中，我經常碰到的一個問題是，變換高階張量到向量。例如，我可能有一個張量，其中包含一個batch中的N個訓練樣本，每個樣本是一個長度為T的K維詞向量序列，我想把詞向量投影到一個不同的維度Q。如果將這個張量記作T ∈ ?N × T × K，將投影矩陣記作W ∈ ?K × Q，那么所需計算可以用einsum表達為：

最后一個例子，比方說有一個四階張量T ∈ ?N × T × K × M，我們想要使用之前的投影矩陣將第三維投影至Q維，并累加第二維，然后轉置結果中的第一維和最后一維，最終得到張量C ∈ ?M × Q × N。einsum可以非常簡潔地表達這一切：

注意，我們通過交換下標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寫成equation字符串就是ik,kj -> j。注意這里(i, j, k)的命名是任意的，但需要一致。

PyTorch和TensorFlow像numpy支持einsum的好處之一是einsum可以用于神經網絡架構的任意計算圖，并且可以反向傳播。典型的einsum調用格式如下：

上式中?是占位符，表示張量維度。上面的例子中，arg1和arg3是矩陣，arg2是二階張量，這一einsum運算的結果（result）是矩陣。注意einsum處理的是可變數量的輸入。在上面的例子中，einsum指定了三個參數之上的操作，但它同樣可以用在牽涉一個參數、兩個參數、三個以上參數的操作上。學習einsum的最佳途徑是通過學習一些例子，所以下面我們將展示一下，在許多深度學習模型中常用的庫函數，用einsum該如何表達（以PyTorch為例）。

2.1 矩陣轉置

import torch

a = torch.arange(6).reshape(2, 3)

torch.einsum('ij->ji', [a])

tensor([[ 0., 3.],

[ 1., 4.],

[ 2., 5.]])

2.2 求和

a = torch.arange(6).reshape(2, 3)

torch.einsum('ij->', [a])

tensor(15.)

2.3 列求和

a = torch.arange(6).reshape(2, 3)

torch.einsum('ij->j', [a])

tensor([ 3., 5., 7.])

2.4 行求和

a = torch.arange(6).reshape(2, 3)

torch.einsum('ij->i', [a])

tensor([ 3., 12.])

2.5 矩陣-向量相乘

a = torch.arange(6).reshape(2, 3)

b = torch.arange(3)

torch.einsum('ik,k->i', [a, b])

tensor([ 5., 14.])

2.6 矩陣-矩陣相乘

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.]])

2.7 點積

向量：

a = torch.arange(3)

b = torch.arange(3,6) # [3, 4, 5]

torch.einsum('i,i->', [a, b])

tensor(14.)

矩陣：

a = torch.arange(6).reshape(2, 3)

b = torch.arange(6,12).reshape(2, 3)

torch.einsum('ij,ij->', [a, b])

tensor(145.)

2.8 哈達瑪積

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.]])

2.9 外積

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.]])

2.10 batch矩陣相乘

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]]])

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

a = torch.randn(2,3,5,7)

b = torch.randn(11,13,3,17,5)

torch.einsum('pqrs,tuqvr->pstuv', [a, b]).shape

torch.Size([2, 7, 11, 13, 17])

2.12 雙線性變換

如前所述，einsum可用于超過兩個張量的計算。這里舉一個這方面的例子，雙線性變換。

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. 案例

3.1 TreeQN

我曾經在實現TreeQN（ arXiv:1710.11417）的等式6時使用了einsum：給定網絡層l上的低維狀態表示zl，和激活a上的轉換函數Wa，我們想要計算殘差鏈接的下一層狀態表示。

在實踐中，我們想要高效地計算大小為B的batch中的K維狀態表示Z ∈ ?B × K，并同時計算所有轉換函數（即，所有激活A）。我們可以將這些轉換函數安排為一個張量W ∈ ?A × K × K，并使用einsum高效地計算下一層狀態表示。

import torch.nn.functional as F

def random_tensors(shape, num=1, requires_grad=False):

tensors = [torch.randn(shape, requires_grad=requires_grad) for i in range(0, num)]

return tensors[0] if num == 1else tensors

# 參數

# -- [激活數 x 隱藏層維度]

b = random_tensors([5, 3], requires_grad=True)

# -- [激活數 x 隱藏層維度 x 隱藏層維度]

W = random_tensors([5, 3, 3], requires_grad=True)

def transition(zl):

# -- [batch大小 x 激活數 x 隱藏層維度]

return zl.unsqueeze(1) + F.tanh(torch.einsum("bk,aki->bai", [zl, W]) + b)

# 隨機取樣仿造輸入

# -- [batch大小 x 隱藏層維度]

zl = random_tensors([2, 3])

transition(zl)

3.2 注意力

讓我們再看一個使用einsum的真實例子，實現注意力機制的等式11-13（arXiv:1509.06664）：

用傳統寫法實現這些可要費不少力氣，特別是考慮batch實現。einsum是我們的救星！

# 參數

# -- [隱藏層維度]

bM, br, w = random_tensors([7], num=3, requires_grad=True)

# -- [隱藏層維度 x 隱藏層維度]

WY, Wh, Wr, Wt = random_tensors([7, 7], num=4, requires_grad=True)

# 注意力機制的單次應用

def attention(Y, ht, rt1):

# -- [batch大小 x 隱藏層維度]

tmp = torch.einsum("ik,kl->il", [ht, Wh]) + torch.einsum("ik,kl->il", [rt1, Wr])

Mt = F.tanh(torch.einsum("ijk,kl->ijl", [Y, WY]) + tmp.unsqueeze(1).expand_as(Y) + bM)

# -- [batch大小 x 序列長度]

at = F.softmax(torch.einsum("ijk,k->ij", [Mt, w]))

# -- [batch大小 x 隱藏層維度]

rt = torch.einsum("ijk,ij->ik", [Y, at]) + F.tanh(torch.einsum("ij,jk->ik", [rt1, Wt]) + br)

# -- [batch大小 x 隱藏層維度], [batch大小 x 序列維度]

return rt, at

# 取樣仿造輸入

# -- [batch大小 x 序列長度 x 隱藏層維度]

Y = random_tensors([3, 5, 7])

# -- [batch大小 x 隱藏層維度]

ht, rt1 = random_tensors([3, 7], num=2)

rt, at = attention(Y, ht, rt1)

4. 總結

einsum是一個函數走天下，是處理各種張量操作的瑞士軍刀。話雖如此，“einsum滿足你一切需要”顯然夸大其詞了。從上面的真實用例可以看到，我們仍然需要在einsum之外應用非線性和構造額外維度（unsqueeze）。類似地，分割、連接、索引張量仍然需要應用其他庫函數。

使用einsum的麻煩之處是你需要手動實例化參數，操心它們的初始化，并在模型中注冊這些參數。不過我仍然強烈建議你在實現模型時，考慮下有哪些情況適合使用einsum.