[논문 리뷰] Pre-training via Denoising for Molecular Property Prediction

Sheheryar Zaidi, Michael Schaarschmidt외 7명|arXiv (Cornell University)|2022. 05. 31.

Machine Learning in Materials Science인용 수 30

한 줄 요약

저자들은 3D 분자 구조의 노이즈 제거를 통해 표현을 학습하는 자기지도 사전 학습 방법을 제안하며 QM9 및 다른 벤치마크에서 최첨단 결과를 달성한다.

ABSTRACT

Many important problems involving molecular property prediction from 3D structures have limited data, posing a generalization challenge for neural networks. In this paper, we describe a pre-training technique based on denoising that achieves a new state-of-the-art in molecular property prediction by utilizing large datasets of 3D molecular structures at equilibrium to learn meaningful representations for downstream tasks. Relying on the well-known link between denoising autoencoders and score-matching, we show that the denoising objective corresponds to learning a molecular force field -- arising from approximating the Boltzmann distribution with a mixture of Gaussians -- directly from equilibrium structures. Our experiments demonstrate that using this pre-training objective significantly improves performance on multiple benchmarks, achieving a new state-of-the-art on the majority of targets in the widely used QM9 dataset. Our analysis then provides practical insights into the effects of different factors -- dataset sizes, model size and architecture, and the choice of upstream and downstream datasets -- on pre-training.

연구 동기 및 목표

• 레이블 데이터가 희소한 3D 구조로부터 분자 특성 예측을 위한 사전 학습 방식을 동기 부여하고 개발한다.
• 평균 구조로부터 의미 있는 포스 필드와 같은 표현을 학습하기 위해 노이즈 제거/점수 매칭을 활용한다.
• 노이즈 제거 기반 사전 학습이 다양한 데이터셋과 구조에서 다운스트림 작업 성능을 향상시킨다는 것을 입증한다.
• 업스트림 데이터셋 크기, 모델 크기, 아키텍처 및 데이터셋 유사성이 전이성에 어떤 영향을 미치는지에 대해 실용적 통찰을 제공한다.

제안 방법

• 가우시안 노이즈로 원자 좌표를 섞어 3D 분자 구조에 대한 노이즈 제거 사전 학습 목적을 형식화하고 GNN을 학습시켜 노이즈를 예측한다.
• 노이즈 제거 목적을 평형 구조 주위의 볼츠만 분포의 가우시안 혼합 근사를 나타내는 포스 필드를 학습하는 것으로 해석한다.
• 이 사전 학습을 Graph Net Simulator(GNS)와 Tailored Activation Transform(TAT) 강화 변형(GNS-TAT)에 적용한다.
• 불변/동등성 구성요소를 갖춘 트랜스포머 기반 모델인 TorchMD-NET에도 노이즈 제거 사전 학습을 적용하여 아키텍처에 구애받지 않음을 시연한다.
• 사전 학습을 보완하기 위한 보조 손실로 Noisy Nodes를 사용하고 그 영향력을 분석한다.
• 전이 성능에 대한 데이터셋 크기, 모델 크기 및 업스트림/다운스트림 데이터셋 간 관계를 분석한다.

실험 결과

연구 질문

• RQ1 노이즈 제거를 통한 사전 학습이 무작위 초기화에 비해 다운스트림 분자 특성 예측을 개선하는가?
• RQ2 업스트림(사전 학습)과 다운스트림 데이터셋 간의 관계가 노이즈 제거 사전 학습의 이점에 얼마나 의존하는가?
• RQ3 노이즈 제거 사전 학습이 다양한 아키텍처(GNN 및 트랜스포머)와 작업(QM9, OC20, DES15K)에서 효과적일 수 있는가?

주요 결과

• 노이즈 제거를 통한 사전 학습은 여러 데이터셋과 타깃에서 다운스트림 성능을 크게 향상시키며 QM9의 타깃 다수에서 새로운 최첨단 성능을 달성한다.
• 노이즈 제거 사전 학습 후 Noisy Nodes를 적용한 GNS-TAT가 QM9의 12개 타깃 중 10개에서 최첨단 결과를 산출하되 일부 타깃은 노이즈 스케일에 민감할 수 있다.
• PCQM4Mv2(3.4M 3D 구조)에서의 사전 학습은 일반적으로 QM9 및 DES15K로 잘 전이되나 원소 겹침 및 작업 유사도에 따라 전이 정도가 달라진다.
• OC20에서 사전 학습한 경우 PCQM4Mv2로 사전 학습했을 때 OC20 IS2RE 성능이 향상되지 않았으며 이는 데이터셋 간 이질성으로 인한 전이 한계를 시사하지만 OC20 자체를 사전 학습하면 수렴 속도가 빨랐다.
• 업스트림 데이터를 늘리면 일반적으로 다운스트림 표현이 개선되고 더 큰 모델은 사전 학습의 혜택을 받아 처음부터 학습된 모델을 능가할 수 있다.
• 사전 학습은 또한 힘 예측(예: MD17 아스피린)을 개선하고 디코더 구성요소만 미세조정할 때 학습을 가속화한다.