Einsum Notation
Note
Pytorch나 Tensorflow 내의 많은 글들이 외우기 너무너무너무 진짜 외우기도 어렵고, 쓰기도 어려워서, 쉽게 표현할 방법이 없나 찾아보다 정리한 글입니다. 기본적으로, Einsum is All You Need 이 글을 많이 참조했습니다.
Introduction
PyTorch, Tensorflow 내의 다양한 함수(Dot Products, Outer Products, Transposes ,matrix-vector, 아니면 matrix-matrix multiplication)들의 name과 signature을 외우기 어렵지 않은가? 이 글을 읽는 사람이 나와 비슷하다면, 분명 이를 어렵게 느낄 것이다.
Einsum 표기법은 특수한 Domain Specific Language를 이용해 이 모든 행렬,(사실은 텐서) 연산을 표기하는 방법이다.
처음 보면
# Numpy
np.einsum('ix,jx->ij', A, B)
# PyTorch
torch.einsum('ix,jx->ij', [A, B])
이런 괴상한 표기법에 당황하겠지만, 조금만 익숙해지면 각종 함수를 이용하는 것보다 훨씬 훨씬 편하다.
이러한 장점이 있어요!!
- 다양한 연산에 대한 통일된 표기법
- 외울 것도 없고, 까먹을 것도 없다. 상황에 따라 여러 Framework를 사용해야 하는 경우가 많은데, 대부분의 Framework에서 Einsum을 지원하기 때문에, 외울게 적다.
- 더 간결하며, 쉬운 표기법
- 거의 모든 연산에 대해 중간 연산 없이 계산을 할 수 있다.
- 많은 프레임워크가 Einsum에 대한 최적화가 잘 되어있다.
- PyTorch는 아직 좀 불편한 것 같다
- Domain Specific Language의 장점인 것 같다.
현재, 많은 라이브러리에서 Einsum을 지원한다.
일반적으로 Python을 써서 프로그래밍하기 때문에 다른 언어는 잘 모르겠지만, 수치 연산 라이브러리라면 대부분 포함하고 있는 것 같다.
Operations
Basics
기본적으로 많은 Einsum Framework는 다음과 같은 함수 형태를 갖는다
Result = einsum("dimension notation of A, dimension notation of B,...->Result Dimension", A, B, ...)
임의의 갯수의 Tenosr(Scalar, Vector, Matrix, …)를 입력으로 받아, 임의의 차원을 갖는 Tensor를 결과로 돌려주는 방법을 말한다. 설명을 길게 하는 것보다, 계산의 결과를 보는 것이 훨씬 쉽게 Einsum Notation에 익숙해질 것 같다. 바로 들어가보자 ‘ㅅ’.. 바로 복사-붙여넣기를 하면 실행이 되도록 만들었으니, python console을 켜고 복사붙여넣기를 해보면서 확인해보자!
Unary Operation
Transpose
\(R = A^T \\ R_{ij} = A{ji}\)
import numpy as np
A = np.array([[1,2,3], [4,5,6]])
R = np.einsum("ij->ji", A)
print(R)
diagonal, Trace
\(r_i = A_{ii} \\ r = \sum_i A_{ii}\)
import numpy as np
A = np.eye(10)
diag = np.einsum('ii->i', A)
trace =np.einsum('ii->', A)
print(diag)
print(trace)
2.Summation
matrix sum to scalar
\(R = \sum_i \sum_jA_{ij}\)
import numpy as np
A = np.array([[1,2,3], [4,5,6]])
R = np.einsum("ij->", A)
print(R)
matrix column or row sum (to vector)
\(r_i = \sum_j A_{ij} r_j = \sum_i A_{ij}\)
import numpy as np
A = np.array([[1,2,3], [4,5,6]])
row_sum = np.einsum("ij->i", A)
col_sum =np.einsum("ij->j", A)
print(row_sum)
print(col_sum)
3. Multiplication
Dot Product, Outer product of two vectors
\(r = x^Ty, r_i = x_i y_i \\ R_{ij} = x_iy_j\)
import numpy as np
x = np.array([-1, -10, -100])
y = np.array([1, 10, 100])
dot = np.einsum('i,i->', x, y )
outer = np.einsum('i,j->ij', x,y)
print(dot)
print(outer)
Hadamard(element-wise) product of vector or matrix
\(R_{ij} = A_{ij}B_{ij} \\ r_{i} = a_i b_i\)
import numpy as np
x = np.array([-1, -10, -100])
y = np.array([1, 10, 100])
elemwise_vec = np.einsum('i,i->i', x, y)
print(elemwise_vec)
A = np.arange(6).reshape((2, 3))
B = np.arange(6).reshape((2, 3))
elemwise_mat = np.einsum('ij,ij->', A, B)
print(elemwise_mat)
Matrix-Vector multiplication
\[b = Ax, b_i = \sum_j A_{ij} b_j \\ b_i = \sum_j A_{ij}b_j\]import numpy as np
A = np.array([[1,2,3], [4,5,6]])
x = np.array([-1, -10, -100])
b = np.einsum('ij,j->i', A, x)
print(b)
Matrix-Matrix Multiplication and Batched Matrix multiplication
\(R_{ij} = \sum_k A_{ik}B_{ki} \\ B_{bjj} = \sum_k B_{bik} B_{bkj}\)
import numpy as np
## Matrix-Matrix Multiplication
A = np.array([[1,2,3], [4,5,6]])
B = A.transpose()
R = np.einsum('ik,kj->ij', A, b)
## Batched Matrix Multiplication
A = np.random.random(size=(3,10,4))
B = np.random.random(size=(3,4, 8))
R = np.einsum('bik,bkj->bij',A, B)
마지막으로, 2개를 초과하는 변수에 대해 einsum을 할 수 있음을 보이기 위해, Quadratic Form에 대해 계산해보자.
Quadritc Form, or Matrix norm, or Distance with respect to Matrix(Mahalanobis distance)
\(r = x^TAx \\ r = a^TAb\)
import numpy as np
x = np.array([1,2,3])
y = np.array([-1,-2,-3])
A = np.random.random(size=(3, 3))
r = np.einsum('i,ij,j->', x, A, y)
마땅히 예시로 들만한게 없어서 생략하지만, 더 고차원의 Tensor Contraction(텐서 연산)도 손쉽게 할 수 있다. 임의의 곱-합으로 이어지는 연산은 쉽게 할 수 있다.(연습해보자!)
Examples 1: Linear Regression with Einsum
Note
This assume that you have basic knowledge in PyTorch
Imports
import math
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.autograd as autograd
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
Examples 1: Linear Regression with Einsum
class LinearModule(nn.Module):
def __init__(self, idim, odim):
super(LinearModule, self).__init__()
self.W = nn.Parameter(torch.Tensor(idim, odim))
self.W.data.uniform_(-(1./ math.sqrt(idim+odim)), 1. / math.sqrt(idim+odim))
self.b = nn.Parameter(torch.Tensor(np.random.normal(scale=0.001, size=(odim, ))))
def forward(self, x):
"""
Args:
x: [batch_size x dim]
"""
p = torch.einsum("bt,to->bo", [x, self.W])
p += self.b
q = p.squeeze(1)
return p, q
def loss(self, x, y):
_, pred_y = self.forward(x)
diff = (pred_y - y) / y.size(0)
return torch.einsum("x,x->", [diff, diff])
Examples 2: Pointer Attention mechanism With Einsum
class ptr_att(nn.Module):
def __init__(self, hidden_size, name='PointerAttention', use_cuda=False):
super(ptr_att, self).__init__()
self.W_enc = nn.Parameter(torch.FloatTensor(hidden_size, hidden_size))
self.W_ref = nn.Parameter(torch.FloatTensor(hidden_size, hidden_size))
self.V = nn.Parameter(torch.FloatTensor(hidden_size))
self.V.data.uniform_(-(1. / math.sqrt(hidden_size)) , 1. / math.sqrt(hidden_size))
self.W_enc.data.uniform_(-(1. / math.sqrt(hidden_size)) , 1. / math.sqrt(hidden_size))
self.W_ref.data.uniform_(-(1. / math.sqrt(hidden_size)) , 1. / math.sqrt(hidden_size))
def forward(self, enc, ref):
batch_size = enc.size(0)
seq_len = enc.size(1)
"""
Args:
enc: [batch_size x seq_len x hidden_size] (actually, seq_len is different by each)
ref: [batch_size x hidden_size]
"""
Wenc = torch.einsum("ak,bjk->bja", [self.W_enc, enc])
Wref = torch.einsum("ak,bk->ba", [self.W_ref, ref]).unsqueeze(1).repeat(1,seq_len,1)
# [batch_size x seq_len x hidden_size] reference vector multiplied by w_enc
W = torch.einsum("k,ijk->ij", [self.V, F.tanh(Wenc + Wref)])
# [batch_size x seq_len],
#return W
return W
Example 3: Attention Mechanism brought from here
# Parameters
# -- [hidden_dimension]
bM, br, w = random_tensors([7], num=3, requires_grad=True)
# -- [hidden_dimension x hidden_dimension]
WY, Wh, Wr, Wt = random_tensors([7, 7], num=4, requires_grad=True)
# Single application of attention mechanism
def attention(Y, ht, rt1):
# -- [batch_size x hidden_dimension]
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_size x sequence_length]
at = F.softmax(torch.einsum("ijk,k->ij", [Mt, w]))
# -- [batch_size x hidden_dimension]
rt = torch.einsum("ijk,ij->ik", [Y, at]) + F.tanh(torch.einsum("ij,jk->ik", [rt1, Wt]) + br)
# -- [batch_size x hidden_dimension], [batch_size x sequence_dimension]
return rt, at
# Sampled dummy inputs
# -- [batch_size x sequence_length x hidden_dimension]
Y = random_tensors([3, 5, 7])
# -- [batch_size x hidden_dimension]
ht, rt1 = random_tensors([3, 7], num=2)
rt, at = attention(Y, ht, rt1)
at # -- print attention weights
Summary
여러 라이브러리 Tensorflow, PyTorch, Numpy같은 여러 라이브러리를 같이 써야 하는 경우가 많고, Notation이 미묘하게 달라서 외우기 어려운 경우가 많았는데, (읽을 땐 문제가 아니지만 쓸 때) 이 표기법을 알고, 그런 고민에서 많이 자유로워졌다.
여러분도 그랬으면 좋겠다. 이미 fluent하게 구사하고 있다면….. 에… 부럽당…
Quick reference for Einsum from here
Vector inner product: "a,a->" (Assumes two vectors of same length)
Vector element-wise product: "a,a->a" (Assumes two vectors of same length)
Vector outer product: "a,b->ab" (Vectors not necessarily same length.)
Matrix transposition: "ab->ba"
Matrix diagonal: "ii->i"
Matrix trace: "ii->"
1-D Sum: "a->"
2-D Sum: "ab->"
3-D Sum: "abc->"
Matrix inner product "ab,ab->" (If you pass twice the same argument, it becomes a matrix L2 norm)
Left-multiplication Matrix-Vector: "ab,b->a"
Right-multiplication Vector-Matrix: "a,ab->b"
Matrix Multiply: "ab,bc->ac"
Batch Matrix Multiply: "Yab,Ybc->Yac"
Quadratic form / Mahalanobis Distance: "a,ab,b->"