vox2vec: A Framework for Self-supervised Contrastive Learning of Voxel-level Representations in Medical Images

TL;DR

vox2vec is a self-supervised contrastive learning framework that produces multi-scale voxel representations in medical images, improving segmentation performance with fewer trainable parameters.

Contribution

It introduces a novel contrastive SSL method for voxel-level representations using a Feature Pyramid Network in medical imaging.

Findings

Outperforms existing SSL techniques in segmentation tasks

Achieves competitive performance with significantly fewer trainable parameters

Effective in linear, non-linear, and end-to-end training setups

Abstract

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 .