# Keypoint-Augmented Self-Supervised Learning for Medical Image Segmentation with Limited Annotation

arXiv: https://arxiv.org/abs/2310.01680 · HTML: https://mlsanity.com/paper/2310.01680/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 57 of its 77 references that are in the map.
- First posted on arXiv 2 Oct 2023.

## Builds on (most cited first, 30 of 57)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [Unsupervised Representation Learning by Predicting Image Rotations](https://mlsanity.com/paper/1803.07728/) (2018)
- [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://mlsanity.com/paper/1911.11763/) (2019)
- [RePaint: Inpainting using Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2201.09865/) (2022)
- [Context Encoders: Feature Learning by Inpainting](https://mlsanity.com/paper/1604.07379/) (2016)
- [Big Self-Supervised Models are Strong Semi-Supervised Learners](https://mlsanity.com/paper/2006.10029/) (2020)
- [SuperPoint: Self-Supervised Interest Point Detection and Description](https://mlsanity.com/paper/1712.07629/) (2017)
- [Contrastive Multiview Coding](https://mlsanity.com/paper/1906.05849/) (2019)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [A Survey on Deep Learning in Medical Image Analysis](https://mlsanity.com/paper/1702.05747/) (2017)
- [Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images](https://mlsanity.com/paper/2201.01266/) (2022)
- [Self-Supervised Learning of Pretext-Invariant Representations](https://mlsanity.com/paper/1912.01991/) (2019)
- [UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation](https://mlsanity.com/paper/1912.05074/) (2019)
- [Self-Supervised Pre-Training of Swin Transformers for 3D Medical Image Analysis](https://mlsanity.com/paper/2111.14791/) (2021)
- [Contrastive learning of global and local features for medical image segmentation with limited annotations](https://mlsanity.com/paper/2006.10511/) (2020)
- [Discriminative Unsupervised Feature Learning with Exemplar Convolutional Neural Networks](https://mlsanity.com/paper/1406.6909/) (2014)
- [Local Aggregation for Unsupervised Learning of Visual Embeddings](https://mlsanity.com/paper/1903.12355/) (2019)

## Built on by (most cited first, 7 of 7)

- [SegICL: A Multimodal In-context Learning Framework for Enhanced Segmentation in Medical Imaging](https://mlsanity.com/paper/2403.16578/) (2024)
- [Structure-aware Semantic Discrepancy and Consistency for 3D Medical Image Self-supervised Learning](https://mlsanity.com/paper/2507.02581/) (2025)
- [Retinal IPA: Iterative KeyPoints Alignment for Multimodal Retinal Imaging](https://mlsanity.com/paper/2407.18362/) (2024)
- [Devil is in Details: Locality-Aware 3D Abdominal CT Volume Generation for Self-Supervised Organ Segmentation](https://mlsanity.com/paper/2409.20332/) (2024)
- [K-Stain: Keypoint-Driven Correspondence for H&E-to-IHC Virtual Staining](https://mlsanity.com/paper/2511.06709/) (2025)
- [Coordinative Learning with Ordinal and Relational Priors for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/2511.11276/) (2025)
- [Learning from Noisy Prompts: Saliency-Guided Prompt Distillation for Robust Segmentation with SAM](https://mlsanity.com/paper/2604.23314/) (2026)

## Abstract (arXiv)

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.