3D Meta-Registration: Meta-learning 3D Point Cloud Registration Functions

1 INTRODUCTION . The point cloud registration is defined as a process to determine the spatial geometric transformations ( i.e . rigid and non-rigid transformation ) that can optimally register the source point cloud towards the target one . In comparison to classical registration methods Besl & McKay ( 1992 ) ; Yang et al . ( 2015 ) ; Myronenko et al . ( 2007 ) , learning-based registration methods Liu et al . ( 2019 ) ; Balakrishnan et al . ( 2018 ) usually leverage a neural network-based structure to directly predict the desired transformation for a given pair of source and target point clouds . Recently based on the PointNet Qi et al . ( 2017a ) structure , Liu et al . proposed FlowNet3D Liu et al . ( 2019 ) to learn the points flow field to register two point clouds together . Balakrishnan et al . proposed VoxelMorph Balakrishnan et al . ( 2018 ) for aligning two volumetric 3D shapes . These methods achieved impressive performance for the registration task of 3D shapes/scenes . In comparison to iterative registration methods , learning-based methods have advantages in dealing with a large number of datasets since learning-based methods can transfer the registration pattern from one dataset to another one . However , there are two main challenges for the learning-based methods . Firstly , learning-based networks often require a large volume of data and a long learning period to acquire the ability to predict the desired geometric transformation to register 3D point clouds . Secondly , The generalization capacity can be greatly degraded if the distribution of the dataset in practice differs from the training dataset . As shown in Figure 1 , previous learning-based methods tend to learn a single 3D registration function ( learner ) for any pair of source and target point clouds . In this way , the 3D registration learner performs well from extensive training on a large number of labeled data . In contrast , we proposed to formalize the learning of a 3D registration function space as a meta-learning problem . More specifically , we define each task as learning a unique 3D registration function ( learner ) for each given 3D object . Besides , we design an auxiliary deep neural network as a meta-learner that can predict the prior over the respective 3D registration function space . The meta-learner is responsible for providing the optimal initialization of a 3D registration learner . In general , our meta-learning-based approach gains competitive advantages over existing generalization methods that our method can uniquely parameterize the 3D registration function for each pair of shapes to provide 3D point cloud registration . As shown in Figure 2 , our model includes two modules : 3D registration learner and 3D registration meta-learner . We observe that different regions of a 3D shape vary in their geometric structures which makes it more sense that we have a region-conditioned transformation instead of the shapeconditioned one . In this way , the proposed 3D registration learner includes shape embedding and region aware flow regression which uses multiple non-linear multi-layer perceptron ( MLP ) -based function to predict the transforms and weights for different regions respectively . The 3D registration meta-learner includes two parts . The first part is a variational auto-encoder that maps shape representations to a distribution over 3D registration tasks , which gives the priors over the 3D registration function space and the second part is to sample from the 3D registration function space to predict optimal initialization for the parameters of the 3D shape registration function . Our contributions are listed as below : • In this paper , to the best of our knowledge , it is the first time to formalize learning of a 3D point cloud registration function space as a meta-learning problem 3D computer vision . Under this circumstance , the 3D registration model can be quickly adapted to new point clouds with no or limited labeled training data . • In this paper , we propose a novel variational encoder that maps shape representations to a distribution over 3D registration tasks , contributing to robust learning of priors over the 3D registration function space . • In this paper , we observe that different regions of a 3D shape vary in their geometric structures which enables us to propose a region-aware flow regression module instead of the shape-conditioned one . • In this paper , we compared our 3D Meta-Registration to other state-of-the-art ones on widely used benchmark datasets and demonstrated superior registration performance over both seen and unseen data . 2 RELATED WORKS . 2.1 POINT CLOUD REGISTRATION . In comparison to classical methods Besl & McKay ( 1992 ) ; Yang et al . ( 2015 ) ; Myronenko et al . ( 2007 ) , learning-based methods have significant advantages in dealing with a large number of datasets by transferring the “ knowledge ” from registering training pairs to testing pairs . Based on the feature learning structure proposed by PointNet Qi et al . ( 2017a ) , Aoki et al . Aoki et al . ( 2019 ) proposed PointNetLK for rigid point cloud registration by leveraging Lucas & Kanade algorithm . Based on DGCNN Wang et al . ( 2019 ) , Wang et al . proposed Deep Closest Point Wang & Solomon ( 2019a ) for learning rigid point cloud registration and PR-Net Wang & Solomon ( 2019b ) for learning partial shapes registration . For non-rigid point cloud registration , Liu et al . proposed FlowNet3D Liu et al . ( 2019 ) to learn the points flow field for non-rigid point cloud registration . Wang et al . proposed FlowNet3D++ Wang et al . ( 2020 ) on top of FlowNet3D by adding geometric constraints and angular alignment to dramatically improve the alignment performance . 2.2 META-LEARNING METHODS . Meta-learning Andrychowicz et al . ( 2016 ) ; Schmidhuber ( 1992 ) ; Hospedales et al . ( 2020 ) refers to a subfield of machine learning that learning new concepts and skills much faster and more efficiently given only a small amount of dataset . Parameters prediction Finn et al . ( 2017 ; 2018 ) ; Lee & Choi ( 2018 ) is one of the strategies in meta-learning , which refers to a network trained to predict the parameters of another network so that the first network can encode the related information to the second network which makes the overall network more flexible and adaptive to a particular task . Recently , meta-learning approaches are widely used in computer vision tasks . MANN was proposed by Santoro et al . Santoro et al . ( 2016 ) to use an explicit storage buffer which is easier for the network to rapidly incorporate new information and not to forget in the future . Ravi et al . Ravi & Larochelle ( 2016 ) use the LSTM to learn an update rule for training a neural network in few-shot learning . Snell et al . Snell et al . ( 2017 ) proposed Prototypical networks for few-shot classification task which map the sample data of each class into a metric space by calculating the euclidean distance of prototype representations of each class . In the field of 3D computer vision , Littwin et al . Littwin & Wolf ( 2019 ) firstly use a deep neural network to map the input point cloud to the parameters of another network for the 3D shape representation task . Yang et al . Yang et al . ( 2020 ) proposed Meta3D which an external memory is used to store image features and their corresponding volumes . In this paper , we first propose a meta-learning based method with a 3D registration meta-learner , which can learn to learn the registration pattern more efficiently . 3 METHODS . We introduce our approach in the following sections . In section 3.1 , we state the problem of learningbased registration . Section 3.2 illustrates the 3D registration learner . The 3D registration metalearner is explained in section 3.3 . The definition of the loss function is discussed in section 3.4 . 3.1 PROBLEM STATEMENT . We need to define the optimization task firstly . For a given training dataset D = { ( Si , Gi ) } , where Si , Gi ⊂ R3 , Si is the source point cloud and Gi is the target point cloud . We assume the existence of a parametric function gθ ( Si , Gi ) = φ using a neural network structure , where φ is the transformation function ( point flow in this paper ) which deforms source point cloud towards the target point cloud . We call gθ 3D registration learner in this paper and θ is the set of weights in the 3D registration learner . For previous learning-based network structure , the θ is optimized using stochastic gradient descent based algorithm for a given dataset : θoptimal = argmin θ [ E ( Si , Gi ) ∼D [ L ( Si , Gi , gθ ( Si , Gi ) ) ] ] ( 1 ) where L represents a similarity measure . In comparison , we assume that the 3D registration learner g includes two sets of parameters : θt and θm . θt is pre-trained from training dataset , but θm is predicted by another parametric function fσ which is called 3D registration meta-learner in this paper . Similarly , we have the desired transformation function φ = g ( θt , θm ) ( Si , Gi ) and we have θm = fσ ( Si , Gi ) . For a given training data set , we have : θoptimalt , σ optimal = argmin θt , σ [ E ( Si , Gi ) ∼D [ L ( Si , Gi , g ( θt , fσ ( Si , Gi ) ) ( Si , Gi ) ) ] ] ( 2 ) 3.2 3D REGISTRATION LEARNER . The 3D registration learner includes two modules : shape embedding ( 3.2.1 ) and region aware flow regression ( 3.2.2 ) . We discuss them in the following subsections . 3.2.1 SHAPE EMBEDDING . For a given pair of input point clouds , we firstly learn to extract the shape features that captures representative and deformation-insensitive geometric information . Let ( Si , Gi ) denotes the source and target point clouds and Si , Gi ⊂ R3 . ∀x ∈ Si , we denote the feature of x as lx ∈ Rc . Following the recent architecture from PointNet++ Qi et al . ( 2017b ) and FlowNet3D Liu et al . ( 2019 ) , the first set convolution g1 : R3 → Rc is a non-linear MLP-based function g1 : lx = Maxpool { g1 ( xj ) } ||xj−x||≤r∧xj∈Si ( 3 ) where r is a predefined distance and Maxpool is an element-wise max pooling function . ∀x ∈ Si , we further concatenate the coordinates x with the learned feature lx and we denote it as [ x , lx ] ∈ R ( c+3 ) . Similarly , we can learn the feature of each point in the target point cloud Gi . The shape descriptor for source point cloud Si is : { [ x , lx ] } x∈Si and the shape descriptor for target point cloud Gi is : { [ x , lx ] } x∈Gi . 3.2.2 REGION AWARE FLOW REGRESSION . Based on the learned shape descriptors for both source point cloud { [ x , lx ] } x∈Si and target point cloud { [ x , lx ] } x∈Gi from previous section , in this section we introduce two more set convolution structures for point flow regression . We define the second set convolution g2 : R ( 2c+9 ) → R ( c+3 ) to learn the relation information between descriptors of source and target point clouds . g2 is a non-linear MLP-based function . ∀x ∈ Si , we denote relation tensor px as : px = Maxpool { g2 ( [ x , lx , yj , lyj , x− yj ] ) } ||yj−x||≤r∧yj∈Gi ( 4 ) where [ , ] denotes concatenation . Based on the learned relation feature { [ x , px ] } x∈Si from source and target descriptors , we define a set of region aware decoders { g3 , k } k=1,2 , .. , K , where K is the pre-defined region numbers in dataset D and g3 , k : R ( c+3 ) → R3 is a non-learn MLP-based function . We note the pre-trained weights in g3 as θt and the meta-learned weights in g3 as θm . The weights for g3 is the element-wise summation of θt and θm . ∀x ∈ Si , we have the estimated region probability score px , k and flow vx , k for each region . Notice that we use a softmax function to normalize the output to the estimated region probability score px , k . We have : px , k , vx , k = g3 , k ( [ x , px ] ) ( 5 ) Then we use the estimated region probability score px , k as the weight to further balance among the point flow vx , k for each region . We define the final point flow vx as : vx = K∑ k=1 px , kvx , k ( 6 ) Therefore , the transformed source shape S′i = { x+ vx } x∈Si .

Summary: This paper presents a new architecture for point cloud registration. Different from existing methods, the proposed architecture consists of two stages. The first one is called meta learner, which is used to predict a task distribution and sample the key parameters for the second stage - 3D registration learner. By separating the network into two different parts, the learned model may be better at generalization.   Experiments were conducted on ModelNet40, FlyingThings3D, and KITTI datasets. The results show the proposed method outperforms some existing methods, including FlowNet3D, FlowNet3D++, and HPLFlowNet.