A diffusion-conditioned representation learning framework for disease classification in medical imaging

Area of Science:

• Medical Imaging
• Deep Learning
• Artificial Intelligence

Background:

• Deep learning models in medical imaging show advancements but often overfit and are overconfident.
• Deterministic discriminative models struggle with distributional shifts and uncertainty quantification.

Purpose of the Study:

• To propose diffusion-conditioned representation learning for improved disease classification and distributional shift detection.
• To leverage internal diffusion model dynamics for noise-level awareness and diverse semantic representations.
• To encode aleatoric uncertainty for regularizing predictions during training.

Main Methods:

• Developed a diffusion-conditioned representation learning framework.
• Extracted noise-level awareness and diverse semantic representations from diffusion model dynamics.
• Sampled features from multiple denoising time steps to encode aleatoric uncertainty.
• Utilized entropy-regularized loss for prediction regularization.

Main Results:

• Framework demonstrated competitive performance across five diverse clinical datasets.
• Achieved reduced prediction entropy and improved robustness to input degradation.
• Showcased effectiveness in classification tasks via negative entropy-accuracy correlation, expected calibration error (ECE), and dynamic uncertainty.
• Obtained AUCs of 83.81%, 97.45%, 86.15%, and 96.75% on datasets 1-4, respectively.
• Demonstrated calibrated performance with reduced ECE values.

Conclusions:

• Diffusion-conditioned representation learning offers a more efficient strategy for medical imaging tasks.
• The proposed method enhances disease classification accuracy and robustness.
• The framework effectively quantifies uncertainty and improves model calibration.