[论文解读] Orb-v3: atomistic simulation at scale

Bradley J. Rhodes, Sander Vandenhaute+5 more|ArXiv.org|Apr 8, 2025

Machine Learning in Materials Science被引用 16

一句话总结

Orb-v3 引入了一套通用、可扩展的原子间势能集合，在保持高精度的同时极大提升了速度和内存使用效率；非保守直接模型在声子和性质基准测试上可达到或超过最先进水平。

ABSTRACT

We introduce Orb-v3, the next generation of the Orb family of universal interatomic potentials. Models in this family expand the performance-speed-memory Pareto frontier, offering near SoTA performance across a range of evaluations with a >10x reduction in latency and > 8x reduction in memory. Our experiments systematically traverse this frontier, charting the trade-off induced by roto-equivariance, conservatism and graph sparsity. Contrary to recent literature, we find that non-equivariant, non-conservative architectures can accurately model physical properties, including those which require higher-order derivatives of the potential energy surface. This model release is guided by the principle that the most valuable foundation models for atomic simulation will excel on all fronts: accuracy, latency and system size scalability. The reward for doing so is a new era of computational chemistry driven by high-throughput and mesoscale all-atom simulations.

研究动机与目标

• 激发需要将准确性与大规模、高通量模拟相结合的通用机器学习势的需求。
• 将 Orb-v3 家族呈现为在性能-速度-内存帕累托前沿上具可扩展性、接近状态-of-the-art 的解决方案。
• 研究设计选择（保守性、邻居数量限制、数据集）如何影响准确性和可扩展性。
• 展示提升旋转不变性并为实际分子动力学工作流提供不确定性估计的技术。

提出的方法

• 提出 Orb-v3 作为一组具有与 Orb-v2 相同基础架构但为速度和内存效率而编译和扩展的模型族。
• 同时使用直接（非保守）和保守的力/势表达来探索帕累托前沿。
• 引入 equigrad，一种基于梯度的正则化，用以在保守模型中引入旋转等变性。
• 在大规模数据集上训练（OMat24 AIMD 子集；为兼容性使用 mpa），并将保守模型向直接模型进行蒸馏，以提高高阶导数的精度。
• 实现高效的图构建、受限的邻居数量，以及基于GPU的最近邻算法以实现高吞吐量。
• 提供本征式逐原子置信度头，用于预测力误差以进行主动学习和数据筛选。

Figure 1 : The Pareto frontier for a range of universal Machine Learning Interatomic Potentials. The $K_{SRME}$ metric assesses a model’s ability to predict thermal conductivity via the Wigner formulation of heat transport [ 31 ] and requires accurate geometry optimizations as well as second and thi

实验结果

研究问题

• RQ1 非等变、非保守的架构是否能够在提供卓越速度和内存效率的同时，在热力学和声子相关预测方面达到具有竞争力的水平？
• RQ2 保守性、邻居数量限制和训练数据集如何影响通用 MLIPs 在多种性质上的准确性与可扩展性？
• RQ3 在 Orb-v3 模型中，几何优化、高阶导数精度与分子动力学稳定性之间有哪些权衡？
• RQ4 equigrad 正则化是否能够在不损失性能的前提下提升旋转不变性和下游对称性工作流？

主要发现

• Orb-v3-direct-20-omat 与 Orb-v3-direct-inf-omat 实现显著的速度/内存提升，能够实现每秒数百次前向计算，并在大规模系统上实现<1 秒级。
• 非保守的直接模型在若干基准测试上可达到或超过最先进水平，包括声子与热导等任务。
• Equigrad 正则化显著提升旋转不变性和对称性引导工作流的鲁棒性。
• Conservative-inf-omat 模型在多项物理性质基准上达到最高精度，同时保持比多数竞争者更快的运行时间。
• Orb-v3 模型显示出出色的可扩展性，direct-20-omat 在 100k 原子测试中需要 32.8 GB GPU 内存，且完成时间 <0.5 s，凸显吞吐量的跃升。
• 本征逐原子置信度估计与力误差相关，可帮助主动学习与数据整理。

Figure 2 : Speed + max GPU memory allocated on an NVIDIA H200 for the computation of energies, forces and stress. The batch size is fixed to 1, but we vary the number of atoms across the subplots. Relative times are computed with respect to the fastest model: orb-v3 Direct (20 neighbors) . Times inc