GAUGE EQUIVARIANT SPHERICAL CNNS
Year: 2019

Abstract:
Spherical CNNs are convolutional neural networks that can process signals on the sphere, such as global climate and weather patterns or omnidirectional images. Over the last few years, a number of spherical convolution methods have been proposed, based on generalized spherical FFTs, graph convolutions, and other ideas. However, none of these methods is simultaneously equivariant to 3D rotations, able to detect anisotropic patterns, computationally efficient, agnostic to the type of sample grid used, and able to deal with signals defined on only a part of the sphere. To address these limitations, we introduce the Gauge Equivariant Spherical CNN. Our method is based on the recently proposed theory of Gauge Equivariant CNNs, which is in principle applicable to signals on any manifold, and which can be computed on any set of local charts covering all of the manifold or only part of it. In this paper we show how this method can be implemented efficiently for the sphere, and show that the resulting method is fast, numerically accurate, and achieves good results on the widely used benchmark problems of climate pattern segmentation and omnidirectional semantic segmentation.

text: Spherical CNNs are convolutional neural networks that can process signals on the sphere, such as global climate and weather patterns or omnidirectional images. Over the last few years, a number of spherical convolution methods have been proposed, based on generalized spherical FFTs, graph convolutions, and other ideas. However, none of these methods is simultaneously equivariant to 3D rotations, able to detect anisotropic patterns, computationally efficient, agnostic to the type of sample grid used, and able to deal with signals defined on only a part of the sphere. To address these limitations, we introduce the Gauge Equivariant Spherical CNN. Our method is based on the recently proposed theory of Gauge Equivariant CNNs, which is in principle applicable to signals on any manifold, and which can be computed on any set of local charts covering all of the manifold or only part of it. In this paper we show how this method can be implemented efficiently for the sphere, and show that the resulting method is fast, numerically accurate, and achieves good results on the widely used benchmark problems of climate pattern segmentation and omnidirectional semantic segmentation.

heading: 1 INTRODUCTION
text: In many disciplines of science and engineering, spherical signals emerge naturally. In the earth and climate sciences, globally distributed sensor arrays collect measurements like temperature, pressure, wind directions, and many other variables. Cosmologists are interested in identifying physical model parameters from real and simulated cosmic microwave background measurements sampled on spherical sky maps. In robotics, especially in applications like SLAM and visual odometry, omnidirectional and fish-eye cameras are widely used. Thus, it is clear that efficient CNNs that can directly operate on spherical signals are necessary.
When designing Spherical CNNs for these applications, there are a number of practical considerations that need to be taken into account. Firstly, whereas planar images are always sampled on a square grid, there is a wide variety of spherical grids, and the choice of sampling grid is usually dictated by hardware or task-dependent requirements. Secondly, several applications require the analysis of very high resolution spherical maps. Most existing algorithms cannot handle this scale. Thirdly, in some cases data is only available on a part of the sphere. Fourthly, whereas several graph-based and spectral methods can only learn isotropic blob-like filters, some problems require the detection of anisotropic patterns. Finally, in almost all cases, equivariance to rotations is a sought-after property since results should never depend on how we orient the sphere.
The method we propose in this paper is based on the recently proposed theory of Gauge CNNs (Cohen et al., 2019). This framework has several appealing properties: (i) it provides a principled way to achieve spatial and rotational weight-sharing on the sphere, (ii) it naturally allows for processing geometric features such as scalar, vector and tensor fields in an anisotropic and equivariant manner, and (iii) it makes it possible in principle to process a signal on overlapping local charts, in case data is not available for the whole sphere or when the whole signal does not fit in memory. To the best of our knowledge, we present the first implementation of a Gauge CNNs on a manifold with non-trivial curvature. We describe an efficient implementation which is, in principle, applicable to any grid on the sphere. Experiments on spherical segmentation of climate patterns and omnidirectional images demonstrate the effectiveness of our method. We show that our method is scalable and fast while being as numerically accurate as the best existing Spherical CNNs.

heading: 2 RELATED WORK
text: Conventional CNNs for Spherical Images. A line of prior work (Su & Grauman, 2017; Coors et al., 2018) considers using planar CNNs on intrinsically spherical tasks such as omnidirectional image segmentation. These approaches represent the spherical data on R2 via equirectangular projection which introduces position-dependent distortion rates due to curvature. In order to counteract the varying distortion rates, Su & Grauman (2017) uses kernels with increasing support towards the poles thus limiting parameter-sharing to only longitudes. Contrastingly, Coors et al. (2018) use fixed-sized kernels sampled on the tangent plane of spherical signal for better parameter-sharing and distortion invariance. Both methods assume a preferred orientation and are not SO(3) equivariant.
Spherical CNNs. Another line of work operates directly on S2 alleviating the issues related to distortions caused by projection. Boomsma & Frellsen (2017) proposes spherical convolutions for molecular modeling in both volumetric and mesh representations. Perraudin et al. (2018) proposes a graph-based spherical convolution for cosmological model classification. A subset of this line (Kondor et al., 2018; Cohen et al., 2018; Esteves et al., 2018) extends group equivariance (Cohen & Welling (2016)) over SO(3). Esteves et al. (2018) uses spherical convolutions computed by spherical FFTs. Similarly, Cohen et al. (2018) proposes a spherical correlation operation based on the SO(3)FFT. As a result both methods have inbuilt rotation equivariance property. However, the Generalized FFT algorithms used only work on inhomogeneous grids that over-sample the poles, and in practice the FFT-based methods are slow.
Recently, a number of spherical CNNs that parametrize the sphere with icosahedral tiling have been proposed (Jiang et al., 2019; Zhang et al., 2019; Liu et al., 2019; Cohen et al., 2019). These methods are tailored for the icosphere grid since this grid is quite regular and its charts can be easily mapped on R2 allowing highly optimized convolution routines to be used. Out of these methods, Gauge Equivariant Icosahedral CNN of Cohen et al. (2019) is the most similar to our method. However, their model operates on the icosahedron, which is only an approximation of the sphere, while ours operates on the sphere. Even though their method is fast and accurate, it is only equivariant to discrete icosahedral symmetries, while our method is fully SO(3)-equivariant. In comparison to all methods in this category, our method is agnostic to sampling grid used and better suited for larger scale problems.
Geometric Deep Learning. There have been a number of attempts to generalize convolution operator to manifolds in the geometric deep learning field (Bronstein et al., 2017). The key issue with manifold convolution is the lack of globally consistent reference frames attached to points over a manifold. Therefore, unlike shifting a filter over a flat image grid with a clear sense of up/down and left/right, it is not clear how to place the convolution kernel. To overcome this issue, Bruna et al. (2014); Boscaini et al. (2015) have used isotropic filters at the expense of kernel expressivity. Masci et al. (2015) have applied filters in fixed number of orientations and accumulated the responses via max pooling, losing the orientation information. Recently proposed Gauge Equivariant CNNs (Cohen et al., 2019) have shown that it is possible to have both expressive kernels and orientation information by allowing the network response change equivariantly with respect to arbitrarily chosen local reference frames (i.e. gauges).

heading: 3 CONTINUOUS THEORY OF GAUGE CNNS
text: In this section we will review the mathematical theory of gauge CNNs on general manifolds, as presented in Cohen et al. (2019). Specifically, we will define a mathematical model of feature spaces as fields, and show how one can define a convolution-like operation that only makes use of the intrinsic structure of the manifold.

heading: 3.1 GEOMETRICAL FEATURES & GAUGE TRANSFORMATIONS
text: The feature spaces in Gauge CNNs are modelled as fields f over a manifold M . For example, the input data could be a vector field of wind directions on earth, or a scalar field of intensity values on the plane (a grayscale image), or a field of diffusion tensors on R3. We will refer to such quantities (scalars, vectors, tensors, and others) as geometrical features and speak of a field of geometrical features or geometrical feature field. Even if the input data consists only of scalars, one might
want to use other kinds of fields for the internal representation learned by the network, so we will describe Gauge CNNs for general fields. Because we are primarily interested in the 2-sphere S2, we will specialize to the case of d = 2 dimensional Riemannian manifolds.
In computer science it is common to think of a vector or tensor as a list or array of numbers, but from a physical or mathematical perspective these are geometrical quantities that exist independent of a coordinatization / choice of basis. To represent a geometrical feature numerically however, we need to choose a frame for the tangent space TpM at each position p ∈ M . A smooth choice of frame is also known as a gauge. Mathematically, a gauge can be defined as a smoothly parameterized set of linear maps wp : Rd → TpM (see Fig. 1).
Since the choice of gauge is arbitrary, we should consider what happens to the coefficients of geometric features if we change it (i.e. apply a gauge transformation). Consider first the coefficients f(p) = v of a tangent vector V in the tangent space TpM at p ∈M , expressed as a pair of numbers v = (v1, v2) relative to an orthogonal frame (wp(e1), wp(e2)) at p. If we rotate the frame at p by r ∈ SO(2), i.e. wp 7→ wp ◦ r, then the coefficient vector transforms as v 7→ r−1v. The vector itself is an abstract geometrical quantity, invariant to gauge transformations: V = (wpr)r−1v = wpv.
Note that we are free to change the gauge not just at one point, but at all positions simultaneously in an arbitrary (smooth) way. However, since we will only want to work with right-handed orthogonal frames, we only need to consider rotations of the frames. Thus, for our purposes we can define a gauge transformation as a smoothly varying choice of rotation rp ∈ SO(2). Beyond scalars (which are invariant to gauge transformations) and vectors (which transform like f(p) 7→ r−1p f(p)), we will want to consider more general kinds of geometrical features. A (2, 0)- tensor, for instance, is (a linear combination of) tensor products V ⊗W of vectors V,W ∈ TpM . Given a frame, such a tensor is represented as a d × d matrix. Under a change of frame, a matrix f(p) will transform like f(p) 7→ rpf(p)r−1p . We can also flatten the matrix into a d2-dimensional coordinate vector f(p), and write the transformation as f(p) 7→ (rp⊗ rp)f(p), where rp⊗ rp is the Kronecker product.
The tensor product ρ(r) = r ⊗ r is an example of a group representation. This is a map ρ : G → GL(C,R) taking each element r of G (the rotation group SO(2) in our case) to an invertible matrix ρ(r) that acts on an C-dimensional feature vector. To be called a representation, it has to satisfy ρ(rr′) = ρ(r)ρ(r′), which is easily checked for the tensor / Kronecker product.
Thus, we can generalize to geometric feature fields that transform like f(p) 7→ ρ(r−1p )f(p) under gauge transformations, for any group representation ρ of SO(2). We will refer to such fields as a ρ-field or a field of type ρ. In a gauge equivariant CNN, one chooses for each feature space of the network such a representation ρ that determines the kind of features learned by that layer. The network is constructed such that a gauge transformation applied to the input will result in a corresponding gauge transformation in each feature space. The dimension C of ρ is equal to what is normally called the number of channels in the feature space. Typically, one would choose ρ to be block-diagonal, containing for instance a number of scalar fields (1×1 blocks ρi(r) = 1), a number of vector fields, etc. The number of copies of each type of feature is called its multiplicity.

heading: 3.2 GAUGE EQUIVARIANT CONVOLUTION
text: For each layer of the network, we want to interpret both the input and output as fields of geometrical features. If we apply a gauge transformation, the input coefficients change (f(p) 7→ ρ(r−1p )f(p)), and we want the same to happen to the output (gauge equivariance), so that we may interpret it as the coefficients of a geometrical quantity relative to a gauge. In this section we will define a convolution-like operation that has this property.
The classical convolution operation involves summing or integrating the product of a filter and the input signal over a local region. If we want to generalize this to fields on a manifold, we run into a difficulty: the geometric feature f(p) and f(q) at different points p and q in M live in a different vector space, and so we cannot directly compare them or add them up. For instance, if we have two tangent vectors V ∈ TpM and W ∈ TqM , how can we say that they are “the same” or add them up? Having chosen a frame for TpM and TqM we could add the coordinate vectors v ∈ R2 and w ∈ R2, but because we can change the gauge of both tangent spaces independently, the result is not the coefficient vector of any invariant geometrical quantity.
The solution is to apply parallel transport to the feature vectors before adding them up. Given a curve from q to p, we can transport a vectorW ∈ TqM to TpM by applying a rotation rp←q ∈ SO(2) to its coefficient vectorw. Since we can interpret rp←qw as a vector in TpM , the quantity v+rp←qw is well-defined. In general, for other kinds of geometrical features, parallel transport acts via ρ, i.e. we can add v + ρ(rp←q)w. We will use this in the definition of the gauge equivariant convolution.
Following Masci et al. (2015), we parameterize a local neighborhood around p ∈M by the tangent plane TpM ' R2 via exponential map (see Fig. 1). That is, we index nearby points q by tangent vectors using the exponential map, by defining qv = expp wpv for v ∈ R2 (“Riemannian normal coordinates”). The convolution is then defined by transporting for each nearby point qv the feature vector f(qv) to p by computing ρ(rp←qv )f(qv), transforming the resulting features at p using a learned kernel K : R2 → RCout×Cin , and integrating the result over the support of K in R2:
ψ ? f(p) = ∫ R2 K(v)ρin(rp←qv )f(qv)dv. (1)
As shown by Cohen et al. (2019), this operation is gauge equivariant if and only if K(v) satisfies
K(r−1v) = ρout(r −1)K(v)ρin(r). (2)
In section 4.2 we show how we can parameterize such a kernel via rotational weight-sharing.
3.3 EQUIVARIANCE TO SO(3)
In addition to gauge equivariance, SO(3) equivariance is a desirable property for a Spherical CNN (Cohen et al., 2018). This means that if we apply a 3D rotation to the input of the network, the output is also rotated. In this section we show that the gauge equivariant convolution as defined is also equivariant to SO(3).
Consider a local patch on the sphere (e.g. the support of the kernel), and the signal defined there. When we rotate the sphere, the patch is moved to another place, and it may change its orientation. Moving the patch is not a problem: at the new position we apply the same kernel K, so one expects that the convolution result at the new position equals the convolution result of the original signal at the old position. However, since the orientation of the kernel is determined by the gauge (which is arbitrary but fixed) and because we can arbitrarily change the orientation of the patch by rotating around its center, the kernel and the patch may be matched in a different relative orientation after applying the rotation. Fortunately, because the kernel satisfies Eq. 2, the result will be equivalent up to a gauge transformation acting by ρout, and so we have SO(3) equivariance.
Thus, in the continuous theory, the gauge equivariant convolution is also SO(3) equivariant. However, as we will see, making sure that this holds in a discrete implementation is not entirely trivial.

heading: 4 DISCRETE IMPLEMENTATION OF SPHERICAL GAUGE CNNS
text: The theory covered so far tells us how, mathematically, we can define a convolution-like operation that is gauge equivariant. However, it does not tell us exactly how to implement it on a computer, in order to process discretely sampled signals on a manifold. How we do this exactly can have a large effect on the efficiency and numerical accuracy of the method.
A signal is represented as a list of values fi = f(pi) associated with a finite number of points pi ∈ V ⊂ S2 (see Sec. 4.1). We assume that the kernel K(v) has local support, so K(v) = 0 whenever ‖v‖ > R for some radius R. Equivalently, we can say that q ∈ S2 only contributes to the
convolution result at p ∈ S2 if the geodesic distance between p and q is smaller thanR. Accordingly, we define the set of neighbors N (p) of p as the set of points q within radius R from p. A simple way of discretizing the gauge convolution (Eq. 1) is to replace the integral over R2 (identified with TpM ) by a sum over neighbors of p. Each neighbor can be associated with a tangent vector via the logarithmic map: vpq = logp q. This yields the following approximation:
ψ ? f(p) = ∑
q∈N (p)
K(vpq)ρin(rp←q)f(q) (3)
At this point it is worth comparing Eq. 3 to the simplest form of message-passing based graph convolution (Kipf & Welling, 2017; Gilmer et al., 2017). In both cases, the result of convolution is computed as a sum of messages comming from neighbors. In the case of graph CNNs, these messages are computed as Af(q), where A is a weight matrix that is shared by all neighbors. So graph convolutions use isotropic filters that cannot distinguish where the message comes from.
The gauge convolution sums messages of the form K(vpq)ρin(rp←q)f(q). Thus, the feature vectors f(q) of neighbors q are transformed in a way that depends i) on the intrinsic geometry of the manifold via rp←q and vpq , and ii) by a non-isotropic (but gauge-equivariant) learnable kernel K(vpq). We can thus think of gauge convolutions (implemented in this way) as a more powerful version of graph convolution that leverages the additional topological and metric structure of a Riemannian manifold to process geometrical data in a more flexible manner.
The discrete gauge convolution is computed in a few steps, some of which are done during precomputation and some during the forward pass: i) the logarithmic map vpq = logp q, ii) the parallel transporter rp←q , iii) the construction / parameterization of the kernel K(v), iv) the linear contraction of the kernel and the signal.
Steps (i) and (ii) can be complicated for a general manifold or mesh, but for the sphere are easily and exactly computable in closed form with standard operations, which are detailed in Appendix B. We will discuss our method for selecting the finite number of points on the sphere, as well as steps (iii) and (iv) in the following sections.
4.1 THE ICOSPHERE GRID
Even though our method uses the exact geometry of the sphere, a grid of points on the sphere must nevertheless be chosen on which the features will live. The sphere does not admit perfectly symmetrical and homogeneous high resolution grids. A grid that is fairly homogeneous and has been used successfully in Spherical CNNs before is what we call the icosphere grid (Jiang et al., 2019; Liu et al., 2019; Cohen et al., 2019). The icosphere grid can be computed at different levels of resolution. The lowest resolution s = 0 has as points pi the 12 corners of the icosahedron. Higher resolutions are obtained by repeated subdivision of the triangular faces of the icosahedron into 4 sub-triangles, followed by a projection of all points to the sphere. The result is a grid Hs with 5 × 22s+1 + 2 points at subdivision level s. We emphasize that our method is not in any way tailored to this grid, and other options such as HEALPix (Górski et al., 2005) could easily be substituted.

heading: 4.2 PARAMETERIZATION OF THE KERNEL
text: The kernel K(v) is defined as a continuous matrix-valued function of R2 that satisfies the kernel constraint (Eq. 2). In a classical CNN, where we are dealing with a homogeneous grid of pixels in R2, we can define a small (e.g. 3 × 3) set of neighboring pixels N (p) = {p + v(i)}i so that we only ever need to evaluate the kernel at a small number (e.g. 9) of points v(i). This results in a parameterization of K as an array with Cout × Cin × 3× 3 learnable coefficients. On the sphere there are no perfectly homogeneous grids, so depending on the point p ∈ V where we are evaluating the convolution ψ ? f , the neighborhood structureN (p) may look quite different.
Hence, the points vpq ∈ R2 where we need to evaluate K will differ as well. For this reason, we parameterize K as a linear combination of analytically determined continuous basis kernels. The linear coefficients will be learned.
We assume that ρin and ρout are block-diagonal with irreducible representations (irreps) as blocks (any SO(2) representation can be brought to this form by a change of basis). In this case the kernel takes on a block structure as well, with each block corresponding to a particular input/output irrep (Cohen & Welling, 2017), with irreps labelled by integer frequency n ≥ 0 (Worrall et al., 2017) (see Appendix D). So we will focus on the case where both input and output representation consist of a single irrep, and construct the full kernel block-wise as described in Appendix A.
As derived in Appendix D, the analytical solutions to Eq. 2 can be split in a independent radial part and angular part. The solutions for the angular part K(θ) are shown in Table 1, while the radial part is unconstrained. So if a set of radial functions {Ra(r)} are chosen, and {Kb(θ)} is the complete set of angular solutions, the parameterized kernel is: K(r, θ) = ∑ ab wabRl(a)Kb(θ) for weights
w. Hereafter, we denote one such solution as Ki, so that the parameterized kernel is ∑ i wiKi. The number of basis-kernels is called num_basis.
ρin → ρout Linearly independent solutions for K(θ) ρ0 → ρ0 1 ρn → ρ0 (cosnθ sinnθ) , (sinnθ − cosnθ)
ρ0 → ρm (
cosmθ sinmθ
) , ( sinmθ − cosmθ ) ρn → ρm ( c− −s− s− c− ) , ( s− c− −c− s− ) , ( c+ s+ s+ −c+ ) , ( −s+ c+ c+ s+ )
nals.

heading: 4.3 COMPUTING THE CONVOLUTION
text: Having computed the basis kernels at each vpq , we can compute the discretized gauge convolution (Eq. 3) as a linear contraction. This is done in two steps. Initially, we expand the signal f(p), which has shape num_v× c_in, to f̂ of shape num_v× num_neigh× c_in. This is done so that f̂pq is the value of the signal at the q-th neighbor of p.
Subsequently, we contract the signal f̂ with basis kernelsKi(vpq)ρin(rp←q) and weightswi to obtain the convolution result ψ ? f of shape num_v× c_out. Since a basis-kernel Ki only acts on one in/out irrep pair, it is mostly zero. In Appendix A, we detail how this block-sparsity can get exploited for computational efficiency. We note that ψ ? f can easily be computed for a subset V ⊂ V .

heading: 4.4 NONLINEARITIES
text: For the network to be gauge equivariant, every layer should be gauge equivariant, including nonlinearities. Irrep features do not commute with pointwise nonlinearities (Worrall et al., 2017; Thomas et al., 2018; Weiler et al., 2018; Kondor et al., 2018). However, we can perform a basis transformation to a basis in which pointwise non linearities are approximately gauge equivariant. Afterwards, we transform the basis back to the irreps.
For simplicity, we assume that the representation isU copies of ρ0⊕ρ1⊕...⊕ρM . One such copy can be treated as the discrete Fourier modes of a circular signal with band limit M . An inverse Discrete Fourier Transform (DFT) matrix can map these modes to N spatial samples. Under a gauge transformation of a multiple of 2π/N , the samples are cyclically shifted. The resulting representation can thus be called a regular representation and hence our procedure a RegularNonlinearity. Non linearities that act pointwise on these samples, such as the ReLU, commute with such gauge transformations. The procedure is however only approximately gauge equivariant under gauge transformations of angles that are not multiples of 2π/N . Nevertheless, we prove in Appendix E that in
the limit N → ∞, exact equivariance is restored and we provide a finite error bound, which can help selecting N .

heading: 5 QUADRATURE INTEGRATION
text: The implementation described in the previous section satisfies gauge equivariance exactly. In addition to gauge symmetry, the sphere has global rotational / SO(3) symmetry, and like previous methods we would like our method to be equivariant to 3D rotations as well. In the continuous case, gauge equivariance implies SO(3) equivariance, but the scheme above is not exactly SO(3) equivariant due to discretization (see Figure 3a). In this section we describe a numerical method that improves SO(3) equivariance at no additional computational cost during training.
Our general approach for computing ψ ? f(p) will be to interpolate the sample values at N (p) to obtain a continuous function on R2, and then use quadrature integration to get a more precise value for the integral. Quadrature is a general numerical technique for approximating integrals with finite sums. For some region A and function g, the integral ∫ A g(x)dx can be approximated by∑
x∈I ωxg(x), where I ⊂ A is some finite set of quadrature points, each with a weight ωx. The goal is selecting I and ωx such that the approximation is accurate (or even exact), for functions g satisfying some regularity assumptions (e.g. being band-limited). In our case, the region A is a disk with as radius the support radius R of the kernel. The quadrature rule is described in Appendix C.
The signal at c ∈ I is inferred from the signals at N (p) by interpolation:
f̃p(c) = 1
Z(p, c) ∑ q∈N (p) k(c, q)ρin(gp←q)f(q) (4)
where k(c, q) = exp (−||c− logp(q)||2/σ2) is a Gaussian kernel with scale σ, measuring distance between c and q in the tangent space, and Z(p, c) = ∑ q∈N (p) k(c, q) is a normalizing constant.
We can now compute the integral over R2 by quadrature integration:
ψ ? f(p) = ∑ c∈I ωcK(c)f̃p(c) (5)
The convolution Eq. 5 sums over a homogenized neighborhood and is thus more equivariant to rotations of the sphere. See Figure 3b for an illustration. Equivariance improves if a large number of quadrature points are used, which increases the computational cost. Luckily, since the composition
of linear operations is linear, we simplify:
ψ ? f(p) = ∑ c∈Ip ωcK(c)f̃p(c) = ∑ q∈N (p) ∑ c∈Ip ωck(c, q) Z(p, c) K(c)ρin(gp←q)f(q) = ∑ q∈N (p) K̂(p, q)f(q)
(6) for a new kernel K̂(p, q) = ∑ c∈Ip ωck(c,q) Z(p,c) K(c)ρin(gp←q). The new kernel K̂ can be pre-computed once, so that the convolution during run-time involves only a sum over the neighbors, just as in the naive convolution (Eq. 3). The interpolation thus does not affect computational cost.

heading: 6 EXPERIMENTS
text: In this section, we present experimental results on three benchmark problems. Throughout this section, we will refer to the gauge equivariant convolutions described in 4.3 as IrrepConvolution and similarly refer to tailored batch normalization and nonlinearities as IrrepBatchNorm (normalizing irreps instead of channels) and RegularNonlinearity as we name in our implementation. An IrrepConvolution and other operations are defined by specifying the number of input and output multiplicities for each irrep order, in analogy to number of channels in regular CNNs. In addition, we need to specify the types of irreps by their angular frequency. We do that by max freq argument which determines the cutoff frequency. We will describe ρ by a parameter max freq and multiplicity, and include each frequency 0, . . . ,max freq with the same multiplicity.

heading: 6.1 SPHERICAL MNIST
text: We perform a series of experiments on a toy dataset to validate the equivariance and generalization properties of our gauge-equivariant spherical CNN. Specifically, we compare our method with different filter sizes (1-ring, 2-ring) with and without interpolation, to three prior methods, as well as an isotropic baseline. The isotropic baseline uses only scalar features and isotropic kernels, but is similar in number of parameters to our model.
We follow a setup similar to that of Cohen et al. (2019) and generate Spherical MNIST in randomly Rotated (R) and Non-rotated (N) conditions. For all experiments, we represent the signals on an icospherical grid at level s = 4. For network architecture and training procedure, see Appendix F.1.
We present our results in Table 2 on three configurations of the Spherical MNIST dataset: N/N, N/R and R/R, where X/Y means we train on X and test on Y. In addition, we run an ablation study to test the effect of quadrature integration on equivariance. Whether interpolation is used or not, our method performs similarly to two relevant spherical CNN baselines in terms of classification score on N/N. We observe that our method automatically generalizes to 3D rotations without data augmentation and thus compares favorably against IcoCNN in the N/R condition. However, in the ablation study, we observe that our method’s generalization capability against 3D rotations drops in the absence of interpolation, particularly when smaller (1-ring) filters are used. This result shows the usefulness of quadrature integration when using small filters.
Our model outperforms FFS2CNN in all categories. Furthermore, it performs better than the isotropic baseline, which illustrates the importance of anisotropic filters and non-scalar features for the MNIST dataset. This is to be expected, as detection of line direction is naturally expressed in anisotropic filters.
As a final note, our results confirm the theoretical result that gauge equivariance implies SO(3)equivariance (Sec. 3.3).

heading: 6.2 CLIMATE PATTERN SEGMENTATION
text: We put our model to the test on a climate pattern segmentation dataset proposed by Mudigonda et al. (2017). The dataset is collection of simulated climate variables over 20 years from the Community Atmosphere Model CAM-5.0 (Neale et al., 2010). The task is to segment out regions of Atmospheric Rivers (AR) and Tropical Cyclones (TC) from background (BG). We use the preprocessed data
1Results from (Cohen et al., 2019).
released by Jiang et al. (2019) as is. The data consist of 16 channels sampled on an icospherical grid at level s = 5 (i.e. 10242 vertices).
For this experiment, we use a residual U-Net architecture. Our residual blocks contain two blocks of [IrrepConvolution, IrrepBatchNorm, RegularNonlinearity] and a skip connection. We set the multiplicity to 2 at the input resolution s = 5 and fix max freq = 2 throughout. The rest of the encoder stream progressively reduces the grid resolution down to s = 0 while doubling the multiplicities (and thus the number of channels) every time the resolution is decreased, resulting in 196 channels. The decoder stream is mirror-symmetrical to the encoder, replacing downsampling layers with upsampling layers. We use 2 residual blocks at each resolution level. The network is trained for 60 epochs with an initial learning rate of 0.01 which is reduced by a factor of 0.4 every 30 epochs. Batch size is 64. We use a weighted cross-entropy loss, due to class imbalance following the setup of Jiang et al. (2019), with Adam optimizer (Kingma & Ba, 2015).
In Table 3, we present our results. Our model compares favorably against the two spherical CNNs that have reported results on this dataset (Jiang et al., 2019; Zhang et al., 2019). Both methods operate on the same icospherical grid as our method. Our method is on par with the Icosahedral CNN of Cohen et al. (2019) on accuracy while significantly outperforming on the more relevant mAP metric. We attribute this gain to the fact that our model reflects the geometry of the signal more faithfully than IcoCNN which operates on an approximated sphere, namely, icosahedron.

heading: 6.3 OMNIDIRECTIONAL SEMANTIC SEGMENTATION
text: Next, we run our model in omnidirectional semantic segmentation task. For this experiment, we use Stanford 2D3DS dataset which consists of 1413 omnidirectional RGBD images acquired in 6 different locations with labels of 13 semantic categories. We are following the experimental setups in (Jiang et al., 2019; Cohen et al., 2019) and report averages over 3 cross-validation splits.
We use a residual U-Net similar to what is described in 6.2. The only difference is that we use 4 residual blocks per resolution instead of 2. We train the network using weighted cross-entropy loss and Adam optimizer for 100 epochs. We set batch size to 16 and initial learning rate to 0.01. Learning rate is reduced by a factor of 0.7 every 20 epochs.
We report our results in Table 4. We observe that our rotation equivariant network performs similar to the orientation-aware network of Zhang et al. (2019) while outperforming IcoCNN which is only equivariant up to discrete subgroup, A5 (i.e. icosahedral group), of SO(3). We note that 2D3DS is acquired with a preferred camera orientation rendering an additional challenge to our equivariant model which assumes no preferred orientation. Yet, the comparison against orientation-aware model of Zhang et al. (2019) reflects the expressive capability of ours despite the disadvantage.
Method BG TC AR Average (%) mAP (%)
Jiang et al. (2019) 97.0 94.0 93.0 94.7 N/A Zhang et al. (2019) 97.3 96.3 97.5 97.0 55.5 Cohen et al. (2019) 97.4 97.9 97.3 97.7 75.9 Ours 97.0 96.7 97.6 97.1 80.6
Table 3: Results on Climate Pattern Segmentation
Method Mean Acc. (%) Mean IoU (%)
Jiang et al. (2019) 54.7 38.3 Cohen et al. (2019) 55.9 39.4 Zhang et al. (2019) 58.6 43.3 Ours 58.2 39.7
Table 4: Results on Stanford 2D3DS

heading: 6.4 ATOMIZATION ENERGY PREDICTION
text: In this experiment, we test our method on molecular energy regression task given the molecular geometry and charges associated with each atom. To this end, we use QM7 dataset (Blum & Reymond, 2009; Rupp et al., 2012) which consists of molecules with varying number of atoms N (≤ 23). The molecules considered in this dataset are made up of T = 5 different elements (i.e. H, C, N, O and S). Each molecule is defined by a tuple (pi, zi) with i indexing the atoms. We follow the setup for signal representation outlined in Cohen et al. (2018) and represent geometry of molecules with spheres, si, of constant radius centered at atomic positions, pi, in a careful arrangement such that no two spheres intersect in a molecule. Finally, each atom i is associated with T -channel spherical potential functions describing atomic interactions within the molecule with Uz(x)x∈si = ∑ j 6=i,zj=z(ziz)/ |x− pi|.
We use a similar architecture to that of Cohen et al. (2018) in terms of parameters. The details of network and training is available in Appendix F.2. Our method achieves the lowest RMSE error in comparison to the spherical baselines as shown in Table 5.

heading: 6.5 COSMOLOGICAL MODEL CLASSIFICATION
text: We test our model on cosmological model classification task which is an interesting test bed to demonstrate our algorithm’s flexibility in admitted sampling scheme and scalability aspects. To this end, we use a dataset released by Perraudin et al. (2018). This dataset consists of cosmological convergence maps which are generated by whole-sky N-body simulations for two different parameter settings of ΛCDM cosmological model. Both cosmological models are simulated 30 times. For testing, 20 simulations are held out. The remaining 40 simulations are further split into training and validation sets of size 32 and 8, respectively. The signals are represented on HEALPix grid.We refer the reader to the original article for more information.
We follow the experimental setup described in (Perraudin et al., 2018) for evaluation. We split the maps into 12 × o2 patches where o = 4 indexes the hierarchical level of the patches in HEALPix sampling scheme. We then obtain signals of dimensionality (Nside/o)2 where Nside, i.e. number of pixels on one side of HEALPix patch, is set to 210. During training, we apply additive Gaussian noise with varying standard deviation, σnoise ∈ {0, 0.5, 1, 1.5, 2} and report results separately on each setting. The details of network architecture and training are available in Appendix F.3. The results in Table 6 show that our model performs similarly to the graph-based approach.

heading: 7 CONCLUSION
text: We have introduced the gauge equivariant spherical CNN, which is simultaneously efficient, numerically accurate, flexible with respect to the pixel grid used, and which can in principle be applied to data that is defined on local parts of the sphere only. Additionally, it is the first implementation of the theory of gauge CNNs for a manifold that is not locally flat. Our experimental results show that the method is accurate and achieves strong performance on various benchmark problems. Moreover, our method can be adapted to work on meshes, and we plan to work on this in the future.
1Results from running code in repository https://github.com/jonas-koehler/s2cnn.

heading: B SPHERICAL GEOMETRY
text: To compute the logarithmic map at p whose exponential map lands at q, and to compute the parallel transporter from q to p along a geodesic, we do the following.
B.1 PARALLEL TRANSPORTERS
Since rp←q is a planar rotation, it is determined by where it sends a single (non-zero) vector. We take the first basis vector bq1 = wq((1, 0)) and express it in 3D Euclidean coordinates. We then rotate this vector by the angle ∠(p, q) = arccos 〈p, q〉 between p and q around the axis p× q which is orthogonal to the pq plane. The resulting vector lies in the tangent plane at p. Then rp←q is determined as the angle between this vector and the first basis vector bp1 = wp((1, 0)) in TpS 2.
We precompute the transport angles for every point p in the grid V and every q ∈ N (p). This results in an array of angles of size num_v × num_neigh where num_v = |V| and num_neigh = maxp∈V |N (p)| is the maximum neighborhood size. For nodes with a non-maximal number of neighbors, we pad with zeros.
B.2 LOGARITHMIC MAPS
For each p ∈ V and q ∈ N (p) we need to compute vpq = logp q, which is the vector in TpS2 that points in the direction of q and has length equal to the geodesic distance between p and q. One way to compute the log map is to project the 3 dimensional vector pointing from p to q to the tangent plane. This is done by subtracting from q−p the component parallel to the p vector, as the p vector is orthogonal to the tangent plane at p. This produces a vector ṽ ∝ (q− p)−〈p, q− p〉p = q−〈p, q〉p which has the right direction. Then one can scale the length of ṽ so that it matches the geodesic distance d(p, q) (the arclength):
logp q = d(p, q) q − 〈p, q〉p ‖q − 〈p, q〉p‖
(7)
We express the result vpq = logp q in polar coordinates relative to the gauge at p. This gives two arrays log_map_r (the length / radial coordinate of v) and log_map_angle (the angular part of v, relative to the gauge at p). Both are shaped num_v×num_neigh as before. Since the geometry and grid are fixed, these arrays are computed only once before training.

heading: C QUADRATURE
text: The goal of selecting the quadrature points and weights is to approximate an integral ∫ D g(x)dx,
where D ⊂ R2 is the unit disk, with a finite sum ∑ x∈I wxg(x), for a set of points I ⊂ D and each with weight wx. We use a simple method of selecting I and wx. We desire a rotational symmetry, so the points lie on NR radial rings, each with NΘ points.
First, we use Gauss-Legendre quadrature (Hildebrand, 1987) to obtain NR quadrature points u with weights wu over the interval [−1, 1]. We map this to radial weights and points by r =√
(u+ 1)/2, wr = wu/2. For the angular component, we take a uniform grid of NΘ points on (−π, π]. The integration points on the disk are then x = (r, θ) with weights wx = wr/NΘ. This gives accurate integration, because points (u + 1)/2 and weights wu/2 create a quadrature scheme on [0, 1]. The square root creates a uniform polar grid, as the integration measure for the disk is rdrdθ = dudθ/2.
D SOLVING THE SO(2) KERNEL CONSTRAINT
Consider a kernel K a function on the tangent space TpM , such that for all v ∈ TpM , K(v) is a linear map from features in representation ρout to features in a representation ρout. This kernel is equivariant if ∀v ∈ TpM and ∀g ∈ SO(2) the following kernel constraint is satisfied:
K(g−1v) = ρout(g −1)K(v)ρin(g). ∀v ∈ TpM, g ∈ SO(2) (8)
We immediately see that this constraint is independent for each radius, so we can solve the constraint independently for each radial ring, giving rise to the angular constraint on each ring:
K(θ − g) = ρout(−g)K(θ)ρin(g). ∀θ ∈ S1, g ∈ SO(2) (9)
where we note that the circle equals SO(2) and the action of SO(2) on itself is simply an addition of angles. Any representation ρ can be written as equivalent to a direct sum of the irreducible representations (irreps) of SO(2), which are the following, indexed by their order n ∈ N:
ρ0(g) = 1, ρn(g) = ( cosng − sinng sinng cosng ) , n ∈ N>0
This means that for any representation ρ, we can always find an invertible matrix A such: ∀g : ρ(g) = A−1 block diag(ρn1 , ...ρnN )A, where {n1, ..., nN} indicate the orders of the constituent irreps. Hence, we can solve the angular kernel constraint (Eq. 9) for two irreps, ρin = ρn and ρout = ρm, and construct the solution for general representations ρin and ρout from the obtained solutions. Noting that all irreps are orthogonal in our chosen basis are orthogonal, so that ρn(g)T = ρ(−g), we see that, introducing explicit indices:
K(θ − g)ij = ρm(−g)ilK(θ)lkρn(g)kj = ρm(−g)ilρn(−g)jkK(θ)lk = (ρm ⊗ ρn)(−g)(ij),(lk)K(θ)lk
As this equation is ought to be satisfied for any g ∈ SO(2) and SO(2) is a one-dimensional Lie group, we can equivalently require it to be satisfied by an infinitesimal element of SO(2) close to the identity, creating an ordinary differential equation:
K̇(θ) =d(ρm ⊗ ρn)K(θ), where d(ρm ⊗ ρn) := ( ∂
∂g (ρm ⊗ ρn)(g)
) |g=0,
For variousm,n, we can construct (ρm⊗ρn)(g), take the derivative with respect to g and set g to 0 to obtain the matrix d(ρm ⊗ ρn). Subsequently, any computer algebra system can solve the ODE. The
solutions can be found in Table 7. Now, for a general ρin(g) = A−1 block diag(ρn1 , ...ρnN )(g)A, ρout(g) = B −1 block diag(ρm1 , ...ρmM )(g)B, a general solution has the form:
K(θ) = B−1 Km1,n1(θ) · · · Km1,nN (θ)... . . . ... KmM ,n1(θ) · · · KmM ,nN (θ) A where Km,n(θ) is a solution from Table 7 mapping from irrep ρn to ρm. For each in/out irrep pair, we get a linear solution space of 1, 2 or 4 dimensions. For general ρin, ρout, the solution space is the product of the solution spaces of each in/out irrep pair. A parameterized kernel can thus be constructed by assigning a weight to each independent solution of each in/out irrep pair.

heading: E EQUIVARIANCE ERROR BOUNDS ON REGULAR NON-LINEARITY
text: The regular non-linearity acts on each point on the sphere in the following way. For simplicity, we assume that the representation is U copies of ρ0 ⊕ ρ1 ⊕ ... ⊕ ρM . One such copy can be treated as the discrete Fourier modes of a circular signal with band limit M . We map these Fourier modes to N spatial samples with an inverse Discrete Fourier Transform (DFT) matrix. Then apply to those samples a point-wise non-linearity, like ReLU, and map back to the Fourier modes with a Discrete Fourier Transform Matrix.
This procedure is exactly equivariant for gauge transformation with angles multiple of 2π/N , but approximately equivariant for small rotations in between.
In equations, we start with Fourier modes x0, (xα(m), xβ(m))Bm=1 at some point on the sphere and result in Fourier modes z0, (zα(m), zβ(m))Bm=1. We let t = 0, ..., N − 1 index the spatial samples.
x(t) = x0 + ∑ m xα(m) cos ( 2π N mt ) + ∑ m xβ(m) sin ( 2π N mt ) y(t) = f(x(t))
z0 = 1
N ∑ t y(t)
zα(m) = 2
N ∑ t cos ( 2π N mt ) y(t)
zβ(m) = 2
N ∑ t sin ( 2π N mt ) y(t)
Note that Nyquist’s sampling theorem requires us to pick N ≥ 2B + 1, as otherwise information is always lost. The normalization is chosen so that zα(m) = xα(m) if f is the identity.
Now we are interested in the equivariance error between the following two terms, for small rotation δ ∈ [0, 1). Any larger rotation can be expressed in a rotation by a multiple of 2π/N , which is exactly equivariant, followed by a smaller rotation. We let zFTα (m) be the resulting Fourier mode if first the input is gauge-transformed and then the regular non-linearity is applied, and let zTFα (m) be the result of first applying the regular non-linearity, followed by the gauge transformation.
zFTα (m) = 2
N ∑ t cos ( 2π N mt ) y(t+ δ) = 2 N ∑ t cm(t)y(t+ δ)
zTFα (m) = 2
N ∑ t cos ( 2π N m(t− δ) ) y(t) = 2 N ∑ t cm(t− δ)y(t)
where we defined for convenience cm(t) = cos(2πmt/N). We define norms ||x||1 = |x0| +∑ m(|xα(m)|+ |xβ(m)|) and ||∂x||1 = ∑ mm(|xα(m)|+ |xβ(m)|).
Theorem 1. If the input x is band limited by B, the output z is band limited by B′, N samples are used and the non-linearity has Lipschitz constant Lf , then the error to the gauge equivariance of the regular non-linearity bounded by:
||zFT − zTF ||1 ≤ 4πLf N
( (2B′ + 1
2 )||∂x||1 +B′(B′ + 1)||x||1
)
which goes to zero as N →∞.
Proof. First, we note, since the Lipschitz constant of the cosine and sine is 1:
|cm(t− δ)− cm(t)| ≤ 2πmδ N ≤ 2πm N
|x(t+ δ)− x(t)| ≤ 2π N ∑ m m(|xα(m)|+ |xβ(m)|) = 2π N ||∂x||1
|y(t+ δ)− y(t)| ≤ Lf 2π
N ||∂x||1
|cm(t)| ≤ 1 |x(t)| ≤ |x0|+ ∑ m (|xα(m)|+ |xβ(m)|) = ||x||1
|y(t)| ≤ Lf ||x||1 Then:
|cm(t)y(t+ δ)− cm(t− δ)y(t)| =|cm(t) [y(t+ δ)− y(t)]− y(t) [cm(t− δ)− cm(t)] | ≤|cm(t)||y(t+ δ)− y(t)|+ |y(t)||cm(t− δ)− cm(t)|
≤Lf 2π
N ||∂x||1 + Lf ||x||1
2πm
N
= 2πLf N (||∂x||1 +m||x||1)
So that finally:
|zFTα (m)− zTFα (m)|
≤ 2 N ∑ t |cm(t)y(t+ δ)− cm(t− δ)y(t)|
≤4πLf N (||∂x||1 +m||x||1)
The sinus component |zFTβ (m)− zTFβ (m)| has the same bound, while |zFT0 − zTF0 | = |y(t+ δ)− y(t)|, which is derived above. So if z is band-limited by B′:
||zFT − zTF ||1 = |zFT0 − zTF0 |+ B′∑ m=1 |zFTα (m)− zTFα (m)|+ |zFTβ (m)− zTFβ (m)|
≤ 4πLf N (2B′ + 1 2 )||∂x||1 + B′∑ m=1 2m||x||1  =
4πLf N
( (2B′ + 1
2 )||∂x||1 +B′(B′ + 1)||x||1
)
Since ||∂x||1 = O(B||x||1), we get ||zFT −zTF ||1 = O(BB ′+B′2
N ||x||1), which obviously vanishes as N →∞.

heading: F EXPERIMENTAL DETAILS
text: In the following subsections, we provide with more details regarding experiments on MNIST, Atomization Energy Prediction and Cosmological Model Classification.
F.1 MNIST
For the MNIST experiments, we disambiguate two variants, both used with interpolation and without. The first variant uses only the 1-ring of 7 nearest neighbors and representations ρ0 ⊕ ρ1. This version uses three radial functions: one that covers the self-interaction and one that covers the other neighbors.
The second variant uses the 2-ring of 19 nearest neighbors and representations ρ0 ⊕ ρ1 ⊕ ρ2. This version uses three radial functions: one for the self-interaction, one for the neighbors on the 1-ring and one for the neighbors on the 2-ring.
For either variant, we build a network consisting of 8 convolutional layers, with stride after each 2 layers. A stride layer selects the vertices that are present in the IcoSphere with one fewer subdivision, and divides the number of vertices by roughly 4. The input is a scalar feature. After convolution, the irrep features with multiplicity after each layer: [16, 16, 32, 32, 64, 64, 64, 64]. The last convolution layer maps to a scalar feature of multiplicity 64. Then the signal is averaged over the vertices, creating a gauge invariant network. This is then processed by 2 layer of MLP mapping to multiplicity 50 and then to the 10 classification logits.
Nonlinearities are ReLU for scalar features and the RegularNonlinearity for all other features. Batchnorm is used, as well as dropout with drop probability 10%.
F.2 ATOMIZATION ENERGY PREDICTION
Our network architecture used in this experiment is depicted in Table 8. We use four blocks of gauge equivariant convolutions followed by batchnorm and regular nonlinearities. We also deploy pooling / downsampling layers in blocks (1-3) to reduce the resolution of the spherical maps progressively. Rest of the network consists of linear layers followed by batchnorm and ReLU. We sum-pool over the atoms before the loss function (see block-7) to become invariant against permutations of them.
We train the model with batch size 32 and an initial learning rate of 1e−3 for 30 epochs using mean-squared-error (MSE) as loss function. The learning rate is reduced by 0.1 every tenth epoch.
F.3 COSMOLOGICAL MODEL CLASSIFICATION
For this experiment, we use HEALpix sampling on split level o = 4 following Perraudin et al. (2018) which covers 0.5% of the spherical manifold and has signal dimensionality of 65k.
For the experiments, we use a network architecture identical to the one described in (Perraudin et al., 2018) by replacing their graph convolutions with our gauge equivariant convolutions. We elaborate the architecture in Table 9.
For each noise factor, we train an instance of this model with batch size 16 and an initial learning rate of 1e−4 for 120 epochs. Finally, we use Binary Cross Entropy as the loss function.

Review summaries:
The paper proposes SO(3) equivariant layer, derived from the recently introduced Gauge equivariant CNN framework. The novel contributions are in taking the Gauge equivariance CNN and finding efficient ways to perform logarithmic mapping, parallel transport, and convolution by the equivariant kernel when applied to the sphere. An interpolation scheme for an improved approximation to global SO(3) symmetry is also discussed. Experimental results on spherical MNIST, climate pattern segmentation, and omnidirectional semantic segmentation demonstrate the usefulness of the proposed method for prediction on a sphere.