Learning without relearning physical symmetry: an SO(3)-equivariant convolutional network (CNN) on the sphere

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

Machine Learning: Science and Technology Crossmark RECEIVED

1 October 2026 REVISED

dd Month yyyy

PAPER

Learning without relearning physical symmetry: an SO(3)-equivariant convolutional network (CNN) on the sphere Joshua Williamson1,∗ , Niall Jeffrey2,1 , Lorne Whiteway1 1 Department of Physics & Astronomy, University College London, Gower Street, London, WC1E 6BT, UK 2 Department of Physics & King’s Institute for Artificial Intelligence, King’s College London, Strand, London WC2R

arXiv:2610.23821v1 [stat.ML] 16 Oct 2026

2LS, United Kingdom ∗ Author to whom any correspondence should be addressed. E-mail: joshua.williamson.21@ucl.ac.uk Keywords: rotation equivariance, spherical convolutional neural networks, simulation-based inference

Abstract Many of the most important datasets in the physical sciences live natively on the sphere, and these data often carry no preferred position or orientation. Using methods that do not reflect this symmetry leads to far higher simulation requirements, which can be catastrophic for scientific analyses (forcing reliance on cheap, imperfect models that bias results). We present a novel spherical implementation of a gauge equivariant mesh architecture that is specifically designed for the Hierarchical Equal Area isoLatitude Pixelization (HEALPix) of the sphere. The SO(3)-equivariant convolution is built from steerable filters, parallel transport and spin-weighted features. The architecture supplies the exact symmetry, so the network does not spend training data learning it. The convolution operations are constructed to give provably exact equivariance, given a band-limit. Our full CNNs can achieve SO(3)-invariance error at per-cent-level and lower, which is controlled through well-understood hyperparameters. On benchmarks, our technique gives the most accurate dense gradient regression and competitive cosmic microwave background foreground removal and denoising. The improved simulation efficiency matters most in simulation-based inference (SBI), where high-fidelity simulations are expensive and neural compressors require large training sets, so a fixed budget forces analyses with less accurate simulators. In cosmology lensing SBI, our technique yields the tightest posteriors at matched simulation count. This approach is a novel method that will enable new convolutional neural network analyses, including generative and inference tasks, for processes and data on the sphere.

1 Introduction Many of the most important datasets in the physical sciences live natively on the sphere. The cosmic microwave background (CMB), a probe that is central to our standard model of cosmology, is measured over the full celestial sphere (Planck Collaboration et al., 2020); weak gravitational lensing surveys map the large-scale structure of the Universe (Jeffrey et al., 2025); and weather forecasting and climate science operate on fields covering the Earth’s sphere (Bonev et al., 2023). Increasingly, extracting scientific information from these high-dimensional datasets relies on large suites of numerical simulations, most directly in simulation-based inference (SBI). High-fidelity simulations are often extremely expensive, however, forcing practical analyses to rely on cheaper, imperfect models in order to generate sufficient training data. Any mismatch between these models and reality propagates directly into the inferred posteriors. Reducing the number of simulations required can reduce this dependence on imperfect models. The more efficient the training, the less imperfect the models, as the simulation budget is spent on fidelity rather than simulation count. We provide one route to such a reduction: by encoding known physical structure as an inductive bias, the learning problem requires fewer examples. For spherical data, that structure is symmetry: these

1

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

m

-0.02

= 0.206

m

0.12

Williamson et al

Noiseless CMB

= 0.417

-0.02

0.12

-300

K

Observation, 100 GHz

300

-300

K

1000

(a) Weak-lensing convergence κ for two values of Ωm . (b) Noiseless CMB map and a corresponding 100 GHz mock observation.

Figure 1: Spherical data from cosmology for the benchmark tasks. data often have no preferred position or orientation, motivating SO(3)-equivariant architectures; here SO(3) is the group of three-dimensional rotations. Planar convolutional neural networks (CNNs) are poorly matched to this geometry: projections of the sphere onto the plane distort areas and distances, and translational weight sharing corresponds to no symmetry of the sphere. A network without rotation-equivariant filters must learn the same feature at every position and orientation, seeing it repeated across the training data through augmentation or natural variability. SO(3)-equivariant representations remove this redundancy, sharply reducing training data requirements. When every training example is an expensive simulation, exploiting this symmetry offers a direct route to doing more science with fewer simulations. For SBI this matters twice over, since the simulation suite is both the training set and the sole source of the likelihood. Convolutions that are both rotationally equivariant and scalable have proven difficult to construct (Table 1). Group-based spherical CNNs built on Wigner transforms (Cohen et al., 2018; Esteves et al., 2018; Kondor et al., 2018; Cobb et al., 2021) are exactly equivariant but costly, and their scalable spin-weighted descendants (Esteves et al., 2020, 2023) remain bound to spectral transforms on full sky maps. Spatially localized alternatives, hybrid discrete-continuous (DISCO) convolutions (Ocampo et al., 2023) and graph constructions on the Hierarchical Equal Area isoLatitude Pixelization (HEALPix) (Perraudin et al., 2019; Defferrard et al., 2020) aim to be approximately rotationally equivariant, but in practice are not truly equivariant. CNNs have been used already in SBI on probes of weak lensing and CMB maps (Jeffrey et al. 2025; Williamson & Makinen et al., 2026; Thomsen et al. 2026; Wolz et al. 2023), and global weather forecasting (Hu et al., 2025), yet few have attempted architectures that enforce strict SO(3) equivariance. In this paper we demonstrate that exact SO(3) equivariance can be carried onto the HEALPix grid at practical cost, and that in a realistic weak lensing SBI task it yields tighter posteriors from the same simulation suite compared to planar, graph and axisymmetric alternatives. 2 SO(3) convolution method This section gives a self-contained account of the convolution layers used in this work. The ingredients are established: convolution on the rotation group (Cohen and Welling, 2016; Cohen et al., 2018), its evaluation by quadrature with the filter kept continuous (Ocampo et al., 2023), filters built from circular harmonics (Freeman and Adelson, 1991; Worrall et al., 2017), and the treatment of local frames on a curved surface by parallel transport (de Haan et al., 2021). Our contribution is their combination on the HEALPix grid in a form that is exactly, rather than approximately, equivariant; derivations beyond those given here can be found in these papers and references therein. 2.1 Rotations as positions and orientations Let S2 be the unit sphere, with points labelled by colatitude θ ∈ [0, π] and longitude ϕ ∈ [0, 2π), and let SO(3) be the group of rotations of three-dimensional space. In the zyz Euler convention every rotation can be written as R = Z(ϕ)Y (θ)Z(γ), where Z(·) and Y (·) denote rotations about the z and y axes. The first two factors carry the north pole ẑ to the point x = (θ, ϕ), and γ ∈ [0, 2π) is a rotation about the axis through x, so a rotation specifies a position x on the sphere together with an orientation γ at that position. We write g = (x, γ) for the rotation with these coordinates and call g a pose.

2

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

(a) pose g = (x, )

Williamson et al

(c) frame transport angle Q(x)

(b) relative pose g 1h z

z e(xj )

e(x)

j

xj

x g

i

e(xi ) i

g

u

z e(Qx)

e(x)

j

xi

ij v

x

Q

Qe(x)

Qx

h

Figure 2: Coordinates of the SO(3) convolution: a pose, the relative pose of two points, and the transport angle. The angle γ is measured from a reference direction in the tangent plane at x, which the Euler convention takes to be the meridian direction, i.e. local north. No such direction can be chosen continuously over the whole sphere, and this one is undefined at the two poles. This suffices for integration and, since no HEALPix pixel centre lies at a pole, for every sample point that we use. The arbitrariness of north is the origin of the transport angle of Section 2.3. Integration over the group uses the Haar measure µ, the unique measure up to normalization that is invariant under rotation, µ(QS) = µ(S) for every rotation Q and every measurable set S. In Euler coordinates dµ = sin θ dθ dϕ dγ, the product of the area element dµS2 = sin θ dθ dϕ on the sphere and the uniform measure dγ on the circle of orientations. A signal on the group is a function F (x, γ) and a signal on the sphere a function f (x); channel indices are suppressed. 2.2 Convolution on the rotation group Given a signal F on SO(3) and a filter ψ, which is also a function on SO(3), their group convolution is Z H(g) ≡ (F ⋆ ψ)(g) = F (h) ψ(g −1 h) dµ(h). (1) SO(3)

Here g is the pose at which the output H is evaluated, h is the integration variable running over all input poses, and g −1 h is the pose of h relative to g, i.e. the rotation that carries the frame g to the frame h. Evaluating the filter on the relative pose is what makes equation (1) equivariant. For any rotation Q define the rotated signal (Q · F )(h) = F (Q−1 h); substituting h = Qh′ and using the invariance of µ, Z ((Q · F ) ⋆ ψ)(g) = F (h′ ) ψ(g −1 Qh′ ) dµ(h′ ) = (F ⋆ ψ)(Q−1 g) = (Q · H)(g), (2) so rotating the input rotates the output by the same rotation (Cohen and Welling, 2016). This is equivariance, the desired property. Equation (1) maps signals on SO(3) to signals on SO(3) and is the form of every layerof the convolutional neural network, except for the first layer. The first layer receives a signal f on the sphere, and its lift to the group is Z H(g) = f (x) ψ(g −1 x) dµS2 (x), (3) S2

where ψ is now a function on the sphere and g −1 x is the point x expressed in the frame g (Cohen et al., 2018; Esteves et al., 2018). The output depends on the orientation of g because the filter is anisotropic, and equivariance follows as in equation (2). 2.3 Relative pose in coordinates We refer the reader to Fig. 2 to aid geometric understanding. To evaluate ψ(g −1 h) we need coordinates for the relative pose. Let g = (xj , γ) and h = (xi , γ ′ ). The separation θij ∈ [0, π] is the length of the shorter great-circle arc joining xi and xj . This is unique unless the points are antipodal, though our filters have a cutoff θc < π meaning that this never arises. 3

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

Writing ê(x) for the unit tangent vector at x pointing along the meridian towards the north pole, the bearing αj→i ∈ [0, 2π) is the angle at xj from ê(xj ) to the direction of the arc towards xi . The bearing αi→j is defined in the same way at xi . These are two are independent quantities. The bearings depend on the choice of north, but the differences v = αi→j − γ ′

u = αj→i − γ,

(4)

do not: u is the direction of xi measured from the reference direction of the frame g, and v the direction of xj measured from that of the frame h. Since g −1 h moves the pole to colatitude θij and azimuth u in the frame g, with the remaining frame rotation given by v (up to an offset absorbed into the filter), ψ(g −1 h) is a function ψ(θij , u, v) of these three numbers. A rotation Q carries x to Qx and carries ê(x) to a tangent vector at Qx which in general is not ê(Qx); the two differ by an angle χQ (x), the transport angle (Cohen et al., 2019; de Haan et al., 2021). In the coordinates of Section 2.1 the rotation therefore acts as Q · (x, γ) = (Qx, γ + χQ (x)),

(5)

and since every bearing at x is measured from the same ê(x), it shifts by the same χQ (x) as γ does. The differences in equation (4) are unchanged, as is θij , which is why the integrand of equation (1) is invariant and the convolution equivariant. 2.4 Filter basis The filter is a function on the cap θ ≤ θc , periodic in u and in v, which we expand as ψ(θ, u, v) =

P M X X

N X

cpmn Rp (θ) eimu einv .

(6)

p=1 m=−M n=−N

The Rp , p = 1, . . . , P , are fixed radial profiles vanishing for θ > θc (we use Zernike radial polynomials), m and n are the angular orders in u and v, truncated at M and N , and the complex coefficients cpmn are the only learnable parameters; smoothness at θ = 0 requires Rp (0) = 0 whenever m = ̸ 0. These are the circular harmonic, or steerable, filters of planar rotation-equivariant networks (Freeman and Adelson, 1991; Worrall et al., 2017) transferred to the sphere, adopted in place of the interpolated filters of Ocampo et al. (2023) for one reason: by equation (4), eimu = e−imγ eimαj→i , so changing the output orientation γ multiplies each term by a phase, and the filter is steered to any orientation exactly, without resampling. 2.5 Sampling and quadrature 2 Position is sampled at the Npix = 12Nside HEALPix pixel centres xi (Górski et al., 2005), Nside being the resolution parameter, and orientation at Nγ equispaced angles γa = 2πa/Nγ , a = 0, . . . , Nγ − 1; the product is the equivolumetric SO(3) grid of Yershova et al. (2010), and a discrete set of values F (xi , γa ). The integral in equation (1) is replaced by the weighted P signal is theP sum i wi (2π/Nγ ) a over these points, with weights wi for the sphere and 2π/Nγ for the circle of orientations, while the filter and the group action are kept continuous. This is the discrete-continuous approach of Ocampo et al. (2023), and it is what allows exact equivariance to be recovered rather than approximated. The orientation samples integrate each Fourier mode exactly: for integer q with |q| < Nγ , Nγ −1 1 X iqγa e = δq0 . Nγ a=0

(7)

For the sphere we require the wi to integrate exactly every function whose spherical harmonic expansion contains √ no degree above a band limit Lw , i.e. to satisfy the moment conditions P w Ȳ (x ) = 4π δℓ0 δm0 for all ℓ ≤ Lw and |m| ≤ ℓ, where Yℓm are the spherical harmonics. i ℓm i i HEALPix pixels have equal area, so the natural weights wi = 4π/Npix satisfy these conditions only approximately. We solve for quadrature weights using an initialisation of 4π/Npix for all pixels (Keiner and Potts, 2008) at the band limit Lw = 2Nside , which produces small residuals of order 10−15 .

4

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

A signal on the group is 2π-periodic with respect to orientation and may be written as a Fourier series, N X F (x, γ) = F̂n (x) einγ , n=−N

(8)

Nγ −1

1 X F̂n (x) = F (x, γa ) e−inγa , Nγ a=0 where F̂n (x) is the orientation component of order n and N the orientation band limit; the second expression is exact, by equation (7), provided Nγ > 2N . We store signals as their components rather than their samples (compare Esteves et al., 2020), and take the range of the filter order n in equation (6) to be the same N . The component F̂0 is the average of F over orientation; it transforms as a scalar field on the sphere, so its average over the sphere is rotation invariant and provides the invariant readout of the network, as explained in Section 2.8. 2.6 The discrete layers Inserting equations (6) and (8) into the discretized equation (1), with u and equation (4), the P v from ′ sum over the input orientation samples γa′ contains the factor (2π/Nγ ) a′ ei(n −n)γa′ , where n′ is the order of the input component and n the order of the filter in v. By equation (7) this equals 2π if n′ = n and vanishes otherwise, since |n′ − n| ≤ 2N < Nγ , so the orientation integral can be evaluated in closed form. This is how we achieve exact SO(3)-equivariance without onerous computational demands. The output of the convolution, in real-space, is give as H(xj , γ) =

M X

e−imγ 2π

P N X X

cpmn ×

wi Rp (θij ) eimαj→i einαi→j F̂n (xi ).

(9)

i: θij ≤θc

p=1 n=−N

m=−M

X

Comparing with equation (8), H is band limited in orientation and its component of order k, with |k| ≤ M , is the coefficient of eikγ , i.e. the terms with m = −k. The convolution output, in Fourier-space, is calculated via Ĥk (xj ) = 2π

P N X X

cp,−k,n ×

X

wi Rp (θij ) e−ikαj→i einαi→j F̂n (xi ).

(10)

i: θij ≤θc

p=1 n=−N

This is the SO(3) to SO(3) layer. We note that we take M = N so that the output can be passed to the next layer. The lifting layer is the case noted after equation (3): a spherical input f is the signal whose only non-zero component is F̂0 = f , so equation (10) retains only the n = 0 term and the coefficients cp,−k,0 . The innermost sum over neighbouring pixels contains no learnable parameters and depends only on the HEALPix geometry, which we regenerate on the fly rather than store. Pointwise non-linearities act on the Nγ orientation samples given by the first expression in equation (8), after which the components are recovered by the second expression. 2.7 Equivariance of the discrete layers The continuous convolution of equation (1) is equivariant by construction. Here we show that the discrete layer of equation (10) inherits this, exactly, whenever the two quadrature rules of Section 2.5 reproduce the integral. There are two conditions. The first relates to the orientation sum, in that the integrand contains frequencies up to 2N , so by equation (7) it is integrated exactly provided Nγ > 2N , which we ensure by our choice of Nγ . The second relates to the sum over the sphere. At fixed output pixel j, the integrand is the product of the input component, band limited at some Lf , and the filter factor, which we take to be band limited at some effective Lψ , so the sum equals the integral provided Lf + Lψ ≤ Lw . When both these conditions hold, the layer satisfies equation (2) exactly. It is worth seeing how this works in the components, since this is what allows layers to be stacked. By equation (5), (Q · F )(Qx, γ) = F (x, γ − χQ (x)), and comparing Fourier coefficients in equation (8), \ (Q · F )n (Qx) = e−inχQ (x) F̂n (x). (11) Each component therefore transforms as a spin-weighted field of order n, picking up a phase set by the transport angle. Under the conditions above the sum over i in equation (10) equals the corresponding integral, so the rotated quantities can be substituted directly. The phase e−inχQ (xi ) acquired by the input component is then cancelled by the phase einχQ (xi ) acquired by the back

5

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

Table 1: Comparison of convolution methods for spherical data. Columns 1–2 are described in Section 3.8. Columns 3–5 are validation RMSE for CMB foreground removal and denoising, dense gradient regression and weak lensing parameter estimation respectively; ⟨FOM⟩ is the mean joint posterior figure of merit for (Ωm , S8 ) inference. We use Nγ = 8, achieving per-cent-level SO(3) invariance error at Nside = 256 – our error can be further reduced to 10−4 by introducing a spherical harmonic band limit (we leave the 1.39% value as a simpler direct comparison with the other methods). Method

Invariance SO(3) err. (%)

Planar

107.22

|∇δ| RMSE

{Ωm , S8 } RMSE

{Ωm , S8 } ⟨FOM⟩

0.71±0.01 7.490×10−1

4.37×10−2

4.862×10−2

561

−1

−2

−2

1724

2.291×10−2

3467

−2

5026

Time (ms)

CMB RMSE

DISCO

2.36

1.93±0.02

DeepSphere

71.56

54.86±0.02 1.755×10−1 9.80×10−3

This work

1.39

4.212×10

−1

27.91±1.13 2.261×10

4.08×10

4.43×10

−3

2.623×10

1.651×10

bearing αi→j , which is only possible because the two carry the same order n; the phase from the forward bearing αj→i depends only on the output pixel and factors out, so the output transforms as equation (11) requires. In other words, the orthogonality of the orientation modes that collapsed the orientation integral in Section 2.6 is also what makes the output an appropriate input to the next layer. This is the kernel constraint for rotation-equivariant features, solved for the plane by Weiler and Cesa (2019) and on curved surfaces by de Haan et al. (2021); here it is not imposed on the filter but falls out of the construction. The sampling conditions above are analogous to those of the SO(3) sampling theorems (Kostelec and Rockmore, 2008; McEwen et al., 2015), with HEALPix and the weights of Section 2.5 playing the role of the rule for the polar angle. These conditions apply to the convolution alone. Pointwise non-linearities, activation functions, broaden the spectrum of the signal in both orientation and position, so that the signal is no longer band limited. This is likely the source of the residual per-cent-level error in Table 1, and it can be mitigated by raising Nγ and the pixel resolution. We have tested this empirically – decreasing Nγ increases the error as expected. 2.8 Invariance for regression/classification tasks For tasks like denoising, the convolutional neural network with equivariant layers, is itself equivariant. For tasks like regression or classification, where the network outputs a vector (so not matching the input spherical geometry), it does not make sense to consider equivariance of the entire network. Instead, we want a network F network to be invariant under a transformation Q for an input S: F(QS) = F(S). To achieve this, the final network architecture’s layers, or head, must be constructed to preserve SO(3)-invariance. We achieve this through global average pooling of each channel (i.e. averaging all elements in a channel) in advance of the final multi-layer perception (MLP) – this pooling erases any position or rotation information, as desired. 3 Benchmarking and data 3.1 Simulation-based inference For a network summary of a weak lensing map the likelihood p(d | θ) has no closed form, but a simulator can draw d ∼ p(d | θ). Simulation-based inference (SBI) learns the posterior from such simulated pairs (Cranmer et al., 2020). Density estimation in O(Npix ) dimensions is not feasible, so a network first compresses each map to a low-dimensional summary t = F (d). This compressor is where the convolution of Section 2 enters. We use neural posterior estimation (Papamakarios and Murray, 2016): a masked autoregressive flowPqφ (θ | t) (MAF; Papamakarios et al., 2017) is fitted to the simulated pairs (θi , ti ) by minimizing − i ln qφ (θi | ti ). The fitted flow is amortized, yielding a posterior for any new observation with no further simulation and no MCMC. 3.2 Dense gradient-magnitude regression task Each input is a Gaussian random field δ on S 2 , generated in harmonic space from a smooth power spectrum
Cℓ = A (ℓ + 10−3 )−α exp − (ℓ/ℓbl )2 ,

C0 := 0,

(12)

with α =P 1.8, ℓbl = 2 Nside , and the amplitude A fixed in closed form by requiring Var[δ] = ℓ 2ℓ+1 4π Cℓ = 1, so every map has unit variance. Maps are band-limited at ℓmax = 2 Nside , 6

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

Planar CNN DISCO This work DeepSphere

0.85 0.80

S8

IOP Publishing

0.75 0.70 0.65

0.2

0.3 m

0.4

0.70 0.75 0.80 0.85 S8

Figure 3: Cosmological parameter Ωm –S8 posteriors from simulation-based inference with lensing data, comparing a CNN from this work with a planar CNN, DISCO and DeepSphere compressors (i.e. embeddings) using the same number of expensive training simulations. The poor results from models not optimised for spherical data (especially the Planar CNN) given a fixed simulation budget forces a user to rely on imperfect simulators to increase the simulation count – our method avoids this by including a spherical symmetry as inductive bias. the same pband limit as our HEALPix quadrature. The target is the smoothed gradient magnitude |∇δ| = (∂θ δ)2 + (∂ϕ δ)2 , with the partial derivatives taken analytically in harmonic space, rather than by finite differencing. 3.3 Cosmic Microwave Background denoising We evaluate on the CMB-ML benchmark (Amato et al., 2025), which provides simulated Planck-like sky maps at nine frequencies (Nside = 256): a lensed CMB realization combined with Galactic and extragalactic foreground emission and correlated, frequency-dependent instrumental noise. The task is to recover the underlying CMB temperature map from these nine contaminated observations (Fig. 1b), a spherical analogue of denoising under structured, spatially correlated corruption rather than i.i.d. pixel noise. 3.4 Cosmological parameters from weak lensing maps via simulation-based inference We use the Gower Street suite of full-sky simulations (Jeffrey et al., 2025), from which we build a weak lensing convergence map at Nside = 256 for each realization (Fig. 1a), with a single source plane and a fixed source redshift; the construction is described in Section 3.5. Unlike the pixel-wise targets of the previous two tasks, the regression target here is a pair of numbers: the matter density Ωm and the clustering amplitude S8 , which is the combination of parameters that weak lensing constrains most tightly. The summary given to the flow is the two-dimensional point estimate t = (Ω̂m , Ŝ8 ) returned by the regression network under its mean squared error loss, so the density estimator is conditioned on two numbers per map rather than on the κ map itself. We fit a single MAF, with four networks in

7

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

the ensemble, using two autoregressive transforms and eight hidden features per transform, implemented with the sbi package (Tejero-Cantero et al., 2020). It is trained on the parameter-summary pairs from the simulations, with the 78 maps used for the NDE validation loss. This choice comes from the limited number of simulations; as no overfitting is observed in any architecture, this does not induce a bias in our results. The joint figure of merit (FOM) in Table 1 is computed from the resulting average FOM of the 78 held out validation posteriors. 3.5 Weak lensing convergence maps This section gives the details of the data used for the cosmological parameter regression task of Section 3.4. The Gower Street suite is a set of 791 full-sky N -body simulations of a seven-parameter νwCDM cosmology, run with PKDGRAV3 (Potter et al., 2017) and described in Jeffrey et al. (2025). Each simulation provides pixelized matter overdensity shells δ(n̂, χ) along the observer’s lightcone, where n̂ is a direction on the sky and χ is comoving distance. From these shells we build a Born-approximated convergence map for a single source plane fixed at z = 0.8, Z 3Ωm H02 χs χ(χs − χ) δ(n̂, χ) κ(n̂) = dχ, (13) 2c2 χs a(χ) 0 with χs the comoving distance to the source plane, a the scale factor, H0 the Hubble constant and c the speed of light. Every map is then downgraded to Nside = 256. The two regression targets are the matter density Ωm and the clustering amplitude S8 ≡ σ8 (Ωm /0.3)0.5 . The five remaining parameters varied across the suite, namely w, Ωb h2 , mν , h and ns , are not predicted by the network and are therefore marginalized over implicitly. 3.6 Networks We compare the convolution of Section 2 against three alternatives. The architectures and network configurations for the weak lensing parameter estimation task are given in Section 3.7. The network of this work uses band limits P = 4, M = N = 1 and Nγ = 8 fibre samples, with radial profiles Rp,|m| taken as Zernike radial polynomials of order |m| + 2p in the haversine distance, multiplied by a quartic taper at the cutoff θc = 0.015590 rad, and all computations in single precision. DeepSphere: The HEALPix grid is treated as a graph, with pixel centres as vertices and edge weights given by a Gaussian kernel of the Euclidean distance between the k nearest neighbours, and the convolution is a polynomial in the resulting graph Laplacian (Defferrard et al., 2016; Perraudin et al., 2019; Defferrard et al., 2020). It costs O(Npix ) and pools along the HEALPix nested pixelization hierarchy, its filters are isotropic and the graph Laplacian produces asymptotic equivariance only as Npix → ∞, so equivariance is approximate at finite resolution. DISCO: The signal and the quadrature are discretized on the sphere, while the filter and the rotation acting on it stay continuous (Ocampo et al., 2023), as in equation (10), and we use the DISCO implementation of Bonev et al. (2023), where we port the convolutions to a HEALPix sampling scheme. Rotations are restricted to the quotient space SO(3)/SO(2), with a non-directional axisymmetric filter, which is SO(3) equivariant. Planar CNN: The sphere is unrolled onto a single equirectangular projection whose rows are the 4Nside − 1 iso-latitude HEALPix rings and whose columns are the 4Nside pixels of the equator, so that near-pole rings holding as few as four distinct pixels are stretched across every column. An ordinary two-dimensional convolution stack is applied to this raster, periodic in longitude, zero padded in latitude and read out through a solid-angle weighted pool. It encodes no rotational symmetry and heavily distorts the true spherical geometry. Every architecture is trained either to convergence or with early stopping on the validation loss to prevent overfitting. We use an 80/20 training-validation split for all datasets, apart from the weak lensing maps of Section 3.4, where the split is 90/10. 3.7 Network architectures This section gives the parameter budget of the four networks compared on the weak lensing task of Section 3.4, as totals in Table 2 and layer by layer in Tables 3–6. All four use 16 hidden channels at a single resolution, Nside = 256, and read out the two parameters (Ωm , S8 ) through a linear head. Operations carrying no learnable parameters are omitted: activations, pooling, normalization by quadrature, and the fixed geometric quantities of each convolution, which are held as non-persistent buffers (the resampling indices and solid-angle pooling weights of the planar raster, the graph Laplacian of DeepSphere, and the separations, bearings and radial profiles of equation (10)).

8

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

Table 2: Total learnable parameters. The convolution of Section 2 carries an order of magnitude more parameters than DISCO, since its filter coefficients are complex and indexed by radial profile and by both bearing orders, yet it also gives the lowest validation RMSE and the tightest posteriors of the four (Table 1), so the gain is in how the parameters are constrained rather than in how many there are. Network

Total params

This work Planar CNN DeepSphere DISCO

27,938 11,794 10,770 2,450

Table 3: Learnable parameters of the planar CNN. Each convolution is padded manually, periodically in longitude and with zeros in colatitude. Layer

Type

Weight shape

Bias shape

blocks.0 Conv2d (3×3) blocks.1 Conv2d (3×3) blocks.2 Conv2d (3×3) blocks.3 Conv2d (3×3) blocks.4 Conv2d (3×3) blocks.5 Conv2d (3×3) head Linear

(16, 1, 3, 3) (16, 16, 3, 3) (16, 16, 3, 3) (16, 16, 3, 3) (16, 16, 3, 3) (16, 16, 3, 3) (2, 16)

(16) (16) (16) (16) (16) (16) (2)

# Params 160 2,320 2,320 2,320 2,320 2,320 34

Total

11,794

Table 4: Learnable parameters of DISCO, with an axisymmetric filter of three radial taps. Layer

Type

lift Discrete-continuous conv blocks.0 Discrete-continuous conv blocks.1 Discrete-continuous conv blocks.2 Discrete-continuous conv head Linear

Weight shape

Bias shape

(16, 1, 3) (16, 16, 3) (16, 16, 3) (16, 16, 3) (2, 16)

(16) (16) (16) (16) (2)

Total

# Params 64 784 784 784 34 2,450

Table 5: Learnable parameters of DeepSphere, with Chebyshev order K = 4 and a 1→16→32→64 channel schedule, pooling between stages. Layer names are relative to the encoder, whose stages wrap the convolution and its normalization together. Layer

Type

Weight shape

Bias shape

blocks.0 ChebConv (K=4) blocks.0 BatchNorm1d blocks.1 ChebConv (K=4) blocks.1 BatchNorm1d blocks.2 ChebConv (K=4) blocks.2 BatchNorm1d head Linear

(4, 1, 16) (16) (4, 16, 32) (32) (4, 32, 64) (64) (2, 64)

(16) (16) (32) (32) (64) (64) (2)

Total

# Params 80 32 2,080 64 8,256 128 130 10,770

9

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

Table 6: Learnable parameters of the network of this work. The band limits of equation (6) give four radial profiles, up to two orders and three taps per channel pair; the coefficients cpmn are complex, so real and imaginary parts are stored separately and the imaginary part of the m = 0 order is absent. Layer

Type

Weight shapes (re, im) 2

lift Lift S → SO(3) (16, 1, 4, 2), (16, 1, 4, 1) blocks.0 SO(3) conv (16, 16, 4, 2, 3), (16, 16, 4, 1, 3) blocks.1 SO(3) conv (16, 16, 4, 2, 3), (16, 16, 4, 1, 3) blocks.2 SO(3) conv (16, 16, 4, 2, 3), (16, 16, 4, 1, 3) head Linear (2, 16) Total

Bias shape (16) (16) (16) (16) (2)

# Params 208 9,232 9,232 9,232 34 27,938

3.8 Equivariance, invariance, and time test protocols The invariance error of Table 1 is measured end-to-end on the trained networks, which are comprised of our equivariant convolutional layers, of the weak lensing task. Each validation map is rotated on the sphere by a random Euler rotation, and both the original and rotated maps are normalized by the training-set statistics and passed through the same frozen network. The error is the percentage deviation of the predicted summaries under rotation, and Table 1 quotes the mean over rotations and validation maps. The rotations are drawn once from a fixed seed and reused across every architecture. The time column of Table 1 times the first layer of each network, which maps the raw HEALPix scalar map into the working representation of the given architecture, rather than the full forward pass, since the networks differ in depth and a whole-model timing would largely measure those unrelated choices. All four architectures are timed sequentially in a single process on one dedicated, otherwise idle GPU, since contention on shared nodes inflates latencies by tens to hundreds of per cent. A single Nside = 256 map is passed at batch size one; warm-up iterations are discarded to settle kernel selection, then 50 forward passes and 50 forward–backward passes are timed individually, each bracketed by paired CUDA events with device synchronization, so the quoted numbers are per-call latencies rather than throughput. Table 1 reports the mean and standard deviation of the forward pass in single precision. 3.9 Results Referring to Table 1, as expected the planar architecture has poor SO(3) equivariance. DeepSphere also exhibits large equivariance error, due to the relatively poor resolution of the input maps. Our implementation, along with DISCO, offers near per-cent-level equivariance for the entire network. Under random rotations of the map, the inferred summaries shift on average only around a couple of per cent. Therefore any learned summary has SO(3) equivariant features extracted from it. The usefulness of this is evidenced by the weak lensing parameter estimation task, where our architecture both provides the best validation loss and the largest average FOM, thereby resulting in the most precise constraints of all the methods. Fig. 3 shows the posteriors of each method for a map at the prior mean of Ωm and S8 from the Gower Street simulations. We attribute this success not only to the exact equivariance of the network but also to the directionality of its filters: the non-zero bearing orders m and n of equation (6) carry oriented, spin-weighted features that an axisymmetric filter cannot represent. We note, however, that the κ maps are of relatively low resolution (Nside = 256), so the amount of non-linear (non-Gaussian) information available to be extracted beyond the power spectrum is limited. The planar CNN collapses to the trivial solution of predicting the ensemble mean of the dataset, so its posteriors in Ωm are prior driven. The planar CNN cannot predict Ωm at the resolution of the input maps. In the remaining tasks of CMB denoising and dense gradient regression on Gaussian random fields, the spherical convolution methods all perform significantly better than the naive planar method. Our method performs particularly well in the Gaussian random field dense gradient regression task. DeepSphere performs best in the CMB denoising task. We note that this task is not strictly SO(3) equivariant, as the instrumental noise has a preferred orientation on the sky, so an exact symmetry inductive bias is not necessarily optimal here. 4 Discussion The method has limitations. It is slower than the planar and DISCO baselines (Table 1). Its equivariance is exact only under the assumptions of Section 2, which real observations satisfy only approximately.

10

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

In this paper we have presented a SO(3)-equivariant convolutional neural network and introduced four distinct contributions. First, in Section 2 we have introduced a convolution for data on the sphere, defined directly on the HEALPix grid, that is exactly equivariant under SO(3), the group of three-dimensional rotations: rotating the input map rotates the network response in the same way. The symmetry is built into the structure of the network itself, through steerable filters, parallel transport of local reference frames and spin-weighted features, rather than being learned from training data. Second, this exactness does not come at prohibitive cost: each filter acts only over a small neighbourhood of each pixel, no transform of the full sky is ever required, and in the tests of Section 3.8 the predictions of the trained networks shift by only around one per cent when the input maps are rotated, the smallest deviation of all the methods compared. Third, on tasks where the true answer is known, the method recovers the gradient magnitude of Gaussian random fields (Section 3.2) more accurately than the planar, graph-based and axisymmetric alternatives, and removes noise and foreground contamination from cosmic microwave background maps (Section 3.3) with accuracy comparable to other methods in the literature. Fourth, in Section 3.4 we have shown what the symmetry is worth for inference: given the same number of training simulations, our network yields the tightest and most accurate constraints on the cosmological parameters Ωm and S8 , the matter density and the amplitude of matter clustering from cosmology data. A given measurement precision therefore requires fewer training simulations, and a fixed budget can be spent making each simulation more faithful to reality rather than simply making more of them, reducing the reliance on imperfect models in scientific analyses. We have introduced a scalable, exactly SO(3) equivariant convolution on the HEALPix grid and demonstrated it in SBI from spherical cosmological data. The symmetry is supplied by the architecture rather than learned from data, so a given precision requires fewer training simulations. Acknowledgments The authors thank Ofer Lahav for useful discussions. JW is supported by the STFC UCL Centre for Doctoral Training in Data Intensive Science (grant ST/W00674X/1) and by departmental and industry contributions. NJ acknowledges support from the King’s AI+ Academic Fellowship at King’s College London. NJ also acknowledges support from the ERC-selected UKRI Frontier Research Grant EP/Y03015X/1. The authors declare no competing interests. Data availability The data and code will be made publicly available upon acceptance. References Amato J, Xie Y, Medina-Varela L, Aljerwi A, McCutcheon A, Rippentrop T S, Gonzalez K, Delabrouille J, Ishak M and Ruozzi N 2025 in Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) pp. 9418–9430 Bonev B, Kurth T, Hundt C, Pathak J, Baust M, Kashinath K and Anandkumar A 2023 in Proceedings of the 40th International Conference on Machine Learning vol. 202 of Proceedings of Machine Learning Research pp. 2806–2823 Cobb O J, Wallis C G R, Mavor-Parker A N, Marignier A, Price M A, d’Avezac M and McEwen J D 2021 in International Conference on Learning Representations Cohen T S, Geiger M, Köhler J and Welling M 2018 in International Conference on Learning Representations Cohen T S, Weiler M, Kicanaoglu B and Welling M 2019 in International Conference on Machine Learning (ICML) pp. 1321–1330 Cohen T S and Welling M 2016 in International Conference on Machine Learning (ICML) pp. 2990–2999 Cranmer K, Brehmer J and Louppe G 2020 Proceedings of the National Academy of Sciences 117(48) 30055–30062 de Haan P, Weiler M, Cohen T and Welling M 2021 in International Conference on Learning Representations (ICLR)

11

IOP Publishing

Mach. Learn.: Sci. Technol. vv (yyyy) aaaaaa

Williamson et al

Defferrard M, Bresson X and Vandergheynst P 2016 in Advances in Neural Information Processing Systems 1606.09375 Defferrard M, Milani M, Gusset F and Perraudin N 2020 in International Conference on Learning Representations Esteves C, Allen-Blanchette C, Makadia A and Daniilidis K 2018 in Proceedings of the European Conference on Computer Vision (ECCV) pp. 52–68 Esteves C, Makadia A and Daniilidis K 2020 in Advances in Neural Information Processing Systems vol. 33 pp. 8614–8625 Esteves C, Slotine J J and Makadia A 2023 in Proceedings of the 40th International Conference on Machine Learning vol. 202 of Proceedings of Machine Learning Research pp. 9396–9411 Freeman W T and Adelson E H 1991 IEEE Transactions on Pattern Analysis and Machine Intelligence 13(9) 891–906 Górski K M, Hivon E, Banday A J, Wandelt B D, Hansen F K, Reinecke M and Bartelmann M 2005 The Astrophysical Journal 622 759–771 Hu Y, Yin F, Zhang W, Ren K, Song J, Deng K and Zhang D 2025 Scientific Reports 15 11522 Jeffrey N, Whiteway L, Gatti M, Williamson J, Alsing J, Porredon A et al. 2025 Monthly Notices of the Royal Astronomical Society 536 1303 2403.02314 Keiner J and Potts D 2008 Mathematics of Computation 77(261) 397–419 Kondor R, Lin Z and Trivedi S 2018 in Advances in Neural Information Processing Systems vol. 31 Kostelec P J and Rockmore D N 2008 Journal of Fourier Analysis and Applications 14(2) 145–179 McEwen J D, Büttner M, Leistedt B, Peiris H V and Wiaux Y 2015 IEEE Signal Processing Letters 22(12) 2425–2429 Ocampo J, Price M A and McEwen J D 2023 in International Conference on Learning Representations Papamakarios G and Murray I 2016 in Advances in Neural Information Processing Systems vol. 29 1605.06376 Papamakarios G, Pavlakou T and Murray I 2017 in Advances in Neural Information Processing Systems vol. 30 1705.07057 Perraudin N, Defferrard M, Kacprzak T and Sgier R 2019 Astronomy and Computing 27 130–146 Planck Collaboration, Aghanim N, Akrami Y et al. 2020 Astronomy & Astrophysics 641 A6 Potter D, Stadel J and Teyssier R 2017 Computational Astrophysics and Cosmology 4(1) 2 1609.08621 Tejero-Cantero A, Boelts J, Deistler M, Lueckmann J M, Durkan C, Gonçalves P J, Greenberg D S and Macke J H 2020 Journal of Open Source Software 5(52) 2505 Thomsen A et al. 2026 Physical Review D 2511.04681 Weiler M and Cesa G 2019 in Advances in Neural Information Processing Systems (NeurIPS) vol. 32 Williamson J, Makinen T L et al. 2026 arXiv e-prints p. arXiv:2606.11309 2606.11309 Wolz K, Krachmalnicoff N and Pagano L 2023 Astronomy & Astrophysics 676 A30 2301.09634 Worrall D E, Garbin S J, Turmukhambetov D and Brostow G J 2017 in IEEE Conference on Computer Vision and Pattern Recognition (CVPR) pp. 5028–5037 Yershova A, Jain S, LaValle S M and Mitchell J C 2010 International Journal of Robotics Research 29(7) 801–812

12

Related documents

• On Making $\mathrm{SE}(2)$-Invariant Networks Optimal#1342641
• Rotational Equivariance in Machine Learning: A Comprehensive Tutorial#819075
• Rotational Equivariance in Machine Learning: A Comprehensive Tutorial#809083
• Optimal Transportation and Alignment Between Gaussian Measures#1216211
• Why Geometric Continuity Emerges in Deep Neural Networks: Residual Connections and Rotational Symmetry Breaking#158567