이 논문은 분자 특성 예측이라는 중요한 과제를 해결하기 위한 새로운 분자 표현 학습 방법인 ChemRL-GEM을 제안합니다. 분자 특성 예측은 신약 및 재료 산업에서 분자를 평가, 선택, 생성하는 데 매우 중요한 근본적인 작업입니다. 최근 그래프 신경망(GNN)이 분자 표현 학습에 유망한 성능을 보였으며, 레이블이 부족한 분자 문제를 극복하기 위해 자체 지도 학습 방법을 활용한 사전 학습도 성공적으로 적용되었습니다.
문제점: 기존 GNN 및 사전 학습 전략은 분자를 원자와 결합의 위상학적 그래프 데이터로만 취급하며, 분자의 3차원(3D) 공간 구조, 즉 분자 기하 정보를 충분히 활용하지 못했습니다. 하지만 분자 기하 정보는 분자의 물리적, 화학적, 생물학적 특성을 결정하는 데 가장 중요한 요소 중 하나입니다. 예를 들어, 동일한 위상 구조를 가졌지만 기하 구조가 다른 기하 이성질체(Geometric Isomerism)는 기존 그래프 기반 표현 방법으로는 구분할 수 없습니다. Figure 1에 제시된 두 분자(cis-1,2-DCE 및 trans-1,2-DCE)는 동일한 위상 구조를 가지지만, 염소 원자의 위치가 다르기 때문에 물에 대한 용해도와 같은 특성이 다릅니다. 또 다른 예로, 동일한 위상을 가진 cis-플라틴과 trans-플라틴 중 cis-플라틴만 항암 활성이 있습니다. 또한 기존의 자체 지도 학습 방법은 주로 위상학적 특징에 기반한 마스크 및 예측 작업에 집중하여, "어떤 원자/그룹이 이중 결합에 연결될 수 있는가"와 같은 표면적인 화학 법칙 학습에는 도움이 되지만, Figure 1의 두 분자 간의 결합 각도 차이와 같은 기하 지식을 학습하는 데는 부족합니다.
제안 방법 (ChemRL-GEM): 이 논문은 이러한 문제점을 해결하기 위해 다음과 같은 ChemRL-GEM을 제안합니다:
- 기하 기반 GNN (GeoGNN): 메시지 전달이 기하 정보에 민감하도록 설계된 GeoGNN 아키텍처를 제안합니다. GeoGNN은 원자, 결합, 결합 각도를 동시에 모델링합니다. 이를 위해 분자에 대해 두 개의 그래프를 사용합니다:
- 원자-결합 그래프 (Atom-bond graph): 원자를 노드로, 결합을 엣지로 사용하여 원자와 결합 간의 관계를 나타냅니다.
- 결합-각도 그래프 (Bond-angle graph): 결합을 노드로, 결합 각도를 엣지로 사용하여 결합과 결합 각도 간의 관계를 나타냅니다. GeoGNN은 이 두 그래프를 연결하기 위해 결합 표현 벡터를 사용합니다. 결합-각도 그래프에서 이웃 결합 및 해당 결합 각도로부터 메시지를 집계하여 결합 표현 벡터를 학습합니다. 그 다음, 학습된 결합 표현 벡터는 원자-결합 그래프의 엣지 특징으로 사용되어 원자 표현 벡터를 학습하는 데 도움을 줍니다. 최종적으로 원자 표현 벡터를 통합하여 분자 표현 벡터를 얻고 이를 특성 예측에 사용합니다.
- 기하 수준 자체 지도 학습 전략: GeoGNN의 일반화 능력을 향상시키기 위해 다양한 기하 수준 자체 지도 학습 작업을 설계하여 공간 지식을 학습하도록 사전 학습합니다. 이러한 작업은 다음과 같습니다:
- 결합 길이 예측: 두 연결된 원자 사이의 거리인 결합 길이를 예측합니다. 이는 원자 간의 결합 강도를 반영하며 국부적인 공간 구조를 나타냅니다.
- 결합 각도 예측: 두 연속적인 결합(세 원자 포함)을 연결하는 각도인 결합 각도를 예측합니다. 이는 분자의 국부적인 공간 구조를 나타냅니다.
- 원자 간 거리 행렬 예측: 원자의 3D 좌표를 기반으로 각 분자에 대한 원자 간 거리 행렬을 구성하고, 행렬의 요소를 예측합니다. 위상 구조가 동일하더라도 원자 간 거리는 크게 다를 수 있으므로, 이 작업은 회귀 문제가 아닌 다중 클래스 분류 문제로 취급하여 원자 간 거리를 이산화합니다. 이 작업은 분자의 전역적인 공간 구조를 학습합니다. 이러한 기하 수준 작업은 국부 공간 구조(결합 길이, 결합 각도)와 전역 공간 구조(원자 간 거리 행렬)를 모두 고려합니다. 사전 학습 과정에서는 이러한 기하 수준 작업의 손실 함수를 합산하여 사용합니다.
실험 결과: ChemRL-GEM의 효과를 검증하기 위해 MoleculeNet의 12개 벤치마크 데이터셋(회귀 및 분류)에서 다양한 최신 방법들과 비교 실험을 수행했습니다. 사전 학습에는 Zinc15에서 샘플링된 2000만 개의 레이블 없는 분자를 사용했습니다. 3D 좌표는 RDKit 패키지의 Merck 분자력장 함수를 사용하여 시뮬레이션했습니다. 실험 결과 ChemRL-GEM이 12개 데이터셋 중 11개에서 최신 결과(SOTA)를 달성하며 모든 기준 모델보다 우수한 성능을 보였습니다.
- 회귀 작업: ESOL, FreeSolv, Lipo, QM7, QM8, QM9 데이터셋에서 평가되었습니다. ChemRL-GEM은 각 데이터셋의 이전 SOTA 결과에 비해 평균적으로 8.8%의 전체적인 상대적 개선을 달성했습니다.
- 분류 작업: BACE, BBBP, ClinTox, SIDER, Tox21, ToxCast 데이터셋에서 평가되었으며, 평균 ROC-AUC로 평가되었습니다. ChemRL-GEM은 이전 SOTA 결과(D-MPNN)에 비해 평균 ROC-AUC에서 3.7%의 전체적인 상대적 개선을 달성했습니다. 회귀 데이터셋은 물 용해도 예측(ESOL)이나 전자 특성 예측(QM7)과 같이 분자 기하 정보와 더 밀접하게 관련된 특성을 다루기 때문에, 분류 작업에 비해 ChemRL-GEM이 더 상당한 개선을 이루었습니다. ClinTox 데이터셋에서는 SOTA 결과를 달성하지 못했는데, 이는 해당 데이터셋의 심각한 클래스 불균형(테스트 세트에 긍정 샘플이 9개뿐) 때문일 수 있습니다. 다른 분할 방법을 사용한 실험에서도 ChemRL-GEM은 SOTA를 달성했습니다.
Ablation 연구: GeoGNN의 기여도를 평가하기 위해 다양한 GNN 아키텍처와 비교한 결과, GeoGNN이 모든 회귀 데이터셋에서 다른 GNN 아키텍처보다 우수한 성능을 보였습니다. 이는 3D 좌표가 시뮬레이션된 경우에도 기하 매개변수를 통합하는 GeoGNN의 구조 덕분입니다. QM9 데이터셋에서 더 정확한 3D 좌표를 사용했을 때 GeoGNN의 성능이 더욱 향상되었는데, 이는 GeoGNN이 더 정확한 3D 좌표로부터 분자 표현을 학습하는 능력이 증가함을 보여줍니다. 제안된 기하 수준 자체 지도 학습 작업의 효과를 연구한 결과, 기하 수준 작업이 포함된 방법들이 그렇지 않은 방법들보다 더 나은 성능을 보였습니다. 특히, 기하 수준 작업만 사용한 경우(Geometry)가 기하 수준 작업과 그래프 수준 작업을 함께 사용한 경우(Geometry+Graph)보다 회귀 작업에서 더 나은 성능을 보였는데, 이는 분자 지문(그래프 수준 작업에 사용됨)과 회귀 작업 간의 약한 연결성 때문일 수 있습니다.
결론 및 향후 연구: 이 논문은 기존 GNN 및 사전 학습 방법이 분자 기하 정보를 충분히 활용하지 못한다는 점을 지적하며, 이를 해결하기 위해 기하 기반 GNN인 GeoGNN과 기하 수준 자체 지도 학습 방법을 제안했습니다. GeoGNN은 원자-결합 그래프와 결합-각도 그래프를 사용하여 결합 각도 정보를 활용하고, 기하 수준 자체 지도 학습 작업은 분자 기하 정보(결합 길이, 결합 각도, 원자 간 거리)를 예측하여 공간 지식을 포착합니다. 광범위한 실험을 통해 ChemRL-GEM이 12개 벤치마크에서 기존 방법들보다 상당히 우수한 성능을 보임을 입증했습니다. 향후 연구에서는 토션 각도와 같은 다른 기하 매개변수를 통합하여 분자 표현 능력을 더욱 향상시키고, 분자 표현 학습을 약물-표적 상호작용 예측 및 분자 간 상호작용과 같은 다른 분자 관련 문제에 적용하는 연구를 진행할 계획입니다.