arxiv202209 | ChemBERTa-2+：Towards Chemical Foundation Models

原文标题：ChemBERTa-2：Towards Chemical Foundation Models

链接：https://arxiv.org/abs/2209.01712v1

 ChemBERTa-1是2020年的，可以看看：

NeurIPS 2020 | ChemBERTa+：用于分子性质预测的大规模自监督预训练_羊飘的博客-CSDN博客

从PubChem获得的数据集可达77M。迄今为止用于分子预训练的最大数据集之一。采用了mask language modeling(MLM)和multi-task regression(MTR)对7700万SMILES字符串的大型语料进行预训练。

一、模型

数据集：数据集划分为5M、10M和77M的采样子集(完整集)，组成三个数据集，用于两个预训练任务。

1、Pretraining Strategies and Setup

MLM：vocab_size为591标记的词汇表，max_seq_len为512个token。

MTR：为训练数据集中的每个分子计算200个分子属性（RDKit计算得出）。然后，训练一个多任务回归(MTR)体系结构，以同时预测这些属性。因为预测属性有非常不同的规模和范围，在训练之前对每个任务进行规范化标签。

2、 Hyperparameter Search

选择50个随机的超参数配置：hidden_size, heads, dropout, intermediate size, number of hidden layers, learning rate（参数量在5M和46M之间）。每个配置都在每个MLM和MTR的预训练任务上训练，使用5M数据集。选择了五种配置，不同的验证损失值，在10M和77M集上训练。

3、Finetuning on MoleculeNet

对来自MoleculeNet的几个回归和分类任务进行了评估。选择单一的代表性任务:ClinTox的临床毒性(CT_TOX)任务和Tox21的p53应激反应通路激活(SR-p53)任务。对于每个数据集，使用DeepChem中的scaffold splitter生成80/10/10的train/valid/test拆分。基于验证损失对模型进行了多达100个epoch的微调

二、实验结果

在不同任务(MLM和MTR)和不同数据集大小(5M、10M和77M)上预训练的ChemBERTa-2与选定的MoleculeNet任务上现有架构的比较。

在每一个下游精调任务中，MTR任务预训练的模型往往比MLM任务预训练的模型表现得更好。然而，在我们当前的实现中，MTR训练比MLM慢得多（200个元素的标签向量增加了数据集的大小）。

对于给定的架构，MLM预训练损失与MTR预训练损失非常吻合：

在三个数据集(5M, 10M和77M)上训练的5种模型配置的每一种的预训练损失：

 使用：

from transformers import AutoModelForMaskedLM
encoder = AutoModelForMaskedLM.from_pretrained("DeepChem/ChemBERTa-77M-MLM")
encoder.resize_token_embeddings(768)    # 设置token_embedding_dim
# 对token embedding，经过encoder，得到的是MaskedLM Object，需要经过ouput.logits得到变量