SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Benchmark results on CHD and ACDC datasets under both random initialized weights, and pretrained weights from SSL.  $M$  is the number of patients used in supervised training. We perform 5-fold cross-validation and the mean (standard deviation) dice scores are reported.
Init.	Method	CHD (68 patients in total)
M=2	M=6	M=10	M=15	M=20	M=30	M=51
Random	UNet [52]	0.184(.06)	0.508(.06)	0.584(.05)	0.627(.05)	0.658(.04)	0.693(.04)	0.754(.02)
Swin-Unet [5]	0.291(.07)	0.543(.07)	0.624(.04)	0.675(.05)	0.717(.04)	0.732(.05)	0.784(.03)
SwinUNETR [58]	0.345(.07)	0.565(.06)	0.638(.05)	0.682(.06)	0.711(.05)	0.725(.06)	0.785(.03)
Ours	0.344(.05)	0.576(.07)	0.646(.03)	0.686(.03)	0.706(.03)	0.728(.04)	0.778(.03)
SSL pretrain	Rotation [18]	0.171(.06)	0.488(.07)	0.575(.04)	0.625(.04)	0.651(.04)	0.691(.04)	0.749(.03)
PIRL [45]	0.196(.07)	0.504(.08)	0.617(.05)	0.658(.03)	0.674(.04)	0.714(.04)	0.761(.03)
SimCLR [9]	0.192(.06)	0.515(.06)	0.599(.06)	0.631(.05)	0.666(.05)	0.699(.05)	0.756(.03)
GLCL-global [7]	0.255(.10)	0.564(.04)	0.646(.03)	0.669(.04)	0.697(.04)	0.725(.04)	0.766(.03)
GLCL-full [7]	0.286(.06)	0.555(.07)	0.614(.06)	0.666(.04)	0.694(.04)	0.733(.04)	0.772(.03)
CAiD [56]	0.265(.08)	0.581(.06)	0.647(.04)	0.684(.04)	0.700(.04)	0.737(.04)	0.771(.02)
PCL [72]	0.356(.08)	0.600(.06)	0.661(.05)	0.686(.05)	0.716(.04)	0.735(.05)	0.774(.03)
Ours	0.392(.06)	0.636(.06)	0.693(.03)	0.712(.03)	0.728(.04)	0.754(.04)	0.788(.03)
Init.	ACDC (100 patients in total)
Method	M=2	M=6	M=10	M=15	M=20	M=30	M=80
Random	UNet [52]	0.588(.07)	0.782(.03)	0.840(.03)	0.876(.01)	0.894(.01)	0.909(.01)	0.928(.00)
Swin-Unet [5]	0.181(.07)	0.483(.09)	0.610(.05)	0.700(.04)	0.735(.02)	0.772(.01)	0.870(.01)
SwinUNETR [58]	0.567(.04)	0.740(.07)	0.799(.05)	0.850(.02)	0.881(.01)	0.904(.01)	0.922(.00)
Ours	0.655(.05)	0.827(.05)	0.871(.02)	0.897(.01)	0.901(.01)	0.915(.00)	0.927(.00)
SSL pretrain	Rotation [18]	0.572(.08)	0.809(.03)	0.868(.02)	0.886(.01)	0.898(.01)	0.910(.01)	0.925(.00)
PIRL [45]	0.492(.03)	0.823(.04)	0.865(.01)	0.880(.02)	0.896(.02)	0.912(.01)	0.927(.00)
SimCLR [9]	0.352(.06)	0.725(.08)	0.824(.04)	0.869(.02)	0.894(.01)	0.913(.01)	0.927(.00)
GLCL-global [7]	0.636(.05)	0.803(.04)	0.872(.01)	0.891(.01)	0.902(.01)	0.913(.01)	0.927(.01)
GLCL-full [7]	0.642(.06)	0.802(.03)	0.877(.01)	0.891(.01)	0.904(.02)	0.912(.00)	0.927(.00)
CAiD [56]	0.483(.11)	0.822(.02)	0.879(.02)	0.896(.01)	0.905(.00)	0.914(.00)	0.926(.00)
PCL [72]	0.671(.06)	0.850(.01)	0.885(.01)	0.904(.01)	0.909(.01)	0.919(.00)	0.929(.00)
Ours	0.741(.03)	0.873(.01)	0.895(.01)	0.908(.01)	0.915(.00)	0.921(.00)	0.930(.00)
ACDC (Automatic Cardiac Diagnosis Challenge) [2] is a MRI dataset consists of cardiac images from 100 patients under several different pathologies. Manual expert segmentation of the RV, LV cavities, and Myo are conducted for the volumes from end-diastolic and end-systolic phase. For each cross-validation fold,  $N = 20$  images are used as the validation set, and  $M$  images from the remaining 80 images are used for supervised finetuning.
Baselines and Evaluation Settings. We evaluate segmentation performance trained with limited annotation under both (1) random initialized network backbones (UNet v.s. Ours) to verify the architectural advantages of the proposed keypoint-augmented fusion layer. (2) pretrained weights from our self-supervision compared with various self-supervised pretraining methods including pre-text tasks [18, 45] and contrastive learning [9, 7]. We quantify the segmentation performance via the Dice coefficient with five-fold cross validation following the set up in [72]. Additionally, we conducted ablation studies to isolate each component in the pipeline to assess individual effects.
Implementation Details. For network configuration, we build our framework on 2D UNet, and set the starting number of channels of the network as 32 for CHD, and 48 for ACDC. Five convolutional blocks are included in the encoder and four blocks are used within the decoder. Each block contains two convolution layers, followed by batch normalization, and ReLU activation. Each encoder block downsample the feature map by a factor of 2, while simultaneously doubling the number of channels. In between the first four encoder blocks, we further append the proposed keypoint-augmented fusion layer, and concatenate the output to the original convolutional feature map (Fig. 1). Each Transformer inside KAF layer consists of six self-attention layers, and the comparison among different number of self-attention layers are presented in ablation study below. For pretraining, we assign loss weights  $w_{1}, w_{2}, w_{3}$  to 1.0, 1.0, and 0.01, respectively. We employ the SGD optimizer with a learning rate of 0.002 and batch sizes of 32 for CHD and ACDC. A cosine learning rate scheduler is utilized, with the minimum learning rate set to 0. We pretrain the model for 50 epochs. For finetuning, we use the standard cross-entropy loss with the Adam [33] optimizer, with learning rates of  $\overline{5 \times 10^{-5}}$  for CHD,  $5 \times 10^{-4}$  for ACDC. The batch size is set to 10, and we finetune on CHD for 100 epochs and on ACDC for 200 epochs. Additional implementation details are provided in the appendix.
Figure 2: Few-shot segmentation results. Colonn (a)-(c) show a comparison among the UNet backbone, SwinUNETR [58], and our proposed backbone trained on limited dataset with random weight initialization. Our architecture allows for reduced false positive and improved segmentation accuracy. Column (c)-(e) compares our pretraining strategy with the existing state-of-the-art SSL methods [7] and [72]. When finetuned with small number of labeled dataset, our method presents higher coherence with the reference image (column (f)). The number beneath each prediction represents the dice value for the displayed slice. For better readability, we scale the value by 100.
Segmentation Results without Pretraining. To verify the benefits of the proposed layer, we start with analyzing the performance gain without any pretraining, to isolate and evaluate the impact solely attributed to the architectural change on top of the UNet backbone. Specifically, we train the randomly initialized UNet and our backbone with varying limited number of training data, and the results are shown in Table 1 under Init.Random. Compared with the UNet backbone with an identical number of convolutional blocks and network configurations, introducing KAF layer consistently and significantly improves the segmentation results on both datasets across various training sample sizes. The empirical findings strongly support our assumption that incorporating long-range dependencies into the segmentation backbone is crucial and helps better modeling and utilizing regional image information when only limited annotations are available. Notably, our network demonstrates superior performance over the UNet baselines when the training size is extremely small. For instance, when the sample size is reduced to two ( $M = 2$ ) on the CHD dataset, the inclusion of KAF layer leads to  $87\%$  performance gain in dice score over UNet. We also include Transformer-based UNets [5, 58], which also consider long-range spatial dependencies. Our observation is consistent with the conclusion in [5], where Transformers may be more severely affected by the initialization and require large  $M$  in order to perform well, while method obtain much better results when only limited data is available.
Fig. 2 Random Initialization shows a visual comparison of the segmentation between UNet and Ours. While both methods suffer from sub-optimal generalization ability on the input image from the unseen validation set, our method largely reduces false positive prediction and presents higher similarity with the reference image on the right. Such improvement can be attributed to the utilization of keypoints, which enhances the feature learning process by directing attention to a more constrained set of regions rather than the entire image space. We deduce that this focused awareness contributes to better localization of important features thus leading to better semantic segmentation.
Segmentation Results with Pretraining and Finetuning. To further enhance the segmentation performance and promote better generalization of our model, we pretrain the network with the proposed SS objectives and leverage the pretrained weights as the initialization for fine-tuning on the same amount of labeled dataset. Quantitative results are shown in Table 1 under the tab SSL pretrain initialization, where benchmarking is conducted across existing state-of-the-art SSL pretraining methods. Notably, our pretraining technique yields significant improvements in few-shot
Table 2: Ablation study of our framework over (1) architecture design: the number of self-attention layers within the transformer (#T), scales to insert KAF layer  $(l_{1},\dots ,l_{4})$ ; and (2) pretraining hyperparameters  $(w_{1},w_{2},w_{3})$ . Five-fold cross-validation results on both datasets are reported.
Init.	Exp	Architecture design	Pretraining	Dice
#T	l1	l2	l3	l4	w1	w2	w3	CHD (M=15)	ACDC (M=6)
Random	A	9	0	0	0	0	-	-	-	0.627(.05)	0.782(.03)
B	9	1	1	1	0	-	-	-	0.658(.04)	0.806(.04)
C	9	1	1	0	1	-	-	-	0.677(.04)	0.817(.04)
D	9	1	0	1	1	-	-	-	0.667(.04)	0.814(.04)
E	9	0	1	1	1	-	-	-	0.666(.04)	0.811(.04)
F	9	1	1	1	1	-	-	-	0.686(.03)	0.827(.05)
G	6	1	1	1	1	-	-	-	0.690(.03)	0.824(.03)
H	3	1	1	1	1	-	-	-	0.677(.03)	0.814(.03)
SSL pretrain	I	9	1	1	1	1	1	0	0	0.699(.04)	0.867(.02)
J	9	1	1	1	1	0	1	0	0.701(.03)	0.865(.02)
K	9	1	1	1	1	1	1	0	0.711(.03)	0.865(.01)
L	9	1	1	1	1	1	1	0.01	0.712(.03)	0.873(.01)
segmentation tasks compared to the random initialization. Moreover, in comparison to alternative pretraining strategies that do not incorporate long-range dependencies, our method consistently achieves state-of-the-art dice scores across different numbers of training subjects for both datasets, demonstrating the benefits of the proposed keypoint-augmented SSL objectives. In Fig. 2 SSL initialization, we present visual comparisons between our method and the best-performing SSL methods GLCL [7] and PCL [72]. Visually, we observe that GLCL tends to produce false positives or incorrectly labels the anatomy, as evidenced by the CHD result. Conversely, PCL tends to generate false negatives, as observed in the ACDC example. In contrast, our proposed method achieves the highest level of similarity with the reference image, accurately capturing the desired anatomical structures. These visual comparisons further reinforce the effectiveness and superior performance of our approach in comparison to existing SSL methods.
Ablation Study. As our proposed framework consists of both architectural change on the UNet building blocks and pretraining strategy, we conduct ablation analysis over different architectural configurations and pretraining objectives, with results presented in Table 2. Additional ablations on keypoint detection method and threshold for identifying correspondence are provided in the appendix.
Architecture (Exp A-H). We start with architectural analysis by training different randomly initialized backbone networks under different configurations over the number of self-attention (#T) within the KAF layer, and the injection of KAF layer at different encoder scales  $l_{1}, \dots, l_{4}$ .
Exp A represents a UNet backbone without the use of KAF layer. With the same #T, Exp F inserts four KAF layer after each of the encoder layer. A significant performance gain was obtained on CHD by 0.06, and on ACDC by 0.04. To analyze layerwise effects, we remove one KAF layer each time from Exp B-E, and observe a degradation in performance, indicating the importance of utilizing a multiscale setting. We observe that the last layer plays the most crucial role in achieving superior segmentation results. In addition, we investigate the effect of the number of self-attention within the Transformer used in our architecture in Exp F-H. We observe that the performance reaches a plateau with  $\# T = 9$ , and when it reduces to 3, the performance further degrades. Nevertheless, Exp H still overperforms Exp A, indicating the benefits of introducing long-range self-attention.
Pretraining Strategy (Exp I-L). On the architectures of Exp F, we further add the proposed pretraining strategies, and isolate the different loss weights during the training.  $w_{1}, w_{2}$  indicate the global PCL loss applied on the convolutional feature and keypoint features respectively,  $w_{3}$  indicates the local loss weight (See Eq. 3). Compared with single global loss on either the global CNN feature (Exp I) or the keypoint feature (Exp J), a combination of the two terms (Exp K) shows improved results on CHD. The best performance is achieved with a combination of both global and local loss (Exp L). Additional weight tuning details are provided in appendix.
Qualitative Analysis of the Learned Self-Attention. In Fig. 3, we qualitatively compare the learned self-similarity from our pretraining, with the established PCL [72] (global), and GLCL [7] (global and local) pretraining. Given an input image  $x$ , we perform two random transformations  $T_{1}$  and  $T_{2}$  to obtain a simulated positive pairs in contrastive set up. A query point is randomly sampled from
Figure 3: Learned multiscale self-similarity between the query point feature (star) and all other feature points within the same feature map from the UNet encoder. The comparison between (a) and (b), (d) and (e) indicates that our method maintains better invariance of the local feature self-similarity under different transformations. Additionally, the coherence between (b) and (c), (e) and (f) verify that our method learns better equivariance. Details are elaborated in Sec. 4.
the  $x$  and gets transformed to  $T_{1}(x)$  and  $T_{2}(x)$  respectively (stars in Fig. 3). The image is fed into the pretrained network, and the feature map from the  $i$ -th layer is obtained. The self-similarity  $S$  is computed as the dot product between the query point feature and other feature points within the same map. Fig. 3 visualizes the self-attention from the third and fourth layer of the pretrained UNet  $(l_{3}, l_{4})$ .
We observe two major differences between our method and PCL and GLCL: (1) our keypoint-augmented fusion layer helps the network to concentrate its attention on more constrained sub-regions. For example, both PCL and GLCL tend to learn high response values between the query point and background regions (e.g., Fig. 3 (b), (e)), whereas our method exhibits a more refined attention mechanism, directing the focus of the query point towards anatomically relevant regions, which helps the model to capture and leverage important features within the desired regions, leading to improved segmentation performance. (2) our proposed self-supervision helps maintain better local equivariance of the self-attention, i.e., with the same query point location, its interaction with other points remains identical no matter how the image is transformed. This is verified by a comparison between Fig. 3 (a) v.s. (b), and (d) v.s. (e), where our method achieves the best consistency of the local similarity when features are extracted from the image with different transforms.
Additionally, we calculate the self-attention based on the input image  $x$  and transform the resulting similarity map using  $T_{2}$ . The transformed similarity map is expected to exhibit a high level of coherence with the similarity map obtained from the transformed input. Compared with PCL and GLCL, our method demonstrates stronger consistency between (b) and (c), as well as (e) and (f), which validates the advantage of introducing the keypoint-augmented self-supervised learning.
In this work, we present a keypoint-augmented fusion layer to incorporate long-range dependencies into the UNet-based segmentation framework, accompanied by global & local SSL objectives for pretraining. Open questions exist and will be included in future work. Currently, our backbone is built upon 2D UNet following [72, 7], while 3D UNet naturally serves as a better baseline in many biomedical segmentation tasks [51]. Although our method is generic and can be scaled up to 3D data, it may require different architectural configurations to accommodate the increasing number of keypoints within the 3D volume. Besides, our assumption for identifying correspondence may fail in 3D volume with the sparse acquisition, when the neighboring slices no longer maintain semantically
similar structures. Therefore, more robust and complicated criteria for obtaining the correspondence is required.
Co-authors Zhangsihao Yang and Yalin Wang are grateful to R21AG065942, R01EY032125, and R01DE030286.
More comparison with Transformers As we pointed out in the main text, Transformer-based architecture may be extremely sensitive to weight initialization, and may benefit from effective pretraining. Therefore, we further pretrain the SwinUNETR with two SSL strategies: PCL [72] and the self-supervised loss proposed in [22] (in Tabel 3). Both of these pretraining methods led to observable improvements over random initialization, while our proposed method still achieved the best performance.
Table 3: Comparison with Transformer-based methods on CHD and ACDC datasets under both random initialized weights, and pretrained weights from SSL.  $M$  is the number of patients used in supervised training. We perform 5-fold cross-validation and the mean (standard deviation) dice scores are reported.
Init.	Method	CHD (68 patients in total)
M=2	M=6	M=10	M=15	M=20	M=30	M=51
Random	SwinUNETR-3d [58]	0.040(.01)	0.189(.03)	0.332(.08)	0.420(.06)	0.472(.05)	0.504(.05)	0.577(.06)
Ours	0.344(.05)	0.576(.07)	0.646(.03)	0.686(.03)	0.706(.03)	0.728(.04)	0.778(.03)
SSL	PCL (Swin-Unet) [5]	0.321(.05)	0.556(.06)	0.658(.03)	0.692(.04)	0.722(.04)	0.747(.04)	0.783(.03)
PCL (SwinUNETR) [58]	0.343(.06)	0.576(.07)	0.644(.05)	0.692(.06)	0.714(.05)	0.730(.05)	0.786(.03)
Ours	0.392(.06)	0.636(.06)	0.693(.03)	0.712(.03)	0.728(.04)	0.754(.04)	0.788(.03)
Init.	Method	ACDC (100 patients in total)
M=2	M=6	M=10	M=15	M=20	M=30	M=80
Random	SwinUNETR-3d [58]	0.391(.07)	0.466(.05)	0.545(.05)	0.603(.03)	0.629(.03)	0.679(.02)	0.735(.01)
Ours	0.655(.05)	0.827(.05)	0.871(.02)	0.897(.01)	0.901(.01)	0.915(.00)	0.927(.00)
SSL	PCL (Swin-Unet) [5]	0.292(.09)	0.577(.07)	0.702(.04)	0.742(.03)	0.798(.02)	0.825(.02)	0.878(.01)
PCL (SwinUNETR) [58]	0.547(.05)	0.759(.05)	0.808(.05)	0.850(.03)	0.882(.01)	0.904(.01)	0.923(.00)
Ours	0.741(.03)	0.873(.01)	0.895(.01)	0.908(.01)	0.915(.00)	0.921(.00)	0.930(.00)
Ablation on keypoint detection method While we use SIFT as keypoint detector, it can be replaced with other learning based SoTA methods as well. We used a pretrained keypoint detection model Superpoint [13] as an alternative to SIFT. Segmentation results on CHD and ACDC dataset are reported below comparing two methods, and the results indicate that our method is not sensitive to different keypoint detection algorithms. Overall, Superpoint leads to slightly lower results, and we speculate this was due to the pretraining was done on natural images (COCO dataset).
Table 4: Segmentation results (dice scores) with different keypoint detection methods  $\left( {M = 2}\right)$  .
Init.	Keypoint detector	CHD	ACDC
Random Init.	SuperPoint [13]	0.643(.04)	0.810(.02)
SIFT	0.686(.03)	0.827(.05)
SSL Pretrain	SuperPoint [13]	0.703(.03)	0.865(.02)
SIFT	0.712(.03)	0.873(.01)
Sensitivity of the keypoint correspondence To verify how sensitive the method is to the number of matching keypoints, we perform an ablation on different threshold values, to pretrain our model on the CHD dataset , and finetune it on  $M = 15$  labeled data. Global losses  $(w_{1} = w_{2} = 0)$  are turned off to isolate the effects of the threshold for the local SSL loss. The dice scores are reported in Table 5, indicating that our model is not sensitive to the threshold setting and remains robust across different values. We also visualize the correspondence and the number of matching keypoints under different threshold values in Fig 4.
Table 5: Dice scores under different keypoint correspondence threshold values. Overall, we observe that the results are not sensitive to specific threshold values.
Threshold	Dice
5	0.689 (0.043)
10	0.690 (0.031)
15	0.684 (0.035)
20	0.689 (0.036)
25	0.684 (0.032)
30	0.685 (0.033)
35	0.690 (0.037)
40	0.687 (0.037)
Figure 4: Sensitivity analysis of the Manhattan distance, illustrating the correspondence found under the specified Manhattan distance threshold.
Additional dataset We include experiments on a non-cardiac multi-organ CT segmentation dataset to evaluate the generalization of our method. Results are reported below. With 2 training subjects, we tested the performance of the model under both random and SSL pretrained initialization. In both scenarios, our method outperformed the existing works.
Table 6: Segmentation results on Synapse dataset, where finetuning is done on  $M = 2$  subjects out of a total of 18 subjects.
Init.	Method	Dice (M=2)
Random	Unet [52]	0.253(.06)
Swin-Unet [5]	0.198(.04)
SwinUNETR [58]	0.279(.06)
Ours	0.289(.06)
SSL	PCL [72]	0.306(.05)
Swin-Unet (with [72])	0.210(.07)
SwinUNETR (with [72])	0.304(.06)
Ours	0.322(.06)
Ablation on UNet Channels (Double Channels). In our implementation, we augment the convolutional UNet features by concatenating them with features learned from KAF. Introducing KAF layers leads to an increase in the total number of parameters of the final KAF-enhanced UNet. Specifically, the number of channels for the second, third, and fourth blocks in the KAF-enhanced UNet become twice as large as the original UNet. For a more fair comparison, we construct a non-KAF UNet ('UNet(c2)') in Tab. 7 with the same amount of parameters in convolutional layers as our model by duplicating the features of the first, second, third, and fourth blocks from the UNet baseline ('UNet(c1)') in Tab. 7 and concatenated each with itself. The segmentation results are presented in Table 7. UNet(c1) indicates the UNet backbone, and UNet(c2) denotes the larger UNet whose convolutional parameters matched our model. They indicate that widening the network's architecture by increasing the input channel size can improve its performance. However, the performance enhancement is even more substantial with our modified layer. This suggests that incorporating features beyond simple convolution into the network architecture can further enhance the network's performance.
Table 7: Performance comparison among standard UNet (UNet (c1)), a larger UNet with duplicated input channels (UNet (c2)), and UNet augmented with features derived from KAF (Ours). The results indicate that using a larger UNet slightly improves the segmentation performance. Our proposed KAF-enhanced UNet further boosts the performance significantly compared with UNet (c2).
Sample M	dataset	Method	mean/std	#params
15	CHD	UNet(c1)	0.627(.05)	7.8 M
15	CHD	UNet(c2)	0.646(.04)	27.9 M
15	CHD	Ours	0.712(.03)	71.7 M
6	ACDC	UNet(c1)	0.782(.03)	17.5 M
6	ACDC	UNet(c2)	0.796(.03)	62.8 M
6	ACDC	Ours	0.873(.01)	106.8 M
Ablation on Correspondence Weights. To further investigate the contribution of the correspondence loss to the performance of the pretraining weights, we conducted a study on the various combinations of weights  $w_{1}$ ,  $w_{2}$ , and  $w_{3}$ , as defined in eq. 3, and results are reported in Table 8.