bispectrum: Selective G-Bispectra Made Practical

Many machine learning tasks are invariant under the action of a group G of transformations: signal classification can be invariant under translations, image classification under 2D rotations, and spherical-image classification under 3D rotations. The G-bispectrum is a principled complete invariant of a signal (retaining all all signal's information up to the group action) with proven benefits in machine learning and as a pooling layer in deep networks. However, its deployment has been hampered by high computational cost and a patchwork of group-specific implementations. We present bispectrum, an open-source, fully unit-tested PyTorch library that implements selective G-bispectra for seven different group actions, as differentiable modules that can be directly incorporated into machine learning pipelines and deep learning architectures. For finite groups G, selectivity reduces the computational cost from O(|G|2) to O(|G|). For planar rotations, we leverage the disk bispectrum. For spherical 3D rotations, we introduce an augmented selective bispectrum at band-limit L which reduces the cost from O(L3) to Theta(L2) coefficients. We profile the entire library (for which we implemented various compute optimizations), showing that it delivers near-exact G-invariance with its selective G-bispectra computed in sub-millisecond time on GPU (up to commonly used bandlimits). We evaluate the benefits of incorporating G-bispectra as pooling layers into deep learning architectures on three classical benchmark datasets --comparing against norm pooling, gated pooling, Fourier-ELU pooling, max pooling, and (non-equivariant) data-augmented convolutional baselines. Results show that G-bispectra consistently outperform alternatives in the low-data, moderate-capacity regime.