Reliable Uncertainty Under Class Imbalance and Distribution Shift: Class-Conditional Conformal Prediction of Multiple Sclerosis

Objectives: To evaluate whether class-conditional conformal prediction (CP) can provide reliable uncertainty quantification (UQ) under severe class imbalance and distribution shift, using multiple sclerosis (MS) diagnosis from magnetic resonance imaging (MRI) as a clinical exemplar. Methods: We evaluated marginal and class-conditional CP using 720 T2-weighted MRI scans (142 MS, 578 controls). A convolutional neural network trained on 3 T data was evaluated under distribution shift (1.5 T acquisitions and synthetic image degradations). Through 100 Monte Carlo experiments, we assessed coverage guarantees, class-specific performance, and relationships between calibration set size, coverage variance, and uncertainty. Results: Marginal CP severely under-covered the minority MS class (16.9% mean coverage at 1.5 T vs. 95.2% for controls) despite valid population-level guarantees. Class-conditional CP dramatically improved MS coverage to 77.5% at 1.5 T and 85.8% at 3 T, significantly reducing severe undercoverage (89% control coverage. Minority class coverage variance increased due to limited calibration samples, matching theoretical Beta-binomial predictions. CP maintained validity under distribution shift; prediction set sizes scaled monotonically with shift severity, yielding clinically interpretable UQ. Conclusions: Class-conditional CP successfully mitigates systematic undercoverage of minority disease classes while maintaining validity under distribution shift. The approach offers a practical, model-agnostic solution for uncertainty quantification applicable across clinical AI systems, though increased coverage variance for less represented conditions reflects fundamental statistical constraints. By characterizing these variance trade-offs, this framework enables more reliable deployment of diagnostic AI in heterogeneous clinical environments across diverse medical domains where minority disease class detection is critical.