Medical contrastive learning stops forcing invariance: ProbSSL models embeddings as distributions and reaches 84.89 mean Dice on BTCV
Synopsis
Medical self-supervised contrastive learning assumes two augmented views of the same image should share one semantic representation, but rotation and intensity perturbations can erase subtle yet critical anatomical and pathological cues, so rigid alignment suppresses task-relevant information. ProbSSL models each image embedding as a Gaussian, with the mean carrying shared semantics and the variance recording augmentation-induced variability, and aligns distributions by KL divergence so dimensions with larger variability receive weaker alignment pressure. After pre-training on CT, transfer to BTCV, MM-WHS, CC-CCII and BraTS gives higher average scores than the compared general and medical self-supervised methods on all four downstream tasks, including 84.89 mean Dice on BTCV.
Interpretation
Medical contrastive learning no longer collapses two augmented views of one image into a single deterministic vector; each embedding now carries a variance that tracks augmentation. Conventional contrastive learning forces positive pairs into strict alignment, and that invariance assumption can become spurious in medical images because spatial transforms distort anatomy and intensity perturbations mask pathology. In a contrastive-only controlled ablation, raising augmentation strength from A0 to A1 and A2 lowered the deterministic baseline on BTCV from 81.07 to 80.75 and then 80.65 Dice, while ProbCL rose from 83.02 to 83.20 and then 83.27.
Alignment happens at the distribution level: means are pulled together weighted by inverse variance, and variances are matched too, so dimensions made uncertain by augmentation are constrained less. Treating embeddings as deterministic points allowed only whole-vector alignment, with no way to separate reliable dimensions from augmentation-corrupted ones. Across BTCV, MM-WHS, CC-CCII and BraTS, each loss scored higher in its probabilistic setting than in its deterministic version, for example contrastive loss rising from 81.07 to 83.02 on BTCV and from 75.98 to 78.95 on BraTS.
Contrastive, predictive and generative self-supervised objectives share one probabilistic latent space, constraining global semantic consistency, geometric transformation prediction and masked-region reconstruction respectively. Earlier multi-task self-supervised methods added pretext tasks side by side, each acting on deterministic embeddings. Adding DCL, SCL, the prior term, then ProbPL and ProbGL one at a time on BTCV and MM-WHS gave the full ProbSSL the highest score, with 84.89 mean Dice on BTCV and 91.57 mean Dice on MM-WHS.
Perspective
The demonstrated scope is 3D medical imaging: pre-training on 1,589 unlabeled CT scans, then transfer to three CT tasks (BTCV, MM-WHS, CC-CCII) and one MRI task (BraTS), scored by segmentation Dice, HD95 and classification accuracy.
Performance can still degrade under low signal-to-noise input, motion or streak artifacts, large morphological variability and long-tail conditions; rotation prediction and masked reconstruction may not help equally across modalities and organs; diagonal covariance ignores cross-dimension correlations; and compute and data scale limited the breadth of pre-training and ablations.