Self-Supervised Learning for Out-of-Distribution Detection in Medical Imaging: A Comparative Study of JEPA, MAE, and Supervised Learning

Nikhileswara Rao Sulake*, Likith Vinay Kumar Busam*, Siva Teja Reddy Annapureddy*, Harsha Allapati*

* Equal contribution

Code
Architecture overview

System Architecture. A shared ViT-Small encoder backbone is pretrained with JEPA, MAE, and supervised objectives on in-distribution chest X-rays, then evaluated via linear probing, OOD detection with energy and Mahalanobis scores, and robustness analysis under corruption shifts.

Abstract

Out-of-distribution (OOD) detection is a critical capability for the safe deployment of machine learning models in medical imaging, where systems inevitably encounter data from diverse sources that differ from the training distribution. This study presents a comprehensive comparative analysis of three representation learning paradigms—Joint Embedding Predictive Architecture (JEPA), Masked Autoencoder (MAE), and supervised learning—for chest X-ray out-of-distribution detection.

Using the TB Chest Radiography Database as in-distribution data and the Montgomery and Shenzhen tuberculosis datasets as out-of-distribution test sets, we evaluate OOD detection performance, cross-dataset generalization, and robustness to common image corruptions. Our experimental results demonstrate that self-supervised methods significantly outperform supervised learning on energy-based OOD detection, achieving 77–87% AUROC compared to 58–61% for supervised approaches. Furthermore, MAE exhibits exceptional robustness to contrast perturbations, maintaining 98.3% classification AUC under severe corruption where JEPA and supervised methods degrade to 67–76%. Mahalanobis distance-based detection achieves near-perfect performance (95–99% AUROC) across all methods, suggesting that distribution-aware scoring can compensate for representation differences.

OOD Detection Self-Supervised Learning Medical Imaging Chest X-ray JEPA MAE Vision Transformer Tuberculosis

Key Findings

🎯

Self-supervised methods outperform supervised for energy-based OOD

JEPA and MAE achieve 77–87% AUROC on energy-based OOD detection, compared to only 58–61% for supervised learning. The supervised model's FPR@95 reaches 93.1% on Shenzhen—nearly all OOD samples pass undetected.

🛡️

MAE is remarkably robust to contrast perturbations

Under severe contrast corruption (severity 5), MAE maintains 98.3% AUC while JEPA and supervised learning degrade to 67–76%. This is directly attributable to MAE's pixel-reconstruction objective.

📐

Mahalanobis distance rescues supervised representations

Replacing energy scoring with Mahalanobis distance boosts all methods to 95–99% AUROC. Supervised learning even achieves 99.7% AUROC with 0% FPR@95 on Shenzhen— distribution-aware scoring can compensate for representation weaknesses.

🔄

MAE generalises best across hospital domains

Cross-dataset classification shows MAE outperforms JEPA and supervised learning by 10–26 percentage points when applied to OOD hospital sources, suggesting reconstruction pretraining encodes transferable anatomical features.

Experimental Results

In-Distribution Classification

All methods achieve strong in-distribution classification on the TB Chest Radiography Database (n = 840 validation samples).

Method AUROC Accuracy ECE ↓ NLL ↓
JEPA 0.939 0.917 0.029 0.206
MAE 0.993 0.971 0.013 0.085
Supervised 0.996 0.980 0.007 0.064

OOD Detection — Energy Score

Self-supervised methods substantially outperform supervised learning, especially on the Shenzhen dataset.

Method Montgomery Shenzhen
AUROC FPR@95 ↓ AUROC FPR@95 ↓
JEPA 0.771 0.540 0.718 0.633
MAE 0.769 0.577 0.866 0.358
Supervised 0.607 0.814 0.582 0.931

OOD Detection — Mahalanobis Distance

Distribution-aware scoring achieves excellent performance across all methods (≥95% AUROC), compensating for representation differences.

Method Montgomery Shenzhen
AUROC FPR@95 ↓ AUROC FPR@95 ↓
JEPA 0.956 0.186 0.985 0.077
MAE 0.994 0.030 0.984 0.085
Supervised 0.991 0.024 0.997 0.000

Robustness to Image Corruptions (Severity 5)

MAE demonstrates exceptional robustness to contrast perturbations, maintaining 98.3% AUC where other methods degrade significantly.

Method Clean Gaussian Noise Blur Contrast ⚡
JEPA 0.990 0.983 0.981 0.672
MAE 0.992 0.984 0.992 0.983
Supervised 0.996 0.992 0.994 0.758

Cross-Dataset Generalisation

MAE demonstrates superior transfer to OOD hospital datasets, suggesting reconstruction-based pretraining learns more generalizable features.

Method ID (AUC) Montgomery (AUC) Shenzhen (AUC)
JEPA 0.992 0.554 0.430
MAE 0.995 0.607 0.695
Supervised 0.996 0.471 0.592

Visualizations

Energy score distributions

Figure 1: Energy score distributions for in-distribution and out-of-distribution samples across all pretraining methods. Self-supervised methods (JEPA, MAE) show better separation between distributions, enabling more reliable detection.

Robustness plot

Figure 2: Classification AUC vs. corruption severity for all corruption types. MAE maintains stable performance across all severity levels, while JEPA and Supervised methods show progressive degradation under contrast perturbations.

Reliability diagram

Figure 3: Reliability diagram showing calibration of probability estimates. Perfect calibration would follow the diagonal. Supervised learning shows the best calibration on in-distribution data.

Embedding visualization

Figure 4: Detailed embedding visualization showing the distribution of samples from in-distribution and out-of-distribution datasets across the learned representation spaces.

Figure 5: 3D Embedding Space Visualizations

t-SNE visualizations colored by dataset source and disease label reveal qualitatively different geometric organizations across pretraining methods.

JEPA 3D embedding

(a) JEPA
Intermediate structure with moderate class separation; OOD samples not distinctly clustered.

MAE 3D embedding

(b) MAE
OOD samples occupy systematically different regions, enabling effective Mahalanobis detection.

Supervised 3D embedding

(c) Supervised
Compact ID clusters but OOD samples mixed in, explaining poor energy-based detection.

Method Overview

JEPA

Joint Embedding Predictive Architecture predicts latent representations of masked image regions from context patches. It uses a momentum-updated target encoder (EMA) and a lightweight predictor transformer to learn high-level semantic features without pixel reconstruction. Training objective: L₂ distance between predicted and actual target embeddings.

85–100% context masking 4 target blocks

MAE

Masked Autoencoder uses an asymmetric encoder-decoder architecture. The encoder processes only 25% of visible patches; a lightweight decoder reconstructs masked pixel values. Training objective: MSE on normalized pixel values within masked patches, implicitly learning full image intensity distributions.

75% masking ratio 4-block decoder

Supervised

Supervised Pretraining trains the ViT-S encoder end-to-end with a linear classification head using cross-entropy loss. This directly optimises for the binary normal/tuberculosis classification task, achieving the highest in-distribution accuracy but with limited OOD generalisation.

Binary CE loss Linear head

Shared Backbone: ViT-Small (21.6M params)

All experiments use an identical Vision Transformer Small architecture for fair comparison: 12 transformer blocks · 384-dim embeddings · 6 attention heads · 1536 MLP dim. Input: 224×224 images → 196 × 16×16 patch tokens + CLS token.

50
pretraining epochs
3,360
training samples
ViT-S
backbone

Datasets

In-Distribution

TB Chest Radiography Database

4,200 chest X-ray images with binary labels (Normal vs. Tuberculosis). Class-imbalanced: 700 normal + 3,500 TB images, reflecting clinical screening prevalence.

4,200 images Binary labels 80/20 split

Out-of-Distribution

Montgomery County TB Dataset

138 chest X-rays from Montgomery County, Maryland, USA. Acquired on a Eureka stationary X-ray machine — different equipment and geographic origin.

138 images USA (Maryland) Eureka X-ray

Out-of-Distribution

Shenzhen Hospital TB Dataset

662 chest X-rays from Shenzhen No.3 People's Hospital, China. Acquired on a Philips DR Digital Diagnost system — different country, demographics, and equipment.

662 images China (Shenzhen) Philips DR

Conclusion

This work presents a systematic comparison of self-supervised learning methods for chest X-ray out-of-distribution detection. Our comprehensive evaluation yields several findings with direct practical implications:

  • Self-supervised pretraining provides 15–25 percentage-point advantages over supervised learning for energy-based OOD detection, capturing distributional characteristics beyond discriminative boundaries.
  • Mahalanobis distance achieves excellent performance (≥95% AUROC) regardless of pretraining, suggesting that detection methodology may be as important as representation choice.
  • MAE is the recommended default for medical imaging deployment under distribution shift, offering superior robustness (especially contrast), better cross-dataset transfer, and consistent performance across both energy and Mahalanobis scoring.

Future work extends this framework to echocardiogram segmentation using EchoNet-Dynamic (in-distribution) and EchoNet-Pediatric (OOD), testing whether these patterns generalise to video-based medical imaging and dense prediction tasks.

Paper

paper.pdf