提出首个真正渐进加速的克莱布什-戈尔丹张量积算法,显著提升3D神经网络效率。
Asymptotically Fast Clebsch-Gordan Tensor Products with Vector Spherical Harmonics
- 基于向量球谐函数推广快速傅里叶卷积,构建新型张量球谐函数
- 将计算复杂度从原生 $O(L^6)$ 降至 $O(L^4/log^2 L)$,接近理论下界 $O(L^4)$
- 适用于需高精度对称性保持的3D分子或物理模拟场景
E(3)等变神经网络在多种3D建模任务中表现优异。其核心操作是张量积,用于不同特征类型间的交互。由于该操作计算开销大,已有大量工作致力于加速。然而,近期研究表明,多数提速源于表达能力下降而非真正的算法改进。已有方法如Gaunt张量积虽能实现渐进加速,但不完整且遗漏大量交互。本文首次提出完整的、真正具有渐进优势的克莱布什-戈尔丹张量积算法。对于全量级CGTP,我们将时间复杂度从朴素的 $O(L^6)$ 降低至 $O(L^4\log^2 L)$,逼近理论下界 $O(L^4)$。我们首先揭示,推广基于快速傅里叶的卷积自然导出先前提出的Gaunt张量积。为解决反对称性问题,我们将信号从标量推广至不可约表示值信号,从而引入张量球谐函数,并证明了其广义的Gaunt公式。最后,我们证明仅需至向量值信号即可恢复原始Gaunt张量积中缺失的交互项。
原文摘要 · Abstract (English)
$E(3)$-equivariant neural networks have proven to be effective in a wide range of 3D modeling tasks. A fundamental operation of such networks is the tensor product, which allows interaction between different feature types. Because this operation scales poorly, there has been considerable work towards accelerating this interaction. However, recently \citet{xieprice} have pointed out that most speedups come from a reduction in expressivity rather than true algorithmic improvements on computing Clebsch-Gordan tensor products. A modification of Gaunt tensor product \citep{gaunt} can give a true asymptotic speedup but is incomplete and misses many interactions. In this work, we provide the first complete algorithm which truly provides asymptotic benefits Clebsch-Gordan tensor products. For full CGTP, our algorithm brings runtime complexity from the naive $O(L^6)$ to $O(L^4\log^2 L)$, close to the lower bound of $O(L^4)$. We first show how generalizing fast Fourier based convolution naturally leads to the previously proposed Gaunt tensor product \citep{gaunt}. To remedy antisymmetry issues, we generalize from scalar signals to irrep valued signals, giving us tensor spherical harmonics. We prove a generalized Gaunt formula for the tensor harmonics. Finally, we show that we only need up to vector valued signals to recover the missing interactions of Gaunt tensor product.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。