Einsum (Einstein Summation)

einsum 연산은 Einstein Summation Convention에 따라 연산을 진행하는 방법이라고 한다. 위키에서 Einstein Notation에 대한 내용을 살펴보면 "한 항에 동일한 첨자가 윗첨자와 아랫첨자로 한 번씩 짝을 지어 나타날 경우, (마치 합의 기호가 항의 앞에 있을 때처럼) 해당 첨자가 가질 수 있는 모든 값에 대해 항의 값을 전부 더하는 것으로 이해한다."고 나타나 있다. 즉, 특정 index의 집합에 대한 합 연산(일반적인 ∑index set 연산) 을 간결하게 표시하는 방법이다.

 

einsum 연산을 통해 행렬, 벡터의 내적(dot products), 외적(outer products), 전치(transpose), 대각합(trace), 행렬곱(multiplication) 등을 일관성있게 표현할 수 있다.

 

einsum 연산은 numpy(np.einsum), torch(torch.einsum), tensorflow(tf.einsum)과 같이 자주 사용하는 연산 라이브러리에 모두 구현되어 있으며, 세 경우 모두 einsum(equation, operands)와 같이 인자로 equation과 operands를 받는다.

 

result = numpy.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', optimize=False)

 

• subscripts(str): 쉼표로 구분된 인덱스 레이블 목록으로 합계를 위한 인덱스를 지정한다. 명시적으로 '->' 와 정확한 출력 형식의 인덱스 레이블이 포함되지 않는 한 암시적(고전적인 아인슈타인 합산) 계산이 수행된다. 
• *operands: 해당 연산을 수행할 대상들

예를들어 np.einsum("ij,jk->ik", A, B)라고 하면, 텐서 A는 i,j 인덱스를 갖고, B는 j,k 인덱스를 갖으며, "ij, jk->ik"는 텐서 A와 B의 연산의 결과가 인덱스 i,k가 된다는 의미이며, 결과적으로 A와 B의 곱 연산을 수행하라는 뜻이다. 만약 명시적으로 연산 결과에 대한 인덱스를 지정하지 않았다면, np.einsum("ij,jk->", A, B)은 아인슈타인 합산이 계산되어 텐서의 원소에 대한 합이 결과로 반환된다.

 

In [1]:

import numpy as np

a = np.arange(25).reshape(5,5)
a

 

Out[1]:

array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

 

행렬의 합

 

In [2]:

np.einsum('ij->', a)  # 행렬의 모든 원소의 합을 구한다.

 

Out[2]:

300

 

'ij'는 행렬의 모든 행과 열을 뜻하며 화살표를 사용하면 행렬의 모든 원소의 합을 계산한다.

 

In [3]:

a.sum()

 

Out[3]:

300

 

In [4]:

np.einsum('ij->i', a)  # 각 행의 합을 구한다.

 

Out[4]:

array([ 10,  35,  60,  85, 110])

 

'ij->i'는 행렬의 모든 행과 열에 대해 행을 기준으로 합을 구하라는 의미이다. 앞서 설명했듯이 Einstein notation은 특정 인덱스 집합에 대한 합연산을 간결하게 표시한 것이라고 하였다.

 

In [5]:

a.sum(axis=1)   # np.sum(a, axis=1)

 

Out[5]:

array([ 10,  35,  60,  85, 110])

 

행렬 곱

두 행렬의 곱을 계산한다. np.dot(a,b)와 np.matmul(a,b)와 동일하다.

 

In [6]:

a = np.array([[1,2],[3,4]])
b = np.array([[2,3],[4,5]])

 

In [7]:

np.matmul(a,b)  # 두 행렬의 곱

 

Out[7]:

array([[10, 13],
       [22, 29]])

 

In [8]:

np.dot(a,b)

 

Out[8]:

array([[10, 13],
       [22, 29]])

 

In [9]:

np.einsum('ij,jk->ik',a,b)  # np.einsum('ij,jk',a,b)와 동일

 

Out[9]:

array([[10, 13],
       [22, 29]])

 

np.einsum('ij,jk->ik',a,b)에서 ij는 행렬 a의 행과 열의 인덱스 표현이며, jk는 행렬 b의 행과 열의 인덱스 표현이다. 행렬의 곱을 shape로 나타내면 (i,j) x (j,k) = (i,k)이므로 'ij,jk->ik'은 행렬 곱을 표현한 것이다.
화살표는 출력의 인덱스 표현을 나타내는 것인데, 화살표와 출력 인덱스를 생략한 형태인 np.einsum('ij,jk',a,b)도 동일한 결과를 얻을 수 있다.

 

In [10]:

np.einsum('ij,jk',a,b)

 

Out[10]:

array([[10, 13],
       [22, 29]])

 

In [11]:

np.einsum('ij,jk->i',a,b)  # 행렬 곱의 결과에서 각 행의 합을 구한다.

 

Out[11]:

array([23, 51])

 

다시말하지만, Einstein notation은 특정 인덱스 집합에 대한 합연산을 간결하게 표시한 것이라고 하였다. 출력 인덱스를 하나만 표시하면 해당 인덱스 요소의 합을 구하라는 의미이다. 출력 인덱스는 'ik'이므로, 'ij,jk->i'는 행의 원소를 더한 결과를 반환한다. 열의 원소를 더한 결과는 'ij,jk->k'를 사용한다.

 

아다마르 행렬 곱 (Hadamard product, Element-wise product)

shape이 같은 두 행렬의 원소 대 원소 곱을 연산한다.

 

In [12]:

np.einsum('ij,ij->ij',a,b)

 

Out[12]:

array([[ 2,  6],
       [12, 20]])

 

In [13]:

a*b  # 원소 대 원소 곱셈

 

Out[13]:

array([[ 2,  6],
       [12, 20]])

 

 

대각행렬 (diagonal matrix)

행렬의 대각성분을 가져온다. np.diag(a)와 동일하다.

 

In [14]:

c = np.arange(9).reshape(3,3)
np.diag(c)

 

Out[14]:

array([0, 4, 8])

 

In [15]:

np.einsum('ii->i', c)  # 행렬의 대각 행렬을 구한다.

 

Out[15]:

array([0, 4, 8])

 

np.einsum('ii->i', a)에서 ii->i에서 ii는 행렬 a의 행과 열이 같은 성분인 (i,i)에 해당하는 원소를 뜻하는 것이며, 화살표(->) 뒤의 i는 행 또는 열의 성분을 뜻하는데, 대각행렬이므로 행과 열이 같으므로 i가 된다.
즉, "행렬 a의 행 인덱스와 열 인덱스가 같은 원소들을 구하여라"와 같은 의미를 표현한 것이다.

 

전치행렬 (Transpose)

행렬의 행과 열을 바꾼다. ixj 행렬을 jxi 행렬로 변환한다.

 

In [16]:

a.T

 

Out[16]:

array([[1, 3],
       [2, 4]])

 

In [17]:

np.einsum('ij->ji',a)

 

Out[17]:

array([[1, 3],
       [2, 4]])

 

내적

두 벡터 내적은 원소 대 원소의 곱의 합이고, 행렬의 경우 행렬 곱에 해당한다.

 

In [18]:

v1 = np.array([1,2,3])
v2 = np.array([4,5,6])

In [19]:

np.einsum('i,i->',v1,v2)

 

Out[19]:

32

 

In [20]:

np.dot(v1,v2)

 

Out[20]:

32

 

 

Reference

• https://baekyeongmin.github.io/dev/einsum/
• https://ita9naiwa.github.io/numeric%20calculation/2018/11/10/Einsum.html
• https://numpy.org/doc/stable/reference/generated/numpy.einsum.html