arXiv:2501.13986v4 宣布类型：替换-交叉 摘要: 旋转不变图神经网络，即设计用于确保存在某些几何关系的输入和输出的网络，在空间深度学习任务中达到了最先进的性能。与经典方法相比，这些模型在训练过程中表现出高度的数据效率，并且在原子间势能计算中的推理时间显著减少。这些模型的核心是Clebsch-Gordon (CG) 张量积，这是一种将两个密集特征向量与一个高度结构化的稀疏张量相结合以产生一个密集输出向量的核。该操作在典型等变模型中可能需要重复数百万次，是一个代价高昂且效率低下的瓶颈。我们引入了一种GPU稀疏核生成器来实现CG张量积，其在最佳现有开源和闭源实现的基础上提供了显著的加速。我们的实现通过在模型编译时进行静态分析，精细管理有限的GPU共享内存，从而减少对全局内存的读写操作，从而达到高性能。我们将其张量积分解为一系列较小的核，操作数完全容纳在寄存器中，这使我们能够生成最大限度地利用指令级并行性的长算术指令流。通过将CG张量积与后续的图卷积相结合，我们减少了两种简单方法中由于重复输入数据而产生的中间存储和全局内存流量。我们还提供了CG张量积梯度的优化核以及用于预测原子间力所需的更高阶偏导数的新型恒等式。我们的核在NVIDIA的闭源cuEquivariance包上提供了高达1.3倍的加速，在广泛使用的e3nn包上提供了高达10倍的加速。在FP64精度下，我们针对MACE化学基础模型提供了高达6.2倍的推理时间加速，相比于原版本。