arXiv:2501.13986cs.LGcs.AI2025-01被引 29

优化旋转等变网络的核心计算,显著提升分子模拟推理速度。

An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks

  • 通过静态分析优化GPU共享内存使用,分解计算为小核函数提升并行度。
  • 相比e3nn快10倍,对MACE模型在FP64下提速达6.2倍。
  • 适配分子建模、量子化学等需高精度等变网络的研究者使用。

旋转等变图神经网络在空间深度学习任务中表现优异,训练数据效率高,且在原子间势能计算中推理速度远超传统方法。其核心是克莱布希-戈登(CG)张量积,该操作将两个稠密特征向量与高度结构化的稀疏张量相乘,生成稠密输出,但重复百万次导致性能瓶颈。本文提出一种高效的GPU稀疏核生成器,通过编译时静态分析精细管理有限的共享内存,减少全局内存读写。将张量积分解为多个小核函数,使操作完全进入寄存器,可发射长指令流以最大化指令级并行。通过融合CG张量积与后续图卷积,减少中间存储和全局内存流量。还提供CG张量积的梯度优化核及用于预测原子受力的高阶偏导数新恒等式。在闭源cuEquivariance包上实现最高1.3倍加速,在广泛使用的e3nn上提速10倍;在FP64精度下,对MACE化学基础模型的推理速度提升最高达6.2倍。

原文摘要 · Abstract (English)

Rotation equivariant graph neural networks, i.e. networks designed to guarantee certain geometric relations between their inputs and outputs, yield state of the art performance on spatial deep learning tasks. They exhibit high data efficiency during training and significantly reduced inference time for interatomic potential calculations compared to classical approaches. Key to these models is the Clebsch-Gordon (CG) tensor product, a kernel that contracts two dense feature vectors with a highly-structured sparse tensor to produce a dense output vector. The operation, which may be repeated millions of times for typical equivariant models, is a costly and inefficient bottleneck. We introduce a GPU sparse kernel generator for the CG tensor product that provides significant speedups over the best existing open and closed-source implementations. Our implementation achieves high performance by carefully managing the limited GPU shared memory through static analysis at model compile-time, minimizing reads and writes to global memory. We break the tensor product into a series of smaller kernels with operands that fit entirely into registers, enabling us to emit long arithmetic instruction streams that maximize instruction-level parallelism. By fusing the CG tensor product with a subsequent graph convolution, we reduce both intermediate storage and global memory traffic over naive approaches that duplicate input data. We also provide optimized kernels for the gradient of the CG tensor product and a novel identity for the higher partial derivatives required to predict interatomic forces. Our kernels offer up to 1.3x speedup over NVIDIA's closed-source cuEquivariance package, as well as 10x speedup over the widely-used e3nn package. In FP64 precision, we offer up to 6.2x inference-time speedup for the MACE chemistry foundation model over the original unoptimized version.

等变网络分子模拟GPU优化张量积

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。