[논문 리뷰] SPMM : Bidirectional generation of structure and properties through a single molecular foundation model

This entry is part 3 of 3 in the series Multimodal Model

[21′ PMLR] Learning Transferable Visual Models From Natural Language Supervision (CLIP)
[22′ NIPS] Flamingo: a Visual Language Model for Few-Shot Learning
[24′ Nature Communications] Bidirectional generation of structure and properties through a single molecular foundation model

1. 들어가며

최근 인공지능과 딥러닝의 발전은 다양한 분야에서 혁신을 일으키고 있습니다. 특히, 화학 분야에서도 이러한 기술의 도입이 활발히 이루어지고 있는데, 분자 생성(Molecule Generation)과 분자 특성 예측(Molecule Property Prediction)과 같은 태스크에서 딥러닝 모델들이 큰 성과를 보이고 있습니다. 이번 포스팅에서는 KAIST에서 발표한 논문 “Bidirectional generation of structure and properties through a single molecular foundation model”을 리뷰하고자 합니다. 이 논문에서는 화학 연구에서 중요한 두 가지 문제를 동시에 해결할 수 있는 SPMM(Structure-Property Multi-Modal foundation model)이라는 새로운 모델을 제안합니다.

기존의 화학 모델들은 주로 단일 모달리티를 활용하여 분자의 구조 또는 특성 중 하나만을 예측하거나 생성하는 데 집중해왔습니다. 그러나 이러한 접근 방식은 두 모달리티 간의 상호작용을 충분히 고려하지 못한다는 한계를 가지고 있습니다. 이번 논문에서는 이러한 문제를 해결하기 위해 멀티모달 학습 방식을 도입하였습니다. 이를 통해 분자 구조와 특성 간의 복잡한 상호작용을 더 잘 이해하고, 이를 기반으로 더 정확하고 유연한 예측 및 생성 모델을 개발할 수 있었습니다.

SPMM 모델은 두 가지 주요 기능을 제공합니다. 첫째, 분자의 구조를 입력으로 받아 해당 분자의 다양한 특성을 예측할 수 있습니다. 둘째, 특정 특성을 입력으로 받아 그 특성에 부합하는 분자 구조를 생성할 수 있습니다. 이러한 양방향 예측 및 생성 능력은 SPMM 모델의 큰 강점 중 하나로, 다양한 화학적 문제를 해결하는 데 매우 유용합니다. 예를 들어, 신약 개발 과정에서 특정 약리학적 특성을 가진 새로운 화합물을 설계하거나, 기존 화합물의 특성을 정확히 예측하는 데 사용할 수 있습니다.

이 논문에서 제안된 모델의 학습 과정도 매우 흥미롭습니다. SPMM은 분자의 구조와 특성을 언어 모델의 형태로 변환하여 처리합니다. 이를 위해 SMILES와 PV(Property Vector)를 각각의 인코더를 통해 임베딩 벡터로 변환하고, 이를 트랜스포머 기반의 퓨전 인코더를 통해 통합합니다. 이러한 접근 방식은 자연어 처리 분야에서 성공적으로 사용되고 있는 멀티모달 학습 기법을 화학 데이터에 적용한 것으로, 모델의 성능을 크게 향상시켰습니다.

논문의 실험 결과 또한 주목할 만합니다. 다양한 조건 하에서 SPMM 모델은 높은 정확도와 유연성을 보여주었으며, 기존의 단일 모달리티 모델들과 비교하여 우수한 성능을 입증하였습니다. 특히, 입력된 특성 값에 따라 유효하면서도 독창적인 분자 구조를 생성하는 능력은 매우 인상적이었습니다. 이러한 결과는 SPMM 모델이 실제 화학 연구 및 산업 응용에서 큰 잠재력을 가지고 있음을 시사합니다.

이제부터 본격적으로 논문에서 제안한 SPMM 모델의 구조와 학습 방법, 그리고 실험 결과를 자세히 살펴보겠습니다. 이를 통해 이번 연구의 혁신성과 실용성을 더욱 깊이 이해할 수 있을 것입니다.

2. 기존 방법의 문제점

먼저 기존 방법의 문제점을 살펴보겠습니다. 이때 기존 방법이란, Molecule Generation 또는 Molecule Property Prediction Task를 위한 방법들을 의미하는데요. 저자들이 주장하는 기존 방법의 문제점은 바로 Unimodality Model만의 사용입니다. 즉 Molecule Structure와 Molecule Property를 각각의 Modality로 나누어 Multimodal Model을 구축하지 않고, Unimodal Model만을 고집하는 방법을 문제로 지적하고 있죠.

한편 Image, Language 도메인에서는 Multimodal Model 연구가 활발하게 진행되고 있습니다. 대표적으로 ChatGPT에 적용되어 있는 GPT-4 모델이 있죠. 뿐만 아니라 LLaVA, Flamingo 등 다양한 Vision Language Model이 발표되었습니다. 이들은 사전 학습된 Image Encoder와 Language Encoder가 있으면, 이 둘을 Align 하도록 학습해주면 훌륭한 Multimodel Model이 만들어짐을 증명했습니다. 이렇게 만들어진 Vision Language Model들은 단일 Modality 입력을 받아, 다른 Modality로 출력할 수 있는 능력을 보여주었죠.

이에 저자들은 화학 분야의 Multimodal Model을 제안합니다. Molecule Structure와 Modelcule Property 데이터를 각각의 Encoder로 학습하는 방법을 사용합니다. 이렇게 학습한 모델은 Molecule Structure를 입력 받으면 해당 Molecule의 Property를 예측할 수 있고, 반대로 Property를 입력 받으면 해당하는 Molecule Structure를 출력해주는 능력을 보여주게 됩니다.

이렇게 구성된 모델을 저자들은 SPMM(molecule Structure Property MultiModal foundation model) 이라고 부릅니다. 앞으로의 글에서는 이번 논문의 제안 모델을 SPMM이라고 부르겠습니다.

3. SPMM Method

이러한 능력을 어떻게 구현했는지 자세한 방법을 살펴보겠습니다. 먼저 모델을 어떻게 구성해서 학습했는지 살펴보고요. 이렇게 학습할때 Molecule Structure와 Molecure Property는 어떻게 Language 형태로 변환했는지 보겠습니다. 마지막으로 이 모델을 학습하기 위해 Loss Function은 어떠한 형태로 구성했는지 살펴보겠습니다.

3-1. Training

먼저 SPMM 모델을 보겠습니다.

SPMM 모델은 위와 같이 구성되어 있는데요. 먼저 크게 보면 PV(Property Vector)를 처리하는 부분과 SMILES(분자 구조) 를 처리하는 부분으로 나뉘어 있는 모양을 볼 수 있습니다. 즉 각 Modality별로 전문적으로 처리하는 구조를 만들고, 위에서 정보를 섞어주는거죠. 이렇게 각 Modality별 정보를 처리하는 부분은 PV encoder/SMILES encoder로 처리되어 있습니다. Language Model과 유사한 컨셉으로 구성했기 때문에 Transformer Encoder라고 봐도 무방합니다.

이제 이렇게 나온 Modality 별 정보를 섞어줘야 하는데요. 이 역할은 Fusion Encoder가 수행합니다. PV/SMILES Encoder에서는 Self Attention 연산을 사용하여 Modality 별 정보의 관계성을 파악했는데요. Fusion Encoder에서는 Cross Attention 연산을 사용하여 서로 다른 Modality에서의 정보를 섞어줍니다. Cross Attention의 상세 연산 과정은 위 그림에 첨부했으니 참고해주세요!

이제 마지막으로 Fusion 까지 완료된 Feature들은 Next Word/Property Prediction을 위한 Regression/Classification Layer를 거치게 됩니다. 이는 마치 ChatGPT 등의 LLM 모델들이 다음 단어를 학습하는것과 동일하다고 할 수 있습니다.

3-2. 입력값 Language 처럼 변환하기

지금까지 설명을 보면 분자구조와 특성 벡터를 마치 언어처럼 처리하는 모습을 볼 수 있는데요. 구조도 마치 LLM을 보는것과 유사하죠. 이렇게 언어 모델과 유사하게 처리하기 위해서는 분자구조와 특성 벡터를 언어와 유사한 구조로 변형해주는 작업이 필요합니다. LLM에서는 이 작업을 Word Embedding이 수행해주는데요. 즉 각각의 단어 Token들을 해당 Embedding으로 변환해주는 작업이 필요합니다.

위 그림은 SMILES와 PV를 Transformer가 처리할 수 있는 Embedding 형태로 변환하는 과정을 보여주고 있습니다.

먼저 SMILES 변환 과정을 살펴볼게요. SMILES는 분자 구조를 ‘영문’ 형태로 변환해놓은 데이터입니다. 예를 들어, 위 그림에서 볼 수 있듯 분자 구조가 ‘c1cnccl’과 같은 ‘영문’으로 표현해놓았죠. 이런 경우 LLM으로 처리하기는 수월한데요. 각각의 ‘영문’을 LLM이 하는 방식 그대로 Word Embedding을 사용하여 Embedding 변환을 수행해주면 됩니다. 물론 각 위치별로 Positional Encoding도 추가해줘야 합니다.

문제는 PV인데요. PV는 이름에서 알 수 있듯 특성값을 저장해놓은 벡터입니다. 각각의 특성은 ‘값’으로 표현될텐데, 어떻게 Embedding으로 변환하면 좋을까요? 저자들은 Linear Layer를 붙여 학습해줍니다. 특성 값이 Linear Layer를 거쳐 고차원의 Embedding으로 매핑되는것이죠. 물론 SMILES 때와 동일하게, Positional Encoding도 추가해줘야 합니다. 이때 재미있는 포인트는 이렇게 만들어진 PV Embedding에 대해 약 50% 정도 랜덤으로 마스킹을 처리해주는데요. 저자들의 주장에 따르면 이는 일종의 Augmentation 역할을 수행한다고 합니다. 그리고 이런 효과보다 더 좋은 점은, Inference 단계에서 PV 값에 ‘Unknown’ 값이 들어와도 동작할 수 있다는 점입니다. 이 PV Encoder는 애초에 50% 정도는 ‘Unknown’ 형태로 입력 받아 학습했기 때문이죠.

이 포인트가 중요한 이유는 실제 사용 시나리오 때문입니다. 아마도 이 모델을 가장 많이 사용하는 시나리오는 특성 값을 입력 받아 분자 구조를 생성해내는 태스크일텐데요. 상상해보면, 이 시나리오에서는 거의 필연적으로 일부 특성값은 ‘Unknown’ 형태로 들어가게 됩니다. 예를 들어, “질량은 1g을 만족해야 하는데 밀도는 너 마음대로 만들어줘” 라고 요청할 수 있으니까요. 따라서 ‘Unknown’ 입력값에 대응할 수 있도록 구성한 부분은 매우 중요하다고 할 수 있습니다.

3-3. Loss Functions

이제 이렇게 구성된 모델을 학습하기 위한 Loss Function들을 살펴보겠습니다. 위의 SPMM 모델 그림에서 잘 표현되어 있었는데요. 정리해보면 다음과 같습니다.

먼저 Modality 별 출력된 [CLS] 토큰에 해당하는 Feature는 Contrastive Loss를 적용해줍니다. 이는 각 Modality별 Encoder를 Align 해주고 Feature Space를 맞춰주는 효과가 있습니다.

크게 두가지 Similarity를 측정하는데요. 먼저 서로 다른 Modality끼리 측정하는 Inter Modal Similarity입니다. 나와 매칭 되는 데이터의 서로 다른 Modality Feature가 가까워지도록 당겨주는 역할을 합니다.

반대로 Modality끼리 측정하는 Intra Modal Similarity가 있습니다. 나 이외의 다른 Feature는 멀어지도록 밀어주는 역할을 합니다.

위 수식들에서 n과 m은 각 Modality별 배치 사이즈를 의미합니다.

두 번째로 Next Word/Property Prediction Loss가 있습니다. 이는 각 Modality별 관계를 잘 학습하도록 유도하는 역할을 하는데요. ChatGPT가 다음 단어를 예측하도록 학습하는것과 유사하다고 볼 수 있습니다. 이때 SMILES는 Next Word를 예측해야 하므로 Cross Entropy로 학습하고요. PV는 Next Property Value를 예측해야 하므로 MSE로 학습합니다.

세 번째로 SMILES PV Matching Loss가 있습니다. SMILES와 PV Feature에 대해 서로 매칭되면 1로, 아니면 0으로 Classification을 학습하는데요. 구체적으로는 Fusion Encoder Output을 Concat하여 Classification 해주면 됩니다. Minor한 역할을 할 것으로 보이는데, 어쨌든 두개의 Encoder가 서로 Align 되는데 일조해줍니다.

3-4. Inference

이렇게 학습한 SPMM이 어떻게 Inference를 수행할지 궁금한데요. 크게 두 가지 Task를 수행할겁니다.

첫 번째는 Property를 입력받아 Molecule Structure를 Generation 해줄겁니다. 거의 Main Task라고 생각되는데요. 동작 방법은 아래와 같습니다.

우선 학습 완료된 PV Encoder와 SMILES Encoder를 모두 사용해줄겁니다. 그리고 Fusion Encoder부터 상위 모든 Layer는 SMILES 것만 사용해줍니다. 입력은 PV로 받고, 최종 출력은 SMILES에서 해줘야 하니까요. 이제 입력으로는 순차적으로 PV값이 들어가고, SMILES와는 Fusion Encoder에서 Cross Attention으로 섞일겁니다. 이제 SMILES Encoder에 [CLS] 토큰이 입력 되면서 SMILES Output이 출력되는데요. 이렇게 출력되는 Output은 Autoregressive하게 다시 SMILES Encoder 입력으로 들어가줍니다. 마치 LLM이 서로 다른 언어를 번역하는 과정과 유사하게 돌아간다고 보시면 됩니다.

두 번째는 반대로 SMILES를 입력받아 Property를 Prediction하는 문제입니다. 동작 방법은 아래와 같습니다.

SMILES Generation 방법과 정 반대로 구성된 모습을 볼 수 있습니다. SMILES Encoder로부터는 SMILES 입력까지만 처리해주고요, 나머지 모든 부분은 PV에서 처리해줍니다. 출력도 PV쪽에서 나오게 되고, 마찬가지로 Autoregressive하게 출력된 PV값은 다시 PV Encoder 입력으로 들어가는 모습을 볼 수 있습니다.

실용적인 부분은, 이렇게 두가지 메인 Task를 수행하기 위해 따로 Downstream Task를 학습할 필요가 없다는 점입니다. 애초에 학습한 내용이 SMILES Generation과 PV Property였기 때문에, Inference 단계에서는 구조만 살짝 변형해주면 원하는대로 생성 또는 예측을 풀어낼 수 있다는점이 장점입니다.

4. Experiment

이제 앞서 설명한 두 가지 Task에 대한 실험 결과를 살펴보겠습니다.

4-1. SMILES Generation

먼저 PV를 입력 받아 SMILES를 생성하는 실험 결과를 살펴보겠습니다. 결과를 살펴보기 전에 평가 Metric을 정리해야되는데요. 크게 세 가지 Metric이 있습니다.

Validity는 생성한 SMILES가 얼마나 ‘진짜 같은지’를 의미합니다. 생성된 SMILES 중 Valid한 비율로 평가합니다
Uniqueness는 생성한 SMILES가 얼마나 겹치지 않는지 평가합니다. Valid한 생성 SMILES 중 겹치지 않는 비율로 평가합니다.
Novelty는 생성한 SMILES가 얼마나 학습 데이터에 없는지를 평가합니다. Unique SMILES 중 학습 데이터에 없는 비율로 평가합니다.

위 표는 이렇게 세 가지 Metric으로 평가한 결과를 보여줍니다.
첫 번째 행은 PubCehm 데이터셋에서 1000개의 분자를 선택하여, 이들의 PV값을 입력해주고 생성한 결과를 평가한 결과입니다.
두 번째 행은 특정 분자의 PV를 입력하여 1000개 SMILES를 생성한 결과입니다.
세 번째 행은 PV중 질량만 입력하여 1000개 SMILES를 생성한 결과입니다.
네 번째 행은 네 가지 특성만 입력하여 1000개 SMILES를 생성한 결과입니다.
다섯 번째 행은 아무 조건 없이 1000개 SMILES를 생성한 결과입니다.

결과 값을 봤을때 대체로 0.9 이상을 보여, 다 잘했다라고 보여지는데요. 조금 더 자세히 보겠습니다.

위 그래프는 각 시나리오 별 생성된 SMILES 특성을 분석한 그래프입니다.

먼저 a는 53개의 모든 PV값을 입력했을때 생성한 SMILES 결과를 보여주고 있는데요. 생성된 SMILES는 모두 입력된 특성값에서 크게 벗어나지 않는 수준으로 잘 생성해준 모습을 볼 수 있습니다.
b에서는 질량값만 입력해준 시나리오인데요. 이렇게 생성된 1000개의 SMILES들은 질량값에 대해서는 입력값에 맞추어 잘 생성했고, 나머지 특성들은 random하게 잘 분포되도록 생성해주는 모습을 볼 수 있습니다.
c에서는 네개 PV 값을 입력해준 시나리오인데요. 마찬가지로 입력해준 특성은 최대한 만족하며, 나머지 특성들에서는 Random하게 잘 분포되는 모습을 볼 수 있습니다.
마지막으로 d에서는 모든 PV를 ‘Unknown’으로 입력해준 시나리오인데요. 모든 특성들이 잘 Random하게 분포하도록 생성된 모습을 볼 수 있습니다.

여기까지 봤을때 SPMM은 다양한 조건의 PV를 입력 받을 수 있고, 유효하면서도 독창적인 SMILES를 잘 생성한다고 결론 내릴 수 있습니다.

4-2. PV Prediction

이번에는 반대로 SMILES를 입력 받아 Property를 예측하는 실험 결과를 보겠습니다.

위 표는 모델별 Property Prediction 성능 결과를 비교하고 있는데요. 가장 아래쪽에 SPMM 성능을 확인할 수 있습니다. Pretrain (SMILES, Property 모두 학습)을 진행했을 때 성능이 훨씬 좋은 모습을 볼 수 있습니다.

5. 장단점

이번 논문에서 제안한 SPMM(Structure-Property Multi-Modal foundation model)은 화학 분야의 멀티모달 학습 모델로서 여러 가지 장점과 단점을 가지고 있습니다. 이를 구체적으로 살펴보겠습니다.

첫 번째 장점은 멀티모달 학습의 도입입니다. SPMM은 Molecule Structure와 Molecule Property를 각각의 모달리티로 간주하여 학습합니다. 이를 통해 두 모달리티 간의 상호작용을 효과적으로 모델링할 수 있으며, 기존의 단일 모달리티 모델보다 더 풍부한 정보를 활용할 수 있습니다.

두 번째 장점은 양방향 예측이 가능하다는 점입니다. SPMM은 분자 구조를 입력 받아 특성을 예측할 수 있을 뿐만 아니라, 특성을 입력 받아 분자 구조를 생성할 수도 있습니다. 이는 모델의 유연성과 활용도를 크게 높여줍니다.

세 번째 장점은 효율적인 학습과 일반화 능력입니다. SPMM은 50%의 특성 벡터를 마스킹하여 학습하는 기법을 도입해 데이터 증강과 유사한 효과를 얻습니다. 이로 인해 모델은 실제 사용 시나리오에서 일부 특성이 주어지지 않아도 안정적으로 작동할 수 있습니다.

마지막으로, SPMM은 다양한 Downstream Task에 적용할 수 있다는 장점이 있습니다. 별도의 추가 학습 없이도 특성 예측, 분자 생성, 반응 예측 등 다양한 화학적 문제를 해결할 수 있습니다.

반면 단점도 존재하는데요. 첫 번째 단점은 데이터 준비의 복잡성입니다. SPMM은 분자 구조와 특성 데이터를 모두 필요로 하기 때문에, 두 종류의 데이터를 동시에 준비해야 합니다. 이는 데이터 수집과 전처리 과정에서 추가적인 노력을 요구할 수 있습니다.

두 번째 단점은 모델의 복잡성입니다. 멀티모달 학습 모델이기 때문에, 단일 모달리티 모델에 비해 더 많은 계산 자원과 시간이 필요합니다.

마지막으로, SPMM은 현재 2D 구조 정보에만 의존하고 있어 3D 구조 정보를 충분히 반영하지 못한다는 한계가 있습니다. 이는 입체화학(stereochemistry)과 같은 복잡한 화학적 특성을 다룰 때 성능 저하를 초래할 수 있습니다.

이와 같이 SPMM은 멀티모달 학습의 장점을 활용하여 다양한 화학적 문제를 해결할 수 있는 강력한 도구이지만, 데이터 준비와 모델 복잡성 등의 단점도 가지고 있습니다. 앞으로 이러한 단점을 보완하는 연구가 지속된다면 더욱 발전된 모델이 될 것으로 기대됩니다.

6. 마치며

이번 포스팅에서는 KAIST에서 발표한 논문 “Bidirectional generation of structure and properties through a single molecular foundation model”을 통해 제안된 SPMM(Structure-Property Multi-Modal foundation model)을 리뷰하였습니다. SPMM 모델은 화학 분야에서 멀티모달 학습의 가능성을 보여주며, 분자 구조와 특성 간의 복잡한 상호작용을 효과적으로 모델링할 수 있는 새로운 접근 방식을 제시하였습니다.

SPMM의 가장 큰 혁신 중 하나는 분자 구조와 특성을 각각의 모달리티로 간주하고, 이를 통합하여 양방향 예측과 생성을 가능하게 한 점입니다. 이러한 접근 방식은 기존의 단일 모달리티 모델들이 가진 한계를 극복하고, 더 풍부한 정보를 활용하여 정확하고 유연한 예측 및 생성 결과를 제공합니다. 특히, SPMM은 신약 개발과 같은 실질적인 응용 분야에서 큰 잠재력을 가지고 있습니다.

또한, SPMM 모델의 학습 과정에서 도입된 여러 기술적 요소들, 예를 들어 특성 벡터의 마스킹과 같은 기법들은 모델의 일반화 능력을 향상시키고, 실제 사용 시나리오에서의 유연성을 높이는 데 기여하였습니다. 이러한 요소들은 모델의 실용성을 높이며, 다양한 조건 하에서도 안정적으로 동작할 수 있게 합니다.

논문의 실험 결과는 SPMM 모델의 우수한 성능을 입증하였으며, 다양한 조건 하에서 높은 Validity, Uniqueness, Novelty를 보여주었습니다. 이는 SPMM이 실제 화학 연구와 산업 응용에서 유용한 도구가 될 수 있음을 시사합니다. 특히, 특성 입력에 따라 유효한 분자 구조를 생성하는 능력은 매우 인상적이었으며, 이는 화학적 디자인과 예측 작업에 큰 도움을 줄 수 있습니다.

하지만 SPMM 모델에도 몇 가지 한계가 존재합니다. 데이터 준비의 복잡성, 모델의 계산 자원 요구량, 3D 구조 정보 반영의 한계 등은 앞으로 개선해야 할 부분입니다. 이러한 단점들을 보완하는 연구가 지속된다면, SPMM은 더욱 강력한 도구로 발전할 수 있을 것입니다.

결론적으로, 이번 논문은 화학 분야에서 멀티모달 학습의 가능성을 보여주며, SPMM 모델의 혁신성과 실용성을 입증하였습니다. 앞으로 이러한 접근 방식이 더욱 발전하여, 화학 연구와 산업 응용에 더 큰 기여를 할 수 있기를 기대합니다. 이 포스팅을 통해 SPMM 모델의 중요성과 잠재력을 이해하는 데 도움이 되었기를 바랍니다. 감사합니다.