Distribution-aware Fairness Learning in Medical Image Segmentation From A Control-Theoretic Perspective
Authors: Yujin Oh
Pengfei Jin
Sangjoon Park
Sekeun Kim
Siyeop yoon
Jin Sung Kim
Kyungsang Kim
Xiang Li
Quanzheng Li
Venue: ICML 2025 spotlightposter (2025)
Keywords: Fairness Learning
Medical Image Segmentation
Distribution
Optimal Control
Primary area: applications->health_medicine
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.

Distribution-aware Fairness Learning in Medical Image Segmentation From A Control-Theoretie Perspective
Yujin Oh
Pengfei Jin
Sangjoon Park 22 Sekeun Kim Siyeop Yoon Quanzheng Li
Kyungsang Kim
Jin Sung Kim 24 Xiang Li
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., dis- ease severity). To address these challenges, we introduce Distribution-awan Mixture of Experts (dMoE), inspired by optimal control theory. We provide comprehensive analysis of its underly- ing mechanisms and clarify dMoE's role in adapt- ing to heterogeneous distributions in medical im- age segmentation. Furthermore, we integrate dMoE into multiple network architectures, demon- strating broad applicability across diverse med- ical image analysis tasks. By incorporating demo- graphic and clinical factors, dMoE achieves state- of-the-art performance
Imbalanced Clinical Data Distribution
Presade Conce Radisthenay Targes
Hanard Ophihilmoscopy
HAAR10300 Blinn Lesion
Minor
T3
74
Asiin
Stack Race
Whiti
T-stage
Age
Poor Equilty
inproved Expatty
Peformnn Poat
T-stage
Teslage
Distribution-awe
Experts
Equity for Medical Image Segmentation Parformance
two 2D benchmark datasets and 3D in-house dataset. Our results highlight the effectiveness of dMoB in mitigat- ing biases from imbalanced distributions, offering a promising approach to bridging control theory and medical image segmentation within faireess learning paradigms. The source code is available at https: / htt///gtttuu/b com/ wtt.///vee//////eee// dMME..
Figure I. The influence of clinical data distribution on medical image segmentation and the role of dMoE a distribution-and aware control mechanism to address inequity challenges. Transparent bluc lines within the violin plots connect the most densely popu- lated regions of each attribute, visually representing overall equity.
tributes such as gender and race, as well as clinical factors, e.g., disease severity, emphasize the need for fairness-aware learning approaches in medical data analysis (Ktena et al., 2024; Jin et al., 2024). Deep neural networks, due to their data-driven optimization processes, often overfit preva- lent patterns while failing to adequately learn from under- represented ones, leading to potentially biased predictions and weakening fairness across demographic subgroups, as shown in Figure 1. Although advanced fairness training strategies have been actively developed (Tian et al., 2025; 2024; et al., 2024), distributional approaches remain un- derexplored despite the fact that clinical and demographic factors play significant role in many clinical decision- making (Chester et al., 2020; Theodore al., 2023). Clini- cians consider demographic factors such as age, gender, and race alongside patient's 's unique condition to assess both individual and subgroup characteristics. This process not only informs clinical factors like disease stage and decision- making but also implicitly shapes region/population-speee specific
1. Introduction
The inherent imbalance and ill-posed nature of data col- lected in clinical practice, impacted by demographic at-
"Equal contribution 'Center for Advanced Medical Comput- ing and Analysis (CAMCA). Department of Radiology. Mas- sachusetts General Hospital (MGH) and Harvard Medical School, MA 02114, USA 'Department of Radiation Oncology, Yon- sei University College of Medicine, Yonsei University, Scoul 03772, Republic of Korea 'Institute for Innovation in Digital Healthcare. Yonsei University. Seoul 03772, Republic of Ko- "Oncosoft Inc. Scoul 03776, Republic of Korea. Correspon- dence to: Xiang <xli60@ hddu...dddddd.. Quanzheng Li <li cli.quanzheng26 mgh www.hhharvrd..d.......
Proceedings of the nl International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Distribution-aware Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
practice patterns based on cumulative experience and lo- cal demographics. However, current fairnsss learning ap- proaches primarily focus on explicit factors such as de- mographic attributes but neglect implicis/contextual factors such as disease progression patterns or severity, despite the fact that these factors can be potential influence by re- gion/population characteristics.
2. Related Work
2.1. Fairness Learning in Medical Image Segmentation
Faimess--oriented medical image segmentation datasets with demographic information (Tschandl et al., 2018; Tian et al., 2024) have enabled research bias mitigation, driving the development of robust fairness training strategies. Advanced generative approaches (Li et al., 2024; Ktena et al., 2024) aim to generate diverse samples from skewed distributions to mitigate bias. However, challenges remain in addressing the computational demands and achieving high-quality sam- ple generation for high-dimensional medical images, such as 3D whole-body computed tomography (CT), for use in aug- mentation datasets. On the other hand, crrness-focused loss function modifications, such as distributionally robust opti- mization (DRO) (Sagawa et al., 2019) and fair bound scaling (FEBS) (Tian et al., 2024), have been proposed to integrate fairness considerations directly into the optimiza- tion process. While partially effective, these methods are vulnerable to the data distribution within the training batch, limiting their applicability in 3D medical image segmenta- tion, where large batch sizes constrained. Furthermore, existing benchmarks and studies primarily focus on demo- graphic attributes (Tian et al., 2025; Jin et al., 2024), often overlooking critical clinical factors, such as tumor progres- sion and metastasis, which contribute to regional variations in elinical practice patterns due to inherent biases. this study, we propose general method mitigate bias arising from both demographic and clinical factors.
Building on these distributional insights and the need for broader applicability across diverse attributes, we propose Distribution-aware Mixture of Experts (dMoE), a frame- work that adaptively incorporates individual characteristics and distributional patterns into deep neural network train- ing. This architecture is grounded in an in-depth analysis of the Mixture of Experts (MoF) framework (Shazeer et al., 2017) with the adaptation of an optimal control perspective system. By analyzing the structural parallels between MoE and traditional control systems (Aström & Murray, 2021), we reinterpret MoE. as a feedback control mechanism. We further enhance its gating mechanism to incorporate distribu- tional information as switchnw control for adaptive parameter selection, resulting in dMoF. (Figure 2). This approach ensures robust and fair performance across diverse demographic and clinical subgroups.
Moreover, to ensure broad applicability, we integrate dMoE with various network architectures, including transformers (Vaswani et al., 2017) and convolutional neural networks (CNNs), demonstrating the generalizability of its architec- ture and making the framework well-suited for a wide range of medical image analysis tasks. We validate the effective- ness of dMoE based on the segmentation task multiple clinical datasets with diverse segmentation masks for diag- nosis and treatment planning tasks. Experimental results demonstrate that dMoF not only advances state-of-the-art (SOTA) fairness learning approaches but also presents a way to incorporate distributional attributes to provide robust and equitable diagnosis and clinical decision-making across diverse demographic and clinical attributes.
2.2. Mixture of Expert in Multi-distribution Learning
Recent advances in the MoE framework (Shazeer al., 2017) have demonstrated remarkable potential for adapting Al models to diverse data distributions, particularly within continual learning paradigms. MoE enhances robustness and adaptability when confronted with previously unseen data patterns (Rypešé et al., 2024; Yu et al., 2024). In the medical domain, MoE has been effectively extended to address challenges such as multimodal integration (Jiang & Shen, 2024), scanning modality heterogeneity (Zhang et al., 2024), and catastrophic forgetting issues in continual learning (Chen et al., 2024; Wang et al., 2024), unifying these approaches into a cohesive framework that enhances performance. However, theoretical insights into how MoE facilitates the adaptation of disparate distributions to a target distribution remain limited. This study clarifies the under- lying mechanism of MoE dynamic parameter selection, and integrates environmental attributes into its gating mech- anism, enabling distributional adaptation for medical image segmentation.
Our core contributions are summarized as follows:
We reinterpret MoE from the perspective of feedback control mechanism and extend it into dMoE for fairness learning by incorporating distributional awareness.
dMoE operates seamlessly within transformers and CNNs, enabling its use in 2D and 3D segmentation tasks across demographic and clinical attributes.
Extensive medical image segmentation experiments for diagnosis and treatment planning demonstrate d.MEE's robustness, demonstrating its effectiveness in mitigat- ing biases from imbalanced medical data distributions.
2

---

Distribation-wwar Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
2.3. Training Neural Networks Optimal Control
where is the total number of experts, (x) represents the output of the i-th expert network for the input x, and G(x) is the output of the gating network, sparse n-dimensional vector that determines which experts are activated. The MoR approach fundamentally relies on sparsity, where the gating network engages only a limited subset of experts for each input, substantially lowering computational overhead by excluding inactive experts from processing. The gat- ing function calculates G(x) using methods such as Noisy Top-K Gating, which prioritizes the most relevant k experts while nullifying contributions from the remaining Both the gating network and the experts are optimized si- multaneously through backpropagation, ensuring seamless integration and fine-tuning of the MoE layer. This design facilitates a significant expansion of model capacity without proportionally increasing computational demands, making it particularly well-suited for large-scale tasks like language modeling and machine translation.
Neural networks, especially those with shorteut connec- tions, perform complex transformations through successive modifications of hidden state. These networks can be con- ceptualized as undergoing continuous dynamic process, which can be described using ordinary differential equations (ODEs) (Weinan, 2017; Lu et al., 2018). Training these networks resembles solving an optimal control problem, where the objective is to adjust the network parameters, i.e., weights, to minimize loss function (Chen al., 2018; Sun et al., 2024).
k.
Fixed-architectur feedforward neural networks can be re- garded as operating under a non-feedback control. Em- ploying a consistent strategy, however, can be restrictive in dynamic environments (Aström, 1995). To address this limitation, feedback control mechanisms (Doyle et al., 2013; Áström& & Murray, 2021) offer alternative by enabling continuous monitoring and adjustment based on the sys- tem's 's outputs and desired targets. Furthermore, mode- switching control (Yamaguchi et al., 1996; Boskovic & Mehra, 2000) introduces additional flexibility by enabling the policy to alternate between multiple optimized opera- tional modes in response to external inputs. This capacity is essential for managing complex systems across different conditions and offers robust response to diverse opera- tional challenges (Yu et al., 2017). Likewise, in the context of image segmentation, adopting varying strategies- such as leveraging distributional attributes of the input- ppooves beneficial and serves as the motivation for our method.
3.2. Distribution-aware MoE
Now, explain our proposed Distribution-awan Mixture of Experts (dMoE), as illustrated in Figure 2(a). We fur- ther provide detailed network architecture in Appendix A.I. Unlike traditional sparse MoE, our proposed dMoE mod- ule integrates multiple distribution-wise router networks G and set of n expert modules, consisting of shallow multi-layer perceptron (MLP) neural networks, defined as E |Exper, Expert, .... ,xperr,,,, 1- We start by flattening intermediate image embeddings from 1-th layer block out- put, represented as h2 E ,,,,,,, where the total number of embeddings is N₁ WI,,DD. Here, H₁, W₁. D,, and Ch₁ correspond to the height, width, depth, and channel dimensions of the intermediate image embeddings, respectively. For the transformer backbone, which already has a flattened dimension, we omit this pro- cess. Given these flattend embeddings h, and an attribute flag attr, the activated router G identifies the top-k ex- perts. The final output is computed as a weighted sum of the outputs of these selected experts:
into h, E R. (eNi,,h,,
3. Method
In the methods section, we provide theoretical insights and the motivation for proposing the distribution- aware MoE as a solution for fairness learning. We begin by revisiting MoE in Section 3.1 and then expand it into dMoB in Section 3.2. After defining dMoE, in Section 3.3, we further elaborate on the principles of optimal control and demonstrate how mode-switching optimal con-
dMoE can be formulated trol problem, by providing conceptual connections between them, as illustrated in Figure 2(a) and Figure 2(b) (bottom).
k h, = h4 + Σ (h)) E,Ch))
(2)
3.1. Revisiting MoE
where Gatir() outputs a weight matrix prioritizing each expert's contribution in a attr-specific manner. The re- sulting weighted output is combined with h, represent- ing the shared path, to yield the dMoE image embedding h E NxCh, For con-transformer backbones, this rep- resentation is reshaped to match spatial dimensions as h, E E " W,W,,,DC,, where N H,W,D,,
The MoE framework (Shazeer et al., 2017), which serves as the backbone structure for our proposed dMoE, originally leverages sparse gating to achieve computational efficiency while allowing for large model capacity. The output of the MoE layer is defined as follows:
In specific, the router network G computes sparse weights H using Gaussian noise, as follows:
y= Σ E(x),
G Sofmax((eepPppx(((H)))))
(3)
3

---

Distribution-aware Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
Data Attribute
Input
(a) dMoE Segmentation Network
Output
(b) Control System
Attrixune-weie
Non-toodbeck Control
Data Distr bution
nztr
Layer Blocks...
iMME
Ni-1
Controller
Network
Rou Ro (Ch,
Reckallonyyy Tarppl / Tumarr
Attribute Flag (21.72.7..9)
Layer Block
dMoE
Feedback Cantrol
Foodback
Mce Controlter
Neural Network
dMoE (Distribution-awan Mixture Experts)
mclk)
70h, M/h())
Optice/ Race
Distribution-wwar Router Networks
Top K Gating
Mode- switching Control
Eopert
Foodback
Flstlen
Aggent
Unfisiten
Stlal Controlior
Neural Network
Mode swetstiiiiingg
Gt
Skin Lesion Age
actr
of (h)
Pia Ay f(% 10 (k))
Figure 2. (a) Schematic of the dMoE segmentation network for faimess learning, and (b) its interpretation through a control system.
H(x), = (x". W) +O( 1). Softplus(xx" -oeoos----miiii W ),), (4)
into neural network design, including architectural princi- ples, optimization techniques, and regularization strategies. Building on this interpretation, we further show that dMoE naturally corresponds to mode-switching control a classi- cal paradigm in control theory- and highlight its potential to promote fairness and improve generalization in deep learning systems.
D if e; is in top k elements of D,
KeepTop-k(u, k),
-00 otherwise.
(5) where, Softmaax(- function normalizes the selected weights, KeepTop-k(-) retrains only the top-k expert contributions, and W and noiss are trainable weight matrices, and Softplus(-) is a smooth alternative activation function.
Neural networks create complex transformations by repeat- edly updating their internal states (known as hidden fea- tures). A common mathematical description of these trans- formations is given by:
Following the dMoE modules, the final dMoE image em- bedding h is passed to the decoder, which predicts the final output, J. The network is optimized using the segmentation loss:
h141 = h1 + [ht, 0,)
(8)
min )) = PE (y; log p(;)). M
(6)
where h, represents the hidden feature vector at layer 1, and f(h,, 0,) represents the transformation applied at layer 1, parameterized by weights 0,. This iterative updating formula resembles the Fuler discretization method used to approximate continuous-time dynamical systems described by ordinary differential equations (ODEs). The continuous- time counterpart of the above discrete transformation can be represented as:
where M represents any 2D-1o-3D neural network archi- tecture equipped with our proposed dMoE module, y e HWWD D is the ground-truth label, and the predicted output E HWD, which is computed as follows:
$ M(x attr),
where x the input image and attr the attribute flag.
dh,
dt = ((,
3.3. Interpreting dMoE Through Optimal Control
where h, represents the hidden state at continuous time r, and u, denotes a control input or parameters guiding the system dynamics at each instant. To clarify further, the discrete layer index in the neural network corresponds to specific discrete time points sampled from continuous time Similarly, parameters 0, at discrete layers analogous to continuous-time control inputs 4, which represent con- tinuously adjustable parameters that influence the evolution of hidden states over time.
In this section, we provide a control-theoretic interpretation of deep neural networks by viewing their layer-wise wise com- putations as dynamical systems and formulating training as an optimal control problem. This viewpoint offers two key benefits. First, it provides unified framework for ana- lyzing MoE architectures and yields concrete insights into their adaptive shhavior-specifically how baddback-assed control enables more flexible and expressive representations than fixed expert assignments. Second, it facilitates the in- tegration of estalisst concepts from control theory
In supervised learning applications, the training process of neural network aims minimize loss function given
4

---

Distribution-aware Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
labeled data pairs (x, y):
form the columns of a learnable parameter matrix W. and o² denoting the transpose of Φ. Furthermore, for compu- tational efficiency and robustness, we have incorporated Noisy Top-K Gating, corresponding to Eqs. (3), (4), and (5), Specifically, if the function f is linear with respect to its second argument and can be interpreted as the strategy of an expert, where each expert E, acts as a controller with fixed parameters 0,, then the following equation holds:
arg min / E (x.,) (xx,)) (()
(01)
(10)
h/41 h, + f(h,, 0j), h1 x. x h-)-
where f is the prediction at the final layer L., starting from the input x. Here, optimizing over parameters 0, corre- sponds to choosing the best transformations at each discrete step to minimize the discrepancy between predictions and true labels. From the continuous-time viewpoint, this opti- mization corresponds terminal optimal control problem:
S(h, Σ k(h,,)00 = Σ K(h, ((,n,,
(16)
G
arg min E (x.y] c( (uf) dh, 50 (,), hg = x. hy
(11)
where G,(h)) = ((h,, serve as gating weights dynami- cally allocating the influence of each expert E,- This for- mulation naturally leads classical concept in optimal control that aligns with the de- sign philosophy of dMoF. In this framework, the system selects between multiple control strategies based on both the current state h, and external contextual parameters attr. Concretely, we model this behavior by introducing a switch- ing logic governed by external conditions attr. The system dynamics are defined as:
mode-switching control,
Here, P is the system state at the final time 7 starting from in- put x and evolving under control inputs 4,- Standard feedfor- ward neural networks correspond to Non-feedback control scenarios, meaning the control inputs are determined upfront and do not adapt based on intermediate states. In contrast, the MoE framework introduces adaptive decisions u, that se- lect transformations based on current hidden states h, is precisely a feedback (closed-loop) control mechanism, as illustrated in Figure 2(b) (middle). The corresponding dynamics are given by:
dhh = Uh,, " = K ssaarrr (h,),
(17)
where x denotes a family of control strategies and s(attr) is a switching function that selects the appropriate mode based on environmental attributes. Each mode K2 corresponds to a distinct sub-policy activated under specific conditions. To prevent overfitting and share knowledge learned from images of different distributions, we utilize a shared expert E, across different flags aftr. For each attribute, distinct G is trained, implying that for different attr, we adopt dif- ferent functions & in Eq. (15), analogous to using different matrices W in Eq. (4):
dh, = (,,))d.
(12)
Given the complexity of optimizing policy function through neural networks, kernel methods (Hofmann et al., 2008) are often utilized to approximate the control function u,
u,(h,) Σ K (h, (K).
(13)
where K is a kernel function, h are predefined anchor points, and (h) correspond fixed control parameters at these anchor points. This ce-parameterization transforms the optimization process from directly optimizing a poten- tially complex continuous function u, into optimizing a set of fixed parameters 02- = u/(h), significantly simplifying the learning task:
& Σ
Gha E,(h).
(18)
Considering that the neural network's 's layerwise structure represents a discrete form of the control leads to Eq. (2). Therefore, dMoE can be interpreted as attribute-wise mode- switching control variant of optimal control for fairness learning, as reflected in the structural resemblance illustrated in Figure 2(a) and (b) (bottom).
(h,,,(())) ss fh, Σ K K(h, K/00).
(14)
We consider the following kernel function
((p((())).((())))
K(h, ht)= Σ, (xp(((h))) ((h)))
(15)
4. Experimental Results
4.1. Datasets
To facilitate the learning of the kernel function df we in- troduce a simplified linear reparameterization. The inner product is computed by: (d(h)). (d())) (h, "h..b"'((k ((h()) = (hT W)p, where the transformed anchor features o= (h!))
To demonstrate the effectiveness of our proposed dMoE framework as an optimal control approach, we conduct extensive experiments on two benchmark datasets and an
5

---

Distribation-wware Fairness Learning in Medical Image Segmentation From A Control-Theoretie Perspective
in-house dataset. For each dataset, we visualize the data distribution for each attribute in Figure and provide further details of the training and test datasets in Appendix A.2. For the 2D segmentation experiments, we utilize two datasets: 1) Harvard-FairSeg (Tian et al., 2024) and 2) HAM10000 (Tschandl et al., 2018). For the 3D segmentation experi- ments with clinical attributes, we utilize our in-house radio- therapy target dataset for prostate cancer patients.
composed of 132 primary prostate cancer patients from Yon- gin Severance Hospital, Yongin, South Korea and 143 test samples from Gangnam Severance Hospital, Seoul, South Korea. The data collected for this study has been ethically approved by the IRB of the Department of Radiation Oncol- ogy at Yonsei Cancer Center, Yongin Severance Hospital, Gangnam Severance Hospital (IRB numbers 4-2023-0179, 9-2023-0161, and 3-2023-0396).
Harvard-FairSeg is a scanning laser ophthalmoscopy (SLO) fundus image dataset comprising 10,000 samples with pixel-wise optic cup and outer neuroretinal rim seg- mentation masks for diagnosing glaucoma. It includes six key demographic attributes- aage, gender, race, ethnicity, language preference, and marital status-enabling compre- hensive studies of fairness. In this study, we focus on the race attribute, i.c. race E [Black, Asian, White), as individu- als from minor attribute subgroups face a risk of developing glaucoma compared to other groups, yet the segmentation accuracy is often lowest for this demographic (Tian et al., 2024). Fairness and segmentation performance are evalu- ated on the test benchmark, which consists of 2,000 samples.
4.2. Implementation Details
For all experiments, we set the dMoE module hyperparame- ters with Top- as 2 and the number of experts n as 8. Each expert layer consists of a standard MI.P with two linear layers, a ReLU activation, and dropout layer. For 2D seg- mentation tasks, we use TransUNet (Chen et al., 2021) as the backbone with the standard ViT-B architecture. The net- work is trained following the setup used in previous studies. All the input images are center-cropped and resized into 2D patches of size 224 x 224 pixels with a batch size of 42, The network is trained with learning rate of 0.01 for 300 epochs the Harvard-Fairse
dataset and 100 epochs on HAM10000, following the benchmark setting. The best per- formance is selected from checkpoints saved at 100-epoch intervals. For 3D radiotherapy target segmentation task, we adopt the 3D Residual U-Net (Çiçek et al., 2016) as the backbone architecture, since reported as an effective architecture for radiotherapy target segmentation (Oh et al., 2024). The network is trained using randomly cropped 3D patches of size 384 x 384 128 voxels and batch size of 4. During evaluation, the entire 3D CT volumes are processed using a sliding window approach. The training is conducted with a learning rate of 5 x 10-5 over 100 epochs and early stopping based on the validation dataset. The additional computational costs introduced each backbone are com- pared in Appendix A.3. We implement the networks using PyTorch (Paszke et al., 2019) in Python with CUDA 11.8 The Adam (Loshchilov & Hutter, 2017) optimizer with exponential learning rate decay is used for all experiments. For 2D segmentation tasks, we use a single NVIDIA A100 80GB GPU, while for the 3D segmentation task, we use a single NVIDIA RTX A6000 48GB GPU.
HAM10000 is dermatology image dataset comprising 10,015 2D RGB samples with binary segmentation masks for diagnosing skin lesions. It includes demographic at- tributes such as sex and age, enabling targeted analysis of distributional disparities. In this study, we focus on the age attribute, as younger and older populations are underrepre- sented in the dataset. For age categorization, patients are divided into four groups at 20-year intervals, with the test benchmark consisting of 1,061 samples.
Radiotherapy Target Dataset comprises pelvic CT scans of prostate cancer patients, accompanied by clinical target volume (CTV) segmentation masks for radiotherapy plan- ning in radiation oncology. The dataset includes clinical factors, such as tumor staging and histopathological findings. In this study, we focus on the Tumor (T)-stage attribute, i.e. T-stage (TI, T2, T3, T4], as the dataset exhibits an imbal- ance in population distribution. For example, radiotherapy can be applied across all T-stages in prostate cancer, but its aim may vary depending on clinical practices. In certain regions, radiotherapy alone is used as definitive treatment for both early and advanced stages without surgery. In contrast, other regions favor radiotherapy alone for early stages, while combining surgery with adjuvant radiother- apy for advanced stages, reflecting institutional variations in treatment patterns. To address biases arising from the imbalanced T-stage distribution when training deep neural networks for radiotherapy target segmentation, we utilize a training dataset comprising 721 primary prostate cancer pa- tients from Yonsei Cancer Center, Scoul, South Korea, and validate network performance using an independent test set,
4.3. Bascline Method and Evaluation Metrics
We evaluate our method against baseline approaches for fairness learning. On the 2D Harvard- -airSee dataset, we compare our approach with reported metrics from adversar- ial (ADV) training (Madras et al., 2018), distributionally robust optimization (DRO) (Sagawa et al., 2019), fair-error- bound scaling (FEBS) (Tian et al., 2024), and generative model-based augmentation (FairDiff). We implement our method under the same training conditions as the benchmark methods. For further experiments on the 2D H0000000 dataset and the 3D radiotherapy target dataset, we compare
6

---

Distribation-aware Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
our method with FEBS among SOTA methods. We exclude the generative model-based approach, FairDiff, it requires additional image generation and evaluation tailored to 2D and 3D datasets, which makes it suboptimal for direct com- parison. Moreover, we use the default MoF as an additional baseline, replacing our proposed dMoE modules by default MoE modules without distribution-aware router. We imple- ment all methods under the same training conditions as our method.
performance for the Asian attribute subgroup, achieving a Dice score of 0.844 compared to FairDiff (0.832) and FEBS (0.825). Although MoE, which removes the distribution- aware gating network from our proposed method, achieves comparable performance (0.845), dMoE outperforms MoF across other subgroups for both rim and cup segmenta- tion. Furthermore, the shared structure within dMoe module across subgroups enables synergistic effects, maintaining or even improving performance for the major attribute sub- group (White). For the HAM10000 skin lesion segmenta- tion task, our method also demonstrates SOTA performance in ES-Dice and ES-IoU, with scores of 0.801 and 0.725, respectively, outperforming MoE (0.796 and 0.721) and FEBS (0.757 and 0.667), as shown in Table 2. Specifically, loss function modification methods, such as FEBS, exhibit inferior performance across attribute subgroups compared (19) to the baseline method. This suggests that when the train- ing dataset for each subgroup is insufficient, these methods may sacrifice overall performance. While MoE demon- strates promising performance for minor attribute subgroups, dMoE further improves fairness.
To evaluate performance, we employ the Dice Similarity Coefficient, Intersection over Union (IoU), as well as equity- scaled segmentation performance (ESSP) metrics for both Dice and SSU-denoted as ES-Dice and ES- Lo---oU--loiiinn (Tian et al., 2024):
ESSP I(1(
1+A A
A = Σ - atreE
-
altr (20)
where E (Dice, JoU] and A represents the set of all demo- graphic groups. ESSP is fundamentally aligned with both the demographic parity and equalized odds principles (Hardt et al., 2016). While ESSP evaluates performance fairness across diverse demographic subgroups, it has limitations in capturing substantial performance gains for specific sub- groups, particularly in the 3D radiotherapy target segmenta- tion task. We complement the limitation of ESSP metrics by providing quantitative analysis with violin plots, which of- fer more detailed visualization of equity and performance distribution across diverse subgroups. We also report worst- group accuracy metrics to ensure that the model's 's lowest performance is properly accounted for.
4.4.2. 3D RADIOTHERAPY TARGET SEGMENTATION WITH CLINICAL ATTRIBUTES
We further evaluate dMoE for 3D radiotherapy target seg- mentation, using T-stage as distributional attribute. Table 3 summarizes the segmentation performance. Despite the test set being acquired from a different hospital with a distinct data distribution, dMoE demonstrates promising general- ization performance. The strength of dMoE is particularly evident in underrepresented subgroups, such as Tl and T4 where the frequency is significantly lower than that of T2 and T3, by resulting the most notable performance improve- ments in these subgroups. Specifically, dMoE demonstrates robust and consistent performance, particularly excelling in the T4 subgroup with a Dice score of 0.778, significantly outperforming FEBS (0.685) and MoE (0.708). Moreover, despite the limited sample size of the TI subgroup, dMoE achieves superior performance. This consistent improve- ment demonstrates dMoE's generalizability and robustness, even when applied to datasets from different centers with varying data distributions.
To support statistical analysis for fairesss evaluation, we apply bootstrapped 95% confidence intervals (CIs) when computing both ESSP and all attribute-average metrics, re- sampling each subgroup and all groups with replacement over 1,000 iterations, respectively.
4.4. Results
1. 2D SEGMENTATION WITH DEMOGRAPLIIC ATTRIBUTES
However, as noted in Section 4.3, ESSP metrics have limita- tions in capturing significant performance gains for minor attribute subgroups. To address this, we further illustrate violin plots in Figure 3(a). Comparison on violin plots, with equity measures indicated by transparent blue lines, confirms dMoE's s ability to maintain equity across attribute- wise subgroups. Qualitative results for each subgroup are presented in Figure 3(b). As observed in the ground truth labels, the radiotherapy target volume tends to expand with the progression of the tumor's clinical stage. This trend is distinctly captured by dMoE, which demonstrates superior
We first evaluate dMoE for 2D segmentation on benchmark datasets. For the Harvard-FairSeg neuroretinal rim and optic nerve segmentation tasks, our method consistently achieves SOTA performance in terms of ES-Dice and ES-IoU, sur- passing previous faireess learning approaches, as shown in Table 1. Specifically, our method demonstrates supe- rior performance for the minor attribute subgroup (Black), achieving a Dice score of 0.776 for rim segmentation, com- pared to FairDiff (0.743) and FEBS (0.733). Similarly, for optic cup segmentation, our method significantly improves
7

---

Distribution-wware Fairness Learning in Medical Image Segmentation From A Control-Theoretie Perspective
Table 1. Comparison on 2D Harvard- Frirrrd dataset with race as the distribution attribute.
All (m=2000)
Asian (a=169)
Black (u=2919)
White (n=1532)
Method
ES-Dice (CIs)
Dice (Cls)
BS-ToU (CK))
ToUl (CIs)
Dice
IoU
Dice
IoL.
Dice
IoU
Rim Segmentation
TransUNet" (Chen 2021)
0.703
0,793
0.585
0.671
0.746
0.616
0.731
0.599
0.811
0.691
+ ADV" (Madua al., 2018)
0,700
0,791
0,583
0.668
0.741
0.612
0.729
0.598
0.809
0.689
+ DRO (Sagawa et al., 2019)
0,700
0.791
0.581
0.667
0.747
0.618
0.723
0.590
0.808
0.689
+ FEBS' (Tan al., 2024)
0.705
0.795
0.587
0.673
0.748
0.733
0.602
0.813
0.694
+ FairDif (L. et al., 2024)
0,716
0.800
0.596
0.680
0.757
0.628
0.743
0.611
0.817
0.699
+ MoB
0.738 (0.713-0.752)
0.804 (0. (0.799-0. 809)
0.614 (0.596-0 633)
0.685 (0 (0.680-0.00 691)
0,760
0.635
0.763
0.635
0.817
0.701
JMoE
0.743 (0.723-0.763)
0.813 (0.808 (0.188
0.627 (0.608-9.6455)
0.698 (0.692-0.704)
0.769
0.645
0.776
0.652
0.825
0.713
Cup Segmentation
TransUNet' (Chen et 2021)
0.753
0.827
0.728
0.849
0.850
0.755
ADV' (Madras al., 2018)
0,826
0.841
0.727
0.743
0.825
0.726
0.842
0.748
0.843
0.744
+ DRO' (Sagawa et al., 2019)
0,820
0,844
0.725
0,748
0.820
0.723
0.847
0.753
0.846
0,750
+ PEBS' (Tian al., 2024)
0.825
0.750
0.825
0.725
0.848
0.755
0.848
0.751
+ DairDin (Li al., 2024)
0,825
0.848
0.736
0.753
0.832
0.735
0.848
0.757
0.850
0.754
+ Moll
0.830 (0.809. 0.847)
0.854 (0.849 (1. (68)))
0.739 (0.720-0.754)
0.762 (0.755.0.766))
0.845
0.757
0.842
0.748
0.857
0.765
dlMaF
0.832 (0.810-0.853)
0.862 (0.856-0.867)
0.745 (0.722-0.765)
0.773 (0.76660..779
0.844
0,755
0.851
0.761
0.866
0.777
Metric reported from (Tan al., 2024). | metrics are rocalculated using Eq. (19), based on metrics reported the original paper (L. at. 2024),
fair comparison.
Table 2. Comparison on 2D HAM10000 dataset for skin lesion segmentation with age as the distribution attribute.
All (o-10
Age 80 (n-121)
Age 2: 60 (o-169)
Age 2 40 (B-328)
Age 20 (S-120)
Age 20
Method
(n-24)
ES-Dice
Dioz (Cls)
ES-bol (CBl)
Toll
1ol1
Dice
B:07
Dice
100
[o0)
Truns iChen 2021)
0.792 (0.757 (1841)
0.876
0.889)
0,714
0.766)
(0) 0.838)
0.787
0,868
0.588
0,895
0:857
0.839
FEDS (Tan al., 2024)
0.757 (0 0.8877
0.858 (0 845-8--88 0 872)
0.667 (0.613-0.7199
0,796 (0, (0,783- 0.812)
0.831
0.247
0.841
0,774
0.584
0.837
0.871
0.827
0 0,369
0.830
MaF.
0.796 (0 741-0 844)
0.882 (0 888-888 895)
0.721 (0.87) -0.770
0.833 (0.818-0 846)
0.864
0.794
0.875
0.820
0.589
0.851
0.904
0.809
0.882
0,850
dMMoF
0.801 (0.745-0U847)
0.884 (0.870-8.896)
0,725 (0.673-8.776)
0,834 (0.8221-0.847)
0,864
0.791
0.881
0.824
0.890
0,850
0.501
0.856
0.580
0.846
Table 3. Comparison on 3D radiotherapy target segmentation with tumor stage as the distrihution attribute.
(s=275)
Ti (m=11)
(m=129)
T3 (s=114)
T4 (m=21)
Mothod
ES-Diee (CIs)
Dice (CIs)
ES-1ot) (CIs)
loll (CIs)
Dice
1oU
Diee
Iell
Dioe
10l1
Dice
3D ResUNet (Çiçek et al., 2016)
0.487 (0.447-0,35 0.599)
0.610 (0.589-00)) 0.630)
0,367 (0.336 0.399)
0.162 (0.440-0.182)
0.193
0.341
0.569
0.420
0.659
0.511
0.656
0.505
FEBS (Tian ol., 2014) MoE
0.434 (0.406-0.467)
0.586 (..66700.044)
(0.302-0.346)
(0.414-D:452)
0.442
0.374
0.652
0.501
0.452 (0.415 0. (92)
0.608 (0.586 0.628)
0.342 (0.314
0.461 (0.439-0.8)) 0.442)
0.492
0.345
0,542
0.393
0.674
0.532
0.708
0.557
dMoE
0.499 (3.3)))0.531)
0,650 (0.628-0.671)
0.384 (0.358-0.410)
0.506 (0.483-0.528)
0,718
0.585
0.435
0.693
0.556
0.778
0.641
Note. The underlined value indicates the wonst group accuracy among distribution attiibttes for each method.
(a)
Prostate Cancer Radiotherapy Segmentation
Abrbbuaacoiee Trannat Dattouton
Daokire
rEnS
Mat:
ditc
244
02
23-
62
T1
12
13
T4
T1
T4
T4
-50ge
78
T9
78
T9
74
T4
74
T4
T4
Figure 3. (a) Violin plots depicts attribute- wise segmentation performance. Transparent blue lines within the plots connect the most densely populated regions for each attribute, visually representing overall equity. (b) Qualitative comparison across different subgroups.
8

---

Distribation-wwar Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
adaptability to distributioonsspecfii characteristics, thereby achieving notable faimess in performance.
ical decision-making that aligns with clinicians' perspec- tives. Notably, ensuring distributional equity is essential for real-world clinical AI deployment, where training and de- ployment dataset distributions often differ. Our distribution- aware dMoF. holds promise in adapting trained models to unknown distributions, thereby improving the success of clinical Al integration across diverse hospitals.
We also present detailed comparison between the dMoE attribute-wise gating mechanism and multiple independently trained networks in terms of computational complexity and performance, as shown in Table of Appendix A.3. The results demonstrate that dMoE not only achieves higher segmentation accuracy but also maintains greater computa- tional efficiency, outperforming specialized models with single unified network.
However, several limitations remain for this study:
First, while we explore diverse attributes across three datasets, the performance improvement trends vary depend- ing on the task and dataset characteristics. Moreover, the optimal configuration of the dMoE module varies between the two different architectures. Future research will focus on uncovering the underlying characteristics of sttribute-wise subgroups within each dataset to identify the factors driving these variations and will aim to develop a more generalized module capable of delivering task-agnostic performance improvements.
4.4.3. RADIOTHERAPY TARGET SEGMENTATION WITH DIVERSE CLINICAL ATTRIBUTES
In clinical practice, radiotherapy targets are determined not solely by tumor stage or visible tumor imaging but through the integration of anatomical imaging with clinical param- eters such T-stage, Gleason Grade Group (GG), and prostate-specific antigen (PSA) levels to account for po- tential microscopic spread (National Comprehensive Cancer Network, 2024). To evaluate the broader applicability of our method, we further tested additional clinical factors which reflect pathological tumor differentiation but do not indi- cate the precise tumor location or extent. As shown in Ap- pendix A.4, our results demonstrate faimess improvements comparable to those observed with T-stage. These find- ing show that incorporating clinical indicators- even when not directly related to visible anatomical structures can enhance segmentation accuracy across subpopulations.
Second, our focus on a single attribute per task limits the exploration of combined attribute imbalances. For instance, while patients with T2-stage cancer may be prevalent, fur- ther subgroup imbalances, such as age metastasis status, could exist. Addressing this issue may require integrating multiple attributes within a hierarchical dMoE framework, which we plan to explore in future studies.
Last, this study adopts an optimal control perspective to design efficient deep neural network architectures, leverag- ing back back- propagatiin for parameter optimization. Future research will explore how advanced numerical methods from optimal control theory can inspire novel optimization algorithms specifically tailored to enhance fairness learning.
4.4.4. ABLATION STUDIES
We further conduct additional analyses explore the ef- fects of various configurations of the dMoE module in Ap- pendix A.5. Specifically, we vary the placement of the dMoE module within the encoder and decoder and inves- tigate whether its parameters should be shared or sepa- rated across layers to determine the optimal architecture. We further ablate each component of the optimal control framework to assess the contribution of control theory to distribution faimess learning. Detailed results of the ablation study are provided in Table 11 and Table 12.
Despite these limitations, we believe our distribution-awan approach represents a step forward in advancing fairness learning for diverse fata-imbalanced clinical scenarios.
A Acknowledgements
This work was supported by the National Institutes of Health, USA ROIHL159183, and also supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (RS- (RS-2024-003458544)
5. Conclusion, Limitations & Future Work
Adopting an optimal control perspective, we design distribution-aware Mixture of Experts (dMoE) architectures to address data imbalance issues. By modeling distribution as external factor influencing control, we integrate dMoE into medical image segmentation for more effective han- dling of imbalanced clinical datasets. Given the nature of clinical practice, which accounts for distributional charac- teristics, our proposed algorithm offers a promising solution to data imbalance and provides robust framework for clin-
Impact Statement
This paper aims to advance faireess in medical image seg- mentation by addressing biases arising from imbalanced clinical data distributions. Our work demonstrates robust- ness across diverse datasets and network architectures, con- tributing to equitable and reliable AI-driven healthcare ap- plications.
9

---

Distribution-aware Fairness Learning in Medical Image Segmentation From A Thetrol-Theorctic Perspective
References
Jin, R., Xu, Z., Zhong. Y., Yao, Q. Dou, Q. Zhou, S. K., and Li, X. Fairmedfm: fairness benchmarking for medical imaging foundation models. arXiv preprint arXiv:2407.0098. 2024.
Aström, K. J. Adaptive control. In Mathematical Sys- tem Theory: The Influence of RE Kalman, pp. 437-450. Springer, 1995.
Åström, K. J. and Murray, R. Feedback systems: intro- duction for scientists and engineers. Princeton university press, 2021.
Ktena, I., Wiles, O., Albuquerque, I., Rebuffi, S.-A., Tanno, R., Roy, A. G., Azizi, S., Belgrave, D., Kohli, P., Cemgil, T., et al. Generative models improve fairness of medical classifiers under distribution shifts, Nature Medicine, pp.
Boskovic, J. D. and Mchra, R. K. Multi-mode switching in flight control. In 19th DASC. 19th Digital Avionics Systems Conference. Proceedings (Cat. No. 00CH37126), volume 2, pp. 6F2-1. IEEE, 2000.
1-8, 2024.
Langley, Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference Machine Learning (ICMI,2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Chen, J., Lu, Y., Yu, Q. Luo, X., Adeli, E., Wang. Y., Lu, L.. Yuille, A. L., and Zhou, Y. Transunet: Transformers make strong encoders for medical image segmentation. arXiv preprint arXiv:2102.0445
Li, W. Xu, H., Zhang, G., Gao, H.-a., Gao, M., Wang, M. and Zhao. H. Fairdiff: Fair segmentation with point- image diffusion. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 617-628. Springer, 2024.
2021.
Chen, Q. Zhu, L., He, H., Zhang, X., Zeng, S., Ren, Q., and Lu, Y. Low-rank mixture-ou mixture-of-expertt ex-erts for continual medical image segmentation. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 382-392. Springer, 2024.
Loshchilov, and Hutter, Decoupled weight decay regu- larization. arXiv preprint
2017.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Lu, Y., Zhong, A., Li, Q., and Dong, B. Beyond finite layer neural networks: Bridging deep architectures and numer- ical differential equations. In International Conference on Machine Learning, pp. 3276-3285. PMLR, 2018.
Chester, A., Koh, S., Wicker, J., Sun, Q., and Lee, Balancing utility and fairness against privacy in medical data. In 2020 IEEE Symposium Series on Computational Intelligence (SSCI). pp. RCCC448803.220.930332
Madras, D.. Creager, E., Pitassi, T., and Zemel. Learning adversarially fair and transferable representations. In International Conference on Machine Learning, pp. 3384- 3393. PMLR, 2018.
2020. doi: 10.1109)
Çięek, Ö., Abdulkadir, A., Lienkamp, S. S., Brox, T., and Ronneberger, O. 3D U-Net: learning dense volumetric segmentation from sparse annotation. In Medical Image Computing and Computer-Assisted Intervention, pp. 424 432. Springer, 2016.
National Comprehensive Cancer Network. NCCN Clinical Practice Guidelines in Oncology: Prostate Cancer (Version 4.2024), 2024.
https:/ hols::/wwwr neen www.ao.......ooososooo..oo
Doyle, J. C., Francis, B. A., and Tannenbaum, A. R. Feed- back control theory. Courier Corporation, 2013.
Oh, Y., Park, S., Byun, H. K., Cho, Y., Lee, J., Kim, J. S.,
and Ye, C. LIm-driven multimodal target volume con- touring radiation oncology. Nature Communications. 15(1):1186, 2024.
Hardt, M., Price, E.,, and Srebro, N. Equality of opportunity in supervised learning. In Proceedings of the 30th Inter- national Conference on Neural Information Processing Systems, NIPS' 16, pp. 3323-3331, Red Hook, NY, USA, 2016. Curran Associates Inc. ISBN 9781510838819,
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z.,, Gimelshein, N., Antiga, L., al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Hofmann, T., Schölkopf, B., and Smola, A. Kernel meth- ods in machine learning. 2008
Jiang, Y. and Shen, Y. M4oe: A foundation model for medical multimodal image segmentation with mixture of experts. In International Conference on Medical Image Computing and Computer-Assin Assistee Intervention, pp. 621- 631. Springer, 2024.
Rypešć, G., Cygert, S., Khan, V., Trzciński, T., Zieliński, B., and Twardowski, B. Divide and not forget: Ensemble of selectively trained experts in continual learning. arXiv preprint arXiv:2401.10/9/ 2024.
10

---

Distribation-wware Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case gener- alization. arXiv preprint
Yamaguchi, T., Shishida, K., Tohyama, S., and Hirai, H. Mode switching control design with initial value compen- sation and its application to head positioning control on magnetic disk drives. IEEE Transactions on Industrial Electronics, 43( 43(1):65-73, 1996.
2019.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture- -ffteereof-eepperts layer. arXiv preprint arXiv: 011.7700...5550 2017.
Yu, Zhuge, Y., Zhang, L.,, Hu, P., Wang, D., Lu, H., and He, Y. Boosting continual learning of vision-language models via mixture-of-exper adapters. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 23219-232336 2024.
Sun, Q. Dong, H., Chen, Z., Sun, J., Li, Z., and Dong, B. Layer-parallel training of residual networks with auxil- iary variable networks. Numerical Methods for Partial Differential Equations, 40(6)):23147,, 2024.
Yu, L., Huang. J., and Fei, S. Sliding mode switching control of manipulators based on disturbance observer. Circuits, Systems, and Signal Processing, 36:2574-2585. 2017.
Theodore, D. A., Branche, A. R., Zhang, L.., Graciaa, D. S., Choudhary, M., Hatlen, T. J., Osman, R., Babu, T. M., Robinson, S. T., Gilbert, P. B.. et al. Clinical and demo- graphic factors associated with covid-19, severe covid-19, and sars-cov-2 infection adults: secondary cross- protocol analysis of 4 randomized clinical trials. JAMA network open, 6(7)o2332339-<2222334 2023.
Zhang, X., Ou, N., Basaran, B. D., Visentin, M., Qiao, M., Gu, R., Ouyang, C., Liu, Y., Maithews, P. M., Ye, C., et al. A foundation model for brain lesion segmentation with mixture of modality experts. International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 379-389. Springer, 2024.
Tian, Y., Shi, M., Luo, Y., Kouhana, A., Elze, T., and Wang, M. Fairseg: A large-scale medical image seg- mentation dataset for fairness learning using segment anything model with fair error-bound scaling. In The Twelfth International Conference on Learning Represen- tations, 2024.
Tian, Y., Wen, C., Shi, M., Afzal, M. M., Huang, H., Khan, M. O., Luo, Y., Fang, Y., and Wang, M. Fairdomain: Achieving fairness in cross-domain medical image seg- mentation and classification. In European Conference Computer Vision, pp. 251-271. Springer, 2025.
Tschandl, P., Rosendahl, C., and Kittler, H. The ham 10000 dataset, a large collection of multi-source dermatoscopic images of common pigmented skin lesions. Scientific data, 5(1):1-9, 2018 2018.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention all you need. Advances in neural information processing systems, 30, 2017.
Wang, G., Ye, J.. Cheng, J., Li, T., Chen, Z., Cai, J., He, J., and Zhuang, B. :am-m: Towards forgetting segment anything model via mixture of experts for 3d medical image segmentation. In International Conference on Medical Image Computing and Computer- Assisted Intervention, pp. 552-561. Springer, 2024.
Weinan, E. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 155:1-11, 2017.

---

Distribution-aware Fairness Learning in Medical Image Segmentation From A Control-Theoreti Perspective
A. Appendix
A.1. Detailed Network Architecture.
We provide the network architecture for a better understanding of the dMoE modules location within each Transformer- and CNN-based architecture in Table 4 and Table 5, respectively.
Table dMoE within Transformer (TransUNet).
Data dimension (C H W)
Module
Layer Block
Resample
dMoB
Chia 224 x 224 1 x 14
In
Conv
AltentionBlock, AttentionBlockg
768 (14 14)
768 (14 x 14)
Encoder
dMoE
AttentionBlockj AttentionBlnck
768 x (14 x 14)
768 x (14 x 14)
UpResBlock,
Up
256 28 x 28
UpResBlock,
Up
128 56 x 56
Decoder
UpResBlock,
Up
x 112 x 112
UpResBlock,
Up
16 224 x 224
Out
Conv
Chas 224 224
Table 5. dMoE within 3D CNN (3D ResUNet).
Skip-
Data dimension (C XxxHxX W D)
Module
Layer Block
Resample
dMoB
Connection
Conv
Ch., x 384
128
ResBlock,
Down
dMaF,
48 192 x 192 x
ResBlockg
Down dMol:
48 96
Encoder
ResBlocks
Down dMoE,
7
96
ResBlock,
Down dMoE,
192
ResBlocks
Down
dMoF
384 12
UpResBlock,
Up
192 24 x 8
UpResBlocks
Up
96 x 48 x 48 x 16
Decoder
UpResBlockg
Up
48 x 96 x 96 x 32
UpResBlock,
Up
48 192 192 x 64
Out
TransposeConv
Up
Cha.. x 384 x 384 x 128
A.2. Dataset Details.
We further provide the trainset and testset for each dataset, along with the attribute subgroup-wise data distribution and percentiles for the trainset in Table 6.
Table 6. Detailed distribution of data across attribute subgroups.
Harvard-FairSeg
HAM10M0
Radiotherapy Target Dataset
Attribute (Race)
Attribute (Age)
Attribute (T-stage)
Dataset
Total
Total
Total
Asian
Black
White
2 80
2 60
40
20
<02
T1
T2
T3
T4
Trainset
7945
750
1174
6021
8137
191
1324
3693
2356
573
721
26
227
425
43
(%)
(100)
(9)
(15)
(76)
(100)
(2)
(16)
(45)
(31)
(7)
(100)
(4)
(31)
(59)
(6)
Testset
2000
169
299
1532
1061
121
469
328
120
24
275
11
129
114
21
12

---

Distribation-wware Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
A.3. Comparison on Computational Efficiency.
We compare the computational complexity of incorporating the dMoE module into the backbone in Table 7. Additionally, Table 8 presents comparison between the dMoE attribute-wise gating mechanism and multiple independently trained networks, in terms of both computational complexity and performance, focusing the 3D radiotherapy target segmentation task using tumor stage as the distribution attribute. The results show that dMoE achieves superior performance while maintaining greater computational efficiency, outperforming individually specialized networks with single unified model.
TransUNett
+MoE
+dMoE
3D ResUNet
+MoE
+dMoE
All
Ti
72
T3
T4
Method
Gillpee
(n=2275)
[n=11]
(S5129)
(nc)
(e:21)
Input
224
384
D)
ES-DicetD)
Dise
Dioe
Diss
Dase
Dise
GPlops
45.84
90.28
90.28
1542.36 13.28
1761.30
1761.30
dMol: (Uues)
1760.30
0.499
0.718
0.585
0,93
0.778
Params (M)
91.67
129.46
129.51
Malippic udtworks for carbattriute
$729:44
0.457
0.606
0.999
0.515
0.681
0.760
Table 7. Computational complexity comparison.
Table 8. Comparison to multiple networks for each attribute.
A.4. Further Clinical Attribute Analysis on Radiotherapy Target Segmentation.
To further evaluate the ability of the proposed dMoF. module to capture the diversity of real-world clinical settings, we incorporate an additional clinical parameter as a distribution attribute in the prostate cancer study. Specifically, we use the Gleason Grade Group (GG) and prostate-specife antigen (PSA) level to represent pathological differentiation and disease progression, respectively. As shown in Table 9, our method demonstrates robust performance across various subgroups, particularly in underrepresented groups such as GG 6, 9, and 10. Similarly, Table 10 shows that the proposed dMoE module enhances equity across PSA subgroups, with notable improvements in underrepresented groups such as PSA levels and 4.
Table 9. Comparison 3D radiotherapy target segmentation with Gleason Grade Groups (GG) the distribution attribute.
All (n=275)
GG level 6 (n=31)
GG 7 (n= 125)
GG 8 (n=62)
GG 9 (n=47)
GG 10 (0)
Method
ES. Dice
Dice
FS-InU
loU]
Dice
IoU
Dice
IoU]
Dice
IoU
Dice
loU
Dice
loU
3D ResUNet (Çiçek et al., 2016)
0.512
0.610
0.389
0.461
0,562
0.417
0.578
0.429
0.650
0,501
0.669
0.517
0.623
0.474
+ FEBS (Tian et al., 2024)
0.451
0.593
0.337
0.441
0,501
0.349
0.557
0,406
0.628
0.473
0.686
0.534
0.650
0.494
MoF.
0.447
0.514
0.366
0.565
0.419
0.653
0.505
0.704
0.689
0.536
dMoE
0.473
0.638
0.361
0,494
0.672
0.533
0.566
0.419
0.657
0.511
0.750
0.614
0.750
0.612
Table 10. Comparison on 3D radiotherapy target sogmentation with prostate-specific antigen (PSA) level as the distribution attribute.
All (n=275)
PSA level (n=51)
level (n=84)
level (n=60)
level (n=29)
level (n=52)
Method
Dice
Dice
ES loU
loU
IoU
Dice
loU
loU
Dice
IoUl
Dice
loU
3D Res (Çiçek et al., 2016)
0.504
0,608
0.380
0.460
0.552
0.402
0.578
0.429
0.643
0.498
0.643
0.495
0.659
0.508
FEBS (Tian et al., 2024)
0.561
0.588
0.412
0.434
0.571
0.422
0.584
0.429
0.609
0.455
0.592
0.441
0.585
0.427
+ MoF:
0.461
0.606
0.348
0.459
0.532
0,383
0,553
0.407
0.632
0.486
0.667
0.519
0.709
0.563
+ dMoE
0.575
0.654
0.442
0,510
0.664
0,530
0.612
0.465
0.643
0,494
0.660
0.513
0.723
0.582
A.5. Ablation Study Results.
We perform ablation studies to examine the effects of various configurations of the dMoE module. Specifically, we investigate performance changes based on 1) dMoE Location within the network, by inserting the dMoE module at different layers within the encoder, the decoder, or both for comparison. We also examine whether 2) dMoE Parameters should be shared or separated across layers to determine the optimal architecture. In the case of parameter sharing for CNNs, we incorporate additional linear layers before and after the dMoE module to match the channel dimensions across different layer blocks. We further analyze the contribution of 3) Optimal Control Components by defining the mode-switching feedback control as the dMoE module. To assess its impact, we ablate both the mode-switching and feedback control components. For performing experiments, we use the HAM10000 dataset for Transformer- based architecture UNet), while for CNN-based architecture (3D ResUNet), we utilize a radiotherapy target segmentation task, by reporting performance using the ES-Dice and Dice score as the metric.
13

---

Distribation-aware Fairness Learning in Medical Image Segmentation From A Control-Theorctic Perspective
As shown in Table 11 and Table 12, inserting the dMoE module within the encoder, decoder, both yields comparable performance for both architectures, particularly for minority groups (The subgroup with Age 20 for HAM10000 dataset and the T4-stage subgroup for radiotherapy target dataset). Therefore, retain the dMoE module within the encoder only to optimize network training efficiency. In the Transformer-basee architecture, sharing dMoE parameters across layers enhances performance due to consistent dimensionality between layers. Conversely, in CNN-based architecture, differences in layer-wise dimensionality reduce the effectiveness of parameter sharing. This necessitates adding dimension dimension-matching linear layers before and after the dMoE modules; however, performance remains comparable. As a result, we adopt a layer-wise dMoE configuration specifically tailored for CNNs. In the analysis of optimal control components, naive performance degradation depending the dataset.
adaptation of feedback control resulted in marginal improvements In contrast, our proposed mode-switching control consistently improved both ES-Dice and Dice scores across experiments.
Table 11. Ablation study on HAM10000 datsaset for Transformer-base architecture.
Age 80
Age (n=169)
Age 2 (n=328)
Age 2 (n=120)
Age (n=24)
Method
(n= 1061)
(n=121)
ES-Dice
Dice
Dice
Dice
Dice
Encoder
Decoder
1) dMoE Location
((urs)
0.841
0.884
0.864
0.881
0.890
0.901
0.880
0.840
0.881
0.871
0.872
0.892
0.899
0.879
0.842
0.885
0.873
0.877
0.893
0.905
0.882
Layer- Wise
Shared
2) dMoE Parameters
0.812
0.881
0.870
0.875
0.893
0.890
0.834
((uus)
0.841
0.884
0.864
0.881
0,890
0.901
0.880
Control
Mode-switching
3) Optimal Control Components
Feedback
/ (Ours)
0.841
0,884
0.864
0.881
0.890
0.901
0,880
Feedback
0.836
0.882
0.864
0.875
0.889
0.904
0.882
Non- feedback
0.826
0.882
0.863
0.871
0.897
0.906
0.882
Table 12. Ablation study on radiotherapy target segmentation for CNN-based architecture.
T3 (n=114)
Method
(n=275)
(n=1
(n=129)
(n=21)
ES-Dice
Dice
Dice
Dice
Dice
Dice
Encoder
Decoder
dMoE Location
((Ous)
0.546
0.711
0.828
0.634
0.765
0.767
0.569
0.704
0.798
0,644
0.747
0.746
0.543
0.701
0.840
0.635
0.747
0.741
Layer-Wise
Shared
2) dMoE Parameters
/(Ours)
0.546
0.711
0.828
0.634
0.765
0.767
/
0.546
0.710
0.806
0.626
0.769
0.769
Control
Mode-switching
Optimal Control Components
Feedback
(Ours)
0.499
0.650
0.718
0.585
0.693
0.778
Feedback
0.451
0.608
0.492
0.542
0.674
0.708
Non-feedback
0.509
0.615
0.524
0.573
0.668
0.637
14


Reviews:
[{"review_id": "ojSQSV8vJw", "reviewer": "Reviewer_Fj34", "summary": "The paper proposes a distribution-aware image segmentation framework inspired by the control theory in mode switching and closed loop control. The framework incorporates the mixture of expert to address the heterogeneous distributions in medical images. Experiments in two 2D image benchmarks and a 3D in-house dataset shows superior performance in mitigating the bias.\nThanks for the detailed response. The topic is promising, and the interpretation of the method is sound. The rebuttal is complete, and most concerns are well addressed. Thus, I raise the score to weak accept given the update after rebuttal.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**R1-W1.1** How is dMoE more distribution aware compared to a normal MoE?\n- Thanks for the reviewer’s insightful point. The improved \"distribution-awareness\" of dMoE stems from its subgroup-aware gating, whereas a normal MoE employs per-sample gating. This design enables each gating to better capture shared patterns within subgroups, enabling subdistribution-aware gating.\n\n>**R1-W1.2/Q1** In what scenarios would a user know the tumor severity in advance and only need segmentation? Typically, when users know the severity of the tumor, the user already knows the tumor area. I may change my mind if reasonable scenarios are described.\n- Thanks for the reviewer's thoughtful comment, and we appreciate the chance to clarify the clinical rationale behind our study design. While it may seem intuitive that knowing tumor severity implies knowledge of tumor location and radiotherapy (RT) target, this assumption does not fully align with standard RT practice. In clinical settings, RT target delineation occurs after initial tumor diagnosis and is not determined solely by tumor severity nor visible tumor imaging. Instead, it integrates anatomical imaging and clinical parameters, such as T stage and Gleason Grade Group (GG), to account for potential microscopic spread. \n- For instance, in prostate cancer, for early stage (T1–T2), even when tumors appear localized on imaging, the entire prostate gland is typically included due to potential microscopic disease beyond visible boundaries. For advanced stage (T3–T4), RT volumes expand further to cover extracapsular extension or adjacent organ invasion such as bladder or rectum, with possible elective nodal irradiation depending on clinical risk.\n- Additionally, clinical factors such as GG represent pathological tumor differentiation but do not specify exact tumor location or extent. Our experiment by subgrouping with GG (Resp_Tables 2 in response to R2-W1) demonstrated fairness improvements, similar to T stage. This emphasizes that integrating clinical indicators—even if not directly related to visible structures—enhances segmentation accuracy by identifying shared visual features within severity groups.\n\n>**R1-S1** Highlighting the framework's computational efficiency compared to training multiple networks for each attribute would be beneficial.\n-  Due to the character limit, we kindly refer the reviewer to Resp_Table 7 in response to R3-W5, which demonstrates dMoE's efficiency.\n\n>**R1-W2** Why incorporate W and W_noise?\n- W_noise further imposes a controlled level of randomness into the expert selection process into the gating mechanism, for avoiding convergence to a few dominant experts.\n\n>**R1-W3** The description for control theory is confusing.\n-  Please see inline-answers bellow.\n\n> [line129-131, right] shape of the h and \\tilde{h} are same? [line 145, right] why does the patched embedding, after going through gating*expert network, become unpatched?\n- Please refer to R3-W1.1.\n\n> [line 156] Normal() is a weird notation. \n- It should be updated to N(0,1), the standard normal distribution.\n\n> [E.q 8] What is f, \\theta [line 208, right] the shape of \\theta? [E.q 9, 12] What is u? t? \n- Please refer to R3-W1.2 and R3-W3.\n\n> [line 205, right] What is u_i^i(h_t^i)? \n- It should be u_t(h_t^i), where h_t^i​ is the i-th anchor point, and u_t(h_t^i) is the value taken at h_t^i​.\n\n> [line 240] What is the system parameter? manual?\n- Instead of manually designing a system, we interpret the neural networks (NN)'s operations on hidden states as a dynamical system. That said, these NN parameters serve as the system parameters, governing hidden state dynamics in accordance with dynamical systems theory.\n\n>**R1-W4** The experiments cannot justify the additional benefit in considering both demographic and contextual factors.\n- dMoE is designed to adapt to each dataset by effectively incorporating attributes that influence performance and distribution, unlike traditional methods that are working on specific datasets. However, as noted in Section 5, we plan to integrate both demographic and contextual factors in future studies.\n\n>**R1-W5** dMoE does not outperform MoE in all subgroups. \n- As discussed in R1-W1.1, MoE, as a distribution-aware mechanism, can outperform in major subgroups. However, the performance gains dMoE provides for minor groups are particularly valuable, and the strength of dMoE lies in its ability to promote balanced performance gains.\n\n>**R1-W6** Interpreting dMoE via optimal control doesn’t bring additional insights.\n- Interpreting NNs through the lens of dynamic processes (Section 2.3) is an active research area with both theoretical and practical benefits: (1) It enhances understanding of the mechanisms behind NN, such as why a MoE outperforms a fixed experts-based NN. (2) It enables the transfer of well-established concepts from control theory to NNs, including architectural structures, optimization, and regularization techniques such as mode-switching mechanisms.", "claims_and_evidence": "1.\tThe paper claims the framework to be distribution aware. But section 3.2 doesn’t highlight how is the proposed dMoE more distribution aware compared to a normal MoE: a normal MoE also has an input dependent gate [equation 1] in choosing the experts and thus also “distribution-aware”. The “Distribution-wise router” for each attribute seem to leverage fine-grained attribute annotations of the dataset, which means the user knows in advance which group the input belongs to. This reduces the method contribution of “distribution-aware” in the network as users are required to be aware of it first.\n2.\tThere is no explanation on the design expressed in equation 4: why incorporate two different learnable matrices W and W_noise?\n3.\tThe paper claims to get inspiration from the optimal control theory in the framework design. But the description related to the control theory is confusing and miss explanations on many used notations, making the contribution of the control theory to the proposed framework hard to follow.", "methods_and_evaluation_criteria": "The paper criticizes that “current fairness learning approaches primarily focus on explicit factors such as demographic attributes but neglect implicit/contextual factors such as disease progression patterns or severity” [line 47-50, right column]. However, experiments conducted in dataset Harvard-Fairseg only considers race (a demopraphic attribute), HAM10000 only considers age (a demographic attribute), and Radiotherapy Target Dataset only considers tumor stage attribute (a clinical factor indicating severity). Thus the experiments cannot justify additional benefit of the proposed approach in considering both demographic and implicit/contextual factors.", "theoretical_claims": "Many notations are not explained in the first time they are used and some notations are never explained. The confusing description makes it hard to build the connection between the control theory and the proposed framework.\n[line129-131, right column]: shape of the h and \\tilde{h} are the same?\n[line 145, right column] why does the patched embedding, after going through gating network*expert network, suddenly become unpatched? \n[line 156] Normal() is a weird notation. Is this a real value distributed in gaussian distribution?\n[equation 8] what is \\theta? What is f?\n[equation 9, 12] what is u? t?\n[line 205, right column] what is u_i^i(h_t^i) ?\n[line 208, right column] what is the shape of \\theta? \n[line 240] what is the system parameter? Manual setting?", "experimental_designs_or_analyses": "I checked experiments in section 4.4. In table 2 and table 3, dMoE do not outperform MoE in all subgroups. Since the manually chosen attribute wise router in dMoE is actually overfitting to each subgroup’s images, this result is a bit surprising. Further discussions on these results would be valuable.", "relation_to_broader_scientific_literature": "The paper is broadly related to fairness-learning [1] where balancing strategies are reviewed and control theory [2] where mode switching is used based on the control signal.\n[1] Xu, Zikang, et al. \"Addressing fairness issues in deep learning-based medical image analysis: a systematic review.\" npj Digital Medicine 7.1 (2024): 286.\n[2] Yamaguchi, T., Shishida, K., Tohyama, S., and Hirai, H. Mode switching control design with initial value compensation and its application to head positioning control on magnetic disk drives. IEEE Transactions on Industrial Electronics, 43(1):65–73, 1996.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths: It could be potentially a good application approach when attribute annotations are available as a unified framework handling different types of biases individually.\nWeakness: The section on interpreting dMoE through optimal control doesn’t seem to bring additional insights.", "other_comments_or_suggestions": "Maybe motivating the framework to be more computationally efficient than training multiple networks in each individual attribute would be beneficial for the writing.", "questions_for_authors": "1.In what scenario will a user know the severity of a tumor in advance and only wants to segment the tumor area? Typically, when users know the severity of the tumor, the user already knows the tumor area. This setting seems not very likely to happen in real world. I may change my mind if reasonable scenarios are described.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "NA."}}, {"review_id": "APwbltOa4q", "reviewer": "Reviewer_qTb4", "summary": "The paper \"Distribution-aware Fairness Learning in Medical Image Segmentation From A Control-Theoretic Perspective\" explores the issue of fairness in medical image segmentation, particularly in cases where demographic and clinical factors contribute to biased model performance. The authors argue that biases in deep learning models arise due to the inherent imbalance in clinical data acquisition, often skewed along dimensions such as age, sex, race, and disease severity. As discussed in the introduction, existing fairness-aware training strategies mainly focus on demographic attributes while overlooking clinical factors that influence medical decision-making. This paper proposes a novel approach, termed Distribution-aware Mixture of Experts (dMoE), which adapts deep learning models to heterogeneous distributions in medical imaging.\n\nThe study builds upon the Mixture of Experts (MoE) framework. The authors reinterpret MoE as a feedback control mechanism (Section 3.3), where distributional attributes are incorporated into the gating function. This enables the model to adaptively select experts based on demographic and clinical contexts, thereby improving fairness in segmentation tasks. Unlike previous fairness-learning methods, dMoE integrates distributional awareness at the architectural level rather than as a post-hoc correction. As shown in Equations (1)–(5), the dMoE gating mechanism operates using a mode-switching control paradigm, dynamically selecting the optimal expert networks based on subgroup attributes.\n\nThe authors validate their approach through experiments on three medical imaging datasets: Harvard-FairSeg for ophthalmology, HAM10000 for skin lesion segmentation, and a 3D radiotherapy target dataset for prostate cancer segmentation. The results, summarized in Tables 1–3, show that dMoE outperforms baseline methods in terms of fairness and segmentation accuracy across underrepresented subgroups.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**R2-W1** Limited external validation and dataset diversity, which may not fully capture the diversity of real-world clinical settings. The prostate cancer test set in particular is relatively small, with only 132 test samples.\n- Thanks for your constructive feedback. For the prostate cancer test set, we collected an additional 143 test samples from a different hospital, more than doubling the total test sample size. These samples were scanned using a different CT manufacturer (SIEMENS) than the training data (Canon). As shown in Resp_Table 1, dMoE demonstrated the most promising and robust fairness performance on the expanded dataset.\n\n***Resp_Table 1.*** Radiotherapy target segmentation with **tumor stage** on the expanded testset.\n|Methods|All (n=275)|T1 (n=11)|T2 (n=129)|T3 (n=114)|T3 (n=21)|\n|-|-|-|-|-|-|\n||ES-Dice(D)/D|D|D|D|D|D|\n|RedUNet|0.487/0.610|0.493|0.569|0.659|0.656|\n|+ FEBS|0.432/0.590|0.442|0.528|0.652|0.685|\n|+ MoE\t|0.451/0.608|0.492|0.542|0.674|0.708|\n|+ dMoE|**0.499/0.650**|**0.718**|**0.585**|**0.693**|**0.778**|\n\n- Moreover, our used three datasets do not fully capture the diversity of real-world clinical settings. However, datasets containing fairness-related attributes for segmentation remain scarce [1]. Therefore, for the existing dataset, we incorporated another clinical parameter. For prostate cancer, we used Gleason Grade Group (GG), which reflects pathological differentiation and impacts both patient distribution and radiotherapy target patterns. As shown in Resp_Table 2, our method demonstrated robust performance across different subgroups, particularly in underrepresented subgroups such as GG 6, 9, and 10.\n- **Reference**: [1] Yu Tian et al., ICLR 2024, https://arxiv.org/abs/2311.02189\n\n***Resp_Table 2.*** Radiotherapy target segmentation with **Gleason Grade Groups (GG)**.\n|Methods|All (n=275)|GG 6 (n=31)|GG 7 (n=125)|GG 8 (n=62)|GG 9 (n=47)|GG 10 (n=10)|\n|-|-|-|-|-|-|-|\n||ES-Dice(D)/D|D|D|D|D|D|\n|RedUNet|**0.512**/0.610|0.562|**0.578**|0.650|0.669|0.623|\n|+ FEBS|0.451/0.593|0.501|0.557|0.628|0.686|0.650|\n|+ MoE|0.447/0.608|0.514|0.565|0.653|0.704|0.689|\n|+ dMoE|0.473/**0.638**|**0.672**|0.566|**0.657**|**0.750**|**0.750**|\n\n- Additionally, we incorporated the gender attribute for the other datasets in Resp_Tables 3 and 4. However, due to the relatively balanced distribution and the absence of established evidence indicating an effect of gender attribute on segmentation patterns, the performance gains of dMoE were less pronounced. Nevertheless, dMoE outperformed other methods for Harvard-FairSeg while showing comparable performance to MoE for HAM10000.\n\n***Resp_Table 3.*** Harvard-FairSeg dataset with **gender**.\n|Methods|All (n=2000)|Female (n=1229)|Male (n=771)| \n|-|-|-|-|\n||ES-Dice(D)/D|D|D|\n|TransUNet|0.844/0.848|0.851|0.846|\n|+ FEBS|0.846/0.849|0.851|0.0.849|\n|+ MoE|0.845/0.854|0.850|**0.860**|\n|+ dMoE|**0.856/0.858**|**0.857**|0.859|\n\n***Resp_Table 4.*** HAM10000 dataset with **gender**.\n|Methods|All (n=1061)|Female (n=496)|Male (n=566)|\n|-|-|-|-|\n||ES-Dice(D)/D|D|D|\n|TransUNet|0.862/0.879|0.890|0.869|\n|+ FEBS|0.846/0.860|0.869|0.853|\n|+ MoE|**0.880**/0.882|0.881|**0.882**|\n|+ dMoE|0.871/**0.883**|**0.891**|0.877|\n\n>**R2-W2** The lack of statistical significance in the fairness evaluations.\n- To address the lack of statistical analysis in fairness evaluation, we have adapted Bootstrapping Confidence Intervals (CIs) when calculating ESSP metrics, by resampling each subgroup sample with replacement for 1,000 iterations. We will update all the metrics with the 95% CIs as exemplified in Resp_Table 5.\n\n***Resp_Table 5.*** Statistical significance analyzed Dice metric and an additional  Worst-group accuracy metric for radiotherapy target segmentation.\n|Metric|ES-Dice (CIs)|Dice (CIs)|Worst-group Accuracy|\n|-|-|-|-|\n|RedUNet|0.487 (0.447-0.529)|0.610 (0.589-0.630)|0.493|\n|+ FEBS|0.434 (0.406-0.467)|0.586 (0.567-0.604)|0.438|\n|+ MoE|0.452 (0.415-0.492)|0.608 (0.586-0.628)|0.492|\n|+ dMoE|**0.499 (0.469-0.531)**|**0.650 (0.628-0.671)**|**0.585**|\n\n>**R2-W3** Could you clarify how ESSP relates to common fairness metrics, such as demographic parity, equalized odds, and worst-group accuracy?\n- Thanks for your informative comment. We will explain the ESSP metric in detail in the appendix of the final version. (1) Demographic parity ensures that the probability of a positive outcome is the same for all demographic groups, and (2) Equalized odds ensures false positive and false negative rates to be equal across demographic groups. ESSP aligns with both (1,2) principles by maintaining comprehensive segmentation performance, such as Dice or IoU, across different groups. Whereas, (3) Worst-group accuracy ensures the model's lowest performance is still adequately addressed. ESSP emphasizes equity across worst- to best-group performance, which does not fully align with worst-group accuracy. Therefore, we will include Worst-group accuracy in the main paper, as exemplified in Resp_Table 5.", "claims_and_evidence": "The claims made in the submission are generally well-supported. The claim that dMoE improves fairness in medical image segmentation is strongly supported by quantitative results in Tables 1–3, which compare its performance against existing fairness-learning approaches. The experiments consistently demonstrate that dMoE achieves state-of-the-art performance, particularly for underrepresented subgroups such as Black patients in ophthalmology and older patients in dermatology. Furthermore, Figure 3 illustrates how dMoE achieves a more balanced segmentation performance across demographic and clinical attributes. The study provides some evidence for generalization in the 3D segmentation experiment, where the test set is sourced from a different hospital than the training data. However, this dataset remains relatively small, with only 132 test cases. To substantiate this claim more convincingly, further external validation across multiple independent datasets would be necessary.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in the paper are generally well-aligned with the problem of fairness in medical image segmentation. The choice of using a distribution-aware mixture of experts (dMoE) as an architectural modification is sensible for this problem. By integrating demographic and clinical attributes into the expert selection process, dMoE offers a more structured approach to fairness that dynamically adjusts to subgroup-specific biases. The fairness evaluation metric, equity-scaled segmentation performance (ESSP), is a suitable to qunatify performance.", "theoretical_claims": "There are no proofs in this paper.", "experimental_designs_or_analyses": "The experimental design and analyses in the paper appear to be overall sound. The experimental design includes tests on three distinct medical imaging datasets. This selection allows for an evaluation of fairness across both demographic attributes (race, age) and clinical attributes (tumor stage). The inclusion of a 3D dataset is a strength of the study. However, the sample sizes of the test datasets are relatively small, particularly for the prostate cancer segmentation task (n=132 test cases).", "relation_to_broader_scientific_literature": "To the best of my knowledge, most previous applications of MoE in medical imaging have focused on multimodal learning (Jiang & Shen, 2024) and heterogeneous scanning modalities (Zhang et al., 2024), rather than fairness. The novelty of dMoE lies in integrating fairness as a gating criterion, which has not been explicitly explored in prior MoE applications. However, I cannot be more specific, since I am not very familiar with the literature pertaining to this work.", "essential_references_not_discussed": "I am not aware of missing references.", "other_strengths_and_weaknesses": "To the best of my knowledge, the integration of MoE with fairness-aware medical image segmentation appears to be novel. Unlike traditional MoE models, which select experts based solely on feature space partitioning, the proposed algorithm introduces attribute-aware gating functions (Equation 3) that enable the network to adjust its expert selection dynamically based on fairness-sensitive factors.\n\nOne of the main weaknesses is the limited external validation and dataset diversity. While the paper evaluates dMoE on three medical imaging datasets, these datasets may not fully capture the diversity of real-world clinical settings. The prostate cancer test dataset in particular is relatively small, with only 132 test samples. Another limitation is the lack of statistical significance testing in the fairness evaluations.", "other_comments_or_suggestions": "No further comments.", "questions_for_authors": "Could you clarify how ESSP relates to common fairness metrics in Machine Learning, such as demographic parity, equalized odds, and worst-group accuracy?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xH2m7n9Ixk", "reviewer": "Reviewer_V3Mq", "summary": "This paper proposes a novel method to address the fairness issue in medical image segmentation by incorporating control theory to handle data distribution disparities. By introducing distribution-aware fairness learning, the method is able to reduce unfairness among different groups while maintaining model performance. Fairness is particularly crucial in the medical field, where the representation of different racial,\ngender, and other groups is vital. Therefore, the proposed approach holds significant research value.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**R3-W1.1** The Patchify (Section 3.2) fails to clearly explain the implementation.\n- Thanks for your careful comments. Patchify flattens the 4D intermediate image embeddings ‘h’ in [H, W, Z, Ch] from CNN blocks to 2D flattened embeddings ‘\\tilde{h}’ in [N, Ch], where N corresponds to H x W x Z, ensuring compatibility of the embeddings with the transformer-based dMoE module. After ‘tilde{h}’ going through the gating x expert network, it needs to be reverted to 4D shape to be further computed with consecutive CNN blocks. This is only necessary for CNN-based architecture, as highlighted in line 134-right, however, we will clarify it in Section 3.2 as well as Figure 2.\n\n>**R3-W1.2** There are also formatting and notation errors (such as incorrect subscripts for 'h' in the formula on page 5).\n- The potential ambiguity of the subscript 'h' might stem from the distinction between 'l' and 't'. In this context, we follow the work on NeuralODE (Section 2.3) which establishes a connection between discrete neural network (NN) structures and continuous dynamic processes. Therefore, in the NN structures and parameters, 'l' is used as the subscript. In discussions of the dynamic process, it serves as the subscript. Moreover, 'l' is treated as a discretization of 't', as seen in E.q 10 and 11. We will provide a more detailed derivation and clarify notations to eliminate any potential misunderstandings.\n\n>**R3-W2.1** Lack of detailed explanation of the Softplus. \n- We will update the explanation: “Softplus is a smooth alternative activation function to ReLU, which is defined as SoftPlus(x) = log(1+e^x)”.\n\n>**R3-W2.2** Section 3.2 seems to be a minor modification of the MOE, lacking sufficient innovation or explanation.\n- Our innovation, grounded in an in-depth analysis of MoE, bridges MoE and optimal control theory and leverages well-established mode-switching control theory for fairness learning. While our modification is simple, it is insightful (based on above analysis) and significant, which addresses a critical unmet need in clinical AI: mitigating performance degradation caused by unbalanced data distribution.\n\n>**R3-W3**  The derivation of the optimal control theory in Section 3.3 appears abrupt, with insufficient transitions and explanation. The authors are encouraged to provide a more complete mathematical derivation, especially regarding how control theory is integrated with the model, to enhance clarity and persuasiveness. \n- We appreciate your constructive point and would like to clarify that our formulation aligns with related work that interprets neural networks through the lens of dynamic processes (Section 2.3). However, we will clarify some points that might have caused ambiguity and add preliminary background and detailed derivations in the final version: (1) The distinction between 'l' and 't' has been explained in R3-W1.2. (2) In our model, 'u' represents the control input: In non-feedback control, it corresponds to NN parameters (\\theta), while in feedback control, 'u' becomes a function of 'h', forming an ensemble of experts dependent on 'h' through a kernel method.\n\n>**R3-W4** Conduct additional experiments on the same dataset using different attribute distributions to further validate the method's effectiveness. Furthermore, expanding the dataset for more comprehensive experiments would be beneficial.\n- Due to the character limit, we kindly refer the reviewer to our responses in Resp_Tables 2-4 for additional experimental results using different attributes and Resp_Table 1 for expanded test dataset, both provided in response to R2-W1 of Reviewer qTb4.\n\n>**R3-W5** The ablation study section is underdeveloped. The authors should conduct separate ablation experiments to validate the effectiveness of the optimal control theory component and what extent the proposed innovations contribute.\n- We performed an ablation study on Optimal Control components for radiotherapy target segmentation, as shown in Resp_Table 6.\n\n***Resp_Table 6.*** Ablation study on Optimal Control components.\n|Methods|Optimal Control|All (n=275)|T1 (n=11)|T2 (n=129)|T3 (n=114)|T3 (n=21)|\n|-|-|-|-|-|-|-|\n|||ES-Dice(D)/D|D|D|D|D|D|\n|dMoE (Ours)|Mode-switching Feedback|0.499/**0.650**|**0.718**|**0.585**|**0.693**|**0.778**|\n|(a)|Feedback|0.451/0.608|0.492|0.542|0.674|0.708|\n|(b)|Non-feedback|**0.509**/0.615|0.524|0.573|0.668|0.637|\n\n- To further evaluate our innovation, we compare dMoE’s attribute-wise gating mechanism with multiple networks trained separately for each attribute. As shown in Resp_Table 7, dMoE demonstrates superior performance and computational efficiency.\n\n***Resp_Table 7.*** Comparison to multiple networks for each individual attribute.\n|Methods||All (n=275)|T1 (n=11)|T2 (n=129)|T3 (n=114)|T3 (n=21)|\n|-|-|-|-|-|-|-|\n||GFlops↓|ES-Dice(D)/D|D|D|D|D|D|\n|dMoE (Ours)|**1761.30**|**0.499**/**0.650**|**0.718**|**0.585**|**0.693**|**0.778**|\n|Multiple networks for each attribute|5729.44|0.457/0.606|0.599|0.515|0.681|0.760|", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "The paper extends prior frameworks MOE (Mixture of Experts) to address unresolved limitations in fairness learning in the medical image segmentation field", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "1、Strengthsd:\n（1）. Innovation in the Use of MOE (Mixture of Experts): The use of MOE to improve fairness learning in the medical image segmentation field is an innovative angle, particularly in addressing the challenge of imbalanced data distributions across different groups.\n（2）. Application of Control Theory: Control theory, typically used for optimizing and regulating system behaviors, has been applied here to fairness learning in medical image segmentation, providing a novel perspective.\n\n 2、Weaknesses：\n（1）. Insufficient Description of the Patchify Method: The description of the Patchify method (Section 3.2) is rather brief and fails to clearly explain the implementation. There are also formatting and notation errors (such as incorrect subscripts for 'h' in the formula on page 5. It is suggested that the authors provide a more detailed description of the method and correct the formatting issues.\n（2）. Lack of Detailed Explanation of the Softplus Function: The Softplus function is mentioned but not explained in detail. Additionally, the content in Section 3.2 seems to be a minor modification of the MOE method, lacking sufficient innovation or explanation.\n（3）. Inadequate Derivation of Control Theory Formulation: The derivation of the optimal control theory in Section 3.3 appears abrupt, with insufficient transitions and explanation. The authors are encouraged to provide a more complete mathematical derivation, especially regarding how control theory is integrated with the model, to enhance clarity and persuasiveness.\n（4）. Limitations of Experimental Datasets and Design: The experimental datasets are limited in number and relatively small in size. It is recommended that the authors conduct additional experiments on the same dataset using different attribute distributions to further validate the method's effectiveness. Furthermore, expanding the dataset for more comprehensive experiments would be beneficial.\n（5）. Insufficient Ablation Studies: The ablation study section is underdeveloped, and it is unclear to what extent the proposed innovations contribute to improvements, particularly regarding the optimal control theory component. The authors should conduct separate ablation experiments to validate the effectiveness of this part.", "other_comments_or_suggestions": "See the Other Strengths And Weaknesses", "questions_for_authors": "See the Other Strengths And Weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This manuscript proposes a Distribution-aware Mixture of Experts (dMoE) approach for fairness-aware medical image segmentation, reinterpreting the MoE framework through control theory. The work demonstrates how incorporating demographic and clinical attributes into expert selection can improve segmentation fairness across underrepresented groups. The reviewers highlighted several key strengths, including the novel integration of control theory principles, comprehensive evaluation across multiple medical imaging domains, and demonstrated improvements in fairness metrics. The authors effectively addressed initial concerns by expanding their test dataset for prostate cancer segmentation, incorporating additional clinical attributes like Gleason Grade Groups, and providing statistical validation through bootstrapped confidence intervals. While reviewers initially expressed concerns about limited dataset sizes and theoretical foundations, the authors' detailed rebuttal, including expanded experiments and clarified mathematical derivations, largely satisfied these concerns. The discussion phase showed evolving reviewer perspectives, with Reviewer Fj34 raising their score to weak accept despite some remaining questions about the technical distinction between dMoE and standard MoE, while Reviewer V3Mq maintained a strong accept after their concerns were addressed. Reviewer qTb4 acknowledged the authors' thorough response regarding statistical validation and expanded dataset analysis. The consensus among reviewers improved following the rebuttal, with particular appreciation for the practical clinical relevance and theoretical grounding of the approach.
Decision: Accept (spotlight poster)