MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models

Motivation

• Known: Discrete diffusion model은 분자 그래프의 범주형 특성과 자연스럽게 부합하지만, 기존 그래프 확산 모델들은 화학적 유효성이 낮고 1D SMILES 모델에 비해 성능이 뒤떨어진다.
• Gap: 기존 그래프 확산 모델은 (1) 원자 범주들 간의 화학적 관계를 무시한 균일한 전이 프로세스를 사용하고, (2) 원자의 형식 전하와 방향성 같은 화학적 역할을 인코딩하지 못해 원자 표현이 불완전하다.
• Why: 약물 발견과 신소재 과학에서 AI 기반 분자 생성의 잠재력을 실현하려면 높은 화학적 유효성과 구조적 혁신성을 동시에 달성해야 하는데, 기존 접근법은 이 두 목표 사이의 트레이드오프에 갇혀 있다.
• Approach: Hierarchical Discrete Diffusion Model을 도입하여 화학적 선행지식을 인코딩한 중간 상태를 추가하고, Decoupled Atom Encoding으로 원자를 형식 전하와 방향성에 따라 분리하여 원자 표현의 정보 손실을 해결한다.

Achievement

Figure 1. MolHIT achieves SOTA result on MOSES dataset.

• SOTA 성능 달성: MOSES 벤치마크에서 99.1%의 거의 완벽한 유효성으로 기존 그래프 확산 모델(DiGress 92.8%, DisCo 90.5% 등)을 큰 폭으로 초과하고 1D 모델도 능가
• 품질-혁신성 파레토 최적: 높은 유효성을 유지하면서도 구조적 혁신성에서도 최고 수준 달성, 기존 방법들의 트레이드오프 극복
• 하위 과제 일반화: 다중 속성 조건부 생성과 scaffold extension 등 실제 적용 과제에서도 최고 성능 달성
• HDDM 수학적 근거: 계층적 확산 모델이 수학적으로 보장된 ELBO를 가지며 이론적으로 정당화됨

How

Figure 2. Overview of MolHIT. (a) Markov chain of Hierarchical Discrete Diffusion Model (HDDM). Clean states (S0) are tr

• Hierarchical Discrete Diffusion Model: 원자를 S0(깨끗한 상태) → S1(중간 상태, 화학적 그룹) → S2(마스크 상태)로 점진적으로 노이즈화하는 3단계 마르코프 체인 설계
• Decoupled Atom Encoding (DAE): 원자 타입을 형식 전하(charged/neutral)와 방향성(aromatic/non-aromatic)에 따라 분리하여 독립적으로 표현, 예를 들어 [nH]같은 특수 원자 타입을 명시적으로 인코딩
• 화학적 선행지식 인코딩: 원자 간의 치환 가능성 등 확립된 화학 관계를 확산 프로세스의 전이 확률에 반영
• Coarse-to-fine 생성: 마스크된 상태에서 시작해 먼저 중간 수준의 화학적 정체성을 결정한 후 구체적인 원자 타입으로 정제하는 방식으로 구조적 일관성 보장

Originality

• HDDM의 계층적 설계: 기존 discrete diffusion의 균일한(uniform) 또는 흡수(absorbing) 전이를 화학적 선행지식을 명시적으로 인코딩한 다단계 구조로 일반화
• DAE 문제 정의 및 해결: 기존 분자 그래프 모델의 원자 인코딩이 '병적(ill-posed)'임을 실증적으로 증명하고(Fig. 3 재구성 실험), 이를 해결하는 간단하면서도 효과적인 방법 제시", '수학적 엄밀성: HDDM이 ELBO 하한이 보장되는 정당한 확산 모델임을 이론적으로 입증
• 통합적 평가: 무조건부 생성(MOSES, GuacaMol), 조건부 생성(scaffold extension, 다중 속성 가이딩) 등 다양한 벤치마크에서 종합적으로 검증

Limitation & Further Study

• 계산 복잡도: 계층적 확산 프로세스가 기존 discrete diffusion 모델보다 추가 연산을 필요로 하는지 여부와 추론 시간에 미치는 영향 미분석
• 더 큰 분자: 평가가 MOSES(평균 ~16개 원자) 규모의 약물형 분자에 집중되었으므로, 더 큰 단백질이나 고분자 같은 대규모 구조에 대한 확장성 미확인
• 화학적 선행지식의 한계: HDDM의 중간 상태와 전이 확률이 현재의 화학적 규칙에 기반하므로, 미발견된 화학 공간을 탐색하는 데 제약이 있을 수 있음
• 후속 연구 방향: (1) 더 큰 분자와 다양한 화학 공간에 대한 확장성 검증, (2) 다른 discrete diffusion 프레임워크와의 통합, (3) 합성 가능성(synthetic accessibility) 등 추가 약물 특성 최적화

Evaluation

총평: MolHIT는 계층적 확산 모델과 원자 인코딩의 재설계를 통해 분자 그래프 생성의 오래된 문제를 우아하게 해결하며, 강력한 실험적 증거와 이론적 정당성으로 뒷받침되는 견고한 기여를 제시한다.