[논문 리뷰] Scalable Autoregressive 3D Molecule Generation

Austin H. Cheng, Chong Sun외 1명|ArXiv.org|2025. 05. 20.

Nanofabrication and Lithography Techniques인용 수 3

한 줄 요약

Quetzal은 연쇄적 자기회귀 모델로, 3D 분자를 원자 하나씩 구축하며 인과 트랜스포머와 Diffusion MLP를 사용해 이산 원자 유형과 연속 위치를 예측하고, 확산 모델과의 품질 경쟁력을 가지면서 생성 속도가 더 빠릅니다.

ABSTRACT

Generative models of 3D molecular structure play a rapidly growing role in the design and simulation of molecules. Diffusion models currently dominate the space of 3D molecule generation, while autoregressive models have trailed behind. In this work, we present Quetzal, a simple but scalable autoregressive model that builds molecules atom-by-atom in 3D. Treating each molecule as an ordered sequence of atoms, Quetzal combines a causal transformer that predicts the next atom's discrete type with a smaller Diffusion MLP that models the continuous next-position distribution. Compared to existing autoregressive baselines, Quetzal achieves substantial improvements in generation quality and is competitive with the performance of state-of-the-art diffusion models. In addition, by reducing the number of expensive forward passes through a dense transformer, Quetzal enables significantly faster generation speed, as well as exact divergence-based likelihood computation. Finally, without any architectural changes, Quetzal natively handles variable-size tasks like hydrogen decoration and scaffold completion. We hope that our work motivates a perspective on scalability and generality for generative modelling of 3D molecules.

연구 동기 및 목표

• 원자 단위의 고정 크기 확산 모델을 넘어서는 3D 분자 자기회귀 생성의 동기 부여.
• 정확한 가능도 계산을 가능하게 하면서 임의 크기의 분자 출력을 가지는 확장 가능한 아키텍처를 개발하는 것.
• 3D 구조를 생성하기 위해 이산 원자 유형 예측기와 연속 위치 모델을 활용.
• QM9 및 GEOM 데이터셋에서 확산 모델과의 경쟁력 있는 샘플 품질을 입증.
• 아키텍처 변경 없이 수소 장식 및 골격 완성 등의 유연한 작업을 보여주는 것.

제안 방법

• 모델 아키텍처는 다음 원자 유형을 예측하기 위한 인과 트랜스포머와 다음 원자 위치를 모델링하기 위한 Diffusion MLP를 결합합니다.
• 다음 위치 분포 p_pos는 DiffLoss 목표로 학습된 조건부 확산 모델입니다.
• 원자 유형 예측은 접두사 임베딩에서 출력된 로그잇에 대해 소프트맥스 softmax를 사용하는 MLP를 통해 이뤄집니다.
• 위치 조건 벡터 z_i는 다음 원자 유형과 접두사 표현을 소비하는 두 번째 트랜스포머에 의해 생성됩니다.
• 훈련은 DiffMLP의 시간에 따른 배치 확장을 포함한 다중 타임스텝의 전체 시퀀스 감독을 사용합니다.
• 생성은 자기회귀적으로 진행되며, 중지 토큰이 생성될 때까지 원자 유형을 샘플링한 후 위치를 샘플링하고, O(n)번의 트랜스포머 통과와 O(n N_diff) 좌표 업데이트를 수행합니다.

실험 결과

연구 질문

• RQ1 표준 벤치마크(QM9, GEOM)에서 확산 모델의 생성 품질에 비견될 수 있는 3D 분자 모델을 자동회귀 방식으로 만들 수 있는가?
• RQ2 제안된 Quetzal 아키텍처가 더 빠른 생성 속도를 가능하게 하면서 분자의 높은 유효성(Validity)과 고유성을 유지하는가?
• RQ3 아키텍처를 변경하지 않고 수소 장식 및 골격 완성과 같은 가변 크기 생성 작업을 모델이 지원하는가?
• RQ4 이 자동회귀 확산 프레임워크에서 원자별 좌표에 대한 정확한 가능도 계산이 가능한가?
• RQ5 생성 순서가 성능과 일반화에 어떤 영향을 미치며, 일반적인 순서 선택에 대해 모델이 얼마나 강건한가?

주요 결과

• Quetzal은 이전의 자기회귀 기준선을 능가하고 QM9 및 GEOM에서 샘플 품질 지표상 최신 확산 모델과 경쟁합니다.
• 생성은 원자당 하나의 트랜스포머 패스와 좌표에 대한 DiffMLP를 사용함으로써 확산 기반 접근 방식보다 현저히 빠릅니다.
• 모델은 확산만 방법들에는 일반적이지 않은, 정확한 발산 기반 가능도 계산을 가능하게 합니다.
• Quetzal은 아키텍처 변경 없이 수소 장식 및 골격 완성과 같은 가변 크기 작업을 기본적으로 지원합니다.
• 수소 장식 실험에서 학습 데이터 순서와 원자 주문이 일치할 때 높은 정확도를 보입니다.