vox2vec: A Framework for Self-supervised Contrastive Learning of Voxel-level Representations in Medical Images
This work addresses the need for efficient and effective representation learning in medical imaging, particularly for segmentation tasks, but it is incremental as it builds on existing SSL and FPN methods.
The paper tackles the problem of learning voxel-level representations in medical images by introducing vox2vec, a self-supervised contrastive learning framework that uses a Feature Pyramid Network to produce multi-scale representations, and it shows that vox2vec outperforms existing SSL techniques in segmentation tasks while achieving competitive performance with 50 times fewer trainable parameters.
This paper introduces vox2vec - a contrastive method for self-supervised learning (SSL) of voxel-level representations. vox2vec representations are modeled by a Feature Pyramid Network (FPN): a voxel representation is a concatenation of the corresponding feature vectors from different pyramid levels. The FPN is pre-trained to produce similar representations for the same voxel in different augmented contexts and distinctive representations for different voxels. This results in unified multi-scale representations that capture both global semantics (e.g., body part) and local semantics (e.g., different small organs or healthy versus tumor tissue). We use vox2vec to pre-train a FPN on more than 6500 publicly available computed tomography images. We evaluate the pre-trained representations by attaching simple heads on top of them and training the resulting models for 22 segmentation tasks. We show that vox2vec outperforms existing medical imaging SSL techniques in three evaluation setups: linear and non-linear probing and end-to-end fine-tuning. Moreover, a non-linear head trained on top of the frozen vox2vec representations achieves competitive performance with the FPN trained from scratch while having 50 times fewer trainable parameters. The code is available at https://github.com/mishgon/vox2vec .