Equivariant Evidential Deep Learning for Interatomic Potentials

Motivation

• Known: Evidential deep learning (EDL)은 앙상블의 비용 없이 단일 모델로 불확실성을 정량화할 수 있으나, 기존 방법들은 스칼라 대상에 초점을 맞추었거나 대각 공분산으로 제한되어 회전 변환 시 통계적 자기 일관성을 잃는다.
• Gap: 벡터 값 원자 힘(atomic forces)에 대한 불확실성 모델링에서 회전 등변성(SO(3) equivariance)을 만족하면서 full covariance를 구성하는 방법이 부재했고, 대각 공분산 방식은 축 간 상관관계를 포착하지 못한다.
• Why: 원자 힘은 등변 벡터이므로 불확실성 텐서도 Σ' = RΣR⊤를 만족해야 하며, 이는 물리적 가정이 아닌 분포 자기 일관성의 요구사항이다. 신뢰할 수 있는 분자 동역학 시뮬레이션을 위해 계산 효율적인 UQ가 필수적이다.
• Approach: Normal-Inverse-Wishart (NIW) evidential 프레임워크 위에서, Riemannian exponential map과 tangent-space parameterization을 통해 S3++ 매니폴드에서 SO(3) 등변성과 positive definiteness를 동시에 보장하는 SPD 공분산을 구성하고, equivariance-preserving spectral stabilizer로 수치적 안정성을 확보한다.

Achievement

Figure 2. Uncertainty diagnostics on liquid water. Top-left:

• SO(3)-등변 full 3×3 SPD 공분산 구성: tangent-space parameterization과 Riemannian exponential map을 활용하여 positive definiteness와 회전 등변성을 동시에 보장하는 첫 번째 방법 제시
• 단일 forward pass에서 불확실성 분해: Normal-Inverse-Wishart 기반 evidential 모델로 aleatoric 및 epistemic 불확실성을 명시적으로 분리
• 우월한 보정 성능: liquid water, MD22, rMD17, silica glass 벤치마크에서 deep ensembles(heteroskedastic 변형 포함) 대비 더 나은 calibration error (CE), energy score (ES), negative log likelihood (NLL) 달성
• 4.5배 추론 속도 향상: 단일 모델 기반으로 앙상블 대비 현저한 계산 효율성 확보

How

Figure 1. e2IP framework. An SE(3)-equivariant GNN predicts the force mean and NIW evidential parameters, and constructs

• SE(3)-등변 GNN 백본으로부터 angular momentum order별 SO(3) irrep 특징 추출
• SO(3) 표현론을 활용하여 symmetric matrix를 isotropic (0e) 및 traceless (2e) 성분으로 분해
• Lie algebra의 tangent space에서 symmetric matrix를 매개변수화한 후 Riemannian exponential map exp_I(·)를 통해 S3++로 매핑하여 Σ₀ 구성
• NIW 선행 분포 하에서 증거 매개변수 ν, κ 예측으로 Student-t 예측 분포 도출
• Equivariance-preserving spectral stabilizer: coefficient-space damping으로 고유값 크기를 제한하면서 등변성 보존

Originality

• vectorial regression에 대한 equivariant SPD covariance의 첫 번째 통합 구성 — 기존 eIP의 대각 공분산 제한을 극복
• Riemannian geometry (exponential map on S3++)를 deep learning의 불확실성 정량화에 적용한 혁신적 접근
• Equivariance-preserving spectral stabilization 기법으로 NIW 훈련의 수치적 퇴화 문제 해결
• SO(3) irrep 분해를 systematic하게 활용하여 backbone-agnostic framework 설계

Limitation & Further Study

• 현재 구현은 NequIP, MACE 등 특정 SE(3)-등변 백본에 의존하며, 다른 등변 아키텍처로의 일반화는 미검증
• NIW 모델의 가정(Gaussian likelihood)이 strong하며, 비정규 분포를 따르는 원자 환경에서의 성능 저하 가능성
• Spectral stabilizer의 hyperparameter (threshold τ, ceiling c)에 대한 민감도 분석 부족
• 후속: multi-atom covariance correlation 모델링, graph-level uncertainty 확장, 실험 검증(ab initio 계산과 직접 비교)

Evaluation

총평: SE(3)-등변 evidential learning 분야의 기술적 혁신으로, Riemannian geometry를 통한 엄밀한 이론적 기초와 광범위한 벤치마크 검증을 결합하여 분자 동역학의 신뢰할 수 있는 배포 문제를 해결한 고수준의 연구이다.