MolFM-Lite: Multi-Modal Molecular Property Prediction with Conformer Ensemble Attention and Cross-Modal Fusion

이 논문은 SELFIES 시퀀스, 분자 그래프, 그리고 볼츠만 가중치를 적용한 컨포머 앙상블을 교차 어텐션과 FiLM 을 통해 융합하여 분자 특성 예측 성능을 획기적으로 향상시킨 경량 다중 모달 모델 'MolFM-Lite'를 제안하고, MoleculeNet 벤치마크에서 단일 모달 기반 모델 대비 우수한 성능을 입증합니다.

핵심

이 논문은 MolFM-Lite이라는 새로운 인공지능 모델을 소개합니다. 이 모델은 약을 개발할 때 가장 중요한 '분자의 성질'을 예측하는 일을 훨씬 더 정확하게 해내도록 도와줍니다.

기존의 방법들이 분자를 볼 때 '한 가지 눈'만 썼다면, MolFM-Lite 는 '세 가지 눈'을 동시에 뜨고 분자를 바라본다는 점이 핵심입니다.

이 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

🧪 1. 분자를 보는 세 가지 눈 (멀티모달)

약학 연구자들은 분자를 분석할 때 보통 세 가지 방식으로 정보를 얻습니다. 하지만 기존 AI 는 이 중 하나만 선택해서 사용했죠. MolFM-Lite 는 이 세 가지를 모두 합칩니다.

1. 1 차원 (SELFIES): 분자의 문자열입니다. 마치 분자의 이름이나 주소처럼 "C-C-O-N"처럼 글자로 된 정보예요. (예: 분자의 구성 성분)
2. 2 차원 (그래프): 분자의 연결도입니다. 누가 누구와 손잡고 있는지, 어떤 고리를 이루고 있는지 보여주는 지도 같은 거예요. (예: 분자의 구조)
3. 3 차원 (입체 구조): 분자의 실제 모양입니다. 분자는 종이 위가 아니라 3D 공간에서 구부러지고 뒤틀립니다. 이 모양에 따라 약이 몸속 수용체에 잘 끼울지 결정되죠. (예: 분자의 실제 형태)

🍕 비유:
마치 피자를 주문할 때 생각해보세요.

• 1 차원: 메뉴판에 적힌 "페페로니 피자"라는 이름만 보고 주문하는 것.
• 2 차원: 피자가 어떻게 잘려 있는지, 토핑이 어디에 있는지 보는 것.
• 3 차원: 피자가 오븐에서 어떻게 부풀어 오르고, 치즈가 어떻게 녹아내리는지 실제 모양을 보는 것.

기존 AI 는 이름이나 연결도만 보고 맛을 예측했지만, MolFM-Lite 는 이 모든 정보를 한눈에 보고 "아, 이 피자는 정말 맛있겠구나!"라고 예측합니다.

---

🤝 2. 세 눈이 대화하는 방식 (크로스-어텐션)

단순히 세 가지 정보를 나란히 쌓아두는 게 아니라, 서로 대화하게 만듭니다.

• 기존 방식: "나는 이름만 봤어", "나는 구조만 봤어", "나는 모양만 봤어"라고 각자 따로 말하고 합치는 거예요.
• MolFM-Lite 방식: "이름을 보니 이 구조가 더 중요할 것 같아", "이 모양을 보니 이 이름이 더 의미 있네"라고 서로 정보를 주고받으며 결론을 내립니다.

👥 비유:
세 명의 전문가가 회의를 한다고 상상해보세요.

• 언어학자 (1 차원): "이 단어는 보통 이런 뜻이야."
• 지도학자 (2 차원): "그런데 이 지역 지도를 보면 그 뜻이 달라져."
• 건축가 (3 차원): "그리고 건물의 실제 구조를 보면 더 명확해져."

기존 모델은 세 사람이 따로따로 의견을 내서 합쳤다면, MolFM-Lite 는 세 사람이 서로의 의견을 들으며 **"아, 우리가 합치면 이 문제가 해결되네!"**라고 협력하는 것입니다.

---

🔄 3. 분자는 고정된 게 아니야! (Conformer Ensemble)

가장 중요한 발견 중 하나는 **"분자는 딱딱하게 고정된 게 아니라, 여러 가지 모양으로 흔들린다"**는 점입니다.

• 기존 모델: 분자를 딱 한 가지 모양 (가장 에너지가 낮은 상태) 으로만 보았습니다. 마치 사진을 한 장만 찍은 것처럼요.
• MolFM-Lite: 분자가 가질 수 있는 **5 가지 다른 모양 (Conformer)**을 모두 만들어서 봅니다. 그리고 "이 모양이 가장 안정적이니까 60% 확률로 믿고, 저 모양은 20% 확률로 참고하자"라고 물리 법칙 (열역학) 을 적용해서 합칩니다.

🕺 비유:
사람이 춤을 추는 모습을 생각해보세요.

• 기존 모델: 춤추는 사람의 정지된 사진 한 장만 보고 "이 사람은 춤을 잘 춰"라고 판단합니다.
• MolFM-Lite: 춤추는 사람의 동영상 여러 컷을 보고, "이 동작이 가장 자연스럽지만, 저 동작도 가끔 하네"라고 전체적인 흐름을 파악합니다.

이 덕분에 분자가 실제로 몸속에서 어떻게 움직일지 더 정확하게 예측할 수 있습니다.

---

🧪 4. 실험실 환경도 고려한다 (FiLM)

약의 효과는 실험 조건 (온도, 세포 종류 등) 에 따라 달라질 수 있습니다. MolFM-Lite 는 이 환경 정보도 모델에 넣어주도록 설계되었습니다.

• 비유: 같은 요리라도 "매운맛을 좋아하는 사람"에게 먹일지, "싱거운 맛을 좋아하는 사람"에게 먹일지에 따라 레시피를 살짝 바꾸는 것과 같습니다. 이 모델은 실험 조건을 알면 그에 맞춰 예측을 조정할 수 있습니다.

---

🏆 5. 결과는 어떨까? (성공 사례)

이 모델은 유명한 약물 개발 데이터셋 (BBBP, BACE 등) 에서 기존 최고의 모델들보다 7~11% 더 높은 정확도를 기록했습니다.

• 비용: 놀랍게도 이 모든 실험을 하는 데 든 비용은 약 47 달러 (약 6 만 원) 정도였습니다. 거대 기업들이 수백만 달러를 쓰는 대형 모델과 달리, 일반 연구실에서도 쉽게 따라 할 수 있는 '가벼운' 모델입니다.

💡 요약

MolFM-Lite는 분자를 이해할 때 "한 가지 정보만 믿지 말고, 문자, 구조, 3D 모양을 모두 보고 서로 대화하게 하라"는 철학을 가진 모델입니다. 특히 분자가 여러 가지 모양으로 움직인다는 사실을 반영해서, 더 정확하고 저렴하게 새로운 약을 찾아낼 수 있게 도와줍니다.

이 기술은 앞으로 새로운 약을 개발하는 시간을 단축하고, 실패할 확률을 줄여 더 많은 생명을 구하는 데 기여할 것으로 기대됩니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

분자 특성 예측은 컴퓨테이셔널 약물 발견의 핵심 과제이나, 기존 머신러닝 모델들은 다음과 같은 한계를 가지고 있습니다.

• 단일 모달리티 의존성: 대부분의 모델이 분자의 1 차원 서열 (SMILES/SELFIES), 2 차원 그래프, 또는 3 차원 구조 중 하나의 표현 방식만 사용합니다. 각 관점은 상호 보완적인 정보를 담고 있음에도 불구하고 이를 통합하지 못합니다.
• 정적 기하학 가정: 분자는 정적인 구조가 아니라 열역학적 평형 상태의 콘포머 (conformer) 앙상블로 존재합니다. 기존 모델 (Uni-Mol 등 포함) 은 대부분 분자당 단일 콘포머만 사용하여 분자의 유연성과 다양한 형태를 무시합니다.
• 실험 컨텍스트 부재: 측정된 분자 특성은 실험 조건 (assay type, 세포주, 온도 등) 에 크게 의존하지만, 이를 고려하지 않은 모델은 실험 조건이 다른 데이터에 대해 일반화하기 어렵습니다.

2. 제안 방법론: MolFM-Lite (Methodology)

저자들은 MolFM-Lite라는 새로운 멀티모달 모델을 제안하며, 1D, 2D, 3D 정보를 통합하고 실험 컨텍스트를 반영하는 아키텍처를 설계했습니다.

A. 아키텍처 개요

모델은 네 가지 주요 모듈로 구성됩니다:

1. 모달리티별 인코더 (Modality Encoders):
   • 1D (SELFIES): 4 레이어 트랜스포머 (Transformer) 를 사용하여 SELFIES 시퀀스를 인코딩합니다.
   • 2D (Graph): 4 레이어 GIN (Graph Isomorphism Network) 을 사용하여 분자 그래프의 위상 정보를 추출합니다.
   • 3D (Conformers): 경량화된 SchNet (SchNet-Lite) 을 사용하여 원자 좌표를 기반으로 3D 구조를 인코딩합니다.
2. 콘포머 앙상블 어텐션 (Conformer Ensemble Attention):
   • RDKit 의 ETKDG 알고리즘을 통해 분자당 $K=5$개의 콘포머를 생성합니다.
   • 학습 가능한 어텐션 점수와 **볼츠만 가중치 (Boltzmann-weighted prior, MMFF94 힘장 기반)**를 결합하여 각 콘포머의 가중치를 계산합니다. 이는 열역학적으로 유리한 형태를 우선시하면서도 태스크에 따라 학습된 어텐션이 이를 수정할 수 있게 합니다.
3. 크로스-모달 퓨전 (Cross-Modal Fusion):
   • 단순 연결 (Concatenation) 이 아닌 크로스 어텐션 (Cross-Attention) 메커니즘을 사용하여 각 모달리티가 다른 모달리티의 정보를 선택적으로 통합하도록 합니다.
4. 컨텍스트 조건부 (Context Conditioning):
   • 실험 메타데이터 (assay type 등) 를 FiLM (Feature-wise Linear Modulation) 모듈을 통해 모델에 주입합니다. (MoleculeNet 벤치마크에는 메타데이터가 없어 0 벡터로 처리되지만, 아키텍처적 유연성을 제공합니다.)

B. 사전 학습 (Pre-training)

• ZINC250K (약 25 만 개 분자) 에서 30 에포크 동안 사전 학습을 수행합니다.
• 목표 함수:
  • 크로스-모달 대비 학습 (Cross-Modal Contrastive Loss): 동일한 분자의 서로 다른 모달리티 표현을 정렬 (InfoNCE).
  • 마스크된 원자 예측 (Masked Atom Prediction): 2D 그래프에서 원자 유형을 마스킹하고 복원.
• 이 사전 학습은 대규모 모델 (Uni-Mol 등) 에 비해 계산 비용은 낮지만, 멀티모달 파인튜닝의 안정성을 높이는 역할을 합니다.

3. 주요 기여 (Key Contributions)

1. 물리 기반 콘포머 앙상블 어텐션: 학습된 어텐션과 열역학적 볼츠만 분포를 결합하여 분자의 유연성을 효과적으로 모델링합니다.
2. 크로스-모달 퓨전: 1D, 2D, 3D 정보가 상호 보완적으로 작용하도록 크로스 어텐션 레이어를 도입하여, 단순 연결 방식보다 우수한 성능을 달성했습니다.
3. 통제된 평가 프로토콜: 모든 베이스라인 모델 (ChemBERTa, GIN, GROVER, SchNet 등) 을 동일한 스캐폴드 분할 (scaffold split) 과 하이퍼파라미터로 재평가하여 공정한 비교를 가능하게 했습니다.
4. 효율성과 접근성: 약 10M 파라미터 크기의 모델로, 전체 실험 비용 (AWS Spot 인스턴스 기준) 이 약 47 달러에 불과하여 학술 연구실에서도 접근 가능한 비용 효율적인 솔루션을 제시했습니다.

4. 실험 결과 (Results)

MoleculeNet 의 4 가지 벤치마크 (BBBP, BACE, Tox21, Lipophilicity) 에서 스캐폴드 분할 기준으로 평가되었습니다.

• 성능 향상:
  • MolFM-Lite 는 모든 단일 모달리티 베이스라인을 압도적으로 능가했습니다.
  • BBBP: 0.956 AUC (기존 최상위 모델 Uni-Mol 대비 0.916, 단일 모달리티 대비 약 7~11% 향상).
  • BACE: 0.902 AUC.
  • Tox21: 0.848 AUC.
  • Lipophilicity: 0.570 RMSE.
• 애블레이션 연구 (Ablation Study) 결과:
  • 3 모달리티 퓨전: 1D+2D+3D 통합이 가장 강력하며, 단일 모달리티 제거 시 AUC 가 4~11% 감소했습니다.
  • 콘포머 앙상블: 단일 콘포머 ($K=1$) 대비 $K=5$ 사용 시 BBBP 에서 1.8%, BACE 에서 1.5% 정도의 성능 향상을 보였습니다. 특히 볼츠만 사전 확률 (Boltzmann prior) 이 학습된 어텐션의 안정화에 기여했습니다.
  • 크로스 어텐션: 단순 연결 방식보다 2.0~2.7% 더 높은 성능을 기록했습니다.
  • 사전 학습: ZINC250K 기반 사전 학습은 BBBP 에서 3.3% 정도의 성능 향상을 가져왔습니다.

5. 의의 및 결론 (Significance)

• 구조적 다중모달 융합의 유효성: 분자의 1D 서열, 2D 위상, 3D 구조가 서로 다른 정보를 제공하며, 이를 물리 기반 어텐션과 크로스 어텐션을 통해 융합할 때 예측 정확도가 크게 향상됨을 입증했습니다.
• 비용 효율성: 대규모 사전 학습 (수억 개 분자) 없이도, 적절한 아키텍처 설계와 소규모 대비 학습을 통해 최첨단 (SOTA) 성능을 달성할 수 있음을 보여주었습니다.
• 실용성: 불확실성 추정 (MC Dropout) 을 통해 가상 스크리닝에서 신뢰도가 낮은 예측을 식별할 수 있으며, 실험 컨텍스트를 고려할 수 있는 아키텍처적 기반을 마련했습니다.

이 논문은 분자 특성 예측 분야에서 단일 모달리티의 한계를 넘어, 물리 법칙을 반영한 멀티모달 접근법이 계산 비용이 적게 들면서도 높은 성능을 낼 수 있음을 보여주는 중요한 사례입니다.