Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SO(3)-Equivariant Neural Networks for Learning Vector Fields on Spheres

Francesco Ballerin, Nello Blaser, Erlend Grong

TL;DR

This work develops SO(3)-equivariant neural networks for learning vector and tensor fields on the sphere, addressing both model and signal symmetries by performing group convolutions on SO(3) and employing a smoothing operator to retain equivariance while enabling expressive activations. The architecture supports scalar and vector fields through a UNet-like structure with spectral pooling, leveraging Wigner-D Fourier representations and carefully designed loss on S^2. Empirical results on ERA5 wind and temperature data show improved robustness to rotations and competitive performance relative to CNNs and spherical CNNs, with notable gains in rotated settings and vector-to-vector tasks. The approach provides a principled, geometry-aware alternative for geophysical prediction and vector-field autoencoding, highlighting the benefits of integrating group-theoretic convolutions with flexible nonlinearities in non-Euclidean domains.

Abstract

Analyzing vector fields on the sphere, such as wind speed and direction on Earth, is a difficult task. Models should respect both the rotational symmetries of the sphere and the inherent symmetries of the vector fields. In this paper, we introduce a deep learning architecture that respects both symmetry types using novel techniques based on group convolutions in the 3-dimensional rotation group. This architecture is suitable for scalar and vector fields on the sphere as they can be described as equivariant signals on the 3-dimensional rotation group. Experiments show that our architecture achieves lower prediction and reconstruction error when tested on rotated data compared to both standard CNNs and spherical CNNs.

SO(3)-Equivariant Neural Networks for Learning Vector Fields on Spheres

TL;DR

This work develops SO(3)-equivariant neural networks for learning vector and tensor fields on the sphere, addressing both model and signal symmetries by performing group convolutions on SO(3) and employing a smoothing operator to retain equivariance while enabling expressive activations. The architecture supports scalar and vector fields through a UNet-like structure with spectral pooling, leveraging Wigner-D Fourier representations and carefully designed loss on S^2. Empirical results on ERA5 wind and temperature data show improved robustness to rotations and competitive performance relative to CNNs and spherical CNNs, with notable gains in rotated settings and vector-to-vector tasks. The approach provides a principled, geometry-aware alternative for geophysical prediction and vector-field autoencoding, highlighting the benefits of integrating group-theoretic convolutions with flexible nonlinearities in non-Euclidean domains.

Abstract

Analyzing vector fields on the sphere, such as wind speed and direction on Earth, is a difficult task. Models should respect both the rotational symmetries of the sphere and the inherent symmetries of the vector fields. In this paper, we introduce a deep learning architecture that respects both symmetry types using novel techniques based on group convolutions in the 3-dimensional rotation group. This architecture is suitable for scalar and vector fields on the sphere as they can be described as equivariant signals on the 3-dimensional rotation group. Experiments show that our architecture achieves lower prediction and reconstruction error when tested on rotated data compared to both standard CNNs and spherical CNNs.

Paper Structure

This paper contains 31 sections, 2 theorems, 49 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Learning vector fields on the sphere
  3. Differential tensors on the sphere and equivariant functions
  4. Vector field on $S^2$, signal equivariance, and model equivariance
  5. Equivariant layers and group-convolutions
  6. Layers of our model
  7. The proposed network architecture
  8. Filters
  9. Nonlinearities / activation functions
  10. Spectral pooling
  11. UNet-like Architecture
  12. Loss function on $S^2$
  13. Equivariance by rotation
  14. Experiments
  15. Performance measure
  16. ...and 16 more sections

Key Result

Lemma B.1

Let $Y(\beta)$ denote the matrix corresponding to positive rotation around the $y$-axis by an angle $\beta$. Define a matrix $\Delta^l = D^l(Y(\frac{\pi}{2}))$.

Figures (3)

  • Figure 1: The proposed $\mathop{\mathrm{SO}}\nolimits(3)$-equivariant layer, mapping a $p$-equivariant signal to a $p$-equivariant signal. It consists of a group-convolution, inverse Fourier transformation ($\mathop{\mathrm{SO}}\nolimits(3)$), nonlinearity, Fourier transformation ($\mathop{\mathrm{SO}}\nolimits(3)$) and smoothing in $p$. A scalar field is encoded with $p=0$ while $p=1$ corresponds to a vector field.
  • Figure 2: The proposed $\mathop{\mathrm{SO}}\nolimits(3)$-equivariant UNet-type architecture. CONV layers as presented in Section \ref{['subsec:layers']} and as visualized in Figure \ref{['fig:layer']}. Yellow dotted lines (RES) correspond to the residual blocks. Pool and Unpool are performed in the spectral domain.
  • Figure 3: Comparison in performance on the test dataset (2022, 2023) between a convolutional UNet (CNN), our implementation of spherical CNN network (sCNN) , and our proposed architecture (Ours), trained on a variable number of years in the training set (years from 2000 to 2009). Our proposed architecture generalizes to unseen rotations even without data augmentation. Increasing the size of the training dataset does not improve the performance of a classic CNN on rotated test samples if data augmentation has not been used during training.

Theorems & Definitions (15)

  • Example 2.1
  • Remark 2.2
  • Remark 2.3
  • Remark 2.4
  • Example A.1: Functions
  • Example A.2: Vector fields
  • Example A.3: Symmetric tensors
  • Lemma B.1
  • proof
  • Remark B.2: Further simplifications
  • ...and 5 more