[논문 리뷰] Structured Multi-task Learning for Molecular Property Prediction

Shengchao Liu, Meng Qu외 3명|arXiv (Cornell University)|2022. 02. 22.

Machine Learning in Materials Science인용 수 7

한 줄 요약

이 논문은 구조적 다중작업 학습 프레임워크인 SGNN-EBM를 제안하며, 분자의 성질 예측에 대해 명시적인 작업 관계 그래프를 활용하여 성능을 향상시킨다. 잠재 작업 표현 학습을 위한 상태 그래프 신경망(SGNN)과 구조적 출력 예측을 위한 에너지 기반 모델(EBM)을 결합함으로써, SGNN-EBM는 382개의 작업을 포함하는 새로 구축된 ChEMBL-STRING 데이터셋에서 최신 기술 수준의 성능을 달성한다. 노이즈 대비 추정(NCE)을 통해 효율적인 훈련을 실현하며, 잠재 공간과 출력 공간에서 모두 작업 관계를 공동으로 모델링하여 강력한 베이스라인을 능가한다.

ABSTRACT

Multi-task learning for molecular property prediction is becoming increasingly important in drug discovery. However, in contrast to other domains, the performance of multi-task learning in drug discovery is still not satisfying as the number of labeled data for each task is too limited, which calls for additional data to complement the data scarcity. In this paper, we study multi-task learning for molecular property prediction in a novel setting, where a relation graph between tasks is available. We first construct a dataset (ChEMBL-STRING) including around 400 tasks as well as a task relation graph. Then to better utilize such relation graph, we propose a method called SGNN-EBM to systematically investigate the structured task modeling from two perspectives. (1) In the \emph{latent} space, we model the task representations by applying a state graph neural network (SGNN) on the relation graph. (2) In the \emph{output} space, we employ structured prediction with the energy-based model (EBM), which can be efficiently trained through noise-contrastive estimation (NCE) approach. Empirical results justify the effectiveness of SGNN-EBM. Code is available on https://github.com/chao1224/SGNN-EBM.

연구 동기 및 목표

• 분야 지식에서 유도된 구조적 작업 관계를 활용하여 분자의 성질 예측에서 데이터 부족 문제를 해결하기 위해.
• 명시적인 작업 관계 그래프가 약물 발견 분야의 다중작업 학습 성능에 어떻게 기여할 수 있는지 조사하기 위해.
• 일관된 프레임워크를 개발하여 잠재 공간과 출력 공간 양쪽에서 작업 관계를 모델링함으로써 일반화 능력을 향상시키기 위해.
• 382개의 작업을 포함하고 단백질-단백질 상호작용 기반의 작업 관계 그래프를 갖는 대규모 도메인 특화 다중작업 데이터셋(ChEMBL-STRING)을 구축하기 위해.

제안 방법

• ChEMBL 생물활성 아시이와 STRING 단백질-단백질 상호작용 네트워크를 통합하여 ChEMBL-STRING 데이터셋을 구축하고, 단백질 타겟 기반으로 382개의 이진 분류 작업을 정의한다.
• 작업 관계 그래프를 따라 특징을 전파함으로써 잠재 공간에서 작업 표현을 학습하기 위해 상태 그래프 신경망(SGNN)을 사용한다.
• 에너지 기반 모델(EBM)을 활용하여 다중작업 예측을 구조적 예측으로 공식화하며, 에너지 함수에 SGNN가 학습한 작업 표현을 통합한다.
• 노이즈 대비 추정(NCE)을 통해 EBM를 효율적으로 훈련시켜 실제 레이블 구성과 노이즈 샘플을 구별함으로써 연합 레이블 분포를 최적화한다.
• 두 가지 노이즈 분포 전략을 탐색한다: 고정된 균일 노이즈와 사전 훈련된 SGNN 모델에서 초기화된 적응형 노이즈.
• EBM의 추론을 위해 지속적 샘플링(Gibbs sampling)을 적용하고, SGNN의 층 수에 대한 제거 분석을 통해 장거리 종속성 모델링 능력을 평가한다.

실험 결과

연구 질문

• RQ1 생물학적 지식에서 유도된 명시적인 작업 관계 그래프가 분자의 성질 예측에서 다중작업 학습 성능을 향상시킬 수 있는가?
• RQ2 잠재 공간뿐 아니라 출력 공간에서도 작업 관계를 공동으로 모델링할 경우, 단지 잠재 공간에서만 모델링하는 것과 비교해 성능에 어떤 영향을 미치는가?
• RQ3 EBM의 NCE 기반 훈련에서 다양한 노이즈 분포 선택이 최종 성능에 어떤 영향을 미치는가?
• RQ4 SGNN의 깊이(층 수)가 관계 그래프 내에서 장거리 작업 종속성을 포착하는 데 어떤 영향을 미치는가?
• RQ5 EBM를 활용한 구조적 예측이 각 작업 간 독립 예측에 비해 성능 향상에 얼마나 기여하는가?

주요 결과

• SGNN-EBM는 평가된 모든 데이터셋에서 평균 ROC-AUC 성능이 최고로 기록되었으며, 모든 강력한 MTL 기준선을 뛰어넘었다.
• SGNN 구성 요소만으로도 단일작업 학습 및 표준 MTL 기준선보다 성능 향상을 보였으며, 구조적 작업 표현 학습의 가치를 입증했다.
• SGNN와 EBM를 결합한 SGNN-EBM는 SGNN만 사용할 경우보다 일관되고 뚜렷한 성능 향상을 보였으며, 구조적 출력 모델링의 효과를 입증했다.
• 사전 훈련된 SGNN 모델에서 초기화된 적응형 업데이트 노이즈 분포를 사용할 경우, 고정되거나 균일한 노이즈 분포보다 더 높은 성능을 기록했다.
• SGNN의 층 수를 늘릴수록 성능 향상이 지속되다가 (L=4에서 정점에 도달함) 더 이상 증가하면 성능 저하가 발생했으며, 이는 지속적 샘플링이 장거리 종속성을 포착하는 데 한계가 있음을 시사한다.
• 제거 분석 결과, EBM를 통한 구조적 예측은 모든 사전 훈련된 사전 분포에서 일관되게, 비록 중간 정도의 향상이지만 안정적인 성능 향상을 제공함을 확인하여 이 방법의 강건성을 뒷받침한다.