A Point in the Right Direction: Vector Prediction for Spatially-aware Self-supervised Volumetric Representation Learning

TL;DR

VectorPOSE introduces two novel self-supervised tasks, Vector Prediction and Boundary-Focused Reconstruction, to enhance spatial understanding in 3D medical imaging, leading to improved segmentation performance with limited labels.

Contribution

The paper proposes VectorPOSE, a new self-supervised method that incorporates spatially-aware pretext tasks for better 3D medical image representation learning.

Findings

Outperforms state-of-the-art methods on 3D segmentation tasks

Enhances learning in limited annotation scenarios

Improves both global and local spatial feature understanding

Abstract

High annotation costs and limited labels for dense 3D medical imaging tasks have recently motivated an assortment of 3D self-supervised pretraining methods that improve transfer learning performance. However, these methods commonly lack spatial awareness despite its centrality in enabling effective 3D image analysis. More specifically, position, scale, and orientation are not only informative but also automatically available when generating image crops for training. Yet, to date, no work has proposed a pretext task that distills all key spatial features. To fulfill this need, we develop a new self-supervised method, VectorPOSE, which promotes better spatial understanding with two novel pretext tasks: Vector Prediction (VP) and Boundary-Focused Reconstruction (BFR). VP focuses on global spatial concepts (i.e., properties of 3D patches) while BFR addresses weaknesses of recent reconstruction methods to learn more effective local representations. We evaluate VectorPOSE on three 3D medical image segmentation tasks, showing that it often outperforms state-of-the-art methods, especially in limited annotation settings.