[논문 리뷰] A Large Encoder-Decoder Family of Foundation Models For Chemical Language

Eduardo Soares, Victor Shirasuna외 4명|arXiv (Cornell University)|2024. 07. 24.

Advanced Computational Techniques and Applications인용 수 4

한 줄 요약

이 논문은 PubChem에서 추려낸 9100만 개의 정제된 SMILES(40억 개의 분자 토큰)를 기반으로 사전 훈련한 대규모 인코더-디코더 기초 모델인 SMI-TED289M을 소개한다. 이 모델은 분자 성질 예측 및 분자 재구성에서 최신 기술 수준의 성능을 달성하였으며, 잠재 공간에서 강력한 조합적 구조를 보이며, 분자 조각 간의 선형 관계를 고려한 소수의 예시를 통한 추론과 생성 작업에서 높은 Tanimoto 유사도를 보였다.

ABSTRACT

Large-scale pre-training methodologies for chemical language models represent a breakthrough in cheminformatics. These methods excel in tasks such as property prediction and molecule generation by learning contextualized representations of input tokens through self-supervised learning on large unlabeled corpora. Typically, this involves pre-training on unlabeled data followed by fine-tuning on specific tasks, reducing dependence on annotated datasets and broadening chemical language representation understanding. This paper introduces a large encoder-decoder chemical foundation models pre-trained on a curated dataset of 91 million SMILES samples sourced from PubChem, which is equivalent to 4 billion of molecular tokens. The proposed foundation model supports different complex tasks, including quantum property prediction, and offer flexibility with two main variants (289M and $8 imes289M$). Our experiments across multiple benchmark datasets validate the capacity of the proposed model in providing state-of-the-art results for different tasks. We also provide a preliminary assessment of the compositionality of the embedding space as a prerequisite for the reasoning tasks. We demonstrate that the produced latent space is separable compared to the state-of-the-art with few-shot learning capabilities.

연구 동기 및 목표

• 대규모 비라벨링 분자 데이터를 기반으로 자기지도 학습을 통한 사전 훈련을 통해 확장 가능하고 오픈소스 기초 모델을 개발한다.
• 인코더-디코더 트랜스포머 아키텍처를 통해 문맥화되고 조합적인 표현을 학습함으로써 분자 성질 예측 및 생성 성능을 향상시킨다.
• 화학적 추론과 소수의 예시에 기반한 일반화 능력의 전제 조건이 되는 학습된 분자 임베딩의 조합적 구조를 조사한다.
• 약물 발견 및 재료 과학 분야의 다양한 과학적 응용에 유연하고 고성능인 모델 패밀리(289M 및 8×289M 파라미터 변종)를 제공한다.
• 향후 화학 기초 모델의 사전 훈련을 위한 9100만 개의 고품질 SMILES로 구성된 정제된 데이터셋을 공개한다.

제안 방법

• PubChem의 9100만 개의 정제된 SMILES를 기반으로, 40억 개의 분자 토큰에 해당하는 대규모 사전 훈련을 수행하며, 자기지도 학습 기반의 마스킹 언어 모델링을 사용한다.
• 공통 어휘를 공유하고 계층적 토크나이제이션을 적용한 이중 인코더 메커니즘을 도입하여 SMILES 문자열을 통해 분자 구조와 결합성을 모델링한다.
• 더 큰 변종(SMI-TED8x289M)에 Mixture-of-Experts(MoE) 아키텍처를 도입하여 희소 활성화를 통해 파라미터 효율적인 확장성을 달성한다.
• 임베딩 투영에 대한 t-SNE 시각화와 선형 회귀 분석을 통해 잠재 공간 내 조합적 구조를 정량화하며, 분자 조각의 덧셈을 나타내는 SMILES 삼중세트를 사용한다.
• 소수의 예시 삼중세트로부터 임베딩을 생성하고 Tanimoto 유사도를 사용하여 목표 SMILES를 재구성함으로써 소수의 예시 추론 성능을 평가한다.
• 표준 벤치마크(QM9, MOSES)와 회귀/분류 작업을 통해 양자역학적, 물리적, 생물물리적, 생리학적 성질에 걸쳐 성능을 검증한다.

실험 결과

연구 질문

• RQ1 대규모 인코더-디코더 기초 모델은 비라벨링 SMILES 데이터를 기반으로 자기지도 학습을 통해 조합적 분자 표현을 학습할 수 있는가?
• RQ2 SMI-TED289M의 학습된 잠재 공간은 소수의 예시 추론과 분자 재구성에 얼마나 잘 기여하는가?
• RQ3 모델의 임베딩 공간은 분자 조각 간의 계층적 및 선형 관계를 어느 정도 반영하는가?
• RQ4 정제된 대규모 SMILES 코퍼스를 기반으로 사전 훈련하면 이전 방법에 비해 분자 성질 예측 성능이 향상되는가?
• RQ5 최소한의 미세조정으로 다양한 화학 작업에 일반화할 수 있는가? 즉, 강력한 소수의 예시 추론 능력을 보이는가?

주요 결과

• SMI-TED289M는 QM9 벤치마크에서 양자역학적 성질 예측의 12개 작업 중 11개에서 최신 기술 수준의 성능을 달성하였다.
• 분자 조각 덧셈 작업에 대한 선형 회귀 분석에서 R²는 0.99, MSE는 0.002를 기록하여 잠재 공간 내 강력한 조합적 구조를 입증하였다.
• 소수의 예시 생성 작업에서 평균 Tanimoto 유사도 점수는 0.52를 기록하였으며, 정확한 SMILES 재구성 결과 중 최고 점수는 0.92에 달했다.
• SMI-TED289M 임베딩의 t-SNE 시각화 결과는 분자 가족 및 계층적 거리에 따라 명확하게 분리되어 있었으며, MoLFormer(R²=0.55, MSE=0.237)를 뛰어넘는 성능을 보였다.
• 소수의 삼중세트만을 입력으로 사용하여도 임베딩 생성에 대해 강력한 소수의 예시 일반화 능력을 보였다.
• PubChem에서 유래한 9100만 개의 SMILES로 구성된 정제된 데이터셋이 사전 훈련에 효과적이며, 모델 성능 향상과 조합적 추론 능력 향상에 기여하는 것으로 입증되었다.