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: This paper discusses the application of optimal transport theory in and contributed several useful practices to shape matching. First, the proposed RobOT method implements the traditional entropy-regularized unbalanced optimal transport using the kernel method and reduces the memory footprint to linear, enabling the processing of over 60k points. Second, the paper introduces several effective strategies to regularize/project the transport vectors given a deformation model (RobOTP). Third, the deep deformation framework (DRobOT) is proposed by sandwiching the parameter prediction module with pre-alignment and post-optimization steps, which lead to SoTA performance on a new lung dataset and scene flow datasets. Apart from the algorithmic side, the authors also give a clear summary which cover most of the vision algorithms using optimal transport.