SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhangsihao Yang*
zshyang1106@gmail.com
Mengwei Ren*
mengwei. ren@nyu.edu
Kaize Ding
kaize.ding@northwestern.edu
Guido Gerig
gerig@nyu.edu
Yalin Wang
ylwang@asu.edu
Pretraining CNN models (i.e., UNet) through self-supervision has become a powerful approach to facilitate medical image segmentation under low annotation regimes. Recent contrastive learning methods encourage similar global representations when the same image undergoes different transformations, or enforce invariance across different image/patch features that are intrinsically correlated. However, CNN-extracted global and local features are limited in capturing long-range spatial dependencies that are essential in biological anatomy. To this end, we present a keypoint-augmented fusion layer that extracts representations preserving both short- and long-range self-attention. In particular, we augment the CNN feature map at multiple scales by incorporating an additional input that learns long-range spatial self-attention among localized keypoint features. Further, we introduce both global and local self-supervised pretraining for the framework. At the global scale, we obtain global representations from both the bottleneck of the UNet, and by aggregating multiscale keypoint features. These global features are subsequently regularized through image-level contrastive objectives. At the local scale, we define a distance-based criterion to first establish correspondences among keypoints and encourage similarity between their features. Through extensive experiments on both MRI and CT segmentation tasks, we demonstrate the architectural advantages of our proposed method in comparison to both CNN and Transformer-based UNets, when all architectures are trained with randomly initialized weights. With our proposed pretraining strategy, our method further outperforms existing SSL methods by producing more robust self-attention and achieving state-of-the-art segmentation results. The code is available at https://github.com/zshyang/kaf.git.
Large-scale and diverse source of training data has significantly empowered the generalization ability of supervised segmentation models [52, 77, 34]. However, in biomedical image analysis, manual delineation of images/volumes requires extensive domain knowledge and can be extremely time-consuming and error-prone. Recently, self-supervised learning (SSL) [9, 25, 18, 45, 7, 72, 51] has emerged in medical vision to pretrain a segmentation backbone (e.g. UNet) leveraging large scale unlabelled data, to serve as a better initialization before finetuning the model with limited annotation.
As one of the most successful learning paradigms in SSL, contrastive learning pulls closer the representations of positive pairs (i.e. the same image under different augmentations), while pushing apart the negative pairs (i.e. different images) [9, 28, 10, 25, 46, 11, 12]. However, in medical images, strong anatomical similarities often exist across different images, resulting in an increased number of false negative samples which is detrimental to the representation learning process. Therefore, heuristics about relationship across different images are utilized to constitute positive pairs, such as encouraging the global representations of 2D slices at similar volumetric positions [7, 72] to be close. To further benefit pixel-level tasks such as segmentation, local representation learning methods [7, 51] learn distinctive local representations by attracting patch features at the same position from two augmented views [7], or across registered intra-subject volumes [51]. Nevertheless, the potential dependencies among patch features that are spatially distant are not properly taken into account. Such region-awareness can be enhanced by attracting spatial features points with the same semantics [35, 27, 38, 75, 1]. However, their training requires segmentation ground truth, which limits their applicability in medical vision when adequate annotation is unavailable.
In this paper, we incorporate the long-range spatial dependencies into the UNet backbone with a plug-and-play keypoint-augmented fusion layer (KAF layer), accompanied by keypoint-enhanced global and local objectives for pretraining the network with self-supervision. To do so, we first identify a set of sparse keypoint locations from the input image. On the output feature map of each UNet encoder block, the keypoint features are sampled from the convolutional feature map, and the attention among them are modeled through a Vision Transformer. The output is then scattered back to the feature grid, and gets fused to the original CNN feature, to explicitly provide clues on the spatial long-range dependencies. Further, we propose the global and local self-supervised learning objectives to pretrain the keypoint-augmented fusion layer -enhanced network. The global contrastive loss is applied on both the UNet bottleneck feature, as well as keypoint-augmented global feature obtained by aggregating the multiscale keypoint features. To further benefit pixel-level down-stream tasks, we identify the correspondence among keypoints across different image slices and maximize their feature similarities with a local self-supervised objective.
Our contributions are threefold. (1) We develop a plug-and-play keypoint-augmented fusion layer that augments the convolutional feature map with long-range dependencies among keypoint features. When trained with only limited annotation, our proposed KAF layer achieves noticeably better results than the CNN-only and/or Transformer-based backbone. (2) We further propose the keypoint-aware global and local self-supervised learning objectives to pretrain our model before finetuning it with annotations. (3) We conduct extensive experiments on three MRI and CT datasets, and achieve state-of-the-art few-shot segmentation results, verifying both architecture advantage of the proposed layer and our pretraining strategies.
Biomedical Image Segmentation is the process of dividing images into various segments or regions and isolating certain structures of interest in modalities such as magnetic resonance imaging (MRI) and computed tomography (CT) [50, 37]. UNet [52] is the pioneering method for addressing biomedical image segmentation tasks using deep learning. Subsequent research aims to enhance or apply the UNet model in various ways, such as extending to 3D volumes [44, 64], auto-configuration [29], altering skip connections [77, 66], and replacing and/or combining the convolutional backbone layers with Transformers [23, 5, 22, 30]. However, a common limitation shared by all end-to-end methods lies in the requirement of large-scale and high-quality biomedical annotations which are both time-consuming and requires expert knowledge [52]. This limitation underscores the need for novel methods that can learn effectively with fewer or without annotations.
Self-Supervised Learning (SSL) learns meaningful representations from unlabeled datasets with self-supervised objectives, and the pretrained network constitutes a better initialization than random initialization when it is finetuned on downstream tasks with limited label supervision, e.g., classification or segmentation. Existing SSL algorithms either define pretext tasks with heuristics, or employ self-supervised objectives to learn invariance and/or equivariance of image features. Early pre-text tasks based pretraining designs objective that does not require additional labels, such as rotation prediction [18], context restoration [48], reconstruction [32], among others [24, 42, 15]. Recently, contrastive learning has become a prevailing paradigm in SSL and has shown great success in both natural image [6, 59, 78, 63, 76] and medical vision [7, 72, 51, 58, 71, 70, 56, 49, 8, 62, 16, 47, 73, 68].
It enforces similar representations between positive pairs and distinct representations between negative pairs. Typically, similarity is determined in an unsupervised way, i.e., the same image under different transformations are recognized as positive pairs, whose global representations should remain invariant, and features from different images are pulled apart. In medical images, heuristics about relationships among different images are also utilized to define positives, e.g., 2D slices at similar volumetric positions [7, 72], or local patches at the same location from intra-subject volumes [51].
Long-Range Dependency generally cannot be well captured in CNN due to its locality nature, thus neglecting the potential relationship across spatially distant local regions that possess high semantic consistency. To tackle the limitation, existing works focus on either (1) architectural design that injects long-range attention, or (2) SSL objectives that explicitly enforce regional similarity. For the first line of research, a series of non-local networks have been proposed. For instance, Vision Transformer (ViT) [14] enables self-attention on uniformly sampled image patches and has been incorporated into CNN backbones for segmentation tasks [22, 5, 36, 26]; [54, 43, 20, 65] adopt graph networks to model dependencies among context, keypoints or boundaries. The other line of research captures regional correlation with customized contrastive formulation objective, typically in a label-supervised manner [35, 27, 75, 1, 38], where local features within the same semantic label form the positive pairs. In our setting, we eliminate the need of segmentation labels via establishing correspondence among detected keypoints and apply local self-supervision on the keypoint features.
Keypoint Descriptors have been successfully exploited in various domains to improve the performance on tasks such as pose estimation [19, 4], image generation [57] and object detection [79, 17]. Specifically, in biomedical imaging, keypoint descriptors have found applications in a variety of areas, including anatomy object classification [31], medical image retrieval [74], segmentation tasks [69, 61], motion estimation of organs [53], and medical image registrations [21, 3]. They aim to incorporate sparse and localized key points into the training process, in order to attend the network to the most important features within the image. This is typically done via first localizing a sparse set of keypoint locations in the image, where keypoint features are obtained and their interactions are modeled. Finally, the feature is aggregated and used for downstream task prediction. In this work, we propose to inject learnable keypoint descriptors into the UNet building block so that the network is able to model complex long-range attention without a significant increase in computational cost.
Fig. 1 illustrates an overview of the proposed method. Built upon a segmentation backbone (i.e., UNet), we attach our proposed keypoint-augmented fusion layer (Sec. 3.1) after each convolutional block, to inject long-range self-attention into the original convolutional feature map. Further, we propose both global and local self-supervision (Sec. 3.2) to pretrain the network before finetuning with limited annotation.
3.1 Keypoint-Augmented Fusion (KAF) Layer
The proposed KAF layer is diagrammed in the red dashed box in Fig. 1. Given an input image  $\mathcal{X} \in \mathbb{R}^{W \times H \times C}$ , we start with detecting keypoints  $\mathcal{K} \in \mathbb{R}^{N \times 2}$  using a keypoint detector  $\mathcal{D}_k: \mathbb{R}^{W \times H \times C} \to \mathbb{R}^{N \times 2}$ . In particular, we employ the Scale-Invariant Feature Transform (SIFT) [41]. We note that alternative keypoint detection methods can also be applied, and we have included the ablation in the appendix. At  $l$ -th layer of the 2D UNet, we acquire  $N$  CNN features  $\mathcal{F}_k^l \in \mathbb{R}^{N \times C^l}$  from the dense convolutional feature map  $\mathcal{F}^l \in \mathbb{R}^{W^l \times H^l \times C^l}$ , based on the keypoint coordinates  $\mathcal{K}^l$ . Note that at each scale,  $\mathcal{K}^l$  is acquired by re-scaling  $\mathcal{K}$  with the resolution ratio between the input image and the feature map. Then, we map the keypoint features to the embedding space with dimension  $E \in \mathbb{R}$  with a projector  $\phi: \mathbb{R}^{N \times C^l} \to \mathbb{R}^{N \times E}$ . We use a single-layer MLP in our implementation as the projector. The embedded features are then fed into a Transformer  $\mathcal{T}: \mathbb{R}^{N \times E} \to \mathbb{R}^{N \times E}$  to learn the self-attention among keypoints. and we define the transformed feature as  $U^l \in \mathbb{R}^{N \times E}$  (which is further used as the keypoint feature at layer  $l$  for calculating the SSL losses). Additionally, to properly propagate the keypoint features learned at the current scale  $l$  to  $l + 1$  within the convolutional UNet,  $U^l$  is scattered back into the feature grid indexed by the rescaled keypoint position  $\mathcal{K}^l$ . The output is sparse feature map  $\mathcal{F}_s^l$  which retains information exclusively at the keypoint positions, and the rest of the feature map are assigned with zeros.
Figure 1: Top: The overview of the proposed SSL framework that incorporates both local and global self-supervision with the keypoint guidance. The UNet decoder is omitted for better readability. Bottom: The proposed Keypoint-Augmented Fusion Layer (KAF layer), which learns long-range spatial dependencies among localized keypoint features. We insert the KAF layer after each encoder block of the UNet backbone to augment the original convolutional features.
We further attach a two-layer CNN  $\mathcal{D}_l$  to diffuse the sparse keypoint feature map into a dense image feature. Lastly, we concatenate the diffused keypoint feature map with the input feature map  $F^l$ , constituting the final output of the proposed layer, represented as  $\mathcal{F}_o^l = \mathrm{Concat}(F^l,\mathcal{D}_l(\mathcal{F}_s^l))$ .
3.2 Keypoint-Augmented Self-Supervised Learning
Furthermore, we design the global and local self-supervised learning algorithms to enhance the keypoint-augmented feature learning, thus the learned feature representations can be well generalized to different downstream tasks through finetuning.
Global SSL Loss. In order to apply a global contrastive loss (e.g., [9, 72]) that learns image wise feature similarity, we first aggregate multiscale keypoint features and extract a keypoint-enhanced global representation for each input image. Specifically, keypoint features from four encoder blocks are first concatenated as  $\text{Concat}(U^1, U^2, U^3, U^4)$ , which is then fed into a Multilayer Perceptron (MLP) to get a global feature  $g_i$ , where  $i$  is the index of the input image from the dataset. In our training, we adopt the global contrastive loss from PCL [72] on the keypoint-enhanced global feature, which assumes 2D slices within a certain positional distance (within the 3D volume) threshold to be anatomically similar and constitute similar pairs. Formally, the global self-supervised learning loss function can be formulated as:
$$
L _ {g l o b a l} ^ {i} = - \frac {1}{\left| \Omega_ {i} ^ {+} \right|} \sum_ {j \in \Omega_ {i} ^ {+}} \log \frac {e ^ {\frac {\operatorname {s i m} (g _ {i} , g _ {j})}{\tau}}}{\sum_ {k = 1} ^ {2 N} 1 _ {i \neq k} \cdot e ^ {\frac {\operatorname {s i m} (g _ {i} , g _ {k})}{\tau}}}, \tag {1}
$$
where  $\text{sim}(\cdot, \cdot)$  is the cosine similarity between two vectors in the embedding space,  $\tau$  indicates the temperature term.  $\Omega^{+}$  is the set of positive images to the input  $x_{i}$ .
We additionally keep the global contrastive loss applied on the features from the last layer of the UNet encoder, as commonly done in existing SSL literature [72, 7, 9]. To do so, a global pooling layer transforms the feature  $\mathcal{F}^{last}$  from  $\mathbb{R}^{W^{last} \times H^{last} \times C^{last}}$  to  $\mathbb{R}^{1 \times C^{last}}$ . Afterward, an MLP projector is applied to obtain the global embedding for contrastive objectives, denoted as  $c_{i}$ . Similar to Eq. 1, the loss can be formulated by substituting  $g_{i}$  with  $c_{i}$  and  $g_{j}$  with  $c_{j}$  in Eq. 1. As we adopt PCL loss [72] to the global CNN feature, we denote this loss as  $L_{PCL}^{i}$  in the following context.
Local SSL Loss. While [72] only considers image-wise similarity based on the slice positions, we claim that pixel-wise similarity across slices should also be exploited to supervise the keypoint feature learning, enhancing fine-grained and localized control of the learned representation. We
leverage the intrinsic structural correlation in 3D medical imaging, where nearby slices display strong local similarities. For instance, in cardiac imaging, anatomical structures such as blood vessels and ventricles span across multiple 2D slices within the 3D volume. Therefore, keypoints within a certain spatial distance from adjacent slices are likely to be semantically correlated and form positive pairs.
In order to locate the correspondence for each keypoint, we define two major criteria based on heuristics: (1) the spatial distance in the 2D plane between two correspondences should be within a threshold; (2) the distance between their SIFT features should be similar. Correspondingly, we first sample two positive slices along with their keypoints  $(\mathcal{X}_i,\mathcal{K}_i)$  and  $(\mathcal{X}_j,\mathcal{K}_j)$ , where the positional distance between  $\mathcal{X}_i$  and  $\mathcal{X}_j$  fall within the positive threshold defined in [72]. For each keypoint  $a\in \mathcal{X}_i$ , we aim to find its correspondence from  $\mathcal{X}_j$ . We do so by first filtering out a set of keypoints whose Manhattan distances with  $a$  is less than a predefined threshold. Then, we find the keypoint  $b$  as the correspondence of  $a$  based on the closest L2 distance of the SIFT features. Once matched,  $(a,b)$  is added into the groundtruth match set  $\mathcal{M}\subset \mathcal{A}\times \mathcal{B}$ . However, if no points in slice  $\mathcal{X}_j$  meet the criteria, the candidate point  $a$  is treated as a negative sample and added to the unmatched set  $\mathcal{I}\subseteq \mathcal{A}$ . Similarly, for slice  $\mathcal{X}_j$ , the candidate point without match is added to the unmatched set  $\mathcal{J}\subseteq \mathcal{B}$ .  $\mathcal{M},\mathcal{I},\mathcal{J}$  are then used as the groundtruth for the keypoint feature matching detailed below.
We repurpose SuperGlue [54] as the backbone method to learn the correspondence between keypoint features on slice  $\mathcal{X}_i$  and keypoints on slice  $\mathcal{X}_j$ , supervised by the correspondence extracted above. We first compute the feature associated with keypoints  $U_i = \mathrm{Concat}(U_i^1, U_i^2, U_i^3, U_i^4)$ , and  $U_j = \mathrm{Concat}(U_j^1, U_j^2, U_j^3, U_j^4)$ , along with their groundtruth matches  $\mathcal{M}$  and unmatched sets  $\mathcal{I}$  and  $\mathcal{J}$ . Then, SuperGlue takes  $U_i, U_j, \mathcal{K}_i$ , and  $\mathcal{K}_j$  as inputs and computes an output  $\bar{P} \in \mathbb{R}^{(N_i + 1) \times (N_j + 1)}$  through self- and cross-graph message passing. Here,  $\bar{P}_{a,b}$  represents the probability that keypoint  $a \in \mathcal{K}_i$  matches with  $b \in \mathcal{K}_j$ . Given  $\mathcal{M}, \mathcal{I}$ , and  $\mathcal{J}$ , we minimize the negative log-likelihood of the assignment  $\bar{P}$ , formulated as follows:
$$
L _ {\text {l o c a l}} ^ {i} = - \sum_ {(a, b) \in M} \log \bar {P} a, b - \sum_ {a \in I} \log \bar {P} _ {a, N _ {i} + 1} - \sum_ {b \in J} \log \bar {P} _ {N _ {j} + 1, b}. \tag {2}
$$
By establishing such localized correlation among keypoint features, we assume the network is able to better recognize local correspondence even when the input image undergoes various transformations. Such localized correlation ensures that the network can maintain a consistent understanding of the relationships between keypoint features, regardless of changes in the image appearance or orientation. Our assumption is empirically verified from the computed equivariant self-attention map shown in Fig. 3.
Total SSL Loss. With the proposed keypoint-augmented fusion layer integrated into existing CNN backbone, the model can be pretrained with the above global and local losses, optionally with existing global contrastive loss (in our case, PCL [72]). The final objective is:
$$
L _ {t o t a l} = w _ {1} \cdot L _ {P C L} + w _ {2} \cdot L _ {g l o b a l} + w _ {3} \cdot L _ {l o c a l}, \tag {3}
$$
where the coefficients  $w_{1}, w_{2}$ , and  $w_{3}$  are employed to balance the contributions of each term. In our experiments, we use a grid search method to optimize these coefficients.
Datasets. We conduct experiments on two publicly accessible cardiac MRI datasets for the task of segmentation under limited annotation. For fair benchmarking, we follow the same preprocessing steps as [72]. For each cross-validation fold,  $N$  images are held-out for validation, and varying  $M$  images from the remaining images are used for the few-shot training. We additionally include results on a non-cardiac CT dataset Synapse<sup>1</sup> in the appendix to illustrate the generalization of our method.
CHD (Congenital Heart Disease) [67] is a CT dataset consisting of 68 3D cardiac images, covering patent ages ranging from 1 month to 21 year. It includes 14 types of congenital heart disease and the segmentation labels consists of seven distinct substructures: left ventricle (LV), right ventricle (RV), left atrium (LA), right atrium (RA), myocardium (Myo), aorta (Ao), and pulmonary artery (PA). For each cross-validation fold,  $N = 18$  images are used for validation, and  $M$  images are selected from the remaining 50 images for few-shot training.
<https://www.synapse.org/#!Synapse:syn3193805/wiki/217789>