Accurate Point Cloud Registration with Robust Optimal Transport

1 Introduction . Shape registration is a fundamental but difficult problem in computer vision . The task is to determine plausible spatial correspondences between pairs of shapes , with use cases that range from pose estimation for noisy point clouds [ 14 ] to the nonparametric registration of high-resolution medical images [ 17 ] . As illustrated in Fig . 1 , most existing approaches to this problem consist of a combination of three steps , possibly fused together by some deep learning ( DL ) methods : ( 1 ) feature extraction ; ( 2 ) feature matching ; and ( 3 ) regularization using a class of acceptable transformations that is specified through a parametric or nonparametric model . This work discusses how tools derived from ∗Equal Contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . optimal transport ( OT ) theory [ 87 ] can improve the second step of this pipeline ( feature matching ) on challenging problems . To put these results in context , we first present an overview of related methods . 1 . Feature extraction . To establish spatial correspondences , one first computes descriptive local features . When dealing with ( possibly annotated ) point clouds , a simple choice is to rely on Cartesian coordinates ( x , y , z ) [ 3 , 26 ] . Going further , stronger descriptors capture local geometric and topological properties : examples include shape orientation and curvatures [ 21 , 96 ] , shape contexts [ 6 ] , spectral eigenvalues [ 84 , 70 ] and annotations such as color [ 76 ] or chemical fingerprints [ 43 , 113 ] . Recently , expressive feature representations have also been learned using deep neural networks ( DNN ) : see [ 16 ] and subsequent works on geometric deep learning . Generally , feature extractors are designed to make shape registration as unambiguous as possible . In order to get closer to the ideal case of landmark matching [ 11 ] , we associate discriminative features to the salient points of our shapes : this increases the robustness of the subsequent matching and regularization steps . 2 . Feature matching . Once computed on both of the source and target shapes , feature vectors are put in correspondence with each other . This assignment is often encoded as an explicit mapping between the two shapes ; alternatively , the vector field relating the shapes can be defined implicitly as the gradient of a geometric loss function that quantifies discrepancies between two distributions of features [ 35 ] : a ) A first major approach is to rely on nearest neighbor projections and the related chamfer [ 12 ] and Hausdorff distances [ 13 ] , as in the Iterative Closest Point ( ICP ) algorithm [ 7 ] . This method can be softened through the use of a softmax ( log-sum-exp ) operator as in the many variants of the Coherent Point Drift ( CPD ) method [ 81 , 73 , 72 , 44 ] , or made robust to outliers in the specific context of rigid and affine registrations [ 41 , 129 , 128 , 14 ] . b ) Alternatively , a second approach is to rely on convolutional kernel norms such as the Energy Distance [ 94 ] , which are also known as Maximum Mean Discrepancies ( MMD ) in statistics [ 48 ] . These loss functions are common in imaging science [ 88 ] and computational anatomy [ 115 , 21 ] but are prone to vanishing gradients [ 40 , 39 ] . c ) Finally , a third type of approach is to rely on optimal transport ( OT ) theory [ 87 ] and solutions of the earth mover ’ s problem [ 97 ] . This method is equivalent to a nearest neighbor projection under a global constraint of bijectivity that enforces consistency in the matching . On the one hand , OT has been known to provide reliable correspondences in computer vision for more than two decades [ 26 , 47 , 64 ] . On the other hand , it has often faced major issues of scalability and robustness to outliers on noisy data . As detailed below , the main purpose of this work is to overcome these limitations and enable the widespread use of OT tools for challenging registration problems . 3 . Regularization with a deformation model . The output of the two steps above is a non-smooth vector field that may not be suitable for downstream tasks due to e.g . tears and compression artifacts . As a third step , most registration methods thus rely on regularization to obtain plausible deformations . This process is task-specific , with applications that range from rigid registration [ 133 , 46 , 30 , 2 , 121 , 122 ] to free-form motion estimation [ 91 , 126 , 69 , 49 ] . In Sec . 3 , we address the interaction of OT matching layers with a varied collection of regularization strategies – from optimization-based spline and diffeomorphic models to DNNs . Recent progresses . Research works on shape registration combine ideas from the three paragraphs above to best fit the characteristics of computer vision problems [ 71 , 31 , 107 ] . Over the past few years , significant progress has been made on all fronts . On the one hand , ( geometric ) deep learning networks have been used to define data-driven feature maps [ 92 , 92 , 123 ] and multiscale regularization modules [ 126 , 68 , 108 ] , sometimes fused within end-to-end architectures [ 133 , 91 , 132 , 30 ] . On the other hand , nearest neighbor projections , kernel convolutions and transport-based matching strategies have all been generalized to take advantage of these modern descriptors : they can now be used in high-dimensional feature spaces [ 59 , 37 ] . Challenges . Nevertheless , state-of-the-art ( SOTA ) methods in the field still have important limitations . First , modern deep learning pipelines are often hard to train to “ pixel-perfect ” accuracy on nonsmooth shapes , with diminishing returns in terms of model size and training data [ 2 ] . Second , scaling up point neural networks to finely sampled shapes ( N > 10k points ) remains a challenging research topic [ 49 , 135 , 37 ] . Third , the impact of the choice of a specific feature matching method on the performance of deep learning models remains only partially understood [ 58 ] . Related works . Following major progress on computational OT in the mathematical literature [ 79 , 27 , 66 , 103 ] , improved modules for feature matching have attracted interest as a possible solution to these challenges . Works on sliced partial OT [ 9 ] and dustbin-OT [ 28 ] have shown that outliers can be handled effectively by OT methods for rigid registration , beyond the classic Robust Point Matching method ( RPM ) [ 47 , 26 ] . Going further , the Sinkhorn algorithm for entropy-regularized OT [ 27 , 64 , 104 , 65 ] has been studied extensively for shape registration in computational anatomy [ 36 , 45 ] and computer graphics [ 74 , 33 , 85 ] . The Gromov–Wasserstein distance has also been used for shape analysis [ 110 , 118 ] , albeit at a higher computational cost . These applications have driven interest in the development of a complete theory for Robust Optimal Transport ( RobOT ) , outlined in Sec . 2 , which handles both sampling artifacts and outliers [ 25 , 105 , 67 , 24 , 106 , 80 ] . Most recently , this framework has started to be used in shape analysis with applications to shape matching [ 36 ] , the segmentation of brain tractograms [ 38 ] and deep deformation estimation with the FLOT architecture [ 91 ] . Contributions . We build upon the work above to tackle challenging point cloud registration problems for scene flow estimation and computational anatomy . Our key contributions are : 1 . Accurate feature matching with scalable OT solvers . For the first time , we scale up RobOT for deep feature matching to high-resolution shapes with more than 10k points . To this end , we leverage the latest generation of OT solvers [ 39 , 35 ] and overcome significant issues of memory usage and numerical stability . This allows us to handle fine-grained details effectively , which is key for e.g . most medical applications . 2 . Interaction with task-specific regularization strategies . We show how to interface RobOT matchings with advanced deformation models . This is in contrast with e.g . the FLOT architecture , which focuses on the direct prediction of a vector field and can not be used for applications that require guarantees on the smoothness of the registration . 3 . Challenging new dataset . We release a large dataset of lung vascular trees that should be registered between inhalation and exhalation . This relevant medical problem involves large and complex deformations of high-resolution 3D point clouds . As a new benchmark for the community , we provide two strong baselines that rely respectively on global feature matching and on deep deformation estimation . 4 . Consistent SOTA performance . Our proposed models achieve SOTA results for scene flow on Kitti [ 78 , 77 ] and for point-cloud-based lung registration on DirLab-COPDGene [ 19 ] . Notably , we show that RobOT is highly suited to fine-tuning tasks : it consistently turns “ good ” matchings into nearly perfect registrations at an affordable numerical cost . Main experimental observations . Is OT relevant in the deep learning era ? To answer this question decisively , we perform extensive numerical experiments and ablation studies . We fully document “ underwhelming ” results in the Supplementary Material and distill the key lessons that we learned in the Deep-RobOT architecture ( Section 3.2 ) . This model relies on fast RobOT layers to cover for the main weaknesses of point neural networks for shape registration . It is remarkably easy to deploy and generalizes well from synthetic training data to real test samples . We thus believe that it will have a stimulating impact on both of the computer vision and medical imaging literature . 2 Robust optimal transport . This section introduces the mathematical foundations of our work . After a brief overview of Robust Optimal Transport ( RobOT ) theory , we discuss the main challenges that one encounters when using this framework for computer vision . To avoid memory overflows and numerical stability issues , we introduce the weighted “ RobOT matching ” : a vector field that summarizes the information of a full transport plan with a linear memory footprint . As detailed in the next sections , this representation lets us scale up to high-resolution shapes without compromising on accuracy . 2.1 Mathematical background . The assignment problem . If A = ( x1 , . . . , xN ) and B = ( y1 , . . . , yM ) are two point clouds in R3 with N = M , the assignment problem between A and B reads : Assignment ( A , B ) = min s : [ [ 1 , N ] ] → [ [ 1 , N ] ] 1 2N N∑ i=1 ‖xi − ys ( i ) ‖2R3 , where s is a permutation . ( 1 ) This problem generalizes sorting to R3 : if the points xi and yj all belong to a line , the optimal permutation s∗ corresponds to a non-decreasing re-ordering of the point sets A and B [ 87 ] . Robust optimal transport . Further , OT theory allows us to consider problems where N 6= M. Non-negative weights α1 , . . . , αN , β1 , . . . , βM > 0 are attached to the points xi , yj and account for variations of the sampling densities , while feature vectors p1 , . . . , pN and q1 , . . . , qM in RD may advantageously replace raw point coordinates xi and yj in R3 . Following [ 67 , 25 ] , the robust OT problem between the shapes A = ( αi , xi , pi ) and B = ( βj , yj , qj ) reads : OTσ , τ ( A , B ) = min ( πi , j ) ∈RN×M > 0 N∑ i=1 M∑ j=1 πi , j · 12‖pi − qj‖ 2 RD ( 2 ) + σ2 KL ( πi , j ||αi ⊗ βj ) ︸ ︷︷ ︸ Entropic blur at scale σ . + τ2 KL ( ∑ j πi , j ||αi ) ︸ ︷︷ ︸ π should match A. . . + τ2 KL ( ∑ i πi , j ||βj ) ︸ ︷︷ ︸ . . . onto B. , for any choice of the regularization parameters σ > 0 and τ > 0 . In the equation above , the Kullback-Leibler divergence KL ( ai||bi ) = ∑ ai log ( ai/bi ) − ai + bi is a relative entropy that penalizes deviations of a non-negative vector of weights ( ai ) to a reference measure ( bi ) . Parameters . The first regularization term is scaled by the square of a blur radius σ . This characteristic length quantifies the fuzziness of the probabilistic assignment ( πi , j ) between points xi and yj [ 35 ] . The last two regularization terms promote the matching of the full distribution of points A onto the target shape B : they generalize the constraints of injectivity and surjectivity of Eq . ( 1 ) to the probabilistic setting . They are scaled by the square of a maximum reach distance τ : this parameter acts as a soft upper bound on the distance between feature vectors pi and qj that should be matched with each other [ 38 , 105 ] . For shape registration , we use simple heuristics for the values of these two characteristic scales : the blur σ should be equal to the average sampling distance in feature space RD while the reach τ should be equal to the largest plausible displacement for any given feature vector pi . These rules are easy to follow if point features correspond to Cartesian coordinates xi and yj in R3 but may lead to unexpected behaviors if features are output by a DNN . In the latter case , we thus normalize our feature vectors so that ‖pi‖RD = ‖qj‖RD = 1 and pick values for σ and τ between 0 and 2 . Working with a soft , probabilistic transport plan . As detailed in [ 35 ] , scalable OT solvers for Eq . ( 2 ) return a pair of dual vectors ( fi ) ∈ RN and ( gj ) ∈ RM that encode implicitly an optimal transport plan ( πi , j ) ∈ RN×M with coefficients : πi , j = αiβj · exp 1σ2 [ fi + gj − 12‖pi − qj‖ 2 RD ] > 0 . ( 3 ) In the limit case where σ tends to 0 and τ tends to +∞ , for generic point clouds ( xi ) and ( yj ) with N = M and equal weights αi = βj = 1/N , πi , j is a permutation matrix [ 87 ] . We retrieve the simple assignment problem of Eq . ( 1 ) : πi , j = 1/N if j = s∗ ( i ) and 0 otherwise . However , in general the transport plan must be understood as a probabilistic map between the point distributions A and B that assigns a weight πi , j to the coupling “ xi ↔ yj ” . For shape registration , this implies that the main difficulties for using robust OT are two-fold : first , the coupling π is not one-to-one , but one-to-many ; second , the lines and columns of the transport plan π do not sum up to one . Notably , this implies that when τ < +∞ , the gradient of the OT cost with respect to the point positions xi is not homogeneous : we observe vanishing and inflated values across the domain [ 105 ] .

Summary: The paper proposes a computationally efficient point cloud registration method based on Optimal Transport. The DRobOT pipeline relies on a deep module that predicts parameters of a registration model, plus pre and post processing steps based on RobOT. The paper shows results for scene flow and point cloud lung registration.