Gauge Equivariant Spherical CNNs

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 . 1 INTRODUCTION . 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 can not 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 . 2 RELATED WORK . 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 ) . 3 CONTINUOUS THEORY OF GAUGE CNNS . 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 . 3.1 GEOMETRICAL FEATURES & GAUGE TRANSFORMATIONS . 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 .

Summary: 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.