一种高效稀疏核生成器用于O(3)等变深层网络
机器学习
2025-05-12 v4 人工智能
摘要
旋转等变图神经网络即设计为保证其输入与输出之间 Certain 几何关系的网络，在空间深度学习任务中表现出优异性能。它们在训练期间具有高数据效率，并且在计算原子间相势时显著减少推理时间，相较于传统方法效果更佳。关键在于Clebsch-Gordan（CG）张量积，这是一种与两个稠密特征向量进行张量积的核，通过与高度结构稀疏张量进行协商以产生稠密输出向量。该操作在典型等变模型中可能重复执行数百万次，是昂贵且效率低下的瓶颈。我们引入了一种GPU稀疏核生成器，用于CG张量积，相较于现有最佳开源和闭源实现实现了显著加速。我们的实现通过在模型编译时进行静态分析，精确管理有限的GPU共享内存，最小化对全局内存的读写。我们将张量积分解为一系列较小的内核，其操作数完全适合寄存器，从而发射长期的算术指令流以最大化指令级并行度。通过将CG张量积与随后的图卷积融合，我们减少了中间存储和全局内存流量，优于采用复制输入数据的笨拙方法。我们还提供了CG张量积梯度的优化内核以及一种用于预测原子间力的更高阶导数的新颖公式。我们的内核在NVIDIA闭源cuEquivolution软件包上实现了最高1.3倍加速，以及在广泛使用的e3nn软件包上实现了10倍加速。在FP64精度下，我们为MACE化学基础模型在推理时间上实现了最高6.2倍加速。
引用
@article{arxiv.2501.13986,
title = {An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks},
author = {Vivek Bharadwaj and Austin Glover and Aydin Buluc and James Demmel},
journal= {arXiv preprint arXiv:2501.13986},
year = {2025}
}
备注
To appear in the Proceedings of the 2025 SIAM Conference on Applied and Computational Discrete Algorithms. 15 pages, 10 figures, 4 tables