Search

Basic Algebra for E(3)-equivariance

1. Group (군)

어떤 set (집합)에 대해, 어떤 binary operation (이항연산)을 수행하여도 항상 그 set 안에 속하는 경우, 이 set과 binary operation의 묶음을 group이라고 한다.
Group을 이루기 위한 조건은 아래와 같다.
1.
Closure
2.
Associativity (결합법칙)
(ab)c=a(bc)(a * b) * c = a * (b * c)
3.
Identity (항등원) 존재
ae=a=eaa * e = a = e * a
4.
Inverse (역원)의 존재
ax=e=xaa * x = e = x * a
Example of groups
Euclidean Groups E(n)E(n)
nn-차원 공간 상의 모든 rigid motion (translation, rotation, reflection) 의 집합은 group을 이룬다.
Special Euclidean Group SE(n)SE(n)
nn-차원 공간 상의 translation과 rotation
Orthogonal Group O(n)O(n)
nn-차원 공간 상의 rotation과 reflection
Special Orthogonal Group SO(n)SO(n)
nn-차원 공간 상의 rotation

2. Linear map

Rigid motion은 linear mapping 이고, linear map은 matrix 형태로 표현할 수 있다.
예를 들어, 3차원 상에서 rotation matrix는 아래와 같이 표현한다.
x축을 중심으로 α\alpha 만큼, y축을 중심으로 β\beta 만큼, z축을 중심으로 γ\gamma 만큼 회전

3. Invariance & Equivariance

어떤 operation 이자 group GG에 대해,
GG-invariant
f(gx)=f(x),xX,gGf(g \cdot x) = f(x), \quad \forall x \in X, \forall g \in G
GG-equivariant
f(gx)=gf(x),xX,gGf(g \cdot x) = g \cdot f(x), \quad \forall x \in X, \forall g \in G
즉, invariance 는 equivariance의 special case (gy=y,yYg \cdot y = y, \quad \forall y \in Y를 만족하는) 이다.
우리가 원하는 E(3)E(3)-equivariant 한 함수 ff는 어떻게 만들까?
그 전에, vector space에 대해 잠시 복습해보자.

4. Vector Space

Linear algebra에서 배웠듯이, 모든 vector space는 basis를 가진다.
예를 들어, V=R3V=\mathbb{R}^3의 경우, {(1,0,0),(0,1,0),(0,0,1)}\{ (1, 0, 0), (0, 1, 0), (0,0,1)\}은 basis 이다. 또, {(cosθ,sinθ,0),(sinθ,cosθ,0),(0,0,1)}\{ (\cos\theta, -\sin\theta, 0), (\sin \theta, \cos \theta, 0), (0, 0, 1)\}도 basis 이다. 즉, vector space R3\mathbb{R}^3 상의 임의의 점을 이 basis 의 조합으로 항상 표현할 수 있다.
그런데 만약 vector space VV가 함수들의 집합인 경우에는 어떨까? 그 집합에 속하는 임의의 함수를 basis로 구분하여 나타낼 수 있는 방법이 있을까?

5. Fourier Series

Periodic functions can be expressed as a sum of trigonometric functions.
임의의 함수를 삼각함수 또는 지수함수의 일차결합으로 나타내는 것을 의미한다.
더 formal 하게 정의해보면,
xS1x \in S^1에 대해, fL2(S1)f \in L^2(S^1)인 임의의 ff에 대해
f(x)=k=0akcoskx+k=0bksinkxf(x) = \sum^\infin_{k=0} a_k \cos kx + \sum^\infin_{k=0}b_k \sin kx
→ 즉, L2(S1)L^2(S^1)의 basis는 {1,sinx,cosx,sin2x,cos2x,...}\{1, \sin x, \cos x, \sin 2x, \cos 2x, ...\} 이다.
참고: S1S^1은 one-dimensional sphere, 즉 circle 이다.

6. Laplace Series

Fourier series를 L2(S2)L^2(S^2)로 연장하여 Laplace series를 정의해보자.
xS2x\in S^2에 대해, fL2(S2)f \in L^2(S^2)인 임의의 ff에 대해
f(x)=l=0m=llam,lYml(x)f(x) = \sum^\infin_{l=0}\sum^l_{m=-l}a_{m,l}Y^l_m(x)
→ 즉, L2(S2)L^2(S^2)의 basis는 {Y00,Y11,Y01,Y11,...}\{Y^0_0, Y^1_{-1}, Y^1_0, Y^1_{1}, ...\} 이다.
이때, Yml:S2RY^l_m: S^2 \to \mathbb{R}.
참고: S2S^2는 unit sphere 이다.
이제 SO(3)SE(3)E(3)SO(3) \to SE(3) \to E(3) 순으로 equivariance 를 확장시켜 나가보자.

7. SO(3)SO(3)-equivariance

Vector space L2(S2)L^2(S^2)에 작용하는 SO(3)SO(3)에 대해 π\pi를 아래와 같이 정의하면, equivariant한 π\pi를 얻을 수 있다.
π:SO(3)GL(L2(S2))s.t.(π(R)f)(x)=f(R1x)\pi: SO(3) \to GL(L^2(S^2)) \quad s.t. \quad (\pi(R)f)(\vec{x}) = f(R^{-1}\vec{x})
이때, RSO(3),xS2R \in SO(3), \vec{x} \in S^2 이다.
π(R)=(π0(R)00...0π1(R)0...00π2(R)...............) with (Y0Y1Y2...)\pi(R) = \begin{pmatrix} \pi_0(R) & \mathbf{0} & \mathbf{0} & ...\\ \mathbf{0} & \pi_1(R) & \mathbf{0} & ... \\ \mathbf{0} & \mathbf{0} & \pi_2(R) & ... \\ ... & ... & ... & ... \end{pmatrix} \text{ with } \begin{pmatrix}Y^0 \\ Y^1 \\ Y^2 \\ ...\end{pmatrix}
Yl=[Yll,...,Yll]:S2R2l+1Y^l = [Y^l_{-l}, ..., Y^l_{l}]: S^2 \to \mathbb{R}^{2l+1}
πl(R)\pi_l(R)(2l+1)×(2l+1)(2l+1) \times (2l+1) 차원의 matrix이고, 아래의 중요한 성질을 가진다.
Yl(Rx)=πl(R)Yl(x)Y^l(R\vec{x}) = \pi_l(R)Y^l(\vec{x})
이때, π0(R)=1,π1(R)=R\pi_0(R) = 1, \pi_1(R)=R 이다.
위에서 equivariance에 대한 정의를 참고하면, 어떤 SO(3)SO(3) 상의 임의의 operation RR, 즉, RSO(3)R \in SO(3)에 대해서,
Yl(Rx)=Yl(Rx)=πl(R)Yl(x)=RYl(x)Y^l(R \cdot \vec{x}) = Y^l (R \vec{x}) = \pi_l(R)Y^l(\vec{x}) = R \cdot Y^l(\vec{x})
를 만족하면, SO(3)SO(3) equivariance 를 만족한다고 할 수 있다.
이때 YlY^lll-irreps (irreducible representation)라고 부르고, lele, 즉 0e,1e,2e,...0e, 1e, 2e, ... 등으로 표현한다.
예를 들어보자. R=(1000cosθsinθ0sinθcosθ)R=\begin{pmatrix} 1 & 0 & 0\\ 0 & \cos\theta & -\sin\theta \\ 0 & \sin\theta & \cos\theta \end{pmatrix}에 대해,
Y0(x)=1Y^0(\vec{x})=1, Y1(x)=[x,y,z]Y^1(\vec{x}) = [x, y, z] 이라고 정의하면,
RY0(x)=π0(R)1=11=1=Y0(Rx)R \cdot Y^0 (\vec{x}) = \pi_0(R) \cdot 1 = 1 \cdot 1 = 1 = Y^0(R\vec{x})
RY1(x)=π1(R)Y1(xyz)=R(xyz)=(xycosθzsinθysinθ+zcosθ)=Y1(Rx)R \cdot Y^1(\vec{x}) = \pi_1(R) Y^1 \begin{pmatrix}x \\ y \\ z\end{pmatrix} = R \begin{pmatrix}x \\ y \\ z\end{pmatrix} = \begin{pmatrix}x \\ y \cos\theta - z\sin\theta \\ y\sin\theta + z\cos\theta\end{pmatrix} = Y^1(R\vec{x})
이므로, equivariant function Y0Y^0Y1Y^1을 얻었다.

7.1. More equivariant functions

그렇다면 Y2,Y3Y^2, Y^3 등 더 많은 equivariant function은 어떻게 얻을 수 있을까?
아래 세 operation을 통해 얻을 수 있다.
Composition
Linear sum
Rule of tensor product between Yl1Y^{l_1} and Yl2Y^{l_2} (thanks to Universal property)
(gf)(gh)=g(fh)(g \cdot f) \otimes (g \cdot h) = g \cdot (f \otimes h)
두 함수를 조합하여 어떤 함수 p:S2×S2R2l3+1p: S^2 \times S^2 \to \mathbb{R}^{2l_3 + 1}를 얻을 수 있다. 이때, l1l2l3l1+l2|l_1 - l_2| \leq l_3 \leq l_1 + l_2 이다.
예시

7.2. Learnable parameter

이제 여러 SO(3)SO(3)-equivariant function을 얻을 수 있다는 것은 확인했다.
그럼 neural network의 핵심인 learnable parameter는 어떻게 추가할 수 있을까?
각 irreps 에 scalar weight를 곱해주면 된다.
예를 들어, x=(x1,x2),y=(y1,y2)x=(x_1, x_2), y=(y_1, y_2)로 정의하고 이때 x1,x2,y2x_1,x_2, y_21e1e이고 y1y_10e0e라고 해보자.
이를 x=x=”2x1e”, y=y=”1x0e + 1x1e” 라고 표기한다.
그리고 output을 lw=(w1x1y2+w2x2y2w3x1y1+w4x2y1)l_w = \begin{pmatrix} w_1x_1 \cdot y_2 + w_2x_2 \cdot y_2 \\ w_3 x_1y_1 + w_4 x_2 y_1 \end{pmatrix} 이라고 정의하면, 총 4개의 learnable parameter가 있다.
x1x_1y2y_2, 그리고 x2x_2y2y_2를 조합하면 0e, 1e, 2e 가 가능한데, \cdot 으로 표시하였으므로 그 중 0e에 해당하는 것만 사용한다는 의미이다. 따라서 lwl_w의 첫 번째 원소는 0e 이다.
x1x_1y1y_1, x2x_2y1y_1을 조합하면 1e만 가능하다.
따라서 lwl_w는 “1x0e + 1x1e” 로 표기할 수 있다.

7.3. Non-linearity

그럼 non-linearity는 어떻게 부여할까?
대부분의 non-linear function은 그 결과값이 equivariance를 유지하지 못 하는 경우가 많으므로, non-linear function에 들어가는 input으로는 invariant한 function만 넣어주어야 한다.
non-linear_function(invariant polynomials)\text{non-linear\_function}( \text{invariant polynomials)}
따라서, 방금 전 예시에서 lwl_w의 0e 부분인 w1x1y2+w2x2y2w_1 x_1 \cdot y_2 + w_2x_2 \cdot y_2에 sigmoid 등을 붙인 σ(w1x1y2+w2x2y2)\sigma(w_1 x_1 \cdot y_2 + w_2x_2 \cdot y_2) 로 non-linearity 를 부여할 수 있다. 1e, 2e 등 다른 것들은 equivariant 하지만 invariant 하지는 않아서 non-linearity 를 부여하면 equivariance 성질을 잃게 된다.
최종적으로 우리가 원하는 feature는 invariant 한 형태가 대부분일텐데, 처음부터 invariant한 연산만 수행하면 가능한 연산이 너무 제한적이기 때문에 equivariant 연산들로 neural net을 구성하고, 마지막에 invariant feature를 얻도록 하는 방향이라고 이해하면 된다.

SE(3)SE(3)-equivariance

앞서 본 SO(3)SO(3)-equivariance 를 만족하는 연산과 적절한 parametrization과 non-linearity 가 만족되는 네트워크를 여러 층 쌓으면 SO(3)SO(3)-equivariant neural network를 만들 수 있다.
SE(3)SE(3)SO(3)SO(3)에 translation 연산이 추가된 개념인데, 이는 각 SO(3)SO(3)-equivariant layer에 input을 normalized relative vector (아래 수식 참고) 형태로 넣어주게 되면 만족할 수 있다.
hi=1NijNilθ(hj,xjxixjxi)h_i = \frac{1}{N_i} \sum_{j \in N_i} l_\theta \bigg(h_j, \frac{x_j - x_i}{\|x_j - x_i\|}\bigg)

E(3)E(3)-equivariance

E(3)E(3)-equivariance 는 SE(3)SE(3)에 reflection 연산이 추가된 개념이다. 이는 parity={1,+1}\{-1, +1\}라는 개념을 도입하여 확장할 수 있다.
Parity는 쉽게 말해 어떤 object가 홀수인지, 짝수인지를 나타내는 개념이라고 봐도 될 것 같다.