[논문 리뷰] Learning Invariant Molecular Representation in Latent Discrete Space

Xiang Zhuang, Qiang Zhang외 6명|arXiv (Cornell University)|2023. 10. 22.

Computational Drug Discovery Methods인용 수 5

한 줄 요약

이 논문은 분자의 표현을 불변으로 학습하기 위한 새로운 프레임워크인 iMoLD를 제안한다. 먼저 GNN를 통해 분자를 인코딩한 후 잔차 벡터 양자화(RVQ) 모듈을 사용해 이산 잠재 공간에서 불변 특징을 분리한다. 이 방법은 임의의 작업에 의존하지 않는 자기지도 학습 목표를 활용하여 원인적이고 분포 이동에 강인한 특징을 식별하며, 18개의 분자 데이터셋에서 분포 이동 상황에서도 최신 기술 수준(SOTA)의 OOD 일반화 성능을 달성한다.

ABSTRACT

Molecular representation learning lays the foundation for drug discovery. However, existing methods suffer from poor out-of-distribution (OOD) generalization, particularly when data for training and testing originate from different environments. To address this issue, we propose a new framework for learning molecular representations that exhibit invariance and robustness against distribution shifts. Specifically, we propose a strategy called ``first-encoding-then-separation'' to identify invariant molecule features in the latent space, which deviates from conventional practices. Prior to the separation step, we introduce a residual vector quantization module that mitigates the over-fitting to training data distributions while preserving the expressivity of encoders. Furthermore, we design a task-agnostic self-supervised learning objective to encourage precise invariance identification, which enables our method widely applicable to a variety of tasks, such as regression and multi-label classification. Extensive experiments on 18 real-world molecular datasets demonstrate that our model achieves stronger generalization against state-of-the-art baselines in the presence of various distribution shifts. Our code is available at https://github.com/HICAI-ZJU/iMoLD.

연구 동기 및 목표

• 분자 표현 학습에서 분포 이동에 따른 열악한 OOD 일반화 문제를 해결하기 위해, 특히 다양한 데이터 소스에서 발생하는 분포 이동 상황을 대비한다.
• 기존의 '먼저 분리하고 나서 인코딩'하는 방법들이 분자 구조의 엉키는 특성과 환경에 대한 가정에 의존하는 한계를 극복한다.
• 다양한 분자 작업, 예를 들어 회귀 및 다중 레이블 분류에 적용 가능한 작업에 의존하지 않는 강건한 프레임워크를 개발한다.
• 분포 이동 동안에도 안정적인 유지가 가능한 불변의 원인적 분자 특징을 식별함으로써 일반화 성능을 향상시킨다.
• 작업 특화 지도 없이도 정확한 불변성 식별이 가능한 자기지도 학습 목표를 설계한다.

제안 방법

• GNN이 먼저 전체 분자 그래프를 잠재 표현으로 인코딩한 후 특징 분리를 수행하는 '먼저 인코딩하고 나서 분리' 전략을 제안한다.
• 잠재 공간의 이산화를 통해 과적합을 줄이고 인코더의 표현 능력을 유지하기 위해 잔차 벡터 양자화(RVQ) 모듈을 도입한다.
• 스코어러 GNN를 사용해 이산 잠재 표현에서 불변 구조를 식별하고 추출한다.
• 다양한 환경 간 불변성을 장려하기 위해 작업에 의존하지 않는 자기지도 학습 목표를 설계한다.
• 엔코더 출력이 학습된 코드북과 일치하도록 유도하기 위해 커밋먼트 손실을 사용하여 안정적인 학습을 보장한다.
• 아키텍처 변경 없이도 회귀 및 다중 레이블 분류를 포함한 다양한 후행 작업에 이 프레임워크를 적용한다.

Figure 1: (a) First-Separation-Then-Encoding: the input is separated using a subgraph generator, then each subgraph is encoded respectively. (b): First-Encoding-Then-Separation: the input is encoded, then the representation is separated by a scorer.

실험 결과

연구 질문

• RQ1 불변 분자 특징을 보다 효과적으로 포착하기 위해 '먼저 인코딩하고 나서 분리' 전략이 기존의 '먼저 분리하고 나서 인코딩' 접근 방식을 능가할 수 있는가?
• RQ2 잔차 벡터 양자화(RVQ)는 표현 능력을 훼손하지 않으면서 OOD 일반화를 얼마나 효과적으로 향상시키는가?
• RQ3 제안된 작업에 의존하지 않는 자기지도 학습 목표가 다양한 분자 작업에서 불변성 식별 능력을 얼마나 향상시키는가?
• RQ4 모델은 코바리에이트 이동 및 개념 이동을 포함한 다양한 분포 이동 상황에서 어떻게 성능을 발휘하는가?
• RQ5 각 구성 요소(예: RVQ, 불변 목표)가 최종 OOD 일반화 성능에 기여하는 정도는 어떠한가?

주요 결과

• 제안된 iMoLD 프레임워크는 코바리에이트 이동 및 개념 이동을 포함한 다양한 분포 이동 상황에서 18개의 실제 분자 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성한다.
• GOOD-HIV-Scaffold 데이터셋에서 VQ 모듈을 제거하면 성능이 심각하게 악화되며, 이는 OOD 일반화에서 그 역할이 핵심적임을 입증한다.
• 더 복잡한 GOOD-PCBA 데이터셋에서는 잔차 연결을 제거할 경우 성능 저하가 뚜렷하게 발생하여, 모델 용량 유지에 있어 그 중요성이 확인된다.
• 제거 실험 결과, 불변 학습 목표 $\mathcal{L}_{\textrm{inv}}$ 와 커밋먼트 손실 $\mathcal{L}_{\textrm{cmt}}$ 둘 다 강건한 성능을 확보하는 데 필수적임을 확인한다.
• t-SNE 시각화 결과, iMoLD는 상대적으로 균일한 특징 분포와 더 작은 환경 간 거리(Wasserstein 거리)를 보이며, 더 나은 불변성을 확인한다.
• iMoLD는 높은 검증 점수를 확보하면서도 높은 학습 점수를 기록하여, ERM 기반 베이스라인 대비 과적합이 적고 일반화 능력이 뛰어나다는 것을 나타낸다.

Figure 2: An overview of distribution shifts in molecules. Distribution shifts occur when molecules originate from different scaffold, size or assay environments.