FlashCart 如何让等变原子间势更快更准？

FlashCart 用笛卡尔内核替代球谐耦合

等变原子间势把原子模拟扩展到电子结构方法无法触及的长度和时间尺度，但其计算成本限制了它在实践中能表示的局域关联，进而限制了可达到的精度。FlashCart 直接针对这一成本：它把张量积写成独立的笛卡尔分量，对乘积及其导数做符号化简，并把结果融合进生成的 GPU 内核。论文报告这些内核往往优于经过优化的球谐实现。

该架构递归地构建等变特征，并在每一步把它们压缩到固定宽度，这正是让高阶关联变得可负担的原因。论文的核心主张并不是笛卡尔数学在运算次数上更省；作者指出，运算次数相比稀疏球谐耦合并无算术优势。运行速度的优势来自完整的内核实现，包括内存访问和边贡献的聚合。

提升关联阶比加宽通道更省算力

论文把提高关联阶与增加通道宽度、消息传递深度或最大张量秩进行比较。精度用归一化误差衡量，它把能量和力的平均绝对误差各自除以一个固定参考值后再取平均。在对数坐标轴上，拟合斜率描述误差随运行时间或内存的相对变化。

在所测试的配置中，关联阶的运行时间斜率在 -6.5 到 -2.6 之间，而其他架构选择的斜率在 -1.5 到 -0.8 之间，因此提高关联阶带来的每单位新增运行时间的误差下降最大。在其余架构固定的情况下，把关联阶从 2 提高到 5 使归一化误差下降 19% 到 60%，运行时间增加 4% 到 37%，峰值 GPU 内存最多增加 1.7%。处于观测到的精度—效率前沿上的模型大多使用 4 和 5 的关联阶。

更高张量秩放大关联阶的收益

关联阶的收益取决于允许的张量秩。当关联阶梯中的方向张量和中间张量都被限制在秩 1 时，把阶从 2 提高到 5 使能量误差下降 20%，运行时间增加 4%。在最大秩为 4 时，同样的变化使能量误差下降 58%，运行时间增加 37%。因此更高的张量秩允许提高阶带来更大的相对误差下降，但计算成本也更高。

提高关联阶带来的能量和力误差相对下降，通常随消息传递块增多或最大传播秩提高而变小，但随通道宽度变化不大。由于提高阶也会增加参数，作者还比较了参数规模大致固定在 93 万左右的模型，做法是随阶提高而减小通道宽度。在最大边秩和乘积秩为 3 时，一个 64 通道的 5 阶模型在 3.94 微秒/原子下达到 8.29 meV/Å 的力 MAE，而 108 通道的 2 阶模型为 5.49 微秒/原子下 11.53 meV/Å。

生成内核在高阶下胜过球谐基线

在 1024 个节点、128 个特征通道、每节点 128 个近邻的边卷积微基准中，在 30 种配置的网格上，FlashCart 的前向、向量—雅可比积以及向量—雅可比积的向量—雅可比积平均分别比 cuEquivariance 快 1.6 倍、1.1 倍和 1.9 倍。相比 OpenEquivariance，它平均快 1.9 倍、2.1 倍和 3.1 倍，不过在秩为 1 和 2 的某些配置中 OpenEquivariance 更快。在向量—雅可比积的向量—雅可比积阶段，FlashCart 的峰值 GPU 内存也约为两个基线的一半。

对于 4096 个节点、128 个特征通道、最大张量秩为 3 和 4 的节点局域乘积，差距随阶提高而扩大。在秩 3、阶 4 时，前向评估对路径约简乘积耗时 0.40 毫秒，对完整阶梯耗时 2.07 毫秒，而 cuEquivariance 为 280.61 毫秒，分别约快 695 倍和 136 倍。FlashCart 在每个测试秩下都能达到 6 阶，而某些 cuEquivariance 收缩无法构建或其内核启动失败。作者指出，这些运行时间差异也反映了关联的构建方式不同：cuEquivariance 为完整的对称收缩分配学习权重，而 FlashCart 逐步施加加权乘积，并在每一步之后只保留不超过所选最大值的张量秩。

560 万参数模型胜过 1.89 亿参数 transformer

在 SPICE-MACE-OFF 上，最大的 FlashCart 模型有 560 万参数，能量和力 MAE 分别为 0.072±0.002 meV/原子和 1.99±0.04 meV/Å。对 1000 个金刚石原子评估能量和力耗时 21.5 微秒/原子，峰值 GPU 内存为 5.2 MiB/原子。这优于 PET-L 的 0.106 meV/原子和 2.52 meV/Å，尽管参数少得多且运行速度约快十倍。它也接近 DPA4-Plus 的 1.80 meV/Å 力 MAE，同时优于其 0.091 meV/原子的能量 MAE，评估速度快 7.7 倍、内存少 3.1 倍。

这些计时使用各模型自身的截断半径。在金刚石中，FlashCart 每原子处理 158 条边，而 PET-L 为 70 条，使其每条边的运行时间优势约为 23 倍。作者指出，这种归一化考虑了不同的近邻数，但比较中仍包含架构差异。为分离生成内核的贡献，论文用同一个 560 万模型检查点，把它们与对应的 PyTorch 张量积和边卷积实现进行比较。在两种实现都编译的情况下，生成内核把推理时间从 175.5 微秒/原子降到 21.5 微秒/原子，把训练基准运行时间从 442.6 微秒/原子降到 82.2 微秒/原子，分别快 8.2 倍和 5.4 倍。

230 万参数模型在 85 种元素上保持竞争力

论文还评估了一个在 MAD 上训练的 230 万参数 FlashCart 模型，MAD 包含不到 10 万个结构、覆盖 85 种元素，并与 280 万参数的 PET-MAD 势及其他机器学习原子间势比较。在报告力误差的全部六个基准上，FlashCart 的力 MAE 都低于 PET-MAD，在七个基准中的六个上能量 MAE 更低。在完整 MAD 测试集上，FlashCart 的能量和力 MAE 分别为 11.4±0.1 meV/原子和 63.9±0.5 meV/Å，而 PET-MAD 为 14.7 meV/原子和 72.9 meV/Å。

对 1000 个金刚石原子评估能量和力时，FlashCart 和 PET-MAD 分别耗时 8.4 和 45.1 微秒/原子，峰值 GPU 内存分别为 2.0 和 3.3 MiB/原子。作者称这些结果确立了 FlashCart 作为通用势的竞争力架构，在多样的材料和分子环境中达到与最先进模型相当的精度。

水模拟能量守恒但密度偏高

为展示对原子模拟的适用性，作者用精度比较中的五个模型研究体相液态水，在 LAMMPS 中结合 Kokkos 加速模拟 256 个水分子。所有模型的能量漂移在 -0.002 到 0.008 meV/原子/纳秒之间，围绕均值的能量涨落标准差为 0.012 到 0.013 meV/原子，最大绝对偏差低于 0.056 meV/原子。在 NPT 模拟中，模型预测的水密度随温度下降，且始终高于实验值，这与 MACE-OFF 及一些在同一电子结构水平上训练的相关变体所报告的高估一致。

论文还定位了编译的一项成本。在不编译模型的情况下，用生成的 Triton 内核替换 PyTorch 张量积和边卷积，使同一模型推理快 26.8 倍、训练基准快 23.8 倍。用 Triton 内核编译模型会进一步降低运行时间，但增加训练时的峰值 GPU 内存。在 512 个原子、禁用 CUDA 图的情况下比较内存分配，能量和力误差合并进损失后，GPU 分配在编译时增加 5.21 GiB，不编译时增加 5.72 GiB；在参数梯度计算期间，编译时峰值又比前一数值高出 2.26 GiB，而不编译时为 0.17 GiB。因此额外的峰值内存出现在模型编译时的参数梯度计算阶段，两种情况使用相同的生成 Triton 内核。

先调关联阶，再加宽或加深

与决策相关的结果是一个优先级排序，而不是单一的精度数字。在固定的算力或内存预算下，论文的拟合斜率和受控比较表明，提高关联阶比增加通道宽度、消息传递深度或最大张量秩，每单位新增运行时间带来的归一化误差下降更多。当张量秩也允许增大时收益最大：在最大秩为 4 时，把阶从 2 提高到 5 使能量误差下降 58%，运行时间增加 37%；而在秩 1 时，同样的阶提升只使能量误差下降 20%，运行时间增加 4%。

有两条边界限定这一排序。第一，相对收益随消息传递深度或最大传播秩提高而缩小，因此阶的优势在不同架构间并不一致。第二，运行时间和内存比较使用各模型自身的截断半径并包含架构差异，且水模拟显示密度高于实验值，因此精度—效率前沿应被理解为在论文基准上测得的结果，而非普遍保证。在这些限制之内，论文最清晰的指引是：当目标是在给定成本下降低能量和力误差时，关联阶而非宽度或深度，是首先要调整的架构旋钮。

本文引用的文献