[논문 리뷰] Learning Equivariant Energy Based Models with Equivariant Stein Variational Gradient Descent

Priyank Jaini, Lars Holdijk외 1명|arXiv (Cornell University)|2021. 06. 15.

Adversarial Robustness in Machine Learning참고 문헌 36인용 수 5

한 줄 요약

이 논문은 대칭성 인variant 확률 밀도에서 다양하고 고품질의 샘플을 효율적으로 생성하기 위해 등변 핵함수를 사용하는 대칭성 인지 샘플링 방법인 등변 스틸 변분 기울기 하강법(E-SVGD)을 소개한다. 또한 E-SVGD를 통해 훈련된 등변 에너지 기반 모델(E-EBMs)을 제안하며, 기하학적 딥러닝과 군 등변 신경망을 활용하여 분자의 생성, 이미지 모델링, 입자 시스템에서 샘플 효율성, 일반화 능력 및 수렴성 향상을 입증한다.

ABSTRACT

We focus on the problem of efficient sampling and learning of probability densities by incorporating symmetries in probabilistic models. We first introduce Equivariant Stein Variational Gradient Descent algorithm -- an equivariant sampling method based on Stein's identity for sampling from densities with symmetries. Equivariant SVGD explicitly incorporates symmetry information in a density through equivariant kernels which makes the resultant sampler efficient both in terms of sample complexity and the quality of generated samples. Subsequently, we define equivariant energy based models to model invariant densities that are learned using contrastive divergence. By utilizing our equivariant SVGD for training equivariant EBMs, we propose new ways of improving and scaling up training of energy based models. We apply these equivariant energy models for modelling joint densities in regression and classification tasks for image datasets, many-body particle systems and molecular structure generation.

연구 동기 및 목표

• 대칭성 불변성을 명시적으로 통합하여, 불변 확률 밀도에서 샘플의 다양성과 효율성을 향상시키는 샘플링 방법을 개발한다.
• 표준 SVGD가 대칭 분포에서 샘플링하는 데에 한계가 있음을 해결하기 위해, 군의 구조를 유지하는 등변 핵함수를 도입한다.
• 에너지 함수와 샘플링 절차가 데이터의 기하학적 대칭성을 모두 존중하는 등변 에너지 기반 모델(E-EBMs)을 설계한다.
• 대칭성 인지 학습과 대칭성 인지 샘플링을 조합함으로써 생성 모델링 과제에서 더 빠른 수렴과 더 나은 일반화 성능을 달성함을 입증한다.
• 분자의 구조 생성, 회전 대칭성을 가진 이미지 데이터셋, 다체 입자 시스템을 포함한 다양한 영역에서 프레임워크를 검증한다.

제안 방법

• 군 작용에 대해 분포가 불변하도록 입자를 진화시키는 입자 기반 샘플링 방법인 등변 SVGD(E-SVGD)를 제안하며, 등변 핵함수를 사용한다.
• 스테인의 항등식을 활용해 재생 핵 힐버트 공간(RKHS) 내에서 커널 기반 기울기 유량을 유도하며, 핵함수를 군 작용에 대해 일관되게 변환하도록 구성한다.
• 에너지 함수를 등변 신경망(예: E-GNN, 기울기 조절 가능 CNN, 군 등변 네트워크 등)으로 매개변수화하여 데이터의 대칭성을 존중하는 등변 에너지 기반 모델(E-EBMs)을 정의한다.
• E-SVGD를 통해 음성 샘플을 생성함으로써 대trastive divergence를 사용해 E-EBMs를 훈련하며, 샘플링 과정이 에너지 함수와 동일한 대칭성을 존중하도록 보장한다.
• 군 표현(예: SO(3), C4, Sn)을 사용해 입자와 함수가 대칭 연산에 대해 어떻게 변환되는지 정의함으로써, 모델과 추론 과정에서 일관된 등변성을 확보한다.
• 오비트 기반 밀도 인수분해를 활용해 몫 공간에서의 분포를 모델링함으로써 복잡성을 감소시키고 샘플 효율성을 향상시킨다.

실험 결과

연구 질문

• RQ1 목표 분포의 대칭성을 명시적으로 유지하는 샘플링 알고리즘을 설계할 수 있을까? 이는 더 효율적이고 다양한 샘플링을 가능하게 할까?
• RQ2 에너지 함수와 샘플링 절차 양쪽에 등변성을 통합할 경우, 에너지 기반 모델의 훈련 안정성과 일반화 능력에 어떤 영향을 미칠까?
• RQ3 기존 SVGD에 비해 E-SVGD를 훈련 샘플러로 사용할 경우, 대칭적인 생성 모델링 과제에서 수렴 속도와 샘플 품질은 어느 정도 향상될까?
• RQ4 등변 아키텍처를 가진 E-EBMs는 분자의 구조 생성 및 대칭 제약 하에서의 이미지 생성과 같은 과제에서 훈련 데이터를 초월해 일반화할 수 있을까?
• RQ5 대칭성 인지 모델은 대칭 데이터를 포함한 분류 및 생성 과제에서 비등변 기반 모델에 비해 어떻게 성능을 발휘할까?

주요 결과

• E-SVGD는 특히 고차원 및 장거리 상호작용 상황에서 대칭 밀도에서 표본의 다양성과 품질에서 표준 SVGD를 능가한다.
• E-SVGD로 훈련된 등변 에너지 기반 모델는 표준 EBM보다 수렴 속도가 빠르고 일반화 능력이 뛰어나며, 특히 자료가 부족하거나 대칭성이 뚜렷한 환경에서 유의미한 성능 향상을 보인다.
• QM9 데이터셋에서 훈련된 E-EBMs로부터 생성된 분자들은 정확한 3D 결합 배열을 보이며, 아로마틱 고리와 구조 이성체를 포함한다. 반경 분포 함수는 진짜 데이터 분포와 매우 유사하다.
• C4 대칭성을 가진 FashionMNIST에서 E-SVGD로 훈련된 등변 EBM은 일반 SVGD나 비등변 아키텍처로 훈련된 모델에 비해 더 높은 테스트 정확도와 더 빠른 수렴 속도를 기록한다.
• E-SVGD를 훈련 샘플러로 사용함으로써, 다체 입자 시스템과 같은 복잡하고 대칭성이 뚜렷한 영역에서도 안정적이고 효율적인 대비 분산 훈련이 가능해졌다.
• 실증 결과에 따르면, E-SVGD로 훈련된 E-EBMs는 이미지 생성, 분자 설계, 입자 시스템 모델링을 포함한 모든 평가 벤치마크에서 더 현실적이며 구조적으로 일관된 샘플을 생성한다.