VEDA: 3D Molecular Generation via Variance-Exploding Diffusion with Annealing

본 논문은 분산 발산 확산(variance-exploding diffusion)과 어닐링(annealing)을 결합하여 3D 분자 생성에서의 효율성-정확도 트레이드오프를 극복하고, 단 100회의 샘플링 단계만으로 최첨단 화학적 타당성과 놀라운 형태적 안정성을 달나성하는 통합된 SE(3)-등변 프레임워크인 VEDA를 제안한다.

핵심

당신이 로봇에게 복잡한 3D 분자, 예를 들어 아주 작고 정교한 레고 조각품 같은 것을 만드는 법을 가르치려 한다고 상상해 보세요. 로봇은 안정적이고 제대로 작동하는 분자를 만들기 위해 모든 원자(레고 브릭)가 공간의 정확히 어디에 위치해야 하는지를 알아내야 합니다.

이 논문은 이 로봇이 업무를 더 잘, 더 빠르게 수행할 수 있도록 돕는 VEDA라는 새로운 방법을 소개합니다. VEDA가 어떻게 작동하는지 쉬운 개념으로 나누어 설명하겠습니다.

문제점: "빠르지만 엉성한" 대 "느리지만 완벽한" 딜레마

현재 로봇이 이러한 분자를 만드는 데는 두 가지 주요 방식이 있습니다:

1. 스피드스터 (Flow-based models): 이들은 마치 빨리 감기 비디오와 같습니다. 분자를 매우 빠르게 만들어내지만, 종종 결과물이 엉망인 브릭 더미처럼 보여서 서로 결합되지 않거나 기하학적으로 틀려 보입니다. 이들은 분자가 뒤틀리고 회전하는 다양한 방식을 포착하는 데 어려움을 겪습니다.
2. 완벽주의자 (Denoising Diffusion models): 이들은 돌을 깎아내는 조각가와 같습니다. 노이즈 덩어리에서 시작하여 서서히 완벽한 모양을 깎아 나갑니다. 결과는 매우 정확하지만, 수천 번의 미세한 단계를 거쳐야 하기 때문에 조각을 완성하는 데 시간이 매우 오래 걸립니다.

저자들은 빠르면서도(스피드스터처럼) 정확한(완벽주의자처럼) 로봇을 만들고자 했습니다.

해결책: VEDA (Variance-Exploding Diffusion with Annealing)

VEDA는 두 방식의 장점을 결합한 새로운 프레임워크입니다. 이것을 "시뮬레이티드 어닐링(Simulated Annealing, 담금질)" 과정이라고 생각해보세요.

비유: 퍼즐 조각이 든 상자를 흔들기
퍼즐 조각(원자)들이 무작위로 흩어져 있는 상자가 있다고 상상해 보세요. 당신은 이 조각들이 올바른 그림으로 딱 맞게 끼워지도록 만들고 싶습니다.

• 기존 방식: 당신은 조각들을 제자리로 살짝 밀어 넣으려고 할 수도 있습니다. 만약 너무 살살 밀면, 조각들이 잘못된 위치에 갇혀 버립니다(지역적 함정). 반대로 너무 세게 밀면, 조각들이 부서져 버립니다.
• VEDA 방식: VEDA는 상자를 격렬하게 흔드는 것부터 시작합니다. 조각들을 멀리 던져버림으로써 구조를 사실상 "녹여서", 잘못된 연결 관계가 방해가 되지 않도록 합니다. 그런 다음, 상자를 천천히 식힙니다(이것이 "어닐링" 부분입니다). 식어감에 따라 조각들은 가장 안정적이고 에너지 효율적인 위치로 자리 잡게 됩니다.

엄청난 양의 "노이즈"(흔들림)에서 시작하여 이 노이즈를 조절하는 방식을 정교하게 제어함으로써, VEDA는 분자가 최적의 형태를 찾도록 도와주며, "잘못된 위치에 갇히는" 문제를 피하게 해줍니다.

VEDA가 사용하는 세 가지 핵심 기술

1. "어닐링" 흔들기 (Variance-Exploding)
단순히 약간의 노이즈를 더하는 대신, VEDA는 시작 단계에서 엄청난 양의 노이즈를 추가합니다. 이것은 구겨진 종이를 공중에 던져서 완전히 펼쳐 놓은 뒤에 다시 매끄럽게 펴려고 노력하는 것과 같습니다. 이는 분자가 초기에 나쁜 형태에 갇히지 않도록 보장합니다.

2. "안티-아이덴티티" 교정 (Preconditioning)
VEDA가 사용하는 AI 두뇌(신경망)에는 나쁜 습관이 하나 있습니다. 바로 자신이 보는 것을 그대로 복사하려는 경향입니다. 만약 노이즈가 섞인 분자를 보여주면, AI는 그것을 어떻게 고칠지 고민하는 대신 단순히 "여기에 노이즈 섞인 분자가 있습니다"라고 말해버리는 식입니다.

• 해결책: VEDA는 AI가 예측을 하기 전에 자신의 "복사하기" 성향을 먼저 빼도록 강제합니다. 이는 마치 예술가에게 "밑그림을 따라 그리기만 하지 말고, 스케치와 실제 그림 사이의 '차이점'이 무엇인지 말해달라"고 요구하는 것과 같습니다. 이를 통해 AI는 실제 구조를 훨씬 빠르게 학습할 수 있습니다.

3. "골든 아워" 스케줄러 (Arcsin Scheduler)
로봇이 분자를 만들 때 여러 단계를 거칩니다. 저자들은 모든 단계가 똑같이 중요한 것은 아니라는 점을 깨달았습니다.

• 비유: 케이크를 굽는 과정을 생각해 보세요. 처음 10분(반죽하기)과 마지막 10분(식히기)도 중요하지만, 케이크가 부풀어 오르는 중간 단계가 가장 결정적입니다.
• 해결책: VEDA는 (arcsin이라는 수학 함수에 기반한) 특별한 스케줄을 사용하여 분자의 형태가 실제로 형성되는 "중간" 단계에 더 많은 시간과 노력을 쏟습니다. 지루한 부분은 무시하고 가장 중요한 순간에 에너지를 집중합니다.

결과: 빠르고, 안정적이며, 정확함

논문은 QM9과 GEOM-DRUGS라는 두 가지 큰 분자 데이터셋을 통해 VEDA를 테스트했습니다.

• 속도: VEDA는 "스피드스터" 모델만큼 빠릅니다. 단 100단계 만에 분자를 생성할 수 있으며, 기존의 "완벽주의자" 모델들은 1,000단계가 필요했습니다.
• 정확도: VEDA가 만드는 분자는 믿을 수 없을 정도로 안정적입니다. 과학자들이 이 분자들을 테스트했을 때, 분자를 물리적으로 현실적이게 만드는 데 필요한 에너지(수정 작업)가 이전의 최고 방법보다 90%나 낮았습니다.
  • 비유: 만약 기존 방식이 서 있기 위해 32 단위의 접착제가 필요한 흔들거리는 탑을 만들었다면, VEDA는 단 1.7 단위의 접착제만으로도 충분히 서 있을 수 있는 탑을 만든 것입니다.

요 요약

VEDA는 "흔들고 안착시키는" 전략을 사용하는 3D 분자 생성 방식입니다. 혼돈의 상태에서 시작하여, AI가 게으름을 피우지 못하도록 스마트한 수학을 사용하며, 분자 형성이 일어나는 가장 결정적인 순간에 노력을 집중합니다. 그 결과, 현재 사용 가능한 가장 빠른 방법들만큼 빠르게 화학적으로 정확하고 안정적인 분자를 구축하는 시스템을 만들어냈습니다.

기술 요약

기술 요약: VEDA – 어닐링을 결한 분산 폭발형 확산 기반의 3D 분자 생성

문제 정의

확산 모델(Diffusion models)은 3D 분자 생성 분야에서 지배적인 도구로 부상했으나, 샘플링 효율성과 컨포메이션 정확도(conformational accuracy) 사이의 근본적인 트레이드오프에 직면해 있습니다.

• 플로우 기반 모델(예: EquiFM, SemlaFlow)은 빠른 샘플링을 제공하지만, 분자 컨포메이션의 다중 모드 분포(multimodal distributions)를 포착하는 데 어려움을 겪어 기하학적으로 부정확한 구조를 생성하는 경우가 많습니다.
• 디노이징 확산 모델(예: EDM, GeoLDM)은 높은 정확도를 달로는 하지만, 샘pling 속도가 느리다는 단점이 있습니다. 이러한 한계는 확산 역학과 SE(3)-등변(equivariant) 아키텍처 간의 최적화되지 않은 통합에 기인합니다.

핵심적인 미해결 과제는 SE(3)-등변 신경망(예: EGNN)의 **귀납적 편향(inductive bias)**입니다. 이러한 네트워크는 메시 전달 메커니ism과 엄격한 기하학적 제약으로 인해 **항등 함수와 유사한 매핑(identity-like mappings)**을 학습하려는 강한 경향을 보입니다. 표준 확산 프레임워크에서 이러한 편향은 잔차 기반 학습 목표(residual-based learning objectives)와 충돌하여 학습 불안정성과 최적화되지 않은 생성을 초래합니다. 또한, 기존 방식들은 네트워크 아키텍처를 개선하거나 도메인 지식을 주입하는 데 집중할 뿐, 아키텍처의 편향에 맞추기 위한 원칙적인 확산 프로세스 재설계는 간과하고 있습니다.

방법론

저자들은 VEDA(Variance-Exploring Diffusion with Annealing, 어닐링을 결합한 분산 폭발형 확산)를 제안합니다. 이는 분산 폭발(VE) 확산과 시뮬레이티드 어닐링(simulated annealing) 원리를 결합하여 컨포메이션 정확도가 높은 3D 분자 구조를 생성하는 통합된 SE(3)-등변 프레임워크입니다. 이 프레임워크는 단일 프로세스 내에서 연속적인 좌표와 이산적인 원자 특징을 모두 처리합니다.

1. 어닐링을 결합한 분산 폭발(VE) 확산

VEDA는 노이즈가 시뮬레이티드 어닐링과 기능적으로 유사하게 주입되는 VE 스케줄을 채택합니다.

• 순방향 과정(Forward Process): 깨끗한 데이터 $x_0$에 가우시안 노이즈를 더하여 $x_t = x_0 + t\epsilon$ 형태로 오염시킵니다. 노이즈 레벨 $t$는 로그-노멀 분포에서 샘플링됩니다.
• 어닐링 비유: 노이즈 레벨 $t$는 온도 파라미터 역할을 합니다. 샘플링 중 $t$가 감소함에 따라, 시스템은 에너지 지형의 광범위한 탐색(높은 노이즈)에서 저에너지의 안정적인 컨포메이션으로의 수렴으로 전이됩니다.
• 무한 분리(Infinite Separation): VE 정식화는 원자 간의 무한 분리 상태($t \to \infty$)를 이상적인 노이즈 상태로 취급하며, 이는 분자 생성 시 표준 가우시안 사전 확률보다 우수합니다.

2. SE(3)-등변 네트워크를 위한 프리컨디셔닝(Preconditioning)

네트워크의 항등 편향과 잔차 학습 목표 사이의 충돌을 해결하기 위해, VEDA는 새로운 프리컨디셔닝 스킴을 도입합니다.

• 문제점: 표준 좌표 예측 네트워크는 입력 $x_t$와 매우 상관관계가 높은 $F_\theta(x_t)$를 출력하며, 이는 네트워크가 입력과 상관관계가 없는 잔차를 예측해야 한다는 가정에 위배됩니다.
• 해결책: 저자들은 네트워크 출력에서 스케일링된 항등 성분을 명시적으로 뺄셈합니다. 수정된 디노이저는 다음과 같이 정의됩니다:
  $$D_\theta(x_t; t) = c_{skip}x_t + c_{out} (F_\theta(c_{in}x_t; c_{noise}) - \alpha_t c_{in} x_t)$$
• 최적 계수 ($\alpha_t$): 계수 $\alpha_t$는 선형 최소 평균 제곱 오차(LMMSE) 솔루션으로 도출되며, 이는 입력에 따른 실제 노이즈의 최적 선형 예측값을 나타냅니다. 이는 다음과 같이 계산됩니다:
  $$\alpha_t = \frac{\sigma_d t}{\sigma_d^2 + t^2}$$
  여기서 $\sigma_d$는 데이터 분포의 표준 편차입니다. 이는 학습 목표를 네트워크의 귀납적 편향과 일치시킵니다.

3. Arcsin 기반 노이즈 스케줄러

저자들은 로그 신호 대 잡음비(log-SNR)의 임계 구간에 샘플링 단계를 집중시키기 위한 새로운 스케줄러를 제안합니다.

• 경험적 분석에 따르면, log-SNR $\approx$ 0 근처의 단계들이 최종 분자 구조 형성에 매우 중요합니다.
• 이 스케줄러는 정규화된 단계 인덱스 $u$와 농도를 조절하는 $\rho$를 사용하여 다음과 같은 arcsin 함수를 통해 샘플링 단계를 중간 범위에 클러스터링합니다:
  $$w(u) = (1 - \rho) u + \rho \frac{2}{\pi} \arcsin(\sqrt{u})$$
  이는 특히 화학적으로 민감한 구성에서 구조적 충실도를 향상시킵니다.

4. 샘플링 전략: 스토캐스틱 어닐링(Stochastic Annealing)

샘플링 프로세스는 하이퍼파라미터 $\gamma > 0$에 의해 제어되는 증폭된 노이즈 주입 전략을 사용합니다.

• 섭동(Perturbation): 디노이징 전에 증폭된 노이즈가 현재 샘플에 주입됩니다 ($\hat{x}_i = x_i + \sqrt{\hat{t}_i^2 - t_i^2} \cdot \epsilon$).
• 효과: 이는 분자 퍼텐셜 에너지 표면의 가우시안 평활화(Gaussian smoothing)와 동일하며, 국소 최솟값(local minima)과 표면 거칠기를 억제하여 궤적이 전역 저에너지 베이슨(global low-energy basin)을 찾도록 돕습니다.
• 이산적 정교화(Discrete Refinement): 이산적 특징(원자 유형, 결합)의 경우, V는 연속적 노이즈 스케줄과 정렬된 시간 의존적 마스크 비율을 사용하는 이산 플로우 매칭(DFM) 기반의 마스크 토큰 정교화 프로세스를 사용합니다.

주요 기여

1. 통합된 VE 프레임워크: VEDA는 3D 분자의 하이브리드 이산-연속 영역에 분산 폭발 패러다임을 적용하여, 원자 유형과 좌표를 시뮬레이티드 어닐링과 유사한 프로세스로 통합한 최초의 프레임워크입니다.
2. 이론적으로 근거한 프리컨디셔닝: 최적의 선형 항등 성분을 명시적으로 제거함으로써 좌표 예측 SE(3) 네트워크의 귀납적 편향을 교정하고, 학습 안정성과 정렬을 개선하는 새로운 프리컨디셔닝 스킴을 제공합니다.
3. Arcsin 스케줄러: 샘플링을 임계 log-SNR $\approx$ 0 영역에 집중시켜 초기 탐색과 후기 정교화 사이의 균형을 맞추는 새로운 노이즈 스케줄러를 제안합니다.
4. 성능: VEDA는 단 100번의 샘플링 단계만으로 플로우 기반 모델의 효율성을 유지하면서도 기하학적 정확도 면에서 이를 크게 능가하는 최첨단(SOTA) 원자가(valency) 안정성과 유효성을 달성했습니다.

실험 결과

프레임워크는 QM9 및 GEOM-DRUGS 데이터셋에 대해 평가되었습니다.

• QM9 (결합 암시형 & 결합 명시형):

  • VEDA-E (암시형): 단 50번의 샘플링 단계만으로 **97.9%**의 유효 및 고유(valid and unique) 비율을 달성하여, EDM(1000단계에서 90.9%)과 GeoLDM(1000단계에서 92.6%)을 크게 앞질렀습니다.
  • VEDA-S (명시형): 거의 완벽한 원자 및 분자 안정성을 보였으며, 100단계에서 **98.9%**의 유효 및 고유 점수를 기록했습니다.

• GEOM-DRUGS (물리적 실재성):

  • 이완 에너지 ($\Delta E_{relax}$): GFN2-xTB 최적화 과정 중의 에너지 변화를 측정하는 지표입니다. VEDA-S는 중앙값 $\Delta E_{relax}$ 1.72 kcal/mol을 달성하여, 아키텍처 베이스라인인 SemlaFlow(32.3 kcal/mol) 대비 90% 감소된 수치를 보였습니다.
  • 효율성 vs 정확도: VED-S는 생성 품질과 계산 비용(NFE) 사이의 트레이드오프에서 최적의 영역을 차지합니다. 이는 경쟁 모델들보다 한 자릿수 더 나은 기하학적 정확도를 제공하면서도 플로우 기반 모델의 효율성을 유지합니다.
  • 어블레이션 연구(Ablation Studies): 프리컨디셔닝, 노이즈 주입 또는 특정 스케줄러를 제거하면 유효성 저하와 이완 에너지 증가가 나타났습니다. 특히, SemlaFlow에서 사용된 최적 운송(OT) 정렬을 추가하면 에너지 지표가 악화되었는데, 이는 확산에 필수적인 독립성 가정을 위반함을 시사합니다.

의의 및 주장

본 논문은 VE 확산과 SE(3)-등변 아키텍처의 원칙적인 통합이 높은 화학적 정확도와 계산 효율성을 동시에 달성할 수 있음을 입증한다고 주장합니다.

• 트레이드오프 극복: 본 연구는 플로우 기반 모델의 속도와 확산 모델의 정확도 중 하나를 선택해야 한다는 기존의 통념에 도전합니다. 확산 프로세스(VE 스케일, 프리컨디셔닝, 스케줄러)를 등변 네트워크의 아키텍처 편향에 맞게 재설계함으로써, VEDA는 이 트레이드오프를 깨뜨립니다.
• 안정성: 생성된 구조는 매우 안정적이며 에너지 이완이 최소한으로 필요합니다. 이는 기하학적 유효성이 매우 중요한 약물 발견과 같은 다운스트림 응용 분야에서 결정적입니다.
• 미래 설계를 위한 토대: 무조건부 생성(unconditional generation)에서의 VEDA의 성공은 속성 가이드 분자 설계(예: 결합 친화도 조건부 생성)를 위한 강력한 토대를 마련합니다. 저자들은 현재 모델의 암시적 속도 예측(implicit velocity prediction)에 대한 한계를 식별하고, 향-단계 통합(one-step integration)을 가능하게 하기 위해 미래 아키텍처가 명시적으로 시간 의존적 벡터 필드를 출력해야 한다고 제안합니다.

요약하자면, VEDA는 제어된 노이즈 역학과 기하학적 제약을 융합함으로써 분자 설계를 위한 효과적인 청사진을 제공하며, 단백질 포켓과 같이 제약이 있는 환경에서 기능적인 분자를 높은 충실도로 생성할 수 있는 경로를 제시합니다.