Nat. Mach. Intell. | 深度化学模型的神经缩放

今天为大家介绍的是来自Nathan C. Frey团队的一篇论文。在数据可用性和计算量方面的大规模扩展，使得深度学习在自然语言处理和计算机视觉等关键应用领域实现了重要突破。越来越多的证据表明，在科学深度学习中，规模可能是一个关键因素，但科学领域中物理先验的重要性使得规模扩展的策略和益处尚不确定。在这里，作者通过改变模型和数据集的大小（跨越多个数量级），研究了大型化学模型中的神经缩放行为，研究对象包括拥有超过十亿参数的模型，这些模型在高达一千万数据点的数据集上进行了预训练。

深度学习在计算机视觉和自然语言处理（NLP）等领域的有效性依赖于深度神经网络利用不断增加的计算能力、数据量和模型容量的能力。包括BERT和DALL-E在内的大型模型，通过自我监督的预训练成功地从大型数据集中提取信息，并且在很少甚至不需要微调的情况下就能执行各种下游任务，以至于大多数在NLP和计算机视觉领域的最先进模型都是从一小组大型预训练模型改编而来。自然地，我们可能会期望在科学领域中，大规模模型和数据集的扩展将是实现深度学习巨大成功的先决条件。最近的研究，如AlphaFold、ChemBERTa表明，更大的数据集和模型、预训练以及自我监督学习——这些在计算机视觉和NLP中的关键要素——为化学领域的深度学习解锁了新的能力。然而，与计算机视觉和NLP不同，化学网络的扩展路径及其潜在收益尚不明确。化学深度学习可以纳入基于物理的先验知识，这可能减轻在其他领域看到的陡峭资源需求。此外，由于化学空间的异质性和复杂性，以及分子机器学习任务的多样性，培训通用且鲁棒的模型在各种下游任务上表现良好仍然是一个迫切的挑战。化学空间的巨大性和这些任务的异质性促使了对大规模化学模型的研究，因为这些模型非常适合未标记的多模态数据集。最近，神经缩放法则作为一种描述模型性能随着模型大小、数据集大小和计算量在多个数量级上显著提高的趋势的方法出现，然而，这些实验需要巨大的计算资源，并依赖于在传统深度学习应用领域之外不适用的众所周知的领域特定模型训练程序。

图 1

在这篇文章中作者为扩展深度化学模型开发了策略，并研究了大型语言模型（LLM）在生成化学建模和图神经网络（GNN）在机器学习原子间势能方面的神经缩放行为。作者引入了ChemGPT，这是一个针对小分子自回归语言建模的生成预训练transformer。作者用超过10亿参数训练了ChemGPT模型，使用的数据集包含高达1000万个独特分子。作者还研究了在分子动力学轨迹上训练的大型、不变和等变GNN，并探讨了基于物理的先验知识如何影响缩放行为。论文的工作流程如图1所示。

加速超参优化

图 2

图 3

为了进行广泛的缩放实验，作者首先需要找到合理的超参数和训练设置。与自然语言处理（NLP）和计算机视觉不同，大规模化学深度学习没有默认的模型架构、数据集、任务、超参数设置或训练设置。简单地从其他深度学习领域或小规模实验转移经验结果将导致次优结果。为此，作者首先解决在一般设置中加速超参数优化（HPO）的问题，适用于新的模型架构、异质数据集以及之前未曾研究过的规模。图2显示了对Molecular Sets (MOSES)数据集中200万分子训练的ChemGPT模型的训练性能估计（TPE）结果。MOSES比PubChem小，代表了通常用于训练化学生成模型的数据集。作者在这里使用MOSES来演示如何使用TPE快速发现像ChemGPT这样的化学语言大型模型（LLM）的最优设置。为了实现缩放实验，作者主要关注与学习动态相关的设置（例如，批量大小和学习率），这些设置会影响大规模训练，并根据模型类型和数据集特征而波动。为了展示TPE的有效性，作者使用HuggingFace中因果语言建模的默认学习率和批量大小初始化ChemGPT。然后，改变学习率和批量大小，并用不同的超参数训练模型50个周期。图2显示了50个周期后的真实损失与仅经过10个周期后使用TPE预测的损失。线性回归的R2 = 0.98，而Spearman相关ρ = 1.0。仅用20%的总训练预算，作者能够识别出比HuggingFace默认设置更好的模型配置。这个过程可以轻松重复用于新数据集，并实现加速HPO。

虽然像ChemGPT这样的LLM的训练程序已经很成熟，但将神经力场（NFF）数据扩展到更大的数据集和更具表现力的模型需要新的、可扩展的训练程序。通过数据并行实现大批量训练是加速训练的一种方法，但存在已知的局限性，且不同领域的批量大小差异很大。对于NFF来说，这个问题尤为严重，因为每个数据点实际上包含3N + 1个能量和原子力的标签，其中N是原子的数量，不同批次易变较大。因此，已经观察到小批量大小（甚至是1的小批量）在不同的NFF架构中表现良好。TPE提供了一种快速评估不同批量大小和学习率组合的速度-精度权衡的方法，这些组合是相互依赖的，必须一起变化以实现大批量训练。

表 1

TPE对于图神经网络（GNN）同样表现良好。作者重复了TPE流程，改变了学习率和批量大小，对SchNet、PaiNN和SpookyNet进行了训练。数据集包含10种小有机分子的100,000个结构，这些结构取自原始MD17数据集，并在PBE/def2-SVP理论水平下重新计算了能量和力。仅使用总训练预算的20%，作者通过TPE为SchNet和PaiNN预测了其建模能力（见图3）。由于SpookyNet是一个复杂的架构，包含非局部相互作用和经验修正，它表现出缓慢的收敛速度，其训练速度与收敛模型损失的相关性低于SchNet和PaiNN。然而，TPE对SpookyNet模型配置的排序仍然是稳健的，这允许我们在训练早期就淘汰非最优模型配置。使用TPE进行线性回归的拟合度指标在表1中给出。

神经缩放与模型损失的关系

图 4

接下来，作者通过使用TPE有效地扩大实验规模，并研究了ChemGPT和神经力场（NFF）的神经缩放。对于每个模型，作者首先执行TPE以识别在不同模型和数据集大小范围内预计表现良好的超参数选择。然后系统地改变数据集大小（d）和模型大小（m），并进行全面实验以确定收敛的损失，L(m, d)。由于计算资源的限制，作者训练ChemGPT模型固定的周期数（十个）以确定损失。图4显示了作为模型和数据集大小函数的预训练损失，覆盖了多个数量级。模型在自我监督的因果语言建模设置中进行训练，并针对固定验证集的下一个标记预测进行评估。令人惊讶的是，随着规模的增加，没有看到损失改善的限制。预训练损失随着数据集大小的增加而单调改善，直到近1000万分子。此外，对于固定的数据预算，增加模型大小会持续改善预训练损失，直到模型达到10亿+非嵌入参数。对于这里考虑的最大数据集，模型在超过1亿非嵌入参数时看到损失改善的递减回报。有趣的是，相比于更大的数据集，较小的数据集对于增加模型大小时看到更大的损失改善。对于考虑的最大数据集，模型损失在超过1亿参数后迅速饱和。然而，对于考虑的最小数据集，模型损失在10到10^7参数之间的模型大小时呈现先平缓态势，然后显著改善。这表明，对于固定的小预训练数据预算，通过扩大模型大小可以显著改善预训练损失。不管模型大小如何，增加数据集大小都会持续改善损失，且没有迹象表明作者设置的数据集大小存在递减回报。

GNN的实验结果如何呢

图 5

图神经网络（GNN）表现出稳健的神经缩放行为。与大型语言模型（LLM）相比，大规模GNN的潜在益处不那么明确，因为深度和宽度对于GNN来说并不等同，还有一些影响性能的其他参数超出了模型大小的概念，例如在图构建中的最近邻截断。为了简化GNN缩放实验，作者通过系统改变网络宽度和卷积数（深度）来改变GNN的容量（深度×宽度）。作者训练GNN来预测来自ANI-1x数据集（500万个小有机分子的密度泛函理论计算）的原子力，这是目前公开可用的最大的小分子能量和力的数据集。神经力场（NFF）模型使用学习率调度器进行训练，该调度器在验证损失没有改进的情况下每50个周期降低学习率，直到学习率达到10^-7。损失是L1损失，实验如图5所示，覆盖了四个数量级的数据集大小。等变GNN，PaiNN的神经缩放结果（图5）显示，随着数据集大小的增加，损失持续改善。对于固定的数据集大小，收敛的损失与总训练时间（计算）和模型容量强相关。除了10^3数据点（一些小型模型很快达到收敛）外，收敛的损失与模型容量的Spearman相关系数ρ≥0.88，与总训练时间的ρ≥0.75。这意味着最佳模型是那些具有最佳容量并能够在验证损失不平稳的情况下训练最长时间的模型。最佳容量和深度与宽度的比例随数据集大小而变化，即理想的GNN容量取决于数据集大小，这些选择会影响收敛的损失。这些效应也可能是随机初始化的偶然结果，可能会随着重复试验而减弱。有趣的是，在10^4数据点处出现了明显的变化——此时收敛的损失几乎与模型容量完美相关（Spearman的ρ≥0.93）。这可能表明训练集和验证集之间存在大量重叠，以至于高容量模型仅仅表现出比低容量模型更好的记忆能力。

编译 | 曾全晨

审稿 | 王建民

参考资料

Frey, N.C., Soklaski, R., Axelrod, S. et al. Neural scaling of deep chemical models. Nat Mach Intell 5, 1297–1305 (2023).

 https://doi.org/10.1038/s42256-023-00740-3