[논문 리뷰] From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction

Nima Shoghi, Adeesh Kolluru외 4명|arXiv (Cornell University)|2023. 10. 25.

Machine Learning in Materials Science인용 수 21

한 줄 요약

이 연구는 Joint Multi-domain Pre-training (JMP)을 도입하여 다양한 화학 도메인에서 단일 모델을 학습하고, 처음부터 학습하는 것 대비 평균 59% 향상을 달성하며 40개 과제 중 34개에서 최첨단 수준과 일치하거나 이를 능가합니다.

ABSTRACT

Foundation models have been transformational in machine learning fields such as natural language processing and computer vision. Similar success in atomic property prediction has been limited due to the challenges of training effective models across multiple chemical domains. To address this, we introduce Joint Multi-domain Pre-training (JMP), a supervised pre-training strategy that simultaneously trains on multiple datasets from different chemical domains, treating each dataset as a unique pre-training task within a multi-task framework. Our combined training dataset consists of $\sim$120M systems from OC20, OC22, ANI-1x, and Transition-1x. We evaluate performance and generalization by fine-tuning over a diverse set of downstream tasks and datasets including: QM9, rMD17, MatBench, QMOF, SPICE, and MD22. JMP demonstrates an average improvement of 59% over training from scratch, and matches or sets state-of-the-art on 34 out of 40 tasks. Our work highlights the potential of pre-training strategies that utilize diverse data to advance property prediction across chemical domains, especially for low-data tasks. Please visit https://nima.sh/jmp for further information.

연구 동기 및 목표

• 다양한 화학 도메인에 걸친 원자 특성 예측의 일반화 도전과제를 동기 부여하고 해결합니다.
• 크고 이질적인 데이터셋을 활용하는 확장 가능한 사전 학습 전략을 개발합니다.
• 한정된 데이터로 하위 작업에 대한 파인튜닝을 가능하게 하면서도 강력한 성능을 유지합니다.
• 사전 학습 도메인을 넘어서는 미지의 도메인(대분자 및 재료)으로의 일반화 가능성을 입증합니다.

제안 방법

• Joint Multi-domain Pre-training (JMP)을 다중 작업 감독 학습 프레임워크로 제안합니다.
• 에너지 및 힘에 대해 데이터세트별 예측 헤드를 갖춘 단일 백본 모델(GemNet-OC)을 사용합니다.
• 데이터세트별 선형 에너지 기준화 및 힘 정규화를 통해 타깃을 정규화합니다.
• 배치 구성 중 데이터셋 크기의 균형을 맞추기 위해 온도 기반 샘플링을 적용합니다.
• 다른 시스템 크기를 가진 데이터셋의 기여를 균형 맞추기 위해 구조별 손실 감소를 도입합니다.
• 정규화(가중치 감소, 에지 드롭아웃, EMA)를 포함하는 다중 작업 손실에 대한 단위 스칼라화(unitary scalarization)를 채택합니다.
• 사전 학습 헤드를 작업 특화 헤드로 교체하고 선택적으로 에너지 기울기를 통해 힘을 계산하여 파인튜닝합니다.

실험 결과

연구 질문

• RQ1 다양한 화학 도메인에서 학습된 단일 사전 학습 모델이 소분자, 대분자, 재료에 걸친 하위 작업으로 얼마나 잘 일반화합니까?
• RQ2 공통 다중 도메인 사전 학습이 scratch 학습 및 이전의 단일 도메인 또는 자기지도 학습 방식보다 다양한 벤치마크에서 더 우수합니까?
• RQ3 데이터 균형, 손실 공식화, 정규화 전략이 다중 작업 사전 학습 성능에 미치는 영향은 무엇입니까?
• RQ4 JMP가 제한된 다운스트림 데이터로 대형 모델의 파인튜닝을 가능하게 하고 미지의 도메인으로의 이전을 개선할 수 있습니까?

주요 결과

• JMP는 파인튜닝 작업에서 scratch 학습 대비 평균 59% 향상을 제공합니다.
• QM9, rMD17, MD22, SPICE, MatBench, QMOF에 걸친 40개 파인튜닝 작업 중 34개에서 최첨단과 일치하거나 이를 상회합니다.
• 235M 매개변수 JMP 모델이 다수의 저데이터 벤치마크에서 최첨단 성능을 달성합니다.
• JMP-L의 파인튜닝은 GN-OC-L의 성능에 약 1/12의 학습 시간으로 도달하여 더 빠른 적응을 나타냅니다.
• 다양한 화학 데이터로의 사전 학습은 비평형 구성 및 도메인 밖 목표에 일반화되는 전이 가능한 표현을 제공합니다(예: MatBench 및 QMOF의 재료 특성).
• 전체 JMP 사전 학습 비용은 scratch 대비 다운스트림 파인튜닝이 12배 이상 빠르다는 이점으로 상쇄됩니다.