编辑 | 绿萝
在数据可用性和计算方面，大规模使得自然语言处理和计算机视觉等深度学习关键应用领域取得了重要突破。越来越多的证据表明，规模可能是科学深度学习的关键因素，但物理先验在科学领域的重要性使得规模化的策略和收益变得不确定。
近日，来自 MIT 的研究团队通过将模型和数据集大小改变多个数量级来研究大型化学模型中的神经尺度（neural-scaling）行为，研究具有超过 10 亿个参数的模型，并在多达 1000 万个数据点的数据集上进行预训练。
研究考虑用于生成化学的大型语言模型和用于机器学习原子间势的图神经网络。研究了物理先验和尺度之间的相互作用，并发现了化学语言模型的经验神经尺度关系，所考虑的最大数据集大小的标度指数（scaling exponent）为 0.17，等变图神经网络原子间势的标度指数为 0.26。
该研究以《Neural scaling of deep chemical models》为题，于 2023 年 10 月 23 日发布在《Nature Machine Intelligence》上。
深度学习在计算机视觉和自然语言处理（NLP）等领域的「不合理有效性」（unreasonable effectiveness）依赖于深度神经网络利用不断增加的计算、数据和模型容量的能力。
大型模型，包括来自 Transformers (BERT) 和 DALL-E 的双向编码器表示，已经非常成功地通过自监督预训练从大型数据集中合成信息，并执行各种下游任务，几乎不需要任何精细处理。调整 NLP 和计算机视觉中大多数最先进的模型都是从一小组大型预训练模型改编而来的。当然，我们可能期望大模型和数据集尺度将是科学深度学习取得巨大成功的先决条件。
AlphaFold、Open Catalyst Project 和 ChemBERTa 等近期工作表明，更大的数据集和模型、预训练和自监督学习可以释放化学深度学习的新功能。然而，与计算机视觉和自然语言处理不同，深度尺度化学网络的路径和潜在好处尚不清楚。
化学深度学习可以结合基于物理的先验，这可能会改善其他领域的巨大资源需求。此外，由于化学空间和分子机器学习任务的异质性和复杂性，训练在各种下游任务上表现良好的通用且稳健的模型仍然是一个紧迫的挑战。化学空间的巨大和这些任务的异质性激发了对化学中大规模模型的研究，因为此类模型非常适合未标记的多模式数据集。
最近，神经尺度定律（neural-scaling laws）作为一种描述模型性能在模型大小、数据集大小和计算等多个数量级上显著提高的趋势的方法而出现；然而，这些实验需要大量的计算资源，并依赖于众所周知的、特定领域的模型训练程序，这些程序不适用于传统的深度学习应用领域之外。
由于开发和部署大模型的成本过高，因此很难研究科学深度学习模型的神经尺度行为，这需要昂贵的超参数优化（HPO）和实验。适用于小型模型和小型数据集的架构和超参数无法 transfer 到更大的规模。随着资源需求的增加，科学深度学习将变得越来越难以获得。
在此，MIT 研究人员开发了深度（deep）化学模型的神经尺度策略，并研究了用于生成化学建模的大语言模型（LLM）和用于机器学习原子间势的图神经网络（GNN）中的神经尺度行为。
图 1：发现深度化学模型的神经尺度关系。（来源：论文）
ChemGPT，一种用于小分子自回归语言建模的生成式预训练 Transformer。使用多达 1000 万个独特分子的数据集来训练具有超过 10 亿个参数的 ChemGPT 模型。研究人员还研究了在分子动力学轨迹上训练的大型、不变和等变 GNN，并研究基于物理的先验如何影响尺度行为。
为了克服新领域中大规模超参数调整的挑战，研究人员扩展了加速神经架构搜索的技术，以在 HPO 和神经架构选择过程中将总时间和计算预算减少高达 90%。研究确定了化学模型在模型容量和数据集大小方面的扩展趋势，并展示了随着规模的增加所看到的预训练损失性能的改进。
图 2 显示了对来自分子集 (Molecular Sets，MOSES) 数据集的 200 万个分子进行训练的 ChemGPT 模型的训练性能估计 (TPE) 结果。
使用 MOSES 来演示如何使用 TPE 快速发现化学 LLM（例如 ChemGPT）的最佳设置。为了实现扩展实验，主要关注与学习动态相关的设置。为了证明 TPE 的有效性，使用 HuggingFace 中因果语言建模的默认学习率和 batch 大小初始化 ChemGPT。然后，改变学习率和 batch 大小，并使用不同的超参数训练模型 50 epochs。图 2 为 50 epochs 后的真实损失与仅 10 epochs 后使用 TPE 的预测损失。线性回归 R^2 = 0.98，Spearman 秩相关 ρ = 1.0。只需总训练预算的 20%，就能够识别出优于 HuggingFace 默认设置的模型配置。对于新数据集，该过程可轻松重复，并可加速 HPO。
图 2：在使用 TPE 进行训练的早期就确定了最佳模型，并停止非最佳模型的训练，以节省 80% 以上的总计算消耗。（来源：论文）
据观察，小 batch（甚至是 1 batch）在不同的 NFF 架构中运行良好。TPE 提供了一种快速评估 batch 大小和学习率不同组合的速度与准确度权衡的方法，这些组合是相互依赖的，并且必须一起变化才能实现大 batch 训练。
TPE 对于 GNN 的表现同样出色。使用整个训练预算的模型损失的方差很重要，表明适当的 HPO 的重要性。
图 3：在使用 TPE 进行训练的早期就确定了最佳模型，并停止非最佳模型的训练，以节省 80% 以上的总计算消耗。（来源：论文）
接下来，通过使用 TPE 有效扩展实验的策略，研究了 ChemGPT 和 NFF 中的神经尺度。
图 4：ChemGPT 模型性能（验证损失）的神经尺度作为模型（非嵌入参数数量）和数据集（标记数量）大小的函数。（来源：论文）
该研究的核心贡献是发现化学深度学习极其多样化领域的神经尺度定律：语言模型和神经原子间势。研究结果为科学深度学习中的尺度研究提供了动力和实践指导，并为大规模和物理深度学习的交叉点提供了许多富有成效的新研究方向。
该研究的一个重要发现是，对于大型化学语言模型和 NFF，在模型大小、数据集大小或计算方面都没有饱和模型损失。在研究化学尺度的局限性方面仍有许多进一步的工作要做。
论文链接：https://www.nature.com/articles/s42256-023-00740-3
相关文章
