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

R3-W1.1 The Patchify (Section 3.2) fails to clearly explain the implementation.

• 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.

R3-W1.2 There are also formatting and notation errors (such as incorrect subscripts for 'h' in the formula on page 5).

• 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.

R3-W2.1 Lack of detailed explanation of the Softplus.

• We will update the explanation: “Softplus is a smooth alternative activation function to ReLU, which is defined as SoftPlus(x) = log(1+e^x)”.

R3-W2.2 Section 3.2 seems to be a minor modification of the MOE, lacking sufficient innovation or explanation.

• 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.

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.

• 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.

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.

• 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.

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.

• We performed an ablation study on Optimal Control components for radiotherapy target segmentation, as shown in Resp_Table 6.

Resp_Table 6. Ablation study on Optimal Control components.

• 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.

Resp_Table 7. Comparison to multiple networks for each individual attribute.

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.

• 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.

Resp_Table 1. Radiotherapy target segmentation with tumor stage on the expanded testset.

• 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.
• Reference: [1] Yu Tian et al., ICLR 2024, https://arxiv.org/abs/2311.02189

Resp_Table 2. Radiotherapy target segmentation with Gleason Grade Groups (GG).

• 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.

Resp_Table 3. Harvard-FairSeg dataset with gender.

Resp_Table 4. HAM10000 dataset with gender.

R2-W2 The lack of statistical significance in the fairness evaluations.

• 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.

Resp_Table 5. Statistical significance analyzed Dice metric and an additional Worst-group accuracy metric for radiotherapy target segmentation.

R2-W3 Could you clarify how ESSP relates to common fairness metrics, such as demographic parity, equalized odds, and worst-group accuracy?

• 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.

R1-W1.1 How is dMoE more distribution aware compared to a normal MoE?

• 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.

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.

• 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.
• 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.
• 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.

R1-S1 Highlighting the framework's computational efficiency compared to training multiple networks for each attribute would be beneficial.

• Due to the character limit, we kindly refer the reviewer to Resp_Table 7 in response to R3-W5, which demonstrates dMoE's efficiency.

R1-W2 Why incorporate W and W_noise?

• 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.

R1-W3 The description for control theory is confusing.

• Please see inline-answers bellow.

[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?

• Please refer to R3-W1.1.

[line 156] Normal() is a weird notation.

• It should be updated to N(0,1), the standard normal distribution.

[E.q 8] What is f, \theta [line 208, right] the shape of \theta? [E.q 9, 12] What is u? t?

• Please refer to R3-W1.2 and R3-W3.

[line 205, right] What is u_i^i(h_t^i)?

• 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​.

[line 240] What is the system parameter? manual?

• 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.

R1-W4 The experiments cannot justify the additional benefit in considering both demographic and contextual factors.

• 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.

R1-W5 dMoE does not outperform MoE in all subgroups.

• 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.

R1-W6 Interpreting dMoE via optimal control doesn’t bring additional insights.

• 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.