Shape registration in the time of transformers
Year: 2022

Abstract:
In this paper, we propose a transformer-based procedure for the efficient registration of non-rigid 3D point clouds. The proposed approach is data-driven and adopts for the first time the transformer architecture in the registration task. Our method is general and applies to different settings. Given a fixed template with some desired properties (e.g. skinning weights or other animation cues), we can register raw acquired data to it, thereby transferring all the template properties to the input geometry. Alternatively, given a pair of shapes, our method can register the first onto the second (or vice-versa), obtaining a high-quality dense correspondence between the two. In both contexts, the quality of our results enables us to target real applications such as texture transfer and shape interpolation. Furthermore, we also show that including an estimation of the underlying density of the surface eases the learning process. By exploiting the potential of this architecture, we can train our model requiring only a sparse set of ground truth correspondences (10 ∼ 20% of the total points). The proposed model and the analysis that we perform pave the way for future exploration of transformer-based architectures for registration and matching applications. Qualitative and quantitative evaluations demonstrate that our pipeline outperforms state-of-the-art methods for deformable and unordered 3D data registration on different datasets and scenarios.

heading: 1 Introduction
text: Recent technological advancements of 3D acquisition pipelines have produced an abundance of available data. The direct consequence is the non-standardization of the acquisition process. Such technological democratization brings along a disparate amount of different representations, discretization, and arbitrary resolution. Given so, the request to align such data has become urgent. Furthermore, data-driven statistical approaches require aligned data to relate feature changes across the population, inferring underlying patterns.
The Computer Vision community has devoted an extraordinary effort in the last decades to address 3D objects analysis. A common way to approach this problem is to align the geometry of one
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
source S target T output
known shape to an incoming one. Such methodology is referred to as registration. Many different axiomatic pipelines have been proposed that address different kind of objects and domains. While many methods rely on the assumption that the shapes used in the registration task differ just by a rigid transformation, the non-rigid domain is far more complex and interesting. This category pertains organic objects (e.g., humans, animals, internal organs), which are particularly interesting as well. Non-rigid registration aims to align two geometries that may differ by bending and stretching of the geometry, which may also significantly modify its metric. This problem is even more complicated if the geometry representation is given just by a sparse point cloud.
However, an emerging field merges data with classical algorithmic problems, exploiting such statistics as regularization. Among the different learning approaches, recently, the use of the Attention mechanism has become significantly popular in NLP domain, being later transferred to Computer Vision applications. Such architecture is called Transformers, and they represent one of the most significant groundbreaking methodological advancement since the introduction of CNNs.
In this work, we aim to let non-rigid registration meet the transformers. Intuitively, we aim to use the transformer as a geometrical translator between two non-rigid point clouds. As the first element, we modified the attention mechanism, proposing to make it aware of the underlying density of the geometry. Hence, we apply such a mechanism in an autoencoder–like architecture, which takes a template point cloud as input and aims to modify its geometry to fit the target point cloud.
The proposed method achieves better results than several state-of-the-art competitors in the shape matching task. We show results on humans case, but also inter-class objects. In this second case, our method is trained in an unsupervised manner, showing the power of our attention mechanism to infer the underlying geometry. Also, thanks to the attention mechanism, we are able to interpret what the network considers relevant for the registration. Finally, we can target texture transfer and shape interpolation showing applicability in real tasks as demonstrated in 1.
Our contributions could be summarized as follows: (a) We propose the first transformer for non-rigid registration task, showing the advantages of translation paradigm; (b) We modified the attention mechanism to make it aware of the point cloud density and of the underlying geometry of a shape; (c) We significantly improve the state-of-the-art performances on different datasets and challenging scenarios.
All code and data is publicly available 1.
1https://github.com/GiovanniTRA/transmatching

heading: 2 Related work
text: Shape matching is a problem with a tradition of decades. For a complete overview, we refer to the surveys [54, 51]; below we cover the literature that more closely relates to our work.
Surface matching Early attempts to match non-rigid objects work under the assumption of near isometry. Such is the case, for instance, of blended intrinsic maps [27], which combine multiple conformal mappings with an additional penalty to preserve local areas. Similarly, several variants and applications of the functional maps framework [41] implicitly assume near isometries by requiring a special structure of the functional representation, or by means of dedicated regularizers [42, 40, 18, 38, 49] also designed for handling partial shapes [50, 11, 45]. A common drawback to these works is that they do not disambiguate the intrinsic shape symmetries; further, the surface connectivity introduces a structural bias which may affect the performance, as recently shown in [36]. An attempt to overcome these issues was proposed in [37], but with the extra assumption that the shapes to match are in the same pose. More recently, SmoothShells [15] proposes an iterative algorithm to recover dense correspondences by an alignment of intrinsic information. While these methods do not address 3D shape registration directly, they rely on the general idea that a correspondence can be recovered by aligning specialized, possibly high-dimensional embeddings of the shapes at hand.
Template-free registration A popular approach to solve for a matching between 3D objects is to align their geometries extrinsically via ICP-like procedures [5, 29, 2]. In fact, registration and matching are intimately related problems with different goals. While the matching problem aims to find a combinatorial solution, which indicates for each point its image on the target shape, registration looks for a spatial transformation of the geometry. If the two shapes have a significantly different discretization, the latter problem is less ambiguous than the former. ICP-based approaches iteratively solve the two subproblems in an alternating fashion, by finding a point-to-point correspondence and the best transformation that adheres to such correspondence. These methods do not convergence to a good solution if the input shapes are significantly misaligned. Similarly, Coherent Point Drift [39] and variants [26, 22] rephrase the registration problem as an alignment of probability densities.
Template-based registration
A different family of approaches make use of a given template, which is known a priori and is possibly parametric, toward which a given input shape is to be matched. Taking as an example the case of human bodies, it is common to model the surface deformation related to the subject identity using PCA [1, 3, 31, 44], while recent advancements in statistical data-science suggested that non-linear methods are more expressive to catch fine details of humans [48, 59, 9]. Similarly, the pose can be modeled by simple paradigms like Linear-Blend Skinning [31, 44], triangle deformations [3, 23], but also learning methods [59]. Efforts to register such templates to arbitrary target models have been carried out extensively by the community [23, 61, 32, 33]. However, the requirement of a template is not always easy to satisfy.
Learning methods With the rise of learning methods, several attempts have been made to introduce a statistical prior to the matching and registration process. For example, several extensions have been proposed to bring the functional maps formalism into a learning paradigm [30, 13, 53]. SmoothShells has also been extended to be data-driven [16]. The point cloud representation has received comparably less attention, mainly in a rigid alignment setting [24, 43, 52]. In the non-rigid domain, a seminal work is 3DCoded [19] that proposes a proper registration using an autoencoder architecture. However, having a fixed template forbids inter–class operations and limits the use of the point cloud structure. The constraint of having a fixed template has been relaxed in [21] but it still requires couple of isomorphic shapes in training. Recently, it has been proposed to learn a linearly-invariant embedding [34], but the method requires training two separate networks and relies on simple PointNets [46] which are not able to catch the fine details of the objects. Finally, a recent trend in geometric deep learning suggests that implicit representations may also be used for shape matching [6].
Transformer-based architectures Transformers have been first introduced in the context of neural machine translation by the pioneering work [56], and later procedeed to revolutionize the field of natural language processing [12, 47]. The success obtained in NLP inspired further work to employ transformers in computer vision [14], where they managed to outperform convolutional networks. Exploiting the input invariance property characterising the attention mechanism, many works have
naturally extended transformers to handle point clouds [20, 60, 17]. These works show promising results, but work only in the context of object classification and segmentation. Instead, [57] proposes a network to find the rigid alignment between two point clouds imitating ICP, and using a transformer architecture to infer the residual term. Differently, we aim to solve for non-rigid registration, which is a more general case.

heading: 3 Notation and general objective
text: 3D shapes We model 3D shapes as compact 2-dimensional Riemannian manifoldsM, possibly with boundary ∂M. In the discrete setting, we represent the manifoldM as an unorganized point cloud of nM vertices embedded in R3, and encoded in a coordinate matrix XM ∈ RnM×3.
source target registration
Shape registration The main objective of this paper is to introduce a data-driven approach to perform shape registration. Given a source shape S and a target shape T , respectively represented by the sets of vertices XS and XT , our goal is to find a corresponding 3D position for each point in S on the surface of T . An example is shown in the inset figure, where the underlying surface is visualized for reference only.
Our method does not assume the source S to be a fixed template, but can generalize to arbitrary shape pairs S and T ; this is in contrast, e.g., with [19], where the objective is to learn how to deform a fixed known template into another shape.
Attention One of the most influential ideas in the recent years, attention originated in the realm of natural language processing but has since gained traction in other fields, such as computer vision and signal processing, due to the vast increase in performance and interpretability exhibited in several tasks. At its heart, the attention mechanism allows learning models to encode latent relations between inputs, assigning higher importance, or "attention", to the parts deemed more relevant. In this sense, attention allows to efficiently capture context information as well as higher order dependencies.
Formally, given two generic input sequences X1 ∈ Rn×d and X2 ∈ Rm×d (for clarity of exposition we assume a constant embedding dimension d, but it is not a necessary assumption), the attention mechanism models the linearly encoded representation of the inputs as triplets of query, key, and value matrices, respectively Q ∈ Rn×d, K ∈ Rm×d, and V ∈ Rm×d. The attention score W is then defined as W = softmax ( QKT d−0.5 ) and used to compute the weighted mean of the value vectors, resulting in an output feature matrix A = WV . We refer to self–attention whenever X1 and X2 are the same object, and use the term cross–attention otherwise.

heading: 4 Method
text: Our method takes as input two point clouds XT and XS , with nT and nS points respectively, and deforms the points of the source shape XS to fit the geometry described by the target point cloud XT . To do so, we rely on a novel attention mechanism which considers the underlying geometry conveyed by the point cloud, rather than treating the points simply as elements of a set.
Surface Attention The classic attention definition, as introduced in Section 3, looks at the input data points simply as elements of a set, and uses the computed attention scores to perform a weighted sum of the value vectors. When the value vectors represent a sampling of a signal over a surface, however, the natural domain for the integration should be the surface itself, of which the weighted sum is just an approximation highly sensitive to the specific surface sampling (depending, for instance, on the acquisition method).
To overcome this limitation we propose to modify the attention mechanism to consider the portion of surface represented by each point, weighting the attention score by an estimated local area element.
In practice, for each point xi ∈ X , we estimate its area contribution as the inverse of the local point density: A(X)i = (|{xj ∈ X s.t. ‖xj − xi‖2 < r}|)−1, where | · | denotes the cardinality of a set, and r is a local radius (r = 0.05 in our experiments). Note that, we are not interested in the absolute value of the area elements, rather on the relative contribution of each point. The surface attention
score is thus defined as:
W̃i,j = eWi,jA(X)j∑ t e Wi,tA(X)t . (1)
As in the classical attention, output features are computed as W̃V . Figure 3 shows that the surface attention mechanism results in a more stable localization of the attention scores across different samplings of the surface.
In our architecture, each time attention is computed over point clouds, we use such formulation.
Architecture The architecture we propose, portrayed in Figure 2, is an iteratively conditioned autoencoder. The two main ingredients are an encoder, that maps a target point cloud XT into a latent space LST , and a decoder, that deforms a source point cloud XS to resemble the geometry of the target point cloud XT .
The encoder draws inspiration from [25], featuring a set of learnable parameters, we refer to them as latent probes LP . It presents multiple layers of cross attention which iteratively condition the latent probes LP with the embedding of XT . After each conditioning, the resulting latent space is further transformed by layers of self attention, feed forwards and residual connections. The encoder output is a set of latent vectors LST containing relevant information collected from the target point cloud XT .
The decoder is analogous to the encoder but the relationship between the latent space and the point cloud in the cross attention is reversed. That is, the embedding of the source point cloud XS is transformed and iteratively conditioned with the latent space LST produced by the encoder. This procedure induces a deformation of XS that, after a final MLP layer, aligns to the points of XT .
Training Given a training dataset equipped with a ground-truth correspondence, we train our network in a supervised setting. Starting from a set of shapes in correspondence, we minimize the reconstruction error using the standard reconstruction loss:
Lsup(XS , XT ) = ‖XT −D(XS , E(XT ))‖22 . (2)
Furthermore, in case a ground-truth correspondence is not available, our network can be trained in an unsupervised way using the Chamfer distance, between XT and YT = D(XS , E(XT )), as defined below: ∑
s∈XT
min t∈YT ‖s− t‖22 + ∑ t∈YT min s∈XT ‖t− s‖22 . (3)
Testing At test time, our network can register a point cloud to another instantaneously, with a single forward pass. In the following experiments that involve the computation of matching, the correspondence can be obtained just by looking for the Euclidean nearest-neighbor between XT and the output D(XS , E(XT )) in the 3D space.
Refinement The peculiar structure of our architecture allows us to refine the output during the testing procedure. This is achieved by minimizing the energy function chamfer(XT , D(XS , E(XT ))) with
respect to the latent vectors LST . In practice, we use the latent vectors produced by a single forward pass as initial guess for LST , and minimize the previous energy function using Adam optimizer.
As we show in Section 5, this step can significantly improve the shape matching results. We remark that this is possible only thanks to the registration formulation of our approach. Shape alignment in 3D space can be optimized continuously, while a point-to-point assignment is non-differentiable due to the combinatorial nature of the problem.

heading: 5 Experiments & Results
text: We evaluate the effectiveness of our architecture on a number of challenges. We begin by analyzing the key components of our model, motivating our architectural choices through ablation studies. Finally, we present our results in the context of matching, registration, and inter–class registration.

heading: 5.1 Experimental settings
text: Training data For all our experiments in the humans domain, we trained our method on the same shapes from the SURREAL dataset [55] used in [34]. It consists of 10000 point clouds for training. Each point cloud has 1000 points, which simulate a significantly sparse sampling of the original shape. During training we augment the data by randomly rotating shapes along the second axis. We also trained our model on the ShapeNet dataset [8]. This dataset is composed of 16881 point clouds representing 3D shapes from 16 different objects categories(from chairs to airplanes), that also do not share a ground truth correspondence. Hence, we trained our method in an unsupervised manner for this particular data. We sample 1024 points for each object and use the same train/test split as in [58].
We also compare our model on other human datasets, thus assessing the ability to generalize to data out of the training distribution. A popular dataset to analyze real identities and poses is FAUST [7], which is composed by ten subjects in ten different poses. We also used a 1000 points version of it, which we refer to as FAUST1K. To simulate the noise produced by a 3D acquisition pipeline, we considered the same data from [34] in which the points are perturbated by a Gaussian noise. Also, we challenge our method on SHREC’19 [36]. Such dataset is composed by 44 shapes which have different connectivities, poses and densities. In all the experiments we refer to our method as Our.
In our comparisons we train our model for 10000 epochs using Adam optimizer [28]. We use 32 latent probes of dimension 64, and 8 layers for both the encoder and the decoder.
Competitors We consider 3DCoded [19] (3DC) as our principal competitor. Similarly to us, it aims at deforming a shape into another using an autoencoder architecture. However, it assumes one of the two shapes to be a predefined template, limiting its generalization capability. Further, we consider the Linearly-Invariant Embedding approach of [34] (LinInv), which learns an high-dimensional embedding in which the shapes can be aligned by a linear transformation. It is based on PointNet, and do not exploit any local structural mechanism. Finally, we considered the Geometric Functional Maps definition [13], using DiffusionNet [53] as feature extractor (DiffNet). Similarly to [34], it learns to embed each shape point into a common higher-dimensional.
Both our method and 3DC can improve the registration with a post-processing refinement, we refer to these as 3DCR and OurR.

heading: 5.2 Ablation and analysis
text: Here we justify our choices on hyperparameters through an ablation study, and we present some insight on the key properties of our method with an in depth analysis.
Ablation study To investigate the different components of the proposed architecture, we run a batch of experiments in a reduced version of the SURREAL dataset, using the first 1000 shapes for training, and a different set of 2700 shapes for testing. We also test on FAUST1K.
We report the results in Table 1. We test for three possible ablations: number of latent space vectors, their dimension, and the number of encoder–decoder blocks. We remark that the while SURREAL share discretization and density with the training data, FAUST1K does not. This may explain the difference in performance we observe, and their almost inverse relationship. In fact, it seems that a low number of latent vectors tend to overfit the specific sampling seen during training, while a higher number seems to provide better generalization. A similar behavior can be observed for the other hyperparameters considered. Finally, to reach a good trade off between bias and variance, we choose the configuration that produces the minimum mean error on these two evaluation sets.
Analysis of the Surface attention Quantitative and qualitative results showing the importance of the surface attention mechanism compared to the standard point attention. This novel kind of attention we propose has the merit of being much more agnostic to the particular discretization and density of the given point cloud. We can observe this behavior clearly in Figure 3. Here we visualize the attention across different densities and discretization strategies and two different settings, one with surface attention and one with regular attention. With the regular attention mechanism the part of the point cloud attended shows erratic behavior, with different intensities and often even different part that gets attended. Surface attention completely solves this issue, enabling the architecture to achieve
greater generalization capacity and enjoying increased robustness, and allowing us to decrease the error on the full version of FAUST of more that 50%.
Multivariate Latent Space
Inspecting the cross attention of the decoder layer we can seize the impact of the latent probes in our learning. In the top row of Figure 4 we can see how the attention in the first layer of the decoder captures global information of the shape, while in the last layer (bottom row), attention puts its focus on small details of the shape.
The possibility of directly visualizing attention maps and the multivariate nature of the latent space permits further analysis. In Figure 5, in particular, we explore the structure of our latent space and how attention influences it. Similarly to recent non-rigid shape autoencoders [10, 35, 4] we can manipulate the latent vectors to generate new shapes. Specifically, we grab the two latent spaces output of the encoder by registering a source shape S to two different target shapes T1, T2. Starting from the left triad, the left most and the right most shape represents the registration of S to T1 and T2 respectively; the one in the middle is a obtained by linearly interpolating the two latent spaces. In the middle triad we add a twist to this procedure, namely we locate the latent vectors attending the most to the upper body of the shape and keep them fixed. A similar procedure is undertaken in the right triad, but this time focusing on the legs. From this experiments, it follows that our latent space is not only linearly navigable, meaning that a linear interpolation of the shapes encoded in the latent space preserve a reasonable semantics, but also, and most interestingly we might say, attention characterizes this space and directly allow for meaningful alteration, or preservation, of selected chosen characteristics.

heading: 5.3 Results and Applications
text: In this section we present results and application of our method. In particular, we show state of the art performance on the shape matching task as well as shape registration.
Matching One of the task we consider is that of matching. Given two generic point clouds we want to find correspondences between them. Our model approach this task in a natural and elegant way, by registrating one shape onto the other. It becomes trivial then to obtain correspondences through a nearest point search. Results are reported in Table 2. Our method consistently outperforms the state of the art by a solid margin. Furthermore, we notice that our method is endowed with a much greater ability to generalize. This can be noted on the SHREC’19 dataset, as visualzied in Figures 6. The quality of our matching enable us to achieve high quality texture transfer as shown in Figures 1 and 7.
Template Registration A classical problem in Computer Graphics is to register a given template, usually a triangular or polygonal mesh, to some acquired point cloud. This setup is a special instance our method, in which the point cloud to be given as input to the decoder remains constant. Regarding our competing methods, even if DiffusionNet and LinInv are not proper registration algorithm, we can move a template point to the corresponding point (as found by the matching algorithm) on the input point cloud. This partially explains why, even though performing the worst overall, achieve lower chamfer distance, since their error is somewhat bound. On the other hand, 3DC is trained
exactly in this fashion. Note also that 3DC is the only method that sees the template shape during the training phase. Even though we train our method on a different task, we manage to improve on the state of the art as can be seen in Table 3, without the need for any, altough possible, fine-tuning.

heading: 5.4 Unsupervised Registration and Interpolation
text: One of our main advantages is that we do not require a template. Fixing a common template is not trivial, if not possible at all, when dealing with very different objects. To show the ability of our method of dealing with this challenging scenario we trained a model to register pair of shapes belonging to possibly different object categories of ShapeNet, using the chamfer loss defined in Section 4.
We show in Figure 1 (2nd and 3rd row) two interpolation sequences between two airplanes and between a chair and a table, showing that our method is able to register different objects preserving a meaningful correspondence, represented by similar colors. The interpolated reconstructions are obtained by embedding the outermost shapes in the latent space through the encoder and then using the linearly interpolated latent vectors and the left-most shape as input to the decoder.

heading: 6 Conclusion
text: We propose the first transformer based architecture to tackle the problem of non–rigid registration. We introduce a novel surface attention mechanism better suited to exploit the local geometric priors of the underlying structure. Our method reaches state of the art performance in shape matching and shape registration without assuming any fixed template, and generalizes also to different and complex geometries, e.g. handling multiple classes of ShapeNet [8] simultaneously. The attention mechanism at the core of our architecture has the potential to enforce local control of the interpolation, as seen in Figure 5.
Limitations and future works. Our method shares a common drawback with most transformedbased architectures, requiring long training and post-processing time due to the nature of the refinement procedure. Further investigation is needed to explore the possibility to introduce additional priors on the attention to force a semantically meaningful localization and interpolation behavior.

heading: 7 Acknowledgment
text: This work is supported by the ERC Grant No. 802554 (SPECGEO) and the MIUR under grant “Dipartimenti di eccellenza 2018-2022”.

Review summaries:
This paper proposes a Transformer-based point cloud registration scheme. Given two point clouds, the target and an initial latent feature go through the encoder of a Transformer-like network, and the resulting latent vector is combined with the source and goes through the decoder to yield a transformed point cloud (that resembles the target). In the Transformer-like structure, iterative cross attention is used, which is inspired by [21]. There are some other details to handle the particular problem efficiently, e.g., density aware attention for handling the density differences in point clouds. The decoder and the loss function can also be used to refine the result with further optimization. The experiments show that the proposed method mostly achieves state-of-the-art results (sometimes with large gaps) for well-known datasets.

This paper introduces the use of transformers architecture for shape matching and the method is applied to 3D point clouds. The main innovation with respect to standard attention architecture on point clouds is the weight introduced in the attention mechanism which accounts for the local density of the shape. Experiments show a clear gain of the method compared to other state of the art method, both quantitatively and qualitatively.

The authors propose a novel transformer based approach for 3D shape registration for non-rigid objects. Transformers have been very successful in language modelling and recently also in vision applications. This paper introduces transformers for 3D shape registration which is an interesting step forward. The approach also outperforms competing baselines.