On the Universality of Rotation Equivariant Point Cloud Networks

1 INTRODUCTION . Designing neural networks that respect data symmetry is a powerful approach for obtaining efficient deep models . Prominent examples being convolutional networks which respect the translational invariance of images , graph neural networks which respect the permutation invariance of graphs ( Gilmer et al. , 2017 ; Maron et al. , 2019b ) , networks such as ( Zaheer et al. , 2017 ; Qi et al. , 2017a ) which respect the permutation invariance of sets , and networks which respect 3D rotational symmetries ( Cohen et al. , 2018 ; Weiler et al. , 2018 ; Esteves et al. , 2018 ; Worrall & Brostow , 2018 ; Kondor et al. , 2018a ) . While the expressive power of equivariant models is reduced by design to include only equivariant functions , a desirable property of equivariant networks is universality : the ability to approximate any continuous equivariant function . This is not always the case : while convolutional networks and networks for sets are universal ( Yarotsky , 2018 ; Segol & Lipman , 2019 ) , popular graph neural networks are not ( Xu et al. , 2019 ; Morris et al. , 2018 ) . In this paper , we consider the universality of networks that respect the symmetries of 3D point clouds : translations , rotations , and permutations . Designing such networks is a popular paradigm in recent years ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ; Poulenard et al. , 2019 ; Zhao et al. , 2019 ) . While there have been many works on the universality of permutation invariant networks ( Zaheer et al. , 2017 ; Maron et al. , 2019c ; Keriven & Peyré , 2019 ) , and a recent work discussing the universality of rotation equivariant networks ( Bogatskiy et al. , 2020 ) , this is a first paper which discusses the universality of networks which combine rotations , permutations and translations . We start the paper with a general , architecture-agnostic , discussion , and derive two sufficient conditions for universality . These conditions are a result of a novel characterization of equivariant polynomials for the symmetry group of interest . We use these conditions in order to prove universality of the prominent Tensor Field Networks ( TFN ) architecture ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ) . The following is a weakened and simplified statement of Theorem 2 stated later on in the paper : Theorem ( Simplification of Theorem 2 ) . Any continuous equivariant function on point clouds can be approximated uniformly on compact sets by a composition of TFN layers . We use our general discussion to prove the universality of two additional equivariant models : the first is a simple modification of the TFN architecture which allows for universality using only low dimensional filters . The second is a minimal architecture which is based on tensor product representations , rather than the more commonly used irreducible representations of SO ( 3 ) . We discuss the advantages and disadvantages of both approaches . To summarize , the contributions of this paper are : ( 1 ) A general approach for proving the universality of rotation equivariant models for point clouds ; ( 2 ) A proof that two recent equivariant models ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ) are universal ; ( 3 ) Two additional simple and novel universal architectures . 2 PREVIOUS WORK . Deep learning on point clouds . ( Qi et al. , 2017a ; Zaheer et al. , 2017 ) were the first to apply neural networks directly to the raw point cloud data , by using pointwise functions and pooling operations . Many subsequent works used local neighborhood information ( Qi et al. , 2017b ; Wang et al. , 2019b ; Atzmon et al. , 2018 ) . We refer the reader to a recent survey for more details ( Guo et al. , 2020 ) . In contrast with the aforementioned works which focused solely on permutation invariance , more related to this paper are works that additionally incorporated invariance to rigid motions . ( Thomas et al. , 2018 ) proposed Tensor Field Networks ( TFN ) and showed their efficacy on physics and chemistry tasks . ( Kondor et al. , 2018b ) also suggested an equivariant model for continuous rotations . ( Li et al. , 2019 ) suggested models that are equivariant to discrete subgroups of SO ( 3 ) . ( Poulenard et al. , 2019 ) suggested an invariant model based on spherical harmonics . ( Fuchs et al. , 2020 ) followed TFN and added an attention mechanism . Recently , ( Zhao et al. , 2019 ) proposed a quaternion equivariant point capsule network that also achieves rotation and translation invariance . Universal approximation for invariant networks . Understanding the approximation power of invariant models is a popular research goal . Most of the current results assume that the symmetry group is a permutation group . ( Zaheer et al. , 2017 ; Qi et al. , 2017a ; Segol & Lipman , 2019 ; Maron et al. , 2020 ; Serviansky et al. , 2020 ) proved universality for several Sn-invariant and equivariant models . ( Maron et al. , 2019b ; a ; Keriven & Peyré , 2019 ; Maehara & NT , 2019 ) studied the approximation power of high-order graph neural networks . ( Maron et al. , 2019c ; Ravanbakhsh , 2020 ) targeted universality of networks that use high-order representations for permutation groups ( Yarotsky , 2018 ) provided several theoretical constructions of universal equivariant neural network models based on polynomial invariants , including an SE ( 2 ) equivariant model . In a recent work ( Bogatskiy et al. , 2020 ) presented a universal approximation theorem for networks that are equivariant to several Lie groups including SO ( 3 ) . The main difference from our paper is that we prove a universality theorem for a more complex group that besides rotations also includes translations and permutations . 3 A FRAMEWORK FOR PROVING UNIVERSALITY . In this section , we describe a framework for proving the universality of equivariant networks . We begin with some mathematical preliminaries : 3.1 MATHEMATICAL SETUP . An action of a group G on a real vector space W is a collection of maps ρ ( g ) : W → W defined for any g ∈ G , such that ρ ( g1 ) ◦ ρ ( g2 ) = ρ ( g1g2 ) for all g1 , g2 ∈ G , and the identity element of G is mapped to the identity mapping on W . We say ρ is a representation of G if ρ ( g ) is a linear map for every g ∈ G. As is customary , when it does not cause confusion we often say that W itself is a representation of G . In this paper , we are interested in functions on point clouds . Point clouds are sets of vectors in R3 arranged as matrices : X = ( x1 , . . . , xn ) ∈ R3×n . Many machine learning tasks on point clouds , such as classification , aim to learn a function which is invariant to rigid motions and relabeling of the points . Put differently , such functions are required to be invariant to the action of G = R3 o SO ( 3 ) × Sn on R3×n via ρG ( t , R , P ) ( X ) = R ( X − t1Tn ) PT , ( 1 ) where t ∈ R3 defines a translation , R is a rotation and P is a permutation matrix . Equivariant functions are generalizations of invariant functions : If G acts on W1 via some action ρ1 ( g ) , and on W2 via some other group action ρ2 ( g ) , we say that a function f : W1 → W2 is equivariant if f ( ρ1 ( g ) w ) = ρ2 ( g ) f ( w ) , ∀w ∈W1 and g ∈ G. Invariant functions correspond to the special case where ρ2 ( g ) is the identity mapping for all g ∈ G. In some machine learning tasks on point clouds , the functions learned are not invariant but rather equivariant . For example , segmentation tasks assign a discrete label to each point . They are invariant to translations and rotations but equivariant to permutations – in the sense that permuting the input causes a corresponding permutation of the output . Another example is predicting a normal for each point of a point cloud . This task is invariant to translations but equivariant to both rotations and permutations . In this paper , we are interested in learning equivariant functions from point clouds into WnT , where WT is some representation of SO ( 3 ) . The equivariance of these functions is with respect to the action ρG on point clouds defined in equation 1 , and the action of G on WnT defined by applying the rotation action from the left and permutation action from the right as in 1 , but ‘ ignoring ’ the translation component . Thus , G-equivariant functions will be translation invariant . This formulation of equivariance includes the normals prediction example by taking WT = R3 , as well as the segmentation case by setting WT = R with the trivial identity representation . We focus on the harder case of functions into WnT which are equivariant to permutations , since it easily implies the easier case of permutation invariant functions to WT . Notation . We use the notation N+ = N ∪ { 0 } and N∗+ = ⋃ r∈N Nr+ . We set [ D ] = { 1 , . . . , D } and [ D ] 0 = { 0 , . . . , D } . Proofs . Proofs appear in the appendices , arranged according to sections . 3.2 CONDITIONS FOR UNIVERSALITY . The semi-lifted approach In general , highly expressive equivariant neural networks can be achieved by using a ‘ lifted approach ’ , where intermediate features in the network belong to high dimensional representations of the group . In the context of point clouds where typically n 3 , many papers , e.g. , ( Thomas et al. , 2018 ; Kondor , 2018 ; Bogatskiy et al. , 2020 ) use a ‘ semi-lifted ’ approach , where hidden layers hold only higher dimensional representations of SO ( 3 ) , but not high order permutation representations . In this subsection , we propose a strategy for achieving universality with the semi-lifted approach . We begin by an axiomatic formulation of the semi-lifted approach ( see illustration in inset ) : we assume that our neural networks are composed of two main components : the first component is a family Ffeat of parametric continuous G-equivariant functions ffeat which map the original point cloud R3×n to a semi-lifted point cloud Wnfeat = ⊕n i=1Wfeat , where Wfeat is a lifted ( i.e. , high-order ) representation of SO ( 3 ) . The second component is a family of parametric linear SO ( 3 ) -equivariant functions Fpool , which map from the high order representation Wfeat down to the target representation WT . Each such SO ( 3 ) –equivariant function Λ : Wfeat →WT can be extended to a SO ( 3 ) ×Sn equivariant function Λ̂ : Wnfeat → WnT by applying Λ elementwise . For every positive integer C , these two families of functions induce a family of functions FC obtained by summing C different compositions of these functions : FC ( Ffeat , Fpool ) = { f |f ( X ) = C∑ c=1 Λ̂c ( gc ( X ) ) , ( Λc , gc ) ∈ Fpool ×Ffeat } . ( 2 ) Conditions for universality We now describe two conditions that guarantee universality using the semi-lifted approach . The first step is showing , as in ( Yarotsky , 2018 ) , that continuousG-equivariant functions CG ( R3×n , WnT ) can be approximated by G-equivariant polynomials PG ( R3×n , WnT ) . Lemma 1 . Any continuous G-equivariant function in CG ( R3×n , WnT ) can be approximated uniformly on compact sets by G-equivariant polynomials in PG ( R3×n , WnT ) . Universality is now reduced to the approximation of G-equivariant polynomials . Next , we provide two conditions which guarantee that G-equivariant polynomials of degree D can be expressed by function spaces FC ( Ffeat , Fpool ) as defined in equation 2 . The idea behind these conditions is that explicit characterizations of polynomials equivariant to the joint action of translations , rotations and permutations is challenging . However , it is possible to explicitly characterize polynomials equivariant to translations and permutations ( but not rotations ) . The key observation is then that this characterization can be rewritten as a sum of functions to Wnfeat , a high dimensional representations of SO ( 3 ) which is equivariant to translations , permutations and rotations , composed with a linear map which is permutation equivariant ( but does not respect rotations ) . Accordingly , our first condition is that Ffeat contains a spanning set of such functions to Wnfeat . We call this condition D-spanning : Definition 1 ( D-spanning ) . For D ∈ N+ , let Ffeat be a subset of CG ( R3×n , Wnfeat ) . We say that Ffeat is D-spanning , if there exist f1 , . . . , fK ∈ Ffeat , such that every polynomial p : R3×n → Rn of degree D which is invariant to translations and equivariant to permutations , can be written as p ( X ) = K∑ k=1 Λ̂k ( fk ( X ) ) , ( 3 ) where Λk : Wfeat → R are all linear functionals , and Λ̂k : Wnfeat → Rn are the functions defined by elementwise applications of Λk . In Lemma 4 we explicitly construct a D-spanning family of functions . This provides us with a concrete condition which implies D-spanning for other function families as well . The second condition is that Fpool contains all SO ( 3 ) linear equivariant layers . We call this condition Linear universality . Intuitively , taking linear rotation equivariant Λk in equation 3 ensures that the resulting function p will be rotation equivariant and thus fully G-equivariant , and linear universality guarantees the ability to express all such G invariant functions . Definition 2 ( Linear universality ) . We say that a collection Fpool of equivariant linear functionals between two representations Wfeat and WT of SO ( 3 ) is linearly universal , if it contains all linear SO ( 3 ) -equivariant mappings between the two representations . When these two conditions apply , a rather simple symmetrization arguments leads to the following theorem : Theorem 1 . If Ffeat is D-spanning and Fpool is linearly universal , then there exists some C ( D ) ∈ N such that for all C ≥ C ( D ) the function space FC ( Ffeat , Fpool ) contains all G-equivariant polynomials of degree ≤ D. Proof idea . By the D-spanning assumption , there exist f1 , . . . , fK ∈ Ffeat such that any vector valued polynomial invariant to translations and equivariant to permutations is of the form p ( X ) = K∑ k=1 Λ̂k ( fk ( X ) ) , ( 4 ) While by definition this holds for functions p whose image is Rn , this is easily extended to functions to WnT as well . It remains to show that when p is also SO ( 3 ) -equivariant , we can choose Λk to be SO ( 3 ) equivariant . This is accomplished by averaging over SO ( 3 ) . As a result of Theorem 1 and Lemma 1 we obtain our universality result ( see inset for illustration ) Corollary 1 . For all C , D ∈ N+ , let FC , D denote function spaces generated by a pair of functions spaces which are D-spanning and linearly universal as in equation 2 . Then any continuous G-equivariant function in CG ( R3×n , WnT ) can be approximated uniformly on compact sets by equivariant functions in F = ⋃ D∈N FC ( D ) , D .

Summary: The authors introduce a framework for sufficient conditions for proving universality of a general class of neural networks that operate on point clouds which takes as input a set of coordinates of points and as output a feature for each point, such that the network is invariant to joint translation of the coordinates, equivariant to permutation of the points and equivariant to joint SO(3) transformations of the coordinates and output features of all points. Notably, this class contains Tensor Field Networks (TFN). The authors accomplish this by writing the network as a composition of an equivariant function from a class F_feat and followed by a linear pooling layer. When the F_feat class satisfies a “D-spanning” criterion and the pooling layer is universal, the network is universal. For a simple class of networks and for TFNs, the authors prove D-spanning. Linear universality of the pooling layer follows from simple representation theory.