벡터 필드 기반 3D 분자 생성

초록

VecMol은 분자를 3차원 공간에 정의된 연속 벡터 필드로 표현하고, 이를 신경 필드와 잠재 확산 모델을 이용해 생성한다. 기존의 그래프·좌표 기반 방식이 겪는 원자 타입·좌표 혼합 학습 문제와 원자 수 제한을 탈피하여, 해상도에 구애받지 않는 구조 표현과 효율적인 샘플링을 가능하게 한다.

상세 분석

본 논문은 3D 분자 생성 문제를 “연속 벡터 필드”라는 새로운 표현으로 재정의한다. 기존 방법들은 원자 종류라는 이산적 속성과 원자 좌표라는 연속적 속성을 동시에 학습해야 하는데, 이는 모달리티 간 얽힘과 화학·기하학 일관성 제약을 초래한다. VecMol은 공간의 임의 점 q∈ℝ³에 대해 K개의 원자 종류별 3차원 벡터 v_k(q)를 출력한다. 각 벡터는 해당 점에서 가장 가까운 원자 종류 k의 중심을 향하는 방향과 크기를 포함하며, 소프트맥스와 가우시안 클리핑을 통해 근접 원자에 가중치를 집중시킨다. 이렇게 정의된 필드는 원자 존재와 위치 정보를 연속적으로 내포하므로, 해상도와 원자 수에 독립적인 표현이 가능하다.

벡터 필드의 파라미터화는 신경 필드(auto‑encoder)로 수행된다. 인코더는 원자 좌표와 타입을 입력으로 받아 3D 격자(L³) 위에 위치한 고정 차원의 잠재 코드 z∈ℝ^{L³×d}를 생성한다. 여기서는 원자‑그리드 간 교차 그래프 메시지 패싱을 도입해 가변 크기 그래프를 고정 격자 형태로 매핑한다. 디코더는 E(n)‑equivariant Graph Neural Network(EGNN)를 사용해 임의의 query point와 격자 잠재 코드를 결합, 각 원자 종류별 가상 소스 위치 s_k(i)를 예측하고 v_k(q_i)=s_k(i)−q_i 형태로 벡터를 재구성한다. EGNN 구조는 회전 equivariance와 평행 이동 invariance를 보장해 물리적 일관성을 유지한다.

학습 단계에서는 실제 벡터 필드와 디코더가 출력한 필드 사이의 MSE 손실을 최소화한다. 이후 잠재 공간에 대해 DDPM(Latent Diffusion Probabilistic Model)을 학습해, 가우시안 노이즈에서 시작해 점진적으로 denoise 함으로써 새로운 잠재 코드 z₀를 샘플링한다. 샘플된 z₀는 동일한 디코더를 통해 벡터 필드 V를 생성하고, ODE 기반 최적화(gradient ascent)와 원자 병합 절차를 통해 최종 원자 좌표와 타입을 복원한다. 이 과정은 사전에 원자 수를 지정할 필요가 없으며, 원하는 해상도로 임의의 위치에서 원자 위치를 추출할 수 있다.

실험에서는 QM9와 GEOM‑Drugs 두 벤치마크에서 기존 E(3)‑equivariant diffusion 모델(EDM 등)과 비교해 구조적 정확도와 화학적 유효성 측면에서 경쟁력 있는 결과를 보였다. 특히, 벡터 필드가 제공하는 방향 정보는 원자 위치 복원 시 수렴 속도를 높이고, 복잡한 분자에 대한 스케일링을 개선한다. 한계점으로는 현재 K(원자 종류)마다 별도 벡터를 학습해야 하는 점과, ODE 기반 복원 단계가 추가적인 계산 비용을 유발한다는 점이 있다. 향후 연구에서는 벡터 필드와 스칼라 밀도 필드를 결합하거나, 더 효율적인 원자 추출 알고리즘을 도입해 속도·정확도 트레이드오프를 최적화할 여지가 있다.