3D Molecule Generation from Rigid Motifs via $\mathrm{SE}(3)$ Flows

Motivation

• Known: 기존 3D 분자 생성은 E(3)/SE(3)-equivariant diffusion 모델을 이용해 원자 단위 point cloud로 좌표를 직접 생성하는 방식이 주류이며, 분자 그래프 생성 분야에서는 fragment 기반 방법이 화학적 모듈성을 활용해 널리 쓰인다. 단백질 구조 생성에서는 AlphaFold2, FrameDiff처럼 아미노산 잔기를 SE(3) rigid frame으로 취급하는 방식이 효과적임이 입증되었다.
• Gap: 단백질과 달리 소분자는 고정된 잔기 집합의 선형 사슬이 아니라 임의의 분기 구조와 다양한 위상, 방대한 motif 어휘를 가지므로, rigid frame 기반 표현을 일반 drug-like 분자에 적용하는 것은 아직 다뤄지지 않은 문제였다.
• Why: 원자 수준 표현을 화학적으로 의미 있는 rigid motif로 압축하면 생성 공간의 차원과 스텝 수를 크게 줄이면서도 구조적 유효성과 안정성을 유지할 수 있어, 신약 설계 및 재료 설계를 위한 3D 분자 생성의 효율성과 확장성을 동시에 개선할 잠재력이 크다.
• Approach: 분자를 데이터 기반 fragmentation·canonicalisation 전략으로 K개의 rigid motif {(T_i, m_i)}로 분해한 뒤, SE(3) 상의 연속적 프레임(회전·이동)과 이산적 motif 타입을 multimodal flow matching(Campbell et al., 2024의 프레임워크를 확장)으로 공동 생성한다.

Achievement

Figure 2. Comparison of molecular representations on QMUGS

1. MOTIFLOW 프레임워크 제안: 분자를 개별 원자가 아닌 SE(3) 상의 rigid-body motif 집합으로 파라미터화하는 3D 분자 생성 프레임워크를 제시했다.
2. 데이터 기반 fragmentation·canonicalisation: 임의의 분자 위상을 처리하고 motif 대칭성을 고려하는 프래그멘테이션 및 정규화 전략을 개발했다.
3. multimodal flow의 적응: 이산 motif identity와 연속 SE(3) 기하 구성을 자기회귀 단계나 학습된 원자 수준 디코딩 없이 자연스럽게 공동 생성하도록 flow matching을 de novo 3D 분자 구조 생성에 맞게 확장했다.
4. 성능 및 효율성 입증: GEOM-DRUGS 벤치마크에서 atom stability 기준 SOTA를 능가하고, QMUGS의 더 큰 분자에도 확장 가능함을 보였으며, 생성 스텝을 2~10배 줄이고 표현을 3.5배 압축했다. 또한 조건부 생성 작업에서도 우수한 성능을 보였다.

How

Figure 4. Resulting sets of rigid motifs for a GEOM-DRUGS

• 분자를 {(y_j, h_j)} 형태의 원자 point cloud로부터 K개의 rigid motif {(T_i, m_i)}로 재파라미터화하며, 각 프레임 T_i=(R_i,x_i)∈SE(3)는 motif의 canonical pose에 대한 회전과 중심 좌표를 인코딩한다.
• SO(3) 회전에는 FoldFlow-Base(Bose et al., 2023) 기반 Riemannian flow matching을 적용하여 geodesic 보간과 closed-form conditional vector field(Rodrigues' formula 활용)로 학습 안정성과 속도를 확보한다.
• R^3 이동 성분에는 선형 보간 기반의 constant vector field를 사용하는 표준 flow matching을 적용한다.
• 이산 motif 타입은 continuous-time Markov chain(CTMC) 기반 discrete flow matching(Campbell et al., 2024)으로 모델링하며, 신경망 denoiser p_θ(m1|mt)를 cross-entropy loss로 학습하고 추론 시 학습된 rate matrix Q_θ_t로 CTMC를 시뮬레이션한다.
• 최종적으로 연속(R^x_θ, v^x_θ)과 이산(Q^θ_t) 벡터장/전이율을 결합한 joint multimodal flow로 T와 m을 동시에 적분하여 분자를 생성한다.

Originality

• 단백질 구조 생성의 frame 기반 표현(FrameDiff, FoldFlow)을 일반 drug-like 소분자로 처음 확장하여, 분자를 SE(3) rigid motif 집합으로 재구성하는 새로운 표현 방식을 제안했다.
• 분자 그래프 생성 분야의 fragment 기반 아이디어와 3D equivariant flow matching을 결합하여, 이산 motif 타입과 연속 SE(3) 프레임을 동시에 다루는 multimodal flow를 설계했다.
• 임의의 분자 위상(분기, 다양한 결합 구조)을 처리할 수 있는 데이터 기반 fragmentation·canonicalisation 전략을 도입하여 motif 대칭성 문제를 해결했다.

Limitation & Further Study

• motif 어휘(vocabulary)의 크기와 구성 방식이 데이터셋에 의존적일 가능성이 있어, 매우 다양한 화학 공간(예: 특이한 고리 구조, 신규 스캐폴드)에 대한 일반화 성능이 명확히 검증되지 않았다.
• fragmentation 및 canonicalisation 전략이 특정 motif 대칭성 처리에 의존하므로, 대칭성이 복잡하거나 모호한 motif에서의 안정성에 대한 추가 분석이 필요하다.
• 논문 발췌본에서는 조건부 생성 성능과 대형 분자(QMUGS)에서의 확장성에 대한 구체적 실패 사례나 실패 모드 분석이 제시되지 않아, 실제 약물 설계 파이프라인에 적용 시의 한계(예: 특정 화학적 제약 조건 반영)에 대한 후속 연구가 필요하다.
• rigid motif 표현으로 인해 motif 내부의 국소적 유연성(local flexibility)이나 미세한 원자 수준 왜곡을 포착하지 못할 가능성이 있다.

Evaluation

총평: 단백질 구조 생성의 rigid frame 개념을 일반 소분자로 확장한 참신하고 실용적인 접근으로, 효율성과 성능을 동시에 개선한 의미 있는 기여이나 motif 어휘 설계의 일반화 가능성에 대한 추가 검증이 필요하다.