VQ-SAD: Vector Quantized Structure Aware Diffusion For Molecule Generation

"VQ-SAD: 분자 생성을 위한 벡터 양자화 구조 인식 확산" 논문에 대한 상세한 기술적 요약입니다.

1. 문제 제기

현재 확산 기반 분자 생성 방법은 세 가지 주요 한계에 직면해 있습니다:

• 상징적 정보의 손실: 대부분의 모델은 원자와 결합을 단순한 원-핫 (one-hot) 인코딩으로 표현합니다. 이는 산소와 인접한 탄소와 황과 인접한 탄소와 같은 서로 다른 화학적 맥락을 동일한 표현으로 축소시켜, 원자가나 작용기와 같은 중요한 상징적 구조 규칙을 무시합니다.
• 지문 (Fingerprinting) 의 한계: 모건 지문 (Morgan fingerprints) 을 사용하는 대안적 접근법은 해시 충돌 (서로 다른 분자가 동일한 지문으로 매핑됨) 에 시달리며, 정보 손실 없이 연속 공간에 임베딩하기 어렵습니다.
• 클래스 불균형 및 상태 충돌: 실제 분자 데이터셋은 종종 희소한 클래스 유형을 가집니다 (일부 원자/결합은 드뭅니다). 이러한 불균형은 소수 클래스 유형을 재구성하는 디노이저 (denoiser) 의 능력을 저해합니다. 또한, 표준 확산 과정은 "상태 충돌"로 고통받을 수 있는데, 이는 서로 다른 분자가 순방향 과정에서 공통 상태로 붕괴하여 모드 붕괴를 초래하고 다양성을 감소시킵니다.

2. 방법론: VQ-SAD

저자들은 VQ-VAE(벡터 양자화 변분 오토인코더) 와 구조 인식 확산 (SAD) 을 통합하는 신경 - 상징적 (neuro-symbolic) 프레임워크인 VQ-SAD를 제안합니다. 방법론은 두 가지 주요 구성 요소로 이루어집니다:

A. 신경 - 상징적 토큰화 (VQ-VAE)

VQ-SAD 는 원자와 결합을 원시 원-핫 벡터로 취급하는 대신, 사전 훈련된 VQ-VAE 를 토크나이저로 사용합니다.

• 과정: VQ-VAE 는 먼저 연속적인 노드 및 엣지 표현을 이산 잠재 공간 (코드북) 으로 매핑하도록 훈련됩니다.
• 이점: 이는 반복되는 구조적 패턴의 "학습된 어휘"를 생성합니다. 코드북의 각 코드는 상징적 엔티티에 해당하여, 명시적 열거 없이도 화학적으로 다른 맥락 (예: 서로 다른 탄소 환경) 을 구별할 수 있게 합니다.
• 동결: 훈련이 완료되면 VQ-VAE 인코더와 디코더는 동결됩니다. 이산 코드는 하류 확산 모델의 입력 토큰으로 사용되어 원자 및 결합 유형의 분포를 효과적으로 균형 잡습니다.

B. 구조 인식 확산 (SAD)

확산 과정 자체는 "구조 인식"적이고 "학습 가능"하도록 강화됩니다.

• 구조 임베딩: 모델은 노드와 엣지의 구조적 맥락을 인코딩하기 위해 **상대적 랜덤 워크 확률 (RRWP)**을 활용합니다. 단순한 차수 (degree) 카운트와 달리, RRWP 는 복잡한 위상적 특징 (예: 사이클, 거리) 을 포착합니다.
• 적응형 노이즈 스케줄링: 정적 확산 스케줄과 달리, VQ-SAD 는 순방향 과정 매개변수 (노이즈 분산 $\beta_t$ 및 교체 확률 $\gamma_t$) 를 동적으로 학습합니다.
  • 학습 가능한 네트워크는 구조 임베딩 (RRWP), 노드/엣지 유형, 시간 단계를 입력으로 받습니다.
  • 각 노드와 엣지에 대한 특정 노이즈 스케줄을 출력합니다.
  • 이는 구조적으로 다른 노드가 서로 다른 노이즈 스케줄을 받아 확산 과정에서 고유한 정체성을 유지하도록 보장함으로써 "상태 충돌"을 방지합니다.
• 이산 전이: 순방향 과정은 세 가지 확률을 포함하는 학습 가능한 마르코프 커널을 사용합니다:
  1. 유지 ($\alpha_t$): 노드가 변경되지 않음.
  2. 마스크 ($\beta_t$): 노드가 마스크됨.
  3. 교체 ($\gamma_t$): 노드가 다른 유형으로 교체됨 (마스크되지 않은 오류를 수정하는 데 중요함).

C. 훈련 및 생성

훈련은 두 단계 파이프라인을 따릅니다:

1. 단계 1: 분자 데이터셋에서 VQ-VAE 토크나이저를 훈련합니다.
2. 단계 2: 이산 잠재 공간에서 확산 모델 (디노이저) 을 훈련합니다. 디노이저는 구조 임베딩을 조건으로 하여 노이즈가 있는 입력에서 원래 이산 코드를 예측합니다.

• 손실 함수: 모델은 학습된 노이즈 스케줄로 가중치가 부여된 음수 증거 하한 (NELBO) 에서 유도된 노이즈 조건부 손실을 최적화합니다.

3. 주요 기여

1. 신경 - 상징적 패러다임: 원자와 결합 코드를 VQ-VAE 의 잠재 변수로 취급하는 새로운 접근법을 도입하여, 상징적 화학 (이산 규칙) 과 신경 구조적 정보를 연결합니다.
2. 적응형 구조 인식 스케줄링: 노이즈 스케줄링이 시간 단계나 노드 유형뿐만 아니라 구조적 특징 (RRWP) 에 조건부로 결정되는 학습 가능한 순방향 과정을 개발합니다. 이는 상태 충돌을 완화하고 소수 클래스의 재구성을 개선합니다.
3. 균형 잡힌 생성: 이산 코드북을 사용하여 원자 및 결합 유형 분포를 자연스럽게 균형 있게 만들어, 생성된 그래프의 품질과 다양성을 높입니다.
4. 최첨단 성능: 기존 확산 모델에 비해 표준 벤치마크에서 우수한 성능을 입증했습니다.

4. 실험 결과

모델은 두 가지 표준 데이터셋인 QM9(소형 유기 분자) 와 ZINC250k(약물 유사 분자) 에서 평가되었습니다.

• 무조건부 생성:
  • 유효성: VQ-SAD 는 DiGress, MELD, 토큰화되지 않은 SAD 와 같은 베이스라인을 능가하는 가장 높은 유효성 점수 (QM9 에서 97.31%, ZINC250k 에서 93.84%) 를 달성했습니다.
  • 고유성: QM9 에서 98.51% 로 가장 구조적으로 구별되는 분자를 생성했습니다.
  • 분포 유사성: 더 낮은 프라체트 체넷 거리 (FCD) 와 NSPDK 점수를 달성하여 생성된 분자가 실제 데이터 분포에 통계적으로 더 가깝음을 나타냈습니다.
• 충돌률: VQ-SAD 는 MELD 에 비해 "충돌률"(서로 다른 분자가 동일한 표현으로 매핑되는 경우) 을 QM9 에서 0.35 에서 0.21 로 크게 감소시켜, 맥락적 차이를 보존하는 능력을 입증했습니다.
• 조건부 생성: 특성 유도 생성 (예: 열용량, 쌍극자 모멘트) 에서 VQ-SAD 는 유효성과 고유성이 향상되었으나, 저자들은 모델이 고확률 모드로 붕괴하는 경향이 있어 고도로 제약된 조건부 작업에서 다양성이 감소한다고 지적했습니다.

5. 의의

VQ-SAD 는 순수 신경 표현에서 신경 - 상징적 접근법으로 전환함으로써 분자 생성 모델링에서 중요한 변화를 나타냅니다.

• 화학적 유효성: VQ-VAE 토크나이저를 통해 상징적 제약을 존중함으로써, 모델은 화학적으로 유효한 분자를 더 신뢰성 있게 생성합니다.
• 확장성 및 견고성: 적응형 스케줄링 메커니즘은 "상태 충돌" 문제를 해결하여, 불균형한 클래스 분포를 가진 복잡한 그래프에 대해 확산 모델을 더 견고하게 만듭니다.
• 향후 방향: 이 연구는 이산적이고 해석 가능한 상징적 표현을 연속 확산 과정과 통합하는 것이 신약 개발 및 재료 과학에 유망한 길임을 시사하며, 원시 원-핫 데이터만으로는 학습하기 어려운 복잡한 화학 규칙을 인코딩할 수 있는 방법을 제공합니다.