编辑丨&
扩展一直是提高机器学习各个领域的模型性能和泛化的关键因素。尽管在扩展其他类型的机器学习模型方面取得了成功，但对神经网络原子间电位 （NNIP） 扩展的研究仍然有限。
该领域的主要范式是将许多物理域约束合并到模型中，例如旋转等方差等对称性约束。来自美国加州大学伯克利分校和劳伦斯伯克利国家实验室 （Berkeley Lab） 的研究团队认为，这些日益复杂的领域约束抑制了 NNIP 的扩展能力，从长远来看，这种策略可能会导致模型性能停滞不前。
为了系统地研究 NNIP 扩展属性和策略，团队提出了一种全新的专为可扩展性而设计的 NNIP 架构：高效缩放注意力原子间势 （EScAIP），意在通过注意力机制扩展模型有效提高模型表达能力。
该方法以「The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains」为题，于 2024 年 10 月 31 日发布于 Arxiv 预印平台。
与现有的 NNIP 模型相比，EScAIP 因使用高度优化的注意力 GPU 内核实现，效率大幅提高，推理时间至少加快 10 倍，内存使用量减少 5 倍。
团队强调，他们的方法应被视为一种哲学而非特定模型。它代表了开发通用 NNIP 的概念验证，这些 NNIP 通过扩展实现更好的表达性，并通过增加计算资源和训练数据继续有效扩展。
关于神经网络通用进化的思考
近年来，扩展模型大小、数据和计算的原理已成为提高机器学习 （ML） 性能和泛化的关键因素，涵盖从自然语言处理 （NLP）到计算机视觉 （CV）。ML 中的扩展在很大程度上取决于是否能最好地利用 GPU 计算能力。这通常涉及如何有效地将模型大小增加到较大的参数量级，以及优化模型训练和推理以实现最佳计算效率。
与这类发展并行的方向，涵盖了原子模拟，解决药物设计、催化、材料等方面的问题。其中，机器学习原子间势，尤其是神经网络原子间势 （NNIP），作为密度泛函论等计算密集型量子力学计算的替代模型而广受欢迎。
NNIP 旨在高效、准确地预测分子系统的能量和力，允许在难以用密度泛函理论直接模拟的系统上执行几何弛豫或分子动力学等下游任务。
当前的 NNIP 主要基于图神经网络 （GNN）。该领域的许多有效模型越来越多地尝试将受物理启发的约束嵌入到模型中，这些约束包括将预定义的对称性（例如旋转等方差）合并到 NN 架构中，以及使用复杂的输入特征集。
团队认为，这些日益复杂的领域约束抑制了 ML 模型的扩展能力，并且可能会随着时间的推移在模型性能方面趋于平稳。随着模型规模的增加，可以假设施加这些约束会阻碍有效表示的学习，限制模型的泛化能力，并阻碍有效的优化。其中许多功能工程方法并未针对 GPU 上的高效并行化进行优化，进一步限制了它们的可扩展性和效率。
基于这些理解，他们开发了高效缩放注意力原子电位 （EScAIP）。该模型在各种化学应用中实现了最佳性能，包括在 Open Catalyst 2020、Open Catalyst 2022、SPICE 分子和材料项目 （MPTrj） 数据集上的最佳性能。
图 1：EScAIP 与 Open Catalyst 数据集（OC20）上的基线模型之间的效率、性能和扩展比较。
EScAIP 可以很好地与计算一起扩展，并且其设计方式将随着 GPU 计算的不断进步而进一步提高效率。
模型的种子与未来的枝芽
这些模型通常经过训练，可以根据系统属性（包括原子序数和位置）预测系统能量和每原子力。模型分为两类：基于组表示节点特征的模型，以及基于笛卡尔坐标表示的节点特征的模型。
NNIP 领域也越来越关注使用量子力学模拟生成更大的数据集，并使用它来训练模型。有一种趋势是将物理启发的约束纳入 NNIP 模型架构，例如所有将对称约束合并到模型中的组。然而，还有其他工作路线并没有试图直接在 NNIP 中构建对称性，而是尝试“近似”对称性。
通过消融研究，该团队系统地研究了缩放神经网络原子间势 （NNIP） 模型的策略。在确认了高阶对称性（旋转阶数 𝐿）对扩展效率的影响后，他们也得出了增加模型参数的最佳方法。
图 2：EquiformerV2 在 OC20 2M 数据集上的消融研究结果。
提高具有组表示特征的 NNIP 模型能力的一种流行方法是增加表示的顺序。如果没有控制模型中可训练参数的总数，会在模拟中引入差异，这可能混淆 𝐿 对模型的性能影响。
为了明确增加的影响𝐿在模型性能上，并确定在 NNIP 模型中增加参数的最有效策略，团队将不同值的可训练参数数量标准化，并系统地添加进模型。
比较分析揭示了不同参数扩展策略的性能增益具有明显的层次结构。一旦跨模型的参数数量得到控制，许多模型与原始模型就会产生相当的误差。增加注意力机制的参数是最有利的，并且比简单地在所有组件中添加更多参数有着更实质性的突破。
EScAIP
为了避免冗杂的张量积，团队对旋转和平移不变的标量特征进行调整，以利用自然语言处理中优化的自我注意机制。这种处理使该模型比 EquiformerV2 等等变群表示模型的时间和内存效率更高。
图 3：高效缩放注意力原子间势 （EScAIP） 模型架构的图示。
通过引导模型进行精细的力大小预测，模型可以学习更好的系统表示，从而帮助它更准确地预测系统能量。力的大小是原子的局部属性，而能量预测是分子系统的全局属性。这可能解释了为何节点读出特征预测力的大小对于能量预测很有帮助。
甜蜜的教训与轻快的步伐
团队注意到，他们的调查路线遵循了一些苦涩教训的原则。专注于扩展和计算的策略往往优于那些试图将领域知识嵌入模型的策略。但是他们却相信这个教训是甜蜜的，因其使大规模培训大众化，并使其可供更广泛的社区使用。
原子系统比嵌入到模型中的特定领域信息要复杂得多。预定义的对称约束和手工制作的特征仅提供了这种复杂性的简单表示。后续改进最好伴随着相关的评估指标，使 NNIP 能够通过缩放获得表达能力来学习其余信息。
随着数据集的不断增长，在小型数据集上从头开始训练模型可能变得没有必要。虽然约束在非常小的数据制度中可能会提供一些帮助，但也可以利用预先训练的大型模型的表示作为在较小数据集上进行微调的起点。
除了专注于数据生成之外，其他技术在 NNIP 领域可能会越来越重要。其中包括模型蒸馏、通用训练和推理策略，这些策略与模型无关，可以应用于任何 NNIP，以及更好地与实验结果联系起来的方法。这表明更全面的策略对于 NNIP 的准确性与效用极为重要。
原文链接：https://arxiv.org/abs/2410.24169
相关代码：https://github.com/ASK-Berkeley/EScAIP