einsum滿足你一切需要:深度學習中的愛因斯坦求和約定

編者按：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）。

如果你像我一樣，發現記住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.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 == 1 else 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)

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

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