Differentiable and accelerated spherical harmonic and Wigner transforms

要約

科学と工学の多くの分野では、球面多様体で定義されたデータに遭遇します。
球面データのモデリングと分析では、高度な球面調和関数変換が必要になることが多く、機械学習やその他の微分可能なプログラミング タスクのための勾配の効率的な計算がますます必要になります。
球 $\mathbb{S}^2$ と回転群 $\text{SO}(3)$ 上の一般化フーリエ変換、つまりそれぞれ球面調和関数とウィグナー変換の高速かつ微分可能な計算のための新しいアルゴリズム構造を開発します。
我々は、高調波次数まで安定であり、極めて並列化可能な Wigner $d$ 関数の計算のための再帰的アルゴリズムを提案します。
これを分離可能な球面変換と緊密に結合することにより、最新のハードウェア アクセラレータ (GPU など) の高スループット コンピューティングに適した、非常に並列化可能な構造を示すアルゴリズムが得られます。
また、勾配を効率的に計算できるように、自動および手動のハイブリッド微分アプローチも開発します。
私たちのアルゴリズムは、S2FFT ソフトウェア コードの JAX 微分可能プログラミング フレームワーク内に実装されています。
等角サンプリングや HEALPix サンプリングなど、球体の多数のサンプリングがサポートされています。
サンプリング定理が認められる球面サンプリングでは、計算誤差は機械精度程度です。
代替の C コードに対してベンチマークを行った場合、最大 400 倍の加速が観察されました。
さらに、複数の GPU に分散する場合、アルゴリズムが高度に並列化されバランスが取れているため、GPU の数が増加しても最適に近い線形スケーリングが実現します。
十分な数の GPU へのアクセスが提供されるため、変換は前例のない効果的な線形時間計算量を示します。

要約(オリジナル)

Many areas of science and engineering encounter data defined on spherical manifolds. Modelling and analysis of spherical data often necessitates spherical harmonic transforms, at high degrees, and increasingly requires efficient computation of gradients for machine learning or other differentiable programming tasks. We develop novel algorithmic structures for accelerated and differentiable computation of generalised Fourier transforms on the sphere $\mathbb{S}^2$ and rotation group $\text{SO}(3)$, i.e. spherical harmonic and Wigner transforms, respectively. We present a recursive algorithm for the calculation of Wigner $d$-functions that is both stable to high harmonic degrees and extremely parallelisable. By tightly coupling this with separable spherical transforms, we obtain algorithms that exhibit an extremely parallelisable structure that is well-suited for the high throughput computing of modern hardware accelerators (e.g. GPUs). We also develop a hybrid automatic and manual differentiation approach so that gradients can be computed efficiently. Our algorithms are implemented within the JAX differentiable programming framework in the S2FFT software code. Numerous samplings of the sphere are supported, including equiangular and HEALPix sampling. Computational errors are at the order of machine precision for spherical samplings that admit a sampling theorem. When benchmarked against alternative C codes we observe up to a 400-fold acceleration. Furthermore, when distributing over multiple GPUs we achieve very close to optimal linear scaling with increasing number of GPUs due to the highly parallelised and balanced nature of our algorithms. Provided access to sufficiently many GPUs our transforms thus exhibit an unprecedented effective linear time complexity.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google