Fairness-Aware Checkpoint Screening for Neural Models via Multi-Task Learning and Monte Carlo Dropout

TMLR Paper8797 Authors

Abstract: Machine learning models deployed in high-stakes domains often exhibit trade-offs between predictive performance and group fairness, and identifying models that navigate this trade-off remains challenging in practice. We present a neural in-processing framework that combines multi-task learning and Monte Carlo (MC) dropout to support uncertainty-aware checkpoint selection for fairness-aware prediction. Our approach jointly predicts a primary target and a protected attribute using a shared representation, then evaluates saved training checkpoints using predictive performance and a group-fairness objective based on disparate impact ratio. We use MC dropout to characterize checkpoint-level predictive variability and perform Pareto-based screening over fairness–performance trade-offs on a validation set, enabling selection of candidate checkpoints that better balance these competing objectives. We evaluate the approach on three datasets: ADULT, MIMIC-III, and SNAPSHOT, and compare against standard fairness baselines including reweighing, adversarial reweighted learning, and FairRF where applicable. Across these settings, the proposed selection strategy often identifies checkpoints with improved demographic-parity trade-offs relative to baseline models, while maintaining competitive predictive performance. We further provide qualitative saliency-map analyses to illustrate how feature emphasis may shift across selected checkpoints. Our results suggest that uncertainty-aware checkpoint screening can serve as a practical mechanism for navigating fairness–performance trade-offs in neural prediction pipelines. We discuss limitations, including dependence on neural architectures with MC dropout and the current focus on a demographic-parity-style fairness criterion.

Submission Type: Long submission (more than 12 pages of main content)

Assigned Action Editor: ~Novi_Quadrianto1

Submission Number: 8797