PaperlayerView exact revision on arXiv

Source-linked AI summary

Spin Spherical Fourier Neural Operators: Efficient Equivariant Learning for Vector and Spin Fields on the Sphere

Kevin Mulder, Matthew A. Price, Eirini Leftaki, Jason D. McEwen

arXiv:2611.04389v1cs.CVastro-ph.IMcs.LG

TL;DR

Spin fields such as wind and CMB polarisation change phase with local-frame rotations, while common scalar-channel spherical models do not respect general rotations. This paper extends SFNOs to arbitrary-spin fields using lifted convolutions computed in harmonic space. Its convolutional layers are equivariant to 10^-15, and its networks achieve state-of-the-art results on three new benchmarks.

  • Problem

    Common scalar spherical CNNs and SFNOs treat spin-field components as unrelated channels, which does not respect general rotations.

    Source: §1 · p. 1

  • Method

    S2FNOs extend SFNOs to arbitrary-spin fields using lifted spherical convolutions, with Mackey-function lifting as the framework and harmonic-space computation.

    Source: §4.3 · p. 4, §8 · p. 8

  • Results

    S2FNO networks achieve state-of-the-art results on three new benchmarks, and their convolutional layers are equivariant to 10^-15.

    Source: §6.2 · p. 6, §8 · p. 8

  • Takeaways & Limitations

    The three new benchmarks span classification, weather, and cosmology, with inputs and targets of definite spin, including spin 2.

    Source: §7.1 · p. 7

    Directional kernels, a single-spin spherical output, and rotational equivariance cannot be achieved simultaneously.

    Source: §3.2 · p. 3

Abstract

from arXiv · show

Many fields of physical interest on the sphere, such as wind, the polarisation of the cosmic microwave background and the shear of weak gravitational lensing, are spin fields, whose values change phase when the local frame rotates. We extend spherical Fourier neural operators (SFNOs) to fields of any spin, through lifted spherical convolutions that are exactly rotationally equivariant. Lifting spin fields to Mackey functions on the rotation group provides the conceptual framework, while computation is performed exactly and efficiently through their harmonic representation and by appealing to sampling theorems on the sphere and rotation group. We construct spin SFNOs (S$^2$FNOs) from these building blocks. We demonstrate empirically that rotational equivariance holds to machine precision for these layers. We achieve state-of-the-art performance on existing benchmarks, and introduce new benchmarks on which spin-aware models outperform their scalar counterparts.

1. Introduction

Spin fields encode quantities such as wind and polarisation relative to a local frame, but common scalar-channel models do not respect general rotations. The paper extends SFNOs to arbitrary spin and reports machine-precision equivariance and state-of-the-art results on new benchmarks.

Motivation

  • Standard scalar spherical CNNs and SFNOs treat spin-field components as unrelated channels, an approach that does not respect general rotations.Wind is one example of a field handled this way in data-driven weather models.Source: §1 · p. 1

Gap

  • Existing spin-equivariant networks can be computationally and parametrically costly, while equivariance is often not quantified and spin-field benchmarks are scarce.The group CNN of Ballerin et al. (2025) has about 2 million parameters on spherical MNIST, compared with 50 thousand for similar performance here.Source: §1 · p. 1

Contributions

  • S2FNOs extend SFNOs to fields of any spin, with rotationally equivariant convolutional layers and continuum equivariance for every activation.They are described through Mackey functions on the rotation group and computed in harmonic space.Source: §1 · p. 1
  • The authors introduce three public spin-field benchmarks spanning classification, weather, and cosmology, and report state-of-the-art performance on them.The benchmark inputs and targets have definite spin, including spin 2.Source: §1 · p. 2

2. Background and related work

The background defines spin fields and their harmonic representation, then situates SFNOs among spherical neural operators and equivariant networks. Prior spin-field approaches include harmonic filters and rotation-group lifting, but the cited network experiments use spins of at most 1.

2.1. Spin fields and their harmonic analysis

  • A spin-s field is defined relative to a local tangent frame, and its harmonic coefficients and Wigner representation are transformed exactly using sampling-theorem-based methods.Spin spherical harmonic transforms have complexity O(L^3), while Wigner transforms have complexity O(NL^3).Source: §2.1 · p. 2

2.2. Neural operators and equivariant networks on the sphere

  • SFNOs learn weights by spherical harmonic degree and represent continuous spherical operators, while prior spin-equivariant networks use harmonic filters or rotation-group operations.The cited prior network experiments use spins of at most 1; this paper's experiments include spin 2.Source: §2.2 · p. 2

3. Equivariant spin convolutions by lifting

The lifting framework represents spin fields as constrained functions on the rotation group and builds equivariant convolutions by lifting, convolving, and projecting. Harmonic-space forms provide spin-changing and directional operations, subject to a constraint on directional kernels and single-spin spherical outputs.

3.1. Lifting spin fields to the rotation group

  • A spin field lifts to an equivariant Mackey function on SO(3), whose harmonic representation occupies the single Wigner column n = -s.The adjoint projection returns the function to a spin field on the sphere.Source: §3.1 · p. 2

3.2. Lifted convolutions in harmonic space

  • The generalized lifted convolution lifts, convolves with a kernel on SO(3), and projects as specified by four lift-and-projection cases, all equivariant for square-integrable kernels.The spin spherical convolution uses one multiplier per degree and spans bounded equivariant complex-linear maps between fixed spins.Source: §3.2 · p. 3
  • A directional layer can be represented as a tuple of at most 2N - 1 spin-changing layers, and generalized lifted convolutions can recover several established spherical convolution forms.The paper restricts network filters through interpolated profiles, a small internal spin set, and a kernel diagonal on that set.Source: §3.2 · p. 3
  • Directional kernels, a spherical output of one spin, and rotational equivariance cannot all be achieved simultaneously.Directional information must remain on the rotation group or be represented as a tuple of spin fields.Source: §3.2 · p. 3

4. Spin spherical Fourier neural operators

S2FNOs combine spin spherical and rotation-group convolutions with equivariant nonlinearities and harmonic channel-and-spin mixing. The resulting blocks are assembled into classifier and dense-prediction networks.

4.1. Convolutional building blocks

  • The spin spherical convolution mixes spin fields by degree-specific multipliers without explicitly forming a lift, while the rotation-group convolution couples spins through nonlinearities on SO(3).The rotation-group implementation computes only the Wigner columns needed for projection, at O(|S|(L^3 + NL^2)) rather than O(NL^3).Source: §4.1 · p. 3

4.2. Equivariant non-linearities

  • For nonzero spins, spherical pointwise activations must commute with phase rotations, whereas pointwise maps on the lifted rotation-group functions are equivariant without that constraint.The paper uses a modulus nonlinearity for nonzero-spin fields on the sphere and standard activations for scalar fields.Source: §4.2 · pp. 3–4

4.3. S2FNO blocks

  • S2FNO blocks combine a convolution branch with a pointwise branch and use harmonic 1 × 1 convolutions to enrich channel and spin mixing.The harmonic maps mix channels and spins in coefficient space; the block has no identity residual connection.Source: §4.3 · p. 4, Figure 2 · p. 4

4.4. Networks

  • The paper builds classifiers and UNets from either convolutional or S2FNO blocks, with spherical or rotation-group constructions.Classifiers stack six blocks and pool rotation-invariant features; UNets use encoder-decoder skip connections at equal bandlimit and spin set.Source: §4.4 · pp. 4–5

5. Efficient representation, computation and parameterisation

S2FNOs describe spin-field layers as lifts to Mackey functions on SO(3), but use harmonic representations for computation. Their design handles spin coupling, exact transforms, and bandlimit-independent parameter counts without generally forming the lift.

5. Efficient representation, computation and parameterisation

  • Mackey functions provide the conceptual account of the layers, while harmonic representations enable efficient computation without forming the lift except when it adds expressivity.A spin-s lift occupies a single Wigner column and can be stored in O(L^2) spin spherical harmonic coefficients, rather than O(L^3) for a general SO(3) function.Source: §5 · p. 5, §5.1 · p. 5

5.1. Representation and spin mixing

  • Every bounded equivariant complex-linear map between spins is a spin spherical convolution; the rotation-group construction couples spins nonlinearly through the lift.Its group-convolution equivalent can be represented as an array of spin spherical convolutions, while the network uses only the array diagonal.Source: §5.1 · p. 5
  • The explicit lift enables general nonlinear spin coupling, whereas the spherical construction's phase-preserving nonlinearity couples no spins.On SO(3), one pointwise nonlinearity after lifting couples all spins.Source: §5.1 · p. 5
  • Real profiles make spin multipliers commute with antipodal reflection; classifiers use complex profiles, while real profiles may suit fields with antipodal symmetry.Representing fields as spin fields also avoids a restriction of scalar-field networks, which cannot represent maps that anticommute with the antipodal map.Source: §5.1 · p. 5

5.2. Computation

  • Sampling theorems make harmonic transforms exact for bandlimited inputs, and the spin spherical convolution costs O(|S|L^3), compared with O(NL^3) for a rotation-group convolution at azimuthal bandlimit N.The full SO(3) convolution costs O(L^4) when N = L; the rotation-group construction samples SO(3) only for its nonlinearity.Source: §5.2 · p. 5

5.3. Parameterisation

  • Interpolating multipliers from ten knots makes parameter counts independent of bandlimit and sampling grid, allowing a trained layer to apply on any grid supporting exact transforms at its bandlimit.The knots reduce the O(L) multipliers per spin-and-channel pair to ten.Source: §5.3 · pp. 5–6

6. Equivariance

The convolutional layers are exactly rotation-equivariant, and activation-free blocks and UNets reach transform round-off. The measured errors arise from activations and normalisation, with aliasing and nonzero-spin mean subtraction contributing distinct sources.

6.1. Sources of equivariance error

  • Pointwise activations introduce aliasing because their outputs are not bandlimited; group normalisation aliases, and sphere normalisation also creates structural error for nonzero spins.Subtracting a non-rotation-invariant mean causes the sphere-specific error; harmonic-space statistics without mean subtraction for nonzero spins would be exactly equivariant.Source: §6.1 · p. 6

6.2. Equivariance tests

  • The mean relative equivariance error compares rotated inputs and outputs over 100 random bandlimited signals and 100 random rotations at L = 16 in double precision.The activation-free blocks and UNets were tested separately at L = 32.Source: §6.2 · p. 6, Table 1 · p. 6
  • At L = 32, activation- and normalisation-free S2FNO blocks and UNets reach the transform round-off floor, with errors of order 10^-14.Linear layers attain errors of order 10^-15; the table caption specifies that these blocks and UNets omit activation and normalisation.Source: §6.2 · p. 6, Table 1 · p. 6
  • Activation errors are of order 0.1 to 1, while group-normalisation errors are of order 10^-2 on the sphere and 10^-3 on SO(3).The complex ReLU is an exception at initialisation, where it is the identity up to epsilon.Source: §6.2 · p. 6

7. Spin benchmarks and experiments

The section introduces three spin-field benchmarks and compares spin-aware networks with scalar counterparts and published results. Across the tasks, spin-aware models show strong performance, including lower errors than scalar counterparts on spin-changing prediction, wind forecasting, and lensing.

7.1. New benchmarks

  • The three public benchmarks span classification and regression on synthetic, atmospheric, and cosmological fields, addressing the scarcity of public spin-field benchmarks.Their inputs and targets have definite spin, and spin spherical MNIST includes any spin.Source: §7.1 · pp. 6–7
  • Spin spherical MNIST synthesises fields of any spin and evaluates classification and spin-changing dense prediction under training on unrotated and testing on randomly rotated data.It includes harmonic spin-1 and spin-2 fields alongside the Sobel spin-1 field.Source: §7.1 · p. 7
  • The global wind benchmark predicts 18:00 UTC wind at 50 hPa from that day's 00:00, 06:00, and 12:00 UTC winds, testing within-year interpolation rather than later-year generalisation.It uses ERA5 reanalysis from 1980 to 2022 at L = 128.Source: §7.1 · p. 7
  • The lensing benchmark reconstructs full-sky spin-0 convergence from noisy spin-2 shear using simulated training, validation, and evaluation maps.The simulations use GLASS and Stage IV survey shape noise.Source: §7.1 · p. 7

7.2. Spin spherical MNIST

  • All four classifiers exceed 99% accuracy on harmonic spin-1 and spin-2 fields, while the best Sobel-field classifier comes within 0.15% of the published GsCNN accuracy.The classifiers generalise to rotations not seen in training.Source: §7.2 · p. 7
  • Across all four spin-changing MNIST mappings, each rotation-group UNet has fewer parameters and lower error than its spherical counterpart.Conv-SO(3) has the lowest errors from images to fields, while S2FNO-SO(3) has the lowest errors from fields to images; Figure 3 shows both predictions following the rotated target frame.Source: §7.2 · p. 7, Table 3 · p. 8, Figure 3 · p. 8

7.3. Global wind forecasting

  • The wind spin UNet reduces MSE by 44.9% with 34.9% fewer parameters and improves every other metric, outperforming persistence on every metric.The scalar UNet is worse than persistence on every metric except correlation.Source: §7.3 · pp. 7–8

7.4. Weak gravitational lensing

  • The lensing spin UNet reduces MSE by 37.9% with 13.4% fewer parameters and improves every other metric over a scalar UNet.The spin model lifts spin-2 shear and returns spin-0 convergence; the scalar model takes the shear components as scalar channels.Source: §7.4 · p. 8

8. Conclusions

S2FNOs extend spherical Fourier neural operators to fields of any spin, with exactly rotationally equivariant convolutional layers, and achieve state-of-the-art results on three new benchmarks.

8. Conclusions

  • S2FNOs extend SFNOs to any spin, with exact rotational equivariance in each convolutional layer and in the continuum for every activation.Source: §8 · p. 8
  • Mackey-function-described S2FNOs are computed in harmonic space and are efficient in representation, computation, and parameters, with far fewer parameters than full rotation-group convolution networks.Source: §8 · p. 8
  • For a held-out wind-forecast day, the scalar UNet leaves larger, coherent patches of vector error than the spin UNet.Source: Figure 4 · p. 8
  • In mass mapping, the scalar UNet misses convergence tails and leaves large-scale residual structure, whereas the spin residual is noise-like.Source: Figure 5 · p. 8
  • The networks achieve state-of-the-art results on three new benchmarks.Table 4 covers held-out 50 hPa wind forecasting and L=256 mass mapping averaged over 600 maps.Source: §8 · p. 8, Table 4 · p. 8
  • Convolutional layers are equivariant to 10^-15; blocks and UNets without activations or normalisation are equivariant to round-off.Source: §8 · p. 8

Impact Statement

The paper’s supporting sections develop the harmonic and lifting framework for spin fields, establish equivariant convolution and non-linearity constructions, and discuss constraints from discretisation and parity. The impact statement identifies geophysical and astrophysical modelling as the intended application area.

Impact Statement

  • The work aims to advance machine learning for geophysical and astrophysical fields on the sphere, with improved spin-field treatment potentially benefiting weather, climate, and cosmology.The authors state that LLMs assisted manuscript writing, while the theory was worked out manually.Source: Impact Statement · p. 9, Declaration on the use of large… · p. 9

A.2. Spin fields, Mackey functions and the lift

  • A spin field’s frame-dependent values are represented on SO(3) by a lift whose choice of section encodes the local frame and its spin phase.Changing the section changes the coordinate expression by the spin-dependent frame law.Source: §A.2 · pp. 12, 14
  • The lift intertwines spin rotations with left translations on SO(3), connecting the sphere’s spin-field representation to the group’s regular representation.The induced spin representation is identified with the spin rotation operator.Source: §A.1 · p. 12, §A.2 · p. 14

A.3. Lifted convolutions

  • Lifted convolutions are bounded and SO(3)-equivariant for arbitrary kernels, and their harmonic form selects the input and output Wigner columns associated with the spins.At finite band-limit, an assembled kernel is square integrable; an output may be represented by at most 2N − 1 spin-changing components.Source: §A.3 · p. 15
  • For a fixed input/output spin pair, equivariance and spherical output preclude directional kernels beyond the frame phases prescribed by the spins.The kernel is radial apart from those phases.Source: §A.3 · pp. 15–16

B. Spherical layers as neural operators

  • Spherical-layer neural-operator parameters can be independent of sphere sampling, but resolution-independent parameters alone do not guarantee exact sampled non-linearity equivariance or learned weights at unseen degrees.The comparison includes the S2FNO’s real profile in degree for each spin pair, parameterised by 10 knots.Source: B. Spherical layers as neural operators · p. 16, Table 5 · p. 16

C. Equivariant non-linearities

  • On the sphere, a pointwise activation for nonzero spin must commute with every frame-rotation phase; the modulus non-linearity satisfies this constraint, while scalar fields allow any pointwise map.On SO(3), pointwise maps commute with rotations because rotations translate the function’s argument, enabling a lifted activation followed by projection.Source: §C.1 · p. 16, §C.2 · pp. 16–17

C.3. Discretisation; C.4. Parity

  • Sampling a generally non-bandlimited activation can alias its higher-frequency content, which oversampling reduces; meanwhile, spin fields escape a parity restriction that applies to scalar networks built from isotropic multipliers and pointwise maps.The scalar restriction prevents representing maps that anticommute with the antipodal map, such as the stated vorticity advection example.Source: §C.3 · p. 17, §C.4 · p. 17

D. Network architectures

The appendix specifies the S2FNO layers and their assembly into classifiers and UNets for the evaluated tasks. It also notes that group normalisation is not exactly rotationally equivariant for nonzero-spin fields on the sphere.

D.1. Layers

  • The implemented S2FNO block runs a spin spherical convolution and harmonic 1×1 convolution in parallel, then applies a second harmonic convolution, modulus activation, and group normalisation.Only the rotation-group construction uses lift-and-sum.Source: Figure 9 · p. 18
  • Group normalisation yields approximate rather than exact rotational equivariance for nonzero-spin fields, with reported errors of order 10^-2 on the sphere and 10^-3 on the rotation group.The appendix notes that a different harmonic-statistics normalisation without mean subtraction for nonzero spin would be exactly equivariant.Source: §D.1 · p. 19

D.2. Blocks and spin spherical MNIST networks

  • The four block families combine spherical or rotation-group convolutions with optional lift-and-sum, harmonic channel/spin mixing, activation, and normalisation.The UNet input layer maps image spin 0 or vector-field spin 1 to the internal set {−1, 0, 1} in the rotation-group construction.Source: §D.2 · p. 19
  • The MNIST UNets use a four-level encoder and two-level decoder, while S2FNO UNets retain bandlimit L = 32 throughout.Decoder levels concatenate encoder features at matching bandlimit and spin set.Source: §D.2 · pp. 19–20, Figure 11 · p. 20

D.3. Networks for wind forecasting and weak lensing

  • For mass mapping, the five-level spin UNet lifts spin-2 shear to internal spins and returns spin-0 convergence, while the scalar comparator uses DISCO convolutions and residual layers.The architectures span bandlimits L = 256 to L = 16.Source: §D.3 · p. 20, Figure 12 · p. 21

E. Equivariance tests

The tests assess rotational equivariance with a round-trip comparison of outputs for rotated and unrotated inputs. Linear blocks and UNets reach the transform round-off floor, while trained-network equivariance with learned activations and normalisations is not measured.

E. Equivariance tests

  • The round-trip test compares outputs for rotated and unrotated inputs across 100 test signals and 100 Haar-drawn rotations.Signals have complex Gaussian random harmonic coefficients bandlimited at L; outputs are compared after transforming and rotating in harmonic space.Source: E. Equivariance tests · p. 21
  • Pointwise activations commute with rotations in the continuum, but aliasing contributes to their measured errors; trained-network equivariance with learned activations and normalisations is not measured.The complex ReLU begins as the identity up to 10^-8, with error of order 10^-9; that error is the round-trip error and would grow once its gain and offset are learned.Source: E. Equivariance tests · p. 21
  • Every tested linear block and UNet has equivariance error of order 10^-14, the round-off floor on the sphere; the positive control yields order 10^-1.The test covers models with several spins and channels, including UNets tested in both 0 → 1 and 1 → 0 directions.Source: E. Equivariance tests · p. 21, Table 7 · p. 22

F. Datasets

The datasets cover spin-field prediction, classification, wind forecasting, and weak-lensing mass mapping. Their construction and evaluation settings define the scope of the benchmarks.

F.1. Spin spherical MNIST

  • The spin spherical MNIST variants include projected scalar digits, grid-derived Sobel spin-1 fields, and harmonic spin-1 and spin-2 fields.The Sobel derivatives are combined as Xθ + iXφ, while the harmonic construction rotates scalar harmonic coefficients before inverse spin synthesis.Source: §F.1 · p. 22
  • Prediction uses 60,000 samples at L = 32 with an 85/15 training–held-out split; rotated held-out samples use uniform Euler angles, not Haar-distributed rotations.Classification uses a separate regenerated split and reports results on randomly rotated test samples after training on unrotated samples.Source: §F.1 · p. 22

F.2. Global wind forecasting

  • Wind forecasting uses 50 hPa ERA5 inputs at 00:00, 06:00, and 12:00 UTC to predict 18:00 UTC wind, represented as the spin-1 field u + iv.Data are sampled at L = 128; eight days spanning all seasons are held out from every year for evaluation.Source: §F.2 · p. 23

F.3. Weak gravitational lensing

  • Recovering convergence from shear is ill-posed in practice, so the setup calls for priors learned from simulations.The shear determines convergence only on degrees ℓ ≥ 2.Source: §F.3 · p. 23
  • Weak-lensing maps pair simulated shear with convergence; the dataset has 10,000 training/validation pairs and 600 evaluation pairs from the same pipeline.The maps use Stage IV survey shape noise and are sampled at L = 256.Source: §F.3 · p. 23

G. Training details and additional results

The appendix reports training and evaluation details for the experiments and provides additional prediction examples. It also notes limits on run-to-run estimates and on the selection of prediction checkpoints.

G. Training details and additional results

G.1. Spin spherical MNIST

  • For spin spherical MNIST prediction, the reported checkpoint minimizes loss on the rotated held-out set, which is also used to evaluate the reported errors.The prediction UNets are trained in double precision for up to 2,000 epochs; the displayed examples are fixed in advance rather than selected by error.Source: §G.1 · p. 23
  • The spin spherical MNIST prediction figure shows Conv-S2 and Conv-SO(3) UNet outputs for spin-0-to-spin-1 Sobel prediction and the reverse task.The examples are at L = 32 and show the first three test digits under one rotation.Source: Figure 13 · p. 24, §G.1 · p. 24

G.2. Wind forecasting

  • Wind UNets are trained for 200 epochs, and Table 8 reports configurations and metrics for 50 hPa forecasting on held-out days.The reported checkpoint has the lowest loss on those held-out days; the figure presents a one-day forecast and regional error comparisons.Source: §G.2 · p. 25, Table 8 · p. 25, Figure 14 · p. 26

G.3. Weak gravitational lensing

  • On weak-lensing maps, the scalar UNet residual retains high-amplitude large-scale structure, while the spin UNet reconstruction follows the target distribution tails more closely.Table 9 reports configurations and metrics averaged over 600 evaluation maps; the figure compares reconstructions and residuals for one map.Source: §G.3 · p. 25, Table 9 · p. 25, Figure 15 · p. 27
Loading 2611.04389v1…