[논문 리뷰] Diffusion-based Molecule Generation with Informative Prior Bridges

Lemeng Wu, Chengyue Gong외 3명|arXiv (Cornell University)|2022. 09. 02.

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology인용 수 19

한 줄 요약

이 논문은 물리적 및 통계적 사전 지식을 정보적인 사전 다리(informative prior bridges)를 통해 통합함으로써, 확산 기반 분자의 3D 포인트 클러스터 생성을 위한 새로운 프레임워크를 제안한다. 이 사전 다리는 확산 과정의 종료 시점에 목표 데이터와 일치하는 출력을 보장하는 스토케스틱 프로세스이다. 리아푸노프 기반 방법을 사용하여 이러한 다리를 구축함으로써, 생성 품질, 안정성, 포인트 클러스터의 균일성이 향상되어, 아키텍처 제약 없이 효율적인 사전 지식 통합과 더불어 더 적은 샘플링 단계로 최신 기술 수준의 성능을 달성한다.

ABSTRACT

AI-based molecule generation provides a promising approach to a large area of biomedical sciences and engineering, such as antibody design, hydrolase engineering, or vaccine development. Because the molecules are governed by physical laws, a key challenge is to incorporate prior information into the training procedure to generate high-quality and realistic molecules. We propose a simple and novel approach to steer the training of diffusion-based generative models with physical and statistics prior information. This is achieved by constructing physically informed diffusion bridges, stochastic processes that guarantee to yield a given observation at the fixed terminal time. We develop a Lyapunov function based method to construct and determine bridges, and propose a number of proposals of informative prior bridges for both high-quality molecule generation and uniformity-promoted 3D point cloud generation. With comprehensive experiments, we show that our method provides a powerful approach to the 3D generation task, yielding molecule structures with better quality and stability scores and more uniformly distributed point clouds of high qualities.

연구 동기 및 목표

• 확산 기반 생성 모델에 강력한 물리적 및 통계적 사전 지식을 통합하는 데 도전하는 것.
• 모델 아키텍처를 수정하지 않고도 생성 품질과 안정성을 향상시킬 수 있는 유연하고 아키텍처에 구애받지 않는 사전 지식 통합 방법을 개발하는 것.
• 확산 과정에 균일성 촉진 힘을 통합하여 균일하고 고품질의 3D 포인트 클러스터 생성을 가능하게 하는 것.
• 종료 출력을 보장하는 체계적인 리아푸노프 기반 방법을 통해 확산 브릿지(bridge)를 구축하는 것.
• 샘플링 단계를 줄이며도 분자 생성 및 포인트 클러스터 균일성에서 최신 기술 수준의 성능을 입증하는 것.

제안 방법

• 고정된 종료 시간에 주어진 데이터 포인트를 결정적으로 생성하는 스토케스틱 프로세스로 확산 브릿지를 구축하여, 사전 지식 기반 생성을 보장한다.
• 리아푸노프 함수 기반 접근법을 사용하여 다리의 드리프트 항을 시스템적으로 설계하고 결정함으로써, 안정적이고 해석 가능한 사전 지식 통합을 가능하게 한다.
• 3D 생성 작업에서 균일한 포인트 분포를 촉진하기 위해 리에즈 포텐셜(Riesz potential)과 통계 갭 에너지(statistic gap energy)를 도입한다.
• 결합 길이, 각도, 토크 포텐셜 등의 물리 기반 에너지 함수를 사용하여 분자 생성을 안내하면서도 생성 유연성을 유지한다.
• 신경 확산 모델이 사전 다리를 모방하도록 훈련시켜, 역방향 과정이 현실적이고 고품질의 출력으로 유도되도록 한다.
• 브릿지 파라미터(예: α)를 확산 모델과 함께 엔드 투 엔드로 최적화하는 공동 학습 체계를 적용한다.

실험 결과

연구 질문

• RQ1 정보적인 사전 다리가 아키텍처 수정 없이도 확산 기반 분자 생성의 품질과 안정성을 크게 향상시킬 수 있는가?
• RQ2 물리적 및 통계적 사전 지식를 어떻게 체계적으로 확산 과정에 통합하여 현실적이고 균일한 출력으로 유도할 수 있는가?
• RQ3 사전 지식 기반 브릿지가 3D 포인트 클러스터 생성에서 샘플링 단계를 줄여도 생성 품질을 유지하거나 향상시킬 수 있는가?
• RQ4 다양한 사전 힘(예: 리에즈 vs. 통계 갭 에너지)이 포인트 클러스터의 균일성과 형태의 현실성에 어떤 영향을 미치는가?
• RQ5 기존의 확산 모델에 비해 사전 지식 통합을 브릿지 방식으로 구현했을 때 훈련 효율성과 수렴 속도에 어떤 영향을 미치는가?

주요 결과

• 제안된 방법은 분자 생성에서 최신 기술 수준의 성능을 달성하여, 기존의 물리 기반 사전 지식을 통합한 기준 모델보다 품질 및 안정성 지표에서 모두 뛰어난 성능을 보였다.
• 단지 10개의 확산 단계만으로도 기존의 100단계로 훈련된 표준 확산 모델과 동일하거나 이를 초월하는 성능을 달성하여 더 빠른 수렴을 입증하였다.
• 브릿지 과정에 리에즈 에너지와 통계 갭 에너지를 통합함으로써, 균일도가 크게 향상된 포인트 클러스터를 생성하였으며, 특히 통계 갭 에너지가 더 뛰어난 강건성을 보였다.
• 100단계 설정에서, 사전 지식 기반 브릿지가 모든 평가 카테고리에서 기준 확산 모델보다 MMD 및 COV 지표에서 모두 뛰어난 성능을 보였다.
• 모델 학습 과정에 직접적으로 사전 지식를 통합함으로써, 아키텍처 제약 없이 고품질이고 안정적인 분자 생성을 달성하였다.
• 시각화 결과는 사전 지식 기반 브릿지가 더 현실적이고 균일하게 분포된 포인트 클러스터를 생성함을 확인하였으며, 리에즈 에너지는 강한 반발력으로 인해 가끔 외곽치(outlier)를 유발하는 경향이 있었다.