Concentric Spherical GNN for 3D Representation Learning

Learning 3D representations of point clouds that generalize well to arbitrary orientations is a challenge of practical importance in problems ranging from computer vision to molecular modeling . The proposed approach is based on a concentric spherical representation of 3D space , formed by nesting spatially-sampled spheres resulting from the highly regular icosahedral discretization . We propose separate intra-sphere and inter-sphere convolutions over the resulting concentric spherical grid , which are combined into a convolutional framework for learning volumetric and rotationally equivariant representations over point clouds . We demonstrate the effectiveness of our approach for 3D object classification , and towards resolving the electronic structure of atomistic systems . 1 INTRODUCTION . 3D point cloud data appear in domains ranging from computer vision , geographic information systems , and molecular modeling . Learning suitable representations of point clouds for data-driven modeling is well-motivated by applications like automonous vehicles and molecular simulation . It is important and yet challenging to learn representations that generalize well to arbitrary orientations of point clouds , in an efficient and accurate manner . Spherical Convolutional Neural Networks ( CNNs ) were introduced in Cohen et al . ( 2018 ) and Esteves et al . ( 2018 ) to address the challenge of rotations for spherical images , by defining rotationally equivariant convolutions in harmonic-space . Convolutions equivariant to transformations ( such as translation ) underpins recognizable aspects of 2D CNNs such as effective filter sharing and feature localization , a key reason for their success Cohen & Welling ( 2016 ) . However , Spherical CNNs operate over data projected onto the spherical manifold , which is lossy for general 3D point cloud data . It is desirable to learn features volumetrically , in order to accurately detect patterns in complex 3D shapes , or incorporate spatial relationships between atoms in describing atomic environments . We show that it is more expressive and general to operate over an underlying spatial representation of concentric spheres , demonstrated experimentally . We propose a new spatial structure consisting of multiple nested spheres , each discretized by the icosahedral grid . This structure is used to record both the angular and radial distribution of input points , as volumetric information . The icosahedral grid produces a highly regular sampling of the sphere , which permits efficient convolutions scaling linearly with spherical resolution , explored in Jiang et al . ( 2019 ) , Cohen et al . ( 2019 ) , and Defferrard et al . ( 2020 ) . We then propose two types of convolution for learning over the concentric spherical structure , by separately learning intra-sphere and inter-sphere features . We formulate intra-sphere convolution in terms of graphbased convolution over localized intra-sphere neighborhoods , and inter-sphere convolution as 1D convolution over co-radial grid points . The resulting convolutions are both rotationally equivariant and scalable , each scaling ( near ) linearly with respect to the size of the grid . Finally , we incorporate the proposed convolutions into a hierarchical and multi-resolution architecture , CSGNN , for learning over concentric spherical feature maps , and demonstrate its applicability to point cloud data . We first apply our approach to the problem of classifying rotated 3D objects sampled by point clouds , achieving state-of-the-art performance on the task . We further apply our approach to a problem of molecular modeling , namely predicting the electronic density of states of materials . The density of states is a fundamental property of electronic structure , used in determining total energy contributions . Our approach is applied to learn localized descriptions of atomic environments , enabling more accurate resolution of the band energy of carbon-based materials compared to previous approaches . The implementation of our methods and experiments are publicly available 1 . In summary , our primary contributions are as follows : 1 . We propose a new volumetric representation as the basis for convolutional learning . This representation consists of multiple nested spheres , each discretized by the icosahedral grid . 2 . We introduce a novel architecture for learning volumetric representation over concentric spheres by combining intra-sphere and inter-sphere convolutions . The proposed convolutions are rotationally equivariant , and also scale ( near ) linearly with grid resolution . 3 . We demonstrate the applicability of our approach through experiments in 3D object classification , and resolving electronic structure of atomistic systems . 2 BACKGROUND AND RELATED WORK . The goal of learning representations of general 3D point cloud data has led to a diverse body of work for structured learning ( Maturana & Scherer , 2015 ; Qi et al. , 2017a ; b ; Wang et al. , 2019 ; Thomas et al. , 2019 ; Zhang et al. , 2019b ) . However , the main problem shared by these methods is that they do not generalize well to general rotations of the data , which can lead to catastrophic loss of performance when they are encountered . Augmenting training with rotated data helps bridge the gap somewhat , but a significant performance gap remains . The key missing piece in many earlier work is their lack of rotationally equivariant model design . A model layer is equivariant to rotation if it commutes with rotation . In other words , feeding a rotated input to the model layer is same as feeding the original input to the layer and rotating its output . A rotationally invariant layer is a special case of an equivariant layer where rotation does not affect the output space of the layer . Several models have been proposed which exclusively use rotationally invariant layers ( Schütt et al. , 2017 ; Gilmer et al. , 2017 ; Schütt et al. , 2017 ; Chen et al. , 2019 ; Zhang et al. , 2019a ; Poulenard et al. , 2019 ; Kim et al. , 2020 ) . However , using invariant layers through the entirety of model is unnecessarily restrictive , as important information about the underlying spatial structure of the data may be lost . Our work focuses on designing rotationally equivariant layers as primary building blocks , while invariant layer ( s ) can be used before final output to achieve overall invariance of the model . This design has already seen extension to other structures and learning strategies relevant to point cloud representation . Thomas et al . ( 2018 ) proposed rotation equivariant point-wise convolutions for 3D graphs , but their approach has difficulty scaling to point clouds beyond the scale of small molecules . Equivariant design has also seen extension to spherical images through Spherical CNNs ( Cohen et al. , 2018 ; Esteves et al. , 2018 ; Jiang et al. , 2019 ; Cohen et al. , 2019 ; Rao et al. , 2019 ; Defferrard et al. , 2020 ; Yang et al. , 2020 ) . However , they are not well-suited for direct application to point clouds . Their key limitation is loss of information in constraining spatial representation from 3D domain to a 2D ( spherical ) manifold . Our work overcomes this limitation by proposing convolutional learning over concentric spheres , achieving rotationally equivariant 3D feature learning with scalable convolutions . 3 ARCHITECTURE DESIGN . The primary goal of our proposed approach is to learn volumetric representations of 3D point clouds in a rotationally equivariant and also scalable manner . To achieve this goal , we propose concentric spheres at different radii , each discretized by the icosahedral grid . The proposed construction naturally organizes 3D data by angular and radial distribution , where the resolution of each component can be controlled independently . We use the icosahedral grid as it results in a highly regular sampling of the sphere . The former permits efficient use of spatial resolution , and the later results in design of efficient and rotationally-equivariant convolutions . We propose using two separate convolutions together to learn volumetric features over concentric spheres : ( 1 ) graph-based convolution to incorporate information within spheres , and ( 2 ) radial convolutions to separately incorporate information between spheres . The proposed convolutions are extended to different spatial scales via pooling based on regular properties of the icosahedral grid , resulting in the proposed hierarchical convolutional architecture of Fig . 1 . We explain each component of our model in more detail in subsequent sections , and pooling in A.5 of Appendix . 1https : //github.com/anonymous10521/CSGNN 3.1 CONCENTRIC SPHERICAL DISCRETIZATION AND POINT CLOUD MAPPING . In this section we explain in detail our method of volumetric discretization by concentric spheres . We further present our approach to converting arbitrary point cloud data to initial feature channels over this spatial structure . Concentric Icosahedral Spheres . The initial icosahedron has 12 vertices forming 20 equilateral triangular faces . To increase grid resolution , each face can be sub-divided , with resolution scaling as |V | = 10 ∗ 4l + 2 ( l is target discretization level ) . See Fig . 3a for an illustration of this process . We implement concentric spheres by stacking R identical icosahedral grids to form the radial dimension , shown in Fig . 3b . Assuming normalization to unit radius , we uniformly assign concentric spheres at radii [ 1R , 2 R , ... , 1 ] . Assuming single-channel feature map , the resulting grid is the matrix H ∈ RR×|V | , where each vertex is indexed by the sphere it belongs to , and its position on the sphere . The resulting volumetric representation has several noteworthy properties . First , the icosahedral discretization results in a highly regular spatial sampling within each sphere . This allows more efficient use of spatial resolution compared to polar grids used in earlier works ( Cohen et al. , 2018 ; Esteves et al. , 2018 ; You et al. , 2020 ) , which have have resolution bias towards the polar regions . Second , spatial resolution is not uniform across different spheres , as the proposed construction results in higher sampling density closer to the center . However , this non-uniformity difference is largely accounted for as a function of the radius of the sphere , and does not inhibit the design of efficient and rotationally equivariant convolutions . Point Cloud to Concentric Spheres . We now consider the problem of describing a point cloud P ∈ RN×3 in terms of concentric spherical feature map H ∈ RR×|V |×C , where C is number of channels . While the concentric grid representation is defined discretely , the space point positions are continuous , which we would like to capture . To do so , we summarize the contribution of points using the Gaussian radial basis function ( RBF ) : f ( x ) = N∑ j=1 φ ( ||x− Pj ||22 ) ( 1 ) N is the number of data points , and the function φ = exp ( −γr2 ) is parameterized by the bandwidth γ . In practice we limit computation to a local neighborhood ( instead of considering all points ) and choose γ accordingly . We refer to Fig . 2 for illustrative example of the conversion , and Sec . A.6 for implementation details as applied in experiments . We determine the input channels C by grouping concentric spheres and concatenating features of co-radial vertices , also shown in Fig . 2 . Suppose converting the point cloud results in a singlechannel feature tensor H′ ∈ RR′×|V | , where R′ is the initial number of spheres . Grouping the features from co-radial vertices across R groups results in feature tensor H ∈ RR×|V |×R ′ R , where R is the number of spheres represented spatially , and R ′ R is the number of spheres represented via input channels . The proposed grouping mechanism gives flexibility to significantly increase radial resolution without correspondingly increasing the size of the spatial grid , which persists across convolutions . This enables balancing between computational efficiency and the representational boost of having additional spheres .

Summary: The paper presents a concentric sphere representation and icosahedron-based spherical CNNs for 3D point clouds. The contributions are two folds. First, the concentric sphere representation enables learning features volumetrically. Second, two types of convolutions: intra-sphere and inter-sphere are combined towards rotation equivariant and scalable computations.