kNN-Res: Residual Neural Network with kNN-Graph coherence for point
  cloud registration
In this paper, we present a residual neural network-based method for point
set registration that preserves the topological structure of the target point
set. Similar to coherent point drift (CPD), the registration (alignment)
problem is viewed as the movement of data points sampled from a target
distribution along a regularized displacement vector field. While the coherence
constraint in CPD is stated in terms of local motion coherence, the proposed
regularization term relies on a global smoothness constraint as a proxy for
preserving local topology. This makes CPD less flexible when the deformation is
locally rigid but globally non-rigid as in the case of multiple objects and
articulate pose registration. A Jacobian-based cost function and
geometric-aware statistical distances are proposed to mitigate these issues.
The latter allows for measuring misalignment between the target and the
reference. The justification for the k-Nearest Neighbour(kNN) graph
preservation of target data, when the Jacobian cost is used, is also provided.
Further, to tackle the registration of high-dimensional point sets, a constant
time stochastic approximation of the Jacobian cost is introduced. The proposed
method is illustrated on several 2-dimensional toy examples and tested on
high-dimensional flow Cytometry datasets where the task is to align two
distributions of cells whilst preserving the kNN-graph in order to preserve the
biological signal of the transformed data. The implementation of the proposed
approach is available at https://github.com/MuhammadSaeedBatikh/kNN-Res_Demo/
under the MIT license.
Muhammad S. Battikh, Dillon Hammill, Matthew Cook, Artem Lensky
# kNN-Res: Residual Neural Network with kNN-Graph coherence for point cloud registration

###### Abstract

In this paper, we present a residual neural network-based method for point set registration that preserves the topological structure of the target point set. Similar to coherent point drift (CPD), the registration (alignment) problem is viewed as the movement of data points sampled from a target distribution along a regularized displacement vector field. While the coherence constraint in CPD is stated in terms of local motion coherence, the proposed regularization term relies on a global smoothness constraint as a proxy for preserving local topology. This makes CPD less flexible when the deformation is locally rigid but globally non-rigid as in the case of multiple objects and articulate pose registration. A Jacobian-based cost function and geometric-aware statistical distances are proposed to mitigate these issues. The latter allows for measuring misalignment between the target and the reference. The justification for the k-Nearest Neighbour(kNN) graph preservation of target data, when the Jacobian cost is used, is also provided. Further, to tackle the registration of high-dimensional point sets, a constant time stochastic approximation of the Jacobian cost is introduced. The proposed method is illustrated on several 2-dimensional toy examples and tested on high-dimensional flow Cytometry datasets where the task is to align two distributions of cells
whilst preserving the kNN-graph in order to preserve the biological signal of the transformed data. The implementation of the proposed approach is available at [https://github.com/MuhammadSaeedBatikh/kNN-Res_Demo/](https://github.com/MuhammadSaeedBatikh/kNN-Res_Demo/) under the MIT license.

## 1 Introduction

Point set registration is a widely studied problem in the field of computer vision but also arises in other fields e.g. bioinformatics as is discussed below. The problem involves aligning a deformed target set of \(d\)-dimensional points to another reference point set by applying a constrained transformation. This alignment allows for improved comparison and analysis of the two sets of points and is used in a variety of fields including object tracking, body shape modeling, human pose estimation, and removal of batch effects in biological data. [1, 2, 3, 4, 5]

Point set registration techniques are typically categorized based on two main properties, first, whether the technique is a correspondence-based or a correspondence-free technique, and second, whether the estimated transformation is rigid or non-rigid. Correspondence-based techniques require the availability of correspondence information (e.g. labels) between the two point sets, while correspondence-free, sometimes called simultaneous pose and correspondence registration, does not require such information and therefore is considered a significantly more difficult problem. Rigid registration techniques are also generally simpler. A rigid transformation is an isometric transformation that preserves the pairwise distance between points and such transformation is typically modeled as a combination of rotation and translation. Several rigid registration techniques have been proposed in [6, 7, 8, 9, 10, 11, 12, 13, 14]. Assuming the transformation is rigid, however, makes the types of deformations that could be handled quite limited. Non-rigid transformations allow for more flexibility; however, this makes the problem ill-posed as there are an infinite number of transformations that could align two point sets, thus, non-rigid registration techniques employ additional constraints.

### Problem Formulation

In this section, we formulate the alignment problem. Inspired by CPD [15], we view an alignment method as finding a map \(\phi\) that transforms data points sampled from an underlying distribution \(Q\) to distribution \(P\) in such a way that preserves the topological structure of data sampled from \(Q\). This is an ill-posed density estimation problem, therefore, we require an additional desiderium for \(\phi\) to be as simple as possible. In this context, we call a map \(\phi\) simple if it is close to the identity transformation. Importantly, this could be visualized as data points sampled from \(Q\) moving along a regularized displacement vector field \(F\).

More formally, we denote two sets of \(d\)-dimensional vectors (points), a ref
erence point set \(\mathbf{R}=\{\mathbf{x}_{1},\mathbf{x}_{2},...,\mathbf{x}_{n}\}\), and target point set \(\mathbf{T}=\{\mathbf{y}_{1},\mathbf{y}_{2},...,\mathbf{y}_{m}\}\), generated by a probability distributions \(P\) and \(Q\) respectively. Additionally, a \(k\)-Nearest Neighbour (kNN) graph is associated with (or constructed from) the set \(\mathbf{T}\) which must be preserved after transformation. A kNN graph for set \(\mathbf{T}\) is a directed graph such that there is an edge from node \(i\) to \(j\) if and only if \(\mathbf{y}_{j}\) is among \(\mathbf{y}_{i}\)'s \(k\) most similar items in \(\mathbf{T}\) under some similarity measure \(\rho\).

Thus, the goal of an alignment method, given the sets \(\mathbf{R}\) and \(\mathbf{T}\) in a matrix form of \(X\in\mathbf{R}^{n\times d}\) and \(Y\in\mathbf{R}^{m\times d}\) respectively, is finding a transformation \(\phi\) parameterized by \(\theta\) such that:

\[\hat{\theta}=\arg\max_{\theta}D(\phi(Y;\theta),X) \tag{1}\]

subject to the constraints:

\[\texttt{kNN}_{g}(\phi(Y;\theta))=\texttt{kNN}_{g}(y) \tag{2}\]

where \(D\) is a statistical distance that measures the difference between two probability distributions.

### Limitations of existing approaches

A classic example of a such constraint is found in CPD [15] and its extensions [16, 17, 18]. CPD uses a Gaussian Mixture Model to induce a displacement field from the target to source points and uses local motion coherence to constrain the field such that nearby target points move together. CPD achieves this however via a global smoothing constraint which makes it locally inflexible, and therefore unsuitable for articulated deformations in 3D human data, scenes with multiple objects, and biological data [19].

In this work, we introduce a Jacobian orthogonality loss and show that it is a sufficient condition for preserving the kNN graph of the data. Jacobian orthogonality introduced as a penalty \(|\mathbf{J}_{\mathbf{X}}^{\top}\mathbf{J}_{\mathbf{X}}-\mathbf{I}_{d}|\) where \(\mathbf{J}_{\mathbf{X}}\) is the Jacobian matrix at a point \(\mathbf{x}\) and \(\mathbf{I}_{d}\) is the \(d\times d\) identity matrix. The penalty has been proposed in other contexts as well, such as unsupervised disentanglement [20] and medical image registration [21, 22].

In [21], the finite difference method is employed to compute the Jacobian penalty for the B-splines warping transformation, and mutual information of corresponding voxel intensities is used as the similarity measure. Instead of using finite difference for the Jacobian penalty, which produces a numerical approximation of first-order derivatives, the authors of [22] derive an analytical derivative specific to the multidimensional B-splines case. Such approaches however are limited to low dimensions by the nature of the transformations used, the way in which the Jacobian penalty is computed, and their proposed similarity measures.
### Contributions

To address these limitations, we use Hutchinson's estimator [20, 23] for fast computation of the Jacobian loss for high-dimensional point clouds, a scalable residual neural network (ResNet) [24] architecture as our warping transformation, and geometry-aware statistical distances. The choice of ResNet with identity block \(\phi(x)=x+\delta(x)\) is natural since we view alignment similar to CPD as a regularized movement of data points along a displacement vector field; which in our case is simply \(\phi(x)-x=\delta(x)\). It is also worth mentioning that ResNets can learn identity mapping more easily. Further discussion on this choice is given in section 2.2. Moment-matching ResNet(MM-Res) [5] use a similar ResNet architecture with RBF kernel maximum-mean discrepancy as its similarity measure [25, 26], however, no topological constraints are provided to preserve the topological structure of the transformed data nor to limit the nature of the learned transformation as shown in Figure 1. Additionally, while maximum-mean discrepancy is a geometry-aware distance, we address some limitations by incorporating Sinkhorn divergence into our framework [27].

Figure 1: Stanford Bunny example showing the effect of the Jacobian penalty on the learned transformation.
To elaborate further, we first start by defining Maximum Mean Discrepancy (MMD):

\[\texttt{MMD}(\alpha,\beta):=\frac{1}{2}\int_{X^{2}}k(\mathbf{x},\mathbf{y})d \zeta(\mathbf{x})d\zeta(\mathbf{y}) \tag{3}\]

where \(\alpha,\beta\in M_{1}^{+}(X)\) are unit-mass positive empirical distributions on a feature space \(X\), \(\zeta=\alpha-\beta\), and \(k(\mathbf{x},\mathbf{y})\) is a kernel function. MM-Res uses an RBF kernel which is suitable for high-dimensional Euclidean feature spaces (e.g. to represent \(X\subset\mathbb{R}^{n}\)) and makes training complexity low as it scales up to large batches, nonetheless, such kernel blinds the model to details smaller than its standard deviation, and the networks' gradient suffers from the well-known vanishing gradient problem. One simple solution is to decrease the standard deviation of the kernel; however, this introduces another issue, namely, the target points will not be properly attracted to source points [28]. In practice, this makes such a framework incapable of learning simple deformations with sizable translations as we show in section 4. Optimal transport (OT) losses do not typically suffer from this issue and produce more stable gradients; however, such losses require solving computationally costly linear programs. A well-known efficient approximation of the OT problem is entropic regularized \(OT_{\epsilon}\)[29], for \(\epsilon>0\), it is defined as:

\[\texttt{OT}_{\epsilon}(\alpha,\beta):=\min_{\pi_{1}=\alpha,\pi_{2}=\beta}\int _{X^{2}}C(\mathbf{x},\mathbf{y})d\pi+\epsilon\texttt{KL}(\pi|\alpha\times\beta) \tag{4}\]

where \(C\) is a cost function (typically the Euclidean distance), \((\pi_{1},\pi_{2})\) denotes the two marginals of the coupling measure \(\pi\in M_{1}^{+}\) and KL is the KL-divergence. The solution for this formulation could be efficiently computed using the Sinkhorn algorithm as long as \(\epsilon>0\). It is clear that by setting \(\epsilon\) to 0, this minimization problem reduces back to standard OT. Sinkhorn divergence combines the advantages of MMD and OT and is defined as:

\[S_{\epsilon}(\alpha,\beta)=\texttt{OT}_{\epsilon}(\alpha,\beta)-\frac{1}{2}( \texttt{OT}_{\epsilon}(\alpha,\alpha)+\texttt{OT}_{\epsilon}(\beta,\beta)) \tag{5}\]

The authors of [29] show that:

\[\lim_{\epsilon\to 0}S_{\epsilon}(\alpha,\beta)=\texttt{OT}(\alpha,\beta) \tag{6}\]

and

\[\lim_{\epsilon\rightarrow\infty}S_{\epsilon}(\alpha,\beta)=\frac{1}{2} \texttt{MDD}_{-C}^{2}(\alpha,\beta) \tag{7}\]

where \(C\) is the kernel used by MMD.

In the following section, we review other related methods.

### Other related work

Several point cloud registration approaches have been proposed. Thin plate spline functions-based techniques preserve the local topology via local rigidity on the surface of a deformed shape; however, these approaches are not scalable
to large datasets and are typically limited to 3-dimensional point clouds [30, 31, 32, 33, 34, 35]. To address these limitations, a deep learning paradigm for point cloud registration has been adopted. Deep learning-based approaches can be divided into two categories, namely, features-based, and end-to-end learning. In features-based methods, a neural network is used as a feature extraction. By developing sophisticated network architectures or loss functions, they aim to estimate robust correspondences by the learned distinctive feature [30, 36, 37, 38]. While feature-based learning typically involves elaborate pipelines with various steps such as feature extraction, correspondence estimation, and registration, end-to-end learning methods combine various steps in one objective and try to solve the registration problem directly by converting it to a regression problem [39, 40]. For example, [39] employs a key point detection method while simultaneously estimating relative pose.

Another class of methods is Graph Matching techniques, which are quadratic assignment problems (QAP) [40]. The main challenge for such methods is finding efficient approximate methods to the otherwise NP-hard QAP. Congruent Sets Gaussian Mixture (CSGM) [41] uses a linear program to solve the graph-matching problem and apply it to solve the cross-source point cloud registration task. Another approach is a high-order graph [42] that uses an integer projection algorithm to optimize the objective function in the integer domain. Finally, Factorized Graph Matching (FGM) method [43] factorizes the large pairwise affinity matrix into some smaller matrices. Then, the graph-matching problem is solved with a simple path following optimization algorithm.

## 2 Proposed model

### Methodology

In our case, we parametrize the transformation \(\phi\) as a residual neural network and formulate the optimization problem as:

\[\mathcal{L}(\theta)=\mathcal{L}_{1}+\lambda\mathcal{L}_{2} \tag{8}\]

where \(\mathcal{L}_{1}\) is the alignment loss \(D(\theta(Y;\theta),X)\) and \(\lambda\) is a hyperparamater, and \(\mathcal{L}_{2}\) is the topology preserving loss:

\[\mathcal{L}_{2}=\frac{1}{m}\sum_{\mathbf{y}\in T}|\mathbf{J}_{X}^{\top} \mathbf{J}_{X}-\mathbf{I}_{d}| \tag{9}\]

where \(\mathbf{J}_{y}\) is the Jacobian matrix at points \(y\) and \(\mathbf{I}_{d}\) is the \(d\times d\) identity matrix. In section 2.4 we prove that the orthogonality of the Jacobian matrix is indeed a sufficient condition for preserving the kNN graph of the data. We use two statistical distances, namely, Sinkhorn divergences, and maximum mean discrepancy. Sinkhorn divergences is a computationally efficient approximation for the Wasserstein distance in high dimensions and converge to the maximum mean discrepancy.
\[\mathcal{L}_{1}(\theta)=S(\alpha,\beta)=\texttt{OT}_{2}(\alpha,\beta)-\frac{1}{2}( \texttt{OT}_{2}(\alpha,\alpha)+\texttt{OT}_{2}(\beta,\beta)) \tag{10}\]

where \(\texttt{OT}_{\epsilon}\) is the optimal transport with \(\mathcal{L}_{2}\)-norm cost, and \(\alpha\) and \(\beta\) are measures over reference and target distributions respectively. The measures \(\alpha\) and \(\beta\) are unknown and are only known via samples from \(\mathbf{R}\) and \(\mathbf{T}\) respectively. Although \(S_{\epsilon}(\alpha,\beta)\) interpolates to MMD as \(\epsilon\) goes to infinity, we still maintain an efficient standalone MMD distance for data where MMD performs better than the Wasserstein distance and therefore no need for the interpolation overhead. Specifically, we use Gaussian-based MMD:

\[\texttt{MMD}(\alpha,\beta):=\frac{1}{2}\int_{X^{2}}k(\mathbf{x},\mathbf{y})d \zeta(\mathbf{x})d\zeta(\mathbf{y}) \tag{11}\]

### Architecture

We use a simple ResNet identity block with a skip connection as our transformation where the output dimension is equal to the input dimension, and the output is calculated as such: \(\phi(\mathbf{y};\theta)=\mathbf{y}+\delta(\mathbf{y};\theta)\), where \(\delta\) is a standard multi-layer perceptron (MLP) with LeakyRelu activation functions and \(\theta\) represents the trainable weights of the network. The ResNet identity block has been chosen for two reasons: biasing \(\theta\) to have small values via weight decay or initializing the output layer using a distribution with mean zero and a small standard deviation minimizes the contribution of \(\delta(\mathbf{y};\theta)\) to the final transformation which makes \(\phi(\mathbf{y};\theta)\) close to the identity. Additionally, this follows the same recipe from CPD of viewing the alignment function as a smooth displacement field.

The ResNet identity block is chosen for the following two reasons. Biasing \(\theta\) to have small values via weight decay or initialization using a distribution with close to zero values minimizes the contribution of \(\delta(\mathbf{x}:\theta)\) to the final transformation which in turn makes \(\phi(\mathbf{x}:\theta)\) close to the identity by design. Additionally, since we take a similar approach to CPD by viewing the alignment transformation as a regularized movement of data point along displacement vector field \(F\); having a ResNet identity block is mathematically convenient since a displacement vector is a difference between the final position \(\phi(\mathbf{x}:\theta)\) (transformed point) and the initial position (data point) \(\mathbf{x}\) such that \(F(\mathbf{x})=\phi(\mathbf{x}:\theta)-\mathbf{x}=\delta(\mathbf{x}:\theta)\), therefore, we only need to worry about \(\delta(\mathbf{x}:\theta)\) instead of \((\phi(\mathbf{x}:\theta)-\mathbf{x})\) absent skip connection.

### Orthogonal Jacobian preserves kNN graph

In this section, we show that the orthogonality of the Jacobian matrix evaluated at data points is a sufficient condition for preserving the kNN graph of the data. A vector-valued function \(\mathcal{F}:\mathbb{R}_{n}\rightarrow\mathbb{R}_{n}\) preserves the kNN graph of data points \(X\in\mathbb{R}_{n}\) if, for every two points \(\mathbf{v}\) and \(\mathbf{w}\) that are in some small \(\epsilon\)-neighborhood of \(\mathbf{u}\), the following holds:

\[||\mathbf{u}-\mathbf{v}||_{2}^{2}<||\mathbf{u}-\mathbf{w}||_{2}^{2} \rightarrow||F(\mathbf{u}),F(\mathbf{v})||_{2}^{2}<||F(\mathbf{u}),F( \mathbf{w})||_{2}^{2}, \tag{12}\]
where \(||\cdot||_{2}^{2}\) is the squared Euclidian distance. Without loss of generality, we choose two points \(\mathbf{w}\), \(\mathbf{v}\) that lie in \(\epsilon\) neighborhood of point \(\mathbf{u}\) and linearize the vector field \(F\) around point \(\mathbf{u}\) such that:

\[F(\mathbf{x};\mathbf{u})\approx F(\mathbf{u})+\mathbf{J}_{\mathbf{u}}(\mathbf{x }-\mathbf{u}), \tag{13}\]

where \(\mathbf{J}_{\mathbf{u}}\) is the Jacobian matrix evaluated at point \(\mathbf{u}\).

The squared distance of \(\mathbf{u}\) and \(\mathbf{v}\) is:

\[||\mathbf{u}-\mathbf{v}||_{2}^{2}=(\mathbf{u}-\mathbf{v})^{\top}(\mathbf{u}- \mathbf{v})=\sum_{i}^{n}\left(\mathbf{u}_{i}-\mathbf{v}_{i}\right)^{2} \tag{14}\]

Similarly, the squared distance between \(F(\mathbf{u};\mathbf{u})\) and \(F(\mathbf{v};\mathbf{u})\) computes as follows

\[\begin{array}{rcl}||F(\mathbf{u};\mathbf{u})-F(\mathbf{v};\mathbf{u})||_{2 }^{2}&=&(F(\mathbf{u};\mathbf{u})-F(\mathbf{v};\mathbf{u})^{\top}(F(\mathbf{u} ;\mathbf{u})-F(\mathbf{v};\mathbf{u}))\\ &=&F(\mathbf{u})-F(\mathbf{u})-\mathbf{J}_{\mathbf{u}}(\mathbf{v}-\mathbf{u})^ {\top}(F(\mathbf{u})-F(\mathbf{u})-\mathbf{J}_{\mathbf{u}}(\mathbf{v}- \mathbf{u}))\\ &=&(\mathbf{J}_{\mathbf{u}}(\mathbf{v}-\mathbf{u}))^{\top}(\mathbf{J}_{ \mathbf{u}}(\mathbf{v}-\mathbf{u}))\\ &=&(\mathbf{v}-\mathbf{u})^{\top}\mathbf{J}_{\mathbf{u}}^{\top}\mathbf{J}_{ \mathbf{u}}(\mathbf{v}-\mathbf{u})\\ &=&(\mathbf{v}-\mathbf{u})^{\top}(\mathbf{v}-\mathbf{u})\end{array}\]

The last step follows from the orthogonality of \(\mathbf{J}_{\mathbf{u}}\) i.e. \((\mathbf{J}_{\mathbf{u}}^{\top}\mathbf{J}_{\mathbf{u}}=\mathbf{I})\)

### Jacobian Loss Via Finite Difference

Given a vector-valued function \(F:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}\), a data batch \(X\in\mathbb{R}^{m\times d}\), and the Jacobian \(\mathbf{J}_{X}\) of \(F\) at points \(\mathbf{X}\) is an \(\mathbb{R}^{m\times d\times d}\) tensor, it is possible to compute \(\mathbf{J}_{\mathbf{X}}\) analytically using autodifferentiation modules, however, such computation is highly inefficient, thus, we use numerical approximation.
Given a \(d\)-dimensional vector \(\mathbf{x}=[x_{1},...,x_{d}]\), the partial first derivative of \(F\) with respect to \(x_{i}\) is:

\[\frac{\partial F}{\partial x_{i}}=\lim_{\epsilon\to 0}\frac{F(\mathbf{x}+ \epsilon e_{i})-F(\mathbf{x})}{\epsilon}, \tag{15}\]

where \(e_{i}\) is a standard basis vector (i.e. only the \(i\)th component equals 1 and the rest are zero). This could be approximated numerically using a small \(\epsilon\). The Jacobian matrix \(\mathbf{J}_{\mathbf{x}}\) is simply \(\lfloor\frac{\partial F}{\partial x_{i}},...,\frac{\partial F}{\partial x_{d}}\rfloor\). To ensure the orthogonality of the Jacobian at \(\mathbf{X}\), we minimize the following loss:

\[\mathcal{L}_{2}=\frac{1}{m}\sum_{\mathbf{x}\in\mathbf{X}}|\mathbf{J}_{\mathbf{ x}}^{\top}\mathbf{J}_{\mathbf{x}}-\mathbf{I}_{d}| \tag{16}\]

This process could be computed efficiently in a few lines of code as indicated in algorithm 1.

### Training

The training process (algorithm 2) takes advantage of two sets of \(d\)-dimensional vectors (points), a reference point set \(\mathbf{R}=\{\mathbf{x}_{1},\mathbf{x}_{2},...,\mathbf{x}_{n}\}\), and target point set \(\mathbf{T}=\{\mathbf{y}_{1},\mathbf{y}_{2},...,\mathbf{y}_{m}\}\). First, we sample points from \(\mathbf{R}\) and \(\mathbf{T}\) and create two matrices \(\mathbf{X}\) and \(\mathbf{Y}\). We feed \(\mathbf{Y}\) to the model and obtain \(\hat{\mathbf{Y}}\). Under the GMM assumption, we compute the GMM posterior probability as a similarity matrix and estimate \(\mathcal{L}_{1}\) as the negative log-likelihood. For the Sinkhorn divergence approach, we compute equation (10). We use the SoftkNN operator to construct the kNN graph for both the input \(\mathbf{Y}\) and the output \(\hat{\mathbf{Y}}\) and compute \(\mathcal{L}_{2}\) as the mean squared error between the two. Finally, we use backpropagation by minimizing the loss \(\mathcal{L}=\mathcal{L}_{1}+\lambda\mathcal{L}_{2}\) until convergence.

### Stochastic Approximation of Orthogonal Jacobian Loss

Using finite difference to compute the Jacobian for low-dimensional point clouds is efficient, however, the computational cost increases linearly with the dimension of the data. Thus, an approximate estimate with the constant computational cost is introduced.

Given a vector-valued function \(F\), and a sample \(\mathbf{x}\), we would like to minimize the following:

\[\mathcal{L}_{\mathbf{J}}(F)=|\mathbf{J}^{\top}\mathbf{J}\circ(1-\mathbf{I})| _{2}=\sum_{i\neq j}\frac{\partial F_{i}}{\partial x_{j}}\frac{\partial F_{j}} {\partial x_{i}} \tag{17}\]

Following [20, 23], the Hutchinson's estimator of \(\mathcal{L}_{\mathbf{J}}(F)\) can be approximated as such:

\[\mathcal{L}_{\mathbf{J}}(F)=\texttt{Var}_{r}(r_{\epsilon}^{\top}(\mathbf{J}^{ \top}\mathbf{J})r_{\epsilon})=\texttt{Var}_{r}((\mathbf{J}r_{\epsilon})^{\top }(\mathbf{J}r_{\epsilon})) \tag{18}\]

where \(r_{\epsilon}\) denotes a scaled Rademacher vector (each entry is either \(-\epsilon\) or \(+\epsilon\) with equal probability) where \(\epsilon>0\) is a hyperparameter that controls the granularity of the first directional derivative estimate and \(\texttt{Var}_{r}\) is the variance. It
is worth noting that this does not guarantee orthonormality, only orthogonality. In practice, however, we find that such an estimator produces comparable results to the standard finite difference method and could be efficiently implemented in Pytorch as shown in algorithm 3.

```
Input: \(\mathbf{R}\), and \(\mathbf{T}\) pointsets, blurring factor \(\sigma\), step size \(\epsilon\), regularisation \(\lambda\), and batch size \(b\); Output: Trained model \(\triangleright\) Simple mini-batches of size \(b\) from \(\mathbf{R}\) and \(\mathbf{T}\) while\((\mathbf{X},\mathbf{Y})\in(\mathbf{R},\mathbf{T})\) until convergencedo \(\phi(\mathbf{Y})\leftarrow\mathbf{Y}+\delta(\mathbf{Y})\); ifloss=="sinkhorn"then \(\mathcal{L}_{1}=\mathtt{S}(\mathbf{X},\phi(\mathbf{Y});\sigma^{2})\); else \(\mathcal{L}_{1}=\mathtt{MMD}(\mathbf{X},\phi(\mathbf{Y});\sigma^{2})\); \(\mathbf{J}_{\mathbf{Y}}[i,:]=\frac{\delta(\mathbf{Y}+\epsilon\epsilon_{i})- \delta(\mathbf{Y})}{\epsilon}\); \(\mathcal{L}_{2}=\frac{1}{m}\sum_{\mathbf{x}\in\mathbf{X}}|\mathbf{J}_{\mathbf{ x}}^{\top}\mathbf{J}_{\mathbf{x}}-\mathbf{I}_{d}|\); \(\mathcal{L}=\mathcal{L}_{1}+\lambda\mathcal{L}_{2}\); \(\triangleright\) backpropogation step Minimize(\(\mathcal{L}\));
```

**Algorithm 2**Training kNN-Resnet

### Parameters Selection

The proposed model has three main hyperparameters, namely: \(\sigma\), \(\epsilon\), and \(\lambda\). In the case of Sinkhorn divergence, \(\sigma>0\) is the blur (interpolation) parameter between OT and MMD, with a default value of \(0.01\) for datasets that lie in the first quadrant of the unit hypercube (minmax normalized data). Decreasing \(\sigma\) has the effect of solving for an exact OT, which typically produces very accurate registration, however, this comes at a slower convergence cost. In the cases where it is more advantageous to use MMD, \(\sigma\) represents the standard deviation of the Gaussian kernel. \(\epsilon>0\) represents the finite difference step size and controls the radius of topology preservation around each point. It is worth noting that a large epsilon value that covers all data tends to produce globally isomorphic transformations. \(\lambda>0\) is simply a regularization parameter that prioritizes regularization over alignment and is typically less than \(0.01\). An additional hyperparameter \(k\) is introduced when using a stochastic approximation of Jacobian orthogonality for high-dimensional data. This hyperparameter determines the number of Rademacher vectors sampled to estimate the Jacobian orthogonality penalty. Generally, a large \(k\) tends to produce a more accurate estimation, however; in practice, \(k=5\) seems to be a sufficient number for the datasets we experimented with.
```
1defstochastic_orth_jacobian(G,z,epsilon=0.01):
2''
3InputG:FunctiontocomputetheJacobianPenaltyfor.
4Inputz:(batchsize,d)InputtoGthattheJacobianis
5computedw.r.t.
6Inputk:numberofdirectionstosample(default5)
7Inputepsilon:(default0.1)
8Output:mean(\(|\mathbf{J}_{X}^{T}\mathbf{J}_{X}-\mathbf{I}_{d}|\))
9'
10r=torch.randint(0,2,size=torch.Size((k,*z.size()),))
11#r:rademacherrandomvector
12r[r==0]=-1
13vs=epsilon*r
14diffs=[G(z+v)-Gzforvinvs]
15#std:stochasticfinitediffs
16sfd=torch.stack(diffs)/epsilon
17loss=torch.var(sfd,dim=0).max()
18returnloss
19
```

**Algorithm 3**PyTorch code for Hutchinson approximation for Jacobian off-diagonal elements at data points \(z\).

## 3 Experiments

In this section, we provide experimental results on several datasets, namely, Chui-Rangarajan synthesized dataset used in [31, 44, 45], and single-cell RNA data used in [5]. The Chui-Rangarajan synthesized dataset is comprised of two shapes; a fish shape, and a Chinese character shape. Each shape is subjected to 5 increasing levels of deformations using an RBF kernel, and each deformation contains 100 different samples. The samples are generated using different RBF coefficients which are sampled from a zero-mean normal distribution with standard deviation \(\sigma\), whereby increasing \(\sigma\) leads to generally larger deformation.

### Results on 2D Data

We use the root-mean-squared error (RMSE) between the transformed data \(\hat{y_{i}}\) and the ground truth \(y_{i}\) available from the Chui-Rangarajan synthesized dataset: \(error=\sqrt{\frac{1}{m}\sum_{i=0}^{m}{(\hat{y_{i}}-y_{i})^{2}}}\).

It is important to note that such ground-truth correspondence is absent during training time and is only available during test time. Figures 2 and 3 show the initial point set distributions and their corresponding aligned versions for the Chinese character and the fish examples respectively. We also report results for our kNN-Res, MM-Res[5], CPD [15], TRS-RPM [31], RPM-LNS [45], and GMMREG [32] over 5 deformation levels and 100 samples per level. Figures 4b and 4b show results for tested models for the Chinese character, and Fish datasets respectively. We notice that after a certain level of
non-rigid deformation, MM-Res is unable to converge. For our kNN-Res, we set \(\epsilon=.005,\lambda=10^{-5},\sigma=.001\) and number of hidden units = 50. We start with a relatively high learning rate (0.01) for ADAM [46] optimizer and use a reduce-on-plateau scheduler with a reduction factor of 0.7 and minimum learning rate of \(5\times 10^{-5}\). Qualitatively, the grid-warp representations from the second column in figures 2 and 3 indicate that our estimated transformations are, at least visually, "simple" and "coherent". Furthermore, to quantitatively assess neighborhood preservation we use the hamming loss \(\mathcal{L}_{H}\) to estimate the difference between the kNN graph before and after transformation:

\[\mathcal{L}_{H}=\sum_{i=0}^{m}\sum_{j=0}^{m}I(\hat{p}_{i,j}^{k}\neq p_{i,j}^{k})\]

where \(p_{i,j}^{k}\) is the \(i\),\(j\) element of the k-NN graph matrix before transformation, \(\hat{p}_{i,j}^{k}\) is the corresponding element after transformation, and \(I\) is the indicator function. Figures 5b and 5a show the difference in neighborhood preservation between MM-Res and our kNN-Res for the Chinese character, and Fish datasets respectively for three different levels of deformations.

Moreover, despite the additional topology regularization term, our kNN-Res generally incurred smaller alignment errors and was able to converge under large deformation levels.

### Results on High-Dimensional CyTOF Data

Cytometry by time of flight (CyTOF) provides the means for the quantification of multiple cellular components data, however, is susceptible to the so-called batch effect problem, where systematic non-biological variations during

Figure 2: The Chinese character deformation example: Top row represents original and deformed sets, Mid row represents the vector field, and Bottom row is the final alignment.
the measuring process result in a distributional shift of otherwise similar samples. This effect breaks the intra-comparability of samples which is a crucial component of downstream tasks such as disease diagnosis and typically requires the intervention of human experts to remove these batch effects. The CyTOF dataset used in our experiments was curated by the Yale New Haven Hospital. There are two patients, and two conditions were measured on two different days. All eight samples have 25 markers each representing a separate dimension ('CD45', 'CD19', 'CD127', 'CD4', 'CD8a', 'CD20', 'CD25', 'CD278', 'TNFa', 'Tim3', 'CD27', 'CD14', 'CCR7', 'CD28','CD152', 'FOXP3', 'CD45RO', 'INFg', 'CD223', 'GzB', 'CD3', 'CD274', 'HLADR', 'PD1', 'CD11b'), and a range of cells (points) between 1800 to 5000 cells per sample. The split is done such that samples collected on day 1 are the target, and samples collected on day 2 are the target, and samples collected on day 3 are the target, and samples collected on day 4 are the target, and samples collected on day 5 are the target, and samples collected on day 6 are the target, and samples collected on day 7 are the target, and samples collected on day 8 are the target, and samples collected on
Figure 5: The figures show Hamming loss for the following levels of deformations: (left) level 1, (mid) level 2, (right) level 3.

Figure 6: The blue and red dots represent 1st and 2nd principal components of reference (patient #2 on day 2) and the target samples (patient #2 on day 1) correspondingly.
the reference, resulting in four alignment experiments.

We follow the exact preprocessing steps described in [5]. To adjust the dynamic range of samples, a standard pre-preprocessing step of CyTOF data is applying a log transformation [47]. Additionally, CyTOF data typically contains a large number of zero values (40%) due to instrumental instability which are not considered biological signals. Thus, a denoising autoencoder (DAE) is used to remove these zero-values [48]. The Encoder of the DAE is comprised of two fully-connected layers with ReLU activation functions. The decoder (output) is a single linear layer, with no activation function. All layers of the DAE have the same number of neurons as the dimensionality of the data. Next, each cell is multiplied by an independent Bernoulli random vector with probability =.8, and the DAE is trained to reconstruct the original cell using an MSE loss. Furthermore, the DAE is optimized via RMSprop and weight decay regularization. The zero values in both reference and target are then removed using the trained DAE. Finally, each feature in both target and reference samples is independently standardized to have a zero-mean and unit variance. For our kNN-Res, we set \(\epsilon=0.05,\lambda=0.1,\sigma=0.04\), \(k=5\) for Hutchinson's estimator, and the number of hidden units to 50. We start with a relatively high learning rate (0.01) for the ADAM optimizer and use a reduce-on-plateau scheduler with a reduction factor of.7, and a minimum learning rate of \(5\times 10^{-5}\). Figure 6 shows the first two principal components of data before and after alignment using two kNN-Res models with different lambdas. Although the two samples appear less aligned when using a large \(\lambda\), this comes with the benefit of preserving the topology of the data as shown by the learned transformation in figure 7 where points (cells) are moved in a coherent way.

This becomes more clearer when looking at the marginals in figure 13 in the appendix. In this experiment, we trained five models with five different

Figure 7: Point set transformation(alignment) of patient #2 sample on day 1 and day 2, shown in space of 1st and 2nd principal components.
lambdas ranging from 0 to 1. It is clear that having a small \(\lambda\) favors alignment over faithfulness to the original distribution, however, increasing \(\lambda\) preserves the shape of the original data after transformation, which is desirable in biological settings. For results of other experiments see Appendix.

## 4 Discussion

### Implications

Point-set registration methods are typically used for problems in computer vision to align point clouds produced by either stereo vision or by Light Detection and Ranging devices (e.g. Velodyne scanner) for instance to stitch scenes and align objects. These datasets are usually of 2 or 3 dimensions and hence the methods had limited exposure to high-dimensional datasets. Biological data, on the other hand, is usually of high dimension and hence methods from point-set registration do not directly translate to biological data. The proposed method in this study was tested on a 25-dimensional CyTOF dataset. However, in flow and mass cytometry, data could easily go beyond 50 dimensions (markers). For instance, methods that combine protein marker detection with unbiased transcriptome profiling of single cells provide an even higher number of markers. These methods show that multimodal data analysis can achieve a more detailed characterization of cellular phenotypes than transcriptome measurements alone [49, 50] and hence recently gained significant traction. Unfortunately, these methods require more sophisticated batch normalization algorithms, since manual gating and normalization using marginal distributions become infeasible. It is worth mentioning that even though the experts are making sure that the marginal distributions are aligned, there is still no guarantee that the samples are aligned in the higher dimensional space. Moreover, the alignment might result in such nonlinear and non-smooth transformations that break biological relationships or introduce non-existing biological variabilities. The proposed method mitigates these issues and guarantees smooth transformations.

### Limitations

It is clear from the last step of the proof that the orthogonal Jacobian is too strong a condition for preserving the kNN graph:

\[(\mathbf{v}-\mathbf{u})^{\top}\mathbf{J}_{\mathbf{u}}^{\top}\mathbf{J}_{ \mathbf{u}}(\mathbf{v}-\mathbf{u})=(\mathbf{v}-\mathbf{u})^{\top}(\mathbf{v}- \mathbf{u}) \tag{19}\]

The objective is satisfied by preserving inequality and not equality. In other words, it is only necessary and sufficient for \(\mathbf{J}\) to preserve the kNN graph if the following holds:

\[\mathbf{u}^{\top}\mathbf{u}\leq\mathbf{v}^{\top}\mathbf{v}\rightarrow\mathbf{ u}^{\top}\mathbf{J}^{\top}\mathbf{J}\mathbf{u}\leq\mathbf{v}^{\top}\mathbf{J} ^{\top}\mathbf{J}\mathbf{v} \tag{20}\]

or

\[\langle\mathbf{u},\mathbf{u}\rangle\leq\langle\mathbf{v},\mathbf{v}\rangle \rightarrow\langle\mathbf{J}\mathbf{u},\mathbf{J}\mathbf{u}\rangle\leq \langle\mathbf{J}\mathbf{v},\mathbf{J}\mathbf{v}\rangle \tag{21}\]
Having strict equality puts a limitation on the kind of transformations the model is capable of learning. Furthermore, even if the deformation could theoretically be expressed, such a penalty makes convergence unnecessarily slower. On the empirical side, we only have a limited number of experiments to test the proposed method. More experimentation and ablation are required to better understand the limits of our current approach and to learn how it fairs on a wider selection of real-world data such as RNA-Seq.

### Future Work

An important future direction is incorporating local or partial matching using modified alignment losses such as Gromov-Wasserstein distance. This should lead to a much more robust solution than global matching, especially in the case of outliers and missing objects. We also posit that solving point set registration under topological constraints such as preserving the kNN graph is naturally extendable to dimensionality reduction.

## 5 Conclusion

This paper presents a simple, scalable framework for point cloud registration. At its core, it consists of three components, namely (a) residual neural network with identity blocks as a parametrized displacement field, (b) Jacobian penalty as a topology-preserving loss, and (c) Sinkhorn Divergence as a sample-based, geometry-aware statistical distance. Additionally, by incorporating Hutchinson's estimator for the Jacobian loss, we show that our model is easily extensible to high dimensions with constant complexity. Furthermore, we offer both qualitative and quantitative analysis for synthetic and CyTOF datasets showing the flexibility and applicability of our model in multiple domains.
