A Program to Build E(N)-Equivariant Steerable CNNs 

1 INTRODUCTION . In machine learning , it is common for tasks to present a number of symmetries . A symmetry in the data occurs , for example , when some property ( e.g. , the label ) does not change if a set of transformations is applied to the data itself , e.g . translations or rotations of images . Symmetries are algebraically described by groups . If prior knowledge about the symmetries of a task is available , it is usually beneficial to encode them in the models used ( Shawe-Taylor , 1989 ; Cohen & Welling , 2016a ) . The property of such models is referred to as equivariance and is obtained by introducing some equivariance constraints in the architecture ( see Eq . 2 ) . A classical example are convolutional neural networks ( CNNs ) , which achieve translation equivariance by constraining linear layers to be convolution operators . A wider class of equivariant models are Euclidean steerable CNNs ( Cohen & Welling , 2016b ; Weiler et al. , 2018a ; Weiler & Cesa , 2019 ; Jenner & Weiler , 2021 ) , which guarantee equivariance to isometries of a Euclidean space Rn , i.e. , to translations and a group G of originpreserving transformations , such as rotations and reflections . As proven in Weiler et al . ( 2018a ; 2021 ) ; Jenner & Weiler ( 2021 ) , this requires convolutions with G-steerable ( equivariant ) kernels . Our goal is developing a program to implement any semi-direct product group Rn oG equivariant CNN , with compact G , minimizing the requirements to implement new G-steerable kernels . Lang & Weiler ( 2020 ) provides a first step in this direction by generalizing the Wigner-Eckart theorem from quantum mechanics to obtain a general technique to parametrize G-steerable kernel spaces over orbits of a compact G. The theorem reduces the task of building steerable kernel bases to that of finding some pure representation theoretic ingredients . Since the equivariance constraint only relates points g.x ∈ Rn in the same orbit G.x ⊂ Rn , a kernel can take independent values on different orbits . Fig . 1 shows examples of orbits in R2 under the action of continuous planar rotations G = SO ( 2 ) ( orbits are rings of different radii ) and of discrete rotations by π2 , G = C4 ( orbits are sets of four points ) . Hence , Lang & Weiler ( 2020 ) suggest independently parameterizing orbits of G in Rn ( colored regions in Fig.1 ) . While this guarantees a complete parameterization of G-steerable functions , convolution requires kernels to be integrable , and the numerical discretization of the filters requires them to be sufficiently band-limited1 . Indeed , particular care should be taken to avoid aliasing effects while discretizing filters on a grid ( e.g. , to process voxelized data ) . Lang & Weiler ( 2020 ) naturally support band-limiting along each orbit of G but do not restrict the kernel across different orbits . Intuitively , this means that a filter can be decomposed into a finite number of orbits but these should then be patched together with some smoothness considerations . This introduces an important requirement and a sensitive design choice . First , building a new G-steerable CNNs requires one to identify the orbits of G in Rn . Second , the choice of orbits to consider affect the band-limitation of a filter , but this relation is hard to explicitly quantify in a general setting . This problem is more severe for smaller groups G ( particularly , finite ones ) , as Rn decomposes into a larger number of orbits . For instance , if G = SO ( 2 ) in Fig . 1a , band-limiting along rings is controlled by the Wigner-Eckart theorem but is not enforced along the radial component . If G = C4 in Fig . 1b , each ring is partitioned into discrete orbits : because a larger number of orbits increases the maximum angular frequency , aliasing can be sensitive to this design choice . Contributions and Outline In this work , we solve this issue and propose a general solution strategy to build arbitrary CNNs with G- steerable kernels , for any compact group G. Theorem 2.1 generalizes the result of Lang & Weiler ( 2020 ) from G-homogeneous spaces to more general spaces X carrying a G-action . Our theorem relies on an initial basis for the unconstrained kernel space which is then restricted to a G-steerable subspace . Aliasing effects are thereby controlled by the choice of this initial basis on the general G-space X . In practice , in this work we focus on Euclidean spaces X = Rn , although the theory developed more generally supports semi-direct product groups , i.e . X oG . To support any compact group G with minimal requirements , we propose to choose an initial basis spanning a kernel space closed under the action of any such G , i.e . under G′ = O ( n ) . The benefits are two-fold : 1 ) it allows direct control on aliasing through the initial band-limited basis and 2 ) it completely disentangles the discretization issues from the choice of G , minimizing the requirements to implement equivariance to new groups . For example , this enables the parameterization of C4 filters without discretizing rings into a finite number of points as in Fig . 1b . Algorithm 1 summarizes our method and its requirements . Because of the large variety of symmetries in R3 , and since the planar case was extensively studied in Weiler & Cesa ( 2019 ) , in our experiments , we put particular emphasis on the 3D setting . To illustrate the generality of our method , we instantiate it on many different subgroups of O ( 3 ) , the group of 3D rotations and reflections , and compare them experimentally . This results in many new 3D convolution networks equivariant , for example , to icosahedral symmetries , axial symmetries , cylindrical symmetries or conical symmetries . Axial symmetries are particularly relevant since natural 3D data generally has vertical orientation ; see Sec . G. In Sec.5.1 , we discuss many design choices for 3D equivariant networks , and compare them in the experiments in Sec . 5.3 . In particular , we find that designs based on different discretizations of the group and pointwise non-linearities to be beneficial when working with volumetric data . Additionally , Sec.5.2 compares the stability of our basis and of one using finite orbits of G and the relative performance of two models using these bases . In the Appendix , we derive some results for working with representation theoretic objects to reduce the work needed to implement new equivariant CNNs . This includes real-valued representations of compact groups , harmonic bases for induced representations , numerical irreps decomposition and representation theory of direct product groups . Finally , we implement the program described in this work in the form of a general purpose library based on PyTorch which we plan to release publicly . 2 THE THEORY OF STEERABLE CNNS AND G-STEERABLE KERNEL SPACES In Sec . 2.1 , we discuss preliminaries on steerable CNNs and the equivariance constraints on their filters . Sec . 2.2 describes the generalization of Lang & Weiler ( 2020 ) we use to parametrize kernels . 1A band-limited signal ’ s spectrum has bounded support . Nyquist-Shannon sampling theorem states that only sufficiently band-limited signals are faithfully represented by discrete samples . Otherwise , aliasing effects occur . 2.1 EUCLIDEAN STEERABLE CNNS . Equivariant neural networks are characterized by their property to commute with the action of a given symmetry group on their input , intermediate feature spaces , and output . We focus on convolutional networks on Euclidean spaces Rn , acted on by the group ( Rn , + ) oG , which is a semidirect product of translations in ( Rn , + ) and origin-preserving transformations G ≤ O ( n ) . A general framework of equivariant convolutional models is that of steerable CNNs , which we briefly recapitulate here . For more details , see Cohen & Welling ( 2016b ) ; Weiler et al . ( 2018a ; 2021 ) ; Weiler & Cesa ( 2019 ) . The fundamental design choice underlying steerable CNNs is that they operate on feature vector fields , which are similar to feature maps but differ in that they are associated with a well defined action of ( Rn , + ) oG . The geometric type of a feature field is prescribed by an orthogonal group representation ρ : G→ Rdρ×dρ which determines the transformation law of dρ-dimensional feature vectors under the action of G. Specifically , a feature field of type ρ is a map f : Rn → Rdρ which transforms under ( Rn , + ) oG according to the induced representation Ind ( R n , + ) oG G ρ : ( [ Ind ( Rn , + ) oG G ρ ] ( tg ) f ) ( x ) : = ρ ( g ) f ( g−1 ( x− t ) ) g ∈ G , t ∈ ( Rn , + ) . ( 1 ) Intuitively , the induced representation acts by moving feature vectors spatially from g−1 ( x− t ) to x and by transforming each individual feature according to ρ ( g ) ( see the visualizations in Weiler & Cesa ( 2019 ) ) . Prominent examples are scalar or tangent vector fields , which transform according to the trivial representation ρ ( g ) = 1 or the standard representation ρ ( g ) = g , respectively . Groupconvolution networks ( Cohen & Welling , 2016a ; Weiler et al. , 2018b ; Kondor & Trivedi , 2018 ) are covered as a special case with ρ being the regular representation of G. It is common to define the full feature fields of a network as a direct sum ( concatenation ) f : = ⊕ i fi of multiple individual feature fields fi . The full feature space transforms according to the direct sum representation ρ : = ⊕ i ρi of the individual fields , which ensures that they transform independently from each other . Weiler et al . ( 2018a ) ; Jenner & Weiler ( 2021 ) proved that , under mild assumptions , the most general linear equivariant maps between spaces of feature fields are convolutions with G-steerable kernels . If the input and output fields have types ρin : G→ Rdin×din and ρout : G→ Rdout×dout , G-steerable kernels are convolution kernels K : Rn → Rdout×din satisfying the linear G-steerability constraint K ( g.x ) = ρout ( g ) K ( x ) ρin ( g ) −1 ∀ g ∈ G , x ∈ Rn . ( 2 ) Thus , designing steerable CNNs only requires finding a basis of the vector space of G-steerable kernels , then used to parameterize a conventional Euclidean convolution . Note that Eq . 2 relates the kernel values on all points g.x in the orbit G.x , but leaves their values on different orbits unrelated . 2.2 A WIGNER-ECKART THEOREM FOR SUBGROUP EQUIVARIANT CONVOLUTION KERNELS . Since we study equivariance with respect toRnoG , it is necessary to reinterpret a convolution filter in a CNN as the discretization of a continuous function . To build a basis for G-steerable filters , we start from a finite basis B for band-limited filters on Rn , spanning a subspace of L2 ( Rn ) closed under G ; then , we find its subspace of filters satisfying the G-steerability constraint in Eq . 2 . Before stating our main result , we need to introduce a few concepts and define some notation . Since G is compact , we can assume all finite-dimensional representations to be orthogonal , i.e. , ρ ( g−1 ) = ρ ( g ) −1 = ρ ( g ) T . To parametrize a kernel κ : Rn → Rdout×din satisfying the steerability constraint in Eq . 2 , it is convenient to use the vectorized form K ( · ) = vec ( κ ( · ) ) : Rn → Rdout·din . The constraint2 becomes : K ( g.x ) = [ ( ρin ⊗ ρout ) ( g ) ] K ( x ) ∀g ∈ G , x ∈ Rn . ( 3 ) The matrix ( ρin ⊗ ρout ) ( g ) is the Kronecker product of the two matrices ρin ( g ) and ρout ( g ) . Irrep Decomposition A useful property is that there exists a set 3 Ĝ of special representations of G , called irreducible representations or irreps , such that any representation ρ of G can be decomposed as a direct sum of them : ρ ( g ) = QT ( ⊕ i∈I ρi ( g ) ) Q , where ρi ∈ Ĝ , I is an index set ranging over the elements of Ĝ ( possibly , with repetition ) andQ is an orthogonal matrix . The direct sum ρ1 ( g ) ⊕ρ2 ( g ) of two matrices is a block diagonal matrix containing the two matrices in its blocks . W.l.o.g. , we can 2Note that vec ( ABC ) = ( CT ⊗A ) vec ( B ) , where vec ( · ) stacks matrix columns into a vector . 3More precisely , irreps come in equivalence classes and the set Ĝ contains a representative for each class . assume ρin and ρout in the kernel constraint from Eq . 2 to be irreps . The solutions for arbitrary ρin and ρout can be recovered from their irreps decomposition , as explained in Weiler & Cesa ( 2019 ) . Hence , from now , we will assume ρin = ρl and ρout = ρJ , where l and J are two indices over Ĝ . Tensor Product The matrix ρin ⊗ ρout built with the Kronecker product in Eq . 3 is itself a representation of G , namely a tensor product representation . As such , we can decompose it in terms of G irreps : ρl ⊗ ρJ = [ CGlJ ] T ( ⊕ j ⊕ [ j ( lJ ) ] s ρj ) CGlJ , where [ j ( lJ ) ] indicates the multiplicity of the irrep ρj , while CGlJ is the change of basis matrix.4 The block diagonal structure of the direct sum allows one to distinguish blocks of rows in CGlJ , acted on by individual irreps in the direct sum . The block associated with the s-th occurrence of ρj is denoted as CGj ( lJ ) s ∈ Rdj×dldJ . Endomorphisms We are interested in the matrices commuting with an irrep ρj ∈ Ĝ , i.e . its endomorphism space . This is a vector space and we assume we know a basis5 { cjr ∈ Rdj×dj } r for it . Steerable Basis To parametrize our kernels , we first need a basis for square-integrable functions on Rn , i.e . for L2 ( Rn ) . A G-steerable basis for L2 ( X ) is a collection of orthogonal functions { Y mji : X → R } j∈Ĝ , m≤dj , i≤mj , with mj a positive ( possibly infinite ) integer . Denoting by Yji : X → R dj the stack { Y mji } dj m=1 , Yji has the defining property that ∀g ∈ G , x ∈ X , Yji ( g.x ) = ρj ( g ) Yji ( x ) . Theorem 2.1 ( Basis for G-Steerable Kernels ) . Let G be a compact group acting on a space X . Let B = { Yj′i′ } j′i′ be a G-steerable basis for X . Assume ρin = ρl and ρout = ρJ in Ĝ . A basis for ( vectorized ) G-steerable kernels over X is given by { Kjisr } jisr , which are defined as : Kjisr ( x ) = [ CG j ( lJ ) s ] T · cjr · Yji ( x ) . ( 4 ) The proof is in Appendix B ( including the case of complex valued representations and functions ) . Note that this relies on a G-steerable basis B for L2 ( X ) . Since designing a new G-steerable basis can be laborious , we propose reusing a G′-steerable basis for all G ≤ G′ via group restriction . Group Restriction Given a group G′ and a subgroup G < G′ , one can turn a representation ρ′ of G′ into a representation ResG ′ G ρ ′ of G by restricting its domain to G. It decomposes as : ResG ′ G ρj′ = [ ID j′ ] T ( ⊕ j ⊕ [ jj′ ] t ρj ) IDj ′ where [ jj′ ] is the multiplicity of ρj and IDj ′ is the change of basis . As earlier , the block-diagonal structure of the direct sum distinguishes blocks of rows of the matrix IDj ′ . The block associated with the t-th occurrence of ρj is ID jj′ t ∈ Rdj×dj′ . Theorem 2.2 ( Basis for G-Steerable Kernels through group Restriction ) . Let G ≤ G′ be compact groups acting on X . Let B′ = { Yj′i′ } j′i′ be a G′-steerable basis for X . Assume ρin = ρl and ρout = ρJ in Ĝ . A basis for ( vectorized ) G-steerable kernels over X is given by { Kj′i′jtsr } j′i′jtsr , which are defined as : Kj′i′jtsr ( x ) = [ CG j ( lJ ) s ] T · cjr · ID jj′ t ·Yj′i′ ( x ) . ( 5 ) See proof in Appendix B.3 . To give an intuition for Theorem 2.2 , observe that performing convolution with Kj′i′jtsr over a scalar input field is equivalent to a first convolution with Yj′i′ , followed by multiplication with the other matrices . The convolution with Yj′i′ is a projection on the space spanned by { Y mj′i′ } dj′ m=1 , which transforms with ρj′ under G ′ . Through IDjj ′ t , we project to its t-th subspace which transforms as ρj under G. The result can be transformed using any endomorphism of ρj ; hence , we have a different basis kernel for each element of the basis { cjr } r. Finally , we embed the result into a space transforming under the tensor product ρl ⊗ ρJ using [ CGj ( lJ ) s ] T . We now present two examples on the circle X = S1 for simplicity . Example 1 : SO ( 2 ) -Steerable Kernels In Fig . 2 , we show a simple example of the basis for a G = SO ( 2 ) -steerable kernel space . In this case , we use G′ = G , so our parameterization is equivalent to Lang & Weiler ( 2020 ) . We consider steerable kernels with input irrep ρin = ρl=0 and output irrep ρout one of ρJ=0 , ρJ=1 or ρJ=2.6 Hence , the kernels have form , respectively , 4CG stands for Clebsh-Gordan since this matrix contains the so-called Clebsh-Gordan coefficients . 5In the complex case , such matrices are multiples of the identity , but this space might be larger in the real case . 6For θ∈ [ 0 , 2π ) , ρ0 : θ 7→1 is the trivial irrep ; if k > 0 , ρk maps θ to the 2×2 rotation matrix of k · θ radians . j ′ = 0 r = 1 t = 1 ( a ) J = 0 j ′ = 1 r = 1 t = 1 j ′ = 1 r = 2 t = 1 ( b ) J = 1 j ′ = 2 r = 1 t = 1 j ′ = 2 r = 2 t = 1 ( c ) J = 2 Figure 2 : Basis for G = SO ( 2 ) -steerable kernels according to Eq . 5. ρin = ρl=0 is 1-dimensional . ρout is either J = 0 , 1 or J = 2 , which are respectively 1 , 2 and 2 dimensional . Each column is an element Kj′i′jtsr of the basis , while each row is an output channel . Since ρ0 ⊗ ρJ = ρJ , j = J . For circular harmonics , i′ = 1 ; since G = G′ , j = j′ and t = 1 . K0 : S1 → R1×1 , K1 : S1 → R2×1 and K2 : S1 → R2×1 . An SO ( 2 ) -steerable basis on X = S1 is given by the circular harmonics : Y 0j′i′ ( θ ) = cos ( j ′θ ) and Y 1j′i′ ( θ ) = sin ( j ′θ ) , for j′ ∈ N and i′ = 1. ρ0 ⊗ ρJ = ρJ , so we always have j = J and s = 1 . Since G = G′ , ResG ′ G ρj′ = ρj′ = ρj , hence t = 1 and ID jj′ t is the identity . Additionally , any 2-dimensional irrep ρj ( j > 0 ) of SO ( 2 ) has a 2-dimensional endomorphism space spanned by cjr=0 = [ 1 0 0 1 ] and cjr=1 = [ 0 1 −1 0 ] .7 This can be verified visually : in Fig . 2b and 2c , the second column is obtained by swapping the rows of the first and by changing the sign of the second row . Example 2 : CN -Steerable Kernels In Fig . 3 , we show the basis for G = C4 steerable kernels obtained using Theorem 2.2 by choosing G′ = SO ( 2 ) . We use the trivial irrep l = 0 in the input and J = 0 , 1 , and 2 in the output . ρj occurs in a restricted representation Res SO ( 2 ) C4 ρj′ whenever j′ = |j + 4k| , k ∈ Z . Indeed , the basis prescribed in Theorem 2.2 has infinitely many elements , but each element is associated with a circular harmonic of a specific frequency j′ . It is natural to use only a band-limited finite subset of this basis : in Fig . 3 , we use only frequencies up to j′ = 4 . Because both j = 0 and j = 2 are 1-dimensional irreps of C4 but j′ > 0 is a 2-dimensional representation of SO ( 2 ) , ResG ′ G ρj′ contains two copies of ρj ; in such cases , t ∈ { 1 , 2 } . This happens for j′ = 4 and j′ = 2 in Fig . 3a . ρj=1 has a 2-dimensional space of endomorphisms , spanned by the same basis described in the SO ( 2 ) example above , so r ∈ { 1 , 2 } in this case . Note that the parametrization in Fig . 3 is continuous along the ring and does not suffer from the discretization issue in Fig . 1b .

Summary: The authors present a general algorithm for constructing steerable equivariant neural networks. The main idea is that by constructing a steerable basis for a large symmetry group such as O(3), using the ideas in the paper, we can readily construct equivariant networks which are equivariants to smaller subgroups, such as SO(2). Experimental results on subgroups of isometries of $\mathbb{R}^3$ show the usefulness of the method for 3D shape recognition for voxelized data and a regression task related to molecular binding affinity.