Large-scale chemical language representations
capture molecular structure and properties
https://www.nature.com/articles/s42256-022-00580-7
2022, Nature machine intelligence, 512 citation
약물 셋이 다양하고 적은 경우의 학습효율을 높이기위해 pre-trained 모델을 써보고자 읽게됨.
weight를 공개한 MoLFormer 모델에 대한 내용이다.
머신러닝 기반 모델은 정확하고 신속한 분자 특성 예측을 가능하게 하며, 이는 신약 개발 및 소재 디자인 분야에서 큰 관심을 받고 있습니다. 다양한 지도학습(supervised) 머신러닝 모델들이 유망한 성능을 보여주었으나, 방대한 화학적 공간(chemical space)과 제한적인 특성 레이블(property labels)로 인해 지도학습은 까다로운 과제였습니다. 최근에는 대규모 라벨이 없는 말뭉치(corpus)로 사전 학습된 비지도 트랜스포머(transformer) 기반 언어 모델들이 여러 하위 자연어 처리(NLP) 작업에서 최고 수준(state-of-the-art)의 결과를 만들어냈습니다.
이러한 발전에 영감을 받아, 본 연구에서는 회전식 위치 임베딩(rotary positional embeddings)을 사용하는 효율적인 트랜스포머 인코더 모델인 MoLFormer를 훈련하여 얻은 분자 임베딩을 제안합니다. 이 모델은 PubChem 및 ZINC 데이터셋에서 추출한 11억 개의 라벨이 없는 분자 SMILES 시퀀스를 대상으로, 고도로 분산된 훈련 방식과 선형 어텐션(linear attention) 메커니즘을 결합하여 적용했습니다.
우리가 학습시킨 분자 표현력은 10개의 벤치마크 데이터셋을 활용한 여러 하위 작업에서 지도학습 및 자가지도학습(self-supervised) 그래프 신경망(GNN), 그리고 기존 언어 모델을 포함한 기존 베이스라인 모델들의 성능을 능가함을 보여줍니다. 또한, 다른 2개의 데이터셋에서도 경쟁력 있는 성능을 나타냈습니다.
나아가 어텐션(attention)의 관점에서 추가 분석을 진행한 결과, 화학적 SMILES로 훈련된 MoLFormer가 실제로 분자 내 원자 간의 공간적 관계를 학습한다는 것을 입증했습니다. 이러한 결과는 대규모 분자 언어 모델이 양자화학적 특성을 포함한 다양하고 뚜렷한 분자 특성들을 예측할 수 있을 만큼 충분한 화학적·구조적 정보를 포착할 수 있다는 고무적인 증거를 제공합니다.
1. MoLFormer 프레임워크의 개요 및 특징
- 목적: 대규모 화학 SMILES 데이터를 통해 범용적인 분자 표현(Representation)을 학습하고, 다양한 하위 분자 특성 예측 작업에 활용합니다.
- 학습 방식 (Masked Language Model): SMILES 시퀀스 내의 특정 토큰을 무작위로 마스킹한 뒤 이를 예측하는 자가지도학습(Self-supervision) 방식을 사용하여 문맥을 학습합니다.
- 기술적 개선: 절대적 위치 임베딩 대신 회전식 위치 임베딩(Rotary Positional Embedding)과 선형 어텐션(Linear Attention)을 도입하여, 학습의 안정성을 높이고 수렴 속도를 단축했습니다.
- 모델 성능은 상관없었고, loss가 안정적이고 빠르게 수렴했음.
2. MoLFormer 임베딩 도출 및 Downstream analysis
- 임베딩 추출: 화학 SMILES를 입력받아, 인코더 모델의 마지막 은닉 상태(Last hidden state)들의 평균값을 추출하여 분자 임베딩으로 사용합니다.
- 하위 작업 적용 방식 (두 가지 전략):
- Frozen (고정) 방식: MoLFormer 인코더 가중치를 그대로 고정한 채, 추출된 임베딩을 기반으로 완전 연결(Fully Connected) 모델만 새롭게 훈련합니다.
- Fine-tuned (미세조정) 방식: 하위 작업에 맞춰 인코더 모델의 가중치와 후속 완전 연결 모델의 가중치를 동시에 함께 업데이트합니다.
- Fine-tuned 모델 구조: 인코더 크기(768차원)와 일치하는 2개 층의 완전 연결 신경망을 쌓고, 드롭아웃(0.1) 및 GELU 활성화 함수를 적용하여 최종 출력을 도출
- GELU활성화 함수는 modern transformer구조 공식
*MoLFormer-XL
= PubChem과 ZINC 데이터셋에서 수집한 약 11억 개(≈1.1 billion)에 달하는 분자 데이터 전체를 사용하여 사전 학습
+ 회전식 위치 임베딩(Rotary positional embeddings)과 선형 어텐션(Linear attention) 메커니즘을 적용한 모델임.
흥미로운 점은 " 중간 어텐션 레이어에서 학습된 평균 어텐션 계수(attention coefficients)"를 시각화 했을 때,
중간 회전식 어텐션 레이어에 있는 헤드들의 집합(aggregation of heads)이 공유 결합 패턴과 잘 일치하는 동시에,
분자 내에서 결합되지 않은 원자들 사이의 공간적 관계의 특징(signature)도 포착하고 있음
The linear attention variant captures (encircled in red) the medium-range 3D distance better than does its counterpart.
