Differentiable and accelerated wavelet transforms on the sphere and ball

TL;DR

This paper introduces highly efficient, differentiable wavelet transforms on the sphere and ball, enabling faster computations and new data-driven analysis techniques for signals on these domains, with open-source implementations.

Contribution

The authors develop new accelerated, differentiable wavelet transforms on the sphere and ball, with significant speedups and automatic differentiation capabilities for advanced analysis.

Findings

Up to 300-fold acceleration on the sphere

Up to 21800-fold acceleration on the ball

Open-source JAX libraries for the transforms

Abstract

Directional wavelet dictionaries are hierarchical representations which efficiently capture and segment information across scale, location and orientation. Such representations demonstrate a particular affinity to physical signals, which often exhibit highly anisotropic, localised multiscale structure. Many physically important signals are observed over spherical domains, such as the celestial sky in cosmology. Leveraging recent advances in computational harmonic analysis, we design new highly distributable and automatically differentiable directional wavelet transforms on the $2$-dimensional sphere $\mathbb{S}^2$ and $3$-dimensional ball $\mathbb{B}^3 = \mathbb{R}^+ \times \mathbb{S}^2$ (the space formed by augmenting the sphere with the radial half-line). We observe up to a $300$-fold and $21800$-fold acceleration for signals on the sphere and ball, respectively, compared to existing software, whilst maintaining 64-bit machine precision. Not only do these algorithms dramatically accelerate existing spherical wavelet transforms, the gradient information afforded by automatic differentiation unlocks many data-driven analysis techniques previously not possible for these spaces. We publicly release both S2WAV and S2BALL, open-sourced JAX libraries for our transforms that are automatically differentiable and readily deployable both on and over clusters of hardware accelerators (e.g. GPUs & TPUs).