E(n) Equivariant Graph Neural Networks
E(n) 등변 그래프 신경망
Victor Garcia Satorras, Emiel Hoogeboom, Max Welling (2021)
구면 조화 함수(spherical harmonics)나 고차 텐서 없이도 좌표 업데이트를 통해 E(n)(유클리드 이동, 회전, 반사) 등변성을 달성하는 단순하면서도 효과적인 EGNN을 제안하여, 분자 동역학 및 N-body 시뮬레이션에서 우수한 성능을 보였다.
배경
물리 시스템과 분자 구조는 공간의 이동, 회전, 반사에 대해 대칭적이다. 즉, 분자를 회전시켜도 그 성질은 변하지 않으며(불변), 힘과 같은 벡터량은 함께 회전해야 한다(등변). TFN(Tensor Field Networks), SE(3)-Transformers 등 기존 등변 GNN은 구면 조화 함수와 Clebsch-Gordan 계수를 사용하여 등변성을 보장했으나, 구현이 복잡하고 계산 비용이 높았다. 불변 모델(SchNet 등)은 거리 정보만 사용하므로 방향 정보를 충분히 활용하지 못했다.
핵심 아이디어
EGNN의 핵심 혁신은 노드 좌표를 명시적으로 업데이트하면서도 E(n) 등변성을 유지하는 간단한 메커니즘을 제안한 것이다. 기존 등변 GNN이 고차 텐서(구면 조화)를 사용한 것과 달리, EGNN은 두 노드 간의 상대 위치 벡터(x_i - x_j)에 스칼라 가중치를 곱하여 좌표를 업데이트한다. 상대 위치 벡터 자체가 이미 등변적이므로, 스칼라 가중치만 불변이면 전체 업데이트가 등변적이 된다. 이 스칼라 가중치는 두 노드의 특징과 거리에 의존하는 학습 가능한 함수로 계산된다. 이를 통해 구면 조화, Wigner-D 행렬 등의 복잡한 수학적 장치 없이도 등변성을 달성한다.
방법론
각 레이어에서 세 가지 업데이트를 수행한다. (1) 메시지 계산: m_{ij} = phi_e(h_i, h_j, ||x_i - x_j||^2, a_{ij}), 여기서 a_{ij}는 엣지 속성. (2) 좌표 업데이트: x_i = x_i + C * sum_j (x_i - x_j) * phi_x(m_{ij}), 여기서 phi_x는 스칼라 가중치를 출력하는 MLP. (3) 노드 특징 업데이트: h_i = h_i + phi_h(h_i, sum_j m_{ij}). phi_e, phi_x, phi_h는 모두 간단한 MLP이다. 좌표 업데이트가 상대 위치 벡터에 스칼라를 곱하는 형태이므로, 회전/이동에 대한 등변성이 자동으로 보장된다. 반사 등변성도 마찬가지이다.
주요 결과
N-body 시뮬레이션(하전 입자 궤적 예측)에서 EGNN은 SE(3)-Transformer, TFN, Radial Field 등 기존 등변 모델을 MSE 기준으로 크게 능가했다. QM9 분자 성질 예측에서 SchNet, DimeNet 등 불변 모델과 경쟁력 있는 성능을 보이면서 계산 비용은 크게 낮았다. 모델 그래프 오토인코더(graph autoencoder)에서 분자 3D 구조 생성에도 적용하여, 등변 생성 모델의 가능성을 시연했다. TFN/SE(3)-Transformer 대비 학습 속도가 수배 빨랐다.
임팩트
EGNN은 등변 GNN의 민주화에 기여했다. 구면 조화 없이도 등변성을 달성할 수 있다는 것을 보여줌으로써, 등변 모델의 구현 및 활용 장벽을 크게 낮추었다. 이후 GVP(Geometric Vector Perceptrons), SEGNN, PaiNN 등 다양한 간소화된 등변 아키텍처의 발전에 직접적 영감을 주었다. 분자 생성(EDM, GeoDiff), 단백질 구조 예측, 점군(point cloud) 처리 등 다양한 3D 과학 응용에서 기본 구성 요소로 채택되고 있으며, DiffDock 등 분자 도킹 모델의 기초가 되었다.