A collaborative constrained graph diffusion model for the generation of realistic synthetic molecules

Motivation

• Known: Diffusion model을 기반으로 한 분자 생성 모델들(DiGress, CDGS, GDSS 등)이 개발되었으나, 화학적 유효성 확보, 계산 효율성, 일반화 능력에서 여전히 한계가 있다. Graph neural network 기반 생성 모델들은 분자 공간의 광활함으로 인해 제약조건 학습이 어렵다.
• Gap: 기존 확산 모델들은 화학적 제약조건을 학습해야 하므로 모델 크기가 크고 실제 분자 분포와의 일치도가떨어진다. 또한 100% 화학적 유효성을 보장하지 못하고 세부적인 화학 특성(원자가, 분자식)을 완벽하게 학습하지 못한다.
• Why: 새로운 분자 화합물 개발은 신약 개발, 첨단 소재, 환경친화적 냉매 등 많은 분야에서 중요하며, ~10^60개의 약물 유사 분자 공간을 효율적으로 탐색하기 위해서는 화학적으로 타당하고 사실적인 분자를 빠르게 생성할 수 있어야 한다.
• Approach: Double Edge Swapping을 통해 각 noising/denoising 단계에서 원자 개수와 원자가를 보존하는 discrete diffusion process를 설계하고, Diffusion model과 Time model 두 개의 협력 모델을 훈련하여 확산 진행 상황에 대한 더 정확한 정보를 제공한다.

Achievement

Fig. 2 | Performance comparison of CoCoGraph FPS on GuacaMol benchmark

• 100% 화학적 유효성: CoCoGraph는 가장 포괄적인 벤치마크에서 100% 화학적으로 유효한 분자만 생성하며 기존 모델들을 능가한다.
• 우수한 화학 성질 분포: 36개 화학 물성에 대해 생성된 분자의 분포가 실제 분자의 분포와 통계적으로 가장 가깝다.
• 경량 모델 구조: 기존 최첨단 모델에 비해 최대 10배 적은 매개변수로 더 나은 성능을 달성하며 샘플링 효율이 우수하다.
• 높은 사실성: 유기화학 전문가를 대상으로 한 Turing-like test에서 학부 이상 교육을 받은 화학자가 생성 분자를 62% 정확도로만 구분할 수 있어 현실성을 입증한다.
• 대규모 합성 분자 데이터베이스: 820만 개의 합성 분자로 구성된 데이터베이스를 생성하여 분자 발견에 활용할 수 있다.

How

Fig. 1 | Constrained collaborative graph diffusion model, CoCoGraph.

• Double Edge Swapping(DES): 각 noising 단계에서 두 개의 결합 AB와 CD를 선택하여 제거한 후 새로운 결합 AC와 BD를 형성하여 분자식과 원자가를 보존한다.
• Diffusion model: 입력 시간 t와 분자 그래프를 받아 3개의 GNN 레이어를 통해 노드 임베딩을 생성하고, feedforward network로 각 edge swap의 그럴듯함을 추정한다.
• Time model: 분자 그래프가 실제 분자에서 얼마나 멀어졌는지를 추정하는 보조 모델로, 정규화된 시간 정보를 diffusion model에 제공하여 실제 확산 진행 상황을 반영한다.
• Collaborative mechanism: 샘플링 시 예상 시간 단계 대신 time model의 정규화된 시간을 사용하고, 전체 궤적에서 가장 낮은 예측 시간을 가진 분자를 선택한다.
• 제약 기반 설계: 화학적 제약조건을 확산 과정에 직접 내장하여 무효한 분자 생성을 방지하고 모델이 학습해야 할 내용을 감소시킨다.

Originality

• Double Edge Swapping을 통한 새로운 discrete diffusion process: 기존 diffusion model의 제약조건 학습 필요성을 완전히 제거하고 Molloy-Reed 분포로 수렴하는 구조적 혁신이다.
• Time model을 통한 협력 메커니즘: 확산 진행의 비균일성 문제를 식별하고 이를 해결하기 위해 두 모델 간 협력 구조를 도입한 창의적 접근이다.
• 매개변수 감소 역설: 제약조건 내장으로 더 작은 모델이 더 큰 기존 모델보다 우수한 성능을 달성하는 것은 설계 철학의 우수성을 보여준다.
• 실제 적용성 강화: Inpainting 기반 응용으로 기존 분자 구조를 보존하면서 새로운 화학 단위체를 추가할 수 있도록 확장하여 실제 분자 설계에 활용 가능하다.

Limitation & Further Study

• Turing-like test의 해석 제한: 62% 정확도는 '구별 불가능'을 완전히 의미하지 않으며, 특정 유형 분자에서만 50% 수준의 성능을 보여 일반화의 한계를 시사한다.", '분자 특성 최적화 부재: 모델이 특정 원하는 화학 성질(약효성, 독성 등)을 직접 제어하여 생성하지 못하며, 생성 후 필터링 단계가 필요하다.
• 실제 합성 가능성 검증 미흡: 화학적 유효성과 사실성을 갖춘 분자가 실제로 합성 가능한지에 대한 검증이 논문에서 충분히 다루어지지 않는다.
• 확산 과정의 비균일성 완전 해결 부족: Time model로 개선되었으나, 서로 다른 분자 유형 간 확산 속도 차이의 근본적 원인에 대한 분석이 부족하다.
• 후속 연구 방향: (1) 특정 약물 성질 최적화를 위한 조건부 생성 메커니즘 개발, (2) 실제 합성 경로 가능성을 고려한 모델 확장, (3) 더 다양한 분자 클래스에 대한 성능 평가 필요.

Evaluation

총평: CoCoGraph는 Double Edge Swapping을 통한 제약 기반 확산 모델로 100% 화학적 유효성을 달성하면서도 기존 모델보다 적은 매개변수로 더 현실적인 분자를 생성하는 획기적인 접근법을 제시한다. 약물 발견과 분자 설계의 실제 응용 가능성이 높으며, 구조적 우아함과 실험적 검증이 우수하다.