SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
E _ {\text {a r a p}} = \sum_ {h \in [ H ]} \sum_ {l \in \psi (h)} \left( \right.\left\| d _ {h, l} (\mathbf {X}) \right\| _ {2} ^ {2} + \alpha \left\| \right.\left( \right.R \left(\Theta_ {h}\right) - R \left(\Theta_ {l}\right)\left. \right\| _ {2} ^ {2}\left. \right) \tag {10}
$$
Algorithm 1: Shape registration pipeline.
Input: Source mesh  $S = \{\mathcal{V},\mathcal{E}\}$  and target point cloud  $\mathcal{T}$ ; Trained point feature extractor  $\mathbf{F}$   
Output:  $\mathbf{X}^*$  converging to a local minimum of  $E_{\mathrm{total}}$ ; Deformed source model  $\{\mathcal{V}^*,\mathcal{E}\}$ ; Correspondence  $\Pi_{\mathcal{S}\mathcal{T}}^{*},\Pi_{\mathcal{T}\mathcal{S}}^{*}$  between  $\mathcal{S}$  and  $\mathcal{T}$ .
1 Initialize deformation graph  $\mathcal{DG}$  and  $\mathbf{X}^{(0)}$  with identity transformations;  $F_{\mathcal{T}} = \mathbf{F}(\mathcal{T})$ ;  $\mathrm{k} = 0$ ; 2 while True do
Update source vertex  $\nu^{(k)}$  by Eqn.(7);  
if  $k\% 100 == 0$  and Flag  $==$  Stage-I then  
 $\begin{array}{r}\big{\lfloor}F_{S}^{(k)} = \mathbf{F}(\mathcal{V}^{(k)});\Pi_{ST}^{(k)} = \mathbf{NN}(F_{S}^{(k)},F_{T});\Pi_{TS}^{(k)} = \mathbf{NN}(F_{T},F_{S}^{(k)}) \end{array}$   
if  $k\% 100 == 0$  and Flag  $==$  Stage-II then  
 $\begin{array}{r}\big{\lfloor}\Pi_{ST}^{(k)} = \mathbf{NN}(\mathcal{V}^{(k)},\mathcal{T});\Pi_{TS}^{(k)} = \mathbf{NN}(\mathcal{T},\mathcal{V}^{(k)});\end{array}$   
Compute the set of filtered correspondences  $\mathcal{C}^k$  .  
 $\mathbf{X}^{(k + 1)} = \arg \min E_{\mathrm{total}}$  by Eqn.(12)  
if converged and Flag  $==$  Stage-I then Flag  $=$  Stage-II;  
if converged and Flag  $==$  Stage-II then return  $\Pi_{ST}^{*};\Pi_{TS}^{*};\mathcal{V}^{*};$ $k = k + 1$
$$
d _ {h, l} (\mathbf {X}) = d _ {h, l} (\Theta , \Delta) = R \left(\Theta_ {h}\right) \left(g _ {l} - g _ {h}\right) + \Delta_ {k} + g _ {k} - \left(g _ {l} + \Delta_ {l}\right). \tag {11}
$$
Here,  $g \in R^{H \times 3}$  are the original positions of the nodes in the deformation graph  $\mathcal{DG}$ , and  $\psi(h)$  is the 1-ring neighborhood of the  $h$ -th deformation node.  $R(\cdot)$  is Rodrigues' rotation formula that computes a rotation matrix from an axis-angle representation and  $\alpha$  is the weight of the smooth rotation regularization term.
Total cost function: The total cost function  $E_{\mathrm{total}}$  combines the above terms with the weighting factors  $\lambda_{cd}$ ,  $\lambda_{corr}$ , and  $\lambda_{arap}$  to balance them:
$$
E _ {\text {t o t a l}} = \lambda_ {\mathrm {c d}} E _ {\mathrm {c d}} + \lambda_ {\text {c o r r}} E _ {\text {c o r r}} + \lambda_ {\text {a r a p}} E _ {\text {a r a p}} \tag {12}
$$
Now we are ready to describe our algorithm for minimizing  $E_{\mathrm{total}}$ , which is shown in Alg. 1.
Two-stage Registration: Finally, we observe that solely depending on learned features to infer correspondence is sub-optimal. At the converging point, the deformed source shape is often at the right pose but has deficiencies in shape style. To this end, we perform a second-stage registration, based on the coordinates of deforming source and target. As shown in Sec. 4.2, such a design is clearly beneficial.
Datasets: We evaluate our method and several state-of-the-art techniques for estimating correspondences between deformable shapes on an array of benchmarks as follows: FAUST_r: The remeshed version of FAUST dataset [4], which consists of 100 human shapes (10 individuals performing the same 10 actions). We split the first 80 as training shapes and the rest as testing shapes; SCAPE_r: The remeshed version of SCAPE dataset [2], which consists 71 human shapes (same individual in varying poses). We split the first 51 as training shapes and the rest as testing shapes; SHREC19_r: The remehsed version of SHREC19 dataset [36], which consists of 44 shapes of different identities and poses. We use it solely in test, and follow the test pair list provided by [36]; SHREC07-H: A subset of SHREC07 dataset [18], which consists of 20 heterogeneous human shapes of the varying number of vertices. We use it solely in test, and use the accompanied sparse landmark annotations to quantify all the pairwise maps among them; DT4D-H: A dataset proposed in [34], which consists of 10 categories of heterogeneous humanoid shapes. We use it solely in testing, and evaluating the inter-class maps split in [29]; TOPKIDS: This is a challenging dataset [25] consisting of 26 shapes of a kid in different poses, which manifest significant topological perturbations in meshes.
Baselines: We compare our method with an array of competitive baselines, including axiomatic shape registration methods: Smooth Shells [14], NDP [31], AMM [55]; learning-based registration
methods: 3D-CODED [19], Deep Shells [16], TransMatch [51], SyNoRiM [21]; deep functional maps frameworks: SURFMNet [45], WSupFMNet [46], GeomFMaps [13], DiffFMaps [35], NIE [24], ConsistFMaps [7], SSMSM [8]. According to input requirements, we put those relying on pure mesh input on the top, and the rest, suitable for point clouds, in the bottom of both tables.
Train/Test Cases: Throughout this section, we consider a highly challenging scenario: for each learning-based method, we train two models respectively with FAUST_r and SCAPE_r dataset, and then we run tests in a range of test cases including FAUST_r, SCAPE_r, SHREC19_r, SHREC07-H, DT4D-H and TOPKIDS. The test pairs of each case have been described above. In the following,  $A / B$  means we train on dataset  $A$  and test on dataset  $B$ .
Datasets Alignment: There exist extrinsically aligned versions for all the involved datasets but SHREC07-H. We equally feed in aligned datasets to all the baselines when available. On the other hand, the original version of SHREC07-H manifests significant variability of orientations. For the sake of fairness and simplicity, we apply our orientation regressor to it and provide baselines for our automatically aligned data. Note that, all the aligned datasets, as well as the synthetic dataset on which we train our orientation regressor, roughly share the same canonical orientation, which is defined by the SMPL [33] model. Finally, we always apply automatic alignment before inputting shapes into our pipeline, whatever the original orientation they are.
Choice of Source Shape: As mentioned at the beginning of Sec. 3, we compute correspondences between two raw input point clouds by associating them via a common source mesh  $S$ . In Tab. 1 and Tab. 2, we indicate our choice of source mesh by Ours-name, where name indicates the origin of the source mesh. For simplicity, we fix the source mesh from each dataset and visualize them in the appendix. On the other hand, when implementing axiomatic shape registration methods, we only consider deforming the same template we pick from FAUST_r (resp. SCAPE_r) to every test point cloud in the test set of FAUST_r (resp. SCAPE_r), and establish correspondence by map composition, in the same manner as ours.
Metric: Though we primarily focus on matching point clouds, we adopt the commonly used geodesic error normalized by the square root of the total area of the mesh, to evaluate all methods, either for mesh or point cloud.
Hyper-Parameters: We remark that all the hyper-parameters are fixed for all experiments in this section. In particular, we settle them by performing a grid search with respect to the weights used in the final optimization to seek for the combination that leads to the best registration results (quantitatively in terms of Chamfer distance and qualitatively by visual inspection) on a few training shapes. We provide more detailed discussion and ablation of the choice of hyper-parameters in the appendix.
Near-isometric Benchmarks: As shown in Fig. 1, in the presence of large deformation between the template and the target, NDP[31] and AMM[55] fail completely while our method delivers high-quality deformation. Moreover, as illustrated in Table 1, our method achieves the best performance in 5 out of 6 test cases. Remarkably, in the two most challenging tasks, FAUST_r/SHREC19_r and FAUST_r/SHREC19_r, our method indeed outperforms all of the baselines, including the state-of-the-art methods that take meshes as input. Regarding point-based methods, SSMSM [8] performs well in the standard case and outperforms ours in FASUT_r/FAUST_r, but generalizes poorly to unseen shapes. Another important observation is that our method manifests robustness with respect to the choice of template shapes. In fact, the above observations remain true no matter which template we select.
Non-isometric Benchmarks: We stress test our method on challenging non-isometric datasets including SHREC07-H and DT4D-H. We emphasize that these test shapes are unseen during training.
1) SHREC07-H contains 20 heterogeneous human shapes, whose number of vertices ranges from 3000 to 15000. Moreover, there exists some topological noise (e.g., the hands of the rightmost shape in Fig. 2 is attached to the thigh in mesh representation). As shown in Fig. 2, SSMSM [8] barely delivers reasonable results, which might be due to the sensitivity of graph Laplacian construction on point clouds. Topological noise, on the other hand, degrades mesh-based methods like Deep Shells [16]. Meanwhile, as shown in Tab. 2, our method achieves a performance improvement of
Table 1: Mean geodesic errors  $(\times 100)$  on FAUST_r, SCAPE_r and SHREC19_r. The best is highlighted.
Method	Train Test	FAUST_r	SCAPE_r
FAUST_r	SCAPE_r	SHREC19_r	SCAPE_r	FAUST_r	SHREC19_r
Smooth Shells [14]	2.5	\	\	4.7	\	\
SURFMNet(U) [45]	15.0	32.0	\	12.0	32.0	\
NeuroMorph(U) [15]	8.5	29.0	\	30.0	18.0	\
WSupFMNet(W) [46]	3.3	12.0	\	7.3	6.2	\
GeomFMaps(S) [13]	mesh	3.1	11.0	9.6	4.4	6.0	11.4
Deep Shells(W) [16]	1.7	5.4	26.6	2.5	2.7	21.4
ConsistFMaps(U) [7]	1.5	7.3	20.9	2.0	8.6	28.7
AttentiveFMaps(U) [29]	1.9	2.6	6.2	2.2	2.2	9.3
NDP [31]	20.4	\	\	16.2	\	\
AMM [55]	14.2	\	\	13.1	\	\
CorrNet3D(U) [59]	63.0	58.0	\	58.0	63.0	\
3D-CODED(S) [19]	2.5	31.0	\	31.0	33.0	\
SyNoRiM(S) [21]	7.9	21.9	\	9.5	24.6	\
TransMatch(S) [51]	pcd	2.7	33.6	21.0	18.6	18.3	38.8
DPC(S) [26]	11.1	17.5	31.0	17.3	11.2	28.7
DiffMaps(S) [35]	3.6	19.0	16.4	12.0	12.0	17.6
NIE(W) [24]	5.5	15.0	15.1	11.0	8.7	15.6
SSMSM(W) [8]	2.4	11.0	9.0	4.1	8.5	7.3
Ours-SCAPE	3.4	5.1	5.4	2.6	4.0	5.1
Ours-FAUST	3.0	6.3	5.9	2.9	4.0	4.8
Table 2: Mean geodesic errors  $\left( {\times {100}}\right)$  on SHREC’07 and DT4D-H. The best is highlighted.
Method	Train Test	FAUST_r	SCAPE_r
SHREC07-H(*)	DT4D-H	SHREC07-H(*)	DT4D-H
GeomFMaps [13]	30.5	38.5	28.9	28.6
Deep Shells [16]	30.6	35.9	31.3	25.8
ConsistFMaps [7]	mesh	36.2	33.5	37.3	38.6
AttentiveFMaps [29]	16.4	11.0	21.1	21.4
TransMatch [51]	25.3	26.7	31.2	25.3
DiffFMaps [35]	16.8	18.5	15.4	15.9
NIE [24]	15.3	13.3	13.4	12.1
SSMSM [8]	pcd	42.2	11.8	37.7	8.0
Ours-SCAPE	9.3	9.8	5.9	5.7
Ours-FAUST	8.5	\	6.1	\
Ours-CRYPTO	\	6.9	\	5.7
over  $40\%$  compared to the previous SOTA approaches (8.5 vs. 15.3; 5.9 vs. 13.4), which again confirms the robustness of our approach. 2) Regarding DT4D-H, We follow the test setting of AttentiveFMaps [29], and only consider the more challenging inter-class mapping. Interestingly, our method outperforms the state-of-the-art point cloud-based approach by approximately  $30\%$  and exceeds the performance of mesh-based state-of-the-art methods by over  $70\%$ . Furthermore, even in the worst case, training on FAUST_r and using SCAPE template, our error is still the lowest compared to external baselines.
Topologically Perturbed Benchmark: The shapes in TOPKIDS present various adhesions – hand/arm adhered to abdomen/thigh (see Fig. 6 in the appendix). Given the ground-truth maps from each topologically perturbed shape to the reference shape, we evaluate all 25 maps (excluding the trivial one from reference to itself) with models trained on FAUST_r and SCAPE_r respectively. Since the orientation of the TOPKIDS dataset is not originally agreeable with the aligned version of the two training sets, we align the input with our orientation regressor and feed them into the baselines depending on extrinsic information [16; 35; 24; 8]. Regarding the template shape, we adopt the reference shape as our source mesh in registration. We report the quantitative comparisons in Tab. 3. It is obvious that topological noise poses great challenges for methods based on pure intrinsic information [13; 7; 29], while the counterparts perform relatively well. Especially, among the point-based methods, our method outperforms the latest SOTA method by a large margin, namely, over  $40\%$  relative error reduction (7.1 vs. 12.3). We refer readers to the appendix for a qualitative comparison, which agrees with the quantitative results above.
Table 3: Mean geodesic errors  $(\times 100)$  on TOPKIDS which trained on FAUST_r and SCAPE_r. The best is highlighted.
GeomFMaps [13]	Deep Shells [16]	ConsistFMaps [7]	AttentiveFMaps [29]	DiffFMaps [35]	NIE [24]	SSMSM [8]	Ours
FAUST_nTOPKIDS	26.2	14.7	35.9	31.7	20.5	18.9	14.2	8.9
SCAPE_nTOPKIDS	21.7	15.3	33.1	39.4	18.0	16.2	12.3	7.1
Table 4: Mean geodesic errors  $(\times 100)$  on different ablated settings, the models are all trained on SCAPE r and test on SHREC'07.
w/o Registration	w/o Stage I	w/o Stage II	w/o updating corre.	w/o cons. filter	Full
11.5	10.6	10.1	8.1	7.2	5.9
Table 5: Mean geodesic errors  $(\times 100)$  of Ours, NDP, AMM based on the same initial maps.
SCAPE_r	SHREC19_r	SHREC07-H
Ini.	5.5	8.1	11.5
NDP	5.4	11.4	8.9
AMM	11.4	10.7	8.8
Ours	2.6	5.1	5.9
We report ablation studies in Tab. 4, in which we verify the effectiveness of each core design formulated in Sec. 3. Throughout this part, we train on SCAPE_r and test on SHREC07-H. In the registration stage, we use the SCAPE_r template as the source mesh and deform it to each point cloud from SHREC07-H. It is evident that each ablated module contributes to the final performance. In particular, both stages of registration play an equally important role in our framework.
Finally, to validate the superiority of our registration scheme, we report in Tab. 5 the average errors of the initial maps computed by our point-based DFM, and that of output maps of Ours, NDP [31], AMM [55], which are all based on the former. It is evident that, across three different test sets, our method consistently improves the initial maps, while NDP and AMM can even lead to deteriorated maps than the initial input.
In this paper, we propose a novel learning-based shape registration framework without correspondence supervision. Our framework incorporates learning-based shape matching and optimization-based shape registration. Based on the former, our framework can perform registration between shapes undergoing significant intrinsic deformations; On the other hand, thanks to the latter, our method exhibit superior generalizability over the learning-based competitors. Apart from several designs tailored for our intuitive pipeline, we also introduce a data-driven solution to facilitate the burden of extrinsically aligning non-rigid points. We verify our framework through a series of challenging tests, in which it shows superior performance but also remarkable robustness.
Limitation & Future Work We identify two main limitations of our method: 1) Our optimization procedure involves iteratively updating correspondences between the intermediate point clouds and the target one, leaving sufficient room for improvement for efficiency; 2) Our pipeline is designed for full shapes, while it is applicable to partial targets, the registration quality is sub-optimal. In the future, we also plan to incorporate more existing techniques (e.g., various regularization terms) from shape registration to further enhance the overall performance.
Acknowledgement The authors thank anonymous reviewers for their constructive comments and suggestions. This work was supported in part by the National Natural Science Foundation of China under contract No. 62171256, 62331006, in part by Shenzhen Key Laboratory of next-generation interactive media innovative technology (No. ZDSYS20210623092001004).
[1] Brian Amberg, Sami Romdhani, and Thomas Vetter. Optimal step nonrigid icp algorithms for surface registration. 2007.
In this appendix, we provide more technical details and experimental results. In Sec. A, we provide 1) A detailed description of our orientation regressor in Sec. A.1; 2) More details on the registration method presented in Alg. 1 of the main text in Sec. A.3; 3) Implementation details of the whole pipeline in Sec. A.4.
In the second part, we provide more experimental results. We first provide more qualitative results in Sec. B.1, as well as quantitative results on training on large dataset and on animal dataset in Sec. B.2; Then we demonstrate the robustness of our method with respect to several perturbations in Sec. B.3; In Sec. B.4, we showcase the scalability of our method in matching real scans. Finally, we report a run-time decomposition and analysis in Sec. B.5.
Note that we also include a short video showcasing the registration process of our pipeline in the Supplemental Material.
A Technical Details
A.1 Orientation Regressor
Orientation plays an important role in point cloud processing. Classic point-based feature extractors [42; 43; 53] are not rotation-invariant/equivariant by design and, therefore are limited. A common practice is to leverage data augmentation [42]. However, such can only robustify the model under small perturbations, while falling short of handling arbitrary rotations in  $SO(3)$ . On the other hand, recent advances in designing rotation-invariant/equivariant feature extractor [41; 10] seem to shed light on alleviating alignment requirements in non-rigid point matching. However, we empirically observe that such a design generally degrades the expressive power of the feature extractor, leading to inferior performance of the trained DFM.
This issue also attracts attention from the community of shape matching. In [11], the rotation prediction, instead of directly estimating the transformation, is formulated as a classification problem. However, this method can only address the rotation around a fixed axis and exhibits relatively low accuracy. Other works [56; 57] transform the input point cloud into handcrafted features. However, it should be noted that this approach often increases the computational complexity due to the additional feature extraction and processing steps involved.
In this paper, we propose a data-driven solution to this problem, accompanied by a registration module that is robust with the estimated alignment (see Tab.3 in the main text). As mentioned in the main text, we adopt the regressor model proposed in [9]. We use the synthetic SURREAL shapes [52], which are implicitly aligned by the corresponding generative codes to train our orientation regressor. We follow [9] and use PointNet++ [43] as the feature extractor. As we assume to be given full shape, either mesh or point cloud, for each input cloud, we translate it so its mass lies at the origin point. Then we estimate the orientation deviation,  $R$ , of it from the orientation implicitly defined by the SURREAL dataset. Finally, we apply the inverse rotation  $R^T$  on the input to obtain the aligned shape. We emphasize again that all the shapes, either train or test, go through this alignment in our pipeline.
A.2 Deep Functional Maps (DFM)
We illustrate the DFM, both for meshes and point clouds in Fig. 4.
A.3 Registration Algorithm
We provide a more detailed version of our registration algorithm in Alg. 2. given a source mesh  $S$ , target point cloud  $\mathcal{T}$  and a trained point feature extractor  $\mathbf{F}$ , we first initialize the deformation graph  $\mathcal{DG}$  with QSlim method [17] and set  $\mathbf{X}^{(0)}$  with identity transformations. After that, we compute the target point feature  $F_{\mathcal{T}}$ . During registration, we first deform the source model with a deformation graph, then we update the correspondence by computing the nearest neighbor between the source feature and target feature in stage-I or the source point cloud and target point cloud in stage-II. After we compute the set of filtered correspondences  $\mathcal{C}^k$ , we update the total cost function  $E_{\mathrm{total}}$  with the ADMM optimizer. Finally, we follow [31] to implement an early stop strategy in both stages.
(b)
Figure 4: (a) DFM architecture for pre-training on meshes; (b) Our proposed architecture for training a point-based DFM, which is self-supervised by its mesh-based counterpart.
Algorithm 2: Shape registration pipeline.
Input: Source mesh  $\mathcal{S} = \{\mathcal{V},\mathcal{E}\}$  and target point cloud  $\mathcal{T}$ ; Trained point feature extractor  $\mathbf{F}$   
Output:  $\mathbf{X}^*$  converging to a local minimum of  $E_{\mathrm{total}}$ ; Deformed source model  $\{\mathcal{V}^*,\mathcal{E}\}$ ; Correspondence  $\Pi_{\mathcal{S}\mathcal{T}}^{*},\Pi_{\mathcal{T}\mathcal{S}}^{*}$  between  $\mathcal{S}$  and  $\mathcal{T}$ .
1 Initialize deformation graph  $\mathcal{DG}$  and  $\mathbf{X}^{(0)}$  with identity transformations;  $F_{\mathcal{T}} = \mathbf{F}(\mathcal{T})$ ;  $\mathrm{k} = 0$ ;
3 while True do
2  $eps = 10^{-8}$ ; count = 0;
$\mathcal{V}^k = \mathcal{D}\mathcal{G}(\mathbf{X}^{k - 1},\mathcal{V}^{k - 1})$    
if  $k \% 100 = = 0$  and Flag  $= =$  Stage-I then  
 $\begin{array}{r}\big{\lfloor}F_{S}^{(k)} = \mathbf{F}(\mathcal{V}^{(k)});\Pi_{ST}^{(k)} = \mathbf{NN}(F_{S}^{(k)},F_{T});\Pi_{TS}^{(k)} = \mathbf{NN}(F_{T},F_{S}^{(k)}) \end{array}$    
if  $k \% 100 = = 0$  and Flag  $= =$  Stage-II then  
 $\begin{array}{rl} & {\big{\lfloor}\Pi_{ST}^{(k)} = \mathbf{NN}(\mathcal{V}^{(k)},\mathcal{T});\Pi_{TS}^{(k)} = \mathbf{NN}(\mathcal{T},\mathcal{V}^{(k)});}\\ & {\mathrm{Compute~the~set~of~filtered~correspondences~}\mathcal{C}^k;}\\ & {E_{\mathrm{total}}^{(k)} = \lambda_{\mathrm{cd}}E_{\mathrm{cd}}(\mathcal{V}^{(k)},\mathcal{T}) + \lambda_{\mathrm{corr}}E_{\mathrm{corr}}(\mathcal{V}_{\mathcal{C}^k}^{(k)},\mathcal{T}_{\mathcal{C}^k}) + \lambda_{\mathrm{arap}}E_{\mathrm{arap}}(\mathcal{D}\mathcal{G});}\\ & {\mathbf{X}^{(k)} = \arg \min E_{\mathrm{total}}^{(k)};}\\ & {\mathrm{if~}(E_{\mathrm{total}}^{(k)} - E_{\mathrm{total}}^{(k - 1)}) <   \mathrm{eps~then}}\\ & {\mathrm{count = count + 1};}\\ & {\mathrm{if~count > 15~then}}\\ & {\mathrm{if~Flag == Stage - I~then~Flag =~Stage - II;~count = 0;~eps = 10^{-7};}}\\ & {\mathrm{if~Flag == Stage - II~then~return~}\Pi_{ST}^*;\Pi_{TS}^*;\mathcal{V}^*;}\\ & {\mathrm{k = k + 1};} \end{array}$
Figure 5: More non-rigid shape registration results. The original source mesh and its deformed counterpart are represented in blue, while the target point cloud and the ground truth mesh are depicted in orange.
We implement our framework in PyTorch. We use modified DGCNN [24] the backbone of our feature extractor. The output feature dimension is set to 128. In the training related to deep functional maps, we always use the first 50 eigenfunctions of Laplace-Beltrami operators. For training the DFM network, in Eqn.(6) of the main text, we empirically set  $\lambda_{\mathrm{bij}} = 1.0$ ,  $\lambda_{\mathrm{orth}} = 1.0$ ,  $\lambda_{\mathrm{align}} = 1\mathrm{e - }4$ ,  $\lambda_{\mathrm{NCE}} = 1.0$ . We train our feature extractor with the Adam optimizer with a learning rate equal to 2e-3. The batch size is chosen to be 4 for all datasets.