关于ChemBERTa-3开源训练框架的研究报告

主要作者与发表信息

本研究由Riya Singh、Aryan Amit Barsainyan、Rida Irfan等人作为共同第一作者，Bharath Ramsundar作为通讯作者完成。研究机构包括Deep Forest Sciences、劳伦斯利弗莫尔国家实验室（Lawrence Livermore National Laboratory）、密歇根大学（University of Michigan）和加州大学伯克利分校（University of California, Berkeley）。该论文于2026年1月19日发表在期刊《Digital Discovery》上。

研究背景与目的

药物发现是一个复杂且耗时漫长的过程，涉及识别潜在的治疗化合物并评估其生物功效。分子性质预测（molecular property prediction）模型在药物发现过程中具有重要影响。近年来，深度学习已被广泛用于分子性质预测，包含图神经网络（Graph Neural Networks, GNNs）和基于Transformer的架构等一系列方法。其中，大规模预训练的Transformer架构，即化学基础模型（chemical foundation models），因其能够直接从海量未标记化合物数据库中学习基础化学知识而备受关注。

然而，尽管化学基础模型的发展取得了快速进展，但对不同预训练方法在不同模型架构上的表现缺乏系统性的比较。造成这一现象的关键原因是预训练和评估大规模化学基础模型所需的庞大基础设施带来的挑战。此外，许多现有模型并未完全开源，或不同模型间的基准测试（benchmarking）并未采用统一标准，导致报告的比较结果不够准确。为了解决这些问题，本研究的作者团队推出了ChemBERTa-3，一个开源、可扩展的平台，旨在为化学基础模型的训练和基准测试提供统一的框架。

研究的详细工作流程

1. 框架与基础设施建设 本研究首先构建了ChemBERTa-3框架，该框架完全集成到开源的DeepChem库中。为了实现多GPU上的数据并行训练，ChemBERTa-3利用了Ray的分布式训练基础设施，并提供了专门用于高效预训练和微调大规模化学基础模型的工具。该框架支持基于Transformer和基于图模型的预训练。作者在DeepChem中引入了新的核心基类ModularTorchModel，支持灵活的预训练、微调以及中间损失计算。他们还增加了HuggingFace DeepChem封装器（wrapper），用于无缝集成来自Transformers生态的模型和标记器（tokenizer）。此外，新添加了如RDKitConformerFeaturizer和GROVERFeaturizer等特征化器，以支持3D构象和功能团感知的图表示。

2. 预训练数据与实验对象 研究使用了大规模的预训练数据集ZINC20，这是一个包含14亿个化合物的化学库。为了探究数据规模的影响，模型在不同规模的ZINC数据集上进行了基准测试。具体模型和数据量如下：Molformer模型变体分别在1000万、1亿、5.5亿和11亿个分子上进行预训练。其中，C3-Molformer-550M使用5亿ZINC和5千万PubChem分子进行预训练，而C3-Molformer-1.1B则使用了10亿ZINC和1亿PubChem分子。ChemBERTa模型分别在1000万和1亿个分子上进行预训练。图基模型（如InfoGraph、InfoMax3D和GROVER）由于扩展困难，仅在25万个分子的ZINC子集上进行了预训练。作为基线，随机森林（RF）、图卷积网络（GCN）和D-MPNN仅在下游任务的微调数据上训练。

3. 模型架构与预训练方法 该框架对多种架构进行了基准测试。Transformer架构方面，使用了基于BERT的ChemBERTa和基于Molformer架构的C3-Molformer，两者都采用掩码语言模型（Masked Language Modeling, MLM）进行预训练，即随机遮蔽输入SMILES序列中的标记，并训练模型正确预测它们。图基架构方面，测试了InfoGraph（通过最大化图级和子结构表示间的互信息来预训练）、InfoMax3D（在InfoGraph基础上利用3D分子数据，最大化2D图与3D表示间的互信息）以及GROVER（一种整合了消息传递网络和Transformer风格的架构）。D-MPNN和GCN等作为监督学习基线。

4. 下游评估与基准测试 预训练完成后，所有模型在MoleculeNet基准测试套件中选择的分类和回归任务上进行微调和评估，以覆盖药物化学应用。数据集包括BACE、BBBP、Tox21、HIV、SIDER、ClinTox、ESOL、FreeSolv和Lipo等。研究特别强调了数据集划分（dataset splitting）的重要性。研究团队在两种不同的骨架划分（scaffold split）上进行了评估：一种是DeepChem实现的严格按照Bemis-Murcko骨架的划分法，另一种是Molformer论文提供的划分法。作者发现这两种方法存在显著差异。通过比较最小谷本距离（Minimum Tanimoto Distance, MTD）分析，他们发现DeepChem的划分会导致训练集和测试集分子结构差异更大，从而使评估任务更具挑战性和真实性；而Molformer的划分则将结构更相似的分子分在一起，导致模型性能虚高。

5. 大规模部署与稳定性处理 为测试可复现性，作者在两种不同环境中训练了大型Molformer模型：一是基于AWS的Ray部署（使用40个T4 GPU的竞价实例，耗时约10天），二是本地高性能计算（HPC）集群（使用AMD MI300A APU，耗时约4天）。在扩展训练中，团队遇到了一些数值不稳定问题，如因学习率过高导致损失突然飙升（loss spike）和因梯度爆炸导致的NaN值。他们采用的解决策略包括从最新的稳定检查点（checkpoint）重启训练，以及一种两阶段训练策略：先在小数据集上用极低学习率进行预训练以稳定初始权重，然后再在全量数据上以正常学习率进行训练。

主要研究成果

1. 框架的可复现性与统一性 研究成功展示了ChemBERTa-3基础设施能够在云计算和本地HPC两种截然不同的环境中进行有意义的部署。两者的训练结果直接可比，验证了框架和结果的可复现性。

2. Transformer与图模型的性能比较 基准测试结果显示，在小规模数据（25万分子）上，图基预训练方法与Transformer基方法性能大致相当。然而，Transformer模型在扩展到海量数据时明显更容易，而图基方法则面临严重的工程挑战，例如GROVER的特征化处理在100万数据集上就消耗了100GB的磁盘空间。尽管如此，图基模型在小规模上展现的强大性能表明，进一步投资于图预训练基础设施可能是值得的。

3. 数据集划分对评估指标的显著影响 在Molformer骨架划分下（较容易），RF基线在六个分类任务中的三个（BACE, HIV, SIDER）上排名第一，C3-Molformer-1.1B的性能低于已发表的Molformer结果。在DeepChem骨架划分下（较困难），所有模型的绝对AUC值普遍下降。Molformer-LHPC模型在BACE、BBBP、Tox21和ClinTox四个分类任务上获得第一名，展示了其强大性能；而ChemBERTa-MLM-100M依然保持稳定。在回归任务上，Transformer模型同样优于传统基线，D-MPNN在Lipo等任务上表现强劲。这一对比强调了不同划分算法会极大地影响结构分离程度，进而导致跨论文直接比较结果可能具有误导性。

4. 预训练数据规模的效益递减 实验表明，将预训练数据从1000万扩展到11亿个SMILES，对下游性能的提升有限，呈现出边际收益递减的现象。作者推测，当前的MLM预训练目标可能无法完全捕获有效表示学习所需的内在化学性质。例如，MLM倾向于预测高频原子标记（如“C”），鼓励“捷径学习”而非化学上有意义的表示。

结论、价值与未来展望

本研究的核心结论是推出了ChemBERTa-3框架，这是一个集成了DeepChem的可扩展、开源训练与基准测试平台。其科学价值在于为化学基础模型领域提供了统一、可复现的评估协议，解决了以往因评估标准不一（特别是骨架划分算法不同）而导致的不可靠比较问题。应用价值在于该框架支持研究者轻松预训练和微调模型，并开放了全部模型权重、训练代码和部署工作流，有望成为下一代化学基础模型和更广泛的科学开源大型语言模型（LLMs）的基础构建块。作者建议采用DeepChem/ChemBERTa-3作为未来基准测试的标准化框架，以确保一致性。

研究亮点

1. 开源框架与统一测试：提供了一个完整的、可扩展的开源框架，首次系统性地在统一标准下比较了Transformer和图模型及多种预训练方法。
2. 揭示基准测试陷阱：通过实验和MTD分析，明确揭示了Molformer和DeepChem两种骨架划分方法的根本差异及其对模型性能评估的重大影响，倡导标准化基准测试。
3. 模型规模与效益研究：通过不同数据规模的详实实验，揭示了当前预训练方法下数据量增加带来的边际效益递减现象，为未来研究方向提供了数据支持。
4. 深度工具集成：向DeepChem库贡献了大量新基础设施（如ModularTorchModel、HuggingFace封装器等），极大地增强了该库对基础模型研究的支持能力。