[논문 리뷰] All-atom Diffusion Transformers: Unified generative modelling of molecules and materials

Chaitanya K. Joshi, Xiang Fu외 5명|ArXiv.org|2025. 03. 05.

Machine Learning in Materials Science인용 수 3

한 줄 요약

ADiT는 비주기성 분자와 주기성 결정 모두를 공동으로 생성하기 위해 단일 잠재 확산 Transformer 프레이크워크를 도입하여 최첨단 결과와 도메인 간 전이 학습을 성공적으로 달성합니다.

ABSTRACT

Diffusion models are the standard toolkit for generative modelling of 3D atomic systems. However, for different types of atomic systems -- such as molecules and materials -- the generative processes are usually highly specific to the target system despite the underlying physics being the same. We introduce the All-atom Diffusion Transformer (ADiT), a unified latent diffusion framework for jointly generating both periodic materials and non-periodic molecular systems using the same model: (1) An autoencoder maps a unified, all-atom representations of molecules and materials to a shared latent embedding space; and (2) A diffusion model is trained to generate new latent embeddings that the autoencoder can decode to sample new molecules or materials. Experiments on MP20, QM9 and GEOM-DRUGS datasets demonstrate that jointly trained ADiT generates realistic and valid molecules as well as materials, obtaining state-of-the-art results on par with molecule and crystal-specific models. ADiT uses standard Transformers with minimal inductive biases for both the autoencoder and diffusion model, resulting in significant speedups during training and inference compared to equivariant diffusion models. Scaling ADiT up to half a billion parameters predictably improves performance, representing a step towards broadly generalizable foundation models for generative chemistry. Open source code: https://github.com/facebookresearch/all-atom-diffusion-transformer

연구 동기 및 목표

• 기초 물리학을 공유하는 분자 및 결정 시스템 전반에 걸친 통합 확산 접근법의 필요성을 제고한다.
• 분자와 재료 모두에 대해 VAE를 통해 공유 잠재 표현을 개발한다.
• 잠재 공간에서 확산 모델을 학습하여 새로운 원자를 샘플링하고 유효한 구조로 복원한다.
• 주기적 시스템과 비주기적 시스템 간의 전이 학습을 시연하고 확장성을 보인다.

제안 방법

• 원자를 일원화된 속성 집합으로 표현(범주형 원자 타입, 3D 좌표, 분수 좌표, 격자 파라미터).
• 분자와 결정의 입력을 공유 잠재 공간으로 매핑하고 재구성하기 위해 Variational Autoencoder (VAE)를 사용한다.
• 잠재 공간에서 Diffusion Transformer (DiT)를 학습하여 새로운 잠재 임베딩을 생성하고 샘플링 시 classifier-free guidance를 사용한다.
• VAE 디코더로 잠재를 디코딩하여 RDKit으로는 비주기성 유효한 분자, PyMatGen으로는 주기성 결정과 같은 유효한 구조를 생성한다.
• 두 단계 학습 채택: 1단계 VAE 재구성, 2단계 잠재 확산으로 클래스 조건이 있는 DiT 및 로토-트랜스레이션 및 주기 대칭성에 대한 데이터 증강.
• 화학 중심 지표(타당성, 안정성, 고유성, 참신성) 및 Posebusters/DFT 체크를 통한 타당성 확인을 평가하고 데이터셋별 기준선과 비교한다.

실험 결과

연구 질문

• RQ1 단일 통합 확산 모델이 높은 타당성과 현실감을 가진 분자와 결정 모두를 생성할 수 있는가?
• RQ2 분자와 결정 데이터를 결합한 공동 학습이 전이 학습을 통해 생성 품질을 향상시키는가?
• RQ3 모델 크기와 학습 규모가 등가 불변 기반 모델에 비해 성능과 속도에 어떤 영향을 미치는가?
• RQ4 MOF와 같은 더 복잡한 하이브리드 및 다른 주기적/비주기적 하이브리드로 확장 가능한가?
• RQ5 도메인 특정 확산 모델과 비교한 학습 및 추론의 효율성은 어느 정도인가?

주요 결과

• 공동 학습된 ADiT는 QM9 분자와 MP20 결정 모두에서 최첨단 또는 경쟁력 있는 결과를 달성한다.
• 결정의 경우 ADiT는 10,000 샘플에서 여러 기준선에 비해 타당성, 안정성 및 참신성이 향상되었다.
• 분자의 경우 ADiT는 타당성 및 고유성과 일치하거나 이를 능가하며 물리적 현실성에 대한 Posebusters 검사도 충족한다.
• 공동 학습은 QM9만 또는 MP20만 버전에 비해 더 높은 타당성과 안정성을 제공하며 전이 학습의 이점을 시사한다.
• ADiTs는 등가 확산 모델에 비해 추론 속도를 크게 향상시켜 단일 V100 GPU에서 20분 미만에 10,000개의 샘플 생성을 가능하게 한다.
• DiT 디노이저 크기(DiT-S에서 DiT-L로)를 확장하면 성능 증가가 예측 가능하고, 더 많은 데이터와 더 큰 모델에서 추가 이득이 있을 것을 시사한다.