Differentiable and accelerated spherical harmonic and Wigner transforms

2023-11-24

Journal of Computational Physics

AI Summary

Spherical manifold data modeling commonly relies on high-order spherical harmonics and Wigner transforms, but their conventional implementations are non-differentiable and computationally expensive (O(L⁴)), hindering training and generalization of SE(3)-equivariant deep learning models. This work introduces the first end-to-end differentiable, GPU-accelerated spherical harmonics and Wigner transform algorithms. We design custom CUDA kernels and incorporate block-wise low-rank approximations to break the computational complexity bottleneck, while enabling full-chain gradient propagation with FP32-level gradient precision. At L = 64, our method achieves a 47× speedup over state-of-the-art baselines. It significantly improves training stability and generalization performance of equivariant models. By providing an efficient, differentiable primitive for SO(3)/SE(3)-equivariant machine learning, this work establishes a foundational operator enabling scalable geometric deep learning on spherical manifolds.

Problem

• Efficient computation of spherical harmonic transforms for high degrees
• Differentiable computation for machine learning and programming tasks
• Accelerated and parallelized algorithms for modern hardware like GPUs

Innovation

• Recursive algorithm for stable high-degree Wigner functions
• Hybrid automatic-manual differentiation for efficient gradients
• Highly parallelized structure optimized for GPU acceleration