[论文解读] ChemBERTa-2: Towards Chemical Foundation Models

Walid Ahmad, Elana P. Simon+3 more|arXiv (Cornell University)|Sep 5, 2022

Machine Learning in Materials Science被引用 133

一句话总结

ChemBERTa-2 预训练了一个基于 SMILES 的大规模 Transformer（77M 化合物），并比较 MLM 与 MTR 预训练；它在 MoleculeNet 结果上具竞争力，并分析预训练如何扩展到下游任务。

ABSTRACT

Large pretrained models such as GPT-3 have had tremendous impact on modern natural language processing by leveraging self-supervised learning to learn salient representations that can be used to readily finetune on a wide variety of downstream tasks. We investigate the possibility of transferring such advances to molecular machine learning by building a chemical foundation model, ChemBERTa-2, using the language of SMILES. While labeled data for molecular prediction tasks is typically scarce, libraries of SMILES strings are readily available. In this work, we build upon ChemBERTa by optimizing the pretraining process. We compare multi-task and self-supervised pretraining by varying hyperparameters and pretraining dataset size, up to 77M compounds from PubChem. To our knowledge, the 77M set constitutes one of the largest datasets used for molecular pretraining to date. We find that with these pretraining improvements, we are competitive with existing state-of-the-art architectures on the MoleculeNet benchmark suite. We analyze the degree to which improvements in pretraining translate to improvement on downstream tasks.

研究动机与目标

• 通过将大语言模型预训练的概念迁移到基于 SMILES 的表示来激发化学基础模型。
• 探索预训练数据集规模和目标（MLM vs MTR）对下游分子预测任务的影响。
• 识别预训练改进在 MoleculeNet 基准测试中的性能提升的转化方式。
• 提供大规模分子模型的预训练策略与超参数的实用指导。

提出的方法

• 在 PubChem 的 77M SMILES 上使用基于 RoBERTa 的 Transformer（ChemBERTa-2）进行 MLM 预训练。
• 比较 MLM 与多任务回归（MTR）预训练目标。
• 在 5M、10M 和 77M 数据集上对架构与训练设置进行超参数搜索。
• 对预训练模型在 MoleculeNet 任务上进行微调，采用 scaffold-based 拆分以及在适当情况下进行任务级归一化或类别加权。
• 分析预训练损失与下游性能之间的迁移关系。
• 开源训练模型以实现可重复性。

实验结果

研究问题

• RQ1 更大规模的 SMILES 预训练是否能提升 MoleculeNet 任务的分子性质预测性能？
• RQ2 在下游性能和训练效率方面，MLM 与多任务回归预训练有何差异？
• RQ3 预训练损失在特定下游数据集（如溶解度、毒性）上的提升程度如何体现？

主要发现

• 使用更大规模数据集进行预训练（最多 77M SMILES）可获得更低的预训练损失并提升若干任务的下游性能。
• MTR 预训练在下游任务上通常优于 MLM，但基于 MLM 的架构搜索可以通过相关损失有效引导 MTR 的预训练选择。
• 以 MLM 或 MTR 在 77M 数据上进行预训练的 ChemBERTa-2 配置在表 1 的 8 个任务中有 6 项超过 D-MPNN，达到具竞争力的结果。
• 预训练损失的改进与某些任务（如 Lipophilicity）上的收益相关，但跨数据集与任务的转移效率存在差异。
• 在 commitment 于较慢的 MTR 训练之前使用 MLM 来选择超参数可能成为一种实际的工作流。
• 降维分析（UMAP）显示 ChemBERTa 的嵌入可以作为下游分类任务的强大先验。