MolHIT: 계층적 이산 확산 모델을 통한 분자 그래프 생성의 진보

기존의 그래프 기반 AI 모델은 분자의 복잡한 구조를 잘 탐색하지만, 화학적으로 불가능한 구조를 만드는 실수가 잦았습니다. 이 논문은 화학적 지식을 모델 구조에 직접 녹여내어, 사람이 설계한 것과 다름없는 정확한 분자를 생성하면서도 새로운 구조를 찾아내는 능력을 동시에 확보했습니다.

분자 생성에서 기존 그래프 확산 모델은 원자와 결합을 독립적인 범주형 데이터로 취급한다. 하지만 실제 화학에서는 원자의 종류가 같아도 주변 환경(고리 구조 여부 등)에 따라 성질이 완전히 달라진다. 기존 모델은 이러한 '화학적 맥락'을 무시하고 무작위로 노이즈를 섞고 제거하기 때문에, 결과적으로 화학 법칙에 어긋나는 구조를 생성하는 한계가 있었다.

MolHIT은 이를 해결하기 위해 '계층적 구조'를 도입한다. 확산 과정에서 원자를 바로 지우는 대신, 먼저 할로겐이나 지방족 같은 '화학적 그룹'으로 뭉뚱그린 중간 상태를 거치게 한다. 이는 모델이 세부 원자를 결정하기 전에 분자의 대략적인 화학적 정체성을 먼저 학습하도록 유도하는 원리다.

또한 원자 인코딩 방식을 개선하여 질소(N) 하나라도 그것이 고리 안에 있는지, 전하를 띠고 있는지 등을 명시적으로 구분한다. 이를 통해 모델은 모호함 없이 정확한 화학적 역할을 이해하게 되며, 결과적으로 기존 모델들이 어려워하던 복잡한 약물 구조를 매우 정확하게 재현하고 생성할 수 있게 된다.

HDDM(Hierarchical Discrete Diffusion Model)은 상태 공간 T를 깨끗한 상태(S0), 중간 계층 상태(S1), 마스크 상태(S2)로 구분한다. 전방 확산 과정에서 S0의 원자들은 행렬 Φ를 통해 S1의 화학적 그룹으로 전이된 후 최종적으로 마스크 상태가 된다. 이 과정은 Qt = αtI + (βt - αt)Q(1) + (1 - βt)Q(2) 수식으로 정의되는 누적 전이 커널을 통해 계산된다.

DAE(Decoupled Atom Encoding)는 원자 번호에만 의존하던 기존 방식에서 벗어나 방향족성(Aromaticity)과 형식 전하(Formal Charge)를 별도의 속성으로 분리한다. 예를 들어 MOSES 데이터셋의 12개 원자 타입을 4개의 의미론적 그룹으로 묶어 계층 구조를 형성한다. 이는 원자 토큰과 실제 물리적 상태 사이의 일대다 매핑 문제를 해결하여 재구성 성공률을 비약적으로 높인다.

PN-sampler(Project-and-Noise)는 샘플링 시 모델의 예측값 pθ(G0|Gt)를 원-핫 벡터 형태의 깨끗한 매니폴드 M으로 투영한 뒤, 다시 이전 단계로 노이즈를 추가하는 방식을 취한다. 이는 단순한 사후 확률 제약에서 벗어나 생성된 그래프의 구조적 다양성을 극대화하는 역할을 한다.

MOSES 벤치마크에서 MolHIT은 99.1%의 유효성(Validity)을 기록하며 기존 SOTA 모델인 DeFoG(92.8%)를 크게 앞질렀다. 특히 Scaffold Novelty 점수에서 0.39를 기록하여, 높은 유효성을 유지하면서도 훈련 데이터에 없는 새로운 골격 구조를 찾아내는 능력이 탁월함을 입증했다.

GuacaMol 데이터셋 실험에서도 기존 모델들이 전하를 띤 원자 재구성에 거의 실패(1.88% 성공)한 것과 달리, MolHIT은 DAE를 통해 100%의 재구성 성공률을 달성했다. 이는 복잡한 헤테로 고리 화합물이나 쯔비터이온(Zwitterionic) 종을 생성하는 데 필수적인 능력이다.

Ablation Study 결과, DAE와 HDDM, PN-sampler가 결합되었을 때 품질(Quality)과 유효성 지표가 모두 최적화됨이 확인되었다. 특히 온도 샘플링(Temperature Sampling)을 원자 예측에만 적용했을 때 가장 우수한 성능을 보였다.

MolHIT의 아키텍처는 DiGress의 Graph Transformer를 기반으로 하며, 조건부 생성을 위해 적응형 레이어 정규화(adaLN)를 노드 어텐션 레이어에 추가했다. 이를 통해 QED, SA, logP, MW 등 다중 화학적 특성을 가이드로 삼아 분자를 생성할 수 있다.

HDDM의 수학적 기반은 마르코프 체인을 통한 계층적 전이 커널 설계에 있다. 전방 과정의 전이 행렬 Q(1)은 S0에서 S1로의 매핑을, Q(2)는 S0 ∪ S1에서 S2로의 마스킹을 담당한다. 이는 연속 시간 NELBO(Negative ELBO) 유도를 통해 이론적으로 보장된 학습 목표를 제공한다.

원자와 결합의 확산 역학을 분리하여 적용했다. 원자 타입에는 HDDM을 적용하여 화학적 계층 구조를 활용하는 반면, 결합(Edge) 생성에는 균일 전이(Uniform Transition) 커널이 더 효과적임을 실험적으로 확인하여 이를 채택했다.

모델 크기 확장이나 아키텍처 개선에 따른 성능 변화를 아직 테스트하지 못했으며, GuacaMol 데이터셋에서 성능 포화 지점까지 충분히 학습시키지 못한 한계가 있음.