Distribution-aware Fairness Learning in Medical Image Segmentation From A Control-Theoretic Perspective

Yujin Oh, Pengfei Jin, Sangjoon Park, Sekeun Kim, Siyeop yoon, Jin Sung Kim, Kyungsang Kim, Xiang Li, Quanzheng Li

TL;DR: We introduce Distribution-aware Mixture of Experts (dMoE), a fairness learning-driven approach inspired by optimal control theory, to mitigate demographic and clinical biases in medical image segmentation.

Abstract: Ensuring fairness in medical image segmentation is critical due to biases in imbalanced clinical data acquisition caused by demographic attributes (e.g., age, sex, race) and clinical factors (e.g., disease severity). To address these challenges, we introduce Distribution-aware Mixture of Experts (dMoE), inspired by optimal control theory. We provide a comprehensive analysis of its underlying mechanisms and clarify dMoE's role in adapting to heterogeneous distributions in medical image segmentation. Furthermore, we integrate dMoE into multiple network architectures, demonstrating its broad applicability across diverse medical image analysis tasks. By incorporating demographic and clinical factors, dMoE achieves state-of-the-art performance on two 2D benchmark datasets and a 3D in-house dataset. Our results highlight the effectiveness of dMoE in mitigating biases from imbalanced distributions, offering a promising approach to bridging control theory and medical image segmentation within fairness learning paradigms. The source code is available at https://github.com/tvseg/dMoE.

Lay Summary:

Medical images play a key role in helping doctors diagnose and treat patients. However, AI systems that analyze these images can sometimes show biased performance—especially when they are trained on imbalanced data, such as having more images from certain age groups, genders, or disease stages. This means they may work better for some patient groups than others. To address this problem, we developed a new approach called Distribution-aware Mixture of Experts (dMoE), inspired by concepts from control theory—a field that focuses on managing and guiding complex systems.

This method helps AI systems adjust to differences in patient data, making the image analysis more fair and reliable. We tested dMoE with several types of AI models and found that it worked well across different tasks. It performed better than other methods on both publicly available 2D datasets and our own 3D medical image data. By including information like age, sex, and disease severity, our approach reduces bias and improves performance.

This work shows how tools from other areas of science, like control theory, can help make AI in healthcare more fair and effective. We believe our distribution-aware approach is a meaningful step toward fairer AI systems in healthcare, especially in real-world situations where patient data can be unbalanced or varied.

Link To Code: https://github.com/tvseg/dMoE

Primary Area: Applications->Health / Medicine

Keywords: Fairness Learning, Medical Image Segmentation, Distribution, Optimal Control

Submission Number: 5156