기술 요약: ElemeNet

문제 정의

딥러닝이 분자 특성 예측 분야에 혁신을 가져왔음에도 불구하고, 최첨단(SOTA) 모델들은 여전히 다양한 화학 공간 전반에 걸친 일반화 능력이 부족한 특화된 코드베이스에 국한되어 있습니다. Chemprop와 같은 기존의 범용 도구들은 주로 2D 표현(SMILES 문자열)과 유기 화학에 제한되어 있어, 입체 화학, 유기 금속 화합물의 복잡한 결합, 중원소 또는 생물학적 거대 분자를 네이티브로 지원하지 못합니다. 또한, 이러한 도구들은 3D 등변(equivariant) 아키텍처, 전하 및 스핀 상태에 대한 명시적 조건화, 그리고 견고한 불확실성 정량화(UQ)를 지원하는 경우가 많습니다. 따라서 원소 번호 1~100(수소부터 페르뮴까지) 전체를 지원하고, 수백만 개의 분자로 확장 가능하며, 머신러닝에 대한 깊은 전문 지식 없이도 보정된 불확실성 추정치를 제공할 수 있는 통합 소프트웨어 프레임워크가 절실히 필요합니다.

방법론

저자들은 이러한 격차를 해소하기 위해 설계된 모듈형, 범용 분자 머신러닝 소프트웨어 패키지인 ElemeNet을 소개합니다.

1. 표현 및 입력

• 원소 범위: 원자 번호 1~100(수소에서 페르뮴까지)을 지원합니다.
• 입력 형식: 다양한 2D 및 3D 형식(cif, cml, mol, mol2, sdf, SMILES, xyz)을 수용합니다.
• 그래프 구축:
  • 2D: 위상적 연결성으로부터 유도되며, 좌표는 오직 결합 해석을 위해서만 사용되고 폐기됩니다.
  • 3D: 원자 좌표로부터 구축되며, 에지는 조절 가능한 컷오프 반경에 의해 정의됩니다. 이는 E(3) 등변성을 부과하기 위해 쌍별 거리를 활용합니다.
• 특징 공학(Feature Engineering):
  • 노드 특징: 연속적인 표 형식의 원자 속성(질량, 반경, 전기 음성도 등)과 범주형 특징(족, 주기, 원자가 전자)의 원-핫 인코딩을 포함하는 균형 잡힌 초기 표현을 구성합니다.
  • 에지 특징: 결합 유형, 추정 결합 길이 및 금속 포함 여부를 인코딩합니다.
  • 그래프 특징: 전하 및 스핀 다중도를 전용 신경망을 통해 명시적으로 인코딩하여, 모델이 전자 상태의 함수로서 속성을 학습할 수 있도록 합니다.
  • 유연성: 사용자 제공 특징(예: 부분 전하, Wiberg 지수) 및 암시적/명시적 수소 처리를 지원합니다.

2. 모델 아키텍처

ElemeNet은 모듈형 인코더-리드아웃(encoder-readout) 아키텍처를 채택합니다.

• 인코더:
  • 2D GNN: 방향성 메시지 전달(D-MPNN), 그래프 컨볼루션, 그래프 동형성(graph isomorphism), 그리고 그래프 어텐션 네트워크를 지원합니다.
  • 3D 등변 GNN (EGNN): 불변 스칼라 특징과 3D 좌표를 위한 별도의 메시지 전달 레이어를 사용하며, 구면 조화 함수(spherical harmonics) 없이 E(3) 등변성을 보장하기 위해 가우시안 기저 함수를 통해 원자 간 거리를 투영합니다.
• 리드아웃(Readouts):
  • MLP: 표준 다층 퍼셉트론.
  • Transformer: 그래프 구성 요소의 동적 가중치 부여를 위한 멀티 헤드 셀프 어텐션 메커니즘.
• 하이퍼파라미터 최적화: 네트워크 형태(상수형, 증가형, 감소형, 아워글래스형, 피라미드형), 활성화 함수, 컨볼루션 연산자를 포함한 방대한 설계 공간에 대해 베이지안 최적화를 수행하는 Optuna를 활용합니다.

3. 다중 스케일 학습 태스크

본 프레임워크는 단일 그래프 내에서 여러 스케일의 예측을 지원합니다.

• 그래프 레벨: 분자 속성(예: 총 에너지, 독성).
• 노드 레벨: 원자 중심 속성(예: 부분 전하).
• 에지 레벨: 결합 속성(예: 결합 길이).
• 모이어티(Moiety) 레벨(신규): 노드를 중심으로 하는 특정 서브그래프(k-hop 이웃)에 국한된 속성을 예측하는 새로운 기능으로, 손실 계산 시 나머지 분자 부분을 마스킹합니다. 이는 작용기나 활성 부위를 타겟팅합니다.

4. 불확실성 정량화 (UQ)

ElemeNet은 두 단계의 UQ를 구현합니다.

• 휴리스틱 척도: 잠재 공간 거리(유클리드 및 코사인)와 분류를 위한 정보 엔트로피를 사용합니다. 이는 비용은 낮지만 오차 크기에 대해 공식적으로 보정되지는 않았습니다.
• 앙상블 기반 척도:
  • 회귀: 최종 리드아웃 레이어를 $N$번 복제하되 인코더 트렁크는 공유하는 **얕은 앙상블(shallow ensembles)**을 사용합니다. 모델은 앙상블 평균과 분산으로 매개변수화된 가우시안 분포의 음의 로그 가능도(NLL)를 최소화합니다.
  • 분류: 평균 확률(Mean-of-Probabilities, MoP) 기법을 사용하여 얕은 앙상블을 확장합니다. 활성화 함수(sigmoid/softmax)를 적용한 후 헤드 간의 로짓(logits)을 평균 내며, 확률 공간에서의 앙상블 표준 편차가 불확실성 지표 역할을 합니다.

주요 결과

저자들은 네 가지 다양한 데이터셋에 대해 ElemeNet을 문헌 베이스라인과 비교 벤치마킹했습니다.

1. QM9 (유기 화학):

   • 3D EGNN 모델이 2D GNN보다 성능이 뛰어났으며, 특히 쌍극자 모멘트에서 MAE가 6배 감소하여 기하학적 의존적 속성에 대한 3D 정보의 필요성을 확인했습니다.
   • ElemeNet 모델(2D 및 3D)은 Chemprop v1/v2의 성능과 일치하거나 이를 능가했습니다.
   • 얕은 앙상블은 예측 정확도를 크게 저해하지 않으면서 불확실성 보정(MAE와의 Spearman 순위 상관관계)을 개선했습니다.

2. tmQMg (무기/전이 금속 화학):

   • EGNN 모델이 HOMO-LUMO 갭, 쌍극자, 편극률에서 2D GNN보다 우수한 성능을 보였습니다.
   • 최상의 ElemeNet 모델은 Transformer 리드아웃을 사용했음에도 불구하고, 문헌에서 벤치마킹된 기존 3D 아키텍처들의 평균 성능을 능가했습니다.
   • 멀티태스크 EGNN 모델은 세 가지 타겟 모두에서 경쟁력 있는 성능을 달성했습니다.

3. pydentate (배위 화학):

   • ElemeNet은 배위 원자 및 배위수 예측에 대해 기존 2D 모델과 대등한 정확도를 달성했습니다.
   • Transformer 리드아웃은 배위 원자 예측에 대한 균형 정확도(balanced accuracy)를 개선했습니다(MLP의 92.4% 대비 96.1%).
   • 이전에 보고되지 않았던 에지 레벨(결합 길이) 및 서브그래프 레벨(금속-리간드 결합 길이) 예측에 대한 새로운 벤치마크를 수립했습니다.

4. GEMS (생물학적 화학, 270만 개 분자):

   • 수백만 개의 항목이 포함된 데이터셋에 대한 확장성을 입증했습니다.
   • 기본 하이퍼파라미터로 훈련된 모델은 전자 에너지 및 쌍극자 크기에 대해 경쟁력 있는 MAE를 달성했습니다.
   • 전하를 띤 종을 중성으로 잘못 라벨링하는 것이 예측 오차를 크게 증가시킨다는 점을 통해 명시적 전하 및 스핀 인코딩의 중요성을 보여주었습니다.

의의 및 주장

논문은 ElemeNet이 전체 주기율표에 걸쳐 고성능 딥러닝 모델을 훈련할 수 있는 실용적이고 통합된 플랫폼 역할을 한다고 주장합니다. 주요 의의는 다음과 같습니다.

• 민주화: 간결한 커맨드 라인 인터페이스와 자동화된 하이퍼파라미터 최적화를 통해 비컴퓨터 과학 연구자의 진입 장벽을 낮춥니다.
• 일반화: 복잡한 결합과 전자 상태를 가진 유기 금속, 무기 및 생물학적 시스템을 지원함으로써 유기 화학을 넘어 확장합니다.
• 다중 스케일 능력: 모이어티 레벨 예측을 도입하고 노드, 에지, 그래프 및 서브그래프 타겟을 단일 프레임워크 내에서 지원합니다.
• 견고성: 얕은 앙상블을 통해 잘 보정된 불확실성 추정치를 제공하며, 이는 능동 학습(active learning) 및 신뢰할 수 있는 배포를 위한 필수 전제 조건입니다.
• 성능: 최소한의 사용자 노력과 적절한 계산 예산으로 다양한 화학 영역에서 SOTA 또는 경쟁력 있는 성능을 달성합니다.

저자들은 ElemeNet이 연구자들이 화학 및 물리 과학 전반에 걸친 분자 발견 및 최적화를 위해 현대적인 딥러닝 방법을 활용할 수 있도록 지원한다고 결론짓습니다.