기술 요약: FlashMol

문제 제기

화학적 유효성이 보장된 3 분자 입체구조 (conformations) 생성은 컴퓨팅 기반 약물 발견에서의 핵심 병목 현상입니다. GeoLDM 과 EDM 과 같은 확산 기반 모델들은 고품질이고 안정적인 분자를 생성하지만, 일반적으로 샘플 당 수백에서 수천 개의 순차적 네트워크 평가 (NFE) 를 필요로 합니다 (예: 1000 NFE). 이러한 계산 비용은 대규모 in silico 스크리닝을 비현실적으로 만듭니다. 최근 몇 단계 (few-step) 방법들은 NFE 를 12~50 범위로 줄였지만, 종종 샘플의 안정성과 유효성이 급격히 떨어지는 심각한 트레이드오프를 겪습니다. 현재까지 이미지 및 비디오 생성에서 표준이 되어 가는 10 단계 이하의 영역에서 신뢰할 수 있고 고품질의 분자 생성을 달성한 기존 방법은 없습니다.

방법론

저자들은 FlashMol을 제안합니다. 이는 4 단계와 같이 극히 적은 단계로도 고품질 3 분자 입체구조를 생성할 수 있는 초고속 생성 모델입니다. FlashMol 은 분자 도메인에 특화되어 적용된 **분포 매칭 증류 (Distribution Matching Distillation, DMD)**를 사용하여 1000 단계의 교사 모델 (GeoLDM) 로부터 몇 단계의 학생 생성기를 증류합니다. 이 접근법은 분자 생성에 DMD 를 적용할 때 내재된 두 가지 주요 과제를 해결합니다:

1. 재배열된 시간 단계 (Respaced Timesteps) 를 통한 열악한 초기 모드 탈출

표준 DMD 는 학생 모델이 유용한 최적화 신호를 제공하기 위해 합리적인 초기 샘플을 생성하는 것에 의존합니다. 그러나 분자 구조는 잡음에 매우 민감합니다. 결합 길이와 각도의 작은 편차만으로도 분자가 화학적으로 유효하지 않게 될 수 있습니다. 균일한 시간 단계로 DMD 를 직접 적용하면, 학생 모델이 생성하는 안정적인 분자가 거의 0% 에 달하게 되어 기울기 폭발 (gradient explosions) 이 발생하고 훈련이 비효율적으로 됩니다.

이를 해결하기 위해 FlashMol 은 EDM 의 설계 공간 분석에서 영감을 받은 시간 단계 재배열 전략을 도입합니다. 균일한 시간 단계 대신, 이 방법은 고잡음 수준에서 큰 탈잡음 (denoising) 단계를 앞당기고, 정밀한 기하학적 세부 사항이 복원되는 저잡음 영역을 위해 많은 작은 단계를 남겨둡니다.

• 메커니즘: 잡음 스케줄 $\sigma_i$는 $\sigma_i = (\sigma_{max}^{1/\rho} + \frac{i}{N-1}(\sigma_{min}^{1/\rho} - \sigma_{max}^{1/\rho}))^\rho$로 정의됩니다.
• 영향: 이 전략은 몇 단계의 학생 모델에 실행 가능한 초기화를 제공하여, 증류가 시작되기 전에 안정적이지만 완벽하지는 않은 분자를 생성할 수 있게 합니다. 이는 DMD 훈련 과정을 안정화시켜 기울기 폭발을 방지하고 효과적인 최적화를 가능하게 합니다.

2. $f$-발산 정규화를 통한 다양성 개선

DMD 는 역 Kullback-Leibler (KL) 발산 ($KL(p_{fake} || p_{real})$) 을 최소화하여 학생이 교사의 고확률 모드를 따르도록 하지만, 종종 "모드 추구 (mode-seeking)" 행동을 유발하여 지지 영역을 축소시키고 샘플 다양성을 감소시킵니다.

이를 완화하기 위해 FlashMol 은 DMD 목적 함수에 Jensen-Shannon (JS) 발산 항을 정규화합니다.

• 메커니즘: JS 발산은 순방향 및 역방향 KL 구성 요소를 모두 포함합니다. 순방향 KL 구성 요소는 "평균 추구 (mean-seeking)" 행동을 도입하여 역방향 KL 의 모드 추구 경향을 상쇄합니다.
• 구현: JS 항은 실제 분자와 생성된 분자로 훈련된 GAN 판별자를 사용하여 근사화됩니다. 최종 손실 함수는 $L_{FlashMol} = L_{DMD} + \lambda \cdot L_{JS}$입니다.

주요 기여

1. 분자 생성으로의 DMD 통합: 저자들은 분자 잠재 생성 모델인 GeoLDM 에 Distribution Matching Distillation 을 성공적으로 적용하여 극히 적은 단계로 고충실도 샘플을 달성했습니다.
2. 시간 단계 재배열 전략: 저자들은 분자에서의 분포 증류에 있어 추론 시 시간 단계 스케줄링이 핵심 병목 현상임을 규명했습니다. EDM 에서 영감을 받은 재배열 전략을 도입함으로써 추가 훈련 없이 초기 몇 단계 생성기의 품질을 크게 향상시켜 효과적인 증류를 가능하게 했습니다.
3. 통합 발산 목적 함수: 저자들은 vanilla DMD 목적 함수에 Jensen-Shannon 발산 항을 통합하여 확장함으로써, 고충실도 생성 (역 KL 을 통해) 과 다양성 보존 (순방향 KL 구성 요소를 통해) 을 균형 있게 잡는 통합 프레임워크를 제공했습니다.
4. 성능: 이 방법은 4~8 단계만으로 최첨단 (state-of-the-art) 결과를 달성하여 원래의 1000 단계 교사 모델을 능가하거나与之 맞먹는 성능을 보입니다.

실험 결과

실험은 QM9(13 만 개의 작은 분자) 와 GEOM-DRUG(45 만 개의 더 큰 분자) 데이터셋에서 수행되었습니다.

• 안정성과 유효성:
  • QM9에서 FlashMol 은 4 단계로 **83.89%**의 분자 안정성을 달성하고 8 단계에서는 **94.87%**를 달성합니다. 이는 재배열된 GeoLDM 베이스라인 (4 단계에서 12.43%) 과 GeoRCG(50 단계), MOLTD(12 단계) 와 같은 다른 몇 단계 방법들보다 훨씬 뛰어납니다.
  • GEOM-DRUG에서 FlashMol 은 4 단계로 **84.72%**의 원자 안정성을, 8 단계로 **89.98%**를 달성하여 데이터셋 수준의 원자 안정성인 86.5% 를 초과합니다.
• 가속화: FlashMol 은 고품질 분자 품질을 유지하면서 1000 단계 GeoLDM 교사 모델 대비 최대 250 배의 샘플링 속도 가속화를 달성합니다.
• 조건부 생성: 조건부 생성 작업 ( $\epsilon_{HOMO}$, $\epsilon_{LUMO}$ 등의 특성 예측) 에서 FlashMol(32 단계) 은 1000 단계 베이스라인과 비교하거나 더 나은 평균 절대 오차 (MAE) 를 달성합니다.
• Ablation 연구:
  • 재배열: 재배열 없이 4 단계 생성은 거의 0% 의 안정성을 보입니다. 재배열을 적용하면 증류 없이도 안정성이 약 94% 로 급증합니다.
  • 발산: DMD 만 사용하면 안정성은 향상되지만 다양성 (유효하고 고유한 샘플) 이 감소합니다. JS 항을 추가하면 높은 안정성을 유지하면서 다양성을 회복합니다.

의미 및 주장

이 논문은 FlashMol 이 확산 기반 분자 설계를 고처리량 스크리닝에 실용적으로 만드는 중요한 진전이라고 주장합니다. 샘플링 비용을 1000 단계에서 4 단계까지 줄임으로써 FlashMol 은 산업용 약물 발견 파이프라인에서 이러한 모델의 배포를 방해해 온 주요 계산적 장벽을 제거합니다.

저자들은 그들의 접근 방식이 교사 모델을 넘어선 새로운 생성 능력을 도입하는 것이 아니라, 효율적인 증류를 통해 기존 모델의 잠재력을 unlocking 한다는 점을 강조합니다. 그들은 다양성과 안정성 사이의 트레이드오프 (점수 증류 목적 함수의 알려진 특성) 가 존재하지만, 그들의 방법은 신뢰할 수 있는 유효한 화학 구조 생성을 가능하게 하는 정도로 모드 붕괴를 성공적으로 완화했다고 지적합니다. 이 연구는 단백질 - 리간드 공동 설계 및 신소재 발견과 같은 도메인에서 다른 등변 잠재 확산 모델에 유사한 증류 기법이 적용될 수 있음을 시사합니다.