SE(3)-Equivariant Flow Matching with Gaussian Process Priors for Geometric Trajectory Prediction

Motivation

• Known: flow matching과 diffusion model 같은 생성모델이 N-body 궤적 예측의 확률성과 복잡한 분포를 모델링하는 데 효과적임이 알려져 있으며, SE(3) 대칭성을 반영하는 equivariant 신경망 구조들도 다수 연구되어 왔다.
• Gap: 기존 flow matching 기반 궤적 예측 모델들은 isotropic Gaussian이나 고정된 CoM Gaussian 같은 관측 독립적(observation-agnostic)이고 trivial한 prior 분포를 사용하여, 미래 궤적이 갖는 시간적 연속성·관성 등의 temporal correlation과 SE(3) 공간 대칭성을 동시에 반영하지 못하는 한계가 있다.
• Why: 보다 데이터 분포에 가까운 equivariant prior를 설계하면 transport cost를 줄이고 vector field를 더 직선화하여, 물리학·화학·생물학 등 다양한 분야에서 활용되는 N-body 궤적 예측의 정확도와 샘플링 효율성을 동시에 향상시킬 수 있어 실용적, 이론적으로 중요하다.
• Approach: 저자들은 관측 궤적으로부터 vector-valued Gaussian process를 이용해 SE(3)-equivariant한 데이터 종속적 prior 분포를 구성하고, 이를 equivariant 메시지 패싱 GNN과 temporal cross-attention 기반 vector field 모델과 결합한 conditional flow matching 프레임워크인 GP-EquiFlow를 제안한다.

Achievement

Figure 1 provides the overview of GP-EquiFlow.

1. GP-EquiFlow 프레임워크 제안: 관측 궤적에 조건화된 N-body 미래 궤적의 SE(3)-equivariant 분포를 모델링하는 통합 flow matching 프레임워크를 제시하고, equivariant message-passing GNN과 temporal cross-attention으로 ground-truth vector field를 학습한다.
2. VVGP 기반 equivariant prior 설계: 관측 궤적의 시간적 상관관계를 활용하면서 SE(3)-equivariance를 보장하는 data-dependent prior 분포를 새롭게 설계하여 prior와 데이터 분포 간 결합을 자연스럽게 강화했다.
3. 실험적 우수성 입증: N-body 물리 시뮬레이션 및 분자 동역학 벤치마크에서 GP-EquiFlow가 기존 SOTA 대비 더 적은 샘플링 스텝으로 더 정확한 예측 성능을 보임을 광범위한 실험으로 입증했다.

How

Figure 1 provides the overview of GP-EquiFlow.

• geometric trajectory graph G=(X,H,E)를 정의하고, 관측 구간 Xobs와 예측 구간 Xpred로 분리하여 conditional 생성모델 pθ(Xpred|Xobs)를 SE(3)-equivariant prior q0와 equivariant generative process로 분해
• 각 입자의 3D 위치 궤적에 대해 vector-valued Gaussian process를 적용해 시간적 상관구조를 반영하는 prior N(μ(Xobs), Σ)를 구성하고, 이 과정이 SE(3) 변환에 대해 equivariant함을 보장
• conditional flow matching(CFM) 이론(Lipman et al., 2023)을 따라 prior p0와 target 분포 p1을 잇는 conditional probability path를 정의하고, 이에 대응하는 ground-truth vector field uτ(xτ|x1)를 유도
• equivariant message-passing GNN과 temporal cross-attention을 결합한 신경망으로 시간에 따른 vector field uθτ(xτ)를 파라미터화하여 MSE 기반 CFM 손실로 학습
• 추론 시에는 VVGP로 구성한 prior에서 샘플링한 초기값으로부터 ODE를 적은 스텝으로 적분하여 미래 궤적을 생성

Originality

• 기존 flow matching 기반 궤적 예측 연구들이 간과했던 prior 분포 설계 문제를 정면으로 다루며, temporal correlation과 spatial symmetry를 동시에 satisfy하는 prior를 최초로 제안
• vector-valued Gaussian process를 SE(3)-equivariant 생성모델의 prior 구성에 결합한 시도로, 기존 isotropic/CoM Gaussian prior의 한계를 극복하는 새로운 관점 제시
• prior 자체의 equivariance와 generative process의 equivariance를 함께 증명하여 전체 프레임워크의 SE(3)-equivariance를 이론적으로 보장

Limitation & Further Study

• VVGP 기반 prior 구성 시 커널 함수나 하이퍼파라미터 선택에 따른 민감도 및 계산 비용에 대한 심층 분석이 본문 발췌에서 명확히 드러나지 않음
• 관측 궤적이 짧거나 노이즈가 큰 경우 VVGP의 시간적 상관관계 추정이 부정확해질 수 있는데 이에 대한 강건성 검증이 필요
• N이 매우 큰 대규모 시스템에서의 확장성(계산 복잡도) 검증이 추가로 요구됨
• 후속 연구로는 다양한 커널 설계, 이질적 시스템(heterogeneous particle types)으로의 확장, 실제 관측 노이즈가 있는 실세계 데이터에 대한 검증이 필요해 보임

Evaluation

총평: Prior 분포 설계라는 상대적으로 덜 주목받은 문제를 SE(3)-equivariance와 Gaussian process를 결합해 우아하게 해결한 견고한 연구로, 이론적 타당성과 실험적 성능 개선을 모두 보여준다.