Bayesian analysis of the prevalence bias: learning and predicting from imbalanced data

TL;DR

This paper introduces a Bayesian framework to address prevalence bias in imbalanced datasets, providing a theoretically grounded loss function and predictive rules that improve model robustness in real-world applications.

Contribution

It develops a novel bias-corrected loss function and predictive methods based on Bayesian risk minimization, specifically targeting prevalence bias in training data.

Findings

Bias-corrected loss improves model calibration.

Framework integrates seamlessly with deep learning.

Provides principled alternative to heuristic methods.

Abstract

Datasets are rarely a realistic approximation of the target population. Say, prevalence is misrepresented, image quality is above clinical standards, etc. This mismatch is known as sampling bias. Sampling biases are a major hindrance for machine learning models. They cause significant gaps between model performance in the lab and in the real world. Our work is a solution to prevalence bias. Prevalence bias is the discrepancy between the prevalence of a pathology and its sampling rate in the training dataset, introduced upon collecting data or due to the practioner rebalancing the training batches. This paper lays the theoretical and computational framework for training models, and for prediction, in the presence of prevalence bias. Concretely a bias-corrected loss function, as well as bias-corrected predictive rules, are derived under the principles of Bayesian risk minimization. The loss exhibits a direct connection to the information gain. It offers a principled alternative to heuristic training losses and complements test-time procedures based on selecting an operating point from summary curves. It integrates seamlessly in the current paradigm of (deep) learning using stochastic backpropagation and naturally with Bayesian models.