DriftingMol: Decoder-Coupled Drift for One-Pass Property-Conditional Molecular Generation

기술 요약: DriftingMol

문제 정의
본 논문은 연속적인 목표 속성 값 (예: QED, LogP) 에 반응하는 유효하고 다양한 분자를 낮은 샘플링 비용으로 생성해야 하는 속성 조건부 분자 생성의 과제를 다룹니다. 기존 방법들은 트레이드오프에 직면해 있습니다: 반복적 샘플링 접근법 (확산 또는 강화 학습 등) 은 추론 시 계산 비용이 많이 들며, 잠재 공간 또는 그래프 기반 디코더는 이산적 디코딩이 인근 잠재 상태를 붕괴시킬 때 연속적인 속성 신호를 종종 상실합니다. 또한, 분자의 유효성 (구문 및 원자가 정확성) 을 보장하는 것은 연속적인 속성 제어를 복잡하게 만드는 이산적 제약으로 남아 있습니다.

방법론
DriftingMol 은 SELFIES 잠재 분자 공간에 "드리프트 모델 (drifting models)"을 적용하는 2 단계 프레임워크를 도입합니다.

1. 2 단계 파이프라인:

   • 1 단계: ZINC250K 데이터셋으로 SELFIES $\beta$-VAE 를 학습합니다. 인코더는 토큰 시퀀스를 256 차원 잠재 공간 ($z$) 으로 매핑하고, 원샷 (one-shot) Transformer 디코더는 SELFIES 위치를 병렬로 예측합니다. 이 단계는 생성기가 학습하는 것이 아니라 SELFIES 표현에서 유래된 유효성이 부여된 정규화된 연속 잠재 공간을 확립합니다.
   • 2 단계: VAE 를 고정합니다. DiT 스타일 생성기 ($f_\theta$) 를 학습시켜 가우스 노이즈, 선택적 속성 조건 ($c$), 그리고 임베딩된 가이드 스케일 ($\alpha$) 을 단일 순전파 (forward pass) 로 잠재 벡터 $\hat{z}$에 매핑합니다. 최종 분자는 고정된 SELFIES 디코더를 통해 $\hat{z}$를 디코딩하여 얻습니다.

2. 디코더 결합 드리프트 (Decoder-Coupled Drift):
   핵심 혁신은 드리프트 필드 (drift field) 의 정의에 있습니다. 외부 특징 추출기를 사용하는 대신, DriftingMol 은 고정된 VAE 디코더의 숨겨진 표현 ($\phi$) 을 드리프트 특징 맵으로 활용합니다.

   • 드리프트 필드: 생성된 잠재 ($\hat{z}_i$) 와 참조 잠재 ($z_{ref, j}$) 간의 커널 드리프트 필드는 디코더 특징을 기반으로 구성됩니다: $V_{\tau, i} = \sum_j W^\tau_{ij} (\phi(z_{ref, j}) - \phi(\hat{z}_i))$.
   • 그라디언트 결합: 결정적으로, 디코더 가중치는 고정되지만, 드리프트 손실에서 디코더 특징 맵 $\phi(\hat{z})$를 거쳐 생성기 매개변수 $\theta$로 되돌아가는 그라디언트 경로는 개방된 상태로 유지됩니다. 이는 생성기 업데이트가 디코더의 국소적 구조 민감도에 의해 가중되는 "풀백 메트릭 (pullback metric)"을 생성합니다.
   • 다양성: 모드 붕괴를 방지하고 고유성을 보장하기 위해 $k$-NN 잠재 반발 항 ($L_{zdiv}$) 을 손실 함수에 추가합니다.

3. 학습 목적 함수:
   생성기는 속성 일치 참조로 정의된 목표 분포를 향해 생성된 잠재의 디코더 특징을 끌어당기는 드리프트 손실 ($L_{drift}$) 을 최소화하며, 여기에 다양성 정규화가 결합됩니다.

주요 기여

• SELFIES 잠재 드리프트 파이프라인: 유효성 (SELFIES/VAE 가 처리) 과 제어 가능성을 분리하는 속성 조건부 생성을 위한 저비용 프레임워크.
• 디코더 결합 드리프트: 고정된 VAE 디코더를 드리프트 특징 공간으로 사용하면서 디코더를 거쳐 생성기로 가는 그라디언트 경로를 보존하는 메커니즘. 이는 드리프트 그라디언트를 실제 분자 디코딩 과정과 정렬시킵니다.
• 어블레이션 증거: 디코더 결합이 속성 제어의 주된 동인임을 입증하는 포괄적인 15 조건 어블레이션 연구로, 잠재 공간, 무작위 특징, 외부 특징 드리프트 변형들을 능가합니다.
• 효율성: 이 방법은 반복적 탈노이즈 또는 ODE 통합을 피하면서 생성기 평가 한 번과 고정된 디코더 통과 한 번만으로 속성 편향을 달성합니다.

결과
ZINC250K 에서 QED 및 다중 속성 조건부로 평가한 결과:

• QED 제어: 기본 DriftingMol 설정은 0.493의 QED 스피어만 상관관계 ($\rho$) 와 **94.7%**의 고유성을 달성합니다. 가장 강력한 디코더 결합 변형 (다양성 정규화 없이) 은 $\rho = 0.510$에 도달하지만 고유성은 감소합니다 (74.1%).
• 다중 속성 조건부: 4 가지 속성 (QED, SA, LogP, MolWt) 조건부 프로토콜 하에서 디코더 결합 드리프트는 평균 스피어만 상관관계 0.598을 달성합니다.
• 어블레이션 발견:
  • 결합의 필요성: 디코더를 분리하거나 스톱-그라디언트를 사용하는 모델은 거의 0 에 가까운 QED 상관관계 ($\rho \approx 0.01$) 와 매우 낮은 고유성을 보입니다.
  • 특징 품질: 외부 학습 특징 (예: LatentMAE) 이나 무작위 특징은 디코더의 야코비안 없이는 예측 특징만으로는 불충분함을 나타내는 현저히 낮은 상관관계 ($\rho \approx 0.28\text{--}0.33$) 를 초래합니다.
  • 다양성: $z$-다양성 항을 제거하면 상관관계는 약간 증가하지만 고유성은 극적으로 감소하여 모드 붕괴를 방지하는 해당 항의 역할을 확인시켜 줍니다.
• 베이스라인: 동일한 잠재 캐시에서 학습된 동일한 백본 베이스라인 (VAE, WGAN-GP, DDPM, Flow-Matching) 은 현저히 낮은 상관관계 ($\rho < 0.16$) 를 달성했습니다.

의의 및 주장
본 논문은 디코더 결합 드리프트가 속성 편향 분자 생성을 위한 효율적인 메커니즘이라고 주장합니다. 그 의의는 고정된 디코더가 단순히 특징 공간을 제공하는 것이 아니라, 생성기가 최적화되는 기하학을 정의하는 특정 풀백 메트릭 ($J_\phi^\top J_\phi$) 을 제공한다는 수학적 통찰에 있습니다. 이 메트릭은 생성기의 업데이트를 분자 디코더의 구조적 민감도와 정렬시킵니다.

저자들은 다음을 강조합니다:

1. 유효성은 표현 수준: 유효성은 조건부 생성기가 학습하는 것이 아니라 SELFIES 표현과 고정된 디코더에 의해 보장됩니다.
2. 단일 통과 효율성: 이 방법은 반복적 샘플링의 높은 추론 비용을 피합니다.
3. 속성 편향 대 정확성: 이 시스템은 목표 영역을 향한 라이브러리 풍부화를 위한 속성 편향 생성을 위해 설계되었으며, 속성 정확 최적화를 위한 것이 아닙니다. 이는 합성 가능성이나 생물학적 활성을 보장하지 않으며, 생성된 분자는 하류 필터링이 필요합니다.
4. 메커니즘 특이성: 성능 향상은 단순히 사전 학습된 VAE 나 외부 특징 추출기의 사용이 아니라, 디코더를 통한 그라디언트 결합에 구체적으로 기인합니다.