Debiased Noise Editing on Foundation Models for Fair Medical Image Classification

Fairness Issue: Fig. 1 (a) shows there exists an association between sensitive attributes $\mathcal{A}$ and targets $\mathcal{Y}$ (dashed line). In the Empirical Risk Minimization (ERM) [22], the model $f(\phi(\cdot))$ superficially consider $\mathcal{A}$ as a proxy of $\mathcal{Y}$, causing the issue of fairness. This phenomenon persists even when there is an equal number of data points for each class in both $\mathcal{Y}$ and $\mathcal{A}$, as demonstrated in Sec. 3.1. Our approach aims to address these disparities using debiased noise edit.

Debiased Noise Editing. Given the spurious correlation between SA $\mathcal{A}$ and disease $\mathcal{Y}$, we aim to remove this link through intervention on $\mathcal{X}$ by adding minimal editing vector $\epsilon$, known as DNE, concealing the spurious correlations while preserving image utility. As shown in Fig. 1 (b), we first obtain a pre-trained SA classifier for medical images by either leveraging an existing one (if available) or using the collected $[X_{A},A]$ to train the classifier $g$ with cross-entropy loss. Then, we freeze the classifier and leverage gradient ascent to learn the minimal DNE noise $\epsilon$, which is constrained by certain threshold to preserve image fidelity. $\epsilon$ is trained to deceive the SA classifier $g$ without compromising useful information by solving the following optimization problem:

where $\lambda$ regularizes the magnitude of $\epsilon$. For example, larger $\lambda$ enforces the small L2-norm of the $\epsilon$ in the process of gradient descent. Sec. 3.4 explores its effect in detail. Importantly, we also offer a zero-order strategy for black-box APIs, as detailed in Sec. 2.3.

Fair Disease Classification Having obtained the DNE noise $\epsilon$ to conceal the SA $\mathcal{A}$, we then freeze and add it to disease classification images $X_{T}$. As shown in Fig. 1 (c), this action guarantees that the FM generates a fair embedding devoid of sensitive information, denoted as $\hat{z}_{i}=\phi(x_{i}+\epsilon)$. Finally, these fair embeddings are leveraged by the subsequent disease classifier $f$ for prediction. To train the classifier $f$, we employ a cross-entropy loss:

where classifier $f$ is the only trainable parameter and $M$ is the size of $X_{T}$.

Dataset. To demonstrate the generalizability of DNE. we adopt the CheXpert dataset [7], a chest X-ray dataset with multiple disease labels, to predict the binary label for Pleural Effusion, Pneumonia and Edema individually in chest radiographs. All the ablation studies are performed on Pleural Effusion classification task for space limit. We take gender bias (male and female) as an example due to its broad impact on society and medical imaging analysis. To demonstrate the effectiveness of bias mitigation methods, we follow [4] to amplify the training data bias for each disease by (1) firstly dividing the data into different groups according to the SA; (2) secondly calculating the positive rate of each subgroup; (3) sampling subsets from the original training dataset and increase each subgroup’s bias gap (more positive sample in a subgroup). Then, we sample testing data with the same procedure, but achieve an equal subgroup bias gap. The detailed data distribution is shown in Table 3 in App. 0.C.

Evaluation metrics. We use the classification accuracy to evaluate the utility of classifiers on the test set. To measure fairness, we employ equal opportunity (EO) [6] and disparate impact (DI) [13] metrics. EO aims to ensure equitable prediction probabilities across different groups, defined by the sensitive attribute $a^{j}$, for a given class $y^{j}$. It quantifies the disparity in true positive rates between groups: $\text{EO}_{Y=y_{1}}=P(\hat{Y}=y^{1}|Y=y^{1},A=a^{1})-P(\hat{Y}=y^{1}|Y=y^{1},A=a^{2})$, where a smaller gap signifies greater equality of opportunity. DI evaluates the presence of indirect discrimination by measuring the ratio of positive predictions across different groups: $\text{DI}=\frac{P(\hat{Y}=1|A=a^{1})}{P(\hat{Y}=1|A=a^{2})}$. DI closing to 1 indicates the minimal disparity in positive prediction rates between the groups. To align with EO metric, we employ the $|1-\textsc{DI}|$ to quantify the fairness in percentage, with smaller values denoting greater fairness.

Architecture. In our implementation, all methods use the same architecture. To simulate the FM API, we utilize a publicly available self-supervised pre-trained ViT-base model ²²2https://github.com/lambert-x/Medical_MAE, optimized on X-ray images, given its superior performance metrics among all other architectures [26]. Within our study, we fix the encoder of the ViT model and only fine-tune the classification head, simulating the configuration detailed in Sec. 2.1.

Debiased Noise Editing. For DNE, we first fine-tune the SA classifier, $g$, using the feature encoded by the FM ($\phi(x_{i})$) with the training set of Pleural Effusion in Table 3. We update it using Adam with a learning rate (lr) of $0.0001$ for $50$ epochs. Second, we update the DNE through implementing Eq. 1 by initializing $\epsilon$ as a trainable PyTorch parameter with all entries initially set to zero. The parameter matches the input image dimensions of $224\times 224$. The DNE can be optimized using classic gradient descent or GeZO. Finally, we fine-tune the disease classifier, $f$, with $\epsilon$ added on the input data as delineated in Eq. 2. We update it using AdamW with lr of $1.25\times 10^{-4}$ for $50$ epochs. We provide the visualization of the magnitude of DNE as an interpretation of DNE in App. 0.B Fig. 3, where DNE is smoothed by the Gaussian kernel. As shown, larger noises are added to the bottom to discriminate the gender-related features.

GeZO. Here, we keep all parts for DNE the same as above, except using GeZO to update the edit rather than the Adam, given that the gradient is not accessible. The only key parameter that we are interested in is the local iteration $T$, where we investigated it in Sec. 3.4. Furthermore, the detailed implementation of the algorithm and hyperparameter setting for GeZO can be found in App. 0.D.

Effect of Regularization Coefficient. To investigate the effect of different regularization coefficients’ ($\lambda$) effect, we vary $\lambda$ from 0 to 1. Fig. 2 (a) depicts the EO and the Acc for different $\lambda$s. The metrics for ERM are labeled as horizontal dashed lines for convenience. As shown, both the fairness and utility metrics remain relatively stable for $\lambda<1$, consistently surpassing the ERM baseline. However, as $\lambda$ increases to 1, we observe a marked decline in accuracy below the ERM benchmark, along with a notable increase in the EO metric for both classes.

Effect of Local Iterations in GeZO. In black-box API, total local iterations $R$ in GeZO affect the optimization performance, wherein larger $R$ leads to more accurate optimization at the cost of efficiency. Here, we examine the effect of changing $R$, as shown in Fig. 2 (b). For fairness metrics, increasing $R$ from 2 to 20 significantly reduces the EO score for both classes, demonstrating a considerable debiasing impact with more local epochs. This effect is attributed to increased perturbation sampling that expands the search space with more local epochs, as introduced in Sec. 2.3. Meanwhile, accuracy remains relatively stable, with minor fluctuations between 72% and 73%.