제2회 AI 신약개발 경진대회, 우수상 후기

제 2회 AI 신약 개발 경진 대회: 도전

안녕하세요, 히츠 AI 연구팀의 배성한입니다.

지난 7월, 한국제약바이오협회 주관으로 제2회 AI 신약개발 경진대회가 열렸습니다. 작년 1회 대회가 제가 히츠에 합류한 후 맡은 첫 프로젝트였는데요. 그때는 저희 AI 연구팀과 함께 팀을 이루어 참여했었습니다. 팀원들과의 긴밀한 토론과 협력을 통해 순위를 서서히 올려, 700여 개 팀 중 최종 4위를 달성한 좋은 기억이 있습니다.

지난 1년 동안 제가 히츠에서 얼마나 성장했는지 스스로 시험해보고자, 올해는 혼자 대회에 참여하게 되었습니다. 그리고 이번 대회에서도 4위라는 성적으로 우수상을 수상하게 되어, 지난 10월 31일에 시상식이 열리는 AI Pharma Korea 2024 컨퍼런스에 참석하게 되었습니다. 이번 포스팅에서는 제가 이번 대회에서 시도한 여러 AI 기반 신약 개발 기술들에 대해 소개하고, 수상 후기에 대해서도 짤막하게 공유하고자 합니다.

데이터 부족, 어떻게 극복할 수 있을까?

이번 제2회 AI 신약개발 경진대회의 과제는 면역 반응 신호 전달에 관여하는 단백질 카이네이즈(kinase)의 일종인 IRAK4를 대상으로 한 약물의 억제 능력을 나타내는 IC50 값을 예측하는 AI 모델을 구축하는 것이었습니다. 총 1,952종의 약물에 대한 IC50 값만이 학습 데이터로 제공되었는데, 이는 머신러닝 학습에 사용되기에는 매우 적은 양으로, 작년 대회와 마찬가지로 이번 경진대회의 난이도를 높이는 주요 요인이 되었습니다. 적은 데이터로 학습된 AI 모델은 학습에 사용된 데이터의 패턴만 기억하게 되어, 학습 데이터 외의 테스트 케이스들에 대해서는 잘 예측하지 못하는 현상인 '과적합(Over-fitting)'이 일어날 가능성이 크기 때문입니다.

따라서 이번 대회의 핵심은 데이터의 부족을 극복하는 것이었다고 할 수 있습니다. 저는 아래 3가지 주요 전략을 활용하여 데이터 부족을 극복하고자 했습니다.

GNN (Graph Neural Network) 형태의 사전학습 모델인 MolCLR 기반 모델 구축
IC50 라벨의 불균형을 해소하기 위한 LDS (Label Distribution Smoothing)과 FDS (Feature Distribution Smoothing) 기법 활용
IC50 예측에 힌트가 되는 Binding affinity 예측 값 활용

분자 데이터를 학습하다: MolCLR의 활용

먼저 이번 경진대회에서 저의 AI 모델의 Backbone이 되는 MolCLR에 대해 소개하고자 합니다. MolCLR는 방대한 양의 분자 데이터를 그래프 형태로 사전 학습(Pre-training)한 GNN 모델입니다. 사전 학습이란 예측 목표에 대한 학습 이전에, 해당 예측 목표와 연관되어 있는 도메인의 데이터를 AI 모델에 미리 학습시켜 놓는 방법입니다. 사전 학습으로 습득한 해당 도메인에 대한 일반적인 지식은 AI 모델이 상대적으로 적은 데이터로도 과적합이 일어나지 않게 학습되도록 도와주는 효과가 있습니다. 특히 컴퓨터를 통해 얼마든지 생성할 수 있는 이미지 데이터나 텍스트 데이터와는 달리, 실험을 통해 얻어야 하는 과학적인 데이터는 매우 한정적이기 때문에 사전학습 모델의 구축과 활용이 활발히 이루어지고 있습니다.

그 중 대표적인 분자 구조 사전 학습 모델 중 하나인 MolCLR는 세계 최대 화학 데이터베이스인 PubChem에 속한 천만 개 이상의 분자들을 미리 사전 학습하였습니다. MolCLR는 PubChem에서 가져온 진짜 분자와, 원자나 결합 또는 작용기를 무작위로 가리거나 없앰으로써 만든 가짜 분자를 구분하는 대조 학습(Contrastive learning)을 통해 분자 자체가 가지는 구조적, 화학적 특성을 사전 학습한 모델입니다. 이렇게 AI 모델이 습득한 사전 지식들은 학습 데이터의 부족한 구조적, 화학적 다양성을 보충함으로써 학습 과정에서 과적합을 방지해주는 효과가 있습니다.

MolCLR의 분자 사전학습 개요 [출처] Wang, Yuyang, et al. "Molecular contrastive learning of representations via graph neural networks." Nature Machine Intelligence 4.3 (2022): 279-287.

과적합을 방지하기 위해 활용한 또 다른 전략은 바로 학습 데이터의 불균형한 라벨(IC50)의 분포를 완화하는 것입니다. 라벨의 불균형은 학습 데이터 자체의 다양성만큼이나 모델의 과적합에 큰 영향을 미칩니다. 아무리 구조적, 화학적으로 다양한 데이터를 확보한다 하더라도 그 IC50의 분포가 낮은 쪽으로 치우쳐 있다면, 해당 데이터로 학습된 AI 모델은 높은 IC50 값을 예측하지 못할 것입니다. 이러한 라벨의 불균형으로 인한 과적합을 방지하기 위해 모델이 학습 과정에서 보지 못한 라벨을 예측할 수 있도록 보완해주는 LDS & FDS 방법을 도입하였습니다.

LDS와 FDS로 불균형을 넘어서다

LDS와 FDS는 지난 2021년에 발표된 Delving into Deep Imbalanced Regression 논문에서 제안한 딥 러닝 모델 학습 기법으로, 연속적인 라벨 값을 예측하는 회귀(Regression) 문제에서 라벨의 분포가 불균형할 경우 이를 보정함으로써 AI 모델이 편향된 값을 예측하는 것을 방지해주는 방법입니다. 이름 그대로 두 방법 모두 학습 데이터의 라벨 분포에 따라 데이터를 매끈하게(Smoothing) 해주는 방법인데요. 그 대상이 라벨 그 자체인 경우 LDS, 딥러닝 모델이 학습하는 데이터의 feature인 경우 FDS가 됩니다.

LDS와 FDS의 개념에 대해 간략하게 설명드리겠습니다. 불균형한 학습 데이터의 분포를 조정하여 부드럽게 이어지는 분포를 근사한 뒤, 모델이 학습하는 feature와 근사한 라벨의 분포가 유사하다는 가정 하에 이를 보정해주는 방식이라고 할 수 있습니다. LDS와 FDS의 도입으로 라벨 불균형으로 인한 과적합이 완화되어 학습 과정에서의 예측 성능과 테스트에서의 예측 성능의 차이가 크게 감소하였으며, 이는 제가 대회 기간 중 이룬 두 번의 jump-up 중 하나였습니다.

LDS & FDS 방법의 개요 [출처] Yang, Yuzhe, et al. "Delving into deep imbalanced regression." International conference on machine learning. PMLR, 2021.

또 다른 jump-up은 바로 AI 모델에 inductive bias를 도입하는 것이었습니다. Inductive bias는 AI 모델이 특정 태스크에 대해 학습될 때 주어지는, 해당 태스크의 이해에 도움이 될만한 특수한 가정을 뜻합니다. 가장 대표적인 예가 앞서 언급한 MolCLR 모델의 구조인 GNN 입니다. GNN은 노드(node)와 엣지(edge)로 구성된 그래프의 구조를 모사하여 모델의 구조를 구성함으로써 그래프 형태의 데이터 패턴을 잘 학습할 수 있도록 inductive bias가 적용된 딥러닝 모델의 형태입니다. 이러한 inductive bias가 잘 작동하였기 때문에, GNN은 그래프 형태의 분자 데이터 학습에 효과적인 모델로 여겨져 활발히 활용되고 있습니다.

Binding Affinity로 IC50의 실마리를 찾다

이번 경진대회 때 모델에 적용한 inductive bias는 binding affinity 예측 값을 모델이 활용할 feature로 제공하는 것이었습니다. 이번 태스크의 목적인 IC50의 개념에 대해 생각해 봅시다. IC50란 특정 표적 단백질 (이번 태스크에서는 IRAK4)의 활성을 50%로 억제하는데 필요한 약물 농도를 의미합니다. 즉, IC50는 해당 표적 단백질이 속해있는 시스템 (e.g. 세포나 조직) 전체에 약물을 투여하였을 때 반응을 정량적으로 측정한 값으로, 거시 단위에서 약물의 억제 능력의 척도라고 할 수 있습니다.

하지만 좀 더 미시적인 관점으로 보면, 약물이 표적 단백질을 억제하는 기작은 해당 약물 분자가 표적 단백질 내의 특정 위치에 화학적 작용으로 인해 결합함으로써 이루어집니다. 이 때 약물이 단백질에 결합하는 세기를 binding affinity라고 합니다. 즉, 미시세계에서 약물-단백질 간의 결합력인 binding affinity와 거시세계에서 해당 단백질의 억제도인 IC50는 매우 밀접한 관련이 있을 것이라고 추측할 수 있습니다. binding affinity가 IC50 예측 태스크에 힌트가 되는 일종의 inductive bias인 셈입니다. 그러나 Binding affinity 역시도 IC50와 마찬가지로 실험을 통해 측정해야 하는 값입니다. 그렇다고 AI 모델 학습을 위해 1,952종의 약물에 대해 실험을 할 수는 없는 노릇입니다.

Hyper Lab, 도약의 기반이 되다

하지만 저희 HITS에서 제공하는 AI 신약개발 플랫폼 Hyper Lab을 통해서라면 표적 단백질과 후보 약물의 구조만 있다면 AI가 예측해주는 약물-단백질 결합 구조와 Binding affinity를 손쉽게 계산할 수 있습니다. 히츠의 단백질-약물 결합 예측 AI 모델은 실제 미시 세계에서 단백질 및 약물 분자간의 상호작용을 모사한 네트워크 구조와 알고리즘을 통해 설계되었기 때문에 매우 높은 정확도를 보유하고 있습니다.

실제로 저희 Hyper Lab 플랫폼에 내장된 binding affinity 예측 AI 모듈인 PIGNet은 binding affinity 예측 벤치마크인 CASF-2016 scoring에서 Pearson 상관 계수 기준 0.77 내외의 성능을 보여 줍니다. 즉, Hyper Lab을 통해 다량의 약물에 대한 실험 값과 유사한 binding affinity 예측 값을 빠른 시간 안에 얻을 수 있고, 이를 IC50 예측을 위한 feature로 활용할 수 있습니다. 실제로 이번 경진 대회에서도 Hyper Lab에서 얻은 binding affinity 예측값을 feature로 추가함으로써 리더보드 순위에서 두 번째 jump-up이 있었습니다.

PIGNet CASF benchmark 성능 비교 [참조] Moon, Seokhyun, et al. "PIGNet: a physics-informed deep learning model toward generalized drug–target interaction predictions." Chemical Science 13.13 (2022): 3661-3673.

HyperLab을 통해 얻은 IRAK4에 대한 약물 분자의 binding 구조 및 에너지 예측 예시

히츠와 함께 도전과 성장을 이어가겠습니다

지금까지 제가 이번 경진대회에서 활용한 여러 AI 신약 개발 기술들에 대해 소개하였습니다. 작년 경진대회에서는 첫 프로젝트로 참여했다보니 팀원들의 덕이 컸지만, 올해는 지난 1년 동안 배운 경험과 노하우를 활용해 스스로의 힘으로 입상하게 되어 뿌듯했습니다.

또한 훈련 데이터의 한계(데이터의 부족과 라벨 분포의 불균형)와 태스크에 대한 이해(IC50의 개념과 binding affinity와의 관계)를 파악했습니다. 이를 해결할 수 있는 가설을 수립하고, 그에 따른 실험을 진행하여 실제로 리더보드에서 jump-up을 경험한 것은 연구자로서 매우 뜻깊은 경험이었습니다. 뿐만 아니라, 실험을 통해서 얻어야 해서 AI 예측을 위한 feature로 사용하기 어려운 binding affinity를 Hyper Lab을 통해 예측할 수 있는 저희 히츠의 기술력에도 큰 자부심을 느낄 수 있었습니다. 저의 도전과 성장은 여기서 멈추지 않을 것이며, 히츠와 함께 세상을 더 나은 방향으로 이끄는 데 최선을 다하겠습니다.