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: This paper addresses a challenging problem in 3D representation learning of point clouds, which aims to generalize the representations well to arbitrary orientations. The proposed CSGNN uses a convolutional framework to learn over concentric spherical feature maps, which is incorporated with both intra-sphere and inter-sphere information. The idea is interesting and reasonably well presented. In general, the proposed CSGNN is practical on spherical representation learning. However, there still remains some issues that should be addressed.