论文

FlashCart:等变原子间势的快速笛卡尔张量积

FlashCart: Fast Cartesian Tensor Products for Equivariant Interatomic Potentials

模型推理推理加速

摘要

机器学习的原子间势将原子模拟扩展到电子结构方法可访问的长度和时间尺度之外。然而,等变架构的计算成本限制了它们在实践中可以表示的局部相关性,从而限制了它们可实现的准确性。在这里,我们介绍 FlashCart,它通过将生成的 GPU 内核与递归构建等变特征并在每一步将它们压缩到固定宽度的架构相结合,使高阶相关性变得经济实惠。我们用独立的笛卡尔分量表示张量积,并象征性地简化它们及其导数,产生通常优于优化的球形对应物的融合内核。然后我们证明,增加相关阶数比增加宽度、深度或张量秩更有效地提高准确性。在 SPICE-MACE-OFF 上,FlashCart 模型推进了测量的精度-效率前沿:具有 $5.6$ 百万参数的模型比具有 $189$ 百万参数的变压器实现了更低的能量和力误差,并且 $10\times$ 的推理速度更快。