Keypoints Tracking via Transformer Networks

For the image features extraction, we used the CNN model architecture from SuperPoint(SP) [2]. This model proved to be one of the best for representative feature extraction. For an image of size $H\times W$ it extracts $\nicefrac{{1}}{{8}}\times\nicefrac{{1}}{{8}}\times 256$ features. Each feature characterizes a $8\times 8$ patch around it, so we set its position to be the center of that patch. After the features are flattened, we form a set of $\nicefrac{{W\times H}}{{64}}$ pairs, where each entry $(f_{i},p_{i})$ represents the $i^{th}$ patch’s descriptor, and position. The process is visualized at fig. 3.

PUT FIGURE

For points’ feature descriptors, we firstly use SuperPoint to extract features from $Im_{1}$, and next based on points’ position we use bi-linear interpolation to get their descriptors. As a results we obtain sparse features descriptors from the $Im_{1}$ and a dense feature representation of the second image in its entirety.

During the next step, previously extracted features are passed to a positional encoding module, as well as to the Attentional aggregation modules. To provide a better understanding of Attention Layers, we firstly introduce the concept of transformer networks, and explain the intuition behind them. A Transformer encoder consists of sequentially connected encoder layers, the core element of the encoder layer is the Attention layer. Traditionally the names of inputs to it are : $Q$ - Queries, $K$ - Keys, $V$ - Values. Assuming that the size of all matrices is Nd, we can define attention as :

$$\text{Attention}(Q,K,V)=\operatorname{softmax}(QK^{T})V.$$

PUT FIGURE The intuition at this point is that attention operation learns feature level association by calculating similarity between Query and Key matrices. Another possible explanation is that by taking the $QK^{T}$ product we dynamically generate weights based on input of $Q$ and $K$ for linear layer with input $V$. In such a way that it is possible to increase the generalizability of the model as certain weights are generated on the fly and depend on the input.

The main problem of attention layers in Computer Vision field is that their complexity grows quadratically $(O(N^{2}))$ as the length of the input increases (this can be clearly seen from fig. 4:

the size of the output from the first MatMul operation is basically $N^{2}$). In computer vision the input mainly consists of pixels, patches, or their descriptors, so the complexity of the attention layer for images of size $(H,W)$ becomes proportional to $(HW)^{2}$. For example, even for $32\times 32$ images this number becomes   $10^{6}$. Recently there were several successful attempts to challenge this problem [3] by expressing self-attention as a linear dot-product of kernel feature maps. In this project, we are taking advantage of this method to speed-up, and reduce the memory consumption of the model. The difference (see fig. 4)

PUT FIGURE

is that in linear transformers instead of computing similarity between $Q$, and $K$ as $\operatorname{softmax}(QK^{T})$,it is approximated as $\phi(Q)\cdot\phi(K)^{T},\text{where }\phi(\cdot)=\operatorname{elu}(\cdot)+1$. Therefore, there is no need to compute a costly $\operatorname{softmax}$ operation, instead $\phi(K)^{T}\cdot V$ is computed, which allows us to avoid large $N\times N$ matrices, while preserving the order of matrix multiplication.

Positional encoding is of utter importance for any transformer-based model. In this work, we used Multi Layer Perceptron (MLP) to embed points’ 2d positions into vectors with the dimension same to descriptors’ dimension, and then add them together, as you can see at fig. 5.

PUT FIGURE

$$f_{i}=f_{i}+\operatorname{MLP}(p_{i})$$

After features from the second set are concatenated with a special OCL token, it can be treated as an additional feature from the $Im_{2}$. OCL token stands for occlusion token, it is used for finding occluded points, and outliers. This will be discussed in detail in further sections.

The attention aggregation module consists of several ($N_{L}$) self-attention, and cross-attention layers stacked together (see fig. 6).

The structure of the cross attention layer is shown on the fig (see fig. 4) .

The input is two sets of features after positional encoding $F_{1}$ and $F_{2}$ or $F_{2}$ and $F_{1}$ depending on the direction of cross-attention. Self-attention layer has the exact same structure, the only difference is that we use 2 identical sets of features: $F_{1}$ and $F_{1}$ or $F_{2}$ and $F_{2}$ ( note that we consider OCL token as a part of $F_{2}$). The primary goal of the self-attention layer is to share the information inside the same set of features, and the goal of the cross-attention layer is to share information between $F_{1}$ and $F_{2}$.

In order to establish the coarse matching between $F_{1}$ and $F_{2}$, we calculate the similarity score $s_{ij}$ for each pair of features $f_{1}^{i}$ and $f_{2}^{j}$ by computing their dot product. In such a way we build a similarity matrix $S$ of the size $M\times\frac{HW}{64}+1$, and next apply softmax along the X direction (along direction with $\frac{HW}{64}+1$ points ) to obtain the probabilities of matching for each of $M$ keypoints.

$$S_{ij}=<f_{1}^{i},f_{2}^{j}>,\forall(i,j)\in M\times\frac{HW}{64}+1$$

Now the problem became similar to a multi label classification problem, we aim to classify each point from $F_{1}$ to one of $\frac{HW}{64}+1$ classes. We select matches based on the highest similarity score. Optionally, it is possible to consider the match if the similarity score is higher than some threshold. Point is discarded if its confidence score is lower than threshold or if it’s best match is OCL token.

After establishing the coarse matching, we know the approximate position of the points on $Im_{2}$, and we assume that the exact match is located in the $8\times 8$ region around the matched feature. Since we are using bi-linear interpolation to extract the point descriptors, we can write $f_{1}^{j}$ as linear combination of $F_{1}^{i}$, where $F_{1}^{i}$ are the nearest descriptors to the point on the $Im_{1}$ (see fig. 7).

Combination of $\alpha_{i}$ - coefficients of bi-lienar interpolation, uniquely determines the relative position of the $f_{1}^{j}$ :

$$f_{1}^{j}=\sum_{i\in(1,4)}\alpha_{i}F_{i}^{i},\quad s.t.\sum{}\alpha_{i}=1$$

Thus we can assume that it is also possible to express $f_{1}^{j}$ as a linear combination of the nearest descriptors from $Im_{2}$ with some noise $n_{j}$:

$$f_{1}^{j}=\sum_{i\in(1,9)}\alpha_{i}F_{i}^{i}+n_{j}$$

Here we use 9 $F_{2}^{i}$ (see fig. 7)

FIG REFERENCE

as we do not know the exact location of the point on the second image. Again, we can assume that coefficients $\alpha_{i}$ determine the points’ positions. Thus based on these assumptions, for each keypoint we use its original descriptor with representation after AAM, and take 9 nearest image descriptors from Im2. Next we process them using the Attention aggregation module, calculate the similarity scores, and use them to regress the relative position of the point (see fig. 8 for details)

that allows us to move to sub-pixel matching.

One of the biggest challenges of points and images matching is to deal with outliers, and occluded points. The popular approach to deal with this problem, in the context of deep learning based keypoints matching, is the Sinkhorn Algorithm[21]. After computing the similarity matrix, this approach formulates the matching problem as a bipartite assignment solving the optimal transport problem. However, it can not be used for point tracking, as there is a chance that two different points from $F_{1}$ are matched to the same patch from $F_{2}$ (multiple to one matching case). Therefore, we employ another strategy to discard occlusions: we concatenate the $F_{2}$ set with a learnable OCL token. Simply speaking, we just add one additional feature to $F_{2}$ that stands for patch outside the $Im_{2}$ . During the training stage, the model is taught to match all occluded points to OCL token. Intuitively, the coarse module processes the features of the OCL token in a way that it becomes quite similar to all features in $F_{1}$ and different from all features from F2. Thus, after computing the similarity matrix, all features $f_{1}$ that lack similar features $f_{2}$ are classified as outliers. As we discussed in the Coarse matching part, after calculating the similarity matrix, the problem can be formulated as a classification. Let assume that 20% of the keypoints are occluded, then the probability that one point is matched with an OCL token is 0.2, while its probability to be matched with any other feature is:

$$\text{Prob}(p_{i}\text{ is not occluded})=\frac{0.8*M*64}{HW},M<<WH$$

This means that our dataset is extremely unbalanced, and direct training will lead to classifying all points as occluded. Thus we divided the training into several stages, this will be discussed in detail in the Training strategy section.

The Synthetic dataset consists of image pairs. Specifically, each image consists of a background full of stripes, triangles, stars, etc., and a big cube located in the center of the image. The first, and second images differ by small displacement of background( 16 pixels), and huge displacement of the cube(  50 pixels ).

For this dataset, we select the corners of the geometric primitives as points to track. While the structure of the images is rather simple, it is still a considerably difficult task, as there is a need to track points with different levels of displacement simultaneously. Thus, such an approach disables the model to learn homographies instead of point tracking. Since the displacement of the cube is large compared to the displacement of the background, it is likely that some points on the second image naturally become occluded. Hence, we use those to train OCL token.

For training on the real images we select the COCO14[23] image dataset. For each image we 1) applied random projective transformation to the image, 2) selected regions with huge concentration of key points, 3) cropped patches of the same size that are centered at corresponding points, 4) selected key points on the first patch, and used them as P1; those of them that do not appear on the second patch are treated as outliers.

Firstly we trained the model on the Synthetic dataset without occluded points in order to “warm up” the model. For each keypoint $p_{1}^{i}$, the ground truth consists of the index of the patch where this keypoint is located - $y^{i}$, and its position on the second image - $p_{2}^{i}$. Based on the predicted similarity score $s^{i}$, it is possible to select predicted position on the second patch - $\tilde{p_{2}}^{i}$.

$$s_{i}=<f_{1}^{i},f_{2}^{j}>,\forall f_{2}^{j}\in F$$

Thus the total loss consists of Cross Entropy Loss between $s^{i}$ and $y^{i}$, also, we added the L2 loss between $p_{2}^{i}$ and $p_{2}^{i}$ in order to speed up the convergence time:

$$\mathcal{L}_{total}=\sum_{i\in(1,M)}CEL(s^{i},y^{i})+L2(p_{2}^{i},\tilde{p_{2}}^{i})$$

At this stage we used the same dataset, however now occluded points are also used, and the loss consists only of CEL:

$$\mathcal{L}_{total}=\sum_{i\in(1,M)}CEL(s^{i},y^{i})$$

This transition stage is extremely important, as ‘occluded class’ should be added only after the model is taught how to distinguish between all inliers due to reasons we discussed in the Coarse matching part.

Finally, we train the model on the dataset of real images. At this stage we applied random color jitter to the images to mimic the change of illumination, and used projective transforms of different difficulty levels. The loss again consists only of CEL:

$$\mathcal{L}_{total}=\sum_{i\in(1,M)}CEL(s^{i},y^{i})$$

For training the fine matching module, we freeze the weights of the Coarse module. Based on the results of coarse module, for each point $p_{1}^{i}$ and position of the patch where it is located ($p_{2}^{i}$), we predict their relative distance $d^{*}\in(-4,4)$, having the ground truth $d$. For this we use L2 loss:

$$\mathcal{L}_{total}=\sum_{i\in(1,M)}L2(d,d^{*})$$

Thus the result position of the keypoint becomes: $\tilde{p_{2}}^{i}+d^{*}$ .

In order to provide a fair comparison, we conducted the experiments in the following way: for each image in the dataset 1) we created a corresponding image by applying projective transforms, and color jitter 2) sampled 512 most representative keypoints from the first image - $kp_{1}$, and all keypoints from the second image - $kp_{2}$ 3) matched $kp_{1}$ to $kp_{2}$ using SuperGlue, and matched $kp_{1}$ to the entire $Im_{2}$ 4) compare the results. The table below summarizes the outcomes of our results (Table 1, Table 2).

During testing of our model without a fine matching module, we considered the predicted patch’s center as the prediction for keypoints’ positions. Since we put the distance threshold to be 6 pixels, the accuracy of the coarse module represents the accuracy of the matching point to the correct patch. The fine module on average adds  1.5% accuracy, the reason is that for the cases when the coarse module predicts a nearby patch instead of the correct one, the fine module may refine the predicted position in a way that it becomes closer to the correct patch. The results clearly show that our model reaches the same, and sometimes higher level of accuracy compared to SuperGlue, while producing a significantly larger number of correctly matched points

Given two images and input points from one of them, our current method finds the position of these points on the second image. To further improve the matching accuracy, we can additionally try to find location of the predicted points on the first image, in such a way it is possible to increase robustness, and establish semi-supervised learning technique. Another way to improve the current results is to experiment with different numbers of attention layers, and train models on image sequences from Megadepth or ScanNet Datasets. In such a way the models can learn different sets of parameters for indoor, and outdoor environments. Additionally, instead of separate training of coarse and fine modules, we could integrate it in a single training pipeline that can be used out of the box in SLAM. We plan to extend this work to journal or conference publication.