[논문 리뷰] Frame Averaging for Invariant and Equivariant Network Design

Omri Puny, Matan Atzmon외 5명|arXiv (Cornell University)|2021. 10. 07.

Machine Learning in Materials Science참고 문헌 62인용 수 10

한 줄 요약

이 논문은 계산적으로 비가능한 군 평균을 작은 데이터에 의존하는 군 원소의 부분집합(즉, 프레임)에 대한 평균으로 대체하는 체계적인 프레임 평균화(FA)를 소개한다. 이를 통해 신경망에서 정확한 불변성 또는 등변성을 달성할 수 있다. FA는 백본 아키텍처의 표현력을 유지하면서도 점군 법선 추정, 그래프 분리, n체 역학 예측 등에서 최신 기술(SOTA) 수준의 성능을 달성한다.

ABSTRACT

Many machine learning tasks involve learning functions that are known to be invariant or equivariant to certain symmetries of the input data. However, it is often challenging to design neural network architectures that respect these symmetries while being expressive and computationally efficient. For example, Euclidean motion invariant/equivariant graph or point cloud neural networks. We introduce Frame Averaging (FA), a general purpose and systematic framework for adapting known (backbone) architectures to become invariant or equivariant to new symmetry types. Our framework builds on the well known group averaging operator that guarantees invariance or equivariance but is intractable. In contrast, we observe that for many important classes of symmetries, this operator can be replaced with an averaging operator over a small subset of the group elements, called a frame. We show that averaging over a frame guarantees exact invariance or equivariance while often being much simpler to compute than averaging over the entire group. Furthermore, we prove that FA-based models have maximal expressive power in a broad setting and in general preserve the expressive power of their backbone architectures. Using frame averaging, we propose a new class of universal Graph Neural Networks (GNNs), universal Euclidean motion invariant point cloud networks, and Euclidean motion invariant Message Passing (MP) GNNs. We demonstrate the practical effectiveness of FA on several applications including point cloud normal estimation, beyond $2$-WL graph separation, and $n$-body dynamics prediction, achieving state-of-the-art results in all of these benchmarks.

연구 동기 및 목표

• 복잡한 대칭성(예: 순열 및 유클리드 운동)에 대해 정확히 불변 또는 등변인 신경망을 설계하면서도 표현력과 계산 효율성을 유지하는 데 도전하는 것.
• 크거나 연속적인 대칭군에서 전체 군 평균이 비가능한 문제를 해결하기 위해, 전체 군 원소의 평균을 작은, 효율적으로 계산 가능한 군 원소 부분집합으로 대체하는 것.
• GNN, 점군 네트워크, MLP와 같은 다양한 아키텍처에 적용 가능한 일반 목적의 체계적 프레임워크를 제공하는 것.
• FA 기반 모델이 백본 아키텍처의 최대 표현력을 유지함을 증명하여 핵심 설정에서 보편 근사가 가능하게 하는 것.
• 다양한 벤치마크에서 FA를 경험적으로 검증하여, 근사적 또는 전체 군 평균 기반의 기준 모델에 비해 뛰어난 성능과 불변성 안정성을 보여주는 것.

제안 방법

• 프레임 평균화는 대칭군 G의 모든 원소에 대한 계산적으로 비가능한 군 평균을, 입력 X에 따라 달라지는 유한한 부분집합인 프레임 F(X)에 대한 평균으로 대체한다.
• 프레임 F(X)는 군 작용에 대한 집합 등변성 성질을 만족하도록 구성되어 있어, F(X)에 대한 평균이 전체 군 작용 하에서 정확한 불변성 또는 등변성을 유지함을 보장한다.
• 이 프레임워크는 다양한 아키텍처에 적용된다: 순열 불변성에 대한 MLP, 그래프 수준의 불변성에 대한 GNN, 유클리드 운동 불변성/등변성에 대한 메시지 전달 GNN.
• 유클리드 운동 대칭성(E(d))의 경우, 프레임은 안정자 부분군에서 유도되며 기하학적 불변량을 활용하여 효율적인 계산을 가능하게 한다.
• 특징이 각 레이어에서 프레임을 기준으로 평균화되는 수정된 메시지 전달 메커니즘을 사용하여, 등변성을 유지하면서도 표현력을 유지한다.
• 이론적 분석을 통해 FA 기반 모델이 백본 모델이 보편 근사가 가능할 경우 보편 근사를 달성하며, 프레임 선택이 표현 능력을 감소시키지 않는다는 것을 증명한다.

실험 결과

연구 질문

• RQ1 전체 군 평균의 계산 비용을 지불하지 않고도 정확한 불변성 또는 등변성을 신경망에 강제로 적용할 수 있는 체계적 프레임워크를 개발할 수 있는가?
• RQ2 어떤 조건에서 작은, 입력에 의존하는 군 원소 부분집합(즉, 프레임)이 전체 군 평균을 대체하면서도 정확한 대칭성 성질을 유지할 수 있는가?
• RQ3 프레임 평균화는 대칭 인식 학습 과제에서 기반 백본 아키텍처의 표현력을 유지하거나 향상시키는가?
• RQ4 프레임 평균화는 근사 평균화 방법(예: 몬테카를로) 또는 전체 군 평균화에 비해 불변성 정확도와 모델 성능 측면에서 어떻게 비교되는가?
• RQ5 프레임 평균화는 순열, 유클리드 운동과 같은 다양한 대칭군에 일반화되고, GNN 및 점군 네트워크를 포함한 다양한 아키텍처에 적용 가능한가?

주요 결과

• 프레임 평균화는 작은 입력에 의존하는 프레임에 대한 평균을 취하여 전체 군 평균의 계산 비가능성 문제를 피하면서도 정확한 불변성과 등변성을 달성한다.
• FA 기반 모델은 백본 아키텍처의 최대 표현력을 유지하여 그래프 및 점군 학습 과제에서 보편 근사를 가능하게 한다.
• n체 역학 예측 과제에서 FA-GNN는 테스트 MSE 0.0057을 기록하여 파rameter 수가 유사한 SOTA인 EGNN(0.0071)보다 20% 이상 뛰어난 성능을 보였다.
• 2-WL 초월 그래프 분리 벤치마크에서 FA-MLP와 FA-GIN+ID는 완벽한 분리를 달성하여 보편적 표현력을 입증했다.
• 최소한 k=1개의 프레임 샘플을 사용한 근사적 FA는 k=1일 때 전체 군 평균화(GA)보다 훨씬 낮은 불변성 오차를 보였으며, 이는 더 뛰어난 안정성과 일반화 능력을 시사한다.
• 경험적 결과는 FA가 근사적 및 전체 군 평균화보다 불변성과 효율성이 뛰어나며, 특히 샘플 수가 적거나 고대칭 환경에서 두드러진 성능을 보임을 보여준다.