[논문 리뷰] Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3D

Bo Qiang, Yuxuan Song외 6명|arXiv (Cornell University)|2023. 05. 05.

Innovative Microfluidic and Catalytic Techniques Innovation인용 수 5

한 줄 요약

이 논문은 3D 분자의 구조적 타당성과 품질을 햖थ기 위해 군집에서 원자로의 전이 전략을 사용하는 계층적 확산 모델인 HierDiff를 제안한다. 먼저 등변성 확산 과정을 통해 조각 수준의 표현을 생성한 후, 메시지 전달과 반복적 샘플링을 통해 원자 수준의 정확한 3D 구조를 정밀 조정한다. 이는 기존 방법들보다 약물 유사성, 구조적 정확도 및 안정성에서 뛰어난 성능을 보이며, 특히 큰 분자에서 두각을 나타낸다.

ABSTRACT

Generating desirable molecular structures in 3D is a fundamental problem for drug discovery. Despite the considerable progress we have achieved, existing methods usually generate molecules in atom resolution and ignore intrinsic local structures such as rings, which leads to poor quality in generated structures, especially when generating large molecules. Fragment-based molecule generation is a promising strategy, however, it is nontrivial to be adapted for 3D non-autoregressive generations because of the combinational optimization problems. In this paper, we utilize a coarse-to-fine strategy to tackle this problem, in which a Hierarchical Diffusion-based model (i.e.~HierDiff) is proposed to preserve the validity of local segments without relying on autoregressive modeling. Specifically, HierDiff first generates coarse-grained molecule geometries via an equivariant diffusion process, where each coarse-grained node reflects a fragment in a molecule. Then the coarse-grained nodes are decoded into fine-grained fragments by a message-passing process and a newly designed iterative refined sampling module. Lastly, the fine-grained fragments are then assembled to derive a complete atomic molecular structure. Extensive experiments demonstrate that HierDiff consistently improves the quality of molecule generation over existing methods

연구 동기 및 목표

• 원자 수준의 3D 분자 생성의 한계, 특히 큰 분자에서의 구조적 타당성 저하와 오류 누적 문제를 해결한다.
• 조각 기반 프레임워크에서 비자기적 3D 생성의 과제를 극복한다. 이 경우 조각 간 연결성과 기하학적 구조를 유지해야 한다.
• 원자 수준의 확산이 아닌 조각 수준의 모델링을 통해 화학적 인도적 편향을 인코딩하여 생성 품질을 향상시킨다.
• 실제 분자의 기하학적 및 위상적 제약 조건을 준수하는 안정적이고 고품질의 3D 구조 생성을 가능하게 한다.
• 최신의 확산 및 자기적 모델들과 비교해도 약물 유사성 및 구조적 정확도에서 더 뛰어난 성능을 달성한다.

제안 방법

• 두 단계로 구성된 계층적 확산 프레임워크를 제안한다: 군집 수준의 표현 생성 후 원자 수준의 구조 복원.
• 기하학적 및 화학적 제약 조건을 유지하면서 비자기적 방식으로 등변성 확산 모델을 사용해 조각 수준의 좌표와 특징을 생성한다.
• 화학적으로 해석 가능한 조각 수준의 특징 두 가지 유형을 도입한다: 원소 기반 히스토그램과 성질 기반 기술자.
• 조각 복원 과정에서 위상 일관성과 연결성을 확보하기 위해 등변성 메시지 전달 네트워크를 적용한다.
• 기하학적 편향을 교정하고 조각 및 원자 수준의 구조적 정확도를 향상시키기 위해 반복적 정밀 조정 샘플링 모듈을 설계한다.
• 결합된 정밀 조정된 조각들로부터 최종 원자 수준의 3D 구조를 조립하며, 결합 Valence 및 공간 제약 조건을 유지한다.

실험 결과

연구 질문

• RQ1 원자 수준의 확산 모델과 비교해 볼 때, 군집에서 원자로의 전이 전략을 사용하는 계층적 확산 프레임워크는 3D 분자 생성의 구조적 타당성과 품질을 향상시킬 수 있는가?
• RQ2 조각 기반 모델링은 큰 분자 생성에서 오류 누적 문제를 어떻게 완화하고 일반화 능력을 향상시키는가?
• RQ3 화학적으로 해석 가능한 군집 수준의 특징(원소 기반 대비 성질 기반)은 약물 유사성 및 안정된 구조 생성에 얼마나 큰 영향을 미치는가?
• RQ4 조각 수준의 추상화를 통한 비자기적 3D 생성은 자기적 기준 모델보다 더 높은 구조적 정확도를 달성할 수 있는가?
• RQ5 제안된 방법의 계산 효율성은 어떠한가? 품질을 손상시키지 않고도 더 적은 확산 단계로 가속화할 수 있는가?

주요 결과

• HierDiff는 QM9와 GEOM_DRUG 데이터셋 모두에서 최신 기술을 초월하며, 비자기적 설정에서 최대 8개의 무거운 원자를 가진 분자에 대해 100%의 타당성을 달성한다.
• GEOM_DRUG 데이터셋에서 HierDiff는 조각 수준의 RMSD 0.609와 원자 수준의 RMSD 0.632를 기록하여 기준 모델들을 크게 앞서간다.
• HierDiff-P(성질 기반 특징)는 GEOM_DRUG에서 더 높은 약물 유사성 점수를 기록했고, 반면 HierDiff-E(원소 기반 특징)는 MCF 및 화학적 안전성 지표에서 더 우수한 성능을 보였다.
• 큰 분자에서 오류 누적 문제를 줄였으며, 20개 이상의 무거운 원자를 가진 분자에서도 높은 타당성을 유지한다. 반면 자기적 모델은 20개 원자를 넘어서면 실패한다.
• 샘플링 효율성은 경쟁력 있다: HierDiff는 250개의 확산 단계를 사용해 1개 샘플당 3.02초가 소요되며, 잠재 공간 차원 수가 감소해 더 큰 배치 크기를 지원한다.
• 시각화 결과에 따르면, HierDiff는 원자 수준 기준 모델 대비 더 안정적인 구조(예: 충만 고리)를 생성하며 기하학적 왜곡을 최소화한다.